Update app.py

app.py

@@ -1,26 +1,29 @@
 from diffusers import StableDiffusionPipeline
 import torch
-
-modelieo=[
- 'naclbit/trinart_stable_diffusion_v2',
-]
-
-
-def TextToImage(Prompt,model):
-  model_id = model
-  pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-  pipe = pipe.to("cpu")
-
-  prompt = Prompt
-  image = pipe(prompt).images[0]
-
-  return image
-
-
 import gradio as gr
-interface = gr.Interface(fn=TextToImage, 
-                        inputs=["text", gr.Dropdown(modelieo)],
-                         outputs="image", 
-                        title='Text to Image')
 
-interface.launch()
+# Load the model
+model_id = "CompVis/stable-diffusion-v-1-4-original"
+pipe = StableDiffusionPipeline.from_pretrained(model_id)
+pipe = pipe.to("cuda")  # Use GPU if available
+
+def generate_image(prompt):
+    image = pipe(prompt).images[0]
+    return image
+
+# Define the chatbot function
+def chatbot(prompt):
+    image = generate_image(prompt)
+    return image
+
+# Create the Gradio interface
+interface = gr.Interface(
+    fn=chatbot,
+    inputs="text",
+    outputs="image",
+    title="Text to Image Chatbot",
+    description="Generate images from text using Stable Diffusion"
+)
+
+# Launch the interface
+interface.launch()