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

Added ability to use short captions from json caption file

Browse Source
This commit is contained in:
Jaret Burkett
2025-04-16 08:32:28 -06:00
parent 68018c908e
commit c12036df95
2 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
toolkit/dataloader_mixins.py
View File

@@ -308,6 +308,8 @@ class CaptionProcessingDTOMixin:
self.raw_caption = caption_dict[self.path]["caption"]
if 'caption_short' in caption_dict[self.path]:
self.raw_caption_short = caption_dict[self.path]["caption_short"]
if self.dataset_config.use_short_captions:
self.raw_caption = caption_dict[self.path]["caption_short"]
else:
# see if prompt file exists
path_no_ext = os.path.splitext(self.path)[0]
@@ -330,7 +332,8 @@ class CaptionProcessingDTOMixin:
prompt = prompt_json['caption']
if 'caption_short' in prompt_json:
short_caption = prompt_json['caption_short']
if self.dataset_config.use_short_captions:
prompt = short_caption
if 'extra_values' in prompt_json:
self.extra_values = prompt_json['extra_values']