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