John6666/votepurchase-multiple-model
141
1import spaces2import json3import gradio as gr4import os5import re6from pathlib import Path7from PIL import Image8import numpy as np9import shutil10import requests11from requests.adapters import HTTPAdapter12from urllib3.util import Retry13import urllib.parse14import pandas as pd15from typing import Any16from huggingface_hub import HfApi, hf_hub_download, snapshot_download17from translatepy import Translator18from unidecode import unidecode19import copy20from datetime import datetime, timezone, timedelta21FILENAME_TIMEZONE = timezone(timedelta(hours=9)) # JST22import torch23from safetensors import safe_open24import gc25import html as html_lib26import subprocess27import tempfile28import time29 30from env import (HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,31 HF_MODEL_USER_EX, HF_MODEL_USER_LIKES, DIFFUSERS_FORMAT_LORAS,32 DIRECTORY_LORAS, HF_READ_TOKEN, HF_TOKEN, CIVITAI_API_KEY)33 34OUTPUT_CACHE_DIR = Path(os.getenv("OUTPUT_CACHE_DIR", "outputs"))35OUTPUT_CACHE_MAX_FILES = max(16, int(os.getenv("OUTPUT_CACHE_MAX_FILES", "256")))36OUTPUT_CACHE_MAX_BYTES = max(512 * 1024**2, int(float(os.getenv("OUTPUT_CACHE_MAX_GB", "4")) * 1024**3))37CIVITAI_LORA_CACHE_MAX_FILES = max(16, int(os.getenv("CIVITAI_LORA_CACHE_MAX_FILES", "128")))38CIVITAI_LORA_CACHE_MAX_BYTES = max(4 * 1024**3, int(float(os.getenv("CIVITAI_LORA_CACHE_MAX_GB", "32")) * 1024**3))39CIVITAI_LORA_CACHE_INDEX = Path(DIRECTORY_LORAS) / ".civitai_lora_lru.json"40CIVITAI_ALLOWED_LORA_BASE_MODELS = ['Flux.1 D', 'Flux.1 S', 'Flux.1 Kontext', 'Flux.1 Krea', 'Flux.2 D', 'Flux.2 Klein 4B-base', 'Flux.2 Klein 9B', 'Flux.2 Klein 9B-base', 'SD 1.5', 'SD 1.5 Hyper', 'SD 1.5 LCM', 'SDXL 0.9', 'SDXL 1.0', 'SDXL Hyper', 'SDXL Lightning', 'SDXL 1.0 LCM', 'Pony', 'Illustrious', 'NoobAI']41 42 43def _prune_paths_lru(paths, max_files: int, max_bytes: int, protect=None):44 protect = {str(Path(p).resolve()) for p in (protect or []) if p}45 entries = []46 for raw in paths:47 try:48 p = Path(raw)49 if not p.is_file():50 continue51 st = p.stat()52 entries.append((p, int(st.st_size), float(st.st_mtime)))53 except Exception:54 continue55 total = sum(size for _, size, _ in entries)56 entries.sort(key=lambda item: item[2])57 while entries and (len(entries) > max_files or total > max_bytes):58 victim, size, _ = entries.pop(0)59 if str(victim.resolve()) in protect:60 entries.append((victim, size, float("inf")))61 entries.sort(key=lambda item: item[2])62 if all(str(p.resolve()) in protect for p, _, _ in entries):63 break64 continue65 try:66 victim.unlink()67 total -= size68 print(f"[cache] pruned {victim}")69 except Exception as e:70 print(f"[cache] prune failed {victim} {type(e).__name__}: {e}")71 return total72 73 74def _prune_generated_outputs(protect=None):75 OUTPUT_CACHE_DIR.mkdir(parents=True, exist_ok=True)76 _prune_paths_lru(OUTPUT_CACHE_DIR.glob("*.png"), OUTPUT_CACHE_MAX_FILES, OUTPUT_CACHE_MAX_BYTES, protect=protect)77 try:78 return {str(path.resolve()) for path in OUTPUT_CACHE_DIR.glob("*.png") if path.is_file()}79 except Exception:80 return set()81 82 83def _load_civitai_lora_index():84 try:85 data = json.loads(CIVITAI_LORA_CACHE_INDEX.read_text(encoding="utf-8"))86 return data if isinstance(data, dict) else {}87 except Exception:88 return {}89 90 91def _save_civitai_lora_index(data):92 try:93 CIVITAI_LORA_CACHE_INDEX.parent.mkdir(parents=True, exist_ok=True)94 temp = CIVITAI_LORA_CACHE_INDEX.with_suffix(".tmp")95 temp.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")96 os.replace(temp, CIVITAI_LORA_CACHE_INDEX)97 except Exception as e:98 print(f"[civitai] lora cache index save failed: {type(e).__name__}: {e}")99 100 101def _record_civitai_lora_cache(path: str):102 try:103 current = Path(path).resolve()104 lora_root = Path(DIRECTORY_LORAS).resolve()105 if not current.is_file() or current.parent != lora_root:106 return107 data = _load_civitai_lora_index()108 normalized = {}109 for raw, touched in data.items():110 try:111 p = Path(raw).resolve()112 if p.is_file() and p.parent == lora_root:113 normalized[str(p)] = float(touched)114 except Exception:115 pass116 normalized[str(current)] = time.time()117 paths = list(normalized.keys())118 _prune_paths_lru(paths, CIVITAI_LORA_CACHE_MAX_FILES, CIVITAI_LORA_CACHE_MAX_BYTES, protect=[str(current)])119 normalized = {raw: ts for raw, ts in normalized.items() if Path(raw).is_file()}120 _save_civitai_lora_index(normalized)121 except Exception as e:122 print(f"[civitai] lora cache accounting failed: {type(e).__name__}: {e}")123 124 125def is_allowed_civitai_lora_base_model(base_model: str) -> bool:126 value = str(base_model or "").strip()127 return value in CIVITAI_ALLOWED_LORA_BASE_MODELS128 129MODEL_TYPE_DICT = {130 "diffusers:StableDiffusionPipeline": "SD 1.5",131 "diffusers:StableDiffusionXLPipeline": "SDXL",132 "diffusers:FluxPipeline": "FLUX",133}134 135def log_info(message: str):136 print(str(message))137 138def log_warning(message: str):139 print(str(message))140 141def log_error(message: str):142 print(str(message))143 144def get_user_agent():145 return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0'146 147def to_list(s):148 return [x.strip() for x in s.split(",") if not s == ""]149 150def list_uniq(l):151 return sorted(set(l), key=l.index)152 153def list_sub(a, b):154 return [e for e in a if e not in b]155 156def is_repo_name(s):157 return re.fullmatch(r'^[^/]+?/[^/]+?$', s)158 159DEFAULT_STATE = {160 "show_diffusers_model_list_detail": False,161}162 163def get_state(state: dict, key: str):164 if key in state:165 return state[key]166 if key in DEFAULT_STATE:167 log_info(f"State '{key}' not found. Use default value.")168 return DEFAULT_STATE[key]169 log_warning(f"State '{key}' not found.")170 return None171 172def set_state(state: dict, key: str, value: Any):173 state[key] = value174 175translator = Translator()176def translate_to_en(input: str):177 try:178 output = str(translator.translate(input, 'English'))179 except Exception as e:180 output = input181 log_warning(e)182 return output183 184def get_local_model_list(dir_path):185 model_list = []186 valid_extensions = ('.ckpt', '.pt', '.pth', '.safetensors', '.bin')187 dir_path = Path(dir_path)188 for file in dir_path.glob("*"):189 if file.suffix in valid_extensions:190 file_path = str(dir_path / file.name)191 model_list.append(file_path)192 #print('\033[34mFILE: ' + file_path + '\033[0m')193 return model_list194 195HF_FOLDER_TOKEN = ""196 197def get_token():198 return HF_FOLDER_TOKEN199 200def set_token(token):201 global HF_FOLDER_TOKEN202 HF_FOLDER_TOKEN = token203 204set_token(HF_TOKEN)205 206def get_hf_api(token: str = ""):207 return HfApi(token=token) if token else HfApi()208 209HF_HOST_ALIASES = frozenset({"huggingface.co", "www.huggingface.co", "hf.co"})210 211def parse_hf_file_url(url: str):212 raw = str(url or "").strip()213 if not raw:214 return {}215 try:216 parts = urllib.parse.urlsplit(raw)217 except Exception:218 return {}219 if str(parts.netloc or "").strip().lower() not in HF_HOST_ALIASES:220 return {}221 222 path_segments = [seg for seg in str(parts.path or "").split("/") if seg]223 if not path_segments:224 return {}225 226 repo_type = "model"227 if path_segments[0] in ["datasets", "spaces"]:228 repo_type = "dataset" if path_segments[0] == "datasets" else "space"229 path_segments = path_segments[1:]230 231 if len(path_segments) < 5:232 return {}233 234 namespace, repo_name, action, revision = path_segments[:4]235 if action not in ["resolve", "blob"]:236 return {}237 238 file_segments = [urllib.parse.unquote(seg) for seg in path_segments[4:]]239 if not file_segments:240 return {}241 242 filename = file_segments[-1]243 subfolder = "/".join(file_segments[:-1]) if len(file_segments) > 1 else None244 return {245 "repo_id": f"{namespace}/{repo_name}",246 "filename": filename,247 "subfolder": subfolder,248 "repo_type": repo_type,249 "revision": urllib.parse.unquote(revision),250 }251 252def split_hf_url(url: str):253 parsed = parse_hf_file_url(url)254 if not parsed:255 return "", "", "", ""256 return parsed["repo_id"], parsed["filename"], parsed["subfolder"], parsed["repo_type"]257 258def download_hf_file(directory, url, force_filename="", hf_token="", progress=gr.Progress(track_tqdm=True)):259 parsed = parse_hf_file_url(url)260 if not parsed:261 log_download_error("hf", "parse_url", url=url)262 return None263 264 kwargs = {}265 if parsed["subfolder"] is not None:266 kwargs["subfolder"] = parsed["subfolder"]267 if parsed.get("revision"):268 kwargs["revision"] = parsed["revision"]269 try:270 print(271 f"Start HF download: repo={parsed['repo_id']} rev={parsed.get('revision') or '-'} "272 f"file={parsed['filename']} to {directory}"273 )274 path = hf_hub_download(275 repo_id=parsed["repo_id"],276 filename=parsed["filename"],277 repo_type=parsed["repo_type"],278 local_dir=directory,279 token=hf_token,280 **kwargs,281 )282 forced_path = str(Path(directory) / force_filename) if force_filename else ""283 if forced_path:284 return move_downloaded_file_to_target(path, forced_path)285 return path286 except Exception as e:287 log_download_error("hf", "hub_download", url=url, error=e)288 forced_path = str(Path(directory) / force_filename) if force_filename else ""289 if forced_path and Path(forced_path).exists():290 return forced_path291 return None292 293USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0'294CIVITAI_DEFAULT_ORIGIN = "https://civitai.com"295CIVITAI_CANONICAL_WEB_ORIGIN = CIVITAI_DEFAULT_ORIGIN296CIVITAI_RED_ORIGIN = "https://civitai.red"297CIVITAI_GREEN_HOST_ALIASES = frozenset({"civitai.green", "www.civitai.green"})298CIVITAI_RED_HOST_ALIASES = frozenset({"civitai.red", "www.civitai.red"})299CIVITAI_HOST_ALIASES = frozenset({"civitai.com", "www.civitai.com", *CIVITAI_GREEN_HOST_ALIASES, *CIVITAI_RED_HOST_ALIASES})300CIVITAI_API_ORIGIN_CANDIDATES = (CIVITAI_RED_ORIGIN, CIVITAI_DEFAULT_ORIGIN)301CIVITAI_REFERER = f"{CIVITAI_CANONICAL_WEB_ORIGIN}/"302CIVITAI_RETRY_TOTAL = 5303CIVITAI_RETRY_BACKOFF = 1.0304CIVITAI_RESOLVE_RETRY_TOTAL = 4305CIVITAI_RESOLVE_RETRY_BACKOFF = 0.8306CIVITAI_STATUS_FORCELIST = [429, 500, 502, 503, 504]307CIVITAI_RESOLVE_TIMEOUT = (7.0, 25.0)308CIVITAI_METADATA_TIMEOUT = (3.0, 15.0)309CIVITAI_SEARCH_TIMEOUT = (3.0, 30.0)310CIVITAI_NEGATIVE_CACHE_LIMIT = 256311CIVITAI_RESOLVE_CACHE: dict[str, str] = {}312CIVITAI_RESOLVE_NEGATIVE_CACHE: dict[str, str] = {}313CIVITAI_VERSION_JSON_CACHE: dict[str, dict] = {}314CIVITAI_VERSION_NEGATIVE_CACHE: dict[str, str] = {}315CIVITAI_WGET_FRESH_RETRY_LIMIT = 1316CIVITAI_METADATA_RECONNECT_ATTEMPTS = 3317CIVITAI_METADATA_RECONNECT_BACKOFF = 0.8318CIVITAI_API_PROBE_TIMEOUT = (3.0, 8.0)319CIVITAI_API_RETRYABLE_STATUSES = frozenset([404, 405, 429, 500, 502, 503, 504])320CIVITAI_ACTIVE_API_ORIGIN = ""321CIVITAI_ACTIVE_API_BASE = ""322 323def create_retry_session(total=CIVITAI_RETRY_TOTAL, backoff_factor=CIVITAI_RETRY_BACKOFF):324 session = requests.Session()325 retries = Retry(total=total, backoff_factor=backoff_factor, status_forcelist=CIVITAI_STATUS_FORCELIST)326 session.mount("https://", HTTPAdapter(max_retries=retries))327 session.mount("http://", HTTPAdapter(max_retries=retries))328 return session329 330def cache_put(cache: dict, key: str, value):331 key = str(key or "").strip()332 if not key:333 return334 if key in cache:335 cache.pop(key, None)336 elif len(cache) >= CIVITAI_NEGATIVE_CACHE_LIMIT:337 try:338 cache.pop(next(iter(cache)))339 except Exception:340 cache.clear()341 cache[key] = value342 343def get_civitai_headers(api_key: str = ""):344 headers = {'User-Agent': USER_AGENT, 'content-type': 'application/json'}345 if api_key:346 headers['Authorization'] = f'Bearer {api_key}'347 return headers348 349def get_civitai_url_parts(url: str):350 try:351 return urllib.parse.urlsplit(str(url or "").strip())352 except Exception:353 return urllib.parse.urlsplit("")354 355def sanitize_url_for_log(url: str):356 raw = str(url or "").strip()357 if not raw:358 return raw359 parts = get_civitai_url_parts(raw)360 if not parts.netloc:361 return raw362 pairs = [363 (k, v)364 for k, v in urllib.parse.parse_qsl(parts.query, keep_blank_values=True)365 if str(k).lower() != "token"366 ]367 query = urllib.parse.urlencode(pairs)368 return urllib.parse.urlunsplit((parts.scheme, parts.netloc, parts.path, query, parts.fragment))369 370def canonicalize_civitai_netloc(netloc: str):371 host = str(netloc or "").strip().lower()372 if host in CIVITAI_GREEN_HOST_ALIASES:373 return "civitai.com"374 if host == "www.civitai.com":375 return "civitai.com"376 if host == "www.civitai.red":377 return "civitai.red"378 return host379 380def is_civitai_host(netloc: str):381 return canonicalize_civitai_netloc(netloc) in {"civitai.com", "civitai.red"}382 383def is_civitai_url(url: str):384 return is_civitai_host(get_civitai_url_parts(url).netloc)385 386def get_civitai_canonical_web_origin():387 return CIVITAI_CANONICAL_WEB_ORIGIN388 389def build_civitai_api_base(origin: str):390 raw = str(origin or "").strip().rstrip("/")391 return f"{raw}/api/v1" if raw else ""392 393def set_civitai_active_api_origin(origin: str):394 global CIVITAI_ACTIVE_API_ORIGIN, CIVITAI_ACTIVE_API_BASE395 raw = str(origin or "").strip().rstrip("/")396 if not raw:397 raw = CIVITAI_DEFAULT_ORIGIN398 CIVITAI_ACTIVE_API_ORIGIN = raw399 CIVITAI_ACTIVE_API_BASE = build_civitai_api_base(raw)400 return CIVITAI_ACTIVE_API_BASE401 402def probe_civitai_api_origin(session, origin: str, api_key: str = ""):403 headers = get_civitai_headers(api_key or CIVITAI_API_KEY)404 base_url = build_civitai_api_base(origin)405 if not base_url:406 return False407 response = None408 try:409 response = session.get(410 f"{base_url}/tags",411 params={"limit": 1},412 headers=headers,413 stream=True,414 timeout=CIVITAI_API_PROBE_TIMEOUT,415 )416 if not response.ok:417 return False418 content_type = str(response.headers.get("content-type") or "").lower()419 if "json" not in content_type:420 return False421 data = response.json()422 return isinstance(data, dict) and "items" in data423 except Exception:424 return False425 finally:426 try:427 if response is not None:428 response.close()429 except Exception:430 pass431 432def get_civitai_active_api_origin(force_refresh: bool = False, api_key: str = ""):433 global CIVITAI_ACTIVE_API_ORIGIN434 if CIVITAI_ACTIVE_API_ORIGIN and not force_refresh:435 return CIVITAI_ACTIVE_API_ORIGIN436 session = create_retry_session(total=2, backoff_factor=0.5)437 try:438 for origin in CIVITAI_API_ORIGIN_CANDIDATES:439 if probe_civitai_api_origin(session, origin, api_key=api_key):440 set_civitai_active_api_origin(origin)441 print(f"[civitai] selected api origin: {CIVITAI_ACTIVE_API_ORIGIN}")442 return CIVITAI_ACTIVE_API_ORIGIN443 set_civitai_active_api_origin(CIVITAI_DEFAULT_ORIGIN)444 print(f"[civitai] api probe fallback origin: {CIVITAI_ACTIVE_API_ORIGIN}")445 return CIVITAI_ACTIVE_API_ORIGIN446 finally:447 try:448 session.close()449 except Exception:450 pass451 452def get_civitai_active_api_base(force_refresh: bool = False, api_key: str = ""):453 if CIVITAI_ACTIVE_API_BASE and not force_refresh:454 return CIVITAI_ACTIVE_API_BASE455 get_civitai_active_api_origin(force_refresh=force_refresh, api_key=api_key)456 return CIVITAI_ACTIVE_API_BASE457 458def iter_civitai_api_bases(api_key: str = ""):459 preferred = get_civitai_active_api_base(api_key=api_key)460 bases = [preferred] if preferred else []461 for origin in CIVITAI_API_ORIGIN_CANDIDATES:462 base = build_civitai_api_base(origin)463 if base and base not in bases:464 bases.append(base)465 return bases466 467def is_retryable_civitai_api_status(status):468 try:469 return int(status) in CIVITAI_API_RETRYABLE_STATUSES470 except Exception:471 return False472 473def request_civitai_api_response(path: str, params=None, headers=None, timeout=CIVITAI_METADATA_TIMEOUT,474 api_key: str = "", session=None, stream: bool = True, allow_not_found: bool = False):475 effective_api_key = api_key or CIVITAI_API_KEY476 request_headers = headers or get_civitai_headers(effective_api_key)477 request_session = session or create_retry_session()478 last_response = None479 last_url = ""480 last_error = None481 bases = iter_civitai_api_bases(api_key=effective_api_key)482 for idx, base_url in enumerate(bases):483 url = f"{base_url}/{str(path or '').lstrip('/')}"484 try:485 response = request_session.get(url, params=params, headers=request_headers, stream=stream, timeout=timeout)486 if response.ok or (allow_not_found and response.status_code == 404):487 set_civitai_active_api_origin(base_url.rsplit('/api/v1', 1)[0])488 return response, url489 last_response = response490 last_url = url491 if idx + 1 < len(bases) and is_retryable_civitai_api_status(response.status_code):492 try:493 response.close()494 except Exception:495 pass496 continue497 return response, url498 except Exception as e:499 last_error = e500 last_url = url501 if idx + 1 < len(bases):502 continue503 raise504 if last_response is not None:505 return last_response, last_url506 if last_error is not None:507 raise last_error508 raise RuntimeError(f"Failed to request Civitai API path: {path}")509 510def get_civitai_api_origin_from_url(url: str):511 raw = str(url or "").strip()512 if not raw:513 return ""514 if raw.endswith('/api/v1'):515 raw = raw.rsplit('/api/v1', 1)[0]516 parts = get_civitai_url_parts(raw)517 netloc = canonicalize_civitai_netloc(parts.netloc)518 if not netloc:519 return ""520 scheme = parts.scheme or 'https'521 return f"{scheme}://{netloc}"522 523def request_civitai_api_json(path: str, params=None, headers=None, timeout=CIVITAI_METADATA_TIMEOUT,524 api_key: str = "", session=None, stream: bool = True, allow_not_found: bool = False,525 non_json_fallback_origin: str = CIVITAI_DEFAULT_ORIGIN):526 effective_api_key = api_key or CIVITAI_API_KEY527 request_headers = headers or get_civitai_headers(effective_api_key)528 request_session = session or create_retry_session()529 result, endpoint_url = request_civitai_api_response(530 path,531 params=params,532 headers=request_headers,533 timeout=timeout,534 api_key=effective_api_key,535 session=request_session,536 stream=stream,537 allow_not_found=allow_not_found,538 )539 if allow_not_found and result.status_code == 404:540 return None, endpoint_url, result541 result.raise_for_status()542 try:543 return result.json(), endpoint_url, result544 except Exception:545 current_origin = get_civitai_api_origin_from_url(endpoint_url)546 fallback_origin = str(non_json_fallback_origin or CIVITAI_DEFAULT_ORIGIN).strip().rstrip('/')547 if fallback_origin and current_origin and current_origin != fallback_origin:548 print(f"[retry] Civitai API non-json response from {current_origin}: {str(path or '').lstrip('/')}")549 try:550 result.close()551 except Exception:552 pass553 fallback_base = build_civitai_api_base(fallback_origin)554 fallback_url = f"{fallback_base}/{str(path or '').lstrip('/')}"555 fallback_response = request_session.get(556 fallback_url,557 params=params,558 headers=request_headers,559 stream=stream,560 timeout=timeout,561 )562 if allow_not_found and fallback_response.status_code == 404:563 return None, fallback_url, fallback_response564 fallback_response.raise_for_status()565 fallback_json = fallback_response.json()566 set_civitai_active_api_origin(fallback_origin)567 return fallback_json, fallback_url, fallback_response568 raise569 570try:571 get_civitai_active_api_base()572except Exception as e:573 print(f"[civitai] startup api probe failed: {type(e).__name__}: {e}")574 set_civitai_active_api_origin(CIVITAI_DEFAULT_ORIGIN)575 576def is_civitai_download_api_path(path: str):577 return re.match(r'^/api/download/models/\d+$', str(path or "").strip()) is not None578 579def extract_civitai_model_version_id(url: str):580 try:581 parts = get_civitai_url_parts(url)582 for pattern in [r'^/api/download/models/(\d+)$', r'^/api/v1/model-versions/(\d+)$']:583 m = re.match(pattern, str(parts.path or "").strip())584 if m:585 return m.group(1)586 qs = urllib.parse.parse_qs(parts.query)587 for key in ["modelVersionId", "modelversionid", "versionId", "versionid"]:588 values = qs.get(key, [])589 if not values:590 continue591 value = str(values[0]).strip()592 if value.isdigit():593 return value594 except Exception:595 return ""596 return ""597 598def extract_civitai_file_id(url: str):599 try:600 parts = get_civitai_url_parts(url)601 qs = urllib.parse.parse_qs(parts.query)602 for key, values in qs.items():603 if str(key).casefold() != "fileid" or not values:604 continue605 value = str(values[0] or "").strip()606 if value.isdigit():607 return value608 except Exception:609 return ""610 return ""611 612def get_civitai_query_filters(url: str):613 try:614 parts = get_civitai_url_parts(url)615 qs = urllib.parse.parse_qs(parts.query)616 except Exception:617 return {}618 filters = {}619 for key in ["type", "format", "size", "fp"]:620 values = qs.get(key, [])621 if values:622 filters[key] = str(values[0]).strip()623 return filters624 625def normalize_civitai_filter_value(key: str, value):626 if value is None:627 return ""628 text = str(value).strip()629 if not text:630 return ""631 if key == "fp":632 return text.replace("-", "").replace("_", "").replace(" ", "").lower()633 return text.lower()634 635def describe_civitai_file_for_log(file_info):636 if not isinstance(file_info, dict):637 return ""638 parts = []639 for key in ["name", "type", "format", "size", "fp"]:640 value = file_info.get(key)641 if value is None:642 continue643 text = str(value).strip()644 if text:645 parts.append(f"{key}={text}")646 hashes = file_info.get("hashes") if isinstance(file_info.get("hashes"), dict) else {}647 sha256 = str(hashes.get("SHA256") or "").strip()648 if sha256:649 parts.append(f"sha256={sha256[:12]}...")650 return ", ".join(parts)651 652def build_civitai_download_query_from_url(url: str):653 try:654 parts = get_civitai_url_parts(url)655 except Exception:656 return ""657 blocked = {"modelversionid", "versionid"}658 pairs = [659 (k, v)660 for k, v in urllib.parse.parse_qsl(parts.query, keep_blank_values=True)661 if str(k).lower() not in blocked662 ]663 return urllib.parse.urlencode(pairs)664 665def to_civitai_default_download_url(version_id: str, query: str = ""):666 if not str(version_id or "").isdigit():667 return ""668 base = f"{get_civitai_active_api_origin()}/api/download/models/{version_id}"669 return f"{base}?{query}" if query else base670 671def normalize_civitai_download_api_url(url: str):672 parts = get_civitai_url_parts(url)673 if not is_civitai_host(parts.netloc) or not is_civitai_download_api_path(parts.path):674 return str(url or "").strip()675 host = canonicalize_civitai_netloc(parts.netloc)676 return urllib.parse.urlunsplit(("https", host, parts.path, parts.query, ""))677 678def extract_first_civitai_download_url_from_html(html: str):679 if not html:680 return ""681 page = html_lib.unescape(str(html))682 patterns = [683 r'https?://(?:www\.)?(?:civitai\.com|civitai\.green|civitai\.red)/api/download/models/\d+[^\s\'\"<>\)\]\}]*',684 r'["\'](/api/download/models/\d+[^"\']*)["\']',685 ]686 for pattern in patterns:687 try:688 m = re.search(pattern, page, flags=re.IGNORECASE)689 except re.error:690 m = None691 if not m:692 continue693 candidate = m.group(1) if m.lastindex else m.group(0)694 candidate = str(candidate or "").strip("\"'")695 if candidate.startswith("/"):696 candidate = urllib.parse.urljoin(get_civitai_canonical_web_origin(), candidate)697 return normalize_civitai_download_api_url(candidate)698 return ""699 700def resolve_civitai_model_page_to_download_url(url: str, api_key: str = ""):701 raw = str(url or "").strip()702 if not raw:703 return raw704 cached = CIVITAI_RESOLVE_CACHE.get(raw)705 if cached:706 return cached707 if raw in CIVITAI_RESOLVE_NEGATIVE_CACHE:708 return raw709 parts = get_civitai_url_parts(raw)710 if not is_civitai_host(parts.netloc):711 return raw712 if is_civitai_download_api_path(parts.path):713 normalized = normalize_civitai_download_api_url(raw)714 cache_put(CIVITAI_RESOLVE_CACHE, raw, normalized)715 return normalized716 if not re.match(r'^/models/\d+(?:/[^/?#]+)?/?$', parts.path or ""):717 return raw718 version_id = extract_civitai_model_version_id(raw)719 if version_id:720 normalized = to_civitai_default_download_url(version_id, query=build_civitai_download_query_from_url(raw))721 cache_put(CIVITAI_RESOLVE_CACHE, raw, normalized)722 return normalized723 headers = get_civitai_headers(api_key if parts.netloc.lower().endswith("civitai.com") else "")724 headers['Referer'] = f"{parts.scheme or 'https'}://{parts.netloc}/"725 session = create_retry_session(total=CIVITAI_RESOLVE_RETRY_TOTAL, backoff_factor=CIVITAI_RESOLVE_RETRY_BACKOFF)726 r = None727 try:728 r = session.get(raw, headers=headers, timeout=CIVITAI_RESOLVE_TIMEOUT)729 if not r.ok:730 print(f"Civitai model page resolve failed: {sanitize_url_for_log(raw)} status={r.status_code}")731 if r.status_code in [400, 401, 403, 404]:732 cache_put(CIVITAI_RESOLVE_NEGATIVE_CACHE, raw, f"status={r.status_code}")733 return raw734 extracted = extract_first_civitai_download_url_from_html(r.text)735 if extracted:736 normalized = normalize_civitai_download_api_url(extracted)737 cache_put(CIVITAI_RESOLVE_CACHE, raw, normalized)738 return normalized739 return raw740 except Exception as e:741 print(f"Failed to resolve Civitai model page URL: {sanitize_url_for_log(raw)} {type(e).__name__}: {sanitize_sensitive_log_text(e)}")742 return raw743 finally:744 try:745 if r is not None:746 r.close()747 except Exception:748 pass749 try:750 session.close()751 except Exception:752 pass753 754def normalize_civitai_input_url(url: str, api_key: str = ""):755 raw = str(url or "").strip()756 if not raw or not is_civitai_url(raw):757 return raw758 normalized = resolve_civitai_model_page_to_download_url(raw, api_key=api_key)759 if normalized != raw:760 print(f"Normalized Civitai URL: {sanitize_url_for_log(raw)} -> {sanitize_url_for_log(normalized)}")761 return normalized762 763def append_civitai_token(url: str, api_key: str = ""):764 raw = str(url or "").strip()765 if not raw or not api_key:766 return raw767 parts = get_civitai_url_parts(raw)768 pairs = [(k, v) for k, v in urllib.parse.parse_qsl(parts.query, keep_blank_values=True) if k.lower() != "token"]769 pairs.append(("token", api_key))770 query = urllib.parse.urlencode(pairs)771 return urllib.parse.urlunsplit((parts.scheme or "https", parts.netloc, parts.path, query, parts.fragment))772 773def get_civitai_request_context(url: str, api_key: str = ""):774 raw_url = str(url or "").strip()775 normalized_url = normalize_civitai_input_url(raw_url, api_key=api_key)776 model_version_id = extract_civitai_model_version_id(normalized_url) or extract_civitai_model_version_id(raw_url)777 return {778 "raw_url": raw_url,779 "normalized_url": normalized_url,780 "model_version_id": model_version_id,781 "filters": get_civitai_query_filters(raw_url),782 }783 784def resolve_civitai_download_url(url: str, civitai_api_key: str = "", max_tries: int = 3):785 raw = normalize_civitai_download_api_url(str(url or "").strip())786 if not raw:787 return raw788 headers = get_civitai_headers(civitai_api_key)789 headers["Referer"] = CIVITAI_REFERER790 dl_url = append_civitai_token(raw, civitai_api_key)791 last_error = None792 for attempt in range(1, max_tries + 1):793 response = None794 session = create_retry_session(total=3, backoff_factor=1.0)795 try:796 response = session.get(797 dl_url,798 headers=headers,799 allow_redirects=False,800 stream=True,801 timeout=CIVITAI_RESOLVE_TIMEOUT,802 )803 status = int(response.status_code)804 location = str(response.headers.get("Location") or "").strip()805 resolved_url = str(location or response.url or dl_url).strip()806 resolved_host = get_civitai_url_parts(resolved_url).netloc807 print(808 f"[civitai] resolve signed url attempt={attempt}/{max_tries} status={status} "809 f"host={resolved_host or '-'} url={sanitize_url_for_log(raw)}"810 )811 if status in (301, 302, 303, 307, 308) and location:812 return resolved_url813 if response.ok and resolved_url and not is_civitai_host(resolved_host):814 return resolved_url815 last_error = RuntimeError(f"status={status}")816 except Exception as e:817 last_error = e818 print(819 f"[civitai] resolve signed url failed attempt={attempt}/{max_tries} "820 f"url={sanitize_url_for_log(raw)} error={type(e).__name__}: {sanitize_sensitive_log_text(e)}"821 )822 finally:823 try:824 if response is not None:825 response.close()826 except Exception:827 pass828 try:829 session.close()830 except Exception:831 pass832 if attempt < max_tries:833 time.sleep(min(3.0, 0.8 * attempt))834 if last_error is not None:835 raise last_error836 raise RuntimeError("Failed to resolve Civitai signed download URL")837 838def pick_civitai_file_from_version_json(json_data, source_url: str = ""):839 files = json_data.get("files", []) if isinstance(json_data, dict) else []840 if not isinstance(files, list) or not files:841 return {}842 explicit_file_id = extract_civitai_file_id(source_url)843 if explicit_file_id:844 for file_info in files:845 if not isinstance(file_info, dict):846 continue847 candidate_id = str(file_info.get("id") or file_info.get("fileId") or "").strip()848 if candidate_id == explicit_file_id:849 return dict(file_info)850 print(f"[civitai] explicit fileId={explicit_file_id} not present in model version metadata")851 return {}852 version_id = str((json_data or {}).get("id") or "")853 filters = get_civitai_query_filters(source_url)854 candidates = []855 fallback = []856 for idx, file_info in enumerate(files):857 if not isinstance(file_info, dict):858 continue859 mismatch = False860 matched_filter_count = 0861 for key, expected in filters.items():862 actual = file_info.get(key)863 expected_norm = normalize_civitai_filter_value(key, expected)864 actual_norm = normalize_civitai_filter_value(key, actual)865 if actual_norm:866 if actual_norm != expected_norm:867 mismatch = True868 break869 matched_filter_count += 1870 download_url = str(file_info.get("downloadUrl") or "")871 score = 0872 if matched_filter_count:873 score += matched_filter_count * 3874 if version_id and version_id in download_url:875 score += 4876 if download_url:877 score += 2878 if file_info.get("name"):879 score += 1880 hashes = file_info.get("hashes") if isinstance(file_info.get("hashes"), dict) else {}881 if str(hashes.get("SHA256") or "").strip():882 score += 1883 target = fallback if mismatch else candidates884 target.append((score, idx, file_info))885 pool = candidates if candidates else fallback886 if not pool:887 return {}888 pool.sort(key=lambda item: (item[0], item[1]), reverse=True)889 return dict(pool[0][2])890 891def move_downloaded_file_to_target(downloaded_path: str, target_path: str):892 source = Path(str(downloaded_path or "")).expanduser()893 target = Path(str(target_path or "")).expanduser()894 if not str(target):895 return str(source)896 if not source.exists():897 return str(target) if target.exists() else str(source)898 try:899 if source.resolve() == target.resolve():900 return str(target)901 except Exception:902 pass903 904 try:905 target.parent.mkdir(parents=True, exist_ok=True)906 if target.exists() and target.is_file():907 target.unlink()908 shutil.move(str(source), str(target))909 return str(target)910 except Exception as e:911 print(f"HF local rename failed: {source} -> {target} {type(e).__name__}: {sanitize_sensitive_log_text(e)}")912 return str(source)913 914def request_json_data(url, api_key: str = ""):915 effective_api_key = api_key or CIVITAI_API_KEY916 context = get_civitai_request_context(url, api_key=effective_api_key)917 raw_url = context["raw_url"]918 normalized_url = context["normalized_url"]919 model_version_id = context["model_version_id"]920 if not model_version_id:921 print(f"Civitai metadata lookup skipped: modelVersionId not found for {sanitize_url_for_log(raw_url)}")922 cache_put(CIVITAI_RESOLVE_NEGATIVE_CACHE, raw_url, "missing_model_version_id")923 return None924 925 cached_json = CIVITAI_VERSION_JSON_CACHE.get(model_version_id)926 if cached_json:927 return copy.deepcopy(cached_json)928 if model_version_id in CIVITAI_VERSION_NEGATIVE_CACHE:929 return None930 931 endpoint_path = f"/model-versions/{model_version_id}"932 last_error = None933 for attempt in range(1, CIVITAI_METADATA_RECONNECT_ATTEMPTS + 1):934 session = create_retry_session()935 headers = get_civitai_headers(effective_api_key)936 if attempt > 1:937 headers["Connection"] = "close"938 endpoint_url = ""939 result = None940 try:941 json_data, endpoint_url, result = request_civitai_api_json(942 endpoint_path,943 headers=headers,944 timeout=CIVITAI_METADATA_TIMEOUT,945 api_key=effective_api_key,946 session=session,947 stream=True,948 allow_not_found=True,949 )950 if result.status_code == 404:951 print(f"Civitai metadata lookup status=404: {endpoint_url}")952 cache_put(CIVITAI_VERSION_NEGATIVE_CACHE, model_version_id, "status=404")953 return None954 if not json_data:955 print(f"Civitai metadata lookup returned empty JSON: {endpoint_url}")956 cache_put(CIVITAI_VERSION_NEGATIVE_CACHE, model_version_id, "empty_json")957 return None958 cache_put(CIVITAI_VERSION_JSON_CACHE, model_version_id, copy.deepcopy(json_data))959 if normalized_url and normalized_url != raw_url:960 cache_put(CIVITAI_RESOLVE_CACHE, raw_url, normalized_url)961 return json_data962 except Exception as e:963 last_error = e964 print(965 f"[civitai] metadata reconnect attempt={attempt}/{CIVITAI_METADATA_RECONNECT_ATTEMPTS} "966 f"url={sanitize_url_for_log(endpoint_url or endpoint_path)} "967 f"error={type(e).__name__}: {sanitize_sensitive_log_text(e)}"968 )969 finally:970 try:971 if result is not None:972 result.close()973 except Exception:974 pass975 try:976 session.close()977 except Exception:978 pass979 if attempt < CIVITAI_METADATA_RECONNECT_ATTEMPTS:980 time.sleep(min(2.5, CIVITAI_METADATA_RECONNECT_BACKOFF * attempt))981 982 if last_error is not None:983 print(f"Civitai metadata lookup failed after reconnects: {type(last_error).__name__}: {sanitize_sensitive_log_text(last_error)}")984 return None985 986class ModelInformation:987 def __init__(self, json_data, source_url: str = ""):988 selected_file = pick_civitai_file_from_version_json(json_data, source_url=source_url)989 explicit_file_id = extract_civitai_file_id(source_url)990 self.model_version_id = json_data.get("id", "")991 self.model_id = json_data.get("modelId", "")992 self.download_url = selected_file.get("downloadUrl", "") or ("" if explicit_file_id else json_data.get("downloadUrl", ""))993 self.model_url = f"{get_civitai_canonical_web_origin()}/models/{self.model_id}?modelVersionId={self.model_version_id}"994 self.filename_url = selected_file.get("name", "") or ""995 self.description = json_data.get("description", "")996 if self.description is None:997 self.description = ""998 self.model_name = json_data.get("model", {}).get("name", "")999 self.model_type = json_data.get("model", {}).get("type", "")1000 self.base_model = json_data.get("baseModel", "")1001 self.nsfw = json_data.get("model", {}).get("nsfw", False)1002 self.poi = json_data.get("model", {}).get("poi", False)1003 self.images = [img.get("url", "") for img in json_data.get("images", [])]1004 self.example_prompt = json_data.get("trainedWords", [""])[0] if json_data.get("trainedWords") else ""1005 self.original_json = copy.deepcopy(json_data)1006 self.selected_file = copy.deepcopy(selected_file)1007 1008def retrieve_model_info(url, api_key: str = ""):1009 json_data = request_json_data(url, api_key=api_key)1010 if not json_data:1011 return None1012 model_descriptor = ModelInformation(json_data, source_url=url)1013 filters = get_civitai_query_filters(url)1014 if filters:1015 selected_summary = describe_civitai_file_for_log(model_descriptor.selected_file)1016 if selected_summary:1017 print(f"Civitai selected file: filters={filters} {selected_summary}")1018 else:1019 print(f"Civitai selected file: filters={filters} using model-level downloadUrl")1020 return model_descriptor1021 1022def list_downloaded_candidate_files(directory):1023 try:1024 return {1025 str(path.resolve())1026 for path in Path(directory).iterdir()1027 if path.is_file()1028 }1029 except Exception:1030 return set()1031 1032def sanitize_civitai_log_text(text: str):1033 output = str(text or "")1034 if not output:1035 return output1036 output = re.sub(r"([?&]token=)[^&\s\"']+", r"\1***", output, flags=re.IGNORECASE)1037 output = re.sub(r"([?&]Authorization=)[^&\s\"']+", r"\1***", output, flags=re.IGNORECASE)1038 return output1039 1040def sanitize_sensitive_log_text(text):1041 output = sanitize_civitai_log_text(text)1042 if not output:1043 return output1044 output = re.sub(r"(authorization:\s*bearer\s+)[^\s\"']+", r"\1***", output, flags=re.IGNORECASE)1045 output = re.sub(r"(bearer\s+)[^\s\"']+", r"\1***", output, flags=re.IGNORECASE)1046 return output1047 1048def log_download_error(scope: str, kind: str, url: str = "", status=None, error=None, detail: str = ""):1049 parts = [f"[{scope}] error={kind}"]1050 if status is not None:1051 parts.append(f"status={status}")1052 if url:1053 parts.append(f"url={sanitize_url_for_log(url)}")1054 if error is not None:1055 parts.append(f"exc={type(error).__name__}: {sanitize_sensitive_log_text(error)}")1056 elif detail:1057 parts.append(str(detail))1058 print(" ".join(parts))1059 1060def terminate_subprocess_safely(process, label: str = "subprocess"):1061 if process is None:1062 return1063 try:1064 if process.poll() is not None:1065 return1066 process.terminate()1067 process.wait(timeout=3)1068 except subprocess.TimeoutExpired:1069 try:1070 process.kill()1071 process.wait(timeout=3)1072 except Exception as e:1073 print(f"[{label}] kill failed: {type(e).__name__}: {sanitize_sensitive_log_text(e)}")1074 except Exception as e:1075 print(f"[{label}] terminate failed: {type(e).__name__}: {sanitize_sensitive_log_text(e)}")1076 1077def run_subprocess_capture(args, cwd=None, label: str = "subprocess"):1078 process = subprocess.Popen(1079 list(args),1080 cwd=str(cwd) if cwd else None,1081 stdout=subprocess.PIPE,1082 stderr=subprocess.PIPE,1083 text=True,1084 )1085 try:1086 stdout, stderr = process.communicate()1087 except BaseException:1088 terminate_subprocess_safely(process, label=label)1089 raise1090 output = "\n".join([part for part in [stdout, stderr] if part]).strip()1091 return int(process.returncode or 0), output1092 1093def build_civitai_wget_args(directory, download_url: str, filename: str = ""):1094 args = [1095 "wget",1096 "-c",1097 "-nv",1098 "--user-agent", USER_AGENT,1099 "--referer", CIVITAI_REFERER,1100 ]1101 if filename:1102 args.extend(["-O", str(Path(directory) / filename)])1103 else:1104 args.extend(["-P", str(directory)])1105 args.append(str(download_url))1106 return args1107 1108def run_civitai_wget(directory, download_url: str, filename: str = ""):1109 args = build_civitai_wget_args(directory, download_url, filename=filename)1110 return run_subprocess_capture(args, cwd=None, label="civitai-wget")1111 1112def build_generic_wget_args(directory, download_url: str):1113 return [1114 "wget",1115 "-c",1116 "-nv",1117 "-P", str(directory),1118 str(download_url),1119 ]1120 1121def run_generic_wget(directory, download_url: str):1122 args = build_generic_wget_args(directory, download_url)1123 return run_subprocess_capture(args, cwd=None, label="generic-wget")1124 1125def classify_civitai_download_failure(output_text: str):1126 text = str(output_text or "")1127 lower = text.lower()1128 if "status=403" in lower and "b2.civitai.com" in lower:1129 return "b2_403"1130 if "status=403" in lower and "civitai.com/api/download/models/" in lower:1131 return "api_403"1132 if "status=403" in lower:1133 return "http_403"1134 if "timed out" in lower or "timeout" in lower:1135 return "timeout"1136 return "other"1137 1138def cleanup_civitai_download_artifacts(directory, filename: str = ""):1139 removed = []1140 if not filename:1141 return removed1142 target = Path(directory) / filename1143 for candidate in [target]:1144 try:1145 if candidate.exists() and candidate.is_file():1146 candidate.unlink()1147 removed.append(str(candidate))1148 except Exception as e:1149 print(f"[civitai] cleanup failed path={candidate} {type(e).__name__}: {e}")1150 return removed1151 1152def guess_downloaded_file_path(directory, before_files, expected_filename=""):1153 expected_path = str(Path(directory) / expected_filename) if expected_filename else ""1154 if expected_path and Path(expected_path).exists():1155 return expected_path1156 1157 after_files = list_downloaded_candidate_files(directory)1158 new_files = sorted(list(after_files - set(before_files)))1159 if len(new_files) == 1:1160 return new_files[0]1161 1162 if expected_filename:1163 expected_name = str(expected_filename).strip()1164 stem = Path(expected_name).stem1165 suffix = Path(expected_name).suffix.lower()1166 matched = []1167 for path_str in new_files:1168 path_obj = Path(path_str)1169 if suffix and path_obj.suffix.lower() != suffix:1170 continue1171 if stem and (path_obj.stem == stem or path_obj.name == expected_name):1172 matched.append(path_str)1173 if len(matched) == 1:1174 return matched[0]1175 1176 return None1177 1178def get_civitai_expected_size_bytes(file_info):1179 if not isinstance(file_info, dict):1180 return 01181 raw_size_kb = file_info.get("sizeKB")1182 if raw_size_kb is None:1183 raw_size_kb = file_info.get("sizeKb")1184 try:1185 size_kb = float(raw_size_kb)1186 except (TypeError, ValueError):1187 return 01188 if size_kb <= 0:1189 return 01190 return max(1, int(round(size_kb * 1024.0)))1191 1192def is_civitai_file_complete(path, file_info=None, *, log_mismatch=False):1193 candidate = Path(path)1194 if not candidate.exists() or not candidate.is_file():1195 return False1196 expected_size = get_civitai_expected_size_bytes(file_info)1197 if expected_size <= 0:1198 return True1199 try:1200 actual_size = int(candidate.stat().st_size)