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})
" f"Model Name: {base_item['name']}
" f"Version: {base_item['model_name']}
" f"Creator: {base_item['creator']}
" f"Tags: {', '.join(base_item['tags'])}
" f"Base Model: {base_item['base_model']}
" 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))