(self)
| 195 | return self.model.training_losses(x1=x1, x0=x0, **kwargs).mean() |
| 196 | |
| 197 | def add_lora_to_unet(self): |
| 198 | unet = self.model |
| 199 | freeze(unet) |
| 200 | self.params_to_optimize = [] |
| 201 | self.params_names = [] |
| 202 | lora_cfg = self.lora_cfg |
| 203 | |
| 204 | assert lora_cfg is not None, "LoRA config cannot be None" |
| 205 | do_lora_conv = "lora_conv" in lora_cfg |
| 206 | do_first_full_conv = lora_cfg.get("do_full_first_conv", False) |
| 207 | do_lora_self_attn = "lora_self_attn" in lora_cfg |
| 208 | do_lora_cross_attn = "lora_cross_attn" in lora_cfg |
| 209 | do_lora_full_attn = "lora_full_attn" in lora_cfg |
| 210 | do_lora_attn = do_lora_self_attn or do_lora_cross_attn or do_lora_full_attn |
| 211 | do_lora_mlp = "lora_mlp" in lora_cfg |
| 212 | assert not (do_lora_full_attn and (do_lora_self_attn or do_lora_cross_attn)), "LoRA full attn cannot be used with self or cross attn" |
| 213 | assert do_lora_conv or do_lora_attn, "LoRA config must contain either 'lora_conv' or 'lora_attn'" |
| 214 | |
| 215 | lora_scale = lora_cfg.get("lora_scale", 1.0) |
| 216 | if self.lora_type == "lora_adapter": |
| 217 | self.conv_data_provider_names = [] |
| 218 | lora_cdim = lora_cfg.get("lora_cdim", 4) |
| 219 | |
| 220 | # LEGACY: _rank and _ratio keys shouldnt be in the general config file! |
| 221 | lora_cfg_cp = {k: v for k, v in lora_cfg.items() if not k.endswith("_ratio") and not k.endswith("_rank")} |
| 222 | # Populate keys |
| 223 | for k in list(lora_cfg_cp.keys()): |
| 224 | # if k not in ["lora_conv", "lora_self_attn", "lora_cross_attn", "lora_full_attn", "do_full_first_conv", "lora_scale", "do_full_last_conv"]: |
| 225 | # raise ValueError(f"Unknown LoRA config key {k}") |
| 226 | # if not isinstance(lora_cfg_cp[k], (float, int)): |
| 227 | # raise ValueError(f"LoRA config key must be float or int") |
| 228 | if isinstance(lora_cfg_cp[k], float): |
| 229 | lora_cfg_cp[k+"_ratio"] = lora_cfg_cp[k] |
| 230 | elif isinstance(lora_cfg_cp[k], int): |
| 231 | lora_cfg_cp[k+"_rank"] = lora_cfg_cp[k] |
| 232 | |
| 233 | # train last layer fully |
| 234 | do_full_last_conv = lora_cfg_cp.get("do_full_last_conv", False) |
| 235 | if do_full_last_conv: |
| 236 | self.params_to_optimize.extend([p for p in self.model.net.out.parameters()]) |
| 237 | self.model.net.out[0].weight.requires_grad = True |
| 238 | self.model.net.out[0].bias.requires_grad = True |
| 239 | self.model.net.out[2].weight.requires_grad = True |
| 240 | self.model.net.out[2].bias.requires_grad = True |
| 241 | |
| 242 | for path, w in unet.state_dict().items(): |
| 243 | if "." not in path: |
| 244 | continue |
| 245 | # Determine whether we want to finetune the first full convolutional layer |
| 246 | if "net.input_blocks.0.0" in path and do_first_full_conv: |
| 247 | self.params_to_optimize.append(w) |
| 248 | self.model.net.input_blocks[0][0].weight.requires_grad = True |
| 249 | self.model.net.input_blocks[0][0].bias.requires_grad = True |
| 250 | |
| 251 | # Add LoRA layers to the full attention modules |
| 252 | elif "attn" in path and do_lora_full_attn: |
| 253 | if path.split(".")[-2] not in ["to_q", "to_k", "to_v", "query", "key", "value"]: |
| 254 | continue |
no test coverage detected