Update app.py

app.py

@@ -15,7 +15,6 @@ from diffusers import ShapEImg2ImgPipeline
 from diffusers.utils import export_to_obj
 from huggingface_hub import snapshot_download
 from flask_cors import CORS
-import signal
 import functools
 
 app = Flask(__name__)
@@ -51,29 +50,41 @@ model_loading = False
 # Configuration for processing
 TIMEOUT_SECONDS = 300  # 5 minutes max for processing
 MAX_DIMENSION = 512    # Max image dimension to process
+MAX_INFERENCE_STEPS = 64  # Maximum allowed inference steps to prevent the index error
 
-# Timeout handler for long-running processes
+# TimeoutError for handling timeouts
 class TimeoutError(Exception):
     pass
 
-def timeout_handler(signum, frame):
-    raise TimeoutError("Processing timed out")
-
-def with_timeout(timeout):
-    def decorator(func):
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs):
-            # Set the timeout handler
-            signal.signal(signal.SIGALRM, timeout_handler)
-            signal.alarm(timeout)
-            try:
-                result = func(*args, **kwargs)
-            finally:
-                # Disable the alarm
-                signal.alarm(0)
-            return result
-        return wrapper
-    return decorator
+# Thread-safe timeout implementation
+def process_with_timeout(function, args, timeout):
+    result = [None]
+    error = [None]
+    completed = [False]
+    
+    def target():
+        try:
+            result[0] = function(*args)
+            completed[0] = True
+        except Exception as e:
+            error[0] = e
+    
+    thread = threading.Thread(target=target)
+    thread.daemon = True
+    thread.start()
+    
+    thread.join(timeout)
+    
+    if not completed[0]:
+        if thread.is_alive():
+            return None, TimeoutError(f"Processing timed out after {timeout} seconds")
+        elif error[0]:
+            return None, error[0]
+    
+    if error[0]:
+        return None, error[0]
+    
+    return result[0], None
 
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
@@ -223,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 = int(request.form.get('num_inference_steps', 64))
+        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
@@ -232,8 +243,8 @@ 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 < 32 or num_inference_steps > 128:
-        return jsonify({"error": "Number of inference steps must be between 32 and 128"}), 400
+    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
     if output_format not in ['obj', 'glb']:
@@ -260,21 +271,6 @@ def convert_image_to_3d():
         'created_at': time.time()
     }
     
-    # Process function with timeout
-    @with_timeout(TIMEOUT_SECONDS)
-    def process_with_timeout(image, steps, scale, format):
-        # Load model
-        pipe = load_model()
-        processing_jobs[job_id]['progress'] = 30
-        
-        # Generate 3D model
-        return pipe(
-            image,
-            guidance_scale=scale,
-            num_inference_steps=steps,
-            output_type="mesh",
-        ).images
-    
     # Start processing in a separate thread
     def process_image():
         thread = threading.current_thread()
@@ -286,50 +282,87 @@ def convert_image_to_3d():
             image = preprocess_image(filepath)
             processing_jobs[job_id]['progress'] = 10
             
-            # Process image with timeout
+            # Load model
             try:
-                images = process_with_timeout(image, num_inference_steps, guidance_scale, output_format)
-                processing_jobs[job_id]['progress'] = 80
-            except TimeoutError:
+                pipe = load_model()
+                processing_jobs[job_id]['progress'] = 30
+            except Exception as e:
                 processing_jobs[job_id]['status'] = 'error'
-                processing_jobs[job_id]['error'] = f"Processing timed out after {TIMEOUT_SECONDS} seconds"
+                processing_jobs[job_id]['error'] = f"Error loading model: {str(e)}"
                 return
             
-            # Export based on requested format
-            if output_format == 'obj':
-                obj_path = os.path.join(output_dir, "model.obj")
-                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")
-                
-                processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
-                processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
-                
-            elif output_format == 'glb':
-                from trimesh import Trimesh
-                mesh = images[0]
-                vertices = mesh.verts
-                faces = mesh.faces
-                
-                # Create a trimesh object
-                trimesh_obj = Trimesh(vertices=vertices, faces=faces)
+            # Process image with thread-safe timeout
+            try:
+                def generate_mesh():
+                    return pipe(
+                        image,
+                        guidance_scale=guidance_scale,
+                        num_inference_steps=num_inference_steps,
+                        output_type="mesh",
+                    ).images
                 
-                # Export as GLB
-                glb_path = os.path.join(output_dir, "model.glb")
-                trimesh_obj.export(glb_path)
+                images, error = process_with_timeout(generate_mesh, [], TIMEOUT_SECONDS)
                 
-                processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
-                processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
+                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
             
-            # Update job status
-            processing_jobs[job_id]['status'] = 'completed'
-            processing_jobs[job_id]['progress'] = 100
+            # Export based on requested format
+            try:
+                if output_format == 'obj':
+                    obj_path = os.path.join(output_dir, "model.obj")
+                    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")
+                    
+                    processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
+                    processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
+                    
+                elif output_format == 'glb':
+                    from trimesh import Trimesh
+                    mesh = images[0]
+                    vertices = mesh.verts
+                    faces = mesh.faces
+                    
+                    # Create a trimesh object
+                    trimesh_obj = 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):