[modular] fill a missing input from the first block that reads it, keep a passed None

src/diffusers/modular_pipelines/modular_pipeline.py

@@ -588,13 +588,12 @@ 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 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
+                if input_param.name in state.values:
+                    # a value in the state, None included, is what the user passed or an earlier block set
+                    value = state.get(input_param.name)
+                else:
+                    # the first block that reads an input fills in its own default; `set_block_state` records it so
+                    # the blocks after it read the same value
                     value = input_param.default
                 if input_param.required and value is None:
                     raise ValueError(f"Required input '{input_param.name}' is missing")
@@ -624,9 +623,10 @@ def set_block_state(self, state: PipelineState, block_state: BlockState):
         for input_param in self.inputs:
             if input_param.name and hasattr(block_state, input_param.name):
                 param = getattr(block_state, input_param.name)
-                # Only add if the value is different from what's in the state
+                # Only add if the value is different from what's in the state, or the state did not hold it yet (the
+                # default this block filled in, see `get_block_state`)
                 current_value = state.get(input_param.name)
-                if current_value is not param:  # Using identity comparison to check if object was modified
+                if input_param.name not in state.values or current_value is not param:  # identity: object modified?
                     state.set(input_param.name, param, input_param.kwargs_type)
 
             elif input_param.kwargs_type:
@@ -3006,6 +3006,9 @@ def __call__(self, state: PipelineState = None, output: str | list[str] = None,
 
         # Add inputs to state, using defaults if not provided in the kwargs or the state
         # if same input already in the state, will override it if provided in the kwargs
+        # a None passed by the user is stored as None. A None default is not prefilled: the first block that reads
+        # the input fills in its own default and records it, see `get_block_state` and `set_block_state`. That is how
+        # a conditional branch gets its own default (`defaults_by_block`)
         for expected_input_param in self._blocks.inputs:
             name = expected_input_param.name
             default = expected_input_param.default
@@ -3016,7 +3019,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

src/diffusers/modular_pipelines/modular_pipeline_utils.py

@@ -1091,8 +1091,8 @@ def combine_inputs(*named_input_lists: list[tuple[str, list[InputParam]]]) -> li
     """
     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.
+    recorded in `defaults_by_block`; the branch that runs fills in its own default when it reads the state
+    (`get_block_state`), and the docstring formatters render the per-block defaults.
 
     Example:
 

tests/modular_pipelines/test_conditional_pipeline_blocks.py

@@ -19,6 +19,7 @@
     InputParam,
     ModularPipelineBlocks,
     OutputParam,
+    SequentialPipelineBlocks,
 )
 from diffusers.modular_pipelines.modular_pipeline_utils import combine_inputs
 
@@ -320,6 +321,27 @@ def select_block(self, mask=None) -> str | None:
         return None
 
 
+class NoDefaultStrengthBlock(ModularPipelineBlocks):
+    @property
+    def inputs(self):
+        return [InputParam(name="strength", type_hint=float)]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam(name="strength_seen_first")]
+
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+        block_state.strength_seen_first = block_state.strength
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class NoDefaultThenImg2Img(SequentialPipelineBlocks):
+    block_classes = [NoDefaultStrengthBlock, ImageToImageBlock]
+    block_names = ["no_default", "img2img"]
+
+
 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
@@ -372,3 +394,44 @@ def test_nested_defaults_prefixed_with_sub_block_name(self):
         # 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
+
+    def test_explicit_none_reaches_every_branch(self):
+        # a passed None is a value, not "use the default": the block keeps it instead of filling in its default
+        pipe = ImageToImageBlock().init_pipeline()
+        state = pipe(prompt="p", image="i", strength=None)
+        assert state.get("strength") is None
+
+        # and no branch of a conditional block swaps in its own default either
+        pipe = AutoImageBlocks().init_pipeline()
+        state = pipe(prompt="p", image="i", strength=None)
+        assert state.get("workflow") == "img2img"
+        assert state.get("strength") is None
+
+        state = pipe(prompt="p", image="i", mask="m", strength=None)
+        assert state.get("workflow") == "inpaint"
+        assert state.get("strength") is None
+
+    def test_nested_branch_resolves_own_default(self):
+        pipe = NestedImageBlocks().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_branch_defaults_do_not_affect_sequential_blocks(self):
+        # in a sequence an input is one value for every block. The first block that reads `strength` fills in its
+        # default (None) and records it; the img2img block then reads that None instead of its own 0.3, so the value
+        # the first block saw and the value left in the state are the same
+        pipe = NoDefaultThenImg2Img().init_pipeline()
+        state = pipe(prompt="p", image="i")
+        assert state.get("workflow") == "img2img"
+        assert state.get("strength_seen_first") is None
+        assert state.get("strength") is None
+
+        # a passed value reaches every block
+        state = pipe(prompt="p", image="i", strength=0.7)
+        assert state.get("strength_seen_first") == 0.7
+        assert state.get("strength") == 0.7