From 5fdf01ff058f6b1e13011dd7727ac33f062ba781 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 8 Aug 2026 04:10:22 +0100 Subject: [PATCH] feat(video): stage labels for modular generation progress Modular pipelines run every stage inside one pipeline call, leaving the ui on a single inference label. Forward hooks on the text encoder, transformer and vae decoders now surface the current stage through shared state, and the interrupt check runs in every stage so encodes and tiled decodes abort promptly. - saving a model registers a job instead of appearing idle - group offload logs each component before the pin step instead of only after completion --- modules/sd_models.py | 3 +++ modules/sd_offload.py | 1 + modules/video_models/video_modular.py | 28 +++++++++++++++++++++++++++ 3 files changed, 32 insertions(+) diff --git a/modules/sd_models.py b/modules/sd_models.py index f5f38f1ea..c9c495781 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1644,6 +1644,7 @@ def save_model(name: str, path: str | None = None, shard: str = "5GB", overwrite torch.cuda.synchronize() except Exception: pass + jobid = shared.state.begin('Save model') try: t0 = time.time() log.info(f'Save model: path="{model_name}" cls={shared.sd_model.__class__.__name__} start') @@ -1660,6 +1661,8 @@ def save_model(name: str, path: str | None = None, shard: str = "5GB", overwrite log.error(f'Save model: path="{model_name}" {e}') errors.display(e, 'Save model') return f'Error: {e}' + finally: + shared.state.end(jobid) def list_hfcache(): diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 4c94ae081..89eddae81 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -164,6 +164,7 @@ def apply_group_offload_component(module, module_name: str, main: bool, op: str module = accelerate.hooks.remove_hook_from_module(module, recurse=True) remove_group_offload_component(module) module.requires_grad_(False) + log.debug(f'Setting {op}: offload=group op=apply module={module_name} pin={cfg["use_stream"] and not cfg["low_cpu_mem_usage"]}') # before the apply: pinning large components takes a while and would otherwise run silently apply_group_offloading(module, onload_device=devices.device, offload_device=devices.cpu, **cfg) module.sdnext_group_offload_sig = sig return True diff --git a/modules/video_models/video_modular.py b/modules/video_models/video_modular.py index 3424695aa..3229fac01 100644 --- a/modules/video_models/video_modular.py +++ b/modules/video_models/video_modular.py @@ -165,7 +165,16 @@ def install_state_hook(pipe): if not any(isinstance(f, InterruptLogFilter) for f in runner_log.filters): runner_log.addFilter(InterruptLogFilter()) + def set_phase(phase: str): + # every stage runs inside one pipeline call, so the forward hooks are the only + # place the current stage is visible; state.begin clears the label per job + if getattr(pipe, 'sdnext_phase', None) != phase: + pipe.sdnext_phase = phase + shared.state.textinfo = phase + log.debug(f'Video modular: cls={pipe.__class__.__name__} phase="{phase}"') + def state_hook(module, args): # pylint: disable=unused-argument + set_phase('Generate') if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0: shared.state.sampling_steps = pipe.num_timesteps if shared.state.paused: @@ -178,8 +187,27 @@ def install_state_hook(pipe): if shared.state.interrupted or shared.state.skipped: raise AssertionError('Interrupted...') + def encode_hook(module, args): # pylint: disable=unused-argument + set_phase('Text encode') + if shared.state.interrupted or shared.state.skipped: + raise AssertionError('Interrupted...') + + def decode_hook(module, args): # pylint: disable=unused-argument + set_phase('Decode') + if shared.state.interrupted or shared.state.skipped: # fires per tile, so tiled decodes abort promptly + raise AssertionError('Interrupted...') + for name in ('transformer', 'transformer_ref'): module = getattr(pipe, name, None) if module is None or getattr(module, 'sdnext_state_hook', None) is not None: continue module.sdnext_state_hook = module.register_forward_pre_hook(state_hook) + text_encoder = getattr(pipe, 'text_encoder', None) + if text_encoder is not None: + target = getattr(text_encoder, 'model', text_encoder) # 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(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(decode_hook)