Automated PR - 2026-01-15

packages/ltx-trainer/configs/ltx2_av_lora_low_vram.yaml

@@ -176,7 +176,7 @@ acceleration:
 
   # Load text encoder in 8-bit precision to save memory
   # Useful when GPU memory is limited
-  load_text_encoder_in_8bit: false
+  load_text_encoder_in_8bit: true
 
 # -----------------------------------------------------------------------------
 # Data Configuration

packages/ltx-trainer/docs/configuration-reference.md

@@ -25,7 +25,8 @@ sub-configurations:
 Check out our example configurations in the `configs` directory:
 
 - 📄 [Audio-Video LoRA Training](../configs/ltx2_av_lora.yaml) - Joint audio-video generation training
-- 📄 [Audio-Video LoRA Training (Low VRAM)](../configs/ltx2_av_lora_low_vram.yaml) - Memory-optimized config for 32GB GPUs (uses 8-bit optimizer, INT8 quantization, and reduced LoRA rank)
+- 📄 [Audio-Video LoRA Training (Low VRAM)](../configs/ltx2_av_lora_low_vram.yaml) - Memory-optimized config for 32GB
+  GPUs (uses 8-bit optimizer, INT8 quantization, and reduced LoRA rank)
 - 📄 [IC-LoRA Training](../configs/ltx2_v2v_ic_lora.yaml) - Video-to-video transformation training
 
 ## ⚙️ Configuration Sections
@@ -300,10 +301,10 @@ checkpoints:
 
 **Key parameters:**
 
-| Parameter     | Description                                                            |
-|---------------|------------------------------------------------------------------------|
-| `interval`    | Steps between intermediate checkpoint saves (set to `null` to disable) |
-| `keep_last_n` | Number of most recent checkpoints to keep (-1 = keep all)              |
+| Parameter     | Description                                                                   |
+|---------------|-------------------------------------------------------------------------------|
+| `interval`    | Steps between intermediate checkpoint saves (set to `null` to disable)        |
+| `keep_last_n` | Number of most recent checkpoints to keep (-1 = keep all)                     |
 | `precision`   | Precision for saved checkpoint weights: `"bfloat16"` (default) or `"float32"` |
 
 ### HubConfig

packages/ltx-trainer/scripts/decode_latents.py

@@ -15,6 +15,7 @@ import torch
 import torchaudio
 import torchvision.utils
 import typer
+from einops import rearrange
 from rich.console import Console
 from rich.progress import (
     BarColumn,
@@ -27,10 +28,16 @@ from rich.progress import (
 )
 from transformers.utils.logging import disable_progress_bar
 
+from ltx_core.model.video_vae import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
 from ltx_trainer import logger
 from ltx_trainer.model_loader import load_audio_vae_decoder, load_video_vae_decoder, load_vocoder
 from ltx_trainer.video_utils import save_video
 
+DEFAULT_TILE_SIZE_PIXELS = 512  # Spatial tile size in pixels (must be ≥64 and divisible by 32)
+DEFAULT_TILE_OVERLAP_PIXELS = 128  # Spatial tile overlap in pixels (must be divisible by 32)
+DEFAULT_TILE_SIZE_FRAMES = 128  # Temporal tile size in frames (must be ≥16 and divisible by 8)
+DEFAULT_TILE_OVERLAP_FRAMES = 24  # Temporal tile overlap in frames (must be divisible by 8)
+
 disable_progress_bar()
 console = Console()
 app = typer.Typer(
@@ -60,19 +67,15 @@ class LatentsDecoder:
         self.vae = None
         self.audio_vae = None
         self.vocoder = None
-        self._load_model(model_path, vae_tiling, with_audio)
+        self.vae_tiling = vae_tiling
 
-    def _load_model(self, model_path: str, vae_tiling: bool, with_audio: bool = False) -> None:
+        self._load_model(model_path, with_audio)
+
+    def _load_model(self, model_path: str, with_audio: bool = False) -> None:
         """Initialize and load the VAE model(s)."""
         with console.status(f"[bold]Loading video VAE decoder from {model_path}...", spinner="dots"):
             self.vae = load_video_vae_decoder(model_path, device=self.device, dtype=torch.bfloat16)
 
-            if vae_tiling:
-                logger.warning(
-                    "VAE tiling is not yet implemented in this script. "
-                    "Continuing without tiling - this may cause OOM errors for large resolutions."
-                )
-
         if with_audio:
             with console.status(f"[bold]Loading audio VAE decoder from {model_path}...", spinner="dots"):
                 self.audio_vae = load_audio_vae_decoder(model_path, device=self.device, dtype=torch.bfloat16)
@@ -125,61 +128,6 @@ class LatentsDecoder:
 
         logger.info(f"Decoding complete! Videos saved to {output_dir}")
 
-    def _process_file(self, latent_file: Path, output_dir: Path, seed: int | None) -> None:
-        """Process a single latent file."""
-        # Load the latent data
-        data = torch.load(latent_file, map_location=self.device, weights_only=False)
-
-        # Get latents - handle both old patchified [seq_len, C] and new [C, F, H, W] formats
-        latents = data["latents"]
-        num_frames = data["num_frames"]
-        height = data["height"]
-        width = data["width"]
-
-        # Check if latents need reshaping (old patchified format)
-        if latents.dim() == 2:
-            # Old format: [seq_len, C] -> reshape to [C, F, H, W]
-            _seq_len, channels = latents.shape
-            latents = latents.reshape(num_frames, height, width, channels)
-            latents = latents.permute(3, 0, 1, 2)  # [F, H, W, C] -> [C, F, H, W]
-
-        # Add batch dimension: [C, F, H, W] -> [1, C, F, H, W]
-        latents = latents.unsqueeze(0).to(device=self.device, dtype=torch.bfloat16)
-
-        # Create generator only if seed is provided
-        generator = None
-        if seed is not None:
-            generator = torch.Generator(device=self.device)
-            generator.manual_seed(seed)
-
-        # Decode the video (VAE decoder uses forward/call, not decode method)
-        video = self.vae(latents)  # [B, C, F, H, W]
-
-        # Convert to [F, C, H, W] format and normalize to [0, 1]
-        video = video[0]  # Remove batch dimension -> [C, F, H, W]
-        video = video.permute(1, 0, 2, 3)  # [C, F, H, W] -> [F, C, H, W]
-        video = (video + 1) / 2  # Denormalize from [-1, 1] to [0, 1]
-        video = video.clamp(0, 1)
-
-        # Determine output format and save
-        is_image = video.shape[0] == 1
-        if is_image:
-            # Save as PNG for single frame
-            output_path = output_dir / f"{latent_file.stem}.png"
-            torchvision.utils.save_image(
-                video[0],  # [C, H, W] in [0, 1]
-                str(output_path),
-            )
-        else:
-            # Save as MP4 for video using PyAV-based save_video
-            output_path = output_dir / f"{latent_file.stem}.mp4"
-            fps = data.get("fps", 24)  # Use stored FPS or default to 24
-            save_video(
-                video_tensor=video,  # [F, C, H, W] in [0, 1]
-                output_path=output_path,
-                fps=fps,
-            )
-
     @torch.inference_mode()
     def decode_audio(self, latents_dir: Path, output_dir: Path) -> None:
         """Decode all audio latent files in the directory recursively.
@@ -229,6 +177,87 @@ class LatentsDecoder:
 
         logger.info(f"Audio decoding complete! Audio files saved to {output_dir}")
 
+    def _process_file(self, latent_file: Path, output_dir: Path, seed: int | None) -> None:
+        """Process a single latent file."""
+        # Load the latent data
+        data = torch.load(latent_file, map_location=self.device, weights_only=False)
+
+        # Get latents - handle both old patchified [seq_len, C] and new [C, F, H, W] formats
+        latents = data["latents"]
+        num_frames = data["num_frames"]
+        height = data["height"]
+        width = data["width"]
+
+        # Check if latents need reshaping (old patchified format)
+        if latents.dim() == 2:
+            # Old format: [seq_len, C] -> reshape to [C, F, H, W]
+            latents = rearrange(latents, "(f h w) c -> c f h w", f=num_frames, h=height, w=width)
+
+        # Add batch dimension: [C, F, H, W] -> [1, C, F, H, W]
+        latents = latents.unsqueeze(0).to(device=self.device, dtype=torch.bfloat16)
+
+        # Create generator only if seed is provided
+        generator = None
+        if seed is not None:
+            generator = torch.Generator(device=self.device)
+            generator.manual_seed(seed)
+
+        # Decode the video
+        video = self._decode_video(latents, generator)
+
+        # Determine output format and save
+        is_image = video.shape[0] == 1
+        if is_image:
+            # Save as PNG for single frame
+            output_path = output_dir / f"{latent_file.stem}.png"
+            torchvision.utils.save_image(
+                video[0],  # [C, H, W] in [0, 1]
+                str(output_path),
+            )
+        else:
+            # Save as MP4 for video using PyAV-based save_video
+            output_path = output_dir / f"{latent_file.stem}.mp4"
+            fps = data.get("fps", 24)  # Use stored FPS or default to 24
+            save_video(
+                video_tensor=video,  # [F, C, H, W] in [0, 1]
+                output_path=output_path,
+                fps=fps,
+            )
+
+    def _decode_video(self, latents: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
+        """Decode latents to video frames."""
+        if self.vae_tiling:
+            # Use tiled decoding for reduced VRAM
+            tiling_config = TilingConfig(
+                spatial_config=SpatialTilingConfig(
+                    tile_size_in_pixels=DEFAULT_TILE_SIZE_PIXELS,
+                    tile_overlap_in_pixels=DEFAULT_TILE_OVERLAP_PIXELS,
+                ),
+                temporal_config=TemporalTilingConfig(
+                    tile_size_in_frames=DEFAULT_TILE_SIZE_FRAMES,
+                    tile_overlap_in_frames=DEFAULT_TILE_OVERLAP_FRAMES,
+                ),
+            )
+            chunks = list(
+                self.vae.tiled_decode(
+                    latents,
+                    tiling_config=tiling_config,
+                    generator=generator,
+                )
+            )
+            # Concatenate along temporal dimension
+            video = torch.cat(chunks, dim=2)  # [B, C, F, H, W]
+        else:
+            # Standard full decoding
+            video = self.vae(latents, generator=generator)  # [B, C, F, H, W]
+
+        # Convert to [F, C, H, W] format and normalize to [0, 1]
+        video = rearrange(video, "1 c f h w -> f c h w")
+        video = (video + 1) / 2  # Denormalize from [-1, 1] to [0, 1]
+        video = video.clamp(0, 1)
+
+        return video
+
     def _process_audio_file(self, latent_file: Path, output_dir: Path) -> None:
         """Process a single audio latent file."""
         # Load the latent data
@@ -242,8 +271,7 @@ class LatentsDecoder:
         if latents.dim() == 2:
             # Old format: [seq_len, channels] where seq_len = time * freq
             # Reshape to [C, T, F]
-            latents = latents.reshape(num_time_steps, freq_bins, -1)  # [T, F, C]
-            latents = latents.permute(2, 0, 1)  # [T, F, C] -> [C, T, F]
+            latents = rearrange(latents, "(t f) c -> c t f", t=num_time_steps, f=freq_bins)
 
         # Add batch dimension: [C, T, F] -> [1, C, T, F]
         latents = latents.unsqueeze(0)

packages/ltx-trainer/scripts/process_captions.py

@@ -220,7 +220,7 @@ class CaptionsDataset(Dataset):
                     break
 
 
-def compute_captions_embeddings(
+def compute_captions_embeddings(  # noqa: PLR0913
     dataset_file: str | Path,
     output_dir: str,
     model_path: str,
@@ -231,6 +231,7 @@ def compute_captions_embeddings(
     remove_llm_prefixes: bool = False,
     batch_size: int = 8,
     device: str = "cuda",
+    load_in_8bit: bool = False,
 ) -> None:
     """
     Process captions and save text embeddings.
@@ -245,6 +246,7 @@ def compute_captions_embeddings(
         remove_llm_prefixes: Whether to remove common LLM-generated prefixes
         batch_size: Batch size for processing
         device: Device to use for computation
+        load_in_8bit: Whether to load the Gemma text encoder in 8-bit precision
     """
 
     console = Console()
@@ -264,7 +266,13 @@ def compute_captions_embeddings(
 
     # Load text encoder
     with console.status("[bold]Loading Gemma text encoder...", spinner="dots"):
-        text_encoder = load_text_encoder(model_path, text_encoder_path, device=device, dtype=torch.bfloat16)
+        text_encoder = load_text_encoder(
+            model_path,
+            text_encoder_path,
+            device=device,
+            dtype=torch.bfloat16,
+            load_in_8bit=load_in_8bit,
+        )
 
     logger.info("Text encoder loaded successfully")
 
@@ -326,7 +334,7 @@ def compute_captions_embeddings(
 
 
 @app.command()
-def main(
+def main(  # noqa: PLR0913
     dataset_file: str = typer.Argument(
         ...,
         help="Path to metadata file (CSV/JSON/JSONL) containing captions and media paths",
@@ -368,6 +376,10 @@ def main(
         default=False,
         help="Remove common LLM-generated prefixes from captions",
     ),
+    load_text_encoder_in_8bit: bool = typer.Option(
+        default=False,
+        help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
+    ),
 ) -> None:
     """Process text captions and save embeddings for video generation training.
     This script processes captions from metadata files and saves text embeddings
@@ -408,6 +420,7 @@ def main(
         remove_llm_prefixes=remove_llm_prefixes,
         batch_size=batch_size,
         device=device,
+        load_in_8bit=load_text_encoder_in_8bit,
     )
 
 

packages/ltx-trainer/scripts/process_dataset.py

@@ -20,6 +20,7 @@ from process_videos import compute_latents, parse_resolution_buckets
 from rich.console import Console
 
 from ltx_trainer import logger
+from ltx_trainer.gpu_utils import free_gpu_memory_context
 
 console = Console()
 app = typer.Typer(
@@ -46,15 +47,9 @@ def preprocess_dataset(  # noqa: PLR0913
     remove_llm_prefixes: bool = False,
     reference_column: str | None = None,
     with_audio: bool = False,
+    load_text_encoder_in_8bit: bool = False,
 ) -> None:
     """Run the preprocessing pipeline with the given arguments."""
-    # VAE tiling is not yet implemented
-    if vae_tiling:
-        logger.warning(
-            "VAE tiling is not yet implemented in this script. "
-            "Continuing without tiling - this may cause OOM errors for large resolutions."
-        )
-
     # Validate dataset file
     _validate_dataset_file(dataset_file)
 
@@ -66,19 +61,21 @@ def preprocess_dataset(  # noqa: PLR0913
     if lora_trigger:
         logger.info(f'LoRA trigger word "{lora_trigger}" will be prepended to all captions')
 
-    # Process captions using the dedicated function
-    compute_captions_embeddings(
-        dataset_file=dataset_file,
-        output_dir=str(conditions_dir),
-        model_path=model_path,
-        text_encoder_path=text_encoder_path,
-        caption_column=caption_column,
-        media_column=video_column,
-        lora_trigger=lora_trigger,
-        remove_llm_prefixes=remove_llm_prefixes,
-        batch_size=batch_size,
-        device=device,
-    )
+    with free_gpu_memory_context():
+        # Process captions using the dedicated function
+        compute_captions_embeddings(
+            dataset_file=dataset_file,
+            output_dir=str(conditions_dir),
+            model_path=model_path,
+            text_encoder_path=text_encoder_path,
+            caption_column=caption_column,
+            media_column=video_column,
+            lora_trigger=lora_trigger,
+            remove_llm_prefixes=remove_llm_prefixes,
+            batch_size=batch_size,
+            device=device,
+            load_in_8bit=load_text_encoder_in_8bit,
+        )
 
     # Process videos using the dedicated function
     audio_latents_dir = None
@@ -86,36 +83,37 @@ def preprocess_dataset(  # noqa: PLR0913
         logger.info("Audio preprocessing enabled - will extract and encode audio from videos")
         audio_latents_dir = output_base / "audio_latents"
 
-    compute_latents(
-        dataset_file=dataset_file,
-        video_column=video_column,
-        resolution_buckets=resolution_buckets,
-        output_dir=str(latents_dir),
-        model_path=model_path,
-        batch_size=batch_size,
-        device=device,
-        vae_tiling=vae_tiling,
-        with_audio=with_audio,
-        audio_output_dir=str(audio_latents_dir) if audio_latents_dir else None,
-    )
-
-    # Process reference videos if reference_column is provided
-    if reference_column:
-        logger.info("Processing reference videos for IC-LoRA training...")
-        reference_latents_dir = output_base / "reference_latents"
-
+    with free_gpu_memory_context():
         compute_latents(
             dataset_file=dataset_file,
-            main_media_column=video_column,
-            video_column=reference_column,
+            video_column=video_column,
             resolution_buckets=resolution_buckets,
-            output_dir=str(reference_latents_dir),
+            output_dir=str(latents_dir),
             model_path=model_path,
             batch_size=batch_size,
             device=device,
             vae_tiling=vae_tiling,
+            with_audio=with_audio,
+            audio_output_dir=str(audio_latents_dir) if audio_latents_dir else None,
         )
 
+        # Process reference videos if reference_column is provided
+        if reference_column:
+            logger.info("Processing reference videos for IC-LoRA training...")
+            reference_latents_dir = output_base / "reference_latents"
+
+            compute_latents(
+                dataset_file=dataset_file,
+                main_media_column=video_column,
+                video_column=reference_column,
+                resolution_buckets=resolution_buckets,
+                output_dir=str(reference_latents_dir),
+                model_path=model_path,
+                batch_size=batch_size,
+                device=device,
+                vae_tiling=vae_tiling,
+            )
+
     # Handle decoding if requested (for verification)
     if decode:
         logger.info("Decoding latents for verification...")
@@ -224,6 +222,10 @@ def main(  # noqa: PLR0913
         default=False,
         help="Extract and encode audio from video files",
     ),
+    load_text_encoder_in_8bit: bool = typer.Option(
+        default=False,
+        help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
+    ),
 ) -> None:
     """Preprocess a video dataset by computing and saving latents and text embeddings.
     The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.
@@ -269,6 +271,7 @@ def main(  # noqa: PLR0913
         remove_llm_prefixes=remove_llm_prefixes,
         reference_column=reference_column,
         with_audio=with_audio,
+        load_text_encoder_in_8bit=load_text_encoder_in_8bit,
     )
 
 

packages/ltx-trainer/scripts/process_videos.py

@@ -59,6 +59,9 @@ VAE_TEMPORAL_FACTOR = 8
 AUDIO_LATENT_CHANNELS = 8
 AUDIO_FREQUENCY_BINS = 16
 
+DEFAULT_TILE_SIZE = 512  # Spatial tile size in pixels (must be ≥64 and divisible by 32)
+DEFAULT_TILE_OVERLAP = 128  # Spatial tile overlap in pixels (must be divisible by 32)
+
 app = typer.Typer(
     pretty_exceptions_enable=False,
     no_args_is_help=True,
@@ -488,12 +491,6 @@ def compute_latents(  # noqa: PLR0913, PLR0915
     with console.status(f"[bold]Loading video VAE encoder from [cyan]{model_path}[/]...", spinner="dots"):
         vae = load_video_vae_encoder(model_path, device=torch_device, dtype=torch.bfloat16)
 
-    if vae_tiling:
-        logger.warning(
-            "VAE tiling is not yet implemented in this script. "
-            "Continuing without tiling - this may cause OOM errors for large resolutions."
-        )
-
     # Load audio VAE encoder and audio processor if needed
     audio_vae_encoder = None
     audio_processor = None
@@ -543,7 +540,7 @@ def compute_latents(  # noqa: PLR0913, PLR0915
 
             # Encode video
             with torch.inference_mode():
-                video_latent_data = encode_video(vae=vae, video=video)
+                video_latent_data = encode_video(vae=vae, video=video, use_tiling=vae_tiling)
 
             # Save latents for each item in batch
             for i in range(len(batch["relative_path"])):
@@ -611,6 +608,9 @@ def encode_video(
     vae: torch.nn.Module,
     video: torch.Tensor,
     dtype: torch.dtype | None = None,
+    use_tiling: bool = False,
+    tile_size: int = DEFAULT_TILE_SIZE,
+    tile_overlap: int = DEFAULT_TILE_OVERLAP,
 ) -> dict[str, torch.Tensor | int]:
     """Encode video into non-patchified latent representation.
     Args:
@@ -618,6 +618,9 @@ def encode_video(
         video: Input tensor of shape [B, C, F, H, W] (batch, channels, frames, height, width)
                This is the format expected by the VAE encoder.
         dtype: Target dtype for output latents
+        use_tiling: Whether to use spatial tiling for memory efficiency
+        tile_size: Tile size in pixels (must be divisible by 32)
+        tile_overlap: Overlap between tiles in pixels (must be divisible by 32)
     Returns:
         Dict containing non-patchified latents and shape information:
         {
@@ -636,8 +639,17 @@ def encode_video(
 
     video = video.to(device=device, dtype=vae_dtype)
 
-    # Encode video - VAE expects [B, C, F, H, W], returns [B, C, F', H', W']
-    latents = vae(video)
+    # Choose encoding method based on tiling flag
+    if use_tiling:
+        latents = tiled_encode_video(
+            vae=vae,
+            video=video,
+            tile_size=tile_size,
+            tile_overlap=tile_overlap,
+        )
+    else:
+        # Encode video - VAE expects [B, C, F, H, W], returns [B, C, F', H', W']
+        latents = vae(video)
 
     if dtype is not None:
         latents = latents.to(dtype=dtype)
@@ -652,6 +664,161 @@ def encode_video(
     }
 
 
+def tiled_encode_video(  # noqa: PLR0912, PLR0915
+    vae: torch.nn.Module,
+    video: torch.Tensor,
+    tile_size: int = DEFAULT_TILE_SIZE,
+    tile_overlap: int = DEFAULT_TILE_OVERLAP,
+) -> torch.Tensor:
+    """Encode video using spatial tiling for memory efficiency.
+    Splits the video into overlapping spatial tiles, encodes each tile separately,
+    and blends the results using linear feathering in the overlap regions.
+    Args:
+        vae: Video VAE encoder model
+        video: Input tensor of shape [B, C, F, H, W]
+        tile_size: Tile size in pixels (must be divisible by 32)
+        tile_overlap: Overlap between tiles in pixels (must be divisible by 32)
+    Returns:
+        Encoded latent tensor [B, C_latent, F_latent, H_latent, W_latent]
+    """
+    batch, _channels, frames, height, width = video.shape
+    device = video.device
+    dtype = video.dtype
+
+    # Validate tile parameters
+    if tile_size % VAE_SPATIAL_FACTOR != 0:
+        raise ValueError(f"tile_size must be divisible by {VAE_SPATIAL_FACTOR}, got {tile_size}")
+    if tile_overlap % VAE_SPATIAL_FACTOR != 0:
+        raise ValueError(f"tile_overlap must be divisible by {VAE_SPATIAL_FACTOR}, got {tile_overlap}")
+    if tile_overlap >= tile_size:
+        raise ValueError(f"tile_overlap ({tile_overlap}) must be less than tile_size ({tile_size})")
+
+    # If video fits in a single tile, use regular encoding
+    if height <= tile_size and width <= tile_size:
+        return vae(video)
+
+    # Calculate output dimensions
+    # VAE compresses: H -> H/32, W -> W/32, F -> 1 + (F-1)/8
+    output_height = height // VAE_SPATIAL_FACTOR
+    output_width = width // VAE_SPATIAL_FACTOR
+    output_frames = 1 + (frames - 1) // VAE_TEMPORAL_FACTOR
+
+    # Latent channels (128 for LTX-2)
+    # Get from a small test encode or assume 128
+    latent_channels = 128
+
+    # Initialize output and weight tensors
+    output = torch.zeros(
+        (batch, latent_channels, output_frames, output_height, output_width),
+        device=device,
+        dtype=dtype,
+    )
+    weights = torch.zeros(
+        (batch, 1, output_frames, output_height, output_width),
+        device=device,
+        dtype=dtype,
+    )
+
+    # Calculate tile positions with overlap
+    # Step size is tile_size - tile_overlap
+    step_h = tile_size - tile_overlap
+    step_w = tile_size - tile_overlap
+
+    h_positions = list(range(0, max(1, height - tile_overlap), step_h))
+    w_positions = list(range(0, max(1, width - tile_overlap), step_w))
+
+    # Ensure last tile covers the edge
+    if h_positions[-1] + tile_size < height:
+        h_positions.append(height - tile_size)
+    if w_positions[-1] + tile_size < width:
+        w_positions.append(width - tile_size)
+
+    # Remove duplicates and sort
+    h_positions = sorted(set(h_positions))
+    w_positions = sorted(set(w_positions))
+
+    # Overlap in latent space
+    overlap_out_h = tile_overlap // VAE_SPATIAL_FACTOR
+    overlap_out_w = tile_overlap // VAE_SPATIAL_FACTOR
+
+    # Process each tile
+    for h_pos in h_positions:
+        for w_pos in w_positions:
+            # Calculate tile boundaries in input space
+            h_start = max(0, h_pos)
+            w_start = max(0, w_pos)
+            h_end = min(h_start + tile_size, height)
+            w_end = min(w_start + tile_size, width)
+
+            # Ensure tile dimensions are divisible by VAE_SPATIAL_FACTOR
+            tile_h = ((h_end - h_start) // VAE_SPATIAL_FACTOR) * VAE_SPATIAL_FACTOR
+            tile_w = ((w_end - w_start) // VAE_SPATIAL_FACTOR) * VAE_SPATIAL_FACTOR
+
+            if tile_h < VAE_SPATIAL_FACTOR or tile_w < VAE_SPATIAL_FACTOR:
+                continue
+
+            # Adjust end positions
+            h_end = h_start + tile_h
+            w_end = w_start + tile_w
+
+            # Extract tile
+            tile = video[:, :, :, h_start:h_end, w_start:w_end]
+
+            # Encode tile
+            encoded_tile = vae(tile)
+
+            # Get actual encoded dimensions
+            _, _, tile_out_frames, tile_out_height, tile_out_width = encoded_tile.shape
+
+            # Calculate output positions
+            out_h_start = h_start // VAE_SPATIAL_FACTOR
+            out_w_start = w_start // VAE_SPATIAL_FACTOR
+            out_h_end = min(out_h_start + tile_out_height, output_height)
+            out_w_end = min(out_w_start + tile_out_width, output_width)
+
+            # Trim encoded tile if necessary
+            actual_tile_h = out_h_end - out_h_start
+            actual_tile_w = out_w_end - out_w_start
+            encoded_tile = encoded_tile[:, :, :, :actual_tile_h, :actual_tile_w]
+
+            # Create blending mask with linear feathering at edges
+            mask = torch.ones(
+                (1, 1, tile_out_frames, actual_tile_h, actual_tile_w),
+                device=device,
+                dtype=dtype,
+            )
+
+            # Apply feathering at edges (linear blend in overlap regions)
+            # Left edge
+            if h_pos > 0 and overlap_out_h > 0 and overlap_out_h < actual_tile_h:
+                fade_in = torch.linspace(0.0, 1.0, overlap_out_h + 2, device=device, dtype=dtype)[1:-1]
+                mask[:, :, :, :overlap_out_h, :] *= fade_in.view(1, 1, 1, -1, 1)
+
+            # Right edge (bottom in height dimension)
+            if h_end < height and overlap_out_h > 0 and overlap_out_h < actual_tile_h:
+                fade_out = torch.linspace(1.0, 0.0, overlap_out_h + 2, device=device, dtype=dtype)[1:-1]
+                mask[:, :, :, -overlap_out_h:, :] *= fade_out.view(1, 1, 1, -1, 1)
+
+            # Top edge (left in width dimension)
+            if w_pos > 0 and overlap_out_w > 0 and overlap_out_w < actual_tile_w:
+                fade_in = torch.linspace(0.0, 1.0, overlap_out_w + 2, device=device, dtype=dtype)[1:-1]
+                mask[:, :, :, :, :overlap_out_w] *= fade_in.view(1, 1, 1, 1, -1)
+
+            # Bottom edge (right in width dimension)
+            if w_end < width and overlap_out_w > 0 and overlap_out_w < actual_tile_w:
+                fade_out = torch.linspace(1.0, 0.0, overlap_out_w + 2, device=device, dtype=dtype)[1:-1]
+                mask[:, :, :, :, -overlap_out_w:] *= fade_out.view(1, 1, 1, 1, -1)
+
+            # Accumulate weighted results
+            output[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += encoded_tile * mask
+            weights[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += mask
+
+    # Normalize by weights (avoid division by zero)
+    output = output / (weights + 1e-8)
+
+    return output
+
+
 def encode_audio(
     audio_vae_encoder: torch.nn.Module,
     audio_processor: torch.nn.Module,

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

@@ -5,7 +5,7 @@ This module provides captioning capabilities for videos with audio using:
 - Gemini Flash: Cloud-based API for audio-visual captioning
 Requirements:
 - Qwen2.5-Omni: transformers>=4.50, torch
-- Gemini Flash: google-generativeai (pip install google-generativeai)
+- Gemini Flash: google-generativeai (uv pip install google-generativeai)
   Set GEMINI_API_KEY or GOOGLE_API_KEY environment variable
 """
 

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

@@ -0,0 +1,135 @@
+# ruff: noqa: PLC0415
+
+"""
+8-bit Gemma text encoder loading utilities.
+This module provides functionality for loading the Gemma text encoder in 8-bit precision
+using bitsandbytes, which significantly reduces GPU memory usage.
+Example usage:
+    from ltx_trainer.gemma_8bit import load_8bit_gemma
+    text_encoder = load_8bit_gemma(
+        checkpoint_path="/path/to/ltx2.safetensors",
+        gemma_model_path="/path/to/gemma",
+    )
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Generator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import torch
+
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
+from ltx_core.text_encoders.gemma.embeddings_connector import Embeddings1DConnectorConfigurator
+from ltx_core.text_encoders.gemma.encoders.av_encoder import (
+    AV_GEMMA_TEXT_ENCODER_KEY_OPS,
+    AVGemmaTextEncoderModel,
+)
+from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
+from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
+
+if TYPE_CHECKING:
+    from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
+
+
+def load_8bit_gemma(
+    checkpoint_path: str | Path,
+    gemma_model_path: str | Path,
+    dtype: torch.dtype = torch.bfloat16,
+) -> "AVGemmaTextEncoderModel":
+    """Load the Gemma text encoder in 8-bit precision using bitsandbytes.
+    This function bypasses ltx-core's standard loading path to enable 8-bit quantization
+    via the bitsandbytes library. The Gemma model is loaded with load_in_8bit=True and
+    torch_dtype=bfloat16, while the feature extractor and connector weights are loaded
+    from the LTX-2 checkpoint.
+    Args:
+        checkpoint_path: Path to the LTX-2 safetensors checkpoint file
+        gemma_model_path: Path to Gemma model directory
+        dtype: Data type for non-quantized model weights (feature extractor, connectors)
+    Returns:
+        Loaded AVGemmaTextEncoderModel with 8-bit quantized Gemma backbone
+    Raises:
+        ImportError: If bitsandbytes is not installed
+        FileNotFoundError: If required model files are not found
+    """
+    try:
+        from transformers import BitsAndBytesConfig, Gemma3ForConditionalGeneration
+    except ImportError as e:
+        raise ImportError(
+            "8-bit text encoder loading requires bitsandbytes. Install it with: uv pip install bitsandbytes"
+        ) from e
+
+    # Find paths within gemma_model_path
+    gemma_path = _find_gemma_subpath(gemma_model_path, "model*.safetensors")
+    tokenizer_path = _find_gemma_subpath(gemma_model_path, "tokenizer.model")
+
+    quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+    with _suppress_accelerate_memory_warnings():
+        gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
+            gemma_path,
+            quantization_config=quantization_config,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            local_files_only=True,
+        )
+
+    # Load tokenizer
+    tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024)
+
+    # Load config and weights from the LTX-2 checkpoint
+    loader = SafetensorsModelStateDictLoader()
+    config = loader.metadata(str(checkpoint_path))
+    sd = loader.load(str(checkpoint_path), sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS)
+
+    # Helper to extract state dict for a given prefix
+    def extract_state_dict(prefix: str) -> dict[str, torch.Tensor]:
+        return {k.replace(prefix, ""): v for k, v in sd.sd.items() if k.startswith(prefix)}
+
+    # Create and load feature extractor
+    feature_extractor = GemmaFeaturesExtractorProjLinear()
+    feature_extractor.load_state_dict(extract_state_dict("feature_extractor_linear."))
+    feature_extractor = feature_extractor.to(device=gemma_model.device, dtype=dtype)
+
+    # Create and load video embeddings connector
+    embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
+    embeddings_connector.load_state_dict(extract_state_dict("embeddings_connector."))
+    embeddings_connector = embeddings_connector.to(device=gemma_model.device, dtype=dtype)
+
+    # Create and load audio embeddings connector
+    audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
+    audio_embeddings_connector.load_state_dict(extract_state_dict("audio_embeddings_connector."))
+    audio_embeddings_connector = audio_embeddings_connector.to(device=gemma_model.device, dtype=dtype)
+
+    # Construct the text encoder
+    text_encoder = AVGemmaTextEncoderModel(
+        feature_extractor_linear=feature_extractor,
+        embeddings_connector=embeddings_connector,
+        audio_embeddings_connector=audio_embeddings_connector,
+        tokenizer=tokenizer,
+        model=gemma_model,
+    )
+
+    return text_encoder
+
+
+def _find_gemma_subpath(root_path: str | Path, pattern: str) -> str:
+    """Find a file matching a glob pattern and return its parent directory."""
+    matches = list(Path(root_path).rglob(pattern))
+    if not matches:
+        raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}")
+    return str(matches[0].parent)
+
+
+@contextmanager
+def _suppress_accelerate_memory_warnings() -> Generator[None, None, None]:
+    """Temporarily suppress INFO warnings from accelerate about memory allocation."""
+    accelerate_logger = logging.getLogger("accelerate.utils.modeling")
+    old_level = accelerate_logger.level
+    accelerate_logger.setLevel(logging.WARNING)
+    try:
+        yield
+    finally:
+        accelerate_logger.setLevel(old_level)

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

@@ -0,0 +1,90 @@
+"""GPU memory management utilities for training and inference."""
+
+import functools
+import gc
+import subprocess
+from typing import Callable, TypeVar
+
+import torch
+
+from ltx_trainer import logger
+
+F = TypeVar("F", bound=Callable)
+
+
+def free_gpu_memory(log: bool = False) -> None:
+    """Free GPU memory by running garbage collection and emptying CUDA cache.
+    Args:
+        log: If True, log memory stats after clearing
+    """
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        if log:
+            allocated = torch.cuda.memory_allocated() / 1024**3
+            reserved = torch.cuda.memory_reserved() / 1024**3
+            logger.debug(f"GPU memory freed. Allocated: {allocated:.2f}GB, Reserved: {reserved:.2f}GB")
+
+
+class free_gpu_memory_context:  # noqa: N801
+    """Context manager and decorator to free GPU memory before and/or after execution.
+    Can be used as a decorator:
+        @free_gpu_memory_context(after=True)
+        def my_function():
+            ...
+    Or as a context manager:
+        with free_gpu_memory_context():
+            heavy_operation()
+    Args:
+        before: Free memory before execution (default: False)
+        after: Free memory after execution (default: True)
+        log: Log memory stats when freeing (default: False)
+    """
+
+    def __init__(self, *, before: bool = False, after: bool = True, log: bool = False) -> None:
+        self.before = before
+        self.after = after
+        self.log = log
+
+    def __enter__(self) -> "free_gpu_memory_context":
+        if self.before:
+            free_gpu_memory(log=self.log)
+        return self
+
+    def __exit__(self, exc_type: type | None, exc_val: Exception | None, exc_tb: object) -> None:
+        if self.after:
+            free_gpu_memory(log=self.log)
+
+    def __call__(self, func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs) -> object:
+            with self:
+                return func(*args, **kwargs)
+
+        return wrapper  # type: ignore
+
+
+def get_gpu_memory_gb(device: torch.device) -> float:
+    """Get current GPU memory usage in GB using nvidia-smi.
+    Args:
+        device: torch.device to get memory usage for
+    Returns:
+        Current GPU memory usage in GB
+    """
+    try:
+        device_id = device.index if device.index is not None else 0
+        result = subprocess.check_output(
+            [
+                "nvidia-smi",
+                "--query-gpu=memory.used",
+                "--format=csv,nounits,noheader",
+                "-i",
+                str(device_id),
+            ],
+            encoding="utf-8",
+        )
+        return float(result.strip()) / 1024  # Convert MB to GB
+    except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e:
+        logger.error(f"Failed to get GPU memory from nvidia-smi: {e}")
+        # Fallback to torch
+        return torch.cuda.memory_allocated(device) / 1024**3

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

@@ -191,6 +191,7 @@ def load_text_encoder(
     gemma_model_path: str | Path,
     device: Device = "cpu",
     dtype: torch.dtype = torch.bfloat16,
+    load_in_8bit: bool = False,
 ) -> "AVGemmaTextEncoderModel":
     """Load the Gemma text encoder.
     Args:
@@ -198,9 +199,22 @@ def load_text_encoder(
         gemma_model_path: Path to Gemma model directory
         device: Device to load model on
         dtype: Data type for model weights
+        load_in_8bit: Whether to load the Gemma model in 8-bit precision using bitsandbytes.
+            When True, the model is loaded with device_map="auto" and the device argument
+            is ignored for the Gemma backbone (feature extractor still uses dtype).
     Returns:
         Loaded AVGemmaTextEncoderModel
     """
+    if not Path(gemma_model_path).is_dir():
+        raise ValueError(f"Gemma model path is not a directory: {gemma_model_path}")
+
+    # Use 8-bit loading path if requested
+    if load_in_8bit:
+        from ltx_trainer.gemma_8bit import load_8bit_gemma
+
+        return load_8bit_gemma(checkpoint_path, gemma_model_path, dtype)
+
+    # Standard loading path
     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,
@@ -208,9 +222,6 @@ def load_text_encoder(
     )
     from ltx_core.text_encoders.gemma.encoders.base_encoder import module_ops_from_gemma_root
 
-    if not Path(gemma_model_path).is_dir():
-        raise ValueError(f"Gemma model path is not a directory: {gemma_model_path}")
-
     torch_device = _to_torch_device(device)
     text_encoder = SingleGPUModelBuilder(
         model_path=str(checkpoint_path),

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

@@ -1,8 +1,9 @@
-# Adapted from: https://github.com/bghira/SimpleTuner/blob/main/helpers/training/quantisation/__init__.py
+# Adapted from: https://github.com/bghira/SimpleTuner
+# With improvements from: https://github.com/ostris/ai-toolkit
 from typing import Literal
 
 import torch
-from optimum.quanto import qtype
+from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
 
 from ltx_trainer import logger
 
@@ -14,55 +15,161 @@ QuantizationOptions = Literal[
     "fp8uz-quanto",
 ]
 
+# Modules to exclude from quantization.
+# These are glob patterns passed to quanto's `exclude` parameter.
+# When quantizing the full model at once, these patterns match against full module paths.
+# When quantizing block-by-block, we also use SKIP_ROOT_MODULES for top-level modules.
+EXCLUDE_PATTERNS = [
+    # Input/output projection layers
+    "patchify_proj",
+    "audio_patchify_proj",
+    "proj_out",
+    "audio_proj_out",
+    # Timestep embedding layers - int4 tinygemm requires strict bfloat16 input
+    # and these receive float32 sinusoidal embeddings that are cast to bfloat16
+    "*adaln*",
+    "time_proj",
+    "timestep_embedder*",
+    # Caption/text projection layers
+    "caption_projection*",
+    "audio_caption_projection*",
+    # Normalization layers (usually excluded from quantization)
+    "*norm*",
+]
+
+# Top-level modules to skip entirely during block-by-block quantization.
+# These are exact matches against model.named_children() names.
+# (Needed because quanto's exclude patterns don't work when calling quantize() directly on a module)
+SKIP_ROOT_MODULES = {
+    "patchify_proj",
+    "audio_patchify_proj",
+    "proj_out",
+    "audio_proj_out",
+    "audio_caption_projection",
+}
+
 
 def quantize_model(
     model: torch.nn.Module,
     precision: QuantizationOptions,
     quantize_activations: bool = False,
+    device: torch.device | str | None = None,
 ) -> torch.nn.Module:
     """
-    Quantize a model using the specified precision settings.
+    Quantize a model using optimum-quanto.
+    For large models with transformer_blocks, this function quantizes block-by-block
+    on GPU then moves back to CPU, which is much faster than quantizing on CPU and
+    uses less peak VRAM than loading the entire model to GPU at once.
     Args:
         model: The model to quantize.
-        precision: The precision level to quantize to (e.g. "int8-quanto", "fp8-quanto").
+        precision: The quantization precision (e.g. "int8-quanto", "fp8-quanto").
         quantize_activations: Whether to quantize activations in addition to weights.
+        device: Device to use for quantization. If None, uses CUDA if available, else CPU.
     Returns:
-        The quantized model, or the original model if no quantization is performed.
+        The quantized model.
     """
     from optimum.quanto import freeze, quantize  # noqa: PLC0415
 
-    weight_quant = _quanto_type_map(precision)
-    extra_quanto_args = {
-        "exclude": [
-            # Input/output projection layers
-            "patchify_proj",
-            "audio_patchify_proj",
-            "proj_out",
-            "audio_proj_out",
-            # Timestep embedding layers - int4 tinygemm requires strict bfloat16 input
-            # and these receive float32 sinusoidal embeddings that are cast to bfloat16
-            "*adaln*",
-            "time_proj",
-            "timestep_embedder*",
-            # Caption/text projection layers
-            "caption_projection*",
-            "audio_caption_projection*",
-            # Normalization layers (usually excluded from quantization)
-            "*norm*",
-        ]
-    }
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    elif isinstance(device, str):
+        device = torch.device(device)
+
+    weight_quant = _get_quanto_dtype(precision)
+
     if quantize_activations:
         logger.debug("Quantizing model weights and activations")
-        extra_quanto_args["activations"] = weight_quant
+        activations_quant = weight_quant
     else:
-        logger.debug("Quantizing model weights only")
+        activations_quant = None
+
+    # Remember original device to restore after quantization
+    original_device = next(model.parameters()).device
+
+    # Check if model has transformer_blocks for block-by-block quantization
+    if hasattr(model, "transformer_blocks"):
+        logger.debug("Quantizing model using block-by-block approach for memory efficiency")
+        _quantize_blockwise(
+            model,
+            weight_quant=weight_quant,
+            activations_quant=activations_quant,
+            device=device,
+        )
+    else:
+        # Fallback: quantize entire model at once
+        model.to(device)
+        quantize(model, weights=weight_quant, activations=activations_quant, exclude=EXCLUDE_PATTERNS)
+        freeze(model)
+
+    # Restore model to original device
+    model.to(original_device)
 
-    quantize(model, weights=weight_quant, **extra_quanto_args)
-    freeze(model)
     return model
 
 
-def _quanto_type_map(precision: QuantizationOptions) -> torch.dtype | qtype | None:
+def _quantize_blockwise(
+    model: torch.nn.Module,
+    weight_quant: torch.dtype,
+    activations_quant: torch.dtype | None,
+    device: torch.device,
+) -> None:
+    """Quantize a model block-by-block using optimum-quanto.
+    This approach:
+    1. Moves each transformer block to GPU
+    2. Quantizes on GPU (fast!)
+    3. Freezes the quantized weights
+    4. Moves back to CPU
+    This is much faster than quantizing on CPU and uses less peak VRAM
+    than loading the entire model to GPU.
+    """
+    from optimum.quanto import freeze, quantize  # noqa: PLC0415
+
+    original_dtype = next(model.parameters()).dtype
+    transformer_blocks = list(model.transformer_blocks)
+
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        transient=True,
+    ) as progress:
+        task = progress.add_task("Quantizing transformer blocks", total=len(transformer_blocks))
+
+        for block in transformer_blocks:
+            # Move block to GPU
+            block.to(device, dtype=original_dtype, non_blocking=True)
+
+            # Quantize on GPU
+            quantize(block, weights=weight_quant, activations=activations_quant, exclude=EXCLUDE_PATTERNS)
+            freeze(block)
+
+            # Move back to CPU to free up VRAM for next block
+            block.to("cpu", non_blocking=True)
+
+            progress.advance(task)
+
+    # Quantize remaining non-transformer-block modules (e.g., embeddings, timestep projections)
+    # Skip modules that should not be quantized (patchify_proj, proj_out, etc.)
+    logger.debug("Quantizing remaining model components")
+
+    for name, module in model.named_children():
+        if name == "transformer_blocks":
+            continue  # Already quantized
+
+        if name in SKIP_ROOT_MODULES:
+            logger.debug(f"Skipping quantization for module: {name}")
+            continue  # Don't quantize these modules
+
+        # Move to device, quantize, freeze, move back
+        module.to(device, dtype=original_dtype, non_blocking=True)
+        quantize(module, weights=weight_quant, activations=activations_quant, exclude=EXCLUDE_PATTERNS)
+        freeze(module)
+        module.to("cpu", non_blocking=True)
+
+
+def _get_quanto_dtype(precision: QuantizationOptions) -> torch.dtype:
+    """Map precision string to quanto dtype."""
     from optimum.quanto import (  # noqa: PLC0415
         qfloat8,
         qfloat8_e4m3fnuz,
@@ -79,14 +186,10 @@ def _quanto_type_map(precision: QuantizationOptions) -> torch.dtype | qtype | No
         return qint8
     elif precision in ("fp8-quanto", "fp8uz-quanto"):
         if torch.backends.mps.is_available():
-            logger.warning(
-                "MPS doesn't support dtype float8. "
-                "you must select another precision level such as int2, int8, or int8.",
-            )
-            return None
+            raise ValueError("FP8 quantization is not supported on MPS devices. Use int2, int4, or int8 instead.")
         if precision == "fp8-quanto":
             return qfloat8
         elif precision == "fp8uz-quanto":
             return qfloat8_e4m3fnuz
 
-    raise ValueError(f"Invalid quantisation level: {precision}")
+    raise ValueError(f"Invalid quantization precision: {precision}")

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

@@ -31,6 +31,7 @@ from ltx_trainer import logger
 from ltx_trainer.config import LtxTrainerConfig
 from ltx_trainer.config_display import print_config
 from ltx_trainer.datasets import PrecomputedDataset
+from ltx_trainer.gpu_utils import free_gpu_memory, free_gpu_memory_context, get_gpu_memory_gb
 from ltx_trainer.hf_hub_utils import push_to_hub
 from ltx_trainer.model_loader import load_model as load_ltx_model
 from ltx_trainer.model_loader import load_text_encoder
@@ -38,7 +39,7 @@ from ltx_trainer.progress import TrainingProgress
 from ltx_trainer.quantization import quantize_model
 from ltx_trainer.timestep_samplers import SAMPLERS
 from ltx_trainer.training_strategies import get_training_strategy
-from ltx_trainer.utils import get_gpu_memory_gb, open_image_as_srgb, save_image
+from ltx_trainer.utils import open_image_as_srgb, save_image
 from ltx_trainer.validation_sampler import CachedPromptEmbeddings, GenerationConfig, ValidationSampler
 from ltx_trainer.video_utils import read_video, save_video
 
@@ -330,6 +331,7 @@ class LtxvTrainer:
 
         return loss
 
+    @free_gpu_memory_context(after=True)
     def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
         """Load text encoder, computes and returns validation embeddings."""
 
@@ -342,17 +344,13 @@ class LtxvTrainer:
 
         # Load text encoder on GPU
         logger.debug("Loading text encoder...")
-        if self._config.acceleration.load_text_encoder_in_8bit:
-            logger.warning(
-                "⚠️ load_text_encoder_in_8bit is set to True but 8-bit text encoder loading "
-                "is not currently implemented. The text encoder will be loaded in bfloat16 precision."
-            )
 
         self._text_encoder = load_text_encoder(
             checkpoint_path=self._config.model.model_path,
             gemma_model_path=self._config.model.text_encoder_path,
             device="cuda",
             dtype=torch.bfloat16,
+            load_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
         )
 
         # Cache validation embeddings if prompts are configured
@@ -378,7 +376,6 @@ class LtxvTrainer:
         self._text_encoder.model = None
         self._text_encoder.tokenizer = None
         self._text_encoder.feature_extractor_linear = None
-        torch.cuda.empty_cache()
 
         logger.debug("Validation prompt embeddings cached. Gemma model unloaded")
         return cached_embeddings
@@ -724,6 +721,7 @@ class LtxvTrainer:
 
     # Note: Use @torch.no_grad() instead of @torch.inference_mode() to avoid FSDP inplace update errors after validation
     @torch.no_grad()
+    @free_gpu_memory_context(after=True)
     def _sample_videos(self, progress: TrainingProgress) -> list[Path] | None:
         """Run validation by generating videos from validation prompts."""
         use_images = self._config.validation.images is not None
@@ -731,10 +729,9 @@ class LtxvTrainer:
         generate_audio = self._config.validation.generate_audio
         inference_steps = self._config.validation.inference_steps
 
-        # Free up GPU memory before validation sampling.
-        # Zero gradients and empty the cache to reclaim memory.
+        # Zero gradients and free GPU memory to reclaim memory before validation sampling
         self._optimizer.zero_grad(set_to_none=True)
-        torch.cuda.empty_cache()
+        free_gpu_memory()
 
         # Start sampling progress tracking
         sampling_ctx = progress.start_sampling(
@@ -831,9 +828,6 @@ class LtxvTrainer:
         # Clean up progress tasks
         sampling_ctx.cleanup()
 
-        # Clear GPU cache after validation
-        torch.cuda.empty_cache()
-
         rel_outputs_path = output_dir.relative_to(self._config.output_dir)
         logger.info(f"🎥 Validation samples for step {self._global_step} saved in {rel_outputs_path}")
         return video_paths

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

@@ -1,5 +1,4 @@
 import io
-import subprocess
 from pathlib import Path
 
 import numpy as np
@@ -7,35 +6,6 @@ import torch
 from PIL import ExifTags, Image, ImageCms, ImageOps
 from PIL.Image import Image as PilImage
 
-from ltx_trainer import logger
-
-
-def get_gpu_memory_gb(device: torch.device) -> float:
-    """
-    Get current GPU memory usage in GB using nvidia-smi
-    Args:
-        device: torch.device to get memory usage for
-    Returns:
-        Current GPU memory usage in GB
-    """
-    try:
-        device_id = device.index if device.index is not None else 0
-        result = subprocess.check_output(
-            [
-                "nvidia-smi",
-                "--query-gpu=memory.used",
-                "--format=csv,nounits,noheader",
-                "-i",
-                str(device_id),
-            ],
-            encoding="utf-8",
-        )
-        return float(result.strip()) / 1024  # Convert MB to GB
-    except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e:
-        logger.error(f"Failed to get GPU memory from nvidia-smi: {e}")
-        # Fallback to torch
-        return torch.cuda.memory_allocated(device) / 1024**3
-
 
 def open_image_as_srgb(image_path: str | Path | io.BytesIO) -> PilImage:
     """