ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 11:11:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #304 from hameerabbasi/fix-caption-loads

Browse Source 
Fix caption loading
This commit is contained in:
Jaret Burkett
2025-06-12 07:44:12 -06:00
committed by GitHub
parent ffb5fe0667 5e86139e0a
commit cf11f128b9
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
toolkit/dataloader_mixins.py
View File

@@ -111,12 +111,12 @@ class CaptionMixin:
if not hasattr(self, 'file_list'): if not hasattr(self, 'file_list'):
raise Exception('file_list not found on class instance') raise Exception('file_list not found on class instance')
img_path_or_tuple = self.file_list[index] img_path_or_tuple = self.file_list[index]
ext = self.dataset_config.caption_ext
if isinstance(img_path_or_tuple, tuple): if isinstance(img_path_or_tuple, tuple):
img_path = img_path_or_tuple[0] if isinstance(img_path_or_tuple[0], str) else img_path_or_tuple[0].path img_path = img_path_or_tuple[0] if isinstance(img_path_or_tuple[0], str) else img_path_or_tuple[0].path
# check if either has a prompt file # check if either has a prompt file
path_no_ext = os.path.splitext(img_path)[0] path_no_ext = os.path.splitext(img_path)[0]
prompt_path = None prompt_path = None
ext = self.dataset_config.caption_ext
prompt_path = path_no_ext + ext prompt_path = path_no_ext + ext
else: else:
img_path = img_path_or_tuple if isinstance(img_path_or_tuple, str) else img_path_or_tuple.path img_path = img_path_or_tuple if isinstance(img_path_or_tuple, str) else img_path_or_tuple.path
@@ -315,7 +315,7 @@ class CaptionProcessingDTOMixin:
# see if prompt file exists # see if prompt file exists
path_no_ext = os.path.splitext(self.path)[0] path_no_ext = os.path.splitext(self.path)[0]
prompt_ext = self.dataset_config.caption_ext prompt_ext = self.dataset_config.caption_ext
prompt_path = f"{path_no_ext}.{prompt_ext}" prompt_path = path_no_ext + prompt_ext
short_caption = None short_caption = None
if os.path.exists(prompt_path): if os.path.exists(prompt_path):