Update app.py

app.py

@@ -14,7 +14,7 @@ from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
-from sf3d.system import StableFast3D
+from lgm.pipeline import LGMPipeline
 
 # Force CPU usage
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -46,12 +46,12 @@ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 processing_jobs = {}
 
 # Global model variables
-sf3d = None
+lgm_pipeline = None
 model_loaded = False
 model_loading = False
 
 # Configuration for processing
-TIMEOUT_SECONDS = 180  # 3 minutes max for SF3D on CPU
+TIMEOUT_SECONDS = 240  # 4 minutes max for LGM on CPU
 MAX_DIMENSION = 256
 
 # TimeoutError for handling timeouts
@@ -106,21 +106,21 @@ def preprocess_image(image_path):
         return img
 
 def load_model():
-    global sf3d, model_loaded, model_loading
+    global lgm_pipeline, model_loaded, model_loading
     
     if model_loaded:
-        return sf3d
+        return lgm_pipeline
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return sf3d
+        return lgm_pipeline
     
     try:
         model_loading = True
         print("Starting model loading...")
         
-        model_name = "stabilityai/stable-fast-3d"
+        model_name = "open-mmlab/LGM"
         
         # Download model with retry mechanism
         max_retries = 3
@@ -141,18 +141,18 @@ def load_model():
                 else:
                     raise
         
-        # Load Stable Fast 3D
-        sf3d = StableFast3D.from_pretrained(
+        # Load LGM pipeline
+        lgm_pipeline = LGMPipeline.from_pretrained(
             model_name,
             use_safetensors=True,
             torch_dtype=torch.float16,
             cache_dir=CACHE_DIR,
-            device="cpu"
+            device_map="cpu"
         )
         
         model_loaded = True
         print("Model loaded successfully on CPU")
-        return sf3d
+        return lgm_pipeline
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -165,7 +165,7 @@ def load_model():
 def health_check():
     return jsonify({
         "status": "healthy",
-        "model": "Stable Fast 3D Generator",
+        "model": "LGM 3D Generator",
         "device": "cpu"
     }), 200
 
@@ -263,8 +263,8 @@ def convert_image_to_3d():
             try:
                 def generate_3d():
                     # Adjust settings based on detail level
-                    steps = {'low': 30, 'medium': 50, 'high': 75}
-                    texture_res = {'low': 512, 'medium': 1024, 'high': 2048}
+                    steps = {'low': 40, 'medium': 60, 'high': 80}
+                    resolution = {'low': 512, 'medium': 1024, 'high': 2048}
                     
                     # Convert PIL image to numpy
                     img_array = np.array(image)
@@ -272,9 +272,10 @@ def convert_image_to_3d():
                     # Generate mesh
                     mesh = model(
                         image=img_array,
-                        num_steps=steps[detail_level],
-                        texture_resolution=texture_res[detail_level],
-                        seed=12345
+                        num_inference_steps=steps[detail_level],
+                        texture_resolution=resolution[detail_level],
+                        generator=torch.manual_seed(12345),
+                        output_type="trimesh"
                     )
                     return mesh
                 
@@ -414,7 +415,7 @@ def model_info(job_id):
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API (Stable Fast 3D)",
+        "message": "Image to 3D API (LGM)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -426,7 +427,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 Stable Fast 3D"
+        "description": "This API creates full 3D models from 2D images using LGM"
     }), 200
 
 if __name__ == '__main__':