diff options
author | brkirch <brkirch@users.noreply.github.com> | 2022-10-04 08:25:27 +0000 |
---|---|---|
committer | GitHub <noreply@github.com> | 2022-10-04 08:25:27 +0000 |
commit | 35a00b013007c16c97ed79bee4cefcb798181dea (patch) | |
tree | fd3bf9f74b3b549d20836116eea759c5b06f48b8 /modules/textual_inversion/dataset.py | |
parent | b88e4ea7d6536ddc814a9101a258f87b06230019 (diff) | |
parent | 556c36b9607e3f4eacdddc85f8e7a78b29476ea7 (diff) | |
download | stable-diffusion-webui-gfx803-35a00b013007c16c97ed79bee4cefcb798181dea.tar.gz stable-diffusion-webui-gfx803-35a00b013007c16c97ed79bee4cefcb798181dea.tar.bz2 stable-diffusion-webui-gfx803-35a00b013007c16c97ed79bee4cefcb798181dea.zip |
Merge branch 'AUTOMATIC1111:master' into cpu-cmdline-opt
Diffstat (limited to 'modules/textual_inversion/dataset.py')
-rw-r--r-- | modules/textual_inversion/dataset.py | 7 |
1 files changed, 5 insertions, 2 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)
|