[Modular]improve loop block

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,

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,

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"