From 12827a56bad89d81a3014dae56175e5162eb8bbe Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 10 Jul 2026 05:07:39 +0000 Subject: [PATCH 1/2] 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 From e080c4fca67a95eb9be5d2005bbc2bfa75acee55 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 14 Jul 2026 18:48:10 +0000 Subject: [PATCH 2/2] Address review feedback: drop 'bag' term, assert on returned state, test pipeline call params Co-Authored-By: Claude Fable 5 --- .ai/modular.md | 12 +-- .../test_modular_pipelines_custom_blocks.py | 93 ++++++++++--------- 2 files changed, 57 insertions(+), 48 deletions(-) diff --git a/.ai/modular.md b/.ai/modular.md index 3f8997993a68..31891608be82 100644 --- a/.ai/modular.md +++ b/.ai/modular.md @@ -102,9 +102,9 @@ 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`) +## Key pattern: `kwargs_type` inputs (`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: +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; the denoiser then declares a single input with that `kwargs_type` and receives every tagged value collected into one dict. 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 @@ -117,7 +117,7 @@ OutputParam( description="Action conditioning fed into the transformer.", ) -# consumer side (the loop denoiser): declare the bag once +# consumer side (the loop denoiser): declare the kwargs_type input once InputParam.template("denoiser_input_fields") # inside the denoiser __call__: every tagged value arrives in one dict — @@ -125,13 +125,13 @@ InputParam.template("denoiser_input_fields") 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). +The denoiser typically filters this dict 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 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. +- Users can always pass all the tagged values as a dict under the kwargs_type name — `pipe(denoiser_input_fields={"prompt_embeds": ...})` — and every entry gets tagged. In a full pipeline this is rarely needed: named inputs and tagged block outputs get tagged on their own; the dict form matters mainly for standalone runs (below). +- **Gotcha — standalone runs:** a named input declared *without* the kwargs_type lands in state by name but never gets tagged, so it never reaches the consumer's dict. So when a denoise block runs standalone (without the upstream blocks whose tagged outputs normally supply these values), passing them 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 diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py index eb7d35061038..882dde5ae361 100644 --- a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -600,29 +600,32 @@ def __call__(self, components, state: PipelineState) -> PipelineState: 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`).""" + """Consumes the values tagged `typea` (as the `typea` dict) and the named input `b`, and + records what it received as outputs (`received_*`) so tests can assert on the returned state.""" @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), + ] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [ + OutputParam("received_typea"), + OutputParam("received_a"), + OutputParam("received_b"), ] 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 + block_state.received_typea = block_state.typea + # tagged values delivered through the `typea` dict are also set individually on + # block_state; `a` only exists here if it was delivered as a tagged value + block_state.received_a = getattr(block_state, "a", "") + block_state.received_b = block_state.b + self.set_block_state(state, block_state) return components, state @@ -635,7 +638,7 @@ class TestBlockKwargsTypeInputs: 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. + therefore never reaches the consumer's dict, even though it is available in state by name. """ def test_tagged_block_outputs_are_delivered_to_consumer(self): @@ -645,41 +648,47 @@ def test_tagged_block_outputs_are_delivered_to_consumer(self): 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): + # the consumer through the `typea` dict; `b` also goes through the producer, but its + # output is untagged: it reaches the consumer only as the named input + received = pipe(a="testa", b="testb", output=["received_typea", "received_a", "received_b"]) + assert received["received_typea"] == {"a": "testa-producer"} + assert received["received_a"] == "testa-producer" + assert received["received_b"] == "testb-producer" + + def test_user_inputs_passed_by_name_are_not_tagged(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={}, - ) + # the consumer only knows `a` as a tagged value: 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 `typea` dict. + received = pipe(a="testa", b="testb", output=["received_typea", "received_a", "received_b"]) + assert received["received_typea"] == {} + assert received["received_a"] == "" + assert received["received_b"] == "testb" 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"}, + # tagged values can be passed as a dict under the kwargs_type name: every entry is + # tagged individually, so `a` now reaches the consumer through the `typea` dict + received = pipe(typea={"a": "testa"}, output=["received_typea", "received_a", "received_b"]) + assert received["received_typea"] == {"a": "testa"} + assert received["received_a"] == "testa" + assert received["received_b"] is None + + def test_kwargs_type_input_in_pipeline_call_params(self): + blocks = SequentialPipelineBlocks.from_blocks_dict( + {"producer": DummyKwargsProducerStep(), "consumer": DummyKwargsConsumerStep()} ) + pipe = blocks.init_pipeline() + + # the kwargs_type input is exposed as a single nameless param alongside the named + # inputs, and renders as a `**typea` kwargs-style param in the docstring + named = [inp.name for inp in pipe.blocks.inputs if inp.name is not None] + kwargs_inputs = [inp.kwargs_type for inp in pipe.blocks.inputs if inp.name is None] + assert named == ["a", "b"] + assert kwargs_inputs == ["typea"] + assert "**typea" in pipe.blocks.doc @slow