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>

docs/plan.md

@@ -45,8 +45,12 @@
 - `VideoConditionByMaskChannels`：(K+1) 語意 mask → 空間下採樣 + 時間 8× 堆疊 → 寫入尾端 driving token（target 保持零，符合論文）。
 - 驗證通過（`verify_mask_channels.py`）：向後相容、zero-init 加寬 forward == baseline、mask pipeline 不 crash。**僅驗證 plumbing，畫質需 Phase 3 微調。**
 
-### Phase 3 — 訓練整合（ltx-trainer）⬜ 未開始
+### Phase 3 — 訓練整合（ltx-trainer）✅ 已完成（程式碼路徑；實訓需 GPU）
-- dataset 產出 (target, driving, mask)，接上 Phase 1/2 conditioning，設微調 loss 與凍結策略。
+- **決策**：完整整合到 `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 包裝 ⬜ 未開始
 - 仿 `lipdub.py` 寫 `scail_animation.py` pipeline + arg parser。

docs/tasks.md

@@ -30,10 +30,19 @@
 
 ## Phase 3 — 訓練整合（ltx-trainer）
 
+> **決策**：完整整合到 `FlexibleStrategy`；訓練方式 = **LoRA + 解凍 patchify_proj**（新 mask 欄位無法純 LoRA 訓練）。**本機無 GPU/Linux/checkpoint，只做到 CPU 單元驗證**，實訓需在 GPU 機器跑。
+
 | # | 任務 | 狀態 | 備註 |
 |---|---|---|---|
-| 3.1 | dataset 產出 (target, driving, mask) | ⬜ | |
+| 3.1 | 新增 Driving/MaskChannels ConditionConfig | ✅ | `flexible.py`：`DrivingConditionConfig`(latents_dir/mode/width_offset) + `MaskChannelsConditionConfig`(mask_dir/num_slots)，加入 union + `get_data_sources` |
-| 3.2 | 接上 conditioning + 微調 loss / 凍結策略 | ⬜ | |
+| 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） | ⬜ | 需 process_dataset 產出 `driving_latents/`（同 target 形狀）與 `char_masks/`(mask=[K+1,F_pix,H,W])；語意 mask 需分割模型，屬資料工程，未做 |
+| 3.8 | validation runner 接 driving/mask | ⬜ | 驗證期取樣尚未接 SCAIL 條件（config 內 validation 先停用），與 Phase 4 一起 |
+| 3.9 | 實機訓練跑通 | ⬜ | 需 Linux + GPU + checkpoint，本機無法 |
 
 ## Phase 4 — Pipeline + CLI
 

packages/ltx-core/src/ltx_core/model/transformer/mask_channels_checkpoint.py

@@ -22,6 +22,41 @@ _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,

packages/ltx-trainer/configs/README.md

@@ -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).

packages/ltx-trainer/configs/scail_animation_lora.yaml

@@ -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"

packages/ltx-trainer/docs/training-modes.md

@@ -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) |
 
 ---
 

packages/ltx-trainer/src/ltx_trainer/config.py

@@ -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. "

packages/ltx-trainer/src/ltx_trainer/model_loader.py

@@ -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,

packages/ltx-trainer/src/ltx_trainer/trainer.py

@@ -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]

packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py

@@ -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,37 @@ class FlexibleStrategy(TrainingStrategy):
                     modality_key=modality_key,
                 )
 
+        # Step 5b: Apply SCAIL-2 driving conditioning (concatenation with a RoPE width offset).
+        # Driving tokens are prepended (cond-first, like reference) so the target stays at the tail
+        # for loss slicing; ``driving_token_count`` marks how many leading tokens are driving.
+        driving_token_count = 0
+        for cond in modality_config.conditions:
+            if isinstance(cond, DrivingConditionConfig) and modality_key == "video":
+                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,
+                )
+
+        # Step 5c: Build SCAIL-2 in-context mask channels on the driving tokens (target stays zero).
+        cond_channels = None
+        for cond in modality_config.conditions:
+            if isinstance(cond, MaskChannelsConditionConfig) and modality_key == "video":
+                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,
+                )
+
         # Step 6: Build Modality
         modality = Modality(
             enabled=True,
@@ -437,6 +510,7 @@ class FlexibleStrategy(TrainingStrategy):
             positions=positions,
             context=prompt_embeds,
             context_mask=prompt_attention_mask,
+            cond_channels=cond_channels,
         )
 
         return ModalityProcessingResult(
@@ -669,6 +743,100 @@ class FlexibleStrategy(TrainingStrategy):
 
         return combined_latents, combined_positions, combined_timesteps, combined_loss_mask, targets
 
+    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,]."""