| |
| """Gradio UI for CivitAI/Hugging Face downloads and Hugging Face uploads.""" |
|
|
| import os |
| import shutil |
| import subprocess |
| import threading |
| import time |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
| from typing import Callable, Dict, Iterator, List, Optional |
| from urllib.parse import parse_qs, urlencode, urlparse, urlunparse |
| from urllib.request import Request, urlopen |
|
|
| import gradio as gr |
| from huggingface_hub import HfApi, upload_file |
|
|
| FOLDER_CHOICES = [ |
| "checkpoints", |
| "loras", |
| "clip", |
| "text_encoders", |
| "diffusion_models", |
| "vae", |
| "embeddings", |
| "custom", |
| ] |
|
|
| |
| PREDEFINED_HF_MODEL_SETS: Dict[str, List[Dict[str, str]]] = { |
| "Qwen Edit": [ |
| { |
| "url": "https://huggingface.co/Arunk25/Qwen-Image-Edit-Rapid-AIO-GGUF/resolve/main/v23/Qwen-Rapid-NSFW-v23_Q8_0.gguf?download=true", |
| "folder": "diffusion_models", |
| }, |
| { |
| "url": "https://huggingface.co/mradermacher/Qwen2.5-VL-7B-Instruct-abliterated-GGUF/resolve/main/Qwen2.5-VL-7B-Instruct-abliterated.mmproj-Q8_0.gguf?download=true", |
| "folder": "text_encoders", |
| }, |
| { |
| "url": "https://huggingface.co/mradermacher/Qwen2.5-VL-7B-Instruct-abliterated-GGUF/resolve/main/Qwen2.5-VL-7B-Instruct-abliterated.Q8_0.gguf?download=true", |
| "folder": "text_encoders", |
| }, |
| { |
| "url": "https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/vae/qwen_image_vae.safetensors", |
| "folder": "vae", |
| }, |
| ], |
| "WAN": [ |
| { |
| "url": "https://huggingface.co/jmew1989/CMFUI/resolve/main/wan22EnhancedNSFWCameraPrompt_nsfwFASTMOVEV2Q8L.gguf", |
| "folder": "diffusion_models", |
| }, |
| { |
| "url": "https://huggingface.co/Cassanovason69/Wan22nsfwenhanced/resolve/main/wan22EnhancedNSFWCameraPrompt_nsfwFASTMOVEV2Q8H.gguf", |
| "folder": "diffusion_models", |
| }, |
| { |
| "url": "https://huggingface.co/NSFW-API/NSFW-Wan-UMT5-XXL/resolve/main/nsfw_wan_umt5-xxl_bf16.safetensors?download=true", |
| "folder": "text_encoders", |
| }, |
| { |
| "url": "https://huggingface.co/NSFW-API/NSFW-Wan-UMT5-XXL/resolve/main/nsfw_wan_umt5-xxl_fp8_scaled.safetensors?download=true", |
| "folder": "text_encoders", |
| }, |
| { |
| "url": "https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/resolve/main/split_files/vae/wan_2.1_vae.safetensors?download=true", |
| "folder": "vae", |
| }, |
| ], |
| "Z Image Turbo": [ |
| { |
| "url": "https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors", |
| "folder": "diffusion_models", |
| }, |
| { |
| "url": "https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/text_encoders/qwen_3_4b.safetensors", |
| "folder": "text_encoders", |
| }, |
| { |
| "url": "https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/vae/ae.safetensors", |
| "folder": "vae", |
| }, |
| ],"Flux Klein 9B NSFW": [ |
| { |
| "url": "https://huggingface.co/evag3/pornmaster-flux/resolve/main/pornmasterFlux2Klein_v4TurboFp8.safetensors", |
| "folder": "diffusion_models", |
| }, |
| { |
| "url": "https://huggingface.co/Comfy-Org/flux2-klein-9B/resolve/main/split_files/text_encoders/qwen_3_8b_fp8mixed.safetensors", |
| "folder": "text_encoders", |
| }, |
| { |
| "url": "https://huggingface.co/black-forest-labs/FLUX.2-small-decoder/resolve/main/full_encoder_small_decoder.safetensors", |
| "folder": "vae", |
| }, |
| { |
| "url": "https://huggingface.co/Comfy-Org/flux2-dev/resolve/main/split_files/vae/flux2-vae.safetensors", |
| "folder": "vae", |
| }, |
| ], "Krea 2": [ |
| { |
| "url": "https://huggingface.co/Comfy-Org/Krea-2/resolve/main/diffusion_models/krea2_raw_bf16.safetensors?download=true", |
| "folder": "diffusion_models", |
| }, |
| { |
| "url": "https://huggingface.co/Comfy-Org/Krea-2/resolve/main/vae/qwen_image_vae.safetensors?download=true", |
| "folder": "vae", |
| }, |
| { |
| "url": "https://huggingface.co/artsyww/KREA2REALVAE/resolve/main/krea2RealVae_v10.safetensors", |
| "folder": "vae", |
| }, |
| { |
| "url": "https://huggingface.co/ahmed22xa/Huihui-Qwen3-VL-4B-Instruct-abliterated-comfy/resolve/main/Huihui-Qwen3-VL-4B-Instruct-abliterated.safetensors?download=true", |
| "folder": "text_encoders", |
| }, |
| { |
| "url": "https://huggingface.co/Comfy-Org/Krea-2/resolve/main/loras/krea2_turbo_lora_rank_64_bf16.safetensors?download=true", |
| "folder": "loras", |
| }, |
| { |
| "url": "https://huggingface.co/gtaayush010/Krea2-RawLoras/resolve/main/Krea2_TextFusion_Refusal_Reduction.safetensors?download=true", |
| "folder": "loras", |
| }, |
| { |
| "url": "https://huggingface.co/gtaayush010/Krea2-RawLoras/resolve/main/realism_engine_krea2_v3.1.safetensors?download=true", |
| "folder": "loras", |
| }, |
| ] |
| } |
|
|
| GLOBAL_STATES: Dict[str, List[Dict[str, object]]] = {"civitai": [], "hf": [], "hf_set": []} |
| GLOBAL_USED_NAMES: Dict[str, int] = {} |
| GLOBAL_UI_UPDATER: Dict[str, object] = {"civitai": None, "hf": None, "hf_set": None} |
| GLOBAL_LOCK = threading.Lock() |
|
|
| def to_positive_int(value, default: int) -> int: |
| try: |
| parsed = int(value) |
| except (TypeError, ValueError): |
| return default |
| return parsed if parsed > 0 else default |
|
|
|
|
| def resolve_parallel_limit(value, total_jobs: int, default: int) -> int: |
| try: |
| parsed = int(value) |
| except (TypeError, ValueError): |
| parsed = default |
|
|
| if parsed == -1: |
| return max(1, total_jobs) |
| if parsed > 0: |
| return parsed |
| return default |
|
|
|
|
| def human_size(size_bytes: float) -> str: |
| units = ["B", "KB", "MB", "GB", "TB"] |
| value = float(size_bytes) |
| idx = 0 |
| while value >= 1024 and idx < len(units) - 1: |
| value /= 1024 |
| idx += 1 |
| return f"{value:.2f} {units[idx]}" |
|
|
|
|
| def safe_filename_from_url(url: str) -> str: |
| name = os.path.basename(urlparse(url).path) |
| return name or "downloaded_file" |
|
|
|
|
| def parse_urls(urls_text: str) -> List[str]: |
| return [line.strip() for line in urls_text.splitlines() if line.strip()] |
|
|
|
|
| def has_aria2c() -> bool: |
| return shutil.which("aria2c") is not None |
|
|
|
|
| def add_civitai_token(url: str, token: str) -> str: |
| if not token.strip(): |
| return url |
| parsed = urlparse(url) |
| params = parse_qs(parsed.query) |
| params["token"] = [token.strip()] |
| return urlunparse( |
| ( |
| parsed.scheme, |
| parsed.netloc, |
| parsed.path, |
| parsed.params, |
| urlencode(params, doseq=True), |
| parsed.fragment, |
| ) |
| ) |
|
|
|
|
| def resolve_destination(base_dir: str, folder: str, custom_folder: str, filename: str) -> Path: |
| folder_name = custom_folder.strip() if folder == "custom" else folder |
| if not folder_name: |
| raise ValueError("Custom folder cannot be empty when folder is set to custom.") |
| destination_dir = Path(base_dir).expanduser().resolve() / folder_name |
| destination_dir.mkdir(parents=True, exist_ok=True) |
| return destination_dir / filename |
|
|
|
|
| import re |
|
|
| def parse_aria2_size(size_str: str) -> int: |
| size_str = size_str.upper().replace("I", "") |
| if size_str.endswith("B"): |
| size_str = size_str[:-1] |
| units = {"K": 1024, "M": 1024**2, "G": 1024**3, "T": 1024**4} |
| for u, mult in units.items(): |
| if size_str.endswith(u): |
| return int(float(size_str[:-1]) * mult) |
| try: |
| return int(float(size_str)) |
| except ValueError: |
| return 0 |
|
|
| def resolve_url_info(url: str, headers: Optional[Dict[str, str]] = None) -> tuple: |
| req = Request(url, headers=headers or {"User-Agent": "Mozilla/5.0"}, method="HEAD") |
| total = 0 |
| filename = "" |
| try: |
| with urlopen(req) as response: |
| total_raw = response.headers.get("Content-Length") |
| total = int(total_raw) if total_raw and total_raw.isdigit() else 0 |
| |
| cd = response.headers.get("Content-Disposition", "") |
| if "filename=" in cd: |
| parts = cd.split("filename=") |
| if len(parts) > 1: |
| name = parts[1].split(";")[0].strip("\"' ") |
| filename = os.path.basename(name) |
| |
| if not filename and response.url != url: |
| filename = safe_filename_from_url(response.url) |
| except Exception: |
| pass |
| |
| if not filename: |
| filename = safe_filename_from_url(url) |
| |
| return total, filename |
|
|
|
|
| def stream_download( |
| url: str, |
| destination: Path, |
| headers: Optional[Dict[str, str]] = None, |
| on_progress: Optional[Callable[[int, int], None]] = None, |
| ) -> int: |
| req = Request(url, headers=headers or {"User-Agent": "Mozilla/5.0"}) |
|
|
| with urlopen(req) as response: |
| total_raw = response.headers.get("Content-Length") |
| total = int(total_raw) if total_raw and total_raw.isdigit() else 0 |
|
|
| chunk_size = 1024 * 512 |
| downloaded = 0 |
| start = time.time() |
|
|
| with open(destination, "wb") as output: |
| while True: |
| chunk = response.read(chunk_size) |
| if not chunk: |
| break |
|
|
| output.write(chunk) |
| downloaded += len(chunk) |
|
|
| if on_progress is not None: |
| on_progress(downloaded, total) |
|
|
| return downloaded |
|
|
|
|
| def download_with_aria2c( |
| url: str, |
| destination: Path, |
| headers: Optional[Dict[str, str]] = None, |
| per_file_connections: int = 8, |
| on_progress: Optional[Callable[[int, int], None]] = None, |
| ) -> int: |
| cmd = [ |
| "aria2c", |
| url, |
| "--dir", |
| str(destination.parent), |
| "--out", |
| destination.name, |
| "--continue=true", |
| "--allow-overwrite=true", |
| "--summary-interval=1", |
| "--console-log-level=notice", |
| "--max-connection-per-server", |
| str(per_file_connections), |
| "--split", |
| str(per_file_connections), |
| "--min-split-size=1M", |
| ] |
| for key, value in (headers or {}).items(): |
| cmd.extend(["--header", f"{key}: {value}"]) |
|
|
| proc = subprocess.Popen( |
| cmd, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT, |
| text=True, |
| bufsize=1, |
| ) |
|
|
| prog_pattern = re.compile(r"\[#.*?\s+([\d\.]+[KMG]?i?B)/([\d\.]+[KMG]?i?B)\((\d+)%\)") |
|
|
| if proc.stdout is not None: |
| for line in proc.stdout: |
| for match in prog_pattern.finditer(line): |
| dl_size = parse_aria2_size(match.group(1)) |
| tot_size = parse_aria2_size(match.group(2)) |
| if on_progress is not None: |
| on_progress(dl_size, tot_size) |
|
|
| proc.wait() |
|
|
| if proc.returncode != 0: |
| raise RuntimeError(f"aria2c failed with exit code {proc.returncode}") |
|
|
| downloaded = destination.stat().st_size if destination.exists() else 0 |
| if on_progress is not None: |
| on_progress(downloaded, downloaded) |
| return downloaded |
|
|
|
|
| def render_html_progress(states: List[Dict[str, object]]) -> str: |
| completed = sum(1 for state in states if state["status"] in ("done", "error")) |
| active = sum(1 for state in states if state["status"] == "running") |
| total_files = len(states) |
|
|
| downloaded_sum = sum(int(state["downloaded"]) for state in states) |
| total_sum = sum(int(state["total"]) for state in states) |
|
|
| desc = ( |
| f"Completed {completed}/{total_files} | Active {active} | " |
| f"Downloaded {human_size(downloaded_sum)}" |
| ) |
| if total_sum > 0: |
| ratio = min(downloaded_sum / total_sum, 1.0) |
| desc += f" / {human_size(total_sum)} - {ratio * 100:.1f}%" |
| else: |
| ratio = 0.0 |
|
|
| percent = ratio * 100 |
| html = f""" |
| <div style="padding: 20px; background: #1f2937; border-radius: 8px; color: white; font-family: monospace; display: flex; flex-direction: column; justify-content: center; height: 100%;"> |
| <div style="margin-bottom: 10px; text-align: center; font-size: 14px; font-weight: bold;">{desc}</div> |
| <div style="width: 100%; background: #374151; border-radius: 4px; overflow: hidden; height: 24px;"> |
| <div style="width: {percent}%; background: #f97316; height: 100%; transition: width 0.3s ease;"></div> |
| </div> |
| </div> |
| """ |
| return html |
|
|
|
|
| def build_text_progress_bar(downloaded: int, total: int, width: int = 18) -> str: |
| if total > 0: |
| ratio = max(0.0, min(1.0, downloaded / total)) |
| filled = int(width * ratio) |
| return f"[{'#' * filled}{'-' * (width - filled)}] {ratio * 100:5.1f}%" |
| spinner_width = max(0, width - 1) |
| pulse = downloaded % (spinner_width + 1) if spinner_width > 0 else 0 |
| return f"[{'=' * pulse}{'.' * (spinner_width - pulse)}>] ..." |
|
|
|
|
| def render_batch_status_lines( |
| states: List[Dict[str, object]], |
| backend: str, |
| filename_note: str, |
| ) -> str: |
| done_count = sum(1 for state in states if state["status"] == "done") |
| error_count = sum(1 for state in states if state["status"] == "error") |
| running_count = sum(1 for state in states if state["status"] == "running") |
|
|
| lines = [ |
| f"Backend: {backend}", |
| f"Files total: {len(states)} | Running: {running_count} | Success: {done_count} | Failed: {error_count}", |
| ] |
| if filename_note: |
| lines.append(filename_note) |
|
|
| for state in states: |
| downloaded = int(state["downloaded"]) |
| total = int(state["total"]) |
| bar = build_text_progress_bar(downloaded, total) |
| if state["status"] == "done": |
| prefix = "OK " |
| elif state["status"] == "error": |
| prefix = "ERR" |
| elif state["status"] == "running": |
| prefix = "RUN" |
| else: |
| prefix = "..." |
|
|
| line = f"{prefix} {state['name']} | {bar} | {human_size(downloaded)}" |
| if total > 0: |
| line += f"/{human_size(total)}" |
| if state["status"] == "error" and state["error"]: |
| line += f" | {state['error']}" |
| lines.append(line) |
|
|
| return "\n".join(lines) |
|
|
|
|
| def run_parallel_downloads( |
| urls_text: str, |
| base_dir: str, |
| folder: str, |
| custom_folder: str, |
| custom_filename: str, |
| max_parallel_files, |
| per_file_connections, |
| transform_url: Callable[[str], str], |
| build_headers: Callable[[], Dict[str, str]], |
| tab_name: str, |
| ) -> Iterator[tuple]: |
| urls = parse_urls(urls_text) |
| if not urls: |
| yield "", "Please enter one or more URLs (one per line)." |
| return |
|
|
| if not base_dir.strip(): |
| base_dir = os.getcwd() |
|
|
| max_parallel = resolve_parallel_limit(max_parallel_files, len(urls), 3) |
| per_file = to_positive_int(per_file_connections, 8) |
| backend = "aria2c" if has_aria2c() else "python" |
|
|
| if custom_filename.strip() and len(urls) > 1: |
| filename_note = "Note: custom filename is ignored for multi-file batches." |
| else: |
| filename_note = "" |
|
|
| jobs = [] |
| for idx, source_url in enumerate(urls): |
| jobs.append( |
| { |
| "index": idx, |
| "source_url": source_url, |
| "final_url": transform_url(source_url), |
| "custom_filename": custom_filename if len(urls) == 1 else "", |
| } |
| ) |
|
|
| global GLOBAL_STATES, GLOBAL_USED_NAMES, GLOBAL_LOCK, GLOBAL_UI_UPDATER |
|
|
| my_id = object() |
|
|
| with GLOBAL_LOCK: |
| start_idx = len(GLOBAL_STATES.get(tab_name, [])) |
| if tab_name not in GLOBAL_STATES: |
| GLOBAL_STATES[tab_name] = [] |
| |
| for idx in range(len(jobs)): |
| GLOBAL_STATES[tab_name].append({ |
| "status": "pending", |
| "downloaded": 0, |
| "total": 0, |
| "name": f"URL {start_idx + idx + 1}", |
| "error": "", |
| }) |
|
|
| def worker(job: Dict[str, object], global_idx: int) -> None: |
| url = str(job["final_url"]) |
| headers = build_headers() |
|
|
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["status"] = "running" |
|
|
| total_hint, resolved_name = resolve_url_info(url, headers) |
| |
| custom_name = str(job.get("custom_filename") or "").strip() |
| base_name = custom_name if custom_name else resolved_name |
|
|
| with GLOBAL_LOCK: |
| if base_name in GLOBAL_USED_NAMES: |
| GLOBAL_USED_NAMES[base_name] += 1 |
| stem, ext = os.path.splitext(base_name) |
| filename = f"{stem}_{GLOBAL_USED_NAMES[base_name]}{ext}" |
| else: |
| GLOBAL_USED_NAMES[base_name] = 0 |
| filename = base_name |
|
|
| destination = resolve_destination(base_dir, folder, custom_folder, filename) |
|
|
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["total"] = total_hint |
| GLOBAL_STATES[tab_name][global_idx]["name"] = destination.name |
|
|
| def on_file_progress(downloaded: int, total: int) -> None: |
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["downloaded"] = downloaded |
| if total > 0: |
| GLOBAL_STATES[tab_name][global_idx]["total"] = total |
|
|
| try: |
| if backend == "aria2c": |
| download_with_aria2c( |
| url=url, |
| destination=destination, |
| headers=headers, |
| per_file_connections=per_file, |
| on_progress=on_file_progress, |
| ) |
| else: |
| stream_download( |
| url=url, |
| destination=destination, |
| headers=headers, |
| on_progress=on_file_progress, |
| ) |
|
|
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["status"] = "done" |
| except Exception as exc: |
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["status"] = "error" |
| GLOBAL_STATES[tab_name][global_idx]["error"] = str(exc) |
|
|
| with ThreadPoolExecutor(max_workers=min(max_parallel, len(jobs))) as executor: |
| futures = [] |
| for idx, job in enumerate(jobs): |
| global_idx = start_idx + idx |
| futures.append(executor.submit(worker, job, global_idx)) |
|
|
| while any(not future.done() for future in futures): |
| with GLOBAL_LOCK: |
| if GLOBAL_UI_UPDATER.get(tab_name) is None: |
| GLOBAL_UI_UPDATER[tab_name] = my_id |
| is_updater = (GLOBAL_UI_UPDATER[tab_name] == my_id) |
| snapshot = [dict(state) for state in GLOBAL_STATES[tab_name]] |
|
|
| if is_updater: |
| html_str = render_html_progress(snapshot) |
| text_str = render_batch_status_lines(snapshot, backend, filename_note) |
| yield html_str, text_str |
| else: |
| |
| if hasattr(gr, "skip"): |
| try: |
| yield gr.skip() |
| except Exception: |
| pass |
| |
| time.sleep(0.5) |
|
|
| for future in futures: |
| future.result() |
|
|
| with GLOBAL_LOCK: |
| if GLOBAL_UI_UPDATER.get(tab_name) == my_id: |
| GLOBAL_UI_UPDATER[tab_name] = None |
| final_states = [dict(state) for state in GLOBAL_STATES[tab_name]] |
| |
| html_str = render_html_progress(final_states) |
| text_str = render_batch_status_lines(final_states, backend, filename_note) |
| yield html_str, text_str |
|
|
|
|
| def civitai_download( |
| model_urls: str, |
| api_key: str, |
| base_dir: str, |
| folder: str, |
| custom_folder: str, |
| custom_filename: str, |
| max_parallel_files, |
| per_file_connections, |
| ): |
| yield from run_parallel_downloads( |
| urls_text=model_urls, |
| base_dir=base_dir, |
| folder=folder, |
| custom_folder=custom_folder, |
| custom_filename=custom_filename, |
| max_parallel_files=max_parallel_files, |
| per_file_connections=per_file_connections, |
| transform_url=lambda raw_url: add_civitai_token(raw_url, api_key), |
| build_headers=lambda: {"User-Agent": "Mozilla/5.0"}, |
| tab_name="civitai", |
| ) |
|
|
|
|
| def hf_download( |
| file_urls: str, |
| hf_token: str, |
| base_dir: str, |
| folder: str, |
| custom_folder: str, |
| custom_filename: str, |
| max_parallel_files, |
| per_file_connections, |
| ): |
| def headers_builder() -> Dict[str, str]: |
| headers = {"User-Agent": "Mozilla/5.0"} |
| if hf_token.strip(): |
| headers["Authorization"] = f"Bearer {hf_token.strip()}" |
| return headers |
|
|
| yield from run_parallel_downloads( |
| urls_text=file_urls, |
| base_dir=base_dir, |
| folder=folder, |
| custom_folder=custom_folder, |
| custom_filename=custom_filename, |
| max_parallel_files=max_parallel_files, |
| per_file_connections=per_file_connections, |
| transform_url=lambda raw_url: raw_url, |
| build_headers=headers_builder, |
| tab_name="hf", |
| ) |
|
|
|
|
| def predefined_hf_set_download( |
| set_name: str, |
| hf_token: str, |
| base_dir: str, |
| max_parallel_files, |
| per_file_connections, |
| ): |
| selected_set = PREDEFINED_HF_MODEL_SETS.get((set_name or "").strip()) |
| if not selected_set: |
| yield "", "Please select a valid predefined model set." |
| return |
|
|
| if not base_dir.strip(): |
| base_dir = os.getcwd() |
|
|
| max_parallel = resolve_parallel_limit(max_parallel_files, len(selected_set), 3) |
| per_file = to_positive_int(per_file_connections, 8) |
| backend = "aria2c" if has_aria2c() else "python" |
| filename_note = f"Model set: {set_name} ({len(selected_set)} files)" |
| tab_name = "hf_set" |
|
|
| global GLOBAL_STATES, GLOBAL_USED_NAMES, GLOBAL_LOCK, GLOBAL_UI_UPDATER |
| my_id = object() |
|
|
| with GLOBAL_LOCK: |
| start_idx = len(GLOBAL_STATES.get(tab_name, [])) |
| if tab_name not in GLOBAL_STATES: |
| GLOBAL_STATES[tab_name] = [] |
| for idx in range(len(selected_set)): |
| GLOBAL_STATES[tab_name].append( |
| { |
| "status": "pending", |
| "downloaded": 0, |
| "total": 0, |
| "name": f"SET URL {start_idx + idx + 1}", |
| "error": "", |
| } |
| ) |
|
|
| def headers_builder() -> Dict[str, str]: |
| headers = {"User-Agent": "Mozilla/5.0"} |
| if hf_token.strip(): |
| headers["Authorization"] = f"Bearer {hf_token.strip()}" |
| return headers |
|
|
| def worker(item: Dict[str, str], global_idx: int) -> None: |
| url = item["url"] |
| headers = headers_builder() |
|
|
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["status"] = "running" |
|
|
| total_hint, resolved_name = resolve_url_info(url, headers) |
| custom_name = item.get("filename", "").strip() |
| base_name = custom_name if custom_name else resolved_name |
|
|
| with GLOBAL_LOCK: |
| if base_name in GLOBAL_USED_NAMES: |
| GLOBAL_USED_NAMES[base_name] += 1 |
| stem, ext = os.path.splitext(base_name) |
| filename = f"{stem}_{GLOBAL_USED_NAMES[base_name]}{ext}" |
| else: |
| GLOBAL_USED_NAMES[base_name] = 0 |
| filename = base_name |
|
|
| destination = resolve_destination( |
| base_dir=base_dir, |
| folder=item.get("folder", "checkpoints"), |
| custom_folder=item.get("custom_folder", ""), |
| filename=filename, |
| ) |
|
|
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["total"] = total_hint |
| GLOBAL_STATES[tab_name][global_idx]["name"] = destination.name |
|
|
| def on_file_progress(downloaded: int, total: int) -> None: |
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["downloaded"] = downloaded |
| if total > 0: |
| GLOBAL_STATES[tab_name][global_idx]["total"] = total |
|
|
| try: |
| if backend == "aria2c": |
| download_with_aria2c( |
| url=url, |
| destination=destination, |
| headers=headers, |
| per_file_connections=per_file, |
| on_progress=on_file_progress, |
| ) |
| else: |
| stream_download( |
| url=url, |
| destination=destination, |
| headers=headers, |
| on_progress=on_file_progress, |
| ) |
|
|
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["status"] = "done" |
| except Exception as exc: |
| with GLOBAL_LOCK: |
| GLOBAL_STATES[tab_name][global_idx]["status"] = "error" |
| GLOBAL_STATES[tab_name][global_idx]["error"] = str(exc) |
|
|
| with ThreadPoolExecutor(max_workers=min(max_parallel, len(selected_set))) as executor: |
| futures = [] |
| for idx, item in enumerate(selected_set): |
| futures.append(executor.submit(worker, item, start_idx + idx)) |
|
|
| while any(not future.done() for future in futures): |
| with GLOBAL_LOCK: |
| if GLOBAL_UI_UPDATER.get(tab_name) is None: |
| GLOBAL_UI_UPDATER[tab_name] = my_id |
| is_updater = GLOBAL_UI_UPDATER[tab_name] == my_id |
| snapshot = [dict(state) for state in GLOBAL_STATES[tab_name]] |
|
|
| if is_updater: |
| html_str = render_html_progress(snapshot) |
| text_str = render_batch_status_lines(snapshot, backend, filename_note) |
| yield html_str, text_str |
| else: |
| if hasattr(gr, "skip"): |
| try: |
| yield gr.skip(), gr.skip() |
| except Exception: |
| pass |
|
|
| time.sleep(0.5) |
|
|
| for future in futures: |
| future.result() |
|
|
| with GLOBAL_LOCK: |
| if GLOBAL_UI_UPDATER.get(tab_name) == my_id: |
| GLOBAL_UI_UPDATER[tab_name] = None |
| final_states = [dict(state) for state in GLOBAL_STATES[tab_name]] |
|
|
| html_str = render_html_progress(final_states) |
| text_str = render_batch_status_lines(final_states, backend, filename_note) |
| yield html_str, text_str |
|
|
|
|
| def hf_upload( |
| token: str, |
| repo_id: str, |
| local_dir: str, |
| repo_subdir: str, |
| only_safetensors: bool, |
| ): |
| if not token.strip(): |
| return "Please provide a Hugging Face token." |
| if not repo_id.strip(): |
| return "Please provide the Hugging Face repo ID (example: username/my-model)." |
| if not local_dir.strip() or not Path(local_dir).expanduser().is_dir(): |
| return "Please provide a valid local directory." |
|
|
| base_path = Path(local_dir).expanduser().resolve() |
| files = [] |
| for root, _, names in os.walk(base_path): |
| for name in names: |
| if only_safetensors and not name.endswith(".safetensors"): |
| continue |
| files.append(Path(root) / name) |
|
|
| if not files: |
| suffix = " .safetensors" if only_safetensors else "" |
| return f"No{suffix} files found in {base_path}." |
|
|
| api = HfApi(token=token.strip()) |
| remote_prefix = repo_subdir.strip().strip("/") |
|
|
| uploaded = 0 |
| failed = 0 |
| lines = [f"Uploading {len(files)} file(s) to {repo_id.strip()}..."] |
|
|
| for file_path in sorted(files): |
| file_name = file_path.name |
| remote_path = f"{remote_prefix}/{file_name}" if remote_prefix else file_name |
| try: |
| upload_file( |
| path_or_fileobj=str(file_path), |
| path_in_repo=remote_path, |
| repo_id=repo_id.strip(), |
| repo_type="model", |
| token=token.strip(), |
| ) |
| uploaded += 1 |
| lines.append(f"✅ {file_name} -> {remote_path}") |
| except Exception as exc: |
| failed += 1 |
| lines.append(f"❌ {file_name} failed: {exc}") |
|
|
| lines.append(f"Done. Uploaded: {uploaded}, Failed: {failed}") |
| return "\n".join(lines) |
|
|
|
|
| def build_ui() -> gr.Blocks: |
| with gr.Blocks(title="Comfy Model Manager") as app: |
| gr.Markdown("# Comfy Model Manager") |
| gr.Markdown("Download models with live progress, and upload to Hugging Face without terminal logs.") |
|
|
| with gr.Tabs(): |
| with gr.Tab("CivitAI Downloader"): |
| gr.Markdown("### CivitAI Download") |
| civitai_url = gr.Textbox( |
| label="Model URLs (one per line)", |
| lines=6, |
| placeholder="https://civitai.com/api/download/models/...", |
| ) |
| civitai_token = gr.Textbox(label="CivitAI API Key (optional)", type="password") |
|
|
| with gr.Row(): |
| civitai_base_dir = gr.Textbox(label="Base Models Directory", value=os.getcwd()) |
| civitai_folder = gr.Dropdown(label="Destination Folder", choices=FOLDER_CHOICES, value="checkpoints") |
|
|
| with gr.Row(): |
| civitai_max_parallel = gr.Number(label="Max Parallel Files (-1 = unlimited)", value=-1, precision=0) |
| civitai_per_file_connections = gr.Number(label="Per-File Connections (aria2)", value=8, precision=0) |
|
|
| civitai_custom_folder = gr.Textbox(label="Custom Folder Name (only if Destination Folder is custom)") |
| civitai_custom_filename = gr.Textbox(label="Custom Filename (optional)") |
| civitai_download_btn = gr.Button("Start CivitAI Batch Download", variant="primary") |
| with gr.Row(): |
| with gr.Column(scale=1): |
| gr.Markdown("#### Global Progress") |
| civitai_progress_html = gr.HTML() |
| with gr.Column(scale=2): |
| civitai_status = gr.Textbox(label="Status", lines=16, interactive=False) |
|
|
| civitai_download_btn.click( |
| fn=civitai_download, |
| inputs=[ |
| civitai_url, |
| civitai_token, |
| civitai_base_dir, |
| civitai_folder, |
| civitai_custom_folder, |
| civitai_custom_filename, |
| civitai_max_parallel, |
| civitai_per_file_connections, |
| ], |
| outputs=[civitai_progress_html, civitai_status], |
| concurrency_limit=None, |
| trigger_mode="multiple", |
| ) |
|
|
| with gr.Tab("Hugging Face Downloader"): |
| gr.Markdown("### Hugging Face Download") |
| hf_url = gr.Textbox( |
| label="File URLs (one per line)", |
| lines=6, |
| placeholder="https://huggingface.co/.../resolve/main/model.safetensors", |
| ) |
| hf_token = gr.Textbox(label="HF Token (optional for private repos)", type="password") |
|
|
| with gr.Row(): |
| hf_base_dir = gr.Textbox(label="Base Models Directory", value="/workspace/swarmui/comfyui/ComfyUI/models/") |
| hf_folder = gr.Dropdown(label="Destination Folder", choices=FOLDER_CHOICES, value="checkpoints") |
|
|
| with gr.Row(): |
| hf_max_parallel = gr.Number(label="Max Parallel Files (-1 = unlimited)", value=-1, precision=0) |
| hf_per_file_connections = gr.Number(label="Per-File Connections (aria2)", value=8, precision=0) |
|
|
| hf_custom_folder = gr.Textbox(label="Custom Folder Name (only if Destination Folder is custom)") |
| hf_custom_filename = gr.Textbox(label="Custom Filename (optional)") |
| hf_download_btn = gr.Button("Start Hugging Face Batch Download", variant="primary") |
| with gr.Row(): |
| with gr.Column(scale=1): |
| gr.Markdown("#### Global Progress") |
| hf_progress_html = gr.HTML() |
| with gr.Column(scale=2): |
| hf_download_status = gr.Textbox(label="Download Status", lines=16, interactive=False) |
|
|
| hf_download_btn.click( |
| fn=hf_download, |
| inputs=[ |
| hf_url, |
| hf_token, |
| hf_base_dir, |
| hf_folder, |
| hf_custom_folder, |
| hf_custom_filename, |
| hf_max_parallel, |
| hf_per_file_connections, |
| ], |
| outputs=[hf_progress_html, hf_download_status], |
| concurrency_limit=None, |
| trigger_mode="multiple", |
| ) |
|
|
| gr.Markdown("### Predefined Model Sets") |
| set_names = sorted(PREDEFINED_HF_MODEL_SETS.keys()) |
| hf_model_set = gr.Dropdown( |
| label="Model Set", |
| choices=set_names, |
| value=set_names[0] if set_names else None, |
| ) |
| hf_set_download_btn = gr.Button("Download Selected Model Set", variant="primary") |
| with gr.Row(): |
| with gr.Column(scale=1): |
| gr.Markdown("#### Set Progress") |
| hf_set_progress_html = gr.HTML() |
| with gr.Column(scale=2): |
| hf_set_status = gr.Textbox(label="Set Download Status", lines=12, interactive=False) |
|
|
| hf_set_download_btn.click( |
| fn=predefined_hf_set_download, |
| inputs=[ |
| hf_model_set, |
| hf_token, |
| hf_base_dir, |
| hf_max_parallel, |
| hf_per_file_connections, |
| ], |
| outputs=[hf_set_progress_html, hf_set_status], |
| concurrency_limit=None, |
| trigger_mode="multiple", |
| ) |
|
|
| with gr.Tab("Hugging Face Uploader"): |
| gr.Markdown("### Hugging Face Upload") |
|
|
| upload_token = gr.Textbox(label="HF Token", type="password") |
| upload_repo = gr.Textbox(label="Repo ID", placeholder="username/my-model") |
|
|
| with gr.Row(): |
| upload_local_dir = gr.Textbox(label="Local Directory", value=os.getcwd()) |
| upload_remote_subdir = gr.Textbox(label="Repo Subfolder (optional)", placeholder="models/v1") |
|
|
| upload_only_safetensors = gr.Checkbox(label="Upload only .safetensors files", value=True) |
| upload_btn = gr.Button("Upload Files", variant="secondary") |
| upload_status = gr.Textbox(label="Upload Status", lines=14, interactive=False) |
|
|
| upload_btn.click( |
| fn=hf_upload, |
| inputs=[ |
| upload_token, |
| upload_repo, |
| upload_local_dir, |
| upload_remote_subdir, |
| upload_only_safetensors, |
| ], |
| outputs=upload_status, |
| ) |
|
|
| return app |
|
|
|
|
| def main() -> None: |
| app = build_ui() |
| app.queue(default_concurrency_limit=None) |
| app.launch(server_name="0.0.0.0", server_port=7860, show_error=True, share=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
