Update app.py

app.py

@@ -4,94 +4,94 @@ import sys
 import gradio as gr
 import numpy as np
 import torch
+from PIL import Image
 import trimesh
 import random
+from huggingface_hub import snapshot_download
 import shutil
+import subprocess
 
-from huggingface_hub import snapshot_download
-from PIL import Image
+# ---------------------------------------------------------------------
+# Basic setup
+# ---------------------------------------------------------------------
 
-# --------------------
-# Device
-# --------------------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16
+
 print("DEVICE:", DEVICE)
 
-# --------------------
-# Constants
-# --------------------
 DEFAULT_FACE_NUMBER = 100000
 MAX_SEED = np.iinfo(np.int32).max
 
 TRIPOSG_REPO_URL = "https://github.com/VAST-AI-Research/TripoSG.git"
-
-RMBG_PRETRAINED_MODEL = "checkpoints/RMBG-1.4"
 TRIPOSG_PRETRAINED_MODEL = "checkpoints/TripoSG"
+RMBG_PRETRAINED_MODEL = "checkpoints/RMBG-1.4"
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 TMP_DIR = os.path.join(BASE_DIR, "tmp")
 os.makedirs(TMP_DIR, exist_ok=True)
 
-# --------------------
-# Clone TripoSG
-# --------------------
+# ---------------------------------------------------------------------
+# Clone TripoSG code (runtime-safe)
+# ---------------------------------------------------------------------
+
 TRIPOSG_CODE_DIR = os.path.join(BASE_DIR, "triposg")
+
 if not os.path.exists(TRIPOSG_CODE_DIR):
     os.system(f"git clone {TRIPOSG_REPO_URL} {TRIPOSG_CODE_DIR}")
 
-sys.path.append(TRIPOSG_CODE_DIR)
+# ---------------------------------------------------------------------
+# 🔑 CRITICAL FIX: make TripoSG imports visible BEFORE importing
+# ---------------------------------------------------------------------
 
-# --------------------
-# UI Header
-# --------------------
-HEADER = """
-# 🔮 Image to 3D with TripoSG
+sys.path.insert(0, TRIPOSG_CODE_DIR)
+sys.path.insert(0, os.path.join(TRIPOSG_CODE_DIR, "scripts"))
 
-Upload an image → get a clean 3D mesh (GLB).
+# ---------------------------------------------------------------------
+# Now imports work
+# ---------------------------------------------------------------------
 
-**Texture generation intentionally disabled.**
-"""
-
-# --------------------
-# TripoSG + RMBG
-# --------------------
 from image_process import prepare_image
 from briarmbg import BriaRMBG
+from triposg.pipelines.pipeline_triposg import TripoSGPipeline
+from utils import simplify_mesh
+
+# ---------------------------------------------------------------------
+# Load models
+# ---------------------------------------------------------------------
 
 snapshot_download("briaai/RMBG-1.4", local_dir=RMBG_PRETRAINED_MODEL)
 rmbg_net = BriaRMBG.from_pretrained(RMBG_PRETRAINED_MODEL).to(DEVICE)
 rmbg_net.eval()
 
-from triposg.pipelines.pipeline_triposg import TripoSGPipeline
-
 snapshot_download("VAST-AI/TripoSG", local_dir=TRIPOSG_PRETRAINED_MODEL)
 triposg_pipe = TripoSGPipeline.from_pretrained(
     TRIPOSG_PRETRAINED_MODEL
 ).to(DEVICE, DTYPE)
 
-# --------------------
+# ---------------------------------------------------------------------
 # Helpers
-# --------------------
+# ---------------------------------------------------------------------
+
 def get_random_hex():
     return os.urandom(8).hex()
 
-def get_random_seed(randomize_seed, seed):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
+def get_random_seed(randomize, seed):
+    return random.randint(0, MAX_SEED) if randomize else seed
 
 def start_session(req: gr.Request):
-    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(save_dir, exist_ok=True)
+    path = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(path, exist_ok=True)
 
 def end_session(req: gr.Request):
-    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shutil.rmtree(save_dir, ignore_errors=True)
+    path = os.path.join(TMP_DIR, str(req.session_hash))
+    if os.path.exists(path):
+        shutil.rmtree(path)
+
+# ---------------------------------------------------------------------
+# GPU functions
+# ---------------------------------------------------------------------
 
-# --------------------
-# GPU Functions
-# --------------------
 @spaces.GPU()
 @torch.no_grad()
 def run_segmentation(image_path: str):
@@ -106,27 +106,27 @@ def run_segmentation(image_path: str):
 def image_to_3d(
     image: Image.Image,
     seed: int,
-    num_inference_steps: int,
-    guidance_scale: float,
+    steps: int,
+    guidance: float,
     simplify: bool,
-    target_face_num: int,
+    target_faces: int,
     req: gr.Request,
 ):
     outputs = triposg_pipe(
         image=image,
         generator=torch.Generator(device=DEVICE).manual_seed(seed),
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
+        num_inference_steps=steps,
+        guidance_scale=guidance,
     ).samples[0]
 
     mesh = trimesh.Trimesh(
         outputs[0].astype(np.float32),
         np.ascontiguousarray(outputs[1]),
+        process=False,
     )
 
     if simplify:
-        from utils import simplify_mesh
-        mesh = simplify_mesh(mesh, target_face_num)
+        mesh = simplify_mesh(mesh, target_faces)
 
     save_dir = os.path.join(TMP_DIR, str(req.session_hash))
     mesh_path = os.path.join(save_dir, f"triposg_{get_random_hex()}.glb")
@@ -135,35 +135,40 @@ def image_to_3d(
     torch.cuda.empty_cache()
     return mesh_path
 
-# --------------------
+# ---------------------------------------------------------------------
 # UI
-# --------------------
+# ---------------------------------------------------------------------
+
+HEADER = """
+# 🔮 Image → 3D (TripoSG)
+
+Mesh-only demo (no texture, no MV-Adapter).
+"""
+
 with gr.Blocks(title="TripoSG") as demo:
     gr.Markdown(HEADER)
 
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(label="Input Image", type="filepath")
+            input_image = gr.Image(label="Input Image", type="filepath")
             seg_image = gr.Image(label="Segmentation", type="pil")
 
-            with gr.Accordion("Generation Settings", open=True):
-                seed = gr.Slider(0, MAX_SEED, value=0, label="Seed")
-                randomize_seed = gr.Checkbox(value=True, label="Randomize Seed")
-                num_inference_steps = gr.Slider(8, 50, value=50, step=1, label="Inference Steps")
-                guidance_scale = gr.Slider(0, 20, value=7.5, step=0.1, label="CFG Scale")
-                simplify = gr.Checkbox(value=True, label="Simplify Mesh")
-                target_face_num = gr.Slider(
-                    10_000, 1_000_000, value=DEFAULT_FACE_NUMBER, label="Target Face Count"
-                )
+            seed = gr.Slider(0, MAX_SEED, value=0, label="Seed")
+            randomize_seed = gr.Checkbox(value=True, label="Randomize Seed")
+            steps = gr.Slider(8, 50, value=50, step=1, label="Inference Steps")
+            guidance = gr.Slider(0.0, 20.0, value=7.5, step=0.1, label="CFG Scale")
+
+            simplify = gr.Checkbox(value=True, label="Simplify Mesh")
+            target_faces = gr.Slider(10_000, 1_000_000, value=DEFAULT_FACE_NUMBER)
 
-            gen_button = gr.Button("Generate 3D", variant="primary")
+            gen_btn = gr.Button("Generate 3D", variant="primary")
 
         with gr.Column():
-            model_output = gr.Model3D(label="Generated GLB")
+            model_out = gr.Model3D(label="Generated GLB")
 
-    gen_button.click(
+    gen_btn.click(
         run_segmentation,
-        inputs=image_input,
+        inputs=input_image,
         outputs=seg_image,
     ).then(
         get_random_seed,
@@ -171,15 +176,8 @@ with gr.Blocks(title="TripoSG") as demo:
         outputs=seed,
     ).then(
         image_to_3d,
-        inputs=[
-            seg_image,
-            seed,
-            num_inference_steps,
-            guidance_scale,
-            simplify,
-            target_face_num,
-        ],
-        outputs=model_output,
+        inputs=[seg_image, seed, steps, guidance, simplify, target_faces],
+        outputs=model_out,
     )
 
     demo.load(start_session)