diff --git a/py/api/lora_integration.py b/py/api/lora_integration.py index 82171eda..20d7ee4b 100644 --- a/py/api/lora_integration.py +++ b/py/api/lora_integration.py @@ -76,6 +76,7 @@ class LoraIntegrationMixin: "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, @@ -94,6 +95,7 @@ class LoraIntegrationMixin: 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 @@ -115,6 +117,7 @@ class LoraIntegrationMixin: 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( @@ -248,7 +251,11 @@ class LoraIntegrationMixin: get_preview_images_from_metadata, metadata, meta_file ) civitai_paths = await self._run_in_executor( - download_civitai_images, metadata, meta_file, cache_dir + download_civitai_images, + metadata, + meta_file, + cache_dir, + IntegrationConfig.CIVITAI_API_KEY, ) # Merge, local first, dedup diff --git a/py/config.py b/py/config.py index 360b2e2e..88ef6b93 100644 --- a/py/config.py +++ b/py/config.py @@ -212,6 +212,7 @@ class IntegrationConfig: LORA_MANAGER_ENABLED = False LORA_MANAGER_PATH = "" # Auto-detected if empty LORA_TRIGGER_WORDS_ENABLED = False # Auto-inject trigger words into prompts + CIVITAI_API_KEY = "" # Required to download NSFW example images @classmethod def get_config(cls) -> Dict[str, Any]: @@ -220,6 +221,7 @@ class IntegrationConfig: "enabled": cls.LORA_MANAGER_ENABLED, "path": cls.LORA_MANAGER_PATH, "trigger_words_enabled": cls.LORA_TRIGGER_WORDS_ENABLED, + "civitai_api_key": cls.CIVITAI_API_KEY, }, } @@ -232,6 +234,8 @@ class IntegrationConfig: cls.LORA_MANAGER_PATH = lora["path"] if "trigger_words_enabled" in lora: cls.LORA_TRIGGER_WORDS_ENABLED = lora["trigger_words_enabled"] + if "civitai_api_key" in lora: + cls.CIVITAI_API_KEY = lora["civitai_api_key"] class PromptManagerConfig: diff --git a/py/lora_utils.py b/py/lora_utils.py index dab96f6a..0eb90d93 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -297,13 +297,16 @@ def get_preview_image_from_metadata( def download_civitai_images( - metadata: Dict, metadata_path: Path, cache_dir: Path + metadata: Dict, metadata_path: Path, cache_dir: Path, api_key: str = "" ) -> List[str]: """Download civitai example images to a local cache directory. Images are stored as ``//.jpg``. Already-downloaded files are skipped. + Args: + api_key: CivitAI API key for authenticated downloads (NSFW content). + Returns: List of absolute paths to downloaded image files. """ @@ -340,9 +343,10 @@ def download_civitai_images( continue try: - req = urllib.request.Request( - url, headers={"User-Agent": "ComfyUI-PromptManager/1.0"} - ) + headers = {"User-Agent": "ComfyUI-PromptManager/1.0"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + req = urllib.request.Request(url, headers=headers) with urllib.request.urlopen(req, timeout=15) as resp: local_path.write_bytes(resp.read()) downloaded.append(str(local_path.resolve())) diff --git a/web/admin.html b/web/admin.html index ea58c478..a810f205 100644 --- a/web/admin.html +++ b/web/admin.html @@ -439,6 +439,13 @@

Automatically append trigger words when <lora:name:weight> is detected in prompts.

+
+ + +

Get your key from civitai.com/user/account

+
+