Update app.py

app.py

@@ -14,8 +14,18 @@ from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
-from triposr import TripoSR
 from rembg import remove
+import sys
+
+# Add SF3D code to sys.path
+sys.path.append('/app/stable-fast-3d')
+
+# Import SF3D components (adjust based on actual module structure)
+try:
+    from sf3d.system import SF3D
+except ImportError as e:
+    print(f"Failed to import SF3D: {e}")
+    raise
 
 # Force CPU usage
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -50,13 +60,13 @@ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 processing_jobs = {}
 
 # Global model variables
-triposr = None
+sf3d_model = None
 model_loaded = False
 model_loading = False
 
 # Configuration for processing
-TIMEOUT_SECONDS = 180  # 3 minutes max for TripoSR on CPU
-MAX_DIMENSION = 256
+TIMEOUT_SECONDS = 240  # 4 minutes max for SF3D on CPU
+MAX_DIMENSION = 512  # SF3D expects 512x512
 
 # TimeoutError for handling timeouts
 class TimeoutError(Exception):
@@ -95,36 +105,30 @@ def process_with_timeout(function, args, timeout):
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
-# Simplified image preprocessing
+# Image preprocessing for SF3D (512x512)
 def preprocess_image(image_path):
     with Image.open(image_path) as img:
         img = img.convert("RGB")
-        if img.width > MAX_DIMENSION or img.height > MAX_DIMENSION:
-            if img.width > img.height:
-                new_width = MAX_DIMENSION
-                new_height = int(img.height * (MAX_DIMENSION / img.width))
-            else:
-                new_height = MAX_DIMENSION
-                new_width = int(img.width * (MAX_DIMENSION / img.height))
-            img = img.resize((new_width, new_height), Image.LANCZOS)
+        # SF3D requires 512x512
+        img = img.resize((512, 512), Image.LANCZOS)
         return img
 
 def load_model():
-    global triposr, model_loaded, model_loading
+    global sf3d_model, model_loaded, model_loading
     
     if model_loaded:
-        return triposr
+        return sf3d_model
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return triposr
+        return sf3d_model
     
     try:
         model_loading = True
         print("Starting model loading...")
         
-        model_name = "VAST-AI-Research/TripoSR"
+        model_name = "stabilityai/stable-fast-3d"
         
         # Download model with retry mechanism
         max_retries = 3
@@ -145,19 +149,17 @@ def load_model():
                 else:
                     raise
         
-        # Load TripoSR
-        triposr = TripoSR.from_pretrained(
+        # Load SF3D model
+        sf3d_model = SF3D.from_pretrained(
             model_name,
-            use_safetensors=True,
+            cache_dir=CACHE_DIR,
             torch_dtype=torch.float16,
-            cache_dir=CACHE_DIR
+            device_map="cpu"
         )
-        # Explicitly move to CPU
-        triposr.to("cpu")
         
         model_loaded = True
         print("Model loaded successfully on CPU")
-        return triposr
+        return sf3d_model
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -170,7 +172,7 @@ def load_model():
 def health_check():
     return jsonify({
         "status": "healthy",
-        "model": "TripoSR 3D Generator",
+        "model": "Stable Fast 3D (SF3D)",
         "device": "cpu"
     }), 200
 
@@ -268,8 +270,8 @@ def convert_image_to_3d():
             try:
                 def generate_3d():
                     # Adjust settings based on detail level
-                    foreground_ratios = {'low': 0.6, 'medium': 0.8, 'high': 1.0}
-                    num_samples = {'low': 5000, 'medium': 10000, 'high': 20000}
+                    texture_sizes = {'low': 512, 'medium': 1024, 'high': 2048}
+                    remesh_options = {'low': 0.5, 'medium': 1.0, 'high': 2.0}
                     
                     # Remove background
                     image_no_bg = remove(image)
@@ -280,8 +282,8 @@ def convert_image_to_3d():
                     # Generate mesh
                     mesh = model(
                         image=img_array,
-                        foreground_ratio=foreground_ratios[detail_level],
-                        num_samples=num_samples[detail_level],
+                        texture_size=texture_sizes[detail_level],
+                        remesh_option=remesh_options[detail_level],
                         seed=12345
                     )
                     return mesh
@@ -422,7 +424,7 @@ def model_info(job_id):
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API (TripoSR)",
+        "message": "Image to 3D API (Stable Fast 3D)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -432,9 +434,9 @@ def index():
         ],
         "parameters": {
             "output_format": "glb",
-            "detail_level": "low, medium, or high - controls mesh detail"
+            "detail_level": "low, medium, or high - controls texture and mesh detail"
         },
-        "description": "This API creates full 3D models from 2D images using TripoSR"
+        "description": "This API creates full 3D models from 2D images using Stable Fast 3D"
     }), 200
 
 if __name__ == '__main__':