mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-18 03:21:26 -03:00
Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 8e45c22d7a | |||
| 191c4e03cd |
@@ -2,7 +2,8 @@ import json
|
|||||||
import os
|
import os
|
||||||
import re
|
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):
|
def _store_checkpoint_metadata(metadata, node_id, model_name):
|
||||||
@@ -1233,14 +1234,7 @@ class MetadataOverwriteExtractor(NodeMetadataExtractor):
|
|||||||
if not inputs:
|
if not inputs:
|
||||||
return
|
return
|
||||||
|
|
||||||
overwrite_params = {}
|
overwrite_params = collect_overwrite_params(inputs)
|
||||||
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
|
|
||||||
|
|
||||||
if overwrite_params:
|
if overwrite_params:
|
||||||
metadata.setdefault(OVERWRITE, {})
|
metadata.setdefault(OVERWRITE, {})
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -9,10 +9,8 @@ but users may wire 0 to express "no clip skip / default".
|
|||||||
|
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from ..metadata_collector.constants import (
|
from ..metadata_collector.constants import CLIP_SKIP_SENTINEL as _CLIP_SKIP_SENTINEL
|
||||||
CLIP_SKIP_SENTINEL as _CLIP_SKIP_SENTINEL,
|
from ..metadata_collector.overwrite_utils import collect_overwrite_params
|
||||||
METADATA_OVERWRITE_FIELDS,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class MetadataOverwriteLM:
|
class MetadataOverwriteLM:
|
||||||
@@ -87,12 +85,16 @@ class MetadataOverwriteLM:
|
|||||||
},
|
},
|
||||||
),
|
),
|
||||||
"model": (
|
"model": (
|
||||||
"STRING",
|
"STRING,MODEL",
|
||||||
{
|
{
|
||||||
"default": "",
|
"default": "",
|
||||||
|
"widgetType": "STRING",
|
||||||
"tooltip": (
|
"tooltip": (
|
||||||
"The checkpoint or diffusion model (UNet) used "
|
"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"
|
For most fields, a falsy value (empty string, 0) means "not set"
|
||||||
and is skipped. clip_skip uses a dedicated sentinel (-25) so that
|
and is skipped. clip_skip uses a dedicated sentinel (-25) so that
|
||||||
a wired value of 0 is preserved and reaches the metadata pipeline.
|
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] = {}
|
return (collect_overwrite_params(kwargs),)
|
||||||
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,)
|
|
||||||
|
|||||||
@@ -7,6 +7,21 @@ from ..utils.utils import get_checkpoint_info_absolute, _format_model_name_for_c
|
|||||||
logger = logging.getLogger(__name__)
|
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:
|
class UNETLoaderLM:
|
||||||
"""UNET Loader with support for extra folder paths
|
"""UNET Loader with support for extra folder paths
|
||||||
|
|
||||||
@@ -196,6 +211,12 @@ class UNETLoaderLM:
|
|||||||
# Wrap with GGUFModelPatcher
|
# Wrap with GGUFModelPatcher
|
||||||
model = GGUFModelPatcher.clone(model)
|
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,)
|
return (model,)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
+48
-1
@@ -1,7 +1,7 @@
|
|||||||
from difflib import SequenceMatcher
|
from difflib import SequenceMatcher
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
from typing import Dict
|
from typing import Any, Dict, List, Optional
|
||||||
from ..services.service_registry import ServiceRegistry
|
from ..services.service_registry import ServiceRegistry
|
||||||
from ..config import config
|
from ..config import config
|
||||||
from ..services.settings_manager import get_settings_manager
|
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)
|
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:
|
def fuzzy_match(text: str, pattern: str, threshold: float = 0.85) -> bool:
|
||||||
"""
|
"""
|
||||||
Check if text matches pattern using fuzzy matching.
|
Check if text matches pattern using fuzzy matching.
|
||||||
|
|||||||
@@ -90,7 +90,7 @@ export class BulkManager {
|
|||||||
moveAll: true,
|
moveAll: true,
|
||||||
autoOrganize: false,
|
autoOrganize: false,
|
||||||
deleteAll: true,
|
deleteAll: true,
|
||||||
setContentRating: false,
|
setContentRating: true,
|
||||||
skipMetadataRefresh: false,
|
skipMetadataRefresh: false,
|
||||||
setFavorite: true,
|
setFavorite: true,
|
||||||
unfavorite: true,
|
unfavorite: true,
|
||||||
@@ -1528,14 +1528,18 @@ export class BulkManager {
|
|||||||
let failureCount = 0;
|
let failureCount = 0;
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const apiClient = getModelApiClient();
|
const isRecipesPage = state.currentPageType === 'recipes';
|
||||||
for (const filePath of targets) {
|
for (const filePath of targets) {
|
||||||
if (cancelled) {
|
if (cancelled) {
|
||||||
showToast('toast.api.operationCancelled', {}, 'info');
|
showToast('toast.api.operationCancelled', {}, 'info');
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
try {
|
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++;
|
successCount++;
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
failureCount++;
|
failureCount++;
|
||||||
|
|||||||
@@ -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();
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user