aboutsummaryrefslogtreecommitdiffstats
path: root/modules/textual_inversion/dataset.py
diff options
context:
space:
mode:
authorJairo Correa <jn.j41r0@gmail.com>2022-10-04 22:53:52 +0000
committerJairo Correa <jn.j41r0@gmail.com>2022-10-04 22:53:52 +0000
commit1f50971fb8c83c255c2819dd0b3f29a46b74f7d9 (patch)
treefd57f40a1ffa2b28105ec0bb3f7f3ab4a742681a /modules/textual_inversion/dataset.py
parentad0cc85d1f0bd52877963f296eb1257a0c2b012b (diff)
parentef40e4cd4d383a3405e03f1da3f5b5a1820a8f53 (diff)
downloadstable-diffusion-webui-gfx803-1f50971fb8c83c255c2819dd0b3f29a46b74f7d9.tar.gz
stable-diffusion-webui-gfx803-1f50971fb8c83c255c2819dd0b3f29a46b74f7d9.tar.bz2
stable-diffusion-webui-gfx803-1f50971fb8c83c255c2819dd0b3f29a46b74f7d9.zip
Merge branch 'master' into fix-vram
Diffstat (limited to 'modules/textual_inversion/dataset.py')
-rw-r--r--modules/textual_inversion/dataset.py7
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)