upload updated files

lora_diffusion/cli_lora_add.py

@@ -1,35 +1,73 @@
 from typing import Literal, Union, Dict
-
+import os
+import shutil
 import fire
 from diffusers import StableDiffusionPipeline
 
 import torch
 from .lora import tune_lora_scale, weight_apply_lora
+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 = "./merged_lora.pt",
+    output_path: str,
     alpha: float = 0.5,
-    mode: Literal["lpl", "upl"] = "lpl",
+    mode: Literal[
+        "lpl",
+        "upl",
+        "upl-ckpt-v2",
+    ] = "lpl",
+    with_text_lora: bool = False,
 ):
+    print("Lora Add, mode " + mode)
     if mode == "lpl":
-        out_list = []
-        l1 = torch.load(path_1)
-        l2 = torch.load(path_2)
+        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])
+            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 * x1.data + (1 - alpha) * x2.data
-            y1.data = alpha * y1.data + (1 - alpha) * y2.data
+            for (x1, y1), (x2, y2) in zip(l1pairs, l2pairs):
+                # print("Merging", x1.shape, y1.shape, x2.shape, y2.shape)
+                x1.data = alpha * x1.data + (1 - alpha) * x2.data
+                y1.data = alpha * y1.data + (1 - alpha) * y2.data
 
-            out_list.append(x1)
-            out_list.append(y1)
+                out_list.append(x1)
+                out_list.append(y1)
 
-        torch.save(out_list, output_path)
+            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 mode == "upl":
 
@@ -38,12 +76,43 @@ def add(
         ).to("cpu")
 
         weight_apply_lora(loaded_pipeline.unet, torch.load(path_2), alpha=alpha)
+        if with_text_lora:
 
-        if output_path.endswith(".pt"):
-            output_path = output_path[:-3]
+            weight_apply_lora(
+                loaded_pipeline.text_encoder,
+                torch.load(_text_lora_path(path_2)),
+                alpha=alpha,
+                target_replace_module=["CLIPAttention"],
+            )
 
         loaded_pipeline.save_pretrained(output_path)
 
+    elif mode == "upl-ckpt-v2":
+
+        loaded_pipeline = StableDiffusionPipeline.from_pretrained(
+            path_1,
+        ).to("cpu")
+
+        weight_apply_lora(loaded_pipeline.unet, torch.load(path_2), alpha=alpha)
+        if with_text_lora:
+            weight_apply_lora(
+                loaded_pipeline.text_encoder,
+                torch.load(_text_lora_path(path_2)),
+                alpha=alpha,
+                target_replace_module=["CLIPAttention"],
+            )
+
+        _tmp_output = output_path + ".tmp"
+
+        loaded_pipeline.save_pretrained(_tmp_output)
+        convert_to_ckpt(_tmp_output, output_path, as_half=True)
+        # remove the tmp_output folder
+        shutil.rmtree(_tmp_output)
+
+    else:
+        print("Unknown mode", mode)
+        raise ValueError(f"Unknown mode {mode}")
+
 
 def main():
     fire.Fire(add)

lora_diffusion/lora.py

@@ -10,14 +10,20 @@ import torch.nn as nn
 
 
 class LoraInjectedLinear(nn.Module):
-    def __init__(self, in_features, out_features, bias=False):
+    def __init__(self, in_features, out_features, bias=False, r=4):
         super().__init__()
+
+        if r > min(in_features, out_features):
+            raise ValueError(
+                f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
+            )
+
         self.linear = nn.Linear(in_features, out_features, bias)
-        self.lora_down = nn.Linear(in_features, 4, bias=False)
-        self.lora_up = nn.Linear(4, out_features, bias=False)
+        self.lora_down = nn.Linear(in_features, r, bias=False)
+        self.lora_up = nn.Linear(r, out_features, bias=False)
         self.scale = 1.0
 
-        nn.init.normal_(self.lora_down.weight, std=1 / 16)
+        nn.init.normal_(self.lora_down.weight, std=1 / r**2)
         nn.init.zeros_(self.lora_up.weight)
 
     def forward(self, input):
@@ -25,7 +31,10 @@ class LoraInjectedLinear(nn.Module):
 
 
 def inject_trainable_lora(
-    model: nn.Module, target_replace_module: List[str] = ["CrossAttention", "Attention"]
+    model: nn.Module,
+    target_replace_module: List[str] = ["CrossAttention", "Attention"],
+    r: int = 4,
+    loras=None,  # path to lora .pt
 ):
     """
     inject lora into model, and returns lora parameter groups.
@@ -34,6 +43,9 @@ def inject_trainable_lora(
     require_grad_params = []
     names = []
 
+    if loras != None:
+        loras = torch.load(loras)
+
     for _module in model.modules():
         if _module.__class__.__name__ in target_replace_module:
 
@@ -46,6 +58,7 @@ def inject_trainable_lora(
                         _child_module.in_features,
                         _child_module.out_features,
                         _child_module.bias is not None,
+                        r,
                     )
                     _tmp.linear.weight = weight
                     if bias is not None:
@@ -61,10 +74,13 @@ def inject_trainable_lora(
                         _module._modules[name].lora_down.parameters()
                     )
 
+                    if loras != None:
+                        _module._modules[name].lora_up.weight = loras.pop(0)
+                        _module._modules[name].lora_down.weight = loras.pop(0)
+
                     _module._modules[name].lora_up.weight.requires_grad = True
                     _module._modules[name].lora_down.weight.requires_grad = True
                     names.append(name)
-
     return require_grad_params, names
 
 
@@ -82,9 +98,13 @@ def extract_lora_ups_down(model, target_replace_module=["CrossAttention", "Atten
     return loras
 
 
-def save_lora_weight(model, path="./lora.pt"):
+def save_lora_weight(
+    model, path="./lora.pt", target_replace_module=["CrossAttention", "Attention"]
+):
     weights = []
-    for _up, _down in extract_lora_ups_down(model):
+    for _up, _down in extract_lora_ups_down(
+        model, target_replace_module=target_replace_module
+    ):
         weights.append(_up.weight)
         weights.append(_down.weight)
 
@@ -125,7 +145,7 @@ def weight_apply_lora(
 
 
 def monkeypatch_lora(
-    model, loras, target_replace_module=["CrossAttention", "Attention"]
+    model, loras, target_replace_module=["CrossAttention", "Attention"], r: int = 4
 ):
     for _module in model.modules():
         if _module.__class__.__name__ in target_replace_module:
@@ -138,6 +158,44 @@ def monkeypatch_lora(
                         _child_module.in_features,
                         _child_module.out_features,
                         _child_module.bias is not None,
+                        r=r,
+                    )
+                    _tmp.linear.weight = weight
+
+                    if bias is not None:
+                        _tmp.linear.bias = bias
+
+                    # switch the module
+                    _module._modules[name] = _tmp
+
+                    up_weight = loras.pop(0)
+                    down_weight = loras.pop(0)
+
+                    _module._modules[name].lora_up.weight = nn.Parameter(
+                        up_weight.type(weight.dtype)
+                    )
+                    _module._modules[name].lora_down.weight = nn.Parameter(
+                        down_weight.type(weight.dtype)
+                    )
+
+                    _module._modules[name].to(weight.device)
+
+
+def monkeypatch_replace_lora(
+    model, loras, target_replace_module=["CrossAttention", "Attention"], r: int = 4
+):
+    for _module in model.modules():
+        if _module.__class__.__name__ in target_replace_module:
+            for name, _child_module in _module.named_modules():
+                if _child_module.__class__.__name__ == "LoraInjectedLinear":
+
+                    weight = _child_module.linear.weight
+                    bias = _child_module.linear.bias
+                    _tmp = LoraInjectedLinear(
+                        _child_module.linear.in_features,
+                        _child_module.linear.out_features,
+                        _child_module.linear.bias is not None,
+                        r=r,
                     )
                     _tmp.linear.weight = weight
 
@@ -160,7 +218,138 @@ def monkeypatch_lora(
                     _module._modules[name].to(weight.device)
 
 
+def monkeypatch_add_lora(
+    model,
+    loras,
+    target_replace_module=["CrossAttention", "Attention"],
+    alpha: float = 1.0,
+    beta: float = 1.0,
+):
+    for _module in model.modules():
+        if _module.__class__.__name__ in target_replace_module:
+            for name, _child_module in _module.named_modules():
+                if _child_module.__class__.__name__ == "LoraInjectedLinear":
+
+                    weight = _child_module.linear.weight
+
+                    up_weight = loras.pop(0)
+                    down_weight = loras.pop(0)
+
+                    _module._modules[name].lora_up.weight = nn.Parameter(
+                        up_weight.type(weight.dtype).to(weight.device) * alpha
+                        + _module._modules[name].lora_up.weight.to(weight.device) * beta
+                    )
+                    _module._modules[name].lora_down.weight = nn.Parameter(
+                        down_weight.type(weight.dtype).to(weight.device) * alpha
+                        + _module._modules[name].lora_down.weight.to(weight.device)
+                        * beta
+                    )
+
+                    _module._modules[name].to(weight.device)
+
+
 def tune_lora_scale(model, alpha: float = 1.0):
     for _module in model.modules():
         if _module.__class__.__name__ == "LoraInjectedLinear":
             _module.scale = alpha
+
+
+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 _ti_lora_path(path: str) -> str:
+    assert path.endswith(".pt"), "Only .pt files are supported"
+    return ".".join(path.split(".")[:-1] + ["ti", "pt"])
+
+
+def load_learned_embed_in_clip(
+    learned_embeds_path, text_encoder, tokenizer, token=None, idempotent=False
+):
+    loaded_learned_embeds = torch.load(learned_embeds_path, map_location="cpu")
+
+    # separate token and the embeds
+    trained_token = list(loaded_learned_embeds.keys())[0]
+    embeds = loaded_learned_embeds[trained_token]
+
+    # cast to dtype of text_encoder
+    dtype = text_encoder.get_input_embeddings().weight.dtype
+
+    # add the token in tokenizer
+    token = token if token is not None else trained_token
+    num_added_tokens = tokenizer.add_tokens(token)
+    i = 1
+    if num_added_tokens == 0 and idempotent:
+        return token
+
+    while num_added_tokens == 0:
+        print(f"The tokenizer already contains the token {token}.")
+        token = f"{token[:-1]}-{i}>"
+        print(f"Attempting to add the token {token}.")
+        num_added_tokens = tokenizer.add_tokens(token)
+        i += 1
+
+    # resize the token embeddings
+    text_encoder.resize_token_embeddings(len(tokenizer))
+
+    # get the id for the token and assign the embeds
+    token_id = tokenizer.convert_tokens_to_ids(token)
+    text_encoder.get_input_embeddings().weight.data[token_id] = embeds
+    return token
+
+
+def patch_pipe(
+    pipe,
+    unet_path,
+    token,
+    alpha: float = 1.0,
+    r: int = 4,
+    patch_text=False,
+    patch_ti=False,
+    idempotent_token=True,
+):
+
+    ti_path = _ti_lora_path(unet_path)
+    text_path = _text_lora_path(unet_path)
+
+    unet_has_lora = False
+    text_encoder_has_lora = False
+
+    for _module in pipe.unet.modules():
+        if _module.__class__.__name__ == "LoraInjectedLinear":
+            unet_has_lora = True
+
+    for _module in pipe.text_encoder.modules():
+        if _module.__class__.__name__ == "LoraInjectedLinear":
+            text_encoder_has_lora = True
+
+    if not unet_has_lora:
+        monkeypatch_lora(pipe.unet, torch.load(unet_path), r=r)
+    else:
+        monkeypatch_replace_lora(pipe.unet, torch.load(unet_path), r=r)
+
+    if patch_text:
+        if not text_encoder_has_lora:
+            monkeypatch_lora(
+                pipe.text_encoder,
+                torch.load(text_path),
+                target_replace_module=["CLIPAttention"],
+                r=r,
+            )
+        else:
+
+            monkeypatch_replace_lora(
+                pipe.text_encoder,
+                torch.load(text_path),
+                target_replace_module=["CLIPAttention"],
+                r=r,
+            )
+    if patch_ti:
+        token = load_learned_embed_in_clip(
+            ti_path,
+            pipe.text_encoder,
+            pipe.tokenizer,
+            token,
+            idempotent=idempotent_token,
+        )

lora_diffusion/to_ckpt_v2.py

@@ -0,0 +1,232 @@
+# from https://gist.github.com/jachiam/8a5c0b607e38fcc585168b90c686eb05
+# Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint.
+# *Only* converts the UNet, VAE, and Text Encoder.
+# Does not convert optimizer state or any other thing.
+# Written by jachiam
+import argparse
+import os.path as osp
+
+import torch
+
+
+# =================#
+# UNet Conversion #
+# =================#
+
+unet_conversion_map = [
+    # (stable-diffusion, HF Diffusers)
+    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
+    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
+    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
+    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
+    ("input_blocks.0.0.weight", "conv_in.weight"),
+    ("input_blocks.0.0.bias", "conv_in.bias"),
+    ("out.0.weight", "conv_norm_out.weight"),
+    ("out.0.bias", "conv_norm_out.bias"),
+    ("out.2.weight", "conv_out.weight"),
+    ("out.2.bias", "conv_out.bias"),
+]
+
+unet_conversion_map_resnet = [
+    # (stable-diffusion, HF Diffusers)
+    ("in_layers.0", "norm1"),
+    ("in_layers.2", "conv1"),
+    ("out_layers.0", "norm2"),
+    ("out_layers.3", "conv2"),
+    ("emb_layers.1", "time_emb_proj"),
+    ("skip_connection", "conv_shortcut"),
+]
+
+unet_conversion_map_layer = []
+# hardcoded number of downblocks and resnets/attentions...
+# would need smarter logic for other networks.
+for i in range(4):
+    # loop over downblocks/upblocks
+
+    for j in range(2):
+        # loop over resnets/attentions for downblocks
+        hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+        sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+        unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+        if i < 3:
+            # no attention layers in down_blocks.3
+            hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+            sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+            unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+    for j in range(3):
+        # loop over resnets/attentions for upblocks
+        hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+        sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+        unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+        if i > 0:
+            # no attention layers in up_blocks.0
+            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+    if i < 3:
+        # no downsample in down_blocks.3
+        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+        sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
+        unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+        # no upsample in up_blocks.3
+        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+        sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
+        unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+hf_mid_atn_prefix = "mid_block.attentions.0."
+sd_mid_atn_prefix = "middle_block.1."
+unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+for j in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{j}."
+    sd_mid_res_prefix = f"middle_block.{2*j}."
+    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+
+def convert_unet_state_dict(unet_state_dict):
+    # buyer beware: this is a *brittle* function,
+    # and correct output requires that all of these pieces interact in
+    # the exact order in which I have arranged them.
+    mapping = {k: k for k in unet_state_dict.keys()}
+    for sd_name, hf_name in unet_conversion_map:
+        mapping[hf_name] = sd_name
+    for k, v in mapping.items():
+        if "resnets" in k:
+            for sd_part, hf_part in unet_conversion_map_resnet:
+                v = v.replace(hf_part, sd_part)
+            mapping[k] = v
+    for k, v in mapping.items():
+        for sd_part, hf_part in unet_conversion_map_layer:
+            v = v.replace(hf_part, sd_part)
+        mapping[k] = v
+    new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
+    return new_state_dict
+
+
+# ================#
+# VAE Conversion #
+# ================#
+
+vae_conversion_map = [
+    # (stable-diffusion, HF Diffusers)
+    ("nin_shortcut", "conv_shortcut"),
+    ("norm_out", "conv_norm_out"),
+    ("mid.attn_1.", "mid_block.attentions.0."),
+]
+
+for i in range(4):
+    # down_blocks have two resnets
+    for j in range(2):
+        hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
+        sd_down_prefix = f"encoder.down.{i}.block.{j}."
+        vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
+
+    if i < 3:
+        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
+        sd_downsample_prefix = f"down.{i}.downsample."
+        vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
+
+        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+        sd_upsample_prefix = f"up.{3-i}.upsample."
+        vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    # up_blocks have three resnets
+    # also, up blocks in hf are numbered in reverse from sd
+    for j in range(3):
+        hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
+        sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
+        vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
+
+# this part accounts for mid blocks in both the encoder and the decoder
+for i in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{i}."
+    sd_mid_res_prefix = f"mid.block_{i+1}."
+    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+
+vae_conversion_map_attn = [
+    # (stable-diffusion, HF Diffusers)
+    ("norm.", "group_norm."),
+    ("q.", "query."),
+    ("k.", "key."),
+    ("v.", "value."),
+    ("proj_out.", "proj_attn."),
+]
+
+
+def reshape_weight_for_sd(w):
+    # convert HF linear weights to SD conv2d weights
+    return w.reshape(*w.shape, 1, 1)
+
+
+def convert_vae_state_dict(vae_state_dict):
+    mapping = {k: k for k in vae_state_dict.keys()}
+    for k, v in mapping.items():
+        for sd_part, hf_part in vae_conversion_map:
+            v = v.replace(hf_part, sd_part)
+        mapping[k] = v
+    for k, v in mapping.items():
+        if "attentions" in k:
+            for sd_part, hf_part in vae_conversion_map_attn:
+                v = v.replace(hf_part, sd_part)
+            mapping[k] = v
+    new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
+    weights_to_convert = ["q", "k", "v", "proj_out"]
+    for k, v in new_state_dict.items():
+        for weight_name in weights_to_convert:
+            if f"mid.attn_1.{weight_name}.weight" in k:
+                print(f"Reshaping {k} for SD format")
+                new_state_dict[k] = reshape_weight_for_sd(v)
+    return new_state_dict
+
+
+# =========================#
+# Text Encoder Conversion #
+# =========================#
+# pretty much a no-op
+
+
+def convert_text_enc_state_dict(text_enc_dict):
+    return text_enc_dict
+
+
+def convert_to_ckpt(model_path, checkpoint_path, as_half):
+
+    assert model_path is not None, "Must provide a model path!"
+
+    assert checkpoint_path is not None, "Must provide a checkpoint path!"
+
+    unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.bin")
+    vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.bin")
+    text_enc_path = osp.join(model_path, "text_encoder", "pytorch_model.bin")
+
+    # Convert the UNet model
+    unet_state_dict = torch.load(unet_path, map_location="cpu")
+    unet_state_dict = convert_unet_state_dict(unet_state_dict)
+    unet_state_dict = {
+        "model.diffusion_model." + k: v for k, v in unet_state_dict.items()
+    }
+
+    # Convert the VAE model
+    vae_state_dict = torch.load(vae_path, map_location="cpu")
+    vae_state_dict = convert_vae_state_dict(vae_state_dict)
+    vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()}
+
+    # Convert the text encoder model
+    text_enc_dict = torch.load(text_enc_path, map_location="cpu")
+    text_enc_dict = convert_text_enc_state_dict(text_enc_dict)
+    text_enc_dict = {
+        "cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()
+    }
+
+    # Put together new checkpoint
+    state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict}
+    if as_half:
+        state_dict = {k: v.half() for k, v in state_dict.items()}
+    state_dict = {"state_dict": state_dict}
+    torch.save(state_dict, checkpoint_path)