Update app.py

app.py

@@ -14,11 +14,12 @@ from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
-from diffusers import DiffusionPipeline
-import cv2
+from trellis.pipelines import TrellisImageTo3DPipeline
 
 # Force CPU usage
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
+os.environ["ATTN_BACKEND"] = "native"  # Disable xformers/flash-attn
+os.environ["SPCONV_ALGO"] = "native"  # Optimize for CPU
 torch.set_default_device("cpu")
 torch.cuda.is_available = lambda: False
 torch.cuda.device_count = lambda: 0
@@ -49,13 +50,13 @@ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 processing_jobs = {}
 
 # Global model
-tripo_pipeline = None
+trellis_pipeline = None
 model_loaded = False
 model_loading = False
 
 # Configuration
-TIMEOUT_SECONDS = 300  # 5 minutes for TripoSG
-MAX_DIMENSION = 256  # TripoSG works with smaller images
+TIMEOUT_SECONDS = 360  # 6 minutes for TRELLIS
+MAX_DIMENSION = 256  # TRELLIS works with smaller images
 
 class TimeoutError(Exception):
     pass
@@ -100,33 +101,26 @@ def preprocess_image(image_path):
                 img = img.convert('RGB')
             # Resize to 256x256
             img = img.resize((256, 256), Image.LANCZOS)
-            
-            # Basic cv2 cleanup
-            img_array = np.array(img)
-            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
-            _, mask = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)
-            img_array = cv2.bitwise_and(img_array, img_array, mask=mask)
-            
-            return Image.fromarray(img_array)
+            return img
     except Exception as e:
         raise Exception(f"Error preprocessing image: {str(e)}")
 
 def load_model():
-    global tripo_pipeline, model_loaded, model_loading
+    global trellis_pipeline, model_loaded, model_loading
     
     if model_loaded:
-        return tripo_pipeline
+        return trellis_pipeline
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return tripo_pipeline
+        return trellis_pipeline
     
     try:
         model_loading = True
-        print("Loading TripoSG model...")
+        print("Loading TRELLIS-image-large...")
         
-        model_name = "tripo3d/tripo-sg-3d"
+        model_name = "JeffreyXiang/TRELLIS-image-large"
         
         # Download model
         max_retries = 3
@@ -147,17 +141,17 @@ def load_model():
                 else:
                     raise
         
-        # Load TripoSG pipeline
-        tripo_pipeline = DiffusionPipeline.from_pretrained(
+        # Load TRELLIS pipeline
+        trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
             model_name,
             cache_dir=CACHE_DIR,
             torch_dtype=torch.float32,
         )
-        tripo_pipeline.to("cpu")
+        trellis_pipeline.to("cpu")
         
         model_loaded = True
-        print("TripoSG loaded successfully on CPU")
-        return tripo_pipeline
+        print("TRELLIS loaded successfully on CPU")
+        return trellis_pipeline
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -169,20 +163,17 @@ def load_model():
 def generate_3d_model(image, detail_level):
     try:
         # Parameters
-        num_steps = {'low': 20, 'medium': 30, 'high': 40}
+        num_steps = {'low': 50, 'medium': 75, 'high': 100}
         steps = num_steps[detail_level]
         
         # Generate 3D model
         with torch.no_grad():
-            result = tripo_pipeline(image, num_inference_steps=steps)
+            result = trellis_pipeline(image, num_inference_steps=steps, output_type="mesh")
         
         # Extract mesh
-        mesh = result.meshes[0]
-        
-        # Convert to trimesh
-        vertices = np.array(mesh.vertices)
-        faces = np.array(mesh.faces)
-        vertex_colors = np.array(mesh.vertex_colors) if mesh.vertex_colors is not None else None
+        vertices = np.array(result.vertices)
+        faces = np.array(result.faces)
+        vertex_colors = np.array(result.vertex_colors) if result.vertex_colors is not None else None
         
         trimesh_mesh = trimesh.Trimesh(
             vertices=vertices,
@@ -201,7 +192,7 @@ def generate_3d_model(image, detail_level):
 def health_check():
     return jsonify({
         "status": "healthy",
-        "model": "TripoSG",
+        "model": "TRELLIS-image-large",
         "device": "cpu"
     }), 200
 
@@ -442,7 +433,7 @@ def model_info(job_id):
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API (TripoSG)",
+        "message": "Image to 3D API (TRELLIS-image-large)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -454,7 +445,7 @@ def index():
             "output_format": "glb or obj",
             "detail_level": "low, medium, or high - controls inference steps"
         },
-        "description": "Creates 3D models from 2D images using TripoSG. Use transparent PNGs for best results."
+        "description": "Creates 3D models from 2D images using TRELLIS-image-large. Use transparent PNGs for best results."
     }), 200
 
 if __name__ == '__main__':