Update app.py

app.py

@@ -1,17 +1,15 @@
-import os
 import gradio as gr
 from diffusers import StableDiffusionUpscalePipeline
 from diffusers.utils import load_image
 import torch
 from PIL import Image
 import base64
-import requests
 from io import BytesIO
 
 # Load model and scheduler
 model_id = "stabilityai/stable-diffusion-x4-upscaler"
-pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipeline = pipeline.to("cuda")
+pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id)
+pipeline = pipeline.to("cpu")  # Use CPU instead of GPU
 
 def upscale_image(image, prompt):
     image = image.resize((128, 128))  # Resize to the expected input size
@@ -27,10 +25,10 @@ def base64_to_image(base64_str):
     image_data = base64.b64decode(base64_str)
     return Image.open(BytesIO(image_data))
 
-def process_image(image):
-    prompt = "a white cat"
+def handle_upload(image, prompt):
     upscaled_image = upscale_image(image, prompt)
-    return image_to_base64(upscaled_image)
+    base64_str = image_to_base64(upscaled_image)
+    return base64_str, upscaled_image
 
 def main():
     with gr.Blocks() as demo:
@@ -45,11 +43,6 @@ def main():
                 base64_output = gr.Textbox(label="Base64 Encoded Image")
                 upscaled_image_output = gr.Image(type="pil", label="Upscaled Image")
 
-                def handle_upload(image, prompt):
-                    upscaled_image = upscale_image(image, prompt)
-                    base64_str = image_to_base64(upscaled_image)
-                    return base64_str, upscaled_image
-
                 upload_btn.click(fn=handle_upload, inputs=[image_input, prompt_input], outputs=[base64_output, upscaled_image_output])
             
     demo.launch()