Fix modular branch input defaults

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

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