From b272422fc5d6bc78578d27a8aeb4b381ec06f820 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 09:34:34 +0000 Subject: [PATCH 01/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] =?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/24] 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/24] 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/24] 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): From 5dadd8b0044563b7abcd95e47a12bf2f249a22d3 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 25 Aug 2026 19:29:26 +0000 Subject: [PATCH 11/24] update docs on interactive loop block --- .ai/modular.md | 70 +++-- docs/source/en/_toctree.yml | 2 + .../api/modular_diffusers/pipeline_blocks.md | 8 + .../api/modular_diffusers/pipeline_states.md | 6 +- .../en/modular_diffusers/custom_blocks.md | 2 +- .../iterative_pipeline_blocks.md | 286 ++++++++++++++++++ .../loop_sequential_pipeline_blocks.md | 3 + .../en/modular_diffusers/modular_pipeline.md | 2 +- docs/source/en/modular_diffusers/overview.md | 3 +- .../sequential_pipeline_blocks.md | 6 + 10 files changed, 365 insertions(+), 23 deletions(-) create mode 100644 docs/source/en/modular_diffusers/iterative_pipeline_blocks.md diff --git a/.ai/modular.md b/.ai/modular.md index 3353b878e0ad..465a55e1957e 100644 --- a/.ai/modular.md +++ b/.ai/modular.md @@ -4,7 +4,8 @@ Shared reference for modular pipeline conventions, patterns, and gotchas. ## Common modular conventions -When adding a new modular pipeline (or reviewing one), skim `src/diffusers/modular_pipelines/qwenimage/`, `src/diffusers/modular_pipelines/flux2/`, `src/diffusers/modular_pipelines/wan/`, and `src/diffusers/modular_pipelines/helios/` first to establish the pattern. Most conventions (file split between `encoders.py` / `before_denoise.py` / `denoise.py` / `decoders.py`, how `expected_components` / `inputs` / `intermediate_outputs` are declared, the denoise-loop wrapping with `LoopSequentialPipelineBlocks`, top-level assembly via `AutoPipelineBlocks` / `SequentialPipelineBlocks` in `modular_blocks_.py`, the `ModularPipeline` subclass shape, the guider-abstracted denoise body, `kwargs_type="denoiser_input_fields"` plumbing) are easiest to internalize by comparison rather than from a fixed list. +# should we include minimax, maybe replace wan? +When adding a new modular pipeline (or reviewing one), skim `src/diffusers/modular_pipelines/qwenimage/`, `src/diffusers/modular_pipelines/flux2/`, `src/diffusers/modular_pipelines/wan/`, and `src/diffusers/modular_pipelines/helios/` first to establish the pattern. Most conventions (file split between `encoders.py` / `before_denoise.py` / `denoise.py` / `decoders.py`, how `expected_components` / `inputs` / `intermediate_outputs` are declared, the denoise-loop wrapping with `IterativePipelineBlocks` (see `flux2/denoise.py`; `wan_animate_2/denoise.py` for a nested chunk loop), top-level assembly via `AutoPipelineBlocks` / `SequentialPipelineBlocks` in `modular_blocks_.py`, the `ModularPipeline` subclass shape, the guider-abstracted denoise body, `kwargs_type="denoiser_input_fields"` plumbing) are easiest to internalize by comparison rather than from a fixed list. ## Running a modular pipeline @@ -51,8 +52,12 @@ Is this a single operation? Does it run multiple blocks in sequence? YES -> SequentialPipelineBlocks - Does it iterate (e.g. chunk loop)? - YES -> LoopSequentialPipelineBlocks + Does it iterate (denoising loop)? + YES -> IterativePipelineBlocks (steps are ModularLoopPipelineBlocks) + Does each iteration run a loop of its own (autoregressive chunk loop: per chunk, a denoise loop)? + YES -> nest them: the inner IterativePipelineBlocks is a step of the outer one, and its + __call__/stream and its steps also take the outer loop variable (see wan_animate_2/denoise.py) + LoopSequentialPipelineBlocks is the legacy loop type — don't use it for new pipelines Does it choose ONE block based on which input is present? Is the selection 1:1 with trigger inputs? @@ -121,25 +126,48 @@ for i, t in enumerate(timesteps): ## Key pattern: Denoising loop -All models use `LoopSequentialPipelineBlocks` for the denoising loop (iterating over timesteps): -```python -class MyModelDenoiseLoopWrapper(LoopSequentialPipelineBlocks): - block_classes = [LoopBeforeDenoiser, LoopDenoiser, LoopAfterDenoiser] -``` +The denoising loop is an `IterativePipelineBlocks` whose steps are `ModularLoopPipelineBlocks` (canonical: `flux2/denoise.py`). The wrapper declares the loop variables it passes to its steps, the inputs its own loop logic reads, and the loop in `__call__`; `stream` is the same loop as a generator over `stream_step` and is what makes `pipe.stream(...)` work — implement both: -Autoregressive video models (e.g. Helios) also use it for an outer chunk loop: ```python -class HeliosChunkDenoiseStep(HeliosChunkLoopWrapper): - block_classes = [ - HeliosChunkHistorySliceStep, - HeliosChunkNoiseGenStep, - HeliosChunkSchedulerResetStep, - HeliosChunkDenoiseInner, - HeliosChunkUpdateStep, - ] +class MyModelDenoiseLoopWrapper(IterativePipelineBlocks): + @property + 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), InputParam.template("num_inference_steps", 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 + + @torch.no_grad() + 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 + + +class MyModelDenoiseStep(MyModelDenoiseLoopWrapper): + block_classes = [MyModelLoopDenoiser, MyModelLoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] ``` -Note: sub-blocks inside `LoopSequentialPipelineBlocks` receive `(components, block_state, i, t)` for denoise loops or `(components, block_state, k)` for chunk loops. +Loop steps are regular blocks (own `inputs` / `intermediate_outputs`, `get_block_state` / `set_block_state` on the full `PipelineState`) whose `__call__` takes the loop variables as arguments — `def __call__(self, components, state, i, t)`. Every step of a loop must accept exactly that loop's variables (validated at construction), even if it ignores some of them. + +Autoregressive video models nest loops: an outer chunk loop (`loop_variables = ["k"]`) whose steps prepare the chunk, run the full inner denoising loop, and update the history (canonical: `wan_animate_2/denoise.py`). The inner loop is a step of the outer one, so its `__call__` / `stream` accept the outer variable — `def __call__(self, components, state, k)` — and pass its own `i`, `t` to its own steps. + +Conventions that fall out of this: +- A value the loop logic itself produces (e.g. the collected per-chunk frames) goes in `loop_intermediate_outputs` and is written with `set_block_state`, like a leaf output. Step outputs are read from the state (`state.get("out_frames")`) — the loop's block state holds only its `loop_inputs` / `loop_intermediate_outputs`. +- A loop-carried value (written by a step at the end of iteration `k`, read by a step at the start of `k + 1` — a decoder cache, the previous segment's frames) just flows through the state. Because the reader comes before the writer, it would surface as a pipeline input; if seeding it on the first iteration is not meaningful, have the prepare step before the loop declare it as an output and set its initial value (`None`). Don't filter it out of `inputs` in the wrapper. +- Components the loop logic itself uses (e.g. `scheduler.order` for the progress bar) are added by overriding `expected_components`. + +Existing pipelines still use `LoopSequentialPipelineBlocks` (steps receive a shared flattened `block_state`, no nesting, no streaming). Leave them alone unless you are porting the pipeline; don't use it for new ones. ## Key pattern: `kwargs_type` inputs (`denoiser_input_fields`) @@ -315,7 +343,11 @@ ComponentSpec( 9. **Serving a checkpoint variant through a config flag in a shared block.** `ConfigSpec(name="is_distilled")` plus `if components.config.is_distilled:` bundles two checkpoints' behavior into one blockset — and it can't change the input surface at all (the distilled variant would still accept `negative_prompt`). Suggest a separate blockset for the variant instead (see Key pattern: Checkpoint variants). -10. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators (what the test mixins pass) work, and the CUDA-generator path is bit-identical to `torch.randn`. +10. **Writing loop-level values with `state.set()` inside a loop wrapper.** Same rule as leaf blocks: declare them in `loop_intermediate_outputs` and write through `set_block_state`. `state.get(...)` to *read* a step's output inside the loop logic is fine — those live in the `PipelineState`, not in the loop's block state. + +11. **Hiding a loop-carried input by overriding `inputs` on the loop wrapper.** Seed it from the block that runs before the loop instead (declare it as that block's output, set it to `None` / its initial value); the sequential aggregation then hides it on its own. See Key pattern: Denoising loop. + +12. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators (what the test mixins pass) work, and the CUDA-generator path is bit-identical to `torch.randn`. ## Conversion checklist diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 6c6e3a8e7882..d0b693359beb 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -104,6 +104,8 @@ title: ModularPipelineBlocks - local: modular_diffusers/sequential_pipeline_blocks title: SequentialPipelineBlocks + - local: modular_diffusers/iterative_pipeline_blocks + title: IterativePipelineBlocks - local: modular_diffusers/loop_sequential_pipeline_blocks title: LoopSequentialPipelineBlocks - local: modular_diffusers/auto_pipeline_blocks diff --git a/docs/source/en/api/modular_diffusers/pipeline_blocks.md b/docs/source/en/api/modular_diffusers/pipeline_blocks.md index 4808f2cf3bbe..268de620c9ea 100644 --- a/docs/source/en/api/modular_diffusers/pipeline_blocks.md +++ b/docs/source/en/api/modular_diffusers/pipeline_blocks.md @@ -8,6 +8,14 @@ [[autodoc]] diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks +## ModularLoopPipelineBlocks + +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.ModularLoopPipelineBlocks + +## IterativePipelineBlocks + +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.IterativePipelineBlocks + ## LoopSequentialPipelineBlocks [[autodoc]] diffusers.modular_pipelines.modular_pipeline.LoopSequentialPipelineBlocks diff --git a/docs/source/en/api/modular_diffusers/pipeline_states.md b/docs/source/en/api/modular_diffusers/pipeline_states.md index 341d18ecb41c..c3cd454f2642 100644 --- a/docs/source/en/api/modular_diffusers/pipeline_states.md +++ b/docs/source/en/api/modular_diffusers/pipeline_states.md @@ -6,4 +6,8 @@ ## BlockState -[[autodoc]] diffusers.modular_pipelines.modular_pipeline.BlockState \ No newline at end of file +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.BlockState + +## StreamEvent + +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.StreamEvent diff --git a/docs/source/en/modular_diffusers/custom_blocks.md b/docs/source/en/modular_diffusers/custom_blocks.md index f5e957a08530..a5f31edeac11 100644 --- a/docs/source/en/modular_diffusers/custom_blocks.md +++ b/docs/source/en/modular_diffusers/custom_blocks.md @@ -319,7 +319,7 @@ This guide covered creating a single custom block. Learn how to compose multiple - [SequentialPipelineBlocks](./sequential_pipeline_blocks): Chain blocks to execute in sequence - [ConditionalPipelineBlocks](./auto_pipeline_blocks): Create conditional blocks that select different execution paths -- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks): Define an iterative workflows like the denoising loop +- [IterativePipelineBlocks](./iterative_pipeline_blocks): Define an iterative workflow like the denoising loop diff --git a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md new file mode 100644 index 000000000000..6a10ff4ca4b4 --- /dev/null +++ b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md @@ -0,0 +1,286 @@ + + +# IterativePipelineBlocks + +[`~modular_pipelines.IterativePipelineBlocks`] is a multi-block type that runs its sub-blocks multiple times. It is what we use to build a denoising loop: the sub-blocks predict the noise and step the scheduler, the loop runs them once per timestep. You can also nest one [`~modular_pipelines.IterativePipelineBlocks`] under another to build an autoregressive video pipeline that generates chunk after chunk. Every iteration can be [streamed](./modular_pipeline#streaming) to the caller as it completes. + +This guide shows you how to write the loop steps, the loop itself, how to nest loops, and how values travel from one iteration to the next. + +> [!TIP] +> [`~modular_pipelines.IterativePipelineBlocks`] replaces [`~modular_pipelines.LoopSequentialPipelineBlocks`]; see [the last section](#loopsequentialpipelineblocks) for the differences. + +## Loop steps + +A loop step is a [`~modular_pipelines.ModularLoopPipelineBlocks`]. It is a regular [`~modular_pipelines.ModularPipelineBlocks`] — it declares `inputs` and `intermediate_outputs`, and reads and writes the [`~modular_pipelines.PipelineState`] through `get_block_state` / `set_block_state` — with one difference: its `__call__` also receives the loop's *loop variables* as arguments. For example, a denoising loop can pass the step index `i` and the timestep `t`. + +Loop variables are the loop's own bookkeeping, not pipeline data. They are local to the loop: the loop hands them to each step as plain call arguments for that one iteration, and they are never written to the [`~modular_pipelines.PipelineState`]. Anything that has to outlive the iteration goes through the state instead, like `noise_pred` and `latents` below. (A streaming consumer does see them: each [`~modular_pipelines.StreamEvent`] carries that iteration's values in `event.loop_kwargs`.) + +```py +from diffusers.modular_pipelines import ModularLoopPipelineBlocks, InputParam, OutputParam + +class DenoiserStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "predicts the noise for one timestep" + + @property + def inputs(self): + return [InputParam(name="latents", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="noise_pred")] + + def __call__(self, components, state, i, t): + block_state = self.get_block_state(state) + block_state.noise_pred = block_state.latents * 0 + t # stands in for the denoiser + self.set_block_state(state, block_state) + return components, state + + +class SchedulerStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "updates the latents with the noise prediction" + + @property + def inputs(self): + return [InputParam(name="latents", required=True), InputParam(name="noise_pred", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="latents")] + + def __call__(self, components, state, i, t): + block_state = self.get_block_state(state) + block_state.latents = block_state.latents + block_state.noise_pred + self.set_block_state(state, block_state) + return components, state +``` + +Because each step works on the [`~modular_pipelines.PipelineState`], the values it writes (`noise_pred`, the updated `latents`) are visible to the next step in the same iteration and to the next iteration — `latents` is read at the start of every iteration and written at the end of it. + +## Loop wrapper + +The loop itself is a subclass of [`~modular_pipelines.IterativePipelineBlocks`]. It declares: + +- `loop_variables`, the names of the variables it passes to its steps on every iteration. Every step's `__call__` must accept exactly these after `(components, state)`; this is validated when the loop is constructed. +- `loop_inputs`, the inputs the loop logic itself reads — here the `timesteps` it iterates. They join the inputs aggregated from the steps (see below), and they are what `get_block_state` returns for the loop block. +- `__call__`, the loop logic: read the loop's block state, and call `loop_step` once per iteration with the loop variables. `loop_step` runs every step once. + +```py +import torch +from diffusers.modular_pipelines import IterativePipelineBlocks + +class DenoiseLoop(IterativePipelineBlocks): + model_name = "test" + block_classes = [DenoiserStep, SchedulerStep] + block_names = ["denoiser", "scheduler"] + + @property + def description(self): + return "denoises the latents over the timesteps" + + @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 +``` + +An [`~modular_pipelines.IterativePipelineBlocks`] is a [`~modular_pipelines.SequentialPipelineBlocks`], so it [aggregates](./sequential_pipeline_blocks#aggregated-inputs-and-outputs) its steps' `inputs` and `intermediate_outputs` the same way any assembled block does, and adds `loop_inputs` / `loop_intermediate_outputs` on top. `DenoiseLoop.inputs` is therefore `timesteps` (its own) and `latents` (what the steps need from outside the loop) — but not `noise_pred`, which `DenoiserStep` produces before `SchedulerStep` reads it, so it is satisfied inside the loop. The assembled loop ends up with the same kind of input/output contract as a single block, which is what lets it be dropped into a [`~modular_pipelines.SequentialPipelineBlocks`], or into another loop. Outputs the steps write persist in the state after the loop. Run it like any other block: + +```py +pipeline = DenoiseLoop().init_pipeline() +state = pipeline(latents=torch.tensor(0.0), timesteps=torch.tensor([1.0, 2.0, 3.0])) +state.get("latents") # tensor(6.) — 0 + 1 + 2 + 3 +``` + +Steps can also be attached after the fact with [`~modular_pipelines.IterativePipelineBlocks.from_blocks_dict`], which keeps the loop logic separate from what runs inside it: + +```py +loop = DenoiseLoop.from_blocks_dict({"denoiser": DenoiserStep(), "scheduler": SchedulerStep()}) +``` + +You can also change what runs inside a loop you already have: `sub_blocks` is an ordered dict, so any loop step that takes the same loop variables can be inserted into it. + +```py +loop.sub_blocks.insert("log", LogStep(), 2) # after the denoiser and the scheduler +``` + +A step inserted this way isn't signature-checked the way the ones a loop is constructed with are — a mismatch raises a `TypeError` on the first iteration. + +If the loop logic produces a value of its own — an autoregressive loop collecting the decoded frames of every chunk, for example — declare it in `loop_intermediate_outputs` and write it back with `set_block_state`, exactly as a leaf block would: + +```py +@property +def loop_intermediate_outputs(self): + return [OutputParam(name="history")] + +@torch.no_grad() +def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.history = [] + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + block_state.history.append(state.get("latents")) + self.set_block_state(state, block_state) + return components, state +``` + +The loop's block state holds only its `loop_inputs` and `loop_intermediate_outputs`. The values the steps produce live in the [`~modular_pipelines.PipelineState`] — read them with `state.get(...)`, as above — so the loop never works from a stale copy. + +## Nesting loops + +An [`~modular_pipelines.IterativePipelineBlocks`] can be a step of another one. An autoregressive video pipeline generates a chunk of frames at a time, so it is an outer loop over chunks: each of its iterations prepares the chunk's latents from the frames generated so far, runs a full denoising loop over them, and appends the result to the history. + +The inner denoising loop is a step of the outer loop, so its `__call__` must accept the outer loop's variables, and it declares `loop_variables` of its own for its own steps. The two sets are independent: `k` arrives as a call argument and the inner loop is free to use it — the wan-animate-2 denoise loop puts the chunk index in its progress bar — but it is not forwarded to the steps, which are passed the inner loop's `i` and `t`. + +```py +class ChunkDenoiseLoop(IterativePipelineBlocks): + model_name = "test" + block_classes = [DenoiserStep, SchedulerStep] + block_names = ["denoiser", "scheduler"] + + @property + def description(self): + return "denoises one chunk over the timesteps" + + @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, k): # `k` comes from the chunk loop + 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 + + +class PrepareChunkStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "prepares this chunk's latents from the history" + + @property + def inputs(self): + return [InputParam(name="history", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="latents")] + + def __call__(self, components, state, k): + block_state = self.get_block_state(state) + block_state.latents = block_state.history + k + self.set_block_state(state, block_state) + return components, state + + +class UpdateHistoryStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "records the denoised chunk" + + @property + def inputs(self): + return [InputParam(name="latents", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="history")] + + def __call__(self, components, state, k): + block_state = self.get_block_state(state) + block_state.history = block_state.latents + self.set_block_state(state, block_state) + return components, state + + +class ChunkLoop(IterativePipelineBlocks): + model_name = "test" + block_classes = [PrepareChunkStep, ChunkDenoiseLoop, UpdateHistoryStep] + block_names = ["prepare", "denoise", "update"] + + @property + def description(self): + return "generates the video chunk by chunk" + + @property + def loop_variables(self): + return ["k"] + + @property + def loop_inputs(self): + return [InputParam(name="num_chunks", required=True)] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + for k in range(block_state.num_chunks): + components, state = self.loop_step(components, state, k=k) + return components, state +``` + +`ChunkDenoiseLoop` is `DenoiseLoop` with `k` added to `__call__`, spelled out in full here so the whole loop is visible in one place. When two loops share their logic and differ only in what runs inside them, write the logic once in a wrapper class and subclass it to attach `block_classes` / `block_names` — that is how `WanAnimate2DenoiseLoopWrapper` serves both the regular and the distilled denoise step. + +```py +pipeline = ChunkLoop().init_pipeline() +state = pipeline(num_chunks=2, timesteps=torch.tensor([1.0, 2.0]), history=torch.tensor(0.0)) +state.get("history") # tensor(7.) — chunk 0: 0 + 0 + 3 = 3, chunk 1: 3 + 1 + 3 = 7 +``` + +`history` is carried from one iteration to the next: `UpdateHistoryStep` writes it at the end of one, `PrepareChunkStep` reads it at the start of the next. Because the reader comes before the writer, it is one of the loop's inputs — which is why `pipeline(history=...)` works above — and like any input it has to come from either the user or an earlier block. If seeding it isn't meaningful (a decoder cache, the previous chunk's frames), have the block that runs before the loop declare it as an output and set its initial value; it then drops out of the pipeline's signature. + +## Streaming + +To let [`~ModularPipeline.stream`] hand back the live state 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 [`~modular_pipelines.StreamEvent`] for it (after the events of any nested loop): + +```py +class DenoiseLoop(IterativePipelineBlocks): + ... + + 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. Streaming is opt-in per loop: `pipeline.blocks.supports_streaming` tells you whether every loop on the path implements it. See [Streaming](./modular_pipeline#streaming) for the consumer side, including how to run a single iteration at a time with `loop_step` when a serving engine or a real-time input source needs to own the loop. + +## LoopSequentialPipelineBlocks + +[`~modular_pipelines.LoopSequentialPipelineBlocks`] is the earlier loop type and is still used by existing pipelines. It differs in three ways: its steps share one flattened [`~modular_pipelines.BlockState`] that the wrapper extracts before the loop (instead of each step reading the [`~modular_pipelines.PipelineState`] itself), it cannot contain another loop, and it cannot stream. Use [`~modular_pipelines.IterativePipelineBlocks`] for new pipelines. diff --git a/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md index 74a868922799..3e8c03a8d355 100644 --- a/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md @@ -12,6 +12,9 @@ specific language governing permissions and limitations under the License. # LoopSequentialPipelineBlocks +> [!WARNING] +> [`~modular_pipelines.LoopSequentialPipelineBlocks`] is superseded by [`~modular_pipelines.IterativePipelineBlocks`], which lets loop steps work on the [`~modular_pipelines.PipelineState`] like regular blocks, can be nested, and supports streaming. Use [IterativePipelineBlocks](./iterative_pipeline_blocks) for new pipelines; this page documents the earlier type that existing pipelines still use. + [`~modular_pipelines.LoopSequentialPipelineBlocks`] are a multi-block type that composes other [`~modular_pipelines.ModularPipelineBlocks`] together in a loop. Data flows circularly, using `inputs` and `intermediate_outputs`, and each block is run iteratively. This is typically used to create a denoising loop which is iterative by default. This guide shows you how to create [`~modular_pipelines.LoopSequentialPipelineBlocks`]. diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 1908065f760c..401d8a0d29d6 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -413,7 +413,7 @@ for event in pipeline.stream(...): 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: +Streaming is opt-in per loop block. An [`IterativePipelineBlocks`] (see the [IterativePipelineBlocks](./iterative_pipeline_blocks) guide) 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): diff --git a/docs/source/en/modular_diffusers/overview.md b/docs/source/en/modular_diffusers/overview.md index f80fff3061de..1c8d91a782ba 100644 --- a/docs/source/en/modular_diffusers/overview.md +++ b/docs/source/en/modular_diffusers/overview.md @@ -28,7 +28,8 @@ The Modular Diffusers docs are organized as shown below. - [States](./modular_diffusers_states) explains how data is shared and communicated between blocks and [`ModularPipeline`]. - [ModularPipelineBlocks](./pipeline_block) is the most basic unit of a [`ModularPipeline`] and this guide shows you how to create one. - [SequentialPipelineBlocks](./sequential_pipeline_blocks) is a type of block that chains multiple blocks so they run one after another, passing data along the chain. This guide shows you how to create [`~modular_pipelines.SequentialPipelineBlocks`] and how they connect and work together. -- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks) is a type of block that runs a series of blocks in a loop. This guide shows you how to create [`~modular_pipelines.LoopSequentialPipelineBlocks`]. +- [IterativePipelineBlocks](./iterative_pipeline_blocks) is a type of block that runs a series of blocks in a loop — a denoising loop, or an autoregressive chunk loop with a denoising loop nested inside. This guide shows you how to create [`~modular_pipelines.IterativePipelineBlocks`]. +- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks) is the earlier loop type, still used by existing pipelines. - [AutoPipelineBlocks](./auto_pipeline_blocks) is a type of block that automatically chooses which blocks to run based on the input. This guide shows you how to create [`~modular_pipelines.AutoPipelineBlocks`]. - [Building Custom Blocks](./custom_blocks) shows you how to create your own custom blocks and share them on the Hub. diff --git a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md index 1bd67e17b8bf..d9c093b68703 100644 --- a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md @@ -113,6 +113,12 @@ class ImageProcessingStep(SequentialPipelineBlocks): When you create a [`~modular_pipelines.SequentialPipelineBlocks`], properties like `inputs`, `intermediate_outputs`, and `expected_components` are automatically aggregated from the sub-blocks, so there is no need to define them again. +### Aggregated inputs and outputs + +Aggregation follows the order the sub-blocks run in. `inputs` walks the sub-blocks in order and collects each one's inputs, skipping any that an earlier sub-block already declares in its `intermediate_outputs` — that value is produced inside the assembled block, so it isn't asked of the caller. `intermediate_outputs` is the union of what the sub-blocks write. + +This is what makes blocks composable: an assembled block has the same kind of input/output contract as a single [`~modular_pipelines.ModularPipelineBlocks`], so it can in turn be a sub-block of another one, at any depth. It is also why an input disappears from a pipeline's signature once some earlier block produces it — see [nesting loops](./iterative_pipeline_blocks#nesting-loops) for a case where you use that deliberately. + There are a few properties you should set: - `description`: We recommend adding a description for the assembled block to explain what the combined step does. From 1c17e84a005b0391a8118a045c2f7c6e73e8e306 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Tue, 25 Aug 2026 09:43:29 -1000 Subject: [PATCH 12/24] Apply suggestion from @yiyixuxu --- .ai/modular.md | 1 - 1 file changed, 1 deletion(-) diff --git a/.ai/modular.md b/.ai/modular.md index 465a55e1957e..95229cd9cb6d 100644 --- a/.ai/modular.md +++ b/.ai/modular.md @@ -4,7 +4,6 @@ Shared reference for modular pipeline conventions, patterns, and gotchas. ## Common modular conventions -# should we include minimax, maybe replace wan? When adding a new modular pipeline (or reviewing one), skim `src/diffusers/modular_pipelines/qwenimage/`, `src/diffusers/modular_pipelines/flux2/`, `src/diffusers/modular_pipelines/wan/`, and `src/diffusers/modular_pipelines/helios/` first to establish the pattern. Most conventions (file split between `encoders.py` / `before_denoise.py` / `denoise.py` / `decoders.py`, how `expected_components` / `inputs` / `intermediate_outputs` are declared, the denoise-loop wrapping with `IterativePipelineBlocks` (see `flux2/denoise.py`; `wan_animate_2/denoise.py` for a nested chunk loop), top-level assembly via `AutoPipelineBlocks` / `SequentialPipelineBlocks` in `modular_blocks_.py`, the `ModularPipeline` subclass shape, the guider-abstracted denoise body, `kwargs_type="denoiser_input_fields"` plumbing) are easiest to internalize by comparison rather than from a fixed list. ## Running a modular pipeline From d7ec4e6952150c1a9657ce91181a29f1aa2ff51f Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 1 Sep 2026 06:47:17 +0200 Subject: [PATCH 13/24] keep loop logic minimal: collect segment frames in a loop step, let loop steps take only the loop variables they use (plus **kwargs), update docs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .ai/references/modular.md | 49 ++++++--- .../iterative_pipeline_blocks.md | 87 ++++++++-------- .../modular_pipelines/modular_pipeline.py | 33 +++++-- .../wan_animate_2/before_denoise.py | 11 ++- .../wan_animate_2/denoise.py | 99 ++++++++++++------- .../test_iterative_pipeline_blocks.py | 84 ++++++++++++---- 6 files changed, 238 insertions(+), 125 deletions(-) diff --git a/.ai/references/modular.md b/.ai/references/modular.md index 2d8cd45115dd..71ae6d4cc42c 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -125,16 +125,18 @@ for i, t in enumerate(timesteps): ## Key pattern: Denoising loop -The denoising loop is an `IterativePipelineBlocks` whose steps are `ModularLoopPipelineBlocks` (canonical: `flux2/denoise.py`). The wrapper declares the loop variables it passes to its steps, the inputs its own loop logic reads, and the loop in `__call__`; `stream` is the same loop as a generator over `stream_step` and is what makes `pipe.stream(...)` work — implement both: +The denoising loop is an `IterativePipelineBlocks` whose steps are `ModularLoopPipelineBlocks` (see `flux2/denoise.py` as example). The wrapper declares the loop variables it passes to its steps, the inputs its own loop logic reads (in addition to the inputs/outputs of the loop steps), and the loop in `__call__`; `stream` is the same loop as a generator over `stream_step` and is what makes `pipe.stream(...)` work — implement both: ```python class MyModelDenoiseLoopWrapper(IterativePipelineBlocks): @property - def loop_variables(self): + def loop_variables(self): # passed to every loop step return ["i", "t"] @property - def loop_inputs(self): # what the loop logic itself reads; `get_block_state` returns exactly these + def loop_inputs(self): + # what the loop logic itself reads — `get_block_state` returns exactly these + # (the loop steps declare their own `inputs` on top of this) return [InputParam("timesteps", required=True), InputParam.template("num_inference_steps", required=True)] @torch.no_grad() @@ -157,14 +159,37 @@ class MyModelDenoiseStep(MyModelDenoiseLoopWrapper): block_names = ["denoiser", "after_denoiser"] ``` -Loop steps are regular blocks (own `inputs` / `intermediate_outputs`, `get_block_state` / `set_block_state` on the full `PipelineState`) whose `__call__` takes the loop variables as arguments — `def __call__(self, components, state, i, t)`. Every step of a loop must accept exactly that loop's variables (validated at construction), even if it ignores some of them. +Loop steps are regular blocks (own `inputs` / `intermediate_outputs`, `get_block_state` / `set_block_state` on the full `PipelineState`) whose `__call__` takes the loop variables as extra keyword arguments — name the ones the step uses, declare `**kwargs` for any it ignores (validated at construction: unknown names raise, and so does a missing variable without a `**kwargs` catch-all). A minimal one: -Autoregressive video models nest loops: an outer chunk loop (`loop_variables = ["k"]`) whose steps prepare the chunk, run the full inner denoising loop, and update the history (canonical: `wan_animate_2/denoise.py`). The inner loop is a step of the outer one, so its `__call__` / `stream` accept the outer variable — `def __call__(self, components, state, k)` — and pass its own `i`, `t` to its own steps. +```python +class MyModelLoopAfterDenoiser(ModularLoopPipelineBlocks): + @property + def expected_components(self): + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self): + return [InputParam("latents", required=True), InputParam("noise_pred", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam("latents")] + + @torch.no_grad() + def __call__(self, components, state, i, t): + block_state = self.get_block_state(state) + block_state.latents = components.scheduler.step(block_state.noise_pred, t, block_state.latents, return_dict=False)[0] + self.set_block_state(state, block_state) + return components, state +``` -Conventions that fall out of this: -- A value the loop logic itself produces (e.g. the collected per-chunk frames) goes in `loop_intermediate_outputs` and is written with `set_block_state`, like a leaf output. Step outputs are read from the state (`state.get("out_frames")`) — the loop's block state holds only its `loop_inputs` / `loop_intermediate_outputs`. -- A loop-carried value (written by a step at the end of iteration `k`, read by a step at the start of `k + 1` — a decoder cache, the previous segment's frames) just flows through the state. Because the reader comes before the writer, it would surface as a pipeline input; if seeding it on the first iteration is not meaningful, have the prepare step before the loop declare it as an output and set its initial value (`None`). Don't filter it out of `inputs` in the wrapper. -- Components the loop logic itself uses (e.g. `scheduler.order` for the progress bar) are added by overriding `expected_components`. +Autoregressive video models nest loops: an outer segment loop (`loop_variables = ["k"]`) whose steps prepare the segment, run the full inner denoising loop, and update the history (see `wan_animate_2/denoise.py` as an example). The inner loop is a step of the outer one, so its `__call__` / `stream` accept the outer variable — `def __call__(self, components, state, k)` — and pass its own `i`, `t` to its own steps. + +The wrapper contains only the loop logic — how to iterate through the steps. Its `loop_inputs` are just what that takes (`timesteps`, `num_segments`); all data flows through the steps, which read and write the pipeline state directly. If the loop logic seems to need to do more than iterate — collect results, carry something to the next iteration — add a loop step for it instead (in `wan_animate_2`, a small `collect` step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back). + +Two small notes: +- A value written late in iteration `k` and read early in iteration `k + 1` would surface as a pipeline input (the first read has no writer before it). Seed it in a block that runs before the loop, declare it as an output, set it to `None`, and it stays internal. +- Components the loop logic itself uses (e.g. `flux2`'s wrapper reads `scheduler.order` to compute the progress-bar warmup steps) are added by overriding `expected_components`. Existing pipelines still use `LoopSequentialPipelineBlocks` (steps receive a shared flattened `block_state`, no nesting, no streaming). Leave them alone unless you are porting the pipeline; don't use it for new ones. @@ -342,11 +367,9 @@ ComponentSpec( 9. **Serving a checkpoint variant through a config flag in a shared block.** `ConfigSpec(name="is_distilled")` plus `if components.config.is_distilled:` bundles two checkpoints' behavior into one blockset — and it can't change the input surface at all (the distilled variant would still accept `negative_prompt`). Suggest a separate blockset for the variant instead (see Key pattern: Checkpoint variants). -10. **Writing loop-level values with `state.set()` inside a loop wrapper.** Same rule as leaf blocks: declare them in `loop_intermediate_outputs` and write through `set_block_state`. `state.get(...)` to *read* a step's output inside the loop logic is fine — those live in the `PipelineState`, not in the loop's block state. - -11. **Hiding a loop-carried input by overriding `inputs` on the loop wrapper.** Seed it from the block that runs before the loop instead (declare it as that block's output, set it to `None` / its initial value); the sequential aggregation then hides it on its own. See Key pattern: Denoising loop. +10. **Hiding a loop-carried input by overriding `inputs` on the loop wrapper.** Seed it from the block that runs before the loop instead (declare it as that block's output, set it to `None` / its initial value); the sequential aggregation then hides it on its own. See Key pattern: Denoising loop. -12. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators (what the test mixins pass) work, and the CUDA-generator path is bit-identical to `torch.randn`. +11. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators (what the test mixins pass) work, and the CUDA-generator path is bit-identical to `torch.randn`. ## Conversion checklist diff --git a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md index 6a10ff4ca4b4..c1669108303e 100644 --- a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md @@ -78,18 +78,18 @@ Because each step works on the [`~modular_pipelines.PipelineState`], the values The loop itself is a subclass of [`~modular_pipelines.IterativePipelineBlocks`]. It declares: -- `loop_variables`, the names of the variables it passes to its steps on every iteration. Every step's `__call__` must accept exactly these after `(components, state)`; this is validated when the loop is constructed. +- `loop_variables`, the names of the variables it passes to its steps on every iteration (as keyword arguments). A step's `__call__` names the ones it uses after `(components, state)` and declares `**kwargs` for any it ignores — naming all of them works too. This is validated when the loop is constructed: a named parameter that isn't a loop variable, or a missing one without a `**kwargs` catch-all, raises. - `loop_inputs`, the inputs the loop logic itself reads — here the `timesteps` it iterates. They join the inputs aggregated from the steps (see below), and they are what `get_block_state` returns for the loop block. - `__call__`, the loop logic: read the loop's block state, and call `loop_step` once per iteration with the loop variables. `loop_step` runs every step once. +Note that the wrapper defines only the loop logic - which steps run inside it should be attached separately, in a subclass: + ```py import torch from diffusers.modular_pipelines import IterativePipelineBlocks -class DenoiseLoop(IterativePipelineBlocks): +class DenoiseLoopWrapper(IterativePipelineBlocks): model_name = "test" - block_classes = [DenoiserStep, SchedulerStep] - block_names = ["denoiser", "scheduler"] @property def description(self): @@ -109,61 +109,55 @@ class DenoiseLoop(IterativePipelineBlocks): for i, t in enumerate(block_state.timesteps): components, state = self.loop_step(components, state, i=i, t=t) return components, state -``` -An [`~modular_pipelines.IterativePipelineBlocks`] is a [`~modular_pipelines.SequentialPipelineBlocks`], so it [aggregates](./sequential_pipeline_blocks#aggregated-inputs-and-outputs) its steps' `inputs` and `intermediate_outputs` the same way any assembled block does, and adds `loop_inputs` / `loop_intermediate_outputs` on top. `DenoiseLoop.inputs` is therefore `timesteps` (its own) and `latents` (what the steps need from outside the loop) — but not `noise_pred`, which `DenoiserStep` produces before `SchedulerStep` reads it, so it is satisfied inside the loop. The assembled loop ends up with the same kind of input/output contract as a single block, which is what lets it be dropped into a [`~modular_pipelines.SequentialPipelineBlocks`], or into another loop. Outputs the steps write persist in the state after the loop. Run it like any other block: -```py -pipeline = DenoiseLoop().init_pipeline() -state = pipeline(latents=torch.tensor(0.0), timesteps=torch.tensor([1.0, 2.0, 3.0])) -state.get("latents") # tensor(6.) — 0 + 1 + 2 + 3 +class DenoiseLoop(DenoiseLoopWrapper): + block_classes = [DenoiserStep, SchedulerStep] + block_names = ["denoiser", "scheduler"] ``` -Steps can also be attached after the fact with [`~modular_pipelines.IterativePipelineBlocks.from_blocks_dict`], which keeps the loop logic separate from what runs inside it: +This separation means the same loop logic can work with different combination of loop steps: subclass the wrapper again with different `block_classes` (this is exactly how `Flux2DenoiseLoopWrapper` serves the base and klein denoise steps). Steps can also be attached to the loop with [`~modular_pipelines.IterativePipelineBlocks.from_blocks_dict`]: ```py -loop = DenoiseLoop.from_blocks_dict({"denoiser": DenoiserStep(), "scheduler": SchedulerStep()}) +loop = DenoiseLoopWrapper.from_blocks_dict({"denoiser": DenoiserStep(), "scheduler": SchedulerStep()}) ``` -You can also change what runs inside a loop you already have: `sub_blocks` is an ordered dict, so any loop step that takes the same loop variables can be inserted into it. +You can also change what runs inside a loop you already have: add a step, reorder, swap one out. e.g. you can insert a logging step like this: ```py -loop.sub_blocks.insert("log", LogStep(), 2) # after the denoiser and the scheduler +loop = DenoiseLoopWrapper.from_blocks_dict(loop.sub_blocks.copy().insert("log", LogStep(), 1)) ``` -A step inserted this way isn't signature-checked the way the ones a loop is constructed with are — a mismatch raises a `TypeError` on the first iteration. - -If the loop logic produces a value of its own — an autoregressive loop collecting the decoded frames of every chunk, for example — declare it in `loop_intermediate_outputs` and write it back with `set_block_state`, exactly as a leaf block would: +An [`~modular_pipelines.IterativePipelineBlocks`] is a [`~modular_pipelines.SequentialPipelineBlocks`], so it [aggregates](./sequential_pipeline_blocks#aggregated-inputs-and-outputs) its steps' `inputs` and `intermediate_outputs` the same way any assembled block does, and adds `loop_inputs` / `loop_intermediate_outputs` on top. `DenoiseLoop.inputs` is therefore `timesteps` (its own) and `latents` (what the steps need from outside the loop) — but not `noise_pred`, which `DenoiserStep` produces before `SchedulerStep` reads it, so it is satisfied inside the loop. The assembled loop ends up with the same kind of input/output contract as a single block, which is what lets it be dropped into a [`~modular_pipelines.SequentialPipelineBlocks`], or into another loop. Outputs the steps write persist in the state after the loop. Run it like any other block: ```py -@property -def loop_intermediate_outputs(self): - return [OutputParam(name="history")] - -@torch.no_grad() -def __call__(self, components, state): - block_state = self.get_block_state(state) - block_state.history = [] - for i, t in enumerate(block_state.timesteps): - components, state = self.loop_step(components, state, i=i, t=t) - block_state.history.append(state.get("latents")) - self.set_block_state(state, block_state) - return components, state +pipeline = DenoiseLoop().init_pipeline() +state = pipeline(latents=torch.tensor(0.0), timesteps=torch.tensor([1.0, 2.0, 3.0])) +state.get("latents") # tensor(6.) — 0 + 1 + 2 + 3 ``` -The loop's block state holds only its `loop_inputs` and `loop_intermediate_outputs`. The values the steps produce live in the [`~modular_pipelines.PipelineState`] — read them with `state.get(...)`, as above — so the loop never works from a stale copy. +The wrapper contains only the loop logic, i.e. how to iterate through its steps, so its `loop_inputs` should be just what that takes (the `timesteps` above). All data flows through the steps, which read and write the pipeline state directly. If the loop logic seems to need to do more than iterate: e.g. collect results, you should add a loop step for it instead: in `wan_animate_2`, a small collect step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back to condition on; under [streaming](#streaming), the partial collection is visible after every iteration. ## Nesting loops An [`~modular_pipelines.IterativePipelineBlocks`] can be a step of another one. An autoregressive video pipeline generates a chunk of frames at a time, so it is an outer loop over chunks: each of its iterations prepares the chunk's latents from the frames generated so far, runs a full denoising loop over them, and appends the result to the history. -The inner denoising loop is a step of the outer loop, so its `__call__` must accept the outer loop's variables, and it declares `loop_variables` of its own for its own steps. The two sets are independent: `k` arrives as a call argument and the inner loop is free to use it — the wan-animate-2 denoise loop puts the chunk index in its progress bar — but it is not forwarded to the steps, which are passed the inner loop's `i` and `t`. +The inner denoising loop is a step of the outer loop, so its `__call__` must accept the outer loop's variables (or take `**kwargs` if it ignores them), and it declares `loop_variables` of its own for its own steps. The two sets are independent: `k` arrives as a call argument and the inner loop is free to use it — the wan-animate-2 denoise loop puts the chunk index in its progress bar — but it is not forwarded to the steps, which are passed the inner loop's `i` and `t`. + +This is the structure this section builds: + +``` +VideoLoop outer loop — loop_variables ["k"], iterates over chunks +├─ prepare PrepareChunkStep (leaf step, takes k) +├─ denoise ChunkDenoiseLoop (the INNER loop — a step of the outer one) +│ ├─ denoiser DenoiserStep (takes the inner loop's i, t) +│ └─ scheduler SchedulerStep +└─ update UpdateHistoryStep (leaf step, takes k) +``` ```py -class ChunkDenoiseLoop(IterativePipelineBlocks): +class ChunkDenoiseLoopWrapper(IterativePipelineBlocks): model_name = "test" - block_classes = [DenoiserStep, SchedulerStep] - block_names = ["denoiser", "scheduler"] @property def description(self): @@ -178,13 +172,21 @@ class ChunkDenoiseLoop(IterativePipelineBlocks): return [InputParam(name="timesteps", required=True)] @torch.no_grad() - def __call__(self, components, state, k): # `k` comes from the chunk loop + def __call__(self, components, state, k): # `k` comes from the outer video loop 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 +class ChunkDenoiseLoop(ChunkDenoiseLoopWrapper): + block_classes = [DenoiserStep, SchedulerStep] + block_names = ["denoiser", "scheduler"] +``` + +`ChunkDenoiseLoop` is the *inner* loop: it denoises a single chunk. It becomes a step of the outer loop below, which iterates over the chunks — at each `k` it prepares the chunk's latents from the history, runs the full inner denoising loop over them, and records the result: + +```py class PrepareChunkStep(ModularLoopPipelineBlocks): model_name = "test" @@ -229,10 +231,8 @@ class UpdateHistoryStep(ModularLoopPipelineBlocks): return components, state -class ChunkLoop(IterativePipelineBlocks): +class VideoLoopWrapper(IterativePipelineBlocks): model_name = "test" - block_classes = [PrepareChunkStep, ChunkDenoiseLoop, UpdateHistoryStep] - block_names = ["prepare", "denoise", "update"] @property def description(self): @@ -252,12 +252,17 @@ class ChunkLoop(IterativePipelineBlocks): for k in range(block_state.num_chunks): components, state = self.loop_step(components, state, k=k) return components, state + + +class VideoLoop(VideoLoopWrapper): + block_classes = [PrepareChunkStep, ChunkDenoiseLoop, UpdateHistoryStep] + block_names = ["prepare", "denoise", "update"] ``` -`ChunkDenoiseLoop` is `DenoiseLoop` with `k` added to `__call__`, spelled out in full here so the whole loop is visible in one place. When two loops share their logic and differ only in what runs inside them, write the logic once in a wrapper class and subclass it to attach `block_classes` / `block_names` — that is how `WanAnimate2DenoiseLoopWrapper` serves both the regular and the distilled denoise step. +Run the outer loop like any other block: ```py -pipeline = ChunkLoop().init_pipeline() +pipeline = VideoLoop().init_pipeline() state = pipeline(num_chunks=2, timesteps=torch.tensor([1.0, 2.0]), history=torch.tensor(0.0)) state.get("history") # tensor(7.) — chunk 0: 0 + 0 + 3 = 3, chunk 1: 3 + 1 + 3 = 7 ``` diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index b9ee1346d0a1..9ece2b9a2241 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -1487,8 +1487,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 variables are passed to sub-blocks as keyword call arguments: a sub-block's `__call__` names the loop + variables it uses after `(components, state)` and declares `**kwargs` for any it ignores (naming all of them and + omitting `**kwargs` is fine too). This is validated at construction: a named parameter that is not a loop + variable, or a missing loop variable without a `**kwargs` catch-all, raises. 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: @@ -1510,8 +1512,8 @@ def __call__(self, components, state, k): # accepts the OUTER chunk loop's vari 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`. + 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 @@ -1597,7 +1599,8 @@ def from_blocks_dict(cls, blocks_dict, description: str | None = None) -> "Itera 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)`.""" + and accept the loop variables after `(components, state)` — either all of them by name, or the ones the + step uses plus a `**kwargs` catch-all for the rest.""" expected = set(self.loop_variables) for block_name, block in self.sub_blocks.items(): if not isinstance(block, (ModularLoopPipelineBlocks, IterativePipelineBlocks)): @@ -1606,13 +1609,23 @@ def _validate_sub_blocks(self): "a `ModularLoopPipelineBlocks` (a loop step) or an `IterativePipelineBlocks` (a nested loop); " f"got `{block.__class__.__bases__[0].__name__}`." ) - params = list(inspect.signature(block.__call__).parameters) - extra = set(params[2:]) - if extra != expected: + params = inspect.signature(block.__call__).parameters + has_var_keyword = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()) + # `block.__call__` is bound, so `self` is already excluded; `[2:]` skips `components` and `state`, + # leaving only the loop-variable parameters + named = {name for name in list(params)[2:] if params[name].kind is not inspect.Parameter.VAR_KEYWORD} + if not named <= 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)}." + f"accepts {sorted(named - expected)}, which are not loop variables of this loop " + f"({sorted(expected)})." + ) + if named != expected and not has_var_keyword: + 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)` — either " + f"all of them by name, or the ones it uses plus `**kwargs`; its `__call__` accepts " + f"{sorted(named)}." ) def loop_step(self, components, state: PipelineState, **loop_kwargs) -> PipelineState: 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 6d84e5d8d20f..7df13e151360 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py @@ -80,10 +80,11 @@ def intermediate_outputs(self) -> list[OutputParam]: 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)", + "segment_frames", + type_hint=list, + description="The decoded frames collected by the segment loop, carried across iterations (the " + "next segment conditions on the last entry's tail); starts as `None` (nothing has been " + "collected yet, and the first segment has no previous segment to condition on)", ), ] @@ -113,7 +114,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 + block_state.segment_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 d1798e042277..541592449265 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -170,9 +170,9 @@ def inputs(self) -> list[InputParam]: 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", + "segment_frames", + type_hint=list, + description="The decoded segments collected so far (on CPU), written by the collect step; the previous segment's tail is read from the last entry. `None` for the first segment", ), InputParam( "segment_frame_length", @@ -210,7 +210,8 @@ def __call__(self, components, state: PipelineState, k: int): num_frames = block_state.segment_frame_length + 1 mask_len = block_state.prev_segment_conditioning_frames if k > 0 else 0 if mask_len > 0: - prev_frames = block_state.out_frames[0, :, -mask_len:].clone().detach() + # slice the tail on CPU first so only the few conditioning frames are transferred + prev_frames = block_state.segment_frames[-1][0, :, -mask_len:].to(device) prev_frames = F.interpolate(prev_frames.permute(1, 0, 2, 3), size=(height, width), mode="bicubic").permute( 1, 0, 2, 3 ) @@ -279,7 +280,7 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState, k: int): + def __call__(self, components, state: PipelineState, **kwargs): block_state = self.get_block_state(state) device = components._execution_device @@ -331,7 +332,7 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState, k: int): + def __call__(self, components, state: PipelineState, **kwargs): block_state = self.get_block_state(state) device = components._execution_device @@ -415,7 +416,7 @@ def inputs(self) -> list[InputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState, k: int): + def __call__(self, components, state: PipelineState, **kwargs): block_state = self.get_block_state(state) device = components._execution_device transformer_dtype = components.transformer.dtype @@ -806,7 +807,7 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "out_frames", type_hint=torch.Tensor, - description="This segment's decoded frames on device; the next segment conditions on its tail", + description="This segment's decoded frames on device, consumed (and freed) by the collect step", ), ] @@ -832,6 +833,53 @@ def __call__(self, components, state: PipelineState, k: int): return components, state +class WanAnimate2SegmentCollectStep(ModularLoopPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step within the segment loop that appends this segment's decoded frames to the collected video, " + "moving them to CPU, and frees the on-device copy — the next segment's conditioning reads the tail " + "of the last collected entry. This block should be used to compose the `sub_blocks` attribute of an " + "`IterativePipelineBlocks` object (e.g. `WanAnimate2SegmentLoopWrapper`)." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "out_frames", + required=True, + type_hint=torch.Tensor, + description="This segment's decoded frames on device, from the decode step", + ), + InputParam( + "segment_frames", + type_hint=list, + description="The segments collected so far; `None` before the first segment", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "segment_frames", + type_hint=list[torch.Tensor], + description="Per-segment decoded frames on CPU, each `[1, 3, T, H, W]`", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, **kwargs): + block_state = self.get_block_state(state) + block_state.segment_frames = [*(block_state.segment_frames or []), block_state.out_frames.cpu()] + block_state.out_frames = None # release the on-device copy; nothing reads it after this + self.set_block_state(state, block_state) + return components, state + + # ======================================== # Segment Loop Wrapper # ======================================== @@ -848,8 +896,8 @@ def loop_variables(self) -> list[str]: def description(self) -> str: return ( "Pipeline block that iterates over the driving video's segments. At each segment it runs sub-blocks " - "for per-segment encoding, preparation, reference extraction, denoising, and decoding; each segment " - "conditions on the previous one's decoded tail frames." + "for per-segment encoding, preparation, reference extraction, denoising, decoding, and collection; " + "each segment conditions on the previous one's decoded tail frames." ) @property @@ -863,41 +911,18 @@ def loop_inputs(self) -> list[InputParam]: ), ] - @property - 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], - description="Per-segment decoded frames on CPU, each `[1, 3, T, H, W]`", - ), - ] - @torch.no_grad() def __call__(self, components, state: PipelineState): block_state = self.get_block_state(state) - - # `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. - block_state.segment_frames = [] for k in range(block_state.num_segments): components, state = self.loop_step(components, state, k=k) - block_state.segment_frames.append(state.get("out_frames").cpu()) - 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) - - block_state.segment_frames = [] for k in range(block_state.num_segments): components, state = yield from self.stream_step(components, state, k=k) - block_state.segment_frames.append(state.get("out_frames").cpu()) - self.set_block_state(state, block_state) - return components, state @@ -915,6 +940,7 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2RefExtractStep, WanAnimate2SegmentDenoiseStep, WanAnimate2SegmentDecodeStep, + WanAnimate2SegmentCollectStep, ] block_names = [ "vae_encoder", @@ -924,6 +950,7 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): "ref_extract", "denoise_inner", "decode", + "collect", ] @property @@ -931,7 +958,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 (a nested denoising loop over this segment's timesteps) -> decode." + "denoise_inner (a nested denoising loop over this segment's timesteps) -> decode -> collect." ) @@ -946,6 +973,7 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2RefExtractStep, WanAnimate2DistilledSegmentDenoiseStep, WanAnimate2SegmentDecodeStep, + WanAnimate2SegmentCollectStep, ] block_names = [ "vae_encoder", @@ -955,6 +983,7 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): "ref_extract", "denoise_inner", "decode", + "collect", ] @property @@ -963,5 +992,5 @@ def description(self) -> str: "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 (a nested denoising loop over this segment's timesteps, no classifier-free guidance) " - "-> decode." + "-> decode -> collect." ) diff --git a/tests/modular_pipelines/test_iterative_pipeline_blocks.py b/tests/modular_pipelines/test_iterative_pipeline_blocks.py index 9054bb0d535e..b0090445b517 100644 --- a/tests/modular_pipelines/test_iterative_pipeline_blocks.py +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -144,7 +144,7 @@ def intermediate_outputs(self): def description(self): return "records the denoised chunk and updates the history" - def __call__(self, components, state, k): + def __call__(self, components, state, **kwargs): # ignores the loop's `k`: a catch-all is enough 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)] @@ -180,7 +180,13 @@ def __call__(self, components, state): class CollectingChunkLoop(ChunkLoop): - """Chunk loop whose loop logic has an output of its own, written through `set_block_state`.""" + """Chunk loop whose loop logic has an output of its own, written through `set_block_state`, and which + observes a step output (`history`) by declaring it in `loop_inputs` and calling `get_block_state` again + after each `loop_step` — the fresh snapshot picks up what the steps just wrote.""" + + @property + def loop_inputs(self): + return [*super().loop_inputs, InputParam(name="history")] @property def loop_intermediate_outputs(self): @@ -189,10 +195,12 @@ def loop_intermediate_outputs(self): @torch.no_grad() def __call__(self, components, state): block_state = self.get_block_state(state) - block_state.chunk_history = [] + 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"))) + block_state = self.get_block_state(state) + chunk_history.append(float(block_state.history)) + block_state.chunk_history = chunk_history self.set_block_state(state, block_state) return components, state @@ -266,6 +274,8 @@ def test_block_state_is_loop_scoped(self): 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)) + # `history` is re-read with `get_block_state` after every iteration — without the re-read this + # would collect the stale iteration-0 seed, [0.0, 0.0, 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] @@ -294,26 +304,16 @@ def description(self): 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 at construction - class WrongSigStep(ModularLoopPipelineBlocks): + @staticmethod + def _loop_over(step_cls): + class Loop(IterativePipelineBlocks): model_name = "test" + block_classes = [step_cls] + block_names = ["step"] @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" + return "loop over i, t" @property def loop_variables(self): @@ -330,8 +330,50 @@ def __call__(self, components, state): components, state = self.loop_step(components, state, i=i, t=t) return components, state + return Loop + + def test_leaf_signature_is_validated(self): + # a named parameter that isn't a loop variable fails at construction, even with a catch-all + class UnknownVarStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "loop step naming a variable the loop doesn't have" + + def __call__(self, components, state, k, **kwargs): + return components, state + + with pytest.raises(ValueError, match="not loop variables"): + self._loop_over(UnknownVarStep)() + + # naming only some loop variables without a `**kwargs` catch-all fails at construction + class MissingVarStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "loop step naming only one of the loop variables, without a catch-all" + + def __call__(self, components, state, t): + return components, state + with pytest.raises(ValueError, match="must accept the loop variables"): - BadSigLoop() + self._loop_over(MissingVarStep)() + + def test_leaf_signature_subset_with_kwargs_is_valid(self): + # naming only the used loop variables plus `**kwargs` is accepted + class SubsetStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "loop step using only `t`, ignoring the rest via a catch-all" + + def __call__(self, components, state, t, **kwargs): + return components, state + + self._loop_over(SubsetStep)() # does not raise def test_loop_leaf_standalone_raises(self): # outside a loop, a leaf block with loop variables in its signature cannot run From 8bc93c97d9646a954dc582084d9896df282479d6 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 1 Sep 2026 07:02:10 +0200 Subject: [PATCH 14/24] remove loop gotchas covered by the updated docs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .ai/references/modular.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.ai/references/modular.md b/.ai/references/modular.md index 71ae6d4cc42c..907bd4ba3dc1 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -367,9 +367,7 @@ ComponentSpec( 9. **Serving a checkpoint variant through a config flag in a shared block.** `ConfigSpec(name="is_distilled")` plus `if components.config.is_distilled:` bundles two checkpoints' behavior into one blockset — and it can't change the input surface at all (the distilled variant would still accept `negative_prompt`). Suggest a separate blockset for the variant instead (see Key pattern: Checkpoint variants). -10. **Hiding a loop-carried input by overriding `inputs` on the loop wrapper.** Seed it from the block that runs before the loop instead (declare it as that block's output, set it to `None` / its initial value); the sequential aggregation then hides it on its own. See Key pattern: Denoising loop. - -11. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators (what the test mixins pass) work, and the CUDA-generator path is bit-identical to `torch.randn`. +10. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators (what the test mixins pass) work, and the CUDA-generator path is bit-identical to `torch.randn`. ## Conversion checklist From 37ff5588037b731476a4651db22ba9d475d34d29 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 1 Sep 2026 09:58:22 +0200 Subject: [PATCH 15/24] clarify streaming docs: nested-loop events, drop the drive-the-loop-yourself section Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .../iterative_pipeline_blocks.md | 2 +- .../en/modular_diffusers/modular_pipeline.md | 23 +++---------------- 2 files changed, 4 insertions(+), 21 deletions(-) diff --git a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md index c1669108303e..ec347ace19ac 100644 --- a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md @@ -284,7 +284,7 @@ class DenoiseLoop(IterativePipelineBlocks): return components, state ``` -A nested loop's `stream` takes the outer loop's variables exactly like its `__call__` does. Streaming is opt-in per loop: `pipeline.blocks.supports_streaming` tells you whether every loop on the path implements it. See [Streaming](./modular_pipeline#streaming) for the consumer side, including how to run a single iteration at a time with `loop_step` when a serving engine or a real-time input source needs to own the loop. +A nested loop's `stream` takes the outer loop's variables exactly like its `__call__` does. Streaming is opt-in per loop: `pipeline.blocks.supports_streaming` tells you whether every loop on the path implements it. See [Streaming](./modular_pipeline#streaming) for the consumer side. ## LoopSequentialPipelineBlocks diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 3778dffa1c27..e0a2e8402201 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -394,7 +394,7 @@ If pipeline stages share components (e.g., the same VAE used for encoding and de ## 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. +[`~ModularPipeline.stream`] runs the same pipeline as a generator. It yields a [`StreamEvent`] after every iteration of every loop block, e.g. 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) @@ -403,11 +403,11 @@ for event in generator: 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: +`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 (e.g. an autoregressive video that denoises one chunk at a time), you receive events from both loops: one after every denoising step of the inner loop (path `"denoise.denoise_inner"`), and one after each completed chunk of the outer loop (path `"denoise"`). Check `event.path` to tell them apart — for example, to react only when a chunk is finished: ```py for event in pipeline.stream(...): - if event.path == "denoise": # the chunk loop, not "denoise.denoise_inner" + if event.path == "denoise": # an outer-loop event: a whole chunk is done show(event.state.get("out_frames")) ``` @@ -441,23 +441,6 @@ class DenoiseLoop(IterativePipelineBlocks): `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. From 532265485c1e9bbb37083822515620fe4a3af0fe Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 1 Sep 2026 10:34:50 +0200 Subject: [PATCH 16/24] polish loop docstrings and comments Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .ai/references/modular.md | 2 +- .../modular_pipelines/modular_pipeline.py | 22 ++++++++++--------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/.ai/references/modular.md b/.ai/references/modular.md index 907bd4ba3dc1..e57d1376efbc 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -188,7 +188,7 @@ Autoregressive video models nest loops: an outer segment loop (`loop_variables = The wrapper contains only the loop logic — how to iterate through the steps. Its `loop_inputs` are just what that takes (`timesteps`, `num_segments`); all data flows through the steps, which read and write the pipeline state directly. If the loop logic seems to need to do more than iterate — collect results, carry something to the next iteration — add a loop step for it instead (in `wan_animate_2`, a small `collect` step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back). Two small notes: -- A value written late in iteration `k` and read early in iteration `k + 1` would surface as a pipeline input (the first read has no writer before it). Seed it in a block that runs before the loop, declare it as an output, set it to `None`, and it stays internal. +- A value written late in iteration `k` and read early in iteration `k + 1` would become a pipeline input (the first read has no writer before it). Seed it in a block that runs before the loop, declare it as an output, set it to `None`, and it stays internal. - Components the loop logic itself uses (e.g. `flux2`'s wrapper reads `scheduler.order` to compute the progress-bar warmup steps) are added by overriding `expected_components`. Existing pipelines still use `LoopSequentialPipelineBlocks` (steps receive a shared flattened `block_state`, no nesting, no streaming). Leave them alone unless you are porting the pipeline; don't use it for new ones. diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 9ece2b9a2241..18d36336a582 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -679,9 +679,10 @@ 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. + state)`, the block receives the enclosing loop's variables as keyword call arguments — its signature names the + loop variables it uses and declares `**kwargs` for any it ignores (e.g. `def __call__(self, components, state, t, + **kwargs)`; naming all of them without `**kwargs` works too). The loop validates this at construction: a named + parameter that is not a loop variable, or a missing loop variable without a `**kwargs` catch-all, raises. > [!WARNING] > This is an experimental feature and is likely to change in the future. """ @@ -883,8 +884,8 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState: 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. + # Same branch selection as `__call__`. The branch does not add a level to event paths: "core_denoise.denoise", + # not "core_denoise.t2v.denoise"/"core_denoise.i2v.denoise" — so `event.path` filters keep working whichever branch runs. 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) @@ -1509,8 +1510,8 @@ def __call__(self, components, state, k): # accepts the OUTER chunk loop's vari ``` 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 + inputs (e.g. `timesteps`) and outputs are declared in `loop_inputs` / `loop_intermediate_outputs`: they join + 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`. @@ -1593,7 +1594,8 @@ def __init__(self): @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 + # `super().from_blocks_dict` runs `__init__` first and fills `sub_blocks` from `blocks_dict` after, + # so the validation in `__init__` never saw these blocks — validate them now instance._validate_sub_blocks() return instance @@ -3311,8 +3313,8 @@ def __call__(self, state: PipelineState = None, output: str | list[str] = None, 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 + Run the pipeline as a generator that yields a [`StreamEvent`] after every iteration of every loop block + (e.g. each denoising step, each segment of a chunked video), with the live [`PipelineState`] attached. The generator's return value is the final state, the same one `__call__` returns. Args: From 91c6aa50cf71b047d629624f0b1d9704cc1befd7 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 1 Sep 2026 10:49:44 +0200 Subject: [PATCH 17/24] rewrite iterative-blocks tests with simple arithmetic dummies Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .../test_iterative_pipeline_blocks.py | 246 ++++++++---------- 1 file changed, 114 insertions(+), 132 deletions(-) diff --git a/tests/modular_pipelines/test_iterative_pipeline_blocks.py b/tests/modular_pipelines/test_iterative_pipeline_blocks.py index b0090445b517..4700c32d88f7 100644 --- a/tests/modular_pipelines/test_iterative_pipeline_blocks.py +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -26,142 +26,140 @@ ) -# 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. 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. +# Dummy blocks with trivially checkable arithmetic, in the same nested shape as an autoregressive +# video pipeline: an outer loop (variable `k`) whose steps add 1 to `x` and record the result, with a +# nested inner loop (variable `i`) that multiplies `x` by 10 on each of its iterations. +# +# OuterLoop loop over k in range(num_outer_steps) +# ├─ add_one AddOneStep x += 1 +# ├─ times_ten TimesTenLoop the inner loop, over i in range(num_inner_steps) +# │ ├─ compute_delta delta = x * 9 +# │ └─ apply_delta x += delta (net effect of one inner iteration: x *= 10) +# └─ record RecordStep xs.append(x) -class ChunkNoiseGenStep(ModularLoopPipelineBlocks): +class AddOneStep(ModularLoopPipelineBlocks): model_name = "test" @property def inputs(self): - return [InputParam(name="history", required=True)] + return [InputParam(name="x", required=True)] @property def intermediate_outputs(self): - return [OutputParam(name="chunk_latents")] + return [OutputParam(name="x")] @property def description(self): - return "prepares this chunk's latents from the history" + return "adds 1 to x" def __call__(self, components, state, k): block_state = self.get_block_state(state) - block_state.chunk_latents = block_state.history + k + block_state.x = block_state.x + 1 self.set_block_state(state, block_state) return components, state -class LoopDenoiserStep(ModularLoopPipelineBlocks): +class ComputeDeltaStep(ModularLoopPipelineBlocks): model_name = "test" @property def inputs(self): - return [InputParam(name="chunk_latents", required=True)] + return [InputParam(name="x", required=True)] @property def intermediate_outputs(self): - return [OutputParam(name="noise_pred")] + return [OutputParam(name="delta")] @property def description(self): - return "predicts the noise for one timestep" + return "computes this inner iteration's increment" - def __call__(self, components, state, i, t): + def __call__(self, components, state, i): block_state = self.get_block_state(state) - block_state.noise_pred = block_state.chunk_latents * 0 + t + block_state.delta = block_state.x * 9 self.set_block_state(state, block_state) return components, state -class LoopSchedulerStep(ModularLoopPipelineBlocks): +class ApplyDeltaStep(ModularLoopPipelineBlocks): model_name = "test" @property def inputs(self): - return [InputParam(name="chunk_latents", required=True), InputParam(name="noise_pred", required=True)] + return [InputParam(name="x", required=True), InputParam(name="delta", required=True)] @property def intermediate_outputs(self): - return [OutputParam(name="chunk_latents")] + return [OutputParam(name="x")] @property def description(self): - return "updates the chunk latents with the noise prediction" + return "applies the increment to x" - def __call__(self, components, state, i, t): + def __call__(self, components, state, i): block_state = self.get_block_state(state) - block_state.chunk_latents = block_state.chunk_latents + block_state.noise_pred + block_state.x = block_state.x + block_state.delta 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. - - 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. - """ - +class InnerLoopWrapper(IterativePipelineBlocks): model_name = "test" - block_classes = [LoopDenoiserStep, LoopSchedulerStep] - block_names = ["denoiser", "scheduler"] @property def description(self): - return "inner timestep loop" + return "inner loop over num_inner_steps" @property def loop_variables(self): - return ["i", "t"] + return ["i"] @property def loop_inputs(self): - return [InputParam(name="timesteps", required=True)] + return [InputParam(name="num_inner_steps", required=True)] @torch.no_grad() - def __call__(self, components, state, k): + def __call__(self, components, state, **kwargs): # ignores the outer loop's `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) + for i in range(block_state.num_inner_steps): + components, state = self.loop_step(components, state, i=i) return components, state -class ChunkUpdateStep(ModularLoopPipelineBlocks): +class TimesTenLoop(InnerLoopWrapper): + block_classes = [ComputeDeltaStep, ApplyDeltaStep] + block_names = ["compute_delta", "apply_delta"] + + +class RecordStep(ModularLoopPipelineBlocks): model_name = "test" @property def inputs(self): - return [InputParam(name="chunk_latents", required=True), InputParam(name="latent_chunks", default=None)] + return [InputParam(name="x", required=True), InputParam(name="xs", default=None)] @property def intermediate_outputs(self): - return [OutputParam(name="history"), OutputParam(name="latent_chunks")] + return [OutputParam(name="xs")] @property def description(self): - return "records the denoised chunk and updates the history" + return "records x after this outer iteration" def __call__(self, components, state, **kwargs): # ignores the loop's `k`: a catch-all is enough 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)] + block_state.xs = [*(block_state.xs or []), float(block_state.x)] 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.""" - +class OuterLoopWrapper(IterativePipelineBlocks): model_name = "test" - block_classes = [ChunkNoiseGenStep, InnerDenoiseLoop, ChunkUpdateStep] - block_names = ["noise_gen", "denoise", "update"] @property def description(self): - return "outer autoregressive chunk loop" + return "outer loop over num_outer_steps" @property def loop_variables(self): @@ -169,116 +167,119 @@ def loop_variables(self): @property def loop_inputs(self): - return [InputParam(name="num_latent_chunk", required=True)] + return [InputParam(name="num_outer_steps", required=True)] @torch.no_grad() def __call__(self, components, state): block_state = self.get_block_state(state) - for k in range(block_state.num_latent_chunk): + for k in range(block_state.num_outer_steps): components, state = self.loop_step(components, state, k=k) return components, state -class CollectingChunkLoop(ChunkLoop): - """Chunk loop whose loop logic has an output of its own, written through `set_block_state`, and which - observes a step output (`history`) by declaring it in `loop_inputs` and calling `get_block_state` again - after each `loop_step` — the fresh snapshot picks up what the steps just wrote.""" +class OuterLoop(OuterLoopWrapper): + block_classes = [AddOneStep, TimesTenLoop, RecordStep] + block_names = ["add_one", "times_ten", "record"] + + +class TimestepLoopWrapper(IterativePipelineBlocks): + """Loop wrapper over `i`, `t` — the signature-validation tests assemble it with different steps.""" + + model_name = "test" @property - def loop_inputs(self): - return [*super().loop_inputs, InputParam(name="history")] + def description(self): + return "loop over i, t" + + @property + def loop_variables(self): + return ["i", "t"] @property - def loop_intermediate_outputs(self): - return [OutputParam(name="chunk_history")] + 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) - chunk_history = [] - for k in range(block_state.num_latent_chunk): - components, state = self.loop_step(components, state, k=k) - block_state = self.get_block_state(state) - chunk_history.append(float(block_state.history)) - block_state.chunk_history = chunk_history - self.set_block_state(state, block_state) + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) return components, state +# OuterLoop loop over k in range(num_outer_steps) +# ├─ add_one AddOneStep x += 1 +# ├─ times_ten TimesTenLoop the inner loop, over i in range(num_inner_steps) +# │ ├─ compute_delta delta = x * 9 +# │ └─ apply_delta x += delta (net effect of one inner iteration: x *= 10) +# └─ record RecordStep xs.append(x) + + class TestIterativePipelineBlocksStructure: def test_inputs_aggregation(self): - loop = ChunkLoop() + loop = OuterLoop() input_names = [p.name for p in loop.inputs] - # 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 + # the outer loop logic's own input, declared in its `loop_inputs` + assert "num_outer_steps" in [p.name for p in loop.loop_inputs] + assert "num_outer_steps" in input_names + # the nested inner loop's `loop_inputs` entry is aggregated too + assert "num_inner_steps" in [p.name for p in loop.sub_blocks["times_ten"].loop_inputs] + assert "num_inner_steps" in input_names # 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 - # cross-chunk carries surface as (optional) iteration-0 seeds - assert "history" in input_names - assert "latent_chunks" in input_names + # `x` is read by the first step (`add_one`) before any step writes it -> a pipeline input + assert "x" in input_names + # `xs` is the accumulator: written by `record` at the end of iteration k, read by it again at + # k + 1 — the read comes first, so it is an (optional, default None) pipeline input + assert "xs" in input_names + # `delta` is written by `compute_delta` before `apply_delta` reads it -> satisfied inside the loop + assert "delta" not in input_names def test_sub_block_outputs_are_aggregated(self): - loop = ChunkLoop() + loop = OuterLoop() output_names = [o.name for o in loop.intermediate_outputs] - 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 + assert "x" in output_names + assert "xs" in output_names + assert "delta" 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"] + loop = OuterLoop() + assert isinstance(loop.sub_blocks["times_ten"], IterativePipelineBlocks) + assert list(loop.sub_blocks["times_ten"].sub_blocks) == ["compute_delta", "apply_delta"] class TestIterativePipelineBlocksExecution: def _make_pipeline(self): - return SequentialPipelineBlocks.from_blocks_dict({"chunks": ChunkLoop()}).init_pipeline() + return SequentialPipelineBlocks.from_blocks_dict({"loop": OuterLoop()}).init_pipeline() - def test_nested_chunk_loop(self): + def test_nested_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)) + # per outer step: x += 1, then the inner loop doubles the digits (x *= 10 per inner step), then record + # k=0: (0 + 1) * 10 * 10 = 100 ; k=1: (100 + 1) * 10 * 10 = 10100 + state = pipe(x=torch.tensor(0.0), num_outer_steps=2, num_inner_steps=2) - 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 + assert state.get("xs") == [100.0, 10100.0] + # the carried value persists as a declared output + assert float(state.get("x")) == 10100.0 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)) + state = pipe(x=torch.tensor(0.0), num_outer_steps=2, num_inner_steps=1) - for name in ("k", "i", "t"): + for name in ("k", "i"): 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 + assert state.get("delta") 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)) - # `history` is re-read with `get_block_state` after every iteration — without the re-read this - # would collect the stale iteration-0 seed, [0.0, 0.0, 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] + state = pipe(x=torch.tensor(0.0), num_outer_steps=2, num_inner_steps=1) + block_state = pipe.blocks.sub_blocks["loop"].get_block_state(state) + assert block_state.as_dict().keys() == {"num_outer_steps"} def test_sub_block_type_is_validated(self): # a regular ModularPipelineBlocks cannot be a loop sub-block: fails at construction @@ -306,30 +307,11 @@ def description(self): @staticmethod def _loop_over(step_cls): - class Loop(IterativePipelineBlocks): - model_name = "test" + # assemble the shared `TimestepLoopWrapper` with the given step + class Loop(TimestepLoopWrapper): block_classes = [step_cls] block_names = ["step"] - @property - def description(self): - return "loop over i, t" - - @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 - return Loop def test_leaf_signature_is_validated(self): @@ -377,6 +359,6 @@ def __call__(self, components, state, t, **kwargs): 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() + pipe = SequentialPipelineBlocks.from_blocks_dict({"compute_delta": ComputeDeltaStep()}).init_pipeline() with pytest.raises(TypeError): - pipe(chunk_latents=torch.tensor(1.0)) + pipe(x=torch.tensor(1.0)) From fab047627e76118a93ca77cfeb5bb1a7cda2298a Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 1 Sep 2026 10:54:00 +0200 Subject: [PATCH 18/24] make style Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .../modular_pipelines/modular_pipeline.py | 30 +++++++++---------- src/diffusers/pipelines/ltx2/__init__.py | 2 +- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 18d36336a582..30ad5f0704c5 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -679,8 +679,8 @@ 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 receives the enclosing loop's variables as keyword call arguments — its signature names the - loop variables it uses and declares `**kwargs` for any it ignores (e.g. `def __call__(self, components, state, t, + state)`, the block receives the enclosing loop's variables as keyword call arguments — its signature names the loop + variables it uses and declares `**kwargs` for any it ignores (e.g. `def __call__(self, components, state, t, **kwargs)`; naming all of them without `**kwargs` works too). The loop validates this at construction: a named parameter that is not a loop variable, or a missing loop variable without a `**kwargs` catch-all, raises. @@ -1490,10 +1490,10 @@ def __call__(self, components, state): Loop variables are passed to sub-blocks as keyword call arguments: a sub-block's `__call__` names the loop variables it uses after `(components, state)` and declares `**kwargs` for any it ignores (naming all of them and - omitting `**kwargs` is fine too). This is validated at construction: a named parameter that is not a loop - variable, or a missing loop variable without a `**kwargs` catch-all, raises. 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: + omitting `**kwargs` is fine too). This is validated at construction: a named parameter that is not a loop variable, + or a missing loop variable without a `**kwargs` catch-all, raises. 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): @@ -1510,11 +1510,11 @@ def __call__(self, components, state, k): # accepts the OUTER chunk loop's vari ``` 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 join - 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`. + inputs (e.g. `timesteps`) and outputs are declared in `loop_inputs` / `loop_intermediate_outputs`: they join 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 @@ -1601,8 +1601,8 @@ def from_blocks_dict(cls, blocks_dict, description: str | None = None) -> "Itera def _validate_sub_blocks(self): """Sub-blocks must be loop steps (`ModularLoopPipelineBlocks`) or nested loops (`IterativePipelineBlocks`) - and accept the loop variables after `(components, state)` — either all of them by name, or the ones the - step uses plus a `**kwargs` catch-all for the rest.""" + and accept the loop variables after `(components, state)` — either all of them by name, or the ones the step + uses plus a `**kwargs` catch-all for the rest.""" expected = set(self.loop_variables) for block_name, block in self.sub_blocks.items(): if not isinstance(block, (ModularLoopPipelineBlocks, IterativePipelineBlocks)): @@ -3313,8 +3313,8 @@ def __call__(self, state: PipelineState = None, output: str | list[str] = None, def stream(self, state: PipelineState = None, **kwargs): """ - Run the pipeline as a generator that yields a [`StreamEvent`] after every iteration of every loop block - (e.g. each denoising step, each segment of a chunked video), with the live [`PipelineState`] attached. The + Run the pipeline as a generator that yields a [`StreamEvent`] after every iteration of every loop block (e.g. + each denoising step, each segment of a chunked video), with the live [`PipelineState`] attached. The generator's return value is the final state, the same one `__call__` returns. Args: diff --git a/src/diffusers/pipelines/ltx2/__init__.py b/src/diffusers/pipelines/ltx2/__init__.py index d4aa35127403..d48d890f4cb6 100644 --- a/src/diffusers/pipelines/ltx2/__init__.py +++ b/src/diffusers/pipelines/ltx2/__init__.py @@ -30,12 +30,12 @@ _import_structure["pipeline_ltx2_condition"] = ["LTX2ConditionPipeline", "LTX2VideoCondition"] _import_structure["pipeline_ltx2_dfr"] = ["LTX2DFRPipeline"] _import_structure["pipeline_ltx2_dfr_temporal_refine"] = ["LTX2DFRTemporalRefinePipeline"] - _import_structure["pipeline_output"] = ["LTX2DFRPipelineOutput", "LTX2PipelineOutput", "LTX2VideoDecodeOutput"] _import_structure["pipeline_ltx2_diffusion_decode"] = ["LTX2VideoDiffusionDecodePipeline"] _import_structure["pipeline_ltx2_hdr_lora"] = ["LTX2HDRPipeline", "LTX2HDRReferenceCondition"] _import_structure["pipeline_ltx2_ic_lora"] = ["LTX2InContextPipeline", "LTX2ReferenceCondition"] _import_structure["pipeline_ltx2_image2video"] = ["LTX2ImageToVideoPipeline"] _import_structure["pipeline_ltx2_latent_upsample"] = ["LTX2LatentUpsamplePipeline"] + _import_structure["pipeline_output"] = ["LTX2DFRPipelineOutput", "LTX2PipelineOutput", "LTX2VideoDecodeOutput"] _import_structure["vocoder"] = ["LTX2Vocoder", "LTX2VocoderWithBWE"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: From b6274dc216e4d0cf566e2af250f5635dc26ba1b6 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 1 Sep 2026 11:16:58 +0200 Subject: [PATCH 19/24] loop docs: three-part recipe, concrete data-flow examples, fewer dashes Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .ai/references/modular.md | 75 ++++++++++++------- .../iterative_pipeline_blocks.md | 13 +++- 2 files changed, 60 insertions(+), 28 deletions(-) diff --git a/.ai/references/modular.md b/.ai/references/modular.md index e57d1376efbc..ba0638571466 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -125,7 +125,33 @@ for i, t in enumerate(timesteps): ## Key pattern: Denoising loop -The denoising loop is an `IterativePipelineBlocks` whose steps are `ModularLoopPipelineBlocks` (see `flux2/denoise.py` as example). The wrapper declares the loop variables it passes to its steps, the inputs its own loop logic reads (in addition to the inputs/outputs of the loop steps), and the loop in `__call__`; `stream` is the same loop as a generator over `stream_step` and is what makes `pipe.stream(...)` work — implement both: +The denoising loop is an `IterativePipelineBlocks` whose steps are `ModularLoopPipelineBlocks` (see `flux2/denoise.py` as example). Building one takes three parts: write the loop steps, write the wrapper, assemble. + +**1. Write the loop steps.** Regular blocks (own `inputs` / `intermediate_outputs`, `get_block_state` / `set_block_state` on the full `PipelineState`) whose `__call__` takes the loop variables as extra keyword arguments — name the ones the step uses, declare `**kwargs` for any it ignores (validated at construction: unknown names raise, and so does a missing variable without a `**kwargs` catch-all): + +```python +class MyModelLoopAfterDenoiser(ModularLoopPipelineBlocks): + @property + def expected_components(self): + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self): + return [InputParam("latents", required=True), InputParam("noise_pred", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam("latents")] + + @torch.no_grad() + def __call__(self, components, state, i, t): + block_state = self.get_block_state(state) + block_state.latents = components.scheduler.step(block_state.noise_pred, t, block_state.latents, return_dict=False)[0] + self.set_block_state(state, block_state) + return components, state +``` + +**2. Write the wrapper.** The loop logic only: `loop_variables` (passed to every step), `loop_inputs` (what the loop logic itself reads), and the loop in `__call__`; `stream` is the same loop as a generator over `stream_step` and is what makes `pipe.stream(...)` work — implement both: ```python class MyModelDenoiseLoopWrapper(IterativePipelineBlocks): @@ -152,44 +178,39 @@ class MyModelDenoiseLoopWrapper(IterativePipelineBlocks): for i, t in enumerate(block_state.timesteps): components, state = yield from self.stream_step(components, state, i=i, t=t) return components, state +``` +**3. Assemble.** Attach the steps in a subclass (or at runtime with `from_blocks_dict`): +```python class MyModelDenoiseStep(MyModelDenoiseLoopWrapper): block_classes = [MyModelLoopDenoiser, MyModelLoopAfterDenoiser] block_names = ["denoiser", "after_denoiser"] ``` -Loop steps are regular blocks (own `inputs` / `intermediate_outputs`, `get_block_state` / `set_block_state` on the full `PipelineState`) whose `__call__` takes the loop variables as extra keyword arguments — name the ones the step uses, declare `**kwargs` for any it ignores (validated at construction: unknown names raise, and so does a missing variable without a `**kwargs` catch-all). A minimal one: - -```python -class MyModelLoopAfterDenoiser(ModularLoopPipelineBlocks): - @property - def expected_components(self): - return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] - - @property - def inputs(self): - return [InputParam("latents", required=True), InputParam("noise_pred", required=True)] +Autoregressive video models nest loops: an outer segment loop (`loop_variables = ["k"]`) whose steps prepare the segment, run the full inner denoising loop, and update the history (see `wan_animate_2/denoise.py` as an example). The inner loop is a step of the outer one, so its `__call__` / `stream` accept the outer variable (`def __call__(self, components, state, k)`) and pass its own `i`, `t` to its own steps. - @property - def intermediate_outputs(self): - return [OutputParam("latents")] +The wrapper should contain only the loop logic, i.e. how to iterate through the steps, and its `loop_inputs` are just what that takes (e.g. `timesteps`, `num_segments`). All data flows through the steps, which read and write the pipeline state directly: in the example above, the denoiser step writes `noise_pred` to the pipeline state and `MyModelLoopAfterDenoiser` reads it back and writes the updated `latents`; the wrapper touches none of them. If the loop logic seems to need to do more than just iterate (e.g. collect results, carry something to the next iteration), add a loop step for it instead. In `wan_animate_2`, a small `collect` step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back. What we don't want is the wrapper doing it inline: - @torch.no_grad() - def __call__(self, components, state, i, t): - block_state = self.get_block_state(state) - block_state.latents = components.scheduler.step(block_state.noise_pred, t, block_state.latents, return_dict=False)[0] - self.set_block_state(state, block_state) - return components, state +```python +# don't: loop logic collecting results itself +def __call__(self, components, state): + block_state = self.get_block_state(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")) # reaching into the state: make this a `collect` loop step + ... ``` -Autoregressive video models nest loops: an outer segment loop (`loop_variables = ["k"]`) whose steps prepare the segment, run the full inner denoising loop, and update the history (see `wan_animate_2/denoise.py` as an example). The inner loop is a step of the outer one, so its `__call__` / `stream` accept the outer variable — `def __call__(self, components, state, k)` — and pass its own `i`, `t` to its own steps. - -The wrapper contains only the loop logic — how to iterate through the steps. Its `loop_inputs` are just what that takes (`timesteps`, `num_segments`); all data flows through the steps, which read and write the pipeline state directly. If the loop logic seems to need to do more than iterate — collect results, carry something to the next iteration — add a loop step for it instead (in `wan_animate_2`, a small `collect` step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back). - Two small notes: -- A value written late in iteration `k` and read early in iteration `k + 1` would become a pipeline input (the first read has no writer before it). Seed it in a block that runs before the loop, declare it as an output, set it to `None`, and it stays internal. -- Components the loop logic itself uses (e.g. `flux2`'s wrapper reads `scheduler.order` to compute the progress-bar warmup steps) are added by overriding `expected_components`. +- A value written late in iteration `k` and read early in iteration `k + 1` becomes a pipeline input, since its first read has no writer before it. if a user-supplied value isn't meaningful, seed it in a block that runs before the loop instead: declare it as an output and set it to `None`, and it stays internal. + + ``` + loop, iteration k: step_1: x = xs[-1] <- first read of `xs`, no writer before it: + step_2: xs.append(x * 10) `xs` becomes a pipeline input unless a previous block produce it + ``` +- Components the loop logic itself uses are added by overriding `expected_components` (e.g. `flux2`'s wrapper reads `scheduler.order` to compute the progress-bar warmup steps). Existing pipelines still use `LoopSequentialPipelineBlocks` (steps receive a shared flattened `block_state`, no nesting, no streaming). Leave them alone unless you are porting the pipeline; don't use it for new ones. diff --git a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md index ec347ace19ac..824678b47ad6 100644 --- a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md @@ -136,7 +136,18 @@ state = pipeline(latents=torch.tensor(0.0), timesteps=torch.tensor([1.0, 2.0, 3. state.get("latents") # tensor(6.) — 0 + 1 + 2 + 3 ``` -The wrapper contains only the loop logic, i.e. how to iterate through its steps, so its `loop_inputs` should be just what that takes (the `timesteps` above). All data flows through the steps, which read and write the pipeline state directly. If the loop logic seems to need to do more than iterate: e.g. collect results, you should add a loop step for it instead: in `wan_animate_2`, a small collect step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back to condition on; under [streaming](#streaming), the partial collection is visible after every iteration. +The wrapper contains only the loop logic, i.e. how to iterate through its steps, so its `loop_inputs` should be just what that takes (the `timesteps` above). All data flows through the steps, which read and write the pipeline state directly: in the example above, `DenoiserStep` writes `noise_pred` to the state and `SchedulerStep` reads it back and writes the updated `latents`; the wrapper touches none of them. If the loop logic seems to need to do more than iterate: e.g. collect results, you should add a loop step for it instead: in `wan_animate_2`, a small collect step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back to condition on; under [streaming](#streaming), the partial collection is visible after every iteration. What we don't want is the wrapper doing it inline: + +```py +# don't: loop logic collecting results itself +def __call__(self, components, state): + block_state = self.get_block_state(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")) # reaching into the state: make this a collect loop step + ... +``` ## Nesting loops From 1c85e0ce7a7484932852ad547f209baa1e7b9a19 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Wed, 2 Sep 2026 13:17:09 -1000 Subject: [PATCH 20/24] Update docs/source/en/modular_diffusers/sequential_pipeline_blocks.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/modular_diffusers/sequential_pipeline_blocks.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md index d9c093b68703..d9c061e85950 100644 --- a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md @@ -115,7 +115,7 @@ When you create a [`~modular_pipelines.SequentialPipelineBlocks`], properties li ### Aggregated inputs and outputs -Aggregation follows the order the sub-blocks run in. `inputs` walks the sub-blocks in order and collects each one's inputs, skipping any that an earlier sub-block already declares in its `intermediate_outputs` — that value is produced inside the assembled block, so it isn't asked of the caller. `intermediate_outputs` is the union of what the sub-blocks write. +Aggregation follows the order the sub-blocks run in. `inputs` walks the sub-blocks in order and omits an input when an earlier sub-block declares an output with the same name, because the assembled block produces that value internally. `intermediate_outputs` is the union of what the sub-blocks write. This is what makes blocks composable: an assembled block has the same kind of input/output contract as a single [`~modular_pipelines.ModularPipelineBlocks`], so it can in turn be a sub-block of another one, at any depth. It is also why an input disappears from a pipeline's signature once some earlier block produces it — see [nesting loops](./iterative_pipeline_blocks#nesting-loops) for a case where you use that deliberately. From 75857a4837b6bbdcca59e5e210ee85d777f1bfbd Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Wed, 2 Sep 2026 13:17:58 -1000 Subject: [PATCH 21/24] Update docs/source/en/modular_diffusers/modular_pipeline.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/modular_diffusers/modular_pipeline.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index e0a2e8402201..a5744c2edc95 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -394,7 +394,7 @@ If pipeline stages share components (e.g., the same VAE used for encoding and de ## Streaming -[`~ModularPipeline.stream`] runs the same pipeline as a generator. It yields a [`StreamEvent`] after every iteration of every loop block, e.g. 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. +[`~ModularPipeline.stream`] runs the same pipeline as a generator. It yields a [`StreamEvent`] after every iteration of every loop block, e.g. 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. When exhausted, the generator returns the final [`PipelineState`], the default return value of `__call__`. `stream()` does not support `output= selection`. ```py generator = pipeline.stream(prompt="a cat", num_inference_steps=20) From 13fe60d35e38ba4d8dd0c98e3f1b0596e2f3c8e2 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Wed, 2 Sep 2026 13:18:23 -1000 Subject: [PATCH 22/24] Update docs/source/en/modular_diffusers/modular_pipeline.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/modular_diffusers/modular_pipeline.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index a5744c2edc95..bfda81a9e501 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -403,7 +403,7 @@ for event in generator: 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 (e.g. an autoregressive video that denoises one chunk at a time), you receive events from both loops: one after every denoising step of the inner loop (path `"denoise.denoise_inner"`), and one after each completed chunk of the outer loop (path `"denoise"`). Check `event.path` to tell them apart — for example, to react only when a chunk is finished: +`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), you receive events from both loops. One after every denoising step of the inner loop (path `"denoise.denoise_inner"`), and one after each completed chunk of the outer loop (path `"denoise"`). Check `event.path` to tell them apart, for example, to react only when a chunk is finished: ```py for event in pipeline.stream(...): From 1e97ff039e93a8034578b2d37bbbd9e53d8bd7f8 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Wed, 2 Sep 2026 13:34:58 -1000 Subject: [PATCH 23/24] Apply batched suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- .ai/references/modular.md | 4 +- .../iterative_pipeline_blocks.md | 42 +++++++++++-------- .../en/modular_diffusers/modular_pipeline.md | 6 +-- 3 files changed, 30 insertions(+), 22 deletions(-) diff --git a/.ai/references/modular.md b/.ai/references/modular.md index ba0638571466..062f4bf69ed1 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -127,7 +127,9 @@ for i, t in enumerate(timesteps): The denoising loop is an `IterativePipelineBlocks` whose steps are `ModularLoopPipelineBlocks` (see `flux2/denoise.py` as example). Building one takes three parts: write the loop steps, write the wrapper, assemble. -**1. Write the loop steps.** Regular blocks (own `inputs` / `intermediate_outputs`, `get_block_state` / `set_block_state` on the full `PipelineState`) whose `__call__` takes the loop variables as extra keyword arguments — name the ones the step uses, declare `**kwargs` for any it ignores (validated at construction: unknown names raise, and so does a missing variable without a `**kwargs` catch-all): +**1. Write the loop steps.** Loop steps are regular blocks. They declare `inputs` and `intermediate_outputs` and use `get_block_state` / `set_block_state` with the full `PipelineState`. + +Their `__call__` also accepts the loop variables they use as keyword arguments. If a step ignores some loop variables, declare `**kwargs`. Otherwise, name every loop variable. The loop validates these signatures when it is constructed and raises for unknown or missing variables. ```python class MyModelLoopAfterDenoiser(ModularLoopPipelineBlocks): diff --git a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md index 824678b47ad6..a19ce667c5e6 100644 --- a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md @@ -12,18 +12,18 @@ specific language governing permissions and limitations under the License. # IterativePipelineBlocks -[`~modular_pipelines.IterativePipelineBlocks`] is a multi-block type that runs its sub-blocks multiple times. It is what we use to build a denoising loop: the sub-blocks predict the noise and step the scheduler, the loop runs them once per timestep. You can also nest one [`~modular_pipelines.IterativePipelineBlocks`] under another to build an autoregressive video pipeline that generates chunk after chunk. Every iteration can be [streamed](./modular_pipeline#streaming) to the caller as it completes. +[`~modular_pipelines.IterativePipelineBlocks`] is a multi-block type that runs its sub-blocks multiple times. It is used to build a denoising loop where the sub-blocks predict the noise and step the scheduler, then the loop runs them once per timestep. You can also nest one [`~modular_pipelines.IterativePipelineBlocks`] under another to build an autoregressive video pipeline that generates chunk after chunk. A loop can stream each iteration to the caller when its `stream` method is implemented. -This guide shows you how to write the loop steps, the loop itself, how to nest loops, and how values travel from one iteration to the next. +Use this block type when a pipeline must repeat a sequence of blocks while carrying state between iterations and, when implemented, streaming progress to the caller. > [!TIP] -> [`~modular_pipelines.IterativePipelineBlocks`] replaces [`~modular_pipelines.LoopSequentialPipelineBlocks`]; see [the last section](#loopsequentialpipelineblocks) for the differences. +> [`~modular_pipelines.IterativePipelineBlocks`] replaces [`~modular_pipelines.LoopSequentialPipelineBlocks`]; see the [LoopSequentialPipelineBlocks](#loopsequentialpipelineblocks) guide for the differences. ## Loop steps -A loop step is a [`~modular_pipelines.ModularLoopPipelineBlocks`]. It is a regular [`~modular_pipelines.ModularPipelineBlocks`] — it declares `inputs` and `intermediate_outputs`, and reads and writes the [`~modular_pipelines.PipelineState`] through `get_block_state` / `set_block_state` — with one difference: its `__call__` also receives the loop's *loop variables* as arguments. For example, a denoising loop can pass the step index `i` and the timestep `t`. +A loop step is a [`~modular_pipelines.ModularLoopPipelineBlocks`]. It is a regular [`~modular_pipelines.ModularPipelineBlocks`] — it declares `inputs` and `intermediate_outputs`, and reads and writes the [`~modular_pipelines.PipelineState`] through `get_block_state` / `set_block_state` — with one difference. Its `__call__` also receives the loop's *loop variables* as arguments. For example, a denoising loop can pass the step index `i` and the timestep `t`. -Loop variables are the loop's own bookkeeping, not pipeline data. They are local to the loop: the loop hands them to each step as plain call arguments for that one iteration, and they are never written to the [`~modular_pipelines.PipelineState`]. Anything that has to outlive the iteration goes through the state instead, like `noise_pred` and `latents` below. (A streaming consumer does see them: each [`~modular_pipelines.StreamEvent`] carries that iteration's values in `event.loop_kwargs`.) +Loop variables are local and acts as its own bookkeeping, not pipeline data. The loop hands them to each step as plain call arguments for that one iteration, and they are never written to the [`~modular_pipelines.PipelineState`]. Anything that has to outlive the iteration goes through the state instead, like `noise_pred` and `latents` below. A streaming consumer does see them, each [`~modular_pipelines.StreamEvent`] carries that iteration's values in `event.loop_kwargs`. ```py from diffusers.modular_pipelines import ModularLoopPipelineBlocks, InputParam, OutputParam @@ -72,17 +72,17 @@ class SchedulerStep(ModularLoopPipelineBlocks): return components, state ``` -Because each step works on the [`~modular_pipelines.PipelineState`], the values it writes (`noise_pred`, the updated `latents`) are visible to the next step in the same iteration and to the next iteration — `latents` is read at the start of every iteration and written at the end of it. +Because each step works on the [`~modular_pipelines.PipelineState`], the values it writes (`noise_pred`, the updated `latents`) are visible to the next step in the same iteration and to the next iteration. `latents` is read at the start of every iteration and written at the end of it. ## Loop wrapper The loop itself is a subclass of [`~modular_pipelines.IterativePipelineBlocks`]. It declares: -- `loop_variables`, the names of the variables it passes to its steps on every iteration (as keyword arguments). A step's `__call__` names the ones it uses after `(components, state)` and declares `**kwargs` for any it ignores — naming all of them works too. This is validated when the loop is constructed: a named parameter that isn't a loop variable, or a missing one without a `**kwargs` catch-all, raises. +- `loop_variables`, the names of the variables passed to each step on every iteration. Steps receive loop variables as keyword arguments. A step can name only the variables it uses if its `__call__` accepts `**kwargs`. Otherwise, it must name every loop variable. The loop validates these signatures when it is constructed and raises for unknown or missing variables. - `loop_inputs`, the inputs the loop logic itself reads — here the `timesteps` it iterates. They join the inputs aggregated from the steps (see below), and they are what `get_block_state` returns for the loop block. - `__call__`, the loop logic: read the loop's block state, and call `loop_step` once per iteration with the loop variables. `loop_step` runs every step once. -Note that the wrapper defines only the loop logic - which steps run inside it should be attached separately, in a subclass: +The wrapper defines only the loop logic. Which steps run inside it should be attached separately, in a subclass: ```py import torch @@ -116,13 +116,13 @@ class DenoiseLoop(DenoiseLoopWrapper): block_names = ["denoiser", "scheduler"] ``` -This separation means the same loop logic can work with different combination of loop steps: subclass the wrapper again with different `block_classes` (this is exactly how `Flux2DenoiseLoopWrapper` serves the base and klein denoise steps). Steps can also be attached to the loop with [`~modular_pipelines.IterativePipelineBlocks.from_blocks_dict`]: +This separation means the same loop logic can work with different combinations of loop steps. For example, subclass the wrapper again with different `block_classes`. Steps can also be attached to the loop with [`~modular_pipelines.IterativePipelineBlocks.from_blocks_dict`]: ```py loop = DenoiseLoopWrapper.from_blocks_dict({"denoiser": DenoiserStep(), "scheduler": SchedulerStep()}) ``` -You can also change what runs inside a loop you already have: add a step, reorder, swap one out. e.g. you can insert a logging step like this: +You can also change what runs inside a loop you already have: add a step, reorder, swap one out. For example, you can insert a logging step like this: ```py loop = DenoiseLoopWrapper.from_blocks_dict(loop.sub_blocks.copy().insert("log", LogStep(), 1)) @@ -136,7 +136,7 @@ state = pipeline(latents=torch.tensor(0.0), timesteps=torch.tensor([1.0, 2.0, 3. state.get("latents") # tensor(6.) — 0 + 1 + 2 + 3 ``` -The wrapper contains only the loop logic, i.e. how to iterate through its steps, so its `loop_inputs` should be just what that takes (the `timesteps` above). All data flows through the steps, which read and write the pipeline state directly: in the example above, `DenoiserStep` writes `noise_pred` to the state and `SchedulerStep` reads it back and writes the updated `latents`; the wrapper touches none of them. If the loop logic seems to need to do more than iterate: e.g. collect results, you should add a loop step for it instead: in `wan_animate_2`, a small collect step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back to condition on; under [streaming](#streaming), the partial collection is visible after every iteration. What we don't want is the wrapper doing it inline: +The wrapper contains only the loop logic, i.e. how to iterate through its steps, so its `loop_inputs` should be just what that takes (the `timesteps` above). All data flows through the steps, which read and write the pipeline state directly. In the example above, `DenoiserStep` writes `noise_pred` to the state and `SchedulerStep` reads it back and writes the updated `latents`. The wrapper touches none of them. If the loop logic seems to need to do more than iterate, for example, collect results, you should add a loop step for it instead. In `wan_animate_2`, a small collect step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back to condition on. Under [streaming](#streaming), the partial collection is visible after every iteration. What we don't want is the wrapper doing it inline: ```py # don't: loop logic collecting results itself @@ -151,9 +151,9 @@ def __call__(self, components, state): ## Nesting loops -An [`~modular_pipelines.IterativePipelineBlocks`] can be a step of another one. An autoregressive video pipeline generates a chunk of frames at a time, so it is an outer loop over chunks: each of its iterations prepares the chunk's latents from the frames generated so far, runs a full denoising loop over them, and appends the result to the history. +An [`~modular_pipelines.IterativePipelineBlocks`] can be a step of another one. An autoregressive video pipeline generates a chunk of frames at a time, so it is an outer loop over chunks. Each of its iterations prepares the chunk's latents from the frames generated so far, runs a full denoising loop over them, and appends the result to the history. -The inner denoising loop is a step of the outer loop, so its `__call__` must accept the outer loop's variables (or take `**kwargs` if it ignores them), and it declares `loop_variables` of its own for its own steps. The two sets are independent: `k` arrives as a call argument and the inner loop is free to use it — the wan-animate-2 denoise loop puts the chunk index in its progress bar — but it is not forwarded to the steps, which are passed the inner loop's `i` and `t`. +The inner denoising loop is a step of the outer loop, so its `__call__` must accept the outer loop's variables (or take `**kwargs` if it ignores them), and it declares `loop_variables` of its own for its own steps. The two sets are independent. `k` arrives as a call argument and the inner loop is free to use it — the wan-animate-2 denoise loop puts the chunk index in its progress bar — but it is not forwarded to the steps, which are passed the inner loop's `i` and `t`. This is the structure this section builds: @@ -195,7 +195,7 @@ class ChunkDenoiseLoop(ChunkDenoiseLoopWrapper): block_names = ["denoiser", "scheduler"] ``` -`ChunkDenoiseLoop` is the *inner* loop: it denoises a single chunk. It becomes a step of the outer loop below, which iterates over the chunks — at each `k` it prepares the chunk's latents from the history, runs the full inner denoising loop over them, and records the result: +`ChunkDenoiseLoop` is the *inner* loop that denoises a single chunk. It becomes a step of the outer loop below, which iterates over the chunks. At each `k`, it prepares the chunk's latents from the history, runs the full inner denoising loop over them, and records the result: ```py class PrepareChunkStep(ModularLoopPipelineBlocks): @@ -278,11 +278,11 @@ state = pipeline(num_chunks=2, timesteps=torch.tensor([1.0, 2.0]), history=torch state.get("history") # tensor(7.) — chunk 0: 0 + 0 + 3 = 3, chunk 1: 3 + 1 + 3 = 7 ``` -`history` is carried from one iteration to the next: `UpdateHistoryStep` writes it at the end of one, `PrepareChunkStep` reads it at the start of the next. Because the reader comes before the writer, it is one of the loop's inputs — which is why `pipeline(history=...)` works above — and like any input it has to come from either the user or an earlier block. If seeding it isn't meaningful (a decoder cache, the previous chunk's frames), have the block that runs before the loop declare it as an output and set its initial value; it then drops out of the pipeline's signature. +`history` is carried from one iteration to the next. `UpdateHistoryStep` writes it at the end of one, and `PrepareChunkStep` reads it at the start of the next. Because the reader comes before the writer, it is one of the loop's inputs — which is why `pipeline(history=...)` works above — and like any input it has to come from either the user or an earlier block. If seeding it isn't meaningful (a decoder cache, the previous chunk's frames), have the block that runs before the loop declare it as an output and set its initial value. It then drops out of the pipeline's signature. ## Streaming -To let [`~ModularPipeline.stream`] hand back the live state 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 [`~modular_pipelines.StreamEvent`] for it (after the events of any nested loop): +To let [`~ModularPipeline.stream`] hand back the live state after every iteration, also implement `stream`. This is the same loop written as a generator over `stream_step`, which runs one iteration like `loop_step` and additionally yields a [`~modular_pipelines.StreamEvent`] for it (after the events of any nested loop): ```py class DenoiseLoop(IterativePipelineBlocks): @@ -295,8 +295,14 @@ class DenoiseLoop(IterativePipelineBlocks): return components, state ``` -A nested loop's `stream` takes the outer loop's variables exactly like its `__call__` does. Streaming is opt-in per loop: `pipeline.blocks.supports_streaming` tells you whether every loop on the path implements it. See [Streaming](./modular_pipeline#streaming) for the consumer side. +A nested loop's `stream` takes the outer loop's variables exactly like its `__call__` does. Streaming is opt-in per loop.`pipeline.blocks.supports_streaming` tells you whether every loop on the path implements it. See [Streaming](./modular_pipeline#streaming) for the consumer side. ## LoopSequentialPipelineBlocks -[`~modular_pipelines.LoopSequentialPipelineBlocks`] is the earlier loop type and is still used by existing pipelines. It differs in three ways: its steps share one flattened [`~modular_pipelines.BlockState`] that the wrapper extracts before the loop (instead of each step reading the [`~modular_pipelines.PipelineState`] itself), it cannot contain another loop, and it cannot stream. Use [`~modular_pipelines.IterativePipelineBlocks`] for new pipelines. +[`~modular_pipelines.LoopSequentialPipelineBlocks`] is the earlier loop type and is still used by existing pipelines. It differs in three ways: + +1. Steps share one flattened [`~modular_pipelines.BlockState`] that the wrapper extracts before the loop (instead of each step reading the [`~modular_pipelines.PipelineState`] itself). +2. It cannot contain another loop. +3. It cannot stream. + +Use [`~modular_pipelines.IterativePipelineBlocks`] for new pipelines. diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index bfda81a9e501..864687644054 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -411,9 +411,9 @@ for event in pipeline.stream(...): 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. +To stop early, stop iterating (or call `generator.close()`). Blocks without loops run to completion and yield nothing. -Streaming is opt-in per loop block. An [`IterativePipelineBlocks`] (see the [IterativePipelineBlocks](./iterative_pipeline_blocks) guide) 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: +Streaming is opt-in per loop block. An [`IterativePipelineBlocks`] (see the [IterativePipelineBlocks](./iterative_pipeline_blocks) guide) implements its loop in `__call__` as usual and, to support streaming, also implements `stream`. This is 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): @@ -439,7 +439,7 @@ class DenoiseLoop(IterativePipelineBlocks): 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`). +`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`). ## Modular repository From c406d5b3af3400f1bcc1db4f7925dd48e59731b5 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 3 Sep 2026 01:50:26 +0200 Subject: [PATCH 24/24] address more feedbacks Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .../iterative_pipeline_blocks.md | 16 +++++++++++++++- .../en/modular_diffusers/modular_pipeline.md | 4 ++-- docs/source/en/modular_diffusers/overview.md | 2 +- 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md index a19ce667c5e6..26fe281246a2 100644 --- a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md @@ -16,6 +16,20 @@ specific language governing permissions and limitations under the License. Use this block type when a pipeline must repeat a sequence of blocks while carrying state between iterations and, when implemented, streaming progress to the caller. +Two classes are involved in building a loop. + +| Class | Role | State it works with | +|---|---|---| +| [`~modular_pipelines.ModularLoopPipelineBlocks`] | A *loop step*: a regular block that runs inside a loop, once per iteration | The full [`~modular_pipelines.PipelineState`], through its own `get_block_state` / `set_block_state`; also receives the loop variables as call arguments | +| [`~modular_pipelines.IterativePipelineBlocks`] | The *loop* itself: holds the loop logic and runs its steps once per iteration; loops can nest and stream | Its block state holds only its own `loop_inputs`; all data flows between its steps through the [`~modular_pipelines.PipelineState`] | + +This guide uses a few closely related terms, so let's define them upfront. + +- **Pipeline state** is the shared [`~modular_pipelines.PipelineState`] every block reads and writes. It is the only place where data crosses blocks and survives an iteration. +- **Block state** is one block's declared view of the pipeline state: what its `get_block_state` returns and its `set_block_state` writes back. For a loop wrapper, it holds only the loop's `loop_inputs`. +- **Loop inputs** are the loop wrapper's own input declaration: what the loop logic itself reads to drive the iteration (the `timesteps` it iterates over, for example). +- **Loop variables** are the per-iteration values the loop passes to its steps as call arguments (`i`, `t`). They are never written to the pipeline state. + > [!TIP] > [`~modular_pipelines.IterativePipelineBlocks`] replaces [`~modular_pipelines.LoopSequentialPipelineBlocks`]; see the [LoopSequentialPipelineBlocks](#loopsequentialpipelineblocks) guide for the differences. @@ -23,7 +37,7 @@ Use this block type when a pipeline must repeat a sequence of blocks while carry A loop step is a [`~modular_pipelines.ModularLoopPipelineBlocks`]. It is a regular [`~modular_pipelines.ModularPipelineBlocks`] — it declares `inputs` and `intermediate_outputs`, and reads and writes the [`~modular_pipelines.PipelineState`] through `get_block_state` / `set_block_state` — with one difference. Its `__call__` also receives the loop's *loop variables* as arguments. For example, a denoising loop can pass the step index `i` and the timestep `t`. -Loop variables are local and acts as its own bookkeeping, not pipeline data. The loop hands them to each step as plain call arguments for that one iteration, and they are never written to the [`~modular_pipelines.PipelineState`]. Anything that has to outlive the iteration goes through the state instead, like `noise_pred` and `latents` below. A streaming consumer does see them, each [`~modular_pipelines.StreamEvent`] carries that iteration's values in `event.loop_kwargs`. +Loop variables are local to the loop and act as its bookkeeping, not pipeline data. The loop hands them to each step as plain call arguments for that one iteration, and they are never written to the [`~modular_pipelines.PipelineState`]. Anything that has to outlive the iteration goes through the state instead, like `noise_pred` and `latents` below. A streaming consumer does see them, each [`~modular_pipelines.StreamEvent`] carries that iteration's values in `event.loop_kwargs`. ```py from diffusers.modular_pipelines import ModularLoopPipelineBlocks, InputParam, OutputParam diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 864687644054..87d503c52318 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -394,7 +394,7 @@ If pipeline stages share components (e.g., the same VAE used for encoding and de ## Streaming -[`~ModularPipeline.stream`] runs the same pipeline as a generator. It yields a [`StreamEvent`] after every iteration of every loop block, e.g. 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. When exhausted, the generator returns the final [`PipelineState`], the default return value of `__call__`. `stream()` does not support `output= selection`. +[`~ModularPipeline.stream`] runs the same pipeline as a generator. It yields a [`StreamEvent`] after every iteration of every loop block, e.g. 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. Exactly one event is emitted per loop iteration — the loop steps inside an iteration don't emit their own events, and blocks without loops emit none. When exhausted, the generator returns the final [`PipelineState`], the default return value of `__call__`. `stream()` does not support `output=` selection. ```py generator = pipeline.stream(prompt="a cat", num_inference_steps=20) @@ -411,7 +411,7 @@ for event in pipeline.stream(...): show(event.state.get("out_frames")) ``` -To stop early, stop iterating (or call `generator.close()`). Blocks without loops run to completion and yield nothing. +To stop early, stop iterating (or call `generator.close()`). Streaming is opt-in per loop block. An [`IterativePipelineBlocks`] (see the [IterativePipelineBlocks](./iterative_pipeline_blocks) guide) implements its loop in `__call__` as usual and, to support streaming, also implements `stream`. This is the same loop written as a generator over `stream_step`, which runs one iteration like `loop_step` and additionally yields the event for it: diff --git a/docs/source/en/modular_diffusers/overview.md b/docs/source/en/modular_diffusers/overview.md index 1c8d91a782ba..2530ace49a88 100644 --- a/docs/source/en/modular_diffusers/overview.md +++ b/docs/source/en/modular_diffusers/overview.md @@ -29,7 +29,7 @@ The Modular Diffusers docs are organized as shown below. - [ModularPipelineBlocks](./pipeline_block) is the most basic unit of a [`ModularPipeline`] and this guide shows you how to create one. - [SequentialPipelineBlocks](./sequential_pipeline_blocks) is a type of block that chains multiple blocks so they run one after another, passing data along the chain. This guide shows you how to create [`~modular_pipelines.SequentialPipelineBlocks`] and how they connect and work together. - [IterativePipelineBlocks](./iterative_pipeline_blocks) is a type of block that runs a series of blocks in a loop — a denoising loop, or an autoregressive chunk loop with a denoising loop nested inside. This guide shows you how to create [`~modular_pipelines.IterativePipelineBlocks`]. -- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks) is the earlier loop type, still used by existing pipelines. +- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks) is the earlier loop type, still used by existing pipelines. Prefer [`~modular_pipelines.IterativePipelineBlocks`] for new pipelines. - [AutoPipelineBlocks](./auto_pipeline_blocks) is a type of block that automatically chooses which blocks to run based on the input. This guide shows you how to create [`~modular_pipelines.AutoPipelineBlocks`]. - [Building Custom Blocks](./custom_blocks) shows you how to create your own custom blocks and share them on the Hub.