Upload 2 files

app.py

@@ -1,200 +1,245 @@
-import gradio as gr
-import os
-import shutil
-import tempfile
-import datetime
-import numpy as np
-import torch
-import imageio
-import trimesh
-
-from PIL import Image
-from typing import *
-from gradio_litmodel3d import LitModel3D
-from trellis.pipelines import TrellisImageTo3DPipeline
-from trellis.utils import render_utils
-
-os.environ['SPCONV_ALGO'] = 'native'
-
-MAX_SEED = np.iinfo(np.int32).max
-TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
-os.makedirs(TMP_DIR, exist_ok=True)
-
-def preprocess_mesh(mesh_prompt):
-    print("Processing mesh")
-    trimesh_mesh = trimesh.load_mesh(mesh_prompt)
-    trimesh_mesh.export(mesh_prompt+'.glb')
-    return mesh_prompt+'.glb'
-
-def preprocess_image(image):
-    if image is None:
-        return None
-    image = pipeline.preprocess_image(image, resolution=1024)
-    return image
-
-def generate_3d(image, seed=-1,  
-                ss_guidance_strength=3, ss_sampling_steps=50,
-                slat_guidance_strength=3, slat_sampling_steps=6,):
-    if image is None:
-        return None, None, None, None
-
-    if seed == -1:
-        seed = np.random.randint(0, MAX_SEED)
-    
-    image = pipeline.preprocess_image(image, resolution=1024)
-    normal_image = normal_predictor(image, resolution=768, match_input_resolution=True, data_type='object')
-
-    outputs = pipeline.run(
-        normal_image,
-        seed=seed,
-        formats=["mesh",],
-        preprocess_image=False,
-        sparse_structure_sampler_params={
-            "steps": ss_sampling_steps,
-            "cfg_strength": ss_guidance_strength,
-        },
-        slat_sampler_params={
-            "steps": slat_sampling_steps,
-            "cfg_strength": slat_guidance_strength,
-        },
-    )
-    generated_mesh = outputs['mesh'][0]
-    
-    output_id = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
-    os.makedirs(os.path.join(TMP_DIR, output_id), exist_ok=True)
-    mesh_path = f"{TMP_DIR}/{output_id}/mesh.glb"
-    
-    render_results = render_utils.render_video(generated_mesh, resolution=1024, ssaa=1, num_frames=8, pitch=0.25, inverse_direction=True)
-    
-    def combine_diagonal(color_np, normal_np):
-        h, w, c = color_np.shape
-        mask = np.fromfunction(lambda y, x: x > y, (h, w)).astype(bool)
-        mask = np.stack([mask] * c, axis=-1)
-        combined_np = np.where(mask, color_np, normal_np)
-        return Image.fromarray(combined_np)
-
-    preview_images = [combine_diagonal(c, n) for c, n in zip(render_results['color'], render_results['normal'])]
-    
-    trimesh_mesh = generated_mesh.to_trimesh(transform_pose=True)
-    trimesh_mesh.export(mesh_path)
-
-    return preview_images, normal_image, mesh_path, mesh_path
-
-def convert_mesh(mesh_path, export_format):
-    if not mesh_path:
-        return None
-    temp_file = tempfile.NamedTemporaryFile(suffix=f".{export_format}", delete=False)
-    mesh = trimesh.load_mesh(mesh_path)
-    mesh.export(temp_file.name)
-    return temp_file.name
-
-with gr.Blocks(css="footer {visibility: hidden}") as demo:
-    gr.Markdown("""
-        <h1 style='text-align: center;'>Hi3DGen: High-fidelity 3D Geometry Generation from Images via Normal Bridging</h1>
-        <p style='text-align: center;'>
-            <strong>V0.1, Introduced By 
-            <a href="https://gaplab.cuhk.edu.cn/" target="_blank">GAP Lab</a> from CUHKSZ and 
-            <a href="https://www.nvsgames.cn/" target="_blank">Game-AIGC Team</a> from ByteDance</strong>
-        </p>
-    """)
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            with gr.Tabs():
-                with gr.Tab("Single Image"):
-                    with gr.Row():
-                        image_prompt = gr.Image(label="Image Prompt", image_mode="RGBA", type="pil")
-                        normal_output = gr.Image(label="Normal Bridge", image_mode="RGBA", type="pil")
-                with gr.Tab("Multiple Images"):
-                    gr.Markdown("<div style='text-align: center; padding: 40px; font-size: 24px;'>Multiple Images functionality is coming soon!</div>")
-            with gr.Accordion("Advanced Settings", open=False):
-                seed = gr.Slider(-1, MAX_SEED, label="Seed", value=0, step=1)
-                gr.Markdown("#### Stage 1: Sparse Structure Generation")
-                with gr.Row():
-                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3, step=0.1)
-                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=50, step=1)
-                gr.Markdown("#### Stage 2: Structured Latent Generation")
-                with gr.Row():
-                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
-                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=6, step=1)
-            with gr.Group():
-                with gr.Row():
-                    gen_shape_btn = gr.Button("Generate Shape", size="lg", variant="primary")
-
-        with gr.Column(scale=1):
-            with gr.Tabs():
-                with gr.Tab("Preview"):
-                    output_gallery = gr.Gallery(label="Examples", columns=4, rows=2, object_fit="contain", height="auto", show_label=False)
-                with gr.Tab("3D Model"):
-                    with gr.Column():
-                        model_output = gr.Model3D(label="3D Model Preview (Each model is approx. 40MB)")
-                    with gr.Column():
-                        export_format = gr.Dropdown(
-                            choices=["obj", "glb", "ply", "stl"],
-                            value="glb",
-                            label="File Format"
-                        )
-                        download_btn = gr.DownloadButton(label="Export Mesh", interactive=False)
-
-    image_prompt.upload(
-        preprocess_image,
-        inputs=[image_prompt],
-        outputs=[image_prompt]
-    )
-
-    gen_shape_btn.click(
-        generate_3d,
-        inputs=[
-            image_prompt, seed,
-            ss_guidance_strength, ss_sampling_steps,
-            slat_guidance_strength, slat_sampling_steps
-        ],
-        outputs=[output_gallery, normal_output, model_output, download_btn]
-    ).then(
-        lambda: gr.Button(interactive=True),
-        outputs=[download_btn],
-    )
-
-    def update_download_button(mesh_path, export_format):
-        if not mesh_path:
-            return gr.File.update(value=None, interactive=False)
-        download_path = convert_mesh(mesh_path, export_format)
-        return download_path
-
-    export_format.change(
-        update_download_button,
-        inputs=[model_output, export_format],
-        outputs=[download_btn]
-    ).then(
-        lambda: gr.Button(interactive=True),
-        outputs=[download_btn],
-    )
-
-    examples = gr.Examples(
-        examples=[
-            f'assets/example_image/{image}'
-            for image in os.listdir("assets/example_image")
-        ],
-        inputs=image_prompt,
-    )
-
-    gr.Markdown("""
-    **Acknowledgments**: Hi3DGen is built on the shoulders of giants. We acknowledge contributions from:
-    - [Trellis 3D](https://github.com/microsoft/TRELLIS)
-    - [StableNormal](https://github.com/hugoycj/StableNormal)
-    """)
-
-if __name__ == "__main__":
-    # ✅ 强制使用 CPU
-    pipeline = TrellisImageTo3DPipeline.from_pretrained("Stable-X/trellis-normal-v0-1")
-    pipeline.to("cpu")  # <-- 强制使用 CPU
-
-    normal_predictor = torch.hub.load(
-        "hugoycj/StableNormal",
-        "StableNormal_turbo",
-        trust_repo=True,
-        yoso_version="yoso-normal-v1-8-1"
-    )
-    normal_predictor.to("cpu")  # <-- 也强制使用 CPU
-
-    demo.launch()
+import gradio as gr
+import spaces
+from gradio_litmodel3d import LitModel3D
+
+import os
+import shutil
+os.environ['SPCONV_ALGO'] = 'native'
+from typing import *
+import torch
+import numpy as np
+import imageio
+from PIL import Image
+from trellis.pipelines import TrellisImageTo3DPipeline
+from trellis.utils import render_utils
+import trimesh
+import tempfile
+
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+os.makedirs(TMP_DIR, exist_ok=True)
+
+def preprocess_mesh(mesh_prompt):
+    print("Processing mesh")
+    trimesh_mesh = trimesh.load_mesh(mesh_prompt)
+    trimesh_mesh.export(mesh_prompt+'.glb')
+    return mesh_prompt+'.glb'
+
+def preprocess_image(image):
+    if image is None:
+        return None
+    image = pipeline.preprocess_image(image, resolution=1024)
+    return image
+
+# Removed @spaces.GPU decorator to allow CPU execution
+def generate_3d(image, seed=-1,  
+                ss_guidance_strength=3, ss_sampling_steps=50,
+                slat_guidance_strength=3, slat_sampling_steps=6,):
+    if image is None:
+        return None, None, None
+
+    if seed == -1:
+        seed = np.random.randint(0, MAX_SEED)
+    
+    image = pipeline.preprocess_image(image, resolution=1024)
+    normal_image = normal_predictor(image, resolution=768, match_input_resolution=True, data_type='object')
+
+    outputs = pipeline.run(
+        normal_image,
+        seed=seed,
+        formats=["mesh",],
+        preprocess_image=False,
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps,
+            "cfg_strength": ss_guidance_strength,
+        },
+        slat_sampler_params={
+            "steps": slat_sampling_steps,
+            "cfg_strength": slat_guidance_strength,
+        },
+    )
+    generated_mesh = outputs['mesh'][0]
+    
+    # Save outputs
+    import datetime
+    output_id = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+    os.makedirs(os.path.join(TMP_DIR, output_id), exist_ok=True)
+    mesh_path = f"{TMP_DIR}/{output_id}/mesh.glb"
+    
+    render_results = render_utils.render_video(generated_mesh, resolution=1024, ssaa=1, num_frames=8, pitch=0.25, inverse_direction=True)
+    def combine_diagonal(color_np, normal_np):
+        # Convert images to numpy arrays
+        h, w, c = color_np.shape
+        # Create a boolean mask that is True for pixels where x > y (diagonally)
+        mask = np.fromfunction(lambda y, x: x > y, (h, w))
+        mask = mask.astype(bool)
+        mask = np.stack([mask] * c, axis=-1)
+        # Where mask is True take color, else normal
+        combined_np = np.where(mask, color_np, normal_np)
+        return Image.fromarray(combined_np)
+
+    preview_images = [combine_diagonal(c, n) for c, n in zip(render_results['color'], render_results['normal'])]
+    
+    # Export mesh
+    trimesh_mesh = generated_mesh.to_trimesh(transform_pose=True)
+
+    trimesh_mesh.export(mesh_path)
+
+    return preview_images, normal_image, mesh_path, mesh_path
+
+def convert_mesh(mesh_path, export_format):
+    """Download the mesh in the selected format."""
+    if not mesh_path:
+        return None
+    
+    # Create a temporary file to store the mesh data
+    temp_file = tempfile.NamedTemporaryFile(suffix=f".{export_format}", delete=False)
+    temp_file_path = temp_file.name
+    
+    new_mesh_path = mesh_path.replace(".glb", f".{export_format}")
+    mesh = trimesh.load_mesh(mesh_path)
+    mesh.export(temp_file_path)  # Export to the temporary file
+    
+    return temp_file_path # Return the path to the temporary file
+
+# Create the Gradio interface with improved layout
+with gr.Blocks(css="footer {visibility: hidden}") as demo:
+    gr.Markdown(
+        """
+        <h1 style='text-align: center;'>Hi3DGen: High-fidelity 3D Geometry Generation from Images via Normal Bridging</h1>
+        <p style='text-align: center;'>
+            <strong>V0.1, Introduced By 
+            <a href="https://gaplab.cuhk.edu.cn/" target="_blank">GAP Lab</a> from CUHKSZ and 
+            <a href="https://www.nvsgames.cn/" target="_blank">Game-AIGC Team</a> from ByteDance</strong>
+        </p>
+        """
+    )
+    
+    with gr.Row():
+        gr.Markdown("""
+                    <p align="center">
+                    <a title="Website" href="https://stable-x.github.io/Hi3DGen/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                        <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
+                    </a>
+                    <a title="arXiv" href="https://stable-x.github.io/Hi3DGen/hi3dgen_paper.pdf" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                        <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
+                    </a>
+                    <a title="Github" href="https://github.com/Stable-X/Hi3DGen" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                        <img src="https://img.shields.io/github/stars/Stable-X/Hi3DGen?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
+                    </a>
+                    <a title="Social" href="https://x.com/ychngji6" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                        <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
+                    </a>
+                    </p>
+                    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            with gr.Tabs():
+                
+                with gr.Tab("Single Image"):
+                    with gr.Row():
+                        image_prompt = gr.Image(label="Image Prompt", image_mode="RGBA", type="pil")
+                        normal_output = gr.Image(label="Normal Bridge", image_mode="RGBA", type="pil")
+                        
+                with gr.Tab("Multiple Images"):
+                    gr.Markdown("<div style='text-align: center; padding: 40px; font-size: 24px;'>Multiple Images functionality is coming soon!</div>")
+                        
+            with gr.Accordion("Advanced Settings", open=False):
+                seed = gr.Slider(-1, MAX_SEED, label="Seed", value=0, step=1)
+                gr.Markdown("#### Stage 1: Sparse Structure Generation")
+                with gr.Row():
+                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3, step=0.1)
+                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=50, step=1)
+                gr.Markdown("#### Stage 2: Structured Latent Generation")
+                with gr.Row():
+                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
+                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=6, step=1)
+                    
+            with gr.Group():
+                with gr.Row():
+                    gen_shape_btn = gr.Button("Generate Shape", size="lg", variant="primary")
+                        
+        # Right column - Output
+        with gr.Column(scale=1):
+            with gr.Tabs():
+                with gr.Tab("Preview"):
+                    output_gallery = gr.Gallery(label="Examples", columns=4, rows=2, object_fit="contain", height="auto",show_label=False)
+                with gr.Tab("3D Model"):
+                    with gr.Column():
+                        model_output = gr.Model3D(label="3D Model Preview (Each model is approximately 40MB, may take around 1 minute to load)")
+                    with gr.Column():
+                        export_format = gr.Dropdown(
+                            choices=["obj", "glb", "ply", "stl"],
+                            value="glb",
+                            label="File Format"
+                        )
+                        download_btn = gr.DownloadButton(label="Export Mesh", interactive=False)
+    
+    image_prompt.upload(
+        preprocess_image,
+        inputs=[image_prompt],
+        outputs=[image_prompt]
+    )
+    
+    gen_shape_btn.click(
+        generate_3d,
+        inputs=[
+            image_prompt, seed,  
+            ss_guidance_strength, ss_sampling_steps,
+            slat_guidance_strength, slat_sampling_steps
+        ],
+        outputs=[output_gallery, normal_output, model_output, download_btn]
+    ).then(
+        lambda: gr.Button(interactive=True),
+        outputs=[download_btn],
+    )
+    
+    
+    def update_download_button(mesh_path, export_format):
+        if not mesh_path:
+            return gr.File.update(value=None, interactive=False)
+        
+        download_path = convert_mesh(mesh_path, export_format)
+        return download_path
+    
+    export_format.change(
+        update_download_button,
+        inputs=[model_output, export_format],
+        outputs=[download_btn]
+    ).then(
+        lambda: gr.Button(interactive=True),
+        outputs=[download_btn],
+    )
+    
+    examples = gr.Examples(
+        examples=[
+            f'assets/example_image/{image}'
+            for image in os.listdir("assets/example_image")
+        ],
+        inputs=image_prompt,
+    )
+
+    gr.Markdown(
+        """
+        **Acknowledgments**: Hi3DGen is built on the shoulders of giants. We would like to express our gratitude to the open-source research community and the developers of these pioneering projects:
+        - **3D Modeling:** Our 3D Model is finetuned from the SOTA open-source 3D foundation model [Trellis](https://github.com/microsoft/TRELLIS) and we draw inspiration from the teams behind [Rodin](https://hyperhuman.deemos.com/rodin), [Tripo](https://www.tripo3d.ai/app/home), and [Dora](https://github.com/Seed3D/Dora).
+        - **Normal Estimation:** Our Normal Estimation Model builds on the leading normal estimation research such as [StableNormal](https://github.com/hugoycj/StableNormal) and [GenPercept](https://github.com/aim-uofa/GenPercept).
+        
+        **Your contributions and collaboration push the boundaries of 3D modeling!**
+        """
+    )
+
+if __name__ == "__main__":
+    # Initialize pipeline
+    pipeline = TrellisImageTo3DPipeline.from_pretrained("Stable-X/trellis-normal-v0-1")
+    # Use CPU instead of GPU
+    pipeline.to("cpu")
+
+    # Initialize normal predictor
+    normal_predictor = torch.hub.load("hugoycj/StableNormal", "StableNormal_turbo", trust_repo=True, yoso_version='yoso-normal-v1-8-1')
+    # Ensure normal predictor is on CPU
+    normal_predictor.to("cpu")
+    
+    # Launch the app
+    demo.launch()

requirements.txt

@@ -1,31 +1,30 @@
---extra-index-url https://download.pytorch.org/whl/cu121
-huggingface_hub==0.25.0
-diffusers==0.28.0
-accelerate==1.2.1
-kornia==0.8.0
-timm==0.6.7
-
-torch==2.4.0
-torchvision==0.19.0
-pillow==10.4.0
-imageio==2.36.1
-imageio-ffmpeg==0.5.1
-tqdm==4.67.1
-easydict==1.13
-opencv-python-headless==4.10.0.84
-scipy==1.14.1
-rembg==2.0.60
-onnxruntime==1.20.1
-trimesh==4.5.3
-xatlas==0.0.9
-pyvista==0.44.2
-pymeshfix==0.17.0
-igraph==0.11.8
-git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
-xformers==0.0.27.post2
-spconv-cu120==2.3.6
-transformers==4.46.3
-gradio_litmodel3d==0.0.1
-https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.0.post2/flash_attn-2.7.0.post2+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
-https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl?download=true
-https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/nvdiffrast-0.3.3-cp310-cp310-linux_x86_64.whl?download=true
+huggingface_hub==0.25.0
+diffusers==0.28.0
+accelerate==1.2.1
+kornia==0.8.0
+timm==0.6.7
+
+# CPU versions of PyTorch packages
+torch==2.4.0+cpu
+torchvision==0.19.0+cpu
+--extra-index-url https://download.pytorch.org/whl/cpu
+
+pillow==10.4.0
+imageio==2.36.1
+imageio-ffmpeg==0.5.1
+tqdm==4.67.1
+easydict==1.13
+opencv-python-headless==4.10.0.84
+scipy==1.14.1
+rembg==2.0.60
+onnxruntime==1.20.1
+trimesh==4.5.3
+xatlas==0.0.9
+pyvista==0.44.2
+pymeshfix==0.17.0
+igraph==0.11.8
+git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
+
+# Remove GPU-specific packages
+transformers==4.46.3
+gradio_litmodel3d==0.0.1