fred-dev/comfy_ui_ali
0
1"""2 This file is part of ComfyUI.3 Copyright (C) 2024 Comfy4 5 This program is free software: you can redistribute it and/or modify6 it under the terms of the GNU General Public License as published by7 the Free Software Foundation, either version 3 of the License, or8 (at your option) any later version.9 10 This program is distributed in the hope that it will be useful,11 but WITHOUT ANY WARRANTY; without even the implied warranty of12 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the13 GNU General Public License for more details.14 15 You should have received a copy of the GNU General Public License16 along with this program. If not, see <https://www.gnu.org/licenses/>.17"""18 19from __future__ import annotations20import comfy.utils21import comfy.model_management22import comfy.model_base23import logging24import torch25 26LORA_CLIP_MAP = {27 "mlp.fc1": "mlp_fc1",28 "mlp.fc2": "mlp_fc2",29 "self_attn.k_proj": "self_attn_k_proj",30 "self_attn.q_proj": "self_attn_q_proj",31 "self_attn.v_proj": "self_attn_v_proj",32 "self_attn.out_proj": "self_attn_out_proj",33}34 35 36def load_lora(lora, to_load, log_missing=True):37 patch_dict = {}38 loaded_keys = set()39 for x in to_load:40 alpha_name = "{}.alpha".format(x)41 alpha = None42 if alpha_name in lora.keys():43 alpha = lora[alpha_name].item()44 loaded_keys.add(alpha_name)45 46 dora_scale_name = "{}.dora_scale".format(x)47 dora_scale = None48 if dora_scale_name in lora.keys():49 dora_scale = lora[dora_scale_name]50 loaded_keys.add(dora_scale_name)51 52 reshape_name = "{}.reshape_weight".format(x)53 reshape = None54 if reshape_name in lora.keys():55 try:56 reshape = lora[reshape_name].tolist()57 loaded_keys.add(reshape_name)58 except:59 pass60 61 regular_lora = "{}.lora_up.weight".format(x)62 diffusers_lora = "{}_lora.up.weight".format(x)63 diffusers2_lora = "{}.lora_B.weight".format(x)64 diffusers3_lora = "{}.lora.up.weight".format(x)65 mochi_lora = "{}.lora_B".format(x)66 transformers_lora = "{}.lora_linear_layer.up.weight".format(x)67 A_name = None68 69 if regular_lora in lora.keys():70 A_name = regular_lora71 B_name = "{}.lora_down.weight".format(x)72 mid_name = "{}.lora_mid.weight".format(x)73 elif diffusers_lora in lora.keys():74 A_name = diffusers_lora75 B_name = "{}_lora.down.weight".format(x)76 mid_name = None77 elif diffusers2_lora in lora.keys():78 A_name = diffusers2_lora79 B_name = "{}.lora_A.weight".format(x)80 mid_name = None81 elif diffusers3_lora in lora.keys():82 A_name = diffusers3_lora83 B_name = "{}.lora.down.weight".format(x)84 mid_name = None85 elif mochi_lora in lora.keys():86 A_name = mochi_lora87 B_name = "{}.lora_A".format(x)88 mid_name = None89 elif transformers_lora in lora.keys():90 A_name = transformers_lora91 B_name ="{}.lora_linear_layer.down.weight".format(x)92 mid_name = None93 94 if A_name is not None:95 mid = None96 if mid_name is not None and mid_name in lora.keys():97 mid = lora[mid_name]98 loaded_keys.add(mid_name)99 patch_dict[to_load[x]] = ("lora", (lora[A_name], lora[B_name], alpha, mid, dora_scale, reshape))100 loaded_keys.add(A_name)101 loaded_keys.add(B_name)102 103 104 ######## loha105 hada_w1_a_name = "{}.hada_w1_a".format(x)106 hada_w1_b_name = "{}.hada_w1_b".format(x)107 hada_w2_a_name = "{}.hada_w2_a".format(x)108 hada_w2_b_name = "{}.hada_w2_b".format(x)109 hada_t1_name = "{}.hada_t1".format(x)110 hada_t2_name = "{}.hada_t2".format(x)111 if hada_w1_a_name in lora.keys():112 hada_t1 = None113 hada_t2 = None114 if hada_t1_name in lora.keys():115 hada_t1 = lora[hada_t1_name]116 hada_t2 = lora[hada_t2_name]117 loaded_keys.add(hada_t1_name)118 loaded_keys.add(hada_t2_name)119 120 patch_dict[to_load[x]] = ("loha", (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2, dora_scale))121 loaded_keys.add(hada_w1_a_name)122 loaded_keys.add(hada_w1_b_name)123 loaded_keys.add(hada_w2_a_name)124 loaded_keys.add(hada_w2_b_name)125 126 127 ######## lokr128 lokr_w1_name = "{}.lokr_w1".format(x)129 lokr_w2_name = "{}.lokr_w2".format(x)130 lokr_w1_a_name = "{}.lokr_w1_a".format(x)131 lokr_w1_b_name = "{}.lokr_w1_b".format(x)132 lokr_t2_name = "{}.lokr_t2".format(x)133 lokr_w2_a_name = "{}.lokr_w2_a".format(x)134 lokr_w2_b_name = "{}.lokr_w2_b".format(x)135 136 lokr_w1 = None137 if lokr_w1_name in lora.keys():138 lokr_w1 = lora[lokr_w1_name]139 loaded_keys.add(lokr_w1_name)140 141 lokr_w2 = None142 if lokr_w2_name in lora.keys():143 lokr_w2 = lora[lokr_w2_name]144 loaded_keys.add(lokr_w2_name)145 146 lokr_w1_a = None147 if lokr_w1_a_name in lora.keys():148 lokr_w1_a = lora[lokr_w1_a_name]149 loaded_keys.add(lokr_w1_a_name)150 151 lokr_w1_b = None152 if lokr_w1_b_name in lora.keys():153 lokr_w1_b = lora[lokr_w1_b_name]154 loaded_keys.add(lokr_w1_b_name)155 156 lokr_w2_a = None157 if lokr_w2_a_name in lora.keys():158 lokr_w2_a = lora[lokr_w2_a_name]159 loaded_keys.add(lokr_w2_a_name)160 161 lokr_w2_b = None162 if lokr_w2_b_name in lora.keys():163 lokr_w2_b = lora[lokr_w2_b_name]164 loaded_keys.add(lokr_w2_b_name)165 166 lokr_t2 = None167 if lokr_t2_name in lora.keys():168 lokr_t2 = lora[lokr_t2_name]169 loaded_keys.add(lokr_t2_name)170 171 if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None):172 patch_dict[to_load[x]] = ("lokr", (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2, dora_scale))173 174 #glora175 a1_name = "{}.a1.weight".format(x)176 a2_name = "{}.a2.weight".format(x)177 b1_name = "{}.b1.weight".format(x)178 b2_name = "{}.b2.weight".format(x)179 if a1_name in lora:180 patch_dict[to_load[x]] = ("glora", (lora[a1_name], lora[a2_name], lora[b1_name], lora[b2_name], alpha, dora_scale))181 loaded_keys.add(a1_name)182 loaded_keys.add(a2_name)183 loaded_keys.add(b1_name)184 loaded_keys.add(b2_name)185 186 w_norm_name = "{}.w_norm".format(x)187 b_norm_name = "{}.b_norm".format(x)188 w_norm = lora.get(w_norm_name, None)189 b_norm = lora.get(b_norm_name, None)190 191 if w_norm is not None:192 loaded_keys.add(w_norm_name)193 patch_dict[to_load[x]] = ("diff", (w_norm,))194 if b_norm is not None:195 loaded_keys.add(b_norm_name)196 patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (b_norm,))197 198 diff_name = "{}.diff".format(x)199 diff_weight = lora.get(diff_name, None)200 if diff_weight is not None:201 patch_dict[to_load[x]] = ("diff", (diff_weight,))202 loaded_keys.add(diff_name)203 204 diff_bias_name = "{}.diff_b".format(x)205 diff_bias = lora.get(diff_bias_name, None)206 if diff_bias is not None:207 patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (diff_bias,))208 loaded_keys.add(diff_bias_name)209 210 set_weight_name = "{}.set_weight".format(x)211 set_weight = lora.get(set_weight_name, None)212 if set_weight is not None:213 patch_dict[to_load[x]] = ("set", (set_weight,))214 loaded_keys.add(set_weight_name)215 216 if log_missing:217 for x in lora.keys():218 if x not in loaded_keys:219 logging.warning("lora key not loaded: {}".format(x))220 221 return patch_dict222 223def model_lora_keys_clip(model, key_map={}):224 sdk = model.state_dict().keys()225 for k in sdk:226 if k.endswith(".weight"):227 key_map["text_encoders.{}".format(k[:-len(".weight")])] = k #generic lora format without any weird key names228 229 text_model_lora_key = "lora_te_text_model_encoder_layers_{}_{}"230 clip_l_present = False231 clip_g_present = False232 for b in range(32): #TODO: clean up233 for c in LORA_CLIP_MAP:234 k = "clip_h.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)235 if k in sdk:236 lora_key = text_model_lora_key.format(b, LORA_CLIP_MAP[c])237 key_map[lora_key] = k238 lora_key = "lora_te1_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c])239 key_map[lora_key] = k240 lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora241 key_map[lora_key] = k242 243 k = "clip_l.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)244 if k in sdk:245 lora_key = text_model_lora_key.format(b, LORA_CLIP_MAP[c])246 key_map[lora_key] = k247 lora_key = "lora_te1_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #SDXL base248 key_map[lora_key] = k249 clip_l_present = True250 lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora251 key_map[lora_key] = k252 253 k = "clip_g.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)254 if k in sdk:255 clip_g_present = True256 if clip_l_present:257 lora_key = "lora_te2_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #SDXL base258 key_map[lora_key] = k259 lora_key = "text_encoder_2.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora260 key_map[lora_key] = k261 else:262 lora_key = "lora_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #TODO: test if this is correct for SDXL-Refiner263 key_map[lora_key] = k264 lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora265 key_map[lora_key] = k266 lora_key = "lora_prior_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #cascade lora: TODO put lora key prefix in the model config267 key_map[lora_key] = k268 269 for k in sdk:270 if k.endswith(".weight"):271 if k.startswith("t5xxl.transformer."):#OneTrainer SD3 and Flux lora272 l_key = k[len("t5xxl.transformer."):-len(".weight")]273 t5_index = 1274 if clip_g_present:275 t5_index += 1276 if clip_l_present:277 t5_index += 1278 if t5_index == 2:279 key_map["lora_te{}_{}".format(t5_index, l_key.replace(".", "_"))] = k #OneTrainer Flux280 t5_index += 1281 282 key_map["lora_te{}_{}".format(t5_index, l_key.replace(".", "_"))] = k283 elif k.startswith("hydit_clip.transformer.bert."): #HunyuanDiT Lora284 l_key = k[len("hydit_clip.transformer.bert."):-len(".weight")]285 lora_key = "lora_te1_{}".format(l_key.replace(".", "_"))286 key_map[lora_key] = k287 288 289 k = "clip_g.transformer.text_projection.weight"290 if k in sdk:291 key_map["lora_prior_te_text_projection"] = k #cascade lora?292 # key_map["text_encoder.text_projection"] = k #TODO: check if other lora have the text_projection too293 key_map["lora_te2_text_projection"] = k #OneTrainer SD3 lora294 295 k = "clip_l.transformer.text_projection.weight"296 if k in sdk:297 key_map["lora_te1_text_projection"] = k #OneTrainer SD3 lora, not necessary but omits warning298 299 return key_map300 301def model_lora_keys_unet(model, key_map={}):302 sd = model.state_dict()303 sdk = sd.keys()304 305 for k in sdk:306 if k.startswith("diffusion_model."):307 if k.endswith(".weight"):308 key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")309 key_map["lora_unet_{}".format(key_lora)] = k310 key_map["{}".format(k[:-len(".weight")])] = k #generic lora format without any weird key names311 else:312 key_map["{}".format(k)] = k #generic lora format for not .weight without any weird key names313 314 diffusers_keys = comfy.utils.unet_to_diffusers(model.model_config.unet_config)315 for k in diffusers_keys:316 if k.endswith(".weight"):317 unet_key = "diffusion_model.{}".format(diffusers_keys[k])318 key_lora = k[:-len(".weight")].replace(".", "_")319 key_map["lora_unet_{}".format(key_lora)] = unet_key320 key_map["lycoris_{}".format(key_lora)] = unet_key #simpletuner lycoris format321 322 diffusers_lora_prefix = ["", "unet."]323 for p in diffusers_lora_prefix:324 diffusers_lora_key = "{}{}".format(p, k[:-len(".weight")].replace(".to_", ".processor.to_"))325 if diffusers_lora_key.endswith(".to_out.0"):326 diffusers_lora_key = diffusers_lora_key[:-2]327 key_map[diffusers_lora_key] = unet_key328 329 if isinstance(model, comfy.model_base.StableCascade_C):330 for k in sdk:331 if k.startswith("diffusion_model."):332 if k.endswith(".weight"):333 key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")334 key_map["lora_prior_unet_{}".format(key_lora)] = k335 336 if isinstance(model, comfy.model_base.SD3): #Diffusers lora SD3337 diffusers_keys = comfy.utils.mmdit_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")338 for k in diffusers_keys:339 if k.endswith(".weight"):340 to = diffusers_keys[k]341 key_lora = "transformer.{}".format(k[:-len(".weight")]) #regular diffusers sd3 lora format342 key_map[key_lora] = to343 344 key_lora = "base_model.model.{}".format(k[:-len(".weight")]) #format for flash-sd3 lora and others?345 key_map[key_lora] = to346 347 key_lora = "lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_")) #OneTrainer lora348 key_map[key_lora] = to349 350 key_lora = "lycoris_{}".format(k[:-len(".weight")].replace(".", "_")) #simpletuner lycoris format351 key_map[key_lora] = to352 353 if isinstance(model, comfy.model_base.AuraFlow): #Diffusers lora AuraFlow354 diffusers_keys = comfy.utils.auraflow_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")355 for k in diffusers_keys:356 if k.endswith(".weight"):357 to = diffusers_keys[k]358 key_lora = "transformer.{}".format(k[:-len(".weight")]) #simpletrainer and probably regular diffusers lora format359 key_map[key_lora] = to360 361 if isinstance(model, comfy.model_base.PixArt):362 diffusers_keys = comfy.utils.pixart_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")363 for k in diffusers_keys:364 if k.endswith(".weight"):365 to = diffusers_keys[k]366 key_lora = "transformer.{}".format(k[:-len(".weight")]) #default format367 key_map[key_lora] = to368 369 key_lora = "base_model.model.{}".format(k[:-len(".weight")]) #diffusers training script370 key_map[key_lora] = to371 372 key_lora = "unet.base_model.model.{}".format(k[:-len(".weight")]) #old reference peft script373 key_map[key_lora] = to374 375 if isinstance(model, comfy.model_base.HunyuanDiT):376 for k in sdk:377 if k.startswith("diffusion_model.") and k.endswith(".weight"):378 key_lora = k[len("diffusion_model."):-len(".weight")]379 key_map["base_model.model.{}".format(key_lora)] = k #official hunyuan lora format380 381 if isinstance(model, comfy.model_base.Flux): #Diffusers lora Flux382 diffusers_keys = comfy.utils.flux_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")383 for k in diffusers_keys:384 if k.endswith(".weight"):385 to = diffusers_keys[k]386 key_map["transformer.{}".format(k[:-len(".weight")])] = to #simpletrainer and probably regular diffusers flux lora format387 key_map["lycoris_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #simpletrainer lycoris388 key_map["lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #onetrainer389 390 if isinstance(model, comfy.model_base.GenmoMochi):391 for k in sdk:392 if k.startswith("diffusion_model.") and k.endswith(".weight"): #Official Mochi lora format393 key_lora = k[len("diffusion_model."):-len(".weight")]394 key_map["{}".format(key_lora)] = k395 396 if isinstance(model, comfy.model_base.HunyuanVideo):397 for k in sdk:398 if k.startswith("diffusion_model.") and k.endswith(".weight"):399 # diffusion-pipe lora format400 key_lora = k401 key_lora = key_lora.replace("_mod.lin.", "_mod.linear.").replace("_attn.qkv.", "_attn_qkv.").replace("_attn.proj.", "_attn_proj.")402 key_lora = key_lora.replace("mlp.0.", "mlp.fc1.").replace("mlp.2.", "mlp.fc2.")403 key_lora = key_lora.replace(".modulation.lin.", ".modulation.linear.")404 key_lora = key_lora[len("diffusion_model."):-len(".weight")]405 key_map["transformer.{}".format(key_lora)] = k406 key_map["diffusion_model.{}".format(key_lora)] = k # Old loras407 408 return key_map409 410 411def weight_decompose(dora_scale, weight, lora_diff, alpha, strength, intermediate_dtype, function):412 dora_scale = comfy.model_management.cast_to_device(dora_scale, weight.device, intermediate_dtype)413 lora_diff *= alpha414 weight_calc = weight + function(lora_diff).type(weight.dtype)415 weight_norm = (416 weight_calc.transpose(0, 1)417 .reshape(weight_calc.shape[1], -1)418 .norm(dim=1, keepdim=True)419 .reshape(weight_calc.shape[1], *[1] * (weight_calc.dim() - 1))420 .transpose(0, 1)421 )422 423 weight_calc *= (dora_scale / weight_norm).type(weight.dtype)424 if strength != 1.0:425 weight_calc -= weight426 weight += strength * (weight_calc)427 else:428 weight[:] = weight_calc429 return weight430 431def pad_tensor_to_shape(tensor: torch.Tensor, new_shape: list[int]) -> torch.Tensor:432 """433 Pad a tensor to a new shape with zeros.434 435 Args:436 tensor (torch.Tensor): The original tensor to be padded.437 new_shape (List[int]): The desired shape of the padded tensor.438 439 Returns:440 torch.Tensor: A new tensor padded with zeros to the specified shape.441 442 Note:443 If the new shape is smaller than the original tensor in any dimension,444 the original tensor will be truncated in that dimension.445 """446 if any([new_shape[i] < tensor.shape[i] for i in range(len(new_shape))]):447 raise ValueError("The new shape must be larger than the original tensor in all dimensions")448 449 if len(new_shape) != len(tensor.shape):450 raise ValueError("The new shape must have the same number of dimensions as the original tensor")451 452 # Create a new tensor filled with zeros453 padded_tensor = torch.zeros(new_shape, dtype=tensor.dtype, device=tensor.device)454 455 # Create slicing tuples for both tensors456 orig_slices = tuple(slice(0, dim) for dim in tensor.shape)457 new_slices = tuple(slice(0, dim) for dim in tensor.shape)458 459 # Copy the original tensor into the new tensor460 padded_tensor[new_slices] = tensor[orig_slices]461 462 return padded_tensor463 464def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, original_weights=None):465 for p in patches:466 strength = p[0]467 v = p[1]468 strength_model = p[2]469 offset = p[3]470 function = p[4]471 if function is None:472 function = lambda a: a473 474 old_weight = None475 if offset is not None:476 old_weight = weight477 weight = weight.narrow(offset[0], offset[1], offset[2])478 479 if strength_model != 1.0:480 weight *= strength_model481 482 if isinstance(v, list):483 v = (calculate_weight(v[1:], v[0][1](comfy.model_management.cast_to_device(v[0][0], weight.device, intermediate_dtype, copy=True), inplace=True), key, intermediate_dtype=intermediate_dtype), )484 485 if len(v) == 1:486 patch_type = "diff"487 elif len(v) == 2:488 patch_type = v[0]489 v = v[1]490 491 if patch_type == "diff":492 diff: torch.Tensor = v[0]493 # An extra flag to pad the weight if the diff's shape is larger than the weight494 do_pad_weight = len(v) > 1 and v[1]['pad_weight']495 if do_pad_weight and diff.shape != weight.shape:496 logging.info("Pad weight {} from {} to shape: {}".format(key, weight.shape, diff.shape))497 weight = pad_tensor_to_shape(weight, diff.shape)498 499 if strength != 0.0:500 if diff.shape != weight.shape:501 logging.warning("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, diff.shape, weight.shape))502 else:503 weight += function(strength * comfy.model_management.cast_to_device(diff, weight.device, weight.dtype))504 elif patch_type == "set":505 weight.copy_(v[0])506 elif patch_type == "model_as_lora":507 target_weight: torch.Tensor = v[0]508 diff_weight = comfy.model_management.cast_to_device(target_weight, weight.device, intermediate_dtype) - \509 comfy.model_management.cast_to_device(original_weights[key][0][0], weight.device, intermediate_dtype)510 weight += function(strength * comfy.model_management.cast_to_device(diff_weight, weight.device, weight.dtype))511 elif patch_type == "lora": #lora/locon512 mat1 = comfy.model_management.cast_to_device(v[0], weight.device, intermediate_dtype)513 mat2 = comfy.model_management.cast_to_device(v[1], weight.device, intermediate_dtype)514 dora_scale = v[4]515 reshape = v[5]516 517 if reshape is not None:518 weight = pad_tensor_to_shape(weight, reshape)519 520 if v[2] is not None:521 alpha = v[2] / mat2.shape[0]522 else:523 alpha = 1.0524 525 if v[3] is not None:526 #locon mid weights, hopefully the math is fine because I didn't properly test it527 mat3 = comfy.model_management.cast_to_device(v[3], weight.device, intermediate_dtype)528 final_shape = [mat2.shape[1], mat2.shape[0], mat3.shape[2], mat3.shape[3]]529 mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1), mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1)530 try:531 lora_diff = torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1)).reshape(weight.shape)532 if dora_scale is not None:533 weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength, intermediate_dtype, function)534 else:535 weight += function(((strength * alpha) * lora_diff).type(weight.dtype))536 except Exception as e:537 logging.error("ERROR {} {} {}".format(patch_type, key, e))538 elif patch_type == "lokr":539 w1 = v[0]540 w2 = v[1]541 w1_a = v[3]542 w1_b = v[4]543 w2_a = v[5]544 w2_b = v[6]545 t2 = v[7]546 dora_scale = v[8]547 dim = None548 549 if w1 is None:550 dim = w1_b.shape[0]551 w1 = torch.mm(comfy.model_management.cast_to_device(w1_a, weight.device, intermediate_dtype),552 comfy.model_management.cast_to_device(w1_b, weight.device, intermediate_dtype))553 else:554 w1 = comfy.model_management.cast_to_device(w1, weight.device, intermediate_dtype)555 556 if w2 is None:557 dim = w2_b.shape[0]558 if t2 is None:559 w2 = torch.mm(comfy.model_management.cast_to_device(w2_a, weight.device, intermediate_dtype),560 comfy.model_management.cast_to_device(w2_b, weight.device, intermediate_dtype))561 else:562 w2 = torch.einsum('i j k l, j r, i p -> p r k l',563 comfy.model_management.cast_to_device(t2, weight.device, intermediate_dtype),564 comfy.model_management.cast_to_device(w2_b, weight.device, intermediate_dtype),565 comfy.model_management.cast_to_device(w2_a, weight.device, intermediate_dtype))566 else:567 w2 = comfy.model_management.cast_to_device(w2, weight.device, intermediate_dtype)568 569 if len(w2.shape) == 4:570 w1 = w1.unsqueeze(2).unsqueeze(2)571 if v[2] is not None and dim is not None:572 alpha = v[2] / dim573 else:574 alpha = 1.0575 576 try:577 lora_diff = torch.kron(w1, w2).reshape(weight.shape)578 if dora_scale is not None:579 weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength, intermediate_dtype, function)580 else:581 weight += function(((strength * alpha) * lora_diff).type(weight.dtype))582 except Exception as e:583 logging.error("ERROR {} {} {}".format(patch_type, key, e))584 elif patch_type == "loha":585 w1a = v[0]586 w1b = v[1]587 if v[2] is not None:588 alpha = v[2] / w1b.shape[0]589 else:590 alpha = 1.0591 592 w2a = v[3]593 w2b = v[4]594 dora_scale = v[7]595 if v[5] is not None: #cp decomposition596 t1 = v[5]597 t2 = v[6]598 m1 = torch.einsum('i j k l, j r, i p -> p r k l',599 comfy.model_management.cast_to_device(t1, weight.device, intermediate_dtype),600 comfy.model_management.cast_to_device(w1b, weight.device, intermediate_dtype),601 comfy.model_management.cast_to_device(w1a, weight.device, intermediate_dtype))602 603 m2 = torch.einsum('i j k l, j r, i p -> p r k l',604 comfy.model_management.cast_to_device(t2, weight.device, intermediate_dtype),605 comfy.model_management.cast_to_device(w2b, weight.device, intermediate_dtype),606 comfy.model_management.cast_to_device(w2a, weight.device, intermediate_dtype))607 else:608 m1 = torch.mm(comfy.model_management.cast_to_device(w1a, weight.device, intermediate_dtype),609 comfy.model_management.cast_to_device(w1b, weight.device, intermediate_dtype))610 m2 = torch.mm(comfy.model_management.cast_to_device(w2a, weight.device, intermediate_dtype),611 comfy.model_management.cast_to_device(w2b, weight.device, intermediate_dtype))612 613 try:614 lora_diff = (m1 * m2).reshape(weight.shape)615 if dora_scale is not None:616 weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength, intermediate_dtype, function)617 else:618 weight += function(((strength * alpha) * lora_diff).type(weight.dtype))619 except Exception as e:620 logging.error("ERROR {} {} {}".format(patch_type, key, e))621 elif patch_type == "glora":622 dora_scale = v[5]623 624 old_glora = False625 if v[3].shape[1] == v[2].shape[0] == v[0].shape[0] == v[1].shape[1]:626 rank = v[0].shape[0]627 old_glora = True628 629 if v[3].shape[0] == v[2].shape[1] == v[0].shape[1] == v[1].shape[0]:630 if old_glora and v[1].shape[0] == weight.shape[0] and weight.shape[0] == weight.shape[1]:631 pass632 else:633 old_glora = False634 rank = v[1].shape[0]635 636 a1 = comfy.model_management.cast_to_device(v[0].flatten(start_dim=1), weight.device, intermediate_dtype)637 a2 = comfy.model_management.cast_to_device(v[1].flatten(start_dim=1), weight.device, intermediate_dtype)638 b1 = comfy.model_management.cast_to_device(v[2].flatten(start_dim=1), weight.device, intermediate_dtype)639 b2 = comfy.model_management.cast_to_device(v[3].flatten(start_dim=1), weight.device, intermediate_dtype)640 641 if v[4] is not None:642 alpha = v[4] / rank643 else:644 alpha = 1.0645 646 try:647 if old_glora:648 lora_diff = (torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1).to(dtype=intermediate_dtype), a2), a1)).reshape(weight.shape) #old lycoris glora649 else:650 if weight.dim() > 2:651 lora_diff = torch.einsum("o i ..., i j -> o j ...", torch.einsum("o i ..., i j -> o j ...", weight.to(dtype=intermediate_dtype), a1), a2).reshape(weight.shape)652 else:653 lora_diff = torch.mm(torch.mm(weight.to(dtype=intermediate_dtype), a1), a2).reshape(weight.shape)654 lora_diff += torch.mm(b1, b2).reshape(weight.shape)655 656 if dora_scale is not None:657 weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength, intermediate_dtype, function)658 else:659 weight += function(((strength * alpha) * lora_diff).type(weight.dtype))660 except Exception as e:661 logging.error("ERROR {} {} {}".format(patch_type, key, e))662 else:663 logging.warning("patch type not recognized {} {}".format(patch_type, key))664 665 if old_weight is not None:666 weight = old_weight667 668 return weight669