File size: 1,440 Bytes
c2f314c 0d29e9f c2f314c 0d29e9f bb347bc c4c4b83 0d29e9f dac85aa 0d29e9f dac85aa b314fb6 0d29e9f dac85aa 0d29e9f dac85aa 0d29e9f dac85aa 0d29e9f 3f7bb18 c4c4b83 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 45 46 47 48 49 50 |
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 = 30, prng_seed: int = 0):
rng = jax.random.PRNGKey(int(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=True,
).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)),
title="Stable Diffusion Nano",
description=(
"Based on stable diffusion and fine-tuned on 128x128 images, "
"Stable Diffusion Nano allows for fast prototyping of diffusion models, "
"enabling quick experimentation with easily available hardware."
),
examples=[["A watercolor painting of a bird", 30, 0]],
)
app.launch()
|