indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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:
indigo
2026-07-09 09:50:27 +08:00
parent baa6646fd1
commit 110adc781e
17 changed files with 300 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/plan.md
+7 -6
View File
@@ -37,12 +37,13 @@
- 用現有 checkpoint 驗證資料流正確性（序列長度、座標偏移、attention mask、denoise_mask、不 wrap）。
- **PoC 僅驗證 plumbing**；LTX-2 未經此訓練，畫面不會是正確動畫。
### Phase 2 — In-context mask channel（機制 2，checkpoint 手術）⬜ 未開始
- `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（機制 2，checkpoint 手術）✅ 已完成（plumbing）
- **決策**：時間編碼採忠實堆疊，`8×(K+1)=56` channel（LTX 時間因子 8，K=6），非 Wan 的 28。
- `LTXModel(mask_conditioning_channels=0)` config-gated 加寬 `patchify_proj`；預設不變。
- `LatentState`/`Modality` 加 `cond_channels` 欄位，跟著 clone/clear/append 流動；投影前在 `_apply_patchify_proj` concat（None 補零）。
- 新輸入欄位 **zero-init**：`widen_patchify_proj_for_mask_channels` 轉換舊 checkpoint，行為不變、待微調才生效。
- `VideoConditionByMaskChannels`：(K+1) 語意 mask → 空間下採樣 + 時間 8× 堆疊 → 寫入尾端 driving token（target 保持零，符合論文）。
- 驗證通過（`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 與凍結策略。
docs/tasks.md
+10 -6
View File
@@ -15,14 +15,18 @@
## Phase 2 — In-context mask channel（checkpoint 手術）
> **決策**：mask channel 時間編碼採**忠實堆疊**——每 latent 幀對應 8 個 pixel 子幀沿 channel 堆疊，`8×(K+1)=56` channel（LTX 時間因子 8，K=6）。非 Wan 的 4×(K+1)=28。
| # | 任務 | 狀態 | 備註 |
|---|---|---|---|
| 2.1 | `patchify_proj` 加寬 `128 → 128+28` | ⬜ | `model.py:158` `_init_video`；`proj_out` 不動 |
| 2.2 | `LatentState` 帶額外 conditioning-channel 欄位 | ⬜ | 同步改 `tools.py` patchify/unpatchify/clear、`Modality` |
| 2.3 | 投影前 concat mask channel 到 `x` | ⬜ | `transformer_args.py:209` |
| 2.4 | checkpoint zero-init 轉換 script | ⬜ | 新輸入欄位 zero-init，載入舊權重行為不變 |
| 2.5 | 產生 28-channel mask 的 conditioning item | ⬜ | 環境開關 + K=6 角色綁定槽，`4(K+1)=28` |
| 2.6 | Replacement Mode z_ref 高度位移 ΔH_ref | ⬜ | Phase 1 暫緩項 |
| 2.1 | `patchify_proj` 加寬（config-gated） | ✅ | `LTXModel(mask_conditioning_channels=0)`；>0 時 `patchify_proj=Linear(in+mask, inner)`。預設不變 |
| 2.2 | `LatentState`/`Modality` 帶 `cond_channels` 欄位 | ✅ | `types.py`、`modality.py` 加 `cond_channels: Tensor\|None=None`（patchified [B,T,C]）；`tools.clear_conditioning` 裁切；`helpers.modality_from_latent_state` 帶入 |
| 2.3 | 投影前 concat/zero-pad cond_channels | ✅ | `transformer_args.py` `_apply_patchify_proj`：寬度不足時用 cond_channels 補齊，None 則補零 |
| 2.4 | checkpoint zero-init 加寬轉換 | ✅ | `mask_channels_checkpoint.py` `widen_patchify_proj_for_mask_channels`：尾端補零欄，載入舊權重行為不變 |
| 2.5 | `VideoConditionByMaskChannels`（56ch） | ✅ | `mask_channels_cond.py`：1 環境開關 + K=6 綁定槽 → 空間下採樣 + 時間 8× 堆疊(causal 首幀複製)=56ch，寫入尾端 driving token，target 保持零 |
| 2.6 | token-append conditioning 延伸 cond_channels | ✅ | `cond_channels.py` `extend_cond_channels`；reference_video/reference_audio/driving/keyframe 皆接上 |
| 2.7 | Phase 2 驗證（免訓練） | ✅ | `verify_mask_channels.py`：mask=0 向後相容；zero-init 加寬 forward == baseline（任意 cond_channels）；driving+mask pipeline forward 不 crash、cond_channels 形狀/placement 正確、`clear_conditioning` 剝除 |
| 2.8 | Replacement Mode z_ref 高度位移 ΔH_ref | ⬜ | 仍暫緩（Phase 1 divergence，需獨立 reference token group） |
## Phase 3 — 訓練整合（ltx-trainer）