Will Miao
0b6348be5c
Importing an image whose metadata references a model by CivitAI id (ComfyUI URN names, A1111 "Civitai resources", recipe checkpoint blocks) always queried the metadata provider, even when the model was already in the library. That made large imports network-bound and prone to vendor rate limiting for no benefit. Parsers now resolve the id against the scanner's version index first and fill the entry from the local cache (model name, base model, hash, preview, size). The provider is only consulted on a miss, matching the existing hash-based local-first path. - RecipeScanner: get_local_lora_by_version_id / get_local_checkpoint_by_version_id - RecipeMetadataParser: populate_checkpoint_from_local + resolver helpers - comfy / automatic / recipe_format parsers use the local lookup first
451 lines
15 KiB
Python
451 lines
15 KiB
Python
import pytest
|
|
import json
|
|
from py.recipes.parsers.comfy import ComfyMetadataParser
|
|
|
|
|
|
class LocalRecipeScanner:
|
|
class LoraScanner:
|
|
@staticmethod
|
|
def has_hash(model_hash):
|
|
return False
|
|
|
|
def __init__(self, models, models_by_version=None, checkpoints_by_version=None):
|
|
self.models = models
|
|
self.models_by_version = models_by_version or {}
|
|
self.checkpoints_by_version = checkpoints_by_version or {}
|
|
self.queries = []
|
|
self.base_models = []
|
|
self.version_queries = []
|
|
self.checkpoint_version_queries = []
|
|
self._lora_scanner = self.LoraScanner()
|
|
|
|
async def get_local_lora(self, name, base_model=None):
|
|
self.queries.append(name)
|
|
self.base_models.append(base_model)
|
|
return self.models.get(name)
|
|
|
|
async def get_local_lora_by_version_id(self, version_id):
|
|
self.version_queries.append(version_id)
|
|
return self.models_by_version.get(int(version_id))
|
|
|
|
async def get_local_checkpoint_by_version_id(self, version_id):
|
|
self.checkpoint_version_queries.append(version_id)
|
|
return self.checkpoints_by_version.get(int(version_id))
|
|
|
|
|
|
def local_lora(file_name):
|
|
return {
|
|
"file_path": f"/models/loras/{file_name}.safetensors",
|
|
"file_name": file_name.rsplit("/", 1)[-1],
|
|
"model_name": file_name.rsplit("/", 1)[-1],
|
|
"sha256": file_name[0] * 64,
|
|
"size": 4096,
|
|
"base_model": "SDXL 1.0",
|
|
"preview_url": "",
|
|
"civitai": None,
|
|
}
|
|
|
|
|
|
def local_lora_for_version(model_version_id):
|
|
return {
|
|
"file_path": "/models/loras/urn_local_lora.safetensors",
|
|
"file_name": "urn_local_lora",
|
|
"model_name": "URN Local LoRA",
|
|
"sha256": "e" * 64,
|
|
"size": 8192,
|
|
"base_model": "SDXL 1.0",
|
|
"preview_url": "/models/loras/urn_local_lora.preview.png",
|
|
"civitai": {"id": model_version_id, "modelId": 123, "name": "v2"},
|
|
}
|
|
|
|
|
|
def local_checkpoint_for_version(model_version_id):
|
|
return {
|
|
"file_path": "/models/checkpoints/juggernaut_xl.safetensors",
|
|
"file_name": "juggernaut_xl",
|
|
"model_name": "Juggernaut XL",
|
|
"sha256": "d" * 64,
|
|
"size": 6_000_000_000,
|
|
"base_model": "SDXL 1.0",
|
|
"preview_url": "/models/checkpoints/juggernaut_xl.preview.png",
|
|
"civitai": {"id": model_version_id, "modelId": 1908679, "name": "v9"},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parse_metadata_without_loras(monkeypatch):
|
|
checkpoint_info = {
|
|
"id": 2224012,
|
|
"modelId": 1908679,
|
|
"model": {"name": "SDXL Checkpoint", "type": "checkpoint"},
|
|
"name": "v1.0",
|
|
"images": [{"url": "https://image.civitai.com/checkpoints/original=true"}],
|
|
"baseModel": "sdxl",
|
|
"downloadUrl": "https://civitai.com/api/download/checkpoint",
|
|
}
|
|
|
|
async def fake_metadata_provider():
|
|
class Provider:
|
|
async def get_model_version_info(self, version_id):
|
|
assert version_id == "2224012"
|
|
return checkpoint_info, None
|
|
return Provider()
|
|
|
|
monkeypatch.setattr(
|
|
"py.recipes.parsers.comfy.get_default_metadata_provider",
|
|
fake_metadata_provider,
|
|
)
|
|
|
|
parser = ComfyMetadataParser()
|
|
|
|
# User provided metadata
|
|
metadata_json = {
|
|
"resource-stack": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "urn:air:sdxl:checkpoint:civitai:1908679@2224012"}
|
|
},
|
|
"6": {
|
|
"class_type": "smZ CLIPTextEncode",
|
|
"inputs": {"text": "Positive prompt content"},
|
|
"_meta": {"title": "Positive"}
|
|
},
|
|
"7": {
|
|
"class_type": "smZ CLIPTextEncode",
|
|
"inputs": {"text": "Negative prompt content"},
|
|
"_meta": {"title": "Negative"}
|
|
},
|
|
"11": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"sampler_name": "euler_ancestral",
|
|
"scheduler": "normal",
|
|
"seed": 904124997,
|
|
"steps": 35,
|
|
"cfg": 6,
|
|
"denoise": 0.1,
|
|
"model": ["resource-stack", 0],
|
|
"positive": ["6", 0],
|
|
"negative": ["7", 0],
|
|
"latent_image": ["21", 0]
|
|
},
|
|
"_meta": {"title": "KSampler"}
|
|
},
|
|
"extraMetadata": json.dumps({
|
|
"prompt": "One woman, (solo:1.3), ...",
|
|
"negativePrompt": "lowres, worst quality, ...",
|
|
"steps": 35,
|
|
"cfgScale": 6,
|
|
"sampler": "euler_ancestral",
|
|
"seed": 904124997,
|
|
"width": 1024,
|
|
"height": 1024
|
|
})
|
|
}
|
|
|
|
result = await parser.parse_metadata(json.dumps(metadata_json))
|
|
|
|
assert "error" not in result
|
|
assert result["loras"] == []
|
|
assert result["checkpoint"] is not None
|
|
assert int(result["checkpoint"]["modelId"]) == 1908679
|
|
assert int(result["checkpoint"]["id"]) == 2224012
|
|
assert result["gen_params"]["prompt"] == "One woman, (solo:1.3), ..."
|
|
assert result["gen_params"]["steps"] == 35
|
|
assert result["gen_params"]["size"] == "1024x1024"
|
|
assert result["from_comfy_metadata"] is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parse_metadata_resolves_standard_and_manager_local_loras(monkeypatch):
|
|
async def fake_metadata_provider():
|
|
class Provider:
|
|
async def get_model_version_info(self, version_id):
|
|
raise AssertionError("Local LoRAs must not query Civitai")
|
|
|
|
return Provider()
|
|
|
|
monkeypatch.setattr(
|
|
"py.recipes.parsers.comfy.get_default_metadata_provider",
|
|
fake_metadata_provider,
|
|
)
|
|
scanner = LocalRecipeScanner({
|
|
"styles/standard.safetensors": local_lora("standard"),
|
|
"manager": local_lora("manager"),
|
|
})
|
|
metadata_json = {
|
|
"1": {
|
|
"class_type": "LoraLoader",
|
|
"inputs": {
|
|
"lora_name": "styles/standard.safetensors",
|
|
"strength_model": 0.55,
|
|
},
|
|
},
|
|
"2": {
|
|
"class_type": "LoraLoaderLM",
|
|
"inputs": {
|
|
"loras": {
|
|
"__value__": [
|
|
{"name": "manager", "strength": "0.80", "active": True},
|
|
{"name": "disabled", "strength": 1.0, "active": False},
|
|
{"name": "dummy", "strength": 1.0, "active": True, "_isDummy": True},
|
|
]
|
|
}
|
|
},
|
|
},
|
|
}
|
|
|
|
result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json), scanner)
|
|
|
|
assert [entry["file_name"] for entry in result["loras"]] == ["standard", "manager"]
|
|
assert [entry["weight"] for entry in result["loras"]] == [0.55, 0.8]
|
|
assert all(isinstance(entry["weight"], float) for entry in result["loras"])
|
|
assert all(entry["existsLocally"] is True for entry in result["loras"])
|
|
assert all(entry["isDeleted"] is False for entry in result["loras"])
|
|
assert scanner.queries == ["styles/standard.safetensors", "manager"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parse_metadata_defaults_malformed_weight_and_passes_checkpoint_base_model(monkeypatch):
|
|
checkpoint_info = {
|
|
"id": 456,
|
|
"modelId": 123,
|
|
"model": {"name": "Checkpoint", "type": "checkpoint"},
|
|
"name": "v1",
|
|
"baseModel": "SDXL 1.0",
|
|
}
|
|
|
|
async def fake_metadata_provider():
|
|
class Provider:
|
|
async def get_model_version_info(self, version_id):
|
|
return checkpoint_info, None
|
|
|
|
return Provider()
|
|
|
|
monkeypatch.setattr(
|
|
"py.recipes.parsers.comfy.get_default_metadata_provider",
|
|
fake_metadata_provider,
|
|
)
|
|
scanner = LocalRecipeScanner({"style": local_lora("style")})
|
|
metadata_json = {
|
|
"1": {"class_type": "LoraLoader", "inputs": {"lora_name": "style", "strength_model": "invalid"}},
|
|
"2": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "civitai:123@456"}},
|
|
}
|
|
|
|
result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json), scanner)
|
|
|
|
assert result["loras"][0]["weight"] == 1.0
|
|
assert scanner.base_models == ["SDXL 1.0"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parse_metadata_keeps_civitai_urn_with_local_lora(monkeypatch):
|
|
remote_info = {
|
|
"id": 456,
|
|
"modelId": 123,
|
|
"model": {"name": "Remote LoRA", "type": "LORA"},
|
|
"name": "v1",
|
|
"files": [
|
|
{
|
|
"type": "Model",
|
|
"primary": True,
|
|
"name": "remote.safetensors",
|
|
"hashes": {"SHA256": "c" * 64},
|
|
}
|
|
],
|
|
}
|
|
|
|
async def fake_metadata_provider():
|
|
class Provider:
|
|
async def get_model_version_info(self, version_id):
|
|
assert version_id == "456"
|
|
return remote_info, None
|
|
|
|
return Provider()
|
|
|
|
monkeypatch.setattr(
|
|
"py.recipes.parsers.comfy.get_default_metadata_provider",
|
|
fake_metadata_provider,
|
|
)
|
|
scanner = LocalRecipeScanner({"local": local_lora("local")})
|
|
metadata_json = {
|
|
"1": {
|
|
"class_type": "LoraLoader",
|
|
"inputs": {"lora_name": "local", "strength_model": 0.4},
|
|
},
|
|
"2": {
|
|
"class_type": "LoraLoader",
|
|
"inputs": {
|
|
"lora_name": "urn:air:sdxl:lora:civitai:123@456",
|
|
"strength_model": 0.9,
|
|
},
|
|
},
|
|
}
|
|
|
|
result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json), scanner)
|
|
|
|
assert [entry["name"] for entry in result["loras"]] == ["local", "Remote LoRA"]
|
|
assert [entry["weight"] for entry in result["loras"]] == [0.4, 0.9]
|
|
assert result["loras"][0]["existsLocally"] is True
|
|
assert result["loras"][1]["id"] == 456
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parse_metadata_without_extra_metadata(monkeypatch):
|
|
async def fake_metadata_provider():
|
|
class Provider:
|
|
async def get_model_version_info(self, version_id):
|
|
return {"model": {"name": "Test"}, "id": version_id}, None
|
|
return Provider()
|
|
|
|
monkeypatch.setattr(
|
|
"py.recipes.parsers.comfy.get_default_metadata_provider",
|
|
fake_metadata_provider,
|
|
)
|
|
|
|
parser = ComfyMetadataParser()
|
|
|
|
metadata_json = {
|
|
"node_1": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "urn:air:sdxl:checkpoint:civitai:123@456"}
|
|
}
|
|
}
|
|
|
|
result = await parser.parse_metadata(json.dumps(metadata_json))
|
|
|
|
assert "error" not in result
|
|
assert result["loras"] == []
|
|
assert result["checkpoint"]["id"] == "456"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parse_metadata_with_list_ckpt_name(monkeypatch):
|
|
"""ckpt_name serialized as a single-element list must not crash re.search."""
|
|
checkpoint_info = {
|
|
"id": 456,
|
|
"modelId": 123,
|
|
"model": {"name": "Checkpoint", "type": "checkpoint"},
|
|
"name": "v1",
|
|
"baseModel": "SDXL 1.0",
|
|
}
|
|
|
|
async def fake_metadata_provider():
|
|
class Provider:
|
|
async def get_model_version_info(self, version_id):
|
|
assert version_id == "456"
|
|
return checkpoint_info, None
|
|
|
|
return Provider()
|
|
|
|
monkeypatch.setattr(
|
|
"py.recipes.parsers.comfy.get_default_metadata_provider",
|
|
fake_metadata_provider,
|
|
)
|
|
|
|
metadata_json = {
|
|
"1": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": ["urn:air:sdxl:checkpoint:civitai:123@456"]},
|
|
}
|
|
}
|
|
|
|
result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json))
|
|
|
|
assert "error" not in result
|
|
assert result["checkpoint"] is not None
|
|
assert int(result["checkpoint"]["id"]) == 456
|
|
assert int(result["checkpoint"]["modelId"]) == 123
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parse_metadata_with_none_ckpt_name(monkeypatch):
|
|
"""Missing (None) ckpt_name must not crash re.search with a TypeError."""
|
|
|
|
async def fake_metadata_provider():
|
|
class Provider:
|
|
async def get_model_version_info(self, version_id):
|
|
raise AssertionError("Checkpoint lookup must be skipped")
|
|
|
|
return Provider()
|
|
|
|
monkeypatch.setattr(
|
|
"py.recipes.parsers.comfy.get_default_metadata_provider",
|
|
fake_metadata_provider,
|
|
)
|
|
|
|
metadata_json = {
|
|
"1": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": None},
|
|
}
|
|
}
|
|
|
|
result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json))
|
|
|
|
assert "error" not in result
|
|
assert result["checkpoint"] is None
|
|
assert result["loras"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parse_metadata_resolves_urn_models_from_local_library(monkeypatch):
|
|
"""URN models that exist locally must not trigger a CivitAI request."""
|
|
|
|
async def failing_metadata_provider():
|
|
class Provider:
|
|
async def get_model_version_info(self, version_id):
|
|
raise AssertionError(
|
|
f"Unexpected CivitAI lookup for local model {version_id}"
|
|
)
|
|
|
|
return Provider()
|
|
|
|
monkeypatch.setattr(
|
|
"py.recipes.parsers.comfy.get_default_metadata_provider",
|
|
failing_metadata_provider,
|
|
)
|
|
|
|
scanner = LocalRecipeScanner(
|
|
{},
|
|
models_by_version={456: local_lora_for_version(456)},
|
|
checkpoints_by_version={2224012: local_checkpoint_for_version(2224012)},
|
|
)
|
|
metadata_json = {
|
|
"1": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {
|
|
"ckpt_name": "urn:air:sdxl:checkpoint:civitai:1908679@2224012"
|
|
},
|
|
},
|
|
"2": {
|
|
"class_type": "LoraLoader",
|
|
"inputs": {
|
|
"lora_name": "urn:air:sdxl:lora:civitai:123@456",
|
|
"strength_model": 0.8,
|
|
},
|
|
},
|
|
}
|
|
|
|
result = await ComfyMetadataParser().parse_metadata(
|
|
json.dumps(metadata_json), scanner
|
|
)
|
|
|
|
assert scanner.checkpoint_version_queries == ["2224012"]
|
|
assert scanner.version_queries == ["456"]
|
|
|
|
checkpoint = result["checkpoint"]
|
|
assert checkpoint["existsLocally"] is True
|
|
assert checkpoint["name"] == "Juggernaut XL"
|
|
assert checkpoint["file_name"] == "juggernaut_xl"
|
|
assert checkpoint["baseModel"] == "SDXL 1.0"
|
|
assert checkpoint["hash"] == "d" * 64
|
|
assert checkpoint["localPath"] == "/models/checkpoints/juggernaut_xl.safetensors"
|
|
|
|
lora = result["loras"][0]
|
|
assert lora["existsLocally"] is True
|
|
assert lora["name"] == "URN Local LoRA"
|
|
assert lora["file_name"] == "urn_local_lora"
|
|
assert lora["weight"] == 0.8
|
|
assert lora["localPath"] == "/models/loras/urn_local_lora.safetensors"
|
|
assert result["base_model"] == "SDXL 1.0"
|
|
|