Skip to content

[modular] ConditionalPipelineBlocks should declare its own trigger inputs #15101

Open

Description

@yiyixuxu

When an input exists only to decide which branch a conditional block runs, no sub-block reads it, so nothing declares it. ConditionalPipelineBlocks.inputs only merges the sub-blocks' inputs and does not add the block's own block_trigger_inputs. The repo has three patterns for this:

  1. Declare it nowhere. Ernie's use_pe is only listed in block_trigger_inputs:

    class ErnieImageAutoPromptEnhancerStep(ConditionalPipelineBlocks):
        block_trigger_inputs = ["use_pe"]
    
        def select_block(self, use_pe=None):
            return "prompt_enhancer" if use_pe else None

    use_pe is missing from inputs, from the generated docstring and from default_call_parameters. And it does not work at call time either: ModularPipeline.__call__ only copies declared inputs into the state and warns about the rest, so pipe(prompt=..., use_pe=True) prints "Unexpected input 'dict_keys(['use_pe'])' provided. This input will be ignored.", the conditional block reads None from the state, and the prompt enhancer never runs. Only the static get_execution_blocks(prompt=True, use_pe=True) resolves to it, because that takes kwargs directly.

  2. Override inputs on the wrapper. Cosmos 3's enable_sound:

    class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks):
        block_trigger_inputs = ["action", "enable_sound", "control_videos"]
    
        @property
        def inputs(self):
            inputs = super().inputs
            inputs.append(InputParam(name="enable_sound", type_hint=bool, default=False, description="..."))
            return inputs

    The input is listed at the top level with a type, a default and a description. And because enable_sound is also a workflow input, get_workflow does the right thing with it:

    >>> blocks = Cosmos3OmniBlocks()
    >>> wf = blocks.get_workflow("text2video_with_sound")
    >>> [name for name in wf.sub_blocks if "sound" in name]
    ['denoise.prepare_sound_latents', 'denoise.pack_sound_sequence', 'denoise.prepare_sound_denoiser_inputs', 'decode.sound']
    >>> "enable_sound" in [p.name for p in wf.inputs]
    False

    The workflow resolves to the sound blocks, and the resolved workflow does not list enable_sound as an input, which is correct: none of the resolved blocks reads it, the workflow entry has already decided it. get_execution_blocks works the same way, get_execution_blocks(prompt=True, enable_sound=True) resolves to the same sound blocks, since get_workflow is just get_execution_blocks called with the workflow entry. If enable_sound were not a workflow input, get_workflow would resolve the wrapper with the default and the input would be lost together with it, see [modular] conditional block triggers should be workflow inputs #15100.

  3. Declare it on every leaf. LTX-2's enable_prompt_enhancement is read only by LTX2AutoPromptEnhancerStep.select_block, but it is declared as an InputParam on each of the three prompt-enhancer leaves it chooses between:

    class LTX2PromptEnhancerStep(ModularPipelineBlocks):
        @property
        def inputs(self):
            return [
                InputParam.template("prompt", required=True),
                InputParam("enable_prompt_enhancement", type_hint=bool, default=False, description="..."),
                ...,
            ]

    The leaves carry an input they never read, each new leaf has to repeat the declaration, and a resolved workflow that keeps such a leaf lists the input although nothing on that workflow acts on it.

I think the Cosmos 3 pattern is the right one: the wrapper is the block that reads the input, so the wrapper declares it, with a type, a default and a description like any other input, so it shows up in inputs, in the generated docstring and in default_call_parameters. The framework should make this the documented way instead of leaving it to each blockset to work out. The Cosmos 3 inputs override is the shape of it; the base class could offer a small hook for it so the override does not have to re-merge the sub-blocks' inputs, but that is a detail.


Posted by Claude Code on behalf of @yiyixuxu, who reviewed and approved this text.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions