| 1 | from dataclasses import replace |
| 2 | |
| 3 | import torch |
| 4 | |
| 5 | from ltx_core.loader import SDOps |
| 6 | from ltx_core.loader.primitives import LoraPathStrengthAndSDOps |
| 7 | from ltx_core.loader.registry import DummyRegistry, Registry |
| 8 | from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder |
| 9 | from ltx_core.model.audio_vae import ( |
| 10 | AUDIO_VAE_DECODER_COMFY_KEYS_FILTER, |
| 11 | AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER, |
| 12 | VOCODER_COMFY_KEYS_FILTER, |
| 13 | AudioDecoder, |
| 14 | AudioDecoderConfigurator, |
| 15 | AudioEncoder, |
| 16 | AudioEncoderConfigurator, |
| 17 | Vocoder, |
| 18 | VocoderConfigurator, |
| 19 | ) |
| 20 | from ltx_core.model.transformer import ( |
| 21 | LTXV_MODEL_COMFY_RENAMING_MAP, |
| 22 | LTXModelConfigurator, |
| 23 | X0Model, |
| 24 | ) |
| 25 | from ltx_core.model.upsampler import LatentUpsampler, LatentUpsamplerConfigurator |
| 26 | from ltx_core.model.video_vae import ( |
| 27 | VAE_DECODER_COMFY_KEYS_FILTER, |
| 28 | VAE_ENCODER_COMFY_KEYS_FILTER, |
| 29 | VideoDecoder, |
| 30 | VideoDecoderConfigurator, |
| 31 | VideoEncoder, |
| 32 | VideoEncoderConfigurator, |
| 33 | ) |
| 34 | from ltx_core.quantization import QuantizationPolicy |
| 35 | from ltx_core.text_encoders.gemma import ( |
| 36 | EMBEDDINGS_PROCESSOR_KEY_OPS, |
| 37 | GEMMA_LANGUAGE_ONLY_KEY_OPS, |
| 38 | GEMMA_LANGUAGE_ONLY_MODEL_OPS, |
| 39 | GEMMA_LLM_KEY_OPS, |
| 40 | GEMMA_MODEL_OPS, |
| 41 | EmbeddingsProcessor, |
| 42 | EmbeddingsProcessorConfigurator, |
| 43 | GemmaTextEncoder, |
| 44 | GemmaTextEncoderConfigurator, |
| 45 | module_ops_from_gemma_root, |
| 46 | tokenizer_module_ops_from_gemma_root, |
| 47 | ) |
| 48 | from ltx_core.utils import find_matching_file |
| 49 | |
| 50 | |
| 51 | class ModelLedger: |
| 52 | """ |
| 53 | Central coordinator for loading and building models used in an LTX pipeline. |
| 54 | The ledger wires together multiple model builders (transformer, video VAE encoder/decoder, |
| 55 | audio VAE decoder, vocoder, text encoder, and optional latent upsampler) and exposes |
| 56 | factory methods for constructing model instances. |
| 57 | ### Model Building |
| 58 | Each model method (e.g. :meth:`transformer`, :meth:`video_decoder`, :meth:`text_encoder`) |
| 59 | constructs a new model instance on each call. The builder uses the |
| 60 | :class:`~ltx_core.loader.registry.Registry` to load weights from the checkpoint, |
| 61 | instantiates the model with the configured ``dtype``, and moves it to ``self.device``. |
| 62 | .. note:: |
| 63 | Models are **not cached**. Each call to a model method creates a new instance. |
| 64 | Callers are responsible for storing references to models they wish to reuse |
| 65 | and for freeing GPU memory (e.g. by deleting references and calling |
| 66 | ``torch.cuda.empty_cache()``). |
| 67 | ### Constructor parameters |
| 68 | dtype: |
| 69 | Torch dtype used when constructing all models (e.g. ``torch.bfloat16``). |
| 70 | device: |
| 71 | Target device to which models are moved after construction (e.g. ``torch.device("cuda")``). |
| 72 | checkpoint_path: |
| 73 | Path to a checkpoint directory or file containing the core model weights |
| 74 | (transformer, video VAE, audio VAE, text encoder, vocoder). If ``None``, the |
| 75 | corresponding builders are not created and calling those methods will raise |
| 76 | a :class:`ValueError`. |
| 77 | gemma_root_path: |
| 78 | Base path to Gemma-compatible CLIP/text encoder weights. Required to |
| 79 | initialize the text encoder builder; if omitted, :meth:`text_encoder` cannot be used. |
| 80 | spatial_upsampler_path: |
| 81 | Optional path to a latent upsampler checkpoint. If provided, the |
| 82 | :meth:`spatial_upsampler` method becomes available; otherwise calling it raises |
| 83 | a :class:`ValueError`. |
| 84 | loras: |
| 85 | Tuple of LoRA configurations (path, strength, sd_ops) applied on top of the base |
| 86 | transformer weights. Use ``()`` for none. |
| 87 | registry: |
| 88 | Optional :class:`Registry` instance for weight caching across builders. |
| 89 | Defaults to :class:`DummyRegistry` which performs no cross-builder caching. |
| 90 | quantization: |
| 91 | Optional :class:`QuantizationPolicy` controlling how transformer weights |
| 92 | are stored and how matmul is executed. Defaults to None, which means no quantization. |
| 93 | ### Creating Variants |
| 94 | Use :meth:`with_additional_loras` to create a new ``ModelLedger`` instance that |
| 95 | includes additional LoRA configurations or :meth:`with_loras` to replace existing |
| 96 | lora configurations while sharing the same registry for weight caching. |
| 97 | """ |
| 98 | |
| 99 | def __init__( |
| 100 | self, |
| 101 | dtype: torch.dtype, |
| 102 | device: torch.device, |
| 103 | checkpoint_path: str | None = None, |
| 104 | gemma_root_path: str | None = None, |
| 105 | spatial_upsampler_path: str | None = None, |
| 106 | loras: tuple[LoraPathStrengthAndSDOps, ...] = (), |
| 107 | registry: Registry | None = None, |
| 108 | quantization: QuantizationPolicy | None = None, |
| 109 | ): |
| 110 | self.dtype = dtype |
| 111 | self.device = device |
| 112 | self.checkpoint_path = checkpoint_path |
| 113 | self.gemma_root_path = gemma_root_path |
| 114 | self.spatial_upsampler_path = spatial_upsampler_path |
| 115 | self.loras = loras |
| 116 | self.registry = registry or DummyRegistry() |
| 117 | self.quantization = quantization |
| 118 | self.build_model_builders() |
| 119 | |
| 120 | def build_model_builders(self) -> None: |
| 121 | if self.checkpoint_path is not None: |
| 122 | self.transformer_builder = Builder( |
| 123 | model_path=self.checkpoint_path, |
| 124 | model_class_configurator=LTXModelConfigurator, |
| 125 | model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP, |
| 126 | loras=tuple(self.loras), |
| 127 | registry=self.registry, |
| 128 | ) |
| 129 | |
| 130 | self.vae_decoder_builder = Builder( |
| 131 | model_path=self.checkpoint_path, |
| 132 | model_class_configurator=VideoDecoderConfigurator, |
| 133 | model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER, |
| 134 | registry=self.registry, |
| 135 | ) |
| 136 | |
| 137 | self.vae_encoder_builder = Builder( |
| 138 | model_path=self.checkpoint_path, |
| 139 | model_class_configurator=VideoEncoderConfigurator, |
| 140 | model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER, |
| 141 | registry=self.registry, |
| 142 | ) |
| 143 | |
| 144 | self.audio_encoder_builder = Builder[AudioEncoder]( |
| 145 | model_path=self.checkpoint_path, |
| 146 | model_class_configurator=AudioEncoderConfigurator, |
| 147 | model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER, |
| 148 | registry=self.registry, |
| 149 | ) |
| 150 | |
| 151 | self.audio_decoder_builder = Builder( |
| 152 | model_path=self.checkpoint_path, |
| 153 | model_class_configurator=AudioDecoderConfigurator, |
| 154 | model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER, |
| 155 | registry=self.registry, |
| 156 | ) |
| 157 | |
| 158 | self.vocoder_builder = Builder( |
| 159 | model_path=self.checkpoint_path, |
| 160 | model_class_configurator=VocoderConfigurator, |
| 161 | model_sd_ops=VOCODER_COMFY_KEYS_FILTER, |
| 162 | registry=self.registry, |
| 163 | ) |
| 164 | |
| 165 | # Embeddings processor only needs the LTX checkpoint (no Gemma weights) |
| 166 | self.embeddings_processor_builder = Builder( |
| 167 | model_path=self.checkpoint_path, |
| 168 | model_class_configurator=EmbeddingsProcessorConfigurator, |
| 169 | model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS, |
| 170 | registry=self.registry, |
| 171 | ) |
| 172 | |
| 173 | if self.gemma_root_path is not None: |
| 174 | module_ops = module_ops_from_gemma_root(self.gemma_root_path) |
| 175 | model_folder = find_matching_file(self.gemma_root_path, "model*.safetensors").parent |
| 176 | weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")] |
| 177 | |
| 178 | self.text_encoder_builder = Builder( |
| 179 | model_path=tuple(weight_paths), |
| 180 | model_class_configurator=GemmaTextEncoderConfigurator, |
| 181 | model_sd_ops=GEMMA_LLM_KEY_OPS, |
| 182 | registry=self.registry, |
| 183 | module_ops=(GEMMA_MODEL_OPS, *module_ops), |
| 184 | ) |
| 185 | |
| 186 | self.language_only_text_encoder_builder = Builder( |
| 187 | model_path=tuple(weight_paths), |
| 188 | model_class_configurator=GemmaTextEncoderConfigurator, |
| 189 | model_sd_ops=GEMMA_LANGUAGE_ONLY_KEY_OPS, |
| 190 | registry=self.registry, |
| 191 | module_ops=( |
| 192 | GEMMA_LANGUAGE_ONLY_MODEL_OPS, |
| 193 | *tokenizer_module_ops_from_gemma_root(self.gemma_root_path), |
| 194 | ), |
| 195 | ) |
| 196 | |
| 197 | if self.spatial_upsampler_path is not None: |
| 198 | self.upsampler_builder = Builder( |
| 199 | model_path=self.spatial_upsampler_path, |
| 200 | model_class_configurator=LatentUpsamplerConfigurator, |
| 201 | registry=self.registry, |
| 202 | ) |
| 203 | |
| 204 | def _target_device(self) -> torch.device: |
| 205 | if isinstance(self.registry, DummyRegistry) or self.registry is None: |
| 206 | return self.device |
| 207 | else: |
| 208 | return torch.device("cpu") |
| 209 | |
| 210 | def with_additional_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger": |
| 211 | """Add new lora configurations to the existing ones.""" |
| 212 | return self.with_loras((*self.loras, *loras)) |
| 213 | |
| 214 | def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger": |
| 215 | """Replace existing lora configurations with new ones.""" |
| 216 | return ModelLedger( |
| 217 | dtype=self.dtype, |
| 218 | device=self.device, |
| 219 | checkpoint_path=self.checkpoint_path, |
| 220 | gemma_root_path=self.gemma_root_path, |
| 221 | spatial_upsampler_path=self.spatial_upsampler_path, |
| 222 | loras=loras, |
| 223 | registry=self.registry, |
| 224 | quantization=self.quantization, |
| 225 | ) |
| 226 | |
| 227 | def transformer(self) -> X0Model: |
| 228 | if not hasattr(self, "transformer_builder"): |
| 229 | raise ValueError( |
| 230 | "Transformer not initialized. Please provide a checkpoint path to the ModelLedger constructor." |
| 231 | ) |
| 232 | |
| 233 | if self.quantization is None: |
| 234 | return ( |
| 235 | X0Model(self.transformer_builder.build(device=self._target_device(), dtype=self.dtype)) |
| 236 | .to(self.device) |
| 237 | .eval() |
| 238 | ) |
| 239 | else: |
| 240 | sd_ops = self.transformer_builder.model_sd_ops |
| 241 | if self.quantization.sd_ops is not None: |
| 242 | sd_ops = SDOps( |
| 243 | name=f"sd_ops_chain_{sd_ops.name}+{self.quantization.sd_ops.name}", |
| 244 | mapping=(*sd_ops.mapping, *self.quantization.sd_ops.mapping), |
| 245 | ) |
| 246 | builder = replace( |
| 247 | self.transformer_builder, |
| 248 | module_ops=(*self.transformer_builder.module_ops, *self.quantization.module_ops), |
| 249 | model_sd_ops=sd_ops, |
| 250 | ) |
| 251 | return X0Model(builder.build(device=self._target_device())).to(self.device).eval() |
| 252 | |
| 253 | def video_decoder(self) -> VideoDecoder: |
| 254 | if not hasattr(self, "vae_decoder_builder"): |
| 255 | raise ValueError( |
| 256 | "Video decoder not initialized. Please provide a checkpoint path to the ModelLedger constructor." |
| 257 | ) |
| 258 | |
| 259 | return self.vae_decoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval() |
| 260 | |
| 261 | def video_encoder(self) -> VideoEncoder: |
| 262 | if not hasattr(self, "vae_encoder_builder"): |
| 263 | raise ValueError( |
| 264 | "Video encoder not initialized. Please provide a checkpoint path to the ModelLedger constructor." |
| 265 | ) |
| 266 | |
| 267 | return self.vae_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval() |
| 268 | |
| 269 | def text_encoder(self) -> GemmaTextEncoder: |
| 270 | if not hasattr(self, "text_encoder_builder"): |
| 271 | raise ValueError( |
| 272 | "Text encoder not initialized. Please provide a checkpoint path and gemma root path to the " |
| 273 | "ModelLedger constructor." |
| 274 | ) |
| 275 | |
| 276 | return self.text_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval() |
| 277 | |
| 278 | def language_only_text_encoder(self) -> GemmaTextEncoder: |
| 279 | if not hasattr(self, "language_only_text_encoder_builder"): |
| 280 | raise ValueError( |
| 281 | "Language-only text encoder not initialized. Please provide a checkpoint path and gemma root path " |
| 282 | "to the ModelLedger constructor." |
| 283 | ) |
| 284 | |
| 285 | return ( |
| 286 | self.language_only_text_encoder_builder.build(device=self._target_device(), dtype=self.dtype) |
| 287 | .to(self.device) |
| 288 | .eval() |
| 289 | ) |
| 290 | |
| 291 | def gemma_embeddings_processor(self) -> EmbeddingsProcessor: |
| 292 | if not hasattr(self, "embeddings_processor_builder"): |
| 293 | raise ValueError( |
| 294 | "Embeddings processor not initialized. Please provide a checkpoint path to the ModelLedger constructor." |
| 295 | ) |
| 296 | |
| 297 | return ( |
| 298 | self.embeddings_processor_builder.build(device=self._target_device(), dtype=self.dtype) |
| 299 | .to(self.device) |
| 300 | .eval() |
| 301 | ) |
| 302 | |
| 303 | def audio_encoder(self) -> AudioEncoder: |
| 304 | if not hasattr(self, "audio_encoder_builder"): |
| 305 | raise ValueError( |
| 306 | "Audio encoder not initialized. Please provide a checkpoint path to the ModelLedger constructor." |
| 307 | ) |
| 308 | |
| 309 | return self.audio_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval() |
| 310 | |
| 311 | def audio_decoder(self) -> AudioDecoder: |
| 312 | if not hasattr(self, "audio_decoder_builder"): |
| 313 | raise ValueError( |
| 314 | "Audio decoder not initialized. Please provide a checkpoint path to the ModelLedger constructor." |
| 315 | ) |
| 316 | |
| 317 | return self.audio_decoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval() |
| 318 | |
| 319 | def vocoder(self) -> Vocoder: |
| 320 | if not hasattr(self, "vocoder_builder"): |
| 321 | raise ValueError( |
| 322 | "Vocoder not initialized. Please provide a checkpoint path to the ModelLedger constructor." |
| 323 | ) |
| 324 | |
| 325 | return self.vocoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval() |
| 326 | |
| 327 | def spatial_upsampler(self) -> LatentUpsampler: |
| 328 | if not hasattr(self, "upsampler_builder"): |
| 329 | raise ValueError("Upsampler not initialized. Please provide upsampler path to the ModelLedger constructor.") |
| 330 | |
| 331 | return self.upsampler_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval() |
| 332 |