Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added warning on training without captions #533

Merged
merged 4 commits into from
May 28, 2023
Merged

Added warning on training without captions #533

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5a1a14f
Update train_util.py
TingTingin May 23, 2023
d859a3a
Update train_util.py
TingTingin May 23, 2023
061e157
Update train_util.py
TingTingin May 23, 2023
db756e9
Update train_util.py
TingTingin May 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 19 additions & 1 deletion library/train_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -348,6 +348,8 @@ def __init__(
self.is_reg = is_reg
self.class_tokens = class_tokens
self.caption_extension = caption_extension
if self.caption_extension and not self.caption_extension.startswith("."):
self.caption_extension = "." + self.caption_extension

def __eq__(self, other) -> bool:
if not isinstance(other, DreamBoothSubset):
Expand Down Expand Up @@ -1081,16 +1083,32 @@ def load_dreambooth_dir(subset: DreamBoothSubset):

# 画像ファイルごとにプロンプトを読み込み、もしあればそちらを使う
captions = []
missing_captions = []
for img_path in img_paths:
cap_for_img = read_caption(img_path, subset.caption_extension)
if cap_for_img is None and subset.class_tokens is None:
print(f"neither caption file nor class tokens are found. use empty caption for {img_path}")
kohya-ss marked this conversation as resolved.
Show resolved Hide resolved
captions.append("")
else:
captions.append(subset.class_tokens if cap_for_img is None else cap_for_img)
if cap_for_img is None:
captions.append(subset.class_tokens)
missing_captions.append(img_path)
kohya-ss marked this conversation as resolved.
Show resolved Hide resolved
else:
captions.append(cap_for_img)

self.set_tag_frequency(os.path.basename(subset.image_dir), captions) # タグ頻度を記録

if missing_captions:
number_of_missing_captions = len(missing_captions)
number_of_missing_captions_to_show = 5
remaining_missing_captions = number_of_missing_captions - number_of_missing_captions_to_show

print(f"No caption file found for {number_of_missing_captions} images. Training will continue without captions for these images")
for i, missing_caption in enumerate(missing_captions):
if i >= number_of_missing_captions_to_show:
print(missing_caption+f"... and {remaining_missing_captions} more")
break
print(missing_caption)
return img_paths, captions

print("prepare images.")
Expand Down