Merge branch 'master' into inpaint_textual_inversion

1 files changed, 94 insertions, 55 deletions

diff --git a/modules/extras.py b/modules/extras.py
index 8e2ab35c..d665440a 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,6 +1,8 @@
 from __future__ import annotations

 import math

 import os

+import sys

+import traceback

 

 import numpy as np

 from PIL import Image

@@ -12,7 +14,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

+from modules import processing, shared, images, devices, sd_models, sd_samplers

 from modules.shared import opts

 import modules.gfpgan_model

 from modules.ui import plaintext_to_html

@@ -20,7 +22,7 @@ import modules.codeformer_model
 import piexif

 import piexif.helper

 import gradio as gr

-

+import safetensors.torch

 

 class LruCache(OrderedDict):

     @dataclass(frozen=True)

@@ -53,14 +55,17 @@ class LruCache(OrderedDict):
 cached_images: LruCache = LruCache(max_size=5)

 

 

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

+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:

@@ -92,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)

 

@@ -102,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)

 

@@ -112,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:

@@ -136,12 +144,13 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
 

     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=hash(np.array(image.getdata()).tobytes()),

+            cache_key = LruCache.Key(image_hash=image_hash,

                                      info_hash=hash(info),

-                                     args_hash=hash((upscale_args, upscale_first)))

+                                     args_hash=hash(upscale_args))

             cached_entry = cached_images.get(cache_key)

             if cached_entry is None:

                 res = upscale(image, *upscale_args)

@@ -177,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")

@@ -185,18 +197,25 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_
         for op in extras_ops:

             image, info = op(image, info)

 

-        if opts.use_original_name_batch and image_name != None:

+        if opts.use_original_name_batch and image_name is not None:

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

         else:

             basename = ''

 

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

-

-        if opts.enable_pnginfo:

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

 

@@ -212,25 +231,8 @@ def run_pnginfo(image):
     if image is None:

         return '', '', ''

 

-    items = image.info

-    geninfo = ''

-

-    if "exif" in image.info:

-        exif = piexif.load(image.info["exif"])

-        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b'')

-        try:

-            exif_comment = piexif.helper.UserComment.load(exif_comment)

-        except ValueError:

-            exif_comment = exif_comment.decode('utf8', errors="ignore")

-

-        items['exif comment'] = exif_comment

-        geninfo = exif_comment

-

-        for field in ['jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',

-                      'loop', 'background', 'timestamp', 'duration']:

-            items.pop(field, None)

-

-    geninfo = items.get('parameters', geninfo)

+    geninfo, items = images.read_info_from_image(image)

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

 

     info = ''

     for key, text in items.items():

@@ -248,7 +250,10 @@ def run_pnginfo(image):
     return '', geninfo, info

 

 

-def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_name, interp_method, multiplier, save_as_half, custom_name):

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

 

@@ -260,23 +265,8 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
 

     primary_model_info = sd_models.checkpoints_list[primary_model_name]

     secondary_model_info = sd_models.checkpoints_list[secondary_model_name]

-    teritary_model_info = sd_models.checkpoints_list.get(teritary_model_name, None)

-

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

-    primary_model = torch.load(primary_model_info.filename, map_location='cpu')

-    theta_0 = sd_models.get_state_dict_from_checkpoint(primary_model)

-

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

-    secondary_model = torch.load(secondary_model_info.filename, map_location='cpu')

-    theta_1 = sd_models.get_state_dict_from_checkpoint(secondary_model)

-

-    if teritary_model_info is not None:

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

-        teritary_model = torch.load(teritary_model_info.filename, map_location='cpu')

-        theta_2 = sd_models.get_state_dict_from_checkpoint(teritary_model)

-    else:

-        teritary_model = None

-        theta_2 = None

+    tertiary_model_info = sd_models.checkpoints_list.get(tertiary_model_name, None)

+    result_is_inpainting_model = False

 

     theta_funcs = {

         "Weighted sum": (None, weighted_sum),

@@ -284,9 +274,19 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
     }

     theta_func1, theta_func2 = theta_funcs[interp_method]

 

-    print(f"Merging...")

+    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')

 

     if theta_func1:

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

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

+

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

             if 'model' in key:

                 if key in theta_2:

@@ -294,12 +294,33 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
                     theta_1[key] = theta_func1(theta_1[key], t2)

                 else:

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

-    del theta_2, teritary_model

+        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')

+

+    print("Merging...")

 

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

         if 'model' in key and key in theta_1:

+            a = theta_0[key]

+            b = theta_1[key]

 

-            theta_0[key] = theta_func2(theta_0[key], theta_1[key], multiplier)

+            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

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

 

             if save_as_half:

                 theta_0[key] = theta_0[key].half()

@@ -310,17 +331,35 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
             theta_0[key] = theta_1[key]

             if save_as_half:

                 theta_0[key] = theta_0[key].half()

+    del theta_1

 

     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.ckpt'

-    filename = filename if custom_name == '' else (custom_name + '.ckpt')

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

+

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

 

+    shared.state.textinfo = f"Saving to {output_modelname}..."

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

-    torch.save(primary_model, 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()

 

-    print(f"Checkpoint saved.")

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