Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import os | |
| os.environ['SPCONV_ALGO'] = 'native' | |
| os.environ['ATTN_BACKEND'] = 'xformers' | |
| os.environ['SPARSE_ATTN_BACKEND'] = 'xformers' | |
| import spaces | |
| # Build the CUDA extensions that used to ship as torch-2.4-ABI prebuilt wheels | |
| # (nvdiffrast, diff_gaussian_rasterization) and the torch-2.4-cu121-only | |
| # torch-scatter, all from source against the actually-installed torch. | |
| import subprocess, sys, tempfile, ctypes | |
| CUDA_HOME = "/cuda-image/usr/local/cuda-13.0" | |
| CUDA_LIBDIR = os.path.join(CUDA_HOME, "lib64") | |
| def _first_gpu_setup(): | |
| need = {} | |
| for name, modname in [ | |
| ("nvdiffrast", "nvdiffrast"), | |
| ("diff_gaussian_rasterization", "diff_gaussian_rasterization"), | |
| ("torch_scatter", "torch_scatter"), | |
| ("detectron2", "detectron2"), | |
| ]: | |
| try: | |
| __import__(modname) | |
| except ImportError: | |
| need[name] = True | |
| if not need: | |
| print("CUDA extensions already present.") | |
| return | |
| if not os.path.exists(os.path.join(CUDA_HOME, "bin", "nvcc")): | |
| raise RuntimeError(f"nvcc not at {CUDA_HOME}/bin/nvcc; update CUDA_HOME.") | |
| patch_dir = tempfile.mkdtemp(prefix="torch_cuda_patch_") | |
| with open(os.path.join(patch_dir, "sitecustomize.py"), "w") as f: | |
| f.write( | |
| "try:\n" | |
| " import torch.utils.cpp_extension as _c\n" | |
| " _c._check_cuda_version = lambda *a, **k: None\n" | |
| "except Exception:\n" | |
| " pass\n" | |
| ) | |
| env = os.environ.copy() | |
| env["CUDA_HOME"] = CUDA_HOME | |
| env["CUDA_PATH"] = CUDA_HOME | |
| env["PATH"] = os.path.join(CUDA_HOME, "bin") + os.pathsep + env.get("PATH", "") | |
| env["PYTHONPATH"] = patch_dir + os.pathsep + env.get("PYTHONPATH", "") | |
| env["TORCH_CUDA_ARCH_LIST"] = "12.0" | |
| subprocess.check_call( | |
| [sys.executable, "-m", "pip", "install", "--no-deps", "setuptools", "wheel", "ninja"], | |
| ) | |
| if "nvdiffrast" in need: | |
| subprocess.check_call( | |
| [sys.executable, "-m", "pip", "install", "--no-build-isolation", | |
| "git+https://github.com/NVlabs/nvdiffrast/"], | |
| env=env, | |
| ) | |
| if "diff_gaussian_rasterization" in need: | |
| mip_dir = tempfile.mkdtemp(prefix="mip_splatting_") | |
| subprocess.check_call( | |
| ["git", "clone", "--recursive", "--depth=1", | |
| "https://github.com/autonomousvision/mip-splatting.git", mip_dir], | |
| ) | |
| subprocess.check_call( | |
| [sys.executable, "-m", "pip", "install", "--no-build-isolation", | |
| os.path.join(mip_dir, "submodules", "diff-gaussian-rasterization")], | |
| env=env, | |
| ) | |
| if "torch_scatter" in need: | |
| # Prebuilt wheel for torch 2.10.0 + cu128 (matches the torch wheel's | |
| # CUDA), avoids the CUDA-version mismatch you get when building against | |
| # the container's CUDA 13 toolkit. | |
| subprocess.check_call( | |
| [sys.executable, "-m", "pip", "install", "--no-deps", | |
| "torch-scatter", | |
| "-f", "https://data.pyg.org/whl/torch-2.10.0+cu128.html"], | |
| ) | |
| if "detectron2" in need: | |
| subprocess.check_call( | |
| [sys.executable, "-m", "pip", "install", "--no-build-isolation", | |
| "git+https://github.com/facebookresearch/detectron2.git"], | |
| env=env, | |
| ) | |
| _first_gpu_setup() | |
| ctypes.CDLL(os.path.join(CUDA_LIBDIR, "libcudart.so.13"), mode=ctypes.RTLD_GLOBAL) | |
| os.environ["LD_LIBRARY_PATH"] = CUDA_LIBDIR + os.pathsep + os.environ.get("LD_LIBRARY_PATH", "") | |
| # torch 2.6 flipped torch.load weights_only default to True; the bbox_gen | |
| # partfield checkpoint pickles yacs.config.CfgNode and other non-tensor | |
| # objects. Trusted upstream checkpoint — fall back to weights_only=False. | |
| import torch as _torch | |
| _orig_load = _torch.load | |
| _torch.load = lambda *a, **k: _orig_load(*a, **{**k, "weights_only": k.get("weights_only", False)}) | |
| # xformers on Blackwell (sm_120) has no operator for fp32 memory_efficient_attention: | |
| # FA3 / Cutlass cap at compute capability 9.0, FA2 only supports fp16/bf16. DINOv2 | |
| # (loaded via torch.hub) calls xops.memory_efficient_attention with fp32 → no | |
| # dispatch. Route through torch SDPA, which handles sm_120 fp32 natively. | |
| import xformers.ops as _xops | |
| import torch.nn.functional as _F | |
| def _mea_via_sdpa(q, k, v, attn_bias=None, p=0.0, scale=None, op=None): | |
| q = q.transpose(1, 2); k = k.transpose(1, 2); v = v.transpose(1, 2) | |
| attn_mask = None | |
| if attn_bias is not None: | |
| if hasattr(attn_bias, "materialize"): | |
| try: | |
| attn_mask = attn_bias.materialize((q.shape[-2], k.shape[-2]), device=q.device, dtype=q.dtype) | |
| except Exception: | |
| attn_mask = None | |
| elif isinstance(attn_bias, _torch.Tensor): | |
| attn_mask = attn_bias | |
| out = _F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=p, scale=scale) | |
| return out.transpose(1, 2).contiguous() | |
| _xops.memory_efficient_attention = _mea_via_sdpa | |
| import shutil | |
| from huggingface_hub import hf_hub_download | |
| from app_utils import ( | |
| generate_parts, | |
| prepare_models, | |
| process_image, | |
| apply_merge, | |
| DEFAULT_SIZE_TH, | |
| TMP_ROOT, | |
| ) | |
| EXAMPLES = [ | |
| ["assets/example_data/knight.png", 1800, "6,0,26,20,7;13,1,22,11,12,2,21,27,3,24,23;5,18;4,17;19,16,14,25,28", 42], | |
| ["assets/example_data/car.png", 2000, "12,10,2,11;1,7", 42], | |
| ["assets/example_data/warhammer.png", 1800, "7,1,0,8", 0], | |
| ["assets/example_data/snake.png", 3000, "2,3;0,1;4,5,6,7", 42], | |
| ["assets/example_data/Batman.png", 1800, "4,5", 42], | |
| ["assets/example_data/robot1.jpeg", 1600, "0,5;10,14,3;1,12,2;13,11,4;7,15", 42], | |
| ["assets/example_data/astronaut.png", 2000, "0,4,6;1,8,9,7;2,5", 42], | |
| ["assets/example_data/crossbow.jpg", 2000, "2,9;10,12,0,7,11,8,13;4,3", 42], | |
| ["assets/example_data/robot.jpg", 1600, "7,19;15,0;6,18", 42], | |
| ["assets/example_data/robot_dog.jpg", 1000, "21,9;2,12,10,15,17;11,7;1,0;13,19;4,16", 0], | |
| ["assets/example_data/crossbow.jpg", 1600, "9,2;10,15,13;7,14,8,11;0,12,16;5,3,1", 42], | |
| ["assets/example_data/robot.jpg", 1800, "1,2,3,5,4,16,17;11,7,19;10,14;18,6,0,15;13,9;12,8", 0], | |
| ["assets/example_data/robot_dog.jpg", 1000, "2,12,10,15,17,8,3,5,13,19,6,14;11,7;1,0,21,9,11;4,16", 0], | |
| ] | |
| HEADER = """ | |
| # OmniPart: Part-Aware 3D Generation with Semantic Decoupling and Structural Cohesion | |
| 🔮 Generate **part-aware 3D content** from a single 2D image with **2D mask control**. | |
| ## How to Use | |
| **🚀 Quick Start**: Select an example below and click **"▶️ Run Example"** | |
| **📋 Custom Image Processing**: | |
| 1. **Upload Image** - Select your image file | |
| 2. **Click "Segment Image"** - Get initial 2D segmentation | |
| 3. **Merge Segments** - Enter merge groups like `0,1;3,4` and click **"Apply Merge"** (Recommend keeping **2-15 parts**) | |
| 4. **Click "Generate 3D Model"** - Create the final 3D results | |
| """ | |
| def start_session(req: gr.Request): | |
| user_dir = os.path.join(TMP_ROOT, str(req.session_hash)) | |
| os.makedirs(user_dir, exist_ok=True) | |
| def end_session(req: gr.Request): | |
| user_dir = os.path.join(TMP_ROOT, str(req.session_hash)) | |
| shutil.rmtree(user_dir) | |
| with gr.Blocks(title="OmniPart") as demo: | |
| gr.Markdown(HEADER) | |
| state = gr.State({}) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("<div style='text-align: center'>\n\n## Input\n\n</div>") | |
| input_image = gr.Image(label="Upload Image", type="filepath", height=250, width=250) | |
| with gr.Row(): | |
| segment_btn = gr.Button("Segment Image", variant="primary", size="lg") | |
| run_example_btn = gr.Button("▶️ Run Example", variant="secondary", size="lg") | |
| size_threshold = gr.Slider( | |
| minimum=600, | |
| maximum=4000, | |
| value=DEFAULT_SIZE_TH, | |
| step=200, | |
| label="Minimum Segment Size (pixels)", | |
| info="Segments smaller than this will be ignored" | |
| ) | |
| gr.Markdown("### Merge Controls") | |
| merge_input = gr.Textbox( | |
| label="Merge Groups", | |
| placeholder="0,1;3,4", | |
| lines=2, | |
| info="Specify which segments to merge (e.g., '0,1;3,4' merges segments 0&1 together and 3&4 together)" | |
| ) | |
| merge_btn = gr.Button("Apply Merge", variant="primary", size="lg") | |
| gr.Markdown("### 3D Generation Controls") | |
| seed_slider = gr.Slider( | |
| minimum=0, | |
| maximum=10000, | |
| value=42, | |
| step=1, | |
| label="Generation Seed", | |
| info="Random seed for 3D model generation" | |
| ) | |
| cfg_slider = gr.Slider( | |
| minimum=0.0, | |
| maximum=15.0, | |
| value=7.5, | |
| step=0.5, | |
| label="CFG Strength", | |
| info="Classifier-Free Guidance strength" | |
| ) | |
| generate_mesh_btn = gr.Button("Generate 3D Model", variant="secondary", size="lg") | |
| with gr.Column(scale=2): | |
| gr.Markdown("<div style='text-align: center'>\n\n## Results Display\n\n</div>") | |
| with gr.Row(): | |
| initial_seg = gr.Image(label="Init Seg", height=220, width=220) | |
| pre_merge_vis = gr.Image(label="Pre-merge", height=220, width=220) | |
| merged_seg = gr.Image(label="Merged Seg", height=220, width=220) | |
| with gr.Row(): | |
| bbox_mesh = gr.Model3D(label="Bounding Boxes", height=350) | |
| whole_mesh = gr.Model3D(label="Combined Parts", height=350) | |
| exploded_mesh = gr.Model3D(label="Exploded Parts", height=350) | |
| with gr.Row(): | |
| combined_gs = gr.Model3D(label="Combined 3D Gaussians", clear_color=(0.0, 0.0, 0.0, 0.0), height=350) | |
| exploded_gs = gr.Model3D(label="Exploded 3D Gaussians", clear_color=(0.0, 0.0, 0.0, 0.0), height=350) | |
| with gr.Row(): | |
| examples = gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=[input_image, size_threshold, merge_input, seed_slider], | |
| cache_examples=False, | |
| ) | |
| demo.load(start_session) | |
| demo.unload(end_session) | |
| segment_btn.click( | |
| process_image, | |
| inputs=[input_image, size_threshold], | |
| outputs=[initial_seg, pre_merge_vis, state] | |
| ) | |
| merge_btn.click( | |
| apply_merge, | |
| inputs=[merge_input, state], | |
| outputs=[merged_seg, state] | |
| ) | |
| generate_mesh_btn.click( | |
| generate_parts, | |
| inputs=[state, seed_slider, cfg_slider], | |
| outputs=[bbox_mesh, whole_mesh, exploded_mesh, combined_gs, exploded_gs] | |
| ) | |
| run_example_btn.click( | |
| fn=process_image, | |
| inputs=[input_image, size_threshold], | |
| outputs=[initial_seg, pre_merge_vis, state] | |
| ).then( | |
| fn=apply_merge, | |
| inputs=[merge_input, state], | |
| outputs=[merged_seg, state] | |
| ).then( | |
| fn=generate_parts, | |
| inputs=[state, seed_slider, cfg_slider], | |
| outputs=[bbox_mesh, whole_mesh, exploded_mesh, combined_gs, exploded_gs] | |
| ) | |
| if __name__ == "__main__": | |
| os.makedirs("ckpt", exist_ok=True) | |
| sam_ckpt_path = hf_hub_download(repo_id="omnipart/OmniPart_modules", filename="sam_vit_h_4b8939.pth", local_dir="ckpt") | |
| partfield_ckpt_path = hf_hub_download(repo_id="omnipart/OmniPart_modules", filename="partfield_encoder.ckpt", local_dir="ckpt") | |
| bbox_gen_ckpt_path = hf_hub_download(repo_id="omnipart/OmniPart_modules", filename="bbox_gen.ckpt", local_dir="ckpt") | |
| prepare_models(sam_ckpt_path, partfield_ckpt_path, bbox_gen_ckpt_path) | |
| demo.launch() |