diff options
Diffstat (limited to 'modules/processing.py')
| -rw-r--r-- | modules/processing.py | 33 |
1 files changed, 33 insertions, 0 deletions
diff --git a/modules/processing.py b/modules/processing.py index f902b9df..94fe2625 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -29,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
@@ -471,6 +478,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter 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,
@@ -489,6 +497,8 @@ 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,
@@ -522,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():
@@ -977,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
|