mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-29 15:04:09 -03:00
feat(recipes): preserve embedded ComfyUI workflow on remote imports
CivitAI serves a re-encoded, metadata-free optimized rendition as the recipe preview, so the ComfyUI workflow embedded in the original image was dropped: imported recipes reported has_workflow=false and never offered "Send Workflow to ComfyUI" even when the source image carried one. Recover the workflow from the original rendition and carry it to the save step as data, so the stored preview stays the small optimized image: - ExifUtils: embed a caller-supplied workflow during optimize_image's single encode pass, and add embed_workflow() to patch WebP EXIF in place (used by the verbatim skip_optimize branch and as a safety net). - RecipePersistenceService.save_recipe: embed metadata["workflow"] before detecting has_workflow. - analyze_remote_image: return the workflow recovered from the original rendition it already downloads for EXIF parsing. - RecipeManagementHandler: add _fetch_original_media() and workflow helpers; _do_import_from_url reuses them, and _do_import_remote_recipe fetches the original only when CivitAI reports a ComfyUI payload (meta.comfy) so workflow-less images pay no extra bandwidth. - Batch URL imports and the import modal forward the recovered workflow. Verified against the reported image: has_workflow flips from false to true and the recovered workflow matches the original (25 nodes, same graph id).
This commit is contained in:
@@ -0,0 +1,101 @@
|
||||
import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
|
||||
|
||||
const { showToastMock, translateMock } = vi.hoisted(() => ({
|
||||
showToastMock: vi.fn(),
|
||||
translateMock: vi.fn((key, params, fallback) =>
|
||||
typeof fallback === 'string' ? fallback : key
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/uiHelpers.js', () => ({
|
||||
showToast: showToastMock,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/i18nHelpers.js', () => ({
|
||||
translate: translateMock,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/api/modelApiFactory.js', () => ({
|
||||
getModelApiClient: vi.fn(() => ({})),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/api/apiConfig.js', () => ({
|
||||
MODEL_TYPES: { LORA: 'loras', CHECKPOINT: 'checkpoints', EMBEDDING: 'embeddings' },
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/storageHelpers.js', () => ({
|
||||
getStorageItem: vi.fn(() => null),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/state/index.js', () => ({
|
||||
state: { virtualScroller: null },
|
||||
}));
|
||||
|
||||
import { DownloadManager } from '../../../static/js/managers/import/DownloadManager.js';
|
||||
|
||||
function buildImportManager(recipeData) {
|
||||
return {
|
||||
recipeId: null,
|
||||
recipeName: 'Test Recipe',
|
||||
recipeImage: null,
|
||||
recipeTags: [],
|
||||
downloadableLoRAs: [],
|
||||
importMode: 'url',
|
||||
recipeData,
|
||||
loadingManager: { showSimpleLoading: vi.fn(), hide: vi.fn() },
|
||||
};
|
||||
}
|
||||
|
||||
async function saveAndReadMetadata(recipeData) {
|
||||
let capturedBody = null;
|
||||
const fetchMock = vi.fn(async (url, options) => {
|
||||
capturedBody = options?.body ?? null;
|
||||
return { ok: true, json: async () => ({ success: true }) };
|
||||
});
|
||||
globalThis.fetch = fetchMock;
|
||||
window.fetch = fetchMock;
|
||||
|
||||
const manager = new DownloadManager(buildImportManager(recipeData));
|
||||
await manager.saveRecipe(true);
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledWith('/api/lm/recipes/save', expect.anything());
|
||||
expect(capturedBody).toBeInstanceOf(FormData);
|
||||
return JSON.parse(capturedBody.get('metadata'));
|
||||
}
|
||||
|
||||
describe('recipe import save payload', () => {
|
||||
beforeEach(() => {
|
||||
globalThis.modalManager = { closeModal: vi.fn() };
|
||||
globalThis.window.recipeManager = { loadRecipes: vi.fn() };
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
delete globalThis.modalManager;
|
||||
delete globalThis.window.recipeManager;
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('forwards a workflow recovered from the original rendition', async () => {
|
||||
const workflow = '{"nodes": [{"id": 1}]}';
|
||||
const metadata = await saveAndReadMetadata({
|
||||
image_base64: 'AAAA',
|
||||
base_model: 'sd',
|
||||
loras: [],
|
||||
gen_params: {},
|
||||
workflow,
|
||||
});
|
||||
|
||||
expect(metadata.workflow).toBe(workflow);
|
||||
});
|
||||
|
||||
it('omits the workflow key when analysis recovered none', async () => {
|
||||
const metadata = await saveAndReadMetadata({
|
||||
image_base64: 'AAAA',
|
||||
base_model: 'sd',
|
||||
loras: [],
|
||||
gen_params: {},
|
||||
});
|
||||
|
||||
expect('workflow' in metadata).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,262 @@
|
||||
"""Workflow preservation for remote recipe imports.
|
||||
|
||||
CivitAI serves a re-encoded, metadata-free ``optimized`` rendition as the
|
||||
recipe preview, so an embedded ComfyUI workflow only exists in the
|
||||
``original=true`` image. These tests pin the recovery and transport of that
|
||||
workflow through the remote import paths.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
from py.routes.handlers.recipe_handlers import RecipeManagementHandler
|
||||
from py.services.recipes.persistence_service import PersistenceResult
|
||||
from py.utils.exif_utils import ExifUtils
|
||||
|
||||
|
||||
async def _noop_ensure() -> None:
|
||||
return None
|
||||
|
||||
|
||||
class CapturingPersistence:
|
||||
"""Persistence service double recording the save payload."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def save_recipe(self, **kwargs: Any) -> PersistenceResult:
|
||||
self.calls.append(kwargs)
|
||||
return PersistenceResult({"success": True, "recipe_id": "recipe-1"})
|
||||
|
||||
|
||||
class StubScanner:
|
||||
"""Scanner double exposing only what the remote import paths touch."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.recipes_dir = "/tmp/recipes"
|
||||
|
||||
async def build_local_hash_cache(self) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
async def get_local_lora(self, name, base_model=None):
|
||||
return None
|
||||
|
||||
|
||||
def _make_handler(
|
||||
persistence: CapturingPersistence,
|
||||
*,
|
||||
downloader_factory=None,
|
||||
) -> RecipeManagementHandler:
|
||||
async def default_downloader_factory():
|
||||
class Downloader:
|
||||
async def download_file(self, url, path, use_auth=False):
|
||||
Path(path).write_bytes(b"downloaded")
|
||||
return True, "ok"
|
||||
|
||||
return Downloader()
|
||||
|
||||
analysis_service = SimpleNamespace(
|
||||
_recipe_parser_factory=SimpleNamespace(create_parser=lambda metadata: None)
|
||||
)
|
||||
|
||||
return RecipeManagementHandler(
|
||||
ensure_dependencies_ready=_noop_ensure,
|
||||
recipe_scanner_getter=lambda: StubScanner(),
|
||||
logger=logging.getLogger(__name__),
|
||||
persistence_service=persistence, # pyright: ignore[reportArgumentType]
|
||||
analysis_service=analysis_service, # pyright: ignore[reportArgumentType]
|
||||
downloader_factory=downloader_factory or default_downloader_factory,
|
||||
civitai_client_getter=lambda: None,
|
||||
)
|
||||
|
||||
|
||||
def _meta_with_comfy() -> dict[str, Any]:
|
||||
return {
|
||||
"id": 143518055,
|
||||
"meta": {"prompt": "p", "comfy": '{"prompt": {"1": {"class_type": "KSampler"}}}'},
|
||||
}
|
||||
|
||||
|
||||
def test_meta_indicates_comfy_workflow() -> None:
|
||||
assert RecipeManagementHandler._meta_indicates_comfy_workflow(
|
||||
{"meta": {"comfy": "{}"}}
|
||||
)
|
||||
assert RecipeManagementHandler._meta_indicates_comfy_workflow({"comfy": "{}"})
|
||||
assert not RecipeManagementHandler._meta_indicates_comfy_workflow({"meta": {}})
|
||||
assert not RecipeManagementHandler._meta_indicates_comfy_workflow(
|
||||
{"meta": {"comfy": None}}
|
||||
)
|
||||
assert not RecipeManagementHandler._meta_indicates_comfy_workflow(None)
|
||||
assert not RecipeManagementHandler._meta_indicates_comfy_workflow("comfy")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_original_media_reads_workflow_and_cleans_up(tmp_path, monkeypatch):
|
||||
workflow = json.dumps({"nodes": [{"id": 1}], "last_node_id": 1})
|
||||
source = tmp_path / "original.png"
|
||||
png_info = PngImagePlugin.PngInfo()
|
||||
png_info.add_text("workflow", workflow)
|
||||
png_info.add_text("prompt", '{"1": {"class_type": "KSampler"}}')
|
||||
Image.new("RGB", (32, 32), color="red").save(source, pnginfo=png_info)
|
||||
|
||||
written: list[str] = []
|
||||
|
||||
async def downloader_factory():
|
||||
class Downloader:
|
||||
async def download_file(self, url, path, use_auth=False):
|
||||
written.append(str(path))
|
||||
Path(path).write_bytes(source.read_bytes())
|
||||
return True, "ok"
|
||||
|
||||
return Downloader()
|
||||
|
||||
handler = _make_handler(CapturingPersistence(), downloader_factory=downloader_factory)
|
||||
|
||||
raw_metadata, recovered = await handler._fetch_original_media(
|
||||
"https://image.civitai.com/x/original=true/x.png"
|
||||
)
|
||||
|
||||
assert recovered == workflow
|
||||
# extract_image_metadata prefers the prompt chunk over the workflow.
|
||||
assert raw_metadata is not None and "class_type" in raw_metadata
|
||||
assert written and not os.path.exists(written[0])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_original_media_degrades_on_download_failure():
|
||||
async def downloader_factory():
|
||||
class Downloader:
|
||||
async def download_file(self, url, path, use_auth=False):
|
||||
return False, "boom"
|
||||
|
||||
return Downloader()
|
||||
|
||||
handler = _make_handler(CapturingPersistence(), downloader_factory=downloader_factory)
|
||||
|
||||
assert await handler._fetch_original_media("https://image.civitai.com/x.png") == (
|
||||
None,
|
||||
None,
|
||||
)
|
||||
assert await handler._fetch_original_media(None) == (None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_import_transports_workflow_to_save(monkeypatch):
|
||||
workflow = json.dumps({"nodes": [{"id": 4}]})
|
||||
persistence = CapturingPersistence()
|
||||
handler = _make_handler(persistence)
|
||||
|
||||
async def fake_download_remote_media(image_url):
|
||||
return (
|
||||
b"optimized-preview",
|
||||
".jpg",
|
||||
_meta_with_comfy(),
|
||||
12345,
|
||||
"https://image.civitai.com/x/original=true/x.png",
|
||||
)
|
||||
|
||||
fetched: list[str] = []
|
||||
|
||||
async def fake_fetch_original_media(original_url):
|
||||
fetched.append(original_url)
|
||||
return None, workflow
|
||||
|
||||
handler._download_remote_media = fake_download_remote_media # type: ignore[method-assign]
|
||||
handler._fetch_original_media = fake_fetch_original_media # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
ExifUtils, "extract_image_metadata", staticmethod(lambda path: None)
|
||||
)
|
||||
|
||||
response = await handler._do_import_remote_recipe(
|
||||
image_url="https://civitai.red/images/143518055",
|
||||
name="Recipe",
|
||||
lora_entries=[],
|
||||
checkpoint_entry=None,
|
||||
gen_params_request={},
|
||||
tags=[],
|
||||
base_model="Krea 2",
|
||||
source_path="https://civitai.red/images/143518055",
|
||||
)
|
||||
|
||||
assert response.status == 200
|
||||
assert fetched == ["https://image.civitai.com/x/original=true/x.png"]
|
||||
assert persistence.calls[0]["metadata"]["workflow"] == workflow
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remote_import_skips_original_without_comfy_meta(monkeypatch):
|
||||
persistence = CapturingPersistence()
|
||||
handler = _make_handler(persistence)
|
||||
|
||||
async def fake_download_remote_media(image_url):
|
||||
return (
|
||||
b"optimized-preview",
|
||||
".jpg",
|
||||
{"id": 1, "meta": {"prompt": "p"}},
|
||||
None,
|
||||
"https://image.civitai.com/x/original=true/x.png",
|
||||
)
|
||||
|
||||
async def fail_fetch(original_url): # pragma: no cover - must not be called
|
||||
raise AssertionError("original rendition should not be fetched")
|
||||
|
||||
handler._download_remote_media = fake_download_remote_media # type: ignore[method-assign]
|
||||
handler._fetch_original_media = fail_fetch # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
ExifUtils, "extract_image_metadata", staticmethod(lambda path: None)
|
||||
)
|
||||
|
||||
response = await handler._do_import_remote_recipe(
|
||||
image_url="https://civitai.red/images/1",
|
||||
name="Recipe",
|
||||
lora_entries=[],
|
||||
checkpoint_entry=None,
|
||||
gen_params_request={},
|
||||
tags=[],
|
||||
base_model="SDXL 1.0",
|
||||
source_path="https://civitai.red/images/1",
|
||||
)
|
||||
|
||||
assert response.status == 200
|
||||
assert "workflow" not in persistence.calls[0]["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_url_import_transports_workflow_to_save(monkeypatch):
|
||||
workflow = json.dumps({"nodes": [{"id": 5}]})
|
||||
persistence = CapturingPersistence()
|
||||
handler = _make_handler(persistence)
|
||||
|
||||
async def fake_download_remote_media(image_url):
|
||||
return (
|
||||
b"optimized-preview",
|
||||
".jpg",
|
||||
{"id": 9, "meta": {"prompt": "p"}},
|
||||
None,
|
||||
"https://image.civitai.com/x/original=true/x.png",
|
||||
)
|
||||
|
||||
async def fake_fetch_original_media(original_url):
|
||||
return None, workflow
|
||||
|
||||
handler._download_remote_media = fake_download_remote_media # type: ignore[method-assign]
|
||||
handler._fetch_original_media = fake_fetch_original_media # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
ExifUtils, "extract_image_metadata", staticmethod(lambda path: None)
|
||||
)
|
||||
|
||||
response = await handler._do_import_from_url(
|
||||
"https://civitai.red/images/143518055", StubScanner()
|
||||
)
|
||||
|
||||
assert response.status == 200
|
||||
assert persistence.calls[0]["metadata"]["workflow"] == workflow
|
||||
@@ -589,6 +589,49 @@ class TestBatchImportServiceEdgeCases:
|
||||
assert "batch-import" in persistence_service.saved_recipes[0]["tags"]
|
||||
assert "test" in persistence_service.saved_recipes[0]["tags"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workflow_from_analysis_is_passed_to_persistence(self, tmp_path):
|
||||
"""A workflow recovered from the source's original rendition travels in
|
||||
the analysis payload and must reach save_recipe as metadata."""
|
||||
workflow = '{"nodes": [{"id": 1}]}'
|
||||
ws_manager = MockWebSocketManager()
|
||||
analysis_service = MockAnalysisService(
|
||||
{
|
||||
"https://civitai.red/images/1": MockAnalysisResult(
|
||||
{
|
||||
"loras": [{"name": "test-lora"}],
|
||||
"workflow": workflow,
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
persistence_service = MockPersistenceService()
|
||||
logger = logging.getLogger("test")
|
||||
|
||||
service = BatchImportService(
|
||||
analysis_service=analysis_service, # pyright: ignore[reportArgumentType]
|
||||
persistence_service=persistence_service, # pyright: ignore[reportArgumentType]
|
||||
ws_manager=ws_manager,
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
recipe_scanner_getter = lambda: SimpleNamespace(
|
||||
find_recipes_by_fingerprint=lambda x: [],
|
||||
)
|
||||
civitai_client_getter = lambda: SimpleNamespace()
|
||||
|
||||
await service.start_batch_import(
|
||||
recipe_scanner_getter=recipe_scanner_getter,
|
||||
civitai_client_getter=civitai_client_getter,
|
||||
items=[{"source": "https://civitai.red/images/1"}],
|
||||
tags=[],
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
assert persistence_service.saved_recipes
|
||||
assert persistence_service.saved_recipes[0]["metadata"]["workflow"] == workflow
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skip_duplicates_parameter(self, service):
|
||||
recipe_scanner_getter = lambda: SimpleNamespace()
|
||||
|
||||
@@ -27,14 +27,29 @@ class DummyExifUtils:
|
||||
self.appended = None
|
||||
self.optimized_calls = 0
|
||||
self.workflow_value = None
|
||||
self.optimized_workflow = None
|
||||
self.embedded_workflows = []
|
||||
|
||||
def optimize_image(self, image_data, target_width, format, quality, preserve_metadata):
|
||||
def optimize_image(
|
||||
self,
|
||||
image_data,
|
||||
target_width,
|
||||
format,
|
||||
quality,
|
||||
preserve_metadata,
|
||||
workflow=None,
|
||||
):
|
||||
self.optimized_calls += 1
|
||||
self.optimized_workflow = workflow
|
||||
return image_data, ".webp"
|
||||
|
||||
def append_recipe_metadata(self, image_path, recipe_data, pixel_preserving=False):
|
||||
self.appended = (image_path, recipe_data, pixel_preserving)
|
||||
|
||||
def embed_workflow(self, image_path, workflow):
|
||||
self.embedded_workflows.append((image_path, workflow))
|
||||
return image_path
|
||||
|
||||
def extract_image_metadata(self, path):
|
||||
return {}
|
||||
|
||||
@@ -131,6 +146,55 @@ async def test_save_recipe_skip_optimize_preserves_image_bytes(tmp_path):
|
||||
assert exif_utils.appended[2] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_recipe_skip_optimize_still_embeds_recovered_workflow(tmp_path):
|
||||
"""The verbatim branch bypasses optimize_image, so the recovered workflow
|
||||
has to be embedded by the explicit safety net."""
|
||||
image_buffer = BytesIO()
|
||||
Image.new("RGB", (96, 48), color="olive").save(
|
||||
image_buffer, format="WEBP", quality=85
|
||||
)
|
||||
|
||||
class DummyScanner:
|
||||
def __init__(self, root):
|
||||
self.recipes_dir = str(root / "recipes")
|
||||
|
||||
async def add_recipe(self, recipe_data):
|
||||
return None
|
||||
|
||||
async def find_recipes_by_fingerprint(self, fingerprint):
|
||||
return []
|
||||
|
||||
service = RecipePersistenceService(
|
||||
exif_utils=ExifUtils,
|
||||
card_preview_width=512,
|
||||
logger=logging.getLogger("test"),
|
||||
)
|
||||
|
||||
workflow = {"nodes": [{"id": 8}]}
|
||||
result = await service.save_recipe(
|
||||
recipe_scanner=DummyScanner(tmp_path),
|
||||
image_bytes=image_buffer.getvalue(),
|
||||
image_base64=None,
|
||||
name="Verbatim Workflow",
|
||||
tags=[],
|
||||
metadata={"base_model": "sd", "loras": [], "workflow": workflow},
|
||||
extension=".webp",
|
||||
skip_optimize=True,
|
||||
)
|
||||
|
||||
image_path = Path(result.payload["image_path"])
|
||||
with Image.open(image_path) as img:
|
||||
assert img.size == (96, 48)
|
||||
assert img.format == "WEBP"
|
||||
assert ExifUtils._load_structured_metadata(str(image_path))["workflow"] == (
|
||||
json.dumps(workflow)
|
||||
)
|
||||
|
||||
stored = json.loads(Path(result.payload["json_path"]).read_text())
|
||||
assert stored["has_workflow"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_recipe_skip_optimize_default_optimizes(tmp_path):
|
||||
"""Normal saves must keep optimizing; only re-import opts out."""
|
||||
@@ -650,6 +714,88 @@ async def test_save_recipe_preserves_workflow_when_png_is_converted_to_webp(tmp_
|
||||
assert "Recipe metadata:" in decoded_comment
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_recipe_embeds_workflow_recovered_from_source(tmp_path):
|
||||
"""Import paths hand the workflow over as metadata when their preview bytes
|
||||
are metadata-free (CivitAI's optimized rendition); save_recipe must embed
|
||||
it so the recipe reports has_workflow and can send it to ComfyUI."""
|
||||
class DummyScanner:
|
||||
def __init__(self, root):
|
||||
self.recipes_dir = str(root)
|
||||
|
||||
async def find_recipes_by_fingerprint(self, fingerprint):
|
||||
return []
|
||||
|
||||
async def add_recipe(self, recipe_data):
|
||||
return None
|
||||
|
||||
image_buffer = BytesIO()
|
||||
Image.new("RGB", (96, 48), color="teal").save(
|
||||
image_buffer, format="WEBP", quality=85
|
||||
)
|
||||
|
||||
service = RecipePersistenceService(
|
||||
exif_utils=ExifUtils,
|
||||
card_preview_width=64,
|
||||
logger=logging.getLogger("test"),
|
||||
)
|
||||
|
||||
workflow = {"nodes": [{"id": 1}], "last_node_id": 1}
|
||||
result = await service.save_recipe(
|
||||
recipe_scanner=DummyScanner(tmp_path),
|
||||
image_bytes=image_buffer.getvalue(),
|
||||
image_base64=None,
|
||||
name="Recovered Workflow",
|
||||
tags=["workflow"],
|
||||
metadata={"base_model": "sd", "loras": [], "workflow": workflow},
|
||||
extension=".webp",
|
||||
)
|
||||
|
||||
image_path = Path(result.payload["image_path"])
|
||||
assert ExifUtils._load_structured_metadata(str(image_path))["workflow"] == (
|
||||
json.dumps(workflow)
|
||||
)
|
||||
|
||||
stored = json.loads(Path(result.payload["json_path"]).read_text())
|
||||
assert stored["has_workflow"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_recipe_passes_recovered_workflow_to_optimizer(tmp_path):
|
||||
"""The workflow travels through optimize_image (single encode pass) rather
|
||||
than being patched in afterwards."""
|
||||
exif_utils = DummyExifUtils()
|
||||
|
||||
class DummyScanner:
|
||||
def __init__(self, root):
|
||||
self.recipes_dir = str(root)
|
||||
|
||||
async def find_recipes_by_fingerprint(self, fingerprint):
|
||||
return []
|
||||
|
||||
async def add_recipe(self, recipe_data):
|
||||
return None
|
||||
|
||||
workflow = '{"nodes": [{"id": 2}]}'
|
||||
service = RecipePersistenceService(
|
||||
exif_utils=exif_utils,
|
||||
card_preview_width=512,
|
||||
logger=logging.getLogger("test"),
|
||||
)
|
||||
|
||||
await service.save_recipe(
|
||||
recipe_scanner=DummyScanner(tmp_path),
|
||||
image_bytes=b"image-bytes",
|
||||
image_base64=None,
|
||||
name="Recovered Workflow",
|
||||
tags=[],
|
||||
metadata={"base_model": "sd", "loras": [], "workflow": workflow},
|
||||
extension=".webp",
|
||||
)
|
||||
|
||||
assert exif_utils.optimized_workflow == workflow
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_recipe_strips_checkpoint_local_fields(tmp_path):
|
||||
exif_utils = DummyExifUtils()
|
||||
@@ -1071,6 +1217,83 @@ async def test_analyze_remote_image_supports_civitai_red():
|
||||
assert result.payload["loras"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_remote_image_returns_workflow_from_original_rendition():
|
||||
"""CivitAI's optimized rendition is re-encoded and metadata-free, so an
|
||||
embedded workflow only exists in the original. Analysis must surface it so
|
||||
the save step can embed it while the stored preview stays the optimized
|
||||
image."""
|
||||
workflow = json.dumps({"nodes": [{"id": 1}], "last_node_id": 1})
|
||||
|
||||
class FakeExif:
|
||||
def extract_image_metadata(self, path):
|
||||
# The optimized rendition carries no metadata at all.
|
||||
return None
|
||||
|
||||
def _load_structured_metadata(self, path):
|
||||
# Only the original rendition (fetched to a .png temp file)
|
||||
# carries the embedded workflow.
|
||||
return {
|
||||
"parameters": None,
|
||||
"prompt": None,
|
||||
"workflow": workflow if str(path).endswith(".png") else None,
|
||||
"comment": None,
|
||||
}
|
||||
|
||||
downloaded: list[str] = []
|
||||
|
||||
async def downloader_factory():
|
||||
class Downloader:
|
||||
async def download_file(self, url, path, use_auth=False):
|
||||
downloaded.append(url)
|
||||
Path(path).write_bytes(b"fake-image")
|
||||
return True, "success"
|
||||
|
||||
return Downloader()
|
||||
|
||||
class DummyFactory:
|
||||
def create_parser(self, metadata):
|
||||
async def parse_metadata(m, recipe_scanner=None, civitai_client=None):
|
||||
return {"loras": [], "gen_params": {"prompt": "p"}}
|
||||
|
||||
return SimpleNamespace(parse_metadata=parse_metadata)
|
||||
|
||||
service = RecipeAnalysisService(
|
||||
exif_utils=FakeExif(),
|
||||
recipe_parser_factory=DummyFactory(),
|
||||
downloader_factory=downloader_factory,
|
||||
metadata_collector=None,
|
||||
metadata_processor_cls=None,
|
||||
metadata_registry_cls=None,
|
||||
standalone_mode=False,
|
||||
logger=logging.getLogger("test"),
|
||||
)
|
||||
|
||||
class DummyClient:
|
||||
async def get_image_info(self, image_id, source_url=None):
|
||||
return {
|
||||
"url": "https://image.civitai.com/x/original=true/sample.jpeg",
|
||||
"type": "image",
|
||||
"meta": {"prompt": "p"},
|
||||
}
|
||||
|
||||
class DummyScanner:
|
||||
async def find_recipes_by_fingerprint(self, fingerprint):
|
||||
return []
|
||||
|
||||
result = await service.analyze_remote_image(
|
||||
url="https://civitai.red/images/143518055",
|
||||
recipe_scanner=DummyScanner(),
|
||||
civitai_client=DummyClient(),
|
||||
)
|
||||
|
||||
assert result.payload["workflow"] == workflow
|
||||
# The optimized rendition is used as the preview, the original only as the
|
||||
# metadata/workflow fallback.
|
||||
assert any("width=450,optimized=true" in url for url in downloaded)
|
||||
assert any("original=true" in url for url in downloaded)
|
||||
|
||||
|
||||
def _exif_utils_returning(metadata):
|
||||
class MetadataExifUtils(DummyExifUtils):
|
||||
def extract_image_metadata(self, path):
|
||||
|
||||
@@ -211,6 +211,104 @@ def test_update_image_metadata_preserves_png_workflow(tmp_path):
|
||||
)
|
||||
|
||||
|
||||
def test_optimize_image_embeds_supplied_workflow_when_source_has_none(tmp_path):
|
||||
"""Import paths hand the workflow over as data when the preview source is
|
||||
metadata-free (CivitAI's optimized rendition); optimize_image must embed
|
||||
it while re-encoding, otherwise the recipe loses has_workflow."""
|
||||
image_path = tmp_path / "optimized.webp"
|
||||
Image.new("RGB", (64, 32), color="red").save(image_path, format="WEBP", quality=85)
|
||||
|
||||
workflow = {"nodes": [{"id": 1}], "last_node_id": 1}
|
||||
optimized_data, extension = ExifUtils.optimize_image(
|
||||
str(image_path),
|
||||
target_width=32,
|
||||
format="webp",
|
||||
quality=85,
|
||||
preserve_metadata=True,
|
||||
workflow=workflow,
|
||||
)
|
||||
|
||||
optimized_path = tmp_path / f"embedded{extension}"
|
||||
optimized_path.write_bytes(optimized_data)
|
||||
|
||||
metadata = ExifUtils._load_structured_metadata(str(optimized_path))
|
||||
assert metadata["workflow"] == json.dumps(workflow)
|
||||
|
||||
|
||||
def test_optimize_image_keeps_source_workflow_over_supplied(tmp_path):
|
||||
image_path = tmp_path / "source.png"
|
||||
png_info = PngImagePlugin.PngInfo()
|
||||
png_info.add_text("workflow", '{"nodes":[{"id":7}]}')
|
||||
Image.new("RGB", (64, 32), color="red").save(image_path, pnginfo=png_info)
|
||||
|
||||
optimized_data, extension = ExifUtils.optimize_image(
|
||||
str(image_path),
|
||||
target_width=32,
|
||||
format="webp",
|
||||
quality=85,
|
||||
preserve_metadata=True,
|
||||
workflow={"nodes": [{"id": 1}]},
|
||||
)
|
||||
|
||||
optimized_path = tmp_path / f"sourcewins{extension}"
|
||||
optimized_path.write_bytes(optimized_data)
|
||||
|
||||
metadata = ExifUtils._load_structured_metadata(str(optimized_path))
|
||||
assert metadata["workflow"] == '{"nodes":[{"id":7}]}'
|
||||
|
||||
|
||||
def test_embed_workflow_adds_workflow_to_metadata_free_webp(tmp_path):
|
||||
image_path = tmp_path / "preview.webp"
|
||||
Image.new("RGB", (32, 32), color="blue").save(image_path, format="WEBP", quality=85)
|
||||
|
||||
workflow = json.dumps({"nodes": [{"id": 1}]})
|
||||
returned = ExifUtils.embed_workflow(str(image_path), workflow)
|
||||
|
||||
assert returned == str(image_path)
|
||||
metadata = ExifUtils._load_structured_metadata(str(image_path))
|
||||
assert metadata["workflow"] == workflow
|
||||
with Image.open(image_path) as img:
|
||||
assert img.size == (32, 32)
|
||||
|
||||
|
||||
def test_embed_workflow_adds_workflow_to_metadata_free_png(tmp_path):
|
||||
image_path = tmp_path / "preview.png"
|
||||
Image.new("RGB", (32, 32), color="blue").save(image_path)
|
||||
|
||||
workflow = {"nodes": [{"id": 3}]}
|
||||
ExifUtils.embed_workflow(str(image_path), workflow)
|
||||
|
||||
metadata = ExifUtils._load_structured_metadata(str(image_path))
|
||||
assert metadata["workflow"] == json.dumps(workflow)
|
||||
|
||||
|
||||
def test_embed_workflow_leaves_existing_workflow_untouched(tmp_path):
|
||||
image_path = tmp_path / "preview.png"
|
||||
png_info = PngImagePlugin.PngInfo()
|
||||
png_info.add_text("workflow", '{"nodes":[{"id":9}]}')
|
||||
Image.new("RGB", (32, 32), color="green").save(image_path, pnginfo=png_info)
|
||||
|
||||
ExifUtils.embed_workflow(str(image_path), {"nodes": [{"id": 1}]})
|
||||
|
||||
with Image.open(image_path) as img:
|
||||
assert img.info["workflow"] == '{"nodes":[{"id":9}]}'
|
||||
|
||||
|
||||
def test_embed_workflow_ignores_unsupported_payloads_and_containers(tmp_path):
|
||||
image_path = tmp_path / "preview.webp"
|
||||
Image.new("RGB", (16, 16), color="black").save(image_path, format="WEBP")
|
||||
|
||||
# Nothing to embed / unsupported payload types are no-ops.
|
||||
assert ExifUtils.embed_workflow(str(image_path), None) == str(image_path)
|
||||
assert ExifUtils.embed_workflow(str(image_path), "") == str(image_path)
|
||||
assert ExifUtils.embed_workflow(str(image_path), 123) == str(image_path)
|
||||
assert ExifUtils._load_structured_metadata(str(image_path))["workflow"] is None
|
||||
|
||||
video_path = tmp_path / "clip.mp4"
|
||||
video_path.write_bytes(b"video")
|
||||
assert ExifUtils.embed_workflow(str(video_path), {"nodes": []}) == str(video_path)
|
||||
|
||||
|
||||
# --- ISOBMFF / brotli extraction tests ---
|
||||
|
||||
import struct
|
||||
|
||||
Reference in New Issue
Block a user