ui.py 32.4 KB
Newer Older
1 2 3 4
import base64
import html
import io
import json
A
AUTOMATIC 已提交
5
import math
6 7
import mimetypes
import os
A
AUTOMATIC 已提交
8
import random
9 10 11 12
import sys
import time
import traceback

A
AUTOMATIC 已提交
13 14
import numpy as np
import torch
15 16 17 18
from PIL import Image

import gradio as gr
import gradio.utils
A
AUTOMATIC 已提交
19
import gradio.routes
20 21 22 23 24 25

from modules.paths import script_path
from modules.shared import opts, cmd_opts
import modules.shared as shared
from modules.sd_samplers import samplers, samplers_for_img2img
import modules.realesrgan_model as realesrgan
A
AUTOMATIC 已提交
26
import modules.scripts
27 28
import modules.gfpgan_model
import modules.codeformer_model
A
AUTOMATIC 已提交
29
import modules.styles
30 31 32 33 34 35

# this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the bowser will not show any UI
mimetypes.init()
mimetypes.add_type('application/javascript', '.js')


36
if not cmd_opts.share and not cmd_opts.listen:
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
    # fix gradio phoning home
    gradio.utils.version_check = lambda: None
    gradio.utils.get_local_ip_address = lambda: '127.0.0.1'


def gr_show(visible=True):
    return {"visible": visible, "__type__": "update"}


sample_img2img = "assets/stable-samples/img2img/sketch-mountains-input.jpg"
sample_img2img = sample_img2img if os.path.exists(sample_img2img) else None

css_hide_progressbar = """
.wrap .m-12 svg { display:none!important; }
.wrap .m-12::before { content:"Loading..." }
.progress-bar { display:none!important; }
.meta-text { display:none!important; }
"""

def plaintext_to_html(text):
57
    text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"
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
    return text


def image_from_url_text(filedata):
    if type(filedata) == list:
        if len(filedata) == 0:
            return None

        filedata = filedata[0]

    if filedata.startswith("data:image/png;base64,"):
        filedata = filedata[len("data:image/png;base64,"):]

    filedata = base64.decodebytes(filedata.encode('utf-8'))
    image = Image.open(io.BytesIO(filedata))
    return image


def send_gradio_gallery_to_image(x):
    if len(x) == 0:
        return None

    return image_from_url_text(x[0])


def save_files(js_data, images):
    import csv

    os.makedirs(opts.outdir_save, exist_ok=True)

    filenames = []

    data = json.loads(js_data)

92
    with open(os.path.join(opts.outdir_save, "log.csv"), "a", encoding="utf8", newline='') as file:
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
        at_start = file.tell() == 0
        writer = csv.writer(file)
        if at_start:
            writer.writerow(["prompt", "seed", "width", "height", "sampler", "cfgs", "steps", "filename"])

        filename_base = str(int(time.time() * 1000))
        for i, filedata in enumerate(images):
            filename = filename_base + ("" if len(images) == 1 else "-" + str(i + 1)) + ".png"
            filepath = os.path.join(opts.outdir_save, filename)

            if filedata.startswith("data:image/png;base64,"):
                filedata = filedata[len("data:image/png;base64,"):]

            with open(filepath, "wb") as imgfile:
                imgfile.write(base64.decodebytes(filedata.encode('utf-8')))

            filenames.append(filename)

        writer.writerow([data["prompt"], data["seed"], data["width"], data["height"], data["sampler"], data["cfg_scale"], data["steps"], filenames[0]])

    return '', '', plaintext_to_html(f"Saved: {filenames[0]}")


def wrap_gradio_call(func):
    def f(*args, **kwargs):
        t = time.perf_counter()

        try:
            res = list(func(*args, **kwargs))
        except Exception as e:
            print("Error completing request", file=sys.stderr)
            print("Arguments:", args, kwargs, file=sys.stderr)
            print(traceback.format_exc(), file=sys.stderr)

A
AUTOMATIC 已提交
127 128 129
            shared.state.job = ""
            shared.state.job_count = 0

130 131 132 133 134 135 136 137 138 139 140 141 142 143
            res = [None, '', f"<div class='error'>{plaintext_to_html(type(e).__name__+': '+str(e))}</div>"]

        elapsed = time.perf_counter() - t

        # last item is always HTML
        res[-1] = res[-1] + f"<p class='performance'>Time taken: {elapsed:.2f}s</p>"

        shared.state.interrupted = False

        return tuple(res)

    return f


A
AUTOMATIC 已提交
144 145 146
def check_progress_call():

    if shared.state.job_count == 0:
A
AUTOMATIC 已提交
147
        return "", gr_show(False), gr_show(False)
A
AUTOMATIC 已提交
148

A
AUTOMATIC 已提交
149 150 151 152
    progress = 0

    if shared.state.job_count > 0:
        progress += shared.state.job_no / shared.state.job_count
A
AUTOMATIC 已提交
153 154 155 156 157
    if shared.state.sampling_steps > 0:
        progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps

    progress = min(progress, 1)

A
AUTOMATIC 已提交
158 159 160 161 162 163 164 165
    progressbar = ""
    if opts.show_progressbar:
        progressbar = f"""<div class='progressDiv'><div class='progress' style="width:{progress * 100}%">{str(int(progress*100))+"%" if progress > 0.01 else ""}</div></div>"""

    image = gr_show(False)
    preview_visibility = gr_show(False)

    if opts.show_progress_every_n_steps > 0:
166
        if shared.parallel_processing_allowed:
A
AUTOMATIC 已提交
167

168 169 170
            if shared.state.sampling_step - shared.state.current_image_sampling_step >= opts.show_progress_every_n_steps and shared.state.current_latent is not None:
                shared.state.current_image = modules.sd_samplers.sample_to_image(shared.state.current_latent)
                shared.state.current_image_sampling_step = shared.state.sampling_step
A
AUTOMATIC 已提交
171

A
AUTOMATIC 已提交
172 173 174 175 176 177
        image = shared.state.current_image

        if image is None or progress >= 1:
            image = gr.update(value=None)
        else:
            preview_visibility = gr_show(True)
A
AUTOMATIC 已提交
178

A
AUTOMATIC 已提交
179
    return f"<span style='display: none'>{time.time()}</span><p>{progressbar}</p>", preview_visibility, image
A
AUTOMATIC 已提交
180 181


A
AUTOMATIC 已提交
182 183 184 185 186 187 188
def roll_artist(prompt):
    allowed_cats = set([x for x in shared.artist_db.categories() if len(opts.random_artist_categories)==0 or x in opts.random_artist_categories])
    artist = random.choice([x for x in shared.artist_db.artists if x.category in allowed_cats])

    return prompt + ", " + artist.name if prompt != '' else artist.name


A
AUTOMATIC 已提交
189 190 191 192 193 194 195
def visit(x, func, path=""):
    if hasattr(x, 'children'):
        for c in x.children:
            visit(c, func, path)
    elif x.label is not None:
        func(path + "/" + str(x.label), x)

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
def create_seed_inputs():
    with gr.Row():
        seed = gr.Number(label='Seed', value=-1)
        subseed = gr.Number(label='Variation seed', value=-1, visible=False)
        seed_checkbox = gr.Checkbox(label="Extra", elem_id="subseed_show", value=False)

    with gr.Row():
        subseed_strength = gr.Slider(label='Variation strength', value=0.0, minimum=0, maximum=1, step=0.01, visible=False)
        seed_resize_from_h = gr.Slider(minimum=0, maximum=2048, step=64, label="Resize seed from height", value=0, visible=False)
        seed_resize_from_w = gr.Slider(minimum=0, maximum=2048, step=64, label="Resize seed from width", value=0, visible=False)

    def change_visiblity(show):

        return {
            subseed: gr_show(show),
            subseed_strength: gr_show(show),
            seed_resize_from_h: gr_show(show),
            seed_resize_from_w: gr_show(show),
        }

    seed_checkbox.change(
        change_visiblity,
        inputs=[seed_checkbox],
        outputs=[
            subseed,
            subseed_strength,
            seed_resize_from_h,
            seed_resize_from_w
        ]
    )

    return seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w


A
AUTOMATIC 已提交
231 232 233 234 235 236 237 238 239 240 241 242 243 244
def add_style(style_name, text):
    if style_name is None:
        return [gr_show(), gr_show()]

    style = modules.styles.PromptStyle(style_name, text)

    modules.styles.save_style(shared.styles_filename, style)

    shared.prompt_styles[style.name] = style

    update = {"visible": True, "choices": [k for k, v in shared.prompt_styles.items()], "__type__": "update"}
    return [update, update]


A
AUTOMATIC 已提交
245
def create_ui(txt2img, img2img, run_extras, run_pnginfo):
246 247
    with gr.Blocks(analytics_enabled=False) as txt2img_interface:
        with gr.Row():
A
AUTOMATIC 已提交
248
            txt2img_prompt = gr.Textbox(label="Prompt", elem_id="txt2img_prompt", show_label=False, placeholder="Prompt", lines=1)
A
AUTOMATIC 已提交
249
            negative_prompt = gr.Textbox(label="Negative prompt", elem_id="txt2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1)
A
AUTOMATIC 已提交
250
            txt2img_prompt_style = gr.Dropdown(label="Style", show_label=False, elem_id="style_index", choices=[k for k, v in shared.prompt_styles.items()], value=next(iter(shared.prompt_styles.keys())), visible=len(shared.prompt_styles) > 1)
A
AUTOMATIC 已提交
251
            roll = gr.Button('Roll', elem_id="txt2img_roll", visible=len(shared.artist_db.artists) > 0)
252
            submit = gr.Button('Generate', elem_id="txt2img_generate", variant='primary')
A
AUTOMATIC 已提交
253
            check_progress = gr.Button('Check progress', elem_id="check_progress", visible=False)
254 255 256 257 258 259 260

        with gr.Row().style(equal_height=False):
            with gr.Column(variant='panel'):
                steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20)
                sampler_index = gr.Radio(label='Sampling method', elem_id="txt2img_sampling", choices=[x.name for x in samplers], value=samplers[0].name, type="index")

                with gr.Row():
A
AUTOMATIC 已提交
261
                    restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)
262
                    tiling = gr.Checkbox(label='Tiling', value=False)
263 264 265 266 267 268 269 270 271 272 273

                with gr.Row():
                    batch_count = gr.Slider(minimum=1, maximum=cmd_opts.max_batch_count, step=1, label='Batch count', value=1)
                    batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1)

                cfg_scale = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='CFG Scale', value=7.0)

                with gr.Group():
                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
                    width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)

274
                seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w = create_seed_inputs()
275

A
AUTOMATIC 已提交
276
                with gr.Group():
A
AUTOMATIC 已提交
277
                    custom_inputs = modules.scripts.scripts_txt2img.setup_ui(is_img2img=False)
278 279 280

            with gr.Column(variant='panel'):
                with gr.Group():
A
AUTOMATIC 已提交
281
                    txt2img_preview = gr.Image(elem_id='txt2img_preview', visible=False)
282
                    txt2img_gallery = gr.Gallery(label='Output', elem_id='txt2img_gallery').style(grid=4)
283

A
AUTOMATIC 已提交
284

285 286 287 288 289 290 291
                with gr.Group():
                    with gr.Row():
                        save = gr.Button('Save')
                        send_to_img2img = gr.Button('Send to img2img')
                        send_to_inpaint = gr.Button('Send to inpaint')
                        send_to_extras = gr.Button('Send to extras')
                        interrupt = gr.Button('Interrupt')
A
AUTOMATIC 已提交
292
                        txt2img_save_style = gr.Button('Save prompt as style')
293

A
AUTOMATIC 已提交
294 295
                progressbar = gr.HTML(elem_id="progressbar")

296 297 298 299
                with gr.Group():
                    html_info = gr.HTML()
                    generation_info = gr.Textbox(visible=False)

A
AUTOMATIC 已提交
300

301 302
            txt2img_args = dict(
                fn=txt2img,
A
AUTOMATIC 已提交
303
                _js="submit",
304
                inputs=[
A
AUTOMATIC 已提交
305
                    txt2img_prompt,
306
                    negative_prompt,
A
AUTOMATIC 已提交
307
                    txt2img_prompt_style,
308 309
                    steps,
                    sampler_index,
A
AUTOMATIC 已提交
310
                    restore_faces,
311
                    tiling,
312 313 314 315
                    batch_count,
                    batch_size,
                    cfg_scale,
                    seed,
316
                    subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
317 318
                    height,
                    width,
A
AUTOMATIC 已提交
319
                ] + custom_inputs,
320 321 322 323 324 325 326
                outputs=[
                    txt2img_gallery,
                    generation_info,
                    html_info
                ]
            )

A
AUTOMATIC 已提交
327
            txt2img_prompt.submit(**txt2img_args)
328 329
            submit.click(**txt2img_args)

A
AUTOMATIC 已提交
330 331
            check_progress.click(
                fn=check_progress_call,
A
AUTOMATIC 已提交
332
                show_progress=False,
A
AUTOMATIC 已提交
333
                inputs=[],
A
AUTOMATIC 已提交
334
                outputs=[progressbar, txt2img_preview, txt2img_preview],
A
AUTOMATIC 已提交
335 336 337
            )


338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356
            interrupt.click(
                fn=lambda: shared.state.interrupt(),
                inputs=[],
                outputs=[],
            )

            save.click(
                fn=wrap_gradio_call(save_files),
                inputs=[
                    generation_info,
                    txt2img_gallery,
                ],
                outputs=[
                    html_info,
                    html_info,
                    html_info,
                ]
            )

A
AUTOMATIC 已提交
357 358 359
            roll.click(
                fn=roll_artist,
                inputs=[
A
AUTOMATIC 已提交
360
                    txt2img_prompt,
A
AUTOMATIC 已提交
361 362
                ],
                outputs=[
A
AUTOMATIC 已提交
363
                    txt2img_prompt,
A
AUTOMATIC 已提交
364 365 366
                ]
            )

367 368
    with gr.Blocks(analytics_enabled=False) as img2img_interface:
        with gr.Row():
A
AUTOMATIC 已提交
369
            img2img_prompt = gr.Textbox(label="Prompt", elem_id="img2img_prompt", show_label=False, placeholder="Prompt", lines=1)
A
AUTOMATIC 已提交
370
            negative_prompt = gr.Textbox(label="Negative prompt", elem_id="img2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1)
A
AUTOMATIC 已提交
371
            img2img_prompt_style = gr.Dropdown(label="Style", show_label=False, elem_id="style_index", choices=[k for k, v in shared.prompt_styles.items()], value=next(iter(shared.prompt_styles.keys())), visible=len(shared.prompt_styles) > 1)
372
            submit = gr.Button('Generate', elem_id="img2img_generate", variant='primary')
A
AUTOMATIC 已提交
373
            check_progress = gr.Button('Check progress', elem_id="check_progress", visible=False)
374 375 376 377 378 379

        with gr.Row().style(equal_height=False):
            with gr.Column(variant='panel'):
                with gr.Group():
                    switch_mode = gr.Radio(label='Mode', elem_id="img2img_mode", choices=['Redraw whole image', 'Inpaint a part of image', 'Loopback', 'SD upscale'], value='Redraw whole image', type="index", show_label=False)
                    init_img = gr.Image(label="Image for img2img", source="upload", interactive=True, type="pil")
380
                    init_img_with_mask = gr.Image(label="Image for inpainting with mask", elem_id="img2maskimg", source="upload", interactive=True, type="pil", tool="sketch", visible=False, image_mode="RGBA")
A
AUTOMATIC 已提交
381
                    init_img_with_mask_comment = gr.HTML(elem_id="mask_bug_info", value="<small>if the editor shows ERROR, switch to another img2img mode above and back</small>", visible=False)
382 383 384 385 386
                    init_mask = gr.Image(label="Mask", source="upload", interactive=True, type="pil", visible=False)

                    with gr.Row():
                        resize_mode = gr.Radio(label="Resize mode", elem_id="resize_mode", show_label=False, choices=["Just resize", "Crop and resize", "Resize and fill"], type="index", value="Just resize")
                        mask_mode = gr.Radio(label="Mask mode", show_label=False, choices=["Draw mask", "Upload mask"], type="index", value="Draw mask")
387 388 389 390

                steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20)
                sampler_index = gr.Radio(label='Sampling method', choices=[x.name for x in samplers_for_img2img], value=samplers_for_img2img[0].name, type="index")
                mask_blur = gr.Slider(label='Mask blur', minimum=0, maximum=64, step=1, value=4, visible=False)
A
AUTOMATIC 已提交
391
                inpainting_fill = gr.Radio(label='Masked content', choices=['fill', 'original', 'latent noise', 'latent nothing'], value='fill', type="index", visible=False)
392 393

                with gr.Row():
A
AUTOMATIC 已提交
394
                    inpaint_full_res = gr.Checkbox(label='Inpaint at full resolution', value=False, visible=False)
A
AUTOMATIC 已提交
395 396 397
                    inpainting_mask_invert = gr.Radio(label='Masking mode', choices=['Inpaint masked', 'Inpaint not masked'], value='Inpaint masked', type="index", visible=False)

                with gr.Row():
A
AUTOMATIC 已提交
398
                    restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)
399
                    tiling = gr.Checkbox(label='Tiling', value=False)
A
AUTOMATIC 已提交
400
                    sd_upscale_overlap = gr.Slider(minimum=0, maximum=256, step=16, label='Tile overlap', value=64, visible=False)
401 402

                with gr.Row():
A
AUTOMATIC 已提交
403
                    sd_upscale_upscaler_name = gr.Radio(label='Upscaler', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index", visible=False)
404 405 406 407 408 409 410

                with gr.Row():
                    batch_count = gr.Slider(minimum=1, maximum=cmd_opts.max_batch_count, step=1, label='Batch count', value=1)
                    batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1)

                with gr.Group():
                    cfg_scale = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='CFG Scale', value=7.0)
411 412
                    denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.75)
                    denoising_strength_change_factor = gr.Slider(minimum=0.9, maximum=1.1, step=0.01, label='Denoising strength change factor', value=1, visible=False)
413 414 415 416 417

                with gr.Group():
                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
                    width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)

418
                seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w = create_seed_inputs()
419

A
AUTOMATIC 已提交
420
                with gr.Group():
A
AUTOMATIC 已提交
421
                    custom_inputs = modules.scripts.scripts_img2img.setup_ui(is_img2img=True)
A
AUTOMATIC 已提交
422

423 424
            with gr.Column(variant='panel'):
                with gr.Group():
A
AUTOMATIC 已提交
425
                    img2img_preview = gr.Image(elem_id='img2img_preview', visible=False)
426
                    img2img_gallery = gr.Gallery(label='Output', elem_id='img2img_gallery').style(grid=4)
427 428 429 430

                with gr.Group():
                    with gr.Row():
                        save = gr.Button('Save')
A
AUTOMATIC 已提交
431 432
                        img2img_send_to_img2img = gr.Button('Send to img2img')
                        img2img_send_to_inpaint = gr.Button('Send to inpaint')
433
                        img2img_send_to_extras = gr.Button('Send to extras')
A
AUTOMATIC 已提交
434
                        interrupt = gr.Button('Interrupt')
A
AUTOMATIC 已提交
435
                        img2img_save_style = gr.Button('Save prompt as style')
436

A
AUTOMATIC 已提交
437
                progressbar = gr.HTML(elem_id="progressbar")
A
AUTOMATIC 已提交
438
                style_dummpy = gr.Textbox(visible=False)
A
AUTOMATIC 已提交
439

440 441 442 443
                with gr.Group():
                    html_info = gr.HTML()
                    generation_info = gr.Textbox(visible=False)

444
            def apply_mode(mode, uploadmask):
445 446 447 448 449 450
                is_classic = mode == 0
                is_inpaint = mode == 1
                is_loopback = mode == 2
                is_upscale = mode == 3

                return {
451 452
                    init_img: gr_show(not is_inpaint or (is_inpaint and uploadmask == 1)),
                    init_img_with_mask: gr_show(is_inpaint and uploadmask == 0),
A
AUTOMATIC 已提交
453
                    init_img_with_mask_comment: gr_show(is_inpaint and uploadmask == 0),
454 455
                    init_mask: gr_show(is_inpaint and uploadmask == 1),
                    mask_mode: gr_show(is_inpaint),
456 457 458 459
                    mask_blur: gr_show(is_inpaint),
                    inpainting_fill: gr_show(is_inpaint),
                    batch_size: gr_show(not is_loopback),
                    sd_upscale_upscaler_name: gr_show(is_upscale),
A
AUTOMATIC 已提交
460
                    sd_upscale_overlap: gr_show(is_upscale),
461
                    inpaint_full_res: gr_show(is_inpaint),
A
AUTOMATIC 已提交
462
                    inpainting_mask_invert: gr_show(is_inpaint),
463
                    denoising_strength_change_factor: gr_show(is_loopback),
464 465 466 467
                }

            switch_mode.change(
                apply_mode,
468
                inputs=[switch_mode, mask_mode],
469 470 471
                outputs=[
                    init_img,
                    init_img_with_mask,
A
AUTOMATIC 已提交
472
                    init_img_with_mask_comment,
473 474
                    init_mask,
                    mask_mode,
475 476 477 478 479 480
                    mask_blur,
                    inpainting_fill,
                    batch_size,
                    sd_upscale_upscaler_name,
                    sd_upscale_overlap,
                    inpaint_full_res,
A
AUTOMATIC 已提交
481
                    inpainting_mask_invert,
482
                    denoising_strength_change_factor,
483 484 485
                ]
            )

486 487 488 489 490 491 492 493 494 495 496 497 498 499
            mask_mode.change(
                lambda mode: {
                    init_img: gr_show(mode == 1),
                    init_img_with_mask: gr_show(mode == 0),
                    init_mask: gr_show(mode == 1),
                },
                inputs=[mask_mode],
                outputs=[
                    init_img,
                    init_img_with_mask,
                    init_mask,
                ],
            )

500 501
            img2img_args = dict(
                fn=img2img,
A
AUTOMATIC 已提交
502
                _js="submit",
503
                inputs=[
A
AUTOMATIC 已提交
504
                    img2img_prompt,
505
                    negative_prompt,
A
AUTOMATIC 已提交
506
                    img2img_prompt_style,
507 508
                    init_img,
                    init_img_with_mask,
509 510
                    init_mask,
                    mask_mode,
511 512 513 514
                    steps,
                    sampler_index,
                    mask_blur,
                    inpainting_fill,
A
AUTOMATIC 已提交
515
                    restore_faces,
516
                    tiling,
517 518 519 520 521
                    switch_mode,
                    batch_count,
                    batch_size,
                    cfg_scale,
                    denoising_strength,
522
                    denoising_strength_change_factor,
523
                    seed,
524
                    subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
525 526 527 528 529 530
                    height,
                    width,
                    resize_mode,
                    sd_upscale_upscaler_name,
                    sd_upscale_overlap,
                    inpaint_full_res,
A
AUTOMATIC 已提交
531
                    inpainting_mask_invert,
A
AUTOMATIC 已提交
532
                ] + custom_inputs,
533 534 535 536 537 538 539
                outputs=[
                    img2img_gallery,
                    generation_info,
                    html_info
                ]
            )

A
AUTOMATIC 已提交
540
            img2img_prompt.submit(**img2img_args)
541 542
            submit.click(**img2img_args)

A
AUTOMATIC 已提交
543 544
            check_progress.click(
                fn=check_progress_call,
A
AUTOMATIC 已提交
545
                show_progress=False,
A
AUTOMATIC 已提交
546
                inputs=[],
A
AUTOMATIC 已提交
547
                outputs=[progressbar, img2img_preview, img2img_preview],
A
AUTOMATIC 已提交
548 549
            )

550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568
            interrupt.click(
                fn=lambda: shared.state.interrupt(),
                inputs=[],
                outputs=[],
            )

            save.click(
                fn=wrap_gradio_call(save_files),
                inputs=[
                    generation_info,
                    img2img_gallery,
                ],
                outputs=[
                    html_info,
                    html_info,
                    html_info,
                ]
            )

A
AUTOMATIC 已提交
569 570 571 572 573 574 575 576
            for button, propmt in zip([txt2img_save_style, img2img_save_style], [txt2img_prompt, img2img_prompt]):
                button.click(
                    fn=add_style,
                    _js="ask_for_style_name",
                    inputs=[style_dummpy, propmt],
                    outputs=[txt2img_prompt_style, img2img_prompt_style],
                )

577 578 579 580 581
    with gr.Blocks(analytics_enabled=False) as extras_interface:
        with gr.Row().style(equal_height=False):
            with gr.Column(variant='panel'):
                with gr.Group():
                    image = gr.Image(label="Source", source="upload", interactive=True, type="pil")
A
AUTOMATIC 已提交
582 583 584 585 586 587 588 589 590 591 592

                upscaling_resize = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Resize", value=2)

                with gr.Group():
                    extras_upscaler_1 = gr.Radio(label='Upscaler 1', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index")

                with gr.Group():
                    extras_upscaler_2 = gr.Radio(label='Upscaler 2', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index")
                    extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=1)

                with gr.Group():
593 594 595 596
                    gfpgan_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="GFPGAN visibility", value=0, interactive=modules.gfpgan_model.have_gfpgan)

                with gr.Group():
                    codeformer_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer visibility", value=0, interactive=modules.codeformer_model.have_codeformer)
597
                    codeformer_weight = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer weight (0 = maximum effect, 1 = minimum effect)", value=0, interactive=modules.codeformer_model.have_codeformer)
598 599 600 601 602 603 604 605 606 607 608 609

                submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')

            with gr.Column(variant='panel'):
                result_image = gr.Image(label="Result")
                html_info_x = gr.HTML()
                html_info = gr.HTML()

        extras_args = dict(
            fn=run_extras,
            inputs=[
                image,
610 611 612
                gfpgan_visibility,
                codeformer_visibility,
                codeformer_weight,
A
AUTOMATIC 已提交
613 614 615 616
                upscaling_resize,
                extras_upscaler_1,
                extras_upscaler_2,
                extras_upscaler_2_visibility,
617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648
            ],
            outputs=[
                result_image,
                html_info_x,
                html_info,
            ]
        )

        submit.click(**extras_args)

    pnginfo_interface = gr.Interface(
        wrap_gradio_call(run_pnginfo),
        inputs=[
            gr.Image(label="Source", source="upload", interactive=True, type="pil"),
        ],
        outputs=[
            gr.HTML(),
            gr.HTML(),
            gr.HTML(),
        ],
        allow_flagging="never",
        analytics_enabled=False,
    )

    def create_setting_component(key):
        def fun():
            return opts.data[key] if key in opts.data else opts.data_labels[key].default

        info = opts.data_labels[key]
        t = type(info.default)

        if info.component is not None:
A
AUTOMATIC 已提交
649 650
            args = info.component_args() if callable(info.component_args) else info.component_args
            item = info.component(label=info.label, value=fun, **(args or {}))
651 652 653 654 655 656 657 658 659 660 661
        elif t == str:
            item = gr.Textbox(label=info.label, value=fun, lines=1)
        elif t == int:
            item = gr.Number(label=info.label, value=fun)
        elif t == bool:
            item = gr.Checkbox(label=info.label, value=fun)
        else:
            raise Exception(f'bad options item type: {str(t)} for key {key}')

        return item

A
AUTOMATIC 已提交
662 663 664 665 666
    components = []
    keys = list(opts.data_labels.keys())
    settings_cols = 3
    items_per_col = math.ceil(len(keys) / settings_cols)

667 668 669
    def run_settings(*args):
        up = []

A
AUTOMATIC 已提交
670
        for key, value, comp in zip(opts.data_labels.keys(), args, components):
671 672 673 674 675
            opts.data[key] = value
            up.append(comp.update(value=value))

        opts.save(shared.config_filename)

A
AUTOMATIC 已提交
676
        return 'Settings applied.'
677

A
AUTOMATIC 已提交
678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695
    with gr.Blocks(analytics_enabled=False) as settings_interface:
        submit = gr.Button(value="Apply settings", variant='primary')
        result = gr.HTML()

        with gr.Row(elem_id="settings").style(equal_height=False):
            for colno in range(settings_cols):
                with gr.Column(variant='panel'):
                    for rowno in range(items_per_col):
                        index = rowno + colno * items_per_col

                        if index < len(keys):
                            components.append(create_setting_component(keys[index]))

        submit.click(
            fn=run_settings,
            inputs=components,
            outputs=[result]
        )
696 697

    interfaces = [
A
AUTOMATIC 已提交
698 699 700 701 702
        (txt2img_interface, "txt2img", "txt2img"),
        (img2img_interface, "img2img", "img2img"),
        (extras_interface, "Extras", "extras"),
        (pnginfo_interface, "PNG Info", "pnginfo"),
        (settings_interface, "Settings", "settings"),
703 704 705 706 707 708 709 710
    ]

    with open(os.path.join(script_path, "style.css"), "r", encoding="utf8") as file:
        css = file.read()

    if not cmd_opts.no_progressbar_hiding:
        css += css_hide_progressbar

A
AUTOMATIC 已提交
711 712 713 714 715 716 717 718 719 720 721 722 723 724
    with gr.Blocks(css=css, analytics_enabled=False, title="Stable Diffusion") as demo:
        with gr.Tabs() as tabs:
            for interface, label, ifid in interfaces:
                with gr.TabItem(label, id=ifid):
                    interface.render()

        tabs.change(
            fn=lambda x: x,
            inputs=[init_img_with_mask],
            outputs=[init_img_with_mask],
        )

        send_to_img2img.click(
            fn=lambda x: image_from_url_text(x),
S
Seki 已提交
725
            _js="extract_image_from_gallery_img2img",
A
AUTOMATIC 已提交
726 727 728 729 730 731
            inputs=[txt2img_gallery],
            outputs=[init_img],
        )

        send_to_inpaint.click(
            fn=lambda x: image_from_url_text(x),
S
Seki 已提交
732
            _js="extract_image_from_gallery_img2img",
A
AUTOMATIC 已提交
733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752
            inputs=[txt2img_gallery],
            outputs=[init_img_with_mask],
        )

        img2img_send_to_img2img.click(
            fn=lambda x: image_from_url_text(x),
            _js="extract_image_from_gallery",
            inputs=[img2img_gallery],
            outputs=[init_img],
        )

        img2img_send_to_inpaint.click(
            fn=lambda x: image_from_url_text(x),
            _js="extract_image_from_gallery",
            inputs=[img2img_gallery],
            outputs=[init_img_with_mask],
        )

        send_to_extras.click(
            fn=lambda x: image_from_url_text(x),
S
Seki 已提交
753
            _js="extract_image_from_gallery_extras",
A
AUTOMATIC 已提交
754 755 756 757 758 759
            inputs=[txt2img_gallery],
            outputs=[image],
        )

        img2img_send_to_extras.click(
            fn=lambda x: image_from_url_text(x),
S
Seki 已提交
760
            _js="extract_image_from_gallery_extras",
A
AUTOMATIC 已提交
761 762 763
            inputs=[img2img_gallery],
            outputs=[image],
        )
764

765
    ui_config_file = cmd_opts.ui_config_file
A
AUTOMATIC 已提交
766 767 768 769 770 771 772 773 774 775 776 777 778 779
    ui_settings = {}
    settings_count = len(ui_settings)
    error_loading = False

    try:
        if os.path.exists(ui_config_file):
            with open(ui_config_file, "r", encoding="utf8") as file:
                ui_settings = json.load(file)
    except Exception:
        error_loading = True
        print("Error loading settings:", file=sys.stderr)
        print(traceback.format_exc(), file=sys.stderr)

    def loadsave(path, x):
780
        def apply_field(obj, field, condition=None):
A
AUTOMATIC 已提交
781 782 783 784 785
            key = path + "/" + field

            saved_value = ui_settings.get(key, None)
            if saved_value is None:
                ui_settings[key] = getattr(obj, field)
786
            elif condition is None or condition(saved_value):
A
AUTOMATIC 已提交
787 788 789 790 791 792 793 794 795
                setattr(obj, field, saved_value)

        if type(x) == gr.Slider:
            apply_field(x, 'value')
            apply_field(x, 'minimum')
            apply_field(x, 'maximum')
            apply_field(x, 'step')

        if type(x) == gr.Radio:
796
            apply_field(x, 'value', lambda val: val in x.choices)
A
AUTOMATIC 已提交
797 798 799 800 801 802 803 804

    visit(txt2img_interface, loadsave, "txt2img")
    visit(img2img_interface, loadsave, "img2img")

    if not error_loading and (not os.path.exists(ui_config_file) or settings_count != len(ui_settings)):
        with open(ui_config_file, "w", encoding="utf8") as file:
            json.dump(ui_settings, file, indent=4)

805 806 807
    return demo


A
AUTOMATIC 已提交
808 809
with open(os.path.join(script_path, "script.js"), "r", encoding="utf8") as jsfile:
    javascript = jsfile.read()
810 811


A
AUTOMATIC 已提交
812 813 814 815 816
def template_response(*args, **kwargs):
    res = gradio_routes_templates_response(*args, **kwargs)
    res.body = res.body.replace(b'</head>', f'<script>{javascript}</script></head>'.encode("utf8"))
    res.init_headers()
    return res
817 818


A
AUTOMATIC 已提交
819 820
gradio_routes_templates_response = gradio.routes.templates.TemplateResponse
gradio.routes.templates.TemplateResponse = template_response