helpers.py

#!/usr/bin/env python3
"""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",
]

# Hardcoded predefined model sets (HF URLs + destination folder per file).
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:
                # keep the generator alive without causing UI conflicts
                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:  # noqa: BLE001
            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()