diff options
| author | AUTOMATIC <16777216c@gmail.com> | 2022-10-04 05:52:11 +0000 |
|---|---|---|
| committer | AUTOMATIC <16777216c@gmail.com> | 2022-10-04 05:52:27 +0000 |
| commit | 5ef0baf5eaec7f21a1666af424405cbee19f3764 (patch) | |
| tree | 44ae1840049bc638ef602ff111d2fb25da1e533f /modules/textual_inversion/dataset.py | |
| parent | 2865ef4b9ab16d56326cc805541bebcf01d099bc (diff) | |
| download | stable-diffusion-webui-gfx803-5ef0baf5eaec7f21a1666af424405cbee19f3764.tar.gz stable-diffusion-webui-gfx803-5ef0baf5eaec7f21a1666af424405cbee19f3764.tar.bz2 stable-diffusion-webui-gfx803-5ef0baf5eaec7f21a1666af424405cbee19f3764.zip | |
add support for gelbooru tags in filenames for textual inversion
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)
|