updated interface to use Blocks

added mode toggle for img2img
added inpainting to readme

README.md

@@ -204,6 +204,10 @@ To use this feature, tick a checkbox in the img2img interface. Original
 image will be upscaled to twice the original width and height, while width and height sliders
 will specify the size of individual tiles. At the moment this method does not support batch size.
 
+Rcommended parameters for upscaling:
+ - Sampling method: Euler a
+ - Denoising strength: 0.2, can go up to 0.4 if you feel adventureous
+
 ![](images/sd-upscale.jpg)
 
 ### User scripts
@@ -261,4 +265,7 @@ compared to normal operation on my RTX 3090.
 This is an independent implementation that does not require any modification to original Stable Diffusion code, and
 with all code concenrated in one place rather than scattered around the program.
 
+### Inpainting
+In img2img tab, draw a mask over a part of image, and that part will be in-painted.
 
+![](images/inpainting.png)

style.css

@@ -0,0 +1,61 @@
+
+button{
+    align-self: stretch !important;
+}
+
+#img2img_mode{
+    padding: 0 0 1em 0;
+    border: none !important;
+}
+
+#img2img_prompt, #txt2img_prompt{
+    padding: 0;
+    border: none !important;
+}
+
+#img2maskimg .h-60{
+    height: 30rem;
+}
+
+.overflow-hidden, .gr-panel{
+    overflow: visible !important;
+}
+
+fieldset span.text-gray-500{
+    position: absolute;
+    top: -0.425em;
+    background-color: rgb(249 250 251);
+    line-height: 0.7em;
+    padding: 0 0.5em;
+}
+.dark fieldset span.text-gray-500 { background-color: rgb(17 24 39); }
+
+.gr-panel div.flex-col div.justify-between label, label.block span{
+    position: absolute;
+    top: -0.4em;
+    background-color: rgb(249 250 251);
+    line-height: 0.7em;
+    padding: 0 0.5em;
+    margin: 0;
+}
+.dark .gr-panel div.flex-col div.justify-between label{ background-color: rgb(17 24 39); }
+
+.gr-panel div.flex-col div.justify-between label span{
+    margin: 0;
+}
+
+.gr-panel div.flex-col div.justify-between div{
+  position: absolute;
+  top: -0.1em;
+  right: 1em;
+  padding: 0 0.5em;
+}
+
+input[type="range"]{
+    margin: 0.5em 0 -0.3em 0;
+}
+
+#txt2img_sampling label{
+    padding-left: 0.6em;
+    padding-right: 0.6em;
+}

webui.py

@@ -8,6 +8,7 @@ import torch
 import torch.nn as nn
 import numpy as np
 import gradio as gr
+import gradio.utils
 from omegaconf import OmegaConf
 from PIL import Image, ImageFont, ImageDraw, PngImagePlugin, ImageFilter, ImageOps
 from torch import autocast
@@ -24,18 +25,16 @@ from ldm.util import instantiate_from_config
 from ldm.models.diffusion.ddim import DDIMSampler
 from ldm.models.diffusion.plms import PLMSSampler
 
-try:
+# fix gradio phoning home
-    # this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.
+gradio.utils.version_check = lambda: None
-
+gradio.utils.get_local_ip_address = lambda: '127.0.0.1'
-    from transformers import logging
-    logging.set_verbosity_error()
-except Exception:
-    pass
 
 # 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')
 
+script_path = os.path.dirname(os.path.realpath(__file__))
+
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
 opt_C = 4
 opt_f = 8
@@ -72,12 +71,12 @@ css_hide_progressbar = """
 SamplerData = namedtuple('SamplerData', ['name', 'constructor'])
 samplers = [
     *[SamplerData(x[0], lambda funcname=x[1]: KDiffusionSampler(funcname)) for x in [
-        ('Euler ancestral', 'sample_euler_ancestral'),
+        ('Euler a', 'sample_euler_ancestral'),
         ('Euler', 'sample_euler'),
         ('LMS', 'sample_lms'),
         ('Heun', 'sample_heun'),
-        ('DPM 2', 'sample_dpm_2'),
+        ('DPM2', 'sample_dpm_2'),
-        ('DPM 2 Ancestral', 'sample_dpm_2_ancestral'),
+        ('DPM2 a', 'sample_dpm_2_ancestral'),
     ] if hasattr(k_diffusion.sampling, x[1])],
     SamplerData('DDIM', lambda: VanillaStableDiffusionSampler(DDIMSampler)),
     SamplerData('PLMS', lambda: VanillaStableDiffusionSampler(PLMSSampler)),
@@ -1083,31 +1082,67 @@ class Flagging(gr.FlaggingCallback):
 
         print("Logged:", filenames[0])
 
+with gr.Blocks(analytics_enabled=False) as txt2img_interface:
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", elem_id="txt2img_prompt", show_label=False, placeholder="Prompt", lines=1)
+        submit = gr.Button('Generate', variant='primary')
+
+    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_for_img2img[0].name, type="index")
+
+            with gr.Row():
+                use_GFPGAN = gr.Checkbox(label='GFPGAN', value=False, visible=have_gfpgan)
+                prompt_matrix = gr.Checkbox(label='Prompt matrix', value=False)
+
+            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='Classifier Free Guidance Scale (how strongly the image should follow the prompt)', 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)
+
+            seed = gr.Number(label='Seed', value=-1)
+
+            code = gr.Textbox(label="Python script", visible=cmd_opts.allow_code, lines=1)
+
+        with gr.Column(variant='panel'):
+            with gr.Group():
+                gallery = gr.Gallery(label='Output')
+                output_seed = gr.Number(label='Seed', visible=False)
+                html_info = gr.HTML()
+
+        txt2img_args = dict(
+            fn=wrap_gradio_call(txt2img),
+            inputs=[
+                prompt,
+                steps,
+                sampler_index,
+                use_GFPGAN,
+                prompt_matrix,
+                batch_count,
+                batch_size,
+                cfg_scale,
+                seed,
+                height,
+                width,
+                code
+            ],
+            outputs=[
+                gallery,
+                output_seed,
+                html_info
+            ]
+        )
+
+        prompt.submit(**txt2img_args)
+        submit.click(**txt2img_args)
+
 
-txt2img_interface = gr.Interface(
-    wrap_gradio_call(txt2img),
-    inputs=[
-        gr.Textbox(label="Prompt", placeholder="A corgi wearing a top hat as an oil painting.", lines=1),
-        gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20),
-        gr.Radio(label='Sampling method', choices=[x.name for x in samplers], value=samplers[0].name, type="index"),
-        gr.Checkbox(label='Fix faces using GFPGAN', value=False, visible=have_gfpgan),
-        gr.Checkbox(label='Create prompt matrix (separate multiple prompts using |, and get all combinations of them)', value=False),
-        gr.Slider(minimum=1, maximum=cmd_opts.max_batch_count, step=1, label='Batch count (how many batches of images to generate)', value=1),
-        gr.Slider(minimum=1, maximum=8, step=1, label='Batch size (how many images are in a batch; memory-hungry)', value=1),
-        gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='Classifier Free Guidance Scale (how strongly the image should follow the prompt)', value=7.5),
-        gr.Number(label='Seed', value=-1),
-        gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512),
-        gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512),
-        gr.Textbox(label="Python script", visible=cmd_opts.allow_code, lines=1)
-    ],
-    outputs=[
-        gr.Gallery(label="Images"),
-        gr.Number(label='Seed'),
-        gr.HTML(),
-    ],
-    title="Stable Diffusion Text-to-Image",
-    flagging_callback=Flagging()
-)
 
 def fill(image, mask):
     image_mod = Image.new('RGBA', (image.width, image.height))
@@ -1223,10 +1258,15 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         return samples_ddim
 
 
-def img2img(prompt: str, init_img, init_img_with_mask, ddim_steps: int, sampler_index: int, mask_blur: int, inpainting_fill: int, use_GFPGAN: bool, prompt_matrix, loopback: bool, sd_upscale: bool, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, height: int, width: int, resize_mode: int):
+def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index: int, mask_blur: int, inpainting_fill: int, use_GFPGAN: bool, prompt_matrix, mode: int, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, height: int, width: int, resize_mode: int):
     outpath = opts.outdir or "outputs/img2img-samples"
 
-    if init_img_with_mask is not None:
+    is_classic = mode == 0
+    is_inpaint = mode == 1
+    is_loopback = mode == 2
+    is_upscale = mode == 3
+
+    if is_inpaint:
         image = init_img_with_mask['image']
         mask = init_img_with_mask['mask']
     else:
@@ -1242,7 +1282,7 @@ def img2img(prompt: str, init_img, init_img_with_mask, ddim_steps: int, sampler_
         sampler_index=sampler_index,
         batch_size=batch_size,
         n_iter=n_iter,
-        steps=ddim_steps,
+        steps=steps,
         cfg_scale=cfg_scale,
         width=width,
         height=height,
@@ -1257,7 +1297,7 @@ def img2img(prompt: str, init_img, init_img_with_mask, ddim_steps: int, sampler_
         extra_generation_params={"Denoising Strength": denoising_strength}
     )
 
-    if loopback:
+    if is_loopback:
         output_images, info = None, None
         history = []
         initial_seed = None
@@ -1286,7 +1326,7 @@ def img2img(prompt: str, init_img, init_img_with_mask, ddim_steps: int, sampler_
 
         processed = Processed(history, initial_seed, initial_info)
 
-    elif sd_upscale:
+    elif is_upscale:
         initial_seed = None
         initial_info = None
 
@@ -1345,36 +1385,103 @@ def img2img(prompt: str, init_img, init_img_with_mask, ddim_steps: int, sampler_
 sample_img2img = "assets/stable-samples/img2img/sketch-mountains-input.jpg"
 sample_img2img = sample_img2img if os.path.exists(sample_img2img) else None
 
-img2img_interface = gr.Interface(
+
-    wrap_gradio_call(img2img),
+with gr.Blocks(analytics_enabled=False) as img2img_interface:
-    inputs=[
+    with gr.Row():
-        gr.Textbox(placeholder="A fantasy landscape, trending on artstation.", lines=1),
+        prompt = gr.Textbox(label="Prompt", elem_id="img2img_prompt", show_label=False, placeholder="Prompt", lines=1)
-        gr.Image(label="Image for img2img", source="upload", interactive=True, type="pil"),
+        submit = gr.Button('Generate', variant='primary')
-        gr.Image(label="Image for inpainting with mask", source="upload", interactive=True, type="pil", tool="sketch"),
+
-        gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20),
+    with gr.Row().style(equal_height=False):
-        gr.Radio(label='Sampling method', choices=[x.name for x in samplers_for_img2img], value=samplers_for_img2img[0].name, type="index"),
+
-        gr.Slider(label='Inpainting: mask blur', minimum=0, maximum=64, step=1, value=4),
+        with gr.Column(variant='panel'):
-        gr.Radio(label='Inpainting: masked content', choices=['fill', 'original', 'latent noise', 'latent nothing'], value='fill', type="index"),
+            with gr.Group():
-        gr.Checkbox(label='Fix faces using GFPGAN', value=False, visible=have_gfpgan),
+                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)
-        gr.Checkbox(label='Create prompt matrix (separate multiple prompts using |, and get all combinations of them)', value=False),
+                init_img = gr.Image(label="Image for img2img", source="upload", interactive=True, type="pil")
-        gr.Checkbox(label='Loopback (use images from previous batch when creating next batch)', value=False),
+                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)
-        gr.Checkbox(label='Stable Diffusion upscale', value=False),
+                resize_mode = gr.Radio(label="Resize mode", show_label=False, choices=["Just resize", "Crop and resize", "Resize and fill"], type="index", value="Just resize")
-        gr.Slider(minimum=1, maximum=cmd_opts.max_batch_count, step=1, label='Batch count (how many batches of images to generate)', value=1),
+
-        gr.Slider(minimum=1, maximum=8, step=1, label='Batch size (how many images are in a batch; memory-hungry)', value=1),
+            steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20)
-        gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='Classifier Free Guidance Scale (how strongly the image should follow the prompt)', value=7.0),
+            sampler_index = gr.Radio(label='Sampling method', choices=[x.name for x in samplers_for_img2img], value=samplers_for_img2img[0].name, type="index")
-        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising Strength', value=0.75),
+            mask_blur = gr.Slider(label='Inpainting: mask blur', minimum=0, maximum=64, step=1, value=4, visible=False)
-        gr.Number(label='Seed', value=-1),
+            inpainting_fill = gr.Radio(label='Inpainting: masked content', choices=['fill', 'original', 'latent noise', 'latent nothing'], value='fill', type="index", visible=False)
-        gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512),
+
-        gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512),
+            with gr.Row():
-        gr.Radio(label="Resize mode", choices=["Just resize", "Crop and resize", "Resize and fill"], type="index", value="Just resize")
+                use_GFPGAN = gr.Checkbox(label='GFPGAN', value=False, visible=have_gfpgan)
-    ],
+                prompt_matrix = gr.Checkbox(label='Prompt matrix', value=False)
-    outputs=[
+
-        gr.Gallery(),
+            with gr.Row():
-        gr.Number(label='Seed'),
+                batch_count = gr.Slider(minimum=1, maximum=cmd_opts.max_batch_count, step=1, label='Batch count', value=1)
-        gr.HTML(),
+                batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1)
-    ],
+
-    allow_flagging="never",
+            with gr.Group():
-)
+                cfg_scale = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='Classifier Free Guidance Scale (how strongly the image should follow the prompt)', value=7.0)
+                denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising Strength', value=0.75)
+
+            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)
+
+            seed = gr.Number(label='Seed', value=-1)
+
+        with gr.Column(variant='panel'):
+            with gr.Group():
+                gallery = gr.Gallery(label='Output')
+                output_seed = gr.Number(label='Seed', visible=False)
+                html_info = gr.HTML()
+
+        def apply_mode(mode):
+            is_classic = mode == 0
+            is_inpaint = mode == 1
+            is_loopback = mode == 2
+            is_upscale = mode == 3
+
+            return {
+                init_img: gr.update(visible=not is_inpaint),
+                init_img_with_mask: gr.update(visible=is_inpaint),
+                mask_blur: gr.update(visible=is_inpaint),
+                inpainting_fill: gr.update(visible=is_inpaint),
+                prompt_matrix: gr.update(visible=is_classic),
+                batch_count: gr.update(visible=not is_upscale),
+                batch_size: gr.update(visible=not is_loopback),
+            }
+
+        switch_mode.change(
+            apply_mode,
+            inputs=[switch_mode],
+            outputs=[init_img, init_img_with_mask, mask_blur, inpainting_fill, prompt_matrix, batch_count, batch_size]
+        )
+
+        img2img_args = dict(
+            fn=wrap_gradio_call(img2img),
+            inputs=[
+                prompt,
+                init_img,
+                init_img_with_mask,
+                steps,
+                sampler_index,
+                mask_blur,
+                inpainting_fill,
+                use_GFPGAN,
+                prompt_matrix,
+                switch_mode,
+                batch_count,
+                batch_size,
+                cfg_scale,
+                denoising_strength,
+                seed,
+                height,
+                width,
+                resize_mode
+            ],
+            outputs=[
+                gallery,
+                output_seed,
+                html_info
+            ]
+        )
+
+        prompt.submit(**img2img_args)
+        submit.click(**img2img_args)
 
 
 def upscale_with_realesrgan(image, RealESRGAN_upscaling, RealESRGAN_model_index):
@@ -1434,6 +1541,7 @@ extras_interface = gr.Interface(
         gr.HTML(),
     ],
     allow_flagging="never",
+    analytics_enabled=False,
 )
 
 
@@ -1463,6 +1571,7 @@ pnginfo_interface = gr.Interface(
         gr.HTML(),
     ],
     allow_flagging="never",
+    analytics_enabled=False,
 )
 
 
@@ -1514,6 +1623,7 @@ settings_interface = gr.Interface(
     title=None,
     description=None,
     allow_flagging="never",
+    analytics_enabled=False,
 )
 
 interfaces = [
@@ -1524,6 +1634,15 @@ interfaces = [
     (settings_interface, "Settings"),
 ]
 
+try:
+    # this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.
+
+    from transformers import logging
+
+    logging.set_verbosity_error()
+except Exception:
+    pass
+
 sd_config = OmegaConf.load(cmd_opts.config)
 sd_model = load_model_from_config(sd_config, cmd_opts.ckpt)
 sd_model = (sd_model if cmd_opts.no_half else sd_model.half())
@@ -1537,13 +1656,17 @@ else:
 model_hijack = StableDiffusionModelHijack()
 model_hijack.hijack(sd_model)
 
+with open(os.path.join(script_path, "style.css"), "r", encoding="utf8") as file:
+    css = file.read()
+
 demo = gr.TabbedInterface(
     interface_list=[x[0] for x in interfaces],
     tab_names=[x[1] for x in interfaces],
     css=("" if cmd_opts.no_progressbar_hiding else css_hide_progressbar) + """
 .output-html p {margin: 0 0.5em;}
 .performance { font-size: 0.85em; color: #444; }
-"""
+""" + css,
+    analytics_enabled=False,
 )
 
 demo.queue(concurrency_count=1)