Update app.py

app.py

@@ -3,24 +3,32 @@ import torch
 import os
 import time
 import threading
-from datetime import datetime
+from datetime import datetime, timedelta
 import cv2
 from werkzeug.utils import secure_filename
 import uuid
 import mimetypes
 import numpy as np
 from PIL import Image
+import schedule
 
 # Configuration
 UPLOAD_FOLDER = '/data/uploads'
 OUTPUT_FOLDER = '/data/outputs'
+CLEANUP_INTERVAL_MINUTES = 10
+FILE_MAX_AGE_HOURS = 1
 
 # Global application state
 app_state = {
     "cuda_available": torch.cuda.is_available(),
     "processing_active": False,
     "logs": [],
-    "processed_files": []
+    "processed_files": [],
+    "cleanup_stats": {
+        "last_cleanup": None,
+        "files_deleted": 0,
+        "space_freed_mb": 0
+    }
 }
 
 def ensure_directories():
@@ -59,6 +67,80 @@ def log_message(message):
         app_state["logs"] = app_state["logs"][-100:]
     print(f"[{timestamp}] {message}")
 
+def cleanup_old_files():
+    """Delete files older than FILE_MAX_AGE_HOURS"""
+    try:
+        current_time = datetime.now()
+        cutoff_time = current_time - timedelta(hours=FILE_MAX_AGE_HOURS)
+        
+        files_deleted = 0
+        space_freed = 0
+        
+        # Clean upload folder
+        for folder_path in [UPLOAD_FOLDER, OUTPUT_FOLDER]:
+            if not os.path.exists(folder_path):
+                continue
+                
+            for filename in os.listdir(folder_path):
+                file_path = os.path.join(folder_path, filename)
+                
+                if os.path.isfile(file_path):
+                    try:
+                        # Get file modification time
+                        file_time = datetime.fromtimestamp(os.path.getmtime(file_path))
+                        
+                        if file_time < cutoff_time:
+                            # Get file size before deletion
+                            file_size = os.path.getsize(file_path)
+                            
+                            # Delete the file
+                            os.remove(file_path)
+                            
+                            files_deleted += 1
+                            space_freed += file_size
+                            
+                            log_message(f"🗑️ Deleted old file: {filename} ({file_size / (1024*1024):.1f}MB)")
+                    
+                    except Exception as e:
+                        log_message(f"⚠️ Error deleting {filename}: {str(e)}")
+        
+        # Update cleanup stats
+        app_state["cleanup_stats"]["last_cleanup"] = current_time.strftime("%Y-%m-%d %H:%M:%S")
+        app_state["cleanup_stats"]["files_deleted"] += files_deleted
+        app_state["cleanup_stats"]["space_freed_mb"] += space_freed / (1024*1024)
+        
+        if files_deleted > 0:
+            log_message(f"🧹 Cleanup completed: {files_deleted} files deleted, {space_freed / (1024*1024):.1f}MB freed")
+        else:
+            log_message(f"🧹 Cleanup completed: No old files to delete")
+            
+        # Clean up processed files list to remove references to deleted files
+        valid_processed_files = []
+        for file_info in app_state["processed_files"]:
+            output_path = os.path.join(OUTPUT_FOLDER, file_info["output_file"])
+            if os.path.exists(output_path):
+                valid_processed_files.append(file_info)
+        
+        app_state["processed_files"] = valid_processed_files
+        
+    except Exception as e:
+        log_message(f"❌ Error during cleanup: {str(e)}")
+
+def run_scheduler():
+    """Run the file cleanup scheduler in background"""
+    def scheduler_worker():
+        while True:
+            try:
+                schedule.run_pending()
+                time.sleep(60)  # Check every minute
+            except Exception as e:
+                log_message(f"❌ Scheduler error: {str(e)}")
+                time.sleep(300)  # Wait 5 minutes before retrying
+    
+    thread = threading.Thread(target=scheduler_worker, daemon=True)
+    thread.start()
+    log_message(f"🕒 File cleanup scheduler started (every {CLEANUP_INTERVAL_MINUTES} minutes)")
+
 def optimize_gpu():
     """Optimize GPU configuration for 4K upscaling"""
     try:
@@ -444,6 +526,9 @@ def process_frame_batch(frame_batch, out, device, target_h, target_w):
 # Initialize directories
 ensure_directories()
 
+# Set up file cleanup scheduler
+schedule.every(CLEANUP_INTERVAL_MINUTES).minutes.do(cleanup_old_files)
+
 app = Flask(__name__)
 
 @app.route('/')
@@ -494,6 +579,12 @@ def api_system():
                 info["storage_outputs"] = f"{output_size / (1024**2):.1f}MB"
                 info["upload_files_count"] = len(upload_files)
                 info["output_files_count"] = len(output_files)
+                
+                # Add cleanup info
+                info["cleanup_stats"] = app_state["cleanup_stats"]
+                info["cleanup_interval"] = f"{CLEANUP_INTERVAL_MINUTES} minutes"
+                info["file_max_age"] = f"{FILE_MAX_AGE_HOURS} hour(s)"
+                
             except Exception as e:
                 info["storage_uploads"] = f"Error: {str(e)}"
                 info["storage_outputs"] = "N/A"
@@ -643,19 +734,76 @@ def api_clear_cache():
     except Exception as e:
         return jsonify({"success": False, "error": str(e)})
 
+@app.route('/api/cleanup-now', methods=['POST'])
+def api_cleanup_now():
+    """Manually trigger file cleanup"""
+    try:
+        cleanup_old_files()
+        return jsonify({"success": True, "message": "Manual cleanup completed"})
+    except Exception as e:
+        return jsonify({"success": False, "error": str(e)})
+
+@app.route('/api/storage-stats')
+def api_storage_stats():
+    """Get detailed storage statistics"""
+    try:
+        stats = {
+            "cleanup_stats": app_state["cleanup_stats"],
+            "current_files": {},
+            "total_storage_mb": 0
+        }
+        
+        for folder_name, folder_path in [("uploads", UPLOAD_FOLDER), ("outputs", OUTPUT_FOLDER)]:
+            if os.path.exists(folder_path):
+                files = []
+                total_size = 0
+                
+                for filename in os.listdir(folder_path):
+                    file_path = os.path.join(folder_path, filename)
+                    if os.path.isfile(file_path):
+                        file_size = os.path.getsize(file_path)
+                        file_time = datetime.fromtimestamp(os.path.getmtime(file_path))
+                        
+                        files.append({
+                            "name": filename,
+                            "size_mb": file_size / (1024*1024),
+                            "created": file_time.strftime("%Y-%m-%d %H:%M:%S"),
+                            "age_hours": (datetime.now() - file_time).total_seconds() / 3600
+                        })
+                        total_size += file_size
+                
+                stats["current_files"][folder_name] = {
+                    "files": files,
+                    "count": len(files),
+                    "total_size_mb": total_size / (1024*1024)
+                }
+                stats["total_storage_mb"] += total_size / (1024*1024)
+        
+        return jsonify({"success": True, "data": stats})
+    except Exception as e:
+        return jsonify({"success": False, "error": str(e)})
+
 if __name__ == '__main__':
     # Initialize system
     log_message("🚀 4K Upscaler starting...")
     
     try:
+        # Start file cleanup scheduler
+        run_scheduler()
+        
         # Optimize GPU if available
         if optimize_gpu():
             log_message("✅ GPU optimized for 4K upscaling")
         else:
             log_message("⚠️ GPU optimization failed, using CPU fallback")
         
+        # Run initial cleanup
+        log_message("🧹 Running initial file cleanup...")
+        cleanup_old_files()
+        
         log_message("✅ 4K Upscaler ready")
         log_message("📤 Upload images or videos to upscale to 4K resolution")
+        log_message(f"🗑️ Files will be automatically deleted after {FILE_MAX_AGE_HOURS} hour(s)")
         
     except Exception as e:
         log_message(f"❌ Initialization error: {str(e)}")