Implement Hunyuan3D-2.1 generation with background removal

- generate_hunyuan() is now fully implemented using the vendored
hy3dshape package (tencent/Hunyuan3D-2.1, diffusers-based, no
custom CUDA wheels, works on torch 2.8.0 + ZeroGPU H200)
- @spaces.GPU(duration=240) wraps the local inference call
- Pipeline is loaded lazily on first GPU request and cached globally
- Background removal via rembg (IS-Net) with fallback to hy3dshape's
BackgroundRemover; toggled by do_rembg parameter
- Multi-view support: >1 image routed as {front/left/back/right} dict
- Working mesh is decimated via PyMeshLab to preset face target; full
mesh saved as high_poly.glb for Stage 2 baking
- Added gen_rembg checkbox to the Generate tab (default: on)
- Hunyuan3D-2.1 is now the default model (moved to top of radio)
- Preset save/load includes rembg flag (n=27 components)
- requirements.txt: added diffusers>=0.30.0, accelerate, einops, rembg

app.py

@@ -81,11 +81,11 @@ def _stub(stage: str) -> str:
     )
 
 
-def handle_generate(images, model, quality, seed, _steps, _octree, tex_size, _symmetry):
+def handle_generate(images, model, quality, seed, _steps, _octree, tex_size, _symmetry, do_rembg):
     """Dispatch to the correct generation backend based on model selection."""
     if "TRELLIS" in model:
         return generate_trellis(images, quality, int(seed), int(tex_size))
-    return generate_hunyuan(images, quality, int(seed), int(tex_size))
+    return generate_hunyuan(images, quality, int(seed), int(tex_size), do_rembg=bool(do_rembg))
 
 
 def run_post_process(
@@ -284,7 +284,7 @@ def ui_refresh_presets() -> gr.Dropdown:
 
 def ui_save_preset(
     name: str,
-    gen_model, gen_quality, gen_seed, gen_tex_size,
+    gen_model, gen_quality, gen_seed, gen_tex_size, gen_rembg,
     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,
@@ -296,7 +296,7 @@ def ui_save_preset(
         return "❌ Preset name cannot be empty.", ui_refresh_presets()
     config = {
         "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)},
+        "stage1": {"model": gen_model, "quality": gen_quality, "seed": int(gen_seed), "tex_size": int(gen_tex_size), "rembg": bool(gen_rembg)},
         "stage2": {
             "repair": pp_repair, "cleanup": pp_cleanup, "decimate": pp_decimate,
             "target_faces": int(pp_target_faces), "symmetry": pp_symmetry,
@@ -315,7 +315,7 @@ def ui_save_preset(
 def ui_load_preset(name: str):
     """Apply a saved preset to all UI components."""
     _no_op = gr.update()
-    n = 26  # number of component outputs after pr_status
+    n = 27  # number of component outputs after pr_status
     if not name:
         return ("❌ Select a preset first.",) + (_no_op,) * n
     try:
@@ -331,10 +331,11 @@ def ui_load_preset(name: str):
     return (
         f"✅ Loaded preset: `{name}`",
         # Stage 1
-        gr.update(value=s1.get("model", "TRELLIS.2 (Hard Surface)")),
+        gr.update(value=s1.get("model", "Hunyuan3D-2.1 (Organic / Characters)")),
         gr.update(value=s1.get("quality", "Balanced (~60s)")),
         gr.update(value=s1.get("seed", 42)),
         gr.update(value=s1.get("tex_size", 2048)),
+        gr.update(value=s1.get("rembg", True)),
         # Stage 2
         gr.update(value=s2.get("repair", True)),
         gr.update(value=s2.get("cleanup", True)),
@@ -424,10 +425,10 @@ def build_ui() -> gr.Blocks:
                         )
                         gen_model = gr.Radio(
                             choices=[
+                                "Hunyuan3D-2.1 (Organic / Characters)",
                                 "TRELLIS.2 (Hard Surface)",
-                                "Hunyuan3D-2 (Organic / Characters)",
                             ],
-                            value="TRELLIS.2 (Hard Surface)",
+                            value="Hunyuan3D-2.1 (Organic / Characters)",
                             label="Generation model",
                         )
                         gen_quality = gr.Radio(
@@ -435,6 +436,12 @@ def build_ui() -> gr.Blocks:
                             value="Balanced (~60s)",
                             label="Quality preset",
                         )
+                        gen_rembg = gr.Checkbox(
+                            value=True,
+                            label="Remove background automatically",
+                            info="Uses rembg IS-Net to strip background before generation. "
+                                 "Disable if your image already has a transparent background.",
+                        )
                         with gr.Accordion("Advanced", open=False):
                             gen_seed = gr.Number(value=42, label="Seed", precision=0)
                             gen_steps = gr.Slider(
@@ -442,7 +449,7 @@ def build_ui() -> gr.Blocks:
                                 label="Inference steps",
                             )
                             gen_octree = gr.Dropdown(
-                                choices=[128, 192, 256], value=192,
+                                choices=[256, 384, 512], value=384,
                                 label="Octree resolution",
                             )
                             gen_tex_size = gr.Dropdown(
@@ -460,7 +467,8 @@ def build_ui() -> gr.Blocks:
                 gen_btn.click(
                     fn=handle_generate,
                     inputs=[gen_images, gen_model, gen_quality, gen_seed,
-                            gen_steps, gen_octree, gen_tex_size, gen_symmetry],
+                            gen_steps, gen_octree, gen_tex_size, gen_symmetry,
+                            gen_rembg],
                     outputs=gen_status,
                 )
 
@@ -627,7 +635,7 @@ def build_ui() -> gr.Blocks:
                     inputs=pr_list,
                     outputs=[
                         pr_status,
-                        gen_model, gen_quality, gen_seed, gen_tex_size,
+                        gen_model, gen_quality, gen_seed, gen_tex_size, gen_rembg,
                         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,
@@ -640,7 +648,7 @@ def build_ui() -> gr.Blocks:
                     fn=ui_save_preset,
                     inputs=[
                         pr_name,
-                        gen_model, gen_quality, gen_seed, gen_tex_size,
+                        gen_model, gen_quality, gen_seed, gen_tex_size, gen_rembg,
                         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,

requirements.txt

@@ -7,12 +7,12 @@ spaces                      # HF ZeroGPU decorator
 numpy
 pillow
 
-# ===== Milestone 2: Stage 1 Generation (via gradio_client) =====
-# NOTE: TRELLIS.2 generation runs on microsoft/TRELLIS.2 via gradio_client.
-# The custom wheels (o_voxel, cumesh, flex_gemm, nvdiffrast) are ABI-incompatible
-# with torch 2.8.0+ — they were compiled for torch 2.6.0 which ZeroGPU no longer
-# supports. Remove them until Jeff Xiang releases torch 2.8+ compatible wheels.
-gradio_client               # call TRELLIS.2 / Hunyuan3D-2.1 Spaces
+# ===== Milestone 2: Stage 1 Generation =====
+# TRELLIS.2 generation: via gradio_client (remote Space) — ABI-incompatible custom
+# wheels (o_voxel, cumesh, flex_gemm) cannot run locally on torch 2.8.0+.
+# Hunyuan3D-2.1 generation: local inference via vendored hy3dshape package —
+# diffusers-based, no custom CUDA extensions, works on torch 2.8.0 + ZeroGPU.
+gradio_client               # call TRELLIS.2 Space
 torch==2.8.0
 torchvision==0.23.0
 imageio==2.37.2
@@ -26,6 +26,12 @@ zstandard==0.25.0
 kornia==0.8.2
 timm==1.0.22
 git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
+# Hunyuan3D-2.1 dependencies
+diffusers>=0.30.0
+accelerate
+einops
+# Background removal (used before shape generation)
+rembg
 
 # ===== Milestone 3: Stage 2A-2C Mesh Cleanup =====
 trimesh[easy]

src/stages/stage1_generate.py

@@ -1,25 +1,34 @@
 """
 Stage 1 — Image to 3D generation.
 
-Generation is offloaded via gradio_client to the official TRELLIS.2 / Hunyuan3D
-HuggingFace Spaces.  All the TRELLIS.2 custom wheels (o_voxel, cumesh, flex_gemm)
-are compiled against torch 2.6.0 and cannot be loaded under torch 2.8.0+ due to
-C++ ABI changes.  This client approach requires no local GPU for Stage 1; the
-remote Space does the heavy lifting, returns two GLB files, and we copy them into
-our workspace for Stage 2 processing.
+TRELLIS.2 generation is offloaded via gradio_client to the official HuggingFace
+Space (microsoft/TRELLIS.2) because the TRELLIS.2 custom wheels (o_voxel, cumesh,
+flex_gemm) are compiled for torch 2.6.0 and are ABI-broken on torch 2.8.0+.
+
+Hunyuan3D-2.1 runs locally using the vendored hy3dshape package.  The shape
+pipeline is diffusers-based with no custom CUDA extensions — it works on torch
+2.8.0 + ZeroGPU H200.  The model is streamed from tencent/Hunyuan3D-2.1 on first
+run and cached by HuggingFace Hub.
 """
 from __future__ import annotations
 
 import shutil
+import sys
 import tempfile
 import time
 from pathlib import Path
 
+import spaces
 from PIL import Image
 
 from src import workspace
 from src.workspace import CURRENT
 
+# Make the vendored hy3dshape package importable.
+_HY3D_PATH = Path(__file__).resolve().parents[2] / "Hunyuan3D-2.1" / "hy3dshape"
+if str(_HY3D_PATH) not in sys.path:
+    sys.path.insert(0, str(_HY3D_PATH))
+
 # ---------------------------------------------------------------------------
 # Quality presets → TRELLIS.2 sampler parameters
 # ---------------------------------------------------------------------------
@@ -180,16 +189,162 @@ def generate_trellis(
 
 
 # ---------------------------------------------------------------------------
-# Hunyuan3D-2.1 (Milestone 2b — stub)
+# Hunyuan3D-2.1 — local inference using vendored hy3dshape
 # ---------------------------------------------------------------------------
 
+# Quality presets for Hunyuan3D-2.1
+_HY_PRESETS = {
+    "Fast (~30s)":     {"steps": 20, "octree_resolution": 256,  "decimation_target": 50_000},
+    "Balanced (~60s)": {"steps": 35, "octree_resolution": 384,  "decimation_target": 100_000},
+    "Hero (~90s)":     {"steps": 50, "octree_resolution": 512,  "decimation_target": 200_000},
+}
+
+# Pipeline loaded lazily on first GPU call to avoid startup cost.
+_hy_pipeline = None
+
+
+def _get_hy_pipeline():
+    global _hy_pipeline
+    if _hy_pipeline is None:
+        from hy3dshape.pipelines import Hunyuan3DDiTFlowMatchingPipeline
+        _hy_pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
+            "tencent/Hunyuan3D-2.1",
+            subfolder="hunyuan3d-dit-v2-1",
+            use_safetensors=False,
+            variant="fp16",
+        )
+    return _hy_pipeline
+
+
+def _remove_background(image: Image.Image) -> Image.Image:
+    """Strip background using rembg (IS-Net model)."""
+    try:
+        from rembg import remove, new_session
+        session = new_session()
+        return remove(image, session=session, bgcolor=[255, 255, 255, 0])
+    except Exception as e:
+        # Fall back to the hy3dshape wrapper if rembg is unavailable.
+        try:
+            from hy3dshape.rembg import BackgroundRemover
+            remover = BackgroundRemover()
+            return remover(image)
+        except Exception:
+            raise RuntimeError(f"Background removal failed: {e}")
+
+
+@spaces.GPU(duration=240)
 def generate_hunyuan(
     images: list | None,
     quality: str,
     seed: int,
     tex_size: int = 2048,
+    do_rembg: bool = True,
 ) -> str:
+    """Generate a 3D mesh with Hunyuan3D-2.1 and save it to the workspace."""
+    import torch
+
+    t0 = time.time()
+
+    if not images:
+        return "❌ Please upload at least one reference image."
+
+    # ------------------------------------------------------------------
+    # Load and optionally remove background from input image(s)
+    # ------------------------------------------------------------------
+    def _open(f) -> Image.Image:
+        path = f.name if hasattr(f, "name") else str(f)
+        return Image.open(path).convert("RGBA")
+
+    pil_images = [_open(f) for f in images]
+
+    if do_rembg:
+        try:
+            pil_images = [_remove_background(img) for img in pil_images]
+        except Exception as e:
+            return f"❌ Background removal failed:\n```\n{e}\n```"
+
+    preset = _HY_PRESETS.get(quality, _HY_PRESETS["Balanced (~60s)"])
+
+    # Multi-view: pass as dict if more than one image is provided.
+    if len(pil_images) == 1:
+        image_input = pil_images[0]
+    else:
+        view_keys = ["front", "left", "back", "right"]
+        image_input = {view_keys[i]: img for i, img in enumerate(pil_images[:4])}
+
+    # ------------------------------------------------------------------
+    # Shape generation
+    # ------------------------------------------------------------------
+    try:
+        pipeline = _get_hy_pipeline()
+        generator = torch.Generator().manual_seed(int(seed))
+        meshes = pipeline(
+            image=image_input,
+            num_inference_steps=preset["steps"],
+            octree_resolution=preset["octree_resolution"],
+            mc_level=0.0,
+            output_type="trimesh",
+            generator=generator,
+        )
+        mesh = meshes[0] if isinstance(meshes, list) else meshes
+    except Exception as e:
+        return (
+            f"❌ Hunyuan3D-2.1 generation failed:\n```\n{type(e).__name__}: {e}\n```\n\n"
+            f"Tip: make sure the image has a clean subject on a white/transparent background."
+        )
+
+    # ------------------------------------------------------------------
+    # Save to workspace
+    # ------------------------------------------------------------------
+    workspace.reset_current()
+    ws = workspace.reset_state()
+
+    high_poly_path = CURRENT / "high_poly.glb"
+    raw_gen_path = CURRENT / "raw_gen.glb"
+    mesh.export(str(high_poly_path))
+    # Decimate for the working copy
+    try:
+        import pymeshlab
+        import trimesh as _tm
+        with tempfile.NamedTemporaryFile(suffix=".ply", delete=False) as tmp:
+            ply_in = tmp.name
+        mesh.export(ply_in)
+        ms = pymeshlab.MeshSet()
+        ms.load_new_mesh(ply_in)
+        target = max(100, preset["decimation_target"])
+        ms.apply_filter(
+            "meshing_decimation_quadric_edge_collapse",
+            targetfacenum=target,
+            preservenormal=True,
+            autoclean=True,
+        )
+        with tempfile.NamedTemporaryFile(suffix=".ply", delete=False) as tmp2:
+            ply_out = tmp2.name
+        ms.save_current_mesh(ply_out)
+        working = _tm.load(ply_out, force="mesh")
+        working.export(str(raw_gen_path))
+    except Exception:
+        # If decimation fails, just copy the full mesh.
+        shutil.copy(str(high_poly_path), str(raw_gen_path))
+
+    face_count = len(mesh.faces)
+    vertex_count = len(mesh.vertices)
+
+    ws.high_poly_glb = high_poly_path
+    ws.raw_gen_glb = raw_gen_path
+    ws.model_used = "Hunyuan3D-2.1"
+    ws.face_count = face_count
+    ws.vertex_count = vertex_count
+
+    elapsed = time.time() - t0
+
+    from src import quota as _quota
+    _quota.record_usage(f"generate_hunyuan_{quality.split()[0].lower()}", elapsed)
+
     return (
-        "🚧 **Hunyuan3D-2.1** — not wired yet.\n\n"
-        "Will be enabled in Milestone 2b. Use TRELLIS.2 for now."
+        f"✅ **Generated with Hunyuan3D-2.1** · {quality}\n\n"
+        f"- Faces: {face_count:,} · Vertices: {vertex_count:,}\n"
+        f"- Time: {elapsed:.1f}s\n\n"
+        f"Scroll down to see the result in the 3D viewer. "
+        f"High-poly saved for Stage 2 baking."
     )