Automated PR - 2026-01-29

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

@@ -207,6 +207,14 @@ class ValidationConfig(ConfigBaseModel):
         "One video path must be provided for each validation prompt",
     )
 
+    reference_downscale_factor: int = Field(
+        default=1,
+        description="Downscale factor for reference videos in IC-LoRA validation. "
+        "When > 1, reference videos are processed at 1/n resolution (e.g., 2 means half resolution). "
+        "Must match the factor used during dataset preprocessing.",
+        ge=1,
+    )
+
     video_dims: tuple[int, int, int] = Field(
         default=(960, 544, 97),
         description="Dimensions of validation videos (width, height, frames). "
@@ -334,6 +342,41 @@ class ValidationConfig(ConfigBaseModel):
 
         return v
 
+    @model_validator(mode="after")
+    def validate_scaled_reference_dimensions(self) -> "ValidationConfig":
+        """Validate that scaled reference dimensions are valid when reference_downscale_factor > 1."""
+        if self.reference_downscale_factor > 1:
+            width, height, _frames = self.video_dims
+
+            # Validate that downscale factor evenly divides the target dimensions
+            if width % self.reference_downscale_factor != 0:
+                raise ValueError(
+                    f"Width {width} is not evenly divisible by reference_downscale_factor "
+                    f"{self.reference_downscale_factor}. Choose a downscale factor that divides {width} evenly."
+                )
+            if height % self.reference_downscale_factor != 0:
+                raise ValueError(
+                    f"Height {height} is not evenly divisible by reference_downscale_factor "
+                    f"{self.reference_downscale_factor}. Choose a downscale factor that divides {height} evenly."
+                )
+
+            scaled_width = width // self.reference_downscale_factor
+            scaled_height = height // self.reference_downscale_factor
+
+            # Validate scaled dimensions are divisible by 32
+            if scaled_width % 32 != 0:
+                raise ValueError(
+                    f"Scaled reference width {scaled_width} (from {width} / {self.reference_downscale_factor}) "
+                    f"is not divisible by 32. Choose a different downscale factor or adjust video_dims."
+                )
+            if scaled_height % 32 != 0:
+                raise ValueError(
+                    f"Scaled reference height {scaled_height} (from {height} / {self.reference_downscale_factor}) "
+                    f"is not divisible by 32. Choose a different downscale factor or adjust video_dims."
+                )
+
+        return self
+
 
 class CheckpointsConfig(ConfigBaseModel):
     """Configuration for model checkpointing during training"""

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

@@ -218,16 +218,22 @@ def load_text_encoder(
     from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
     from ltx_core.text_encoders.gemma.encoders.av_encoder import (
         AV_GEMMA_TEXT_ENCODER_KEY_OPS,
+        GEMMA_MODEL_OPS,
         AVGemmaTextEncoderModelConfigurator,
     )
     from ltx_core.text_encoders.gemma.encoders.base_encoder import module_ops_from_gemma_root
+    from ltx_core.utils import find_matching_file
 
     torch_device = _to_torch_device(device)
+
+    gemma_model_folder = find_matching_file(str(gemma_model_path), "model*.safetensors").parent
+    gemma_weight_paths = [str(p) for p in gemma_model_folder.rglob("*.safetensors")]
+
     text_encoder = SingleGPUModelBuilder(
-        model_path=str(checkpoint_path),
+        model_path=(str(checkpoint_path), *gemma_weight_paths),
         model_class_configurator=AVGemmaTextEncoderModelConfigurator,
         model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
-        module_ops=module_ops_from_gemma_root(str(gemma_model_path)),
+        module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(str(gemma_model_path))),
     ).build(device=torch_device, dtype=dtype)
 
     return text_encoder

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

@@ -795,6 +795,7 @@ class LtxvTrainer:
                 seed=self._config.validation.seed,
                 condition_image=condition_image,
                 reference_video=reference_video,
+                reference_downscale_factor=self._config.validation.reference_downscale_factor,
                 generate_audio=generate_audio,
                 include_reference_in_output=self._config.validation.include_reference_in_output,
                 cached_embeddings=cached_embeddings,
@@ -885,8 +886,11 @@ class LtxvTrainer:
             # Cast to configured precision
             state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in state_dict.items()}
 
-            # Save to disk
-            save_file(state_dict, saved_weights_path)
+            # Build metadata for safetensors file
+            metadata = self._build_checkpoint_metadata()
+
+            # Save to disk with metadata
+            save_file(state_dict, saved_weights_path, metadata=metadata)
         else:
             # Cast to configured precision
             full_state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in full_state_dict.items()}
@@ -913,6 +917,21 @@ class LtxvTrainer:
             # Update the list to only contain kept checkpoints
             self._checkpoint_paths = self._checkpoint_paths[-self._config.checkpoints.keep_last_n :]
 
+    def _build_checkpoint_metadata(self) -> dict[str, str]:
+        """Build metadata dictionary for safetensors checkpoint.
+        Delegates to the training strategy to get strategy-specific metadata
+        that downstream inference pipelines may need.
+        Returns:
+            Dictionary of string key-value pairs for safetensors metadata.
+            Values are converted to strings for safetensors compatibility.
+        """
+        raw_metadata = self._training_strategy.get_checkpoint_metadata()
+        # Convert all values to strings for safetensors compatibility
+        metadata = {k: str(v) for k, v in raw_metadata.items()}
+        if metadata:
+            logger.info(f"Saving checkpoint metadata: {metadata}")
+        return metadata
+
     def _save_config(self) -> None:
         """Save the training configuration as a YAML file in the output directory."""
         if not IS_MAIN_PROCESS:

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

@@ -128,6 +128,15 @@ class TrainingStrategy(ABC):
             Scalar loss tensor
         """
 
+    def get_checkpoint_metadata(self) -> dict[str, Any]:
+        """Get strategy-specific metadata to include in checkpoint files.
+        Override this method in subclasses to add custom metadata,
+        e.g. any parameters that a downstream inference pipeline may need.
+        Returns:
+            Dictionary of metadata key-value pairs (values must be JSON-serializable)
+        """
+        return {}
+
     def _get_video_positions(
         self,
         num_frames: int,

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

@@ -46,9 +46,13 @@ class VideoToVideoStrategy(TrainingStrategy):
     - Reference latents (clean) are concatenated with target latents (noised)
     - Video coordinates handle both reference and target sequences
     - Loss is computed only on the target portion
+    Attributes:
+        reference_downscale_factor: The inferred downscale factor of reference videos.
+            This is computed from the first batch and cached for metadata export.
     """
 
     config: VideoToVideoConfig
+    reference_downscale_factor: int | None
 
     def __init__(self, config: VideoToVideoConfig):
         """Initialize strategy with configuration.
@@ -56,6 +60,7 @@ class VideoToVideoStrategy(TrainingStrategy):
             config: Video-to-video configuration
         """
         super().__init__(config)
+        self.reference_downscale_factor = None  # Will be inferred from first batch
 
     def get_data_sources(self) -> dict[str, str]:
         """IC-LoRA training requires latents, conditions, and reference latents."""
@@ -65,7 +70,7 @@ class VideoToVideoStrategy(TrainingStrategy):
             self.config.reference_latents_dir: "ref_latents",
         }
 
-    def prepare_training_inputs(
+    def prepare_training_inputs(  # noqa: PLR0915
         self,
         batch: dict[str, Any],
         timestep_sampler: TimestepSampler,
@@ -86,6 +91,26 @@ class VideoToVideoStrategy(TrainingStrategy):
         ref_height = ref_latents_info["height"][0].item()
         ref_width = ref_latents_info["width"][0].item()
 
+        # Infer reference downscale factor from dimension ratios
+        # This allows training with downscaled reference videos for efficiency
+        reference_downscale_factor = self._infer_reference_downscale_factor(
+            target_height=height,
+            target_width=width,
+            ref_height=ref_height,
+            ref_width=ref_width,
+        )
+
+        # Cache the scale factor for metadata export (only on first batch)
+        if self.reference_downscale_factor is None:
+            self.reference_downscale_factor = reference_downscale_factor
+        elif self.reference_downscale_factor != reference_downscale_factor:
+            raise ValueError(
+                f"Inconsistent reference downscale factor across batches. "
+                f"First batch had factor={self.reference_downscale_factor}, "
+                f"but current batch has factor={reference_downscale_factor}. "
+                f"All training samples must use the same reference/target resolution ratio."
+            )
+
         # Patchify latents: [B, C, F, H, W] -> [B, seq_len, C]
         target_latents = self._video_patchifier.patchify(target_latents)
         ref_latents = self._video_patchifier.patchify(ref_latents)
@@ -159,6 +184,15 @@ class VideoToVideoStrategy(TrainingStrategy):
             dtype=dtype,
         )
 
+        # Scale reference positions to match target coordinate space
+        # This maps ref positions from (0, ref_H, ref_W) to (0, target_H, target_W)
+        # Position tensor shape: [B, 3, seq_len, 2] where dim 1 is (time, height, width)
+        if reference_downscale_factor != 1:
+            ref_positions = ref_positions.clone()
+            ref_positions[:, 1, ...] *= reference_downscale_factor  # height axis
+            ref_positions[:, 2, ...] *= reference_downscale_factor  # width axis
+            # Time axis (index 0) remains unchanged
+
         target_positions = self._get_video_positions(
             num_frames=num_frames,
             height=height,
@@ -221,3 +255,48 @@ class VideoToVideoStrategy(TrainingStrategy):
         loss = loss.mul(loss_mask).div(loss_mask.mean())
 
         return loss.mean()
+
+    def get_checkpoint_metadata(self) -> dict[str, Any]:
+        """Get metadata for checkpoint files."""
+        metadata: dict[str, Any] = {}
+        # Always include reference_downscale_factor for IC-LoRAs so inference
+        # pipelines know the expected scale factor for reference videos.
+        if self.reference_downscale_factor is not None:
+            metadata["reference_downscale_factor"] = self.reference_downscale_factor
+        return metadata
+
+    @staticmethod
+    def _infer_reference_downscale_factor(
+        target_height: int,
+        target_width: int,
+        ref_height: int,
+        ref_width: int,
+    ) -> int:
+        """Infer the reference downscale factor from target and reference dimensions."""
+        # If dimensions match, no scaling needed
+        if target_height == ref_height and target_width == ref_width:
+            return 1
+
+        # Calculate scale factors for each dimension
+        if target_height % ref_height != 0 or target_width % ref_width != 0:
+            raise ValueError(
+                f"Target dimensions ({target_height}x{target_width}) must be exact multiples "
+                f"of reference dimensions ({ref_height}x{ref_width})"
+            )
+
+        scale_h = target_height // ref_height
+        scale_w = target_width // ref_width
+
+        if scale_h != scale_w:
+            raise ValueError(
+                f"Reference scale must be uniform. Got height scale {scale_h} and width scale {scale_w}. "
+                f"Target: {target_height}x{target_width}, Reference: {ref_height}x{ref_width}"
+            )
+
+        if scale_h < 1:
+            raise ValueError(
+                f"Reference dimensions ({ref_height}x{ref_width}) cannot be larger than "
+                f"target dimensions ({target_height}x{target_width})"
+            )
+
+        return scale_h

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

@@ -85,6 +85,7 @@ class GenerationConfig:
     seed: int = 42  # Random seed for reproducibility
     condition_image: Tensor | None = None  # Optional first frame image for image-to-video
     reference_video: Tensor | None = None  # For IC-LoRA: [F, C, H, W] in [0, 1]
+    reference_downscale_factor: int = 1  # For IC-LoRA: downscale factor (1 = same resolution, 2 = half resolution)
     generate_audio: bool = True  # Whether to generate audio alongside video
     include_reference_in_output: bool = False  # For IC-LoRA: concatenate original reference with generated output
     cached_embeddings: CachedPromptEmbeddings | None = None  # Pre-computed text embeddings (avoids loading Gemma)
@@ -251,6 +252,14 @@ class ValidationSampler:
         ref_latent, ref_positions = self._encode_video(ref_video_preprocessed, config.frame_rate, device)
         ref_seq_len = ref_latent.shape[1]
 
+        # Scale reference positions to match target coordinate space
+        # Position tensor shape: [B, 3, seq_len, 2] where dim 1 is (time, height, width)
+        if config.reference_downscale_factor != 1:
+            ref_positions = ref_positions.clone()
+            ref_positions[:, 1, ...] *= config.reference_downscale_factor  # height axis
+            ref_positions[:, 2, ...] *= config.reference_downscale_factor  # width axis
+            # Time axis (index 0) remains unchanged
+
         # Create target video state
         video_tools = self._create_video_latent_tools(config)
         target_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
@@ -375,13 +384,28 @@ class ValidationSampler:
     @staticmethod
     def _preprocess_reference_video(config: GenerationConfig) -> Tensor:
         """Preprocess reference video: resize, crop, and convert to model input format.
+        When reference_downscale_factor > 1, the reference video is downscaled to a smaller
+        resolution for more efficient inference. The positions will be scaled up later
+        to match the target coordinate space.
         Args:
-            config: Generation configuration with reference_video
+            config: Generation configuration
         Returns:
             Preprocessed video tensor [B, C, F, H, W] in [-1, 1] range
         """
         ref_video = config.reference_video  # [F, C, H, W] in [0, 1]
-        target_height, target_width = config.height, config.width
+        scale_factor = config.reference_downscale_factor
+
+        # Target dimensions for reference (scaled down if scale_factor > 1)
+        target_height = config.height // scale_factor
+        target_width = config.width // scale_factor
+
+        # Validate scaled dimensions
+        if target_height % 32 != 0 or target_width % 32 != 0:
+            raise ValueError(
+                f"Scaled reference dimensions ({target_height}x{target_width}) must be divisible by 32. "
+                f"Original: {config.height}x{config.width}, scale_factor: {scale_factor}"
+            )
+
         current_height, current_width = ref_video.shape[2:]
 
         # Resize maintaining aspect ratio and center crop if needed
@@ -745,11 +769,28 @@ class ValidationSampler:
         If the videos have different frame counts, the shorter one is padded with
         its last frame repeated.
         Args:
-            left_video: Left video tensor [C, F1, H, W] in [0, 1]
-            right_video: Right video tensor [C, F2, H, W] in [0, 1]
+            left_video: Left video tensor [C, F1, H1, W1] in [0, 1]
+            right_video: Right video tensor [C, F2, H2, W2] in [0, 1]
         Returns:
-            Concatenated video tensor [C, max(F1,F2), H, W*2] in [0, 1]
+            Concatenated video tensor [C, max(F1,F2), H2, W1_scaled+W2] in [0, 1]
         """
+        left_height, left_width = left_video.shape[2], left_video.shape[3]
+        right_height = right_video.shape[2]
+
+        # Resize left video to match right video's height if needed
+        if left_height != right_height:
+            # Scale width proportionally to maintain aspect ratio
+            scale = right_height / left_height
+            new_width = int(left_width * scale)
+            # Interpolate expects [N, C, H, W], we have [C, F, H, W]
+            # Reshape to [C*F, 1, H, W] -> interpolate -> reshape back
+            c, f, h, w = left_video.shape
+            left_video = left_video.reshape(c * f, 1, h, w)
+            left_video = torch.nn.functional.interpolate(
+                left_video, size=(right_height, new_width), mode="bilinear", align_corners=False
+            )
+            left_video = left_video.reshape(c, f, right_height, new_width)
+
         left_frames = left_video.shape[1]
         right_frames = right_video.shape[1]