| 1 | import hashlib |
| 2 | import threading |
| 3 | from dataclasses import dataclass, field |
| 4 | from pathlib import Path |
| 5 | from typing import Protocol |
| 6 | |
| 7 | from ltx_core.loader.primitives import StateDict |
| 8 | from ltx_core.loader.sd_ops import SDOps |
| 9 | |
| 10 | |
| 11 | class Registry(Protocol): |
| 12 | """ |
| 13 | Protocol for managing state dictionaries in a registry. |
| 14 | It is used to store state dictionaries and reuse them later without loading them again. |
| 15 | Implementations must provide: |
| 16 | - add: Add a state dictionary to the registry |
| 17 | - pop: Remove a state dictionary from the registry |
| 18 | - get: Retrieve a state dictionary from the registry |
| 19 | - clear: Clear all state dictionaries from the registry |
| 20 | """ |
| 21 | |
| 22 | def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None: ... |
| 23 | |
| 24 | def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ... |
| 25 | |
| 26 | def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ... |
| 27 | |
| 28 | def clear(self) -> None: ... |
| 29 | |
| 30 | |
| 31 | class DummyRegistry(Registry): |
| 32 | """ |
| 33 | Dummy registry that does not store state dictionaries. |
| 34 | """ |
| 35 | |
| 36 | def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None: |
| 37 | pass |
| 38 | |
| 39 | def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: |
| 40 | pass |
| 41 | |
| 42 | def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: |
| 43 | pass |
| 44 | |
| 45 | def clear(self) -> None: |
| 46 | pass |
| 47 | |
| 48 | |
| 49 | @dataclass |
| 50 | class StateDictRegistry(Registry): |
| 51 | """ |
| 52 | Registry that stores state dictionaries in a dictionary. |
| 53 | """ |
| 54 | |
| 55 | _state_dicts: dict[str, StateDict] = field(default_factory=dict) |
| 56 | _lock: threading.Lock = field(default_factory=threading.Lock) |
| 57 | |
| 58 | def _generate_id(self, paths: list[str], sd_ops: SDOps) -> str: |
| 59 | m = hashlib.sha256() |
| 60 | parts = [str(Path(p).resolve()) for p in paths] |
| 61 | if sd_ops is not None: |
| 62 | parts.append(sd_ops.name) |
| 63 | m.update("\0".join(parts).encode("utf-8")) |
| 64 | return m.hexdigest() |
| 65 | |
| 66 | def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> str: |
| 67 | sd_id = self._generate_id(paths, sd_ops) |
| 68 | with self._lock: |
| 69 | if sd_id in self._state_dicts: |
| 70 | raise ValueError(f"State dict retrieved from {paths} with {sd_ops} already added, check with get first") |
| 71 | self._state_dicts[sd_id] = state_dict |
| 72 | return sd_id |
| 73 | |
| 74 | def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: |
| 75 | with self._lock: |
| 76 | return self._state_dicts.pop(self._generate_id(paths, sd_ops), None) |
| 77 | |
| 78 | def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: |
| 79 | with self._lock: |
| 80 | return self._state_dicts.get(self._generate_id(paths, sd_ops), None) |
| 81 | |
| 82 | def clear(self) -> None: |
| 83 | with self._lock: |
| 84 | self._state_dicts.clear() |
| 85 |