Update

app.py

@@ -3,63 +3,65 @@ import torch
 import numpy as np
 from PIL import Image
 from torchvision.transforms import ToTensor, ToPILImage
-from typing import Tuple, Optional
+from urllib.request import urlretrieve
+import os
 
 # Device configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Constants
-SUPPORTED_FORMATS = ["JPEG", "PNG", "WEBP"]
+MODEL_URL = "https://github.com/xinntao/ESRGAN/releases/download/v0.1.1/RRDB_ESRGAN_x4.pth"
+MODEL_PATH = "RRDB_ESRGAN_x4.pth"
 MAX_IMAGE_SIZE = (1024, 1024)
 
+# ESRGAN model architecture
+class RRDBNet(torch.nn.Module):
+    def __init__(self, in_nc=3, out_nc=3, nf=64, nb=23, gc=32):
+        super(RRDBNet, self).__init__()
+        self.model = self._make_network(in_nc, out_nc, nf, nb, gc)
+
+    def _make_network(self, in_nc, out_nc, nf, nb, gc):
+        # [Original architecture implementation here...]
+        # Full implementation: https://github.com/xinntao/ESRGAN/blob/master/RRDBNet_arch.py
+
 def load_model() -> torch.nn.Module:
-    """Load pretrained ESRGAN model"""
-    model = torch.hub.load(
-        "facebookresearch/AnimatedDrawings",
-        "esrgan",
-        pretrained=True,
-        verbose=False,
-        trust_repo=True
-    )
+    """Download and load ESRGAN model"""
+    if not os.path.exists(MODEL_PATH):
+        print("Downloading ESRGAN model...")
+        urlretrieve(MODEL_URL, MODEL_PATH)
+    
+    model = RRDBNet()
+    state_dict = torch.load(MODEL_PATH, map_location=device)
+    model.load_state_dict(state_dict)
     return model.to(device).eval()
 
 def preprocess_image(image: Image.Image) -> torch.Tensor:
-    """Convert PIL image to preprocessed tensor"""
+    """Convert PIL image to normalized tensor"""
     transform = ToTensor()
-    tensor = transform(image).unsqueeze(0).to(device)
-    return tensor * 2.0 - 1.0  # ESRGAN requires [-1,1] normalization
+    return transform(image).unsqueeze(0).to(device)
 
 def postprocess_image(tensor: torch.Tensor) -> Image.Image:
-    """Convert model output tensor to PIL image"""
+    """Convert tensor to PIL image"""
     transform = ToPILImage()
-    tensor = (tensor + 1.0) / 2.0  # Convert back to [0,1]
     tensor = tensor.squeeze(0).detach().cpu().clamp(0, 1)
     return transform(tensor)
 
-def validate_image(image: Image.Image) -> None:
-    """Validate input image dimensions and format"""
+def validate_image(image: Image.Image):
+    """Validate input image constraints"""
     if image.mode not in ["RGB", "RGBA"]:
         raise gr.Error("Only RGB/RGBA images supported")
-    if image.size[0] > MAX_IMAGE_SIZE[0] or image.size[1] > MAX_IMAGE_SIZE[1]:
-        raise gr.Error(f"Max image size {MAX_IMAGE_SIZE} exceeded")
+    if max(image.size) > max(MAX_IMAGE_SIZE):
+        raise gr.Error(f"Max image dimension exceeded ({MAX_IMAGE_SIZE[0]}x{MAX_IMAGE_SIZE[1]})")
 
 def enhance_image(
     input_image: Image.Image,
     scale_factor: float = 2.0
 ) -> Image.Image:
-    """
-    Enhance image using ESRGAN model
-    Args:
-        input_image: PIL Image to process
-        scale_factor: Multiplier for image scaling (2.0 or 4.0)
-    Returns:
-        Enhanced PIL Image
-    """
+    """Main processing function"""
     try:
         validate_image(input_image)
-        original_size = input_image.size
         
-        # Convert RGBA to RGB if needed
+        # Convert RGBA to RGB
         if input_image.mode == 'RGBA':
             input_image = input_image.convert('RGB')
         
@@ -68,62 +70,30 @@ def enhance_image(
             output_tensor = model(input_tensor)
         
         result = postprocess_image(output_tensor)
-        result = result.resize(
-            (int(original_size[0]*scale_factor), 
-             int(original_size[1]*scale_factor)),
+        return result.resize(
+            (int(input_image.width*scale_factor), 
+             int(input_image.height*scale_factor)),
             Image.LANCZOS
         )
-        return result
-    
+        
     except Exception as e:
-        raise gr.Error(f"Image processing error: {str(e)}")
+        raise gr.Error(f"Processing error: {str(e)}")
 
-# Load model once at startup
+# Initialize model
 model = load_model()
 
-# Gradio interface configuration
+# Gradio interface
 interface = gr.Interface(
     fn=enhance_image,
     inputs=[
-        gr.Image(
-            label="Input Image",
-            type="pil",
-            image_mode="RGB",
-            sources=["upload"],
-            elem_id="input_image"
-        ),
-        gr.Slider(
-            minimum=2.0,
-            maximum=4.0,
-            value=2.0,
-            step=2.0,
-            label="Upscale Factor",
-            info="Select 2x or 4x upscaling"
-        )
-    ],
-    outputs=gr.Image(
-        label="Enhanced Image",
-        type="pil",
-        elem_id="output_image"
-    ),
-    title="🖼️ AI Image Enhancer",
-    description="Enhance image quality using ESRGAN super-resolution (2x/4x upscaling)",
-    examples=[
-        ["examples/example1.jpg", 2.0],
-        ["examples/example2.png", 4.0]
+        gr.Image(type="pil", label="Input Image"),
+        gr.Slider(2.0, 4.0, 2.0, step=2.0, label="Scale Factor")
     ],
-    allow_flagging="never",
-    css="""
-    footer {visibility: hidden}
-    .gradio-container {max-width: 800px !important}
-    """
+    outputs=gr.Image(type="pil", label="Enhanced Image"),
+    title="🎨 AI Image Enhancer",
+    examples=[["examples/example1.jpg", 2.0]],
+    css=".gradio-container {max-width: 800px !important}"
 )
 
-# Deployment configuration
 if __name__ == "__main__":
-    interface.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        debug=False
-    )
+    interface.launch(server_name="0.0.0.0")