mirror of
https://github.com/kohya-ss/sd-scripts.git
synced 2026-04-06 13:47:06 +00:00
72 lines
1.8 KiB
Python
72 lines
1.8 KiB
Python
from typing import *
|
|
from huggingface_hub import HfApi
|
|
from pathlib import Path
|
|
import argparse
|
|
import os
|
|
|
|
from library.utils import fire_in_thread
|
|
|
|
|
|
def exists_repo(
|
|
repo_id: str, repo_type: str, revision: str = "main", token: str = None
|
|
):
|
|
api = HfApi(
|
|
token=token,
|
|
)
|
|
try:
|
|
api.repo_info(repo_id=repo_id, revision=revision, repo_type=repo_type)
|
|
return True
|
|
except:
|
|
return False
|
|
|
|
|
|
@fire_in_thread
|
|
def upload(
|
|
src: Union[str, Path, bytes, BinaryIO],
|
|
args: argparse.Namespace,
|
|
dest_suffix: str = "",
|
|
):
|
|
repo_id = args.huggingface_repo_id
|
|
repo_type = args.huggingface_repo_type
|
|
token = args.huggingface_token
|
|
path_in_repo = args.huggingface_path_in_repo + dest_suffix
|
|
private = args.huggingface_repo_visibility == "private"
|
|
api = HfApi(token=token)
|
|
if not exists_repo(repo_id=repo_id, repo_type=repo_type, token=token):
|
|
api.create_repo(repo_id=repo_id, repo_type=repo_type, private=private)
|
|
|
|
is_folder = (type(src) == str and os.path.isdir(src)) or (
|
|
isinstance(src, Path) and src.is_dir()
|
|
)
|
|
if is_folder:
|
|
api.upload_folder(
|
|
repo_id=repo_id,
|
|
repo_type=repo_type,
|
|
folder_path=src,
|
|
path_in_repo=path_in_repo,
|
|
)
|
|
else:
|
|
api.upload_file(
|
|
repo_id=repo_id,
|
|
repo_type=repo_type,
|
|
path_or_fileobj=src,
|
|
path_in_repo=path_in_repo,
|
|
)
|
|
|
|
|
|
def list_dir(
|
|
repo_id: str,
|
|
subfolder: str,
|
|
repo_type: str,
|
|
revision: str = "main",
|
|
token: str = None,
|
|
):
|
|
api = HfApi(
|
|
token=token,
|
|
)
|
|
repo_info = api.repo_info(repo_id=repo_id, revision=revision, repo_type=repo_type)
|
|
file_list = [
|
|
file for file in repo_info.siblings if file.rfilename.startswith(subfolder)
|
|
]
|
|
return file_list
|