Update app.py

app.py

@@ -13,12 +13,13 @@ SUPPORTED_FORMATS = ["JPEG", "PNG", "WEBP"]
 MAX_IMAGE_SIZE = (1024, 1024)
 
 def load_model() -> torch.nn.Module:
-    """Load pretrained ESRGAN model from torch hub"""
+    """Load pretrained ESRGAN model"""
     model = torch.hub.load(
-        "pytorch/vision",
+        "facebookresearch/AnimatedDrawings",
         "esrgan",
         pretrained=True,
-        verbose=False
+        verbose=False,
+        trust_repo=True
     )
     return model.to(device).eval()
 
@@ -26,20 +27,20 @@ def preprocess_image(image: Image.Image) -> torch.Tensor:
     """Convert PIL image to preprocessed tensor"""
     transform = ToTensor()
     tensor = transform(image).unsqueeze(0).to(device)
-    return tensor
+    return tensor * 2.0 - 1.0  # ESRGAN requires [-1,1] normalization
 
 def postprocess_image(tensor: torch.Tensor) -> Image.Image:
     """Convert model output tensor to PIL image"""
     transform = ToPILImage()
-    tensor = tensor.squeeze(0).detach().cpu()
-    tensor = torch.clamp(tensor, 0, 1)
+    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"""
     if image.mode not in ["RGB", "RGBA"]:
         raise gr.Error("Only RGB/RGBA images supported")
-    if image.size > MAX_IMAGE_SIZE:
+    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")
 
 def enhance_image(
@@ -50,23 +51,32 @@ def enhance_image(
     Enhance image using ESRGAN model
     Args:
         input_image: PIL Image to process
-        scale_factor: Multiplier for image scaling (1.0-4.0)
+        scale_factor: Multiplier for image scaling (2.0 or 4.0)
     Returns:
         Enhanced PIL Image
     """
     try:
-        # Input validation
         validate_image(input_image)
+        original_size = input_image.size
+        
+        # Convert RGBA to RGB if needed
+        if input_image.mode == 'RGBA':
+            input_image = input_image.convert('RGB')
         
-        # Model processing
         with torch.no_grad():
             input_tensor = preprocess_image(input_image)
             output_tensor = model(input_tensor)
         
-        return postprocess_image(output_tensor)
+        result = postprocess_image(output_tensor)
+        result = result.resize(
+            (int(original_size[0]*scale_factor), 
+             int(original_size[1]*scale_factor)),
+            Image.LANCZOS
+        )
+        return result
     
     except Exception as e:
-        raise gr.Error(f"Image processing failed: {str(e)}")
+        raise gr.Error(f"Image processing error: {str(e)}")
 
 # Load model once at startup
 model = load_model()
@@ -83,12 +93,12 @@ interface = gr.Interface(
             elem_id="input_image"
         ),
         gr.Slider(
-            minimum=1.0,
+            minimum=2.0,
             maximum=4.0,
             value=2.0,
-            step=0.5,
-            label="Scale Factor",
-            info="Select upscaling multiplier (1x to 4x)"
+            step=2.0,
+            label="Upscale Factor",
+            info="Select 2x or 4x upscaling"
         )
     ],
     outputs=gr.Image(
@@ -97,13 +107,16 @@ interface = gr.Interface(
         elem_id="output_image"
     ),
     title="🖼️ AI Image Enhancer",
-    description="Enhance image quality using ESRGAN super-resolution model (Supports 2x-4x upscaling)",
+    description="Enhance image quality using ESRGAN super-resolution (2x/4x upscaling)",
     examples=[
         ["examples/example1.jpg", 2.0],
         ["examples/example2.png", 4.0]
     ],
     allow_flagging="never",
-    css="footer {visibility: hidden}"
+    css="""
+    footer {visibility: hidden}
+    .gradio-container {max-width: 800px !important}
+    """
 )
 
 # Deployment configuration