Update app.py

app.py

@@ -16,8 +16,7 @@ from flask_cors import CORS
 import numpy as np
 import trimesh
 import cv2
-from tsr.models import TripoSR  # Custom TripoSR model
-from tsr.pipeline import TripoSRPipeline  # Custom pipeline
+from tsr.system import TSR  # Updated import
 import torchvision.transforms as T
 
 app = Flask(__name__)
@@ -47,7 +46,7 @@ processing_jobs = {}
 
 # Global model variables
 u2net_model = None
-triposr_pipeline = None
+triposr_model = None
 model_loaded = False
 model_loading = False
 
@@ -142,15 +141,15 @@ def remove_background(image):
     return Image.fromarray(result.astype('uint8'))
 
 def load_model():
-    global triposr_pipeline, model_loaded, model_loading
+    global triposr_model, model_loaded, model_loading
     
     if model_loaded:
-        return triposr_pipeline
+        return triposr_model
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return triposr_pipeline
+        return triposr_model
     
     try:
         model_loading = True
@@ -176,15 +175,17 @@ def load_model():
                 else:
                     raise
         
-        # Initialize TripoSR pipeline
-        triposr_pipeline = TripoSRPipeline(
-            model_path=os.path.join(CACHE_DIR, "stabilityai/TripoSR"),
-            device='cpu'
+        # Initialize TSR model
+        triposr_model = TSR.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,
+            device="cpu",
+            cache_dir=CACHE_DIR
         )
         
         model_loaded = True
-        print("TripoSR pipeline loaded successfully on CPU")
-        return triposr_pipeline
+        print("TripoSR model loaded successfully on CPU")
+        return triposr_model
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -304,9 +305,9 @@ def convert_image_to_3d():
             clean_image = remove_background(image)
             processing_jobs[job_id]['progress'] = 30
             
-            # Load TripoSR pipeline
+            # Load TripoSR model
             try:
-                pipeline = load_model()
+                model = load_model()
                 processing_jobs[job_id]['progress'] = 40
             except Exception as e:
                 processing_jobs[job_id]['status'] = 'error'
@@ -316,8 +317,8 @@ def convert_image_to_3d():
             # Generate 3D model
             try:
                 def generate_3d():
-                    # TripoSR pipeline expects a PIL image
-                    mesh = pipeline(clean_image)
+                    # TSR expects a PIL image
+                    mesh = model(clean_image)
                     return mesh
                 
                 mesh, error = process_with_timeout(generate_3d, [], TIMEOUT_SECONDS)