Automated PR - 2026-01-15

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:
     """