From b272422fc5d6bc78578d27a8aeb4b381ec06f820 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 09:34:34 +0000 Subject: [PATCH 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 6/6] 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