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'^[^/,\s\"\']+/[^/,\s\"\']+$', s) def split_hf_url(url: str): try: s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.\w+)(?:\?download=true)?$', url)[0]) if len(s) < 4: return "", "", "", "" repo_id = s[1] repo_type = "dataset" if s[0] == "datasets" else "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 # https://huggingface.co/docs/huggingface_hub/v0.25.1/en/package_reference/file_download#huggingface_hub.snapshot_download def download_repo(repo_id, dir_path, progress=gr.Progress(track_tqdm=True)): 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(new_repo_id, dir_path, is_private, progress=gr.Progress(track_tqdm=True)): hf_token = get_token() api = HfApi(token=hf_token) try: progress(0, desc="Start uploading...") api.create_repo(repo_id=new_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=new_repo_id, folder_path=str(path), path_in_repo=path.name, token=hf_token) elif path.is_file(): api.upload_file(repo_id=new_repo_id, path_or_fileobj=str(path), path_in_repo=path.name, token=hf_token) progress(1, desc="Uploaded.") url = f"https://huggingface.co/{new_repo_id}" except Exception as e: print(f"Error: Failed to upload to {new_repo_id}. {e}") return "" return url