From 68d892e6b2f5d81c96855095a8d55014c5536eac Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 28 Aug 2026 13:31:54 +0200 Subject: [PATCH] modular pipelines intercepts and profiling Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 17 +++- extensions-builtin/sdnq | 2 +- modules/modular_load.py | 95 +----------------- modules/processing.py | 4 + modules/processing_diffusers.py | 2 +- modules/processing_info.py | 2 +- modules/sd_hijack_modular.py | 156 +++++++++++++++++++++++++++++ modules/timer.py | 1 + modules/video_models/video_load.py | 4 +- 9 files changed, 183 insertions(+), 100 deletions(-) create mode 100644 modules/sd_hijack_modular.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8390e4e5c..9ddbec6f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,12 +1,21 @@ # Change Log for SD.Next -## Update for 2026-08-27 +## Update for 2026-08-28 - **LoRA** - - new calibration engine that allows lora to be applied with far smaller error when dealing with highly quantized models - new apply engine that allows lora to be applied much faster - - *note*: calibration data is stored once calculated so it can be reused for future runs, location is `models/calibration` folder + - new calibration engine that allows lora to be applied with far smaller error when dealing with highly quantized models + - *TODO*: see [LoRA docs](https://vladmandic.github.io/sdnext-docs/LoRA) for details and usage instructions + - *note*: calibration data is stored once calculated so it can be reused for future runs + location is `models/calibration` folder + - new factor cache that allows lora effects to be pre-calculated and persistently cached for future runs + location is `models/lora-factor-cache` folder +- **Attention** + - new `sparse-attention` method that can be combined with other attention methods + to reduce memory usage and improve performance on large models + *TODO*: see [Attention docs](https://vladmandic.github.io/sdnext-docs/Attention) for details and usage instructions - **Internal** + - modular pipelines intercept and profiling hooks - attention mechanisms decision tree and apply method refactor ## Update for 2026-08-26 @@ -15,7 +24,7 @@ Time for a new release, *this is a large one*! Main focus is improving video workflows which also brings full support for new [MiniMax H3](https://vladmandic.github.io/sdnext-docs/MiniMax) and [LTXVideo-2.5](https://vladmandic.github.io/sdnext-docs/LTX) -and improves general video processing with flexible video upscaling, updated interpolation, etc. +and improvements to general video processing with flexible video upscaling, updated interpolation, etc. *What else?* - [Detailer.next](https://vladmandic.github.io/sdnext-docs/Detailer) with new support for *vision-language models* and *per-class prompts* diff --git a/extensions-builtin/sdnq b/extensions-builtin/sdnq index fee789419..530f4d96c 160000 --- a/extensions-builtin/sdnq +++ b/extensions-builtin/sdnq @@ -1 +1 @@ -Subproject commit fee7894192ef89805584a29d0cbf731d07b1b9aa +Subproject commit 530f4d96c0ccd755246d894e474202a1b16cdf5a diff --git a/modules/modular_load.py b/modules/modular_load.py index 490fbe8f4..06245934e 100644 --- a/modules/modular_load.py +++ b/modules/modular_load.py @@ -1,15 +1,7 @@ import time -import logging import torch -from modules import shared, errors, devices, sd_offload +from modules import shared, errors, devices, sd_hijack_modular from modules.logger import log -from modules.attention import context as attention_context - - -class InterruptLogFilter(logging.Filter): - """Drops the per-block error dumps the modular runner logs when an interrupt raises through it.""" - def filter(self, record): - return 'Interrupted...' not in record.msg def apply_progress_bar_config(block): @@ -32,87 +24,6 @@ def trace_modules(pipe): log.trace(f'Module: name={module_name} cls={module.__class__.__name__} device={next(module.parameters()).device} dtype={next(module.parameters()).dtype}') -def install_state_hook(pipe): - runner_log = logging.getLogger('diffusers.modular_pipelines.modular_pipeline') - if not any(isinstance(f, InterruptLogFilter) for f in runner_log.filters): - runner_log.addFilter(InterruptLogFilter()) - - def set_phase(phase: str, module: torch.nn.Module | None = None): - # every stage runs inside one pipeline call, so the forward hooks are the only place the current stage is visible - if getattr(pipe, 'sdnext_phase', None) != phase: - pipe.sdnext_phase = phase - jobid = getattr(pipe, 'sdnext_phaseid', None) # previous jobid if any - shared.state.end(jobid) # clear the previous job if exists - pipe.sdnext_phaseid = shared.state.begin(phase) # start a new job for the current phase - log.debug(f'Pipeline: phase={phase.replace(" ", "")} cls={pipe.__class__.__name__} module={module.__class__.__name__ if module is not None else None}') - return True - return False - - def _pre_transformer_hook(module, args): # pylint: disable=unused-argument - new_phase = set_phase('Generate', module) - attention_context.set_role('transformer') - if new_phase: - sd_offload.offload_ondemand(pipe, exclude=['transformer', 'transformer_ref'], reason='generate', force=hasattr(pipe, 'sdnext_force_offload')) - if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0: - shared.state.sampling_steps = pipe.num_timesteps - if shared.state.paused: - log.debug('Sampling paused') - while shared.state.paused: - if shared.state.interrupted or shared.state.skipped: - raise AssertionError('Interrupted...') - time.sleep(0.1) - shared.state.step() - attention_context.tick() - if shared.state.interrupted or shared.state.skipped: - raise AssertionError('Interrupted...') - - def _pre_text_encode_hook(module, args): # pylint: disable=unused-argument - new_phase = set_phase('Text Encode', module) - attention_context.set_role('te') - if new_phase: - sd_offload.offload_ondemand(pipe, exclude=['text_encoder'], reason='text encode', force=hasattr(pipe, 'sdnext_force_offload')) - if shared.state.interrupted or shared.state.skipped: - raise AssertionError('Interrupted...') - - def _pre_vae_decode_hook(module, args): # pylint: disable=unused-argument - new_phase = set_phase('Decode', module) - attention_context.set_role('vae') - if new_phase: - sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae decode', force=hasattr(pipe, 'sdnext_force_offload')) - if shared.state.interrupted or shared.state.skipped: # fires per tile, so tiled decodes abort promptly - raise AssertionError('Interrupted...') - - def _pre_vae_encode_hook(module, args): # pylint: disable=unused-argument - new_phase = set_phase('Encode', module) - attention_context.set_role('vae') - if new_phase: - sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae encode', force=hasattr(pipe, 'sdnext_force_offload')) - if shared.state.interrupted or shared.state.skipped: # fires per tile, so tiled encodes abort promptly - raise AssertionError('Interrupted...') - - for name in ('transformer', 'transformer_ref'): - module = getattr(pipe, name, None) - if module is not None: - target = getattr(module, 'model', module) # conditioning calls the inner model directly - if isinstance(target, torch.nn.Module) and getattr(target, 'sdnext_state_hook', None) is None: - target.sdnext_state_hook = target.register_forward_pre_hook(_pre_transformer_hook) - - for name in ('text_encoder', 'text_encoder_2'): - module = getattr(pipe, name, None) - if module is not None: - target = getattr(module, 'model', module) # conditioning calls the inner model directly - if isinstance(target, torch.nn.Module) and getattr(target, 'sdnext_state_hook', None) is None: - target.sdnext_state_hook = target.register_forward_pre_hook(_pre_text_encode_hook) - - for name in ('vae', 'audio_vae'): - decoder = getattr(getattr(pipe, name, None), 'decoder', None) # decode entry points bypass forward, the inner decoder does not - if isinstance(decoder, torch.nn.Module) and getattr(decoder, 'sdnext_state_hook', None) is None: - decoder.sdnext_state_hook = decoder.register_forward_pre_hook(_pre_vae_decode_hook) - encoder = getattr(getattr(pipe, name, None), 'encoder', None) # decode entry points bypass forward, the inner encoder does not - if isinstance(encoder, torch.nn.Module) and getattr(encoder, 'sdnext_state_hook', None) is None: - encoder.sdnext_state_hook = encoder.register_forward_pre_hook(_pre_vae_encode_hook) - - def is_modular(obj) -> bool: if obj is None: return False @@ -220,7 +131,9 @@ def load_modular_pipe(repo_cls, repo: str, workflow: str | None = None, revision # diffusers logger, so the reason is in the log above this line rather than in the exception path log.error(f'Load modular: cls={pipe.__class__.__name__} workflow={workflow} missing={missing} components the workflow requires did not load') - install_state_hook(pipe) + sd_hijack_modular.install_state_hook(pipe) + sd_hijack_modular.register_callbacks(pipe) + apply_progress_bar_config(pipe._blocks) # pylint: disable=protected-access return pipe except Exception as e: diff --git a/modules/processing.py b/modules/processing.py index 39d832802..fbb2f68a1 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -440,6 +440,10 @@ def print_stats(): if dynamo_dct: log.debug(f'Processed: dynamo={dynamo_dct}') + if timer.blocks.get_total() > 0.1: + log.debug(f'Processed: blocks={timer.blocks.dct(min_time=0.1, no_total=True)}') + timer.blocks.reset() + def process_images_inner(p: StableDiffusionProcessing) -> Processed: t0 = time.time() diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 9e71050e3..0f975ea72 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -86,12 +86,12 @@ def process_pre(p: processing.StableDiffusionProcessing): cfgzero.apply(p) linfusion.apply(shared.sd_model) cachedit.apply_cache_dit(shared.sd_model) - # apply-only sd_hijack_freeu.apply_freeu(p) transformer_cache.set_cache() para_attention.apply_first_block_cache() teacache.apply_teacache(p) + except Exception as e: log.error(f'Processing apply: {e}') errors.display(e, 'apply') diff --git a/modules/processing_info.py b/modules/processing_info.py index ad87e6177..dca52112e 100644 --- a/modules/processing_info.py +++ b/modules/processing_info.py @@ -183,7 +183,7 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No args['ToMe'] = _tome if _tome != 0 else None elif _token_method == 'ToDo': args['ToDo'] = _todo if _todo != 0 else None - if hasattr(shared.sd_model, 'embedding_db') and len(shared.sd_model.embedding_db.embeddings_used) > 0: # register used embeddings + if hasattr(shared.sd_model, 'embedding_db') and (shared.sd_model.embedding_db is not None) and len(shared.sd_model.embedding_db.embeddings_used) > 0: # register used embeddings args['Embeddings'] = ', '.join(shared.sd_model.embedding_db.embeddings_used) # samplers diff --git a/modules/sd_hijack_modular.py b/modules/sd_hijack_modular.py new file mode 100644 index 000000000..3c151e2ab --- /dev/null +++ b/modules/sd_hijack_modular.py @@ -0,0 +1,156 @@ +import os +import time +import logging +import torch +import diffusers +from modules.logger import log +from modules import shared, sd_offload, timer +from modules.attention import context as attention_context + + +debug = os.environ.get('SD_MODULAR_DEBUG', None) is not None +intercepted = set() + + +def modular_intercept(self, components, state: diffusers.modular_pipelines.modular_pipeline.BlockState, *args, **kwargs): + t0 = time.time() + block = type(self).__name__ + keys = state if isinstance(state, list) else list(state.__dict__.keys()) + # run code before block call + result = self.__orig_call__(components, state, *args, **kwargs) + t1 = time.time() + timer.blocks.add(block, t1 - t0) + # run code after block call + # TODO modular: intercept latents and set current latents for preview + """ + if 'latents' in keys: + ... + t2 = time.time() + timer.blocks.add('callback', t2 - t1) + """ + if debug: + log.trace(f'Modular intercept: block={block} state={keys} time={t1 - t0:.4f}') + return result + + +def patch_blocks(blocks: diffusers.ModularPipelineBlocks): + """recursively walks the block tree and patches the CLS __call__ method""" + def _patch_recursive(current_block): + block_cls = type(current_block) + if (block_cls not in intercepted) and (block_cls != diffusers.ModularPipelineBlocks): + if callable(block_cls) and not getattr(block_cls, "_is_patched", False): + block_cls.__orig_call__ = block_cls.__call__ # store original call for reference + block_cls.__call__ = modular_intercept + block_cls._is_patched = True # pylint: disable=protected-access + intercepted.add(block_cls) + if debug: + log.trace(f'Modular hijack: {block_cls.__name__}') + for attr in ("sub_blocks", "blocks"): # recurse into child blocks if containers exist + sub = getattr(current_block, attr, None) + if isinstance(sub, dict): + for child in sub.values(): + if isinstance(child, diffusers.ModularPipelineBlocks): + _patch_recursive(child) + elif isinstance(sub, (list, tuple)): + for child in sub: + if isinstance(child, diffusers.ModularPipelineBlocks): + _patch_recursive(child) + + _patch_recursive(blocks) + + +def register_callbacks(pipe: diffusers.ModularPipeline): + intercepted.clear() + if not isinstance(pipe, diffusers.ModularPipeline): + return + try: + patch_blocks(pipe._blocks) # pylint: disable=protected-access + except Exception as e: + log.error(f'Modular intercept: {e}') + + +class InterruptLogFilter(logging.Filter): + """Drops the per-block error dumps the modular runner logs when an interrupt raises through it.""" + def filter(self, record): + return 'Interrupted...' not in record.msg + + +def install_state_hook(pipe): + runner_log = logging.getLogger('diffusers.modular_pipelines.modular_pipeline') + if not any(isinstance(f, InterruptLogFilter) for f in runner_log.filters): + runner_log.addFilter(InterruptLogFilter()) + + def set_phase(phase: str, module: torch.nn.Module | None = None): + # every stage runs inside one pipeline call, so the forward hooks are the only place the current stage is visible + if getattr(pipe, 'sdnext_phase', None) != phase: + pipe.sdnext_phase = phase + jobid = getattr(pipe, 'sdnext_phaseid', None) # previous jobid if any + shared.state.end(jobid) # clear the previous job if exists + pipe.sdnext_phaseid = shared.state.begin(phase) # start a new job for the current phase + log.debug(f'Pipeline: phase={phase.replace(" ", "")} cls={pipe.__class__.__name__} module={module.__class__.__name__ if module is not None else None}') + return True + return False + + def _pre_transformer_hook(module, args): # pylint: disable=unused-argument + new_phase = set_phase('Generate', module) + attention_context.set_role('transformer') + if new_phase: + sd_offload.offload_ondemand(pipe, exclude=['transformer', 'transformer_ref'], reason='generate', force=hasattr(pipe, 'sdnext_force_offload')) + if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0: + shared.state.sampling_steps = pipe.num_timesteps + if shared.state.paused: + log.debug('Sampling paused') + while shared.state.paused: + if shared.state.interrupted or shared.state.skipped: + raise AssertionError('Interrupted...') + time.sleep(0.1) + shared.state.step() + attention_context.tick() + if shared.state.interrupted or shared.state.skipped: + raise AssertionError('Interrupted...') + + def _pre_text_encode_hook(module, args): # pylint: disable=unused-argument + new_phase = set_phase('Text Encode', module) + attention_context.set_role('te') + if new_phase: + sd_offload.offload_ondemand(pipe, exclude=['text_encoder'], reason='text encode', force=hasattr(pipe, 'sdnext_force_offload')) + if shared.state.interrupted or shared.state.skipped: + raise AssertionError('Interrupted...') + + def _pre_vae_decode_hook(module, args): # pylint: disable=unused-argument + new_phase = set_phase('Decode', module) + attention_context.set_role('vae') + if new_phase: + sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae decode', force=hasattr(pipe, 'sdnext_force_offload')) + if shared.state.interrupted or shared.state.skipped: # fires per tile, so tiled decodes abort promptly + raise AssertionError('Interrupted...') + + def _pre_vae_encode_hook(module, args): # pylint: disable=unused-argument + new_phase = set_phase('Encode', module) + attention_context.set_role('vae') + if new_phase: + sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae encode', force=hasattr(pipe, 'sdnext_force_offload')) + if shared.state.interrupted or shared.state.skipped: # fires per tile, so tiled encodes abort promptly + raise AssertionError('Interrupted...') + + for name in ('transformer', 'transformer_ref'): + module = getattr(pipe, name, None) + if module is not None: + target = getattr(module, 'model', module) # conditioning calls the inner model directly + if isinstance(target, torch.nn.Module) and getattr(target, 'sdnext_state_hook', None) is None: + target.sdnext_state_hook = target.register_forward_pre_hook(_pre_transformer_hook) + + for name in ('text_encoder', 'text_encoder_2'): + module = getattr(pipe, name, None) + if module is not None: + target = getattr(module, 'model', module) # conditioning calls the inner model directly + if isinstance(target, torch.nn.Module) and getattr(target, 'sdnext_state_hook', None) is None: + target.sdnext_state_hook = target.register_forward_pre_hook(_pre_text_encode_hook) + + for name in ('vae', 'audio_vae'): + decoder = getattr(getattr(pipe, name, None), 'decoder', None) # decode entry points bypass forward, the inner decoder does not + if isinstance(decoder, torch.nn.Module) and getattr(decoder, 'sdnext_state_hook', None) is None: + decoder.sdnext_state_hook = decoder.register_forward_pre_hook(_pre_vae_decode_hook) + encoder = getattr(getattr(pipe, name, None), 'encoder', None) # decode entry points bypass forward, the inner encoder does not + if isinstance(encoder, torch.nn.Module) and getattr(encoder, 'sdnext_state_hook', None) is None: + encoder.sdnext_state_hook = encoder.register_forward_pre_hook(_pre_vae_encode_hook) diff --git a/modules/timer.py b/modules/timer.py index 2cdecc87e..daacd1e7d 100644 --- a/modules/timer.py +++ b/modules/timer.py @@ -106,4 +106,5 @@ launch = Timer() init = Timer() load = Timer() dynamo = Timer() +blocks = Timer() autotune = Timer(profile=True) diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py index 9fcae71d1..1e4343846 100644 --- a/modules/video_models/video_load.py +++ b/modules/video_models/video_load.py @@ -5,7 +5,7 @@ import time import torch import transformers import diffusers -from modules import shared, errors, sd_models, sd_checkpoint, model_quant, devices, sd_hijack_te, sd_hijack_vae, modular_load +from modules import shared, errors, sd_models, sd_checkpoint, model_quant, devices, modular_load, sd_hijack_te, sd_hijack_vae, sd_hijack_modular from modules.logger import log from modules.video_models import models_def, video_utils, video_overrides, video_cache from pipelines import generic @@ -196,7 +196,7 @@ def load_model(selected: models_def.Model): shared.sd_model = model_quant.do_post_load_quant(shared.sd_model, allow=False) sd_models.set_diffuser_offload(shared.sd_model) if modular_load.is_modular(shared.sd_model): - modular_load.install_state_hook(shared.sd_model) + sd_hijack_modular.install_state_hook(shared.sd_model) loaded_model = selected.name msg = f'Load video: cls={shared.sd_model.__class__.__name__} model="{selected.name}" time={t1-t0:.2f}'