kenyonxu
37f508a2ed
- Enable trust_env=True in aiohttp ClientSession so Civitai API requests respect HTTP_PROXY/HTTPS_PROXY environment variables (required in regions where Civitai is blocked, e.g. mainland China behind mihomo/clash). - Add _err_msg() helper to handle exceptions with empty str() (e.g. TimeoutError, certain aiohttp errors) by falling back to repr(). - Replace all bare str(e) in error responses with _err_msg(e) so users always see a meaningful error description.
234 lines
8.1 KiB
Python
234 lines
8.1 KiB
Python
import asyncio
|
|
import os
|
|
from datetime import datetime, timezone
|
|
|
|
import aiohttp
|
|
import folder_paths
|
|
from aiohttp import web
|
|
|
|
from .civitai_client import compute_file_hash, download_image
|
|
from .cache_manager import cache_manager
|
|
|
|
# Frontend JS directory -- ComfyUI serves all files in this dir to the browser
|
|
WEB_DIRECTORY = os.path.join(os.path.dirname(os.path.abspath(__file__)), "frontend")
|
|
|
|
# Node exports -- required for ComfyUI to discover this node
|
|
from .nodes import KaiCivitaiLoRALoader
|
|
|
|
NODE_CLASS_MAPPINGS = {"KaiCivitaiLoRALoader": KaiCivitaiLoRALoader}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"KaiCivitaiLoRALoader": "Kai Civitai LoRA Loader"}
|
|
|
|
__all__ = [
|
|
"NODE_CLASS_MAPPINGS",
|
|
"NODE_DISPLAY_NAME_MAPPINGS",
|
|
"WEB_DIRECTORY",
|
|
]
|
|
|
|
# In-flight request lock state
|
|
_fetching_locks: dict[str, asyncio.Lock] = {}
|
|
_fetching_lock_refcounts: dict[str, int] = {}
|
|
_fetching_locks_lock = asyncio.Lock()
|
|
|
|
|
|
def _now_iso() -> str:
|
|
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
|
|
|
|
def _err_msg(e: Exception) -> str:
|
|
"""Return a non-empty error message. Some exceptions (e.g. TimeoutError)
|
|
have an empty str() by default; fall back to repr() in that case."""
|
|
msg = str(e)
|
|
return msg if msg else repr(e)
|
|
|
|
|
|
def _normalize_civitai_response(raw: dict, file_hash: str) -> dict:
|
|
"""Convert Civitai API response to our cache schema."""
|
|
model = raw.get("model") or {}
|
|
images = []
|
|
for img in raw.get("images", []):
|
|
images.append({
|
|
"url": img.get("url", ""),
|
|
"local_path": None,
|
|
"type": img.get("type", "image"),
|
|
"nsfw": bool(img.get("nsfw", False)),
|
|
})
|
|
|
|
# Find matching file info for download URL (try uppercase SHA256, then lowercase)
|
|
download_url = ""
|
|
for f in raw.get("files", []):
|
|
hashes = f.get("hashes", {})
|
|
file_hashes = [
|
|
hashes.get("SHA256", "").lower(),
|
|
hashes.get("sha256", "").lower(),
|
|
]
|
|
if file_hash.lower() in file_hashes:
|
|
download_url = f.get("downloadUrl", "")
|
|
break
|
|
|
|
trained_words = raw.get("trainedWords") or []
|
|
|
|
return {
|
|
"hash": file_hash,
|
|
"modelId": raw.get("modelId"),
|
|
"name": raw.get("name", ""),
|
|
"description": model.get("description", ""),
|
|
"triggerWords": trained_words,
|
|
"images": images,
|
|
"baseModel": raw.get("baseModel", ""),
|
|
"downloadUrl": download_url,
|
|
"fetched_at": _now_iso(),
|
|
}
|
|
|
|
|
|
async def _do_fetch(lora_path: str, file_hash: str, force: bool = False) -> dict:
|
|
"""Core fetch logic: check cache -> fetch Civitai -> download images -> cache."""
|
|
if not force:
|
|
cached = cache_manager.read_cache(file_hash)
|
|
if cached is not None:
|
|
return {"success": True, "data": cached, "cached": True}
|
|
|
|
# Fetch from Civitai
|
|
from .civitai_client import fetch_model_by_hash
|
|
|
|
try:
|
|
raw = await fetch_model_by_hash(file_hash)
|
|
except aiohttp.ClientResponseError as e:
|
|
# Only cache placeholder for 404/not-found to avoid repeated 404s
|
|
if e.status == 404:
|
|
placeholder = {
|
|
"hash": file_hash,
|
|
"error": _err_msg(e),
|
|
"fetched_at": _now_iso(),
|
|
"triggerWords": [],
|
|
"images": [],
|
|
}
|
|
cache_manager.write_cache(file_hash, placeholder)
|
|
raise
|
|
except Exception:
|
|
# Don't cache for network errors, rate limits, etc. -- let user retry
|
|
raise
|
|
|
|
# Normalize Civitai response
|
|
data = _normalize_civitai_response(raw, file_hash)
|
|
|
|
# Download images
|
|
images = data.get("images", [])
|
|
for i, img in enumerate(images):
|
|
dest = cache_manager.get_image_path(file_hash, i)
|
|
success = await download_image(img["url"], dest)
|
|
if success:
|
|
img["local_path"] = dest
|
|
else:
|
|
img["local_path"] = None
|
|
|
|
cache_manager.write_cache(file_hash, data)
|
|
|
|
return {"success": True, "data": data, "cached": False}
|
|
|
|
|
|
async def fetch_civitai_info(request):
|
|
try:
|
|
body = await request.json()
|
|
except Exception:
|
|
return web.json_response({"success": False, "error": "Invalid JSON"}, status=400)
|
|
|
|
lora_name = body.get("lora_name")
|
|
force = body.get("force", False)
|
|
|
|
if not lora_name:
|
|
return web.json_response({"success": False, "error": "Missing lora_name"}, status=400)
|
|
|
|
allowed = folder_paths.get_filename_list("loras")
|
|
if lora_name not in allowed:
|
|
return web.json_response({"success": False, "error": "Invalid lora_name"}, status=400)
|
|
|
|
try:
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
except Exception as e:
|
|
return web.json_response({"success": False, "error": str(e)}, status=400)
|
|
|
|
try:
|
|
file_hash = await compute_file_hash(lora_path)
|
|
except Exception as e:
|
|
return web.json_response({"success": False, "error": f"Hash failed: {_err_msg(e)}"}, status=500)
|
|
|
|
# Get or create lock atomically with refcount
|
|
async with _fetching_locks_lock:
|
|
if file_hash not in _fetching_locks:
|
|
_fetching_locks[file_hash] = asyncio.Lock()
|
|
lock = _fetching_locks[file_hash]
|
|
_fetching_lock_refcounts[file_hash] = _fetching_lock_refcounts.get(file_hash, 0) + 1
|
|
|
|
try:
|
|
async with lock:
|
|
try:
|
|
result = await _do_fetch(lora_path, file_hash, force=force)
|
|
except aiohttp.ClientResponseError as e:
|
|
if e.status == 404:
|
|
cached = cache_manager.read_cache(file_hash)
|
|
if cached and "error" in cached:
|
|
return web.json_response({
|
|
"success": False,
|
|
"error": str(e),
|
|
"data": cached,
|
|
"cached": True,
|
|
}, status=404)
|
|
return web.json_response({"success": False, "error": _err_msg(e)}, status=e.status)
|
|
except Exception as e:
|
|
return web.json_response({"success": False, "error": _err_msg(e)}, status=500)
|
|
finally:
|
|
# Clean up lock when refcount reaches zero
|
|
async with _fetching_locks_lock:
|
|
_fetching_lock_refcounts[file_hash] -= 1
|
|
if _fetching_lock_refcounts[file_hash] == 0:
|
|
del _fetching_lock_refcounts[file_hash]
|
|
del _fetching_locks[file_hash]
|
|
|
|
return web.json_response(result)
|
|
|
|
|
|
async def get_cache(request):
|
|
hash_val = request.match_info["hash"]
|
|
cached = cache_manager.read_cache(hash_val)
|
|
if cached is None:
|
|
return web.json_response({"exists": False}, status=404)
|
|
return web.json_response({"exists": True, "data": cached})
|
|
|
|
|
|
async def get_image(request):
|
|
"""Serve a cached preview image securely."""
|
|
hash_val = request.match_info["hash"]
|
|
try:
|
|
index = int(request.match_info["index"])
|
|
except ValueError:
|
|
return web.json_response({"error": "Invalid index"}, status=400)
|
|
|
|
# Validate hash format (hex only, 64 chars) to prevent path traversal
|
|
if len(hash_val) != 64 or not all(c in "0123456789abcdefABCDEF" for c in hash_val):
|
|
return web.json_response({"error": "Invalid hash"}, status=400)
|
|
|
|
image_path = cache_manager.get_image_path(hash_val, index)
|
|
cache_dir = os.path.realpath(cache_manager.cache_dir)
|
|
abs_image_path = os.path.realpath(image_path)
|
|
|
|
# Security: ensure the resolved path is within the cache directory
|
|
if not abs_image_path.startswith(cache_dir + os.sep):
|
|
return web.json_response({"error": "Forbidden"}, status=403)
|
|
|
|
if not os.path.exists(abs_image_path):
|
|
return web.json_response({"error": "Not found"}, status=404)
|
|
|
|
return web.FileResponse(abs_image_path)
|
|
|
|
|
|
# Register HTTP routes only when ComfyUI server is available
|
|
try:
|
|
from server import PromptServer
|
|
except Exception:
|
|
PromptServer = None # type: ignore[misc,assignment]
|
|
|
|
if PromptServer is not None:
|
|
PromptServer.instance.routes.post("/kai/civitai/fetch")(fetch_civitai_info)
|
|
PromptServer.instance.routes.get("/kai/civitai/cache/{hash}")(get_cache)
|
|
PromptServer.instance.routes.get("/kai/civitai/image/{hash}/{index}")(get_image)
|