File size: 1,134 Bytes
c2f314c 0d29e9f c2f314c 0d29e9f dac85aa 971eeee 0d29e9f dac85aa 0d29e9f dac85aa 0d29e9f dac85aa 0d29e9f dac85aa 0d29e9f dac85aa 0d29e9f dac85aa |
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 36 37 38 39 40 41 42 43 44 |
import gradio as gr
import jax
from diffusers import FlaxStableDiffusionPipeline
pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
"bguisard/stable-diffusion-nano",
)
def generate_image(prompt: str, inference_steps: int, prng_seed: int):
rng = jax.random.PRNGKey(prng_seed)
prompt_ids = pipeline.prepare_inputs(prompt)
images = pipeline(
prompt_ids=prompt_ids,
params=pipeline_params,
prng_seed=rng,
height=128,
width=128,
num_inference_steps=int(inference_steps),
jit=False,
).images
pil_imgs = pipeline.numpy_to_pil(images)
return pil_imgs[0]
prompt_input = gr.inputs.Textbox(
label="Prompt", placeholder="A watercolor painting of a bird"
)
inf_steps_input = gr.inputs.Slider(
minimum=1, maximum=100, default=30, step=1, label="Inference Steps"
)
seed_input = gr.inputs.Number(default=0, label="Seed")
app = gr.Interface(
fn=generate_image,
inputs=[prompt_input, inf_steps_input, seed_input],
outputs=gr.Image(shape=(128, 128)),
examples=[["A watercolor painting of a bird"]],
)
app.launch()
|