Spaces:
Configuration error
Configuration error
| import os | |
| import copy | |
| import json | |
| import torch | |
| import comfy.lora | |
| import comfy.model_management | |
| from comfy.model_patcher import ModelPatcher | |
| from .diffusers_convert import convert_lora_state_dict | |
| class EXM_PixArt_ModelPatcher(ModelPatcher): | |
| def calculate_weight(self, patches, weight, key): | |
| """ | |
| This is almost the same as the comfy function, but stripped down to just the LoRA patch code. | |
| The problem with the original code is the q/k/v keys being combined into one for the attention. | |
| In the diffusers code, they're treated as separate keys, but in the reference code they're recombined (q+kv|qkv). | |
| This means, for example, that the [1152,1152] weights become [3456,1152] in the state dict. | |
| The issue with this is that the LoRA weights are [128,1152],[1152,128] and become [384,1162],[3456,128] instead. | |
| This is the best thing I could think of that would fix that, but it's very fragile. | |
| - Check key shape to determine if it needs the fallback logic | |
| - Cut the input into parts based on the shape (undoing the torch.cat) | |
| - Do the matrix multiplication logic | |
| - Recombine them to match the expected shape | |
| """ | |
| for p in patches: | |
| alpha = p[0] | |
| v = p[1] | |
| strength_model = p[2] | |
| if strength_model != 1.0: | |
| weight *= strength_model | |
| if isinstance(v, list): | |
| v = (self.calculate_weight(v[1:], v[0].clone(), key), ) | |
| if len(v) == 2: | |
| patch_type = v[0] | |
| v = v[1] | |
| if patch_type == "lora": | |
| mat1 = comfy.model_management.cast_to_device(v[0], weight.device, torch.float32) | |
| mat2 = comfy.model_management.cast_to_device(v[1], weight.device, torch.float32) | |
| if v[2] is not None: | |
| alpha *= v[2] / mat2.shape[0] | |
| try: | |
| mat1 = mat1.flatten(start_dim=1) | |
| mat2 = mat2.flatten(start_dim=1) | |
| ch1 = mat1.shape[0] // mat2.shape[1] | |
| ch2 = mat2.shape[0] // mat1.shape[1] | |
| ### Fallback logic for shape mismatch ### | |
| if mat1.shape[0] != mat2.shape[1] and ch1 == ch2 and (mat1.shape[0]/mat2.shape[1])%1 == 0: | |
| mat1 = mat1.chunk(ch1, dim=0) | |
| mat2 = mat2.chunk(ch1, dim=0) | |
| weight += torch.cat( | |
| [alpha * torch.mm(mat1[x], mat2[x]) for x in range(ch1)], | |
| dim=0, | |
| ).reshape(weight.shape).type(weight.dtype) | |
| else: | |
| weight += (alpha * torch.mm(mat1, mat2)).reshape(weight.shape).type(weight.dtype) | |
| except Exception as e: | |
| print("ERROR", key, e) | |
| return weight | |
| def clone(self): | |
| n = EXM_PixArt_ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device, weight_inplace_update=self.weight_inplace_update) | |
| n.patches = {} | |
| for k in self.patches: | |
| n.patches[k] = self.patches[k][:] | |
| n.object_patches = self.object_patches.copy() | |
| n.model_options = copy.deepcopy(self.model_options) | |
| n.model_keys = self.model_keys | |
| return n | |
| def replace_model_patcher(model): | |
| n = EXM_PixArt_ModelPatcher( | |
| model = model.model, | |
| size = model.size, | |
| load_device = model.load_device, | |
| offload_device = model.offload_device, | |
| current_device = model.current_device, | |
| weight_inplace_update = model.weight_inplace_update, | |
| ) | |
| n.patches = {} | |
| for k in model.patches: | |
| n.patches[k] = model.patches[k][:] | |
| n.object_patches = model.object_patches.copy() | |
| n.model_options = copy.deepcopy(model.model_options) | |
| return n | |
| def find_peft_alpha(path): | |
| def load_json(json_path): | |
| with open(json_path) as f: | |
| data = json.load(f) | |
| alpha = data.get("lora_alpha") | |
| alpha = alpha or data.get("alpha") | |
| if not alpha: | |
| print(" Found config but `lora_alpha` is missing!") | |
| else: | |
| print(f" Found config at {json_path} [alpha:{alpha}]") | |
| return alpha | |
| # For some weird reason peft doesn't include the alpha in the actual model | |
| print("PixArt: Warning! This is a PEFT LoRA. Trying to find config...") | |
| files = [ | |
| f"{os.path.splitext(path)[0]}.json", | |
| f"{os.path.splitext(path)[0]}.config.json", | |
| os.path.join(os.path.dirname(path),"adapter_config.json"), | |
| ] | |
| for file in files: | |
| if os.path.isfile(file): | |
| return load_json(file) | |
| print(" Missing config/alpha! assuming alpha of 8. Consider converting it/adding a config json to it.") | |
| return 8.0 | |
| def load_pixart_lora(model, lora, lora_path, strength): | |
| k_back = lambda x: x.replace(".lora_up.weight", "") | |
| # need to convert the actual weights for this to work. | |
| if any(True for x in lora.keys() if x.endswith("adaln_single.linear.lora_A.weight")): | |
| lora = convert_lora_state_dict(lora, peft=True) | |
| alpha = find_peft_alpha(lora_path) | |
| lora.update({f"{k_back(x)}.alpha":torch.tensor(alpha) for x in lora.keys() if "lora_up" in x}) | |
| else: # OneTrainer | |
| lora = convert_lora_state_dict(lora, peft=False) | |
| key_map = {k_back(x):f"diffusion_model.{k_back(x)}.weight" for x in lora.keys() if "lora_up" in x} # fake | |
| loaded = comfy.lora.load_lora(lora, key_map) | |
| if model is not None: | |
| # switch to custom model patcher when using LoRAs | |
| if isinstance(model, EXM_PixArt_ModelPatcher): | |
| new_modelpatcher = model.clone() | |
| else: | |
| new_modelpatcher = replace_model_patcher(model) | |
| k = new_modelpatcher.add_patches(loaded, strength) | |
| else: | |
| k = () | |
| new_modelpatcher = None | |
| k = set(k) | |
| for x in loaded: | |
| if (x not in k): | |
| print("NOT LOADED", x) | |
| return new_modelpatcher | |