"""Image -> 3D (TRELLIS.2-4B) -> auto-rigged 3D (SkinTokens / TokenRig) -> Blender-ready GLB/FBX. Combines: - microsoft/TRELLIS.2 (image-to-3D generation, MIT) - VAST-AI/SkinTokens (TokenRig automatic skeleton + skinning, MIT) """ import spaces # must be imported before torch / anything CUDA import os os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" os.environ["ATTN_BACKEND"] = "flash_attn" os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join( os.path.dirname(os.path.abspath(__file__)), "autotune_cache.json" ) os.environ.setdefault("XFORMERS_IGNORE_FLASH_VERSION_CHECK", "1") # --------------------------------------------------------------------------- # Install the bundled `bpy` wheel at runtime if it isn't already importable. # # Why (inherited from the SkinTokens Space): # - Putting the wheel in requirements.txt fails: HF Spaces' Docker build # mounts only requirements.txt BEFORE the repo COPY, so the wheel path # doesn't exist at pip-install time. # - PyPI doesn't ship a bpy wheel matching this exact build. # - The wheel committed in this repo is auto-tracked by LFS, so the # container COPY may only carry the LFS *pointer*; in that case we # re-fetch the real wheel from the Hub at runtime. # --------------------------------------------------------------------------- def _ensure_bpy_installed(): try: import bpy # noqa: F401 return except Exception: pass import glob import sysconfig import zipfile here = os.path.dirname(os.path.abspath(__file__)) wheels = sorted(glob.glob(os.path.join(here, "bpy-*.whl"))) if not wheels: print("[app] WARNING: bpy not importable and no bundled wheel found", flush=True) return wheel = wheels[-1] wheel_name = os.path.basename(wheel) is_real_zip = False try: with open(wheel, "rb") as f: is_real_zip = f.read(4).startswith(b"PK") except Exception: pass if not is_real_zip: print( f"[app] {wheel_name} on disk is an LFS pointer ({os.path.getsize(wheel)} B); " f"fetching real wheel from HF Hub...", flush=True, ) from huggingface_hub import hf_hub_download space_id = os.environ.get("SPACE_ID", "JSCPPProgrammer/image-to-rigged-3d") wheel = hf_hub_download( repo_id=space_id, repo_type="space", filename=wheel_name, token=os.environ.get("HF_TOKEN"), ) print(f"[app] fetched -> {wheel} ({os.path.getsize(wheel)} B)", flush=True) site = sysconfig.get_paths()["purelib"] print(f"[app] Extracting {wheel_name} into {site}", flush=True) with zipfile.ZipFile(wheel) as z: z.extractall(site) print("[app] bpy wheel extracted.", flush=True) _ensure_bpy_installed() # --------------------------------------------------------------------------- # Download the TokenRig / SkinTokens checkpoints and the Qwen3 tokenizer on # first cold start (they live in the public model repo VAST-AI/SkinTokens). # --------------------------------------------------------------------------- def _ensure_rig_models_downloaded(): here = os.path.dirname(os.path.abspath(__file__)) needed_ckpts = [ "experiments/skin_vae_2_10_32768/last.ckpt", "experiments/articulation_xl_quantization_256_token_4/grpo_1400.ckpt", ] qwen_dir = os.path.join(here, "models", "Qwen3-0.6B") all_present = all( os.path.exists(os.path.join(here, p)) for p in needed_ckpts ) and os.path.exists(os.path.join(qwen_dir, "tokenizer.json")) if all_present: return from huggingface_hub import hf_hub_download, snapshot_download for rel in needed_ckpts: if os.path.exists(os.path.join(here, rel)): continue print(f"[app] Downloading rigging checkpoint: {rel}", flush=True) hf_hub_download(repo_id="VAST-AI/SkinTokens", filename=rel, local_dir=here) if not os.path.exists(os.path.join(qwen_dir, "tokenizer.json")): print("[app] Downloading Qwen3-0.6B tokenizer/config", flush=True) snapshot_download( repo_id="Qwen/Qwen3-0.6B", local_dir=qwen_dir, ignore_patterns=["*.bin", "*.safetensors"], ) print("[app] Rigging checkpoints ready.", flush=True) _ensure_rig_models_downloaded() import atexit import shutil import signal import subprocess import sys import time import traceback from datetime import datetime from pathlib import Path from typing import Dict, List, Tuple import gradio as gr import numpy as np import requests import torch from gradio_client import Client, handle_file from PIL import Image from torch import Tensor from trellis2.pipelines import Trellis2ImageTo3DPipeline import o_voxel from src.data.dataset import DatasetConfig, RigDatasetModule from src.data.transform import Transform from src.tokenizer.parse import get_tokenizer from src.server.spec import BPY_SERVER, get_model, object_to_bytes, bytes_to_object MAX_SEED = np.iinfo(np.int32).max TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp") RIG_CKPT = "experiments/articulation_xl_quantization_256_token_4/grpo_1400.ckpt" # --------------------------------------------------------------------------- # Session handling # --------------------------------------------------------------------------- def start_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) def end_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) shutil.rmtree(user_dir, ignore_errors=True) # --------------------------------------------------------------------------- # Image preprocessing (background removal via remote BRIA RMBG space) # --------------------------------------------------------------------------- _rmbg_client = None def _get_rmbg_client(): global _rmbg_client if _rmbg_client is None: _rmbg_client = Client("briaai/BRIA-RMBG-2.0") return _rmbg_client def remove_background(input: Image.Image) -> Image.Image: import tempfile with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f: input.convert("RGB").save(f.name) path = f.name try: output = _get_rmbg_client().predict(handle_file(path), api_name="/image")[0][0] return Image.open(output) finally: try: os.remove(path) except OSError: pass def preprocess_image(input: Image.Image) -> Image.Image: """Resize, remove background (if no alpha), crop to object, premultiply alpha.""" if input is None: return None has_alpha = False if input.mode == "RGBA": alpha = np.array(input)[:, :, 3] if not np.all(alpha == 255): has_alpha = True max_size = max(input.size) scale = min(1, 1024 / max_size) if scale < 1: input = input.resize( (int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS ) if has_alpha: output = input else: output = remove_background(input) output_np = np.array(output) alpha = output_np[:, :, 3] bbox = np.argwhere(alpha > 0.8 * 255) bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0]) center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2 size = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) bbox = ( center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2, ) output = output.crop(bbox) output = np.array(output).astype(np.float32) / 255 output = output[:, :, :3] * output[:, :, 3:4] output = Image.fromarray((output * 255).astype(np.uint8)) return output def get_seed(randomize_seed: bool, seed: int) -> int: return int(np.random.randint(0, MAX_SEED)) if randomize_seed else int(seed) # --------------------------------------------------------------------------- # Per-task progress UI (HTML bars streamed via generator yields) # --------------------------------------------------------------------------- GEN_STEP_NAMES = [ "1. Remove background", "2. Encode image (DINOv3)", "3. Sample sparse 3D structure", "4. Generate mesh shape", "5. Generate PBR textures", "6. Export textured GLB", ] RIG_STEP_NAMES = [ "1. Start Blender server", "2. Load TokenRig model", "3. Load input mesh", "4. Predict skeleton + skinning", "5. Export rigged GLB", "6. Export rigged FBX", ] _STEP_STYLE = { "pending": ("#9ca3af", "○"), "running": ("#7c3aed", "◐"), "done": ("#22c55e", "✓"), "error": ("#ef4444", "✗"), } def _init_steps(names: List[str]) -> List[Dict]: return [{"name": n, "state": "pending", "pct": 0} for n in names] def _render_steps(steps: List[Dict]) -> str: rows = [] for s in steps: color, icon = _STEP_STYLE[s["state"]] pct = s.get("pct", 0) rows.append( f'
' f'
' f"{icon} {s['name']}{pct}%
" f'
' f'
' ) return f'
{"".join(rows)}
' def _mark(steps: List[Dict], idx: int, state: str, pct: int) -> str: steps[idx]["state"] = state steps[idx]["pct"] = pct return _render_steps(steps) def _noop_files(): return gr.update(), gr.update(), gr.update() def preprocess_with_progress(image: Image.Image): """Step 1 runs on CPU (not ZeroGPU) so BRIA RMBG doesn't eat the GPU quota.""" steps = _init_steps(GEN_STEP_NAMES) if image is None: raise gr.Error("Please upload an image first.") yield image, _mark(steps, 0, "running", 10), "Step 1/6: Removing background (BRIA RMBG)…" try: processed = preprocess_image(image) except Exception as e: tb = traceback.format_exc() print(tb, flush=True) yield image, _mark(steps, 0, "error", 100), f"FAILED at step 1: {e}" raise gr.Error(f"Background removal failed: {e}") from e steps[0]["state"] = "done" steps[0]["pct"] = 100 yield processed, _render_steps(steps), "Step 1/6 done. Queuing TRELLIS.2 on ZeroGPU…" def _gpu_duration_generate(*args, **kwargs): # ZeroGPU forwards the full handler arg list (including gr.Progress). resolution = kwargs.get("resolution") if resolution is None and len(args) >= 3: resolution = args[2] return {"512": 300, "1024": 480, "1536": 600}.get(str(resolution), 480) # --------------------------------------------------------------------------- # Stage 1: Image -> textured GLB (TRELLIS.2-4B) # --------------------------------------------------------------------------- @spaces.GPU(duration=_gpu_duration_generate) def generate_3d( image: Image.Image, seed: int, resolution: str, decimation_target: int, texture_size: int, req: gr.Request, progress=gr.Progress(track_tqdm=True), ): steps = [{"name": GEN_STEP_NAMES[0], "state": "done", "pct": 100}] steps += _init_steps(GEN_STEP_NAMES[1:]) ptype = {"512": "512", "1024": "1024_cascade", "1536": "1536_cascade"}[resolution] ss_params = { "steps": 12, "guidance_strength": 7.5, "guidance_rescale": 0.7, "rescale_t": 5.0, } shape_params = { "steps": 12, "guidance_strength": 7.5, "guidance_rescale": 0.5, "rescale_t": 3.0, } tex_params = { "steps": 12, "guidance_strength": 1.0, "guidance_rescale": 0.0, "rescale_t": 3.0, } if image is None: raise gr.Error("Please upload an image first.") try: # Step 2 — encode image features yield (*_noop_files(), _mark(steps, 1, "running", 15), "Step 2/6: Encoding image (DINOv3)…") torch.manual_seed(seed) cond_512 = pipeline.get_cond([image], 512) cond_1024 = pipeline.get_cond([image], 1024) if ptype != "512" else None yield (*_noop_files(), _mark(steps, 1, "done", 100), "Step 2/6 done.") # Step 3 — sparse structure yield (*_noop_files(), _mark(steps, 2, "running", 20), "Step 3/6: Sampling sparse 3D structure…") ss_res = {"512": 32, "1024": 64, "1024_cascade": 32, "1536_cascade": 32}[ptype] coords = pipeline.sample_sparse_structure(cond_512, ss_res, 1, ss_params) yield (*_noop_files(), _mark(steps, 2, "done", 100), "Step 3/6 done.") # Step 4 — mesh shape yield (*_noop_files(), _mark(steps, 3, "running", 30), "Step 4/6: Generating mesh shape…") if ptype == "512": shape_slat = pipeline.sample_shape_slat( cond_512, pipeline.models["shape_slat_flow_model_512"], coords, shape_params ) res = 512 elif ptype == "1024": shape_slat = pipeline.sample_shape_slat( cond_1024, pipeline.models["shape_slat_flow_model_1024"], coords, shape_params ) res = 1024 elif ptype == "1024_cascade": shape_slat, res = pipeline.sample_shape_slat_cascade( cond_512, cond_1024, pipeline.models["shape_slat_flow_model_512"], pipeline.models["shape_slat_flow_model_1024"], 512, 1024, coords, shape_params, 49152, ) else: # 1536_cascade shape_slat, res = pipeline.sample_shape_slat_cascade( cond_512, cond_1024, pipeline.models["shape_slat_flow_model_512"], pipeline.models["shape_slat_flow_model_1024"], 512, 1536, coords, shape_params, 49152, ) yield (*_noop_files(), _mark(steps, 3, "done", 100), "Step 4/6 done.") # Step 5 — PBR textures yield (*_noop_files(), _mark(steps, 4, "running", 50), "Step 5/6: Generating PBR textures…") if ptype == "512": tex_slat = pipeline.sample_tex_slat( cond_512, pipeline.models["tex_slat_flow_model_512"], shape_slat, tex_params ) else: tex_slat = pipeline.sample_tex_slat( cond_1024, pipeline.models["tex_slat_flow_model_1024"], shape_slat, tex_params ) mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0] mesh.simplify(16777216) yield (*_noop_files(), _mark(steps, 4, "done", 100), "Step 5/6 done.") # Step 6 — export GLB yield (*_noop_files(), _mark(steps, 5, "running", 70), "Step 6/6: Exporting textured GLB…") glb = o_voxel.postprocess.to_glb( vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs, coords=mesh.coords, attr_layout=pipeline.pbr_attr_layout, grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], decimation_target=decimation_target, texture_size=texture_size, remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True, ) user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) now = datetime.now() timestamp = now.strftime("%Y-%m-%dT%H%M%S") + f".{now.microsecond // 1000:03d}" glb_path = os.path.join(user_dir, f"model_{timestamp}.glb") glb.export(glb_path, extension_webp=False) torch.cuda.empty_cache() yield ( glb_path, glb_path, glb_path, _mark(steps, 5, "done", 100), f"All 6 steps complete — {os.path.basename(glb_path)} ready. Click Auto-Rig.", ) except Exception as e: tb = traceback.format_exc() print(tb, flush=True) for s in steps: if s["state"] == "running": s["state"] = "error" s["pct"] = 100 yield (*_noop_files(), _render_steps(steps), f"FAILED: {e}") raise gr.Error(f"3D generation failed: {e}") from e # --------------------------------------------------------------------------- # bpy_server lifecycle (internal IPC on 127.0.0.1; Blender import is slow, so # it runs once as a detached subprocess and is reused by later GPU workers). # --------------------------------------------------------------------------- _BPY_SERVER_PROC = None def is_bpy_server_alive(timeout: float = 1.0) -> bool: try: return requests.get(f"{BPY_SERVER}/ping", timeout=timeout).status_code == 200 except Exception: return False def start_bpy_server(): proc = subprocess.Popen( [sys.executable, "bpy_server.py"], stdout=None, stderr=None, preexec_fn=os.setsid, ) print(f"[app] bpy_server.py started (pid={proc.pid})", flush=True) def cleanup(): try: os.killpg(os.getpgid(proc.pid), signal.SIGTERM) except (ProcessLookupError, OSError): pass atexit.register(cleanup) return proc def wait_for_bpy_server(timeout: float = 120): t0 = time.time() last_log = 0.0 while True: try: requests.get(f"{BPY_SERVER}/ping", timeout=1) print(f"[app] bpy_server ready after {time.time() - t0:.1f}s", flush=True) return except Exception: now = time.time() if now - t0 > timeout: raise RuntimeError(f"bpy_server failed to start after {timeout:.0f}s") if now - last_log > 10: print(f"[app] waiting for bpy_server ({now - t0:.0f}s elapsed)", flush=True) last_log = now time.sleep(0.5) def ensure_bpy_server_started(): global _BPY_SERVER_PROC if is_bpy_server_alive(): return if _BPY_SERVER_PROC is not None and _BPY_SERVER_PROC.poll() is None: return _BPY_SERVER_PROC = start_bpy_server() wait_for_bpy_server() # --------------------------------------------------------------------------- # Stage 2: GLB -> rigged GLB + FBX (TokenRig) # --------------------------------------------------------------------------- rig_model = None rig_tokenizer = None rig_transform = None def load_rig_model(): global rig_model, rig_tokenizer, rig_transform if rig_model is not None: return print(f"[app] Loading TokenRig model: {RIG_CKPT}", flush=True) rig_model = get_model(RIG_CKPT, hf_path=None) assert rig_model.tokenizer_config is not None rig_tokenizer = get_tokenizer(**rig_model.tokenizer_config) rig_transform = Transform.parse(**rig_model.transform_config["predict_transform"]) print("[app] TokenRig model loaded.", flush=True) def _noop_rig_files(): return gr.update(), gr.update(), gr.update() @spaces.GPU(duration=480) def rig_3d( glb_path: str, num_beams: int, top_k: int, top_p: float, temperature: float, repetition_penalty: float, req: gr.Request, progress=gr.Progress(track_tqdm=True), ): steps = _init_steps(RIG_STEP_NAMES) if not glb_path or not os.path.exists(glb_path): raise gr.Error("Generate a 3D model first.") try: yield (*_noop_rig_files(), _mark(steps, 0, "running", 10), "Step 1/6: Starting Blender export server…") ensure_bpy_server_started() yield (*_noop_rig_files(), _mark(steps, 0, "done", 100), "Step 1/6 done.") yield (*_noop_rig_files(), _mark(steps, 1, "running", 20), "Step 2/6: Loading TokenRig model…") load_rig_model() yield (*_noop_rig_files(), _mark(steps, 1, "done", 100), "Step 2/6 done.") yield (*_noop_rig_files(), _mark(steps, 2, "running", 30), "Step 3/6: Loading input mesh…") datapath = { "data_name": None, "loader": "bpy_server", "filepaths": {"articulation": [str(glb_path)]}, } dataset_config = DatasetConfig.parse( shuffle=False, batch_size=1, num_workers=0, pin_memory=False, persistent_workers=False, datapath=datapath, ).split_by_cls() module = RigDatasetModule( predict_dataset_config=dataset_config, predict_transform=rig_transform, tokenizer=rig_tokenizer, process_fn=rig_model._process_fn, ) dataloader = module.predict_dataloader()["articulation"] infer_device = rig_model.device if rig_model is not None else "cuda" yield (*_noop_rig_files(), _mark(steps, 2, "done", 100), "Step 3/6 done.") user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) stem = Path(glb_path).stem out_glb = os.path.join(user_dir, f"{stem}_rigged.glb") out_fbx = os.path.join(user_dir, f"{stem}_rigged.fbx") yield ( *_noop_rig_files(), _mark(steps, 3, "running", 45), "Step 4/6: Predicting skeleton + skinning weights…", ) for batch in dataloader: batch = { k: v.to(infer_device) if isinstance(v, Tensor) else v for k, v in batch.items() } batch.pop("skeleton_tokens", None) batch.pop("skeleton_mask", None) batch["generate_kwargs"] = dict( max_length=2048, top_k=int(top_k), top_p=float(top_p), temperature=float(temperature), repetition_penalty=float(repetition_penalty), num_return_sequences=1, num_beams=int(num_beams), do_sample=True, ) preds = rig_model.predict_step(batch, skeleton_tokens=None, make_asset=True)["results"] asset = preds[0].asset assert asset is not None yield (*_noop_rig_files(), _mark(steps, 3, "done", 100), "Step 4/6 done.") yield (*_noop_rig_files(), _mark(steps, 4, "running", 70), "Step 5/6: Exporting rigged GLB…") payload = dict( source_asset=asset, target_path=asset.path, export_path=out_glb, group_per_vertex=4, ) res = bytes_to_object( requests.post(f"{BPY_SERVER}/transfer", data=object_to_bytes(payload)).content ) if res != "ok": raise RuntimeError(f"GLB export failed: {res}") yield (*_noop_rig_files(), _mark(steps, 4, "done", 100), "Step 5/6 done.") yield (*_noop_rig_files(), _mark(steps, 5, "running", 85), "Step 6/6: Exporting rigged FBX…") payload["export_path"] = out_fbx res = bytes_to_object( requests.post(f"{BPY_SERVER}/transfer", data=object_to_bytes(payload)).content ) if res != "ok": raise RuntimeError(f"FBX export failed: {res}") torch.cuda.empty_cache() yield ( out_glb, out_glb, out_fbx, _mark(steps, 5, "done", 100), "All 6 rigging steps complete — download GLB or FBX for Blender.", ) except Exception as e: tb = traceback.format_exc() print(tb, flush=True) for s in steps: if s["state"] == "running": s["state"] = "error" s["pct"] = 100 yield (*_noop_rig_files(), _render_steps(steps), f"FAILED: {e}") raise gr.Error(f"Auto-rigging failed: {e}") from e # --------------------------------------------------------------------------- # UI # --------------------------------------------------------------------------- BLENDER_HELP = """ ### Importing into Blender 1. Download the **rigged GLB** (recommended) or **rigged FBX**. 2. In Blender: **File → Import → glTF 2.0** (or **FBX**) and select the file. 3. The model imports with its armature (skeleton) and skinning weights. Select the armature, switch to **Pose Mode**, and rotate bones to pose or animate. 4. If you see a `glTF_not_exported` placeholder node, you can safely delete it. *Note: glTF is the more reliable format — FBX export from Blender's bpy can lose the rest pose on animated assets.* """ with gr.Blocks(title="Image → 3D → Rigged Model", delete_cache=(600, 600)) as demo: gr.Markdown( """ ## 🖼️ → 🧊 → 🦴 Image to Rigged 3D Model — ready for Blender **Step 1**: Upload an image → **Generate 3D** with [TRELLIS.2-4B](https://huggingface.co/microsoft/TRELLIS.2-4B) (Microsoft, state-of-the-art image-to-3D with PBR textures). **Step 2**: **Auto-Rig** the mesh with [SkinTokens / TokenRig](https://huggingface.co/VAST-AI/SkinTokens) (VAST AI, state-of-the-art automatic skeleton + skinning weights, successor to UniRig). **Step 3**: Download the rigged **GLB / FBX** and import it into Blender. *Each stage takes 1–3 minutes of ZeroGPU time. Works on characters, animals, and creatures — anything that should be posable.* """ ) with gr.Row(): with gr.Column(scale=1, min_width=340): image_prompt = gr.Image( label="Input Image", format="png", image_mode="RGBA", type="pil", sources=["upload", "clipboard"], height=320, ) status = gr.Textbox( label="Status", value="Upload an image, then click Generate.", interactive=False, lines=2, ) task_progress = gr.HTML( label="Task progress", value=_render_steps(_init_steps(GEN_STEP_NAMES)), ) generate_btn = gr.Button("1️⃣ Generate 3D Model", variant="primary") rig_btn = gr.Button("2️⃣ Auto-Rig Model", variant="primary", interactive=False) with gr.Accordion("Generation Settings", open=False): resolution = gr.Radio(["512", "1024", "1536"], label="Resolution", value="1024") seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1) randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) decimation_target = gr.Slider( 100000, 500000, label="Decimation Target (faces)", value=200000, step=10000 ) texture_size = gr.Slider(1024, 4096, label="Texture Size", value=2048, step=1024) with gr.Accordion("Rigging Settings", open=False): num_beams = gr.Slider(1, 20, value=10, step=1, label="Beam search width") top_k = gr.Slider(1, 200, value=5, step=1, label="top_k") top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="top_p") temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="temperature") repetition_penalty = gr.Slider( 0.5, 3.0, value=2.0, step=0.1, label="repetition_penalty" ) with gr.Column(scale=3): with gr.Row(): model_output = gr.Model3D( label="Generated 3D Model", height=420, display_mode="solid", clear_color=(0.25, 0.25, 0.25, 1.0), ) rigged_output = gr.Model3D( label="Rigged 3D Model", height=420, display_mode="solid", clear_color=(0.25, 0.25, 0.25, 1.0), ) with gr.Row(): download_glb = gr.DownloadButton(label="⬇️ Unrigged GLB", interactive=False) download_rigged_glb = gr.DownloadButton( label="⬇️ Rigged GLB (Blender)", interactive=False ) download_rigged_fbx = gr.DownloadButton( label="⬇️ Rigged FBX (Blender)", interactive=False ) gr.Markdown(BLENDER_HELP) with gr.Column(scale=1, min_width=160): examples = gr.Examples( examples=[ f"assets/example_image/{image}" for image in sorted(os.listdir("assets/example_image")) ], inputs=[image_prompt], fn=preprocess_image, outputs=[image_prompt], run_on_click=True, examples_per_page=8, ) glb_state = gr.State() demo.load(start_session) demo.unload(end_session) def _reset_rig_progress(): return _render_steps(_init_steps(RIG_STEP_NAMES)) generate_btn.click( get_seed, inputs=[randomize_seed, seed], outputs=[seed], show_progress="hidden", ).then( preprocess_with_progress, inputs=[image_prompt], outputs=[image_prompt, task_progress, status], show_progress="minimal", ).then( generate_3d, inputs=[image_prompt, seed, resolution, decimation_target, texture_size], outputs=[glb_state, model_output, download_glb, task_progress, status], show_progress="minimal", ).then( lambda: gr.update(interactive=True), outputs=[rig_btn], show_progress="hidden", ) rig_btn.click( _reset_rig_progress, outputs=[task_progress], show_progress="hidden", ).then( rig_3d, inputs=[glb_state, num_beams, top_k, top_p, temperature, repetition_penalty], outputs=[rigged_output, download_rigged_glb, download_rigged_fbx, task_progress, status], show_progress="minimal", ) # Required for loading spinners + ZeroGPU job queue on Spaces. demo.queue(default_concurrency_limit=1) # --------------------------------------------------------------------------- # Module-scope model load (ZeroGPU packs CUDA weights at boot, streams them # into VRAM on first @spaces.GPU call). # --------------------------------------------------------------------------- os.makedirs(TMP_DIR, exist_ok=True) pipeline = Trellis2ImageTo3DPipeline.from_pretrained("microsoft/TRELLIS.2-4B") pipeline.rembg_model = None pipeline.low_vram = False pipeline.cuda() # NOTE: the TokenRig model is loaded lazily inside the GPU worker (see # load_rig_model). Loading it at module scope would trigger CUDA calls during # checkpoint load that fail in the main ZeroGPU process. if __name__ == "__main__": demo.launch(show_error=True, ssr_mode=False)