From 310103a53e16e4841a9068ef2530d57364f88a6e Mon Sep 17 00:00:00 2001 From: sync-bot Date: Thu, 15 Jan 2026 19:19:42 +0000 Subject: [PATCH] Automated PR - 2026-01-15 --- .../configs/ltx2_av_lora_low_vram.yaml | 2 +- .../docs/configuration-reference.md | 11 +- .../ltx-trainer/scripts/decode_latents.py | 158 +++++++++------ .../ltx-trainer/scripts/process_captions.py | 19 +- .../ltx-trainer/scripts/process_dataset.py | 85 ++++---- .../ltx-trainer/scripts/process_videos.py | 185 +++++++++++++++++- .../ltx-trainer/src/ltx_trainer/captioning.py | 2 +- .../ltx-trainer/src/ltx_trainer/gemma_8bit.py | 135 +++++++++++++ .../ltx-trainer/src/ltx_trainer/gpu_utils.py | 90 +++++++++ .../src/ltx_trainer/model_loader.py | 17 +- .../src/ltx_trainer/quantization.py | 175 +++++++++++++---- .../ltx-trainer/src/ltx_trainer/trainer.py | 20 +- packages/ltx-trainer/src/ltx_trainer/utils.py | 30 --- 13 files changed, 722 insertions(+), 207 deletions(-) create mode 100644 packages/ltx-trainer/src/ltx_trainer/gemma_8bit.py create mode 100644 packages/ltx-trainer/src/ltx_trainer/gpu_utils.py diff --git a/packages/ltx-trainer/configs/ltx2_av_lora_low_vram.yaml b/packages/ltx-trainer/configs/ltx2_av_lora_low_vram.yaml index fcaf0ba..811edf3 100644 --- a/packages/ltx-trainer/configs/ltx2_av_lora_low_vram.yaml +++ b/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 diff --git a/packages/ltx-trainer/docs/configuration-reference.md b/packages/ltx-trainer/docs/configuration-reference.md index 12e9449..df5c023 100644 --- a/packages/ltx-trainer/docs/configuration-reference.md +++ b/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 diff --git a/packages/ltx-trainer/scripts/decode_latents.py b/packages/ltx-trainer/scripts/decode_latents.py index a69f72a..f489a91 100755 --- a/packages/ltx-trainer/scripts/decode_latents.py +++ b/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) diff --git a/packages/ltx-trainer/scripts/process_captions.py b/packages/ltx-trainer/scripts/process_captions.py index f511690..49cab23 100755 --- a/packages/ltx-trainer/scripts/process_captions.py +++ b/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, ) diff --git a/packages/ltx-trainer/scripts/process_dataset.py b/packages/ltx-trainer/scripts/process_dataset.py index 29289ed..f13fc0c 100755 --- a/packages/ltx-trainer/scripts/process_dataset.py +++ b/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, ) diff --git a/packages/ltx-trainer/scripts/process_videos.py b/packages/ltx-trainer/scripts/process_videos.py index 3e6bc4f..ca51774 100755 --- a/packages/ltx-trainer/scripts/process_videos.py +++ b/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, diff --git a/packages/ltx-trainer/src/ltx_trainer/captioning.py b/packages/ltx-trainer/src/ltx_trainer/captioning.py index 303d5bb..5e09c30 100644 --- a/packages/ltx-trainer/src/ltx_trainer/captioning.py +++ b/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 """ diff --git a/packages/ltx-trainer/src/ltx_trainer/gemma_8bit.py b/packages/ltx-trainer/src/ltx_trainer/gemma_8bit.py new file mode 100644 index 0000000..5e55ab7 --- /dev/null +++ b/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) diff --git a/packages/ltx-trainer/src/ltx_trainer/gpu_utils.py b/packages/ltx-trainer/src/ltx_trainer/gpu_utils.py new file mode 100644 index 0000000..de385a1 --- /dev/null +++ b/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 diff --git a/packages/ltx-trainer/src/ltx_trainer/model_loader.py b/packages/ltx-trainer/src/ltx_trainer/model_loader.py index c7fd248..3220463 100644 --- a/packages/ltx-trainer/src/ltx_trainer/model_loader.py +++ b/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), diff --git a/packages/ltx-trainer/src/ltx_trainer/quantization.py b/packages/ltx-trainer/src/ltx_trainer/quantization.py index ff319f2..31bc9bc 100644 --- a/packages/ltx-trainer/src/ltx_trainer/quantization.py +++ b/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}") diff --git a/packages/ltx-trainer/src/ltx_trainer/trainer.py b/packages/ltx-trainer/src/ltx_trainer/trainer.py index d7a7a59..b9f5f80 100644 --- a/packages/ltx-trainer/src/ltx_trainer/trainer.py +++ b/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 diff --git a/packages/ltx-trainer/src/ltx_trainer/utils.py b/packages/ltx-trainer/src/ltx_trainer/utils.py index 4530e9b..85bf1dd 100644 --- a/packages/ltx-trainer/src/ltx_trainer/utils.py +++ b/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: """