From eea5e24d13c5c5e2c3db7008cbdcfbd2ae8dce84 Mon Sep 17 00:00:00 2001
From: carefree0910 <heyujian@nolibox.com>
Date: Mon, 12 Jun 2023 09:02:26 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8Supported=20`custom=5Fembeddings`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 examples/carefree_creator/app.py              | 29 +++++++++++++++++++
 .../put your embeddings here                  |  2 ++
 2 files changed, 31 insertions(+)
 create mode 100644 examples/carefree_creator/custom_embeddings/put your embeddings here

diff --git a/examples/carefree_creator/app.py b/examples/carefree_creator/app.py
index 666e273a..d7470e2c 100644
--- a/examples/carefree_creator/app.py
+++ b/examples/carefree_creator/app.py
@@ -1,3 +1,6 @@
+import os
+import json
+
 from PIL import Image
 from typing import Any
 from typing import Dict
@@ -19,6 +22,8 @@
 
 
 TDataModel = TypeVar("TDataModel", bound="BaseModel")
+text_keys = ["text", "prompt", "negative_prompt"]
+CUSTOM_EMBEDDINGS_PATH = os.environ.get("CFDRAW_CFCREATOR_CUSTOM_EMBEDDING_PATH")
 
 
 def inject(
@@ -57,6 +62,30 @@ def inject(
             data.extraData["lora_paths"] = lora_paths
         if lora_scales:
             data.extraData["lora_scales"] = lora_scales
+    # custom embeddings
+    if CUSTOM_EMBEDDINGS_PATH is not None:
+        custom_embeddings_folder = Path(CUSTOM_EMBEDDINGS_PATH)
+    else:
+        custom_embeddings_folder = Path(__file__).parent / "custom_embeddings"
+    if custom_embeddings_folder.is_dir():
+        custom_embedding_paths = [
+            p
+            for p in custom_embeddings_folder.iterdir()
+            if p.is_file() and p.name.endswith(".ce")
+        ]
+        custom_embeddings = {}
+        for key in text_keys:
+            k_text = data.extraData.get(key)
+            if k_text is None:
+                continue
+            for ce_path in custom_embedding_paths:
+                stem = ce_path.stem
+                if stem in k_text and stem not in custom_embeddings:
+                    with open(ce_path, "r") as f:
+                        ce = json.load(f)
+                    custom_embeddings[stem] = ce
+        if custom_embeddings:
+            data.extraData["custom_embeddings"] = custom_embeddings
     # collect
     kw = shallow_copy_dict(data.extraData)
     if extra is not None:
diff --git a/examples/carefree_creator/custom_embeddings/put your embeddings here b/examples/carefree_creator/custom_embeddings/put your embeddings here
new file mode 100644
index 00000000..be41d3ed
--- /dev/null
+++ b/examples/carefree_creator/custom_embeddings/put your embeddings here	
@@ -0,0 +1,2 @@
+1. embeddings should be json file
+2. file names should be `*.ce`
\ No newline at end of file