Update app.py

app.py

@@ -11,11 +11,11 @@ import io
 import zipfile
 import uuid
 import traceback
-from transformers import AutoImageProcessor, AutoModel
 from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
 import trimesh
+from transformers import pipeline
 
 app = Flask(__name__)
 CORS(app)  # Enable CORS for all routes
@@ -43,8 +43,7 @@ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max
 processing_jobs = {}
 
 # Global model variables
-image_processor = None
-model = None
+depth_estimator = None
 model_loaded = False
 model_loading = False
 
@@ -107,23 +106,23 @@ def preprocess_image(image_path):
         return img
 
 def load_model():
-    global image_processor, model, model_loaded, model_loading
+    global depth_estimator, model_loaded, model_loading
     
     if model_loaded:
-        return image_processor, model
+        return depth_estimator
     
     if model_loading:
         # Wait for model to load if it's already in progress
         while model_loading and not model_loaded:
             time.sleep(0.5)
-        return image_processor, model
+        return depth_estimator
     
     try:
         model_loading = True
         print("Starting model loading...")
         
-        # Using a lightweight model: Pictorial 3D Scene Representation
-        model_name = "damo-vilab/text-to-3d-texture-base" # Smaller model than ShapE-img2img
+        # Using DPT-Hybrid which is smaller than other depth estimation models
+        model_name = "Intel/dpt-hybrid-midas"
         
         # Download model with retry mechanism
         max_retries = 3
@@ -147,24 +146,22 @@ def load_model():
         
         # Initialize model with lower precision to save memory
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        dtype = torch.float16 if device == "cuda" else torch.float32
         
-        image_processor = AutoImageProcessor.from_pretrained(model_name, cache_dir=CACHE_DIR)
-        model = AutoModel.from_pretrained(
-            model_name,
-            torch_dtype=dtype,
-            cache_dir=CACHE_DIR,
-            low_cpu_mem_usage=True,
+        # Load depth estimator pipeline
+        depth_estimator = pipeline(
+            "depth-estimation", 
+            model=model_name,
+            device=device if device == "cuda" else -1,
+            cache_dir=CACHE_DIR
         )
-        model = model.to(device)
         
-        # Optimize for inference
+        # Optimize memory usage
         if device == "cuda":
-            model = model.half()  # Use half precision on GPU
+            torch.cuda.empty_cache()
         
         model_loaded = True
         print(f"Model loaded successfully on {device}")
-        return image_processor, model
+        return depth_estimator
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -173,89 +170,62 @@ def load_model():
     finally:
         model_loading = False
 
-# Convert model output to 3D mesh
-def create_mesh_from_output(output, resolution=64):
-    """Create a mesh from model output"""
-    # Extract features from model output and create mesh
-    # This is a simplified implementation - adapt based on your specific model
-    features = output.last_hidden_state.detach().cpu().numpy()[0]
+# Convert depth map to 3D mesh
+def depth_to_mesh(depth_map, image, resolution=100):
+    """Convert depth map to 3D mesh"""
+    # Get dimensions
+    h, w = depth_map.shape
+    
+    # Create a grid of points
+    x = np.linspace(0, w-1, resolution)
+    y = np.linspace(0, h-1, resolution)
+    x_grid, y_grid = np.meshgrid(x, y)
+    
+    # Sample depth at grid points
+    x_indices = x_grid.astype(int)
+    y_indices = y_grid.astype(int)
+    z_values = depth_map[y_indices, x_indices]
+    
+    # Normalize depth values to suitable range
+    z_min, z_max = z_values.min(), z_values.max()
+    z_values = (z_values - z_min) / (z_max - z_min) * 2.0  # Map to 0-2 range
+    
+    # Normalize x and y coordinates
+    x_grid = (x_grid / w - 0.5) * 2.0  # Map to -1 to 1
+    y_grid = (y_grid / h - 0.5) * 2.0  # Map to -1 to 1
+    
+    # Create vertices
+    vertices = np.vstack([x_grid.flatten(), -y_grid.flatten(), -z_values.flatten()]).T
+    
+    # Create faces (triangles)
+    faces = []
+    for i in range(resolution-1):
+        for j in range(resolution-1):
+            p1 = i * resolution + j
+            p2 = i * resolution + (j + 1)
+            p3 = (i + 1) * resolution + j
+            p4 = (i + 1) * resolution + (j + 1)
+            
+            faces.append([p1, p2, p4])
+            faces.append([p1, p4, p3])
     
-    # Create a simple cube mesh as placeholder - replace with actual mesh generation
-    vertices, faces = create_primitive_mesh(features, resolution)
+    faces = np.array(faces)
     
+    # Create mesh
     mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+    
+    # Optional: Apply texture from original image
+    if image:
+        # This is simplified - proper UV mapping would be needed for accurate texturing
+        pass
+    
     return mesh
 
-def create_primitive_mesh(features, resolution=64):
-    """Create a simple primitive mesh based on features"""
-    # Create a mesh using features as modifiers
-    # This is a simplified implementation - adapt based on your specific model's output
-    
-    # Create a cube/sphere mesh as a placeholder
-    use_sphere = True  # Change to False for cube
-    
-    if use_sphere:
-        # Create a sphere
-        u = np.linspace(0, 2 * np.pi, resolution)
-        v = np.linspace(0, np.pi, resolution)
-        
-        # Base radius and modifiers
-        base_radius = 1.0
-        
-        # Use some features to modify the radius (just as an example)
-        feature_sum = np.sum(features[:10])  # Use first 10 features
-        radius_mod = 0.5 + (feature_sum % 1.0)  # Simple modifier between 0.5 and 1.5
-        
-        # Create vertices
-        x = base_radius * radius_mod * np.outer(np.cos(u), np.sin(v))
-        y = base_radius * radius_mod * np.outer(np.sin(u), np.sin(v))
-        z = base_radius * radius_mod * np.outer(np.ones_like(u), np.cos(v))
-        
-        # Reshape to get list of vertices
-        vertices = np.vstack([x.flatten(), y.flatten(), z.flatten()]).T
-        
-        # Create faces (triangles)
-        faces = []
-        for i in range(resolution-1):
-            for j in range(resolution-1):
-                p1 = i * resolution + j
-                p2 = i * resolution + (j + 1)
-                p3 = (i + 1) * resolution + j
-                p4 = (i + 1) * resolution + (j + 1)
-                
-                faces.append([p1, p2, p4])
-                faces.append([p1, p4, p3])
-        
-        faces = np.array(faces)
-    else:
-        # Create a cube
-        vertices = np.array([
-            [-1, -1, -1], [1, -1, -1], [1, 1, -1], [-1, 1, -1],
-            [-1, -1, 1], [1, -1, 1], [1, 1, 1], [-1, 1, 1]
-        ])
-        
-        # Apply some feature-based modifications
-        feature_sum = np.sum(features[:10])  # Use first 10 features
-        scale_factor = 0.5 + (feature_sum % 1.0)  # Simple modifier between 0.5 and 1.5
-        vertices *= scale_factor
-        
-        # Faces (triangles)
-        faces = np.array([
-            [0, 1, 2], [0, 2, 3],  # Bottom face
-            [4, 5, 6], [4, 6, 7],  # Top face
-            [0, 1, 5], [0, 5, 4],  # Front face
-            [2, 3, 7], [2, 7, 6],  # Back face
-            [0, 3, 7], [0, 7, 4],  # Left face
-            [1, 2, 6], [1, 6, 5]   # Right face
-        ])
-    
-    return vertices, faces
-
 @app.route('/health', methods=['GET'])
 def health_check():
     return jsonify({
         "status": "healthy", 
-        "model": "Lightweight 3D Model Generator",
+        "model": "Depth-Based 3D Model Generator",
         "device": "cuda" if torch.cuda.is_available() else "cpu"
     }), 200
 
@@ -313,15 +283,11 @@ def convert_image_to_3d():
     
     # Get optional parameters with defaults
     try:
-        guidance_scale = float(request.form.get('guidance_scale', 3.0))
+        mesh_resolution = min(int(request.form.get('mesh_resolution', 100)), 200)  # Limit max resolution
         output_format = request.form.get('output_format', 'obj').lower()
     except ValueError:
         return jsonify({"error": "Invalid parameter values"}), 400
     
-    # Validate parameters
-    if guidance_scale < 1.0 or guidance_scale > 5.0:
-        return jsonify({"error": "Guidance scale must be between 1.0 and 5.0"}), 400
-    
     # Validate output format
     if output_format not in ['obj', 'glb']:
         return jsonify({"error": "Unsupported output format. Use 'obj' or 'glb'"}), 400
@@ -360,7 +326,7 @@ def convert_image_to_3d():
             
             # Load model
             try:
-                processor, model_instance = load_model()
+                model = load_model()
                 processing_jobs[job_id]['progress'] = 30
             except Exception as e:
                 processing_jobs[job_id]['status'] = 'error'
@@ -369,18 +335,20 @@ def convert_image_to_3d():
             
             # Process image with thread-safe timeout
             try:
-                def generate_3d():
-                    # Process the image
-                    device = model_instance.device
-                    inputs = processor(images=image, return_tensors="pt").to(device)
+                def estimate_depth():
+                    # Get depth map
+                    result = model(image)
+                    depth_map = result["depth"]
                     
-                    # Forward pass through model
-                    with torch.no_grad():
-                        outputs = model_instance(**inputs)
+                    # Convert to numpy array if needed
+                    if isinstance(depth_map, torch.Tensor):
+                        depth_map = depth_map.cpu().numpy()
+                    elif hasattr(depth_map, 'numpy'):
+                        depth_map = depth_map.numpy()
                     
-                    return outputs
+                    return depth_map
                 
-                outputs, error = process_with_timeout(generate_3d, [], TIMEOUT_SECONDS)
+                depth_map, error = process_with_timeout(estimate_depth, [], TIMEOUT_SECONDS)
                 
                 if error:
                     if isinstance(error, TimeoutError):
@@ -390,10 +358,11 @@ def convert_image_to_3d():
                     else:
                         raise error
                         
-                processing_jobs[job_id]['progress'] = 80
+                processing_jobs[job_id]['progress'] = 60
                 
-                # Create mesh from outputs
-                mesh = create_mesh_from_output(outputs)
+                # Create mesh from depth map
+                mesh = depth_to_mesh(depth_map, image, resolution=mesh_resolution)
+                processing_jobs[job_id]['progress'] = 80
                 
             except Exception as e:
                 error_details = traceback.format_exc()