File size: 1,746 Bytes
c2f314c 0d29e9f 680c3bf 0d29e9f c2f314c 0d29e9f bb347bc c4c4b83 54771e9 a9057b2 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 57 58 |
import gradio as gr
import jax
from flax.jax_utils import replicate
from flax.training.common_utils import shard
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(pipeline_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()
|