[Modular]improve loop block by yiyixuxu · Pull Request #14159 · huggingface/diffusers

yiyixuxu · 2026-07-10T09:50:10Z

This PR adds IterativePipelineblocks which is an improved version of LoopSequentialPipelineBlocks:

(1) loop steps receive the full PipelineState and declare their own block_state like regular blocks. Previously, the loop wrapper extracted one flattened BlockState before the loop, and all steps share that

(2)it can be nested and composed: e.g. you can nest a IterativePipelineblocks inside another to created a nested loop (autoregressive chunk loop)

(3) added first-class streaming support, an example

# a preview decoder from the pipeline's own blocks, with a fast VAE swapped in
preview = SequentialPipelineBlocks.from_blocks_dict(
    {"unpack": Flux2UnpackLatentsStep(), "decode": Flux2DecodeStep()}
).init_pipeline()
preview.update_components(vae=tiny_vae)

for event in pipe.stream(prompt="a cat", num_inference_steps=28):
    image = preview(
        latents=event.state.get("latents"),
        latent_ids=event.state.get("latent_ids"),
        output="images",
    )[0]  # show it

I refactored flux2 denoise loop as as example on how to use them

…tate 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 <noreply@anthropic.com>

HuggingFaceDocBuilderDev · 2026-07-10T09:57:25Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

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 <noreply@anthropic.com>

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 <noreply@anthropic.com>

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 <noreply@anthropic.com>

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 <noreply@anthropic.com>

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 <noreply@anthropic.com>

… every loop iteration - `stream(components, state)` generator on every block type; leaves yield nothing, composites re-yield sub-block events with the sub-block name prepended to `event.path` - `IterativePipelineBlocks`: streaming is opt-in — `__call__`/`loop_step` unchanged; a loop that supports it also implements `stream` as a generator over the new `stream_step`; base `stream` raises a clear NotImplementedError - `StreamEvent(path, block, loop_kwargs, state)` dataclass, exported - `supports_streaming` property on blocks (legacy `LoopSequentialPipelineBlocks` sets it to False; to be deprecated once all pipelines are ported) - `ModularPipeline.stream(state=None, **kwargs)` — same seeding as `__call__`, no_grad held only while blocks run; error logging mirrors `__call__` at every level - flux2: `Flux2DenoiseLoopWrapper.stream` next to its `__call__` - tests: `test_stream_matches_call` on `ModularPipelineTesterMixin` (auto-skips when blocks don't support streaming) - docs: Streaming section on the ModularPipeline page; regenerate dummy_pt_objects Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Conflict resolutions: - flux2/denoise.py: keep this branch's Flux2LoopAfterDenoiser — main's #14278 removed its unused inputs/intermediate_inputs, but the class is rewritten here as a ModularLoopPipelineBlocks that declares real inputs (latents, noise_pred) - test_modular_pipelines_common.py: deleted on main by #14444 (tests split up); test_stream_matches_call moved to its new home on ModularPipelineTesterMixin in tests/modular_pipelines/testing_utils/common.py Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- WanAnimate2SegmentLoopWrapper is an IterativePipelineBlocks (loop variable `k`); the per-segment steps become ModularLoopPipelineBlocks operating on the shared PipelineState - the hand-rolled timestep loop in WanAnimate2SegmentDenoiseInner becomes a nested loop: WanAnimate2LoopDenoiser + WanAnimate2LoopAfterDenoiser under WanAnimate2DenoiseLoopWrapper (loop variables `i`, `t`) - loop-carried state made explicit: `out_frames` is a declared input on the prev-frames step but filtered from the wrapper's aggregated inputs; `segment_frames` accumulation moves from the decode step into the segment loop's own logic - both loops implement `stream`: events per denoise step and per segment, so test_stream_matches_call now runs for wan-animate-2 (stream == call) - regenerate stale flux2 modular_blocks docstrings; doc-builder reflow in modular_pipeline.py Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- the 7 loop steps become ModularLoopPipelineBlocks operating on the shared PipelineState; the intra-iteration dataflow is now declared (before-denoiser outputs latent_model_input/timesteps, denoiser outputs noise_pred_video/ noise_pred_audio, after-denoisers output latents/audio_latents) and stays satisfied within the loop, so pipeline inputs are unchanged - LTX2DenoiseLoopWrapper is an IterativePipelineBlocks (loop variables `i`, `t`) and implements `stream`, so LTX-2 and LTX-2.5 stream with both video and audio latents live in every event; test_stream_matches_call now runs for all 8 ltx2/ltx25 testers - regenerate ltx2 modular_blocks docstrings (timesteps correctly optional at the core-step level, produced by set_timesteps) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

github-actions Bot added size/L PR with diff > 200 LOC tests modular-pipelines and removed size/L PR with diff > 200 LOC labels Jul 10, 2026

yiyixuxu and others added 5 commits July 10, 2026 16:06

github-actions Bot added the size/L PR with diff > 200 LOC label Jul 10, 2026

github-actions Bot added documentation Improvements or additions to documentation utils labels Aug 19, 2026

yiyixuxu mentioned this pull request Aug 19, 2026

[RFC]: Modular Diffusers adapter with per-step execution via greenlet vllm-project/vllm-omni#5822

Open

1 task

yiyixuxu and others added 4 commits August 19, 2026 18:33

Merge branch 'main' into refactor-iterative-loop-blocks

b00f89b

Navigation Menu

Uh oh!

[Modular]improve loop block - #14159

[Modular]improve loop block#14159
yiyixuxu wants to merge 11 commits into
mainfrom
refactor-iterative-loop-blocks

yiyixuxu commented Jul 10, 2026 •

edited

Loading

Uh oh!

HuggingFaceDocBuilderDev commented Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

yiyixuxu commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yiyixuxu commented Jul 10, 2026 •

edited

Loading