fix: bring positional args back, add recursive to blip etc

This commit is contained in:
Linaqruf
2023-04-12 05:41:28 +07:00
parent bf8088e225
commit c316c63dff
4 changed files with 35 additions and 16 deletions

View File

@@ -2,6 +2,7 @@ import argparse
import os
import re
from pathlib import Path
from PIL import Image
from tqdm import tqdm
import torch
@@ -65,7 +66,8 @@ def main(args):
GenerationMixin._prepare_input_ids_for_generation = _prepare_input_ids_for_generation_patch
print(f"load images from {args.train_data_dir}")
image_paths = train_util.glob_images(args.train_data_dir)
train_data_dir_path = Path(args.train_data_dir)
image_paths = train_util.glob_images_pathlib(train_data_dir_path, args.recursive)
print(f"found {len(image_paths)} images.")
# できればcacheに依存せず明示的にダウンロードしたい
@@ -140,7 +142,8 @@ def setup_parser() -> argparse.ArgumentParser:
parser.add_argument("--remove_words", action="store_true",
help="remove like `with the words xxx` from caption / `with the words xxx`のような部分をキャプションから削除する")
parser.add_argument("--debug", action="store_true", help="debug mode")
parser.add_argument("--recursive", action="store_true", help="search for images in subfolders recursively")
return parser