| 1 | from dataclasses import dataclass |
| 2 | from typing import NamedTuple, Protocol |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from ltx_core.loader.module_ops import ModuleOps |
| 7 | from ltx_core.loader.sd_ops import SDOps |
| 8 | from ltx_core.model.model_protocol import ModelType |
| 9 | |
| 10 | |
| 11 | @dataclass(frozen=True) |
| 12 | class StateDict: |
| 13 | """ |
| 14 | Immutable container for a PyTorch state dictionary. |
| 15 | Contains: |
| 16 | - sd: Dictionary of tensors (weights, buffers, etc.) |
| 17 | - device: Device where tensors are stored |
| 18 | - size: Total memory footprint in bytes |
| 19 | - dtype: Set of tensor dtypes present |
| 20 | """ |
| 21 | |
| 22 | sd: dict |
| 23 | device: torch.device |
| 24 | size: int |
| 25 | dtype: set[torch.dtype] |
| 26 | |
| 27 | def footprint(self) -> tuple[int, torch.device]: |
| 28 | return self.size, self.device |
| 29 | |
| 30 | |
| 31 | class StateDictLoader(Protocol): |
| 32 | """ |
| 33 | Protocol for loading state dictionaries from various sources. |
| 34 | Implementations must provide: |
| 35 | - metadata: Extract model metadata from a single path |
| 36 | - load: Load state dict from path(s) and apply SDOps transformations |
| 37 | """ |
| 38 | |
| 39 | def metadata(self, path: str) -> dict: |
| 40 | """ |
| 41 | Load metadata from path |
| 42 | """ |
| 43 | |
| 44 | def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict: |
| 45 | """ |
| 46 | Load state dict from path or paths (for sharded model storage) and apply sd_ops |
| 47 | """ |
| 48 | |
| 49 | |
| 50 | class ModelBuilderProtocol(Protocol[ModelType]): |
| 51 | """ |
| 52 | Protocol for building PyTorch models from configuration dictionaries. |
| 53 | Implementations must provide: |
| 54 | - meta_model: Create a model from configuration dictionary and apply module operations |
| 55 | - build: Create and initialize a model from state dictionary and apply dtype transformations |
| 56 | """ |
| 57 | |
| 58 | def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType: |
| 59 | """ |
| 60 | Create a model on the meta device from a configuration dictionary. |
| 61 | This decouples model creation from weight loading, allowing the model |
| 62 | architecture to be instantiated without allocating memory for parameters. |
| 63 | Args: |
| 64 | config: Model configuration dictionary. |
| 65 | module_ops: Optional list of module operations to apply (e.g., quantization). |
| 66 | Returns: |
| 67 | Model instance on meta device (no actual memory allocated for parameters). |
| 68 | """ |
| 69 | ... |
| 70 | |
| 71 | def build(self, dtype: torch.dtype | None = None) -> ModelType: |
| 72 | """ |
| 73 | Build the model |
| 74 | Args: |
| 75 | dtype: Target dtype for the model, if None, uses the dtype of the model_path model |
| 76 | Returns: |
| 77 | Model instance |
| 78 | """ |
| 79 | ... |
| 80 | |
| 81 | |
| 82 | class LoRAAdaptableProtocol(Protocol): |
| 83 | """ |
| 84 | Protocol for models that can be adapted with LoRAs. |
| 85 | Implementations must provide: |
| 86 | - lora: Add a LoRA to the model |
| 87 | """ |
| 88 | |
| 89 | def lora(self, lora_path: str, strength: float) -> "LoRAAdaptableProtocol": |
| 90 | pass |
| 91 | |
| 92 | |
| 93 | class LoraPathStrengthAndSDOps(NamedTuple): |
| 94 | """ |
| 95 | Tuple containing a LoRA path, strength, and SDOps for applying to the LoRA state dict. |
| 96 | """ |
| 97 | |
| 98 | path: str |
| 99 | strength: float |
| 100 | sd_ops: SDOps |
| 101 | |
| 102 | |
| 103 | class LoraStateDictWithStrength(NamedTuple): |
| 104 | """ |
| 105 | Tuple containing a LoRA state dict and strength for applying to the model. |
| 106 | """ |
| 107 | |
| 108 | state_dict: StateDict |
| 109 | strength: float |
| 110 |