Spaces:
Running
Running
| 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) | |