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