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()