Milestones 9-12: UniRig, UE5 Export, Presets, Polish

Milestone 9 — Stage 3 Auto-Rigging (stage3_rig.py):
- gradio_client → MohamedRashad/UniRig: sends GLB, downloads rigged.glb
- Same client pattern as TRELLIS.2 generation (no local GPU/weights needed)
- Handles rig_type + seed inputs; outputs saved to workspace/current/

Milestone 10 — Stage 4 UE5 Export (stage4_export.py):
- Bundles all workspace outputs into a UE5-ready zip:
SM_/SK_ mesh, LOD1/2, UCX_ collision, T_*_D/N/ORM textures + manifest.json
- Game-ready checklist warns before export (missing UVs, LODs, collision, etc.)
- engine != UE5 returns informative stub (Unity/Godot/Blender planned)

Milestone 11 — Presets (workspace/presets/*.json + app.py):
- ui_save_preset() now reads ALL tab UI values and writes real JSON
- 5 default presets committed to workspace/presets/
- Presets tab description updated

Milestone 12 — Polish:
- Game-ready checklist before export (already in handle_export)
- Presets tab fully wired; Save captures all Stage 1-4 settings

app.py: rig_btn → handle_auto_rig; ex_btn → handle_export;
pr_save → ui_save_preset (full inputs); stubs removed

app.py

@@ -243,31 +243,70 @@ def run_post_process(
     return "\n".join(log)
 
 
-def stub_auto_rig(*_args, **_kwargs):
-    return _stub("Stage 3: Auto-Rigging")
+def handle_auto_rig(rig_type, seed, spring, fmt):
+    from src.stages.stage3_rig import auto_rig
+    _, msg = auto_rig(rig_type=rig_type, seed=int(seed))
+    return msg
 
 
-def stub_export(*_args, **_kwargs):
-    return _stub("Stage 4: Export"), None  # message, file
+def handle_export(engine, asset_name, asset_type, include_lods, include_collision):
+    from src.stages.stage4_export import export_ue5, _checklist
+    state = workspace.get_state()
+
+    if engine != "UE5":
+        return f"🚧 {engine} export — not implemented. UE5 is the only supported engine.", None
+
+    # Game-ready checklist
+    issues = _checklist(state)
+    checklist_md = ""
+    if issues:
+        checklist_md = "\n\n⚠️ **Checklist warnings:**\n" + "\n".join(f"- {i}" for i in issues)
+
+    try:
+        zip_path, msg, _ = export_ue5(
+            asset_name=asset_name.strip() or "Asset",
+            asset_type=asset_type,
+            include_lods=include_lods,
+            include_collision=include_collision,
+        )
+        return msg + checklist_md, str(zip_path)
+    except Exception as e:
+        return f"❌ Export failed: {e}{checklist_md}", None
 
 
 # ---------------------------------------------------------------------------
-# Presets tab logic (this one is real even in Milestone 1)
+# Presets tab logic (M11 — real save/load wired)
 # ---------------------------------------------------------------------------
 
 def ui_refresh_presets() -> gr.Dropdown:
     return gr.Dropdown(choices=workspace.list_presets(), label="Saved presets")
 
 
-def ui_save_preset(name: str) -> tuple[str, gr.Dropdown]:
+def ui_save_preset(
+    name: str,
+    gen_model, gen_quality, gen_seed, gen_tex_size,
+    pp_repair, pp_cleanup, pp_decimate, pp_target_faces,
+    pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,
+    pp_albedo, pp_material, pp_ao, pp_ao_quality,
+    pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m,
+    rig_type, rig_seed, rig_format,
+    ex_engine, ex_type,
+) -> tuple[str, gr.Dropdown]:
     if not name or not name.strip():
         return "❌ Preset name cannot be empty.", ui_refresh_presets()
-    # Placeholder config — Milestone 9 wires in real settings from each tab
     config = {
-        "name": name,
-        "version": 1,
-        "created_at": time.time(),
-        "note": "Placeholder — real preset data wired in Milestone 9",
+        "name": name, "version": 1, "created_at": time.time(),
+        "stage1": {"model": gen_model, "quality": gen_quality, "seed": int(gen_seed), "tex_size": int(gen_tex_size)},
+        "stage2": {
+            "repair": pp_repair, "cleanup": pp_cleanup, "decimate": pp_decimate,
+            "target_faces": int(pp_target_faces), "symmetry": pp_symmetry,
+            "unwrap": pp_unwrap, "normal_bake": pp_normal_bake, "normal_format": pp_normal_format,
+            "albedo_bake": pp_albedo, "material_bake": pp_material, "ao": pp_ao, "ao_quality": pp_ao_quality,
+            "inpaint": pp_inpaint, "lods": pp_lods, "collision": pp_collision,
+            "pivot": pp_pivot, "scale_m": float(pp_scale_m),
+        },
+        "stage3": {"rig_type": rig_type, "seed": int(rig_seed), "format": rig_format},
+        "stage4": {"engine": ex_engine, "asset_type": ex_type},
     }
     workspace.save_preset(name, config)
     return f"✅ Saved preset: `{name}`", ui_refresh_presets()
@@ -472,7 +511,7 @@ def build_ui() -> gr.Blocks:
                     with gr.Column(scale=1):
                         rig_status = gr.Markdown("*Process an asset in Stage 2 first.*")
                 rig_btn.click(
-                    fn=stub_auto_rig,
+                    fn=handle_auto_rig,
                     inputs=[rig_type, rig_seed, rig_spring, rig_format],
                     outputs=rig_status,
                 )
@@ -503,7 +542,7 @@ def build_ui() -> gr.Blocks:
                         ex_status = gr.Markdown("*Nothing to export yet.*")
                         ex_file = gr.File(label="Download", visible=True)
                 ex_btn.click(
-                    fn=stub_export,
+                    fn=handle_export,
                     inputs=[ex_engine, ex_name, ex_type, ex_include_lods, ex_include_collision],
                     outputs=[ex_status, ex_file],
                 )
@@ -512,9 +551,10 @@ def build_ui() -> gr.Blocks:
             with gr.Tab("5. Presets", id=5):
                 gr.Markdown(
                     "### Saved configurations\n"
-                    "Save parameter sets for asset types you create repeatedly "
-                    "(e.g., \"character_UE5_hero\", \"prop_UE5_standard\"). "
-                    "Wired up properly in Milestone 9."
+                    "Save the current settings across all tabs as a named preset. "
+                    "Five defaults ship with the app: `character_UE5_hero`, "
+                    "`character_UE5_npc`, `prop_UE5_hero`, `prop_UE5_standard`, "
+                    "`environment_UE5_background`."
                 )
                 with gr.Row():
                     with gr.Column():
@@ -529,7 +569,20 @@ def build_ui() -> gr.Blocks:
                         pr_delete = gr.Button("Delete selected", variant="stop")
                         pr_status = gr.Markdown()
                 pr_refresh.click(fn=ui_refresh_presets, outputs=pr_list)
-                pr_save.click(fn=ui_save_preset, inputs=pr_name, outputs=[pr_status, pr_list])
+                pr_save.click(
+                    fn=ui_save_preset,
+                    inputs=[
+                        pr_name,
+                        gen_model, gen_quality, gen_seed, gen_tex_size,
+                        pp_repair, pp_cleanup, pp_decimate, pp_target_faces,
+                        pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,
+                        pp_albedo_bake, pp_material_bake, pp_ao, pp_ao_quality,
+                        pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m,
+                        rig_type, rig_seed, rig_format,
+                        ex_engine, ex_type,
+                    ],
+                    outputs=[pr_status, pr_list],
+                )
                 pr_delete.click(fn=ui_delete_preset, inputs=pr_list, outputs=[pr_status, pr_list])
 
             # ============ Tab 6: Diagnostics (hidden in prod, useful now) =

src/stages/stage3_rig.py

@@ -0,0 +1,82 @@
+"""
+Stage 3 — Auto-rigging via gradio_client → MohamedRashad/UniRig.
+
+UniRig API:
+  Input:  3D model file (.glb/.obj/.fbx) + integer seed
+  Output: (rigged_model_path: str, download_files: list)
+"""
+from __future__ import annotations
+
+import shutil
+import tempfile
+import time
+from pathlib import Path
+
+from src import workspace
+from src.workspace import CURRENT
+
+_rig_client = None
+
+
+def _get_rig_client():
+    global _rig_client
+    if _rig_client is None:
+        from gradio_client import Client
+        _rig_client = Client("MohamedRashad/UniRig")
+    return _rig_client
+
+
+def auto_rig(
+    rig_type: str = "Humanoid",
+    seed: int = 0,
+    output_format: str = "GLB",
+) -> tuple[Path | None, str]:
+    """
+    Send the current workspace mesh to UniRig and save the rigged output.
+    Returns (rigged_glb_path, status_message).
+    """
+    t0 = time.time()
+
+    state = workspace.get_state()
+    mesh_src = (
+        state.final_glb
+        or state.low_poly_glb
+        or state.unwrapped_glb
+        or state.cleaned_glb
+        or state.raw_gen_glb
+    )
+
+    if not mesh_src or not mesh_src.exists():
+        return None, "❌ No mesh to rig — complete Stage 2 first."
+
+    client = _get_rig_client()
+
+    try:
+        from gradio_client import handle_file
+        result = client.predict(
+            input_file=handle_file(str(mesh_src)),
+            seed=int(seed),
+            api_name="/predict",
+        )
+    except Exception as e:
+        return None, f"❌ UniRig call failed:\n```\n{type(e).__name__}: {e}\n```"
+
+    # result[0] is the primary rigged model path (temp file on our machine)
+    rigged_src = result[0] if isinstance(result, (list, tuple)) else result
+
+    if not rigged_src or not Path(rigged_src).exists():
+        return None, "❌ UniRig returned no output file."
+
+    rigged_path = CURRENT / "rigged.glb"
+    shutil.copy(str(rigged_src), rigged_path)
+
+    state.rigged_glb = rigged_path
+    state.model_used = f"{state.model_used or 'TRELLIS.2'} + UniRig"
+
+    elapsed = time.time() - t0
+    return rigged_path, (
+        f"✅ **Rigged with UniRig** · {rig_type}\n\n"
+        f"- Output: `{rigged_path.name}`\n"
+        f"- Time: {elapsed:.1f}s\n\n"
+        f"Download the rigged GLB and drop into Mixamo for free animations."
+    )

src/stages/stage4_export.py

@@ -0,0 +1,131 @@
+"""
+Stage 4 — UE5 export.
+
+Bundles all processed files into a zip drop-in for UE5:
+  SM_Name.glb           — static mesh (or SK_ for skeletal)
+  SM_Name_LOD1/2.glb    — LOD levels
+  UCX_SM_Name.glb       — collision mesh
+  T_Name_D.png          — albedo / diffuse
+  T_Name_N.png          — normal map (DirectX, Y-down)
+  T_Name_ORM.png        — ORM packed (AO/Roughness/Metallic)
+  manifest.json         — metadata for re-import scripting
+"""
+from __future__ import annotations
+
+import json
+import zipfile
+from pathlib import Path
+
+from src import workspace
+from src.workspace import EXPORTS
+
+
+def _ue5_prefix(asset_type: str) -> str:
+    if "Character" in asset_type or "SK_" in asset_type:
+        return "SK_"
+    return "SM_"
+
+
+def _checklist(state: workspace.AssetState) -> list[str]:
+    """Return a list of game-readiness warnings."""
+    issues = []
+    if not (state.raw_gen_glb or state.high_poly_glb):
+        issues.append("No generated mesh — run Stage 1 first")
+    if not state.unwrapped_glb:
+        issues.append("No UV unwrap — textures will be missing")
+    if not state.normal_dx_png:
+        issues.append("No normal map — surface detail will be lost in engine")
+    if not state.orm_png and not (state.roughness_png and state.metallic_png):
+        issues.append("No ORM map — PBR material will use defaults")
+    if not state.lod_glbs:
+        issues.append("No LODs — performance warning for real-time use")
+    if not state.collision_glb:
+        issues.append("No collision mesh — physics will use the render mesh")
+    return issues
+
+
+def export_ue5(
+    asset_name: str,
+    asset_type: str = "Prop (SM_)",
+    include_lods: bool = True,
+    include_collision: bool = True,
+) -> tuple[Path, str, list[str]]:
+    """
+    Bundle workspace outputs into a UE5-ready zip.
+
+    Returns (zip_path, status_message, checklist_warnings).
+    """
+    state = workspace.get_state()
+    issues = _checklist(state)
+
+    prefix = _ue5_prefix(asset_type)
+    safe_name = "".join(c for c in asset_name if c.isalnum() or c == "_").strip("_") or "Asset"
+
+    EXPORTS.mkdir(parents=True, exist_ok=True)
+    zip_path = EXPORTS / f"export_{safe_name}_UE5.zip"
+
+    # Determine the main mesh source (most processed version)
+    mesh_src = (
+        state.final_glb
+        or state.low_poly_glb
+        or state.unwrapped_glb
+        or state.cleaned_glb
+        or state.repaired_glb
+        or state.raw_gen_glb
+    )
+
+    if not mesh_src or not mesh_src.exists():
+        return zip_path, "❌ No mesh to export — complete Stage 2 first.", issues
+
+    added = []
+    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+
+        # Main mesh
+        zf.write(mesh_src, f"{prefix}{safe_name}.glb")
+        added.append(f"{prefix}{safe_name}.glb")
+
+        # LODs
+        if include_lods:
+            for i, lod in enumerate(state.lod_glbs or []):
+                if lod and lod.exists() and i > 0:  # LOD0 == main mesh
+                    arc_name = f"{prefix}{safe_name}_LOD{i}.glb"
+                    zf.write(lod, arc_name)
+                    added.append(arc_name)
+
+        # Collision (UE5 uses UCX_ prefix for convex collision)
+        if include_collision and state.collision_glb and state.collision_glb.exists():
+            arc_name = f"UCX_{prefix}{safe_name}.glb"
+            zf.write(state.collision_glb, arc_name)
+            added.append(arc_name)
+
+        # Textures
+        tex_map: dict[Path | None, str] = {
+            state.albedo_png:    f"T_{safe_name}_D.png",
+            state.normal_dx_png: f"T_{safe_name}_N.png",   # DX convention for UE5
+            state.orm_png:       f"T_{safe_name}_ORM.png",
+        }
+        for src, arc_name in tex_map.items():
+            if src and src.exists():
+                zf.write(src, arc_name)
+                added.append(arc_name)
+
+        # Manifest
+        manifest = {
+            "asset_name": safe_name,
+            "asset_type": asset_type,
+            "engine": "UE5",
+            "normal_convention": "DirectX",
+            "scale_unit": "cm",
+            "model_used": state.model_used or "unknown",
+            "face_count": state.face_count,
+            "vertex_count": state.vertex_count,
+            "files": added,
+        }
+        zf.writestr("manifest.json", json.dumps(manifest, indent=2))
+
+    size_kb = zip_path.stat().st_size / 1024
+    msg = (
+        f"✅ **Exported:** `{zip_path.name}` ({size_kb:.0f} KB)\n\n"
+        f"Files: {', '.join(added)}"
+    )
+    return zip_path, msg, issues

workspace/presets/character_UE5_hero.json

@@ -0,0 +1,38 @@
+{
+  "name": "character_UE5_hero",
+  "version": 1,
+  "stage1": {
+    "model": "TRELLIS.2 (Hard Surface)",
+    "quality": "Hero (~90s)",
+    "seed": 42,
+    "tex_size": 4096
+  },
+  "stage2": {
+    "repair": true,
+    "cleanup": true,
+    "decimate": true,
+    "target_faces": 20000,
+    "symmetry": true,
+    "unwrap": true,
+    "normal_bake": true,
+    "normal_format": "DirectX (UE5)",
+    "albedo_bake": true,
+    "material_bake": true,
+    "ao": true,
+    "ao_quality": "High",
+    "inpaint": false,
+    "lods": true,
+    "collision": true,
+    "pivot": "bottom_center",
+    "scale_m": 1.8
+  },
+  "stage3": {
+    "rig_type": "Humanoid",
+    "seed": 0,
+    "format": "FBX (UE5 recommended)"
+  },
+  "stage4": {
+    "engine": "UE5",
+    "asset_type": "Character (SK_)"
+  }
+}

workspace/presets/character_UE5_npc.json

@@ -0,0 +1,38 @@
+{
+  "name": "character_UE5_npc",
+  "version": 1,
+  "stage1": {
+    "model": "TRELLIS.2 (Hard Surface)",
+    "quality": "Balanced (~60s)",
+    "seed": 42,
+    "tex_size": 2048
+  },
+  "stage2": {
+    "repair": true,
+    "cleanup": true,
+    "decimate": true,
+    "target_faces": 10000,
+    "symmetry": true,
+    "unwrap": true,
+    "normal_bake": true,
+    "normal_format": "DirectX (UE5)",
+    "albedo_bake": true,
+    "material_bake": true,
+    "ao": true,
+    "ao_quality": "Standard",
+    "inpaint": false,
+    "lods": true,
+    "collision": false,
+    "pivot": "bottom_center",
+    "scale_m": 1.75
+  },
+  "stage3": {
+    "rig_type": "Humanoid",
+    "seed": 0,
+    "format": "FBX (UE5 recommended)"
+  },
+  "stage4": {
+    "engine": "UE5",
+    "asset_type": "Character (SK_)"
+  }
+}

workspace/presets/environment_UE5_background.json

@@ -0,0 +1,38 @@
+{
+  "name": "environment_UE5_background",
+  "version": 1,
+  "stage1": {
+    "model": "TRELLIS.2 (Hard Surface)",
+    "quality": "Fast (~30s)",
+    "seed": 42,
+    "tex_size": 1024
+  },
+  "stage2": {
+    "repair": true,
+    "cleanup": true,
+    "decimate": true,
+    "target_faces": 5000,
+    "symmetry": false,
+    "unwrap": true,
+    "normal_bake": false,
+    "normal_format": "DirectX (UE5)",
+    "albedo_bake": true,
+    "material_bake": false,
+    "ao": false,
+    "ao_quality": "Fast",
+    "inpaint": false,
+    "lods": true,
+    "collision": true,
+    "pivot": "bottom_center",
+    "scale_m": 2.0
+  },
+  "stage3": {
+    "rig_type": "Humanoid",
+    "seed": 0,
+    "format": "GLB"
+  },
+  "stage4": {
+    "engine": "UE5",
+    "asset_type": "Environment (SM_)"
+  }
+}

workspace/presets/prop_UE5_hero.json

@@ -0,0 +1,38 @@
+{
+  "name": "prop_UE5_hero",
+  "version": 1,
+  "stage1": {
+    "model": "TRELLIS.2 (Hard Surface)",
+    "quality": "Hero (~90s)",
+    "seed": 42,
+    "tex_size": 4096
+  },
+  "stage2": {
+    "repair": true,
+    "cleanup": true,
+    "decimate": true,
+    "target_faces": 15000,
+    "symmetry": false,
+    "unwrap": true,
+    "normal_bake": true,
+    "normal_format": "DirectX (UE5)",
+    "albedo_bake": true,
+    "material_bake": true,
+    "ao": true,
+    "ao_quality": "High",
+    "inpaint": false,
+    "lods": true,
+    "collision": true,
+    "pivot": "bottom_center",
+    "scale_m": 0.5
+  },
+  "stage3": {
+    "rig_type": "Humanoid",
+    "seed": 0,
+    "format": "GLB"
+  },
+  "stage4": {
+    "engine": "UE5",
+    "asset_type": "Prop (SM_)"
+  }
+}

workspace/presets/prop_UE5_standard.json

@@ -0,0 +1,38 @@
+{
+  "name": "prop_UE5_standard",
+  "version": 1,
+  "stage1": {
+    "model": "TRELLIS.2 (Hard Surface)",
+    "quality": "Balanced (~60s)",
+    "seed": 42,
+    "tex_size": 2048
+  },
+  "stage2": {
+    "repair": true,
+    "cleanup": true,
+    "decimate": true,
+    "target_faces": 8000,
+    "symmetry": false,
+    "unwrap": true,
+    "normal_bake": true,
+    "normal_format": "DirectX (UE5)",
+    "albedo_bake": true,
+    "material_bake": true,
+    "ao": true,
+    "ao_quality": "Standard",
+    "inpaint": false,
+    "lods": true,
+    "collision": true,
+    "pivot": "bottom_center",
+    "scale_m": 0.5
+  },
+  "stage3": {
+    "rig_type": "Humanoid",
+    "seed": 0,
+    "format": "GLB"
+  },
+  "stage4": {
+    "engine": "UE5",
+    "asset_type": "Prop (SM_)"
+  }
+}