Compare commits

...

3 Commits

Author SHA1 Message Date
Will Miao 303cca0d85 fix(download): accept newer CivitAI file types for primary file selection
Downloads failed with "No suitable file found in metadata" for models whose
only file uses newer CivitAI file types (e.g. 'Enhancement LoRA' for
Anima/AIR image-editing LoRAs) because the primary-file allowlist only
covered legacy types.

- unify the weights-type allowlist as MODEL_WEIGHT_FILE_TYPES
  (py/utils/constants.py) and apply it across download, recipe and
  metadata-refresh lookups
- mirror CivitAI's getPrimaryFile() semantics: prefer weights-type primary,
  fall back to weights files, then trust CivitAI's primary flag (excluding
  non-downloadable artifacts like Config/Archive/Workflow)
- mirror the allowlist in the frontend via shared isModelWeightFile() helper
- add regression tests for the Enhancement LoRA primary-file download,
  primary-flag fallback and weights-over-non-weights-primary preference
2026-08-12 21:14:23 +08:00
Will Miao c2f16784b3 fix(metadata): keep identity selectors from leaking unselected prompts 2026-08-12 19:44:43 +08:00
Will Miao 5bc6d8286c fix(metadata): exclude scalar fields from conditioning provenance inputs 2026-08-12 19:18:46 +08:00
12 changed files with 620 additions and 34 deletions
+43 -13
View File
@@ -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
View File
@@ -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,
)
+2 -1
View File
@@ -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."""
+9 -2
View File
@@ -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
+41 -4
View File
@@ -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"}
+14
View File
@@ -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
+2 -1
View File
@@ -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
+4 -4
View File
@@ -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',
+15
View File
@@ -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
):
+164 -4
View File
@@ -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