Merge branch 'dev' into find_vae

1 files changed, 59 insertions, 7 deletions

diff --git a/modules/processing.py b/modules/processing.py
index e8808beb..94fe2625 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -2,7 +2,6 @@ import json
 import math

 import os

 import sys

-import warnings

 import hashlib

 

 import torch

@@ -11,10 +10,10 @@ from PIL import Image, ImageFilter, ImageOps
 import random

 import cv2

 from skimage import exposure

-from typing import Any, Dict, List, Optional

+from typing import Any, Dict, List

 

 import modules.sd_hijack

-from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, script_callbacks, extra_networks, sd_vae_approx, scripts

+from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, extra_networks, sd_vae_approx, scripts

 from modules.sd_hijack import model_hijack

 from modules.shared import opts, cmd_opts, state

 import modules.shared as shared

@@ -30,6 +29,13 @@ from ldm.models.diffusion.ddpm import LatentDepth2ImageDiffusion
 

 from einops import repeat, rearrange

 from blendmodes.blend import blendLayers, BlendType

+import tomesd

+

+# add a logger for the processing module

+logger = logging.getLogger(__name__)

+# manually set output level here since there is no option to do so yet through launch options

+# logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(name)s %(message)s')

+

 

 # some of those options should not be changed at all because they would break the model, so I removed them from options.

 opt_C = 4

@@ -165,7 +171,7 @@ class StableDiffusionProcessing:
         self.all_subseeds = None

         self.iteration = 0

         self.is_hr_pass = False

-        

+

 

     @property

     def sd_model(self):

@@ -458,10 +464,21 @@ def fix_seed(p):
     p.subseed = get_fixed_seed(p.subseed)

 

 

+def program_version():

+    import launch

+

+    res = launch.git_tag()

+    if res == "<none>":

+        res = None

+

+    return res

+

+

 def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iteration=0, position_in_batch=0):

     index = position_in_batch + iteration * p.batch_size

 

     clip_skip = getattr(p, 'clip_skip', opts.CLIP_stop_at_last_layers)

+    enable_hr = getattr(p, 'enable_hr', False)

 

     generation_params = {

         "Steps": p.steps,

@@ -480,16 +497,19 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Conditional mask weight": getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None,

         "Clip skip": None if clip_skip <= 1 else clip_skip,

         "ENSD": None if opts.eta_noise_seed_delta == 0 else opts.eta_noise_seed_delta,

+        "Token merging ratio": None if opts.token_merging_ratio == 0 else opts.token_merging_ratio,

+        "Token merging ratio hr": None if not enable_hr or opts.token_merging_ratio_hr == 0 else opts.token_merging_ratio_hr,

         "Init image hash": getattr(p, 'init_img_hash', None),

         "RNG": opts.randn_source if opts.randn_source != "GPU" else None,

         "NGMS": None if p.s_min_uncond == 0 else p.s_min_uncond,

+        "Version": program_version() if opts.add_version_to_infotext else None,

     }

 

     generation_params.update(p.extra_generation_params)

 

     generation_params_text = ", ".join([k if k == v else f'{k}: {generation_parameters_copypaste.quote(v)}' for k, v in generation_params.items() if v is not None])

 

-    negative_prompt_text = "\nNegative prompt: " + p.all_negative_prompts[index] if p.all_negative_prompts[index] else ""

+    negative_prompt_text = f"\nNegative prompt: {p.all_negative_prompts[index]}" if p.all_negative_prompts[index] else ""

 

     return f"{all_prompts[index]}{negative_prompt_text}\n{generation_params_text}".strip()

 

@@ -512,9 +532,18 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             if k == 'sd_vae':

                 sd_vae.reload_vae_weights()

 

+        if opts.token_merging_ratio > 0:

+            sd_models.apply_token_merging(sd_model=p.sd_model, hr=False)

+            logger.debug(f"Token merging applied to first pass. Ratio: '{opts.token_merging_ratio}'")

+

         res = process_images_inner(p)

 

     finally:

+        # undo model optimizations made by tomesd

+        if opts.token_merging_ratio > 0:

+            tomesd.remove_patch(p.sd_model)

+            logger.debug('Token merging model optimizations removed')

+

         # restore opts to original state

         if p.override_settings_restore_afterwards:

             for k, v in stored_opts.items():

@@ -653,7 +682,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if not shared.opts.dont_fix_second_order_samplers_schedule:

                 try:

                     step_multiplier = 2 if sd_samplers.all_samplers_map.get(p.sampler_name).aliases[0] in ['k_dpmpp_2s_a', 'k_dpmpp_2s_a_ka', 'k_dpmpp_sde', 'k_dpmpp_sde_ka', 'k_dpm_2', 'k_dpm_2_a', 'k_heun'] else 1

-                except:

+                except Exception:

                     pass

             uc = get_conds_with_caching(prompt_parser.get_learned_conditioning, negative_prompts, p.steps * step_multiplier, cached_uc)

             c = get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, prompts, p.steps * step_multiplier, cached_c)

@@ -769,7 +798,16 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 

     devices.torch_gc()

 

-    res = Processed(p, output_images, p.all_seeds[0], infotext(), comments="".join(["\n\n" + x for x in comments]), subseed=p.all_subseeds[0], index_of_first_image=index_of_first_image, infotexts=infotexts)

+    res = Processed(

+        p,

+        images_list=output_images,

+        seed=p.all_seeds[0],

+        info=infotext(),

+        comments="".join(f"\n\n{comment}" for comment in comments),

+        subseed=p.all_subseeds[0],

+        index_of_first_image=index_of_first_image,

+        infotexts=infotexts,

+    )

 

     if p.scripts is not None:

         p.scripts.postprocess(p, res)

@@ -958,8 +996,22 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         x = None

         devices.torch_gc()

 

+        # apply token merging optimizations from tomesd for high-res pass

+        if opts.token_merging_ratio_hr > 0:

+            # in case the user has used separate merge ratios

+            if opts.token_merging_ratio > 0:

+                tomesd.remove_patch(self.sd_model)

+                logger.debug('Adjusting token merging ratio for high-res pass')

+

+            sd_models.apply_token_merging(sd_model=self.sd_model, hr=True)

+            logger.debug(f"Applied token merging for high-res pass. Ratio: '{opts.token_merging_ratio_hr}'")

+

         samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)

 

+        if opts.token_merging_ratio_hr > 0 or opts.token_merging_ratio > 0:

+            tomesd.remove_patch(self.sd_model)

+            logger.debug('Removed token merging optimizations from model')

+

         self.is_hr_pass = False

 

         return samples