Add checkbox for variable training dims

This commit is contained in:
dan
2023-01-08 01:34:52 +08:00
parent 448b9cedab
commit 669fb18d52
4 changed files with 8 additions and 5 deletions

View File

@@ -29,7 +29,7 @@ class DatasetEntry:
class PersonalizedBase(Dataset):
def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, cond_model=None, device=None, template_file=None, include_cond=False, batch_size=1, gradient_step=1, shuffle_tags=False, tag_drop_out=0, latent_sampling_method='once'):
def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, cond_model=None, device=None, template_file=None, include_cond=False, batch_size=1, gradient_step=1, shuffle_tags=False, tag_drop_out=0, latent_sampling_method='once', varsize=False):
re_word = re.compile(shared.opts.dataset_filename_word_regex) if len(shared.opts.dataset_filename_word_regex) > 0 else None
self.placeholder_token = placeholder_token
@@ -59,7 +59,7 @@ class PersonalizedBase(Dataset):
raise Exception("interrupted")
try:
image = Image.open(path).convert('RGB')
if width < 2000:
if not varsize:
image = image.resize((width, height), PIL.Image.BICUBIC)
else:
assert batch_size == 1, 'variable img size must have batch size 1'