Merge branch 'cpu-cmdline-opt' of https://github.com/brkirch/stable-diffusion-webui into cpu-cmdline-opt

2 files changed, 8 insertions, 3 deletions

diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
index e8394ff6..7c44ea5b 100644
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@@ -9,6 +9,9 @@ from torchvision import transforms
 import random

 import tqdm

 from modules import devices

+import re

+

+re_tag = re.compile(r"[a-zA-Z][_\w\d()]+")

 

 

 class PersonalizedBase(Dataset):

@@ -38,8 +41,8 @@ class PersonalizedBase(Dataset):
             image = image.resize((self.width, self.height), PIL.Image.BICUBIC)

 

             filename = os.path.basename(path)

-            filename_tokens = os.path.splitext(filename)[0].replace('_', '-').replace(' ', '-').split('-')

-            filename_tokens = [token for token in filename_tokens if token.isalpha()]

+            filename_tokens = os.path.splitext(filename)[0]

+            filename_tokens = re_tag.findall(filename_tokens)

 

             npimage = np.array(image).astype(np.uint8)

             npimage = (npimage / 127.5 - 1.0).astype(np.float32)

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index 209e928f..f545a993 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -26,7 +26,9 @@ def preprocess(process_src, process_dst, process_flip, process_split, process_ca
         if process_caption:

             caption = "-" + shared.interrogator.generate_caption(image)

         else:

-            caption = ""

+            caption = filename

+            caption = os.path.splitext(caption)[0]

+            caption = os.path.basename(caption)

 

         image.save(os.path.join(dst, f"{index:05}-{subindex[0]}{caption}.png"))

         subindex[0] += 1