diff --git a/tooling/db/trellis_connector.py b/tooling/db/trellis_connector.py index 3d28b3831..a49035254 100755 --- a/tooling/db/trellis_connector.py +++ b/tooling/db/trellis_connector.py @@ -8,6 +8,32 @@ Pipeline: upload image → start session → image_to_3d → extract_glb → dow Usage: python3 trellis_connector.py health python3 trellis_connector.py generate image.png [--output model.glb] [--simplify 0.95] [--texture-size 1024] [--seed 42] [--timeout 600] + +Gradio API Parameter Reference (TRELLIS v1, microsoft/TRELLIS): + + /image_to_3d — 9 inputs: + 0: image (Image) preprocessed image from /preprocess_image_1 + 1: multiimages (Gallery) [] for single-image mode + 2: is_multiimage (State) False for single-image, True for multi-image + 3: seed (Slider) int, 0-2147483647 + 4: ss_guidance (Slider) float, sparse structure guidance strength (default 7.5) + 5: ss_steps (Slider) int, sparse structure sampling steps (default 12) + 6: slat_guidance (Slider) float, structured latent guidance strength (default 3.0) + 7: slat_steps (Slider) int, structured latent sampling steps (default 12) + 8: multiimage_algo (Radio) "stochastic" or "multidiffusion" + + /extract_glb — 3 inputs: + 0: output_buf (State) None — server uses internal state from image_to_3d + 1: simplify (Slider) float, mesh simplification ratio (default 0.95) + 2: texture_size (Slider) int, texture resolution (default 1024) + + Common failure modes: + - "needed 9, got 8": missing is_multiimage (position 2) — must pass False + - "needed 3, got 2": missing output_buf (position 0) — must pass None + - "'float' cannot be interpreted as int": numpy version issue on server, + Gradio Sliders send all values as float. Fix: patch flow_euler.py on + the server to cast steps to int, or pin numpy < 2.0 + - CUDA device mismatch after crash: restart the container to clear GPU state """ import base64 @@ -55,12 +81,20 @@ def health(): sys.exit(1) -def _call_api(base, endpoint, data, timeout=600): - """Call a Gradio API endpoint. Tries sync /api/ first, falls back to SSE /gradio_api/call/.""" - # Trellis uses the sync /api/ pattern +def _call_api(base, endpoint, data, timeout=600, session_hash=None): + """Call a Gradio API endpoint with optional session tracking. + + Gradio gr.State components are stored per session_hash on the server. + All calls in a pipeline (image_to_3d → extract_glb) must share the same + session_hash so the server can pass state between them. + """ api_url = f"{base}/api{endpoint}" - payload = json.dumps({"data": data}) + body_dict = {"data": data} + if session_hash: + body_dict["session_hash"] = session_hash + + payload = json.dumps(body_dict) req = urllib.request.Request( api_url, data=payload.encode(), @@ -73,7 +107,6 @@ def _call_api(base, endpoint, data, timeout=600): try: with urllib.request.urlopen(req, timeout=timeout) as resp: result = json.loads(resp.read()) - # Sync Gradio returns {"data": [...], "is_generating": false, ...} if isinstance(result, dict) and "data" in result: return result["data"] return result @@ -179,9 +212,16 @@ def generate(image_path, output=None, simplify=0.95, texture_size=1024, name = os.path.splitext(os.path.basename(image_path))[0] output = f"{name}.glb" + # Generate a session hash — Gradio uses this to maintain gr.State between + # separate API calls. Without it, image_to_3d's output state is lost before + # extract_glb can read it. + import random, string + session = ''.join(random.choices(string.ascii_lowercase + string.digits, k=12)) + print(f" Session: {session}", file=sys.stderr) + # Step 1: Start session print("Step 1/5: Starting session...", file=sys.stderr) - _call_api(base, "/start_session", [], timeout=30) + session_result = _call_api(base, "/start_session", [], timeout=30, session_hash=session) # Step 2: Upload and preprocess image print("Step 2/5: Uploading and preprocessing image...", file=sys.stderr) @@ -190,7 +230,7 @@ def generate(image_path, output=None, simplify=0.95, texture_size=1024, "path": uploaded_path, "meta": {"_type": "gradio.FileData"} } - preprocess_result = _call_api(base, "/preprocess_image_1", [file_ref], timeout=60) + preprocess_result = _call_api(base, "/preprocess_image_1", [file_ref], timeout=60, session_hash=session) # _call_api returns the "data" array directly if isinstance(preprocess_result, list) and len(preprocess_result) > 0: @@ -200,28 +240,34 @@ def generate(image_path, output=None, simplify=0.95, texture_size=1024, # Step 3: Get seed print("Step 3/5: Generating 3D model...", file=sys.stderr) - seed_result = _call_api(base, "/get_seed", [True, seed], timeout=10) + seed_result = _call_api(base, "/get_seed", [True, seed], timeout=10, session_hash=session) if isinstance(seed_result, list) and seed_result: actual_seed = seed_result[0] else: actual_seed = seed # Step 4: Image to 3D - # Parameters: image, multiimages, seed, ss_guidance, ss_steps, slat_guidance, slat_steps, algo + # Gradio app has 9 inputs: Image, Gallery, State, Slider×5, Radio + # State is a hidden session component at position 3 — pass None. + # Slider order (from /info): seed, ss_guidance, ss_steps, slat_guidance, slat_steps + # Note: server-side numpy bug requires patching flow_euler.py to cast steps to int. gen_result = _call_api(base, "/image_to_3d", [ preprocessed_ref, # image [], # multiimages (empty) - actual_seed, # seed + False, # is_multiimage (boolean flag, not session state) + actual_seed, # seed (Slider, 0-2147483647) 7.5, # ss_guidance_strength 12, # ss_sampling_steps 3.0, # slat_guidance_strength 12, # slat_sampling_steps "stochastic", # multiimage_algo - ], timeout=timeout) + ], timeout=timeout, session_hash=session) # Step 5: Extract GLB + # State (output_buf) is maintained server-side via session_hash. + # extract_glb reads it automatically. Pass None as placeholder for the State component. print("Step 4/5: Extracting GLB...", file=sys.stderr) - glb_result = _call_api(base, "/extract_glb", [simplify, texture_size], timeout=120) + glb_result = _call_api(base, "/extract_glb", [None, simplify, texture_size], timeout=120, session_hash=session) # _call_api returns the "data" array: [model_viewer_data, download_button_data] glb_url = None diff --git a/tooling/trellis-batch.sh b/tooling/trellis-batch.sh new file mode 100755 index 000000000..56278767e --- /dev/null +++ b/tooling/trellis-batch.sh @@ -0,0 +1,93 @@ +#!/usr/bin/env bash +# Batch Trellis generation — one at a time, gently. +# Usage: ./tooling/trellis-batch.sh +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" +cd "$PROJECT_ROOT" + +INPUT_DIR=".tmp/image-gen/characters/bodies" +OUTPUT_DIR=".tmp/glb-gen/characters/bodies" +COOLDOWN=15 # seconds between jobs +MAX_RETRIES=3 +RETRY_DELAY=60 # seconds before retry + +mkdir -p "$OUTPUT_DIR" + +TYPES=( + slim_m slim_f + average_m average_f + stocky_m stocky_f + tall_lean_m tall_lean_f + short_stout_m short_stout_f + athletic_m athletic_f + heavyset_m heavyset_f + petite_m petite_f +) + +TOTAL=${#TYPES[@]} +SUCCESS=0 +FAILED=0 + +echo "=== Trellis Batch: $TOTAL bodies ===" +echo " Cooldown: ${COOLDOWN}s between jobs" +echo " Retries: $MAX_RETRIES with ${RETRY_DELAY}s delay" +echo "" + +for i in "${!TYPES[@]}"; do + TYPE="${TYPES[$i]}" + NUM=$((i + 1)) + INPUT="$INPUT_DIR/${TYPE}.png" + OUTPUT="$OUTPUT_DIR/${TYPE}.glb" + + # Skip if already generated + if [ -f "$OUTPUT" ]; then + echo "[$NUM/$TOTAL] $TYPE — already exists, skipping" + SUCCESS=$((SUCCESS + 1)) + continue + fi + + if [ ! -f "$INPUT" ]; then + echo "[$NUM/$TOTAL] $TYPE — input not found: $INPUT" + FAILED=$((FAILED + 1)) + continue + fi + + ATTEMPT=0 + DONE=false + while [ "$ATTEMPT" -lt "$MAX_RETRIES" ] && [ "$DONE" = "false" ]; do + ATTEMPT=$((ATTEMPT + 1)) + echo "[$NUM/$TOTAL] $TYPE (attempt $ATTEMPT/$MAX_RETRIES)..." + + if python3 tooling/db/trellis_connector.py generate \ + "$INPUT" \ + --output "$OUTPUT" \ + --simplify 0.95 \ + --texture-size 1024 \ + 2>&1 | tee /dev/stderr | grep -q '"ok": true'; then + echo " OK" + SUCCESS=$((SUCCESS + 1)) + DONE=true + else + echo " FAILED" + if [ "$ATTEMPT" -lt "$MAX_RETRIES" ]; then + echo " Waiting ${RETRY_DELAY}s before retry..." + sleep "$RETRY_DELAY" + else + echo " Giving up on $TYPE after $MAX_RETRIES attempts" + FAILED=$((FAILED + 1)) + fi + fi + done + + # Cooldown between successful jobs + if [ "$DONE" = "true" ] && [ "$NUM" -lt "$TOTAL" ]; then + echo " Cooling down ${COOLDOWN}s..." + sleep "$COOLDOWN" + fi +done + +echo "" +echo "=== Results: $SUCCESS/$TOTAL succeeded, $FAILED failed ===" +ls -la "$OUTPUT_DIR"/*.glb 2>/dev/null || echo "No GLB files generated"