Update app.py

app.py

@@ -47,6 +47,32 @@ def ensure_directories():
         except Exception as e:
             print(f"⚠️ Error creating directory {directory}: {e}")
 
+def allowed_file(filename):
+    """Check if file has allowed extension"""
+    return '.' in filename and \
+           filename.rsplit('.', 1)[1].lower() in ['png', 'jpg', 'jpeg', 'gif', 'mp4', 'avi', 'mov', 'mkv']
+
+def get_file_mimetype(filename):
+    """Get correct mimetype for file"""
+    mimetype, _ = mimetypes.guess_type(filename)
+    if mimetype is None:
+        ext = filename.lower().rsplit('.', 1)[1] if '.' in filename else ''
+        if ext in ['mp4', 'avi', 'mov', 'mkv']:
+            mimetype = f'video/{ext}'
+        elif ext in ['png', 'jpg', 'jpeg', 'gif']:
+            mimetype = f'image/{ext}'
+        else:
+            mimetype = 'application/octet-stream'
+    return mimetype
+
+def log_message(message):
+    """Add message to log with timestamp"""
+    timestamp = datetime.now().strftime("%H:%M:%S")
+    app_state["logs"].append(f"[{timestamp}] {message}")
+    if len(app_state["logs"]) > 100:
+        app_state["logs"] = app_state["logs"][-100:]
+    print(f"[{timestamp}] {message}")
+
 def download_realesrgan_models():
     """Download Real-ESRGAN models if not present"""
     if not REALESRGAN_AVAILABLE:
@@ -116,6 +142,28 @@ def initialize_realesrgan(model_name='RealESRGAN_x4plus', scale=4):
         log_message(f"❌ Error initializing Real-ESRGAN: {str(e)}")
         return None
 
+def optimize_gpu():
+    """Optimize GPU configuration for 4K upscaling"""
+    try:
+        if torch.cuda.is_available():
+            torch.backends.cudnn.benchmark = True
+            torch.backends.cudnn.allow_tf32 = True
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.cuda.empty_cache()
+            
+            # Test GPU
+            test_tensor = torch.randn(100, 100, device='cuda')
+            _ = torch.mm(test_tensor, test_tensor)
+            
+            log_message("✅ GPU optimized for 4K upscaling")
+            return True
+        else:
+            log_message("⚠️ CUDA not available")
+            return False
+    except Exception as e:
+        log_message(f"❌ Error optimizing GPU: {str(e)}")
+        return False
+
 def upscale_image_realesrgan(input_path, output_path):
     """Upscale image using Real-ESRGAN"""
     def process_worker():
@@ -209,7 +257,7 @@ def upscale_image_realesrgan(input_path, output_path):
     thread.start()
 
 def upscale_image_4k_fallback(input_path, output_path):
-    """Fallback upscaling method (your original implementation)"""
+    """Fallback upscaling method"""
     try:
         log_message("🔄 Using fallback upscaling method")
         
@@ -219,7 +267,7 @@ def upscale_image_4k_fallback(input_path, output_path):
         
         if torch.cuda.is_available():
             device = torch.device('cuda')
-            # Your original GPU implementation
+            # GPU implementation
             image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
             image_tensor = torch.from_numpy(image_rgb).float().to(device) / 255.0
             image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0)
@@ -264,85 +312,152 @@ def upscale_image_4k_fallback(input_path, output_path):
     except Exception as e:
         log_message(f"❌ Error in fallback processing: {str(e)}")
 
-# Update your existing upload endpoint
-@app.route('/api/upload', methods=['POST'])
-def api_upload():
-    """Upload and process file for 4K upscaling"""
-    try:
-        if 'file' not in request.files:
-            return jsonify({"success": False, "error": "No file provided"})
-        
-        file = request.files['file']
-        if file.filename == '':
-            return jsonify({"success": False, "error": "No file selected"})
-        
-        if file and allowed_file(file.filename):
-            file_id = str(uuid.uuid4())
-            filename = secure_filename(file.filename)
-            file_ext = filename.rsplit('.', 1)[1].lower()
+def upscale_video_4k(input_path, output_path):
+    """Upscale video to 4K frame by frame"""
+    def process_worker():
+        try:
+            log_message(f"🎬 Starting 4K video upscaling: {os.path.basename(input_path)}")
+            app_state["processing_active"] = True
             
-            input_filename = f"{file_id}_input.{file_ext}"
-            input_path = os.path.join(UPLOAD_FOLDER, input_filename)
-            file.save(input_path)
+            # Open video
+            cap = cv2.VideoCapture(input_path)
+            if not cap.isOpened():
+                log_message("❌ Error: Could not open video")
+                return
             
-            output_filename = f"{file_id}_4k.{file_ext}"
-            output_path = os.path.join(OUTPUT_FOLDER, output_filename)
+            # Get video properties
+            fps = int(cap.get(cv2.CAP_PROP_FPS))
+            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            log_message(f"📹 Video: {w}x{h}, {fps}FPS, {frame_count} frames")
             
-            if file_ext in ['png', 'jpg', 'jpeg', 'gif']:
-                # Use Real-ESRGAN for images if available
-                if REALESRGAN_AVAILABLE:
-                    upscale_image_realesrgan(input_path, output_path)
-                else:
-                    upscale_image_4k_fallback(input_path, output_path)
-                media_type = "image"
-            elif file_ext in ['mp4', 'avi', 'mov', 'mkv']:
-                # Keep your existing video processing
-                upscale_video_4k(input_path, output_path)
-                media_type = "video"
+            # Configure 4K output
+            target_w, target_h = w * 4, h * 4
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            out = cv2.VideoWriter(output_path, fourcc, fps, (target_w, target_h))
             
-            log_message(f"📤 File uploaded: {filename}")
+            if torch.cuda.is_available():
+                device = torch.device('cuda')
+                log_message(f"🚀 Processing with GPU: {torch.cuda.get_device_name()}")
+                process_frames_gpu(cap, out, device, target_h, target_w, frame_count)
+            else:
+                log_message("💻 Processing with CPU (may be slower)")
+                process_frames_cpu(cap, out, target_h, target_w, frame_count)
             
-            return jsonify({
-                "success": True, 
-                "file_id": file_id,
-                "filename": filename,
-                "output_filename": output_filename,
-                "media_type": media_type,
-                "method": "Real-ESRGAN" if REALESRGAN_AVAILABLE and media_type == "image" else "Traditional",
-                "message": "Upload successful, processing started"
+            cap.release()
+            out.release()
+            
+            # Verify the output file was created and has content
+            if os.path.exists(output_path):
+                file_size = os.path.getsize(output_path)
+                if file_size > 0:
+                    log_message(f"✅ 4K video completed: {target_w}x{target_h}")
+                    log_message(f"📁 Output file size: {file_size / (1024**2):.1f}MB")
+                else:
+                    log_message(f"❌ Output file is empty: {output_path}")
+                    raise Exception("Output video file is empty")
+            else:
+                log_message(f"❌ Output file not created: {output_path}")
+                raise Exception("Output video file was not created")
+            
+            # Add to processed files list
+            app_state["processed_files"].append({
+                "input_file": os.path.basename(input_path),
+                "output_file": os.path.basename(output_path),
+                "original_size": f"{w}x{h}",
+                "upscaled_size": f"{target_w}x{target_h}",
+                "frame_count": frame_count,
+                "fps": fps,
+                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
             })
-        else:
-            return jsonify({"success": False, "error": "File type not allowed"})
-    except Exception as e:
-        return jsonify({"success": False, "error": str(e)})
+            
+        except Exception as e:
+            log_message(f"❌ Error processing video: {str(e)}")
+        finally:
+            app_state["processing_active"] = False
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+    
+    thread = threading.Thread(target=process_worker)
+    thread.daemon = True
+    thread.start()
 
-# Add new endpoint for model selection
-@app.route('/api/select-model', methods=['POST'])
-def api_select_model():
-    """Select Real-ESRGAN model"""
-    try:
-        data = request.get_json()
-        model_name = data.get('model', 'RealESRGAN_x4plus')
+def process_frames_cpu(cap, out, target_h, target_w, frame_count):
+    """Process video frames using CPU"""
+    frame_num = 0
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
         
-        if model_name not in ['RealESRGAN_x4plus', 'RealESRGAN_x2plus', 'RealESRNet_x4plus']:
-            return jsonify({"success": False, "error": "Invalid model name"})
+        frame_num += 1
         
-        scale = 4 if '4' in model_name else 2
-        upscaler = initialize_realesrgan(model_name, scale)
+        # Simple CPU upscaling
+        upscaled_frame = cv2.resize(frame, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
+        out.write(upscaled_frame)
         
-        if upscaler:
-            return jsonify({
-                "success": True, 
-                "message": f"Model {model_name} selected",
-                "current_model": app_state["current_model"]
-            })
-        else:
-            return jsonify({"success": False, "error": "Failed to initialize model"})
+        # Progress logging
+        if frame_num % 30 == 0:
+            progress = (frame_num / frame_count) * 100
+            log_message(f"🎞️ Processing frame {frame_num}/{frame_count} ({progress:.1f}%)")
+
+def process_frames_gpu(cap, out, device, target_h, target_w, frame_count):
+    """Process video frames using GPU with PyTorch"""
+    frame_num = 0
+    torch.backends.cudnn.benchmark = True
+    
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        frame_num += 1
+        
+        try:
+            # Convert to tensor
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame_tensor = torch.from_numpy(frame_rgb).float().to(device) / 255.0
+            frame_tensor = frame_tensor.permute(2, 0, 1).unsqueeze(0)
             
-    except Exception as e:
-        return jsonify({"success": False, "error": str(e)})
+            with torch.no_grad():
+                upscaled = torch.nn.functional.interpolate(
+                    frame_tensor,
+                    size=(target_h, target_w),
+                    mode='bicubic',
+                    align_corners=False
+                )
+            
+            # Convert back
+            result_cpu = upscaled.squeeze(0).permute(1, 2, 0).cpu().numpy()
+            result_frame = (result_cpu * 255).astype(np.uint8)
+            result_bgr = cv2.cvtColor(result_frame, cv2.COLOR_RGB2BGR)
+            out.write(result_bgr)
+            
+        except Exception as e:
+            log_message(f"⚠️ GPU processing failed for frame {frame_num}, using CPU fallback")
+            # CPU fallback
+            upscaled_frame = cv2.resize(frame, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
+            out.write(upscaled_frame)
+        
+        # Progress logging
+        if frame_num % 30 == 0:
+            progress = (frame_num / frame_count) * 100
+            log_message(f"🎞️ Processing frame {frame_num}/{frame_count} ({progress:.1f}%)")
+            
+            # Periodic memory cleanup
+            if frame_num % 60 == 0 and torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+# Initialize directories
+ensure_directories()
+
+app = Flask(__name__)
+
+@app.route('/')
+def index():
+    return render_template('index.html')
 
-# Update system info endpoint
 @app.route('/api/system')
 def api_system():
     """Get system information"""
@@ -358,9 +473,16 @@ def api_system():
             info["gpu_memory"] = f"{total_memory / (1024**3):.1f}GB"
             info["gpu_memory_used"] = f"{allocated_memory / (1024**3):.1f}GB"
             info["gpu_memory_free"] = f"{(total_memory - allocated_memory) / (1024**3):.1f}GB"
+            info["cuda_version"] = torch.version.cuda
+            info["pytorch_version"] = torch.__version__
         else:
             info["gpu_available"] = False
             info["gpu_name"] = "CPU Only"
+            info["gpu_memory"] = "N/A"
+            info["gpu_memory_used"] = "N/A"
+            info["gpu_memory_free"] = "N/A"
+            info["cuda_version"] = "Not available"
+            info["pytorch_version"] = torch.__version__
         
         # Real-ESRGAN info
         info["realesrgan_available"] = REALESRGAN_AVAILABLE
@@ -384,18 +506,83 @@ def api_system():
         
         info["available_models"] = available_models
         
+        # Storage info
+        if os.path.exists("/data"):
+            info["persistent_storage"] = True
+            try:
+                upload_files = os.listdir(UPLOAD_FOLDER) if os.path.exists(UPLOAD_FOLDER) else []
+                output_files = os.listdir(OUTPUT_FOLDER) if os.path.exists(OUTPUT_FOLDER) else []
+                
+                upload_size = sum(os.path.getsize(os.path.join(UPLOAD_FOLDER, f)) 
+                                for f in upload_files if os.path.isfile(os.path.join(UPLOAD_FOLDER, f)))
+                output_size = sum(os.path.getsize(os.path.join(OUTPUT_FOLDER, f)) 
+                                for f in output_files if os.path.isfile(os.path.join(OUTPUT_FOLDER, f)))
+                
+                info["storage_uploads"] = f"{upload_size / (1024**2):.1f}MB"
+                info["storage_outputs"] = f"{output_size / (1024**2):.1f}MB"
+                info["upload_files_count"] = len(upload_files)
+                info["output_files_count"] = len(output_files)
+            except Exception as e:
+                info["storage_uploads"] = f"Error: {str(e)}"
+                info["storage_outputs"] = "N/A"
+                info["upload_files_count"] = 0
+                info["output_files_count"] = 0
+        else:
+            info["persistent_storage"] = False
+        
         return jsonify({"success": True, "data": info})
     except Exception as e:
         return jsonify({"success": False, "error": str(e)})
 
-# Initialize the app
-if __name__ == '__main__':
-    ensure_directories()
-    log_message("🚀 4K Upscaler starting...")
-    
-    # Initialize Real-ESRGAN if available
-    if REALESRGAN_AVAILABLE:
-        log_message("🔧 Initializing Real-ESRGAN...")
-        initialize_realesrgan()
-    
-    app.run(host='0.0.0.0', port=7860, debug=False, threaded=True)
+@app.route('/api/upload', methods=['POST'])
+def api_upload():
+    """Upload and process file for 4K upscaling"""
+    try:
+        if 'file' not in request.files:
+            return jsonify({"success": False, "error": "No file provided"})
+        
+        file = request.files['file']
+        if file.filename == '':
+            return jsonify({"success": False, "error": "No file selected"})
+        
+        if file and allowed_file(file.filename):
+            file_id = str(uuid.uuid4())
+            filename = secure_filename(file.filename)
+            file_ext = filename.rsplit('.', 1)[1].lower()
+            
+            input_filename = f"{file_id}_input.{file_ext}"
+            input_path = os.path.join(UPLOAD_FOLDER, input_filename)
+            file.save(input_path)
+            
+            output_filename = f"{file_id}_4k.{file_ext}"
+            output_path = os.path.join(OUTPUT_FOLDER, output_filename)
+            
+            if file_ext in ['png', 'jpg', 'jpeg', 'gif']:
+                # Use Real-ESRGAN for images if available
+                if REALESRGAN_AVAILABLE:
+                    upscale_image_realesrgan(input_path, output_path)
+                else:
+                    upscale_image_4k_fallback(input_path, output_path)
+                media_type = "image"
+            elif file_ext in ['mp4', 'avi', 'mov', 'mkv']:
+                upscale_video_4k(input_path, output_path)
+                media_type = "video"
+            
+            log_message(f"📤 File uploaded: {filename}")
+            
+            return jsonify({
+                "success": True, 
+                "file_id": file_id,
+                "filename": filename,
+                "output_filename": output_filename,
+                "media_type": media_type,
+                "method": "Real-ESRGAN" if REALESRGAN_AVAILABLE and media_type == "image" else "Traditional",
+                "message": "Upload successful, processing started"
+            })
+        else:
+            return jsonify({"success": False, "error": "File type not allowed"})
+    except Exception as e:
+        return jsonify({"success": False, "error": str(e)})
+
+@app.route('/api/select-model', methods=['POST'])
+def api_select_model