Skip to content
25 changes: 25 additions & 0 deletions .ai/modular.md
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,10 @@ Does it choose ONE block based on which input is present?
Is the selection 1:1 with trigger inputs?
YES -> AutoPipelineBlocks (simple trigger mapping)
NO -> ConditionalPipelineBlocks (custom select_block method)

Is it a different CHECKPOINT (distilled / turbo / a variant with its own schedule)?
YES -> its own blocks assembly, unless it behaves literally the same

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"its own block assembly" reads a bit confusing to me. Possible reword?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it may also be nice to use one word for "block assembly" and "blockset" so it can standardize on one term

(see Key pattern: Checkpoint variants)
```

## Build order (easiest first)
Expand All @@ -67,6 +71,17 @@ Does it choose ONE block based on which input is present?
3. `before_denoise.py` -- Timesteps, latent prep, noise setup. Each logical operation = one block
4. `denoise.py` -- The hardest. Convert guidance to guider abstraction

## Growing a pipeline: one workflow at a time

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Possible to also link an example here?


Build one workflow end-to-end first (e.g. t2v), then add the next workflow — and later the next blocks assembly / checkpoint variant — one at a time. Each addition should **introduce new blocks rather than modify existing ones**: existing blocks are already wired into working workflows, and a new leaf block plus a new assembly entry can't break them.

The one good reason to touch an existing block is to make it strictly more **general**. Be honest about which direction the edit goes:

- **Adding a branch is not generalizing — it's specializing.** `if components.config.foo:` or `if block_state.image is not None:` inside an existing block means the block now does two things. Add a new block for the new case instead, and let workflow selection or a variant assembly pick between them.
- **Collapsing duplicates is generalizing.** If two blocks are identical except for which conditioning inputs they pass to the denoiser, don't keep both — rework the one block to take `kwargs_type="denoiser_input_fields"` (see the `kwargs_type` pattern below) so the same block serves every workflow, as the Cosmos3 denoise step does.

Rule of thumb: a generalizing edit *removes* an if/else or a duplicate block. If your edit *adds* an if/else, it's a new block trying to get out.

## Key pattern: Guider abstraction

Original pipeline has guidance baked in:
Expand Down Expand Up @@ -166,6 +181,14 @@ class AutoDenoise(ConditionalPipelineBlocks):
default_block_name = "text2video"
```

## Key pattern: Checkpoint variants

A different checkpoint (distilled / turbo / a variant with its own schedule) can have its own blocks assembly mapped to it: give the variant a `ModularPipeline` subclass carrying its `default_blocks_name`, and checkpoints route to it automatically — via `_class_name` in `modular_model_index.json`, or, for repos that only ship a standard `model_index.json`, a config-keyed map fn in `MODULAR_PIPELINE_MAPPING` (see `_flux2_klein_map_fn`).

Default to taking that option. The only reason not to is when the variant behaves literally the same; if the split buys anything at all — the distilled variant doesn't have to declare `negative_prompt`, doesn't carry a guider, its docs describe exactly what the checkpoint does — make the separate assembly. It costs almost nothing: assemblies compose the same shared leaf blocks, and only the steps that truly differ need new block classes. See `modular_blocks_flux2_klein.py`, which reuses the base flux2 leaf blocks and swaps in just a `negative_prompt`-free text encoder and a guider-free denoise step.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Default to taking that option. The only reason not to is when the variant behaves literally the same; if the split buys anything at all — the distilled variant doesn't have to declare `negative_prompt`, doesn't carry a guider, its docs describe exactly what the checkpoint does — make the separate assembly. It costs almost nothing: assemblies compose the same shared leaf blocks, and only the steps that truly differ need new block classes. See `modular_blocks_flux2_klein.py`, which reuses the base flux2 leaf blocks and swaps in just a `negative_prompt`-free text encoder and a guider-free denoise step.
Default to taking that option. The only reason not to split is when the variant behaves literally the same. If the split buys anything at all — the distilled variant doesn't have to declare `negative_prompt`, doesn't carry a guider, and its docs describe exactly what the checkpoint does — make the separate assembly. It costs almost nothing: assemblies compose the same shared leaf blocks, and only the steps that truly differ need new block classes. See `modular_blocks_flux2_klein.py`, which reuses the base flux2 leaf blocks and swaps in just a `negative_prompt`-free text encoder and a guider-free denoise step.


Don't fall back to the standard-pipeline habit of a config flag branching inside a shared block (`ConfigSpec(name="is_distilled")` + `if components.config.is_distilled:`). That keeps both variants' behavior bundled in one blockset — and the input surface is the one thing it can never fix: a repo can override components and config values per checkpoint, but never which inputs the blocks declare, so the distilled checkpoint would still accept `negative_prompt` and silently ignore it.

## Key pattern: Standalone block reusability

One of the core reason a pipeline is split into blocks at all: each block (text encoder, VAE encoder, prepare-latents, denoise, decoder) must be runnable on its own, and its output must be reusable as the input to a different downstream chain.
Expand Down Expand Up @@ -249,6 +272,8 @@ ComponentSpec(

8. **No-op skip logic inside an optional block.** If a step is conditional (e.g. an optional prompt enhancer), don't have the block check a flag at the top of `__call__` and `return` early. Wrap it in an `AutoPipelineBlocks` with `block_trigger_inputs = ["use_xxx"]` so the block is only assembled into the pipeline when the trigger input is provided. The block's own `__call__` should always assume its components and inputs are present.

9. **Serving a checkpoint variant through a config flag in a shared block.** `ConfigSpec(name="is_distilled")` plus `if components.config.is_distilled:` bundles two checkpoints' behavior into one blockset — and it can't change the input surface at all (the distilled variant would still accept `negative_prompt`). Suggest a separate blocks assembly for the variant instead (see Key pattern: Checkpoint variants).

## Conversion checklist

- [ ] Read original pipeline's `__call__` end-to-end, map stages
Expand Down
Loading