diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 28e485eddbc1..ba9740e23258 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/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 diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py index 07dd94c5693e..733f5bdc8a73 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py +++ b/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: diff --git a/tests/modular_pipelines/test_conditional_pipeline_blocks.py b/tests/modular_pipelines/test_conditional_pipeline_blocks.py index 184f0bfd4535..cc0075d0022f 100644 --- a/tests/modular_pipelines/test_conditional_pipeline_blocks.py +++ b/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