| 1 | import logging |
| 2 | from dataclasses import dataclass, field, replace |
| 3 | from typing import Generic |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from ltx_core.loader.fuse_loras import apply_loras |
| 8 | from ltx_core.loader.module_ops import ModuleOps |
| 9 | from ltx_core.loader.primitives import ( |
| 10 | LoRAAdaptableProtocol, |
| 11 | LoraPathStrengthAndSDOps, |
| 12 | LoraStateDictWithStrength, |
| 13 | ModelBuilderProtocol, |
| 14 | StateDict, |
| 15 | StateDictLoader, |
| 16 | ) |
| 17 | from ltx_core.loader.registry import DummyRegistry, Registry |
| 18 | from ltx_core.loader.sd_ops import SDOps |
| 19 | from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader |
| 20 | from ltx_core.model.model_protocol import ModelConfigurator, ModelType |
| 21 | |
| 22 | logger: logging.Logger = logging.getLogger(__name__) |
| 23 | |
| 24 | |
| 25 | @dataclass(frozen=True) |
| 26 | class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol): |
| 27 | """ |
| 28 | Builder for PyTorch models residing on a single GPU. |
| 29 | Attributes: |
| 30 | model_class_configurator: Class responsible for constructing the model from a config dict. |
| 31 | model_path: Path (or tuple of shard paths) to the model's `.safetensors` checkpoint(s). |
| 32 | model_sd_ops: Optional state-dict operations applied when loading the model weights. |
| 33 | module_ops: Sequence of module-level mutations applied to the meta model before weight loading. |
| 34 | loras: Sequence of LoRA adapters (path, strength, optional sd_ops) to fuse into the model. |
| 35 | model_loader: Strategy for loading state dicts from disk. Defaults to |
| 36 | :class:`SafetensorsModelStateDictLoader`. |
| 37 | registry: Cache for already-loaded state dicts. Defaults to :class:`DummyRegistry` (no caching). |
| 38 | lora_load_device: Device used when loading LoRA weight tensors from disk. Defaults to |
| 39 | ``torch.device("cpu")``, which keeps LoRA weights in CPU memory and transfers them to |
| 40 | the target GPU sequentially during fusion, reducing peak GPU memory usage compared to |
| 41 | loading all LoRA weights directly onto the GPU at once. |
| 42 | """ |
| 43 | |
| 44 | model_class_configurator: type[ModelConfigurator[ModelType]] |
| 45 | model_path: str | tuple[str, ...] |
| 46 | model_sd_ops: SDOps | None = None |
| 47 | module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple) |
| 48 | loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple) |
| 49 | model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader) |
| 50 | registry: Registry = field(default_factory=DummyRegistry) |
| 51 | lora_load_device: torch.device = field(default_factory=lambda: torch.device("cpu")) |
| 52 | |
| 53 | def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder": |
| 54 | return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops))) |
| 55 | |
| 56 | def model_config(self) -> dict: |
| 57 | first_shard_path = self.model_path[0] if isinstance(self.model_path, tuple) else self.model_path |
| 58 | return self.model_loader.metadata(first_shard_path) |
| 59 | |
| 60 | def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType: |
| 61 | with torch.device("meta"): |
| 62 | model = self.model_class_configurator.from_config(config) |
| 63 | for module_op in module_ops: |
| 64 | if module_op.matcher(model): |
| 65 | model = module_op.mutator(model) |
| 66 | return model |
| 67 | |
| 68 | def load_sd( |
| 69 | self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None |
| 70 | ) -> StateDict: |
| 71 | state_dict = registry.get(paths, sd_ops) |
| 72 | if state_dict is None: |
| 73 | state_dict = self.model_loader.load(paths, sd_ops=sd_ops, device=device) |
| 74 | registry.add(paths, sd_ops=sd_ops, state_dict=state_dict) |
| 75 | return state_dict |
| 76 | |
| 77 | def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType: |
| 78 | uninitialized_params = [name for name, param in meta_model.named_parameters() if str(param.device) == "meta"] |
| 79 | uninitialized_buffers = [name for name, buffer in meta_model.named_buffers() if str(buffer.device) == "meta"] |
| 80 | if uninitialized_params or uninitialized_buffers: |
| 81 | logger.warning(f"Uninitialized parameters or buffers: {uninitialized_params + uninitialized_buffers}") |
| 82 | return meta_model |
| 83 | retval = meta_model.to(device) |
| 84 | return retval |
| 85 | |
| 86 | def build(self, device: torch.device | None = None, dtype: torch.dtype | None = None) -> ModelType: |
| 87 | device = torch.device("cuda") if device is None else device |
| 88 | config = self.model_config() |
| 89 | meta_model = self.meta_model(config, self.module_ops) |
| 90 | model_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path] |
| 91 | model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device) |
| 92 | |
| 93 | lora_strengths = [lora.strength for lora in self.loras] |
| 94 | if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0): |
| 95 | sd = model_state_dict.sd |
| 96 | if dtype is not None: |
| 97 | sd = {key: value.to(dtype=dtype) for key, value in model_state_dict.sd.items()} |
| 98 | meta_model.load_state_dict(sd, strict=False, assign=True) |
| 99 | return self._return_model(meta_model, device) |
| 100 | |
| 101 | lora_state_dicts = [ |
| 102 | self.load_sd([lora.path], sd_ops=lora.sd_ops, registry=self.registry, device=self.lora_load_device) |
| 103 | for lora in self.loras |
| 104 | ] |
| 105 | lora_sd_and_strengths = [ |
| 106 | LoraStateDictWithStrength(sd, strength) |
| 107 | for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True) |
| 108 | ] |
| 109 | final_sd = apply_loras( |
| 110 | model_sd=model_state_dict, |
| 111 | lora_sd_and_strengths=lora_sd_and_strengths, |
| 112 | dtype=dtype, |
| 113 | destination_sd=model_state_dict if isinstance(self.registry, DummyRegistry) else None, |
| 114 | ) |
| 115 | meta_model.load_state_dict(final_sd.sd, strict=False, assign=True) |
| 116 | return self._return_model(meta_model, device) |
| 117 |