utils.py

import gradio as gr
from huggingface_hub import HfApi, hf_hub_download, snapshot_download
try:
    from huggingface_hub.utils import disable_progress_bars as hf_disable_progress_bars, enable_progress_bars as hf_enable_progress_bars, are_progress_bars_disabled as hf_are_progress_bars_disabled
except Exception:
    hf_disable_progress_bars = None
    hf_enable_progress_bars = None
    hf_are_progress_bars_disabled = None
import os
from pathlib import Path
import shutil
import gc
import re
import urllib.parse
import subprocess
import random
import requests
from requests.adapters import HTTPAdapter
from urllib3.util import Retry
import time
from typing import Any
from contextvars import ContextVar
from contextlib import contextmanager


@contextmanager
def suppress_hf_hub_progress_bars():
    if hf_disable_progress_bars is None or hf_enable_progress_bars is None:
        yield
        return
    was_disabled = False
    try:
        was_disabled = bool(hf_are_progress_bars_disabled()) if hf_are_progress_bars_disabled is not None else False
    except Exception:
        was_disabled = False
    if not was_disabled:
        try:
            hf_disable_progress_bars()
        except Exception:
            pass
    try:
        yield
    finally:
        if not was_disabled:
            try:
                hf_enable_progress_bars()
            except Exception:
                pass


CIVITAI_KEY_SWITCH_STATUS_CODES = frozenset([401, 403, 429])


HF_FOLDER_TOKEN = os.getenv("HF_TOKEN", False) or False
HF_FOLDER_TOKEN_CONTEXT = ContextVar("hf_folder_token", default=False)
CIVITAI_KEY_STATUS_CONTEXT = ContextVar("civitai_key_status", default={})


def reset_civitai_key_status(raw_keys: Any=None, source: str=""):
    count = len(parse_civitai_api_keys(raw_keys)) if raw_keys not in [None, "", False] else 0
    status = {"count": count, "active_index": 1 if count else 0, "last_switch_reason": "", "last_status": "", "source": source}
    CIVITAI_KEY_STATUS_CONTEXT.set(status)
    return dict(status)


def update_civitai_key_status(*, raw: Any=None, count: int | None=None, active_index: int | None=None, last_switch_reason: str | None=None, last_status: str | None=None, source: str | None=None):
    status = dict(CIVITAI_KEY_STATUS_CONTEXT.get() or {})
    if raw is not None:
        status["count"] = len(parse_civitai_api_keys(raw))
    if count is not None:
        status["count"] = int(count)
    if active_index is not None:
        status["active_index"] = int(active_index)
    if last_switch_reason is not None:
        status["last_switch_reason"] = str(last_switch_reason)
    if last_status is not None:
        status["last_status"] = str(last_status)
    if source is not None:
        status["source"] = str(source)
    CIVITAI_KEY_STATUS_CONTEXT.set(status)
    return dict(status)


def get_civitai_key_status(raw_keys: Any=None):
    status = dict(CIVITAI_KEY_STATUS_CONTEXT.get() or {})
    if raw_keys is not None:
        status["count"] = len(parse_civitai_api_keys(raw_keys))
    count = int(status.get("count") or 0)
    if count == 0:
        status["active_index"] = 0
    elif int(status.get("active_index") or 0) <= 0:
        status["active_index"] = 1
    return status


def get_token(session_state=None):
    if isinstance(session_state, dict):
        token = session_state.get("hf_token", False)
        if token not in [None, "", False]:
            return token
    token = HF_FOLDER_TOKEN_CONTEXT.get()
    if token not in [None, "", False]:
        return token
    env_token = os.getenv("HF_TOKEN", False) or HF_FOLDER_TOKEN
    return env_token if env_token not in [None, "", False] else False


def set_token(token, session_state=None):
    normalized = token if token else False
    HF_FOLDER_TOKEN_CONTEXT.set(normalized)
    if isinstance(session_state, dict):
        session_state["hf_token"] = normalized
    return normalized


def get_state(state: dict, key: str):
    if key in state.keys(): return state[key]
    else:
        print(f"State '{key}' not found.")
        return None


def set_state(state: dict, key: str, value: Any):
    state[key] = value


def get_user_agent():
    return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0'

CIVITAI_REFERER = 'https://civitai.com/'
RETRYABLE_STATUS_CODES = frozenset([408, 409, 425, 429, 500, 502, 503, 504])
WGET_CIVITAI_OPTIONS = [
    "-c",
    "-nv",
    "--content-disposition",
]
WGET_GENERIC_OPTIONS = [
    "-c",
    "-nv",
]



def build_retry(total: int = 5, backoff_factor: float = 1.0, status_forcelist=RETRYABLE_STATUS_CODES):
    kwargs = {
        'total': total,
        'connect': total,
        'read': total,
        'status': total,
        'backoff_factor': backoff_factor,
        'status_forcelist': status_forcelist,
        'respect_retry_after_header': True,
        'raise_on_status': False,
    }
    try:
        return Retry(allowed_methods=frozenset(['HEAD', 'GET', 'OPTIONS']), **kwargs)
    except TypeError:
        return Retry(method_whitelist=frozenset(['HEAD', 'GET', 'OPTIONS']), **kwargs)


def create_retry_session(total: int = 5, backoff_factor: float = 1.0):
    session = requests.Session()
    adapter = HTTPAdapter(max_retries=build_retry(total=total, backoff_factor=backoff_factor))
    session.mount('https://', adapter)
    session.mount('http://', adapter)
    return session


def is_retryable_exception(exc: Exception):
    msg = f"{type(exc).__name__}: {exc}".lower()
    tokens = [
        '429', '408', '425', '500', '502', '503', '504',
        'timed out', 'timeout', 'connection', 'temporarily unavailable',
        'remote end closed', 'reset by peer', 'server error', 'service unavailable',
        'gateway timeout', 'too many requests', 'network',
    ]
    return any(token in msg for token in tokens)


def format_error_short(exc: Exception):
    return f"{type(exc).__name__}: {exc}"


def parse_civitai_api_keys(raw: Any):
    if raw is None: return []
    if isinstance(raw, (list, tuple, set)):
        parts = []
        for item in raw:
            parts.extend(re.split(r'[\s,;]+', str(item or '')))
    else:
        parts = re.split(r'[\s,;]+', str(raw or ''))
    keys = []
    for part in parts:
        key = str(part).strip()
        if key and key not in keys:
            keys.append(key)
    return keys


def should_switch_civitai_key(status_code: int | None):
    return status_code in CIVITAI_KEY_SWITCH_STATUS_CODES


def sanitize_url_for_log(url: str):
    try:
        parts = urllib.parse.urlsplit(url)
        query = urllib.parse.parse_qsl(parts.query, keep_blank_values=True)
        safe_query = []
        for key, value in query:
            if key.lower() in {"token", "authorization", "auth", "key", "api_key", "x-amz-signature", "x-amz-credential", "x-amz-security-token", "policy", "signature", "key-pair-id"}:
                safe_query.append((key, "***"))
            else:
                safe_query.append((key, value))
        return urllib.parse.urlunsplit((parts.scheme, parts.netloc, parts.path, urllib.parse.urlencode(safe_query), parts.fragment))
    except Exception:
        return url


def sanitize_sensitive_log_text(value: Any):
    text = str(value or "")
    def repl(match):
        return f"{match.group(1)}=***"
    return re.sub(r"(?i)(token|authorization|auth|key|api_key|x-amz-signature|x-amz-credential|x-amz-security-token|policy|signature|key-pair-id)=([^\s&]+)", repl, text)


def log_subprocess_tail(label: str, result: subprocess.CompletedProcess, success_tail: int = 1200, failure_tail: int = 4000):
    tail = success_tail if result.returncode == 0 else failure_tail
    if result.stdout:
        print(f"{label} stdout:\n", sanitize_sensitive_log_text(result.stdout[-tail:]))
    if result.stderr:
        print(f"{label} stderr:\n", sanitize_sensitive_log_text(result.stderr[-tail:]))
    print(f"{label} returncode:", result.returncode)


def ensure_wget_available():
    if shutil.which("wget") is None:
        raise FileNotFoundError("wget is required for direct URL downloads but was not found")



def retry_call(func, attempts: int = 4, base_wait: float = 1.0, action: str = 'operation'):
    last_error = None
    for attempt in range(1, attempts + 1):
        try:
            return func()
        except Exception as e:
            last_error = e
            should_retry = is_retryable_exception(e) and attempt < attempts
            if not should_retry:
                print(f"[fail] action={action} error={format_error_short(e)}")
                raise
            delay = min(8.0, base_wait * (2 ** (attempt - 1))) + random.uniform(0.0, 0.3)
            print(f"[retry] action={action} attempt={attempt}/{attempts} wait={round(delay, 2)}s error={format_error_short(e)}")
            time.sleep(delay)
    if last_error is not None:
        raise last_error

HF_UPLOAD_RETRY_POLICY_CHOICES = ["Auto", "Gentle", "Standard", "Patient"]
HF_UPLOAD_RETRY_POLICIES = {
    "auto": {"attempts": 4, "base_wait": 2.0, "max_wait": 30.0, "post_upload_sleep": 1.0},
    "gentle": {"attempts": 2, "base_wait": 2.0, "max_wait": 12.0, "post_upload_sleep": 1.5},
    "standard": {"attempts": 4, "base_wait": 2.0, "max_wait": 25.0, "post_upload_sleep": 1.0},
    "patient": {"attempts": 5, "base_wait": 3.0, "max_wait": 60.0, "post_upload_sleep": 2.5},
}


def normalize_hf_upload_retry_policy(policy: Any):
    value = str(policy or "Auto").strip().lower()
    if value not in HF_UPLOAD_RETRY_POLICIES:
        return "auto"
    return value


def get_hf_upload_retry_policy_config(policy: Any):
    key = normalize_hf_upload_retry_policy(policy)
    config = dict(HF_UPLOAD_RETRY_POLICIES.get(key) or HF_UPLOAD_RETRY_POLICIES["auto"])
    config["key"] = key
    return config


def _get_exception_headers(exc: Exception):
    response = getattr(exc, "response", None)
    headers = getattr(response, "headers", None) if response is not None else None
    return headers or {}


def parse_hf_retry_delay_from_headers(exc: Exception):
    headers = _get_exception_headers(exc)
    retry_after = str(headers.get("Retry-After", "") or "").strip()
    if retry_after:
        try:
            return max(0.0, float(retry_after)), "retry-after"
        except ValueError:
            pass
    rate_limit = str(headers.get("RateLimit", "") or "")
    match = re.search(r"(?:^|[;,\s])t=(\d+)", rate_limit)
    if match:
        try:
            return max(0.0, float(match.group(1))), "ratelimit-reset"
        except ValueError:
            pass
    return None, ""


def format_hf_rate_limit_hint(exc: Exception):
    headers = _get_exception_headers(exc)
    parts = []
    for key in ("RateLimit", "RateLimit-Policy", "Retry-After"):
        value = str(headers.get(key, "") or "").strip()
        if value:
            parts.append(f"{key}={value}")
    return " ".join(parts)


def is_retryable_hf_upload_exception(exc: Exception):
    msg = f"{type(exc).__name__}: {exc}".lower()
    non_retryable = [
        "401 unauthorized", "403 forbidden", "invalid token", "permission",
        "repo not found", "not found", "invalid repo",
    ]
    if any(token in msg for token in non_retryable):
        return False
    retryable = [
        "429", "408", "425", "500", "502", "503", "504",
        "too many requests", "slow down", "rate limit",
        "timed out", "timeout", "connection", "temporarily unavailable",
        "remote end closed", "reset by peer", "server error", "service unavailable",
        "gateway timeout", "network",
        "lfs-verify", "commit endpoint", "internal error hook",
        "/info/lfs/objects/batch", "lfs/objects/batch",
    ]
    return any(token in msg for token in retryable)


def hf_upload_retry_call(func, policy: Any="Auto", action: str="hf_upload"):
    config = get_hf_upload_retry_policy_config(policy)
    attempts = max(1, int(config.get("attempts", 1)))
    base_wait = float(config.get("base_wait", 2.0))
    max_wait = float(config.get("max_wait", 30.0))
    last_error = None
    for attempt in range(1, attempts + 1):
        try:
            return func()
        except Exception as e:
            last_error = e
            if attempt >= attempts or not is_retryable_hf_upload_exception(e):
                hint = format_hf_rate_limit_hint(e)
                suffix = f" {hint}" if hint else ""
                print(f"[fail] action={action} policy={config['key']} error={format_error_short(e)}{suffix}")
                raise
            header_delay, delay_source = parse_hf_retry_delay_from_headers(e)
            if header_delay is not None:
                if header_delay > max_wait:
                    print(f"[fail] action={action} policy={config['key']} wait_hint={round(header_delay, 2)}s exceeds max_wait={round(max_wait, 2)}s error={format_error_short(e)}")
                    raise
                delay = header_delay
            else:
                delay = min(max_wait, base_wait * (2 ** (attempt - 1))) + random.uniform(0.0, 0.5)
                delay_source = "backoff"
            hint = format_hf_rate_limit_hint(e)
            suffix = f" {hint}" if hint else ""
            print(f"[retry] action={action} policy={config['key']} attempt={attempt}/{attempts} wait={round(delay, 2)}s source={delay_source} error={format_error_short(e)}{suffix}")
            time.sleep(delay)
    if last_error is not None:
        raise last_error



def resolve_civitai_download_url(url: str, civitai_api_key: str, max_tries: int = 3):
    user_agent = get_user_agent()
    keys = parse_civitai_api_keys(civitai_api_key)
    if not keys: keys = [""]
    reset_civitai_key_status(civitai_api_key, source="resolve")
    last_error = None
    for key_index, key in enumerate(keys, start=1):
        update_civitai_key_status(raw=civitai_api_key, active_index=key_index, source="resolve")
        dl_url = f"{url}&token={key}" if "?" in url else f"{url}?token={key}"
        headers = {"User-Agent": user_agent, "Referer": CIVITAI_REFERER}
        allow_key_switch = True
        for i in range(max_tries):
            try:
                r = create_retry_session(total=3, backoff_factor=1.0).get(dl_url, headers=headers, allow_redirects=False, stream=True, timeout=(10, 30))
                status = r.status_code
                location = r.headers.get("Location", "")
                resolved_url = location or r.url
                host = urllib.parse.urlparse(resolved_url).netloc
                print(f"Civitai resolve attempt {i + 1} key {key_index}/{len(keys)}: status={status} host={host}")
                update_civitai_key_status(raw=civitai_api_key, active_index=key_index, last_status=str(status), source="resolve")
                if status in (301, 302, 303, 307, 308) and location:
                    r.close()
                    return resolved_url
                if r.ok and resolved_url:
                    r.close()
                    return resolved_url
                last_error = RuntimeError(f"Failed to resolve Civitai download URL. status={status}")
                allow_key_switch = should_switch_civitai_key(status)
                r.close()
                if allow_key_switch and key_index < len(keys):
                    print(f"Switching Civitai key {key_index}/{len(keys)} after resolve status={status}")
                    update_civitai_key_status(raw=civitai_api_key, active_index=min(key_index + 1, len(keys)), last_switch_reason=f"resolve status={status}", last_status=str(status), source="resolve")
                    break
            except Exception as e:
                last_error = e
                print(f"Civitai resolve attempt {i + 1} key {key_index}/{len(keys)} failed: {format_error_short(e)}")
            if i + 1 < max_tries: time.sleep(2)
        if last_error is None or not allow_key_switch:
            break
    raise last_error if last_error else RuntimeError("Failed to resolve Civitai download URL.")

def is_repo_exists(repo_id: str, repo_type: str="model"):
    hf_token = get_token()
    api = HfApi(token=hf_token)
    try:
        if retry_call(lambda: api.repo_exists(repo_id=repo_id, repo_type=repo_type, token=hf_token), action=f'repo_exists {repo_id}'): return True
        else: return False
    except Exception as e:
        print(f"Error: Failed to connect {repo_id} ({repo_type}). {e}")
        return True # for safe


def ensure_repo(api: HfApi, repo_id: str, repo_type: str="model", is_private: bool=True, hf_token=None):
    if hf_token is None: hf_token = get_token()
    return retry_call(
        lambda: api.create_repo(repo_id=repo_id, repo_type=repo_type, token=hf_token, private=is_private, exist_ok=True),
        action=f'create_repo {repo_id}'
    )


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 retry_call(lambda: api.file_exists(repo_id=repo_id, filename=lora_filename, token=hf_token), action=f'file_exists {repo_id}:{lora_filename}'): return "LoRA"
        if not retry_call(lambda: api.file_exists(repo_id=repo_id, filename=diffusers_filename, token=hf_token), action=f'file_exists {repo_id}:{diffusers_filename}'): return "None"
        model = retry_call(lambda: api.model_info(repo_id=repo_id, token=hf_token), action=f'model_info {repo_id}')
        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 = retry_call(
                lambda: hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder, repo_type=repo_type, local_dir=directory, token=hf_token),
                action=f'hf_hub_download {repo_id}/{subfolder}/{filename}'
            )
        else:
            path = retry_call(
                lambda: hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type, local_dir=directory, token=hf_token),
                action=f'hf_hub_download {repo_id}/{filename}'
            )
        return path
    except Exception as e:
        print(f"Failed to download HF file: {url} {e}")
        return None


def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(track_tqdm=True)): # requires wget, gdown
    try:
        url = url.strip()
        if not url:
            print("Skipping empty download URL.")
            return None
        if "drive.google.com" in url:
            original_dir = os.getcwd()
            try:
                os.chdir(directory)
                subprocess.run(["gdown", "--fuzzy", url], check=False)
            finally:
                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:
            ensure_wget_available()
            keys = parse_civitai_api_keys(civitai_api_key)
            if keys:
                user_agent = get_user_agent()
                last_error = None
                success = False
                for i in range(3):
                    for key_index, key in enumerate(keys, start=1):
                        try:
                            signed_url = resolve_civitai_download_url(url, key, max_tries=1 if len(keys) > 1 else 3)
                        except Exception as e:
                            last_error = e
                            if len(keys) > 1:
                                print(f"Switching Civitai key {key_index}/{len(keys)} after resolve failure: {format_error_short(e)}")
                            continue
                        signed_host = urllib.parse.urlparse(signed_url).netloc
                        print(f"Downloading {sanitize_url_for_log(url)} -> {signed_host} (key {key_index}/{len(keys)})")
                        cmd = [
                            "wget",
                            *WGET_CIVITAI_OPTIONS,
                            "--user-agent", user_agent,
                            "--referer", CIVITAI_REFERER,
                            "-P", directory,
                            signed_url,
                        ]
                        print("Running:", " ".join(cmd[:-1]), "[signed URL omitted]")
                        result = subprocess.run(cmd, capture_output=True, text=True)
                        log_subprocess_tail("wget", result)
                        if result.returncode == 0:
                            last_error = None
                            success = True
                            break
                        last_error = RuntimeError(f"wget failed with rc={result.returncode} host={signed_host}")
                        if len(keys) > 1 and key_index < len(keys):
                            print(f"Switching Civitai key {key_index}/{len(keys)} after wget failure.")
                    if success:
                        break
                    if i + 1 < 3:
                        time.sleep(2)
                if last_error is not None and not success:
                    raise last_error
            else:
                print("You need an API key to download Civitai models.")
        else:
            ensure_wget_available()
            cmd = ["wget", *WGET_GENERIC_OPTIONS, "-P", directory, url]
            print("Running:", " ".join(cmd[:-1]), sanitize_url_for_log(cmd[-1]))
            result = subprocess.run(cmd, capture_output=True, text=True)
            log_subprocess_tail("wget", result)
    except Exception as e:
        print(f"Failed to download: {format_error_short(e)}")


def get_local_file_list(dir_path, recursive=False):
    file_list = []
    pattern = "**/*.*" if recursive else "*/*.*"
    for file in Path(dir_path).glob(pattern):
        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)):
    try:
        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 already exists: {url}")
            new_file = f"{temp_dir}/{url.split('/')[-1]}"
        elif "huggingface.co" in url:
            url = url.replace("?download=true", "")
            if "/blob/" in url: url = url.replace("/blob/", "/resolve/")
            new_file = download_hf_file(temp_dir, url)
        else:
            print(f"Start downloading: {url}")
            recursive = False if "huggingface.co" in url else True
            before = get_local_file_list(temp_dir, recursive)
            download_thing(temp_dir, url.strip(), civitai_key)
            after = get_local_file_list(temp_dir, recursive)
            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: {sanitize_url_for_log(url)}")
        return new_file
    except Exception as e:
        print(f"Download failed: {sanitize_url_for_log(url)} {format_error_short(e)}")
        return ""


def download_repo(repo_id: str, dir_path: str, progress=gr.Progress(track_tqdm=True)): # for diffusers repo
    hf_token = get_token()
    try:
        retry_call(
            lambda: 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),
            action=f'snapshot_download {repo_id}'
        )
        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, is_pr: bool=False, progress=gr.Progress(track_tqdm=True)): # for diffusers repo
    hf_token = get_token()
    api = HfApi(token=hf_token)
    try:
        progress(0, desc="Start uploading...")
        ensure_repo(api, repo_id=repo_id, repo_type='model', is_private=is_private, hf_token=hf_token)
        retry_call(lambda: api.upload_folder(repo_id=repo_id, folder_path=dir_path, path_in_repo="", create_pr=is_pr, token=hf_token), action=f'upload_folder {repo_id}')
        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 ""


def gate_repo(repo_id: str, gated_str: str, repo_type: str="model"):
    hf_token = get_token()
    api = HfApi(token=hf_token)
    try:
        if gated_str == "auto": gated = "auto"
        elif gated_str == "manual": gated = "manual"
        else: gated = False
        api.update_repo_settings(repo_id=repo_id, gated=gated, repo_type=repo_type, token=hf_token)
    except Exception as e:
        print(f"Error: Failed to update settings {repo_id}. {e}")


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...")
        ensure_repo(api, repo_id=dst_repo, repo_type=dst_repo_type, is_private=is_private, hf_token=hf_token)
        for path in retry_call(lambda: api.list_repo_files(repo_id=src_repo, repo_type=src_repo_type, token=hf_token), action=f'list_repo_files {src_repo}'):
            file = retry_call(lambda: hf_hub_download(repo_id=src_repo, filename=path, repo_type=src_repo_type, token=hf_token), action=f'hf_hub_download {src_repo}/{path}')
            if not Path(file).exists(): continue
            if Path(file).is_dir(): # unused for now
                retry_call(lambda: 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), action=f'upload_folder {dst_repo}:{path}')
            elif Path(file).is_file():
                retry_call(lambda: 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), action=f'upload_file {dst_repo}:{path}')
            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 ""


BASE_DIR = str(Path(__file__).resolve().parent.resolve())
CIVITAI_API_KEY = os.environ.get("CIVITAI_API_KEY")


def get_file(url: str, path: str): # requires wget, gdown
    print(f"Downloading {url} to {path}...")
    get_download_file(path, url, CIVITAI_API_KEY)


def git_clone(url: str, path: str, pip: bool=False, addcmd: str=""): # requires git
    os.makedirs(str(Path(BASE_DIR, path)), exist_ok=True)
    os.chdir(Path(BASE_DIR, path))
    print(f"Cloning {url} to {path}...")
    cmd = f'git clone {url}'
    print(f'Running {cmd} at {Path.cwd()}')
    i = subprocess.run(cmd, shell=True).returncode
    if i != 0: print(f'Error occured at running {cmd}')
    p = url.split("/")[-1]
    if not Path(p).exists(): return
    if pip:
        os.chdir(Path(BASE_DIR, path, p))
        cmd = f'pip install -r requirements.txt'
        print(f'Running {cmd} at {Path.cwd()}')
        i = subprocess.run(cmd, shell=True).returncode
        if i != 0: print(f'Error occured at running {cmd}')
    if addcmd:
        os.chdir(Path(BASE_DIR, path, p))
        cmd = addcmd
        print(f'Running {cmd} at {Path.cwd()}')
        i = subprocess.run(cmd, shell=True).returncode
        if i != 0: print(f'Error occured at running {cmd}')


def run(cmd: str, timeout: float=0):
    print(f'Running {cmd} at {Path.cwd()}')
    if timeout == 0:
        i = subprocess.run(cmd, shell=True).returncode
        if i != 0: print(f'Error occured at running {cmd}')
    else:
        p = subprocess.Popen(cmd, shell=True)
        time.sleep(timeout)
        p.terminate()
        print(f'Terminated in {timeout} seconds')