jpmschweitzerandClaude Opus 4.6
98f8cbab0d
- Add session_hash to all API calls (maintains gr.State between image_to_3d and extract_glb) - Pass is_multiimage=False at position 2 (9-param image_to_3d) - Pass output_buf=None at position 0 (3-param extract_glb) - Document full API parameter reference and failure modes - Add trellis-batch.sh for gentle sequential batch generation Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
361 lines
13 KiB
Python
Executable File
361 lines
13 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
Trellis 3D model generator connector — Gradio API wrapper.
|
||
|
||
Talks to the Trellis Gradio app at tower-of-joy:11510.
|
||
Pipeline: upload image → start session → image_to_3d → extract_glb → download .glb
|
||
|
||
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
|
||
import json
|
||
import os
|
||
import shutil
|
||
import sys
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
import urllib.parse
|
||
|
||
CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")
|
||
|
||
|
||
def load_config():
|
||
with open(CONFIG_PATH) as f:
|
||
return json.load(f)
|
||
|
||
|
||
def get_base_url():
|
||
config = load_config()
|
||
return config.get("trellis_url", "http://tower-of-joy:11510")
|
||
|
||
|
||
def health():
|
||
"""Check if the Trellis API is reachable."""
|
||
base = get_base_url()
|
||
try:
|
||
req = urllib.request.Request(f"{base}/info", method="GET")
|
||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||
data = json.loads(resp.read())
|
||
endpoints = list(data.get("named_endpoints", {}).keys())
|
||
print(json.dumps({
|
||
"ok": True,
|
||
"url": base,
|
||
"endpoints": endpoints
|
||
}, indent=2))
|
||
except Exception as e:
|
||
print(json.dumps({
|
||
"ok": False,
|
||
"url": base,
|
||
"error": str(e)
|
||
}, indent=2))
|
||
sys.exit(1)
|
||
|
||
|
||
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}"
|
||
|
||
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(),
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST"
|
||
)
|
||
|
||
print(f" Calling {endpoint}...", file=sys.stderr)
|
||
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
result = json.loads(resp.read())
|
||
if isinstance(result, dict) and "data" in result:
|
||
return result["data"]
|
||
return result
|
||
except urllib.error.HTTPError as e:
|
||
body = e.read().decode("utf-8", errors="replace")
|
||
raise RuntimeError(f"{endpoint} failed ({e.code}): {body[:300]}")
|
||
|
||
|
||
def _upload_image(base, image_path):
|
||
"""Upload an image file to the Gradio server and return the file reference."""
|
||
upload_url = f"{base}/upload"
|
||
|
||
with open(image_path, "rb") as f:
|
||
image_data = f.read()
|
||
|
||
filename = os.path.basename(image_path)
|
||
|
||
# Gradio upload expects multipart/form-data with a 'files' field
|
||
boundary = "----TrellisConnectorBoundary"
|
||
body = (
|
||
f"--{boundary}\r\n"
|
||
f'Content-Disposition: form-data; name="files"; filename="{filename}"\r\n'
|
||
f"Content-Type: image/png\r\n"
|
||
f"\r\n"
|
||
).encode() + image_data + f"\r\n--{boundary}--\r\n".encode()
|
||
|
||
req = urllib.request.Request(
|
||
upload_url,
|
||
data=body,
|
||
headers={
|
||
"Content-Type": f"multipart/form-data; boundary={boundary}",
|
||
},
|
||
method="POST"
|
||
)
|
||
|
||
print(f" Uploading {filename}...", file=sys.stderr)
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
result = json.loads(resp.read())
|
||
# Gradio returns a list of uploaded file paths
|
||
if isinstance(result, list) and len(result) > 0:
|
||
return result[0]
|
||
raise RuntimeError(f"Upload failed: {result}")
|
||
|
||
|
||
def _download_file(url, output_path, base):
|
||
"""Download a file from the Gradio server."""
|
||
if url.startswith("/"):
|
||
url = f"{base}{url}"
|
||
elif not url.startswith("http"):
|
||
url = f"{base}/file={url}"
|
||
|
||
print(f" Downloading to {output_path}...", file=sys.stderr)
|
||
req = urllib.request.Request(url, method="GET")
|
||
with urllib.request.urlopen(req, timeout=120) as resp:
|
||
with open(output_path, "wb") as f:
|
||
shutil.copyfileobj(resp, f)
|
||
|
||
return os.path.getsize(output_path)
|
||
|
||
|
||
def _check_available(base):
|
||
"""Quick check if Trellis is reachable. Fail fast with a clear message."""
|
||
try:
|
||
req = urllib.request.Request(f"{base}/info", method="GET")
|
||
urllib.request.urlopen(req, timeout=5)
|
||
except Exception:
|
||
print(json.dumps({
|
||
"ok": False,
|
||
"error": f"Trellis is not available at {base}. The service may be switched off to save system resources. Start it before generating 3D models."
|
||
}, indent=2))
|
||
sys.exit(1)
|
||
|
||
|
||
def generate(image_path, output=None, simplify=0.95, texture_size=1024,
|
||
seed=0, timeout=600):
|
||
"""
|
||
Generate a 3D model from an image.
|
||
|
||
Pipeline:
|
||
1. Start session
|
||
2. Upload and preprocess image
|
||
3. Generate 3D from image
|
||
4. Extract GLB
|
||
5. Download GLB file
|
||
|
||
Args:
|
||
image_path: Path to the input image (PNG recommended)
|
||
output: Output .glb file path (default: auto-named)
|
||
simplify: Mesh simplification factor (0.9-0.98, default 0.95)
|
||
texture_size: Texture resolution (512-2048, default 1024)
|
||
seed: Random seed (default 0)
|
||
timeout: Max wait time per step in seconds
|
||
"""
|
||
base = get_base_url()
|
||
_check_available(base)
|
||
start_time = time.time()
|
||
|
||
if not os.path.isfile(image_path):
|
||
print(json.dumps({"ok": False, "error": f"Image not found: {image_path}"}), indent=2)
|
||
sys.exit(1)
|
||
|
||
if output is None:
|
||
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)
|
||
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)
|
||
uploaded_path = _upload_image(base, image_path)
|
||
file_ref = {
|
||
"path": uploaded_path,
|
||
"meta": {"_type": "gradio.FileData"}
|
||
}
|
||
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:
|
||
preprocessed_ref = preprocess_result[0]
|
||
else:
|
||
preprocessed_ref = preprocess_result
|
||
|
||
# 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, 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
|
||
# 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)
|
||
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, 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", [None, simplify, texture_size], timeout=120, session_hash=session)
|
||
|
||
# _call_api returns the "data" array: [model_viewer_data, download_button_data]
|
||
glb_url = None
|
||
if isinstance(glb_result, list):
|
||
for item in glb_result:
|
||
if isinstance(item, dict):
|
||
url = item.get("url") or item.get("path")
|
||
if url:
|
||
glb_url = url
|
||
break
|
||
|
||
if not glb_url:
|
||
print(json.dumps({
|
||
"ok": False,
|
||
"error": "Could not extract GLB URL from response",
|
||
"response": glb_result
|
||
}, indent=2))
|
||
sys.exit(1)
|
||
|
||
# Step 6: Download
|
||
print("Step 5/5: Downloading GLB...", file=sys.stderr)
|
||
os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True)
|
||
file_size = _download_file(glb_url, output, base)
|
||
|
||
elapsed = round(time.time() - start_time, 1)
|
||
print(json.dumps({
|
||
"ok": True,
|
||
"file": output,
|
||
"size_bytes": file_size,
|
||
"simplify": simplify,
|
||
"texture_size": texture_size,
|
||
"seed": actual_seed,
|
||
"generation_time_s": elapsed,
|
||
"source_image": image_path
|
||
}, indent=2))
|
||
|
||
|
||
def main():
|
||
if len(sys.argv) < 2:
|
||
print("Usage:")
|
||
print(" trellis_connector.py health")
|
||
print(" trellis_connector.py generate image.png [--output model.glb] [--simplify 0.95] [--texture-size 1024] [--seed N] [--timeout N]")
|
||
sys.exit(1)
|
||
|
||
cmd = sys.argv[1]
|
||
|
||
if cmd == "health":
|
||
health()
|
||
elif cmd == "generate":
|
||
if len(sys.argv) < 3:
|
||
print("Error: image path required", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
image_path = sys.argv[2]
|
||
output = None
|
||
simplify = 0.95
|
||
texture_size = 1024
|
||
seed = 0
|
||
timeout = 600
|
||
|
||
i = 3
|
||
while i < len(sys.argv):
|
||
if sys.argv[i] == "--output" and i + 1 < len(sys.argv):
|
||
output = sys.argv[i + 1]
|
||
i += 2
|
||
elif sys.argv[i] == "--simplify" and i + 1 < len(sys.argv):
|
||
simplify = float(sys.argv[i + 1])
|
||
i += 2
|
||
elif sys.argv[i] == "--texture-size" and i + 1 < len(sys.argv):
|
||
texture_size = int(sys.argv[i + 1])
|
||
i += 2
|
||
elif sys.argv[i] == "--seed" and i + 1 < len(sys.argv):
|
||
seed = int(sys.argv[i + 1])
|
||
i += 2
|
||
elif sys.argv[i] == "--timeout" and i + 1 < len(sys.argv):
|
||
timeout = int(sys.argv[i + 1])
|
||
i += 2
|
||
else:
|
||
print(f"Unknown argument: {sys.argv[i]}", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
generate(image_path, output=output, simplify=simplify,
|
||
texture_size=texture_size, seed=seed, timeout=timeout)
|
||
else:
|
||
print(f"Unknown command: {cmd}", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|