Will Miao
19abf7ed03
Distinguish same-version quantized models (e.g. GGUF Q4_K_M vs Q8_0) on cards, which previously showed only the identical CivitAI version name. Detection is content-based and stored as a three-state sidecar field (absent = unchecked, "" = checked-unavailable, value = canonical name), mirroring the autov3 pipeline: - GGUF: minimal pure-struct header parser reads general.file_type - safetensors: dominant tensor dtype F8_* -> FP8 - CivitAI files[].metadata.fp == "fp8" fills in as a fallback only - A one-time backfill service covers existing libraries; each file costs only a few-hundred-byte header read The footer renders the badge after the version name (suppressed in grouped mode), gated by a new show_quantization_on_card setting that follows the show_version_on_card toggle pattern.
204 lines
7.1 KiB
Python
204 lines
7.1 KiB
Python
"""Scanner -> format_response flow tests for the quantization field."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import os
|
|
import struct
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
import pytest
|
|
|
|
from py.services import model_scanner
|
|
from py.services.checkpoint_scanner import CheckpointScanner
|
|
from py.services.checkpoint_service import CheckpointService
|
|
from py.services.model_scanner import ModelScanner
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_model_scanner_singletons():
|
|
ModelScanner._instances.clear()
|
|
ModelScanner._locks.clear()
|
|
yield
|
|
ModelScanner._instances.clear()
|
|
ModelScanner._locks.clear()
|
|
|
|
|
|
def _normalize(path: Path) -> str:
|
|
return str(path).replace(os.sep, "/")
|
|
|
|
|
|
def _write_gguf(path: Path, file_type: int) -> None:
|
|
"""Write a minimal GGUF file whose general.file_type is ``file_type``."""
|
|
blob = io.BytesIO()
|
|
blob.write(b"GGUF")
|
|
blob.write(struct.pack("<I", 3)) # version
|
|
blob.write(struct.pack("<Q", 0)) # tensor_count
|
|
blob.write(struct.pack("<Q", 1)) # metadata_kv_count
|
|
key = b"general.file_type"
|
|
blob.write(struct.pack("<Q", len(key)))
|
|
blob.write(key)
|
|
blob.write(struct.pack("<I", 4)) # UINT32
|
|
blob.write(struct.pack("<I", file_type))
|
|
path.write_bytes(blob.getvalue())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gguf_checkpoint_quantization_flows_to_format_response(tmp_path: Path, monkeypatch):
|
|
"""A Q4_K_M GGUF checkpoint keeps its quantization from scan to API response."""
|
|
checkpoints_root = tmp_path / "checkpoints"
|
|
checkpoints_root.mkdir()
|
|
checkpoint_file = checkpoints_root / "quantized_model.gguf"
|
|
_write_gguf(checkpoint_file, 15) # MOSTLY_Q4_K_M
|
|
|
|
normalized_root = _normalize(checkpoints_root)
|
|
normalized_file = _normalize(checkpoint_file)
|
|
|
|
monkeypatch.setattr(model_scanner.config, "base_models_roots", [normalized_root], raising=False)
|
|
monkeypatch.setattr(model_scanner.config, "checkpoints_roots", [normalized_root], raising=False)
|
|
|
|
scanner = CheckpointScanner()
|
|
model_data = await scanner._process_model_file(normalized_file, normalized_root)
|
|
|
|
assert model_data is not None
|
|
assert model_data["quantization"] == "Q4_K_M"
|
|
|
|
# The checked value is persisted to the sidecar, so future scans never
|
|
# re-read the header.
|
|
sidecar = checkpoints_root / "quantized_model.metadata.json"
|
|
assert sidecar.exists()
|
|
saved = json.loads(sidecar.read_text(encoding="utf-8"))
|
|
assert saved.get("quantization") == "Q4_K_M"
|
|
|
|
service = CheckpointService(scanner)
|
|
response = await service.format_response(model_data)
|
|
assert response is not None
|
|
assert response["quantization"] == "Q4_K_M"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fp8_safetensors_checkpoint_quantization_flows_to_format_response(
|
|
tmp_path: Path, monkeypatch
|
|
):
|
|
"""A safetensors checkpoint with a dominant F8 dtype reports FP8 end to end."""
|
|
torch = pytest.importorskip("torch")
|
|
from safetensors.torch import save_file
|
|
|
|
checkpoints_root = tmp_path / "checkpoints"
|
|
checkpoints_root.mkdir()
|
|
checkpoint_file = checkpoints_root / "fp8_model.safetensors"
|
|
save_file(
|
|
{"weight": torch.zeros(32, 32, dtype=torch.float8_e4m3fn)},
|
|
str(checkpoint_file),
|
|
)
|
|
|
|
normalized_root = _normalize(checkpoints_root)
|
|
normalized_file = _normalize(checkpoint_file)
|
|
|
|
monkeypatch.setattr(model_scanner.config, "base_models_roots", [normalized_root], raising=False)
|
|
monkeypatch.setattr(model_scanner.config, "checkpoints_roots", [normalized_root], raising=False)
|
|
|
|
scanner = CheckpointScanner()
|
|
model_data = await scanner._process_model_file(normalized_file, normalized_root)
|
|
|
|
assert model_data is not None
|
|
assert model_data["quantization"] == "FP8"
|
|
|
|
service = CheckpointService(scanner)
|
|
response = await service.format_response(model_data)
|
|
assert response is not None
|
|
assert response["quantization"] == "FP8"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fp16_safetensors_checkpoint_reports_checked_unavailable(
|
|
tmp_path: Path, monkeypatch
|
|
):
|
|
"""An unquantized safetensors checkpoint stores '' (checked, no value)."""
|
|
torch = pytest.importorskip("torch")
|
|
from safetensors.torch import save_file
|
|
|
|
checkpoints_root = tmp_path / "checkpoints"
|
|
checkpoints_root.mkdir()
|
|
checkpoint_file = checkpoints_root / "plain_model.safetensors"
|
|
save_file({"weight": torch.zeros(32, 32, dtype=torch.float16)}, str(checkpoint_file))
|
|
|
|
normalized_root = _normalize(checkpoints_root)
|
|
normalized_file = _normalize(checkpoint_file)
|
|
|
|
monkeypatch.setattr(model_scanner.config, "base_models_roots", [normalized_root], raising=False)
|
|
monkeypatch.setattr(model_scanner.config, "checkpoints_roots", [normalized_root], raising=False)
|
|
|
|
scanner = CheckpointScanner()
|
|
model_data = await scanner._process_model_file(normalized_file, normalized_root)
|
|
|
|
assert model_data is not None
|
|
assert model_data["quantization"] == ""
|
|
|
|
# JSON null encodes the checked-unavailable state in the sidecar.
|
|
sidecar = checkpoints_root / "plain_model.metadata.json"
|
|
saved = json.loads(sidecar.read_text(encoding="utf-8"))
|
|
assert "quantization" in saved
|
|
assert saved["quantization"] is None
|
|
|
|
service = CheckpointService(scanner)
|
|
response = await service.format_response(model_data)
|
|
assert response is not None
|
|
assert response["quantization"] == ""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_civitai_fp8_fallback_fills_unchecked_lora(tmp_path: Path, monkeypatch):
|
|
"""When content detection is inconclusive and the sidecar carries Civitai
|
|
metadata whose SHA256-matching file is fp8, the scan resolves 'FP8'."""
|
|
from py.services.lora_scanner import LoraScanner
|
|
|
|
loras_root = tmp_path / "loras"
|
|
loras_root.mkdir()
|
|
lora_file = loras_root / "fp8_lora.safetensors"
|
|
content = b"fake lora content - not a real safetensors header"
|
|
lora_file.write_bytes(content)
|
|
sha256 = hashlib.sha256(content).hexdigest()
|
|
|
|
normalized_root = _normalize(loras_root)
|
|
normalized_file = _normalize(lora_file)
|
|
|
|
# Pre-existing sidecar with Civitai data but no quantization key yet.
|
|
sidecar = loras_root / "fp8_lora.metadata.json"
|
|
sidecar.write_text(
|
|
json.dumps({
|
|
"file_name": "fp8_lora",
|
|
"model_name": "fp8_lora",
|
|
"file_path": normalized_file,
|
|
"size": len(content),
|
|
"modified": 1.0,
|
|
"sha256": sha256,
|
|
"base_model": "Unknown",
|
|
"preview_url": "",
|
|
"civitai": {
|
|
"files": [
|
|
{
|
|
"name": "fp8_lora.safetensors",
|
|
"hashes": {"SHA256": sha256.upper()},
|
|
"metadata": {"fp": "fp8"},
|
|
}
|
|
]
|
|
},
|
|
}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
monkeypatch.setattr(model_scanner.config, "loras_roots", [normalized_root], raising=False)
|
|
|
|
scanner = LoraScanner()
|
|
model_data = await scanner._process_model_file(normalized_file, normalized_root)
|
|
|
|
assert model_data is not None
|
|
assert model_data["quantization"] == "FP8"
|
|
|
|
saved = json.loads(sidecar.read_text(encoding="utf-8"))
|
|
assert saved.get("quantization") == "FP8"
|