app.py

from __future__ import annotations

import json
import random
import uuid
import shutil
import tempfile
import time
import zipfile
from pathlib import Path

import gradio as gr
from huggingface_hub import HfApi, fetch_job_logs, hf_hub_download, inspect_job, run_job, snapshot_download, whoami

APP_NAME = "Companion Forge"
VERSION = "7.0-lab-v6.4-prod"
APP_ROOT = Path(__file__).resolve().parent
TMP_ROOT = Path(tempfile.gettempdir()) / "companion-forge-v7"
TMP_ROOT.mkdir(parents=True, exist_ok=True)

L4_FLAVOR = "l4x1"
FLUX2_IMAGE = "hf.co/spaces/black-forest-labs/FLUX.2-klein-4B"
ANIGEN_IMAGE = "hf.co/spaces/VAST-AI/AniGen"
V7_RUNTIME_REPO = "patdev/Companion-Forge-L4-ONNX"

PRESETS = {
    "Desktop Companion": "friendly virtual desktop companion, expressive face, appealing mascot proportions, clean readable silhouette, full body, neutral A-pose, centered, isolated subject, no environment, no text",
    "Chibi Mascot": "cute chibi desktop mascot, large expressive head, compact proportions, full body, neutral A-pose, centered, isolated subject, clean silhouette, no environment, no text",
    "Robot": "small friendly desktop robot companion, articulated limbs, clean hard-surface design, full body, neutral A-pose, centered, isolated subject, no environment, no text",
    "Creature": "stylized fantasy creature companion, friendly and expressive, full body, standing neutral pose, centered, isolated subject, clean silhouette, no environment, no text",
    "Stylized Human": "stylized humanoid virtual companion, full body, neutral A-pose, centered, isolated subject, clean silhouette, practical clothing, no environment, no text",
}
BACKENDS = ["AniGen ONNX/TensorRT L4 — Full Engine Runtime (recommended)", "AniGen Native PyTorch — L4 Job"]
QUALITY = ["Fast", "Balanced", "Quality"]
BEHAVIOR_PROFILES = ["Code Pet", "Calm Assistant", "Energetic Mascot"]


def _require_token(oauth_token: gr.OAuthToken | None) -> str:
    if not oauth_token:
        raise gr.Error("Sign in with Hugging Face first. Companion Forge uses the OAuth `jobs` scope to launch an on-demand L4 under your account.")
    return oauth_token.token


def _expanded_prompt(prompt: str, preset: str) -> str:
    base = (prompt or "").strip()
    if not base:
        raise gr.Error("Enter a prompt first.")
    return (
        f"{base}. {PRESETS.get(preset, PRESETS['Desktop Companion'])}. high quality stylized 3D asset reference, coherent materials, "
        "entire character visible from head to feet, arms clearly separated from torso, legs clearly separated, symmetric neutral A-pose, "
        "simple light studio background, no floating objects, no props hiding limbs, animation-friendly character design"
    )


def _quality_slug(value: str) -> str:
    return str(value or "Balanced").strip().lower()


def _tail_job_logs(job_id: str, token: str, limit: int = 30) -> str:
    try:
        lines = [str(line) for line in fetch_job_logs(job_id=job_id, token=token)]
        return "\n".join(lines[-limit:])
    except Exception as exc:
        return f"Could not fetch Job logs: {exc}"


def _artifact_repo(token: str) -> tuple[HfApi, str]:
    api = HfApi(token=token)
    username = whoami(token=token)["name"]
    repo_id = f"{username}/companion-forge-jobs-artifacts"
    api.create_repo(repo_id=repo_id, repo_type="dataset", private=True, exist_ok=True)
    return api, repo_id


def _run_l4_worker(*, work: Path, worker: str, image: str, model_repo: str, model_mount: str, token: str, name: str, timeout: str, progress: gr.Progress):
    worker_src = APP_ROOT / "jobs" / worker
    runner_src = APP_ROOT / "jobs" / "repo_runner.py"
    if not worker_src.exists() or not runner_src.exists():
        raise gr.Error("Missing HF Jobs worker sources in the Space image")
    shutil.copy2(worker_src, work / "worker.py")
    shutil.copy2(runner_src, work / "runner.py")

    api, repo_id = _artifact_repo(token)
    run_id = f"{int(time.time())}-{random.randint(100000,999999)}"
    prefix = f"runs/{run_id}"
    progress(0.05, desc="Uploading private Job artifacts")
    api.upload_folder(
        repo_id=repo_id,
        repo_type="dataset",
        folder_path=work,
        path_in_repo=prefix,
        commit_message=f"Companion Forge request {run_id}",
    )

    progress(0.10, desc="Scheduling NVIDIA L4 Job")
    bootstrap = (
        "from huggingface_hub import hf_hub_download; import os,runpy; "
        "p=hf_hub_download(os.environ['CF_REPO_ID'],f\"runs/{os.environ['CF_RUN_ID']}/runner.py\","
        "repo_type='dataset',token=os.environ['HF_TOKEN']); runpy.run_path(p,run_name='__main__')"
    )
    job = run_job(
        image=image,
        command=["python", "-c", bootstrap],
        flavor=L4_FLAVOR,
        timeout=timeout,
        name=name,
        labels={"app": "companion-forge", "worker": worker.replace(".", "-").replace("/", "-"), "version": VERSION.replace(".", "-")},
        env={"CF_REPO_ID": repo_id, "CF_RUN_ID": run_id, "HF_XET_HIGH_PERFORMANCE": "1", "CF_RUNTIME_FIX": "flux2-trt" if worker == "flux2_trt_l4.py" else ("anigen-hybrid" if (worker == "anigen_hybrid_l4.py" or worker.startswith("v7_")) else ("anigen" if worker == "anigen_l4.py" else ("trellis" if worker == "trellis_l4.py" else "")))},
        secrets={"HF_TOKEN": token},
        token=token,
    )

    started = time.time()
    last_stage = None
    while True:
        info = inspect_job(job_id=job.id, token=token)
        stage = str(info.status.stage)
        if stage != last_stage:
            print(f"[Companion Forge] Job {job.id}: {stage}", flush=True)
            last_stage = stage
        if stage in {"COMPLETED", "CANCELED", "ERROR", "DELETED"}:
            break
        elapsed = time.time() - started
        if stage == "SCHEDULING":
            progress(min(0.25, 0.12 + elapsed / 600.0), desc=f"L4 scheduling · {int(elapsed)}s")
        else:
            progress(min(0.82, 0.30 + elapsed / 900.0), desc=f"L4 running · {int(elapsed)}s")
        time.sleep(2.5)

    if stage != "COMPLETED":
        raise gr.Error(f"HF Job ended with {stage}.\n\n{_tail_job_logs(job.id, token)}")

    progress(0.88, desc="Downloading private Job outputs")
    downloaded = Path(tempfile.mkdtemp(prefix="job-result-", dir=TMP_ROOT))
    snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        token=token,
        allow_patterns=[f"{prefix}/**"],
        local_dir=downloaded,
    )
    remote_run = downloaded / prefix
    for item in remote_run.iterdir():
        if item.is_file():
            shutil.copy2(item, work / item.name)

    result_file = work / "result.json"
    if not result_file.exists():
        raise gr.Error(f"Job completed but result.json is missing.\n\n{_tail_job_logs(job.id, token)}")
    result = json.loads(result_file.read_text(encoding="utf-8"))
    result.update({"job_id": job.id, "job_url": job.url, "hardware": L4_FLAVOR, "artifact_repo": repo_id, "run_id": run_id})
    try:
        api.delete_folder(
            path_in_repo=prefix, repo_id=repo_id, repo_type="dataset",
            commit_message=f"Clean Companion Forge run {run_id}",
        )
        result["artifact_cleanup"] = True
    except Exception as exc:
        print(f"[Companion Forge] artifact cleanup skipped: {exc}", flush=True)
        result["artifact_cleanup"] = False
    return result


def generate_reference(prompt, preset, seed, randomize_seed, progress=gr.Progress(), oauth_token: gr.OAuthToken | None = None):
    token = _require_token(oauth_token)
    full_prompt = _expanded_prompt(prompt, preset)
    seed = random.randint(0, 2_147_483_647) if randomize_seed else int(seed)
    work = Path(tempfile.mkdtemp(prefix="reference-", dir=TMP_ROOT))
    (work / "request.json").write_text(json.dumps({"prompt": full_prompt, "seed": seed, "width": 512, "height": 512, "steps": 4}, indent=2), encoding="utf-8")
    result = _run_l4_worker(
        work=work, worker="flux2_trt_l4.py", image=FLUX2_IMAGE,
        model_repo="black-forest-labs/FLUX.2-klein-4B", model_mount="/models/flux2",
        token=token, name="companion-forge-flux2-trt512-reference", timeout="25m", progress=progress,
    )
    image_path = work / result.get("file", "reference.png")
    if not image_path.exists():
        raise gr.Error("FLUX.2 Klein Job did not return reference.png")
    progress(1.0, desc="Reference ready")
    job_md = f"**FLUX.2 Reference Job:** [{result['job_id']}]({result['job_url']}) · {result.get('seconds', '?')}s · runtime `{result.get('runtime','?')}` · device peak {result.get('peak_device_vram_gib', result.get('peak_vram_gib','?'))} GiB"
    return str(image_path), seed, full_prompt, job_md


def _animation_pack(profile: str, rigged: bool) -> dict:
    energetic, calm = profile == "Energetic Mascot", profile == "Calm Assistant"
    speed = 0.72 if calm else (1.25 if energetic else 1.0)
    driver = "skeletal-procedural" if rigged else "model-transform"
    clips = {
        "idle": {"loop": True, "duration": 3.4 if calm else (1.8 if energetic else 2.6), "driver": driver},
        "walk": {"loop": True, "duration": 0.72 / speed, "driver": driver},
        "wave": {"loop": False, "duration": 1.35 / speed, "driver": driver},
        "happy": {"loop": False, "duration": 1.0 / speed, "driver": driver},
        "thinking": {"loop": True, "duration": 2.2 / speed, "driver": driver},
        "typing": {"loop": True, "duration": 0.55 / speed, "driver": driver},
        "sleep": {"loop": True, "duration": 3.5, "driver": driver},
        "error": {"loop": False, "duration": 0.65, "driver": driver},
        "celebrate": {"loop": False, "duration": 1.45 / speed, "driver": driver},
    }
    states = {
        "idle": {"clip": "idle", "loop": True, "blend_ms": 180},
        "walking": {"clip": "walk", "loop": True, "blend_ms": 120},
        "waving": {"clip": "wave", "return": "idle", "blend_ms": 120},
        "happy": {"clip": "happy", "return": "idle", "blend_ms": 100},
        "thinking": {"clip": "thinking", "loop": True, "blend_ms": 180},
        "typing": {"clip": "typing", "loop": True, "blend_ms": 100},
        "sleeping": {"clip": "sleep", "loop": True, "blend_ms": 400},
        "error": {"clip": "error", "return": "idle", "blend_ms": 60},
        "celebrating": {"clip": "celebrate", "return": "idle", "blend_ms": 90},
    }
    events = {
        "agent.idle": "idle", "agent.move": "walking", "agent.hello": "waving",
        "agent.message.received": "happy", "agent.thinking": "thinking",
        "agent.tool.start": "typing", "agent.tool.end": "happy",
        "agent.task.complete": "celebrating", "agent.error": "error",
        "agent.sleep": "sleeping", "agent.wake": "idle",
    }
    return {
        "format": "companion-motion/v2", "profile": profile, "rigged": rigged,
        "default_state": "idle", "blend_mode": "crossfade",
        "semantic_bones": {
            "strategy": "geometry-inference", "source_pattern": "joint_*" if rigged else None,
            "targets": ["root", "body", "head", "arm_l", "forearm_l", "hand_l", "arm_r", "forearm_r", "hand_r", "leg_l", "shin_l", "foot_l", "leg_r", "shin_r", "foot_r", "antenna"] if rigged else [],
        },
        "clips": clips, "states": states, "events": events,
    }


def _build_bundle(work: Path, *, image_path: str, worker_result: dict, backend: str, behavior_profile: str, prompt: str = "", expanded_prompt: str = "", preset: str = ""):
    output_dir = Path(tempfile.mkdtemp(prefix="output-", dir=TMP_ROOT))
    source_model = work / worker_result.get("file", "companion.glb")
    if not source_model.exists():
        raise gr.Error("The L4 Job completed but companion.glb is missing")
    reference_out, model_out = output_dir / "reference.png", output_dir / "companion.glb"
    skeleton_out, manifest_out = output_dir / "skeleton.glb", output_dir / "companion.json"
    animations_out, bundle_out = output_dir / "animations.json", output_dir / "companion-bundle.zip"
    shutil.copy2(image_path, reference_out)
    shutil.copy2(source_model, model_out)
    rigged = worker_result.get("kind") == "rigged"
    skeleton_file = None
    src_skeleton = worker_result.get("skeleton_file")
    if src_skeleton and (work / src_skeleton).exists():
        shutil.copy2(work / src_skeleton, skeleton_out)
        skeleton_file = str(skeleton_out)
    animations = _animation_pack(behavior_profile, rigged)
    animations_out.write_text(json.dumps(animations, indent=2), encoding="utf-8")
    manifest = {
        "format": "companion-forge/v5.1",
        "execution": {
            "platform": "Hugging Face Jobs", "hardware": L4_FLAVOR,
            "job_id": worker_result.get("job_id"), "job_url": worker_result.get("job_url"),
            "seconds": worker_result.get("seconds"),
            "runtime": worker_result.get("runtime"),
            "peak_vram_gib": worker_result.get("peak_vram_gib"),
            "peak_device_vram_gib": worker_result.get("peak_device_vram_gib"),
            "baseline_device_vram_gib": worker_result.get("baseline_device_vram_gib"),
            "gpu": worker_result.get("gpu"), "stages": worker_result.get("stages"),
        },
        "prompt": prompt or None, "expanded_prompt": expanded_prompt or None, "preset": preset or None,
        "backend": backend,
        "mesh": {
            "file": "companion.glb", "vertices": worker_result.get("vertices"), "faces": worker_result.get("faces"),
            "resolution": worker_result.get("resolution"), "texture_size": worker_result.get("texture_size"),
            "texture_mode": worker_result.get("texture_mode", "PBR"),
        },
        "rig": {"rigged": rigged, "skeleton_file": "skeleton.glb" if skeleton_file else None, "joints": worker_result.get("joints"), "skinning": rigged},
        "behavior": {"profile": behavior_profile, "animation_file": "animations.json", "states": list(animations["states"])},
    }
    manifest_out.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    with zipfile.ZipFile(bundle_out, "w", zipfile.ZIP_DEFLATED) as archive:
        archive.write(model_out, "companion.glb")
        archive.write(reference_out, "reference.png")
        archive.write(manifest_out, "companion.json")
        archive.write(animations_out, "animations.json")
        if skeleton_file:
            archive.write(skeleton_out, "skeleton.glb")
    return str(reference_out), str(model_out), skeleton_file, str(bundle_out), manifest, animations


def generate_3d_from_image(image_path, backend, behavior_profile, quality, seed, randomize_seed, texture_size, prompt="", expanded_prompt="", preset="", progress=gr.Progress(), oauth_token: gr.OAuthToken | None = None):
    if not image_path:
        raise gr.Error("Provide an image first.")
    token = _require_token(oauth_token)
    seed = random.randint(0, 2_147_483_647) if randomize_seed else int(seed)
    work = Path(tempfile.mkdtemp(prefix="forge-", dir=TMP_ROOT))
    suffix = Path(image_path).suffix.lower() or ".png"
    local_input = work / f"input{suffix}"
    shutil.copy2(image_path, local_input)
    (work / "request.json").write_text(json.dumps({"input_file": local_input.name, "seed": seed, "quality": _quality_slug(quality), "texture_size": int(texture_size)}, indent=2), encoding="utf-8")
    if backend.startswith("AniGen ONNX/TensorRT") or backend.startswith("AniGen Hybrid"):
        result = _run_l4_worker(work=work, worker="anigen_hybrid_l4.py", image=ANIGEN_IMAGE, model_repo="VAST-AI/AniGen", model_mount="/models/anigen", token=token, name="companion-forge-anigen-hybrid", timeout="35m", progress=progress)
    elif backend.startswith("AniGen Native"):
        result = _run_l4_worker(work=work, worker="anigen_l4.py", image=ANIGEN_IMAGE, model_repo="VAST-AI/AniGen", model_mount="/models/anigen", token=token, name="companion-forge-anigen-native", timeout="35m", progress=progress)
    else:
        raise gr.Error("Unknown L4 backend")
    progress(0.93, desc="Building companion bundle")
    outputs = _build_bundle(work, image_path=image_path, worker_result=result, backend=backend, behavior_profile=behavior_profile, prompt=prompt, expanded_prompt=expanded_prompt, preset=preset)
    progress(1.0, desc="Companion ready")
    job_md = f"**3D Job:** [{result['job_id']}]({result['job_url']}) · {result.get('seconds', '?')}s · runtime `{result.get('runtime','?')}` · device peak {result.get('peak_device_vram_gib', result.get('peak_vram_gib','?'))} GiB"
    return (*outputs, seed, job_md)


def text_to_3d_generate(reference_path, prompt, expanded_prompt, preset, backend, behavior_profile, quality, seed, randomize_seed, texture_size, progress=gr.Progress(), oauth_token: gr.OAuthToken | None = None):
    if not reference_path:
        raise gr.Error("Generate a reference first.")
    return generate_3d_from_image(reference_path, backend, behavior_profile, quality, seed, randomize_seed, texture_size, prompt=prompt, expanded_prompt=expanded_prompt, preset=preset, progress=progress, oauth_token=oauth_token)


def _v7_user_repos(token: str) -> tuple[str, str, str]:
    api = HfApi(token=token)
    username = whoami(token=token)["name"]
    cache_repo = f"{username}/companion-forge-v7-teacher-cache"
    students_repo = f"{username}/companion-forge-v7-students"
    api.create_repo(repo_id=cache_repo, repo_type="dataset", private=True, exist_ok=True)
    api.create_repo(repo_id=students_repo, repo_type="model", private=True, exist_ok=True)
    return username, cache_repo, students_repo


def v7_refresh_status(oauth_token: gr.OAuthToken | None = None):
    token = oauth_token.token if oauth_token else None
    try:
        local = hf_hub_download(V7_RUNTIME_REPO, "v7/bench/validation_status.json", repo_type="model", token=token, force_download=True)
        status = json.loads(Path(local).read_text(encoding="utf-8"))
        if token:
            username, cache_repo, students_repo = _v7_user_repos(token)
            status["user"] = {"name": username, "cache_repo": cache_repo, "students_repo": students_repo}
        else:
            status["user"] = {"signed_in": False}
        return status
    except Exception as exc:
        return {"error": str(exc)}


def v7_add_to_cache(image, seed, progress=gr.Progress(), oauth_token: gr.OAuthToken | None = None):
    token = _require_token(oauth_token)
    if not image:
        raise gr.Error("Choose a reference image first.")
    _, cache_repo, _ = _v7_user_repos(token)
    work = Path(tempfile.mkdtemp(prefix="v7-cache-", dir=TMP_ROOT))
    src = work / "reference.png"
    shutil.copy2(Path(image), src)
    key = f"ref-{int(seed)}-{uuid.uuid4().hex[:10]}"
    request = {"input_file": src.name, "seed": int(seed), "ss_steps": 10, "cache_repo": cache_repo, "key": key}
    (work / "request.json").write_text(json.dumps(request, indent=2), encoding="utf-8")
    result = _run_l4_worker(
        work=work, worker="v7_cache_l4.py", image=ANIGEN_IMAGE,
        model_repo="", model_mount="", token=token,
        name="companion-forge-v7-cache", timeout="35m", progress=progress,
    )
    progress(1.0, desc="v7 cache ready")
    return (
        f"**v7 cache:** `{cache_repo}:{result.get('path')}` · "
        f"{result.get('bytes',0)/1e6:.1f} MB · geo coords {result.get('coords')} · skeleton coords {result.get('coords_skl')}",
        result,
    )



def v7_add_multiview_cache(front, left, back, right, seed, progress=gr.Progress(), oauth_token: gr.OAuthToken | None = None):
    token = _require_token(oauth_token)
    views = {"front": front, "left": left, "back": back, "right": right}
    missing = [name for name, value in views.items() if not value]
    if missing:
        raise gr.Error("Missing multi-view images: " + ", ".join(missing))
    _, cache_repo, _ = _v7_user_repos(token)
    work = Path(tempfile.mkdtemp(prefix="v7-multiview-", dir=TMP_ROOT))
    request = {"seed": int(seed), "ss_steps": 10, "cache_repo": cache_repo, "key": f"multiview-{int(seed)}-{uuid.uuid4().hex[:10]}"}
    for name, value in views.items():
        dst = work / f"{name}.png"
        shutil.copy2(Path(value), dst)
        request[f"{name}_file"] = dst.name
    (work / "request.json").write_text(json.dumps(request, indent=2), encoding="utf-8")
    result = _run_l4_worker(
        work=work, worker="v7_multiview_cache_l4.py", image=ANIGEN_IMAGE,
        model_repo="", model_mount="", token=token,
        name="companion-forge-v7-multiview", timeout="45m", progress=progress,
    )
    progress(1.0, desc="v7 multi-view cache ready")
    return (
        f"**v7 multi-view cache:** `{cache_repo}:{result.get('path')}` · fusion `{result.get('fusion')}` · "
        f"geo coords {result.get('coords')} · skeleton coords {result.get('coords_skl')}",
        result,
    )

def v7_train_stage(component, stage_label, run_length, progress=gr.Progress(), oauth_token: gr.OAuthToken | None = None):
    token = _require_token(oauth_token)
    _, cache_repo, students_repo = _v7_user_repos(token)
    stage_map = {"10 → 4": 0, "4 → 2": 1, "2 → 1": 2}
    step_map = {"Smoke · 1 step": 1, "Short · 100 steps": 100, "Configured stage": 0}
    stage = stage_map[stage_label]
    steps = step_map[run_length]
    chain = {
        0: [],
        1: [f"{component}/10to4/adapter.pt"],
        2: [f"{component}/10to4/adapter.pt", f"{component}/4to2/adapter.pt"],
    }[stage]
    api = HfApi(token=token)
    missing = [path for path in chain if not api.file_exists(repo_id=students_repo, filename=path, repo_type="model")]
    if missing:
        raise gr.Error("Previous v7 stage is missing: " + ", ".join(missing))
    request = {
        "component": component, "stage": stage, "steps": steps,
        "cache_repo": cache_repo, "cache_prefix": "cache",
        "students_repo": students_repo, "resume_paths": chain,
    }
    work = Path(tempfile.mkdtemp(prefix="v7-train-", dir=TMP_ROOT))
    (work / "request.json").write_text(json.dumps(request, indent=2), encoding="utf-8")
    result = _run_l4_worker(
        work=work, worker="v7_train_l4.py", image=ANIGEN_IMAGE,
        model_repo="", model_mount="", token=token,
        name=f"companion-forge-v7-{component}-{stage}", timeout="90m", progress=progress,
    )
    metrics = result.get("job_metrics", {})
    progress(1.0, desc="v7 training stage complete")
    return (
        f"**v7 {component} {stage_label}:** `{students_repo}:{result.get('adapter_path')}` · "
        f"{metrics.get('elapsed_s','?')} s · peak {metrics.get('peak_vram_gib','?')} GiB",
        result,
    )


CSS = """
.gradio-container { max-width: 1500px !important; }
#hero { text-align:center; margin: 4px 0 14px; }
#hero h1 { font-size: 2.25rem; margin-bottom: .15rem; }
#hero p { opacity: .72; margin-top: 0; }
"""

with gr.Blocks(title=APP_NAME, delete_cache=(3600, 3600)) as demo:
    gr.HTML("<div id='hero'><h1>🤖 Companion Forge v7 Lab</h1><p>Validated v6.4 production runtime + v7 few-step distillation, FP8, 2:4 sparsity, multi-view, symmetry and MoE research on on-demand NVIDIA L4.</p></div>")
    with gr.Row():
        gr.LoginButton(value="Sign in with Hugging Face", logout_value="Logout ({})", variant="huggingface")
        gr.Markdown("**Compute:** `l4x1` · 8 vCPU · 30 GB RAM · 1× NVIDIA L4 (24 GB class). Jobs stop automatically after generation.")

    with gr.Tabs():
        with gr.Tab("Text → Code Pet"):
            with gr.Row():
                with gr.Column(scale=5, min_width=390):
                    prompt = gr.Textbox(label="Prompt", placeholder="A cute white chibi desktop robot, blue screen face, small antenna, articulated arms and legs...", lines=4)
                    preset = gr.Dropdown(list(PRESETS), value="Robot", label="Character preset")
                    with gr.Row():
                        seed = gr.Number(value=42, precision=0, label="Seed")
                        randomize = gr.Checkbox(value=True, label="Randomize")
                    gen_ref = gr.Button("1. Generate Reference on L4", variant="primary", size="lg")
                    reference = gr.Image(label="Reference — edit/replace before 3D", type="filepath", height=380)
                    expanded = gr.Textbox(visible=False)
                    ref_job = gr.Markdown()
                    backend = gr.Dropdown(BACKENDS, value=BACKENDS[0], label="3D backend")
                    behavior = gr.Dropdown(BEHAVIOR_PROFILES, value="Code Pet", label="Behavior profile")
                    quality = gr.Radio(QUALITY, value="Balanced", label="L4 optimization profile")
                    texture = gr.State(0)
                    forge = gr.Button("2. Forge Companion on L4", variant="primary", size="lg")
                    forge_job = gr.Markdown()
                with gr.Column(scale=7):
                    with gr.Tabs():
                        with gr.Tab("Companion"): model = gr.Model3D(label="companion.glb", height=590, display_mode="solid")
                        with gr.Tab("Skeleton"): skeleton = gr.Model3D(label="skeleton.glb", height=590, display_mode="solid")
                        with gr.Tab("Behavior"): animation = gr.JSON(label="State machine")
                        with gr.Tab("Manifest"): manifest = gr.JSON(label="Generation metadata / benchmark")
                    bundle = gr.File(label="Download Code Pet bundle")
                    final_ref = gr.Image(visible=False)
            gen_ref.click(generate_reference, inputs=[prompt, preset, seed, randomize], outputs=[reference, seed, expanded, ref_job], concurrency_limit=1)
            forge.click(text_to_3d_generate, inputs=[reference, prompt, expanded, preset, backend, behavior, quality, seed, randomize, texture], outputs=[final_ref, model, skeleton, bundle, manifest, animation, seed, forge_job], concurrency_limit=1)

        with gr.Tab("Image → Code Pet"):
            with gr.Row():
                with gr.Column(scale=5, min_width=390):
                    image = gr.Image(label="Input character", type="filepath", image_mode="RGBA", height=470)
                    img_backend = gr.Dropdown(BACKENDS, value=BACKENDS[0], label="3D backend")
                    img_behavior = gr.Dropdown(BEHAVIOR_PROFILES, value="Code Pet", label="Behavior profile")
                    img_quality = gr.Radio(QUALITY, value="Balanced", label="L4 optimization profile")
                    with gr.Row():
                        img_seed = gr.Number(value=42, precision=0, label="Seed")
                        img_random = gr.Checkbox(value=True, label="Randomize")
                    img_texture = gr.State(0)
                    img_forge = gr.Button("Forge on L4", variant="primary", size="lg")
                    img_job = gr.Markdown()
                    gr.Markdown("Best rigging: one full-body character, clean background, separated limbs, neutral A-pose.")
                with gr.Column(scale=7):
                    with gr.Tabs():
                        with gr.Tab("Companion"): img_model = gr.Model3D(label="companion.glb", height=590, display_mode="solid")
                        with gr.Tab("Skeleton"): img_skeleton = gr.Model3D(label="skeleton.glb", height=590, display_mode="solid")
                        with gr.Tab("Behavior"): img_animation = gr.JSON(label="State machine")
                        with gr.Tab("Manifest"): img_manifest = gr.JSON(label="Generation metadata / benchmark")
                    img_bundle = gr.File(label="Download Code Pet bundle")
                    img_final_ref = gr.Image(visible=False)
            img_forge.click(generate_3d_from_image, inputs=[image, img_backend, img_behavior, img_quality, img_seed, img_random, img_texture], outputs=[img_final_ref, img_model, img_skeleton, img_bundle, img_manifest, img_animation, img_seed, img_job], concurrency_limit=1)

        with gr.Tab("V7 Lab"):
            gr.Markdown(
                "### 🧪 Companion Forge v7 — Distillation & Optimization Lab\n"
                "Generation stays on the validated v6.4 TensorRT runtime until a v7 student passes quality, finite-tensor and rig gates. "
                "References enter the training cache **only when you explicitly click the cache button**; the cache and student repos are private to your HF account."
            )
            v7_refresh_btn = gr.Button("Refresh v7 status")
            v7_status = gr.JSON(value={"status": "Click Refresh v7 status"}, label="Technology validation status")
            gr.Markdown("#### 1. Add a training reference")
            with gr.Row():
                v7_cache_image = gr.Image(type="filepath", label="Reference image", sources=["upload", "clipboard"])
                with gr.Column():
                    v7_cache_seed = gr.Number(value=42, precision=0, label="Cache seed")
                    v7_cache_btn = gr.Button("Add to private v7 cache", variant="primary")
                    v7_cache_msg = gr.Markdown()
            v7_cache_result = gr.JSON(label="Cache Job result")
            with gr.Accordion("Optional 4-view cache · front / left / back / right", open=False):
                with gr.Row():
                    v7_mv_front = gr.Image(type="filepath", label="Front", sources=["upload", "clipboard"])
                    v7_mv_left = gr.Image(type="filepath", label="Left", sources=["upload", "clipboard"])
                    v7_mv_back = gr.Image(type="filepath", label="Back", sources=["upload", "clipboard"])
                    v7_mv_right = gr.Image(type="filepath", label="Right", sources=["upload", "clipboard"])
                with gr.Row():
                    v7_mv_seed = gr.Number(value=42, precision=0, label="Multi-view seed")
                    v7_mv_btn = gr.Button("Add 4-view set to private v7 cache")
                v7_mv_msg = gr.Markdown()
                v7_mv_result = gr.JSON(label="Multi-view cache Job result")
                v7_mv_btn.click(v7_add_multiview_cache, [v7_mv_front, v7_mv_left, v7_mv_back, v7_mv_right, v7_mv_seed], [v7_mv_msg, v7_mv_result], concurrency_limit=1)
            gr.Markdown("#### 2. Train a few-step flow student")
            with gr.Row():
                v7_component = gr.Dropdown(["ss_flow", "slat_flow"], value="ss_flow", label="Component")
                v7_stage = gr.Dropdown(["10 → 4", "4 → 2", "2 → 1"], value="10 → 4", label="Progressive stage")
                v7_length = gr.Radio(["Smoke · 1 step", "Short · 100 steps", "Configured stage"], value="Smoke · 1 step", label="Run length")
            v7_train_btn = gr.Button("Launch v7 L4 training", variant="primary")
            v7_train_msg = gr.Markdown()
            v7_train_result = gr.JSON(label="Training Job result")
            v7_refresh_btn.click(v7_refresh_status, inputs=[], outputs=[v7_status])
            v7_cache_btn.click(v7_add_to_cache, [v7_cache_image, v7_cache_seed], [v7_cache_msg, v7_cache_result], concurrency_limit=1)
            v7_train_btn.click(v7_train_stage, [v7_component, v7_stage, v7_length], [v7_train_msg, v7_train_result], concurrency_limit=1)

        with gr.Tab("L4 Architecture"):
            gr.Markdown("""
### Production v6.4 + v7 training architecture
`CPU Space` → private request repo → `HF Job l4x1` → **staged ONNX/TensorRT engines** → minimal sparse topology shell → private output sync → GPU terminates.

**Production runtime order**
- FLUX.2 Klein transformer: **ONNX opset 23 + TensorRT FP16 / SM89**, static 512 profile; BF16 PyTorch fallback.
- DINOv2: **ONNX + TensorRT** → ORT CUDA → PyTorch.
- DSINE: **ONNX + TensorRT** → staged native/ORT fallback.
- AniGen SS Flow: **full ONNX opset 23 + TensorRT**.
- AniGen SS decoder: **full ONNX opset 23 + TensorRT**.
- AniGen SLat Flow: **24-block transformer core ONNX/TensorRT** + only a ~188 MB native sparse IO shell; the old ~2.46 GB PyTorch flow checkpoint is no longer downloaded normally.
- AniGen skin decoder: **dynamic ONNX/TensorRT**, profile up to 300k vertices / 64 joints.
- SLat topology/decoder is now backed by the validated custom ONNX/TensorRT plugin stack (`SparseConv3D`, `SparseWindowAttention`, `SparseDownsample`, `SparseUpsample`, `SparseSubdivide`, `SparseMeshTopologyExtract`); native decoder is lazy fallback only.

**Fast**: 10/10 SS + SLat steps. **Balanced**: 14/14. **Quality**: 20/20.

Validated Fast L4 runtime: `dino-trt + dsine-trt + slat-shell-trt + skin-trt + ss-trt`, GLB + skeleton exported with no PyTorch flow-model fallback.

The rigged path exports **vertex-color GLB** to avoid nvdiffrast CUDA JIT. TensorRT plans are precompiled for NVIDIA L4 / SM89; dense conditionners are freed before 3D generation to limit VRAM overlap.
            """)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch(css=CSS, ssr_mode=False, show_error=True, mcp_server=True)