From 12827a56bad89d81a3014dae56175e5162eb8bbe Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 05:07:39 +0000 Subject: [PATCH] Add tests and agent docs for kwargs_type input delivery in modular pipelines Document how kwargs_type-tagged values flow from block outputs and user inputs to consumer blocks (the mechanism behind denoiser_input_fields), pin the behavior down with tests, and add a key-pattern section to .ai/modular.md. Co-Authored-By: Claude Fable 5 --- .ai/modular.md | 31 ++++++ .../test_modular_pipelines_custom_blocks.py | 104 ++++++++++++++++++ 2 files changed, 135 insertions(+) diff --git a/.ai/modular.md b/.ai/modular.md index 46ccd30031b7..3f8997993a68 100644 --- a/.ai/modular.md +++ b/.ai/modular.md @@ -102,6 +102,37 @@ class HeliosChunkDenoiseStep(HeliosChunkLoopWrapper): Note: sub-blocks inside `LoopSequentialPipelineBlocks` receive `(components, block_state, i, t)` for denoise loops or `(components, block_state, k)` for chunk loops. +## Key pattern: `kwargs_type` bags (`denoiser_input_fields`) + +The conditioning inputs a denoiser needs often vary by workflow — especially for omni models like Cosmos3, where the action workflow requires additional action conditioning, and a workflow that generates sound along with video requires additional sound inputs. Tag these outputs with `kwargs_type="denoiser_input_fields"` when they are written, and have the denoiser declare the bag once and receive everything tagged — this avoids creating a new denoiser block for each workflow just to list its specific inputs: + +```python +# producer side: standard conditioning outputs already carry the tag via their templates +OutputParam.template("prompt_embeds") # kwargs_type="denoiser_input_fields" +# workflow-specific fields declare it explicitly +OutputParam( + "action_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Action conditioning fed into the transformer.", +) + +# consumer side (the loop denoiser): declare the bag once +InputParam.template("denoiser_input_fields") + +# inside the denoiser __call__: every tagged value arrives in one dict — +# and also individually (block_state.prompt_embeds, block_state.action_embeds, ...) +block_state.denoiser_input_fields # {"prompt_embeds": ..., "action_embeds": ...} +``` + +The denoiser typically filters the bag against the transformer's forward signature and forwards the matches — so a new block can add conditioning just by tagging its output (no change to the denoiser), and tagged fields the transformer doesn't accept are silently ignored (see `qwenimage/denoise.py` or `helios/denoise.py`; `z_image/denoise.py` is a minimal bag consumer). + +How the tagging works (behavior is pinned down in `tests/modular_pipelines/test_modular_pipelines_custom_blocks.py::TestBlockKwargsTypeInputs`): + +- A value gets its tag when it is **written** to pipeline state: a block output is tagged if declared with `OutputParam(..., kwargs_type=...)`; a user-passed input is tagged if the pipeline-level `InputParam` it matches declares a kwargs_type. +- Users can always pass the whole bag as a dict under the kwargs_type name — `pipe(denoiser_input_fields={"prompt_embeds": ...})` — and every entry gets tagged. +- **Gotcha — standalone runs:** a named input declared *without* the kwargs_type lands in state by name but never reaches the bag. So when a denoise block runs standalone (without the upstream blocks whose tagged outputs normally fill the bag), passing those values as plain named inputs silently does nothing — they must go through the `denoiser_input_fields={...}` dict, or the block must declare them as named `InputParam(..., kwargs_type="denoiser_input_fields")` inputs. + ## Key pattern: Workflow selection ```python diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index 433d1e872a7d..eb7d35061038 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -578,6 +578,110 @@ def test_loop_block_requirements_save_load(self, tmp_path): assert expected_requirements == config["requirements"] +class DummyKwargsProducerStep(ModularPipelineBlocks): + """Takes `a` and `b` as regular named inputs and passes them through (with a `-producer` + suffix so tests can see the values went through this block), writing `a` back as an output + tagged with kwargs_type `typea` and `b` as an untagged output.""" + + @property + def inputs(self) -> List[InputParam]: + return [InputParam(name="a", default=None), InputParam(name="b", default=None)] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [OutputParam("a", kwargs_type="typea"), OutputParam("b")] + + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.a = f"{block_state.a}-producer" + block_state.b = f"{block_state.b}-producer" + self.set_block_state(state, block_state) + return components, state + + +class DummyKwargsConsumerStep(ModularPipelineBlocks): + """Consumes the `typea` bag and the named input `b`, and verifies what it received against + the expected values passed as inputs (`expected_a`, `expected_b`, `expected_typea`).""" + + @property + def inputs(self) -> List[InputParam]: + return [ + InputParam(kwargs_type="typea"), + InputParam(name="b", default=None), + InputParam(name="expected_a", default=None), + InputParam(name="expected_b", default=None), + InputParam(name="expected_typea", default=None), + ] + + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + assert block_state.typea == block_state.expected_typea + assert block_state.b == block_state.expected_b + # values delivered through the bag are also set individually on block_state; + # `a` only exists here if it was delivered through the bag + if block_state.expected_a is None: + assert not hasattr(block_state, "a") + else: + assert block_state.a == block_state.expected_a + return components, state + + +class TestBlockKwargsTypeInputs: + """Test how `kwargs_type` fields flow from user inputs and block outputs to consumer blocks. + + This is the mechanism behind `denoiser_input_fields`: a block declaring + `InputParam(kwargs_type=...)` receives a dict of every state value *tagged* with that + kwargs_type. A value gets its tag when it is written to the pipeline state: an output + written by a block is tagged if the block declared it with + `OutputParam(..., kwargs_type=...)`, and a user-passed input is tagged if the pipeline's + `InputParam` for it declares a kwargs_type. A named input declared *without* a kwargs_type + therefore never reaches the bag, even though it is available in state by name. + """ + + def test_tagged_block_outputs_are_delivered_to_consumer(self): + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"producer": DummyKwargsProducerStep(), "consumer": DummyKwargsConsumerStep()} + ) + pipe = blocks.init_pipeline() + + # `a` goes through the producer and is written back as a tagged output, so it reaches + # the consumer through the bag; `b` also goes through the producer, but its output is + # untagged: it reaches the consumer only as the named input, never through the bag + pipe( + a="testa", + b="testb", + expected_a="testa-producer", + expected_b="testb-producer", + expected_typea={"a": "testa-producer"}, + ) + + def test_user_inputs_passed_by_name_do_not_reach_the_bag(self): + pipe = DummyKwargsConsumerStep().init_pipeline() + + # the consumer only knows `a` through the `typea` bag: passing it by name does not + # reach the block at all. `b` is declared by name without a kwargs_type: it reaches + # the block as the named input, but never through the bag. + pipe( + a="testa", + b="testb", + expected_a=None, + expected_b="testb", + expected_typea={}, + ) + + def test_kwargs_type_dict_input_is_delivered(self): + pipe = DummyKwargsConsumerStep().init_pipeline() + + # the whole bag can be passed as a dict under the kwargs_type name: every entry is + # tagged individually, so `a` now reaches the consumer through the bag + pipe( + typea={"a": "testa"}, + expected_a="testa", + expected_b=None, + expected_typea={"a": "testa"}, + ) + + @slow @nightly @require_torch