Compare commits

...

2 Commits

Author SHA1 Message Date
Will Miao
8e45c22d7a fix(recipes): enable bulk content rating for selected recipes 2026-08-03 19:31:58 +08:00
Will Miao
191c4e03cd feat(metadata-overwrite): support wired MODEL input on model field
The model field now accepts either a manual string or a MODEL connection.
When wired, the model name is extracted from the patcher's
cached_patcher_init (registered by core loaders load_checkpoint_guess_config
and load_diffusion_model, preserved through LoRA clones) and converted to a
ComfyUI-style relative name via config model roots.

- model input declared as "STRING,MODEL" with widgetType STRING, so the
  text widget and the dual-type connection slot coexist; non-STRING/MODEL
  links are rejected by frontend and backend type validation
- UNETLoaderLM GGUF branch now registers a custom cached_patcher_init reload
  factory so GGUF models participate in name extraction and ModelPatcher
  deepclone/dynamic machinery
- shared collect_overwrite_params() helper keeps the node and the metadata
  extractor conversion logic in sync; extraction failures are logged instead
  of silently dropping the overwrite
2026-08-03 16:44:03 +08:00
7 changed files with 268 additions and 28 deletions

View File

@@ -2,7 +2,8 @@ import json
import os
import re
from .constants import CLIP_SKIP_SENTINEL, MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER, OVERWRITE, METADATA_OVERWRITE_FIELDS
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER, OVERWRITE
from .overwrite_utils import collect_overwrite_params
def _store_checkpoint_metadata(metadata, node_id, model_name):
@@ -1233,14 +1234,7 @@ class MetadataOverwriteExtractor(NodeMetadataExtractor):
if not inputs:
return
overwrite_params = {}
for key in METADATA_OVERWRITE_FIELDS:
value = inputs.get(key)
if key == "clip_skip":
if value != CLIP_SKIP_SENTINEL:
overwrite_params[key] = value
elif value: # truthy — only overwrite when user provided a real value
overwrite_params[key] = value
overwrite_params = collect_overwrite_params(inputs)
if overwrite_params:
metadata.setdefault(OVERWRITE, {})

View File

@@ -0,0 +1,42 @@
"""Shared helpers for Metadata Overwrite node metadata collection.
Used by both the MetadataOverwriteLM node (execution time) and the
MetadataOverwriteExtractor (hook time) so the conversion/filtering logic
cannot drift between the two paths.
"""
import logging
from typing import Any, Dict
from ..utils.utils import model_patcher_to_name
from .constants import CLIP_SKIP_SENTINEL, METADATA_OVERWRITE_FIELDS
logger = logging.getLogger(__name__)
def collect_overwrite_params(values: Dict[str, Any]) -> Dict[str, Any]:
"""Convert node input values into non-default overwrite parameters.
For most fields, a falsy value (empty string, 0) means "not set" and is
skipped. clip_skip uses a dedicated sentinel (-25) so that a wired value
of 0 is preserved. The ``model`` field accepts either a manual string or
a wired MODEL (ModelPatcher) connection; in the latter case the source
model name is extracted from the patcher's ``cached_patcher_init`` and
stored as a ComfyUI-style relative path.
"""
result: Dict[str, Any] = {}
for key in METADATA_OVERWRITE_FIELDS:
value = values.get(key)
if key == "model" and not isinstance(value, str):
value = model_patcher_to_name(value)
if value is None:
logger.warning(
"Could not extract model name from wired MODEL input "
"(no cached_patcher_init); model metadata overwrite skipped"
)
if key == "clip_skip":
if value != CLIP_SKIP_SENTINEL:
result[key] = value
elif value:
result[key] = value
return result

View File

@@ -9,10 +9,8 @@ but users may wire 0 to express "no clip skip / default".
from typing import Any
from ..metadata_collector.constants import (
CLIP_SKIP_SENTINEL as _CLIP_SKIP_SENTINEL,
METADATA_OVERWRITE_FIELDS,
)
from ..metadata_collector.constants import CLIP_SKIP_SENTINEL as _CLIP_SKIP_SENTINEL
from ..metadata_collector.overwrite_utils import collect_overwrite_params
class MetadataOverwriteLM:
@@ -87,12 +85,16 @@ class MetadataOverwriteLM:
},
),
"model": (
"STRING",
"STRING,MODEL",
{
"default": "",
"widgetType": "STRING",
"tooltip": (
"The checkpoint or diffusion model (UNet) used "
"for generation. Only overwrites when non-empty."
"for generation. Fill in the name manually or "
"connect a MODEL output — the model name is then "
"extracted automatically. Only overwrites when "
"non-empty."
),
},
),
@@ -158,13 +160,10 @@ class MetadataOverwriteLM:
For most fields, a falsy value (empty string, 0) means "not set"
and is skipped. clip_skip uses a dedicated sentinel (-25) so that
a wired value of 0 is preserved and reaches the metadata pipeline.
The ``model`` field accepts either a manual string or a wired MODEL
(ModelPatcher) connection; in the latter case the underlying model
name is extracted from the patcher's ``cached_patcher_init`` and
stored as a ComfyUI-style relative path.
"""
result: dict[str, Any] = {}
for key in METADATA_OVERWRITE_FIELDS:
value = kwargs.get(key)
if key == "clip_skip":
if value != _CLIP_SKIP_SENTINEL:
result[key] = value
elif value:
result[key] = value
return (result,)
return (collect_overwrite_params(kwargs),)

View File

@@ -7,6 +7,21 @@ from ..utils.utils import get_checkpoint_info_absolute, _format_model_name_for_c
logger = logging.getLogger(__name__)
def _reload_gguf_unet(
unet_path: str, weight_dtype: str, disable_dynamic: bool = False
) -> object:
"""Reload a GGUF diffusion model from disk (cached_patcher_init factory).
Mirrors the GGUF branch of UNETLoaderLM.load_unet so ModelPatcher
deepclone/dynamic machinery can rebuild GGUF models with the correct
GGMLOps. ``disable_dynamic`` is accepted for signature compatibility
with core ComfyUI loaders.
"""
loader = UNETLoaderLM()
model, = loader._load_gguf_unet(unet_path, unet_path, weight_dtype)
return model
class UNETLoaderLM:
"""UNET Loader with support for extra folder paths
@@ -196,6 +211,12 @@ class UNETLoaderLM:
# Wrap with GGUFModelPatcher
model = GGUFModelPatcher.clone(model)
# Register a reload factory so the MODEL carries its source path
# (cached_patcher_init) like core ComfyUI loaders do — required
# for model-name extraction downstream and for ModelPatcher
# deepclone/dynamic machinery.
model.cached_patcher_init = (_reload_gguf_unet, (unet_path, weight_dtype))
return (model,)
except Exception as e:

View File

@@ -1,7 +1,7 @@
from difflib import SequenceMatcher
import os
import re
from typing import Dict
from typing import Any, Dict, List, Optional
from ..services.service_registry import ServiceRegistry
from ..config import config
from ..services.settings_manager import get_settings_manager
@@ -294,6 +294,53 @@ def _format_model_name_for_comfyui(file_path: str, model_roots: list) -> str:
return os.path.basename(file_path)
def model_patcher_to_name(model_patcher: Any) -> Optional[str]:
"""Extract a ComfyUI-style model name from a MODEL (ModelPatcher) object.
Core ComfyUI loaders record the absolute weight file path on the patcher's
``cached_patcher_init`` attribute:
- load_checkpoint_guess_config -> (fn, (ckpt_path, ...), index)
- load_diffusion_model -> (fn, (unet_path, model_options))
Patcher clones (LoRA loaders, model merges, ...) preserve the attribute,
so the name is recoverable anywhere downstream of a core loader — including
from LoRA Manager's own loaders (CheckpointLoaderLM / UNETLoaderLM), which
call the same core load functions.
The absolute path is converted to the ComfyUI-style relative name used by
the metadata pipeline (covering standard ComfyUI roots and LoRA Manager
extra folder paths).
Returns None when the path cannot be recovered (e.g. third-party loaders
that never set ``cached_patcher_init``).
"""
init = getattr(model_patcher, "cached_patcher_init", None)
if not isinstance(init, (tuple, list)) or len(init) < 2:
return None
args = init[1]
abs_path = args[0] if args else None
if not isinstance(abs_path, str) or not abs_path:
return None
return _abs_model_path_to_name(abs_path)
def _abs_model_path_to_name(abs_path: str) -> str:
"""Convert an absolute model path to a ComfyUI-style relative name.
Tries standard ComfyUI model roots plus LoRA Manager extra folder paths;
falls back to the bare filename.
"""
try:
roots: List[str] = list(config.base_models_roots or [])
roots.extend(config.extra_checkpoints_roots or [])
roots.extend(config.extra_unet_roots or [])
formatted = _format_model_name_for_comfyui(abs_path, roots)
if formatted:
return formatted
except Exception:
pass
return os.path.basename(abs_path)
def fuzzy_match(text: str, pattern: str, threshold: float = 0.85) -> bool:
"""
Check if text matches pattern using fuzzy matching.

View File

@@ -90,7 +90,7 @@ export class BulkManager {
moveAll: true,
autoOrganize: false,
deleteAll: true,
setContentRating: false,
setContentRating: true,
skipMetadataRefresh: false,
setFavorite: true,
unfavorite: true,
@@ -1528,14 +1528,18 @@ export class BulkManager {
let failureCount = 0;
try {
const apiClient = getModelApiClient();
const isRecipesPage = state.currentPageType === 'recipes';
for (const filePath of targets) {
if (cancelled) {
showToast('toast.api.operationCancelled', {}, 'info');
break;
}
try {
await apiClient.saveModelMetadata(filePath, { preview_nsfw_level: level });
if (isRecipesPage) {
await updateRecipeMetadata(filePath, { preview_nsfw_level: level });
} else {
await getModelApiClient().saveModelMetadata(filePath, { preview_nsfw_level: level });
}
successCount++;
} catch (error) {
failureCount++;

View File

@@ -0,0 +1,133 @@
import { describe, it, beforeEach, expect, vi } from 'vitest';
const showToastMock = vi.fn();
const translateMock = vi.fn((key, params, fallback) => (typeof fallback === 'string' ? fallback : key));
const getNSFWLevelNameMock = vi.fn((level) => {
if (level >= 16) return 'XXX';
if (level >= 8) return 'X';
if (level >= 4) return 'R';
if (level >= 2) return 'PG13';
if (level >= 1) return 'PG';
return 'Unknown';
});
const loadingManagerStub = {
showSimpleLoading: vi.fn(),
showCancelButton: vi.fn(),
hide: vi.fn(),
};
const stateStub = {
currentPageType: 'recipes',
bulkMode: false,
selectedModels: new Set(),
loadingManager: loadingManagerStub,
virtualScroller: { updateSingleItem: vi.fn() },
global: { settings: {} },
};
const saveModelMetadataMock = vi.fn();
const getModelApiClientMock = vi.fn(() => ({ saveModelMetadata: saveModelMetadataMock }));
const updateRecipeMetadataMock = vi.fn(() => Promise.resolve({ success: true }));
vi.mock('../../../static/js/state/index.js', () => ({
state: stateStub,
getCurrentPageState: vi.fn(),
}));
vi.mock('../../../static/js/utils/uiHelpers.js', () => ({
showToast: showToastMock,
copyToClipboard: vi.fn(),
sendLoraToWorkflow: vi.fn(),
sendEmbeddingToWorkflow: vi.fn(),
buildLoraSyntax: vi.fn(),
getNSFWLevelName: getNSFWLevelNameMock,
}));
vi.mock('../../../static/js/api/modelApiFactory.js', () => ({
getModelApiClient: getModelApiClientMock,
resetAndReload: vi.fn(),
}));
vi.mock('../../../static/js/api/recipeApi.js', () => ({
RecipeSidebarApiClient: class {},
updateRecipeMetadata: updateRecipeMetadataMock,
extractRecipeId: vi.fn(),
}));
vi.mock('../../../static/js/api/apiConfig.js', () => ({
MODEL_TYPES: { LORA: 'loras', CHECKPOINT: 'checkpoints', EMBEDDING: 'embeddings' },
MODEL_CONFIG: {},
}));
vi.mock('../../../static/js/managers/ModalManager.js', () => ({
modalManager: { showModal: vi.fn(), closeModal: vi.fn() },
}));
vi.mock('../../../static/js/components/shared/ModelCard.js', () => ({
updateCardsForBulkMode: vi.fn(),
}));
vi.mock('../../../static/js/utils/i18nHelpers.js', () => ({
translate: translateMock,
}));
vi.mock('../../../static/js/utils/priorityTagHelpers.js', () => ({
getPriorityTagSuggestions: vi.fn(),
}));
vi.mock('../../../static/js/components/shared/NsfwLevelSelector.js', () => ({
getNsfwLevelSelector: vi.fn(),
}));
describe('BulkManager bulk content rating', () => {
beforeEach(() => {
vi.clearAllMocks();
stateStub.currentPageType = 'recipes';
stateStub.bulkMode = false;
stateStub.selectedModels.clear();
saveModelMetadataMock.mockResolvedValue(undefined);
updateRecipeMetadataMock.mockResolvedValue({ success: true });
});
async function createBulkManager() {
const { BulkManager } = await import('../../../static/js/managers/BulkManager.js');
return new BulkManager();
}
it('exposes the content rating action on the recipes page action config', async () => {
const bulk = await createBulkManager();
expect(bulk.actionConfig.recipes.setContentRating).toBe(true);
});
it('persists the rating through the recipe API when on the recipes page', async () => {
const bulk = await createBulkManager();
stateStub.currentPageType = 'recipes';
stateStub.selectedModels.add('/recipes/test.webp');
const ok = await bulk.setBulkContentRating(4, ['/recipes/test.webp']);
expect(ok).toBe(true);
expect(updateRecipeMetadataMock).toHaveBeenCalledWith('/recipes/test.webp', { preview_nsfw_level: 4 });
expect(updateRecipeMetadataMock).toHaveBeenCalledTimes(1);
expect(saveModelMetadataMock).not.toHaveBeenCalled();
expect(showToastMock).toHaveBeenCalledWith(
'toast.models.bulkContentRatingSet',
{ count: 1, level: 'R' },
'success'
);
});
it('persists the rating through the model API on model pages', async () => {
const bulk = await createBulkManager();
stateStub.currentPageType = 'loras';
stateStub.selectedModels.add('/models/test.safetensors');
const ok = await bulk.setBulkContentRating(8, ['/models/test.safetensors']);
expect(ok).toBe(true);
expect(saveModelMetadataMock).toHaveBeenCalledWith('/models/test.safetensors', { preview_nsfw_level: 8 });
expect(saveModelMetadataMock).toHaveBeenCalledTimes(1);
expect(updateRecipeMetadataMock).not.toHaveBeenCalled();
});
});