1280 lines
50 KiB
Python
1280 lines
50 KiB
Python
"""
|
|
CWK Checkpoints Preset Manager — REST / SSE API routes.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import hashlib
|
|
import re
|
|
from typing import Optional
|
|
|
|
import aiohttp
|
|
from aiohttp import web
|
|
|
|
import folder_paths
|
|
from .nodes import (
|
|
load_presets, save_presets, default_preset,
|
|
get_clip_list, get_vae_list,
|
|
resolve_sampler, resolve_scheduler,
|
|
get_last_used_model, save_last_used_model,
|
|
RESOLUTION_PRESETS,
|
|
)
|
|
|
|
# ─── Constants ────────────────────────────────────────────────────────────────
|
|
|
|
CIVITAI_API_BASE = "https://civitai.com/api/v1"
|
|
NODE_DIR = os.path.dirname(__file__)
|
|
THUMBNAILS_CACHE = os.path.join(NODE_DIR, "thumbnails_cache.json")
|
|
HASH_CACHE_FILE = os.path.join(NODE_DIR, "hash_cache.json")
|
|
LOCAL_THUMBS_DIR = os.path.join(NODE_DIR, "local_thumbnails")
|
|
MODEL_META_DIR = os.path.join(NODE_DIR, "model_metadata")
|
|
|
|
_CONCURRENCY = 5
|
|
_USER_AGENT = "CWK-PresetManager/2.0 (ComfyUI custom node)"
|
|
_REQ_TIMEOUT = aiohttp.ClientTimeout(total=20)
|
|
_DL_TIMEOUT = aiohttp.ClientTimeout(total=3600)
|
|
|
|
|
|
# ─── NSFW normalisation ───────────────────────────────────────────────────────
|
|
|
|
def _normalize_nsfw_level(raw) -> int:
|
|
if isinstance(raw, str):
|
|
mapping = {"none": 0, "soft": 1, "mature": 2, "x": 3, "explicit": 3, "xxx": 4}
|
|
return mapping.get(raw.lower().strip(), 0)
|
|
if isinstance(raw, int):
|
|
if raw <= 1: return 0
|
|
if raw <= 2: return 1
|
|
if raw <= 4: return 2
|
|
if raw <= 8: return 3
|
|
return 4
|
|
return 0
|
|
|
|
|
|
# ─── JSON helpers ─────────────────────────────────────────────────────────────
|
|
|
|
def _load_json(path: str) -> dict:
|
|
if os.path.exists(path):
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
pass
|
|
return {}
|
|
|
|
|
|
def _save_json(path: str, data: dict) -> None:
|
|
try:
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=2)
|
|
except Exception as e:
|
|
print(f"[CWK] save_json error ({os.path.basename(path)}): {e}")
|
|
|
|
|
|
# ─── Per-model metadata ───────────────────────────────────────────────────────
|
|
|
|
def _safe_filename(model_name: str) -> str:
|
|
safe = re.sub(r'[<>:"/\\|?*]', '_', model_name)
|
|
return safe + ".json"
|
|
|
|
|
|
def _meta_path(model_name: str) -> str:
|
|
return os.path.join(MODEL_META_DIR, _safe_filename(model_name))
|
|
|
|
|
|
def _load_meta(model_name: str) -> dict:
|
|
path = _meta_path(model_name)
|
|
if os.path.exists(path):
|
|
return _load_json(path)
|
|
legacy = _load_json(THUMBNAILS_CACHE)
|
|
if model_name in legacy:
|
|
data = legacy[model_name]
|
|
_save_meta(model_name, data)
|
|
return data
|
|
return {}
|
|
|
|
|
|
def _save_meta(model_name: str, data: dict) -> None:
|
|
os.makedirs(MODEL_META_DIR, exist_ok=True)
|
|
_save_json(_meta_path(model_name), data)
|
|
|
|
|
|
def _delete_meta(model_name: str) -> None:
|
|
path = _meta_path(model_name)
|
|
if os.path.exists(path):
|
|
try:
|
|
os.remove(path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
# ─── Hash cache ───────────────────────────────────────────────────────────────
|
|
|
|
def _hash_cache_get(path: str) -> Optional[str]:
|
|
entry = _load_json(HASH_CACHE_FILE).get(path)
|
|
if not entry:
|
|
return None
|
|
try:
|
|
st = os.stat(path)
|
|
if st.st_mtime == entry["mtime"] and st.st_size == entry["size"]:
|
|
return entry["sha256"]
|
|
except OSError:
|
|
pass
|
|
return None
|
|
|
|
|
|
def _hash_cache_set(path: str, sha256: str) -> None:
|
|
cache = _load_json(HASH_CACHE_FILE)
|
|
try:
|
|
st = os.stat(path)
|
|
cache[path] = {"mtime": st.st_mtime, "size": st.st_size, "sha256": sha256}
|
|
except OSError:
|
|
cache[path] = {"mtime": 0, "size": 0, "sha256": sha256}
|
|
_save_json(HASH_CACHE_FILE, cache)
|
|
|
|
|
|
# ─── Model file helpers ───────────────────────────────────────────────────────
|
|
|
|
def _all_models() -> list:
|
|
models = []
|
|
seen = set()
|
|
for name in folder_paths.get_filename_list("checkpoints"):
|
|
model_type = "gguf" if name.lower().endswith(".gguf") else "checkpoint"
|
|
models.append({"name": name, "type": model_type})
|
|
seen.add(name)
|
|
try:
|
|
for name in folder_paths.get_filename_list("diffusion_models"):
|
|
if name in seen:
|
|
continue
|
|
model_type = "gguf" if name.lower().endswith(".gguf") else "diffusion_model"
|
|
models.append({"name": name, "type": model_type})
|
|
seen.add(name)
|
|
except Exception:
|
|
pass
|
|
# Also include .gguf files from city96's "unet_gguf" folder type
|
|
try:
|
|
for name in folder_paths.get_filename_list("unet_gguf"):
|
|
if name in seen:
|
|
continue
|
|
models.append({"name": name, "type": "gguf"})
|
|
seen.add(name)
|
|
except Exception:
|
|
pass
|
|
return models
|
|
|
|
|
|
def _full_path(model_name: str) -> Optional[str]:
|
|
# Check all folder types including unet_gguf
|
|
for folder_type in ("checkpoints", "diffusion_models", "unet_gguf"):
|
|
try:
|
|
result = folder_paths.get_full_path(folder_type, model_name)
|
|
if result and os.path.exists(result):
|
|
return result
|
|
except Exception:
|
|
pass
|
|
|
|
basename = os.path.basename(model_name)
|
|
for folder_type in ("checkpoints", "diffusion_models"):
|
|
try:
|
|
roots = folder_paths.get_folder_paths(folder_type)
|
|
except Exception:
|
|
continue
|
|
for root in roots:
|
|
for dirpath, _, filenames in os.walk(root):
|
|
if basename in filenames:
|
|
return os.path.join(dirpath, basename)
|
|
return None
|
|
|
|
def _clean_name(filename: str) -> str:
|
|
return os.path.splitext(os.path.basename(filename))[0]
|
|
|
|
|
|
def _model_save_dir(model_name: str) -> str:
|
|
p = _full_path(model_name)
|
|
if p:
|
|
return os.path.dirname(p)
|
|
dirs = folder_paths.get_folder_paths("checkpoints")
|
|
return dirs[0] if dirs else os.getcwd()
|
|
|
|
def _bust_folder_cache() -> None:
|
|
"""Force ComfyUI to rescan model folders on the next list request."""
|
|
try:
|
|
for folder_type in ("checkpoints", "diffusion_models", "unet_gguf"):
|
|
folder_paths.filename_list_cache.pop(folder_type, None)
|
|
except Exception as e:
|
|
print(f"[CWK] Warning: could not bust folder_paths cache: {e}")
|
|
|
|
|
|
# ─── Hash computation ─────────────────────────────────────────────────────────
|
|
|
|
def _sha256_sync(path: str) -> Optional[str]:
|
|
try:
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
while chunk := f.read(1024 * 1024):
|
|
h.update(chunk)
|
|
result = h.hexdigest().upper()
|
|
_hash_cache_set(path, result)
|
|
return result
|
|
except Exception as e:
|
|
print(f"[CWK] SHA256 error {path}: {e}")
|
|
return None
|
|
|
|
|
|
async def _get_hash(model_name: str, loop) -> Optional[str]:
|
|
path = _full_path(model_name)
|
|
if not path or not os.path.exists(path):
|
|
return None
|
|
cached = _hash_cache_get(path)
|
|
if cached:
|
|
return cached
|
|
print(f"[CWK] Hashing (first time): {model_name}")
|
|
return await loop.run_in_executor(None, _sha256_sync, path)
|
|
|
|
|
|
async def _get_hash_by_path(abs_path: str, loop) -> Optional[str]:
|
|
"""Hash by absolute path — used for freshly downloaded files not yet
|
|
registered in folder_paths."""
|
|
if not abs_path or not os.path.exists(abs_path):
|
|
return None
|
|
cached = _hash_cache_get(abs_path)
|
|
if cached:
|
|
return cached
|
|
print(f"[CWK] Hashing new file: {abs_path}")
|
|
return await loop.run_in_executor(None, _sha256_sync, abs_path)
|
|
|
|
|
|
# ─── CivitAI HTTP helpers ─────────────────────────────────────────────────────
|
|
|
|
class CivitAIAuthError(Exception):
|
|
pass
|
|
|
|
|
|
def _headers(api_key: str) -> dict:
|
|
h = {"User-Agent": _USER_AGENT}
|
|
if api_key:
|
|
h["Authorization"] = f"Bearer {api_key}"
|
|
return h
|
|
|
|
|
|
async def _cget(session: aiohttp.ClientSession, url: str) -> dict:
|
|
async with session.get(url, timeout=_REQ_TIMEOUT) as r:
|
|
if r.status == 403:
|
|
raise CivitAIAuthError("CivitAI returned 403 — API key required or invalid.")
|
|
if r.status == 404:
|
|
return {}
|
|
if r.status == 429:
|
|
raise aiohttp.ClientResponseError(
|
|
r.request_info, r.history, status=429, message="Rate limited by CivitAI"
|
|
)
|
|
r.raise_for_status()
|
|
return await r.json(content_type=None)
|
|
|
|
|
|
def _parse_version(data: dict, filename: str) -> dict:
|
|
if not data:
|
|
return {
|
|
"thumbnail": None,
|
|
"civitai_name": _clean_name(filename),
|
|
"nsfw_level": 0,
|
|
"tags": [],
|
|
}
|
|
|
|
images = data.get("images", [])
|
|
model_blk = data.get("model", {})
|
|
|
|
norm_images = []
|
|
for img in images:
|
|
lvl = _normalize_nsfw_level(img.get("nsfwLevel", 0))
|
|
norm_images.append({**img, "nsfwLevel": lvl})
|
|
|
|
non_video = [img for img in norm_images
|
|
if not img.get("url", "").lower().endswith(".mp4")]
|
|
candidates = non_video if non_video else norm_images
|
|
thumb_img = min(candidates, key=lambda i: i["nsfwLevel"]) if candidates else None
|
|
nsfw_level = thumb_img["nsfwLevel"] if thumb_img else 0
|
|
|
|
raw_tags = model_blk.get("tags", [])
|
|
tags = [t for t in raw_tags if isinstance(t, str)] if isinstance(raw_tags, list) else []
|
|
|
|
return {
|
|
"thumbnail": thumb_img["url"] if thumb_img else None,
|
|
"nsfw_level": nsfw_level,
|
|
"civitai_name": model_blk.get("name") or data.get("name") or _clean_name(filename),
|
|
"base_model": data.get("baseModel", ""),
|
|
"description": (model_blk.get("description") or "")[:300],
|
|
"tags": tags,
|
|
"images": norm_images,
|
|
"model_id": data.get("modelId"),
|
|
"version_id": data.get("id"),
|
|
"version_name": data.get("name", ""),
|
|
}
|
|
|
|
|
|
async def _lookup(model_name: str, session, loop) -> dict:
|
|
file_hash = await _get_hash(model_name, loop)
|
|
if not file_hash:
|
|
return {"thumbnail": None, "civitai_name": _clean_name(model_name), "nsfw_level": 0, "tags": []}
|
|
data = await _cget(session, f"{CIVITAI_API_BASE}/model-versions/by-hash/{file_hash}")
|
|
return _parse_version(data, model_name)
|
|
|
|
|
|
async def _lookup_by_path(abs_path: str, model_name: str, session, loop) -> dict:
|
|
"""Lookup by hashing an absolute path directly — for freshly downloaded files."""
|
|
file_hash = await _get_hash_by_path(abs_path, loop)
|
|
if not file_hash:
|
|
return {"thumbnail": None, "civitai_name": _clean_name(model_name), "nsfw_level": 0, "tags": []}
|
|
data = await _cget(session, f"{CIVITAI_API_BASE}/model-versions/by-hash/{file_hash}")
|
|
return _parse_version(data, model_name)
|
|
|
|
|
|
def _keep_manual_overrides(new_info: dict, existing: dict) -> dict:
|
|
# Preserve local thumbnail
|
|
existing_thumb = existing.get("thumbnail") or ""
|
|
if existing_thumb.startswith("/cwk/local_thumbnails/"):
|
|
new_info["thumbnail"] = existing["thumbnail"]
|
|
# Preserve manual NSFW override
|
|
if existing.get("nsfw_manual") is not None:
|
|
new_info["nsfw_manual"] = existing["nsfw_manual"]
|
|
# Preserve favorite flag
|
|
if existing.get("favorite"):
|
|
new_info["favorite"] = True
|
|
# Preserve tags if new data doesn't have them
|
|
existing_tags = existing.get("tags")
|
|
if isinstance(existing_tags, list) and existing_tags and not new_info.get("tags"):
|
|
new_info["tags"] = existing_tags
|
|
# Preserve update check results
|
|
if existing.get("update_available"):
|
|
new_info["update_available"] = existing["update_available"]
|
|
new_info["latest_version_id"] = existing.get("latest_version_id")
|
|
|
|
# ── NEW: Preserve manually-edited fields ──────────────────────────────────
|
|
manual_fields = existing.get("_manual_overrides", [])
|
|
if isinstance(manual_fields, list):
|
|
for field in manual_fields:
|
|
if field in existing:
|
|
new_info[field] = existing[field]
|
|
if manual_fields:
|
|
new_info["_manual_overrides"] = manual_fields
|
|
|
|
return new_info
|
|
|
|
|
|
# ─── Preset auto-population from CivitAI image metadata ───────────────────────
|
|
|
|
_PRESET_META_KEYS = ["sampler", "scheduler", "cfgScale", "steps", "clipSkip"]
|
|
_PRESET_SIZE_KEYS = ["Size", "size"]
|
|
|
|
|
|
def _extract_best_image_metadata(images: list) -> dict:
|
|
"""
|
|
Pick the example image with the most generation metadata available
|
|
and return a preset-compatible dict.
|
|
"""
|
|
if not images:
|
|
return {}
|
|
|
|
best_img = None
|
|
best_count = 0
|
|
|
|
for img in images:
|
|
meta = img.get("meta") or img.get("metadata") or {}
|
|
if not meta or not isinstance(meta, dict):
|
|
continue
|
|
count = sum(1 for k in _PRESET_META_KEYS if meta.get(k) is not None)
|
|
# Also count size
|
|
for sk in _PRESET_SIZE_KEYS:
|
|
if meta.get(sk):
|
|
count += 1
|
|
break
|
|
if count > best_count:
|
|
best_count = count
|
|
best_img = meta
|
|
|
|
if not best_img or best_count == 0:
|
|
return {}
|
|
|
|
preset = {}
|
|
|
|
# Sampler — resolve through fallback system
|
|
raw_sampler = best_img.get("sampler")
|
|
if raw_sampler and isinstance(raw_sampler, str) and raw_sampler.strip():
|
|
preset["sampler_name"] = resolve_sampler(raw_sampler.strip())
|
|
|
|
# Scheduler — resolve through fallback system
|
|
raw_scheduler = best_img.get("scheduler")
|
|
if raw_scheduler and isinstance(raw_scheduler, str) and raw_scheduler.strip():
|
|
preset["scheduler"] = resolve_scheduler(raw_scheduler.strip())
|
|
|
|
# CFG
|
|
cfg_val = best_img.get("cfgScale")
|
|
if cfg_val is not None:
|
|
try:
|
|
preset["cfg"] = round(float(cfg_val), 2)
|
|
except (ValueError, TypeError):
|
|
pass
|
|
|
|
# Steps
|
|
steps_val = best_img.get("steps")
|
|
if steps_val is not None:
|
|
try:
|
|
preset["steps"] = int(steps_val)
|
|
except (ValueError, TypeError):
|
|
pass
|
|
|
|
# Clip skip
|
|
clip_val = best_img.get("clipSkip")
|
|
if clip_val is not None:
|
|
try:
|
|
cs = int(clip_val)
|
|
preset["clip_skip"] = -cs if cs > 0 else cs
|
|
except (ValueError, TypeError):
|
|
pass
|
|
|
|
# Resolution from "Size" field (e.g. "832x1216")
|
|
for sk in _PRESET_SIZE_KEYS:
|
|
size_str = best_img.get(sk)
|
|
if size_str and isinstance(size_str, str) and "x" in size_str.lower():
|
|
parts = size_str.lower().split("x")
|
|
if len(parts) == 2:
|
|
try:
|
|
w = int(parts[0].strip())
|
|
h = int(parts[1].strip())
|
|
if 64 <= w <= 8192 and 64 <= h <= 8192:
|
|
preset["width"] = w
|
|
preset["height"] = h
|
|
except (ValueError, TypeError):
|
|
pass
|
|
break
|
|
|
|
return preset
|
|
|
|
|
|
def _auto_populate_preset(model_name: str, images: list) -> bool:
|
|
"""
|
|
Extract generation metadata from CivitAI example images and save
|
|
as a preset for the model. Returns True if a preset was written.
|
|
"""
|
|
extracted = _extract_best_image_metadata(images)
|
|
if not extracted:
|
|
return False
|
|
|
|
presets = load_presets()
|
|
existing = presets.get(model_name, default_preset())
|
|
existing.update(extracted)
|
|
presets[model_name] = existing
|
|
save_presets(presets)
|
|
print(f"[CWK] Auto-populated preset for {model_name}: {extracted}")
|
|
return True
|
|
|
|
|
|
# ─── SSE helper ───────────────────────────────────────────────────────────────
|
|
|
|
async def _sse(resp: web.StreamResponse, payload: dict) -> None:
|
|
try:
|
|
await resp.write(f"data: {json.dumps(payload)}\n\n".encode())
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════════
|
|
# Route handlers
|
|
# ═══════════════════════════════════════════════════════════════════════════════
|
|
|
|
async def handle_list_models(req: web.Request) -> web.Response:
|
|
models = _all_models()
|
|
presets = load_presets()
|
|
result = []
|
|
for m in models:
|
|
full_path = _full_path(m["name"])
|
|
try:
|
|
file_size = os.path.getsize(full_path) if full_path else 0
|
|
except OSError:
|
|
file_size = 0
|
|
meta = _load_meta(m["name"])
|
|
result.append({
|
|
"name": m["name"],
|
|
"type": m["type"],
|
|
"preset": presets.get(m["name"], default_preset()),
|
|
"civitai": meta if meta else None,
|
|
"file_path": full_path or "",
|
|
"file_size": file_size,
|
|
})
|
|
return web.json_response(result)
|
|
|
|
|
|
async def handle_get_preset(req: web.Request) -> web.Response:
|
|
name = req.rel_url.query.get("model", "")
|
|
if not name:
|
|
return web.json_response({"error": "model parameter required"}, status=400)
|
|
presets = load_presets()
|
|
return web.json_response({"model": name, "preset": presets.get(name, default_preset())})
|
|
|
|
|
|
async def handle_save_preset(req: web.Request) -> web.Response:
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
name = body.get("model", "")
|
|
if not name:
|
|
return web.json_response({"error": "model required"}, status=400)
|
|
presets = load_presets()
|
|
existing = presets.get(name, default_preset())
|
|
existing.update(body.get("preset", {}))
|
|
presets[name] = existing
|
|
save_presets(presets)
|
|
return web.json_response({"ok": True, "model": name, "preset": existing})
|
|
|
|
|
|
async def handle_validate_key(req: web.Request) -> web.Response:
|
|
key = req.rel_url.query.get("key", "")
|
|
if not key:
|
|
return web.json_response({"ok": False, "error": "No key provided."})
|
|
try:
|
|
async with aiohttp.ClientSession(headers=_headers(key)) as s:
|
|
await _cget(s, f"{CIVITAI_API_BASE}/models?limit=1")
|
|
return web.json_response({"ok": True})
|
|
except CivitAIAuthError as e:
|
|
return web.json_response({"ok": False, "error": str(e)})
|
|
except Exception as e:
|
|
return web.json_response({"ok": False, "error": str(e)})
|
|
|
|
|
|
# ── Favourite ─────────────────────────────────────────────────────────────────
|
|
|
|
async def handle_set_favorite(req: web.Request) -> web.Response:
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
name = body.get("model", "")
|
|
favorite = bool(body.get("favorite", False))
|
|
if not name:
|
|
return web.json_response({"error": "model required"}, status=400)
|
|
meta = _load_meta(name)
|
|
if favorite:
|
|
meta["favorite"] = True
|
|
else:
|
|
meta.pop("favorite", None)
|
|
_save_meta(name, meta)
|
|
return web.json_response({"ok": True, "favorite": favorite})
|
|
|
|
|
|
# ── Bulk SSE fetch ─────────────────────────────────────────────────────────────
|
|
|
|
async def _run_sse_fetch(req, names, api_key, force, rebuild=False):
|
|
resp = web.StreamResponse(headers={
|
|
"Content-Type": "text/event-stream",
|
|
"Cache-Control": "no-cache",
|
|
"X-Accel-Buffering": "no",
|
|
})
|
|
await resp.prepare(req)
|
|
|
|
if not api_key:
|
|
await _sse(resp, {
|
|
"error": "api_key_required",
|
|
"message": "A CivitAI API key is required. Click 🔑 in the browser panel.",
|
|
"done": True,
|
|
})
|
|
await resp.write_eof()
|
|
return resp
|
|
|
|
if not names:
|
|
await _sse(resp, {"done": True, "total": 0})
|
|
await resp.write_eof()
|
|
return resp
|
|
|
|
total = len(names)
|
|
loop = asyncio.get_event_loop()
|
|
sem = asyncio.Semaphore(_CONCURRENCY)
|
|
completed = [0]
|
|
auth_failed = [False]
|
|
presets_populated = [0]
|
|
|
|
try:
|
|
async with aiohttp.ClientSession(headers=_headers(api_key)) as session:
|
|
|
|
async def _bounded(name: str):
|
|
if auth_failed[0]:
|
|
return
|
|
|
|
existing = _load_meta(name)
|
|
|
|
if not force and not rebuild and existing.get("thumbnail"):
|
|
completed[0] += 1
|
|
# Even for cached models, try to auto-populate preset if
|
|
# the model doesn't have one yet
|
|
images = existing.get("images", [])
|
|
if images:
|
|
current_presets = load_presets()
|
|
if name not in current_presets:
|
|
try:
|
|
if _auto_populate_preset(name, images):
|
|
presets_populated[0] += 1
|
|
except Exception as e:
|
|
print(f"[CWK] Preset auto-populate error for {name}: {e}")
|
|
await _sse(resp, {
|
|
"model": name, "info": existing,
|
|
"index": completed[0], "total": total,
|
|
"done": False,
|
|
})
|
|
return
|
|
|
|
async with sem:
|
|
if auth_failed[0]:
|
|
return
|
|
try:
|
|
info = await _lookup(name, session, loop)
|
|
except CivitAIAuthError as e:
|
|
auth_failed[0] = True
|
|
await _sse(resp, {"error": "api_key_invalid", "message": str(e), "done": True})
|
|
return
|
|
except Exception as e:
|
|
print(f"[CWK] Lookup error for {name}: {e}")
|
|
info = existing if existing else {
|
|
"thumbnail": None,
|
|
"civitai_name": _clean_name(name),
|
|
"nsfw_level": 0,
|
|
"tags": [],
|
|
}
|
|
|
|
try:
|
|
info = _keep_manual_overrides(info, existing)
|
|
except Exception as e:
|
|
print(f"[CWK] _keep_manual_overrides error for {name}: {e}")
|
|
|
|
_save_meta(name, info)
|
|
|
|
# Auto-populate preset from example images
|
|
images = info.get("images", [])
|
|
if images:
|
|
try:
|
|
if _auto_populate_preset(name, images):
|
|
presets_populated[0] += 1
|
|
except Exception as e:
|
|
print(f"[CWK] Preset auto-populate error for {name}: {e}")
|
|
|
|
completed[0] += 1
|
|
await _sse(resp, {
|
|
"model": name, "info": info,
|
|
"index": completed[0], "total": total,
|
|
"done": False,
|
|
})
|
|
|
|
results = await asyncio.gather(
|
|
*[_bounded(n) for n in names],
|
|
return_exceptions=True,
|
|
)
|
|
# Log any unexpected exceptions that slipped through
|
|
for i, result in enumerate(results):
|
|
if isinstance(result, Exception):
|
|
print(f"[CWK] Unhandled error in _bounded for {names[i]}: {result}")
|
|
|
|
except Exception as e:
|
|
print(f"[CWK] SSE fetch outer error: {e}")
|
|
|
|
# Always send the terminal done event and close cleanly
|
|
if not auth_failed[0]:
|
|
await _sse(resp, {
|
|
"done": True, "total": total,
|
|
"presets_populated": presets_populated[0],
|
|
})
|
|
|
|
try:
|
|
await resp.write_eof()
|
|
except Exception:
|
|
pass
|
|
|
|
return resp
|
|
|
|
|
|
async def handle_fetch_stream(req: web.Request) -> web.StreamResponse:
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.Response(status=400, text="invalid JSON")
|
|
names = body.get("models") or [m["name"] for m in _all_models()]
|
|
api_key = body.get("api_key", "")
|
|
force = bool(body.get("force", False))
|
|
rebuild = bool(body.get("rebuild", False))
|
|
return await _run_sse_fetch(req, names, api_key, force=force, rebuild=rebuild)
|
|
|
|
|
|
async def handle_refresh_all_stream(req: web.Request) -> web.StreamResponse:
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.Response(status=400, text="invalid JSON")
|
|
names = [m["name"] for m in _all_models()]
|
|
api_key = body.get("api_key", "")
|
|
return await _run_sse_fetch(req, names, api_key, force=True, rebuild=False)
|
|
|
|
|
|
async def handle_get_cache(req: web.Request) -> web.Response:
|
|
result = {}
|
|
for m in _all_models():
|
|
meta = _load_meta(m["name"])
|
|
if meta:
|
|
result[m["name"]] = meta
|
|
return web.json_response(result)
|
|
|
|
|
|
async def handle_clear_cache(req: web.Request) -> web.Response:
|
|
deleted = 0
|
|
if os.path.isdir(MODEL_META_DIR):
|
|
for fname in os.listdir(MODEL_META_DIR):
|
|
if fname.endswith(".json"):
|
|
try:
|
|
os.remove(os.path.join(MODEL_META_DIR, fname))
|
|
deleted += 1
|
|
except OSError:
|
|
pass
|
|
return web.json_response({"ok": True, "deleted": deleted})
|
|
|
|
|
|
async def handle_list_images(req: web.Request) -> web.Response:
|
|
name = req.rel_url.query.get("model", "")
|
|
if not name:
|
|
return web.json_response({"error": "model required"}, status=400)
|
|
entry = _load_meta(name)
|
|
return web.json_response({"model": name, "images": entry.get("images", [])})
|
|
|
|
|
|
async def handle_set_thumbnail(req: web.Request) -> web.Response:
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
name = body.get("model", "")
|
|
url = body.get("url", "")
|
|
if not name or not url:
|
|
return web.json_response({"error": "model and url required"}, status=400)
|
|
meta = _load_meta(name)
|
|
meta["thumbnail"] = url
|
|
_save_meta(name, meta)
|
|
return web.json_response({"ok": True, "thumbnail": url})
|
|
|
|
|
|
async def handle_set_local_thumbnail(req: web.Request) -> web.Response:
|
|
try:
|
|
os.makedirs(LOCAL_THUMBS_DIR, exist_ok=True)
|
|
reader = await req.multipart()
|
|
model_name = None
|
|
saved_url = None
|
|
async for part in reader:
|
|
if part.name == "model":
|
|
model_name = (await part.read()).decode("utf-8").strip()
|
|
elif part.name == "file":
|
|
filename = part.filename or "thumb.jpg"
|
|
safe_stem = (model_name or "unknown").replace("/", "_").replace("\\", "_")
|
|
ext = os.path.splitext(filename)[1] or ".jpg"
|
|
dest = os.path.join(LOCAL_THUMBS_DIR, f"{safe_stem}{ext}")
|
|
with open(dest, "wb") as f:
|
|
while True:
|
|
chunk = await part.read_chunk(65536)
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
saved_url = f"/cwk/local_thumbnails/{safe_stem}{ext}"
|
|
if not model_name or not saved_url:
|
|
return web.json_response({"error": "model and file required"}, status=400)
|
|
meta = _load_meta(model_name)
|
|
meta["thumbnail"] = saved_url
|
|
_save_meta(model_name, meta)
|
|
return web.json_response({"ok": True, "thumbnail": saved_url})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
|
|
async def handle_list_versions(req: web.Request) -> web.Response:
|
|
name = req.rel_url.query.get("model", "")
|
|
api_key = req.rel_url.query.get("api_key", "")
|
|
if not name:
|
|
return web.json_response({"error": "model required"}, status=400)
|
|
entry = _load_meta(name)
|
|
model_id = entry.get("model_id")
|
|
if not model_id:
|
|
return web.json_response(
|
|
{"error": "No model_id cached — run Fetch Thumbnails first."}, status=404
|
|
)
|
|
try:
|
|
async with aiohttp.ClientSession(headers=_headers(api_key)) as s:
|
|
data = await _cget(s, f"{CIVITAI_API_BASE}/models/{model_id}")
|
|
if not data:
|
|
return web.json_response({"error": "Model not found on CivitAI."}, status=404)
|
|
|
|
versions = sorted(data.get("modelVersions", []),
|
|
key=lambda v: v.get("createdAt", ""), reverse=True)
|
|
current_vid = entry.get("version_id")
|
|
|
|
# ── Build set of ALL locally installed version IDs for this model ─────
|
|
# This detects versions installed under different filenames
|
|
installed_version_ids = set()
|
|
for m in _all_models():
|
|
m_meta = _load_meta(m["name"])
|
|
m_model_id = m_meta.get("model_id")
|
|
m_version_id = m_meta.get("version_id")
|
|
if m_model_id == model_id and m_version_id:
|
|
installed_version_ids.add(m_version_id)
|
|
|
|
# ── Fetch per-version details in parallel to get earlyAccessEndsAt ───
|
|
sem = asyncio.Semaphore(_CONCURRENCY)
|
|
|
|
async def _fetch_version_detail(version_id: int) -> dict:
|
|
async with sem:
|
|
try:
|
|
async with aiohttp.ClientSession(headers=_headers(api_key)) as s:
|
|
detail = await _cget(s, f"{CIVITAI_API_BASE}/model-versions/{version_id}")
|
|
return detail or {}
|
|
except Exception as e:
|
|
print(f"[CWK] version detail fetch error for {version_id}: {e}")
|
|
return {}
|
|
|
|
version_ids = [v.get("id") for v in versions if v.get("id")]
|
|
version_details = await asyncio.gather(*[_fetch_version_detail(vid) for vid in version_ids])
|
|
detail_map = {d.get("id"): d for d in version_details if d.get("id")}
|
|
|
|
result = []
|
|
for v in versions:
|
|
files = v.get("files", [])
|
|
mfile = next(
|
|
(f for f in files if f.get("type") == "Model" and f.get("primary")),
|
|
files[0] if files else {},
|
|
)
|
|
detail = detail_map.get(v.get("id"), {})
|
|
ea_ends = detail.get("earlyAccessEndsAt") or detail.get("earlyAccessDeadline") or None
|
|
|
|
vid = v.get("id")
|
|
result.append({
|
|
"id": vid,
|
|
"name": v.get("name", ""),
|
|
"created_at": v.get("createdAt", ""),
|
|
"base_model": v.get("baseModel", ""),
|
|
"download_url": mfile.get("downloadUrl", ""),
|
|
"filename": mfile.get("name", ""),
|
|
"size_kb": mfile.get("sizeKB", 0),
|
|
"is_installed": vid in installed_version_ids,
|
|
"is_current": vid == current_vid,
|
|
"images": v.get("images", [])[:1],
|
|
"early_access_ends": ea_ends,
|
|
})
|
|
|
|
return web.json_response({"model_name": name, "versions": result})
|
|
except CivitAIAuthError as e:
|
|
return web.json_response({"error": str(e)}, status=403)
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
|
|
async def handle_model_description(req: web.Request) -> web.Response:
|
|
name = req.rel_url.query.get("model", "")
|
|
api_key = req.rel_url.query.get("api_key", "")
|
|
if not name:
|
|
return web.json_response({"error": "model required"}, status=400)
|
|
entry = _load_meta(name)
|
|
model_id = entry.get("model_id")
|
|
if not model_id:
|
|
return web.json_response(
|
|
{"error": "No model_id cached — run Fetch Thumbnails first."}, status=404
|
|
)
|
|
try:
|
|
async with aiohttp.ClientSession(headers=_headers(api_key)) as s:
|
|
data = await _cget(s, f"{CIVITAI_API_BASE}/models/{model_id}")
|
|
if not data:
|
|
return web.json_response({"error": "Model not found on CivitAI."}, status=404)
|
|
description = data.get("description") or ""
|
|
raw_tags = data.get("tags", [])
|
|
tags = [t for t in raw_tags if isinstance(t, str)] if isinstance(raw_tags, list) else []
|
|
if tags:
|
|
entry["tags"] = tags
|
|
_save_meta(name, entry)
|
|
return web.json_response({"description": description, "tags": tags})
|
|
except CivitAIAuthError as e:
|
|
return web.json_response({"error": str(e)}, status=403)
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
|
|
async def handle_check_updates(req: web.Request) -> web.Response:
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
api_key = body.get("api_key", "")
|
|
if not api_key:
|
|
return web.json_response({"error": "api_key required"}, status=400)
|
|
|
|
models = _all_models()
|
|
results = []
|
|
sem = asyncio.Semaphore(_CONCURRENCY)
|
|
|
|
# ── Pre-build a map: model_id → set of installed version_ids ──────────────
|
|
# This lets us detect that the "latest" version is already installed
|
|
# under a different filename.
|
|
installed_by_model_id = {}
|
|
for m in models:
|
|
m_meta = _load_meta(m["name"])
|
|
mid = m_meta.get("model_id")
|
|
vid = m_meta.get("version_id")
|
|
if mid and vid:
|
|
installed_by_model_id.setdefault(mid, set()).add(vid)
|
|
|
|
async def _check_one(m):
|
|
meta = _load_meta(m["name"])
|
|
model_id = meta.get("model_id")
|
|
current_id = meta.get("version_id")
|
|
if not model_id or not current_id:
|
|
return
|
|
async with sem:
|
|
try:
|
|
async with aiohttp.ClientSession(headers=_headers(api_key)) as s:
|
|
data = await _cget(s, f"{CIVITAI_API_BASE}/models/{model_id}")
|
|
if not data:
|
|
return
|
|
versions = sorted(data.get("modelVersions", []),
|
|
key=lambda v: v.get("createdAt", ""), reverse=True)
|
|
if not versions:
|
|
return
|
|
latest_id = versions[0].get("id")
|
|
# Check if the latest version is installed under ANY local file
|
|
all_installed = installed_by_model_id.get(model_id, set())
|
|
has_update = (latest_id not in all_installed)
|
|
meta["update_available"] = has_update
|
|
meta["latest_version_id"] = latest_id if has_update else None
|
|
_save_meta(m["name"], meta)
|
|
if has_update:
|
|
results.append({
|
|
"name": m["name"],
|
|
"current_version": meta.get("version_name", ""),
|
|
"latest_version": versions[0].get("name", ""),
|
|
"latest_version_id": latest_id,
|
|
})
|
|
except Exception as e:
|
|
print(f"[CWK] Update check error for {m['name']}: {e}")
|
|
|
|
await asyncio.gather(*[_check_one(m) for m in models])
|
|
return web.json_response({"ok": True, "updates": results, "count": len(results)})
|
|
|
|
|
|
async def handle_delete_model(req: web.Request) -> web.Response:
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
name = body.get("model", "")
|
|
if not name:
|
|
return web.json_response({"error": "model required"}, status=400)
|
|
errors = []
|
|
full = _full_path(name)
|
|
if full and os.path.exists(full):
|
|
try:
|
|
os.remove(full)
|
|
_bust_folder_cache()
|
|
except OSError as e:
|
|
errors.append(str(e))
|
|
_delete_meta(name)
|
|
presets = load_presets()
|
|
if name in presets:
|
|
del presets[name]
|
|
save_presets(presets)
|
|
return web.json_response({"ok": not errors, "errors": errors})
|
|
|
|
|
|
async def handle_refresh_civitai(req: web.Request) -> web.Response:
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
name = body.get("model", "")
|
|
api_key = body.get("api_key", "")
|
|
if not name:
|
|
return web.json_response({"error": "model required"}, status=400)
|
|
loop = asyncio.get_event_loop()
|
|
try:
|
|
async with aiohttp.ClientSession(headers=_headers(api_key)) as session:
|
|
info = await _lookup(name, session, loop)
|
|
existing = _load_meta(name)
|
|
info = _keep_manual_overrides(info, existing)
|
|
_save_meta(name, info)
|
|
|
|
# Auto-populate preset from example images
|
|
images = info.get("images", [])
|
|
preset_populated = False
|
|
if images:
|
|
preset_populated = _auto_populate_preset(name, images)
|
|
|
|
return web.json_response({
|
|
"ok": True, "info": info,
|
|
"preset_populated": preset_populated,
|
|
})
|
|
except CivitAIAuthError as e:
|
|
return web.json_response({"ok": False, "error": str(e)}, status=403)
|
|
except Exception as e:
|
|
return web.json_response({"ok": False, "error": str(e)}, status=500)
|
|
|
|
|
|
async def handle_get_meta(req: web.Request) -> web.Response:
|
|
name = req.rel_url.query.get("model", "")
|
|
if not name:
|
|
return web.json_response({}, status=400)
|
|
return web.json_response(_load_meta(name))
|
|
|
|
|
|
# ─── CLIP / VAE list endpoints ────────────────────────────────────────────────
|
|
|
|
async def handle_list_clips(req: web.Request) -> web.Response:
|
|
"""Return list of available CLIP models (with 'embedded' first)."""
|
|
return web.json_response({"clips": get_clip_list()})
|
|
|
|
|
|
async def handle_list_vaes(req: web.Request) -> web.Response:
|
|
"""Return list of available VAE models (with 'embedded' first)."""
|
|
return web.json_response({"vaes": get_vae_list()})
|
|
|
|
|
|
# ─── Download handler (SSE progress stream) ───────────────────────────────────
|
|
|
|
async def handle_download_version(req: web.Request) -> web.StreamResponse:
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.Response(status=400, text="invalid JSON")
|
|
|
|
model_name = body.get("model", "")
|
|
download_url = body.get("download_url", "")
|
|
filename = body.get("filename", "")
|
|
api_key = body.get("api_key", "")
|
|
|
|
if not download_url or not filename:
|
|
return web.Response(status=400, text="download_url and filename required")
|
|
|
|
save_dir = _model_save_dir(model_name)
|
|
save_path = os.path.join(save_dir, filename)
|
|
|
|
resp = web.StreamResponse(headers={
|
|
"Content-Type": "text/event-stream",
|
|
"Cache-Control": "no-cache",
|
|
"X-Accel-Buffering": "no",
|
|
})
|
|
await resp.prepare(req)
|
|
|
|
try:
|
|
dl_headers = _headers(api_key)
|
|
url = download_url
|
|
if api_key and "token=" not in url:
|
|
sep = "&" if "?" in url else "?"
|
|
url = f"{url}{sep}token={api_key}"
|
|
|
|
async with aiohttp.ClientSession(headers=dl_headers, timeout=_DL_TIMEOUT) as session:
|
|
async with session.get(url) as r:
|
|
if r.status == 401 or r.status == 403:
|
|
reason = "check your API key."
|
|
try:
|
|
err_data = await r.json(content_type=None)
|
|
err_msg = err_data.get("message") or err_data.get("error") or ""
|
|
if err_msg:
|
|
reason = err_msg
|
|
elif r.status == 401:
|
|
reason = "This model requires early access — you need to purchase access on CivitAI before downloading."
|
|
except Exception:
|
|
if r.status == 401:
|
|
reason = "This model requires early access — you need to purchase access on CivitAI before downloading."
|
|
await _sse(resp, {"error": f"⚠ Download not available: {reason}"})
|
|
await resp.write_eof()
|
|
return resp
|
|
|
|
if not r.ok:
|
|
await _sse(resp, {"error": f"CivitAI returned HTTP {r.status}"})
|
|
await resp.write_eof()
|
|
return resp
|
|
|
|
total_bytes = int(r.headers.get("Content-Length", 0))
|
|
received_bytes = 0
|
|
last_progress = -1
|
|
|
|
os.makedirs(save_dir, exist_ok=True)
|
|
with open(save_path, "wb") as f:
|
|
async for chunk in r.content.iter_chunked(1024 * 256):
|
|
f.write(chunk)
|
|
received_bytes += len(chunk)
|
|
if total_bytes > 0:
|
|
progress = int(received_bytes / total_bytes * 100)
|
|
if progress != last_progress:
|
|
last_progress = progress
|
|
await _sse(resp, {"progress": progress})
|
|
|
|
print(f"[CWK] Downloaded: {filename} → {save_path}")
|
|
|
|
# ── Bust folder cache so the new file is visible immediately ──────────
|
|
_bust_folder_cache()
|
|
|
|
# ── Fetch CivitAI metadata right now using the absolute path ──────────
|
|
loop = asyncio.get_event_loop()
|
|
try:
|
|
async with aiohttp.ClientSession(headers=_headers(api_key)) as session:
|
|
info = await _lookup_by_path(save_path, filename, session, loop)
|
|
if info.get("thumbnail"):
|
|
registered_name = filename
|
|
for folder_type in ("checkpoints", "diffusion_models"):
|
|
try:
|
|
roots = folder_paths.get_folder_paths(folder_type)
|
|
except Exception:
|
|
continue
|
|
for root in roots:
|
|
root_norm = os.path.normpath(root)
|
|
dir_norm = os.path.normpath(save_dir)
|
|
if dir_norm.startswith(root_norm):
|
|
rel = os.path.relpath(save_path, root_norm)
|
|
registered_name = rel
|
|
break
|
|
_save_meta(registered_name, info)
|
|
|
|
# Auto-populate preset for the newly downloaded model
|
|
images = info.get("images", [])
|
|
if images:
|
|
_auto_populate_preset(registered_name, images)
|
|
|
|
await _sse(resp, {"progress": 100, "done": True, "civitai": info, "registered_name": registered_name})
|
|
else:
|
|
await _sse(resp, {"progress": 100, "done": True})
|
|
except Exception as e:
|
|
print(f"[CWK] Post-download metadata fetch error: {e}")
|
|
await _sse(resp, {"progress": 100, "done": True})
|
|
|
|
except asyncio.CancelledError:
|
|
if os.path.exists(save_path):
|
|
try:
|
|
os.remove(save_path)
|
|
except OSError:
|
|
pass
|
|
except Exception as e:
|
|
print(f"[CWK] Download error: {e}")
|
|
if os.path.exists(save_path):
|
|
try:
|
|
os.remove(save_path)
|
|
except OSError:
|
|
pass
|
|
await _sse(resp, {"error": str(e)})
|
|
|
|
try:
|
|
await resp.write_eof()
|
|
except Exception:
|
|
pass
|
|
|
|
return resp
|
|
|
|
|
|
async def handle_last_model(req: web.Request) -> web.Response:
|
|
"""GET /cwk/last_model — Return the last-used model name and its preset + meta."""
|
|
model_name = get_last_used_model()
|
|
if not model_name:
|
|
return web.json_response({"model_name": None})
|
|
presets = load_presets()
|
|
preset = presets.get(model_name, default_preset())
|
|
meta = _load_meta(model_name)
|
|
return web.json_response({
|
|
"model_name": model_name,
|
|
"preset": preset,
|
|
"meta": meta if meta else None,
|
|
})
|
|
|
|
|
|
async def handle_save_last_model(req: web.Request) -> web.Response:
|
|
"""POST /cwk/last_model — Persist the last-used model name."""
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
model_name = body.get("model_name", "")
|
|
if not model_name:
|
|
return web.json_response({"error": "model_name required"}, status=400)
|
|
save_last_used_model(model_name)
|
|
return web.json_response({"ok": True})
|
|
|
|
|
|
async def handle_resolution_presets(req: web.Request) -> web.Response:
|
|
"""GET /cwk/resolution_presets — Return all resolution preset labels and values."""
|
|
result = []
|
|
for label, (w, h) in RESOLUTION_PRESETS.items():
|
|
result.append({"label": label, "width": w, "height": h})
|
|
return web.json_response(result)
|
|
|
|
# ─── Manual metadata edit endpoint ────────────────────────────────────────
|
|
|
|
async def handle_edit_meta(req: web.Request) -> web.Response:
|
|
"""POST /cwk/civitai/meta/edit — Manually set metadata fields for models
|
|
that don't have CivitAI data (or override existing values).
|
|
Accepted fields: civitai_name, version_name, base_model.
|
|
These are flagged as manual overrides so they survive CivitAI refreshes."""
|
|
try:
|
|
body = await req.json()
|
|
except Exception:
|
|
return web.json_response({"error": "invalid JSON"}, status=400)
|
|
|
|
model_name = body.get("model", "")
|
|
if not model_name:
|
|
return web.json_response({"error": "model required"}, status=400)
|
|
|
|
edits = body.get("edits", {})
|
|
if not edits or not isinstance(edits, dict):
|
|
return web.json_response({"error": "edits required"}, status=400)
|
|
|
|
ALLOWED_FIELDS = {"civitai_name", "version_name", "base_model"}
|
|
filtered = {k: v for k, v in edits.items() if k in ALLOWED_FIELDS and isinstance(v, str)}
|
|
if not filtered:
|
|
return web.json_response({"error": "no valid fields to edit"}, status=400)
|
|
|
|
meta = _load_meta(model_name)
|
|
|
|
# Apply edits
|
|
for key, val in filtered.items():
|
|
meta[key] = val.strip()
|
|
|
|
# Track which fields were manually overridden so they survive CivitAI refreshes
|
|
manual = meta.get("_manual_overrides", [])
|
|
if not isinstance(manual, list):
|
|
manual = []
|
|
for key in filtered:
|
|
if key not in manual:
|
|
manual.append(key)
|
|
meta["_manual_overrides"] = manual
|
|
|
|
_save_meta(model_name, meta)
|
|
|
|
print(f"[CWK] Manual metadata edit for {model_name}: {filtered}")
|
|
return web.json_response({"ok": True, "meta": meta})
|
|
|
|
# ─── Route registration ───────────────────────────────────────────────────────
|
|
|
|
def register_routes(app: web.Application) -> None:
|
|
os.makedirs(LOCAL_THUMBS_DIR, exist_ok=True)
|
|
os.makedirs(MODEL_META_DIR, exist_ok=True)
|
|
r = app.router
|
|
r.add_get ("/cwk/models", handle_list_models)
|
|
r.add_get ("/cwk/preset", handle_get_preset)
|
|
r.add_post ("/cwk/preset", handle_save_preset)
|
|
r.add_post ("/cwk/civitai/fetch/stream", handle_fetch_stream)
|
|
r.add_post ("/cwk/civitai/refresh/all", handle_refresh_all_stream)
|
|
r.add_get ("/cwk/civitai/cache", handle_get_cache)
|
|
r.add_delete("/cwk/civitai/cache", handle_clear_cache)
|
|
r.add_get ("/cwk/civitai/validate", handle_validate_key)
|
|
r.add_get ("/cwk/civitai/images", handle_list_images)
|
|
r.add_post ("/cwk/civitai/thumbnail/set", handle_set_thumbnail)
|
|
r.add_post ("/cwk/civitai/thumbnail/local", handle_set_local_thumbnail)
|
|
r.add_get ("/cwk/civitai/versions", handle_list_versions)
|
|
r.add_get ("/cwk/civitai/model-description", handle_model_description)
|
|
r.add_post ("/cwk/civitai/check-updates", handle_check_updates)
|
|
r.add_post ("/cwk/civitai/download", handle_download_version)
|
|
r.add_post ("/cwk/model/favorite", handle_set_favorite)
|
|
r.add_delete("/cwk/model", handle_delete_model)
|
|
r.add_post ("/cwk/civitai/refresh", handle_refresh_civitai)
|
|
r.add_static("/cwk/local_thumbnails", LOCAL_THUMBS_DIR, show_index=False)
|
|
r.add_get ("/cwk/civitai/meta", handle_get_meta)
|
|
r.add_get ("/cwk/clips", handle_list_clips)
|
|
r.add_get ("/cwk/vaes", handle_list_vaes)
|
|
r.add_get ("/cwk/last_model", handle_last_model)
|
|
r.add_post ("/cwk/last_model", handle_save_last_model)
|
|
r.add_get ("/cwk/resolution_presets", handle_resolution_presets)
|
|
r.add_post ("/cwk/civitai/meta/edit", handle_edit_meta)
|
|
print("[CWK_PresetManager] Routes registered.") |