120 lines
4.8 KiB
Python
120 lines
4.8 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
|
|
import os
|
|
|
|
dir_path = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
class HFDownLoad_Tool:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"repo_id": ("STRING", {"default": "ByteDance/Hyper-SD"}),
|
|
"local_dir": ("STRING", {"default": "models/diffusers"}),
|
|
"ignore_patterns": (["none", "big_files", "safetensors", "bin", "safetensors,bin",
|
|
"pth", "safetensors,bin,pth", "model", "msgpack", "onnx_data", "onnx", ],),
|
|
"max_workers": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1, "display": "slider"}),
|
|
"download_single_file": ("STRING", {"default": ""}),
|
|
"use_default_cache_dir": ("BOOLEAN", {"default": False},),
|
|
"use_hfmirror": ("BOOLEAN", {"default": False},),
|
|
"use_subdir": ("BOOLEAN", {"default": False},),
|
|
"use_localproxy": ("BOOLEAN", {"default": False},),
|
|
"localproxy": ("STRING", {"default": "http://127.0.0.1:1081"},),
|
|
|
|
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("model_path",)
|
|
FUNCTION = "pipeline_tool"
|
|
CATEGORY = "HFDownLoad_Tool"
|
|
|
|
def pipeline_tool(self, repo_id, local_dir,
|
|
ignore_patterns, max_workers, download_single_file, use_default_cache_dir,use_hfmirror,use_subdir,use_localproxy,localproxy):
|
|
if use_default_cache_dir == True:
|
|
cache_dir = None
|
|
local_dir = None
|
|
local_dir_use_symlinks = True
|
|
else:
|
|
local_dir_use_symlinks = False
|
|
path_dir = os.path.dirname(dir_path)
|
|
path = os.path.dirname(path_dir)
|
|
repo_list = repo_id.split('/')
|
|
dir_list = local_dir.split('/')
|
|
if use_subdir ==True:
|
|
path = os.path.join(path, f"{dir_list[0]}", f"{dir_list[1]}", f"{repo_list[0]}", f"{repo_list[1]}")
|
|
else:
|
|
if len(dir_list)==2:
|
|
path = os.path.join(path, f"{dir_list[0]}", f"{dir_list[1]}")
|
|
else:
|
|
path = local_dir
|
|
|
|
cache_dir = os.path.join(path, "cache")
|
|
|
|
local_dir = os.path.normpath(path)
|
|
print(local_dir)
|
|
|
|
if ignore_patterns == "none":
|
|
ignore_patterns = None
|
|
elif ignore_patterns == "big_files":
|
|
ignore_patterns = ["*.safetensors", "*.bin", "*.pth", "*.model", "*.msgpack", "*.onnx_data", "*.onnx",
|
|
"*.gguf", "*.xml"]
|
|
elif ignore_patterns == "safetensors":
|
|
ignore_patterns = ["*.safetensors"]
|
|
elif ignore_patterns == "bin":
|
|
ignore_patterns = ["*.bin"]
|
|
elif ignore_patterns == "safetensors,bin":
|
|
ignore_patterns = ["*.safetensors", "*.bin"]
|
|
elif ignore_patterns == "pth":
|
|
ignore_patterns = ["*.pth"]
|
|
elif ignore_patterns == "safetensors,bin,pth":
|
|
ignore_patterns = ["*.safetensors", "*.bin", "*.pth"]
|
|
elif ignore_patterns == "model":
|
|
ignore_patterns = ["*.model"]
|
|
elif ignore_patterns == "msgpack":
|
|
ignore_patterns = ["*.msgpack"]
|
|
elif ignore_patterns == "onnx_data":
|
|
ignore_patterns = ["*.onnx_data"]
|
|
else:
|
|
ignore_patterns = ["*.onnx"]
|
|
|
|
s = len(download_single_file)
|
|
if use_hfmirror ==True:
|
|
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
|
|
else:
|
|
os.environ['HF_ENDPOINT'] = ''
|
|
if use_localproxy == True:
|
|
os.environ['http_proxy'] = localproxy
|
|
os.environ['https_proxy'] = localproxy
|
|
print("set proxy")
|
|
if s > 0:
|
|
from huggingface_hub import hf_hub_download
|
|
model_path = hf_hub_download(repo_id=repo_id, filename=download_single_file, cache_dir=cache_dir,
|
|
local_dir=local_dir,
|
|
local_dir_use_symlinks=local_dir_use_symlinks, resume_download=True
|
|
)
|
|
return (model_path,)
|
|
else:
|
|
from huggingface_hub import snapshot_download
|
|
model_path = snapshot_download(repo_id=repo_id, cache_dir=cache_dir, local_dir=local_dir,
|
|
local_dir_use_symlinks=local_dir_use_symlinks,
|
|
ignore_patterns=ignore_patterns,
|
|
max_workers=max_workers
|
|
)
|
|
return (model_path,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"HFDownLoad_Tool": HFDownLoad_Tool
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"HFDownLoad_Tool": "HFDownLoad_Tool"
|
|
}
|