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)