packages/ltx-trainer/pyproject.toml

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

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

@@ -470,36 +470,18 @@ 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 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(
@@ -743,6 +725,56 @@ 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,