import gradio as gr
import jax
from diffusers import FlaxStableDiffusionPipeline

pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
    "bguisard/stable-diffusion-nano",
)

prng_seed = jax.random.PRNGKey(0)
inference_steps = 50


def generate_image(prompt: str):
    prompt_ids = pipeline.prepare_inputs(prompt)
    images = pipeline(
        prompt_ids=prompt_ids,
        params=pipeline_params,
        prng_seed=prng_seed,
        height=128,
        width=128,
        num_inference_steps=inference_steps,
        jit=False,
    ).images
    pil_imgs = pipeline.numpy_to_pil(images)
    return pil_imgs[0]


app = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs=gr.Image(shape=(128, 128)),
    examples=[["A watercolor painting of a bird"]],
)

app.launch()