Update app.py

app.py

@@ -14,17 +14,7 @@ from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
-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
+from diffusers import DiffusionPipeline
 
 # Force CPU usage
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -59,13 +49,13 @@ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
 processing_jobs = {}
 
 # Global model variables
-sf3d_model = None
+zero123plus_pipeline = None
 model_loaded = False
 model_loading = False
 
 # Configuration for processing
-TIMEOUT_SECONDS = 240  # 4 minutes max for SF3D on CPU
-MAX_DIMENSION = 512  # SF3D expects 512x512
+TIMEOUT_SECONDS = 240  # 4 minutes max for Zero123++ on CPU
+MAX_DIMENSION = 512  # Zero123++ expects 512x512
 
 # TimeoutError for handling timeouts
 class TimeoutError(Exception):
@@ -104,30 +94,30 @@ def process_with_timeout(function, args, timeout):
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
-# Image preprocessing for SF3D (512x512, no background removal)
+# Image preprocessing for Zero123++ (512x512, no background removal)
 def preprocess_image(image_path):
     with Image.open(image_path) as img:
         img = img.convert("RGB")
-        # SF3D requires 512x512
+        # Zero123++ requires 512x512
         img = img.resize((512, 512), Image.LANCZOS)
         return img
 
 def load_model():
-    global sf3d_model, model_loaded, model_loading
+    global zero123plus_pipeline, model_loaded, model_loading
     
     if model_loaded:
-        return sf3d_model
+        return zero123plus_pipeline
     
     if model_loading:
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return sf3d_model
+        return zero123plus_pipeline
     
     try:
         model_loading = True
         print("Starting model loading...")
         
-        model_name = "stabilityai/stable-fast-3d"
+        model_name = "sudo-ai/zero123plus-v1.2"
         
         # Download model with retry mechanism
         max_retries = 3
@@ -148,8 +138,8 @@ def load_model():
                 else:
                     raise
         
-        # Load SF3D model
-        sf3d_model = SF3D.from_pretrained(
+        # Load Zero123++ pipeline
+        zero123plus_pipeline = DiffusionPipeline.from_pretrained(
             model_name,
             cache_dir=CACHE_DIR,
             torch_dtype=torch.float16,
@@ -158,7 +148,7 @@ def load_model():
         
         model_loaded = True
         print("Model loaded successfully on CPU")
-        return sf3d_model
+        return zero123plus_pipeline
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -171,7 +161,7 @@ def load_model():
 def health_check():
     return jsonify({
         "status": "healthy",
-        "model": "Stable Fast 3D (SF3D)",
+        "model": "Zero123++",
         "device": "cpu"
     }), 200
 
@@ -259,7 +249,7 @@ def convert_image_to_3d():
             processing_jobs[job_id]['progress'] = 10
             
             try:
-                model = load_model()
+                pipeline = load_model()
                 processing_jobs[job_id]['progress'] = 30
             except Exception as e:
                 processing_jobs[job_id]['status'] = 'error'
@@ -269,19 +259,22 @@ def convert_image_to_3d():
             try:
                 def generate_3d():
                     # Adjust settings based on detail level
-                    texture_sizes = {'low': 512, 'medium': 1024, 'high': 2048}
-                    remesh_options = {'low': 0.5, 'medium': 1.0, 'high': 2.0}
+                    num_steps = {'low': 20, 'medium': 50, 'high': 75}
                     
-                    # Convert image to numpy (no background removal)
-                    img_array = np.array(image)
+                    # Convert image to tensor
+                    img_array = np.array(image) / 255.0
+                    img_tensor = torch.from_numpy(img_array).permute(2, 0, 1).float()
                     
-                    # Generate mesh
-                    mesh = model(
-                        image=img_array,
-                        texture_size=texture_sizes[detail_level],
-                        remesh_option=remesh_options[detail_level],
+                    # Generate multi-view images and reconstruct mesh
+                    pipeline_output = pipeline(
+                        img_tensor.unsqueeze(0),
+                        num_inference_steps=num_steps[detail_level],
+                        guidance_scale=7.5,
                         seed=12345
                     )
+                    
+                    # Extract mesh (Zero123++ outputs mesh directly)
+                    mesh = pipeline_output.meshes[0]
                     return mesh
                 
                 mesh, error = process_with_timeout(generate_3d, [], TIMEOUT_SECONDS)
@@ -420,7 +413,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 (Zero123++)",
         "endpoints": [
             "/convert",
             "/progress/<job_id>",
@@ -430,9 +423,9 @@ def index():
         ],
         "parameters": {
             "output_format": "glb",
-            "detail_level": "low, medium, or high - controls texture and mesh detail"
+            "detail_level": "low, medium, or high - controls inference steps"
         },
-        "description": "This API creates full 3D models from 2D images using Stable Fast 3D. Images should have transparent backgrounds."
+        "description": "This API creates full 3D models from 2D images using Zero123++. Images should have transparent backgrounds."
     }), 200
 
 if __name__ == '__main__':