Update app.py

app.py

@@ -22,28 +22,28 @@ def resize_image(image, max_size=2048):
 
 # 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().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)
+    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)
 
 @spaces.GPU
 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")
-    
     # Resize the input image
     image = resize_image(image)
     
@@ -56,24 +56,33 @@ def main(image, model_choice, save_as_jpg=True):
     # 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)
     
-    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"
+    # 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."
 
 # Gradio interface
 def gradio_interface(image, model_choice, save_as_jpg):
-    return main(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
 
 # Create a Gradio interface
 interface = gr.Interface(
@@ -87,9 +96,12 @@ interface = gr.Interface(
         ),
         gr.Checkbox(value=True, label="Save as JPEG"),
     ],
-    outputs=gr.File(label="Download Upscaled Image"),
+    outputs=[
+        gr.File(label="Download Upscaled Image"),
+        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.",
+    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.",
 )
 
 # Launch the interface