From baa6646fd1fc3ad029eae9ed73bcb33e18196c11 Mon Sep 17 00:00:00 2001 From: indigo Date: Thu, 9 Jul 2026 09:06:35 +0800 Subject: [PATCH] Add SCAIL-2 driving-latent conditioning (Phase 1, inference PoC) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- docs/plan.md | 61 +++++++ docs/tasks.md | 42 +++++ .../src/ltx_core/conditioning/__init__.py | 4 + .../ltx_core/conditioning/types/__init__.py | 3 + .../conditioning/types/driving_video_cond.py | 160 ++++++++++++++++++ 5 files changed, 270 insertions(+) create mode 100644 docs/plan.md create mode 100644 docs/tasks.md create mode 100644 packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py diff --git a/docs/plan.md b/docs/plan.md new file mode 100644 index 0000000..6c57ff7 --- /dev/null +++ b/docs/plan.md @@ -0,0 +1,61 @@ +# SCAIL-2 → LTX-2 移植計畫 + +將 SCAIL-2（arXiv: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,Tv−1], H=[0,Hv), W=[0,Wv) | +| z_driv | T=[1,Tv], H=[0,Hv), **W=[ΔW, ΔW+Wv)** | T=[0,Tv−1], 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（機制 2，checkpoint 手術）⬜ 未開始 +- `model.py:158` `patchify_proj` 由 `Linear(128, inner)` 加寬到 `Linear(128+28, inner)`。 +- `LatentState` 增加 optional conditioning-channel 欄位，跟著 patchify/concat/clear 流動。 +- 投影前 concat（`transformer_args.py:209`）。 +- 新輸入欄位 **zero-init**，載入舊 checkpoint 行為不變；寫 checkpoint 轉換 script。 +- 新增產生 28-channel mask（環境開關 + 角色槽）的 conditioning item。 + +### Phase 3 — 訓練整合（ltx-trainer）⬜ 未開始 +- dataset 產出 (target, driving, mask)，接上 Phase 1/2 conditioning，設微調 loss 與凍結策略。 + +### Phase 4 — Pipeline + CLI 包裝 ⬜ 未開始 +- 仿 `lipdub.py` 寫 `scail_animation.py` pipeline + arg parser。 + +## 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 2，enum 先保留佔位。 + +## 參考 +- 論文：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` diff --git a/docs/tasks.md b/docs/tasks.md new file mode 100644 index 0000000..2fa0d7f --- /dev/null +++ b/docs/tasks.md @@ -0,0 +1,42 @@ +# 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 channel（checkpoint 手術） + +| # | 任務 | 狀態 | 備註 | +|---|---|---|---| +| 2.1 | `patchify_proj` 加寬 `128 → 128+28` | ⬜ | `model.py:158` `_init_video`；`proj_out` 不動 | +| 2.2 | `LatentState` 帶額外 conditioning-channel 欄位 | ⬜ | 同步改 `tools.py` patchify/unpatchify/clear、`Modality` | +| 2.3 | 投影前 concat mask channel 到 `x` | ⬜ | `transformer_args.py:209` | +| 2.4 | checkpoint zero-init 轉換 script | ⬜ | 新輸入欄位 zero-init，載入舊權重行為不變 | +| 2.5 | 產生 28-channel mask 的 conditioning item | ⬜ | 環境開關 + K=6 角色綁定槽，`4(K+1)=28` | +| 2.6 | Replacement Mode z_ref 高度位移 ΔH_ref | ⬜ | Phase 1 暫緩項 | + +## Phase 3 — 訓練整合（ltx-trainer） + +| # | 任務 | 狀態 | 備註 | +|---|---|---|---| +| 3.1 | dataset 產出 (target, driving, mask) | ⬜ | | +| 3.2 | 接上 conditioning + 微調 loss / 凍結策略 | ⬜ | | + +## Phase 4 — Pipeline + CLI + +| # | 任務 | 狀態 | 備註 | +|---|---|---|---| +| 4.1 | `scail_animation.py` pipeline + arg parser | ⬜ | 仿 `lipdub.py` | + +## 決議紀錄 +- **範圍**：先只做 Phase 1（推論期 PoC）。Phase 2+ 待 Phase 1 驗證後再討論。 +- **架構前提**：SCAIL-2 建構於 Wan 2.1，本移植為跨架構移植。 diff --git a/packages/ltx-core/src/ltx_core/conditioning/__init__.py b/packages/ltx-core/src/ltx_core/conditioning/__init__.py index 8d3eb9f..f229a1a 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/__init__.py +++ b/packages/ltx-core/src/ltx_core/conditioning/__init__.py @@ -5,6 +5,8 @@ from ltx_core.conditioning.item import ConditioningItem from ltx_core.conditioning.types import ( AudioConditionByReferenceLatent, ConditioningItemAttentionStrengthWrapper, + DrivingMode, + VideoConditionByDrivingLatent, VideoConditionByKeyframeIndex, VideoConditionByLatentIndex, VideoConditionByMask, @@ -16,6 +18,8 @@ __all__ = [ "ConditioningError", "ConditioningItem", "ConditioningItemAttentionStrengthWrapper", + "DrivingMode", + "VideoConditionByDrivingLatent", "VideoConditionByKeyframeIndex", "VideoConditionByLatentIndex", "VideoConditionByMask", diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py b/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py index bd8d962..e94564c 100644 --- a/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py +++ b/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py @@ -1,6 +1,7 @@ """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_cond import VideoConditionByMask @@ -10,6 +11,8 @@ from ltx_core.conditioning.types.reference_video_cond import VideoConditionByRef __all__ = [ "AudioConditionByReferenceLatent", "ConditioningItemAttentionStrengthWrapper", + "DrivingMode", + "VideoConditionByDrivingLatent", "VideoConditionByKeyframeIndex", "VideoConditionByLatentIndex", "VideoConditionByMask", diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py b/packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py new file mode 100644 index 0000000..4fb7bc0 --- /dev/null +++ b/packages/ltx-core/src/ltx_core/conditioning/types/driving_video_cond.py @@ -0,0 +1,160 @@ +"""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.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, + )