| 1 | from dataclasses import dataclass, replace |
| 2 | from typing import NamedTuple, Protocol |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | |
| 7 | @dataclass(frozen=True, slots=True) |
| 8 | class ContentReplacement: |
| 9 | """ |
| 10 | Represents a content replacement operation. |
| 11 | Used to replace a specific content with a replacement in a state dict key. |
| 12 | """ |
| 13 | |
| 14 | content: str |
| 15 | replacement: str |
| 16 | |
| 17 | |
| 18 | @dataclass(frozen=True, slots=True) |
| 19 | class ContentMatching: |
| 20 | """ |
| 21 | Represents a content matching operation. |
| 22 | Used to match a specific prefix and suffix in a state dict key. |
| 23 | """ |
| 24 | |
| 25 | prefix: str = "" |
| 26 | suffix: str = "" |
| 27 | |
| 28 | |
| 29 | class KeyValueOperationResult(NamedTuple): |
| 30 | """ |
| 31 | Represents the result of a key-value operation. |
| 32 | Contains the new key and value after the operation has been applied. |
| 33 | """ |
| 34 | |
| 35 | new_key: str |
| 36 | new_value: torch.Tensor |
| 37 | |
| 38 | |
| 39 | class KeyValueOperation(Protocol): |
| 40 | """ |
| 41 | Protocol for key-value operations. |
| 42 | Used to apply operations to a specific key and value in a state dict. |
| 43 | """ |
| 44 | |
| 45 | def __call__(self, tensor_key: str, tensor_value: torch.Tensor) -> list[KeyValueOperationResult]: ... |
| 46 | |
| 47 | |
| 48 | @dataclass(frozen=True, slots=True) |
| 49 | class SDKeyValueOperation: |
| 50 | """ |
| 51 | Represents a key-value operation. |
| 52 | Used to apply operations to a specific key and value in a state dict. |
| 53 | """ |
| 54 | |
| 55 | key_matcher: ContentMatching |
| 56 | kv_operation: KeyValueOperation |
| 57 | |
| 58 | |
| 59 | @dataclass(frozen=True, slots=True) |
| 60 | class SDOps: |
| 61 | """Immutable class representing state dict key operations.""" |
| 62 | |
| 63 | name: str |
| 64 | mapping: tuple[ |
| 65 | ContentReplacement | ContentMatching | SDKeyValueOperation, ... |
| 66 | ] = () # Immutable tuple of (key, value) pairs |
| 67 | |
| 68 | def with_replacement(self, content: str, replacement: str) -> "SDOps": |
| 69 | """Create a new SDOps instance with the specified replacement added to the mapping.""" |
| 70 | |
| 71 | new_mapping = (*self.mapping, ContentReplacement(content, replacement)) |
| 72 | return replace(self, mapping=new_mapping) |
| 73 | |
| 74 | def with_matching(self, prefix: str = "", suffix: str = "") -> "SDOps": |
| 75 | """Create a new SDOps instance with the specified prefix and suffix matching added to the mapping.""" |
| 76 | |
| 77 | new_mapping = (*self.mapping, ContentMatching(prefix, suffix)) |
| 78 | return replace(self, mapping=new_mapping) |
| 79 | |
| 80 | def with_kv_operation( |
| 81 | self, |
| 82 | operation: KeyValueOperation, |
| 83 | key_prefix: str = "", |
| 84 | key_suffix: str = "", |
| 85 | ) -> "SDOps": |
| 86 | """Create a new SDOps instance with the specified value operation added to the mapping.""" |
| 87 | key_matcher = ContentMatching(key_prefix, key_suffix) |
| 88 | sd_kv_operation = SDKeyValueOperation(key_matcher, operation) |
| 89 | new_mapping = (*self.mapping, sd_kv_operation) |
| 90 | return replace(self, mapping=new_mapping) |
| 91 | |
| 92 | def apply_to_key(self, key: str) -> str | None: |
| 93 | """Apply the mapping to the given name.""" |
| 94 | matchers = [content for content in self.mapping if isinstance(content, ContentMatching)] |
| 95 | valid = any(key.startswith(f.prefix) and key.endswith(f.suffix) for f in matchers) |
| 96 | if not valid: |
| 97 | return None |
| 98 | |
| 99 | for replacement in self.mapping: |
| 100 | if not isinstance(replacement, ContentReplacement): |
| 101 | continue |
| 102 | if replacement.content in key: |
| 103 | key = key.replace(replacement.content, replacement.replacement) |
| 104 | return key |
| 105 | |
| 106 | def apply_to_key_value(self, key: str, value: torch.Tensor) -> list[KeyValueOperationResult]: |
| 107 | """Apply the value operation to the given name and associated value.""" |
| 108 | for operation in self.mapping: |
| 109 | if not isinstance(operation, SDKeyValueOperation): |
| 110 | continue |
| 111 | if key.startswith(operation.key_matcher.prefix) and key.endswith(operation.key_matcher.suffix): |
| 112 | return operation.kv_operation(key, value) |
| 113 | return [KeyValueOperationResult(key, value)] |
| 114 | |
| 115 | |
| 116 | # Predefined SDOps instances |
| 117 | LTXV_LORA_COMFY_RENAMING_MAP = ( |
| 118 | SDOps("LTXV_LORA_COMFY_PREFIX_MAP").with_matching().with_replacement("diffusion_model.", "") |
| 119 | ) |
| 120 | |
| 121 | LTXV_LORA_COMFY_TARGET_MAP = ( |
| 122 | SDOps("LTXV_LORA_COMFY_TARGET_MAP") |
| 123 | .with_matching() |
| 124 | .with_replacement("diffusion_model.", "") |
| 125 | .with_replacement(".lora_A.weight", ".weight") |
| 126 | .with_replacement(".lora_B.weight", ".weight") |
| 127 | ) |
| 128 |