alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-05-06 16:36:45 -03:00
Code Issues Packages Projects Releases Wiki Activity

fix(recipes): sanitize remote import gen params

Browse Source
This commit is contained in:
Will Miao
2026-04-12 20:29:01 +08:00
parent 0253d001e6
commit 55e9e4bb6f
6 changed files with 244 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
py/recipes/constants.py
View File

@@ -13,4 +13,5 @@ GEN_PARAM_KEYS = [
'seed',
'size',
'clip_skip',
'denoising_strength',
]

75
py/recipes/merger.py
View File

@@ -1,27 +1,33 @@
from typing import Any, Dict, Optional
import logging
from .constants import GEN_PARAM_KEYS
logger = logging.getLogger(__name__)
class GenParamsMerger:
"""Utility to merge generation parameters from multiple sources with priority."""
ALLOWED_KEYS = set(GEN_PARAM_KEYS)
BLACKLISTED_KEYS = {
"id", "url", "userId", "username", "createdAt", "updatedAt", "hash", "meta",
"draft", "extra", "width", "height", "process", "quantity", "workflow",
"baseModel", "resources", "disablePoi", "aspectRatio", "Created Date",
"experimental", "civitaiResources", "civitai_resources", "Civitai resources",
"modelVersionId", "modelId", "hashes", "Model", "Model hash", "checkpoint_hash",
"checkpoint", "checksum", "model_checksum"
"checkpoint", "checksum", "model_checksum", "raw_metadata",
}
NORMALIZATION_MAPPING = {
# Civitai specific
"cfg": "cfg_scale",
"cfgScale": "cfg_scale",
"clipSkip": "clip_skip",
"negativePrompt": "negative_prompt",
# Case variations
"Sampler": "sampler",
"sampler_name": "sampler",
"scheduler": "sampler",
"Steps": "steps",
"Seed": "seed",
"Size": "size",
@@ -36,63 +42,40 @@ class GenParamsMerger:
def merge(
request_params: Optional[Dict[str, Any]] = None,
civitai_meta: Optional[Dict[str, Any]] = None,
embedded_metadata: Optional[Dict[str, Any]] = None
embedded_metadata: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""
Merge generation parameters from three sources.
Priority: request_params > civitai_meta > embedded_metadata
Args:
request_params: Params provided directly in the import request
civitai_meta: Params from Civitai Image API 'meta' field
embedded_metadata: Params extracted from image EXIF/embedded metadata
Returns:
Merged parameters dictionary
"""
result = {}
# 1. Start with embedded metadata (lowest priority)
Priority: request_params > civitai_meta > embedded_metadata
"""
result: Dict[str, Any] = {}
if embedded_metadata:
# If it's a full recipe metadata, we use its gen_params
if "gen_params" in embedded_metadata and isinstance(embedded_metadata["gen_params"], dict):
if "gen_params" in embedded_metadata and isinstance(
embedded_metadata["gen_params"], dict
):
GenParamsMerger._update_normalized(result, embedded_metadata["gen_params"])
else:
# Otherwise assume the dict itself contains gen_params
GenParamsMerger._update_normalized(result, embedded_metadata)
# 2. Layer Civitai meta (medium priority)
if civitai_meta:
GenParamsMerger._update_normalized(result, civitai_meta)
# 3. Layer request params (highest priority)
if request_params:
GenParamsMerger._update_normalized(result, request_params)
# Filter out blacklisted keys and also the original camelCase keys if they were normalized
final_result = {}
for k, v in result.items():
if k in GenParamsMerger.BLACKLISTED_KEYS:
continue
if k in GenParamsMerger.NORMALIZATION_MAPPING:
continue
final_result[k] = v
return final_result
return result
@staticmethod
def _update_normalized(target: Dict[str, Any], source: Dict[str, Any]) -> None:
"""Update target dict with normalized keys from source."""
for k, v in source.items():
normalized_key = GenParamsMerger.NORMALIZATION_MAPPING.get(k, k)
target[normalized_key] = v
# Also keep the original key for now if it's not the same,
# so we can filter at the end or avoid losing it if it wasn't supposed to be renamed?
# Actually, if we rename it, we should probably NOT keep both in 'target'
# because we want to filter them out at the end anyway.
if normalized_key != k:
# If we are overwriting an existing snake_case key with a camelCase one's value,
# that's fine because of the priority order of calls to _update_normalized.
pass
target[k] = v
"""Update target dict with normalized, persistence-safe keys from source."""
for key, value in source.items():
if key in GenParamsMerger.BLACKLISTED_KEYS:
continue
normalized_key = GenParamsMerger.NORMALIZATION_MAPPING.get(key, key)
if normalized_key not in GenParamsMerger.ALLOWED_KEYS:
continue
target[normalized_key] = value

8
py/routes/handlers/recipe_handlers.py
View File

@@ -756,6 +756,14 @@ class RecipeManagementHandler:
)
gen_params_request = self._parse_gen_params(params.get("gen_params"))
self._logger.info(
"Remote recipe import received: url=%s, request_gen_params_keys=%s, lora_count=%d, checkpoint_keys=%s",
image_url,
sorted(gen_params_request.keys()) if gen_params_request else [],
len(lora_entries),
sorted(checkpoint_entry.keys()) if isinstance(checkpoint_entry, dict) else [],
)
# 2. Initial Metadata Construction
metadata: Dict[str, Any] = {
"base_model": params.get("base_model", "") or "",

44
py/services/recipes/persistence_service.py
View File

@@ -12,6 +12,7 @@ from dataclasses import dataclass
from typing import Any, Dict, Iterable, Optional
from ...config import config
from ...recipes.constants import GEN_PARAM_KEYS
from ...utils.utils import calculate_recipe_fingerprint
from .errors import RecipeNotFoundError, RecipeValidationError
@@ -90,23 +91,7 @@ class RecipePersistenceService:
current_time = time.time()
loras_data = [self._normalise_lora_entry(lora) for lora in (metadata.get("loras") or [])]
checkpoint_entry = self._sanitize_checkpoint_entry(self._extract_checkpoint_entry(metadata))
gen_params = metadata.get("gen_params") or {}
if not gen_params and "raw_metadata" in metadata:
raw_metadata = metadata.get("raw_metadata", {})
gen_params = {
"prompt": raw_metadata.get("prompt", ""),
"negative_prompt": raw_metadata.get("negative_prompt", ""),
"steps": raw_metadata.get("steps", ""),
"sampler": raw_metadata.get("sampler", ""),
"cfg_scale": raw_metadata.get("cfg_scale", ""),
"seed": raw_metadata.get("seed", ""),
"size": raw_metadata.get("size", ""),
"clip_skip": raw_metadata.get("clip_skip", ""),
}
# Drop checkpoint duplication from generation parameters to store it only at top level
gen_params.pop("checkpoint", None)
gen_params = self._sanitize_gen_params_for_storage(metadata)
fingerprint = calculate_recipe_fingerprint(loras_data)
recipe_data: Dict[str, Any] = {
@@ -133,6 +118,7 @@ class RecipePersistenceService:
json_filename = f"{recipe_id}.recipe.json"
json_path = os.path.join(recipes_dir, json_filename)
json_path = os.path.normpath(json_path)
with open(json_path, "w", encoding="utf-8") as file_obj:
json.dump(recipe_data, file_obj, indent=4, ensure_ascii=False)
@@ -152,6 +138,30 @@ class RecipePersistenceService:
}
)
@staticmethod
def _sanitize_gen_params_for_storage(metadata: dict[str, Any]) -> dict[str, Any]:
gen_params = metadata.get("gen_params")
if isinstance(gen_params, dict) and gen_params:
source = gen_params
else:
source = metadata.get("raw_metadata")
if not isinstance(source, dict):
return {}
allowed_keys = set(GEN_PARAM_KEYS)
sanitized: dict[str, Any] = {}
for key in allowed_keys:
if key not in source:
continue
value = source.get(key)
if value in (None, ""):
continue
sanitized[key] = value
sanitized.pop("checkpoint", None)
return sanitized
async def delete_recipe(self, *, recipe_scanner, recipe_id: str) -> PersistenceResult:
"""Delete an existing recipe."""