Update app.py

app.py

@@ -11,11 +11,11 @@ import io
 import zipfile
 import uuid
 import traceback
+from diffusers import ShapEImg2ImgPipeline
+from diffusers.utils import export_to_obj
 from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import functools
-import numpy as np
-import trimesh
 
 app = Flask(__name__)
 CORS(app)  # Enable CORS for all routes
@@ -43,14 +43,14 @@ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max
 processing_jobs = {}
 
 # Global model variable
-neus_model = None
+pipe = None
 model_loaded = False
 model_loading = False
 
 # Configuration for processing
-TIMEOUT_SECONDS = 180  # 3 minutes max for processing
+TIMEOUT_SECONDS = 300  # 5 minutes max for processing
 MAX_DIMENSION = 512    # Max image dimension to process
-MAX_INFERENCE_STEPS = 32  # Maximum allowed inference steps
+MAX_INFERENCE_STEPS = 64  # Maximum allowed inference steps to prevent the index error
 
 # TimeoutError for handling timeouts
 class TimeoutError(Exception):
@@ -104,310 +104,65 @@ def preprocess_image(image_path):
                 new_width = int(img.width * (MAX_DIMENSION / img.height))
             img = img.resize((new_width, new_height), Image.LANCZOS)
         
-        # Convert to RGB and convert to tensor
-        img_array = np.array(img) / 255.0  # Normalize to [0, 1]
-        img_tensor = torch.from_numpy(img_array).float().permute(2, 0, 1).unsqueeze(0)  # [1, 3, H, W]
-        return img_tensor
-
-# Simple NeuS2-inspired implementation for reconstructing 3D surfaces from images
-class NeuS2Model:
-    def __init__(self, device="cuda" if torch.cuda.is_available() else "cpu"):
-        self.device = device
-        self.encoder = self._create_encoder().to(device)
-        self.volume_network = self._create_volume_network().to(device)
-        
-    def _create_encoder(self):
-        # Simple convolutional encoder
-        return torch.nn.Sequential(
-            torch.nn.Conv2d(3, 32, 3, stride=2, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(32, 64, 3, stride=2, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(64, 128, 3, stride=2, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.AdaptiveAvgPool2d((8, 8)),
-            torch.nn.Flatten(),
-            torch.nn.Linear(8192, 512)
-        )
-    
-    def _create_volume_network(self):
-        # MLP to predict occupancy and SDF values
-        return torch.nn.Sequential(
-            torch.nn.Linear(515, 256),  # 512 features + 3 coordinates
-            torch.nn.ReLU(),
-            torch.nn.Linear(256, 256),
-            torch.nn.ReLU(),
-            torch.nn.Linear(256, 1)  # SDF value
-        )
-    
-    def extract_features(self, image):
-        with torch.no_grad():
-            return self.encoder(image.to(self.device))
-    
-    def query_points(self, points, features):
-        # points shape: [batch, num_points, 3]
-        # features shape: [batch, 512]
-        batch_size, num_points, _ = points.shape
-        
-        # Expand features to match points
-        features = features.unsqueeze(1).expand(-1, num_points, -1)  # [batch, num_points, 512]
-        
-        # Concatenate points with features
-        points_features = torch.cat([points, features], dim=-1)  # [batch, num_points, 515]
-        points_features = points_features.reshape(-1, 515)  # [batch*num_points, 515]
-        
-        # Query network
-        with torch.no_grad():
-            sdf = self.volume_network(points_features.to(self.device))
-        
-        return sdf.reshape(batch_size, num_points, 1)
-    
-    def generate_mesh(self, image, resolution=64, threshold=0.0, num_steps=16):
-        # Extract image features
-        features = self.extract_features(image)  # [1, 512]
-        
-        # Create grid points
-        x = torch.linspace(-1, 1, resolution)
-        y = torch.linspace(-1, 1, resolution)
-        z = torch.linspace(-1, 1, resolution)
-        grid_x, grid_y, grid_z = torch.meshgrid(x, y, z, indexing='ij')
-        points = torch.stack([grid_x, grid_y, grid_z], dim=-1).reshape(1, -1, 3)  # [1, res^3, 3]
-        
-        # Process in batches to avoid OOM
-        batch_size = 32768  # Adjust based on available memory
-        sdf_values = []
-        
-        for i in range(0, points.shape[1], batch_size):
-            batch_points = points[:, i:i+batch_size]
-            batch_sdf = self.query_points(batch_points, features)
-            sdf_values.append(batch_sdf)
-        
-        sdf_volume = torch.cat(sdf_values, dim=1).reshape(resolution, resolution, resolution).cpu().numpy()
-        
-        # Extract mesh - alternative to marching cubes since we don't have scipy
-        vertices, faces = self._simple_mesh_extraction(sdf_volume, threshold)
-        
-        # Create a mesh object with vertices and faces
-        mesh = type('Mesh', (), {'verts': vertices, 'faces': faces})
-        
-        return [mesh]  # Returning in list format to match ShapE's output format
-    
-    def _simple_mesh_extraction(self, sdf_volume, threshold=0.0):
-        """Simple mesh extraction without scipy dependency"""
-        resolution = sdf_volume.shape[0]
-        
-        # Find surface points (approximate)
-        surface_points = []
-        surface_normals = []
-        
-        # Sample points on three orthogonal grids
-        for axis in range(3):
-            for i in range(resolution):
-                for j in range(resolution):
-                    # Create a line along the current axis
-                    line = np.zeros((resolution, 3), dtype=int)
-                    for k in range(resolution):
-                        if axis == 0:
-                            line[k] = [k, i, j]
-                        elif axis == 1:
-                            line[k] = [i, k, j]
-                        else:
-                            line[k] = [i, j, k]
-                    
-                    # Get SDF values along this line
-                    sdf_line = np.array([sdf_volume[tuple(idx)] for idx in line])
-                    
-                    # Find zero crossings
-                    signs = np.sign(sdf_line)
-                    zero_crossings = np.where(np.diff(signs) != 0)[0]
-                    
-                    for idx in zero_crossings:
-                        # Linear interpolation to find more accurate zero crossing
-                        t = sdf_line[idx] / (sdf_line[idx] - sdf_line[idx + 1])
-                        point = line[idx] * (1 - t) + line[idx + 1] * t
-                        
-                        # Normalize to [-1, 1] range
-                        normalized_point = 2 * point / resolution - 1
-                        surface_points.append(normalized_point)
-                        
-                        # Compute normal (gradient of SDF)
-                        normal = np.zeros(3)
-                        idx_3d = tuple(np.round(point).astype(int).clip(0, resolution - 1))
-                        
-                        # Compute gradient using central differences where possible
-                        for d in range(3):
-                            if 0 < idx_3d[d] < resolution - 1:
-                                idx_plus = list(idx_3d)
-                                idx_minus = list(idx_3d)
-                                idx_plus[d] += 1
-                                idx_minus[d] -= 1
-                                normal[d] = (sdf_volume[tuple(idx_plus)] - sdf_volume[tuple(idx_minus)]) / 2
-                            else:
-                                # Forward or backward difference at boundaries
-                                idx_curr = list(idx_3d)
-                                if idx_3d[d] == 0:
-                                    idx_other = list(idx_3d)
-                                    idx_other[d] = 1
-                                    normal[d] = sdf_volume[tuple(idx_other)] - sdf_volume[tuple(idx_curr)]
-                                else:
-                                    idx_other = list(idx_3d)
-                                    idx_other[d] = resolution - 2
-                                    normal[d] = sdf_volume[tuple(idx_curr)] - sdf_volume[tuple(idx_other)]
-                        
-                        if np.linalg.norm(normal) > 0:
-                            normal = normal / np.linalg.norm(normal)
-                            surface_normals.append(normal)
-        
-        # Limit the number of points to avoid OOM
-        max_points = 5000
-        if len(surface_points) > max_points:
-            indices = np.random.choice(len(surface_points), max_points, replace=False)
-            surface_points = [surface_points[i] for i in indices]
-            surface_normals = [surface_normals[i] for i in indices]
-        
-        if len(surface_points) < 4:
-            # Not enough points found, create a simple cube
-            vertices = np.array([
-                [-0.5, -0.5, -0.5],
-                [0.5, -0.5, -0.5],
-                [0.5, 0.5, -0.5],
-                [-0.5, 0.5, -0.5],
-                [-0.5, -0.5, 0.5],
-                [0.5, -0.5, 0.5],
-                [0.5, 0.5, 0.5],
-                [-0.5, 0.5, 0.5]
-            ])
-            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
-        
-        # Convert points to numpy array
-        points = np.array(surface_points)
-        
-        # Create a simple mesh using ball-pivoting like algorithm
-        vertices = points
-        
-        # For simplicity, create faces from nearest neighbors
-        # This is a very simple approach, not as good as Delaunay but doesn't require scipy
-        faces = []
-        
-        # Use a simple approach to creating faces
-        # We'll use a greedy algorithm to connect nearby points
-        n_points = len(vertices)
-        
-        # Create a simple connectivity graph
-        # For each point, find the N closest points
-        n_neighbors = min(12, n_points - 1)
-        adjacency = [[] for _ in range(n_points)]
-        
-        # Compute all pairwise distances (this is O(n²) but should be ok for small point clouds)
-        for i in range(n_points):
-            distances = []
-            for j in range(n_points):
-                if i != j:
-                    dist = np.linalg.norm(vertices[i] - vertices[j])
-                    distances.append((dist, j))
-            distances.sort()
-            for k in range(min(n_neighbors, len(distances))):
-                adjacency[i].append(distances[k][1])
-        
-        # Create triangles from the adjacency list
-        added_edges = set()
-        for i in range(n_points):
-            for j in adjacency[i]:
-                if j > i:  # Avoid duplicates
-                    edge_ij = (i, j)
-                    added_edges.add(edge_ij)
-                    
-                    # Find common neighbors between i and j to form triangles
-                    common_neighbors = set(adjacency[i]) & set(adjacency[j])
-                    for k in common_neighbors:
-                        if k != i and k != j:
-                            edge_ik = (i, k) if i < k else (k, i)
-                            edge_jk = (j, k) if j < k else (k, j)
-                            
-                            # Check if the other two edges exist
-                            if edge_ik in added_edges and edge_jk in added_edges:
-                                # Ensure consistent winding
-                                normal = np.cross(vertices[j] - vertices[i], vertices[k] - vertices[i])
-                                center_normal = np.mean(surface_normals, axis=0)
-                                
-                                if np.dot(normal, center_normal) > 0:
-                                    faces.append([i, j, k])
-                                else:
-                                    faces.append([i, k, j])
-        
-        # If we couldn't create enough faces, create a simple shape
-        if len(faces) < 4:
-            # Create a convex hull-like shape
-            center = np.mean(vertices, axis=0)
-            n_points = len(vertices)
-            
-            # Sort points by distance from center
-            dists = np.linalg.norm(vertices - center, axis=1)
-            sorted_indices = np.argsort(dists)
-            
-            # Create a star-like structure connecting to center
-            center_idx = sorted_indices[0]
-            faces = []
-            
-            for i in range(1, min(n_points, 10)):
-                if i + 1 < n_points:
-                    faces.append([center_idx, sorted_indices[i], sorted_indices[i + 1]])
-                else:
-                    faces.append([center_idx, sorted_indices[i], sorted_indices[1]])
-            
-            # If still not enough, create a simple cube
-            if len(faces) < 4:
-                vertices = np.array([
-                    [-0.5, -0.5, -0.5],
-                    [0.5, -0.5, -0.5],
-                    [0.5, 0.5, -0.5],
-                    [-0.5, 0.5, -0.5],
-                    [-0.5, -0.5, 0.5],
-                    [0.5, -0.5, 0.5],
-                    [0.5, 0.5, 0.5],
-                    [-0.5, 0.5, 0.5]
-                ])
-                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 np.array(vertices), np.array(faces)
+        # Convert to RGB and return
+        return img
 
 def load_model():
-    global neus_model, model_loaded, model_loading
+    global pipe, model_loaded, model_loading
     
     if model_loaded:
-        return neus_model
+        return pipe
     
     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 neus_model
+        return pipe
     
     try:
         model_loading = True
         print("Starting model loading...")
         
+        model_name = "openai/shap-e-img2img"
+        
+        # Download model with retry mechanism
+        max_retries = 3
+        retry_delay = 5
+        
+        for attempt in range(max_retries):
+            try:
+                snapshot_download(
+                    repo_id=model_name,
+                    cache_dir=CACHE_DIR,
+                    resume_download=True,
+                )
+                break
+            except Exception as e:
+                if attempt < max_retries - 1:
+                    print(f"Download attempt {attempt+1} failed: {str(e)}. Retrying in {retry_delay} seconds...")
+                    time.sleep(retry_delay)
+                    retry_delay *= 2
+                else:
+                    raise
+        
+        # Initialize pipeline with lower precision to save memory
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        neus_model = NeuS2Model(device=device)
+        dtype = torch.float16 if device == "cuda" else torch.float32
+        
+        pipe = ShapEImg2ImgPipeline.from_pretrained(
+            model_name,
+            torch_dtype=dtype,
+            cache_dir=CACHE_DIR,
+        )
+        pipe = pipe.to(device)
+        
+        # Optimize for inference
+        if device == "cuda":
+            pipe.enable_model_cpu_offload()
         
         model_loaded = True
         print(f"Model loaded successfully on {device}")
-        return neus_model
+        return pipe
     
     except Exception as e:
         print(f"Error loading model: {str(e)}")
@@ -420,7 +175,7 @@ def load_model():
 def health_check():
     return jsonify({
         "status": "healthy", 
-        "model": "NeuS2 Image to 3D",
+        "model": "Shap-E Image to 3D",
         "device": "cuda" if torch.cuda.is_available() else "cpu"
     }), 200
 
@@ -463,291 +218,6 @@ def progress(job_id):
     
     return Response(stream_with_context(generate()), mimetype='text/event-stream')
 
-def export_to_obj(mesh, obj_path):
-    """Export mesh to OBJ file format"""
-    vertices = mesh.verts
-    faces = mesh.faces
-    
-    with open(obj_path, 'w') as f:
-        # Write vertices
-        for v in vertices:
-            f.write(f"v {v[0]} {v[1]} {v[2]}\n")
-        
-        # Write faces (OBJ uses 1-indexed vertices)
-        for face in faces:
-            f.write(f"f {face[0]+1} {face[1]+1} {face[2]+1}\n")
-    
-    # Create a simple MTL file
-    mtl_path = obj_path.replace('.obj', '.mtl')
-    with open(mtl_path, 'w') as f:
-        f.write("newmtl material0\n")
-        f.write("Ka 1.0 1.0 1.0\n")  # ambient color
-        f.write("Kd 0.8 0.8 0.8\n")  # diffuse color
-        f.write("Ks 0.0 0.0 0.0\n")  # specular color
-        f.write("Ns 0.0\n")          # specular exponent
-        f.write("illum 2\n")         # illumination model
-    
-    return obj_path, mtl_path
-
-@app.route('/convert', methods=['POST'])
-def convert_image_to_3d():
-    # Check if image is in the request
-    if 'image' not in request.files:
-        return jsonify({"error": "No image provided"}), 400
-    
-    file = request.files['image']
-    if file.filename == '':
-        return jsonify({"error": "No image selected"}), 400
-    
-    if not allowed_file(file.filename):
-        return jsonify({"error": f"File type not allowed. Supported types: {', '.join(ALLOWED_EXTENSIONS)}"}), 400
-    
-    # Get optional parameters with defaults
-    try:
-        guidance_scale = float(request.form.get('guidance_scale', 3.0))
-        num_inference_steps = min(int(request.form.get('num_inference_steps', 32)), MAX_INFERENCE_STEPS)
-        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
-    
-    if num_inference_steps < 16 or num_inference_steps > MAX_INFERENCE_STEPS:
-        num_inference_steps = min(num_inference_steps, MAX_INFERENCE_STEPS)
-    
-    # Validate output format
-    if output_format not in ['obj', 'glb']:
-        return jsonify({"error": "Unsupported output format. Use 'obj' or 'glb'"}), 400
-    
-    # Create a job ID
-    job_id = str(uuid.uuid4())
-    output_dir = os.path.join(RESULTS_FOLDER, job_id)
-    os.makedirs(output_dir, exist_ok=True)
-    
-    # Save the uploaded file
-    filename = secure_filename(file.filename)
-    filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{job_id}_{filename}")
-    file.save(filepath)
-    
-    # Initialize job tracking
-    processing_jobs[job_id] = {
-        'status': 'processing',
-        'progress': 0,
-        'result_url': None,
-        'preview_url': None,
-        'error': None,
-        'output_format': output_format,
-        'created_at': time.time()
-    }
-    
-    # Start processing in a separate thread
-    def process_image():
-        thread = threading.current_thread()
-        processing_jobs[job_id]['thread_alive'] = lambda: thread.is_alive()
-        
-        try:
-            # Preprocess image (resize if needed)
-            processing_jobs[job_id]['progress'] = 5
-            image_tensor = preprocess_image(filepath)
-            processing_jobs[job_id]['progress'] = 10
-            
-            # Load model
-            try:
-                model = load_model()
-                processing_jobs[job_id]['progress'] = 30
-            except Exception as e:
-                processing_jobs[job_id]['status'] = 'error'
-                processing_jobs[job_id]['error'] = f"Error loading model: {str(e)}"
-                return
-            
-            # Process image with thread-safe timeout
-            try:
-                def generate_mesh():
-                    return model.generate_mesh(
-                        image_tensor, 
-                        resolution=min(32 + num_inference_steps // 2, 64),  # Adjust resolution based on steps
-                        threshold=0.0,
-                        num_steps=num_inference_steps
-                    )
-                
-                images, error = process_with_timeout(generate_mesh, [], TIMEOUT_SECONDS)
-                
-                if error:
-                    if isinstance(error, TimeoutError):
-                        processing_jobs[job_id]['status'] = 'error'
-                        processing_jobs[job_id]['error'] = f"Processing timed out after {TIMEOUT_SECONDS} seconds"
-                        return
-                    else:
-                        raise error
-                        
-                processing_jobs[job_id]['progress'] = 80
-            except Exception as e:
-                error_details = traceback.format_exc()
-                processing_jobs[job_id]['status'] = 'error'
-                processing_jobs[job_id]['error'] = f"Error during processing: {str(e)}"
-                print(f"Error processing job {job_id}: {str(e)}")
-                print(error_details)
-                return
-            
-            # Export based on requested format
-            try:
-                if output_format == 'obj':
-                    obj_path = os.path.join(output_dir, "model.obj")
-                    obj_path, mtl_path = export_to_obj(images[0], obj_path)
-                    
-                    # Create a zip file with OBJ and MTL
-                    zip_path = os.path.join(output_dir, "model.zip")
-                    with zipfile.ZipFile(zip_path, 'w') as zipf:
-                        zipf.write(obj_path, arcname="model.obj")
-                        if os.path.exists(mtl_path):
-                            zipf.write(mtl_path, arcname="model.mtl")
-                    
-                    processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
-                    processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
-                    
-                elif output_format == 'glb':
-                    # Convert to trimesh format
-                    vertices = images[0].verts
-                    faces = images[0].faces
-                    
-                    # Create a trimesh object
-                    trimesh_obj = trimesh.Trimesh(vertices=vertices, faces=faces)
-                    
-                    # Export as GLB
-                    glb_path = os.path.join(output_dir, "model.glb")
-                    trimesh_obj.export(glb_path)
-                    
-                    processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
-                    processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
-                
-                # Update job status
-                processing_jobs[job_id]['status'] = 'completed'
-                processing_jobs[job_id]['progress'] = 100
-                print(f"Job {job_id} completed successfully")
-            except Exception as e:
-                error_details = traceback.format_exc()
-                processing_jobs[job_id]['status'] = 'error'
-                processing_jobs[job_id]['error'] = f"Error exporting model: {str(e)}"
-                print(f"Error exporting model for job {job_id}: {str(e)}")
-                print(error_details)
-            
-            # Clean up temporary file
-            if os.path.exists(filepath):
-                os.remove(filepath)
-            
-            # Force garbage collection to free memory
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            
-        except Exception as e:
-            # Handle errors
-            error_details = traceback.format_exc()
-            processing_jobs[job_id]['status'] = 'error'
-            processing_jobs[job_id]['error'] = f"{str(e)}\n{error_details}"
-            print(f"Error processing job {job_id}: {str(e)}")
-            print(error_details)
-            
-            # Clean up on error
-            if os.path.exists(filepath):
-                os.remove(filepath)
-    
-    # Start processing thread
-    processing_thread = threading.Thread(target=process_image)
-    processing_thread.daemon = True
-    processing_thread.start()
-    
-    # Return job ID immediately
-    return jsonify({"job_id": job_id}), 202  # 202 Accepted
-
-@app.route('/download/<job_id>', methods=['GET'])
-def download_model(job_id):
-    if job_id not in processing_jobs or processing_jobs[job_id]['status'] != 'completed':
-        return jsonify({"error": "Model not found or processing not complete"}), 404
-    
-    # Get the output directory for this job
-    output_dir = os.path.join(RESULTS_FOLDER, job_id)
-    
-    # Determine file format from the job data
-    output_format = processing_jobs[job_id].get('output_format', 'obj')
-    
-    if output_format == 'obj':
-        zip_path = os.path.join(output_dir, "model.zip")
-        if os.path.exists(zip_path):
-            return send_file(zip_path, as_attachment=True, download_name="model.zip")
-    else:  # glb
-        glb_path = os.path.join(output_dir, "model.glb")
-        if os.path.exists(glb_path):
-            return send_file(glb_path, as_attachment=True, download_name="model.glb")
-    
-    return jsonify({"error": "File not found"}), 404
-@app.route('/progress/<job_id>', methods=['GET'])
-def progress(job_id):
-    def generate():
-        if job_id not in processing_jobs:
-            yield f"data: {json.dumps({'error': 'Job not found'})}\n\n"
-            return
-            
-        job = processing_jobs[job_id]
-        
-        # Send initial progress
-        yield f"data: {json.dumps({'status': 'processing', 'progress': job['progress']})}\n\n"
-        
-        # Wait for job to complete or update
-        last_progress = job['progress']
-        check_count = 0
-        while job['status'] == 'processing':
-            if job['progress'] != last_progress:
-                yield f"data: {json.dumps({'status': 'processing', 'progress': job['progress']})}\n\n"
-                last_progress = job['progress']
-            
-            time.sleep(0.5)
-            check_count += 1
-            
-            # If client hasn't received updates for a while, check if job is still running
-            if check_count > 60:  # 30 seconds with no updates
-                if 'thread_alive' in job and not job['thread_alive']():
-                    job['status'] = 'error'
-                    job['error'] = 'Processing thread died unexpectedly'
-                    break
-                check_count = 0
-        
-        # Send final status
-        if job['status'] == 'completed':
-            yield f"data: {json.dumps({'status': 'completed', 'progress': 100, 'result_url': job['result_url'], 'preview_url': job['preview_url']})}\n\n"
-        else:
-            yield f"data: {json.dumps({'status': 'error', 'error': job['error']})}\n\n"
-    
-    return Response(stream_with_context(generate()), mimetype='text/event-stream')
-
-def export_to_obj(mesh, obj_path):
-    """Export mesh to OBJ file format"""
-    vertices = mesh.verts
-    faces = mesh.faces
-    
-    with open(obj_path, 'w') as f:
-        # Write vertices
-        for v in vertices:
-            f.write(f"v {v[0]} {v[1]} {v[2]}\n")
-        
-        # Write faces (OBJ uses 1-indexed vertices)
-        for face in faces:
-            f.write(f"f {face[0]+1} {face[1]+1} {face[2]+1}\n")
-    
-    # Create a simple MTL file
-    mtl_path = obj_path.replace('.obj', '.mtl')
-    with open(mtl_path, 'w') as f:
-        f.write("newmtl material0\n")
-        f.write("Ka 1.0 1.0 1.0\n")  # ambient color
-        f.write("Kd 0.8 0.8 0.8\n")  # diffuse color
-        f.write("Ks 0.0 0.0 0.0\n")  # specular color
-        f.write("Ns 0.0\n")          # specular exponent
-        f.write("illum 2\n")         # illumination model
-    
-    return obj_path, mtl_path
-
 @app.route('/convert', methods=['POST'])
 def convert_image_to_3d():
     # Check if image is in the request
@@ -764,7 +234,7 @@ def convert_image_to_3d():
     # Get optional parameters with defaults
     try:
         guidance_scale = float(request.form.get('guidance_scale', 3.0))
-        num_inference_steps = min(int(request.form.get('num_inference_steps', 32)), MAX_INFERENCE_STEPS)
+        num_inference_steps = min(int(request.form.get('num_inference_steps', 64)), MAX_INFERENCE_STEPS)
         output_format = request.form.get('output_format', 'obj').lower()
     except ValueError:
         return jsonify({"error": "Invalid parameter values"}), 400
@@ -773,7 +243,7 @@ def convert_image_to_3d():
     if guidance_scale < 1.0 or guidance_scale > 5.0:
         return jsonify({"error": "Guidance scale must be between 1.0 and 5.0"}), 400
     
-    if num_inference_steps < 16 or num_inference_steps > MAX_INFERENCE_STEPS:
+    if num_inference_steps < 32 or num_inference_steps > MAX_INFERENCE_STEPS:
         num_inference_steps = min(num_inference_steps, MAX_INFERENCE_STEPS)
     
     # Validate output format
@@ -809,12 +279,12 @@ def convert_image_to_3d():
         try:
             # Preprocess image (resize if needed)
             processing_jobs[job_id]['progress'] = 5
-            image_tensor = preprocess_image(filepath)
+            image = preprocess_image(filepath)
             processing_jobs[job_id]['progress'] = 10
             
             # Load model
             try:
-                model = load_model()
+                pipe = load_model()
                 processing_jobs[job_id]['progress'] = 30
             except Exception as e:
                 processing_jobs[job_id]['status'] = 'error'
@@ -824,12 +294,12 @@ def convert_image_to_3d():
             # Process image with thread-safe timeout
             try:
                 def generate_mesh():
-                    return model.generate_mesh(
-                        image_tensor, 
-                        resolution=min(32 + num_inference_steps, 64),  # Adjust resolution based on steps
-                        threshold=0.0,
-                        num_steps=num_inference_steps
-                    )
+                    return pipe(
+                        image,
+                        guidance_scale=guidance_scale,
+                        num_inference_steps=num_inference_steps,
+                        output_type="mesh",
+                    ).images
                 
                 images, error = process_with_timeout(generate_mesh, [], TIMEOUT_SECONDS)
                 
@@ -854,12 +324,13 @@ def convert_image_to_3d():
             try:
                 if output_format == 'obj':
                     obj_path = os.path.join(output_dir, "model.obj")
-                    obj_path, mtl_path = export_to_obj(images[0], obj_path)
+                    export_to_obj(images[0], obj_path)
                     
                     # Create a zip file with OBJ and MTL
                     zip_path = os.path.join(output_dir, "model.zip")
                     with zipfile.ZipFile(zip_path, 'w') as zipf:
                         zipf.write(obj_path, arcname="model.obj")
+                        mtl_path = os.path.join(output_dir, "model.mtl")
                         if os.path.exists(mtl_path):
                             zipf.write(mtl_path, arcname="model.mtl")
                     
@@ -867,12 +338,13 @@ def convert_image_to_3d():
                     processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
                     
                 elif output_format == 'glb':
-                    # Convert to trimesh format
-                    vertices = images[0].verts
-                    faces = images[0].faces
+                    from trimesh import Trimesh
+                    mesh = images[0]
+                    vertices = mesh.verts
+                    faces = mesh.faces
                     
                     # Create a trimesh object
-                    trimesh_obj = trimesh.Trimesh(vertices=vertices, faces=faces)
+                    trimesh_obj = Trimesh(vertices=vertices, faces=faces)
                     
                     # Export as GLB
                     glb_path = os.path.join(output_dir, "model.glb")
@@ -996,7 +468,7 @@ def cleanup_old_jobs():
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API using NeuS2 is running", 
+        "message": "Image to 3D API is running", 
         "endpoints": ["/convert", "/progress/<job_id>", "/download/<job_id>", "/preview/<job_id>"]
     }), 200