Update app.py

app.py

@@ -78,6 +78,22 @@ class JobStatusResponse(BaseModel):
     created_at: float
     updated_at: float
 
+class MemoryClearanceRequest(BaseModel):
+    clear_models: bool = True
+    clear_jobs: bool = False
+    clear_local_images: bool = False
+    force_gc: bool = True
+
+class MemoryStatusResponse(BaseModel):
+    memory_used_mb: float
+    memory_percent: float
+    models_loaded: int
+    active_jobs: int
+    local_images_count: int
+    gpu_memory_allocated_mb: Optional[float] = None
+    gpu_memory_cached_mb: Optional[float] = None
+    status: str
+
 # HIGH-QUALITY MODEL SELECTION - ANIME FOCUSED & WORKING
 MODEL_CHOICES = {
     "dreamshaper-8": "lykon/dreamshaper-8",        # Great all-rounder
@@ -97,6 +113,103 @@ current_model_name = None
 current_pipe = None
 model_lock = threading.Lock()
 
+# MEMORY MANAGEMENT FUNCTIONS
+def get_memory_usage():
+    """Get current memory usage statistics"""
+    process = psutil.Process()
+    memory_info = process.memory_info()
+    memory_used_mb = memory_info.rss / (1024 * 1024)
+    memory_percent = process.memory_percent()
+    
+    # GPU memory if available
+    gpu_memory_allocated_mb = None
+    gpu_memory_cached_mb = None
+    
+    if torch.cuda.is_available():
+        gpu_memory_allocated_mb = torch.cuda.memory_allocated() / (1024 * 1024)
+        gpu_memory_cached_mb = torch.cuda.memory_reserved() / (1024 * 1024)
+    
+    return {
+        "memory_used_mb": round(memory_used_mb, 2),
+        "memory_percent": round(memory_percent, 2),
+        "gpu_memory_allocated_mb": round(gpu_memory_allocated_mb, 2) if gpu_memory_allocated_mb else None,
+        "gpu_memory_cached_mb": round(gpu_memory_cached_mb, 2) if gpu_memory_cached_mb else None,
+        "models_loaded": len(model_cache),
+        "active_jobs": len(job_storage),
+        "local_images_count": len(refresh_local_images())
+    }
+
+def clear_memory(clear_models=True, clear_jobs=False, clear_local_images=False, force_gc=True):
+    """Clear memory by unloading models and cleaning up resources"""
+    results = []
+    
+    # Clear model cache
+    if clear_models:
+        with model_lock:
+            models_cleared = len(model_cache)
+            for model_name, pipe in model_cache.items():
+                try:
+                    # Move to CPU first if it's on GPU
+                    if hasattr(pipe, 'to'):
+                        pipe.to('cpu')
+                    
+                    # Delete the pipeline
+                    del pipe
+                    results.append(f"Unloaded model: {model_name}")
+                except Exception as e:
+                    results.append(f"Error unloading {model_name}: {str(e)}")
+            
+            model_cache.clear()
+            global current_pipe, current_model_name
+            current_pipe = None
+            current_model_name = None
+            results.append(f"Cleared {models_cleared} models from cache")
+    
+    # Clear completed jobs
+    if clear_jobs:
+        jobs_to_clear = []
+        for job_id, job_data in job_storage.items():
+            if job_data["status"] in [JobStatus.COMPLETED, JobStatus.FAILED]:
+                jobs_to_clear.append(job_id)
+        
+        for job_id in jobs_to_clear:
+            del job_storage[job_id]
+            results.append(f"Cleared job: {job_id}")
+        
+        results.append(f"Cleared {len(jobs_to_clear)} completed/failed jobs")
+    
+    # Clear local images
+    if clear_local_images:
+        try:
+            storage_info = get_local_storage_info()
+            deleted_count = 0
+            if "images" in storage_info:
+                for image_info in storage_info["images"]:
+                    success, _ = delete_local_image(image_info["path"])
+                    if success:
+                        deleted_count += 1
+            results.append(f"Deleted {deleted_count} local images")
+        except Exception as e:
+            results.append(f"Error clearing local images: {str(e)}")
+    
+    # Force garbage collection
+    if force_gc:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+            results.append("GPU cache cleared")
+        results.append("Garbage collection forced")
+    
+    # Get memory status after cleanup
+    memory_status = get_memory_usage()
+    
+    return {
+        "status": "success",
+        "actions_performed": results,
+        "memory_after_cleanup": memory_status
+    }
+
 def load_model(model_name="dreamshaper-8"):
     """Thread-safe model loading with HIGH-QUALITY settings and better error handling"""
     global model_cache, current_model_name, current_pipe
@@ -732,6 +845,62 @@ async def api_health():
         "oci_api_connected": OCI_API_BASE_URL
     }
 
+# NEW MEMORY MANAGEMENT ENDPOINTS
+@app.get("/api/memory-status")
+async def get_memory_status():
+    """Get current memory usage and system status"""
+    memory_info = get_memory_usage()
+    return MemoryStatusResponse(
+        memory_used_mb=memory_info["memory_used_mb"],
+        memory_percent=memory_info["memory_percent"],
+        models_loaded=memory_info["models_loaded"],
+        active_jobs=memory_info["active_jobs"],
+        local_images_count=memory_info["local_images_count"],
+        gpu_memory_allocated_mb=memory_info["gpu_memory_allocated_mb"],
+        gpu_memory_cached_mb=memory_info["gpu_memory_cached_mb"],
+        status="healthy"
+    )
+
+@app.post("/api/clear-memory")
+async def clear_memory_endpoint(request: MemoryClearanceRequest):
+    """Clear memory by unloading models and cleaning up resources"""
+    try:
+        result = clear_memory(
+            clear_models=request.clear_models,
+            clear_jobs=request.clear_jobs,
+            clear_local_images=request.clear_local_images,
+            force_gc=request.force_gc
+        )
+        
+        return {
+            "status": "success",
+            "message": "Memory clearance completed",
+            "details": result
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Memory clearance failed: {str(e)}")
+
+@app.post("/api/auto-cleanup")
+async def auto_cleanup():
+    """Automatic cleanup - clears completed jobs and forces GC"""
+    try:
+        result = clear_memory(
+            clear_models=False,  # Don't clear models by default
+            clear_jobs=True,     # Clear completed jobs
+            clear_local_images=False,  # Don't clear images by default
+            force_gc=True        # Force garbage collection
+        )
+        
+        return {
+            "status": "success",
+            "message": "Automatic cleanup completed",
+            "details": result
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Auto cleanup failed: {str(e)}")
+
 @app.get("/api/local-images")
 async def get_local_images():
     """API endpoint to get locally saved test images"""
@@ -793,12 +962,30 @@ def create_gradio_interface():
             lines=2
         )
         
+        # Memory status display
+        memory_status = gr.Textbox(
+            label="🧠 Memory Status",
+            interactive=False,
+            lines=3
+        )
+        
         def update_storage_info():
             info = get_local_storage_info()
             if "error" not in info:
                 return f"📁 Local Storage: {info['total_files']} images, {info['total_size_mb']} MB used"
             return "📁 Local Storage: Unable to calculate"
         
+        def update_memory_status():
+            memory_info = get_memory_usage()
+            status_text = f"🧠 Memory Usage: {memory_info['memory_used_mb']} MB ({memory_info['memory_percent']}%)\n"
+            status_text += f"📦 Models Loaded: {memory_info['models_loaded']}\n"
+            status_text += f"⚡ Active Jobs: {memory_info['active_jobs']}"
+            
+            if memory_info['gpu_memory_allocated_mb']:
+                status_text += f"\n🎮 GPU Memory: {memory_info['gpu_memory_allocated_mb']} MB allocated"
+            
+            return status_text
+        
         with gr.Row():
             with gr.Column(scale=1):
                 gr.Markdown("### 🎯 Quality Settings")
@@ -828,12 +1015,25 @@ def create_gradio_interface():
                 delete_btn = gr.Button("🗑️ Delete This Image", variant="stop")
                 delete_status = gr.Textbox(label="Delete Status", interactive=False, lines=2)
                 
+                # Memory management section
+                gr.Markdown("### 🧠 Memory Management")
+                with gr.Row():
+                    auto_cleanup_btn = gr.Button("🔄 Auto Cleanup", size="sm")
+                    clear_models_btn = gr.Button("🗑️ Clear Models", variant="stop", size="sm")
+                
+                memory_clear_status = gr.Textbox(label="Memory Clear Status", interactive=False, lines=2)
+                
                 gr.Markdown("### 📚 API Usage for n8n")
                 gr.Markdown("""
                 **For complete storybooks (OCI bucket):**
                 - Endpoint: `POST /api/generate-storybook`
                 - Input: `story_title`, `scenes[]`, `characters[]`
                 - Output: Uses pure prompts only from your script
+                
+                **Memory Management APIs:**
+                - `GET /api/memory-status` - Check memory usage
+                - `POST /api/clear-memory` - Clear memory
+                - `POST /api/auto-cleanup` - Auto cleanup jobs
                 """)
             
             with gr.Column(scale=2):
@@ -888,6 +1088,32 @@ def create_gradio_interface():
                 return f"✅ Deleted {deleted_count} images", updated_files
             except Exception as e:
                 return f"❌ Error: {str(e)}", refresh_local_images()
+        
+        def perform_auto_cleanup():
+            """Perform automatic cleanup"""
+            try:
+                result = clear_memory(
+                    clear_models=False,
+                    clear_jobs=True,
+                    clear_local_images=False,
+                    force_gc=True
+                )
+                return f"✅ Auto cleanup completed: {len(result['actions_performed'])} actions"
+            except Exception as e:
+                return f"❌ Auto cleanup failed: {str(e)}"
+        
+        def clear_models():
+            """Clear all loaded models"""
+            try:
+                result = clear_memory(
+                    clear_models=True,
+                    clear_jobs=False,
+                    clear_local_images=False,
+                    force_gc=True
+                )
+                return f"✅ Models cleared: {len(result['actions_performed'])} actions"
+            except Exception as e:
+                return f"❌ Model clearance failed: {str(e)}"
 
         # Connect buttons to functions
         generate_btn.click(
@@ -900,6 +1126,9 @@ def create_gradio_interface():
         ).then(
             fn=update_storage_info,
             outputs=storage_info
+        ).then(
+            fn=update_memory_status,
+            outputs=memory_status
         )
         
         delete_btn.click(
@@ -909,6 +1138,9 @@ def create_gradio_interface():
         ).then(
             fn=update_storage_info,
             outputs=storage_info
+        ).then(
+            fn=update_memory_status,
+            outputs=memory_status
         )
         
         refresh_btn.click(
@@ -917,6 +1149,9 @@ def create_gradio_interface():
         ).then(
             fn=update_storage_info,
             outputs=storage_info
+        ).then(
+            fn=update_memory_status,
+            outputs=memory_status
         )
         
         clear_all_btn.click(
@@ -925,11 +1160,32 @@ def create_gradio_interface():
         ).then(
             fn=update_storage_info,
             outputs=storage_info
+        ).then(
+            fn=update_memory_status,
+            outputs=memory_status
+        )
+        
+        # Memory management buttons
+        auto_cleanup_btn.click(
+            fn=perform_auto_cleanup,
+            outputs=memory_clear_status
+        ).then(
+            fn=update_memory_status,
+            outputs=memory_status
+        )
+        
+        clear_models_btn.click(
+            fn=clear_models,
+            outputs=memory_clear_status
+        ).then(
+            fn=update_memory_status,
+            outputs=memory_status
         )
         
         # Initialize on load
         demo.load(fn=refresh_local_images, outputs=file_gallery)
         demo.load(fn=update_storage_info, outputs=storage_info)
+        demo.load(fn=update_memory_status, outputs=memory_status)
     
     return demo
 
@@ -945,11 +1201,15 @@ async def root():
             "health_check": "GET /api/health",
             "generate_storybook": "POST /api/generate-storybook",
             "check_job_status": "GET /api/job-status/{job_id}",
-            "local_images": "GET /api/local-images"
+            "local_images": "GET /api/local-images",
+            "memory_status": "GET /api/memory-status",
+            "clear_memory": "POST /api/clear-memory",
+            "auto_cleanup": "POST /api/auto-cleanup"
         },
         "features": {
             "pure_prompts": "✅ Enabled - No automatic enhancements",
-            "n8n_integration": "✅ Enabled"
+            "n8n_integration": "✅ Enabled",
+            "memory_management": "✅ Enabled"
         },
         "web_interface": "GET /ui"
     }
@@ -961,6 +1221,7 @@ async def test_endpoint():
         "status": "success", 
         "message": "API with pure prompts is working correctly",
         "pure_prompts": "✅ Enabled - Using exact prompts from Telegram",
+        "memory_management": "✅ Enabled - Memory clearance available",
         "timestamp": datetime.now().isoformat()
     }
 
@@ -980,6 +1241,7 @@ if __name__ == "__main__":
         print("📚 API endpoints available at: /api/*")
         print("🎨 Web interface available at: /ui")
         print("📝 PURE PROMPTS enabled - no automatic enhancements")
+        print("🧠 MEMORY MANAGEMENT enabled - automatic cleanup available")
         
         # Mount Gradio without reassigning app
         gr.mount_gradio_app(app, demo, path="/ui")
@@ -997,6 +1259,7 @@ if __name__ == "__main__":
         print("📚 API endpoints: http://localhost:8000/api/*")
         print("🎨 Web interface: http://localhost:7860/ui")
         print("📝 PURE PROMPTS enabled - no automatic enhancements")
+        print("🧠 MEMORY MANAGEMENT enabled - automatic cleanup available")
         
         def run_fastapi():
             """Run FastAPI on port 8000 for API calls"""
@@ -1010,6 +1273,19 @@ if __name__ == "__main__":
         
         def run_gradio():
             """Run Gradio on port 7860 for web interface"""
-            demo.launch(
-                server_name="0.0.0.0",
-            )
+            demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+        
+        # Run both servers in separate threads
+        import threading
+        fastapi_thread = threading.Thread(target=run_fastapi, daemon=True)
+        gradio_thread = threading.Thread(target=run_gradio, daemon=True)
+        
+        fastapi_thread.start()
+        gradio_thread.start()
+        
+        try:
+            # Keep main thread alive
+            while True:
+                time.sleep(1)
+        except KeyboardInterrupt:
+            print("🛑 Shutting down servers...")