import torch
from diffusers import AutoPipelineForText2Image
from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
import gradio as gr

if torch.cuda.is_available():
    pipe = AutoPipelineForText2Image.from_pretrained("warp-ai/wuerstchen", torch_dtype=torch.float16).to("cuda")
else:
    pipe = AutoPipelineForText2Image.from_pretrained("warp-ai/wuerstchen")

def gen_image(caption):
    # caption = "Anthropomorphic fox dressed as a fire fighter"
    # fantasy art of a band of brothers human cleric by greg rutkowski
    images = pipe(
        caption,
        width=1024,
        height=1536,
        prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS,
        prior_guidance_scale=4.0,
        num_images_per_prompt=1,
    ).images
    return images[0]


#  Once primed, 1024 x 1536 images are generated in ~6 seconds on my machine; quality is so-so but similar to SD 1.5
with gr.Blocks() as demo:
    gr.Markdown("# Test of the Wueurstchen Model")
    with gr.Row():
        with gr.Column():
            in_text = gr.Textbox(value="Enter a prompt here")
            run_button = gr.Button(variant="primary")
        with gr.Column():
            out_image = gr.Image(label="Image Output")

    run_button.click(gen_image, [in_text], [out_image], None)

if __name__ == '__main__':
    demo.launch()