Create app.py

app.py

@@ -0,0 +1,58 @@
+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")
+
+def upscale_image(image, prompt):
+    image = image.resize((128, 128))  # Resize to the expected input size
+    upscaled_image = pipeline(prompt=prompt, image=image).images[0]
+    return upscaled_image
+
+def image_to_base64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    return base64.b64encode(buffered.getvalue()).decode()
+
+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"
+    upscaled_image = upscale_image(image, prompt)
+    return image_to_base64(upscaled_image)
+
+def main():
+    with gr.Blocks() as demo:
+        gr.Markdown("# Stable Diffusion Upscaler")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.Image(type="pil", label="Low-Resolution Image", tool="editor")
+                prompt_input = gr.Textbox(label="Prompt", value="a white cat")
+
+                upload_btn = gr.Button("Upload and Upscale")
+                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()
+
+if __name__ == "__main__":
+    main()