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