init

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: HoloPart
+emoji: 🔮
+colorFrom: indigo
+colorTo: indigo
+sdk: gradio
+sdk_version: 5.24.0
+app_file: app.py
+pinned: false
+license: mit
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,243 @@
+import spaces
+import os
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+import trimesh
+import random
+from transformers import AutoModelForImageSegmentation
+from torchvision import transforms
+from huggingface_hub import hf_hub_download, snapshot_download, login
+import subprocess
+import shutil
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.float16
+
+print("DEVICE: ", DEVICE)
+
+DEFAULT_PART_FACE_NUMBER = 10000
+MAX_SEED = np.iinfo(np.int32).max
+HOLOPART_REPO_URL = "https://github.com/VAST-AI-Research/HoloPart"
+HOLOPART_PRETRAINED_MODEL = "checkpoints/HoloPart"
+
+TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
+os.makedirs(TMP_DIR, exist_ok=True)
+
+HOLOPART_CODE_DIR = "./holopart"
+if not os.path.exists(HOLOPART_REPO_URL):
+    os.system(f"git clone {HOLOPART_REPO_URL} {HOLOPART_CODE_DIR}")
+
+import sys
+sys.path.append(HOLOPART_CODE_DIR)
+sys.path.append(os.path.join(HOLOPART_CODE_DIR, "scripts"))
+
+EXAMPLES = [
+    ["./holopart/assets/example_data/000.glb", "./holopart/assets/example_data/000.png"],
+    ["./holopart/assets/example_data/001.glb", "./holopart/assets/example_data/001.png"],
+    ["./holopart/assets/example_data/002.glb", "./holopart/assets/example_data/002.png"],
+    ["./holopart/assets/example_data/003.glb", "./holopart/assets/example_data/003.png"],
+    ["./holopart/assets/example_data/004.glb", "./holopart/assets/example_data/004.png"],
+]
+
+HEADER = """
+# 🔮 Decompose a 3D shape into complete parts with [HoloPart](https://github.com/VAST-AI-Research/HoloPart).
+### Step 1: Prepare Your Segmented Mesh
+Upload a mesh with part segmentation. We recommend using these segmentation tools:
+- [SAMPart3D](https://github.com/Pointcept/SAMPart3D)
+- [SAMesh](https://github.com/gtangg12/samesh)
+For a mesh file `mesh.glb` and corresponding face mask `mask.npy`, prepare your input using this Python code:
+```python
+import trimesh
+import numpy as np
+mesh = trimesh.load("mesh.glb", force="mesh")
+mesh_parts = []
+for part_id in np.unique(mask_npy):
+    mesh_part = mesh.submesh([mask_npy == part_id], append=True)
+    mesh_parts.append(mesh_part)
+mesh_parts = trimesh.Scene(mesh_parts).export(input_mesh.glb)
+```
+The resulting **input_mesh.glb** is your prepared input for HoloPart.
+### Step 2: Click the Decompose Parts button to begin the decomposition process.
+"""
+
+from inference_holopart import prepare_data, run_holopart
+from holopart.pipelines.pipeline_holopart import HoloPartPipeline
+
+snapshot_download("VAST-AI/HoloPart", local_dir=HOLOPART_PRETRAINED_MODEL)
+holopart_pipe = HoloPartPipeline.from_pretrained(HOLOPART_PRETRAINED_MODEL).to(DEVICE, DTYPE)
+
+def start_session(req: gr.Request):
+    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(save_dir, exist_ok=True)
+    print("start session, mkdir", save_dir)
+
+def end_session(req: gr.Request):
+    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    shutil.rmtree(save_dir)
+
+def get_random_hex():
+    random_bytes = os.urandom(8)
+    random_hex = random_bytes.hex()
+    return random_hex
+
+def get_random_seed(randomize_seed, seed):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+def explode_mesh(mesh: trimesh.Scene, explode_factor: float = 0.5):
+    center = mesh.centroid
+    exploded_mesh = trimesh.Scene()
+    for geometry_name, geometry in mesh.geometry.items():
+        transform = mesh.graph[geometry_name][0]
+        vertices_global = trimesh.transformations.transform_points(
+            geometry.vertices, transform)
+        part_center = np.mean(vertices_global, axis=0)
+        direction = part_center - center
+        direction_length = np.linalg.norm(direction)
+        if direction_length > 0:
+            direction = direction / direction_length
+        displacement = direction * explode_factor
+        new_transform = np.copy(transform)
+        new_transform[:3, 3] += displacement
+        exploded_mesh.add_geometry(geometry, transform=new_transform, geom_name=geometry_name)
+    return exploded_mesh
+
+
+@spaces.GPU(duration=600)
+def run_full(data_path, seed=42, num_inference_steps=25, guidance_scale=3.5):
+
+    batch_size = 30
+    parts_data = prepare_data(data_path)
+
+    part_scene = run_holopart(
+        holopart_pipe,
+        batch=parts_data,
+        batch_size=batch_size,
+        seed=seed,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        num_chunks=1000000,
+    )
+    print("mesh extraction done")
+    
+    save_dir = os.path.join(TMP_DIR, "examples")
+    os.makedirs(save_dir, exist_ok=True)
+    mesh_path = os.path.join(save_dir, f"holorpart_{get_random_hex()}.glb")
+    part_scene.export(mesh_path)
+    print("save to ", mesh_path)
+    exploded_mesh = explode_mesh(part_scene, 0.7)
+    exploded_mesh_path = os.path.join(save_dir, f"holorpart_exploded_{get_random_hex()}.glb")
+    exploded_mesh.export(exploded_mesh_path)
+
+    torch.cuda.empty_cache()
+
+    return mesh_path, exploded_mesh_path
+
+
+@spaces.GPU(duration=600)
+def run_example(data_path: str, example_image_path, seed=42, num_inference_steps=25, guidance_scale=3.5):
+
+    batch_size = 30
+    parts_data = prepare_data(data_path)
+
+    part_scene = run_holopart(
+        holopart_pipe,
+        batch=parts_data,
+        batch_size=batch_size,
+        seed=seed,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        num_chunks=1000000,
+    )
+    print("mesh extraction done")
+
+    
+    save_dir = os.path.join(TMP_DIR, "examples")
+    os.makedirs(save_dir, exist_ok=True)
+    mesh_path = os.path.join(save_dir, f"holorpart_{get_random_hex()}.glb")
+    part_scene.export(mesh_path)
+    print("save to ", mesh_path)
+    exploded_mesh = explode_mesh(part_scene, 0.5)
+    exploded_mesh_path = os.path.join(save_dir, f"holorpart_exploded_{get_random_hex()}.glb")
+    exploded_mesh.export(exploded_mesh_path)
+
+    torch.cuda.empty_cache()
+
+    return mesh_path, exploded_mesh_path
+
+
+with gr.Blocks(title="HoloPart") as demo:
+    gr.Markdown(HEADER)
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                input_mesh = gr.Model3D(label="Input Mesh")
+                example_image = gr.Image(label="Example Image", type="filepath", interactive=False, visible=False)
+                # seg_image = gr.Image(
+                #     label="Segmentation Result", type="pil", format="png", interactive=False
+                # )
+                
+            with gr.Accordion("Generation Settings", open=True):
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=0,
+                    value=0
+                )
+                # randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=8,
+                    maximum=50,
+                    step=1,
+                    value=25,
+                )
+                guidance_scale = gr.Slider(
+                    label="CFG scale",
+                    minimum=0.0,
+                    maximum=20.0,
+                    step=0.1,
+                    value=3.5,
+                )
+
+                with gr.Row():
+                    reduce_face = gr.Checkbox(label="Simplify Mesh", value=True, interactive=False)
+                    # target_face_num = gr.Slider(maximum=1000000, minimum=10000, value=DEFAULT_FACE_NUMBER, label="Target Face Number")
+
+                gen_button = gr.Button("Decompose Parts", variant="primary")
+
+        with gr.Column():
+            model_output = gr.Model3D(label="Decomposed GLB", interactive=False)
+            exploded_parts_output = gr.Model3D(label="Exploded Parts", interactive=False)
+
+    with gr.Row():
+        examples = gr.Examples(
+            examples=EXAMPLES,
+            fn=run_example,
+            inputs=[input_mesh, example_image],
+            outputs=[model_output, exploded_parts_output],
+            cache_examples=True,
+        )
+
+
+    gen_button.click(
+        run_full,
+        inputs=[
+            input_mesh, 
+            seed, 
+            num_inference_steps, 
+            guidance_scale
+        ],
+        outputs=[model_output, exploded_parts_output],
+    )
+
+    demo.load(start_session)
+    demo.unload(end_session)
+
+demo.launch()

requirements.txt

@@ -0,0 +1,25 @@
+torchvision
+diffusers
+transformers==4.49.0
+einops
+huggingface_hub
+opencv-python
+trimesh==4.5.3
+omegaconf
+scikit-image
+numpy
+peft
+scipy==1.11.4
+jaxtyping
+typeguard
+pymeshlab==2022.2.post4
+open3d
+timm
+kornia
+ninja
+https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/nvdiffrast-0.3.3-cp310-cp310-linux_x86_64.whl?download=true
+cvcuda_cu12
+gltflib
+https://huggingface.co/spaces/VAST-AI/TripoSG/resolve/main/diso-0.1.4-cp310-cp310-linux_x86_64.whl?download=true
+--find-links https://data.pyg.org/whl/torch-2.6.0+cu124.html
+torch-cluster