Update app.py

app.py

@@ -11,9 +11,10 @@ import io
 import zipfile
 import uuid
 import traceback
-from huggingface_hub import snapshot_download
 from flask_cors import CORS
 import numpy as np
+import subprocess
+import sys
 
 app = Flask(__name__)
 CORS(app)  # Enable CORS for all routes
@@ -33,7 +34,7 @@ os.makedirs(CACHE_DIR, exist_ok=True)
 os.environ['HF_HOME'] = CACHE_DIR
 os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
 os.environ['HF_DATASETS_CACHE'] = os.path.join(CACHE_DIR, 'datasets')
-os.environ['NUMBA_THREADING_LAYER'] = 'omp'
+os.environ['NUMBA_THREADING_LAYER'] = 'workqueue'  # Changed from 'omp'
 
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max
@@ -48,8 +49,8 @@ model_loaded = False
 model_loading = False
 
 # Configuration for processing
-TIMEOUT_SECONDS = 240  # 4 minutes max for processing
-MAX_DIMENSION = 512    # Max image dimension to process
+TIMEOUT_SECONDS = 480  # 8 minutes max for processing (increased for CPU)
+MAX_DIMENSION = 384    # Max image dimension to process (reduced for CPU)
 
 # TimeoutError for handling timeouts
 class TimeoutError(Exception):
@@ -58,22 +59,40 @@ class TimeoutError(Exception):
 # Install necessary dependencies
 def install_dependencies():
     try:
-        import subprocess
-        # Install core dependencies
+        # Install core dependencies in smaller batches to avoid memory issues
+        # First batch: essential packages
         subprocess.check_call([
-            "pip", "install", 
-            "torch>=2.0.0", 
+            "pip", "install", "--no-cache-dir", 
+            "torch==2.0.1+cpu", 
+            "torchvision==0.15.2+cpu", 
+            "--extra-index-url", 
+            "https://download.pytorch.org/whl/cpu"
+        ])
+        
+        # Second batch: OpenLRM dependencies 1
+        subprocess.check_call([
+            "pip", "install", "--no-cache-dir", 
             "lpips", 
             "omegaconf", 
             "transformers",
-            "safetensors", 
+            "safetensors"
+        ])
+        
+        # Third batch: OpenLRM dependencies 2
+        subprocess.check_call([
+            "pip", "install", "--no-cache-dir", 
             "accelerate", 
             "imageio[ffmpeg]",
             "PyMCubes", 
-            "trimesh", 
-            "opencv-python", 
-            "rembg[gpu,cli]", 
-            "httpx[socks]", 
+            "trimesh"
+        ])
+        
+        # Fourth batch: Image processing
+        subprocess.check_call([
+            "pip", "install", "--no-cache-dir", 
+            "opencv-python-headless",  # Use headless variant
+            "rembg[cpu]",  # CPU version instead of GPU
+            "httpx", 
             "tensorboard"
         ])
         
@@ -84,12 +103,14 @@ def install_dependencies():
         # Add OpenLRM to python path
         if not "OpenLRM" in os.getenv("PYTHONPATH", ""):
             os.environ["PYTHONPATH"] = f"{os.getenv('PYTHONPATH', '')}:OpenLRM"
+            sys.path.append("OpenLRM")  # Add to sys.path directly
             
         print("Successfully installed dependencies")
+        return True
     except Exception as e:
         print(f"Error installing dependencies: {str(e)}")
         print(traceback.format_exc())
-        raise
+        return False
 
 # Thread-safe timeout implementation
 def process_with_timeout(function, args, timeout):
@@ -127,7 +148,14 @@ def allowed_file(filename):
 # Image preprocessing with automatic background removal
 def preprocess_image(image_path):
     try:
-        from rembg import remove
+        # Try to use rembg for bg removal
+        try:
+            from rembg import remove
+            use_rembg = True
+        except ImportError:
+            use_rembg = False
+            print("Rembg not available, skipping background removal")
+            
         with Image.open(image_path) as img:
             img = img.convert("RGBA")
             
@@ -144,8 +172,12 @@ def preprocess_image(image_path):
                 # Use high-quality Lanczos resampling for better detail preservation
                 img = img.resize((new_width, new_height), Image.LANCZOS)
             
-            # Remove background automatically
-            img_no_bg = remove(img)
+            # Remove background automatically if rembg is available
+            if use_rembg:
+                img_no_bg = remove(img)
+            else:
+                # Create a fallback "no bg" version by just using original
+                img_no_bg = img.copy()
             
             # Save both versions for flexibility
             img_path = image_path.replace(".jpg", ".png").replace(".jpeg", ".png")
@@ -158,10 +190,10 @@ def preprocess_image(image_path):
     except Exception as e:
         print(f"Error in image preprocessing: {str(e)}")
         print(traceback.format_exc())
-        # Return original if rembg fails
+        # Return original if preprocessing fails
         return image_path, image_path
 
-# Initialize OpenLRM model
+# Initialize OpenLRM model with special handling for CPU-only environments
 def load_model():
     global openlrm_model, openlrm_processor, model_loaded, model_loading
     
@@ -178,59 +210,55 @@ def load_model():
         model_loading = True
         print("Starting OpenLRM model loading...")
         
-        # Import OpenLRM components
+        # First check if dependencies are installed, if not install them
+        try:
+            from openlrm.utils.preprocess import Preprocessor
+        except ImportError:
+            success = install_dependencies()
+            if not success:
+                raise ImportError("Failed to install OpenLRM dependencies")
+        
+        # Now try importing needed components
         try:
             from openlrm.utils.preprocess import Preprocessor
             from openlrm.utils.config import load_config
             from openlrm.models.registry import get_model
             from openlrm.pipelines.inference import InferencePipeline
             
-            # Use the small model variant for HF free tier
-            model_name = "zxhezexin/openlrm-mix-small-1.1"  # Smallest model for HF free tier
+            # Use the tiny model variant for CPU environment
+            model_name = "zxhezexin/openlrm-mix-tiny-1.0"  # Tiniest model for CPU
             
             # Load configuration for inference
-            config_path = "OpenLRM/configs/infer-s.yaml"  # Small model config
+            config_path = "OpenLRM/configs/infer-xs.yaml"  # Extra small model config
             config = load_config(config_path)
             config.model_name = model_name
             
+            # CPU optimizations
+            config.half_precision = False  # Disable half precision on CPU
+            config.mesh_resolution = 64    # Lower mesh resolution
+            
             # Initialize preprocessor
             openlrm_processor = Preprocessor()
             
-            # Initialize model and inference pipeline
-            device = "cuda" if torch.cuda.is_available() else "cpu"
+            # Initialize model and inference pipeline - force CPU
+            device = "cpu"  # Force CPU
+            torch.set_num_threads(4)  # Limit CPU threads to avoid OOM
+            
+            print(f"Using device: {device}")
             openlrm_model = InferencePipeline(config, device)
             
             print(f"OpenLRM model loaded successfully on {device}")
             model_loaded = True
             
-            # Optimize memory usage
-            if device == "cuda":
-                torch.cuda.empty_cache()
+            # Clean up memory
+            gc.collect()
                 
             return openlrm_model, openlrm_processor
             
-        except ImportError as e:
-            print(f"ImportError: {str(e)}")
-            print("Installing OpenLRM dependencies...")
-            install_dependencies()
-            # Try loading again after installing dependencies
-            from openlrm.utils.preprocess import Preprocessor
-            from openlrm.utils.config import load_config
-            from openlrm.models.registry import get_model
-            from openlrm.pipelines.inference import InferencePipeline
-            
-            model_name = "zxhezexin/openlrm-mix-small-1.1"
-            config_path = "OpenLRM/configs/infer-s.yaml"
-            config = load_config(config_path)
-            config.model_name = model_name
-            
-            openlrm_processor = Preprocessor()
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            openlrm_model = InferencePipeline(config, device)
-            
-            model_loaded = True
-            print(f"OpenLRM model loaded successfully on {device} after installing dependencies")
-            return openlrm_model, openlrm_processor
+        except Exception as e:
+            print(f"Error importing OpenLRM components: {str(e)}")
+            print(traceback.format_exc())
+            raise
             
     except Exception as e:
         print(f"Error loading OpenLRM model: {str(e)}")
@@ -239,25 +267,73 @@ def load_model():
     finally:
         model_loading = False
 
-# Fallback to original depth-based implementation if OpenLRM fails
-def depth_based_fallback(image_path, output_dir, detail_level='high'):
+# Simple depth-based fallback implementation if OpenLRM fails
+def depth_based_fallback(image_path, output_dir, detail_level='medium'):
     try:
-        # This uses your original depth estimation implementation as a fallback
-        # [Implementation would go here]
-        print("Using depth-based fallback implementation")
-        # Your original implementation could be added here
-        pass
+        import numpy as np
+        from PIL import Image
+        import trimesh
+        
+        print("Using simple depth-based fallback implementation")
+        
+        # Create a simple mesh using a heightmap approach
+        img = Image.open(image_path).convert('L')  # Convert to grayscale
+        
+        # Resize to manageable size based on detail level
+        if detail_level == 'low':
+            size = 32
+        elif detail_level == 'medium':
+            size = 64
+        else:  # high
+            size = 96
+            
+        img = img.resize((size, size))
+        
+        # Create a heightmap from image intensity
+        heightmap = np.array(img) / 255.0
+        
+        # Create a 3D mesh grid
+        x, y = np.meshgrid(np.linspace(-1, 1, size), np.linspace(-1, 1, size))
+        z = heightmap * 0.5  # Scale the height
+        
+        # Create vertices for the mesh
+        vertices = np.zeros((size * size, 3))
+        vertices[:, 0] = x.flatten()
+        vertices[:, 1] = y.flatten()
+        vertices[:, 2] = z.flatten()
+        
+        # Create face indices
+        faces = []
+        for i in range(size - 1):
+            for j in range(size - 1):
+                v0 = i * size + j
+                v1 = i * size + (j + 1)
+                v2 = (i + 1) * size + j
+                v3 = (i + 1) * size + (j + 1)
+                faces.append([v0, v1, v2])
+                faces.append([v1, v3, v2])
+        
+        faces = np.array(faces)
+        
+        # Create a mesh using trimesh
+        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+        
+        # Export as PLY
+        mesh_path = os.path.join(output_dir, "output.ply")
+        mesh.export(mesh_path)
+        
+        return True
     except Exception as e:
         print(f"Fallback also failed: {str(e)}")
+        print(traceback.format_exc())
         return False
-    return True
 
 @app.route('/health', methods=['GET'])
 def health_check():
     return jsonify({
         "status": "healthy", 
-        "model": "OpenLRM Image-to-3D Model Generator",
-        "device": "cuda" if torch.cuda.is_available() else "cpu"
+        "model": "OpenLRM Image-to-3D Model Generator (CPU Edition)",
+        "device": "cpu"
     }), 200
 
 @app.route('/progress/<job_id>', methods=['GET'])
@@ -315,7 +391,7 @@ def convert_image_to_3d():
     # Get optional parameters with defaults
     try:
         output_format = request.form.get('output_format', 'obj').lower()
-        detail_level = request.form.get('detail_level', 'medium').lower()
+        detail_level = request.form.get('detail_level', 'low').lower()  # Default to low for CPU
         source_cam_dist = float(request.form.get('source_cam_dist', 2.0))
         remove_bg = request.form.get('remove_bg', 'true').lower() == 'true'
     except ValueError:
@@ -357,150 +433,151 @@ def convert_image_to_3d():
             img_path, img_no_bg_path = preprocess_image(filepath) if remove_bg else (filepath, filepath)
             processing_jobs[job_id]['progress'] = 10
             
-            # Load model
+            # Try to run with OpenLRM
+            use_fallback = False
             try:
+                # Load model
                 openlrm_model, openlrm_processor = 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
+                processing_jobs[job_id]['status'] = 'processing'  # Keep processing, we'll use fallback
+                processing_jobs[job_id]['error'] = f"OpenLRM model failed to load: {str(e)}. Using fallback."
+                use_fallback = True
             
-            # Process image with thread-safe timeout
-            try:
-                def generate_3d():
-                    # Import here to ensure it's within the thread
-                    import os
-                    from openlrm.pipelines.inference import InferencePipeline
-                    
-                    # Process with OpenLRM
-                    image_to_use = img_no_bg_path if remove_bg else img_path
+            # Process image with appropriate method
+            if not use_fallback:
+                try:
+                    def generate_3d():
+                        # Use OpenLRM
+                        image_to_use = img_no_bg_path if remove_bg else img_path
+                        
+                        # Configure export paths
+                        dump_video_path = os.path.join(output_dir, "output.mp4")
+                        dump_mesh_path = os.path.join(output_dir, "output.ply")
+                        
+                        # Process with OpenLRM
+                        openlrm_model.infer_single(
+                            image_path=image_to_use,
+                            source_cam_dist=source_cam_dist,
+                            export_video=True,
+                            export_mesh=True,
+                            dump_video_path=dump_video_path,
+                            dump_mesh_path=dump_mesh_path,
+                        )
+                        
+                        return dump_video_path, dump_mesh_path
                     
-                    # Configure export paths
-                    dump_video_path = os.path.join(output_dir, "output.mp4")
-                    dump_mesh_path = os.path.join(output_dir, "output.ply")  # OpenLRM uses .ply format
+                    # Run with timeout
+                    (video_path, mesh_path), error = process_with_timeout(generate_3d, [], TIMEOUT_SECONDS)
                     
-                    # Process with OpenLRM
-                    openlrm_model.infer_single(
-                        image_path=image_to_use,
-                        source_cam_dist=source_cam_dist,
-                        export_video=True,
-                        export_mesh=True,
-                        dump_video_path=dump_video_path,
-                        dump_mesh_path=dump_mesh_path,
-                    )
+                    if error:
+                        if isinstance(error, TimeoutError):
+                            use_fallback = True
+                            processing_jobs[job_id]['error'] = f"OpenLRM processing timed out after {TIMEOUT_SECONDS} seconds. Using fallback."
+                        else:
+                            use_fallback = True
+                            processing_jobs[job_id]['error'] = f"OpenLRM processing failed: {str(error)}. Using fallback."
                     
-                    return dump_video_path, dump_mesh_path
-                
-                (video_path, mesh_path), error = process_with_timeout(generate_3d, [], TIMEOUT_SECONDS)
+                except Exception as e:
+                    use_fallback = True
+                    processing_jobs[job_id]['error'] = f"Error during OpenLRM processing: {str(e)}. Using fallback."
+            
+            # Use fallback method if needed
+            if use_fallback:
+                processing_jobs[job_id]['progress'] = 35
+                print("Using fallback method for 3D generation")
                 
-                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:
-                        # Try fallback implementation if OpenLRM fails
-                        processing_jobs[job_id]['progress'] = 35
-                        processing_jobs[job_id]['error'] = f"Primary method failed: {str(error)}. Trying fallback..."
-                        
-                        # Use fallback depth-based implementation
-                        if depth_based_fallback(img_path, output_dir, detail_level):
-                            processing_jobs[job_id]['progress'] = 60
-                            processing_jobs[job_id]['error'] = None  # Clear error if fallback succeeded
-                        else:
-                            raise Exception(f"Both primary and fallback 3D generation methods failed: {str(error)}")
-                        
+                # Use our simple depth-based implementation
+                if depth_based_fallback(img_path, output_dir, detail_level):
+                    processing_jobs[job_id]['progress'] = 60
+                else:
+                    raise Exception("Both primary and fallback 3D generation methods failed")
+            else:
                 processing_jobs[job_id]['progress'] = 60
+            
+            # Check if the mesh was created
+            mesh_path_orig = os.path.join(output_dir, "output.ply")
+            if not os.path.exists(mesh_path_orig):
+                raise Exception("3D model generation failed to produce a mesh")
                 
-                # Convert PLY to requested format if needed
-                mesh_path_orig = os.path.join(output_dir, "output.ply")
-                if os.path.exists(mesh_path_orig):
-                    if output_format == 'obj':
-                        # Convert PLY to OBJ
-                        import trimesh
-                        mesh = trimesh.load(mesh_path_orig)
-                        obj_path = os.path.join(output_dir, "model.obj")
-                        mesh.export(obj_path, file_type='obj')
-                        
-                        # 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")
-                            
-                            # Include texture file if it exists
-                            texture_path = os.path.join(output_dir, "model.png")
-                            if os.path.exists(texture_path):
-                                zipf.write(texture_path, arcname="model.png")
-                        
-                        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 PLY to GLB
-                        import trimesh
-                        mesh = trimesh.load(mesh_path_orig)
-                        glb_path = os.path.join(output_dir, "model.glb")
-                        mesh.export(glb_path, file_type='glb')
-                        
-                        processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
-                        processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
+            # Convert to requested format
+            if output_format == 'obj':
+                # Convert PLY to OBJ
+                import trimesh
+                mesh = trimesh.load(mesh_path_orig)
+                obj_path = os.path.join(output_dir, "model.obj")
+                mesh.export(obj_path, file_type='obj')
+                
+                # 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")
                     
-                    else:  # Keep as PLY format
-                        import shutil
-                        ply_path = os.path.join(output_dir, "model.ply")
-                        shutil.copy(mesh_path_orig, ply_path)
-                        
-                        processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
-                        processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
+                    # Include texture file if it exists
+                    texture_path = os.path.join(output_dir, "model.png")
+                    if os.path.exists(texture_path):
+                        zipf.write(texture_path, arcname="model.png")
                 
-                processing_jobs[job_id]['progress'] = 90
+                processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
+                processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
                 
-                # Also save the video preview
-                video_path_orig = os.path.join(output_dir, "output.mp4")
-                if os.path.exists(video_path_orig):
-                    preview_path = os.path.join(output_dir, "preview.mp4")
-                    import shutil
-                    shutil.copy(video_path_orig, preview_path)
-                    processing_jobs[job_id]['preview_video'] = f"/preview-video/{job_id}"
+            elif output_format == 'glb':
+                # Convert PLY to GLB
+                import trimesh
+                mesh = trimesh.load(mesh_path_orig)
+                glb_path = os.path.join(output_dir, "model.glb")
+                mesh.export(glb_path, file_type='glb')
                 
-                # Update job status
-                processing_jobs[job_id]['status'] = 'completed'
-                processing_jobs[job_id]['progress'] = 100
-                processing_jobs[job_id]['completed_at'] = time.time()
-                print(f"Job {job_id} completed successfully")
+                processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
+                processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
+            
+            else:  # Keep as PLY format
+                import shutil
+                ply_path = os.path.join(output_dir, "model.ply")
+                shutil.copy(mesh_path_orig, ply_path)
                 
-            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
+                processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
+                processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
             
-            # Clean up temporary file
-            if os.path.exists(filepath):
-                os.remove(filepath)
+            processing_jobs[job_id]['progress'] = 90
             
-            # Force garbage collection to free memory
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
+            # Create or check for video preview
+            video_path_orig = os.path.join(output_dir, "output.mp4")
+            if os.path.exists(video_path_orig):
+                preview_path = os.path.join(output_dir, "preview.mp4")
+                import shutil
+                shutil.copy(video_path_orig, preview_path)
+                processing_jobs[job_id]['preview_video'] = f"/preview-video/{job_id}"
+            
+            # Update job status
+            processing_jobs[job_id]['status'] = 'completed'
+            processing_jobs[job_id]['progress'] = 100
+            processing_jobs[job_id]['completed_at'] = time.time()
+            print(f"Job {job_id} completed successfully")
             
         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}"
+            processing_jobs[job_id]['error'] = f"Error during processing: {str(e)}"
             print(f"Error processing job {job_id}: {str(e)}")
             print(error_details)
-            
-            # Clean up on error
+            return
+        
+        finally:
+            # Clean up temporary file
             if os.path.exists(filepath):
-                os.remove(filepath)
+                try:
+                    os.remove(filepath)
+                except:
+                    pass
+            
+            # Force garbage collection to free memory
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
     
     # Start processing thread
     processing_thread = threading.Thread(target=process_image)
@@ -652,21 +729,26 @@ def model_info(job_id):
         "preview_url": job['preview_url'],
         "preview_video": job.get('preview_video'),
         "model_stats": model_stats,
-        "created_at": job.get('created_at'),
-        "completed_at": job.get('completed_at')
+    "created_at": job.get('created_at'),
+        "completed_at": job.get('completed_at'),
+        "cpu_only": True,  # Flag to indicate this was processed on CPU
+        "used_fallback": job.get('error') and "Using fallback" in job.get('error')
     }), 200
 
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "OpenLRM Image-to-3D Model Generator API", 
+        "message": "OpenLRM Image-to-3D Model Generator API (CPU Edition)", 
+        "status": "running",
+        "device": "cpu",
         "endpoints": [
             "/convert", 
             "/progress/<job_id>", 
             "/download/<job_id>", 
             "/preview/<job_id>",
             "/preview-video/<job_id>",
-            "/model-info/<job_id>"
+            "/model-info/<job_id>",
+            "/health"
         ],
         "parameters": {
             "output_format": "obj, glb, or ply",
@@ -674,13 +756,67 @@ def index():
             "source_cam_dist": "Camera distance from object (1.0-3.5, default 2.0)",
             "remove_bg": "true/false - automatically remove background"
         },
-        "description": "This API creates high-quality 3D models from 2D images with full structural completion from all angles"
+        "notes": "This CPU-only version has reduced capabilities compared to the GPU version. It uses a smaller model and may fall back to a simpler 3D generation method when needed.",
+        "description": "Creates 3D models from 2D images with structural completion"
     }), 200
 
+# Add a status endpoint for monitoring compute resources
+@app.route('/status', methods=['GET'])
+def system_status():
+    # Collect system information
+    import psutil
+    
+    try:
+        memory = psutil.virtual_memory()
+        disk = psutil.disk_usage('/')
+        
+        # Count active and completed jobs
+        active_jobs = sum(1 for j in processing_jobs.values() if j['status'] == 'processing')
+        completed_jobs = sum(1 for j in processing_jobs.values() if j['status'] == 'completed')
+        failed_jobs = sum(1 for j in processing_jobs.values() if j['status'] == 'error')
+        
+        return jsonify({
+            "status": "operational",
+            "system": {
+                "memory_used_percent": memory.percent,
+                "memory_available_mb": memory.available // (1024 * 1024),
+                "disk_used_percent": disk.percent,
+                "disk_free_gb": disk.free // (1024 * 1024 * 1024),
+                "cpu_percent": psutil.cpu_percent(interval=0.1)
+            },
+            "jobs": {
+                "active": active_jobs,
+                "completed": completed_jobs,
+                "failed": failed_jobs,
+                "total": len(processing_jobs)
+            },
+            "model_loaded": model_loaded,
+            "model_loading": model_loading
+        }), 200
+    except Exception as e:
+        return jsonify({
+            "status": "operational",
+            "error": str(e),
+            "jobs": {
+                "total": len(processing_jobs)
+            },
+            "model_loaded": model_loaded
+        }), 200
+
 if __name__ == '__main__':
     # Start the cleanup thread
     cleanup_old_jobs()
     
+    # Try installing dependencies at startup
+    try:
+        # Check if we need to install psutil
+        try:
+            import psutil
+        except ImportError:
+            subprocess.check_call(["pip", "install", "--no-cache-dir", "psutil"])
+    except:
+        print("Warning: Could not install psutil. Status endpoint may have limited functionality.")
+    
     # Use port 7860 which is standard for Hugging Face Spaces
     port = int(os.environ.get('PORT', 7860))
-    app.run(host='0.0.0.0', port=port)
+    app.run(host='0.0.0.0', port=port)