| from typing import Literal, Union, Dict |
| import os |
| import shutil |
| import fire |
| from diffusers import StableDiffusionPipeline |
| from safetensors.torch import safe_open, save_file |
|
|
| import torch |
| from .lora import ( |
| tune_lora_scale, |
| patch_pipe, |
| collapse_lora, |
| monkeypatch_remove_lora, |
| ) |
| from .lora_manager import lora_join |
| from .to_ckpt_v2 import convert_to_ckpt |
|
|
|
|
| def _text_lora_path(path: str) -> str: |
| assert path.endswith(".pt"), "Only .pt files are supported" |
| return ".".join(path.split(".")[:-1] + ["text_encoder", "pt"]) |
|
|
|
|
| def add( |
| path_1: str, |
| path_2: str, |
| output_path: str, |
| alpha_1: float = 0.5, |
| alpha_2: float = 0.5, |
| mode: Literal[ |
| "lpl", |
| "upl", |
| "upl-ckpt-v2", |
| ] = "lpl", |
| with_text_lora: bool = False, |
| ): |
| print("Lora Add, mode " + mode) |
| if mode == "lpl": |
| if path_1.endswith(".pt") and path_2.endswith(".pt"): |
| for _path_1, _path_2, opt in [(path_1, path_2, "unet")] + ( |
| [(_text_lora_path(path_1), _text_lora_path(path_2), "text_encoder")] |
| if with_text_lora |
| else [] |
| ): |
| print("Loading", _path_1, _path_2) |
| out_list = [] |
| if opt == "text_encoder": |
| if not os.path.exists(_path_1): |
| print(f"No text encoder found in {_path_1}, skipping...") |
| continue |
| if not os.path.exists(_path_2): |
| print(f"No text encoder found in {_path_1}, skipping...") |
| continue |
|
|
| l1 = torch.load(_path_1) |
| l2 = torch.load(_path_2) |
|
|
| l1pairs = zip(l1[::2], l1[1::2]) |
| l2pairs = zip(l2[::2], l2[1::2]) |
|
|
| for (x1, y1), (x2, y2) in zip(l1pairs, l2pairs): |
| |
| x1.data = alpha_1 * x1.data + alpha_2 * x2.data |
| y1.data = alpha_1 * y1.data + alpha_2 * y2.data |
|
|
| out_list.append(x1) |
| out_list.append(y1) |
|
|
| if opt == "unet": |
|
|
| print("Saving merged UNET to", output_path) |
| torch.save(out_list, output_path) |
|
|
| elif opt == "text_encoder": |
| print("Saving merged text encoder to", _text_lora_path(output_path)) |
| torch.save( |
| out_list, |
| _text_lora_path(output_path), |
| ) |
|
|
| elif path_1.endswith(".safetensors") and path_2.endswith(".safetensors"): |
| safeloras_1 = safe_open(path_1, framework="pt", device="cpu") |
| safeloras_2 = safe_open(path_2, framework="pt", device="cpu") |
|
|
| metadata = dict(safeloras_1.metadata()) |
| metadata.update(dict(safeloras_2.metadata())) |
|
|
| ret_tensor = {} |
|
|
| for keys in set(list(safeloras_1.keys()) + list(safeloras_2.keys())): |
| if keys.startswith("text_encoder") or keys.startswith("unet"): |
|
|
| tens1 = safeloras_1.get_tensor(keys) |
| tens2 = safeloras_2.get_tensor(keys) |
|
|
| tens = alpha_1 * tens1 + alpha_2 * tens2 |
| ret_tensor[keys] = tens |
| else: |
| if keys in safeloras_1.keys(): |
|
|
| tens1 = safeloras_1.get_tensor(keys) |
| else: |
| tens1 = safeloras_2.get_tensor(keys) |
|
|
| ret_tensor[keys] = tens1 |
|
|
| save_file(ret_tensor, output_path, metadata) |
|
|
| elif mode == "upl": |
|
|
| print( |
| f"Merging UNET/CLIP from {path_1} with LoRA from {path_2} to {output_path}. Merging ratio : {alpha_1}." |
| ) |
|
|
| loaded_pipeline = StableDiffusionPipeline.from_pretrained( |
| path_1, |
| ).to("cpu") |
|
|
| patch_pipe(loaded_pipeline, path_2) |
|
|
| collapse_lora(loaded_pipeline.unet, alpha_1) |
| collapse_lora(loaded_pipeline.text_encoder, alpha_1) |
|
|
| monkeypatch_remove_lora(loaded_pipeline.unet) |
| monkeypatch_remove_lora(loaded_pipeline.text_encoder) |
|
|
| loaded_pipeline.save_pretrained(output_path) |
|
|
| elif mode == "upl-ckpt-v2": |
|
|
| assert output_path.endswith(".ckpt"), "Only .ckpt files are supported" |
| name = os.path.basename(output_path)[0:-5] |
|
|
| print( |
| f"You will be using {name} as the token in A1111 webui. Make sure {name} is unique enough token." |
| ) |
|
|
| loaded_pipeline = StableDiffusionPipeline.from_pretrained( |
| path_1, |
| ).to("cpu") |
|
|
| tok_dict = patch_pipe(loaded_pipeline, path_2, patch_ti=False) |
|
|
| collapse_lora(loaded_pipeline.unet, alpha_1) |
| collapse_lora(loaded_pipeline.text_encoder, alpha_1) |
|
|
| monkeypatch_remove_lora(loaded_pipeline.unet) |
| monkeypatch_remove_lora(loaded_pipeline.text_encoder) |
|
|
| _tmp_output = output_path + ".tmp" |
|
|
| loaded_pipeline.save_pretrained(_tmp_output) |
| convert_to_ckpt(_tmp_output, output_path, as_half=True) |
| |
| shutil.rmtree(_tmp_output) |
|
|
| keys = sorted(tok_dict.keys()) |
| tok_catted = torch.stack([tok_dict[k] for k in keys]) |
| ret = { |
| "string_to_token": {"*": torch.tensor(265)}, |
| "string_to_param": {"*": tok_catted}, |
| "name": name, |
| } |
|
|
| torch.save(ret, output_path[:-5] + ".pt") |
| print( |
| f"Textual embedding saved as {output_path[:-5]}.pt, put it in the embedding folder and use it as {name} in A1111 repo, " |
| ) |
| elif mode == "ljl": |
| print("Using Join mode : alpha will not have an effect here.") |
| assert path_1.endswith(".safetensors") and path_2.endswith( |
| ".safetensors" |
| ), "Only .safetensors files are supported" |
|
|
| safeloras_1 = safe_open(path_1, framework="pt", device="cpu") |
| safeloras_2 = safe_open(path_2, framework="pt", device="cpu") |
|
|
| total_tensor, total_metadata, _, _ = lora_join([safeloras_1, safeloras_2]) |
| save_file(total_tensor, output_path, total_metadata) |
|
|
| else: |
| print("Unknown mode", mode) |
| raise ValueError(f"Unknown mode {mode}") |
|
|
|
|
| def main(): |
| fire.Fire(add) |
|
|
