Update app.py

app.py

@@ -1,20 +1,15 @@
 """
-AI 3D Generation Pipeline — Hugging Face Space (TripoSG + rembg)
+AI 3D Generator - TripoSG + rembg (Fixed for HF Space)
 """
-from __future__ import annotations
 import os
-import re
 import sys
-import json
 import uuid
 import time
-import shutil
-import base64
 import logging
 import tempfile
-import zipfile
 from pathlib import Path
 from typing import List
+
 import gradio as gr
 import numpy as np
 import torch
@@ -22,7 +17,7 @@ from PIL import Image
 import trimesh
 from rembg import remove
 
-# HuggingFace ZeroGPU
+# ZeroGPU support
 try:
     import spaces
 except ImportError:
@@ -33,31 +28,25 @@ except ImportError:
                 return lambda f: f
             return fn
 
-# ---------------------------------------------------------------------------
-# Configuration
-# ---------------------------------------------------------------------------
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+# Config
 MODEL_ID = "VAST-AI/TripoSG"
 TRIPOSG_REPO_URL = "https://github.com/VAST-AI-Research/TripoSG.git"
-TRIPOSG_CODE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "triposg_repo")
-MAX_FACES = 90000
-DEFAULT_FACES = 50000
+TRIPOSG_CODE_DIR = "triposg_repo"
 OUTPUT_DIR = Path(tempfile.gettempdir()) / "triposg_output"
 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-logger = logging.getLogger(__name__)
-
-# Clone repo if needed
+# Clone repo
 if not os.path.exists(TRIPOSG_CODE_DIR):
-    logger.info("Cloning TripoSG repository...")
+    logger.info("Cloning TripoSG...")
     os.system(f"git clone {TRIPOSG_REPO_URL} {TRIPOSG_CODE_DIR}")
 
 sys.path.insert(0, TRIPOSG_CODE_DIR)
 sys.path.insert(0, os.path.join(TRIPOSG_CODE_DIR, "scripts"))
 
-# ---------------------------------------------------------------------------
-# Global models
-# ---------------------------------------------------------------------------
+# Global
 triposg_pipeline = None
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -65,61 +54,60 @@ def load_models():
     global triposg_pipeline
     if triposg_pipeline is not None:
         return
-    logger.info("Loading TripoSG pipeline...")
+    logger.info("Loading TripoSG...")
     try:
         from huggingface_hub import snapshot_download
         from triposg.pipelines.pipeline_triposg import TripoSGPipeline
         model_path = snapshot_download(MODEL_ID)
-        triposg_pipeline = TripoSGPipeline.from_pretrained(model_path).to(
-            device, torch.float16 if device == "cuda" else torch.float32
-        )
-        logger.info("TripoSG loaded successfully")
+        triposg_pipeline = TripoSGPipeline.from_pretrained(model_path).to(device, torch.float16 if device == "cuda" else torch.float32)
+        logger.info("TripoSG loaded")
     except Exception as e:
-        logger.error("Failed to load TripoSG: %s", e)
+        logger.error(f"Load failed: {e}")
         triposg_pipeline = None
 
-# ---------------------------------------------------------------------------
-# Background Removal
-# ---------------------------------------------------------------------------
 def remove_background(image: Image.Image) -> Image.Image:
     try:
-        output = remove(image)
-        logger.info("Background removed with rembg")
-        return output
+        return remove(image)
     except Exception as e:
-        logger.warning("rembg failed: %s", e)
+        logger.warning(f"rembg failed: {e}")
         return image
 
-# ---------------------------------------------------------------------------
-# 其他原有函數 (select_best_view, create_composite_view, extract_dimensions..., repair_mesh, scale_mesh, segment_mesh_parts, export_mesh)
-# 請把你原本這些函數貼返入嚟（我下面留空位，你 copy 返落去）
-# ---------------------------------------------------------------------------
+# ================== 下面貼返你原本其他函數 ==================
+# 請你從原本 app.py copy 下面這些函數貼入去：
+# select_best_view, create_composite_view, extract_dimensions_from_prompt, repair_mesh, scale_mesh, segment_mesh_parts, export_mesh
 
-# ... paste your original functions here (from select_best_view to export_mesh) ...
+# (如果你冇，我可以提供最小版 placeholder)
 
-# ---------------------------------------------------------------------------
-# Main Generation (keep your original logic)
-# ---------------------------------------------------------------------------
+# ================== Main Pipeline ==================
 @spaces.GPU(duration=300)
-def generate_3d_model(
-    input_images: list,
-    text_prompt: str,
-    num_faces: int = DEFAULT_FACES,
-    export_format: str = "glb",
-    enable_modular: bool = False,
-    progress=gr.Progress(track_tqdm=True),
-):
-    # ... 你原本 generate_3d_model 嘅 code（load_models, process images, remove bg, triposg inference, post-process, export）...
-    # 只需要確保第一��� load_models() 仲喺度
-    pass   # ← 暫時 placeholder，請貼返你原本內容
-
-# ---------------------------------------------------------------------------
-# Gradio UI
-# ---------------------------------------------------------------------------
-with gr.Blocks(title="AI 3D Model Generator") as demo:
-    # ... 你原本 UI code（HEADER_MD, inputs, outputs, button click 等）...
-    pass   # ← placeholder，請貼返你原本 UI 部分
-
-# Launch
+def generate_3d_model(input_images: list, text_prompt: str, num_faces: int = 50000, export_format: str = "glb", progress=gr.Progress()):
+    load_models()
+    if triposg_pipeline is None:
+        raise gr.Error("Model load failed. Check logs.")
+
+    # ... 你原本 generate 邏輯 ...
+    # 暫時用 placeholder
+    return None, None, "Pipeline running... (add your full logic here)"
+
+# ================== UI ==================
+with gr.Blocks(title="AI 3D Generator") as demo:
+    gr.Markdown("# 🚀 TripoSG 3D Generator (rembg version)")
+    with gr.Row():
+        with gr.Column():
+            input_images = gr.File(label="Upload Images", file_count="multiple", file_types=["image"])
+            text_prompt = gr.Textbox(label="Text Prompt (dimensions etc)", lines=3)
+            num_faces = gr.Slider(5000, 90000, value=50000, step=5000, label="Target Faces")
+            btn = gr.Button("Generate 3D Model", variant="primary")
+        with gr.Column():
+            preview = gr.Model3D(label="3D Preview")
+            download = gr.File(label="Download Model")
+            status = gr.Textbox(label="Status", lines=8)
+
+    btn.click(
+        generate_3d_model,
+        inputs=[input_images, text_prompt, num_faces],
+        outputs=[preview, download, status]
+    )
+
 demo.queue(max_size=5)
 demo.launch(show_api=True)