"""LoraManager integration API routes for PromptManager.""" import json import os from pathlib import Path from aiohttp import web class LoraIntegrationMixin: """Mixin providing LoraManager detection, scanning, and trigger word endpoints.""" def _register_lora_routes(self, routes): @routes.get("/prompt_manager/lora/detect") async def lora_detect_route(request): return await self.lora_detect(request) @routes.get("/prompt_manager/lora/status") async def lora_status_route(request): return await self.lora_status(request) @routes.post("/prompt_manager/lora/enable") async def lora_enable_route(request): return await self.lora_enable(request) @routes.post("/prompt_manager/lora/scan") async def lora_scan_route(request): return await self.lora_scan(request) @routes.get("/prompt_manager/lora/trigger-words") async def lora_trigger_words_route(request): return await self.lora_trigger_words(request) @routes.post("/prompt_manager/lora/refresh-cache") async def lora_refresh_cache_route(request): return await self.lora_refresh_cache(request) # ── Detection ──────────────────────────────────────────────────── async def lora_detect(self, request): """Auto-detect LoraManager installation.""" try: from ..lora_utils import detect_lora_manager path = await self._run_in_executor(detect_lora_manager) return web.json_response( { "success": True, "detected": path is not None, "path": path or "", } ) except Exception as e: self.logger.error(f"LoraManager detection failed: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) # ── Status ─────────────────────────────────────────────────────── async def lora_status(self, request): """Get current LoraManager integration status.""" try: from ..config import IntegrationConfig from ..lora_utils import detect_lora_manager, get_trigger_cache config = IntegrationConfig.get_config()["lora_manager"] cache = get_trigger_cache() # Check if the configured path is still valid detected_path = await self._run_in_executor( detect_lora_manager, config.get("path", "") ) return web.json_response( { "success": True, "enabled": config["enabled"], "path": config["path"], "trigger_words_enabled": config["trigger_words_enabled"], "civitai_api_key": config.get("civitai_api_key", ""), "detected": detected_path is not None, "detected_path": detected_path or "", "trigger_cache_loaded": cache.is_loaded, } ) except Exception as e: self.logger.error(f"LoraManager status check failed: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) # ── Enable / Disable ───────────────────────────────────────────── async def lora_enable(self, request): """Enable or disable LoraManager integration and save to config.json.""" try: data = await request.json() enabled = data.get("enabled", False) path = data.get("path", "") trigger_words = data.get("trigger_words_enabled", False) civitai_key = data.get("civitai_api_key", "") from ..config import IntegrationConfig, PromptManagerConfig from ..lora_utils import detect_lora_manager, get_trigger_cache # If enabling, validate the path if enabled: resolved = await self._run_in_executor(detect_lora_manager, path) if not resolved: return web.json_response( { "success": False, "error": "LoraManager not found at the specified path", }, status=400, ) path = resolved # Update in-memory config IntegrationConfig.LORA_MANAGER_ENABLED = enabled IntegrationConfig.LORA_MANAGER_PATH = path IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED = trigger_words IntegrationConfig.CIVITAI_API_KEY = civitai_key # Persist to config.json config_dir = os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) config_file = os.path.join(config_dir, "config.json") PromptManagerConfig.save_to_file(config_file) # Load trigger word cache if enabling cache = get_trigger_cache() if enabled and trigger_words and path: count = await self._run_in_executor(cache.load, path) self.logger.info(f"Trigger word cache loaded: {count} LoRAs") elif not enabled: cache.clear() return web.json_response( { "success": True, "enabled": enabled, "path": path, "trigger_words_enabled": trigger_words, } ) except Exception as e: self.logger.error(f"LoraManager enable/disable failed: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) # ── Scan LoRA example images ───────────────────────────────────── async def lora_scan(self, request): """Scan LoraManager metadata and import LoRA info + preview images. Streams progress as SSE, matching the existing scan pattern. """ try: from ..config import IntegrationConfig from ..lora_utils import ( download_civitai_images, find_lora_directories, get_example_prompt_from_metadata, get_lora_image_cache_dir, get_preview_images_from_metadata, get_trigger_words_from_metadata, get_model_name_from_metadata, read_lora_metadata, ) if not IntegrationConfig.LORA_MANAGER_ENABLED: return web.json_response( { "success": False, "error": "LoraManager integration is not enabled", }, status=400, ) lm_path = IntegrationConfig.LORA_MANAGER_PATH if not lm_path: return web.json_response( {"success": False, "error": "LoraManager path not configured"}, status=400, ) response = web.StreamResponse( status=200, reason="OK", headers={"Content-Type": "text/event-stream"}, ) await response.prepare(request) async def send_progress(data): line = f"data: {json.dumps(data)}\n\n" await response.write(line.encode("utf-8")) await send_progress( { "type": "progress", "status": "Clearing previous lora-manager imports...", "progress": 0, } ) # Clear previous imports so reimport is always clean await self._run_in_executor( self.db.delete_prompts_by_category, "lora-manager" ) await send_progress( { "type": "progress", "status": "Finding LoRA directories...", "progress": 2, } ) lora_dirs = await self._run_in_executor(find_lora_directories, lm_path) # Collect all metadata files meta_files = [] for d in lora_dirs: dir_path = Path(d) meta_files.extend(dir_path.rglob("*.metadata.json")) total = len(meta_files) imported = 0 skipped = 0 await send_progress( { "type": "progress", "status": f"Found {total} LoRA metadata files", "progress": 5, "total": total, } ) cache_dir = get_lora_image_cache_dir() for i, meta_file in enumerate(meta_files): metadata = await self._run_in_executor(read_lora_metadata, meta_file) if not metadata: skipped += 1 continue model_name = get_model_name_from_metadata(metadata) trigger_words = get_trigger_words_from_metadata(metadata) # Collect all images: local previews + downloaded civitai examples preview_paths = await self._run_in_executor( get_preview_images_from_metadata, metadata, meta_file ) civitai_paths = await self._run_in_executor( download_civitai_images, metadata, meta_file, cache_dir, IntegrationConfig.CIVITAI_API_KEY, ) # Merge, local first, dedup seen = set(preview_paths) all_images = list(preview_paths) for cp in civitai_paths: if cp not in seen: all_images.append(cp) seen.add(cp) # Build prompt text: prefer example prompt, then model name example_prompt = get_example_prompt_from_metadata(metadata) prompt_text = example_prompt or model_name # Build tags tags = ["lora-manager", f"lora:{model_name}"] tags.extend(trigger_words) # Save to database via existing mechanism try: import hashlib prompt_hash = hashlib.sha256( prompt_text.strip().lower().encode("utf-8") ).hexdigest() existing = await self._run_in_executor( self.db.get_prompt_by_hash, prompt_hash ) if existing: # Link all images for pp in all_images: await self._run_in_executor( self.db.link_image_to_prompt, existing["id"], pp, ) skipped += 1 else: prompt_id = await self._run_in_executor( self.db.save_prompt, prompt_text, "lora-manager", # category tags, None, # rating None, # notes prompt_hash, ) if prompt_id: for pp in all_images: await self._run_in_executor( self.db.link_image_to_prompt, prompt_id, pp, ) imported += 1 else: skipped += 1 except Exception as e: self.logger.warning(f"Failed to import LoRA {model_name}: {e}") skipped += 1 # Progress update for every LoRA progress = int(5 + (90 * (i + 1) / max(total, 1))) img_count = len(all_images) status = f"{model_name}" if img_count: status += f" ({img_count} images)" await send_progress( { "type": "progress", "status": status, "progress": progress, "processed": i + 1, "imported": imported, "skipped": skipped, } ) await send_progress( { "type": "complete", "progress": 100, "total": total, "imported": imported, "skipped": skipped, } ) await response.write_eof() return response except Exception as e: self.logger.error(f"LoRA scan failed: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) # ── Trigger word endpoints ─────────────────────────────────────── async def lora_trigger_words(self, request): """Look up trigger words for a specific LoRA name.""" try: from ..config import IntegrationConfig from ..lora_utils import get_trigger_cache if not IntegrationConfig.LORA_MANAGER_ENABLED: return web.json_response( {"success": False, "error": "LoraManager integration not enabled"}, status=400, ) lora_name = request.query.get("name", "") if not lora_name: return web.json_response( {"success": False, "error": "Missing 'name' query parameter"}, status=400, ) cache = get_trigger_cache() if not cache.is_loaded: lm_path = IntegrationConfig.LORA_MANAGER_PATH if lm_path: await self._run_in_executor(cache.load, lm_path) words = cache.get_trigger_words(lora_name) return web.json_response( {"success": True, "lora": lora_name, "trigger_words": words} ) except Exception as e: self.logger.error(f"Trigger word lookup failed: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) async def lora_refresh_cache(self, request): """Force-refresh the trigger word cache from disk.""" try: from ..config import IntegrationConfig from ..lora_utils import get_trigger_cache if not IntegrationConfig.LORA_MANAGER_ENABLED: return web.json_response( {"success": False, "error": "LoraManager integration not enabled"}, status=400, ) lm_path = IntegrationConfig.LORA_MANAGER_PATH if not lm_path: return web.json_response( {"success": False, "error": "LoraManager path not configured"}, status=400, ) cache = get_trigger_cache() count = await self._run_in_executor(cache.load, lm_path) return web.json_response( {"success": True, "loras_with_trigger_words": count} ) except Exception as e: self.logger.error(f"Trigger cache refresh failed: {e}") return web.json_response({"success": False, "error": str(e)}, status=500)