diff --git a/modules/attention/__init__.py b/modules/attention/__init__.py index 7b923ca6b..359ab81ca 100644 --- a/modules/attention/__init__.py +++ b/modules/attention/__init__.py @@ -1,12 +1,12 @@ -"""Attention backends: one scaled_dot_product_attention router over the registered backends, plus the diffusers-side processor and dispatcher setup.""" +"""Attention backends: one scaled_dot_product_attention router over the registered backends, the per-generation context, and the diffusers-side processor and dispatcher setup.""" from modules.attention.registry import AttentionBackend, AttentionCall, Constraints, Platform, Registry, registry from modules.attention.router import Plan, PlanEntry, build_plan, get_plan, install_router from modules.attention.dispatcher import set_diffusers_attention, set_attention_dispatcher, hijack_kernels, get_kernel_hijack, get_hf_api_hijack -from modules.attention import backends +from modules.attention import backends, context __all__ = [ 'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry', 'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', 'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack', - 'backends', + 'backends', 'context', ] diff --git a/modules/attention/context.py b/modules/attention/context.py new file mode 100644 index 000000000..24aa8300a --- /dev/null +++ b/modules/attention/context.py @@ -0,0 +1,77 @@ +"""Per-generation state for attention consumers: the component running, the denoiser forward about to run, and the model.""" +from contextlib import contextmanager +from dataclasses import dataclass +import torch + + +@dataclass +class GenerationContext: + active: bool = False + role: str | None = None # 'transformer', 'te' or 'vae' while a generation runs, None outside one + step: int = 0 # index of the denoiser forward about to run + steps: int = 0 # forwards in the current pass + forwards: int = 0 + model_key: tuple[str, str | None] | None = None # pipeline class and denoiser class, telemetry only + step_buffer: torch.Tensor | None = None # the step as a device scalar updated in place, so compiled readers keep their graph + + +current = GenerationContext() + + +def denoiser_name(pipe) -> str | None: + for name in ('transformer', 'unet'): + module = getattr(pipe, name, None) + if module is not None: + return module.__class__.__name__ + return None + + +def begin(pipe, steps: int = 0) -> None: + from modules import devices + current.active = True + current.role = 'transformer' + current.model_key = (pipe.__class__.__name__, denoiser_name(pipe)) if pipe is not None else None + device = devices.device if devices.device is not None else torch.device('cpu') + if current.step_buffer is None or current.step_buffer.device != device: + current.step_buffer = torch.zeros((), dtype=torch.int64, device=device) + new_pass(steps) + + +def new_pass(steps: int = 0) -> None: + """Restart the step count for a denoising pass: base, hires or refiner.""" + current.steps = int(steps or 0) + current.forwards = 0 + set_step(0) + + +def set_step(step: int) -> None: + current.step = int(step) + if current.step_buffer is not None: + current.step_buffer.fill_(current.step) + + +def tick(step: int | None = None) -> None: + """Advance to the next forward: the classic callback passes the completed step plus one, the modular pre-hook passes nothing and counts forwards.""" + set_step(current.forwards if step is None else step) + current.forwards = current.step + 1 + + +def end() -> None: + current.active = False + current.role = None + current.model_key = None + new_pass(0) + + +def set_role(name: str | None) -> None: + current.role = name + + +@contextmanager +def role(name: str): + previous = current.role + current.role = name + try: + yield + finally: + current.role = previous diff --git a/modules/modular_load.py b/modules/modular_load.py index b8e8b8cc0..b5c784c39 100644 --- a/modules/modular_load.py +++ b/modules/modular_load.py @@ -3,6 +3,7 @@ import logging import torch from modules import shared, errors, devices, sd_offload from modules.logger import log +from modules.attention import context as attention_context class InterruptLogFilter(logging.Filter): @@ -41,6 +42,7 @@ def install_state_hook(pipe): 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: @@ -52,11 +54,13 @@ def install_state_hook(pipe): 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: @@ -64,6 +68,7 @@ def install_state_hook(pipe): 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 @@ -71,6 +76,7 @@ def install_state_hook(pipe): 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 diff --git a/modules/processing.py b/modules/processing.py index b72c8e151..0e919cd23 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -5,6 +5,7 @@ import numpy as np from PIL import Image, ImageOps from modules import shared, devices, errors, images, scripts_manager, memstats, script_callbacks, extra_networks, sd_models, sd_checkpoint, sd_vae, processing_helpers, processing_grading, timer, masking from modules.logger import log +from modules.attention import context as attention_context from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet from modules.processing_info import create_infotext from modules.processing_class import ( # pylint: disable=unused-import @@ -199,6 +200,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed | None: script_callbacks.before_process_callback(p) timer.process.record('pre') + attention_context.begin(shared.sd_model, p.steps) if shared.cmd_opts.profile: timer.startup.profile = True @@ -232,6 +234,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed | None: results = process_images_inner(p) finally: + attention_context.end() script_callbacks.after_process_callback(p) if p.override_settings_restore_afterwards: # restore opts to original state diff --git a/modules/processing_args.py b/modules/processing_args.py index 94a66153f..115bc6c61 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -8,6 +8,7 @@ import numpy as np from PIL import Image from modules import shared, sd_models, processing, processing_vae, processing_helpers, sd_hijack_hypertile, sd_vae from modules.logger import log +from modules.attention import context as attention_context from modules.processing_callbacks import diffusers_callback_legacy, diffusers_callback, set_callbacks_p from modules.processing_helpers import get_generator, apply_circular # pylint: disable=unused-import from modules.processing_prompt import set_prompt @@ -366,6 +367,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:l args['callback_steps'] = 1 set_callbacks_p(p) + attention_context.new_pass(steps) if 'prior_callback_on_step_end' in possible: # Wuerstchen / Cascade args['prior_callback_on_step_end'] = diffusers_callback if 'prior_callback_on_step_end_tensor_inputs' in possible: diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index 48ac6e75b..9b8728695 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -4,6 +4,7 @@ import torch import numpy as np from modules import shared, devices, processing_correction, timer, prompt_parser_diffusers from modules.logger import log +from modules.attention import context as attention_context p = None @@ -87,6 +88,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0: shared.state.sampling_steps = pipe.num_timesteps shared.state.step() + attention_context.tick(step + 1) if shared.state.interrupted or shared.state.skipped: raise AssertionError('Interrupted...') if latents is None or p is None: diff --git a/modules/sd_hijack_te.py b/modules/sd_hijack_te.py index 0e9111c09..06826cd33 100644 --- a/modules/sd_hijack_te.py +++ b/modules/sd_hijack_te.py @@ -2,6 +2,7 @@ import os import time from modules import shared, errors, timer, sd_models from modules.logger import log +from modules.attention import context as attention_context class PromptCache: @@ -65,10 +66,11 @@ def hijack_encode_prompt(*args, **kwargs): res = cached else: log.debug(f'Encode: prompt="{prompt}" hijack=True') - if hasattr(shared.sd_model, 'orig_encode_prompt'): - res = shared.sd_model.orig_encode_prompt(*args_copy, **kwargs) - else: - res = shared.sd_model.encode_prompt(*args_copy, **kwargs) + with attention_context.role('te'): + if hasattr(shared.sd_model, 'orig_encode_prompt'): + res = shared.sd_model.orig_encode_prompt(*args_copy, **kwargs) + else: + res = shared.sd_model.encode_prompt(*args_copy, **kwargs) prompt_cache.set(prompt, res) if hasattr(shared.sd_model, 'after_prompt_encode'): diff --git a/modules/sd_hijack_vae.py b/modules/sd_hijack_vae.py index cc6b83919..990495782 100644 --- a/modules/sd_hijack_vae.py +++ b/modules/sd_hijack_vae.py @@ -3,6 +3,7 @@ import time import torch from modules import shared, sd_models, devices, timer, errors from modules.logger import log +from modules.attention import context as attention_context debug = log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -25,20 +26,22 @@ def hijack_vae_decode(*args, **kwargs): sd_models.move_model(shared.sd_model.vae, devices.device) if torch.is_tensor(args[0]): latents = args[0].to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype - if hasattr(shared.sd_model.vae, '_asymmetric_upscale_vae'): - res = hijack_vae_upscale(latents, *args[1:], **kwargs) - elif getattr(shared.sd_model, 'sdnext_vae_type', None) == 'Tiny': - from modules.video_models import video_vae - res = video_vae.vae_decode_tiny(latents) # None when the model has no tiny counterpart, and it says so - if res is None: - res = shared.sd_model.vae.orig_decode(latents, *args[1:], **kwargs) + with attention_context.role('vae'): + if hasattr(shared.sd_model.vae, '_asymmetric_upscale_vae'): + res = hijack_vae_upscale(latents, *args[1:], **kwargs) + elif getattr(shared.sd_model, 'sdnext_vae_type', None) == 'Tiny': + from modules.video_models import video_vae + res = video_vae.vae_decode_tiny(latents) # None when the model has no tiny counterpart, and it says so + if res is None: + res = shared.sd_model.vae.orig_decode(latents, *args[1:], **kwargs) t1 = time.time() try: log.debug(f'Decode: vae={shared.sd_model.vae.__class__.__name__} dtype={latents.dtype} latents={list(latents.shape)}:{latents.device} decoded={list(res[0].shape)} slicing={getattr(shared.sd_model.vae, "use_slicing", None)} tiling={getattr(shared.sd_model.vae, "use_tiling", None)} time={t1-t0:.3f}') except Exception: pass else: - res = shared.sd_model.vae.orig_decode(*args, **kwargs) + with attention_context.role('vae'): + res = shared.sd_model.vae.orig_decode(*args, **kwargs) except Exception as e: log.error(f'Decode: vae={shared.sd_model.vae.__class__.__name__} {e}') errors.display(e, 'vae') @@ -58,11 +61,13 @@ def hijack_vae_encode(*args, **kwargs): sd_models.move_model(shared.sd_model.vae, devices.device) if torch.is_tensor(args[0]): latents = args[0].to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype - res = shared.sd_model.vae.orig_encode(latents, *args[1:], **kwargs) + with attention_context.role('vae'): + res = shared.sd_model.vae.orig_encode(latents, *args[1:], **kwargs) t1 = time.time() log.debug(f'Encode: vae={shared.sd_model.vae.__class__.__name__} slicing={getattr(shared.sd_model.vae, "use_slicing", None)} tiling={getattr(shared.sd_model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}') else: - res = shared.sd_model.vae.orig_encode(*args, **kwargs) + with attention_context.role('vae'): + res = shared.sd_model.vae.orig_encode(*args, **kwargs) except Exception as e: log.error(f'Encode: vae={shared.sd_model.vae.__class__.__name__} {e}') errors.display(e, 'vae') diff --git a/test/test-attention-router.py b/test/test-attention-router.py index 44a95c150..91c56145b 100644 --- a/test/test-attention-router.py +++ b/test/test-attention-router.py @@ -15,6 +15,8 @@ Covers: - a backend whose prepare raises is skipped without disturbing the rest - install_router leaves the original sdpa in place for an empty plan - the dynamic backend pins the pre-dynamic sdpa the sliced path reads +- the generation context: step normalized to the forward about to run on both the classic + callback and the modular pre-hook, per-pass resets, the in-place step buffer, role scopes No running server required. Nothing is moved to the accelerator. @@ -309,6 +311,57 @@ def test_dynamic_backend_pins_pre_dynamic_sdpa(): return True +def test_context_classic_ticks_follow_the_callback(): + ctx = attention.context + + class Pipe: + transformer = object() + + ctx.begin(Pipe(), steps=4) + assert ctx.current.active and ctx.current.role == 'transformer' and ctx.current.step == 0 and ctx.current.steps == 4 + assert ctx.current.model_key == ('Pipe', 'object'), ctx.current.model_key + buffer = ctx.current.step_buffer + for completed in range(4): + ctx.tick(completed + 1) # the diffusers callback reports the step just completed + assert ctx.current.step == completed + 1 + assert ctx.current.step_buffer is buffer and int(buffer.item()) == completed + 1 + ctx.new_pass(2) # hires or refiner pass + assert ctx.current.step == 0 and ctx.current.steps == 2 and int(buffer.item()) == 0 + ctx.end() + assert not ctx.current.active and ctx.current.role is None and ctx.current.model_key is None and ctx.current.step == 0 + return True + + +def test_context_modular_ticks_count_forwards(): + ctx = attention.context + ctx.begin(None, steps=3) + assert ctx.current.model_key is None + for expected in range(3): + ctx.tick() # the modular pre-hook fires before each forward + assert ctx.current.step == expected, ctx.current.step + ctx.end() + return True + + +def test_context_roles_nest_and_stick(): + ctx = attention.context + ctx.begin(None) + with ctx.role('te'): + assert ctx.current.role == 'te' + with ctx.role('vae'): + assert ctx.current.role == 'vae' + assert ctx.current.role == 'te' + assert ctx.current.role == 'transformer' + ctx.set_role('vae') + assert ctx.current.role == 'vae' + ctx.end() + assert ctx.current.role is None + with ctx.role('te'): # outside a generation the scope still restores what it found + assert ctx.current.role == 'te' + assert ctx.current.role is None + return True + + def run_all(): log.warning('=== attention router ===') cat = category('router') @@ -324,6 +377,15 @@ def run_all(): ]: run_test(cat, fn) + log.warning('=== generation context ===') + cat = category('context') + for fn in [ + test_context_classic_ticks_follow_the_callback, + test_context_modular_ticks_count_forwards, + test_context_roles_nest_and_stick, + ]: + run_test(cat, fn) + log.warning('=== Results ===') total_passed = 0 total_failed = 0