flashdreams-integrations

# flashdreams integration architecture A map of how `flashdreams/` is organized and how a single rollout flows through the framework. Read once before adding a built-in recipe under `flashdreams/flashdreams/recipes/`, adding a workspace integration under `integrations/`, or restructuring an existing one. Keep docstrings consistent with the `python-docstring-style` skill. > For the **step-by-step procedure** of doing an integration end-to-end (scope → scaffold → recipe → checkpoint remap → conditioners → runner → verify → perf), see the **`integrate-a-model`** skill. This skill is the *map*; that one is the *route*. > **The fastest way to learn this codebase is to clone the structure of `flashdreams/flashdreams/recipes/template/`.** It is the reference integration — every contract this skill describes is wired up there in its minimal form. Skim it side-by-side with this document. ## TL;DR - Three layers, strict dependency direction: `core` -> `infra` -> recipes/integrations. `infra` and `core` never import from `integrations`. Recipes and integrations may reuse sibling recipe code when the dependency stays explicit and model-owned. - An integration = a `Pipeline` that owns a `DiffusionModel` + optional `Encoder` / `StreamingDecoder`. The `DiffusionModel` owns a `Transformer` + a `Scheduler`. You author the integration-specific subclasses of these and ship one **module-level literal** `StreamInferencePipelineConfig` per variant in `config.py`. No `build_*(...)` factories — variants derive from a base via `derive_config(BASE, ...)`. - Every config sets `name: str` (a stable slug). The per-integration `<NAME>_CONFIGS: dict[str, StreamInferencePipelineConfig]` dict is keyed by `name`. There is no central pipeline-config registry — pipelines are reachable via direct per-integration imports. Integrations that ship a `Runner` (see §5) self-register their slugs into `flashdreams.configs.registry._SUPPORTED_RUNNERS` (read it via `supported_runners()`), which is the only registry the `flashdreams-run` CLI dispatches over. - Per-rollout state lives in nested `*Cache` dataclasses that mirror the same containment tree. - Lifecycle: `pipeline.initialize_cache(...)` once, then a loop of `pipeline.generate(ar_idx, ...)` + `pipeline.finalize(ar_idx, ...)`. - Two shape regimes, separated by `transformer.patchify_and_maybe_split_cp`: pre-patchify `[B, C, T, H, W]` outside, post-patchify `[B, L/cp, C]` inside. ## 1. Codebase layout ``` flashdreams/ ├── core/ reusable numerical primitives (no integration-specific code, no infra deps) ├── infra/ framework contracts + orchestration (ABCs, base configs, pipeline glue) └── integrations/ concrete model bindings that satisfy the infra contracts ``` | Layer | Owns | Imports from | |----------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------| | `core/` | `attention/` (`NativeAttention`, `RingAttention`, `BlockKVCache`, `RotaryPositionEmbedding3D`, `apply_rope_freqs`), `checkpoint/load.py`, `distributed/` (`split_inputs_cp`, `cat_outputs_cp`, `*_object_list`), `io/` | nothing in flashdreams | | `infra/` | `config` (`InstantiateConfig`, `derive_config`), `pipeline` (`StreamInferencePipeline*`), `diffusion.{model, scheduler, transformer}` (ABCs + base impls), `encoder` (`Encoder` + `StreamingEncoder` + `StreamingVideoEncoder` + `NullEncoder`), `decoder` (`StreamingDecoder` + `StreamingVideoDecoder`), `compile`, `cuda_graph`, `profiler` | `core` | | `integrations/<name>/` | concrete model: `transformer/`, optional `encoder.py` / `decoder.py` / `pipeline.py`, `config.py` builders | `core`, `infra` | ### Where does this code go? | Question | Layer | |---------------------------------------------------------|------------------------------------| | New attention kernel or shared CUDA utility | `core/` | | Reusable text/CLIP encoder any integration could use | `infra/encoder/<kind>/` | | New ABC or generic orchestrator | `infra/` | | Model-specific DiT, control encoder, or VAE | `integrations/<name>/` | | CLI runner config + ``run()`` body | `integrations/<name>/runner.py` | If you're tempted to add an integration-specific branch in `infra/` or `core/` — expose a config slot or override hook instead. ## 2. What a pipeline contains The whole framework is built around three nested objects: pipeline, diffusion model, transformer. Each layer (a) holds the next layer down and (b) holds a per-rollout cache that mirrors the same shape. ```mermaid flowchart TB classDef infra fill:#d6f5d6,stroke:#2e7d32,color:#1b5e20 classDef integration fill:#fddede,stroke:#c62828,color:#b71c1c subgraph PIPE["StreamInferencePipeline (infra)"] ENC["encoder<br/>(per-AR-step, optional)"]:::integration DM["DiffusionModel (infra)"]:::infra DEC["decoder<br/>(optional)"]:::integration end subgraph DM_INNER[" "] TF["YourTransformer (integration)<br/>+ context_encoder slot<br/>+ rope_adapter (built per rollout)"]:::integration SCH["Scheduler (infra)<br/>FlowMatch / UniPC"]:::infra end subgraph TF_INNER[" "] NET["YourDiT network (integration)<br/>blocks + KV caches"]:::integration end PIPE --> DM DM --> TF DM --> SCH TF --> NET ``` **Containment, top-down:** - `StreamInferencePipeline` (use as-is in most cases) - `encoder: StreamingEncoder | None` (optional; per-AR-step control like HDMap, camera, first-frame VAE) - `diffusion_model: DiffusionModel` - `transformer: YourTransformer` ← you write this - `network: YourDiT` ← you write this - `context_encoder: Encoder` (one-shot encoder slot — text / CLIP-image / `NullEncoder`) - `rope_adapter: RotaryPositionEmbedding3D` (built per rollout, lives on the cache) - `scheduler: FlowMatchScheduler | UniPCScheduler` (pick from `infra.diffusion.scheduler`) - `decoder: StreamingDecoder | None` (optional; latent → pixels). Use `StreamingVideoDecoder` when the decoder is a pixel-video VAE. **The per-rollout cache mirrors that tree** (`StreamInferencePipelineCache` → `transformer_cache` → `network_cache`). Each level forwards `before_update` / `after_update` to the level below. ### One-shot context vs per-AR-step control input There are **two encoder slots**, and they take different base classes. Confusing them is the most common pitfall. | Slot | Runs | Base class | Input | Disable | |-----------------------------------------------|-------------------------------------|-----------------------------|------------------------------------|---------------------| | `transformer.context_encoder` (one-shot) | once, in `initialize_autoregressive_cache` | `Encoder` (stateless) | text prompts, reference image | `NullEncoderConfig()` | | `pipeline.encoder` (per-AR-step) | every AR step, in `pipeline.generate` | `StreamingEncoder` (stateful, has cache) | per-step control (HDMap, camera, hand-crafted control latent) | `encoder=None` | Text encoders (subclass `Encoder`) go on `context_encoder`. Per-AR-step controls (subclass `StreamingEncoder`) go on `pipeline.encoder`. Putting a text encoder on the per-AR-step slot reruns it every step; putting a streaming encoder on the one-shot slot drops its cache. The decoder slot (`pipeline.decoder`) takes a `StreamingDecoder` (stateful, `forward(input, ar_idx, cache)`). Use `StreamingVideoDecoder` for pixel-video VAEs (WAN VAE, TAEHV) — it adds the spatial / temporal compression contracts the pipeline needs to size pixel I/O. Stateless decoders just return an empty `StreamingDecoderCache` from `initialize_autoregressive_cache` and ignore `autoregressive_index` / `cache` in `forward` (see `template/decoder.py`). **Where the per-AR-step control tensor flows.** This is the path a new control input (HDMap, camera trajectory, ...) takes through the framework. Defining a new control = author one `StreamingEncoder` subclass under `integrations/<name>/encoder.py` and consume the `control` arg inside your network's forward. ``` user passes raw control as `pipeline.generate(ar_idx, cache, input=hdmap)` │ [B, C_ctrl, T, H, W] ▼ pipeline.encoder.forward(input, ar_idx, cache.encoder_cache) ← integrations/<name>/encoder.py │ [B, C_latent, T, H, W] (still pre-patchify; same T/H/W as the noisy latent) ▼ diffusion_model.generate(ar_idx, transformer_cache, input=encoded) │ ├── transformer.patchify_and_maybe_split_cp(encoded) │ [B, L/cp, C] │ └── scheduler loop: transformer.predict_flow(noisy, t, cache, input=patchified_control) └── network.forward(noisy, ..., control=patchified_control) └── x = input_proj(noisy) + input_proj(control) # additive bias ``` Two corollaries: - **The encoder's output shape must match the noisy latent's pre-patchify shape** so the same `patchify_and_maybe_split_cp` call works on both, and so the network can fuse them as an additive bias on the per-token channel dim. - **`encoder=None` round-trips `input=None` end-to-end.** Your network's `forward` should treat `control=None` as "skip the control bias" — `flashdreams/flashdreams/recipes/template/transformer/network.py` is the reference. This lets the same integration support both controlled and uncontrolled rollouts without a separate config. ## 3. Anatomy of an integration A minimum viable built-in recipe (what `flashdreams/flashdreams/recipes/template/` ships) is **3 files and 4 classes**: ``` flashdreams/flashdreams/recipes/<name>/ ├── transformer/ │ ├── __init__.py YourTransformerConfig + YourTransformerCache + YourTransformer │ └── network.py YourDiTConfig + YourDiTCache + YourDiT └── config.py module-level literal StreamInferencePipelineConfig per variant <NAME>_CONFIGS: dict[str, StreamInferencePipelineConfig] ``` Add files only when you actually need them: | File | When to add | |---------------------|---------------------------------------------------------------------| | `encoder.py` | integration needs a per-AR-step control input | | `decoder.py` | integration owns the latent → pixel stage | | `pipeline.py` | rare — only when `pipeline.initialize_cache(...)` needs a custom signature (e.g. derive per-rollout `(height, width)` from an input image, accept text strings instead of pre-encoded embeddings) | | `transformer/impl/` | network is large enough to split (`modules.py`, `network.py`, ...) | | `config/` | many shipped variants — split `config.py` into a package |