Update app.py

app.py

@@ -1,16 +1,33 @@
+import gradio as gr
 import torch
 from diffusers import StableDiffusionPipeline
+from PIL import Image
+import io
 
+# Load the model
 pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
-
 pipe = pipe.to("cuda")
 
-prompt = "a photograph of an astronaut riding a horse"
-
-image = pipe(prompt).images[0]  # image here is in [PIL format](https://pillow.readthedocs.io/en/stable/)
+def generate_image(prompt):
+    image = pipe(prompt).images[0]
+    # Save the image to a bytes buffer
+    img_buffer = io.BytesIO()
+    image.save(img_buffer, format="PNG")
+    img_buffer.seek(0)
+    return img_buffer
 
-# Now to display an image you can either save it such as:
-image.save(f"astronaut_rides_horse.png")
+iface = gr.Interface(
+    fn=generate_image,
+    inputs="text",
+    outputs="image",
+    layout="vertical",
+    inputs_label="Enter Prompt",
+    outputs_label="Generated Image",
+    examples=[
+        ["a photograph of an astronaut riding a horse"],
+        ["a beautiful sunset over the mountains"]
+    ]
+)
 
-# or if you're in a google colab you can directly display it with 
-image
+if __name__ == "__main__":
+    iface.launch()