Update app.py

app.py

@@ -1,24 +1,29 @@
 import gradio as gr
-from transformers import pipeline
+from diffusers import StableDiffusionPipeline
+import torch
 
-# Load Stable Diffusion model from Hugging Face
-generator = pipeline("text-to-image", model="stabilityai/stable-diffusion-3-medium-diffusers")
+# Load the model (you can replace this with the specific Kwai-Kolors/Kolors model if available)
+pipe = StableDiffusionPipeline.from_pretrained("Kwai-Kolors/Kolors")
+pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Define the chatbot function
-def generate_image(prompt: str):
-    # Generate the image from the prompt
-    image = generator(prompt)[0]["image"]
+# Define function to generate image from text
+def generate_image(prompt):
+    # Generate an image based on the provided text prompt
+    image = pipe(prompt).images[0]
     return image
 
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=generate_image, 
-    inputs=gr.Textbox(label="Enter a description of the image"), 
-    outputs=gr.Image(label="Generated Image"),
-    live=True
+# Create Gradio interface for chatbot
+interface = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Textbox(label="Enter your prompt:"),
+    outputs=gr.Image(type="pil"),
+    live=True,
+    title="Text-to-Image Generation",
+    description="Enter a text prompt to generate a photorealistic image."
 )
 
-# Launch the Gradio interface
-iface.launch(share=True)
+# Launch the interface
+interface.launch()
+