app.py

"""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'<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=_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)