diff --git a/.ai/modular.md b/.ai/modular.md index 016b099b0c8b..b68129d08536 100644 --- a/.ai/modular.md +++ b/.ai/modular.md @@ -235,6 +235,19 @@ OutputParam( If a template's predefined description doesn't fit (e.g. the `"latents"` output template means "Denoised latents", which is wrong for the noisy latents out of a prepare-latents step) — drop the template and declare the field directly with an accurate description. See gotcha #5. +**Declare defaults in the `InputParam`, not inside `__call__`.** + +```python +# yes +InputParam(name="num_frames", type_hint=int, default=189) + +# no — works, but the assembled pipeline is not aware of it +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. + ## ComponentSpec patterns ```python diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index b62b6d7fe7b4..d7c6c846fc8e 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -123,6 +123,9 @@ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): Components: video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) + Configs: + default_use_system_prompt (default: True) enable_safety_checker (default: True) + Inputs: control_videos (`dict`, *optional*): Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. @@ -140,9 +143,8 @@ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): The text prompt that guides Cosmos3 generation. negative_prompt (`str`, *optional*): The negative text prompt used for classifier-free guidance. - use_system_prompt (`bool`, *optional*): - Whether to prepend the system prompt. Defaults to the pipeline configuration for standard and action - workflows and to True for transfer. + use_system_prompt (`bool`, *optional*, defaults to True or None, depending on the workflow): + Whether to prepend the Cosmos3 transfer system prompt. action (`CosmosActionCondition`, *optional*): Action-conditioning metadata and its reference visual input. fps (`float`, *optional*, defaults to 24.0): @@ -368,6 +370,9 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks): Components: transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + Configs: + use_native_flow_schedule (default: False) + Inputs: cond_input_ids (`None`): Token IDs for the conditional prompt. @@ -436,6 +441,9 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks): Components: transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + Configs: + use_native_flow_schedule (default: False) + Inputs: cond_input_ids (`None`): Token IDs for the conditional prompt. @@ -517,6 +525,9 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks): Components: transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + Configs: + use_native_flow_schedule (default: False) + Inputs: cond_input_ids (`None`): Token IDs for the conditional prompt. @@ -602,6 +613,9 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks): Components: transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + Configs: + use_native_flow_schedule (default: False) + Inputs: cond_input_ids (`None`): Token IDs for the conditional prompt. @@ -962,6 +976,9 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler (`UniPCMultistepScheduler`) + Configs: + use_native_flow_schedule (default: False) + Inputs: cond_input_ids (`None`): Token IDs for the conditional prompt. @@ -1149,6 +1166,10 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) + Configs: + default_use_system_prompt (default: True) enable_safety_checker (default: True) use_native_flow_schedule + (default: False) + Inputs: control_videos (`dict`, *optional*): Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. @@ -1166,9 +1187,8 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): The text prompt that guides Cosmos3 generation. negative_prompt (`str`, *optional*): The negative text prompt used for classifier-free guidance. - use_system_prompt (`bool`, *optional*): - Whether to prepend the system prompt. Defaults to the pipeline configuration for standard and action - workflows and to True for transfer. + use_system_prompt (`bool`, *optional*, defaults to True or None, depending on the workflow): + Whether to prepend the Cosmos3 transfer system prompt. action (`CosmosActionCondition`, *optional*): Action-conditioning metadata and its reference visual input. fps (`float`, *optional*, defaults to 24.0): diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py index ff2be89f2241..e168cc24d8cd 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py @@ -165,7 +165,8 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`) Configs: - is_distilled (default: True) distilled_sigmas (default: None) + default_use_system_prompt (default: True) enable_safety_checker (default: True) is_distilled (default: True) + distilled_sigmas (default: None) Inputs: prompt (`str`): diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index d43825860d8e..00b51501210d 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -502,6 +502,13 @@ def get_block_state(self, state: PipelineState) -> dict: for input_param in state_inputs: if input_param.name: value = state.get(input_param.name) + if value is None: + # if the value is None (not passed, or passed as None), the first block that reads it sets the + # default at call time. For sequential blocks the default is already resolved at compile time + # (first declaring block wins, see _get_inputs), but for conditional blocks the selected branch + # is only known at runtime: disagreeing branch defaults merge to None (see combine_inputs) and + # the block that actually runs applies its own declared default here + value = input_param.default if input_param.required and value is None: raise ValueError(f"Required input '{input_param.name}' is missing") elif value is not None or (value is None and input_param.name not in data): diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py index 7aa61c55ef1f..2c859cc0ff96 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py +++ b/src/diffusers/modular_pipelines/modular_pipeline_utils.py @@ -14,9 +14,8 @@ import inspect import re -import warnings from collections import OrderedDict -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from types import UnionType from typing import Any, Literal, Type, Union, get_args, get_origin @@ -556,6 +555,10 @@ class InputParam: description: str = "" kwargs_type: str = None metadata: dict[str, Any] = None + # set by `combine_inputs` when sub-blocks of a ConditionalPipelineBlocks declare different defaults for this + # input: maps block name -> that block's declared default; `default` is None in that case and each block + # resolves its own default at runtime in `get_block_state` + defaults_by_block: dict[str, Any] = None def __repr__(self): return f"<{self.name}: {'required' if self.required else 'optional'}, default={self.default}>" @@ -751,7 +754,6 @@ def wrap_text(text, indent, max_length): for param in params: # Format parameter name and type type_str = get_type_str(param.type_hint) if param.type_hint != Any else "" - # YiYi Notes: remove this line if we remove kwargs_type name = f"**{param.kwargs_type}" if param.name is None and param.kwargs_type is not None else param.name param_str = f"{param_indent}{name} (`{type_str}`" @@ -759,7 +761,13 @@ def wrap_text(text, indent, max_length): if hasattr(param, "required"): if not param.required: param_str += ", *optional*" - if param.default is not None: + if param.defaults_by_block is not None: + # e.g. ", defaults to None or 189, depending on the workflow" + distinct_defaults = list(dict.fromkeys(param.defaults_by_block.values())) + param_str += ( + f", defaults to {' or '.join(str(v) for v in distinct_defaults)}, depending on the workflow" + ) + elif param.default is not None: param_str += f", defaults to {param.default}" param_str += "):" @@ -833,7 +841,13 @@ def get_type_str(type_hint): if hasattr(param, "required") and not param.required: param_str += ", *optional*" - if param.default is not None: + if param.defaults_by_block is not None: + # e.g. ", defaults to `None` or `189`, depending on the workflow" + distinct_defaults = list(dict.fromkeys(param.defaults_by_block.values())) + param_str += ( + f", defaults to {' or '.join(f'`{v}`' for v in distinct_defaults)}, depending on the workflow" + ) + elif param.default is not None: param_str += f", defaults to `{param.default}`" param_str += ")" @@ -1102,9 +1116,25 @@ def _accumulate(mapping: dict[str, Any]): def combine_inputs(*named_input_lists: list[tuple[str, list[InputParam]]]) -> list[InputParam]: """ - Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if current - default value is None and new default value is not None. Warns if multiple non-None default values exist for the - same input. + Combines multiple lists of InputParam objects from different blocks. Duplicate inputs keep the first occurrence. If + duplicate inputs declare different defaults, the combined input's default is None and the per-block defaults are + recorded in `defaults_by_block`; each block resolves its own default at runtime in `get_block_state`, and the + docstring formatters render the per-block defaults. + + Example: + + ```python + combine_inputs( + ("img2img", [InputParam("prompt"), InputParam("strength", default=0.3)]), + ("inpaint", [InputParam("prompt"), InputParam("strength", default=0.9999)]), + ) + # returns: + # InputParam("prompt") # no disagreement -> first occurrence kept as-is + # InputParam("strength", default=None, defaults_by_block={"img2img": 0.3, "inpaint": 0.9999}) + ``` + + See `TestConditionalBlocksInputs` in `tests/modular_pipelines/test_conditional_pipeline_blocks.py` for the full + behavior. Args: named_input_lists: List of tuples containing (block_name, input_param_list) pairs @@ -1112,8 +1142,8 @@ def combine_inputs(*named_input_lists: list[tuple[str, list[InputParam]]]) -> li Returns: List[InputParam]: Combined list of unique InputParam objects """ - combined_dict = {} # name -> InputParam - value_sources = {} # name -> block_name + combined_dict = {} # name -> InputParam, e.g. {"strength": InputParam("strength", default=0.3)} + defaults_by_block = {} # name -> {block_name: default}, e.g. {"strength": {"img2img": 0.3, "inpaint": 0.9999}} for block_name, inputs in named_input_lists: for input_param in inputs: @@ -1121,24 +1151,29 @@ def combine_inputs(*named_input_lists: list[tuple[str, list[InputParam]]]) -> li input_name = "*_" + input_param.kwargs_type else: input_name = input_param.name - if input_name in combined_dict: - current_param = combined_dict[input_name] - if ( - current_param.default is not None - and input_param.default is not None - and current_param.default != input_param.default - ): - warnings.warn( - f"Multiple different default values found for input '{input_name}': " - f"{current_param.default} (from block '{value_sources[input_name]}') and " - f"{input_param.default} (from block '{block_name}'). Using {current_param.default}." - ) - if current_param.default is None and input_param.default is not None: - combined_dict[input_name] = input_param - value_sources[input_name] = block_name + + if input_param.defaults_by_block: + # nested conditional block: prefix its block names with the sub-block name, + # e.g. {"img2img": 0.3, "inpaint": 0.9999} from sub-block "image" -> {"image.img2img": 0.3, "image.inpaint": 0.9999} + new_defaults = {f"{block_name}.{k}": v for k, v in input_param.defaults_by_block.items()} else: + new_defaults = {block_name: input_param.default} + + if input_name not in combined_dict: combined_dict[input_name] = input_param - value_sources[input_name] = block_name + defaults_by_block[input_name] = new_defaults + continue + + # duplicate input: accumulate this block's default, + # e.g. "strength" was {"img2img": 0.3}, after inpaint's turn it is {"img2img": 0.3, "inpaint": 0.9999} + defaults_by_block[input_name].update(new_defaults) + defaults = list(defaults_by_block[input_name].values()) + if any(d != defaults[0] for d in defaults[1:]): + # blocks disagree (0.3 vs 0.9999): combined default becomes None and the per-block defaults are + # recorded on the param; `replace` creates a fresh copy so the blocks' own params are never mutated + combined_dict[input_name] = replace( + combined_dict[input_name], default=None, defaults_by_block=dict(defaults_by_block[input_name]) + ) return list(combined_dict.values()) diff --git a/tests/modular_pipelines/test_conditional_pipeline_blocks.py b/tests/modular_pipelines/test_conditional_pipeline_blocks.py index fa5f64f8b31b..184f0bfd4535 100644 --- a/tests/modular_pipelines/test_conditional_pipeline_blocks.py +++ b/tests/modular_pipelines/test_conditional_pipeline_blocks.py @@ -18,7 +18,60 @@ ConditionalPipelineBlocks, InputParam, ModularPipelineBlocks, + OutputParam, ) +from diffusers.modular_pipelines.modular_pipeline_utils import combine_inputs + + +class TestConditionalBlocksInputs: + def test_combine_inputs_agreeing_defaults_stay_untouched(self): + combined = combine_inputs( + ("block-a", [InputParam(name="x", default=5)]), + ("block-b", [InputParam(name="x", default=5)]), + ) + assert combined == [InputParam(name="x", default=5)] + + def test_combine_inputs_first_occurrence_wins(self): + combined = combine_inputs( + ("block-a", [InputParam(name="x", default=5, description="from block-a")]), + ("block-b", [InputParam(name="x", default=5, description="from block-b")]), + ) + # duplicate inputs keep the first occurrence: block-b's description is dropped + assert combined == [InputParam(name="x", default=5, description="from block-a")] + + def test_combine_inputs_none_default_counts_as_disagreement(self): + # a None default is a sentinel ("user didn't pass this"), it must not be overridden by a sibling's default + combined = combine_inputs( + ("block-a", [InputParam(name="x", default=None)]), + ("block-b", [InputParam(name="x", default=5)]), + ) + assert combined == [InputParam(name="x", default=None, defaults_by_block={"block-a": None, "block-b": 5})] + + def test_combine_inputs_disagreement_records_every_block(self): + combined = combine_inputs( + ("block-a", [InputParam(name="x", default=1)]), + ("block-b", [InputParam(name="x", default=2)]), + ("block-c", [InputParam(name="x", default=1)]), + ) + # any disagreement records every block's default, including the ones that agree with each other + assert combined == [ + InputParam(name="x", default=None, defaults_by_block={"block-a": 1, "block-b": 2, "block-c": 1}) + ] + + def test_combine_inputs_nested_defaults_prefixed(self): + # block-b is itself a conditional block whose branches already disagree, so its param carries + # defaults_by_block; merging at the outer level prefixes those entries with the sub-block name + combined = combine_inputs( + ("block-a", [InputParam(name="x", default=3)]), + ("block-b", [InputParam(name="x", default=None, defaults_by_block={"branch-a": 1, "branch-b": 2})]), + ) + assert combined == [ + InputParam( + name="x", + default=None, + defaults_by_block={"block-a": 3, "block-b.branch-a": 1, "block-b.branch-b": 2}, + ) + ] class TextToImageBlock(ModularPipelineBlocks): @@ -30,7 +83,7 @@ def inputs(self): @property def intermediate_outputs(self): - return [] + return [OutputParam(name="workflow")] @property def description(self): @@ -48,11 +101,15 @@ class ImageToImageBlock(ModularPipelineBlocks): @property def inputs(self): - return [InputParam(name="prompt"), InputParam(name="image")] + return [ + InputParam(name="prompt"), + InputParam(name="image"), + InputParam(name="strength", type_hint=float, default=0.3), + ] @property def intermediate_outputs(self): - return [] + return [OutputParam(name="workflow")] @property def description(self): @@ -70,11 +127,16 @@ class InpaintBlock(ModularPipelineBlocks): @property def inputs(self): - return [InputParam(name="prompt"), InputParam(name="image"), InputParam(name="mask")] + return [ + InputParam(name="prompt"), + InputParam(name="image"), + InputParam(name="mask"), + InputParam(name="strength", type_hint=float, default=0.9999), + ] @property def intermediate_outputs(self): - return [] + return [OutputParam(name="workflow")] @property def description(self): @@ -240,3 +302,73 @@ def test_sub_block_types(self): def test_description(self): blocks = ConditionalImageBlocks() assert "Conditional" in blocks.description + + +class NestedImageBlocks(ConditionalPipelineBlocks): + block_classes = [InpaintBlock, AutoImageBlocks] + block_names = ["refine", "image"] + block_trigger_inputs = ["mask"] + default_block_name = "image" + + @property + def description(self): + return "Nested conditional blocks: refine when `mask` is provided, auto image blocks otherwise" + + def select_block(self, mask=None) -> str | None: + if mask is not None: + return "refine" + return None + + +class TestConditionalBlocksBranchDefaults: + def test_conflicting_defaults_merge_to_none(self): + # inpaint (0.9999) and img2img (0.3) disagree on strength: the merged input's default becomes None + # and the per-block defaults are recorded in defaults_by_block + strength_params = [p for p in AutoImageBlocks().inputs if p.name == "strength"] + assert len(strength_params) == 1 + assert strength_params[0] == InputParam( + name="strength", + type_hint=float, + default=None, + defaults_by_block={"inpaint": 0.9999, "img2img": 0.3}, + ) + + # the docstring renders the distinct per-block defaults instead of a single value + doc = " ".join(AutoImageBlocks().doc.split()) + assert "strength (`float`, *optional*, defaults to 0.9999 or 0.3, depending on the workflow):" in doc + + def test_branch_resolves_own_default(self): + pipe = AutoImageBlocks().init_pipeline() + state = pipe(prompt="p", image="i") + assert state.get("workflow") == "img2img" + assert state.get("strength") == 0.3 + + state = pipe(prompt="p", image="i", mask="m") + assert state.get("workflow") == "inpaint" + assert state.get("strength") == 0.9999 + + def test_explicit_value_overrides_branch_default(self): + pipe = AutoImageBlocks().init_pipeline() + state = pipe(prompt="p", image="i", strength=0.7) + assert state.get("workflow") == "img2img" + assert state.get("strength") == 0.7 + + def test_standalone_branch_keeps_default(self): + pipe = ImageToImageBlock().init_pipeline() + state = pipe(prompt="p", image="i") + assert state.get("workflow") == "img2img" + assert state.get("strength") == 0.3 + + def test_nested_defaults_prefixed_with_sub_block_name(self): + strength_params = [p for p in NestedImageBlocks().inputs if p.name == "strength"] + assert len(strength_params) == 1 + assert strength_params[0] == InputParam( + name="strength", + type_hint=float, + default=None, + defaults_by_block={"refine": 0.9999, "image.inpaint": 0.9999, "image.img2img": 0.3}, + ) + + # the docstring renders distinct values only: refine and image.inpaint both use 0.9999, shown once + doc = " ".join(NestedImageBlocks().doc.split()) + assert "strength (`float`, *optional*, defaults to 0.9999 or 0.3, depending on the workflow):" in doc