From 5a7b40c2f1c12dcb4e346e338d6e43fe50628c90 Mon Sep 17 00:00:00 2001 From: LuqP2 Date: Mon, 1 Jun 2026 11:43:04 -0300 Subject: [PATCH] fix: harden MetaHub save node output handling --- README.md | 10 ++ metadata_utils.py | 192 +++++++++++++++++++++++++---- metahub_save_node.py | 67 ++++++---- metahub_save_video_node.py | 32 +++++ tests/test_metadata_utils.py | 65 ++++++++++ tests/test_video_metadata_utils.py | 46 +++++++ tests/test_workflow_extractor.py | 65 ++++++++++ video_metadata_utils.py | 36 +++--- workflow_extractor.py | 6 +- 9 files changed, 451 insertions(+), 68 deletions(-) create mode 100644 tests/test_video_metadata_utils.py diff --git a/README.md b/README.md index 4b0d854..e37a8be 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,10 @@ All overrides use `forceInput: True` (hidden from UI, connection-only). ### Filename Pattern Use `filename_pattern` to customize names (default `ComfyUI_%counter%`). Invalid filename characters are sanitized and missing values become `unknown`. `filename_prefix` is deprecated but still supported. +`filename_pattern` can include subfolders. Each path segment is sanitized separately, so `project/%date:yyyyMMdd%/image_%time%` creates a dated folder under the selected output directory. Unsafe segments such as `..`, drive letters, absolute paths, empty path segments, and Windows reserved names are blocked. + +`output_path` controls the base directory. Leave it empty for ComfyUI's default output folder. Relative paths are created under the ComfyUI output folder; absolute paths are used as-is. + | Placeholder | Description | Example | | --- | --- | --- | | `%counter%` | Auto-increment (00001, 00002...) | 00042 | @@ -136,6 +140,8 @@ Use `filename_pattern` to customize names (default `ComfyUI_%counter%`). Invalid | `%date%` | Date YYYY-MM-DD | 2025-01-15 | | `%time%` | Time HH-MM-SS | 14-30-22 | | `%datetime%` | Date + time | 2025-01-15_14-30-22 | +| `%date:yyyyMMdd%` | Custom date format | 20250115 | +| `%time:HHmmss%` | Custom time format | 143022 | | `%model%` | Model name (no extension) | dreamshaper_v8 | | `%sampler%` | Sampler name | euler_ancestral | | `%steps%` | Steps | 20 | @@ -152,6 +158,8 @@ Use `filename_pattern` to customize names (default `ComfyUI_%counter%`). Invalid If structured metadata cannot be written, the image is still saved. +If no output file can be written, the node now raises a visible ComfyUI error. In batch saves, successful files still appear in the preview and failed files are listed in the console. + ### Multi-Sampler Heuristic When multiple samplers exist, the node prefers the sampler that feeds the latent used by the VAEDecode connected to the saved images; it falls back to the first sampler found. @@ -207,6 +215,8 @@ export COMFYUI_LORA_PATH="/path/to/loras" export COMFYUI_VAE_PATH="/path/to/vae" ``` +Checkpoint hashes are searched in both `models/checkpoints` and `models/diffusion_models` by default, matching newer Flux-style ComfyUI layouts. + ## Metadata Formats ### A1111/Civitai Format (PNG: tEXt "parameters", JPEG/WebP: EXIF UserComment) diff --git a/metadata_utils.py b/metadata_utils.py index 7d4eae7..b3450b1 100644 --- a/metadata_utils.py +++ b/metadata_utils.py @@ -6,6 +6,7 @@ Handles hash calculation, metadata formatting, and metadata injection. import hashlib import json import os +import re from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional, Tuple @@ -56,19 +57,27 @@ def get_model_search_paths(model_type: str) -> List[str]: import folder_paths type_map = { "checkpoint": "checkpoints", + "diffusion_model": "diffusion_models", "lora": "loras", "vae": "vae" } - if model_type in type_map: - comfy_paths = folder_paths.get_folder_paths(type_map[model_type]) - paths.extend(comfy_paths) + comfy_types = [type_map[model_type]] if model_type in type_map else [] + if model_type == "checkpoint": + comfy_types.append("diffusion_models") + for comfy_type in comfy_types: + try: + comfy_paths = folder_paths.get_folder_paths(comfy_type) + paths.extend(comfy_paths) + except Exception: + pass except (ImportError, Exception): pass # Add default relative paths if no paths found if not paths: default_paths = { - "checkpoint": ["models/checkpoints"], + "checkpoint": ["models/checkpoints", "models/diffusion_models"], + "diffusion_model": ["models/diffusion_models"], "lora": ["models/loras"], "vae": ["models/vae"], } @@ -88,18 +97,19 @@ def find_model_file(model_name: str, model_type: str) -> Optional[Path]: Returns: Path to model file or None if not found """ - # Ensure .safetensors extension - if not model_name.endswith('.safetensors'): - model_name += '.safetensors' + candidate_names = [model_name] + if "." not in Path(model_name).name: + candidate_names.append(f"{model_name}.safetensors") # Get search paths search_paths = get_model_search_paths(model_type) # Search for file for base_path in search_paths: - potential_path = Path(base_path) / model_name - if potential_path.exists(): - return potential_path + for candidate_name in candidate_names: + potential_path = Path(base_path) / candidate_name + if potential_path.exists(): + return potential_path return None @@ -383,6 +393,8 @@ def build_imh_metadata(params: dict, workflow_json: dict) -> dict: return { # CRITICAL: Required field for IMH detection "generator": "ComfyUI", + "metadata_status": params.get("metadata_status", "partial"), + "metadata_sources": _sanitize(params.get("metadata_sources", {})), # Main fields (compatible with comfyUIParser.ts) "prompt": params['positive'], @@ -443,6 +455,30 @@ def build_imh_metadata(params: dict, workflow_json: dict) -> dict: } +def build_metadata_sources(manual_inputs: dict, extracted: dict, fields: List[str]) -> Dict[str, str]: + sources: Dict[str, str] = {} + for field in fields: + manual_value = manual_inputs.get(field) + extracted_value = extracted.get(field) + if manual_value is not None and not (isinstance(manual_value, str) and not manual_value.strip()): + sources[field] = "manual_override" + elif extracted_value is not None and not (isinstance(extracted_value, str) and not extracted_value.strip()): + sources[field] = "detected" + else: + sources[field] = "default" + return sources + + +def build_metadata_status(sources: Dict[str, str], important_fields: Optional[List[str]] = None) -> str: + important = important_fields or ["positive", "model_name", "seed", "steps", "sampler_name"] + relevant = [sources.get(field, "unknown") for field in important] + if relevant and all(source in {"detected", "manual_override"} for source in relevant): + return "complete" + if any(source in {"detected", "manual_override"} for source in sources.values()): + return "partial" + return "fallback" + + def _serialize_metadata_json(value: Any) -> Optional[str]: if value is None: return None @@ -828,6 +864,12 @@ def save_webp_with_metadata( _INVALID_FILENAME_CHARS = '<>:"/\\|?*' _KNOWN_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp") +_WINDOWS_RESERVED_NAMES = { + "CON", "PRN", "AUX", "NUL", + "COM1", "COM2", "COM3", "COM4", "COM5", "COM6", "COM7", "COM8", "COM9", + "LPT1", "LPT2", "LPT3", "LPT4", "LPT5", "LPT6", "LPT7", "LPT8", "LPT9", +} +_TOKEN_RE = re.compile(r"%(date|time|datetime)(?::([^%]+))?%") def _normalize_file_format(file_format: str) -> str: @@ -866,6 +908,15 @@ def sanitize_filename(name: str) -> str: return cleaned or "unknown" +def sanitize_path_segment(name: str) -> str: + cleaned = sanitize_filename(name) + if cleaned in {".", ".."}: + raise ValueError(f"Unsafe path segment: {name!r}") + if cleaned.split(".", 1)[0].upper() in _WINDOWS_RESERVED_NAMES: + cleaned = f"_{cleaned}" + return cleaned + + def _format_placeholder_value(value: Any) -> str: if value is None: return "unknown" @@ -874,6 +925,39 @@ def _format_placeholder_value(value: Any) -> str: return str(value) +def _translate_datetime_format(fmt: str) -> str: + translated = fmt + replacements = ( + ("yyyy", "%Y"), + ("yy", "%y"), + ("MM", "%m"), + ("dd", "%d"), + ("HH", "%H"), + ("mm", "%M"), + ("ss", "%S"), + ) + for source, target in replacements: + translated = translated.replace(source, target) + return translated + + +def resolve_time_tokens(value: str, now: Optional[datetime] = None) -> str: + now = now or datetime.now() + + def replace(match: re.Match) -> str: + token = match.group(1) + fmt = match.group(2) + if fmt: + return now.strftime(_translate_datetime_format(fmt)) + if token == "date": + return now.strftime("%Y-%m-%d") + if token == "time": + return now.strftime("%H-%M-%S") + return now.strftime("%Y-%m-%d_%H-%M-%S") + + return _TOKEN_RE.sub(replace, value or "") + + def resolve_placeholders(pattern: str, params: dict, counter: int) -> str: now = datetime.now() model_name = params.get("model_name") or "" @@ -893,12 +977,40 @@ def resolve_placeholders(pattern: str, params: dict, counter: int) -> str: "%height%": _format_placeholder_value(params.get("height")), } - resolved = pattern or "ComfyUI_%counter%" + resolved = resolve_time_tokens(pattern or "ComfyUI_%counter%", now) for placeholder, value in replacements.items(): resolved = resolved.replace(placeholder, value) return resolved +def _get_default_output_directory() -> Path: + try: + import folder_paths + return Path(folder_paths.get_output_directory()) + except ImportError: + return Path("output") + + +def _has_drive_or_root(path_text: str) -> bool: + normalized = path_text.replace("\\", "/") + return bool(re.match(r"^[A-Za-z]:", normalized)) or normalized.startswith("/") + + +def _sanitize_relative_segments(path_text: str) -> Path: + normalized = path_text.replace("\\", "/") + if _has_drive_or_root(normalized): + raise ValueError(f"Absolute paths are not allowed here: {path_text!r}") + raw_segments = normalized.split("/") + if any(segment == "" for segment in raw_segments): + raise ValueError(f"Empty path segments are not allowed: {path_text!r}") + safe_segments = [] + for segment in raw_segments: + if segment in {".", ".."}: + raise ValueError(f"Unsafe path segment: {segment!r}") + safe_segments.append(sanitize_path_segment(segment)) + return Path(*safe_segments) if safe_segments else Path() + + def get_output_directory(output_path: str) -> Path: """ Determines output directory with automatic creation (mkdir -p). @@ -911,15 +1023,16 @@ def get_output_directory(output_path: str) -> Path: Path object for output directory (created if doesn't exist) """ if output_path and output_path.strip(): - output_dir = Path(output_path) + resolved_output = resolve_time_tokens(str(output_path).strip()) + candidate = Path(resolved_output) + if candidate.is_absolute(): + output_dir = candidate + elif _has_drive_or_root(resolved_output): + raise RuntimeError(f"Output path must be absolute or relative, not drive-relative: {output_path!r}") + else: + output_dir = _get_default_output_directory() / _sanitize_relative_segments(resolved_output) else: - # Use ComfyUI's default output directory - try: - import folder_paths - output_dir = Path(folder_paths.get_output_directory()) - except ImportError: - # Fallback if not in ComfyUI environment - output_dir = Path("output") + output_dir = _get_default_output_directory() # Create directory if it doesn't exist (mkdir -p) try: @@ -931,6 +1044,18 @@ def get_output_directory(output_path: str) -> Path: return output_dir +def resolve_filename_pattern_path(filename_pattern: str, params: dict, counter: int, file_format: str) -> Path: + resolved = resolve_placeholders(filename_pattern, params, counter) + if _has_drive_or_root(resolved): + raise ValueError(f"Filename pattern cannot contain an absolute path: {filename_pattern!r}") + relative = _sanitize_relative_segments(resolved) + if not relative.parts: + relative = Path("unknown") + parent = relative.parent + filename = sanitize_path_segment(_strip_known_extension(relative.name)) + _get_extension(file_format) + return parent / filename + + def get_next_filename(output_dir: Path, filename_pattern: str, params: dict, file_format: str) -> Path: """ Generates next available filename with auto-increment counter. @@ -949,13 +1074,11 @@ def get_next_filename(output_dir: Path, filename_pattern: str, params: dict, fil """ counter = 1 uses_counter = "%counter%" in (filename_pattern or "") - extension = _get_extension(file_format) while True: - resolved = resolve_placeholders(filename_pattern, params, counter) + candidate_pattern = filename_pattern if not uses_counter and counter > 1: - resolved = f"{resolved}_{counter:05d}" - filename = sanitize_filename(_strip_known_extension(resolved)) + extension - full_path = output_dir / filename + candidate_pattern = f"{filename_pattern}_{counter:05d}" + full_path = output_dir / resolve_filename_pattern_path(candidate_pattern, params, counter, file_format) if not full_path.exists(): return full_path counter += 1 @@ -1043,9 +1166,10 @@ def save_image_batch( imh_metadata: IMH metadata dict Returns: - List of saved file paths + List of saved file paths. Raises if none were saved. """ saved_paths = [] + failures: List[Tuple[Path, str]] = [] for image_tensor in images: # Convert to PIL @@ -1053,10 +1177,13 @@ def save_image_batch( # Get next filename file_path = get_next_filename(output_dir, filename_pattern, params, file_format) + file_path.parent.mkdir(parents=True, exist_ok=True) # Save image with metadata try: save_image_with_metadata(pil_image, file_path, file_format, quality, a1111_metadata, imh_metadata or {}) + if not file_path.exists() or file_path.stat().st_size <= 0: + raise RuntimeError("save completed but output file is missing or empty") saved_paths.append(file_path) except Exception as e: print(f"[MetaHub] Warning: Could not save image {file_path}: {e}") @@ -1070,9 +1197,24 @@ def save_image_batch( pil_image.save(file_path, "WEBP", quality=quality) else: pil_image.save(file_path, "PNG", compress_level=4) + if not file_path.exists() or file_path.stat().st_size <= 0: + raise RuntimeError("fallback save completed but output file is missing or empty") saved_paths.append(file_path) except Exception as save_error: print(f"[MetaHub] Warning: Fallback save failed for {file_path}: {save_error}") + failures.append((file_path, str(save_error))) + + if not saved_paths: + details = "; ".join(f"{path}: {error}" for path, error in failures) or "no images were processed" + raise RuntimeError(f"No images were saved. {details}") + + if failures: + print( + "[MetaHub] WARNING: Partial save failure - " + f"{len(saved_paths)} saved, {len(failures)} failed" + ) + for failed_path, error in failures: + print(f"[MetaHub] FAILED: {failed_path} ({error})") return saved_paths diff --git a/metahub_save_node.py b/metahub_save_node.py index 0560bc7..1e2150a 100644 --- a/metahub_save_node.py +++ b/metahub_save_node.py @@ -53,6 +53,19 @@ class MetaHubSaveNode: "default": "", "tooltip": "Custom output directory (empty = ComfyUI default)" }), + "user_tags": ("STRING", { + "default": "", + "tooltip": "Optional Image MetaHub tags (comma-separated)" + }), + "notes": ("STRING", { + "multiline": True, + "default": "", + "tooltip": "Optional Image MetaHub notes" + }), + "project_name": ("STRING", { + "default": "", + "tooltip": "Optional Image MetaHub project name" + }), "seed": ("INT", { "default": None, "forceInput": True, @@ -105,19 +118,6 @@ class MetaHubSaveNode: "forceInput": True, "tooltip": "Override VAE filename" }), - "user_tags": ("STRING", { - "default": "", - "tooltip": "IMH Pro: User tags (comma-separated)" - }), - "notes": ("STRING", { - "multiline": True, - "default": "", - "tooltip": "IMH Pro: Notes" - }), - "project_name": ("STRING", { - "default": "", - "tooltip": "IMH Pro: Project name" - }), "generation_time": ("FLOAT", { "default": 0.0, "min": 0.0, @@ -241,6 +241,32 @@ class MetaHubSaveNode: negative_value = resolve_value(negative, extracted.get("negative"), "") denoise_value = normalize_float(resolve_value(denoise, extracted.get("denoise"), 1.0), 1.0) vae_name_value = resolve_value(vae_name, extracted.get("vae_name"), "") + metadata_fields = [ + "seed", + "steps", + "cfg", + "sampler_name", + "scheduler", + "model_name", + "positive", + "negative", + "denoise", + "vae_name", + ] + manual_inputs = { + "seed": seed, + "steps": steps, + "cfg": cfg, + "sampler_name": sampler_name, + "scheduler": scheduler, + "model_name": model_name, + "positive": positive, + "negative": negative, + "denoise": denoise, + "vae_name": vae_name, + } + metadata_sources = utils.build_metadata_sources(manual_inputs, extracted, metadata_fields) + metadata_status = utils.build_metadata_status(metadata_sources) quality_value = normalize_int(quality if quality is not None else 95, 95) quality_value = max(1, min(100, quality_value)) @@ -323,6 +349,8 @@ class MetaHubSaveNode: "comfyui_version": version_info.get("comfyui_version"), "torch_version": version_info.get("torch_version"), "python_version": version_info.get("python_version"), + "metadata_status": metadata_status, + "metadata_sources": metadata_sources, } a1111_metadata = utils.build_a1111_metadata(params) @@ -355,16 +383,6 @@ class MetaHubSaveNode: "sampler_name": "sampler", "model_name": "model", } - manual_inputs = { - "seed": seed, - "steps": steps, - "cfg": cfg, - "sampler_name": sampler_name, - "scheduler": scheduler, - "model_name": model_name, - "positive": positive, - "negative": negative, - } missing_warn = [] for field in missing_fields: if field not in warn_fields: @@ -409,8 +427,7 @@ class MetaHubSaveNode: except Exception as e: print(f"[ImageMetaHub-Save] Warning: Save failed: {e}") - # Return empty UI structure on error - return {"ui": {"images": []}} + raise RuntimeError(f"MetaHub Save Image failed: {e}") from e NODE_CLASS_MAPPINGS = {"MetaHubSaveNode": MetaHubSaveNode} diff --git a/metahub_save_video_node.py b/metahub_save_video_node.py index 30df437..47cea63 100644 --- a/metahub_save_video_node.py +++ b/metahub_save_video_node.py @@ -240,6 +240,32 @@ class MetaHubSaveVideoNode: negative_value = resolve_value(negative, extracted.get("negative"), "") denoise_value = normalize_float(resolve_value(denoise, extracted.get("denoise"), 1.0), 1.0) vae_name_value = resolve_value(vae_name, extracted.get("vae_name"), "") + metadata_fields = [ + "seed", + "steps", + "cfg", + "sampler_name", + "scheduler", + "model_name", + "positive", + "negative", + "denoise", + "vae_name", + ] + manual_inputs = { + "seed": seed, + "steps": steps, + "cfg": cfg, + "sampler_name": sampler_name, + "scheduler": scheduler, + "model_name": model_name, + "positive": positive, + "negative": negative, + "denoise": denoise, + "vae_name": vae_name, + } + metadata_sources = utils.build_metadata_sources(manual_inputs, extracted, metadata_fields) + metadata_status = utils.build_metadata_status(metadata_sources) # Calculate model hash if model_name_value: @@ -292,6 +318,8 @@ class MetaHubSaveVideoNode: gpu_metrics=gpu_metrics, version_info=version_info, workflow_json=workflow_json, + metadata_sources=metadata_sources, + metadata_status=metadata_status, ) except Exception as e: print(f"[MetaHub Video] Warning: Failed to add metadata to {video_path.name}: {e}") @@ -325,6 +353,8 @@ class MetaHubSaveVideoNode: gpu_metrics: Dict[str, Any], version_info: Dict[str, str], workflow_json: Dict, + metadata_sources: Dict[str, str], + metadata_status: str, ) -> None: """ Processes a single video file and injects metadata. @@ -383,6 +413,8 @@ class MetaHubSaveVideoNode: "comfyui_version": version_info.get("comfyui_version"), "torch_version": version_info.get("torch_version"), "python_version": version_info.get("python_version"), + "metadata_sources": metadata_sources, + "metadata_status": metadata_status, } # Build metadata diff --git a/tests/test_metadata_utils.py b/tests/test_metadata_utils.py index 80eb766..242d0d8 100644 --- a/tests/test_metadata_utils.py +++ b/tests/test_metadata_utils.py @@ -1,14 +1,20 @@ import json import math +import sys from pathlib import Path import copy from metadata_utils import ( build_a1111_metadata, build_imh_metadata, + build_metadata_sources, + build_metadata_status, ensure_metahub_save_node, ensure_prompt_in_workflow, + get_next_filename, + get_output_directory, get_workflow_json, + resolve_filename_pattern_path, save_png_with_metadata, ) from PIL import Image @@ -37,6 +43,61 @@ def _save_simple_png(path: Path) -> None: image.save(path, "PNG") +def test_output_path_relative_resolves_under_comfy_output(tmp_path, monkeypatch): + class FakeFolderPaths: + @staticmethod + def get_output_directory(): + return str(tmp_path / "output") + + monkeypatch.setitem(sys.modules, "folder_paths", FakeFolderPaths) + + output = get_output_directory("project/%date:yyyyMMdd%") + + assert output.parent.name == "project" + assert output.parent.parent == tmp_path / "output" + assert output.exists() + + +def test_filename_pattern_supports_subfolders_and_tokens(): + params = {"seed": 123, "model_name": "model.safetensors"} + + relative = resolve_filename_pattern_path("%date:yyyyMMdd%/image_%seed%", params, 1, "PNG") + + assert len(relative.parts) == 2 + assert relative.parts[1] == "image_123.png" + + +def test_filename_pattern_rejects_unsafe_segments(): + try: + resolve_filename_pattern_path("../escape", {}, 1, "PNG") + except ValueError as error: + assert "Unsafe path segment" in str(error) + else: + raise AssertionError("Expected unsafe pattern to fail") + + +def test_get_next_filename_creates_nested_collision_path(tmp_path): + existing = tmp_path / "set" / "image.png" + existing.parent.mkdir() + existing.write_bytes(b"exists") + + path = get_next_filename(tmp_path, "set/image", {}, "PNG") + + assert path == tmp_path / "set" / "image_00002.png" + + +def test_metadata_status_marks_defaulted_fields_as_partial(): + sources = build_metadata_sources( + {"positive": None, "model_name": None, "seed": None}, + {"positive": "prompt"}, + ["positive", "model_name", "seed"], + ) + + assert sources["positive"] == "detected" + assert sources["model_name"] == "default" + assert build_metadata_status(sources, ["positive", "model_name", "seed"]) == "partial" + + def test_ensure_prompt_in_workflow_adds_prompt(): workflow = {"workflow": {"nodes": []}} prompt = {"1": {"class_type": "KSampler", "inputs": {}}} @@ -149,6 +210,8 @@ def test_build_imh_metadata_includes_workflow_and_prompt(): "generation_type": "img2img", "parent_image": {"fileName": "selected.png", "relativePath": "library/selected.png"}, "source_image": {"fileName": "base.png", "relativePath": "inputs/base.png"}, + "metadata_status": "partial", + "metadata_sources": {"model_name": "default"}, } workflow_json = {"workflow": {"nodes": []}, "prompt": {"1": {"class_type": "KSampler"}}} @@ -160,6 +223,8 @@ def test_build_imh_metadata_includes_workflow_and_prompt(): assert imh["generation_type"] == "img2img" assert imh["parent_image"]["fileName"] == "selected.png" assert imh["source_image"]["fileName"] == "base.png" + assert imh["metadata_status"] == "partial" + assert imh["metadata_sources"] == {"model_name": "default"} def test_build_imh_metadata_sanitizes_nan(): diff --git a/tests/test_video_metadata_utils.py b/tests/test_video_metadata_utils.py new file mode 100644 index 0000000..6961c4d --- /dev/null +++ b/tests/test_video_metadata_utils.py @@ -0,0 +1,46 @@ +from video_metadata_utils import build_video_metahub_metadata + + +def test_video_metadata_omits_full_workflow_payloads(): + payload = build_video_metahub_metadata( + { + "positive": "prompt", + "negative": "negative", + "width": 640, + "height": 480, + "frame_rate": 24, + "frame_count": 48, + "metadata_status": "complete", + "metadata_sources": {"positive": "detected"}, + }, + {"workflow": {"nodes": [{"id": 1}]}, "prompt": {"1": {"class_type": "KSampler"}}}, + ) + + assert payload["generator"] == "ComfyUI" + assert payload["media_type"] == "video" + assert payload["video"]["duration_seconds"] == 2 + assert payload["metadata_status"] == "complete" + assert "workflow" not in payload + assert "prompt_api" not in payload + + +def test_video_metadata_nulls_defaulted_canonical_fields(): + payload = build_video_metahub_metadata( + { + "positive": "", + "model_name": "", + "seed": 0, + "width": 640, + "height": 480, + "metadata_sources": { + "positive": "default", + "model_name": "default", + "seed": "default", + }, + }, + {}, + ) + + assert payload["prompt"] is None + assert payload["model"] is None + assert payload["seed"] is None diff --git a/tests/test_workflow_extractor.py b/tests/test_workflow_extractor.py index 51e0827..e25f5c2 100644 --- a/tests/test_workflow_extractor.py +++ b/tests/test_workflow_extractor.py @@ -70,6 +70,71 @@ def test_workflow_extractor_lora_detection(): assert data["lora_list"] == [{"name": "detail.safetensors", "weight": 0.8}] +def test_workflow_extractor_detects_flux_unet_model_and_vae(): + prompt = { + "1": { + "class_type": "UNETLoader", + "inputs": {"unet_name": "flux-model.safetensors"}, + }, + "2": { + "class_type": "VAELoader", + "inputs": {"vae_name": "flux-vae.safetensors"}, + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "positive", "clip": ["9", 0]}, + }, + "4": { + "class_type": "BasicScheduler", + "inputs": {"model": ["1", 0], "steps": 8, "scheduler": "normal"}, + }, + "5": { + "class_type": "KSampler", + "inputs": { + "seed": 123, + "steps": 8, + "cfg": 1, + "sampler_name": "euler", + "scheduler": "normal", + "model": ["1", 0], + "positive": ["3", 0], + "negative": ["3", 0], + "latent_image": ["8", 0], + }, + }, + "6": { + "class_type": "VAEDecode", + "inputs": {"samples": ["5", 0], "vae": ["2", 0]}, + }, + "7": { + "class_type": "MetaHubSaveNode", + "inputs": {"images": ["6", 0]}, + }, + "8": {"class_type": "EmptyLatentImage", "inputs": {}}, + "9": {"class_type": "CLIPLoader", "inputs": {}}, + } + + extractor = WorkflowExtractor(prompt) + data, _missing = extractor.extract(save_node_id="7") + + assert data["model_name"] == "flux-model.safetensors" + assert data["vae_name"] == "flux-vae.safetensors" + + +def test_workflow_extractor_detects_lora_alternate_keys(): + prompt = { + "1": { + "class_type": "PowerLoraLoader", + "inputs": {"lora_name_1": "style.safetensors", "strength": 0.65}, + } + } + + extractor = WorkflowExtractor(prompt) + data, _missing = extractor.extract() + + assert data["lora_list"] == [{"name": "style.safetensors", "weight": 0.65}] + + def test_workflow_extractor_detects_img2img_lineage(): prompt = { diff --git a/video_metadata_utils.py b/video_metadata_utils.py index c50632e..eee7616 100644 --- a/video_metadata_utils.py +++ b/video_metadata_utils.py @@ -199,9 +199,6 @@ def build_video_metahub_metadata(params: dict, workflow_json: dict) -> dict: return [_sanitize(v) for v in value] return value - safe_workflow = _sanitize(workflow_json.get('workflow', {})) - safe_prompt = _sanitize(workflow_json.get('prompt', {})) - # Calculate duration if we have frame count and rate duration_seconds = None frame_count = params.get('frame_count') @@ -209,23 +206,32 @@ def build_video_metahub_metadata(params: dict, workflow_json: dict) -> dict: if frame_count and frame_rate and frame_rate > 0: duration_seconds = round(frame_count / frame_rate, 2) + sources = params.get("metadata_sources") if isinstance(params.get("metadata_sources"), dict) else {} + + def canonical(field: str, param_key: str, fallback: Any = None) -> Any: + if sources.get(field) == "default": + return None + return params.get(param_key, fallback) + return { # CRITICAL: Required field for detection "generator": "ComfyUI", "media_type": "video", + "metadata_status": params.get("metadata_status", "partial"), + "metadata_sources": _sanitize(params.get("metadata_sources", {})), # Main generation fields - "prompt": params.get('positive', ''), - "negativePrompt": params.get('negative', ''), - "seed": params.get('seed', 0), - "steps": params.get('steps', 20), - "cfg": params.get('cfg', 7.0), - "sampler_name": params.get('sampler', 'euler'), - "scheduler": params.get('scheduler', 'normal'), - "model": params.get('model_name', ''), + "prompt": canonical("positive", "positive"), + "negativePrompt": canonical("negative", "negative"), + "seed": canonical("seed", "seed"), + "steps": canonical("steps", "steps"), + "cfg": canonical("cfg", "cfg"), + "sampler_name": canonical("sampler_name", "sampler"), + "scheduler": canonical("scheduler", "scheduler"), + "model": canonical("model_name", "model_name"), "model_hash": params.get('model_hash', ''), - "vae": params.get('vae_name', ''), - "denoise": params.get('denoise', 1.0), + "vae": canonical("vae_name", "vae_name"), + "denoise": canonical("denoise", "denoise"), "width": params.get('width', 512), "height": params.get('height', 512), @@ -270,9 +276,7 @@ def build_video_metahub_metadata(params: dict, workflow_json: dict) -> dict: "python_version": params.get('python_version'), }, - # Complete workflow data - "workflow": safe_workflow, - "prompt_api": safe_prompt, + # Full workflow JSON intentionally omitted for video containers. } diff --git a/workflow_extractor.py b/workflow_extractor.py index bbd12e8..5924b87 100644 --- a/workflow_extractor.py +++ b/workflow_extractor.py @@ -24,6 +24,8 @@ class WorkflowExtractor: "CheckpointLoader", "UNETLoader", "UnetLoaderGGUF", + "DualDiffusionLoader", + "DiffusionModelLoader", ] VAE_NODES = [ @@ -228,7 +230,7 @@ class WorkflowExtractor: continue has_lora_nodes = True inputs = node.get("inputs", {}) - lora_name = self._get_first_literal(inputs, ("lora_name", "lora")) + lora_name = self._get_first_literal(inputs, ("lora_name", "lora", "lora_name_1", "lora_1")) if not lora_name: continue strength_model = self._get_first_literal( @@ -407,7 +409,7 @@ class WorkflowExtractor: def _get_checkpoint_name(self, node: Dict[str, Any]) -> Optional[str]: inputs = node.get("inputs", {}) - for key in ("ckpt_name", "checkpoint", "model_name", "unet_name"): + for key in ("ckpt_name", "checkpoint", "model_name", "unet_name", "diffusion_model_name", "model"): value = self._get_literal_input(inputs, key) if value: return str(value)