Dariusz L
794083b212
- Add paste handler (Ctrl+V) in Metadata tab for workflow JSON, images, and Civitai image URLs/HTML - Resolve authentic .safetensors filenames and local ComfyUI model paths via Civitai modelVersionId - Normalize prompt <lora:...> tags to authentic filenames and enrich transfer/loaded chip labels
460 lines
15 KiB
Python
460 lines
15 KiB
Python
import asyncio
|
|
import base64
|
|
import io
|
|
import json
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
from aiohttp import web
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
from core.routes.context import RouteContext
|
|
from core.services.image_metadata_service import (
|
|
ImageMetadataService,
|
|
inspect_civitai_image,
|
|
inspect_image_bytes,
|
|
inspect_workflow_json,
|
|
)
|
|
from core.workflow.formats import (
|
|
API_WORKFLOW_FORMAT,
|
|
UI_WORKFLOW_FORMAT,
|
|
detect_workflow_format,
|
|
normalize_workflow_payload,
|
|
)
|
|
|
|
|
|
def api_workflow():
|
|
return {
|
|
"1": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "models/checkpoint.safetensors"},
|
|
"_meta": {"title": "Checkpoint"},
|
|
},
|
|
"2": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["1", 0],
|
|
"seed": 123,
|
|
"steps": 20,
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def png_with_metadata(metadata):
|
|
image = Image.new("RGB", (32, 16), color=(20, 30, 40))
|
|
png_info = PngInfo()
|
|
for key, value in metadata.items():
|
|
png_info.add_text(key, value)
|
|
output = io.BytesIO()
|
|
image.save(output, format="PNG", pnginfo=png_info)
|
|
return output.getvalue()
|
|
|
|
|
|
def jpeg_with_nested_user_comment(comment):
|
|
image = Image.new("RGB", (32, 16), color=(20, 30, 40))
|
|
exif = Image.Exif()
|
|
exif[34665] = {
|
|
37510: b"UNICODE\x00\x00" + comment.encode("utf-16le"),
|
|
}
|
|
output = io.BytesIO()
|
|
image.save(output, format="JPEG", exif=exif)
|
|
return output.getvalue()
|
|
|
|
|
|
def test_api_workflow_is_detected_and_normalized_for_existing_analyzer():
|
|
api = api_workflow()
|
|
|
|
assert detect_workflow_format(api) == API_WORKFLOW_FORMAT
|
|
normalized, workflow_format = normalize_workflow_payload(api)
|
|
|
|
assert workflow_format == API_WORKFLOW_FORMAT
|
|
assert [node["type"] for node in normalized["nodes"]] == [
|
|
"CheckpointLoaderSimple",
|
|
"KSampler",
|
|
]
|
|
assert normalized["nodes"][0]["widgets_values"] == [
|
|
"models/checkpoint.safetensors"
|
|
]
|
|
assert normalized["nodes"][0]["outputs"][0]["links"] == [1]
|
|
assert normalized["nodes"][1]["inputs"][0]["link"] == 1
|
|
|
|
|
|
def test_image_metadata_reads_comfyui_api_prompt():
|
|
result = inspect_image_bytes(
|
|
png_with_metadata({"prompt": json.dumps(api_workflow())}),
|
|
"generated.png",
|
|
)
|
|
|
|
assert result["filename"] == "generated.png"
|
|
assert result["image"] == {
|
|
"format": "PNG",
|
|
"width": 32,
|
|
"height": 16,
|
|
"mode": "RGB",
|
|
}
|
|
assert result["workflow_format"] == API_WORKFLOW_FORMAT
|
|
assert result["workflow_source"] == "prompt"
|
|
assert result["workflow_node_count"] == 2
|
|
assert result["metadata"]["has_workflow"] is True
|
|
|
|
|
|
def test_image_without_workflow_still_returns_safe_image_metadata():
|
|
result = inspect_image_bytes(
|
|
png_with_metadata({"parameters": "seed: 123, steps: 20"}),
|
|
"plain.png",
|
|
)
|
|
|
|
assert result["workflow"] is None
|
|
assert result["workflow_format"] == ""
|
|
assert result["metadata"]["parameters"] == "seed: 123, steps: 20"
|
|
|
|
|
|
def test_jpeg_nested_exif_reads_civitai_metadata_and_resources():
|
|
named_resources = [
|
|
{
|
|
"type": "checkpoint",
|
|
"modelVersionId": 3072332,
|
|
"modelName": "Krea 2 Turbo",
|
|
"modelVersionName": "v1.0",
|
|
},
|
|
{
|
|
"type": "lora",
|
|
"weight": 0.5,
|
|
"modelVersionId": 3220691,
|
|
"modelName": "SNOFS",
|
|
"modelVersionName": "Krea 2 v1.3D",
|
|
},
|
|
]
|
|
civitai_metadata = {
|
|
"workflow": "txt2img",
|
|
"prompt": "a pilot",
|
|
"resources": [
|
|
{
|
|
"modelVersionId": 3072332,
|
|
"strength": 1,
|
|
"type": "Checkpoint",
|
|
},
|
|
{
|
|
"modelVersionId": 3220691,
|
|
"strength": 0.5,
|
|
"type": "LORA",
|
|
},
|
|
],
|
|
}
|
|
comment = (
|
|
"a pilot\n"
|
|
"Negative prompt: \n"
|
|
"Steps: 8, Sampler: euler, CFG scale: 1\n"
|
|
f"Civitai resources: {json.dumps(named_resources)}\n"
|
|
f"Civitai metadata: {json.dumps(civitai_metadata)}"
|
|
)
|
|
|
|
result = inspect_image_bytes(jpeg_with_nested_user_comment(comment), "civitai.jpg")
|
|
|
|
assert "EXIF:UserComment" in result["metadata"]["keys"]
|
|
assert "Civitai metadata" in result["metadata"]["keys"]
|
|
assert "Civitai resources" in result["metadata"]["keys"]
|
|
assert result["metadata"]["parameters"].startswith("a pilot")
|
|
assert result["civitai_workflow"] == "txt2img"
|
|
assert result["workflow"] is None
|
|
assert result["workflow_format"] == ""
|
|
assert result["civitai_metadata"]["prompt"] == "a pilot"
|
|
assert [resource["name"] for resource in result["civitai_resources"]] == [
|
|
"Krea 2 Turbo",
|
|
"SNOFS",
|
|
]
|
|
assert [resource["category"] for resource in result["civitai_resources"]] == [
|
|
"checkpoints",
|
|
"loras",
|
|
]
|
|
assert result["civitai_resources"][1]["strength"] == 0.5
|
|
|
|
|
|
def test_json_inspection_accepts_ui_and_api_workflows():
|
|
ui = {"nodes": [{"id": 1, "type": "CheckpointLoaderSimple", "widgets_values": []}]}
|
|
|
|
ui_result = inspect_workflow_json(ui, filename="workflow.json")
|
|
api_result = inspect_workflow_json(json.dumps(api_workflow()))
|
|
|
|
assert ui_result["workflow_format"] == UI_WORKFLOW_FORMAT
|
|
assert ui_result["workflow_node_count"] == 1
|
|
assert api_result["workflow_format"] == API_WORKFLOW_FORMAT
|
|
assert api_result["workflow_node_count"] == 2
|
|
|
|
|
|
def test_json_inspection_rejects_unrelated_json():
|
|
with pytest.raises(ValueError, match="Unsupported workflow JSON"):
|
|
inspect_workflow_json({"settings": {"theme": "dark"}})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_metadata_service_accepts_image_data_url():
|
|
raw = png_with_metadata({"prompt": json.dumps(api_workflow())})
|
|
data_url = "data:image/png;base64," + base64.b64encode(raw).decode("ascii")
|
|
request = SimpleNamespace(
|
|
json=AsyncMock(
|
|
return_value={
|
|
"source_type": "image",
|
|
"filename": "generated.png",
|
|
"data_url": data_url,
|
|
}
|
|
)
|
|
)
|
|
service = ImageMetadataService(
|
|
RouteContext({"asyncio": asyncio, "web": web})
|
|
)
|
|
|
|
response = await service.inspect_metadata(request)
|
|
body = json.loads(response.text)
|
|
|
|
assert response.status == 200
|
|
assert body["workflow_format"] == API_WORKFLOW_FORMAT
|
|
assert body["workflow_node_count"] == 2
|
|
|
|
|
|
def civitai_html_payload(meta=None, resources=None, img_data=None):
|
|
if meta is None:
|
|
meta = {
|
|
"prompt": "cyberpunk woman, glowing neon",
|
|
"negativePrompt": "blurry, low quality",
|
|
"cfgScale": 7.0,
|
|
"steps": 25,
|
|
"sampler": "Euler a",
|
|
"seed": 42,
|
|
"Size": "1024x1024",
|
|
"civitaiResources": [
|
|
{
|
|
"type": "checkpoint",
|
|
"modelId": 12345,
|
|
"modelVersionId": 67890,
|
|
"modelName": "CyberRealistic",
|
|
"versionName": "v4.0",
|
|
"strength": 1.0,
|
|
}
|
|
],
|
|
}
|
|
if resources is None:
|
|
resources = [
|
|
{
|
|
"id": 9999,
|
|
"modelId": 54321,
|
|
"name": "NeonGlow LoRA",
|
|
"modelType": "LORA",
|
|
"strength": 0.8,
|
|
}
|
|
]
|
|
if img_data is None:
|
|
img_data = {
|
|
"id": 145110824,
|
|
"name": "cyberpunk_image.jpeg",
|
|
"width": 1024,
|
|
"height": 1024,
|
|
"mimeType": "image/jpeg",
|
|
}
|
|
|
|
next_data = {
|
|
"props": {
|
|
"pageProps": {
|
|
"trpcState": {
|
|
"json": {
|
|
"queries": [
|
|
{
|
|
"queryKey": [["image", "get"], {"input": {"id": 145110824}}],
|
|
"state": {"data": img_data},
|
|
},
|
|
{
|
|
"queryKey": [["image", "getGenerationData"], {"input": {"id": 145110824}}],
|
|
"state": {
|
|
"data": {
|
|
"meta": meta,
|
|
"resources": resources,
|
|
}
|
|
},
|
|
},
|
|
]
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return (
|
|
'<!DOCTYPE html><html><head>'
|
|
f'<script id="__NEXT_DATA__" type="application/json">{json.dumps(next_data)}</script>'
|
|
'</head><body></body></html>'
|
|
)
|
|
|
|
|
|
def mock_civitai_version_request(name, url, **kwargs):
|
|
if "67890" in url:
|
|
return {
|
|
"name": "v4.0",
|
|
"model": {"name": "CyberRealistic"},
|
|
"files": [
|
|
{
|
|
"name": "cyberrealistic_v40.safetensors",
|
|
"primary": True,
|
|
"type": "Model",
|
|
"hashes": {"SHA256": "ABC123"},
|
|
}
|
|
],
|
|
}
|
|
if "9999" in url or "54321" in url:
|
|
return {
|
|
"name": "v1.0",
|
|
"model": {"name": "NeonGlow LoRA"},
|
|
"files": [
|
|
{
|
|
"name": "neon_glow_lora_v1.safetensors",
|
|
"primary": True,
|
|
"type": "Model",
|
|
"hashes": {"SHA256": "DEF456"},
|
|
}
|
|
],
|
|
}
|
|
return None
|
|
|
|
|
|
def test_inspect_civitai_image_extracts_metadata_and_resources():
|
|
html = civitai_html_payload()
|
|
with patch(
|
|
"core.services.image_metadata_service.request_source_response",
|
|
return_value=SimpleNamespace(status_code=200, text=html),
|
|
), patch(
|
|
"core.services.image_metadata_service.execute_provider_json_request",
|
|
side_effect=mock_civitai_version_request,
|
|
):
|
|
result = inspect_civitai_image(145110824, url="https://civitai.red/images/145110824")
|
|
|
|
assert result["source_type"] == "civitai"
|
|
assert result["civitai_image_id"] == 145110824
|
|
assert result["civitai_metadata"]["prompt"] == "cyberpunk woman, glowing neon"
|
|
assert result["civitai_metadata"]["negativePrompt"] == "blurry, low quality"
|
|
assert len(result["civitai_resources"]) == 2
|
|
categories = {res["category"] for res in result["civitai_resources"]}
|
|
assert categories == {"checkpoints", "loras"}
|
|
assert "cyberpunk woman" in result["metadata"]["parameters"]
|
|
assert "Civitai metadata" in result["metadata"]["keys"]
|
|
assert "Civitai resources" in result["metadata"]["keys"]
|
|
|
|
# Verify filenames were intelligently resolved from version details
|
|
ckpt_res = next(r for r in result["civitai_resources"] if r["category"] == "checkpoints")
|
|
assert ckpt_res["filename"] == "cyberrealistic_v40.safetensors"
|
|
assert ckpt_res["model_name"] == "CyberRealistic"
|
|
lora_res = next(r for r in result["civitai_resources"] if r["category"] == "loras")
|
|
assert lora_res["filename"] == "neon_glow_lora_v1.safetensors"
|
|
assert lora_res["model_name"] == "NeonGlow LoRA"
|
|
assert "Model: cyberrealistic_v40.safetensors" in result["metadata"]["parameters"]
|
|
|
|
|
|
def test_inspect_civitai_image_normalizes_prompt_lora_tags():
|
|
meta = {
|
|
"prompt": "portrait of woman <lora:NeonGlow LoRA:0.85> in neon city",
|
|
"civitaiResources": [
|
|
{
|
|
"type": "lora",
|
|
"modelId": 54321,
|
|
"modelVersionId": 9999,
|
|
"modelName": "NeonGlow LoRA",
|
|
"strength": 0.85,
|
|
}
|
|
],
|
|
}
|
|
html = civitai_html_payload(meta=meta, resources=[])
|
|
with patch(
|
|
"core.services.image_metadata_service.request_source_response",
|
|
return_value=SimpleNamespace(status_code=200, text=html),
|
|
), patch(
|
|
"core.services.image_metadata_service.execute_provider_json_request",
|
|
side_effect=mock_civitai_version_request,
|
|
):
|
|
result = inspect_civitai_image(145110824)
|
|
|
|
assert "<lora:neon_glow_lora_v1:0.85>" in result["civitai_metadata"]["prompt"]
|
|
|
|
|
|
def test_inspect_civitai_image_with_embedded_comfy_workflow():
|
|
meta = {
|
|
"prompt": "prompt with workflow",
|
|
"comfy": api_workflow(),
|
|
}
|
|
html = civitai_html_payload(meta=meta)
|
|
with patch(
|
|
"core.services.image_metadata_service.request_source_response",
|
|
return_value=SimpleNamespace(status_code=200, text=html),
|
|
), patch(
|
|
"core.services.image_metadata_service.execute_provider_json_request",
|
|
side_effect=mock_civitai_version_request,
|
|
):
|
|
result = inspect_civitai_image(145110824)
|
|
|
|
assert result["workflow_format"] == API_WORKFLOW_FORMAT
|
|
assert result["workflow_node_count"] == 2
|
|
assert result["workflow_source"] == "Civitai generation data"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_metadata_service_accepts_civitai_source_type():
|
|
html = civitai_html_payload()
|
|
request = SimpleNamespace(
|
|
json=AsyncMock(
|
|
return_value={
|
|
"source_type": "civitai",
|
|
"image_id": 145110824,
|
|
"url": "https://civitai.red/images/145110824",
|
|
}
|
|
)
|
|
)
|
|
service = ImageMetadataService(
|
|
RouteContext({"asyncio": asyncio, "web": web})
|
|
)
|
|
|
|
with patch(
|
|
"core.services.image_metadata_service.request_source_response",
|
|
return_value=SimpleNamespace(status_code=200, text=html),
|
|
), patch(
|
|
"core.services.image_metadata_service.execute_provider_json_request",
|
|
side_effect=mock_civitai_version_request,
|
|
):
|
|
response = await service.inspect_metadata(request)
|
|
|
|
body = json.loads(response.text)
|
|
assert response.status == 200
|
|
assert body["source_type"] == "civitai"
|
|
assert body["civitai_image_id"] == 145110824
|
|
assert len(body["civitai_resources"]) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_metadata_service_workflow_source_detects_civitai_url():
|
|
html = civitai_html_payload()
|
|
request = SimpleNamespace(
|
|
json=AsyncMock(
|
|
return_value={
|
|
"source_type": "workflow",
|
|
"json_text": "https://civitai.com/images/145110824",
|
|
}
|
|
)
|
|
)
|
|
service = ImageMetadataService(
|
|
RouteContext({"asyncio": asyncio, "web": web})
|
|
)
|
|
|
|
with patch(
|
|
"core.services.image_metadata_service.request_source_response",
|
|
return_value=SimpleNamespace(status_code=200, text=html),
|
|
), patch(
|
|
"core.services.image_metadata_service.execute_provider_json_request",
|
|
side_effect=mock_civitai_version_request,
|
|
):
|
|
response = await service.inspect_metadata(request)
|
|
|
|
body = json.loads(response.text)
|
|
assert response.status == 200
|
|
assert body["source_type"] == "civitai"
|
|
assert body["civitai_image_id"] == 145110824
|
|
|