diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 68177c7e7fb8..1908065f760c 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -392,6 +392,72 @@ 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 — 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. The generator's return value is the final state, exactly what `__call__` returns. + +```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 — the inner loop's events surface too, so a consumer that only wants finished chunks filters on the outer path: + +```py +for event in pipeline.stream(...): + if event.path == "denoise": # the chunk loop, not "denoise.denoise_inner" + show(event.state.get("out_frames")) +``` + +To stop early, stop iterating (or call `generator.close()`); nothing needs cleaning up. Blocks without loops run to completion and yield nothing. + +Streaming is opt-in per loop block. An [`IterativePipelineBlocks`] implements its loop in `__call__` as usual and, to support streaming, also implements `stream` — 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`). + +If you need to own the loop yourself — a serving engine that advances every request by one denoising step per tick, or a real-time pipeline fed one chunk of input at a time — run the blocks before the loop, then call the loop block's `loop_step` once per iteration. Anything you write into the state between calls is seen by the next iteration: + +```py +from diffusers.modular_pipelines import PipelineState + +loop = pipeline.blocks.sub_blocks["denoise"] + +state = PipelineState() +for param in pipeline.blocks.inputs: # seed the declared defaults + state.set(param.name, param.default) +state.set("prompt", "a cat") +state.set("num_inference_steps", 20) +# ... run the blocks before `denoise` on `state` ... +for i, t in enumerate(state.get("timesteps")): + _, state = loop.loop_step(pipeline, state, i=i, t=t) +``` + ## Modular repository A repository is required if the pipeline blocks use *pretrained components*. The repository supplies loading specifications and metadata. diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index b7d79b8ee97d..88ede2fd009a 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -376,11 +376,14 @@ "ConditionalPipelineBlocks", "ConfigSpec", "InputParam", + "IterativePipelineBlocks", "LoopSequentialPipelineBlocks", + "ModularLoopPipelineBlocks", "ModularPipeline", "ModularPipelineBlocks", "OutputParam", "SequentialPipelineBlocks", + "StreamEvent", ] ) _import_structure["optimization"] = [ @@ -1244,11 +1247,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 5cc5a4e57abc..972a40b48d83 100644 --- a/src/diffusers/modular_pipelines/ltx2/denoise.py +++ b/src/diffusers/modular_pipelines/ltx2/denoise.py @@ -22,12 +22,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 from .guider import LTX2Guidance @@ -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 @@ -425,10 +506,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 @@ -450,10 +533,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 ) @@ -464,10 +568,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 @@ -492,6 +598,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( @@ -505,8 +623,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 @@ -534,10 +661,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 @@ -564,6 +693,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, @@ -578,8 +719,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). @@ -596,12 +746,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 ( @@ -610,11 +766,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]: @@ -627,19 +785,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 3aa2c854dfe6..5d14dfd4e10b 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, @@ -574,6 +594,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] @@ -601,6 +640,48 @@ 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 accepts the enclosing loop's variables as call arguments — its signature must name exactly the + loop's `loop_variables` (e.g. `def __call__(self, components, state, i, t)`), which the loop validates before the + first iteration. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + def __call__(self, components, state: PipelineState, **kwargs) -> PipelineState: + # Subclasses name the enclosing loop's variables explicitly, e.g. + # `def __call__(self, components, state, i, t)`. + raise NotImplementedError(f"`__call__` method must be implemented in {self.__class__.__name__}") + class ConditionalPipelineBlocks(ModularPipelineBlocks): """ @@ -792,12 +873,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 is transparent in event paths, as it is in + # `get_execution_blocks`: events carry the name this conditional block has in its parent, not the branch name. + 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 @@ -820,7 +929,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 @@ -1153,6 +1262,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): """ @@ -1206,13 +1337,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"): @@ -1322,6 +1453,227 @@ def _requirements(self) -> dict[str, str]: return requirements +class IterativePipelineBlocks(SequentialPipelineBlocks): + """ + A pipeline blocks that runs its sub-blocks multiple times. Subclasses declare their loop-variable names in + `loop_variables` and implement `__call__` with their loop logic — the same way leaf blocks implement `__call__` + around `get_block_state` — calling `loop_step` once per iteration with the loop variables: + + ```python + @property + def loop_variables(self): + return ["i", "t"] + + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + return components, state + ``` + + Unlike [`LoopSequentialPipelineBlocks`], sub-blocks operate on the full [`PipelineState`] with the regular + `get_block_state`/`set_block_state` behavior, so an `IterativePipelineBlocks` can itself be a sub-block of another + one — loops can be nested and composed freely. Sub-blocks must be [`ModularLoopPipelineBlocks`] (loop steps) or + nested `IterativePipelineBlocks`, which is validated at construction. + + Loop variables are passed to sub-blocks as call arguments: every sub-block must have the signature `__call__(self, + components, state, )`, which is validated against `loop_variables` at construction. A nested + loop accepts the outer loop's variables in its own hand-written `__call__` (ignoring or forwarding them) and passes + its own `loop_variables` to its own sub-blocks: + + ```python + class InnerDenoiseLoop(IterativePipelineBlocks): + @property + def loop_variables(self): + return ["i", "t"] # what it passes to ITS sub-blocks + + @torch.no_grad() + def __call__(self, components, state, k): # accepts the OUTER chunk loop's variable + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + return components, state + ``` + + Sub-block outputs are written to the pipeline state as usual and persist after the loop. The loop logic's own + inputs (e.g. `timesteps`) and outputs are declared in `loop_inputs` / `loop_intermediate_outputs`: they are + surfaced alongside 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) + # sub_blocks are assigned after __init__ on this path, so validate again + instance._validate_sub_blocks() + return instance + + def _validate_sub_blocks(self): + """Sub-blocks must be loop steps (`ModularLoopPipelineBlocks`) or nested loops (`IterativePipelineBlocks`) + and accept exactly the loop variables after `(components, state)`.""" + expected = set(self.loop_variables) + for block_name, block in self.sub_blocks.items(): + if not isinstance(block, (ModularLoopPipelineBlocks, IterativePipelineBlocks)): + raise ValueError( + f"Sub-block '{block_name}' ({block.__class__.__name__}) of {self.__class__.__name__} must be " + "a `ModularLoopPipelineBlocks` (a loop step) or an `IterativePipelineBlocks` (a nested loop); " + f"got `{block.__class__.__bases__[0].__name__}`." + ) + params = list(inspect.signature(block.__call__).parameters) + extra = set(params[2:]) + if extra != expected: + raise ValueError( + f"Loop sub-block '{block_name}' ({block.__class__.__name__}) of {self.__class__.__name__} " + f"must accept the loop variables {sorted(expected)} after `(components, state)`; " + f"its `__call__` accepts {sorted(extra)}." + ) + + def loop_step(self, components, state: PipelineState, **loop_kwargs) -> PipelineState: + """Run all sub-blocks once over the pipeline state (one loop iteration), passing the loop variables.""" + for block_name, block in self.sub_blocks.items(): + try: + components, state = block(components, state, **loop_kwargs) + except Exception as e: + error_msg = ( + f"\nError in block: ({block_name}, {block.__class__.__name__})\n" + f"Error details: {str(e)}\n" + f"Traceback:\n{traceback.format_exc()}" + ) + logger.error(error_msg) + raise + return components, state + + def __call__(self, components, state: PipelineState, **kwargs) -> PipelineState: + # Subclasses implement their loop logic here. When the loop is nested inside another + # IterativePipelineBlocks, the signature must also accept the outer loop's variables, + # e.g. `def __call__(self, components, state, k)`. + raise NotImplementedError("`__call__` method needs to be implemented by the subclass") + + 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 @@ -1339,6 +1691,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.""" @@ -1594,25 +1950,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) @@ -2940,3 +3277,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 — each + denoising step, each segment of a chunked video, and so on — 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..6d84e5d8d20f 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,12 @@ def intermediate_outputs(self) -> list[OutputParam]: type_hint=int, description="Packed sequence length of the reference tokens", ), + OutputParam( + "out_frames", + type_hint=torch.Tensor, + description="The previous segment's decoded frames, carried across the segment loop; starts as " + "`None` (the first segment has no previous segment to condition on)", + ), ] @torch.no_grad() @@ -107,6 +113,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.out_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..d1798e042277 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( + "out_frames", + type_hint=torch.Tensor, + description="The previous segment's decoded frames on device, written by the decode step of the previous iteration; `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:] @@ -226,10 +234,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 +246,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 +279,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 block_state.latents = randn_tensor( @@ -286,10 +296,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 +308,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 +331,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, k: int): + 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 +351,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 +415,8 @@ def inputs(self) -> list[InputParam]: ] @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 transformer_dtype = components.transformer.dtype @@ -417,31 +433,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 +489,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 +520,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 +528,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 +550,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) + + 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] - return components, block_state + components.guider.cleanup_models(components.transformer) + block_state.noise_pred = components.guider(guider_state)[0] -class WanAnimate2DistilledSegmentDenoiseInner(WanAnimate2SegmentDenoiseInner): + 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 +606,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 @@ -673,7 +811,9 @@ 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) + 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 +821,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,7 +828,8 @@ def __call__(self, components, block_state: BlockState, k: int): block_state.latents = None torch.cuda.empty_cache() - return components, block_state + self.set_block_state(state, block_state) + return components, state # ======================================== @@ -697,9 +837,13 @@ def __call__(self, components, block_state: BlockState, k: int): # ======================================== -class WanAnimate2SegmentLoopWrapper(LoopSequentialPipelineBlocks): +class WanAnimate2SegmentLoopWrapper(IterativePipelineBlocks): model_name = "wan-animate-2" + @property + def loop_variables(self) -> list[str]: + return ["k"] + @property def description(self) -> str: return ( @@ -721,6 +865,7 @@ def loop_inputs(self) -> list[InputParam]: @property def loop_intermediate_outputs(self) -> list[OutputParam]: + # the loop logic collects each segment's decoded frames return [ OutputParam( "segment_frames", @@ -730,19 +875,29 @@ def loop_intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState) -> PipelineState: + def __call__(self, components, state: PipelineState): block_state = self.get_block_state(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. + # `segment_frames` collects each segment's decoded frames on CPU; `out_frames` (this segment's frames, + # on device) stays in the state for the prev-frames step of the next iteration to condition on. block_state.segment_frames = [] - block_state.out_frames = None - 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) + block_state.segment_frames.append(state.get("out_frames").cpu()) + 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) + block_state.segment_frames = [] + for k in range(block_state.num_segments): + components, state = yield from self.stream_step(components, state, k=k) + block_state.segment_frames.append(state.get("out_frames").cpu()) self.set_block_state(state, block_state) + return components, state @@ -758,7 +913,7 @@ class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, WanAnimate2RefExtractStep, - WanAnimate2SegmentDenoiseInner, + WanAnimate2SegmentDenoiseStep, WanAnimate2SegmentDecodeStep, ] block_names = [ @@ -776,7 +931,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." ) @@ -789,7 +944,7 @@ class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper): WanAnimate2SegmentPrepareStep, WanAnimate2SegmentSchedulerResetStep, WanAnimate2RefExtractStep, - WanAnimate2DistilledSegmentDenoiseInner, + WanAnimate2DistilledSegmentDenoiseStep, WanAnimate2SegmentDecodeStep, ] block_names = [ @@ -807,5 +962,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." ) 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/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index f34008252ab6..d4691a6a3a76 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2599,6 +2599,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"] @@ -2614,6 +2629,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"] @@ -2674,6 +2704,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..9054bb0d535e --- /dev/null +++ b/tests/modular_pipelines/test_iterative_pipeline_blocks.py @@ -0,0 +1,340 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest +import torch + +from diffusers.modular_pipelines import ( + InputParam, + IterativePipelineBlocks, + ModularLoopPipelineBlocks, + ModularPipelineBlocks, + OutputParam, + SequentialPipelineBlocks, +) + + +# Dummy blocks modeled on the Helios chunk-loop use case: an outer autoregressive chunk loop +# (history carried across chunks) containing a full inner timestep denoising loop. Loop variables +# (`k` for the chunk loop, `i`/`t` for the timestep loop) are passed to leaf sub-blocks as call +# arguments; every leaf sub-block of a loop must accept its loop's variables. + + +class ChunkNoiseGenStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="history", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="chunk_latents")] + + @property + def description(self): + return "prepares this chunk's latents from the history" + + def __call__(self, components, state, k): + block_state = self.get_block_state(state) + block_state.chunk_latents = block_state.history + k + self.set_block_state(state, block_state) + return components, state + + +class LoopDenoiserStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="chunk_latents", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="noise_pred")] + + @property + def description(self): + return "predicts the noise for one timestep" + + def __call__(self, components, state, i, t): + block_state = self.get_block_state(state) + block_state.noise_pred = block_state.chunk_latents * 0 + t + self.set_block_state(state, block_state) + return components, state + + +class LoopSchedulerStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="chunk_latents", required=True), InputParam(name="noise_pred", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="chunk_latents")] + + @property + def description(self): + return "updates the chunk latents with the noise prediction" + + def __call__(self, components, state, i, t): + block_state = self.get_block_state(state) + block_state.chunk_latents = block_state.chunk_latents + block_state.noise_pred + self.set_block_state(state, block_state) + return components, state + + +class InnerDenoiseLoop(IterativePipelineBlocks): + """Inner timestep loop — itself an assembled loop block, nested inside the chunk loop. + + Like every sub-block of the chunk loop, it accepts the outer loop variable `k` (and ignores it); + its own sub-blocks accept its own loop variables `i` / `t` instead. + """ + + model_name = "test" + block_classes = [LoopDenoiserStep, LoopSchedulerStep] + block_names = ["denoiser", "scheduler"] + + @property + def description(self): + return "inner timestep loop" + + @property + def loop_variables(self): + return ["i", "t"] + + @property + def loop_inputs(self): + return [InputParam(name="timesteps", required=True)] + + @torch.no_grad() + def __call__(self, components, state, k): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + return components, state + + +class ChunkUpdateStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(name="chunk_latents", required=True), InputParam(name="latent_chunks", default=None)] + + @property + def intermediate_outputs(self): + return [OutputParam(name="history"), OutputParam(name="latent_chunks")] + + @property + def description(self): + return "records the denoised chunk and updates the history" + + def __call__(self, components, state, k): + block_state = self.get_block_state(state) + block_state.history = block_state.chunk_latents + block_state.latent_chunks = [*(block_state.latent_chunks or []), float(block_state.chunk_latents)] + self.set_block_state(state, block_state) + return components, state + + +class ChunkLoop(IterativePipelineBlocks): + """Outer chunk loop containing the inner timestep loop as a sub-block.""" + + model_name = "test" + block_classes = [ChunkNoiseGenStep, InnerDenoiseLoop, ChunkUpdateStep] + block_names = ["noise_gen", "denoise", "update"] + + @property + def description(self): + return "outer autoregressive chunk loop" + + @property + def loop_variables(self): + return ["k"] + + @property + def loop_inputs(self): + return [InputParam(name="num_latent_chunk", required=True)] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + for k in range(block_state.num_latent_chunk): + components, state = self.loop_step(components, state, k=k) + return components, state + + +class CollectingChunkLoop(ChunkLoop): + """Chunk loop whose loop logic has an output of its own, written through `set_block_state`.""" + + @property + def loop_intermediate_outputs(self): + return [OutputParam(name="chunk_history")] + + @torch.no_grad() + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.chunk_history = [] + for k in range(block_state.num_latent_chunk): + components, state = self.loop_step(components, state, k=k) + block_state.chunk_history.append(float(state.get("history"))) + self.set_block_state(state, block_state) + return components, state + + +class TestIterativePipelineBlocksStructure: + def test_inputs_aggregation(self): + loop = ChunkLoop() + input_names = [p.name for p in loop.inputs] + + # inputs of the loop logic itself and of the nested loop are surfaced + assert "num_latent_chunk" in input_names + assert "timesteps" in input_names + # loop variables are call arguments, not inputs + assert "k" not in input_names + assert "i" not in input_names + assert "t" not in input_names + # cross-chunk carries surface as (optional) iteration-0 seeds + assert "history" in input_names + assert "latent_chunks" in input_names + + def test_sub_block_outputs_are_aggregated(self): + loop = ChunkLoop() + output_names = [o.name for o in loop.intermediate_outputs] + assert "history" in output_names + assert "latent_chunks" in output_names + + def test_loop_outputs_are_aggregated(self): + loop = CollectingChunkLoop() + output_names = [o.name for o in loop.intermediate_outputs] + assert "chunk_history" in output_names + assert "history" in output_names + + def test_loop_block_can_nest_assembled_blocks(self): + # the nested inner loop stays an assembled IterativePipelineBlocks sub-block + loop = ChunkLoop() + assert isinstance(loop.sub_blocks["denoise"], IterativePipelineBlocks) + assert list(loop.sub_blocks["denoise"].sub_blocks) == ["denoiser", "scheduler"] + + +class TestIterativePipelineBlocksExecution: + def _make_pipeline(self): + return SequentialPipelineBlocks.from_blocks_dict({"chunks": ChunkLoop()}).init_pipeline() + + def test_nested_chunk_loop(self): + pipe = self._make_pipeline() + # per chunk: chunk_latents = history + k, then += t for every timestep (1.0 + 2.0), + # then history <- chunk_latents + # chunk 0: 0 + 0 + 3 = 3 ; chunk 1: 3 + 1 + 3 = 7 ; chunk 2: 7 + 2 + 3 = 12 + state = pipe(num_latent_chunk=3, timesteps=torch.tensor([1.0, 2.0]), history=torch.tensor(0.0)) + + assert state.get("latent_chunks") == [3.0, 7.0, 12.0] + # the cross-chunk carry persists as a declared output + assert float(state.get("history")) == 12.0 + + def test_loop_variables_do_not_leak_into_state(self): + pipe = self._make_pipeline() + state = pipe(num_latent_chunk=2, timesteps=torch.tensor([1.0]), history=torch.tensor(0.0)) + + for name in ("k", "i", "t"): + assert state.get(name) is None + # declared sub-block outputs persist after the loop (last iteration's value) + assert state.get("noise_pred") is not None + + def test_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(num_latent_chunk=2, timesteps=torch.tensor([1.0]), history=torch.tensor(0.0)) + block_state = pipe.blocks.sub_blocks["chunks"].get_block_state(state) + assert block_state.as_dict().keys() == {"num_latent_chunk"} + + def test_loop_output_via_set_block_state(self): + pipe = SequentialPipelineBlocks.from_blocks_dict({"chunks": CollectingChunkLoop()}).init_pipeline() + state = pipe(num_latent_chunk=3, timesteps=torch.tensor([1.0, 2.0]), history=torch.tensor(0.0)) + assert state.get("chunk_history") == [3.0, 7.0, 12.0] + # sub-block outputs are untouched by the loop's own write-back + assert state.get("latent_chunks") == [3.0, 7.0, 12.0] + + def test_sub_block_type_is_validated(self): + # a regular ModularPipelineBlocks cannot be a loop sub-block: fails at construction + class PlainStep(ModularPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "regular block, not a loop step" + + def __call__(self, components, state): + return components, state + + class BadTypeLoop(IterativePipelineBlocks): + model_name = "test" + block_classes = [PlainStep] + block_names = ["plain"] + + @property + def description(self): + return "loop with a non-loop sub-block" + + with pytest.raises(ValueError, match="must be a `ModularLoopPipelineBlocks`"): + BadTypeLoop() + + def test_leaf_signature_is_validated(self): + # a loop step whose signature doesn't match the loop's variables fails at construction + class WrongSigStep(ModularLoopPipelineBlocks): + model_name = "test" + + @property + def description(self): + return "loop step with the wrong loop variables" + + def __call__(self, components, state, k): + return components, state + + class BadSigLoop(IterativePipelineBlocks): + model_name = "test" + block_classes = [WrongSigStep] + block_names = ["wrong"] + + @property + def description(self): + return "loop whose sub-block names the wrong loop variables" + + @property + def loop_variables(self): + return ["i", "t"] + + @property + def 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 + + with pytest.raises(ValueError, match="must accept the loop variables"): + BadSigLoop() + + def test_loop_leaf_standalone_raises(self): + # outside a loop, a leaf block with loop variables in its signature cannot run + pipe = SequentialPipelineBlocks.from_blocks_dict({"denoiser": LoopDenoiserStep()}).init_pipeline() + with pytest.raises(TypeError): + pipe(chunk_latents=torch.tensor(1.0)) 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"