Update app.py

app.py

@@ -1,5 +1,21 @@
-import gradio as gr
-
+import torch
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+import io
 
 # Load the model
-gr.load("CompVis/stable-diffusion-v1-4").launch()
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
+pipe = pipe.to("cuda")
+
+def generate_image(prompt):
+    image = pipe(prompt).images[0]
+    return image
+
+# Take prompt as input
+prompt = input("Enter a prompt: ")
+
+# Generate the image
+generated_image = generate_image(prompt)
+
+# Display the generated image
+generated_image.show()