diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index dcccf5cd2de3..776b848eb307 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -342,6 +342,8 @@ "ConditionalPipelineBlocks", "ConfigSpec", "InputParam", + "IterativePipelineBlocks", + "ModularLoopPipelineBlocks", "LoopSequentialPipelineBlocks", "ModularPipeline", "ModularPipelineBlocks", @@ -1212,7 +1214,9 @@ ConditionalPipelineBlocks, ConfigSpec, InputParam, + IterativePipelineBlocks, LoopSequentialPipelineBlocks, + ModularLoopPipelineBlocks, ModularPipeline, ModularPipelineBlocks, OutputParam, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 25db2ef3bee2..09420b614beb 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -34,6 +34,8 @@ "AutoPipelineBlocks", "SequentialPipelineBlocks", "ConditionalPipelineBlocks", + "IterativePipelineBlocks", + "ModularLoopPipelineBlocks", "LoopSequentialPipelineBlocks", "PipelineState", "BlockState", @@ -160,7 +162,9 @@ AutoPipelineBlocks, BlockState, 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 1a782e70de33..8d7716fda436 100644 --- a/src/diffusers/modular_pipelines/flux2/denoise.py +++ b/src/diffusers/modular_pipelines/flux2/denoise.py @@ -22,9 +22,8 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging from ..modular_pipeline import ( - BlockState, - LoopSequentialPipelineBlocks, - ModularPipelineBlocks, + IterativePipelineBlocks, + ModularLoopPipelineBlocks, PipelineState, ) from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam @@ -42,7 +41,7 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -class Flux2LoopDenoiser(ModularPipelineBlocks): +class Flux2LoopDenoiser(ModularLoopPipelineBlocks): model_name = "flux2" @property @@ -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 @@ -103,10 +102,16 @@ def inputs(self) -> list[tuple[str, Any]]: ), ] + @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 + self, components: Flux2ModularPipeline, state: PipelineState, i: int, t: torch.Tensor ) -> 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 @@ -133,11 +138,12 @@ 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 -class Flux2KleinLoopDenoiser(ModularPipelineBlocks): +class Flux2KleinLoopDenoiser(ModularLoopPipelineBlocks): model_name = "flux2-klein" @property @@ -148,8 +154,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 @@ -192,10 +198,16 @@ def inputs(self) -> list[tuple[str, Any]]: ), ] + @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 + self, components: Flux2KleinModularPipeline, state: PipelineState, i: int, t: torch.Tensor ) -> 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 @@ -222,11 +234,12 @@ 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 -class Flux2KleinBaseLoopDenoiser(ModularPipelineBlocks): +class Flux2KleinBaseLoopDenoiser(ModularLoopPipelineBlocks): model_name = "flux2-klein" @property @@ -251,8 +264,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 +319,24 @@ 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.", + ), ] + @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 + self, components: Flux2KleinModularPipeline, state: PipelineState, i: int, t: torch.Tensor ) -> 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 @@ -356,10 +382,11 @@ 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): +class Flux2LoopAfterDenoiser(ModularLoopPipelineBlocks): model_name = "flux2" @property @@ -370,24 +397,36 @@ 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.", + ), + ] @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, 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, @@ -400,12 +439,26 @@ 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_variables(self) -> list[str]: + return ["i", "t"] + + @property + 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 description(self) -> str: return ( @@ -414,15 +467,11 @@ def description(self) -> str: ) @property - def loop_expected_components(self) -> list[ComponentSpec]: - return [ - ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), - ComponentSpec("transformer", Flux2Transformer2DModel), - ] - - @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, @@ -436,28 +485,27 @@ 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: 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() if XLA_AVAILABLE: xm.mark_step() - self.set_block_state(state, block_state) return components, state @@ -469,7 +517,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 +533,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 +549,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..c30201801ea3 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -550,6 +550,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] @@ -577,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): """ @@ -775,7 +815,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 +838,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 +1219,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 +1335,118 @@ def _requirements(self) -> dict[str, str]: return requirements +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: + + ```python + @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 + ``` + + 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. + + 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): + @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. 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. + + 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_variables(self) -> list[str]: + """Names of the loop variables `loop_step` passes to leaf sub-blocks each iteration (e.g. `["i", "t"]`).""" + return [] + + def __init__(self): + super().__init__() + 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_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( + 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__}`." + ) + 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.""" + for block_name, block in self.sub_blocks.items(): + try: + 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, **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") + + class LoopSequentialPipelineBlocks(ModularPipelineBlocks): """ A Pipeline blocks that combines multiple pipeline block classes into a For Loop. When called, it will call each @@ -1569,25 +1721,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..ad74c649d62c --- /dev/null +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -0,0 +1,302 @@ +# 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, + ModularLoopPipelineBlocks, + 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. 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(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="history", required=True)] + + @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, k): + block_state = self.get_block_state(state) + block_state.chunk_latents = block_state.history + k + self.set_block_state(state, block_state) + return components, state + + +class LoopDenoiserStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="chunk_latents", required=True)] + + @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, i, t): + block_state = self.get_block_state(state) + block_state.noise_pred = block_state.chunk_latents * 0 + t + self.set_block_state(state, block_state) + return components, state + + +class LoopSchedulerStep(ModularLoopPipelineBlocks): + 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, 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) + 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. + """ + + model_name = "test" + block_classes = [LoopDenoiserStep, LoopSchedulerStep] + block_names = ["denoiser", "scheduler"] + + @property + def description(self): + return "inner timestep loop" + + @property + def loop_variables(self): + return ["i", "t"] + + @property + def inputs(self): + return [InputParam(name="timesteps", required=True), *super().inputs] + + @torch.no_grad() + 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) + return components, state + + +class ChunkUpdateStep(ModularLoopPipelineBlocks): + 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, 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)] + 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_variables(self): + return ["k"] + + @property + def inputs(self): + return [InputParam(name="num_latent_chunk", required=True), *super().inputs] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(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_inputs_aggregation(self): + loop = ChunkLoop() + 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 + # 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 + + 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_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)) + + 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_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, not a loop step" + + def __call__(self, components, state): + return components, state + + class BadTypeLoop(IterativePipelineBlocks): + model_name = "test" + block_classes = [PlainStep] + block_names = ["plain"] + + @property + def description(self): + 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 at construction + 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): + return ["i", "t"] + + @property + def inputs(self): + return [InputParam(name="timesteps", required=True), *super().inputs] + + @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 + + with pytest.raises(ValueError, match="must accept the loop variables"): + BadSigLoop() + + 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(TypeError): + pipe(chunk_latents=torch.tensor(1.0))