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>

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,]."""