# This Gradio app uses the KolorsPipeline from the diffusers library to generate images based on a given prompt.
import gradio as gr
import torch
from diffusers import KolorsPipeline

# Load the KolorsPipeline model
pipe = KolorsPipeline.from_pretrained(
    "Kwai-Kolors/Kolors-diffusers", 
    torch_dtype=torch.float16, 
    variant="fp16"
).to("cuda")

# Define the function to generate an image based on the prompt
def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps, seed):
    generator = torch.Generator(pipe.device).manual_seed(seed)
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator,
    ).images[0]
    return image

# Create the Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt", value="一张瓢虫的照片，微距，变焦，高质量，电影，拿着一个牌子，写着'可图'")
        negative_prompt_input = gr.Textbox(label="Negative Prompt", value="")
    with gr.Row():
        guidance_scale_slider = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, value=5.0, step=0.1)
        num_inference_steps_slider = gr.Slider(label="Number of Inference Steps", minimum=1, maximum=100, value=50, step=1)
        seed_slider = gr.Slider(label="Seed", minimum=0, maximum=100000, value=66, step=1)
    generate_button = gr.Button("Generate Image")
    output_image = gr.Image(label="Generated Image")

    # Define the event listener for the generate button
    generate_button.click(
        fn=generate_image,
        inputs=[prompt_input, negative_prompt_input, guidance_scale_slider, num_inference_steps_slider, seed_slider],
        outputs=output_image
    )

# Launch the Gradio app
demo.launch(show_error=True)