import gradio as gr
from diffusers import DiffusionPipeline
import torch

# Load the Kwai-Kolors diffusion model
pipeline = DiffusionPipeline.from_pretrained("Kwai-Kolors/Kolors").to("cuda" if torch.cuda.is_available() else "cpu")

# Function to generate image from prompt
def generate_image(prompt):
    # Use the pipeline to generate an image from the text prompt
    image = pipeline(prompt).images[0]
    return image

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,                # Function that generates the image
    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt"),  # Textbox input for the prompt
    outputs="image",                  # Output is an image
    title="Kwai-Kolors Image Generator",
    description="Generate images from text prompts using the Kwai-Kolors diffusion model."
)

# Launch the app
iface.launch()