Add SCAIL-2 in-context mask channels (Phase 2, plumbing + zero-init surgery)

Port mechanism 2 of SCAIL-2 (arXiv:2606.10804) to LTX-2: extra per-token
in-context conditioning channels (1 environment switch + K=6 character binding
slots) concatenated onto the latent before the first projection.

Faithful temporal encoding for LTX's VAE (temporal factor 8): each latent frame
stacks its 8 pixel sub-frames along the channel dim, giving 8*(K+1)=56 channels
(vs the paper's 4*(K+1)=28 on Wan 2.1).

Plumbing (backward compatible; cond_channels=None leaves every existing pipeline
unchanged):
- LatentState/Modality gain an optional cond_channels field (patchified [B,T,C]).
- LTXModel(mask_conditioning_channels=0) config-gates a widened patchify_proj;
  TransformerArgsPreprocessor concatenates cond_channels (or zero-pads) before it.
- tools.clear_conditioning trims it; token-appending conditioning items
  (reference video/audio, driving, keyframe) extend it via extend_cond_channels.

New:
- VideoConditionByMaskChannels: encodes (K+1) pixel masks -> 56 channels, written
  onto the trailing driving tokens (noisy target stays all-zero, per the paper).
- widen_patchify_proj_for_mask_channels: zero-init checkpoint surgery so a
  converted model reproduces the base output exactly until finetuned.

Verified (CPU, random weights): backward compat, zero-init widened forward is
bit-identical to baseline for any cond_channels, and the driving+mask pipeline
forwards without crashing with correct placement/clipping. Visual quality
requires Phase 3 finetuning; Replacement-mode z_ref height shift still deferred.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-07-09 09:50:27 +08:00
parent baa6646fd1
commit 110adc781e
17 changed files with 300 additions and 15 deletions
+7 -6
View File
@@ -37,12 +37,13 @@
- 用現有 checkpoint 驗證資料流正確性(序列長度、座標偏移、attention mask、denoise_mask、不 wrap)。
- **PoC 僅驗證 plumbing**LTX-2 未經此訓練,畫面不會是正確動畫。
### Phase 2 — In-context mask channel(機制 2checkpoint 手術)⬜ 未開始
- `model.py:158` `patchify_proj``Linear(128, inner)` 加寬到 `Linear(128+28, inner)`
- `LatentState` 增加 optional conditioning-channel 欄位,跟著 patchify/concat/clear 流動
- 投影前 concat`transformer_args.py:209`)。
- 新輸入欄位 **zero-init**,載入舊 checkpoint 行為不變;寫 checkpoint 轉換 script
- 新增產生 28-channel mask(環境開關 + 角色槽)的 conditioning item
### Phase 2 — In-context mask channel(機制 2checkpoint 手術)✅ 已完成(plumbing
- **決策**:時間編碼採忠實堆疊,`8×(K+1)=56` channelLTX 時間因子 8K=6),非 Wan 的 28
- `LTXModel(mask_conditioning_channels=0)` config-gated 加寬 `patchify_proj`;預設不變
- `LatentState`/`Modality``cond_channels` 欄位,跟著 clone/clear/append 流動;投影前在 `_apply_patchify_proj` concatNone 補零)。
- 新輸入欄位 **zero-init**`widen_patchify_proj_for_mask_channels` 轉換舊 checkpoint,行為不變、待微調才生效
- `VideoConditionByMaskChannels`(K+1) 語意 mask → 空間下採樣 + 時間 8× 堆疊 → 寫入尾端 driving tokentarget 保持零,符合論文)
- 驗證通過(`verify_mask_channels.py`):向後相容、zero-init 加寬 forward == baseline、mask pipeline 不 crash。**僅驗證 plumbing,畫質需 Phase 3 微調。**
### Phase 3 — 訓練整合(ltx-trainer)⬜ 未開始
- dataset 產出 (target, driving, mask),接上 Phase 1/2 conditioning,設微調 loss 與凍結策略。