(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT_ENCODER, alpha)
| 24 | |
| 25 | |
| 26 | def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT_ENCODER, alpha): |
| 27 | # load base model |
| 28 | pipeline = StableDiffusionPipeline.from_pretrained(base_model_path, torch_dtype=torch.float32) |
| 29 | |
| 30 | # load LoRA weight from .safetensors |
| 31 | state_dict = load_file(checkpoint_path) |
| 32 | |
| 33 | visited = [] |
| 34 | |
| 35 | # directly update weight in diffusers model |
| 36 | for key in state_dict: |
| 37 | # it is suggested to print out the key, it usually will be something like below |
| 38 | # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight" |
| 39 | |
| 40 | # as we have set the alpha beforehand, so just skip |
| 41 | if ".alpha" in key or key in visited: |
| 42 | continue |
| 43 | |
| 44 | if "text" in key: |
| 45 | layer_infos = key.split(".")[0].split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_") |
| 46 | curr_layer = pipeline.text_encoder |
| 47 | else: |
| 48 | layer_infos = key.split(".")[0].split(LORA_PREFIX_UNET + "_")[-1].split("_") |
| 49 | curr_layer = pipeline.unet |
| 50 | |
| 51 | # find the target layer |
| 52 | temp_name = layer_infos.pop(0) |
| 53 | while len(layer_infos) > -1: |
| 54 | try: |
| 55 | curr_layer = curr_layer.__getattr__(temp_name) |
| 56 | if len(layer_infos) > 0: |
| 57 | temp_name = layer_infos.pop(0) |
| 58 | elif len(layer_infos) == 0: |
| 59 | break |
| 60 | except Exception: |
| 61 | if len(temp_name) > 0: |
| 62 | temp_name += "_" + layer_infos.pop(0) |
| 63 | else: |
| 64 | temp_name = layer_infos.pop(0) |
| 65 | |
| 66 | pair_keys = [] |
| 67 | if "lora_down" in key: |
| 68 | pair_keys.append(key.replace("lora_down", "lora_up")) |
| 69 | pair_keys.append(key) |
| 70 | else: |
| 71 | pair_keys.append(key) |
| 72 | pair_keys.append(key.replace("lora_up", "lora_down")) |
| 73 | |
| 74 | # update weight |
| 75 | if len(state_dict[pair_keys[0]].shape) == 4: |
| 76 | weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32) |
| 77 | weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32) |
| 78 | curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3) |
| 79 | else: |
| 80 | weight_up = state_dict[pair_keys[0]].to(torch.float32) |
| 81 | weight_down = state_dict[pair_keys[1]].to(torch.float32) |
| 82 | curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down) |
| 83 |
no test coverage detected
searching dependent graphs…