diff --git a/toolkit/basic.py b/toolkit/basic.py index 378444571..554226ee7 100644 --- a/toolkit/basic.py +++ b/toolkit/basic.py @@ -2,6 +2,24 @@ import os import torch +from PIL import Image + +# Maps the string stored on DatasetConfig.resize_method to the PIL resample +# filter used for all image resize calls in data_loader.py, dataloader_mixins.py, +# and control_generator.py. +RESIZE_METHODS = { + 'bicubic': Image.BICUBIC, + 'lanczos': Image.LANCZOS, +} + + +def get_resize_method(name: str = 'lanczos'): + try: + return RESIZE_METHODS[name] + except KeyError: + raise ValueError( + f"Unknown resize_method '{name}'. Valid options are: {list(RESIZE_METHODS.keys())}" + ) def value_map(inputs, min_in, max_in, min_out, max_out): diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index a0ef97177..01eca6a1b 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -931,6 +931,9 @@ def __init__(self, **kwargs): self.random_scale: bool = kwargs.get('random_scale', False) self.random_crop: bool = kwargs.get('random_crop', False) self.resolution: int = kwargs.get('resolution', 512) + # PIL resample filter used for all image resize/scale ops on this dataset. + # One of: 'bicubic', 'lanczos' + self.resize_method: str = kwargs.get('resize_method', 'lanczos') self.scale: float = kwargs.get('scale', 1.0) self.buckets: bool = kwargs.get('buckets', True) self.bucket_tolerance: int = kwargs.get('bucket_tolerance', 64) diff --git a/toolkit/control_generator.py b/toolkit/control_generator.py index 88c03ee31..eb491a452 100644 --- a/toolkit/control_generator.py +++ b/toolkit/control_generator.py @@ -9,6 +9,8 @@ from torchvision import transforms +from toolkit.basic import get_resize_method + # supress all warnings import warnings @@ -76,7 +78,7 @@ def load_image(self, img_path): scale = math.sqrt(max_size / (w * h)) w = int(w * scale) h = int(h * scale) - image = image.resize((w, h), Image.BICUBIC) + image = image.resize((w, h), get_resize_method('lanczos')) return image def control_save_path(self, img_path, control_type): diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index ee90e2631..1e8ad09f8 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -17,6 +17,7 @@ import albumentations as A from toolkit import image_utils +from toolkit.basic import get_resize_method from toolkit.buckets import get_bucket_for_image_size, BucketResolution from toolkit.config_modules import DatasetConfig, preprocess_dataset_raw_config from toolkit.dataloader_mixins import CaptionMixin, BucketsMixin, LatentCachingMixin, Augments, CLIPCachingMixin, ControlCachingMixin, TextEmbeddingCachingMixin @@ -97,6 +98,7 @@ def __init__(self, config): self.random_crop = self.random_scale if self.random_scale else self.get_config('random_crop', False) self.resolution = self.get_config('resolution', 256) + self.resize_method = get_resize_method(self.get_config('resize_method', 'lanczos')) self.file_list = [os.path.join(self.path, file) for file in os.listdir(self.path) if file.lower().endswith(('.jpg', '.jpeg', '.png', '.webp'))] @@ -150,7 +152,7 @@ def __getitem__(self, index): img = Image.fromarray(np.random.randint(0, 255, (1024, 1024, 3), dtype=np.uint8)) # Downscale the source image first - img = img.resize((int(img.size[0] * self.scale), int(img.size[1] * self.scale)), Image.BICUBIC) + img = img.resize((int(img.size[0] * self.scale), int(img.size[1] * self.scale)), self.resize_method) min_img_size = min(img.size) if self.random_crop: @@ -164,11 +166,11 @@ def __getitem__(self, index): scaler = scale_size / min_img_size scale_width = int((img.width + 5) * scaler) scale_height = int((img.height + 5) * scaler) - img = img.resize((scale_width, scale_height), Image.BICUBIC) + img = img.resize((scale_width, scale_height), self.resize_method) img = transforms.RandomCrop(self.resolution)(img) else: img = transforms.CenterCrop(min_img_size)(img) - img = img.resize((self.resolution, self.resolution), Image.BICUBIC) + img = img.resize((self.resolution, self.resolution), self.resize_method) img = self.transform(img) @@ -236,6 +238,7 @@ def __init__(self, config): self.network_weight = self.get_config('network_weight', 1.0) self.pos_weight = self.get_config('pos_weight', self.network_weight) self.neg_weight = self.get_config('neg_weight', self.network_weight) + self.resize_method = get_resize_method(self.get_config('resize_method', 'lanczos')) supported_exts = ('.jpg', '.jpeg', '.png', '.webp', '.JPEG', '.JPG', '.PNG', '.WEBP') @@ -357,9 +360,9 @@ def __getitem__(self, index): img2_crop_width = bucket_resolution["width"] # scale then center crop images - img1 = img1.resize((img1_scale_to_width, img1_scale_to_height), Image.BICUBIC) + img1 = img1.resize((img1_scale_to_width, img1_scale_to_height), self.resize_method) img1 = transforms.CenterCrop((img1_crop_height, img1_crop_width))(img1) - img2 = img2.resize((img2_scale_to_width, img2_scale_to_height), Image.BICUBIC) + img2 = img2.resize((img2_scale_to_width, img2_scale_to_height), self.resize_method) img2 = transforms.CenterCrop((img2_crop_height, img2_crop_width))(img2) # combine them side by side @@ -374,7 +377,7 @@ def __getitem__(self, index): width = int(img.size[0] * height / img.size[1]) # Downscale the source image first - img = img.resize((width, height), Image.BICUBIC) + img = img.resize((width, height), self.resize_method) prompt = self.get_prompt_item(index) img = self.transform(img) diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 0c5ac31d6..c419020cb 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -17,7 +17,7 @@ from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection, SiglipImageProcessor from toolkit.audio.preserve_pitch import time_stretch_preserve_pitch -from toolkit.basic import flush, value_map +from toolkit.basic import flush, value_map, get_resize_method from toolkit.buckets import get_bucket_for_image_size, get_resolution from toolkit.config_modules import ControlTypes from toolkit.control_generator import ControlGenerator @@ -619,7 +619,7 @@ def load_and_process_video( img = img.transpose(Image.FLIP_TOP_BOTTOM) # Apply bucketing - img = img.resize((self.scale_to_width, self.scale_to_height), Image.BICUBIC) + img = img.resize((self.scale_to_width, self.scale_to_height), get_resize_method(self.dataset_config.resize_method)) img = img.crop(( self.crop_x, self.crop_y, @@ -829,7 +829,7 @@ def load_and_process_image( if self.dataset_config.buckets: # scale and crop based on file item - img = img.resize((self.scale_to_width, self.scale_to_height), Image.BICUBIC) + img = img.resize((self.scale_to_width, self.scale_to_height), get_resize_method(self.dataset_config.resize_method)) # crop to x_crop, y_crop, x_crop + crop_width, y_crop + crop_height if img.width < self.crop_x + self.crop_width or img.height < self.crop_y + self.crop_height: # todo look into this. This still happens sometimes @@ -847,7 +847,7 @@ def load_and_process_image( # TODO this is nto right img = img.resize( (int(img.size[0] * self.dataset_config.scale), int(img.size[1] * self.dataset_config.scale)), - Image.BICUBIC) + get_resize_method(self.dataset_config.resize_method)) min_img_size = min(img.size) if self.dataset_config.random_crop: if self.dataset_config.random_scale and min_img_size > self.dataset_config.resolution: @@ -860,11 +860,11 @@ def load_and_process_image( scaler = scale_size / min_img_size scale_width = int((img.width + 5) * scaler) scale_height = int((img.height + 5) * scaler) - img = img.resize((scale_width, scale_height), Image.BICUBIC) + img = img.resize((scale_width, scale_height), get_resize_method(self.dataset_config.resize_method)) img = transforms.RandomCrop(self.dataset_config.resolution)(img) else: img = transforms.CenterCrop(min_img_size)(img) - img = img.resize((self.dataset_config.resolution, self.dataset_config.resolution), Image.BICUBIC) + img = img.resize((self.dataset_config.resolution, self.dataset_config.resolution), get_resize_method(self.dataset_config.resize_method)) if self.augments is not None and len(self.augments) > 0: # do augmentations @@ -943,7 +943,7 @@ def load_inpaint_image(self: 'FileItemDTO'): if self.dataset_config.buckets: # scale and crop based on file item - img = img.resize((self.scale_to_width, self.scale_to_height), Image.BICUBIC) + img = img.resize((self.scale_to_width, self.scale_to_height), get_resize_method(self.dataset_config.resize_method)) # img = transforms.CenterCrop((self.crop_height, self.crop_width))(img) # crop img = img.crop(( @@ -1061,7 +1061,7 @@ def load_control_image(self: 'FileItemDTO'): if not self.full_size_control_images: # we just scale them to 512x512: w, h = img.size - img = img.resize((512, 512), Image.BICUBIC) + img = img.resize((512, 512), get_resize_method(self.dataset_config.resize_method)) elif not self.use_raw_control_images: w, h = img.size @@ -1074,7 +1074,7 @@ def load_control_image(self: 'FileItemDTO'): if self.dataset_config.buckets: # scale and crop based on file item - img = img.resize((self.scale_to_width, self.scale_to_height), Image.BICUBIC) + img = img.resize((self.scale_to_width, self.scale_to_height), get_resize_method(self.dataset_config.resize_method)) # img = transforms.CenterCrop((self.crop_height, self.crop_width))(img) # crop img = img.crop(( @@ -1295,7 +1295,7 @@ def load_clip_image(self: 'FileItemDTO'): else: # image must be square. If it is not, we will resize/squish it so it is, that way we don't crop out data # resize to the smallest dimension - img = img.resize((min_size, min_size), Image.BICUBIC) + img = img.resize((min_size, min_size), get_resize_method(self.dataset_config.resize_method)) if self.has_clip_augmentations: self.clip_image_tensor = self.augment_clip_image(img, transform=None) @@ -1522,7 +1522,7 @@ def load_mask_image(self: 'FileItemDTO'): if self.dataset_config.buckets: # scale and crop based on file item - img = img.resize((self.scale_to_width, self.scale_to_height), Image.BICUBIC) + img = img.resize((self.scale_to_width, self.scale_to_height), get_resize_method(self.dataset_config.resize_method)) # img = transforms.CenterCrop((self.crop_height, self.crop_width))(img) # crop img = img.crop(( @@ -1597,7 +1597,7 @@ def load_unconditional_image(self: 'FileItemDTO'): if self.dataset_config.buckets: # scale and crop based on file item - img = img.resize((self.scale_to_width, self.scale_to_height), Image.BICUBIC) + img = img.resize((self.scale_to_width, self.scale_to_height), get_resize_method(self.dataset_config.resize_method)) # img = transforms.CenterCrop((self.crop_height, self.crop_width))(img) # crop img = img.crop(( diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index 36ec7fd13..4f994d7d0 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -1199,6 +1199,16 @@ export default function SimpleJob({ { value: 'caption', label: 'caption' }, ]} /> + setJobConfig(value, `config.process[0].datasets[${i}].resize_method`)} + options={[ + { value: 'lanczos', label: 'Lanczos' }, + { value: 'bicubic', label: 'Bicubic' }, + ]} + /> {modelArch?.additionalSections?.includes('datasets.num_frames') && !dataset.auto_frame_count && (