Update app.py

app.py

@@ -16,8 +16,8 @@ from flask_cors import CORS
 import numpy as np
 import trimesh
 import cv2
-from transformers import AutoModel, AutoProcessor  # For TripoSR
-from u2net import U2NET  # For background removal; install from https://github.com/xuebinqin/U-2-Net
+from tsr.models import TripoSR  # Custom TripoSR model
+from tsr.pipeline import TripoSRPipeline  # Custom pipeline
 import torchvision.transforms as T
 
 app = Flask(__name__)
@@ -47,8 +47,7 @@ processing_jobs = {}
 
 # Global model variables
 u2net_model = None
-triposr_model = None
-triposr_processor = None
+triposr_pipeline = None
 model_loaded = False
 model_loading = False
 
@@ -121,6 +120,8 @@ def preprocess_image(image_path):
 def remove_background(image):
     global u2net_model
     if u2net_model is None:
+        # Dynamically import U2NET to avoid circular import issues
+        from u2net import U2NET
         u2net_model = U2NET()
         u2net_model.load_state_dict(torch.load('u2net.pth', map_location='cpu'))
         u2net_model.eval()
@@ -141,15 +142,15 @@ def remove_background(image):
     return Image.fromarray(result.astype('uint8'))
 
 def load_model():
-    global triposr_model, triposr_processor, model_loaded, model_loading
+    global triposr_pipeline, model_loaded, model_loading
     
     if model_loaded:
-        return triposr_model, triposr_processor
+        return triposr_pipeline
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return triposr_model, triposr_processor
+        return triposr_pipeline
     
     try:
         model_loading = True
@@ -175,13 +176,15 @@ def load_model():
                 else:
                     raise
         
-        triposr_processor = AutoProcessor.from_pretrained(model_name, cache_dir=CACHE_DIR)
-        triposr_model = AutoModel.from_pretrained(model_name, cache_dir=CACHE_DIR)
-        triposr_model.to('cpu')
+        # Initialize TripoSR pipeline
+        triposr_pipeline = TripoSRPipeline(
+            model_path=os.path.join(CACHE_DIR, "stabilityai/TripoSR"),
+            device='cpu'
+        )
         
         model_loaded = True
-        print("TripoSR model loaded successfully on CPU")
-        return triposr_model, triposr_processor
+        print("TripoSR pipeline loaded successfully on CPU")
+        return triposr_pipeline
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -301,9 +304,9 @@ def convert_image_to_3d():
             clean_image = remove_background(image)
             processing_jobs[job_id]['progress'] = 30
             
-            # Load TripoSR model
+            # Load TripoSR pipeline
             try:
-                model, processor = load_model()
+                pipeline = load_model()
                 processing_jobs[job_id]['progress'] = 40
             except Exception as e:
                 processing_jobs[job_id]['status'] = 'error'
@@ -313,10 +316,8 @@ def convert_image_to_3d():
             # Generate 3D model
             try:
                 def generate_3d():
-                    inputs = processor(images=clean_image, return_tensors="pt").to('cpu')
-                    with torch.no_grad():
-                        outputs = model(**inputs)
-                    mesh = outputs.mesh  # TripoSR outputs a trimesh object
+                    # TripoSR pipeline expects a PIL image
+                    mesh = pipeline(clean_image)
                     return mesh
                 
                 mesh, error = process_with_timeout(generate_3d, [], TIMEOUT_SECONDS)