split oversize extras.py to postprocessing.py

5 files changed, 18 insertions, 474 deletions

diff --git a/modules/extras.py b/modules/extras.py
index 385430dc..f04ddfc2 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,231 +1,16 @@
-from __future__ import annotations

-import math

 import os

 import re

-import sys

-import traceback

 import shutil

 

-import numpy as np

-from PIL import Image

 

 import torch

 import tqdm

 

-from typing import Callable, List, OrderedDict, Tuple

-from functools import partial

-from dataclasses import dataclass

-

-from modules import processing, shared, images, devices, sd_models, sd_samplers, sd_vae

-from modules.shared import opts

-import modules.gfpgan_model

+from modules import shared, images, sd_models, sd_vae

 from modules.ui import plaintext_to_html

-import modules.codeformer_model

 import gradio as gr

 import safetensors.torch

 

-class LruCache(OrderedDict):

-    @dataclass(frozen=True)

-    class Key:

-        image_hash: int

-        info_hash: int

-        args_hash: int

-

-    @dataclass

-    class Value:

-        image: Image.Image

-        info: str

-

-    def __init__(self, max_size: int = 5, *args, **kwargs):

-        super().__init__(*args, **kwargs)

-        self._max_size = max_size

-

-    def get(self, key: LruCache.Key) -> LruCache.Value:

-        ret = super().get(key)

-        if ret is not None:

-            self.move_to_end(key)  # Move to end of eviction list

-        return ret

-

-    def put(self, key: LruCache.Key, value: LruCache.Value) -> None:

-        self[key] = value

-        while len(self) > self._max_size:

-            self.popitem(last=False)

-

-

-cached_images: LruCache = LruCache(max_size=5)

-

-

-def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, upscale_first: bool, save_output: bool = True):

-    devices.torch_gc()

-

-    shared.state.begin()

-    shared.state.job = 'extras'

-

-    imageArr = []

-    # Also keep track of original file names

-    imageNameArr = []

-    outputs = []

-

-    if extras_mode == 1:

-        #convert file to pillow image

-        for img in image_folder:

-            image = Image.open(img)

-            imageArr.append(image)

-            imageNameArr.append(os.path.splitext(img.orig_name)[0])

-    elif extras_mode == 2:

-        assert not shared.cmd_opts.hide_ui_dir_config, '--hide-ui-dir-config option must be disabled'

-

-        if input_dir == '':

-            return outputs, "Please select an input directory.", ''

-        image_list = shared.listfiles(input_dir)

-        for img in image_list:

-            try:

-                image = Image.open(img)

-            except Exception:

-                continue

-            imageArr.append(image)

-            imageNameArr.append(img)

-    else:

-        imageArr.append(image)

-        imageNameArr.append(None)

-

-    if extras_mode == 2 and output_dir != '':

-        outpath = output_dir

-    else:

-        outpath = opts.outdir_samples or opts.outdir_extras_samples

-

-    # Extra operation definitions

-

-    def run_gfpgan(image: Image.Image, info: str) -> Tuple[Image.Image, str]:

-        shared.state.job = 'extras-gfpgan'

-        restored_img = modules.gfpgan_model.gfpgan_fix_faces(np.array(image, dtype=np.uint8))

-        res = Image.fromarray(restored_img)

-

-        if gfpgan_visibility < 1.0:

-            res = Image.blend(image, res, gfpgan_visibility)

-

-        info += f"GFPGAN visibility:{round(gfpgan_visibility, 2)}\n"

-        return (res, info)

-

-    def run_codeformer(image: Image.Image, info: str) -> Tuple[Image.Image, str]:

-        shared.state.job = 'extras-codeformer'

-        restored_img = modules.codeformer_model.codeformer.restore(np.array(image, dtype=np.uint8), w=codeformer_weight)

-        res = Image.fromarray(restored_img)

-

-        if codeformer_visibility < 1.0:

-            res = Image.blend(image, res, codeformer_visibility)

-

-        info += f"CodeFormer w: {round(codeformer_weight, 2)}, CodeFormer visibility:{round(codeformer_visibility, 2)}\n"

-        return (res, info)

-

-    def upscale(image, scaler_index, resize, mode, resize_w, resize_h, crop):

-        shared.state.job = 'extras-upscale'

-        upscaler = shared.sd_upscalers[scaler_index]

-        res = upscaler.scaler.upscale(image, resize, upscaler.data_path)

-        if mode == 1 and crop:

-            cropped = Image.new("RGB", (resize_w, resize_h))

-            cropped.paste(res, box=(resize_w // 2 - res.width // 2, resize_h // 2 - res.height // 2))

-            res = cropped

-        return res

-

-    def run_prepare_crop(image: Image.Image, info: str) -> Tuple[Image.Image, str]:

-        # Actual crop happens in run_upscalers_blend, this just sets upscaling_resize and adds info text

-        nonlocal upscaling_resize

-        if resize_mode == 1:

-            upscaling_resize = max(upscaling_resize_w/image.width, upscaling_resize_h/image.height)

-            crop_info = " (crop)" if upscaling_crop else ""

-            info += f"Resize to: {upscaling_resize_w:g}x{upscaling_resize_h:g}{crop_info}\n"

-        return (image, info)

-

-    @dataclass

-    class UpscaleParams:

-        upscaler_idx: int

-        blend_alpha: float

-

-    def run_upscalers_blend(params: List[UpscaleParams], image: Image.Image, info: str) -> Tuple[Image.Image, str]:

-        blended_result: Image.Image = None

-        image_hash: str = hash(np.array(image.getdata()).tobytes())

-        for upscaler in params:

-            upscale_args = (upscaler.upscaler_idx, upscaling_resize, resize_mode,

-                            upscaling_resize_w, upscaling_resize_h, upscaling_crop)

-            cache_key = LruCache.Key(image_hash=image_hash,

-                                     info_hash=hash(info),

-                                     args_hash=hash(upscale_args))

-            cached_entry = cached_images.get(cache_key)

-            if cached_entry is None:

-                res = upscale(image, *upscale_args)

-                info += f"Upscale: {round(upscaling_resize, 3)}, visibility: {upscaler.blend_alpha}, model:{shared.sd_upscalers[upscaler.upscaler_idx].name}\n"

-                cached_images.put(cache_key, LruCache.Value(image=res, info=info))

-            else:

-                res, info = cached_entry.image, cached_entry.info

-

-            if blended_result is None:

-                blended_result = res

-            else:

-                blended_result = Image.blend(blended_result, res, upscaler.blend_alpha)

-        return (blended_result, info)

-

-    # Build a list of operations to run

-    facefix_ops: List[Callable] = []

-    facefix_ops += [run_gfpgan] if gfpgan_visibility > 0 else []

-    facefix_ops += [run_codeformer] if codeformer_visibility > 0 else []

-

-    upscale_ops: List[Callable] = []

-    upscale_ops += [run_prepare_crop] if resize_mode == 1 else []

-

-    if upscaling_resize != 0:

-        step_params: List[UpscaleParams] = []

-        step_params.append(UpscaleParams(upscaler_idx=extras_upscaler_1, blend_alpha=1.0))

-        if extras_upscaler_2 != 0 and extras_upscaler_2_visibility > 0:

-            step_params.append(UpscaleParams(upscaler_idx=extras_upscaler_2, blend_alpha=extras_upscaler_2_visibility))

-

-        upscale_ops.append(partial(run_upscalers_blend, step_params))

-

-    extras_ops: List[Callable] = (upscale_ops + facefix_ops) if upscale_first else (facefix_ops + upscale_ops)

-

-    for image, image_name in zip(imageArr, imageNameArr):

-        if image is None:

-            return outputs, "Please select an input image.", ''

-

-        shared.state.textinfo = f'Processing image {image_name}'

-        

-        existing_pnginfo = image.info or {}

-

-        image = image.convert("RGB")

-        info = ""

-        # Run each operation on each image

-        for op in extras_ops:

-            image, info = op(image, info)

-

-        if opts.use_original_name_batch and image_name is not None:

-            basename = os.path.splitext(os.path.basename(image_name))[0]

-        else:

-            basename = ''

-

-        if opts.enable_pnginfo: # append info before save

-            image.info = existing_pnginfo

-            image.info["extras"] = info

-

-        if save_output:

-            # Add upscaler name as a suffix.

-            suffix = f"-{shared.sd_upscalers[extras_upscaler_1].name}" if shared.opts.use_upscaler_name_as_suffix else ""

-            # Add second upscaler if applicable.

-            if suffix and extras_upscaler_2 and extras_upscaler_2_visibility:

-                suffix += f"-{shared.sd_upscalers[extras_upscaler_2].name}"

-

-            images.save_image(image, path=outpath, basename=basename, seed=None, prompt=None, extension=opts.samples_format, info=info, short_filename=True,

-                            no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=existing_pnginfo, forced_filename=None, suffix=suffix)

-

-        if extras_mode != 2 or show_extras_results :

-            outputs.append(image)

-

-    devices.torch_gc()

-

-    return outputs, plaintext_to_html(info), ''

-

-def clear_cache():

-    cached_images.clear()

-

 

 def run_pnginfo(image):

     if image is None:

diff --git a/modules/postprocessing.py b/modules/postprocessing.py
index 385430dc..cb85720b 100644
--- a/modules/postprocessing.py
+++ b/modules/postprocessing.py
@@ -1,28 +1,18 @@
 from __future__ import annotations

-import math

 import os

-import re

-import sys

-import traceback

-import shutil

 

 import numpy as np

 from PIL import Image

 

-import torch

-import tqdm

-

 from typing import Callable, List, OrderedDict, Tuple

 from functools import partial

 from dataclasses import dataclass

 

-from modules import processing, shared, images, devices, sd_models, sd_samplers, sd_vae

+from modules import shared, images, devices, ui_components

 from modules.shared import opts

 import modules.gfpgan_model

-from modules.ui import plaintext_to_html

 import modules.codeformer_model

-import gradio as gr

-import safetensors.torch

+

 

 class LruCache(OrderedDict):

     @dataclass(frozen=True)

@@ -55,7 +45,7 @@ class LruCache(OrderedDict):
 cached_images: LruCache = LruCache(max_size=5)

 

 

-def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, upscale_first: bool, save_output: bool = True):

+def run_postprocessing(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, upscale_first: bool, save_output: bool = True):

     devices.torch_gc()

 

     shared.state.begin()

@@ -221,246 +211,9 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
 

     devices.torch_gc()

 

-    return outputs, plaintext_to_html(info), ''

+    return outputs, ui_components.plaintext_to_html(info), ''

+

 

 def clear_cache():

     cached_images.clear()

 

-

-def run_pnginfo(image):

-    if image is None:

-        return '', '', ''

-

-    geninfo, items = images.read_info_from_image(image)

-    items = {**{'parameters': geninfo}, **items}

-

-    info = ''

-    for key, text in items.items():

-        info += f"""

-<div>

-<p><b>{plaintext_to_html(str(key))}</b></p>

-<p>{plaintext_to_html(str(text))}</p>

-</div>

-""".strip()+"\n"

-

-    if len(info) == 0:

-        message = "Nothing found in the image."

-        info = f"<div><p>{message}<p></div>"

-

-    return '', geninfo, info

-

-

-def create_config(ckpt_result, config_source, a, b, c):

-    def config(x):

-        res = sd_models.find_checkpoint_config(x) if x else None

-        return res if res != shared.sd_default_config else None

-

-    if config_source == 0:

-        cfg = config(a) or config(b) or config(c)

-    elif config_source == 1:

-        cfg = config(b)

-    elif config_source == 2:

-        cfg = config(c)

-    else:

-        cfg = None

-

-    if cfg is None:

-        return

-

-    filename, _ = os.path.splitext(ckpt_result)

-    checkpoint_filename = filename + ".yaml"

-

-    print("Copying config:")

-    print("   from:", cfg)

-    print("     to:", checkpoint_filename)

-    shutil.copyfile(cfg, checkpoint_filename)

-

-

-checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]

-

-

-def to_half(tensor, enable):

-    if enable and tensor.dtype == torch.float:

-        return tensor.half()

-

-    return tensor

-

-

-def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights):

-    shared.state.begin()

-    shared.state.job = 'model-merge'

-

-    def fail(message):

-        shared.state.textinfo = message

-        shared.state.end()

-        return [*[gr.update() for _ in range(4)], message]

-

-    def weighted_sum(theta0, theta1, alpha):

-        return ((1 - alpha) * theta0) + (alpha * theta1)

-

-    def get_difference(theta1, theta2):

-        return theta1 - theta2

-

-    def add_difference(theta0, theta1_2_diff, alpha):

-        return theta0 + (alpha * theta1_2_diff)

-

-    def filename_weighted_sum():

-        a = primary_model_info.model_name

-        b = secondary_model_info.model_name

-        Ma = round(1 - multiplier, 2)

-        Mb = round(multiplier, 2)

-

-        return f"{Ma}({a}) + {Mb}({b})"

-

-    def filename_add_difference():

-        a = primary_model_info.model_name

-        b = secondary_model_info.model_name

-        c = tertiary_model_info.model_name

-        M = round(multiplier, 2)

-

-        return f"{a} + {M}({b} - {c})"

-

-    def filename_nothing():

-        return primary_model_info.model_name

-

-    theta_funcs = {

-        "Weighted sum": (filename_weighted_sum, None, weighted_sum),

-        "Add difference": (filename_add_difference, get_difference, add_difference),

-        "No interpolation": (filename_nothing, None, None),

-    }

-    filename_generator, theta_func1, theta_func2 = theta_funcs[interp_method]

-    shared.state.job_count = (1 if theta_func1 else 0) + (1 if theta_func2 else 0)

-

-    if not primary_model_name:

-        return fail("Failed: Merging requires a primary model.")

-

-    primary_model_info = sd_models.checkpoints_list[primary_model_name]

-

-    if theta_func2 and not secondary_model_name:

-        return fail("Failed: Merging requires a secondary model.")

-

-    secondary_model_info = sd_models.checkpoints_list[secondary_model_name] if theta_func2 else None

-

-    if theta_func1 and not tertiary_model_name:

-        return fail(f"Failed: Interpolation method ({interp_method}) requires a tertiary model.")

-

-    tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if theta_func1 else None

-

-    result_is_inpainting_model = False

-

-    if theta_func2:

-        shared.state.textinfo = f"Loading B"

-        print(f"Loading {secondary_model_info.filename}...")

-        theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu')

-    else:

-        theta_1 = None

-

-    if theta_func1:

-        shared.state.textinfo = f"Loading C"

-        print(f"Loading {tertiary_model_info.filename}...")

-        theta_2 = sd_models.read_state_dict(tertiary_model_info.filename, map_location='cpu')

-

-        shared.state.textinfo = 'Merging B and C'

-        shared.state.sampling_steps = len(theta_1.keys())

-        for key in tqdm.tqdm(theta_1.keys()):

-            if key in checkpoint_dict_skip_on_merge:

-                continue

-

-            if 'model' in key:

-                if key in theta_2:

-                    t2 = theta_2.get(key, torch.zeros_like(theta_1[key]))

-                    theta_1[key] = theta_func1(theta_1[key], t2)

-                else:

-                    theta_1[key] = torch.zeros_like(theta_1[key])

-

-            shared.state.sampling_step += 1

-        del theta_2

-

-        shared.state.nextjob()

-

-    shared.state.textinfo = f"Loading {primary_model_info.filename}..."

-    print(f"Loading {primary_model_info.filename}...")

-    theta_0 = sd_models.read_state_dict(primary_model_info.filename, map_location='cpu')

-

-    print("Merging...")

-    shared.state.textinfo = 'Merging A and B'

-    shared.state.sampling_steps = len(theta_0.keys())

-    for key in tqdm.tqdm(theta_0.keys()):

-        if theta_1 and 'model' in key and key in theta_1:

-

-            if key in checkpoint_dict_skip_on_merge:

-                continue

-

-            a = theta_0[key]

-            b = theta_1[key]

-

-            # this enables merging an inpainting model (A) with another one (B);

-            # where normal model would have 4 channels, for latenst space, inpainting model would

-            # have another 4 channels for unmasked picture's latent space, plus one channel for mask, for a total of 9

-            if a.shape != b.shape and a.shape[0:1] + a.shape[2:] == b.shape[0:1] + b.shape[2:]:

-                if a.shape[1] == 4 and b.shape[1] == 9:

-                    raise RuntimeError("When merging inpainting model with a normal one, A must be the inpainting model.")

-

-                assert a.shape[1] == 9 and b.shape[1] == 4, f"Bad dimensions for merged layer {key}: A={a.shape}, B={b.shape}"

-

-                theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)

-                result_is_inpainting_model = True

-            else:

-                theta_0[key] = theta_func2(a, b, multiplier)

-

-            theta_0[key] = to_half(theta_0[key], save_as_half)

-

-        shared.state.sampling_step += 1

-

-    del theta_1

-

-    bake_in_vae_filename = sd_vae.vae_dict.get(bake_in_vae, None)

-    if bake_in_vae_filename is not None:

-        print(f"Baking in VAE from {bake_in_vae_filename}")

-        shared.state.textinfo = 'Baking in VAE'

-        vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename, map_location='cpu')

-

-        for key in vae_dict.keys():

-            theta_0_key = 'first_stage_model.' + key

-            if theta_0_key in theta_0:

-                theta_0[theta_0_key] = to_half(vae_dict[key], save_as_half)

-

-        del vae_dict

-

-    if save_as_half and not theta_func2:

-        for key in theta_0.keys():

-            theta_0[key] = to_half(theta_0[key], save_as_half)

-

-    if discard_weights:

-        regex = re.compile(discard_weights)

-        for key in list(theta_0):

-            if re.search(regex, key):

-                theta_0.pop(key, None)

-

-    ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path

-

-    filename = filename_generator() if custom_name == '' else custom_name

-    filename += ".inpainting" if result_is_inpainting_model else ""

-    filename += "." + checkpoint_format

-

-    output_modelname = os.path.join(ckpt_dir, filename)

-

-    shared.state.nextjob()

-    shared.state.textinfo = "Saving"

-    print(f"Saving to {output_modelname}...")

-

-    _, extension = os.path.splitext(output_modelname)

-    if extension.lower() == ".safetensors":

-        safetensors.torch.save_file(theta_0, output_modelname, metadata={"format": "pt"})

-    else:

-        torch.save(theta_0, output_modelname)

-

-    sd_models.list_models()

-

-    create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info)

-

-    print(f"Checkpoint saved to {output_modelname}.")

-    shared.state.textinfo = "Checkpoint saved"

-    shared.state.end()

-

-    return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "Checkpoint saved to " + output_modelname]

diff --git a/modules/ui.py b/modules/ui.py
index eb4b7e6b..4116e167 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -20,7 +20,7 @@ import numpy as np
 from PIL import Image, PngImagePlugin

 from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call

 

-from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru, sd_vae, extra_networks

+from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru, sd_vae, extra_networks, postprocessing, ui_components

 from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML

 from modules.paths import script_path

 

@@ -95,8 +95,8 @@ extra_networks_symbol = '\U0001F3B4'  # 🎴
 

 

 def plaintext_to_html(text):

-    text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"

-    return text

+    return ui_components.plaintext_to_html(text)

+

 

 def send_gradio_gallery_to_image(x):

     if len(x) == 0:

@@ -1152,7 +1152,7 @@ def create_ui():
             result_images, html_info_x, html_info, html_log = create_output_panel("extras", opts.outdir_extras_samples)

 

         submit.click(

-            fn=wrap_gradio_gpu_call(modules.extras.run_extras, extra_outputs=[None, '']),

+            fn=wrap_gradio_gpu_call(postprocessing.run_postprocessing, extra_outputs=[None, '']),

             _js="get_extras_tab_index",

             inputs=[

                 dummy_component,

@@ -1183,7 +1183,7 @@ def create_ui():
         parameters_copypaste.add_paste_fields("extras", extras_image, None)

 

         extras_image.change(

-            fn=modules.extras.clear_cache,

+            fn=postprocessing.clear_cache,

             inputs=[], outputs=[]

         )

 

diff --git a/modules/ui_components.py b/modules/ui_components.py
index 46324425..989cc87b 100644
--- a/modules/ui_components.py
+++ b/modules/ui_components.py
@@ -1,3 +1,5 @@
+import html

+

 import gradio as gr

 

 

@@ -47,3 +49,8 @@ class FormColorPicker(gr.ColorPicker, gr.components.FormComponent):
 

     def get_block_name(self):

         return "colorpicker"

+

+

+def plaintext_to_html(text):

+    text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"

+    return text

diff --git a/webui.py b/webui.py
index d235da74..7cf5885e 100644
--- a/webui.py
+++ b/webui.py
@@ -22,7 +22,6 @@ if ".dev" in torch.__version__ or "+git" in torch.__version__:
 

 from modules import shared, devices, sd_samplers, upscaler, extensions, localization, ui_tempdir, ui_extra_networks

 import modules.codeformer_model as codeformer

-import modules.extras

 import modules.face_restoration

 import modules.gfpgan_model as gfpgan

 import modules.img2img