Automated PR - 2026-01-15

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,
     )