Skip to content

[modular] ensure branch-specific input defaults #14132

Description

@yiyixuxu

What API design would you like to have changed or added to the library? Why?

Currently, if a AutoPipelineBlocks contains 2 branch blocks; if both have the same input, one branch block declares a default while the other does not; both blocks will have the same default, so that the default value set in one branch will leak into the sibling one, which is unexpected

What use case would this enable or better enable? Can you give us a code example?

"""
Repro: in an AutoPipelineBlocks, a default declared by one branch leaks into the
sibling branch via `combine_inputs` (non-None default overwrites None) + global
default-filling of PipelineState in `ModularPipeline.__call__`.

The sibling branch uses `None` as a "user didn't pass this" sentinel, so it can
reject explicitly-passed values. The leaked default makes that check fire on
every run even when the user never passed the input.
"""

from diffusers.modular_pipelines import AutoPipelineBlocks, ModularPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam


class PlainStep(ModularPipelineBlocks):
    model_name = "dummy"

    @property
    def description(self):
        return "Plain branch: declares a real default for num_frames."

    @property
    def inputs(self):
        return [InputParam(name="num_frames", default=189)]

    @property
    def intermediate_outputs(self):
        return [OutputParam("resolved_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 ActionStep(ModularPipelineBlocks):
    model_name = "dummy"

    @property
    def description(self):
        return "Action branch: num_frames must NOT be passed (derived from 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")]

    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  # pretend: derived from action
        self.set_block_state(state, block_state)
        return components, state


class AutoStep(AutoPipelineBlocks):
    model_name = "dummy"
    block_classes = [ActionStep, PlainStep]
    block_names = ["action", "plain"]
    block_trigger_inputs = ["action", None]

    @property
    def description(self):
        return "Runs ActionStep when `action` is provided, PlainStep otherwise."


auto = AutoStep()
merged = {p.name: p.default for p in auto.inputs}
print(f"merged pipeline-level inputs: {merged}")
print(f"  -> ActionStep declared num_frames default=None, but merged default is {merged['num_frames']}\n")

pipe = auto.init_pipeline()

print("1) plain path, no num_frames passed:")
state = pipe()
print(f"   resolved_num_frames = {state.get('resolved_num_frames')}  (default worked)\n")

print("2) action path, user does NOT pass num_frames:")
try:
    pipe(action="dummy-action")
except ValueError as e:
    print(f"   ValueError: {e}")
    print("   ^ spurious! the user never passed num_frames — PlainStep's default=189")
    print("     was filled into state before branch selection, so ActionStep cannot")
    print("     tell 'user passed 189' from 'default filled 189'.")

you can get away from the unexpected behavior by not setting a default in InputParam, and handle the default value assignment inside call of each branch
like in #14110, but I think we should fix this on our end

Metadata

Metadata

Assignees

Labels

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions