feat: requirements, cleanup

Dockerfile

@@ -0,0 +1,22 @@
+# Dockerfile Public T4
+
+FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04
+ENV DEBIAN_FRONTEND noninteractive
+
+WORKDIR /content
+
+RUN apt-get update -y && apt-get upgrade -y && apt-get install -y libgl1 libglib2.0-0 wget git git-lfs python3-pip python-is-python3 && pip3 install --upgrade pip
+
+RUN pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchsde --extra-index-url https://download.pytorch.org/whl/cu113
+RUN pip install https://github.com/camenduru/stable-diffusion-webui-colab/releases/download/0.0.16/xformers-0.0.16+814314d.d20230118-cp310-cp310-linux_x86_64.whl
+RUN pip install --pre triton
+RUN pip install numexpr einops diffusers transformers k_diffusion safetensors gradio
+
+ADD . .
+RUN adduser --disabled-password --gecos '' user
+RUN chown -R user:user /content
+RUN chmod -R 777 /content
+USER user
+
+EXPOSE 7860
+CMD python /content/app.py

README.md

@@ -1,13 +1,13 @@
 ---
 title: Sd Diffusers Webui
-emoji: 👀
-colorFrom: red
-colorTo: blue
-sdk: gradio
-sdk_version: 3.16.2
-app_file: app.py
+emoji: 🐳
+colorFrom: purple
+colorTo: gray
+sdk: docker
+sdk_version: 3.9
 pinned: false
 license: openrail
+app_port: 7860
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -4,6 +4,7 @@ import time
 import gradio as gr
 import numpy as np
 import torch
+import math
 
 from gradio import inputs
 from diffusers import (
@@ -14,7 +15,6 @@ from diffusers import (
 from modules.model import (
     CrossAttnProcessor,
     StableDiffusionPipeline,
-    load_lora_attn_procs,
 )
 from torchvision import transforms
 from transformers import CLIPTokenizer, CLIPTextModel
@@ -22,16 +22,17 @@ from PIL import Image
 from pathlib import Path
 from safetensors.torch import load_file
 import modules.safe as _
+from modules.lora import LoRANetwork
 
 models = [
-    ("AbyssOrangeMix2", "Korakoe/AbyssOrangeMix2-HF"),
-    ("Basil Mix", "nuigurumi/basil_mix"),
-    ("Pastal Mix", "andite/pastel-mix"),
-    ("ACertainModel", "JosephusCheung/ACertainModel"),
+    # format: name, model_path, clip_skip
+    ("AbyssOrangeMix2", "Korakoe/AbyssOrangeMix2-HF", 2),
+    ("Basil Mix", "nuigurumi/basil_mix", 2),
+    ("Pastal Mix", "andite/pastel-mix", 2),
+    ("ACertainModel", "JosephusCheung/ACertainModel", 2),
 ]
 
-base_name, base_model = models[0]
-clip_skip = 2
+base_name, base_model, clip_skip = models[0]
 
 samplers_k_diffusion = [
     ("Euler a", "sample_euler_ancestral", {}),
@@ -103,6 +104,10 @@ unet_cache = {
     base_name: unet
 }
 
+lora_cache = {
+    base_name: LoRANetwork(text_encoder, unet)
+}
+
 def get_model(name):
     keys = [k[0] for k in models]
     if name not in unet_cache:
@@ -114,11 +119,21 @@ def get_model(name):
                 subfolder="unet",
                 torch_dtype=torch.float16,
             )
+            if torch.cuda.is_available():
+                unet.to("cuda")
+
             unet_cache[name] = unet
+            lora_cache[name] = LoRANetwork(lora_cache[base_name].text_encoder_loras, unet)
 
     g_unet = unet_cache[name]
-    g_unet.set_attn_processor(None)
-    return g_unet
+    g_lora = lora_cache[name]
+    g_unet.set_attn_processor(CrossAttnProcessor())
+    g_lora.reset()
+    return g_unet, g_lora
+
+# precache on huggingface
+# for model in get_model_list():
+#     get_model(model[0])
 
 def error_str(error, title="Error"):
     return (
@@ -129,18 +144,46 @@ def error_str(error, title="Error"):
     )
 
 
-te_base_weight = text_encoder.get_input_embeddings().weight.data.detach().clone()
+te_base_weight_length = text_encoder.get_input_embeddings().weight.data.shape[0]
+original_prepare_for_tokenization = tokenizer.prepare_for_tokenization
+
+def make_token_names(embs):
+    all_tokens = []
+    for name, vec in embs.items():
+        tokens = [f'emb-{name}-{i}' for i in range(len(vec))]
+        all_tokens.append(tokens)
+    return all_tokens
+
+def setup_tokenizer(embs):
+    reg_match = [re.compile(fr"(?:^|(?<=\s|,)){k}(?=,|\s|$)") for k in embs.keys()]
+    clip_keywords = [' '.join(s) for s in make_token_names(embs)]
+
+    def parse_prompt(prompt: str):
+        for m, v in zip(reg_match, clip_keywords):
+            prompt = m.sub(v, prompt)
+        return prompt
 
 
 def restore_all():
     global te_base_weight, tokenizer
-    text_encoder.get_input_embeddings().weight.data = te_base_weight
+    tokenizer.prepare_for_tokenization = original_prepare_for_tokenization
+
+    embeddings = text_encoder.get_input_embeddings()
+    text_encoder.get_input_embeddings().weight.data = embeddings.weight.data[:te_base_weight_length]
     tokenizer = CLIPTokenizer.from_pretrained(
         base_model,
         subfolder="tokenizer",
         torch_dtype=torch.float16,
     )
 
+def convert_size(size_bytes):
+    if size_bytes == 0:
+        return "0B"
+    size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
+    i = int(math.floor(math.log(size_bytes, 1024)))
+    p = math.pow(1024, i)
+    s = round(size_bytes / p, 2)
+    return "%s %s" % (s, size_name[i])
 
 def inference(
     prompt,
@@ -167,13 +210,15 @@ def inference(
     global pipe, unet, tokenizer, text_encoder
     if seed is None or seed == 0:
         seed = random.randint(0, 2147483647)
+
+    start_time = time.time()
+    restore_all()
     generator = torch.Generator("cuda").manual_seed(int(seed))
 
-    local_unet = get_model(model)
+    local_unet, local_lora = get_model(model)
     if lora_state is not None and lora_state != "":
-        load_lora_attn_procs(lora_state, local_unet, lora_scale)
-    else:
-        local_unet.set_attn_processor(CrossAttnProcessor())
+        local_lora.load(lora_state, lora_scale)
+        local_lora.to(local_unet.device, dtype=local_unet.dtype)
 
     pipe.setup_unet(local_unet)
     sampler_name, sampler_opt = None, None
@@ -182,23 +227,23 @@ def inference(
             sampler_name, sampler_opt = funcname, options
 
     if embs is not None and len(embs) > 0:
-        delta_weight = []
+        ti_embs = {}
         for name, file in embs.items():
             if str(file).endswith(".pt"):
                 loaded_learned_embeds = torch.load(file, map_location="cpu")
             else:
                 loaded_learned_embeds = load_file(file, device="cpu")
             loaded_learned_embeds = loaded_learned_embeds["string_to_param"]["*"]
-            added_length = tokenizer.add_tokens(name)
+            ti_embs[name] = loaded_learned_embeds
 
-            assert added_length == loaded_learned_embeds.shape[0]
-            delta_weight.append(loaded_learned_embeds)
+        if len(ti_embs) > 0:
+            tokens = setup_tokenizer(ti_embs)
+            added_tokens = tokenizer.add_tokens(tokens)
+            delta_weight = torch.cat([val for val in ti_embs.values()], dim=0)
 
-        delta_weight = torch.cat(delta_weight, dim=0)
-        text_encoder.resize_token_embeddings(len(tokenizer))
-        text_encoder.get_input_embeddings().weight.data[
-            -delta_weight.shape[0] :
-        ] = delta_weight
+            assert added_tokens == delta_weight.shape[0]
+            text_encoder.resize_token_embeddings(len(tokenizer))
+            text_encoder.get_input_embeddings().weight.data[-delta_weight.shape[0]:] = delta_weight
 
     config = {
         "negative_prompt": neg_prompt,
@@ -234,6 +279,10 @@ def inference(
     # restore
     if embs is not None and len(embs) > 0:
         restore_all()
+
+    end_time = time.time()
+    vram_free, vram_total = torch.cuda.mem_get_info()
+    print(f"done: res={width}x{height}, step={steps}, time={round(end_time-start_time, 2)}s, vram_alloc={convert_size(vram_total-vram_free)}/{convert_size(vram_total)}")
     return gr.Image.update(result[0][0], label=f"Initial Seed: {seed}")
 
 
@@ -513,7 +562,7 @@ with gr.Blocks(css=css) as demo:
                             label="Guidance scale", value=7.5, maximum=15
                         )
                         steps = gr.Slider(
-                            label="Steps", value=25, minimum=2, maximum=75, step=1
+                            label="Steps", value=25, minimum=2, maximum=50, step=1
                         )
 
                     with gr.Row():
@@ -704,7 +753,7 @@ with gr.Blocks(css=css) as demo:
                     step=0.01,
                     value=0.5,
                 )
-                
+
 
                 sk_update.click(
                     detect_text,
@@ -739,7 +788,7 @@ with gr.Blocks(css=css) as demo:
                     source="upload",
                     shape=(512, 512),
                 )
-                
+
                 mask_outsides2 = gr.Checkbox(
                     label="Mask other areas", 
                     value=False
@@ -803,4 +852,4 @@ with gr.Blocks(css=css) as demo:
 
 print(f"Space built in {time.time() - start_time:.2f} seconds")
 # demo.launch(share=True)
-demo.launch()
+demo.launch(enable_queue=True, server_name="0.0.0.0", server_port=7860)

modules/lora.py

@@ -0,0 +1,181 @@
+# LoRA network module
+# reference:
+# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
+# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
+# https://github.com/bmaltais/kohya_ss/blob/master/networks/lora.py#L48
+
+import math
+import os
+import torch
+import modules.safe as _
+from safetensors.torch import load_file
+
+
+class LoRAModule(torch.nn.Module):
+    """
+    replaces forward method of the original Linear, instead of replacing the original Linear module.
+    """
+
+    def __init__(
+            self,
+            lora_name,
+            org_module: torch.nn.Module,
+            multiplier=1.0,
+            lora_dim=4,
+            alpha=1,
+    ):
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
+        super().__init__()
+        self.lora_name = lora_name
+        self.lora_dim = lora_dim
+
+        if org_module.__class__.__name__ == "Conv2d":
+            in_dim = org_module.in_channels
+            out_dim = org_module.out_channels
+            self.lora_down = torch.nn.Conv2d(in_dim, lora_dim, (1, 1), bias=False)
+            self.lora_up = torch.nn.Conv2d(lora_dim, out_dim, (1, 1), bias=False)
+        else:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            self.lora_down = torch.nn.Linear(in_dim, lora_dim, bias=False)
+            self.lora_up = torch.nn.Linear(lora_dim, out_dim, bias=False)
+
+        if type(alpha) == torch.Tensor:
+            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+
+        alpha = lora_dim if alpha is None or alpha == 0 else alpha
+        self.scale = alpha / self.lora_dim
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
+
+        # same as microsoft's
+        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.lora_up.weight)
+
+        self.multiplier = multiplier
+        self.org_module = org_module  # remove in applying
+        self.enable = False
+
+    def resize(self, rank, alpha):
+        self.alpha = torch.tensor(alpha)
+        self.scale = alpha / rank
+        if self.lora_down.__class__.__name__ == "Conv2d":
+            in_dim = self.lora_down.in_channels
+            out_dim = self.lora_up.out_channels
+            self.lora_down = torch.nn.Conv2d(in_dim, rank, (1, 1), bias=False)
+            self.lora_up = torch.nn.Conv2d(rank, out_dim, (1, 1), bias=False)
+        else:
+            in_dim = self.lora_down.in_features
+            out_dim = self.lora_up.out_features
+            self.lora_down = torch.nn.Linear(in_dim, rank, bias=False)
+            self.lora_up = torch.nn.Linear(rank, out_dim, bias=False)
+
+    def apply(self):
+        if hasattr(self, "org_module"):
+            self.org_forward = self.org_module.forward
+            self.org_module.forward = self.forward
+            del self.org_module
+
+    def forward(self, x):
+        if self.enable:
+            return (
+        self.org_forward(x)
+        + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+        )
+        return self.org_forward(x)
+
+
+class LoRANetwork(torch.nn.Module):
+    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
+    LORA_PREFIX_UNET = "lora_unet"
+    LORA_PREFIX_TEXT_ENCODER = "lora_te"
+
+    def __init__(self, text_encoder, unet, multiplier=1.0, lora_dim=4, alpha=1) -> None:
+        super().__init__()
+        self.multiplier = multiplier
+        self.lora_dim = lora_dim
+        self.alpha = alpha
+
+        # create module instances
+        def create_modules(prefix, root_module: torch.nn.Module, target_replace_modules):
+            loras = []
+            for name, module in root_module.named_modules():
+                if module.__class__.__name__ in target_replace_modules:
+                    for child_name, child_module in module.named_modules():
+                        if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)):
+                            lora_name = prefix + "." + name + "." + child_name
+                            lora_name = lora_name.replace(".", "_")
+                            lora = LoRAModule(lora_name, child_module, self.multiplier, self.lora_dim, self.alpha,)
+                            loras.append(lora)
+            return loras
+
+        if isinstance(text_encoder, list):
+            self.text_encoder_loras = text_encoder
+        else:
+            self.text_encoder_loras = create_modules(LoRANetwork.LORA_PREFIX_TEXT_ENCODER, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
+            print(f"Create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+
+        self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
+        print(f"Create LoRA for U-Net: {len(self.unet_loras)} modules.")
+
+        self.weights_sd = None
+
+        # assertion
+        names = set()
+        for lora in self.text_encoder_loras + self.unet_loras:
+            assert (lora.lora_name not in names), f"duplicated lora name: {lora.lora_name}"
+            names.add(lora.lora_name)
+
+            lora.apply()
+            self.add_module(lora.lora_name, lora)
+
+    def reset(self):
+        for lora in self.text_encoder_loras + self.unet_loras:
+            lora.enable = False
+
+    def load(self, file, scale):
+
+        weights = None
+        if os.path.splitext(file)[1] == ".safetensors":
+            weights = load_file(file)
+        else:
+            weights = torch.load(file, map_location="cpu")
+
+        if not weights:
+            return
+
+        network_alpha = None
+        network_dim = None
+        for key, value in weights.items():
+            if network_alpha is None and "alpha" in key:
+                network_alpha = value
+            if network_dim is None and "lora_down" in key and len(value.size()) == 2:
+                network_dim = value.size()[0]
+
+        if network_alpha is None:
+            network_alpha = network_dim
+
+        weights_has_text_encoder = weights_has_unet = False
+        weights_to_modify = []
+
+        for key in weights.keys():
+            if key.startswith(LoRANetwork.LORA_PREFIX_TEXT_ENCODER):
+                weights_has_text_encoder = True
+
+            if key.startswith(LoRANetwork.LORA_PREFIX_UNET):
+                weights_has_unet = True
+
+        if weights_has_text_encoder:
+            weights_to_modify += self.text_encoder_loras
+
+        if weights_has_unet:
+            weights_to_modify += self.unet_loras
+
+        for lora in self.text_encoder_loras + self.unet_loras:
+            lora.resize(network_dim, network_alpha)
+            if lora in weights_to_modify:
+                lora.enable = True
+
+        info = self.load_state_dict(weights, False)
+        print(f"Weights are loaded. Unexpect keys={info.unexpected_keys}")
+        

modules/model.py

@@ -68,79 +68,6 @@ def get_attention_scores(attn, query, key, attention_mask=None):
     return attention_scores
 
 
-def load_lora_attn_procs(model_file, unet, scale=1.0):
-
-    if Path(model_file).suffix == ".pt":
-        state_dict = torch.load(model_file, map_location="cpu")
-    else:
-        state_dict = load_file(model_file, device="cpu")
-
-    if any("lora_unet_down_blocks" in k for k in state_dict.keys()):
-        # convert ldm format lora
-        df_lora = {}
-        attn_numlayer = re.compile(r"_attn(\d)_to_([qkv]|out).lora_")
-        alpha_numlayer = re.compile(r"_attn(\d)_to_([qkv]|out).alpha")
-        for k, v in state_dict.items():
-            if "attn" not in k or "lora_te" in k:
-                # currently not support: ff, clip-attn
-                continue
-            k = k.replace("lora_unet_down_blocks_", "down_blocks.")
-            k = k.replace("lora_unet_up_blocks_", "up_blocks.")
-            k = k.replace("lora_unet_mid_block_", "mid_block_")
-            k = k.replace("_attentions_", ".attentions.")
-            k = k.replace("_transformer_blocks_", ".transformer_blocks.")
-            k = k.replace("to_out_0", "to_out")
-            k = attn_numlayer.sub(r".attn\1.processor.to_\2_lora.", k)
-            k = alpha_numlayer.sub(r".attn\1.processor.to_\2_lora.alpha", k)
-            df_lora[k] = v
-        state_dict = df_lora
-
-    # fill attn processors
-    attn_processors = {}
-
-    is_lora = all("lora" in k for k in state_dict.keys())
-
-    if is_lora:
-        lora_grouped_dict = defaultdict(dict)
-        for key, value in state_dict.items():
-            if "alpha" in key:
-                attn_processor_key, sub_key = ".".join(key.split(".")[:-2]), ".".join(
-                    key.split(".")[-2:]
-                )
-            else:
-                attn_processor_key, sub_key = ".".join(key.split(".")[:-3]), ".".join(
-                    key.split(".")[-3:]
-                )
-            lora_grouped_dict[attn_processor_key][sub_key] = value
-
-        for key, value_dict in lora_grouped_dict.items():
-            rank = value_dict["to_k_lora.down.weight"].shape[0]
-            cross_attention_dim = value_dict["to_k_lora.down.weight"].shape[1]
-            hidden_size = value_dict["to_k_lora.up.weight"].shape[0]
-
-            attn_processors[key] = LoRACrossAttnProcessor(
-                hidden_size=hidden_size,
-                cross_attention_dim=cross_attention_dim,
-                rank=rank,
-                scale=scale,
-            )
-            attn_processors[key].load_state_dict(value_dict, strict=False)
-
-    else:
-        raise ValueError(
-            f"{model_file} does not seem to be in the correct format expected by LoRA training."
-        )
-
-    # set correct dtype & device
-    attn_processors = {
-        k: v.to(device=unet.device, dtype=unet.dtype)
-        for k, v in attn_processors.items()
-    }
-
-    # set layers
-    unet.set_attn_processor(attn_processors)
-
-
 class CrossAttnProcessor(nn.Module):
     def __call__(
         self,
@@ -148,7 +75,6 @@ class CrossAttnProcessor(nn.Module):
         hidden_states,
         encoder_hidden_states=None,
         attention_mask=None,
-        qkvo_bias=None,
     ):
         batch_size, sequence_length, _ = hidden_states.shape
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)
@@ -166,11 +92,6 @@ class CrossAttnProcessor(nn.Module):
         key = attn.to_k(encoder_states)
         value = attn.to_v(encoder_states)
 
-        if qkvo_bias is not None:
-            query += qkvo_bias["q"](hidden_states)
-            key += qkvo_bias["k"](encoder_states)
-            value += qkvo_bias["v"](encoder_states)
-
         query = attn.head_to_batch_dim(query)
         key = attn.head_to_batch_dim(key)
         value = attn.head_to_batch_dim(value)
@@ -219,76 +140,11 @@ class CrossAttnProcessor(nn.Module):
         # linear proj
         hidden_states = attn.to_out[0](hidden_states)
 
-        if qkvo_bias is not None:
-            hidden_states += qkvo_bias["o"](hidden_states)
-
         # dropout
         hidden_states = attn.to_out[1](hidden_states)
 
         return hidden_states
 
-
-class LoRACrossAttnProcessor(CrossAttnProcessor):
-    def __init__(self, hidden_size, cross_attention_dim=None, rank=4, scale=1.0):
-        super().__init__()
-
-        self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank)
-        self.to_k_lora = LoRALinearLayer(
-            cross_attention_dim or hidden_size, hidden_size, rank
-        )
-        self.to_v_lora = LoRALinearLayer(
-            cross_attention_dim or hidden_size, hidden_size, rank
-        )
-        self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank)
-        self.scale = scale
-
-    def __call__(
-        self,
-        attn,
-        hidden_states,
-        encoder_hidden_states=None,
-        attention_mask=None,
-    ):
-        scale = self.scale
-        qkvo_bias = {
-            "q": lambda inputs: scale * self.to_q_lora(inputs),
-            "k": lambda inputs: scale * self.to_k_lora(inputs),
-            "v": lambda inputs: scale * self.to_v_lora(inputs),
-            "o": lambda inputs: scale * self.to_out_lora(inputs),
-        }
-        return super().__call__(
-            attn, hidden_states, encoder_hidden_states, attention_mask, qkvo_bias
-        )
-
-
-class LoRALinearLayer(nn.Module):
-    def __init__(self, in_features, out_features, rank=4):
-        super().__init__()
-
-        if rank > min(in_features, out_features):
-            raise ValueError(
-                f"LoRA rank {rank} must be less or equal than {min(in_features, out_features)}"
-            )
-
-        self.down = nn.Linear(in_features, rank, bias=False)
-        self.up = nn.Linear(rank, out_features, bias=False)
-        self.scale = 1.0
-        self.alpha = rank
-
-        nn.init.normal_(self.down.weight, std=1 / rank)
-        nn.init.zeros_(self.up.weight)
-
-    def forward(self, hidden_states):
-        orig_dtype = hidden_states.dtype
-        dtype = self.down.weight.dtype
-        rank = self.down.out_features
-
-        down_hidden_states = self.down(hidden_states.to(dtype))
-        up_hidden_states = self.up(down_hidden_states) * (self.alpha / rank)
-
-        return up_hidden_states.to(orig_dtype)
-
-
 class ModelWrapper:
     def __init__(self, model, alphas_cumprod):
         self.model = model

requirements.txt

@@ -1,8 +0,0 @@
-torch
-einops
-diffusers
-transformers
-k_diffusion
-safetensors
-gradio
-torch