import argparse
import os
from waifuc.action import HeadCountAction, AlignMinSizeAction, CCIPAction, ThreeStageSplitAction, ModeConvertAction, ClassFilterAction, PersonSplitAction, TaggingAction, RatingFilterAction, NoMonochromeAction, RandomFilenameAction, FirstNSelectAction, FilterSimilarAction, FileExtAction
from waifuc.export import SaveExporter, TextualInversionExporter
from waifuc.source import DanbooruSource, PixivSearchSource, ZerochanSource, LocalSource, GcharAutoSource
from cyberharem.dataset.crawler import crawl_dataset_to_huggingface


def main():
    os.environ['ONNX_MODE'] = 'CPUExecutionProvider'
    parser = argparse.ArgumentParser()
    parser.add_argument('--char', type=str, help='角色列表')
    parser.add_argument('--artist', type=str, help='画师列表')
    parser.add_argument('--token', type=str, help='token')
    os.environ['HF_TOKEN'] = args.token
    args = parser.parse_args()
    if args.artist:
        char_list = args.artist.split(',')
        for ch in char_list:
            crawl_dataset_to_huggingface(ch,DanbooruSource)
            print(ch+"完成")
        print("全部完成")
    else:
        char_list = args.char.split(',')
        for ch in char_list:
            crawl_dataset_to_huggingface(ch)
            print(ch+"完成")
        print("全部完成")


if __name__ == "__main__":
    main()