From 84dd7e8e2495c4fc2997e97f8267aa831eb90d11 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 2 Jan 2023 20:30:02 +0300
Subject: error out with a readable message in chwewckpoint merger for
 incompatible tensor shapes (ie when trying to merge SD1.5 with SD2.0)

---
 modules/extras.py | 2 ++
 1 file changed, 2 insertions(+)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 68939dea..5e270250 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -303,6 +303,8 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
                 theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)
                 result_is_inpainting_model = True
             else:
+                assert a.shape == b.shape, f'Incompatible shapes for layer {key}: A is {a.shape}, and B is {b.shape}'
+
                 theta_0[key] = theta_func2(a, b, multiplier)
 
             if save_as_half:
-- 
cgit v1.2.3


From 1d9dc48efda2e8da6d13fc62e65500198a9b041c Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Tue, 3 Jan 2023 10:21:51 -0500
Subject: init job and add info to model merge

---
 modules/extras.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 5e270250..7e222313 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -242,6 +242,9 @@ def run_pnginfo(image):
 
 
 def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format):
+    shared.state.begin()
+    shared.state.job = 'model-merge'
+
     def weighted_sum(theta0, theta1, alpha):
         return ((1 - alpha) * theta0) + (alpha * theta1)
 
@@ -263,8 +266,11 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     theta_func1, theta_func2 = theta_funcs[interp_method]
 
     if theta_func1 and not tertiary_model_info:
+        shared.state.textinfo = "Failed: Interpolation method requires a tertiary model."
+        shared.state.end()
         return ["Failed: Interpolation method requires a tertiary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
 
+    shared.state.textinfo = f"Loading {secondary_model_info.filename}..."
     print(f"Loading {secondary_model_info.filename}...")
     theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu')
 
@@ -281,6 +287,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
                     theta_1[key] = torch.zeros_like(theta_1[key])
         del theta_2
 
+    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')
 
@@ -291,6 +298,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
             a = theta_0[key]
             b = theta_1[key]
 
+            shared.state.textinfo = f'Merging layer {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
@@ -303,8 +311,6 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
                 theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)
                 result_is_inpainting_model = True
             else:
-                assert a.shape == b.shape, f'Incompatible shapes for layer {key}: A is {a.shape}, and B is {b.shape}'
-
                 theta_0[key] = theta_func2(a, b, multiplier)
 
             if save_as_half:
@@ -332,6 +338,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
 
     output_modelname = os.path.join(ckpt_dir, filename)
 
+    shared.state.textinfo = f"Saving to {output_modelname}..."
     print(f"Saving to {output_modelname}...")
 
     _, extension = os.path.splitext(output_modelname)
@@ -343,4 +350,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     sd_models.list_models()
 
     print("Checkpoint saved.")
+    shared.state.textinfo = "Checkpoint saved to " + output_modelname
+    shared.state.end()
+
     return ["Checkpoint saved to " + output_modelname] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
-- 
cgit v1.2.3


From 192ddc04d6de0d780f73aa5fbaa8c66cd4642e1c Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Tue, 3 Jan 2023 10:34:51 -0500
Subject: add job info to modules

---
 modules/extras.py                              | 17 +++++++++++++----
 modules/hypernetworks/hypernetwork.py          |  1 +
 modules/textual_inversion/preprocess.py        |  1 +
 modules/textual_inversion/textual_inversion.py |  1 +
 4 files changed, 16 insertions(+), 4 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 7e222313..d665440a 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -58,6 +58,9 @@ 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 = []
@@ -94,6 +97,7 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
     # 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)
 
@@ -104,6 +108,7 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
         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)
 
@@ -114,6 +119,7 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
         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:
@@ -180,6 +186,9 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
     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")
@@ -193,6 +202,10 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
         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 ""
@@ -203,10 +216,6 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
             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 opts.enable_pnginfo:
-            image.info = existing_pnginfo
-            image.info["extras"] = info
-
         if extras_mode != 2 or show_extras_results :
             outputs.append(image)
 
diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py
index 109e8078..450fecac 100644
--- a/modules/hypernetworks/hypernetwork.py
+++ b/modules/hypernetworks/hypernetwork.py
@@ -417,6 +417,7 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, gradient_step,
     shared.loaded_hypernetwork = Hypernetwork()
     shared.loaded_hypernetwork.load(path)
 
+    shared.state.job = "train-hypernetwork"
     shared.state.textinfo = "Initializing hypernetwork training..."
     shared.state.job_count = steps
 
diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index 56b9b2eb..feb876c6 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -124,6 +124,7 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
 
     files = listfiles(src)
 
+    shared.state.job = "preprocess"
     shared.state.textinfo = "Preprocessing..."
     shared.state.job_count = len(files)
 
diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index fd253477..2c1251d6 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -245,6 +245,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
     create_image_every = create_image_every or 0
     validate_train_inputs(embedding_name, learn_rate, batch_size, gradient_step, data_root, template_file, steps, save_embedding_every, create_image_every, log_directory, name="embedding")
 
+    shared.state.job = "train-embedding"
     shared.state.textinfo = "Initializing textual inversion training..."
     shared.state.job_count = steps
 
-- 
cgit v1.2.3


From 79c682ad4f2d982b26fa1a15044582d1005134f9 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Wed, 4 Jan 2023 08:20:42 -0500
Subject: fix jpeg

---
 modules/extras.py         | 2 --
 modules/images.py         | 2 ++
 requirements_versions.txt | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index d665440a..7407bfe3 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -19,8 +19,6 @@ from modules.shared import opts
 import modules.gfpgan_model
 from modules.ui import plaintext_to_html
 import modules.codeformer_model
-import piexif
-import piexif.helper
 import gradio as gr
 import safetensors.torch
 
diff --git a/modules/images.py b/modules/images.py
index c3a5fc8b..a73be3fa 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -22,6 +22,8 @@ from modules.shared import opts, cmd_opts
 
 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
 
+Image.init() # initialize once all known file format handlers
+
 
 def image_grid(imgs, batch_size=1, rows=None):
     if rows is None:
diff --git a/requirements_versions.txt b/requirements_versions.txt
index 975102d9..7ae118cb 100644
--- a/requirements_versions.txt
+++ b/requirements_versions.txt
@@ -5,7 +5,7 @@ basicsr==1.4.2
 gfpgan==1.3.8
 gradio==3.15.0
 numpy==1.23.3
-Pillow==9.2.0
+Pillow==9.3.0
 realesrgan==0.3.0
 torch
 omegaconf==2.2.3
-- 
cgit v1.2.3


From 954091697fce7a1b7997d5f3d73551f793f6bebc Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Wed, 11 Jan 2023 09:10:07 +0300
Subject: add an option to copy config from one of models in checkpoint merger

---
 modules/extras.py | 30 +++++++++++++++++++++++++++++-
 modules/ui.py     |  9 ++++++---
 2 files changed, 35 insertions(+), 4 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 7407bfe3..a03d558e 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -3,6 +3,7 @@ import math
 import os
 import sys
 import traceback
+import shutil
 
 import numpy as np
 from PIL import Image
@@ -248,7 +249,32 @@ def run_pnginfo(image):
     return '', geninfo, info
 
 
-def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format):
+def create_config(ckpt_result, config_source, a, b, c):
+    def config(x):
+        return sd_models.find_checkpoint_config(x) if x 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)
+
+
+def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source):
     shared.state.begin()
     shared.state.job = 'model-merge'
 
@@ -356,6 +382,8 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
 
     sd_models.list_models()
 
+    create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info)
+
     print("Checkpoint saved.")
     shared.state.textinfo = "Checkpoint saved to " + output_modelname
     shared.state.end()
diff --git a/modules/ui.py b/modules/ui.py
index 3c458ce8..82f5dd7c 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1129,7 +1129,7 @@ def create_ui():
             with gr.Column(variant='panel'):
                 gr.HTML(value="<p>A merger of the two checkpoints will be generated in your <b>checkpoint</b> directory.</p>")
 
-                with gr.Row():
+                with FormRow():
                     primary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_primary_model_name", label="Primary model (A)")
                     create_refresh_button(primary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_A")
 
@@ -1143,11 +1143,13 @@ def create_ui():
                 interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Multiplier (M) - set to 0 to get model A', value=0.3, elem_id="modelmerger_interp_amount")
                 interp_method = gr.Radio(choices=["Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method", elem_id="modelmerger_interp_method")
 
-                with gr.Row():
+                with FormRow():
                     checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="ckpt", label="Checkpoint format", elem_id="modelmerger_checkpoint_format")
                     save_as_half = gr.Checkbox(value=False, label="Save as float16", elem_id="modelmerger_save_as_half")
 
-                modelmerger_merge = gr.Button(elem_id="modelmerger_merge", label="Merge", variant='primary')
+                config_source = gr.Radio(choices=["A, B or C", "B", "C", "Don't"], value="A, B or C", label="Copy config from", type="index", elem_id="modelmerger_config_method")
+
+                modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')
 
             with gr.Column(variant='panel'):
                 submit_result = gr.Textbox(elem_id="modelmerger_result", show_label=False)
@@ -1703,6 +1705,7 @@ def create_ui():
                 save_as_half,
                 custom_name,
                 checkpoint_format,
+                config_source,
             ],
             outputs=[
                 submit_result,
-- 
cgit v1.2.3


From 5f8685237ed6427c9a8e502124074c740ea7696a Mon Sep 17 00:00:00 2001
From: bbc_mc <bbc_mc@hotmail.co.jp>
Date: Sat, 14 Jan 2023 20:00:00 +0900
Subject: Exclude clip index from merge

---
 modules/extras.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index a03d558e..22668fcd 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -326,8 +326,14 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
 
     print("Merging...")
 
+    chckpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
+
     for key in tqdm.tqdm(theta_0.keys()):
         if 'model' in key and key in theta_1:
+
+            if key in chckpoint_dict_skip_on_merge:
+                continue
+
             a = theta_0[key]
             b = theta_1[key]
 
@@ -352,6 +358,10 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     # I believe this part should be discarded, but I'll leave it for now until I am sure
     for key in theta_1.keys():
         if 'model' in key and key not in theta_0:
+
+            if key in chckpoint_dict_skip_on_merge:
+                continue
+
             theta_0[key] = theta_1[key]
             if save_as_half:
                 theta_0[key] = theta_0[key].half()
-- 
cgit v1.2.3


From 99207bc816d027b522e1c49001748c63fd426b53 Mon Sep 17 00:00:00 2001
From: EllangoK <ellangok@bxscience.edu>
Date: Wed, 18 Jan 2023 19:13:15 -0500
Subject: check model name values are set before merging

---
 modules/extras.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 22668fcd..29eb1f07 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -287,10 +287,19 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     def add_difference(theta0, theta1_2_diff, alpha):
         return theta0 + (alpha * theta1_2_diff)
 
+    if not primary_model_name:
+        shared.state.textinfo = "Failed: Merging requires a primary model."
+        shared.state.end()
+        return ["Failed: Merging requires a primary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
+
     primary_model_info = sd_models.checkpoints_list[primary_model_name]
+
+    if not secondary_model_name:
+        shared.state.textinfo = "Failed: Merging requires a secondary model."
+        shared.state.end()
+        return ["Failed: Merging requires a secondary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
+    
     secondary_model_info = sd_models.checkpoints_list[secondary_model_name]
-    tertiary_model_info = sd_models.checkpoints_list.get(tertiary_model_name, None)
-    result_is_inpainting_model = False
 
     theta_funcs = {
         "Weighted sum": (None, weighted_sum),
@@ -298,10 +307,15 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     }
     theta_func1, theta_func2 = theta_funcs[interp_method]
 
-    if theta_func1 and not tertiary_model_info:
+    tertiary_model_info = None
+    if theta_func1 and not tertiary_model_name:
         shared.state.textinfo = "Failed: Interpolation method requires a tertiary model."
         shared.state.end()
-        return ["Failed: Interpolation method requires a tertiary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
+        return [f"Failed: Interpolation method ({interp_method}) requires a tertiary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
+    else:
+        tertiary_model_info = sd_models.checkpoints_list.get(tertiary_model_name, None)
+
+    result_is_inpainting_model = False
 
     shared.state.textinfo = f"Loading {secondary_model_info.filename}..."
     print(f"Loading {secondary_model_info.filename}...")
-- 
cgit v1.2.3


From 26a6a78b16f88a6f88f4cca3f378db3b83fc94f8 Mon Sep 17 00:00:00 2001
From: EllangoK <karun.ellango7@gmail.com>
Date: Wed, 18 Jan 2023 21:21:52 -0500
Subject: only lookup tertiary model if theta_func1 is set

---
 modules/extras.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 29eb1f07..88eea22e 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -307,13 +307,12 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     }
     theta_func1, theta_func2 = theta_funcs[interp_method]
 
-    tertiary_model_info = None
     if theta_func1 and not tertiary_model_name:
         shared.state.textinfo = "Failed: Interpolation method requires a tertiary model."
         shared.state.end()
         return [f"Failed: Interpolation method ({interp_method}) requires a tertiary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
-    else:
-        tertiary_model_info = sd_models.checkpoints_list.get(tertiary_model_name, None)
+    
+    tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if theta_func1 else None
 
     result_is_inpainting_model = False
 
-- 
cgit v1.2.3


From 7cfc6450305125683799208fb7bc27c0b12586b3 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Thu, 19 Jan 2023 08:53:50 +0300
Subject: eliminate repetition of code in #6910

---
 modules/extras.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 88eea22e..367c15cc 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -278,6 +278,11 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     shared.state.begin()
     shared.state.job = 'model-merge'
 
+    def fail(message):
+        shared.state.textinfo = message
+        shared.state.end()
+        return [message, *[gr.update() for _ in range(4)]]
+
     def weighted_sum(theta0, theta1, alpha):
         return ((1 - alpha) * theta0) + (alpha * theta1)
 
@@ -288,16 +293,12 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
         return theta0 + (alpha * theta1_2_diff)
 
     if not primary_model_name:
-        shared.state.textinfo = "Failed: Merging requires a primary model."
-        shared.state.end()
-        return ["Failed: Merging requires a primary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
+        return fail("Failed: Merging requires a primary model.")
 
     primary_model_info = sd_models.checkpoints_list[primary_model_name]
 
     if not secondary_model_name:
-        shared.state.textinfo = "Failed: Merging requires a secondary model."
-        shared.state.end()
-        return ["Failed: Merging requires a secondary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
+        return fail("Failed: Merging requires a secondary model.")
     
     secondary_model_info = sd_models.checkpoints_list[secondary_model_name]
 
@@ -308,9 +309,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     theta_func1, theta_func2 = theta_funcs[interp_method]
 
     if theta_func1 and not tertiary_model_name:
-        shared.state.textinfo = "Failed: Interpolation method requires a tertiary model."
-        shared.state.end()
-        return [f"Failed: Interpolation method ({interp_method}) requires a tertiary model."] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
+        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
 
-- 
cgit v1.2.3


From c7e50425f63c07242068f8dcccce70a4ef28a17f Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Thu, 19 Jan 2023 09:25:37 +0300
Subject: add progress bar to modelmerger

---
 javascript/ui.js    | 11 +++++++++++
 modules/extras.py   | 18 +++++++++++++++---
 modules/progress.py |  2 +-
 modules/ui.py       | 13 ++++++++-----
 style.css           |  5 +++++
 5 files changed, 40 insertions(+), 9 deletions(-)

(limited to 'modules/extras.py')

diff --git a/javascript/ui.js b/javascript/ui.js
index 7d3d57a3..428375d4 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -172,6 +172,17 @@ function submit_img2img(){
     return res
 }
 
+function modelmerger(){
+    var id = randomId()
+    requestProgress(id, gradioApp().getElementById('modelmerger_results_panel'), null, function(){})
+
+    gradioApp().getElementById('modelmerger_result').innerHTML = ''
+
+    var res = create_submit_args(arguments)
+    res[0] = id
+    return res
+}
+
 
 function ask_for_style_name(_, prompt_text, negative_prompt_text) {
     name_ = prompt('Style name:')
diff --git a/modules/extras.py b/modules/extras.py
index 367c15cc..034f28e4 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -274,14 +274,15 @@ def create_config(ckpt_result, config_source, a, b, c):
     shutil.copyfile(cfg, checkpoint_filename)
 
 
-def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source):
+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):
     shared.state.begin()
     shared.state.job = 'model-merge'
+    shared.state.job_count = 1
 
     def fail(message):
         shared.state.textinfo = message
         shared.state.end()
-        return [message, *[gr.update() for _ in range(4)]]
+        return [*[gr.update() for _ in range(4)], message]
 
     def weighted_sum(theta0, theta1, alpha):
         return ((1 - alpha) * theta0) + (alpha * theta1)
@@ -320,9 +321,12 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu')
 
     if theta_func1:
+        shared.state.job_count += 1
+
         print(f"Loading {tertiary_model_info.filename}...")
         theta_2 = sd_models.read_state_dict(tertiary_model_info.filename, map_location='cpu')
 
+        shared.state.sampling_steps = len(theta_1.keys())
         for key in tqdm.tqdm(theta_1.keys()):
             if 'model' in key:
                 if key in theta_2:
@@ -330,8 +334,12 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
                     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')
@@ -340,6 +348,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
 
     chckpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
 
+    shared.state.sampling_steps = len(theta_0.keys())
     for key in tqdm.tqdm(theta_0.keys()):
         if 'model' in key and key in theta_1:
 
@@ -367,6 +376,8 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
             if save_as_half:
                 theta_0[key] = theta_0[key].half()
 
+        shared.state.sampling_step += 1
+
     # I believe this part should be discarded, but I'll leave it for now until I am sure
     for key in theta_1.keys():
         if 'model' in key and key not in theta_0:
@@ -393,6 +404,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
 
     output_modelname = os.path.join(ckpt_dir, filename)
 
+    shared.state.nextjob()
     shared.state.textinfo = f"Saving to {output_modelname}..."
     print(f"Saving to {output_modelname}...")
 
@@ -410,4 +422,4 @@ def run_modelmerger(primary_model_name, secondary_model_name, tertiary_model_nam
     shared.state.textinfo = "Checkpoint saved to " + output_modelname
     shared.state.end()
 
-    return ["Checkpoint saved to " + output_modelname] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]
+    return [*[gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)], "Checkpoint saved to " + output_modelname]
diff --git a/modules/progress.py b/modules/progress.py
index f9e005d3..c69ecf3d 100644
--- a/modules/progress.py
+++ b/modules/progress.py
@@ -72,7 +72,7 @@ def progressapi(req: ProgressRequest):
 
     if job_count > 0:
         progress += job_no / job_count
-    if sampling_steps > 0:
+    if sampling_steps > 0 and job_count > 0:
         progress += 1 / job_count * sampling_step / sampling_steps
 
     progress = min(progress, 1)
diff --git a/modules/ui.py b/modules/ui.py
index 09a3c92e..aeee7853 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1208,8 +1208,9 @@ def create_ui():
                 with gr.Row():
                     modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')
 
-            with gr.Column(variant='panel'):
-                submit_result = gr.Textbox(elem_id="modelmerger_result", show_label=False)
+            with gr.Column(variant='compact', elem_id="modelmerger_results_container"):
+                with gr.Group(elem_id="modelmerger_results_panel"):
+                    modelmerger_result = gr.HTML(elem_id="modelmerger_result", show_label=False)
 
     with gr.Blocks(analytics_enabled=False) as train_interface:
         with gr.Row().style(equal_height=False):
@@ -1753,12 +1754,14 @@ def create_ui():
                 print("Error loading/saving model file:", file=sys.stderr)
                 print(traceback.format_exc(), file=sys.stderr)
                 modules.sd_models.list_models()  # to remove the potentially missing models from the list
-                return [f"Error merging checkpoints: {e}"] + [gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)]
+                return [*[gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"]
             return results
 
         modelmerger_merge.click(
-            fn=modelmerger,
+            fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),
+            _js='modelmerger',
             inputs=[
+                dummy_component,
                 primary_model_name,
                 secondary_model_name,
                 tertiary_model_name,
@@ -1770,11 +1773,11 @@ def create_ui():
                 config_source,
             ],
             outputs=[
-                submit_result,
                 primary_model_name,
                 secondary_model_name,
                 tertiary_model_name,
                 component_dict['sd_model_checkpoint'],
+                modelmerger_result,
             ]
         )
 
diff --git a/style.css b/style.css
index a6abd93d..32ba4753 100644
--- a/style.css
+++ b/style.css
@@ -737,6 +737,11 @@ footer {
     line-height: 2.4em;
 }
 
+#modelmerger_results_container{
+    margin-top: 1em;
+    overflow: visible;
+}
+
 /* The following handles localization for right-to-left (RTL) languages like Arabic.
 The rtl media type will only be activated by the logic in javascript/localization.js.
 If you change anything above, you need to make sure it is RTL compliant by just running
-- 
cgit v1.2.3


From 0f5dbfffd0b7202a48e404d8e74b5cc9a3e5b135 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Thu, 19 Jan 2023 10:39:51 +0300
Subject: allow baking in VAE in checkpoint merger tab do not save config if
 it's the default for checkpoint merger tab change file naming scheme for
 checkpoint merger tab allow just saving A without any merging for checkpoint
 merger tab some stylistic changes for UI in checkpoint merger tab

---
 javascript/hints.js |   1 +
 javascript/ui.js    |   2 -
 modules/extras.py   | 112 +++++++++++++++++++++++++++++++---------------------
 modules/sd_vae.py   |   9 ++++-
 modules/shared.py   |   3 +-
 modules/ui.py       |  17 ++++++--
 style.css           |  15 ++++---
 7 files changed, 101 insertions(+), 58 deletions(-)

(limited to 'modules/extras.py')

diff --git a/javascript/hints.js b/javascript/hints.js
index fa5e5ae8..e746e20d 100644
--- a/javascript/hints.js
+++ b/javascript/hints.js
@@ -92,6 +92,7 @@ titles = {
 
     "Weighted sum": "Result = A * (1 - M) + B * M",
     "Add difference": "Result = A + (B - C) * M",
+    "No interpolation": "Result = A",
 
 	"Initialization text": "If the number of tokens is more than the number of vectors, some may be skipped.\nLeave the textbox empty to start with zeroed out vectors",
     "Learning rate": "How fast should training go. Low values will take longer to train, high values may fail to converge (not generate accurate results) and/or may break the embedding (This has happened if you see Loss: nan in the training info textbox. If this happens, you need to manually restore your embedding from an older not-broken backup).\n\nYou can set a single numeric value, or multiple learning rates using the syntax:\n\n   rate_1:max_steps_1, rate_2:max_steps_2, ...\n\nEG:   0.005:100, 1e-3:1000, 1e-5\n\nWill train with rate of 0.005 for first 100 steps, then 1e-3 until 1000 steps, then 1e-5 for all remaining steps.",
diff --git a/javascript/ui.js b/javascript/ui.js
index 428375d4..37788a3e 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -176,8 +176,6 @@ function modelmerger(){
     var id = randomId()
     requestProgress(id, gradioApp().getElementById('modelmerger_results_panel'), null, function(){})
 
-    gradioApp().getElementById('modelmerger_result').innerHTML = ''
-
     var res = create_submit_args(arguments)
     res[0] = id
     return res
diff --git a/modules/extras.py b/modules/extras.py
index 034f28e4..fe701a0e 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -15,7 +15,7 @@ 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
+from modules import processing, shared, images, devices, sd_models, sd_samplers, sd_vae
 from modules.shared import opts
 import modules.gfpgan_model
 from modules.ui import plaintext_to_html
@@ -251,7 +251,8 @@ def run_pnginfo(image):
 
 def create_config(ckpt_result, config_source, a, b, c):
     def config(x):
-        return sd_models.find_checkpoint_config(x) if x else None
+        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)
@@ -274,10 +275,12 @@ def create_config(ckpt_result, config_source, a, b, c):
     shutil.copyfile(cfg, checkpoint_filename)
 
 
-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):
+chckpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
+
+
+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):
     shared.state.begin()
     shared.state.job = 'model-merge'
-    shared.state.job_count = 1
 
     def fail(message):
         shared.state.textinfo = message
@@ -293,41 +296,68 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     def add_difference(theta0, theta1_2_diff, alpha):
         return theta0 + (alpha * theta1_2_diff)
 
+    def filename_weighed_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_differnece():
+        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_weighed_sum, None, weighted_sum),
+        "Add difference": (filename_add_differnece, 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 not secondary_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]
 
-    theta_funcs = {
-        "Weighted sum": (None, weighted_sum),
-        "Add difference": (get_difference, add_difference),
-    }
-    theta_func1, theta_func2 = theta_funcs[interp_method]
+    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
 
-    shared.state.textinfo = f"Loading {secondary_model_info.filename}..."
-    print(f"Loading {secondary_model_info.filename}...")
-    theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu')
+    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.job_count += 1
-
+        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 chckpoint_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]))
@@ -345,12 +375,10 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     theta_0 = sd_models.read_state_dict(primary_model_info.filename, map_location='cpu')
 
     print("Merging...")
-
-    chckpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
-
+    shared.state.textinfo = 'Merging A and B'
     shared.state.sampling_steps = len(theta_0.keys())
     for key in tqdm.tqdm(theta_0.keys()):
-        if 'model' in key and key in theta_1:
+        if theta_1 and 'model' in key and key in theta_1:
 
             if key in chckpoint_dict_skip_on_merge:
                 continue
@@ -358,7 +386,6 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
             a = theta_0[key]
             b = theta_1[key]
 
-            shared.state.textinfo = f'Merging layer {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
@@ -378,34 +405,31 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
 
         shared.state.sampling_step += 1
 
-    # I believe this part should be discarded, but I'll leave it for now until I am sure
-    for key in theta_1.keys():
-        if 'model' in key and key not in theta_0:
+    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')
 
-            if key in chckpoint_dict_skip_on_merge:
-                continue
+        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] = vae_dict[key].half() if save_as_half else vae_dict[key]
 
-            theta_0[key] = theta_1[key]
-            if save_as_half:
-                theta_0[key] = theta_0[key].half()
-    del theta_1
+        del vae_dict
 
     ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path
 
-    filename = \
-        primary_model_info.model_name + '_' + str(round(1-multiplier, 2)) + '-' + \
-        secondary_model_info.model_name + '_' + str(round(multiplier, 2)) + '-' + \
-        interp_method.replace(" ", "_") + \
-        '-merged.' +  \
-        ("inpainting." if result_is_inpainting_model else "") + \
-        checkpoint_format
-
-    filename = filename if custom_name == '' else (custom_name + '.' + checkpoint_format)
+    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 = f"Saving to {output_modelname}..."
+    shared.state.textinfo = "Saving"
     print(f"Saving to {output_modelname}...")
 
     _, extension = os.path.splitext(output_modelname)
@@ -418,8 +442,8 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
 
     create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info)
 
-    print("Checkpoint saved.")
-    shared.state.textinfo = "Checkpoint saved to " + output_modelname
+    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/sd_vae.py b/modules/sd_vae.py
index da1bf15c..4ce238b8 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -120,6 +120,12 @@ def resolve_vae(checkpoint_file):
     return None, None
 
 
+def load_vae_dict(filename, map_location):
+    vae_ckpt = sd_models.read_state_dict(filename, map_location=map_location)
+    vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys}
+    return vae_dict_1
+
+
 def load_vae(model, vae_file=None, vae_source="from unknown source"):
     global vae_dict, loaded_vae_file
     # save_settings = False
@@ -137,8 +143,7 @@ def load_vae(model, vae_file=None, vae_source="from unknown source"):
             print(f"Loading VAE weights {vae_source}: {vae_file}")
             store_base_vae(model)
 
-            vae_ckpt = sd_models.read_state_dict(vae_file, map_location=shared.weight_load_location)
-            vae_dict_1 = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss" and k not in vae_ignore_keys}
+            vae_dict_1 = load_vae_dict(vae_file, map_location=shared.weight_load_location)
             _load_vae_dict(model, vae_dict_1)
 
             if cache_enabled:
diff --git a/modules/shared.py b/modules/shared.py
index 77e5e91c..29b28bff 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -20,10 +20,11 @@ from modules.paths import models_path, script_path, sd_path
 
 demo = None
 
+sd_default_config = os.path.join(script_path, "configs/v1-inference.yaml")
 sd_model_file = os.path.join(script_path, 'model.ckpt')
 default_sd_model_file = sd_model_file
 parser = argparse.ArgumentParser()
-parser.add_argument("--config", type=str, default=os.path.join(script_path, "configs/v1-inference.yaml"), help="path to config which constructs model",)
+parser.add_argument("--config", type=str, default=sd_default_config, help="path to config which constructs model",)
 parser.add_argument("--ckpt", type=str, default=sd_model_file, help="path to checkpoint of stable diffusion model; if specified, this checkpoint will be added to the list of checkpoints and loaded",)
 parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints")
 parser.add_argument("--vae-dir", type=str, default=None, help="Path to directory with VAE files")
diff --git a/modules/ui.py b/modules/ui.py
index aeee7853..4e381a49 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
+from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions, deepbooru, sd_vae
 from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML
 from modules.paths import script_path
 
@@ -1185,7 +1185,7 @@ def create_ui():
             with gr.Column(variant='compact'):
                 gr.HTML(value="<p style='margin-bottom: 2.5em'>A merger of the two checkpoints will be generated in your <b>checkpoint</b> directory.</p>")
 
-                with FormRow():
+                with FormRow(elem_id="modelmerger_models"):
                     primary_model_name = gr.Dropdown(modules.sd_models.checkpoint_tiles(), elem_id="modelmerger_primary_model_name", label="Primary model (A)")
                     create_refresh_button(primary_model_name, modules.sd_models.list_models, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, "refresh_checkpoint_A")
 
@@ -1197,13 +1197,20 @@ def create_ui():
 
                 custom_name = gr.Textbox(label="Custom Name (Optional)", elem_id="modelmerger_custom_name")
                 interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Multiplier (M) - set to 0 to get model A', value=0.3, elem_id="modelmerger_interp_amount")
-                interp_method = gr.Radio(choices=["Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method", elem_id="modelmerger_interp_method")
+                interp_method = gr.Radio(choices=["No interpolation", "Weighted sum", "Add difference"], value="Weighted sum", label="Interpolation Method", elem_id="modelmerger_interp_method")
 
                 with FormRow():
                     checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="ckpt", label="Checkpoint format", elem_id="modelmerger_checkpoint_format")
                     save_as_half = gr.Checkbox(value=False, label="Save as float16", elem_id="modelmerger_save_as_half")
 
-                config_source = gr.Radio(choices=["A, B or C", "B", "C", "Don't"], value="A, B or C", label="Copy config from", type="index", elem_id="modelmerger_config_method")
+                with FormRow():
+                    with gr.Column():
+                        config_source = gr.Radio(choices=["A, B or C", "B", "C", "Don't"], value="A, B or C", label="Copy config from", type="index", elem_id="modelmerger_config_method")
+
+                    with gr.Column():
+                        with FormRow():
+                            bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None", label="Bake in VAE", elem_id="modelmerger_bake_in_vae")
+                            create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["None"] + list(sd_vae.vae_dict)}, "modelmerger_refresh_bake_in_vae")
 
                 with gr.Row():
                     modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')
@@ -1757,6 +1764,7 @@ def create_ui():
                 return [*[gr.Dropdown.update(choices=modules.sd_models.checkpoint_tiles()) for _ in range(4)], f"Error merging checkpoints: {e}"]
             return results
 
+        modelmerger_merge.click(fn=lambda: '', inputs=[], outputs=[modelmerger_result])
         modelmerger_merge.click(
             fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)]),
             _js='modelmerger',
@@ -1771,6 +1779,7 @@ def create_ui():
                 custom_name,
                 checkpoint_format,
                 config_source,
+                bake_in_vae,
             ],
             outputs=[
                 primary_model_name,
diff --git a/style.css b/style.css
index 32ba4753..c10e32a1 100644
--- a/style.css
+++ b/style.css
@@ -641,6 +641,16 @@ canvas[key="mask"] {
     margin: 0.6em 0em 0.55em 0;
 }
 
+#modelmerger_results_container{
+    margin-top: 1em;
+    overflow: visible;
+}
+
+#modelmerger_models{
+    gap: 0;
+}
+
+
 #quicksettings .gr-button-tool{
     margin: 0;
 }
@@ -737,11 +747,6 @@ footer {
     line-height: 2.4em;
 }
 
-#modelmerger_results_container{
-    margin-top: 1em;
-    overflow: visible;
-}
-
 /* The following handles localization for right-to-left (RTL) languages like Arabic.
 The rtl media type will only be activated by the logic in javascript/localization.js.
 If you change anything above, you need to make sure it is RTL compliant by just running
-- 
cgit v1.2.3


From 54674674b813894b908283531ddaab4ccfeac721 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Thu, 19 Jan 2023 12:12:09 +0300
Subject: allow having at half precision when there is only one checkpoint in
 merger tab

---
 modules/extras.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index fe701a0e..d03f976e 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -278,6 +278,13 @@ def create_config(ckpt_result, config_source, a, b, c):
 chckpoint_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):
     shared.state.begin()
     shared.state.job = 'model-merge'
@@ -400,8 +407,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
             else:
                 theta_0[key] = theta_func2(a, b, multiplier)
 
-            if save_as_half:
-                theta_0[key] = theta_0[key].half()
+            theta_0[key] = to_half(theta_0[key], save_as_half)
 
         shared.state.sampling_step += 1
 
@@ -416,10 +422,14 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
         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] = vae_dict[key].half() if save_as_half else vae_dict[key]
+                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)
+
     ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path
 
     filename = filename_generator() if custom_name == '' else custom_name
-- 
cgit v1.2.3


From f2ae2529877072874ebaac0257fe4af48c5855a4 Mon Sep 17 00:00:00 2001
From: EllangoK <karun.ellango7@gmail.com>
Date: Thu, 19 Jan 2023 10:24:17 -0500
Subject: fixes minor typos around run_modelmerger

---
 modules/extras.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index d03f976e..1218f88f 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -275,7 +275,7 @@ def create_config(ckpt_result, config_source, a, b, c):
     shutil.copyfile(cfg, checkpoint_filename)
 
 
-chckpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
+checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"]
 
 
 def to_half(tensor, enable):
@@ -303,7 +303,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     def add_difference(theta0, theta1_2_diff, alpha):
         return theta0 + (alpha * theta1_2_diff)
 
-    def filename_weighed_sum():
+    def filename_weighted_sum():
         a = primary_model_info.model_name
         b = secondary_model_info.model_name
         Ma = round(1 - multiplier, 2)
@@ -311,7 +311,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
 
         return f"{Ma}({a}) + {Mb}({b})"
 
-    def filename_add_differnece():
+    def filename_add_difference():
         a = primary_model_info.model_name
         b = secondary_model_info.model_name
         c = tertiary_model_info.model_name
@@ -323,8 +323,8 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
         return primary_model_info.model_name
 
     theta_funcs = {
-        "Weighted sum": (filename_weighed_sum, None, weighted_sum),
-        "Add difference": (filename_add_differnece, get_difference, add_difference),
+        "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]
@@ -362,7 +362,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
         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 chckpoint_dict_skip_on_merge:
+            if key in checkpoint_dict_skip_on_merge:
                 continue
 
             if 'model' in key:
@@ -387,7 +387,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     for key in tqdm.tqdm(theta_0.keys()):
         if theta_1 and 'model' in key and key in theta_1:
 
-            if key in chckpoint_dict_skip_on_merge:
+            if key in checkpoint_dict_skip_on_merge:
                 continue
 
             a = theta_0[key]
-- 
cgit v1.2.3


From 112416d04171e4bee673f0adc9bd3aeba87ec71a Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 22 Jan 2023 10:17:12 +0300
Subject: add option to discard weights in checkpoint merger UI

---
 modules/extras.py | 9 ++++++++-
 modules/ui.py     | 4 ++++
 2 files changed, 12 insertions(+), 1 deletion(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 1218f88f..385430dc 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 import math
 import os
+import re
 import sys
 import traceback
 import shutil
@@ -285,7 +286,7 @@ def to_half(tensor, enable):
     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):
+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'
 
@@ -430,6 +431,12 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
         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
diff --git a/modules/ui.py b/modules/ui.py
index af6dfb21..eb4b7e6b 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1248,6 +1248,9 @@ def create_ui():
                             bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None", label="Bake in VAE", elem_id="modelmerger_bake_in_vae")
                             create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["None"] + list(sd_vae.vae_dict)}, "modelmerger_refresh_bake_in_vae")
 
+                with FormRow():
+                    discard_weights = gr.Textbox(value="", label="Discard weights with matching name", elem_id="modelmerger_discard_weights")
+
                 with gr.Row():
                     modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')
 
@@ -1838,6 +1841,7 @@ def create_ui():
                 checkpoint_format,
                 config_source,
                 bake_in_vae,
+                discard_weights,
             ],
             outputs=[
                 primary_model_name,
-- 
cgit v1.2.3


From b238b14ee459486c4734cc2899b83f547813a467 Mon Sep 17 00:00:00 2001
From: Andrey <16777216c@gmail.com>
Date: Sun, 22 Jan 2023 15:26:40 +0300
Subject: Split history extras.py to postprocessing.py

---
 modules/extras.py | 466 ------------------------------------------------------
 modules/temp      | 466 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 466 insertions(+), 466 deletions(-)
 delete mode 100644 modules/extras.py
 create mode 100644 modules/temp

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
deleted file mode 100644
index 385430dc..00000000
--- a/modules/extras.py
+++ /dev/null
@@ -1,466 +0,0 @@
-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.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:
-        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/temp b/modules/temp
new file mode 100644
index 00000000..385430dc
--- /dev/null
+++ b/modules/temp
@@ -0,0 +1,466 @@
+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.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:
+        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]
-- 
cgit v1.2.3


From c56b36712289020a98f0c77794b9045a251ecd55 Mon Sep 17 00:00:00 2001
From: Andrey <16777216c@gmail.com>
Date: Sun, 22 Jan 2023 15:26:41 +0300
Subject: Split history extras.py to postprocessing.py

---
 modules/extras.py | 466 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 modules/temp      | 466 ------------------------------------------------------
 2 files changed, 466 insertions(+), 466 deletions(-)
 create mode 100644 modules/extras.py
 delete mode 100644 modules/temp

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
new file mode 100644
index 00000000..385430dc
--- /dev/null
+++ b/modules/extras.py
@@ -0,0 +1,466 @@
+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.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:
+        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/temp b/modules/temp
deleted file mode 100644
index 385430dc..00000000
--- a/modules/temp
+++ /dev/null
@@ -1,466 +0,0 @@
-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.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:
-        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]
-- 
cgit v1.2.3


From 68303c96e5ab31576a8238a24bf5b6191cf16ed1 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 22 Jan 2023 15:38:39 +0300
Subject: split oversize extras.py to postprocessing.py

---
 modules/extras.py         | 217 +-------------------------------------
 modules/postprocessing.py | 257 +---------------------------------------------
 modules/ui.py             |  10 +-
 modules/ui_components.py  |   7 ++
 webui.py                  |   1 -
 5 files changed, 18 insertions(+), 474 deletions(-)

(limited to 'modules/extras.py')

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
-- 
cgit v1.2.3


From 41265a026de699cc223ca5b76c69b4e8e74aa7c1 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 23 Jan 2023 14:50:20 +0300
Subject: third time's the charm

---
 modules/extras.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index f04ddfc2..36123aa5 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -7,7 +7,7 @@ import torch
 import tqdm
 
 from modules import shared, images, sd_models, sd_vae
-from modules.ui import plaintext_to_html
+from modules.ui_common import plaintext_to_html
 import gradio as gr
 import safetensors.torch
 
-- 
cgit v1.2.3


From f4ec411f2c9d6bc6817a2eca8a2c00f255ffb386 Mon Sep 17 00:00:00 2001
From: "ULTRANOX\\Chris" <boston@udel.edu>
Date: Thu, 26 Jan 2023 03:45:16 -0500
Subject: Allow checkpoint merger to merge pix2pix models in the same way that
 it currently supports inpainting models.

---
 modules/extras.py | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 36123aa5..67ffdee3 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -132,6 +132,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if theta_func1 else None
 
     result_is_inpainting_model = False
+    result_is_pix2pix_model = False
 
     if theta_func2:
         shared.state.textinfo = f"Loading B"
@@ -186,13 +187,17 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
                 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
+                if a.shape[1] == 8 and b.shape[1] == 4:#If we have an InstructPix2Pix model...
+                    print("Detected possible merge of instruct model with non-instruct model.")
+                    theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)#Merge only the vectors the models have in common.  Otherwise we get an error due to dimension mismatch.
+                    result_is_pix2pix_model = True
+                else:
+                    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
@@ -226,6 +231,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
 
     filename = filename_generator() if custom_name == '' else custom_name
     filename += ".inpainting" if result_is_inpainting_model else ""
+    filename += ".pix2pix" if result_is_pix2pix_model else ""
     filename += "." + checkpoint_format
 
     output_modelname = os.path.join(ckpt_dir, filename)
-- 
cgit v1.2.3


From f90798c6b6cc48e514acb08ce02bdb5874bf74d8 Mon Sep 17 00:00:00 2001
From: "ULTRANOX\\Chris" <boston@udel.edu>
Date: Thu, 26 Jan 2023 04:38:04 -0500
Subject: Added error check for the rare case a user merges a pix2pix model
 with a normal model using weighted sum.  Also removed bad print message that
 interfered with merging progress bar.

---
 modules/extras.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 67ffdee3..badd13c7 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -186,9 +186,10 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
             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.")
+                if a.shape[1] == 4 and b.shape[1] == 8:
+                    raise RuntimeError("When merging pix2pix model with a normal one, A must be the pix2pix model.")
 
                 if a.shape[1] == 8 and b.shape[1] == 4:#If we have an InstructPix2Pix model...
-                    print("Detected possible merge of instruct model with non-instruct model.")
                     theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)#Merge only the vectors the models have in common.  Otherwise we get an error due to dimension mismatch.
                     result_is_pix2pix_model = True
                 else:
-- 
cgit v1.2.3


From 9e72dc743480c8b1ca6aeb8ced3af03f3e3243a3 Mon Sep 17 00:00:00 2001
From: "ULTRANOX\\Chris" <boston@udel.edu>
Date: Thu, 26 Jan 2023 06:05:40 -0500
Subject: Changed all references to "pix2pix" to the more precise name
 "instruct pix2pix".  Also changed extension to instrpix2pix at least for now.

---
 modules/extras.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index badd13c7..2bf0d17e 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -132,7 +132,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
     tertiary_model_info = sd_models.checkpoints_list[tertiary_model_name] if theta_func1 else None
 
     result_is_inpainting_model = False
-    result_is_pix2pix_model = False
+    result_is_instruct_pix2pix_model = False
 
     if theta_func2:
         shared.state.textinfo = f"Loading B"
@@ -187,11 +187,11 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
                 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.")
                 if a.shape[1] == 4 and b.shape[1] == 8:
-                    raise RuntimeError("When merging pix2pix model with a normal one, A must be the pix2pix model.")
+                    raise RuntimeError("When merging instruct-pix2pix model with a normal one, A must be the instruct-pix2pix model.")
 
-                if a.shape[1] == 8 and b.shape[1] == 4:#If we have an InstructPix2Pix model...
+                if a.shape[1] == 8 and b.shape[1] == 4:#If we have an Instruct-Pix2Pix model...
                     theta_0[key][:, 0:4, :, :] = theta_func2(a[:, 0:4, :, :], b, multiplier)#Merge only the vectors the models have in common.  Otherwise we get an error due to dimension mismatch.
-                    result_is_pix2pix_model = True
+                    result_is_instruct_pix2pix_model = True
                 else:
                     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)
@@ -232,7 +232,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
 
     filename = filename_generator() if custom_name == '' else custom_name
     filename += ".inpainting" if result_is_inpainting_model else ""
-    filename += ".pix2pix" if result_is_pix2pix_model else ""
+    filename += ".instrpix2pix" if result_is_instruct_pix2pix_model else ""
     filename += "." + checkpoint_format
 
     output_modelname = os.path.join(ckpt_dir, filename)
-- 
cgit v1.2.3


From cdc2fa209a3efdc71a90643a5e7a1df49869cd5f Mon Sep 17 00:00:00 2001
From: "ULTRANOX\\Chris" <boston@udel.edu>
Date: Thu, 26 Jan 2023 11:27:07 -0500
Subject: Changed filename addition from "instrpix2pix" to the more readable
 ".instruct-pix2pix" for newly generated instruct pix2pix models.

---
 modules/extras.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 2bf0d17e..466ecc15 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -232,7 +232,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
 
     filename = filename_generator() if custom_name == '' else custom_name
     filename += ".inpainting" if result_is_inpainting_model else ""
-    filename += ".instrpix2pix" if result_is_instruct_pix2pix_model else ""
+    filename += ".instruct-pix2pix" if result_is_instruct_pix2pix_model else ""
     filename += "." + checkpoint_format
 
     output_modelname = os.path.join(ckpt_dir, filename)
-- 
cgit v1.2.3


From cc8c9b7474d917888a0bd069fcd59a458c67ae4b Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Fri, 27 Jan 2023 22:43:08 +0300
Subject: fix broken calls to find_checkpoint_config

---
 modules/extras.py         | 4 ++--
 modules/sd_hijack_ip2p.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 36123aa5..4f842be9 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -6,7 +6,7 @@ import shutil
 import torch
 import tqdm
 
-from modules import shared, images, sd_models, sd_vae
+from modules import shared, images, sd_models, sd_vae, sd_models_config
 from modules.ui_common import plaintext_to_html
 import gradio as gr
 import safetensors.torch
@@ -37,7 +37,7 @@ def run_pnginfo(image):
 
 def create_config(ckpt_result, config_source, a, b, c):
     def config(x):
-        res = sd_models.find_checkpoint_config(x) if x else None
+        res = sd_models_config.find_checkpoint_config_near_filename(x) if x else None
         return res if res != shared.sd_default_config else None
 
     if config_source == 0:
diff --git a/modules/sd_hijack_ip2p.py b/modules/sd_hijack_ip2p.py
index 635f015f..3c727d3b 100644
--- a/modules/sd_hijack_ip2p.py
+++ b/modules/sd_hijack_ip2p.py
@@ -5,9 +5,9 @@ import gc
 import time
 
 def should_hijack_ip2p(checkpoint_info):
-    from modules import sd_models
+    from modules import sd_models_config
 
     ckpt_basename = os.path.basename(checkpoint_info.filename).lower()
-    cfg_basename = os.path.basename(sd_models.find_checkpoint_config(checkpoint_info)).lower()
+    cfg_basename = os.path.basename(sd_models_config.find_checkpoint_config_near_filename(checkpoint_info)).lower()
 
     return "pix2pix" in ckpt_basename and not "pix2pix" in cfg_basename
-- 
cgit v1.2.3