File size: 1,655 Bytes
c2f314c 0d29e9f c2f314c 0d29e9f bb347bc c4c4b83 54771e9 0d29e9f 54771e9 dac85aa 0d29e9f dac85aa b314fb6 0d29e9f dac85aa 54771e9 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 51 52 53 54 55 56 |
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))
rng = jax.random.split(rng, jax.device_count())
p_params = replicate(params)
num_samples = 1
prompt_ids = pipeline.prepare_inputs([prompt] * num_samples)
prompt_ids = shard(prompt_ids)
images = pipeline(
prompt_ids=prompt_ids,
params=p_params,
prng_seed=rng,
height=128,
width=128,
num_inference_steps=int(inference_steps),
jit=True,
).images
images = images.reshape((num_samples,) + output.shape[-3:])
images = pipeline.numpy_to_pil(images)
return images
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()
|