From b272422fc5d6bc78578d27a8aeb4b381ec06f820 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 09:34:34 +0000 Subject: [PATCH 01/10] Add IterativePipelineBlocks: composable loop blocks with loop-local state scopes Adds a loop composite whose sub-blocks are ordinary state blocks, so loops can nest and compose freely (e.g. an autoregressive chunk loop containing a timestep denoise loop). Loop variables like the current timestep are provided through a loop-local scope on PipelineState (`loop_scope()` / `set_local`): they resolve sub-blocks' declared inputs while the loop runs and are discarded when it exits, so they never surface as pipeline inputs. Sub-blocks declare everything they consume; declared outputs persist as usual. The loop block declares its own surface symmetrically to LoopSequentialPipelineBlocks: loop_inputs, loop_locals (names it provides via the scope), loop_intermediate_outputs, loop_expected_components/configs. Subclasses hand-write `__call__` around `loop_step()`, same idiom as leaf blocks around `get_block_state`. Ports the flux2 denoise loops (flux2, klein, klein-base) as the reference example, moves `progress_bar` to the ModularPipelineBlocks base, and treats the new class as a leaf in workflow traversal like LoopSequential. Adds structure/execution/nesting tests modeled on the helios chunk-loop use case. Co-Authored-By: Claude Fable 5 --- src/diffusers/__init__.py | 2 + src/diffusers/modular_pipelines/__init__.py | 2 + .../modular_pipelines/flux2/denoise.py | 182 +++++++++---- .../modular_pipelines/modular_pipeline.py | 186 ++++++++++++-- .../test_iterative_pipeline_blocks.py | 243 ++++++++++++++++++ 5 files changed, 535 insertions(+), 80 deletions(-) create mode 100644 tests/modular_pipelines/test_iterative_pipeline_blocks.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index dcccf5cd2de3..25cfd4f93a31 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -342,6 +342,7 @@ "ConditionalPipelineBlocks", "ConfigSpec", "InputParam", + "IterativePipelineBlocks", "LoopSequentialPipelineBlocks", "ModularPipeline", "ModularPipelineBlocks", @@ -1212,6 +1213,7 @@ ConditionalPipelineBlocks, ConfigSpec, InputParam, + IterativePipelineBlocks, LoopSequentialPipelineBlocks, ModularPipeline, ModularPipelineBlocks, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 25db2ef3bee2..7a11405f3317 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -34,6 +34,7 @@ "AutoPipelineBlocks", "SequentialPipelineBlocks", "ConditionalPipelineBlocks", + "IterativePipelineBlocks", "LoopSequentialPipelineBlocks", "PipelineState", "BlockState", @@ -160,6 +161,7 @@ AutoPipelineBlocks, BlockState, ConditionalPipelineBlocks, + IterativePipelineBlocks, LoopSequentialPipelineBlocks, ModularPipeline, ModularPipelineBlocks, diff --git a/src/diffusers/modular_pipelines/flux2/denoise.py b/src/diffusers/modular_pipelines/flux2/denoise.py index 1a782e70de33..f455223dde86 100644 --- a/src/diffusers/modular_pipelines/flux2/denoise.py +++ b/src/diffusers/modular_pipelines/flux2/denoise.py @@ -22,8 +22,7 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging from ..modular_pipeline import ( - BlockState, - LoopSequentialPipelineBlocks, + IterativePipelineBlocks, ModularPipelineBlocks, PipelineState, ) @@ -53,8 +52,8 @@ def expected_components(self) -> list[ComponentSpec]: def description(self) -> str: return ( "Step within the denoising loop that denoises the latents for Flux2. " - "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " - "object (e.g. `Flux2DenoiseLoopWrapper`)" + "This block should be used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`); it reads the current timestep `t` from the loop scope." ) @property @@ -101,12 +100,22 @@ def inputs(self) -> list[tuple[str, Any]]: type_hint=torch.Tensor, description="4D position IDs for latent tokens (T, H, W, L)", ), + InputParam( + "t", + required=True, + type_hint=torch.Tensor, + description="The current timestep, provided by the denoise loop scope.", + ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted noise for this step")] + @torch.no_grad() - def __call__( - self, components: Flux2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor - ) -> PipelineState: + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + latents = block_state.latents latent_model_input = latents.to(components.transformer.dtype) img_ids = block_state.latent_ids @@ -117,7 +126,7 @@ def __call__( image_latent_ids = block_state.image_latent_ids img_ids = torch.cat([img_ids, image_latent_ids], dim=1) - timestep = t.expand(latents.shape[0]).to(latents.dtype) + timestep = block_state.t.expand(latents.shape[0]).to(latents.dtype) noise_pred = components.transformer( hidden_states=latent_model_input, @@ -133,7 +142,8 @@ def __call__( noise_pred = noise_pred[:, : latents.size(1)] block_state.noise_pred = noise_pred - return components, block_state + self.set_block_state(state, block_state) + return components, state # same as Flux2LoopDenoiser but guidance=None @@ -148,8 +158,8 @@ def expected_components(self) -> list[ComponentSpec]: def description(self) -> str: return ( "Step within the denoising loop that denoises the latents for Flux2. " - "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " - "object (e.g. `Flux2DenoiseLoopWrapper`)" + "This block should be used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`); it reads the current timestep `t` from the loop scope." ) @property @@ -190,12 +200,22 @@ def inputs(self) -> list[tuple[str, Any]]: type_hint=torch.Tensor, description="4D position IDs for latent tokens (T, H, W, L)", ), + InputParam( + "t", + required=True, + type_hint=torch.Tensor, + description="The current timestep, provided by the denoise loop scope.", + ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted noise for this step")] + @torch.no_grad() - def __call__( - self, components: Flux2KleinModularPipeline, block_state: BlockState, i: int, t: torch.Tensor - ) -> PipelineState: + def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + latents = block_state.latents latent_model_input = latents.to(components.transformer.dtype) img_ids = block_state.latent_ids @@ -206,7 +226,7 @@ def __call__( image_latent_ids = block_state.image_latent_ids img_ids = torch.cat([img_ids, image_latent_ids], dim=1) - timestep = t.expand(latents.shape[0]).to(latents.dtype) + timestep = block_state.t.expand(latents.shape[0]).to(latents.dtype) noise_pred = components.transformer( hidden_states=latent_model_input, @@ -222,7 +242,8 @@ def __call__( noise_pred = noise_pred[:, : latents.size(1)] block_state.noise_pred = noise_pred - return components, block_state + self.set_block_state(state, block_state) + return components, state # support CFG for Flux2-Klein base model @@ -251,8 +272,9 @@ def expected_configs(self) -> list[ConfigSpec]: def description(self) -> str: return ( "Step within the denoising loop that denoises the latents for Flux2. " - "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " - "object (e.g. `Flux2DenoiseLoopWrapper`)" + "This block should be used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`); it reads the current timestep `t` and step index `i` " + "from the loop scope." ) @property @@ -305,12 +327,34 @@ def inputs(self) -> list[tuple[str, Any]]: type_hint=torch.Tensor, description="4D position IDs for latent tokens (T, H, W, L)", ), + InputParam( + "num_inference_steps", + required=True, + type_hint=int, + description="The number of inference steps, used to set the guider state.", + ), + InputParam( + "t", + required=True, + type_hint=torch.Tensor, + description="The current timestep, provided by the denoise loop scope.", + ), + InputParam( + "i", + required=True, + type_hint=int, + description="The current step index, provided by the denoise loop scope.", + ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted noise for this step")] + @torch.no_grad() - def __call__( - self, components: Flux2KleinModularPipeline, block_state: BlockState, i: int, t: torch.Tensor - ) -> PipelineState: + def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + latents = block_state.latents latent_model_input = latents.to(components.transformer.dtype) img_ids = block_state.latent_ids @@ -321,6 +365,7 @@ def __call__( image_latent_ids = block_state.image_latent_ids img_ids = torch.cat([img_ids, image_latent_ids], dim=1) + t = block_state.t timestep = t.expand(latents.shape[0]).to(latents.dtype) guider_inputs = { @@ -334,7 +379,9 @@ def __call__( ), } - components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + components.guider.set_state( + step=block_state.i, num_inference_steps=block_state.num_inference_steps, timestep=t + ) guider_state = components.guider.prepare_inputs(guider_inputs) for guider_state_batch in guider_state: @@ -356,7 +403,8 @@ def __call__( # perform guidance block_state.noise_pred = components.guider(guider_state)[0] - return components, block_state + self.set_block_state(state, block_state) + return components, state class Flux2LoopAfterDenoiser(ModularPipelineBlocks): @@ -370,28 +418,46 @@ def expected_components(self) -> list[ComponentSpec]: def description(self) -> str: return ( "Step within the denoising loop that updates the latents after denoising. " - "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " - "object (e.g. `Flux2DenoiseLoopWrapper`)" + "This block should be used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`); it reads `noise_pred` and the current timestep `t` " + "from the loop scope." ) @property def inputs(self) -> list[tuple[str, Any]]: - return [] - - @property - def intermediate_inputs(self) -> list[str]: - return [InputParam("generator")] + return [ + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The latents to update. Shape: (B, seq_len, C)", + ), + InputParam( + "noise_pred", + required=True, + type_hint=torch.Tensor, + description="The predicted noise for this step.", + ), + InputParam( + "t", + required=True, + type_hint=torch.Tensor, + description="The current timestep, provided by the denoise loop scope.", + ), + ] @property def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] @torch.no_grad() - def __call__(self, components: Flux2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components: Flux2ModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + latents_dtype = block_state.latents.dtype block_state.latents = components.scheduler.step( block_state.noise_pred, - t, + block_state.t, block_state.latents, return_dict=False, )[0] @@ -400,12 +466,22 @@ def __call__(self, components: Flux2ModularPipeline, block_state: BlockState, i: if torch.backends.mps.is_available(): block_state.latents = block_state.latents.to(latents_dtype) - return components, block_state + self.set_block_state(state, block_state) + return components, state -class Flux2DenoiseLoopWrapper(LoopSequentialPipelineBlocks): +class Flux2DenoiseLoopWrapper(IterativePipelineBlocks): model_name = "flux2" + @property + def loop_locals(self) -> list[str]: + return ["i", "t"] + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + # the loop logic itself reads `scheduler.order` for the warmup-step computation + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + @property def description(self) -> str: return ( @@ -413,13 +489,6 @@ def description(self) -> str: "The specific steps within each iteration can be customized with `sub_blocks` attribute" ) - @property - def loop_expected_components(self) -> list[ComponentSpec]: - return [ - ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), - ComponentSpec("transformer", Flux2Transformer2DModel), - ] - @property def loop_inputs(self) -> list[InputParam]: return [ @@ -440,24 +509,25 @@ def loop_inputs(self) -> list[InputParam]: @torch.no_grad() def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - - block_state.num_warmup_steps = max( + num_warmup_steps = max( len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 ) - with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: - for i, t in enumerate(block_state.timesteps): - components, block_state = self.loop_step(components, block_state, i=i, t=t) + with state.loop_scope(): + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + state.set_local("i", i) + state.set_local("t", t) + components, state = self.loop_step(components, state) - if i == len(block_state.timesteps) - 1 or ( - (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 - ): - progress_bar.update() + if i == len(block_state.timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() - if XLA_AVAILABLE: - xm.mark_step() + if XLA_AVAILABLE: + xm.mark_step() - self.set_block_state(state, block_state) return components, state @@ -469,7 +539,7 @@ class Flux2DenoiseStep(Flux2DenoiseLoopWrapper): def description(self) -> str: return ( "Denoise step that iteratively denoises the latents for Flux2. \n" - "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n" + "Its loop logic is defined in `IterativePipelineBlocks.__call__` method \n" "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" " - `Flux2LoopDenoiser`\n" " - `Flux2LoopAfterDenoiser`\n" @@ -485,7 +555,7 @@ class Flux2KleinDenoiseStep(Flux2DenoiseLoopWrapper): def description(self) -> str: return ( "Denoise step that iteratively denoises the latents for Flux2. \n" - "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n" + "Its loop logic is defined in `IterativePipelineBlocks.__call__` method \n" "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" " - `Flux2KleinLoopDenoiser`\n" " - `Flux2LoopAfterDenoiser`\n" @@ -501,7 +571,7 @@ class Flux2KleinBaseDenoiseStep(Flux2DenoiseLoopWrapper): def description(self) -> str: return ( "Denoise step that iteratively denoises the latents for Flux2. \n" - "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n" + "Its loop logic is defined in `IterativePipelineBlocks.__call__` method \n" "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" " - `Flux2KleinBaseLoopDenoiser`\n" " - `Flux2LoopAfterDenoiser`\n" diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index d43825860d8e..c20ac2746ae2 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -18,6 +18,7 @@ import traceback import warnings from collections import OrderedDict +from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass, field from typing import Any @@ -151,6 +152,32 @@ class PipelineState: values: dict[str, Any] = field(default_factory=dict) kwargs_mapping: dict[str, list[str]] = field(default_factory=dict) + # stack of loop-local namespaces managed by `IterativePipelineBlocks`; values in the active scopes are visible + # on every block_state without being declared and are discarded when the loop that created them exits + scopes: list[dict[str, Any]] = field(default_factory=list) + + @contextmanager + def loop_scope(self): + """Context manager for a loop-local scope: values set with `set_local` (e.g. the current timestep) resolve + blocks' declared inputs while the scope is active and are discarded when it exits.""" + self.scopes.append({}) + try: + yield self + finally: + self.scopes.pop() + + def set_local(self, key: str, value: Any): + """Set a value in the innermost loop-local scope (requires an active scope).""" + if not self.scopes: + raise RuntimeError(f"set_local('{key}') called with no active scope; use set() instead.") + self.scopes[-1][key] = value + + def local_values(self) -> dict[str, Any]: + """All values visible from the active scopes, innermost scope winning on name collisions.""" + merged = {} + for scope in self.scopes: + merged.update(scope) + return merged def set(self, key: str, value: Any, kwargs_type: str = None): """ @@ -497,11 +524,15 @@ def get_block_state(self, state: PipelineState) -> dict: """Get all inputs and intermediates in one dictionary""" data = {} state_inputs = self.inputs + local_values = state.local_values() # Check inputs for input_param in state_inputs: if input_param.name: - value = state.get(input_param.name) + if input_param.name in local_values: + value = local_values[input_param.name] + else: + value = state.get(input_param.name) if input_param.required and value is None: raise ValueError(f"Required input '{input_param.name}' is missing") elif value is not None or (value is None and input_param.name not in data): @@ -521,6 +552,8 @@ def get_block_state(self, state: PipelineState) -> dict: return BlockState(**data) def set_block_state(self, state: PipelineState, block_state: BlockState): + local_values = state.local_values() + for output_param in self.intermediate_outputs: if not hasattr(block_state, output_param.name): raise ValueError(f"Intermediate output '{output_param.name}' is missing in block state") @@ -530,6 +563,11 @@ def set_block_state(self, state: PipelineState, block_state: BlockState): for input_param in self.inputs: if input_param.name and hasattr(block_state, input_param.name): param = getattr(block_state, input_param.name) + if input_param.name in local_values: + # the value was read from a loop-local scope; keep updates loop-local + if local_values[input_param.name] is not param: + state.set_local(input_param.name, param) + continue # Only add if the value is different from what's in the state current_value = state.get(input_param.name) if current_value is not param: # Using identity comparison to check if object was modified @@ -550,6 +588,25 @@ def set_block_state(self, state: PipelineState, block_state: BlockState): if current_value is not param: # Using identity comparison to check if object was modified state.set(param_name, param, input_param.kwargs_type) + @torch.compiler.disable + def progress_bar(self, iterable=None, total=None): + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + elif not isinstance(self._progress_bar_config, dict): + raise ValueError( + f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." + ) + + if iterable is not None: + return tqdm(iterable, **self._progress_bar_config) + elif total is not None: + return tqdm(total=total, **self._progress_bar_config) + else: + raise ValueError("Either `total` or `iterable` has to be defined.") + + def set_progress_bar_config(self, **kwargs): + self._progress_bar_config = kwargs + @property def input_names(self) -> list[str]: return [input_param.name for input_param in self.inputs if input_param.name is not None] @@ -775,7 +832,7 @@ def get_execution_blocks(self, **kwargs) -> ModularPipelineBlocks | None: Get the block(s) that would execute given the inputs. Recursively resolves nested ConditionalPipelineBlocks until reaching either: - - A leaf block (no sub_blocks or LoopSequentialPipelineBlocks) → returns single `ModularPipelineBlocks` + - A leaf block (no sub_blocks, or a loop block: IterativePipelineBlocks / LoopSequentialPipelineBlocks) → returns single `ModularPipelineBlocks` - A `SequentialPipelineBlocks` → delegates to its `get_execution_blocks()` which returns a `SequentialPipelineBlocks` containing the resolved execution blocks @@ -798,7 +855,7 @@ def get_execution_blocks(self, **kwargs) -> ModularPipelineBlocks | None: block = self.sub_blocks[block_name] # Recursively resolve until we hit a leaf block - if block.sub_blocks and not isinstance(block, LoopSequentialPipelineBlocks): + if block.sub_blocks and not isinstance(block, (IterativePipelineBlocks, LoopSequentialPipelineBlocks)): return block.get_execution_blocks(**kwargs) return block @@ -1179,13 +1236,13 @@ def fn_recursive_traverse(block, block_name, active_inputs): return result_blocks # Has sub_blocks (SequentialPipelineBlocks/ConditionalPipelineBlocks) - if block.sub_blocks and not isinstance(block, LoopSequentialPipelineBlocks): + if block.sub_blocks and not isinstance(block, (IterativePipelineBlocks, LoopSequentialPipelineBlocks)): for sub_block_name, sub_block in block.sub_blocks.items(): nested_blocks = fn_recursive_traverse(sub_block, sub_block_name, active_inputs) nested_blocks = {f"{block_name}.{k}": v for k, v in nested_blocks.items()} result_blocks.update(nested_blocks) else: - # Leaf block: single ModularPipelineBlocks or LoopSequentialPipelineBlocks + # Leaf block: single ModularPipelineBlocks or a loop block (IterativePipelineBlocks / LoopSequentialPipelineBlocks) result_blocks[block_name] = block # Add outputs to active_inputs so subsequent blocks can use them as triggers if hasattr(block, "intermediate_outputs"): @@ -1295,6 +1352,106 @@ def _requirements(self) -> dict[str, str]: return requirements +class IterativePipelineBlocks(SequentialPipelineBlocks): + """ + A pipeline blocks that runs its sub-blocks multiple times. Subclasses implement `__call__` with their loop + logic — the same way leaf blocks implement `__call__` around `get_block_state` — calling `loop_step` once per + iteration inside a `state.loop_scope()`: + + ```python + @property + def loop_locals(self): + return ["i", "t"] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + with state.loop_scope(): + for i, t in enumerate(block_state.timesteps): + state.set_local("i", i) + state.set_local("t", t) + components, state = self.loop_step(components, state) + return components, state + ``` + + Unlike [`LoopSequentialPipelineBlocks`], sub-blocks are ordinary blocks operating on the full + [`PipelineState`] — leaf or assembled (`SequentialPipelineBlocks`, `ConditionalPipelineBlocks`, another + `IterativePipelineBlocks`, ...) — so loops can be nested and composed freely. + + Sub-blocks declare every input they consume, including loop variables like the current timestep: values set + with `state.set_local` resolve declared inputs while the scope is active and are discarded when it exits. + The loop block lists the names it provides through the scope in the `loop_locals` property so they are + excluded from its own aggregated `inputs`. Sub-block outputs are written to the pipeline state as usual and + persist after the loop. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + + Attributes: + block_classes: list of block classes to be used (same as `SequentialPipelineBlocks`) + block_names: list of names for each block (same as `SequentialPipelineBlocks`) + """ + + @property + def loop_inputs(self) -> list[InputParam]: + """Inputs consumed by the loop logic in `__call__` itself (e.g. `timesteps`).""" + return [] + + @property + def loop_locals(self) -> list[str]: + """Names the loop provides to its sub-blocks through the loop scope via `set_local` (e.g. `["i", "t"]`).""" + return [] + + @property + def loop_intermediate_outputs(self) -> list[OutputParam]: + """Outputs written to the pipeline state by the loop logic in `__call__` itself.""" + return [] + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + """Components used by the loop logic in `__call__` itself (e.g. the scheduler).""" + return [] + + @property + def loop_expected_configs(self) -> list[ConfigSpec]: + """Configs used by the loop logic in `__call__` itself.""" + return [] + + @property + def inputs(self) -> list[InputParam]: + inputs = [p for p in self._get_inputs() if p.name not in self.loop_locals] + names = {p.name for p in inputs} + return [p for p in self.loop_inputs if p.name not in names] + inputs + + @property + def intermediate_outputs(self) -> list[OutputParam]: + outputs = super().intermediate_outputs + names = {output.name for output in outputs} + return outputs + [output for output in self.loop_intermediate_outputs if output.name not in names] + + @property + def expected_components(self) -> list[ComponentSpec]: + expected_components = super().expected_components + for component in self.loop_expected_components: + if component not in expected_components: + expected_components.append(component) + return expected_components + + @property + def expected_configs(self) -> list[ConfigSpec]: + expected_configs = super().expected_configs + for config in self.loop_expected_configs: + if config not in expected_configs: + expected_configs.append(config) + return expected_configs + + def loop_step(self, components, state: PipelineState) -> PipelineState: + """Run all sub-blocks once over the pipeline state (one loop iteration).""" + return super().__call__(components, state) + + def __call__(self, components, state: PipelineState) -> PipelineState: + raise NotImplementedError("`__call__` method needs to be implemented by the subclass") + + class LoopSequentialPipelineBlocks(ModularPipelineBlocks): """ A Pipeline blocks that combines multiple pipeline block classes into a For Loop. When called, it will call each @@ -1569,25 +1726,6 @@ def __repr__(self): return result - @torch.compiler.disable - def progress_bar(self, iterable=None, total=None): - if not hasattr(self, "_progress_bar_config"): - self._progress_bar_config = {} - elif not isinstance(self._progress_bar_config, dict): - raise ValueError( - f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." - ) - - if iterable is not None: - return tqdm(iterable, **self._progress_bar_config) - elif total is not None: - return tqdm(total=total, **self._progress_bar_config) - else: - raise ValueError("Either `total` or `iterable` has to be defined.") - - def set_progress_bar_config(self, **kwargs): - self._progress_bar_config = kwargs - # YiYi TODO: # 1. look into the serialization of modular_model_index.json, make sure the items are properly ordered like model_index.json (currently a mess) diff --git a/tests/modular_pipelines/test_iterative_pipeline_blocks.py b/tests/modular_pipelines/test_iterative_pipeline_blocks.py new file mode 100644 index 000000000000..96934d3cb473 --- /dev/null +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -0,0 +1,243 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest +import torch + +from diffusers.modular_pipelines import ( + InputParam, + IterativePipelineBlocks, + ModularPipelineBlocks, + OutputParam, + SequentialPipelineBlocks, +) + + +# Dummy blocks modeled on the Helios chunk-loop use case: an outer autoregressive chunk loop +# (history carried across chunks) containing a full inner timestep denoising loop. + + +class ChunkNoiseGenStep(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [ + InputParam(name="history", required=True), + InputParam(name="k", required=True, description="Chunk index, provided by the chunk loop scope."), + ] + + @property + def intermediate_outputs(self): + return [OutputParam(name="chunk_latents")] + + @property + def description(self): + return "prepares this chunk's latents from the history" + + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.chunk_latents = block_state.history + block_state.k + self.set_block_state(state, block_state) + return components, state + + +class LoopDenoiserStep(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [ + InputParam(name="chunk_latents", required=True), + InputParam(name="t", required=True, description="Current timestep, provided by the denoise loop scope."), + ] + + @property + def intermediate_outputs(self): + return [OutputParam(name="noise_pred")] + + @property + def description(self): + return "predicts the noise for one timestep" + + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.noise_pred = block_state.chunk_latents * 0 + block_state.t + self.set_block_state(state, block_state) + return components, state + + +class LoopSchedulerStep(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="chunk_latents", required=True), InputParam(name="noise_pred", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="chunk_latents")] + + @property + def description(self): + return "updates the chunk latents with the noise prediction" + + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.chunk_latents = block_state.chunk_latents + block_state.noise_pred + self.set_block_state(state, block_state) + return components, state + + +class InnerDenoiseLoop(IterativePipelineBlocks): + """Inner timestep loop — itself an assembled loop block, nested inside the chunk loop.""" + + model_name = "test" + block_classes = [LoopDenoiserStep, LoopSchedulerStep] + block_names = ["denoiser", "scheduler"] + + @property + def description(self): + return "inner timestep loop" + + @property + def loop_inputs(self): + return [InputParam(name="timesteps", required=True)] + + @property + def loop_locals(self): + return ["i", "t"] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + with state.loop_scope(): + for i, t in enumerate(block_state.timesteps): + state.set_local("i", i) + state.set_local("t", t) + components, state = self.loop_step(components, state) + return components, state + + +class ChunkUpdateStep(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="chunk_latents", required=True), InputParam(name="latent_chunks", default=None)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="history"), OutputParam(name="latent_chunks")] + + @property + def description(self): + return "records the denoised chunk and updates the history" + + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.history = block_state.chunk_latents + block_state.latent_chunks = [*(block_state.latent_chunks or []), float(block_state.chunk_latents)] + self.set_block_state(state, block_state) + return components, state + + +class ChunkLoop(IterativePipelineBlocks): + """Outer chunk loop containing the inner timestep loop as a sub-block.""" + + model_name = "test" + block_classes = [ChunkNoiseGenStep, InnerDenoiseLoop, ChunkUpdateStep] + block_names = ["noise_gen", "denoise", "update"] + + @property + def description(self): + return "outer autoregressive chunk loop" + + @property + def loop_inputs(self): + return [InputParam(name="num_latent_chunk", required=True)] + + @property + def loop_locals(self): + return ["k"] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + with state.loop_scope(): + for k in range(block_state.num_latent_chunk): + state.set_local("k", k) + components, state = self.loop_step(components, state) + return components, state + + +class TestIterativePipelineBlocksStructure: + def test_loop_inputs_and_locals_aggregation(self): + loop = ChunkLoop() + input_names = [p.name for p in loop.inputs] + + # loop_inputs of the loop itself and of the nested loop are surfaced + assert "num_latent_chunk" in input_names + assert "timesteps" in input_names + # values provided through the loop scopes are not user inputs + assert "k" not in input_names + assert "i" not in input_names + assert "t" not in input_names + # cross-chunk carries surface as (optional) iteration-0 seeds + assert "history" in input_names + assert "latent_chunks" in input_names + + def test_sub_block_outputs_are_aggregated(self): + loop = ChunkLoop() + output_names = [o.name for o in loop.intermediate_outputs] + assert "history" in output_names + assert "latent_chunks" in output_names + + def test_loop_block_can_nest_assembled_blocks(self): + # the nested inner loop stays an assembled IterativePipelineBlocks sub-block + loop = ChunkLoop() + assert isinstance(loop.sub_blocks["denoise"], IterativePipelineBlocks) + assert list(loop.sub_blocks["denoise"].sub_blocks) == ["denoiser", "scheduler"] + + +class TestIterativePipelineBlocksExecution: + def _make_pipeline(self): + return SequentialPipelineBlocks.from_blocks_dict({"chunks": ChunkLoop()}).init_pipeline() + + def test_nested_chunk_loop(self): + pipe = self._make_pipeline() + # per chunk: chunk_latents = history + k, then += t for every timestep (1.0 + 2.0), + # then history <- chunk_latents + # chunk 0: 0 + 0 + 3 = 3 ; chunk 1: 3 + 1 + 3 = 7 ; chunk 2: 7 + 2 + 3 = 12 + state = pipe(num_latent_chunk=3, timesteps=torch.tensor([1.0, 2.0]), history=torch.tensor(0.0)) + + assert state.get("latent_chunks") == [3.0, 7.0, 12.0] + # the cross-chunk carry persists as a declared output + assert float(state.get("history")) == 12.0 + + def test_loop_locals_do_not_leak_into_state(self): + pipe = self._make_pipeline() + state = pipe(num_latent_chunk=2, timesteps=torch.tensor([1.0]), history=torch.tensor(0.0)) + + for name in ("k", "i", "t"): + assert state.get(name) is None + # declared sub-block outputs persist after the loop (last iteration's value) + assert state.get("noise_pred") is not None + + def test_loop_sub_block_standalone_requires_loop_locals(self): + # outside a loop scope, a block that declares a loop-provided input fails with a clear error + pipe = SequentialPipelineBlocks.from_blocks_dict({"denoiser": LoopDenoiserStep()}).init_pipeline() + with pytest.raises(ValueError, match="Required input 't' is missing"): + pipe(chunk_latents=torch.tensor(1.0)) From 2e92d11500f57c654b4eed835eaca9210de2ecfc Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 16:06:57 +0000 Subject: [PATCH 02/10] Pass loop variables as call arguments instead of state scopes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Loop variables (i, t, k, ...) now ride the call signature: leaf sub-blocks of an IterativePipelineBlocks accept them after (components, state), the loop declares their names in `loop_variables`, and `loop_step` validates every leaf's signature against it before the first iteration. Assembled sub-blocks (nested loops, sequential/conditional groups) are called with the regular (components, state) interface and pass their own loop variables to their own sub-blocks. This removes the PipelineState scope machinery entirely — PipelineState and get/set_block_state are unchanged from main — and loop sub-blocks keep the familiar LoopSequentialPipelineBlocks authoring style, now with full composability. Co-Authored-By: Claude Fable 5 --- .../modular_pipelines/flux2/denoise.py | 78 ++++------- .../modular_pipelines/modular_pipeline.py | 123 +++++++++--------- .../test_iterative_pipeline_blocks.py | 106 +++++++++------ 3 files changed, 151 insertions(+), 156 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux2/denoise.py b/src/diffusers/modular_pipelines/flux2/denoise.py index f455223dde86..57d8b63f2d59 100644 --- a/src/diffusers/modular_pipelines/flux2/denoise.py +++ b/src/diffusers/modular_pipelines/flux2/denoise.py @@ -100,12 +100,6 @@ def inputs(self) -> list[tuple[str, Any]]: type_hint=torch.Tensor, description="4D position IDs for latent tokens (T, H, W, L)", ), - InputParam( - "t", - required=True, - type_hint=torch.Tensor, - description="The current timestep, provided by the denoise loop scope.", - ), ] @property @@ -113,7 +107,9 @@ def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted noise for this step")] @torch.no_grad() - def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + def __call__( + self, components: Flux2ModularPipeline, state: PipelineState, i: int, t: torch.Tensor + ) -> PipelineState: block_state = self.get_block_state(state) latents = block_state.latents @@ -126,7 +122,7 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi image_latent_ids = block_state.image_latent_ids img_ids = torch.cat([img_ids, image_latent_ids], dim=1) - timestep = block_state.t.expand(latents.shape[0]).to(latents.dtype) + timestep = t.expand(latents.shape[0]).to(latents.dtype) noise_pred = components.transformer( hidden_states=latent_model_input, @@ -200,12 +196,6 @@ def inputs(self) -> list[tuple[str, Any]]: type_hint=torch.Tensor, description="4D position IDs for latent tokens (T, H, W, L)", ), - InputParam( - "t", - required=True, - type_hint=torch.Tensor, - description="The current timestep, provided by the denoise loop scope.", - ), ] @property @@ -213,7 +203,9 @@ def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted noise for this step")] @torch.no_grad() - def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState: + def __call__( + self, components: Flux2KleinModularPipeline, state: PipelineState, i: int, t: torch.Tensor + ) -> PipelineState: block_state = self.get_block_state(state) latents = block_state.latents @@ -226,7 +218,7 @@ def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) image_latent_ids = block_state.image_latent_ids img_ids = torch.cat([img_ids, image_latent_ids], dim=1) - timestep = block_state.t.expand(latents.shape[0]).to(latents.dtype) + timestep = t.expand(latents.shape[0]).to(latents.dtype) noise_pred = components.transformer( hidden_states=latent_model_input, @@ -333,18 +325,6 @@ def inputs(self) -> list[tuple[str, Any]]: type_hint=int, description="The number of inference steps, used to set the guider state.", ), - InputParam( - "t", - required=True, - type_hint=torch.Tensor, - description="The current timestep, provided by the denoise loop scope.", - ), - InputParam( - "i", - required=True, - type_hint=int, - description="The current step index, provided by the denoise loop scope.", - ), ] @property @@ -352,7 +332,9 @@ def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted noise for this step")] @torch.no_grad() - def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState: + def __call__( + self, components: Flux2KleinModularPipeline, state: PipelineState, i: int, t: torch.Tensor + ) -> PipelineState: block_state = self.get_block_state(state) latents = block_state.latents @@ -365,7 +347,6 @@ def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) image_latent_ids = block_state.image_latent_ids img_ids = torch.cat([img_ids, image_latent_ids], dim=1) - t = block_state.t timestep = t.expand(latents.shape[0]).to(latents.dtype) guider_inputs = { @@ -379,9 +360,7 @@ def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) ), } - components.guider.set_state( - step=block_state.i, num_inference_steps=block_state.num_inference_steps, timestep=t - ) + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) guider_state = components.guider.prepare_inputs(guider_inputs) for guider_state_batch in guider_state: @@ -438,12 +417,6 @@ def inputs(self) -> list[tuple[str, Any]]: type_hint=torch.Tensor, description="The predicted noise for this step.", ), - InputParam( - "t", - required=True, - type_hint=torch.Tensor, - description="The current timestep, provided by the denoise loop scope.", - ), ] @property @@ -451,13 +424,13 @@ def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] @torch.no_grad() - def __call__(self, components: Flux2ModularPipeline, state: PipelineState): + def __call__(self, components: Flux2ModularPipeline, state: PipelineState, i: int, t: torch.Tensor): block_state = self.get_block_state(state) latents_dtype = block_state.latents.dtype block_state.latents = components.scheduler.step( block_state.noise_pred, - block_state.t, + t, block_state.latents, return_dict=False, )[0] @@ -474,7 +447,7 @@ class Flux2DenoiseLoopWrapper(IterativePipelineBlocks): model_name = "flux2" @property - def loop_locals(self) -> list[str]: + def loop_variables(self) -> list[str]: return ["i", "t"] @property @@ -513,20 +486,17 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 ) - with state.loop_scope(): - with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: - for i, t in enumerate(block_state.timesteps): - state.set_local("i", i) - state.set_local("t", t) - components, state = self.loop_step(components, state) + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) - if i == len(block_state.timesteps) - 1 or ( - (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0 - ): - progress_bar.update() + if i == len(block_state.timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() - if XLA_AVAILABLE: - xm.mark_step() + if XLA_AVAILABLE: + xm.mark_step() return components, state diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index c20ac2746ae2..9b52ad3f8416 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -18,7 +18,6 @@ import traceback import warnings from collections import OrderedDict -from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass, field from typing import Any @@ -152,32 +151,6 @@ class PipelineState: values: dict[str, Any] = field(default_factory=dict) kwargs_mapping: dict[str, list[str]] = field(default_factory=dict) - # stack of loop-local namespaces managed by `IterativePipelineBlocks`; values in the active scopes are visible - # on every block_state without being declared and are discarded when the loop that created them exits - scopes: list[dict[str, Any]] = field(default_factory=list) - - @contextmanager - def loop_scope(self): - """Context manager for a loop-local scope: values set with `set_local` (e.g. the current timestep) resolve - blocks' declared inputs while the scope is active and are discarded when it exits.""" - self.scopes.append({}) - try: - yield self - finally: - self.scopes.pop() - - def set_local(self, key: str, value: Any): - """Set a value in the innermost loop-local scope (requires an active scope).""" - if not self.scopes: - raise RuntimeError(f"set_local('{key}') called with no active scope; use set() instead.") - self.scopes[-1][key] = value - - def local_values(self) -> dict[str, Any]: - """All values visible from the active scopes, innermost scope winning on name collisions.""" - merged = {} - for scope in self.scopes: - merged.update(scope) - return merged def set(self, key: str, value: Any, kwargs_type: str = None): """ @@ -524,15 +497,11 @@ def get_block_state(self, state: PipelineState) -> dict: """Get all inputs and intermediates in one dictionary""" data = {} state_inputs = self.inputs - local_values = state.local_values() # Check inputs for input_param in state_inputs: if input_param.name: - if input_param.name in local_values: - value = local_values[input_param.name] - else: - value = state.get(input_param.name) + value = state.get(input_param.name) if input_param.required and value is None: raise ValueError(f"Required input '{input_param.name}' is missing") elif value is not None or (value is None and input_param.name not in data): @@ -552,8 +521,6 @@ def get_block_state(self, state: PipelineState) -> dict: return BlockState(**data) def set_block_state(self, state: PipelineState, block_state: BlockState): - local_values = state.local_values() - for output_param in self.intermediate_outputs: if not hasattr(block_state, output_param.name): raise ValueError(f"Intermediate output '{output_param.name}' is missing in block state") @@ -563,11 +530,6 @@ def set_block_state(self, state: PipelineState, block_state: BlockState): for input_param in self.inputs: if input_param.name and hasattr(block_state, input_param.name): param = getattr(block_state, input_param.name) - if input_param.name in local_values: - # the value was read from a loop-local scope; keep updates loop-local - if local_values[input_param.name] is not param: - state.set_local(input_param.name, param) - continue # Only add if the value is different from what's in the state current_value = state.get(input_param.name) if current_value is not param: # Using identity comparison to check if object was modified @@ -1354,35 +1316,33 @@ def _requirements(self) -> dict[str, str]: class IterativePipelineBlocks(SequentialPipelineBlocks): """ - A pipeline blocks that runs its sub-blocks multiple times. Subclasses implement `__call__` with their loop - logic — the same way leaf blocks implement `__call__` around `get_block_state` — calling `loop_step` once per - iteration inside a `state.loop_scope()`: + A pipeline blocks that runs its sub-blocks multiple times. Subclasses declare their loop-variable names in + `loop_variables` and implement `__call__` with their loop logic — the same way leaf blocks implement + `__call__` around `get_block_state` — calling `loop_step` once per iteration with the loop variables: ```python @property - def loop_locals(self): + def loop_variables(self): return ["i", "t"] @torch.no_grad() def __call__(self, components, state): block_state = self.get_block_state(state) - with state.loop_scope(): - for i, t in enumerate(block_state.timesteps): - state.set_local("i", i) - state.set_local("t", t) - components, state = self.loop_step(components, state) + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) return components, state ``` - Unlike [`LoopSequentialPipelineBlocks`], sub-blocks are ordinary blocks operating on the full - [`PipelineState`] — leaf or assembled (`SequentialPipelineBlocks`, `ConditionalPipelineBlocks`, another - `IterativePipelineBlocks`, ...) — so loops can be nested and composed freely. + Unlike [`LoopSequentialPipelineBlocks`], sub-blocks operate on the full [`PipelineState`] with the regular + `get_block_state`/`set_block_state` behavior, and can be leaf or assembled blocks + (`SequentialPipelineBlocks`, `ConditionalPipelineBlocks`, another `IterativePipelineBlocks`, ...) — so loops + can be nested and composed freely. - Sub-blocks declare every input they consume, including loop variables like the current timestep: values set - with `state.set_local` resolve declared inputs while the scope is active and are discarded when it exits. - The loop block lists the names it provides through the scope in the `loop_locals` property so they are - excluded from its own aggregated `inputs`. Sub-block outputs are written to the pipeline state as usual and - persist after the loop. + Loop variables are passed to leaf sub-blocks as call arguments: every leaf sub-block must have the signature + `__call__(self, components, state, )`, which is validated against `loop_variables` before + the first iteration. Assembled sub-blocks are called with the regular `(components, state)` interface and do + not receive the loop variables — a nested loop passes its own `loop_variables` to its own sub-blocks. + Sub-block outputs are written to the pipeline state as usual and persist after the loop. > [!WARNING] > This is an experimental feature and is likely to change in the future. @@ -1392,13 +1352,13 @@ def __call__(self, components, state): """ @property - def loop_inputs(self) -> list[InputParam]: - """Inputs consumed by the loop logic in `__call__` itself (e.g. `timesteps`).""" + def loop_variables(self) -> list[str]: + """Names of the loop variables `loop_step` passes to leaf sub-blocks each iteration (e.g. `["i", "t"]`).""" return [] @property - def loop_locals(self) -> list[str]: - """Names the loop provides to its sub-blocks through the loop scope via `set_local` (e.g. `["i", "t"]`).""" + def loop_inputs(self) -> list[InputParam]: + """Inputs consumed by the loop logic in `__call__` itself (e.g. `timesteps`).""" return [] @property @@ -1418,7 +1378,7 @@ def loop_expected_configs(self) -> list[ConfigSpec]: @property def inputs(self) -> list[InputParam]: - inputs = [p for p in self._get_inputs() if p.name not in self.loop_locals] + inputs = self._get_inputs() names = {p.name for p in inputs} return [p for p in self.loop_inputs if p.name not in names] + inputs @@ -1444,9 +1404,44 @@ def expected_configs(self) -> list[ConfigSpec]: expected_configs.append(config) return expected_configs - def loop_step(self, components, state: PipelineState) -> PipelineState: - """Run all sub-blocks once over the pipeline state (one loop iteration).""" - return super().__call__(components, state) + def _validate_loop_step_signatures(self): + """Every leaf sub-block must accept exactly the loop variables after `(components, state)`.""" + expected = set(self.loop_variables) + for block_name, block in self.sub_blocks.items(): + if block.sub_blocks: + # assembled sub-blocks are called with the regular (components, state) interface + continue + params = list(inspect.signature(block.__call__).parameters) + extra = set(params[2:]) + if extra != expected: + raise ValueError( + f"Loop sub-block '{block_name}' ({block.__class__.__name__}) of {self.__class__.__name__} " + f"must accept the loop variables {sorted(expected)} after `(components, state)`; " + f"its `__call__` accepts {sorted(extra)}." + ) + + def loop_step(self, components, state: PipelineState, **loop_kwargs) -> PipelineState: + """Run all sub-blocks once over the pipeline state (one loop iteration), passing the loop variables to + leaf sub-blocks.""" + if not getattr(self, "_loop_signatures_validated", False): + self._validate_loop_step_signatures() + self._loop_signatures_validated = True + + for block_name, block in self.sub_blocks.items(): + try: + if block.sub_blocks: + components, state = block(components, state) + else: + components, state = block(components, state, **loop_kwargs) + except Exception as e: + error_msg = ( + f"\nError in block: ({block_name}, {block.__class__.__name__})\n" + f"Error details: {str(e)}\n" + f"Traceback:\n{traceback.format_exc()}" + ) + logger.error(error_msg) + raise + return components, state def __call__(self, components, state: PipelineState) -> PipelineState: raise NotImplementedError("`__call__` method needs to be implemented by the subclass") diff --git a/tests/modular_pipelines/test_iterative_pipeline_blocks.py b/tests/modular_pipelines/test_iterative_pipeline_blocks.py index 96934d3cb473..c306eda17c75 100644 --- a/tests/modular_pipelines/test_iterative_pipeline_blocks.py +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -26,7 +26,9 @@ # Dummy blocks modeled on the Helios chunk-loop use case: an outer autoregressive chunk loop -# (history carried across chunks) containing a full inner timestep denoising loop. +# (history carried across chunks) containing a full inner timestep denoising loop. Loop variables +# (`k` for the chunk loop, `i`/`t` for the timestep loop) are passed to leaf sub-blocks as call +# arguments; every leaf sub-block of a loop must accept its loop's variables. class ChunkNoiseGenStep(ModularPipelineBlocks): @@ -34,10 +36,7 @@ class ChunkNoiseGenStep(ModularPipelineBlocks): @property def inputs(self): - return [ - InputParam(name="history", required=True), - InputParam(name="k", required=True, description="Chunk index, provided by the chunk loop scope."), - ] + return [InputParam(name="history", required=True)] @property def intermediate_outputs(self): @@ -47,9 +46,9 @@ def intermediate_outputs(self): def description(self): return "prepares this chunk's latents from the history" - def __call__(self, components, state): + def __call__(self, components, state, k): block_state = self.get_block_state(state) - block_state.chunk_latents = block_state.history + block_state.k + block_state.chunk_latents = block_state.history + k self.set_block_state(state, block_state) return components, state @@ -59,10 +58,7 @@ class LoopDenoiserStep(ModularPipelineBlocks): @property def inputs(self): - return [ - InputParam(name="chunk_latents", required=True), - InputParam(name="t", required=True, description="Current timestep, provided by the denoise loop scope."), - ] + return [InputParam(name="chunk_latents", required=True)] @property def intermediate_outputs(self): @@ -72,9 +68,9 @@ def intermediate_outputs(self): def description(self): return "predicts the noise for one timestep" - def __call__(self, components, state): + def __call__(self, components, state, i, t): block_state = self.get_block_state(state) - block_state.noise_pred = block_state.chunk_latents * 0 + block_state.t + block_state.noise_pred = block_state.chunk_latents * 0 + t self.set_block_state(state, block_state) return components, state @@ -94,7 +90,7 @@ def intermediate_outputs(self): def description(self): return "updates the chunk latents with the noise prediction" - def __call__(self, components, state): + def __call__(self, components, state, i, t): block_state = self.get_block_state(state) block_state.chunk_latents = block_state.chunk_latents + block_state.noise_pred self.set_block_state(state, block_state) @@ -113,21 +109,18 @@ def description(self): return "inner timestep loop" @property - def loop_inputs(self): - return [InputParam(name="timesteps", required=True)] + def loop_variables(self): + return ["i", "t"] @property - def loop_locals(self): - return ["i", "t"] + def loop_inputs(self): + return [InputParam(name="timesteps", required=True)] @torch.no_grad() def __call__(self, components, state): block_state = self.get_block_state(state) - with state.loop_scope(): - for i, t in enumerate(block_state.timesteps): - state.set_local("i", i) - state.set_local("t", t) - components, state = self.loop_step(components, state) + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) return components, state @@ -146,7 +139,7 @@ def intermediate_outputs(self): def description(self): return "records the denoised chunk and updates the history" - def __call__(self, components, state): + def __call__(self, components, state, k): block_state = self.get_block_state(state) block_state.history = block_state.chunk_latents block_state.latent_chunks = [*(block_state.latent_chunks or []), float(block_state.chunk_latents)] @@ -166,32 +159,30 @@ def description(self): return "outer autoregressive chunk loop" @property - def loop_inputs(self): - return [InputParam(name="num_latent_chunk", required=True)] + def loop_variables(self): + return ["k"] @property - def loop_locals(self): - return ["k"] + def loop_inputs(self): + return [InputParam(name="num_latent_chunk", required=True)] @torch.no_grad() def __call__(self, components, state): block_state = self.get_block_state(state) - with state.loop_scope(): - for k in range(block_state.num_latent_chunk): - state.set_local("k", k) - components, state = self.loop_step(components, state) + for k in range(block_state.num_latent_chunk): + components, state = self.loop_step(components, state, k=k) return components, state class TestIterativePipelineBlocksStructure: - def test_loop_inputs_and_locals_aggregation(self): + def test_loop_inputs_aggregation(self): loop = ChunkLoop() input_names = [p.name for p in loop.inputs] # loop_inputs of the loop itself and of the nested loop are surfaced assert "num_latent_chunk" in input_names assert "timesteps" in input_names - # values provided through the loop scopes are not user inputs + # loop variables are call arguments, not inputs assert "k" not in input_names assert "i" not in input_names assert "t" not in input_names @@ -227,7 +218,7 @@ def test_nested_chunk_loop(self): # the cross-chunk carry persists as a declared output assert float(state.get("history")) == 12.0 - def test_loop_locals_do_not_leak_into_state(self): + def test_loop_variables_do_not_leak_into_state(self): pipe = self._make_pipeline() state = pipe(num_latent_chunk=2, timesteps=torch.tensor([1.0]), history=torch.tensor(0.0)) @@ -236,8 +227,47 @@ def test_loop_locals_do_not_leak_into_state(self): # declared sub-block outputs persist after the loop (last iteration's value) assert state.get("noise_pred") is not None - def test_loop_sub_block_standalone_requires_loop_locals(self): - # outside a loop scope, a block that declares a loop-provided input fails with a clear error + def test_leaf_signature_is_validated(self): + class PlainStep(ModularPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "regular block without the loop variables" + + def __call__(self, components, state): + return components, state + + class BadLoop(IterativePipelineBlocks): + model_name = "test" + block_classes = [PlainStep] + block_names = ["plain"] + + @property + def description(self): + return "loop with a mismatched leaf signature" + + @property + def loop_variables(self): + return ["i", "t"] + + @property + def loop_inputs(self): + return [InputParam(name="timesteps", required=True)] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + return components, state + + pipe = SequentialPipelineBlocks.from_blocks_dict({"loop": BadLoop()}).init_pipeline() + with pytest.raises(ValueError, match="must accept the loop variables"): + pipe(timesteps=torch.tensor([1.0])) + + def test_loop_leaf_standalone_raises(self): + # outside a loop, a leaf block with loop variables in its signature cannot run pipe = SequentialPipelineBlocks.from_blocks_dict({"denoiser": LoopDenoiserStep()}).init_pipeline() - with pytest.raises(ValueError, match="Required input 't' is missing"): + with pytest.raises(TypeError): pipe(chunk_latents=torch.tensor(1.0)) From 6f529775ba977ba4951567ebf6cb3bf697351308 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 16:15:03 +0000 Subject: [PATCH 03/10] Enforce uniform loop-variable signatures for all loop sub-blocks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove the composite exemption in loop_step: every sub-block of an IterativePipelineBlocks — including a nested loop — must accept the loop variables after (components, state), validated before the first iteration. A nested loop accepts the outer variables in its hand-written __call__ (ignoring or forwarding them) and passes its own loop_variables to its own sub-blocks. Plain Sequential/Conditional groups are not supported as loop sub-blocks (flatten instead). Co-Authored-By: Claude Fable 5 --- .../modular_pipelines/modular_pipeline.py | 26 +++++++------------ .../test_iterative_pipeline_blocks.py | 8 ++++-- 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 9b52ad3f8416..fe8ef4ec65fe 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -1334,15 +1334,14 @@ def __call__(self, components, state): ``` Unlike [`LoopSequentialPipelineBlocks`], sub-blocks operate on the full [`PipelineState`] with the regular - `get_block_state`/`set_block_state` behavior, and can be leaf or assembled blocks - (`SequentialPipelineBlocks`, `ConditionalPipelineBlocks`, another `IterativePipelineBlocks`, ...) — so loops - can be nested and composed freely. + `get_block_state`/`set_block_state` behavior, so an `IterativePipelineBlocks` can itself be a sub-block of + another one — loops can be nested and composed freely. - Loop variables are passed to leaf sub-blocks as call arguments: every leaf sub-block must have the signature + Loop variables are passed to sub-blocks as call arguments: every sub-block must have the signature `__call__(self, components, state, )`, which is validated against `loop_variables` before - the first iteration. Assembled sub-blocks are called with the regular `(components, state)` interface and do - not receive the loop variables — a nested loop passes its own `loop_variables` to its own sub-blocks. - Sub-block outputs are written to the pipeline state as usual and persist after the loop. + the first iteration. A nested loop accepts the outer loop's variables in its own hand-written `__call__` + (ignoring or forwarding them) and passes its own `loop_variables` to its own sub-blocks. Sub-block outputs + are written to the pipeline state as usual and persist after the loop. > [!WARNING] > This is an experimental feature and is likely to change in the future. @@ -1405,12 +1404,9 @@ def expected_configs(self) -> list[ConfigSpec]: return expected_configs def _validate_loop_step_signatures(self): - """Every leaf sub-block must accept exactly the loop variables after `(components, state)`.""" + """Every sub-block must accept exactly the loop variables after `(components, state)`.""" expected = set(self.loop_variables) for block_name, block in self.sub_blocks.items(): - if block.sub_blocks: - # assembled sub-blocks are called with the regular (components, state) interface - continue params = list(inspect.signature(block.__call__).parameters) extra = set(params[2:]) if extra != expected: @@ -1421,18 +1417,14 @@ def _validate_loop_step_signatures(self): ) def loop_step(self, components, state: PipelineState, **loop_kwargs) -> PipelineState: - """Run all sub-blocks once over the pipeline state (one loop iteration), passing the loop variables to - leaf sub-blocks.""" + """Run all sub-blocks once over the pipeline state (one loop iteration), passing the loop variables.""" if not getattr(self, "_loop_signatures_validated", False): self._validate_loop_step_signatures() self._loop_signatures_validated = True for block_name, block in self.sub_blocks.items(): try: - if block.sub_blocks: - components, state = block(components, state) - else: - components, state = block(components, state, **loop_kwargs) + components, state = block(components, state, **loop_kwargs) except Exception as e: error_msg = ( f"\nError in block: ({block_name}, {block.__class__.__name__})\n" diff --git a/tests/modular_pipelines/test_iterative_pipeline_blocks.py b/tests/modular_pipelines/test_iterative_pipeline_blocks.py index c306eda17c75..9bff50138623 100644 --- a/tests/modular_pipelines/test_iterative_pipeline_blocks.py +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -98,7 +98,11 @@ def __call__(self, components, state, i, t): class InnerDenoiseLoop(IterativePipelineBlocks): - """Inner timestep loop — itself an assembled loop block, nested inside the chunk loop.""" + """Inner timestep loop — itself an assembled loop block, nested inside the chunk loop. + + Like every sub-block of the chunk loop, it accepts the outer loop variable `k` (and ignores it); + its own sub-blocks accept its own loop variables `i` / `t` instead. + """ model_name = "test" block_classes = [LoopDenoiserStep, LoopSchedulerStep] @@ -117,7 +121,7 @@ def loop_inputs(self): return [InputParam(name="timesteps", required=True)] @torch.no_grad() - def __call__(self, components, state): + def __call__(self, components, state, k): block_state = self.get_block_state(state) for i, t in enumerate(block_state.timesteps): components, state = self.loop_step(components, state, i=i, t=t) From 7d697b2be1fd18ceb5377000a9acbfbc7cd45859 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 16:29:34 +0000 Subject: [PATCH 04/10] Document nested-loop __call__ signature on IterativePipelineBlocks The abstract __call__ placeholder now accepts **kwargs and the docstring shows the nested case: a loop nested inside another accepts the outer loop's variables in its hand-written __call__. Co-Authored-By: Claude Fable 5 --- .../modular_pipelines/modular_pipeline.py | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index fe8ef4ec65fe..dc2b12e9eb45 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -1340,8 +1340,23 @@ def __call__(self, components, state): Loop variables are passed to sub-blocks as call arguments: every sub-block must have the signature `__call__(self, components, state, )`, which is validated against `loop_variables` before the first iteration. A nested loop accepts the outer loop's variables in its own hand-written `__call__` - (ignoring or forwarding them) and passes its own `loop_variables` to its own sub-blocks. Sub-block outputs - are written to the pipeline state as usual and persist after the loop. + (ignoring or forwarding them) and passes its own `loop_variables` to its own sub-blocks: + + ```python + class InnerDenoiseLoop(IterativePipelineBlocks): + @property + def loop_variables(self): + return ["i", "t"] # what it passes to ITS sub-blocks + + @torch.no_grad() + def __call__(self, components, state, k): # accepts the OUTER chunk loop's variable + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + return components, state + ``` + + Sub-block outputs are written to the pipeline state as usual and persist after the loop. > [!WARNING] > This is an experimental feature and is likely to change in the future. @@ -1435,7 +1450,10 @@ def loop_step(self, components, state: PipelineState, **loop_kwargs) -> Pipeline raise return components, state - def __call__(self, components, state: PipelineState) -> PipelineState: + def __call__(self, components, state: PipelineState, **kwargs) -> PipelineState: + # Subclasses implement their loop logic here. When the loop is nested inside another + # IterativePipelineBlocks, the signature must also accept the outer loop's variables, + # e.g. `def __call__(self, components, state, k)`. raise NotImplementedError("`__call__` method needs to be implemented by the subclass") From 2f3f3f54ca11a78e3ca7d61fa6199210cc1754ec Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 17:33:12 +0000 Subject: [PATCH 05/10] Add ModularLoopPipelineBlocks base and __call__ contracts ModularPipelineBlocks now defines an abstract __call__ raising a clear NotImplementedError. Loop steps get their own base class, ModularLoopPipelineBlocks, whose only difference is the __call__ contract (accepts the enclosing loop's variables after (components, state)). IterativePipelineBlocks validates at construction that every sub-block is a ModularLoopPipelineBlocks or a nested IterativePipelineBlocks, in addition to the signature validation before the first iteration. Co-Authored-By: Claude Fable 5 --- src/diffusers/__init__.py | 2 + src/diffusers/modular_pipelines/__init__.py | 2 + .../modular_pipelines/flux2/denoise.py | 10 ++--- .../modular_pipelines/modular_pipeline.py | 40 ++++++++++++++++- .../test_iterative_pipeline_blocks.py | 44 +++++++++++++++---- 5 files changed, 83 insertions(+), 15 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 25cfd4f93a31..776b848eb307 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -343,6 +343,7 @@ "ConfigSpec", "InputParam", "IterativePipelineBlocks", + "ModularLoopPipelineBlocks", "LoopSequentialPipelineBlocks", "ModularPipeline", "ModularPipelineBlocks", @@ -1215,6 +1216,7 @@ InputParam, IterativePipelineBlocks, LoopSequentialPipelineBlocks, + ModularLoopPipelineBlocks, ModularPipeline, ModularPipelineBlocks, OutputParam, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 7a11405f3317..09420b614beb 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -35,6 +35,7 @@ "SequentialPipelineBlocks", "ConditionalPipelineBlocks", "IterativePipelineBlocks", + "ModularLoopPipelineBlocks", "LoopSequentialPipelineBlocks", "PipelineState", "BlockState", @@ -163,6 +164,7 @@ ConditionalPipelineBlocks, IterativePipelineBlocks, LoopSequentialPipelineBlocks, + ModularLoopPipelineBlocks, ModularPipeline, ModularPipelineBlocks, PipelineState, diff --git a/src/diffusers/modular_pipelines/flux2/denoise.py b/src/diffusers/modular_pipelines/flux2/denoise.py index 57d8b63f2d59..2c3f08257da1 100644 --- a/src/diffusers/modular_pipelines/flux2/denoise.py +++ b/src/diffusers/modular_pipelines/flux2/denoise.py @@ -23,7 +23,7 @@ from ...utils import is_torch_xla_available, logging from ..modular_pipeline import ( IterativePipelineBlocks, - ModularPipelineBlocks, + ModularLoopPipelineBlocks, PipelineState, ) from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam @@ -41,7 +41,7 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -class Flux2LoopDenoiser(ModularPipelineBlocks): +class Flux2LoopDenoiser(ModularLoopPipelineBlocks): model_name = "flux2" @property @@ -143,7 +143,7 @@ def __call__( # same as Flux2LoopDenoiser but guidance=None -class Flux2KleinLoopDenoiser(ModularPipelineBlocks): +class Flux2KleinLoopDenoiser(ModularLoopPipelineBlocks): model_name = "flux2-klein" @property @@ -239,7 +239,7 @@ def __call__( # support CFG for Flux2-Klein base model -class Flux2KleinBaseLoopDenoiser(ModularPipelineBlocks): +class Flux2KleinBaseLoopDenoiser(ModularLoopPipelineBlocks): model_name = "flux2-klein" @property @@ -386,7 +386,7 @@ def __call__( return components, state -class Flux2LoopAfterDenoiser(ModularPipelineBlocks): +class Flux2LoopAfterDenoiser(ModularLoopPipelineBlocks): model_name = "flux2" @property diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index dc2b12e9eb45..81ddb86b7f7f 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -596,6 +596,27 @@ def doc(self): expected_configs=self.expected_configs, ) + def __call__(self, components, state: PipelineState) -> PipelineState: + raise NotImplementedError(f"`__call__` method must be implemented in {self.__class__.__name__}") + + +class ModularLoopPipelineBlocks(ModularPipelineBlocks): + """ + Base class for leaf blocks that run inside an [`IterativePipelineBlocks`] loop. + + The only difference from [`ModularPipelineBlocks`] is the `__call__` contract: in addition to + `(components, state)`, the block accepts the enclosing loop's variables as call arguments — its signature + must name exactly the loop's `loop_variables` (e.g. `def __call__(self, components, state, i, t)`), which + the loop validates before the first iteration. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + def __call__(self, components, state: PipelineState, **kwargs) -> PipelineState: + # Subclasses name the enclosing loop's variables explicitly, e.g. + # `def __call__(self, components, state, i, t)`. + raise NotImplementedError(f"`__call__` method must be implemented in {self.__class__.__name__}") + class ConditionalPipelineBlocks(ModularPipelineBlocks): """ @@ -1335,7 +1356,8 @@ def __call__(self, components, state): Unlike [`LoopSequentialPipelineBlocks`], sub-blocks operate on the full [`PipelineState`] with the regular `get_block_state`/`set_block_state` behavior, so an `IterativePipelineBlocks` can itself be a sub-block of - another one — loops can be nested and composed freely. + another one — loops can be nested and composed freely. Sub-blocks must be [`ModularLoopPipelineBlocks`] + (loop steps) or nested `IterativePipelineBlocks`, which is validated at construction. Loop variables are passed to sub-blocks as call arguments: every sub-block must have the signature `__call__(self, components, state, )`, which is validated against `loop_variables` before @@ -1418,6 +1440,20 @@ def expected_configs(self) -> list[ConfigSpec]: expected_configs.append(config) return expected_configs + def __init__(self): + super().__init__() + self._validate_sub_block_types() + + def _validate_sub_block_types(self): + """Sub-blocks must be loop steps (`ModularLoopPipelineBlocks`) or nested loops (`IterativePipelineBlocks`).""" + for block_name, block in self.sub_blocks.items(): + if not isinstance(block, (ModularLoopPipelineBlocks, IterativePipelineBlocks)): + raise ValueError( + f"Sub-block '{block_name}' ({block.__class__.__name__}) of {self.__class__.__name__} must be " + "a `ModularLoopPipelineBlocks` (a loop step) or an `IterativePipelineBlocks` (a nested loop); " + f"got `{block.__class__.__bases__[0].__name__}`." + ) + def _validate_loop_step_signatures(self): """Every sub-block must accept exactly the loop variables after `(components, state)`.""" expected = set(self.loop_variables) @@ -1434,6 +1470,8 @@ def _validate_loop_step_signatures(self): def loop_step(self, components, state: PipelineState, **loop_kwargs) -> PipelineState: """Run all sub-blocks once over the pipeline state (one loop iteration), passing the loop variables.""" if not getattr(self, "_loop_signatures_validated", False): + # re-validate types here to cover sub_blocks assigned after __init__ (e.g. from_blocks_dict) + self._validate_sub_block_types() self._validate_loop_step_signatures() self._loop_signatures_validated = True diff --git a/tests/modular_pipelines/test_iterative_pipeline_blocks.py b/tests/modular_pipelines/test_iterative_pipeline_blocks.py index 9bff50138623..27db8b2bf855 100644 --- a/tests/modular_pipelines/test_iterative_pipeline_blocks.py +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -19,6 +19,7 @@ from diffusers.modular_pipelines import ( InputParam, IterativePipelineBlocks, + ModularLoopPipelineBlocks, ModularPipelineBlocks, OutputParam, SequentialPipelineBlocks, @@ -31,7 +32,7 @@ # arguments; every leaf sub-block of a loop must accept its loop's variables. -class ChunkNoiseGenStep(ModularPipelineBlocks): +class ChunkNoiseGenStep(ModularLoopPipelineBlocks): model_name = "test" @property @@ -53,7 +54,7 @@ def __call__(self, components, state, k): return components, state -class LoopDenoiserStep(ModularPipelineBlocks): +class LoopDenoiserStep(ModularLoopPipelineBlocks): model_name = "test" @property @@ -75,7 +76,7 @@ def __call__(self, components, state, i, t): return components, state -class LoopSchedulerStep(ModularPipelineBlocks): +class LoopSchedulerStep(ModularLoopPipelineBlocks): model_name = "test" @property @@ -128,7 +129,7 @@ def __call__(self, components, state, k): return components, state -class ChunkUpdateStep(ModularPipelineBlocks): +class ChunkUpdateStep(ModularLoopPipelineBlocks): model_name = "test" @property @@ -231,25 +232,50 @@ def test_loop_variables_do_not_leak_into_state(self): # declared sub-block outputs persist after the loop (last iteration's value) assert state.get("noise_pred") is not None - def test_leaf_signature_is_validated(self): + def test_sub_block_type_is_validated(self): + # a regular ModularPipelineBlocks cannot be a loop sub-block: fails at construction class PlainStep(ModularPipelineBlocks): model_name = "test" @property def description(self): - return "regular block without the loop variables" + return "regular block, not a loop step" def __call__(self, components, state): return components, state - class BadLoop(IterativePipelineBlocks): + class BadTypeLoop(IterativePipelineBlocks): model_name = "test" block_classes = [PlainStep] block_names = ["plain"] @property def description(self): - return "loop with a mismatched leaf signature" + return "loop with a non-loop sub-block" + + with pytest.raises(ValueError, match="must be a `ModularLoopPipelineBlocks`"): + BadTypeLoop() + + def test_leaf_signature_is_validated(self): + # a loop step whose signature doesn't match the loop's variables fails before the first iteration + class WrongSigStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "loop step with the wrong loop variables" + + def __call__(self, components, state, k): + return components, state + + class BadSigLoop(IterativePipelineBlocks): + model_name = "test" + block_classes = [WrongSigStep] + block_names = ["wrong"] + + @property + def description(self): + return "loop whose sub-block names the wrong loop variables" @property def loop_variables(self): @@ -266,7 +292,7 @@ def __call__(self, components, state): components, state = self.loop_step(components, state, i=i, t=t) return components, state - pipe = SequentialPipelineBlocks.from_blocks_dict({"loop": BadLoop()}).init_pipeline() + pipe = SequentialPipelineBlocks.from_blocks_dict({"loop": BadSigLoop()}).init_pipeline() with pytest.raises(ValueError, match="must accept the loop variables"): pipe(timesteps=torch.tensor([1.0])) From 281333e0a358385f7c57ee90b820e95b57c71b8b Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 17:50:02 +0000 Subject: [PATCH 06/10] Drop loop_* declaration properties; validate sub-blocks at construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove loop_inputs / loop_intermediate_outputs / loop_expected_components / loop_expected_configs from IterativePipelineBlocks — when the loop logic in __call__ consumes inputs or components beyond what sub-blocks declare, the subclass overrides the aggregated inputs / expected_components properties directly (see Flux2DenoiseLoopWrapper). Sub-block validation (type + loop-variable signature) now runs at construction (__init__ and from_blocks_dict) instead of lazily in loop_step. Co-Authored-By: Claude Fable 5 --- .../modular_pipelines/flux2/denoise.py | 16 +++- .../modular_pipelines/modular_pipeline.py | 82 ++++--------------- .../test_iterative_pipeline_blocks.py | 21 +++-- 3 files changed, 39 insertions(+), 80 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux2/denoise.py b/src/diffusers/modular_pipelines/flux2/denoise.py index 2c3f08257da1..8d7716fda436 100644 --- a/src/diffusers/modular_pipelines/flux2/denoise.py +++ b/src/diffusers/modular_pipelines/flux2/denoise.py @@ -451,9 +451,13 @@ def loop_variables(self) -> list[str]: return ["i", "t"] @property - def loop_expected_components(self) -> list[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: + expected_components = super().expected_components # the loop logic itself reads `scheduler.order` for the warmup-step computation - return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + scheduler = ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler) + if scheduler not in expected_components: + expected_components.append(scheduler) + return expected_components @property def description(self) -> str: @@ -463,8 +467,11 @@ def description(self) -> str: ) @property - def loop_inputs(self) -> list[InputParam]: - return [ + def inputs(self) -> list[InputParam]: + inputs = super().inputs + names = {param.name for param in inputs} + # inputs consumed by the loop logic itself, on top of what the sub-blocks declare + loop_inputs = [ InputParam( "timesteps", required=True, @@ -478,6 +485,7 @@ def loop_inputs(self) -> list[InputParam]: description="The number of inference steps to use for the denoising process.", ), ] + return [param for param in loop_inputs if param.name not in names] + inputs @torch.no_grad() def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 81ddb86b7f7f..c30201801ea3 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -1360,8 +1360,8 @@ def __call__(self, components, state): (loop steps) or nested `IterativePipelineBlocks`, which is validated at construction. Loop variables are passed to sub-blocks as call arguments: every sub-block must have the signature - `__call__(self, components, state, )`, which is validated against `loop_variables` before - the first iteration. A nested loop accepts the outer loop's variables in its own hand-written `__call__` + `__call__(self, components, state, )`, which is validated against `loop_variables` at + construction. A nested loop accepts the outer loop's variables in its own hand-written `__call__` (ignoring or forwarding them) and passes its own `loop_variables` to its own sub-blocks: ```python @@ -1378,7 +1378,9 @@ def __call__(self, components, state, k): # accepts the OUTER chunk loop's return components, state ``` - Sub-block outputs are written to the pipeline state as usual and persist after the loop. + Sub-block outputs are written to the pipeline state as usual and persist after the loop. If the loop logic in + `__call__` itself consumes inputs (e.g. `timesteps`) or uses components (e.g. the scheduler) beyond what the + sub-blocks declare, override the aggregated `inputs` / `expected_components` / ... properties to add them. > [!WARNING] > This is an experimental feature and is likely to change in the future. @@ -1392,60 +1394,21 @@ def loop_variables(self) -> list[str]: """Names of the loop variables `loop_step` passes to leaf sub-blocks each iteration (e.g. `["i", "t"]`).""" return [] - @property - def loop_inputs(self) -> list[InputParam]: - """Inputs consumed by the loop logic in `__call__` itself (e.g. `timesteps`).""" - return [] - - @property - def loop_intermediate_outputs(self) -> list[OutputParam]: - """Outputs written to the pipeline state by the loop logic in `__call__` itself.""" - return [] - - @property - def loop_expected_components(self) -> list[ComponentSpec]: - """Components used by the loop logic in `__call__` itself (e.g. the scheduler).""" - return [] - - @property - def loop_expected_configs(self) -> list[ConfigSpec]: - """Configs used by the loop logic in `__call__` itself.""" - return [] - - @property - def inputs(self) -> list[InputParam]: - inputs = self._get_inputs() - names = {p.name for p in inputs} - return [p for p in self.loop_inputs if p.name not in names] + inputs - - @property - def intermediate_outputs(self) -> list[OutputParam]: - outputs = super().intermediate_outputs - names = {output.name for output in outputs} - return outputs + [output for output in self.loop_intermediate_outputs if output.name not in names] - - @property - def expected_components(self) -> list[ComponentSpec]: - expected_components = super().expected_components - for component in self.loop_expected_components: - if component not in expected_components: - expected_components.append(component) - return expected_components - - @property - def expected_configs(self) -> list[ConfigSpec]: - expected_configs = super().expected_configs - for config in self.loop_expected_configs: - if config not in expected_configs: - expected_configs.append(config) - return expected_configs - def __init__(self): super().__init__() - self._validate_sub_block_types() + self._validate_sub_blocks() + + @classmethod + def from_blocks_dict(cls, blocks_dict, description: str | None = None) -> "IterativePipelineBlocks": + instance = super().from_blocks_dict(blocks_dict, description) + # sub_blocks are assigned after __init__ on this path, so validate again + instance._validate_sub_blocks() + return instance - def _validate_sub_block_types(self): - """Sub-blocks must be loop steps (`ModularLoopPipelineBlocks`) or nested loops (`IterativePipelineBlocks`).""" + def _validate_sub_blocks(self): + """Sub-blocks must be loop steps (`ModularLoopPipelineBlocks`) or nested loops (`IterativePipelineBlocks`) + and accept exactly the loop variables after `(components, state)`.""" + expected = set(self.loop_variables) for block_name, block in self.sub_blocks.items(): if not isinstance(block, (ModularLoopPipelineBlocks, IterativePipelineBlocks)): raise ValueError( @@ -1453,11 +1416,6 @@ def _validate_sub_block_types(self): "a `ModularLoopPipelineBlocks` (a loop step) or an `IterativePipelineBlocks` (a nested loop); " f"got `{block.__class__.__bases__[0].__name__}`." ) - - def _validate_loop_step_signatures(self): - """Every sub-block must accept exactly the loop variables after `(components, state)`.""" - expected = set(self.loop_variables) - for block_name, block in self.sub_blocks.items(): params = list(inspect.signature(block.__call__).parameters) extra = set(params[2:]) if extra != expected: @@ -1469,12 +1427,6 @@ def _validate_loop_step_signatures(self): def loop_step(self, components, state: PipelineState, **loop_kwargs) -> PipelineState: """Run all sub-blocks once over the pipeline state (one loop iteration), passing the loop variables.""" - if not getattr(self, "_loop_signatures_validated", False): - # re-validate types here to cover sub_blocks assigned after __init__ (e.g. from_blocks_dict) - self._validate_sub_block_types() - self._validate_loop_step_signatures() - self._loop_signatures_validated = True - for block_name, block in self.sub_blocks.items(): try: components, state = block(components, state, **loop_kwargs) diff --git a/tests/modular_pipelines/test_iterative_pipeline_blocks.py b/tests/modular_pipelines/test_iterative_pipeline_blocks.py index 27db8b2bf855..ad74c649d62c 100644 --- a/tests/modular_pipelines/test_iterative_pipeline_blocks.py +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -118,8 +118,8 @@ def loop_variables(self): return ["i", "t"] @property - def loop_inputs(self): - return [InputParam(name="timesteps", required=True)] + def inputs(self): + return [InputParam(name="timesteps", required=True), *super().inputs] @torch.no_grad() def __call__(self, components, state, k): @@ -168,8 +168,8 @@ def loop_variables(self): return ["k"] @property - def loop_inputs(self): - return [InputParam(name="num_latent_chunk", required=True)] + def inputs(self): + return [InputParam(name="num_latent_chunk", required=True), *super().inputs] @torch.no_grad() def __call__(self, components, state): @@ -180,11 +180,11 @@ def __call__(self, components, state): class TestIterativePipelineBlocksStructure: - def test_loop_inputs_aggregation(self): + def test_inputs_aggregation(self): loop = ChunkLoop() input_names = [p.name for p in loop.inputs] - # loop_inputs of the loop itself and of the nested loop are surfaced + # inputs of the loop logic itself and of the nested loop are surfaced assert "num_latent_chunk" in input_names assert "timesteps" in input_names # loop variables are call arguments, not inputs @@ -257,7 +257,7 @@ def description(self): BadTypeLoop() def test_leaf_signature_is_validated(self): - # a loop step whose signature doesn't match the loop's variables fails before the first iteration + # a loop step whose signature doesn't match the loop's variables fails at construction class WrongSigStep(ModularLoopPipelineBlocks): model_name = "test" @@ -282,8 +282,8 @@ def loop_variables(self): return ["i", "t"] @property - def loop_inputs(self): - return [InputParam(name="timesteps", required=True)] + def inputs(self): + return [InputParam(name="timesteps", required=True), *super().inputs] @torch.no_grad() def __call__(self, components, state): @@ -292,9 +292,8 @@ def __call__(self, components, state): components, state = self.loop_step(components, state, i=i, t=t) return components, state - pipe = SequentialPipelineBlocks.from_blocks_dict({"loop": BadSigLoop()}).init_pipeline() with pytest.raises(ValueError, match="must accept the loop variables"): - pipe(timesteps=torch.tensor([1.0])) + BadSigLoop() def test_loop_leaf_standalone_raises(self): # outside a loop, a leaf block with loop variables in its signature cannot run From 2cf97ca5d171d34eafbe4bd568b20edc09e78873 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 19 Aug 2026 08:41:39 +0000 Subject: [PATCH 07/10] =?UTF-8?q?Modular:=20opt-in=20streaming=20=E2=80=94?= =?UTF-8?q?=20pipe.stream()=20yields=20the=20live=20state=20after=20every?= =?UTF-8?q?=20loop=20iteration?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - `stream(components, state)` generator on every block type; leaves yield nothing, composites re-yield sub-block events with the sub-block name prepended to `event.path` - `IterativePipelineBlocks`: streaming is opt-in — `__call__`/`loop_step` unchanged; a loop that supports it also implements `stream` as a generator over the new `stream_step`; base `stream` raises a clear NotImplementedError - `StreamEvent(path, block, loop_kwargs, state)` dataclass, exported - `supports_streaming` property on blocks (legacy `LoopSequentialPipelineBlocks` sets it to False; to be deprecated once all pipelines are ported) - `ModularPipeline.stream(state=None, **kwargs)` — same seeding as `__call__`, no_grad held only while blocks run; error logging mirrors `__call__` at every level - flux2: `Flux2DenoiseLoopWrapper.stream` next to its `__call__` - tests: `test_stream_matches_call` on `ModularPipelineTesterMixin` (auto-skips when blocks don't support streaming) - docs: Streaming section on the ModularPipeline page; regenerate dummy_pt_objects Co-Authored-By: Claude Fable 5 --- .../en/modular_diffusers/modular_pipeline.md | 62 +++++ src/diffusers/__init__.py | 4 +- src/diffusers/modular_pipelines/__init__.py | 2 + .../modular_pipelines/flux2/denoise.py | 9 + .../modular_pipelines/modular_pipeline.py | 243 +++++++++++++++++- src/diffusers/utils/dummy_pt_objects.py | 45 ++++ .../test_modular_pipelines_common.py | 26 ++ 7 files changed, 378 insertions(+), 13 deletions(-) diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 27bc61634805..f23da58b7c2d 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -380,6 +380,68 @@ output = pipeline( If pipeline stages share components (e.g., the same VAE used for encoding and decoding), you can use [`~ModularPipeline.update_components`] to pass an already-loaded component to another pipeline instead of loading it again. +## Streaming + +[`~ModularPipeline.stream`] runs the same pipeline as a generator. It yields a [`StreamEvent`] after every iteration of every loop block — each denoising step, each segment of a chunked video — with the live [`PipelineState`] attached, so you can show progress, decode a preview, or stop early. The generator's return value is the final state, exactly what `__call__` returns. + +```py +generator = pipeline.stream(prompt="a cat", num_inference_steps=20) +for event in generator: + print(event.path, event.loop_kwargs) # "denoise.denoise" {"i": 0, "t": tensor(1000.)} + latents = event.state.get("latents") # the live state — clone anything you keep +``` + +`event.path` is the loop block's dotted name from the top of the pipeline, and `event.loop_kwargs` its loop variables for that iteration. When loops are nested — an autoregressive video that denoises one chunk at a time — the inner loop's events surface too, so a consumer that only wants finished chunks filters on the outer path: + +```py +for event in pipeline.stream(...): + if event.path == "denoise": # the chunk loop, not "denoise.denoise_inner" + show(event.state.get("out_frames")) +``` + +To stop early, stop iterating (or call `generator.close()`); nothing needs cleaning up. Blocks without loops run to completion and yield nothing. + +Streaming is opt-in per loop block. An [`IterativePipelineBlocks`] implements its loop in `__call__` as usual and, to support streaming, also implements `stream` — the same loop written as a generator over `stream_step`, which runs one iteration like `loop_step` and additionally yields the event for it: + +```py +class DenoiseLoop(IterativePipelineBlocks): + @property + def loop_variables(self): + return ["i", "t"] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + return components, state + + def stream(self, components, state): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + return components, state +``` + +`pipeline.stream(...)` raises `NotImplementedError` if a loop on its path doesn't implement `stream`. Check `pipeline.blocks.supports_streaming` to find out ahead of time — it is `True` unless the blocks contain a loop that can't yield per iteration (an `IterativePipelineBlocks` that doesn't implement `stream`, or a legacy `LoopSequentialPipelineBlocks`). + +If you need to own the loop yourself — a serving engine that advances every request by one denoising step per tick, or a real-time pipeline fed one chunk of input at a time — run the blocks before the loop, then call the loop block's `loop_step` once per iteration. Anything you write into the state between calls is seen by the next iteration: + +```py +from diffusers.modular_pipelines import PipelineState + +loop = pipeline.blocks.sub_blocks["denoise"] + +state = PipelineState() +for param in pipeline.blocks.inputs: # seed the declared defaults + state.set(param.name, param.default) +state.set("prompt", "a cat") +state.set("num_inference_steps", 20) +# ... run the blocks before `denoise` on `state` ... +for i, t in enumerate(state.get("timesteps")): + _, state = loop.loop_step(pipeline, state, i=i, t=t) +``` + ## Modular repository A repository is required if the pipeline blocks use *pretrained components*. The repository supplies loading specifications and metadata. diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 776b848eb307..64d5efe9f493 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -343,12 +343,13 @@ "ConfigSpec", "InputParam", "IterativePipelineBlocks", - "ModularLoopPipelineBlocks", "LoopSequentialPipelineBlocks", + "ModularLoopPipelineBlocks", "ModularPipeline", "ModularPipelineBlocks", "OutputParam", "SequentialPipelineBlocks", + "StreamEvent", ] ) _import_structure["optimization"] = [ @@ -1221,6 +1222,7 @@ ModularPipelineBlocks, OutputParam, SequentialPipelineBlocks, + StreamEvent, ) from .optimization import ( get_constant_schedule, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 09420b614beb..54abd64d7c4b 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -39,6 +39,7 @@ "LoopSequentialPipelineBlocks", "PipelineState", "BlockState", + "StreamEvent", ] _import_structure["modular_pipeline_utils"] = [ "ComponentSpec", @@ -169,6 +170,7 @@ ModularPipelineBlocks, PipelineState, SequentialPipelineBlocks, + StreamEvent, ) from .modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, InsertableDict, OutputParam from .qwenimage import ( diff --git a/src/diffusers/modular_pipelines/flux2/denoise.py b/src/diffusers/modular_pipelines/flux2/denoise.py index 8d7716fda436..fa40cf7c7523 100644 --- a/src/diffusers/modular_pipelines/flux2/denoise.py +++ b/src/diffusers/modular_pipelines/flux2/denoise.py @@ -508,6 +508,15 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi return components, state + @torch.no_grad() + def stream(self, components: Flux2ModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + if XLA_AVAILABLE: + xm.mark_step() + return components, state + class Flux2DenoiseStep(Flux2DenoiseLoopWrapper): block_classes = [Flux2LoopDenoiser, Flux2LoopAfterDenoiser] diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index c30201801ea3..eea3a4e3aefc 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -305,6 +305,26 @@ def format_value(v): return f"BlockState(\n{attributes}\n)" +@dataclass +class StreamEvent: + """ + One iteration of a loop block, as yielded by `stream()`. + + Attributes: + path: + Dotted name of the loop block from the top of the pipeline, e.g. `"denoise"` or, for a loop nested inside + another, `"denoise.denoise_inner"`. + block: The [`IterativePipelineBlocks`] that just finished the iteration. + loop_kwargs: The loop variables of that iteration, e.g. `{"i": 3, "t": tensor(...)}`. + state: The live [`PipelineState`] after the iteration — not a copy. Clone anything you keep. + """ + + path: str + block: "IterativePipelineBlocks" + loop_kwargs: dict[str, Any] + state: PipelineState + + class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): """ Base class for all Pipeline Blocks: ConditionalPipelineBlocks, AutoPipelineBlocks, SequentialPipelineBlocks, @@ -599,15 +619,36 @@ def doc(self): def __call__(self, components, state: PipelineState) -> PipelineState: raise NotImplementedError(f"`__call__` method must be implemented in {self.__class__.__name__}") + def stream(self, components, state: PipelineState): + """ + Run the block as a generator that yields a [`StreamEvent`] after every iteration of every loop block it + contains, and returns `(components, state)` when done. A block with no loops runs to completion and yields + nothing; composite blocks re-yield their sub-blocks' events with the sub-block name prepended to `event.path`. + """ + yield from () + return self(components, state) + + @property + def supports_streaming(self) -> bool: + """ + Whether the block can be streamed: `True` unless it contains a loop that can't yield per iteration — an + [`IterativePipelineBlocks`] that doesn't implement `stream`, or a legacy [`LoopSequentialPipelineBlocks`]. + """ + for block in self.sub_blocks.values(): + if not block.supports_streaming: + return False + # a leaf block (no sub_blocks) always streams: it runs to completion and yields nothing + return True + class ModularLoopPipelineBlocks(ModularPipelineBlocks): """ Base class for leaf blocks that run inside an [`IterativePipelineBlocks`] loop. - The only difference from [`ModularPipelineBlocks`] is the `__call__` contract: in addition to - `(components, state)`, the block accepts the enclosing loop's variables as call arguments — its signature - must name exactly the loop's `loop_variables` (e.g. `def __call__(self, components, state, i, t)`), which - the loop validates before the first iteration. + The only difference from [`ModularPipelineBlocks`] is the `__call__` contract: in addition to `(components, + state)`, the block accepts the enclosing loop's variables as call arguments — its signature must name exactly the + loop's `loop_variables` (e.g. `def __call__(self, components, state, i, t)`), which the loop validates before the + first iteration. > [!WARNING] > This is an experimental feature and is likely to change in the future. """ @@ -810,12 +851,40 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState: logger.error(error_msg) raise + def stream(self, pipeline, state: PipelineState): + # Same branch selection as `__call__`. The branch is transparent in event paths, as it is in + # `get_execution_blocks`: events carry the name this conditional block has in its parent, not the branch name. + trigger_kwargs = {name: state.get(name) for name in self.block_trigger_inputs if name is not None} + block_name = self.select_block(**trigger_kwargs) + + if block_name is None: + block_name = self.default_block_name + + if block_name is None: + logger.info(f"skipping conditional block: {self.__class__.__name__}") + return pipeline, state + + block = self.sub_blocks[block_name] + + try: + logger.info(f"Running block: {block.__class__.__name__}") + return (yield from block.stream(pipeline, state)) + except Exception as e: + error_msg = ( + f"\nError in block: {block.__class__.__name__}\n" + f"Error details: {str(e)}\n" + f"Traceback:\n{traceback.format_exc()}" + ) + logger.error(error_msg) + raise + def get_execution_blocks(self, **kwargs) -> ModularPipelineBlocks | None: """ Get the block(s) that would execute given the inputs. Recursively resolves nested ConditionalPipelineBlocks until reaching either: - - A leaf block (no sub_blocks, or a loop block: IterativePipelineBlocks / LoopSequentialPipelineBlocks) → returns single `ModularPipelineBlocks` + - A leaf block (no sub_blocks, or a loop block: IterativePipelineBlocks / LoopSequentialPipelineBlocks) → + returns single `ModularPipelineBlocks` - A `SequentialPipelineBlocks` → delegates to its `get_execution_blocks()` which returns a `SequentialPipelineBlocks` containing the resolved execution blocks @@ -1166,6 +1235,28 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState: raise return pipeline, state + def stream(self, pipeline, state: PipelineState): + for block_name, block in self.sub_blocks.items(): + # re-yield the sub-block's events with its name prepended to the path; its return value is the new state + generator = block.stream(pipeline, state) + while True: + try: + event = next(generator) + except StopIteration as e: + pipeline, state = e.value + break + except Exception as e: + error_msg = ( + f"\nError in block: ({block_name}, {block.__class__.__name__})\n" + f"Error details: {str(e)}\n" + f"Traceback:\n{traceback.format_exc()}" + ) + logger.error(error_msg) + raise + event.path = f"{block_name}.{event.path}" if event.path else block_name + yield event + return pipeline, state + # used for `__repr__` def _get_trigger_inputs(self): """ @@ -1346,6 +1437,7 @@ class IterativePipelineBlocks(SequentialPipelineBlocks): def loop_variables(self): return ["i", "t"] + @torch.no_grad() def __call__(self, components, state): block_state = self.get_block_state(state) @@ -1355,23 +1447,23 @@ def __call__(self, components, state): ``` Unlike [`LoopSequentialPipelineBlocks`], sub-blocks operate on the full [`PipelineState`] with the regular - `get_block_state`/`set_block_state` behavior, so an `IterativePipelineBlocks` can itself be a sub-block of - another one — loops can be nested and composed freely. Sub-blocks must be [`ModularLoopPipelineBlocks`] - (loop steps) or nested `IterativePipelineBlocks`, which is validated at construction. + `get_block_state`/`set_block_state` behavior, so an `IterativePipelineBlocks` can itself be a sub-block of another + one — loops can be nested and composed freely. Sub-blocks must be [`ModularLoopPipelineBlocks`] (loop steps) or + nested `IterativePipelineBlocks`, which is validated at construction. Loop variables are passed to sub-blocks as call arguments: every sub-block must have the signature `__call__(self, components, state, )`, which is validated against `loop_variables` at - construction. A nested loop accepts the outer loop's variables in its own hand-written `__call__` - (ignoring or forwarding them) and passes its own `loop_variables` to its own sub-blocks: + construction. A nested loop accepts the outer loop's variables in its own hand-written `__call__` (ignoring or + forwarding them) and passes its own `loop_variables` to its own sub-blocks: ```python class InnerDenoiseLoop(IterativePipelineBlocks): @property def loop_variables(self): - return ["i", "t"] # what it passes to ITS sub-blocks + return ["i", "t"] # what it passes to ITS sub-blocks @torch.no_grad() - def __call__(self, components, state, k): # accepts the OUTER chunk loop's variable + def __call__(self, components, state, k): # accepts the OUTER chunk loop's variable block_state = self.get_block_state(state) for i, t in enumerate(block_state.timesteps): components, state = self.loop_step(components, state, i=i, t=t) @@ -1382,6 +1474,20 @@ def __call__(self, components, state, k): # accepts the OUTER chunk loop's `__call__` itself consumes inputs (e.g. `timesteps`) or uses components (e.g. the scheduler) beyond what the sub-blocks declare, override the aggregated `inputs` / `expected_components` / ... properties to add them. + Streaming is opt-in: to let `pipe.stream(...)` hand back the live [`PipelineState`] after every iteration, also + implement `stream` — the same loop, written as a generator over `stream_step` (which runs one iteration like + `loop_step` and additionally yields a [`StreamEvent`] for it, after any events of a nested loop): + + ```python + def stream(self, components, state): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + return components, state + ``` + + A nested loop's `stream` takes the outer loop's variables exactly like its `__call__` does. + > [!WARNING] > This is an experimental feature and is likely to change in the future. Attributes: @@ -1446,6 +1552,52 @@ def __call__(self, components, state: PipelineState, **kwargs) -> PipelineState: # e.g. `def __call__(self, components, state, k)`. raise NotImplementedError("`__call__` method needs to be implemented by the subclass") + def stream_step(self, components, state: PipelineState, **loop_kwargs): + """ + The streaming counterpart of `loop_step`: a generator that runs all sub-blocks once, re-yields the events of + any nested loop, then yields one [`StreamEvent`] for this iteration and returns `(components, state)`. Call it + with `yield from` inside `stream`. + """ + for block_name, block in self.sub_blocks.items(): + try: + if isinstance(block, IterativePipelineBlocks): + # a nested loop streams too: re-yield its events with its name prepended to the path + generator = block.stream(components, state, **loop_kwargs) + while True: + try: + event = next(generator) + except StopIteration as e: + components, state = e.value + break + event.path = f"{block_name}.{event.path}" if event.path else block_name + yield event + else: + components, state = block(components, state, **loop_kwargs) + except Exception as e: + error_msg = ( + f"\nError in block: ({block_name}, {block.__class__.__name__})\n" + f"Error details: {str(e)}\n" + f"Traceback:\n{traceback.format_exc()}" + ) + logger.error(error_msg) + raise + yield StreamEvent(path="", block=self, loop_kwargs=dict(loop_kwargs), state=state) + return components, state + + def stream(self, components, state: PipelineState, **kwargs): + # Optional. Subclasses that support streaming implement their loop logic here a second time, as a generator + # running `yield from self.stream_step(...)` once per iteration, with the same signature as their `__call__` + # (`**kwargs` stands for the outer loop's variables when this loop is nested, e.g. `stream(self, components, + # state, k)`). + raise NotImplementedError( + f"{self.__class__.__name__} does not support streaming: implement `stream` (the loop written as a " + "generator over `stream_step`) to use it with `pipe.stream(...)`." + ) + + @property + def supports_streaming(self) -> bool: + return type(self).stream is not IterativePipelineBlocks.stream and super().supports_streaming + class LoopSequentialPipelineBlocks(ModularPipelineBlocks): """ @@ -1466,6 +1618,10 @@ class LoopSequentialPipelineBlocks(ModularPipelineBlocks): block_classes = [] block_names = [] + # this legacy loop runs all iterations in one call and cannot yield per iteration; it will be deprecated in + # favor of `IterativePipelineBlocks` once all pipelines are ported to it + supports_streaming = False + @property def description(self) -> str: """Description of the block. Must be implemented by subclasses.""" @@ -2980,3 +3136,66 @@ def __call__(self, state: PipelineState = None, output: str | list[str] = None, return state.get(output) else: raise ValueError(f"Output '{output}' is not a valid output type") + + def stream(self, state: PipelineState = None, **kwargs): + """ + Run the pipeline as a generator that yields a [`StreamEvent`] after every iteration of every loop block — each + denoising step, each segment of a chunked video, and so on — with the live [`PipelineState`] attached. The + generator's return value is the final state, the same one `__call__` returns. + + Args: + state (`PipelineState`, optional): + Same as in `__call__`. + **kwargs: + Same as in `__call__`. + + Examples: + ```python + for event in pipeline.stream(prompt="A beautiful sunset", num_inference_steps=20): + print(event.path, event.loop_kwargs) # "denoise" {"i": 0, "t": tensor(...)} + latents = event.state.get("latents") # live state — clone anything you keep + + # stop early: just stop iterating (or call `.close()` on the generator) + ``` + """ + if state is None: + state = PipelineState() + else: + state = deepcopy(state) + + # Make a copy of the input kwargs + passed_kwargs = kwargs.copy() + + # Add inputs to state, using defaults if not provided in the kwargs or the state + # if same input already in the state, will override it if provided in the kwargs + for expected_input_param in self._blocks.inputs: + name = expected_input_param.name + default = expected_input_param.default + kwargs_type = expected_input_param.kwargs_type + if name in passed_kwargs: + state.set(name, passed_kwargs.pop(name), kwargs_type) + elif kwargs_type is not None and kwargs_type in passed_kwargs: + kwargs_dict = passed_kwargs.pop(kwargs_type) + for k, v in kwargs_dict.items(): + state.set(k, v, kwargs_type) + elif name is not None and name not in state.values: + state.set(name, default, kwargs_type) + + # Warn about unexpected inputs + if len(passed_kwargs) > 0: + warnings.warn(f"Unexpected input '{passed_kwargs.keys()}' provided. This input will be ignored.") + + # `torch.no_grad()` is held only while blocks run, not while the caller holds an event. + generator = self._blocks.stream(self, state) + while True: + with torch.no_grad(): + try: + event = next(generator) + except StopIteration as e: + _, state = e.value + return state + except Exception: + error_msg = f"Error in block: ({self._blocks.__class__.__name__}):\n" + logger.error(error_msg) + raise + yield event diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 9035efb3e6e2..407bd6f4a6bf 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2419,6 +2419,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class IterativePipelineBlocks(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class LoopSequentialPipelineBlocks(metaclass=DummyObject): _backends = ["torch"] @@ -2434,6 +2449,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class ModularLoopPipelineBlocks(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class ModularPipeline(metaclass=DummyObject): _backends = ["torch"] @@ -2494,6 +2524,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class StreamEvent(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + def get_constant_schedule(*args, **kwargs): requires_backends(get_constant_schedule, ["torch"]) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index 223a25e436fa..d5a418c43305 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -469,6 +469,32 @@ def test_workflow_map(self): f"{actual_block.__class__.__name__}, expected {expected_class_name}" ) + def test_stream_matches_call(self, expected_max_diff=1e-4): + pipe = self.get_pipeline().to(torch_device) + + if not pipe.blocks.supports_streaming: + pytest.skip("Skipping test as blocks do not support streaming.") + + inputs = self.get_dummy_inputs() + inputs["generator"] = self.get_generator(0) + output = pipe(**inputs, output=self.output_name) + + inputs = self.get_dummy_inputs() + inputs["generator"] = self.get_generator(0) + num_events = 0 + generator = pipe.stream(**inputs) + while True: + try: + next(generator) + except StopIteration as e: + state = e.value + break + num_events += 1 + + assert num_events > 0, "stream() yielded no events" + max_diff = torch.abs(state.get(self.output_name) - output).max() + assert max_diff < expected_max_diff, "stream() results different from __call__ results" + class ModularGuiderTesterMixin: def test_guider_cfg(self, expected_max_diff=1e-2): From 0f1245a6ac3ec1a6345151701414c728c477459b Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 19 Aug 2026 19:41:38 +0000 Subject: [PATCH 08/10] Port wan-animate-2 to IterativePipelineBlocks with nested streaming - WanAnimate2SegmentLoopWrapper is an IterativePipelineBlocks (loop variable `k`); the per-segment steps become ModularLoopPipelineBlocks operating on the shared PipelineState - the hand-rolled timestep loop in WanAnimate2SegmentDenoiseInner becomes a nested loop: WanAnimate2LoopDenoiser + WanAnimate2LoopAfterDenoiser under WanAnimate2DenoiseLoopWrapper (loop variables `i`, `t`) - loop-carried state made explicit: `out_frames` is a declared input on the prev-frames step but filtered from the wrapper's aggregated inputs; `segment_frames` accumulation moves from the decode step into the segment loop's own logic - both loops implement `stream`: events per denoise step and per segment, so test_stream_matches_call now runs for wan-animate-2 (stream == call) - regenerate stale flux2 modular_blocks docstrings; doc-builder reflow in modular_pipeline.py Co-Authored-By: Claude Fable 5 --- .../flux2/modular_blocks_flux2.py | 4 +- .../flux2/modular_blocks_flux2_klein.py | 8 +- .../flux2/modular_blocks_flux2_klein_base.py | 4 +- .../modular_pipelines/modular_pipeline.py | 12 +- .../wan_animate_2/denoise.py | 404 ++++++++++++------ .../modular_blocks_wan_animate_2.py | 4 +- .../modular_blocks_wan_animate_2_distilled.py | 4 +- 7 files changed, 303 insertions(+), 137 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py index 2bbb7975a983..80800c1d5936 100644 --- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py @@ -316,9 +316,9 @@ class Flux2AutoBlocks(SequentialPipelineBlocks): TODO: Add description. latents (`Tensor | NoneType`): TODO: Add description. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 50): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py index 689cf808c4ba..8756adf6c3ae 100644 --- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py @@ -284,9 +284,9 @@ class Flux2KleinAutoCoreDenoiseStep(AutoPipelineBlocks): TODO: Add description. generator (`None`, *optional*): TODO: Add description. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 50): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. @@ -357,9 +357,9 @@ class Flux2KleinAutoBlocks(SequentialPipelineBlocks): TODO: Add description. latents (`Tensor | NoneType`): TODO: Add description. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 50): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py index f3108bdadeac..7c43be332047 100644 --- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py @@ -299,7 +299,7 @@ class Flux2KleinBaseAutoCoreDenoiseStep(AutoPipelineBlocks): TODO: Add description. num_inference_steps (`None`): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. @@ -373,7 +373,7 @@ class Flux2KleinBaseAutoBlocks(SequentialPipelineBlocks): TODO: Add description. num_inference_steps (`None`): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 8e936c827118..d3e703d6bd85 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -1456,8 +1456,8 @@ def _requirements(self) -> dict[str, str]: class IterativePipelineBlocks(SequentialPipelineBlocks): """ A pipeline blocks that runs its sub-blocks multiple times. Subclasses declare their loop-variable names in - `loop_variables` and implement `__call__` with their loop logic — the same way leaf blocks implement - `__call__` around `get_block_state` — calling `loop_step` once per iteration with the loop variables: + `loop_variables` and implement `__call__` with their loop logic — the same way leaf blocks implement `__call__` + around `get_block_state` — calling `loop_step` once per iteration with the loop variables: ```python @property @@ -1478,10 +1478,10 @@ def __call__(self, components, state): one — loops can be nested and composed freely. Sub-blocks must be [`ModularLoopPipelineBlocks`] (loop steps) or nested `IterativePipelineBlocks`, which is validated at construction. - Loop variables are passed to sub-blocks as call arguments: every sub-block must have the signature - `__call__(self, components, state, )`, which is validated against `loop_variables` at - construction. A nested loop accepts the outer loop's variables in its own hand-written `__call__` (ignoring or - forwarding them) and passes its own `loop_variables` to its own sub-blocks: + Loop variables are passed to sub-blocks as call arguments: every sub-block must have the signature `__call__(self, + components, state, )`, which is validated against `loop_variables` at construction. A nested + loop accepts the outer loop's variables in its own hand-written `__call__` (ignoring or forwarding them) and passes + its own `loop_variables` to its own sub-blocks: ```python class InnerDenoiseLoop(IterativePipelineBlocks): diff --git a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py index d96b8f814239..8cd34c7cfc5c 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -25,7 +25,7 @@ from ...schedulers.scheduling_utils import SchedulerMixin from ...utils import logging from ...utils.torch_utils import randn_tensor -from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState +from ..modular_pipeline import IterativePipelineBlocks, ModularLoopPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .encoders import encode_vae, get_i2v_mask @@ -47,11 +47,11 @@ def decode_vae(vae: AutoencoderKLWan, latents: torch.Tensor) -> torch.Tensor: # ======================================== -# Segment Loop Leaf Blocks +# Segment Loop Steps # ======================================== -class WanAnimate2SegmentVaeEncoderStep(ModularPipelineBlocks): +class WanAnimate2SegmentVaeEncoderStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -61,7 +61,8 @@ def description(self) -> str: "the i2v conditioning mask on top. The Wan VAE is causal in time, so encoding the whole video once " "and slicing the latents would not be equivalent — each segment restarts the temporal convolution on " "its own slice. A streaming mode would replace this block with one fed segments incrementally. This " - "block should be used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + "block should be used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` object " + "(e.g. `WanAnimate2SegmentLoopWrapper`); it reads the current segment index `k` from the loop scope." ) @property @@ -115,7 +116,8 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) device = components._execution_device latent_height, latent_width = block_state.reference_image_latents.shape[-2:] @@ -134,10 +136,11 @@ def __call__(self, components, block_state: BlockState, k: int): ).to(block_state.driving_video_latents.dtype) block_state.driving_video_condition = torch.cat([condition_mask, block_state.driving_video_latents[0]], dim=0) - return components, block_state + self.set_block_state(state, block_state) + return components, state -class WanAnimate2SegmentPrevFramesStep(ModularPipelineBlocks): +class WanAnimate2SegmentPrevFramesStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -146,8 +149,9 @@ def description(self) -> str: "Step within the segment loop that builds the generation-side conditioning tensor `reference_latents`: the previous " "segment's tail frames (zeros for the first segment) are VAE-encoded, masked, and stacked under the " "reference half `reference_image_latents`. This is how motion continuity crosses segment boundaries — in pixel space, " - "not latent space. This block should be used to compose the `sub_blocks` attribute of " - "`WanAnimate2SegmentLoopWrapper`." + "not latent space. This block should be used to compose the `sub_blocks` attribute of an " + "`IterativePipelineBlocks` object (e.g. `WanAnimate2SegmentLoopWrapper`); it reads the current segment " + "index `k` from the loop scope." ) @property @@ -165,6 +169,11 @@ def inputs(self) -> list[InputParam]: type_hint=torch.Tensor, description="i2v mask + reference image latents `[20, 1, latent_height, latent_width]`, from the image VAE encoder step", ), + InputParam( + "out_frames", + type_hint=torch.Tensor, + description="The previous segment's decoded frames on device, written by the decode step of the previous iteration; `None` for the first segment", + ), InputParam( "segment_frame_length", type_hint=int, @@ -190,9 +199,8 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): - # `block_state.out_frames` is seeded by the loop wrapper and written by the decode step of the - # previous iteration. + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) device = components._execution_device latent_height, latent_width = block_state.reference_image_latents.shape[-2:] @@ -226,10 +234,11 @@ def __call__(self, components, block_state: BlockState, k: int): [block_state.reference_image_latents, prev_segment_cond_latents], dim=1 ) - return components, block_state + self.set_block_state(state, block_state) + return components, state -class WanAnimate2SegmentPrepareStep(ModularPipelineBlocks): +class WanAnimate2SegmentPrepareStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -237,7 +246,7 @@ def description(self) -> str: return ( "Step within the segment loop that draws this segment's initial noise and allocates a fresh KV cache " "for the reference-extraction pass. This block should be used to compose the `sub_blocks` attribute " - "of `WanAnimate2SegmentLoopWrapper`." + "of an `IterativePipelineBlocks` object (e.g. `WanAnimate2SegmentLoopWrapper`)." ) @property @@ -270,7 +279,8 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) device = components._execution_device block_state.latents = randn_tensor( @@ -286,10 +296,11 @@ def __call__(self, components, block_state: BlockState, k: int): ) block_state.kv_cache = WanAnimate2KVCache(components.transformer.config.num_layers) - return components, block_state + self.set_block_state(state, block_state) + return components, state -class WanAnimate2SegmentSchedulerResetStep(ModularPipelineBlocks): +class WanAnimate2SegmentSchedulerResetStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -297,7 +308,8 @@ def description(self) -> str: return ( "Step within the segment loop that resets the scheduler: each segment is an independent denoising " "trajectory, so the solver state and timesteps are re-prepared per segment. This block should be used " - "to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + "to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` object " + "(e.g. `WanAnimate2SegmentLoopWrapper`)." ) @property @@ -319,16 +331,18 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) device = components._execution_device components.scheduler.set_timesteps(block_state.num_inference_steps, device=device) block_state.timesteps = components.scheduler.timesteps - return components, block_state + self.set_block_state(state, block_state) + return components, state -class WanAnimate2RefExtractStep(ModularPipelineBlocks): +class WanAnimate2RefExtractStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -337,7 +351,8 @@ def description(self) -> str: "Step within the segment loop that runs the transformer's reference-extraction pass " '(`kv_cache_mode="extract"`): the driving-video segment is encoded once and every layer\'s reference ' "K/V is stored in the KV cache, which the denoising forwards then attend over. This block should be " - "used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + "used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` object " + "(e.g. `WanAnimate2SegmentLoopWrapper`)." ) @property @@ -400,7 +415,8 @@ def inputs(self) -> list[InputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) device = components._execution_device transformer_dtype = components.transformer.dtype @@ -417,31 +433,33 @@ def __call__(self, components, block_state: BlockState, k: int): offset_grid_sizes=block_state.grid_sizes_ref, ) - return components, block_state + self.set_block_state(state, block_state) + return components, state # ======================================== -# Inner Denoising Blocks +# Denoising Loop Steps # ======================================== -class WanAnimate2SegmentDenoiseInner(ModularPipelineBlocks): +class WanAnimate2LoopDenoiser(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property def description(self) -> str: return ( - "Inner timestep loop that denoises one segment with guidance, attending over the segment's cached " - "reference K/V. The unconditional branch passes `is_uncondtion=True` to the transformer (it skips a " - "dedicated layer on that branch), routed through the guider as a per-branch input. This block should " - "be used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + "Step within the segment's denoising loop that predicts the noise with guidance, attending over the " + "segment's cached reference K/V. The unconditional branch passes `is_uncondtion=True` to the " + "transformer (it skips a dedicated layer on that branch), routed through the guider as a per-branch " + "input. This block should be used to compose the `sub_blocks` attribute of an " + "`IterativePipelineBlocks` object (e.g. `WanAnimate2DenoiseLoopWrapper`); it reads the current " + "timestep `t` and step index `i` from the loop scope." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", WanAnimate2Transformer3DModel), - ComponentSpec("scheduler", SchedulerMixin), ComponentSpec( "guider", ClassifierFreeGuidance, @@ -471,19 +489,7 @@ def inputs(self) -> list[InputParam]: type_hint=WanAnimate2KVCache, description="Per-segment cache holding every layer's reference K/V", ), - InputParam( - "timesteps", - required=True, - type_hint=torch.Tensor, - description="This segment's denoising timesteps", - ), InputParam.template("num_inference_steps", default=40), - InputParam( - "num_segments", - required=True, - type_hint=int, - description="Total number of segments in the driving video, from the video preprocess step", - ), InputParam( "max_seq_len", required=True, @@ -514,7 +520,6 @@ def inputs(self) -> list[InputParam]: type_hint=int, description="The resolved frame width in pixels", ), - InputParam.template("generator"), InputParam.template("prompt_embeds"), InputParam.template("negative_prompt_embeds"), InputParam.template("denoiser_input_fields"), @@ -523,11 +528,12 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam.template("latents"), + OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted noise for this step"), ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) transformer_dtype = components.transformer.dtype guider_inputs = { @@ -544,68 +550,53 @@ def __call__(self, components, block_state: BlockState, k: int): if name in transformer_args and name not in guider_inputs } - with tqdm( - total=len(block_state.timesteps), desc=f"Segment {k + 1}/{block_state.num_segments}" - ) as progress_bar: - for i, t in enumerate(block_state.timesteps): - timestep = torch.stack([t]) - - components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) - guider_state = components.guider.prepare_inputs(guider_inputs) - - for guider_state_batch in guider_state: - components.guider.prepare_models(components.transformer) - - guider_state_batch.noise_pred = components.transformer( - [block_state.latents.to(transformer_dtype)], - timestep=timestep, - encoder_hidden_states=[guider_state_batch.encoder_hidden_states[0].to(transformer_dtype)], - condition_latents=[block_state.reference_latents.to(transformer_dtype)], - kv_cache=block_state.kv_cache, - kv_cache_mode="cached", - seq_len=block_state.max_seq_len, - reference_grid_sizes=block_state.grid_sizes_ref, - origin_len=block_state.segment_frame_length, - origin_area=[block_state.height, block_state.width], - is_uncondtion=guider_state_batch.is_uncondtion, - **shared_kwargs, - ).sample[0] - - components.guider.cleanup_models(components.transformer) - - noise_pred = components.guider(guider_state)[0] - - latents = components.scheduler.step( - noise_pred.unsqueeze(0), - t, - block_state.latents.unsqueeze(0), - return_dict=False, - generator=block_state.generator, - )[0] - block_state.latents = latents.squeeze(0) + timestep = torch.stack([t]) - progress_bar.update() + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + + guider_state_batch.noise_pred = components.transformer( + [block_state.latents.to(transformer_dtype)], + timestep=timestep, + encoder_hidden_states=[guider_state_batch.encoder_hidden_states[0].to(transformer_dtype)], + condition_latents=[block_state.reference_latents.to(transformer_dtype)], + kv_cache=block_state.kv_cache, + kv_cache_mode="cached", + seq_len=block_state.max_seq_len, + reference_grid_sizes=block_state.grid_sizes_ref, + origin_len=block_state.segment_frame_length, + origin_area=[block_state.height, block_state.width], + is_uncondtion=guider_state_batch.is_uncondtion, + **shared_kwargs, + ).sample[0] - return components, block_state + components.guider.cleanup_models(components.transformer) + block_state.noise_pred = components.guider(guider_state)[0] -class WanAnimate2DistilledSegmentDenoiseInner(WanAnimate2SegmentDenoiseInner): + self.set_block_state(state, block_state) + return components, state + + +class WanAnimate2DistilledLoopDenoiser(WanAnimate2LoopDenoiser): model_name = "wan-animate-2-distilled" @property def description(self) -> str: return ( - "Inner timestep loop that denoises one segment for the distilled model, which is trained for few-step " - "sampling without classifier-free guidance — the guider defaults to `guidance_scale=1.0`, so only the " - "conditional branch runs. This block should be used to compose the `sub_blocks` attribute of " - "`WanAnimate2SegmentLoopWrapper`." + "Step within the segment's denoising loop that predicts the noise for the distilled model, which is " + "trained for few-step sampling without classifier-free guidance — the guider defaults to " + "`guidance_scale=1.0`, so only the conditional branch runs. This block should be used to compose the " + "`sub_blocks` attribute of an `IterativePipelineBlocks` object (e.g. `WanAnimate2DenoiseLoopWrapper`)." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", WanAnimate2Transformer3DModel), - ComponentSpec("scheduler", SchedulerMixin), ComponentSpec( "guider", ClassifierFreeGuidance, @@ -615,22 +606,173 @@ def expected_components(self) -> list[ComponentSpec]: ] +class WanAnimate2LoopAfterDenoiser(ModularLoopPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step within the segment's denoising loop that updates the latents after denoising. " + "This block should be used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` " + "object (e.g. `WanAnimate2DenoiseLoopWrapper`); it reads `noise_pred` and the current timestep `t` " + "from the loop scope." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", SchedulerMixin), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="This segment's latents", + ), + InputParam( + "noise_pred", + required=True, + type_hint=torch.Tensor, + description="The predicted noise for this step", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("latents"), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + + latents = components.scheduler.step( + block_state.noise_pred.unsqueeze(0), + t, + block_state.latents.unsqueeze(0), + return_dict=False, + generator=block_state.generator, + )[0] + block_state.latents = latents.squeeze(0) + + self.set_block_state(state, block_state) + return components, state + + +class WanAnimate2DenoiseLoopWrapper(IterativePipelineBlocks): + model_name = "wan-animate-2" + + @property + def loop_variables(self) -> list[str]: + return ["i", "t"] + + @property + def description(self) -> str: + return ( + "Pipeline block that iteratively denoises one segment's latents over `timesteps`, attending over the " + "segment's cached reference K/V. The specific steps within each iteration can be customized with the " + "`sub_blocks` attribute. It runs inside the segment loop and reads the current segment index `k` from " + "the loop scope." + ) + + @property + def inputs(self) -> list[InputParam]: + inputs = super().inputs + names = {param.name for param in inputs} + # inputs consumed by the loop logic itself, on top of what the sub-blocks declare + loop_inputs = [ + InputParam( + "timesteps", + required=True, + type_hint=torch.Tensor, + description="This segment's denoising timesteps", + ), + InputParam( + "num_segments", + required=True, + type_hint=int, + description="Total number of segments in the driving video, from the video preprocess step", + ), + ] + return [param for param in loop_inputs if param.name not in names] + inputs + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) + + with tqdm( + total=len(block_state.timesteps), desc=f"Segment {k + 1}/{block_state.num_segments}" + ) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + progress_bar.update() + + return components, state + + @torch.no_grad() + def stream(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + return components, state + + +class WanAnimate2SegmentDenoiseStep(WanAnimate2DenoiseLoopWrapper): + block_classes = [WanAnimate2LoopDenoiser, WanAnimate2LoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises one segment's latents with guidance, attending over the " + "segment's cached reference K/V. \n" + "Its loop logic is defined in `WanAnimate2DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `WanAnimate2LoopDenoiser`\n" + " - `WanAnimate2LoopAfterDenoiser`\n" + ) + + +class WanAnimate2DistilledSegmentDenoiseStep(WanAnimate2DenoiseLoopWrapper): + model_name = "wan-animate-2-distilled" + + block_classes = [WanAnimate2DistilledLoopDenoiser, WanAnimate2LoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises one segment's latents for the distilled model, which is " + "trained for few-step sampling without classifier-free guidance. \n" + "Its loop logic is defined in `WanAnimate2DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `WanAnimate2DistilledLoopDenoiser`\n" + " - `WanAnimate2LoopAfterDenoiser`\n" + ) + + # ======================================== # Post-Denoise # ======================================== -class WanAnimate2SegmentDecodeStep(ModularPipelineBlocks): +class WanAnimate2SegmentDecodeStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property def description(self) -> str: return ( "Step within the segment loop that VAE-decodes the denoised segment. Decoding happens inside the loop " - "because the next segment conditions on this segment's decoded pixels. Finished frames move to CPU and " - "the per-segment KV cache and latents are freed — holding them across segments fragments the " - "allocator enough to OOM at high resolution. This block should be used to compose the `sub_blocks` " - "attribute of `WanAnimate2SegmentLoopWrapper`." + "because the next segment conditions on this segment's decoded pixels. The per-segment KV cache and " + "latents are freed — holding them across segments fragments the allocator enough to OOM at high " + "resolution. This block should be used to compose the `sub_blocks` attribute of an " + "`IterativePipelineBlocks` object (e.g. `WanAnimate2SegmentLoopWrapper`)." ) @property @@ -673,7 +815,9 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) + latents = block_state.latents.to(torch.float32) # The first latent frame is the reference image's slot, not video content. out_frames = decode_vae(components.vae, latents[:, 1:]) @@ -681,7 +825,6 @@ def __call__(self, components, block_state: BlockState, k: int): if k > 0: out_frames = out_frames[:, :, block_state.prev_segment_conditioning_frames :] - block_state.segment_frames.append(out_frames.cpu()) block_state.out_frames = out_frames block_state.kv_cache.clear() @@ -689,7 +832,8 @@ def __call__(self, components, block_state: BlockState, k: int): block_state.latents = None torch.cuda.empty_cache() - return components, block_state + self.set_block_state(state, block_state) + return components, state # ======================================== @@ -697,9 +841,13 @@ def __call__(self, components, block_state: BlockState, k: int): # ======================================== -class WanAnimate2SegmentLoopWrapper(LoopSequentialPipelineBlocks): +class WanAnimate2SegmentLoopWrapper(IterativePipelineBlocks): model_name = "wan-animate-2" + @property + def loop_variables(self) -> list[str]: + return ["k"] + @property def description(self) -> str: return ( @@ -709,8 +857,13 @@ def description(self) -> str: ) @property - def loop_inputs(self) -> list[InputParam]: - return [ + def inputs(self) -> list[InputParam]: + # `out_frames` is loop-carried — written by the decode step of each iteration and read by the prev-frames + # step of the next — never user-provided, so it is removed from the aggregated inputs. + inputs = [param for param in super().inputs if param.name != "out_frames"] + names = {param.name for param in inputs} + # inputs consumed by the loop logic itself, on top of what the sub-blocks declare + loop_inputs = [ InputParam( "num_segments", required=True, @@ -718,10 +871,12 @@ def loop_inputs(self) -> list[InputParam]: description="Total number of segments in the driving video, from the video preprocess step", ), ] + return [param for param in loop_inputs if param.name not in names] + inputs @property - def loop_intermediate_outputs(self) -> list[OutputParam]: - return [ + def intermediate_outputs(self) -> list[OutputParam]: + # produced by the loop logic itself, which collects each segment's decoded frames + return super().intermediate_outputs + [ OutputParam( "segment_frames", type_hint=list[torch.Tensor], @@ -730,19 +885,29 @@ def loop_intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState) -> PipelineState: + def __call__(self, components, state: PipelineState): block_state = self.get_block_state(state) - # Seed the loop-carried state: `segment_frames` collects each segment's decoded frames (the decode step - # appends to it); `out_frames` is the previous segment's decoded frames — written by the decode step, read - # by the prev-frames step of the next iteration. `None` marks "no previous segment" for the first iteration. - block_state.segment_frames = [] - block_state.out_frames = None + # `segment_frames` collects each segment's decoded frames on CPU; `out_frames` (this segment's frames, + # on device) stays in the state for the prev-frames step of the next iteration to condition on. + segment_frames = [] + for k in range(block_state.num_segments): + components, state = self.loop_step(components, state, k=k) + segment_frames.append(state.get("out_frames").cpu()) + state.set("segment_frames", segment_frames) + + return components, state + @torch.no_grad() + def stream(self, components, state: PipelineState): + block_state = self.get_block_state(state) + + segment_frames = [] for k in range(block_state.num_segments): - components, block_state = self.loop_step(components, block_state, k=k) + components, state = yield from self.stream_step(components, state, k=k) + segment_frames.append(state.get("out_frames").cpu()) + state.set("segment_frames", segment_frames) - self.set_block_state(state, block_state) return components, state @@ -758,7 +923,7 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, WanAnimate2RefExtractStep, - WanAnimate2SegmentDenoiseInner, + WanAnimate2SegmentDenoiseStep, WanAnimate2SegmentDecodeStep, ] block_names = [ @@ -776,7 +941,7 @@ def description(self) -> str: return ( "Segment denoise step that iterates over the driving video's segments.\n" "At each segment: vae_encoder -> prev_frames -> prepare -> scheduler_reset -> ref_extract -> " - "denoise_inner -> decode." + "denoise_inner (a nested denoising loop over this segment's timesteps) -> decode." ) @@ -789,7 +954,7 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, WanAnimate2RefExtractStep, - WanAnimate2DistilledSegmentDenoiseInner, + WanAnimate2DistilledSegmentDenoiseStep, WanAnimate2SegmentDecodeStep, ] block_names = [ @@ -807,5 +972,6 @@ def description(self) -> str: return ( "Segment denoise step for the distilled model that iterates over the driving video's segments.\n" "At each segment: vae_encoder -> prev_frames -> prepare -> scheduler_reset -> ref_extract -> " - "denoise_inner (no classifier-free guidance) -> decode." + "denoise_inner (a nested denoising loop over this segment's timesteps, no classifier-free guidance) " + "-> decode." ) diff --git a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py index f77eb378c15b..1d569b328701 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py @@ -175,8 +175,6 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): The reference conditioning tensor `[20, 1, latent_height, latent_width]`; provides the latent grid driving_video_pixels (`Tensor`): The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step - num_segments (`int`): - Total number of segments in the driving video, from the video preprocess step effective_segment (`int`): Frames each segment advances: `segment_frame_length - prev_segment_conditioning_frames`, from the video preprocess step @@ -190,6 +188,8 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): CLIP vision features of the driving video's first frame prompt_ref_embeds (`Tensor`): Text embeddings of the reference prompt, guiding the reference-extraction pass + num_segments (`int`): + Total number of segments in the driving video, from the video preprocess step height (`int`): The resolved frame height in pixels width (`int`): diff --git a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py index 8eab815897da..2f531ead16bd 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py @@ -175,8 +175,6 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): The reference conditioning tensor `[20, 1, latent_height, latent_width]`; provides the latent grid driving_video_pixels (`Tensor`): The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step - num_segments (`int`): - Total number of segments in the driving video, from the video preprocess step effective_segment (`int`): Frames each segment advances: `segment_frame_length - prev_segment_conditioning_frames`, from the video preprocess step @@ -190,6 +188,8 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): CLIP vision features of the driving video's first frame prompt_ref_embeds (`Tensor`): Text embeddings of the reference prompt, guiding the reference-extraction pass + num_segments (`int`): + Total number of segments in the driving video, from the video preprocess step height (`int`): The resolved frame height in pixels width (`int`): From 3282034a05e400b81cea52ab06b353fb078540d0 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 19 Aug 2026 23:36:19 +0000 Subject: [PATCH 09/10] Port LTX-2/2.5 to IterativePipelineBlocks with streaming - the 7 loop steps become ModularLoopPipelineBlocks operating on the shared PipelineState; the intra-iteration dataflow is now declared (before-denoiser outputs latent_model_input/timesteps, denoiser outputs noise_pred_video/ noise_pred_audio, after-denoisers output latents/audio_latents) and stays satisfied within the loop, so pipeline inputs are unchanged - LTX2DenoiseLoopWrapper is an IterativePipelineBlocks (loop variables `i`, `t`) and implements `stream`, so LTX-2 and LTX-2.5 stream with both video and audio latents live in every event; test_stream_matches_call now runs for all 8 ltx2/ltx25 testers - regenerate ltx2 modular_blocks docstrings (timesteps correctly optional at the core-step level, produced by set_timesteps) Co-Authored-By: Claude Fable 5 --- .../modular_pipelines/ltx2/denoise.py | 242 +++++++++++++++--- .../ltx2/modular_blocks_ltx2.py | 4 +- .../ltx2/modular_blocks_ltx25.py | 2 +- 3 files changed, 208 insertions(+), 40 deletions(-) diff --git a/src/diffusers/modular_pipelines/ltx2/denoise.py b/src/diffusers/modular_pipelines/ltx2/denoise.py index 5cc5a4e57abc..7e3394fd9818 100644 --- a/src/diffusers/modular_pipelines/ltx2/denoise.py +++ b/src/diffusers/modular_pipelines/ltx2/denoise.py @@ -22,12 +22,11 @@ from ...models import LTX2VideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ..modular_pipeline import ( - BlockState, - LoopSequentialPipelineBlocks, - ModularPipelineBlocks, + IterativePipelineBlocks, + ModularLoopPipelineBlocks, PipelineState, ) -from ..modular_pipeline_utils import ComponentSpec, InputParam +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .guider import LTX2Guidance @@ -71,7 +70,7 @@ def _unpack_latents( return latents -class LTX2LoopBeforeDenoiser(ModularPipelineBlocks): +class LTX2LoopBeforeDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -93,17 +92,32 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latent_model_input", type_hint=torch.Tensor, description="Video latents cast to model dtype"), + OutputParam( + "audio_latent_model_input", type_hint=torch.Tensor, description="Audio latents cast to model dtype" + ), + OutputParam("video_timestep", type_hint=torch.Tensor, description="This step's video timestep"), + OutputParam("audio_timestep", type_hint=torch.Tensor, description="This step's audio timestep"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + block_state.latent_model_input = block_state.latents.to(block_state.dtype) block_state.audio_latent_model_input = block_state.audio_latents.to(block_state.dtype) timestep = t.expand(block_state.latents.shape[0]) block_state.video_timestep = timestep block_state.audio_timestep = timestep - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2Image2VideoLoopBeforeDenoiser(ModularPipelineBlocks): +class LTX2Image2VideoLoopBeforeDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -129,17 +143,32 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latent_model_input", type_hint=torch.Tensor, description="Video latents cast to model dtype"), + OutputParam( + "audio_latent_model_input", type_hint=torch.Tensor, description="Audio latents cast to model dtype" + ), + OutputParam("video_timestep", type_hint=torch.Tensor, description="This step's masked video timestep"), + OutputParam("audio_timestep", type_hint=torch.Tensor, description="This step's audio timestep"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + block_state.latent_model_input = block_state.latents.to(block_state.dtype) block_state.audio_latent_model_input = block_state.audio_latents.to(block_state.dtype) timestep = t.expand(block_state.latents.shape[0]) block_state.video_timestep = timestep.unsqueeze(-1) * (1 - block_state.conditioning_mask) block_state.audio_timestep = timestep - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2ConditionLoopBeforeDenoiser(ModularPipelineBlocks): +class LTX2ConditionLoopBeforeDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -171,14 +200,29 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latent_model_input", type_hint=torch.Tensor, description="Video latents cast to model dtype"), + OutputParam( + "audio_latent_model_input", type_hint=torch.Tensor, description="Audio latents cast to model dtype" + ), + OutputParam("video_timestep", type_hint=torch.Tensor, description="This step's masked video timestep"), + OutputParam("audio_timestep", type_hint=torch.Tensor, description="This step's audio timestep"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + block_state.latent_model_input = block_state.latents.to(block_state.dtype) block_state.audio_latent_model_input = block_state.audio_latents.to(block_state.dtype) timestep = t.expand(block_state.latents.shape[0]) block_state.video_timestep = timestep.unsqueeze(-1) * (1 - block_state.conditioning_mask.squeeze(-1)) block_state.audio_timestep = timestep - return components, block_state + + self.set_block_state(state, block_state) + return components, state # Default per-pass conditioning map for `LTX2LoopDenoiser`: transformer argument -> block-state attribute names @@ -212,7 +256,7 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) } -class LTX2LoopDenoiser(ModularPipelineBlocks): +class LTX2LoopDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" def __init__(self, guider_input_fields: dict[str, Any] = _DEFAULT_GUIDER_INPUT_FIELDS): @@ -293,6 +337,30 @@ def inputs(self) -> list[InputParam]: required=True, description="Packed noisy audio latents to denoise.", ), + InputParam( + "latent_model_input", + type_hint=torch.Tensor, + required=True, + description="Video latents cast to model dtype, from the before-denoiser step.", + ), + InputParam( + "audio_latent_model_input", + type_hint=torch.Tensor, + required=True, + description="Audio latents cast to model dtype, from the before-denoiser step.", + ), + InputParam( + "video_timestep", + type_hint=torch.Tensor, + required=True, + description="This step's video timestep, from the before-denoiser step.", + ), + InputParam( + "audio_timestep", + type_hint=torch.Tensor, + required=True, + description="This step's audio timestep, from the before-denoiser step.", + ), InputParam("audio_scheduler", required=True), # `audio_num_frames`, `video_coords`, `audio_coords` arrive tagged `denoiser_input_fields` upstream and # are collected from the tagged dict (filtered against the transformer signature) in `__call__`. @@ -336,8 +404,21 @@ def inputs(self) -> list[InputParam]: ) return inputs + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "noise_pred_video", type_hint=torch.Tensor, description="Guided x0 prediction for the video latents" + ), + OutputParam( + "noise_pred_audio", type_hint=torch.Tensor, description="Guided x0 prediction for the audio latents" + ), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 latent_height = block_state.height // components.vae_spatial_compression_ratio latent_width = block_state.width // components.vae_spatial_compression_ratio @@ -425,10 +506,12 @@ def _combine(guider, field): block_state.noise_pred_video = _combine(components.guider, "video_pred") block_state.noise_pred_audio = _combine(components.audio_guider, "audio_pred") - return components, block_state + self.set_block_state(state, block_state) + return components, state -class LTX2LoopAfterDenoiser(ModularPipelineBlocks): + +class LTX2LoopAfterDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -450,10 +533,31 @@ def inputs(self) -> list[InputParam]: description="Packed noisy audio latents to denoise.", ), InputParam("audio_scheduler", required=True), + InputParam( + "noise_pred_video", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the video latents, from the denoiser step.", + ), + InputParam( + "noise_pred_audio", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the audio latents, from the denoiser step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="The denoised video latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="The denoised audio latents"), ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + noise_pred_video = convert_x0_to_velocity( block_state.latents, block_state.noise_pred_video, i, components.scheduler ) @@ -464,10 +568,12 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) block_state.audio_latents = block_state.audio_scheduler.step( noise_pred_audio, t, block_state.audio_latents, return_dict=False )[0] - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2Image2VideoLoopAfterDenoiser(ModularPipelineBlocks): +class LTX2Image2VideoLoopAfterDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -492,6 +598,18 @@ def inputs(self) -> list[InputParam]: description="Packed noisy audio latents to denoise.", ), InputParam("audio_scheduler", required=True), + InputParam( + "noise_pred_video", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the video latents, from the denoiser step.", + ), + InputParam( + "noise_pred_audio", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the audio latents, from the denoiser step.", + ), InputParam.template("height", default=512), InputParam.template("width", default=704), InputParam( @@ -505,8 +623,17 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="The denoised video latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="The denoised audio latents"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + spatial_patch = components.transformer_spatial_patch_size temporal_patch = components.transformer_temporal_patch_size latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 @@ -534,10 +661,12 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) block_state.audio_latents = block_state.audio_scheduler.step( noise_pred_audio, t, block_state.audio_latents, return_dict=False )[0] - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2ConditionLoopAfterDenoiser(ModularPipelineBlocks): +class LTX2ConditionLoopAfterDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -564,6 +693,18 @@ def inputs(self) -> list[InputParam]: description="Packed noisy audio latents to denoise.", ), InputParam("audio_scheduler", required=True), + InputParam( + "noise_pred_video", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the video latents, from the denoiser step.", + ), + InputParam( + "noise_pred_audio", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the audio latents, from the denoiser step.", + ), InputParam( "conditioning_mask", type_hint=torch.Tensor, @@ -578,8 +719,17 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="The denoised video latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="The denoised audio latents"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + # Conditioning strengths run from 0 (always use the denoised sample) to 1 (always use the condition), with # intermediate values specifying how strongly to follow the condition. Applied in x0 space, not velocity # space (which is what the transformer outputs). @@ -596,12 +746,18 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) block_state.audio_latents = block_state.audio_scheduler.step( noise_pred_audio, t, block_state.audio_latents, return_dict=False )[0] - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2DenoiseLoopWrapper(LoopSequentialPipelineBlocks): +class LTX2DenoiseLoopWrapper(IterativePipelineBlocks): model_name = "ltx2" + @property + def loop_variables(self) -> list[str]: + return ["i", "t"] + @property def description(self) -> str: return ( @@ -610,36 +766,48 @@ def description(self) -> str: ) @property - def loop_expected_components(self) -> list[ComponentSpec]: - return [ - ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), - ComponentSpec("transformer", LTX2VideoTransformer3DModel), - ] + def expected_components(self) -> list[ComponentSpec]: + expected_components = super().expected_components + # the loop logic itself reads `scheduler.order` for the warmup-step computation + scheduler = ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler) + if scheduler not in expected_components: + expected_components.append(scheduler) + return expected_components @property - def loop_inputs(self) -> list[InputParam]: - return [ + def inputs(self) -> list[InputParam]: + inputs = super().inputs + names = {param.name for param in inputs} + # inputs consumed by the loop logic itself, on top of what the sub-blocks declare + loop_inputs = [ InputParam("timesteps", type_hint=torch.Tensor, required=True), InputParam.template("num_inference_steps", required=True), ] + return [param for param in loop_inputs if param.name not in names] + inputs @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - block_state.num_warmup_steps = max( + num_warmup_steps = max( len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 ) with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: for i, t in enumerate(block_state.timesteps): - components, block_state = self.loop_step(components, block_state, i=i, t=t) + components, state = self.loop_step(components, state, i=i, t=t) if i == len(block_state.timesteps) - 1 or ( - (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 + (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0 ): progress_bar.update() - self.set_block_state(state, block_state) + return components, state + + @torch.no_grad() + def stream(self, components, state: PipelineState): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) return components, state diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py index 86428328a5a6..1cba0784856e 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py @@ -975,7 +975,7 @@ class LTX2AutoCoreDenoiseStep(ConditionalPipelineBlocks): Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. num_inference_steps (`int`): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. audio_latents (`Tensor`): Optional pre-encoded audio latents; random noise is used when not provided. @@ -1822,7 +1822,7 @@ class LTX2AutoBlocks(SequentialPipelineBlocks): Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. num_inference_steps (`int`): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. audio_latents (`Tensor`): Optional pre-encoded audio latents; random noise is used when not provided. diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py index 7c77aba94a74..dd4d7c003d5d 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py @@ -318,7 +318,7 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. num_inference_steps (`int`): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. audio_latents (`Tensor`): Optional pre-encoded audio latents; random noise is used when not provided. From e758ae0107f86de67cce2c78d8559905bb352871 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sat, 22 Aug 2026 00:38:31 +0000 Subject: [PATCH 10/10] IterativePipelineBlocks: declare loop-level inputs/outputs; scope block state to the loop logic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Re-add `loop_inputs` / `loop_intermediate_outputs` on IterativePipelineBlocks. The base class merges them into the aggregated `inputs` / `intermediate_outputs`, and `get_block_state` / `set_block_state` on a loop block now read and write only these — sub-block values live in the pipeline state, not in the loop's (pre-loop) snapshot. This lets a loop write its own outputs through `set_block_state` like a leaf, and removes the hand-written inputs merge from every wrapper. - flux2 / ltx2 / wan-animate-2 wrappers and test fixtures move their loop inputs to `loop_inputs`; the wan segment loop declares `segment_frames` in `loop_intermediate_outputs` and writes it via `set_block_state`. - wan-animate-2: `out_frames` (loop-carried, previous segment's frames) is seeded as `None` by the prepare-segments step instead of being filtered out of the loop's inputs, so the sequential aggregation hides it at pipeline level on its own. - tests: loop outputs aggregate; loop block state holds only `loop_inputs`; a loop output written through `set_block_state` leaves sub-block outputs untouched. Co-Authored-By: Claude Fable 5 --- .../en/modular_diffusers/modular_pipeline.md | 4 ++ .../modular_pipelines/flux2/denoise.py | 8 +-- .../modular_pipelines/ltx2/denoise.py | 8 +-- .../modular_pipelines/modular_pipeline.py | 54 +++++++++++++++++-- .../wan_animate_2/before_denoise.py | 7 +++ .../wan_animate_2/denoise.py | 36 +++++-------- .../test_iterative_pipeline_blocks.py | 50 ++++++++++++++--- 7 files changed, 123 insertions(+), 44 deletions(-) diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 93013d743aaa..1908065f760c 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -421,6 +421,10 @@ class DenoiseLoop(IterativePipelineBlocks): def loop_variables(self): return ["i", "t"] + @property + def loop_inputs(self): # what the loop logic itself reads; `get_block_state` returns exactly these + return [InputParam("timesteps", required=True)] + @torch.no_grad() def __call__(self, components, state): block_state = self.get_block_state(state) diff --git a/src/diffusers/modular_pipelines/flux2/denoise.py b/src/diffusers/modular_pipelines/flux2/denoise.py index fa40cf7c7523..b36b06fcff19 100644 --- a/src/diffusers/modular_pipelines/flux2/denoise.py +++ b/src/diffusers/modular_pipelines/flux2/denoise.py @@ -467,11 +467,8 @@ def description(self) -> str: ) @property - def inputs(self) -> list[InputParam]: - inputs = super().inputs - names = {param.name for param in inputs} - # inputs consumed by the loop logic itself, on top of what the sub-blocks declare - loop_inputs = [ + def loop_inputs(self) -> list[InputParam]: + return [ InputParam( "timesteps", required=True, @@ -485,7 +482,6 @@ def inputs(self) -> list[InputParam]: description="The number of inference steps to use for the denoising process.", ), ] - return [param for param in loop_inputs if param.name not in names] + inputs @torch.no_grad() def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: diff --git a/src/diffusers/modular_pipelines/ltx2/denoise.py b/src/diffusers/modular_pipelines/ltx2/denoise.py index 7e3394fd9818..972a40b48d83 100644 --- a/src/diffusers/modular_pipelines/ltx2/denoise.py +++ b/src/diffusers/modular_pipelines/ltx2/denoise.py @@ -775,15 +775,11 @@ def expected_components(self) -> list[ComponentSpec]: return expected_components @property - def inputs(self) -> list[InputParam]: - inputs = super().inputs - names = {param.name for param in inputs} - # inputs consumed by the loop logic itself, on top of what the sub-blocks declare - loop_inputs = [ + def loop_inputs(self) -> list[InputParam]: + return [ InputParam("timesteps", type_hint=torch.Tensor, required=True), InputParam.template("num_inference_steps", required=True), ] - return [param for param in loop_inputs if param.name not in names] + inputs @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index d3e703d6bd85..5d14dfd4e10b 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -1497,9 +1497,12 @@ def __call__(self, components, state, k): # accepts the OUTER chunk loop's vari return components, state ``` - Sub-block outputs are written to the pipeline state as usual and persist after the loop. If the loop logic in - `__call__` itself consumes inputs (e.g. `timesteps`) or uses components (e.g. the scheduler) beyond what the - sub-blocks declare, override the aggregated `inputs` / `expected_components` / ... properties to add them. + Sub-block outputs are written to the pipeline state as usual and persist after the loop. The loop logic's own + inputs (e.g. `timesteps`) and outputs are declared in `loop_inputs` / `loop_intermediate_outputs`: they are + surfaced alongside the sub-blocks' in the aggregated `inputs` / `intermediate_outputs`, and they are what + `get_block_state` / `set_block_state` read and write for the loop block itself — sub-block values live in the + pipeline state, not in the loop's block state. A component used by the loop logic itself (e.g. the scheduler) is + added by overriding `expected_components`. Streaming is opt-in: to let `pipe.stream(...)` hand back the live [`PipelineState`] after every iteration, also implement `stream` — the same loop, written as a generator over `stream_step` (which runs one iteration like @@ -1527,6 +1530,51 @@ def loop_variables(self) -> list[str]: """Names of the loop variables `loop_step` passes to leaf sub-blocks each iteration (e.g. `["i", "t"]`).""" return [] + @property + def loop_inputs(self) -> list[InputParam]: + """Inputs read by the loop logic in `__call__` itself (e.g. `timesteps`), beyond what the sub-blocks declare.""" + return [] + + @property + def loop_intermediate_outputs(self) -> list[OutputParam]: + """Outputs written to the pipeline state by the loop logic in `__call__` itself.""" + return [] + + @property + def inputs(self) -> list[InputParam]: + inputs = super().inputs + names = {param.name for param in inputs} + return [param for param in self.loop_inputs if param.name not in names] + inputs + + @property + def intermediate_outputs(self) -> list[OutputParam]: + outputs = super().intermediate_outputs + names = {output.name for output in outputs} + return outputs + [output for output in self.loop_intermediate_outputs if output.name not in names] + + def get_block_state(self, state: PipelineState) -> BlockState: + """The loop logic's own inputs (`loop_inputs`); sub-block values are read from the pipeline state.""" + data = {} + for input_param in self.loop_inputs: + value = state.get(input_param.name) + if value is None: + value = input_param.default + if input_param.required and value is None: + raise ValueError(f"Required input '{input_param.name}' is missing") + data[input_param.name] = value + return BlockState(**data) + + def set_block_state(self, state: PipelineState, block_state: BlockState): + """Write the loop logic's own outputs (`loop_intermediate_outputs`) and modified inputs back to the state.""" + for output_param in self.loop_intermediate_outputs: + if not hasattr(block_state, output_param.name): + raise ValueError(f"Intermediate output '{output_param.name}' is missing in block state") + state.set(output_param.name, getattr(block_state, output_param.name), output_param.kwargs_type) + for input_param in self.loop_inputs: + value = getattr(block_state, input_param.name) + if state.get(input_param.name) is not value: + state.set(input_param.name, value, input_param.kwargs_type) + def __init__(self): super().__init__() self._validate_sub_blocks() diff --git a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py index 0ad038e8ccaa..6d84e5d8d20f 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py @@ -79,6 +79,12 @@ def intermediate_outputs(self) -> list[OutputParam]: type_hint=int, description="Packed sequence length of the reference tokens", ), + OutputParam( + "out_frames", + type_hint=torch.Tensor, + description="The previous segment's decoded frames, carried across the segment loop; starts as " + "`None` (the first segment has no previous segment to condition on)", + ), ] @torch.no_grad() @@ -107,6 +113,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: latent_noise_frames = latent_segment_frames + 1 block_state.max_seq_len = int(math.ceil(np.prod([latent_noise_frames, latent_height // 2, latent_width // 2]))) block_state.max_seq_len_ref = int(math.ceil(np.prod(ref_shape) // 4)) + block_state.out_frames = None self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py index 8cd34c7cfc5c..d1798e042277 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -682,11 +682,8 @@ def description(self) -> str: ) @property - def inputs(self) -> list[InputParam]: - inputs = super().inputs - names = {param.name for param in inputs} - # inputs consumed by the loop logic itself, on top of what the sub-blocks declare - loop_inputs = [ + def loop_inputs(self) -> list[InputParam]: + return [ InputParam( "timesteps", required=True, @@ -700,7 +697,6 @@ def inputs(self) -> list[InputParam]: description="Total number of segments in the driving video, from the video preprocess step", ), ] - return [param for param in loop_inputs if param.name not in names] + inputs @torch.no_grad() def __call__(self, components, state: PipelineState, k: int): @@ -857,13 +853,8 @@ def description(self) -> str: ) @property - def inputs(self) -> list[InputParam]: - # `out_frames` is loop-carried — written by the decode step of each iteration and read by the prev-frames - # step of the next — never user-provided, so it is removed from the aggregated inputs. - inputs = [param for param in super().inputs if param.name != "out_frames"] - names = {param.name for param in inputs} - # inputs consumed by the loop logic itself, on top of what the sub-blocks declare - loop_inputs = [ + def loop_inputs(self) -> list[InputParam]: + return [ InputParam( "num_segments", required=True, @@ -871,12 +862,11 @@ def inputs(self) -> list[InputParam]: description="Total number of segments in the driving video, from the video preprocess step", ), ] - return [param for param in loop_inputs if param.name not in names] + inputs @property - def intermediate_outputs(self) -> list[OutputParam]: - # produced by the loop logic itself, which collects each segment's decoded frames - return super().intermediate_outputs + [ + def loop_intermediate_outputs(self) -> list[OutputParam]: + # the loop logic collects each segment's decoded frames + return [ OutputParam( "segment_frames", type_hint=list[torch.Tensor], @@ -890,11 +880,11 @@ def __call__(self, components, state: PipelineState): # `segment_frames` collects each segment's decoded frames on CPU; `out_frames` (this segment's frames, # on device) stays in the state for the prev-frames step of the next iteration to condition on. - segment_frames = [] + block_state.segment_frames = [] for k in range(block_state.num_segments): components, state = self.loop_step(components, state, k=k) - segment_frames.append(state.get("out_frames").cpu()) - state.set("segment_frames", segment_frames) + block_state.segment_frames.append(state.get("out_frames").cpu()) + self.set_block_state(state, block_state) return components, state @@ -902,11 +892,11 @@ def __call__(self, components, state: PipelineState): def stream(self, components, state: PipelineState): block_state = self.get_block_state(state) - segment_frames = [] + block_state.segment_frames = [] for k in range(block_state.num_segments): components, state = yield from self.stream_step(components, state, k=k) - segment_frames.append(state.get("out_frames").cpu()) - state.set("segment_frames", segment_frames) + block_state.segment_frames.append(state.get("out_frames").cpu()) + self.set_block_state(state, block_state) return components, state diff --git a/tests/modular_pipelines/test_iterative_pipeline_blocks.py b/tests/modular_pipelines/test_iterative_pipeline_blocks.py index ad74c649d62c..9054bb0d535e 100644 --- a/tests/modular_pipelines/test_iterative_pipeline_blocks.py +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -118,8 +118,8 @@ def loop_variables(self): return ["i", "t"] @property - def inputs(self): - return [InputParam(name="timesteps", required=True), *super().inputs] + def loop_inputs(self): + return [InputParam(name="timesteps", required=True)] @torch.no_grad() def __call__(self, components, state, k): @@ -168,8 +168,8 @@ def loop_variables(self): return ["k"] @property - def inputs(self): - return [InputParam(name="num_latent_chunk", required=True), *super().inputs] + def loop_inputs(self): + return [InputParam(name="num_latent_chunk", required=True)] @torch.no_grad() def __call__(self, components, state): @@ -179,6 +179,24 @@ def __call__(self, components, state): return components, state +class CollectingChunkLoop(ChunkLoop): + """Chunk loop whose loop logic has an output of its own, written through `set_block_state`.""" + + @property + def loop_intermediate_outputs(self): + return [OutputParam(name="chunk_history")] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.chunk_history = [] + for k in range(block_state.num_latent_chunk): + components, state = self.loop_step(components, state, k=k) + block_state.chunk_history.append(float(state.get("history"))) + self.set_block_state(state, block_state) + return components, state + + class TestIterativePipelineBlocksStructure: def test_inputs_aggregation(self): loop = ChunkLoop() @@ -201,6 +219,12 @@ def test_sub_block_outputs_are_aggregated(self): assert "history" in output_names assert "latent_chunks" in output_names + def test_loop_outputs_are_aggregated(self): + loop = CollectingChunkLoop() + output_names = [o.name for o in loop.intermediate_outputs] + assert "chunk_history" in output_names + assert "history" in output_names + def test_loop_block_can_nest_assembled_blocks(self): # the nested inner loop stays an assembled IterativePipelineBlocks sub-block loop = ChunkLoop() @@ -232,6 +256,20 @@ def test_loop_variables_do_not_leak_into_state(self): # declared sub-block outputs persist after the loop (last iteration's value) assert state.get("noise_pred") is not None + def test_block_state_is_loop_scoped(self): + # the loop's block state holds only the loop logic's own inputs; sub-block values live in the pipeline state + pipe = self._make_pipeline() + state = pipe(num_latent_chunk=2, timesteps=torch.tensor([1.0]), history=torch.tensor(0.0)) + block_state = pipe.blocks.sub_blocks["chunks"].get_block_state(state) + assert block_state.as_dict().keys() == {"num_latent_chunk"} + + def test_loop_output_via_set_block_state(self): + pipe = SequentialPipelineBlocks.from_blocks_dict({"chunks": CollectingChunkLoop()}).init_pipeline() + state = pipe(num_latent_chunk=3, timesteps=torch.tensor([1.0, 2.0]), history=torch.tensor(0.0)) + assert state.get("chunk_history") == [3.0, 7.0, 12.0] + # sub-block outputs are untouched by the loop's own write-back + assert state.get("latent_chunks") == [3.0, 7.0, 12.0] + def test_sub_block_type_is_validated(self): # a regular ModularPipelineBlocks cannot be a loop sub-block: fails at construction class PlainStep(ModularPipelineBlocks): @@ -282,8 +320,8 @@ def loop_variables(self): return ["i", "t"] @property - def inputs(self): - return [InputParam(name="timesteps", required=True), *super().inputs] + def loop_inputs(self): + return [InputParam(name="timesteps", required=True)] @torch.no_grad() def __call__(self, components, state):