Update app.py

app.py

@@ -1,23 +1,13 @@
 import os
 import sys
-import subprocess
+import subprocess  # <--- 确保这行在这里！
 import importlib
 import site
 import time
-import uuid
-import shutil
-import glob
-from types import ModuleType
-
-# ========================================================
-# 1. 核心修复：路径环境变量与内存级伪造 diso
-# ========================================================
-# 解决 KeyError: 'PARTCRAFTER_PROCESSED'
-os.environ["PARTCRAFTER_PROCESSED"] = os.environ.get("PARTCRAFTER_PROCESSED", "outputs")
-os.makedirs(os.environ["PARTCRAFTER_PROCESSED"], exist_ok=True)
-os.environ['PYOPENGL_PLATFORM'] = 'egl'
 
+# --- 🧪 1. 内存级伪造 diso (必须在任何业务 import 之前) ---
 def mock_diso():
+    from types import ModuleType
     print("🧪 Creating emergency mock for diso...")
     diso = ModuleType("diso")
     class FakeDiffDMC:
@@ -27,51 +17,155 @@ def mock_diso():
     sys.modules["diso"] = diso
     sys.modules["diso._C"] = ModuleType("diso._C")
     sys.modules["diso.diso_native"] = ModuleType("diso.diso_native")
-    print("✅ diso mocked.")
+    print("✅ diso has been mocked successfully!")
 
 mock_diso()
 
-# ========================================================
-# 2. 极速环境安装 (避开编译，解决超时)
-# ========================================================
-def setup_environment():
-    print("🚀 Initializing Optimized Environment...")
-    # 强制降级 torch 确保与预编译包匹配
-    try:
-        import torch
-        if "2.9" in torch.__version__:
-            print("🔄 Downgrading torch to 2.4.0...")
-            subprocess.run([sys.executable, "-m", "pip", "install", "torch==2.4.0+cu121", "torchvision==0.19.0+cu121", "--extra-index-url", "https://download.pytorch.org/whl/cu121"], check=True)
-            importlib.invalidate_caches()
-            os.execv(sys.executable, ['python'] + sys.argv)
-    except: pass
-
-    # 极速安装 PyG 扩展和渲染工具
+# --- 🚀 2. 极速环境安装 (已经成功的 scatter/sparse) ---
+def install_essential_packages():
+    print("📦 Checking core dependencies...")
+    # 确保基础环境正确
+    subprocess.run([sys.executable, "-m", "pip", "install", "ninja", "setuptools", "wheel", "-q"])
+    
+    # 极速安装 PyG 扩展
     subprocess.run([
         sys.executable, "-m", "pip", "install", 
-        "torch-scatter", "torch-sparse", "torch-cluster", "ninja", "pyrender", "pyopengl==3.1.0", "trimesh", "accelerate",
-        "-f", "https://data.pyg.org/whl/torch-2.4.0+cu121.html", "--no-cache-dir"
+        "torch-scatter", "torch-sparse", "torch-cluster",
+        "-f", "https://data.pyg.org/whl/torch-2.4.0+cu121.html",
+        "--no-cache-dir"
+    ])
+    
+    # 安装剩下的渲染工具
+    subprocess.run([
+        sys.executable, "-m", "pip", "install", 
+        "pyrender", "pyopengl==3.1.0", "pyyaml", "trimesh", "accelerate", "-q"
     ])
     
     importlib.invalidate_caches()
     site.main()
-    print("🎉 Environment ready!")
+    print("🎉 Environment Installation Phase Finished.")
+
+install_essential_packages()
+
+
+# ... 之前的 mock_diso 和安装逻辑 ...
+
+
+# 1. 核心路径保护
+os.environ["PARTCRAFTER_PROCESSED"] = os.environ.get("PARTCRAFTER_PROCESSED", "outputs")
+os.makedirs(os.environ["PARTCRAFTER_PROCESSED"], exist_ok=True)
 
-# 只有在环境未就绪时运行一次
-if "torch-scatter" not in str(subprocess.check_output([sys.executable, "-m", "pip", "freeze"])):
-    setup_environment()
+# 2. 模型权重下载路径确认 (确保这些目录也存在)
+os.makedirs("pretrained_weights/PartCrafter", exist_ok=True)
+os.makedirs("pretrained_weights/RMBG-1.4", exist_ok=True)
 
-# ========================================================
-# 3. 业务代码导入 (保留你原本的所有引用)
-# ========================================================
+# ... 继续执行 snapshot_download ...
+
+# --- 3. 正式导入业务逻辑 (现在开始这几百行代码就不会报错了) ---
 import spaces
 import gradio as gr
 import numpy as np
 import torch
+import uuid
+import shutil
 from huggingface_hub import snapshot_download
 from PIL import Image
 from accelerate.utils import set_seed
+
+# 从这里往下，粘贴你原本所有的业务逻辑代码 (PartCrafterPipeline 等)
+# ...
+
+# --- 🚀 核心修复：强制版本回退以避开编译 ---
+def pre_install_check():
+    try:
+        import torch
+        # 如果是 2.9+ 版本，强制降级到有预编译包的 2.4.0
+        if "2.9" in torch.__version__:
+            print(f"🔄 Current torch {torch.__version__} is too new. Downgrading to 2.4.0 for speed...")
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "ninja", "setuptools", "wheel", "-q"])
+            subprocess.check_call([
+                sys.executable, "-m", "pip", "install", 
+                "torch==2.4.0+cu121", "torchvision==0.19.0+cu121",
+                "--extra-index-url", "https://download.pytorch.org/whl/cu121"
+            ])
+            # 刷新路径
+            importlib.invalidate_caches()
+            os.execv(sys.executable, ['python'] + sys.argv) # 重启进程以加载新版本
+    except Exception as e:
+        print(f"Pre-install check note: {e}")
+
+pre_install_check()
+
 import trimesh
+import glob
+import importlib, site
+
+# Re-discover all .pth/.egg-link files
+for sitedir in site.getsitepackages():
+    site.addsitedir(sitedir)
+
+importlib.invalidate_caches()
+
+# --- 简化的 CUDA 环境配置 ---
+def setup_cuda_env():
+    cuda_path = "/usr/local/cuda"
+    if os.path.exists(cuda_path):
+        os.environ["CUDA_HOME"] = cuda_path
+        os.environ["PATH"] = f"{cuda_path}/bin:{os.environ['PATH']}"
+        os.environ["LD_LIBRARY_PATH"] = f"{cuda_path}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"
+        print(f"==> Using system CUDA at {cuda_path}")
+
+setup_cuda_env()
+
+# --- 🚀 针对 PyTorch 2.9.1 的优化源码编译方案 ---
+# --- 🚀 暴力整合版：攻克 diso 最后的防线 ---
+def install_heavy_packages():
+    os.environ['PYOPENGL_PLATFORM'] = 'egl'
+    
+    # 1. PyG 扩展（这部分已经稳了，保持不动）
+    print("📦 Installing PyG extensions...")
+    subprocess.run([
+        sys.executable, "-m", "pip", "install", 
+        "torch-scatter", "torch-sparse", "torch-cluster",
+        "-f", "https://data.pyg.org/whl/torch-2.4.0+cu121.html"
+    ], check=True)
+
+    # 2. 暴力解决 diso：克隆源码 -> 强行导入
+    print("🔥 Attempting D-Plan: Manual diso injection...")
+    diso_path = os.path.join(os.getcwd(), "diso_source")
+    if not os.path.exists(diso_path):
+        subprocess.run(["git", "clone", "https://github.com/SarahWeiii/diso.git", diso_path])
+    
+    # 将 diso 的源码路径直接加入系统搜索路径
+    # 这样即使没有编译成功 .so 文件，Python 也能找到包结构
+    if diso_path not in sys.path:
+        sys.path.insert(0, diso_path)
+
+    # 3. 安装渲染和其他轻量级依赖
+    print("📦 Installing rendering tools...")
+    subprocess.run([sys.executable, "-m", "pip", "install", "pyrender", "pyopengl==3.1.0", "pyyaml", "-q"], check=True)
+    
+    importlib.invalidate_caches()
+    print("🎉 Environment Installation Phase Finished.")
+
+# 执行安装
+install_heavy_packages()
+
+# --- 🛰️ 关键：diso 导入补丁 ---
+try:
+    import diso
+    print("✅ diso imported successfully!")
+except ImportError:
+    # 如果还是报错，尝试将 diso 内部的包直接暴露出来
+    print("⚠️ diso import failed, applying emergency mock...")
+    diso_src_path = os.path.join(os.getcwd(), "diso_source")
+    sys.path.insert(0, diso_src_path)
+    # 强制让 Python 识别 diso 目录
+    importlib.invalidate_caches()
+
+# ... 后续代码保持不变 ...
+
+
 
 from src.utils.data_utils import get_colored_mesh_composition, scene_to_parts, load_surfaces
 from src.utils.render_utils import render_views_around_mesh, render_normal_views_around_mesh, make_grid_for_images_or_videos, export_renderings, explode_mesh
@@ -79,13 +173,12 @@ from src.pipelines.pipeline_partcrafter import PartCrafterPipeline
 from src.utils.image_utils import prepare_image
 from src.models.briarmbg import BriaRMBG
 
-# ========================================================
-# 4. 业务逻辑 (100% 保留你代码中的参数与函数)
-# ========================================================
+# Constants
 MAX_NUM_PARTS = 16
 DEVICE = "cuda" 
 DTYPE = torch.float16
 
+# Download and initialize models
 partcrafter_weights_dir = "pretrained_weights/PartCrafter"
 rmbg_weights_dir = "pretrained_weights/RMBG-1.4"
 snapshot_download(repo_id="wgsxm/PartCrafter", local_dir=partcrafter_weights_dir)
@@ -99,34 +192,119 @@ def first_file_from_dir(directory, ext):
     files = glob.glob(os.path.join(directory, f"*.{ext}"))
     return sorted(files)[0] if files else None
 
-def get_duration(image_path, num_parts, seed, num_tokens, num_steps, guidance, flash, rmbg, session, progress):
+
+
+def get_duration(
+    image_path,
+    num_parts,
+    seed,
+    num_tokens,
+    num_inference_steps,
+    guidance_scale,
+    use_flash_decoder,
+    rmbg,
+    session_id,
+    progress,
+    ):
+
     duration_seconds = 75
-    if num_parts > 10: duration_seconds = 120
-    elif num_parts > 5: duration_seconds = 90
+
+    if num_parts > 10:
+        duration_seconds = 120
+    elif num_parts > 5:
+        duration_seconds = 90
+    
     return int(duration_seconds)
+        
 
 @spaces.GPU(duration=140)
-def gen_model_n_video(image_path, num_parts, progress=gr.Progress(track_tqdm=True)):
+def gen_model_n_video(image_path: str,
+                      num_parts: int,
+                      progress=gr.Progress(track_tqdm=True),):
+
     model_path = run_partcrafter(image_path, num_parts=num_parts, progress=progress)
     video_path = gen_video(model_path)
+
     return model_path, video_path
 
 @spaces.GPU()
 def gen_video(model_path):
+
     if model_path is None:
         gr.Info("You must craft the 3d parts first")
+
         return None
+        
     export_dir = os.path.dirname(model_path)
+
     merged = trimesh.load(model_path)
+
     preview_path = os.path.join(export_dir, "rendering.gif")
-    rendered_images = render_views_around_mesh(merged, num_views=36, radius=4)
-    export_renderings(rendered_images, preview_path, fps=7)
+
+    num_views = 36
+    radius = 4
+    fps = 7
+    rendered_images = render_views_around_mesh(
+        merged,
+        num_views=num_views,
+        radius=radius,
+    )
+
+    export_renderings(
+        rendered_images,
+        preview_path,
+        fps=fps,
+    )
     return preview_path
 
 @spaces.GPU(duration=get_duration)
 @torch.no_grad()
-def run_partcrafter(image_path, num_parts=1, seed=0, num_tokens=1024, num_inference_steps=50, guidance_scale=7.0, use_flash_decoder=False, rmbg=True, session_id=None, progress=gr.Progress(track_tqdm=True)):
-    if session_id is None: session_id = uuid.uuid4().hex
+def run_partcrafter(image_path: str,
+                num_parts: int = 1,
+                seed: int = 0,
+                num_tokens: int = 1024,
+                num_inference_steps: int = 50,
+                guidance_scale: float = 7.0,
+                use_flash_decoder: bool = False,
+                rmbg: bool = True,
+                session_id = None,
+                progress=gr.Progress(track_tqdm=True),):
+
+    """
+    Generate structured 3D meshes from a 2D image using the PartCrafter pipeline.
+
+    This function takes a single 2D image as input and produces a set of part-based 3D meshes,
+    using compositional latent diffusion with attention to structure and part separation.
+    Optionally removes the background using a pretrained background removal model (RMBG),
+    and outputs a merged object mesh.
+
+    Args:
+        image_path (str): Path to the input image file on disk.
+        num_parts (int, optional): Number of distinct parts to decompose the object into. Defaults to 1.
+        seed (int, optional): Random seed for reproducibility. Defaults to 0.
+        num_tokens (int, optional): Number of tokens used during latent encoding. Higher values yield finer detail. Defaults to 1024.
+        num_inference_steps (int, optional): Number of diffusion inference steps. More steps improve quality but increase runtime. Defaults to 50.
+        guidance_scale (float, optional): Classifier-free guidance scale. Higher values emphasize adherence to conditioning. Defaults to 7.0.
+        use_flash_decoder (bool, optional): Whether to use FlashAttention in the decoder for performance. Defaults to False.
+        rmbg (bool, optional): Whether to apply background removal before processing. Defaults to True.
+        session_id (str, optional): Optional session ID to manage export paths. If not provided, a random UUID is generated.
+        progress (gr.Progress, optional): Gradio progress object for visual feedback. Automatically handled by Gradio.
+
+    Returns:
+        Tuple[str, str, str, str]: 
+            - `merged_path` (str): File path to the merged full object mesh (`object.glb`).
+
+    Notes:
+        - This function utilizes HuggingFace pretrained weights for both part generation and background removal.
+        - The final output includes merged model parts to visualize object structure.
+        - Generation time depends on the number of parts and inference parameters.
+    """
+
+    max_num_expanded_coords = 1e9
+
+    if session_id is None:
+        session_id = uuid.uuid4().hex
+        
     if rmbg:
         img_pil = prepare_image(image_path, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
     else:
@@ -141,47 +319,92 @@ def run_partcrafter(image_path, num_parts=1, seed=0, num_tokens=1024, num_infere
         generator=torch.Generator(device=pipe.device).manual_seed(seed),
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
-        max_num_expanded_coords=1e9,
+        max_num_expanded_coords=max_num_expanded_coords,
         use_flash_decoder=use_flash_decoder,
     ).meshes
-    print(f"Generation time: {time.time() - start_time:.2f}s")
+    duration = time.time() - start_time
+    print(f"Generation time: {duration:.2f}s")
 
+    # Ensure no None outputs
     for i, mesh in enumerate(outputs):
-        if mesh is None: outputs[i] = trimesh.Trimesh(vertices=[[0,0,0]], faces=[[0,0,0]])
+        if mesh is None:
+            outputs[i] = trimesh.Trimesh(vertices=[[0,0,0]], faces=[[0,0,0]])
+
 
     export_dir = os.path.join(os.environ["PARTCRAFTER_PROCESSED"], session_id)
-    if os.path.exists(export_dir): shutil.rmtree(export_dir)
+
+    # If it already exists, delete it (and all its contents)
+    if os.path.exists(export_dir):
+        shutil.rmtree(export_dir)
+    
     os.makedirs(export_dir, exist_ok=True)
 
+    parts = []
+    
     for idx, mesh in enumerate(outputs):
-        mesh.export(os.path.join(export_dir, f"part_{idx:02}.glb"))
+        part = os.path.join(export_dir, f"part_{idx:02}.glb")
+        mesh.export(part)
+        parts.append(part)
     
+    # Merge and color
     merged = get_colored_mesh_composition(outputs)
+    split_mesh = explode_mesh(merged)
+    
     merged_path = os.path.join(export_dir, "object.glb")
     merged.export(merged_path)
+
     return merged_path
 
-# ========================================================
-# 5. UI 界面逻辑 (完全保留你原来的 CSS 和 Examples)
-# ========================================================
 def cleanup(request: gr.Request):
+
     sid = request.session_hash
     if sid:
-        shutil.rmtree(os.path.join(os.environ["PARTCRAFTER_PROCESSED"], sid), ignore_errors=True)
+        d1 = os.path.join(os.environ["PARTCRAFTER_PROCESSED"], sid)
+        shutil.rmtree(d1, ignore_errors=True)
+        
+def start_session(request: gr.Request):
 
+    return request.session_hash
+    
 def build_demo():
-    css = "#col-container { margin: 0 auto; max-width: 1560px; }"
-    with gr.Blocks(css=css, theme=gr.themes.Ocean()) as demo:
+    css = """
+        #col-container {
+            margin: 0 auto;
+            max-width: 1560px;
+        }
+        """
+    theme = gr.themes.Ocean()
+    
+    with gr.Blocks(css=css, theme=theme) as demo:
         session_state = gr.State()
-        demo.load(lambda r: r.session_hash, outputs=[session_state])
+        demo.load(start_session, outputs=[session_state])
+
         with gr.Column(elem_id="col-container"):
-            gr.HTML("<div style='text-align:center;'><strong>PartCrafter</strong> – Structured 3D Mesh Generation</div>")
+            gr.HTML(
+                """
+                <div style="text-align: center;">
+                    <p style="font-size:16px; display: inline; margin: 0;">
+                        <strong>PartCrafter</strong> – Structured 3D Mesh Generation via Compositional Latent Diffusion Transformers
+                    </p>
+                    <a href="https://github.com/wgsxm/PartCrafter" style="display: inline-block; vertical-align: middle; margin-left: 0.5em;">
+                        <img src="https://img.shields.io/badge/GitHub-Repo-blue" alt="GitHub Repo">
+                    </a>
+                </div>
+                <div style="text-align: center;">
+                    HF Space by :<a href="https://twitter.com/alexandernasa/" style="display: inline-block; vertical-align: middle; margin-left: 0.5em;">
+                        <img src="https://img.shields.io/twitter/url/https/twitter.com/cloudposse.svg?style=social&label=Follow Me" alt="GitHub Repo">
+                    </a>
+                </div>
+                """
+            )
             with gr.Row():
                 with gr.Column(scale=1):
+                    
                     input_image = gr.Image(type="filepath", label="Input Image", height=256)
                     num_parts = gr.Slider(1, MAX_NUM_PARTS, value=4, step=1, label="Number of Parts")
                     run_button = gr.Button("Step 1 - 🧩 Craft 3D Parts", variant="primary")
                     video_button = gr.Button("Step 2 - 🎥 Generate Split Preview Gif (Optional)")
+                    
                     with gr.Accordion("Advanced Settings", open=False):
                         seed = gr.Number(value=0, label="Random Seed", precision=0)
                         num_tokens = gr.Slider(256, 2048, value=1024, step=64, label="Num Tokens")
@@ -189,19 +412,59 @@ def build_demo():
                         guidance = gr.Slider(1.0, 20.0, value=7.0, step=0.1, label="Guidance Scale")
                         flash_decoder = gr.Checkbox(value=False, label="Use Flash Decoder")
                         remove_bg = gr.Checkbox(value=True, label="Remove Background (RMBG)")
+
                 with gr.Column(scale=2):
-                    output_model = gr.Model3D(label="Merged 3D Object", height=512)
-                    video_output = gr.Image(label="Split Preview", height=512)
-            gr.Examples(
-                examples=[["assets/images/np5_b81f29e567ea4db48014f89c9079e403.png", 5]],
-                inputs=[input_image, num_parts],
-                outputs=[output_model, video_output],
-                fn=gen_model_n_video,
-                cache_examples=True
-            )
-        run_button.click(fn=run_partcrafter, inputs=[input_image, num_parts, seed, num_tokens, num_steps, guidance, flash_decoder, remove_bg, session_state], outputs=[output_model])
-        video_button.click(fn=gen_video, inputs=[output_model], outputs=[video_output])
-    return demo
+                    gr.HTML(
+                        """
+                        <p style="opacity: 0.6; font-style: italic;">
+                          The 3D Preview might take a few seconds to load the 3D model
+                        </p>
+                        """
+                    )
+                    with gr.Row():
+                        output_model = gr.Model3D(label="Merged 3D Object", height=512, interactive=False)
+                        video_output = gr.Image(label="Split Preview", height=512)
+            with gr.Row():
+                with gr.Column():
+                    examples = gr.Examples(
+                        
+                        examples=[
+                            [
+                                "assets/images/np5_b81f29e567ea4db48014f89c9079e403.png", 
+                                5,
+                            ], 
+                            [
+                                "assets/images/np7_1c004909dedb4ebe8db69b4d7b077434.png", 
+                                7,
+                            ], 
+                            [
+                                "assets/images/np16_dino.png", 
+                                16,
+                            ], 
+                            [
+                                "assets/images/np13_39c0fa16ed324b54a605dcdbcd80797c.png", 
+                                13,
+                            ], 
+                            
+                        ],
+                        inputs=[input_image, num_parts],
+                        outputs=[output_model, video_output],
+                        fn=gen_model_n_video,
+                        cache_examples=True
+                    )
+    
+            run_button.click(fn=run_partcrafter,
+                             inputs=[input_image, num_parts, seed, num_tokens, num_steps,
+                                     guidance, flash_decoder, remove_bg, session_state],
+                             outputs=[output_model])
+            video_button.click(fn=gen_video,
+                             inputs=[output_model],
+                             outputs=[video_output])
+            
+        return demo
 
 if __name__ == "__main__":
-    build_demo().launch()
+    demo = build_demo()
+    demo.unload(cleanup)
+    demo.queue()
+    demo.launch(mcp_server=True, ssr_mode=False)