willmiao-ComfyUI-Lora-Manager/tests/services/test_automatic_metadata_parser.py

import pytest

from py.recipes.parsers.automatic import AutomaticMetadataParser


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.hash_queries = []
        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)
        return self.models.get(name)

    async def get_local_lora_by_hash(self, hash_value):
        self.hash_queries.append(hash_value)
        return next((model for model in self.models.values() if model.get("sha256") == hash_value), None)

    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="local_only"):
    return {
        "file_path": f"/models/loras/styles/{file_name}.safetensors",
        "file_name": file_name,
        "model_name": "Local Only",
        "sha256": "a" * 64,
        "size": 123456,
        "base_model": "Flux.1 D",
        "preview_url": f"/models/loras/styles/{file_name}.preview.png",
        "civitai": None,
    }


@pytest.mark.asyncio
async def test_parse_metadata_extracts_checkpoint_from_civitai_resources(monkeypatch):
    checkpoint_info = {
        "id": 2442439,
        "modelId": 123456,
        "model": {"name": "Z Image", "type": "checkpoint"},
        "name": "Turbo",
        "images": [{"url": "https://image.civitai.com/checkpoints/original=true"}],
        "baseModel": "sdxl",
        "downloadUrl": "https://civitai.com/api/download/checkpoint",
        "files": [
            {
                "type": "Model",
                "primary": True,
                "sizeKB": 2048,
                "name": "Z_Image_Turbo.safetensors",
                "hashes": {"SHA256": "ABC123FF"},
            }
        ],
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_version_info(self, version_id):
                assert version_id == "2442439"
                return checkpoint_info, None

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )

    parser = AutomaticMetadataParser()

    metadata_text = (
        "Negative space, fog, BLACK blue color GRADIENT BACKGROUND, a vintage car in the middle, "
        "FOG, and a silhouetted figure near the car, in the style of the Blade Runner movie "
        "Negative prompt: Steps: 23, Sampler: Undefined, CFG scale: 3.5, Seed: 1760020955, "
        "Size: 832x1216, Clip skip: 2, Created Date: 2025-11-28T09:18:43.5269343Z, "
        'Civitai resources: [{"type":"checkpoint","modelVersionId":2442439,"modelName":"Z Image","modelVersionName":"Turbo"}], '
        "Civitai metadata: {}"
    )

    result = await parser.parse_metadata(metadata_text)

    checkpoint = result.get("checkpoint")
    assert checkpoint is not None
    assert checkpoint["name"] == "Z Image"
    assert checkpoint["version"] == "Turbo"
    assert checkpoint["type"] == "checkpoint"
    assert checkpoint["modelId"] == 123456
    assert checkpoint["hash"] == "abc123ff"
    assert checkpoint["file_name"] == "Z_Image_Turbo"
    assert checkpoint["thumbnailUrl"].endswith("width=450,optimized=true")
    assert result["model"] == checkpoint
    assert result["base_model"] == "sdxl"
    assert result["loras"] == []


@pytest.mark.asyncio
async def test_parse_metadata_merges_lora_hashes_over_empty_hashes_json(monkeypatch):
    """When Hashes JSON has empty lora hashes but Lora hashes text field has
    real ones, the real hashes should be used and those LoRAs resolved
    correctly; entries with empty hashes in both sources should be skipped."""
    lora_version_info = {
        "id": 947620,
        "modelId": 98765,
        "model": {"name": "cfg_scale_boost", "type": "LORA"},
        "name": "v1",
        "images": [{"url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/original=true"}],
        "baseModel": "illustrious",
        "downloadUrl": "https://civitai.com/api/download/models/947620",
        "files": [
            {
                "type": "Model",
                "primary": True,
                "sizeKB": 1024,
                "name": "cfg_scale_boost.safetensors",
                "hashes": {"SHA256": "4605b2de07"},
            }
        ],
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_by_hash(self, model_hash):
                assert model_hash == "4605b2de07"
                return lora_version_info, None

            async def get_model_version_info(self, version_id):
                raise AssertionError("get_model_version_info should not be called")

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )

    parser = AutomaticMetadataParser()

    metadata_text = (
        "a cyberpunk portrait <lora:cfg_scale_boost:0.6>\n"
        "Negative prompt: low quality\n"
        "Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 123456, Size: 512x768, "
        "Model hash: abc123, Model: test.safetensors, "
        'Lora hashes: "cfg_scale_boost: 4605b2de07, EmptyLora: ", '
        'Hashes: {"model": "abc123", "lora:cfg_scale_boost": "", "lora:EmptyLora": "", "lora:UnusedLora": ""}'
    )

    result = await parser.parse_metadata(metadata_text)

    # cfg_scale_boost should be resolved (hash from Lora hashes overrode empty Hashes JSON)
    loras = result.get("loras", [])
    assert len(loras) == 1, f"Expected 1 LoRA, got {len(loras)}"
    lora = loras[0]
    assert lora["name"] == "cfg_scale_boost", f"Expected cfg_scale_boost, got {lora['name']}"
    assert lora["hash"] == "4605b2de07", f"Expected hash 4605b2de07, got {lora['hash']}"
    assert lora.get("isDeleted") in (None, False), f"LoRA should not be deleted"
    assert lora["weight"] == 0.6, f"Expected weight 0.6, got {lora['weight']}"

    # EmptyLora and UnusedLora should be skipped (no hash in either source)
    lora_names = [l["name"] for l in loras]
    assert "EmptyLora" not in lora_names, "EmptyLora should have been skipped"
    assert "UnusedLora" not in lora_names, "UnusedLora should have been skipped"


@pytest.mark.asyncio
async def test_parse_metadata_lora_hashes_override_conflicting_hashes_json(monkeypatch):
    """When Hashes JSON carries a non-empty but stale hash and the Lora
    hashes text field carries the real 12-char AutoV3 hash, the Lora hashes
    value must win: CivitAI is queried with it and the entry is resolved
    instead of being poisoned by the stale hash."""
    lora_version_info = {
        "id": 359072,
        "modelId": 320224,
        "model": {"name": "Daphne Blake Cosplay (Scooby Doo)", "type": "LORA"},
        "name": "v1.0",
        "images": [{"url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/original=true"}],
        "baseModel": "SD 1.5",
        "downloadUrl": "https://civitai.com/api/download/models/359072",
        "files": [
            {
                "type": "Model",
                "primary": True,
                "sizeKB": 1024,
                "name": "Daphne Blake Cosplay_v1.safetensors",
                "hashes": {"SHA256": "533317d3f7d269f9f504bdc432514774d3ada3738ebd80f3f1a37ff848e88276"},
            }
        ],
    }

    queried_hashes = []

    async def fake_metadata_provider():
        class Provider:
            async def get_model_by_hash(self, model_hash):
                queried_hashes.append(model_hash)
                if model_hash == "e67ebd5e315f":
                    return lora_version_info, None
                return None, "Model not found"

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )

    parser = AutomaticMetadataParser()

    metadata_text = (
        "woman, natural blonde hair, ice blue eyes, <lora:Daphne Blake Cosplay_v1:1> "
        "daphne blake cosplay, upper body\n"
        "Negative prompt: low quality\n"
        "Steps: 20, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 4140408634, "
        "Size: 512x768, Model hash: 3c8530cb22, Model: cyberrealistic_v33, "
        'Lora hashes: "Daphne Blake Cosplay_v1: e67ebd5e315f", '
        'Hashes: {"lora:Daphne Blake Cosplay_v1": "a2a12bfa01"}'
    )

    result = await parser.parse_metadata(metadata_text)

    assert "e67ebd5e315f" in queried_hashes, (
        f"CivitAI must be queried with the Lora hashes value, got {queried_hashes}"
    )
    assert "a2a12bfa01" not in queried_hashes, (
        "the stale Hashes JSON value must never be used for CivitAI lookup"
    )
    loras = result.get("loras", [])
    assert len(loras) == 1
    lora = loras[0]
    assert lora["hash"] == "533317d3f7d269f9f504bdc432514774d3ada3738ebd80f3f1a37ff848e88276"
    assert lora["id"] == 359072
    assert lora["modelId"] == 320224
    assert lora.get("isDeleted") in (None, False)


@pytest.mark.asyncio
async def test_parse_metadata_resolves_local_lora_with_empty_hash(monkeypatch):
    async def fake_metadata_provider():
        class Provider:
            async def get_model_by_hash(self, model_hash):
                raise AssertionError("Local and empty-hash LoRAs must not query Civitai")

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )
    scanner = LocalRecipeScanner({"local_only": local_lora()})
    metadata_text = (
        "portrait <lora:local_only:0.65> <lora:missing:0.4>\n"
        "Steps: 20, Sampler: Euler, CFG scale: 7, Seed: 1, "
        'Hashes: {"lora:local_only": "", "lora:missing": ""}'
    )

    result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner)

    assert len(result["loras"]) == 1
    entry = result["loras"][0]
    assert entry["name"] == "Local Only"
    assert entry["file_name"] == "local_only"
    assert entry["weight"] == 0.65
    assert entry["hash"] == "a" * 64
    assert entry["localPath"].endswith("local_only.safetensors")
    assert entry["size"] == 123456
    assert entry["baseModel"] == "Flux.1 D"
    assert entry["existsLocally"] is True
    assert entry["isDeleted"] is False
    assert scanner.queries == ["local_only", "missing"]


@pytest.mark.asyncio
async def test_parse_metadata_resolves_prompt_lora_without_hashes(monkeypatch):
    async def fake_metadata_provider():
        return None

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )
    model = local_lora()
    scanner = LocalRecipeScanner({"styles/local_only": model})
    metadata_text = "portrait <lora:styles/local_only:0.7>\nSteps: 20, Seed: 1"

    result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner)

    assert len(result["loras"]) == 1
    assert result["loras"][0]["weight"] == 0.7
    assert result["loras"][0]["hash"] == "a" * 64
    assert scanner.queries == ["styles/local_only"]


@pytest.mark.asyncio
async def test_parse_metadata_prefers_hash_over_colliding_local_name(monkeypatch):
    remote_info = {
        "id": 100,
        "modelId": 200,
        "model": {"name": "Hash Match", "type": "LORA"},
        "name": "v1",
        "files": [{"type": "Model", "primary": True, "name": "hash_match.safetensors", "hashes": {"SHA256": "b" * 64}}],
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_by_hash(self, model_hash):
                assert model_hash == "deadbeef00"
                return remote_info, None

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )
    scanner = LocalRecipeScanner({"local_only": local_lora()})
    metadata_text = (
        "portrait <lora:local_only:0.8>\n"
        "Steps: 20, Seed: 1, "
        'Hashes: {"lora:local_only": "deadbeef00"}'
    )

    result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner)

    assert len(result["loras"]) == 1
    assert result["loras"][0]["id"] == 100
    assert result["loras"][0]["weight"] == 0.8
    assert scanner.queries == []
    assert scanner.hash_queries == ["deadbeef00"]


@pytest.mark.asyncio
async def test_parse_metadata_falls_back_to_name_when_hash_is_unresolved(monkeypatch):
    async def fake_metadata_provider():
        class Provider:
            async def get_model_by_hash(self, model_hash):
                return None, "Model not found"

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )
    scanner = LocalRecipeScanner({"local_only": local_lora()})
    metadata_text = (
        "portrait <lora:local_only:0.8>\nSteps: 20, Seed: 1, "
        'Hashes: {"lora:local_only": "deadbeef00"}'
    )

    result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner)

    assert result["loras"][0]["hash"] == "a" * 64
    assert scanner.queries == ["local_only"]


@pytest.mark.asyncio
async def test_parse_metadata_uses_prompt_weight_for_civitai_resource(monkeypatch):
    remote_info = {
        "id": 100,
        "modelId": 200,
        "model": {"name": "local_only", "type": "LORA"},
        "name": "v1",
        "files": [
            {
                "type": "Model",
                "primary": True,
                "name": "remote_file.safetensors",
                "hashes": {"SHA256": "b" * 64},
            }
        ],
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_version_info(self, version_id):
                assert version_id == 100
                return remote_info, None

            async def get_model_by_hash(self, model_hash):
                raise AssertionError("The Civitai resource should not be fetched again by hash")

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )
    scanner = LocalRecipeScanner({"local_only": local_lora()})
    metadata_text = (
        "portrait <lora:remote_file:0.35> <lora:local_only:0.6>\n"
        "Steps: 20, Seed: 1, "
        'Civitai resources: [{"type":"lora","modelVersionId":100,"modelName":"local_only"}]'
    )

    result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner)

    assert len(result["loras"]) == 2
    assert [entry["file_name"] for entry in result["loras"]] == ["remote_file", "local_only"]
    assert [entry["weight"] for entry in result["loras"]] == [0.35, 0.6]
    assert result["loras"][1]["existsLocally"] is True
    assert scanner.queries == ["local_only"]


@pytest.mark.asyncio
async def test_parse_metadata_keeps_mixed_local_and_civitai_loras(monkeypatch):
    remote_info = {
        "id": 100,
        "modelId": 200,
        "model": {"name": "Remote LoRA", "type": "LORA"},
        "name": "v1",
        "files": [
            {
                "type": "Model",
                "primary": True,
                "name": "remote.safetensors",
                "hashes": {"SHA256": "b" * 64},
            }
        ],
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_by_hash(self, model_hash):
                assert model_hash == "bbbbbbbbbb"
                return remote_info, None

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )
    scanner = LocalRecipeScanner({"local_only": local_lora()})
    metadata_text = (
        "portrait <lora:local_only:0.6> <lora:remote:0.9>\n"
        "Steps: 20, Seed: 1, "
        'Hashes: {"lora:local_only": "", "lora:remote": "bbbbbbbbbb"}'
    )

    result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner)

    assert [entry["name"] for entry in result["loras"]] == ["Local Only", "Remote LoRA"]
    assert [entry["weight"] for entry in result["loras"]] == [0.6, 0.9]
    assert result["loras"][0]["existsLocally"] is True
    assert result["loras"][1]["existsLocally"] is False


@pytest.mark.asyncio
async def test_parse_metadata_extracts_checkpoint_from_model_hash(monkeypatch):
    checkpoint_info = {
        "id": 98765,
        "modelId": 654321,
        "model": {"name": "Flux Illustrious", "type": "checkpoint"},
        "name": "v1",
        "images": [{"url": "https://image.civitai.com/checkpoints/original=true"}],
        "baseModel": "flux",
        "downloadUrl": "https://civitai.com/api/download/checkpoint",
        "files": [
            {
                "type": "Model",
                "primary": True,
                "sizeKB": 1024,
                "name": "FluxIllustrious_v1.safetensors",
                "hashes": {"SHA256": "C3688EE04C"},
            }
        ],
    }

    async def fake_metadata_provider():
        class Provider:
            async def get_model_by_hash(self, model_hash):
                assert model_hash == "c3688ee04c"
                return checkpoint_info, None

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )

    parser = AutomaticMetadataParser()

    metadata_text = (
        "A cyberpunk portrait with neon highlights.\n"
        "Negative prompt: low quality\n"
        "Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 123456, Size: 832x1216, "
        "Model hash: c3688ee04c, Model: models/waiNSFWIllustrious_v110.safetensors"
    )

    result = await parser.parse_metadata(metadata_text)

    checkpoint = result.get("checkpoint")
    assert checkpoint is not None
    assert checkpoint["hash"] == "c3688ee04c"
    assert checkpoint["name"] == "Flux Illustrious"
    assert checkpoint["version"] == "v1"
    assert checkpoint["file_name"] == "FluxIllustrious_v1"
    assert result["model"] == checkpoint
    assert result["base_model"] == "flux"
    assert result["loras"] == []


@pytest.mark.asyncio
async def test_parse_metadata_keeps_empty_placeholder_hash_lora_unresolved(monkeypatch):
    """A LoRA hash equal to the SHA256("") placeholder must never be resolved
    against CivitAI or the local hash index, but the LoRA item itself must be
    kept: matched by filename locally when present, otherwise kept as an
    unresolved entry (no hash) instead of being dropped."""
    queried_hashes = []

    async def fake_metadata_provider():
        class Provider:
            async def get_model_by_hash(self, model_hash):
                queried_hashes.append(model_hash)
                return None, "Model not found"

            async def get_model_version_info(self, version_id):
                raise AssertionError("get_model_version_info should not be called")

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        fake_metadata_provider,
    )

    parser = AutomaticMetadataParser()
    metadata_text = (
        "photo of a DeLorean DMC12, <lora:dmc12bttf:1.2>, at night\n"
        "Steps: 20, Sampler: Euler, CFG scale: 1, Seed: 2242760352, Size: 1280x720, "
        "Model: flux1-dev, Model hash: 3f97fdc57a, "
        'Lora hashes: "dmc12bttf: e3b0c44298fc"'
    )

    # Local file with the same name: the item is matched by filename.
    scanner_with_local = LocalRecipeScanner({"dmc12bttf": local_lora("dmc12bttf")})
    result = await parser.parse_metadata(metadata_text, recipe_scanner=scanner_with_local)

    assert "e3b0c44298fc" not in queried_hashes
    assert "e3b0c44298" not in queried_hashes
    assert scanner_with_local.hash_queries == []
    assert scanner_with_local.queries == ["dmc12bttf"]
    assert len(result["loras"]) == 1
    assert result["loras"][0]["file_name"] == "dmc12bttf"
    assert result["loras"][0]["weight"] == 1.2
    assert result["loras"][0]["existsLocally"] is True
    assert result["loras"][0]["isDeleted"] is False

    # No local file: the item is kept as unresolved (empty hash, flagged
    # hashInvalid so the UI renders the unresolvable-hash badge).
    scanner_without_local = LocalRecipeScanner({})
    result = await parser.parse_metadata(metadata_text, recipe_scanner=scanner_without_local)

    assert len(result["loras"]) == 1
    lora = result["loras"][0]
    assert lora["file_name"] == "dmc12bttf"
    assert lora["weight"] == 1.2
    assert lora["hash"] == ""
    assert lora["hashInvalid"] is True
    assert lora["existsLocally"] is False
    assert lora["isDeleted"] is False


@pytest.mark.asyncio
async def test_parse_metadata_resolves_civitai_resources_from_local_library(monkeypatch):
    """Resources whose version id is in the library must not hit CivitAI."""

    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}"
                )

            async def get_model_by_hash(self, model_hash):
                raise AssertionError(
                    f"Unexpected CivitAI hash lookup for local model {model_hash}"
                )

        return Provider()

    monkeypatch.setattr(
        "py.recipes.parsers.automatic.get_default_metadata_provider",
        failing_metadata_provider,
    )

    checkpoint = {
        "file_path": "/models/checkpoints/z_image_turbo.safetensors",
        "file_name": "z_image_turbo",
        "model_name": "Z Image",
        "sha256": "b" * 64,
        "size": 2048,
        "base_model": "sdxl",
        "preview_url": "/models/checkpoints/z_image_turbo.preview.png",
        "civitai": {"id": 2442439, "modelId": 123456, "name": "Turbo"},
    }
    lora = {
        "file_path": "/models/loras/styles/cfg_boost.safetensors",
        "file_name": "cfg_boost",
        "model_name": "CFG Boost",
        "sha256": "c" * 64,
        "size": 1024,
        "base_model": "illustrious",
        "preview_url": "/models/loras/styles/cfg_boost.preview.png",
        "civitai": {"id": 947620, "modelId": 98765, "name": "v1"},
    }
    scanner = LocalRecipeScanner(
        {},
        models_by_version={947620: lora},
        checkpoints_by_version={2442439: checkpoint},
    )

    metadata_text = (
        "a vintage car at night\n"
        "Negative prompt: fog\n"
        "Steps: 23, Sampler: Euler, CFG scale: 3.5, Seed: 1760020955, "
        "Size: 832x1216, "
        'Civitai resources: ['
        '{"type":"checkpoint","modelVersionId":2442439,"modelName":"Z Image","modelVersionName":"Turbo"},'
        '{"type":"lora","modelVersionId":947620,"modelName":"cfg_scale_boost","weight":0.7}'
        "]"
    )

    result = await AutomaticMetadataParser().parse_metadata(
        metadata_text, recipe_scanner=scanner
    )

    assert scanner.checkpoint_version_queries == [2442439]
    assert scanner.version_queries == [947620]

    parsed_checkpoint = result["checkpoint"]
    assert parsed_checkpoint["existsLocally"] is True
    assert parsed_checkpoint["name"] == "Z Image"
    assert parsed_checkpoint["file_name"] == "z_image_turbo"
    assert parsed_checkpoint["hash"] == "b" * 64
    assert parsed_checkpoint["localPath"] == "/models/checkpoints/z_image_turbo.safetensors"

    assert len(result["loras"]) == 1
    parsed_lora = result["loras"][0]
    assert parsed_lora["existsLocally"] is True
    assert parsed_lora["name"] == "CFG Boost"
    assert parsed_lora["file_name"] == "cfg_boost"
    assert parsed_lora["baseModel"] == "illustrious"
    assert parsed_lora["weight"] == 0.7
    assert result["base_model"] == "sdxl"