paulpanwang/PartCrafter
5
1import spaces2import gradio as gr3import os4import sys5from glob import glob6import time7from typing import Any, Union8 9import numpy as np10import torch11import uuid12import shutil13 14print(f'torch version:{torch.__version__}')15 16import trimesh17import glob18from huggingface_hub import snapshot_download19from PIL import Image20from accelerate.utils import set_seed21 22import subprocess23import importlib, site, sys24 25# Re-discover all .pth/.egg-link files26for sitedir in site.getsitepackages():27 site.addsitedir(sitedir)28 29# Clear caches so importlib will pick up new modules30importlib.invalidate_caches()31 32def sh(cmd): subprocess.check_call(cmd, shell=True)33 34def install_cuda_toolkit():35 CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.6.0/local_installers/cuda_12.6.0_560.28.03_linux.run"36 CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)37 subprocess.check_call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])38 subprocess.check_call(["chmod", "+x", CUDA_TOOLKIT_FILE])39 subprocess.check_call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])40 41 os.environ["CUDA_HOME"] = "/usr/local/cuda"42 os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])43 os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (44 os.environ["CUDA_HOME"],45 "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],46 )47 # add for compiler header lookup48 os.environ["CPATH"] = f"{os.environ['CUDA_HOME']}/include" + (49 f":{os.environ['CPATH']}" if "CPATH" in os.environ else ""50 )51 # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range52 os.environ["TORCH_CUDA_ARCH_LIST"] = "8.9;9.0"53 print("==> finished installation")54 55print("installing cuda toolkit")56install_cuda_toolkit()57print("finished")58 59os.environ["PARTCRAFTER_PROCESSED"] = f"{os.getcwd()}/proprocess_results"60 61 62def sh(cmd_list, extra_env=None):63 env = os.environ.copy()64 if extra_env:65 env.update(extra_env)66 subprocess.check_call(cmd_list, env=env)67 68# install with FORCE_CUDA=169sh(["pip", "install", "diso"], {"FORCE_CUDA": "1"})70# sh(["pip", "install", "torch-cluster", "-f", "https://data.pyg.org/whl/torch-2.7.0+126.html"])71 72 73 74# tell Python to re-scan site-packages now that the egg-link exists75import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches()76 77 78from src.utils.data_utils import get_colored_mesh_composition, scene_to_parts, load_surfaces79from src.utils.render_utils import render_views_around_mesh, render_normal_views_around_mesh, make_grid_for_images_or_videos, export_renderings80from src.pipelines.pipeline_partcrafter import PartCrafterPipeline81from src.utils.image_utils import prepare_image82from src.models.briarmbg import BriaRMBG83 84# Constants85MAX_NUM_PARTS = 1686DEVICE = "cuda" 87DTYPE = torch.float1688 89# Download and initialize models90partcrafter_weights_dir = "pretrained_weights/PartCrafter"91rmbg_weights_dir = "pretrained_weights/RMBG-1.4"92snapshot_download(repo_id="wgsxm/PartCrafter", local_dir=partcrafter_weights_dir)93snapshot_download(repo_id="briaai/RMBG-1.4", local_dir=rmbg_weights_dir)94 95rmbg_net = BriaRMBG.from_pretrained(rmbg_weights_dir).to(DEVICE)96rmbg_net.eval()97pipe: PartCrafterPipeline = PartCrafterPipeline.from_pretrained(partcrafter_weights_dir).to(DEVICE, DTYPE)98 99def first_file_from_dir(directory, ext):100 files = glob.glob(os.path.join(directory, f"*.{ext}"))101 return sorted(files)[0] if files else None102 103 104def explode_mesh(mesh, explosion_scale=0.4): 105 106 if isinstance(mesh, trimesh.Scene):107 scene = mesh108 elif isinstance(mesh, trimesh.Trimesh):109 print("Warning: Single mesh provided, can't create exploded view")110 scene = trimesh.Scene(mesh)111 return scene112 else:113 print(f"Warning: Unexpected mesh type: {type(mesh)}")114 scene = mesh115 116 if len(scene.geometry) <= 1:117 print("Only one geometry found - nothing to explode")118 return scene119 120 print(f"[EXPLODE_MESH] Starting mesh explosion with scale {explosion_scale}")121 print(f"[EXPLODE_MESH] Processing {len(scene.geometry)} parts")122 123 exploded_scene = trimesh.Scene()124 125 part_centers = []126 geometry_names = []127 128 for geometry_name, geometry in scene.geometry.items():129 if hasattr(geometry, 'vertices'):130 transform = scene.graph[geometry_name][0]131 vertices_global = trimesh.transformations.transform_points(132 geometry.vertices, transform)133 center = np.mean(vertices_global, axis=0)134 part_centers.append(center)135 geometry_names.append(geometry_name)136 print(f"[EXPLODE_MESH] Part {geometry_name}: center = {center}")137 138 if not part_centers:139 print("No valid geometries with vertices found")140 return scene141 142 part_centers = np.array(part_centers)143 global_center = np.mean(part_centers, axis=0)144 145 print(f"[EXPLODE_MESH] Global center: {global_center}")146 147 for i, (geometry_name, geometry) in enumerate(scene.geometry.items()):148 if hasattr(geometry, 'vertices'):149 if i < len(part_centers):150 part_center = part_centers[i]151 direction = part_center - global_center152 153 direction_norm = np.linalg.norm(direction)154 if direction_norm > 1e-6:155 direction = direction / direction_norm156 else:157 direction = np.random.randn(3)158 direction = direction / np.linalg.norm(direction)159 160 offset = direction * explosion_scale161 else:162 offset = np.zeros(3)163 164 original_transform = scene.graph[geometry_name][0].copy()165 166 new_transform = original_transform.copy()167 new_transform[:3, 3] = new_transform[:3, 3] + offset168 169 exploded_scene.add_geometry(170 geometry, 171 transform=new_transform, 172 geom_name=geometry_name173 )174 175 print(f"[EXPLODE_MESH] Part {geometry_name}: moved by {np.linalg.norm(offset):.4f}")176 177 print("[EXPLODE_MESH] Mesh explosion complete")178 return exploded_scene179 180 181def get_duration(182 image_path,183 num_parts,184 seed,185 num_tokens,186 num_inference_steps,187 guidance_scale,188 use_flash_decoder,189 rmbg,190 session_id,191 progress,192 ):193 194 duration_seconds = 60195 196 if num_parts > 5:197 duration_seconds = 75198 elif num_parts > 10:199 duration_seconds = 90200 return int(duration_seconds)201 202 203@spaces.GPU(duration=get_duration)204@torch.no_grad()205def run_triposg(image_path: str,206 num_parts: int = 1,207 seed: int = 0,208 num_tokens: int = 1024,209 num_inference_steps: int = 50,210 guidance_scale: float = 7.0,211 use_flash_decoder: bool = False,212 rmbg: bool = True,213 session_id = None,214 progress=gr.Progress(track_tqdm=True),):215 216 """217 Generate 3D part meshes from an input image.218 """219 220 max_num_expanded_coords = 1e9221 222 if session_id is None:223 session_id = uuid.uuid4().hex224 225 if rmbg:226 img_pil = prepare_image(image_path, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)227 else:228 img_pil = Image.open(image_path)229 230 set_seed(seed)231 start_time = time.time()232 outputs = pipe(233 image=[img_pil] * num_parts,234 attention_kwargs={"num_parts": num_parts},235 num_tokens=num_tokens,236 generator=torch.Generator(device=pipe.device).manual_seed(seed),237 num_inference_steps=num_inference_steps,238 guidance_scale=guidance_scale,239 max_num_expanded_coords=max_num_expanded_coords,240 use_flash_decoder=use_flash_decoder,241 ).meshes242 duration = time.time() - start_time243 print(f"Generation time: {duration:.2f}s")244 245 # Ensure no None outputs246 for i, mesh in enumerate(outputs):247 if mesh is None:248 outputs[i] = trimesh.Trimesh(vertices=[[0,0,0]], faces=[[0,0,0]])249 250 251 export_dir = os.path.join(os.environ["PARTCRAFTER_PROCESSED"], session_id)252 253 # If it already exists, delete it (and all its contents)254 if os.path.exists(export_dir):255 shutil.rmtree(export_dir)256 257 os.makedirs(export_dir, exist_ok=True)258 259 parts = []260 261 for idx, mesh in enumerate(outputs):262 part = os.path.join(export_dir, f"part_{idx:02}.glb")263 mesh.export(part)264 parts.append(part)265 266 zip_path = os.path.join(os.environ["PARTCRAFTER_PROCESSED"], f"{session_id}.zip")267 268 # shutil.make_archive wants the base name without extension:269 base_name = zip_path[:-4] # strip off '.zip'270 shutil.make_archive(base_name, 'zip', export_dir)271 272 # Merge and color273 merged = get_colored_mesh_composition(outputs)274 split_mesh = explode_mesh(merged)275 276 merged_path = os.path.join(export_dir, "object.glb")277 merged.export(merged_path)278 279 split_preview_path = os.path.join(export_dir, "split.glb")280 split_mesh.export(split_preview_path)281 282 return merged_path, split_preview_path, export_dir, zip_path283 284def cleanup(request: gr.Request):285 286 sid = request.session_hash287 if sid:288 d1 = os.path.join(os.environ["PARTCRAFTER_PROCESSED"], sid)289 shutil.rmtree(d1, ignore_errors=True)290 291def start_session(request: gr.Request):292 293 return request.session_hash294 295def build_demo():296 css = """297 #col-container {298 margin: 0 auto;299 max-width: 1280px;300 }301 """302 theme = gr.themes.Ocean()303 304 with gr.Blocks(css=css, theme=theme) as demo:305 session_state = gr.State()306 demo.load(start_session, outputs=[session_state])307 308 with gr.Column(elem_id="col-container"):309 gr.HTML(310 """311 <div style="text-align: center;">312 <p style="font-size:16px; display: inline; margin: 0;">313 <strong>PartCrafter</strong> – Structured 3D Mesh Generation via Compositional Latent Diffusion Transformers314 </p>315 <a href="https://github.com/wgsxm/PartCrafter" style="display: inline-block; vertical-align: middle; margin-left: 0.5em;">316 <img src="https://img.shields.io/badge/GitHub-Repo-blue" alt="GitHub Repo">317 </a>318 </div>319 """320 )321 with gr.Row():322 with gr.Column(scale=1):323 gr.Markdown(324 """ 325 • We would like to acknowledge : [@alexandernasa](https://twitter.com/alexandernasa/) for the contribution of the Hugging Face Space. """326 )327 input_image = gr.Image(type="filepath", label="Input Image", height=256)328 num_parts = gr.Slider(1, MAX_NUM_PARTS, value=4, step=1, label="Number of Parts")329 run_button = gr.Button("🧩 Generate 3D Parts", variant="primary")330 331 with gr.Accordion("Advanced Settings", open=False):332 seed = gr.Number(value=0, label="Random Seed", precision=0)333 num_tokens = gr.Slider(256, 2048, value=1024, step=64, label="Num Tokens")334 num_steps = gr.Slider(1, 100, value=50, step=1, label="Inference Steps")335 guidance = gr.Slider(1.0, 20.0, value=7.0, step=0.1, label="Guidance Scale")336 flash_decoder = gr.Checkbox(value=False, label="Use Flash Decoder")337 remove_bg = gr.Checkbox(value=True, label="Remove Background (RMBG)")338 339 with gr.Column(scale=2):340 gr.HTML(341 """342 <p style="opacity: 0.6; font-style: italic;">343 The 3D Preview might take a few seconds to load the 3D model344 </p>345 """346 )347 with gr.Row():348 output_model = gr.Model3D(label="Merged 3D Object")349 split_model = gr.Model3D(label="Split Preview")350 output_dir = gr.Textbox(label="Export Directory", visible=False)351 download_zip = gr.File(label="Download All Parts (zip)", visible=False)352 with gr.Row():353 with gr.Column():354 examples = gr.Examples(355 356 examples=[357 [358 "assets/images/np5_b81f29e567ea4db48014f89c9079e403.png", 359 5,360 ], 361 [362 "assets/images/np7_1c004909dedb4ebe8db69b4d7b077434.png", 363 7,364 ], 365 [366 "assets/images/np2_tree.png", 367 3,368 ], 369 370 ],371 inputs=[input_image, num_parts],372 outputs=[output_model, split_model, output_dir, download_zip],373 fn=run_triposg,374 cache_examples=True,375 )376 377 run_button.click(fn=run_triposg,378 inputs=[input_image, num_parts, seed, num_tokens, num_steps,379 guidance, flash_decoder, remove_bg, session_state],380 outputs=[output_model, split_model, output_dir, download_zip])381 return demo382 383if __name__ == "__main__":384 demo = build_demo()385 demo.unload(cleanup)386 demo.queue()387 demo.launch()