feat(textualInversion): very early support

Author: gadicc

.gitignore

@@ -14,7 +14,7 @@ dist/
 downloads/
 eggs/
 .eggs/
-lib/
+/lib/
 lib64/
 parts/
 sdist/

api/app.py

@@ -29,23 +29,23 @@
 from threading import Timer
 import extras
 
+from lib.textual_inversions import handle_textual_inversions
+from lib.vars import (
+    RUNTIME_DOWNLOADS,
+    USE_DREAMBOOTH,
+    MODEL_ID,
+    PIPELINE,
+    HF_AUTH_TOKEN,
+    HOME,
+    MODELS_DIR,
+)
 
-RUNTIME_DOWNLOADS = os.getenv("RUNTIME_DOWNLOADS") == "1"
-USE_DREAMBOOTH = os.getenv("USE_DREAMBOOTH") == "1"
 if USE_DREAMBOOTH:
     from train_dreambooth import TrainDreamBooth
-
-MODEL_ID = os.environ.get("MODEL_ID")
-PIPELINE = os.environ.get("PIPELINE")
-HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
-HOME = os.path.expanduser("~")
-MODELS_DIR = os.path.join(HOME, ".cache", "diffusers-api")
-
 print(os.environ.get("USE_PATCHMATCH"))
 if os.environ.get("USE_PATCHMATCH") == "1":
     from PyPatchMatch import patch_match
 
-
 torch.set_grad_enabled(False)
 always_normalize_model_id = None
 
@@ -332,6 +332,9 @@ def sendStatus():
     is_url = call_inputs.get("is_url", False)
     image_decoder = getFromUrl if is_url else decodeBase64Image
 
+    textual_inversions = call_inputs.get("textual_inversions", [])
+    handle_textual_inversions(textual_inversions, model)
+
     # Better to use new lora_weights in next section
     attn_procs = call_inputs.get("attn_procs", None)
     if attn_procs is not last_attn_procs:
@@ -388,7 +391,7 @@ def sendStatus():
                 cache_fname = "lora_weights--" + fname
                 path = os.path.join(MODELS_DIR, cache_fname)
                 if not os.path.exists(path):
-                    storage.download_and_extract(path)
+                    storage.download_and_extract(path, status=status)
                 print("Load lora_weights `" + lora_weights + "` from `" + path + "`")
                 pipeline.load_lora_weights(
                     MODELS_DIR, weight_name=cache_fname, local_files_only=True

api/lib/__init__.py

@@ -0,0 +1,75 @@
+import json
+import re
+import os
+from utils import Storage
+from .vars import MODELS_DIR
+
+last_textual_inversions = None
+last_textual_inversion_model = None
+loaded_textual_inversion_tokens = []
+
+tokenRe = re.compile(
+    r"[#&]{1}fname=(?P<fname>[^\.]+)\.(?:pt|safetensors)(&token=(?P<token>[^&]+))?$"
+)
+
+
+def strMap(str: str):
+    match = re.search(tokenRe, str)
+    print(match)
+    if match:
+        return match.group("token") or match.group("fname")
+
+
+def extract_tokens_from_list(textual_inversions: list):
+    return list(map(strMap, textual_inversions))
+
+
+def handle_textual_inversions(textual_inversions: list, model):
+    global last_textual_inversions
+    global last_textual_inversion_model
+    global loaded_textual_inversion_tokens
+
+    textual_inversions_str = json.dumps(textual_inversions)
+    if (
+        textual_inversions_str is not last_textual_inversions
+        or model is not last_textual_inversion_model
+    ):
+        if (model is not last_textual_inversion_model):
+            loaded_textual_inversion_tokens = []
+            last_textual_inversion_model = model
+        # print({"textual_inversions": textual_inversions})
+        # tokens_to_load = extract_tokens_from_list(textual_inversions)
+        # print({"tokens_loaded": loaded_textual_inversion_tokens})
+        # print({"tokens_to_load": tokens_to_load})
+        #
+        # for token in loaded_textual_inversion_tokens:
+        #     if token not in tokens_to_load:
+        #         print("[TextualInversion] Removing uneeded token: " + token)
+        #         del pipeline.tokenizer.get_vocab()[token]
+        #         # del pipeline.text_encoder.get_input_embeddings().weight.data[token]
+        #         pipeline.text_encoder.resize_token_embeddings(len(pipeline.tokenizer))
+        #
+        # loaded_textual_inversion_tokens = tokens_to_load
+
+        last_textual_inversions = textual_inversions_str
+        for textual_inversion in textual_inversions:
+            storage = Storage(textual_inversion, no_raise=True)
+            if storage:
+                storage_query_fname = storage.query.get("fname")
+                if storage_query_fname:
+                    fname = storage_query_fname[0]
+                else:
+                    fname = textual_inversion.split("/").pop()
+                path = os.path.join(MODELS_DIR, "textual_inversion--" + fname)
+                if not os.path.exists(path):
+                    storage.download_file(path)
+                print("Load textual inversion " + path)
+                token = storage.query.get("token", None)
+                if token not in loaded_textual_inversion_tokens:
+                    model.load_textual_inversion(
+                        path, token=token, local_files_only=True
+                    )
+                    loaded_textual_inversion_tokens.append(token)
+            else:
+                print("Load textual inversion " + textual_inversion)
+                model.load_textual_inversion(textual_inversion)

api/lib/textual_inversions.py

@@ -0,0 +1,16 @@
+import unittest
+from .textual_inversions import extract_tokens_from_list
+
+
+class TextualInversionsTest(unittest.TestCase):
+    def test_extract_tokens_query_fname(self):
+        tis = ["https://civitai.com/api/download/models/106132#fname=4nj0lie.pt"]
+        tokens = extract_tokens_from_list(tis)
+        self.assertEqual(tokens[0], "4nj0lie")
+
+    def test_extract_tokens_query_token(self):
+        tis = [
+            "https://civitai.com/api/download/models/106132#fname=4nj0lie.pt&token=4nj0lie"
+        ]
+        tokens = extract_tokens_from_list(tis)
+        self.assertEqual(tokens[0], "4nj0lie")

api/lib/textual_inversions_test.py

@@ -0,0 +1,9 @@
+import os
+
+RUNTIME_DOWNLOADS = os.getenv("RUNTIME_DOWNLOADS") == "1"
+USE_DREAMBOOTH = os.getenv("USE_DREAMBOOTH") == "1"
+MODEL_ID = os.environ.get("MODEL_ID")
+PIPELINE = os.environ.get("PIPELINE")
+HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
+HOME = os.path.expanduser("~")
+MODELS_DIR = os.path.join(HOME, ".cache", "diffusers-api")