Update app.py

app.py

@@ -11,7 +11,7 @@ import mimetypes
 import numpy as np
 from PIL import Image
 
-# Real-ESRGAN imports with error handling
+# Real-ESRGAN imports with better error handling
 try:
     from realesrgan import RealESRGANer
     from basicsr.archs.rrdbnet_arch import RRDBNet
@@ -20,7 +20,6 @@ try:
 except ImportError as e:
     REALESRGAN_AVAILABLE = False
     print(f"⚠️ Real-ESRGAN not available: {e}")
-    print("📝 Using fallback bicubic upscaling methods")
 
 # Configuration
 UPLOAD_FOLDER = '/data/uploads'
@@ -90,8 +89,12 @@ def download_realesrgan_models():
             model_path = os.path.join(MODEL_FOLDER, f"{model_name}.pth")
             if not os.path.exists(model_path):
                 log_message(f"📥 Downloading {model_name}...")
-                urllib.request.urlretrieve(url, model_path)
-                log_message(f"✅ Downloaded {model_name}")
+                try:
+                    urllib.request.urlretrieve(url, model_path)
+                    log_message(f"✅ Downloaded {model_name}")
+                except Exception as e:
+                    log_message(f"❌ Failed to download {model_name}: {e}")
+                    return False
         return True
     except Exception as e:
         log_message(f"❌ Error downloading models: {str(e)}")
@@ -104,12 +107,15 @@ def initialize_realesrgan(model_name='RealESRGAN_x4plus', scale=4):
         return None
     
     try:
+        log_message(f"🔧 Initializing Real-ESRGAN with {model_name}...")
+        
         model_path = os.path.join(MODEL_FOLDER, f"{model_name}.pth")
         
         # Check if model exists, download if not
         if not os.path.exists(model_path):
             log_message(f"📥 Model {model_name} not found, downloading...")
             if not download_realesrgan_models():
+                log_message("❌ Failed to download models")
                 return None
         
         # Initialize model architecture
@@ -123,18 +129,19 @@ def initialize_realesrgan(model_name='RealESRGAN_x4plus', scale=4):
             log_message(f"❌ Unknown model: {model_name}")
             return None
         
-        # Determine device
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        # Always use CPU for compatibility
+        device = torch.device('cpu')
+        log_message(f"🖥️ Using device: {device}")
         
         # Initialize upscaler with conservative settings
         upscaler = RealESRGANer(
             scale=netscale,
             model_path=model_path,
             model=model,
-            tile=400 if device.type == 'cuda' else 200,  # Conservative tile sizes
+            tile=200,  # Small tile size for CPU
             tile_pad=10,
             pre_pad=0,
-            half=False,  # Disable half precision for stability
+            half=False,  # No half precision on CPU
             device=device
         )
         
@@ -145,6 +152,8 @@ def initialize_realesrgan(model_name='RealESRGAN_x4plus', scale=4):
         
     except Exception as e:
         log_message(f"❌ Error initializing Real-ESRGAN: {str(e)}")
+        app_state["upscaler"] = None
+        app_state["current_model"] = None
         return None
 
 def optimize_gpu():
@@ -163,28 +172,20 @@ def optimize_gpu():
             log_message("✅ GPU optimized")
             return True
         else:
-            log_message("⚠️ CUDA not available")
+            log_message("⚠️ CUDA not available, using CPU")
             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 upscale_image_4k(input_path, output_path):
+    """Main upscaling function - tries Real-ESRGAN first, falls back to enhanced bicubic"""
     def process_worker():
         try:
-            log_message(f"🎨 Starting Real-ESRGAN upscaling: {os.path.basename(input_path)}")
+            log_message(f"🎨 Starting 4K upscaling: {os.path.basename(input_path)}")
             app_state["processing_active"] = True
             
-            # Initialize upscaler if not done
-            if app_state["upscaler"] is None:
-                upscaler = initialize_realesrgan()
-                if upscaler is None:
-                    log_message("❌ Real-ESRGAN initialization failed, using fallback")
-                    upscale_image_fallback(input_path, output_path)
-                    return
-            else:
-                upscaler = app_state["upscaler"]
+            start_time = time.time()
             
             # Read image
             img = cv2.imread(input_path, cv2.IMREAD_COLOR)
@@ -195,51 +196,146 @@ def upscale_image_realesrgan(input_path, output_path):
             h, w = img.shape[:2]
             log_message(f"📏 Original resolution: {w}x{h}")
             
-            # Check for very large images
-            if w * h > 4000 * 4000:
-                log_message("⚠️ Very large image detected, using tiled processing")
-                upscaler.tile = 200
-                upscaler.tile_pad = 5
+            success = False
+            method_used = "Unknown"
             
-            # Perform upscaling
-            log_message("🧠 Applying Real-ESRGAN neural upscaling...")
-            start_time = time.time()
+            # Try Real-ESRGAN first if available
+            if REALESRGAN_AVAILABLE:
+                try:
+                    if app_state["upscaler"] is None:
+                        log_message("🔧 Initializing Real-ESRGAN...")
+                        upscaler = initialize_realesrgan()
+                    else:
+                        upscaler = app_state["upscaler"]
+                    
+                    if upscaler is not None:
+                        log_message("🧠 Applying Real-ESRGAN neural upscaling...")
+                        output, _ = upscaler.enhance(img, outscale=4)
+                        cv2.imwrite(output_path, output)
+                        method_used = f"Real-ESRGAN ({app_state['current_model']})"
+                        success = True
+                        log_message("✅ Real-ESRGAN upscaling successful")
+                    else:
+                        log_message("⚠️ Real-ESRGAN initialization failed")
+                        
+                except Exception as e:
+                    log_message(f"⚠️ Real-ESRGAN failed: {str(e)}")
+                    log_message("🔄 Falling back to enhanced bicubic...")
             
-            try:
-                output, _ = upscaler.enhance(img, outscale=4)
-                
-                # Save result
-                cv2.imwrite(output_path, output)
+            # Fallback to enhanced bicubic if Real-ESRGAN failed or not available
+            if not success:
+                log_message("🔄 Using enhanced bicubic upscaling...")
                 
-                processing_time = time.time() - start_time
-                final_h, final_w = output.shape[:2]
+                target_h, target_w = h * 4, w * 4
                 
-                log_message(f"✅ Real-ESRGAN upscaling completed: {final_w}x{final_h}")
-                log_message(f"📈 Scale factor: {final_w/w:.1f}x")
-                log_message(f"⏱️ Processing time: {processing_time:.1f}s")
-                
-                # 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"{final_w}x{final_h}",
-                    "method": f"Real-ESRGAN ({app_state['current_model']})",
-                    "processing_time": f"{processing_time:.1f}s",
-                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-                })
+                if torch.cuda.is_available():
+                    try:
+                        # GPU enhanced bicubic
+                        device = torch.device('cuda')
+                        image_rgb = cv2.cvtColor(img, 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)
+                        
+                        with torch.no_grad():
+                            # Progressive upscaling for better quality
+                            intermediate = torch.nn.functional.interpolate(
+                                image_tensor, 
+                                size=(h * 2, w * 2), 
+                                mode='bicubic', 
+                                align_corners=False,
+                                antialias=True
+                            )
+                            
+                            upscaled = torch.nn.functional.interpolate(
+                                intermediate, 
+                                size=(target_h, target_w), 
+                                mode='bicubic', 
+                                align_corners=False,
+                                antialias=True
+                            )
+                            
+                            # Enhanced sharpening
+                            kernel = torch.tensor([
+                                [-0.5, -1, -0.5],
+                                [-1, 7, -1],
+                                [-0.5, -1, -0.5]
+                            ], dtype=torch.float32, device=device).unsqueeze(0).unsqueeze(0)
+                            
+                            enhanced_channels = []
+                            for i in range(3):
+                                channel = upscaled[:, i:i+1, :, :]
+                                padded = torch.nn.functional.pad(channel, (1, 1, 1, 1), mode='reflect')
+                                enhanced = torch.nn.functional.conv2d(padded, kernel)
+                                enhanced_channels.append(enhanced)
+                            
+                            enhanced = torch.cat(enhanced_channels, dim=1)
+                            final_result = torch.clamp(enhanced, 0, 1)
+                        
+                        result_cpu = final_result.squeeze(0).permute(1, 2, 0).cpu().numpy()
+                        result_image = (result_cpu * 255).astype(np.uint8)
+                        result_bgr = cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR)
+                        
+                        cv2.imwrite(output_path, result_bgr)
+                        method_used = "Enhanced Bicubic (GPU)"
+                        success = True
+                        log_message("✅ GPU enhanced bicubic completed")
+                        
+                    except Exception as e:
+                        log_message(f"⚠️ GPU processing failed: {e}")
                 
-            except RuntimeError as e:
-                if "out of memory" in str(e).lower():
-                    log_message("⚠️ GPU memory insufficient, using fallback")
-                    upscale_image_fallback(input_path, output_path)
-                else:
-                    raise e
+                if not success:
+                    # CPU enhanced bicubic
+                    log_message("💻 Using CPU enhanced bicubic...")
+                    
+                    # Progressive upscaling
+                    intermediate = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_CUBIC)
+                    upscaled = cv2.resize(intermediate, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
+                    
+                    # Apply sharpening
+                    kernel = np.array([
+                        [-0.5, -1, -0.5],
+                        [-1, 7, -1],
+                        [-0.5, -1, -0.5]
+                    ])
+                    sharpened = cv2.filter2D(upscaled, -1, kernel)
+                    
+                    # Blend for final result
+                    result = cv2.addWeighted(upscaled, 0.7, sharpened, 0.3, 0)
+                    
+                    cv2.imwrite(output_path, result)
+                    method_used = "Enhanced Bicubic (CPU)"
+                    success = True
+                    log_message("✅ CPU enhanced bicubic completed")
+            
+            if success:
+                # Verify output
+                final_img = cv2.imread(output_path)
+                if final_img is not None:
+                    final_h, final_w = final_img.shape[:2]
+                    processing_time = time.time() - start_time
                     
+                    log_message(f"✅ Upscaling completed: {final_w}x{final_h}")
+                    log_message(f"📈 Scale factor: {final_w/w:.1f}x")
+                    log_message(f"⏱️ Processing time: {processing_time:.1f}s")
+                    log_message(f"🔧 Method used: {method_used}")
+                    
+                    # Add to processed files
+                    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"{final_w}x{final_h}",
+                        "method": method_used,
+                        "processing_time": f"{processing_time:.1f}s",
+                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                    })
+                else:
+                    log_message("❌ Error: Output file could not be read")
+            else:
+                log_message("❌ All upscaling methods failed")
+                
         except Exception as e:
-            log_message(f"❌ Error in Real-ESRGAN processing: {str(e)}")
-            log_message("🔄 Falling back to traditional upscaling")
-            upscale_image_fallback(input_path, output_path)
+            log_message(f"❌ Critical error in upscaling: {str(e)}")
         finally:
             app_state["processing_active"] = False
             if torch.cuda.is_available():
@@ -249,70 +345,6 @@ def upscale_image_realesrgan(input_path, output_path):
     thread.daemon = True
     thread.start()
 
-def upscale_image_fallback(input_path, output_path):
-    """Fallback upscaling method using bicubic interpolation"""
-    try:
-        log_message("🔄 Using fallback bicubic upscaling")
-        
-        # Read original image
-        image = cv2.imread(input_path)
-        if image is None:
-            log_message("❌ Error: Could not read image")
-            return
-            
-        h, w = image.shape[:2]
-        target_h, target_w = h * 4, w * 4
-        
-        if torch.cuda.is_available():
-            try:
-                # GPU implementation
-                device = torch.device('cuda')
-                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)
-                
-                with torch.no_grad():
-                    upscaled = torch.nn.functional.interpolate(
-                        image_tensor, 
-                        size=(target_h, target_w), 
-                        mode='bicubic', 
-                        align_corners=False,
-                        antialias=True
-                    )
-                    upscaled = torch.clamp(upscaled, 0, 1)
-                
-                result_cpu = upscaled.squeeze(0).permute(1, 2, 0).cpu().numpy()
-                result_image = (result_cpu * 255).astype(np.uint8)
-                result_bgr = cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR)
-                
-                cv2.imwrite(output_path, result_bgr)
-                log_message(f"✅ GPU fallback completed: {target_w}x{target_h}")
-                
-            except Exception as e:
-                log_message(f"⚠️ GPU fallback failed: {e}, using CPU")
-                # CPU fallback
-                upscaled = cv2.resize(image, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
-                cv2.imwrite(output_path, upscaled)
-                log_message(f"✅ CPU fallback completed: {target_w}x{target_h}")
-        else:
-            # CPU fallback
-            upscaled = cv2.resize(image, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
-            cv2.imwrite(output_path, upscaled)
-            log_message(f"✅ CPU fallback completed: {target_w}x{target_h}")
-        
-        # Add to processed files
-        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}",
-            "method": "Fallback (Bicubic)",
-            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        })
-        
-    except Exception as e:
-        log_message(f"❌ Error in fallback processing: {str(e)}")
-
 def upscale_video_4k(input_path, output_path):
     """Upscale video to 4K frame by frame"""
     def process_worker():
@@ -340,6 +372,8 @@ def upscale_video_4k(input_path, output_path):
             
             # Process frames
             frame_num = 0
+            start_time = time.time()
+            
             while True:
                 ret, frame = cap.read()
                 if not ret:
@@ -348,14 +382,22 @@ def upscale_video_4k(input_path, output_path):
                 frame_num += 1
                 
                 try:
-                    # Simple bicubic upscaling for video (faster)
+                    # Enhanced bicubic for video frames
                     upscaled_frame = cv2.resize(frame, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
-                    out.write(upscaled_frame)
+                    
+                    # Light sharpening
+                    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
+                    sharpened = cv2.filter2D(upscaled_frame, -1, kernel)
+                    final_frame = cv2.addWeighted(upscaled_frame, 0.8, sharpened, 0.2, 0)
+                    
+                    out.write(final_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}%)")
+                        elapsed = time.time() - start_time
+                        eta = (elapsed / frame_num) * (frame_count - frame_num)
+                        log_message(f"🎞️ Frame {frame_num}/{frame_count} ({progress:.1f}%) - ETA: {eta:.0f}s")
                         
                 except Exception as e:
                     log_message(f"⚠️ Error processing frame {frame_num}: {e}")
@@ -367,8 +409,10 @@ def upscale_video_4k(input_path, output_path):
             # Verify output
             if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
                 file_size = os.path.getsize(output_path)
+                total_time = time.time() - start_time
                 log_message(f"✅ Video upscaling completed: {target_w}x{target_h}")
-                log_message(f"📁 Output file size: {file_size / (1024**2):.1f}MB")
+                log_message(f"📁 Output size: {file_size / (1024**2):.1f}MB")
+                log_message(f"⏱️ Total time: {total_time:.1f}s")
                 
                 # Add to processed files
                 app_state["processed_files"].append({
@@ -378,7 +422,8 @@ def upscale_video_4k(input_path, output_path):
                     "upscaled_size": f"{target_w}x{target_h}",
                     "frame_count": frame_count,
                     "fps": fps,
-                    "method": "Bicubic (Video)",
+                    "method": "Enhanced Bicubic (Video)",
+                    "processing_time": f"{total_time:.1f}s",
                     "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                 })
             else:
@@ -395,194 +440,22 @@ def upscale_video_4k(input_path, output_path):
     thread.daemon = True
     thread.start()
 
-# Initialize directories
+# Initialize directories and try to set up Real-ESRGAN
 ensure_directories()
 
+# Try to initialize Real-ESRGAN on startup
+if REALESRGAN_AVAILABLE:
+    try:
+        log_message("🚀 Attempting to initialize Real-ESRGAN on startup...")
+        initialize_realesrgan()
+    except Exception as e:
+        log_message(f"⚠️ Could not initialize Real-ESRGAN on startup: {e}")
+
 app = Flask(__name__)
 
 @app.route('/')
 def index():
-    return """
-    <!DOCTYPE html>
-    <html>
-    <head>
-        <title>4K AI Upscaler</title>
-        <meta charset="utf-8">
-        <meta name="viewport" content="width=device-width, initial-scale=1">
-        <style>
-            body { font-family: Arial, sans-serif; margin: 40px; background-color: #f5f5f5; }
-            .container { max-width: 800px; margin: 0 auto; background: white; padding: 30px; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); }
-            h1 { color: #333; text-align: center; }
-            .upload-area { border: 2px dashed #ccc; padding: 40px; text-align: center; margin: 20px 0; border-radius: 10px; }
-            .upload-area.dragover { border-color: #007bff; background-color: #f0f8ff; }
-            .btn { background-color: #007bff; color: white; border: none; padding: 10px 20px; border-radius: 5px; cursor: pointer; }
-            .btn:hover { background-color: #0056b3; }
-            .progress { width: 100%; height: 20px; background-color: #f0f0f0; border-radius: 10px; overflow: hidden; margin: 20px 0; }
-            .progress-bar { height: 100%; background-color: #007bff; width: 0%; transition: width 0.3s; }
-            .logs { background-color: #f8f9fa; border: 1px solid #dee2e6; border-radius: 5px; padding: 15px; height: 200px; overflow-y: auto; font-family: monospace; font-size: 12px; }
-            .system-info { background-color: #e9ecef; padding: 15px; border-radius: 5px; margin: 20px 0; }
-            .hidden { display: none; }
-        </style>
-    </head>
-    <body>
-        <div class="container">
-            <h1>🚀 4K AI Upscaler</h1>
-            
-            <div class="system-info" id="system-info">
-                Loading system information...
-            </div>
-            
-            <div class="upload-area" id="upload-area">
-                <p>📤 Drop your image or video here, or click to select</p>
-                <input type="file" id="file-input" accept=".png,.jpg,.jpeg,.gif,.mp4,.avi,.mov,.mkv" style="display: none;">
-                <button class="btn" onclick="document.getElementById('file-input').click()">Select File</button>
-            </div>
-            
-            <div class="progress hidden" id="progress">
-                <div class="progress-bar" id="progress-bar"></div>
-            </div>
-            
-            <div id="result" class="hidden">
-                <h3>✅ Processing Complete!</h3>
-                <button class="btn" id="download-btn">Download Result</button>
-            </div>
-            
-            <h3>📊 System Logs</h3>
-            <div class="logs" id="logs"></div>
-            
-            <button class="btn" onclick="clearLogs()" style="margin-top: 10px;">Clear Logs</button>
-        </div>
-        
-        <script>
-            let currentFileId = null;
-            
-            // Load system info
-            function loadSystemInfo() {
-                fetch('/api/system')
-                    .then(response => response.json())
-                    .then(data => {
-                        if (data.success) {
-                            const info = data.data;
-                            document.getElementById('system-info').innerHTML = `
-                                <strong>🖥️ System Status:</strong><br>
-                                GPU: ${info.gpu_name}<br>
-                                PyTorch: ${info.pytorch_version}<br>
-                                Real-ESRGAN: ${info.realesrgan_available ? '✅ Available' : '❌ Not Available'}<br>
-                                Current Model: ${info.current_model || 'None'}
-                            `;
-                        }
-                    })
-                    .catch(error => {
-                        document.getElementById('system-info').innerHTML = '❌ Error loading system info';
-                    });
-            }
-            
-            // Load logs
-            function loadLogs() {
-                fetch('/api/logs')
-                    .then(response => response.json())
-                    .then(data => {
-                        if (data.success) {
-                            const logsDiv = document.getElementById('logs');
-                            logsDiv.innerHTML = data.logs.join('<br>');
-                            logsDiv.scrollTop = logsDiv.scrollHeight;
-                        }
-                    });
-            }
-            
-            // Clear logs
-            function clearLogs() {
-                fetch('/api/clear-logs', { method: 'POST' })
-                    .then(() => loadLogs());
-            }
-            
-            // File upload handling
-            const uploadArea = document.getElementById('upload-area');
-            const fileInput = document.getElementById('file-input');
-            
-            uploadArea.addEventListener('dragover', (e) => {
-                e.preventDefault();
-                uploadArea.classList.add('dragover');
-            });
-            
-            uploadArea.addEventListener('dragleave', () => {
-                uploadArea.classList.remove('dragover');
-            });
-            
-            uploadArea.addEventListener('drop', (e) => {
-                e.preventDefault();
-                uploadArea.classList.remove('dragover');
-                const files = e.dataTransfer.files;
-                if (files.length > 0) {
-                    uploadFile(files[0]);
-                }
-            });
-            
-            fileInput.addEventListener('change', (e) => {
-                if (e.target.files.length > 0) {
-                    uploadFile(e.target.files[0]);
-                }
-            });
-            
-            function uploadFile(file) {
-                const formData = new FormData();
-                formData.append('file', file);
-                
-                document.getElementById('progress').classList.remove('hidden');
-                document.getElementById('result').classList.add('hidden');
-                
-                fetch('/api/upload', {
-                    method: 'POST',
-                    body: formData
-                })
-                .then(response => response.json())
-                .then(data => {
-                    if (data.success) {
-                        currentFileId = data.file_id;
-                        checkProcessingStatus();
-                    } else {
-                        alert('Upload failed: ' + data.error);
-                    }
-                })
-                .catch(error => {
-                    alert('Upload error: ' + error);
-                });
-            }
-            
-            function checkProcessingStatus() {
-                fetch('/api/processing-status')
-                    .then(response => response.json())
-                    .then(data => {
-                        if (data.success) {
-                            if (data.processing) {
-                                // Still processing
-                                setTimeout(checkProcessingStatus, 2000);
-                            } else {
-                                // Processing complete
-                                document.getElementById('progress').classList.add('hidden');
-                                document.getElementById('result').classList.remove('hidden');
-                                
-                                // Find the latest processed file
-                                if (data.processed_files.length > 0) {
-                                    const latest = data.processed_files[data.processed_files.length - 1];
-                                    document.getElementById('download-btn').onclick = () => {
-                                        window.open(`/api/download/${latest.output_file}`, '_blank');
-                                    };
-                                }
-                            }
-                        }
-                        loadLogs(); // Update logs
-                    });
-            }
-            
-            // Initialize
-            loadSystemInfo();
-            loadLogs();
-            setInterval(loadLogs, 5000); // Update logs every 5 seconds
-        </script>
-    </body>
-    </html>
-    """
+    return render_template('index.html')
 
 @app.route('/api/system')
 def api_system():
@@ -600,7 +473,6 @@ def api_system():
             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"
@@ -608,35 +480,32 @@ def api_system():
             info["gpu_memory_used"] = "N/A"
             info["gpu_memory_free"] = "N/A"
             info["cuda_version"] = "Not available"
-            info["pytorch_version"] = torch.__version__
+        
+        info["pytorch_version"] = torch.__version__
         
         # Real-ESRGAN info
-        info["realesrgan_available"] = REALESRGAN_AVAILABLE
+        info["realesrgan_available"] = REALESRGAN_AVAILABLE and app_state["upscaler"] is not None
         info["current_model"] = app_state.get("current_model", "None")
         
         # 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
+        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"] = "Error"
+            info["storage_outputs"] = "Error"
+            info["upload_files_count"] = 0
+            info["output_files_count"] = 0
         
         return jsonify({"success": True, "data": info})
     except Exception as e:
@@ -666,21 +535,14 @@ def api_upload():
             output_path = os.path.join(OUTPUT_FOLDER, output_filename)
             
             if file_ext in ['png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff', 'webp']:
-                # Use Real-ESRGAN for images if available
-                if REALESRGAN_AVAILABLE:
-                    upscale_image_realesrgan(input_path, output_path)
-                    method = "Real-ESRGAN"
-                else:
-                    upscale_image_fallback(input_path, output_path)
-                    method = "Bicubic Fallback"
+                upscale_image_4k(input_path, output_path)
                 media_type = "image"
             elif file_ext in ['mp4', 'avi', 'mov', 'mkv']:
                 upscale_video_4k(input_path, output_path)
-                method = "Video Upscaling"
                 media_type = "video"
             
             log_message(f"📤 File uploaded: {filename}")
-            log_message(f"🎯 Starting upscaling with {method}...")
+            log_message(f"🎯 Starting 4K upscaling process...")
             
             return jsonify({
                 "success": True, 
@@ -688,7 +550,6 @@ def api_upload():
                 "filename": filename,
                 "output_filename": output_filename,
                 "media_type": media_type,
-                "method": method,
                 "message": "Upload successful, processing started"
             })
         else:
@@ -756,16 +617,28 @@ def api_optimize_gpu():
     """Optimize GPU for processing"""
     try:
         success = optimize_gpu()
-        if success:
-            return jsonify({"success": True, "message": "GPU optimized"})
+        return jsonify({"success": success})
+    except Exception as e:
+        return jsonify({"success": False, "error": str(e)})
+
+@app.route('/api/init-realesrgan', methods=['POST'])
+def api_init_realesrgan():
+    """Initialize Real-ESRGAN manually"""
+    try:
+        if not REALESRGAN_AVAILABLE:
+            return jsonify({"success": False, "error": "Real-ESRGAN not available"})
+        
+        upscaler = initialize_realesrgan()
+        if upscaler:
+            return jsonify({"success": True, "message": "Real-ESRGAN initialized successfully"})
         else:
-            return jsonify({"success": False, "message": "GPU optimization failed"})
+            return jsonify({"success": False, "error": "Failed to initialize Real-ESRGAN"})
     except Exception as e:
         return jsonify({"success": False, "error": str(e)})
 
 @app.route('/api/clear-cache', methods=['POST'])
 def api_clear_cache():
-    """Clear GPU cache and processed files"""
+    """Clear cache and processed files"""
     try:
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -777,47 +650,18 @@ def api_clear_cache():
     except Exception as e:
         return jsonify({"success": False, "error": str(e)})
 
-@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_name', 'RealESRGAN_x4plus')
-        
-        if not REALESRGAN_AVAILABLE:
-            return jsonify({"success": False, "error": "Real-ESRGAN not available"})
-        
-        # Initialize new model
-        upscaler = initialize_realesrgan(model_name)
-        if upscaler:
-            return jsonify({"success": True, "message": f"Model {model_name} selected"})
-        else:
-            return jsonify({"success": False, "error": "Failed to initialize model"})
-    except Exception as e:
-        return jsonify({"success": False, "error": str(e)})
-
 if __name__ == '__main__':
     # Initialize system
-    log_message("🚀 4K AI Upscaler starting...")
+    log_message("🚀 4K Upscaler starting...")
     
     try:
         # Optimize GPU if available
         if optimize_gpu():
-            log_message("✅ GPU optimized")
-        else:
-            log_message("⚠️ GPU optimization failed, using CPU fallback")
-        
-        # Initialize Real-ESRGAN if available
-        if REALESRGAN_AVAILABLE:
-            log_message("🧠 Initializing Real-ESRGAN...")
-            if initialize_realesrgan():
-                log_message("✅ Real-ESRGAN ready")
-            else:
-                log_message("⚠️ Real-ESRGAN initialization failed, using fallback methods")
+            log_message("✅ GPU optimization completed")
         else:
-            log_message("📝 Real-ESRGAN not available, using bicubic upscaling")
+            log_message("⚠️ Using CPU mode")
         
-        log_message("✅ 4K AI Upscaler ready")
+        log_message("✅ 4K Upscaler ready")
         log_message("📤 Upload images or videos to upscale to 4K resolution")
         
     except Exception as e: