Update app.py

app.py

@@ -10,41 +10,38 @@ import spaces
 def upscale_image(image, model, processor, device):
     # Convert the image to RGB format
     image = image.convert("RGB")
-
     # Process the image for the model
     inputs = processor(image, return_tensors="pt")
-
     # Move inputs to the same device as model
     inputs = {k: v.to(device) for k, v in inputs.items()}
-
     # Perform inference (upscale)
     with torch.no_grad():
         outputs = model(**inputs)
-
     # Move output back to CPU for further processing
     output = outputs.reconstruction.data.squeeze().cpu().float().clamp_(0, 1).numpy()
     output = np.moveaxis(output, source=0, destination=-1)
     output_image = (output * 255.0).round().astype(np.uint8)  # Convert from float32 to uint8
-
     # Remove 32 pixels from the bottom and right of the image
     output_image = output_image[:-32, :-32]
-
     return Image.fromarray(output_image)
 
 @spaces.GPU
-def main(image, save_as_jpg=True):
+def main(image, model_choice, save_as_jpg=True):
     # Check if GPU is available and set the device accordingly
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
-    realworld_model = "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr"
-
-    # Load the Swin2SR model and processor for 4x upscaling
-    processor = AutoImageProcessor.from_pretrained(realworld_model)
-    model = Swin2SRForImageSuperResolution.from_pretrained(realworld_model)
-
+    # Define model paths
+    model_paths = {
+        "Pixel Perfect": "caidas/swin2SR-classical-sr-x4-64",
+        "PSNR Match (Recommended)": "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr"
+    }
+    
+    # Load the selected Swin2SR model and processor for 4x upscaling
+    processor = AutoImageProcessor.from_pretrained(model_paths[model_choice])
+    model = Swin2SRForImageSuperResolution.from_pretrained(model_paths[model_choice])
     # Move the model to the device (GPU or CPU)
     model.to(device)
-
+    
     # Upscale the image
     upscaled_image = upscale_image(image, model, processor, device)
     
@@ -58,19 +55,24 @@ def main(image, save_as_jpg=True):
         return "upscaled_image.png"
 
 # Gradio interface
-def gradio_interface(image, save_as_jpg):
-    return main(image, save_as_jpg)
+def gradio_interface(image, model_choice, save_as_jpg):
+    return main(image, model_choice, save_as_jpg)
 
 # Create a Gradio interface
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
+        gr.Dropdown(
+            choices=["PSNR Match (Recommended)", "Pixel Perfect"],
+            label="Select Model",
+            value="PSNR Match (Recommended)"
+        ),
         gr.Checkbox(value=True, label="Save as JPEG"),
     ],
     outputs=gr.File(label="Download Upscaled Image"),
     title="Image Upscaler",
-    description="Upload an image, upscale it, and download the new image.",
+    description="Upload an image, select a model, upscale it, and download the new image.",
 )
 
 # Launch the interface