f31c4ab57b53cde8341c10ec5dc18ad37ce7ce0c5359a832639c64914cfc2c52

extensions/addtional/models/lora/sukuna.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f15673b5ea25ed5f1e5cee150ddf4e0d094a65e30f08b0813566f9580abdb627
+size 78018750

extensions/addtional/models/lora/vacuum fellatio1.1【vacuum fellatio】.civitai.info

@@ -0,0 +1,85 @@
+{
+    "id": 22489,
+    "modelId": 18735,
+    "name": "v1.1",
+    "createdAt": "2023-03-13T09:00:46.954Z",
+    "updatedAt": "2023-03-18T10:17:13.538Z",
+    "trainedWords": [
+        "vacuum fellatio"
+    ],
+    "baseModel": "SD 1.5",
+    "earlyAccessTimeFrame": 0,
+    "description": "<p>There is no significant change, but considering that the previous TriggerWords cannot be saved as a file name, I modified TriggerWords. </p><p>By the way, I added several pictures of POV deepthroat, which may have better performance</p>",
+    "stats": {
+        "downloadCount": 17393,
+        "ratingCount": 11,
+        "rating": 5
+    },
+    "model": {
+        "name": "concept Vacuum fellatio(:>=)",
+        "type": "LORA",
+        "nsfw": true,
+        "poi": false
+    },
+    "files": [
+        {
+            "id": 18716,
+            "url": "https://5ac0637cfd0766c97916cefa3764fbdf.r2.cloudflarestorage.com/civitai-prod-settled/105182/model/vacuum20fellatio11E380.pABG.safetensors",
+            "sizeKB": 36985.2275390625,
+            "name": "vacuum fellatio1.1\u3010vacuum fellatio\u3011.safetensors",
+            "type": "Model",
+            "metadata": {
+                "fp": "fp16",
+                "size": "full",
+                "format": "SafeTensor"
+            },
+            "pickleScanResult": "Success",
+            "pickleScanMessage": "No Pickle imports",
+            "virusScanResult": "Success",
+            "virusScanMessage": null,
+            "scannedAt": "2023-03-13T09:01:34.994Z",
+            "hashes": {
+                "AutoV1": "378523B1",
+                "AutoV2": "04D50C1BA2",
+                "SHA256": "04D50C1BA2BD0A1B5F87118C29A8C4C596CEA0BC0991A3EAE7952AA86ED540EB",
+                "CRC32": "EAD3612B",
+                "BLAKE3": "F9652C1CC61EA0B42904CF95C9544EAF792EC7AF3B7B7A4F7CD01FAE67052962"
+            },
+            "primary": true,
+            "downloadUrl": "https://civitai.com/api/download/models/22489"
+        }
+    ],
+    "images": [
+        {
+            "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/3bbcf1a0-0058-4d4b-9d2f-857163335400/width=450/241954.jpeg",
+            "nsfw": "X",
+            "width": 512,
+            "height": 640,
+            "hash": "UEHL-;Q800N@00E,wukX00Tykr^iDO~B?c4:",
+            "meta": {
+                "seed": 2965610135,
+                "steps": 20,
+                "prompt": "masterpiece,1girl,vacuum fellatio,1boy,penis,fellatio,oral,from side, surprise deepthroat,  <lora:vacuum fellatio1.1\u3010vacuum fellatio\u3011:1>",
+                "sampler": "Euler a",
+                "cfgScale": 7,
+                "negativePrompt": "(worst quality:1.4), (low quality:1.4), (normal quality:1.4), lowres"
+            }
+        },
+        {
+            "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/1f78213f-2c53-4e52-7b11-a946bb98d400/width=450/241953.jpeg",
+            "nsfw": "None",
+            "width": 512,
+            "height": 640,
+            "hash": "UAEoGm0N00%E^%E3nOxU00?E~CM*DR-.%f9b",
+            "meta": {
+                "seed": 3859538840,
+                "steps": 20,
+                "prompt": "masterpiece,1girl,vacuum fellatio,1boy,penis,fellatio,oral,pov, surprise deepthroat,  <lora:vacuum fellatio1.1\u3010vacuum fellatio\u3011:1>",
+                "sampler": "Euler a",
+                "cfgScale": 7,
+                "negativePrompt": "(worst quality:1.4), (low quality:1.4), (normal quality:1.4), lowres"
+            }
+        }
+    ],
+    "downloadUrl": "https://civitai.com/api/download/models/22489"
+}

extensions/addtional/models/lora/vacuum fellatio1.1【vacuum fellatio】.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04d50c1ba2bd0a1b5f87118c29a8c4c596cea0bc0991a3eae7952aa86ed540eb
+size 37872873

extensions/addtional/preload.py

@@ -0,0 +1,6 @@
+import os
+from modules import paths
+
+
+def preload(parser):
+    parser.add_argument("--addnet-max-model-count", type=int, help="The maximum number of additional network model can be used.", default=5)

extensions/addtional/scripts/additional_networks.py

@@ -0,0 +1,399 @@
+import os
+
+import torch
+import numpy as np
+
+import modules.scripts as scripts
+from modules import shared, script_callbacks
+import gradio as gr
+
+import modules.ui
+from modules.ui_components import ToolButton, FormRow
+
+from scripts import addnet_xyz_grid_support, lora_compvis, model_util, metadata_editor
+from scripts.model_util import lora_models, MAX_MODEL_COUNT
+
+
+memo_symbol = "\U0001F4DD"  # 📝
+addnet_paste_params = {"txt2img": [], "img2img": []}
+
+
+class Script(scripts.Script):
+    def __init__(self) -> None:
+        super().__init__()
+        self.latest_params = [(None, None, None, None)] * MAX_MODEL_COUNT
+        self.latest_networks = []
+        self.latest_model_hash = ""
+
+    def title(self):
+        return "Additional networks for generating"
+
+    def show(self, is_img2img):
+        return scripts.AlwaysVisible
+
+    def ui(self, is_img2img):
+        global addnet_paste_params
+        # NOTE: Changing the contents of `ctrls` means the XY Grid support may need
+        # to be updated, see xyz_grid_support.py
+        ctrls = []
+        weight_sliders = []
+        model_dropdowns = []
+
+        tabname = "txt2img"
+        if is_img2img:
+            tabname = "img2img"
+
+        paste_params = addnet_paste_params[tabname]
+        paste_params.clear()
+
+        self.infotext_fields = []
+        self.paste_field_names = []
+
+        with gr.Group():
+            with gr.Accordion("Additional Networks", open=False):
+                with gr.Row():
+                    enabled = gr.Checkbox(label="Enable", value=False)
+                    ctrls.append(enabled)
+                    self.infotext_fields.append((enabled, "AddNet Enabled"))
+                    separate_weights = gr.Checkbox(label="Separate UNet/Text Encoder weights", value=False)
+                    ctrls.append(separate_weights)
+                    self.infotext_fields.append((separate_weights, "AddNet Separate Weights"))
+
+                for i in range(MAX_MODEL_COUNT):
+                    with FormRow(variant="compact"):
+                        module = gr.Dropdown(["LoRA"], label=f"Network module {i+1}", value="LoRA")
+                        model = gr.Dropdown(list(lora_models.keys()), label=f"Model {i+1}", value="None")
+                        with gr.Row(visible=False):
+                            model_path = gr.Textbox(value="None", interactive=False, visible=False)
+                        model.change(
+                            lambda module, model, i=i: model_util.lora_models.get(model, "None"),
+                            inputs=[module, model],
+                            outputs=[model_path],
+                        )
+
+                        # Sending from the script UI to the metadata editor has to bypass
+                        # gradio since this button will exit the gr.Blocks context by the
+                        # time the metadata editor tab is created, so event handlers can't
+                        # be registered on it by then.
+                        model_info = ToolButton(value=memo_symbol, elem_id=f"additional_networks_send_to_metadata_editor_{i}")
+                        model_info.click(fn=None, _js="addnet_send_to_metadata_editor", inputs=[module, model_path], outputs=[])
+
+                        module.change(
+                            lambda module, model, i=i: addnet_xyz_grid_support.update_axis_params(i, module, model),
+                            inputs=[module, model],
+                            outputs=[],
+                        )
+                        model.change(
+                            lambda module, model, i=i: addnet_xyz_grid_support.update_axis_params(i, module, model),
+                            inputs=[module, model],
+                            outputs=[],
+                        )
+
+                        # perhaps there is no user to train Text Encoder only, Weight A is U-Net
+                        # The name of label will be changed in future (Weight A and B), but UNet and TEnc for now for easy understanding
+                        with gr.Column() as col:
+                            weight = gr.Slider(label=f"Weight {i+1}", value=1.0, minimum=-1.0, maximum=2.0, step=0.05, visible=True)
+                            weight_unet = gr.Slider(
+                                label=f"UNet Weight {i+1}", value=1.0, minimum=-1.0, maximum=2.0, step=0.05, visible=False
+                            )
+                            weight_tenc = gr.Slider(
+                                label=f"TEnc Weight {i+1}", value=1.0, minimum=-1.0, maximum=2.0, step=0.05, visible=False
+                            )
+
+                        weight.change(lambda w: (w, w), inputs=[weight], outputs=[weight_unet, weight_tenc])
+                        weight.release(lambda w: (w, w), inputs=[weight], outputs=[weight_unet, weight_tenc])
+                        paste_params.append({"module": module, "model": model})
+
+                    ctrls.extend((module, model, weight_unet, weight_tenc))
+                    weight_sliders.extend((weight, weight_unet, weight_tenc))
+                    model_dropdowns.append(model)
+
+                    self.infotext_fields.extend(
+                        [
+                            (module, f"AddNet Module {i+1}"),
+                            (model, f"AddNet Model {i+1}"),
+                            (weight, f"AddNet Weight {i+1}"),
+                            (weight_unet, f"AddNet Weight A {i+1}"),
+                            (weight_tenc, f"AddNet Weight B {i+1}"),
+                        ]
+                    )
+
+                for _, field_name in self.infotext_fields:
+                    self.paste_field_names.append(field_name)
+
+                def update_weight_sliders(separate, *sliders):
+                    updates = []
+                    for w, w_unet, w_tenc in zip(*(iter(sliders),) * 3):
+                        if not separate:
+                            w_unet = w
+                            w_tenc = w
+                        updates.append(gr.Slider.update(visible=not separate))  # Combined
+                        updates.append(gr.Slider.update(visible=separate, value=w_unet))  # UNet
+                        updates.append(gr.Slider.update(visible=separate, value=w_tenc))  # TEnc
+                    return updates
+
+                separate_weights.change(update_weight_sliders, inputs=[separate_weights] + weight_sliders, outputs=weight_sliders)
+
+                def refresh_all_models(*dropdowns):
+                    model_util.update_models()
+                    updates = []
+                    for dd in dropdowns:
+                        if dd in lora_models:
+                            selected = dd
+                        else:
+                            selected = "None"
+                        update = gr.Dropdown.update(value=selected, choices=list(lora_models.keys()))
+                        updates.append(update)
+                    return updates
+
+                # mask for regions
+                with gr.Accordion("Extra args", open=False):
+                    with gr.Row():
+                        mask_image = gr.Image(label="mask image:")
+                        ctrls.append(mask_image)
+
+                refresh_models = gr.Button(value="Refresh models")
+                refresh_models.click(refresh_all_models, inputs=model_dropdowns, outputs=model_dropdowns)
+                ctrls.append(refresh_models)
+
+        return ctrls
+
+    def set_infotext_fields(self, p, params):
+        for i, t in enumerate(params):
+            module, model, weight_unet, weight_tenc = t
+            if model is None or model == "None" or len(model) == 0 or (weight_unet == 0 and weight_tenc == 0):
+                continue
+            p.extra_generation_params.update(
+                {
+                    "AddNet Enabled": True,
+                    f"AddNet Module {i+1}": module,
+                    f"AddNet Model {i+1}": model,
+                    f"AddNet Weight A {i+1}": weight_unet,
+                    f"AddNet Weight B {i+1}": weight_tenc,
+                }
+            )
+
+    def restore_networks(self, sd_model):
+        unet = sd_model.model.diffusion_model
+        text_encoder = sd_model.cond_stage_model
+
+        if len(self.latest_networks) > 0:
+            print("restoring last networks")
+            for network, _ in self.latest_networks[::-1]:
+                network.restore(text_encoder, unet)
+            self.latest_networks.clear()
+
+    def process_batch(self, p, *args, **kwargs):
+        unet = p.sd_model.model.diffusion_model
+        text_encoder = p.sd_model.cond_stage_model
+
+        if not args[0]:
+            self.restore_networks(p.sd_model)
+            return
+
+        params = []
+        for i, ctrl in enumerate(args[2:]):
+            if i % 4 == 0:
+                param = [ctrl]
+            else:
+                param.append(ctrl)
+                if i % 4 == 3:
+                    params.append(param)
+
+        models_changed = len(self.latest_networks) == 0  # no latest network (cleared by check-off)
+        models_changed = models_changed or self.latest_model_hash != p.sd_model.sd_model_hash
+        if not models_changed:
+            for (l_module, l_model, l_weight_unet, l_weight_tenc), (module, model, weight_unet, weight_tenc) in zip(
+                self.latest_params, params
+            ):
+                if l_module != module or l_model != model or l_weight_unet != weight_unet or l_weight_tenc != weight_tenc:
+                    models_changed = True
+                    break
+
+        if models_changed:
+            self.restore_networks(p.sd_model)
+            self.latest_params = params
+            self.latest_model_hash = p.sd_model.sd_model_hash
+
+            for module, model, weight_unet, weight_tenc in self.latest_params:
+                if model is None or model == "None" or len(model) == 0:
+                    continue
+                if weight_unet == 0 and weight_tenc == 0:
+                    print(f"ignore because weight is 0: {model}")
+                    continue
+
+                model_path = lora_models.get(model, None)
+                if model_path is None:
+                    raise RuntimeError(f"model not found: {model}")
+
+                if model_path.startswith('"') and model_path.endswith('"'):  # trim '"' at start/end
+                    model_path = model_path[1:-1]
+                if not os.path.exists(model_path):
+                    print(f"file not found: {model_path}")
+                    continue
+
+                print(f"{module} weight_unet: {weight_unet}, weight_tenc: {weight_tenc}, model: {model}")
+                if module == "LoRA":
+                    if os.path.splitext(model_path)[1] == ".safetensors":
+                        from safetensors.torch import load_file
+
+                        du_state_dict = load_file(model_path)
+                    else:
+                        du_state_dict = torch.load(model_path, map_location="cpu")
+
+                    network, info = lora_compvis.create_network_and_apply_compvis(
+                        du_state_dict, weight_tenc, weight_unet, text_encoder, unet
+                    )
+                    # in medvram, device is different for u-net and sd_model, so use sd_model's
+                    network.to(p.sd_model.device, dtype=p.sd_model.dtype)
+
+                    print(f"LoRA model {model} loaded: {info}")
+                    self.latest_networks.append((network, model))
+            if len(self.latest_networks) > 0:
+                print("setting (or sd model) changed. new networks created.")
+
+        # apply mask: currently only top 3 networks are supported
+        if len(self.latest_networks) > 0:
+            mask_image = args[-2]
+            if mask_image is not None:
+                mask_image = mask_image.astype(np.float32) / 255.0
+                print(f"use mask image to control LoRA regions.")
+                for i, (network, model) in enumerate(self.latest_networks[:3]):
+                    if not hasattr(network, "set_mask"):
+                        continue
+                    mask = mask_image[:, :, i]  # R,G,B
+                    if mask.max() <= 0:
+                        continue
+                    mask = torch.tensor(mask, dtype=p.sd_model.dtype, device=p.sd_model.device)
+
+                    network.set_mask(mask, height=p.height, width=p.width, hr_height=p.hr_upscale_to_y, hr_width=p.hr_upscale_to_x)
+                    print(f"apply mask. channel: {i}, model: {model}")
+            else:
+                for network, _ in self.latest_networks:
+                    if hasattr(network, "set_mask"):
+                        network.set_mask(None)
+
+        self.set_infotext_fields(p, self.latest_params)
+
+
+def on_script_unloaded():
+    if shared.sd_model:
+        for s in scripts.scripts_txt2img.alwayson_scripts:
+            if isinstance(s, Script):
+                s.restore_networks(shared.sd_model)
+                break
+
+
+def on_ui_tabs():
+    global addnet_paste_params
+    with gr.Blocks(analytics_enabled=False) as additional_networks_interface:
+        metadata_editor.setup_ui(addnet_paste_params)
+
+    return [(additional_networks_interface, "Additional Networks", "additional_networks")]
+
+
+def on_ui_settings():
+    section = ("additional_networks", "Additional Networks")
+    shared.opts.add_option(
+        "additional_networks_extra_lora_path",
+        shared.OptionInfo(
+            "",
+            """Extra paths to scan for LoRA models, comma-separated. Paths containing commas must be enclosed in double quotes. In the path, " (one quote) must be replaced by "" (two quotes).""",
+            section=section,
+        ),
+    )
+    shared.opts.add_option(
+        "additional_networks_sort_models_by",
+        shared.OptionInfo(
+            "name",
+            "Sort LoRA models by",
+            gr.Radio,
+            {"choices": ["name", "date", "path name", "rating", "has user metadata"]},
+            section=section,
+        ),
+    )
+    shared.opts.add_option(
+        "additional_networks_reverse_sort_order", shared.OptionInfo(False, "Reverse model sort order", section=section)
+    )
+    shared.opts.add_option(
+        "additional_networks_model_name_filter", shared.OptionInfo("", "LoRA model name filter", section=section)
+    )
+    shared.opts.add_option(
+        "additional_networks_xy_grid_model_metadata",
+        shared.OptionInfo(
+            "",
+            'Metadata to show in XY-Grid label for Model axes, comma-separated (example: "ss_learning_rate, ss_num_epochs")',
+            section=section,
+        ),
+    )
+    shared.opts.add_option(
+        "additional_networks_hash_thread_count",
+        shared.OptionInfo(1, "# of threads to use for hash calculation (increase if using an SSD)", section=section),
+    )
+    shared.opts.add_option(
+        "additional_networks_back_up_model_when_saving",
+        shared.OptionInfo(True, "Make a backup copy of the model being edited when saving its metadata.", section=section),
+    )
+    shared.opts.add_option(
+        "additional_networks_show_only_safetensors",
+        shared.OptionInfo(False, "Only show .safetensors format models", section=section),
+    )
+    shared.opts.add_option(
+        "additional_networks_show_only_models_with_metadata",
+        shared.OptionInfo(
+            "disabled",
+            "Only show models that have/don't have user-added metadata",
+            gr.Radio,
+            {"choices": ["disabled", "has metadata", "missing metadata"]},
+            section=section,
+        ),
+    )
+    shared.opts.add_option(
+        "additional_networks_max_top_tags", shared.OptionInfo(20, "Max number of top tags to show", section=section)
+    )
+    shared.opts.add_option(
+        "additional_networks_max_dataset_folders", shared.OptionInfo(20, "Max number of dataset folders to show", section=section)
+    )
+
+
+def on_infotext_pasted(infotext, params):
+    if "AddNet Enabled" not in params:
+        params["AddNet Enabled"] = "False"
+
+    # TODO changing "AddNet Separate Weights" does not seem to work
+    if "AddNet Separate Weights" not in params:
+        params["AddNet Separate Weights"] = "False"
+
+    for i in range(MAX_MODEL_COUNT):
+        # Convert combined weight into new format
+        if f"AddNet Weight {i+1}" in params:
+            params[f"AddNet Weight A {i+1}"] = params[f"AddNet Weight {i+1}"]
+            params[f"AddNet Weight B {i+1}"] = params[f"AddNet Weight {i+1}"]
+
+        if f"AddNet Module {i+1}" not in params:
+            params[f"AddNet Module {i+1}"] = "LoRA"
+        if f"AddNet Model {i+1}" not in params:
+            params[f"AddNet Model {i+1}"] = "None"
+        if f"AddNet Weight A {i+1}" not in params:
+            params[f"AddNet Weight A {i+1}"] = "0"
+        if f"AddNet Weight B {i+1}" not in params:
+            params[f"AddNet Weight B {i+1}"] = "0"
+
+        params[f"AddNet Weight {i+1}"] = params[f"AddNet Weight A {i+1}"]
+
+        if params[f"AddNet Weight A {i+1}"] != params[f"AddNet Weight B {i+1}"]:
+            params["AddNet Separate Weights"] = "True"
+
+        # Convert potential legacy name/hash to new format
+        params[f"AddNet Model {i+1}"] = str(model_util.find_closest_lora_model_name(params[f"AddNet Model {i+1}"]))
+
+        addnet_xyz_grid_support.update_axis_params(i, params[f"AddNet Module {i+1}"], params[f"AddNet Model {i+1}"])
+
+
+addnet_xyz_grid_support.initialize(Script)
+
+
+script_callbacks.on_script_unloaded(on_script_unloaded)
+script_callbacks.on_ui_tabs(on_ui_tabs)
+script_callbacks.on_ui_settings(on_ui_settings)
+script_callbacks.on_infotext_pasted(on_infotext_pasted)

extensions/addtional/scripts/addnet_xyz_grid_support.py

@@ -0,0 +1,189 @@
+import os
+import os.path
+from modules import shared
+import modules.scripts as scripts
+from scripts import model_util, util
+from scripts.model_util import MAX_MODEL_COUNT
+
+
+LORA_TRAIN_METADATA_NAMES = {
+    "ss_session_id": "Session ID",
+    "ss_training_started_at": "Training started at",
+    "ss_output_name": "Output name",
+    "ss_learning_rate": "Learning rate",
+    "ss_text_encoder_lr": "Text encoder LR",
+    "ss_unet_lr": "UNet LR",
+    "ss_num_train_images": "# of training images",
+    "ss_num_reg_images": "# of reg images",
+    "ss_num_batches_per_epoch": "Batches per epoch",
+    "ss_num_epochs": "Total epochs",
+    "ss_epoch": "Epoch",
+    "ss_batch_size_per_device": "Batch size/device",
+    "ss_total_batch_size": "Total batch size",
+    "ss_gradient_checkpointing": "Gradient checkpointing",
+    "ss_gradient_accumulation_steps": "Gradient accum. steps",
+    "ss_max_train_steps": "Max train steps",
+    "ss_lr_warmup_steps": "LR warmup steps",
+    "ss_lr_scheduler": "LR scheduler",
+    "ss_network_module": "Network module",
+    "ss_network_dim": "Network dim",
+    "ss_network_alpha": "Network alpha",
+    "ss_mixed_precision": "Mixed precision",
+    "ss_full_fp16": "Full FP16",
+    "ss_v2": "V2",
+    "ss_resolution": "Resolution",
+    "ss_clip_skip": "Clip skip",
+    "ss_max_token_length": "Max token length",
+    "ss_color_aug": "Color aug",
+    "ss_flip_aug": "Flip aug",
+    "ss_random_crop": "Random crop",
+    "ss_shuffle_caption": "Shuffle caption",
+    "ss_cache_latents": "Cache latents",
+    "ss_enable_bucket": "Enable bucket",
+    "ss_min_bucket_reso": "Min bucket reso.",
+    "ss_max_bucket_reso": "Max bucket reso.",
+    "ss_seed": "Seed",
+    "ss_keep_tokens": "Keep tokens",
+    "ss_dataset_dirs": "Dataset dirs.",
+    "ss_reg_dataset_dirs": "Reg dataset dirs.",
+    "ss_sd_model_name": "SD model name",
+    "ss_vae_name": "VAE name",
+    "ss_training_comment": "Comment",
+}
+
+
+xy_grid = None  # XY Grid module
+script_class = None  # additional_networks scripts.Script class
+axis_params = [{}] * MAX_MODEL_COUNT
+
+
+def update_axis_params(i, module, model):
+    axis_params[i] = {"module": module, "model": model}
+
+
+def get_axis_model_choices(i):
+    module = axis_params[i].get("module", "None")
+    model = axis_params[i].get("model", "None")
+
+    if module == "LoRA":
+        if model != "None":
+            sort_by = shared.opts.data.get("additional_networks_sort_models_by", "name")
+            return ["None"] + model_util.get_model_list(module, model, "", sort_by)
+
+    return [f"select `Model {i+1}` in `Additional Networks`. models in same folder for selected one will be shown here."]
+
+
+def update_script_args(p, value, arg_idx):
+    global script_class
+    for s in scripts.scripts_txt2img.alwayson_scripts:
+        if isinstance(s, script_class):
+            args = list(p.script_args)
+            # print(f"Changed arg {arg_idx} from {args[s.args_from + arg_idx - 1]} to {value}")
+            args[s.args_from + arg_idx] = value
+            p.script_args = tuple(args)
+            break
+
+
+def confirm_models(p, xs):
+    for x in xs:
+        if x in ["", "None"]:
+            continue
+        if not model_util.find_closest_lora_model_name(x):
+            raise RuntimeError(f"Unknown LoRA model: {x}")
+
+
+def apply_module(p, x, xs, i):
+    update_script_args(p, True, 0)  # set Enabled to True
+    update_script_args(p, x, 2 + 4 * i)  # enabled, separate_weights, ({module}, model, weight_unet, weight_tenc), ...
+
+
+def apply_model(p, x, xs, i):
+    name = model_util.find_closest_lora_model_name(x)
+    update_script_args(p, True, 0)
+    update_script_args(p, name, 3 + 4 * i)  # enabled, separate_weights, (module, {model}, weight_unet, weight_tenc), ...
+
+
+def apply_weight(p, x, xs, i):
+    update_script_args(p, True, 0)
+    update_script_args(p, x, 4 + 4 * i)  # enabled, separate_weights, (module, model, {weight_unet, weight_tenc}), ...
+    update_script_args(p, x, 5 + 4 * i)
+
+
+def apply_weight_unet(p, x, xs, i):
+    update_script_args(p, True, 0)
+    update_script_args(p, x, 4 + 4 * i)  # enabled, separate_weights, (module, model, {weight_unet}, weight_tenc), ...
+
+
+def apply_weight_tenc(p, x, xs, i):
+    update_script_args(p, True, 0)
+    update_script_args(p, x, 5 + 4 * i)  # enabled, separate_weights, (module, model, weight_unet, {weight_tenc}), ...
+
+
+def format_lora_model(p, opt, x):
+    global xy_grid
+    model = model_util.find_closest_lora_model_name(x)
+    if model is None or model.lower() in ["", "none"]:
+        return "None"
+
+    value = xy_grid.format_value(p, opt, model)
+
+    model_path = model_util.lora_models.get(model)
+    metadata = model_util.read_model_metadata(model_path, "LoRA")
+    if not metadata:
+        return value
+
+    metadata_names = util.split_path_list(shared.opts.data.get("additional_networks_xy_grid_model_metadata", ""))
+    if not metadata_names:
+        return value
+
+    for name in metadata_names:
+        name = name.strip()
+        if name in metadata:
+            formatted_name = LORA_TRAIN_METADATA_NAMES.get(name, name)
+            value += f"\n{formatted_name}: {metadata[name]}, "
+
+    return value.strip(" ").strip(",")
+
+
+def initialize(script):
+    global xy_grid, script_class
+    xy_grid = None
+    script_class = script
+    for scriptDataTuple in scripts.scripts_data:
+        if os.path.basename(scriptDataTuple.path) == "xy_grid.py" or os.path.basename(scriptDataTuple.path) == "xyz_grid.py":
+            xy_grid = scriptDataTuple.module
+            for i in range(MAX_MODEL_COUNT):
+                model = xy_grid.AxisOption(
+                    f"AddNet Model {i+1}",
+                    str,
+                    lambda p, x, xs, i=i: apply_model(p, x, xs, i),
+                    format_lora_model,
+                    confirm_models,
+                    cost=0.5,
+                    choices=lambda i=i: get_axis_model_choices(i),
+                )
+                weight = xy_grid.AxisOption(
+                    f"AddNet Weight {i+1}",
+                    float,
+                    lambda p, x, xs, i=i: apply_weight(p, x, xs, i),
+                    xy_grid.format_value_add_label,
+                    None,
+                    cost=0.5,
+                )
+                weight_unet = xy_grid.AxisOption(
+                    f"AddNet UNet Weight {i+1}",
+                    float,
+                    lambda p, x, xs, i=i: apply_weight_unet(p, x, xs, i),
+                    xy_grid.format_value_add_label,
+                    None,
+                    cost=0.5,
+                )
+                weight_tenc = xy_grid.AxisOption(
+                    f"AddNet TEnc Weight {i+1}",
+                    float,
+                    lambda p, x, xs, i=i: apply_weight_tenc(p, x, xs, i),
+                    xy_grid.format_value_add_label,
+                    None,
+                    cost=0.5,
+                )
+                xy_grid.axis_options.extend([model, weight, weight_unet, weight_tenc])

extensions/addtional/scripts/lora_compvis.py

@@ -0,0 +1,634 @@
+# LoRA network module
+# reference:
+# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
+# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
+
+import copy
+import math
+import re
+from typing import NamedTuple
+import torch
+
+
+class LoRAInfo(NamedTuple):
+    lora_name: str
+    module_name: str
+    module: torch.nn.Module
+    multiplier: float
+    dim: int
+    alpha: float
+
+
+class LoRAModule(torch.nn.Module):
+    """
+    replaces forward method of the original Linear, instead of replacing the original Linear module.
+    """
+
+    def __init__(self, lora_name, org_module: torch.nn.Module, multiplier=1.0, lora_dim=4, alpha=1):
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
+        super().__init__()
+        self.lora_name = lora_name
+        self.lora_dim = lora_dim
+
+        if org_module.__class__.__name__ == "Conv2d":
+            in_dim = org_module.in_channels
+            out_dim = org_module.out_channels
+
+            # self.lora_dim = min(self.lora_dim, in_dim, out_dim)
+            # if self.lora_dim != lora_dim:
+            #   print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
+
+            kernel_size = org_module.kernel_size
+            stride = org_module.stride
+            padding = org_module.padding
+            self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
+            self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
+        else:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
+            self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
+
+        if type(alpha) == torch.Tensor:
+            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+        alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
+        self.scale = alpha / self.lora_dim
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
+
+        # same as microsoft's
+        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.lora_up.weight)
+
+        self.multiplier = multiplier
+        self.org_forward = org_module.forward
+        self.org_module = org_module  # remove in applying
+        self.mask_dic = None
+        self.mask = None
+        self.mask_area = -1
+
+    def apply_to(self):
+        self.org_forward = self.org_module.forward
+        self.org_module.forward = self.forward
+        del self.org_module
+
+    def set_mask_dic(self, mask_dic):
+        # called before every generation
+
+        # check this module is related to h,w (not context and time emb)
+        if "attn2_to_k" in self.lora_name or "attn2_to_v" in self.lora_name or "emb_layers" in self.lora_name:
+            # print(f"LoRA for context or time emb: {self.lora_name}")
+            self.mask_dic = None
+        else:
+            self.mask_dic = mask_dic
+
+        self.mask = None
+
+    def forward(self, x):
+        """
+        may be cascaded.
+        """
+        if self.mask_dic is None:
+            return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+
+        # regional LoRA
+
+        # calculate lora and get size
+        lx = self.lora_up(self.lora_down(x))
+
+        if len(lx.size()) == 4:  # b,c,h,w
+            area = lx.size()[2] * lx.size()[3]
+        else:
+            area = lx.size()[1]  # b,seq,dim
+
+        if self.mask is None or self.mask_area != area:
+            # get mask
+            # print(self.lora_name, x.size(), lx.size(), area)
+            mask = self.mask_dic[area]
+            if len(lx.size()) == 3:
+                mask = torch.reshape(mask, (1, -1, 1))
+            self.mask = mask
+            self.mask_area = area
+
+        return self.org_forward(x) + lx * self.multiplier * self.scale * self.mask
+
+
+def create_network_and_apply_compvis(du_state_dict, multiplier_tenc, multiplier_unet, text_encoder, unet, **kwargs):
+    # get device and dtype from unet
+    for module in unet.modules():
+        if module.__class__.__name__ == "Linear":
+            param: torch.nn.Parameter = module.weight
+            # device = param.device
+            dtype = param.dtype
+            break
+
+    # get dims (rank) and alpha from state dict
+    modules_dim = {}
+    modules_alpha = {}
+    for key, value in du_state_dict.items():
+        if "." not in key:
+            continue
+
+        lora_name = key.split(".")[0]
+        if "alpha" in key:
+            modules_alpha[lora_name] = float(value.detach().to(torch.float).cpu().numpy())
+        elif "lora_down" in key:
+            dim = value.size()[0]
+            modules_dim[lora_name] = dim
+
+    # support old LoRA without alpha
+    for key in modules_dim.keys():
+        if key not in modules_alpha:
+            modules_alpha[key] = modules_dim[key]
+
+    print(
+        f"dimension: {set(modules_dim.values())}, alpha: {set(modules_alpha.values())}, multiplier_unet: {multiplier_unet}, multiplier_tenc: {multiplier_tenc}"
+    )
+
+    # if network_dim is None:
+    #   print(f"The selected model is not LoRA or not trained by `sd-scripts`?")
+    #   network_dim = 4
+    #   network_alpha = 1
+
+    # create, apply and load weights
+    network = LoRANetworkCompvis(text_encoder, unet, multiplier_tenc, multiplier_unet, modules_dim, modules_alpha)
+    state_dict = network.apply_lora_modules(du_state_dict)  # some weights are applied to text encoder
+    network.to(dtype)  # with this, if error comes from next line, the model will be used
+    info = network.load_state_dict(state_dict, strict=False)
+
+    # remove redundant warnings
+    if len(info.missing_keys) > 4:
+        missing_keys = []
+        alpha_count = 0
+        for key in info.missing_keys:
+            if "alpha" not in key:
+                missing_keys.append(key)
+            else:
+                if alpha_count == 0:
+                    missing_keys.append(key)
+                alpha_count += 1
+        if alpha_count > 1:
+            missing_keys.append(
+                f"... and {alpha_count-1} alphas. The model doesn't have alpha, use dim (rannk) as alpha. You can ignore this message."
+            )
+
+        info = torch.nn.modules.module._IncompatibleKeys(missing_keys, info.unexpected_keys)
+
+    return network, info
+
+
+class LoRANetworkCompvis(torch.nn.Module):
+    # UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
+    # TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+    UNET_TARGET_REPLACE_MODULE = ["SpatialTransformer", "ResBlock", "Downsample", "Upsample"]  # , "Attention"]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["ResidualAttentionBlock", "CLIPAttention", "CLIPMLP"]
+
+    LORA_PREFIX_UNET = "lora_unet"
+    LORA_PREFIX_TEXT_ENCODER = "lora_te"
+
+    @classmethod
+    def convert_diffusers_name_to_compvis(cls, v2, du_name):
+        """
+        convert diffusers's LoRA name to CompVis
+        """
+        cv_name = None
+        if "lora_unet_" in du_name:
+            m = re.search(r"_down_blocks_(\d+)_attentions_(\d+)_(.+)", du_name)
+            if m:
+                du_block_index = int(m.group(1))
+                du_attn_index = int(m.group(2))
+                du_suffix = m.group(3)
+
+                cv_index = 1 + du_block_index * 3 + du_attn_index  # 1,2, 4,5, 7,8
+                cv_name = f"lora_unet_input_blocks_{cv_index}_1_{du_suffix}"
+                return cv_name
+
+            m = re.search(r"_mid_block_attentions_(\d+)_(.+)", du_name)
+            if m:
+                du_suffix = m.group(2)
+                cv_name = f"lora_unet_middle_block_1_{du_suffix}"
+                return cv_name
+
+            m = re.search(r"_up_blocks_(\d+)_attentions_(\d+)_(.+)", du_name)
+            if m:
+                du_block_index = int(m.group(1))
+                du_attn_index = int(m.group(2))
+                du_suffix = m.group(3)
+
+                cv_index = du_block_index * 3 + du_attn_index  # 3,4,5, 6,7,8, 9,10,11
+                cv_name = f"lora_unet_output_blocks_{cv_index}_1_{du_suffix}"
+                return cv_name
+
+            m = re.search(r"_down_blocks_(\d+)_resnets_(\d+)_(.+)", du_name)
+            if m:
+                du_block_index = int(m.group(1))
+                du_res_index = int(m.group(2))
+                du_suffix = m.group(3)
+                cv_suffix = {
+                    "conv1": "in_layers_2",
+                    "conv2": "out_layers_3",
+                    "time_emb_proj": "emb_layers_1",
+                    "conv_shortcut": "skip_connection",
+                }[du_suffix]
+
+                cv_index = 1 + du_block_index * 3 + du_res_index  # 1,2, 4,5, 7,8
+                cv_name = f"lora_unet_input_blocks_{cv_index}_0_{cv_suffix}"
+                return cv_name
+
+            m = re.search(r"_down_blocks_(\d+)_downsamplers_0_conv", du_name)
+            if m:
+                block_index = int(m.group(1))
+                cv_index = 3 + block_index * 3
+                cv_name = f"lora_unet_input_blocks_{cv_index}_0_op"
+                return cv_name
+
+            m = re.search(r"_mid_block_resnets_(\d+)_(.+)", du_name)
+            if m:
+                index = int(m.group(1))
+                du_suffix = m.group(2)
+                cv_suffix = {
+                    "conv1": "in_layers_2",
+                    "conv2": "out_layers_3",
+                    "time_emb_proj": "emb_layers_1",
+                    "conv_shortcut": "skip_connection",
+                }[du_suffix]
+                cv_name = f"lora_unet_middle_block_{index*2}_{cv_suffix}"
+                return cv_name
+
+            m = re.search(r"_up_blocks_(\d+)_resnets_(\d+)_(.+)", du_name)
+            if m:
+                du_block_index = int(m.group(1))
+                du_res_index = int(m.group(2))
+                du_suffix = m.group(3)
+                cv_suffix = {
+                    "conv1": "in_layers_2",
+                    "conv2": "out_layers_3",
+                    "time_emb_proj": "emb_layers_1",
+                    "conv_shortcut": "skip_connection",
+                }[du_suffix]
+
+                cv_index = du_block_index * 3 + du_res_index  # 1,2, 4,5, 7,8
+                cv_name = f"lora_unet_output_blocks_{cv_index}_0_{cv_suffix}"
+                return cv_name
+
+            m = re.search(r"_up_blocks_(\d+)_upsamplers_0_conv", du_name)
+            if m:
+                block_index = int(m.group(1))
+                cv_index = block_index * 3 + 2
+                cv_name = f"lora_unet_output_blocks_{cv_index}_{bool(block_index)+1}_conv"
+                return cv_name
+
+        elif "lora_te_" in du_name:
+            m = re.search(r"_model_encoder_layers_(\d+)_(.+)", du_name)
+            if m:
+                du_block_index = int(m.group(1))
+                du_suffix = m.group(2)
+
+                cv_index = du_block_index
+                if v2:
+                    if "mlp_fc1" in du_suffix:
+                        cv_name = (
+                            f"lora_te_wrapped_model_transformer_resblocks_{cv_index}_{du_suffix.replace('mlp_fc1', 'mlp_c_fc')}"
+                        )
+                    elif "mlp_fc2" in du_suffix:
+                        cv_name = (
+                            f"lora_te_wrapped_model_transformer_resblocks_{cv_index}_{du_suffix.replace('mlp_fc2', 'mlp_c_proj')}"
+                        )
+                    elif "self_attn":
+                        # handled later
+                        cv_name = f"lora_te_wrapped_model_transformer_resblocks_{cv_index}_{du_suffix.replace('self_attn', 'attn')}"
+                else:
+                    cv_name = f"lora_te_wrapped_transformer_text_model_encoder_layers_{cv_index}_{du_suffix}"
+
+        assert cv_name is not None, f"conversion failed: {du_name}. the model may not be trained by `sd-scripts`."
+        return cv_name
+
+    @classmethod
+    def convert_state_dict_name_to_compvis(cls, v2, state_dict):
+        """
+        convert keys in state dict to load it by load_state_dict
+        """
+        new_sd = {}
+        for key, value in state_dict.items():
+            tokens = key.split(".")
+            compvis_name = LoRANetworkCompvis.convert_diffusers_name_to_compvis(v2, tokens[0])
+            new_key = compvis_name + "." + ".".join(tokens[1:])
+
+            new_sd[new_key] = value
+
+        return new_sd
+
+    def __init__(self, text_encoder, unet, multiplier_tenc=1.0, multiplier_unet=1.0, modules_dim=None, modules_alpha=None) -> None:
+        super().__init__()
+        self.multiplier_unet = multiplier_unet
+        self.multiplier_tenc = multiplier_tenc
+        self.latest_mask_info = None
+
+        # check v1 or v2
+        self.v2 = False
+        for _, module in text_encoder.named_modules():
+            for _, child_module in module.named_modules():
+                if child_module.__class__.__name__ == "MultiheadAttention":
+                    self.v2 = True
+                    break
+            if self.v2:
+                break
+
+        # convert lora name to CompVis and get dim and alpha
+        comp_vis_loras_dim_alpha = {}
+        for du_lora_name in modules_dim.keys():
+            dim = modules_dim[du_lora_name]
+            alpha = modules_alpha[du_lora_name]
+            comp_vis_lora_name = LoRANetworkCompvis.convert_diffusers_name_to_compvis(self.v2, du_lora_name)
+            comp_vis_loras_dim_alpha[comp_vis_lora_name] = (dim, alpha)
+
+        # create module instances
+        def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules, multiplier):
+            loras = []
+            replaced_modules = []
+            for name, module in root_module.named_modules():
+                if module.__class__.__name__ in target_replace_modules:
+                    for child_name, child_module in module.named_modules():
+                        # enumerate all Linear and Conv2d
+                        if child_module.__class__.__name__ == "Linear" or child_module.__class__.__name__ == "Conv2d":
+                            lora_name = prefix + "." + name + "." + child_name
+                            lora_name = lora_name.replace(".", "_")
+                            if "_resblocks_23_" in lora_name:  # ignore last block in StabilityAi Text Encoder
+                                break
+                            if lora_name not in comp_vis_loras_dim_alpha:
+                                continue
+
+                            dim, alpha = comp_vis_loras_dim_alpha[lora_name]
+                            lora = LoRAModule(lora_name, child_module, multiplier, dim, alpha)
+                            loras.append(lora)
+
+                            replaced_modules.append(child_module)
+                        elif child_module.__class__.__name__ == "MultiheadAttention":
+                            # make four modules: not replacing forward method but merge weights later
+                            for suffix in ["q_proj", "k_proj", "v_proj", "out_proj"]:
+                                module_name = prefix + "." + name + "." + child_name  # ~.attn
+                                module_name = module_name.replace(".", "_")
+                                if "_resblocks_23_" in module_name:  # ignore last block in StabilityAi Text Encoder
+                                    break
+
+                                lora_name = module_name + "_" + suffix
+                                if lora_name not in comp_vis_loras_dim_alpha:
+                                    continue
+                                dim, alpha = comp_vis_loras_dim_alpha[lora_name]
+                                lora_info = LoRAInfo(lora_name, module_name, child_module, multiplier, dim, alpha)
+                                loras.append(lora_info)
+
+                                replaced_modules.append(child_module)
+            return loras, replaced_modules
+
+        self.text_encoder_loras, te_rep_modules = create_modules(
+            LoRANetworkCompvis.LORA_PREFIX_TEXT_ENCODER,
+            text_encoder,
+            LoRANetworkCompvis.TEXT_ENCODER_TARGET_REPLACE_MODULE,
+            self.multiplier_tenc,
+        )
+        print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+
+        self.unet_loras, unet_rep_modules = create_modules(
+            LoRANetworkCompvis.LORA_PREFIX_UNET, unet, LoRANetworkCompvis.UNET_TARGET_REPLACE_MODULE, self.multiplier_unet
+        )
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+
+        # make backup of original forward/weights, if multiple modules are applied, do in 1st module only
+        backed_up = False  # messaging purpose only
+        for rep_module in te_rep_modules + unet_rep_modules:
+            if (
+                rep_module.__class__.__name__ == "MultiheadAttention"
+            ):  # multiple MHA modules are in list, prevent to backed up forward
+                if not hasattr(rep_module, "_lora_org_weights"):
+                    # avoid updating of original weights. state_dict is reference to original weights
+                    rep_module._lora_org_weights = copy.deepcopy(rep_module.state_dict())
+                    backed_up = True
+            elif not hasattr(rep_module, "_lora_org_forward"):
+                rep_module._lora_org_forward = rep_module.forward
+                backed_up = True
+        if backed_up:
+            print("original forward/weights is backed up.")
+
+        # assertion
+        names = set()
+        for lora in self.text_encoder_loras + self.unet_loras:
+            assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
+            names.add(lora.lora_name)
+
+    def restore(self, text_encoder, unet):
+        # restore forward/weights from property for all modules
+        restored = False  # messaging purpose only
+        modules = []
+        modules.extend(text_encoder.modules())
+        modules.extend(unet.modules())
+        for module in modules:
+            if hasattr(module, "_lora_org_forward"):
+                module.forward = module._lora_org_forward
+                del module._lora_org_forward
+                restored = True
+            if hasattr(
+                module, "_lora_org_weights"
+            ):  # module doesn't have forward and weights at same time currently, but supports it for future changing
+                module.load_state_dict(module._lora_org_weights)
+                del module._lora_org_weights
+                restored = True
+
+        if restored:
+            print("original forward/weights is restored.")
+
+    def apply_lora_modules(self, du_state_dict):
+        # conversion 1st step: convert names in state_dict
+        state_dict = LoRANetworkCompvis.convert_state_dict_name_to_compvis(self.v2, du_state_dict)
+
+        # check state_dict has text_encoder or unet
+        weights_has_text_encoder = weights_has_unet = False
+        for key in state_dict.keys():
+            if key.startswith(LoRANetworkCompvis.LORA_PREFIX_TEXT_ENCODER):
+                weights_has_text_encoder = True
+            elif key.startswith(LoRANetworkCompvis.LORA_PREFIX_UNET):
+                weights_has_unet = True
+            if weights_has_text_encoder and weights_has_unet:
+                break
+
+        apply_text_encoder = weights_has_text_encoder
+        apply_unet = weights_has_unet
+
+        if apply_text_encoder:
+            print("enable LoRA for text encoder")
+        else:
+            self.text_encoder_loras = []
+
+        if apply_unet:
+            print("enable LoRA for U-Net")
+        else:
+            self.unet_loras = []
+
+        # add modules to network: this makes state_dict can be got from LoRANetwork
+        mha_loras = {}
+        for lora in self.text_encoder_loras + self.unet_loras:
+            if type(lora) == LoRAModule:
+                lora.apply_to()  # ensure remove reference to original Linear: reference makes key of state_dict
+                self.add_module(lora.lora_name, lora)
+            else:
+                # SD2.x MultiheadAttention merge weights to MHA weights
+                lora_info: LoRAInfo = lora
+                if lora_info.module_name not in mha_loras:
+                    mha_loras[lora_info.module_name] = {}
+
+                lora_dic = mha_loras[lora_info.module_name]
+                lora_dic[lora_info.lora_name] = lora_info
+                if len(lora_dic) == 4:
+                    # calculate and apply
+                    module = lora_info.module
+                    module_name = lora_info.module_name
+                    w_q_dw = state_dict.get(module_name + "_q_proj.lora_down.weight")
+                    if w_q_dw is not None:  # corresponding LoRA module exists
+                        w_q_up = state_dict[module_name + "_q_proj.lora_up.weight"]
+                        w_k_dw = state_dict[module_name + "_k_proj.lora_down.weight"]
+                        w_k_up = state_dict[module_name + "_k_proj.lora_up.weight"]
+                        w_v_dw = state_dict[module_name + "_v_proj.lora_down.weight"]
+                        w_v_up = state_dict[module_name + "_v_proj.lora_up.weight"]
+                        w_out_dw = state_dict[module_name + "_out_proj.lora_down.weight"]
+                        w_out_up = state_dict[module_name + "_out_proj.lora_up.weight"]
+                        q_lora_info = lora_dic[module_name + "_q_proj"]
+                        k_lora_info = lora_dic[module_name + "_k_proj"]
+                        v_lora_info = lora_dic[module_name + "_v_proj"]
+                        out_lora_info = lora_dic[module_name + "_out_proj"]
+
+                        sd = module.state_dict()
+                        qkv_weight = sd["in_proj_weight"]
+                        out_weight = sd["out_proj.weight"]
+                        dev = qkv_weight.device
+
+                        def merge_weights(l_info, weight, up_weight, down_weight):
+                            # calculate in float
+                            scale = l_info.alpha / l_info.dim
+                            dtype = weight.dtype
+                            weight = (
+                                weight.float()
+                                + l_info.multiplier
+                                * (up_weight.to(dev, dtype=torch.float) @ down_weight.to(dev, dtype=torch.float))
+                                * scale
+                            )
+                            weight = weight.to(dtype)
+                            return weight
+
+                        q_weight, k_weight, v_weight = torch.chunk(qkv_weight, 3)
+                        if q_weight.size()[1] == w_q_up.size()[0]:
+                            q_weight = merge_weights(q_lora_info, q_weight, w_q_up, w_q_dw)
+                            k_weight = merge_weights(k_lora_info, k_weight, w_k_up, w_k_dw)
+                            v_weight = merge_weights(v_lora_info, v_weight, w_v_up, w_v_dw)
+                            qkv_weight = torch.cat([q_weight, k_weight, v_weight])
+
+                            out_weight = merge_weights(out_lora_info, out_weight, w_out_up, w_out_dw)
+
+                            sd["in_proj_weight"] = qkv_weight.to(dev)
+                            sd["out_proj.weight"] = out_weight.to(dev)
+
+                            lora_info.module.load_state_dict(sd)
+                        else:
+                            # different dim, version mismatch
+                            print(f"shape of weight is different: {module_name}. SD version may be different")
+
+                        for t in ["q", "k", "v", "out"]:
+                            del state_dict[f"{module_name}_{t}_proj.lora_down.weight"]
+                            del state_dict[f"{module_name}_{t}_proj.lora_up.weight"]
+                            alpha_key = f"{module_name}_{t}_proj.alpha"
+                            if alpha_key in state_dict:
+                                del state_dict[alpha_key]
+                    else:
+                        # corresponding weight not exists: version mismatch
+                        pass
+
+        # conversion 2nd step: convert weight's shape (and handle wrapped)
+        state_dict = self.convert_state_dict_shape_to_compvis(state_dict)
+
+        return state_dict
+
+    def convert_state_dict_shape_to_compvis(self, state_dict):
+        # shape conversion
+        current_sd = self.state_dict()  # to get target shape
+        wrapped = False
+        count = 0
+        for key in list(state_dict.keys()):
+            if key not in current_sd:
+                continue  # might be error or another version
+            if "wrapped" in key:
+                wrapped = True
+
+            value: torch.Tensor = state_dict[key]
+            if value.size() != current_sd[key].size():
+                # print(f"convert weights shape: {key}, from: {value.size()}, {len(value.size())}")
+                count += 1
+                if len(value.size()) == 4:
+                    value = value.squeeze(3).squeeze(2)
+                else:
+                    value = value.unsqueeze(2).unsqueeze(3)
+                state_dict[key] = value
+            if tuple(value.size()) != tuple(current_sd[key].size()):
+                print(
+                    f"weight's shape is different: {key} expected {current_sd[key].size()} found {value.size()}. SD version may be different"
+                )
+                del state_dict[key]
+        print(f"shapes for {count} weights are converted.")
+
+        # convert wrapped
+        if not wrapped:
+            print("remove 'wrapped' from keys")
+            for key in list(state_dict.keys()):
+                if "_wrapped_" in key:
+                    new_key = key.replace("_wrapped_", "_")
+                    state_dict[new_key] = state_dict[key]
+                    del state_dict[key]
+
+        return state_dict
+
+    def set_mask(self, mask, height=None, width=None, hr_height=None, hr_width=None):
+        if mask is None:
+            # clear latest mask
+            # print("clear mask")
+            self.latest_mask_info = None
+            for lora in self.unet_loras:
+                lora.set_mask_dic(None)
+            return
+
+        # check mask image and h/w are same
+        if (
+            self.latest_mask_info is not None
+            and torch.equal(mask, self.latest_mask_info[0])
+            and (height, width, hr_height, hr_width) == self.latest_mask_info[1:]
+        ):
+            # print("mask not changed")
+            return
+
+        self.latest_mask_info = (mask, height, width, hr_height, hr_width)
+
+        org_dtype = mask.dtype
+        if mask.dtype == torch.bfloat16:
+            mask = mask.to(torch.float)
+
+        mask_dic = {}
+        mask = mask.unsqueeze(0).unsqueeze(1)  # b(1),c(1),h,w
+
+        def resize_add(mh, mw):
+            # print(mh, mw, mh * mw)
+            m = torch.nn.functional.interpolate(mask, (mh, mw), mode="bilinear")  # doesn't work in bf16
+            m = m.to(org_dtype)
+            mask_dic[mh * mw] = m
+
+        for h, w in [(height, width), (hr_height, hr_width)]:
+            if not h or not w:
+                continue
+
+            h = h // 8
+            w = w // 8
+            for i in range(4):
+                resize_add(h, w)
+                if h % 2 == 1 or w % 2 == 1:  # add extra shape if h/w is not divisible by 2
+                    resize_add(h + h % 2, w + w % 2)
+                h = (h + 1) // 2
+                w = (w + 1) // 2
+
+        for lora in self.unet_loras:
+            lora.set_mask_dic(mask_dic)
+        return

extensions/addtional/scripts/metadata_editor.py

@@ -0,0 +1,622 @@
+import os
+import json
+import sys
+import io
+import base64
+import platform
+import subprocess as sp
+from PIL import PngImagePlugin, Image
+
+from modules import shared
+import gradio as gr
+
+import modules.ui
+from modules.ui_components import ToolButton
+import modules.extras
+import modules.generation_parameters_copypaste as parameters_copypaste
+
+from scripts import safetensors_hack, model_util
+from scripts.model_util import MAX_MODEL_COUNT
+
+
+folder_symbol = "\U0001f4c2"  # 📂
+keycap_symbols = [
+    "\u0031\ufe0f\u20e3",  # 1️⃣
+    "\u0032\ufe0f\u20e3",  # 2️⃣
+    "\u0033\ufe0f\u20e3",  # 3️⃣
+    "\u0034\ufe0f\u20e3",  # 4️⃣
+    "\u0035\ufe0f\u20e3",  # 5️⃣
+    "\u0036\ufe0f\u20e3",  # 6️⃣
+    "\u0037\ufe0f\u20e3",  # 7️⃣
+    "\u0038\ufe0f\u20e3",  # 8️
+    "\u0039\ufe0f\u20e3",  # 9️
+    "\u1f51f",  # 🔟
+]
+
+
+def write_webui_model_preview_image(model_path, image):
+    basename, ext = os.path.splitext(model_path)
+    preview_path = f"{basename}.png"
+
+    # Copy any text-only metadata
+    use_metadata = False
+    metadata = PngImagePlugin.PngInfo()
+    for key, value in image.info.items():
+        if isinstance(key, str) and isinstance(value, str):
+            metadata.add_text(key, value)
+            use_metadata = True
+
+    image.save(preview_path, "PNG", pnginfo=(metadata if use_metadata else None))
+
+
+def delete_webui_model_preview_image(model_path):
+    basename, ext = os.path.splitext(model_path)
+    preview_paths = [f"{basename}.preview.png", f"{basename}.png"]
+
+    for preview_path in preview_paths:
+        if os.path.isfile(preview_path):
+            os.unlink(preview_path)
+
+
+def decode_base64_to_pil(encoding):
+    if encoding.startswith("data:image/"):
+        encoding = encoding.split(";")[1].split(",")[1]
+    return Image.open(io.BytesIO(base64.b64decode(encoding)))
+
+
+def encode_pil_to_base64(image):
+    with io.BytesIO() as output_bytes:
+        # Copy any text-only metadata
+        use_metadata = False
+        metadata = PngImagePlugin.PngInfo()
+        for key, value in image.info.items():
+            if isinstance(key, str) and isinstance(value, str):
+                metadata.add_text(key, value)
+                use_metadata = True
+
+        image.save(output_bytes, "PNG", pnginfo=(metadata if use_metadata else None))
+        bytes_data = output_bytes.getvalue()
+    return base64.b64encode(bytes_data)
+
+
+def open_folder(f):
+    if not os.path.exists(f):
+        print(f'Folder "{f}" does not exist. After you create an image, the folder will be created.')
+        return
+    elif not os.path.isdir(f):
+        print(
+            f"""
+WARNING
+An open_folder request was made with an argument that is not a folder.
+This could be an error or a malicious attempt to run code on your computer.
+Requested path was: {f}
+""",
+            file=sys.stderr,
+        )
+        return
+
+    if not shared.cmd_opts.hide_ui_dir_config:
+        path = os.path.normpath(f)
+        if platform.system() == "Windows":
+            os.startfile(path)
+        elif platform.system() == "Darwin":
+            sp.Popen(["open", path])
+        elif "microsoft-standard-WSL2" in platform.uname().release:
+            sp.Popen(["wsl-open", path])
+        else:
+            sp.Popen(["xdg-open", path])
+
+
+def copy_metadata_to_all(module, model_path, copy_dir, same_session_only, missing_meta_only, cover_image):
+    """
+    Given a model with metadata, copies that metadata to all models in copy_dir.
+
+    :str module: Module name ("LoRA")
+    :str model: Model key in lora_models ("MyModel(123456abcdef)")
+    :str copy_dir: Directory to copy to
+    :bool same_session_only: Only copy to modules with the same ss_session_id
+    :bool missing_meta_only: Only copy to modules that are missing user metadata
+    :Optional[Image] cover_image: Cover image to embed in the file as base64
+    :returns: gr.HTML.update()
+    """
+    if model_path == "None":
+        return "No model selected."
+
+    if not os.path.isfile(model_path):
+        return f"Model path not found: {model_path}"
+
+    model_path = os.path.realpath(model_path)
+
+    if os.path.splitext(model_path)[1] != ".safetensors":
+        return "Model is not in .safetensors format."
+
+    if not os.path.isdir(copy_dir):
+        return "Please provide a directory containing models in .safetensors format."
+
+    print(f"[MetadataEditor] Copying metadata to models in {copy_dir}.")
+    metadata = model_util.read_model_metadata(model_path, module)
+    count = 0
+    for entry in os.scandir(copy_dir):
+        if entry.is_file():
+            path = os.path.realpath(os.path.join(copy_dir, entry.name))
+            if path != model_path and model_util.is_safetensors(path):
+                if same_session_only:
+                    other_metadata = safetensors_hack.read_metadata(path)
+                    if missing_meta_only and other_metadata.get("ssmd_display_name", "").strip():
+                        print(f"[MetadataEditor] Skipping {path} as it already has metadata")
+                        continue
+
+                    session_id = metadata.get("ss_session_id", None)
+                    other_session_id = other_metadata.get("ss_session_id", None)
+                    if session_id is None or other_session_id is None or session_id != other_session_id:
+                        continue
+
+                updates = {
+                    "ssmd_cover_images": "[]",
+                    "ssmd_display_name": "",
+                    "ssmd_version": "",
+                    "ssmd_keywords": "",
+                    "ssmd_author": "",
+                    "ssmd_source": "",
+                    "ssmd_description": "",
+                    "ssmd_rating": "0",
+                    "ssmd_tags": "",
+                }
+
+                for k, v in metadata.items():
+                    if k.startswith("ssmd_") and k != "ssmd_cover_images":
+                        updates[k] = v
+
+                model_util.write_model_metadata(path, module, updates)
+                count += 1
+
+    print(f"[MetadataEditor] Updated {count} models in directory {copy_dir}.")
+    return f"Updated {count} models in directory {copy_dir}."
+
+
+def load_cover_image(model_path, metadata):
+    """
+    Loads a cover image either from embedded metadata or an image file with
+    .preview.png/.png format
+    """
+    cover_images = json.loads(metadata.get("ssmd_cover_images", "[]"))
+    cover_image = None
+    if len(cover_images) > 0:
+        print("[MetadataEditor] Loading embedded cover image.")
+        cover_image = decode_base64_to_pil(cover_images[0])
+    else:
+        basename, ext = os.path.splitext(model_path)
+
+        preview_paths = [f"{basename}.preview.png", f"{basename}.png"]
+
+        for preview_path in preview_paths:
+            if os.path.isfile(preview_path):
+                print(f"[MetadataEditor] Loading webui preview image: {preview_path}")
+                cover_image = Image.open(preview_path)
+
+    return cover_image
+
+
+# Dummy value since gr.Dataframe cannot handle an empty list
+# https://github.com/gradio-app/gradio/issues/3182
+unknown_folders = ["(Unknown)", 0, 0, 0]
+
+
+def refresh_metadata(module, model_path):
+    """
+    Reads metadata from the model on disk and updates all Gradio components
+    """
+    if model_path == "None":
+        return {}, None, "", "", "", "", "", 0, "", "", "", "", "", {}, [unknown_folders]
+
+    if not os.path.isfile(model_path):
+        return (
+            {"info": f"Model path not found: {model_path}"},
+            None,
+            "",
+            "",
+            "",
+            "",
+            "",
+            0,
+            "",
+            "",
+            "",
+            "",
+            "",
+            {},
+            [unknown_folders],
+        )
+
+    if os.path.splitext(model_path)[1] != ".safetensors":
+        return (
+            {"info": "Model is not in .safetensors format."},
+            None,
+            "",
+            "",
+            "",
+            "",
+            "",
+            0,
+            "",
+            "",
+            "",
+            "",
+            "",
+            {},
+            [unknown_folders],
+        )
+
+    metadata = model_util.read_model_metadata(model_path, module)
+
+    if metadata is None:
+        training_params = {}
+        metadata = {}
+    else:
+        training_params = {k: v for k, v in metadata.items() if k.startswith("ss_")}
+
+    cover_image = load_cover_image(model_path, metadata)
+
+    display_name = metadata.get("ssmd_display_name", "")
+    author = metadata.get("ssmd_author", "")
+    # version = metadata.get("ssmd_version", "")
+    source = metadata.get("ssmd_source", "")
+    keywords = metadata.get("ssmd_keywords", "")
+    description = metadata.get("ssmd_description", "")
+    rating = int(metadata.get("ssmd_rating", "0"))
+    tags = metadata.get("ssmd_tags", "")
+    model_hash = metadata.get("sshs_model_hash", model_util.cache("hashes").get(model_path, {}).get("model", ""))
+    legacy_hash = metadata.get("sshs_legacy_hash", model_util.cache("hashes").get(model_path, {}).get("legacy", ""))
+
+    top_tags = {}
+    if "ss_tag_frequency" in training_params:
+        tag_frequency = json.loads(training_params.pop("ss_tag_frequency"))
+        count_max = 0
+        for dir, frequencies in tag_frequency.items():
+            for tag, count in frequencies.items():
+                tag = tag.strip()
+                existing = top_tags.get(tag, 0)
+                top_tags[tag] = count + existing
+        if len(top_tags) > 0:
+            top_tags = dict(sorted(top_tags.items(), key=lambda x: x[1], reverse=True))
+
+            count_max = max(top_tags.values())
+            top_tags = {k: float(v / count_max) for k, v in top_tags.items()}
+
+    dataset_folders = []
+    if "ss_dataset_dirs" in training_params:
+        dataset_dirs = json.loads(training_params.pop("ss_dataset_dirs"))
+        for dir, counts in dataset_dirs.items():
+            img_count = int(counts["img_count"])
+            n_repeats = int(counts["n_repeats"])
+            dataset_folders.append([dir, img_count, n_repeats, img_count * n_repeats])
+    if dataset_folders:
+        dataset_folders.append(
+            ["(Total)", sum(r[1] for r in dataset_folders), sum(r[2] for r in dataset_folders), sum(r[3] for r in dataset_folders)]
+        )
+    else:
+        dataset_folders.append(unknown_folders)
+
+    return (
+        training_params,
+        cover_image,
+        display_name,
+        author,
+        source,
+        keywords,
+        description,
+        rating,
+        tags,
+        model_hash,
+        legacy_hash,
+        model_path,
+        os.path.dirname(model_path),
+        top_tags,
+        dataset_folders,
+    )
+
+
+def save_metadata(module, model_path, cover_image, display_name, author, source, keywords, description, rating, tags):
+    """
+    Writes metadata from the Gradio components to the model file
+    """
+    if model_path == "None":
+        return "No model selected.", "", ""
+
+    if not os.path.isfile(model_path):
+        return f"file not found: {model_path}", "", ""
+
+    if os.path.splitext(model_path)[1] != ".safetensors":
+        return "Model is not in .safetensors format", "", ""
+
+    metadata = safetensors_hack.read_metadata(model_path)
+    model_hash = safetensors_hack.hash_file(model_path)
+    legacy_hash = model_util.get_legacy_hash(metadata, model_path)
+
+    # TODO: Support multiple images
+    # Blocked on gradio not having a gallery upload option
+    # https://github.com/gradio-app/gradio/issues/1379
+    cover_images = []
+    if cover_image is not None:
+        cover_images.append(encode_pil_to_base64(cover_image).decode("ascii"))
+
+    # NOTE: User-specified metadata should NOT be prefixed with "ss_". This is
+    # to maintain backwards compatibility with the old hashing method. "ss_"
+    # should be used for training parameters that will never be manually
+    # updated on the model.
+    updates = {
+        "ssmd_cover_images": json.dumps(cover_images),
+        "ssmd_display_name": display_name,
+        "ssmd_author": author,
+        # "ssmd_version": version,
+        "ssmd_source": source,
+        "ssmd_keywords": keywords,
+        "ssmd_description": description,
+        "ssmd_rating": rating,
+        "ssmd_tags": tags,
+        "sshs_model_hash": model_hash,
+        "sshs_legacy_hash": legacy_hash,
+    }
+
+    model_util.write_model_metadata(model_path, module, updates)
+    if cover_image is None:
+        delete_webui_model_preview_image(model_path)
+    else:
+        write_webui_model_preview_image(model_path, cover_image)
+
+    model_name = os.path.basename(model_path)
+    return f"Model saved: {model_name}", model_hash, legacy_hash
+
+
+model_name_filter = ""
+
+
+def get_filtered_model_paths(s):
+    # newer Gradio seems to show None in the list?
+    # if not s:
+    #     return ["None"] + list(model_util.lora_models.values())
+    # return ["None"] + [v for v in model_util.lora_models.values() if v and s in v.lower()]
+    if not s:
+        l =  list(model_util.lora_models.values())
+    else:
+        l =  [v for v in model_util.lora_models.values() if v and s in v.lower()]
+    l = [v for v in l if v]     # remove None
+    l = ["None"] + l
+    return l
+
+def get_filtered_model_paths_global():
+    global model_name_filter
+    return get_filtered_model_paths(model_name_filter)
+
+
+def setup_ui(addnet_paste_params):
+    """
+    :dict addnet_paste_params: Dictionary of txt2img/img2img controls for each model weight slider,
+                               for sending module and model to them from the metadata editor
+    """
+    can_edit = False
+
+    with gr.Row().style(equal_height=False):
+        # Lefthand column
+        with gr.Column(variant="panel"):
+            # Module and model selector
+            with gr.Row():
+                model_filter = gr.Textbox("", label="Model path filter", placeholder="Filter models by path name")
+
+                def update_model_filter(s):
+                    global model_name_filter
+                    model_name_filter = s.strip().lower()
+
+                model_filter.change(update_model_filter, inputs=[model_filter], outputs=[])
+            with gr.Row():
+                module = gr.Dropdown(
+                    ["LoRA"],
+                    label="Network module",
+                    value="LoRA",
+                    interactive=True,
+                    elem_id="additional_networks_metadata_editor_module",
+                )
+                model = gr.Dropdown(
+                    get_filtered_model_paths_global(),
+                    label="Model",
+                    value="None",
+                    interactive=True,
+                    elem_id="additional_networks_metadata_editor_model",
+                )
+                modules.ui.create_refresh_button(
+                    model, model_util.update_models, lambda: {"choices": get_filtered_model_paths_global()}, "refresh_lora_models"
+                )
+
+                def submit_model_filter(s):
+                    global model_name_filter
+                    model_name_filter = s
+                    paths = get_filtered_model_paths(s)
+                    return gr.Dropdown.update(choices=paths, value="None")
+
+                model_filter.submit(submit_model_filter, inputs=[model_filter], outputs=[model])
+
+            # Model hashes and path
+            with gr.Row():
+                model_hash = gr.Textbox("", label="Model hash", interactive=False)
+                legacy_hash = gr.Textbox("", label="Legacy hash", interactive=False)
+            with gr.Row():
+                model_path = gr.Textbox("", label="Model path", interactive=False)
+                open_folder_button = ToolButton(
+                    value=folder_symbol,
+                    elem_id="hidden_element" if shared.cmd_opts.hide_ui_dir_config else "open_folder_metadata_editor",
+                )
+
+            # Send to txt2img/img2img buttons
+            for tabname in ["txt2img", "img2img"]:
+                with gr.Row():
+                    with gr.Box():
+                        with gr.Row():
+                            gr.HTML(f"Send to {tabname}:")
+                            for i in range(MAX_MODEL_COUNT):
+                                send_to_button = ToolButton(
+                                    value=keycap_symbols[i], elem_id=f"additional_networks_send_to_{tabname}_{i}"
+                                )
+                                send_to_button.click(
+                                    fn=lambda modu, mod: (modu, model_util.find_closest_lora_model_name(mod) or "None"),
+                                    inputs=[module, model],
+                                    outputs=[addnet_paste_params[tabname][i]["module"], addnet_paste_params[tabname][i]["model"]],
+                                )
+                                send_to_button.click(fn=None, _js=f"addnet_switch_to_{tabname}", inputs=None, outputs=None)
+
+            # "Copy metadata to other models" panel
+            with gr.Row():
+                with gr.Column():
+                    gr.HTML(value="Copy metadata to other models in directory")
+                    copy_metadata_dir = gr.Textbox(
+                        "",
+                        label="Containing directory",
+                        placeholder="All models in this directory will receive the selected model's metadata",
+                    )
+                    with gr.Row():
+                        copy_same_session = gr.Checkbox(True, label="Only copy to models with same session ID")
+                        copy_no_metadata = gr.Checkbox(True, label="Only copy to models with no metadata")
+                    copy_metadata_button = gr.Button("Copy Metadata", variant="primary")
+
+        # Center column, metadata viewer/editor
+        with gr.Column():
+            with gr.Row():
+                display_name = gr.Textbox(value="", label="Name", placeholder="Display name for this model", interactive=can_edit)
+                author = gr.Textbox(value="", label="Author", placeholder="Author of this model", interactive=can_edit)
+            with gr.Row():
+                keywords = gr.Textbox(
+                    value="", label="Keywords", placeholder="Activation keywords, comma-separated", interactive=can_edit
+                )
+            with gr.Row():
+                description = gr.Textbox(
+                    value="",
+                    label="Description",
+                    placeholder="Model description/readme/notes/instructions",
+                    lines=15,
+                    interactive=can_edit,
+                )
+            with gr.Row():
+                source = gr.Textbox(
+                    value="", label="Source", placeholder="Source URL where this model could be found", interactive=can_edit
+                )
+            with gr.Row():
+                rating = gr.Slider(minimum=0, maximum=10, step=1, label="Rating", value=0, interactive=can_edit)
+                tags = gr.Textbox(
+                    value="",
+                    label="Tags",
+                    placeholder='Comma-separated list of tags ("artist, style, character, 2d, 3d...")',
+                    lines=2,
+                    interactive=can_edit,
+                )
+            with gr.Row():
+                editing_enabled = gr.Checkbox(label="Editing Enabled", value=can_edit)
+                with gr.Row():
+                    save_metadata_button = gr.Button("Save Metadata", variant="primary", interactive=can_edit)
+            with gr.Row():
+                save_output = gr.HTML("")
+
+        # Righthand column, cover image and training parameters view
+        with gr.Column():
+            # Cover image
+            with gr.Row():
+                cover_image = gr.Image(
+                    label="Cover image",
+                    elem_id="additional_networks_cover_image",
+                    source="upload",
+                    interactive=can_edit,
+                    type="pil",
+                    image_mode="RGBA",
+                ).style(height=480)
+
+            # Image parameters
+            with gr.Accordion("Image Parameters", open=False):
+                with gr.Row():
+                    info2 = gr.HTML()
+            with gr.Row():
+                try:
+                    send_to_buttons = parameters_copypaste.create_buttons(["txt2img", "img2img", "inpaint", "extras"])
+                except:
+                    pass
+
+            # Training info, below cover image
+            with gr.Accordion("Training info", open=False):
+                # Top tags used
+                with gr.Row():
+                    max_top_tags = int(shared.opts.data.get("additional_networks_max_top_tags", 20))
+                    most_frequent_tags = gr.Label(value={}, label="Most frequent tags in captions", num_top_classes=max_top_tags)
+
+                # Dataset folders
+                with gr.Row():
+                    max_dataset_folders = int(shared.opts.data.get("additional_networks_max_dataset_folders", 20))
+                    dataset_folders = gr.Dataframe(
+                        headers=["Name", "Image Count", "Repeats", "Total Images"],
+                        datatype=["str", "number", "number", "number"],
+                        label="Dataset folder structure",
+                        max_rows=max_dataset_folders,
+                        col_count=(4, "fixed"),
+                    )
+
+                # Training Parameters
+                with gr.Row():
+                    metadata_view = gr.JSON(value={}, label="Training parameters")
+
+            # Hidden/internal
+            with gr.Row(visible=False):
+                info1 = gr.HTML()
+                img_file_info = gr.Textbox(label="Generate Info", interactive=False, lines=6)
+
+    open_folder_button.click(fn=lambda p: open_folder(os.path.dirname(p)), inputs=[model_path], outputs=[])
+    copy_metadata_button.click(
+        fn=copy_metadata_to_all,
+        inputs=[module, model, copy_metadata_dir, copy_same_session, copy_no_metadata, cover_image],
+        outputs=[save_output],
+    )
+
+    def update_editing(enabled):
+        """
+        Enable/disable components based on "Editing Enabled" status
+        """
+        updates = [gr.Textbox.update(interactive=enabled)] * 6
+        updates.append(gr.Image.update(interactive=enabled))
+        updates.append(gr.Slider.update(interactive=enabled))
+        updates.append(gr.Button.update(interactive=enabled))
+        return updates
+
+    editing_enabled.change(
+        fn=update_editing,
+        inputs=[editing_enabled],
+        outputs=[display_name, author, source, keywords, description, tags, cover_image, rating, save_metadata_button],
+    )
+
+    cover_image.change(fn=modules.extras.run_pnginfo, inputs=[cover_image], outputs=[info1, img_file_info, info2])
+
+    try:
+        parameters_copypaste.bind_buttons(send_to_buttons, cover_image, img_file_info)
+    except:
+        pass
+
+    model.change(
+        refresh_metadata,
+        inputs=[module, model],
+        outputs=[
+            metadata_view,
+            cover_image,
+            display_name,
+            author,
+            source,
+            keywords,
+            description,
+            rating,
+            tags,
+            model_hash,
+            legacy_hash,
+            model_path,
+            copy_metadata_dir,
+            most_frequent_tags,
+            dataset_folders,
+        ],
+    )
+    save_metadata_button.click(
+        save_metadata,
+        inputs=[module, model, cover_image, display_name, author, source, keywords, description, rating, tags],
+        outputs=[save_output, model_hash, legacy_hash],
+    )

extensions/addtional/scripts/model_util.py

@@ -0,0 +1,338 @@
+import os
+import os.path
+import re
+import shutil
+import json
+import stat
+import tqdm
+from collections import OrderedDict
+from multiprocessing.pool import ThreadPool as Pool
+
+from modules import shared, sd_models, hashes
+from scripts import safetensors_hack, model_util, util
+import modules.scripts as scripts
+
+
+# MAX_MODEL_COUNT = shared.cmd_opts.addnet_max_model_count or 5
+MAX_MODEL_COUNT = shared.cmd_opts.addnet_max_model_count if hasattr(shared.cmd_opts, "addnet_max_model_count") else 5
+LORA_MODEL_EXTS = [".pt", ".ckpt", ".safetensors"]
+re_legacy_hash = re.compile("\(([0-9a-f]{8})\)$")  # matches 8-character hashes, new hash has 12 characters
+lora_models = {}  # "My_Lora(abcdef123456)" -> "C:/path/to/model.safetensors"
+lora_model_names = {}  # "my_lora" -> "My_Lora(My_Lora(abcdef123456)"
+legacy_model_names = {}
+lora_models_dir = os.path.join(scripts.basedir(), "models/lora")
+os.makedirs(lora_models_dir, exist_ok=True)
+
+
+def is_safetensors(filename):
+    return os.path.splitext(filename)[1] == ".safetensors"
+
+
+def read_model_metadata(model_path, module):
+    if model_path.startswith('"') and model_path.endswith('"'):  # trim '"' at start/end
+        model_path = model_path[1:-1]
+    if not os.path.exists(model_path):
+        return None
+
+    metadata = None
+    if module == "LoRA":
+        if os.path.splitext(model_path)[1] == ".safetensors":
+            metadata = safetensors_hack.read_metadata(model_path)
+
+    return metadata
+
+
+def write_model_metadata(model_path, module, updates):
+    if model_path.startswith('"') and model_path.endswith('"'):  # trim '"' at start/end
+        model_path = model_path[1:-1]
+    if not os.path.exists(model_path):
+        return None
+
+    from safetensors.torch import save_file
+
+    back_up = shared.opts.data.get("additional_networks_back_up_model_when_saving", True)
+    if back_up:
+        backup_path = model_path + ".backup"
+        if not os.path.exists(backup_path):
+            print(f"[MetadataEditor] Backing up current model to {backup_path}")
+            shutil.copyfile(model_path, backup_path)
+
+    metadata = None
+    tensors = {}
+    if module == "LoRA":
+        if os.path.splitext(model_path)[1] == ".safetensors":
+            tensors, metadata = safetensors_hack.load_file(model_path, "cpu")
+
+            for k, v in updates.items():
+                metadata[k] = str(v)
+
+            save_file(tensors, model_path, metadata)
+            print(f"[MetadataEditor] Model saved: {model_path}")
+
+
+def get_model_list(module, model, model_dir, sort_by):
+    if model_dir == "":
+        # Get list of models with same folder as this one
+        model_path = lora_models.get(model, None)
+        if model_path is None:
+            return []
+        model_dir = os.path.dirname(model_path)
+
+    if not os.path.isdir(model_dir):
+        return []
+
+    found, _ = get_all_models([model_dir], sort_by, "")
+    return list(found.keys())  # convert dict_keys to list
+
+
+def traverse_all_files(curr_path, model_list):
+    f_list = [(os.path.join(curr_path, entry.name), entry.stat()) for entry in os.scandir(curr_path)]
+    for f_info in f_list:
+        fname, fstat = f_info
+        if os.path.splitext(fname)[1] in LORA_MODEL_EXTS:
+            model_list.append(f_info)
+        elif stat.S_ISDIR(fstat.st_mode):
+            model_list = traverse_all_files(fname, model_list)
+    return model_list
+
+
+def get_model_hash(metadata, filename):
+    if metadata is None:
+        return hashes.calculate_sha256(filename)
+
+    if "sshs_model_hash" in metadata:
+        return metadata["sshs_model_hash"]
+
+    return safetensors_hack.hash_file(filename)
+
+
+def get_legacy_hash(metadata, filename):
+    if metadata is None:
+        return sd_models.model_hash(filename)
+
+    if "sshs_legacy_hash" in metadata:
+        return metadata["sshs_legacy_hash"]
+
+    return safetensors_hack.legacy_hash_file(filename)
+
+
+import filelock
+
+cache_filename = os.path.join(scripts.basedir(), "hashes.json")
+cache_data = None
+
+
+def cache(subsection):
+    global cache_data
+
+    if cache_data is None:
+        with filelock.FileLock(cache_filename + ".lock"):
+            if not os.path.isfile(cache_filename):
+                cache_data = {}
+            else:
+                with open(cache_filename, "r", encoding="utf8") as file:
+                    cache_data = json.load(file)
+
+    s = cache_data.get(subsection, {})
+    cache_data[subsection] = s
+
+    return s
+
+
+def dump_cache():
+    with filelock.FileLock(cache_filename + ".lock"):
+        with open(cache_filename, "w", encoding="utf8") as file:
+            json.dump(cache_data, file, indent=4)
+
+
+def get_model_rating(filename):
+    if not model_util.is_safetensors(filename):
+        return 0
+
+    metadata = safetensors_hack.read_metadata(filename)
+    return int(metadata.get("ssmd_rating", "0"))
+
+
+def has_user_metadata(filename):
+    if not model_util.is_safetensors(filename):
+        return False
+
+    metadata = safetensors_hack.read_metadata(filename)
+    return any(k.startswith("ssmd_") for k in metadata.keys())
+
+
+def hash_model_file(finfo):
+    filename = finfo[0]
+    stat = finfo[1]
+    name = os.path.splitext(os.path.basename(filename))[0]
+
+    # Prevent a hypothetical "None.pt" from being listed.
+    if name != "None":
+        metadata = None
+
+        cached = cache("hashes").get(filename, None)
+        if cached is None or stat.st_mtime != cached["mtime"]:
+            if metadata is None and model_util.is_safetensors(filename):
+                try:
+                    metadata = safetensors_hack.read_metadata(filename)
+                except Exception as ex:
+                    return {"error": ex, "filename": filename}
+            model_hash = get_model_hash(metadata, filename)
+            legacy_hash = get_legacy_hash(metadata, filename)
+        else:
+            model_hash = cached["model"]
+            legacy_hash = cached["legacy"]
+
+    return {"model": model_hash, "legacy": legacy_hash, "fileinfo": finfo}
+
+
+def get_all_models(paths, sort_by, filter_by):
+    fileinfos = []
+    for path in paths:
+        if os.path.isdir(path):
+            fileinfos += traverse_all_files(path, [])
+
+    show_only_safetensors = shared.opts.data.get("additional_networks_show_only_safetensors", False)
+    show_only_missing_meta = shared.opts.data.get("additional_networks_show_only_models_with_metadata", "disabled")
+
+    if show_only_safetensors:
+        fileinfos = [x for x in fileinfos if is_safetensors(x[0])]
+
+    if show_only_missing_meta == "has metadata":
+        fileinfos = [x for x in fileinfos if has_user_metadata(x[0])]
+    elif show_only_missing_meta == "missing metadata":
+        fileinfos = [x for x in fileinfos if not has_user_metadata(x[0])]
+
+    print("[AddNet] Updating model hashes...")
+    data = []
+    thread_count = max(1, int(shared.opts.data.get("additional_networks_hash_thread_count", 1)))
+    p = Pool(processes=thread_count)
+    with tqdm.tqdm(total=len(fileinfos)) as pbar:
+        for res in p.imap_unordered(hash_model_file, fileinfos):
+            pbar.update()
+            if "error" in res:
+                print(f"Failed to read model file {res['filename']}: {res['error']}")
+            else:
+                data.append(res)
+    p.close()
+
+    cache_hashes = cache("hashes")
+
+    res = OrderedDict()
+    res_legacy = OrderedDict()
+    filter_by = filter_by.strip(" ")
+    if len(filter_by) != 0:
+        data = [x for x in data if filter_by.lower() in os.path.basename(x["fileinfo"][0]).lower()]
+    if sort_by == "name":
+        data = sorted(data, key=lambda x: os.path.basename(x["fileinfo"][0]))
+    elif sort_by == "date":
+        data = sorted(data, key=lambda x: -x["fileinfo"][1].st_mtime)
+    elif sort_by == "path name":
+        data = sorted(data, key=lambda x: x["fileinfo"][0])
+    elif sort_by == "rating":
+        data = sorted(data, key=lambda x: get_model_rating(x["fileinfo"][0]), reverse=True)
+    elif sort_by == "has user metadata":
+        data = sorted(
+            data, key=lambda x: os.path.basename(x["fileinfo"][0]) if has_user_metadata(x["fileinfo"][0]) else "", reverse=True
+        )
+
+    reverse = shared.opts.data.get("additional_networks_reverse_sort_order", False)
+    if reverse:
+        data = reversed(data)
+
+    for result in data:
+        finfo = result["fileinfo"]
+        filename = finfo[0]
+        stat = finfo[1]
+        model_hash = result["model"]
+        legacy_hash = result["legacy"]
+
+        name = os.path.splitext(os.path.basename(filename))[0]
+
+        # Commas in the model name will mess up infotext restoration since the
+        # infotext is delimited by commas
+        name = name.replace(",", "_")
+
+        # Prevent a hypothetical "None.pt" from being listed.
+        if name != "None":
+            full_name = name + f"({model_hash[0:12]})"
+            res[full_name] = filename
+            res_legacy[legacy_hash] = full_name
+            cache_hashes[filename] = {"model": model_hash, "legacy": legacy_hash, "mtime": stat.st_mtime}
+
+    return res, res_legacy
+
+
+def find_closest_lora_model_name(search: str):
+    if not search or search == "None":
+        return None
+
+    # Match name and hash, case-sensitive
+    # "MyModel-epoch00002(abcdef123456)"
+    if search in lora_models:
+        return search
+
+    # Match model path, case-sensitive (from metadata editor)
+    # "C:/path/to/mymodel-epoch00002.safetensors"
+    if os.path.isfile(search):
+        import json
+
+        find = os.path.normpath(search)
+        value = next((k for k in lora_models.keys() if lora_models[k] == find), None)
+        if value:
+            return value
+
+    search = search.lower()
+
+    # Match full name, case-insensitive
+    # "mymodel-epoch00002"
+    if search in lora_model_names:
+        return lora_model_names.get(search)
+
+    # Match legacy hash (8 characters)
+    # "MyModel(abcd1234)"
+    result = re_legacy_hash.search(search)
+    if result is not None:
+        model_hash = result.group(1)
+        if model_hash in legacy_model_names:
+            new_model_name = legacy_model_names[model_hash]
+            return new_model_name
+
+    # Use any model with the search term as the prefix, case-insensitive, sorted
+    # by name length
+    # "mymodel"
+    applicable = [name for name in lora_model_names.keys() if search in name.lower()]
+    if not applicable:
+        return None
+    applicable = sorted(applicable, key=lambda name: len(name))
+    return lora_model_names[applicable[0]]
+
+
+def update_models():
+    global lora_models, lora_model_names, legacy_model_names
+    paths = [lora_models_dir]
+    extra_lora_paths = util.split_path_list(shared.opts.data.get("additional_networks_extra_lora_path", ""))
+    for path in extra_lora_paths:
+        path = path.lstrip()
+        if os.path.isdir(path):
+            paths.append(path)
+
+    sort_by = shared.opts.data.get("additional_networks_sort_models_by", "name")
+    filter_by = shared.opts.data.get("additional_networks_model_name_filter", "")
+    res, res_legacy = get_all_models(paths, sort_by, filter_by)
+
+    lora_models.clear()
+    lora_models["None"] = None
+    lora_models.update(res)
+
+    for name_and_hash, filename in lora_models.items():
+        if filename == None:
+            continue
+        name = os.path.splitext(os.path.basename(filename))[0].lower()
+        lora_model_names[name] = name_and_hash
+
+    legacy_model_names = res_legacy
+    dump_cache()
+
+
+update_models()

extensions/addtional/scripts/safetensors_hack.py

@@ -0,0 +1,116 @@
+import io
+import os
+import mmap
+import torch
+import json
+import hashlib
+import safetensors
+import safetensors.torch
+
+from modules import sd_models
+
+# PyTorch 1.13 and later have _UntypedStorage renamed to UntypedStorage
+UntypedStorage = torch.storage.UntypedStorage if hasattr(torch.storage, 'UntypedStorage') else torch.storage._UntypedStorage
+
+def read_metadata(filename):
+    """Reads the JSON metadata from a .safetensors file"""
+    with open(filename, mode="r", encoding="utf8") as file_obj:
+        with mmap.mmap(file_obj.fileno(), length=0, access=mmap.ACCESS_READ) as m:
+            header = m.read(8)
+            n = int.from_bytes(header, "little")
+            metadata_bytes = m.read(n)
+            metadata = json.loads(metadata_bytes)
+
+    return metadata.get("__metadata__", {})
+
+
+def load_file(filename, device):
+    """"Loads a .safetensors file without memory mapping that locks the model file.
+    Works around safetensors issue: https://github.com/huggingface/safetensors/issues/164"""
+    with open(filename, mode="r", encoding="utf8") as file_obj:
+        with mmap.mmap(file_obj.fileno(), length=0, access=mmap.ACCESS_READ) as m:
+            header = m.read(8)
+            n = int.from_bytes(header, "little")
+            metadata_bytes = m.read(n)
+            metadata = json.loads(metadata_bytes)
+
+    size = os.stat(filename).st_size
+    storage = UntypedStorage.from_file(filename, False, size)
+    offset = n + 8
+    md = metadata.get("__metadata__", {})
+    return {name: create_tensor(storage, info, offset) for name, info in metadata.items() if name != "__metadata__"}, md
+
+
+def hash_file(filename):
+    """Hashes a .safetensors file using the new hashing method.
+    Only hashes the weights of the model."""
+    hash_sha256 = hashlib.sha256()
+    blksize = 1024 * 1024
+
+    with open(filename, mode="r", encoding="utf8") as file_obj:
+        with mmap.mmap(file_obj.fileno(), length=0, access=mmap.ACCESS_READ) as m:
+            header = m.read(8)
+            n = int.from_bytes(header, "little")
+
+    with open(filename, mode="rb") as file_obj:
+        offset = n + 8
+        file_obj.seek(offset)
+        for chunk in iter(lambda: file_obj.read(blksize), b""):
+            hash_sha256.update(chunk)
+
+    return hash_sha256.hexdigest()
+
+
+def legacy_hash_file(filename):
+    """Hashes a model file using the legacy `sd_models.model_hash()` method."""
+    hash_sha256 = hashlib.sha256()
+
+    metadata = read_metadata(filename)
+
+    # For compatibility with legacy models: This replicates the behavior of
+    # sd_models.model_hash as if there were no user-specified metadata in the
+    # .safetensors file. That leaves the training parameters, which are
+    # immutable. It is important the hash does not include the embedded user
+    # metadata as that would mean the hash could change every time the user
+    # updates the name/description/etc. The new hashing method fixes this
+    # problem by only hashing the region of the file containing the tensors.
+    if any(not k.startswith("ss_") for k in metadata):
+      # Strip the user metadata, re-serialize the file as if it were freshly
+      # created from sd-scripts, and hash that with model_hash's behavior.
+      tensors, metadata = load_file(filename, "cpu")
+      metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")}
+      model_bytes = safetensors.torch.save(tensors, metadata)
+
+      hash_sha256.update(model_bytes[0x100000:0x110000])
+      return hash_sha256.hexdigest()[0:8]
+    else:
+      # This should work fine with model_hash since when the legacy hashing
+      # method was being used the user metadata system hadn't been implemented
+      # yet.
+      return sd_models.model_hash(filename)
+
+
+DTYPES = {
+    "F64": torch.float64,
+    "F32": torch.float32,
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+    "I64": torch.int64,
+    # "U64": torch.uint64,
+    "I32": torch.int32,
+    # "U32": torch.uint32,
+    "I16": torch.int16,
+    # "U16": torch.uint16,
+    "I8": torch.int8,
+    "U8": torch.uint8,
+    "BOOL": torch.bool
+}
+
+
+def create_tensor(storage, info, offset):
+    """Creates a tensor without holding on to an open handle to the parent model
+    file."""
+    dtype = DTYPES[info["dtype"]]
+    shape = info["shape"]
+    start, stop = info["data_offsets"]
+    return torch.asarray(storage[start + offset : stop + offset], dtype=torch.uint8).view(dtype=dtype).reshape(shape).clone().detach()

extensions/addtional/scripts/util.py

@@ -0,0 +1,12 @@
+import csv
+from io import StringIO
+from typing import List
+
+def split_path_list(path_list: str) -> List[str]:
+    pl = []
+    with StringIO() as f:
+        f.write(path_list)
+        f.seek(0)
+        for r in csv.reader(f):
+            pl += r
+    return pl

extensions/addtional/style.css

@@ -0,0 +1,9 @@
+#additional_networks_cover_image,
+#additional_networks_cover_image > .h-60,
+#additional_networks_cover_image > .h-60 > div,
+#additional_networks_cover_image > .h-60 > div > img
+{
+    height: 480px !important;
+    max-height: 480px !important;
+    min-height: 480px !important;
+}

extensions/adetailer/.github/ISSUE_TEMPLATE/bug_report.yaml

@@ -0,0 +1,23 @@
+name: Bug report
+description: Create a report
+title: "[Bug]: "
+
+body:
+  - type: textarea
+    attributes:
+      label: Describe the bug
+      description: A clear and concise description of what the bug is.
+
+  - type: textarea
+    attributes:
+      label: Full console logs
+      description: |
+        The full console log of your terminal.
+        From `Python 3.10.*, Version: v1.*, Commit hash: *` to the end.
+      render: Shell
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: List of installed extensions

extensions/adetailer/.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.

extensions/adetailer/.github/workflows/stale.yml

@@ -0,0 +1,13 @@
+name: 'Close stale issues and PRs'
+on:
+  schedule:
+    - cron: '30 1 * * *'
+
+jobs:
+  stale:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/stale@v8
+        with:
+          days-before-stale: 30
+          days-before-close: 5

extensions/adetailer/.gitignore

@@ -0,0 +1,196 @@
+# Created by https://www.toptal.com/developers/gitignore/api/python,visualstudiocode
+# Edit at https://www.toptal.com/developers/gitignore?templates=python,visualstudiocode
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+### VisualStudioCode ###
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+!.vscode/*.code-snippets
+
+# Local History for Visual Studio Code
+.history/
+
+# Built Visual Studio Code Extensions
+*.vsix
+
+### VisualStudioCode Patch ###
+# Ignore all local history of files
+.history
+.ionide
+
+# End of https://www.toptal.com/developers/gitignore/api/python,visualstudiocode
+*.ipynb

extensions/adetailer/.pre-commit-config.yaml

@@ -0,0 +1,24 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
+    hooks:
+      - id: trailing-whitespace
+        args: [--markdown-linebreak-ext=md]
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+
+  - repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: "v0.0.270"
+    hooks:
+      - id: ruff
+        args: [--fix, --exit-non-zero-on-fix]
+
+  - repo: https://github.com/psf/black
+    rev: 23.3.0
+    hooks:
+      - id: black

extensions/adetailer/CHANGELOG.md

@@ -0,0 +1,171 @@
+# Changelog
+
+## 2023-06-01
+
+- v23.6.1
+- `inpaint, scribble, lineart, openpose, tile` 5가지 컨트롤넷 모델 지원 (PR #107)
+- controlnet guidance start, end 인자 추가 (PR #107)
+- `modules.extensions`를 사용하여 컨트롤넷 확장을 불러오고 경로를 알아내로록 변경
+- ui에서 컨트롤넷을 별도 함수로 분리
+
+## 2023-05-30
+
+- v23.6.0
+- 스크립트의 이름을 `After Detailer`에서 `ADetailer`로 변경
+  - API 사용자는 변경 필요함
+- 몇몇 설정 변경
+  - `ad_conf` → `ad_confidence`. 0~100 사이의 int → 0.0~1.0 사이의 float
+  - `ad_inpaint_full_res` → `ad_inpaint_only_masked`
+  - `ad_inpaint_full_res_padding` → `ad_inpaint_only_masked_padding`
+- mediapipe face mesh 모델 추가
+  - mediapipe 최소 버전 `0.10.0`
+
+- rich traceback 제거함
+- huggingface 다운로드 실패할 때 에러가 나지 않게 하고 해당 모델을 제거함
+
+## 2023-05-26
+
+- v23.5.19
+- 1번째 탭에도 `None` 옵션을 추가함
+- api로 ad controlnet model에 inpaint가 아닌 다른 컨트롤넷 모델을 사용하지 못하도록 막음
+- adetailer 진행중에 total tqdm 진행바 업데이트를 멈춤
+- state.inturrupted 상태에서 adetailer 과정을 중지함
+- 컨트롤넷 process를 각 batch가 끝난 순간에만 호출하도록 변경
+
+### 2023-05-25
+
+- v23.5.18
+- 컨트롤넷 관련 수정
+  - unit의 `input_mode`를 `SIMPLE`로 모두 변경
+  - 컨트롤넷 유넷 훅과 하이잭 함수들을 adetailer를 실행할 때에만 되돌리는 기능 추가
+  - adetailer 처리가 끝난 뒤 컨트롤넷 스크립트의 process를 다시 진행함. (batch count 2 이상일때의 문제 해결)
+- 기본 활성 스크립트 목록에서 컨트롤넷을 뺌
+
+### 2023-05-22
+
+- v23.5.17
+- 컨트롤넷 확장이 있으면 컨트롤넷 스크립트를 활성화함. (컨트롤넷 관련 문제 해결)
+- 모든 컴포넌트에 elem_id 설정
+- ui에 버전을 표시함
+
+
+### 2023-05-19
+
+- v23.5.16
+- 추가한 옵션
+  - Mask min/max ratio
+  - Mask merge mode
+  - Restore faces after ADetailer
+- 옵션들을 Accordion으로 묶음
+
+### 2023-05-18
+
+- v23.5.15
+- 필요한 것만 임포트하도록 변경 (vae 로딩 오류 없어짐. 로딩 속도 빨라짐)
+
+### 2023-05-17
+
+- v23.5.14
+- `[SKIP]`으로 ad prompt 일부를 건너뛰는 기능 추가
+- bbox 정렬 옵션 추가
+- sd_webui 타입힌트를 만들어냄
+- enable checker와 관련된 api 오류 수정?
+
+### 2023-05-15
+
+- v23.5.13
+- `[SEP]`으로 ad prompt를 분리하여 적용하는 기능 추가
+- enable checker를 다시 pydantic으로 변경함
+- ui 관련 함수를 adetailer.ui 폴더로 분리함
+- controlnet을 사용할 때 모든 controlnet unit 비활성화
+- adetailer 폴더가 없으면 만들게 함
+
+### 2023-05-13
+
+- v23.5.12
+- `ad_enable`을 제외한 입력이 dict타입으로 들어오도록 변경
+  - web api로 사용할 때에 특히 사용하기 쉬움
+  - web api breaking change
+- `mask_preprocess` 인자를 넣지 않았던 오류 수정 (PR #47)
+- huggingface에서 모델을 다운로드하지 않는 옵션 추가 `--ad-no-huggingface`
+
+### 2023-05-12
+
+- v23.5.11
+- `ultralytics` 알람 제거
+- 필요없는 exif 인자 더 제거함
+- `use separate steps` 옵션 추가
+- ui 배치를 조정함
+
+### 2023-05-09
+
+- v23.5.10
+- 선택한 스크립트만 ADetailer에 적용하는 옵션 추가, 기본값 `True`. 설정 탭에서 지정가능.
+  - 기본값: `dynamic_prompting,dynamic_thresholding,wildcards,wildcard_recursive`
+- `person_yolov8s-seg.pt` 모델 추가
+- `ultralytics`의 최소 버전을 `8.0.97`로 설정 (C:\\ 문제 해결된 버전)
+
+### 2023-05-08
+
+- v23.5.9
+- 2가지 이상의 모델을 사용할 수 있음. 기본값: 2, 최대: 5
+- segment 모델을 사용할 수 있게 함. `person_yolov8n-seg.pt` 추가
+
+### 2023-05-07
+
+- v23.5.8
+- 프롬프트와 네거티브 프롬프트에 방향키 지원 (PR #24)
+- `mask_preprocess`를 추가함. 이전 버전과 시드값이 달라질 가능성 있음!
+- 이미지 처리가 일어났을 때에만 before이미지를 저장함
+- 설정창의 레이블을 ADetailer 대신 더 적절하게 수정함
+
+### 2023-05-06
+
+- v23.5.7
+- `ad_use_cfg_scale` 옵션 추가. cfg 스케일을 따로 사용할지 말지 결정함.
+- `ad_enable` 기본값을 `True`에서 `False`로 변경
+- `ad_model`의 기본값을 `None`에서 첫번째 모델로 변경
+- 최소 2개의 입력(ad_enable, ad_model)만 들어오면 작동하게 변경.
+
+- v23.5.7.post0
+- `init_controlnet_ext`을 controlnet_exists == True일때에만 실행
+- webui를 C드라이브 바로 밑에 설치한 사람들에게 `ultralytics` 경고 표시
+
+### 2023-05-05 (어린이날)
+
+- v23.5.5
+- `Save images before ADetailer` 옵션 추가
+- 입력으로 들어온 인자와 ALL_ARGS의 길이가 다르면 에러메세지
+- README.md에 설치방법 추가
+
+- v23.5.6
+- get_args에서 IndexError가 발생하면 자세한 에러메세지를 �� 수 있음
+- AdetailerArgs에 extra_params 내장
+- scripts_args를 딥카피함
+- postprocess_image를 약간 분리함
+
+- v23.5.6.post0
+- `init_controlnet_ext`에서 에러메세지를 자세히 볼 수 있음
+
+### 2023-05-04
+
+- v23.5.4
+- use pydantic for arguments validation
+- revert: ad_model to `None` as default
+- revert: `__future__` imports
+- lazily import yolo and mediapipe
+
+### 2023-05-03
+
+- v23.5.3.post0
+- remove `__future__` imports
+- change to copy scripts and scripts args
+
+- v23.5.3.post1
+- change default ad_model from `None`
+
+### 2023-05-02
+
+- v23.5.3
+- Remove `None` from model list and add `Enable ADetailer` checkbox.
+- install.py `skip_install` fix.

extensions/adetailer/LICENSE.md

@@ -0,0 +1,662 @@
+
+                    GNU AFFERO GENERAL PUBLIC LICENSE
+                       Version 3, 19 November 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <http://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+                            Preamble
+
+  The GNU Affero General Public License is a free, copyleft license for
+software and other kinds of works, specifically designed to ensure
+cooperation with the community in the case of network server software.
+
+  The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works.  By contrast,
+our General Public Licenses are intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users.
+
+  When we speak of free software, we are referring to freedom, not
+price.  Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+  Developers that use our General Public Licenses protect your rights
+with two steps: (1) assert copyright on the software, and (2) offer
+you this License which gives you legal permission to copy, distribute
+and/or modify the software.
+
+  A secondary benefit of defending all users' freedom is that
+improvements made in alternate versions of the program, if they
+receive widespread use, become available for other developers to
+incorporate.  Many developers of free software are heartened and
+encouraged by the resulting cooperation.  However, in the case of
+software used on network servers, this result may fail to come about.
+The GNU General Public License permits making a modified version and
+letting the public access it on a server without ever releasing its
+source code to the public.
+
+  The GNU Affero General Public License is designed specifically to
+ensure that, in such cases, the modified source code becomes available
+to the community.  It requires the operator of a network server to
+provide the source code of the modified version running there to the
+users of that server.  Therefore, public use of a modified version, on
+a publicly accessible server, gives the public access to the source
+code of the modified version.
+
+  An older license, called the Affero General Public License and
+published by Affero, was designed to accomplish similar goals.  This is
+a different license, not a version of the Affero GPL, but Affero has
+released a new version of the Affero GPL which permits relicensing under
+this license.
+
+  The precise terms and conditions for copying, distribution and
+modification follow.
+
+                       TERMS AND CONDITIONS
+
+  0. Definitions.
+
+  "This License" refers to version 3 of the GNU Affero General Public License.
+
+  "Copyright" also means copyright-like laws that apply to other kinds of
+works, such as semiconductor masks.
+
+  "The Program" refers to any copyrightable work licensed under this
+License.  Each licensee is addressed as "you".  "Licensees" and
+"recipients" may be individuals or organizations.
+
+  To "modify" a work means to copy from or adapt all or part of the work
+in a fashion requiring copyright permission, other than the making of an
+exact copy.  The resulting work is called a "modified version" of the
+earlier work or a work "based on" the earlier work.
+
+  A "covered work" means either the unmodified Program or a work based
+on the Program.
+
+  To "propagate" a work means to do anything with it that, without
+permission, would make you directly or secondarily liable for
+infringement under applicable copyright law, except executing it on a
+computer or modifying a private copy.  Propagation includes copying,
+distribution (with or without modification), making available to the
+public, and in some countries other activities as well.
+
+  To "convey" a work means any kind of propagation that enables other
+parties to make or receive copies.  Mere interaction with a user through
+a computer network, with no transfer of a copy, is not conveying.
+
+  An interactive user interface displays "Appropriate Legal Notices"
+to the extent that it includes a convenient and prominently visible
+feature that (1) displays an appropriate copyright notice, and (2)
+tells the user that there is no warranty for the work (except to the
+extent that warranties are provided), that licensees may convey the
+work under this License, and how to view a copy of this License.  If
+the interface presents a list of user commands or options, such as a
+menu, a prominent item in the list meets this criterion.
+
+  1. Source Code.
+
+  The "source code" for a work means the preferred form of the work
+for making modifications to it.  "Object code" means any non-source
+form of a work.
+
+  A "Standard Interface" means an interface that either is an official
+standard defined by a recognized standards body, or, in the case of
+interfaces specified for a particular programming language, one that
+is widely used among developers working in that language.
+
+  The "System Libraries" of an executable work include anything, other
+than the work as a whole, that (a) is included in the normal form of
+packaging a Major Component, but which is not part of that Major
+Component, and (b) serves only to enable use of the work with that
+Major Component, or to implement a Standard Interface for which an
+implementation is available to the public in source code form.  A
+"Major Component", in this context, means a major essential component
+(kernel, window system, and so on) of the specific operating system
+(if any) on which the executable work runs, or a compiler used to
+produce the work, or an object code interpreter used to run it.
+
+  The "Corresponding Source" for a work in object code form means all
+the source code needed to generate, install, and (for an executable
+work) run the object code and to modify the work, including scripts to
+control those activities.  However, it does not include the work's
+System Libraries, or general-purpose tools or generally available free
+programs which are used unmodified in performing those activities but
+which are not part of the work.  For example, Corresponding Source
+includes interface definition files associated with source files for
+the work, and the source code for shared libraries and dynamically
+linked subprograms that the work is specifically designed to require,
+such as by intimate data communication or control flow between those
+subprograms and other parts of the work.
+
+  The Corresponding Source need not include anything that users
+can regenerate automatically from other parts of the Corresponding
+Source.
+
+  The Corresponding Source for a work in source code form is that
+same work.
+
+  2. Basic Permissions.
+
+  All rights granted under this License are granted for the term of
+copyright on the Program, and are irrevocable provided the stated
+conditions are met.  This License explicitly affirms your unlimited
+permission to run the unmodified Program.  The output from running a
+covered work is covered by this License only if the output, given its
+content, constitutes a covered work.  This License acknowledges your
+rights of fair use or other equivalent, as provided by copyright law.
+
+  You may make, run and propagate covered works that you do not
+convey, without conditions so long as your license otherwise remains
+in force.  You may convey covered works to others for the sole purpose
+of having them make modifications exclusively for you, or provide you
+with facilities for running those works, provided that you comply with
+the terms of this License in conveying all material for which you do
+not control copyright.  Those thus making or running the covered works
+for you must do so exclusively on your behalf, under your direction
+and control, on terms that prohibit them from making any copies of
+your copyrighted material outside their relationship with you.
+
+  Conveying under any other circumstances is permitted solely under
+the conditions stated below.  Sublicensing is not allowed; section 10
+makes it unnecessary.
+
+  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
+  No covered work shall be deemed part of an effective technological
+measure under any applicable law fulfilling obligations under article
+11 of the WIPO copyright treaty adopted on 20 December 1996, or
+similar laws prohibiting or restricting circumvention of such
+measures.
+
+  When you convey a covered work, you waive any legal power to forbid
+circumvention of technological measures to the extent such circumvention
+is effected by exercising rights under this License with respect to
+the covered work, and you disclaim any intention to limit operation or
+modification of the work as a means of enforcing, against the work's
+users, your or third parties' legal rights to forbid circumvention of
+technological measures.
+
+  4. Conveying Verbatim Copies.
+
+  You may convey verbatim copies of the Program's source code as you
+receive it, in any medium, provided that you conspicuously and
+appropriately publish on each copy an appropriate copyright notice;
+keep intact all notices stating that this License and any
+non-permissive terms added in accord with section 7 apply to the code;
+keep intact all notices of the absence of any warranty; and give all
+recipients a copy of this License along with the Program.
+
+  You may charge any price or no price for each copy that you convey,
+and you may offer support or warranty protection for a fee.
+
+  5. Conveying Modified Source Versions.
+
+  You may convey a work based on the Program, or the modifications to
+produce it from the Program, in the form of source code under the
+terms of section 4, provided that you also meet all of these conditions:
+
+    a) The work must carry prominent notices stating that you modified
+    it, and giving a relevant date.
+
+    b) The work must carry prominent notices stating that it is
+    released under this License and any conditions added under section
+    7.  This requirement modifies the requirement in section 4 to
+    "keep intact all notices".
+
+    c) You must license the entire work, as a whole, under this
+    License to anyone who comes into possession of a copy.  This
+    License will therefore apply, along with any applicable section 7
+    additional terms, to the whole of the work, and all its parts,
+    regardless of how they are packaged.  This License gives no
+    permission to license the work in any other way, but it does not
+    invalidate such permission if you have separately received it.
+
+    d) If the work has interactive user interfaces, each must display
+    Appropriate Legal Notices; however, if the Program has interactive
+    interfaces that do not display Appropriate Legal Notices, your
+    work need not make them do so.
+
+  A compilation of a covered work with other separate and independent
+works, which are not by their nature extensions of the covered work,
+and which are not combined with it such as to form a larger program,
+in or on a volume of a storage or distribution medium, is called an
+"aggregate" if the compilation and its resulting copyright are not
+used to limit the access or legal rights of the compilation's users
+beyond what the individual works permit.  Inclusion of a covered work
+in an aggregate does not cause this License to apply to the other
+parts of the aggregate.
+
+  6. Conveying Non-Source Forms.
+
+  You may convey a covered work in object code form under the terms
+of sections 4 and 5, provided that you also convey the
+machine-readable Corresponding Source under the terms of this License,
+in one of these ways:
+
+    a) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by the
+    Corresponding Source fixed on a durable physical medium
+    customarily used for software interchange.
+
+    b) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by a
+    written offer, valid for at least three years and valid for as
+    long as you offer spare parts or customer support for that product
+    model, to give anyone who possesses the object code either (1) a
+    copy of the Corresponding Source for all the software in the
+    product that is covered by this License, on a durable physical
+    medium customarily used for software interchange, for a price no
+    more than your reasonable cost of physically performing this
+    conveying of source, or (2) access to copy the
+    Corresponding Source from a network server at no charge.
+
+    c) Convey individual copies of the object code with a copy of the
+    written offer to provide the Corresponding Source.  This
+    alternative is allowed only occasionally and noncommercially, and
+    only if you received the object code with such an offer, in accord
+    with subsection 6b.
+
+    d) Convey the object code by offering access from a designated
+    place (gratis or for a charge), and offer equivalent access to the
+    Corresponding Source in the same way through the same place at no
+    further charge.  You need not require recipients to copy the
+    Corresponding Source along with the object code.  If the place to
+    copy the object code is a network server, the Corresponding Source
+    may be on a different server (operated by you or a third party)
+    that supports equivalent copying facilities, provided you maintain
+    clear directions next to the object code saying where to find the
+    Corresponding Source.  Regardless of what server hosts the
+    Corresponding Source, you remain obligated to ensure that it is
+    available for as long as needed to satisfy these requirements.
+
+    e) Convey the object code using peer-to-peer transmission, provided
+    you inform other peers where the object code and Corresponding
+    Source of the work are being offered to the general public at no
+    charge under subsection 6d.
+
+  A separable portion of the object code, whose source code is excluded
+from the Corresponding Source as a System Library, need not be
+included in conveying the object code work.
+
+  A "User Product" is either (1) a "consumer product", which means any
+tangible personal property which is normally used for personal, family,
+or household purposes, or (2) anything designed or sold for incorporation
+into a dwelling.  In determining whether a product is a consumer product,
+doubtful cases shall be resolved in favor of coverage.  For a particular
+product received by a particular user, "normally used" refers to a
+typical or common use of that class of product, regardless of the status
+of the particular user or of the way in which the particular user
+actually uses, or expects or is expected to use, the product.  A product
+is a consumer product regardless of whether the product has substantial
+commercial, industrial or non-consumer uses, unless such uses represent
+the only significant mode of use of the product.
+
+  "Installation Information" for a User Product means any methods,
+procedures, authorization keys, or other information required to install
+and execute modified versions of a covered work in that User Product from
+a modified version of its Corresponding Source.  The information must
+suffice to ensure that the continued functioning of the modified object
+code is in no case prevented or interfered with solely because
+modification has been made.
+
+  If you convey an object code work under this section in, or with, or
+specifically for use in, a User Product, and the conveying occurs as
+part of a transaction in which the right of possession and use of the
+User Product is transferred to the recipient in perpetuity or for a
+fixed term (regardless of how the transaction is characterized), the
+Corresponding Source conveyed under this section must be accompanied
+by the Installation Information.  But this requirement does not apply
+if neither you nor any third party retains the ability to install
+modified object code on the User Product (for example, the work has
+been installed in ROM).
+
+  The requirement to provide Installation Information does not include a
+requirement to continue to provide support service, warranty, or updates
+for a work that has been modified or installed by the recipient, or for
+the User Product in which it has been modified or installed.  Access to a
+network may be denied when the modification itself materially and
+adversely affects the operation of the network or violates the rules and
+protocols for communication across the network.
+
+  Corresponding Source conveyed, and Installation Information provided,
+in accord with this section must be in a format that is publicly
+documented (and with an implementation available to the public in
+source code form), and must require no special password or key for
+unpacking, reading or copying.
+
+  7. Additional Terms.
+
+  "Additional permissions" are terms that supplement the terms of this
+License by making exceptions from one or more of its conditions.
+Additional permissions that are applicable to the entire Program shall
+be treated as though they were included in this License, to the extent
+that they are valid under applicable law.  If additional permissions
+apply only to part of the Program, that part may be used separately
+under those permissions, but the entire Program remains governed by
+this License without regard to the additional permissions.
+
+  When you convey a copy of a covered work, you may at your option
+remove any additional permissions from that copy, or from any part of
+it.  (Additional permissions may be written to require their own
+removal in certain cases when you modify the work.)  You may place
+additional permissions on material, added by you to a covered work,
+for which you have or can give appropriate copyright permission.
+
+  Notwithstanding any other provision of this License, for material you
+add to a covered work, you may (if authorized by the copyright holders of
+that material) supplement the terms of this License with terms:
+
+    a) Disclaiming warranty or limiting liability differently from the
+    terms of sections 15 and 16 of this License; or
+
+    b) Requiring preservation of specified reasonable legal notices or
+    author attributions in that material or in the Appropriate Legal
+    Notices displayed by works containing it; or
+
+    c) Prohibiting misrepresentation of the origin of that material, or
+    requiring that modified versions of such material be marked in
+    reasonable ways as different from the original version; or
+
+    d) Limiting the use for publicity purposes of names of licensors or
+    authors of the material; or
+
+    e) Declining to grant rights under trademark law for use of some
+    trade names, trademarks, or service marks; or
+
+    f) Requiring indemnification of licensors and authors of that
+    material by anyone who conveys the material (or modified versions of
+    it) with contractual assumptions of liability to the recipient, for
+    any liability that these contractual assumptions directly impose on
+    those licensors and authors.
+
+  All other non-permissive additional terms are considered "further
+restrictions" within the meaning of section 10.  If the Program as you
+received it, or any part of it, contains a notice stating that it is
+governed by this License along with a term that is a further
+restriction, you may remove that term.  If a license document contains
+a further restriction but permits relicensing or conveying under this
+License, you may add to a covered work material governed by the terms
+of that license document, provided that the further restriction does
+not survive such relicensing or conveying.
+
+  If you add terms to a covered work in accord with this section, you
+must place, in the relevant source files, a statement of the
+additional terms that apply to those files, or a notice indicating
+where to find the applicable terms.
+
+  Additional terms, permissive or non-permissive, may be stated in the
+form of a separately written license, or stated as exceptions;
+the above requirements apply either way.
+
+  8. Termination.
+
+  You may not propagate or modify a covered work except as expressly
+provided under this License.  Any attempt otherwise to propagate or
+modify it is void, and will automatically terminate your rights under
+this License (including any patent licenses granted under the third
+paragraph of section 11).
+
+  However, if you cease all violation of this License, then your
+license from a particular copyright holder is reinstated (a)
+provisionally, unless and until the copyright holder explicitly and
+finally terminates your license, and (b) permanently, if the copyright
+holder fails to notify you of the violation by some reasonable means
+prior to 60 days after the cessation.
+
+  Moreover, your license from a particular copyright holder is
+reinstated permanently if the copyright holder notifies you of the
+violation by some reasonable means, this is the first time you have
+received notice of violation of this License (for any work) from that
+copyright holder, and you cure the violation prior to 30 days after
+your receipt of the notice.
+
+  Termination of your rights under this section does not terminate the
+licenses of parties who have received copies or rights from you under
+this License.  If your rights have been terminated and not permanently
+reinstated, you do not qualify to receive new licenses for the same
+material under section 10.
+
+  9. Acceptance Not Required for Having Copies.
+
+  You are not required to accept this License in order to receive or
+run a copy of the Program.  Ancillary propagation of a covered work
+occurring solely as a consequence of using peer-to-peer transmission
+to receive a copy likewise does not require acceptance.  However,
+nothing other than this License grants you permission to propagate or
+modify any covered work.  These actions infringe copyright if you do
+not accept this License.  Therefore, by modifying or propagating a
+covered work, you indicate your acceptance of this License to do so.
+
+  10. Automatic Licensing of Downstream Recipients.
+
+  Each time you convey a covered work, the recipient automatically
+receives a license from the original licensors, to run, modify and
+propagate that work, subject to this License.  You are not responsible
+for enforcing compliance by third parties with this License.
+
+  An "entity transaction" is a transaction transferring control of an
+organization, or substantially all assets of one, or subdividing an
+organization, or merging organizations.  If propagation of a covered
+work results from an entity transaction, each party to that
+transaction who receives a copy of the work also receives whatever
+licenses to the work the party's predecessor in interest had or could
+give under the previous paragraph, plus a right to possession of the
+Corresponding Source of the work from the predecessor in interest, if
+the predecessor has it or can get it with reasonable efforts.
+
+  You may not impose any further restrictions on the exercise of the
+rights granted or affirmed under this License.  For example, you may
+not impose a license fee, royalty, or other charge for exercise of
+rights granted under this License, and you may not initiate litigation
+(including a cross-claim or counterclaim in a lawsuit) alleging that
+any patent claim is infringed by making, using, selling, offering for
+sale, or importing the Program or any portion of it.
+
+  11. Patents.
+
+  A "contributor" is a copyright holder who authorizes use under this
+License of the Program or a work on which the Program is based.  The
+work thus licensed is called the contributor's "contributor version".
+
+  A contributor's "essential patent claims" are all patent claims
+owned or controlled by the contributor, whether already acquired or
+hereafter acquired, that would be infringed by some manner, permitted
+by this License, of making, using, or selling its contributor version,
+but do not include claims that would be infringed only as a
+consequence of further modification of the contributor version.  For
+purposes of this definition, "control" includes the right to grant
+patent sublicenses in a manner consistent with the requirements of
+this License.
+
+  Each contributor grants you a non-exclusive, worldwide, royalty-free
+patent license under the contributor's essential patent claims, to
+make, use, sell, offer for sale, import and otherwise run, modify and
+propagate the contents of its contributor version.
+
+  In the following three paragraphs, a "patent license" is any express
+agreement or commitment, however denominated, not to enforce a patent
+(such as an express permission to practice a patent or covenant not to
+sue for patent infringement).  To "grant" such a patent license to a
+party means to make such an agreement or commitment not to enforce a
+patent against the party.
+
+  If you convey a covered work, knowingly relying on a patent license,
+and the Corresponding Source of the work is not available for anyone
+to copy, free of charge and under the terms of this License, through a
+publicly available network server or other readily accessible means,
+then you must either (1) cause the Corresponding Source to be so
+available, or (2) arrange to deprive yourself of the benefit of the
+patent license for this particular work, or (3) arrange, in a manner
+consistent with the requirements of this License, to extend the patent
+license to downstream recipients.  "Knowingly relying" means you have
+actual knowledge that, but for the patent license, your conveying the
+covered work in a country, or your recipient's use of the covered work
+in a country, would infringe one or more identifiable patents in that
+country that you have reason to believe are valid.
+
+  If, pursuant to or in connection with a single transaction or
+arrangement, you convey, or propagate by procuring conveyance of, a
+covered work, and grant a patent license to some of the parties
+receiving the covered work authorizing them to use, propagate, modify
+or convey a specific copy of the covered work, then the patent license
+you grant is automatically extended to all recipients of the covered
+work and works based on it.
+
+  A patent license is "discriminatory" if it does not include within
+the scope of its coverage, prohibits the exercise of, or is
+conditioned on the non-exercise of one or more of the rights that are
+specifically granted under this License.  You may not convey a covered
+work if you are a party to an arrangement with a third party that is
+in the business of distributing software, under which you make payment
+to the third party based on the extent of your activity of conveying
+the work, and under which the third party grants, to any of the
+parties who would receive the covered work from you, a discriminatory
+patent license (a) in connection with copies of the covered work
+conveyed by you (or copies made from those copies), or (b) primarily
+for and in connection with specific products or compilations that
+contain the covered work, unless you entered into that arrangement,
+or that patent license was granted, prior to 28 March 2007.
+
+  Nothing in this License shall be construed as excluding or limiting
+any implied license or other defenses to infringement that may
+otherwise be available to you under applicable patent law.
+
+  12. No Surrender of Others' Freedom.
+
+  If conditions are imposed on you (whether by court order, agreement or
+otherwise) that contradict the conditions of this License, they do not
+excuse you from the conditions of this License.  If you cannot convey a
+covered work so as to satisfy simultaneously your obligations under this
+License and any other pertinent obligations, then as a consequence you may
+not convey it at all.  For example, if you agree to terms that obligate you
+to collect a royalty for further conveying from those to whom you convey
+the Program, the only way you could satisfy both those terms and this
+License would be to refrain entirely from conveying the Program.
+
+  13. Remote Network Interaction; Use with the GNU General Public License.
+
+  Notwithstanding any other provision of this License, if you modify the
+Program, your modified version must prominently offer all users
+interacting with it remotely through a computer network (if your version
+supports such interaction) an opportunity to receive the Corresponding
+Source of your version by providing access to the Corresponding Source
+from a network server at no charge, through some standard or customary
+means of facilitating copying of software.  This Corresponding Source
+shall include the Corresponding Source for any work covered by version 3
+of the GNU General Public License that is incorporated pursuant to the
+following paragraph.
+
+  Notwithstanding any other provision of this License, you have
+permission to link or combine any covered work with a work licensed
+under version 3 of the GNU General Public License into a single
+combined work, and to convey the resulting work.  The terms of this
+License will continue to apply to the part which is the covered work,
+but the work with which it is combined will remain governed by version
+3 of the GNU General Public License.
+
+  14. Revised Versions of this License.
+
+  The Free Software Foundation may publish revised and/or new versions of
+the GNU Affero General Public License from time to time.  Such new versions
+will be similar in spirit to the present version, but may differ in detail to
+address new problems or concerns.
+
+  Each version is given a distinguishing version number.  If the
+Program specifies that a certain numbered version of the GNU Affero General
+Public License "or any later version" applies to it, you have the
+option of following the terms and conditions either of that numbered
+version or of any later version published by the Free Software
+Foundation.  If the Program does not specify a version number of the
+GNU Affero General Public License, you may choose any version ever published
+by the Free Software Foundation.
+
+  If the Program specifies that a proxy can decide which future
+versions of the GNU Affero General Public License can be used, that proxy's
+public statement of acceptance of a version permanently authorizes you
+to choose that version for the Program.
+
+  Later license versions may give you additional or different
+permissions.  However, no additional obligations are imposed on any
+author or copyright holder as a result of your choosing to follow a
+later version.
+
+  15. Disclaimer of Warranty.
+
+  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
+APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
+HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
+OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
+IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+  16. Limitation of Liability.
+
+  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
+WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
+THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
+GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
+USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
+DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
+PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
+EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGES.
+
+  17. Interpretation of Sections 15 and 16.
+
+  If the disclaimer of warranty and limitation of liability provided
+above cannot be given local legal effect according to their terms,
+reviewing courts shall apply local law that most closely approximates
+an absolute waiver of all civil liability in connection with the
+Program, unless a warranty or assumption of liability accompanies a
+copy of the Program in return for a fee.
+
+                     END OF TERMS AND CONDITIONS
+
+            How to Apply These Terms to Your New Programs
+
+  If you develop a new program, and you want it to be of the greatest
+possible use to the public, the best way to achieve this is to make it
+free software which everyone can redistribute and change under these terms.
+
+  To do so, attach the following notices to the program.  It is safest
+to attach them to the start of each source file to most effectively
+state the exclusion of warranty; and each file should have at least
+the "copyright" line and a pointer to where the full notice is found.
+
+    <one line to give the program's name and a brief idea of what it does.>
+    Copyright (C) <year>  <name of author>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU Affero General Public License as published
+    by the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU Affero General Public License for more details.
+
+    You should have received a copy of the GNU Affero General Public License
+    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+Also add information on how to contact you by electronic and paper mail.
+
+  If your software can interact with users remotely through a computer
+network, you should also make sure that it provides a way for users to
+get its source.  For example, if your program is a web application, its
+interface could display a "Source" link that leads users to an archive
+of the code.  There are many ways you could offer source, and different
+solutions will be better for different programs; see section 13 for the
+specific requirements.
+
+  You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU AGPL, see
+<http://www.gnu.org/licenses/>.

extensions/adetailer/README.md

@@ -0,0 +1,105 @@
+# !After Detailer
+
+!After Detailer is a extension for stable diffusion webui, similar to Detection Detailer, except it uses ultralytics instead of the mmdet.
+
+## Install
+
+(from Mikubill/sd-webui-controlnet)
+
+1. Open "Extensions" tab.
+2. Open "Install from URL" tab in the tab.
+3. Enter `https://github.com/Bing-su/adetailer.git` to "URL for extension's git repository".
+4. Press "Install" button.
+5. Wait 5 seconds, and you will see the message "Installed into stable-diffusion-webui\extensions\adetailer. Use Installed tab to restart".
+6. Go to "Installed" tab, click "Check for updates", and then click "Apply and restart UI". (The next time you can also use this method to update extensions.)
+7. Completely restart A1111 webui including your terminal. (If you do not know what is a "terminal", you can reboot your computer: turn your computer off and turn it on again.)
+
+You can now install it directly from the Extensions tab.
+
+![image](https://i.imgur.com/g6GdRBT.png)
+
+You **DON'T** need to download any model from huggingface.
+
+## Options
+
+| Model, Prompts                    |                                       |                                                   |
+| --------------------------------- | ------------------------------------- | ------------------------------------------------- |
+| ADetailer model                   | Determine what to detect.             | `None` = disable                                  |
+| ADetailer prompt, negative prompt | Prompts and negative prompts to apply | If left blank, it will use the same as the input. |
+
+| Detection                            |                                                                                              |     |
+| ------------------------------------ | -------------------------------------------------------------------------------------------- | --- |
+| Detection model confidence threshold | Only objects with a detection model confidence above this threshold are used for inpainting. |     |
+| Mask min/max ratio                   | Only use masks whose area is between those ratios for the area of the entire image.          |     |
+
+If you want to exclude objects in the background, try setting the min ratio to around `0.01`.
+
+| Mask Preprocessing              |                                                                                                                                     |                                                                                         |
+| ------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------- |
+| Mask x, y offset                | Moves the mask horizontally and vertically by                                                                                       |                                                                                         |
+| Mask erosion (-) / dilation (+) | Enlarge or reduce the detected mask.                                                                                                | [opencv example](https://docs.opencv.org/4.7.0/db/df6/tutorial_erosion_dilatation.html) |
+| Mask merge mode                 | `None`: Inpaint each mask<br/>`Merge`: Merge all masks and inpaint<br/>`Merge and Invert`: Merge all masks and Invert, then inpaint |                                                                                         |
+
+Applied in this order: x, y offset → erosion/dilation → merge/invert.
+
+#### Inpainting
+
+![image](https://i.imgur.com/wyWlT1n.png)
+
+Each option corresponds to a corresponding option on the inpaint tab.
+
+## ControlNet Inpainting
+
+You can use the ControlNet extension if you have ControlNet installed and ControlNet models.
+
+Support `inpaint, scribble, lineart, openpose, tile` controlnet models. Once you choose a model, the preprocessor is set automatically.
+
+## Model
+
+| Model                 | Target                | mAP 50                        | mAP 50-95                     |
+| --------------------- | --------------------- | ----------------------------- | ----------------------------- |
+| face_yolov8n.pt       | 2D / realistic face   | 0.660                         | 0.366                         |
+| face_yolov8s.pt       | 2D / realistic face   | 0.713                         | 0.404                         |
+| hand_yolov8n.pt       | 2D / realistic hand   | 0.767                         | 0.505                         |
+| person_yolov8n-seg.pt | 2D / realistic person | 0.782 (bbox)<br/>0.761 (mask) | 0.555 (bbox)<br/>0.460 (mask) |
+| person_yolov8s-seg.pt | 2D / realistic person | 0.824 (bbox)<br/>0.809 (mask) | 0.605 (bbox)<br/>0.508 (mask) |
+| mediapipe_face_full   | realistic face        | -                             | -                             |
+| mediapipe_face_short  | realistic face        | -                             | -                             |
+| mediapipe_face_mesh   | realistic face        | -                             | -                             |
+
+The yolo models can be found on huggingface [Bingsu/adetailer](https://huggingface.co/Bingsu/adetailer).
+
+### User Model
+
+Put your [ultralytics](https://github.com/ultralytics/ultralytics) model in `webui/models/adetailer`. The model name should end with `.pt` or `.pth`.
+
+It must be a bbox detection or segment model and use all label.
+
+### Dataset
+
+Datasets used for training the yolo models are:
+
+#### Face
+
+- [Anime Face CreateML](https://universe.roboflow.com/my-workspace-mph8o/anime-face-createml)
+- [xml2txt](https://universe.roboflow.com/0oooooo0/xml2txt-njqx1)
+- [AN](https://universe.roboflow.com/sed-b8vkf/an-lfg5i)
+- [wider face](http://shuoyang1213.me/WIDERFACE/index.html)
+
+#### Hand
+
+- [AnHDet](https://universe.roboflow.com/1-yshhi/anhdet)
+- [hand-detection-fuao9](https://universe.roboflow.com/catwithawand/hand-detection-fuao9)
+
+#### Person
+
+- [coco2017](https://cocodataset.org/#home) (only person)
+- [AniSeg](https://github.com/jerryli27/AniSeg)
+- [skytnt/anime-segmentation](https://huggingface.co/datasets/skytnt/anime-segmentation)
+
+## Example
+
+![image](https://i.imgur.com/38RSxSO.png)
+![image](https://i.imgur.com/2CYgjLx.png)
+
+[![ko-fi](https://ko-fi.com/img/githubbutton_sm.svg)](https://ko-fi.com/F1F1L7V2N)

extensions/adetailer/adetailer/__init__.py

@@ -0,0 +1,20 @@
+from .__version__ import __version__
+from .args import AD_ENABLE, ALL_ARGS, ADetailerArgs, EnableChecker
+from .common import PredictOutput, get_models
+from .mediapipe import mediapipe_predict
+from .ultralytics import ultralytics_predict
+
+AFTER_DETAILER = "ADetailer"
+
+__all__ = [
+    "__version__",
+    "AD_ENABLE",
+    "ADetailerArgs",
+    "AFTER_DETAILER",
+    "ALL_ARGS",
+    "EnableChecker",
+    "PredictOutput",
+    "get_models",
+    "mediapipe_predict",
+    "ultralytics_predict",
+]

extensions/adetailer/adetailer/__version__.py

@@ -0,0 +1 @@
+__version__ = "23.6.1.post0"

extensions/adetailer/adetailer/args.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+from collections import UserList
+from functools import cached_property, partial
+from typing import Any, Literal, NamedTuple, Union
+
+import pydantic
+from pydantic import (
+    BaseModel,
+    Extra,
+    NonNegativeFloat,
+    NonNegativeInt,
+    PositiveInt,
+    confloat,
+    constr,
+)
+
+cn_model_regex = r".*(inpaint|tile|scribble|lineart|openpose).*|^None$"
+
+
+class Arg(NamedTuple):
+    attr: str
+    name: str
+
+
+class ArgsList(UserList):
+    @cached_property
+    def attrs(self) -> tuple[str]:
+        return tuple(attr for attr, _ in self)
+
+    @cached_property
+    def names(self) -> tuple[str]:
+        return tuple(name for _, name in self)
+
+
+class ADetailerArgs(BaseModel, extra=Extra.forbid):
+    ad_model: str = "None"
+    ad_prompt: str = ""
+    ad_negative_prompt: str = ""
+    ad_confidence: confloat(ge=0.0, le=1.0) = 0.3
+    ad_mask_min_ratio: confloat(ge=0.0, le=1.0) = 0.0
+    ad_mask_max_ratio: confloat(ge=0.0, le=1.0) = 1.0
+    ad_dilate_erode: int = 32
+    ad_x_offset: int = 0
+    ad_y_offset: int = 0
+    ad_mask_merge_invert: Literal["None", "Merge", "Merge and Invert"] = "None"
+    ad_mask_blur: NonNegativeInt = 4
+    ad_denoising_strength: confloat(ge=0.0, le=1.0) = 0.4
+    ad_inpaint_only_masked: bool = True
+    ad_inpaint_only_masked_padding: NonNegativeInt = 0
+    ad_use_inpaint_width_height: bool = False
+    ad_inpaint_width: PositiveInt = 512
+    ad_inpaint_height: PositiveInt = 512
+    ad_use_steps: bool = False
+    ad_steps: PositiveInt = 28
+    ad_use_cfg_scale: bool = False
+    ad_cfg_scale: NonNegativeFloat = 7.0
+    ad_restore_face: bool = False
+    ad_controlnet_model: constr(regex=cn_model_regex) = "None"
+    ad_controlnet_weight: confloat(ge=0.0, le=1.0) = 1.0
+    ad_controlnet_guidance_start: confloat(ge=0.0, le=1.0) = 0.0
+    ad_controlnet_guidance_end: confloat(ge=0.0, le=1.0) = 1.0
+
+    @staticmethod
+    def ppop(
+        p: dict[str, Any],
+        key: str,
+        pops: list[str] | None = None,
+        cond: Any = None,
+    ):
+        if pops is None:
+            pops = [key]
+        if key not in p:
+            return
+        value = p[key]
+        cond = (not bool(value)) if cond is None else value == cond
+
+        if cond:
+            for k in pops:
+                p.pop(k, None)
+
+    def extra_params(self, suffix: str = ""):
+        if self.ad_model == "None":
+            return {}
+
+        p = {name: getattr(self, attr) for attr, name in ALL_ARGS}
+        ppop = partial(self.ppop, p)
+
+        ppop("ADetailer prompt")
+        ppop("ADetailer negative prompt")
+        ppop("ADetailer mask min ratio", cond=0.0)
+        ppop("ADetailer mask max ratio", cond=1.0)
+        ppop("ADetailer x offset", cond=0)
+        ppop("ADetailer y offset", cond=0)
+        ppop("ADetailer mask merge/invert", cond="None")
+        ppop("ADetailer inpaint only masked", ["ADetailer inpaint padding"])
+        ppop(
+            "ADetailer use inpaint width/height",
+            [
+                "ADetailer use inpaint width/height",
+                "ADetailer inpaint width",
+                "ADetailer inpaint height",
+            ],
+        )
+        ppop(
+            "ADetailer use separate steps",
+            ["ADetailer use separate steps", "ADetailer steps"],
+        )
+        ppop(
+            "ADetailer use separate CFG scale",
+            ["ADetailer use separate CFG scale", "ADetailer CFG scale"],
+        )
+        ppop("ADetailer restore face")
+        ppop(
+            "ADetailer ControlNet model",
+            [
+                "ADetailer ControlNet model",
+                "ADetailer ControlNet weight",
+                "ADetailer ControlNet guidance start",
+                "ADetailer ControlNet guidance end",
+            ],
+            cond="None",
+        )
+        ppop("ADetailer ControlNet weight", cond=1.0)
+        ppop("ADetailer ControlNet guidance start", cond=0.0)
+        ppop("ADetailer ControlNet guidance end", cond=1.0)
+
+        if suffix:
+            p = {k + suffix: v for k, v in p.items()}
+
+        return p
+
+
+class EnableChecker(BaseModel):
+    a0: Union[bool, dict]
+    a1: Any
+
+    def is_enabled(self) -> bool:
+        ad_model = ALL_ARGS[0].attr
+        if isinstance(self.a0, dict):
+            return self.a0.get(ad_model, "None") != "None"
+        if not isinstance(self.a1, dict):
+            return False
+        return self.a0 and self.a1.get(ad_model, "None") != "None"
+
+
+_all_args = [
+    ("ad_enable", "ADetailer enable"),
+    ("ad_model", "ADetailer model"),
+    ("ad_prompt", "ADetailer prompt"),
+    ("ad_negative_prompt", "ADetailer negative prompt"),
+    ("ad_confidence", "ADetailer confidence"),
+    ("ad_mask_min_ratio", "ADetailer mask min ratio"),
+    ("ad_mask_max_ratio", "ADetailer mask max ratio"),
+    ("ad_x_offset", "ADetailer x offset"),
+    ("ad_y_offset", "ADetailer y offset"),
+    ("ad_dilate_erode", "ADetailer dilate/erode"),
+    ("ad_mask_merge_invert", "ADetailer mask merge/invert"),
+    ("ad_mask_blur", "ADetailer mask blur"),
+    ("ad_denoising_strength", "ADetailer denoising strength"),
+    ("ad_inpaint_only_masked", "ADetailer inpaint only masked"),
+    ("ad_inpaint_only_masked_padding", "ADetailer inpaint padding"),
+    ("ad_use_inpaint_width_height", "ADetailer use inpaint width/height"),
+    ("ad_inpaint_width", "ADetailer inpaint width"),
+    ("ad_inpaint_height", "ADetailer inpaint height"),
+    ("ad_use_steps", "ADetailer use separate steps"),
+    ("ad_steps", "ADetailer steps"),
+    ("ad_use_cfg_scale", "ADetailer use separate CFG scale"),
+    ("ad_cfg_scale", "ADetailer CFG scale"),
+    ("ad_restore_face", "ADetailer restore face"),
+    ("ad_controlnet_model", "ADetailer ControlNet model"),
+    ("ad_controlnet_weight", "ADetailer ControlNet weight"),
+    ("ad_controlnet_guidance_start", "ADetailer ControlNet guidance start"),
+    ("ad_controlnet_guidance_end", "ADetailer ControlNet guidance end"),
+]
+
+AD_ENABLE = Arg(*_all_args[0])
+_args = [Arg(*args) for args in _all_args[1:]]
+ALL_ARGS = ArgsList(_args)
+
+BBOX_SORTBY = [
+    "None",
+    "Position (left to right)",
+    "Position (center to edge)",
+    "Area (large to small)",
+]
+MASK_MERGE_INVERT = ["None", "Merge", "Merge and Invert"]

extensions/adetailer/adetailer/common.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional, Union
+
+from huggingface_hub import hf_hub_download
+from PIL import Image, ImageDraw
+
+repo_id = "Bingsu/adetailer"
+
+
+@dataclass
+class PredictOutput:
+    bboxes: list[list[int | float]] = field(default_factory=list)
+    masks: list[Image.Image] = field(default_factory=list)
+    preview: Optional[Image.Image] = None
+
+
+def hf_download(file: str):
+    try:
+        path = hf_hub_download(repo_id, file)
+    except Exception:
+        path = "INVALID"
+    return path
+
+
+def get_models(
+    model_dir: Union[str, Path], huggingface: bool = True
+) -> OrderedDict[str, Optional[str]]:
+    model_dir = Path(model_dir)
+    if model_dir.is_dir():
+        model_paths = [
+            p
+            for p in model_dir.rglob("*")
+            if p.is_file() and p.suffix in (".pt", ".pth")
+        ]
+    else:
+        model_paths = []
+
+    models = OrderedDict()
+    if huggingface:
+        models.update(
+            {
+                "face_yolov8n.pt": hf_download("face_yolov8n.pt"),
+                "face_yolov8s.pt": hf_download("face_yolov8s.pt"),
+                "hand_yolov8n.pt": hf_download("hand_yolov8n.pt"),
+                "person_yolov8n-seg.pt": hf_download("person_yolov8n-seg.pt"),
+                "person_yolov8s-seg.pt": hf_download("person_yolov8s-seg.pt"),
+            }
+        )
+    models.update(
+        {
+            "mediapipe_face_full": None,
+            "mediapipe_face_short": None,
+            "mediapipe_face_mesh": None,
+        }
+    )
+
+    invalid_keys = [k for k, v in models.items() if v == "INVALID"]
+    for key in invalid_keys:
+        models.pop(key)
+
+    for path in model_paths:
+        if path.name in models:
+            continue
+        models[path.name] = str(path)
+
+    return models
+
+
+def create_mask_from_bbox(
+    bboxes: list[list[float]], shape: tuple[int, int]
+) -> list[Image.Image]:
+    """
+    Parameters
+    ----------
+        bboxes: list[list[float]]
+            list of [x1, y1, x2, y2]
+            bounding boxes
+        shape: tuple[int, int]
+            shape of the image (width, height)
+
+    Returns
+    -------
+        masks: list[Image.Image]
+        A list of masks
+
+    """
+    masks = []
+    for bbox in bboxes:
+        mask = Image.new("L", shape, 0)
+        mask_draw = ImageDraw.Draw(mask)
+        mask_draw.rectangle(bbox, fill=255)
+        masks.append(mask)
+    return masks
+
+
+def create_bbox_from_mask(
+    masks: list[Image.Image], shape: tuple[int, int]
+) -> list[list[int]]:
+    """
+    Parameters
+    ----------
+        masks: list[Image.Image]
+            A list of masks
+        shape: tuple[int, int]
+            shape of the image (width, height)
+
+    Returns
+    -------
+        bboxes: list[list[float]]
+        A list of bounding boxes
+
+    """
+    bboxes = []
+    for mask in masks:
+        mask = mask.resize(shape)
+        bbox = mask.getbbox()
+        if bbox is not None:
+            bboxes.append(list(bbox))
+    return bboxes

extensions/adetailer/adetailer/mask.py

@@ -0,0 +1,247 @@
+from __future__ import annotations
+
+from enum import IntEnum
+from functools import partial, reduce
+from math import dist
+
+import cv2
+import numpy as np
+from PIL import Image, ImageChops
+
+from adetailer.args import MASK_MERGE_INVERT
+from adetailer.common import PredictOutput
+
+
+class SortBy(IntEnum):
+    NONE = 0
+    LEFT_TO_RIGHT = 1
+    CENTER_TO_EDGE = 2
+    AREA = 3
+
+
+class MergeInvert(IntEnum):
+    NONE = 0
+    MERGE = 1
+    MERGE_INVERT = 2
+
+
+def _dilate(arr: np.ndarray, value: int) -> np.ndarray:
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (value, value))
+    return cv2.dilate(arr, kernel, iterations=1)
+
+
+def _erode(arr: np.ndarray, value: int) -> np.ndarray:
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (value, value))
+    return cv2.erode(arr, kernel, iterations=1)
+
+
+def dilate_erode(img: Image.Image, value: int) -> Image.Image:
+    """
+    The dilate_erode function takes an image and a value.
+    If the value is positive, it dilates the image by that amount.
+    If the value is negative, it erodes the image by that amount.
+
+    Parameters
+    ----------
+        img: PIL.Image.Image
+            the image to be processed
+        value: int
+            kernel size of dilation or erosion
+
+    Returns
+    -------
+        PIL.Image.Image
+            The image that has been dilated or eroded
+    """
+    if value == 0:
+        return img
+
+    arr = np.array(img)
+    arr = _dilate(arr, value) if value > 0 else _erode(arr, -value)
+
+    return Image.fromarray(arr)
+
+
+def offset(img: Image.Image, x: int = 0, y: int = 0) -> Image.Image:
+    """
+    The offset function takes an image and offsets it by a given x(→) and y(↑) value.
+
+    Parameters
+    ----------
+        mask: Image.Image
+            Pass the mask image to the function
+        x: int
+            →
+        y: int
+            ↑
+
+    Returns
+    -------
+        PIL.Image.Image
+            A new image that is offset by x and y
+    """
+    return ImageChops.offset(img, x, -y)
+
+
+def is_all_black(img: Image.Image) -> bool:
+    arr = np.array(img)
+    return cv2.countNonZero(arr) == 0
+
+
+def bbox_area(bbox: list[float]):
+    return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+
+
+def mask_preprocess(
+    masks: list[Image.Image],
+    kernel: int = 0,
+    x_offset: int = 0,
+    y_offset: int = 0,
+    merge_invert: int | MergeInvert | str = MergeInvert.NONE,
+) -> list[Image.Image]:
+    """
+    The mask_preprocess function takes a list of masks and preprocesses them.
+    It dilates and erodes the masks, and offsets them by x_offset and y_offset.
+
+    Parameters
+    ----------
+        masks: list[Image.Image]
+            A list of masks
+        kernel: int
+            kernel size of dilation or erosion
+        x_offset: int
+            →
+        y_offset: int
+            ↑
+
+    Returns
+    -------
+        list[Image.Image]
+            A list of processed masks
+    """
+    if not masks:
+        return []
+
+    if x_offset != 0 or y_offset != 0:
+        masks = [offset(m, x_offset, y_offset) for m in masks]
+
+    if kernel != 0:
+        masks = [dilate_erode(m, kernel) for m in masks]
+        masks = [m for m in masks if not is_all_black(m)]
+
+    masks = mask_merge_invert(masks, mode=merge_invert)
+
+    return masks
+
+
+# Bbox sorting
+def _key_left_to_right(bbox: list[float]) -> float:
+    """
+    Left to right
+
+    Parameters
+    ----------
+    bbox: list[float]
+        list of [x1, y1, x2, y2]
+    """
+    return bbox[0]
+
+
+def _key_center_to_edge(bbox: list[float], *, center: tuple[float, float]) -> float:
+    """
+    Center to edge
+
+    Parameters
+    ----------
+    bbox: list[float]
+        list of [x1, y1, x2, y2]
+    image: Image.Image
+        the image
+    """
+    bbox_center = ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2)
+    return dist(center, bbox_center)
+
+
+def _key_area(bbox: list[float]) -> float:
+    """
+    Large to small
+
+    Parameters
+    ----------
+    bbox: list[float]
+        list of [x1, y1, x2, y2]
+    """
+    return -bbox_area(bbox)
+
+
+def sort_bboxes(
+    pred: PredictOutput, order: int | SortBy = SortBy.NONE
+) -> PredictOutput:
+    if order == SortBy.NONE or len(pred.bboxes) <= 1:
+        return pred
+
+    if order == SortBy.LEFT_TO_RIGHT:
+        key = _key_left_to_right
+    elif order == SortBy.CENTER_TO_EDGE:
+        width, height = pred.preview.size
+        center = (width / 2, height / 2)
+        key = partial(_key_center_to_edge, center=center)
+    elif order == SortBy.AREA:
+        key = _key_area
+    else:
+        raise RuntimeError
+
+    items = len(pred.bboxes)
+    idx = sorted(range(items), key=lambda i: key(pred.bboxes[i]))
+    pred.bboxes = [pred.bboxes[i] for i in idx]
+    pred.masks = [pred.masks[i] for i in idx]
+    return pred
+
+
+# Filter by ratio
+def is_in_ratio(bbox: list[float], low: float, high: float, orig_area: int) -> bool:
+    area = bbox_area(bbox)
+    return low <= area / orig_area <= high
+
+
+def filter_by_ratio(pred: PredictOutput, low: float, high: float) -> PredictOutput:
+    if not pred.bboxes:
+        return pred
+
+    w, h = pred.preview.size
+    orig_area = w * h
+    items = len(pred.bboxes)
+    idx = [i for i in range(items) if is_in_ratio(pred.bboxes[i], low, high, orig_area)]
+    pred.bboxes = [pred.bboxes[i] for i in idx]
+    pred.masks = [pred.masks[i] for i in idx]
+    return pred
+
+
+# Merge / Invert
+def mask_merge(masks: list[Image.Image]) -> list[Image.Image]:
+    arrs = [np.array(m) for m in masks]
+    arr = reduce(cv2.bitwise_or, arrs)
+    return [Image.fromarray(arr)]
+
+
+def mask_invert(masks: list[Image.Image]) -> list[Image.Image]:
+    return [ImageChops.invert(m) for m in masks]
+
+
+def mask_merge_invert(
+    masks: list[Image.Image], mode: int | MergeInvert | str
+) -> list[Image.Image]:
+    if isinstance(mode, str):
+        mode = MASK_MERGE_INVERT.index(mode)
+
+    if mode == MergeInvert.NONE or not masks:
+        return masks
+
+    if mode == MergeInvert.MERGE:
+        return mask_merge(masks)
+
+    if mode == MergeInvert.MERGE_INVERT:
+        merged = mask_merge(masks)
+        return mask_invert(merged)
+
+    raise RuntimeError

extensions/adetailer/adetailer/mediapipe.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from functools import partial
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+from adetailer import PredictOutput
+from adetailer.common import create_bbox_from_mask, create_mask_from_bbox
+
+
+def mediapipe_predict(
+    model_type: str, image: Image.Image, confidence: float = 0.3
+) -> PredictOutput:
+    mapping = {
+        "mediapipe_face_short": partial(mediapipe_face_detection, 0),
+        "mediapipe_face_full": partial(mediapipe_face_detection, 1),
+        "mediapipe_face_mesh": mediapipe_face_mesh,
+    }
+    if model_type in mapping:
+        func = mapping[model_type]
+        return func(image, confidence)
+    raise RuntimeError(f"[-] ADetailer: Invalid mediapipe model type: {model_type}")
+
+
+def mediapipe_face_detection(
+    model_type: int, image: Image.Image, confidence: float = 0.3
+) -> PredictOutput:
+    import mediapipe as mp
+
+    img_width, img_height = image.size
+
+    mp_face_detection = mp.solutions.face_detection
+    draw_util = mp.solutions.drawing_utils
+
+    img_array = np.array(image)
+
+    with mp_face_detection.FaceDetection(
+        model_selection=model_type, min_detection_confidence=confidence
+    ) as face_detector:
+        pred = face_detector.process(img_array)
+
+    if pred.detections is None:
+        return PredictOutput()
+
+    preview_array = img_array.copy()
+
+    bboxes = []
+    for detection in pred.detections:
+        draw_util.draw_detection(preview_array, detection)
+
+        bbox = detection.location_data.relative_bounding_box
+        x1 = bbox.xmin * img_width
+        y1 = bbox.ymin * img_height
+        w = bbox.width * img_width
+        h = bbox.height * img_height
+        x2 = x1 + w
+        y2 = y1 + h
+
+        bboxes.append([x1, y1, x2, y2])
+
+    masks = create_mask_from_bbox(bboxes, image.size)
+    preview = Image.fromarray(preview_array)
+
+    return PredictOutput(bboxes=bboxes, masks=masks, preview=preview)
+
+
+def mediapipe_face_mesh(image: Image.Image, confidence: float = 0.3) -> PredictOutput:
+    import mediapipe as mp
+    from scipy.spatial import ConvexHull
+
+    mp_face_mesh = mp.solutions.face_mesh
+    draw_util = mp.solutions.drawing_utils
+    drawing_styles = mp.solutions.drawing_styles
+
+    w, h = image.size
+
+    with mp_face_mesh.FaceMesh(
+        static_image_mode=True, max_num_faces=20, min_detection_confidence=confidence
+    ) as face_mesh:
+        arr = np.array(image)
+        pred = face_mesh.process(arr)
+
+        if pred.multi_face_landmarks is None:
+            return PredictOutput()
+
+        preview = arr.copy()
+        masks = []
+
+        for landmarks in pred.multi_face_landmarks:
+            draw_util.draw_landmarks(
+                image=preview,
+                landmark_list=landmarks,
+                connections=mp_face_mesh.FACEMESH_TESSELATION,
+                landmark_drawing_spec=None,
+                connection_drawing_spec=drawing_styles.get_default_face_mesh_tesselation_style(),
+            )
+
+            points = np.array([(land.x * w, land.y * h) for land in landmarks.landmark])
+            hull = ConvexHull(points)
+            vertices = hull.vertices
+            outline = list(zip(points[vertices, 0], points[vertices, 1]))
+
+            mask = Image.new("L", image.size, "black")
+            draw = ImageDraw.Draw(mask)
+            draw.polygon(outline, fill="white")
+            masks.append(mask)
+
+        bboxes = create_bbox_from_mask(masks, image.size)
+        preview = Image.fromarray(preview)
+        return PredictOutput(bboxes=bboxes, masks=masks, preview=preview)

extensions/adetailer/adetailer/ui.py

@@ -0,0 +1,451 @@
+from __future__ import annotations
+
+from functools import partial
+from types import SimpleNamespace
+from typing import Any
+
+import gradio as gr
+
+from adetailer import AFTER_DETAILER, __version__
+from adetailer.args import AD_ENABLE, ALL_ARGS, MASK_MERGE_INVERT
+from controlnet_ext import controlnet_exists, get_cn_models
+
+
+class Widgets(SimpleNamespace):
+    def tolist(self):
+        return [getattr(self, attr) for attr in ALL_ARGS.attrs]
+
+
+def gr_interactive(value: bool = True):
+    return gr.update(interactive=value)
+
+
+def ordinal(n: int) -> str:
+    d = {1: "st", 2: "nd", 3: "rd"}
+    return str(n) + ("th" if 11 <= n % 100 <= 13 else d.get(n % 10, "th"))
+
+
+def suffix(n: int, c: str = " ") -> str:
+    return "" if n == 0 else c + ordinal(n + 1)
+
+
+def on_widget_change(state: dict, value: Any, *, attr: str):
+    state[attr] = value
+    return state
+
+
+def on_generate_click(state: dict, *values: Any):
+    for attr, value in zip(ALL_ARGS.attrs, values):
+        state[attr] = value
+    return state
+
+
+def elem_id(item_id: str, n: int, is_img2img: bool) -> str:
+    tap = "img2img" if is_img2img else "txt2img"
+    suf = suffix(n, "_")
+    return f"script_{tap}_adetailer_{item_id}{suf}"
+
+
+def adui(
+    num_models: int,
+    is_img2img: bool,
+    model_list: list[str],
+    t2i_button: gr.Button,
+    i2i_button: gr.Button,
+):
+    states = []
+    infotext_fields = []
+    eid = partial(elem_id, n=0, is_img2img=is_img2img)
+
+    with gr.Accordion(AFTER_DETAILER, open=False, elem_id=eid("ad_main_accordion")):
+        with gr.Row():
+            with gr.Column(scale=6):
+                ad_enable = gr.Checkbox(
+                    label="Enable ADetailer",
+                    value=False,
+                    visible=True,
+                    elem_id=eid("ad_enable"),
+                )
+
+            with gr.Column(scale=1, min_width=180):
+                gr.Markdown(
+                    f"v{__version__}",
+                    elem_id=eid("ad_version"),
+                )
+
+        infotext_fields.append((ad_enable, AD_ENABLE.name))
+
+        with gr.Group(), gr.Tabs():
+            for n in range(num_models):
+                with gr.Tab(ordinal(n + 1)):
+                    state, infofields = one_ui_group(
+                        n=n,
+                        is_img2img=is_img2img,
+                        model_list=model_list,
+                        t2i_button=t2i_button,
+                        i2i_button=i2i_button,
+                    )
+
+                states.append(state)
+                infotext_fields.extend(infofields)
+
+    # components: [bool, dict, dict, ...]
+    components = [ad_enable] + states
+    return components, infotext_fields
+
+
+def one_ui_group(
+    n: int,
+    is_img2img: bool,
+    model_list: list[str],
+    t2i_button: gr.Button,
+    i2i_button: gr.Button,
+):
+    w = Widgets()
+    state = gr.State({})
+    eid = partial(elem_id, n=n, is_img2img=is_img2img)
+
+    with gr.Row():
+        model_choices = model_list + ["None"] if n == 0 else ["None"] + model_list
+
+        w.ad_model = gr.Dropdown(
+            label="ADetailer model" + suffix(n),
+            choices=model_choices,
+            value=model_choices[0],
+            visible=True,
+            type="value",
+            elem_id=eid("ad_model"),
+        )
+
+    with gr.Group():
+        with gr.Row(elem_id=eid("ad_toprow_prompt")):
+            w.ad_prompt = gr.Textbox(
+                label="ad_prompt" + suffix(n),
+                show_label=False,
+                lines=3,
+                placeholder="ADetailer prompt" + suffix(n),
+                elem_id=eid("ad_prompt"),
+            )
+
+        with gr.Row(elem_id=eid("ad_toprow_negative_prompt")):
+            w.ad_negative_prompt = gr.Textbox(
+                label="ad_negative_prompt" + suffix(n),
+                show_label=False,
+                lines=2,
+                placeholder="ADetailer negative prompt" + suffix(n),
+                elem_id=eid("ad_negative_prompt"),
+            )
+
+    with gr.Group():
+        with gr.Accordion(
+            "Detection", open=False, elem_id=eid("ad_detection_accordion")
+        ):
+            detection(w, n, is_img2img)
+
+        with gr.Accordion(
+            "Mask Preprocessing",
+            open=False,
+            elem_id=eid("ad_mask_preprocessing_accordion"),
+        ):
+            mask_preprocessing(w, n, is_img2img)
+
+        with gr.Accordion(
+            "Inpainting", open=False, elem_id=eid("ad_inpainting_accordion")
+        ):
+            inpainting(w, n, is_img2img)
+
+    with gr.Group():
+        controlnet(w, n, is_img2img)
+
+    for attr in ALL_ARGS.attrs:
+        widget = getattr(w, attr)
+        on_change = partial(on_widget_change, attr=attr)
+        widget.change(
+            fn=on_change, inputs=[state, widget], outputs=[state], queue=False
+        )
+
+    all_inputs = [state] + w.tolist()
+    target_button = i2i_button if is_img2img else t2i_button
+    target_button.click(
+        fn=on_generate_click, inputs=all_inputs, outputs=state, queue=False
+    )
+
+    infotext_fields = [(getattr(w, attr), name + suffix(n)) for attr, name in ALL_ARGS]
+
+    return state, infotext_fields
+
+
+def detection(w: Widgets, n: int, is_img2img: bool):
+    eid = partial(elem_id, n=n, is_img2img=is_img2img)
+
+    with gr.Row():
+        with gr.Column():
+            w.ad_confidence = gr.Slider(
+                label="Detection model confidence threshold" + suffix(n),
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.3,
+                visible=True,
+                elem_id=eid("ad_confidence"),
+            )
+
+        with gr.Column(variant="compact"):
+            w.ad_mask_min_ratio = gr.Slider(
+                label="Mask min area ratio" + suffix(n),
+                minimum=0.0,
+                maximum=1.0,
+                step=0.001,
+                value=0.0,
+                visible=True,
+                elem_id=eid("ad_mask_min_ratio"),
+            )
+            w.ad_mask_max_ratio = gr.Slider(
+                label="Mask max area ratio" + suffix(n),
+                minimum=0.0,
+                maximum=1.0,
+                step=0.001,
+                value=1.0,
+                visible=True,
+                elem_id=eid("ad_mask_max_ratio"),
+            )
+
+
+def mask_preprocessing(w: Widgets, n: int, is_img2img: bool):
+    eid = partial(elem_id, n=n, is_img2img=is_img2img)
+
+    with gr.Group():
+        with gr.Row():
+            with gr.Column(variant="compact"):
+                w.ad_x_offset = gr.Slider(
+                    label="Mask x(→) offset" + suffix(n),
+                    minimum=-200,
+                    maximum=200,
+                    step=1,
+                    value=0,
+                    visible=True,
+                    elem_id=eid("ad_x_offset"),
+                )
+                w.ad_y_offset = gr.Slider(
+                    label="Mask y(↑) offset" + suffix(n),
+                    minimum=-200,
+                    maximum=200,
+                    step=1,
+                    value=0,
+                    visible=True,
+                    elem_id=eid("ad_y_offset"),
+                )
+
+            with gr.Column(variant="compact"):
+                w.ad_dilate_erode = gr.Slider(
+                    label="Mask erosion (-) / dilation (+)" + suffix(n),
+                    minimum=-128,
+                    maximum=128,
+                    step=4,
+                    value=4,
+                    visible=True,
+                    elem_id=eid("ad_dilate_erode"),
+                )
+
+        with gr.Row():
+            w.ad_mask_merge_invert = gr.Radio(
+                label="Mask merge mode" + suffix(n),
+                choices=MASK_MERGE_INVERT,
+                value="None",
+                elem_id=eid("ad_mask_merge_invert"),
+            )
+
+
+def inpainting(w: Widgets, n: int, is_img2img: bool):
+    eid = partial(elem_id, n=n, is_img2img=is_img2img)
+
+    with gr.Group():
+        with gr.Row():
+            w.ad_mask_blur = gr.Slider(
+                label="Inpaint mask blur" + suffix(n),
+                minimum=0,
+                maximum=64,
+                step=1,
+                value=4,
+                visible=True,
+                elem_id=eid("ad_mask_blur"),
+            )
+
+            w.ad_denoising_strength = gr.Slider(
+                label="Inpaint denoising strength" + suffix(n),
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.4,
+                visible=True,
+                elem_id=eid("ad_denoising_strength"),
+            )
+
+        with gr.Row():
+            with gr.Column(variant="compact"):
+                w.ad_inpaint_only_masked = gr.Checkbox(
+                    label="Inpaint only masked" + suffix(n),
+                    value=True,
+                    visible=True,
+                    elem_id=eid("ad_inpaint_full_res"),
+                )
+                w.ad_inpaint_only_masked_padding = gr.Slider(
+                    label="Inpaint only masked padding, pixels" + suffix(n),
+                    minimum=0,
+                    maximum=256,
+                    step=4,
+                    value=32,
+                    visible=True,
+                    elem_id=eid("ad_inpaint_full_res_padding"),
+                )
+
+                w.ad_inpaint_only_masked.change(
+                    gr_interactive,
+                    inputs=w.ad_inpaint_only_masked,
+                    outputs=w.ad_inpaint_only_masked_padding,
+                    queue=False,
+                )
+
+            with gr.Column(variant="compact"):
+                w.ad_use_inpaint_width_height = gr.Checkbox(
+                    label="Use separate width/height" + suffix(n),
+                    value=False,
+                    visible=True,
+                    elem_id=eid("ad_use_inpaint_width_height"),
+                )
+
+                w.ad_inpaint_width = gr.Slider(
+                    label="inpaint width" + suffix(n),
+                    minimum=64,
+                    maximum=2048,
+                    step=4,
+                    value=512,
+                    visible=True,
+                    elem_id=eid("ad_inpaint_width"),
+                )
+
+                w.ad_inpaint_height = gr.Slider(
+                    label="inpaint height" + suffix(n),
+                    minimum=64,
+                    maximum=2048,
+                    step=4,
+                    value=512,
+                    visible=True,
+                    elem_id=eid("ad_inpaint_height"),
+                )
+
+                w.ad_use_inpaint_width_height.change(
+                    lambda value: (gr_interactive(value), gr_interactive(value)),
+                    inputs=w.ad_use_inpaint_width_height,
+                    outputs=[w.ad_inpaint_width, w.ad_inpaint_height],
+                    queue=False,
+                )
+
+        with gr.Row():
+            with gr.Column(variant="compact"):
+                w.ad_use_steps = gr.Checkbox(
+                    label="Use separate steps" + suffix(n),
+                    value=False,
+                    visible=True,
+                    elem_id=eid("ad_use_steps"),
+                )
+
+                w.ad_steps = gr.Slider(
+                    label="ADetailer steps" + suffix(n),
+                    minimum=1,
+                    maximum=150,
+                    step=1,
+                    value=28,
+                    visible=True,
+                    elem_id=eid("ad_steps"),
+                )
+
+                w.ad_use_steps.change(
+                    gr_interactive,
+                    inputs=w.ad_use_steps,
+                    outputs=w.ad_steps,
+                    queue=False,
+                )
+
+            with gr.Column(variant="compact"):
+                w.ad_use_cfg_scale = gr.Checkbox(
+                    label="Use separate CFG scale" + suffix(n),
+                    value=False,
+                    visible=True,
+                    elem_id=eid("ad_use_cfg_scale"),
+                )
+
+                w.ad_cfg_scale = gr.Slider(
+                    label="ADetailer CFG scale" + suffix(n),
+                    minimum=0.0,
+                    maximum=30.0,
+                    step=0.5,
+                    value=7.0,
+                    visible=True,
+                    elem_id=eid("ad_cfg_scale"),
+                )
+
+                w.ad_use_cfg_scale.change(
+                    gr_interactive,
+                    inputs=w.ad_use_cfg_scale,
+                    outputs=w.ad_cfg_scale,
+                    queue=False,
+                )
+
+        with gr.Row():
+            w.ad_restore_face = gr.Checkbox(
+                label="Restore faces after ADetailer" + suffix(n),
+                value=False,
+                elem_id=eid("ad_restore_face"),
+            )
+
+
+def controlnet(w: Widgets, n: int, is_img2img: bool):
+    eid = partial(elem_id, n=n, is_img2img=is_img2img)
+    cn_models = ["None"] + get_cn_models()
+
+    with gr.Row(variant="panel"):
+        with gr.Column(variant="compact"):
+            w.ad_controlnet_model = gr.Dropdown(
+                label="ControlNet model" + suffix(n),
+                choices=cn_models,
+                value="None",
+                visible=True,
+                type="value",
+                interactive=controlnet_exists,
+                elem_id=eid("ad_controlnet_model"),
+            )
+
+            w.ad_controlnet_weight = gr.Slider(
+                label="ControlNet weight" + suffix(n),
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=1.0,
+                visible=True,
+                interactive=controlnet_exists,
+                elem_id=eid("ad_controlnet_weight"),
+            )
+
+        with gr.Column(variant="compact"):
+            w.ad_controlnet_guidance_start = gr.Slider(
+                label="ControlNet guidance start" + suffix(n),
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.0,
+                visible=True,
+                interactive=controlnet_exists,
+                elem_id=eid("ad_controlnet_guidance_start"),
+            )
+
+            w.ad_controlnet_guidance_end = gr.Slider(
+                label="ControlNet guidance end" + suffix(n),
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=1.0,
+                visible=True,
+                interactive=controlnet_exists,
+                elem_id=eid("ad_controlnet_guidance_end"),
+            )

extensions/adetailer/adetailer/ultralytics.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import cv2
+from PIL import Image
+
+from adetailer import PredictOutput
+from adetailer.common import create_mask_from_bbox
+
+
+def ultralytics_predict(
+    model_path: str | Path,
+    image: Image.Image,
+    confidence: float = 0.3,
+    device: str = "",
+) -> PredictOutput:
+    from ultralytics import YOLO
+
+    model_path = str(model_path)
+
+    model = YOLO(model_path)
+    pred = model(image, conf=confidence, device=device)
+
+    bboxes = pred[0].boxes.xyxy.cpu().numpy()
+    if bboxes.size == 0:
+        return PredictOutput()
+    bboxes = bboxes.tolist()
+
+    if pred[0].masks is None:
+        masks = create_mask_from_bbox(bboxes, image.size)
+    else:
+        masks = mask_to_pil(pred[0].masks.data, image.size)
+    preview = pred[0].plot()
+    preview = cv2.cvtColor(preview, cv2.COLOR_BGR2RGB)
+    preview = Image.fromarray(preview)
+
+    return PredictOutput(bboxes=bboxes, masks=masks, preview=preview)
+
+
+def mask_to_pil(masks, shape: tuple[int, int]) -> list[Image.Image]:
+    """
+    Parameters
+    ----------
+    masks: torch.Tensor, dtype=torch.float32, shape=(N, H, W).
+        The device can be CUDA, but `to_pil_image` takes care of that.
+
+    shape: tuple[int, int]
+        (width, height) of the original image
+    """
+    from torchvision.transforms.functional import to_pil_image
+
+    n = masks.shape[0]
+    return [to_pil_image(masks[i], mode="L").resize(shape) for i in range(n)]

extensions/adetailer/controlnet_ext/__init__.py

@@ -0,0 +1,7 @@
+from .controlnet_ext import ControlNetExt, controlnet_exists, get_cn_models
+
+__all__ = [
+    "ControlNetExt",
+    "controlnet_exists",
+    "get_cn_models",
+]

extensions/adetailer/controlnet_ext/controlnet_ext.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import importlib
+import re
+from functools import lru_cache
+from pathlib import Path
+
+from modules import extensions, sd_models, shared
+from modules.paths import data_path, models_path, script_path
+
+ext_path = Path(data_path, "extensions")
+ext_builtin_path = Path(script_path, "extensions-builtin")
+controlnet_exists = False
+controlnet_path = None
+cn_base_path = ""
+
+for extension in extensions.active():
+    if not extension.enabled:
+        continue
+    # For cases like sd-webui-controlnet-master
+    if "sd-webui-controlnet" in extension.name:
+        controlnet_exists = True
+        controlnet_path = Path(extension.path)
+        cn_base_path = ".".join(controlnet_path.parts[-2:])
+        break
+
+cn_model_module = {
+    "inpaint": "inpaint_global_harmonious",
+    "scribble": "t2ia_sketch_pidi",
+    "lineart": "lineart_coarse",
+    "openpose": "openpose_full",
+    "tile": None,
+}
+cn_model_regex = re.compile("|".join(cn_model_module.keys()))
+
+
+class ControlNetExt:
+    def __init__(self):
+        self.cn_models = ["None"]
+        self.cn_available = False
+        self.external_cn = None
+
+    def init_controlnet(self):
+        import_path = cn_base_path + ".scripts.external_code"
+
+        self.external_cn = importlib.import_module(import_path, "external_code")
+        self.cn_available = True
+        models = self.external_cn.get_models()
+        self.cn_models.extend(m for m in models if cn_model_regex.search(m))
+
+    def update_scripts_args(
+        self, p, model: str, weight: float, guidance_start: float, guidance_end: float
+    ):
+        if (not self.cn_available) or model == "None":
+            return
+
+        module = None
+        for m, v in cn_model_module.items():
+            if m in model:
+                module = v
+                break
+
+        cn_units = [
+            self.external_cn.ControlNetUnit(
+                model=model,
+                weight=weight,
+                control_mode=self.external_cn.ControlMode.BALANCED,
+                module=module,
+                guidance_start=guidance_start,
+                guidance_end=guidance_end,
+                pixel_perfect=True,
+            )
+        ]
+
+        self.external_cn.update_cn_script_in_processing(p, cn_units)
+
+
+def get_cn_model_dirs() -> list[Path]:
+    cn_model_dir = Path(models_path, "ControlNet")
+    if controlnet_path is not None:
+        cn_model_dir_old = controlnet_path.joinpath("models")
+    else:
+        cn_model_dir_old = None
+    ext_dir1 = shared.opts.data.get("control_net_models_path", "")
+    ext_dir2 = shared.opts.data.get("controlnet_dir", "")
+
+    dirs = [cn_model_dir]
+    for ext_dir in [cn_model_dir_old, ext_dir1, ext_dir2]:
+        if ext_dir:
+            dirs.append(Path(ext_dir))
+
+    return dirs
+
+
+@lru_cache
+def _get_cn_models() -> list[str]:
+    """
+    Since we can't import ControlNet, we use a function that does something like
+    controlnet's `list(global_state.cn_models_names.values())`.
+    """
+    cn_model_exts = (".pt", ".pth", ".ckpt", ".safetensors")
+    dirs = get_cn_model_dirs()
+    name_filter = shared.opts.data.get("control_net_models_name_filter", "")
+    name_filter = name_filter.strip(" ").lower()
+
+    model_paths = []
+
+    for base in dirs:
+        if not base.exists():
+            continue
+
+        for p in base.rglob("*"):
+            if (
+                p.is_file()
+                and p.suffix in cn_model_exts
+                and cn_model_regex.search(p.name)
+            ):
+                if name_filter and name_filter not in p.name.lower():
+                    continue
+                model_paths.append(p)
+    model_paths.sort(key=lambda p: p.name)
+
+    models = []
+    for p in model_paths:
+        model_hash = sd_models.model_hash(p)
+        name = f"{p.stem} [{model_hash}]"
+        models.append(name)
+    return models
+
+
+def get_cn_models() -> list[str]:
+    if controlnet_exists:
+        return _get_cn_models()
+    return []