diff --git a/.ai/references/modular.md b/.ai/references/modular.md index 8bfb39ebb999..062f4bf69ed1 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -4,7 +4,7 @@ Shared reference for modular pipeline conventions, patterns, and gotchas. ## Common modular conventions -When adding a new modular pipeline (or reviewing one), skim `src/diffusers/modular_pipelines/qwenimage/`, `src/diffusers/modular_pipelines/flux2/`, `src/diffusers/modular_pipelines/wan/`, and `src/diffusers/modular_pipelines/helios/` first to establish the pattern. Most conventions (file split between `encoders.py` / `before_denoise.py` / `denoise.py` / `decoders.py`, how `expected_components` / `inputs` / `intermediate_outputs` are declared, the denoise-loop wrapping with `LoopSequentialPipelineBlocks`, top-level assembly via `AutoPipelineBlocks` / `SequentialPipelineBlocks` in `modular_blocks_.py`, the `ModularPipeline` subclass shape, the guider-abstracted denoise body, `kwargs_type="denoiser_input_fields"` plumbing) are easiest to internalize by comparison rather than from a fixed list. +When adding a new modular pipeline (or reviewing one), skim `src/diffusers/modular_pipelines/qwenimage/`, `src/diffusers/modular_pipelines/flux2/`, `src/diffusers/modular_pipelines/wan/`, and `src/diffusers/modular_pipelines/helios/` first to establish the pattern. Most conventions (file split between `encoders.py` / `before_denoise.py` / `denoise.py` / `decoders.py`, how `expected_components` / `inputs` / `intermediate_outputs` are declared, the denoise-loop wrapping with `IterativePipelineBlocks` (see `flux2/denoise.py`; `wan_animate_2/denoise.py` for a nested chunk loop), top-level assembly via `AutoPipelineBlocks` / `SequentialPipelineBlocks` in `modular_blocks_.py`, the `ModularPipeline` subclass shape, the guider-abstracted denoise body, `kwargs_type="denoiser_input_fields"` plumbing) are easiest to internalize by comparison rather than from a fixed list. ## Running a modular pipeline @@ -51,8 +51,12 @@ Is this a single operation? Does it run multiple blocks in sequence? YES -> SequentialPipelineBlocks - Does it iterate (e.g. chunk loop)? - YES -> LoopSequentialPipelineBlocks + Does it iterate (denoising loop)? + YES -> IterativePipelineBlocks (steps are ModularLoopPipelineBlocks) + Does each iteration run a loop of its own (autoregressive chunk loop: per chunk, a denoise loop)? + YES -> nest them: the inner IterativePipelineBlocks is a step of the outer one, and its + __call__/stream and its steps also take the outer loop variable (see wan_animate_2/denoise.py) + LoopSequentialPipelineBlocks is the legacy loop type — don't use it for new pipelines Does it choose ONE block based on which input is present? Is the selection 1:1 with trigger inputs? @@ -121,25 +125,96 @@ for i, t in enumerate(timesteps): ## Key pattern: Denoising loop -All models use `LoopSequentialPipelineBlocks` for the denoising loop (iterating over timesteps): +The denoising loop is an `IterativePipelineBlocks` whose steps are `ModularLoopPipelineBlocks` (see `flux2/denoise.py` as example). Building one takes three parts: write the loop steps, write the wrapper, assemble. + +**1. Write the loop steps.** Loop steps are regular blocks. They declare `inputs` and `intermediate_outputs` and use `get_block_state` / `set_block_state` with the full `PipelineState`. + +Their `__call__` also accepts the loop variables they use as keyword arguments. If a step ignores some loop variables, declare `**kwargs`. Otherwise, name every loop variable. The loop validates these signatures when it is constructed and raises for unknown or missing variables. + ```python -class MyModelDenoiseLoopWrapper(LoopSequentialPipelineBlocks): - block_classes = [LoopBeforeDenoiser, LoopDenoiser, LoopAfterDenoiser] +class MyModelLoopAfterDenoiser(ModularLoopPipelineBlocks): + @property + def expected_components(self): + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self): + return [InputParam("latents", required=True), InputParam("noise_pred", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam("latents")] + + @torch.no_grad() + def __call__(self, components, state, i, t): + block_state = self.get_block_state(state) + block_state.latents = components.scheduler.step(block_state.noise_pred, t, block_state.latents, return_dict=False)[0] + self.set_block_state(state, block_state) + return components, state ``` -Autoregressive video models (e.g. Helios) also use it for an outer chunk loop: +**2. Write the wrapper.** The loop logic only: `loop_variables` (passed to every step), `loop_inputs` (what the loop logic itself reads), and the loop in `__call__`; `stream` is the same loop as a generator over `stream_step` and is what makes `pipe.stream(...)` work — implement both: + ```python -class HeliosChunkDenoiseStep(HeliosChunkLoopWrapper): - block_classes = [ - HeliosChunkHistorySliceStep, - HeliosChunkNoiseGenStep, - HeliosChunkSchedulerResetStep, - HeliosChunkDenoiseInner, - HeliosChunkUpdateStep, - ] +class MyModelDenoiseLoopWrapper(IterativePipelineBlocks): + @property + def loop_variables(self): # passed to every loop step + return ["i", "t"] + + @property + def loop_inputs(self): + # what the loop logic itself reads — `get_block_state` returns exactly these + # (the loop steps declare their own `inputs` on top of this) + return [InputParam("timesteps", required=True), InputParam.template("num_inference_steps", 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 + + @torch.no_grad() + def stream(self, components, state): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + return components, state +``` + +**3. Assemble.** Attach the steps in a subclass (or at runtime with `from_blocks_dict`): + +```python +class MyModelDenoiseStep(MyModelDenoiseLoopWrapper): + block_classes = [MyModelLoopDenoiser, MyModelLoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] +``` + +Autoregressive video models nest loops: an outer segment loop (`loop_variables = ["k"]`) whose steps prepare the segment, run the full inner denoising loop, and update the history (see `wan_animate_2/denoise.py` as an example). The inner loop is a step of the outer one, so its `__call__` / `stream` accept the outer variable (`def __call__(self, components, state, k)`) and pass its own `i`, `t` to its own steps. + +The wrapper should contain only the loop logic, i.e. how to iterate through the steps, and its `loop_inputs` are just what that takes (e.g. `timesteps`, `num_segments`). All data flows through the steps, which read and write the pipeline state directly: in the example above, the denoiser step writes `noise_pred` to the pipeline state and `MyModelLoopAfterDenoiser` reads it back and writes the updated `latents`; the wrapper touches none of them. If the loop logic seems to need to do more than just iterate (e.g. collect results, carry something to the next iteration), add a loop step for it instead. In `wan_animate_2`, a small `collect` step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back. What we don't want is the wrapper doing it inline: + +```python +# don't: loop logic collecting results itself +def __call__(self, components, state): + block_state = self.get_block_state(state) + segment_frames = [] + for k in range(block_state.num_segments): + components, state = self.loop_step(components, state, k=k) + segment_frames.append(state.get("out_frames")) # reaching into the state: make this a `collect` loop step + ... ``` -Note: sub-blocks inside `LoopSequentialPipelineBlocks` receive `(components, block_state, i, t)` for denoise loops or `(components, block_state, k)` for chunk loops. +Two small notes: +- A value written late in iteration `k` and read early in iteration `k + 1` becomes a pipeline input, since its first read has no writer before it. if a user-supplied value isn't meaningful, seed it in a block that runs before the loop instead: declare it as an output and set it to `None`, and it stays internal. + + ``` + loop, iteration k: step_1: x = xs[-1] <- first read of `xs`, no writer before it: + step_2: xs.append(x * 10) `xs` becomes a pipeline input unless a previous block produce it + ``` +- Components the loop logic itself uses are added by overriding `expected_components` (e.g. `flux2`'s wrapper reads `scheduler.order` to compute the progress-bar warmup steps). + +Existing pipelines still use `LoopSequentialPipelineBlocks` (steps receive a shared flattened `block_state`, no nesting, no streaming). Leave them alone unless you are porting the pipeline; don't use it for new ones. ## Key pattern: `kwargs_type` inputs (`denoiser_input_fields`) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f05667986f11..1cf72913209c 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -104,6 +104,8 @@ title: ModularPipelineBlocks - local: modular_diffusers/sequential_pipeline_blocks title: SequentialPipelineBlocks + - local: modular_diffusers/iterative_pipeline_blocks + title: IterativePipelineBlocks - local: modular_diffusers/loop_sequential_pipeline_blocks title: LoopSequentialPipelineBlocks - local: modular_diffusers/auto_pipeline_blocks diff --git a/docs/source/en/api/modular_diffusers/pipeline_blocks.md b/docs/source/en/api/modular_diffusers/pipeline_blocks.md index 4808f2cf3bbe..268de620c9ea 100644 --- a/docs/source/en/api/modular_diffusers/pipeline_blocks.md +++ b/docs/source/en/api/modular_diffusers/pipeline_blocks.md @@ -8,6 +8,14 @@ [[autodoc]] diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks +## ModularLoopPipelineBlocks + +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.ModularLoopPipelineBlocks + +## IterativePipelineBlocks + +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.IterativePipelineBlocks + ## LoopSequentialPipelineBlocks [[autodoc]] diffusers.modular_pipelines.modular_pipeline.LoopSequentialPipelineBlocks diff --git a/docs/source/en/api/modular_diffusers/pipeline_states.md b/docs/source/en/api/modular_diffusers/pipeline_states.md index 341d18ecb41c..c3cd454f2642 100644 --- a/docs/source/en/api/modular_diffusers/pipeline_states.md +++ b/docs/source/en/api/modular_diffusers/pipeline_states.md @@ -6,4 +6,8 @@ ## BlockState -[[autodoc]] diffusers.modular_pipelines.modular_pipeline.BlockState \ No newline at end of file +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.BlockState + +## StreamEvent + +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.StreamEvent diff --git a/docs/source/en/modular_diffusers/custom_blocks.md b/docs/source/en/modular_diffusers/custom_blocks.md index fbc05b16d77e..3eac4aed31ce 100644 --- a/docs/source/en/modular_diffusers/custom_blocks.md +++ b/docs/source/en/modular_diffusers/custom_blocks.md @@ -319,7 +319,7 @@ This guide covered creating a single custom block. Learn how to compose multiple - [SequentialPipelineBlocks](./sequential_pipeline_blocks): Chain blocks to execute in sequence - [ConditionalPipelineBlocks](./auto_pipeline_blocks): Create conditional blocks that select different execution paths -- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks): Define an iterative workflows like the denoising loop +- [IterativePipelineBlocks](./iterative_pipeline_blocks): Define an iterative workflow like the denoising loop diff --git a/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md new file mode 100644 index 000000000000..26fe281246a2 --- /dev/null +++ b/docs/source/en/modular_diffusers/iterative_pipeline_blocks.md @@ -0,0 +1,322 @@ + + +# IterativePipelineBlocks + +[`~modular_pipelines.IterativePipelineBlocks`] is a multi-block type that runs its sub-blocks multiple times. It is used to build a denoising loop where the sub-blocks predict the noise and step the scheduler, then the loop runs them once per timestep. You can also nest one [`~modular_pipelines.IterativePipelineBlocks`] under another to build an autoregressive video pipeline that generates chunk after chunk. A loop can stream each iteration to the caller when its `stream` method is implemented. + +Use this block type when a pipeline must repeat a sequence of blocks while carrying state between iterations and, when implemented, streaming progress to the caller. + +Two classes are involved in building a loop. + +| Class | Role | State it works with | +|---|---|---| +| [`~modular_pipelines.ModularLoopPipelineBlocks`] | A *loop step*: a regular block that runs inside a loop, once per iteration | The full [`~modular_pipelines.PipelineState`], through its own `get_block_state` / `set_block_state`; also receives the loop variables as call arguments | +| [`~modular_pipelines.IterativePipelineBlocks`] | The *loop* itself: holds the loop logic and runs its steps once per iteration; loops can nest and stream | Its block state holds only its own `loop_inputs`; all data flows between its steps through the [`~modular_pipelines.PipelineState`] | + +This guide uses a few closely related terms, so let's define them upfront. + +- **Pipeline state** is the shared [`~modular_pipelines.PipelineState`] every block reads and writes. It is the only place where data crosses blocks and survives an iteration. +- **Block state** is one block's declared view of the pipeline state: what its `get_block_state` returns and its `set_block_state` writes back. For a loop wrapper, it holds only the loop's `loop_inputs`. +- **Loop inputs** are the loop wrapper's own input declaration: what the loop logic itself reads to drive the iteration (the `timesteps` it iterates over, for example). +- **Loop variables** are the per-iteration values the loop passes to its steps as call arguments (`i`, `t`). They are never written to the pipeline state. + +> [!TIP] +> [`~modular_pipelines.IterativePipelineBlocks`] replaces [`~modular_pipelines.LoopSequentialPipelineBlocks`]; see the [LoopSequentialPipelineBlocks](#loopsequentialpipelineblocks) guide for the differences. + +## Loop steps + +A loop step is a [`~modular_pipelines.ModularLoopPipelineBlocks`]. It is a regular [`~modular_pipelines.ModularPipelineBlocks`] — it declares `inputs` and `intermediate_outputs`, and reads and writes the [`~modular_pipelines.PipelineState`] through `get_block_state` / `set_block_state` — with one difference. Its `__call__` also receives the loop's *loop variables* as arguments. For example, a denoising loop can pass the step index `i` and the timestep `t`. + +Loop variables are local to the loop and act as its bookkeeping, not pipeline data. The loop hands them to each step as plain call arguments for that one iteration, and they are never written to the [`~modular_pipelines.PipelineState`]. Anything that has to outlive the iteration goes through the state instead, like `noise_pred` and `latents` below. A streaming consumer does see them, each [`~modular_pipelines.StreamEvent`] carries that iteration's values in `event.loop_kwargs`. + +```py +from diffusers.modular_pipelines import ModularLoopPipelineBlocks, InputParam, OutputParam + +class DenoiserStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "predicts the noise for one timestep" + + @property + def inputs(self): + return [InputParam(name="latents", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="noise_pred")] + + def __call__(self, components, state, i, t): + block_state = self.get_block_state(state) + block_state.noise_pred = block_state.latents * 0 + t # stands in for the denoiser + self.set_block_state(state, block_state) + return components, state + + +class SchedulerStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "updates the latents with the noise prediction" + + @property + def inputs(self): + return [InputParam(name="latents", required=True), InputParam(name="noise_pred", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="latents")] + + def __call__(self, components, state, i, t): + block_state = self.get_block_state(state) + block_state.latents = block_state.latents + block_state.noise_pred + self.set_block_state(state, block_state) + return components, state +``` + +Because each step works on the [`~modular_pipelines.PipelineState`], the values it writes (`noise_pred`, the updated `latents`) are visible to the next step in the same iteration and to the next iteration. `latents` is read at the start of every iteration and written at the end of it. + +## Loop wrapper + +The loop itself is a subclass of [`~modular_pipelines.IterativePipelineBlocks`]. It declares: + +- `loop_variables`, the names of the variables passed to each step on every iteration. Steps receive loop variables as keyword arguments. A step can name only the variables it uses if its `__call__` accepts `**kwargs`. Otherwise, it must name every loop variable. The loop validates these signatures when it is constructed and raises for unknown or missing variables. +- `loop_inputs`, the inputs the loop logic itself reads — here the `timesteps` it iterates. They join the inputs aggregated from the steps (see below), and they are what `get_block_state` returns for the loop block. +- `__call__`, the loop logic: read the loop's block state, and call `loop_step` once per iteration with the loop variables. `loop_step` runs every step once. + +The wrapper defines only the loop logic. Which steps run inside it should be attached separately, in a subclass: + +```py +import torch +from diffusers.modular_pipelines import IterativePipelineBlocks + +class DenoiseLoopWrapper(IterativePipelineBlocks): + model_name = "test" + + @property + def description(self): + return "denoises the latents over the timesteps" + + @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 + + +class DenoiseLoop(DenoiseLoopWrapper): + block_classes = [DenoiserStep, SchedulerStep] + block_names = ["denoiser", "scheduler"] +``` + +This separation means the same loop logic can work with different combinations of loop steps. For example, subclass the wrapper again with different `block_classes`. Steps can also be attached to the loop with [`~modular_pipelines.IterativePipelineBlocks.from_blocks_dict`]: + +```py +loop = DenoiseLoopWrapper.from_blocks_dict({"denoiser": DenoiserStep(), "scheduler": SchedulerStep()}) +``` + +You can also change what runs inside a loop you already have: add a step, reorder, swap one out. For example, you can insert a logging step like this: + +```py +loop = DenoiseLoopWrapper.from_blocks_dict(loop.sub_blocks.copy().insert("log", LogStep(), 1)) +``` + +An [`~modular_pipelines.IterativePipelineBlocks`] is a [`~modular_pipelines.SequentialPipelineBlocks`], so it [aggregates](./sequential_pipeline_blocks#aggregated-inputs-and-outputs) its steps' `inputs` and `intermediate_outputs` the same way any assembled block does, and adds `loop_inputs` / `loop_intermediate_outputs` on top. `DenoiseLoop.inputs` is therefore `timesteps` (its own) and `latents` (what the steps need from outside the loop) — but not `noise_pred`, which `DenoiserStep` produces before `SchedulerStep` reads it, so it is satisfied inside the loop. The assembled loop ends up with the same kind of input/output contract as a single block, which is what lets it be dropped into a [`~modular_pipelines.SequentialPipelineBlocks`], or into another loop. Outputs the steps write persist in the state after the loop. Run it like any other block: + +```py +pipeline = DenoiseLoop().init_pipeline() +state = pipeline(latents=torch.tensor(0.0), timesteps=torch.tensor([1.0, 2.0, 3.0])) +state.get("latents") # tensor(6.) — 0 + 1 + 2 + 3 +``` + +The wrapper contains only the loop logic, i.e. how to iterate through its steps, so its `loop_inputs` should be just what that takes (the `timesteps` above). All data flows through the steps, which read and write the pipeline state directly. In the example above, `DenoiserStep` writes `noise_pred` to the state and `SchedulerStep` reads it back and writes the updated `latents`. The wrapper touches none of them. If the loop logic seems to need to do more than iterate, for example, collect results, you should add a loop step for it instead. In `wan_animate_2`, a small collect step appends each segment's decoded frames to `segment_frames`, and the next segment's prep step reads it back to condition on. Under [streaming](#streaming), the partial collection is visible after every iteration. What we don't want is the wrapper doing it inline: + +```py +# don't: loop logic collecting results itself +def __call__(self, components, state): + block_state = self.get_block_state(state) + segment_frames = [] + for k in range(block_state.num_segments): + components, state = self.loop_step(components, state, k=k) + segment_frames.append(state.get("out_frames")) # reaching into the state: make this a collect loop step + ... +``` + +## Nesting loops + +An [`~modular_pipelines.IterativePipelineBlocks`] can be a step of another one. An autoregressive video pipeline generates a chunk of frames at a time, so it is an outer loop over chunks. Each of its iterations prepares the chunk's latents from the frames generated so far, runs a full denoising loop over them, and appends the result to the history. + +The inner denoising loop is a step of the outer loop, so its `__call__` must accept the outer loop's variables (or take `**kwargs` if it ignores them), and it declares `loop_variables` of its own for its own steps. The two sets are independent. `k` arrives as a call argument and the inner loop is free to use it — the wan-animate-2 denoise loop puts the chunk index in its progress bar — but it is not forwarded to the steps, which are passed the inner loop's `i` and `t`. + +This is the structure this section builds: + +``` +VideoLoop outer loop — loop_variables ["k"], iterates over chunks +├─ prepare PrepareChunkStep (leaf step, takes k) +├─ denoise ChunkDenoiseLoop (the INNER loop — a step of the outer one) +│ ├─ denoiser DenoiserStep (takes the inner loop's i, t) +│ └─ scheduler SchedulerStep +└─ update UpdateHistoryStep (leaf step, takes k) +``` + +```py +class ChunkDenoiseLoopWrapper(IterativePipelineBlocks): + model_name = "test" + + @property + def description(self): + return "denoises one chunk over the timesteps" + + @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, k): # `k` comes from the outer video loop + 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 ChunkDenoiseLoop(ChunkDenoiseLoopWrapper): + block_classes = [DenoiserStep, SchedulerStep] + block_names = ["denoiser", "scheduler"] +``` + +`ChunkDenoiseLoop` is the *inner* loop that denoises a single chunk. It becomes a step of the outer loop below, which iterates over the chunks. At each `k`, it prepares the chunk's latents from the history, runs the full inner denoising loop over them, and records the result: + +```py +class PrepareChunkStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "prepares this chunk's latents from the history" + + @property + def inputs(self): + return [InputParam(name="history", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="latents")] + + def __call__(self, components, state, k): + block_state = self.get_block_state(state) + block_state.latents = block_state.history + k + self.set_block_state(state, block_state) + return components, state + + +class UpdateHistoryStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "records the denoised chunk" + + @property + def inputs(self): + return [InputParam(name="latents", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="history")] + + def __call__(self, components, state, k): + block_state = self.get_block_state(state) + block_state.history = block_state.latents + self.set_block_state(state, block_state) + return components, state + + +class VideoLoopWrapper(IterativePipelineBlocks): + model_name = "test" + + @property + def description(self): + return "generates the video chunk by chunk" + + @property + def loop_variables(self): + return ["k"] + + @property + def loop_inputs(self): + return [InputParam(name="num_chunks", required=True)] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + for k in range(block_state.num_chunks): + components, state = self.loop_step(components, state, k=k) + return components, state + + +class VideoLoop(VideoLoopWrapper): + block_classes = [PrepareChunkStep, ChunkDenoiseLoop, UpdateHistoryStep] + block_names = ["prepare", "denoise", "update"] +``` + +Run the outer loop like any other block: + +```py +pipeline = VideoLoop().init_pipeline() +state = pipeline(num_chunks=2, timesteps=torch.tensor([1.0, 2.0]), history=torch.tensor(0.0)) +state.get("history") # tensor(7.) — chunk 0: 0 + 0 + 3 = 3, chunk 1: 3 + 1 + 3 = 7 +``` + +`history` is carried from one iteration to the next. `UpdateHistoryStep` writes it at the end of one, and `PrepareChunkStep` reads it at the start of the next. Because the reader comes before the writer, it is one of the loop's inputs — which is why `pipeline(history=...)` works above — and like any input it has to come from either the user or an earlier block. If seeding it isn't meaningful (a decoder cache, the previous chunk's frames), have the block that runs before the loop declare it as an output and set its initial value. It then drops out of the pipeline's signature. + +## Streaming + +To let [`~ModularPipeline.stream`] hand back the live state after every iteration, also implement `stream`. This is the same loop written as a generator over `stream_step`, which runs one iteration like `loop_step` and additionally yields a [`~modular_pipelines.StreamEvent`] for it (after the events of any nested loop): + +```py +class DenoiseLoop(IterativePipelineBlocks): + ... + + def stream(self, components, state): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + return components, state +``` + +A nested loop's `stream` takes the outer loop's variables exactly like its `__call__` does. Streaming is opt-in per loop.`pipeline.blocks.supports_streaming` tells you whether every loop on the path implements it. See [Streaming](./modular_pipeline#streaming) for the consumer side. + +## LoopSequentialPipelineBlocks + +[`~modular_pipelines.LoopSequentialPipelineBlocks`] is the earlier loop type and is still used by existing pipelines. It differs in three ways: + +1. Steps share one flattened [`~modular_pipelines.BlockState`] that the wrapper extracts before the loop (instead of each step reading the [`~modular_pipelines.PipelineState`] itself). +2. It cannot contain another loop. +3. It cannot stream. + +Use [`~modular_pipelines.IterativePipelineBlocks`] for new pipelines. diff --git a/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md index 74a868922799..3e8c03a8d355 100644 --- a/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md @@ -12,6 +12,9 @@ specific language governing permissions and limitations under the License. # LoopSequentialPipelineBlocks +> [!WARNING] +> [`~modular_pipelines.LoopSequentialPipelineBlocks`] is superseded by [`~modular_pipelines.IterativePipelineBlocks`], which lets loop steps work on the [`~modular_pipelines.PipelineState`] like regular blocks, can be nested, and supports streaming. Use [IterativePipelineBlocks](./iterative_pipeline_blocks) for new pipelines; this page documents the earlier type that existing pipelines still use. + [`~modular_pipelines.LoopSequentialPipelineBlocks`] are a multi-block type that composes other [`~modular_pipelines.ModularPipelineBlocks`] together in a loop. Data flows circularly, using `inputs` and `intermediate_outputs`, and each block is run iteratively. This is typically used to create a denoising loop which is iterative by default. This guide shows you how to create [`~modular_pipelines.LoopSequentialPipelineBlocks`]. diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 07dc30b078ae..87d503c52318 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -392,6 +392,55 @@ output = pipeline( If pipeline stages share components (e.g., the same VAE used for encoding and decoding), you can use [`~ModularPipeline.update_components`] to pass an already-loaded component to another pipeline instead of loading it again. +## Streaming + +[`~ModularPipeline.stream`] runs the same pipeline as a generator. It yields a [`StreamEvent`] after every iteration of every loop block, e.g. each denoising step, each segment of a chunked video, with the live [`PipelineState`] attached, so you can show progress, decode a preview, or stop early. Exactly one event is emitted per loop iteration — the loop steps inside an iteration don't emit their own events, and blocks without loops emit none. When exhausted, the generator returns the final [`PipelineState`], the default return value of `__call__`. `stream()` does not support `output=` selection. + +```py +generator = pipeline.stream(prompt="a cat", num_inference_steps=20) +for event in generator: + print(event.path, event.loop_kwargs) # "denoise.denoise" {"i": 0, "t": tensor(1000.)} + latents = event.state.get("latents") # the live state — clone anything you keep +``` + +`event.path` is the loop block's dotted name from the top of the pipeline, and `event.loop_kwargs` its loop variables for that iteration. When loops are nested (an autoregressive video that denoises one chunk at a time), you receive events from both loops. One after every denoising step of the inner loop (path `"denoise.denoise_inner"`), and one after each completed chunk of the outer loop (path `"denoise"`). Check `event.path` to tell them apart, for example, to react only when a chunk is finished: + +```py +for event in pipeline.stream(...): + if event.path == "denoise": # an outer-loop event: a whole chunk is done + show(event.state.get("out_frames")) +``` + +To stop early, stop iterating (or call `generator.close()`). + +Streaming is opt-in per loop block. An [`IterativePipelineBlocks`] (see the [IterativePipelineBlocks](./iterative_pipeline_blocks) guide) implements its loop in `__call__` as usual and, to support streaming, also implements `stream`. This is the same loop written as a generator over `stream_step`, which runs one iteration like `loop_step` and additionally yields the event for it: + +```py +class DenoiseLoop(IterativePipelineBlocks): + @property + def loop_variables(self): + return ["i", "t"] + + @property + def loop_inputs(self): # what the loop logic itself reads; `get_block_state` returns exactly these + return [InputParam("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 + + def stream(self, components, state): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + return components, state +``` + +`pipeline.stream(...)` raises `NotImplementedError` if a loop on its path doesn't implement `stream`. Check `pipeline.blocks.supports_streaming` to find out ahead of time. It is `True` unless the blocks contain a loop that can't yield per iteration (an `IterativePipelineBlocks` that doesn't implement `stream`, or a legacy `LoopSequentialPipelineBlocks`). + ## Modular repository A repository is required if the pipeline blocks use *pretrained components*. The repository supplies loading specifications and metadata. diff --git a/docs/source/en/modular_diffusers/overview.md b/docs/source/en/modular_diffusers/overview.md index f80fff3061de..2530ace49a88 100644 --- a/docs/source/en/modular_diffusers/overview.md +++ b/docs/source/en/modular_diffusers/overview.md @@ -28,7 +28,8 @@ The Modular Diffusers docs are organized as shown below. - [States](./modular_diffusers_states) explains how data is shared and communicated between blocks and [`ModularPipeline`]. - [ModularPipelineBlocks](./pipeline_block) is the most basic unit of a [`ModularPipeline`] and this guide shows you how to create one. - [SequentialPipelineBlocks](./sequential_pipeline_blocks) is a type of block that chains multiple blocks so they run one after another, passing data along the chain. This guide shows you how to create [`~modular_pipelines.SequentialPipelineBlocks`] and how they connect and work together. -- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks) is a type of block that runs a series of blocks in a loop. This guide shows you how to create [`~modular_pipelines.LoopSequentialPipelineBlocks`]. +- [IterativePipelineBlocks](./iterative_pipeline_blocks) is a type of block that runs a series of blocks in a loop — a denoising loop, or an autoregressive chunk loop with a denoising loop nested inside. This guide shows you how to create [`~modular_pipelines.IterativePipelineBlocks`]. +- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks) is the earlier loop type, still used by existing pipelines. Prefer [`~modular_pipelines.IterativePipelineBlocks`] for new pipelines. - [AutoPipelineBlocks](./auto_pipeline_blocks) is a type of block that automatically chooses which blocks to run based on the input. This guide shows you how to create [`~modular_pipelines.AutoPipelineBlocks`]. - [Building Custom Blocks](./custom_blocks) shows you how to create your own custom blocks and share them on the Hub. diff --git a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md index 1bd67e17b8bf..d9c061e85950 100644 --- a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md @@ -113,6 +113,12 @@ class ImageProcessingStep(SequentialPipelineBlocks): When you create a [`~modular_pipelines.SequentialPipelineBlocks`], properties like `inputs`, `intermediate_outputs`, and `expected_components` are automatically aggregated from the sub-blocks, so there is no need to define them again. +### Aggregated inputs and outputs + +Aggregation follows the order the sub-blocks run in. `inputs` walks the sub-blocks in order and omits an input when an earlier sub-block declares an output with the same name, because the assembled block produces that value internally. `intermediate_outputs` is the union of what the sub-blocks write. + +This is what makes blocks composable: an assembled block has the same kind of input/output contract as a single [`~modular_pipelines.ModularPipelineBlocks`], so it can in turn be a sub-block of another one, at any depth. It is also why an input disappears from a pipeline's signature once some earlier block produces it — see [nesting loops](./iterative_pipeline_blocks#nesting-loops) for a case where you use that deliberately. + There are a few properties you should set: - `description`: We recommend adding a description for the assembled block to explain what the combined step does. diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 1fc34e6bdbf6..a39dfefb60f8 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -377,11 +377,14 @@ "ConditionalPipelineBlocks", "ConfigSpec", "InputParam", + "IterativePipelineBlocks", "LoopSequentialPipelineBlocks", + "ModularLoopPipelineBlocks", "ModularPipeline", "ModularPipelineBlocks", "OutputParam", "SequentialPipelineBlocks", + "StreamEvent", ] ) _import_structure["optimization"] = [ @@ -1249,11 +1252,14 @@ ConditionalPipelineBlocks, ConfigSpec, InputParam, + IterativePipelineBlocks, LoopSequentialPipelineBlocks, + ModularLoopPipelineBlocks, ModularPipeline, ModularPipelineBlocks, OutputParam, SequentialPipelineBlocks, + StreamEvent, ) from .optimization import ( get_constant_schedule, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 81b93f88f515..572bab22745b 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -28,9 +28,12 @@ "AutoPipelineBlocks", "SequentialPipelineBlocks", "ConditionalPipelineBlocks", + "IterativePipelineBlocks", + "ModularLoopPipelineBlocks", "LoopSequentialPipelineBlocks", "PipelineState", "BlockState", + "StreamEvent", ] _import_structure["modular_pipeline_utils"] = [ "ComponentSpec", @@ -202,11 +205,14 @@ AutoPipelineBlocks, BlockState, ConditionalPipelineBlocks, + IterativePipelineBlocks, LoopSequentialPipelineBlocks, + ModularLoopPipelineBlocks, ModularPipeline, ModularPipelineBlocks, PipelineState, SequentialPipelineBlocks, + StreamEvent, ) from .modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, InsertableDict, OutputParam from .qwenimage import ( diff --git a/src/diffusers/modular_pipelines/flux2/denoise.py b/src/diffusers/modular_pipelines/flux2/denoise.py index 675f14b03c63..b36b06fcff19 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,16 +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 [ + 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, @@ -392,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 ( @@ -405,13 +466,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 [ @@ -432,24 +486,31 @@ 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) + 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 + + @torch.no_grad() + def stream(self, components: Flux2ModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + if XLA_AVAILABLE: + xm.mark_step() return components, state @@ -461,7 +522,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" @@ -477,7 +538,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" @@ -493,7 +554,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/flux2/modular_blocks_flux2.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py index 2bbb7975a983..80800c1d5936 100644 --- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py @@ -316,9 +316,9 @@ class Flux2AutoBlocks(SequentialPipelineBlocks): TODO: Add description. latents (`Tensor | NoneType`): TODO: Add description. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 50): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py index 689cf808c4ba..8756adf6c3ae 100644 --- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py @@ -284,9 +284,9 @@ class Flux2KleinAutoCoreDenoiseStep(AutoPipelineBlocks): TODO: Add description. generator (`None`, *optional*): TODO: Add description. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 50): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. @@ -357,9 +357,9 @@ class Flux2KleinAutoBlocks(SequentialPipelineBlocks): TODO: Add description. latents (`Tensor | NoneType`): TODO: Add description. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 50): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py index f3108bdadeac..7c43be332047 100644 --- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py @@ -299,7 +299,7 @@ class Flux2KleinBaseAutoCoreDenoiseStep(AutoPipelineBlocks): TODO: Add description. num_inference_steps (`None`): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. @@ -373,7 +373,7 @@ class Flux2KleinBaseAutoBlocks(SequentialPipelineBlocks): TODO: Add description. num_inference_steps (`None`): TODO: Add description. - timesteps (`None`): + timesteps (`None`, *optional*): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. diff --git a/src/diffusers/modular_pipelines/ltx2/denoise.py b/src/diffusers/modular_pipelines/ltx2/denoise.py index b1c4657d4d04..43ffead2e39d 100644 --- a/src/diffusers/modular_pipelines/ltx2/denoise.py +++ b/src/diffusers/modular_pipelines/ltx2/denoise.py @@ -23,12 +23,11 @@ from ...models import LTX2VideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ..modular_pipeline import ( - BlockState, - LoopSequentialPipelineBlocks, - ModularPipelineBlocks, + IterativePipelineBlocks, + ModularLoopPipelineBlocks, PipelineState, ) -from ..modular_pipeline_utils import ComponentSpec, InputParam +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam # Velocity-space helpers, mirrored from `diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline` and redefined here @@ -71,7 +70,7 @@ def _unpack_latents( return latents -class LTX2LoopBeforeDenoiser(ModularPipelineBlocks): +class LTX2LoopBeforeDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -93,17 +92,32 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latent_model_input", type_hint=torch.Tensor, description="Video latents cast to model dtype"), + OutputParam( + "audio_latent_model_input", type_hint=torch.Tensor, description="Audio latents cast to model dtype" + ), + OutputParam("video_timestep", type_hint=torch.Tensor, description="This step's video timestep"), + OutputParam("audio_timestep", type_hint=torch.Tensor, description="This step's audio timestep"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + block_state.latent_model_input = block_state.latents.to(block_state.dtype) block_state.audio_latent_model_input = block_state.audio_latents.to(block_state.dtype) timestep = t.expand(block_state.latents.shape[0]) block_state.video_timestep = timestep block_state.audio_timestep = timestep - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2Image2VideoLoopBeforeDenoiser(ModularPipelineBlocks): +class LTX2Image2VideoLoopBeforeDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -129,17 +143,32 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latent_model_input", type_hint=torch.Tensor, description="Video latents cast to model dtype"), + OutputParam( + "audio_latent_model_input", type_hint=torch.Tensor, description="Audio latents cast to model dtype" + ), + OutputParam("video_timestep", type_hint=torch.Tensor, description="This step's masked video timestep"), + OutputParam("audio_timestep", type_hint=torch.Tensor, description="This step's audio timestep"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + block_state.latent_model_input = block_state.latents.to(block_state.dtype) block_state.audio_latent_model_input = block_state.audio_latents.to(block_state.dtype) timestep = t.expand(block_state.latents.shape[0]) block_state.video_timestep = timestep.unsqueeze(-1) * (1 - block_state.conditioning_mask) block_state.audio_timestep = timestep - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2ConditionLoopBeforeDenoiser(ModularPipelineBlocks): +class LTX2ConditionLoopBeforeDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -171,14 +200,29 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latent_model_input", type_hint=torch.Tensor, description="Video latents cast to model dtype"), + OutputParam( + "audio_latent_model_input", type_hint=torch.Tensor, description="Audio latents cast to model dtype" + ), + OutputParam("video_timestep", type_hint=torch.Tensor, description="This step's masked video timestep"), + OutputParam("audio_timestep", type_hint=torch.Tensor, description="This step's audio timestep"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + block_state.latent_model_input = block_state.latents.to(block_state.dtype) block_state.audio_latent_model_input = block_state.audio_latents.to(block_state.dtype) timestep = t.expand(block_state.latents.shape[0]) block_state.video_timestep = timestep.unsqueeze(-1) * (1 - block_state.conditioning_mask.squeeze(-1)) block_state.audio_timestep = timestep - return components, block_state + + self.set_block_state(state, block_state) + return components, state # Default per-pass conditioning map for `LTX2LoopDenoiser`: transformer argument -> block-state attribute names @@ -212,7 +256,7 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) } -class LTX2LoopDenoiser(ModularPipelineBlocks): +class LTX2LoopDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" def __init__(self, guider_input_fields: dict[str, Any] = _DEFAULT_GUIDER_INPUT_FIELDS): @@ -293,6 +337,30 @@ def inputs(self) -> list[InputParam]: required=True, description="Packed noisy audio latents to denoise.", ), + InputParam( + "latent_model_input", + type_hint=torch.Tensor, + required=True, + description="Video latents cast to model dtype, from the before-denoiser step.", + ), + InputParam( + "audio_latent_model_input", + type_hint=torch.Tensor, + required=True, + description="Audio latents cast to model dtype, from the before-denoiser step.", + ), + InputParam( + "video_timestep", + type_hint=torch.Tensor, + required=True, + description="This step's video timestep, from the before-denoiser step.", + ), + InputParam( + "audio_timestep", + type_hint=torch.Tensor, + required=True, + description="This step's audio timestep, from the before-denoiser step.", + ), InputParam("audio_scheduler", required=True), # `audio_num_frames`, `video_coords`, `audio_coords` arrive tagged `denoiser_input_fields` upstream and # are collected from the tagged dict (filtered against the transformer signature) in `__call__`. @@ -336,8 +404,21 @@ def inputs(self) -> list[InputParam]: ) return inputs + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "noise_pred_video", type_hint=torch.Tensor, description="Guided x0 prediction for the video latents" + ), + OutputParam( + "noise_pred_audio", type_hint=torch.Tensor, description="Guided x0 prediction for the audio latents" + ), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 latent_height = block_state.height // components.vae_spatial_compression_ratio latent_width = block_state.width // components.vae_spatial_compression_ratio @@ -433,10 +514,12 @@ def _combine(guider, field): block_state.noise_pred_video = _combine(components.guider, "video_pred") block_state.noise_pred_audio = _combine(components.audio_guider, "audio_pred") - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2LoopAfterDenoiser(ModularPipelineBlocks): +class LTX2LoopAfterDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -458,10 +541,31 @@ def inputs(self) -> list[InputParam]: description="Packed noisy audio latents to denoise.", ), InputParam("audio_scheduler", required=True), + InputParam( + "noise_pred_video", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the video latents, from the denoiser step.", + ), + InputParam( + "noise_pred_audio", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the audio latents, from the denoiser step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="The denoised video latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="The denoised audio latents"), ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + noise_pred_video = convert_x0_to_velocity( block_state.latents, block_state.noise_pred_video, i, components.scheduler ) @@ -472,10 +576,12 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) block_state.audio_latents = block_state.audio_scheduler.step( noise_pred_audio, t, block_state.audio_latents, return_dict=False )[0] - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2Image2VideoLoopAfterDenoiser(ModularPipelineBlocks): +class LTX2Image2VideoLoopAfterDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -500,6 +606,18 @@ def inputs(self) -> list[InputParam]: description="Packed noisy audio latents to denoise.", ), InputParam("audio_scheduler", required=True), + InputParam( + "noise_pred_video", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the video latents, from the denoiser step.", + ), + InputParam( + "noise_pred_audio", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the audio latents, from the denoiser step.", + ), InputParam.template("height", default=512), InputParam.template("width", default=704), InputParam( @@ -513,8 +631,17 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="The denoised video latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="The denoised audio latents"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + spatial_patch = components.transformer_spatial_patch_size temporal_patch = components.transformer_temporal_patch_size latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 @@ -542,10 +669,12 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) block_state.audio_latents = block_state.audio_scheduler.step( noise_pred_audio, t, block_state.audio_latents, return_dict=False )[0] - return components, block_state + + self.set_block_state(state, block_state) + return components, state -class LTX2ConditionLoopAfterDenoiser(ModularPipelineBlocks): +class LTX2ConditionLoopAfterDenoiser(ModularLoopPipelineBlocks): model_name = "ltx2" @property @@ -572,6 +701,18 @@ def inputs(self) -> list[InputParam]: description="Packed noisy audio latents to denoise.", ), InputParam("audio_scheduler", required=True), + InputParam( + "noise_pred_video", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the video latents, from the denoiser step.", + ), + InputParam( + "noise_pred_audio", + type_hint=torch.Tensor, + required=True, + description="Guided x0 prediction for the audio latents, from the denoiser step.", + ), InputParam( "conditioning_mask", type_hint=torch.Tensor, @@ -586,8 +727,17 @@ def inputs(self) -> list[InputParam]: ), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="The denoised video latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="The denoised audio latents"), + ] + @torch.no_grad() - def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + # Conditioning strengths run from 0 (always use the denoised sample) to 1 (always use the condition), with # intermediate values specifying how strongly to follow the condition. Applied in x0 space, not velocity # space (which is what the transformer outputs). @@ -604,12 +754,18 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) block_state.audio_latents = block_state.audio_scheduler.step( noise_pred_audio, t, block_state.audio_latents, return_dict=False )[0] - return components, block_state + self.set_block_state(state, block_state) + return components, state -class LTX2DenoiseLoopWrapper(LoopSequentialPipelineBlocks): + +class LTX2DenoiseLoopWrapper(IterativePipelineBlocks): model_name = "ltx2" + @property + def loop_variables(self) -> list[str]: + return ["i", "t"] + @property def description(self) -> str: return ( @@ -618,11 +774,13 @@ def description(self) -> str: ) @property - def loop_expected_components(self) -> list[ComponentSpec]: - return [ - ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), - ComponentSpec("transformer", LTX2VideoTransformer3DModel), - ] + 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 loop_inputs(self) -> list[InputParam]: @@ -635,19 +793,25 @@ def loop_inputs(self) -> list[InputParam]: def __call__(self, components, 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() - self.set_block_state(state, block_state) + return components, state + + @torch.no_grad() + def stream(self, components, state: PipelineState): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) return components, state diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py index 86428328a5a6..1cba0784856e 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py @@ -975,7 +975,7 @@ class LTX2AutoCoreDenoiseStep(ConditionalPipelineBlocks): Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. num_inference_steps (`int`): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. audio_latents (`Tensor`): Optional pre-encoded audio latents; random noise is used when not provided. @@ -1822,7 +1822,7 @@ class LTX2AutoBlocks(SequentialPipelineBlocks): Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. num_inference_steps (`int`): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. audio_latents (`Tensor`): Optional pre-encoded audio latents; random noise is used when not provided. diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py index 7c77aba94a74..dd4d7c003d5d 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py @@ -318,7 +318,7 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. num_inference_steps (`int`): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. audio_latents (`Tensor`): Optional pre-encoded audio latents; random noise is used when not provided. diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index f2576b99328d..30ad5f0704c5 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -323,6 +323,26 @@ def format_value(v): return f"BlockState(\n{attributes}\n)" +@dataclass +class StreamEvent: + """ + One iteration of a loop block, as yielded by `stream()`. + + Attributes: + path: + Dotted name of the loop block from the top of the pipeline, e.g. `"denoise"` or, for a loop nested inside + another, `"denoise.denoise_inner"`. + block: The [`IterativePipelineBlocks`] that just finished the iteration. + loop_kwargs: The loop variables of that iteration, e.g. `{"i": 3, "t": tensor(...)}`. + state: The live [`PipelineState`] after the iteration — not a copy. Clone anything you keep. + """ + + path: str + block: "IterativePipelineBlocks" + loop_kwargs: dict[str, Any] + state: PipelineState + + class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): """ Base class for all Pipeline Blocks: ConditionalPipelineBlocks, AutoPipelineBlocks, SequentialPipelineBlocks, @@ -583,6 +603,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] @@ -610,6 +649,49 @@ 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__}") + + def stream(self, components, state: PipelineState): + """ + Run the block as a generator that yields a [`StreamEvent`] after every iteration of every loop block it + contains, and returns `(components, state)` when done. A block with no loops runs to completion and yields + nothing; composite blocks re-yield their sub-blocks' events with the sub-block name prepended to `event.path`. + """ + yield from () + return self(components, state) + + @property + def supports_streaming(self) -> bool: + """ + Whether the block can be streamed: `True` unless it contains a loop that can't yield per iteration — an + [`IterativePipelineBlocks`] that doesn't implement `stream`, or a legacy [`LoopSequentialPipelineBlocks`]. + """ + for block in self.sub_blocks.values(): + if not block.supports_streaming: + return False + # a leaf block (no sub_blocks) always streams: it runs to completion and yields nothing + return True + + +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 receives the enclosing loop's variables as keyword call arguments — its signature names the loop + variables it uses and declares `**kwargs` for any it ignores (e.g. `def __call__(self, components, state, t, + **kwargs)`; naming all of them without `**kwargs` works too). The loop validates this at construction: a named + parameter that is not a loop variable, or a missing loop variable without a `**kwargs` catch-all, raises. + + > [!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): """ @@ -801,12 +883,40 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState: logger.error(error_msg) raise + def stream(self, pipeline, state: PipelineState): + # Same branch selection as `__call__`. The branch does not add a level to event paths: "core_denoise.denoise", + # not "core_denoise.t2v.denoise"/"core_denoise.i2v.denoise" — so `event.path` filters keep working whichever branch runs. + trigger_kwargs = {name: state.get(name) for name in self.block_trigger_inputs if name is not None} + block_name = self.select_block(**trigger_kwargs) + + if block_name is None: + block_name = self.default_block_name + + if block_name is None: + logger.info(f"skipping conditional block: {self.__class__.__name__}") + return pipeline, state + + block = self.sub_blocks[block_name] + + try: + logger.info(f"Running block: {block.__class__.__name__}") + return (yield from block.stream(pipeline, state)) + except Exception as e: + error_msg = ( + f"\nError in block: {block.__class__.__name__}\n" + f"Error details: {str(e)}\n" + f"Traceback:\n{traceback.format_exc()}" + ) + logger.error(error_msg) + raise + 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 @@ -829,7 +939,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 @@ -1162,6 +1272,28 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState: raise return pipeline, state + def stream(self, pipeline, state: PipelineState): + for block_name, block in self.sub_blocks.items(): + # re-yield the sub-block's events with its name prepended to the path; its return value is the new state + generator = block.stream(pipeline, state) + while True: + try: + event = next(generator) + except StopIteration as e: + pipeline, state = e.value + break + 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 + event.path = f"{block_name}.{event.path}" if event.path else block_name + yield event + return pipeline, state + # used for `__repr__` def _get_trigger_inputs(self): """ @@ -1215,13 +1347,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"): @@ -1331,6 +1463,241 @@ 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 keyword call arguments: a sub-block's `__call__` names the loop + variables it uses after `(components, state)` and declares `**kwargs` for any it ignores (naming all of them and + omitting `**kwargs` is fine too). This is validated at construction: a named parameter that is not a loop variable, + or a missing loop variable without a `**kwargs` catch-all, raises. 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. The loop logic's own + inputs (e.g. `timesteps`) and outputs are declared in `loop_inputs` / `loop_intermediate_outputs`: they join the + sub-blocks' in the aggregated `inputs` / `intermediate_outputs`, and they are what `get_block_state` / + `set_block_state` read and write for the loop block itself — sub-block values live in the pipeline state, not in + the loop's block state. A component used by the loop logic itself (e.g. the scheduler) is added by overriding + `expected_components`. + + Streaming is opt-in: to let `pipe.stream(...)` hand back the live [`PipelineState`] after every iteration, also + implement `stream` — the same loop, written as a generator over `stream_step` (which runs one iteration like + `loop_step` and additionally yields a [`StreamEvent`] for it, after any events of a nested loop): + + ```python + def stream(self, components, state): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + return components, state + ``` + + A nested loop's `stream` takes the outer loop's variables exactly like its `__call__` does. + + > [!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 [] + + @property + def loop_inputs(self) -> list[InputParam]: + """Inputs read by the loop logic in `__call__` itself (e.g. `timesteps`), beyond what the sub-blocks declare.""" + 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 inputs(self) -> list[InputParam]: + inputs = super().inputs + names = {param.name for param in inputs} + return [param for param in self.loop_inputs if param.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] + + def get_block_state(self, state: PipelineState) -> BlockState: + """The loop logic's own inputs (`loop_inputs`); sub-block values are read from the pipeline state.""" + data = {} + for input_param in self.loop_inputs: + value = state.get(input_param.name) + if value is None: + value = input_param.default + if input_param.required and value is None: + raise ValueError(f"Required input '{input_param.name}' is missing") + data[input_param.name] = value + return BlockState(**data) + + def set_block_state(self, state: PipelineState, block_state: BlockState): + """Write the loop logic's own outputs (`loop_intermediate_outputs`) and modified inputs back to the state.""" + for output_param in self.loop_intermediate_outputs: + if not hasattr(block_state, output_param.name): + raise ValueError(f"Intermediate output '{output_param.name}' is missing in block state") + state.set(output_param.name, getattr(block_state, output_param.name), output_param.kwargs_type) + for input_param in self.loop_inputs: + value = getattr(block_state, input_param.name) + if state.get(input_param.name) is not value: + state.set(input_param.name, value, input_param.kwargs_type) + + 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) + # `super().from_blocks_dict` runs `__init__` first and fills `sub_blocks` from `blocks_dict` after, + # so the validation in `__init__` never saw these blocks — validate them now + instance._validate_sub_blocks() + return instance + + def _validate_sub_blocks(self): + """Sub-blocks must be loop steps (`ModularLoopPipelineBlocks`) or nested loops (`IterativePipelineBlocks`) + and accept the loop variables after `(components, state)` — either all of them by name, or the ones the step + uses plus a `**kwargs` catch-all for the rest.""" + 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 = inspect.signature(block.__call__).parameters + has_var_keyword = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()) + # `block.__call__` is bound, so `self` is already excluded; `[2:]` skips `components` and `state`, + # leaving only the loop-variable parameters + named = {name for name in list(params)[2:] if params[name].kind is not inspect.Parameter.VAR_KEYWORD} + if not named <= expected: + raise ValueError( + f"Loop sub-block '{block_name}' ({block.__class__.__name__}) of {self.__class__.__name__} " + f"accepts {sorted(named - expected)}, which are not loop variables of this loop " + f"({sorted(expected)})." + ) + if named != expected and not has_var_keyword: + 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)` — either " + f"all of them by name, or the ones it uses plus `**kwargs`; its `__call__` accepts " + f"{sorted(named)}." + ) + + 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") + + def stream_step(self, components, state: PipelineState, **loop_kwargs): + """ + The streaming counterpart of `loop_step`: a generator that runs all sub-blocks once, re-yields the events of + any nested loop, then yields one [`StreamEvent`] for this iteration and returns `(components, state)`. Call it + with `yield from` inside `stream`. + """ + for block_name, block in self.sub_blocks.items(): + try: + if isinstance(block, IterativePipelineBlocks): + # a nested loop streams too: re-yield its events with its name prepended to the path + generator = block.stream(components, state, **loop_kwargs) + while True: + try: + event = next(generator) + except StopIteration as e: + components, state = e.value + break + event.path = f"{block_name}.{event.path}" if event.path else block_name + yield event + 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 + yield StreamEvent(path="", block=self, loop_kwargs=dict(loop_kwargs), state=state) + return components, state + + def stream(self, components, state: PipelineState, **kwargs): + # Optional. Subclasses that support streaming implement their loop logic here a second time, as a generator + # running `yield from self.stream_step(...)` once per iteration, with the same signature as their `__call__` + # (`**kwargs` stands for the outer loop's variables when this loop is nested, e.g. `stream(self, components, + # state, k)`). + raise NotImplementedError( + f"{self.__class__.__name__} does not support streaming: implement `stream` (the loop written as a " + "generator over `stream_step`) to use it with `pipe.stream(...)`." + ) + + @property + def supports_streaming(self) -> bool: + return type(self).stream is not IterativePipelineBlocks.stream and super().supports_streaming + + class LoopSequentialPipelineBlocks(ModularPipelineBlocks): """ A Pipeline blocks that combines multiple pipeline block classes into a For Loop. When called, it will call each @@ -1348,6 +1715,10 @@ class LoopSequentialPipelineBlocks(ModularPipelineBlocks): block_classes = [] block_names = [] + # this legacy loop runs all iterations in one call and cannot yield per iteration; it will be deprecated in + # favor of `IterativePipelineBlocks` once all pipelines are ported to it + supports_streaming = False + @property def description(self) -> str: """Description of the block. Must be implemented by subclasses.""" @@ -1603,25 +1974,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) @@ -2958,3 +3310,66 @@ def __call__(self, state: PipelineState = None, output: str | list[str] = None, return state.get(output) else: raise ValueError(f"Output '{output}' is not a valid output type") + + def stream(self, state: PipelineState = None, **kwargs): + """ + Run the pipeline as a generator that yields a [`StreamEvent`] after every iteration of every loop block (e.g. + each denoising step, each segment of a chunked video), with the live [`PipelineState`] attached. The + generator's return value is the final state, the same one `__call__` returns. + + Args: + state (`PipelineState`, optional): + Same as in `__call__`. + **kwargs: + Same as in `__call__`. + + Examples: + ```python + for event in pipeline.stream(prompt="A beautiful sunset", num_inference_steps=20): + print(event.path, event.loop_kwargs) # "denoise" {"i": 0, "t": tensor(...)} + latents = event.state.get("latents") # live state — clone anything you keep + + # stop early: just stop iterating (or call `.close()` on the generator) + ``` + """ + if state is None: + state = PipelineState() + else: + state = deepcopy(state) + + # Make a copy of the input kwargs + passed_kwargs = kwargs.copy() + + # Add inputs to state, using defaults if not provided in the kwargs or the state + # if same input already in the state, will override it if provided in the kwargs + for expected_input_param in self._blocks.inputs: + name = expected_input_param.name + default = expected_input_param.default + kwargs_type = expected_input_param.kwargs_type + if name in passed_kwargs: + state.set(name, passed_kwargs.pop(name), kwargs_type) + elif kwargs_type is not None and kwargs_type in passed_kwargs: + kwargs_dict = passed_kwargs.pop(kwargs_type) + for k, v in kwargs_dict.items(): + state.set(k, v, kwargs_type) + elif name is not None and name not in state.values: + state.set(name, default, kwargs_type) + + # Warn about unexpected inputs + if len(passed_kwargs) > 0: + warnings.warn(f"Unexpected input '{passed_kwargs.keys()}' provided. This input will be ignored.") + + # `torch.no_grad()` is held only while blocks run, not while the caller holds an event. + generator = self._blocks.stream(self, state) + while True: + with torch.no_grad(): + try: + event = next(generator) + except StopIteration as e: + _, state = e.value + return state + except Exception: + error_msg = f"Error in block: ({self._blocks.__class__.__name__}):\n" + logger.error(error_msg) + raise + yield event diff --git a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py index 0ad038e8ccaa..7df13e151360 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/before_denoise.py @@ -79,6 +79,13 @@ def intermediate_outputs(self) -> list[OutputParam]: type_hint=int, description="Packed sequence length of the reference tokens", ), + OutputParam( + "segment_frames", + type_hint=list, + description="The decoded frames collected by the segment loop, carried across iterations (the " + "next segment conditions on the last entry's tail); starts as `None` (nothing has been " + "collected yet, and the first segment has no previous segment to condition on)", + ), ] @torch.no_grad() @@ -107,6 +114,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: latent_noise_frames = latent_segment_frames + 1 block_state.max_seq_len = int(math.ceil(np.prod([latent_noise_frames, latent_height // 2, latent_width // 2]))) block_state.max_seq_len_ref = int(math.ceil(np.prod(ref_shape) // 4)) + block_state.segment_frames = None self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py index d96b8f814239..541592449265 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/denoise.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/denoise.py @@ -25,7 +25,7 @@ from ...schedulers.scheduling_utils import SchedulerMixin from ...utils import logging from ...utils.torch_utils import randn_tensor -from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState +from ..modular_pipeline import IterativePipelineBlocks, ModularLoopPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .encoders import encode_vae, get_i2v_mask @@ -47,11 +47,11 @@ def decode_vae(vae: AutoencoderKLWan, latents: torch.Tensor) -> torch.Tensor: # ======================================== -# Segment Loop Leaf Blocks +# Segment Loop Steps # ======================================== -class WanAnimate2SegmentVaeEncoderStep(ModularPipelineBlocks): +class WanAnimate2SegmentVaeEncoderStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -61,7 +61,8 @@ def description(self) -> str: "the i2v conditioning mask on top. The Wan VAE is causal in time, so encoding the whole video once " "and slicing the latents would not be equivalent — each segment restarts the temporal convolution on " "its own slice. A streaming mode would replace this block with one fed segments incrementally. This " - "block should be used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + "block should be used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` object " + "(e.g. `WanAnimate2SegmentLoopWrapper`); it reads the current segment index `k` from the loop scope." ) @property @@ -115,7 +116,8 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) device = components._execution_device latent_height, latent_width = block_state.reference_image_latents.shape[-2:] @@ -134,10 +136,11 @@ def __call__(self, components, block_state: BlockState, k: int): ).to(block_state.driving_video_latents.dtype) block_state.driving_video_condition = torch.cat([condition_mask, block_state.driving_video_latents[0]], dim=0) - return components, block_state + self.set_block_state(state, block_state) + return components, state -class WanAnimate2SegmentPrevFramesStep(ModularPipelineBlocks): +class WanAnimate2SegmentPrevFramesStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -146,8 +149,9 @@ def description(self) -> str: "Step within the segment loop that builds the generation-side conditioning tensor `reference_latents`: the previous " "segment's tail frames (zeros for the first segment) are VAE-encoded, masked, and stacked under the " "reference half `reference_image_latents`. This is how motion continuity crosses segment boundaries — in pixel space, " - "not latent space. This block should be used to compose the `sub_blocks` attribute of " - "`WanAnimate2SegmentLoopWrapper`." + "not latent space. This block should be used to compose the `sub_blocks` attribute of an " + "`IterativePipelineBlocks` object (e.g. `WanAnimate2SegmentLoopWrapper`); it reads the current segment " + "index `k` from the loop scope." ) @property @@ -165,6 +169,11 @@ def inputs(self) -> list[InputParam]: type_hint=torch.Tensor, description="i2v mask + reference image latents `[20, 1, latent_height, latent_width]`, from the image VAE encoder step", ), + InputParam( + "segment_frames", + type_hint=list, + description="The decoded segments collected so far (on CPU), written by the collect step; the previous segment's tail is read from the last entry. `None` for the first segment", + ), InputParam( "segment_frame_length", type_hint=int, @@ -190,9 +199,8 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): - # `block_state.out_frames` is seeded by the loop wrapper and written by the decode step of the - # previous iteration. + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) device = components._execution_device latent_height, latent_width = block_state.reference_image_latents.shape[-2:] @@ -202,7 +210,8 @@ def __call__(self, components, block_state: BlockState, k: int): num_frames = block_state.segment_frame_length + 1 mask_len = block_state.prev_segment_conditioning_frames if k > 0 else 0 if mask_len > 0: - prev_frames = block_state.out_frames[0, :, -mask_len:].clone().detach() + # slice the tail on CPU first so only the few conditioning frames are transferred + prev_frames = block_state.segment_frames[-1][0, :, -mask_len:].to(device) prev_frames = F.interpolate(prev_frames.permute(1, 0, 2, 3), size=(height, width), mode="bicubic").permute( 1, 0, 2, 3 ) @@ -226,10 +235,11 @@ def __call__(self, components, block_state: BlockState, k: int): [block_state.reference_image_latents, prev_segment_cond_latents], dim=1 ) - return components, block_state + self.set_block_state(state, block_state) + return components, state -class WanAnimate2SegmentPrepareStep(ModularPipelineBlocks): +class WanAnimate2SegmentPrepareStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -237,7 +247,7 @@ def description(self) -> str: return ( "Step within the segment loop that draws this segment's initial noise and allocates a fresh KV cache " "for the reference-extraction pass. This block should be used to compose the `sub_blocks` attribute " - "of `WanAnimate2SegmentLoopWrapper`." + "of an `IterativePipelineBlocks` object (e.g. `WanAnimate2SegmentLoopWrapper`)." ) @property @@ -270,7 +280,8 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, **kwargs): + block_state = self.get_block_state(state) device = components._execution_device block_state.latents = randn_tensor( @@ -286,10 +297,11 @@ def __call__(self, components, block_state: BlockState, k: int): ) block_state.kv_cache = WanAnimate2KVCache(components.transformer.config.num_layers) - return components, block_state + self.set_block_state(state, block_state) + return components, state -class WanAnimate2SegmentSchedulerResetStep(ModularPipelineBlocks): +class WanAnimate2SegmentSchedulerResetStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -297,7 +309,8 @@ def description(self) -> str: return ( "Step within the segment loop that resets the scheduler: each segment is an independent denoising " "trajectory, so the solver state and timesteps are re-prepared per segment. This block should be used " - "to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + "to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` object " + "(e.g. `WanAnimate2SegmentLoopWrapper`)." ) @property @@ -319,16 +332,18 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, **kwargs): + block_state = self.get_block_state(state) device = components._execution_device components.scheduler.set_timesteps(block_state.num_inference_steps, device=device) block_state.timesteps = components.scheduler.timesteps - return components, block_state + self.set_block_state(state, block_state) + return components, state -class WanAnimate2RefExtractStep(ModularPipelineBlocks): +class WanAnimate2RefExtractStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property @@ -337,7 +352,8 @@ def description(self) -> str: "Step within the segment loop that runs the transformer's reference-extraction pass " '(`kv_cache_mode="extract"`): the driving-video segment is encoded once and every layer\'s reference ' "K/V is stored in the KV cache, which the denoising forwards then attend over. This block should be " - "used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + "used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` object " + "(e.g. `WanAnimate2SegmentLoopWrapper`)." ) @property @@ -400,7 +416,8 @@ def inputs(self) -> list[InputParam]: ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, **kwargs): + block_state = self.get_block_state(state) device = components._execution_device transformer_dtype = components.transformer.dtype @@ -417,31 +434,33 @@ def __call__(self, components, block_state: BlockState, k: int): offset_grid_sizes=block_state.grid_sizes_ref, ) - return components, block_state + self.set_block_state(state, block_state) + return components, state # ======================================== -# Inner Denoising Blocks +# Denoising Loop Steps # ======================================== -class WanAnimate2SegmentDenoiseInner(ModularPipelineBlocks): +class WanAnimate2LoopDenoiser(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property def description(self) -> str: return ( - "Inner timestep loop that denoises one segment with guidance, attending over the segment's cached " - "reference K/V. The unconditional branch passes `is_uncondtion=True` to the transformer (it skips a " - "dedicated layer on that branch), routed through the guider as a per-branch input. This block should " - "be used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`." + "Step within the segment's denoising loop that predicts the noise with guidance, attending over the " + "segment's cached reference K/V. The unconditional branch passes `is_uncondtion=True` to the " + "transformer (it skips a dedicated layer on that branch), routed through the guider as a per-branch " + "input. This block should be used to compose the `sub_blocks` attribute of an " + "`IterativePipelineBlocks` object (e.g. `WanAnimate2DenoiseLoopWrapper`); it reads the current " + "timestep `t` and step index `i` from the loop scope." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", WanAnimate2Transformer3DModel), - ComponentSpec("scheduler", SchedulerMixin), ComponentSpec( "guider", ClassifierFreeGuidance, @@ -471,19 +490,7 @@ def inputs(self) -> list[InputParam]: type_hint=WanAnimate2KVCache, description="Per-segment cache holding every layer's reference K/V", ), - InputParam( - "timesteps", - required=True, - type_hint=torch.Tensor, - description="This segment's denoising timesteps", - ), InputParam.template("num_inference_steps", default=40), - InputParam( - "num_segments", - required=True, - type_hint=int, - description="Total number of segments in the driving video, from the video preprocess step", - ), InputParam( "max_seq_len", required=True, @@ -514,7 +521,6 @@ def inputs(self) -> list[InputParam]: type_hint=int, description="The resolved frame width in pixels", ), - InputParam.template("generator"), InputParam.template("prompt_embeds"), InputParam.template("negative_prompt_embeds"), InputParam.template("denoiser_input_fields"), @@ -523,11 +529,12 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam.template("latents"), + OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted noise for this step"), ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) transformer_dtype = components.transformer.dtype guider_inputs = { @@ -544,68 +551,53 @@ def __call__(self, components, block_state: BlockState, k: int): if name in transformer_args and name not in guider_inputs } - with tqdm( - total=len(block_state.timesteps), desc=f"Segment {k + 1}/{block_state.num_segments}" - ) as progress_bar: - for i, t in enumerate(block_state.timesteps): - timestep = torch.stack([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: - components.guider.prepare_models(components.transformer) - - guider_state_batch.noise_pred = components.transformer( - [block_state.latents.to(transformer_dtype)], - timestep=timestep, - encoder_hidden_states=[guider_state_batch.encoder_hidden_states[0].to(transformer_dtype)], - condition_latents=[block_state.reference_latents.to(transformer_dtype)], - kv_cache=block_state.kv_cache, - kv_cache_mode="cached", - seq_len=block_state.max_seq_len, - reference_grid_sizes=block_state.grid_sizes_ref, - origin_len=block_state.segment_frame_length, - origin_area=[block_state.height, block_state.width], - is_uncondtion=guider_state_batch.is_uncondtion, - **shared_kwargs, - ).sample[0] - - components.guider.cleanup_models(components.transformer) - - noise_pred = components.guider(guider_state)[0] - - latents = components.scheduler.step( - noise_pred.unsqueeze(0), - t, - block_state.latents.unsqueeze(0), - return_dict=False, - generator=block_state.generator, - )[0] - block_state.latents = latents.squeeze(0) + timestep = torch.stack([t]) - progress_bar.update() + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) - return components, block_state + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + guider_state_batch.noise_pred = components.transformer( + [block_state.latents.to(transformer_dtype)], + timestep=timestep, + encoder_hidden_states=[guider_state_batch.encoder_hidden_states[0].to(transformer_dtype)], + condition_latents=[block_state.reference_latents.to(transformer_dtype)], + kv_cache=block_state.kv_cache, + kv_cache_mode="cached", + seq_len=block_state.max_seq_len, + reference_grid_sizes=block_state.grid_sizes_ref, + origin_len=block_state.segment_frame_length, + origin_area=[block_state.height, block_state.width], + is_uncondtion=guider_state_batch.is_uncondtion, + **shared_kwargs, + ).sample[0] -class WanAnimate2DistilledSegmentDenoiseInner(WanAnimate2SegmentDenoiseInner): + components.guider.cleanup_models(components.transformer) + + block_state.noise_pred = components.guider(guider_state)[0] + + self.set_block_state(state, block_state) + return components, state + + +class WanAnimate2DistilledLoopDenoiser(WanAnimate2LoopDenoiser): model_name = "wan-animate-2-distilled" @property def description(self) -> str: return ( - "Inner timestep loop that denoises one segment for the distilled model, which is trained for few-step " - "sampling without classifier-free guidance — the guider defaults to `guidance_scale=1.0`, so only the " - "conditional branch runs. This block should be used to compose the `sub_blocks` attribute of " - "`WanAnimate2SegmentLoopWrapper`." + "Step within the segment's denoising loop that predicts the noise for the distilled model, which is " + "trained for few-step sampling without classifier-free guidance — the guider defaults to " + "`guidance_scale=1.0`, so only the conditional branch runs. This block should be used to compose the " + "`sub_blocks` attribute of an `IterativePipelineBlocks` object (e.g. `WanAnimate2DenoiseLoopWrapper`)." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", WanAnimate2Transformer3DModel), - ComponentSpec("scheduler", SchedulerMixin), ComponentSpec( "guider", ClassifierFreeGuidance, @@ -615,22 +607,169 @@ def expected_components(self) -> list[ComponentSpec]: ] +class WanAnimate2LoopAfterDenoiser(ModularLoopPipelineBlocks): + model_name = "wan-animate-2" + + @property + def description(self) -> str: + return ( + "Step within the segment's denoising loop that updates the latents after denoising. " + "This block should be used to compose the `sub_blocks` attribute of an `IterativePipelineBlocks` " + "object (e.g. `WanAnimate2DenoiseLoopWrapper`); it reads `noise_pred` and the current timestep `t` " + "from the loop scope." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", SchedulerMixin), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="This segment's latents", + ), + InputParam( + "noise_pred", + required=True, + type_hint=torch.Tensor, + description="The predicted noise for this step", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("latents"), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + + latents = components.scheduler.step( + block_state.noise_pred.unsqueeze(0), + t, + block_state.latents.unsqueeze(0), + return_dict=False, + generator=block_state.generator, + )[0] + block_state.latents = latents.squeeze(0) + + self.set_block_state(state, block_state) + return components, state + + +class WanAnimate2DenoiseLoopWrapper(IterativePipelineBlocks): + model_name = "wan-animate-2" + + @property + def loop_variables(self) -> list[str]: + return ["i", "t"] + + @property + def description(self) -> str: + return ( + "Pipeline block that iteratively denoises one segment's latents over `timesteps`, attending over the " + "segment's cached reference K/V. The specific steps within each iteration can be customized with the " + "`sub_blocks` attribute. It runs inside the segment loop and reads the current segment index `k` from " + "the loop scope." + ) + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + "timesteps", + required=True, + type_hint=torch.Tensor, + description="This segment's denoising timesteps", + ), + InputParam( + "num_segments", + required=True, + type_hint=int, + description="Total number of segments in the driving video, from the video preprocess step", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) + + with tqdm( + total=len(block_state.timesteps), desc=f"Segment {k + 1}/{block_state.num_segments}" + ) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + progress_bar.update() + + return components, state + + @torch.no_grad() + def stream(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + return components, state + + +class WanAnimate2SegmentDenoiseStep(WanAnimate2DenoiseLoopWrapper): + block_classes = [WanAnimate2LoopDenoiser, WanAnimate2LoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises one segment's latents with guidance, attending over the " + "segment's cached reference K/V. \n" + "Its loop logic is defined in `WanAnimate2DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `WanAnimate2LoopDenoiser`\n" + " - `WanAnimate2LoopAfterDenoiser`\n" + ) + + +class WanAnimate2DistilledSegmentDenoiseStep(WanAnimate2DenoiseLoopWrapper): + model_name = "wan-animate-2-distilled" + + block_classes = [WanAnimate2DistilledLoopDenoiser, WanAnimate2LoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises one segment's latents for the distilled model, which is " + "trained for few-step sampling without classifier-free guidance. \n" + "Its loop logic is defined in `WanAnimate2DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `WanAnimate2DistilledLoopDenoiser`\n" + " - `WanAnimate2LoopAfterDenoiser`\n" + ) + + # ======================================== # Post-Denoise # ======================================== -class WanAnimate2SegmentDecodeStep(ModularPipelineBlocks): +class WanAnimate2SegmentDecodeStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property def description(self) -> str: return ( "Step within the segment loop that VAE-decodes the denoised segment. Decoding happens inside the loop " - "because the next segment conditions on this segment's decoded pixels. Finished frames move to CPU and " - "the per-segment KV cache and latents are freed — holding them across segments fragments the " - "allocator enough to OOM at high resolution. This block should be used to compose the `sub_blocks` " - "attribute of `WanAnimate2SegmentLoopWrapper`." + "because the next segment conditions on this segment's decoded pixels. The per-segment KV cache and " + "latents are freed — holding them across segments fragments the allocator enough to OOM at high " + "resolution. This block should be used to compose the `sub_blocks` attribute of an " + "`IterativePipelineBlocks` object (e.g. `WanAnimate2SegmentLoopWrapper`)." ) @property @@ -668,12 +807,14 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "out_frames", type_hint=torch.Tensor, - description="This segment's decoded frames on device; the next segment conditions on its tail", + description="This segment's decoded frames on device, consumed (and freed) by the collect step", ), ] @torch.no_grad() - def __call__(self, components, block_state: BlockState, k: int): + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) + latents = block_state.latents.to(torch.float32) # The first latent frame is the reference image's slot, not video content. out_frames = decode_vae(components.vae, latents[:, 1:]) @@ -681,7 +822,6 @@ def __call__(self, components, block_state: BlockState, k: int): if k > 0: out_frames = out_frames[:, :, block_state.prev_segment_conditioning_frames :] - block_state.segment_frames.append(out_frames.cpu()) block_state.out_frames = out_frames block_state.kv_cache.clear() @@ -689,38 +829,40 @@ def __call__(self, components, block_state: BlockState, k: int): block_state.latents = None torch.cuda.empty_cache() - return components, block_state - - -# ======================================== -# Segment Loop Wrapper -# ======================================== + self.set_block_state(state, block_state) + return components, state -class WanAnimate2SegmentLoopWrapper(LoopSequentialPipelineBlocks): +class WanAnimate2SegmentCollectStep(ModularLoopPipelineBlocks): model_name = "wan-animate-2" @property def description(self) -> str: return ( - "Pipeline block that iterates over the driving video's segments. At each segment it runs sub-blocks " - "for per-segment encoding, preparation, reference extraction, denoising, and decoding; each segment " - "conditions on the previous one's decoded tail frames." + "Step within the segment loop that appends this segment's decoded frames to the collected video, " + "moving them to CPU, and frees the on-device copy — the next segment's conditioning reads the tail " + "of the last collected entry. This block should be used to compose the `sub_blocks` attribute of an " + "`IterativePipelineBlocks` object (e.g. `WanAnimate2SegmentLoopWrapper`)." ) @property - def loop_inputs(self) -> list[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( - "num_segments", + "out_frames", required=True, - type_hint=int, - description="Total number of segments in the driving video, from the video preprocess step", + type_hint=torch.Tensor, + description="This segment's decoded frames on device, from the decode step", + ), + InputParam( + "segment_frames", + type_hint=list, + description="The segments collected so far; `None` before the first segment", ), ] @property - def loop_intermediate_outputs(self) -> list[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "segment_frames", @@ -730,19 +872,57 @@ def loop_intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState) -> PipelineState: + def __call__(self, components, state: PipelineState, **kwargs): block_state = self.get_block_state(state) + block_state.segment_frames = [*(block_state.segment_frames or []), block_state.out_frames.cpu()] + block_state.out_frames = None # release the on-device copy; nothing reads it after this + self.set_block_state(state, block_state) + return components, state - # Seed the loop-carried state: `segment_frames` collects each segment's decoded frames (the decode step - # appends to it); `out_frames` is the previous segment's decoded frames — written by the decode step, read - # by the prev-frames step of the next iteration. `None` marks "no previous segment" for the first iteration. - block_state.segment_frames = [] - block_state.out_frames = None +# ======================================== +# Segment Loop Wrapper +# ======================================== + + +class WanAnimate2SegmentLoopWrapper(IterativePipelineBlocks): + model_name = "wan-animate-2" + + @property + def loop_variables(self) -> list[str]: + return ["k"] + + @property + def description(self) -> str: + return ( + "Pipeline block that iterates over the driving video's segments. At each segment it runs sub-blocks " + "for per-segment encoding, preparation, reference extraction, denoising, decoding, and collection; " + "each segment conditions on the previous one's decoded tail frames." + ) + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + "num_segments", + required=True, + type_hint=int, + description="Total number of segments in the driving video, from the video preprocess step", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState): + block_state = self.get_block_state(state) for k in range(block_state.num_segments): - components, block_state = self.loop_step(components, block_state, k=k) + components, state = self.loop_step(components, state, k=k) + return components, state - self.set_block_state(state, block_state) + @torch.no_grad() + def stream(self, components, state: PipelineState): + block_state = self.get_block_state(state) + for k in range(block_state.num_segments): + components, state = yield from self.stream_step(components, state, k=k) return components, state @@ -758,8 +938,9 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, WanAnimate2RefExtractStep, - WanAnimate2SegmentDenoiseInner, + WanAnimate2SegmentDenoiseStep, WanAnimate2SegmentDecodeStep, + WanAnimate2SegmentCollectStep, ] block_names = [ "vae_encoder", @@ -769,6 +950,7 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): "ref_extract", "denoise_inner", "decode", + "collect", ] @property @@ -776,7 +958,7 @@ def description(self) -> str: return ( "Segment denoise step that iterates over the driving video's segments.\n" "At each segment: vae_encoder -> prev_frames -> prepare -> scheduler_reset -> ref_extract -> " - "denoise_inner -> decode." + "denoise_inner (a nested denoising loop over this segment's timesteps) -> decode -> collect." ) @@ -789,8 +971,9 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, WanAnimate2RefExtractStep, - WanAnimate2DistilledSegmentDenoiseInner, + WanAnimate2DistilledSegmentDenoiseStep, WanAnimate2SegmentDecodeStep, + WanAnimate2SegmentCollectStep, ] block_names = [ "vae_encoder", @@ -800,6 +983,7 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): "ref_extract", "denoise_inner", "decode", + "collect", ] @property @@ -807,5 +991,6 @@ def description(self) -> str: return ( "Segment denoise step for the distilled model that iterates over the driving video's segments.\n" "At each segment: vae_encoder -> prev_frames -> prepare -> scheduler_reset -> ref_extract -> " - "denoise_inner (no classifier-free guidance) -> decode." + "denoise_inner (a nested denoising loop over this segment's timesteps, no classifier-free guidance) " + "-> decode -> collect." ) diff --git a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py index f77eb378c15b..1d569b328701 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2.py @@ -175,8 +175,6 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): The reference conditioning tensor `[20, 1, latent_height, latent_width]`; provides the latent grid driving_video_pixels (`Tensor`): The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step - num_segments (`int`): - Total number of segments in the driving video, from the video preprocess step effective_segment (`int`): Frames each segment advances: `segment_frame_length - prev_segment_conditioning_frames`, from the video preprocess step @@ -190,6 +188,8 @@ class WanAnimate2CoreDenoiseStep(SequentialPipelineBlocks): CLIP vision features of the driving video's first frame prompt_ref_embeds (`Tensor`): Text embeddings of the reference prompt, guiding the reference-extraction pass + num_segments (`int`): + Total number of segments in the driving video, from the video preprocess step height (`int`): The resolved frame height in pixels width (`int`): diff --git a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py index 8eab815897da..2f531ead16bd 100644 --- a/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py +++ b/src/diffusers/modular_pipelines/wan_animate_2/modular_blocks_wan_animate_2_distilled.py @@ -175,8 +175,6 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): The reference conditioning tensor `[20, 1, latent_height, latent_width]`; provides the latent grid driving_video_pixels (`Tensor`): The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step - num_segments (`int`): - Total number of segments in the driving video, from the video preprocess step effective_segment (`int`): Frames each segment advances: `segment_frame_length - prev_segment_conditioning_frames`, from the video preprocess step @@ -190,6 +188,8 @@ class WanAnimate2DistilledCoreDenoiseStep(SequentialPipelineBlocks): CLIP vision features of the driving video's first frame prompt_ref_embeds (`Tensor`): Text embeddings of the reference prompt, guiding the reference-extraction pass + num_segments (`int`): + Total number of segments in the driving video, from the video preprocess step height (`int`): The resolved frame height in pixels width (`int`): diff --git a/src/diffusers/pipelines/ltx2/__init__.py b/src/diffusers/pipelines/ltx2/__init__.py index d4aa35127403..d48d890f4cb6 100644 --- a/src/diffusers/pipelines/ltx2/__init__.py +++ b/src/diffusers/pipelines/ltx2/__init__.py @@ -30,12 +30,12 @@ _import_structure["pipeline_ltx2_condition"] = ["LTX2ConditionPipeline", "LTX2VideoCondition"] _import_structure["pipeline_ltx2_dfr"] = ["LTX2DFRPipeline"] _import_structure["pipeline_ltx2_dfr_temporal_refine"] = ["LTX2DFRTemporalRefinePipeline"] - _import_structure["pipeline_output"] = ["LTX2DFRPipelineOutput", "LTX2PipelineOutput", "LTX2VideoDecodeOutput"] _import_structure["pipeline_ltx2_diffusion_decode"] = ["LTX2VideoDiffusionDecodePipeline"] _import_structure["pipeline_ltx2_hdr_lora"] = ["LTX2HDRPipeline", "LTX2HDRReferenceCondition"] _import_structure["pipeline_ltx2_ic_lora"] = ["LTX2InContextPipeline", "LTX2ReferenceCondition"] _import_structure["pipeline_ltx2_image2video"] = ["LTX2ImageToVideoPipeline"] _import_structure["pipeline_ltx2_latent_upsample"] = ["LTX2LatentUpsamplePipeline"] + _import_structure["pipeline_output"] = ["LTX2DFRPipelineOutput", "LTX2PipelineOutput", "LTX2VideoDecodeOutput"] _import_structure["vocoder"] = ["LTX2Vocoder", "LTX2VocoderWithBWE"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 1598814f835a..b7ebd3be0bb0 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2614,6 +2614,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class IterativePipelineBlocks(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class LoopSequentialPipelineBlocks(metaclass=DummyObject): _backends = ["torch"] @@ -2629,6 +2644,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class ModularLoopPipelineBlocks(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class ModularPipeline(metaclass=DummyObject): _backends = ["torch"] @@ -2689,6 +2719,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class StreamEvent(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + def get_constant_schedule(*args, **kwargs): requires_backends(get_constant_schedule, ["torch"]) 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..4700c32d88f7 --- /dev/null +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -0,0 +1,364 @@ +# 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 with trivially checkable arithmetic, in the same nested shape as an autoregressive +# video pipeline: an outer loop (variable `k`) whose steps add 1 to `x` and record the result, with a +# nested inner loop (variable `i`) that multiplies `x` by 10 on each of its iterations. +# +# OuterLoop loop over k in range(num_outer_steps) +# ├─ add_one AddOneStep x += 1 +# ├─ times_ten TimesTenLoop the inner loop, over i in range(num_inner_steps) +# │ ├─ compute_delta delta = x * 9 +# │ └─ apply_delta x += delta (net effect of one inner iteration: x *= 10) +# └─ record RecordStep xs.append(x) + + +class AddOneStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="x", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="x")] + + @property + def description(self): + return "adds 1 to x" + + def __call__(self, components, state, k): + block_state = self.get_block_state(state) + block_state.x = block_state.x + 1 + self.set_block_state(state, block_state) + return components, state + + +class ComputeDeltaStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="x", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="delta")] + + @property + def description(self): + return "computes this inner iteration's increment" + + def __call__(self, components, state, i): + block_state = self.get_block_state(state) + block_state.delta = block_state.x * 9 + self.set_block_state(state, block_state) + return components, state + + +class ApplyDeltaStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="x", required=True), InputParam(name="delta", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="x")] + + @property + def description(self): + return "applies the increment to x" + + def __call__(self, components, state, i): + block_state = self.get_block_state(state) + block_state.x = block_state.x + block_state.delta + self.set_block_state(state, block_state) + return components, state + + +class InnerLoopWrapper(IterativePipelineBlocks): + model_name = "test" + + @property + def description(self): + return "inner loop over num_inner_steps" + + @property + def loop_variables(self): + return ["i"] + + @property + def loop_inputs(self): + return [InputParam(name="num_inner_steps", required=True)] + + @torch.no_grad() + def __call__(self, components, state, **kwargs): # ignores the outer loop's `k` + block_state = self.get_block_state(state) + for i in range(block_state.num_inner_steps): + components, state = self.loop_step(components, state, i=i) + return components, state + + +class TimesTenLoop(InnerLoopWrapper): + block_classes = [ComputeDeltaStep, ApplyDeltaStep] + block_names = ["compute_delta", "apply_delta"] + + +class RecordStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="x", required=True), InputParam(name="xs", default=None)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="xs")] + + @property + def description(self): + return "records x after this outer iteration" + + def __call__(self, components, state, **kwargs): # ignores the loop's `k`: a catch-all is enough + block_state = self.get_block_state(state) + block_state.xs = [*(block_state.xs or []), float(block_state.x)] + self.set_block_state(state, block_state) + return components, state + + +class OuterLoopWrapper(IterativePipelineBlocks): + model_name = "test" + + @property + def description(self): + return "outer loop over num_outer_steps" + + @property + def loop_variables(self): + return ["k"] + + @property + def loop_inputs(self): + return [InputParam(name="num_outer_steps", required=True)] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + for k in range(block_state.num_outer_steps): + components, state = self.loop_step(components, state, k=k) + return components, state + + +class OuterLoop(OuterLoopWrapper): + block_classes = [AddOneStep, TimesTenLoop, RecordStep] + block_names = ["add_one", "times_ten", "record"] + + +class TimestepLoopWrapper(IterativePipelineBlocks): + """Loop wrapper over `i`, `t` — the signature-validation tests assemble it with different steps.""" + + model_name = "test" + + @property + def description(self): + return "loop over i, t" + + @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 + + +# OuterLoop loop over k in range(num_outer_steps) +# ├─ add_one AddOneStep x += 1 +# ├─ times_ten TimesTenLoop the inner loop, over i in range(num_inner_steps) +# │ ├─ compute_delta delta = x * 9 +# │ └─ apply_delta x += delta (net effect of one inner iteration: x *= 10) +# └─ record RecordStep xs.append(x) + + +class TestIterativePipelineBlocksStructure: + def test_inputs_aggregation(self): + loop = OuterLoop() + input_names = [p.name for p in loop.inputs] + + # the outer loop logic's own input, declared in its `loop_inputs` + assert "num_outer_steps" in [p.name for p in loop.loop_inputs] + assert "num_outer_steps" in input_names + # the nested inner loop's `loop_inputs` entry is aggregated too + assert "num_inner_steps" in [p.name for p in loop.sub_blocks["times_ten"].loop_inputs] + assert "num_inner_steps" in input_names + # loop variables are call arguments, not inputs + assert "k" not in input_names + assert "i" not in input_names + # `x` is read by the first step (`add_one`) before any step writes it -> a pipeline input + assert "x" in input_names + # `xs` is the accumulator: written by `record` at the end of iteration k, read by it again at + # k + 1 — the read comes first, so it is an (optional, default None) pipeline input + assert "xs" in input_names + # `delta` is written by `compute_delta` before `apply_delta` reads it -> satisfied inside the loop + assert "delta" not in input_names + + def test_sub_block_outputs_are_aggregated(self): + loop = OuterLoop() + output_names = [o.name for o in loop.intermediate_outputs] + assert "x" in output_names + assert "xs" in output_names + assert "delta" in output_names + + def test_loop_block_can_nest_assembled_blocks(self): + # the nested inner loop stays an assembled IterativePipelineBlocks sub-block + loop = OuterLoop() + assert isinstance(loop.sub_blocks["times_ten"], IterativePipelineBlocks) + assert list(loop.sub_blocks["times_ten"].sub_blocks) == ["compute_delta", "apply_delta"] + + +class TestIterativePipelineBlocksExecution: + def _make_pipeline(self): + return SequentialPipelineBlocks.from_blocks_dict({"loop": OuterLoop()}).init_pipeline() + + def test_nested_loop(self): + pipe = self._make_pipeline() + # per outer step: x += 1, then the inner loop doubles the digits (x *= 10 per inner step), then record + # k=0: (0 + 1) * 10 * 10 = 100 ; k=1: (100 + 1) * 10 * 10 = 10100 + state = pipe(x=torch.tensor(0.0), num_outer_steps=2, num_inner_steps=2) + + assert state.get("xs") == [100.0, 10100.0] + # the carried value persists as a declared output + assert float(state.get("x")) == 10100.0 + + def test_loop_variables_do_not_leak_into_state(self): + pipe = self._make_pipeline() + state = pipe(x=torch.tensor(0.0), num_outer_steps=2, num_inner_steps=1) + + for name in ("k", "i"): + assert state.get(name) is None + # declared sub-block outputs persist after the loop (last iteration's value) + assert state.get("delta") is not None + + def test_block_state_is_loop_scoped(self): + # the loop's block state holds only the loop logic's own inputs; sub-block values live in the pipeline state + pipe = self._make_pipeline() + state = pipe(x=torch.tensor(0.0), num_outer_steps=2, num_inner_steps=1) + block_state = pipe.blocks.sub_blocks["loop"].get_block_state(state) + assert block_state.as_dict().keys() == {"num_outer_steps"} + + 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() + + @staticmethod + def _loop_over(step_cls): + # assemble the shared `TimestepLoopWrapper` with the given step + class Loop(TimestepLoopWrapper): + block_classes = [step_cls] + block_names = ["step"] + + return Loop + + def test_leaf_signature_is_validated(self): + # a named parameter that isn't a loop variable fails at construction, even with a catch-all + class UnknownVarStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "loop step naming a variable the loop doesn't have" + + def __call__(self, components, state, k, **kwargs): + return components, state + + with pytest.raises(ValueError, match="not loop variables"): + self._loop_over(UnknownVarStep)() + + # naming only some loop variables without a `**kwargs` catch-all fails at construction + class MissingVarStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "loop step naming only one of the loop variables, without a catch-all" + + def __call__(self, components, state, t): + return components, state + + with pytest.raises(ValueError, match="must accept the loop variables"): + self._loop_over(MissingVarStep)() + + def test_leaf_signature_subset_with_kwargs_is_valid(self): + # naming only the used loop variables plus `**kwargs` is accepted + class SubsetStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "loop step using only `t`, ignoring the rest via a catch-all" + + def __call__(self, components, state, t, **kwargs): + return components, state + + self._loop_over(SubsetStep)() # does not raise + + 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({"compute_delta": ComputeDeltaStep()}).init_pipeline() + with pytest.raises(TypeError): + pipe(x=torch.tensor(1.0)) diff --git a/tests/modular_pipelines/testing_utils/common.py b/tests/modular_pipelines/testing_utils/common.py index 47614fd51005..af0df62b8750 100644 --- a/tests/modular_pipelines/testing_utils/common.py +++ b/tests/modular_pipelines/testing_utils/common.py @@ -361,3 +361,29 @@ def test_num_images_per_prompt(self, batch_sizes=[1, 2], num_images_per_prompts= images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt, output=self.output_name) assert images.shape[0] == batch_size * num_images_per_prompt + + def test_stream_matches_call(self, expected_max_diff=1e-4): + pipe = self.get_pipeline().to(torch_device) + + if not pipe.blocks.supports_streaming: + pytest.skip("Skipping test as blocks do not support streaming.") + + inputs = self.get_dummy_inputs() + inputs["generator"] = self.get_generator(0) + output = pipe(**inputs, output=self.output_name) + + inputs = self.get_dummy_inputs() + inputs["generator"] = self.get_generator(0) + num_events = 0 + generator = pipe.stream(**inputs) + while True: + try: + next(generator) + except StopIteration as e: + state = e.value + break + num_events += 1 + + assert num_events > 0, "stream() yielded no events" + max_diff = torch.abs(state.get(self.output_name) - output).max() + assert max_diff < expected_max_diff, "stream() results different from __call__ results"