import gradio as gr
import torch
from diffusers import KolorsPipeline

pipeline= KolorsPipeline.from_pretrained(
    "Kwai-Kolors/Kolors-diffusers", 
    torch_dtype=torch.float16, 
    variant="fp16"
).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()