| from typing import List |
| import torch |
| from safetensors import safe_open |
| from diffusers import StableDiffusionPipeline |
| from .lora import ( |
| monkeypatch_or_replace_safeloras, |
| apply_learned_embed_in_clip, |
| set_lora_diag, |
| parse_safeloras_embeds, |
| ) |
|
|
|
|
| def lora_join(lora_safetenors: list): |
| metadatas = [dict(safelora.metadata()) for safelora in lora_safetenors] |
| _total_metadata = {} |
| total_metadata = {} |
| total_tensor = {} |
| total_rank = 0 |
| ranklist = [] |
| for _metadata in metadatas: |
| rankset = [] |
| for k, v in _metadata.items(): |
| if k.endswith("rank"): |
| rankset.append(int(v)) |
|
|
| assert len(set(rankset)) <= 1, "Rank should be the same per model" |
| if len(rankset) == 0: |
| rankset = [0] |
|
|
| total_rank += rankset[0] |
| _total_metadata.update(_metadata) |
| ranklist.append(rankset[0]) |
|
|
| |
| for k, v in _total_metadata.items(): |
| if v != "<embed>": |
| total_metadata[k] = v |
|
|
| tensorkeys = set() |
| for safelora in lora_safetenors: |
| tensorkeys.update(safelora.keys()) |
|
|
| for keys in tensorkeys: |
| if keys.startswith("text_encoder") or keys.startswith("unet"): |
| tensorset = [safelora.get_tensor(keys) for safelora in lora_safetenors] |
|
|
| is_down = keys.endswith("down") |
|
|
| if is_down: |
| _tensor = torch.cat(tensorset, dim=0) |
| assert _tensor.shape[0] == total_rank |
| else: |
| _tensor = torch.cat(tensorset, dim=1) |
| assert _tensor.shape[1] == total_rank |
|
|
| total_tensor[keys] = _tensor |
| keys_rank = ":".join(keys.split(":")[:-1]) + ":rank" |
| total_metadata[keys_rank] = str(total_rank) |
| token_size_list = [] |
| for idx, safelora in enumerate(lora_safetenors): |
| tokens = [k for k, v in safelora.metadata().items() if v == "<embed>"] |
| for jdx, token in enumerate(sorted(tokens)): |
|
|
| total_tensor[f"<s{idx}-{jdx}>"] = safelora.get_tensor(token) |
| total_metadata[f"<s{idx}-{jdx}>"] = "<embed>" |
|
|
| print(f"Embedding {token} replaced to <s{idx}-{jdx}>") |
|
|
| token_size_list.append(len(tokens)) |
|
|
| return total_tensor, total_metadata, ranklist, token_size_list |
|
|
|
|
| class DummySafeTensorObject: |
| def __init__(self, tensor: dict, metadata): |
| self.tensor = tensor |
| self._metadata = metadata |
|
|
| def keys(self): |
| return self.tensor.keys() |
|
|
| def metadata(self): |
| return self._metadata |
|
|
| def get_tensor(self, key): |
| return self.tensor[key] |
|
|
|
|
| class LoRAManager: |
| def __init__(self, lora_paths_list: List[str], pipe: StableDiffusionPipeline): |
|
|
| self.lora_paths_list = lora_paths_list |
| self.pipe = pipe |
| self._setup() |
|
|
| def _setup(self): |
|
|
| self._lora_safetenors = [ |
| safe_open(path, framework="pt", device="cpu") |
| for path in self.lora_paths_list |
| ] |
|
|
| ( |
| total_tensor, |
| total_metadata, |
| self.ranklist, |
| self.token_size_list, |
| ) = lora_join(self._lora_safetenors) |
|
|
| self.total_safelora = DummySafeTensorObject(total_tensor, total_metadata) |
|
|
| monkeypatch_or_replace_safeloras(self.pipe, self.total_safelora) |
| tok_dict = parse_safeloras_embeds(self.total_safelora) |
|
|
| apply_learned_embed_in_clip( |
| tok_dict, |
| self.pipe.text_encoder, |
| self.pipe.tokenizer, |
| token=None, |
| idempotent=True, |
| ) |
|
|
| def tune(self, scales): |
|
|
| assert len(scales) == len( |
| self.ranklist |
| ), "Scale list should be the same length as ranklist" |
|
|
| diags = [] |
| for scale, rank in zip(scales, self.ranklist): |
| diags = diags + [scale] * rank |
|
|
| set_lora_diag(self.pipe.unet, torch.tensor(diags)) |
|
|
| def prompt(self, prompt): |
| if prompt is not None: |
| for idx, tok_size in enumerate(self.token_size_list): |
| prompt = prompt.replace( |
| f"<{idx + 1}>", |
| "".join([f"<s{idx}-{jdx}>" for jdx in range(tok_size)]), |
| ) |
| |
|
|
| return prompt |
|
|
