Fix error handling and add image Resizing for stable diffusion upscaler

app.py

@@ -1,32 +1,36 @@
 import gradio as gr
 from PIL import Image
-import requests
-from io import BytesIO
-
+import torch
 from diffusers import StableDiffusionUpscalePipeline
 
-# Load the upscaling pipeline (this may take some seconds at first run)
-pipe = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler")
-pipe = pipe.to("cpu")  # use cpu, can switch to "cuda" if GPU available
+# Load the pipeline
+pipe = StableDiffusionUpscalePipeline.from_pretrained(
+    "stabilityai/stable-diffusion-x4-upscaler", 
+    revision="fp16", 
+    torch_dtype=torch.float16
+)
+pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
 def upscale_image(input_image):
-    # input_image is a PIL image from Gradio uploader
-    if input_image is None:
-        return None
-    # Run the upscaler pipeline
-    upscaled = pipe(input_image).images[0]
-    return upscaled
+    try:
+        if input_image is None:
+            return "Please upload an image."
+        
+        # Resize image to match what the model expects
+        input_image = input_image.resize((512, 512))  # model requires 512x512 input
 
-title = "Free AI Image Upscaler"
-description = "Upload an image, and this app will upscale it 4x using Stable Diffusion upscaler."
+        # Run the model with a default prompt
+        result = pipe(prompt="a high quality image", image=input_image).images[0]
+        return result
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 iface = gr.Interface(
     fn=upscale_image,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title=title,
-    description=description,
-    allow_flagging="never"
+    outputs="image",
+    title="Free AI Image Upscaler",
+    description="Upload an image to upscale using Stable Diffusion x4 Upscaler."
 )
 
 if __name__ == "__main__":