Update app.py

app.py

@@ -1,12 +1,15 @@
 import gradio as gr
-from transformers import pipeline
+from diffusers import StableDiffusionPipeline
+import torch
 
-# Load the text-to-image model (like Stable Diffusion or DALL·E)
-generator = pipeline("text-to-image", model="CompVis/stable-diffusion-v-1-4-original")
+# Load Stable Diffusion pipeline (make sure to use the right model name)
+model_name = "CompVis/stable-diffusion-v-1-4-original"
+pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
+pipe.to("cuda")  # Move the model to GPU (if available)
 
 def generate_image(description):
     # Generate an image based on the description
-    image = generator(description)[0]["image"]
+    image = pipe(description).images[0]
     return image
 
 # Set up the Gradio interface