| 165 | LORA_PREFIX_TRANSFORMER = "lora_unet" |
| 166 | LORA_PREFIX_TEXT_ENCODER = "lora_te" |
| 167 | def __init__( |
| 168 | self, |
| 169 | text_encoder: Union[List[T5EncoderModel], T5EncoderModel], |
| 170 | unet, |
| 171 | multiplier: float = 1.0, |
| 172 | lora_dim: int = 4, |
| 173 | alpha: float = 1, |
| 174 | dropout: Optional[float] = None, |
| 175 | module_class: Type[object] = LoRAModule, |
| 176 | skip_name: str = None, |
| 177 | target_name: str = None, |
| 178 | varbose: Optional[bool] = False, |
| 179 | ) -> None: |
| 180 | super().__init__() |
| 181 | self.multiplier = multiplier |
| 182 | |
| 183 | self.lora_dim = lora_dim |
| 184 | self.alpha = alpha |
| 185 | self.dropout = dropout |
| 186 | |
| 187 | print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}") |
| 188 | print(f"neuron dropout: p={self.dropout}") |
| 189 | |
| 190 | # create module instances |
| 191 | def create_modules( |
| 192 | is_unet: bool, |
| 193 | root_module: torch.nn.Module, |
| 194 | target_replace_modules: List[torch.nn.Module], |
| 195 | ) -> List[LoRAModule]: |
| 196 | prefix = ( |
| 197 | self.LORA_PREFIX_TRANSFORMER |
| 198 | if is_unet |
| 199 | else self.LORA_PREFIX_TEXT_ENCODER |
| 200 | ) |
| 201 | loras = [] |
| 202 | skipped = [] |
| 203 | for name, module in root_module.named_modules(): |
| 204 | if module.__class__.__name__ in target_replace_modules: |
| 205 | for child_name, child_module in module.named_modules(): |
| 206 | is_linear = child_module.__class__.__name__ == "Linear" or child_module.__class__.__name__ == "LoRACompatibleLinear" |
| 207 | is_conv2d = child_module.__class__.__name__ == "Conv2d" or child_module.__class__.__name__ == "LoRACompatibleConv" |
| 208 | is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1) |
| 209 | |
| 210 | if skip_name is not None and skip_name in child_name: |
| 211 | continue |
| 212 | |
| 213 | if target_name is not None: |
| 214 | target_name_in = False |
| 215 | if isinstance(target_name, str): |
| 216 | target_name_in = target_name in child_name |
| 217 | elif isinstance(target_name, list): |
| 218 | target_name_in = any([_target_name in child_name for _target_name in target_name]) |
| 219 | if not target_name_in: |
| 220 | continue |
| 221 | |
| 222 | if is_linear or is_conv2d: |
| 223 | lora_name = prefix + "." + name + "." + child_name |
| 224 | lora_name = lora_name.replace(".", "_") |