mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 11:11:26 -03:00
Compare commits
3 Commits
3f8381ffee
...
303cca0d85
| Author | SHA1 | Date | |
|---|---|---|---|
| 303cca0d85 | |||
| c2f16784b3 | |||
| 5bc6d8286c |
@@ -82,11 +82,7 @@ class GenericNodeExtractor(NodeMetadataExtractor):
|
||||
text = val.strip()
|
||||
break
|
||||
|
||||
input_conditionings = [
|
||||
value
|
||||
for input_name, value in inputs.items()
|
||||
if input_name.startswith("conditioning") and value is not None
|
||||
]
|
||||
input_conditionings = _collect_conditioning_inputs(inputs)
|
||||
if text or input_conditionings:
|
||||
prompt_metadata = _ensure_prompt_metadata(metadata, node_id)
|
||||
if text:
|
||||
@@ -108,7 +104,11 @@ class GenericNodeExtractor(NodeMetadataExtractor):
|
||||
if not output_tuple or len(output_tuple) < 1:
|
||||
return
|
||||
|
||||
output_conditioning = output_tuple[0]
|
||||
conditioning_index = _first_conditioning_index(return_types)
|
||||
if conditioning_index is None or len(output_tuple) <= conditioning_index:
|
||||
return
|
||||
|
||||
output_conditioning = output_tuple[conditioning_index]
|
||||
if output_conditioning is None:
|
||||
return
|
||||
|
||||
@@ -434,6 +434,34 @@ def _first_output_tuple(outputs):
|
||||
return None
|
||||
|
||||
|
||||
def _first_conditioning_index(return_types):
|
||||
"""Return the index of the first CONDITIONING output slot, or None."""
|
||||
if not return_types:
|
||||
return None
|
||||
for index, return_type in enumerate(return_types):
|
||||
if "CONDITIONING" in str(return_type):
|
||||
return index
|
||||
return None
|
||||
|
||||
|
||||
def _collect_conditioning_inputs(inputs):
|
||||
"""Collect conditioning object inputs (``conditioning*`` keys).
|
||||
|
||||
Primitive values (None, str, int, float, bool) are excluded so scalar
|
||||
fields like ``conditioning_strength`` are not mistaken for conditioning
|
||||
objects during provenance tracking.
|
||||
"""
|
||||
if not inputs:
|
||||
return []
|
||||
return [
|
||||
value
|
||||
for input_name, value in inputs.items()
|
||||
if input_name.startswith("conditioning")
|
||||
and value is not None
|
||||
and not isinstance(value, (str, int, float, bool))
|
||||
]
|
||||
|
||||
|
||||
def _record_conditioning_source(
|
||||
metadata, node_id, output_conditioning, input_conditionings
|
||||
):
|
||||
@@ -446,6 +474,14 @@ def _record_conditioning_source(
|
||||
if not sources:
|
||||
return
|
||||
|
||||
# Identity-preserving selectors return one of their inputs unchanged:
|
||||
# only that input contributed to the output, so record it alone instead
|
||||
# of treating every input as a combination source.
|
||||
for conditioning in sources:
|
||||
if id(conditioning) == id(output_conditioning):
|
||||
sources = [conditioning]
|
||||
break
|
||||
|
||||
prompt_metadata = _ensure_prompt_metadata(metadata, node_id)
|
||||
prompt_metadata.setdefault("conditioning_sources", []).append(
|
||||
{
|
||||
@@ -525,13 +561,7 @@ class ConditioningCombineExtractor(NodeMetadataExtractor):
|
||||
if not inputs:
|
||||
return
|
||||
|
||||
input_conditionings = []
|
||||
for input_name in inputs:
|
||||
if (
|
||||
input_name.startswith("conditioning")
|
||||
and inputs[input_name] is not None
|
||||
):
|
||||
input_conditionings.append(inputs[input_name])
|
||||
input_conditionings = _collect_conditioning_inputs(inputs)
|
||||
|
||||
if input_conditionings:
|
||||
prompt_metadata = _ensure_prompt_metadata(metadata, node_id)
|
||||
|
||||
+14
-4
@@ -11,7 +11,7 @@ import re
|
||||
from typing import Dict, List, Any, Optional, Tuple
|
||||
from abc import ABC, abstractmethod
|
||||
from ..config import config
|
||||
from ..utils.constants import VALID_LORA_TYPES, VALID_CHECKPOINT_SUB_TYPES
|
||||
from ..utils.constants import MODEL_WEIGHT_FILE_TYPES, VALID_LORA_TYPES, VALID_CHECKPOINT_SUB_TYPES
|
||||
from ..utils.civitai_utils import rewrite_preview_url
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -155,9 +155,9 @@ class RecipeMetadataParser(ABC):
|
||||
|
||||
# Process file information if available
|
||||
if 'files' in civitai_info:
|
||||
# Find the primary model file (type="Model" and primary=true) in the files list
|
||||
# Find the primary model file (weights-type and primary=true) in the files list
|
||||
model_file = next((file for file in civitai_info.get('files', [])
|
||||
if file.get('type') == 'Model' and file.get('primary') == True), None)
|
||||
if file.get('type') in MODEL_WEIGHT_FILE_TYPES and file.get('primary') == True), None)
|
||||
|
||||
if model_file:
|
||||
# Get size
|
||||
@@ -261,11 +261,21 @@ class RecipeMetadataParser(ABC):
|
||||
checkpoint['id'] = civitai_data.get('id', 0)
|
||||
|
||||
if 'files' in civitai_data:
|
||||
# Prefer the file CivitAI marked primary; fall back to any
|
||||
# weights-type file (providers without primary flags).
|
||||
model_file = next(
|
||||
(
|
||||
file
|
||||
for file in civitai_data.get('files', [])
|
||||
if file.get('type') == 'Model'
|
||||
if file.get('type') in MODEL_WEIGHT_FILE_TYPES
|
||||
and file.get('primary') is True
|
||||
),
|
||||
None,
|
||||
) or next(
|
||||
(
|
||||
file
|
||||
for file in civitai_data.get('files', [])
|
||||
if file.get('type') in MODEL_WEIGHT_FILE_TYPES
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
@@ -30,6 +30,7 @@ from ..services.websocket_progress_callback import (
|
||||
WebSocketProgressCallback,
|
||||
)
|
||||
from ..utils.exif_utils import ExifUtils
|
||||
from ..utils.constants import MODEL_WEIGHT_FILE_TYPES
|
||||
from ..utils.metadata_manager import MetadataManager
|
||||
from .model_route_registrar import COMMON_ROUTE_DEFINITIONS, ModelRouteRegistrar
|
||||
from .handlers.model_handlers import (
|
||||
@@ -251,7 +252,7 @@ class BaseModelRoutes(ABC):
|
||||
|
||||
def _find_model_file(self, files):
|
||||
"""Find the appropriate model file from the files list - can be overridden by subclasses."""
|
||||
return next((file for file in files if file.get("type") in ("Model", "Diffusion Model") and file.get("primary") is True), None)
|
||||
return next((file for file in files if file.get("type") in MODEL_WEIGHT_FILE_TYPES and file.get("primary") is True), None)
|
||||
|
||||
def get_handler(self, name: str) -> Callable[[web.Request], Awaitable[web.StreamResponse]]:
|
||||
"""Expose handlers for subclasses or tests."""
|
||||
|
||||
@@ -21,6 +21,7 @@ from .model_metadata_provider import (
|
||||
from .downloader import get_downloader
|
||||
from .errors import RateLimitError, ResourceNotFoundError
|
||||
from ..utils.civitai_utils import resolve_license_payload
|
||||
from ..utils.constants import MODEL_WEIGHT_FILE_TYPES
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -538,10 +539,16 @@ class CivitaiClient:
|
||||
return model_versions[0]
|
||||
|
||||
def _extract_primary_model_hash(self, version_entry: Dict[str, Any]) -> Optional[str]:
|
||||
# Prefer the generic "Model" file (most reliable version identity);
|
||||
# fall back to any other weights-type primary.
|
||||
for file_info in version_entry.get("files", []):
|
||||
if file_info.get("type") == "Model" and file_info.get("primary"):
|
||||
hashes = file_info.get("hashes", {})
|
||||
model_hash = hashes.get("SHA256")
|
||||
model_hash = (file_info.get("hashes", {}) or {}).get("SHA256")
|
||||
if model_hash:
|
||||
return model_hash
|
||||
for file_info in version_entry.get("files", []):
|
||||
if file_info.get("type") in MODEL_WEIGHT_FILE_TYPES and file_info.get("primary"):
|
||||
model_hash = (file_info.get("hashes", {}) or {}).get("SHA256")
|
||||
if model_hash:
|
||||
return model_hash
|
||||
return None
|
||||
|
||||
@@ -18,6 +18,7 @@ from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
|
||||
from ..utils.constants import (
|
||||
CARD_PREVIEW_WIDTH,
|
||||
DIFFUSION_MODEL_BASE_MODELS,
|
||||
MODEL_WEIGHT_FILE_TYPES,
|
||||
SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS,
|
||||
VALID_LORA_TYPES,
|
||||
)
|
||||
@@ -46,6 +47,11 @@ CIVITAI_DOWNLOAD_URL_PREFIXES = (
|
||||
)
|
||||
|
||||
|
||||
# File types that are never the intended download target even when CivitAI
|
||||
# marks them primary — configs/archives/workflows are auxiliary artifacts.
|
||||
NON_DOWNLOADABLE_PRIMARY_TYPES = ("Config", "Archive", "Workflow", "Training Data")
|
||||
|
||||
|
||||
class DownloadManager:
|
||||
_instance = None
|
||||
_lock = asyncio.Lock()
|
||||
@@ -1500,7 +1506,7 @@ class DownloadManager:
|
||||
f
|
||||
for f in files
|
||||
if f.get("primary")
|
||||
and f.get("type") in ("Model", "Negative", "Diffusion Model", "UNet")
|
||||
and f.get("type") in MODEL_WEIGHT_FILE_TYPES
|
||||
),
|
||||
None,
|
||||
)
|
||||
@@ -1540,21 +1546,52 @@ class DownloadManager:
|
||||
# Fallback to primary file if no match found
|
||||
if not file_info:
|
||||
logger.debug("[download] Looking for primary file as fallback")
|
||||
# Prefer a weights-type file CivitAI marked primary; then any
|
||||
# weights-type file (providers without primary flags, e.g.
|
||||
# civarchive); then trust CivitAI's primary flag regardless of
|
||||
# type — newer types like 'Enhancement LoRA' are valid primary
|
||||
# files. Weights files are preferred over non-weights primary
|
||||
# files so a Config/Archive primary never replaces a Model.
|
||||
file_info = next(
|
||||
(
|
||||
f
|
||||
for f in files
|
||||
if f.get("primary") and f.get("type") in ("Model", "Negative", "Diffusion Model", "UNet")
|
||||
if f.get("primary") and f.get("type") in MODEL_WEIGHT_FILE_TYPES
|
||||
),
|
||||
None,
|
||||
)
|
||||
if file_info:
|
||||
logger.debug(
|
||||
"[download] Fallback primary file selected: id=%s, name=%s",
|
||||
"[download] Fallback primary file selected (primary + weights): id=%s, name=%s",
|
||||
file_info.get("id"), file_info.get("name"),
|
||||
)
|
||||
else:
|
||||
logger.debug("[download] No primary file found in fallback lookup")
|
||||
file_info = next(
|
||||
(f for f in files if f.get("type") in MODEL_WEIGHT_FILE_TYPES),
|
||||
None,
|
||||
)
|
||||
if file_info:
|
||||
logger.debug(
|
||||
"[download] Fallback primary file selected (weights type, no primary flag): id=%s, name=%s",
|
||||
file_info.get("id"), file_info.get("name"),
|
||||
)
|
||||
else:
|
||||
file_info = next(
|
||||
(
|
||||
f
|
||||
for f in files
|
||||
if f.get("primary")
|
||||
and f.get("type") not in NON_DOWNLOADABLE_PRIMARY_TYPES
|
||||
),
|
||||
None,
|
||||
)
|
||||
if file_info:
|
||||
logger.debug(
|
||||
"[download] Fallback primary file selected (trusting CivitAI primary flag): id=%s, name=%s, type=%s",
|
||||
file_info.get("id"), file_info.get("name"), file_info.get("type"),
|
||||
)
|
||||
else:
|
||||
logger.debug("[download] No primary file found in fallback lookup")
|
||||
|
||||
if not file_info:
|
||||
return {"success": False, "error": "No suitable file found in metadata"}
|
||||
|
||||
@@ -62,6 +62,20 @@ MODEL_FILE_EXTENSIONS = {
|
||||
".gguf",
|
||||
}
|
||||
|
||||
# CivitAI ModelFile.type values eligible as the main download file.
|
||||
# Mirrors CivitAI's getPrimaryFile() (model-helpers.ts): weight types are
|
||||
# preferred, but any file CivitAI marks `primary` is accepted — newer types
|
||||
# like 'Enhancement LoRA' (Anima/AIR image-editing LoRAs) are valid primary
|
||||
# files despite not being in the traditional weights allowlist.
|
||||
MODEL_WEIGHT_FILE_TYPES = (
|
||||
"Model",
|
||||
"Pruned Model",
|
||||
"Negative",
|
||||
"UNet",
|
||||
"Diffusion Model",
|
||||
"Enhancement LoRA",
|
||||
)
|
||||
|
||||
# Valid sub-types for each scanner type
|
||||
VALID_LORA_SUB_TYPES = ["lora", "locon", "dora"]
|
||||
VALID_CHECKPOINT_SUB_TYPES = ["checkpoint", "diffusion_model"]
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import { BaseContextMenu } from './BaseContextMenu.js';
|
||||
import { ModelContextMenuMixin } from './ModelContextMenuMixin.js';
|
||||
import { showToast, copyToClipboard, sendLoraToWorkflow } from '../../utils/uiHelpers.js';
|
||||
import { isModelWeightFile } from '../../utils/modelFileTypes.js';
|
||||
import { setSessionItem, removeSessionItem } from '../../utils/storageHelpers.js';
|
||||
import { updateRecipeMetadata } from '../../api/recipeApi.js';
|
||||
import { state } from '../../state/index.js';
|
||||
@@ -255,7 +256,7 @@ export class RecipeContextMenu extends BaseContextMenu {
|
||||
loras: validLoras.map(lora => {
|
||||
const civitaiInfo = lora.civitaiInfo;
|
||||
const modelFile = civitaiInfo.files ?
|
||||
civitaiInfo.files.find(file => file.type === 'Model') : null;
|
||||
civitaiInfo.files.find(file => isModelWeightFile(file.type)) : null;
|
||||
|
||||
return {
|
||||
// Basic lora info
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
// Recipe Modal Component
|
||||
import { showToast, copyToClipboard, sendLoraToWorkflow, sendModelPathToWorkflow, openCivitaiByMetadata, stripLoraTags, sendPromptToWorkflow, sendGenParamsToWorkflow } from '../utils/uiHelpers.js';
|
||||
import { isModelWeightFile } from '../utils/modelFileTypes.js';
|
||||
import { translate } from '../utils/i18nHelpers.js';
|
||||
import { state } from '../state/index.js';
|
||||
import { setSessionItem, removeSessionItem, getStorageItem, setStorageItem } from '../utils/storageHelpers.js';
|
||||
@@ -1412,7 +1413,7 @@ class RecipeModal {
|
||||
loras: validLoras.map(lora => {
|
||||
const civitaiInfo = lora.civitaiInfo;
|
||||
const modelFile = civitaiInfo.files ?
|
||||
civitaiInfo.files.find(file => file.type === 'Model') : null;
|
||||
civitaiInfo.files.find(file => isModelWeightFile(file.type)) : null;
|
||||
|
||||
return {
|
||||
// Basic lora info
|
||||
|
||||
@@ -3,6 +3,7 @@ import { showToast, setupAutoNewlineOnPaste } from '../utils/uiHelpers.js';
|
||||
import { state } from '../state/index.js';
|
||||
import { LoadingManager } from './LoadingManager.js';
|
||||
import { getModelApiClient, resetAndReload } from '../api/modelApiFactory.js';
|
||||
import { isModelWeightFile } from '../utils/modelFileTypes.js';
|
||||
import { getStorageItem, setStorageItem } from '../utils/storageHelpers.js';
|
||||
import { FolderTreeManager } from '../components/FolderTreeManager.js';
|
||||
import { translate } from '../utils/i18nHelpers.js';
|
||||
@@ -557,8 +558,7 @@ export class DownloadManager {
|
||||
const firstImage = version.images?.find(img => !img.url.endsWith('.mp4'));
|
||||
const thumbnailUrl = firstImage ? firstImage.url : '/loras_static/images/no-preview.png';
|
||||
|
||||
// Count model-type files per version
|
||||
const modelFiles = (version.files || []).filter(f => f.type === 'Model' || f.type === 'UNet' || f.type === 'Diffusion Model');
|
||||
const modelFiles = (version.files || []).filter(f => isModelWeightFile(f.type));
|
||||
const primaryFile = modelFiles.find(f => f.primary) || modelFiles[0] || {};
|
||||
const fileSize = version.modelSizeKB ?
|
||||
(version.modelSizeKB / 1024).toFixed(2) :
|
||||
@@ -685,7 +685,7 @@ export class DownloadManager {
|
||||
if (!version) return;
|
||||
|
||||
this.currentVersion = version;
|
||||
const modelFiles = (version.files || []).filter(f => f.type === 'Model' || f.type === 'UNet' || f.type === 'Diffusion Model');
|
||||
const modelFiles = (version.files || []).filter(f => isModelWeightFile(f.type));
|
||||
|
||||
document.getElementById('versionStep').style.display = 'none';
|
||||
document.getElementById('fileSelectionStep').style.display = 'block';
|
||||
@@ -747,7 +747,7 @@ export class DownloadManager {
|
||||
return;
|
||||
}
|
||||
|
||||
const modelFiles = (version.files || []).filter(f => f.type === 'Model' || f.type === 'UNet' || f.type === 'Diffusion Model');
|
||||
const modelFiles = (version.files || []).filter(f => isModelWeightFile(f.type));
|
||||
this.selectedFile = modelFiles.find(f => f.id.toString() === selectedRadio.value);
|
||||
|
||||
console.log('[download] confirmFileSelection: selected file id=%s, name="%s", type="%s", metadata=%o',
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
// CivitAI ModelFile.type values eligible as the main download file.
|
||||
// Mirrors the backend constant MODEL_WEIGHT_FILE_TYPES (py/utils/constants.py).
|
||||
// Keep both lists in sync when CivitAI introduces new file types.
|
||||
export const MODEL_WEIGHT_FILE_TYPES = [
|
||||
'Model',
|
||||
'Pruned Model',
|
||||
'Negative',
|
||||
'UNet',
|
||||
'Diffusion Model',
|
||||
'Enhancement LoRA',
|
||||
];
|
||||
|
||||
export function isModelWeightFile(type) {
|
||||
return MODEL_WEIGHT_FILE_TYPES.includes(type);
|
||||
}
|
||||
@@ -637,6 +637,316 @@ def test_conditioning_provenance_recovers_transformed_switched_prompts(
|
||||
assert params["negative_prompt"] == "expected negative"
|
||||
|
||||
|
||||
def test_conditioning_provenance_identity_switch_between_encoders(
|
||||
metadata_registry, monkeypatch
|
||||
):
|
||||
"""Lock identity-preserving switches placed directly between encoders.
|
||||
|
||||
A switch returns the selected input conditioning verbatim, so provenance
|
||||
must be recovered through object identity without any transform metadata.
|
||||
"""
|
||||
prompt_graph = {
|
||||
"encode_pos": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "chosen positive", "clip": ["clip", 0]},
|
||||
},
|
||||
"encode_other_pos": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "unchosen positive", "clip": ["clip", 0]},
|
||||
},
|
||||
"positive_switch": {
|
||||
"class_type": "ComfySwitchNode",
|
||||
"inputs": {
|
||||
"switch": True,
|
||||
"on_false": ["encode_other_pos", 0],
|
||||
"on_true": ["encode_pos", 0],
|
||||
},
|
||||
},
|
||||
"sampler": {
|
||||
"class_type": "ClownsharKSampler_Beta",
|
||||
"inputs": {
|
||||
"seed": 123,
|
||||
"steps": 8,
|
||||
"cfg": 1.0,
|
||||
"sampler_name": "linear/euler",
|
||||
"scheduler": "beta57",
|
||||
"denoise": 1.0,
|
||||
"positive": ["positive_switch", 0],
|
||||
"negative": ["encode_other_pos", 0],
|
||||
"latent_image": {
|
||||
"samples": types.SimpleNamespace(shape=(1, 4, 16, 16))
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
prompt = SimpleNamespace(original_prompt=prompt_graph)
|
||||
|
||||
chosen_conditioning = object()
|
||||
unchosen_conditioning = object()
|
||||
|
||||
monkeypatch.setattr(metadata_processor, "standalone_mode", False)
|
||||
|
||||
metadata_registry.start_collection("prompt-identity-switch")
|
||||
metadata_registry.set_current_prompt(prompt)
|
||||
|
||||
metadata_registry.record_node_execution(
|
||||
"encode_pos", "CLIPTextEncode", {"text": "chosen positive"}, None
|
||||
)
|
||||
metadata_registry.update_node_execution(
|
||||
"encode_pos", "CLIPTextEncode", [(chosen_conditioning,)]
|
||||
)
|
||||
metadata_registry.record_node_execution(
|
||||
"encode_other_pos", "CLIPTextEncode", {"text": "unchosen positive"}, None
|
||||
)
|
||||
metadata_registry.update_node_execution(
|
||||
"encode_other_pos", "CLIPTextEncode", [(unchosen_conditioning,)]
|
||||
)
|
||||
metadata_registry.record_node_execution(
|
||||
"positive_switch",
|
||||
"ComfySwitchNode",
|
||||
{
|
||||
"switch": True,
|
||||
"on_false": unchosen_conditioning,
|
||||
"on_true": chosen_conditioning,
|
||||
},
|
||||
None,
|
||||
)
|
||||
metadata_registry.update_node_execution(
|
||||
"positive_switch", "ComfySwitchNode", [(chosen_conditioning,)]
|
||||
)
|
||||
metadata_registry.record_node_execution(
|
||||
"sampler",
|
||||
"ClownsharKSampler_Beta",
|
||||
{
|
||||
"seed": 123,
|
||||
"steps": 8,
|
||||
"cfg": 1.0,
|
||||
"sampler_name": "linear/euler",
|
||||
"scheduler": "beta57",
|
||||
"denoise": 1.0,
|
||||
"positive": chosen_conditioning,
|
||||
"negative": unchosen_conditioning,
|
||||
"latent_image": {
|
||||
"samples": types.SimpleNamespace(shape=(1, 4, 16, 16))
|
||||
},
|
||||
},
|
||||
None,
|
||||
)
|
||||
|
||||
metadata = metadata_registry.get_metadata("prompt-identity-switch")
|
||||
params = MetadataProcessor.extract_generation_params(metadata)
|
||||
|
||||
assert params["prompt"] == "chosen positive"
|
||||
assert params["negative_prompt"] == "unchosen positive"
|
||||
|
||||
|
||||
def test_conditioning_provenance_ignores_scalar_conditioning_fields(
|
||||
metadata_registry, monkeypatch
|
||||
):
|
||||
"""Scalar fields like ``conditioning_strength`` must not be collected as
|
||||
conditioning objects for unregistered transform nodes."""
|
||||
monkeypatch.setattr(metadata_processor, "standalone_mode", False)
|
||||
|
||||
metadata_registry.start_collection("prompt-scalar-filter")
|
||||
metadata_registry.set_current_prompt(SimpleNamespace(original_prompt={}))
|
||||
|
||||
input_conditioning = object()
|
||||
metadata_registry.record_node_execution(
|
||||
"strength_node",
|
||||
"SomeStrengthTransform",
|
||||
{"conditioning": input_conditioning, "conditioning_strength": 0.8},
|
||||
None,
|
||||
return_types=("CONDITIONING",),
|
||||
)
|
||||
|
||||
metadata = metadata_registry.get_metadata("prompt-scalar-filter")
|
||||
assert metadata[PROMPTS]["strength_node"]["orig_conditionings"] == [
|
||||
input_conditioning
|
||||
]
|
||||
|
||||
|
||||
def test_conditioning_provenance_selector_with_conditioning_named_inputs(
|
||||
metadata_registry, monkeypatch
|
||||
):
|
||||
"""An identity selector whose inputs use ``conditioning*`` names must not
|
||||
leak the unselected branch's prompt."""
|
||||
prompt_graph = {
|
||||
"encode_a": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "AAA", "clip": ["clip", 0]},
|
||||
},
|
||||
"encode_b": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "BBB", "clip": ["clip", 0]},
|
||||
},
|
||||
"selector": {
|
||||
"class_type": "ConditioningSelector",
|
||||
"inputs": {
|
||||
"conditioning_a": ["encode_a", 0],
|
||||
"conditioning_b": ["encode_b", 0],
|
||||
},
|
||||
},
|
||||
"sampler": {
|
||||
"class_type": "ClownsharKSampler_Beta",
|
||||
"inputs": {
|
||||
"seed": 123,
|
||||
"steps": 8,
|
||||
"cfg": 1.0,
|
||||
"sampler_name": "linear/euler",
|
||||
"scheduler": "beta57",
|
||||
"denoise": 1.0,
|
||||
"positive": ["selector", 0],
|
||||
"negative": ["encode_b", 0],
|
||||
"latent_image": {
|
||||
"samples": types.SimpleNamespace(shape=(1, 4, 16, 16))
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
prompt = SimpleNamespace(original_prompt=prompt_graph)
|
||||
|
||||
conditioning_a = object()
|
||||
conditioning_b = object()
|
||||
|
||||
monkeypatch.setattr(metadata_processor, "standalone_mode", False)
|
||||
|
||||
metadata_registry.start_collection("prompt-selector")
|
||||
metadata_registry.set_current_prompt(prompt)
|
||||
|
||||
metadata_registry.record_node_execution(
|
||||
"encode_a", "CLIPTextEncode", {"text": "AAA"}, None
|
||||
)
|
||||
metadata_registry.update_node_execution(
|
||||
"encode_a", "CLIPTextEncode", [(conditioning_a,)]
|
||||
)
|
||||
metadata_registry.record_node_execution(
|
||||
"encode_b", "CLIPTextEncode", {"text": "BBB"}, None
|
||||
)
|
||||
metadata_registry.update_node_execution(
|
||||
"encode_b", "CLIPTextEncode", [(conditioning_b,)]
|
||||
)
|
||||
metadata_registry.record_node_execution(
|
||||
"selector",
|
||||
"ConditioningSelector",
|
||||
{"conditioning_a": conditioning_a, "conditioning_b": conditioning_b},
|
||||
None,
|
||||
return_types=("CONDITIONING",),
|
||||
)
|
||||
metadata_registry.update_node_execution(
|
||||
"selector", "ConditioningSelector", [(conditioning_a,)],
|
||||
return_types=("CONDITIONING",),
|
||||
)
|
||||
metadata_registry.record_node_execution(
|
||||
"sampler",
|
||||
"ClownsharKSampler_Beta",
|
||||
{
|
||||
"seed": 123,
|
||||
"steps": 8,
|
||||
"cfg": 1.0,
|
||||
"sampler_name": "linear/euler",
|
||||
"scheduler": "beta57",
|
||||
"denoise": 1.0,
|
||||
"positive": conditioning_a,
|
||||
"negative": conditioning_b,
|
||||
"latent_image": {
|
||||
"samples": types.SimpleNamespace(shape=(1, 4, 16, 16))
|
||||
},
|
||||
},
|
||||
None,
|
||||
)
|
||||
|
||||
metadata = metadata_registry.get_metadata("prompt-selector")
|
||||
params = MetadataProcessor.extract_generation_params(metadata)
|
||||
|
||||
assert params["prompt"] == "AAA"
|
||||
assert params["negative_prompt"] == "BBB"
|
||||
|
||||
|
||||
def test_conditioning_provenance_uses_conditioning_output_slot(
|
||||
metadata_registry, monkeypatch
|
||||
):
|
||||
"""Unregistered nodes whose CONDITIONING output is not the first slot
|
||||
must still be tracked through the correct output position.
|
||||
|
||||
The graph's conditioning chain ends at an unexecuted phantom node so the
|
||||
topology fallback in extract_generation_params cannot mask a runtime
|
||||
provenance failure.
|
||||
"""
|
||||
prompt_graph = {
|
||||
"diag_node": {
|
||||
"class_type": "DiagThenCond",
|
||||
"inputs": {"conditioning": ["phantom_source", 0]},
|
||||
},
|
||||
"sampler": {
|
||||
"class_type": "ClownsharKSampler_Beta",
|
||||
"inputs": {
|
||||
"seed": 123,
|
||||
"steps": 8,
|
||||
"cfg": 1.0,
|
||||
"sampler_name": "linear/euler",
|
||||
"scheduler": "beta57",
|
||||
"denoise": 1.0,
|
||||
"positive": ["diag_node", 1],
|
||||
"latent_image": {
|
||||
"samples": types.SimpleNamespace(shape=(1, 4, 16, 16))
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
prompt = SimpleNamespace(original_prompt=prompt_graph)
|
||||
|
||||
input_conditioning = object()
|
||||
transformed_conditioning = object()
|
||||
|
||||
monkeypatch.setattr(metadata_processor, "standalone_mode", False)
|
||||
|
||||
metadata_registry.start_collection("prompt-output-slot")
|
||||
metadata_registry.set_current_prompt(prompt)
|
||||
|
||||
metadata_registry.record_node_execution(
|
||||
"encode_pos", "CLIPTextEncode", {"text": "AAA"}, None
|
||||
)
|
||||
metadata_registry.update_node_execution(
|
||||
"encode_pos", "CLIPTextEncode", [(input_conditioning,)]
|
||||
)
|
||||
metadata_registry.record_node_execution(
|
||||
"diag_node",
|
||||
"DiagThenCond",
|
||||
{"conditioning": input_conditioning},
|
||||
None,
|
||||
return_types=("STRING", "CONDITIONING"),
|
||||
)
|
||||
metadata_registry.update_node_execution(
|
||||
"diag_node",
|
||||
"DiagThenCond",
|
||||
[("diagnostics", transformed_conditioning)],
|
||||
return_types=("STRING", "CONDITIONING"),
|
||||
)
|
||||
metadata_registry.record_node_execution(
|
||||
"sampler",
|
||||
"ClownsharKSampler_Beta",
|
||||
{
|
||||
"seed": 123,
|
||||
"steps": 8,
|
||||
"cfg": 1.0,
|
||||
"sampler_name": "linear/euler",
|
||||
"scheduler": "beta57",
|
||||
"denoise": 1.0,
|
||||
"positive": transformed_conditioning,
|
||||
"negative": input_conditioning,
|
||||
"latent_image": {
|
||||
"samples": types.SimpleNamespace(shape=(1, 4, 16, 16))
|
||||
},
|
||||
},
|
||||
None,
|
||||
)
|
||||
|
||||
metadata = metadata_registry.get_metadata("prompt-output-slot")
|
||||
params = MetadataProcessor.extract_generation_params(metadata)
|
||||
|
||||
assert params["prompt"] == "AAA"
|
||||
|
||||
|
||||
def test_conditioning_provenance_recovers_kj_set_get_prompts(
|
||||
metadata_registry, monkeypatch
|
||||
):
|
||||
|
||||
@@ -123,10 +123,7 @@ def metadata_provider(monkeypatch):
|
||||
class DummyProvider:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
async def get_model_version(self, model_id, model_version_id):
|
||||
self.calls.append((model_id, model_version_id))
|
||||
return {
|
||||
self.payload = {
|
||||
"id": 42,
|
||||
"model": {"type": "LoRA", "tags": ["fantasy"]},
|
||||
"baseModel": "BaseModel",
|
||||
@@ -141,6 +138,10 @@ def metadata_provider(monkeypatch):
|
||||
],
|
||||
}
|
||||
|
||||
async def get_model_version(self, model_id, model_version_id):
|
||||
self.calls.append((model_id, model_version_id))
|
||||
return self.payload
|
||||
|
||||
provider = DummyProvider()
|
||||
monkeypatch.setattr(
|
||||
download_manager,
|
||||
@@ -233,6 +234,165 @@ async def test_successful_download_uses_defaults(
|
||||
assert captured["download_urls"] == ["https://example.invalid/file.safetensors"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_accepts_enhancement_lora_primary_file(
|
||||
monkeypatch, scanners, metadata_provider, tmp_path
|
||||
):
|
||||
"""A version whose only file has type 'Enhancement LoRA' (Anima/AIR
|
||||
image-editing LoRAs) must download — previously failed with
|
||||
"No suitable file found in metadata" because the type was missing from
|
||||
the primary-file weights allowlist."""
|
||||
manager = DownloadManager()
|
||||
metadata_provider.payload = {
|
||||
"id": 3219121,
|
||||
"model": {"type": "LORA", "tags": ["style"]},
|
||||
"baseModel": "Anima",
|
||||
"creator": {"username": "Deskup"},
|
||||
"files": [
|
||||
{
|
||||
"id": 3100968,
|
||||
"type": "Enhancement LoRA",
|
||||
"primary": True,
|
||||
"name": "deskup-anima-edit-general.safetensors",
|
||||
"sizeKB": 358501.13,
|
||||
"downloadUrl": "https://example.invalid/deskup-anima-edit-general.safetensors",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_execute_download(self, **kwargs):
|
||||
captured.update(
|
||||
{
|
||||
"download_urls": kwargs["download_urls"],
|
||||
"model_type": kwargs["model_type"],
|
||||
}
|
||||
)
|
||||
return {"success": True}
|
||||
|
||||
monkeypatch.setattr(
|
||||
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
||||
)
|
||||
|
||||
result = await manager.download_from_civitai(
|
||||
model_id=2850692,
|
||||
model_version_id=3219121,
|
||||
save_dir=str(tmp_path),
|
||||
use_default_paths=True,
|
||||
progress_callback=None,
|
||||
source=None,
|
||||
)
|
||||
|
||||
assert result["success"] is True
|
||||
assert captured["model_type"] == "lora"
|
||||
assert captured["download_urls"] == [
|
||||
"https://example.invalid/deskup-anima-edit-general.safetensors"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_falls_back_to_civitai_primary_flag_regardless_of_type(
|
||||
monkeypatch, scanners, metadata_provider, tmp_path
|
||||
):
|
||||
"""If no weights-type file exists, trust CivitAI's `primary` flag on any
|
||||
file — mirrors CivitAI's getPrimaryFile() which never excludes a file by
|
||||
type."""
|
||||
manager = DownloadManager()
|
||||
metadata_provider.payload = {
|
||||
"id": 77,
|
||||
"model": {"type": "LORA", "tags": ["concept"]},
|
||||
"baseModel": "Anima",
|
||||
"creator": {"username": "Author"},
|
||||
"files": [
|
||||
{
|
||||
"id": 100,
|
||||
"type": "Other",
|
||||
"primary": True,
|
||||
"name": "custom-type-lora.safetensors",
|
||||
"downloadUrl": "https://example.invalid/custom-type-lora.safetensors",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_execute_download(self, **kwargs):
|
||||
captured["download_urls"] = kwargs["download_urls"]
|
||||
return {"success": True}
|
||||
|
||||
monkeypatch.setattr(
|
||||
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
||||
)
|
||||
|
||||
result = await manager.download_from_civitai(
|
||||
model_version_id=77,
|
||||
save_dir=str(tmp_path),
|
||||
use_default_paths=True,
|
||||
progress_callback=None,
|
||||
source=None,
|
||||
)
|
||||
|
||||
assert result["success"] is True
|
||||
assert captured["download_urls"] == [
|
||||
"https://example.invalid/custom-type-lora.safetensors"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_prefers_weights_file_over_non_weights_primary(
|
||||
monkeypatch, scanners, metadata_provider, tmp_path
|
||||
):
|
||||
"""A Config/Archive-type primary must never replace an existing weights
|
||||
file — the weights file wins even without the primary flag."""
|
||||
manager = DownloadManager()
|
||||
metadata_provider.payload = {
|
||||
"id": 78,
|
||||
"model": {"type": "LORA", "tags": ["concept"]},
|
||||
"baseModel": "BaseModel",
|
||||
"creator": {"username": "Author"},
|
||||
"files": [
|
||||
{
|
||||
"id": 201,
|
||||
"type": "Config",
|
||||
"primary": True,
|
||||
"name": "config.json",
|
||||
"downloadUrl": "https://example.invalid/config.json",
|
||||
},
|
||||
{
|
||||
"id": 202,
|
||||
"type": "Model",
|
||||
"primary": False,
|
||||
"name": "weights.safetensors",
|
||||
"downloadUrl": "https://example.invalid/weights.safetensors",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_execute_download(self, **kwargs):
|
||||
captured["download_urls"] = kwargs["download_urls"]
|
||||
return {"success": True}
|
||||
|
||||
monkeypatch.setattr(
|
||||
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
||||
)
|
||||
|
||||
result = await manager.download_from_civitai(
|
||||
model_version_id=78,
|
||||
save_dir=str(tmp_path),
|
||||
use_default_paths=True,
|
||||
progress_callback=None,
|
||||
source=None,
|
||||
)
|
||||
|
||||
assert result["success"] is True
|
||||
assert captured["download_urls"] == [
|
||||
"https://example.invalid/weights.safetensors"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_keeps_save_dir_when_use_save_dir_as_root(
|
||||
monkeypatch, scanners, metadata_provider, tmp_path
|
||||
|
||||
Reference in New Issue
Block a user