import gradio as gr import jax from diffusers import FlaxStableDiffusionPipeline from flax.jax_utils import replicate from flax.training.common_utils import shard 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,) + images.shape[-3:]) images = pipeline.numpy_to_pil(images) return images[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="image", 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()