ComputeNerd/ltx-2

1# ruff: noqa: PLC04152 3"""4Model loader for LTX-2 trainer using the new ltx-core package.5 6This module provides a unified interface for loading LTX-2 model components7for training, using SingleGPUModelBuilder from ltx-core.8 9Example usage:10    # Load individual components11    vae_encoder = load_video_vae_encoder("/path/to/checkpoint.safetensors", device="cuda")12    vae_decoder = load_video_vae_decoder("/path/to/checkpoint.safetensors", device="cuda")13    text_encoder = load_text_encoder("/path/to/checkpoint.safetensors", "/path/to/gemma", device="cuda")14 15    # Load all components at once16    components = load_model("/path/to/checkpoint.safetensors", text_encoder_path="/path/to/gemma")17"""18 19from __future__ import annotations20 21from dataclasses import dataclass22from pathlib import Path23from typing import TYPE_CHECKING24 25import torch26 27from ltx_trainer import logger28 29# Type alias for device specification30Device = str | torch.device31 32# Type checking imports (not loaded at runtime)33if TYPE_CHECKING:34    from ltx_core.model.audio_vae.audio_vae import Decoder as AudioVAEDecoder35    from ltx_core.model.audio_vae.audio_vae import Encoder as AudioVAEEncoder36    from ltx_core.model.audio_vae.vocoder import Vocoder37    from ltx_core.model.clip.gemma.encoders.av_encoder import AVGemmaTextEncoderModel38    from ltx_core.model.transformer.model import LTXModel39    from ltx_core.model.video_vae.video_vae import Decoder as VideoVAEDecoder40    from ltx_core.model.video_vae.video_vae import Encoder as VideoVAEEncoder41    from ltx_core.pipeline.components.schedulers import LTX2Scheduler42 43 44def _to_torch_device(device: Device) -> torch.device:45    """Convert device specification to torch.device."""46    return torch.device(device) if isinstance(device, str) else device47 48 49# =============================================================================50# Individual Component Loaders51# =============================================================================52 53 54def load_transformer(55    checkpoint_path: str | Path,56    device: Device = "cpu",57    dtype: torch.dtype = torch.bfloat16,58) -> "LTXModel":59    """Load the LTX transformer model.60 61    Args:62        checkpoint_path: Path to the safetensors checkpoint file63        device: Device to load model on64        dtype: Data type for model weights65 66    Returns:67        Loaded LTXModel transformer68    """69    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder70    from ltx_core.model.transformer.model_configurator import (71        LTXV_MODEL_COMFY_RENAMING_MAP,72        LTXModelConfigurator,73    )74 75    return SingleGPUModelBuilder(76        model_path=str(checkpoint_path),77        model_class_configurator=LTXModelConfigurator,78        model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,79    ).build(device=_to_torch_device(device), dtype=dtype)80 81 82def load_video_vae_encoder(83    checkpoint_path: str | Path,84    device: Device = "cpu",85    dtype: torch.dtype = torch.bfloat16,86) -> "VideoVAEEncoder":87    """Load the video VAE encoder (for preprocessing).88 89    Args:90        checkpoint_path: Path to the safetensors checkpoint file91        device: Device to load model on92        dtype: Data type for model weights93 94    Returns:95        Loaded VideoVAEEncoder96    """97    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder98    from ltx_core.model.video_vae.model_configurator import VAE_ENCODER_COMFY_KEYS_FILTER99    from ltx_core.model.video_vae.model_configurator import (100        VAEEncoderConfigurator as VideoVAEEncoderConfigurator,101    )102 103    return SingleGPUModelBuilder(104        model_path=str(checkpoint_path),105        model_class_configurator=VideoVAEEncoderConfigurator,106        model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,107    ).build(device=_to_torch_device(device), dtype=dtype)108 109 110def load_video_vae_decoder(111    checkpoint_path: str | Path,112    device: Device = "cpu",113    dtype: torch.dtype = torch.bfloat16,114) -> "VideoVAEDecoder":115    """Load the video VAE decoder (for inference/validation).116 117    Args:118        checkpoint_path: Path to the safetensors checkpoint file119        device: Device to load model on120        dtype: Data type for model weights121 122    Returns:123        Loaded VideoVAEDecoder124    """125    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder126    from ltx_core.model.video_vae.model_configurator import VAE_DECODER_COMFY_KEYS_FILTER127    from ltx_core.model.video_vae.model_configurator import (128        VAEDecoderConfigurator as VideoVAEDecoderConfigurator,129    )130 131    return SingleGPUModelBuilder(132        model_path=str(checkpoint_path),133        model_class_configurator=VideoVAEDecoderConfigurator,134        model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,135    ).build(device=_to_torch_device(device), dtype=dtype)136 137 138def load_audio_vae_encoder(139    checkpoint_path: str | Path,140    device: Device = "cpu",141    dtype: torch.dtype = torch.bfloat16,142) -> "AudioVAEEncoder":143    """Load the audio VAE encoder (for preprocessing).144 145    Args:146        checkpoint_path: Path to the safetensors checkpoint file147        device: Device to load model on148        dtype: Data type for model weights (default bfloat16, but float32 recommended for quality)149 150    Returns:151        Loaded AudioVAEEncoder152    """153    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder154    from ltx_core.model.audio_vae.model_configurator import AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER155    from ltx_core.model.audio_vae.model_configurator import (156        VAEEncoderConfigurator as AudioVAEEncoderConfigurator,157    )158 159    return SingleGPUModelBuilder(160        model_path=str(checkpoint_path),161        model_class_configurator=AudioVAEEncoderConfigurator,162        model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,163    ).build(device=_to_torch_device(device), dtype=dtype)164 165 166def load_audio_vae_decoder(167    checkpoint_path: str | Path,168    device: Device = "cpu",169    dtype: torch.dtype = torch.bfloat16,170) -> "AudioVAEDecoder":171    """Load the audio VAE decoder.172 173    Args:174        checkpoint_path: Path to the safetensors checkpoint file175        device: Device to load model on176        dtype: Data type for model weights177 178    Returns:179        Loaded AudioVAEDecoder180    """181    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder182    from ltx_core.model.audio_vae.model_configurator import AUDIO_VAE_DECODER_COMFY_KEYS_FILTER183    from ltx_core.model.audio_vae.model_configurator import (184        VAEDecoderConfigurator as AudioVAEDecoderConfigurator,185    )186 187    return SingleGPUModelBuilder(188        model_path=str(checkpoint_path),189        model_class_configurator=AudioVAEDecoderConfigurator,190        model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,191    ).build(device=_to_torch_device(device), dtype=dtype)192 193 194def load_vocoder(195    checkpoint_path: str | Path,196    device: Device = "cpu",197    dtype: torch.dtype = torch.bfloat16,198) -> "Vocoder":199    """Load the vocoder (for audio waveform generation).200 201    Args:202        checkpoint_path: Path to the safetensors checkpoint file203        device: Device to load model on204        dtype: Data type for model weights205 206    Returns:207        Loaded Vocoder208    """209    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder210    from ltx_core.model.audio_vae.model_configurator import VOCODER_COMFY_KEYS_FILTER, VocoderConfigurator211 212    return SingleGPUModelBuilder(213        model_path=str(checkpoint_path),214        model_class_configurator=VocoderConfigurator,215        model_sd_ops=VOCODER_COMFY_KEYS_FILTER,216    ).build(device=_to_torch_device(device), dtype=dtype)217 218 219def load_text_encoder(220    checkpoint_path: str | Path,221    gemma_model_path: str | Path,222    device: Device = "cpu",223    dtype: torch.dtype = torch.bfloat16,224) -> "AVGemmaTextEncoderModel":225    """Load the Gemma text encoder.226 227    Args:228        checkpoint_path: Path to the LTX-2 safetensors checkpoint file229        gemma_model_path: Path to Gemma model directory230        device: Device to load model on231        dtype: Data type for model weights232 233    Returns:234        Loaded AVGemmaTextEncoderModel235    """236    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder237    from ltx_core.model.clip.gemma.encoders.av_encoder import (238        AV_GEMMA_TEXT_ENCODER_KEY_OPS,239        AVGemmaTextEncoderModelConfigurator,240    )241    from ltx_core.model.clip.gemma.encoders.base_encoder import module_ops_from_gemma_root242 243    if not Path(gemma_model_path).is_dir():244        raise ValueError(f"Gemma model path is not a directory: {gemma_model_path}")245 246    torch_device = _to_torch_device(device)247    text_encoder = SingleGPUModelBuilder(248        model_path=str(checkpoint_path),249        model_class_configurator=AVGemmaTextEncoderModelConfigurator,250        model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,251        module_ops=module_ops_from_gemma_root(str(gemma_model_path)),252    ).build(device=torch_device, dtype=dtype)253 254    return text_encoder255 256 257# =============================================================================258# Combined Component Loader259# =============================================================================260 261 262@dataclass263class LtxModelComponents:264    """Container for all LTX-2 model components."""265 266    transformer: "LTXModel"267    video_vae_encoder: "VideoVAEEncoder | None" = None268    video_vae_decoder: "VideoVAEDecoder | None" = None269    audio_vae_decoder: "AudioVAEDecoder | None" = None270    vocoder: "Vocoder | None" = None271    text_encoder: "AVGemmaTextEncoderModel | None" = None272    scheduler: "LTX2Scheduler | None" = None273 274 275def load_model(276    checkpoint_path: str | Path,277    text_encoder_path: str | Path | None = None,278    device: Device = "cpu",279    dtype: torch.dtype = torch.bfloat16,280    with_video_vae_encoder: bool = False,281    with_video_vae_decoder: bool = True,282    with_audio_vae_decoder: bool = True,283    with_vocoder: bool = True,284    with_text_encoder: bool = True,285) -> LtxModelComponents:286    """287    Load LTX-2 model components from a safetensors checkpoint.288 289    This is a convenience function that loads multiple components at once.290    For loading individual components, use the dedicated functions:291    - load_transformer()292    - load_video_vae_encoder()293    - load_video_vae_decoder()294    - load_audio_vae_decoder()295    - load_vocoder()296    - load_text_encoder()297 298    Args:299        checkpoint_path: Path to the safetensors checkpoint file300        text_encoder_path: Path to Gemma model directory (required if with_text_encoder=True)301        device: Device to load models on ("cuda", "cpu", etc.)302        dtype: Data type for model weights303        with_video_vae_encoder: Whether to load the video VAE encoder (for preprocessing)304        with_video_vae_decoder: Whether to load the video VAE decoder (for inference/validation)305        with_audio_vae_decoder: Whether to load the audio VAE decoder306        with_vocoder: Whether to load the vocoder307        with_text_encoder: Whether to load the text encoder308 309    Returns:310        LtxModelComponents containing all loaded model components311    """312    from ltx_core.pipeline.components.schedulers import LTX2Scheduler313 314    checkpoint_path = Path(checkpoint_path)315 316    # Validate checkpoint exists317    if not checkpoint_path.exists():318        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")319 320    logger.info(f"Loading LTX-2 model from {checkpoint_path}")321 322    torch_device = _to_torch_device(device)323 324    # Load transformer325    logger.debug("Loading transformer...")326    transformer = load_transformer(checkpoint_path, torch_device, dtype)327 328    # Load video VAE encoder329    video_vae_encoder = None330    if with_video_vae_encoder:331        logger.debug("Loading video VAE encoder...")332        video_vae_encoder = load_video_vae_encoder(checkpoint_path, torch_device, dtype)333 334    # Load video VAE decoder335    video_vae_decoder = None336    if with_video_vae_decoder:337        logger.debug("Loading video VAE decoder...")338        video_vae_decoder = load_video_vae_decoder(checkpoint_path, torch_device, dtype)339 340    # Load audio VAE decoder341    audio_vae_decoder = None342    if with_audio_vae_decoder:343        logger.debug("Loading audio VAE decoder...")344        audio_vae_decoder = load_audio_vae_decoder(checkpoint_path, torch_device, dtype)345 346    # Load vocoder347    vocoder = None348    if with_vocoder:349        logger.debug("Loading vocoder...")350        vocoder = load_vocoder(checkpoint_path, torch_device, dtype)351 352    # Load text encoder353    text_encoder = None354    if with_text_encoder:355        if text_encoder_path is None:356            raise ValueError("text_encoder_path must be provided when with_text_encoder=True")357        logger.debug("Loading Gemma text encoder...")358        text_encoder = load_text_encoder(checkpoint_path, text_encoder_path, torch_device, dtype)359 360    # Create scheduler (stateless, no loading needed)361    scheduler = LTX2Scheduler()362 363    return LtxModelComponents(364        transformer=transformer,365        video_vae_encoder=video_vae_encoder,366        video_vae_decoder=video_vae_decoder,367        audio_vae_decoder=audio_vae_decoder,368        vocoder=vocoder,369        text_encoder=text_encoder,370        scheduler=scheduler,371    )372