Update app.py

app.py

@@ -1,27 +1,28 @@
 import os
-import spaces
 import random
 import shutil
-import gradio as gr
+import uuid
 from glob import glob
 from pathlib import Path
-import uuid
 import argparse
+
 import torch
+import gradio as gr
 import uvicorn
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
+from PIL import Image
 import trimesh
 from transformers import AutoProcessor, AutoModelForImageClassification
-from PIL import Image
 
+# -------------------- Argumente --------------------
 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')
-parser.add_argument('--device', type=str, default='cuda')
+parser.add_argument('--device', type=str, default=None)
 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')
@@ -30,94 +31,80 @@ 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
 
+# -------------------- Device Setup --------------------
+if args.device is None:
+    if torch.cuda.is_available():
+        args.device = "cuda"
+    elif torch.backends.mps.is_available():  # macOS GPU
+        args.device = "mps"
+    else:
+        args.device = "cpu"
+
+print(f"Using device: {args.device}")
+
+# -------------------- Pfade --------------------
 SAVE_DIR = args.cache_path
 os.makedirs(SAVE_DIR, exist_ok=True)
-
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 
 HTML_HEIGHT = 500
 HTML_WIDTH = 500
+MAX_SEED = int(1e7)
 
-# -------------------- NSFW 检测模型加载 --------------------
+# -------------------- NSFW Modell --------------------
 nsfw_processor = AutoProcessor.from_pretrained("Falconsai/nsfw_image_detection")
-nsfw_model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection").to(args.device)
-# -----------------------------------------------------------
-
+nsfw_model = AutoModelForImageClassification.from_pretrained(
+    "Falconsai/nsfw_image_detection"
+).to(args.device)
 
+# -------------------- Hilfsfunktionen --------------------
 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()]
-
-    # 如果文件夹数量超过 max_size，删除创建时间最久的文件夹
     if len(dirs) >= max_size:
-        # 按创建时间排序，最久的排在前面
         oldest_dir = min(dirs, key=lambda x: x.stat().st_ctime)
         shutil.rmtree(oldest_dir)
         print(f"Removed the oldest folder: {oldest_dir}")
-
-    # 生成一个新的 uuid 文件夹名称
     new_folder = os.path.join(SAVE_DIR, str(uuid.uuid4()))
     os.makedirs(new_folder, exist_ok=True)
     print(f"Created new folder: {new_folder}")
-
     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)
+    filename = f'textured_mesh.{type}' if textured else f'white_mesh.{type}'
+    path = os.path.join(save_folder, filename)
+    mesh.export(path, include_normals=textured)
     return path
 
 def build_model_viewer_html(save_folder, height=660, width=790, textured=False):
-    # Remove first folder from path to make relative path
-    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
+    related_path = "textured_mesh.glb" if textured else "white_mesh.glb"
+    template_name = './assets/modelviewer-textured-template.html' if textured else './assets/modelviewer-template.html'
+    output_html_path = os.path.join(save_folder, f"{'textured' if textured else 'white'}_mesh.html")
+
     with open(os.path.join(CURRENT_DIR, template_name), 'r', encoding='utf-8') as f:
         template_html = f.read()
 
+    offset = 50 if textured else 10
+    template_html = template_html.replace('#height#', str(height - offset))
+    template_html = template_html.replace('#width#', str(width))
+    template_html = template_html.replace('#src#', f'./{related_path}/')
+
     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>'
-    print(
-        f'Find html file {output_html_path}, {os.path.exists(output_html_path)}, relative HTML path is /static/{rel_path}')
 
-    return f"""
-        <div style='height: {height}; width: 100%;'>
-        {iframe_tag}
-        </div>
-    """
+    return f"<div style='height: {height}; width: 100%;'>{iframe_tag}</div>"
 
-
-from hy3dgen.shapegen import FaceReducer, FloaterRemover, DegenerateFaceRemover, MeshSimplifier, \
-    Hunyuan3DDiTFlowMatchingPipeline
+# -------------------- Hy3Dgen Worker --------------------
+from hy3dgen.shapegen import FaceReducer, FloaterRemover, DegenerateFaceRemover, MeshSimplifier, Hunyuan3DDiTFlowMatchingPipeline
 from hy3dgen.shapegen.pipelines import export_to_trimesh
 from hy3dgen.rembg import BackgroundRemover
 
@@ -126,8 +113,9 @@ i23d_worker = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
     args.model_path,
     subfolder=args.subfolder,
     use_safetensors=True,
-    device=args.device,
+    device=args.device
 )
+
 if args.enable_flashvdm:
     mc_algo = 'mc' if args.device in ['cpu', 'mps'] else args.mc_algo
     i23d_worker.enable_flashvdm(mc_algo=mc_algo)
@@ -138,66 +126,39 @@ floater_remove_worker = FloaterRemover()
 degenerate_face_remove_worker = DegenerateFaceRemover()
 face_reduce_worker = FaceReducer()
 
-
+# -------------------- NSFW Detection --------------------
 def detect_nsfw(image: Image.Image, threshold: float = 0.5) -> bool:
-    """Returns True if image is NSFW"""
-    # inputs = nsfw_processor(images=image, return_tensors="pt").to(args.device)
-    # with torch.no_grad():
-    #     outputs = nsfw_model(**inputs)
-    #     probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    #     nsfw_score = probs[0][1].item()  # label 1 = NSFW
-    nsfw_score = 0  # label 1 = NSFW
+    nsfw_score = 0  # Placeholder, optional: implement actual detection
     return nsfw_score > threshold
 
+# -------------------- Mesh Generation --------------------
+progress = gr.Progress()
 
-
-progress=gr.Progress()
-
-# @spaces.GPU(duration=40)
-def _gen_shape_on_gpu(
+def _gen_shape(
     image=None,
-    steps=10, # 50
-    guidance_scale=7.5, # 7.5
+    steps=10,
+    guidance_scale=7.5,
     seed=1234,
-    octree_resolution=128, # 256
-    num_chunks=50000, # 200000
-    target_face_num=2500, # 10000
+    octree_resolution=128,
+    num_chunks=50000,
+    target_face_num=2500,
     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")
-
+    progress(0, desc="Starting")
 
     if image is None:
-        error_info = {
-            "error": "Please provide either a caption or an image.",
-            "status": "failed",
-        }
-        return None,None,None,None,error_info
-
-    rgbImage = image.convert('RGB')
-    
-    # NSFW 检测
-    if nsfw_model and nsfw_processor:
-        if detect_nsfw(rgbImage):
-            error_info = {
-                "error": "The input image contains NSFW content and cannot be used. Please provide a different image and try again.",
-                "status": "failed",
-            }
-            return None,None,None,None,error_info
+        return None, None, None, None, {"error": "Please provide an image.", "status": "failed"}
+
+    rgb_image = image.convert('RGB')
+    if detect_nsfw(rgb_image):
+        return None, None, None, None, {"error": "NSFW content detected.", "status": "failed"}
 
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    octree_resolution = int(octree_resolution)
     save_folder = gen_save_folder()
-    # 先移除背景
-    image = rmbg_worker(rgbImage)
+    image = rmbg_worker(rgb_image)
+
+    generator = torch.Generator(device=args.device).manual_seed(seed)
 
-    # 生成模型
-    generator = torch.Generator()
-    generator = generator.manual_seed(int(seed))
     outputs = i23d_worker(
         image=image,
         num_inference_steps=steps,
@@ -206,94 +167,43 @@ def _gen_shape_on_gpu(
         octree_resolution=octree_resolution,
         num_chunks=num_chunks,
         output_type='mesh',
-        callback=callback,
+        callback=lambda step_idx, timestep, out: progress(((step_idx+1)/steps)*0.5, desc=f"Mesh generating {step_idx+1}/{steps}"),
         callback_steps=1
     )
 
     mesh = export_to_trimesh(outputs)[0]
-
     path = export_mesh(mesh, save_folder, textured=False)
 
-    # model_viewer_html = build_model_viewer_html(save_folder, height=HTML_HEIGHT, width=HTML_WIDTH)
-
-    # return model_viewer_html, path
-
-    if args.low_vram_mode:
-        torch.cuda.empty_cache()
-    
-    if path is None:
-        error_info = {
-            "error": "'Please generate a mesh first.'",
-            "status": "failed",
-        }
-        return None,None,None,None,error_info
-
-    # 简化模型
-    print(f'exporting {path}')
-    print(f'reduce face to {target_face_num}')
-
     mesh = trimesh.load(path)
-    progress(0.5,desc="Optimizing mesh")
-
+    progress(0.5, desc="Optimizing mesh")
     mesh = floater_remove_worker(mesh)
     mesh = degenerate_face_remove_worker(mesh)
-    progress(0.6,desc="Reducing mesh faces")
+    progress(0.6, desc="Reducing faces")
     mesh = face_reduce_worker(mesh, target_face_num)
-    save_folder = gen_save_folder()
 
-    progress(0.9,desc="Converting format")
-    file_type = "obj"
-    sourceObjPath = export_mesh(mesh, save_folder, textured=False, type=file_type)
-    rel_objPath = os.path.relpath(sourceObjPath, SAVE_DIR)
-    objPath = "/static/"+rel_objPath
-
-    # for preview
     save_folder = gen_save_folder()
-    _ = export_mesh(mesh, save_folder, textured=False)
+    source_obj_path = export_mesh(mesh, save_folder, textured=False, type="obj")
     model_viewer_html = build_model_viewer_html(save_folder, height=HTML_HEIGHT, width=HTML_WIDTH, textured=False)
 
+    glb_path = export_mesh(mesh, save_folder, textured=False, type="glb")
+    rel_glb_path = os.path.relpath(glb_path, SAVE_DIR)
+    glb_path = "/static/" + rel_glb_path
+    rel_obj_path = os.path.relpath(source_obj_path, SAVE_DIR)
+    obj_path = "/static/" + rel_obj_path
 
-    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")
-    info = {
-        "status": "success"
-    }
-    return model_viewer_html, gr.update(value=sourceObjPath, interactive=True), glbPath, objPath, info
+    progress(1, desc="Complete")
+    return model_viewer_html, gr.update(value=source_obj_path, interactive=True), glb_path, obj_path, {"status": "success"}
 
-
-def gen_shape(
-    image=None,
-    steps=50,
-    guidance_scale=7.5,
-    seed=1234,
-    octree_resolution=256,
-    num_chunks=50000, # 2000000
-    target_face_num=2500, # 10000
-    randomize_seed: bool = False,
-):
-    # 调用 GPU 函数
-    html_export_mesh,file_export,glbPath_output,objPath_output, info = _gen_shape_on_gpu(
-        image, 
-        steps,
-        guidance_scale, 
-        seed,
-        octree_resolution,
-        num_chunks, 
-        target_face_num,
-        randomize_seed
-    )
-    # 如果出错，抛出异常
+def gen_shape(*args, **kwargs):
+    html, file_export, glb_path, obj_path, info = _gen_shape(*args, **kwargs)
     if info["status"] == "failed":
         raise gr.Error(info["error"])
-    return html_export_mesh, file_export, glbPath_output, objPath_output
-    
+    return html, file_export, glb_path, obj_path
+
+# -------------------- Beispielbilder --------------------
 def get_example_img_list():
-    print('Loading example img list ...')
     return sorted(glob('./assets/example_images/**/*.png', recursive=True))
+
 example_imgs = get_example_img_list()
 
 HTML_OUTPUT_PLACEHOLDER = f"""
@@ -303,72 +213,63 @@ HTML_OUTPUT_PLACEHOLDER = f"""
     </div>
 </div>
 """
-MAX_SEED = 1e7
 
+# -------------------- Gradio UI --------------------
 title = "## AI 3D Model Generator"
-description = "Our Image-to-3D Generator transforms your 2D photos into stunning, AI generated 3D models—ready for games, AR/VR, or 3D printing. Our AI 3D Modeling is based on Hunyuan 2.0. Check more in [imgto3d.ai](https://www.imgto3d.ai)."
+description = "Transforms 2D photos into AI-generated 3D models."
 
 with gr.Blocks().queue() as demo:
     gr.Markdown(title)
     gr.Markdown(description)
     with gr.Row():
         with gr.Column(scale=3):
-            gr.Markdown("#### Image Prompt")
             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,
-                        min_width=100,
-                    )
-                    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=128, 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=50000, minimum=1000, value=2000, label='Number of Chunks') # old maximum=5000000
-                target_face_num = gr.Slider(maximum=1000000, minimum=100, value=2500, label='Target Face Number') # old maximum=1000000
-
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=1234)
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                num_steps = gr.Slider(maximum=100, minimum=1, value=5, step=1, label='Inference Steps')
+                octree_resolution = gr.Slider(maximum=512, minimum=16, value=128, label='Octree Resolution')
+                cfg_scale = gr.Slider(maximum=20.0, minimum=1.0, value=5.5, step=0.1, label='Guidance Scale')
+                num_chunks = gr.Slider(maximum=50000, minimum=1000, value=2000, label='Number of Chunks')
+                target_face_num = gr.Slider(maximum=1000000, minimum=100, value=2500, label='Target Face Number')
         with gr.Column(scale=6):
-            gr.Markdown("#### Generated Mesh")
             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)
-                
+            objPath_output = gr.Text(label="Obj Path", interactive=False)
+            glbPath_output = gr.Text(label="Glb Path", interactive=False)
         with gr.Column(scale=3):
-            gr.Markdown("#### Image Examples")
-            gr.Examples(examples=example_imgs, inputs=[image],
-                        label=None, examples_per_page=18)
-            
+            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]
-    )    
+        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]
+    )
 
+# -------------------- FastAPI + Gradio --------------------
 if __name__ == "__main__":
-    # https://discuss.huggingface.co/t/how-to-serve-an-html-file/33921/2
-    # create a FastAPI app
+    # Device Info
+    print(f"Using device: {args.device}")
+
+    # Optional: FastAPI static files (für Assets)
     app = FastAPI()
-    # create a static directory to store the static files
     static_dir = Path(SAVE_DIR).absolute()
     static_dir.mkdir(parents=True, exist_ok=True)
     app.mount("/static", StaticFiles(directory=static_dir, html=True), name="static")
     shutil.copytree('./assets/env_maps', os.path.join(static_dir, 'env_maps'), dirs_exist_ok=True)
 
-    if args.low_vram_mode:
+    # Low VRAM cleanup
+    if args.low_vram_mode and args.device == "cuda":
         torch.cuda.empty_cache()
 
+    # Gradio Demo starten CPU-kompatibel, funktioniert auch in HF Spaces
+    demo.launch(
+        server_name="0.0.0.0",  # für Spaces oder lokal
+        server_port=args.port,
+        share=True            # erstellt einen öffentlichen Link wie HF Spaces
+    )
     app = gr.mount_gradio_app(app, demo, path="/")
-    # demo.launch()
-    from spaces import zero
-    zero.startup()
+    # from spaces import zero
+    # zero.startup()
     uvicorn.run(app, host=args.host, port=args.port)