Spaces:
Running
Running
| import gradio as gr | |
| from huggingface_hub import HfApi, HfFolder, hf_hub_download, snapshot_download | |
| import os | |
| from pathlib import Path | |
| import shutil | |
| import gc | |
| import re | |
| import urllib.parse | |
| def get_token(): | |
| try: | |
| token = HfFolder.get_token() | |
| except Exception: | |
| token = "" | |
| return token | |
| def set_token(token): | |
| try: | |
| HfFolder.save_token(token) | |
| except Exception: | |
| print(f"Error: Failed to save token.") | |
| def get_user_agent(): | |
| return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0' | |
| def is_repo_exists(repo_id: str, repo_type: str="model"): | |
| hf_token = get_token() | |
| api = HfApi(token=hf_token) | |
| try: | |
| if api.repo_exists(repo_id=repo_id, repo_type=repo_type, token=hf_token): return True | |
| else: return False | |
| except Exception as e: | |
| print(f"Error: Failed to connect {repo_id} ({repo_type}). {e}") | |
| return True # for safe | |
| MODEL_TYPE_CLASS = { | |
| "diffusers:StableDiffusionPipeline": "SD 1.5", | |
| "diffusers:StableDiffusionXLPipeline": "SDXL", | |
| "diffusers:FluxPipeline": "FLUX", | |
| } | |
| def get_model_type(repo_id: str): | |
| hf_token = get_token() | |
| api = HfApi(token=hf_token) | |
| lora_filename = "pytorch_lora_weights.safetensors" | |
| diffusers_filename = "model_index.json" | |
| default = "SDXL" | |
| try: | |
| if api.file_exists(repo_id=repo_id, filename=lora_filename, token=hf_token): return "LoRA" | |
| if not api.file_exists(repo_id=repo_id, filename=diffusers_filename, token=hf_token): return "None" | |
| model = api.model_info(repo_id=repo_id, token=hf_token) | |
| tags = model.tags | |
| for tag in tags: | |
| if tag in MODEL_TYPE_CLASS.keys(): return MODEL_TYPE_CLASS.get(tag, default) | |
| except Exception: | |
| return default | |
| return default | |
| def list_uniq(l): | |
| return sorted(set(l), key=l.index) | |
| def list_sub(a, b): | |
| return [e for e in a if e not in b] | |
| def is_repo_name(s): | |
| return re.fullmatch(r'^[\w_\-\.]+/[\w_\-\.]+$', s) | |
| def get_hf_url(repo_id: str, repo_type: str="model"): | |
| if repo_type == "dataset": url = f"https://huggingface.co/datasets/{repo_id}" | |
| elif repo_type == "space": url = f"https://huggingface.co/spaces/{repo_id}" | |
| else: url = f"https://huggingface.co/{repo_id}" | |
| return url | |
| def split_hf_url(url: str): | |
| try: | |
| s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets|spaces)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.\w+)(?:\?download=true)?$', url)[0]) | |
| if len(s) < 4: return "", "", "", "" | |
| repo_id = s[1] | |
| if s[0] == "datasets": repo_type = "dataset" | |
| elif s[0] == "spaces": repo_type = "space" | |
| else: repo_type = "model" | |
| subfolder = urllib.parse.unquote(s[2]) if s[2] else None | |
| filename = urllib.parse.unquote(s[3]) | |
| return repo_id, filename, subfolder, repo_type | |
| except Exception as e: | |
| print(e) | |
| def download_hf_file(directory, url, progress=gr.Progress(track_tqdm=True)): | |
| hf_token = get_token() | |
| repo_id, filename, subfolder, repo_type = split_hf_url(url) | |
| try: | |
| print(f"Downloading {url} to {directory}") | |
| if subfolder is not None: path = hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder, repo_type=repo_type, local_dir=directory, token=hf_token) | |
| else: path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type, local_dir=directory, token=hf_token) | |
| return path | |
| except Exception as e: | |
| print(f"Failed to download: {e}") | |
| return None | |
| def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(track_tqdm=True)): # requires aria2, gdown | |
| url = url.strip() | |
| if "drive.google.com" in url: | |
| original_dir = os.getcwd() | |
| os.chdir(directory) | |
| os.system(f"gdown --fuzzy {url}") | |
| os.chdir(original_dir) | |
| elif "huggingface.co" in url: | |
| url = url.replace("?download=true", "") | |
| if "/blob/" in url: url = url.replace("/blob/", "/resolve/") | |
| download_hf_file(directory, url) | |
| elif "civitai.com" in url: | |
| if "?" in url: | |
| url = url.split("?")[0] | |
| if civitai_api_key: | |
| url = url + f"?token={civitai_api_key}" | |
| os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}") | |
| else: | |
| print("You need an API key to download Civitai models.") | |
| else: | |
| os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}") | |
| def get_local_file_list(dir_path): | |
| file_list = [] | |
| for file in Path(dir_path).glob("**/*.*"): | |
| if file.is_file(): | |
| file_path = str(file) | |
| file_list.append(file_path) | |
| return file_list | |
| def get_download_file(temp_dir, url, civitai_key, progress=gr.Progress(track_tqdm=True)): | |
| if not "http" in url and is_repo_name(url) and not Path(url).exists(): | |
| print(f"Use HF Repo: {url}") | |
| new_file = url | |
| elif not "http" in url and Path(url).exists(): | |
| print(f"Use local file: {url}") | |
| new_file = url | |
| elif Path(f"{temp_dir}/{url.split('/')[-1]}").exists(): | |
| print(f"File to download alreday exists: {url}") | |
| new_file = f"{temp_dir}/{url.split('/')[-1]}" | |
| else: | |
| print(f"Start downloading: {url}") | |
| before = get_local_file_list(temp_dir) | |
| try: | |
| download_thing(temp_dir, url.strip(), civitai_key) | |
| except Exception: | |
| print(f"Download failed: {url}") | |
| return "" | |
| after = get_local_file_list(temp_dir) | |
| new_file = list_sub(after, before)[0] if list_sub(after, before) else "" | |
| if not new_file: | |
| print(f"Download failed: {url}") | |
| return "" | |
| print(f"Download completed: {url}") | |
| return new_file | |
| def download_repo(repo_id: str, dir_path: str, progress=gr.Progress(track_tqdm=True)): # for diffusers repo | |
| hf_token = get_token() | |
| try: | |
| snapshot_download(repo_id=repo_id, local_dir=dir_path, token=hf_token, allow_patterns=["*.safetensors", "*.bin"], | |
| ignore_patterns=["*.fp16.*", "/*.safetensors", "/*.bin"], force_download=True) | |
| return True | |
| except Exception as e: | |
| print(f"Error: Failed to download {repo_id}. {e}") | |
| gr.Warning(f"Error: Failed to download {repo_id}. {e}") | |
| return False | |
| def upload_repo(repo_id: str, dir_path: str, is_private: bool, progress=gr.Progress(track_tqdm=True)): # for diffusers repo | |
| hf_token = get_token() | |
| api = HfApi(token=hf_token) | |
| try: | |
| progress(0, desc="Start uploading...") | |
| api.create_repo(repo_id=repo_id, token=hf_token, private=is_private, exist_ok=True) | |
| for path in Path(dir_path).glob("*"): | |
| if path.is_dir(): | |
| api.upload_folder(repo_id=repo_id, folder_path=str(path), path_in_repo=path.name, token=hf_token) | |
| elif path.is_file(): | |
| api.upload_file(repo_id=repo_id, path_or_fileobj=str(path), path_in_repo=path.name, token=hf_token) | |
| progress(1, desc="Uploaded.") | |
| return get_hf_url(repo_id, "model") | |
| except Exception as e: | |
| print(f"Error: Failed to upload to {repo_id}. {e}") | |
| return "" | |
| HF_SUBFOLDER_NAME = ["None", "user_repo"] | |
| def duplicate_hf_repo(src_repo: str, dst_repo: str, src_repo_type: str, dst_repo_type: str, | |
| is_private: bool, subfolder_type: str=HF_SUBFOLDER_NAME[1], progress=gr.Progress(track_tqdm=True)): | |
| hf_token = get_token() | |
| api = HfApi(token=hf_token) | |
| try: | |
| if subfolder_type == "user_repo": subfolder = src_repo.replace("/", "_") | |
| else: subfolder = "" | |
| progress(0, desc="Start duplicating...") | |
| api.create_repo(repo_id=dst_repo, repo_type=dst_repo_type, private=is_private, exist_ok=True, token=hf_token) | |
| for path in api.list_repo_files(repo_id=src_repo, repo_type=src_repo_type, token=hf_token): | |
| file = hf_hub_download(repo_id=src_repo, filename=path, repo_type=src_repo_type, token=hf_token) | |
| if not Path(file).exists(): continue | |
| if Path(file).is_dir(): # unused for now | |
| api.upload_folder(repo_id=dst_repo, folder_path=file, path_in_repo=f"{subfolder}/{path}" if subfolder else path, | |
| repo_type=dst_repo_type, token=hf_token) | |
| elif Path(file).is_file(): | |
| api.upload_file(repo_id=dst_repo, path_or_fileobj=file, path_in_repo=f"{subfolder}/{path}" if subfolder else path, | |
| repo_type=dst_repo_type, token=hf_token) | |
| if Path(file).exists(): Path(file).unlink() | |
| progress(1, desc="Duplicated.") | |
| return f"{get_hf_url(dst_repo, dst_repo_type)}/tree/main/{subfolder}" if subfolder else get_hf_url(dst_repo, dst_repo_type) | |
| except Exception as e: | |
| print(f"Error: Failed to duplicate repo {src_repo} to {dst_repo}. {e}") | |
| return "" | |