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 import gradio as gr import os import json 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 from cyberharem.utils import get_hf_client, get_hf_fs from hbutils.system import TemporaryDirectory from cyberharem.utils import download_file as cyber_download_file from huggingface_hub import hf_hub_url, hf_hub_download def start_func(chars, is_cpu, udghs, game_index=None): if not udghs: if is_cpu: os.environ['ONNX_MODE'] = 'CPUExecutionProvider' char_list = chars.split(',') for ch in char_list: crawl_dataset_to_huggingface(ch) print(ch + "完成") return str(chars)+" 上传完成" else: dgrepo = 'deepghs/game_characters' if is_cpu: os.environ['ONNX_MODE'] = 'CPUExecutionProvider' with TemporaryDirectory() as jsondir: print("Downloading jsons..") hf_fs = get_hf_fs() _dgdatas = [file for file in hf_fs.glob(f'datasets/{dgrepo}/*/pixiv_characters.json')] if game_index: name = _dgdatas[game_index-1] os.makedirs(os.path.basename(os.path.dirname(name)), exist_ok=True) # print(f'https://huggingface.co/{dgrepo}/blob/main/{os.path.basename(os.path.dirname(name))}/{os.path.basename(name)}') js = hf_hub_download( # f'https://huggingface.co/{dgrepo}/blob/main/{os.path.basename(os.path.dirname(name))}/{os.path.basename(name)}', # hf_hub_url(dgrepo, filename=os.path.relpath(name, dgrepo)), repo_id=dgrepo, repo_type='dataset', # os.path.join(os.path.basename(os.path.dirname(name)), 'pixiv_characters.json'), filename=os.path.join(os.path.basename(os.path.dirname(name)), 'pixiv_characters.json'), token=os.environ['HF_TOKEN'] ) # with open(os.path.join(os.path.basename(os.path.dirname(name)), 'pixiv_characters.json'), 'r') as f: with open(js, 'r', encoding='utf-8') as f: jt = json.load(f) chs = jt['characters'] for jp in chs: jp = jp['jpname'] print(jp, 'start...') crawl_dataset_to_huggingface(jp) print(jp + "完成") else: for name in _dgdatas: os.makedirs(os.path.basename(os.path.dirname(name)), exist_ok=True) # print(f'https://huggingface.co/{dgrepo}/blob/main/{os.path.basename(os.path.dirname(name))}/{os.path.basename(name)}') js = hf_hub_download( # f'https://huggingface.co/{dgrepo}/blob/main/{os.path.basename(os.path.dirname(name))}/{os.path.basename(name)}', # hf_hub_url(dgrepo, filename=os.path.relpath(name, dgrepo)), repo_id=dgrepo, repo_type='dataset', # os.path.join(os.path.basename(os.path.dirname(name)), 'pixiv_characters.json'), filename=os.path.join(os.path.basename(os.path.dirname(name)), 'pixiv_characters.json'), token=os.environ['HF_TOKEN'] ) # with open(os.path.join(os.path.basename(os.path.dirname(name)), 'pixiv_characters.json'), 'r') as f: with open(js, 'r', encoding='utf-8') as f: jt = json.load(f) chs = jt['characters'] for jp in chs: jp = jp['jpname'] print(jp, 'start...') with open(os.path.join(os.path.basename(os.path.dirname(name)), 'log.txt'), 'w') as log_f: print(f'{jp} is in crawl.', file=log_f) crawl_dataset_to_huggingface(jp) print(jp + "完成") return "完成" def main(): parser = argparse.ArgumentParser() parser.add_argument('--char', type=str, help='角色列表', default=None) parser.add_argument('--index', type=int, default=None) args = parser.parse_args() start_func(args.char, True, False if args.char else True, args.index) print("全部完成") if __name__ == "__main__": main()