michael/stable-diffusion-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages 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
Show stats 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 random
import tqdm import tqdm
from modules import devices from modules import devices
import re
re_tag = re.compile(r"[a-zA-Z][_\w\d()]+")
class PersonalizedBase(Dataset): class PersonalizedBase(Dataset):
@ -38,8 +41,8 @@ class PersonalizedBase(Dataset):
image = image.resize((self.width, self.height), PIL.Image.BICUBIC) image = image.resize((self.width, self.height), PIL.Image.BICUBIC)
filename = os.path.basename(path) filename = os.path.basename(path)
filename_tokens = os.path.splitext(filename)[0].replace('_', '-').replace(' ', '-').split('-') filename_tokens = os.path.splitext(filename)[0]
filename_tokens = [token for token in filename_tokens if token.isalpha()] filename_tokens = re_tag.findall(filename_tokens)
npimage = np.array(image).astype(np.uint8) npimage = np.array(image).astype(np.uint8)
npimage = (npimage / 127.5 - 1.0).astype(np.float32) npimage = (npimage / 127.5 - 1.0).astype(np.float32)

4
modules/textual_inversion/preprocess.py
View file

@ -26,7 +26,9 @@ def preprocess(process_src, process_dst, process_flip, process_split, process_ca
if process_caption: if process_caption:
caption = "-" + shared.interrogator.generate_caption(image) caption = "-" + shared.interrogator.generate_caption(image)
else: 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")) image.save(os.path.join(dst, f"{index:05}-{subindex[0]}{caption}.png"))
subindex[0] += 1 subindex[0] += 1