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("
Validated v6.4 production runtime + v7 few-step distillation, FP8, 2:4 sparsity, multi-view, symmetry and MoE research on on-demand NVIDIA L4.