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 +++++++++++++----
 1 file changed, 13 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)
 
-- 
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 --
 1 file changed, 2 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
 
-- 
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 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

(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()
-- 
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

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

(limited to 'modules/extras.py')

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]
-- 
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

---
 modules/extras.py | 112 +++++++++++++++++++++++++++++++++---------------------
 1 file changed, 68 insertions(+), 44 deletions(-)

(limited to 'modules/extras.py')

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]
-- 
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 ++++++++-
 1 file changed, 8 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
-- 
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 ------------------------------------------------------
 1 file changed, 466 deletions(-)
 delete mode 100644 modules/extras.py

(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]
-- 
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 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 466 insertions(+)
 create mode 100644 modules/extras.py

(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]
-- 
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 +-----------------------------------------------------
 1 file changed, 1 insertion(+), 216 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:
-- 
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 ++--
 1 file changed, 2 insertions(+), 2 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:
-- 
cgit v1.2.3