Create app.py

app.py

@@ -0,0 +1,77 @@
+# Import necessary libraries
+from PIL import Image
+import torch
+import numpy as np
+from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
+import gradio as gr  # Import Gradio for creating the interface
+
+# Function to upscale an image using Swin2SR
+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().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):
+    # Check if GPU is available and set the device accordingly
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    regular_model = "caidas/swin2SR-classical-sr-x4-64"
+    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)
+
+    # Move the model to the device (GPU or CPU)
+    model.to(device)
+
+    # Upscale the image
+    upscaled_image = upscale_image(image, model, processor, device)
+    
+    if save_as_jpg:
+        # Save the upscaled image as JPG with 98% compression
+        upscaled_image.save("upscaled_image.jpg", quality=98)
+        return "upscaled_image.jpg"
+    else:
+        # Save the upscaled image as PNG
+        upscaled_image.save("upscaled_image.png")
+        return "upscaled_image.png"
+
+# Gradio interface
+def gradio_interface(image, save_as_jpg):
+    return main(image, save_as_jpg)
+
+# Create a Gradio interface
+interface = gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.inputs.Image(type="pil", label="Upload Image"),
+        gr.inputs.Checkbox(default=True, label="Save as JPEG"),
+    ],
+    outputs=gr.outputs.File(label="Download Upscaled Image"),
+    title="Image Upscaler",
+    description="Upload an image, upscale it, and download the new image.",
+)
+
+# Launch the interface
+interface.launch()