diff --git "a/civitai_to_hf.py" "b/civitai_to_hf.py" --- "a/civitai_to_hf.py" +++ "b/civitai_to_hf.py" @@ -8,17 +8,73 @@ 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) + 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) +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 +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 +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_GREEN_ORIGIN = "https://civitai.green" +CIVITAI_HOST_ALIASES = frozenset({"civitai.com", "www.civitai.com", "civitai.green", "www.civitai.green"}) +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_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): @@ -43,19 +99,372 @@ def parse_repos(s): 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(parse_urls(s) + parse_repos(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 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), build_run_status_update(session_state), gr.update(value=final_md), remain_update, failed_update, key_status_update, session_state_output(session_state) -def upload_safetensors_to_repo(filename, repo_id, repo_type, is_private, progress=gr.Progress(track_tqdm=True)): +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 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)): output_filename = Path(filename).name - hf_token = get_token() - api = HfApi(token=hf_token) + if hf_token is None: hf_token = get_token() + if api is None: api = HfApi(token=hf_token) try: - if not is_repo_exists(repo_id, repo_type): api.create_repo(repo_id=repo_id, repo_type=repo_type, token=hf_token, private=is_private) + 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}") - 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) + with suppress_hf_hub_progress_bars(): + 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), action=f'upload_file {repo_id}:{output_filename}') progress(1, desc="Uploaded.") url = hf_hub_url(repo_id=repo_id, repo_type=repo_type, filename=output_filename) except Exception as e: @@ -67,13 +476,60 @@ def upload_safetensors_to_repo(filename, repo_id, repo_type, is_private, progres return url -def get_repo_hashes(repo_id: str, repo_type: str="model"): +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() - api = HfApi(token=hf_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 not api.repo_exists(repo_id=repo_id, repo_type=repo_type, token=hf_token): return hashes - tree = api.list_repo_tree(repo_id=repo_id, repo_type=repo_type, token=hf_token) + 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"]) @@ -83,33 +539,131 @@ def get_repo_hashes(repo_id: str, repo_type: str="model"): 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 str(netloc or "").strip().lower() in CIVITAI_HOST_ALIASES + +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"{CIVITAI_DEFAULT_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\.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 parts.netloc.lower().endswith("civitai.com") 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 - user_agent = get_user_agent() - headers = {'User-Agent': user_agent, 'content-type': 'application/json'} - if api_key: headers['Authorization'] = f'Bearer {{{api_key}}}' base_url = 'https://civitai.com/api/v1/model-versions/' params = {} - session = requests.Session() - retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]) - session.mount("https://", HTTPAdapter(max_retries=retries)) + 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 = base_url + m.group(1) qs = urllib.parse.parse_qs(m.group(2)) if "type" not in qs.keys(): qs["type"] = ["Model"] try: - r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 15)) + r = civitai_get(session, url, api_key=api_key, params=params, timeout=(5.0, 15), label='Civitai sha256') if not r.ok: return None json = dict(r.json()) 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 - hash = d["hashes"]["SHA256"].lower() + 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: @@ -130,14 +684,14 @@ def is_same_file(filename: str, cmp_sha256: str, cmp_size: int): else: return False -def get_safe_filename(filename, repo_id, repo_type): - hf_token = get_token() - api = HfApi(token=hf_token) +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 api.file_exists(repo_id=repo_id, filename=Path(new_filename).name, repo_type=repo_type, token=hf_token): - infos = api.get_paths_info(repo_id=repo_id, paths=[Path(new_filename).name], repo_type=repo_type, token=hf_token) + 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 "" @@ -148,23 +702,25 @@ def get_safe_filename(filename, repo_id, repo_type): 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 occured when renaming {filename}. {e}") + print(f"Error occurred when renaming {filename}. {e}") finally: return new_filename -def download_file(dl_url, civitai_key, progress=gr.Progress(track_tqdm=True)): - download_dir = TEMP_DIR +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, dl_url, civitai_key) + output_filename = get_download_file(download_dir, resolved_url, civitai_key) return output_filename -def save_civitai_info(dl_url, filename, civitai_key="", progress=gr.Progress(track_tqdm=True)): - json_str, html_str, image_path = get_civitai_json(dl_url, True, filename, civitai_key) +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(TEMP_DIR, Path(filename).stem + ".json")) - html_path = str(Path(TEMP_DIR, Path(filename).stem + ".html")) + 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) @@ -176,313 +732,1520 @@ def save_civitai_info(dl_url, filename, civitai_key="", progress=gr.Progress(tra return "", "", "" -def upload_info_to_repo(dl_url, filename, repo_id, repo_type, is_private, civitai_key="", progress=gr.Progress(track_tqdm=True)): +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 - api.upload_file(path_or_fileobj=filename, path_in_repo=Path(filename).name, repo_type=repo_type, revision="main", token=hf_token, repo_id=repo_id) + 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() - hf_token = get_token() - api = HfApi(token=hf_token) + if hf_token is None: hf_token = get_token() + if api is None: api = HfApi(token=hf_token) try: - if not is_repo_exists(repo_id, repo_type): api.create_repo(repo_id=repo_id, repo_type=repo_type, token=hf_token, private=is_private) + 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) + 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 - else: upload_file(api, json_path, repo_id, repo_type, hf_token) + 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 + 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, 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, 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, progress=gr.Progress(track_tqdm=True)): - if hf_token: set_token(hf_token) - else: set_token(os.getenv("HF_TOKEN", False)) # default huggingface write token - if not civitai_key: civitai_key = os.environ.get("CIVITAI_API_KEY") # default Civitai API key - if not newrepo_id: newrepo_id = os.environ.get("HF_REPO") # default repo to upload - if not get_token() or not civitai_key: raise gr.Error("HF write token and Civitai API key is required.") - if not urls: urls = [] - dl_urls = parse_urls(dl_url) +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, 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() - hashes = set(get_repo_hashes(newrepo_id, repo_type)) - try: - md = f'### Your repo: [{newrepo_id}]({"https://huggingface.co/datasets/" if repo_type == "dataset" else "https://huggingface.co/"}{newrepo_id})\n' - yield gr.update(value=urls, choices=urls), gr.update(value=md), gr.update(value="\n".join(remain_urls)) - for u in dl_urls: - if get_civitai_sha256(u, civitai_key) in hashes: - print(f"{u} is already exitsts. skipping.") - remain_urls.remove(u) - md += f"- Skipped [{str(u)}]({str(u)})\n" - yield gr.update(value=urls, choices=urls), gr.update(value=md), gr.update(value="\n".join(remain_urls)) - continue - file = download_file(u, civitai_key) - if not Path(file).exists() or not Path(file).is_file(): continue - if is_rename: file = get_safe_filename(file, newrepo_id, repo_type) - url = upload_safetensors_to_repo(file, newrepo_id, repo_type, is_private) - if url: - if is_info: upload_info_to_repo(u, file, newrepo_id, repo_type, is_private, civitai_key) - urls.append(url) - remain_urls.remove(u) - md += f"- Uploaded [{str(u)}]({str(u)})\n" - yield gr.update(value=urls, choices=urls), gr.update(value=md), gr.update(value="\n".join(remain_urls)) - dp_repos = parse_repos(dl_url) - for r in dp_repos: - url = duplicate_hf_repo(r, newrepo_id, "model", repo_type, is_private, HF_SUBFOLDER_NAME[1]) - if url: urls.append(url) - yield gr.update(value=urls, choices=urls), gr.update(value=md), gr.update(value="\n".join(remain_urls), visible=False) + failed_urls = [] + 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_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(), active_run_id=run_id, cancel_requested=False, last_error="", last_failure_summary="", run_started_at=time.time(), run_elapsed_sec=0.0) + 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'}") + 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: + 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}") + current_file = download_file(u, civitai_key, temp_dir=run_temp_dir, progress=progress) + file_ok, file_detail = summarize_downloaded_file(current_file) + if run_mode == "smoke": + smoke_lines.append(smoke_stage_line("Download verify", "ok" if file_ok else "fail", file_detail)) + if not file_ok: + 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 + 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}") + 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) + 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 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: + 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}") + 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: - gr.Info(f"Error occured: {e}") - yield gr.update(value=urls, choices=urls), gr.update(value=md), gr.update(value="\n".join(remain_urls), visible=True) + 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) + final_stage = "Cancelled" if cancelled else ("Failed" if error_message 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) + log_line("cleanup", f"finished {run_mode} run stage={final_stage.lower()} remaining={len(remain_urls)} failed={len(failed_urls)}") + session_state_update(session_state, current_run_mode="idle", 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_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_error=error_message, active_run_id="", cancel_requested=False) gc.collect() + md = build_run_markdown(repo_header if repo_header else "", result_lines, smoke_lines) + if cancelled: + md = build_run_markdown(repo_header if repo_header else "", result_lines + ["- Cancelled by user."], smoke_lines) + elif error_message and not result_lines: + md = build_run_markdown(repo_header if repo_header else "", [f"- Failed ({error_message})"], smoke_lines) + set_stage_progress(progress, 1, 1, "Cancelled" if cancelled else "Done") + 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)) -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 = ["Aura Flow", "Chroma", "CogVideoX", "Flux.1 D", "Flux.1 S", "Flux.1 Kontext", "HiDream", "Hunyuan 1", "Hunyuan Video", "Illustrious", "Imagen4", "Kolors", - "LTXV", "Lumina", "Mochi", "Nano Banana", "NoobAI", "ODOR", "OpenAI", "Other", "PixArt E", "PixArt a", "Playground v2", "Pony", "Qwen", - "SD 1.4", "SD 1.5", "SD 1.5 Hyper", "SD 1.5 LCM", "SD 2.0", "SD 2.0 768", "SD 2.1", "SD 2.1 768", "SD 2.1 Unclip", - "SD 3", "SD 3.5", "SD 3.5 Large", "SD 3.5 Large Turbo", "SD 3.5 Medium", "SDXL 0.9", "SDXL 1.0", - "SDXL 1.0 LCM", "SDXL Distilled", "SDXL Hyper", "SDXL Lightning", "SDXL Turbo", "SVD", "SVD XT", "Seedream", "Stable Cascade", - "Veo 3", "Wan Video", "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", "Wan Video 2.5 T2V"] + +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_url = 'https://civitai.com/api/v1/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 = civitai_get(session, base_url, api_key=api_key, params=params, timeout=(7.0, 30), label=f'Civitai base model refresh sort={sort} page={page_index}') + else: + r = civitai_get(session, next_url, api_key=api_key, timeout=(7.0, 30), label=f'Civitai base model refresh sort={sort} page={page_index}') + if not r.ok: + print(f"Failed to refresh Civitai base models. sort={sort} page={page_index} status={r.status_code}") + break + data = r.json() + 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_green(): + return begin_probe_feedback("civitai.green 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 = "", progress=gr.Progress(track_tqdm=True)): - user_agent = get_user_agent() - headers = {'User-Agent': user_agent, 'content-type': 'application/json'} - if api_key: headers['Authorization'] = f'Bearer {{{api_key}}}' - base_url = 'https://civitai.com/api/v1/models' - params = {'sort': CIVITAI_SORT[0] if sort in CIVITAI_SORT_EXT else sort, 'period': period, 'limit': int(limit), 'nsfw': 'true'} - if len(types) != 0: params["types"] = types - if query: params["query"] = query - if tag: params["tag"] = tag - if user: params["username"] = user - if page != 0: params["page"] = int(page) - session = requests.Session() - retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]) - session.mount("https://", HTTPAdapter(max_retries=retries)) + filetype: list[str] = [], api_key: str = "", base_origin: str = CIVITAI_DEFAULT_ORIGIN, progress=gr.Progress(track_tqdm=False)): + base_url = f'{base_origin}/api/v1/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 = session.get(base_url, params=params | {'page': 1}, headers=headers, stream=True, timeout=(7.0, 30)) + r = civitai_get(session, base_url, 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 = r.json() next_url = json['metadata']['nextPage'] if 'metadata' in json and 'nextPage' in json['metadata'] else None i = 2 - while(next_url is not None): + while next_url is not None: progress(0, desc=f"Searching page {i}...") print(f"Searching page {i}...") - r = session.get(next_url, headers=headers, stream=True, timeout=(7.0, 30)) + r = civitai_get(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 = r.json() next_url = json['metadata']['nextPage'] if 'metadata' in json and 'nextPage' in json['metadata'] else None - else: next_url = 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 = session.get(base_url, params=params, headers=headers, stream=True, timeout=(7.0, 30)) + r = civitai_get(session, base_url, 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 = base_origin.rstrip('/') for r in rs: - if not r.ok: continue + if not r.ok: + continue json = r.json() - if 'items' not in json: continue + if 'items' not in json: + continue for j in json['items']: - for model in j['modelVersions']: - item = {} - if len(allow_model) != 0 and model['baseModel'] not in set(allow_model): continue - item['name'] = j['name'] - item['creator'] = j['creator']['username'] if 'creator' in j.keys() and 'username' in j['creator'].keys() else "" - item['tags'] = j['tags'] if 'tags' in j.keys() else [] - item['model_name'] = model['name'] if 'name' in model.keys() else "" - item['base_model'] = model['baseModel'] if 'baseModel' in model.keys() else "" - item['description'] = model['description'] if 'description' in model.keys() else "" - item['md'] = "" - if 'images' in model.keys() and len(model["images"]) != 0: - item['img_url'] = model["images"][0]["url"] - item['md'] += f'thumbnail
' - else: item['img_url'] = "/home/user/app/null.png" - item['md'] += f'''Model URL: [https://civitai.com/models/{j["id"]}](https://civitai.com/models/{j["id"]})
Model Name: {item["name"]}
- Creator: {item["creator"]}
Tags: {", ".join(item["tags"])}
Base Model: {item["base_model"]}
Description: {item["description"]}''' - if 'files' in model.keys(): - for f in model['files']: - i = item.copy() - i['dl_url'] = f['downloadUrl'] - i['size_kb'] = f['sizeKB'] - if len(filetype) != 0 and f['type'] not in set(filetype): continue - items.append(i) + for model in j.get('modelVersions', []): + if len(allow_model) != 0 and model.get('baseModel', '') not in set(allow_model): + 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: - item['dl_url'] = model['downloadUrl'] + base_item['img_url'] = NULL_IMAGE_PATH + model_url = f"{origin}/models/{j.get('id', '')}" + if model.get('id') is not None: + model_url += f"?modelVersionId={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: + for f in files: + item = base_item.copy() + item['dl_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): + continue + items.append(item) + else: + item = base_item.copy() + item['dl_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.), reverse=True) - elif sort == "Size (from smallest)": items = sorted(items, key=lambda x: x.get('size_kb', 0.)) + 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)) 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="", gallery=[], state={}, progress=gr.Progress(track_tqdm=True)): + 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_gallery", []) set_state(state, "civitai_last_results", civitai_last_results) - results_info = "No item found." - items = search_on_civitai(query, types, base_model, int(limit), sort, period, tag, user, int(page), filetype, api_key) - if not items: return gr.update(choices=[("", "")], value=[], visible=True),\ - gr.update(value="", visible=False), gr.update(), gr.update(), gr.update(), gr.update(), results_info, state + 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 = [] - gallery = [] + ordered_items = [] for item in items: - base_model_name = "Pony🐴" if item['base_model'] == "Pony" else item['base_model'] - name = f"{item['name']} (for {base_model_name} / By: {item['creator']}) ({round(item['size_kb'] / 1000., 2)}MB)" if "size_kb" in item.keys() else f"{item['name']} (for {base_model_name} / By: {item['creator']})" - value = item['dl_url'] - choices.append((name, value)) - gallery.append((item['img_url'], name)) - civitai_last_results[value] = item - if len(choices) >= 1: results_info = f"{int(len(choices))} items found." - else: choices = [("", "")] - md = "" + 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_gallery", gallery) set_state(state, "civitai_last_results", civitai_last_results) - return gr.update(choices=choices, value=[], visible=True), gr.update(value=md, visible=True),\ - gr.update(), gr.update(), gr.update(value=gallery), gr.update(choices=choices, value=[]), results_info, state + 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 get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str="", api_key=""): - if not image_baseurl: image_baseurl = dl_url +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 - user_agent = get_user_agent() - headers = {'User-Agent': user_agent, 'content-type': 'application/json'} - if api_key: headers['Authorization'] = f'Bearer {{{api_key}}}' + if "https://civitai.com/api/download/models/" not in dl_url: + return default base_url = 'https://civitai.com/api/v1/model-versions/' params = {} - session = requests.Session() - retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]) - session.mount("https://", HTTPAdapter(max_retries=retries)) - model_id = re.sub('https://civitai.com/api/download/models/(\\d+)(?:.+)?', '\\1', dl_url) + 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 = base_url + model_id - #url = base_url + str(dl_url.split("/")[-1]) try: - r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 60)) - if not r.ok: return default + reset_civitai_key_status(api_key, source='info') + r = civitai_get(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(r.json()).copy() html = "" image = "" if "modelId" in json.keys(): - url = f"https://civitai.com/models/{json['modelId']}" - r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 60)) - if not r.ok: return json, html, image + 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 = f"https://civitai.com/models/{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 = json["images"][0]["url"] - r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 60)) - if not r.ok: return json, html, image - image_temp = str(Path(TEMP_DIR, "image" + Path(url.split("/")[-1]).suffix)) - image = str(Path(TEMP_DIR, Path(image_baseurl.split("/")[-1]).stem + ".png")) - with open(image_temp, 'wb') as f: - f.write(r.content) - Image.open(image_temp).convert('RGBA').save(image) + 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 get_civitai_tag(): - default = [""] - user_agent = get_user_agent() - headers = {'User-Agent': user_agent, 'content-type': 'application/json'} - base_url = 'https://civitai.com/api/v1/tags' - params = {'limit': 200} - session = requests.Session() - retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]) - session.mount("https://", HTTPAdapter(max_retries=retries)) - url = base_url - try: - r = session.get(url, params=params, headers=headers, stream=True, timeout=(7.0, 15)) - if not r.ok: return default +def _load_civitai_choice_list(base_url: 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=""): + 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 = civitai_get(session, base_url, api_key=api_key, params=params, timeout=timeout, label=label, source=source) + 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(r.json()).copy() - if "items" not in j.keys(): return default - items = [] - for item in j["items"]: - items.append([str(item.get("name", "")), int(item.get("modelCount", 0))]) - df = pd.DataFrame(items) - df.sort_values(1, ascending=False) - tags = df.values.tolist() - tags = [""] + [l[0] for l in tags] - return tags + 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: - print(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( + base_url='https://civitai.com/api/v1/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', + ) + + +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( + base_url='https://civitai.com/api/v1/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): - json = {} - if "http" not in "".join(results) or len(results) == 0: return gr.update(value="", visible=True), gr.update(value=json, visible=False), state - result = get_state(state, "civitai_last_results") - last_selects = get_state(state, "civitai_last_selects") - selects = list_sub(results, last_selects if last_selects else []) - md = result.get(selects[-1]).get('md', "") if result and isinstance(result, dict) and len(selects) > 0 else "" - set_state(state, "civitai_last_selects", results) - return gr.update(value=md, visible=True), gr.update(value=json, visible=False), state + 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) + 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 + 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 select_civitai_all_item(button_name: str, state: dict): - if button_name not in ["Select All", "Deselect All"]: return gr.update(value=button_name), gr.Update(visible=True) - civitai_last_choices = get_state(state, "civitai_last_choices") - selected = [t[1] for t in civitai_last_choices if t[1] != ""] if button_name == "Select All" else [] - new_button_name = "Select All" if button_name == "Deselect All" else "Deselect All" - return gr.update(value=new_button_name), gr.update(value=selected, choices=civitai_last_choices) +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 update_civitai_selection(evt: gr.SelectData, value: list[str], state: dict): +def get_gallery_event_index(evt) -> int | None: try: - civitai_last_choices = get_state(state, "civitai_last_choices") - selected_index = evt.index - selected = list_uniq([v for v in value if v != ""] + [civitai_last_choices[selected_index][1]]) - return gr.update(value=selected) + 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 gr.update() + 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) + + +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 update_civitai_checkbox(selected: list[str]): - return gr.update(value=selected) +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 from_civitai_checkbox(selected: list[str]): - return gr.update(value=selected) +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: + reset_civitai_key_status(effective_api_key, source="probe-keys") + resp = civitai_get( + session, + f"{CIVITAI_DEFAULT_ORIGIN}/api/v1/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): + origin = CIVITAI_DEFAULT_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"- origin: {origin}", f"- query: {probe_query}", f"- key source: {key_source}"] + session = create_retry_session(total=4, backoff_factor=0.8) + anon_ok = False + try: + page_resp = session.get(f"{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 = session.get( + f"{origin}/api/v1/models", + params={"query": probe_query, "limit": 1, "sort": "Newest", "period": "AllTime", "nsfw": "true"}, + headers=get_civitai_headers(""), + timeout=(7.0, 20.0), + ) + lines.append(f"- anonymous api: {'ok' if anon_resp.ok else 'fail'} ({anon_resp.status_code})") + anon_ok = bool(anon_resp.ok) + except Exception as e: + lines.append(f"- anonymous probe: fail ({type(e).__name__}: {e})") + parsed_keys = parse_civitai_api_keys(effective_api_key) + if not parsed_keys: + lines.append("- auth resolve: 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 = civitai_get( + session, + f"{origin}/api/v1/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', + ) + lines.append(f"- auth api: {'ok' if auth_resp.ok else 'fail'} ({auth_resp.status_code})") + if auth_resp.ok: + payload = auth_resp.json() 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") + try: + resolved_url = resolve_civitai_download_url(dl_url, effective_api_key, max_tries=1) + resolved_host = urllib.parse.urlparse(resolved_url).netloc + lines.append(f"- auth resolve: ok ({resolved_host})") + except Exception as e: + lines.append(f"- auth resolve: fail ({type(e).__name__}: {e})") + else: + lines.append("- auth resolve: 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)) +def probe_civitai_green_api(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 = str(query or "").strip() or "lora" + print(f"civitai.green Probe: input={probe_input!r} key_source={key_source}") + lines = ["### civitai.green Probe", f"- input: {probe_input}", f"- key source: {key_source}"] + session = create_retry_session(total=4, backoff_factor=0.8) + try: + parts = get_civitai_url_parts(probe_input) + if is_civitai_host(parts.netloc): + lines.append(f"- url host: {parts.netloc}") + normalized = normalize_civitai_input_url(probe_input, api_key=effective_api_key) + if normalized != probe_input: + lines.append(f"- normalized download url: {normalize_civitai_download_api_url(normalized)}") + if "api/download/models/" in normalized: + lines.append("- mode: direct-download probe") + try: + resolved_url = resolve_civitai_download_url(normalized, effective_api_key, max_tries=1) + resolved_host = urllib.parse.urlparse(resolved_url).netloc + lines.append(f"- resolve: ok ({resolved_host})") + except Exception as e: + lines.append(f"- resolve: fail ({type(e).__name__}: {e})") + else: + lines.append("- mode: page-only probe") + page_resp = session.get(probe_input, headers=get_civitai_headers(effective_api_key), timeout=(7.0, 20.0)) + lines.append(f"- 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"- first model: {first_path}") + first_download = extract_first_civitai_download_url_from_html(page_resp.text) + lines.append(f"- embedded download url: {'found' if first_download else 'not found'}") + else: + origin = CIVITAI_GREEN_ORIGIN + lines.append(f"- origin: {origin}") + lines.append("- mode: query probe") + page_resp = session.get(f"{origin}/models", params={"query": probe_input}, headers=get_civitai_headers(effective_api_key), timeout=(7.0, 20.0)) + lines.append(f"- 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"- first model: {first_path}") + lines.append("- api: skipped (HTML-only probe)") + except Exception as e: + lines.append(f"- 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))