From 41a3f63b72a228919b7222cbab433141416287e8 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sun, 10 Sep 2023 03:43:02 -0600
Subject: [PATCH] allow smaller images in buckets and bucket them

---
 toolkit/buckets.py     | 14 +++++++++++++-
 toolkit/data_loader.py | 10 +++++-----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/toolkit/buckets.py b/toolkit/buckets.py
index 750b7d2..524208a 100644
--- a/toolkit/buckets.py
+++ b/toolkit/buckets.py
@@ -72,15 +72,27 @@ def get_bucket_sizes(resolution: int = 512, divisibility: int = 8) -> List[Bucke
     return bucket_size_list
 
 
+def get_resolution(width, height):
+    num_pixels = width * height
+    # determine same number of pixels for square image
+    square_resolution = int(num_pixels ** 0.5)
+    return square_resolution
+
+
 def get_bucket_for_image_size(
         width: int,
         height: int,
         bucket_size_list: List[BucketResolution] = None,
         resolution: Union[int, None] = None
 ) -> BucketResolution:
+
     if bucket_size_list is None and resolution is None:
-        raise ValueError("Must provide either bucket_size_list or resolution")
+        # get resolution from width and height
+        resolution = get_resolution(width, height)
     if bucket_size_list is None:
+        # if real resolution is smaller, use that instead
+        real_resolution = get_resolution(width, height)
+        resolution = min(resolution, real_resolution)
         bucket_size_list = get_bucket_sizes(resolution=resolution)
 
     # Check for exact match first
diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py
index 90c8171..0cd3784 100644
--- a/toolkit/data_loader.py
+++ b/toolkit/data_loader.py
@@ -331,17 +331,17 @@ class AiToolkitDataset(Dataset, CaptionMixin, BucketsMixin):
                     path=file,
                     dataset_config=dataset_config
                 )
-                if file_item.scale_to_width < self.resolution or file_item.scale_to_height < self.resolution:
-                    bad_count += 1
-                else:
-                    self.file_list.append(file_item)
+                # if file_item.scale_to_width < self.resolution or file_item.scale_to_height < self.resolution:
+                #     bad_count += 1
+                # else:
+                self.file_list.append(file_item)
             except Exception as e:
                 print(f"Error processing image: {file}")
                 print(e)
                 bad_count += 1
 
         print(f"  -  Found {len(self.file_list)} images")
-        print(f"  -  Found {bad_count} images that are too small")
+        # print(f"  -  Found {bad_count} images that are too small")
         assert len(self.file_list) > 0, f"no images found in {self.dataset_path}"
 
         self.setup_epoch()