Update app.py

app.py

@@ -7,11 +7,9 @@ from PIL import Image
 import os
 from huggingface_hub import hf_hub_download
 import cv2
-from pathlib import Path
 import sys
 import warnings
-
-from models.rrdbnet import RRDBNet, process_with_tiling
+import gc
 
 warnings.filterwarnings('ignore', category=FutureWarning)
 warnings.filterwarnings('ignore', category=UserWarning)
@@ -19,6 +17,156 @@ os.environ['PYTHONWARNINGS'] = 'ignore'
 
 sys.path.append(os.path.join(os.path.dirname(__file__), 'models'))
 
+class ResidualDenseBlock(nn.Module):
+    def __init__(self, nf=64, gc=32):
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=True)
+        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=True)
+        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=True)
+        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=True)
+        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=True)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5 * 0.2 + x
+
+class RRDB(nn.Module):
+    def __init__(self, nf, gc=32):
+        super(RRDB, self).__init__()
+        self.rdb1 = ResidualDenseBlock(nf, gc)
+        self.rdb2 = ResidualDenseBlock(nf, gc)
+        self.rdb3 = ResidualDenseBlock(nf, gc)
+        
+    def forward(self, x):
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        return out * 0.2 + x
+
+class RRDBNet(nn.Module):
+    def __init__(self, in_nc=3, out_nc=3, nf=64, nb=23, gc=32, scale=4):
+        super(RRDBNet, self).__init__()
+        self.scale = scale
+        
+        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
+        self.body = nn.ModuleList([RRDB(nf, gc) for _ in range(nb)])
+        self.conv_body = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        
+        self.conv_up1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.conv_up2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        if scale >= 8:
+            self.conv_up3 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        
+        self.conv_hr = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        
+    def forward(self, x):
+        fea = self.conv_first(x)
+        trunk = fea
+        for block in self.body:
+            trunk = block(trunk)
+        trunk = self.conv_body(trunk)
+        fea = fea + trunk
+        del trunk
+        
+        fea = self.lrelu(self.conv_up1(F.interpolate(fea, scale_factor=2, mode='bilinear', align_corners=False)))
+        fea = self.lrelu(self.conv_up2(F.interpolate(fea, scale_factor=2, mode='bilinear', align_corners=False)))
+        if self.scale >= 8:
+            fea = self.lrelu(self.conv_up3(F.interpolate(fea, scale_factor=2, mode='bilinear', align_corners=False)))
+        
+        out = self.conv_last(self.lrelu(self.conv_hr(fea)))
+        del fea
+        return out
+
+def hdr_like(img):
+    mean_val = img.mean(dim=(2, 3), keepdim=True)
+    img = img - mean_val
+    img = img * 1.1
+    img = img + 0.5
+    img = torch.clamp(img, 0, 1)
+    img = img ** 0.85
+    return torch.clamp(img, 0, 1)
+
+def sharpen(img, amount=0.15):
+    blur = F.avg_pool2d(img, kernel_size=3, stride=1, padding=1)
+    sharpened = img + amount * (img - blur)
+    return torch.clamp(sharpened, 0, 1)
+
+def process_with_tiling(model, img_tensor, tile_size=160, tile_overlap=32):
+    device = img_tensor.device
+    b, c, h, w = img_tensor.shape
+    scale = model.scale
+    
+    if device.type == 'cpu':
+        tile_size = min(tile_size, 128)
+        tile_overlap = 16
+    
+    if h <= tile_size and w <= tile_size:
+        with torch.no_grad():
+            output = model(img_tensor)
+        output = hdr_like(output)
+        output = sharpen(output, amount=0.15)
+        return torch.clamp(output, 0, 1)
+    
+    sample_tile = img_tensor[:, :, :min(tile_size, h), :min(tile_size, w)]
+    with torch.no_grad():
+        sample_output = model(sample_tile)
+    
+    output_channels = sample_output.shape[1]
+    sample_scale_h = sample_output.shape[2] / sample_tile.shape[2]
+    sample_scale_w = sample_output.shape[3] / sample_tile.shape[3]
+    del sample_tile, sample_output
+    
+    output_h = int(h * sample_scale_h)
+    output_w = int(w * sample_scale_w)
+    
+    output = torch.zeros((b, output_channels, output_h, output_w), device=device)
+    
+    stride = tile_size - tile_overlap
+    tiles_h = (h - 1) // stride + 1
+    tiles_w = (w - 1) // stride + 1
+    
+    print(f"🔲 Processing {tiles_h}x{tiles_w} = {tiles_h*tiles_w} tiles")
+    print(f"   Input: {c}ch {h}x{w} → Output: {output_channels}ch {output_h}x{output_w}")
+    
+    for i in range(0, h, stride):
+        for j in range(0, w, stride):
+            h_start = i
+            h_end = min(i + tile_size, h)
+            w_start = j
+            w_end = min(j + tile_size, w)
+            
+            tile = img_tensor[:, :, h_start:h_end, w_start:w_end]
+            
+            with torch.no_grad():
+                tile_output = model(tile)
+            
+            actual_h = tile_output.shape[2]
+            actual_w = tile_output.shape[3]
+            
+            output_h_start = int(h_start * sample_scale_h)
+            output_w_start = int(w_start * sample_scale_w)
+            
+            output[:, :, output_h_start:output_h_start+actual_h, output_w_start:output_w_start+actual_w] = tile_output
+            
+            del tile, tile_output
+            
+            if ((i // stride) * tiles_w + (j // stride)) % 4 == 0:
+                gc.collect()
+                if device.type == 'cuda':
+                    torch.cuda.empty_cache()
+    
+    output = hdr_like(output)
+    output = sharpen(output, amount=0.15)
+    
+    return torch.clamp(output, 0, 1)
+
 MODELS = {
     "Classical SR x8 (DIV2K)": {
         "repo": "deepinv/swinir",
@@ -76,7 +224,6 @@ model_cache = {}
 def setup_directories():
     os.makedirs("models", exist_ok=True)
     os.makedirs("temp", exist_ok=True)
-    print("✅ Directories created")
 
 def download_all_models():
     print("🚀 Starting model download...")
@@ -107,13 +254,7 @@ def download_all_models():
             print(f"❌ Failed to download {model_name}: {str(e)}")
             failed.append(model_name)
     
-    print(f"\n📊 Download Summary:")
-    print(f"   ✅ Success: {downloaded}/{len(MODELS)}")
-    if failed:
-        print(f"   ❌ Failed: {len(failed)}")
-        for name in failed:
-            print(f"      - {name}")
-    
+    print(f"\n📊 Download Summary: ✅ {downloaded}/{len(MODELS)}")
     return downloaded, failed
 
 def load_realesrgan_model(model_path, device, scale=4):
@@ -126,126 +267,71 @@ def load_realesrgan_model(model_path, device, scale=4):
         else:
             state_dict = checkpoint
         
-        # ✅ AUTO-DETECT architecture dari state_dict
-        # Detect input channels
         in_nc = 3
         if 'conv_first.weight' in state_dict:
             in_nc = state_dict['conv_first.weight'].shape[1]
         
-        # Detect output channels
         out_nc = 3
         if 'conv_last.weight' in state_dict:
             out_nc = state_dict['conv_last.weight'].shape[0]
         
-        # Detect nf (number of features)
-        nf = 64
-        if 'conv_first.weight' in state_dict:
-            nf = state_dict['conv_first.weight'].shape[0]
-        
-        # Detect nb (number of blocks) by counting body blocks
-        nb = 0
-        for key in state_dict.keys():
-            if key.startswith('body.') and 'rdb1.conv1.weight' in key:
-                block_idx = int(key.split('.')[1])
-                nb = max(nb, block_idx + 1)
-        
-        # Detect gc (growth channels)
-        gc = 32
-        if 'body.0.rdb1.conv1.weight' in state_dict:
-            gc = state_dict['body.0.rdb1.conv1.weight'].shape[0]
-        
-        print(f"🔍 Auto-detected architecture:")
-        print(f"   in_nc={in_nc}, out_nc={out_nc}, nf={nf}, nb={nb}, gc={gc}, scale={scale}")
-        
-        # ✅ Create model with DETECTED architecture (bukan hardcoded!)
-        model = RRDBNet(
-            in_nc=in_nc, 
-            out_nc=out_nc, 
-            nf=nf,
-            nb=nb,
-            gc=gc,
-            scale=scale
-        )
-        
-        # Load state dict
+        model = RRDBNet(in_nc=in_nc, out_nc=out_nc, nf=64, nb=23, gc=32, scale=scale)
         model.load_state_dict(state_dict, strict=True)
         model.eval()
-        
-        # ✅ CPU OPTIMIZATION
-        if device.type == 'cpu':
-            torch.set_num_threads(2)
-            if hasattr(torch.backends, 'mkldnn'):
-                torch.backends.mkldnn.enabled = True
-        
         model = model.to(device)
         
-        print(f"✅ Real-ESRGAN model loaded successfully")
-        print(f"   📊 Model size: ~{sum(p.numel() for p in model.parameters()) / 1e6:.1f}M parameters")
+        print(f"✅ Model loaded: {scale}x | In:{in_nc}ch | Out:{out_nc}ch")
         return model
     except Exception as e:
-        print(f"❌ Error loading Real-ESRGAN: {e}")
-        import traceback
-        traceback.print_exc()
+        print(f"❌ Error loading model: {e}")
         return None
 
-def process_with_realesrgan(image, model_path, device, scale=4, use_tiling=True):
+def process_with_realesrgan(image, model_path, device, scale=4):
     try:
         model = load_realesrgan_model(model_path, device, scale)
-        
         if model is None:
             return None
         
-        # Get input channels from model
         in_nc = model.conv_first.weight.shape[1]
         
         img = np.array(image).astype(np.float32) / 255.0
         if len(img.shape) == 2:
             img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
         
-        # Convert to tensor
         img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
         img = img.unsqueeze(0).to(device)
-        
-        # ✅ Clamp input
         img = torch.clamp(img, 0, 1)
         
-        print(f"📥 Input image shape: {img.shape}")
+        print(f"📥 Input: {img.shape}")
         
-        # Handle different input channel requirements
         if in_nc == 12:
             b, c, h, w = img.shape
-            
-            # Ensure dimensions are divisible by 2
             pad_h = (2 - h % 2) % 2
             pad_w = (2 - w % 2) % 2
             
             if pad_h > 0 or pad_w > 0:
                 img = F.pad(img, (0, pad_w, 0, pad_h), mode='replicate')
-                print(f"🔧 Padded image to: {img.shape}")
+                print(f"🔧 Padded: {img.shape}")
             
-            # Pixel unshuffle
             img = F.pixel_unshuffle(img, 2)
-            print(f"🔄 Applied pixel unshuffle: {img.shape}")
+            print(f"🔄 Pixel unshuffle: {img.shape}")
         
-        # ✅ ALWAYS use tiling untuk HF Spaces (hemat memory)
-        h, w = img.shape[2], img.shape[3]
-        print(f"🔲 Using optimized tiling mode ({h}x{w})")
         output = process_with_tiling(model, img, tile_size=160, tile_overlap=32)
         
-        print(f"📤 Output shape: {output.shape}")
+        print(f"📤 Output: {output.shape}")
         
         output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
         output = np.transpose(output, (1, 2, 0))
         output = (output * 255.0).round().astype(np.uint8)
         
-        # Clean up GPU memory
         del model, img
         if device.type == 'cuda':
             torch.cuda.empty_cache()
+        gc.collect()
         
         return Image.fromarray(output)
     except Exception as e:
-        print(f"❌ Error in Real-ESRGAN processing: {e}")
+        print(f"❌ Processing error: {e}")
         import traceback
         traceback.print_exc()
         return None
@@ -290,35 +376,18 @@ def process_image_simple(image, scale, task_type):
     return Image.fromarray(output)
 
 def upscale_image(image, model_name, output_format="png"):
-    """Process and upscale image with better error handling"""
-    
-    # Validation
     if image is None:
         return None, "❌ Please upload an image first!"
     
+    model_info = MODELS[model_name]
+    model_path = os.path.join("models", model_info["filename"])
+    
+    if not os.path.exists(model_path):
+        return None, f"❌ Model not found: {model_info['filename']}\nPlease restart the app."
+    
     try:
-        # Convert image to PIL if needed
-        if not isinstance(image, Image.Image):
-            try:
-                image = Image.fromarray(image)
-            except Exception as e:
-                return None, f"❌ Could not convert image: {str(e)}"
-        
-        model_info = MODELS[model_name]
-        model_path = os.path.join("models", model_info["filename"])
-        
-        if not os.path.exists(model_path):
-            return None, f"❌ Model not found: {model_info['filename']}\nPlease restart the app to download models."
-        
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         
-        print(f"📥 Processing image: {image.size} | Format: {image.format or 'Unknown'}")
-        
-        # Ensure image is in RGB mode
-        if image.mode != 'RGB':
-            print(f"🔄 Converting from {image.mode} to RGB")
-            image = image.convert('RGB')
-        
         if model_info["type"] == "realesrgan":
             print(f"🔥 Processing with Real-ESRGAN {model_info['scale']}x...")
             result_image = process_with_realesrgan(
@@ -329,50 +398,29 @@ def upscale_image(image, model_name, output_format="png"):
             )
             
             if result_image is None:
-                return None, "❌ Error processing with Real-ESRGAN model. Check console for details."
+                return None, "❌ Error processing with Real-ESRGAN"
         else:
-            print(f"📊 Processing with {model_info['type'].upper()}...")
-            # For SwinIR models, use simple upscaling for now
+            state_dict = load_model(model_path, device)
             result_image = process_image_simple(
                 image, 
                 model_info["scale"], 
                 model_info["task"]
             )
         
-        # Build info message
-        info = f"✅ Successfully Enhanced!\n"
-        info += f"{'='*50}\n"
-        info += f"🎯 Model: {model_name}\n"
-        info += f"📊 Type: {model_info['type'].upper()}\n"
-        info += f"🔍 Task: {model_info['task']}\n"
+        info = f"✅ Model: {model_name}\n"
+        info += f"🎯 Type: {model_info['type'].upper()}\n"
+        info += f"📊 Task: {model_info['task']}\n"
         info += f"🔢 Scale: {model_info['scale']}x\n"
-        info += f"{'='*50}\n"
-        info += f"📐 Input:  {image.size[0]} x {image.size[1]} px\n"
-        info += f"📐 Output: {result_image.size[0]} x {result_image.size[1]} px\n"
-        info += f"{'='*50}\n"
+        info += f"📐 Input: {image.size[0]}x{image.size[1]}\n"
+        info += f"📐 Output: {result_image.size[0]}x{result_image.size[1]}\n"
         info += f"💻 Device: {device}\n"
-        info += f"📁 Format: {output_format.upper()}\n"
-        info += f"📦 Model File: {model_info['filename']}\n"
+        info += f"📁 Format: {output_format.upper()}"
         
         return result_image, info
         
-    except FileNotFoundError as e:
-        error_msg = f"❌ File Error: Could not read uploaded image.\n"
-        error_msg += f"This might be a temporary Gradio issue.\n"
-        error_msg += f"Please try:\n"
-        error_msg += f"  1. Re-uploading the image\n"
-        error_msg += f"  2. Using a different image\n"
-        error_msg += f"  3. Refreshing the page\n\n"
-        error_msg += f"Technical: {str(e)}"
-        return None, error_msg
-        
     except Exception as e:
         import traceback
-        error_msg = f"❌ Error processing image\n"
-        error_msg += f"{'='*50}\n"
-        error_msg += f"Error: {str(e)}\n"
-        error_msg += f"{'='*50}\n"
-        error_msg += f"Traceback:\n{traceback.format_exc()}"
+        error_msg = f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
         return None, error_msg
 
 def get_model_status():
@@ -393,16 +441,12 @@ def get_model_status():
     return status
 
 print("="*60)
-print("🎨 SwinIR & Real-ESRGAN Image Upscaler")
+print("🎨 AI Image Upscaler - Optimized Edition")
 print("="*60)
 downloaded_count, failed_models = download_all_models()
 print("="*60)
 
-# GRADIO INTERFACE
-with gr.Blocks(
-    title="AI Image Upscaler",
-    theme=gr.themes.Soft(),
-) as demo:
+with gr.Blocks(title="AI Image Upscaler", theme=gr.themes.Soft()) as demo:
     
     gr.HTML("""
     <div style="text-align: center; padding: 2rem 0;">
@@ -410,7 +454,7 @@ with gr.Blocks(
             🚀 AI Image Upscaler
         </h1>
         <p style="font-size: 1.1rem; color: #666;">
-            Professional Image Enhancement powered by Real-ESRGAN & SwinIR
+            Enhanced with HDR-like processing & Smart Tiling
         </p>
     </div>
     """)
@@ -422,15 +466,14 @@ with gr.Blocks(
                     input_image = gr.Image(
                         label="📤 Upload Your Image",
                         type="pil",
-                        sources=["upload", "clipboard"],  # Allow clipboard paste
+                        sources=["upload", "clipboard"],
                         height=400
                     )
                     
                     model_dropdown = gr.Dropdown(
                         choices=list(MODELS.keys()),
                         value="🔥 Real-ESRGAN x4 (Best for 4x)",
-                        label="🎯 Choose AI Model",
-                        info="Select the enhancement model"
+                        label="🎯 Choose AI Model"
                     )
                     
                     output_format = gr.Radio(
@@ -446,11 +489,11 @@ with gr.Blocks(
                     )
                     
                     gr.Markdown("""
-                    ### 💡 Pro Tips
-                    - 🔥 **Real-ESRGAN**: Best for real photos
-                    - 🎨 **SwinIR**: Good for general upscaling
-                    - 📊 Higher scale = Larger output file
-                    - 🖼️ Supports JPG, PNG, WEBP formats
+                    ### 💡 New Features
+                    - 🎨 **HDR-like tone mapping**
+                    - 🔪 **Smart sharpening**
+                    - 🔲 **Optimized tiling**
+                    - 🚀 **Better memory management**
                     """)
                     
                 with gr.Column(scale=1):
@@ -462,8 +505,7 @@ with gr.Blocks(
                     
                     output_info = gr.Textbox(
                         label="📊 Processing Details", 
-                        lines=15,
-                        max_lines=20
+                        lines=15
                     )
         
         with gr.Tab("📊 Model Status"):
@@ -473,7 +515,6 @@ with gr.Blocks(
                 label="Model Status",
                 value=get_model_status(),
                 lines=25,
-                max_lines=30,
                 interactive=False
             )
             
@@ -484,37 +525,27 @@ with gr.Blocks(
             gr.Markdown(f"""
             ## About This App
             
-            This application provides state-of-the-art image upscaling using AI models.
-            
             ### 📈 Statistics
             - **Models Available:** {downloaded_count}/{len(MODELS)}
             - **Device:** {'🚀 GPU (CUDA)' if torch.cuda.is_available() else '💻 CPU'}
-            - **PyTorch Version:** {torch.__version__}
-            - **Gradio Version:** {gr.__version__}
+            - **PyTorch:** {torch.__version__}
+            - **Gradio:** {gr.__version__}
             
-            ### 🎯 Supported Tasks
-            1. **Real-ESRGAN 🔥** - Best quality upscaling for real photos
-            2. **Super Resolution** - Increase image resolution up to 8x
-            3. **Denoising** - Remove noise from images (SwinIR)
-            4. **JPEG Artifact Removal** - Fix compression artifacts (SwinIR)
+            ### ✨ Optimizations
+            - Bilinear upsampling for smooth results
+            - HDR-like tone mapping for better contrast
+            - Smart sharpening (DSLR look)
+            - Memory-efficient tiling for large images
+            - Automatic garbage collection
+            
+            ### 🎯 Supported Models
+            1. **Real-ESRGAN 🔥** - Best for real photos (2x, 4x, 8x)
+            2. **SwinIR** - Lightweight super-resolution (2x, 3x, 4x, 8x)
             
             ### 📚 Model Sources
             - **SwinIR:** [deepinv/swinir](https://huggingface.co/deepinv/swinir)
             - **Real-ESRGAN:** [ai-forever/Real-ESRGAN](https://huggingface.co/ai-forever/Real-ESRGAN)
             
-            ### 🔧 Features
-            - ✅ Multiple AI models
-            - ✅ GPU acceleration
-            - ✅ Batch processing ready
-            - ✅ Multiple output formats
-            - ✅ High-quality upscaling
-            
-            ### 🐛 Troubleshooting
-            If you encounter issues:
-            1. **File upload errors**: Try re-uploading or use a different browser
-            2. **Processing errors**: Check the console logs
-            3. **Slow processing**: GPU acceleration requires CUDA
-            
             ---
             Made with ❤️ using Gradio and PyTorch
             """)
@@ -527,9 +558,4 @@ with gr.Blocks(
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=10)
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )
+    demo.launch()