Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from huggingface_hub import HfApi, hf_hub_url | |
| from huggingface_hub.hf_api import RepoFile | |
| import os | |
| from pathlib import Path | |
| import gc | |
| import requests | |
| from requests.adapters import HTTPAdapter | |
| from urllib3.util import Retry | |
| import urllib | |
| import mimetypes | |
| from utils import (get_token, set_token, is_repo_exists, get_user_agent, get_download_file, | |
| list_uniq, list_sub, duplicate_hf_repo, HF_SUBFOLDER_NAME, get_state, set_state, | |
| create_retry_session, retry_call, ensure_repo, parse_civitai_api_keys, | |
| should_switch_civitai_key, resolve_civitai_download_url, suppress_hf_hub_progress_bars, | |
| reset_civitai_key_status, update_civitai_key_status, get_civitai_key_status, | |
| sanitize_url_for_log, HF_UPLOAD_RETRY_POLICY_CHOICES, | |
| get_hf_upload_retry_policy_config, hf_upload_retry_call, | |
| is_retryable_hf_upload_exception, parse_hf_retry_delay_from_headers, | |
| format_hf_rate_limit_hint, format_error_short) | |
| from bucket_ops import (is_bucket_api_available, ensure_bucket, upload_file_to_bucket, | |
| get_safe_bucket_filename, get_bucket_url) | |
| import re | |
| from PIL import Image, ImageOps | |
| import json | |
| import html as html_lib | |
| import pandas as pd | |
| import tempfile | |
| import hashlib | |
| import time | |
| import shutil | |
| import random | |
| import subprocess | |
| import threading | |
| import sys | |
| import platform | |
| import zipfile | |
| from datetime import datetime, timezone | |
| from io import BytesIO | |
| TEMP_DIR = tempfile.mkdtemp() | |
| CIVITAI_TEMP_ROOT = Path(TEMP_DIR) | |
| CIVITAI_BASEMODEL_REFRESH_PAGES_PER_SORT = 4 | |
| CIVITAI_BASEMODEL_MIN_COUNT = 12 | |
| SMOKE_TEST_LIMIT = 30 | |
| SMOKE_TEST_MAX_SIZE_KB = 200000 | |
| SMOKE_TEST_CANDIDATE_POOL = 10 | |
| SEARCH_PAGE_SIZE = 16 | |
| SEARCH_THUMB_SIZE = (320, 432) | |
| SEARCH_DETAIL_SIZE = (960, 1280) | |
| LOAD_ALL_BATCH_SIZE = 8 | |
| NULL_IMAGE_PATH = str(Path(__file__).with_name("null.png")) | |
| CIVITAI_DEFAULT_ORIGIN = "https://civitai.com" | |
| CIVITAI_CANONICAL_WEB_ORIGIN = CIVITAI_DEFAULT_ORIGIN | |
| CIVITAI_RED_ORIGIN = "https://civitai.red" | |
| CIVITAI_GREEN_ORIGIN = "https://civitai.green" | |
| 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) | |
| PREVIEW_VIDEO_EXTS = {".mp4", ".webm", ".mov", ".m4v", ".avi", ".mkv"} | |
| _FFMPEG_PATH = None | |
| _FFMPEG_MISSING_LOGGED = False | |
| _CREATOR_FETCH_WARNED = False | |
| _TAG_FETCH_WARNED = False | |
| CREATOR_SUGGEST_LIMIT = 200 | |
| CREATOR_CACHE_TTL_SEC = 600 | |
| REPO_HASH_CACHE_MAX_AGE_SEC = 900 | |
| RUN_CANCEL_REGISTRY = {} | |
| RUN_CANCEL_LOCK = threading.Lock() | |
| CREATOR_SUGGEST_CACHE = {} | |
| CREATOR_SUGGEST_LOCK = threading.Lock() | |
| CIVITAI_ACTIVE_API_ORIGIN = "" | |
| CIVITAI_ACTIVE_API_BASE = "" | |
| _CIVITAI_API_LOCK = threading.Lock() | |
| REPORT_EVENT_LIMIT = 800 | |
| REPORT_TEXT_LIMIT = 4000 | |
| REPORT_ZIP_PREFIX = "civitai_to_hf_report" | |
| def canonicalize_civitai_netloc(netloc: str): | |
| host = str(netloc or "").strip().lower() | |
| if host in CIVITAI_GREEN_HOST_ALIASES or host == "www.civitai.com": | |
| return "civitai.com" | |
| if host in CIVITAI_RED_HOST_ALIASES: | |
| return "civitai.red" | |
| return host | |
| def canonicalize_civitai_host(netloc: str): | |
| return canonicalize_civitai_netloc(netloc) | |
| def normalize_civitai_origin(value: str): | |
| raw = str(value or "").strip() | |
| if not raw: | |
| return CIVITAI_DEFAULT_ORIGIN | |
| parts = urllib.parse.urlsplit(raw if "://" in raw else f"https://{raw}") | |
| host = canonicalize_civitai_netloc(parts.netloc or parts.path) | |
| if host in {"civitai.com", "civitai.red"}: | |
| return f"https://{host}" | |
| return CIVITAI_DEFAULT_ORIGIN | |
| def get_civitai_canonical_web_origin(): | |
| return CIVITAI_CANONICAL_WEB_ORIGIN | |
| def get_civitai_display_origin(): | |
| return get_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 | |
| normalized = normalize_civitai_origin(origin) | |
| if normalized not in CIVITAI_API_ORIGIN_CANDIDATES: | |
| normalized = CIVITAI_DEFAULT_ORIGIN | |
| base = build_civitai_api_base(normalized) | |
| with _CIVITAI_API_LOCK: | |
| CIVITAI_ACTIVE_API_ORIGIN = normalized | |
| CIVITAI_ACTIVE_API_BASE = base | |
| return base | |
| def set_civitai_api_origin(origin: str): | |
| set_civitai_active_api_origin(origin) | |
| return CIVITAI_ACTIVE_API_ORIGIN | |
| def get_civitai_api_origin(): | |
| with _CIVITAI_API_LOCK: | |
| return CIVITAI_ACTIVE_API_ORIGIN or "" | |
| def probe_civitai_api_origin(session, origin: str, timeout: tuple[float, float] = (3.0, 8.0)): | |
| response = None | |
| try: | |
| base_url = build_civitai_api_base(origin) | |
| if not base_url: | |
| return False | |
| response = session.get( | |
| f"{base_url}/tags", | |
| params={"limit": 1}, | |
| headers=get_civitai_headers(""), | |
| timeout=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) | |
| 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, session=None): | |
| cached_origin = get_civitai_api_origin() | |
| if cached_origin and not force_refresh: | |
| return cached_origin | |
| if session is None: | |
| session = create_retry_session(total=4, backoff_factor=0.8) | |
| for origin in CIVITAI_API_ORIGIN_CANDIDATES: | |
| if probe_civitai_api_origin(session, origin): | |
| set_civitai_active_api_origin(origin) | |
| return get_civitai_api_origin() | |
| set_civitai_active_api_origin(CIVITAI_DEFAULT_ORIGIN) | |
| return get_civitai_api_origin() | |
| def resolve_civitai_api_origin(session=None): | |
| return get_civitai_active_api_origin(session=session) | |
| def get_civitai_active_api_base(force_refresh: bool = False, session=None): | |
| with _CIVITAI_API_LOCK: | |
| cached_base = CIVITAI_ACTIVE_API_BASE or "" | |
| if cached_base and not force_refresh: | |
| return cached_base | |
| get_civitai_active_api_origin(force_refresh=force_refresh, session=session) | |
| with _CIVITAI_API_LOCK: | |
| return CIVITAI_ACTIVE_API_BASE or build_civitai_api_base(CIVITAI_DEFAULT_ORIGIN) | |
| def get_civitai_api_candidate_origins(preferred: str = ""): | |
| origins = list(CIVITAI_API_ORIGIN_CANDIDATES) | |
| preferred_origin = normalize_civitai_origin(preferred) if preferred else "" | |
| if preferred_origin in origins: | |
| return [preferred_origin] + [origin for origin in origins if origin != preferred_origin] | |
| cached_origin = get_civitai_api_origin() | |
| if cached_origin in origins: | |
| return [cached_origin] + [origin for origin in origins if origin != cached_origin] | |
| return origins | |
| def iter_civitai_api_bases(preferred: str = "", session=None): | |
| preferred_origin = normalize_civitai_origin(preferred) if preferred else get_civitai_active_api_origin(session=session) | |
| return [build_civitai_api_base(origin) for origin in get_civitai_api_candidate_origins(preferred_origin)] | |
| def build_civitai_api_url(path: str, origin: str = ""): | |
| clean_path = "/" + str(path or "").strip().lstrip("/") | |
| if not clean_path.startswith("/api/"): | |
| clean_path = f"/api/v1{clean_path}" | |
| base_origin = normalize_civitai_origin(origin) if origin else normalize_civitai_origin(get_civitai_api_origin() or CIVITAI_DEFAULT_ORIGIN) | |
| return f"{base_origin.rstrip('/')}" + clean_path | |
| def build_civitai_model_url(model_id, model_version_id=None): | |
| if model_id is None: | |
| return "" | |
| url = f"{get_civitai_canonical_web_origin().rstrip('/')}/models/{model_id}" | |
| if model_version_id is not None: | |
| url += f"?modelVersionId={model_version_id}" | |
| return url | |
| def should_fallback_civitai_api_response(response): | |
| if response is None: | |
| return True | |
| try: | |
| status_code = int(getattr(response, "status_code", 0) or 0) | |
| except Exception: | |
| return True | |
| return status_code in {404, 405, 408, 429} or status_code >= 500 | |
| def request_civitai_api(session, path: str, api_key: str = "", params=None, timeout: tuple[float, float] = (7.0, 30.0), | |
| label: str = "Civitai API", source: str = "civitai-api", preferred_origin: str = "", | |
| non_json_fallback_origin: str = ""): | |
| preferred = normalize_civitai_origin(preferred_origin) if preferred_origin else get_civitai_active_api_origin(session=session) | |
| non_json_fallback = normalize_civitai_origin(non_json_fallback_origin) if non_json_fallback_origin else "" | |
| last_response = None | |
| last_exception = None | |
| for origin in get_civitai_api_candidate_origins(preferred): | |
| url = build_civitai_api_url(path, origin=origin) | |
| try: | |
| response = civitai_get(session, url, api_key=api_key, params=params, timeout=timeout, label=label, source=source) | |
| last_response = response | |
| if response is not None and response.ok: | |
| try: | |
| get_civitai_response_json(response) | |
| except Exception as e: | |
| last_exception = e | |
| log_line("retry", f"{label}: non-json response from {origin}") | |
| try: | |
| response.close() | |
| except Exception: | |
| pass | |
| if non_json_fallback and origin != non_json_fallback: | |
| fallback_url = build_civitai_api_url(path, origin=non_json_fallback) | |
| try: | |
| fallback_response = civitai_get(session, fallback_url, api_key=api_key, params=params, timeout=timeout, label=label, source=source) | |
| last_response = fallback_response | |
| if fallback_response is not None and fallback_response.ok: | |
| get_civitai_response_json(fallback_response) | |
| set_civitai_active_api_origin(non_json_fallback) | |
| return fallback_response | |
| if fallback_response is not None and not should_fallback_civitai_api_response(fallback_response): | |
| return fallback_response | |
| if fallback_response is not None: | |
| fallback_response.close() | |
| except Exception as fallback_error: | |
| last_exception = fallback_error | |
| continue | |
| set_civitai_active_api_origin(origin) | |
| return response | |
| if not should_fallback_civitai_api_response(response): | |
| return response | |
| if response is not None: | |
| response.close() | |
| except Exception as e: | |
| last_exception = e | |
| continue | |
| if last_exception is not None: | |
| raise last_exception | |
| if last_response is not None: | |
| return last_response | |
| return None | |
| _CIVITAI_JSON_MISSING = object() | |
| def get_civitai_response_json(response, default=_CIVITAI_JSON_MISSING): | |
| if response is None: | |
| if default is not _CIVITAI_JSON_MISSING: | |
| return default | |
| raise ValueError("Civitai response is missing") | |
| cached = getattr(response, "_civitai_json", _CIVITAI_JSON_MISSING) | |
| if cached is not _CIVITAI_JSON_MISSING: | |
| return cached | |
| try: | |
| payload = response.json() | |
| except Exception: | |
| if default is not _CIVITAI_JSON_MISSING: | |
| return default | |
| raise | |
| setattr(response, "_civitai_json", payload) | |
| return payload | |
| def request_civitai_api_url(session, url: str, api_key: str = "", params=None, timeout: tuple[float, float] = (7.0, 30.0), | |
| label: str = "Civitai API", source: str = "civitai-api"): | |
| parts = get_civitai_url_parts(url) | |
| path = str(parts.path or "") | |
| if is_civitai_host(parts.netloc) and path.startswith("/api/"): | |
| api_path = path + (f"?{parts.query}" if parts.query else "") | |
| return request_civitai_api(session, api_path, api_key=api_key, params=params, timeout=timeout, label=label, source=source) | |
| response = civitai_get(session, url, api_key=api_key, params=params, timeout=timeout, label=label, source=source) | |
| get_civitai_response_json(response) | |
| return response | |
| CIVITAI_TYPE = ["Checkpoint", "TextualInversion", "Hypernetwork", "AestheticGradient", "LORA", "LoCon", "DoRA", | |
| "Controlnet", "Upscaler", "MotionModule", "VAE", "Poses", "Wildcards", "Workflows", "Other"] | |
| CIVITAI_FILETYPE = ["Model", "VAE", "Config", "Training Data", "Archive", "Negative"] | |
| 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"] | |
| def parse_urls(s): | |
| url_pattern = "https?://[\\w/:%#\\$&\\?\\(\\)~\\.=\\+\\-]+" | |
| try: | |
| urls = re.findall(url_pattern, s) | |
| return list(urls) | |
| except Exception: | |
| return [] | |
| def parse_repos(s): | |
| repo_pattern = r'[^\w_\-\.]?([\w_\-\.]+/[\w_\-\.]+)[^\w_\-\.]?' | |
| try: | |
| s = re.sub("https?://[\\w/:%#\\$&\\?\\(\\)~\\.=\\+\\-]+", "", s) | |
| repos = re.findall(repo_pattern, s) | |
| return list(repos) | |
| except Exception: | |
| return [] | |
| def to_urls(l: list[str]): | |
| return "\n".join(l) | |
| def normalize_input_token(value: str): | |
| token = str(value or "").strip().strip("\"'") | |
| while token.endswith(",") or token.endswith(";"): | |
| token = token[:-1].rstrip() | |
| return token | |
| def normalize_url_entries(value): | |
| return list_uniq([token for token in [normalize_input_token(url) for url in parse_urls(str(value or ""))] if token]) | |
| def normalize_repo_entries(value): | |
| return list_uniq([token for token in [normalize_input_token(repo) for repo in parse_repos(str(value or ""))] if token]) | |
| def uniq_urls(s): | |
| return to_urls(list_uniq(normalize_url_entries(s) + normalize_repo_entries(s))) | |
| def create_run_temp_dir(): | |
| CIVITAI_TEMP_ROOT.mkdir(parents=True, exist_ok=True) | |
| return tempfile.mkdtemp(prefix="run_", dir=str(CIVITAI_TEMP_ROOT)) | |
| def cleanup_run_temp_dir(path: str): | |
| try: | |
| if path and Path(path).exists(): shutil.rmtree(path, ignore_errors=True) | |
| except Exception as e: | |
| print(f"[cleanup] Failed to cleanup temp dir {path}. {e}") | |
| def is_safe_run_temp_dir(path: str): | |
| try: | |
| if not path: | |
| return False | |
| target = Path(path).resolve() | |
| root = CIVITAI_TEMP_ROOT.resolve() | |
| return str(target).startswith(str(root)) | |
| except Exception: | |
| return False | |
| def log_line(prefix: str, message: str): | |
| tag = str(prefix or "info").strip() or "info" | |
| print(f"[{tag}] {message}") | |
| def utc_timestamp(): | |
| try: | |
| return datetime.now(timezone.utc).isoformat(timespec="seconds") | |
| except Exception: | |
| return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) | |
| def redact_report_value(value): | |
| if isinstance(value, dict): | |
| return {str(k): redact_report_value(v) for k, v in value.items()} | |
| if isinstance(value, (list, tuple, set)): | |
| return [redact_report_value(v) for v in value] | |
| text = str(value if value is not None else "") | |
| if not text: | |
| return "" | |
| text = re.sub(r"hf_[A-Za-z0-9]{20,}", "[redacted-hf-token]", text) | |
| text = re.sub(r"(?i)(token=)[^\s&]+", r"\1[redacted-token]", text) | |
| text = re.sub(r"(?i)(Authorization:\s*Bearer\s+)[^\s]+", r"\1[redacted-token]", text) | |
| text = re.sub(r"(?i)(xet-read-token/)[^\s/?#]+", r"\1[redacted-xet-token]", text) | |
| text = re.sub(r"(?i)(X-Amz-Signature=)[^\s&]+", r"\1[redacted-signature]", text) | |
| text = re.sub(r"(?i)(X-Amz-Credential=)[^\s&]+", r"\1[redacted-credential]", text) | |
| text = re.sub(r"(?i)(Key-Pair-Id=)[^\s&]+", r"\1[redacted-key-pair]", text) | |
| if len(text) > REPORT_TEXT_LIMIT: | |
| return text[:REPORT_TEXT_LIMIT] + "...[truncated]" | |
| return text | |
| def append_report_event(session_state, event: str, **fields): | |
| state = ensure_session_state(session_state) | |
| events = list(state.get("report_events") or []) | |
| session_events = list(state.get("session_report_events") or []) | |
| explicit_run_id = fields.pop("run_id", "") if "run_id" in fields else "" | |
| run_id = str(explicit_run_id or state.get("active_run_id") or state.get("last_run_id") or "") | |
| clean = {"ts": utc_timestamp(), "event": str(event or "event")} | |
| if run_id: | |
| clean["run_id"] = run_id | |
| for key, value in fields.items(): | |
| if key in {"hf_token", "civitai_key", "authorization", "cookie"}: | |
| clean[str(key)] = "[redacted]" if value else "" | |
| else: | |
| clean[str(key)] = redact_report_value(value) | |
| events.append(clean) | |
| session_events.append(clean) | |
| if len(events) > REPORT_EVENT_LIMIT: | |
| events = events[-REPORT_EVENT_LIMIT:] | |
| session_event_limit = max(REPORT_EVENT_LIMIT * 4, REPORT_EVENT_LIMIT) | |
| if len(session_events) > session_event_limit: | |
| session_events = session_events[-session_event_limit:] | |
| session_state_update(state, report_events=events, session_report_events=session_events) | |
| return state | |
| def report_write_text(zipf, name: str, text: str): | |
| zipf.writestr(name, redact_report_value(text)) | |
| def safe_json_dumps(value): | |
| return json.dumps(redact_report_value(value), ensure_ascii=False, indent=2, sort_keys=True) | |
| def list_report_events_for_run(events, run_id: str): | |
| target = str(run_id or "") | |
| if not target: | |
| return list(events or []) | |
| return [ev for ev in list(events or []) if str(ev.get("run_id") or "") == target] | |
| def summarize_report_runs(run_records): | |
| records = list(run_records or []) | |
| summary = { | |
| "runs": len(records), | |
| "done": 0, | |
| "incomplete": 0, | |
| "failed": 0, | |
| "cancelled": 0, | |
| "input_urls": 0, | |
| "downloaded": 0, | |
| "uploaded": 0, | |
| "skipped_duplicate": 0, | |
| "failed_download": 0, | |
| "failed_upload": 0, | |
| "verified_after_error": 0, | |
| "remaining": 0, | |
| "failed_urls": 0, | |
| } | |
| for record in records: | |
| run_summary = record.get("summary") if isinstance(record, dict) else {} | |
| if not isinstance(run_summary, dict): | |
| run_summary = {} | |
| stage = str(run_summary.get("stage") or record.get("stage") or "").lower() | |
| if stage == "done": | |
| summary["done"] += 1 | |
| elif stage == "incomplete": | |
| summary["incomplete"] += 1 | |
| elif stage == "failed": | |
| summary["failed"] += 1 | |
| elif stage == "cancelled": | |
| summary["cancelled"] += 1 | |
| for key in ("input_urls", "downloaded", "uploaded", "skipped_duplicate", "failed_download", "failed_upload", "verified_after_error", "remaining"): | |
| try: | |
| summary[key] += int(run_summary.get(key, 0) or 0) | |
| except Exception: | |
| pass | |
| try: | |
| summary["failed_urls"] += int(run_summary.get("failed", 0) or 0) | |
| except Exception: | |
| pass | |
| return summary | |
| def build_report_run_record(run_id: str, summary: dict, events, remaining, failed, uploaded, smoke_lines=None, failure_reasons=None): | |
| clean_summary = redact_report_value(dict(summary or {})) | |
| clean_events = redact_report_value(list(events or [])) | |
| clean_remaining = redact_report_value(list(remaining or [])) | |
| clean_failed = redact_report_value(list(failed or [])) | |
| clean_uploaded = redact_report_value(list(uploaded or [])) | |
| record = { | |
| "run_id": str(run_id or clean_summary.get("run_id") or ""), | |
| "mode": clean_summary.get("mode") or "", | |
| "stage": clean_summary.get("stage") or "", | |
| "repo_id": clean_summary.get("repo_id") or "", | |
| "repo_type": clean_summary.get("repo_type") or "", | |
| "summary": clean_summary, | |
| "events": clean_events, | |
| "remaining_urls": clean_remaining, | |
| "failed_urls": clean_failed, | |
| "uploaded_urls": clean_uploaded, | |
| "smoke_lines": redact_report_value(list(smoke_lines or [])), | |
| "failure_reasons": redact_report_value(dict(failure_reasons or {})), | |
| "advice": build_report_advice(clean_summary), | |
| } | |
| return record | |
| def append_session_run_record(session_state, record, limit: int = 20): | |
| state = ensure_session_state(session_state) | |
| records = list(state.get("session_run_records") or []) | |
| run_id = str(record.get("run_id") or "") | |
| if run_id: | |
| records = [r for r in records if str(r.get("run_id") or "") != run_id] | |
| records.append(redact_report_value(record)) | |
| if len(records) > limit: | |
| records = records[-limit:] | |
| session_state_update(state, session_run_records=records, last_run_record=record) | |
| return records | |
| def get_package_version(package_name: str): | |
| try: | |
| from importlib import metadata | |
| return metadata.version(package_name) | |
| except Exception: | |
| return "unknown" | |
| def build_report_advice(summary: dict): | |
| summary = summary if isinstance(summary, dict) else {} | |
| failed_upload = int(summary.get("failed_upload", 0) or 0) | |
| failed_download = int(summary.get("failed_download", 0) or 0) | |
| remaining = int(summary.get("remaining", 0) or 0) | |
| skipped = int(summary.get("skipped_duplicate", 0) or 0) | |
| verified = int(summary.get("verified_after_error", 0) or 0) | |
| lines = ["# What to try next", ""] | |
| if remaining: | |
| lines.append(f"- {remaining} URL(s) remained unprocessed. Use the remaining URL list and rerun. This usually means interruption/cancel/timeout before all URLs were processed.") | |
| if failed_upload: | |
| lines.append(f"- {failed_upload} item(s) failed after Civitai download. This points to HF upload/LFS/commit side. Retry failed only; if repeated, wait a few minutes or use the Patient HF upload retry policy.") | |
| if failed_download: | |
| lines.append(f"- {failed_download} item(s) failed before upload. Retry later; if repeated, check Civitai visibility/login/API-key status for those specific items.") | |
| if skipped: | |
| lines.append(f"- {skipped} item(s) were skipped as duplicates by SHA256. This is expected and usually does not need retry.") | |
| if verified: | |
| lines.append(f"- {verified} upload error(s) were recovered because the remote file existed after the API error. This suggests HF commit/LFS returned an error after partial success.") | |
| if not any([remaining, failed_upload, failed_download]): | |
| lines.append("- No retry is needed based on the recorded summary.") | |
| lines.append("") | |
| lines.append("# Notes") | |
| lines.append("- Tokens, signed URLs, Authorization headers, and transient redirect tokens are redacted.") | |
| lines.append("- This report contains structured Space-side state, not raw container stdout.") | |
| return "\n".join(lines) + "\n" | |
| def summarize_failure_text(message: str, url: str=""): | |
| text = str(message or "").strip().replace("\n", " ") | |
| text = re.sub(r"\s+", " ", text) | |
| if url: | |
| text = f"{text} @ {sanitize_url_for_log(url)}" if text else sanitize_url_for_log(url) | |
| return text[:240] | |
| def set_last_failure_summary(session_state, message: str, url: str=""): | |
| summary = summarize_failure_text(message, url=url) | |
| session_state_update(session_state, last_failure_summary=summary) | |
| return summary | |
| def prepare_new_run_state(session_state): | |
| state = ensure_session_state(session_state) | |
| previous_run_id = str(state.get("active_run_id") or "") | |
| previous_temp_dir = str(state.get("current_run_temp_dir") or "") | |
| if previous_run_id: | |
| unregister_run(previous_run_id) | |
| if previous_temp_dir and is_safe_run_temp_dir(previous_temp_dir): | |
| cleanup_run_temp_dir(previous_temp_dir) | |
| session_state_update( | |
| state, | |
| current_run_temp_dir="", | |
| active_run_id="", | |
| cancel_requested=False, | |
| current_stage="", | |
| current_stage_detail="", | |
| current_url="", | |
| current_item_index=0, | |
| current_item_total=0, | |
| current_remaining_urls=[], | |
| current_failed_urls=[], | |
| current_uploaded_urls=[], | |
| current_smoke_lines=[], | |
| last_error="", | |
| last_failure_summary="", | |
| repo_hash_cache={}, | |
| run_started_at=0.0, | |
| run_elapsed_sec=0.0, | |
| ) | |
| return state | |
| def get_civitai_headers(api_key: str=""): | |
| user_agent = get_user_agent() | |
| headers = {'User-Agent': user_agent, 'content-type': 'application/json'} | |
| if api_key: headers['Authorization'] = f'Bearer {api_key}' | |
| return headers | |
| def ensure_session_state(session_state): | |
| return session_state if isinstance(session_state, dict) else {} | |
| def session_state_update(session_state, **kwargs): | |
| state = ensure_session_state(session_state) | |
| for key, value in kwargs.items(): | |
| set_state(state, key, value) | |
| return state | |
| def session_state_output(session_state): | |
| state = ensure_session_state(session_state) | |
| return dict(state) | |
| class RunCancelledError(RuntimeError): | |
| pass | |
| def new_run_id(): | |
| return f"run-{time.time_ns()}-{random.randint(1000, 9999)}" | |
| def register_run(run_id: str): | |
| if not run_id: | |
| return | |
| with RUN_CANCEL_LOCK: | |
| RUN_CANCEL_REGISTRY[run_id] = {"cancel_requested": False, "updated_at": time.time()} | |
| def unregister_run(run_id: str): | |
| if not run_id: | |
| return | |
| with RUN_CANCEL_LOCK: | |
| RUN_CANCEL_REGISTRY.pop(run_id, None) | |
| def request_run_cancel(session_state=None): | |
| session_state = ensure_session_state(session_state) | |
| run_id = str(session_state.get("active_run_id") or "") | |
| if run_id: | |
| with RUN_CANCEL_LOCK: | |
| entry = RUN_CANCEL_REGISTRY.get(run_id) or {} | |
| entry["cancel_requested"] = True | |
| entry["updated_at"] = time.time() | |
| RUN_CANCEL_REGISTRY[run_id] = entry | |
| log_line("cancel", f"requested for {run_id}") | |
| session_state_update(session_state, cancel_requested=True, current_stage="Cancel requested", current_stage_detail="Waiting for a safe stop point.") | |
| return session_state_output(session_state) | |
| def is_run_cancel_requested(run_id: str, session_state=None): | |
| if isinstance(session_state, dict) and session_state.get("cancel_requested"): | |
| return True | |
| if not run_id: | |
| return False | |
| with RUN_CANCEL_LOCK: | |
| entry = RUN_CANCEL_REGISTRY.get(run_id) or {} | |
| return bool(entry.get("cancel_requested")) | |
| def check_run_cancel(run_id: str, session_state=None): | |
| if is_run_cancel_requested(run_id, session_state=session_state): | |
| raise RunCancelledError("Cancelled by user.") | |
| def update_run_stage(session_state, stage: str, detail: str="", index: int=0, total: int=0, current_url: str=""): | |
| state = ensure_session_state(session_state) | |
| started_at = float(state.get("run_started_at") or 0.0) | |
| elapsed_sec = max(0.0, time.time() - started_at) if started_at > 0 else 0.0 | |
| session_state_update( | |
| state, | |
| current_stage=str(stage or ""), | |
| current_stage_detail=str(detail or ""), | |
| current_item_index=int(index or 0), | |
| current_item_total=int(total or 0), | |
| current_url=str(current_url or ""), | |
| run_elapsed_sec=elapsed_sec, | |
| ) | |
| def format_run_status_markdown(session_state): | |
| state = ensure_session_state(session_state) | |
| stage = str(state.get("current_stage") or "").strip() | |
| detail = str(state.get("current_stage_detail") or "").strip() | |
| repo_id = str(state.get("current_repo_id") or "").strip() | |
| repo_type = str(state.get("current_repo_type") or "").strip() | |
| item_index = int(state.get("current_item_index") or 0) | |
| item_total = int(state.get("current_item_total") or 0) | |
| current_url = str(state.get("current_url") or "").strip() | |
| cancel_requested = bool(state.get("cancel_requested")) | |
| last_failure = str(state.get("last_failure_summary") or "").strip() | |
| uploaded_count = len(state.get("current_uploaded_urls") or []) | |
| failed_count = len(state.get("current_failed_urls") or []) | |
| remaining_count = len(state.get("current_remaining_urls") or []) | |
| elapsed_sec = float(state.get("run_elapsed_sec") or 0.0) | |
| key_status = get_civitai_key_status("") | |
| key_count = int(key_status.get("count") or 0) | |
| active_index = int(key_status.get("active_index") or (1 if key_count else 0)) | |
| if not any([stage, detail, repo_id, current_url, cancel_requested, last_failure, uploaded_count, failed_count, remaining_count, elapsed_sec]): | |
| return "" | |
| stage_label = stage or "Idle" | |
| if item_total > 0: | |
| stage_label += f" ({item_index}/{item_total})" | |
| if detail: | |
| stage_label += f" - {detail}" | |
| parts = [f"**Status**: {stage_label}"] | |
| if repo_id: | |
| target = f"{repo_type}:{repo_id}" if repo_type else repo_id | |
| parts.append(f"**Target**: `{target}`") | |
| if current_url: | |
| parts.append(f"**URL**: `{current_url}`") | |
| total_count = uploaded_count + failed_count + remaining_count | |
| if total_count > 0: | |
| counts = f"processed {uploaded_count + failed_count}/{total_count} | uploaded {uploaded_count} | failed {failed_count} | remaining {remaining_count}" | |
| parts.append(f"**Counts**: {counts}") | |
| if elapsed_sec > 0: | |
| parts.append(f"**Elapsed**: {int(elapsed_sec)}s") | |
| if key_count: | |
| parts.append(f"**Key**: {min(max(active_index, 1), key_count)}/{key_count}") | |
| if cancel_requested and stage != "Cancelled": | |
| parts.append("**Cancel**: requested") | |
| if last_failure: | |
| parts.append(f"**Last failure**: {last_failure}") | |
| return " | ".join(parts) | |
| def build_run_status_update(session_state): | |
| status_md = format_run_status_markdown(session_state) | |
| return gr.update(value=status_md, visible=bool(status_md)) | |
| def get_session_repo_hash_cache(session_state, repo_id: str, repo_type: str): | |
| state = ensure_session_state(session_state) | |
| cache = state.get("repo_hash_cache") | |
| if not isinstance(cache, dict): | |
| return None | |
| if str(cache.get("repo_id") or "") != str(repo_id or ""): | |
| return None | |
| if str(cache.get("repo_type") or "") != str(repo_type or ""): | |
| return None | |
| cached_at = float(cache.get("cached_at") or 0.0) | |
| if cached_at <= 0 or (time.time() - cached_at) > REPO_HASH_CACHE_MAX_AGE_SEC: | |
| return None | |
| hashes = cache.get("hashes") or [] | |
| return {str(h) for h in hashes if h} | |
| def store_session_repo_hash_cache(session_state, repo_id: str, repo_type: str, hashes): | |
| state = ensure_session_state(session_state) | |
| payload = { | |
| "repo_id": str(repo_id or ""), | |
| "repo_type": str(repo_type or ""), | |
| "hashes": sorted({str(h) for h in (hashes or []) if h}), | |
| "cached_at": time.time(), | |
| } | |
| state["repo_hash_cache"] = payload | |
| return payload | |
| def smoke_stage_line(name: str, status: str, detail: str=""): | |
| state = "ok" if status == "ok" else "fail" | |
| return f"[{state}] {name}: {detail}" if detail else f"[{state}] {name}" | |
| def build_run_markdown(repo_header: str, result_lines=None, smoke_lines=None): | |
| result_lines = result_lines or [] | |
| smoke_lines = smoke_lines or [] | |
| parts = [] | |
| if smoke_lines: | |
| parts.append("### Smoke Test\n" + "\n".join([f"- {line}" for line in smoke_lines])) | |
| if repo_header: | |
| parts.append(repo_header.rstrip()) | |
| if result_lines: | |
| parts.append("\n".join(result_lines)) | |
| return "\n".join([p for p in parts if p]) + "\n" | |
| def format_civitai_key_status_md(api_key: str=""): | |
| status = get_civitai_key_status(api_key) | |
| count = int(status.get("count") or len(parse_civitai_api_keys(api_key))) | |
| active_index = int(status.get("active_index") or (1 if count else 0)) | |
| source = str(status.get("source") or "") | |
| last_reason = str(status.get("last_switch_reason") or "") | |
| last_status = str(status.get("last_status") or "") | |
| parts = [f"Civitai keys: {count}"] | |
| if count: | |
| parts.append(f"active: {min(max(active_index, 1), count)}/{count}") | |
| if source: | |
| parts.append(f"source: {source}") | |
| if last_status: | |
| parts.append(f"last status: {last_status}") | |
| if last_reason: | |
| parts.append(f"switch: {last_reason[:120]}") | |
| return " | ".join(parts) | |
| def build_run_outputs(urls, md: str, remain_urls, failed_urls, civitai_key, session_state, remain_visible=None, failed_visible=None): | |
| remain_text = "\n".join(remain_urls) if remain_urls else "" | |
| failed_text = "\n".join(failed_urls) if failed_urls else "" | |
| remain_update = gr.update(value=remain_text) if remain_visible is None else gr.update(value=remain_text, visible=remain_visible) | |
| failed_update = gr.update(value=failed_text) if failed_visible is None else gr.update(value=failed_text, visible=failed_visible) | |
| key_status_update = gr.update(value=format_civitai_key_status_md(civitai_key)) | |
| final_md = md.strip() | |
| if final_md: | |
| final_md += "\n" | |
| return gr.update(value=urls, choices=urls), gr.update(value=final_md), remain_update, failed_update, key_status_update, session_state_output(session_state) | |
| def set_stage_progress(progress, current: int, total: int, desc: str): | |
| try: | |
| progress((current, total), desc=desc) | |
| except Exception: | |
| pass | |
| def stage_detail(label: str, enabled: bool): | |
| return "enabled" if enabled else f"skipped ({label} off)" | |
| def new_run_stats(total_urls: int): | |
| return { | |
| "input_urls": int(total_urls or 0), | |
| "downloaded": 0, | |
| "uploaded": 0, | |
| "skipped_duplicate": 0, | |
| "failed_download": 0, | |
| "failed_upload": 0, | |
| "failed_info": 0, | |
| "verified_after_error": 0, | |
| } | |
| def increment_run_stat(stats: dict, key: str, amount: int=1): | |
| if isinstance(stats, dict): | |
| stats[key] = int(stats.get(key, 0) or 0) + int(amount) | |
| return stats | |
| def build_run_summary_lines(stats: dict, remain_urls, failed_urls, final_stage: str): | |
| stats = stats if isinstance(stats, dict) else {} | |
| remain_count = len(remain_urls or []) | |
| failed_count = len(failed_urls or []) | |
| lines = [ | |
| "", | |
| "### Run summary", | |
| f"- Status: **{final_stage}**", | |
| f"- Input URLs: {int(stats.get('input_urls', 0) or 0)}", | |
| f"- Downloaded: {int(stats.get('downloaded', 0) or 0)}", | |
| f"- Uploaded: {int(stats.get('uploaded', 0) or 0)}", | |
| f"- Skipped duplicate: {int(stats.get('skipped_duplicate', 0) or 0)}", | |
| f"- Failed download: {int(stats.get('failed_download', 0) or 0)}", | |
| f"- Failed upload: {int(stats.get('failed_upload', 0) or 0)}", | |
| f"- Remaining: {remain_count}", | |
| f"- Failed URL list: {failed_count}", | |
| ] | |
| verified_after_error = int(stats.get('verified_after_error', 0) or 0) | |
| if verified_after_error: | |
| lines.append(f"- Upload verified after API error: {verified_after_error}") | |
| lines.append("") | |
| lines.append("### What to try next") | |
| if remain_count: | |
| lines.append("- Some URLs are still remaining. Use **Use Remaining URLs** and run again; this usually means the run was interrupted or stopped before all URLs were processed.") | |
| if int(stats.get('failed_upload', 0) or 0): | |
| lines.append("- Some Civitai downloads completed but HF upload failed. Use **Retry Failed Only**; if it repeats, wait a few minutes or switch HF upload retry policy to **Patient**.") | |
| if int(stats.get('failed_download', 0) or 0): | |
| lines.append("- Some Civitai downloads failed or produced no file. Retry them later; if repeated, check Civitai visibility/login/API-key settings for those files.") | |
| if int(stats.get('skipped_duplicate', 0) or 0): | |
| lines.append("- Some files were skipped because matching SHA256 already exists in the target repo. This is expected and does not require retry.") | |
| if not remain_count and not failed_count and not int(stats.get('failed_upload', 0) or 0) and not int(stats.get('failed_download', 0) or 0): | |
| lines.append("- No retry is needed.") | |
| return lines | |
| def log_run_summary(run_mode: str, final_stage: str, stats: dict, remain_urls, failed_urls): | |
| stats = stats if isinstance(stats, dict) else {} | |
| log_line( | |
| "cleanup", | |
| "summary " | |
| f"mode={run_mode} stage={final_stage.lower()} input={int(stats.get('input_urls', 0) or 0)} " | |
| f"downloaded={int(stats.get('downloaded', 0) or 0)} uploaded={int(stats.get('uploaded', 0) or 0)} " | |
| f"skipped_duplicate={int(stats.get('skipped_duplicate', 0) or 0)} " | |
| f"failed_download={int(stats.get('failed_download', 0) or 0)} " | |
| f"failed_upload={int(stats.get('failed_upload', 0) or 0)} " | |
| f"remaining={len(remain_urls or [])} failed={len(failed_urls or [])}" | |
| ) | |
| def verify_repo_upload(repo_id: str, repo_type: str, filename: str, api: HfApi | None = None, hf_token=None): | |
| if hf_token is None: hf_token = get_token() | |
| if api is None: api = HfApi(token=hf_token) | |
| try: | |
| return bool(retry_call(lambda: api.file_exists(repo_id=repo_id, filename=filename, repo_type=repo_type, token=hf_token), action=f'file_exists {repo_id}:{filename}')) | |
| except Exception as e: | |
| print(f"Smoke test upload verify failed for {repo_id}:{filename}. {e}") | |
| return False | |
| def summarize_downloaded_file(filename: str): | |
| path = Path(filename) | |
| if not path.exists() or not path.is_file(): | |
| return False, "missing file" | |
| size = path.stat().st_size | |
| if size <= 0: | |
| return False, f"{path.name} is 0 bytes" | |
| return True, f"{path.name} / {round(size / 1000.0 / 1000.0, 2)}MB" | |
| def civitai_get(session, url: str, *, api_key: str="", params=None, timeout=(7.0, 30), label: str="Civitai GET", source: str="civitai-get"): | |
| keys = parse_civitai_api_keys(api_key) | |
| if not keys: | |
| keys = [""] | |
| last_response = None | |
| total_keys = len(keys) | |
| for index, key in enumerate(keys, start=1): | |
| update_civitai_key_status(raw=api_key, active_index=index if total_keys else 0, source=source) | |
| headers = get_civitai_headers(key) | |
| r = session.get(url, params=params, headers=headers, stream=True, timeout=timeout) | |
| last_response = r | |
| retry_after = str(r.headers.get('Retry-After', '') or '').strip() | |
| if r.status_code == 429: | |
| log_line("retry", f"{label}: key {index}/{total_keys} status=429 retry_after={retry_after or '-'}") | |
| elif not r.ok: | |
| log_line("retry", f"{label}: key {index}/{total_keys} status={r.status_code}") | |
| update_civitai_key_status(raw=api_key, active_index=index if total_keys else 0, last_status=str(r.status_code), source=source) | |
| if r.ok or not should_switch_civitai_key(r.status_code) or index >= total_keys: | |
| return r | |
| reason = f"{source} status={r.status_code} key={index}/{total_keys}" | |
| if retry_after: | |
| reason += f" retry_after={retry_after}" | |
| log_line("retry", f"{label}: switching Civitai key {index}/{total_keys} after status={r.status_code}" + (f" retry_after={retry_after}" if retry_after else "")) | |
| update_civitai_key_status(raw=api_key, active_index=min(index + 1, total_keys), last_switch_reason=reason, last_status=str(r.status_code), source=source) | |
| r.close() | |
| return last_response | |
| def upload_safetensors_to_repo(filename, repo_id, repo_type, is_private, repo_ready=False, api: HfApi | None = None, hf_token=None, progress=gr.Progress(track_tqdm=False), hf_retry_policy="Auto"): | |
| output_filename = Path(filename).name | |
| if hf_token is None: hf_token = get_token() | |
| if api is None: api = HfApi(token=hf_token) | |
| policy_config = get_hf_upload_retry_policy_config(hf_retry_policy) | |
| try: | |
| if not repo_ready and not is_repo_exists(repo_id, repo_type): ensure_repo(api, repo_id=repo_id, repo_type=repo_type, is_private=is_private, hf_token=hf_token) | |
| progress(0, desc=f"Start uploading... {filename} to {repo_id}") | |
| with suppress_hf_hub_progress_bars(): | |
| hf_upload_retry_call(lambda: api.upload_file(path_or_fileobj=filename, path_in_repo=output_filename, repo_type=repo_type, revision="main", token=hf_token, repo_id=repo_id), policy=hf_retry_policy, action=f'upload_file {repo_id}:{output_filename}') | |
| post_sleep = float(policy_config.get("post_upload_sleep", 0.0) or 0.0) | |
| if post_sleep > 0: | |
| time.sleep(post_sleep) | |
| progress(1, desc="Uploaded.") | |
| url = hf_hub_url(repo_id=repo_id, repo_type=repo_type, filename=output_filename) | |
| except Exception as e: | |
| verified_after_error = False | |
| try: | |
| verified_after_error = verify_repo_upload(repo_id, repo_type, output_filename, api=api, hf_token=hf_token) | |
| except Exception: | |
| verified_after_error = False | |
| if verified_after_error: | |
| log_line("retry", f"upload error but remote file exists: {repo_id}:{output_filename}") | |
| progress(1, desc="Uploaded.") | |
| return hf_hub_url(repo_id=repo_id, repo_type=repo_type, filename=output_filename) | |
| print(f"Error: Failed to upload to {repo_id}. {e}") | |
| gr.Warning(f"Error: Failed to upload to {repo_id}. {e}") | |
| return None | |
| finally: | |
| if Path(filename).exists(): Path(filename).unlink() | |
| return url | |
| def upload_safetensors_to_bucket(filename, bucket_id, bucket_ready=False, progress=gr.Progress(track_tqdm=False)): | |
| output_filename = Path(filename).name | |
| hf_token = get_token() | |
| try: | |
| if not is_bucket_api_available(): | |
| raise RuntimeError("Bucket API is unavailable in current huggingface_hub build.") | |
| if not bucket_ready: ensure_bucket(bucket_id=bucket_id, hf_token=hf_token, private=True) | |
| progress(0, desc=f"Start uploading... {filename} to {bucket_id}") | |
| handle, status = upload_file_to_bucket(filename, bucket_id, hf_token, remote_path=output_filename, private=True) | |
| progress(1, desc="Uploaded." if status == "uploaded" else "Skipped.") | |
| return handle | |
| except Exception as e: | |
| print(f"Error: Failed to upload to bucket {bucket_id}. {e}") | |
| gr.Warning(f"Error: Failed to upload to bucket {bucket_id}. {e}") | |
| return None | |
| finally: | |
| if Path(filename).exists(): Path(filename).unlink() | |
| def upload_info_to_bucket(dl_url, filename, bucket_id, civitai_key="", temp_dir="", bucket_ready=False, progress=gr.Progress(track_tqdm=False)): | |
| hf_token = get_token() | |
| uploaded = [] | |
| try: | |
| if not is_bucket_api_available(): | |
| raise RuntimeError("Bucket API is unavailable in current huggingface_hub build.") | |
| if not bucket_ready: ensure_bucket(bucket_id=bucket_id, hf_token=hf_token, private=True) | |
| progress(0, desc=f"Downloading info... {filename}") | |
| json_path, html_path, image_path = save_civitai_info(dl_url, filename, civitai_key, temp_dir=temp_dir) | |
| progress(0, desc=f"Start uploading info... {filename} to {bucket_id}") | |
| for path in [json_path, html_path, image_path]: | |
| if not path or not Path(path).exists(): continue | |
| try: | |
| remote_name = Path(path).name | |
| upload_file_to_bucket(path, bucket_id, hf_token, remote_path=remote_name, private=True) | |
| uploaded.append(remote_name) | |
| finally: | |
| if Path(path).exists(): Path(path).unlink() | |
| progress(1, desc="Info uploaded.") | |
| return uploaded | |
| except Exception as e: | |
| print(f"Error: Failed to upload info to bucket {bucket_id}. {e}") | |
| gr.Warning(f"Error: Failed to upload info to bucket {bucket_id}. {e}") | |
| return uploaded | |
| def get_repo_hashes(repo_id: str, repo_type: str="model", api: HfApi | None = None, hf_token=None, repo_exists: bool | None = None): | |
| if hf_token is None: hf_token = get_token() | |
| if api is None: api = HfApi(token=hf_token) | |
| hashes = [] | |
| try: | |
| if repo_exists is None: | |
| repo_exists = bool(retry_call(lambda: api.repo_exists(repo_id=repo_id, repo_type=repo_type, token=hf_token), action=f'repo_exists {repo_id}')) | |
| if not repo_exists: return hashes | |
| tree = retry_call(lambda: api.list_repo_tree(repo_id=repo_id, repo_type=repo_type, token=hf_token), action=f'list_repo_tree {repo_id}') | |
| for f in tree: | |
| if not isinstance(f, RepoFile) or f.lfs is None or f.lfs.get("sha256", None) is None: continue | |
| hashes.append(f.lfs["sha256"]) | |
| except Exception as e: | |
| print(e) | |
| finally: | |
| return hashes | |
| def get_civitai_url_parts(url: str): | |
| try: | |
| return urllib.parse.urlsplit(str(url or "").strip()) | |
| except Exception: | |
| return urllib.parse.urlsplit("") | |
| def is_civitai_host(netloc: str): | |
| return canonicalize_civitai_host(netloc) in {"civitai.com", "civitai.red"} | |
| def is_civitai_download_api_path(path: str): | |
| return re.match(r'^/api/download/models/\d+$', str(path or "").strip()) is not None | |
| def extract_civitai_model_version_id(url: str): | |
| try: | |
| parts = get_civitai_url_parts(url) | |
| 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 "" | |
| def to_civitai_default_download_url(version_id: str, query: str = ""): | |
| if not str(version_id or "").isdigit(): | |
| return "" | |
| base = f"{get_civitai_display_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() | |
| return urllib.parse.urlunsplit(("https", "civitai.com", 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\.red|civitai\.green)/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(CIVITAI_DEFAULT_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() | |
| parts = get_civitai_url_parts(raw) | |
| if not is_civitai_host(parts.netloc): | |
| return raw | |
| if is_civitai_download_api_path(parts.path): | |
| return normalize_civitai_download_api_url(raw) | |
| if not re.match(r'^/models/\d+(?:/[^/?#]+)?/?$', parts.path or ""): | |
| return raw | |
| version_id = extract_civitai_model_version_id(raw) | |
| if version_id: | |
| return to_civitai_default_download_url(version_id) | |
| headers = get_civitai_headers(api_key if canonicalize_civitai_host(parts.netloc) in {"civitai.com", "civitai.red"} else "") | |
| headers["Referer"] = f"{parts.scheme or 'https'}://{parts.netloc}/" | |
| session = create_retry_session(total=4, backoff_factor=0.8) | |
| try: | |
| r = session.get(raw, headers=headers, timeout=(7.0, 25.0)) | |
| if not r.ok: | |
| return raw | |
| extracted = extract_first_civitai_download_url_from_html(r.text) | |
| return extracted if extracted else raw | |
| except Exception as e: | |
| print(f"Failed to resolve Civitai model page to download URL. {sanitize_url_for_log(raw)} {type(e).__name__}: {e}") | |
| return raw | |
| def normalize_civitai_input_url(url: str, api_key: str = ""): | |
| raw = str(url or "").strip() | |
| if not raw: | |
| return raw | |
| parts = get_civitai_url_parts(raw) | |
| if not is_civitai_host(parts.netloc): | |
| 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 get_civitai_sha256(dl_url: str, api_key=""): | |
| dl_url = normalize_civitai_input_url(dl_url, api_key=api_key) | |
| def is_invalid_file(qs: dict, json: dict, k: str): | |
| return k in qs.keys() and qs[k][0] != json.get(k, None) and json.get(k, None) is not None | |
| if "https://civitai.com/api/download/models/" not in dl_url: return None | |
| base_path = '/model-versions' | |
| params = {} | |
| session = create_retry_session(total=6, backoff_factor=1.0) | |
| m = re.match(r'https://civitai.com/api/download/models/(\d+)\??(.+)?', dl_url) | |
| if m is None: return None | |
| url = f"{base_path}/{m.group(1)}" | |
| qs = urllib.parse.parse_qs(m.group(2)) | |
| if "type" not in qs.keys(): qs["type"] = ["Model"] | |
| try: | |
| r = request_civitai_api(session, url, api_key=api_key, params=params, timeout=(5.0, 15), label='Civitai sha256') | |
| if not r.ok: return None | |
| json = dict(get_civitai_response_json(r, default={}) or {}) | |
| if "files" not in json.keys() or not isinstance(json["files"], list): return None | |
| hash = None | |
| for d in json["files"]: | |
| if is_invalid_file(qs, d, "type") or is_invalid_file(qs, d, "format") or is_invalid_file(qs, d, "size") or is_invalid_file(qs, d, "fp"): continue | |
| hashes = d.get("hashes") if isinstance(d.get("hashes"), dict) else {} | |
| hash_value = str(hashes.get("SHA256") or "").strip() | |
| if not hash_value: | |
| continue | |
| hash = hash_value.lower() | |
| break | |
| return hash | |
| except Exception as e: | |
| print(e) | |
| return None | |
| def is_same_file(filename: str, cmp_sha256: str, cmp_size: int): | |
| if cmp_sha256: | |
| sha256_hash = hashlib.sha256() | |
| with open(filename, "rb") as f: | |
| for byte_block in iter(lambda: f.read(4096), b""): | |
| sha256_hash.update(byte_block) | |
| sha256 = sha256_hash.hexdigest() | |
| else: sha256 = "" | |
| size = os.path.getsize(filename) | |
| if size == cmp_size and sha256 == cmp_sha256: return True | |
| else: return False | |
| def get_safe_filename(filename, repo_id, repo_type, api: HfApi | None = None, hf_token=None): | |
| if hf_token is None: hf_token = get_token() | |
| if api is None: api = HfApi(token=hf_token) | |
| new_filename = filename | |
| try: | |
| i = 1 | |
| while retry_call(lambda: api.file_exists(repo_id=repo_id, filename=Path(new_filename).name, repo_type=repo_type, token=hf_token), action=f'file_exists {repo_id}:{Path(new_filename).name}'): | |
| infos = retry_call(lambda: api.get_paths_info(repo_id=repo_id, paths=[Path(new_filename).name], repo_type=repo_type, token=hf_token), action=f'get_paths_info {repo_id}:{Path(new_filename).name}') | |
| if infos and len(infos) == 1: | |
| repo_fs = infos[0].size | |
| repo_sha256 = infos[0].lfs.sha256 if infos[0].lfs is not None else "" | |
| if is_same_file(filename, repo_sha256, repo_fs): break | |
| new_filename = str(Path(Path(filename).parent, f"{Path(filename).stem}_{i}{Path(filename).suffix}")) | |
| i += 1 | |
| if filename != new_filename: | |
| print(f"{Path(filename).name} is already exists but file content is different. renaming to {Path(new_filename).name}.") | |
| Path(filename).rename(new_filename) | |
| except Exception as e: | |
| print(f"Error occurred when renaming {filename}. {e}") | |
| finally: | |
| return new_filename | |
| def download_file(dl_url, civitai_key, temp_dir="", progress=gr.Progress(track_tqdm=False)): | |
| download_dir = temp_dir if temp_dir else TEMP_DIR | |
| resolved_url = normalize_civitai_input_url(dl_url, api_key=civitai_key) | |
| progress(0, desc=f"Start downloading... {dl_url}") | |
| output_filename = get_download_file(download_dir, resolved_url, civitai_key) | |
| return output_filename | |
| def save_civitai_info(dl_url, filename, civitai_key="", temp_dir="", progress=gr.Progress(track_tqdm=False)): | |
| target_dir = temp_dir if temp_dir else TEMP_DIR | |
| json_str, html_str, image_path = get_civitai_json(dl_url, True, filename, civitai_key, temp_dir=target_dir) | |
| if not json_str: return "", "", "" | |
| json_path = str(Path(target_dir, Path(filename).stem + ".json")) | |
| html_path = str(Path(target_dir, Path(filename).stem + ".html")) | |
| try: | |
| with open(json_path, 'w') as f: | |
| json.dump(json_str, f, indent=2) | |
| with open(html_path, mode='w', encoding="utf-8") as f: | |
| f.write(html_str) | |
| return json_path, html_path, image_path | |
| except Exception as e: | |
| print(f"Error: Failed to save info file {json_path}, {html_path} {e}") | |
| return "", "", "" | |
| def upload_info_to_repo(dl_url, filename, repo_id, repo_type, is_private, civitai_key="", temp_dir="", repo_ready=False, api: HfApi | None = None, hf_token=None, progress=gr.Progress(track_tqdm=False)): | |
| uploaded = [] | |
| def upload_file(api, filename, repo_id, repo_type, hf_token): | |
| if not Path(filename).exists(): return | |
| remote_name = Path(filename).name | |
| with suppress_hf_hub_progress_bars(): | |
| retry_call(lambda: api.upload_file(path_or_fileobj=filename, path_in_repo=remote_name, repo_type=repo_type, revision="main", token=hf_token, repo_id=repo_id), action=f'upload_file {repo_id}:{remote_name}') | |
| uploaded.append(remote_name) | |
| Path(filename).unlink() | |
| if hf_token is None: hf_token = get_token() | |
| if api is None: api = HfApi(token=hf_token) | |
| try: | |
| if not repo_ready and not is_repo_exists(repo_id, repo_type): ensure_repo(api, repo_id=repo_id, repo_type=repo_type, is_private=is_private, hf_token=hf_token) | |
| progress(0, desc=f"Downloading info... {filename}") | |
| json_path, html_path, image_path = save_civitai_info(dl_url, filename, civitai_key, temp_dir=temp_dir) | |
| progress(0, desc=f"Start uploading info... {filename} to {repo_id}") | |
| if not json_path: return uploaded | |
| upload_file(api, json_path, repo_id, repo_type, hf_token) | |
| if html_path: upload_file(api, html_path, repo_id, repo_type, hf_token) | |
| if image_path: upload_file(api, image_path, repo_id, repo_type, hf_token) | |
| progress(1, desc="Info uploaded.") | |
| return uploaded | |
| except Exception as e: | |
| print(f"Error: Failed to upload info to {repo_id}. {e}") | |
| gr.Warning(f"Error: Failed to upload info to {repo_id}. {e}") | |
| return uploaded | |
| def pick_smoke_test_civitai_item(api_key: str = "", progress=gr.Progress(track_tqdm=False)): | |
| search_plans = [("Month", SMOKE_TEST_LIMIT), ("AllTime", SMOKE_TEST_LIMIT)] | |
| for period, limit in search_plans: | |
| progress(0, desc=f"Smoke test: searching small LoRA ({period})...") | |
| items = search_on_civitai("", ["LORA"], [], limit, "Newest", period, "", "", 1, ["Model"], api_key, progress=progress) | |
| if not items: | |
| continue | |
| filtered = [] | |
| for item in items: | |
| dl_url = str(item.get("dl_url", "")).strip() | |
| size_kb = item.get("size_kb") | |
| if not dl_url.startswith("https://civitai.com/api/download/models/"): | |
| continue | |
| try: | |
| size_kb = float(size_kb) | |
| except (TypeError, ValueError): | |
| continue | |
| if size_kb <= 0 or size_kb > SMOKE_TEST_MAX_SIZE_KB: | |
| continue | |
| filtered.append(item | {"size_kb": size_kb}) | |
| if not filtered: | |
| print(f"SMOKE TEST: no LoRA candidates <= {round(SMOKE_TEST_MAX_SIZE_KB / 1000.0, 2)}MB in period={period}.") | |
| continue | |
| filtered = sorted(filtered, key=lambda x: x.get("size_kb", float("inf"))) | |
| pool = filtered[:min(len(filtered), SMOKE_TEST_CANDIDATE_POOL)] | |
| selected = random.choice(pool) | |
| print(f"SMOKE TEST: selected {selected.get('name', '')} / {selected.get('model_name', '')} / {round(selected.get('size_kb', 0.0) / 1000.0, 2)}MB") | |
| return {"selected": selected, "period": period, "candidate_count": len(filtered), "pool_size": len(pool)} | |
| raise gr.Error(f"Smoke test candidate not found within {round(SMOKE_TEST_MAX_SIZE_KB / 1000.0, 2)}MB.") | |
| def smoke_test_civitai(civitai_key, hf_token, urls, repo_type="model", is_private=True, is_info=False, is_rename=True, session_state=None, hf_retry_policy="Auto", progress=gr.Progress(track_tqdm=False)): | |
| session_state = prepare_new_run_state(session_state) | |
| reset_civitai_key_status(civitai_key, source="smoke") | |
| repo_id = str(os.environ.get("HF_REPO", "") or "").strip() | |
| if not repo_id: | |
| raise gr.Error("HF_REPO env var is required for Smoke Test.") | |
| resolved_hf_token = hf_token if hf_token else os.environ.get("HF_TOKEN", "") | |
| if not resolved_hf_token: | |
| raise gr.Error("HF write token is required for Smoke Test.") | |
| api_key = civitai_key if civitai_key else os.environ.get("CIVITAI_API_KEY", "") | |
| parsed_keys = parse_civitai_api_keys(api_key) | |
| if len(parsed_keys) == 0: | |
| raise gr.Error("Civitai API key is required for Smoke Test.") | |
| urls = list(urls) if urls else [] | |
| smoke_lines = [smoke_stage_line("Preflight", "ok", f"repo={repo_id} type={repo_type} info={'on' if is_info else 'off'} rename={'on' if is_rename else 'off'} keys={len(parsed_keys)}")] | |
| selected_url = "" | |
| try: | |
| selected_info = pick_smoke_test_civitai_item(api_key, progress=progress) | |
| selected = selected_info.get("selected", {}) | |
| selected_url = selected.get("dl_url", "") | |
| if not selected_url: | |
| raise RuntimeError("Smoke test candidate is missing download URL.") | |
| smoke_lines.append(smoke_stage_line("Search", "ok", f"period={selected_info.get('period', '')} candidates={selected_info.get('candidate_count', 0)} pool={selected_info.get('pool_size', 0)} selected={selected.get('name', 'LoRA')} {round(float(selected.get('size_kb', 0.0)) / 1000.0, 2)}MB")) | |
| resolved_url = resolve_civitai_download_url(selected_url, api_key, max_tries=2) | |
| resolved_host = urllib.parse.urlparse(resolved_url).netloc | |
| smoke_lines.append(smoke_stage_line("Resolve", "ok", resolved_host)) | |
| session_state_update(session_state, smoke_test_selected_url=selected_url, smoke_test_repo_id=repo_id, smoke_test_resolved_host=resolved_host) | |
| gr.Info(f"Smoke Test target: {selected.get('name', 'LoRA')} / {round(float(selected.get('size_kb', 0.0)) / 1000.0, 2)}MB") | |
| print(f"SMOKE TEST: repo={repo_id} type={repo_type} url={selected_url}") | |
| run_context = {"mode": "smoke", "smoke_lines": smoke_lines, "selected": selected, "resolved_host": resolved_host} | |
| yield from download_civitai(selected_url, api_key, resolved_hf_token, urls, repo_id, repo_type, is_private, is_info, is_rename, session_state=session_state, hf_retry_policy=hf_retry_policy, run_context=run_context, progress=progress) | |
| return | |
| except Exception as e: | |
| detail = f"{type(e).__name__}: {e}" | |
| smoke_lines.append(smoke_stage_line("Resolve", "fail", detail)) | |
| failed_urls = [selected_url] if selected_url else [] | |
| set_last_failure_summary(session_state, detail, url=selected_url) | |
| log_line("fail", f"smoke test failed: {summarize_failure_text(detail, url=selected_url)}") | |
| session_state_update(session_state, current_remaining_urls=[], current_failed_urls=failed_urls.copy(), current_uploaded_urls=urls.copy(), current_smoke_lines=smoke_lines.copy(), last_remaining_urls=[], last_failed_urls=failed_urls.copy()) | |
| md = build_run_markdown("", [], smoke_lines) | |
| yield build_run_outputs(urls, md, [], failed_urls, civitai_key, session_state, remain_visible=False, failed_visible=bool(failed_urls)) | |
| def download_civitai(dl_url, civitai_key, hf_token, urls, | |
| newrepo_id, repo_type="model", is_private=True, is_info=False, is_rename=True, session_state=None, hf_retry_policy="Auto", run_context=None, progress=gr.Progress(track_tqdm=False)): | |
| session_state = prepare_new_run_state(session_state) | |
| run_context = run_context if isinstance(run_context, dict) else {} | |
| run_mode = run_context.get("mode", "manual") | |
| smoke_lines = list(run_context.get("smoke_lines", [])) | |
| reset_civitai_key_status(civitai_key, source=run_mode) | |
| resolved_hf_token = hf_token if hf_token else os.getenv("HF_TOKEN", False) | |
| set_token(resolved_hf_token, session_state) | |
| hf_token_value = get_token(session_state) | |
| if not civitai_key: | |
| civitai_key = os.environ.get("CIVITAI_API_KEY") | |
| if not newrepo_id: | |
| newrepo_id = os.environ.get("HF_REPO") | |
| civitai_keys = parse_civitai_api_keys(civitai_key) | |
| if not hf_token_value or len(civitai_keys) == 0: | |
| raise gr.Error("HF write token and Civitai API key is required.") | |
| if repo_type == "bucket" and not is_bucket_api_available(): | |
| raise gr.Error("Bucket API is unavailable in current huggingface_hub build.") | |
| urls = list(urls) if urls else [] | |
| dl_urls = normalize_url_entries(dl_url) | |
| remain_urls = dl_urls.copy() | |
| failed_urls = [] | |
| failure_reasons = {} | |
| run_stats = new_run_stats(len(dl_urls)) | |
| hf_retry_config = get_hf_upload_retry_policy_config(hf_retry_policy) | |
| result_lines = [] | |
| error_message = "" | |
| cancelled = False | |
| run_temp_dir = create_run_temp_dir() | |
| run_id = new_run_id() | |
| register_run(run_id) | |
| repo_ready = False | |
| bucket_ready = False | |
| repo_header = "" | |
| hashes = set() | |
| api = None | |
| session_state_update(session_state, current_run_mode=run_mode, current_run_id=run_id, current_repo_id=newrepo_id, current_repo_type=repo_type, current_run_temp_dir=run_temp_dir, current_remaining_urls=remain_urls.copy(), current_failed_urls=failed_urls.copy(), current_uploaded_urls=urls.copy(), current_smoke_lines=smoke_lines.copy(), current_hf_retry_policy=hf_retry_config.get('key', 'auto'), report_events=[], active_run_id=run_id, cancel_requested=False, last_error="", last_failure_summary="", run_started_at=time.time(), run_elapsed_sec=0.0) | |
| append_report_event(session_state, "run_started", run_id=run_id, mode=run_mode, target=newrepo_id, repo_type=repo_type, input_urls=len(dl_urls), repo_inputs=len(normalize_repo_entries(dl_url)), info=is_info, rename=is_rename, hf_retry=hf_retry_config.get('key', 'auto')) | |
| log_line("info", f"starting {run_mode} run target={newrepo_id} type={repo_type} urls={len(dl_urls)} repos={len(normalize_repo_entries(dl_url))} info={'on' if is_info else 'off'} rename={'on' if is_rename else 'off'} hf_retry={hf_retry_config.get('key', 'auto')}") | |
| update_run_stage(session_state, "Preparing", f"target {newrepo_id}") | |
| try: | |
| set_stage_progress(progress, 0, max(len(dl_urls), 1), f"Preparing target {newrepo_id}...") | |
| check_run_cancel(run_id, session_state=session_state) | |
| if repo_type == "bucket": | |
| print("Bucket mode: missing buckets are created as private. The privacy checkbox is ignored for new buckets.") | |
| ensure_bucket(bucket_id=newrepo_id, hf_token=hf_token_value, private=True) | |
| bucket_ready = True | |
| repo_header = f"### Your bucket: [{newrepo_id}]({get_bucket_url(newrepo_id)})\n" | |
| hashes = set() | |
| else: | |
| api = HfApi(token=hf_token_value) | |
| repo_ready = is_repo_exists(newrepo_id, repo_type) | |
| if not repo_ready: | |
| ensure_repo(api, repo_id=newrepo_id, repo_type=repo_type, is_private=is_private, hf_token=hf_token_value) | |
| repo_ready = True | |
| hashes = set() | |
| else: | |
| cached_hashes = get_session_repo_hash_cache(session_state, newrepo_id, repo_type) | |
| if cached_hashes is not None: | |
| hashes = set(cached_hashes) | |
| print(f"Using cached repo hash index for {newrepo_id} ({len(hashes)} entries).") | |
| else: | |
| hashes = set(get_repo_hashes(newrepo_id, repo_type, api=api, hf_token=hf_token_value, repo_exists=True)) | |
| store_session_repo_hash_cache(session_state, newrepo_id, repo_type, hashes) | |
| repo_base_url = "https://huggingface.co/datasets/" if repo_type == "dataset" else "https://huggingface.co/" | |
| repo_header = f"### Your repo: [{newrepo_id}]({repo_base_url}{newrepo_id})\n" | |
| md = build_run_markdown(repo_header, result_lines, smoke_lines) | |
| yield build_run_outputs(urls, md, remain_urls, failed_urls, civitai_key, session_state) | |
| total_urls = len(dl_urls) if len(dl_urls) > 0 else 1 | |
| for index, u in enumerate(dl_urls, start=1): | |
| current_file = "" | |
| uploaded_name = "" | |
| try: | |
| check_run_cancel(run_id, session_state=session_state) | |
| update_run_stage(session_state, "Processing", "checking duplicate state", index=index, total=total_urls, current_url=u) | |
| set_stage_progress(progress, index - 1, total_urls, f"Processing {index}/{total_urls}") | |
| civitai_sha256 = get_civitai_sha256(u, civitai_key) if repo_type != "bucket" else None | |
| if repo_type != "bucket" and civitai_sha256 and civitai_sha256 in hashes: | |
| increment_run_stat(run_stats, "skipped_duplicate") | |
| append_report_event(session_state, "skipped_duplicate", url=u, index=index, total=total_urls, sha256=bool(civitai_sha256)) | |
| log_line("retry", f"skip duplicate in target repo: {sanitize_url_for_log(u)}") | |
| if run_mode == "smoke": | |
| smoke_lines.append(smoke_stage_line("Duplicate/skip", "ok", "same SHA256 already exists in target repo")) | |
| if u in remain_urls: | |
| remain_urls.remove(u) | |
| result_lines.append(f"- Skipped [{str(u)}]({str(u)})") | |
| md = build_run_markdown(repo_header, result_lines, smoke_lines) | |
| session_state_update(session_state, current_remaining_urls=remain_urls.copy(), current_failed_urls=failed_urls.copy(), current_uploaded_urls=urls.copy(), current_smoke_lines=smoke_lines.copy()) | |
| yield build_run_outputs(urls, md, remain_urls, failed_urls, civitai_key, session_state) | |
| continue | |
| check_run_cancel(run_id, session_state=session_state) | |
| update_run_stage(session_state, "Downloading", "fetching from Civitai", index=index, total=total_urls, current_url=u) | |
| set_stage_progress(progress, index - 1, total_urls, f"Downloading {index}/{total_urls}") | |
| append_report_event(session_state, "download_started", url=u, index=index, total=total_urls) | |
| current_file = download_file(u, civitai_key, temp_dir=run_temp_dir, progress=progress) | |
| file_ok, file_detail = summarize_downloaded_file(current_file) | |
| append_report_event(session_state, "download_completed" if file_ok else "download_failed", url=u, index=index, total=total_urls, detail=file_detail, file=Path(str(current_file)).name if current_file else "") | |
| if run_mode == "smoke": | |
| smoke_lines.append(smoke_stage_line("Download verify", "ok" if file_ok else "fail", file_detail)) | |
| if not file_ok: | |
| increment_run_stat(run_stats, "failed_download") | |
| append_report_event(session_state, "download_failed", url=u, index=index, total=total_urls, detail=file_detail) | |
| failure_reasons[u] = "download" | |
| if u not in failed_urls: | |
| failed_urls.append(u) | |
| set_last_failure_summary(session_state, "download failed or file missing", url=u) | |
| log_line("fail", f"download failed or file missing: {sanitize_url_for_log(u)}") | |
| result_lines.append(f"- Failed [{str(u)}]({str(u)}) (download)") | |
| md = build_run_markdown(repo_header, result_lines, smoke_lines) | |
| session_state_update(session_state, current_remaining_urls=remain_urls.copy(), current_failed_urls=failed_urls.copy(), current_uploaded_urls=urls.copy(), current_smoke_lines=smoke_lines.copy()) | |
| yield build_run_outputs(urls, md, remain_urls, failed_urls, civitai_key, session_state) | |
| continue | |
| increment_run_stat(run_stats, "downloaded") | |
| check_run_cancel(run_id, session_state=session_state) | |
| if is_rename: | |
| update_run_stage(session_state, "Renaming", "checking target name", index=index, total=total_urls, current_url=u) | |
| if repo_type == "bucket": | |
| current_file = get_safe_bucket_filename(current_file, newrepo_id, hf_token_value) | |
| else: | |
| current_file = get_safe_filename(current_file, newrepo_id, repo_type, api=api, hf_token=hf_token_value) | |
| uploaded_name = Path(current_file).name | |
| update_run_stage(session_state, "Uploading", uploaded_name, index=index, total=total_urls, current_url=u) | |
| set_stage_progress(progress, index - 1, total_urls, f"Uploading {index}/{total_urls}") | |
| append_report_event(session_state, "upload_started", url=u, filename=uploaded_name, index=index, total=total_urls, repo_id=newrepo_id, repo_type=repo_type, hf_retry=hf_retry_config.get('key', 'auto')) | |
| url = upload_safetensors_to_bucket(current_file, newrepo_id, bucket_ready=bucket_ready, progress=progress) if repo_type == "bucket" else upload_safetensors_to_repo(current_file, newrepo_id, repo_type, is_private, repo_ready=repo_ready, api=api, hf_token=hf_token_value, progress=progress, hf_retry_policy=hf_retry_policy) | |
| if url: | |
| upload_verified = True | |
| upload_detail = f"{uploaded_name} -> {newrepo_id}" | |
| if repo_type != "bucket": | |
| upload_verified = verify_repo_upload(newrepo_id, repo_type, uploaded_name, api=api, hf_token=hf_token_value) | |
| upload_detail = uploaded_name | |
| if run_mode == "smoke": | |
| smoke_lines.append(smoke_stage_line("Upload verify", "ok" if upload_verified else "fail", upload_detail)) | |
| if repo_type != "bucket" and not upload_verified: | |
| increment_run_stat(run_stats, "failed_upload") | |
| append_report_event(session_state, "upload_verify_failed", url=u, filename=uploaded_name, index=index, total=total_urls, repo_id=newrepo_id, repo_type=repo_type) | |
| failure_reasons[u] = "upload verify" | |
| if u not in failed_urls: | |
| failed_urls.append(u) | |
| result_lines.append(f"- Failed [{str(u)}]({str(u)}) (upload verify)") | |
| md = build_run_markdown(repo_header, result_lines, smoke_lines) | |
| session_state_update(session_state, current_remaining_urls=remain_urls.copy(), current_failed_urls=failed_urls.copy(), current_uploaded_urls=urls.copy(), current_smoke_lines=smoke_lines.copy()) | |
| yield build_run_outputs(urls, md, remain_urls, failed_urls, civitai_key, session_state) | |
| continue | |
| increment_run_stat(run_stats, "uploaded") | |
| append_report_event(session_state, "upload_completed", url=u, filename=uploaded_name, index=index, total=total_urls, repo_id=newrepo_id, repo_type=repo_type, upload_url=url) | |
| if civitai_sha256: | |
| hashes.add(civitai_sha256) | |
| if repo_type != "bucket": | |
| store_session_repo_hash_cache(session_state, newrepo_id, repo_type, hashes) | |
| if run_mode == "smoke": | |
| if civitai_sha256: | |
| smoke_lines.append(smoke_stage_line("Duplicate/skip", "ok", "same SHA256 would be skipped on immediate retry")) | |
| else: | |
| smoke_lines.append(smoke_stage_line("Duplicate/skip", "fail", "missing SHA256 for duplicate check")) | |
| if is_info: | |
| check_run_cancel(run_id, session_state=session_state) | |
| update_run_stage(session_state, "Uploading info", uploaded_name, index=index, total=total_urls, current_url=u) | |
| info_uploaded = upload_info_to_bucket(u, current_file, newrepo_id, civitai_key, temp_dir=run_temp_dir, bucket_ready=bucket_ready, progress=progress) if repo_type == "bucket" else upload_info_to_repo(u, current_file, newrepo_id, repo_type, is_private, civitai_key, temp_dir=run_temp_dir, repo_ready=repo_ready, api=api, hf_token=hf_token_value, progress=progress) | |
| info_count = len(info_uploaded) if isinstance(info_uploaded, list) else 0 | |
| if info_count == 0: | |
| result_lines.append(f"- Uploaded [{str(u)}]({str(u)}) (info: 0 files)") | |
| else: | |
| result_lines.append(f"- Uploaded [{str(u)}]({str(u)}) (info: {info_count} files)") | |
| if run_mode == "smoke": | |
| if repo_type == "bucket": | |
| info_ok = bool(info_uploaded) | |
| else: | |
| info_ok = bool(info_uploaded) and all(verify_repo_upload(newrepo_id, repo_type, name, api=api, hf_token=hf_token_value) for name in info_uploaded) | |
| detail = f"count={info_count}" if info_uploaded else "no info files uploaded" | |
| smoke_lines.append(smoke_stage_line("Info upload", "ok" if info_ok else "fail", detail)) | |
| else: | |
| result_lines.append(f"- Uploaded [{str(u)}]({str(u)})") | |
| if run_mode == "smoke": | |
| smoke_lines.append(smoke_stage_line("Info upload", "ok", stage_detail("info", False))) | |
| urls.append(url) | |
| if u in remain_urls: | |
| remain_urls.remove(u) | |
| if u in failed_urls: | |
| failed_urls.remove(u) | |
| else: | |
| increment_run_stat(run_stats, "failed_upload") | |
| append_report_event(session_state, "upload_failed", url=u, filename=uploaded_name, index=index, total=total_urls, repo_id=newrepo_id, repo_type=repo_type) | |
| failure_reasons[u] = "upload" | |
| if u not in failed_urls: | |
| failed_urls.append(u) | |
| if run_mode == "smoke": | |
| smoke_lines.append(smoke_stage_line("Upload verify", "fail", f"upload API returned empty for {uploaded_name}")) | |
| result_lines.append(f"- Failed [{str(u)}]({str(u)}) (upload)") | |
| md = build_run_markdown(repo_header, result_lines, smoke_lines) | |
| session_state_update(session_state, current_remaining_urls=remain_urls.copy(), current_failed_urls=failed_urls.copy(), current_uploaded_urls=urls.copy(), current_smoke_lines=smoke_lines.copy()) | |
| yield build_run_outputs(urls, md, remain_urls, failed_urls, civitai_key, session_state) | |
| except RunCancelledError: | |
| cancelled = True | |
| log_line("cancel", f"cancelled while processing {sanitize_url_for_log(u)}") | |
| break | |
| except Exception as e: | |
| log_line("fail", f"error while processing {sanitize_url_for_log(u)}: {type(e).__name__}: {e}") | |
| failure_reasons[u] = f"exception:{type(e).__name__}" | |
| if current_file and Path(str(current_file)).exists(): | |
| increment_run_stat(run_stats, "failed_upload") | |
| append_report_event(session_state, "item_failed", url=u, phase="upload_or_post_download", error_type=type(e).__name__, error=str(e)) | |
| else: | |
| increment_run_stat(run_stats, "failed_download") | |
| append_report_event(session_state, "item_failed", url=u, phase="download_or_pre_download", error_type=type(e).__name__, error=str(e)) | |
| set_last_failure_summary(session_state, f"{type(e).__name__}: {e}", url=u) | |
| if u not in failed_urls: | |
| failed_urls.append(u) | |
| result_lines.append(f"- Failed [{str(u)}]({str(u)}) ({e})") | |
| md = build_run_markdown(repo_header, result_lines, smoke_lines) | |
| session_state_update(session_state, current_remaining_urls=remain_urls.copy(), current_failed_urls=failed_urls.copy(), current_uploaded_urls=urls.copy(), current_smoke_lines=smoke_lines.copy()) | |
| yield build_run_outputs(urls, md, remain_urls, failed_urls, civitai_key, session_state) | |
| if not cancelled: | |
| repo_inputs = normalize_repo_entries(dl_url) | |
| if len(repo_inputs) > 0: | |
| update_run_stage(session_state, "Duplicating repos", f"{len(repo_inputs)} item(s)", index=len(dl_urls), total=max(len(dl_urls) + len(repo_inputs), 1)) | |
| set_stage_progress(progress, len(dl_urls), max(len(dl_urls) + len(repo_inputs), 1), f"Duplicating repos to {newrepo_id}...") | |
| for r in repo_inputs: | |
| check_run_cancel(run_id, session_state=session_state) | |
| try: | |
| if repo_type == "bucket": | |
| log_line("info", f"bucket mode skips direct repo duplicate: {r}") | |
| result_lines.append(f"- Skipped duplicate repo [{str(r)}](https://huggingface.co/{str(r)}) (bucket mode)") | |
| continue | |
| url = duplicate_hf_repo(r, newrepo_id, "model", repo_type, is_private, HF_SUBFOLDER_NAME[1]) | |
| if url: | |
| urls.append(url) | |
| except RunCancelledError: | |
| cancelled = True | |
| break | |
| except Exception as e: | |
| log_line("fail", f"error while duplicating {r}: {type(e).__name__}: {e}") | |
| set_last_failure_summary(session_state, f"duplicate repo failed: {type(e).__name__}: {e}", url=r) | |
| result_lines.append(f"- Failed duplicate repo [{str(r)}](https://huggingface.co/{str(r)}) ({e})") | |
| except RunCancelledError: | |
| cancelled = True | |
| except Exception as e: | |
| error_message = str(e) | |
| failed_urls = list_uniq(failed_urls + remain_urls) | |
| set_last_failure_summary(session_state, f"{type(e).__name__}: {e}") | |
| log_line("fail", f"run failed: {type(e).__name__}: {e}") | |
| gr.Info(f"Error occurred: {e}") | |
| finally: | |
| if repo_type != "bucket" and hashes: | |
| store_session_repo_hash_cache(session_state, newrepo_id, repo_type, hashes) | |
| cleanup_run_temp_dir(run_temp_dir) | |
| unregister_run(run_id) | |
| incomplete = bool(remain_urls) or bool(failed_urls) | |
| final_stage = "Cancelled" if cancelled else ("Failed" if error_message else ("Incomplete" if incomplete else "Done")) | |
| final_detail = f"remaining={len(remain_urls)} failed={len(failed_urls)}" | |
| update_run_stage(session_state, final_stage, final_detail, index=len(dl_urls), total=max(len(dl_urls), 1)) | |
| if run_mode == "smoke": | |
| smoke_lines.append(smoke_stage_line("Retry state", "ok", f"remaining={len(remain_urls)} failed={len(failed_urls)} error={'yes' if error_message else 'no'} cancel={'yes' if cancelled else 'no'}")) | |
| smoke_lines.append(smoke_stage_line("Cleanup", "ok", Path(run_temp_dir).name)) | |
| if cancelled and not str(session_state.get("last_failure_summary") or "").strip(): | |
| set_last_failure_summary(session_state, "cancelled by user") | |
| elif error_message and not str(session_state.get("last_failure_summary") or "").strip(): | |
| set_last_failure_summary(session_state, error_message) | |
| if final_stage == "Incomplete": | |
| log_line("cleanup", f"run incomplete remaining={len(remain_urls)} failed={len(failed_urls)} reasons={json.dumps(failure_reasons, ensure_ascii=False)[:1000]}") | |
| log_line("cleanup", f"finished {run_mode} run stage={final_stage.lower()} remaining={len(remain_urls)} failed={len(failed_urls)}") | |
| log_run_summary(run_mode, final_stage, run_stats, remain_urls, failed_urls) | |
| last_run_summary = dict(run_stats) | |
| last_run_summary.update({"run_id": run_id, "stage": final_stage, "mode": run_mode, "repo_id": newrepo_id, "repo_type": repo_type, "remaining": len(remain_urls), "failed": len(failed_urls), "cancelled": bool(cancelled), "error": bool(error_message), "hf_retry_policy": hf_retry_config.get('key', 'auto')}) | |
| append_report_event(session_state, "run_finished", **last_run_summary) | |
| current_events = list(session_state.get("report_events") or []) | |
| last_run_record = build_report_run_record(run_id, last_run_summary, current_events, remain_urls, failed_urls, urls, smoke_lines=smoke_lines, failure_reasons=failure_reasons) | |
| session_records = append_session_run_record(session_state, last_run_record) | |
| session_summary = summarize_report_runs(session_records) | |
| session_state_update(session_state, current_run_mode="idle", current_run_id=run_id, current_repo_id=newrepo_id, current_repo_type=repo_type, current_run_temp_dir="", current_remaining_urls=remain_urls.copy(), current_failed_urls=failed_urls.copy(), current_uploaded_urls=urls.copy(), current_smoke_lines=smoke_lines.copy(), last_run_id=run_id, last_run_mode=run_mode, last_repo_id=newrepo_id, last_repo_type=repo_type, last_remaining_urls=remain_urls.copy(), last_failed_urls=failed_urls.copy(), last_uploaded_urls=urls.copy(), last_smoke_lines=smoke_lines.copy(), last_run_summary=last_run_summary, last_run_record=last_run_record, session_summary=session_summary, last_failure_reasons=dict(failure_reasons), last_hf_retry_policy=hf_retry_config.get('key', 'auto'), last_error=error_message, active_run_id="", cancel_requested=False) | |
| gc.collect() | |
| summary_lines = build_run_summary_lines(run_stats, remain_urls, failed_urls, final_stage) | |
| md = build_run_markdown(repo_header if repo_header else "", result_lines + summary_lines, smoke_lines) | |
| if cancelled: | |
| md = build_run_markdown(repo_header if repo_header else "", result_lines + ["- Cancelled by user."] + summary_lines, smoke_lines) | |
| elif error_message and not result_lines: | |
| md = build_run_markdown(repo_header if repo_header else "", [f"- Failed ({error_message})"] + summary_lines, smoke_lines) | |
| set_stage_progress(progress, 1, 1, final_stage) | |
| yield build_run_outputs(urls, md, remain_urls, failed_urls, civitai_key, session_state, remain_visible=bool(remain_urls) or bool(error_message) or bool(cancelled), failed_visible=bool(failed_urls)) | |
| def normalize_civitai_basemodel_name(name): | |
| if name is None: return "" | |
| return str(name).strip() | |
| def sort_civitai_basemodels(items: list[str]): | |
| default_index = {name: i for i, name in enumerate(CIVITAI_BASEMODEL_DEFAULT)} | |
| return sorted(items, key=lambda x: (0, default_index[x]) if x in default_index else (1, x.casefold())) | |
| def fetch_civitai_basemodels(api_key: str = "", pages_per_sort: int = CIVITAI_BASEMODEL_REFRESH_PAGES_PER_SORT): | |
| base_path = '/models' | |
| observed = set() | |
| session = create_retry_session(total=6, backoff_factor=1.0) | |
| seeds = [("Newest", "AllTime"), ("Most Downloaded", "AllTime")] | |
| for sort, period in seeds: | |
| next_url = None | |
| for page_index in range(1, pages_per_sort + 1): | |
| params = {'sort': sort, 'period': period, 'limit': 100, 'page': page_index, 'nsfw': 'true'} | |
| try: | |
| if next_url is None: | |
| r = request_civitai_api(session, base_path, api_key=api_key, params=params, timeout=(7.0, 30), label=f'Civitai base model refresh sort={sort} page={page_index}', source='base-models') | |
| else: | |
| r = request_civitai_api_url(session, next_url, api_key=api_key, timeout=(7.0, 30), label=f'Civitai base model refresh sort={sort} page={page_index}', source='base-models') | |
| if not r.ok: | |
| print(f"Failed to refresh Civitai base models. sort={sort} page={page_index} status={r.status_code}") | |
| break | |
| data = get_civitai_response_json(r, default={}) or {} | |
| items = data.get('items', []) | |
| if not isinstance(items, list): break | |
| for item in items: | |
| for model in item.get('modelVersions', []): | |
| name = normalize_civitai_basemodel_name(model.get('baseModel', '')) | |
| if name: observed.add(name) | |
| next_url = data.get('metadata', {}).get('nextPage') | |
| if not next_url: break | |
| time.sleep(0.4) | |
| except Exception as e: | |
| print(f"Failed to refresh Civitai base models. sort={sort} page={page_index} error={e}") | |
| break | |
| return sort_civitai_basemodels(list(observed)) | |
| def get_civitai_basemodels(api_key: str = ""): | |
| observed = fetch_civitai_basemodels(api_key=api_key) | |
| if len(observed) >= CIVITAI_BASEMODEL_MIN_COUNT: | |
| print(f"Loaded {len(observed)} Civitai base models from API at startup.") | |
| print("CIVITAI_BASEMODEL_DEFAULT = " + json.dumps(observed, ensure_ascii=False, separators=(",", ":"))) | |
| return observed | |
| print("Falling back to bundled Civitai base model list.") | |
| return CIVITAI_BASEMODEL_DEFAULT.copy() | |
| CIVITAI_BASEMODEL = get_civitai_basemodels(api_key=os.environ.get("CIVITAI_API_KEY", "")) | |
| #CIVITAI_SORT = ["Highest Rated", "Most Downloaded", "Newest"] | |
| CIVITAI_SORT_EXT = ["Size", "Size (from smallest)"] | |
| CIVITAI_SORT = ["Highest Rated", "Most Downloaded", "Most Liked", "Most Discussed", "Most Collected", "Most Buzz", "Newest"] + CIVITAI_SORT_EXT | |
| CIVITAI_PERIOD = ["AllTime", "Year", "Month", "Week", "Day"] | |
| def build_item_label(item: dict): | |
| base_model_name = "Pony🐴" if item.get('base_model', '') == "Pony" else item.get('base_model', '') | |
| if "size_kb" in item.keys(): | |
| return f"{item.get('name', '')} (for {base_model_name} / By: {item.get('creator', '')}) ({round(float(item.get('size_kb', 0.0)) / 1000.0, 2)}MB)" | |
| return f"{item.get('name', '')} (for {base_model_name} / By: {item.get('creator', '')})" | |
| def shorten_text(text: str, max_len: int = 54): | |
| text = str(text or "").strip() | |
| if len(text) <= max_len: | |
| return text | |
| return text[: max_len - 1] + "…" | |
| def select_all_button_label(state: dict): | |
| items = get_state(state, "civitai_last_items") or [] | |
| selected = get_state(state, "civitai_last_selects") or [] | |
| valid_values = {item.get("dl_url", "") for item in items if item.get("dl_url", "")} | |
| return "Deselect All" if valid_values and valid_values.issubset(set(selected)) else "Select All" | |
| def cleanup_search_preview_dir(state: dict): | |
| state = state if isinstance(state, dict) else {} | |
| preview_dir = str(state.get("civitai_search_preview_dir") or "") | |
| if preview_dir: | |
| cleanup_run_temp_dir(preview_dir) | |
| state["civitai_search_preview_dir"] = "" | |
| set_state(state, "civitai_preview_cache", {}) | |
| def ensure_search_preview_dir(state: dict): | |
| state = state if isinstance(state, dict) else {} | |
| preview_dir = str(state.get("civitai_search_preview_dir") or "") | |
| if preview_dir and Path(preview_dir).exists(): | |
| return preview_dir | |
| preview_dir = create_run_temp_dir() | |
| state["civitai_search_preview_dir"] = preview_dir | |
| return preview_dir | |
| def pil_resample(): | |
| try: | |
| return Image.Resampling.LANCZOS | |
| except Exception: | |
| return Image.LANCZOS | |
| def get_preview_fail_urls(state: dict): | |
| values = get_state(state, "civitai_preview_fail_urls") or [] | |
| return set(str(v) for v in values if v) | |
| def add_preview_fail_url(state: dict, image_url: str): | |
| if not image_url: | |
| return | |
| failed = get_preview_fail_urls(state) | |
| if image_url in failed: | |
| return | |
| failed.add(image_url) | |
| set_state(state, "civitai_preview_fail_urls", sorted(failed)) | |
| def get_preview_cache(state: dict): | |
| cache = get_state(state, "civitai_preview_cache") | |
| return dict(cache) if isinstance(cache, dict) else {} | |
| def set_preview_cache(state: dict, cache: dict): | |
| set_state(state, "civitai_preview_cache", dict(cache or {})) | |
| def get_cached_preview_path(state: dict, cache_key: str): | |
| cache = get_preview_cache(state) | |
| path = str(cache.get(cache_key) or "") | |
| return path if path and Path(path).exists() else "" | |
| def remember_preview_path(state: dict, cache_key: str, preview_path: str): | |
| if not cache_key or not preview_path or preview_path == NULL_IMAGE_PATH: | |
| return | |
| cache = get_preview_cache(state) | |
| cache[cache_key] = preview_path | |
| set_preview_cache(state, cache) | |
| def get_preview_cache_key(item: dict): | |
| return str(item.get('dl_url') or item.get('img_url') or item.get('model_version_id') or item.get('name') or '') | |
| def is_probable_video_url(image_url: str): | |
| try: | |
| path = urllib.parse.urlsplit(str(image_url or "")).path.lower() | |
| except Exception: | |
| path = str(image_url or "").lower() | |
| return any(path.endswith(ext) for ext in PREVIEW_VIDEO_EXTS) | |
| def get_ffmpeg_path(): | |
| global _FFMPEG_PATH, _FFMPEG_MISSING_LOGGED | |
| if _FFMPEG_PATH is None: | |
| _FFMPEG_PATH = shutil.which("ffmpeg") or "" | |
| if not _FFMPEG_PATH and not _FFMPEG_MISSING_LOGGED: | |
| print("ffmpeg not found. Video previews will use fallback image.") | |
| _FFMPEG_MISSING_LOGGED = True | |
| return _FFMPEG_PATH | |
| def save_preview_image(img: Image.Image, output_path: str, size: tuple[int, int], fit_mode: str = "cover"): | |
| path_obj = Path(output_path) | |
| img = ImageOps.exif_transpose(img).convert('RGB') | |
| resample = pil_resample() | |
| if fit_mode == "cover": | |
| img = ImageOps.fit(img, size, method=resample) | |
| else: | |
| img.thumbnail(size, resample) | |
| img.save(path_obj, format='WEBP', quality=80, method=6) | |
| return str(path_obj) | |
| def fetch_video_preview(video_url: str, output_path: str, size: tuple[int, int], state: dict | None = None, fit_mode: str = "cover"): | |
| if not video_url or video_url == NULL_IMAGE_PATH: | |
| return NULL_IMAGE_PATH | |
| path_obj = Path(output_path) | |
| if path_obj.exists(): | |
| return str(path_obj) | |
| if video_url in get_preview_fail_urls(state or {}): | |
| return NULL_IMAGE_PATH | |
| ffmpeg_path = get_ffmpeg_path() | |
| if not ffmpeg_path: | |
| add_preview_fail_url(state or {}, video_url) | |
| return NULL_IMAGE_PATH | |
| temp_frame = path_obj.with_suffix('.png') | |
| cmd = [ffmpeg_path, '-y', '-loglevel', 'error', '-nostdin', '-i', video_url, '-frames:v', '1', str(temp_frame)] | |
| try: | |
| subprocess.run(cmd, check=True, timeout=30) | |
| if not temp_frame.exists(): | |
| add_preview_fail_url(state or {}, video_url) | |
| return NULL_IMAGE_PATH | |
| with Image.open(temp_frame) as img: | |
| return save_preview_image(img, output_path, size, fit_mode=fit_mode) | |
| except Exception as e: | |
| add_preview_fail_url(state or {}, video_url) | |
| print(f"Failed to build video preview from {video_url}. {e}") | |
| return NULL_IMAGE_PATH | |
| finally: | |
| try: | |
| if temp_frame.exists(): | |
| temp_frame.unlink() | |
| except Exception: | |
| pass | |
| def fetch_preview_image(image_url: str, output_path: str, size: tuple[int, int], state: dict | None = None, fit_mode: str = "cover"): | |
| if not image_url or image_url == NULL_IMAGE_PATH: | |
| return NULL_IMAGE_PATH | |
| path_obj = Path(output_path) | |
| if path_obj.exists(): | |
| return str(path_obj) | |
| if image_url in get_preview_fail_urls(state or {}): | |
| return NULL_IMAGE_PATH | |
| if is_probable_video_url(image_url): | |
| return fetch_video_preview(image_url, output_path, size, state=state, fit_mode=fit_mode) | |
| try: | |
| session = create_retry_session(total=4, backoff_factor=0.8) | |
| headers = {'User-Agent': get_user_agent(), 'Referer': 'https://civitai.com/'} | |
| with session.get(image_url, headers=headers, stream=True, timeout=(7.0, 30.0)) as r: | |
| if not r.ok: | |
| add_preview_fail_url(state or {}, image_url) | |
| return NULL_IMAGE_PATH | |
| content_type = str(r.headers.get('content-type') or '').split(';', 1)[0].strip().lower() | |
| if content_type.startswith('video/'): | |
| return fetch_video_preview(image_url, output_path, size, state=state, fit_mode=fit_mode) | |
| with Image.open(BytesIO(r.content)) as img: | |
| return save_preview_image(img, output_path, size, fit_mode=fit_mode) | |
| except Exception as e: | |
| add_preview_fail_url(state or {}, image_url) | |
| print(f"Failed to build preview image from {image_url}. {e}") | |
| return NULL_IMAGE_PATH | |
| def ensure_item_preview(item: dict, state: dict, detail: bool = False): | |
| image_url = str(item.get('img_url') or '').strip() | |
| if not image_url: | |
| return NULL_IMAGE_PATH | |
| preview_dir = ensure_search_preview_dir(state) | |
| mode = 'thumb' | |
| digest = hashlib.sha1(f"{mode}|{image_url}".encode('utf-8', 'ignore')).hexdigest()[:16] | |
| output_path = str(Path(preview_dir, f"{mode}_{digest}.webp")) | |
| return fetch_preview_image(image_url, output_path, SEARCH_THUMB_SIZE, state=state, fit_mode="cover") | |
| def resolve_item_preview(item: dict, state: dict, build_missing: bool = True): | |
| cache_key = get_preview_cache_key(item) | |
| cached_path = get_cached_preview_path(state, cache_key) | |
| if cached_path: | |
| return cached_path | |
| if not build_missing: | |
| return NULL_IMAGE_PATH | |
| preview_path = ensure_item_preview(item, state, detail=False) | |
| if preview_path and preview_path != NULL_IMAGE_PATH: | |
| remember_preview_path(state, cache_key, preview_path) | |
| return preview_path | |
| def is_item_preview_known(state: dict, item: dict): | |
| cache_key = get_preview_cache_key(item) | |
| cached_path = get_cached_preview_path(state, cache_key) | |
| if cached_path: | |
| return True | |
| image_url = str(item.get("img_url") or "").strip() | |
| if not image_url or image_url == NULL_IMAGE_PATH: | |
| return True | |
| return image_url in get_preview_fail_urls(state) | |
| def prune_search_preview_dir(state: dict, keep_paths: list[str]): | |
| state = state if isinstance(state, dict) else {} | |
| preview_dir = str(state.get("civitai_search_preview_dir") or "") | |
| if not preview_dir or not Path(preview_dir).exists(): | |
| return | |
| keep = {str(Path(p)) for p in keep_paths if p and str(p).startswith(preview_dir)} | |
| for path in Path(preview_dir).glob("*.webp"): | |
| if str(path) not in keep: | |
| try: | |
| path.unlink() | |
| except Exception: | |
| pass | |
| def find_item_by_url(state: dict, value: str): | |
| results = get_state(state, "civitai_last_results") or {} | |
| entry = results.get(value, {}) if isinstance(results, dict) else {} | |
| item = entry.get('item') if isinstance(entry, dict) else None | |
| return item if isinstance(item, dict) else None | |
| def is_civitai_gallery_enabled(state: dict): | |
| value = get_state(state, "civitai_gallery_enabled") | |
| return bool(True if value is None else value) | |
| def update_civitai_gallery_mode(enabled: bool, api_key: str, state: dict): | |
| state = state if isinstance(state, dict) else {} | |
| set_state(state, "civitai_gallery_enabled", bool(enabled)) | |
| return render_civitai_state(api_key, state, build_missing=False) | |
| def begin_probe_feedback(title: str): | |
| print(f"{title}: starting") | |
| return gr.update(value=f"### {title}\n- status: probing...", visible=True) | |
| def begin_probe_civitai(): | |
| return begin_probe_feedback("Civitai Probe") | |
| def begin_probe_civitai_keys(): | |
| return begin_probe_feedback("Civitai Key Probe") | |
| def begin_probe_civitai_url(): | |
| return begin_probe_feedback("Civitai URL Probe") | |
| def begin_probe_civitai_sidecar(): | |
| return begin_probe_feedback("Civitai Sidecar Probe") | |
| def get_effective_probe_api_key(api_key: str): | |
| value = str(api_key or "").strip() | |
| return value if value else str(os.environ.get("CIVITAI_API_KEY", "") or "").strip() | |
| def render_civitai_state(api_key: str, state: dict, build_missing: bool = True, info_override: str | None = None, page_label_override: str | None = None): | |
| state = state if isinstance(state, dict) else {} | |
| items = get_state(state, "civitai_last_items") or [] | |
| choices = get_state(state, "civitai_last_choices") or [("", "")] | |
| selected = list_uniq(get_state(state, "civitai_last_selects") or []) | |
| if not items: | |
| cleanup_search_preview_dir(state) | |
| return ( | |
| gr.update(value="Select All"), | |
| gr.update(choices=[("", "")], value=[], visible=True), | |
| gr.update(value="", visible=False), | |
| gr.update(value={}, visible=False), | |
| gr.update(value=[], visible=is_civitai_gallery_enabled(state)), | |
| gr.update(choices=[], value=[]), | |
| gr.update(value="No item found."), | |
| gr.update(value="Showing 0/0"), | |
| gr.update(value=None, visible=False), | |
| gr.update(value=format_civitai_key_status_md(api_key)), | |
| state, | |
| ) | |
| visible_count = int(get_state(state, "civitai_visible_count") or SEARCH_PAGE_SIZE) | |
| total_items = len(items) | |
| visible_count = min(max(visible_count, SEARCH_PAGE_SIZE), total_items) | |
| set_state(state, "civitai_visible_count", visible_count) | |
| gallery_enabled = is_civitai_gallery_enabled(state) | |
| visible_items = items[:visible_count] | |
| visible_values = [str(item.get('dl_url', '')) for item in visible_items if item.get('dl_url', '')] | |
| set_state(state, "civitai_visible_values", visible_values) | |
| all_values = [str(item.get('dl_url', '')) for item in items if item.get('dl_url', '')] | |
| detail_url = str(get_state(state, "civitai_detail_url") or "") | |
| if detail_url not in {item.get('dl_url', '') for item in items}: | |
| detail_url = selected[-1] if selected else (visible_values[0] if visible_values else "") | |
| set_state(state, "civitai_detail_url", detail_url) | |
| keep_paths = [] | |
| gallery = [] | |
| if gallery_enabled: | |
| for item in visible_items: | |
| thumb = resolve_item_preview(item, state, build_missing=build_missing) | |
| keep_paths.append(thumb) | |
| label = shorten_text(str(item.get('model_name') or item.get('name') or ''), 42) | |
| if item.get('dl_url', '') in selected: | |
| label = f"✓ {label}" | |
| gallery.append((thumb, label)) | |
| detail_item = find_item_by_url(state, detail_url) if detail_url else None | |
| detail_path = None | |
| detail_md = "" | |
| if detail_item: | |
| detail_path = resolve_item_preview(detail_item, state, build_missing=(build_missing or not gallery_enabled)) | |
| keep_paths.append(detail_path) | |
| detail_md = detail_item.get('md', '') | |
| prune_search_preview_dir(state, keep_paths) | |
| checkbox_choices = [(str(item.get('choice_name', build_item_label(item))), str(item.get('dl_url', ''))) for item in items] | |
| checkbox_values = [value for value in selected if value in all_values] | |
| info = info_override if info_override is not None else f"{total_items} items found. Showing {visible_count}. Selected {len(selected)}." | |
| if not gallery_enabled and items: | |
| info += " Gallery off." | |
| page_label = page_label_override if page_label_override is not None else f"Showing {visible_count}/{total_items}" | |
| return ( | |
| gr.update(value=select_all_button_label(state)), | |
| gr.update(choices=choices, value=selected, visible=True), | |
| gr.update(value=detail_md, visible=bool(detail_md)), | |
| gr.update(value={}, visible=False), | |
| gr.update(value=gallery, visible=gallery_enabled), | |
| gr.update(choices=checkbox_choices, value=checkbox_values), | |
| gr.update(value=info), | |
| gr.update(value=page_label), | |
| gr.update(value=detail_path, visible=bool(detail_path)), | |
| gr.update(value=format_civitai_key_status_md(api_key)), | |
| state, | |
| ) | |
| def load_more_civitai(state: dict, api_key: str = ""): | |
| state = state if isinstance(state, dict) else {} | |
| items = get_state(state, "civitai_last_items") or [] | |
| current = int(get_state(state, "civitai_visible_count") or SEARCH_PAGE_SIZE) | |
| if items: | |
| set_state(state, "civitai_visible_count", min(len(items), current + SEARCH_PAGE_SIZE)) | |
| else: | |
| set_state(state, "civitai_visible_count", SEARCH_PAGE_SIZE) | |
| return render_civitai_state(api_key, state, build_missing=False) | |
| def load_all_civitai(state: dict, api_key: str = ""): | |
| state = state if isinstance(state, dict) else {} | |
| items = get_state(state, "civitai_last_items") or [] | |
| if not items: | |
| set_state(state, "civitai_visible_count", SEARCH_PAGE_SIZE) | |
| yield render_civitai_state(api_key, state) | |
| return | |
| total_items = len(items) | |
| set_state(state, "civitai_visible_count", total_items) | |
| selected_count = len(get_state(state, 'civitai_last_selects') or []) | |
| if not is_civitai_gallery_enabled(state): | |
| yield render_civitai_state(api_key, state, build_missing=False, info_override=f"{total_items} items found. Showing {total_items}. Selected {selected_count}. Gallery off.", page_label_override=f"Showing {total_items}/{total_items}") | |
| return | |
| missing_items = [item for item in items if not is_item_preview_known(state, item)] | |
| total_missing = len(missing_items) | |
| initial_info = f"{total_items} items found. Showing {total_items}. Selected {selected_count}." | |
| if total_missing > 0: | |
| initial_info += f" Preview queue {total_missing}." | |
| yield render_civitai_state(api_key, state, build_missing=False, info_override=initial_info, page_label_override=f"Showing {total_items}/{total_items}") | |
| if total_missing == 0: | |
| return | |
| built = 0 | |
| for start in range(0, total_missing, LOAD_ALL_BATCH_SIZE): | |
| batch = missing_items[start:start + LOAD_ALL_BATCH_SIZE] | |
| for item in batch: | |
| resolve_item_preview(item, state, build_missing=True) | |
| built += len(batch) | |
| selected_count = len(get_state(state, 'civitai_last_selects') or []) | |
| info = f"{total_items} items found. Showing {total_items}. Selected {selected_count}. Loading previews {built}/{total_missing}." | |
| yield render_civitai_state(api_key, state, build_missing=False, info_override=info, page_label_override=f"Showing {total_items}/{total_items}") | |
| def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [], limit: int = 100, | |
| sort: str = "Highest Rated", period: str = "AllTime", tag: str = "", user: str = "", page: int = 1, | |
| filetype: list[str] = [], api_key: str = "", base_origin: str = CIVITAI_DEFAULT_ORIGIN, progress=gr.Progress(track_tqdm=False)): | |
| base_path = '/models' | |
| api_sort = sort if sort in {"Highest Rated", "Most Downloaded", "Newest"} else CIVITAI_SORT[0] | |
| params = {'sort': api_sort, 'period': period, 'limit': int(limit), 'nsfw': 'true'} | |
| clean_types = [str(t).strip() for t in (types or []) if str(t).strip()] | |
| if clean_types: | |
| params["types"] = clean_types | |
| if query: | |
| params["query"] = query | |
| if tag and str(tag).strip() not in {"", "None"}: | |
| params["tag"] = str(tag).strip() | |
| if user and str(user).strip(): | |
| params["username"] = str(user).strip() | |
| if page != 0: | |
| params["page"] = int(page) | |
| session = create_retry_session(total=6, backoff_factor=1.0) | |
| rs = [] | |
| try: | |
| reset_civitai_key_status(api_key, source="search") | |
| if page == 0: | |
| progress(0, desc="Searching page 1...") | |
| print("Searching page 1...") | |
| r = request_civitai_api(session, base_path, api_key=api_key, params=params | {'page': 1}, timeout=(7.0, 30), label='Civitai search page=1', source='search') | |
| rs.append(r) | |
| if r.ok: | |
| json = get_civitai_response_json(r, default={}) or {} | |
| next_url = json['metadata']['nextPage'] if 'metadata' in json and 'nextPage' in json['metadata'] else None | |
| i = 2 | |
| while next_url is not None: | |
| progress(0, desc=f"Searching page {i}...") | |
| print(f"Searching page {i}...") | |
| r = request_civitai_api_url(session, next_url, api_key=api_key, timeout=(7.0, 30), label=f'Civitai search page={i}', source='search') | |
| rs.append(r) | |
| if r.ok: | |
| json = get_civitai_response_json(r, default={}) or {} | |
| next_url = json['metadata']['nextPage'] if 'metadata' in json and 'nextPage' in json['metadata'] else None | |
| else: | |
| next_url = None | |
| i += 1 | |
| if next_url is not None: | |
| time.sleep(0.8) | |
| else: | |
| progress(0, desc="Searching page 1...") | |
| print("Searching page 1...") | |
| r = request_civitai_api(session, base_path, api_key=api_key, params=params, timeout=(7.0, 30), label='Civitai search page=1', source='search') | |
| rs.append(r) | |
| except requests.exceptions.ConnectTimeout: | |
| print("Request timed out.") | |
| except Exception as e: | |
| print(e) | |
| items = [] | |
| origin = get_civitai_display_origin().rstrip('/') | |
| api_item_count = 0 | |
| version_count = 0 | |
| file_count = 0 | |
| filtered_version_count = 0 | |
| filtered_file_count = 0 | |
| has_next_page = False | |
| for r in rs: | |
| if not r.ok: | |
| continue | |
| json = get_civitai_response_json(r, default={}) or {} | |
| if isinstance(json.get('metadata', {}), dict) and json.get('metadata', {}).get('nextPage'): | |
| has_next_page = True | |
| if 'items' not in json: | |
| continue | |
| api_item_count += len(json.get('items') or []) | |
| for j in json['items']: | |
| for model in j.get('modelVersions', []): | |
| version_count += 1 | |
| if len(allow_model) != 0 and model.get('baseModel', '') not in set(allow_model): | |
| filtered_version_count += 1 | |
| continue | |
| base_item = { | |
| 'name': j.get('name', ''), | |
| 'creator': j.get('creator', {}).get('username', '') if isinstance(j.get('creator', {}), dict) else '', | |
| 'tags': j.get('tags', []) if isinstance(j.get('tags', []), list) else [], | |
| 'model_name': model.get('name', ''), | |
| 'base_model': model.get('baseModel', ''), | |
| 'description': model.get('description', ''), | |
| 'model_id': j.get('id'), | |
| 'model_version_id': model.get('id'), | |
| 'origin': origin, | |
| } | |
| images = model.get('images', []) if isinstance(model.get('images', []), list) else [] | |
| if images: | |
| base_item['img_url'] = images[0].get('url', '') or NULL_IMAGE_PATH | |
| else: | |
| base_item['img_url'] = NULL_IMAGE_PATH | |
| model_url = build_civitai_model_url(j.get('id', ''), model.get('id')) | |
| desc = str(base_item.get('description', '') or '') | |
| base_item['md'] = ( | |
| f"Model URL: [{model_url}]({model_url})<br>" | |
| f"Model Name: {base_item['name']}<br>" | |
| f"Version: {base_item['model_name']}<br>" | |
| f"Creator: {base_item['creator']}<br>" | |
| f"Tags: {', '.join(base_item['tags'])}<br>" | |
| f"Base Model: {base_item['base_model']}<br>" | |
| f"Description: {desc}" | |
| ) | |
| files = model.get('files', []) if isinstance(model.get('files', []), list) else [] | |
| if files: | |
| file_count += len(files) | |
| for f in files: | |
| item = base_item.copy() | |
| item['dl_url'] = normalize_civitai_download_api_url(f.get('downloadUrl', '')) | |
| item['size_kb'] = f.get('sizeKB', 0.0) | |
| item['file_type'] = f.get('type', '') | |
| if len(filetype) != 0 and f.get('type', '') not in set(filetype): | |
| filtered_file_count += 1 | |
| continue | |
| items.append(item) | |
| else: | |
| item = base_item.copy() | |
| item['dl_url'] = normalize_civitai_download_api_url(model.get('downloadUrl', '')) | |
| items.append(item) | |
| if sort in CIVITAI_SORT_EXT: | |
| if sort == "Size": | |
| items = sorted(items, key=lambda x: x.get('size_kb', 0.0), reverse=True) | |
| elif sort == "Size (from smallest)": | |
| items = sorted(items, key=lambda x: x.get('size_kb', 0.0)) | |
| log_line("search", f"summary responses={len(rs)} api_items={api_item_count} versions={version_count} files={file_count} filtered_versions={filtered_version_count} filtered_files={filtered_file_count} selected_files={len(items)} page_mode={'all' if page == 0 else 'single'} has_next_page={str(has_next_page).lower()}") | |
| return items if len(items) > 0 else None | |
| def search_civitai(query, types, base_model=[], sort=CIVITAI_SORT[0], period=CIVITAI_PERIOD[0], tag="", user="", limit=100, page=1, | |
| filetype=[], api_key="", state=None, progress=gr.Progress(track_tqdm=False)): | |
| state = state if isinstance(state, dict) else {} | |
| cleanup_search_preview_dir(state) | |
| civitai_last_results = {} | |
| set_state(state, "civitai_last_choices", [("", "")]) | |
| set_state(state, "civitai_last_results", civitai_last_results) | |
| set_state(state, "civitai_last_selects", []) | |
| set_state(state, "civitai_last_items", []) | |
| set_state(state, "civitai_visible_count", SEARCH_PAGE_SIZE) | |
| set_state(state, "civitai_visible_values", []) | |
| set_state(state, "civitai_detail_url", "") | |
| set_state(state, "civitai_preview_fail_urls", []) | |
| set_state(state, "civitai_preview_cache", {}) | |
| if get_state(state, "civitai_gallery_enabled") is None: | |
| set_state(state, "civitai_gallery_enabled", True) | |
| items = search_on_civitai(query, types, base_model, int(limit), sort, period, tag, user, int(page), filetype, api_key, progress=progress) | |
| if not items: | |
| return render_civitai_state(api_key, state) | |
| choices = [] | |
| ordered_items = [] | |
| for item in items: | |
| value = str(item.get('dl_url', '') or '') | |
| if not value: | |
| continue | |
| choice_name = build_item_label(item) | |
| item['choice_name'] = choice_name | |
| choices.append((choice_name, value)) | |
| civitai_last_results[value] = {'md': item.get('md', ''), 'item': item} | |
| ordered_items.append(item) | |
| if not choices: | |
| return render_civitai_state(api_key, state) | |
| set_state(state, "civitai_last_choices", choices) | |
| set_state(state, "civitai_last_results", civitai_last_results) | |
| set_state(state, "civitai_last_items", ordered_items) | |
| set_state(state, "civitai_last_selects", []) | |
| set_state(state, "civitai_visible_count", SEARCH_PAGE_SIZE) | |
| set_state(state, "civitai_visible_values", [choice[1] for choice in choices[:SEARCH_PAGE_SIZE]]) | |
| set_state(state, "civitai_detail_url", choices[0][1]) | |
| return render_civitai_state(api_key, state) | |
| def save_info_preview_png(source_path: str, output_path: str): | |
| with Image.open(source_path) as img: | |
| ImageOps.exif_transpose(img).convert('RGBA').save(output_path, format='PNG') | |
| return output_path | |
| def extract_video_preview_png(source_path: str, output_path: str): | |
| ffmpeg_path = get_ffmpeg_path() | |
| if not ffmpeg_path: | |
| return "" | |
| cmd = [ffmpeg_path, '-y', '-loglevel', 'error', '-nostdin', '-i', str(source_path), '-frames:v', '1', str(output_path)] | |
| subprocess.run(cmd, check=True, timeout=30) | |
| return output_path if Path(output_path).exists() else "" | |
| def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str="", api_key="", temp_dir=""): | |
| original_dl_url = str(dl_url or "").strip() | |
| dl_url = normalize_civitai_input_url(original_dl_url, api_key=api_key) | |
| if not image_baseurl: | |
| image_baseurl = dl_url | |
| default = ("", "", "") if is_html else "" | |
| if "https://civitai.com/api/download/models/" not in dl_url: | |
| return default | |
| base_path = '/model-versions' | |
| params = {} | |
| session = create_retry_session(total=6, backoff_factor=1.0) | |
| model_id = re.sub(r'https://civitai.com/api/download/models/(\d+)(?:.+)?', r'\1', dl_url) | |
| url = f"{base_path}/{model_id}" | |
| try: | |
| reset_civitai_key_status(api_key, source='info') | |
| r = request_civitai_api(session, url, api_key=api_key, params=params, timeout=(5.0, 60), label='Civitai model-version', source='info') | |
| if not r.ok: | |
| return default | |
| json = dict(get_civitai_response_json(r, default={}) or {}).copy() | |
| html = "" | |
| image = "" | |
| if "modelId" in json.keys(): | |
| original_parts = get_civitai_url_parts(original_dl_url) | |
| if is_civitai_host(original_parts.netloc) and re.match(r'^/models/\d+(?:/[^/?#]+)?/?$', original_parts.path or ""): | |
| url = original_dl_url | |
| else: | |
| url = build_civitai_model_url(json['modelId']) | |
| r = civitai_get(session, url, api_key=api_key, params=params, timeout=(5.0, 60), label='Civitai model page', source='info') | |
| if not r.ok: | |
| return json, html, image | |
| html = r.text | |
| if 'images' in json.keys() and len(json["images"]) != 0: | |
| url = str(json["images"][0].get("url", "") or "") | |
| if url: | |
| r = civitai_get(session, url, api_key=api_key, params=params, timeout=(5.0, 60), label='Civitai preview image', source='info') | |
| if not r.ok: | |
| return json, html, image | |
| target_dir = temp_dir if temp_dir else TEMP_DIR | |
| preview_suffix = Path(urllib.parse.urlsplit(url).path).suffix | |
| content_type = str(r.headers.get('content-type') or '').split(';', 1)[0].strip().lower() | |
| if not preview_suffix: | |
| preview_suffix = mimetypes.guess_extension(content_type or '') or '' | |
| image_temp = str(Path(target_dir, Path(image_baseurl.split("/")[-1]).stem + "_preview" + (preview_suffix or ".bin"))) | |
| image = str(Path(target_dir, Path(image_baseurl.split("/")[-1]).stem + ".png")) | |
| with open(image_temp, 'wb') as f: | |
| f.write(r.content) | |
| try: | |
| is_video = content_type.startswith('video/') or is_probable_video_url(url) | |
| if is_video: | |
| image = extract_video_preview_png(image_temp, image) | |
| else: | |
| image = save_info_preview_png(image_temp, image) | |
| except Exception as e: | |
| image = "" | |
| print(f"Failed to prepare Civitai info preview for {dl_url}. {e}") | |
| finally: | |
| if Path(image_temp).exists(): | |
| Path(image_temp).unlink() | |
| return json, html, image | |
| except Exception as e: | |
| print(e) | |
| return default | |
| def _load_civitai_choice_list(api_path: str, label: str, source: str, value_key: str, count_key: str, | |
| query: str="", api_key: str="", limit: int=200, | |
| timeout: tuple[float, float]=(7.0, 15.0), retry_total: int=6, | |
| backoff_factor: float=1.0, default: list[str] | None=None, | |
| cache: dict | None=None, cache_lock: object | None=None, | |
| cache_ttl_sec: int=0, startup_name: str="", preferred_origin: str="", | |
| non_json_fallback_origin: str=""): | |
| clean_query = str(query or "").strip() | |
| cache_key = clean_query.casefold() | |
| now = time.time() | |
| default = list(default or [""]) | |
| if cache is not None and cache_lock is not None and cache_ttl_sec > 0: | |
| with cache_lock: | |
| cached = cache.get(cache_key) | |
| if isinstance(cached, dict) and (now - float(cached.get("ts") or 0.0)) <= cache_ttl_sec: | |
| return list(cached.get("choices") or default) | |
| params = {'limit': max(1, min(int(limit), 200))} | |
| if clean_query: | |
| params['query'] = clean_query | |
| session = create_retry_session(total=retry_total, backoff_factor=backoff_factor) | |
| try: | |
| r = request_civitai_api( | |
| session, | |
| api_path, | |
| api_key=api_key, | |
| params=params, | |
| timeout=timeout, | |
| label=label, | |
| source=source, | |
| preferred_origin=preferred_origin, | |
| non_json_fallback_origin=non_json_fallback_origin, | |
| ) | |
| if not r.ok: | |
| if not clean_query and startup_name: | |
| print(f"Failed to load {startup_name} from API at startup. HTTP {r.status_code}.") | |
| return default | |
| j = dict(get_civitai_response_json(r, default={}) or {}).copy() | |
| items = j.get('items', []) if isinstance(j.get('items', []), list) else [] | |
| observed = [] | |
| for item in items: | |
| value = str(item.get(value_key, '') or '').strip() | |
| if value: | |
| observed.append((value, int(item.get(count_key, 0) or 0))) | |
| observed = sorted(observed, key=lambda x: (-x[1], x[0].casefold())) | |
| choices = [""] + [name for name, _ in observed] | |
| if cache is not None and cache_lock is not None and cache_ttl_sec > 0: | |
| with cache_lock: | |
| cache[cache_key] = {'choices': choices, 'ts': now} | |
| if not clean_query and startup_name: | |
| print(f"Loaded {len(observed)} {startup_name} from API at startup.") | |
| return choices if choices else default | |
| except Exception as e: | |
| if not clean_query and startup_name: | |
| print(f"Failed to load {startup_name} from API at startup. {type(e).__name__}: {e}") | |
| return default | |
| def get_civitai_creator(query: str="", api_key: str="", limit: int=CREATOR_SUGGEST_LIMIT): | |
| return _load_civitai_choice_list( | |
| api_path='/creators', | |
| label='Civitai creators', | |
| source='creators', | |
| value_key='username', | |
| count_key='modelCount', | |
| query=query, | |
| api_key=api_key, | |
| limit=limit, | |
| timeout=(7.0, 15.0), | |
| retry_total=6, | |
| backoff_factor=1.0, | |
| default=[""], | |
| cache=CREATOR_SUGGEST_CACHE, | |
| cache_lock=CREATOR_SUGGEST_LOCK, | |
| cache_ttl_sec=CREATOR_CACHE_TTL_SEC, | |
| startup_name='Civitai creators', | |
| non_json_fallback_origin=CIVITAI_DEFAULT_ORIGIN, | |
| ) | |
| def refresh_civitai_creators(user_value: str="", api_key: str=""): | |
| value = str(user_value or "").strip() | |
| with CREATOR_SUGGEST_LOCK: | |
| cached = CREATOR_SUGGEST_CACHE.get("") | |
| base_choices = list(cached.get("choices") or [""]) if isinstance(cached, dict) else [""] | |
| if not value: | |
| return gr.update(choices=base_choices, value=value) | |
| folded = value.casefold() | |
| prefix_choices = [choice for choice in base_choices if choice and choice.casefold().startswith(folded)] | |
| contains_choices = [choice for choice in base_choices if choice and choice not in prefix_choices and folded in choice.casefold()] | |
| choices = [""] + prefix_choices + contains_choices | |
| if value not in choices: | |
| choices.insert(1, value) | |
| return gr.update(choices=list_uniq(choices), value=value) | |
| def get_civitai_tag(): | |
| return _load_civitai_choice_list( | |
| api_path='/tags', | |
| label='Civitai tags', | |
| source='tags', | |
| value_key='name', | |
| count_key='modelCount', | |
| limit=200, | |
| timeout=(7.0, 15.0), | |
| retry_total=6, | |
| backoff_factor=1.0, | |
| default=[""], | |
| startup_name='Civitai tags', | |
| ) | |
| def select_civitai_item(results: list[str], state: dict): | |
| state = state if isinstance(state, dict) else {} | |
| set_state(state, "civitai_last_selects", list_uniq(results or [])) | |
| if results: | |
| set_state(state, "civitai_detail_url", results[-1]) | |
| rendered = render_civitai_state("", state) | |
| return rendered[2], rendered[3], state | |
| def add_civitai_item(results: list[str], dl_url: str): | |
| if "http" not in "".join(results): | |
| return gr.update(value=dl_url) | |
| new_url = dl_url if dl_url else "" | |
| for result in results: | |
| if "http" not in result: | |
| continue | |
| new_url += f"\n{result}" if new_url else f"{result}" | |
| new_url = uniq_urls(new_url) | |
| return gr.update(value=new_url) | |
| def from_civitai_dropdown(selected: list[str], api_key: str, state: dict): | |
| state = state if isinstance(state, dict) else {} | |
| selected = list_uniq(selected or []) | |
| set_state(state, "civitai_last_selects", selected) | |
| if selected: | |
| set_state(state, "civitai_detail_url", selected[-1]) | |
| return render_civitai_state(api_key, state, build_missing=False) | |
| def get_gallery_event_index(evt) -> int | None: | |
| try: | |
| data = getattr(evt, "_data", {}) or {} | |
| index = data.get("index") | |
| if isinstance(index, (list, tuple)): | |
| index = index[0] if index else None | |
| return int(index) if index is not None else None | |
| except Exception: | |
| return None | |
| def update_civitai_selection(evt: gr.EventData, value: list[str], api_key: str, state: dict): | |
| state = state if isinstance(state, dict) else {} | |
| selected = list_uniq(value or []) | |
| visible_values = get_state(state, "civitai_visible_values") or [] | |
| selected_index = get_gallery_event_index(evt) | |
| if selected_index is not None and 0 <= selected_index < len(visible_values): | |
| selected_value = visible_values[selected_index] | |
| if selected_value in selected: | |
| selected = [v for v in selected if v != selected_value] | |
| else: | |
| selected.append(selected_value) | |
| set_state(state, "civitai_last_selects", list_uniq(selected)) | |
| set_state(state, "civitai_detail_url", selected_value) | |
| return render_civitai_state(api_key, state, build_missing=False) | |
| def from_civitai_checkbox(selected: list[str], api_key: str, state: dict): | |
| state = state if isinstance(state, dict) else {} | |
| choices = get_state(state, "civitai_last_choices") or [] | |
| allowed_values = {value for _, value in choices if value} | |
| selected_all = [v for v in list_uniq(selected or []) if v in allowed_values] | |
| set_state(state, "civitai_last_selects", selected_all) | |
| if selected_all: | |
| set_state(state, "civitai_detail_url", selected_all[-1]) | |
| return render_civitai_state(api_key, state, build_missing=False) | |
| def select_civitai_all_item_fast(button_name: str, api_key: str, state: dict): | |
| state = state if isinstance(state, dict) else {} | |
| choices = get_state(state, "civitai_last_choices") or [] | |
| if button_name not in ["Select All", "Deselect All"]: | |
| return render_civitai_state(api_key, state, build_missing=False) | |
| selected = [t[1] for t in choices if t[1] != ""] if button_name == "Select All" else [] | |
| set_state(state, "civitai_last_selects", selected) | |
| if selected: | |
| set_state(state, "civitai_detail_url", selected[-1]) | |
| return render_civitai_state(api_key, state, build_missing=False) | |
| class _FakeResponse: | |
| def __init__(self, headers=None): | |
| self.headers = headers or {} | |
| class _FakeHFError(Exception): | |
| def __init__(self, message, headers=None): | |
| super().__init__(message) | |
| self.response = _FakeResponse(headers=headers) | |
| class _FakeUploadApi: | |
| def __init__(self, exists_after_error=False): | |
| self.exists_after_error = bool(exists_after_error) | |
| self.upload_calls = 0 | |
| self.exists_calls = 0 | |
| def upload_file(self, **kwargs): | |
| self.upload_calls += 1 | |
| raise _FakeHFError("Bad request for commit endpoint: Unexpected internal error hook: lfs-verify") | |
| def file_exists(self, **kwargs): | |
| self.exists_calls += 1 | |
| return self.exists_after_error | |
| def safe_retry_probe(hf_retry_policy="Auto"): | |
| config = get_hf_upload_retry_policy_config(hf_retry_policy) | |
| cases = [ | |
| ("429 with Retry-After", _FakeHFError("HTTP Error 429 Too Many Requests", {"Retry-After": "3", "RateLimit": "api|r=0;t=183"})), | |
| ("503 LFS batch", _FakeHFError("HTTP Error 503 while requesting POST /info/lfs/objects/batch")), | |
| ("lfs-verify hook", _FakeHFError("Bad request for commit endpoint: Unexpected internal error hook: lfs-verify")), | |
| ("403 permission", _FakeHFError("403 Forbidden: permission denied")), | |
| ("repo not found", _FakeHFError("Repo not found")), | |
| ] | |
| lines = ["### Safe Retry Probe", "- network: none", f"- policy: {config.get('key')} attempts={config.get('attempts')} base_wait={config.get('base_wait')} max_wait={config.get('max_wait')}"] | |
| for label, exc in cases: | |
| retryable = is_retryable_hf_upload_exception(exc) | |
| delay, delay_source = parse_hf_retry_delay_from_headers(exc) | |
| hint = format_hf_rate_limit_hint(exc) | |
| parts = [f"- {label}: {'retryable' if retryable else 'not retryable'}"] | |
| if delay is not None: | |
| parts.append(f"delay={delay:g}s source={delay_source}") | |
| if hint: | |
| parts.append(hint) | |
| lines.append(" | ".join(parts)) | |
| return gr.update(value="\n".join(lines), visible=True) | |
| def safe_upload_verify_probe(hf_retry_policy="Auto"): | |
| config = get_hf_upload_retry_policy_config(hf_retry_policy) | |
| lines = ["### Safe Upload Verify Probe", "- network: none", "- upload_file failure is simulated", f"- policy: {config.get('key')} (no sleeps, no real retries)"] | |
| for exists_after_error in (True, False): | |
| fake_api = _FakeUploadApi(exists_after_error=exists_after_error) | |
| try: | |
| try: | |
| fake_api.upload_file() | |
| except Exception as e: | |
| retryable = is_retryable_hf_upload_exception(e) | |
| recovered = fake_api.file_exists(repo_id="user/repo", filename="file.safetensors", repo_type="model", token="[redacted]") | |
| error_short = format_error_short(e) | |
| else: | |
| retryable = False | |
| recovered = True | |
| error_short = "" | |
| state = "recovered" if recovered else "failed" | |
| lines.append(f"- remote_exists_after_error={exists_after_error}: {state} retryable={retryable} upload_calls={fake_api.upload_calls} file_exists_calls={fake_api.exists_calls} error={error_short}") | |
| except Exception as e: | |
| lines.append(f"- remote_exists_after_error={exists_after_error}: probe error {type(e).__name__}: {format_error_short(e)}") | |
| return gr.update(value="\n".join(lines), visible=True) | |
| def safe_summary_probe(): | |
| lines = ["### Safe Summary Probe", "- network: none"] | |
| scenarios = [ | |
| ("clean", {"input_urls": 3, "downloaded": 3, "uploaded": 3, "skipped_duplicate": 0, "failed_download": 0, "failed_upload": 0, "verified_after_error": 0}, [], [], "Done"), | |
| ("interrupted", {"input_urls": 5, "downloaded": 2, "uploaded": 2, "skipped_duplicate": 0, "failed_download": 0, "failed_upload": 0, "verified_after_error": 0}, ["https://civitai.com/api/download/models/1"], [], "Incomplete"), | |
| ("hf upload failures", {"input_urls": 4, "downloaded": 4, "uploaded": 2, "skipped_duplicate": 0, "failed_download": 0, "failed_upload": 2, "verified_after_error": 1}, [], ["https://civitai.com/api/download/models/2"], "Incomplete"), | |
| ("duplicates", {"input_urls": 4, "downloaded": 1, "uploaded": 1, "skipped_duplicate": 3, "failed_download": 0, "failed_upload": 0, "verified_after_error": 0}, [], [], "Done"), | |
| ] | |
| for name, stats, remaining, failed, stage in scenarios: | |
| lines.append("") | |
| lines.append(f"#### {name}") | |
| lines.extend(build_run_summary_lines(stats, remaining, failed, stage)) | |
| return gr.update(value="\n".join(lines), visible=True) | |
| def create_report_zip(session_state=None, search_state=None): | |
| state = ensure_session_state(session_state) | |
| search_state = search_state if isinstance(search_state, dict) else {} | |
| report_dir = Path(tempfile.mkdtemp(prefix="civitai_report_", dir=TEMP_DIR)) | |
| timestamp = time.strftime("%Y%m%d_%H%M%S", time.gmtime()) | |
| zip_path = report_dir / f"{REPORT_ZIP_PREFIX}_{timestamp}.zip" | |
| session_events = list(state.get("session_report_events") or []) | |
| current_events = list(state.get("report_events") or []) | |
| run_records = list(state.get("session_run_records") or []) | |
| last_record = state.get("last_run_record") if isinstance(state.get("last_run_record"), dict) else {} | |
| if last_record and not any(str(r.get("run_id") or "") == str(last_record.get("run_id") or "") for r in run_records): | |
| run_records.append(last_record) | |
| remaining = list(state.get("last_remaining_urls") or state.get("current_remaining_urls") or []) | |
| failed = list(state.get("last_failed_urls") or state.get("current_failed_urls") or []) | |
| uploaded = list(state.get("last_uploaded_urls") or state.get("current_uploaded_urls") or []) | |
| smoke_lines = list(state.get("last_smoke_lines") or state.get("current_smoke_lines") or []) | |
| summary = dict(state.get("last_run_summary") or {}) | |
| if not summary: | |
| summary = { | |
| "stage": str(state.get("current_stage") or ""), | |
| "repo_id": str(state.get("last_repo_id") or state.get("current_repo_id") or ""), | |
| "repo_type": str(state.get("last_repo_type") or state.get("current_repo_type") or ""), | |
| "remaining": len(remaining), | |
| "failed": len(failed), | |
| "uploaded": len(uploaded), | |
| } | |
| current_run_id = str(summary.get("run_id") or state.get("last_run_id") or state.get("current_run_id") or "") | |
| if not last_record: | |
| run_events = list_report_events_for_run(session_events or current_events, current_run_id) | |
| last_record = build_report_run_record(current_run_id, summary, run_events or current_events, remaining, failed, uploaded, smoke_lines=smoke_lines, failure_reasons=state.get("last_failure_reasons") or {}) | |
| summary = redact_report_value(summary) | |
| session_summary = summarize_report_runs(run_records) | |
| if not run_records and last_record: | |
| run_records = [last_record] | |
| session_summary = summarize_report_runs(run_records) | |
| search_summary = { | |
| "last_choices": len(search_state.get("civitai_last_choices") or []), | |
| "last_results": len(search_state.get("civitai_last_results") or {}), | |
| "last_selects": len(search_state.get("civitai_last_selects") or []), | |
| "last_items": len(search_state.get("civitai_last_items") or []), | |
| "visible_count": int(search_state.get("civitai_visible_count") or 0), | |
| "gallery_enabled": bool(search_state.get("civitai_gallery_enabled", True)), | |
| "detail_url": search_state.get("civitai_detail_url") or "", | |
| "preview_fail_urls": list(search_state.get("civitai_preview_fail_urls") or []), | |
| } | |
| selected_urls = [] | |
| for value in list(search_state.get("civitai_last_selects") or []): | |
| if value: | |
| selected_urls.append(str(value)) | |
| env = { | |
| "created_at_utc": utc_timestamp(), | |
| "python": sys.version.split()[0], | |
| "platform": platform.platform(), | |
| "gradio": get_package_version("gradio"), | |
| "huggingface_hub": get_package_version("huggingface_hub"), | |
| "requests": get_package_version("requests"), | |
| "civitai_hf_debug": bool(os.environ.get("CIVITAI_HF_DEBUG")), | |
| "hf_upload_retry_policy": str(state.get("last_hf_retry_policy") or state.get("current_hf_retry_policy") or ""), | |
| } | |
| run_table_lines = ["| run_id | mode | stage | input | uploaded | skipped | failed download | failed upload | remaining |", "|---|---:|---:|---:|---:|---:|---:|---:|---:|"] | |
| for record in run_records: | |
| run_summary = record.get("summary") if isinstance(record, dict) else {} | |
| if not isinstance(run_summary, dict): | |
| run_summary = {} | |
| rid = str(record.get("run_id") or run_summary.get("run_id") or "") | |
| run_table_lines.append( | |
| f"| {redact_report_value(rid)} | {redact_report_value(record.get('mode') or run_summary.get('mode') or '')} | {redact_report_value(record.get('stage') or run_summary.get('stage') or '')} | {int(run_summary.get('input_urls', 0) or 0)} | {int(run_summary.get('uploaded', 0) or 0)} | {int(run_summary.get('skipped_duplicate', 0) or 0)} | {int(run_summary.get('failed_download', 0) or 0)} | {int(run_summary.get('failed_upload', 0) or 0)} | {int(run_summary.get('remaining', 0) or 0)} |" | |
| ) | |
| report_md = [ | |
| "# Civitai to HF Diagnostic Report", | |
| "", | |
| "## Current Run Summary", | |
| "```json", | |
| safe_json_dumps(summary), | |
| "```", | |
| "", | |
| build_report_advice(summary), | |
| "", | |
| "## Session Summary", | |
| "```json", | |
| safe_json_dumps(session_summary), | |
| "```", | |
| "", | |
| "## Runs", | |
| "", | |
| *run_table_lines, | |
| ] | |
| if smoke_lines: | |
| report_md.extend(["", "## Smoke / Probe Lines", ""]) | |
| report_md.extend([f"- {redact_report_value(line)}" for line in smoke_lines]) | |
| with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf: | |
| report_write_text(zf, "report.md", "\n".join(report_md).strip() + "\n") | |
| report_write_text(zf, "current_run/summary.json", safe_json_dumps(summary) + "\n") | |
| report_write_text(zf, "current_run/events.jsonl", "\n".join(json.dumps(redact_report_value(ev), ensure_ascii=False, sort_keys=True) for ev in (last_record.get("events") or current_events)) + ("\n" if (last_record.get("events") or current_events) else "")) | |
| report_write_text(zf, "current_run/remaining_urls.txt", "\n".join(redact_report_value(u) for u in (last_record.get("remaining_urls") or remaining)) + ("\n" if (last_record.get("remaining_urls") or remaining) else "")) | |
| report_write_text(zf, "current_run/failed_urls.txt", "\n".join(redact_report_value(u) for u in (last_record.get("failed_urls") or failed)) + ("\n" if (last_record.get("failed_urls") or failed) else "")) | |
| report_write_text(zf, "current_run/uploaded_urls.txt", "\n".join(redact_report_value(u) for u in (last_record.get("uploaded_urls") or uploaded)) + ("\n" if (last_record.get("uploaded_urls") or uploaded) else "")) | |
| report_write_text(zf, "current_run/advice.md", build_report_advice(summary)) | |
| report_write_text(zf, "session_summary.json", safe_json_dumps(session_summary) + "\n") | |
| report_write_text(zf, "session_events.jsonl", "\n".join(json.dumps(redact_report_value(ev), ensure_ascii=False, sort_keys=True) for ev in session_events) + ("\n" if session_events else "")) | |
| report_write_text(zf, "environment.json", safe_json_dumps(env) + "\n") | |
| report_write_text(zf, "search_summary.json", safe_json_dumps(search_summary) + "\n") | |
| report_write_text(zf, "selected_search_urls.txt", "\n".join(redact_report_value(u) for u in selected_urls) + ("\n" if selected_urls else "")) | |
| # Legacy top-level files kept for quick manual inspection. | |
| report_write_text(zf, "summary.json", safe_json_dumps(summary) + "\n") | |
| report_write_text(zf, "events.jsonl", "\n".join(json.dumps(redact_report_value(ev), ensure_ascii=False, sort_keys=True) for ev in current_events) + ("\n" if current_events else "")) | |
| report_write_text(zf, "remaining_urls.txt", "\n".join(redact_report_value(u) for u in remaining) + ("\n" if remaining else "")) | |
| report_write_text(zf, "failed_urls.txt", "\n".join(redact_report_value(u) for u in failed) + ("\n" if failed else "")) | |
| report_write_text(zf, "uploaded_urls.txt", "\n".join(redact_report_value(u) for u in uploaded) + ("\n" if uploaded else "")) | |
| for record in run_records: | |
| run_id = str(record.get("run_id") or "run") or "run" | |
| safe_run_id = re.sub(r"[^A-Za-z0-9_.-]+", "_", run_id)[:80] or "run" | |
| prefix = f"runs/{safe_run_id}" | |
| run_summary = record.get("summary") if isinstance(record.get("summary"), dict) else {} | |
| report_write_text(zf, f"{prefix}/summary.json", safe_json_dumps(run_summary) + "\n") | |
| report_write_text(zf, f"{prefix}/events.jsonl", "\n".join(json.dumps(redact_report_value(ev), ensure_ascii=False, sort_keys=True) for ev in list(record.get("events") or [])) + ("\n" if record.get("events") else "")) | |
| report_write_text(zf, f"{prefix}/remaining_urls.txt", "\n".join(redact_report_value(u) for u in list(record.get("remaining_urls") or [])) + ("\n" if record.get("remaining_urls") else "")) | |
| report_write_text(zf, f"{prefix}/failed_urls.txt", "\n".join(redact_report_value(u) for u in list(record.get("failed_urls") or [])) + ("\n" if record.get("failed_urls") else "")) | |
| report_write_text(zf, f"{prefix}/uploaded_urls.txt", "\n".join(redact_report_value(u) for u in list(record.get("uploaded_urls") or [])) + ("\n" if record.get("uploaded_urls") else "")) | |
| report_write_text(zf, f"{prefix}/advice.md", str(record.get("advice") or build_report_advice(run_summary))) | |
| curated_state = { | |
| "last_run_id": state.get("last_run_id"), | |
| "last_run_mode": state.get("last_run_mode"), | |
| "last_repo_id": state.get("last_repo_id"), | |
| "last_repo_type": state.get("last_repo_type"), | |
| "last_error": state.get("last_error"), | |
| "last_failure_summary": state.get("last_failure_summary"), | |
| "last_run_summary": state.get("last_run_summary"), | |
| "session_summary": session_summary, | |
| "session_run_count": len(run_records), | |
| "current_stage": state.get("current_stage"), | |
| "current_stage_detail": state.get("current_stage_detail"), | |
| "run_elapsed_sec": state.get("run_elapsed_sec"), | |
| } | |
| report_write_text(zf, "session_state_curated.json", safe_json_dumps(curated_state) + "\n") | |
| log_line("probe", f"created report zip: {zip_path}") | |
| return str(zip_path) | |
| def refresh_civitai_key_status(api_key: str = ""): | |
| reset_civitai_key_status(api_key, source="input") | |
| return format_civitai_key_status_md(api_key) | |
| def clear_retry_state(session_state=None): | |
| session_state = ensure_session_state(session_state) | |
| previous_run_id = str(session_state.get("active_run_id") or "") | |
| previous_temp_dir = str(session_state.get("current_run_temp_dir") or "") | |
| if previous_run_id: | |
| unregister_run(previous_run_id) | |
| if previous_temp_dir and is_safe_run_temp_dir(previous_temp_dir): | |
| cleanup_run_temp_dir(previous_temp_dir) | |
| session_state_update(session_state, current_remaining_urls=[], current_failed_urls=[], last_remaining_urls=[], last_failed_urls=[], cancel_requested=False, current_stage="", current_stage_detail="", current_url="", current_item_index=0, current_item_total=0, current_run_temp_dir="", active_run_id="", repo_hash_cache={}, last_error="", last_failure_summary="", run_started_at=0.0, run_elapsed_sec=0.0) | |
| log_line("cleanup", "cleared retry state") | |
| return build_run_status_update(session_state), gr.update(value="", visible=False), gr.update(value="", visible=False), session_state_output(session_state) | |
| def extract_first_model_path(html: str): | |
| if not html: | |
| return "" | |
| m = re.search(r"href=['\"](/models/\d+(?:/[^'\"#?]+)?(?:\?modelVersionId=\d+)?)", html) | |
| return m.group(1) if m else "" | |
| def extract_probe_first_download_url(payload: dict): | |
| items = payload.get('items', []) if isinstance(payload, dict) else [] | |
| for model in items: | |
| versions = model.get('modelVersions', []) if isinstance(model.get('modelVersions', []), list) else [] | |
| for version in versions: | |
| files = version.get('files', []) if isinstance(version.get('files', []), list) else [] | |
| for file_info in files: | |
| dl_url = str(file_info.get('downloadUrl', '') or '').strip() | |
| if dl_url: | |
| return dl_url, model, version, file_info | |
| dl_url = str(version.get('downloadUrl', '') or '').strip() | |
| if dl_url: | |
| return dl_url, model, version, {} | |
| return "", {}, {}, {} | |
| def probe_json_api_step(session, path: str, *, api_key: str = "", params=None, timeout: tuple[float, float] = (7.0, 20.0), | |
| label: str = "Civitai probe", source: str = "probe", preferred_origin: str = "", | |
| non_json_fallback_origin: str = ""): | |
| response = None | |
| try: | |
| response = request_civitai_api( | |
| session, | |
| path, | |
| api_key=api_key, | |
| params=params, | |
| timeout=timeout, | |
| label=label, | |
| source=source, | |
| preferred_origin=preferred_origin, | |
| non_json_fallback_origin=non_json_fallback_origin, | |
| ) | |
| ok = bool(response is not None and getattr(response, "ok", False)) | |
| status = str(getattr(response, "status_code", "-")) | |
| response_url = str(getattr(response, "url", "") or "") | |
| response_host = urllib.parse.urlparse(response_url).netloc | |
| json_ok = False | |
| if ok: | |
| payload = get_civitai_response_json(response, default=_CIVITAI_JSON_MISSING) | |
| json_ok = payload is not _CIVITAI_JSON_MISSING | |
| return {"ok": ok, "status": status, "host": response_host, "json_ok": json_ok, "error": ""} | |
| except Exception as e: | |
| return {"ok": False, "status": "-", "host": "", "json_ok": False, "error": f"{type(e).__name__}: {e}"} | |
| finally: | |
| try: | |
| if response is not None: | |
| response.close() | |
| except Exception: | |
| pass | |
| def probe_download_url_step(session, download_url: str, *, api_key: str = "", timeout: tuple[float, float] = (7.0, 20.0)): | |
| response = None | |
| try: | |
| response = civitai_get(session, download_url, api_key=api_key, timeout=timeout, label='Civitai probe download', source='probe-download') | |
| status = str(getattr(response, "status_code", "-")) | |
| status_code = int(status) if str(status).isdigit() else 0 | |
| response_url = str(getattr(response, "url", "") or download_url or "") | |
| response_host = urllib.parse.urlparse(response_url).netloc | |
| ok = bool(response is not None and (getattr(response, "ok", False) or status_code in {301, 302, 303, 307, 308})) | |
| return {"ok": ok, "status": status, "host": response_host, "error": ""} | |
| except Exception as e: | |
| return {"ok": False, "status": "-", "host": "", "error": f"{type(e).__name__}: {e}"} | |
| finally: | |
| try: | |
| if response is not None: | |
| response.close() | |
| except Exception: | |
| pass | |
| def probe_civitai_keys(api_key: str): | |
| effective_api_key = get_effective_probe_api_key(api_key) | |
| key_source = "input" if str(api_key or "").strip() else ("env" if effective_api_key else "none") | |
| parsed_keys = parse_civitai_api_keys(effective_api_key) | |
| print(f"Civitai Key Probe: keys={len(parsed_keys)} key_source={key_source}") | |
| lines = ["### Civitai Key Probe", f"- key source: {key_source}", f"- parsed keys: {len(parsed_keys)}"] | |
| if not parsed_keys: | |
| lines.append("- auth api: skipped (no Civitai key)") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| session = create_retry_session(total=4, backoff_factor=0.8) | |
| try: | |
| api_origin = resolve_civitai_api_origin(session) | |
| lines.append(f"- api origin: {api_origin}") | |
| reset_civitai_key_status(effective_api_key, source="probe-keys") | |
| resp = request_civitai_api( | |
| session, | |
| "/models", | |
| api_key=effective_api_key, | |
| params={"limit": 1, "sort": "Newest", "period": "AllTime", "nsfw": "true"}, | |
| timeout=(7.0, 20.0), | |
| label='Civitai key probe', | |
| source='probe-keys', | |
| ) | |
| ok = bool(resp is not None and getattr(resp, "ok", False)) | |
| status_code = getattr(resp, "status_code", "-") | |
| lines.append(f"- auth api: {'ok' if ok else 'fail'} ({status_code})") | |
| status = get_civitai_key_status(effective_api_key) | |
| active_index = int(status.get("active_index") or 1) | |
| lines.append(f"- active key: {min(max(active_index, 1), len(parsed_keys))}/{len(parsed_keys)}") | |
| last_status = str(status.get("last_status") or "") | |
| if last_status: | |
| lines.append(f"- last status: {last_status}") | |
| switch_reason = str(status.get("last_switch_reason") or "") | |
| if switch_reason: | |
| lines.append(f"- switch: {switch_reason[:160]}") | |
| except Exception as e: | |
| lines.append(f"- auth probe: fail ({type(e).__name__}: {e})") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| def get_probe_url_candidate(url_input: str, fallback_query: str = ""): | |
| urls = normalize_url_entries(url_input) | |
| if urls: | |
| return str(urls[0]).strip() | |
| candidate = normalize_input_token(fallback_query) | |
| parts = get_civitai_url_parts(candidate) | |
| if parts.scheme in {"http", "https"} and parts.netloc: | |
| return candidate | |
| return "" | |
| def get_civitai_probe_mode(url: str): | |
| parts = get_civitai_url_parts(url) | |
| path = str(parts.path or "") | |
| if not is_civitai_host(parts.netloc): | |
| return "external" | |
| if is_civitai_download_api_path(path): | |
| return "direct-download" | |
| if re.match(r'^/models/\d+(?:/[^/?#]+)?/?$', path): | |
| return "model-page" | |
| return "civitai-other" | |
| def probe_civitai_url(url_input: str, fallback_query: str, api_key: str): | |
| effective_api_key = get_effective_probe_api_key(api_key) | |
| key_source = "input" if str(api_key or "").strip() else ("env" if effective_api_key else "none") | |
| probe_input = get_probe_url_candidate(url_input, fallback_query) | |
| print(f"Civitai URL Probe: input={sanitize_url_for_log(probe_input)} key_source={key_source}") | |
| lines = ["### Civitai URL Probe", f"- key source: {key_source}"] | |
| if not probe_input: | |
| lines.append("- input: skipped (no URL found in Download URL(s) or Query)") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| parts = get_civitai_url_parts(probe_input) | |
| mode = get_civitai_probe_mode(probe_input) | |
| lines.append(f"- input: {sanitize_url_for_log(probe_input)}") | |
| if parts.netloc: | |
| lines.append(f"- input host: {parts.netloc}") | |
| lines.append(f"- mode: {mode}") | |
| if not is_civitai_host(parts.netloc): | |
| lines.append("- normalize: skipped (not a civitai URL)") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| normalized_input = normalize_civitai_input_url(probe_input, api_key=effective_api_key) | |
| normalized_download = normalize_civitai_download_api_url(normalized_input) | |
| if normalized_input != probe_input: | |
| lines.append(f"- normalized: {sanitize_url_for_log(normalized_input)}") | |
| if normalized_download != normalized_input: | |
| lines.append(f"- download api: {sanitize_url_for_log(normalized_download)}") | |
| elif is_civitai_download_api_path(get_civitai_url_parts(normalized_download).path): | |
| lines.append(f"- download api: {sanitize_url_for_log(normalized_download)}") | |
| version_id = extract_civitai_model_version_id(normalized_input) or extract_civitai_model_version_id(probe_input) | |
| if version_id: | |
| lines.append(f"- modelVersionId: {version_id}") | |
| normalized_parts = get_civitai_url_parts(normalized_download) | |
| if is_civitai_download_api_path(normalized_parts.path): | |
| try: | |
| resolved_url = resolve_civitai_download_url(normalized_download, effective_api_key, max_tries=1) | |
| resolved_parts = get_civitai_url_parts(resolved_url) | |
| lines.append(f"- resolve: ok ({resolved_parts.netloc or '-'})") | |
| if resolved_parts.path: | |
| lines.append(f"- resolved path: {resolved_parts.path[:160]}") | |
| except Exception as e: | |
| lines.append(f"- resolve: fail ({type(e).__name__}: {e})") | |
| else: | |
| lines.append("- resolve: skipped (download URL not derivable)") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| def probe_civitai_sidecar(url_input: str, fallback_query: str, api_key: str): | |
| effective_api_key = get_effective_probe_api_key(api_key) | |
| key_source = "input" if str(api_key or "").strip() else ("env" if effective_api_key else "none") | |
| probe_input = get_probe_url_candidate(url_input, fallback_query) | |
| print(f"Civitai Sidecar Probe: input={sanitize_url_for_log(probe_input)} key_source={key_source}") | |
| lines = ["### Civitai Sidecar Probe", f"- key source: {key_source}"] | |
| if not probe_input: | |
| lines.append("- input: skipped (no URL found in Download URL(s) or Query)") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| parts = get_civitai_url_parts(probe_input) | |
| mode = get_civitai_probe_mode(probe_input) | |
| lines.append(f"- input: {sanitize_url_for_log(probe_input)}") | |
| if parts.netloc: | |
| lines.append(f"- input host: {parts.netloc}") | |
| lines.append(f"- mode: {mode}") | |
| if not is_civitai_host(parts.netloc): | |
| lines.append("- sidecar: skipped (not a civitai URL)") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| normalized_input = normalize_civitai_input_url(probe_input, api_key=effective_api_key) | |
| normalized_download = normalize_civitai_download_api_url(normalized_input) | |
| if normalized_input != probe_input: | |
| lines.append(f"- normalized: {sanitize_url_for_log(normalized_input)}") | |
| if normalized_download != normalized_input: | |
| lines.append(f"- download api: {sanitize_url_for_log(normalized_download)}") | |
| elif is_civitai_download_api_path(get_civitai_url_parts(normalized_download).path): | |
| lines.append(f"- download api: {sanitize_url_for_log(normalized_download)}") | |
| if not is_civitai_download_api_path(get_civitai_url_parts(normalized_download).path): | |
| lines.append("- json: skipped (download URL not derivable)") | |
| lines.append("- html: skipped (download URL not derivable)") | |
| lines.append("- preview png: skipped (download URL not derivable)") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| temp_dir = create_run_temp_dir() | |
| try: | |
| json_data, html_text, image_path = get_civitai_json(normalized_download, True, normalized_download, effective_api_key, temp_dir=temp_dir) | |
| json_ok = isinstance(json_data, dict) and bool(json_data) | |
| html_ok = bool(html_text) | |
| image_ok = bool(image_path and Path(image_path).exists()) | |
| lines.append(f"- json: {'ok' if json_ok else 'fail'}") | |
| if json_ok: | |
| model_id = json_data.get('modelId') | |
| version_id = json_data.get('id') | |
| if model_id: | |
| lines.append(f"- modelId: {model_id}") | |
| if version_id: | |
| lines.append(f"- modelVersionId: {version_id}") | |
| lines.append(f"- html: {'ok' if html_ok else 'fail'}") | |
| lines.append(f"- preview png: {'ok' if image_ok else 'fail'}") | |
| if image_ok: | |
| lines.append(f"- preview file: {Path(str(image_path)).name}") | |
| except Exception as e: | |
| lines.append(f"- sidecar: fail ({type(e).__name__}: {e})") | |
| finally: | |
| cleanup_run_temp_dir(temp_dir) | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| def probe_civitai_api(query: str, api_key: str): | |
| page_origin = get_civitai_display_origin() | |
| effective_api_key = get_effective_probe_api_key(api_key) | |
| key_source = "input" if str(api_key or "").strip() else ("env" if effective_api_key else "none") | |
| reset_civitai_key_status(effective_api_key, source="probe") | |
| probe_query = str(query or "lora").strip() or "lora" | |
| print(f"Civitai Probe: query={probe_query!r} key_source={key_source}") | |
| lines = ["### Civitai Probe", f"- page origin: {page_origin}", f"- query: {probe_query}", f"- key source: {key_source}"] | |
| session = create_retry_session(total=4, backoff_factor=0.8) | |
| api_origin = resolve_civitai_api_origin(session) | |
| lines.insert(1, f"- api origin: {api_origin}") | |
| anon_ok = False | |
| try: | |
| page_resp = session.get(f"{page_origin}/models", params={"query": probe_query}, headers=get_civitai_headers(""), timeout=(7.0, 20.0)) | |
| lines.append(f"- anonymous page: {'ok' if page_resp.ok else 'fail'} ({page_resp.status_code})") | |
| if page_resp.ok: | |
| first_path = extract_first_model_path(page_resp.text) | |
| if first_path: | |
| lines.append(f"- anonymous first model: {first_path}") | |
| anon_resp = request_civitai_api( | |
| session, | |
| "/models", | |
| params={"query": probe_query, "limit": 1, "sort": "Newest", "period": "AllTime", "nsfw": "true"}, | |
| timeout=(7.0, 20.0), | |
| label='Civitai anonymous probe models', | |
| source='probe-anon', | |
| ) | |
| anon_status = getattr(anon_resp, "status_code", "-") | |
| lines.append(f"- anonymous api: {'ok' if anon_resp and anon_resp.ok else 'fail'} ({anon_status})") | |
| anon_ok = bool(anon_resp is not None and anon_resp.ok) | |
| except Exception as e: | |
| lines.append(f"- anonymous probe: fail ({type(e).__name__}: {e})") | |
| tags_probe = probe_json_api_step( | |
| session, | |
| "/tags", | |
| params={"limit": 1}, | |
| timeout=(7.0, 20.0), | |
| label='Civitai probe tags', | |
| source='probe-tags', | |
| ) | |
| if tags_probe["error"]: | |
| lines.append(f"- tags api: fail ({tags_probe['error']})") | |
| else: | |
| tags_state = 'ok' if tags_probe["ok"] and tags_probe["json_ok"] else 'fail' | |
| lines.append(f"- tags api: {tags_state} ({tags_probe['status']}, {tags_probe['host'] or '-'})") | |
| creators_probe = probe_json_api_step( | |
| session, | |
| "/creators", | |
| api_key=effective_api_key, | |
| params={"limit": 1, "query": probe_query}, | |
| timeout=(7.0, 20.0), | |
| label='Civitai probe creators', | |
| source='probe-creators', | |
| preferred_origin=CIVITAI_RED_ORIGIN, | |
| non_json_fallback_origin=CIVITAI_DEFAULT_ORIGIN, | |
| ) | |
| if creators_probe["error"]: | |
| lines.append(f"- creators api: fail ({creators_probe['error']})") | |
| else: | |
| creators_state = 'ok' if creators_probe["ok"] and creators_probe["json_ok"] else 'fail' | |
| creator_host = canonicalize_civitai_host(creators_probe["host"]) if creators_probe["host"] else '' | |
| lines.append(f"- creators api: {creators_state} ({creators_probe['status']}, {creators_probe['host'] or '-'})") | |
| if creator_host == 'civitai.com': | |
| lines.append("- creators fallback: ok (civitai.com)") | |
| elif creator_host == 'civitai.red': | |
| lines.append("- creators fallback: not-needed (civitai.red json)") | |
| parsed_keys = parse_civitai_api_keys(effective_api_key) | |
| if not parsed_keys: | |
| lines.append("- auth resolve: skipped (no Civitai key)") | |
| lines.append("- direct download: skipped (no Civitai key)") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |
| try: | |
| reset_civitai_key_status(effective_api_key, source="probe-auth") | |
| auth_resp = request_civitai_api( | |
| session, | |
| "/models", | |
| api_key=effective_api_key, | |
| params={"query": probe_query, "limit": 1, "sort": "Newest", "period": "AllTime", "nsfw": "true"}, | |
| timeout=(7.0, 20.0), | |
| label='Civitai probe models', | |
| source='probe-auth', | |
| ) | |
| auth_status = getattr(auth_resp, "status_code", "-") | |
| lines.append(f"- auth api: {'ok' if auth_resp and auth_resp.ok else 'fail'} ({auth_status})") | |
| if auth_resp is not None and auth_resp.ok: | |
| payload = get_civitai_response_json(auth_resp, default={}) if auth_resp.content else {} | |
| dl_url, model, version, file_info = extract_probe_first_download_url(payload) | |
| if dl_url: | |
| lines.append(f"- auth first file: {model.get('name', '')} / {version.get('name', '')} / {round(float(file_info.get('sizeKB', 0.0) or 0.0) / 1000.0, 2)}MB") | |
| model_url = build_civitai_model_url(model.get('id'), version.get('id')) | |
| normalized_model_download = normalize_civitai_input_url(model_url, api_key=effective_api_key) | |
| normalized_auth_download = normalize_civitai_download_api_url(dl_url) | |
| model_download_parts = get_civitai_url_parts(normalized_model_download) | |
| auth_download_parts = get_civitai_url_parts(normalized_auth_download) | |
| same_download_path = bool( | |
| is_civitai_download_api_path(model_download_parts.path) | |
| and is_civitai_download_api_path(auth_download_parts.path) | |
| and model_download_parts.path == auth_download_parts.path | |
| ) | |
| lines.append( | |
| f"- model page normalize: {'ok' if same_download_path else 'mismatch'} ({sanitize_url_for_log(normalized_model_download)})" | |
| ) | |
| try: | |
| resolved_url = resolve_civitai_download_url(normalized_auth_download, effective_api_key, max_tries=1) | |
| resolved_host = urllib.parse.urlparse(resolved_url).netloc | |
| lines.append(f"- auth resolve: ok ({resolved_host})") | |
| download_probe = probe_download_url_step(session, resolved_url, api_key=effective_api_key, timeout=(7.0, 20.0)) | |
| if download_probe["error"]: | |
| lines.append(f"- direct download: fail ({download_probe['error']})") | |
| else: | |
| lines.append( | |
| f"- direct download: {'ok' if download_probe['ok'] else 'fail'} ({download_probe['status']}, {download_probe['host'] or '-'})" | |
| ) | |
| except Exception as e: | |
| lines.append(f"- auth resolve: fail ({type(e).__name__}: {e})") | |
| lines.append("- direct download: skipped (resolve failed)") | |
| else: | |
| lines.append("- auth resolve: skipped (no downloadUrl in first result)") | |
| lines.append("- direct download: skipped (no downloadUrl in first result)") | |
| except Exception as e: | |
| lines.append(f"- auth probe: fail ({type(e).__name__}: {e})") | |
| if not anon_ok: | |
| lines.append("- note: anonymous probe failed, so auth result may not represent general site health") | |
| probe_text = "\n".join(lines) | |
| print(probe_text) | |
| return gr.update(value=probe_text, visible=True), gr.update(value=format_civitai_key_status_md(effective_api_key)) | |