add resize method for buckets lanczos/bicubic

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

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)

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

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)

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((

ui/src/app/jobs/new/SimpleJob.tsx

@@ -1199,6 +1199,16 @@ export default function SimpleJob({
                           { value: 'caption', label: 'caption' },
                         ]}
                       />
+                      <SelectInput
+                        label="Resize Method"
+                        className="pt-2"
+                        value={dataset.resize_method || 'lanczos'}
+                        onChange={value => 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 && (
                         <NumberInput

ui/src/app/jobs/new/jobConfig.ts

@@ -14,6 +14,7 @@ export const defaultDatasetConfig: DatasetConfig = {
   is_reg: false,
   network_weight: 1,
   resolution: [512, 768, 1024],
+  resize_method: 'lanczos',
   controls: [],
   shrink_video_to_frames: true,
   num_frames: 1,

ui/src/types.ts

@@ -94,6 +94,7 @@ export interface DatasetConfig {
   network_weight: number;
   cache_latents_to_disk?: boolean;
   resolution: number[];
+  resize_method?: 'bicubic' | 'lanczos';
   controls: string[];
   control_path?: string | null;
   num_frames: number;