Update app.py

app.py

@@ -14,11 +14,15 @@ from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
-from lgm.pipeline import LGMPipeline
+from triposr import TripoSR
+from rembg import remove
 
 # Force CPU usage
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 torch.set_default_device("cpu")
+# Patch PyTorch to disable CUDA initialization
+torch.cuda.is_available = lambda: False
+torch.cuda.device_count = lambda: 0
 
 app = Flask(__name__)
 CORS(app)
@@ -46,12 +50,12 @@ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 processing_jobs = {}
 
 # Global model variables
-lgm_pipeline = None
+triposr = None
 model_loaded = False
 model_loading = False
 
 # Configuration for processing
-TIMEOUT_SECONDS = 240  # 4 minutes max for LGM on CPU
+TIMEOUT_SECONDS = 180  # 3 minutes max for TripoSR on CPU
 MAX_DIMENSION = 256
 
 # TimeoutError for handling timeouts
@@ -106,21 +110,21 @@ def preprocess_image(image_path):
         return img
 
 def load_model():
-    global lgm_pipeline, model_loaded, model_loading
+    global triposr, model_loaded, model_loading
     
     if model_loaded:
-        return lgm_pipeline
+        return triposr
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return lgm_pipeline
+        return triposr
     
     try:
         model_loading = True
         print("Starting model loading...")
         
-        model_name = "open-mmlab/LGM"
+        model_name = "VAST-AI-Research/TripoSR"
         
         # Download model with retry mechanism
         max_retries = 3
@@ -141,18 +145,19 @@ def load_model():
                 else:
                     raise
         
-        # Load LGM pipeline
-        lgm_pipeline = LGMPipeline.from_pretrained(
+        # Load TripoSR
+        triposr = TripoSR.from_pretrained(
             model_name,
             use_safetensors=True,
             torch_dtype=torch.float16,
-            cache_dir=CACHE_DIR,
-            device_map="cpu"
+            cache_dir=CACHE_DIR
         )
+        # Explicitly move to CPU
+        triposr.to("cpu")
         
         model_loaded = True
         print("Model loaded successfully on CPU")
-        return lgm_pipeline
+        return triposr
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -165,7 +170,7 @@ def load_model():
 def health_check():
     return jsonify({
         "status": "healthy",
-        "model": "LGM 3D Generator",
+        "model": "TripoSR 3D Generator",
         "device": "cpu"
     }), 200
 
@@ -263,19 +268,21 @@ def convert_image_to_3d():
             try:
                 def generate_3d():
                     # Adjust settings based on detail level
-                    steps = {'low': 40, 'medium': 60, 'high': 80}
-                    resolution = {'low': 512, 'medium': 1024, 'high': 2048}
+                    foreground_ratios = {'low': 0.6, 'medium': 0.8, 'high': 1.0}
+                    num_samples = {'low': 5000, 'medium': 10000, 'high': 20000}
                     
-                    # Convert PIL image to numpy
-                    img_array = np.array(image)
+                    # Remove background
+                    image_no_bg = remove(image)
+                    
+                    # Convert to numpy
+                    img_array = np.array(image_no_bg)
                     
                     # Generate mesh
                     mesh = model(
                         image=img_array,
-                        num_inference_steps=steps[detail_level],
-                        texture_resolution=resolution[detail_level],
-                        generator=torch.manual_seed(12345),
-                        output_type="trimesh"
+                        foreground_ratio=foreground_ratios[detail_level],
+                        num_samples=num_samples[detail_level],
+                        seed=12345
                     )
                     return mesh
                 
@@ -415,7 +422,7 @@ def model_info(job_id):
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API (LGM)",
+        "message": "Image to 3D API (TripoSR)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -427,7 +434,7 @@ def index():
             "output_format": "glb",
             "detail_level": "low, medium, or high - controls mesh detail"
         },
-        "description": "This API creates full 3D models from 2D images using LGM"
+        "description": "This API creates full 3D models from 2D images using TripoSR"
     }), 200
 
 if __name__ == '__main__':