Vito Sansevero
0e41e6d65c
Most civitai example images (especially NSFW) require authentication. Adds API key field to the Integrations settings panel, passed as Bearer token when downloading example images.
417 lines
16 KiB
Python
417 lines
16 KiB
Python
"""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 every 5 items or at the end
|
|
if (i + 1) % 5 == 0 or i == total - 1:
|
|
progress = int(5 + (90 * (i + 1) / max(total, 1)))
|
|
await send_progress(
|
|
{
|
|
"type": "progress",
|
|
"status": f"Processing: {model_name}",
|
|
"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)
|