7 Commits

Author SHA1 Message Date
indigo 5594d49c76 Add char_masks preprocessing script for SCAIL-2 training (task 3.7)
scripts/process_char_masks.py turns per-sample character label-map videos/images
(integer pixel labels: 0 = environment, 1..K = characters -> binding slots) into
the pixel-space semantic-mask tensors the SCAIL training/inference path consumes:
{"mask": [K+1, F_pix, H_pix, W_pix]} (ch0 = environment switch, ch1..K = slots).

- Aligns to the target video's latent grid read from the saved latent metadata
  (F_pix=(F-1)*8+1, H*32, W*32), so char_masks/ lines up file-for-file with
  latents/ / driving_latents/ for PrecomputedDataset.
- Nearest-neighbour resize so integer labels are never blended; labels > K are
  dropped with a warning; ch0 filled uniformly with --environment-switch.
- Reuses process_videos.py helpers (naming, atomic save, VAE factors) and matches
  its typer CLI conventions.

Verified on CPU: a synthetic 2-character label map (plus an out-of-range id)
produces mask (7,17,128,128) with ch0 uniform, slots placed correctly, id>K
dropped, and feeds encode_mask_channels to the 8*(K+1)=56 channels. README +
docs/tasks.md 3.7 updated (upstream label-map generation via SAM/tracking is
dataset-specific and still out of scope).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-13 09:49:25 +08:00
indigo b69eedbd54 Add Modal test environment for SCAIL-2 training/inference
A free-tier-friendly Modal app to validate the SCAIL-2 code path on real
Linux/GPU without a local GPU or the 19B checkpoint.

- modal/checks.py: device-aware (CPU/GPU auto) consolidation of the Phase 1-4
  plumbing checks using tiny random-init models + synthetic data (no checkpoint,
  no dataset): driving concat, zero-init patchify_proj widening (output-preserving),
  a FlexibleStrategy driving+mask training step + loss, and inference conditioning
  assembly. Passes on CPU locally.
- modal/app.py: builds the workspace via `uv sync` (skips CUDA-only ltx-kernels;
  attention falls back to SDPA). Functions: verify (CPU, ~free), smoke (T4, cents),
  train (A10G/A100, paid — runs the real trainer against a checkpoint + data on the
  scail-data Volume).
- modal/README.md: free-tier setup (modal setup), run commands, cost table, the
  scale-up path, and the expected preprocessed dataset layout.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-12 09:37:27 +08:00
indigo a598f89d99 Fix ltx-trainer ruff target-version and a hidden lint finding
The [tool.ruff] target-version was accidentally set to the package version
"1.1.7", which made ruff fail to parse the whole package's pyproject and
silently skip linting. Set it to "py310" to match requires-python >=3.10.

With ruff working again it flagged a too-many-branches finding in the Phase 3
SCAIL wiring: extract the driving + mask-channel loops from _process_modality
into a new _apply_scail_conditions helper. Behavior is unchanged (Phase 3 CPU
verification still passes); full `ruff check .` on the trainer now passes.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-09 21:00:45 +08:00
indigo 06c0870bbb Add SCAIL-2 animation inference pipeline (Phase 4, ltx-pipelines)
Two-stage distilled inference pipeline that animates a character from a driving
video, wiring the Phase 1-3 SCAIL-2 conditioning into a runnable CLI.

- scail_animation.py: ScailAnimationPipeline (mirrors distilled.py) plus a pure,
  testable build_scail_conditionings that assembles VideoConditionByDrivingLatent
  + VideoConditionByMaskChannels (driving appended, mask on the trailing driving
  tokens), and a load_masks helper. main() + scail_animation_arg_parser add
  --driving-video / --mask-path / --mode / --driving-strength on top of the
  standard two-stage distilled parser.
- LTXModelConfigurator now reads config `mask_conditioning_channels`, so a
  SCAIL-trained checkpoint whose config declares it builds the widened
  patchify_proj automatically — no runtime widening wrapper needed at inference.
- CLAUDE.md pipeline table row.

Verified on CPU (verify_phase4_pipeline.py): the module imports, the CLI parses
the SCAIL flags, build_scail_conditionings grows the sequence and places the mask
channels on the driving tokens (target stays zero), driving-only leaves
cond_channels None, and the configurator honors mask_conditioning_channels
(patchify_proj widened from config). End-to-end runs still need a GPU and a
SCAIL-trained checkpoint.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-09 10:22:35 +08:00
indigo e03cc62548 Add SCAIL-2 training integration (Phase 3, ltx-trainer FlexibleStrategy)
Wire SCAIL-2 driving + in-context mask conditioning into the trainer via the
unified FlexibleStrategy, so the widened patchify_proj (Phase 2) can be trained.

- flexible.py: new DrivingConditionConfig (driving-latent concat with a RoPE
  width offset ΔW) and MaskChannelsConditionConfig (semantic masks -> per-token
  channels), added to the condition union and get_data_sources. Driving is
  prepended (cond-first, target stays at the tail for loss slicing); mask
  channels are written onto the driving tokens via Modality.cond_channels,
  reusing ltx-core encode_mask_channels. The noisy target keeps a zero mask.
- model_loader.load_transformer gains mask_conditioning_channels, widening the
  video patchify_proj with zero-init columns via a new live-module helper
  (widen_module_patchify_proj_for_mask_channels). ModelConfig exposes the field.
- trainer unfreezes patchify_proj in LoRA mode when mask channels are active
  (the new input columns are new base params LoRA cannot reach).
- configs/scail_animation_lora.yaml plus README / training-modes table rows.

Verified on CPU (verify_phase3_trainer.py): config round-trips, prepare_training
_inputs builds cond_channels [B,T,56] with the mask on the driving tokens, a tiny
widened model forwards and compute_loss returns a finite [B] loss, and the widen
helper is output-preserving at zero init. Real training needs Linux+GPU+checkpoint;
dataset preprocessing (driving latents + semantic masks) and validation-runner
wiring are left for later.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-09 10:09:40 +08:00
indigo 110adc781e 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>
2026-07-09 09:50:27 +08:00
indigo baa6646fd1 Add SCAIL-2 driving-latent conditioning (Phase 1, inference PoC)
Port mechanisms 1+3 of SCAIL-2 (arXiv:2606.10804) to LTX-2: concatenate a
driving video latent directly into the DiT token sequence with a width-axis
RoPE offset (ΔW) so driving coords stay detached from the target video.

- New VideoConditionByDrivingLatent + DrivingMode in ltx-core conditioning,
  modeled on VideoConditionByReferenceLatent (patchify -> positions -> append
  -> attention mask). Applies ΔW width shift, aligns time to the target, and
  guards against RoPE wrap (max_pos) and target/driving shape mismatch.
- Export both from conditioning packages.
- docs/plan.md and docs/tasks.md track the phased port.

Inference-only: no weight changes. Mechanism 2 (in-context mask channels,
patchify_proj widening) and training are deferred to Phase 2+. Validated via
plumbing checks and a real (random-weight) transformer forward smoke run;
visual quality is not validated (requires Phase 2/3 finetuning).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-09 09:06:35 +08:00
33 changed files with 1953 additions and 7 deletions
+70
View File
@@ -0,0 +1,70 @@
# SCAIL-2 → LTX-2 移植計畫
將 SCAIL-2arXiv:2606.10804`zai-org/SCAIL-2`)的端到端角色動畫手法移植到 LTX-2。
> **注意**SCAIL-2 原始實作建構於 **Wan 2.1**,非 LTX-2。座標慣例與架構需翻譯到 LTX 的資料流。
## 三個核心機制
1. **Driving latent 直接串接** — 把驅動影片 latent 直接接進 DiT token 序列(不經骨架/pose 中介),用 width 軸座標偏移 ΔW 讓 driving 座標跟主 video 座標分開。
2. **In-context mask channel** — 疊加額外輸入 channel(1 個環境開關 + K=6 個角色綁定槽,展開為 `4(K+1)=28` channel)到模型輸入,讓模型知道背景/角色對應。
3. **Mode-specific RoPE** — Animation Mode 與 Replacement Mode 用不同的座標指派規則。
## LTX-2 對應落點
| 機制 | LTX-2 落點 | 改權重? |
|---|---|---|
| 1. Driving 串接 + ΔW | 新 `ConditioningItem`clone `reference_video_cond.py`),改 `positions` 偏移 | 否 |
| 3. Mode-specific RoPE | 上述 item 加 `DrivingMode` enum | 否 |
| 2. In-context mask channel | 加寬 `patchify_proj` 輸入 channel + `LatentState` 帶額外 channel + 投影前 concat | **是**(需微調) |
**關鍵洞察**LTX 的 `ConditioningItem.apply_to(latent_state, latent_tools) -> LatentState` 就是「串接進序列」的天然注入點,機制 1+3 **完全不用動 transformer / rope.py**。RoPE 由 `LatentState.positions` `[B,3,T,2]`pixel 座標,axis1=(time,h,w))驅動;ΔW = 對 `positions[:,2]`width)加常數。width 正規化上界 `max_pos[2]=2048`,超過會 wrap。
## SCAIL-2 座標規則(論文)
序列 `[z_ref; z_t; z_driv]`driving 永遠在 width 軸帶固定偏移 ΔW:
| | Animation | Replacement |
|---|---|---|
| z_ref | T=0, H=[0,Hv), W=[0,Wv) | T=0, **H=[ΔH_ref, ΔH_ref+Hv)**, W=[0,Wv) |
| z_t | T=[1,Tv], H=[0,Hv), W=[0,Wv) | T=[0,Tv1], H=[0,Hv), W=[0,Wv) |
| z_driv | T=[1,Tv], H=[0,Hv), **W=[ΔW, ΔW+Wv)** | T=[0,Tv1], H=[0,Hv), **W=[ΔW, ΔW+Wv)** |
## 分階段計畫
### Phase 1 — Driving 串接 item(機制 1+3,推論期 PoC)✅ 已完成
- 純推論、不改權重、不動 `patchify_proj`、不碰 ltx-trainer。
- 用現有 checkpoint 驗證資料流正確性(序列長度、座標偏移、attention mask、denoise_mask、不 wrap)。
- **PoC 僅驗證 plumbing**LTX-2 未經此訓練,畫面不會是正確動畫。
### 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)✅ 已完成(程式碼路徑;實訓需 GPU)
- **決策**:完整整合到 `FlexibleStrategy`;訓練 = LoRA + 解凍 `patchify_proj`(新 mask 欄位無法純 LoRA 訓練)。
-`DrivingConditionConfig` + `MaskChannelsConditionConfig``flexible.py`);`_apply_driving_condition`(cond-first concat + ΔW) + `_build_mask_channels`(重用 `encode_mask_channels`) → `Modality.cond_channels`
- `load_transformer(mask_conditioning_channels=)``widen_module_patchify_proj_for_mask_channels` 加寬;`ModelConfig.mask_conditioning_channels`trainer 在 LoRA 模式解凍 patchify_proj。
- `configs/scail_animation_lora.yaml` + docs。CPU 單元驗證通過(`verify_phase3_trainer.py`)。
- **本機無 GPU/Linux/checkpoint → 未跑實機訓練**dataset 前處理(driving latents + 語意 mask)與 validation runner 接線未做。
### Phase 4 — Pipeline + CLI 包裝 ✅ 已完成(程式碼路徑;實跑需 GPU)
- `LTXModelConfigurator``mask_conditioning_channels` → SCAIL checkpoint 自描述、載入自動加寬(不需 runtime wrapper)。
- `scail_animation.py``ScailAnimationPipeline`(仿 `distilled.py` 兩階段)+ 可測純函式 `build_scail_conditionings`driving + mask 條件組裝)+ `load_masks` + `main()`
- `utils/args.py` `scail_animation_arg_parser``--driving-video/--mask-path/--mode/--driving-strength`)。
- CPU 驗證通過(`verify_phase4_pipeline.py`);`ltx-pipelines/CLAUDE.md` 表格 row。
- **實機端到端需 GPU + SCAIL-trained checkpoint**config 含 `mask_conditioning_channels` + 訓練好的 patchify_proj + LoRA)。
## Phase 1 簡化取捨(記錄,Phase 2 需回頭處理)
- (a) driving 時間座標直接複製 target 的(token-wise),故 driving 需與 target 同 F/H/W。
- (b) 單一 frozen driving group 的 `attention_mask` 維持 None(= 全連接,target 完全看得到 driving),與 reference cond 一致。
- (c) **ANIMATION 與 REPLACEMENT 在 Phase 1 產生相同 driving 座標** — mode 差異(z_ref 的 ΔH_ref 高度位移、target 時間原點、mask channel)屬 Phase 2enum 先保留佔位。
## 參考
- 論文:arXiv:2606.10804 — *SCAIL-2: Unifying Controlled Character Animation with End-to-end In-Context Conditioning*
- 官方實作:`zai-org/SCAIL-2`GitHub / HuggingFace),建構於 Wan 2.1
- 藍本檔案:`packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py`
+59
View File
@@ -0,0 +1,59 @@
# SCAIL-2 → LTX-2 移植任務追蹤
狀態圖例:✅ 完成 | 🔄 進行中 | ⬜ 未開始 | ⏸️ 暫緩
相關計畫見 [`plan.md`](./plan.md)。
## Phase 1 — Driving 串接 item(推論期 PoC
| # | 任務 | 狀態 | 產出 / 備註 |
|---|---|---|---|
| 1.1 | 新增 `VideoConditionByDrivingLatent` + `DrivingMode` | ✅ | `packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py`,以 `reference_video_cond.py` 為藍本,實作 ΔW width 偏移、時間對齊 target、`max_pos` 防呆、token 數防呆 |
| 1.2 | 匯出新 conditioning 類別 | ✅ | `conditioning/types/__init__.py``conditioning/__init__.py` |
| 1.3 | 免-GPU 資料流驗證腳本 | ✅ | 序列長度、ΔW 不重疊、時間對齊、frozen denoise_mask、attention_mask=None、`clear_conditioning` 剝除、超界/token 數防呆 — 全通過;ruff clean |
| 1.4 | (選配)端到端 smoke run | ✅ | 本機 CPU-only、無 checkpoint,改用小型真實 `LTXModel`(隨機權重)跑 pipeline 實走路徑:`create_noised_state`(含 driving cond)→ `modality_from_latent_state`**真 transformer forward**seq 160=target 80+driving 80)→ `clear_conditioning`(→80)。ANIMATION/REPLACEMENT 皆通過:不 crash、輸出 finite、driving frozen 且正確剝除。**僅驗證整合,不評估畫質** |
## Phase 2 — In-context mask channelcheckpoint 手術)
> **決策**mask channel 時間編碼採**忠實堆疊**——每 latent 幀對應 8 個 pixel 子幀沿 channel 堆疊,`8×(K+1)=56` channelLTX 時間因子 8K=6)。非 Wan 的 4×(K+1)=28。
| # | 任務 | 狀態 | 備註 |
|---|---|---|---|
| 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 tokentarget 保持零 |
| 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
> **決策**:完整整合到 `FlexibleStrategy`;訓練方式 = **LoRA + 解凍 patchify_proj**(新 mask 欄位無法純 LoRA 訓練)。**本機無 GPU/Linux/checkpoint,只做到 CPU 單元驗證**,實訓需在 GPU 機器跑。
| # | 任務 | 狀態 | 備註 |
|---|---|---|---|
| 3.1 | 新增 Driving/MaskChannels ConditionConfig | ✅ | `flexible.py``DrivingConditionConfig`(latents_dir/mode/width_offset) + `MaskChannelsConditionConfig`(mask_dir/num_slots),加入 union + `get_data_sources` |
| 3.2 | strategy 接線 driving concat + cond_channels | ✅ | `_apply_driving_condition`(cond-first concat + ΔW width 偏移) + `_build_mask_channels`(重用 `encode_mask_channels`,寫前 N driving token) → `Modality.cond_channels` |
| 3.3 | model_loader + ModelConfig 支援加寬 | ✅ | `widen_module_patchify_proj_for_mask_channels`(live module zero-init) + `load_transformer(mask_conditioning_channels=)` + `ModelConfig.mask_conditioning_channels` |
| 3.4 | LoRA 模式解凍 patchify_proj | ✅ | `trainer._unfreeze_patchify_proj`mask_channels>0 時把 video patchify_proj 設 trainable,讓新欄位隨 LoRA 一起訓 |
| 3.5 | 範例 config + docs | ✅ | `configs/scail_animation_lora.yaml``configs/README.md``docs/training-modes.md` 表格 row |
| 3.6 | CPU 單元驗證 | ✅ | `verify_phase3_trainer.py`config round-trip、prepare_training_inputs 建 cond_channels[B,T,56]、driving 前置/mask placement、widened model forward + compute_loss finite、widen helper zero-init 等價 |
| 3.7 | dataset 前處理(產 driving latents + 語意 mask | 🟡 | `char_masks/` 前處理已做:`scripts/process_char_masks.py`label-map 影片/圖 → `[K+1,F_pix,H,W]`,對齊 target latentnearest 保留整數 labelch0 環境開關)。`driving_latents/` 沿用既有 `process_videos.py`driving 影片走 video latent 路徑,同 target 形狀)。**仍缺**:從原始影片產生 label-map 的分割/追蹤步驟(SAM 等,資料集特定,未含) |
| 3.8 | validation runner 接 driving/mask | ⬜ | 驗證期取樣尚未接 SCAIL 條件(config 內 validation 先停用),與 Phase 4 一起 |
| 3.9 | 實機訓練跑通 | ⬜ | 需 Linux + GPU + checkpoint,本機無法 |
## Phase 4 — Pipeline + CLI
| # | 任務 | 狀態 | 備註 |
|---|---|---|---|
| 4.1 | configurator 讀 `mask_conditioning_channels` | ✅ | `LTXModelConfigurator` 兩分支 `config.get("mask_conditioning_channels",0)` → SCAIL checkpoint 自描述、載入自動加寬(不需 runtime widen wrapper |
| 4.2 | `ScailAnimationPipeline` + `build_scail_conditionings` | ✅ | `scail_animation.py`:仿 `distilled.py` 兩階段,注入 SCAIL driving+mask conditioning`build_scail_conditionings` 為可測純函式;`load_masks` helper |
| 4.3 | `scail_animation_arg_parser` + `main()` | ✅ | `utils/args.py`:在 `default_2_stage_distilled_arg_parser` 上加 `--driving-video/--mask-path/--mode/--driving-strength` |
| 4.4 | CPU 驗證 + docs | ✅ | `verify_phase4_pipeline.py`import、CLI round-trip、conditioning assembly(driving append + mask 在尾端)、configurator 自描述加寬;`ltx-pipelines/CLAUDE.md` pipeline 表格 row |
| 4.5 | 實機端到端跑通 | ⬜ | 需 GPU + SCAIL-trained checkpoint(含 `mask_conditioning_channels` config + 訓練好的 patchify_proj + LoRA |
## 決議紀錄
- **範圍**:先只做 Phase 1(推論期 PoC)。Phase 2+ 待 Phase 1 驗證後再討論。
- **架構前提**SCAIL-2 建構於 Wan 2.1,本移植為跨架構移植。
+111
View File
@@ -0,0 +1,111 @@
# SCAIL-2 on Modal — free testing environment
Test the SCAIL-2 LTX-2 training/inference **code path** on Modal without a local
GPU. The free path uses tiny random-init models + synthetic data (no checkpoint,
no dataset), so it validates that the SCAIL wiring runs on real Linux/GPU — not
model quality.
> ⚠️ The real 19B LTX-2 model is **not** free to train/run. The `verify`/`smoke`
> targets are near-free; `train` on a real checkpoint uses a paid GPU.
## 1. One-time setup
1. Sign up at [modal.com](https://modal.com) (the free Starter plan includes a
monthly credit allowance — enough for many `verify`/`smoke` runs).
2. Install and authenticate:
```bash
pip install modal
modal setup # opens a browser to link your account / token
```
## 2. Free / near-free checks
From the repo root:
```bash
# CPU: SCAIL-2 Phase 1-4 plumbing (driving concat, mask channels, zero-init widen,
# a FlexibleStrategy training step + loss, inference conditioning assembly).
modal run modal/app.py::verify
# Same checks on a T4 GPU (validates the CUDA path). ~cents.
modal run modal/app.py::smoke
```
The first run builds the image (installs the `ltx-core`/`ltx-pipelines`/`ltx-trainer`
workspace via `uv sync`). `ltx-kernels` (CUDA-compiled) is intentionally skipped —
attention falls back to PyTorch SDPA, so no CUDA toolchain is required.
Expected tail:
```
[OK] Phase 1 driving concat: seq 48 -> 96
[OK] Phase 2 zero-init widen: in_features=72, output preserved
[OK] Phase 3 training step: cond_channels (1, 96, 56), loss ...
[OK] Phase 4 build_scail_conditionings: [driving, mask_channels]
All SCAIL-2 checks passed on cuda # (or cpu)
```
## 3. Scaling up to real training (paid)
`train` runs `packages/ltx-trainer/scripts/train.py` against a real checkpoint.
You must supply the weights + data via the persistent `scail-data` Volume.
1. Create/populate the Volume (checkpoint, Gemma encoder, preprocessed latents):
```bash
modal volume create scail-data # if not auto-created
modal volume put scail-data /local/ltx-2-model.safetensors /model/ltx-2.safetensors
modal volume put scail-data /local/gemma /model/gemma
modal volume put scail-data /local/preprocessed /data/preprocessed
```
2. Copy `configs/scail_animation_lora.yaml`, and point its paths at the mounted
Volume (everything lands under `/data` in the container):
```yaml
model:
model_path: "/data/model/ltx-2.safetensors"
text_encoder_path: "/data/model/gemma"
mask_conditioning_channels: 56 # widens patchify_proj at load
data:
preprocessed_data_root: "/data/preprocessed"
```
(Put your edited config on the Volume too, or bake it into the repo.)
3. Launch (pick a GPU big enough for the model — the 19B needs an A100):
```bash
# edit gpu="A10G" -> "A100" in modal/app.py::train for the full model
modal run modal/app.py::train --config-rel /data/scail_animation_lora.yaml
```
### Dataset preprocessing (not yet automated for SCAIL)
The SCAIL training config expects, under `preprocessed_data_root/`:
```
latents/ # target video latents
conditions/ # text embeddings
driving_latents/ # driving video latents (same F/H/W as target)
char_masks/ # per-sample "mask" = [K+1, F_pix, H, W] (ch0 env switch, 1..K binding slots)
```
`latents/`, `conditions/`, and `driving_latents/` come from the existing
`packages/ltx-trainer/scripts/process_dataset.py` (run it once per video set).
`char_masks/` is produced by `packages/ltx-trainer/scripts/process_char_masks.py`
from per-sample **label-map** videos/images (integer pixel labels: `0` =
environment, `1..K` = characters → binding slots):
```bash
python packages/ltx-trainer/scripts/process_char_masks.py dataset.csv \
--mask-column char_labels --latents-dir ./latents \
--output-dir ./char_masks --num-slots 6 --main-media-column media_path
```
You still need a segmentation/tracking model (e.g. SAM) to *produce* those label
maps from raw video — that upstream step is dataset-specific and not included.
## Cost notes
| Target | GPU | Rough cost | Use |
|---------|-------|-----------|-----|
| `verify`| none | ~free | validate code path on CPU |
| `smoke` | T4 | cents | validate CUDA path |
| `train` | A10G/A100 | paid | real fine-tuning (needs checkpoint + data) |
Free credits are best spent on `verify`/`smoke` to catch integration issues before
committing a paid GPU to a real run. Watch usage in the Modal dashboard.
+84
View File
@@ -0,0 +1,84 @@
"""Modal app for testing the SCAIL-2 LTX-2 training/inference path.
Free-tier friendly: `verify` runs on CPU and `smoke` on a cheap T4, both using
tiny random-init models + synthetic data (no checkpoint, no dataset) so a full run
costs pennies. `train` is the scale-up entrypoint that runs the real trainer once
you upload a checkpoint + preprocessed data to the `scail-data` Volume -- that one
uses a real GPU and is NOT free.
Setup (once):
pip install modal && modal setup
Run:
modal run modal/app.py::verify # CPU plumbing check (~free)
modal run modal/app.py::smoke # same checks on a T4 GPU (cents)
modal run modal/app.py::train --config-rel configs/scail_animation_lora.yaml
"""
import subprocess
from pathlib import Path
import modal
REPO = Path(__file__).parent.parent
# Build the workspace once into the image. ltx-kernels (CUDA-compiled) is excluded
# from the workspace and not needed -- attention falls back to SDPA.
image = (
modal.Image.debian_slim(python_version="3.11")
.apt_install("git")
.pip_install("uv")
.env({"UV_LINK_MODE": "copy", "UV_PROJECT_ENVIRONMENT": "/root/LTX-2/.venv"})
.add_local_dir(
str(REPO),
"/root/LTX-2",
copy=True,
ignore=[".git", ".venv", "**/__pycache__", "**/*.pyc", "outputs", "wandb", "**/.pytest_cache"],
)
.run_commands("cd /root/LTX-2 && uv sync --package ltx-trainer")
)
app = modal.App("scail-ltx2", image=image)
# Persistent volume for the (large) checkpoint + preprocessed dataset used by `train`.
data_volume = modal.Volume.from_name("scail-data", create_if_missing=True)
_UV_PY = ["uv", "run", "--package", "ltx-trainer", "python"]
def _run(args: list[str]) -> None:
subprocess.run([*_UV_PY, *args], cwd="/root/LTX-2", check=True)
@app.function(timeout=1800)
def verify() -> None:
"""Run the SCAIL-2 plumbing checks on CPU (near-free)."""
_run(["modal/checks.py"])
@app.function(gpu="T4", timeout=1800)
def smoke() -> None:
"""Run the same checks on a T4 GPU to validate the CUDA path (cents)."""
_run(["modal/checks.py"])
@app.function(gpu="A10G", timeout=60 * 60 * 6, volumes={"/data": data_volume})
def train(config_rel: str = "configs/scail_animation_lora.yaml") -> None:
"""Run the real SCAIL trainer. NOT free -- needs a real checkpoint + data.
Upload your base checkpoint, Gemma text encoder, and preprocessed data to the
``scail-data`` Volume (mounted at /data), and point the config's paths at /data.
For the full 19B model use a bigger GPU (e.g. gpu="A100") and expect real cost.
"""
_run(["packages/ltx-trainer/scripts/train.py", config_rel])
data_volume.commit()
@app.local_entrypoint()
def main(target: str = "verify", config_rel: str = "configs/scail_animation_lora.yaml") -> None:
if target == "smoke":
smoke.remote()
elif target == "train":
train.remote(config_rel)
else:
verify.remote()
+159
View File
@@ -0,0 +1,159 @@
# ruff: noqa: T201
"""Device-aware SCAIL-2 smoke checks (CPU or GPU), runnable anywhere.
Consolidates the Phase 1-4 plumbing checks into one script that auto-selects CUDA
when available. It uses tiny, randomly-initialised models and synthetic data, so it
needs **no checkpoint and no dataset** -- ideal for a near-free Modal test that
validates the SCAIL training/inference code path end to end in a real Linux/GPU
environment before spending credits on the full 19B model.
Run locally: uv run --package ltx-trainer python modal/checks.py
On Modal: modal run modal/app.py::verify (CPU)
modal run modal/app.py::smoke (T4 GPU)
"""
from __future__ import annotations
from dataclasses import replace
import torch
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.patchifiers import VideoLatentPatchifier
from ltx_core.conditioning import DrivingMode, VideoConditionByDrivingLatent
from ltx_core.model.transformer.mask_channels_checkpoint import (
widen_module_patchify_proj_for_mask_channels,
)
from ltx_core.model.transformer.model import LTXModel, LTXModelType
from ltx_core.tools import VideoLatentTools
from ltx_core.types import SpatioTemporalScaleFactors, VideoLatentShape
from ltx_pipelines.scail_animation import build_scail_conditionings
from ltx_pipelines.utils.helpers import create_noised_state, modality_from_latent_state
from ltx_trainer.timestep_samplers import UniformTimestepSampler
from ltx_trainer.training_strategies.flexible import (
DrivingConditionConfig,
FlexibleStrategy,
FlexibleStrategyConfig,
MaskChannelsConditionConfig,
ModalityConfig,
)
# --- config ---------------------------------------------------------------------
B, C, F, H, W = 1, 16, 3, 4, 4
HEADS, HEAD_DIM, LAYERS, INNER = 4, 16, 2, 64
CTX_LEN, K = 8, 6
SCALE = SpatioTemporalScaleFactors.default()
MASK_CH = SCALE.time * (K + 1) # 56
F_PIX = (F - 1) * SCALE.time + 1
N = F * H * W
def _tiny_model(mask_channels: int, device: torch.device, dtype: torch.dtype) -> LTXModel:
model = LTXModel(
model_type=LTXModelType.VideoOnly,
num_attention_heads=HEADS,
attention_head_dim=HEAD_DIM,
in_channels=C,
out_channels=C,
num_layers=LAYERS,
cross_attention_dim=INNER,
mask_conditioning_channels=mask_channels,
)
for p in model.parameters():
torch.nn.init.normal_(p, std=0.02)
return model.to(device=device, dtype=dtype).eval()
def main() -> None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float32
print(f"== SCAIL-2 checks on device={device} "
f"({torch.cuda.get_device_name(0) if device.type == 'cuda' else 'cpu'}) ==")
torch.manual_seed(0)
tools = VideoLatentTools(
patchifier=VideoLatentPatchifier(patch_size=1),
target_shape=VideoLatentShape(batch=B, channels=C, frames=F, height=H, width=W),
fps=24.0,
scale_factors=SCALE,
)
# Phase 1: driving-latent concat + ΔW RoPE (inference conditioning path).
driving = torch.randn(B, C, F, H, W, device=device, dtype=dtype)
noiser = GaussianNoiser(generator=torch.Generator(device=device).manual_seed(1))
state = create_noised_state(
tools, [VideoConditionByDrivingLatent(driving, mode=DrivingMode.ANIMATION)], noiser, dtype, device, 1.0
)
assert state.latent.shape[1] == 2 * N
print(f"[OK] Phase 1 driving concat: seq {N} -> {state.latent.shape[1]}")
# Phase 2: widen patchify_proj (zero-init) is output-preserving.
base = _tiny_model(0, device, dtype)
ctx = torch.randn(B, CTX_LEN, INNER, device=device, dtype=dtype)
sigma = torch.ones(B, device=device, dtype=dtype)
mod = modality_from_latent_state(state, context=ctx, sigma=sigma)
with torch.inference_mode():
out0, _ = base(video=mod, audio=None, perturbations=None)
widen_module_patchify_proj_for_mask_channels(base, MASK_CH)
mod_cc = replace(mod, cond_channels=torch.randn(B, 2 * N, MASK_CH, device=device, dtype=dtype))
with torch.inference_mode():
out1, _ = base(video=mod_cc, audio=None, perturbations=None)
assert torch.allclose(out0, out1, atol=1e-4), "zero-init widening changed output"
print(f"[OK] Phase 2 zero-init widen: in_features={base.patchify_proj.in_features}, output preserved")
# Phase 3: FlexibleStrategy driving + mask-channels training step.
cfg = FlexibleStrategyConfig(
name="flexible",
video=ModalityConfig(
is_generated=True,
latents_dir="video_latents",
conditions=[
DrivingConditionConfig(type="driving", latents_dir="driving_latents"),
MaskChannelsConditionConfig(type="mask_channels", mask_dir="char_masks", num_slots=K),
],
),
)
strategy = FlexibleStrategy(cfg)
def latents() -> dict:
return {
"latents": torch.randn(B, C, F, H, W, device=device),
"num_frames": torch.tensor([F]),
"height": torch.tensor([H]),
"width": torch.tensor([W]),
"fps": torch.tensor([24.0]),
}
mask_pix = torch.rand(B, K + 1, F_PIX, H * SCALE.height, W * SCALE.width, device=device)
batch = {
"video_latents": latents(),
"driving_latents": latents(),
"char_masks": {"mask": (mask_pix > 0.5).float()},
"conditions": {
"video_prompt_embeds": torch.randn(B, CTX_LEN, INNER, device=device),
"audio_prompt_embeds": torch.randn(B, CTX_LEN, INNER, device=device),
"prompt_attention_mask": torch.ones(B, CTX_LEN, device=device),
},
}
inputs = strategy.prepare_training_inputs(batch, UniformTimestepSampler(0.0, 1.0))
assert inputs.video.cond_channels.shape == (B, 2 * N, MASK_CH)
model = _tiny_model(MASK_CH, device, dtype)
with torch.inference_mode():
vpred, _ = model(video=inputs.video, audio=None, perturbations=None)
loss = strategy.compute_loss(vpred, None, inputs)
assert loss.shape == (B,)
assert torch.isfinite(loss).all()
cc_shape = tuple(inputs.video.cond_channels.shape)
print(f"[OK] Phase 3 training step: cond_channels {cc_shape}, loss {loss.item():.4f}")
# Phase 4: inference conditioning assembly.
masks = batch["char_masks"]["mask"]
conds = build_scail_conditionings(driving, masks, mode=DrivingMode.ANIMATION)
assert len(conds) == 2
print("[OK] Phase 4 build_scail_conditionings: [driving, mask_channels]")
print("\nAll SCAIL-2 checks passed on", device)
if __name__ == "__main__":
main()
@@ -5,10 +5,14 @@ from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.types import (
AudioConditionByReferenceLatent,
ConditioningItemAttentionStrengthWrapper,
DrivingMode,
VideoConditionByDrivingLatent,
VideoConditionByKeyframeIndex,
VideoConditionByLatentIndex,
VideoConditionByMask,
VideoConditionByMaskChannels,
VideoConditionByReferenceLatent,
encode_mask_channels,
)
__all__ = [
@@ -16,8 +20,12 @@ __all__ = [
"ConditioningError",
"ConditioningItem",
"ConditioningItemAttentionStrengthWrapper",
"DrivingMode",
"VideoConditionByDrivingLatent",
"VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex",
"VideoConditionByMask",
"VideoConditionByMaskChannels",
"VideoConditionByReferenceLatent",
"encode_mask_channels",
]
@@ -0,0 +1,28 @@
"""Helpers for in-context conditioning channels (``LatentState.cond_channels``).
Conditioning channels are extra per-token input features (e.g. SCAIL-2 mask
channels) that ride alongside the latent in patchified token space ``(B, T, C)``
and are concatenated onto the latent before the model's first projection. When
present they must stay length-aligned with the token sequence, so any
conditioning item that *appends* tokens must also extend the channels.
"""
from __future__ import annotations
import torch
from ltx_core.types import LatentState
def extend_cond_channels(latent_state: LatentState, num_new_tokens: int) -> torch.Tensor | None:
"""Return ``cond_channels`` extended with ``num_new_tokens`` zero rows, or ``None``.
Appended tokens (reference / driving / keyframe) carry no in-context signal by default, so they
are padded with zeros to preserve the ``T``-alignment invariant. Returns ``None`` unchanged when
the state has no conditioning channels (the standard case), keeping non-SCAIL pipelines untouched.
"""
cond_channels = latent_state.cond_channels
if cond_channels is None:
return None
zeros = cond_channels.new_zeros(cond_channels.shape[0], num_new_tokens, cond_channels.shape[2])
return torch.cat([cond_channels, zeros], dim=1)
@@ -1,8 +1,10 @@
"""Conditioning type implementations."""
from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper
from ltx_core.conditioning.types.driving_video_cond import DrivingMode, VideoConditionByDrivingLatent
from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
from ltx_core.conditioning.types.mask_channels_cond import VideoConditionByMaskChannels, encode_mask_channels
from ltx_core.conditioning.types.mask_cond import VideoConditionByMask
from ltx_core.conditioning.types.reference_audio_cond import AudioConditionByReferenceLatent
from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
@@ -10,8 +12,12 @@ from ltx_core.conditioning.types.reference_video_cond import VideoConditionByRef
__all__ = [
"AudioConditionByReferenceLatent",
"ConditioningItemAttentionStrengthWrapper",
"DrivingMode",
"VideoConditionByDrivingLatent",
"VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex",
"VideoConditionByMask",
"VideoConditionByMaskChannels",
"VideoConditionByReferenceLatent",
"encode_mask_channels",
]
@@ -0,0 +1,162 @@
"""Driving-video conditioning for SCAIL-2-style end-to-end character animation.
Ports mechanism 1 + 3 of SCAIL-2 (arXiv:2606.10804) to LTX-2: the *driving*
video latent is concatenated directly into the DiT token sequence (no skeleton /
pose intermediate), and carries a fixed spatial offset ``width_offset`` (the
paper's ΔW) on the RoPE width axis so its coordinates stay detached from the
main video tokens. This is the inference-only PoC path -- it reuses the existing
frozen-reference-token machinery and does not touch the transformer or RoPE.
This mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent`
(same patchify -> positions -> append -> attention-mask flow); the only new
behaviour is the mode-aware coordinate assignment.
Scope note (Phase 1): SCAIL-2's in-context mask channels (mechanism 2) and the
reference-latent height shift ΔH_ref of Replacement Mode require model-weight
surgery / a separate reference token group and are intentionally out of scope
here. Because LTX handles the reference image as a frame-0 in-place replacement
(not a separate token group), the driving-token placement is identical for both
:class:`DrivingMode` values in Phase 1 -- ``mode`` is stored for forward
compatibility and to document intent, but does not yet alter the driving
coordinates. See the plan for the deferred Phase 2 work.
"""
from __future__ import annotations
from enum import Enum
import torch
from ltx_core.components.patchifiers import get_pixel_coords
from ltx_core.conditioning.cond_channels import extend_cond_channels
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.mask_utils import update_attention_mask
from ltx_core.tools import VideoLatentTools
from ltx_core.types import LatentState, VideoLatentShape
# Default normalization ceiling for the RoPE width axis. Must stay in sync with
# the model's ``positional_embedding_max_pos[2]`` (see rope.py:precompute_freqs_cis
# default ``max_pos=[20, 2048, 2048]`` and model.py `_init_` default). Driving
# width coordinates that reach or exceed this value would wrap under RoPE.
DEFAULT_MAX_WIDTH_POSITION = 2048
class DrivingMode(Enum):
"""SCAIL-2 conditioning mode. See module docstring for the Phase 1 caveat."""
ANIMATION = "animation"
REPLACEMENT = "replacement"
class VideoConditionByDrivingLatent(ConditioningItem):
"""Append driving-video tokens with a width-axis RoPE offset (SCAIL-2 ΔW).
The driving tokens are appended after the target sequence as clean latents
(placeholder zeros in the noisy latent), kept frozen (``denoise_mask =
1 - strength``), temporally aligned to the target, and shifted along the
width axis by ``width_offset`` so they occupy ``[ΔW, ΔW + Wv)`` while the
target stays at ``[0, Wv)``.
Args:
latent: Driving video latents ``[B, C, F, H, W]``. Must match the target
shape (same F/H/W) so tokens align frame-for-frame with the target.
mode: SCAIL-2 mode (reserved for Phase 2; see module docstring).
width_offset: ΔW in RoPE pixel-space width units. ``None`` (default) uses
the target's pixel width (``target_shape.width * scale_factors.width``),
placing the driving tokens immediately to the right of the target.
strength: 1.0 keeps the driving latent fully clean (frozen); 0.0 would
denoise it. Default 1.0.
max_width_position: RoPE width normalization ceiling; validation raises if
the shifted driving coordinates would reach it. Keep in sync with the
model's ``positional_embedding_max_pos[2]``.
"""
def __init__(
self,
latent: torch.Tensor,
mode: DrivingMode = DrivingMode.ANIMATION,
width_offset: float | None = None,
strength: float = 1.0,
max_width_position: int = DEFAULT_MAX_WIDTH_POSITION,
):
self.latent = latent
self.mode = mode
self.width_offset = width_offset
self.strength = strength
self.max_width_position = max_width_position
def apply_to(
self,
latent_state: LatentState,
latent_tools: VideoLatentTools,
) -> LatentState:
"""Append driving tokens with target-aligned time and a ΔW width shift."""
tokens = latent_tools.patchifier.patchify(self.latent)
num_target_tokens = latent_tools.patchifier.get_token_count(latent_tools.target_shape)
if tokens.shape[1] != num_target_tokens:
raise ValueError(
"VideoConditionByDrivingLatent expects the driving latent to match the target shape "
f"(same F/H/W): got {tokens.shape[1]} driving tokens vs {num_target_tokens} target tokens. "
"Resize/resample the driving video to the target resolution and frame count."
)
# Compute the driving tokens' own pixel-space coordinates (same flow as
# the base reference conditioning and create_initial_state).
latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
output_shape=VideoLatentShape.from_torch_shape(self.latent.shape),
device=self.latent.device,
)
positions = get_pixel_coords(
latent_coords=latent_coords,
scale_factors=latent_tools.scale_factors,
causal_fix=latent_tools.causal_fix,
).to(dtype=torch.float32)
# Temporal alignment: copy the target's time coordinates so the driving
# tokens sit on exactly the same time grid as z_t (robust to causal_fix /
# fps nuances). Token order is a flattened (f h w) grid identical to the
# target's, so a token-wise copy is frame-aligned.
positions[:, 0:1, :] = latent_state.positions[:, 0:1, :num_target_tokens].to(dtype=torch.float32)
# ΔW: shift the driving tokens along the width axis so they stay spatially
# detached from the target tokens. Default offset = target pixel width,
# giving target=[0, Wv), driving=[Wv, 2*Wv).
width_offset = self.width_offset
if width_offset is None:
width_offset = float(latent_tools.target_shape.width * latent_tools.scale_factors.width)
positions[:, 2, ...] = positions[:, 2, ...] + width_offset
max_width = positions[:, 2, ...].max().item()
if max_width >= self.max_width_position:
raise ValueError(
f"Driving width coordinate {max_width:.1f} reaches the RoPE ceiling "
f"{self.max_width_position} and would wrap. Reduce width_offset "
f"(currently {width_offset:.1f}) or lower the output width."
)
denoise_mask = torch.full(
size=(*tokens.shape[:2], 1),
fill_value=1.0 - self.strength,
device=self.latent.device,
dtype=self.latent.dtype,
)
new_attention_mask = update_attention_mask(
latent_state=latent_state,
attention_mask=None,
num_noisy_tokens=num_target_tokens,
num_new_tokens=tokens.shape[1],
batch_size=tokens.shape[0],
device=self.latent.device,
dtype=self.latent.dtype,
)
return LatentState(
latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
positions=torch.cat([latent_state.positions, positions], dim=2),
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
attention_mask=new_attention_mask,
cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
)
@@ -1,6 +1,7 @@
import torch
from ltx_core.components.patchifiers import get_pixel_coords
from ltx_core.conditioning.cond_channels import extend_cond_channels
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.mask_utils import update_attention_mask
from ltx_core.tools import VideoLatentTools
@@ -81,4 +82,5 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
positions=torch.cat([latent_state.positions, positions], dim=2),
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
attention_mask=new_attention_mask,
cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
)
@@ -0,0 +1,122 @@
"""SCAIL-2 in-context mask conditioning (mechanism 2) for LTX-2.
Encodes ``K+1`` semantic pixel-space masks (1 environment switch + ``K`` character
binding slots) into per-token conditioning channels and writes them onto the
sequence via :attr:`LatentState.cond_channels`. Following the paper, mask signals
are carried by the *driving* (and reference) tokens while the noisy target keeps
an all-zero mask -- so by default the channels are written onto the last ``N``
tokens of the sequence (the appended driving group), which means **this item must
be applied after the driving conditioning**.
Channel expansion (LTX-2 faithful port of the paper's ``4(K+1)``): each semantic
mask is spatially downsampled to the latent grid and temporally stacked along the
channel dimension by the VAE temporal factor ``t`` (8 for LTX-2, vs 4 for the
paper's Wan-2.1 backbone), giving ``t*(K+1)`` channels -- ``8*7 = 56`` for the
default ``K=6``. The model's ``patchify_proj`` must be built with a matching
``mask_conditioning_channels`` (see :class:`ltx_core.model.transformer.model.LTXModel`).
"""
from __future__ import annotations
from dataclasses import replace
import torch
import torch.nn.functional as F
from einops import rearrange
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.tools import VideoLatentTools
from ltx_core.types import LatentState
def encode_mask_channels(
masks: torch.Tensor,
temporal_factor: int,
height_lat: int,
width_lat: int,
frames_lat: int,
) -> torch.Tensor:
"""Encode ``[B, K+1, F_pix, H_pix, W_pix]`` masks into ``[B, t*(K+1), F_lat, H_lat, W_lat]``.
Each semantic mask is area-downsampled to the latent spatial grid, then the pixel frames that
map to one latent frame are stacked along the channel dimension (temporal factor ``t``). The
causal first latent frame corresponds to a single pixel frame, which is replicated across its
``t`` stacked channels so every latent frame yields a uniform ``t`` channels per semantic class.
"""
b, s, f_pix, _h_pix, _w_pix = masks.shape
t = temporal_factor
expected_f_pix = (frames_lat - 1) * t + 1
if f_pix != expected_f_pix:
raise ValueError(
f"mask pixel frames ({f_pix}) incompatible with latent frames ({frames_lat}) at temporal "
f"factor {t}: expected (F_lat - 1) * {t} + 1 = {expected_f_pix}."
)
# Spatial downsample every (semantic, frame) mask to the latent grid.
flat = rearrange(masks.to(dtype=torch.float32), "b s f h w -> (b s f) 1 h w")
down = F.interpolate(flat, size=(height_lat, width_lat), mode="area")
down = rearrange(down, "(b s f) 1 h w -> b s f h w", b=b, s=s)
# Temporal stacking. Latent frame 0 = pixel frame 0 (causal), replicated across t channels;
# latent frames 1.. group t consecutive pixel frames.
first = down[:, :, :1].repeat(1, 1, t, 1, 1).unsqueeze(2) # [B, S, 1, t, H, W]
rest = rearrange(down[:, :, 1:], "b s (fl t) h w -> b s fl t h w", t=t) # [B, S, F_lat-1, t, H, W]
stacked = torch.cat([first, rest], dim=2) # [B, S, F_lat, t, H, W]
# Fold (semantic, temporal) into a single channel axis, semantic-major: channel = s * t + ti.
return rearrange(stacked, "b s fl t h w -> b (s t) fl h w") # [B, t*(K+1), F_lat, H_lat, W_lat]
class VideoConditionByMaskChannels(ConditioningItem):
"""Write SCAIL-2 in-context mask channels onto the driving/reference tokens.
Args:
masks: Pixel-space semantic masks ``[B, K+1, F_pix, H_pix, W_pix]``. Channel 0 is the
environment switch (whether the environment comes from the reference vs the driving
video); channels ``1..K`` are the character binding slots (regions sharing a slot share
motion). ``F_pix`` must equal ``(F_lat - 1) * temporal_factor + 1``.
applies_to_last_n: Number of trailing tokens to write the mask onto. ``None`` (default) uses
the target token count, i.e. the appended driving group when this item runs right after
the driving conditioning. The noisy target tokens keep an all-zero mask (paper-faithful).
"""
def __init__(self, masks: torch.Tensor, applies_to_last_n: int | None = None):
self.masks = masks
self.applies_to_last_n = applies_to_last_n
def apply_to(self, latent_state: LatentState, latent_tools: VideoLatentTools) -> LatentState:
shape = latent_tools.target_shape
cond = encode_mask_channels(
masks=self.masks,
temporal_factor=latent_tools.scale_factors.time,
height_lat=shape.height,
width_lat=shape.width,
frames_lat=shape.frames,
)
tokens = latent_tools.patchifier.patchify(cond) # [B, T_region, C_mask]
b, n_region, c_mask = tokens.shape
total_tokens = latent_state.latent.shape[1]
n = self.applies_to_last_n if self.applies_to_last_n is not None else n_region
if n != n_region:
raise ValueError(
f"applies_to_last_n ({n}) must equal the encoded mask token count ({n_region})."
)
if n > total_tokens:
raise ValueError(
f"cannot write {n} mask tokens onto a sequence of only {total_tokens} tokens."
)
cond_channels = latent_state.cond_channels
if cond_channels is None:
cond_channels = tokens.new_zeros(b, total_tokens, c_mask)
else:
if cond_channels.shape[2] != c_mask:
raise ValueError(
f"existing cond_channels width ({cond_channels.shape[2]}) != mask channels ({c_mask})."
)
cond_channels = cond_channels.clone()
start = total_tokens - n
cond_channels[:, start : start + n] = tokens.to(dtype=cond_channels.dtype)
return replace(latent_state, cond_channels=cond_channels)
@@ -4,6 +4,7 @@ from __future__ import annotations
import torch
from ltx_core.conditioning.cond_channels import extend_cond_channels
from ltx_core.conditioning.mask_utils import update_attention_mask
from ltx_core.tools import LatentTools
from ltx_core.types import LatentState
@@ -56,4 +57,5 @@ class AudioConditionByReferenceLatent:
positions=torch.cat([latent_state.positions, self.positions], dim=2),
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
attention_mask=new_attention_mask,
cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
)
@@ -3,6 +3,7 @@
import torch
from ltx_core.components.patchifiers import get_pixel_coords
from ltx_core.conditioning.cond_channels import extend_cond_channels
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.mask_utils import update_attention_mask
from ltx_core.tools import VideoLatentTools
@@ -99,4 +100,5 @@ class VideoConditionByReferenceLatent(ConditioningItem):
positions=torch.cat([latent_state.positions, positions], dim=2),
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
attention_mask=new_attention_mask,
cond_channels=extend_cond_channels(latent_state, tokens.shape[1]),
)
@@ -0,0 +1,92 @@
"""Checkpoint surgery to enable SCAIL-2 in-context mask channels on a trained LTX-2 model.
Widening ``LTXModel.mask_conditioning_channels`` from 0 to ``C`` grows the video
``patchify_proj`` weight from ``[inner, in]`` to ``[inner, in + C]``. This helper
appends ``C`` **zero** input columns to that weight so a converted checkpoint is
numerically identical to the original until the new columns are finetuned -- the
extra channels contribute nothing at load time.
Usage::
sd = load_state_dict(path)
widen_patchify_proj_for_mask_channels(sd, mask_channels=56)
model = LTXModel(..., mask_conditioning_channels=56)
model.load_state_dict(sd)
"""
from __future__ import annotations
import torch
_VIDEO_PROJ_SUFFIX = "patchify_proj.weight"
_AUDIO_PROJ_SUFFIX = "audio_patchify_proj.weight"
def widen_module_patchify_proj_for_mask_channels(model: torch.nn.Module, mask_channels: int) -> torch.nn.Module:
"""In place, widen a live ``LTXModel``'s video ``patchify_proj`` to accept ``mask_channels`` extra inputs.
Replaces ``model.patchify_proj`` with a wider ``Linear`` whose original input columns are copied and
whose ``mask_channels`` new columns are zero (so the model reproduces its pre-widening output until
those columns are trained). Also sets ``model.mask_conditioning_channels`` for metadata/consistency.
Idempotent guard: raises if the model is already widened to a different value.
"""
if mask_channels < 0:
raise ValueError(f"mask_channels must be non-negative, got {mask_channels}")
proj = getattr(model, "patchify_proj", None)
if not isinstance(proj, torch.nn.Linear):
raise AttributeError("model has no linear 'patchify_proj' to widen")
already = int(getattr(model, "mask_conditioning_channels", 0))
if mask_channels in (0, already):
model.mask_conditioning_channels = mask_channels
return model
if already != 0:
raise ValueError(f"patchify_proj already widened for {already} mask channels, refusing to re-widen")
new_in = proj.in_features + mask_channels
wider = torch.nn.Linear(new_in, proj.out_features, bias=proj.bias is not None)
wider = wider.to(device=proj.weight.device, dtype=proj.weight.dtype)
with torch.no_grad():
wider.weight.zero_()
wider.weight[:, : proj.in_features].copy_(proj.weight)
if proj.bias is not None:
wider.bias.copy_(proj.bias)
model.patchify_proj = wider
model.mask_conditioning_channels = mask_channels
return model
def widen_patchify_proj_for_mask_channels(
state_dict: dict[str, torch.Tensor],
mask_channels: int,
) -> dict[str, torch.Tensor]:
"""In place, append ``mask_channels`` zero input columns to the video ``patchify_proj`` weight.
Only the video projection is touched (audio ``audio_patchify_proj`` is left alone). The bias,
if present, is unchanged. Returns the same dict for convenience. Idempotency is the caller's
responsibility -- calling twice widens twice.
"""
if mask_channels < 0:
raise ValueError(f"mask_channels must be non-negative, got {mask_channels}")
if mask_channels == 0:
return state_dict
target_keys = [
key
for key in state_dict
if key.endswith(_VIDEO_PROJ_SUFFIX) and not key.endswith(_AUDIO_PROJ_SUFFIX)
]
if not target_keys:
raise KeyError(
f"No '{_VIDEO_PROJ_SUFFIX}' weight found in the state dict; cannot widen for mask channels."
)
for key in target_keys:
weight = state_dict[key] # [inner_dim, in_features]
if weight.ndim != 2:
raise ValueError(f"Expected 2-D weight for '{key}', got shape {tuple(weight.shape)}")
pad = weight.new_zeros(weight.shape[0], mask_channels)
state_dict[key] = torch.cat([weight, pad], dim=1)
return state_dict
@@ -38,6 +38,9 @@ class Modality:
attention. ``None`` means unrestricted (full) attention between
all tokens. Built incrementally by conditioning items; see
:class:`~ltx_core.conditioning.types.attention_strength_wrapper.ConditioningItemAttentionStrengthWrapper`.
cond_channels: Optional in-context conditioning channels, shape ``(B, T, C_cond)``, aligned token-for-token
with ``latent``. Concatenated onto ``latent`` before the first projection (see
:meth:`TransformerArgsPreprocessor.prepare`); never noised. ``None`` for standard models.
"""
latent: (
@@ -53,6 +56,7 @@ class Modality:
enabled: bool = True
context_mask: torch.Tensor | None = None
attention_mask: torch.Tensor | None = None
cond_channels: torch.Tensor | None = None
def split(self, sizes: list[int]) -> list[Modality]:
"""Split along the batch dimension into chunks of the given sizes."""
@@ -73,6 +73,7 @@ class LTXModel(torch.nn.Module):
caption_projection: torch.nn.Module | None = None,
audio_caption_projection: torch.nn.Module | None = None,
cross_attention_adaln: bool = False,
mask_conditioning_channels: int = 0,
):
super().__init__()
# Log the attention backends this transformer is built with. Reading the resolved
@@ -86,6 +87,10 @@ class LTXModel(torch.nn.Module):
)
self._enable_gradient_checkpointing = False
self.cross_attention_adaln = cross_attention_adaln
# Extra per-token input channels (e.g. SCAIL-2 in-context mask channels) concatenated onto the
# video latent before ``patchify_proj``. 0 keeps the standard input width; >0 widens the first
# projection. Fed via ``Modality.cond_channels`` (see TransformerArgsPreprocessor.prepare).
self.mask_conditioning_channels = mask_conditioning_channels
self.use_middle_indices_grid = use_middle_indices_grid
self.rope_type = rope_type
self.double_precision_rope = double_precision_rope
@@ -154,8 +159,11 @@ class LTXModel(torch.nn.Module):
caption_projection: torch.nn.Module | None = None,
) -> None:
"""Initialize video-specific components."""
# Video input components
self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
# Video input components. When ``mask_conditioning_channels > 0`` the first projection is
# widened to also accept the in-context conditioning channels concatenated onto the latent.
self.patchify_proj = torch.nn.Linear(
in_channels + self.mask_conditioning_channels, self.inner_dim, bias=True
)
if caption_projection is not None:
self.caption_projection = caption_projection
@@ -69,6 +69,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
caption_projection=caption_projection,
audio_caption_projection=audio_caption_projection,
cross_attention_adaln=config.get("cross_attention_adaln", False),
mask_conditioning_channels=config.get("mask_conditioning_channels", 0),
)
@@ -120,6 +121,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
apply_gated_attention=config.get("apply_gated_attention", False),
caption_projection=caption_projection,
cross_attention_adaln=config.get("cross_attention_adaln", False),
mask_conditioning_channels=config.get("mask_conditioning_channels", 0),
)
@@ -201,12 +201,46 @@ class TransformerArgsPreprocessor:
)
return pe
def _apply_patchify_proj(self, modality: Modality) -> torch.Tensor:
"""Project patchified latents, concatenating any in-context conditioning channels first.
When ``patchify_proj`` expects more input features than the latent provides (a model built
with ``mask_conditioning_channels > 0``), the extra width is filled by ``modality.cond_channels``
(per-token, never noised). If those channels are absent they default to zeros, so a widened
model still runs and, with zero-initialized new projection columns, reproduces the base model's
output exactly.
"""
latent = modality.latent
expected = self.patchify_proj.in_features
actual = latent.shape[-1]
if expected != actual:
missing = expected - actual
if missing < 0:
raise ValueError(
f"patchify_proj expects {expected} input features but the latent already has {actual}."
)
cond_channels = modality.cond_channels
if cond_channels is None:
cond_channels = latent.new_zeros(latent.shape[0], latent.shape[1], missing)
elif cond_channels.shape[-1] != missing:
raise ValueError(
f"cond_channels has {cond_channels.shape[-1]} channels but patchify_proj needs {missing} "
f"extra input features (latent {actual} + cond {cond_channels.shape[-1]} != {expected})."
)
elif cond_channels.shape[1] != latent.shape[1]:
raise ValueError(
f"cond_channels token length {cond_channels.shape[1]} must match latent token length "
f"{latent.shape[1]}."
)
latent = torch.cat([latent, cond_channels.to(dtype=latent.dtype)], dim=-1)
return self.patchify_proj(latent)
def prepare(
self,
modality: Modality,
cross_modality: Modality | None = None, # noqa: ARG002
) -> TransformerArgs:
x = self.patchify_proj(modality.latent)
x = self._apply_patchify_proj(modality)
batch_size = x.shape[0]
timestep, embedded_timestep = self._prepare_timestep(
modality.timesteps, self.adaln, batch_size, modality.latent.dtype
+4
View File
@@ -75,6 +75,9 @@ class LatentTools(Protocol):
clean_latent = latent_state.clean_latent[:, :num_tokens]
denoise_mask = torch.ones_like(latent_state.denoise_mask)[:, :num_tokens]
positions = latent_state.positions[:, :, :num_tokens]
cond_channels = (
latent_state.cond_channels[:, :num_tokens] if latent_state.cond_channels is not None else None
)
return LatentState(
latent=latent,
@@ -82,6 +85,7 @@ class LatentTools(Protocol):
positions=positions,
clean_latent=clean_latent,
attention_mask=None,
cond_channels=cond_channels,
)
+7
View File
@@ -193,6 +193,11 @@ class LatentState:
clean_latent: Initial state of the latent before denoising, may include conditioning latents.
attention_mask: Optional 2D self-attention mask of shape (B, T, T). Values in [0, 1] where 1 = full attention,
0 = no attention. None means full attention everywhere. Built incrementally by conditioning items.
cond_channels: Optional in-context conditioning channels in patchified token space, shape (B, T, C_cond),
aligned token-for-token with ``latent``. These extra per-token feature channels (e.g. SCAIL-2 mask
channels) are concatenated onto the latent right before the model's first projection; they are never
noised or denoised. ``None`` means no conditioning channels (standard models). When present, the token
length T must always match ``latent``; token-appending conditioning items extend it with zeros.
"""
latent: torch.Tensor
@@ -200,6 +205,7 @@ class LatentState:
positions: torch.Tensor
clean_latent: torch.Tensor
attention_mask: torch.Tensor | None = None
cond_channels: torch.Tensor | None = None
def clone(self) -> "LatentState":
return LatentState(
@@ -208,4 +214,5 @@ class LatentState:
positions=self.positions.clone(),
clean_latent=self.clean_latent.clone(),
attention_mask=self.attention_mask.clone() if self.attention_mask is not None else None,
cond_channels=self.cond_channels.clone() if self.cond_channels is not None else None,
)
+1
View File
@@ -25,6 +25,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
| `DistilledPipeline` | `distilled.py` | 2 | Distilled only | Euler | Fastest inference |
| `ICLoraPipeline` | `ic_lora.py` | 2 | Distilled only | Euler | Video-to-video with IC-LoRA control |
| `LipDubPipeline` | `lipdub.py` | 2 | Distilled only | Euler | Lip dubbing with IC-LoRA + audio ref conditioning |
| `ScailAnimationPipeline` | `scail_animation.py` | 2 | Distilled + SCAIL LoRA | Euler | SCAIL-2 character animation from a driving video (driving latent concat + ΔW RoPE + in-context mask channels) |
| `RetakePipeline` | `retake.py` | 1 | Full or distilled | Euler | Video region regeneration |
## Guidance
@@ -0,0 +1,349 @@
"""SCAIL-2 character-animation inference pipeline.
Two-stage distilled video generation that animates a character from a *driving*
video: the driving latent is concatenated into the token sequence with a RoPE
width offset (ΔW) and in-context mask channels route motion per character. This
mirrors the distilled pipeline (``distilled.py``) and lip-dub pipeline
(``lipdub.py``) structure, adding the SCAIL conditioning assembly.
Requirements:
- A SCAIL-trained checkpoint whose config declares ``mask_conditioning_channels``
(so ``LTXModelConfigurator`` builds the widened ``patchify_proj`` automatically)
plus the SCAIL LoRA adapter.
- A driving video (same target resolution) and a semantic mask tensor
``[K+1, F_pix, H_pix, W_pix]`` (channel 0 = environment switch, 1..K = binding
slots), saved as a ``.pt`` file.
The SCAIL-specific conditioning assembly is factored into
``build_scail_conditionings`` so it can be unit-tested without the heavy models.
"""
from __future__ import annotations
import logging
from collections.abc import Iterator
import torch
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.conditioning import (
ConditioningItem,
DrivingMode,
VideoConditionByDrivingLatent,
VideoConditionByMaskChannels,
)
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import ImageConditioningInput, resolve_cli_params, scail_animation_arg_parser
from ltx_pipelines.utils.blocks import (
DiffusionStage,
ImageConditioner,
PromptEncoder,
VideoDecoder,
VideoUpsampler,
)
from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
from ltx_pipelines.utils.media_io import decode_video_by_frame, encode_video, video_preprocess
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
def build_scail_conditionings(
driving_latent: torch.Tensor,
masks: torch.Tensor | None,
*,
mode: DrivingMode = DrivingMode.ANIMATION,
strength: float = 1.0,
width_offset: float | None = None,
) -> list[ConditioningItem]:
"""Assemble the SCAIL driving + mask-channel conditioning items (order matters).
The driving item is applied first (it appends the driving tokens at the tail); the mask item is
applied second, writing its channels onto those trailing driving tokens (the noisy target keeps a
zero mask). ``masks`` may be ``None`` to run driving-only (e.g. a model without mask channels).
Args:
driving_latent: Driving video latent ``[B, C, F, H, W]`` (same F/H/W as the target).
masks: Semantic masks ``[B, K+1, F_pix, H_pix, W_pix]`` or ``None``.
mode: SCAIL mode (animation/replacement).
strength: Driving conditioning strength (1.0 keeps it clean/frozen).
width_offset: ΔW in RoPE pixel-space width units (None = target pixel width).
"""
conditionings: list[ConditioningItem] = [
VideoConditionByDrivingLatent(
latent=driving_latent,
mode=mode,
width_offset=width_offset,
strength=strength,
)
]
if masks is not None:
conditionings.append(VideoConditionByMaskChannels(masks=masks))
return conditionings
def load_masks(mask_path: str, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
"""Load a semantic mask tensor from a ``.pt`` file and shape it to ``[B, K+1, F_pix, H, W]``."""
masks = torch.load(mask_path, map_location=device, weights_only=True)
if isinstance(masks, dict):
masks = masks["mask"]
if masks.dim() == 4: # [K+1, F, H, W] -> add batch
masks = masks.unsqueeze(0)
return masks.to(device=device, dtype=dtype)
class ScailAnimationPipeline:
"""Two-stage distilled SCAIL-2 character animation (driving video + in-context mask channels)."""
def __init__(
self,
distilled_checkpoint_path: str,
gemma_root: str,
spatial_upsampler_path: str,
scail_lora: LoraPathStrengthAndSDOps,
device: torch.device | None = None,
quantization: QuantizationPolicy | None = None,
registry: Registry | None = None,
compilation_config: CompilationConfig | None = None,
offload_mode: OffloadMode = OffloadMode.NONE,
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
) -> None:
self.device = device or get_device()
self.dtype = torch.bfloat16
self.prompt_encoder = PromptEncoder(
distilled_checkpoint_path,
gemma_root,
self.dtype,
self.device,
registry=registry,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.image_conditioner = ImageConditioner(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.stage = DiffusionStage.from_checkpoint(
distilled_checkpoint_path,
self.dtype,
self.device,
loras=(scail_lora,),
quantization=quantization,
registry=registry,
compilation_config=compilation_config,
offload_mode=offload_mode,
alloc_trim_strategy=alloc_trim_strategy,
)
self.upsampler = VideoUpsampler(
distilled_checkpoint_path,
spatial_upsampler_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
self.video_decoder = VideoDecoder(
distilled_checkpoint_path,
self.dtype,
self.device,
registry=registry,
alloc_trim_strategy=alloc_trim_strategy,
)
def _encode_driving_latent(
self,
driving_video_path: str,
height: int,
width: int,
num_frames: int,
video_encoder: VideoEncoder,
tiling_config: TilingConfig | None,
) -> torch.Tensor:
"""Decode + encode the driving video into a VAE latent at the target resolution."""
frame_gen = decode_video_by_frame(path=driving_video_path, frame_cap=num_frames, device=self.device)
video = video_preprocess(frame_gen, height, width, self.dtype, self.device)
if tiling_config is not None:
return video_encoder.tiled_encode(video, tiling_config)
return video_encoder(video)
def _video_conditionings(
self,
images: list[ImageConditioningInput],
driving_video_path: str,
masks: torch.Tensor | None,
height: int,
width: int,
num_frames: int,
mode: DrivingMode,
driving_strength: float,
video_encoder: VideoEncoder,
tiling_config: TilingConfig | None,
) -> list[ConditioningItem]:
conditionings = combined_image_conditionings(
images=images,
height=height,
width=width,
video_encoder=video_encoder,
dtype=self.dtype,
device=self.device,
)
driving_latent = self._encode_driving_latent(
driving_video_path, height, width, num_frames, video_encoder, tiling_config
)
conditionings.extend(
build_scail_conditionings(driving_latent, masks, mode=mode, strength=driving_strength)
)
return conditionings
@torch.inference_mode()
def __call__( # noqa: PLR0913
self,
prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
driving_video_path: str,
mask_path: str | None,
images: list[ImageConditioningInput] | None = None,
mode: DrivingMode = DrivingMode.ANIMATION,
driving_strength: float = 1.0,
enhance_prompt: bool = False,
tiling_config: TilingConfig | None = None,
stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio | None]:
assert_resolution(height=height, width=width, is_two_stage=True)
images = images or []
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
encode_tiling = TilingConfig.default()
(ctx_p,) = self.prompt_encoder(
[prompt],
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
masks = load_masks(mask_path, self.device, self.dtype) if mask_path is not None else None
def build_video_conditionings(out_height: int, out_width: int) -> list[ConditioningItem]:
return self.image_conditioner(
lambda enc: self._video_conditionings(
images=images,
driving_video_path=driving_video_path,
masks=masks,
height=out_height,
width=out_width,
num_frames=num_frames,
mode=mode,
driving_strength=driving_strength,
video_encoder=enc,
tiling_config=encode_tiling,
)
)
# Stage 1: low resolution.
stage_1_sigmas_t = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
stage_1_conditionings = build_video_conditionings(height // 2, width // 2)
video_state, _ = self.stage(
denoiser=SimpleDenoiser(video_context, audio_context),
sigmas=stage_1_sigmas_t,
noiser=noiser,
width=width // 2,
height=height // 2,
frames=num_frames,
fps=frame_rate,
video=ModalitySpec(context=video_context, conditionings=stage_1_conditionings),
audio=ModalitySpec(context=audio_context),
)
# Stage 2: upsample + refine.
upscaled = self.upsampler(video_state.latent[:1])
stage_2_sigmas_t = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_conditionings = build_video_conditionings(height, width)
video_state, _ = self.stage(
denoiser=SimpleDenoiser(video_context, audio_context),
sigmas=stage_2_sigmas_t,
noiser=noiser,
width=width,
height=height,
frames=num_frames,
fps=frame_rate,
video=ModalitySpec(
context=video_context,
conditionings=stage_2_conditionings,
noise_scale=stage_2_sigmas_t[0].item(),
initial_latent=upscaled,
),
audio=ModalitySpec(context=audio_context),
)
decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
return decoded_video, None
@torch.inference_mode()
def main() -> None:
logging.basicConfig(level=logging.INFO)
params = resolve_cli_params(distilled=True)
parser = scail_animation_arg_parser(params=params)
args = parser.parse_args()
if not args.lora or len(args.lora) != 1:
raise ValueError("SCAIL animation requires exactly one --lora (the SCAIL adapter).")
pipeline = ScailAnimationPipeline(
distilled_checkpoint_path=args.distilled_checkpoint_path,
gemma_root=args.gemma_root,
spatial_upsampler_path=args.spatial_upsampler_path,
scail_lora=args.lora[0],
quantization=args.quantization,
compilation_config=args.compile,
offload_mode=args.offload_mode,
)
tiling_config = TilingConfig.default()
output_shape = VideoPixelShape(
batch=1, frames=args.num_frames, width=args.width, height=args.height, fps=args.frame_rate
)
video_chunks_number = get_video_chunks_number(output_shape.frames, tiling_config)
video, _audio = pipeline(
prompt=args.prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
driving_video_path=args.driving_video,
mask_path=args.mask_path,
images=args.images if hasattr(args, "images") else [],
mode=DrivingMode(args.mode),
driving_strength=args.driving_strength,
enhance_prompt=args.enhance_prompt,
tiling_config=tiling_config,
)
encode_video(
video=video,
fps=int(args.frame_rate),
audio=None,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
if __name__ == "__main__":
main()
@@ -809,3 +809,40 @@ def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS
),
)
return parser
def scail_animation_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
"""Argument parser for the SCAIL-2 character-animation pipeline (distilled, two-stage).
Extends the standard two-stage distilled parser (checkpoint paths, --prompt, --num-frames,
--frame-rate, --height, --width, --image, --lora, --spatial-upsampler-path) with the SCAIL
driving/mask conditioning inputs. The single --lora is the SCAIL adapter.
"""
parser = default_2_stage_distilled_arg_parser(params=params)
parser.add_argument(
"--driving-video",
type=resolve_existing_path,
required=True,
help="Driving video file whose motion is transferred to the animated character.",
)
parser.add_argument(
"--mask-path",
type=resolve_existing_path,
default=None,
help="Optional .pt file with semantic masks [K+1, F_pix, H, W] (ch0 = environment switch, "
"1..K = character binding slots). Omit for a driving-only (mask-free) model.",
)
parser.add_argument(
"--mode",
type=str,
choices=["animation", "replacement"],
default="animation",
help="SCAIL conditioning mode (default: animation).",
)
parser.add_argument(
"--driving-strength",
type=float,
default=1.0,
help="Driving conditioning strength; 1.0 keeps the driving latent clean/frozen (default: 1.0).",
)
return parser
@@ -269,6 +269,7 @@ def modality_from_latent_state(
context=context,
context_mask=None,
attention_mask=state.attention_mask,
cond_channels=state.cond_channels,
)
+5
View File
@@ -23,5 +23,10 @@ adjust paths, dataset, and hyperparameters.
| **Audio Inpainting** | — | Generated | `mask` | [`audio_inpainting_lora.yaml`](./audio_inpainting_lora.yaml) |
| **A2A IC-LoRA** | — | Generated | `reference` | [`a2a_ic_lora.yaml`](./a2a_ic_lora.yaml) |
| **AV2AV IC-LoRA** | Generated | Generated | `reference` (both) | [`av2av_ic_lora.yaml`](./av2av_ic_lora.yaml) |
| **SCAIL Animation** | Generated | — | `driving` + `mask_channels` | [`scail_animation_lora.yaml`](./scail_animation_lora.yaml) |
The [`accelerate/`](./accelerate) directory holds the Accelerate launch configs (FSDP, DDP) for multi-GPU training.
> **SCAIL Animation** (SCAIL-2 character animation) also sets `model.mask_conditioning_channels: 56` to widen the
> video `patchify_proj` for the in-context mask channels; in LoRA mode that projection is unfrozen so the new columns
> train. See [Training Modes Guide](../docs/training-modes.md).
@@ -0,0 +1,132 @@
# =============================================================================
# LTX-2 SCAIL-2 Character Animation (LoRA + mask channels) Training Configuration
# =============================================================================
#
# Trains SCAIL-2-style end-to-end character animation: a driving video latent is
# concatenated into the token sequence with a RoPE width offset (ΔW), and
# in-context mask channels (1 environment switch + K binding slots) are attached
# to the driving tokens to route motion per character.
#
# This combines LoRA on the attention/FFN blocks with an *unfrozen* widened
# patchify_proj (its new mask-channel input columns cannot be reached by LoRA and
# are trained directly). Set `model.mask_conditioning_channels` to the channel
# count = temporal_factor * (K + 1) = 8 * (6 + 1) = 56 for the default K=6.
#
# Dataset structure:
# preprocessed_data_root/
# ├── latents/ # Target video latents (what the model generates)
# ├── conditions/ # Text embeddings for each video
# ├── driving_latents/ # Driving video latents (same F/H/W as target)
# └── char_masks/ # Semantic masks per sample, "mask" = [K+1, F_pix, H_pix, W_pix]
# # channel 0 = environment switch, 1..K = character binding slots
#
# =============================================================================
model:
model_path: "path/to/ltx-2-model.safetensors"
text_encoder_path: "path/to/gemma-text-encoder"
training_mode: "lora"
# Widen the video patchify_proj by 56 zero-init input columns (8 * (K+1), K=6).
# In LoRA mode the trainer additionally unfreezes patchify_proj so these train.
mask_conditioning_channels: 56
load_checkpoint: null
lora:
rank: 32
alpha: 32
dropout: 0.0
target_modules:
- "attn1.to_k"
- "attn1.to_q"
- "attn1.to_v"
- "attn1.to_out.0"
- "attn2.to_k"
- "attn2.to_q"
- "attn2.to_v"
- "attn2.to_out.0"
- "ff.net.0.proj"
- "ff.net.2"
training_strategy:
name: "flexible"
video:
is_generated: true
latents_dir: "latents"
conditions:
# SCAIL-2 driving conditioning: concatenate driving latents with a RoPE width
# offset so they stay spatially detached from the target tokens.
- type: driving
latents_dir: "driving_latents"
mode: "animation" # "animation" | "replacement"
width_offset: null # null = target pixel width (driving sits just to the right)
probability: 1.0
# In-context mask channels attached to the driving tokens (target stays zero-mask).
- type: mask_channels
mask_dir: "char_masks"
num_slots: 6 # K binding slots; channels = 8 * (K + 1) = 56 (match model.mask_conditioning_channels)
optimization:
learning_rate: 2e-4
steps: 3000
batch_size: 1
gradient_accumulation_steps: 1
max_grad_norm: 1.0
optimizer_type: "adamw"
scheduler_type: "linear"
scheduler_params: { }
enable_gradient_checkpointing: true
acceleration:
mixed_precision_mode: "bf16"
quantization: null
load_text_encoder_in_8bit: false
offload_optimizer_during_validation: false
data:
preprocessed_data_root: "/path/to/preprocessed/data"
num_dataloader_workers: 2
validation:
# NOTE: SCAIL driving + mask-channel validation conditions are not yet wired into
# the validation runner (Phase 3 training path only). Keep validation minimal /
# disabled until the inference pipeline (Phase 4) lands.
samples:
- prompt: >-
A person performing an energetic dance routine, matching the motion of the driving
performer, with crisp footwork and expressive arm movements in a bright studio.
conditions: []
negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
video_dims: [ 512, 512, 81 ]
frame_rate: 25.0
seed: 42
inference_steps: 30
interval: null # disabled: driving/mask validation lands in Phase 4
guidance_scale: 4.0
stg_scale: 1.0
stg_blocks: [29]
stg_mode: "stg_v"
generate_audio: false
skip_initial_validation: true
checkpoints:
interval: 250
keep_last_n: 3
precision: "bfloat16"
flow_matching:
timestep_sampling_mode: "shifted_logit_normal"
timestep_sampling_params: { }
hub:
push_to_hub: false
hub_model_id: null
wandb:
enabled: false
project: "ltx-2-trainer"
entity: null
tags: [ "ltx2", "scail-2", "character-animation" ]
log_validation_videos: true
seed: 42
output_dir: "outputs/scail_animation_lora"
@@ -40,6 +40,7 @@ Before diving into individual modes, here are the core ideas behind the flexible
| **Audio Inpainting** | — | Generated | `mask` | [`audio_inpainting_lora`](../configs/audio_inpainting_lora.yaml) |
| **A2A IC-LoRA** | — | Generated | `reference` | [`a2a_ic_lora`](../configs/a2a_ic_lora.yaml) |
| **AV2AV IC-LoRA** | Generated | Generated | `reference` (both) | [`av2av_ic_lora`](../configs/av2av_ic_lora.yaml) |
| **SCAIL Animation** | Generated | — | `driving` + `mask_channels` | [`scail_animation_lora`](../configs/scail_animation_lora.yaml) |
---
+1 -1
View File
@@ -54,7 +54,7 @@ build-backend = "hatchling.build"
[tool.ruff]
target-version = "1.1.7"
target-version = "py310"
line-length = 120
# Restrict isort first-party detection to src/ so stray dirs (e.g. wandb/ run output)
# next to pyproject.toml don't get classified as first-party packages. See ruff#10519.
@@ -0,0 +1,217 @@
#!/usr/bin/env python3
"""Preprocess SCAIL-2 character binding-slot masks into per-sample ``.pt`` tensors.
Each sample's input is a **label-map** video or image whose integer pixel values index
characters: ``0`` = environment/background, ``k`` in ``1..K`` = character *k* (assigned to
binding slot *k*). This script aligns the label map to the target video's latent grid (read
from the saved video-latent metadata), and emits the pixel-space semantic-mask tensor that
``ltx_core.conditioning.encode_mask_channels`` (and the trainer's ``mask_channels`` condition)
consumes:
{"mask": tensor[K+1, F_pix, H_pix, W_pix]} # ch0 = environment switch, ch1..K = binding slots
where ``F_pix = (latent_frames - 1) * 8 + 1``, ``H_pix = latent_h * 32``, ``W_pix = latent_w * 32``
(the SCAIL encoder then spatially downsamples + temporally stacks these to the 8*(K+1)=56 channels).
Label maps are resized with **nearest-neighbour** interpolation so integer labels are never
blended. The environment-switch channel (ch0) is filled uniformly with ``--environment-switch``
(``0.0`` = derive the environment from the reference image, ``1.0`` = from the driving video),
matching the paper's single-bit environment signal.
Output naming mirrors the target latents (same relative path), so ``char_masks/`` lines up
file-for-file with ``latents/`` / ``driving_latents/`` for the trainer's ``PrecomputedDataset``.
Standalone usage::
python scripts/process_char_masks.py dataset.csv \\
--mask-column char_labels --latents-dir ./latents --output-dir ./char_masks --num-slots 6
"""
from pathlib import Path
import numpy as np
import torch
import typer
from PIL import Image
# Sibling scripts (resolved via scripts/ on sys.path), reused so naming + alignment match the other stages.
from process_videos import (
IMAGE_FILE_EXTENSIONS,
VAE_SPATIAL_FACTOR,
VAE_TEMPORAL_FACTOR,
_atomic_save,
_load_paths_from_dataset,
_output_relative,
)
from rich.console import Console
from rich.progress import BarColumn, MofNCompleteColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
from ltx_trainer import logger
from ltx_trainer.video_utils import read_video
app = typer.Typer(
pretty_exceptions_enable=False,
no_args_is_help=True,
help="Preprocess SCAIL-2 character label maps into [K+1, F_pix, H_pix, W_pix] mask tensors.",
)
def _load_label_frames(mask_file: Path, pixel_f: int) -> torch.Tensor:
"""Load a label-map video/image as integer labels ``[F_pix, H, W]`` (float-typed integers).
Images are tiled across ``pixel_f`` frames. Videos are read via the shared ``read_video`` helper
(which returns ``[F, C, H, W]`` in ``[0, 1]``); labels are recovered as ``round(value * 255)``,
so label maps must be stored as small-integer grayscale (label ``k`` -> pixel value ``k``).
"""
if mask_file.suffix.lower() in IMAGE_FILE_EXTENSIONS:
arr = np.array(Image.open(mask_file).convert("L")) # exact integer labels [H, W] (writable copy)
labels = torch.from_numpy(arr).float()
return labels.unsqueeze(0).expand(pixel_f, -1, -1).contiguous()
frames, _ = read_video(str(mask_file), max_frames=pixel_f) # [F, C, H, W] in [0, 1]
labels = frames[:, 0].mul(255.0).round() # channel 0 -> integer labels [F, H, W]
if labels.shape[0] < pixel_f:
# Pad by repeating the last frame so every latent frame has a label map.
pad = labels[-1:].expand(pixel_f - labels.shape[0], -1, -1)
labels = torch.cat([labels, pad], dim=0)
return labels[:pixel_f]
def _labels_to_slot_masks(labels: torch.Tensor, num_slots: int, environment_switch: float) -> torch.Tensor:
"""Convert integer label frames ``[F, H, W]`` to ``[K+1, F, H, W]`` (ch0 env switch, ch1..K slots)."""
f_pix, h, w = labels.shape
out = torch.zeros(num_slots + 1, f_pix, h, w, dtype=torch.float32)
out[0] = environment_switch # uniform environment-switch channel
for k in range(1, num_slots + 1):
out[k] = (labels == k).float()
if bool(((labels > num_slots) & (labels > 0)).any().item()):
logger.warning(
f"Label map contains ids > num_slots ({num_slots}); those pixels are dropped (treated as environment)."
)
return out
def _resize_labels(labels: torch.Tensor, pixel_h: int, pixel_w: int) -> torch.Tensor:
"""Nearest-neighbour resize of integer label frames ``[F, H, W]`` to ``[F, pixel_h, pixel_w]``."""
if labels.shape[1:] == (pixel_h, pixel_w):
return labels
return torch.nn.functional.interpolate(
labels.unsqueeze(1), size=(pixel_h, pixel_w), mode="nearest"
).squeeze(1)
def compute_char_masks(
dataset_file: str | Path,
mask_column: str,
latents_dir: str,
output_dir: str,
num_slots: int = 6,
environment_switch: float = 0.0,
main_media_column: str | None = None,
overwrite: bool = False,
) -> None:
"""Preprocess character label maps into ``[K+1, F_pix, H_pix, W_pix]`` mask tensors.
Args:
dataset_file: Metadata file (CSV/JSON/JSONL) with a column of label-map paths.
mask_column: Column containing the per-sample label-map video/image paths.
latents_dir: Directory of target video latents (read for spatial/temporal alignment).
output_dir: Directory to write ``char_masks`` ``.pt`` files.
num_slots: Number of binding slots K (output has ``K+1`` channels).
environment_switch: Uniform value for ch0 (0.0 = env from reference, 1.0 = from driving video).
main_media_column: Column used for output naming (defaults to ``mask_column``); set it to the
target-video column so masks align with ``latents/`` when label maps live elsewhere.
overwrite: Recompute even if the output already exists.
"""
dataset_path = Path(dataset_file)
data_root = dataset_path.parent
latents_path = Path(latents_dir)
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
naming_column = main_media_column or mask_column
mask_paths = _load_paths_from_dataset(dataset_path, mask_column)
naming_paths = _load_paths_from_dataset(dataset_path, naming_column) if naming_column != mask_column else mask_paths
console = Console()
success = 0
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
MofNCompleteColumn(),
TimeElapsedColumn(),
console=console,
) as progress:
task = progress.add_task("Processing char masks", total=len(mask_paths))
for mask_file, naming_file in zip(mask_paths, naming_paths, strict=True):
progress.advance(task)
rel_path = _output_relative(naming_file, data_root)
latent_file = latents_path / rel_path.with_suffix(".pt")
out_file = output_path / rel_path.with_suffix(".pt")
if not latent_file.exists():
logger.warning(f"No target latent at {latent_file}, skipping mask {mask_file}")
continue
if not overwrite and out_file.is_file():
continue
meta = torch.load(latent_file, map_location="cpu", weights_only=True)
pixel_h = meta["height"] * VAE_SPATIAL_FACTOR
pixel_w = meta["width"] * VAE_SPATIAL_FACTOR
pixel_f = (meta["num_frames"] - 1) * VAE_TEMPORAL_FACTOR + 1
labels = _load_label_frames(mask_file, pixel_f)
labels = _resize_labels(labels, pixel_h, pixel_w)
mask = _labels_to_slot_masks(labels, num_slots, environment_switch)
out_file.parent.mkdir(parents=True, exist_ok=True)
_atomic_save({"mask": mask.contiguous()}, out_file)
success += 1
logger.info(f"Char-mask preprocessing complete: {success} masks saved to {output_path}")
@app.command()
def main(
dataset_file: str = typer.Argument(..., help="Metadata file (CSV/JSON/JSONL) with a label-map column"),
mask_column: str = typer.Option(..., help="Column of per-sample label-map video/image paths"),
latents_dir: str = typer.Option(..., help="Directory of target video latents (for alignment)"),
output_dir: str = typer.Option(..., help="Output directory for char_masks .pt files"),
num_slots: int = typer.Option(6, help="Number of character binding slots K (channels = K+1)"),
environment_switch: float = typer.Option(
0.0, help="Uniform ch0 value: 0.0 = environment from reference, 1.0 = from driving video"
),
main_media_column: str | None = typer.Option(
None, help="Column for output naming (defaults to --mask-column; set to the target-video column to align)"
),
overwrite: bool = typer.Option(False, help="Recompute even if the output already exists"),
) -> None:
"""Preprocess SCAIL-2 character label maps into ``[K+1, F_pix, H_pix, W_pix]`` mask tensors.
Example::
python scripts/process_char_masks.py dataset.csv \\
--mask-column char_labels --latents-dir ./latents \\
--output-dir ./char_masks --num-slots 6 --main-media-column media_path
"""
if not Path(dataset_file).is_file():
raise typer.BadParameter(f"Dataset file not found: {dataset_file}")
if num_slots < 1:
raise typer.BadParameter("--num-slots must be >= 1")
compute_char_masks(
dataset_file=dataset_file,
mask_column=mask_column,
latents_dir=latents_dir,
output_dir=output_dir,
num_slots=num_slots,
environment_switch=environment_switch,
main_media_column=main_media_column,
overwrite=overwrite,
)
if __name__ == "__main__":
app()
@@ -252,6 +252,14 @@ class ModelConfig(ConfigBaseModel):
description="Training mode - either LoRA fine-tuning or full model fine-tuning",
)
mask_conditioning_channels: int = Field(
default=0,
ge=0,
description="SCAIL-2 in-context mask channels. If > 0, the video patchify_proj is widened by this "
"many zero-init input columns at load time. Use with a 'driving' + 'mask_channels' condition and, "
"in LoRA mode, patchify_proj is additionally unfrozen so the new columns can train. 0 disables.",
)
load_checkpoint: str | Path | None = Field(
default=None,
description="Path to a checkpoint file or directory to load from. "
@@ -50,27 +50,39 @@ def load_transformer(
checkpoint_path: str | Path,
device: Device = "cpu",
dtype: torch.dtype = torch.bfloat16,
mask_conditioning_channels: int = 0,
) -> "LTXModel":
"""Load the LTX transformer model.
Args:
checkpoint_path: Path to the safetensors checkpoint file
device: Device to load model on
dtype: Data type for model weights
mask_conditioning_channels: If > 0, widen the video ``patchify_proj`` by this many zero-init
input columns after loading (SCAIL-2 in-context mask channels). The converted model
reproduces the base output exactly until those columns are trained.
Returns:
Loaded LTXModel transformer
"""
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
from ltx_core.model.transformer.mask_channels_checkpoint import (
widen_module_patchify_proj_for_mask_channels,
)
from ltx_core.model.transformer.model_configurator import (
LTXV_MODEL_COMFY_RENAMING_MAP,
LTXModelConfigurator,
)
return SingleGPUModelBuilder(
model = SingleGPUModelBuilder(
model_path=str(checkpoint_path),
model_class_configurator=LTXModelConfigurator,
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
).build(device=_to_torch_device(device), dtype=dtype)
if mask_conditioning_channels > 0:
widen_module_patchify_proj_for_mask_channels(model, mask_conditioning_channels)
return model
def load_video_vae_encoder(
checkpoint_path: str | Path,
@@ -399,6 +399,7 @@ class LtxvTrainer:
checkpoint_path=self._config.model.model_path,
device="cpu",
dtype=torch.bfloat16,
mask_conditioning_channels=self._config.model.mask_conditioning_channels,
)
# DDP-safe: LOCAL_RANK is set by accelerate before trainer init. Loading on bare
@@ -440,9 +441,25 @@ class LtxvTrainer:
else:
raise ValueError(f"Unknown training mode: {self._config.model.training_mode}")
# SCAIL-2: the widened patchify_proj has new mask-channel input columns that LoRA cannot reach
# (they are new base parameters, not a low-rank delta on an existing weight). Unfreeze the whole
# patchify_proj so those columns train alongside the LoRA adapters. Harmless in full mode (already
# trainable). Placed before trainable-param collection so the params below pick it up.
if self._config.model.mask_conditioning_channels > 0:
self._unfreeze_patchify_proj()
self._trainable_params = [p for p in self._transformer.parameters() if p.requires_grad]
logger.debug(f"Trainable params count: {sum(p.numel() for p in self._trainable_params):,}")
def _unfreeze_patchify_proj(self) -> None:
"""Make the video ``patchify_proj`` parameters trainable (for SCAIL-2 mask-channel columns)."""
count = 0
for name, param in self._transformer.named_parameters():
if "patchify_proj" in name and "audio_patchify_proj" not in name:
param.requires_grad_(True)
count += param.numel()
logger.info(f"Unfroze video patchify_proj for mask-channel training ({count:,} params)")
def _init_timestep_sampler(self) -> None:
"""Initialize the timestep sampler based on the config."""
sampler_cls = SAMPLERS[self._config.flow_matching.timestep_sampling_mode]
@@ -15,6 +15,7 @@ import torch
from pydantic import BaseModel, ConfigDict, Field, model_validator
from torch import Tensor
from ltx_core.conditioning import encode_mask_channels
from ltx_core.model.transformer.modality import Modality
from ltx_trainer.timestep_samplers import TimestepSampler
from ltx_trainer.training_strategies.base_strategy import (
@@ -107,6 +108,45 @@ class ReferenceConditionConfig(BaseModel):
probability: float = Field(default=1.0, ge=0.0, le=1.0, description="Probability of applying this condition")
class DrivingConditionConfig(BaseModel):
"""SCAIL-2 driving-video conditioning (concatenation with a RoPE width offset).
Driving latents are concatenated to the sequence like a reference, but their RoPE width
coordinates are shifted by ``width_offset`` (the paper's ΔW) so they stay spatially detached
from the target tokens. Driving tokens are clean (timestep=0) and excluded from loss.
"""
model_config = ConfigDict(extra="forbid")
type: Literal["driving"] = "driving"
latents_dir: str = Field(..., description="Directory for driving-video latents (same F/H/W as target)")
mode: Literal["animation", "replacement"] = Field(
default="animation",
description="SCAIL mode. Reserved for the reference/height-shift distinction; driving-token "
"placement is currently identical for both (see ltx-core VideoConditionByDrivingLatent).",
)
width_offset: float | None = Field(
default=None,
description="ΔW in RoPE pixel-space width units. None = target pixel width (driving sits just to the right).",
)
probability: float = Field(default=1.0, ge=0.0, le=1.0, description="Probability of applying this condition")
class MaskChannelsConditionConfig(BaseModel):
"""SCAIL-2 in-context mask channels attached to the driving tokens.
Encodes ``num_slots + 1`` semantic pixel masks (1 environment switch + K binding slots) into
``temporal_factor * (num_slots + 1)`` per-token conditioning channels (56 for K=6 on LTX-2) and
writes them onto the driving tokens via ``Modality.cond_channels``. The noisy target keeps an
all-zero mask (paper-faithful). Requires a ``driving`` condition and a model built with a matching
``mask_conditioning_channels``.
"""
model_config = ConfigDict(extra="forbid")
type: Literal["mask_channels"] = "mask_channels"
mask_dir: str = Field(..., description="Directory of semantic masks [K+1, F_pix, H_pix, W_pix] per sample")
num_slots: int = Field(default=6, ge=1, description="Number of character binding slots K (channels = t*(K+1))")
# Discriminated union for condition configs
ConditionConfig = Annotated[
Union[
@@ -116,6 +156,8 @@ ConditionConfig = Annotated[
SpatialCropConditionConfig,
MaskConditionConfig,
ReferenceConditionConfig,
DrivingConditionConfig,
MaskChannelsConditionConfig,
],
Field(discriminator="type"),
]
@@ -200,9 +242,9 @@ class FlexibleStrategyConfig(TrainingStrategyConfigBase):
if modality_config is None:
continue
for cond in modality_config.conditions:
if isinstance(cond, ReferenceConditionConfig):
if isinstance(cond, (ReferenceConditionConfig, DrivingConditionConfig)):
sources[cond.latents_dir] = cond.latents_dir
elif isinstance(cond, MaskConditionConfig):
elif isinstance(cond, (MaskConditionConfig, MaskChannelsConditionConfig)):
sources[cond.mask_dir] = cond.mask_dir
return sources
@@ -428,6 +470,19 @@ class FlexibleStrategy(TrainingStrategy):
modality_key=modality_key,
)
# Step 5b: Apply SCAIL-2 driving + in-context mask conditioning (video only).
noisy_latents, positions, timesteps, loss_mask, cond_channels = self._apply_scail_conditions(
modality_config=modality_config,
modality_key=modality_key,
noisy_latents=noisy_latents,
positions=positions,
timesteps=timesteps,
loss_mask=loss_mask,
data=data,
batch=batch,
device=device,
)
# Step 6: Build Modality
modality = Modality(
enabled=True,
@@ -437,6 +492,7 @@ class FlexibleStrategy(TrainingStrategy):
positions=positions,
context=prompt_embeds,
context_mask=prompt_attention_mask,
cond_channels=cond_channels,
)
return ModalityProcessingResult(
@@ -669,6 +725,150 @@ class FlexibleStrategy(TrainingStrategy):
return combined_latents, combined_positions, combined_timesteps, combined_loss_mask, targets
def _apply_scail_conditions(
self,
modality_config: ModalityConfig,
modality_key: str,
noisy_latents: Tensor,
positions: Tensor,
timesteps: Tensor,
loss_mask: Tensor | None,
data: LatentData,
batch: dict[str, Any],
device: torch.device,
) -> tuple[Tensor, Tensor, Tensor, Tensor | None, Tensor | None]:
"""Apply SCAIL-2 driving concatenation then in-context mask channels (video only).
Driving tokens are prepended (cond-first, like reference) so the target stays at the tail for
loss slicing; the mask channels are then written onto those leading driving tokens (the noisy
target keeps a zero mask). Returns the possibly-extended sequence tensors plus ``cond_channels``
(``None`` when no mask condition is present).
"""
if modality_key != "video":
return noisy_latents, positions, timesteps, loss_mask, None
driving_token_count = 0
for cond in modality_config.conditions:
if isinstance(cond, DrivingConditionConfig):
noisy_latents, positions, timesteps, loss_mask, driving_token_count = self._apply_driving_condition(
noisy_latents=noisy_latents,
positions=positions,
timesteps=timesteps,
loss_mask=loss_mask,
target_width=data.width,
batch=batch,
config=cond,
)
cond_channels = None
for cond in modality_config.conditions:
if isinstance(cond, MaskChannelsConditionConfig):
cond_channels = self._build_mask_channels(
config=cond,
batch=batch,
total_tokens=noisy_latents.shape[1],
driving_token_count=driving_token_count,
target_frames=data.num_frames,
target_height=data.height,
target_width=data.width,
device=device,
)
return noisy_latents, positions, timesteps, loss_mask, cond_channels
def _apply_driving_condition(
self,
noisy_latents: Tensor,
positions: Tensor,
timesteps: Tensor,
loss_mask: Tensor | None,
target_width: int,
batch: dict[str, Any],
config: DrivingConditionConfig,
) -> tuple[Tensor, Tensor, Tensor, Tensor | None, int]:
"""Prepend SCAIL-2 driving latents with a RoPE width offset (ΔW).
Driving latents share the target's frame/height/width, so their time and height coordinates
already match the target (same ``_get_video_positions`` grid); only the width axis is shifted
by ``width_offset`` so the driving tokens stay spatially detached. Driving tokens are clean
(timestep=0) and excluded from loss. Returns the driving token count for mask placement.
The apply/skip decision is batch-wide (concatenation changes the sequence length) but drawn
from the torch RNG for reproducibility, mirroring ``_apply_reference_condition``.
"""
if torch.rand((), device=noisy_latents.device).item() >= config.probability:
return noisy_latents, positions, timesteps, loss_mask, 0
cond = self._patchify_latent_data(batch[config.latents_dir], "video")
drv_latents = cond.latents
batch_size, drv_seq_len, _ = drv_latents.shape
device = drv_latents.device
dtype = drv_latents.dtype
drv_positions = self._get_video_positions(
num_frames=cond.num_frames,
height=cond.height,
width=cond.width,
batch_size=batch_size,
fps=cond.fps,
device=device,
).clone()
width_offset = (
config.width_offset
if config.width_offset is not None
else float(target_width * VIDEO_SCALE_FACTORS.width)
)
drv_positions[:, 2, ...] = drv_positions[:, 2, ...] + width_offset
drv_timesteps = torch.zeros(batch_size, drv_seq_len, device=device, dtype=dtype)
drv_loss_mask = torch.zeros(batch_size, drv_seq_len, dtype=torch.bool, device=device)
combined_latents = torch.cat([drv_latents, noisy_latents], dim=1)
combined_positions = torch.cat([drv_positions, positions], dim=2)
combined_timesteps = torch.cat([drv_timesteps, timesteps], dim=1)
combined_loss_mask = torch.cat([drv_loss_mask, loss_mask], dim=1) if loss_mask is not None else None
return combined_latents, combined_positions, combined_timesteps, combined_loss_mask, drv_seq_len
def _build_mask_channels(
self,
config: MaskChannelsConditionConfig,
batch: dict[str, Any],
total_tokens: int,
driving_token_count: int,
target_frames: int,
target_height: int,
target_width: int,
device: torch.device,
) -> Tensor:
"""Encode semantic masks into per-token channels written onto the leading driving tokens.
The mask tensor at ``batch[mask_dir]["mask"]`` is expected as ``[B, K+1, F_pix, H_pix, W_pix]``
(channel 0 = environment switch, ``1..K`` = binding slots). It is encoded to
``temporal_factor * (K+1)`` per-token channels (56 for K=6) via ``encode_mask_channels`` and
placed on the driving tokens; the target tokens keep an all-zero mask (paper-faithful).
"""
masks = batch[config.mask_dir]["mask"].to(device=device)
encoded = encode_mask_channels(
masks=masks,
temporal_factor=VIDEO_SCALE_FACTORS.time,
height_lat=target_height,
width_lat=target_width,
frames_lat=target_frames,
)
tokens = self._video_patchifier.patchify(encoded) # [B, N, C_mask]
batch_size, n_region, c_mask = tokens.shape
cond_channels = tokens.new_zeros(batch_size, total_tokens, c_mask)
if driving_token_count == 0:
# No driving tokens (e.g. driving skipped by its probability): fall back to writing the
# mask onto the leading target tokens so the channel width still matches the model.
cond_channels[:, :n_region] = tokens
else:
if n_region != driving_token_count:
raise ValueError(
f"mask token count ({n_region}) must equal the driving token count "
f"({driving_token_count}); driving and mask must describe the same grid."
)
cond_channels[:, :driving_token_count] = tokens
return cond_channels
@staticmethod
def _compute_modality_loss(pred: Tensor, targets: Tensor, loss_mask: Tensor) -> Tensor:
"""Compute per-element MSE loss for a single modality. Returns [B,]."""