Spaces:
Sleeping
Sleeping
| import os | |
| import random | |
| import shutil | |
| import uuid | |
| from glob import glob | |
| from pathlib import Path | |
| 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 | |
| # -------------------- 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=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') | |
| parser.add_argument('--disable_tex', action='store_true') | |
| 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() | |
| # -------------------- 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 Modell -------------------- | |
| nsfw_processor = AutoProcessor.from_pretrained("Falconsai/nsfw_image_detection") | |
| 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()] | |
| 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}") | |
| 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'): | |
| 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): | |
| 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: | |
| 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>" | |
| # -------------------- Hy3Dgen Worker -------------------- | |
| 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 args.mc_algo | |
| 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() | |
| # -------------------- NSFW Detection -------------------- | |
| def detect_nsfw(image: Image.Image, threshold: float = 0.5) -> bool: | |
| nsfw_score = 0 # Placeholder, optional: implement actual detection | |
| return nsfw_score > threshold | |
| # -------------------- Mesh Generation -------------------- | |
| progress = gr.Progress() | |
| def _gen_shape( | |
| image=None, | |
| steps=10, | |
| guidance_scale=7.5, | |
| seed=1234, | |
| octree_resolution=128, | |
| num_chunks=50000, | |
| target_face_num=2500, | |
| randomize_seed: bool = False, | |
| ): | |
| progress(0, desc="Starting") | |
| if image is None: | |
| 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)) | |
| save_folder = gen_save_folder() | |
| image = rmbg_worker(rgb_image) | |
| if args.device in ["cuda", "cpu"]: | |
| generator = torch.Generator(device=args.device).manual_seed(seed) | |
| else: | |
| generator = torch.Generator().manual_seed(seed) # fallback für MPS | |
| 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=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) | |
| 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 faces") | |
| mesh = face_reduce_worker(mesh, target_face_num) | |
| save_folder = gen_save_folder() | |
| 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 | |
| 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(*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, file_export, glb_path, obj_path | |
| # -------------------- Beispielbilder -------------------- | |
| def get_example_img_list(): | |
| return sorted(glob('./assets/example_images/**/*.png', recursive=True)) | |
| example_imgs = get_example_img_list() | |
| HTML_OUTPUT_PLACEHOLDER = f""" | |
| <div style='height: {500}px; width: 100%; border-radius: 8px; border-color: #e5e7eb; border-style: solid; border-width: 1px; 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> | |
| """ | |
| # -------------------- Gradio UI -------------------- | |
| title = "## AI 3D Model Generator" | |
| 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): | |
| 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): | |
| 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): | |
| html_export_mesh = gr.HTML(HTML_OUTPUT_PLACEHOLDER, label='Output') | |
| file_export = gr.DownloadButton(label="Download", variant='primary', 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.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] | |
| ) | |
| # -------------------- FastAPI + Gradio -------------------- | |
| if __name__ == "__main__": | |
| # Device Info | |
| print(f"Using device: {args.device}") | |
| # Optional: FastAPI static files (für Assets) | |
| app = FastAPI() | |
| 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) | |
| # 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=False # erstellt einen öffentlichen Link wie HF Spaces | |
| ) | |
| app = gr.mount_gradio_app(app, demo, path="/") | |
| # from spaces import zero | |
| # zero.startup() | |
| uvicorn.run(app, host=args.host, port=args.port) |