Compare commits

...
Author SHA1 Message Date
Will Miao faeb66a23d feat(recipes): opt-in workflow embedding for widget recipe saves
Add a "Save Recipe with Workflow" action next to "Save Recipe" in the LoRA
widget context menu. It posts the current UI-format graph alongside the save
request so the stored preview embeds it and the recipe can send the graph back
to ComfyUI. Embedding stays opt-in rather than folded into "Save Recipe": the
workflow is by far the largest metadata field and its widget values may carry
sensitive data.

- web/comfyui: new menu entry; saveRecipeDirectly({ embedWorkflow }) posts the
  UI graph and reports the outcome (embedded / skipped) via toasts.
- save_recipe_from_widget handler: reads an optional JSON workflow field so the
  long-standing body-less POST keeps working, including from cached clients.
- RecipePersistenceService.save_recipe_from_widget: embeds the graph through
  the existing optimize_image workflow path, derives has_workflow by detection,
  and skips graphs above MAX_WORKFLOW_EMBED_BYTES with workflow_skipped.
2026-09-29 09:03:26 +08:00
Will Miao 69691b17a1 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).
2026-09-29 07:20:12 +08:00
Will Miao 0dd8d74032 fix(ui): stop body data-theme from shadowing theme preset tokens
applyTheme() mirrors the active mode onto <body> as data-theme="dark", but
the theme preset is only ever written to <html>. The palette token blocks in
tokens/colors.css and base.css used the bare attribute selector
[data-theme="dark"], so <body> matched them on its own and re-declared the
default dark palette (#1a1a1a / #2d2d2d / ...) directly on the body, where it
shadowed the preset values inherited from <html>. Every non-default preset
therefore painted the selected accent over the default preset's background,
surface, text and border tokens, and flipped into that state ~200ms after
load, when initTheme() first touched <body> — the accent-tinted background
flash seen on reload and nav-tab switches. Reached only in dark mode, since
light mode has no [data-theme="light"] token block.

Scope the palette token blocks to :root so a data-theme attribute on any
descendant (only <body> has one) can no longer re-declare them; descendant
rules such as [data-theme="dark"] .foo still match through <html>. Add a
regression guard that fails on bare attribute token blocks.
2026-09-27 21:48:58 +08:00
22 changed files with 1747 additions and 102 deletions
+149 -40
View File
@@ -1284,6 +1284,21 @@ class RecipeManagementHandler:
_original_image_url, _original_image_url,
) = await self._download_remote_media(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 # Build a version-cached map of local model hashes to cache items so
# CivitaiApiMetadataParser can skip CivitAI API calls for models that # CivitaiApiMetadataParser can skip CivitAI API calls for models that
# exist on disk. Built once and shared by every parse pass below. # exist on disk. Built once and shared by every parse pass below.
@@ -1806,6 +1821,10 @@ class RecipeManagementHandler:
if recipe_scanner is None: if recipe_scanner is None:
raise RuntimeError("Recipe scanner unavailable") raise RuntimeError("Recipe scanner unavailable")
# Opt-in workflow embedding. The widget historically POSTs with no
# body at all, so a missing/empty body is not an error.
workflow = await self._read_optional_json_field(request, "workflow")
analysis = await self._analysis_service.analyze_widget_metadata( analysis = await self._analysis_service.analyze_widget_metadata(
recipe_scanner=recipe_scanner recipe_scanner=recipe_scanner
) )
@@ -1818,6 +1837,7 @@ class RecipeManagementHandler:
recipe_scanner=recipe_scanner, recipe_scanner=recipe_scanner,
metadata=metadata, metadata=metadata,
image_bytes=image_bytes, image_bytes=image_bytes,
workflow=workflow,
) )
return web.json_response(result.payload, status=result.status) return web.json_response(result.payload, status=result.status)
except RecipeValidationError as exc: except RecipeValidationError as exc:
@@ -1882,6 +1902,24 @@ class RecipeManagementHandler:
return [] return []
return [tag.strip() for tag in tag_text.split(",") if tag.strip()] return [tag.strip() for tag in tag_text.split(",") if tag.strip()]
async def _read_optional_json_field(
self, request: web.Request, field: str
) -> Any:
"""Read one field from an optional JSON request body.
Some callers (notably the widget's long-standing "Save Recipe" action)
POST with no body at all, and a stale cached extension may still do so
after a body is introduced. A missing, empty or malformed body is
therefore treated as "no value" rather than a request error.
"""
if not request.can_read_body:
return None
try:
data = await request.json()
except Exception:
return None
return data.get(field) if isinstance(data, dict) else None
async def _count_recipe_loras( async def _count_recipe_loras(
self, recipe_scanner: Any, recipe_id: Optional[str] self, recipe_scanner: Any, recipe_id: Optional[str]
) -> Optional[int]: ) -> Optional[int]:
@@ -2090,6 +2128,90 @@ class RecipeManagementHandler:
except FileNotFoundError: except FileNotFoundError:
pass 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: def _safe_int(self, value: Any) -> int:
try: try:
return int(value) return int(value)
@@ -2295,6 +2417,7 @@ class RecipeManagementHandler:
"Failed to extract embedded metadata: %s", exc "Failed to extract embedded metadata: %s", exc
) )
original_workflow: Optional[str] = None
if not parsed_embedded and original_image_url: if not parsed_embedded and original_image_url:
self._logger.debug( self._logger.debug(
"Optimized image has no embedded metadata, " "Optimized image has no embedded metadata, "
@@ -2302,48 +2425,32 @@ class RecipeManagementHandler:
original_image_url, original_image_url,
) )
try: try:
downloader = await self._downloader_factory() raw_orig, original_workflow = await self._fetch_original_media(
with tempfile.NamedTemporaryFile( original_image_url
suffix=".png", delete=False )
) as tmp: diagnostics["exif_present"] = bool(raw_orig) or bool(
orig_tmp_path = tmp.name diagnostics.get("exif_present")
try: )
success, _ = await downloader.download_file( if raw_orig:
original_image_url, orig_tmp_path, use_auth=False parser = (
) self._analysis_service._recipe_parser_factory.create_parser(
if success: raw_orig
raw_orig = await asyncio.to_thread(
ExifUtils.extract_image_metadata, orig_tmp_path
) )
diagnostics["exif_present"] = bool(raw_orig) )
if raw_orig: if parser:
parser = ( diagnostics["exif_parser"] = parser.__class__.__name__
self._analysis_service._recipe_parser_factory.create_parser( if isinstance(parser, CivitaiApiMetadataParser):
raw_orig parsed_embedded = await parser.parse_metadata(
) raw_orig,
recipe_scanner=recipe_scanner,
local_cache=local_cache,
) )
if parser: else:
diagnostics["exif_parser"] = parser.__class__.__name__ parsed_embedded = await parser.parse_metadata(
if isinstance(parser, CivitaiApiMetadataParser): raw_orig, recipe_scanner=recipe_scanner
parsed_embedded = await parser.parse_metadata( )
raw_orig, if parsed_embedded and "gen_params" in parsed_embedded:
recipe_scanner=recipe_scanner, embedded_gen_params = parsed_embedded["gen_params"]
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)
except Exception as exc: except Exception as exc:
self._logger.warning( self._logger.warning(
"Failed to extract metadata from original image: %s", exc "Failed to extract metadata from original image: %s", exc
@@ -2391,6 +2498,8 @@ class RecipeManagementHandler:
"gen_params": embedded_gen_params or {}, "gen_params": embedded_gen_params or {},
"source_path": image_url, "source_path": image_url,
} }
if original_workflow:
metadata["workflow"] = original_workflow
# Extract preview_nsfw_level from the CivitAI API response # Extract preview_nsfw_level from the CivitAI API response
# (injected into civitai_meta_raw by _download_remote_media). # (injected into civitai_meta_raw by _download_remote_media).
+5
View File
@@ -645,6 +645,11 @@ class BatchImportService:
if payload.get("checkpoint"): if payload.get("checkpoint"):
metadata["checkpoint"] = payload["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") nsfw = payload.get("preview_nsfw_level")
if isinstance(nsfw, int) and nsfw > 0: if isinstance(nsfw, int) and nsfw > 0:
metadata["preview_nsfw_level"] = nsfw metadata["preview_nsfw_level"] = nsfw
+38
View File
@@ -117,6 +117,10 @@ class RecipeAnalysisService:
image_info: Optional[dict[str, Any]] = None image_info: Optional[dict[str, Any]] = None
is_video = False is_video = False
extension = ".jpg" # Default 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 # Diagnostics collected during analysis; surfaced in the payload so
# callers can persist an import_info block explaining empty LoRA lists. # callers can persist an import_info block explaining empty LoRA lists.
diagnostics: dict[str, Any] = {"channel": "url"} diagnostics: dict[str, Any] = {"channel": "url"}
@@ -238,6 +242,9 @@ class RecipeAnalysisService:
exif_metadata = await asyncio.to_thread( exif_metadata = await asyncio.to_thread(
self._exif_utils.extract_image_metadata, temp_path 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 # Fallback: try the original (non-optimized) image for EXIF data
if not exif_metadata and civitai_image_id and image_info: if not exif_metadata and civitai_image_id and image_info:
@@ -255,6 +262,16 @@ class RecipeAnalysisService:
self._exif_utils.extract_image_metadata, self._exif_utils.extract_image_metadata,
orig_temp_path, 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: finally:
self._safe_cleanup(orig_temp_path) self._safe_cleanup(orig_temp_path)
@@ -358,6 +375,8 @@ class RecipeAnalysisService:
diagnostics["is_video"] = is_video diagnostics["is_video"] = is_video
result.payload["diagnostics"] = diagnostics result.payload["diagnostics"] = diagnostics
if recovered_workflow:
result.payload["workflow"] = recovered_workflow
return result return result
finally: finally:
if temp_path: if temp_path:
@@ -545,6 +564,25 @@ class RecipeAnalysisService:
if not success: if not success:
raise RecipeDownloadError(f"Failed to download image from URL: {result}") 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: def _metadata_not_found_response(self, path: str) -> AnalysisResult:
payload: dict[str, Any] = { payload: dict[str, Any] = {
"error": "No metadata found in this image", "error": "No metadata found in this image",
+55 -13
View File
@@ -18,6 +18,7 @@ from ...utils.base_model import (
RELATION_INCOMPATIBLE, RELATION_INCOMPATIBLE,
base_model_relation, base_model_relation,
) )
from ...utils.constants import MAX_WORKFLOW_EMBED_BYTES
from ...utils.utils import calculate_recipe_fingerprint from ...utils.utils import calculate_recipe_fingerprint
from ..pending_delete_service import get_pending_delete_service from ..pending_delete_service import get_pending_delete_service
from .errors import RecipeNotFoundError, RecipeValidationError from .errors import RecipeNotFoundError, RecipeValidationError
@@ -73,6 +74,11 @@ class RecipePersistenceService:
byte-level EXIF update that leaves the pixels untouched). Used byte-level EXIF update that leaves the pixels untouched). Used
by local re-import, where the source is the recipe's own by local re-import, where the source is the recipe's own
already-optimized preview image. 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 = [] missing_fields = []
@@ -87,6 +93,13 @@ class RecipePersistenceService:
assert metadata is not None 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) resolved_image_bytes = self._resolve_image_bytes(image_bytes, image_base64)
recipes_dir = target_dir or recipe_scanner.recipes_dir recipes_dir = target_dir or recipe_scanner.recipes_dir
os.makedirs(recipes_dir, exist_ok=True) os.makedirs(recipes_dir, exist_ok=True)
@@ -108,6 +121,7 @@ class RecipePersistenceService:
format="webp", format="webp",
quality=85, quality=85,
preserve_metadata=True, preserve_metadata=True,
workflow=workflow,
) )
image_filename = f"{recipe_id}{extension}" image_filename = f"{recipe_id}{extension}"
@@ -116,6 +130,12 @@ class RecipePersistenceService:
with open(normalized_image_path, "wb") as file_obj: with open(normalized_image_path, "wb") as file_obj:
file_obj.write(optimized_image) 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() current_time = time.time()
loras_data = [self._normalise_lora_entry(lora) for lora in (metadata.get("loras") or [])] 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)) checkpoint_entry = self._sanitize_checkpoint_entry(self._extract_checkpoint_entry(metadata))
@@ -855,8 +875,15 @@ class RecipePersistenceService:
recipe_scanner, recipe_scanner,
metadata: dict[str, Any], metadata: dict[str, Any],
image_bytes: bytes, image_bytes: bytes,
workflow: Any = None,
) -> PersistenceResult: ) -> PersistenceResult:
"""Save a recipe constructed from widget metadata.""" """Save a recipe constructed from widget metadata.
``workflow`` is the caller's ComfyUI graph (UI or API format) to embed
in the stored preview. Embedding is opt-in because the graph is by far
the largest metadata field and its widget values may contain sensitive
data; an oversized graph is dropped rather than inflating the preview.
"""
if not metadata: if not metadata:
raise RecipeValidationError("No generation metadata found") raise RecipeValidationError("No generation metadata found")
@@ -865,12 +892,25 @@ class RecipePersistenceService:
os.makedirs(recipes_dir, exist_ok=True) os.makedirs(recipes_dir, exist_ok=True)
recipe_id = str(uuid.uuid4()) recipe_id = str(uuid.uuid4())
workflow_json = self._exif_utils.normalise_workflow(workflow)
workflow_skipped: Optional[str] = None
if workflow_json and len(workflow_json.encode("utf-8")) > MAX_WORKFLOW_EMBED_BYTES:
self._logger.warning(
"Widget workflow is %d bytes (limit %d); saving recipe without it",
len(workflow_json),
MAX_WORKFLOW_EMBED_BYTES,
)
workflow_json = None
workflow_skipped = "too_large"
optimized_image, extension = self._exif_utils.optimize_image( optimized_image, extension = self._exif_utils.optimize_image(
image_data=image_bytes, image_data=image_bytes,
target_width=self._card_preview_width, target_width=self._card_preview_width,
format="webp", format="webp",
quality=85, quality=85,
preserve_metadata=True, preserve_metadata=True,
workflow=workflow_json,
) )
image_filename = f"{recipe_id}{extension}" image_filename = f"{recipe_id}{extension}"
image_path = os.path.join(recipes_dir, image_filename) image_path = os.path.join(recipes_dir, image_filename)
@@ -924,9 +964,9 @@ class RecipePersistenceService:
if key not in ["checkpoint", "loras"] if key not in ["checkpoint", "loras"]
}, },
"loras_stack": lora_stack, "loras_stack": lora_stack,
# Widget saves re-encode an in-memory tensor to PNG/WebP with no # Set by detection below: the workflow is embedded during
# embedded metadata chunks, so a workflow can never be present. # re-encoding only when the caller opted in and it fit the cap.
"has_workflow": False, "has_workflow": self._detect_has_workflow(image_path),
# Widget saves read LoRAs straight from the current workflow; an # Widget saves read LoRAs straight from the current workflow; an
# empty list means the workflow used no LoRAs. # empty list means the workflow used no LoRAs.
"import_info": build_import_info(CHANNEL_WIDGET, None, loras_data), "import_info": build_import_info(CHANNEL_WIDGET, None, loras_data),
@@ -942,15 +982,17 @@ class RecipePersistenceService:
self._exif_utils.append_recipe_metadata(image_path, recipe_data) self._exif_utils.append_recipe_metadata(image_path, recipe_data)
await recipe_scanner.add_recipe(recipe_data) await recipe_scanner.add_recipe(recipe_data)
return PersistenceResult( payload: dict[str, Any] = {
{ "success": True,
"success": True, "recipe_id": recipe_id,
"recipe_id": recipe_id, "image_path": image_path,
"image_path": image_path, "json_path": json_path,
"json_path": json_path, "recipe_name": recipe_name,
"recipe_name": recipe_name, "has_workflow": recipe_data["has_workflow"],
} }
) if workflow_skipped:
payload["workflow_skipped"] = workflow_skipped
return PersistenceResult(payload)
# Helper methods --------------------------------------------------- # Helper methods ---------------------------------------------------
+8
View File
@@ -41,6 +41,14 @@ PREVIEW_EXTENSIONS = [
# Card preview image width # Card preview image width
CARD_PREVIEW_WIDTH = 480 CARD_PREVIEW_WIDTH = 480
# Upper bound for a ComfyUI workflow embedded into a recipe preview on the
# opt-in widget save path. The workflow is by far the largest metadata field
# (tens of KB for a simple graph), so an anomalous graph — e.g. one carrying
# base64 blobs in widget values — is skipped instead of inflating the preview.
# Imports are deliberately not capped: their workflow comes from an image the
# user already chose, and preserving it is the point.
MAX_WORKFLOW_EMBED_BYTES = 256 * 1024
# Width for optimized example images # Width for optimized example images
EXAMPLE_IMAGE_WIDTH = 832 EXAMPLE_IMAGE_WIDTH = 832
+134 -23
View File
@@ -341,29 +341,125 @@ class ExifUtils:
metadata_fields = ExifUtils._load_structured_metadata(image_path) metadata_fields = ExifUtils._load_structured_metadata(image_path)
metadata_fields["parameters"] = metadata metadata_fields["parameters"] = metadata
return ExifUtils._write_structured_metadata(image_path, metadata_fields)
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
except Exception as e: except Exception as e:
logger.error(f"Error updating metadata in {image_path}: {e}") logger.error(f"Error updating metadata in {image_path}: {e}")
return image_path 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 @staticmethod
def append_recipe_metadata(image_path, recipe_data, pixel_preserving=False) -> str: def append_recipe_metadata(image_path, recipe_data, pixel_preserving=False) -> str:
"""Append recipe metadata to an image's EXIF data """Append recipe metadata to an image's EXIF data
@@ -550,7 +646,7 @@ class ExifUtils:
return None return None
@staticmethod @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 Optimize an image by resizing and converting to WebP format
@@ -560,10 +656,19 @@ class ExifUtils:
format: Output format (default: webp) format: Output format (default: webp)
quality: Output quality (0-100) quality: Output quality (0-100)
preserve_metadata: Whether to preserve EXIF metadata 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: Returns:
Tuple of (optimized_image_data, extension) 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: try:
if isinstance(image_data, str) and os.path.exists(image_data): if isinstance(image_data, str) and os.path.exists(image_data):
ext = os.path.splitext(image_data)[1].lower() 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}") logger.warning(f"Failed to extract metadata, continuing without it: {e}")
# Continue without metadata # 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 # Calculate new height to maintain aspect ratio
width, height = img.size width, height = img.size
new_height = int(height * (target_width / width)) new_height = int(height * (target_width / width))
@@ -686,8 +797,8 @@ class ExifUtils:
temp_file.write(optimized_data) temp_file.write(optimized_data)
try: try:
ExifUtils.update_image_metadata( ExifUtils._write_structured_metadata(
temp_path, metadata_fields.get("parameters") or "" temp_path, metadata_fields
) )
# Read back the file # Read back the file
with open(temp_path, 'rb') as f: with open(temp_path, 'rb') as f:
+1 -1
View File
@@ -81,7 +81,7 @@ body {
--badge-skip-refresh-glow: var(--color-skip-refresh-glow); --badge-skip-refresh-glow: var(--color-skip-refresh-glow);
} }
[data-theme="dark"] { :root[data-theme="dark"] {
--bg-color: var(--bg-base); --bg-color: var(--bg-base);
--text-color: var(--text-primary); --text-color: var(--text-primary);
--text-muted: var(--text-secondary); --text-muted: var(--text-secondary);
+1 -1
View File
@@ -572,7 +572,7 @@
} }
/* Dark mode: use each preset's dark-mode accent lightness for visibility. /* Dark mode: use each preset's dark-mode accent lightness for visibility.
These match the --color-accent-l values from [data-theme="dark"][data-theme-preset="..."] These match the --color-accent-l values from :root[data-theme="dark"][data-theme-preset="..."]
in tokens/colors.css so the swatch accurately previews what the theme looks like. */ in tokens/colors.css so the swatch accurately previews what the theme looks like. */
[data-theme="dark"] .preset-swatch-default { [data-theme="dark"] .preset-swatch-default {
+11 -11
View File
@@ -78,7 +78,7 @@
--favorite-glow: oklch(65% 0.15 85 / 0.5); --favorite-glow: oklch(65% 0.15 85 / 0.5);
} }
[data-theme="dark"] { :root[data-theme="dark"] {
--bg-base: #1a1a1a; --bg-base: #1a1a1a;
--bg-elevated: oklch(25% 0.02 256 / 0.98); --bg-elevated: oklch(25% 0.02 256 / 0.98);
--bg-overlay: oklch(0% 0 0 / 0.75); --bg-overlay: oklch(0% 0 0 / 0.75);
@@ -118,7 +118,7 @@
/* ── Preset: Nord ──────────────────────────────────────────── */ /* ── Preset: Nord ──────────────────────────────────────────── */
[data-theme-preset="nord"] { :root[data-theme-preset="nord"] {
--color-accent-h: 213; --color-accent-h: 213;
--color-accent-c: 0.18; --color-accent-c: 0.18;
--color-accent-l: 62%; --color-accent-l: 62%;
@@ -152,7 +152,7 @@
--favorite-glow: oklch(72% 0.14 85 / 0.5); --favorite-glow: oklch(72% 0.14 85 / 0.5);
} }
[data-theme="dark"][data-theme-preset="nord"] { :root[data-theme="dark"][data-theme-preset="nord"] {
--color-accent-h: 213; --color-accent-h: 213;
--color-accent-c: 0.18; --color-accent-c: 0.18;
--color-accent-l: 68%; --color-accent-l: 68%;
@@ -188,7 +188,7 @@
/* ── Preset: Midnight ───────────────────────────────────────── */ /* ── Preset: Midnight ───────────────────────────────────────── */
[data-theme-preset="midnight"] { :root[data-theme-preset="midnight"] {
--color-accent-h: 300; --color-accent-h: 300;
--color-accent-c: 0.15; --color-accent-c: 0.15;
--color-accent-l: 52%; --color-accent-l: 52%;
@@ -222,7 +222,7 @@
--favorite-glow: oklch(72% 0.16 85 / 0.5); --favorite-glow: oklch(72% 0.16 85 / 0.5);
} }
[data-theme="dark"][data-theme-preset="midnight"] { :root[data-theme="dark"][data-theme-preset="midnight"] {
--color-accent-h: 300; --color-accent-h: 300;
--color-accent-c: 0.14; --color-accent-c: 0.14;
--color-accent-l: 68%; --color-accent-l: 68%;
@@ -258,7 +258,7 @@
/* ── Preset: Monokai ───────────────────────────────────────── */ /* ── Preset: Monokai ───────────────────────────────────────── */
[data-theme-preset="monokai"] { :root[data-theme-preset="monokai"] {
--color-accent-h: 190; --color-accent-h: 190;
--color-accent-c: 0.24; --color-accent-c: 0.24;
--color-accent-l: 72%; --color-accent-l: 72%;
@@ -291,7 +291,7 @@
--favorite-glow: oklch(72% 0.16 85 / 0.5); --favorite-glow: oklch(72% 0.16 85 / 0.5);
} }
[data-theme="dark"][data-theme-preset="monokai"] { :root[data-theme="dark"][data-theme-preset="monokai"] {
--color-accent-h: 190; --color-accent-h: 190;
--color-accent-c: 0.24; --color-accent-c: 0.24;
--color-accent-l: 72%; --color-accent-l: 72%;
@@ -326,7 +326,7 @@
/* ── Preset: Dracula ───────────────────────────────────────── */ /* ── Preset: Dracula ───────────────────────────────────────── */
[data-theme-preset="dracula"] { :root[data-theme-preset="dracula"] {
--color-accent-h: 265; --color-accent-h: 265;
--color-accent-c: 0.24; --color-accent-c: 0.24;
--color-accent-l: 68%; --color-accent-l: 68%;
@@ -359,7 +359,7 @@
--favorite-glow: oklch(72% 0.16 85 / 0.5); --favorite-glow: oklch(72% 0.16 85 / 0.5);
} }
[data-theme="dark"][data-theme-preset="dracula"] { :root[data-theme="dark"][data-theme-preset="dracula"] {
--color-accent-h: 265; --color-accent-h: 265;
--color-accent-c: 0.24; --color-accent-c: 0.24;
--color-accent-l: 72%; --color-accent-l: 72%;
@@ -394,7 +394,7 @@
/* ── Preset: Solarized ─────────────────────────────────────── */ /* ── Preset: Solarized ─────────────────────────────────────── */
[data-theme-preset="solarized"] { :root[data-theme-preset="solarized"] {
--color-accent-h: 175; --color-accent-h: 175;
--color-accent-c: 0.18; --color-accent-c: 0.18;
--color-accent-l: 55%; --color-accent-l: 55%;
@@ -429,7 +429,7 @@
--favorite-glow: oklch(68% 0.16 75 / 0.5); --favorite-glow: oklch(68% 0.16 75 / 0.5);
} }
[data-theme="dark"][data-theme-preset="solarized"] { :root[data-theme="dark"][data-theme-preset="solarized"] {
--color-accent-h: 175; --color-accent-h: 175;
--color-accent-c: 0.18; --color-accent-c: 0.18;
--color-accent-l: 60%; --color-accent-l: 60%;
@@ -72,6 +72,15 @@ export class DownloadManager {
completeMetadata.diagnostics = diagnostics; 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 // Preserve preview_nsfw_level from analysis so the saved
// recipe applies the correct NSFW blur on the preview image. // recipe applies the correct NSFW blur on the preview image.
const nsfwLevel = this.importManager.recipeData.preview_nsfw_level; const nsfwLevel = this.importManager.recipeData.preview_nsfw_level;
@@ -0,0 +1,79 @@
import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
const { EVENTS_MODULE, API_MODULE, APP_MODULE, COMPONENTS_MODULE, UTILS_MODULE } =
vi.hoisted(() => ({
EVENTS_MODULE: new URL('../../../web/comfyui/loras_widget_events.js', import.meta.url)
.pathname,
API_MODULE: new URL('../../../scripts/api.js', import.meta.url).pathname,
APP_MODULE: new URL('../../../scripts/app.js', import.meta.url).pathname,
COMPONENTS_MODULE: new URL('../../../web/comfyui/loras_widget_components.js', import.meta.url)
.pathname,
UTILS_MODULE: new URL('../../../web/comfyui/loras_widget_utils.js', import.meta.url)
.pathname,
}));
const saveRecipeDirectly = vi.fn();
vi.mock(API_MODULE, () => ({ api: {} }));
vi.mock(APP_MODULE, () => ({ app: {} }));
vi.mock(COMPONENTS_MODULE, () => ({
createMenuItem: (text, icon, onClick) => {
const el = document.createElement('div');
el.className = 'lm-lora-menu-item';
el.textContent = text;
if (onClick) el.addEventListener('click', onClick);
return el;
},
createDropIndicator: vi.fn(),
}));
vi.mock(UTILS_MODULE, () => ({
parseLoraValue: vi.fn(() => []),
formatLoraValue: vi.fn((value) => value),
syncClipStrengthIfCollapsed: vi.fn(),
saveRecipeDirectly,
copyToClipboard: vi.fn(),
showToast: vi.fn(),
moveLoraByDirection: vi.fn(),
getDropTargetIndex: vi.fn(),
getLoraStrengthRange: vi.fn(),
applyStrengthRangeCue: vi.fn(),
}));
function findMenuItem(label) {
return Array.from(document.querySelectorAll('.lm-lora-menu-item')).find(
(item) => item.textContent === label
);
}
describe('LoRA widget context menu save options', () => {
beforeEach(() => {
document.body.innerHTML = '';
});
afterEach(() => {
document.body.innerHTML = '';
vi.clearAllMocks();
});
it('offers workflow embedding as a separate, opt-in action', async () => {
const { createContextMenu } = await import(EVENTS_MODULE);
const widget = { value: [], callback: vi.fn() };
createContextMenu(10, 10, 'lora-a', widget, null, vi.fn());
const plain = findMenuItem('Save Recipe');
const withWorkflow = findMenuItem('Save Recipe with Workflow');
expect(plain).toBeTruthy();
expect(withWorkflow).toBeTruthy();
plain.click();
expect(saveRecipeDirectly).toHaveBeenLastCalledWith();
// Re-open: the first click removed the menu.
createContextMenu(10, 10, 'lora-a', widget, null, vi.fn());
findMenuItem('Save Recipe with Workflow').click();
expect(saveRecipeDirectly).toHaveBeenLastCalledWith({ embedWorkflow: true });
});
});
@@ -0,0 +1,101 @@
import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
const { showToastMock, translateMock } = vi.hoisted(() => ({
showToastMock: vi.fn(),
translateMock: vi.fn((key, params, fallback) =>
typeof fallback === 'string' ? fallback : key
),
}));
vi.mock('../../../static/js/utils/uiHelpers.js', () => ({
showToast: showToastMock,
}));
vi.mock('../../../static/js/utils/i18nHelpers.js', () => ({
translate: translateMock,
}));
vi.mock('../../../static/js/api/modelApiFactory.js', () => ({
getModelApiClient: vi.fn(() => ({})),
}));
vi.mock('../../../static/js/api/apiConfig.js', () => ({
MODEL_TYPES: { LORA: 'loras', CHECKPOINT: 'checkpoints', EMBEDDING: 'embeddings' },
}));
vi.mock('../../../static/js/utils/storageHelpers.js', () => ({
getStorageItem: vi.fn(() => null),
}));
vi.mock('../../../static/js/state/index.js', () => ({
state: { virtualScroller: null },
}));
import { DownloadManager } from '../../../static/js/managers/import/DownloadManager.js';
function buildImportManager(recipeData) {
return {
recipeId: null,
recipeName: 'Test Recipe',
recipeImage: null,
recipeTags: [],
downloadableLoRAs: [],
importMode: 'url',
recipeData,
loadingManager: { showSimpleLoading: vi.fn(), hide: vi.fn() },
};
}
async function saveAndReadMetadata(recipeData) {
let capturedBody = null;
const fetchMock = vi.fn(async (url, options) => {
capturedBody = options?.body ?? null;
return { ok: true, json: async () => ({ success: true }) };
});
globalThis.fetch = fetchMock;
window.fetch = fetchMock;
const manager = new DownloadManager(buildImportManager(recipeData));
await manager.saveRecipe(true);
expect(fetchMock).toHaveBeenCalledWith('/api/lm/recipes/save', expect.anything());
expect(capturedBody).toBeInstanceOf(FormData);
return JSON.parse(capturedBody.get('metadata'));
}
describe('recipe import save payload', () => {
beforeEach(() => {
globalThis.modalManager = { closeModal: vi.fn() };
globalThis.window.recipeManager = { loadRecipes: vi.fn() };
});
afterEach(() => {
delete globalThis.modalManager;
delete globalThis.window.recipeManager;
vi.clearAllMocks();
});
it('forwards a workflow recovered from the original rendition', async () => {
const workflow = '{"nodes": [{"id": 1}]}';
const metadata = await saveAndReadMetadata({
image_base64: 'AAAA',
base_model: 'sd',
loras: [],
gen_params: {},
workflow,
});
expect(metadata.workflow).toBe(workflow);
});
it('omits the workflow key when analysis recovered none', async () => {
const metadata = await saveAndReadMetadata({
image_base64: 'AAAA',
base_model: 'sd',
loras: [],
gen_params: {},
});
expect('workflow' in metadata).toBe(false);
});
});
@@ -0,0 +1,68 @@
import { describe, it, expect } from 'vitest';
import { readFileSync } from 'fs';
import path from 'path';
// Regression guard: theme palette tokens may only be declared on the root element.
//
// `applyTheme()` (static/js/utils/uiHelpers.js, and Header.setThemeMode) mirrors
// the active mode onto <body> as `data-theme="dark"`, but the theme *preset* is
// only ever written to <html> (`data-theme-preset`). While the token blocks in
// tokens/colors.css used the bare attribute selector `[data-theme="dark"]`, the
// body matched that block on its own and re-declared the DEFAULT dark palette
// (#1a1a1a / #2d2d2d), shadowing the preset palette it inherited from <html>.
// The page painted the selected preset's accent but the default preset's
// backgrounds/surfaces/text, and flipped into that state ~200ms after load when
// initTheme() first touched <body> — the accent-tinted background flash on
// reload and nav-tab switches, visible under every preset except "default".
// Root-scoping the token blocks keeps <body>'s data-theme inert.
describe('Theme token scope', () => {
const repoRoot = path.resolve(__dirname, '../../..');
const read = (rel) => readFileSync(path.join(repoRoot, rel), 'utf-8');
const stripComments = (css) => css.replace(/\/\*[\s\S]*?\*\//g, '');
const COLOR_TOKENS = read('static/css/tokens/colors.css');
const BASE_CSS = read('static/css/base.css');
// Selectors that declare a palette token, e.g. `--bg-base:` / `--lora-surface:`.
const tokenDeclaringSelectors = (css) => {
const selectors = [];
const ruleRe = /([^{}]+)\{([^{}]*)\}/g;
let match;
while ((match = ruleRe.exec(stripComments(css)))) {
const selector = match[1].trim();
const declaresPaletteToken = /(^|[;\s])--(?:bg|surface|text|border|color|favorite|lora|badge|card)-[\w-]+\s*:/.test(
match[2]
);
if (declaresPaletteToken) selectors.push(selector);
}
return selectors;
};
const isRootScoped = (selector) =>
selector.split(',').every((part) => /^\s*(:root|html)\b/.test(part));
it('declares every colors.css palette token on :root', () => {
const selectors = tokenDeclaringSelectors(COLOR_TOKENS);
expect(selectors.length).toBeGreaterThan(0);
expect(selectors.filter((selector) => !isRootScoped(selector))).toEqual([]);
});
it('declares every base.css palette alias on :root', () => {
const selectors = tokenDeclaringSelectors(BASE_CSS);
expect(selectors.length).toBeGreaterThan(0);
expect(selectors.filter((selector) => !isRootScoped(selector))).toEqual([]);
});
it('keeps dark/preset token blocks anchored to the root element', () => {
const css = stripComments(COLOR_TOKENS);
expect(css).toContain(':root[data-theme="dark"] {');
for (const preset of ['nord', 'midnight', 'monokai', 'dracula', 'solarized']) {
expect(css).toContain(`:root[data-theme-preset="${preset}"] {`);
expect(css).toContain(`:root[data-theme="dark"][data-theme-preset="${preset}"] {`);
}
// No bare attribute selector may open a rule: <body> carries data-theme
// without the preset, so such a block would re-declare the default palette.
expect(css).not.toMatch(/(^|\n)\s*\[data-theme/);
expect(stripComments(BASE_CSS)).not.toMatch(/(^|\n)\s*\[data-theme/);
});
});
@@ -0,0 +1,92 @@
import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
const { UTILS_MODULE, APP_MODULE, API_MODULE, BASE_PATH_MODULE } = vi.hoisted(() => ({
UTILS_MODULE: new URL('../../../web/comfyui/loras_widget_utils.js', import.meta.url).pathname,
APP_MODULE: new URL('../../../scripts/app.js', import.meta.url).pathname,
API_MODULE: new URL('../../../scripts/api.js', import.meta.url).pathname,
BASE_PATH_MODULE: new URL('../../../web/comfyui/base_path.js', import.meta.url).pathname,
}));
const toastAdd = vi.fn();
const graphToPrompt = vi.fn();
vi.mock(APP_MODULE, () => ({
app: {
graphToPrompt,
extensionManager: { toast: { add: toastAdd } },
},
}));
vi.mock(API_MODULE, () => ({
api: { fetchApi: vi.fn() },
}));
vi.mock(BASE_PATH_MODULE, () => ({
lmUrl: (path) => `/lm${path}`,
}));
async function runSave(options, responseBody) {
let captured = null;
globalThis.fetch = vi.fn(async (url, init) => {
captured = { url, init };
return { json: async () => responseBody };
});
const { saveRecipeDirectly } = await import(UTILS_MODULE);
await saveRecipeDirectly(options);
return captured;
}
describe('saveRecipeDirectly', () => {
beforeEach(() => {
graphToPrompt.mockResolvedValue({
workflow: { nodes: [{ id: 1 }], last_node_id: 1 },
output: { 1: { class_type: 'KSampler' } },
});
});
afterEach(() => {
delete globalThis.fetch;
vi.clearAllMocks();
});
it('posts no workflow by default', async () => {
const captured = await runSave(undefined, { success: true, has_workflow: false });
expect(captured.init.body).toBe('{}');
expect(JSON.parse(captured.init.body)).not.toHaveProperty('workflow');
});
it('embeds the UI-format graph when asked', async () => {
const captured = await runSave(
{ embedWorkflow: true },
{ success: true, has_workflow: true }
);
const body = JSON.parse(captured.init.body);
expect(body.workflow).toEqual({ nodes: [{ id: 1 }], last_node_id: 1 });
expect(captured.init.headers['Content-Type']).toBe('application/json');
});
it('reports a skipped oversized workflow as a warning', async () => {
await runSave(
{ embedWorkflow: true },
{ success: true, has_workflow: false, workflow_skipped: 'too_large' }
);
const lastToast = toastAdd.mock.calls.at(-1)[0];
expect(lastToast.severity).toBe('warn');
expect(lastToast.summary).toBe('Recipe Saved without Workflow');
expect(lastToast.detail).toContain('too large');
});
it('reports a successful embed distinctly from a plain save', async () => {
await runSave(
{ embedWorkflow: true },
{ success: true, has_workflow: true }
);
const lastToast = toastAdd.mock.calls.at(-1)[0];
expect(lastToast.summary).toBe('Recipe Saved with Workflow');
});
});
+262
View File
@@ -0,0 +1,262 @@
"""Workflow preservation for remote recipe imports.
CivitAI serves a re-encoded, metadata-free ``optimized`` rendition as the
recipe preview, so an embedded ComfyUI workflow only exists in the
``original=true`` image. These tests pin the recovery and transport of that
workflow through the remote import paths.
"""
from __future__ import annotations
import json
import logging
import os
from pathlib import Path
from types import SimpleNamespace
from typing import Any
import pytest
from PIL import Image, PngImagePlugin
from py.routes.handlers.recipe_handlers import RecipeManagementHandler
from py.services.recipes.persistence_service import PersistenceResult
from py.utils.exif_utils import ExifUtils
async def _noop_ensure() -> None:
return None
class CapturingPersistence:
"""Persistence service double recording the save payload."""
def __init__(self) -> None:
self.calls: list[dict[str, Any]] = []
async def save_recipe(self, **kwargs: Any) -> PersistenceResult:
self.calls.append(kwargs)
return PersistenceResult({"success": True, "recipe_id": "recipe-1"})
class StubScanner:
"""Scanner double exposing only what the remote import paths touch."""
def __init__(self) -> None:
self.recipes_dir = "/tmp/recipes"
async def build_local_hash_cache(self) -> dict[str, Any]:
return {}
async def get_local_lora(self, name, base_model=None):
return None
def _make_handler(
persistence: CapturingPersistence,
*,
downloader_factory=None,
) -> RecipeManagementHandler:
async def default_downloader_factory():
class Downloader:
async def download_file(self, url, path, use_auth=False):
Path(path).write_bytes(b"downloaded")
return True, "ok"
return Downloader()
analysis_service = SimpleNamespace(
_recipe_parser_factory=SimpleNamespace(create_parser=lambda metadata: None)
)
return RecipeManagementHandler(
ensure_dependencies_ready=_noop_ensure,
recipe_scanner_getter=lambda: StubScanner(),
logger=logging.getLogger(__name__),
persistence_service=persistence, # pyright: ignore[reportArgumentType]
analysis_service=analysis_service, # pyright: ignore[reportArgumentType]
downloader_factory=downloader_factory or default_downloader_factory,
civitai_client_getter=lambda: None,
)
def _meta_with_comfy() -> dict[str, Any]:
return {
"id": 143518055,
"meta": {"prompt": "p", "comfy": '{"prompt": {"1": {"class_type": "KSampler"}}}'},
}
def test_meta_indicates_comfy_workflow() -> None:
assert RecipeManagementHandler._meta_indicates_comfy_workflow(
{"meta": {"comfy": "{}"}}
)
assert RecipeManagementHandler._meta_indicates_comfy_workflow({"comfy": "{}"})
assert not RecipeManagementHandler._meta_indicates_comfy_workflow({"meta": {}})
assert not RecipeManagementHandler._meta_indicates_comfy_workflow(
{"meta": {"comfy": None}}
)
assert not RecipeManagementHandler._meta_indicates_comfy_workflow(None)
assert not RecipeManagementHandler._meta_indicates_comfy_workflow("comfy")
@pytest.mark.asyncio
async def test_fetch_original_media_reads_workflow_and_cleans_up(tmp_path, monkeypatch):
workflow = json.dumps({"nodes": [{"id": 1}], "last_node_id": 1})
source = tmp_path / "original.png"
png_info = PngImagePlugin.PngInfo()
png_info.add_text("workflow", workflow)
png_info.add_text("prompt", '{"1": {"class_type": "KSampler"}}')
Image.new("RGB", (32, 32), color="red").save(source, pnginfo=png_info)
written: list[str] = []
async def downloader_factory():
class Downloader:
async def download_file(self, url, path, use_auth=False):
written.append(str(path))
Path(path).write_bytes(source.read_bytes())
return True, "ok"
return Downloader()
handler = _make_handler(CapturingPersistence(), downloader_factory=downloader_factory)
raw_metadata, recovered = await handler._fetch_original_media(
"https://image.civitai.com/x/original=true/x.png"
)
assert recovered == workflow
# extract_image_metadata prefers the prompt chunk over the workflow.
assert raw_metadata is not None and "class_type" in raw_metadata
assert written and not os.path.exists(written[0])
@pytest.mark.asyncio
async def test_fetch_original_media_degrades_on_download_failure():
async def downloader_factory():
class Downloader:
async def download_file(self, url, path, use_auth=False):
return False, "boom"
return Downloader()
handler = _make_handler(CapturingPersistence(), downloader_factory=downloader_factory)
assert await handler._fetch_original_media("https://image.civitai.com/x.png") == (
None,
None,
)
assert await handler._fetch_original_media(None) == (None, None)
@pytest.mark.asyncio
async def test_remote_import_transports_workflow_to_save(monkeypatch):
workflow = json.dumps({"nodes": [{"id": 4}]})
persistence = CapturingPersistence()
handler = _make_handler(persistence)
async def fake_download_remote_media(image_url):
return (
b"optimized-preview",
".jpg",
_meta_with_comfy(),
12345,
"https://image.civitai.com/x/original=true/x.png",
)
fetched: list[str] = []
async def fake_fetch_original_media(original_url):
fetched.append(original_url)
return None, workflow
handler._download_remote_media = fake_download_remote_media # type: ignore[method-assign]
handler._fetch_original_media = fake_fetch_original_media # type: ignore[method-assign]
monkeypatch.setattr(
ExifUtils, "extract_image_metadata", staticmethod(lambda path: None)
)
response = await handler._do_import_remote_recipe(
image_url="https://civitai.red/images/143518055",
name="Recipe",
lora_entries=[],
checkpoint_entry=None,
gen_params_request={},
tags=[],
base_model="Krea 2",
source_path="https://civitai.red/images/143518055",
)
assert response.status == 200
assert fetched == ["https://image.civitai.com/x/original=true/x.png"]
assert persistence.calls[0]["metadata"]["workflow"] == workflow
@pytest.mark.asyncio
async def test_remote_import_skips_original_without_comfy_meta(monkeypatch):
persistence = CapturingPersistence()
handler = _make_handler(persistence)
async def fake_download_remote_media(image_url):
return (
b"optimized-preview",
".jpg",
{"id": 1, "meta": {"prompt": "p"}},
None,
"https://image.civitai.com/x/original=true/x.png",
)
async def fail_fetch(original_url): # pragma: no cover - must not be called
raise AssertionError("original rendition should not be fetched")
handler._download_remote_media = fake_download_remote_media # type: ignore[method-assign]
handler._fetch_original_media = fail_fetch # type: ignore[method-assign]
monkeypatch.setattr(
ExifUtils, "extract_image_metadata", staticmethod(lambda path: None)
)
response = await handler._do_import_remote_recipe(
image_url="https://civitai.red/images/1",
name="Recipe",
lora_entries=[],
checkpoint_entry=None,
gen_params_request={},
tags=[],
base_model="SDXL 1.0",
source_path="https://civitai.red/images/1",
)
assert response.status == 200
assert "workflow" not in persistence.calls[0]["metadata"]
@pytest.mark.asyncio
async def test_url_import_transports_workflow_to_save(monkeypatch):
workflow = json.dumps({"nodes": [{"id": 5}]})
persistence = CapturingPersistence()
handler = _make_handler(persistence)
async def fake_download_remote_media(image_url):
return (
b"optimized-preview",
".jpg",
{"id": 9, "meta": {"prompt": "p"}},
None,
"https://image.civitai.com/x/original=true/x.png",
)
async def fake_fetch_original_media(original_url):
return None, workflow
handler._download_remote_media = fake_download_remote_media # type: ignore[method-assign]
handler._fetch_original_media = fake_fetch_original_media # type: ignore[method-assign]
monkeypatch.setattr(
ExifUtils, "extract_image_metadata", staticmethod(lambda path: None)
)
response = await handler._do_import_from_url(
"https://civitai.red/images/143518055", StubScanner()
)
assert response.status == 200
assert persistence.calls[0]["metadata"]["workflow"] == workflow
+52 -3
View File
@@ -208,6 +208,11 @@ class StubAnalysisService:
self.remote_calls: List[Optional[str]] = [] self.remote_calls: List[Optional[str]] = []
self.local_calls: List[Optional[str]] = [] self.local_calls: List[Optional[str]] = []
self.local_ignore_recipe_metadata_calls: List[bool] = [] self.local_ignore_recipe_metadata_calls: List[bool] = []
self.widget_analysis_calls: List[Any] = []
self.widget_result = SimpleNamespace(
payload={"metadata": {"loras": ""}, "image_bytes": b"widget-image"},
status=200,
)
self.result = SimpleNamespace(payload={"loras": []}, status=200) self.result = SimpleNamespace(payload={"loras": []}, status=200)
self._recipe_parser_factory: Any = None self._recipe_parser_factory: Any = None
StubAnalysisService.instances.append(self) StubAnalysisService.instances.append(self)
@@ -242,7 +247,8 @@ class StubAnalysisService:
return self.result return self.result
async def analyze_widget_metadata(self, *, recipe_scanner) -> SimpleNamespace: async def analyze_widget_metadata(self, *, recipe_scanner) -> SimpleNamespace:
return SimpleNamespace(payload={"metadata": {}, "image_bytes": b""}, status=200) self.widget_analysis_calls.append(recipe_scanner)
return self.widget_result
class StubPersistenceService: class StubPersistenceService:
@@ -252,6 +258,7 @@ class StubPersistenceService:
def __init__(self, **_: Any) -> None: def __init__(self, **_: Any) -> None:
self.save_calls: List[Dict[str, Any]] = [] self.save_calls: List[Dict[str, Any]] = []
self.widget_calls: List[Dict[str, Any]] = []
self.delete_calls: List[str] = [] self.delete_calls: List[str] = []
self.move_calls: List[Dict[str, str]] = [] self.move_calls: List[Dict[str, str]] = []
self.update_calls: List[Dict[str, Any]] = [] self.update_calls: List[Dict[str, Any]] = []
@@ -359,9 +366,24 @@ class StubPersistenceService:
) )
async def save_recipe_from_widget( async def save_recipe_from_widget(
self, *, recipe_scanner, metadata: Dict[str, Any], image_bytes: bytes self,
*,
recipe_scanner,
metadata: Dict[str, Any],
image_bytes: bytes,
workflow: Any = None,
) -> SimpleNamespace: # pragma: no cover ) -> SimpleNamespace: # pragma: no cover
return SimpleNamespace(payload={"success": True}, status=200) self.widget_calls.append(
{
"recipe_scanner": recipe_scanner,
"metadata": metadata,
"image_bytes": image_bytes,
"workflow": workflow,
}
)
return SimpleNamespace(
payload={"success": True, "has_workflow": workflow is not None}, status=200
)
class StubSharingService: class StubSharingService:
@@ -481,6 +503,33 @@ async def recipe_harness(
StubSharingService.instances.clear() StubSharingService.instances.clear()
async def test_save_from_widget_forwards_workflow_body(monkeypatch, tmp_path: Path) -> None:
"""The opt-in workflow arrives through a real JSON body and is handed to
the persistence layer; the response reports whether it was embedded."""
async with recipe_harness(monkeypatch, tmp_path) as harness:
workflow = {"nodes": [{"id": 1}], "last_node_id": 1}
response = await harness.client.post(
"/api/lm/recipes/save-from-widget", json={"workflow": workflow}
)
payload = await response.json()
assert response.status == 200
assert payload["has_workflow"] is True
assert harness.persistence.widget_calls[0]["workflow"] == workflow
async def test_save_from_widget_without_body_still_saves(monkeypatch, tmp_path: Path) -> None:
"""The long-standing body-less POST must keep working unchanged."""
async with recipe_harness(monkeypatch, tmp_path) as harness:
response = await harness.client.post("/api/lm/recipes/save-from-widget")
payload = await response.json()
assert response.status == 200
assert payload["has_workflow"] is False
assert harness.persistence.widget_calls[0]["workflow"] is None
async def test_list_recipes_provides_file_urls(monkeypatch, tmp_path: Path) -> None: async def test_list_recipes_provides_file_urls(monkeypatch, tmp_path: Path) -> None:
async with recipe_harness(monkeypatch, tmp_path) as harness: async with recipe_harness(monkeypatch, tmp_path) as harness:
recipe_path = harness.tmp_dir / "recipes" / "demo.png" recipe_path = harness.tmp_dir / "recipes" / "demo.png"
+129
View File
@@ -0,0 +1,129 @@
"""Handler tests for the widget "Save Recipe" endpoint.
Covers the opt-in workflow body: the endpoint historically received no body at
all, so a missing, empty or malformed body must degrade to "no workflow"
rather than failing the save.
"""
from __future__ import annotations
import json
import logging
from types import SimpleNamespace
from typing import Any
import pytest
from py.routes.handlers.recipe_handlers import RecipeManagementHandler
async def _noop_ensure() -> None:
return None
class FakeRequest:
"""Minimal request double exposing the optional-body contract."""
def __init__(
self,
*,
body: Any = None,
can_read_body: bool = True,
json_raises: bool = False,
) -> None:
self._body = body
self.can_read_body = can_read_body
self._json_raises = json_raises
async def json(self) -> Any:
if self._json_raises or self._body is None:
raise ValueError("no JSON body")
return self._body
class CapturingPersistence:
def __init__(self) -> None:
self.calls: list[dict[str, Any]] = []
async def save_recipe_from_widget(self, **kwargs: Any) -> SimpleNamespace:
self.calls.append(kwargs)
return SimpleNamespace(
payload={"success": True, "has_workflow": bool(kwargs.get("workflow"))},
status=200,
)
def _make_handler(persistence: CapturingPersistence) -> RecipeManagementHandler:
analysis_service = SimpleNamespace(
analyze_widget_metadata=lambda **kwargs: _analysis_result()
)
return RecipeManagementHandler(
ensure_dependencies_ready=_noop_ensure,
recipe_scanner_getter=lambda: object(),
logger=logging.getLogger(__name__),
persistence_service=persistence, # pyright: ignore[reportArgumentType]
analysis_service=analysis_service, # pyright: ignore[reportArgumentType]
downloader_factory=lambda: None,
civitai_client_getter=lambda: None,
)
async def _analysis_result() -> SimpleNamespace:
return SimpleNamespace(
payload={"metadata": {"loras": ""}, "image_bytes": b"image"}
)
@pytest.mark.asyncio
async def test_widget_save_forwards_workflow_from_json_body():
persistence = CapturingPersistence()
handler = _make_handler(persistence)
workflow = {"nodes": [{"id": 1}]}
response = await handler.save_recipe_from_widget(
FakeRequest(body={"workflow": workflow}) # type: ignore[arg-type]
)
assert response.status == 200
assert persistence.calls[0]["workflow"] == workflow
@pytest.mark.asyncio
async def test_widget_save_without_body_passes_no_workflow():
persistence = CapturingPersistence()
handler = _make_handler(persistence)
response = await handler.save_recipe_from_widget(
FakeRequest(can_read_body=False) # type: ignore[arg-type]
)
assert response.status == 200
assert persistence.calls[0]["workflow"] is None
@pytest.mark.asyncio
async def test_widget_save_tolerates_malformed_body():
persistence = CapturingPersistence()
handler = _make_handler(persistence)
await handler.save_recipe_from_widget(
FakeRequest(json_raises=True) # type: ignore[arg-type]
)
await handler.save_recipe_from_widget(
FakeRequest(body=["not", "an", "object"]) # type: ignore[arg-type]
)
assert [call["workflow"] for call in persistence.calls] == [None, None]
@pytest.mark.asyncio
async def test_widget_save_reports_embedded_workflow_in_response():
persistence = CapturingPersistence()
handler = _make_handler(persistence)
response = await handler.save_recipe_from_widget(
FakeRequest(body={"workflow": {"nodes": []}}) # type: ignore[arg-type]
)
assert json.loads(response.text)["has_workflow"] is True
@@ -589,6 +589,49 @@ class TestBatchImportServiceEdgeCases:
assert "batch-import" in persistence_service.saved_recipes[0]["tags"] assert "batch-import" in persistence_service.saved_recipes[0]["tags"]
assert "test" 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 @pytest.mark.asyncio
async def test_skip_duplicates_parameter(self, service): async def test_skip_duplicates_parameter(self, service):
recipe_scanner_getter = lambda: SimpleNamespace() recipe_scanner_getter = lambda: SimpleNamespace()
+362 -1
View File
@@ -27,14 +27,36 @@ class DummyExifUtils:
self.appended = None self.appended = None
self.optimized_calls = 0 self.optimized_calls = 0
self.workflow_value = None 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_calls += 1
self.optimized_workflow = workflow
return image_data, ".webp" return image_data, ".webp"
def append_recipe_metadata(self, image_path, recipe_data, pixel_preserving=False): def append_recipe_metadata(self, image_path, recipe_data, pixel_preserving=False):
self.appended = (image_path, recipe_data, pixel_preserving) 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 normalise_workflow(self, workflow):
if isinstance(workflow, str):
return workflow or None
if isinstance(workflow, (dict, list)):
return json.dumps(workflow)
return None
def extract_image_metadata(self, path): def extract_image_metadata(self, path):
return {} return {}
@@ -131,6 +153,55 @@ async def test_save_recipe_skip_optimize_preserves_image_bytes(tmp_path):
assert exif_utils.appended[2] is True 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 @pytest.mark.asyncio
async def test_save_recipe_skip_optimize_default_optimizes(tmp_path): async def test_save_recipe_skip_optimize_default_optimizes(tmp_path):
"""Normal saves must keep optimizing; only re-import opts out.""" """Normal saves must keep optimizing; only re-import opts out."""
@@ -650,6 +721,88 @@ async def test_save_recipe_preserves_workflow_when_png_is_converted_to_webp(tmp_
assert "Recipe metadata:" in decoded_comment 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 @pytest.mark.asyncio
async def test_save_recipe_strips_checkpoint_local_fields(tmp_path): async def test_save_recipe_strips_checkpoint_local_fields(tmp_path):
exif_utils = DummyExifUtils() exif_utils = DummyExifUtils()
@@ -835,6 +988,137 @@ async def test_save_recipe_from_widget_enriches_checkpoint_from_local_cache(tmp_
} }
@pytest.mark.asyncio
async def test_save_recipe_from_widget_embeds_opted_in_workflow(tmp_path):
"""Opt-in widget saves embed the live graph so the recipe can send its
workflow back to ComfyUI, mirroring imported recipes."""
class DummyScanner:
def __init__(self, root):
self.recipes_dir = str(root)
self.added = []
async def get_local_lora(self, name): # pragma: no cover - no loras
return None
async def add_recipe(self, recipe_data):
self.added.append(recipe_data)
image_buffer = BytesIO()
Image.new("RGB", (96, 48), color="navy").save(
image_buffer, format="PNG"
)
scanner = DummyScanner(tmp_path)
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_from_widget(
recipe_scanner=scanner,
metadata={"loras": "", "prompt": "a calm scene"},
image_bytes=image_buffer.getvalue(),
workflow=workflow,
)
assert result.payload["has_workflow"] is True
assert "workflow_skipped" not in result.payload
stored = json.loads(Path(result.payload["json_path"]).read_text())
assert stored["has_workflow"] is True
assert ExifUtils._load_structured_metadata(result.payload["image_path"])[
"workflow"
] == json.dumps(workflow)
assert scanner.added[0]["has_workflow"] is True
@pytest.mark.asyncio
async def test_save_recipe_from_widget_without_workflow_stays_unflagged(tmp_path):
"""The default action must keep saving a workflow-free preview."""
class DummyScanner:
def __init__(self, root):
self.recipes_dir = str(root)
async def get_local_lora(self, name): # pragma: no cover - no loras
return None
async def add_recipe(self, recipe_data):
return None
image_buffer = BytesIO()
Image.new("RGB", (96, 48), color="navy").save(
image_buffer, format="PNG"
)
service = RecipePersistenceService(
exif_utils=ExifUtils,
card_preview_width=64,
logger=logging.getLogger("test"),
)
result = await service.save_recipe_from_widget(
recipe_scanner=DummyScanner(tmp_path),
metadata={"loras": "", "prompt": "a calm scene"},
image_bytes=image_buffer.getvalue(),
)
assert result.payload["has_workflow"] is False
assert ExifUtils._load_structured_metadata(result.payload["image_path"])[
"workflow"
] is None
@pytest.mark.asyncio
async def test_save_recipe_from_widget_skips_oversized_workflow(
tmp_path, monkeypatch
):
"""A pathological graph is dropped instead of inflating the preview."""
monkeypatch.setattr(
"py.services.recipes.persistence_service.MAX_WORKFLOW_EMBED_BYTES", 32
)
class DummyScanner:
def __init__(self, root):
self.recipes_dir = str(root)
self.added = []
async def get_local_lora(self, name): # pragma: no cover - no loras
return None
async def add_recipe(self, recipe_data):
self.added.append(recipe_data)
image_buffer = BytesIO()
Image.new("RGB", (96, 48), color="navy").save(
image_buffer, format="PNG"
)
scanner = DummyScanner(tmp_path)
service = RecipePersistenceService(
exif_utils=ExifUtils,
card_preview_width=64,
logger=logging.getLogger("test"),
)
workflow = {"nodes": [{"id": index} for index in range(20)]}
result = await service.save_recipe_from_widget(
recipe_scanner=scanner,
metadata={"loras": "", "prompt": "a calm scene"},
image_bytes=image_buffer.getvalue(),
workflow=workflow,
)
assert result.payload["workflow_skipped"] == "too_large"
assert result.payload["has_workflow"] is False
assert ExifUtils._load_structured_metadata(result.payload["image_path"])[
"workflow"
] is None
assert scanner.added[0]["has_workflow"] is False
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_move_recipe_updates_paths(tmp_path): async def test_move_recipe_updates_paths(tmp_path):
exif_utils = DummyExifUtils() exif_utils = DummyExifUtils()
@@ -1071,6 +1355,83 @@ async def test_analyze_remote_image_supports_civitai_red():
assert result.payload["loras"] == [] 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): def _exif_utils_returning(metadata):
class MetadataExifUtils(DummyExifUtils): class MetadataExifUtils(DummyExifUtils):
def extract_image_metadata(self, path): def extract_image_metadata(self, path):
+98
View File
@@ -211,6 +211,104 @@ def test_update_image_metadata_preserves_png_workflow(tmp_path):
) )
def test_optimize_image_embeds_supplied_workflow_when_source_has_none(tmp_path):
"""Import paths hand the workflow over as data when the preview source is
metadata-free (CivitAI's optimized rendition); optimize_image must embed
it while re-encoding, otherwise the recipe loses has_workflow."""
image_path = tmp_path / "optimized.webp"
Image.new("RGB", (64, 32), color="red").save(image_path, format="WEBP", quality=85)
workflow = {"nodes": [{"id": 1}], "last_node_id": 1}
optimized_data, extension = ExifUtils.optimize_image(
str(image_path),
target_width=32,
format="webp",
quality=85,
preserve_metadata=True,
workflow=workflow,
)
optimized_path = tmp_path / f"embedded{extension}"
optimized_path.write_bytes(optimized_data)
metadata = ExifUtils._load_structured_metadata(str(optimized_path))
assert metadata["workflow"] == json.dumps(workflow)
def test_optimize_image_keeps_source_workflow_over_supplied(tmp_path):
image_path = tmp_path / "source.png"
png_info = PngImagePlugin.PngInfo()
png_info.add_text("workflow", '{"nodes":[{"id":7}]}')
Image.new("RGB", (64, 32), color="red").save(image_path, pnginfo=png_info)
optimized_data, extension = ExifUtils.optimize_image(
str(image_path),
target_width=32,
format="webp",
quality=85,
preserve_metadata=True,
workflow={"nodes": [{"id": 1}]},
)
optimized_path = tmp_path / f"sourcewins{extension}"
optimized_path.write_bytes(optimized_data)
metadata = ExifUtils._load_structured_metadata(str(optimized_path))
assert metadata["workflow"] == '{"nodes":[{"id":7}]}'
def test_embed_workflow_adds_workflow_to_metadata_free_webp(tmp_path):
image_path = tmp_path / "preview.webp"
Image.new("RGB", (32, 32), color="blue").save(image_path, format="WEBP", quality=85)
workflow = json.dumps({"nodes": [{"id": 1}]})
returned = ExifUtils.embed_workflow(str(image_path), workflow)
assert returned == str(image_path)
metadata = ExifUtils._load_structured_metadata(str(image_path))
assert metadata["workflow"] == workflow
with Image.open(image_path) as img:
assert img.size == (32, 32)
def test_embed_workflow_adds_workflow_to_metadata_free_png(tmp_path):
image_path = tmp_path / "preview.png"
Image.new("RGB", (32, 32), color="blue").save(image_path)
workflow = {"nodes": [{"id": 3}]}
ExifUtils.embed_workflow(str(image_path), workflow)
metadata = ExifUtils._load_structured_metadata(str(image_path))
assert metadata["workflow"] == json.dumps(workflow)
def test_embed_workflow_leaves_existing_workflow_untouched(tmp_path):
image_path = tmp_path / "preview.png"
png_info = PngImagePlugin.PngInfo()
png_info.add_text("workflow", '{"nodes":[{"id":9}]}')
Image.new("RGB", (32, 32), color="green").save(image_path, pnginfo=png_info)
ExifUtils.embed_workflow(str(image_path), {"nodes": [{"id": 1}]})
with Image.open(image_path) as img:
assert img.info["workflow"] == '{"nodes":[{"id":9}]}'
def test_embed_workflow_ignores_unsupported_payloads_and_containers(tmp_path):
image_path = tmp_path / "preview.webp"
Image.new("RGB", (16, 16), color="black").save(image_path, format="WEBP")
# Nothing to embed / unsupported payload types are no-ops.
assert ExifUtils.embed_workflow(str(image_path), None) == str(image_path)
assert ExifUtils.embed_workflow(str(image_path), "") == str(image_path)
assert ExifUtils.embed_workflow(str(image_path), 123) == str(image_path)
assert ExifUtils._load_structured_metadata(str(image_path))["workflow"] is None
video_path = tmp_path / "clip.mp4"
video_path.write_bytes(b"video")
assert ExifUtils.embed_workflow(str(video_path), {"nodes": []}) == str(video_path)
# --- ISOBMFF / brotli extraction tests --- # --- ISOBMFF / brotli extraction tests ---
import struct import struct
+14
View File
@@ -873,6 +873,19 @@ export function createContextMenu(x, y, loraName, widget, previewTooltip, render
} }
); );
// Save recipe with the current graph embedded. Kept opt-in rather than
// folded into "Save Recipe": the workflow dwarfs every other metadata field
// and can carry sensitive widget values, so it stays an explicit choice.
const saveWithWorkflowOption = createMenuItem(
'Save Recipe with Workflow',
'<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><circle cx="18" cy="5" r="3"></circle><circle cx="6" cy="12" r="3"></circle><circle cx="18" cy="19" r="3"></circle><line x1="8.59" y1="13.51" x2="15.42" y2="17.49"></line><line x1="15.41" y1="6.51" x2="8.59" y2="10.49"></line></svg>',
() => {
menu.remove();
document.removeEventListener('click', closeMenu);
saveRecipeDirectly({ embedWorkflow: true });
}
);
// Move Up option with arrow up icon // Move Up option with arrow up icon
const moveUpOption = createMenuItem( const moveUpOption = createMenuItem(
'Move Up', 'Move Up',
@@ -982,6 +995,7 @@ export function createContextMenu(x, y, loraName, widget, previewTooltip, render
menu.appendChild(copyTriggerWordsOption); menu.appendChild(copyTriggerWordsOption);
menu.appendChild(separator2); menu.appendChild(separator2);
menu.appendChild(saveOption); menu.appendChild(saveOption);
menu.appendChild(saveWithWorkflowOption);
document.body.appendChild(menu); document.body.appendChild(menu);
+36 -9
View File
@@ -460,15 +460,29 @@ export function syncClipStrengthIfCollapsed(loraData) {
} }
// Function to directly save the recipe without dialog // Function to directly save the recipe without dialog
export async function saveRecipeDirectly() { export async function saveRecipeDirectly({ embedWorkflow = false } = {}) {
try { try {
const prompt = await app.graphToPrompt(); const prompt = await app.graphToPrompt();
console.log('Prompt:', prompt); // for debugging purposes console.log('Prompt:', prompt); // for debugging purposes
// Embedding the graph is opt-in: it is by far the largest metadata field
// and its widget values can contain sensitive data (paths, API keys). The
// UI-format graph is sent rather than the API prompt so node layout and
// groups survive — that is what "Send Workflow to ComfyUI" restores.
const requestBody = {};
if (embedWorkflow) {
if (prompt && prompt.workflow) {
requestBody.workflow = prompt.workflow;
} else {
showToast('No workflow available to embed; saving the recipe without it', 'warning');
}
}
// Show loading toast // Show loading toast
if (app && app.extensionManager && app.extensionManager.toast) { if (app && app.extensionManager && app.extensionManager.toast) {
app.extensionManager.toast.add({ app.extensionManager.toast.add({
severity: 'info', severity: 'info',
summary: 'Saving Recipe', summary: embedWorkflow ? 'Saving Recipe with Workflow' : 'Saving Recipe',
detail: 'Please wait...', detail: 'Please wait...',
life: 2000 life: 2000
}); });
@@ -476,7 +490,9 @@ export async function saveRecipeDirectly() {
// Send the request to the backend API // Send the request to the backend API
const response = await fetch(lmUrl('/api/lm/recipes/save-from-widget'), { const response = await fetch(lmUrl('/api/lm/recipes/save-from-widget'), {
method: 'POST' method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify(requestBody)
}); });
const result = await response.json(); const result = await response.json();
@@ -484,12 +500,23 @@ export async function saveRecipeDirectly() {
// Show result toast // Show result toast
if (app && app.extensionManager && app.extensionManager.toast) { if (app && app.extensionManager && app.extensionManager.toast) {
if (result.success) { if (result.success) {
app.extensionManager.toast.add({ let severity = 'success';
severity: 'success', let summary = embedWorkflow ? 'Recipe Saved with Workflow' : 'Recipe Saved';
summary: 'Recipe Saved', let detail = embedWorkflow
detail: 'Recipe has been saved successfully', ? 'Recipe and the current workflow have been saved'
life: 3000 : 'Recipe has been saved successfully';
});
if (embedWorkflow && result.workflow_skipped === 'too_large') {
severity = 'warn';
summary = 'Recipe Saved without Workflow';
detail = 'The workflow is too large to embed; the recipe was saved without it';
} else if (embedWorkflow && result.has_workflow !== true) {
severity = 'warn';
summary = 'Recipe Saved without Workflow';
detail = 'The workflow could not be embedded in the recipe image';
}
app.extensionManager.toast.add({ severity, summary, detail, life: 5000 });
} else { } else {
app.extensionManager.toast.add({ app.extensionManager.toast.add({
severity: 'error', severity: 'error',