import folder_paths import json import os import hashlib import requests from datetime import datetime from nodes import LoraLoader from server import PromptServer from aiohttp import web import logging import threading CIVITAI_API_BASE = "https://civitai.com/api/v1" # On stocke le cache dans le dossier utilisateur de ComfyUI (comme rgthree) # Ex: ComfyUI/user/nomad/neurad_cache/ CACHE_DIR = os.path.join(folder_paths.user_directory, "neurad_cache") if not os.path.exists(CACHE_DIR): os.makedirs(CACHE_DIR) class VisualLoraLoaderNode: @classmethod def INPUT_TYPES(s): return { "required": { "model": ("MODEL",), "clip": ("CLIP",), "lora_data": ("STRING", {"default": "[]", "multiline": False, "hidden": True}), } } RETURN_TYPES = ("MODEL", "CLIP", "STRING") RETURN_NAMES = ("MODEL", "CLIP", "active_loras_info") FUNCTION = "load_lora" CATEGORY = "loaders" def load_lora(self, model, clip, lora_data): try: active_loras = json.loads(lora_data) except: return (model, clip, "Error") current_model, current_clip = model, clip lora_loader_instance = LoraLoader() applied_info = [] for lora in active_loras: lora_name = lora.get("name") strength_model = float(lora.get("strength_model", 1.0)) strength_clip = float(lora.get("strength_clip", strength_model)) subfolder = lora.get("subfolder", "") full_name = f"{subfolder.replace('\\', '/')}/{lora_name.replace('\\', '/')}".replace("//", "/").strip("/") path = folder_paths.get_full_path("loras", full_name) or folder_paths.get_full_path("loras", lora_name) if path: try: current_model, current_clip = lora_loader_instance.load_lora(current_model, current_clip, full_name, strength_model, strength_clip) applied_info.append(f"{full_name} ({strength_model})") except Exception as e: logging.error(f"[Neurad] Load error: {e}") return (current_model, current_clip, ", ".join(applied_info)) @PromptServer.instance.routes.post("/neurad/reset-cache") async def reset_cache(request): """ PURGE TOTALE : Supprime le fichier de backup ET tous les fichiers de cache par hash. """ try: # 1. Supprimer le backup LocalStorage backup_path = os.path.join(CACHE_DIR, "localStorage_backup.json") if os.path.exists(backup_path): os.remove(backup_path) logging.info("[Neurad] Backup file deleted.") # 2. Supprimer TOUS les fichiers de cache .json (sauf le backup qu'on vient de traiter) # On liste tous les fichiers dans le dossier de cache for filename in os.listdir(CACHE_DIR): if filename.endswith(".json") and filename != "localStorage_backup.json": file_path = os.path.join(CACHE_DIR, filename) try: os.remove(file_path) except Exception as e: logging.warning(f"[Neurad] Could not delete cache file {filename}: {e}") # 3. CRÉER LE DRAPEAU DE PURGE (Fichier vide) flag_path = os.path.join(CACHE_DIR, "purge_flag.lock") with open(flag_path, 'w') as f: f.write("purged") # Ou laisser vide logging.info("[Neurad] Full cache purge completed.") return web.json_response({"success": True, "message": "Cache purged"}) except Exception as e: logging.error(f"[Neurad] Critical purge error: {e}") return web.json_response({"error": str(e)}, status=500) def get_sha256_hash(file_path): """Calcule le hash SHA256 du fichier (comme rgthree).""" if not os.path.exists(file_path): return None sha256_hash = hashlib.sha256() with open(file_path, "rb") as f: for byte_block in iter(lambda: f.read(4096), b""): sha256_hash.update(byte_block) return sha256_hash.hexdigest() def get_cache_file_path(file_hash): return os.path.join(CACHE_DIR, f"{file_hash}.json") def load_from_cache(file_hash): cache_path = get_cache_file_path(file_hash) if os.path.exists(cache_path): try: with open(cache_path, 'r', encoding='utf-8') as f: data = json.load(f) # Vérifier si le cache est trop vieux (optionnel, ici 24h) # if datetime.now().timestamp() - data.get('timestamp', 0) < 86400: return data.get('response') except: pass return None def has_cache_for_file(file_path): """Vérifie simplement si un fichier de cache existe pour ce fichier .safetensors.""" file_hash = get_sha256_hash(file_path) if not file_hash: return False cache_path = get_cache_file_path(file_hash) return os.path.exists(cache_path) def save_to_cache(file_hash, data): cache_path = get_cache_file_path(file_hash) try: with open(cache_path, 'w', encoding='utf-8') as f: json.dump({'timestamp': datetime.now().timestamp(), 'response': data}, f, indent=2) except Exception as e: logging.warning(f"[Neurad] Cache write error (non-critical): {e}") def fetch_civitai_by_hash(file_hash): """Appelle l'API Civitai par HASH (la méthode magique de rgthree).""" url = f"{CIVITAI_API_BASE}/model-versions/by-hash/{file_hash}" try: response = requests.get(url, timeout=10) if response.status_code == 200: data = response.json() # Vérifier que c'est bien un modèle (parfois le hash ne correspond à rien) if data and 'model' in data: return data elif response.status_code == 404: logging.info(f"[Neurad] No model found for hash {file_hash}") else: logging.warning(f"[Neurad] API Error {response.status_code} for hash {file_hash}") except Exception as e: logging.warning(f"[Neurad] Request failed: {e}") return None def parse_civitai_response(data): """Formate la réponse Civitai pour notre frontend.""" if not data: return None # 1. Extraction ROBUSTE de l'ID du MODÈLE (pas la version) model_id = None # Priorité 1: Le champ 'modelId' explicite (souvent à la racine pour cette API) if 'modelId' in data: model_id = data.get('modelId') # Priorité 2: L'objet 'model' imbriqué if not model_id and 'model' in data and isinstance(data['model'], dict): model_id = data['model'].get('id') # Priorité 3: Fallback sur 'id' si on pense que c'est le modèle (rare pour by-hash) # Mais attention, si on n'a que l'ID de version, on doit le gérer différemment pour l'URL version_id = data.get('id') if not model_id and version_id: # Si on n'a trouvé que l'ID de version, on essaie de voir si 'model' existe quand même # Sinon, on devra construire l'URL avec l'ID de version en paramètre pass # 2. Construction de l'URL CORRECTE civitai_url = "https://civitai.com/models" if model_id: civitai_url += f"/{model_id}" # Si on a aussi un ID de version différent, on l'ajoute en paramètre pour tomber sur la bonne version if version_id and version_id != model_id: civitai_url += f"?modelVersionId={version_id}" elif version_id: # Cas rare : on n'a que l'ID de version, on tente l'URL directe (ça redirige souvent) # Mais c'est mieux d'avoir le modelId. Si on est ici, c'est qu'on n'a pas trouvé le modelId. # On laisse l'URL de base ou on tente avec l'ID de version (moins fiable) civitai_url += f"/{version_id}" # Risqué, mais mieux que None # 3. Nom model_name = "" if 'model' in data and isinstance(data['model'], dict): model_name = data['model'].get('name', '') version_name = data.get('name', '') full_name = model_name if version_name and model_name and version_name != model_name: full_name += f" - {version_name}" elif not model_name: full_name = version_name # 4. Images images = [] # 4. Images (Toutes les images, sans limite) images = [] if 'images' in data: for img in data['images']: meta_data = img.get('meta', {}) if isinstance(meta_data, str): try: meta_data = json.loads(meta_data) except: meta_data = {} images.append({ "url": img.get('url'), "nsfwLevel": img.get('nsfwLevel', 0), "width": img.get('width'), "height": img.get('height'), "positive": meta_data.get('prompt', '') if meta_data else "" }) return { "source": "civitai_hash", "model_id": model_id if model_id else version_id, "name": full_name, "trainedWords": (data.get('trainedWords', []) or []) + (data.get('triggerWords', []) or []), "images": images, "strengthMin": None, "strengthMax": None, "civitaiUrl": civitai_url } @PromptServer.instance.routes.get("/neurad/get-loras") async def get_loras(request): """ LISTE ULTRA-RAPIDE. Ne calcule AUCUN hash, ne lit AUCUN cache. Renvoie has_meta: False pour tout le monde. C'est le JS (LocalStorage) qui gère l'affichage des icônes déjà connues. """ lora_folders = folder_paths.get_folder_paths("loras") all_loras = [] for root_folder in lora_folders: if not os.path.exists(root_folder): continue for root, dirs, files in os.walk(root_folder): for filename in files: if filename.endswith(".safetensors"): try: rel_path = os.path.relpath(os.path.join(root, filename), root_folder) except: rel_path = filename all_loras.append({ "name": filename, "display_name": os.path.splitext(filename)[0], "relative_path": rel_path, "subfolder": os.path.dirname(rel_path), "has_meta": False, # Toujours faux ici pour la vitesse "meta": {} # Vide }) all_loras.sort(key=lambda x: (x["subfolder"], x["display_name"].lower())) return web.json_response(all_loras) @PromptServer.instance.routes.post("/neurad/fetch-single-info") async def fetch_single_info(request): """ SECURED: Confines path resolution to allowed lora folders using realpath + commonpath. """ try: data = await request.json() rel_path = data.get("path") if not rel_path: return web.json_response({"error": "No path provided"}, status=400) # Security Fix: Normalize the input path to prevent traversal tricks early # os.path.normpath collapses '..' and '.' but doesn't resolve symlinks yet clean_rel_path = os.path.normpath(rel_path) # Reject absolute paths immediately as we expect a relative path from the client if os.path.isabs(clean_rel_path): return web.json_response({"error": "Absolute paths are not allowed"}, status=400) full_path = None lora_folders = folder_paths.get_folder_paths("loras") for root in lora_folders: if not os.path.exists(root): continue # Construct the candidate path candidate_path = os.path.join(root, clean_rel_path) # Resolve to absolute real path (resolves symlinks and '..') # If the file doesn't exist, realpath still resolves the directory structure up to the last existing component, # but for strict security, we usually want to ensure the target exists and is within bounds. if not os.path.exists(candidate_path): continue real_candidate = os.path.realpath(candidate_path) real_root = os.path.realpath(root) # Ensure the resolved path starts with the resolved root directory # We add os.sep to prevent matching partial folder names (e.g. /loras vs /loras_private) if real_candidate.startswith(real_root + os.sep) or real_candidate == real_root: # Additional check: ensure it's a file, not a directory if os.path.isfile(real_candidate): full_path = real_candidate break else: # Path escapes the root directory logging.warning(f"[Neurad] Blocked path traversal attempt: {rel_path} resolved to {real_candidate}") continue if not full_path: return web.json_response({"error": "File not found or access denied"}, status=404) # 1. Calcul du Hash file_hash = get_sha256_hash(full_path) if not file_hash: return web.json_response({"error": "Could not calculate hash"}, status=500) # 2. Vérifier le cache local cached_data = load_from_cache(file_hash) if cached_data: parsed = parse_civitai_response(cached_data) if parsed: return web.json_response({"success": True, "meta": parsed}) # 3. Fetch API par Hash raw_data = fetch_civitai_by_hash(file_hash) if not raw_data: return web.json_response({"success": False, "error": "No model found on Civitai for this file hash"}, status=404) # 4. Parser et Sauvegarder parsed = parse_civitai_response(raw_data) if parsed: save_to_cache(file_hash, raw_data) return web.json_response({"success": True, "meta": parsed}) return web.json_response({"success": False, "error": "Data parsing failed"}, status=500) except Exception as e: logging.error(f"[Neurad] Critical Error: {e}") return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/neurad/save-localstorage-backup") async def save_backup(request): """Reçoit les données du JS et les sauvegarde sur le disque dur.""" try: data = await request.json() backup_path = os.path.join(CACHE_DIR, "localStorage_backup.json") with open(backup_path, 'w', encoding='utf-8') as f: json.dump(data, f, indent=2) logging.info("[Neurad] LocalStorage backup saved to disk.") return web.json_response({"success": True}) except Exception as e: logging.error(f"[Neurad] Backup save error: {e}") return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.post("/neurad/update-lora-info") async def update_lora_info(request): """ Reçoit les données modifiées par l'utilisateur et met à jour le fichier de cache JSON correspondant. """ try: data = await request.json() lora_path = data.get("path") # Le chemin relatif unique (ex: "subfolder/model.safetensors") updated_meta = data.get("meta") # Les nouvelles données (nom, triggers, etc.) if not lora_path or not updated_meta: return web.json_response({"error": "Missing path or meta data"}, status=400) backup_path = os.path.join(CACHE_DIR, "localStorage_backup.json") backup_data = {} if os.path.exists(backup_path): try: with open(backup_path, 'r', encoding='utf-8') as f: backup_data = json.load(f) except: backup_data = {} # Mise à jour des données pour cette LoRA backup_data[lora_path] = updated_meta # Ajout d'un flag pour indiquer que c'est édité manuellement (optionnel mais utile) backup_data[lora_path]['edited_by_user'] = True # Sauvegarde sur le disque with open(backup_path, 'w', encoding='utf-8') as f: json.dump(backup_data, f, indent=2, ensure_ascii=False) logging.info(f"[Neurad] Updated manual data for: {lora_path}") return web.json_response({"success": True}) except Exception as e: logging.error(f"[Neurad] Update error: {e}") return web.json_response({"error": str(e)}, status=500) @PromptServer.instance.routes.get("/neurad/load-localstorage-backup") async def load_backup(request): """Renvoie les données sauvegardées sur le disque pour restaurer le JS.""" try: flag_path = os.path.join(CACHE_DIR, "purge_flag.lock") # SI LE DRAPEAU EXISTE : On bloque la restauration et on nettoie le drapeau if os.path.exists(flag_path): try: os.remove(flag_path) logging.info("[Neurad] Purge flag detected. Blocking restoration to ensure clean state.") except: pass # On renvoie des données vides, peu importe si le backup existe return web.json_response({"success": False, "data": {}, "blocked_by_purge": True}) backup_path = os.path.join(CACHE_DIR, "localStorage_backup.json") if os.path.exists(backup_path): # Vérification de sécurité : si le fichier est vide ou corrompu, on ne restaure rien try: with open(backup_path, 'r', encoding='utf-8') as f: data = json.load(f) # On ne restaure que si c'est un dictionnaire non vide if isinstance(data, dict) and len(data) > 0: return web.json_response({"success": True, "data": data}) else: # Fichier vide ou invalide, on considère qu'il n'y a rien à restaurer return web.json_response({"success": False, "data": {}}) except json.JSONDecodeError: logging.warning("[Neurad] Backup file corrupted, ignoring.") return web.json_response({"success": False, "data": {}}) else: return web.json_response({"success": False, "data": {}}) except Exception as e: logging.error(f"[Neurad] Backup load error: {e}") return web.json_response({"error": str(e)}, status=500) NODE_CLASS_MAPPINGS = {"NeuradVisualLora": VisualLoraLoaderNode} NODE_DISPLAY_NAME_MAPPINGS = {"NeuradVisualLora": "Visual LoRA Loader - Neurad"}