From 69691b17a107a552dc0e415af139fd39946b2be6 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Tue, 29 Sep 2026 07:20:12 +0800 Subject: [PATCH] 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). --- py/routes/handlers/recipe_handlers.py | 166 ++++++++--- py/services/batch_import_service.py | 5 + py/services/recipes/analysis_service.py | 38 +++ py/services/recipes/persistence_service.py | 19 ++ py/utils/exif_utils.py | 157 +++++++++-- static/js/managers/import/DownloadManager.js | 9 + .../managers/recipeImportWorkflow.test.js | 101 +++++++ tests/routes/test_recipe_import_workflow.py | 262 ++++++++++++++++++ tests/services/test_batch_import_service.py | 43 +++ tests/services/test_recipe_services.py | 225 ++++++++++++++- tests/utils/test_exif_utils.py | 98 +++++++ 11 files changed, 1059 insertions(+), 64 deletions(-) create mode 100644 tests/frontend/managers/recipeImportWorkflow.test.js create mode 100644 tests/routes/test_recipe_import_workflow.py diff --git a/py/routes/handlers/recipe_handlers.py b/py/routes/handlers/recipe_handlers.py index e0661e47..5f4b8d49 100644 --- a/py/routes/handlers/recipe_handlers.py +++ b/py/routes/handlers/recipe_handlers.py @@ -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). diff --git a/py/services/batch_import_service.py b/py/services/batch_import_service.py index 657cd152..f14b75cb 100644 --- a/py/services/batch_import_service.py +++ b/py/services/batch_import_service.py @@ -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 diff --git a/py/services/recipes/analysis_service.py b/py/services/recipes/analysis_service.py index e4b9f60b..7ed8cdfe 100644 --- a/py/services/recipes/analysis_service.py +++ b/py/services/recipes/analysis_service.py @@ -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", diff --git a/py/services/recipes/persistence_service.py b/py/services/recipes/persistence_service.py index 0892f7c0..9027a90b 100644 --- a/py/services/recipes/persistence_service.py +++ b/py/services/recipes/persistence_service.py @@ -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)) diff --git a/py/utils/exif_utils.py b/py/utils/exif_utils.py index f7eca769..c88de9b4 100644 --- a/py/utils/exif_utils.py +++ b/py/utils/exif_utils.py @@ -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: diff --git a/static/js/managers/import/DownloadManager.js b/static/js/managers/import/DownloadManager.js index 1cd3d665..92d85cb4 100644 --- a/static/js/managers/import/DownloadManager.js +++ b/static/js/managers/import/DownloadManager.js @@ -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; diff --git a/tests/frontend/managers/recipeImportWorkflow.test.js b/tests/frontend/managers/recipeImportWorkflow.test.js new file mode 100644 index 00000000..8197238e --- /dev/null +++ b/tests/frontend/managers/recipeImportWorkflow.test.js @@ -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); + }); +}); diff --git a/tests/routes/test_recipe_import_workflow.py b/tests/routes/test_recipe_import_workflow.py new file mode 100644 index 00000000..53acc64b --- /dev/null +++ b/tests/routes/test_recipe_import_workflow.py @@ -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 diff --git a/tests/services/test_batch_import_service.py b/tests/services/test_batch_import_service.py index 9232a4eb..17993800 100644 --- a/tests/services/test_batch_import_service.py +++ b/tests/services/test_batch_import_service.py @@ -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() diff --git a/tests/services/test_recipe_services.py b/tests/services/test_recipe_services.py index 987b4881..410e8f1b 100644 --- a/tests/services/test_recipe_services.py +++ b/tests/services/test_recipe_services.py @@ -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): diff --git a/tests/utils/test_exif_utils.py b/tests/utils/test_exif_utils.py index 09258b1a..c10f92f2 100644 --- a/tests/utils/test_exif_utils.py +++ b/tests/utils/test_exif_utils.py @@ -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