Add step-end callbacks and interrupt to modular pipelines

[fusheng-ji]

What does this PR do?

Fixes #12386.

Modular pipelines had no step callbacks, and set_progress_bar_config didn't reach denoising loops nested inside other blocks (e.g. SDXL). This PR:

• adds callback_on_step_end and callback_on_step_end_tensor_inputs to ModularPipeline.__call__, with the same signature as the standard pipelines;
• stops denoising when pipe.interrupt = True, while still decoding the partial result;
• adds LoopSequentialPipelineBlocks.loop_over_timesteps, which runs the step, the callback and the interrupt check. The 15 timestep loops use it; the chunked loops (Cosmos3, Helios, MiniMax Music 3, Wan Animate 2) call the same pieces directly;
• makes each loop declare the fields a callback may read and replace (pipe.callback_tensor_inputs). step_index counts across chunks and stages;
• makes set_progress_bar_config walk all nested blocks.

The issue's first point (listing call inputs) is already covered by input_names and required_inputs, as noted on the issue.

Coordination: #12386 (@yiyixuxu: "we will add support for callback etc")

def stop_early(pipe, step_index, timestep, callback_kwargs):
    if step_index == 10:
        pipe.interrupt = True
    return {}

image = pipe(prompt="...", callback_on_step_end=stop_early, output="images")

Results

Real checkpoints, same seed, main vs this branch:

	SDXL base	FLUX.1-dev	Wan 2.1 1.3B	Helios-Mid
Output without callback vs main	bitwise identical	bitwise identical	bitwise identical	bitwise identical
Read-only callback	output unchanged, steps 0–29	unchanged, 0–27	unchanged, 0–29	unchanged, 0–59 across 2 chunks
Interrupt after step k (denoiser calls)	60 → 22	28 → 11	60 → 22	120 → 32
set_progress_bar_config(disable=True)	fixed	fixed	fixed	still shows (see below)

• In all four models, replacing latents from the callback changes the result, and the call after an interrupt matches a normal run bitwise.
• A no-op callback adds no measurable time: the median change was −3.5% to +2.6%, within run-to-run noise.
• Helios-Mid uses its published guider and steps.

Tests

pytest tests/modular_pipelines/test_modular_pipeline_callbacks.py
# 11 passed

pytest tests/modular_pipelines
# 1190 passed, 109 skipped, 3 xfailed, 19 failed
# all 19 also fail on main:
#   15 batch/single or fp16 tolerance tests
#   4 Hub-loading tests (TestLoadFromLocalCopy x3, test_custom_block_loads_from_hub) with HTTP errors in my environment

make style && make fix-copies   # no changes

The new SDXL test test_set_progress_bar_config_reaches_nested_loop fails on main and passes here.

Design question for reviewers

loop_over_timesteps puts the step, callback and interrupt logic in a base-class helper instead of repeating it in each loop. The chunked loops can't use it and call the same pieces directly, so there are two patterns. Happy to inline it into each loop instead if you prefer.

Known limitations

• Helios and Wan Animate 2 create their inner progress bars with plain tqdm, so set_progress_bar_config still doesn't reach them.
• Helios, MiniMax Music 3 and Wan Animate 2 expose only latents to callbacks, because their loops build the conditioning once per chunk.
• Unlike the standard pipelines, requesting or returning a field the loop doesn't declare raises an error instead of being ignored.
• The standard CFG-cutoff callbacks don't apply; modular pipelines control guidance through the guider component.

Self-review notes

Final self-review round (verdict: READY)

Blocking: none. Checked request validation, cleanup in finally, that conditioning edits reach the next step, and that the chunked loops handle interrupt and callbacks.

Non-blocking, left for review:

1. set_progress_bar_config doesn't reach the plain tqdm bars in helios/denoise.py and wan_animate_2/denoise.py (existing; noted above).
2. _call_callback_on_step_end writes a returned field into any dict in block_state with that key, not only kwargs_type collections. No current loop is affected, but it could be narrowed.
3. An interrupted Helios pyramid chunk is upsampled with nearest so it can still be decoded.
4. The MiniMax Music 3 multi-chunk test monkeypatches latent_hop_length to get two chunks from the tiny config.
5. test_step_callback_and_interrupt is in the shared modular tester mixin, so a future loop without callback support will fail it.

Dead code: none.

Suggestion for maintainers: a note in .ai/references/modular.md that loops should use loop_over_timesteps and declare _callback_tensor_inputs.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case. (Implement missing features on ModularPipeline #12386)
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Who can review?

@yiyixuxu @DN6

[yiyixuxu]

hi @fusheng-ji
thanks a lot for working on this, sorry i didn't update on the issue earlier
we are going a different direction for modular here:
in #14159 (will be included in next release), we add first-class streaming (pipe.stream(...) yields at every step, so we won't need callbacks etc). Really sorry about the duplicated efforts here.

we'd really welcome it if you are able to do a review that PR and let us know if there is any features are missing or anything can be improved - happy to collaborate more there

[fusheng-ji]

Thanks @yiyixuxu, no worries! Streaming looks like a cleaner fit for modular. I'll close this one and leave some notes on #14159 from what I ran into here.