Update app.py

app.py

@@ -44,11 +44,11 @@ def upscale_chunk(chunk, model, processor, device):
     return Image.fromarray(output_image)
 
 @spaces.GPU
-def main(image, model_choice, save_as_jpg=True):
+def main(image, model_choice, save_as_jpg=True, use_tiling=True, auto_cpu=True):
     # Resize the input image
     image = resize_image(image)
     
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() and not auto_cpu else "cpu")
     
     model_paths = {
         "Pixel Perfect": "caidas/swin2SR-classical-sr-x4-64",
@@ -58,20 +58,24 @@ def main(image, model_choice, save_as_jpg=True):
     processor = AutoImageProcessor.from_pretrained(model_paths[model_choice])
     model = Swin2SRForImageSuperResolution.from_pretrained(model_paths[model_choice]).to(device)
     
-    # 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 use_tiling:
+        # 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)
+    else:
+        # Process the entire image at once
+        upscaled_image = upscale_chunk(image, model, processor, device)
     
     if save_as_jpg:
         upscaled_image.save("upscaled_image.jpg", quality=95)
@@ -80,9 +84,9 @@ def main(image, model_choice, save_as_jpg=True):
         upscaled_image.save("upscaled_image.png")
         return "upscaled_image.png"
 
-def gradio_interface(image, model_choice, save_as_jpg):
+def gradio_interface(image, model_choice, save_as_jpg, use_tiling, auto_cpu):
     try:
-        result = main(image, model_choice, save_as_jpg)
+        result = main(image, model_choice, save_as_jpg, use_tiling, auto_cpu)
         return result, None
     except Exception as e:
         return None, str(e)
@@ -97,13 +101,15 @@ interface = gr.Interface(
             value="PSNR Match (Recommended)"
         ),
         gr.Checkbox(value=True, label="Save as JPEG"),
+        gr.Checkbox(value=True, label="Use Tiling"),
+        gr.Checkbox(value=True, label="Auto CPU"),
     ],
     outputs=[
         gr.File(label="Download Upscaled Image"),
         gr.Textbox(label="Error Message", visible=True)
     ],
     title="Image Upscaler",
-    description="Upload an image, select a model, and upscale it. Images larger than 2048x2048 will be resized while maintaining aspect ratio. The image will be processed in 512x512 pixel chunks for efficient handling.",
+    description="Upload an image, select a model, and upscale it. Images larger than 2048x2048 will be resized while maintaining aspect ratio. Use tiling for efficient processing of large images. Auto CPU will use CPU if GPU is not available.",
 )
 
 interface.launch()