File size: 14,686 Bytes
c2f314c
0d29e9f
90816e9
1f33c8b
680c3bf
 
b7d5b48
c2f314c
3d8eafa
 
0d29e9f
9b38713
3d8eafa
0d29e9f
3d8eafa
 
 
 
 
 
 
 
 
 
 
0d29e9f
 
fedc156
c4c4b83
54771e9
a9057b2
c0c58ec
54771e9
 
 
3d8eafa
e1ad05d
fedc156
e1ad05d
 
 
 
 
 
fedc156
e1ad05d
 
fedc156
 
 
e1ad05d
 
 
 
 
 
 
 
 
fedc156
e1ad05d
 
dac85aa
5be066f
54771e9
651e0b3
fedc156
418dc16
a452886
 
fedc156
b7d5b48
 
 
72d3b5f
 
 
b7d5b48
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
42bdbc7
 
b7d5b48
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fdceb6c
b7d5b48
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72d3b5f
b7d5b48
 
 
 
 
 
 
 
 
 
 
 
 
72d3b5f
 
 
b7d5b48
 
 
 
 
72d3b5f
b7d5b48
0d29e9f
72d3b5f
0d29e9f
2f7b9b1
b7d5b48
 
 
 
 
 
 
 
 
 
 
fdceb6c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b7d5b48
 
 
 
 
 
 
 
 
 
 
 
fdceb6c
72d3b5f
b7d5b48
 
 
72d3b5f
b7d5b48
72d3b5f
b7d5b48
 
 
 
72d3b5f
b7d5b48
72d3b5f
b7d5b48
72d3b5f
b7d5b48
6046757
 
f407d4d
b7d5b48
 
72d3b5f
 
 
 
 
 
 
 
 
 
 
 
 
 
9aa9608
72d3b5f
d59aa40
72d3b5f
 
 
 
b7d5b48
fedc156
 
 
 
 
 
 
b7d5b48
 
 
 
 
 
72d3b5f
b7d5b48
 
fedc156
 
b7d5b48
fedc156
b7d5b48
72d3b5f
b7d5b48
 
 
a452886
fedc156
72d3b5f
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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
import gradio as gr
import jax
import jax.numpy as jnp
from diffusers import FlaxPNDMScheduler, FlaxStableDiffusionPipeline
from flax.jax_utils import replicate
from flax.training.common_utils import shard
from share_btn import community_icon_html, loading_icon_html, share_js

DTYPE = jnp.float16

pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
    "bguisard/stable-diffusion-nano-2-1",
    dtype=DTYPE,
)
if DTYPE != jnp.float32:
    # There is a known issue with schedulers when loading from a pre trained
    # pipeline. We need the schedulers to always use float32.
    # See: https://github.com/huggingface/diffusers/issues/2155
    scheduler, scheduler_params = FlaxPNDMScheduler.from_pretrained(
        pretrained_model_name_or_path="bguisard/stable-diffusion-nano-2-1",
        subfolder="scheduler",
        dtype=jnp.float32,
    )
    pipeline_params["scheduler"] = scheduler_params
    pipeline.scheduler = scheduler


def generate_image(prompt: str, negative_prompt: str = "", inference_steps: int = 25, prng_seed: int = 0, guidance_scale: float = 9):
    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)

    if negative_prompt == "":
        images = pipeline(
            prompt_ids=prompt_ids,
            params=p_params,
            prng_seed=rng,
            height=128,
            width=128,
            num_inference_steps=int(inference_steps),
            guidance_scale=float(guidance_scale),
            jit=True,
        ).images
    else:
        neg_prompt_ids = pipeline.prepare_inputs(
            [negative_prompt] * num_samples)
        neg_prompt_ids = shard(neg_prompt_ids)
        images = pipeline(
            prompt_ids=prompt_ids,
            params=p_params,
            prng_seed=rng,
            height=128,
            width=128,
            num_inference_steps=int(inference_steps),
            neg_prompt_ids=neg_prompt_ids,
            guidance_scale=float(guidance_scale),
            jit=True,
        ).images

    images = images.reshape((num_samples,) + images.shape[-3:])
    images = pipeline.numpy_to_pil(images)
    return images

examples = [
    ["A watercolor painting of a bird"],
    ["A watercolor painting of an otter"]
]
css = """
        .gradio-container {
            font-family: 'IBM Plex Sans', sans-serif;
            max-width: 730px!important;
            margin: auto;
            padding-top: 1.5rem;
        }
        .gr-button {
            color: white;
            border-color: black;
            background: black;
        }
        input[type='range'] {
            accent-color: black;
        }
        .dark input[type='range'] {
            accent-color: #dfdfdf;
        }
        .container {
            max-width: 730px;
            margin: auto;
            padding-top: 1.5rem;
        }
        #gallery {
            min-height: 22rem;
            margin-bottom: 15px;
            margin-left: auto;
            margin-right: auto;
            border-bottom-right-radius: .5rem !important;
            border-bottom-left-radius: .5rem !important;
        }
        #gallery>div>.h-full {
            min-height: 20rem;
        }
        .details:hover {
            text-decoration: underline;
        }
        .gr-button {
            white-space: nowrap;
        }
        .gr-button:focus {
            border-color: rgb(147 197 253 / var(--tw-border-opacity));
            outline: none;
            box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);
            --tw-border-opacity: 1;
            --tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);
            --tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);
            --tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));
            --tw-ring-opacity: .5;
        }
        #advanced-btn {
            font-size: .7rem !important;
            line-height: 19px;
                    cache_examples=True,
                    postprocess=False)
            margin-top: 12px;
            margin-bottom: 12px;
            padding: 2px 8px;
            border-radius: 14px !important;
        }
        #advanced-options {
            display: none;
            margin-bottom: 20px;
        }
        .footer {
            margin-bottom: 45px;
            margin-top: 35px;
            text-align: center;
            border-bottom: 1px solid #e5e5e5;
        }
        .footer>p {
            font-size: .8rem;
            display: inline-block;
            padding: 0 10px;
            transform: translateY(10px);
            background: white;
        }
        .dark .footer {
            border-color: #303030;
        }
        .dark .footer>p {
            background: #0b0f19;
        }
        .acknowledgments h4{
            margin: 1.25em 0 .25em 0;
            font-weight: bold;
            font-size: 115%;
        }
        .animate-spin {
            animation: spin 1s linear infinite;
        }
        @keyframes spin {
            from {
                transform: rotate(0deg);
            }
            to {
                transform: rotate(360deg);
            }
        }
        #share-btn-container {
            display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
            margin-top: 10px;
            margin-left: auto;
        
        #share-btn {
            all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;right:0;
        }
        #share-btn * {
            all: unset;
        }
        #share-btn-container div:nth-child(-n+2){
            width: auto !important;
            min-height: 0px !important;
        }
        #share-btn-container .wrap {
            display: none !important;
        }
        .share_button {
            color:#6366f1!important;
        }
        .gr-form{
            flex: 1 1 50%; border-top-right-radius: 0; border-bottom-right-radius: 0;
        }
        #prompt-text-input, #negative-prompt-text-input{padding: .45rem 0.625rem}
        .image_duplication{position: absolute; width: 100px; left: 50px}

"""

block = gr.Blocks(theme="gradio/soft",css=css)

with block as demo:
    gr.HTML(
        """
            <div style="text-align: center; margin: 0 auto;">
              <div
                style="
                  display: inline-flex;
                  align-items: center;
                  gap: 0.8rem;
                  font-size: 1.75rem;
                "
              >
              <svg
                  width="0.65em"
                  height="0.65em"
                  viewBox="0 0 115 115"
                  fill="none"
                  xmlns="http://www.w3.org/2000/svg"
                >
                  <rect width="23" height="23" fill="white"></rect>
                  <rect y="69" width="23" height="23" fill="white"></rect>
                  <rect x="23" width="23" height="23" fill="#AEAEAE"></rect>
                  <rect x="23" y="69" width="23" height="23" fill="#AEAEAE"></rect>
                  <rect x="46" width="23" height="23" fill="white"></rect>
                  <rect x="46" y="69" width="23" height="23" fill="white"></rect>
                  <rect x="69" width="23" height="23" fill="black"></rect>
                  <rect x="69" y="69" width="23" height="23" fill="black"></rect>
                  <rect x="92" width="23" height="23" fill="#D9D9D9"></rect>
                  <rect x="92" y="69" width="23" height="23" fill="#AEAEAE"></rect>
                  <rect x="115" y="46" width="23" height="23" fill="white"></rect>
                  <rect x="115" y="115" width="23" height="23" fill="white"></rect>
                  <rect x="115" y="69" width="23" height="23" fill="#D9D9D9"></rect>
                  <rect x="92" y="46" width="23" height="23" fill="#AEAEAE"></rect>
                  <rect x="92" y="115" width="23" height="23" fill="#AEAEAE"></rect>
                  <rect x="92" y="69" width="23" height="23" fill="white"></rect>
                  <rect x="69" y="46" width="23" height="23" fill="white"></rect>
                  <rect x="69" y="115" width="23" height="23" fill="white"></rect>
                  <rect x="69" y="69" width="23" height="23" fill="#D9D9D9"></rect>
                  <rect x="46" y="46" width="23" height="23" fill="black"></rect>
                  <rect x="46" y="115" width="23" height="23" fill="black"></rect>
                  <rect x="46" y="69" width="23" height="23" fill="black"></rect>
                  <rect x="23" y="46" width="23" height="23" fill="#D9D9D9"></rect>
                  <rect x="23" y="115" width="23" height="23" fill="#AEAEAE"></rect>
                  <rect x="23" y="69" width="23" height="23" fill="black"></rect>
                </svg>
                <h1 style="font-weight: 900; margin-bottom: 7px;margin-top:5px">
                  Stable Diffusion Nano Demo
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">
                Stable Diffusion Nano was built during the <a style="text-decoration: underline;" href="https://github.com/huggingface/community-events/tree/main/jax-controlnet-sprint">JAX/Diffusers community sprint 🧨</a> based on Stable Diffusion 2.1 and finetuned on 128x128 images for fast prototyping. <br>
              </p>
            </div>
        """
    )
    with gr.Group():
        with gr.Box():
            with gr.Row(elem_id="prompt-container").style(equal_height=True):
                with gr.Column(scale=2):
                    prompt_input = gr.Textbox(
                        label="Enter your prompt",
                        max_lines=1,
                        placeholder="Enter your prompt",
                        elem_id="prompt-text-input",
                        show_label=False,
                    )
                    negative = gr.Textbox(
                        label="Enter your negative prompt",
                        max_lines=1,
                        placeholder="Enter a negative prompt",
                        elem_id="negative-prompt-text-input",
                        show_label=False,
                    )
                btn = gr.Button("Generate image", label="Primary Button", variant="primary")

        gallery = gr.Image(
            label="Generated images", show_label=False, elem_id="gallery",type="pil"
        ).style(height="auto")


        with gr.Row():
            with gr.Column(scale=2):
                with gr.Accordion("Advanced settings"):
                    seed_input = gr.inputs.Number(default=0, label="Seed")
                    inf_steps_input = gr.inputs.Slider(
                        minimum=1, maximum=100, default=25, step=1, label="Inference Steps"
                    )
                    guidance_scale = gr.inputs.Slider(
                        label="Guidance Scale", minimum=0, maximum=50, default=9, step=0.1
                    )
            with gr.Column(scale=1):
                # advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
                ex = gr.Examples(examples=examples,
                    fn=generate_image,
                    inputs=[prompt_input, negative,inf_steps_input, seed_input, guidance_scale],
                    outputs=[gallery],
                    cache_examples=False)
                ex.dataset.headers = [""]

                share_button = gr.Button("Share to community",elem_classes="share_button")
        

        negative.submit(generate_image, inputs=[
                        prompt_input, negative, inf_steps_input, seed_input, guidance_scale], outputs=[gallery], postprocess=False)
        prompt_input.submit(generate_image, inputs=[
                            prompt_input, negative, inf_steps_input, seed_input, guidance_scale], outputs=[gallery], postprocess=False)
        btn.click(generate_image, inputs=[prompt_input, negative, inf_steps_input,
                  seed_input, guidance_scale], outputs=[gallery], postprocess=False)

        share_button.click(
            None,
            [],
            [],
            _js=share_js,
        )
        gr.Markdown("Model by Stable Diffusion Nano Team",elem_classes="footer")
        with gr.Accordion(label="License", open=False):
            gr.HTML(
                """
                <div class="acknowledgments">
                    <p><h4>LICENSE</h4>
                    The model is licensed with a <a href="https://huggingface.co/stabilityai/stable-diffusion-2/blob/main/LICENSE-MODEL" style="text-decoration: underline;" target="_blank">CreativeML OpenRAIL++</a> license. The authors claim no rights on the outputs you generate, you are free to use them and are accountable for their use which must not go against the provisions set in this license. The license forbids you from sharing any content that violates any laws, produce any harm to a person, disseminate any personal information that would be meant for harm, spread misinformation and target vulnerable groups. For the full list of restrictions please <a href="https://huggingface.co/spaces/CompVis/stable-diffusion-license" target="_blank" style="text-decoration: underline;" target="_blank">read the license</a></p>
                    <p><h4>Biases and content acknowledgment</h4>
                    Despite how impressive being able to turn text into image is, beware to the fact that this model may output content that reinforces or exacerbates societal biases, as well as realistic faces, pornography and violence. The model was trained on the <a href="https://huggingface.co/datasets/laion/laion2B-en-aesthetic" style="text-decoration: underline;" target="_blank">LAION-2B Aesthetic dataset</a>, which scraped non-curated image-text-pairs from the internet (the exception being the removal of illegal content) and is meant for research purposes. You can read more in the <a href="https://huggingface.co/bguisard/stable-diffusion-nano-2-1" style="text-decoration: underline;" target="_blank">model card</a></p>
               </div>
                """
            )
demo.queue(concurrency_count=10)
demo.launch()