alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-09-28 22:44:09 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat(recipes): preserve embedded ComfyUI workflow on remote imports

Browse Source 
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:
Will Miao
2026-09-29 07:20:12 +08:00
parent 0dd8d74032
commit 69691b17a1
11 changed files with 1059 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/routes/handlers/recipe_handlers.py
+126 -40
View File
@@ -1284,6 +1284,21 @@ class RecipeManagementHandler:
_original_image_url,
) = await self._download_remote_media(image_url)
# CivitAI's optimized rendition is re-encoded and metadata-free, so an
# embedded ComfyUI workflow only exists in the original. Fetch it
# lazily: unlike the URL import path (which needs the original for
# metadata parsing anyway), this path would download it purely for the
# workflow, so it is skipped unless the API reports one.
original_workflow = None
if _original_image_url and self._meta_indicates_comfy_workflow(
civitai_meta_raw
):
_raw_original, original_workflow = await self._fetch_original_media(
_original_image_url
)
if original_workflow:
metadata["workflow"] = original_workflow
# Build a version-cached map of local model hashes to cache items so
# CivitaiApiMetadataParser can skip CivitAI API calls for models that
# exist on disk. Built once and shared by every parse pass below.
@@ -2090,6 +2105,90 @@ class RecipeManagementHandler:
except FileNotFoundError:
pass
def _read_embedded_workflow(self, image_path: Optional[str]) -> Optional[str]:
"""Return a ComfyUI workflow embedded in ``image_path``, if any.
``ExifUtils.extract_image_metadata`` stops at the generation
parameters, so the UI-format workflow has to be read through the
structured metadata reader. Failures map to ``None``.
"""
if not image_path or not os.path.exists(image_path):
return None
try:
metadata = ExifUtils._load_structured_metadata(image_path)
except Exception as exc:
self._logger.debug(
"Failed to read embedded workflow from %s: %s", image_path, exc
)
return None
workflow = metadata.get("workflow") if isinstance(metadata, dict) else None
return workflow if isinstance(workflow, str) and workflow else None
@staticmethod
def _meta_indicates_comfy_workflow(civitai_meta_raw: Any) -> bool:
"""Whether CivitAI reports an embedded ComfyUI workflow for an image.
``meta.comfy`` is the payload CivitAI captured from the original image,
so its presence is the signal that fetching the original is worth the
bandwidth when the caller does not already need it for metadata
parsing.
"""
if not isinstance(civitai_meta_raw, dict):
return False
inner = civitai_meta_raw.get("meta")
if isinstance(inner, dict) and inner.get("comfy"):
return True
return bool(civitai_meta_raw.get("comfy"))
async def _fetch_original_media(
self, original_image_url: Optional[str]
) -> tuple[Optional[str], Optional[str]]:
"""Download the original rendition and read its embedded media.
CivitAI's optimized renditions are re-encoded and carry no metadata, so
the original is the only source for embedded generation metadata and
for the UI-format ComfyUI workflow (the raw extractor's fallback chain
ends at ``workflow`` only when no prompt is present).
Returns ``(raw_metadata, workflow)``; either element is ``None`` when
unavailable. Failures never raise — imports keep working with the
optimized rendition when the original cannot be fetched.
"""
if not original_image_url:
return None, None
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
temp_path = temp_file.name
try:
downloader = await self._downloader_factory()
success, _result = await downloader.download_file(
original_image_url, temp_path, use_auth=False
)
if not success:
self._logger.warning(
"Failed to download original rendition: %s", original_image_url
)
return None, None
raw_metadata = await asyncio.to_thread(
ExifUtils.extract_image_metadata, temp_path
)
workflow = await asyncio.to_thread(
self._read_embedded_workflow, temp_path
)
return raw_metadata, workflow
except Exception as exc:
self._logger.warning(
"Failed to read original rendition %s: %s", original_image_url, exc
)
return None, None
finally:
try:
if os.path.exists(temp_path):
os.unlink(temp_path)
except OSError:
pass
def _safe_int(self, value: Any) -> int:
try:
return int(value)
@@ -2295,6 +2394,7 @@ class RecipeManagementHandler:
"Failed to extract embedded metadata: %s", exc
)
original_workflow: Optional[str] = None
if not parsed_embedded and original_image_url:
self._logger.debug(
"Optimized image has no embedded metadata, "
@@ -2302,48 +2402,32 @@ class RecipeManagementHandler:
original_image_url,
)
try:
downloader = await self._downloader_factory()
with tempfile.NamedTemporaryFile(
suffix=".png", delete=False
) as tmp:
orig_tmp_path = tmp.name
try:
success, _ = await downloader.download_file(
original_image_url, orig_tmp_path, use_auth=False
)
if success:
raw_orig = await asyncio.to_thread(
ExifUtils.extract_image_metadata, orig_tmp_path
raw_orig, original_workflow = await self._fetch_original_media(
original_image_url
)
diagnostics["exif_present"] = bool(raw_orig) or bool(
diagnostics.get("exif_present")
)
if raw_orig:
parser = (
self._analysis_service._recipe_parser_factory.create_parser(
raw_orig
)
diagnostics["exif_present"] = bool(raw_orig)
if raw_orig:
parser = (
self._analysis_service._recipe_parser_factory.create_parser(
raw_orig
)
)
if parser:
diagnostics["exif_parser"] = parser.__class__.__name__
if isinstance(parser, CivitaiApiMetadataParser):
parsed_embedded = await parser.parse_metadata(
raw_orig,
recipe_scanner=recipe_scanner,
local_cache=local_cache,
)
if parser:
diagnostics["exif_parser"] = parser.__class__.__name__
if isinstance(parser, CivitaiApiMetadataParser):
parsed_embedded = await parser.parse_metadata(
raw_orig,
recipe_scanner=recipe_scanner,
local_cache=local_cache,
)
else:
parsed_embedded = await parser.parse_metadata(
raw_orig, recipe_scanner=recipe_scanner
)
if (
parsed_embedded
and "gen_params" in parsed_embedded
):
embedded_gen_params = parsed_embedded[
"gen_params"
]
finally:
if os.path.exists(orig_tmp_path):
os.unlink(orig_tmp_path)
else:
parsed_embedded = await parser.parse_metadata(
raw_orig, recipe_scanner=recipe_scanner
)
if parsed_embedded and "gen_params" in parsed_embedded:
embedded_gen_params = parsed_embedded["gen_params"]
except Exception as exc:
self._logger.warning(
"Failed to extract metadata from original image: %s", exc
@@ -2391,6 +2475,8 @@ class RecipeManagementHandler:
"gen_params": embedded_gen_params or {},
"source_path": image_url,
}
if original_workflow:
metadata["workflow"] = original_workflow
# Extract preview_nsfw_level from the CivitAI API response
# (injected into civitai_meta_raw by _download_remote_media).
py/services/batch_import_service.py
+5
View File
@@ -645,6 +645,11 @@ class BatchImportService:
if payload.get("checkpoint"):
metadata["checkpoint"] = payload["checkpoint"]
# A workflow recovered from the source's original rendition
# travels as metadata and is embedded into the stored image.
if payload.get("workflow"):
metadata["workflow"] = payload["workflow"]
nsfw = payload.get("preview_nsfw_level")
if isinstance(nsfw, int) and nsfw > 0:
metadata["preview_nsfw_level"] = nsfw
py/services/recipes/analysis_service.py
+38
View File
@@ -117,6 +117,10 @@ class RecipeAnalysisService:
image_info: Optional[dict[str, Any]] = None
is_video = False
extension = ".jpg" # Default
# Workflow recovered from the image. CivitAI's optimized renditions are
# re-encoded and carry no metadata, so for those the workflow only
# exists in the original rendition, fetched below for EXIF extraction.
recovered_workflow: Optional[str] = None
# Diagnostics collected during analysis; surfaced in the payload so
# callers can persist an import_info block explaining empty LoRA lists.
diagnostics: dict[str, Any] = {"channel": "url"}
@@ -238,6 +242,9 @@ class RecipeAnalysisService:
exif_metadata = await asyncio.to_thread(
self._exif_utils.extract_image_metadata, temp_path
)
recovered_workflow = await asyncio.to_thread(
self._read_embedded_workflow, temp_path
)
# Fallback: try the original (non-optimized) image for EXIF data
if not exif_metadata and civitai_image_id and image_info:
@@ -255,6 +262,16 @@ class RecipeAnalysisService:
self._exif_utils.extract_image_metadata,
orig_temp_path,
)
# The original is also the only place a ComfyUI
# workflow survives; carry it so the save step can
# embed it even though the stored preview stays the
# small, metadata-free optimized rendition.
recovered_workflow = (
await asyncio.to_thread(
self._read_embedded_workflow, orig_temp_path
)
or recovered_workflow
)
finally:
self._safe_cleanup(orig_temp_path)
@@ -358,6 +375,8 @@ class RecipeAnalysisService:
diagnostics["is_video"] = is_video
result.payload["diagnostics"] = diagnostics
if recovered_workflow:
result.payload["workflow"] = recovered_workflow
return result
finally:
if temp_path:
@@ -545,6 +564,25 @@ class RecipeAnalysisService:
if not success:
raise RecipeDownloadError(f"Failed to download image from URL: {result}")
def _read_embedded_workflow(self, image_path: Optional[str]) -> Optional[str]:
"""Return a ComfyUI workflow embedded in ``image_path``, if any.
The raw metadata string extractor stops at the generation parameters
(``prompt``/``parameters``), so the UI-format workflow has to be read
through the structured metadata reader. Failures map to ``None``.
"""
if not image_path or not os.path.exists(image_path):
return None
try:
metadata = self._exif_utils._load_structured_metadata(image_path)
except Exception as exc:
self._logger.debug(
"Failed to read embedded workflow from %s: %s", image_path, exc
)
return None
workflow = metadata.get("workflow") if isinstance(metadata, dict) else None
return workflow if isinstance(workflow, str) and workflow else None
def _metadata_not_found_response(self, path: str) -> AnalysisResult:
payload: dict[str, Any] = {
"error": "No metadata found in this image",
py/services/recipes/persistence_service.py
+19
View File
@@ -73,6 +73,11 @@ class RecipePersistenceService:
byte-level EXIF update that leaves the pixels untouched). Used
by local re-import, where the source is the recipe's own
already-optimized preview image.
``metadata`` may carry a ``workflow`` entry (JSON string, dict or
list) recovered from the source's original rendition; it is embedded
into the stored image so the recipe reports ``has_workflow`` and can
send the workflow back to ComfyUI.
"""
missing_fields = []
@@ -87,6 +92,13 @@ class RecipePersistenceService:
assert metadata is not None
# A workflow recovered from a higher-fidelity source (CivitAI's
# original rendition — its optimized preview is re-encoded and carries
# no metadata) travels as data instead of as image bytes. It is
# embedded below so ``has_workflow`` and the "send workflow to ComfyUI"
# action work for imports whose preview pixels are metadata-free.
workflow = metadata.get("workflow")
resolved_image_bytes = self._resolve_image_bytes(image_bytes, image_base64)
recipes_dir = target_dir or recipe_scanner.recipes_dir
os.makedirs(recipes_dir, exist_ok=True)
@@ -108,6 +120,7 @@ class RecipePersistenceService:
format="webp",
quality=85,
preserve_metadata=True,
workflow=workflow,
)
image_filename = f"{recipe_id}{extension}"
@@ -116,6 +129,12 @@ class RecipePersistenceService:
with open(normalized_image_path, "wb") as file_obj:
file_obj.write(optimized_image)
# The optimization branch above embeds the workflow while re-encoding;
# the verbatim (skip_optimize) branch still needs it added, and this is
# also the safety net when re-encoding dropped it.
if workflow and not is_video:
self._exif_utils.embed_workflow(normalized_image_path, workflow)
current_time = time.time()
loras_data = [self._normalise_lora_entry(lora) for lora in (metadata.get("loras") or [])]
checkpoint_entry = self._sanitize_checkpoint_entry(self._extract_checkpoint_entry(metadata))
py/utils/exif_utils.py
+134 -23
View File
@@ -341,29 +341,125 @@ class ExifUtils:
metadata_fields = ExifUtils._load_structured_metadata(image_path)
metadata_fields["parameters"] = metadata
with Image.open(image_path) as img:
img_format = img.format
if img_format == "PNG":
png_info = ExifUtils._build_pnginfo(img, metadata_fields)
img.save(image_path, format="PNG", pnginfo=png_info)
return image_path
exif_bytes = ExifUtils._build_exif_bytes(
metadata_fields, img.info.get("exif")
)
save_kwargs: dict[str, Any] = {"exif": exif_bytes}
if img_format == "WEBP":
save_kwargs["quality"] = 85
img.save(image_path, format=img_format, **save_kwargs)
return image_path
return ExifUtils._write_structured_metadata(image_path, metadata_fields)
except Exception as e:
logger.error(f"Error updating metadata in {image_path}: {e}")
return image_path
@staticmethod
def _write_structured_metadata(
image_path: str, metadata_fields: dict[str, Optional[str]]
) -> str:
"""Write structured metadata fields back into an image.
PNG keeps them as text chunks (``parameters``/``prompt``/``workflow``);
every other supported container stores them in EXIF, where the workflow
travels in ``ImageDescription`` behind a ``Workflow:`` prefix (see
:meth:`_build_exif_bytes`).
"""
with Image.open(image_path) as img:
img_format = img.format
if img_format == "PNG":
png_info = ExifUtils._build_pnginfo(img, metadata_fields)
img.save(image_path, format="PNG", pnginfo=png_info)
return image_path
exif_bytes = ExifUtils._build_exif_bytes(
metadata_fields, img.info.get("exif")
)
save_kwargs: dict[str, Any] = {"exif": exif_bytes}
if img_format == "WEBP":
save_kwargs["quality"] = 85
img.save(image_path, format=img_format, **save_kwargs)
return image_path
@staticmethod
def normalise_workflow(workflow: Any) -> Optional[str]:
"""Coerce a workflow payload into the JSON string metadata form.
Accepts the string form stored in image chunks as well as already
decoded dict/list payloads; anything else yields ``None``.
"""
if isinstance(workflow, str):
return workflow or None
if isinstance(workflow, (dict, list)):
try:
return json.dumps(workflow)
except (TypeError, ValueError):
return None
return None
@staticmethod
def _merge_workflow(
metadata_fields: Optional[dict[str, Optional[str]]], workflow: Any
) -> Optional[dict[str, Optional[str]]]:
"""Add a caller-supplied workflow to extracted metadata fields.
Returns ``metadata_fields`` untouched when there is nothing to add, and
never overwrites a workflow the source image already carries.
"""
workflow_json = ExifUtils.normalise_workflow(workflow)
if not workflow_json:
return metadata_fields
if metadata_fields is None:
metadata_fields = {
"parameters": None,
"prompt": None,
"workflow": None,
"comment": None,
}
if not metadata_fields.get("workflow"):
metadata_fields["workflow"] = workflow_json
return metadata_fields
@staticmethod
def embed_workflow(image_path: str, workflow: Any) -> str:
"""Embed a ComfyUI workflow into an image that does not carry one.
Recipe imports recover the workflow from the source's original
rendition (CivitAI's optimized preview is re-encoded and metadata-free)
and hand it over as data rather than as image bytes. Images that
already embed a workflow are left untouched.
WebP files are patched at the byte level so preview pixels are not
re-encoded a second time.
"""
workflow_json = ExifUtils.normalise_workflow(workflow)
if not image_path or not workflow_json:
return image_path
ext = os.path.splitext(image_path)[1].lower()
if ext in ['.mp4', '.webm', '.avif', '.jxl']:
return image_path
try:
metadata_fields = ExifUtils._load_structured_metadata(image_path)
if metadata_fields.get("workflow"):
return image_path
metadata_fields["workflow"] = workflow_json
if ext == '.webp':
try:
exif_bytes = ExifUtils._build_exif_bytes(metadata_fields)
with open(image_path, "rb") as file_obj:
image_bytes = file_obj.read()
updated = ExifUtils._replace_webp_exif(image_bytes, exif_bytes)
with open(image_path, "wb") as file_obj:
file_obj.write(updated)
return image_path
except ValueError:
# Container without an EXIF chunk: fall through to a full
# rewrite so the workflow is still embedded.
pass
return ExifUtils._write_structured_metadata(image_path, metadata_fields)
except Exception as e:
logger.error(f"Error embedding workflow in {image_path}: {e}")
return image_path
@staticmethod
def append_recipe_metadata(image_path, recipe_data, pixel_preserving=False) -> str:
"""Append recipe metadata to an image's EXIF data
@@ -550,7 +646,7 @@ class ExifUtils:
return None
@staticmethod
def optimize_image(image_data, target_width=250, format='webp', quality=85, preserve_metadata=False):
def optimize_image(image_data, target_width=250, format='webp', quality=85, preserve_metadata=False, workflow=None):
"""
Optimize an image by resizing and converting to WebP format
@@ -560,10 +656,19 @@ class ExifUtils:
format: Output format (default: webp)
quality: Output quality (0-100)
preserve_metadata: Whether to preserve EXIF metadata
workflow: Optional ComfyUI workflow (JSON string, dict or list) to
embed when the source image does not carry one. Used by import
paths that recover the workflow from a higher-fidelity source
(e.g. CivitAI's original rendition) while the preview pixels
come from a metadata-free optimized rendition.
Returns:
Tuple of (optimized_image_data, extension)
"""
# A supplied workflow can only survive when metadata is embedded, so
# treat it as an implicit request for preservation.
if workflow is not None:
preserve_metadata = True
try:
if isinstance(image_data, str) and os.path.exists(image_data):
ext = os.path.splitext(image_data)[1].lower()
@@ -627,6 +732,12 @@ class ExifUtils:
logger.warning(f"Failed to extract metadata, continuing without it: {e}")
# Continue without metadata
# Merge in a workflow recovered elsewhere (e.g. from CivitAI's
# original rendition). The source image wins when it already has
# one, and this is what lets the metadata-free optimized preview
# still end up with the workflow embedded.
metadata_fields = ExifUtils._merge_workflow(metadata_fields, workflow)
# Calculate new height to maintain aspect ratio
width, height = img.size
new_height = int(height * (target_width / width))
@@ -686,8 +797,8 @@ class ExifUtils:
temp_file.write(optimized_data)
try:
ExifUtils.update_image_metadata(
temp_path, metadata_fields.get("parameters") or ""
ExifUtils._write_structured_metadata(
temp_path, metadata_fields
)
# Read back the file
with open(temp_path, 'rb') as f:
static/js/managers/import/DownloadManager.js
+9
View File
@@ -72,6 +72,15 @@ export class DownloadManager {
completeMetadata.diagnostics = diagnostics;
}
// A ComfyUI workflow recovered from the source's original
// rendition: CivitAI's optimized preview is re-encoded and
// metadata-free, so the workflow travels as data and the
// backend embeds it into the stored image.
const workflow = this.importManager.recipeData.workflow;
if (workflow) {
completeMetadata.workflow = workflow;
}
// Preserve preview_nsfw_level from analysis so the saved
// recipe applies the correct NSFW blur on the preview image.
const nsfwLevel = this.importManager.recipeData.preview_nsfw_level;
tests/frontend/managers/recipeImportWorkflow.test.js
+101
View File
@@ -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);
});
});
tests/routes/test_recipe_import_workflow.py
+262
View File
@@ -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
tests/services/test_batch_import_service.py
+43
View File
@@ -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()
tests/services/test_recipe_services.py
+224 -1
View File
@@ -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):
tests/utils/test_exif_utils.py
+98
View File
@@ -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