Unmerge state_dict in LoRANetwork from the pipeline in diffusers.
(pipeline, lora_path, multiplier=1, device="cpu", dtype=torch.float32, sub_transformer_name="transformer")
| 504 | |
| 505 | # TODO: Refactor with merge_lora. |
| 506 | def unmerge_lora(pipeline, lora_path, multiplier=1, device="cpu", dtype=torch.float32, sub_transformer_name="transformer"): |
| 507 | """Unmerge state_dict in LoRANetwork from the pipeline in diffusers.""" |
| 508 | LORA_PREFIX_UNET = "lora_unet" |
| 509 | LORA_PREFIX_TEXT_ENCODER = "lora_te" |
| 510 | state_dict = load_file(lora_path) |
| 511 | |
| 512 | updates = defaultdict(dict) |
| 513 | for key, value in state_dict.items(): |
| 514 | if "diffusion_model" in key: |
| 515 | key = key.replace("diffusion_model.", "lora_unet__") |
| 516 | key = key.replace("blocks.", "blocks_") |
| 517 | key = key.replace(".self_attn.", "_self_attn_") |
| 518 | key = key.replace(".cross_attn.", "_cross_attn_") |
| 519 | key = key.replace(".ffn.", "_ffn_") |
| 520 | if "lora_A" in key or "lora_B" in key: |
| 521 | key = "lora_unet__" + key |
| 522 | key = key.replace("blocks.", "blocks_") |
| 523 | key = key.replace(".self_attn.", "_self_attn_") |
| 524 | key = key.replace(".cross_attn.", "_cross_attn_") |
| 525 | key = key.replace(".ffn.", "_ffn_") |
| 526 | key = key.replace(".lora_A.default.", ".lora_down.") |
| 527 | key = key.replace(".lora_B.default.", ".lora_up.") |
| 528 | layer, elem = key.split('.', 1) |
| 529 | updates[layer][elem] = value |
| 530 | |
| 531 | sequential_cpu_offload_flag = False |
| 532 | if pipeline.transformer.device == torch.device(type="meta"): |
| 533 | pipeline.remove_all_hooks() |
| 534 | sequential_cpu_offload_flag = True |
| 535 | |
| 536 | for layer, elems in updates.items(): |
| 537 | |
| 538 | if "lora_te" in layer: |
| 539 | layer_infos = layer.split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_") |
| 540 | curr_layer = pipeline.text_encoder |
| 541 | else: |
| 542 | layer_infos = layer.split(LORA_PREFIX_UNET + "_")[-1].split("_") |
| 543 | curr_layer = getattr(pipeline, sub_transformer_name) |
| 544 | |
| 545 | try: |
| 546 | curr_layer = curr_layer.__getattr__("_".join(layer_infos[1:])) |
| 547 | except Exception: |
| 548 | temp_name = layer_infos.pop(0) |
| 549 | try: |
| 550 | while len(layer_infos) > -1: |
| 551 | try: |
| 552 | curr_layer = curr_layer.__getattr__(temp_name + "_" + "_".join(layer_infos)) |
| 553 | break |
| 554 | except Exception: |
| 555 | try: |
| 556 | curr_layer = curr_layer.__getattr__(temp_name) |
| 557 | if len(layer_infos) > 0: |
| 558 | temp_name = layer_infos.pop(0) |
| 559 | elif len(layer_infos) == 0: |
| 560 | break |
| 561 | except Exception: |
| 562 | if len(layer_infos) == 0: |
| 563 | print(f'Error loading layer in front search: {layer}. Try it in back search.') |
no test coverage detected