hoelee/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-08-04 19:22:32 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

add support for gelbooru tags in filenames for textual inversion

Browse Source
This commit is contained in:
AUTOMATIC
2022-10-04 08:52:11 +03:00
parent 2865ef4b9a
commit 5ef0baf5ea
2 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
modules/textual_inversion/dataset.py
View File

@@ -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)