integrate-a-model

# Integrate a model into flashdreams The ordered procedure for binding an external video model to the flashdreams framework. Read the **`flashdreams-integrations`** skill first for the architecture (layers, contracts, the cache tree) — this skill is the *route*, that one is the *map*. **Reference examples:** `integrations_v2/wan22/` for the Wan 2.2 TI2V-5B recipe and checkpoint remap, and `integrations_v2/lingbot/` for camera conditioning. Read them side-by-side. Match `python-docstring-style`. ## The core bet: reuse, don't re-implement Most modern video models are DiT-family. Before writing anything, find the closest existing flashdreams recipe (`integrations_v2/wan22`, `wan21`, `self_forcing`, …) and **subclass it**. Keep model-specific conditioner deltas in small subclasses. If your model maps onto an existing backbone, the job is *config + checkpoint remap + deltas + verify*, which is days–weeks. If it needs a novel network/attention/inference loop, it is much longer — say so up front. --- ## Phase 0 — Scope (½–2 days; do this before promising a timeline) **First pick the integration lane** — the `integrations/` directory has several, and they differ a lot in effort: | Lane | What it is | Examples | Effort | |---|---|---|---| | **Config-only recipe** | just `config.py` literals over an existing backbone; no new runner | `wan22` | smallest | | **Runner plugin** | recipe + a `flashdreams-run` runner (+ model deltas) | `flashdreams/flashdreams/recipes/template/runner.py` | small–medium | | **Serving adapter** | adds serving/runtime surfaces on top of a runner | `lingbot` | medium | | **Full native port / builder variants** | real builder helpers, dynamic-resolution variants, a network ported from scratch | `flashvsr` | largest | Then answer these from the upstream repo + model card, and write the answers down: 1. **Backbone family.** Is it a Wan/DiT variant? Diffusion-transformer? → which existing recipe is the closest base. (Decisive for the estimate.) 2. **Checkpoints.** What does upstream publish — native `.pth`/safetensors, a diffusers port, sharded or single-file? Note the HF repo ids. (Drives the remap; see Phase 3.) 3. **Inference shape.** Steps (distilled? e.g. 4-step Euler), scheduler, guidance, resolution, AR/streaming vs one-shot, KV cache. 4. **Conditioners / deltas.** What does it add beyond the base backbone (camera, action, memory, control)? Each is a subclass + (usually) extra checkpoint keys. 5. **Reference for parity.** Can you run upstream to get a ground-truth output to diff against? (You need this for Phase 6.) Output: a one-paragraph scope note + the "closest base recipe" decision. If the answer to (1) is "novel architecture", flag it — the rest of this playbook still applies but Phase 2/4 grow a lot. ## Phase 1 — Scaffold the plugin (pick in-tree or out-of-tree) The package layout is the same either way; only *where it lives* and *how its version is managed* differ. The discovery seam for both is `flashdreams/plugins/registry.py`: runners are found via the `flashdreams.runner_configs` entry point (group `ENTRY_POINT_GROUP`), or the `FLASHDREAMS_RUNNER_CONFIGS` env var during dev. The package body is identical to either reference below. ``` <pkg>/ ├── __init__.py ├── config.py # static PIPELINE_<NAME> + RUNNER_<NAME> + <NAME>_CONFIGS literals ├── runner.py # RunnerConfig + Runner.run() └── _*.py # model-specific subclasses (encoder/transformer/network) tests/ ├── test_smoke.py # ci_cpu: import + static-config assertions └── parity_check/ # GPU parity harness (gitignored heavy deps) ``` **Lane A — in-tree (`integrations/<name>/`)**, for upstreaming into flashdreams (mirror `integrations_v2/wan22/`): - The repo-root `integrations/*` glob auto-adds it to the uv workspace. - `pyproject.toml` `version` must match `flashdreams._version.__version__`; the `sync-version` pre-commit hook enforces it (CI fails otherwise). - `[project.entry-points."flashdreams.runner_configs"]` maps slug → config (see the runner-plugin template): ```toml [project.entry-points."flashdreams.runner_configs"] "my-model-slug" = "my_model.config:RUNNER_MY_MODEL" ``` **Lane B — out-of-tree (your own pip-installable repo)**, the supported path for external contributors who don't want to land in flashdreams. Same package body; standalone `pyproject.toml` that just depends on `flashdreams` and exposes the same entry point: ```toml [project] name = "my-model-flashdreams" dependencies = ["flashdreams"] # no version-sync constraint here [project.entry-points."flashdreams.runner_configs"] "my-model-slug" = "my_model.config:RUNNER_MY_MODEL" ``` `pip install -e .` and `flashdreams-run my-model-slug` discovers it via the entry point — no fork of flashdreams needed. During development before install, point at it without an entry point via `FLASHDREAMS_RUNNER_CONFIGS="my-model-slug=my_model.config:RUNNER_MY_MODEL"`. ## Phase 2 — Recipe config (subclass the base, ship a static literal) In `config.py`, `copy.deepcopy` the closest base pipeline and swap the pieces that differ — encoder / transformer.network / scheduler — into model-specific subclasses. Ship **one module-level literal** `PIPELINE_<NAME>` (no `build_*` factories for the config-only / runner-plugin lanes; the full-native-port lane like `flashvsr` uses real builder helpers for dynamic-resolution variants — see Phase 0) + a `RUNNER_<NAME>` literal + a `<NAME>_CONFIGS` dict keyed by `name`. See the `PIPELINE_WAN22_TI2V_5B` literal in `integrations_v2/wan22/config.py` for the base config. - Subclass `Wan21TransformerConfig` / the network / encoder configs; copy field-by-field so a future base-class field addition surfaces loudly instead of silently dropping. - Set the standard transformer knobs (`len_t`, `window_size_t`, `guidance_scale`, `stamp_image_latent`, …) — see `flashdreams-integrations` §"Standard transformer knobs". - Distilled models: swap the scheduler (e.g. a 4-step `FlowMatchEulerDiscreteScheduler`). ## Phase 3 — Checkpoint loading + key remap (the highest-leverage phase) Upstream weights almost never match flashdreams key names. You write a `state_dict_transform` (regex rename) consumed by the transformer/VAE config. **Prefer the native checkpoint over a diffusers port when both exist.** flashdreams' networks are typically ported from the *native* model, so native keys often match 1:1 (`Wan-AI/Wan2.2-TI2V-5B` native keys = `WanDiTNetwork` keys exactly → **zero** remap, the transform is `lambda sd: sd`; the diffusers port needs ~25 rules). The native VAE needed only 4 rules vs the diffusers ~50. Note the native checkpoint can be **either** a single-file `.pth` **or** sharded safetensors + a `.safetensors.index.json` (the Wan native DiT is the latter, at the repo root; its VAE is a nested `.pth`) — `load_checkpoint` resolves both. Fast pre-check before any set-diff: do the key *counts* even match? (825 == 825 → you likely picked the right source.) **If you must remap (the diffusers port), the renames cluster into a few families.** From the Wan diffusers→native mapping, expect: `attn1.*`→`self_attn.*`, `attn2.*`→`cross_attn.*`, `to_q/to_k/to_v`→`q/k/v`, `to_out.0`→`o`, `condition_embedder.{text,time}_embedder.linear_{1,2}`→`{text,time}_embedding.{0,2}`, `condition_embedder.time_proj`→`time_projection.1`, `ffn.net.0.proj`/`ffn.net.2`→ `ffn.0`/`ffn.2`, `norm2`→`norm3`, `scale_shift_table`→`modulation` (per-block) / `head.modulation` (top), `proj_out`→`head.head`. Write them as ordered regex rules and let unmatched keys fall through (they show up as `unexpected_keys`, which the bijection check below catches). **Verify the remap is a key/shape bijection on CPU — no GPU needed.** This is the single most valuable check. Build the model on `meta` and diff against the checkpoint; any model key the transform doesn't supply stays on `meta` and `.to(device)` later raises "Cannot copy out of meta tensor". Your `state_dict_transform` takes a `{name: tensor}` dict (it renames keys, tensors ride along), so feed it a **zero-memory stand-in**: real key names, `meta` tensors of the real shapes (read from the safetensors headers without loading weights). This runs the *actual* transform and costs no memory: ```python import json, torch from safetensors import safe_open from my_model.config import my_state_dict_transform # the real transform you wrote with torch.device("meta"): net = MyNetworkConfig().setup() model = {k: tuple(v.shape) for k, v in net.state_dict().items()} raw = {} # {name: meta tensor}, no weights index = json.load(open(f"{ckpt_dir}/diffusion_pytorch_model.safetensors.index.json")) for shard in set(index["weight_map"].values()): with safe_open(f"{ckpt_dir}/{shard}", framework="pt") as f: for k in f.keys(): raw[k] = torch.empty(f.get_slice(k).get_shape(), device="meta") ckpt = {k: tuple(v.shape) for k, v in my_state_dict_transform(raw).items()} missing = set(model) - set(ckpt) # would stay on meta — must be empty extra = set(ckpt) - set(model) # unexpected keys — must be empty shapemm = [k for k in model if k in ckpt and model[k] != ckpt[k]] assert not missing and not extra and not shapemm, (missing, extra, shapemm) ``` (For a single-file `.pth`: `raw = torch.load(path, map_location="meta", weights_only=True)` gives the `{name: tensor}` dict directly; skip the safetensors loop.) Codify it as a `ci_cpu` test (`test_*_remap_is_full_bijection`) + spot-checks against real key strings (`test_*_remap_spot_checks_real_keys`). **Before flipping a default checkpoint source, prove weight-equality.** If you switch the production config to a different checkpoint (e.g. native `.pth` instead of diffusers), load *both*, apply each transform, and assert every tensor matches (`max |Δ| == 0`). Identical weights ⇒ identical output, no decode smoke needed. This is how the VAE/DiT defaults were flipped safely (`test_*_weights_identical`, marked `manual` since it downloads checkpoints). **Pitfall — "missing params" is usually a naming mismatch, not absent weights.** If a load fails with missing keys, diff the *names* first; the weights are almost always present under a different convention. ## Phase 4 — Model-specific conditioners / deltas Each delta = a subclass + (usually) extra checkpoint keys. Keep action, camera, and memory conditioners local to the integration. Conventions that make these parity-safe: - **Zero-init new residual heads** so the conditioner is a strict identity until trained weights load (`nn.init.zeros_(head.weight)`). The un-conditioned pipeline then matches the base model exactly. - **Tolerate the extra zero-init keys when loading a base checkpoint** that lacks them. Override `load_state_dict` on the network to allow *exactly* those keys missing (keep it strict for everything else). Without this, a base/un-distilled load raises `Missing key(s)`. - Keep model deltas in the integration — never branch `core/` or `infra/`; expose a config slot or override hook instead. ## Phase 5 — Runner + CLI `runner.py` ships a `RunnerConfig` subclass (I/O fields: image/prompt/output, ckpt override, knobs) + a `Runner` whose `run()` drives `initialize_cache` → per-AR-step `generate`/`finalize` → decode → write mp4. Mirror `flashdreams/flashdreams/recipes/template/runner.py`. Thread an optional `--ckpt-path` through `derive_config` to swap the checkpoint + transform at construction time. Add example-data download helpers if useful for demos. ## Phase 6 — Verify (CPU first, then GPU) In order of cost: 1. **`ci_cpu` smoke** (`test_smoke.py`): imports, the static config is fully swapped, runner slug == pipeline name, entry point registered, remap bijection tests. Run: `uv run --extra dev pytest integrations/<name>/tests/test_smoke.py`. 2. **Checkpoint weight-equality** (Phase 3) — proves the load is correct without a GPU. 3. **GPU rollout smoke** — `flashdreams-run <slug> --ckpt-path <distilled> --num-chunk 1` produces a valid mp4. (Use `--ckpt-path`; a base/un-distilled run gives identity-only output. Keep `num_chunk` small to dodge OOM and short-rollout edge cases.) 4. **Upstream parity** — run upstream on the same input/seed, diff decoded frames, report **mean `|Δ|` / 255** against an agreed model-specific threshold. Account for bf16 FP noise; don't chase bit-exactness across two kernel stacks. ## Phase 7 — Perf + model card (the visible deliverable) - Bench native vs upstream, **stack-matched** (both cuDNN SDPA + `torch.compile`), at the largest `num_chunk` the GPU allows, discarding warmup chunks. Scope = **DiT + VAE enc/dec**, per-stage medians post-warmup. Harnesses: `tests/parity_check/bench.sh` (matched) / `bench_batch.sh` (native-only sample loop). - Author a model-card page mirroring `docs/source/models/lingbot_world.md` (hero + gallery videos, perf table, methodology); register it in `docs/source/models/index.md`. ## Gotchas (hard-won) - **CI-pinned ruff is the source of truth** — `uvx ruff` defaults to a newer version that sorts imports differently and touches unrelated files. Use the pinned version (`uvx ruff@<pinned> …`; check `.pre-commit-config.yaml`). - **`ty` needs the real deps** — a torch-less env can't catch signature/None errors; CI's `cpu` job (full deps) is the real type check. Fix diagnostics, don't `# ty: ignore` what is fixable; remove `ty: ignore` once unneeded (CI flags unused ones). - **`uv sync`/`uv run` builds `block-sparse-attn`** (CUDA ext) → needs `CUDA_HOME`. On a GPU box, use a synced venv; on CPU, run modules with `PYTHONPATH` against a venv that already has torch. - **`expandable_segments:True` breaks CUDA graphs** — scope it to non-graph legs only. - **First AR chunk's `diffuse` time is cold `torch.compile` autotune**, not steady-state — that's why bench discards warmup chunks. - **Diffusers single-file URLs may 404** if the repo is actually sharded — point at the `.safetensors.index.json`; `load_checkpoint` resolves shards from it. - **Keep heavy/scratch out of git** — checkpoints, vendor trees, bench outputs, handoff notes (gitignore them). ## Done criteria - [ ] `ci_cpu` smoke + remap-bijection tests pass. - [ ] Checkpoint weight-equality proven (or remap bijection + a GPU decode smoke). - [ ] GPU rollout produces a valid mp4. - [ ] Upstream parity `mean |Δ|` under the agreed bar. - [ ] Runner registered; `flashdreams-run <slug> --help` works. - [ ] Perf numbers + model-card page (if in scope). - [ ] lint/`ty` green under the CI-pinned tools. ## Evaluating this skill