Update app.py

app.py

@@ -10,6 +10,7 @@ from PIL import Image
 import uuid
 import time
 import sys
+import gc  # For explicit garbage collection
 
 # Set environment variables before anything else
 os.environ['SHAPEE_NO_INTERACTIVE'] = '1'
@@ -118,12 +119,19 @@ def generate_3d():
         prompt = data['prompt']
         print(f"Received prompt: {prompt}")
         
-        # Set parameters for CPU performance (reduced steps)
+        # Set parameters for CPU performance (reduced steps and other optimizations)
         batch_size = 1
         guidance_scale = 15.0
         
+        # *** OPTIMIZATION: Significantly reduce steps for low-memory environments ***
+        karras_steps = 16  # Reduced from 32 to 16 for better performance
+        
+        # *** OPTIMIZATION: Run garbage collection before starting intensive task ***
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        
         # Generate latents with the text-to-3D model
-        print("Starting latent generation...")
+        print("Starting latent generation with reduced steps...")
         latents = sample_latents(
             batch_size=batch_size,
             model=model,
@@ -134,13 +142,17 @@ def generate_3d():
             clip_denoised=True,
             use_fp16=False,  # CPU doesn't support fp16
             use_karras=True,
-            karras_steps=32,  # Reduced steps for CPU
+            karras_steps=karras_steps,  # *** OPTIMIZATION: Reduced steps ***
             sigma_min=1e-3,
             sigma_max=160,
             s_churn=0,
         )
         print("Latent generation complete!")
         
+        # *** OPTIMIZATION: Run garbage collection after intensive step ***
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        
         # Generate a unique filename
         unique_id = str(uuid.uuid4())
         filename = f"{output_dir}/{unique_id}"
@@ -148,9 +160,16 @@ def generate_3d():
         # Convert latent to mesh
         print("Decoding mesh...")
         t0 = time.time()
+        
+        # *** OPTIMIZATION: Use simplified decoding for memory constraints ***
         mesh = decode_latent_mesh(xm, latents[0]).tri_mesh()
         print(f"Mesh decoded in {time.time() - t0:.2f} seconds")
         
+        # *** OPTIMIZATION: Clear latents from memory as they're no longer needed ***
+        del latents
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        
         # Save as GLB
         print("Saving as GLB...")
         glb_path = f"{filename}.glb"
@@ -162,6 +181,10 @@ def generate_3d():
         with open(obj_path, 'w') as f:
             mesh.write_obj(f)
         
+        # *** OPTIMIZATION: Clear mesh from memory ***
+        del mesh
+        gc.collect()
+        
         print("Files saved successfully!")
             
         # Return paths to the generated files
@@ -188,7 +211,19 @@ def download_file(filename):
 @app.route('/health', methods=['GET'])
 def health_check():
     """Simple health check endpoint to verify the app is running"""
-    return jsonify({"status": "ok", "message": "Service is running"})
+    # Check available memory
+    try:
+        import psutil
+        memory_info = psutil.virtual_memory()
+        memory_usage = f"{memory_info.percent}% (Available: {memory_info.available / (1024**3):.2f} GB)"
+    except ImportError:
+        memory_usage = "psutil not installed"
+    
+    return jsonify({
+        "status": "ok", 
+        "message": "Service is running",
+        "memory_usage": memory_usage
+    })
 
 @app.route('/', methods=['GET'])
 def home():
@@ -214,4 +249,6 @@ def home():
     """
 
 if __name__ == '__main__':
+    # Recommended to run with gunicorn for production with increased timeout:
+    # $ gunicorn app:app --bind 0.0.0.0:7860 --timeout 300 --workers 1
     app.run(host='0.0.0.0', port=7860, debug=True)