diff --git "a/modutils.py" "b/modutils.py" --- "a/modutils.py" +++ "b/modutils.py" @@ -13,173 +13,819 @@ from urllib3.util import Retry import urllib.parse import pandas as pd from typing import Any -from huggingface_hub import HfApi, HfFolder, hf_hub_download, snapshot_download +from huggingface_hub import HfApi, hf_hub_download, snapshot_download from translatepy import Translator from unidecode import unidecode import copy from datetime import datetime, timezone, timedelta FILENAME_TIMEZONE = timezone(timedelta(hours=9)) # JST import torch -from safetensors.torch import load_file +from safetensors import safe_open import gc - +import html as html_lib +import subprocess +import tempfile +import time from env import (HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2, HF_MODEL_USER_EX, HF_MODEL_USER_LIKES, DIFFUSERS_FORMAT_LORAS, DIRECTORY_LORAS, HF_READ_TOKEN, HF_TOKEN, CIVITAI_API_KEY) - MODEL_TYPE_DICT = { "diffusers:StableDiffusionPipeline": "SD 1.5", "diffusers:StableDiffusionXLPipeline": "SDXL", "diffusers:FluxPipeline": "FLUX", } +def log_info(message: str): + print(str(message)) + +def log_warning(message: str): + print(str(message)) + +def log_error(message: str): + print(str(message)) def get_user_agent(): return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0' - def to_list(s): return [x.strip() for x in s.split(",") if not s == ""] - def list_uniq(l): return sorted(set(l), key=l.index) - def list_sub(a, b): return [e for e in a if e not in b] - def is_repo_name(s): return re.fullmatch(r'^[^/]+?/[^/]+?$', s) - DEFAULT_STATE = { "show_diffusers_model_list_detail": False, } - def get_state(state: dict, key: str): - if key in state.keys(): return state[key] - elif key in DEFAULT_STATE.keys(): - print(f"State '{key}' not found. Use dedault value.") + if key in state: + return state[key] + if key in DEFAULT_STATE: + log_info(f"State '{key}' not found. Use default value.") return DEFAULT_STATE[key] - else: - print(f"State '{key}' not found.") - return None - + log_warning(f"State '{key}' not found.") + return None def set_state(state: dict, key: str, value: Any): state[key] = value - translator = Translator() def translate_to_en(input: str): try: output = str(translator.translate(input, 'English')) except Exception as e: output = input - print(e) + log_warning(e) return output - def get_local_model_list(dir_path): model_list = [] valid_extensions = ('.ckpt', '.pt', '.pth', '.safetensors', '.bin') - for file in Path(dir_path).glob("*"): + dir_path = Path(dir_path) + for file in dir_path.glob("*"): if file.suffix in valid_extensions: - file_path = str(Path(f"{dir_path}/{file.name}")) + file_path = str(dir_path / file.name) model_list.append(file_path) #print('\033[34mFILE: ' + file_path + '\033[0m') return model_list +HF_FOLDER_TOKEN = "" def get_token(): - try: - token = HfFolder.get_token() - except Exception: - token = "" - return token - + return HF_FOLDER_TOKEN def set_token(token): - try: - HfFolder.save_token(token) - except Exception: - print(f"Error: Failed to save token.") - + global HF_FOLDER_TOKEN + HF_FOLDER_TOKEN = token set_token(HF_TOKEN) +def get_hf_api(token: str = ""): + return HfApi(token=token) if token else HfApi() -def split_hf_url(url: str): +HF_HOST_ALIASES = frozenset({"huggingface.co", "www.huggingface.co", "hf.co"}) + +def parse_hf_file_url(url: str): + raw = str(url or "").strip() + if not raw: + return {} try: - s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.\w+)(?:\?download=true)?$', url)[0]) - if len(s) < 4: return "", "", "", "" - repo_id = s[1] - repo_type = "dataset" if s[0] == "datasets" else "model" - subfolder = urllib.parse.unquote(s[2]) if s[2] else None - filename = urllib.parse.unquote(s[3]) - return repo_id, filename, subfolder, repo_type - except Exception as e: - print(e) + parts = urllib.parse.urlsplit(raw) + except Exception: + return {} + if str(parts.netloc or "").strip().lower() not in HF_HOST_ALIASES: + return {} + + path_segments = [seg for seg in str(parts.path or "").split("/") if seg] + if not path_segments: + return {} + + repo_type = "model" + if path_segments[0] in ["datasets", "spaces"]: + repo_type = "dataset" if path_segments[0] == "datasets" else "space" + path_segments = path_segments[1:] + + if len(path_segments) < 5: + return {} + + namespace, repo_name, action, revision = path_segments[:4] + if action not in ["resolve", "blob"]: + return {} + + file_segments = [urllib.parse.unquote(seg) for seg in path_segments[4:]] + if not file_segments: + return {} + + filename = file_segments[-1] + subfolder = "/".join(file_segments[:-1]) if len(file_segments) > 1 else None + return { + "repo_id": f"{namespace}/{repo_name}", + "filename": filename, + "subfolder": subfolder, + "repo_type": repo_type, + "revision": urllib.parse.unquote(revision), + } +def split_hf_url(url: str): + parsed = parse_hf_file_url(url) + if not parsed: + return "", "", "", "" + return parsed["repo_id"], parsed["filename"], parsed["subfolder"], parsed["repo_type"] def download_hf_file(directory, url, force_filename="", hf_token="", progress=gr.Progress(track_tqdm=True)): - repo_id, filename, subfolder, repo_type = split_hf_url(url) + parsed = parse_hf_file_url(url) + if not parsed: + log_download_error("hf", "parse_url", url=url) + return None + kwargs = {} - if subfolder is not None: kwargs["subfolder"] = subfolder - if force_filename: kwargs["force_filename"] = force_filename + if parsed["subfolder"] is not None: + kwargs["subfolder"] = parsed["subfolder"] + if parsed.get("revision"): + kwargs["revision"] = parsed["revision"] try: - print(f"Start downloading: {url} to {directory}") - path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type, local_dir=directory, token=hf_token, **kwargs) + print( + f"Start HF download: repo={parsed['repo_id']} rev={parsed.get('revision') or '-'} " + f"file={parsed['filename']} to {directory}" + ) + path = hf_hub_download( + repo_id=parsed["repo_id"], + filename=parsed["filename"], + repo_type=parsed["repo_type"], + local_dir=directory, + token=hf_token, + **kwargs, + ) + forced_path = str(Path(directory) / force_filename) if force_filename else "" + if forced_path: + return move_downloaded_file_to_target(path, forced_path) return path except Exception as e: - print(f"Download failed: {url} {e}") + log_download_error("hf", "hub_download", url=url, error=e) + forced_path = str(Path(directory) / force_filename) if force_filename else "" + if forced_path and Path(forced_path).exists(): + return forced_path return None - USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0' +CIVITAI_DEFAULT_ORIGIN = "https://civitai.com" +CIVITAI_CANONICAL_WEB_ORIGIN = CIVITAI_DEFAULT_ORIGIN +CIVITAI_RED_ORIGIN = "https://civitai.red" +CIVITAI_GREEN_HOST_ALIASES = frozenset({"civitai.green", "www.civitai.green"}) +CIVITAI_RED_HOST_ALIASES = frozenset({"civitai.red", "www.civitai.red"}) +CIVITAI_HOST_ALIASES = frozenset({"civitai.com", "www.civitai.com", *CIVITAI_GREEN_HOST_ALIASES, *CIVITAI_RED_HOST_ALIASES}) +CIVITAI_API_ORIGIN_CANDIDATES = (CIVITAI_RED_ORIGIN, CIVITAI_DEFAULT_ORIGIN) +CIVITAI_REFERER = f"{CIVITAI_CANONICAL_WEB_ORIGIN}/" +CIVITAI_RETRY_TOTAL = 5 +CIVITAI_RETRY_BACKOFF = 1.0 +CIVITAI_RESOLVE_RETRY_TOTAL = 4 +CIVITAI_RESOLVE_RETRY_BACKOFF = 0.8 +CIVITAI_STATUS_FORCELIST = [429, 500, 502, 503, 504] +CIVITAI_RESOLVE_TIMEOUT = (7.0, 25.0) +CIVITAI_METADATA_TIMEOUT = (3.0, 15.0) +CIVITAI_SEARCH_TIMEOUT = (3.0, 30.0) +CIVITAI_NEGATIVE_CACHE_LIMIT = 256 +CIVITAI_RESOLVE_CACHE: dict[str, str] = {} +CIVITAI_RESOLVE_NEGATIVE_CACHE: dict[str, str] = {} +CIVITAI_VERSION_JSON_CACHE: dict[str, dict] = {} +CIVITAI_VERSION_NEGATIVE_CACHE: dict[str, str] = {} +CIVITAI_WGET_FRESH_RETRY_LIMIT = 1 +CIVITAI_API_PROBE_TIMEOUT = (3.0, 8.0) +CIVITAI_API_RETRYABLE_STATUSES = frozenset([404, 405, 429, 500, 502, 503, 504]) +CIVITAI_ACTIVE_API_ORIGIN = "" +CIVITAI_ACTIVE_API_BASE = "" + +def create_retry_session(total=CIVITAI_RETRY_TOTAL, backoff_factor=CIVITAI_RETRY_BACKOFF): + session = requests.Session() + retries = Retry(total=total, backoff_factor=backoff_factor, status_forcelist=CIVITAI_STATUS_FORCELIST) + session.mount("https://", HTTPAdapter(max_retries=retries)) + session.mount("http://", HTTPAdapter(max_retries=retries)) + return session + +def cache_put(cache: dict, key: str, value): + key = str(key or "").strip() + if not key: + return + if key in cache: + cache.pop(key, None) + elif len(cache) >= CIVITAI_NEGATIVE_CACHE_LIMIT: + try: + cache.pop(next(iter(cache))) + except Exception: + cache.clear() + cache[key] = value + +def get_civitai_headers(api_key: str = ""): + headers = {'User-Agent': USER_AGENT, 'content-type': 'application/json'} + if api_key: + headers['Authorization'] = f'Bearer {api_key}' + return headers + +def get_civitai_url_parts(url: str): + try: + return urllib.parse.urlsplit(str(url or "").strip()) + except Exception: + return urllib.parse.urlsplit("") + +def sanitize_url_for_log(url: str): + raw = str(url or "").strip() + if not raw: + return raw + parts = get_civitai_url_parts(raw) + if not parts.netloc: + return raw + pairs = [ + (k, v) + for k, v in urllib.parse.parse_qsl(parts.query, keep_blank_values=True) + if str(k).lower() != "token" + ] + query = urllib.parse.urlencode(pairs) + return urllib.parse.urlunsplit((parts.scheme, parts.netloc, parts.path, query, parts.fragment)) + +def canonicalize_civitai_netloc(netloc: str): + host = str(netloc or "").strip().lower() + if host in CIVITAI_GREEN_HOST_ALIASES: + return "civitai.com" + if host == "www.civitai.com": + return "civitai.com" + if host == "www.civitai.red": + return "civitai.red" + return host + +def is_civitai_host(netloc: str): + return canonicalize_civitai_netloc(netloc) in {"civitai.com", "civitai.red"} + +def is_civitai_url(url: str): + return is_civitai_host(get_civitai_url_parts(url).netloc) + +def get_civitai_canonical_web_origin(): + return CIVITAI_CANONICAL_WEB_ORIGIN + +def build_civitai_api_base(origin: str): + raw = str(origin or "").strip().rstrip("/") + return f"{raw}/api/v1" if raw else "" + +def set_civitai_active_api_origin(origin: str): + global CIVITAI_ACTIVE_API_ORIGIN, CIVITAI_ACTIVE_API_BASE + raw = str(origin or "").strip().rstrip("/") + if not raw: + raw = CIVITAI_DEFAULT_ORIGIN + CIVITAI_ACTIVE_API_ORIGIN = raw + CIVITAI_ACTIVE_API_BASE = build_civitai_api_base(raw) + return CIVITAI_ACTIVE_API_BASE + +def probe_civitai_api_origin(session, origin: str, api_key: str = ""): + headers = get_civitai_headers(api_key or CIVITAI_API_KEY) + base_url = build_civitai_api_base(origin) + if not base_url: + return False + response = None + try: + response = session.get( + f"{base_url}/tags", + params={"limit": 1}, + headers=headers, + stream=True, + timeout=CIVITAI_API_PROBE_TIMEOUT, + ) + if not response.ok: + return False + content_type = str(response.headers.get("content-type") or "").lower() + if "json" not in content_type: + return False + data = response.json() + return isinstance(data, dict) and "items" in data + except Exception: + return False + finally: + try: + if response is not None: + response.close() + except Exception: + pass + +def get_civitai_active_api_origin(force_refresh: bool = False, api_key: str = ""): + global CIVITAI_ACTIVE_API_ORIGIN + if CIVITAI_ACTIVE_API_ORIGIN and not force_refresh: + return CIVITAI_ACTIVE_API_ORIGIN + session = create_retry_session(total=2, backoff_factor=0.5) + for origin in CIVITAI_API_ORIGIN_CANDIDATES: + if probe_civitai_api_origin(session, origin, api_key=api_key): + set_civitai_active_api_origin(origin) + print(f"[civitai] selected api origin: {CIVITAI_ACTIVE_API_ORIGIN}") + return CIVITAI_ACTIVE_API_ORIGIN + set_civitai_active_api_origin(CIVITAI_DEFAULT_ORIGIN) + print(f"[civitai] api probe fallback origin: {CIVITAI_ACTIVE_API_ORIGIN}") + return CIVITAI_ACTIVE_API_ORIGIN + +def get_civitai_active_api_base(force_refresh: bool = False, api_key: str = ""): + if CIVITAI_ACTIVE_API_BASE and not force_refresh: + return CIVITAI_ACTIVE_API_BASE + get_civitai_active_api_origin(force_refresh=force_refresh, api_key=api_key) + return CIVITAI_ACTIVE_API_BASE + +def iter_civitai_api_bases(api_key: str = ""): + preferred = get_civitai_active_api_base(api_key=api_key) + bases = [preferred] if preferred else [] + for origin in CIVITAI_API_ORIGIN_CANDIDATES: + base = build_civitai_api_base(origin) + if base and base not in bases: + bases.append(base) + return bases + +def is_retryable_civitai_api_status(status): + try: + return int(status) in CIVITAI_API_RETRYABLE_STATUSES + except Exception: + return False + +def request_civitai_api_response(path: str, params=None, headers=None, timeout=CIVITAI_METADATA_TIMEOUT, + api_key: str = "", session=None, stream: bool = True, allow_not_found: bool = False): + effective_api_key = api_key or CIVITAI_API_KEY + request_headers = headers or get_civitai_headers(effective_api_key) + request_session = session or create_retry_session() + last_response = None + last_url = "" + last_error = None + bases = iter_civitai_api_bases(api_key=effective_api_key) + for idx, base_url in enumerate(bases): + url = f"{base_url}/{str(path or '').lstrip('/')}" + try: + response = request_session.get(url, params=params, headers=request_headers, stream=stream, timeout=timeout) + if response.ok or (allow_not_found and response.status_code == 404): + set_civitai_active_api_origin(base_url.rsplit('/api/v1', 1)[0]) + return response, url + last_response = response + last_url = url + if idx + 1 < len(bases) and is_retryable_civitai_api_status(response.status_code): + try: + response.close() + except Exception: + pass + continue + return response, url + except Exception as e: + last_error = e + last_url = url + if idx + 1 < len(bases): + continue + raise + if last_response is not None: + return last_response, last_url + if last_error is not None: + raise last_error + raise RuntimeError(f"Failed to request Civitai API path: {path}") + +def get_civitai_api_origin_from_url(url: str): + raw = str(url or "").strip() + if not raw: + return "" + if raw.endswith('/api/v1'): + raw = raw.rsplit('/api/v1', 1)[0] + parts = get_civitai_url_parts(raw) + netloc = canonicalize_civitai_netloc(parts.netloc) + if not netloc: + return "" + scheme = parts.scheme or 'https' + return f"{scheme}://{netloc}" + +def request_civitai_api_json(path: str, params=None, headers=None, timeout=CIVITAI_METADATA_TIMEOUT, + api_key: str = "", session=None, stream: bool = True, allow_not_found: bool = False, + non_json_fallback_origin: str = CIVITAI_DEFAULT_ORIGIN): + effective_api_key = api_key or CIVITAI_API_KEY + request_headers = headers or get_civitai_headers(effective_api_key) + request_session = session or create_retry_session() + result, endpoint_url = request_civitai_api_response( + path, + params=params, + headers=request_headers, + timeout=timeout, + api_key=effective_api_key, + session=request_session, + stream=stream, + allow_not_found=allow_not_found, + ) + if allow_not_found and result.status_code == 404: + return None, endpoint_url, result + result.raise_for_status() + try: + return result.json(), endpoint_url, result + except Exception: + current_origin = get_civitai_api_origin_from_url(endpoint_url) + fallback_origin = str(non_json_fallback_origin or CIVITAI_DEFAULT_ORIGIN).strip().rstrip('/') + if fallback_origin and current_origin and current_origin != fallback_origin: + print(f"[retry] Civitai API non-json response from {current_origin}: {str(path or '').lstrip('/')}") + try: + result.close() + except Exception: + pass + fallback_base = build_civitai_api_base(fallback_origin) + fallback_url = f"{fallback_base}/{str(path or '').lstrip('/')}" + fallback_response = request_session.get( + fallback_url, + params=params, + headers=request_headers, + stream=stream, + timeout=timeout, + ) + if allow_not_found and fallback_response.status_code == 404: + return None, fallback_url, fallback_response + fallback_response.raise_for_status() + fallback_json = fallback_response.json() + set_civitai_active_api_origin(fallback_origin) + return fallback_json, fallback_url, fallback_response + raise +try: + get_civitai_active_api_base() +except Exception as e: + print(f"[civitai] startup api probe failed: {type(e).__name__}: {e}") + set_civitai_active_api_origin(CIVITAI_DEFAULT_ORIGIN) -def request_json_data(url): - model_version_id = url.split('/')[-1] - if "?modelVersionId=" in model_version_id: - match = re.search(r'modelVersionId=(\d+)', url) - model_version_id = match.group(1) +def is_civitai_download_api_path(path: str): + return re.match(r'^/api/download/models/\d+$', str(path or "").strip()) is not None - endpoint_url = f"https://civitai.com/api/v1/model-versions/{model_version_id}" +def extract_civitai_model_version_id(url: str): + try: + parts = get_civitai_url_parts(url) + for pattern in [r'^/api/download/models/(\d+)$', r'^/api/v1/model-versions/(\d+)$']: + m = re.match(pattern, str(parts.path or "").strip()) + if m: + return m.group(1) + qs = urllib.parse.parse_qs(parts.query) + for key in ["modelVersionId", "modelversionid", "versionId", "versionid"]: + values = qs.get(key, []) + if not values: + continue + value = str(values[0]).strip() + if value.isdigit(): + return value + except Exception: + return "" + return "" - params = {} - headers = {'User-Agent': USER_AGENT, 'content-type': 'application/json'} - session = requests.Session() - retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]) - session.mount("https://", HTTPAdapter(max_retries=retries)) +def get_civitai_query_filters(url: str): + try: + parts = get_civitai_url_parts(url) + qs = urllib.parse.parse_qs(parts.query) + except Exception: + return {} + filters = {} + for key in ["type", "format", "size", "fp"]: + values = qs.get(key, []) + if values: + filters[key] = str(values[0]).strip() + return filters + +def normalize_civitai_filter_value(key: str, value): + if value is None: + return "" + text = str(value).strip() + if not text: + return "" + if key == "fp": + return text.replace("-", "").replace("_", "").replace(" ", "").lower() + return text.lower() + +def describe_civitai_file_for_log(file_info): + if not isinstance(file_info, dict): + return "" + parts = [] + for key in ["name", "type", "format", "size", "fp"]: + value = file_info.get(key) + if value is None: + continue + text = str(value).strip() + if text: + parts.append(f"{key}={text}") + hashes = file_info.get("hashes") if isinstance(file_info.get("hashes"), dict) else {} + sha256 = str(hashes.get("SHA256") or "").strip() + if sha256: + parts.append(f"sha256={sha256[:12]}...") + return ", ".join(parts) + +def build_civitai_download_query_from_url(url: str): + try: + parts = get_civitai_url_parts(url) + except Exception: + return "" + blocked = {"modelversionid", "versionid"} + pairs = [ + (k, v) + for k, v in urllib.parse.parse_qsl(parts.query, keep_blank_values=True) + if str(k).lower() not in blocked + ] + return urllib.parse.urlencode(pairs) + +def to_civitai_default_download_url(version_id: str, query: str = ""): + if not str(version_id or "").isdigit(): + return "" + base = f"{get_civitai_active_api_origin()}/api/download/models/{version_id}" + return f"{base}?{query}" if query else base + +def normalize_civitai_download_api_url(url: str): + parts = get_civitai_url_parts(url) + if not is_civitai_host(parts.netloc) or not is_civitai_download_api_path(parts.path): + return str(url or "").strip() + host = canonicalize_civitai_netloc(parts.netloc) + return urllib.parse.urlunsplit(("https", host, parts.path, parts.query, "")) + +def extract_first_civitai_download_url_from_html(html: str): + if not html: + return "" + page = html_lib.unescape(str(html)) + patterns = [ + r'https?://(?:www\.)?(?:civitai\.com|civitai\.green|civitai\.red)/api/download/models/\d+[^\s\'\"<>\)\]\}]*', + r'["\'](/api/download/models/\d+[^"\']*)["\']', + ] + for pattern in patterns: + try: + m = re.search(pattern, page, flags=re.IGNORECASE) + except re.error: + m = None + if not m: + continue + candidate = m.group(1) if m.lastindex else m.group(0) + candidate = str(candidate or "").strip("\"'") + if candidate.startswith("/"): + candidate = urllib.parse.urljoin(get_civitai_canonical_web_origin(), candidate) + return normalize_civitai_download_api_url(candidate) + return "" + +def resolve_civitai_model_page_to_download_url(url: str, api_key: str = ""): + raw = str(url or "").strip() + if not raw: + return raw + cached = CIVITAI_RESOLVE_CACHE.get(raw) + if cached: + return cached + if raw in CIVITAI_RESOLVE_NEGATIVE_CACHE: + return raw + parts = get_civitai_url_parts(raw) + if not is_civitai_host(parts.netloc): + return raw + if is_civitai_download_api_path(parts.path): + normalized = normalize_civitai_download_api_url(raw) + cache_put(CIVITAI_RESOLVE_CACHE, raw, normalized) + return normalized + if not re.match(r'^/models/\d+(?:/[^/?#]+)?/?$', parts.path or ""): + return raw + version_id = extract_civitai_model_version_id(raw) + if version_id: + normalized = to_civitai_default_download_url(version_id, query=build_civitai_download_query_from_url(raw)) + cache_put(CIVITAI_RESOLVE_CACHE, raw, normalized) + return normalized + headers = get_civitai_headers(api_key if parts.netloc.lower().endswith("civitai.com") else "") + headers['Referer'] = f"{parts.scheme or 'https'}://{parts.netloc}/" + session = create_retry_session(total=CIVITAI_RESOLVE_RETRY_TOTAL, backoff_factor=CIVITAI_RESOLVE_RETRY_BACKOFF) + try: + r = session.get(raw, headers=headers, timeout=CIVITAI_RESOLVE_TIMEOUT) + if not r.ok: + print(f"Civitai model page resolve failed: {sanitize_url_for_log(raw)} status={r.status_code}") + if r.status_code in [400, 401, 403, 404]: + cache_put(CIVITAI_RESOLVE_NEGATIVE_CACHE, raw, f"status={r.status_code}") + return raw + extracted = extract_first_civitai_download_url_from_html(r.text) + if extracted: + normalized = normalize_civitai_download_api_url(extracted) + cache_put(CIVITAI_RESOLVE_CACHE, raw, normalized) + return normalized + return raw + except Exception as e: + print(f"Failed to resolve Civitai model page URL: {sanitize_url_for_log(raw)} {type(e).__name__}: {sanitize_sensitive_log_text(e)}") + return raw + +def normalize_civitai_input_url(url: str, api_key: str = ""): + raw = str(url or "").strip() + if not raw or not is_civitai_url(raw): + return raw + normalized = resolve_civitai_model_page_to_download_url(raw, api_key=api_key) + if normalized != raw: + print(f"Normalized Civitai URL: {sanitize_url_for_log(raw)} -> {sanitize_url_for_log(normalized)}") + return normalized + +def append_civitai_token(url: str, api_key: str = ""): + raw = str(url or "").strip() + if not raw or not api_key: + return raw + parts = get_civitai_url_parts(raw) + pairs = [(k, v) for k, v in urllib.parse.parse_qsl(parts.query, keep_blank_values=True) if k.lower() != "token"] + pairs.append(("token", api_key)) + query = urllib.parse.urlencode(pairs) + return urllib.parse.urlunsplit((parts.scheme or "https", parts.netloc, parts.path, query, parts.fragment)) + +def get_civitai_request_context(url: str, api_key: str = ""): + raw_url = str(url or "").strip() + normalized_url = normalize_civitai_input_url(raw_url, api_key=api_key) + model_version_id = extract_civitai_model_version_id(normalized_url) or extract_civitai_model_version_id(raw_url) + return { + "raw_url": raw_url, + "normalized_url": normalized_url, + "model_version_id": model_version_id, + "filters": get_civitai_query_filters(raw_url), + } + +def resolve_civitai_download_url(url: str, civitai_api_key: str = "", max_tries: int = 3): + raw = normalize_civitai_download_api_url(str(url or "").strip()) + if not raw: + return raw + headers = get_civitai_headers(civitai_api_key) + headers["Referer"] = CIVITAI_REFERER + dl_url = append_civitai_token(raw, civitai_api_key) + last_error = None + for attempt in range(1, max_tries + 1): + response = None + try: + response = create_retry_session(total=3, backoff_factor=1.0).get( + dl_url, + headers=headers, + allow_redirects=False, + stream=True, + timeout=CIVITAI_RESOLVE_TIMEOUT, + ) + status = int(response.status_code) + location = str(response.headers.get("Location") or "").strip() + resolved_url = str(location or response.url or dl_url).strip() + resolved_host = get_civitai_url_parts(resolved_url).netloc + print( + f"[civitai] resolve signed url attempt={attempt}/{max_tries} status={status} " + f"host={resolved_host or '-'} url={sanitize_url_for_log(raw)}" + ) + if status in (301, 302, 303, 307, 308) and location: + return resolved_url + if response.ok and resolved_url and not is_civitai_host(resolved_host): + return resolved_url + last_error = RuntimeError(f"status={status}") + except Exception as e: + last_error = e + print( + f"[civitai] resolve signed url failed attempt={attempt}/{max_tries} " + f"url={sanitize_url_for_log(raw)} error={type(e).__name__}: {sanitize_sensitive_log_text(e)}" + ) + finally: + try: + if response is not None: + response.close() + except Exception: + pass + if attempt < max_tries: + time.sleep(min(3.0, 0.8 * attempt)) + if last_error is not None: + raise last_error + raise RuntimeError("Failed to resolve Civitai signed download URL") + +def pick_civitai_file_from_version_json(json_data, source_url: str = ""): + files = json_data.get("files", []) if isinstance(json_data, dict) else [] + if not isinstance(files, list) or not files: + return {} + version_id = str((json_data or {}).get("id") or "") + filters = get_civitai_query_filters(source_url) + candidates = [] + fallback = [] + for idx, file_info in enumerate(files): + if not isinstance(file_info, dict): + continue + mismatch = False + matched_filter_count = 0 + for key, expected in filters.items(): + actual = file_info.get(key) + expected_norm = normalize_civitai_filter_value(key, expected) + actual_norm = normalize_civitai_filter_value(key, actual) + if actual_norm: + if actual_norm != expected_norm: + mismatch = True + break + matched_filter_count += 1 + download_url = str(file_info.get("downloadUrl") or "") + score = 0 + if matched_filter_count: + score += matched_filter_count * 3 + if version_id and version_id in download_url: + score += 4 + if download_url: + score += 2 + if file_info.get("name"): + score += 1 + hashes = file_info.get("hashes") if isinstance(file_info.get("hashes"), dict) else {} + if str(hashes.get("SHA256") or "").strip(): + score += 1 + target = fallback if mismatch else candidates + target.append((score, idx, file_info)) + pool = candidates if candidates else fallback + if not pool: + return {} + pool.sort(key=lambda item: (item[0], item[1]), reverse=True) + return dict(pool[0][2]) + +def move_downloaded_file_to_target(downloaded_path: str, target_path: str): + source = Path(str(downloaded_path or "")).expanduser() + target = Path(str(target_path or "")).expanduser() + if not str(target): + return str(source) + if not source.exists(): + return str(target) if target.exists() else str(source) + try: + if source.resolve() == target.resolve(): + return str(target) + except Exception: + pass try: - result = session.get(endpoint_url, params=params, headers=headers, stream=True, timeout=(3.0, 15)) - result.raise_for_status() - json_data = result.json() - return json_data if json_data else None + target.parent.mkdir(parents=True, exist_ok=True) + if target.exists() and target.is_file(): + target.unlink() + shutil.move(str(source), str(target)) + return str(target) except Exception as e: - print(f"Error: {e}") + print(f"HF local rename failed: {source} -> {target} {type(e).__name__}: {sanitize_sensitive_log_text(e)}") + return str(source) + +def request_json_data(url, api_key: str = ""): + effective_api_key = api_key or CIVITAI_API_KEY + context = get_civitai_request_context(url, api_key=effective_api_key) + raw_url = context["raw_url"] + normalized_url = context["normalized_url"] + model_version_id = context["model_version_id"] + if not model_version_id: + print(f"Civitai metadata lookup skipped: modelVersionId not found for {sanitize_url_for_log(raw_url)}") + cache_put(CIVITAI_RESOLVE_NEGATIVE_CACHE, raw_url, "missing_model_version_id") + return None + + cached_json = CIVITAI_VERSION_JSON_CACHE.get(model_version_id) + if cached_json: + return copy.deepcopy(cached_json) + if model_version_id in CIVITAI_VERSION_NEGATIVE_CACHE: return None + endpoint_path = f"/model-versions/{model_version_id}" + headers = get_civitai_headers(effective_api_key) + session = create_retry_session() + + try: + json_data, endpoint_url, result = request_civitai_api_json( + endpoint_path, + headers=headers, + timeout=CIVITAI_METADATA_TIMEOUT, + api_key=effective_api_key, + session=session, + stream=True, + allow_not_found=True, + ) + if result.status_code == 404: + print(f"Civitai metadata lookup status=404: {endpoint_url}") + cache_put(CIVITAI_VERSION_NEGATIVE_CACHE, model_version_id, "status=404") + return None + if not json_data: + print(f"Civitai metadata lookup returned empty JSON: {endpoint_url}") + cache_put(CIVITAI_VERSION_NEGATIVE_CACHE, model_version_id, "empty_json") + return None + cache_put(CIVITAI_VERSION_JSON_CACHE, model_version_id, copy.deepcopy(json_data)) + if normalized_url and normalized_url != raw_url: + cache_put(CIVITAI_RESOLVE_CACHE, raw_url, normalized_url) + return json_data + except Exception as e: + print(f"Civitai metadata lookup failed: {endpoint_url} {type(e).__name__}: {sanitize_sensitive_log_text(e)}") + return None class ModelInformation: - def __init__(self, json_data): + def __init__(self, json_data, source_url: str = ""): + selected_file = pick_civitai_file_from_version_json(json_data, source_url=source_url) self.model_version_id = json_data.get("id", "") self.model_id = json_data.get("modelId", "") - self.download_url = json_data.get("downloadUrl", "") - self.model_url = f"https://civitai.com/models/{self.model_id}?modelVersionId={self.model_version_id}" - self.filename_url = next( - (v.get("name", "") for v in reversed(json_data.get("files", [])) if str(self.model_version_id) in v.get("downloadUrl", "")), "" - ) - self.filename_url = self.filename_url if self.filename_url else "" + self.download_url = selected_file.get("downloadUrl", "") or json_data.get("downloadUrl", "") + self.model_url = f"{get_civitai_canonical_web_origin()}/models/{self.model_id}?modelVersionId={self.model_version_id}" + self.filename_url = selected_file.get("name", "") or "" self.description = json_data.get("description", "") - if self.description is None: self.description = "" + if self.description is None: + self.description = "" self.model_name = json_data.get("model", {}).get("name", "") self.model_type = json_data.get("model", {}).get("type", "") self.nsfw = json_data.get("model", {}).get("nsfw", False) @@ -187,15 +833,198 @@ class ModelInformation: self.images = [img.get("url", "") for img in json_data.get("images", [])] self.example_prompt = json_data.get("trainedWords", [""])[0] if json_data.get("trainedWords") else "" self.original_json = copy.deepcopy(json_data) + self.selected_file = copy.deepcopy(selected_file) - -def retrieve_model_info(url): - json_data = request_json_data(url) +def retrieve_model_info(url, api_key: str = ""): + json_data = request_json_data(url, api_key=api_key) if not json_data: return None - model_descriptor = ModelInformation(json_data) + model_descriptor = ModelInformation(json_data, source_url=url) + filters = get_civitai_query_filters(url) + if filters: + selected_summary = describe_civitai_file_for_log(model_descriptor.selected_file) + if selected_summary: + print(f"Civitai selected file: filters={filters} {selected_summary}") + else: + print(f"Civitai selected file: filters={filters} using model-level downloadUrl") return model_descriptor +def list_downloaded_candidate_files(directory): + try: + return { + str(path.resolve()) + for path in Path(directory).iterdir() + if path.is_file() + } + except Exception: + return set() + +def sanitize_civitai_log_text(text: str): + output = str(text or "") + if not output: + return output + output = re.sub(r"([?&]token=)[^&\s\"']+", r"\1***", output, flags=re.IGNORECASE) + output = re.sub(r"([?&]Authorization=)[^&\s\"']+", r"\1***", output, flags=re.IGNORECASE) + return output + +def sanitize_sensitive_log_text(text): + output = sanitize_civitai_log_text(text) + if not output: + return output + output = re.sub(r"(authorization:\s*bearer\s+)[^\s\"']+", r"\1***", output, flags=re.IGNORECASE) + output = re.sub(r"(bearer\s+)[^\s\"']+", r"\1***", output, flags=re.IGNORECASE) + return output + +def log_download_error(scope: str, kind: str, url: str = "", status=None, error=None, detail: str = ""): + parts = [f"[{scope}] error={kind}"] + if status is not None: + parts.append(f"status={status}") + if url: + parts.append(f"url={sanitize_url_for_log(url)}") + if error is not None: + parts.append(f"exc={type(error).__name__}: {sanitize_sensitive_log_text(error)}") + elif detail: + parts.append(str(detail)) + print(" ".join(parts)) + +def terminate_subprocess_safely(process, label: str = "subprocess"): + if process is None: + return + try: + if process.poll() is not None: + return + process.terminate() + process.wait(timeout=3) + except subprocess.TimeoutExpired: + try: + process.kill() + process.wait(timeout=3) + except Exception as e: + print(f"[{label}] kill failed: {type(e).__name__}: {sanitize_sensitive_log_text(e)}") + except Exception as e: + print(f"[{label}] terminate failed: {type(e).__name__}: {sanitize_sensitive_log_text(e)}") + +def run_subprocess_capture(args, cwd=None, label: str = "subprocess"): + process = subprocess.Popen( + list(args), + cwd=str(cwd) if cwd else None, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + stdout, stderr = process.communicate() + except BaseException: + terminate_subprocess_safely(process, label=label) + raise + output = "\n".join([part for part in [stdout, stderr] if part]).strip() + return int(process.returncode or 0), output + +def build_civitai_wget_args(directory, download_url: str, filename: str = ""): + args = [ + "wget", + "-c", + "-nv", + "--user-agent", USER_AGENT, + "--referer", CIVITAI_REFERER, + ] + if filename: + args.extend(["-O", str(Path(directory) / filename)]) + else: + args.extend(["-P", str(directory)]) + args.append(str(download_url)) + return args + +def run_civitai_wget(directory, download_url: str, filename: str = ""): + args = build_civitai_wget_args(directory, download_url, filename=filename) + return run_subprocess_capture(args, cwd=None, label="civitai-wget") + +def build_generic_wget_args(directory, download_url: str): + return [ + "wget", + "-c", + "-nv", + "-P", str(directory), + str(download_url), + ] + +def run_generic_wget(directory, download_url: str): + args = build_generic_wget_args(directory, download_url) + return run_subprocess_capture(args, cwd=None, label="generic-wget") + +def classify_civitai_download_failure(output_text: str): + text = str(output_text or "") + lower = text.lower() + if "status=403" in lower and "b2.civitai.com" in lower: + return "b2_403" + if "status=403" in lower and "civitai.com/api/download/models/" in lower: + return "api_403" + if "status=403" in lower: + return "http_403" + if "timed out" in lower or "timeout" in lower: + return "timeout" + return "other" + +def cleanup_civitai_download_artifacts(directory, filename: str = ""): + removed = [] + if not filename: + return removed + target = Path(directory) / filename + for candidate in [target]: + try: + if candidate.exists() and candidate.is_file(): + candidate.unlink() + removed.append(str(candidate)) + except Exception as e: + print(f"[civitai] cleanup failed path={candidate} {type(e).__name__}: {e}") + return removed + +def guess_downloaded_file_path(directory, before_files, expected_filename=""): + expected_path = str(Path(directory) / expected_filename) if expected_filename else "" + if expected_path and Path(expected_path).exists(): + return expected_path + + after_files = list_downloaded_candidate_files(directory) + new_files = sorted(list(after_files - set(before_files))) + if len(new_files) == 1: + return new_files[0] + + if expected_filename: + expected_name = str(expected_filename).strip() + stem = Path(expected_name).stem + suffix = Path(expected_name).suffix.lower() + matched = [] + for path_str in new_files: + path_obj = Path(path_str) + if suffix and path_obj.suffix.lower() != suffix: + continue + if stem and (path_obj.stem == stem or path_obj.name == expected_name): + matched.append(path_str) + if len(matched) == 1: + return matched[0] + + return None + +def get_existing_completed_download_path(directory, expected_filename=""): + expected_name = str(expected_filename or "").strip() + if not expected_name: + return "" + + directory_path = Path(directory) + expected_path = directory_path / expected_name + candidate_paths = [expected_path] + + directory_name = directory_path.name.strip() + if directory_name: + legacy_nested_path = directory_path / directory_name / expected_name + if legacy_nested_path != expected_path: + candidate_paths.append(legacy_nested_path) + + for candidate_path in candidate_paths: + if candidate_path.exists() and candidate_path.is_file(): + return str(candidate_path) + + return "" def download_things(directory, url, hf_token="", civitai_api_key="", romanize=False): hf_token = get_token() @@ -203,128 +1032,214 @@ def download_things(directory, url, hf_token="", civitai_api_key="", romanize=Fa downloaded_file_path = None if "drive.google.com" in url: - original_dir = os.getcwd() - os.chdir(directory) - os.system(f"gdown --fuzzy {url}") - os.chdir(original_dir) - elif "huggingface.co" in url: + before_files = list_downloaded_candidate_files(directory) + download_status, download_output = run_subprocess_capture(["gdown", "--fuzzy", str(url)], cwd=directory, label="gdown") + if download_status != 0: + log_download_error("gdown", "command_failed", url=url, status=download_status) + if download_output: + print(sanitize_sensitive_log_text(download_output)) + downloaded_file_path = guess_downloaded_file_path(directory, before_files) + elif "huggingface.co" in url or "hf.co" in url: url = url.replace("?download=true", "") - # url = urllib.parse.quote(url, safe=':/') # fix encoding if "/blob/" in url: url = url.replace("/blob/", "/resolve/") - filename = unidecode(url.split('/')[-1]) if romanize else url.split('/')[-1] - - download_hf_file(directory, url, filename, hf_token) - - downloaded_file_path = os.path.join(directory, filename) - - elif "civitai.com" in url: - + parsed_hf = parse_hf_file_url(url) + filename = parsed_hf.get("filename", "") if parsed_hf else "" + if not filename: + filename = urllib.parse.unquote(url.split('/')[-1]) + if romanize: + filename = unidecode(filename) + + before_files = list_downloaded_candidate_files(directory) + downloaded_file_path = download_hf_file(directory, url, filename, hf_token) or "" + if not downloaded_file_path or not Path(downloaded_file_path).exists(): + downloaded_file_path = guess_downloaded_file_path(directory, before_files, expected_filename=filename) + elif is_civitai_url(url): if not civitai_api_key: - print("\033[91mYou need an API key to download Civitai models.\033[0m") - - model_profile = retrieve_model_info(url) - if model_profile.download_url and model_profile.filename_url: + print("You need an API key to download Civitai models.") + + civitai_context = get_civitai_request_context(url, api_key=civitai_api_key) + normalized_url = civitai_context["normalized_url"] + if normalized_url != url: + print(f"Civitai download URL normalized: {sanitize_url_for_log(url)} -> {sanitize_url_for_log(normalized_url)}") + model_profile = retrieve_model_info(normalized_url, api_key=civitai_api_key) + if model_profile and model_profile.download_url: url = model_profile.download_url - filename = unidecode(model_profile.filename_url) if romanize else model_profile.filename_url + filename = model_profile.filename_url or "" + if filename and romanize: + filename = unidecode(filename) else: - if "?" in url: - url = url.split("?")[0] + url = normalize_civitai_download_api_url(normalized_url) + if not is_civitai_download_api_path(get_civitai_url_parts(url).path): + print(f"Civitai download URL unresolved: {sanitize_url_for_log(normalized_url)}") + return None filename = "" - url_dl = url + f"?token={civitai_api_key}" - print(f"Filename: {filename}") - - param_filename = "" - if filename: - param_filename = f"-o '{filename}'" - - aria2_command = ( - f'aria2c --console-log-level=error --summary-interval=10 -c -x 16 ' - f'-k 1M -s 16 -d "{directory}" {param_filename} "{url_dl}"' - ) - os.system(aria2_command) + signed_url = "" + try: + signed_url = resolve_civitai_download_url(url, civitai_api_key, max_tries=2) + except Exception as e: + print(f"[civitai] failed to resolve signed download url: {sanitize_url_for_log(url)} {type(e).__name__}: {e}") + return None - if param_filename and os.path.exists(os.path.join(directory, filename)): - downloaded_file_path = os.path.join(directory, filename) + signed_host = get_civitai_url_parts(signed_url).netloc + print(f"Filename: {filename}") + print(f"[civitai] resolved signed host={signed_host or '-'} url={sanitize_url_for_log(url)}") + existing_completed_path = get_existing_completed_download_path(directory, expected_filename=filename) + if existing_completed_path: + print(f"[civitai] using existing completed file path={existing_completed_path}") + downloaded_file_path = existing_completed_path + download_status, download_output = 0, "" + else: + before_files = list_downloaded_candidate_files(directory) + download_status, download_output = run_civitai_wget(directory, signed_url, filename=filename) + if download_status != 0: + failure_kind = classify_civitai_download_failure(download_output) + print( + f"[civitai] download failed kind={failure_kind} status={download_status} " + f"filename={filename or '-'} url={sanitize_url_for_log(url)}" + ) + if download_output: + print(sanitize_civitai_log_text(download_output)) + + if failure_kind == "b2_403": + retry_count = 0 + while retry_count < CIVITAI_WGET_FRESH_RETRY_LIMIT and download_status != 0: + retry_count += 1 + removed = cleanup_civitai_download_artifacts(directory, filename=filename) + stale_hint = "yes" if removed or filename else "unknown" + print( + f"[civitai] retrying fresh api/download request after b2_403 " + f"attempt={retry_count}/{CIVITAI_WGET_FRESH_RETRY_LIMIT} stale_resume={stale_hint} " + f"filename={filename or '-'}" + ) + if removed: + print(f"[civitai] removed stale partials: {removed}") + try: + signed_url = resolve_civitai_download_url(url, civitai_api_key, max_tries=2) + signed_host = get_civitai_url_parts(signed_url).netloc + print(f"[civitai] resolved retry signed host={signed_host or '-'} url={sanitize_url_for_log(url)}") + except Exception as e: + print(f"[civitai] retry resolve failed url={sanitize_url_for_log(url)} error={type(e).__name__}: {sanitize_sensitive_log_text(e)}") + break + download_status, download_output = run_civitai_wget(directory, signed_url, filename=filename) + if download_status == 0: + print(f"[civitai] download recovered after fresh retry: {filename or sanitize_url_for_log(url)}") + break + retry_kind = classify_civitai_download_failure(download_output) + print( + f"[civitai] retry failed kind={retry_kind} status={download_status} " + f"filename={filename or '-'} url={sanitize_url_for_log(url)}" + ) + if download_output: + print(sanitize_civitai_log_text(download_output)) + + if download_status != 0: + log_download_error("civitai", "command_failed", url=url, status=download_status) + + if not downloaded_file_path: + downloaded_file_path = guess_downloaded_file_path(directory, before_files, expected_filename=filename) + if not downloaded_file_path: + existing_completed_path = get_existing_completed_download_path(directory, expected_filename=filename) + if existing_completed_path: + print(f"[civitai] using existing completed file path={existing_completed_path}") + downloaded_file_path = existing_completed_path + if not downloaded_file_path: + log_download_error("civitai", "path_unresolved", url=url) else: - os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}") - + before_files = list_downloaded_candidate_files(directory) + download_status, download_output = run_generic_wget(directory, url) + if download_status != 0: + log_download_error("download", "command_failed", url=url, status=download_status) + if download_output: + print(sanitize_sensitive_log_text(download_output)) + downloaded_file_path = guess_downloaded_file_path(directory, before_files) + + if downloaded_file_path and os.path.exists(downloaded_file_path): + print(f"Downloaded file path: {downloaded_file_path}") return downloaded_file_path - def get_download_file(temp_dir, url, civitai_key="", progress=gr.Progress(track_tqdm=True)): + parsed_hf = parse_hf_file_url(url) if ("huggingface.co" in str(url) or "hf.co" in str(url)) else {} + local_name_hint = parsed_hf.get("filename", "") if parsed_hf else urllib.parse.unquote(Path(urllib.parse.urlsplit(str(url or "")).path).name) + cached_local_path = Path(temp_dir) / local_name_hint if local_name_hint else None + if not "http" in url and is_repo_name(url) and not Path(url).exists(): - print(f"Use HF Repo: {url}") + log_info(f"Use HF Repo: {url}") new_file = url elif not "http" in url and Path(url).exists(): - print(f"Use local file: {url}") + log_info(f"Use local file: {url}") new_file = url - elif Path(f"{temp_dir}/{url.split('/')[-1]}").exists(): - print(f"File to download alreday exists: {url}") - new_file = f"{temp_dir}/{url.split('/')[-1]}" + elif cached_local_path and cached_local_path.exists(): + log_info(f"File to download already exists: {url}") + new_file = str(cached_local_path) else: - print(f"Start downloading: {url}") + log_info(f"Start downloading: {url}") before = get_local_model_list(temp_dir) + downloaded_path = "" try: - download_things(temp_dir, url.strip(), HF_TOKEN, civitai_key) + downloaded_path = download_things(temp_dir, url.strip(), HF_TOKEN, civitai_key) or "" except Exception: - print(f"Download failed: {url}") + log_error(f"Download failed: {url}") return "" after = get_local_model_list(temp_dir) - new_file = list_sub(after, before)[0] if list_sub(after, before) else "" + fallback_files = list_sub(after, before) + new_file = downloaded_path if downloaded_path and Path(downloaded_path).exists() else (fallback_files[0] if fallback_files else "") if not new_file: - print(f"Download failed: {url}") + log_error(f"Download failed: {url}") return "" - print(f"Download completed: {url}") + log_info(f"Download completed: {url}") return new_file - -def escape_lora_basename(basename: str): +def normalize_lora_basename(value: str): + basename = str(value or "").strip() return basename.replace(".", "_").replace(" ", "_").replace(",", "") +def escape_lora_basename(basename: str): + return normalize_lora_basename(basename) def to_lora_key(path: str): - return escape_lora_basename(Path(path).stem) - + return normalize_lora_basename(Path(path).stem) def to_lora_path(key: str): if Path(key).is_file(): return key - path = Path(f"{DIRECTORY_LORAS}/{escape_lora_basename(key)}.safetensors") + path = Path(f"{DIRECTORY_LORAS}/{normalize_lora_basename(key)}.safetensors") return str(path) - def safe_float(input): output = 1.0 try: - output = float(input) + value = input.strip() if isinstance(input, str) else input + output = float(value) except Exception: output = 1.0 return output - def valid_model_name(model_name: str): - return model_name.split(" ")[0] + normalized = re.sub(r"\s+", " ", str(model_name or "").strip()) + return normalized.split(" ")[0] if normalized else "" +def create_temp_png_path(prefix: str = "modutils_", suffix: str = ".png"): + fd, temp_path = tempfile.mkstemp(prefix=prefix, suffix=suffix) + os.close(fd) + return str(Path(temp_path).resolve()) def save_images(images: list[Image.Image], metadatas: list[str]): from PIL import PngImagePlugin - import uuid try: output_images = [] for image, metadata in zip(images, metadatas): info = PngImagePlugin.PngInfo() info.add_text("parameters", metadata) - savefile = f"{str(uuid.uuid4())}.png" + savefile = create_temp_png_path(prefix="modimg_") image.save(savefile, "PNG", pnginfo=info) output_images.append(str(Path(savefile).resolve())) return output_images except Exception as e: - print(f"Failed to save image file: {e}") - raise Exception(f"Failed to save image file:") from e - + log_error(f"Failed to save image file: {e}") + raise Exception("Failed to save image file:") from e def save_gallery_images(images, model_name="", progress=gr.Progress(track_tqdm=True)): progress(0, desc="Updating gallery...") @@ -335,19 +1250,25 @@ def save_gallery_images(images, model_name="", progress=gr.Progress(track_tqdm=T for i, image in enumerate(images): filename = f"{basename}{str(i + 1)}.png" oldpath = Path(image[0]) - newpath = oldpath + newpath = oldpath.resolve() if oldpath.exists() else oldpath try: if oldpath.exists(): - newpath = oldpath.resolve().rename(Path(filename).resolve()) + source_path = oldpath.resolve() + target_path = Path(filename).resolve() + if source_path != target_path: + shutil.copy2(str(source_path), str(target_path)) + newpath = target_path + else: + newpath = source_path except Exception as e: - print(e) - finally: + log_error(e) + newpath = oldpath.resolve() if oldpath.exists() else oldpath + finally: output_paths.append(str(newpath)) output_images.append((str(newpath), str(filename))) progress(1, desc="Gallery updated.") return gr.update(value=output_images), gr.update(value=output_paths, visible=True) - def save_gallery_history(images, files, history_gallery, history_files, progress=gr.Progress(track_tqdm=True)): if not images or not files: return gr.update(), gr.update() if not history_gallery: history_gallery = [] @@ -356,42 +1277,49 @@ def save_gallery_history(images, files, history_gallery, history_files, progress output_files = files + history_files return gr.update(value=output_gallery), gr.update(value=output_files, visible=True) - def save_image_history(image, gallery, files, model_name: str, progress=gr.Progress(track_tqdm=True)): if not gallery: gallery = [] if not files: files = [] + temp_path = "" try: basename = f"{model_name.split('/')[-1]}_{datetime.now(FILENAME_TIMEZONE).strftime('%Y%m%d_%H%M%S')}" if image is None or not isinstance(image, (str, Image.Image, np.ndarray, tuple)): return gr.update(), gr.update() filename = f"{basename}.png" if isinstance(image, tuple): image = image[0] - if isinstance(image, str): oldpath = image + if isinstance(image, str): + oldpath = image elif isinstance(image, Image.Image): - oldpath = "temp.png" - image.save(oldpath) + temp_path = create_temp_png_path(prefix="history_") + image.save(temp_path) + oldpath = temp_path elif isinstance(image, np.ndarray): - oldpath = "temp.png" - Image.fromarray(image).convert('RGBA').save(oldpath) + temp_path = create_temp_png_path(prefix="history_") + Image.fromarray(image).convert('RGBA').save(temp_path) + oldpath = temp_path oldpath = Path(oldpath) newpath = oldpath if oldpath.exists(): - shutil.copy(oldpath.resolve(), Path(filename).resolve()) + shutil.copy2(str(oldpath.resolve()), str(Path(filename).resolve())) newpath = Path(filename).resolve() files.insert(0, str(newpath)) gallery.insert(0, (str(newpath), str(filename))) except Exception as e: - print(e) - finally: + log_error(e) + finally: + if temp_path: + try: + safe_clean(temp_path) + except Exception: + pass return gr.update(value=gallery), gr.update(value=files, visible=True) - def download_private_repo(repo_id, dir_path, is_replace): if not HF_READ_TOKEN: return try: snapshot_download(repo_id=repo_id, local_dir=dir_path, allow_patterns=['*.ckpt', '*.pt', '*.pth', '*.safetensors', '*.bin'], token=HF_READ_TOKEN) except Exception as e: - print(f"Error: Failed to download {repo_id}.") - print(e) + log_error(f"Error: Failed to download {repo_id}.") + log_warning(e) return if is_replace: for file in Path(dir_path).glob("*"): @@ -399,30 +1327,29 @@ def download_private_repo(repo_id, dir_path, is_replace): newpath = Path(f'{file.parent.name}/{escape_lora_basename(file.stem)}{file.suffix}') file.resolve().rename(newpath.resolve()) - private_model_path_repo_dict = {} # {"local filepath": "huggingface repo_id", ...} - def get_private_model_list(repo_id, dir_path): global private_model_path_repo_dict - api = HfApi() - if not HF_READ_TOKEN: return [] + if not HF_READ_TOKEN: + return [] + api = get_hf_api(HF_READ_TOKEN) try: - files = api.list_repo_files(repo_id, token=HF_READ_TOKEN) + files = api.list_repo_files(repo_id) except Exception as e: print(f"Error: Failed to list {repo_id}.") - print(e) + log_warning(e) return [] + dir_path_obj = Path(dir_path) model_list = [] for file in files: - path = Path(f"{dir_path}/{file}") - if path.suffix in ['.ckpt', '.pt', '.pth', '.safetensors', '.bin']: - model_list.append(str(path)) + file_path = dir_path_obj / file + if file_path.suffix in ['.ckpt', '.pt', '.pth', '.safetensors', '.bin']: + model_list.append(str(file_path)) for model in model_list: private_model_path_repo_dict[model] = repo_id return model_list - def download_private_file(repo_id, path, is_replace): file = Path(path) newpath = Path(f'{file.parent.name}/{escape_lora_basename(file.stem)}{file.suffix}') if is_replace else file @@ -433,34 +1360,33 @@ def download_private_file(repo_id, path, is_replace): hf_hub_download(repo_id=repo_id, filename=filename, local_dir=dirname, token=HF_READ_TOKEN) except Exception as e: print(f"Error: Failed to download {filename}.") - print(e) + log_warning(e) return if is_replace: file.resolve().rename(newpath.resolve()) - def download_private_file_from_somewhere(path, is_replace): - if not path in private_model_path_repo_dict.keys(): return + if path not in private_model_path_repo_dict: + return repo_id = private_model_path_repo_dict.get(path, None) download_private_file(repo_id, path, is_replace) - model_id_list = [] def get_model_id_list(): global model_id_list - if len(model_id_list) != 0: return model_id_list - api = HfApi() + if model_id_list: return model_id_list + api = get_hf_api() model_ids = [] try: models_likes = [] for author in HF_MODEL_USER_LIKES: - models_likes.extend(api.list_models(author=author, task="text-to-image", cardData=True, sort="likes")) + models_likes.extend(api.list_models(author=author, pipeline_tag="text-to-image", cardData=True, sort="likes")) models_ex = [] for author in HF_MODEL_USER_EX: - models_ex = api.list_models(author=author, task="text-to-image", cardData=True, sort="last_modified") + models_ex = api.list_models(author=author, pipeline_tag="text-to-image", cardData=True, sort="last_modified") except Exception as e: print(f"Error: Failed to list {author}'s models.") - print(e) + log_warning(e) return model_ids for model in models_likes: model_ids.append(model.id) if not model.private else "" @@ -479,43 +1405,62 @@ def get_model_id_list(): model_id_list = model_ids.copy() return model_ids - model_id_list = get_model_id_list() +def is_public_diffusers_model(model) -> bool: + if model is None: + return False + if getattr(model, "private", False) or getattr(model, "gated", False): + return False + tags = getattr(model, "tags", None) + if tags is None: + return False + return "diffusers" in tags + +def get_model_info_tags(model) -> list[str]: + tags = list(getattr(model, "tags", None) or []) + info = [] + for k, v in MODEL_TYPE_DICT.items(): + if k in tags: + info.append(v) + card_data = getattr(model, "card_data", None) + card_tags = getattr(card_data, "tags", None) if card_data else None + if card_tags: + info.extend(list_sub(card_tags, ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl'])) + return info + +def build_tupled_model_name(repo_id: str, info: list[str]) -> str: + info = list(info or []) + if "pony" in info: + info.remove("pony") + return f"{repo_id} (Pony🐴, {', '.join(info)})" + return f"{repo_id} ({', '.join(info)})" def get_t2i_model_info(repo_id: str): - api = HfApi(token=HF_TOKEN) + api = get_hf_api(HF_TOKEN) try: if not is_repo_name(repo_id): return "" model = api.model_info(repo_id=repo_id, timeout=5.0) except Exception as e: print(f"Error: Failed to get {repo_id}'s info.") - print(e) + log_warning(e) return "" - if model.private or model.gated: return "" - tags = model.tags - info = [] + if not is_public_diffusers_model(model): return "" + info = get_model_info_tags(model) url = f"https://huggingface.co/{repo_id}/" - if not 'diffusers' in tags: return "" - for k, v in MODEL_TYPE_DICT.items(): - if k in tags: info.append(v) - if model.card_data and model.card_data.tags: - info.extend(list_sub(model.card_data.tags, ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl'])) info.append(f"DLs: {model.downloads}") info.append(f"likes: {model.likes}") info.append(model.last_modified.strftime("lastmod: %Y-%m-%d")) md = f"Model Info: {', '.join(info)}, [Model Repo]({url})" return gr.update(value=md) - MAX_MODEL_INFO = 100 - def get_tupled_model_list(model_list): if not model_list: return [] #return [(x, x) for x in model_list] # for skipping this function tupled_list = [] - api = HfApi() + api = get_hf_api() for i, repo_id in enumerate(model_list): if i > MAX_MODEL_INFO: tupled_list.append((repo_id, repo_id)) @@ -527,23 +1472,13 @@ def get_tupled_model_list(model_list): print(f"{repo_id}: {e}") tupled_list.append((repo_id, repo_id)) continue - if model.tags is None: continue - tags = model.tags - info = [] - if not 'diffusers' in tags: continue - for k, v in MODEL_TYPE_DICT.items(): - if k in tags: info.append(v) - if model.card_data and model.card_data.tags: - info.extend(list_sub(model.card_data.tags, ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl'])) - if "pony" in info: - info.remove("pony") - name = f"{repo_id} (Pony🐴, {', '.join(info)})" - else: - name = f"{repo_id} ({', '.join(info)})" + if not is_public_diffusers_model(model): + continue + info = get_model_info_tags(model) + name = build_tupled_model_name(repo_id, info) tupled_list.append((name, repo_id)) return tupled_list - private_lora_dict = {} try: with open('lora_dict.json', encoding='utf-8') as f: @@ -560,7 +1495,6 @@ civitai_last_choices = [("", "")] civitai_last_gallery = [] all_lora_list = [] - private_lora_model_list = [] def get_private_lora_model_lists(): global private_lora_model_list @@ -575,60 +1509,20 @@ def get_private_lora_model_lists(): private_lora_model_list = models.copy() return models - private_lora_model_list = get_private_lora_model_lists() - -def get_civitai_info(path): - global civitai_not_exists_list - default = ["", "", "", "", ""] - if path in set(civitai_not_exists_list): return default - if not Path(path).exists(): return None - user_agent = get_user_agent() - headers = {'User-Agent': user_agent, 'content-type': 'application/json'} - base_url = 'https://civitai.com/api/v1/model-versions/by-hash/' - params = {} - session = requests.Session() - retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]) - session.mount("https://", HTTPAdapter(max_retries=retries)) - import hashlib - with open(path, 'rb') as file: - file_data = file.read() - hash_sha256 = hashlib.sha256(file_data).hexdigest() - url = base_url + hash_sha256 - try: - r = session.get(url, params=params, headers=headers, stream=True, timeout=(3.0, 15)) - except Exception as e: - print(e) - return default - if not r.ok: return None - json = r.json() - if not 'baseModel' in json: - civitai_not_exists_list.append(path) - return default - items = [] - items.append(" / ".join(json['trainedWords'])) - items.append(json['baseModel']) - items.append(json['model']['name']) - items.append(f"https://civitai.com/models/{json['modelId']}") - items.append(json['images'][0]['url']) - return items - - def get_lora_model_list(): loras = list_uniq(get_private_lora_model_lists() + DIFFUSERS_FORMAT_LORAS + get_local_model_list(DIRECTORY_LORAS)) loras.insert(0, "None") loras.insert(0, "") return loras - def get_all_lora_list(): global all_lora_list loras = get_lora_model_list() all_lora_list = loras.copy() return loras - def get_all_lora_tupled_list(): global loras_dict models = get_all_lora_list() @@ -639,7 +1533,7 @@ def get_all_lora_tupled_list(): basename = Path(model).stem key = to_lora_key(model) items = None - if key in loras_dict.keys(): + if key in loras_dict: items = loras_dict.get(key, None) else: items = get_civitai_info(model) @@ -655,40 +1549,66 @@ def get_all_lora_tupled_list(): tupled_list.append((name, value)) return tupled_list - def update_lora_dict(path): global loras_dict key = escape_lora_basename(Path(path).stem) - if key in loras_dict.keys(): return + if key in loras_dict: return items = get_civitai_info(path) if items == None: return loras_dict[key] = items +def finalize_downloaded_lora_path(file_path: str, source_url: str = ""): + global loras_url_to_path_dict + if not file_path: + return "" + path = Path(file_path) + if not path.exists(): + return "" + new_path = Path(f'{path.parent.name}/{escape_lora_basename(path.stem)}{path.suffix}') + try: + if path.resolve() != new_path.resolve(): + if new_path.exists(): + new_path = new_path.resolve() + else: + new_path = path.resolve().rename(new_path.resolve()) + else: + new_path = path.resolve() + except Exception as e: + log_error(f"Failed to normalize downloaded lora path: {file_path} {e}") + new_path = path.resolve() + + final_path = str(new_path) + if source_url: + loras_url_to_path_dict[source_url] = final_path + if is_civitai_url(source_url): + normalized_url = get_civitai_request_context(source_url, api_key=CIVITAI_API_KEY).get("normalized_url", "") + if normalized_url: + loras_url_to_path_dict[normalized_url] = final_path + update_lora_dict(final_path) + return final_path def download_lora(dl_urls: str): global loras_url_to_path_dict dl_path = "" - before = get_local_model_list(DIRECTORY_LORAS) - urls = [] - for url in [url.strip() for url in dl_urls.split(',')]: - local_path = f"{DIRECTORY_LORAS}/{url.split('/')[-1]}" - if not Path(local_path).exists(): - download_things(DIRECTORY_LORAS, url, HF_TOKEN, CIVITAI_API_KEY) - urls.append(url) - after = get_local_model_list(DIRECTORY_LORAS) - new_files = list_sub(after, before) - i = 0 - for file in new_files: - path = Path(file) - if path.exists(): - new_path = Path(f'{path.parent.name}/{escape_lora_basename(path.stem)}{path.suffix}') - path.resolve().rename(new_path.resolve()) - loras_url_to_path_dict[urls[i]] = str(new_path) - update_lora_dict(str(new_path)) - dl_path = str(new_path) - i += 1 - return dl_path + for url in [url.strip() for url in dl_urls.split(',') if url.strip()]: + cached_path = loras_url_to_path_dict.get(url, "") + if cached_path and Path(cached_path).exists(): + dl_path = cached_path + continue + + if is_civitai_url(url): + normalized_url = get_civitai_request_context(url, api_key=CIVITAI_API_KEY).get("normalized_url", "") + cached_path = loras_url_to_path_dict.get(normalized_url, "") if normalized_url else "" + if cached_path and Path(cached_path).exists(): + loras_url_to_path_dict[url] = cached_path + dl_path = cached_path + continue + downloaded_path = download_things(DIRECTORY_LORAS, url, HF_TOKEN, CIVITAI_API_KEY) + final_path = finalize_downloaded_lora_path(downloaded_path or "", source_url=url) + if final_path: + dl_path = final_path + return dl_path def copy_lora(path: str, new_path: str): if path == new_path: return new_path @@ -698,14 +1618,13 @@ def copy_lora(path: str, new_path: str): try: shutil.copy(str(cpath.resolve()), str(npath.resolve())) except Exception as e: - print(e) + log_warning(e) return None update_lora_dict(str(npath)) return new_path else: return None - def download_my_lora(dl_urls: str, lora1: str, lora2: str, lora3: str, lora4: str, lora5: str, lora6: str, lora7: str): path = download_lora(dl_urls) if path: @@ -727,12 +1646,11 @@ def download_my_lora(dl_urls: str, lora1: str, lora2: str, lora3: str, lora4: st return gr.update(value=lora1, choices=choices), gr.update(value=lora2, choices=choices), gr.update(value=lora3, choices=choices),\ gr.update(value=lora4, choices=choices), gr.update(value=lora5, choices=choices), gr.update(value=lora6, choices=choices), gr.update(value=lora7, choices=choices) - def get_valid_lora_name(query: str, model_name: str): path = "None" if not query or query == "None": return "None" - if to_lora_key(query) in loras_dict.keys(): return query - if query in loras_url_to_path_dict.keys(): + if to_lora_key(query) in loras_dict: return query + if query in loras_url_to_path_dict: path = loras_url_to_path_dict[query] else: path = to_lora_path(query.strip().split('/')[-1]) @@ -746,17 +1664,21 @@ def get_valid_lora_name(query: str, model_name: str): if dl_file and Path(dl_file).exists(): return dl_file return "None" - def get_valid_lora_path(query: str): path = None - if not query or query == "None": return None - if to_lora_key(query) in loras_dict.keys(): return query - if Path(path).exists(): + if not query or query == "None": + return None + if to_lora_key(query) in loras_dict: + return query + if query in loras_url_to_path_dict: + path = loras_url_to_path_dict[query] + else: + path = to_lora_path(query.strip().split('/')[-1]) + if path and Path(path).exists(): return path else: return None - def get_valid_lora_wt(prompt: str, lora_path: str, lora_wt: float): wt = lora_wt result = re.findall(f'', prompt) @@ -764,6 +1686,10 @@ def get_valid_lora_wt(prompt: str, lora_path: str, lora_wt: float): wt = safe_float(result[0][0]) return wt +LORA_SLOT_COUNT = 7 + +def _choices_only_updates(choices, count=LORA_SLOT_COUNT): + return tuple(gr.update(choices=choices) for _ in range(count)) def set_prompt_loras(prompt, prompt_syntax, model_name, lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt, lora6, lora6_wt, lora7, lora7_wt): if not "Classic" in str(prompt_syntax): return lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt, lora6, lora6_wt, lora7, lora7_wt @@ -800,7 +1726,7 @@ def set_prompt_loras(prompt, prompt_syntax, model_name, lora1, lora1_wt, lora2, wt = result[0][1] path = to_lora_path(key) if not key in loras_dict.keys() or not Path(path).exists(): - path = get_valid_lora_name(path) + path = get_valid_lora_name(path, model_name) if not path or path == "None": continue if path in lora_paths or key in lora_paths: continue @@ -841,7 +1767,6 @@ def set_prompt_loras(prompt, prompt_syntax, model_name, lora1, lora1_wt, lora2, # on7 = True return lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt, lora6, lora6_wt, lora7, lora7_wt - def get_lora_info(lora_path: str): is_valid = False tag = "" @@ -876,7 +1801,6 @@ def get_lora_info(lora_path: str): is_valid = True return is_valid, label, tag, md - def normalize_prompt_list(tags: list[str]): prompts = [] for tag in tags: @@ -885,7 +1809,6 @@ def normalize_prompt_list(tags: list[str]): prompts.append(tag) return prompts - def apply_lora_prompt(prompt: str = "", lora_info: str = ""): if lora_info == "None": return gr.update(value=prompt) tags = prompt.split(",") if prompt else [] @@ -899,7 +1822,6 @@ def apply_lora_prompt(prompt: str = "", lora_info: str = ""): prompt = ", ".join(list_uniq(prompts + lora_prompts) + empty) return gr.update(value=prompt) - def update_loras(prompt, prompt_syntax, lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt, lora6, lora6_wt, lora7, lora7_wt): on1, label1, tag1, md1 = get_lora_info(lora1) on2, label2, tag2, md2 = get_lora_info(lora2) @@ -953,7 +1875,6 @@ def update_loras(prompt, prompt_syntax, lora1, lora1_wt, lora2, lora2_wt, lora3, gr.update(value=lora7, choices=choices), gr.update(value=lora7_wt),\ gr.update(value=tag7, label=label7, visible=on7), gr.update(visible=on7), gr.update(value=md7, visible=on7) - def get_my_lora(link_url, romanize): l_name = "" l_path = "" @@ -994,14 +1915,12 @@ def get_my_lora(link_url, romanize): value=msg_lora ) - def upload_file_lora(files, progress=gr.Progress(track_tqdm=True)): progress(0, desc="Uploading...") file_paths = [file.name for file in files] progress(1, desc="Uploaded.") return gr.update(value=file_paths, visible=True), gr.update() - def move_file_lora(filepaths): for file in filepaths: path = Path(shutil.move(Path(file).resolve(), Path(f"./{DIRECTORY_LORAS}").resolve())) @@ -1011,113 +1930,163 @@ def move_file_lora(filepaths): new_lora_model_list = get_lora_model_list() new_lora_tupled_list = get_all_lora_tupled_list() - - return gr.update( - choices=new_lora_tupled_list, value=new_lora_model_list[-1] - ), gr.update( - choices=new_lora_tupled_list - ), gr.update( - choices=new_lora_tupled_list - ), gr.update( - choices=new_lora_tupled_list - ), gr.update( - choices=new_lora_tupled_list - ), gr.update( - choices=new_lora_tupled_list - ), gr.update( - choices=new_lora_tupled_list - ) - - -CIVITAI_SORT = ["Highest Rated", "Most Downloaded", "Most Liked", "Most Discussed", "Most Collected", "Most Buzz", "Newest"] -CIVITAI_PERIOD = ["AllTime", "Year", "Month", "Week", "Day"] -CIVITAI_BASEMODEL = ["Pony", "Illustrious", "SDXL 1.0", "SD 1.5", "Flux.1 D", "Flux.1 S"] # , "SD 3.5" -CIVITAI_TYPE = ["Checkpoint", "TextualInversion", "Hypernetwork", "AestheticGradient", "LORA", "LoCon", "DoRA", - "Controlnet", "Upscaler", "MotionModule", "VAE", "Poses", "Wildcards", "Workflows", "Other"] -CIVITAI_FILETYPE = ["Model", "VAE", "Config", "Training Data"] - + return gr.update(choices=new_lora_tupled_list, value=new_lora_model_list[-1]), *_choices_only_updates(new_lora_tupled_list, LORA_SLOT_COUNT - 1) def get_civitai_info(path): global civitai_not_exists_list, loras_url_to_path_dict default = ["", "", "", "", ""] - if path in set(civitai_not_exists_list): return default - if not Path(path).exists(): return None - user_agent = get_user_agent() - headers = {'User-Agent': user_agent, 'content-type': 'application/json'} - base_url = 'https://civitai.com/api/v1/model-versions/by-hash/' - params = {} - session = requests.Session() - retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]) - session.mount("https://", HTTPAdapter(max_retries=retries)) + if path in set(civitai_not_exists_list): + return default + if not Path(path).exists(): + return None + + headers = get_civitai_headers(CIVITAI_API_KEY) + endpoint_path = '/model-versions/by-hash/' + session = create_retry_session() + import hashlib + sha256_hash = hashlib.sha256() with open(path, 'rb') as file: - file_data = file.read() - hash_sha256 = hashlib.sha256(file_data).hexdigest() - url = base_url + hash_sha256 + for chunk in iter(lambda: file.read(1024 * 1024), b''): + sha256_hash.update(chunk) + hash_sha256 = sha256_hash.hexdigest() try: - r = session.get(url, params=params, headers=headers, stream=True, timeout=(3.0, 15)) + json_data, url, r = request_civitai_api_json( + endpoint_path + hash_sha256, + headers=headers, + timeout=CIVITAI_METADATA_TIMEOUT, + api_key=CIVITAI_API_KEY, + session=session, + stream=True, + allow_not_found=True, + ) except Exception as e: - print(e) + print(f"Civitai by-hash lookup failed: {path} {type(e).__name__}: {e}") return default - else: - if not r.ok: return None - json = r.json() - if 'baseModel' not in json: + if not r.ok: + print(f"Civitai by-hash lookup status={r.status_code}: {path}") + if r.status_code == 404: civitai_not_exists_list.append(path) return default - items = [] - items.append(" / ".join(json['trainedWords'])) # The words (prompts) used to trigger the model - items.append(json['baseModel']) # Base model (SDXL1.0, Pony, ...) - items.append(json['model']['name']) # The name of the model version - items.append(f"https://civitai.com/models/{json['modelId']}") # The repo url for the model - items.append(json['images'][0]['url']) # The url for a sample image - loras_url_to_path_dict[path] = json['downloadUrl'] # The download url to get the model file for this specific version - return items + return None + if not json_data: + print(f"Civitai by-hash JSON parse failed: {path} empty_json") + return default + if 'baseModel' not in json_data: + civitai_not_exists_list.append(path) + return default + selected_file = pick_civitai_file_from_version_json(json_data, source_url=json_data.get('downloadUrl', '')) + items = [] + items.append(" / ".join(json_data.get('trainedWords', []))) + items.append(json_data.get('baseModel', '')) + items.append(json_data.get('model', {}).get('name', '')) + items.append(f"{get_civitai_canonical_web_origin()}/models/{json_data.get('modelId', '')}") + images = json_data.get('images', []) if isinstance(json_data.get('images'), list) else [] + items.append(images[0].get('url', '') if images else '') + download_url = selected_file.get('downloadUrl', '') or json_data.get('downloadUrl', '') + if download_url: + loras_url_to_path_dict[path] = normalize_civitai_download_api_url(download_url) + return items + +def build_civitai_search_item(item: dict, model: dict) -> dict: + base_model = model.get("baseModel", "") if isinstance(model, dict) else "" + creator = item.get("creator") if isinstance(item, dict) else None + creator_name = creator.get("username", "") if isinstance(creator, dict) else "" + tags = item.get("tags", []) if isinstance(item, dict) else [] + if not isinstance(tags, list): + tags = [] + images = model.get("images", []) if isinstance(model, dict) else [] + image_url = "/home/user/app/null.png" + if isinstance(images, list) and images and isinstance(images[0], dict) and images[0].get("url"): + image_url = images[0]["url"] + page_model_id = item.get("id", "") if isinstance(item, dict) else "" + page_url = f"{get_civitai_canonical_web_origin()}/models/{page_model_id}" if page_model_id else get_civitai_canonical_web_origin() + name = item.get("name", "") if isinstance(item, dict) else "" + model_name = model.get("name", "") if isinstance(model, dict) else "" + desc = model.get("description", "") if isinstance(model, dict) else "" + dl_url = model.get("downloadUrl", "") if isinstance(model, dict) else "" + md = "" + if image_url != "/home/user/app/null.png": + md += f'thumbnail
' + md += ( + f"Model URL: [{page_url}]({page_url})
Model Name: {name}
" + f"Creator: {creator_name}
Tags: {', '.join(tags)}
" + f"Base Model: {base_model}
Description: {desc}" + ) + return { + "name": name, + "creator": creator_name, + "tags": tags, + "model_name": model_name, + "base_model": base_model, + "description": desc, + "img_url": image_url, + "page_url": page_url, + "dl_url": dl_url, + "md": md, + } + +def build_civitai_choice_name(item: dict) -> str: + base_model_name = "Pony🐴" if item.get('base_model') == "Pony" else item.get('base_model', '') + return f"{item.get('name', '')} (for {base_model_name} / By: {item.get('creator', '')} / Tags: {', '.join(item.get('tags', []))})" def search_lora_on_civitai(query: str, allow_model: list[str] = ["Pony", "SDXL 1.0"], limit: int = 100, sort: str = "Highest Rated", period: str = "AllTime", tag: str = "", user: str = "", page: int = 1): - user_agent = get_user_agent() - headers = {'User-Agent': user_agent, 'content-type': 'application/json'} - if CIVITAI_API_KEY: headers['Authorization'] = f'Bearer {{{CIVITAI_API_KEY}}}' - base_url = 'https://civitai.com/api/v1/models' + headers = get_civitai_headers(CIVITAI_API_KEY) + endpoint_path = '/models' params = {'types': ['LORA'], 'sort': sort, 'period': period, 'limit': limit, 'page': int(page), 'nsfw': 'true'} - if query: params["query"] = query - if tag: params["tag"] = tag - if user: params["username"] = user - session = requests.Session() - retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]) - session.mount("https://", HTTPAdapter(max_retries=retries)) + if query: + params["query"] = query + if tag: + params["tag"] = tag + if user: + params["username"] = user + session = create_retry_session() try: - r = session.get(base_url, params=params, headers=headers, stream=True, timeout=(3.0, 30)) + json, _, r = request_civitai_api_json( + endpoint_path, + params=params, + headers=headers, + timeout=CIVITAI_SEARCH_TIMEOUT, + api_key=CIVITAI_API_KEY, + session=session, + stream=True, + ) except Exception as e: - print(e) + print(f"Civitai search failed: query={query!r} page={page} {type(e).__name__}: {e}") return None - else: - if not r.ok: return None - json = r.json() - if 'items' not in json: return None - items = [] - for j in json['items']: - for model in j['modelVersions']: - item = {} - if len(allow_model) != 0 and model['baseModel'] not in set(allow_model): continue - item['name'] = j['name'] - item['creator'] = j['creator']['username'] if 'creator' in j.keys() and 'username' in j['creator'].keys() else "" - item['tags'] = j['tags'] if 'tags' in j.keys() else [] - item['model_name'] = model['name'] if 'name' in model.keys() else "" - item['base_model'] = model['baseModel'] if 'baseModel' in model.keys() else "" - item['description'] = model['description'] if 'description' in model.keys() else "" - item['dl_url'] = model['downloadUrl'] - item['md'] = "" - if 'images' in model.keys() and len(model["images"]) != 0: - item['img_url'] = model["images"][0]["url"] - item['md'] += f'thumbnail
' - else: item['img_url'] = "/home/user/app/null.png" - item['md'] += f'''Model URL: [https://civitai.com/models/{j["id"]}](https://civitai.com/models/{j["id"]})
Model Name: {item["name"]}
- Creator: {item["creator"]}
Tags: {", ".join(item["tags"])}
Base Model: {item["base_model"]}
Description: {item["description"]}''' - items.append(item) - return items + if not r.ok or not json: + print(f"Civitai search status={r.status_code}: query={query!r} page={page}") + return None + if 'items' not in json: + print(f"Civitai search returned no items key: query={query!r} page={page}") + return None + items = [] + allowed_models = set(allow_model) + for j in json['items']: + model_versions = j.get('modelVersions') if isinstance(j, dict) else [] + if not isinstance(model_versions, list): + continue + for model in model_versions: + if not isinstance(model, dict): + continue + base_model = model.get('baseModel', '') + if allowed_models and base_model not in allowed_models: + continue + items.append(build_civitai_search_item(j, model)) + return items + +CIVITAI_SORT = ["Highest Rated", "Most Downloaded", "Most Liked", "Most Discussed", "Most Collected", "Most Buzz", "Newest"] +CIVITAI_PERIOD = ["AllTime", "Year", "Month", "Week", "Day"] +CIVITAI_BASEMODEL_DEFAULT = ["Chroma", "Flux.1 D", "Flux.1 S", "Flux.1 Kontext", "HiDream", "Hunyuan Video", + "Illustrious", "NoobAI", "Other", "Pony", "SD 1.4", "SD 1.5", "SD 1.5 Hyper", + "SD 1.5 LCM", "SD 2.0", "SD 2.1", "SD 2.1 768", "SDXL 0.9", "SDXL 1.0", "SDXL Hyper", + "SDXL Lightning", "Wan Video", "Anima", "Flux.1 Krea", "Flux.2 D", "Flux.2 Klein 4B-base", + "Flux.2 Klein 9B", "Flux.2 Klein 9B-base", "Grok", "LTXV 2.3", "LTXV2", "Qwen", "SDXL 1.0 LCM", + "Wan Video 1.3B t2v", "Wan Video 14B i2v 480p", "Wan Video 14B i2v 720p", "Wan Video 14B t2v", + "Wan Video 2.2 I2V-A14B", "Wan Video 2.2 T2V-A14B", "Wan Video 2.2 TI2V-5B", "ZImageBase", "ZImageTurbo"] +CIVITAI_BASEMODEL = CIVITAI_BASEMODEL_DEFAULT.copy() def search_civitai_lora(query, base_model=[], sort=CIVITAI_SORT[0], period=CIVITAI_PERIOD[0], tag="", user="", gallery=[]): @@ -1132,8 +2101,7 @@ def search_civitai_lora(query, base_model=[], sort=CIVITAI_SORT[0], period=CIVIT choices = [] gallery = [] for item in items: - base_model_name = "Pony🐴" if item['base_model'] == "Pony" else item['base_model'] - name = f"{item['name']} (for {base_model_name} / By: {item['creator']} / Tags: {', '.join(item['tags'])})" + name = build_civitai_choice_name(item) value = item['dl_url'] choices.append((name, value)) gallery.append((item['img_url'], name)) @@ -1147,7 +2115,6 @@ def search_civitai_lora(query, base_model=[], sort=CIVITAI_SORT[0], period=CIVIT return gr.update(choices=choices, value=choices[0][1], visible=True), gr.update(value=md, visible=True),\ gr.update(visible=True), gr.update(visible=True), gr.update(value=gallery) - def update_civitai_selection(evt: gr.SelectData): try: selected_index = evt.index @@ -1156,21 +2123,18 @@ def update_civitai_selection(evt: gr.SelectData): except Exception: return gr.update() - def select_civitai_lora(search_result): if not "http" in search_result: return gr.update(value=""), gr.update(value="None", visible=True) result = civitai_last_results.get(search_result, "None") md = result['md'] if result else "" return gr.update(value=search_result), gr.update(value=md, visible=True) - def download_my_lora_flux(dl_urls: str, lora): path = download_lora(dl_urls) if path: lora = path choices = get_all_lora_tupled_list() return gr.update(value=lora, choices=choices) - def apply_lora_prompt_flux(lora_info: str): if lora_info == "None": return "" lora_tag = lora_info.replace("/",",") @@ -1179,14 +2143,12 @@ def apply_lora_prompt_flux(lora_info: str): prompt = ", ".join(list_uniq(lora_prompts)) return prompt - def update_loras_flux(prompt, lora, lora_wt): on, label, tag, md = get_lora_info(lora) choices = get_all_lora_tupled_list() return gr.update(value=prompt), gr.update(value=lora, choices=choices), gr.update(value=lora_wt),\ gr.update(value=tag, label=label, visible=on), gr.update(value=md, visible=on) - def search_civitai_lora_json(query, base_model): results = {} items = search_lora_on_civitai(query, base_model) @@ -1195,22 +2157,25 @@ def search_civitai_lora_json(query, base_model): results[item['dl_url']] = item return gr.update(value=results) - def get_civitai_tag(): default = [""] user_agent = get_user_agent() headers = {'User-Agent': user_agent, 'content-type': 'application/json'} - base_url = 'https://civitai.com/api/v1/tags' params = {'limit': 200} - session = requests.Session() - retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]) - session.mount("https://", HTTPAdapter(max_retries=retries)) - url = base_url + session = create_retry_session() try: - r = session.get(url, params=params, headers=headers, stream=True, timeout=(3.0, 15)) - if not r.ok: return default - j = dict(r.json()).copy() - if "items" not in j.keys(): return default + json_data, _, r = request_civitai_api_json( + '/tags', + params=params, + headers=headers, + timeout=(3.0, 15), + api_key=CIVITAI_API_KEY, + session=session, + stream=True, + ) + if not r.ok or not json_data: return default + j = dict(json_data).copy() + if "items" not in j: return default items = [] for item in j["items"]: items.append([str(item.get("name", "")), int(item.get("modelCount", 0))]) @@ -1220,30 +2185,27 @@ def get_civitai_tag(): tags = [""] + [l[0] for l in tags] return tags except Exception as e: - print(e) + log_warning(e) return default - LORA_BASE_MODEL_DICT = { "diffusers:StableDiffusionPipeline": ["SD 1.5"], "diffusers:StableDiffusionXLPipeline": ["Pony", "SDXL 1.0"], "diffusers:FluxPipeline": ["Flux.1 D", "Flux.1 S"], } - def get_lora_base_model(model_name: str): - api = HfApi(token=HF_TOKEN) + api = get_hf_api(HF_TOKEN) default = ["Pony", "SDXL 1.0"] try: model = api.model_info(repo_id=model_name, timeout=5.0) tags = model.tags for tag in tags: - if tag in LORA_BASE_MODEL_DICT.keys(): return LORA_BASE_MODEL_DICT.get(tag, default) + if tag in LORA_BASE_MODEL_DICT: return LORA_BASE_MODEL_DICT.get(tag, default) except Exception: return default return default - def find_similar_lora(q: str, model_name: str): from rapidfuzz.process import extractOne from rapidfuzz.utils import default_process @@ -1276,7 +2238,6 @@ def find_similar_lora(q: str, model_name: str): if path and copy_lora(path, new_path): return new_path return None - def change_interface_mode(mode: str): if mode == "Fast": return gr.update(open=False), gr.update(visible=True), gr.update(open=False), gr.update(open=False),\ @@ -1295,7 +2256,6 @@ def change_interface_mode(mode: str): gr.update(visible=True), gr.update(open=False), gr.update(visible=True), gr.update(open=False),\ gr.update(visible=True), gr.update(value="Standard") - quality_prompt_list = [ { "name": "None", @@ -1339,7 +2299,6 @@ quality_prompt_list = [ }, ] - style_list = [ { "name": "None", @@ -1393,7 +2352,6 @@ style_list = [ }, ] - optimization_list = { "None": [28, 7., 'Euler', False, 'None', 1.], "Default": [28, 7., 'Euler', False, 'None', 1.], @@ -1412,28 +2370,23 @@ optimization_list = { "PCM 2step": [2, 1., 'Euler trailing', True, 'loras/pcm_sdxl_smallcfg_2step_converted.safetensors', 1.], } +def build_value_updates(*values): + return tuple(gr.update(value=value) for value in values) def set_optimization(opt, steps_gui, cfg_gui, sampler_gui, clip_skip_gui, lora_gui, lora_scale_gui): - if not opt in list(optimization_list.keys()): opt = "None" + if opt not in optimization_list: opt = "None" def_steps_gui = 28 def_cfg_gui = 7. - steps = optimization_list.get(opt, "None")[0] - cfg = optimization_list.get(opt, "None")[1] - sampler = optimization_list.get(opt, "None")[2] - clip_skip = optimization_list.get(opt, "None")[3] - lora = optimization_list.get(opt, "None")[4] - lora_scale = optimization_list.get(opt, "None")[5] + steps, cfg, sampler, clip_skip, lora, lora_scale = optimization_list.get(opt, optimization_list["None"]) if opt == "None": steps = max(steps_gui, def_steps_gui) cfg = max(cfg_gui, def_cfg_gui) clip_skip = clip_skip_gui - elif opt == "SPO" or opt == "DPO": + elif opt in {"SPO", "DPO"}: steps = max(steps_gui, def_steps_gui) cfg = max(cfg_gui, def_cfg_gui) - return gr.update(value=steps), gr.update(value=cfg), gr.update(value=sampler),\ - gr.update(value=clip_skip), gr.update(value=lora), gr.update(value=lora_scale), - + return build_value_updates(steps, cfg, sampler, clip_skip, lora, lora_scale) # [sampler_gui, steps_gui, cfg_gui, clip_skip_gui, img_width_gui, img_height_gui, optimization_gui] preset_sampler_setting = { @@ -1452,58 +2405,45 @@ preset_sampler_setting = { "Photo 1:1 Heavy": ["DPM++ 2M Karras", 40, 7., False, 1024, 1024, "None"], } - def set_sampler_settings(sampler_setting): - if not sampler_setting in list(preset_sampler_setting.keys()) or sampler_setting == "None": - return gr.update(value="Euler"), gr.update(value=28), gr.update(value=7.), gr.update(value=True),\ - gr.update(value=1024), gr.update(value=1024), gr.update(value="None") + if sampler_setting not in preset_sampler_setting or sampler_setting == "None": + return build_value_updates("Euler", 28, 7., True, 1024, 1024, "None") v = preset_sampler_setting.get(sampler_setting, ["Euler", 28, 7., True, 1024, 1024]) # sampler, steps, cfg, clip_skip, width, height, optimization - return gr.update(value=v[0]), gr.update(value=v[1]), gr.update(value=v[2]), gr.update(value=v[3]),\ - gr.update(value=v[4]), gr.update(value=v[5]), gr.update(value=v[6]) - + return build_value_updates(v[0], v[1], v[2], v[3], v[4], v[5], v[6]) preset_styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list} preset_quality = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in quality_prompt_list} - +ANIMAGINE_PROMPTS = to_list("anime artwork, anime style, vibrant, studio anime, highly detailed, masterpiece, best quality, very aesthetic, absurdres") +ANIMAGINE_NEG_PROMPTS = to_list("lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]") +PONY_PROMPTS = to_list("source_anime, score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres") +PONY_NEG_PROMPTS = to_list("source_pony, source_furry, source_cartoon, score_6, score_5, score_4, busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends") +ALL_STYLE_PROMPTS = list_uniq([item for d in style_list for item in to_list(str(d.get("prompt", "")))]) +ALL_STYLE_NEG_PROMPTS = list_uniq([item for d in style_list for item in to_list(str(d.get("negative_prompt", "")))]) +ALL_QUALITY_PROMPTS = list_uniq([item for d in quality_prompt_list for item in to_list(str(d.get("prompt", "")))]) +ALL_QUALITY_NEG_PROMPTS = list_uniq([item for d in quality_prompt_list for item in to_list(str(d.get("negative_prompt", "")))]) def process_style_prompt(prompt: str, neg_prompt: str, styles_key: str = "None", quality_key: str = "None", type: str = "Auto"): - animagine_ps = to_list("anime artwork, anime style, vibrant, studio anime, highly detailed, masterpiece, best quality, very aesthetic, absurdres") - animagine_nps = to_list("lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]") - pony_ps = to_list("source_anime, score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres") - pony_nps = to_list("source_pony, source_furry, source_cartoon, score_6, score_5, score_4, busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends") prompts = to_list(prompt) neg_prompts = to_list(neg_prompt) - all_styles_ps = [] - all_styles_nps = [] - for d in style_list: - all_styles_ps.extend(to_list(str(d.get("prompt", "")))) - all_styles_nps.extend(to_list(str(d.get("negative_prompt", "")))) - - all_quality_ps = [] - all_quality_nps = [] - for d in quality_prompt_list: - all_quality_ps.extend(to_list(str(d.get("prompt", "")))) - all_quality_nps.extend(to_list(str(d.get("negative_prompt", "")))) - quality_ps = to_list(preset_quality[quality_key][0]) quality_nps = to_list(preset_quality[quality_key][1]) styles_ps = to_list(preset_styles[styles_key][0]) styles_nps = to_list(preset_styles[styles_key][1]) - prompts = list_sub(prompts, animagine_ps + pony_ps + all_styles_ps + all_quality_ps) - neg_prompts = list_sub(neg_prompts, animagine_nps + pony_nps + all_styles_nps + all_quality_nps) + prompts = list_sub(prompts, ANIMAGINE_PROMPTS + PONY_PROMPTS + ALL_STYLE_PROMPTS + ALL_QUALITY_PROMPTS) + neg_prompts = list_sub(neg_prompts, ANIMAGINE_NEG_PROMPTS + PONY_NEG_PROMPTS + ALL_STYLE_NEG_PROMPTS + ALL_QUALITY_NEG_PROMPTS) last_empty_p = [""] if not prompts and type != "None" and type != "Auto" and styles_key != "None" and quality_key != "None" else [] last_empty_np = [""] if not neg_prompts and type != "None" and type != "Auto" and styles_key != "None" and quality_key != "None" else [] if type == "Animagine": - prompts = prompts + animagine_ps - neg_prompts = neg_prompts + animagine_nps + prompts = prompts + ANIMAGINE_PROMPTS + neg_prompts = neg_prompts + ANIMAGINE_NEG_PROMPTS elif type == "Pony": - prompts = prompts + pony_ps - neg_prompts = neg_prompts + pony_nps + prompts = prompts + PONY_PROMPTS + neg_prompts = neg_prompts + PONY_NEG_PROMPTS prompts = prompts + styles_ps + quality_ps neg_prompts = neg_prompts + styles_nps + quality_nps @@ -1513,62 +2453,47 @@ def process_style_prompt(prompt: str, neg_prompt: str, styles_key: str = "None", return gr.update(value=prompt), gr.update(value=neg_prompt), gr.update(value=type) +QUICK_PRESET_STYLE_MAP = { + 'Anime': 'Anime', + 'Photo': 'Photographic', +} + +QUICK_PRESET_SAMPLER_MAP = { + 'Anime': { + '1:1': {'Heavy': 'Anime 1:1 Heavy', 'Fast': 'Anime 1:1 Fast', 'Standard': 'Anime 1:1 Standard'}, + '3:4': {'Heavy': 'Anime 3:4 Heavy', 'Fast': 'Anime 3:4 Fast', 'Standard': 'Anime 3:4 Standard'}, + }, + 'Photo': { + '1:1': {'Heavy': 'Photo 1:1 Heavy', 'Fast': 'Photo 1:1 Fast', 'Standard': 'Photo 1:1 Standard'}, + '3:4': {'Heavy': 'Photo 3:4 Heavy', 'Fast': 'Photo 3:4 Fast', 'Standard': 'Photo 3:4 Standard'}, + }, +} + +QUICK_PRESET_QUALITY_MAP = { + 'Anime': {'Pony': 'Pony Anime Common', 'Animagine': 'Animagine Common'}, + 'Photo': {'Pony': 'Pony Common'}, +} + +def resolve_quick_preset_sampler(genre: str, aspect: str, speed: str): + speed_key = speed if speed in {'Heavy', 'Fast'} else 'Standard' + return QUICK_PRESET_SAMPLER_MAP.get(genre, {}).get(aspect, {}).get(speed_key, 'None') def set_quick_presets(genre:str = "None", type:str = "Auto", speed:str = "None", aspect:str = "None"): quality = "None" style = "None" - sampler = "None" + sampler = resolve_quick_preset_sampler(genre, aspect, speed) opt = "None" - if genre == "Anime": - if type != "None" and type != "Auto": style = "Anime" - if aspect == "1:1": - if speed == "Heavy": - sampler = "Anime 1:1 Heavy" - elif speed == "Fast": - sampler = "Anime 1:1 Fast" - else: - sampler = "Anime 1:1 Standard" - elif aspect == "3:4": - if speed == "Heavy": - sampler = "Anime 3:4 Heavy" - elif speed == "Fast": - sampler = "Anime 3:4 Fast" - else: - sampler = "Anime 3:4 Standard" - if type == "Pony": - quality = "Pony Anime Common" - elif type == "Animagine": - quality = "Animagine Common" - else: - quality = "None" - elif genre == "Photo": - if type != "None" and type != "Auto": style = "Photographic" - if aspect == "1:1": - if speed == "Heavy": - sampler = "Photo 1:1 Heavy" - elif speed == "Fast": - sampler = "Photo 1:1 Fast" - else: - sampler = "Photo 1:1 Standard" - elif aspect == "3:4": - if speed == "Heavy": - sampler = "Photo 3:4 Heavy" - elif speed == "Fast": - sampler = "Photo 3:4 Fast" - else: - sampler = "Photo 3:4 Standard" - if type == "Pony": - quality = "Pony Common" - else: - quality = "None" + if genre in QUICK_PRESET_STYLE_MAP and type not in {"None", "Auto"}: + style = QUICK_PRESET_STYLE_MAP[genre] + quality = QUICK_PRESET_QUALITY_MAP.get(genre, {}).get(type, "None") if speed == "Fast": opt = "DPO Turbo" - if genre == "Anime" and type != "Pony" and type != "Auto": quality = "Animagine Light v3.1" - - return gr.update(value=quality), gr.update(value=style), gr.update(value=sampler), gr.update(value=opt), gr.update(value=type) + if genre == "Anime" and type not in {"Pony", "Auto"}: + quality = "Animagine Light v3.1" + return build_value_updates(quality, style, sampler, opt, type) textual_inversion_dict = {} try: @@ -1578,19 +2503,29 @@ except Exception: pass textual_inversion_file_token_list = [] - def get_tupled_embed_list(embed_list): - global textual_inversion_file_list + global textual_inversion_file_token_list tupled_list = [] + textual_inversion_file_token_list = [] for file in embed_list: - token = textual_inversion_dict.get(Path(file).name, [Path(file).stem.replace(",",""), False])[0] + token = textual_inversion_dict.get(Path(file).name, [Path(file).stem.replace(",", ""), False])[0] + token = str(token).strip() tupled_list.append((token, file)) - textual_inversion_file_token_list.append(token) + if token: + textual_inversion_file_token_list.append(token) return tupled_list +def get_textual_inversion_tokens(): + dict_tokens = [] + for value in textual_inversion_dict.values(): + if isinstance(value, (list, tuple)) and value: + token = str(value[0]).strip() + if token: + dict_tokens.append(token) + return list_uniq(dict_tokens + textual_inversion_file_token_list) def set_textual_inversion_prompt(textual_inversion_gui, prompt_gui, neg_prompt_gui, prompt_syntax_gui): - ti_tags = list(textual_inversion_dict.values()) + textual_inversion_file_token_list + ti_tags = set(get_textual_inversion_tokens()) tags = prompt_gui.split(",") if prompt_gui else [] prompts = [] for tag in tags: @@ -1617,9 +2552,8 @@ def set_textual_inversion_prompt(textual_inversion_gui, prompt_gui, neg_prompt_g neg_prompt = ", ".join(neg_prompts + ti_neg_prompts + empty) return gr.update(value=prompt), gr.update(value=neg_prompt), - def get_model_pipeline(repo_id: str): - api = HfApi(token=HF_TOKEN) + api = get_hf_api(HF_TOKEN) default = "StableDiffusionPipeline" try: if not is_repo_name(repo_id): return default @@ -1638,7 +2572,6 @@ def get_model_pipeline(repo_id: str): else: return default - MODEL_TYPE_KEY = { "model.diffusion_model.output_blocks.1.1.norm.bias": "SDXL", "model.diffusion_model.input_blocks.11.0.out_layers.3.weight": "SD 1.5", @@ -1647,33 +2580,52 @@ MODEL_TYPE_KEY = { "model.diffusion_model.joint_blocks.9.x_block.attn.ln_k.weight": "SD 3.5", } +def is_unsafe_clean_target(path: str): + raw_path = str(path or "").strip() + if not raw_path: + return True + try: + resolved = Path(raw_path).expanduser().resolve() + except Exception: + return True + protected_paths = { + Path(os.getcwd()).resolve(), + Path.home().resolve(), + } + if resolved in protected_paths: + return True + if str(resolved) == resolved.anchor or resolved == resolved.parent: + return True + return False def safe_clean(path: str): + if is_unsafe_clean_target(path): + log_warning(f"Skipped delete: {path}") + return try: if Path(path).exists(): - if Path(path).is_dir(): shutil.rmtree(str(Path(path))) - else: Path(path).unlink() - print(f"Deleted: {path}") - else: print(f"File not found: {path}") + if Path(path).is_dir(): + shutil.rmtree(str(Path(path))) + else: + Path(path).unlink() + log_info(f"Deleted: {path}") + else: + log_info(f"File not found: {path}") except Exception as e: - print(f"Failed to delete: {path} {e}") - + log_error(f"Failed to delete: {path} {e}") def read_safetensors_key(path: str): + keys = [] try: - keys = [] - state_dict = load_file(str(Path(path))) - for k in list(state_dict.keys()): - keys.append(k) - state_dict.pop(k) + with safe_open(str(Path(path)), framework="pt") as f: + keys = list(f.keys()) except Exception as e: - print(e) + log_error(e) finally: - del state_dict - torch.cuda.empty_cache() + if torch.cuda.is_available(): + torch.cuda.empty_cache() gc.collect() - return keys - + return keys def get_model_type_from_key(path: str): default = "SDXL" @@ -1681,14 +2633,13 @@ def get_model_type_from_key(path: str): keys = read_safetensors_key(path) for k, v in MODEL_TYPE_KEY.items(): if k in set(keys): - print(f"Model type is {v}.") + log_info(f"Model type is {v}.") return v - print("Model type could not be identified.") + log_warning("Model type could not be identified.") except Exception: return default return default - def download_link_model(url: str, localdir: str): try: new_file = None @@ -1701,7 +2652,6 @@ def download_link_model(url: str, localdir: str): except Exception as e: raise gr.Error(f"Failed to load single model file: {url} {e}") - EXAMPLES_GUI = [ [ "1girl, souryuu asuka langley, neon genesis evangelion, plugsuit, pilot suit, red bodysuit, sitting, crossing legs, black eye patch, cat hat, throne, symmetrical, looking down, from bottom, looking at viewer, outdoors, masterpiece, best quality, very aesthetic, absurdres", @@ -1770,7 +2720,6 @@ EXAMPLES_GUI = [ ], ] - RESOURCES = ( """### Resources - You can also try the image generator in Colab’s free tier, which provides free GPU [link](https://github.com/R3gm/SD_diffusers_interactive).