From 530e86c725bfab101a21cfbb53cc32b9c55cfdf7 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 27 Oct 2026 05:27:10 +0100 Subject: [PATCH] [modular] derive required inputs and defaults from the ordered block walk Sequential and loop blocks decide each input's `required` flag and default in one ordered walk over their sub-blocks; `required_inputs` is derived from `inputs` as in conditional blocks, and the two overrides are gone. Composites declare their own `InputParam`/`OutputParam`/`ComponentSpec`/`ConfigSpec` for a derived entry with `*_overrides`, and get a FutureWarning until they settle a default two sub-blocks disagree on, a conditional block's per-branch defaults included. Leaf blocks declare what they write back, docstrings regenerated. Fixes #15009. Co-Authored-By: Claude Fable 5.1 --- .ai/references/modular.md | 18 +- .../sequential_pipeline_blocks.md | 1 + .../anima/modular_blocks_anima.py | 8 +- .../modular_pipelines/cosmos/denoise.py | 2 - .../cosmos/modular_blocks_cosmos3.py | 8 +- .../modular_pipelines/echo/denoise.py | 7 +- .../echo/modular_blocks_echo.py | 6 +- .../modular_pipelines/flux/before_denoise.py | 10 + .../flux/modular_blocks_flux.py | 16 +- .../flux/modular_blocks_flux_kontext.py | 12 + .../flux2/modular_blocks_flux2.py | 6 +- .../flux2/modular_blocks_flux2_klein.py | 12 +- .../flux2/modular_blocks_flux2_klein_base.py | 12 +- .../flux3/modular_blocks_flux3.py | 12 +- .../flux3/modular_blocks_flux3_distilled.py | 12 +- .../modular_pipelines/helios/denoise.py | 18 + .../helios/modular_blocks_helios.py | 10 +- .../helios/modular_blocks_helios_pyramid.py | 6 +- ...modular_blocks_helios_pyramid_distilled.py | 16 +- .../modular_blocks_hunyuan_video1_5.py | 14 +- .../modular_pipelines/ideogram4/denoise.py | 2 - .../ideogram4/modular_blocks_ideogram4.py | 4 +- .../modular_pipelines/krea2/before_denoise.py | 14 +- .../modular_pipelines/ltx/before_denoise.py | 4 + .../ltx/modular_blocks_ltx.py | 26 +- .../ltx2/modular_blocks_ltx2.py | 18 +- .../ltx2/modular_blocks_ltx25.py | 8 +- .../minimax_h3/modular_blocks_minimax_h3.py | 10 +- .../modular_pipelines/modular_pipeline.py | 387 +++++++--- .../modular_pipeline_utils.py | 47 ++ .../qwenimage/before_denoise.py | 6 + .../qwenimage/modular_blocks_qwenimage.py | 14 +- .../modular_blocks_qwenimage_edit.py | 10 +- .../qwenimage21/modular_blocks_qwenimage21.py | 8 +- .../modular_blocks_qwenimage21_turbo.py | 8 +- .../modular_blocks_stable_diffusion_3.py | 14 +- .../modular_blocks_stable_diffusion_xl.py | 14 +- .../modular_pipelines/wan/before_denoise.py | 19 + .../wan/modular_blocks_wan22_i2v.py | 2 +- .../wan/modular_blocks_wan_i2v.py | 12 +- .../z_image/before_denoise.py | 14 + .../z_image/modular_blocks_z_image.py | 10 +- .../test_sequential_pipeline_blocks.py | 669 ++++++++++++++++++ 43 files changed, 1270 insertions(+), 256 deletions(-) create mode 100644 tests/modular_pipelines/test_sequential_pipeline_blocks.py diff --git a/.ai/references/modular.md b/.ai/references/modular.md index 8ab88e65cc98..705119969707 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -282,20 +282,22 @@ if block_state.num_frames is None: block_state.num_frames = 189 ``` -A declared default is part of the block's contract, so the assembled pipeline is aware of it: the generated docstring shows it and `default_call_parameters` reports it. Resolved inside the body instead, the input renders as `*optional*` with no default, and nothing at the pipeline level can report what the block will actually do. Don't worry about branches of a conditional blockset declaring different defaults for the same input — each branch resolves its own at runtime. Resolve inside `__call__` only when the default is *computed* — derived from other inputs or component config (`height = components.default_sample_size * components.vae_scale_factor`). And when several blocks in a sequence share an input, declare the same default on each (or only on the first block that reads it): in a sequence the input is one shared value, so disagreeing declarations are silently resolved first-block-wins. +A declared default is part of the block's contract, so the assembled pipeline is aware of it: the generated docstring shows it and `default_call_parameters` reports it. Resolved inside the body instead, the input renders as `*optional*` with no default, and nothing at the pipeline level can report what the block will actually do. Don't worry about branches of a conditional blockset declaring different defaults for the same input — each branch resolves its own at runtime. Resolve inside `__call__` only when the default is *computed* — derived from other inputs or component config (`height = components.default_sample_size * components.vae_scale_factor`). And when several blocks in a sequence share an input, declare the same default on each (or only on the first block that reads it): in a sequence the input is one shared value, so the composite warns about disagreeing declarations until it declares the input itself in `input_overrides` (the first declaration wins meanwhile). A conditional block's per-branch defaults count as its opinion in that comparison; a declaration without a default has none. -**A composed `SequentialPipelineBlocks` can override `inputs` / `outputs`.** Two uses: narrow `outputs` to what downstream actually consumes, so the docstring shows the step's product instead of every internal intermediate (Wan-Animate-2's core denoise exposes only `segment_frames`); and change one input default per preset by mapping over `super().inputs` (the distilled core denoise turns `num_inference_steps` into `default=10`): +**A composite replaces what it derives from its sub-blocks with `*_overrides`.** `input_overrides`, `intermediate_output_overrides`, `component_overrides` and `config_overrides` are `{name: declaration}` class attributes on `ConditionalPipelineBlocks`, `SequentialPipelineBlocks` and `LoopSequentialPipelineBlocks` (also `from_blocks_dict` arguments). The name must already be in the derived list, and the declaration, an `InputParam` / `OutputParam` / `ComponentSpec` / `ConfigSpec`, takes the derived entry's place whole, so it states everything: default, `required`, description. Use it to settle a default two sub-blocks disagree on (the composite gets a warning until it does), to set a preset's default, or to reword a description from the composite's point of view. `InputParam.template(...)` keeps a common input to one line. ```python -@property -def inputs(self): +class DistilledCoreDenoiseStep(SequentialPipelineBlocks): + block_classes = [...] + block_names = [...] # The distilled checkpoint samples in few steps. - return [ - InputParam.template("num_inference_steps", default=10) if param.name == "num_inference_steps" else param - for param in super().inputs - ] + input_overrides = {"num_inference_steps": InputParam.template("num_inference_steps", default=10)} ``` +A composite reads its sub-blocks' declarations and never writes to them: the `required` flag it reports is its own view ("the caller has to pass this"), the leaf's flag is the runtime check in `get_block_state`, and the two legitimately differ. Copy an `InputParam` before changing it. + +Overriding the `outputs` property by hand is still the way to narrow `outputs` to what downstream actually consumes, so the docstring shows the step's product instead of every internal intermediate (Wan-Animate-2's core denoise exposes only `segment_frames`). + ## ComponentSpec patterns ```python diff --git a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md index 1bd67e17b8bf..9c7a7d5892cc 100644 --- a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md @@ -118,6 +118,7 @@ There are a few properties you should set: - `description`: We recommend adding a description for the assembled block to explain what the combined step does. - `model_name`: This is automatically derived from the sub-blocks but isn't always correct, so you may need to override it. - `outputs`: By default this is the same as `intermediate_outputs`, but you can manually set it to control which values appear in the doc. This is useful for showing only the final outputs instead of all intermediate values. +- `input_overrides`, `intermediate_output_overrides`, `component_overrides`, `config_overrides`: The composite's own declarations replacing entries of the aggregated lists, as `{name: declaration}`. For example `input_overrides = {"strength": InputParam.template("strength", default=0.9999)}` sets the default the whole sequence uses when two sub-blocks declare different ones. The declaration replaces the aggregated entry whole, and the name must already be in the aggregated list. These properties, together with the aggregated `inputs`, `intermediate_outputs`, and `expected_components`, are used to automatically generate the `doc` property. diff --git a/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py b/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py index f17538fd258f..4de17619f67a 100644 --- a/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py +++ b/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py @@ -236,7 +236,7 @@ class AnimaAutoCoreDenoiseStep(AutoPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*): The width in pixels of the generated image. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. @@ -244,7 +244,7 @@ class AnimaAutoCoreDenoiseStep(AutoPipelineBlocks): Strength for img2img/inpainting. generator (`Generator`, *optional*): Torch generator for deterministic generation. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. **denoiser_input_fields (`None`, *optional*): The conditional model inputs for the Anima denoiser. @@ -342,13 +342,13 @@ class AnimaAutoBlocks(SequentialPipelineBlocks): The number of images to generate per prompt. image_latents (`Tensor`, *optional*): image latents used to guide the image generation. Can be generated from vae_encoder step. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. strength (`float`, *optional*, defaults to 0.9): Strength for img2img/inpainting. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. **denoiser_input_fields (`None`, *optional*): The conditional model inputs for the Anima denoiser. diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index 2422defe4d46..b193b4eafb54 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -993,8 +993,6 @@ class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper): Timestep interval [lo, hi] over which text guidance is active (None = always). control_guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which control guidance is active (None = always). - latents (`Tensor`): - Noisy target latents to update. condition_latents (`Tensor`): Clean target latents on the conditioned frames (the autoregressive seed). diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index 864e687adad2..e43b48146509 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -1068,7 +1068,7 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): Torch generator for deterministic generation. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. mixed_precision_format (`str`, *optional*): None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is @@ -1099,7 +1099,7 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): Latent-frame indexes fixed by visual conditioning. num_frames (`int`, *optional*): Number of frames to generate. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy vision latents. sound_latents (`Tensor`, *optional*): Pre-generated noisy sound latents. @@ -1280,7 +1280,7 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Number of frames the first chunk reuses from the input video. generator (`Generator`, *optional*): Torch generator for deterministic generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. mixed_precision_format (`str`, *optional*): None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is @@ -1307,7 +1307,7 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Vision latents encoded from the conditioning image or video. vision_condition_frames (`list`, *optional*): Latent-frame indexes fixed by visual conditioning. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy vision latents. sound_latents (`Tensor`, *optional*): Pre-generated noisy sound latents. diff --git a/src/diffusers/modular_pipelines/echo/denoise.py b/src/diffusers/modular_pipelines/echo/denoise.py index cbe80c0ab609..cba80d9bd5fd 100644 --- a/src/diffusers/modular_pipelines/echo/denoise.py +++ b/src/diffusers/modular_pipelines/echo/denoise.py @@ -359,7 +359,8 @@ class EchoDenoiseLoopStep(LoopSequentialPipelineBlocks): transformer (`LTX2VideoTransformer3DModel`) Inputs: - sigmas (`list | tuple`): + sigmas (`list | tuple`, *optional*, defaults to (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, + 0.421875, 0.0)): DMD sigma schedule, including the terminal zero. latents (`Tensor`): Pre-generated noisy latents for image generation. @@ -401,10 +402,6 @@ class EchoDenoiseLoopStep(LoopSequentialPipelineBlocks): Frame rate of the generated video. attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. - audio_latents (`Tensor`): - Packed target audio tokens. - sigmas (`list | tuple`): - DMD sigma schedule including the terminal zero. generator (`Generator`, *optional*): Torch generator for deterministic generation. first_frame_tokens (`Tensor`, *optional*): diff --git a/src/diffusers/modular_pipelines/echo/modular_blocks_echo.py b/src/diffusers/modular_pipelines/echo/modular_blocks_echo.py index 16c95d3372bd..40786e7d49f1 100644 --- a/src/diffusers/modular_pipelines/echo/modular_blocks_echo.py +++ b/src/diffusers/modular_pipelines/echo/modular_blocks_echo.py @@ -120,7 +120,8 @@ class EchoDenoiseStep(SequentialPipelineBlocks): Optional initial audio noise in VAE form (B, C, L, M). generator (`Generator`, *optional*): Torch generator for deterministic generation. - sigmas (`list | tuple`): + sigmas (`list | tuple`, *optional*, defaults to (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, + 0.421875, 0.0)): DMD sigma schedule, including the terminal zero. attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. @@ -263,7 +264,8 @@ class EchoBlocks(SequentialPipelineBlocks): Optional initial audio noise in VAE form (B, C, L, M). generator (`Generator`, *optional*): Torch generator for deterministic generation. - sigmas (`list | tuple`): + sigmas (`list | tuple`, *optional*, defaults to (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, + 0.421875, 0.0)): DMD sigma schedule, including the terminal zero. attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. diff --git a/src/diffusers/modular_pipelines/flux/before_denoise.py b/src/diffusers/modular_pipelines/flux/before_denoise.py index 8e5a285dbe11..1f9e708b8408 100644 --- a/src/diffusers/modular_pipelines/flux/before_denoise.py +++ b/src/diffusers/modular_pipelines/flux/before_denoise.py @@ -355,6 +355,16 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" ), + OutputParam( + "height", + type_hint=int, + description="The height of the generated image, resolved to the default when not passed", + ), + OutputParam( + "width", + type_hint=int, + description="The width of the generated image, resolved to the default when not passed", + ), ] @staticmethod diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py b/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py index a725304c6af6..464d27c1f278 100644 --- a/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py +++ b/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py @@ -162,6 +162,10 @@ class FluxBeforeDenoiseStep(SequentialPipelineBlocks): Outputs: latents (`Tensor`): The initial latents to use for the denoising process + height (`int`): + The height of the generated image, resolved to the default when not passed + width (`int`): + The width of the generated image, resolved to the default when not passed timesteps (`Tensor`): The timesteps to use for inference num_inference_steps (`int`): @@ -227,6 +231,10 @@ class FluxImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks): Outputs: latents (`Tensor`): The initial latents to use for the denoising process + height (`int`): + The height of the generated image, resolved to the default when not passed + width (`int`): + The width of the generated image, resolved to the default when not passed timesteps (`Tensor`): The timesteps to use for inference num_inference_steps (`int`): @@ -268,9 +276,9 @@ class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks): scheduler (`FlowMatchEulerDiscreteScheduler`) Inputs: - height (`int`): + height (`int`, *optional*): TODO: Add description. - width (`int`): + width (`int`, *optional*): TODO: Add description. latents (`Tensor | NoneType`, *optional*): TODO: Add description. @@ -302,6 +310,10 @@ class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks): Outputs: latents (`Tensor`): The initial latents to use for the denoising process + height (`int`): + The height of the generated image, resolved to the default when not passed + width (`int`): + The width of the generated image, resolved to the default when not passed timesteps (`Tensor`): The timesteps to use for inference num_inference_steps (`int`): diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py b/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py index f937d2555b1a..0ac08abfdb65 100644 --- a/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py +++ b/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py @@ -153,6 +153,10 @@ class FluxKontextBeforeDenoiseStep(SequentialPipelineBlocks): Outputs: latents (`Tensor`): The initial latents to use for the denoising process + height (`int`): + The height of the generated image, resolved to the default when not passed + width (`int`): + The width of the generated image, resolved to the default when not passed timesteps (`Tensor`): The timesteps to use for inference num_inference_steps (`int`): @@ -220,6 +224,10 @@ class FluxKontextImageConditionedBeforeDenoiseStep(SequentialPipelineBlocks): Outputs: latents (`Tensor`): The initial latents to use for the denoising process + height (`int`): + The height of the generated image, resolved to the default when not passed + width (`int`): + The width of the generated image, resolved to the default when not passed timesteps (`Tensor`): The timesteps to use for inference num_inference_steps (`int`): @@ -291,6 +299,10 @@ class FluxKontextAutoBeforeDenoiseStep(AutoPipelineBlocks): Outputs: latents (`Tensor`): The initial latents to use for the denoising process + height (`int`): + The height of the generated image, resolved to the default when not passed + width (`int`): + The width of the generated image, resolved to the default when not passed timesteps (`Tensor`): The timesteps to use for inference num_inference_steps (`int`): diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py index 2bbb7975a983..8053a6540803 100644 --- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py @@ -314,11 +314,11 @@ class Flux2AutoBlocks(SequentialPipelineBlocks): TODO: Add description. image_latents (`list`, *optional*): TODO: Add description. - latents (`Tensor | NoneType`): + latents (`Tensor | NoneType`, *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. 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..3004cd73ed90 100644 --- a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py @@ -280,13 +280,13 @@ class Flux2KleinAutoCoreDenoiseStep(AutoPipelineBlocks): TODO: Add description. width (`int`, *optional*): TODO: Add description. - latents (`Tensor | NoneType`): + latents (`Tensor | NoneType`, *optional*): 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. @@ -355,11 +355,11 @@ class Flux2KleinAutoBlocks(SequentialPipelineBlocks): TODO: Add description. image_latents (`list`, *optional*): TODO: Add description. - latents (`Tensor | NoneType`): + latents (`Tensor | NoneType`, *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. 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..c91c51cf84ce 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 @@ -291,15 +291,15 @@ class Flux2KleinBaseAutoCoreDenoiseStep(AutoPipelineBlocks): TODO: Add description. width (`int`, *optional*): TODO: Add description. - latents (`Tensor | NoneType`): + latents (`Tensor | NoneType`, *optional*): TODO: Add description. generator (`None`, *optional*): TODO: Add description. image_latents (`list`, *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. @@ -367,13 +367,13 @@ class Flux2KleinBaseAutoBlocks(SequentialPipelineBlocks): TODO: Add description. num_images_per_prompt (`None`, *optional*, defaults to 1): TODO: Add description. - latents (`Tensor | NoneType`): + latents (`Tensor | NoneType`, *optional*): TODO: Add description. image_latents (`list`, *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. diff --git a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py index 31dcc62f7b0d..0dc01cb5cbc3 100644 --- a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py +++ b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py @@ -286,13 +286,13 @@ class Flux3AutoCoreDenoiseStep(AutoPipelineBlocks): The width in pixels of the generated image; defaults to the first reference image's own width reference_crop (`list | NoneType`, *optional*): Per reference, the box of the reference inside its padded canvas or None, from the preprocess step - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): Torch generator for deterministic generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. (overrides `num_inference_steps` when provided) sigmas (`list`, *optional*): Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided) @@ -379,13 +379,13 @@ class Flux3AutoBlocks(SequentialPipelineBlocks): The width in pixels of the generated image; defaults to the first reference image's own width reference_crop (`list | NoneType`, *optional*): Per reference, the box of the reference inside its padded canvas or None, from the preprocess step - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): Torch generator for deterministic generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. (overrides `num_inference_steps` when provided) sigmas (`list`, *optional*): Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided) diff --git a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py index d23f894c6df7..92e2495dcd26 100644 --- a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py +++ b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py @@ -206,13 +206,13 @@ class Flux3DistilledAutoCoreDenoiseStep(AutoPipelineBlocks): The width in pixels of the generated image; defaults to the first reference image's own width reference_crop (`list | NoneType`, *optional*): Per reference, the box of the reference inside its padded canvas or None, from the preprocess step - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): Torch generator for deterministic generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 16): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. (overrides `num_inference_steps` when provided) sigmas (`list`, *optional*): Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided) @@ -297,13 +297,13 @@ class Flux3DistilledAutoBlocks(SequentialPipelineBlocks): The width in pixels of the generated image; defaults to the first reference image's own width reference_crop (`list | NoneType`, *optional*): Per reference, the box of the reference inside its padded canvas or None, from the preprocess step - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): Torch generator for deterministic generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 16): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. (overrides `num_inference_steps` when provided) sigmas (`list`, *optional*): Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided) diff --git a/src/diffusers/modular_pipelines/helios/denoise.py b/src/diffusers/modular_pipelines/helios/denoise.py index 32a6f0b11b79..8d28c2ac0934 100644 --- a/src/diffusers/modular_pipelines/helios/denoise.py +++ b/src/diffusers/modular_pipelines/helios/denoise.py @@ -1061,6 +1061,15 @@ class HeliosPyramidDistilledChunkDenoiseStep(HeliosChunkLoopWrapper): HeliosChunkUpdateStep, ] block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"] + # the shared noise-gen step declares the non-distilled schedule; the distilled denoiser runs 2 steps per stage + input_overrides = { + "pyramid_num_inference_steps_list": InputParam( + "pyramid_num_inference_steps_list", + default=[2, 2, 2], + type_hint=list, + description="Number of denoising steps per pyramid stage.", + ) + } @property def description(self) -> str: @@ -1080,6 +1089,15 @@ class HeliosPyramidDistilledI2VChunkDenoiseStep(HeliosChunkLoopWrapper): HeliosChunkUpdateStep, ] block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"] + # the shared noise-gen step declares the non-distilled schedule; the distilled denoiser runs 2 steps per stage + input_overrides = { + "pyramid_num_inference_steps_list": InputParam( + "pyramid_num_inference_steps_list", + default=[2, 2, 2], + type_hint=list, + description="Number of denoising steps per pyramid stage.", + ) + } @property def description(self) -> str: diff --git a/src/diffusers/modular_pipelines/helios/modular_blocks_helios.py b/src/diffusers/modular_pipelines/helios/modular_blocks_helios.py index c3d5cb4efc77..eb0e74c6c2e0 100644 --- a/src/diffusers/modular_pipelines/helios/modular_blocks_helios.py +++ b/src/diffusers/modular_pipelines/helios/modular_blocks_helios.py @@ -269,7 +269,7 @@ class HeliosV2VCoreDenoiseStep(SequentialPipelineBlocks): text embeddings used to guide the image generation. Can be generated from text_encoder step. negative_prompt_embeds (`Tensor`, *optional*): negative text embeddings used to guide the image generation. Can be generated from text_encoder step. - image_latents (`Tensor`, *optional*): + image_latents (`Tensor`): image latents used to guide the image generation. Can be generated from vae_encoder step. video_latents (`Tensor`, *optional*): Encoded video latents for V2V generation. @@ -382,13 +382,13 @@ class HeliosAutoCoreDenoiseStep(ConditionalPipelineBlocks): Torch generator for deterministic generation. num_frames (`int`, *optional*, defaults to 132): Total number of video frames to generate. - history_sizes (`list`): + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): Sizes of long/mid/short history buffers for temporal context. keep_first_frame (`bool`, *optional*, defaults to True): Whether to keep the first frame as a prefix in history. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. - sigmas (`list`): + sigmas (`list`, *optional*): Custom sigmas for the denoising process. latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. @@ -496,13 +496,13 @@ class HeliosAutoBlocks(SequentialPipelineBlocks): Maximum sigma for video latent noise. num_frames (`int`, *optional*, defaults to 132): Total number of video frames to generate. - history_sizes (`list`): + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): Sizes of long/mid/short history buffers for temporal context. keep_first_frame (`bool`, *optional*, defaults to True): Whether to keep the first frame as a prefix in history. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. - sigmas (`list`): + sigmas (`list`, *optional*): Custom sigmas for the denoising process. latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. diff --git a/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py b/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py index fea11786de21..7ae868f7e8a7 100644 --- a/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py +++ b/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py @@ -260,7 +260,7 @@ class HeliosPyramidV2VCoreDenoiseStep(SequentialPipelineBlocks): text embeddings used to guide the image generation. Can be generated from text_encoder step. negative_prompt_embeds (`Tensor`, *optional*): negative text embeddings used to guide the image generation. Can be generated from text_encoder step. - image_latents (`Tensor`, *optional*): + image_latents (`Tensor`): image latents used to guide the image generation. Can be generated from vae_encoder step. video_latents (`Tensor`, *optional*): Encoded video latents for V2V generation. @@ -368,7 +368,7 @@ class HeliosPyramidAutoCoreDenoiseStep(ConditionalPipelineBlocks): Torch generator for deterministic generation. num_frames (`int`, *optional*, defaults to 132): Total number of video frames to generate. - history_sizes (`list`): + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): Sizes of long/mid/short history buffers for temporal context. keep_first_frame (`bool`, *optional*, defaults to True): Whether to keep the first frame as a prefix in history. @@ -478,7 +478,7 @@ class HeliosPyramidAutoBlocks(SequentialPipelineBlocks): Maximum sigma for video latent noise. num_frames (`int`, *optional*, defaults to 132): Total number of video frames to generate. - history_sizes (`list`): + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): Sizes of long/mid/short history buffers for temporal context. keep_first_frame (`bool`, *optional*, defaults to True): Whether to keep the first frame as a prefix in history. diff --git a/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py b/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py index 3e0b32f0df7e..3bba29ac4485 100644 --- a/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py +++ b/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py @@ -120,7 +120,7 @@ class HeliosPyramidDistilledCoreDenoiseStep(SequentialPipelineBlocks): Sizes of long/mid/short history buffers for temporal context. keep_first_frame (`bool`, *optional*, defaults to True): Whether to keep the first frame as a prefix in history. - pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [2, 2, 2]): Number of denoising steps per pyramid stage. generator (`Generator`, *optional*): Torch generator for deterministic generation. @@ -193,7 +193,7 @@ class HeliosPyramidDistilledI2VCoreDenoiseStep(SequentialPipelineBlocks): Sizes of long/mid/short history buffers for temporal context. keep_first_frame (`bool`, *optional*, defaults to True): Whether to keep the first frame as a prefix in history. - pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [2, 2, 2]): Number of denoising steps per pyramid stage. latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. @@ -261,7 +261,7 @@ class HeliosPyramidDistilledV2VCoreDenoiseStep(SequentialPipelineBlocks): text embeddings used to guide the image generation. Can be generated from text_encoder step. negative_prompt_embeds (`Tensor`, *optional*): negative text embeddings used to guide the image generation. Can be generated from text_encoder step. - image_latents (`Tensor`, *optional*): + image_latents (`Tensor`): image latents used to guide the image generation. Can be generated from vae_encoder step. video_latents (`Tensor`, *optional*): Encoded video latents for V2V generation. @@ -283,7 +283,7 @@ class HeliosPyramidDistilledV2VCoreDenoiseStep(SequentialPipelineBlocks): Sizes of long/mid/short history buffers for temporal context. keep_first_frame (`bool`, *optional*, defaults to True): Whether to keep the first frame as a prefix in history. - pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [2, 2, 2]): Number of denoising steps per pyramid stage. latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. @@ -370,11 +370,11 @@ class HeliosPyramidDistilledAutoCoreDenoiseStep(ConditionalPipelineBlocks): Torch generator for deterministic generation. num_frames (`int`, *optional*, defaults to 132): Total number of video frames to generate. - history_sizes (`list`): + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): Sizes of long/mid/short history buffers for temporal context. keep_first_frame (`bool`, *optional*, defaults to True): Whether to keep the first frame as a prefix in history. - pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [2, 2, 2]): Number of denoising steps per pyramid stage. latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. @@ -486,11 +486,11 @@ class HeliosPyramidDistilledAutoBlocks(SequentialPipelineBlocks): Maximum sigma for video latent noise. num_frames (`int`, *optional*, defaults to 132): Total number of video frames to generate. - history_sizes (`list`): + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): Sizes of long/mid/short history buffers for temporal context. keep_first_frame (`bool`, *optional*, defaults to True): Whether to keep the first frame as a prefix in history. - pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [2, 2, 2]): Number of denoising steps per pyramid stage. latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. diff --git a/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py b/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py index bdbdba1ecdd9..4164a8fed51f 100644 --- a/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py +++ b/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py @@ -332,10 +332,10 @@ class HunyuanVideo15AutoCoreDenoiseStep(AutoPipelineBlocks): Inputs: prompt_embeds (`Tensor`): text embeddings used to guide the image generation. Can be generated from text_encoder step. - batch_size (`int`): + batch_size (`int`, *optional*): Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. @@ -345,7 +345,7 @@ class HunyuanVideo15AutoCoreDenoiseStep(AutoPipelineBlocks): The width in pixels of the generated image. num_frames (`int`, *optional*, defaults to 121): Number of video frames to generate. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. @@ -418,16 +418,16 @@ class HunyuanVideo15AutoBlocks(SequentialPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*): The width in pixels of the generated image. - batch_size (`int`): + batch_size (`int`, *optional*): Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. num_frames (`int`, *optional*, defaults to 121): Number of video frames to generate. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): Torch generator for deterministic generation. @@ -504,7 +504,7 @@ class HunyuanVideo15Image2VideoBlocks(SequentialPipelineBlocks): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): Torch generator for deterministic generation. - image_latents (`Tensor`): + image_latents (`Tensor`, *optional*): Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V. attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. diff --git a/src/diffusers/modular_pipelines/ideogram4/denoise.py b/src/diffusers/modular_pipelines/ideogram4/denoise.py index db1a708c4315..a173c7af8555 100644 --- a/src/diffusers/modular_pipelines/ideogram4/denoise.py +++ b/src/diffusers/modular_pipelines/ideogram4/denoise.py @@ -236,8 +236,6 @@ class Ideogram4DenoiseStep(LoopSequentialPipelineBlocks): Effective batch size. prompt_embeds (`Tensor`): Packed conditional encoder_hidden_states. - position_ids (`Tensor`): - Conditional 3-axis MRoPE position ids. segment_ids (`Tensor`): Conditional block-diagonal segment ids. indicator (`Tensor`): diff --git a/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py b/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py index 0b788fe236be..80a123296495 100644 --- a/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py +++ b/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py @@ -128,9 +128,9 @@ class Ideogram4AutoBlocks(SequentialPipelineBlocks): If True, rewrite the prompt into Ideogram4's native JSON caption before encoding. prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0): Sampling temperature for prompt upsampling. - height (`int`, *optional*): + height (`int`): The height in pixels of the generated image. - width (`int`, *optional*): + width (`int`): The width in pixels of the generated image. max_sequence_length (`int`, *optional*, defaults to 2048): Maximum sequence length for prompt encoding. diff --git a/src/diffusers/modular_pipelines/krea2/before_denoise.py b/src/diffusers/modular_pipelines/krea2/before_denoise.py index 17ed6a3cd376..3e632b981c86 100644 --- a/src/diffusers/modular_pipelines/krea2/before_denoise.py +++ b/src/diffusers/modular_pipelines/krea2/before_denoise.py @@ -386,6 +386,8 @@ class Krea2SetTimestepsStep(ModularPipelineBlocks): Outputs: timesteps (`Tensor`): The denoising timesteps. + num_inference_steps (`int`): + The number of denoising steps. """ model_name = "krea2" @@ -418,7 +420,10 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: - return [OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps.")] + return [ + OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps."), + OutputParam(name="num_inference_steps", type_hint=int, description="The number of denoising steps."), + ] @torch.no_grad() def __call__( @@ -470,6 +475,8 @@ class Krea2TurboSetTimestepsStep(ModularPipelineBlocks): Outputs: timesteps (`Tensor`): The denoising timesteps. + num_inference_steps (`int`): + The number of denoising steps. """ model_name = "krea2" @@ -496,7 +503,10 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: - return [OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps.")] + return [ + OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps."), + OutputParam(name="num_inference_steps", type_hint=int, description="The number of denoising steps."), + ] @torch.no_grad() def __call__( diff --git a/src/diffusers/modular_pipelines/ltx/before_denoise.py b/src/diffusers/modular_pipelines/ltx/before_denoise.py index 470ea7325ec7..296895f83d7a 100644 --- a/src/diffusers/modular_pipelines/ltx/before_denoise.py +++ b/src/diffusers/modular_pipelines/ltx/before_denoise.py @@ -128,6 +128,10 @@ def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("batch_size", type_hint=int), OutputParam("dtype", type_hint=torch.dtype), + OutputParam.template("prompt_embeds"), + OutputParam.template("prompt_embeds_mask", name="prompt_attention_mask"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("negative_prompt_embeds_mask", name="negative_prompt_attention_mask"), ] @torch.no_grad() diff --git a/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py b/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py index 0bdaf6ec8816..70cb18bdad29 100644 --- a/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py +++ b/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py @@ -289,25 +289,25 @@ class LTXAutoCoreDenoiseStep(AutoPipelineBlocks): text embeddings used to guide the image generation. Can be generated from text_encoder step. prompt_attention_mask (`Tensor`): mask for the text embeddings. Can be generated from text_encoder step. - negative_prompt_embeds (`Tensor`): + negative_prompt_embeds (`Tensor`, *optional*): negative text embeddings used to guide the image generation. Can be generated from text_encoder step. - negative_prompt_attention_mask (`Tensor`): + negative_prompt_attention_mask (`Tensor`, *optional*): mask for the negative text embeddings. Can be generated from text_encoder step. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - height (`int`): + height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. - width (`int`): + width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`): + num_frames (`int`, *optional*, defaults to 161): TODO: Add description. frame_rate (`int`, *optional*, defaults to 25): TODO: Add description. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): Torch generator for deterministic generation. @@ -366,17 +366,17 @@ class LTXAutoBlocks(SequentialPipelineBlocks): Torch generator for deterministic generation. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - num_frames (`int`): + num_frames (`int`, *optional*, defaults to 161): TODO: Add description. frame_rate (`int`, *optional*, defaults to 25): TODO: Add description. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. image_latents (`Tensor`, *optional*): TODO: Add description. @@ -455,7 +455,7 @@ class LTXImage2VideoBlocks(SequentialPipelineBlocks): TODO: Add description. latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. - image_latents (`Tensor`): + image_latents (`Tensor`, *optional*): TODO: Add description. attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py index aa0990f3cc80..16b176ce5c88 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py @@ -971,7 +971,7 @@ class LTX2AutoCoreDenoiseStep(ConditionalPipelineBlocks): RoPE coordinates for the reference tokens. reference_token_counts (`list`, *optional*): Per-reference token counts, in `reference_conditions` order. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. @@ -997,11 +997,11 @@ class LTX2AutoCoreDenoiseStep(ConditionalPipelineBlocks): Torch generator for deterministic generation. reference_cross_mask (`Tensor`, *optional*): Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. - audio_latents (`Tensor`): + audio_latents (`Tensor`, *optional*): Optional pre-encoded audio latents; random noise is used when not provided. dtype (`dtype`): The dtype the model inputs are cast to. @@ -1436,7 +1436,7 @@ class LTX2ImageToVideoBlocks(SequentialPipelineBlocks): noise_scale (`float`, *optional*): Interpolation factor between random noise and any provided latents. `None` (default) resolves to 0.0, which keeps the provided latents. - image_latents (`Tensor`): + image_latents (`Tensor`, *optional*): VAE-encoded reference-image latents used for image-to-video conditioning. audio_latents (`Tensor`, *optional*): Optional pre-encoded audio latents; random noise is used when not provided. @@ -1855,7 +1855,7 @@ class LTX2AutoBlocks(SequentialPipelineBlocks): RoPE coordinates for the reference tokens. reference_token_counts (`list`, *optional*): Per-reference token counts, in `reference_conditions` order. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. noise_scale (`float`, *optional*): Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom @@ -1864,11 +1864,11 @@ class LTX2AutoBlocks(SequentialPipelineBlocks): Custom sigmas for the denoising process. reference_cross_mask (`Tensor`, *optional*): Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. - audio_latents (`Tensor`): + audio_latents (`Tensor`, *optional*): Optional pre-encoded audio latents; random noise is used when not provided. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py index 3280820eab0c..3b89ddc3bbb7 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py @@ -311,7 +311,7 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): RoPE coordinates for the reference tokens. reference_token_counts (`list`, *optional*): Per-reference token counts, in `reference_conditions` order. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. noise_scale (`float`, *optional*): Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom @@ -320,11 +320,11 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): Custom sigmas for the denoising process. reference_cross_mask (`Tensor`, *optional*): Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. - timesteps (`Tensor`): + timesteps (`Tensor`, *optional*): Timesteps for the denoising process. - audio_latents (`Tensor`): + audio_latents (`Tensor`, *optional*): Optional pre-encoded audio latents; random noise is used when not provided. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. diff --git a/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py b/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py index e96454a44fa4..318f55dd1506 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py +++ b/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py @@ -559,10 +559,10 @@ class MiniMaxH3AutoDenoiseStep(ConditionalPipelineBlocks): generator (`Generator`, *optional*): The generator of the request. The conditioning noise is drawn from it first, one draw per condition, before the noise of the generated rows. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used instead of the draw. - audio_latents (`Tensor`): + audio_latents (`Tensor`, *optional*): Pre-generated audio noise of shape `(2, 32, num_audio_latents)`. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. @@ -692,7 +692,7 @@ class MiniMaxH3Blocks(SequentialPipelineBlocks): Height of the generated video in pixels, a multiple of 32. width (`int`, *optional*): Width of the generated video in pixels, a multiple of 32. - num_frames (`int`, *optional*): + num_frames (`int`, *optional*, defaults to None or 124, depending on the workflow): Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video VAE can decode; the resulting duration must stay between 5 and 15 seconds. To generate a video as long as a reference soundtrack, pass `round(samples / sample_rate * 24)`. @@ -716,10 +716,10 @@ class MiniMaxH3Blocks(SequentialPipelineBlocks): generator (`Generator`, *optional*): The generator of the request. The conditioning noise is drawn from it first, one draw per condition, before the noise of the generated rows. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used instead of the draw. - audio_latents (`Tensor`): + audio_latents (`Tensor`, *optional*): Pre-generated audio noise of shape `(2, 32, num_audio_latents)`. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index ba9740e23258..e49e371041d2 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -20,7 +20,7 @@ import warnings from collections import OrderedDict from copy import deepcopy -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any import torch @@ -66,6 +66,7 @@ format_workflow, generate_modular_model_card_content, make_doc_string, + override_attribute, ) @@ -690,12 +691,31 @@ class ConditionalPipelineBlocks(ModularPipelineBlocks): element specifies the trigger input for the corresponding block. default_block_name: Name of the default block to run when no trigger inputs match. If None, this block can be skipped entirely when no trigger inputs are provided. + input_overrides: the composite block's own declarations for inputs it derives from its sub-blocks, as + `{name: InputParam}`, e.g. `{"strength": InputParam.template("strength", default=0.9999)}`. The name must + be a derived input, and the declaration replaces the derived one whole, so it states everything: default, + `required`, description. When the sub-blocks declare different defaults for the same input, each one + keeps its own at runtime (`defaults_by_block`); a declaration here with a default replaces them all with + one value. + intermediate_output_overrides: the composite block's own declarations for intermediate outputs it derives + from its sub-blocks, as `{name: OutputParam}`, e.g. + `{"latents": OutputParam("latents", description="Denoised latents.")}`. The name must be a derived output. + component_overrides: the composite block's own specs for components it derives from its sub-blocks, as + `{name: ComponentSpec}`, e.g. `{"text_encoder": ComponentSpec("text_encoder", T5EncoderModel)}`. The name + must be a derived component. + config_overrides: the composite block's own specs for configs it derives from its sub-blocks, as + `{name: ConfigSpec}`, e.g. `{"vae_scale_factor": ConfigSpec("vae_scale_factor", default=16)}`. The name + must be a derived config. """ block_classes = [] block_names = [] block_trigger_inputs = [] default_block_name = None + input_overrides: dict[str, InputParam] = {} + intermediate_output_overrides: dict[str, OutputParam] = {} + component_overrides: dict[str, ComponentSpec] = {} + config_overrides: dict[str, ConfigSpec] = {} def __init__(self): sub_blocks = InsertableDict() @@ -729,7 +749,9 @@ def expected_components(self): for component in block.expected_components: if component not in expected_components: expected_components.append(component) - return expected_components + return override_attribute( + expected_components, self.component_overrides, "component_overrides", self.__class__.__name__ + ) @property def expected_configs(self): @@ -738,7 +760,7 @@ def expected_configs(self): for config in block.expected_configs: if config not in expected_configs: expected_configs.append(config) - return expected_configs + return override_attribute(expected_configs, self.config_overrides, "config_overrides", self.__class__.__name__) @property def required_inputs(self) -> list[str]: @@ -760,19 +782,25 @@ def required_inputs(self) -> list[str]: def inputs(self) -> list[tuple[str, Any]]: named_inputs = [(name, block.inputs) for name, block in self.sub_blocks.items()] combined_inputs = combine_inputs(*named_inputs) - # mark Required inputs only if that input is required by all the blocks - for input_param in combined_inputs: + # required only if every sub-block requires it; + for position, input_param in enumerate(combined_inputs): if input_param.name in self.required_inputs: - input_param.required = True + # `replace` copies so the sub-blocks' declarations stay untouched + combined_inputs[position] = replace(input_param, required=True) else: - input_param.required = False - return combined_inputs + combined_inputs[position] = replace(input_param, required=False) + return override_attribute(combined_inputs, self.input_overrides, "input_overrides", self.__class__.__name__) @property def intermediate_outputs(self) -> list[str]: named_outputs = [(name, block.intermediate_outputs) for name, block in self.sub_blocks.items()] combined_outputs = combine_outputs(*named_outputs) - return combined_outputs + return override_attribute( + combined_outputs, + self.intermediate_output_overrides, + "intermediate_output_overrides", + self.__class__.__name__, + ) @property def outputs(self) -> list[str]: @@ -1043,10 +1071,30 @@ class SequentialPipelineBlocks(ModularPipelineBlocks): Attributes: block_classes: list of block classes to be used block_names: list of prefixes for each block + input_overrides: the composite block's own declarations for inputs it derives from its sub-blocks, as + `{name: InputParam}`, e.g. `{"strength": InputParam.template("strength", default=0.9999)}`. The name must + be a derived input, and the declaration replaces the derived one whole, so it states everything: default, + `required`, description. When the sub-blocks declare different defaults for the same input, each one + keeps its own at runtime (`defaults_by_block`); a declaration here with a default replaces them all with + one value. Can also be passed to `from_blocks_dict`. + intermediate_output_overrides: the composite block's own declarations for intermediate outputs it derives + from its sub-blocks, as `{name: OutputParam}`, e.g. + `{"latents": OutputParam("latents", description="Denoised latents.")}`. The name must be a derived output. + Can also be passed to `from_blocks_dict`. + component_overrides: the composite block's own specs for components it derives from its sub-blocks, as + `{name: ComponentSpec}`, e.g. `{"text_encoder": ComponentSpec("text_encoder", T5EncoderModel)}`. The name + must be a derived component. Can also be passed to `from_blocks_dict`. + config_overrides: the composite block's own specs for configs it derives from its sub-blocks, as + `{name: ConfigSpec}`, e.g. `{"vae_scale_factor": ConfigSpec("vae_scale_factor", default=16)}`. The name + must be a derived config. Can also be passed to `from_blocks_dict`. """ block_classes = [] block_names = [] + input_overrides: dict[str, InputParam] = {} + intermediate_output_overrides: dict[str, OutputParam] = {} + component_overrides: dict[str, ComponentSpec] = {} + config_overrides: dict[str, ConfigSpec] = {} @property def description(self): @@ -1063,7 +1111,9 @@ def expected_components(self): for component in block.expected_components: if component not in expected_components: expected_components.append(component) - return expected_components + return override_attribute( + expected_components, self.component_overrides, "component_overrides", self.__class__.__name__ + ) @property def expected_configs(self): @@ -1072,7 +1122,7 @@ def expected_configs(self): for config in block.expected_configs: if config not in expected_configs: expected_configs.append(config) - return expected_configs + return override_attribute(expected_configs, self.config_overrides, "config_overrides", self.__class__.__name__) @property def available_workflows(self): @@ -1106,12 +1156,25 @@ def get_workflow(self, workflow_name: str): @classmethod def from_blocks_dict( - cls, blocks_dict: dict[str, Any], description: str | None = None + cls, + blocks_dict: dict[str, Any], + description: str | None = None, + input_overrides: dict[str, InputParam] | None = None, + intermediate_output_overrides: dict[str, OutputParam] | None = None, + component_overrides: dict[str, ComponentSpec] | None = None, + config_overrides: dict[str, ConfigSpec] | None = None, ) -> "SequentialPipelineBlocks": """Creates a SequentialPipelineBlocks instance from a dictionary of blocks. Args: blocks_dict: Dictionary mapping block names to block classes or instances + description: Description of the new block + input_overrides: the composite's own declarations replacing derived `inputs`, as `{name: InputParam}` + intermediate_output_overrides: the composite's own declarations replacing derived + `intermediate_outputs`, as `{name: OutputParam}` + component_overrides: the composite's own specs replacing derived `expected_components`, as + `{name: ComponentSpec}` + config_overrides: the composite's own specs replacing derived `expected_configs`, as `{name: ConfigSpec}` Returns: A new SequentialPipelineBlocks instance @@ -1132,6 +1195,14 @@ def from_blocks_dict( if description is not None: instance.description = description + for name, value in ( + ("input_overrides", input_overrides), + ("intermediate_output_overrides", intermediate_output_overrides), + ("component_overrides", component_overrides), + ("config_overrides", config_overrides), + ): + if value is not None: + setattr(instance, name, value) return instance @@ -1149,63 +1220,124 @@ def __init__(self): ) def _get_inputs(self): - inputs = [] - outputs = set() - - # Go through all blocks in order - for block in self.sub_blocks.values(): - # Add inputs that aren't in outputs yet + # Walk the sub-blocks in order to derive the pipeline inputs. An input is listed unless an earlier block + # produces it. The first declaration is kept; a later declaration of the same input can only make it required + # or disagree on its default. + # + # The comments below follow one img2img pipeline: + # `text_encoder -> image_encoder -> set_timesteps -> prepare_latents -> denoise`, where `image_encoder` is a + # ConditionalPipelineBlocks without a default block: it runs when `image` is passed and produces + # `image_latents`. + inputs = {} # pipeline inputs + outputs = set() # produced by an earlier block + maybe_outputs = set() # produced by an earlier ConditionalPipelineBlocks, if it runs + for block_name, block in self.sub_blocks.items(): for inp in block.inputs: - if inp.name not in outputs and inp.name not in {input.name for input in inputs}: - inputs.append(inp) - - # Only add outputs if the block cannot be skipped - should_add_outputs = True + # kwargs inputs have no name, key them by "*_" like `combine_inputs` does + key = inp.name if inp.name is not None else f"*_{inp.kwargs_type}" + if key in outputs: + # produced by an earlier block, so not a pipeline input: + # e.g. `denoise` needs `prompt_embeds`, but `text_encoder` produces it + continue + # first declaration, listed as declared: + # e.g. `set_timesteps` declares `num_inference_steps` and no earlier block declared or produced it, so + # it becomes a pipeline input with the default `set_timesteps` gave it + if key not in inputs: + # `replace` copies: the flags below are the composite's view and must not reach the sub-block + inputs[key] = replace(inp) + if inp.required and key in maybe_outputs: + # the block requires it, but an earlier conditional block can produce it, so the pipeline lists + # it optional: + # e.g. `prepare_latents` requires `image_latents`, but `image_encoder` produces it when `image` + # is passed, so `image_latents` is optional for the pipeline: the user can pass `image` instead + inputs[key].required = False + continue + if inp.required: + if inputs[key].default is None and not inputs[key].defaults_by_block and key not in maybe_outputs: + # already listed as optional with no default, nothing earlier can produce it, and now a block + # requires it, so the user has to pass it: + # e.g. `prepare_latents` declares `height` optional, `denoise` requires it, so `height` is a + # required pipeline input. Had `prepare_latents` given `height` a default, or had the branches + # of an earlier conditional block each given it theirs, it would stay optional + inputs[key].required = True + elif not inputs[key].required: + if inputs[key].default is None and not inputs[key].defaults_by_block and inp.defaults_by_block: + # the entry has no default and this conditional block's branches each have their own: carry + # them up, prefixed like `combine_inputs` does, so the pipeline docstring shows them. Not a + # conflict: the branch that runs fills in its own default when it reads the state + # (`get_block_state`): + # e.g. `before_encode` lists `num_frames` with no default and the `denoise` branches default + # it to 124 or require it, so the pipeline says "None or 124, depending on the workflow" + inputs[key].defaults_by_block = { + f"{block_name}.{k}": v for k, v in inp.defaults_by_block.items() + } + elif key not in self.input_overrides: + # the defaults each side has an opinion on: its own, or the per-branch ones a conditional block + # recorded. `None` means no default, not an opinion, so it is left out on both sides + existing, declared = inputs[key], inp + existing_defaults = ( + [existing.default] + if existing.default is not None + else [d for d in (existing.defaults_by_block or {}).values() if d is not None] + ) + declared_defaults = ( + [declared.default] + if declared.default is not None + else [d for d in (declared.defaults_by_block or {}).values() if d is not None] + ) + # (`in` compares with `==`, so list defaults such as `[2, 2, 2]` work too) + if (existing_defaults and any(d not in existing_defaults for d in declared_defaults)) or ( + declared_defaults and any(d not in declared_defaults for d in existing_defaults) + ): + # two optional declarations disagree on the default. In a sequence the input is one shared + # value, so the composite has to pick it with `input_overrides`; until 1.0.0 the first + # declaration wins: + # e.g. `strength` is 0.3 on `set_timesteps` and 0.9999 on `prepare_latents`, so the + # pipeline uses 0.3 and warns until it declares `strength` in `input_overrides` + existing_text = f"default {existing.default!r}" + if existing.defaults_by_block: + existing_text = f"per-branch defaults {existing.defaults_by_block!r}" + declared_text = f"default {declared.default!r}" + if declared.defaults_by_block: + declared_text = f"per-branch defaults {declared.defaults_by_block!r}" + deprecate( + "sub-block default conflict", + "1.0.0", + f"{self.__class__.__name__}: sub-block '{block_name}' declares {declared_text} for " + f"input '{key}' but an earlier sub-block declares {existing_text}. The first " + f"declaration currently wins; this will become an error. Declare the input in " + f"`input_overrides` on {self.__class__.__name__} to pick the default explicitly.", + standard_warn=False, + ) if isinstance(block, ConditionalPipelineBlocks) and block.default_block_name is None: - # ConditionalPipelineBlocks without default can be skipped - should_add_outputs = False - - if should_add_outputs: - # Add this block's outputs - block_intermediate_outputs = [out.name for out in block.intermediate_outputs] - outputs.update(block_intermediate_outputs) + # e.g. `image_encoder` may be skipped, so `image_latents` is only produced if it runs: a "maybe" output + maybe_outputs.update(out.name for out in block.intermediate_outputs) + else: + outputs.update(out.name for out in block.intermediate_outputs) - return inputs + return override_attribute( + list(inputs.values()), self.input_overrides, "input_overrides", self.__class__.__name__ + ) - # YiYi TODO: add test for this @property def inputs(self) -> list[tuple[str, Any]]: return self._get_inputs() - @property - def required_inputs(self) -> list[str]: - # Get the first block from the dictionary - first_block = next(iter(self.sub_blocks.values())) - required_by_any = set(getattr(first_block, "required_inputs", set())) - - # Union with required inputs from all other blocks - for block in list(self.sub_blocks.values())[1:]: - block_required = set(getattr(block, "required_inputs", set())) - required_by_any.update(block_required) - - return list(required_by_any) - @property def intermediate_outputs(self) -> list[str]: - named_outputs = [] - for name, block in self.sub_blocks.items(): - inp_names = {inp.name for inp in block.inputs} - # so we only need to list new variables as intermediate_outputs, but if user wants to list these they modified it's still fine (a.k.a we don't enforce) - # filter out them here so they do not end up as intermediate_outputs - if name not in inp_names: - named_outputs.append((name, block.intermediate_outputs)) + named_outputs = [(name, block.intermediate_outputs) for name, block in self.sub_blocks.items()] combined_outputs = combine_outputs(*named_outputs) - return combined_outputs + return override_attribute( + combined_outputs, + self.intermediate_output_overrides, + "intermediate_output_overrides", + self.__class__.__name__, + ) - # YiYi TODO: I think we can remove the outputs property + # `outputs` is what `doc` and the model card list as the block's outputs. By default that is every intermediate + # output; a composite overrides it to show only what the step is for, e.g. just `latents` for the `denoise` block. @property def outputs(self) -> list[str]: - # return next(reversed(self.sub_blocks.values())).intermediate_outputs return self.intermediate_outputs @torch.no_grad() @@ -1403,11 +1535,31 @@ class LoopSequentialPipelineBlocks(ModularPipelineBlocks): Attributes: block_classes: list of block classes to be used block_names: list of prefixes for each block + input_overrides: the composite block's own declarations for inputs it derives from its sub-blocks, as + `{name: InputParam}`, e.g. `{"strength": InputParam.template("strength", default=0.9999)}`. The name must + be a derived input, and the declaration replaces the derived one whole, so it states everything: default, + `required`, description. When the sub-blocks declare different defaults for the same input, each one + keeps its own at runtime (`defaults_by_block`); a declaration here with a default replaces them all with + one value. Can also be passed to `from_blocks_dict`. + intermediate_output_overrides: the composite block's own declarations for intermediate outputs it derives + from its sub-blocks, as `{name: OutputParam}`, e.g. + `{"latents": OutputParam("latents", description="Denoised latents.")}`. The name must be a derived output. + Can also be passed to `from_blocks_dict`. + component_overrides: the composite block's own specs for components it derives from its sub-blocks, as + `{name: ComponentSpec}`, e.g. `{"text_encoder": ComponentSpec("text_encoder", T5EncoderModel)}`. The name + must be a derived component. Can also be passed to `from_blocks_dict`. + config_overrides: the composite block's own specs for configs it derives from its sub-blocks, as + `{name: ConfigSpec}`, e.g. `{"vae_scale_factor": ConfigSpec("vae_scale_factor", default=16)}`. The name + must be a derived config. Can also be passed to `from_blocks_dict`. """ model_name = None block_classes = [] block_names = [] + input_overrides: dict[str, InputParam] = {} + intermediate_output_overrides: dict[str, OutputParam] = {} + component_overrides: dict[str, ComponentSpec] = {} + config_overrides: dict[str, ConfigSpec] = {} @property def description(self) -> str: @@ -1427,14 +1579,6 @@ def loop_inputs(self) -> list[InputParam]: """list of input parameters. Must be implemented by subclasses.""" return [] - @property - def loop_required_inputs(self) -> list[str]: - input_names = [] - for input_param in self.loop_inputs: - if input_param.required: - input_names.append(input_param.name) - return input_names - @property def loop_intermediate_outputs(self) -> list[OutputParam]: """list of intermediate output parameters. Must be implemented by subclasses.""" @@ -1451,7 +1595,9 @@ def expected_components(self): for component in self.loop_expected_components: if component not in expected_components: expected_components.append(component) - return expected_components + return override_attribute( + expected_components, self.component_overrides, "component_overrides", self.__class__.__name__ + ) # modified from SequentialPipelineBlocks to include loop_expected_configs @property @@ -1464,53 +1610,62 @@ def expected_configs(self): for config in self.loop_expected_configs: if config not in expected_configs: expected_configs.append(config) - return expected_configs + return override_attribute(expected_configs, self.config_overrides, "config_overrides", self.__class__.__name__) def _get_inputs(self): - inputs = [] - inputs.extend(self.loop_inputs) - outputs = set() - - for name, block in self.sub_blocks.items(): - # Add inputs that aren't in outputs yet + # Same walk as `SequentialPipelineBlocks._get_inputs`, starting from `loop_inputs`. Loop sub-blocks are leaf + # blocks, so there is no skippable ConditionalPipelineBlocks whose outputs may or may not exist. + inputs = {} + outputs = set() # produced by an earlier block + for inp in self.loop_inputs: + # kwargs inputs have no name, key them by "*_" like `combine_inputs` does + key = inp.name if inp.name is not None else f"*_{inp.kwargs_type}" + inputs[key] = replace(inp) + for block_name, block in self.sub_blocks.items(): for inp in block.inputs: - if inp.name not in outputs and inp not in inputs: - inputs.append(inp) - - # Add this block's outputs - block_intermediate_outputs = [out.name for out in block.intermediate_outputs] - outputs.update(block_intermediate_outputs) - - for input_param in inputs: - if input_param.name in self.required_inputs: - input_param.required = True - else: - input_param.required = False + key = inp.name if inp.name is not None else f"*_{inp.kwargs_type}" + if key in outputs: + continue + if key not in inputs: + # `replace` copies: the flags below are the composite's view and must not reach the sub-block + inputs[key] = replace(inp) + continue + if inp.required: + if inputs[key].default is None: + # nothing earlier defaults or produces it: the caller has to pass it + inputs[key].required = True + elif ( + inp.default is not None + and inputs[key].default is not None + and not inputs[key].required + and inp.default != inputs[key].default + and key not in self.input_overrides + ): + # two optional declarations disagree on the default. In a sequence the input is one shared value, + # so the composite has to pick it with `input_overrides`; until 1.0.0 the first declaration wins. + # `None` means no default, not an opinion. Loop sub-blocks are leaves, so there are no per-branch + # defaults to consider here + deprecate( + "sub-block default conflict", + "1.0.0", + f"{self.__class__.__name__}: sub-block '{block_name}' declares default {inp.default!r} for " + f"input '{key}' but an earlier declaration (a loop input or an earlier sub-block) has " + f"{inputs[key].default!r}. The first declaration currently wins; this will become an error. " + f"Declare the input in `input_overrides` on {self.__class__.__name__} to pick the default " + f"explicitly.", + standard_warn=False, + ) + outputs.update(out.name for out in block.intermediate_outputs) - return inputs + return override_attribute( + list(inputs.values()), self.input_overrides, "input_overrides", self.__class__.__name__ + ) @property # Copied from diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks.inputs def inputs(self): return self._get_inputs() - # modified from SequentialPipelineBlocks, if any additionan input required by the loop is required by the block - @property - def required_inputs(self) -> list[str]: - # Get the first block from the dictionary - first_block = next(iter(self.sub_blocks.values())) - required_by_any = set(getattr(first_block, "required_inputs", set())) - - required_by_loop = set(getattr(self, "loop_required_inputs", set())) - required_by_any.update(required_by_loop) - - # Union with required inputs from all other blocks - for block in list(self.sub_blocks.values())[1:]: - block_required = set(getattr(block, "required_inputs", set())) - required_by_any.update(block_required) - - return list(required_by_any) - # YiYi TODO: this need to be thought about more # modified from SequentialPipelineBlocks to include loop_intermediate_outputs @property @@ -1520,7 +1675,12 @@ def intermediate_outputs(self) -> list[str]: for output in self.loop_intermediate_outputs: if output.name not in {output.name for output in combined_outputs}: combined_outputs.append(output) - return combined_outputs + return override_attribute( + combined_outputs, + self.intermediate_output_overrides, + "intermediate_output_overrides", + self.__class__.__name__, + ) # YiYi TODO: this need to be thought about more @property @@ -1554,12 +1714,25 @@ def __init__(self): ) @classmethod - def from_blocks_dict(cls, blocks_dict: dict[str, Any]) -> "LoopSequentialPipelineBlocks": + def from_blocks_dict( + cls, + blocks_dict: dict[str, Any], + input_overrides: dict[str, InputParam] | None = None, + intermediate_output_overrides: dict[str, OutputParam] | None = None, + component_overrides: dict[str, ComponentSpec] | None = None, + config_overrides: dict[str, ConfigSpec] | None = None, + ) -> "LoopSequentialPipelineBlocks": """ Creates a LoopSequentialPipelineBlocks instance from a dictionary of blocks. Args: blocks_dict: Dictionary mapping block names to block instances + input_overrides: the composite's own declarations replacing derived `inputs`, as `{name: InputParam}` + intermediate_output_overrides: the composite's own declarations replacing derived + `intermediate_outputs`, as `{name: OutputParam}` + component_overrides: the composite's own specs replacing derived `expected_components`, as + `{name: ComponentSpec}` + config_overrides: the composite's own specs replacing derived `expected_configs`, as `{name: ConfigSpec}` Returns: A new LoopSequentialPipelineBlocks instance @@ -1577,6 +1750,14 @@ def from_blocks_dict(cls, blocks_dict: dict[str, Any]) -> "LoopSequentialPipelin instance.block_classes = [block.__class__ for block in blocks_dict.values()] instance.block_names = list(blocks_dict.keys()) instance.sub_blocks = blocks_dict + for name, value in ( + ("input_overrides", input_overrides), + ("intermediate_output_overrides", intermediate_output_overrides), + ("component_overrides", component_overrides), + ("config_overrides", config_overrides), + ): + if value is not None: + setattr(instance, name, value) return instance def loop_step(self, components, state: PipelineState, **kwargs): diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py index 733f5bdc8a73..d4301042e54d 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py +++ b/src/diffusers/modular_pipelines/modular_pipeline_utils.py @@ -1174,6 +1174,53 @@ def combine_outputs(*named_output_lists: list[tuple[str, list[OutputParam]]]) -> return list(combined_dict.values()) +def override_attribute(items: list, overrides: dict[str, Any], attr_name: str, block_name: str) -> list: + """ + Apply a composite block's `*_overrides` to the list it derived from its sub-blocks. + + Args: + items: the derived `InputParam`, `OutputParam`, `ComponentSpec` or `ConfigSpec` list + overrides: `{name: declaration}`; every name must be in `items`, and the declaration, an object of the same + class as the entry it replaces, takes that entry's place whole + attr_name: the overrides attribute being applied, e.g. `"input_overrides"`, for error messages + block_name: the composite's class name, for error messages + + Returns: + `items` with each overridden entry replaced by the composite's own declaration; the sub-blocks' declarations + are not touched. + + Example: + ```python + inputs = [InputParam("strength", default=0.3), InputParam("prompt", required=True)] + override_attribute( + inputs, {"strength": InputParam.template("strength", default=0.9999)}, "input_overrides", "InpaintStep" + ) + # [, ] + ``` + """ + if not overrides: + return items + # kwargs inputs/outputs have no name and are keyed by "*_", like `combine_inputs` does + index = { + item.name if item.name is not None else f"*_{item.kwargs_type}": position + for position, item in enumerate(items) + } + for name, declaration in overrides.items(): + if name not in index: + raise ValueError( + f"{block_name}.{attr_name} has an entry for '{name}', which this block does not derive from its " + f"sub-blocks. Available: {sorted(k for k in index if not k.startswith('*_'))}" + ) + expected = type(items[index[name]]) + if not isinstance(declaration, expected): + raise TypeError( + f"{block_name}.{attr_name}['{name}'] must be of type {expected.__name__}, replacing the derived one; " + f"got {type(declaration).__name__}." + ) + items[index[name]] = declaration + return items + + def generate_modular_model_card_content(blocks) -> dict[str, Any]: """ Generate model card content for a modular pipeline. diff --git a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py index e4720b16244d..f024db511aa2 100644 --- a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py +++ b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py @@ -574,6 +574,8 @@ class QwenImageSetTimestepsStep(ModularPipelineBlocks): Outputs: timesteps (`Tensor`): The timesteps to use for the denoising process + num_inference_steps (`int`): + The number of denoising steps """ model_name = "qwenimage" @@ -607,6 +609,7 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process" ), + OutputParam(name="num_inference_steps", type_hint=int, description="The number of denoising steps"), ] def __call__( @@ -662,6 +665,8 @@ class QwenImageLayeredSetTimestepsStep(ModularPipelineBlocks): Outputs: timesteps (`Tensor`): The timesteps to use for the denoising process. + num_inference_steps (`int`): + The number of denoising steps. """ model_name = "qwenimage-layered" @@ -690,6 +695,7 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process." ), + OutputParam(name="num_inference_steps", type_hint=int, description="The number of denoising steps."), ] @torch.no_grad() diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py index bf87028b2f90..cd66eefa7201 100644 --- a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py +++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py @@ -340,7 +340,7 @@ class QwenImageInpaintInputStep(SequentialPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*): The width in pixels of the generated image. - image_latents (`Tensor`, *optional*): + image_latents (`Tensor`): image latents used to guide the image generation. Can be generated from vae_encoder step. processed_mask_image (`Tensor`, *optional*): The processed mask image @@ -543,7 +543,7 @@ class QwenImageInpaintCoreDenoiseStep(SequentialPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*): The width in pixels of the generated image. - image_latents (`Tensor`, *optional*): + image_latents (`Tensor`): image latents used to guide the image generation. Can be generated from vae_encoder step. processed_mask_image (`Tensor`, *optional*): The processed mask image @@ -789,7 +789,7 @@ class QwenImageControlNetInpaintCoreDenoiseStep(SequentialPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*): The width in pixels of the generated image. - image_latents (`Tensor`, *optional*): + image_latents (`Tensor`): image latents used to guide the image generation. Can be generated from vae_encoder step. processed_mask_image (`Tensor`, *optional*): The processed mask image @@ -1152,17 +1152,17 @@ class QwenImageAutoBlocks(SequentialPipelineBlocks): Control image for ControlNet conditioning. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - prompt_embeds (`Tensor`): + prompt_embeds (`Tensor`, *optional*): text embeddings used to guide the image generation. Can be generated from text_encoder step. - prompt_embeds_mask (`Tensor`): + prompt_embeds_mask (`Tensor`, *optional*): mask for the text embeddings. Can be generated from text_encoder step. negative_prompt_embeds (`Tensor`, *optional*): negative text embeddings used to guide the image generation. Can be generated from text_encoder step. negative_prompt_embeds_mask (`Tensor`, *optional*): mask for the negative text embeddings. Can be generated from text_encoder step. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py index 37b80b69ec7e..e5f8753253f8 100644 --- a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py +++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py @@ -744,17 +744,17 @@ class QwenImageEditAutoBlocks(SequentialPipelineBlocks): Torch generator for deterministic generation. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - height (`int`): + height (`int`, *optional*): The height in pixels of the generated image. - width (`int`): + width (`int`, *optional*): The width in pixels of the generated image. - image_latents (`Tensor`): + image_latents (`Tensor`, *optional*): image latents used to guide the image generation. Can be generated from vae_encoder step. processed_mask_image (`Tensor`, *optional*): The processed mask image - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. diff --git a/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21.py b/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21.py index 569ffe530186..c4c0301604bd 100644 --- a/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21.py +++ b/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21.py @@ -440,13 +440,13 @@ class QwenImage21AutoCoreDenoiseStep(AutoPipelineBlocks): The width in pixels of the generated image. image_latents (`list`, *optional*): Normalized latents of each condition image. Can be generated from vae_encoder step. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. output_resolution (`int`, *optional*, defaults to 1024): Target side length used to derive the output size and to resize condition images. generator (`Generator`, *optional*): Torch generator for deterministic generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 40): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. @@ -528,9 +528,9 @@ class QwenImage21AutoBlocks(SequentialPipelineBlocks): The width in pixels of the generated image. image_latents (`list`, *optional*): Normalized latents of each condition image. Can be generated from vae_encoder step. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 40): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. diff --git a/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21_turbo.py b/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21_turbo.py index e7ba91f01627..2c975b324cc6 100644 --- a/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21_turbo.py +++ b/src/diffusers/modular_pipelines/qwenimage21/modular_blocks_qwenimage21_turbo.py @@ -409,13 +409,13 @@ class QwenImage21TurboAutoCoreDenoiseStep(AutoPipelineBlocks): The width in pixels of the generated image. image_latents (`list`, *optional*): Normalized latents of each condition image. Can be generated from vae_encoder step. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. output_resolution (`int`, *optional*, defaults to 1024): Target side length used to derive the output size and to resize condition images. generator (`Generator`, *optional*): Torch generator for deterministic generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 8): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. @@ -499,9 +499,9 @@ class QwenImage21TurboAutoBlocks(SequentialPipelineBlocks): The width in pixels of the generated image. image_latents (`list`, *optional*): Normalized latents of each condition image. Can be generated from vae_encoder step. - latents (`Tensor`): + latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*, defaults to 8): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. diff --git a/src/diffusers/modular_pipelines/stable_diffusion_3/modular_blocks_stable_diffusion_3.py b/src/diffusers/modular_pipelines/stable_diffusion_3/modular_blocks_stable_diffusion_3.py index dc2aae05577f..12d5346387f1 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_3/modular_blocks_stable_diffusion_3.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_3/modular_blocks_stable_diffusion_3.py @@ -172,7 +172,7 @@ class StableDiffusion3I2ICoreDenoiseStep(SequentialPipelineBlocks): The height in pixels of the generated image. width (`None`, *optional*): The width in pixels of the generated image. - image_latents (`None`, *optional*): + image_latents (`None`): Latent input image_latents to be processed. latents (`Tensor | NoneType`, *optional*): Pre-generated noisy latents to be used as inputs for image generation. @@ -243,13 +243,13 @@ class StableDiffusion3AutoCoreDenoiseStep(AutoPipelineBlocks): The width in pixels of the generated image. image_latents (`None`, *optional*): Latent input image_latents to be processed. - latents (`Tensor | NoneType`): + latents (`Tensor | NoneType`, *optional*): Pre-generated noisy latents to be used as inputs for image generation. generator (`None`, *optional*): One or a list of torch generator(s) to make generation deterministic. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 50): The number of denoising steps. - timesteps (`None`): + timesteps (`None`, *optional*): Custom timesteps to use for the denoising process. sigmas (`None`, *optional*): Custom sigmas to use for the denoising process. @@ -332,11 +332,11 @@ class StableDiffusion3AutoBlocks(SequentialPipelineBlocks): The number of images to generate per prompt. image_latents (`None`, *optional*): Latent input image_latents to be processed. - latents (`Tensor | NoneType`): + latents (`Tensor | NoneType`, *optional*): Pre-generated noisy latents to be used as inputs for image generation. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 50): The number of denoising steps. - timesteps (`None`): + timesteps (`None`, *optional*): Custom timesteps to use for the denoising process. sigmas (`None`, *optional*): Custom sigmas to use for the denoising process. diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py index ebe43e4ff23c..45fbeebdb239 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py @@ -14,7 +14,7 @@ from ...utils import logging from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks -from ..modular_pipeline_utils import OutputParam +from ..modular_pipeline_utils import InputParam, OutputParam from .before_denoise import ( StableDiffusionXLControlNetInputStep, StableDiffusionXLControlNetUnionInputStep, @@ -127,6 +127,8 @@ class StableDiffusionXLInpaintBeforeDenoiseStep(SequentialPipelineBlocks): StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep, ] block_names = ["set_timesteps", "prepare_latents", "prepare_add_cond"] + # set_timesteps is shared with img2img and declares 0.3; inpainting starts from (almost) pure noise + input_overrides = {"strength": InputParam.template("strength", default=0.9999)} @property def description(self): @@ -363,11 +365,11 @@ class StableDiffusionXLAutoBlocks(SequentialPipelineBlocks): TODO: Add description. denoising_end (`None`, *optional*): TODO: Add description. - strength (`None`, *optional*, defaults to 0.3): - TODO: Add description. + strength (`float`, *optional*, defaults to 0.9999 or 0.3, depending on the workflow): + Strength for img2img/inpainting. denoising_start (`None`, *optional*): TODO: Add description. - latents (`None`): + latents (`None`, *optional*): TODO: Add description. image_latents (`Tensor`, *optional*): The latents representing the reference image for image-to-image/inpainting generation. Can be generated @@ -422,6 +424,10 @@ class StableDiffusionXLAutoBlocks(SequentialPipelineBlocks): pooled_prompt_embeds/negative_pooled_prompt_embeds, and ip_adapter_embeds/negative_ip_adapter_embeds (optional).please add `kwargs_type=denoiser_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state + **controlnet_kwargs (`None`, *optional*): + additional kwargs for controlnet (e.g. control_type_idx and control_type from the controlnet union input + step )please add `kwargs_type=controlnet_kwargs` to their parameter spec (`OutputParam`) when they are + created and added to the pipeline state eta (`None`, *optional*, defaults to 0.0): TODO: Add description. output_type (`None`, *optional*, defaults to pil): diff --git a/src/diffusers/modular_pipelines/wan/before_denoise.py b/src/diffusers/modular_pipelines/wan/before_denoise.py index 0e7aec5364fb..4d4a61db66d1 100644 --- a/src/diffusers/modular_pipelines/wan/before_denoise.py +++ b/src/diffusers/modular_pipelines/wan/before_denoise.py @@ -233,6 +233,16 @@ def intermediate_outputs(self) -> list[str]: type_hint=torch.dtype, description="Data type of model tensor inputs (determined by `transformer.dtype`)", ), + OutputParam( + "prompt_embeds", + type_hint=torch.Tensor, + description="Text embeddings expanded to batch_size * num_videos_per_prompt", + ), + OutputParam( + "negative_prompt_embeds", + type_hint=torch.Tensor, + description="Negative text embeddings expanded to batch_size * num_videos_per_prompt", + ), ] def check_inputs(self, components, block_state): @@ -474,6 +484,15 @@ def inputs(self) -> list[InputParam]: InputParam("sigmas"), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process" + ), + OutputParam("num_inference_steps", type_hint=int, description="The number of denoising steps"), + ] + @torch.no_grad() def __call__( self, components: WanModularPipeline, state: PipelineState diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py index 583bcaff1a7c..2570e341f1be 100644 --- a/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py +++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py @@ -112,7 +112,7 @@ class Wan22Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): TODO: Add description. num_frames (`None`, *optional*): TODO: Add description. - image_condition_latents (`None`, *optional*): + image_condition_latents (`None`): TODO: Add description. num_inference_steps (`None`, *optional*, defaults to 50): TODO: Add description. diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py index cb00e9395169..01bbaef38ea7 100644 --- a/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py +++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py @@ -340,7 +340,7 @@ class WanImage2VideoCoreDenoiseStep(SequentialPipelineBlocks): TODO: Add description. num_frames (`None`, *optional*): TODO: Add description. - image_condition_latents (`None`, *optional*): + image_condition_latents (`None`): TODO: Add description. num_inference_steps (`None`, *optional*, defaults to 50): TODO: Add description. @@ -362,6 +362,14 @@ class WanImage2VideoCoreDenoiseStep(SequentialPipelineBlocks): Number of prompts, the final batch size of model inputs should be batch_size * num_videos_per_prompt dtype (`dtype`): Data type of model tensor inputs (determined by `transformer.dtype`) + prompt_embeds (`Tensor`): + Text embeddings expanded to batch_size * num_videos_per_prompt + negative_prompt_embeds (`Tensor`): + Negative text embeddings expanded to batch_size * num_videos_per_prompt + timesteps (`Tensor`): + The timesteps to use for the denoising process + num_inference_steps (`int`): + The number of denoising steps latents (`Tensor`): The initial latents to use for the denoising process """ @@ -441,7 +449,7 @@ class WanImage2VideoAutoBlocks(SequentialPipelineBlocks): TODO: Add description. attention_kwargs (`None`, *optional*): TODO: Add description. - image_embeds (`Tensor`): + image_embeds (`Tensor`, *optional*): TODO: Add description. output_type (`str`, *optional*, defaults to np): The output type of the decoded videos diff --git a/src/diffusers/modular_pipelines/z_image/before_denoise.py b/src/diffusers/modular_pipelines/z_image/before_denoise.py index e96d80b4bfca..63fdf7e4bd57 100644 --- a/src/diffusers/modular_pipelines/z_image/before_denoise.py +++ b/src/diffusers/modular_pipelines/z_image/before_denoise.py @@ -527,6 +527,7 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process" ), + OutputParam("num_inference_steps", type_hint=int, description="The number of denoising steps"), ] @torch.no_grad() @@ -583,6 +584,19 @@ def inputs(self) -> list[InputParam]: InputParam("strength", default=0.6), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "timesteps", type_hint=torch.Tensor, description="The timesteps to use, trimmed by `strength`" + ), + OutputParam( + "num_inference_steps", + type_hint=int, + description="The number of denoising steps, trimmed by `strength`", + ), + ] + def check_inputs(self, components, block_state): if block_state.strength < 0.0 or block_state.strength > 1.0: raise ValueError(f"Strength must be between 0.0 and 1.0, but got {block_state.strength}") diff --git a/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py b/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py index 1440128a3c55..a0bf2e3a8cc3 100644 --- a/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py +++ b/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py @@ -116,7 +116,7 @@ class ZImageImage2ImageCoreDenoiseStep(SequentialPipelineBlocks): TODO: Add description. width (`None`, *optional*): TODO: Add description. - image_latents (`None`, *optional*): + image_latents (`None`): TODO: Add description. latents (`Tensor | NoneType`, *optional*): TODO: Add description. @@ -190,11 +190,11 @@ class ZImageAutoDenoiseStep(AutoPipelineBlocks): TODO: Add description. image_latents (`None`, *optional*): TODO: Add description. - latents (`Tensor | NoneType`): + latents (`Tensor | NoneType`, *optional*): TODO: Add description. generator (`None`, *optional*): TODO: Add description. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 8): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. @@ -295,9 +295,9 @@ class ZImageAutoBlocks(SequentialPipelineBlocks): TODO: Add description. image_latents (`None`, *optional*): TODO: Add description. - latents (`Tensor | NoneType`): + latents (`Tensor | NoneType`, *optional*): TODO: Add description. - num_inference_steps (`None`): + num_inference_steps (`None`, *optional*, defaults to 8): TODO: Add description. sigmas (`None`, *optional*): TODO: Add description. diff --git a/tests/modular_pipelines/test_sequential_pipeline_blocks.py b/tests/modular_pipelines/test_sequential_pipeline_blocks.py new file mode 100644 index 000000000000..3da843d0135c --- /dev/null +++ b/tests/modular_pipelines/test_sequential_pipeline_blocks.py @@ -0,0 +1,669 @@ +# 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 warnings + +import pytest + +from diffusers.modular_pipelines import ( + AutoPipelineBlocks, + ComponentSpec, + ConditionalPipelineBlocks, + ConfigSpec, + InputParam, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + OutputParam, + SequentialPipelineBlocks, +) + + +class EncodeImageBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("image", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam("image_latents")] + + def __call__(self, components, state): + return components, state + + +class PrepareLatentsBlock(ModularPipelineBlocks): + """Takes `latents` as an optional user input and always produces it.""" + + model_name = "test" + + @property + def inputs(self): + return [InputParam("latents"), InputParam("batch_size", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam("latents")] + + def __call__(self, components, state): + return components, state + + +class DenoiseBlock(ModularPipelineBlocks): + """Requires `latents`, which an earlier block normally produces.""" + + model_name = "test" + + @property + def inputs(self): + return [InputParam("latents", required=True), InputParam("prompt_embeds", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam("latents")] + + def __call__(self, components, state): + return components, state + + +class SkippableEncodeBlocks(ConditionalPipelineBlocks): + block_classes = [EncodeImageBlock] + block_names = ["encode"] + block_trigger_inputs = ["image"] + default_block_name = None + + def select_block(self, image=None): + return "encode" if image is not None else None + + +def strength_auto_blocks(img2img_default, inpaint_default): + """An auto block whose two branches declare `strength` with their own default (`None` for no default).""" + + class Img2ImgBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("image", required=True), InputParam("strength", default=img2img_default)] + + def __call__(self, components, state): + return components, state + + class InpaintBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("mask", required=True), InputParam("strength", default=inpaint_default)] + + def __call__(self, components, state): + return components, state + + class Auto(AutoPipelineBlocks): + block_classes = [InpaintBlock, Img2ImgBlock] + block_names = ["inpaint", "img2img"] + block_trigger_inputs = ["mask", "image"] + + return Auto + + +def required_flags(blocks, *names): + return {inp.name: inp.required for inp in blocks.inputs if inp.name in names} + + +class TestSequentialPipelineBlocksRequiredInputs: + def test_input_produced_by_earlier_block_is_not_required(self): + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"prepare_latents": PrepareLatentsBlock(), "denoise": DenoiseBlock()} + ) + + assert set(blocks.required_inputs) == {"batch_size", "prompt_embeds"} + assert required_flags(blocks, "latents", "batch_size", "prompt_embeds") == { + "latents": False, + "batch_size": True, + "prompt_embeds": True, + } + + def test_input_consumed_before_it_is_produced_stays_required(self): + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"denoise": DenoiseBlock(), "prepare_latents": PrepareLatentsBlock()} + ) + + assert set(blocks.required_inputs) == {"latents", "batch_size", "prompt_embeds"} + + def test_later_block_requiring_an_optional_input_makes_it_required(self): + class OptionalPromptBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("prompt")] + + def __call__(self, components, state): + return components, state + + class RequiredPromptBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("prompt", required=True)] + + def __call__(self, components, state): + return components, state + + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"optional": OptionalPromptBlock(), "required": RequiredPromptBlock()} + ) + + assert blocks.required_inputs == ["prompt"] + assert required_flags(blocks, "prompt") == {"prompt": True} + + def test_default_from_earlier_block_satisfies_later_requirement(self): + class DefaultStepsBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("num_inference_steps", default=50)] + + def __call__(self, components, state): + return components, state + + class RequiredStepsBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("num_inference_steps", required=True)] + + def __call__(self, components, state): + return components, state + + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"default": DefaultStepsBlock(), "required": RequiredStepsBlock()} + ) + + assert blocks.required_inputs == [] + steps = next(inp for inp in blocks.inputs if inp.name == "num_inference_steps") + assert steps.required is False + assert steps.default == 50 + + def test_composite_flags_never_reach_the_sub_block_declarations(self): + class OptionalBatchBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("batch_size")] + + def __call__(self, components, state): + return components, state + + class NeedsImageLatentsBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("image_latents", required=True)] + + def __call__(self, components, state): + return components, state + + # a later requirement makes the composite's entry required, the first sub-block still says optional + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"optional_batch": OptionalBatchBlock(), "prepare_latents": PrepareLatentsBlock()} + ) + assert required_flags(blocks, "batch_size") == {"batch_size": True} + assert required_flags(blocks.sub_blocks["optional_batch"], "batch_size") == {"batch_size": False} + + # a skippable producer makes the composite's entry optional, the leaf still requires it at runtime + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"encode": SkippableEncodeBlocks(), "needs_latents": NeedsImageLatentsBlock()} + ) + assert required_flags(blocks, "image_latents") == {"image_latents": False} + assert required_flags(blocks.sub_blocks["needs_latents"], "image_latents") == {"image_latents": True} + + # a conditional block only requires what every branch requires, each branch keeps its own flag + class Auto(AutoPipelineBlocks): + block_classes = [NeedsImageLatentsBlock, EncodeImageBlock] + block_names = ["needs_latents", "encode"] + block_trigger_inputs = ["image_latents", None] + + auto = Auto() + assert required_flags(auto, "image_latents", "image") == {"image_latents": False, "image": False} + assert required_flags(auto.sub_blocks["needs_latents"], "image_latents") == {"image_latents": True} + assert required_flags(auto.sub_blocks["encode"], "image") == {"image": True} + + def test_output_of_skippable_block_keeps_required_input_listed_but_optional(self): + class NeedsImageLatentsBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("image_latents", required=True)] + + def __call__(self, components, state): + return components, state + + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"encode": SkippableEncodeBlocks(), "consume": NeedsImageLatentsBlock()} + ) + + # the caller can pass `image` and let the encoder produce `image_latents`, or pass `image_latents` directly, + # so neither is required on its own and both stay listed + assert required_flags(blocks, "image", "image_latents") == {"image": False, "image_latents": False} + assert blocks.required_inputs == [] + + +class TestSequentialPipelineBlocksKwargsInputs: + def test_kwargs_type_inputs_are_kept_apart(self): + class TwoKwargsBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam(kwargs_type="denoiser_input_fields"), InputParam(kwargs_type="controlnet_kwargs")] + + def __call__(self, components, state): + return components, state + + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"first": PrepareLatentsBlock(), "second": TwoKwargsBlock()} + ) + + assert [inp.kwargs_type for inp in blocks.inputs if inp.name is None] == [ + "denoiser_input_fields", + "controlnet_kwargs", + ] + + +class TestAutoPipelineBlocksRequiredInputs: + def test_inputs_produced_inside_each_workflow_are_not_required(self): + class Text2ImageBlocks(SequentialPipelineBlocks): + block_classes = [PrepareLatentsBlock, DenoiseBlock] + block_names = ["prepare_latents", "denoise"] + + class Image2ImageBlocks(SequentialPipelineBlocks): + block_classes = [EncodeImageBlock, PrepareLatentsBlock, DenoiseBlock] + block_names = ["encode", "prepare_latents", "denoise"] + + class AutoBlocks(AutoPipelineBlocks): + block_classes = [Image2ImageBlocks, Text2ImageBlocks] + block_names = ["img2img", "text2img"] + block_trigger_inputs = ["image", None] + + blocks = AutoBlocks() + + assert set(blocks.required_inputs) == {"batch_size", "prompt_embeds"} + assert required_flags(blocks, "latents", "image", "batch_size", "prompt_embeds") == { + "latents": False, + "image": False, + "batch_size": True, + "prompt_embeds": True, + } + + +class TestLoopSequentialPipelineBlocksRequiredInputs: + def test_loop_inputs_and_unsatisfied_step_inputs_are_required(self): + class DenoiserStep(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("latents", required=True), InputParam(kwargs_type="denoiser_input_fields")] + + @property + def intermediate_outputs(self): + return [OutputParam("noise_pred")] + + def __call__(self, components, state, i, t): + return components, state + + class SchedulerStep(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [ + InputParam("noise_pred", required=True), + InputParam("latents", required=True), + InputParam(kwargs_type="scheduler_kwargs"), + ] + + @property + def intermediate_outputs(self): + return [OutputParam("latents")] + + def __call__(self, components, state, i, t): + return components, state + + class DenoiseLoop(LoopSequentialPipelineBlocks): + model_name = "test" + block_classes = [DenoiserStep, SchedulerStep] + block_names = ["denoiser", "scheduler"] + + @property + def loop_inputs(self): + return [InputParam("timesteps", required=True), InputParam("num_inference_steps")] + + blocks = DenoiseLoop() + + assert set(blocks.required_inputs) == {"timesteps", "latents"} + assert "noise_pred" not in [inp.name for inp in blocks.inputs] + assert required_flags(blocks, "timesteps", "num_inference_steps", "latents") == { + "timesteps": True, + "num_inference_steps": False, + "latents": True, + } + assert [inp.kwargs_type for inp in blocks.inputs if inp.name is None] == [ + "denoiser_input_fields", + "scheduler_kwargs", + ] + + def test_later_step_requirement_makes_an_optional_loop_input_required(self): + class Step(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("num_inference_steps", required=True), InputParam("guidance_scale", required=True)] + + def __call__(self, components, state, i, t): + return components, state + + class Loop(LoopSequentialPipelineBlocks): + model_name = "test" + block_classes = [Step] + block_names = ["step"] + + @property + def loop_inputs(self): + return [InputParam("num_inference_steps"), InputParam("guidance_scale", default=7.5)] + + # no default on the loop declaration, so the step's requirement wins; a default satisfies it + assert required_flags(Loop(), "num_inference_steps", "guidance_scale") == { + "num_inference_steps": True, + "guidance_scale": False, + } + + +class TestDefaultConflicts: + def _blocks(self, first_default, second_default): + class FirstBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("strength", default=first_default)] + + def __call__(self, components, state): + return components, state + + class SecondBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("strength", default=second_default)] + + def __call__(self, components, state): + return components, state + + return FirstBlock, SecondBlock + + def test_conflicting_defaults_are_deprecated_first_wins(self): + FirstBlock, SecondBlock = self._blocks(0.3, 0.9999) + + class Blocks(SequentialPipelineBlocks): + block_classes = [FirstBlock, SecondBlock] + block_names = ["first", "second"] + + with pytest.warns(FutureWarning, match="'second' declares default 0.9999 for input 'strength'"): + inputs = Blocks().inputs + # the first declaration wins until the deprecation turns into an error + assert [inp.default for inp in inputs if inp.name == "strength"] == [0.3] + + def test_later_none_default_is_not_a_conflict(self): + FirstBlock, SecondBlock = self._blocks(0.3, None) + + class Blocks(SequentialPipelineBlocks): + block_classes = [FirstBlock, SecondBlock] + block_names = ["first", "second"] + + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + Blocks().inputs + + def test_default_override_settles_the_conflict(self): + FirstBlock, SecondBlock = self._blocks(0.3, 0.9999) + + class Blocks(SequentialPipelineBlocks): + block_classes = [FirstBlock, SecondBlock] + block_names = ["first", "second"] + input_overrides = {"strength": InputParam("strength", default=0.9999)} + + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + inputs = Blocks().inputs + assert [inp.default for inp in inputs if inp.name == "strength"] == [0.9999] + + def test_loop_input_without_default_is_not_a_conflict(self): + FirstBlock, _ = self._blocks(0.9999, None) + + class Loop(LoopSequentialPipelineBlocks): + model_name = "test" + block_classes = [FirstBlock] + block_names = ["first"] + + @property + def loop_inputs(self): + return [InputParam("timesteps", required=True), InputParam("strength")] + + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + inputs = Loop().inputs + # the loop declaration has no opinion, so the step's default is the one the loop reports + assert [inp.default for inp in inputs if inp.name == "strength"] == [None] + + def test_loop_blocks_deprecate_and_accept_the_override(self): + FirstBlock, SecondBlock = self._blocks(0.3, 0.9999) + + class Loop(LoopSequentialPipelineBlocks): + model_name = "test" + block_classes = [FirstBlock, SecondBlock] + block_names = ["first", "second"] + + @property + def loop_inputs(self): + return [InputParam("timesteps", required=True)] + + with pytest.warns(FutureWarning, match="'strength'"): + Loop().inputs + + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + inputs = Loop.from_blocks_dict( + {"first": FirstBlock(), "second": SecondBlock()}, + input_overrides={"strength": InputParam("strength", default=0.9999)}, + ).inputs + assert [inp.default for inp in inputs if inp.name == "strength"] == [0.9999] + + def test_branch_defaults_against_a_value_are_a_conflict(self): + FirstBlock, _ = self._blocks(0.5, None) + Auto = strength_auto_blocks(0.3, 0.9999) + + class ValueThenAuto(SequentialPipelineBlocks): + block_classes = [FirstBlock, Auto] + block_names = ["first", "auto"] + + with pytest.warns(FutureWarning, match="'auto' declares per-branch defaults .* declares default 0.5"): + inputs = ValueThenAuto().inputs + assert [(inp.default, inp.defaults_by_block) for inp in inputs if inp.name == "strength"] == [(0.5, None)] + + class AutoThenValue(SequentialPipelineBlocks): + block_classes = [Auto, FirstBlock] + block_names = ["auto", "first"] + + with pytest.warns(FutureWarning, match="'first' declares default 0.5 .* declares per-branch defaults"): + AutoThenValue().inputs + + def test_branch_without_default_has_no_opinion(self): + # one branch agrees with the earlier value and the other declares no default: nothing to settle + FirstBlock, _ = self._blocks(0.3, None) + + class Blocks(SequentialPipelineBlocks): + block_classes = [FirstBlock, strength_auto_blocks(0.3, None)] + block_names = ["first", "auto"] + + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + Blocks().inputs + + def test_branch_defaults_are_carried_onto_an_entry_without_default(self): + FirstBlock, _ = self._blocks(None, None) + + class Blocks(SequentialPipelineBlocks): + block_classes = [FirstBlock, strength_auto_blocks(0.3, 0.9999)] + block_names = ["first", "auto"] + + blocks = Blocks() + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + (strength,) = [inp for inp in blocks.inputs if inp.name == "strength"] + assert strength.default is None + assert strength.defaults_by_block == {"auto.inpaint": 0.9999, "auto.img2img": 0.3} + (first,) = [inp for inp in blocks.sub_blocks["first"].inputs if inp.name == "strength"] + assert first.defaults_by_block is None + + def test_branch_defaults_satisfy_a_later_requirement(self): + class RequiresStrengthBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def inputs(self): + return [InputParam("strength", required=True)] + + def __call__(self, components, state): + return components, state + + class Blocks(SequentialPipelineBlocks): + block_classes = [strength_auto_blocks(0.3, 0.9999), RequiresStrengthBlock] + block_names = ["auto", "requires"] + + # the branch that runs fills in its default before `requires` reads it + assert required_flags(Blocks(), "strength") == {"strength": False} + + +class UNetBlock(ModularPipelineBlocks): + model_name = "test" + + @property + def expected_components(self): + return [ComponentSpec("unet", description="the denoiser")] + + @property + def expected_configs(self): + return [ConfigSpec("scale", default=8)] + + @property + def inputs(self): + return [InputParam("latents", description="noisy latents"), InputParam("batch_size", required=True)] + + @property + def intermediate_outputs(self): + return [OutputParam("latents", description="denoised latents")] + + def __call__(self, components, state): + return components, state + + +class TestOverrides: + def _sequential(self, **overrides): + class Blocks(SequentialPipelineBlocks): + block_classes = [PrepareLatentsBlock, UNetBlock] + block_names = ["prepare_latents", "unet"] + + for name, value in overrides.items(): + setattr(Blocks, name, value) + return Blocks() + + def test_input_override_replaces_the_derived_declaration(self): + blocks = self._sequential(input_overrides={"batch_size": InputParam("batch_size", default=1)}) + (batch_size,) = [inp for inp in blocks.inputs if inp.name == "batch_size"] + assert batch_size.default == 1 + assert batch_size.required is False + assert "batch_size" not in blocks.required_inputs + + def test_override_does_not_mutate_the_sub_block_declaration(self): + blocks = self._sequential(input_overrides={"batch_size": InputParam("batch_size", default=1)}) + _ = blocks.inputs + (batch_size,) = [inp for inp in blocks.sub_blocks["prepare_latents"].inputs if inp.name == "batch_size"] + assert batch_size.required is True + assert batch_size.default is None + + def test_output_component_and_config_overrides(self): + blocks = self._sequential( + intermediate_output_overrides={"latents": OutputParam("latents", description="final latents")}, + component_overrides={"unet": ComponentSpec("unet", description="any UNet2DConditionModel")}, + config_overrides={"scale": ConfigSpec("scale", default=16)}, + ) + assert [out.description for out in blocks.intermediate_outputs if out.name == "latents"] == ["final latents"] + assert [c.description for c in blocks.expected_components if c.name == "unet"] == ["any UNet2DConditionModel"] + assert [c.default for c in blocks.expected_configs if c.name == "scale"] == [16] + + def test_unknown_name_raises(self): + blocks = self._sequential(input_overrides={"strenght": InputParam("strenght", default=0.5)}) + with pytest.raises(ValueError, match="strenght"): + blocks.inputs + blocks = self._sequential(component_overrides={"vae": ComponentSpec("vae", description="x")}) + with pytest.raises(ValueError, match="vae"): + blocks.expected_components + + def test_partial_dict_override_raises(self): + blocks = self._sequential(input_overrides={"batch_size": {"default": 1}}) + with pytest.raises(TypeError, match="must be of type InputParam"): + blocks.inputs + + def test_from_blocks_dict_accepts_overrides(self): + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"prepare_latents": PrepareLatentsBlock(), "unet": UNetBlock()}, + input_overrides={"batch_size": InputParam("batch_size", default=2)}, + intermediate_output_overrides={"latents": OutputParam("latents", description="final latents")}, + config_overrides={"scale": ConfigSpec("scale", default=16)}, + ) + assert [inp.default for inp in blocks.inputs if inp.name == "batch_size"] == [2] + assert [out.description for out in blocks.intermediate_outputs if out.name == "latents"] == ["final latents"] + assert [c.default for c in blocks.expected_configs if c.name == "scale"] == [16] + + loop = LoopSequentialPipelineBlocks.from_blocks_dict( + {"unet": UNetBlock()}, + component_overrides={"unet": ComponentSpec("unet", description="any UNet2DConditionModel")}, + ) + assert [c.description for c in loop.expected_components if c.name == "unet"] == ["any UNet2DConditionModel"] + + def test_auto_blocks_override_replaces_the_per_block_defaults(self): + Auto = strength_auto_blocks(0.3, 0.9999) + + (strength,) = [inp for inp in Auto().inputs if inp.name == "strength"] + assert strength.default is None + assert strength.defaults_by_block == {"inpaint": 0.9999, "img2img": 0.3} + + class AutoWithDefault(Auto): + input_overrides = {"strength": InputParam("strength", default=0.5)} + + (strength,) = [inp for inp in AutoWithDefault().inputs if inp.name == "strength"] + assert strength.default == 0.5 + assert strength.defaults_by_block is None