Upload space/app.py with huggingface_hub

space/app.py

@@ -130,9 +130,10 @@ import signal
 import subprocess
 import sys
 import time
+import traceback
 from datetime import datetime
 from pathlib import Path
-from typing import Tuple
+from typing import Dict, List, Tuple
 
 import gradio as gr
 import numpy as np
@@ -240,10 +241,95 @@ 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'<div style="margin:10px 0">'
+            f'<div style="display:flex;justify-content:space-between;font-size:13px;margin-bottom:4px">'
+            f"<span>{icon} {s['name']}</span><span>{pct}%</span></div>"
+            f'<div style="background:#e5e7eb;border-radius:6px;height:10px;overflow:hidden">'
+            f'<div style="width:{pct}%;background:{color};height:10px;border-radius:6px;'
+            f'transition:width .4s"></div></div></div>'
+        )
+    return f'<div style="padding:4px 0">{"".join(rows)}</div>'
+
+
+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=180)
+@spaces.GPU(duration=_gpu_duration_generate)
 def generate_3d(
     image: Image.Image,
     seed: int,
@@ -252,71 +338,117 @@ def generate_3d(
     texture_size: int,
     req: gr.Request,
     progress=gr.Progress(track_tqdm=True),
-) -> str:
+):
+    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.")
 
-    # Always preprocess here so generation works even if the upload handler
-    # hasn't finished yet (and so the GPU worker has a clean RGBA crop).
-    image = preprocess_image(image)
-
-    outputs, latents = pipeline.run(
-        image,
-        seed=seed,
-        preprocess_image=False,
-        sparse_structure_sampler_params={
-            "steps": 12,
-            "guidance_strength": 7.5,
-            "guidance_rescale": 0.7,
-            "rescale_t": 5.0,
-        },
-        shape_slat_sampler_params={
-            "steps": 12,
-            "guidance_strength": 7.5,
-            "guidance_rescale": 0.5,
-            "rescale_t": 3.0,
-        },
-        tex_slat_sampler_params={
-            "steps": 12,
-            "guidance_strength": 1.0,
-            "guidance_rescale": 0.0,
-            "rescale_t": 3.0,
-        },
-        pipeline_type={
-            "512": "512",
-            "1024": "1024_cascade",
-            "1536": "1536_cascade",
-        }[resolution],
-        return_latent=True,
-    )
-    mesh = outputs[0]
-    mesh.simplify(16777216)  # nvdiffrast limit
-    res = latents[2]
-
-    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,
-    )
+    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.")
 
-    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)  # PNG textures: safest for Blender import
-    torch.cuda.empty_cache()
-    return glb_path
+        # 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
 
 
 # ---------------------------------------------------------------------------
@@ -400,7 +532,11 @@ def load_rig_model():
     print("[app] TokenRig model loaded.", flush=True)
 
 
-@spaces.GPU(duration=320)
+def _noop_rig_files():
+    return gr.update(), gr.update(), gr.update()
+
+
+@spaces.GPU(duration=480)
 def rig_3d(
     glb_path: str,
     num_beams: int,
@@ -410,81 +546,113 @@ def rig_3d(
     repetition_penalty: float,
     req: gr.Request,
     progress=gr.Progress(track_tqdm=True),
-) -> Tuple[str, str, str]:
+):
+    steps = _init_steps(RIG_STEP_NAMES)
     if not glb_path or not os.path.exists(glb_path):
-        raise gr.Error("Generate (or upload) a 3D model first.")
-
-    ensure_bpy_server_started()
-    load_rig_model()
+        raise gr.Error("Generate a 3D model first.")
 
-    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"
-
-    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")
-
-    for batch in dataloader:
-        batch = {
-            k: v.to(infer_device) if isinstance(v, Tensor) else v for k, v in batch.items()
+    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)]},
         }
-        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,
+        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,
         )
-
-        preds = rig_model.predict_step(batch, skeleton_tokens=None, make_asset=True)["results"]
-        asset = preds[0].asset
-        assert asset is not None
-
-        # Transfer the predicted rig back onto the original mesh so PBR
-        # textures and world scale survive, then export both formats.
-        for export_path in (out_glb, out_fbx):
-            payload = dict(
-                source_asset=asset,
-                target_path=asset.path,
-                export_path=export_path,
-                group_per_vertex=4,
-            )
-            res = bytes_to_object(
-                requests.post(f"{BPY_SERVER}/transfer", data=object_to_bytes(payload)).content
+        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,
             )
-            if res != "ok":
-                raise gr.Error(f"Rig export failed: {res}")
-            print(f"[app] Exported: {export_path}", flush=True)
-
-    torch.cuda.empty_cache()
-    return out_glb, out_glb, out_fbx
+            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
 
 
 # ---------------------------------------------------------------------------
@@ -528,6 +696,16 @@ with gr.Blocks(title="Image → 3D → Rigged Model", delete_cache=(600, 600)) a
                 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)
 
@@ -591,43 +769,44 @@ with gr.Blocks(title="Image → 3D → Rigged Model", delete_cache=(600, 600)) a
     demo.load(start_session)
     demo.unload(end_session)
 
-    image_prompt.upload(
-        preprocess_image,
-        inputs=[image_prompt],
-        outputs=[image_prompt],
-    )
+    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],
+        outputs=[glb_state, model_output, download_glb, task_progress, status],
+        show_progress="minimal",
     ).then(
-        lambda p: (
-            p,
-            gr.DownloadButton(value=p, interactive=True),
-            gr.Button(interactive=True),
-        ),
-        inputs=[glb_state],
-        outputs=[model_output, download_glb, rig_btn],
+        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],
-    ).then(
-        lambda g, f: (
-            gr.DownloadButton(value=g, interactive=True),
-            gr.DownloadButton(value=f, interactive=True),
-        ),
-        inputs=[download_rigged_glb, download_rigged_fbx],
-        outputs=[download_rigged_glb, download_rigged_fbx],
+        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
@@ -646,4 +825,4 @@ pipeline.cuda()
 
 
 if __name__ == "__main__":
-    demo.launch(show_error=True)
+    demo.launch(show_error=True, ssr_mode=False)