Update app.py

app.py

@@ -1,90 +1,76 @@
-# Import necessary libraries
+import torch
 from PIL import Image
 import numpy as np
-import torch
 from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
 import gradio as gr
 import spaces
 
-# Function to resize image to max 2048x2048 while maintaining aspect ratio
-def resize_image(image, max_size=2048):
+def split_image(image, chunk_size=512):
     width, height = image.size
-    if width > max_size or height > max_size:
-        aspect_ratio = width / height
-        if width > height:
-            new_width = max_size
-            new_height = int(new_width / aspect_ratio)
-        else:
-            new_height = max_size
-            new_width = int(new_height * aspect_ratio)
-        image = image.resize((new_width, new_height), Image.LANCZOS)
-    return image
+    chunks = []
+    for y in range(0, height, chunk_size):
+        for x in range(0, width, chunk_size):
+            chunk = image.crop((x, y, min(x + chunk_size, width), min(y + chunk_size, height)))
+            chunks.append((chunk, x, y))
+    return chunks
 
-# Function to upscale an image using Swin2SR
-def upscale_image(image, model, processor, device):
-    try:
-        # 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), None
-    except RuntimeError as e:
-        return None, str(e)
+def stitch_image(chunks, original_size):
+    result = Image.new('RGB', original_size)
+    for img, x, y in chunks:
+        result.paste(img, (x, y))
+    return result
+
+def upscale_chunk(chunk, model, processor, device):
+    inputs = processor(chunk, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    with torch.no_grad():
+        outputs = model(**inputs)
+    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)
+    return Image.fromarray(output_image)
 
 @spaces.GPU
 def main(image, model_choice, save_as_jpg=True):
-    # Resize the input image
-    image = resize_image(image)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
-    # 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])
+    model = Swin2SRForImageSuperResolution.from_pretrained(model_paths[model_choice]).to(device)
     
-    # Try GPU first, fallback to CPU if there's an error
-    for device in [torch.device("cuda" if torch.cuda.is_available() else "cpu"), torch.device("cpu")]:
-        model.to(device)
-        upscaled_image, error = upscale_image(image, model, processor, device)
-        
-        if upscaled_image is not None:
-            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"
-        
-        if device.type == "cpu":
-            return f"Error: Unable to process the image. {error}"
-
-    return "Error: Unable to process the image on both GPU and CPU."
+    # Split the image into chunks
+    chunks = split_image(image)
+    
+    # Process each chunk
+    upscaled_chunks = []
+    for chunk, x, y in chunks:
+        upscaled_chunk = upscale_chunk(chunk, model, processor, device)
+        # Remove 32 pixels from bottom and right edges
+        upscaled_chunk = upscaled_chunk.crop((0, 0, upscaled_chunk.width - 32, upscaled_chunk.height - 32))
+        upscaled_chunks.append((upscaled_chunk, x * 4, y * 4))  # Multiply coordinates by 4 due to 4x upscaling
+    
+    # Stitch the chunks back together
+    final_size = (image.width * 4 - 32, image.height * 4 - 32)  # Adjust for removed pixels
+    upscaled_image = stitch_image(upscaled_chunks, final_size)
+    
+    if save_as_jpg:
+        upscaled_image.save("upscaled_image.jpg", quality=95)
+        return "upscaled_image.jpg"
+    else:
+        upscaled_image.save("upscaled_image.png")
+        return "upscaled_image.png"
 
-# Gradio interface
 def gradio_interface(image, model_choice, save_as_jpg):
-    result = main(image, model_choice, save_as_jpg)
-    if result.startswith("Error:"):
-        return gr.update(value=None), result
-    return result, None
+    try:
+        result = main(image, model_choice, save_as_jpg)
+        return result, None
+    except Exception as e:
+        return None, str(e)
 
-# Create a Gradio interface
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=[
@@ -101,8 +87,7 @@ interface = gr.Interface(
         gr.Textbox(label="Error Message", visible=True)
     ],
     title="Image Upscaler",
-    description="Upload an image, select a model, upscale it, and download the new image. Images larger than 2048x2048 will be resized while maintaining aspect ratio. If GPU processing fails, it will attempt to process on CPU.",
+    description="Upload an image, select a model, and upscale it. The image will be processed in 512x512 pixel chunks to handle large images efficiently.",
 )
 
-# Launch the interface
 interface.launch()