Update app.py

app.py

@@ -1,63 +1,36 @@
+import os
 import gradio as gr
 import torch
-import torchvision.transforms as transforms
 from PIL import Image
-import os
-import requests
-import numpy as np
-
-# Define the model URL and local path
-MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
-MODEL_PATH = "RealESRGAN_x4_Anime6B.pth"
-
-# Download the model if it doesn't exist
-if not os.path.exists(MODEL_PATH):
-    response = requests.get(MODEL_URL)
-    if response.status_code == 200:
-        with open(MODEL_PATH, 'wb') as f:
-            f.write(response.content)
-    else:
-        raise Exception(f"Could not download model: {response.status_code}")
-
-# Define the Real-ESRGAN class
-class RealESRGAN:
-    def __init__(self, model_path, device):
-        self.model = torch.hub.load('xinntao/Real-ESRGAN', 'real_esrgan', model_path, device=device, trust_repo=True)
-        self.model.eval()  # Set the model to eval mode
-        self.device = device  # Store the device
-    
-    def upscale(self, image_tensor):
-        with torch.no_grad():
-            sr_image_tensor = self.model(image_tensor).clamp(0, 1)
-        return sr_image_tensor
+from torchvision.transforms import ToTensor, ToPILImage
+from torchvision import models
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = RealESRGAN(MODEL_PATH, device)
+# Load the R-ESRGAN Anime model (you need to set up this model properly)
+model = torch.hub.load('xinntao/Real-ESRGAN', 'restorer', model='R-ESRGAN_Anime_X6')
 
-# Define the image upscaling function
-def upscale_image(image):
-    # Pre-process the input image
-    image = image.convert("RGB")
-    transform = transforms.ToTensor()
-    image_tensor = transform(image).unsqueeze(0).to(device)  # Add batch dimension and move to device
+def upscale_image(image, scale_factor):
+    # Ensure image is in PIL format
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
 
-    # Perform upscaling
-    sr_image_tensor = model.upscale(image_tensor)
+    # Upscale using the model
+    with torch.no_grad():
+        upscaled_image = model(image, scale=scale_factor)
 
-    # Post-process the output image and convert back to PIL
-    sr_image = (sr_image_tensor.squeeze(0).cpu().numpy() * 255).astype(np.uint8)
-    sr_image = Image.fromarray(sr_image.transpose(1, 2, 0))  # Change to HWC format
-    return sr_image
+    return upscaled_image
 
-# Create a Gradio interface
-interface = gr.Interface(
+# Create Gradio interface
+iface = gr.Interface(
     fn=upscale_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Image Upscaler with R-ESRGAN Anime 6B",
-    description="Upload an image, and it'll be upscaled using the R-ESRGAN Anime model. Supports inputs in various formats."
+    inputs=[
+        gr.inputs.Image(type="pil", label="Input Image"),
+        gr.inputs.Slider(minimum=1, maximum=4, step=1, default=2, label="Scale Factor")
+    ],
+    outputs=gr.outputs.Image(type="pil", label="Upscaled Image"),
+    title="R-ESRGAN Anime 6B Image Upscaler",
+    description="Upload an image and select a scale factor to upscale the image using R-ESRGAN Anime 6B model."
 )
 
-# Launch the interface
+# Launch the Gradio app
 if __name__ == "__main__":
-    interface.launch()
+    iface.launch()