github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fixed issue with sizes for some images being loaded sideways resulting in squished images.

Browse Source
This commit is contained in:
Jaret Burkett
2024-10-20 11:51:29 -06:00
parent e3ebd73610
commit bedb8197a2
2 changed files with 26 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
toolkit/data_loader.py
View File

@@ -441,16 +441,24 @@ class AiToolkitDataset(LatentCachingMixin, CLIPCachingMixin, BucketsMixin, Capti
if not os.path.isdir(self.dataset_path):
dataset_folder = os.path.dirname(dataset_folder)
dataset_size_file = os.path.join(dataset_folder, '.aitk_size.json')
dataloader_version = "0.1.1"
if os.path.exists(dataset_size_file):
try:
with open(dataset_size_file, 'r') as f:
self.size_database = json.load(f)
if "__version__" not in self.size_database or self.size_database["__version__"] != dataloader_version:
print("Upgrading size database to new version")
# old version, delete and recreate
self.size_database = {}
except Exception as e:
print(f"Error loading size database: {dataset_size_file}")
print(e)
self.size_database = {}
else:
self.size_database = {}
self.size_database["__version__"] = dataloader_version
bad_count = 0
for file in tqdm(file_list):
@@ -461,6 +469,7 @@ class AiToolkitDataset(LatentCachingMixin, CLIPCachingMixin, BucketsMixin, Capti
dataset_config=dataset_config,
dataloader_transforms=self.transform,
size_database=self.size_database,
dataset_root=dataset_folder,
)
self.file_list.append(file_item)
except Exception as e: