263 lines
10 KiB
Python
263 lines
10 KiB
Python
# CivitAI API — model search, download management, and API key
|
|
# configuration endpoints.
|
|
|
|
import asyncio
|
|
import os
|
|
import threading
|
|
import time
|
|
import urllib.parse
|
|
from pathlib import Path
|
|
|
|
import aiohttp
|
|
from aiohttp import web
|
|
import server
|
|
|
|
from . import civitai
|
|
from . import config as promptchain_config
|
|
from . import model_settings
|
|
from .api_utils import error_response, parse_json
|
|
from .shared import send_ws
|
|
|
|
routes = server.PromptServer.instance.routes
|
|
|
|
|
|
# ── offline download test mode ────────────────────────────────────
|
|
# Enabled only by the env var (set by the test launcher, never in code), so a
|
|
# real install can't accidentally be in test mode. When on, the download route
|
|
# redirects model-file fetches to the test-mirror route below, which serves
|
|
# fixtures from PROMPTCHAIN_FIXTURES_DIR — no internet, no second process.
|
|
|
|
def _test_downloads_on() -> bool:
|
|
return os.environ.get("PROMPTCHAIN_TEST_DOWNLOADS", "").strip().lower() not in ("", "0", "false")
|
|
|
|
|
|
def _fixtures_dir() -> Path:
|
|
return Path(os.environ.get("PROMPTCHAIN_FIXTURES_DIR", "").strip() or ".")
|
|
|
|
|
|
@routes.get("/promptchain/test-mirror/{name}")
|
|
async def _api_test_mirror(request):
|
|
"""Serve a fixture file by name (test mode only). Fault injection by name:
|
|
`__fail` -> 500, `__stall` -> hang. Missing fixture -> 404 so the download
|
|
fails cleanly. PROMPTCHAIN_MIRROR_THROTTLE (MB/s) animates the progress bar."""
|
|
if not _test_downloads_on():
|
|
return web.Response(status=404, text="test mode off")
|
|
name = os.path.basename(request.match_info["name"])
|
|
if "__fail" in name:
|
|
return web.Response(status=500, text="injected failure")
|
|
path = _fixtures_dir() / name
|
|
if not path.is_file():
|
|
return web.Response(status=404, text=f"fixture not found: {name}")
|
|
|
|
total = path.stat().st_size
|
|
resp = web.StreamResponse(status=200, headers={
|
|
"Content-Type": "application/octet-stream",
|
|
"Content-Length": str(total),
|
|
})
|
|
await resp.prepare(request)
|
|
if "__stall" in name:
|
|
await asyncio.sleep(600)
|
|
return resp
|
|
|
|
try:
|
|
bps = int(float(os.environ.get("PROMPTCHAIN_MIRROR_THROTTLE", "0")) * 1024 * 1024)
|
|
except ValueError:
|
|
bps = 0
|
|
sent = 0
|
|
start = time.monotonic()
|
|
with open(path, "rb") as f:
|
|
while True:
|
|
buf = f.read(1024 * 1024)
|
|
if not buf:
|
|
break
|
|
try:
|
|
await resp.write(buf)
|
|
except (ConnectionResetError, ConnectionError):
|
|
break
|
|
sent += len(buf)
|
|
if bps:
|
|
delay = start + sent / bps - time.monotonic()
|
|
if delay > 0:
|
|
await asyncio.sleep(delay)
|
|
await resp.write_eof()
|
|
return resp
|
|
|
|
|
|
@routes.get("/promptchain/civitai/search")
|
|
async def _api_civitai_search(request):
|
|
query = request.query.get("q", "").strip()
|
|
if not query:
|
|
return error_response("missing q parameter")
|
|
try:
|
|
limit = min(int(request.query.get("limit", "10")), 20)
|
|
except (ValueError, TypeError):
|
|
limit = 10
|
|
cursor = request.query.get("cursor", "")
|
|
data = await civitai.search_models(query, limit=limit, cursor=cursor)
|
|
return web.json_response(data)
|
|
|
|
|
|
def _installed_version_ids_for_model(model_id: int) -> list[int]:
|
|
"""Walk the in-memory config index once, collecting civitai_version_ids
|
|
whose config points at `model_id`. The index is already deduped by
|
|
config identity so we skip configs we've already counted."""
|
|
seen_configs: set[int] = set()
|
|
vids: list[int] = []
|
|
index = model_settings._get_config_index()
|
|
for cfg in index.get("by_quick_hash", {}).values():
|
|
if id(cfg) in seen_configs:
|
|
continue
|
|
seen_configs.add(id(cfg))
|
|
try:
|
|
mid = cfg.get("civitai_model_id")
|
|
if mid is None or int(mid) != model_id:
|
|
continue
|
|
vid = cfg.get("civitai_version_id")
|
|
if vid is not None:
|
|
vids.append(int(vid))
|
|
except (TypeError, ValueError):
|
|
continue
|
|
return vids
|
|
|
|
|
|
@routes.get("/promptchain/civitai/model-versions")
|
|
async def _api_civitai_model_versions(request):
|
|
"""List all versions of a CivitAI model plus which civitai_version_ids
|
|
the user already has installed locally. `?force=1` bypasses the
|
|
disk-backed TTL cache so the user can refresh on demand."""
|
|
try:
|
|
model_id = int(request.query.get("model_id", "").strip())
|
|
except (ValueError, TypeError):
|
|
return error_response("missing or invalid model_id")
|
|
|
|
force = request.query.get("force", "").strip() in ("1", "true")
|
|
versions = await civitai.fetch_model_versions(model_id, force=force)
|
|
entry = civitai.get_cached_versions_entry(model_id)
|
|
|
|
return web.json_response({
|
|
"versions": versions,
|
|
"installed_version_ids": _installed_version_ids_for_model(model_id),
|
|
"fetched_at": entry.get("fetched_at") if entry else None,
|
|
"deleted": bool(entry.get("deleted")) if entry else False,
|
|
})
|
|
|
|
|
|
@routes.post("/promptchain/civitai/file-size")
|
|
async def _api_civitai_file_size(request):
|
|
"""Best-effort byte size for a download URL. Catalog presets carry no size,
|
|
so the download modal probes Content-Length to display it before download."""
|
|
data, err = await parse_json(request)
|
|
if err:
|
|
return err
|
|
url = (data.get("url") or "").strip()
|
|
size = await asyncio.to_thread(civitai.probe_file_size, url) if url else 0
|
|
return web.json_response({"size_bytes": size})
|
|
|
|
|
|
@routes.post("/promptchain/civitai/download")
|
|
async def _api_civitai_download(request):
|
|
"""Start downloading a model file. Supports CivitAI and HuggingFace URLs."""
|
|
data, err = await parse_json(request)
|
|
if err: return err
|
|
url = data.get("url", "").strip()
|
|
filename = data.get("filename", "").strip()
|
|
folder_type = data.get("folder", "checkpoints").strip()
|
|
# Optional — when present, the versions-cache entry for this model
|
|
# is invalidated on successful download so the next panel open sees
|
|
# the newly-installed version instead of continuing to advertise it.
|
|
civitai_model_id = data.get("civitai_model_id")
|
|
|
|
# Test mode: pull the file from our own mirror route (keyed by filename)
|
|
# instead of the real CDN, so we only need the filename — the real url may
|
|
# be absent (e.g. catalog entries with no download source). request.host
|
|
# gives whatever host:port the client reached us on, so no port to
|
|
# configure; a localhost url also skips the civitai.com key requirement.
|
|
if _test_downloads_on():
|
|
if not filename:
|
|
return error_response("missing filename")
|
|
safe = urllib.parse.quote(filename.rsplit("/", 1)[-1])
|
|
url = f"http://{request.host}/promptchain/test-mirror/{safe}"
|
|
elif not url or not filename:
|
|
return error_response("missing url or filename")
|
|
|
|
state = civitai.get_download_state()
|
|
if state["active"]:
|
|
return error_response("download already in progress", 409)
|
|
|
|
def on_progress(downloaded, total):
|
|
progress = (downloaded / total * 100) if total else 0
|
|
send_ws("promptchain_download_progress", {
|
|
"filename": filename,
|
|
"downloaded": downloaded,
|
|
"total": total,
|
|
"progress": round(progress, 1),
|
|
})
|
|
|
|
def run_download():
|
|
try:
|
|
civitai.download_model(url, filename, on_progress=on_progress, folder_type=folder_type)
|
|
if civitai_model_id is not None:
|
|
try:
|
|
civitai.invalidate_versions_cache(int(civitai_model_id))
|
|
except (TypeError, ValueError):
|
|
pass
|
|
send_ws("promptchain_download_done", {"filename": filename, "status": "completed"})
|
|
except Exception as e:
|
|
send_ws("promptchain_download_done", {"filename": filename, "status": "failed", "error": str(e)})
|
|
|
|
threading.Thread(target=run_download, daemon=True).start()
|
|
return web.json_response({"status": "started", "filename": filename})
|
|
|
|
|
|
@routes.get("/promptchain/civitai/download-status")
|
|
async def _api_download_status(request):
|
|
return web.json_response(civitai.get_download_state())
|
|
|
|
|
|
@routes.post("/promptchain/civitai/download-cancel")
|
|
async def _api_download_cancel(request):
|
|
civitai.cancel_download()
|
|
return web.json_response({"status": "ok"})
|
|
|
|
|
|
@routes.post("/promptchain/civitai/api-key")
|
|
async def _api_civitai_set_key(request):
|
|
"""Save CivitAI API key to config."""
|
|
data, err = await parse_json(request)
|
|
if err: return err
|
|
key = data.get("key", "").strip()
|
|
config = promptchain_config.load()
|
|
config["civitai_api_key"] = key
|
|
promptchain_config.save(config)
|
|
return web.json_response({"status": "ok"})
|
|
|
|
|
|
@routes.get("/promptchain/civitai/api-key")
|
|
async def _api_civitai_get_key(request):
|
|
"""Check if CivitAI API key is configured."""
|
|
key = civitai._get_api_key()
|
|
return web.json_response({"has_key": bool(key)})
|
|
|
|
|
|
@routes.post("/promptchain/civitai/validate-key")
|
|
async def _api_civitai_validate_key(request):
|
|
"""Validate a CivitAI API key by hitting /api/v1/me.
|
|
If no key in payload, validates the stored/env key instead."""
|
|
data, err = await parse_json(request)
|
|
if err: return err
|
|
key = data.get("key", "").strip()
|
|
if not key:
|
|
key = civitai._get_api_key() or ""
|
|
if not key:
|
|
return web.json_response({"valid": False, "error": "No key provided"})
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
headers = {"Authorization": f"Bearer {key}"}
|
|
async with session.get("https://civitai.com/api/v1/me", headers=headers, timeout=civitai._TIMEOUT) as resp:
|
|
if resp.status == 200:
|
|
user = await resp.json()
|
|
return web.json_response({"valid": True, "username": user.get("username", "")})
|
|
return web.json_response({"valid": False, "error": "Invalid API key"})
|
|
except Exception as e:
|
|
return web.json_response({"valid": False, "error": f"Connection failed: {e}"})
|