From 5cb63ddd472a7224e5068ce0ae1f4e83fb85ae93 Mon Sep 17 00:00:00 2001 From: Pawel Date: Tue, 7 Jul 2026 13:00:59 -0700 Subject: [PATCH] Fix modular branch input defaults --- .../modular_pipelines/modular_pipeline.py | 48 +++++- .../test_conditional_pipeline_blocks.py | 152 ++++++++++++++++++ 2 files changed, 196 insertions(+), 4 deletions(-) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 2d661028acf6..9e2af11522d3 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -48,7 +48,6 @@ InsertableDict, OutputParam, _validate_requirements, - combine_inputs, combine_outputs, format_components, format_configs, @@ -500,7 +499,10 @@ def get_block_state(self, state: PipelineState) -> dict: # Check inputs for input_param in state_inputs: if input_param.name: - value = state.get(input_param.name) + if input_param.name in state.values: + value = state.get(input_param.name) + else: + 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): @@ -666,7 +668,45 @@ def required_inputs(self) -> list[str]: @property 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) + + def get_input_key(input_param: InputParam) -> str: + if input_param.name is None and input_param.kwargs_type is not None: + return "*_" + input_param.kwargs_type + return input_param.name + + combined_inputs = [] + combined_by_key = {} + block_count = len(named_inputs) + input_occurrences = {} + input_defaults = {} + + for _, block_inputs in named_inputs: + seen_inputs = set() + for input_param in block_inputs: + input_key = get_input_key(input_param) + if input_key in seen_inputs: + continue + + if input_key not in combined_by_key: + combined_by_key[input_key] = deepcopy(input_param) + combined_inputs.append(combined_by_key[input_key]) + + input_occurrences[input_key] = input_occurrences.get(input_key, 0) + 1 + input_defaults.setdefault(input_key, []).append(input_param.default) + seen_inputs.add(input_key) + + # Conditional blocks should only expose a pipeline-level default when every branch + # declares the input with the same default. Otherwise, keep the pipeline-level + # default unset and let the selected block apply its own local default. + for input_param in combined_inputs: + input_key = get_input_key(input_param) + defaults = input_defaults.get(input_key, []) + + if input_occurrences.get(input_key, 0) != block_count: + input_param.default = None + elif defaults and any(default != defaults[0] for default in defaults[1:]): + input_param.default = None + # mark Required inputs only if that input is required by all the blocks for input_param in combined_inputs: if input_param.name in self.required_inputs: @@ -2821,7 +2861,7 @@ def __call__(self, state: PipelineState = None, output: str | list[str] = None, kwargs_dict = passed_kwargs.pop(kwargs_type) for k, v in kwargs_dict.items(): state.set(k, v, kwargs_type) - elif name is not None and name not in state.values: + elif name is not None and name not in state.values and default is not None: state.set(name, default, kwargs_type) # Warn about unexpected inputs diff --git a/tests/modular_pipelines/test_conditional_pipeline_blocks.py b/tests/modular_pipelines/test_conditional_pipeline_blocks.py index fa5f64f8b31b..e6b9b3f91374 100644 --- a/tests/modular_pipelines/test_conditional_pipeline_blocks.py +++ b/tests/modular_pipelines/test_conditional_pipeline_blocks.py @@ -18,6 +18,7 @@ ConditionalPipelineBlocks, InputParam, ModularPipelineBlocks, + OutputParam, ) @@ -133,6 +134,126 @@ def description(self): return "Auto image blocks for testing" +class PlainFramesStep(ModularPipelineBlocks): + model_name = "plain" + + @property + def inputs(self): + return [InputParam(name="num_frames", default=189)] + + @property + def intermediate_outputs(self): + return [OutputParam("resolved_num_frames")] + + @property + def description(self): + return "Plain branch: declares a real default for num_frames." + + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.resolved_num_frames = block_state.num_frames + self.set_block_state(state, block_state) + return components, state + + +class ActionFramesStep(ModularPipelineBlocks): + model_name = "action" + + @property + def inputs(self): + return [InputParam(name="action", required=True), InputParam(name="num_frames", default=None)] + + @property + def intermediate_outputs(self): + return [OutputParam("resolved_num_frames")] + + @property + def description(self): + return "Action branch: num_frames must not be passed explicitly." + + def __call__(self, components, state): + block_state = self.get_block_state(state) + if block_state.num_frames is not None: + raise ValueError("`num_frames` has to be None if `action` is provided.") + block_state.resolved_num_frames = 100 + self.set_block_state(state, block_state) + return components, state + + +class SharedFramesStep(ModularPipelineBlocks): + model_name = "shared" + + @property + def inputs(self): + return [InputParam(name="num_frames", default=24)] + + @property + def intermediate_outputs(self): + return [OutputParam("resolved_num_frames")] + + @property + def description(self): + return "Branch with a shared default." + + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.resolved_num_frames = block_state.num_frames + self.set_block_state(state, block_state) + return components, state + + +class AlternateFramesStep(ModularPipelineBlocks): + model_name = "alternate" + + @property + def inputs(self): + return [InputParam(name="num_frames", default=32)] + + @property + def intermediate_outputs(self): + return [OutputParam("resolved_num_frames")] + + @property + def description(self): + return "Branch with a conflicting default." + + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.resolved_num_frames = block_state.num_frames + self.set_block_state(state, block_state) + return components, state + + +class BranchDefaultBlocks(AutoPipelineBlocks): + block_classes = [ActionFramesStep, PlainFramesStep] + block_names = ["action", "plain"] + block_trigger_inputs = ["action", None] + + @property + def description(self): + return "Auto blocks used to reproduce branch-local default leakage." + + +class SharedDefaultBlocks(AutoPipelineBlocks): + block_classes = [SharedFramesStep, SharedFramesStep] + block_names = ["first", "second"] + block_trigger_inputs = ["first_trigger", None] + + @property + def description(self): + return "Auto blocks with a genuinely shared default." + + +class ConflictingDefaultBlocks(AutoPipelineBlocks): + block_classes = [AlternateFramesStep, SharedFramesStep] + block_names = ["first", "second"] + block_trigger_inputs = ["first_trigger", None] + + @property + def description(self): + return "Auto blocks with conflicting non-None defaults." + + class TestConditionalPipelineBlocksSelectBlock: def test_select_block_with_mask(self): blocks = ConditionalImageBlocks() @@ -240,3 +361,34 @@ def test_sub_block_types(self): def test_description(self): blocks = ConditionalImageBlocks() assert "Conditional" in blocks.description + + +class TestConditionalPipelineBlockDefaults: + def test_branch_local_default_does_not_leak_into_sibling_branch(self): + pipe = BranchDefaultBlocks().init_pipeline() + + state = pipe(action="dummy-action") + + assert state.get("resolved_num_frames") == 100 + + def test_default_branch_still_receives_its_own_default(self): + pipe = BranchDefaultBlocks().init_pipeline() + + state = pipe() + + assert state.get("resolved_num_frames") == 189 + + def test_default_call_parameters_hide_branch_local_defaults(self): + pipe = BranchDefaultBlocks().init_pipeline() + + assert pipe.default_call_parameters["num_frames"] is None + + def test_shared_defaults_are_preserved(self): + pipe = SharedDefaultBlocks().init_pipeline() + + assert pipe.default_call_parameters["num_frames"] == 24 + + def test_conflicting_non_none_defaults_are_not_promoted(self): + pipe = ConflictingDefaultBlocks().init_pipeline() + + assert pipe.default_call_parameters["num_frames"] is None