Update app.py

app.py

@@ -110,20 +110,20 @@ def get_adaptive_parameters():
     """Adjust parameters based on current system resources"""
     mem = psutil.virtual_memory()
     
-    # Base parameters
+    # Base parameters - more conservative to prevent memory issues
     params = {
-        'karras_steps': 8,
+        'karras_steps': 6,  # Reduced from 8 to 6 as default
         'batch_size': 1,
         'guidance_scale': 15.0
     }
     
     # If memory is tight, reduce steps further
     if mem.percent > 70:
-        params['karras_steps'] = 6
+        params['karras_steps'] = 4  # Even more conservative
     
     # If we have more memory to spare, can be slightly more generous
     if mem.percent < 50:
-        params['karras_steps'] = 10
+        params['karras_steps'] = 8
     
     print(f"Adaptive parameters chosen: karras_steps={params['karras_steps']}, mem={mem.percent}%")
     return params
@@ -131,13 +131,13 @@ def get_adaptive_parameters():
 def check_memory_pressure():
     """Check if memory is getting too high and take action if needed"""
     mem = psutil.virtual_memory()
-    if mem.percent > 85:  # Critical threshold
+    if mem.percent > 80:  # Reduced threshold from 85 to 80
         print("WARNING: Memory pressure critical. Forcing garbage collection.")
         gc.collect()
         torch.cuda.empty_cache() if torch.cuda.is_available() else None
         
         # If still critical, try more aggressive measures
-        if psutil.virtual_memory().percent > 80:
+        if psutil.virtual_memory().percent > 75:
             print("EMERGENCY: Memory still critical. Clearing model cache.")
             # Reset global models to force reload when memory is better
             global xm, model, diffusion
@@ -176,21 +176,87 @@ def model_unloader_thread():
     global xm, model, diffusion, last_usage_time
     
     while True:
-        time.sleep(300)  # Check every 5 minutes
+        time.sleep(180)  # Check more frequently: every 3 minutes instead of 5
         
         if last_usage_time is not None:
             idle_time = time.time() - last_usage_time
             
-            # If models have been idle for more than 10 minutes and no active jobs
-            if idle_time > 600 and active_jobs == 0:
-                # Check memory usage
+            # If models have been idle for more than 5 minutes (reduced from 10) and no active jobs
+            if idle_time > 300 and active_jobs == 0:
+                # Check memory usage - more aggressive unloading
                 mem = psutil.virtual_memory()
-                if mem.percent > 50:  # Only unload if memory usage is significant
+                if mem.percent > 40:  # Lowered threshold from 50 to 40
                     print(f"Models idle for {idle_time:.1f} seconds and memory at {mem.percent}%. Unloading...")
                     xm, model, diffusion = None, None, None
                     gc.collect()
                     torch.cuda.empty_cache() if torch.cuda.is_available() else None
 
+def save_trimesh(mesh, filename_base):
+    """Save mesh in multiple formats using trimesh"""
+    # Convert to trimesh format if needed
+    if not isinstance(mesh, trimesh.Trimesh):
+        try:
+            # Try to convert to trimesh
+            vertices = np.array(mesh.vertices)
+            faces = np.array(mesh.faces)
+            trimesh_obj = trimesh.Trimesh(vertices=vertices, faces=faces)
+        except Exception as e:
+            print(f"Error converting to trimesh: {e}")
+            raise
+    else:
+        trimesh_obj = mesh
+    
+    # Save as GLB
+    glb_path = f"{filename_base}.glb"
+    try:
+        trimesh_obj.export(glb_path, file_type='glb')
+        print(f"Saved GLB file: {glb_path}")
+    except Exception as e:
+        print(f"Error saving GLB: {e}")
+        # Try alternative approach
+        try:
+            scene = trimesh.Scene()
+            scene.add_geometry(trimesh_obj)
+            scene.export(glb_path)
+            print(f"Saved GLB using scene approach: {glb_path}")
+        except Exception as e2:
+            print(f"Alternative GLB export also failed: {e2}")
+            glb_path = None
+    
+    # Save as OBJ - always works more reliably
+    obj_path = f"{filename_base}.obj"
+    try:
+        trimesh_obj.export(obj_path, file_type='obj')
+        print(f"Saved OBJ file: {obj_path}")
+    except Exception as e:
+        print(f"Error saving OBJ: {e}")
+        # Try to write directly
+        try:
+            with open(obj_path, 'w') as f:
+                for v in trimesh_obj.vertices:
+                    f.write(f"v {v[0]} {v[1]} {v[2]}\n")
+                for face in trimesh_obj.faces:
+                    f.write(f"f {face[0]+1} {face[1]+1} {face[2]+1}\n")
+            print(f"Saved OBJ using direct write: {obj_path}")
+        except Exception as e2:
+            print(f"Alternative OBJ export also failed: {e2}")
+            obj_path = None
+    
+    # Also save as PLY as a fallback
+    ply_path = f"{filename_base}.ply"
+    try:
+        trimesh_obj.export(ply_path, file_type='ply')
+        print(f"Saved PLY file: {ply_path}")
+    except Exception as e:
+        print(f"Error saving PLY: {e}")
+        ply_path = None
+    
+    return {
+        "glb": os.path.basename(glb_path) if glb_path else None,
+        "obj": os.path.basename(obj_path) if obj_path else None,
+        "ply": os.path.basename(ply_path) if ply_path else None
+    }
+
 def process_job(job_id, prompt):
     try:
         # Get adaptive parameters
@@ -209,6 +275,7 @@ def process_job(job_id, prompt):
         print(f"Starting latent generation for job {job_id} with {karras_steps} steps...")
         
         # Generate latents
+        latents = None
         with torch.inference_mode():
             latents = sample_latents(
                 batch_size=batch_size,
@@ -227,7 +294,7 @@ def process_job(job_id, prompt):
             )
         print(f"Latent generation complete for job {job_id}!")
         
-        # Optimization: Clear unnecessary memory and load next model
+        # Optimization: Clear unnecessary memory and check pressure
         check_memory_pressure()
         
         # Generate a unique filename
@@ -245,7 +312,7 @@ def process_job(job_id, prompt):
         mem_before = psutil.Process().memory_info().rss / (1024 * 1024)
         print(f"Memory before mesh decoding: {mem_before:.2f} MB")
         
-        # Decode the mesh (fixed: removed 'max_points' parameter from original code)
+        # Decode the mesh
         mesh = decode_latent_mesh(xm, latents[0]).tri_mesh()
         
         print(f"Mesh decoded in {time.time() - t0:.2f} seconds")
@@ -254,50 +321,58 @@ def process_job(job_id, prompt):
         
         # Report mesh complexity if possible
         try:
-            print(f"Mesh complexity: {len(mesh.vertices)} vertices, {len(mesh.faces)} faces")
-        except:
-            print("Could not determine mesh complexity")
-        
-        # Simplify mesh if it's too complex (if supported)
-        try:
-            if hasattr(mesh, 'simplify') and hasattr(mesh, 'faces') and len(mesh.faces) > 5000:
-                target_faces = min(5000, int(len(mesh.faces) * 0.6))
-                print(f"Simplifying mesh to target {target_faces} faces...")
-                t0 = time.time()
-                simplified = mesh.simplify_quadratic_decimation(target_faces)
-                mesh = simplified
-                print(f"Mesh simplified in {time.time() - t0:.2f} seconds")
-                print(f"New complexity: {len(mesh.vertices)} vertices, {len(mesh.faces)} faces")
+            vertices_count = len(mesh.vertices)
+            faces_count = len(mesh.faces)
+            print(f"Mesh complexity: {vertices_count} vertices, {faces_count} faces")
         except Exception as e:
-            print(f"Mesh simplification not available or failed: {e}")
+            print(f"Could not determine mesh complexity: {e}")
+            vertices_count = 0
+            faces_count = 0
         
         # Clear latents from memory
         del latents
         gc.collect()
         
-        # Save as GLB
-        print(f"Saving job {job_id} as GLB...")
-        glb_path = f"{filename}.glb"
-        mesh.write_glb(glb_path)
+        # Convert to trimesh format and save files
+        print(f"Converting and saving mesh for job {job_id}...")
         
-        # Save as OBJ
-        print(f"Saving job {job_id} as OBJ...")
-        obj_path = f"{filename}.obj"
-        with open(obj_path, 'w') as f:
-            mesh.write_obj(f)
+        # Save mesh using the helper function
+        saved_files = save_trimesh(mesh, filename)
         
         # Clear mesh from memory
         del mesh
         gc.collect()
         
-        print(f"Files saved successfully for job {job_id}!")
-            
-        return {
+        # Check which files were successfully saved
+        result = {
             "success": True,
             "message": "3D model generated successfully",
-            "glb_url": f"/download/{os.path.basename(glb_path)}",
-            "obj_url": f"/download/{os.path.basename(obj_path)}"
+            "timestamp": time.time(),
+            "stats": {
+                "vertices": vertices_count,
+                "faces": faces_count
+            }
         }
+        
+        # Add URLs for the files that were saved
+        if saved_files["glb"]:
+            result["glb_url"] = f"/download/{saved_files['glb']}"
+        if saved_files["obj"]:
+            result["obj_url"] = f"/download/{saved_files['obj']}"
+        if saved_files["ply"]:
+            result["ply_url"] = f"/download/{saved_files['ply']}"
+        
+        # If no files were saved, mark as failure
+        if not (saved_files["glb"] or saved_files["obj"] or saved_files["ply"]):
+            result["success"] = False
+            result["message"] = "Failed to save mesh in any format"
+        
+        print(f"Files saved successfully for job {job_id}!")
+        
+        # Force garbage collection again
+        gc.collect()
+        
+        return result
     
     except Exception as e:
         print(f"Error during generation for job {job_id}: {str(e)}")
@@ -305,7 +380,8 @@ def process_job(job_id, prompt):
         traceback.print_exc()
         return {
             "success": False,
-            "error": str(e)
+            "error": str(e),
+            "timestamp": time.time()
         }
 
 def worker_thread():
@@ -326,6 +402,9 @@ def worker_thread():
                 job_results[job_id] = result
                 active_jobs -= 1
                 
+                # Explicit cleanup after job
+                gc.collect()
+                
             except queue.Empty:
                 # No jobs in queue, continue waiting
                 pass
@@ -337,12 +416,43 @@ def worker_thread():
                 if 'job_id' in locals():
                     job_results[job_id] = {
                         "success": False,
-                        "error": str(e)
+                        "error": str(e),
+                        "timestamp": time.time()
                     }
                     active_jobs -= 1
+                    
+                # Force garbage collection to clean up
+                gc.collect()
     finally:
         is_thread_running = False
 
+def purge_old_results_thread():
+    """Thread that periodically cleans up old job results to manage memory"""
+    while True:
+        try:
+            time.sleep(1800)  # Run every 30 minutes
+            
+            # Default threshold: 2 hours
+            threshold_time = time.time() - (2 * 3600)
+            
+            # Track jobs to be removed
+            jobs_to_remove = []
+            for job_id, result in job_results.items():
+                # If the job has a timestamp and it's older than threshold
+                if result.get('timestamp', time.time()) < threshold_time:
+                    jobs_to_remove.append(job_id)
+            
+            # Remove the old jobs
+            for job_id in jobs_to_remove:
+                job_results.pop(job_id, None)
+                
+            if jobs_to_remove:
+                print(f"Auto-purged {len(jobs_to_remove)} old job results")
+                # Force garbage collection
+                gc.collect()
+        except Exception as e:
+            print(f"Error in purge thread: {e}")
+
 def ensure_worker_thread_running():
     global generation_thread, is_thread_running
     
@@ -351,8 +461,13 @@ def ensure_worker_thread_running():
         generation_thread = threading.Thread(target=worker_thread, daemon=True)
         generation_thread.start()
 
-def start_model_unloader():
+def start_monitoring_threads():
+    """Start all monitoring and maintenance threads"""
+    # Start model unloader thread
     threading.Thread(target=model_unloader_thread, daemon=True).start()
+    
+    # Start results purge thread
+    threading.Thread(target=purge_old_results_thread, daemon=True).start()
 
 @app.route('/generate', methods=['POST'])
 def generate_3d():
@@ -394,7 +509,7 @@ def generate_3d():
 def job_status(job_id):
     if job_id in job_results:
         result = job_results[job_id]
-        # Clean up memory if the job is complete and successful
+        # Return the result
         return jsonify(result)
     else:
         # Job is still in progress
@@ -407,9 +522,13 @@ def job_status(job_id):
 @app.route('/download/<filename>', methods=['GET'])
 def download_file(filename):
     try:
-        return send_file(f"{output_dir}/{filename}", as_attachment=True)
+        file_path = os.path.join(output_dir, filename)
+        if not os.path.exists(file_path):
+            return jsonify({"error": "File not found"}), 404
+            
+        return send_file(file_path, as_attachment=True)
     except Exception as e:
-        return jsonify({"error": str(e)}), 404
+        return jsonify({"error": str(e)}), 500
 
 @app.route('/health', methods=['GET'])
 def health_check():
@@ -443,6 +562,9 @@ def health_check():
         if last_usage_time is not None:
             model_inactive = f"{(time.time() - last_usage_time) / 60:.1f} minutes"
         
+        # Number of saved jobs
+        saved_jobs = len(job_results)
+        
         return jsonify({
             "status": "ok",
             "message": "Service is running",
@@ -451,6 +573,7 @@ def health_check():
             "cpu_usage": cpu_usage,
             "queue_size": queue_size,
             "active_jobs": active_jobs,
+            "saved_jobs": saved_jobs,
             "worker_running": is_thread_running,
             "models_loaded": models_loaded,
             "model_inactive_time": model_inactive
@@ -516,12 +639,13 @@ GET /status/123e4567-e89b-12d3-a456-426614174000
     "success": true,
     "message": "3D model generated successfully",
     "glb_url": "/download/abc123.glb",
-    "obj_url": "/download/abc123.obj"
+    "obj_url": "/download/abc123.obj",
+    "ply_url": "/download/abc123.ply"
 }
             </pre>
             
             <h3>Step 3: Download the files</h3>
-            <p>Use the provided URLs to download the GLB and OBJ files.</p>
+            <p>Use the provided URLs to download the GLB, OBJ, and PLY files.</p>
             
             <h2>Health Check:</h2>
             <pre>GET /health</pre>
@@ -563,10 +687,40 @@ def purge_old_results():
             "error": str(e)
         }), 500
 
+@app.route('/force-gc', methods=['POST'])
+def force_garbage_collection():
+    """Endpoint to manually trigger garbage collection"""
+    try:
+        # Get current memory usage
+        before_mem = psutil.Process().memory_info().rss / (1024**3)
+        
+        # Force garbage collection
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        
+        # Get memory usage after GC
+        after_mem = psutil.Process().memory_info().rss / (1024**3)
+        freed = before_mem - after_mem
+        
+        return jsonify({
+            "success": True,
+            "message": f"Garbage collection completed",
+            "before_memory_gb": round(before_mem, 2),
+            "after_memory_gb": round(after_mem, 2),
+            "freed_memory_gb": round(freed, 2) if freed > 0 else 0
+        })
+    except Exception as e:
+        return jsonify({
+            "success": False,
+            "error": str(e)
+        }), 500
+
 if __name__ == '__main__':
-    # Start the worker thread and model unloader
+    # Start all monitoring threads
+    start_monitoring_threads()
+    
+    # Start the worker thread
     ensure_worker_thread_running()
-    start_model_unloader()
     
     # Recommended to run with gunicorn for production with increased timeout:
     # $ gunicorn app:app --bind 0.0.0.0:7860 --timeout 300 --workers 1