fix: simplify app.py for Hugging Face deployment

app.py

@@ -1,30 +1,24 @@
-# tout en haut (cache persistant pour éviter de re-télécharger à chaque boot)
 import os
-os.environ.setdefault("HF_HOME", "/data/.cache/huggingface")
-os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/data/.cache/huggingface")
-os.environ.setdefault("TRANSFORMERS_CACHE", "/data/.cache/huggingface")
-os.environ.setdefault("XDG_CACHE_HOME", "/data/.cache")
-os.environ.setdefault("U2NET_HOME", "/data/.u2net")
-
+import spaces
+import random
+import shutil
+import gradio as gr
+from glob import glob
+from pathlib import Path
+import uuid
 import argparse
 import torch
+import trimesh
 
-# -------- args --------
+# -------------------
+# Arguments & Config
+# -------------------
 parser = argparse.ArgumentParser()
 parser.add_argument("--model_path", type=str, default="tencent/Hunyuan3D-2mini")
 parser.add_argument("--subfolder", type=str, default="hunyuan3d-dit-v2-mini-turbo")
 parser.add_argument("--texgen_model_path", type=str, default="tencent/Hunyuan3D-2")
 parser.add_argument("--port", type=int, default=7860)
 parser.add_argument("--host", type=str, default="0.0.0.0")
-
-# auto-device: cuda > mps > cpu
-_auto_device = (
-    "cuda" if torch.cuda.is_available()
-    else ("mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available()
-          else "cpu")
-)
-parser.add_argument("--device", type=str, default=_auto_device)
-parser.add_argument("--mc_algo", type=str, default="mc")
 parser.add_argument("--cache_path", type=str, default="gradio_cache")
 parser.add_argument("--enable_t23d", action="store_true")
 parser.add_argument("--disable_tex", action="store_true")
@@ -32,44 +26,242 @@ parser.add_argument("--enable_flashvdm", action="store_true")
 parser.add_argument("--compile", action="store_true")
 parser.add_argument("--low_vram_mode", action="store_true")
 args = parser.parse_args()
+args.enable_flashvdm = True
 
-# fp16 sur cuda, sinon fp32
-DTYPE = torch.float16 if args.device == "cuda" else torch.float32
+# Choix auto du device
+if torch.cuda.is_available():
+    args.device = "cuda"
+elif torch.backends.mps.is_available():
+    args.device = "mps"
+else:
+    args.device = "cpu"
 
-# Si flashvdm activé, ne l’utilise pas en CPU/MPS
-args.enable_flashvdm = args.enable_flashvdm or True  # si tu veux garder True par défaut
-if args.device in ["cpu", "mps"]:
-    args.enable_flashvdm = False  # évite les features GPU-only
+SAVE_DIR = args.cache_path
+os.makedirs(SAVE_DIR, exist_ok=True)
 
-# -------- lazy load des workers --------
-from functools import lru_cache
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 
-@lru_cache(maxsize=1)
-def get_workers():
-    from hy3dgen.shapegen import (
-        FaceReducer, FloaterRemover, DegenerateFaceRemover, Hunyuan3DDiTFlowMatchingPipeline
-    )
-    from hy3dgen.rembg import BackgroundRemover
-
-    rmbg_worker = BackgroundRemover()
-    i23d_worker = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
-        args.model_path,
-        subfolder=args.subfolder,
-        use_safetensors=True,
-        device=args.device,        # la pipeline sait gérer cpu/mps/cuda
+HTML_HEIGHT = 500
+HTML_WIDTH = 500
+MAX_SEED = 1e7
+
+# -------------------
+# Utils
+# -------------------
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+def gen_save_folder(max_size=200):
+    os.makedirs(SAVE_DIR, exist_ok=True)
+    dirs = [f for f in Path(SAVE_DIR).iterdir() if f.is_dir()]
+    if len(dirs) >= max_size:
+        oldest_dir = min(dirs, key=lambda x: x.stat().st_ctime)
+        shutil.rmtree(oldest_dir)
+        print(f"Removed oldest folder: {oldest_dir}")
+    new_folder = os.path.join(SAVE_DIR, str(uuid.uuid4()))
+    os.makedirs(new_folder, exist_ok=True)
+    return new_folder
+
+def export_mesh(mesh, save_folder, textured=False, type="glb"):
+    if textured:
+        path = os.path.join(save_folder, f"textured_mesh.{type}")
+    else:
+        path = os.path.join(save_folder, f"white_mesh.{type}")
+    if type not in ["glb", "obj"]:
+        mesh.export(path)
+    else:
+        mesh.export(path, include_normals=textured)
+    return path
+
+def build_model_viewer_html(save_folder, height=660, width=790, textured=False):
+    if textured:
+        related_path = f"./textured_mesh.glb"
+        template_name = "./assets/modelviewer-textured-template.html"
+        output_html_path = os.path.join(save_folder, f"textured_mesh.html")
+    else:
+        related_path = f"./white_mesh.glb"
+        template_name = "./assets/modelviewer-template.html"
+        output_html_path = os.path.join(save_folder, f"white_mesh.html")
+
+    offset = 50 if textured else 10
+    with open(os.path.join(CURRENT_DIR, template_name), "r", encoding="utf-8") as f:
+        template_html = f.read()
+
+    with open(output_html_path, "w", encoding="utf-8") as f:
+        template_html = template_html.replace("#height#", f"{height - offset}")
+        template_html = template_html.replace("#width#", f"{width}")
+        template_html = template_html.replace("#src#", f"{related_path}/")
+        f.write(template_html)
+
+    rel_path = os.path.relpath(output_html_path, SAVE_DIR)
+    iframe_tag = f'<iframe src="/static/{rel_path}" height="{height}" width="100%" frameborder="0"></iframe>'
+    return f"<div style='height: {height}; width: 100%;'>{iframe_tag}</div>"
+
+# -------------------
+# Model loaders
+# -------------------
+from hy3dgen.shapegen import (
+    FaceReducer,
+    FloaterRemover,
+    DegenerateFaceRemover,
+    MeshSimplifier,
+    Hunyuan3DDiTFlowMatchingPipeline,
+)
+from hy3dgen.shapegen.pipelines import export_to_trimesh
+from hy3dgen.rembg import BackgroundRemover
+
+rmbg_worker = BackgroundRemover()
+i23d_worker = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
+    args.model_path,
+    subfolder=args.subfolder,
+    use_safetensors=True,
+    device=args.device,
+)
+if args.enable_flashvdm:
+    mc_algo = "mc" if args.device in ["cpu", "mps"] else "mc"
+    i23d_worker.enable_flashvdm(mc_algo=mc_algo)
+if args.compile:
+    i23d_worker.compile()
+
+floater_remove_worker = FloaterRemover()
+degenerate_face_remove_worker = DegenerateFaceRemover()
+face_reduce_worker = FaceReducer()
+
+progress = gr.Progress()
+
+# -------------------
+# Main function
+# -------------------
+# ⚠️ Désactive si pas de GPU sur ton Space
+# @spaces.GPU(duration=40)
+def gen_shape(
+    image=None,
+    steps=50,
+    guidance_scale=7.5,
+    seed=1234,
+    octree_resolution=256,
+    num_chunks=200000,
+    target_face_num=10000,
+    randomize_seed: bool = False,
+):
+    progress(0, desc="Starting")
+
+    def callback(step_idx, timestep, outputs):
+        progress_value = ((step_idx + 1.0) / steps) * (0.5 / 1.0)
+        progress(progress_value, desc=f"Mesh generating, {step_idx + 1}/{steps} steps")
+
+    if image is None:
+        raise gr.Error("Please provide an image.")
+
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    save_folder = gen_save_folder()
+    image = rmbg_worker(image.convert("RGB"))
+
+    generator = torch.Generator().manual_seed(seed)
+    outputs = i23d_worker(
+        image=image,
+        num_inference_steps=steps,
+        guidance_scale=guidance_scale,
+        generator=generator,
+        octree_resolution=octree_resolution,
+        num_chunks=num_chunks,
+        output_type="mesh",
+        callback=callback,
+        callback_steps=1,
     )
 
-    # cast correct pour le device courant
-    try:
-        i23d_worker.to(args.device, DTYPE)
-    except Exception:
-        # dernier filet de secours pour CPU
-        i23d_worker.to("cpu", torch.float32)
+    mesh = export_to_trimesh(outputs)[0]
+    path = export_mesh(mesh, save_folder, textured=False)
+
+    if args.low_vram_mode:
+        torch.cuda.empty_cache()
+    if path is None:
+        raise gr.Error("Mesh generation failed.")
+
+    mesh = trimesh.load(path)
+    progress(0.5, desc="Optimizing mesh")
+    mesh = floater_remove_worker(mesh)
+    mesh = degenerate_face_remove_worker(mesh)
+    progress(0.6, desc="Reducing mesh faces")
+    mesh = face_reduce_worker(mesh, target_face_num)
+
+    save_folder = gen_save_folder()
+    progress(0.9, desc="Converting format")
+    sourceObjPath = export_mesh(mesh, save_folder, textured=False, type="obj")
+
+    rel_objPath = os.path.relpath(sourceObjPath, SAVE_DIR)
+    objPath = "/static/" + rel_objPath
+
+    save_folder = gen_save_folder()
+    _ = export_mesh(mesh, save_folder, textured=False)
+    model_viewer_html = build_model_viewer_html(save_folder, height=HTML_HEIGHT, width=HTML_WIDTH, textured=False)
+
+    glbPath = os.path.join(save_folder, f"white_mesh.glb")
+    rel_glbPath = os.path.relpath(glbPath, SAVE_DIR)
+    glbPath = "/static/" + rel_glbPath
+
+    progress(1, desc="Complete")
+    return model_viewer_html, gr.update(value=sourceObjPath, interactive=True), glbPath, objPath
+
+# -------------------
+# Gradio UI
+# -------------------
+def get_example_img_list():
+    return sorted(glob("./assets/example_images/**/*.png", recursive=True))
+example_imgs = get_example_img_list()
+
+HTML_OUTPUT_PLACEHOLDER = """
+<div style='height: 500px; width: 100%; border-radius: 8px; border: 1px solid #e5e7eb; display: flex; justify-content: center; align-items: center;'>
+    <div style="text-align: center; font-size: 16px; color: #6b7280;">
+        <p style="color: #8d8d8d;">No mesh here.</p>
+    </div>
+</div>
+"""
+
+title = "## AI 3D Model Generator"
+description = "Upload an image and generate a 3D mesh using Hunyuan 3D. Ready for AR/VR, games or 3D printing."
+
+with gr.Blocks().queue() as demo:
+    gr.Markdown(title)
+    gr.Markdown(description)
+    with gr.Row():
+        with gr.Column(scale=3):
+            image = gr.Image(sources=["upload"], label="Image", type="pil", image_mode="RGBA", height=290)
+            gen_button = gr.Button(value="Generate Shape", variant="primary")
+            with gr.Accordion("Advanced Options", open=False):
+                with gr.Column():
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=1234)
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                with gr.Column():
+                    num_steps = gr.Slider(maximum=100, minimum=1, value=5, step=1, label="Inference Steps")
+                    octree_resolution = gr.Slider(maximum=512, minimum=16, value=256, label="Octree Resolution")
+                with gr.Column():
+                    cfg_scale = gr.Slider(maximum=20.0, minimum=1.0, value=5.5, step=0.1, label="Guidance Scale")
+                    num_chunks = gr.Slider(maximum=5000000, minimum=1000, value=8000, label="Number of Chunks")
+                target_face_num = gr.Slider(maximum=1000000, minimum=100, value=10000, label="Target Face Number")
+
+        with gr.Column(scale=6):
+            html_export_mesh = gr.HTML(HTML_OUTPUT_PLACEHOLDER, label="Output")
+            file_export = gr.DownloadButton(label="Download", variant="primary", interactive=False)
+            with gr.Row():
+                objPath_output = gr.Text(label="Obj Path", interactive=False)
+                glbPath_output = gr.Text(label="Glb Path", interactive=False)
+
+        with gr.Column(scale=3):
+            gr.Examples(examples=example_imgs, inputs=[image], examples_per_page=18)
+
+    gen_button.click(
+        fn=gen_shape,
+        inputs=[image, num_steps, cfg_scale, seed, octree_resolution, num_chunks, target_face_num, randomize_seed],
+        outputs=[html_export_mesh, file_export, glbPath_output, objPath_output],
+    )
 
-    if args.enable_flashvdm and args.device == "cuda":
-        i23d_worker.enable_flashvdm(mc_algo=args.mc_algo)
+# -------------------
+# Hugging Face Spaces entrypoint
+# -------------------
+app = demo
 
-    floater_remove_worker = FloaterRemover()
-    degenerate_face_remove_worker = DegenerateFaceRemover()
-    face_reduce_worker = FaceReducer()
-    return rmbg_worker, i23d_worker, floater_remove_worker, degenerate_face_remove_worker, face_reduce_worker
+if __name__ == "__main__":
+    demo.launch(server_name=args.host, server_port=args.port)