File size: 805 Bytes
c2f314c 0d29e9f c2f314c 0d29e9f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
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() |