Update app.py

app.py

@@ -1,36 +1,33 @@
-import os
 import gradio as gr
 import torch
-from PIL import Image
-from torchvision.transforms import ToTensor, ToPILImage
-from torchvision import models
+from realesrgan import RealESRGAN
 
-# 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')
-
-def upscale_image(image, scale_factor):
-    # Ensure image is in PIL format
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-
-    # Upscale using the model
-    with torch.no_grad():
-        upscaled_image = model(image, scale=scale_factor)
+# Load the model
+model = RealESRGAN('weights/RealESRGAN_x4plus_anime_6B.pth', scale=4)  # Adjust scale if needed
 
+def upscale_image(input_image, scale_factor):
+    # Ensure scale factor is an integer
+    scale_factor = int(scale_factor)
+    
+    # Set the model scale
+    model.scale = scale_factor
+    
+    # Perform the upscaling
+    upscaled_image = model.predict(input_image)
+    
     return upscaled_image
 
-# Create Gradio interface
+# Create the Gradio interface
 iface = gr.Interface(
     fn=upscale_image,
     inputs=[
-        gr.inputs.Image(type="pil", label="Input Image"),
-        gr.inputs.Slider(minimum=1, maximum=4, step=1, default=2, label="Scale Factor")
+        gr.inputs.Image(type="numpy", label="Input Image"),
+        gr.inputs.Slider(minimum=2, maximum=4, step=1, default=4, 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."
+    outputs=gr.outputs.Image(type="numpy", label="Upscaled Image"),
+    title="Image Upscaler using R-ESRGAN Anime 6B",
+    description="Upload an image and select a scale factor to upscale it using the R-ESRGAN Anime 6B model."
 )
 
-# Launch the Gradio app
-if __name__ == "__main__":
-    iface.launch()
+# Launch the app
+iface.launch()