modelscope/DiffSynth-Painter

1from huggingface_hub import hf_hub_download2from modelscope import snapshot_download3import os, shutil4from typing_extensions import Literal, TypeAlias5from typing import List6from ..configs.model_config import preset_models_on_huggingface, preset_models_on_modelscope, Preset_model_id7 8 9def download_from_modelscope(model_id, origin_file_path, local_dir):10    os.makedirs(local_dir, exist_ok=True)11    if os.path.basename(origin_file_path) in os.listdir(local_dir):12        print(f"    {os.path.basename(origin_file_path)} has been already in {local_dir}.")13        return14    else:15        print(f"    Start downloading {os.path.join(local_dir, os.path.basename(origin_file_path))}")16    snapshot_download(model_id, allow_file_pattern=origin_file_path, local_dir=local_dir)17    downloaded_file_path = os.path.join(local_dir, origin_file_path)18    target_file_path = os.path.join(local_dir, os.path.split(origin_file_path)[-1])19    if downloaded_file_path != target_file_path:20        shutil.move(downloaded_file_path, target_file_path)21        shutil.rmtree(os.path.join(local_dir, origin_file_path.split("/")[0]))22 23 24def download_from_huggingface(model_id, origin_file_path, local_dir):25    os.makedirs(local_dir, exist_ok=True)26    if os.path.basename(origin_file_path) in os.listdir(local_dir):27        print(f"    {os.path.basename(origin_file_path)} has been already in {local_dir}.")28        return29    else:30        print(f"    Start downloading {os.path.join(local_dir, os.path.basename(origin_file_path))}")31    hf_hub_download(model_id, origin_file_path, local_dir=local_dir)32 33 34Preset_model_website: TypeAlias = Literal[35    "HuggingFace",36    "ModelScope",37]38website_to_preset_models = {39    "HuggingFace": preset_models_on_huggingface,40    "ModelScope": preset_models_on_modelscope,41}42website_to_download_fn = {43    "HuggingFace": download_from_huggingface,44    "ModelScope": download_from_modelscope,45}46 47 48def download_models(49    model_id_list: List[Preset_model_id] = [],50    downloading_priority: List[Preset_model_website] = ["ModelScope", "HuggingFace"],51):52    print(f"Downloading models: {model_id_list}")53    downloaded_files = []54    for model_id in model_id_list:55        for website in downloading_priority:56            if model_id in website_to_preset_models[website]:57                for model_id, origin_file_path, local_dir in website_to_preset_models[website][model_id]:58                    # Check if the file is downloaded.59                    file_to_download = os.path.join(local_dir, os.path.basename(origin_file_path))60                    if file_to_download in downloaded_files:61                        continue62                    # Download63                    website_to_download_fn[website](model_id, origin_file_path, local_dir)64                    if os.path.basename(origin_file_path) in os.listdir(local_dir):65                        downloaded_files.append(file_to_download)66    return downloaded_files67