From 6184b748092bf7174a2012c17d0a6cd9e885f8a0 Mon Sep 17 00:00:00 2001 From: Artificial Sweetener Date: Mon, 21 Sep 2026 00:17:31 -0400 Subject: [PATCH] feat(loaders): add Krea 2 model loader Add automatic FP8/BF16 Qwen encoder and shared VAE resolution, architecture validation, artifact-aware dropdown deduplication, and renamed artifact discovery across supported loaders. --- README.md | 8 +- simple_syrup/nodes/simple_load_anima.py | 23 +- simple_syrup/nodes_v3/__init__.py | 2 + simple_syrup/nodes_v3/simple_load_flux.py | 33 +- simple_syrup/nodes_v3/simple_load_flux2.py | 22 +- simple_syrup/nodes_v3/simple_load_krea2.py | 166 ++++++++++ simple_syrup/runtime/anima_artifacts.py | 18 +- simple_syrup/runtime/auto_model_artifact.py | 1 + simple_syrup/runtime/auto_model_choices.py | 72 ++++ simple_syrup/runtime/auto_model_resolver.py | 80 ++++- simple_syrup/runtime/clip_type_support.py | 25 ++ .../runtime/diffusion_model_metadata.py | 65 ++++ simple_syrup/runtime/flux_artifacts.py | 7 + simple_syrup/runtime/flux_model_inspector.py | 61 ++-- simple_syrup/runtime/krea2_artifacts.py | 56 ++++ simple_syrup/runtime/qwen_artifacts.py | 27 ++ simple_syrup/services/anima_loader_service.py | 5 +- simple_syrup/services/krea2_loader_service.py | 196 +++++++++++ tests/test_anima_loader.py | 5 +- tests/test_auto_model_choices.py | 112 +++++++ tests/test_auto_model_resolver.py | 109 ++++++ tests/test_clip_type_support.py | 39 +++ tests/test_diffusion_model_metadata.py | 55 +++ tests/test_flux_runtime_loaders.py | 1 + tests/test_krea2_artifacts.py | 39 +++ tests/test_krea2_loader_service.py | 313 ++++++++++++++++++ tests/test_model_catalog.py | 18 +- tests/test_registration.py | 1 + tests/test_simple_load_krea2_node.py | 117 +++++++ 29 files changed, 1574 insertions(+), 102 deletions(-) create mode 100644 simple_syrup/nodes_v3/simple_load_krea2.py create mode 100644 simple_syrup/runtime/auto_model_choices.py create mode 100644 simple_syrup/runtime/clip_type_support.py create mode 100644 simple_syrup/runtime/diffusion_model_metadata.py create mode 100644 simple_syrup/runtime/krea2_artifacts.py create mode 100644 simple_syrup/runtime/qwen_artifacts.py create mode 100644 simple_syrup/services/krea2_loader_service.py create mode 100644 tests/test_auto_model_choices.py create mode 100644 tests/test_clip_type_support.py create mode 100644 tests/test_diffusion_model_metadata.py create mode 100644 tests/test_krea2_artifacts.py create mode 100644 tests/test_krea2_loader_service.py create mode 100644 tests/test_simple_load_krea2_node.py diff --git a/README.md b/README.md index b4f90c8..513c80e 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ The pack now covers model loading, regional prompting and segmentation, high-res ## Highlights -- Loaders that keep checkpoints, Anima, FLUX.1, and FLUX.2 models together with the text encoders, VAE, precision, and device choices they need. +- Loaders that keep checkpoints, Anima, FLUX.1, FLUX.2, and Krea 2 models together with the text encoders, VAE, precision, and device choices they need. - My original Contextual Diffusion method for coherent high-resolution edits, plus MultiDiffusion and Mixture of Diffusers tiled sampling. - Impact-compatible SEGS detection, segmentation, interactive preview, batching, and detailers. - ADetailer-style `[SEP]` prompt batches, masked conditioning, and regional samplers, with optional Prompt Control scheduling and LoRA hooks. @@ -74,7 +74,9 @@ Loading a checkpoint used to feel like choosing one file. Newer model families c **Simple Load FLUX** handles FLUX.1 with CLIP-L, T5-XXL, and its VAE. **Simple Load FLUX.2** inspects the selected diffusion model and chooses the matching text encoder family for FLUX.2 dev, Klein 4B, or Klein 9B/KV conditioning. Both loaders can find or download their known text encoders and VAEs with visible Comfy progress. -The FLUX loaders only download those revision-locked, checksum-pinned support files. You still install and select the diffusion model. They also expose manual component selection, diffusion weight precision, and text-encoder device placement. Moving text encoding to the CPU can save VRAM, although it will take longer. +**Simple Load Krea 2** validates the selected Raw or Turbo diffusion model and loads the required Qwen3-VL 4B encoder with Krea's layered conditioning plus the Qwen Image VAE. Auto uses the official FP8-scaled encoder; the advanced encoder choice can download either the checksum-pinned FP8-scaled or BF16 file. + +The FLUX and Krea 2 loaders only download those revision-locked, checksum-pinned support files. You still install and select the diffusion model. They also expose manual component selection, diffusion weight precision, and text-encoder device placement. Moving text encoding to the CPU can save VRAM, although it will take longer. ## Large images and high-resolution edits @@ -164,7 +166,7 @@ SimpleSyrup adds three ComfyUI settings: - **SimpleSyrup: External LLM endpoint** stores the OpenAI-compatible base URL used to discover provider models and run the external prompt nodes. - **SimpleSyrup: External LLM API key** stores the provider key in OS credential storage. -With downloadable models enabled, selecting a known missing catalog entry lets its loader download the required files. With the setting disabled, the dropdowns contain models SimpleSyrup can verify locally. Anima, FLUX.1, and FLUX.2 support components are resolved by their own loaders and use checksum-pinned automatic choices. +With downloadable models enabled, selecting a known missing catalog entry lets its loader download the required files. With the setting disabled, the dropdowns contain models SimpleSyrup can verify locally. Anima, FLUX.1, FLUX.2, and Krea 2 support components are resolved by their own loaders and use checksum-pinned automatic choices. Automatic resolution checks cached and official paths first, then recognizes renamed files with matching size and checksum inside the appropriate ComfyUI model category. Known local support files are represented by their automatic choice instead of appearing again as manual dropdown entries. Saving the external LLM endpoint and API key refreshes the provider models available in connected SimpleSyrup nodes. Image inputs require a provider model with vision support. diff --git a/simple_syrup/nodes/simple_load_anima.py b/simple_syrup/nodes/simple_load_anima.py index 46aa698..895258b 100644 --- a/simple_syrup/nodes/simple_load_anima.py +++ b/simple_syrup/nodes/simple_load_anima.py @@ -11,10 +11,13 @@ from types import ModuleType from typing import Any from ..domain.anima_quantization import AnimaQuantizationRecipe +from ..runtime.anima_artifacts import ANIMA_QWEN_TEXT_ENCODER +from ..runtime.auto_model_choices import automatic_component_choices from ..runtime.diffusion_model_loader import DIFFUSION_WEIGHT_DTYPES from ..runtime.model_downloads import ComfyProgressReporter from ..runtime.quantization_capabilities import QuantizationCapabilityCatalog from ..runtime.quantization_progress import ComfyQuantizationProgressReporter +from ..runtime.qwen_artifacts import QWEN_IMAGE_VAE from ..runtime.vae_loader import vae_choices from ..services.anima_loader_service import ( AUTO_CHOICE, @@ -84,7 +87,12 @@ class SimpleLoadAnima: }, ), "text_encoder": ( - _choices_with_auto(folder_paths.get_filename_list("text_encoders")), + automatic_component_choices( + installed=folder_paths.get_filename_list("text_encoders"), + artifacts=(ANIMA_QWEN_TEXT_ENCODER,), + leading_choices=(AUTO_CHOICE,), + folder_paths_module=folder_paths, + ), { "default": AUTO_CHOICE, "advanced": True, @@ -106,7 +114,12 @@ class SimpleLoadAnima: }, ), "vae": ( - _choices_with_auto(vae_choices(folder_paths)), + automatic_component_choices( + installed=vae_choices(folder_paths), + artifacts=(QWEN_IMAGE_VAE,), + leading_choices=(AUTO_CHOICE,), + folder_paths_module=folder_paths, + ), { "default": AUTO_CHOICE, "advanced": True, @@ -142,12 +155,6 @@ class SimpleLoadAnima: ) -def _choices_with_auto(choices: list[str]) -> list[str]: - """Return choices with the automatic selection first and deduplicated.""" - - return [AUTO_CHOICE, *(choice for choice in choices if choice != AUTO_CHOICE)] - - def _folder_paths() -> ModuleType: """Import ComfyUI folder paths lazily.""" diff --git a/simple_syrup/nodes_v3/__init__.py b/simple_syrup/nodes_v3/__init__.py index 99af78a..69cc6cc 100644 --- a/simple_syrup/nodes_v3/__init__.py +++ b/simple_syrup/nodes_v3/__init__.py @@ -67,6 +67,7 @@ def get_nodes() -> list[type[object]]: from .simple_load_checkpoint import SimpleLoadCheckpointV3 from .simple_load_flux import SimpleLoadFluxV3 from .simple_load_flux2 import SimpleLoadFlux2V3 + from .simple_load_krea2 import SimpleLoadKrea2V3 from .tag_segs_with_external_llm import TagSEGSWithExternalLLMV3 from .tag_segs_with_wd14 import TagSEGSWithWD14V3 from .tile_and_tag_segs import TileAndTagSEGSV3 @@ -121,6 +122,7 @@ def get_nodes() -> list[type[object]]: SimpleLoadCheckpointV3, SimpleLoadFluxV3, SimpleLoadFlux2V3, + SimpleLoadKrea2V3, SimpleVAEEncodeV3, TagSEGSWithExternalLLMV3, TagSEGSWithWD14V3, diff --git a/simple_syrup/nodes_v3/simple_load_flux.py b/simple_syrup/nodes_v3/simple_load_flux.py index 0c3ef84..f3075d2 100644 --- a/simple_syrup/nodes_v3/simple_load_flux.py +++ b/simple_syrup/nodes_v3/simple_load_flux.py @@ -11,7 +11,9 @@ from types import ModuleType from typing import TYPE_CHECKING, Any, ClassVar from ..nodes import tooltips +from ..runtime.auto_model_choices import automatic_component_choices from ..runtime.diffusion_model_loader import DIFFUSION_WEIGHT_DTYPES +from ..runtime.flux_artifacts import FLUX_CLIP_L, FLUX_T5_XXL, FLUX_VAE from ..runtime.model_downloads import ComfyProgressReporter from ..runtime.text_encoder_loader import TEXT_ENCODER_DEVICES from ..runtime.vae_loader import vae_choices @@ -44,9 +46,7 @@ class SimpleLoadFluxV3(_ComfyNodeBase): """Declare the separate FLUX.1 loader schema.""" folder_paths = _folder_paths() - text_encoder_choices = _choices_with_auto( - list(folder_paths.get_filename_list("text_encoders")) - ) + installed_text_encoders = list(folder_paths.get_filename_list("text_encoders")) return _comfy_io.Schema( node_id="SimpleSyrup.SimpleLoadFlux", display_name="Simple Load FLUX", @@ -77,7 +77,12 @@ class SimpleLoadFluxV3(_ComfyNodeBase): ), _comfy_io.Combo.Input( "clip_l", - options=text_encoder_choices, + options=automatic_component_choices( + installed=installed_text_encoders, + artifacts=(FLUX_CLIP_L,), + leading_choices=(AUTO_CHOICE,), + folder_paths_module=folder_paths, + ), default=AUTO_CHOICE, advanced=True, tooltip=( @@ -87,7 +92,12 @@ class SimpleLoadFluxV3(_ComfyNodeBase): ), _comfy_io.Combo.Input( "t5_xxl", - options=text_encoder_choices, + options=automatic_component_choices( + installed=installed_text_encoders, + artifacts=(FLUX_T5_XXL,), + leading_choices=(AUTO_CHOICE,), + folder_paths_module=folder_paths, + ), default=AUTO_CHOICE, advanced=True, tooltip=( @@ -107,7 +117,12 @@ class SimpleLoadFluxV3(_ComfyNodeBase): ), _comfy_io.Combo.Input( "vae", - options=_choices_with_auto(vae_choices(folder_paths)), + options=automatic_component_choices( + installed=vae_choices(folder_paths), + artifacts=(FLUX_VAE,), + leading_choices=(AUTO_CHOICE,), + folder_paths_module=folder_paths, + ), default=AUTO_CHOICE, advanced=True, tooltip=( @@ -146,12 +161,6 @@ class SimpleLoadFluxV3(_ComfyNodeBase): ) -def _choices_with_auto(choices: list[str]) -> list[str]: - """Return deduplicated choices with automatic selection first.""" - - return [AUTO_CHOICE, *(choice for choice in choices if choice != AUTO_CHOICE)] - - def _folder_paths() -> ModuleType: """Import ComfyUI folder paths lazily for schema declaration.""" diff --git a/simple_syrup/nodes_v3/simple_load_flux2.py b/simple_syrup/nodes_v3/simple_load_flux2.py index 1189fd0..765ee70 100644 --- a/simple_syrup/nodes_v3/simple_load_flux2.py +++ b/simple_syrup/nodes_v3/simple_load_flux2.py @@ -11,7 +11,9 @@ from types import ModuleType from typing import TYPE_CHECKING, Any, ClassVar from ..nodes import tooltips +from ..runtime.auto_model_choices import automatic_component_choices from ..runtime.diffusion_model_loader import DIFFUSION_WEIGHT_DTYPES +from ..runtime.flux_artifacts import FLUX2_TEXT_ENCODERS, FLUX2_VAE from ..runtime.model_downloads import ComfyProgressReporter from ..runtime.text_encoder_loader import TEXT_ENCODER_DEVICES from ..runtime.vae_loader import vae_choices @@ -75,8 +77,11 @@ class SimpleLoadFlux2V3(_ComfyNodeBase): ), _comfy_io.Combo.Input( "text_encoder", - options=_choices_with_auto( - list(folder_paths.get_filename_list("text_encoders")) + options=automatic_component_choices( + installed=list(folder_paths.get_filename_list("text_encoders")), + artifacts=tuple(FLUX2_TEXT_ENCODERS.values()), + leading_choices=(AUTO_CHOICE,), + folder_paths_module=folder_paths, ), default=AUTO_CHOICE, advanced=True, @@ -98,7 +103,12 @@ class SimpleLoadFlux2V3(_ComfyNodeBase): ), _comfy_io.Combo.Input( "vae", - options=_choices_with_auto(vae_choices(folder_paths)), + options=automatic_component_choices( + installed=vae_choices(folder_paths), + artifacts=(FLUX2_VAE,), + leading_choices=(AUTO_CHOICE,), + folder_paths_module=folder_paths, + ), default=AUTO_CHOICE, advanced=True, tooltip=( @@ -135,12 +145,6 @@ class SimpleLoadFlux2V3(_ComfyNodeBase): ) -def _choices_with_auto(choices: list[str]) -> list[str]: - """Return deduplicated choices with automatic selection first.""" - - return [AUTO_CHOICE, *(choice for choice in choices if choice != AUTO_CHOICE)] - - def _folder_paths() -> ModuleType: """Import ComfyUI folder paths lazily for schema declaration.""" diff --git a/simple_syrup/nodes_v3/simple_load_krea2.py b/simple_syrup/nodes_v3/simple_load_krea2.py new file mode 100644 index 0000000..d1acfe3 --- /dev/null +++ b/simple_syrup/nodes_v3/simple_load_krea2.py @@ -0,0 +1,166 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Expose cohesive Krea 2 component loading through Comfy's v3 API.""" + +from __future__ import annotations + +import importlib +from types import ModuleType +from typing import TYPE_CHECKING, Any, ClassVar + +from ..nodes import tooltips +from ..runtime.auto_model_choices import automatic_component_choices +from ..runtime.diffusion_model_loader import DIFFUSION_WEIGHT_DTYPES +from ..runtime.krea2_artifacts import ( + KREA2_AUTO_TEXT_ENCODER, + KREA2_QWEN3_VL_4B_BF16, + KREA2_QWEN3_VL_4B_FP8, +) +from ..runtime.model_downloads import ComfyProgressReporter +from ..runtime.qwen_artifacts import QWEN_IMAGE_VAE +from ..runtime.text_encoder_loader import TEXT_ENCODER_DEVICES +from ..runtime.vae_loader import vae_choices +from ..services.krea2_loader_service import AUTO_CHOICE, Krea2LoaderService + +if TYPE_CHECKING: + + class _ComfyNodeBase: + """Type-checking base for Comfy v3 nodes.""" + + RETURN_TYPES: ClassVar[list[str]] + RETURN_NAMES: ClassVar[list[str]] + +else: + _ComfyNodeBase = importlib.import_module("comfy_api.latest").io.ComfyNode + +_comfy_io: Any = ( + None if TYPE_CHECKING else importlib.import_module("comfy_api.latest").io +) + + +class SimpleLoadKrea2V3(_ComfyNodeBase): + """Load a Krea 2 diffusion model with its Qwen encoder and image VAE.""" + + _service = Krea2LoaderService() + + @classmethod + def define_schema(cls) -> Any: + """Declare the Krea 2 loader schema and downloadable component choices.""" + + folder_paths = _folder_paths() + return _comfy_io.Schema( + node_id="SimpleSyrup.SimpleLoadKrea2", + display_name="Simple Load Krea 2", + category="SimpleSyrup/Loaders", + description=( + "Loads Krea 2 with its Qwen3-VL 4B encoder and Qwen Image VAE; " + "automatic components are downloaded from checksum-pinned " + "Hugging Face files." + ), + search_aliases=["krea", "krea 2", "k2", "load krea"], + inputs=[ + _comfy_io.Combo.Input( + "diffusion_model", + options=list(folder_paths.get_filename_list("diffusion_models")), + tooltip=( + "Krea 2 Raw or Turbo diffusion model to load. This node " + "validates the architecture and never downloads this file." + ), + ), + _comfy_io.Combo.Input( + "diffusion_weight_dtype", + options=list(DIFFUSION_WEIGHT_DTYPES), + default="default", + advanced=True, + tooltip=( + "Load-time diffusion precision; default preserves the " + "selected file's stored BF16, FP8, INT8, MXFP8, or NVFP4 " + "format." + ), + ), + _comfy_io.Combo.Input( + "text_encoder", + options=automatic_component_choices( + installed=list(folder_paths.get_filename_list("text_encoders")), + artifacts=( + KREA2_QWEN3_VL_4B_FP8, + KREA2_QWEN3_VL_4B_BF16, + ), + leading_choices=( + KREA2_AUTO_TEXT_ENCODER, + KREA2_QWEN3_VL_4B_FP8.filename, + KREA2_QWEN3_VL_4B_BF16.filename, + ), + folder_paths_module=folder_paths, + ), + default=KREA2_AUTO_TEXT_ENCODER, + advanced=True, + tooltip=( + "Qwen3-VL 4B encoder loaded with Krea 2's required 12-layer " + "conditioning. Auto uses FP8; selecting official FP8 or BF16 " + "downloads that checksum-pinned file when missing." + ), + ), + _comfy_io.Combo.Input( + "text_encoder_device", + options=list(TEXT_ENCODER_DEVICES), + default="default", + advanced=True, + tooltip=( + "Device for Qwen3-VL; CPU saves GPU memory but makes prompt " + "encoding slower." + ), + ), + _comfy_io.Combo.Input( + "vae", + options=automatic_component_choices( + installed=vae_choices(folder_paths), + artifacts=(QWEN_IMAGE_VAE,), + leading_choices=(AUTO_CHOICE,), + folder_paths_module=folder_paths, + ), + default=AUTO_CHOICE, + advanced=True, + tooltip=( + "VAE used to decode Krea 2 latents. Auto finds or downloads " + "the checksum-pinned Qwen Image VAE with visible progress." + ), + ), + ], + outputs=[ + _comfy_io.Model.Output("model", tooltip=tooltips.MODEL_OUTPUT), + _comfy_io.Clip.Output("clip", tooltip=tooltips.CLIP_OUTPUT), + _comfy_io.Vae.Output("vae", tooltip=tooltips.VAE_OUTPUT), + ], + ) + + @classmethod + def execute( + cls, + diffusion_model: str, + diffusion_weight_dtype: str, + text_encoder: str, + text_encoder_device: str, + vae: str, + ) -> tuple[object, object, object]: + """Load and return validated Krea 2 MODEL, CLIP, and VAE objects.""" + + return cls._service.load_models( + diffusion_model=diffusion_model, + diffusion_weight_dtype=diffusion_weight_dtype, + text_encoder=text_encoder, + text_encoder_device=text_encoder_device, + vae=vae, + progress=ComfyProgressReporter(), + ) + + +def _folder_paths() -> ModuleType: + """Import ComfyUI folder paths lazily for schema declaration.""" + + module: Any = importlib.import_module("folder_paths") + if not isinstance(module, ModuleType): + raise TypeError("folder_paths import did not return a module.") + return module diff --git a/simple_syrup/runtime/anima_artifacts.py b/simple_syrup/runtime/anima_artifacts.py index f0f3746..5005a71 100644 --- a/simple_syrup/runtime/anima_artifacts.py +++ b/simple_syrup/runtime/anima_artifacts.py @@ -7,6 +7,7 @@ from __future__ import annotations from .auto_model_artifact import AutoModelArtifact +from .qwen_artifacts import QWEN_IMAGE_VAE ANIMA_QWEN_TEXT_ENCODER = AutoModelArtifact( cache_id="anima_qwen_text_encoder", @@ -20,20 +21,7 @@ ANIMA_QWEN_TEXT_ENCODER = AutoModelArtifact( source_repo="circlestone-labs/Anima", description="Anima Qwen3 0.6B text encoder", sha256="cd2a512003e2f9f3cd3c32a9c3573f820bb28c940f73c57b1ddaa983d9223eba", + file_size_bytes=1_192_135_096, ) -ANIMA_QWEN_VAE = AutoModelArtifact( - cache_id="anima_qwen_vae", - filename="qwen_image_vae.safetensors", - folder_name="vae", - canonical_subfolder="qwen", - source_url=( - "https://huggingface.co/circlestone-labs/Anima/resolve/main/" - "split_files/vae/qwen_image_vae.safetensors" - ), - source_repo="circlestone-labs/Anima", - description="Anima Qwen Image VAE", - sha256="a70580f0213e67967ee9c95f05bb400e8fb08307e017a924bf3441223e023d1f", -) - -ANIMA_AUTO_ARTIFACTS = (ANIMA_QWEN_TEXT_ENCODER, ANIMA_QWEN_VAE) +ANIMA_AUTO_ARTIFACTS = (ANIMA_QWEN_TEXT_ENCODER, QWEN_IMAGE_VAE) diff --git a/simple_syrup/runtime/auto_model_artifact.py b/simple_syrup/runtime/auto_model_artifact.py index e312464..f8301b4 100644 --- a/simple_syrup/runtime/auto_model_artifact.py +++ b/simple_syrup/runtime/auto_model_artifact.py @@ -21,3 +21,4 @@ class AutoModelArtifact: source_repo: str description: str sha256: str + file_size_bytes: int diff --git a/simple_syrup/runtime/auto_model_choices.py b/simple_syrup/runtime/auto_model_choices.py new file mode 100644 index 0000000..202ba11 --- /dev/null +++ b/simple_syrup/runtime/auto_model_choices.py @@ -0,0 +1,72 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Build component choices without duplicating automatic local artifacts.""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from types import ModuleType +from typing import Any + +from .auto_model_artifact import AutoModelArtifact + + +def automatic_component_choices( + installed: Sequence[str], + artifacts: Sequence[AutoModelArtifact], + leading_choices: Sequence[str], + folder_paths_module: ModuleType, +) -> list[str]: + """Return leading choices plus local files not represented automatically.""" + + if not artifacts: + raise ValueError("Automatic component choices require at least one artifact.") + folder_names = {artifact.folder_name for artifact in artifacts} + if len(folder_names) != 1: + raise ValueError("Automatic component artifacts must share one model category.") + + folder_name = next(iter(folder_names)) + automatic_names = {artifact.filename for artifact in artifacts} + automatic_sizes = {artifact.file_size_bytes for artifact in artifacts} + choices = list(dict.fromkeys(leading_choices)) + seen = set(choices) + for choice in installed: + if choice in seen or choice in automatic_names: + continue + if _installed_file_has_known_size( + folder_paths_module, + folder_name, + choice, + automatic_sizes, + ): + continue + choices.append(choice) + seen.add(choice) + return choices + + +def _installed_file_has_known_size( + folder_paths_module: ModuleType, + folder_name: str, + choice: str, + automatic_sizes: set[int], +) -> bool: + """Identify a likely automatic artifact without hashing during schema creation.""" + + get_full_path: Any = getattr(folder_paths_module, "get_full_path", None) + if not callable(get_full_path): + return False + path_value: Any = get_full_path(folder_name, choice) + if path_value is None: + return False + try: + path = Path(str(path_value)) + return path.is_file() and path.stat().st_size in automatic_sizes + except OSError: + return False + + +__all__ = ["automatic_component_choices"] diff --git a/simple_syrup/runtime/auto_model_resolver.py b/simple_syrup/runtime/auto_model_resolver.py index 040fcba..945a357 100644 --- a/simple_syrup/runtime/auto_model_resolver.py +++ b/simple_syrup/runtime/auto_model_resolver.py @@ -150,8 +150,6 @@ class AutoModelResolver: return False if entry.sha256 != artifact.sha256: return False - if entry.path.name != artifact.filename: - return False if not entry.path.is_file(): return False try: @@ -163,6 +161,8 @@ class AutoModelResolver: except ValueError: return False stat = entry.path.stat() + if stat.st_size != artifact.file_size_bytes: + return False if entry.file_size is not None and entry.modified_time_ns is not None: return ( entry.file_size == stat.st_size @@ -175,19 +175,33 @@ def find_model_artifact( artifact: AutoModelArtifact, folder_paths_module: ModuleType | None = None, ) -> Path | None: - """Return the first same-named local file matching the catalog checksum.""" + """Find an artifact by the cheapest reliable checks within its model category.""" _validate_basename(artifact.filename) - for root in get_model_folder_paths(artifact.folder_name, folder_paths_module): + roots = get_model_folder_paths(artifact.folder_name, folder_paths_module) + canonical = canonical_auto_destination(artifact, folder_paths_module) + if _candidate_matches_artifact(canonical, artifact): + return canonical + + checked: set[Path] = {canonical.resolve()} + for root in roots: if not root.is_dir(): continue matches = sorted( path for path in root.rglob(artifact.filename) if path.is_file() ) for match in matches: - if match.name != artifact.filename or not _path_is_under(match, root): + if not _candidate_has_expected_size(match, artifact): continue - if sha256_file(match).lower() == artifact.sha256.lower(): + resolved_match = match.resolve() + if ( + resolved_match in checked + or match.name != artifact.filename + or not _path_is_under(match, root) + ): + continue + checked.add(resolved_match) + if _candidate_checksum_matches(match, artifact): return match LOGGER.warning( "same-named auto model artifact has a different checksum", @@ -197,9 +211,63 @@ def find_model_artifact( "artifact_filename": artifact.filename, }, ) + + for root in roots: + if not root.is_dir(): + continue + for candidate in root.rglob("*"): + if not _candidate_has_expected_size(candidate, artifact): + continue + if not _path_is_under(candidate, root): + continue + resolved_candidate = candidate.resolve() + if resolved_candidate in checked: + continue + checked.add(resolved_candidate) + if not _candidate_checksum_matches(candidate, artifact): + continue + LOGGER.info( + "auto model artifact found under alternate filename", + extra={ + "cache_id": artifact.cache_id, + "path": str(candidate), + "artifact_filename": artifact.filename, + }, + ) + return candidate return None +def _candidate_matches_artifact( + candidate: Path, + artifact: AutoModelArtifact, +) -> bool: + """Verify a candidate only when its inexpensive file checks match first.""" + + return _candidate_has_expected_size( + candidate, + artifact, + ) and _candidate_checksum_matches(candidate, artifact) + + +def _candidate_has_expected_size( + candidate: Path, + artifact: AutoModelArtifact, +) -> bool: + """Return whether a regular file has the artifact's exact byte size.""" + + return candidate.is_file() and candidate.stat().st_size == artifact.file_size_bytes + + +def _candidate_checksum_matches( + candidate: Path, + artifact: AutoModelArtifact, +) -> bool: + """Return whether an already size-matched candidate has the trusted digest.""" + + return sha256_file(candidate).lower() == artifact.sha256.lower() + + def find_model_by_basename( folder_name: str, basename: str, diff --git a/simple_syrup/runtime/clip_type_support.py b/simple_syrup/runtime/clip_type_support.py new file mode 100644 index 0000000..b427e8f --- /dev/null +++ b/simple_syrup/runtime/clip_type_support.py @@ -0,0 +1,25 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Validate ComfyUI text-encoder type availability before model downloads.""" + +from __future__ import annotations + +import importlib +from typing import Any + + +class ComfyClipTypeSupport: + """Validate named ComfyUI CLIP types at the runtime boundary.""" + + def require(self, clip_type_name: str) -> None: + """Raise an actionable error when ComfyUI lacks a required CLIP type.""" + + comfy_sd: Any = importlib.import_module("comfy.sd") + if hasattr(comfy_sd.CLIPType, clip_type_name): + return + raise RuntimeError( + f"This loader requires ComfyUI CLIP type '{clip_type_name}'. " + "Update ComfyUI before using this node." + ) diff --git a/simple_syrup/runtime/diffusion_model_metadata.py b/simple_syrup/runtime/diffusion_model_metadata.py new file mode 100644 index 0000000..633cd57 --- /dev/null +++ b/simple_syrup/runtime/diffusion_model_metadata.py @@ -0,0 +1,65 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Inspect tensor-derived metadata exposed by loaded ComfyUI diffusion models.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + +from ..shared.logging import get_logger + +LOGGER = get_logger(__name__) + + +@dataclass(frozen=True) +class DiffusionModelMetadata: + """Describe architecture fields relevant to component selection.""" + + image_model: str | None + context_input_dimension: int | None + + +@runtime_checkable +class ModelPatcherBoundary(Protocol): + """Expose the loaded model objects required for architecture inspection.""" + + def get_model_object(self, name: str) -> object: + """Return a named object owned by ComfyUI's model patcher.""" + + +class DiffusionModelMetadataInspector: + """Read normalized architecture metadata from a loaded model patcher.""" + + def inspect(self, model: object) -> DiffusionModelMetadata | None: + """Return narrowed model metadata or None when it is unavailable.""" + + if not isinstance(model, ModelPatcherBoundary): + return None + try: + model_config = model.get_model_object("model_config") + except (AttributeError, KeyError, TypeError, ValueError): + LOGGER.warning( + "loaded model does not expose inspectable model configuration" + ) + return None + + unet_config = getattr(model_config, "unet_config", None) + if not isinstance(unet_config, Mapping): + return None + image_model_value = unet_config.get("image_model") + context_dimension_value = unet_config.get("context_in_dim") + return DiffusionModelMetadata( + image_model=( + image_model_value if isinstance(image_model_value, str) else None + ), + context_input_dimension=( + context_dimension_value + if isinstance(context_dimension_value, int) + and not isinstance(context_dimension_value, bool) + else None + ), + ) diff --git a/simple_syrup/runtime/flux_artifacts.py b/simple_syrup/runtime/flux_artifacts.py index b883b10..591fabc 100644 --- a/simple_syrup/runtime/flux_artifacts.py +++ b/simple_syrup/runtime/flux_artifacts.py @@ -21,6 +21,7 @@ FLUX_CLIP_L = AutoModelArtifact( source_repo="comfyanonymous/flux_text_encoders", description="FLUX CLIP-L text encoder", sha256="660c6f5b1abae9dc498ac2d21e1347d2abdb0cf6c0c0c8576cd796491d9a6cdd", + file_size_bytes=246_144_152, ) FLUX_T5_XXL = AutoModelArtifact( @@ -35,6 +36,7 @@ FLUX_T5_XXL = AutoModelArtifact( source_repo="comfyanonymous/flux_text_encoders", description="FLUX T5-XXL FP16 text encoder", sha256="6e480b09fae049a72d2a8c5fbccb8d3e92febeb233bbe9dfe7256958a9167635", + file_size_bytes=9_787_841_024, ) FLUX_VAE = AutoModelArtifact( @@ -50,6 +52,7 @@ FLUX_VAE = AutoModelArtifact( source_repo="Comfy-Org/Lumina_Image_2.0_Repackaged", description="FLUX autoencoder VAE", sha256="afc8e28272cd15db3919bacdb6918ce9c1ed22e96cb12c4d5ed0fba823529e38", + file_size_bytes=335_304_388, ) FLUX2_DEV_TEXT_ENCODER = AutoModelArtifact( @@ -65,6 +68,7 @@ FLUX2_DEV_TEXT_ENCODER = AutoModelArtifact( source_repo="Comfy-Org/flux2-dev", description="FLUX.2 dev Mistral 3 Small text encoder", sha256="7d79902f60b1aeb3a6de2cfad02f4367b5e300a1387de3d03ac717cfa3df117c", + file_size_bytes=35_584_897_447, ) FLUX2_KLEIN_4B_TEXT_ENCODER = AutoModelArtifact( @@ -80,6 +84,7 @@ FLUX2_KLEIN_4B_TEXT_ENCODER = AutoModelArtifact( source_repo="Comfy-Org/vae-text-encorder-for-flux-klein-4b", description="FLUX.2 Klein 4B Qwen3 text encoder", sha256="6c671498573ac2f7a5501502ccce8d2b08ea6ca2f661c458e708f36b36edfc5a", + file_size_bytes=8_044_982_048, ) FLUX2_KLEIN_9B_TEXT_ENCODER = AutoModelArtifact( @@ -95,6 +100,7 @@ FLUX2_KLEIN_9B_TEXT_ENCODER = AutoModelArtifact( source_repo="Comfy-Org/vae-text-encorder-for-flux-klein-9b", description="FLUX.2 Klein 9B Qwen3 8B FP8-mixed text encoder", sha256="abad16806e0cbabc54e0325d6565847443fe396d5f0be38bb3cd3fe75a1201d6", + file_size_bytes=8_664_848_742, ) FLUX2_VAE = AutoModelArtifact( @@ -110,6 +116,7 @@ FLUX2_VAE = AutoModelArtifact( source_repo="Comfy-Org/flux2-dev", description="FLUX.2 VAE", sha256="d64f3a68e1cc4f9f4e29b6e0da38a0204fe9a49f2d4053f0ec1fa1ca02f9c4b5", + file_size_bytes=336_213_556, ) FLUX2_TEXT_ENCODERS: dict[Flux2TextEncoderProfile, AutoModelArtifact] = { diff --git a/simple_syrup/runtime/flux_model_inspector.py b/simple_syrup/runtime/flux_model_inspector.py index 03928da..8f0e36a 100644 --- a/simple_syrup/runtime/flux_model_inspector.py +++ b/simple_syrup/runtime/flux_model_inspector.py @@ -6,49 +6,42 @@ from __future__ import annotations -from collections.abc import Mapping -from typing import Protocol, runtime_checkable +from typing import Protocol from ..domain.flux_profiles import FluxModelProfile, classify_flux_profile -from ..shared.logging import get_logger - -LOGGER = get_logger(__name__) - - -@runtime_checkable -class ModelPatcherBoundary(Protocol): - """Expose the loaded model objects required for architecture inspection.""" - - def get_model_object(self, name: str) -> object: - """Return a named object owned by ComfyUI's model patcher.""" +from .diffusion_model_metadata import ( + DiffusionModelMetadata, + DiffusionModelMetadataInspector, +) class FluxModelInspector: """Read ComfyUI's tensor-derived model configuration after model loading.""" + def __init__( + self, + metadata_inspector: DiffusionModelMetadataInspectorBoundary | None = None, + ) -> None: + """Create a FLUX classifier over shared metadata inspection.""" + + self._metadata_inspector = ( + metadata_inspector or DiffusionModelMetadataInspector() + ) + def inspect(self, model: object) -> FluxModelProfile | None: """Return a detected FLUX profile or None for unavailable metadata.""" - if not isinstance(model, ModelPatcherBoundary): + metadata = self._metadata_inspector.inspect(model) + if metadata is None: return None - try: - model_config = model.get_model_object("model_config") - except (AttributeError, KeyError, TypeError, ValueError): - LOGGER.warning( - "loaded model does not expose inspectable model configuration" - ) - return None - - unet_config = getattr(model_config, "unet_config", None) - if not isinstance(unet_config, Mapping): - return None - image_model_value = unet_config.get("image_model") - context_dimension_value = unet_config.get("context_in_dim") - image_model = image_model_value if isinstance(image_model_value, str) else None - context_dimension = ( - context_dimension_value - if isinstance(context_dimension_value, int) - and not isinstance(context_dimension_value, bool) - else None + return classify_flux_profile( + metadata.image_model, + metadata.context_input_dimension, ) - return classify_flux_profile(image_model, context_dimension) + + +class DiffusionModelMetadataInspectorBoundary(Protocol): + """Expose normalized loaded diffusion-model metadata.""" + + def inspect(self, model: object) -> DiffusionModelMetadata | None: + """Return normalized metadata when available.""" diff --git a/simple_syrup/runtime/krea2_artifacts.py b/simple_syrup/runtime/krea2_artifacts.py new file mode 100644 index 0000000..325f745 --- /dev/null +++ b/simple_syrup/runtime/krea2_artifacts.py @@ -0,0 +1,56 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Declare trusted automatic artifacts and selections for Krea 2.""" + +from __future__ import annotations + +from .auto_model_artifact import AutoModelArtifact + +KREA2_AUTO_TEXT_ENCODER = "auto" + +KREA2_QWEN3_VL_4B_FP8 = AutoModelArtifact( + cache_id="krea2_qwen3vl_4b_fp8_scaled", + filename="qwen3vl_4b_fp8_scaled.safetensors", + folder_name="text_encoders", + canonical_subfolder="krea2", + source_url=( + "https://huggingface.co/Comfy-Org/Krea-2/resolve/" + "e5ea8b4dd7f38f348b138eb0fe29f92c0e367e96/text_encoders/" + "qwen3vl_4b_fp8_scaled.safetensors" + ), + source_repo="Comfy-Org/Krea-2", + description="Krea 2 Qwen3-VL 4B FP8-scaled text encoder", + sha256="54bd5144df0bbc25dd6ccadfcb826b521445a1b06ae5a42570bdd2974ca87094", + file_size_bytes=5_242_467_968, +) + +KREA2_QWEN3_VL_4B_BF16 = AutoModelArtifact( + cache_id="krea2_qwen3vl_4b_bf16", + filename="qwen3vl_4b_bf16.safetensors", + folder_name="text_encoders", + canonical_subfolder="krea2", + source_url=( + "https://huggingface.co/Comfy-Org/Krea-2/resolve/" + "e5ea8b4dd7f38f348b138eb0fe29f92c0e367e96/text_encoders/" + "qwen3vl_4b_bf16.safetensors" + ), + source_repo="Comfy-Org/Krea-2", + description="Krea 2 Qwen3-VL 4B BF16 text encoder", + sha256="36f3ff447ef59201722e8f9ce6020c9819fdcfba6aa2608c4e09b1c0ce114e34", + file_size_bytes=8_875_719_384, +) + +KREA2_TEXT_ENCODER_ARTIFACTS = { + KREA2_QWEN3_VL_4B_FP8.filename: KREA2_QWEN3_VL_4B_FP8, + KREA2_QWEN3_VL_4B_BF16.filename: KREA2_QWEN3_VL_4B_BF16, +} + + +__all__ = [ + "KREA2_AUTO_TEXT_ENCODER", + "KREA2_QWEN3_VL_4B_BF16", + "KREA2_QWEN3_VL_4B_FP8", + "KREA2_TEXT_ENCODER_ARTIFACTS", +] diff --git a/simple_syrup/runtime/qwen_artifacts.py b/simple_syrup/runtime/qwen_artifacts.py new file mode 100644 index 0000000..2ddc6d3 --- /dev/null +++ b/simple_syrup/runtime/qwen_artifacts.py @@ -0,0 +1,27 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Declare shared checksum-pinned Qwen model artifacts.""" + +from __future__ import annotations + +from .auto_model_artifact import AutoModelArtifact + +QWEN_IMAGE_VAE = AutoModelArtifact( + cache_id="qwen_image_vae", + filename="qwen_image_vae.safetensors", + folder_name="vae", + canonical_subfolder="qwen", + source_url=( + "https://huggingface.co/Comfy-Org/Krea-2/resolve/" + "e5ea8b4dd7f38f348b138eb0fe29f92c0e367e96/vae/" + "qwen_image_vae.safetensors" + ), + source_repo="Comfy-Org/Krea-2", + description="Qwen Image VAE", + sha256="a70580f0213e67967ee9c95f05bb400e8fb08307e017a924bf3441223e023d1f", + file_size_bytes=253_806_246, +) + +__all__ = ["QWEN_IMAGE_VAE"] diff --git a/simple_syrup/services/anima_loader_service.py b/simple_syrup/services/anima_loader_service.py index decf687..2f4afa3 100644 --- a/simple_syrup/services/anima_loader_service.py +++ b/simple_syrup/services/anima_loader_service.py @@ -13,11 +13,12 @@ from typing import Any, Protocol import torch -from ..runtime.anima_artifacts import ANIMA_QWEN_TEXT_ENCODER, ANIMA_QWEN_VAE +from ..runtime.anima_artifacts import ANIMA_QWEN_TEXT_ENCODER from ..runtime.auto_model_artifact import AutoModelArtifact from ..runtime.auto_model_resolver import AutoModelResolution, AutoModelResolver from ..runtime.model_downloads import ProgressReporter from ..runtime.quantization_progress import QuantizationProgressReporter +from ..runtime.qwen_artifacts import QWEN_IMAGE_VAE from ..runtime.vae_loader import VaeLoaderService, load_vae_path from .anima_diffusion_model_service import AnimaDiffusionModelService @@ -125,7 +126,7 @@ class AnimaLoaderService: """Load a VAE using ComfyUI's VAE loader policy.""" if vae == AUTO_CHOICE: - vae_path = self._resolver.resolve(ANIMA_QWEN_VAE, progress).path + vae_path = self._resolver.resolve(QWEN_IMAGE_VAE, progress).path return load_vae_path(vae_path) return self._vae_loader.load_vae(vae) diff --git a/simple_syrup/services/krea2_loader_service.py b/simple_syrup/services/krea2_loader_service.py new file mode 100644 index 0000000..90fb2ea --- /dev/null +++ b/simple_syrup/services/krea2_loader_service.py @@ -0,0 +1,196 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Load and validate the diffusion, text encoder, and VAE for Krea 2.""" + +from __future__ import annotations + +import importlib +from pathlib import Path +from types import ModuleType +from typing import Any, Protocol + +from ..runtime.auto_model_artifact import AutoModelArtifact +from ..runtime.auto_model_resolver import AutoModelResolution, AutoModelResolver +from ..runtime.clip_type_support import ComfyClipTypeSupport +from ..runtime.diffusion_model_loader import DiffusionModelLoader +from ..runtime.diffusion_model_metadata import ( + DiffusionModelMetadata, + DiffusionModelMetadataInspector, +) +from ..runtime.krea2_artifacts import ( + KREA2_AUTO_TEXT_ENCODER, + KREA2_QWEN3_VL_4B_FP8, + KREA2_TEXT_ENCODER_ARTIFACTS, +) +from ..runtime.model_downloads import ProgressReporter +from ..runtime.qwen_artifacts import QWEN_IMAGE_VAE +from ..runtime.text_encoder_loader import TextEncoderLoader +from ..runtime.vae_loader import VaeLoaderService, load_vae_path + +AUTO_CHOICE = "auto" +KREA2_CLIP_TYPE = "KREA2" +KREA2_IMAGE_MODEL = "krea2" + + +class Krea2LoaderService: + """Orchestrate structurally validated Krea 2 component loading.""" + + def __init__( + self, + diffusion_loader: DiffusionModelLoaderBoundary | None = None, + text_encoder_loader: TextEncoderLoaderBoundary | None = None, + model_inspector: DiffusionModelInspectorBoundary | None = None, + clip_type_support: ClipTypeSupportBoundary | None = None, + resolver: AutoModelResolverBoundary | None = None, + vae_loader: VaeLoaderBoundary | None = None, + folder_paths_module: ModuleType | None = None, + ) -> None: + """Create a loader with injectable host and artifact boundaries.""" + + self._folder_paths_module = folder_paths_module + self._diffusion_loader = diffusion_loader or DiffusionModelLoader( + folder_paths_module + ) + self._text_encoder_loader = text_encoder_loader or TextEncoderLoader( + folder_paths_module + ) + self._model_inspector = model_inspector or DiffusionModelMetadataInspector() + self._clip_type_support = clip_type_support or ComfyClipTypeSupport() + self._resolver = resolver or AutoModelResolver( + folder_paths_module=folder_paths_module + ) + self._vae_loader = vae_loader or VaeLoaderService(folder_paths_module) + + def load_models( + self, + diffusion_model: str, + diffusion_weight_dtype: str, + text_encoder: str, + text_encoder_device: str, + vae: str, + progress: ProgressReporter | None = None, + ) -> tuple[object, object, object]: + """Return a validated Krea 2 MODEL, CLIP, and VAE tuple.""" + + model = self._diffusion_loader.load( + diffusion_model, + diffusion_weight_dtype, + ) + self._require_krea2_model(model) + self._clip_type_support.require(KREA2_CLIP_TYPE) + + encoder_path = self._resolve_text_encoder(text_encoder, progress) + clip = self._text_encoder_loader.load( + (encoder_path,), + KREA2_CLIP_TYPE, + text_encoder_device, + ) + loaded_vae = self._load_vae(vae, progress) + return model, clip, loaded_vae + + def _require_krea2_model(self, model: object) -> None: + """Reject non-Krea architectures before resolving large support files.""" + + metadata = self._model_inspector.inspect(model) + if metadata is not None and metadata.image_model == KREA2_IMAGE_MODEL: + return + raise ValueError( + "Simple Load Krea 2 requires a diffusion model ComfyUI recognizes " + "as Krea 2. Select a Krea 2 Raw or Turbo diffusion model." + ) + + def _resolve_text_encoder( + self, + selection: str, + progress: ProgressReporter | None, + ) -> Path: + """Resolve auto and official selections or a local manual encoder.""" + + artifact = ( + KREA2_QWEN3_VL_4B_FP8 + if selection == KREA2_AUTO_TEXT_ENCODER + else KREA2_TEXT_ENCODER_ARTIFACTS.get(selection) + ) + if artifact is not None: + return self._resolver.resolve(artifact, progress).path + path = self._folder_paths().get_full_path_or_raise( + "text_encoders", + selection, + ) + return Path(str(path)) + + def _load_vae( + self, + selection: str, + progress: ProgressReporter | None, + ) -> object: + """Load the shared Qwen Image VAE automatically or a manual VAE.""" + + if selection == AUTO_CHOICE: + path = self._resolver.resolve(QWEN_IMAGE_VAE, progress).path + return load_vae_path(path) + return self._vae_loader.load_vae(selection) + + def _folder_paths(self) -> ModuleType: + """Return the configured ComfyUI folder-path registry.""" + + if self._folder_paths_module is not None: + return self._folder_paths_module + module: Any = importlib.import_module("folder_paths") + if not isinstance(module, ModuleType): + raise TypeError("folder_paths import did not return a module.") + self._folder_paths_module = module + return module + + +class DiffusionModelLoaderBoundary(Protocol): + """Load one selected standalone diffusion model.""" + + def load(self, diffusion_model: str, weight_dtype: str) -> object: + """Return a loaded ComfyUI model patcher.""" + + +class TextEncoderLoaderBoundary(Protocol): + """Load one or more text-encoder files for a named Comfy CLIP type.""" + + def load( + self, + paths: tuple[Path, ...], + clip_type_name: str, + device: str, + ) -> object: + """Return a loaded ComfyUI CLIP object.""" + + +class DiffusionModelInspectorBoundary(Protocol): + """Inspect architecture metadata from a loaded diffusion model.""" + + def inspect(self, model: object) -> DiffusionModelMetadata | None: + """Return normalized metadata when ComfyUI exposes it.""" + + +class ClipTypeSupportBoundary(Protocol): + """Validate that the host supports a required Comfy CLIP type.""" + + def require(self, clip_type_name: str) -> None: + """Raise when a required CLIP type is unavailable.""" + + +class AutoModelResolverBoundary(Protocol): + """Resolve a trusted catalog artifact to a verified local file.""" + + def resolve( + self, + artifact: AutoModelArtifact, + progress: ProgressReporter | None = None, + ) -> AutoModelResolution: + """Return a verified local artifact path.""" + + +class VaeLoaderBoundary(Protocol): + """Load one manually selected ComfyUI VAE.""" + + def load_vae(self, vae_name: str) -> object: + """Return a loaded ComfyUI VAE object.""" diff --git a/tests/test_anima_loader.py b/tests/test_anima_loader.py index 06078c4..c57e684 100644 --- a/tests/test_anima_loader.py +++ b/tests/test_anima_loader.py @@ -232,7 +232,7 @@ def test_loader_uses_auto_resolver_for_auto_choices( progress, ) - assert resolver.requests == ["anima_qwen_text_encoder", "anima_qwen_vae"] + assert resolver.requests == ["anima_qwen_text_encoder", "qwen_image_vae"] assert resolver.progress_reporters == [progress, progress] assert comfy_state.clip_calls[0]["ckpt_paths"] == [str(resolver.text_encoder_path)] assert comfy_state.vae_paths == [str(resolver.vae_path)] @@ -265,7 +265,7 @@ def test_anima_auto_downloads_emit_comfy_node_progress_end_to_end( "ANIMA_QWEN_TEXT_ENCODER", text_artifact, ) - monkeypatch.setattr(anima_loader_module, "ANIMA_QWEN_VAE", vae_artifact) + monkeypatch.setattr(anima_loader_module, "QWEN_IMAGE_VAE", vae_artifact) content_by_url = { text_artifact.source_url: text_content, @@ -494,4 +494,5 @@ def _small_artifact( source_repo="example/progress", description=f"progress test {filename}", sha256=hashlib.sha256(content).hexdigest(), + file_size_bytes=len(content), ) diff --git a/tests/test_auto_model_choices.py b/tests/test_auto_model_choices.py new file mode 100644 index 0000000..e645ec5 --- /dev/null +++ b/tests/test_auto_model_choices.py @@ -0,0 +1,112 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Tests for artifact-aware automatic component dropdown choices.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path +from types import ModuleType + +import pytest + +from simple_syrup.runtime.auto_model_artifact import AutoModelArtifact +from simple_syrup.runtime.auto_model_choices import automatic_component_choices + + +class _FolderPaths(ModuleType): + """Resolve ComfyUI-relative test model choices from one models root.""" + + def __init__(self, models_root: Path) -> None: + """Create a model-path fake rooted at the temporary directory.""" + + super().__init__("folder_paths") + self._models_root = models_root + + def get_full_path(self, folder_name: str, choice: str) -> str | None: + """Return an existing test model path or None.""" + + path = self._models_root / folder_name / choice + return str(path) if path.is_file() else None + + +def test_choices_hide_official_and_renamed_automatic_files(tmp_path: Path) -> None: + """Known local artifacts appear only through their automatic leading choice.""" + + folder_paths = _FolderPaths(tmp_path / "models") + artifact = _artifact("official.safetensors", b"trusted") + folder = tmp_path / "models" / "text_encoders" + folder.mkdir(parents=True) + (folder / "renamed.safetensors").write_bytes(b"trusted") + (folder / "manual.safetensors").write_bytes(b"manual model") + + choices = automatic_component_choices( + installed=( + artifact.filename, + "renamed.safetensors", + "manual.safetensors", + ), + artifacts=(artifact,), + leading_choices=("auto", artifact.filename), + folder_paths_module=folder_paths, + ) + + assert choices == ["auto", artifact.filename, "manual.safetensors"] + + +def test_choices_keep_same_category_files_with_different_sizes(tmp_path: Path) -> None: + """Manual files remain selectable when they cannot be the automatic artifact.""" + + folder_paths = _FolderPaths(tmp_path / "models") + artifact = _artifact("official.safetensors", b"trusted") + folder = tmp_path / "models" / "text_encoders" + folder.mkdir(parents=True) + (folder / "custom.safetensors").write_bytes(b"different size") + + choices = automatic_component_choices( + installed=("custom.safetensors",), + artifacts=(artifact,), + leading_choices=("auto",), + folder_paths_module=folder_paths, + ) + + assert choices == ["auto", "custom.safetensors"] + + +def test_choices_require_artifacts_from_one_model_category() -> None: + """Dropdown filtering cannot accidentally combine unrelated model categories.""" + + folder_paths = _FolderPaths(Path("models")) + + with pytest.raises(ValueError, match="share one model category"): + automatic_component_choices( + installed=(), + artifacts=( + _artifact("encoder.safetensors", b"encoder", "text_encoders"), + _artifact("vae.safetensors", b"vae", "vae"), + ), + leading_choices=("auto",), + folder_paths_module=folder_paths, + ) + + +def _artifact( + filename: str, + content: bytes, + folder_name: str = "text_encoders", +) -> AutoModelArtifact: + """Create a small trusted artifact for dropdown tests.""" + + return AutoModelArtifact( + cache_id=filename, + filename=filename, + folder_name=folder_name, + canonical_subfolder="test", + source_url=f"https://example.invalid/{filename}", + source_repo="example/models", + description=filename, + sha256=hashlib.sha256(content).hexdigest(), + file_size_bytes=len(content), + ) diff --git a/tests/test_auto_model_resolver.py b/tests/test_auto_model_resolver.py index 2e3e72e..2fa0a03 100644 --- a/tests/test_auto_model_resolver.py +++ b/tests/test_auto_model_resolver.py @@ -16,6 +16,7 @@ from simple_syrup.runtime.auto_model_cache import AutoModelCache, AutoModelCache from simple_syrup.runtime.auto_model_resolver import ( AutoModelResolver, canonical_auto_destination, + find_model_artifact, find_model_by_basename, relative_model_name, ) @@ -158,6 +159,112 @@ def test_resolver_ignores_same_named_file_with_wrong_checksum(tmp_path: Path) -> assert len(downloader.requests) == 1 +def test_resolver_finds_renamed_artifact_by_size_and_checksum(tmp_path: Path) -> None: + """A renamed official artifact is reused from its registered model category.""" + + fake = FakeFolderPaths(tmp_path / "models") + artifact = _artifact("text_encoders", "model.safetensors") + renamed_path = ( + tmp_path / "models" / "text_encoders" / "custom" / "my-qwen.safetensors" + ) + renamed_path.parent.mkdir(parents=True) + renamed_path.write_bytes(b"model") + cache = AutoModelCache(fake) + downloader = RecordingDownloader() + + resolved = AutoModelResolver(cache, downloader, fake).resolve(artifact) + cached = AutoModelResolver(cache, downloader, fake).resolve(artifact) + + assert resolved.path == renamed_path + assert resolved.source == "found" + assert cached.path == renamed_path + assert cached.source == "cached" + assert cache.load()[artifact.cache_id].path == renamed_path + assert downloader.requests == [] + + +def test_resolver_does_not_scan_unrelated_model_categories(tmp_path: Path) -> None: + """Checksum discovery stays inside the artifact's registered category.""" + + fake = FakeFolderPaths(tmp_path / "models") + artifact = _artifact("text_encoders", "model.safetensors") + unrelated_path = tmp_path / "models" / "vae" / "renamed.safetensors" + unrelated_path.parent.mkdir(parents=True) + unrelated_path.write_bytes(b"model") + downloader = RecordingDownloader() + + resolved = AutoModelResolver( + AutoModelCache(fake), + downloader, + fake, + ).resolve(artifact) + + assert resolved.source == "downloaded" + assert resolved.path != unrelated_path + assert len(downloader.requests) == 1 + + +def test_resolver_prefers_canonical_filename_before_renamed_match( + tmp_path: Path, +) -> None: + """The direct canonical lookup wins before the broader size-based scan.""" + + fake = FakeFolderPaths(tmp_path / "models") + artifact = _artifact("text_encoders", "model.safetensors") + canonical = canonical_auto_destination(artifact, fake) + canonical.parent.mkdir(parents=True) + canonical.write_bytes(b"model") + renamed = tmp_path / "models" / "text_encoders" / "renamed.safetensors" + renamed.write_bytes(b"model") + + assert find_model_artifact(artifact, fake) == canonical + + +def test_resolver_prefers_recursive_official_filename_before_renamed_match( + tmp_path: Path, +) -> None: + """Official filenames are searched recursively before alternate names.""" + + fake = FakeFolderPaths(tmp_path / "models") + artifact = _artifact("text_encoders", "model.safetensors") + root = tmp_path / "models" / "text_encoders" + renamed = root / "a-renamed.safetensors" + official = root / "nested" / artifact.filename + official.parent.mkdir(parents=True) + renamed.write_bytes(b"model") + official.write_bytes(b"model") + + assert find_model_artifact(artifact, fake) == official + + +def test_artifact_discovery_hashes_only_size_matches( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The broad category scan avoids hashing files with a different byte size.""" + + fake = FakeFolderPaths(tmp_path / "models") + artifact = _artifact("text_encoders", "model.safetensors") + wrong_size = tmp_path / "models" / "text_encoders" / "other.safetensors" + wrong_size.parent.mkdir(parents=True) + wrong_size.write_bytes(b"different size") + hashed_paths: list[Path] = [] + + def record_hash(path: Path) -> str: + """Record unexpected hashing while retaining a valid callable shape.""" + + hashed_paths.append(path) + return artifact.sha256 + + monkeypatch.setattr( + "simple_syrup.runtime.auto_model_resolver.sha256_file", + record_hash, + ) + + assert find_model_artifact(artifact, fake) is None + assert hashed_paths == [] + + def test_find_model_by_basename_respects_folder_priority(tmp_path: Path) -> None: """Recursive search prefers earlier ComfyUI model roots.""" @@ -212,6 +319,7 @@ def test_canonical_destination_rejects_unsafe_subfolder(tmp_path: Path) -> None: source_repo="example/model", description="bad", sha256="abc", + file_size_bytes=1, ) with pytest.raises(ValueError, match="not safe"): @@ -274,4 +382,5 @@ def _artifact(folder_name: str, filename: str) -> AutoModelArtifact: source_repo="example/model", description=f"test {filename}", sha256=hashlib.sha256(b"model").hexdigest(), + file_size_bytes=len(b"model"), ) diff --git a/tests/test_clip_type_support.py b/tests/test_clip_type_support.py new file mode 100644 index 0000000..b37451b --- /dev/null +++ b/tests/test_clip_type_support.py @@ -0,0 +1,39 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Tests for ComfyUI CLIP-type compatibility validation.""" + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from simple_syrup.runtime.clip_type_support import ComfyClipTypeSupport + + +def test_clip_type_support_accepts_installed_krea2( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The current Comfy KREA2 enum satisfies loader preflight.""" + + comfy_sd = ModuleType("comfy.sd") + comfy_sd.CLIPType = SimpleNamespace(KREA2=object()) # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "comfy.sd", comfy_sd) + + ComfyClipTypeSupport().require("KREA2") + + +def test_clip_type_support_reports_actionable_update_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Older Comfy builds fail explicitly before any artifact download.""" + + comfy_sd = ModuleType("comfy.sd") + comfy_sd.CLIPType = SimpleNamespace() # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "comfy.sd", comfy_sd) + + with pytest.raises(RuntimeError, match="Update ComfyUI"): + ComfyClipTypeSupport().require("KREA2") diff --git a/tests/test_diffusion_model_metadata.py b/tests/test_diffusion_model_metadata.py new file mode 100644 index 0000000..2b64331 --- /dev/null +++ b/tests/test_diffusion_model_metadata.py @@ -0,0 +1,55 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Tests for shared loaded diffusion-model metadata inspection.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from simple_syrup.runtime.diffusion_model_metadata import ( + DiffusionModelMetadata, + DiffusionModelMetadataInspector, +) + + +@dataclass +class _ModelConfig: + """Expose one fake tensor-derived UNet configuration.""" + + unet_config: object + + +class _ModelPatcher: + """Expose a model config through ComfyUI's model-patcher surface.""" + + def __init__(self, unet_config: object) -> None: + """Store the fake UNet configuration.""" + + self._config = _ModelConfig(unet_config) + + def get_model_object(self, name: str) -> object: + """Return only the requested model configuration.""" + + assert name == "model_config" + return self._config + + +def test_metadata_inspector_normalizes_krea_architecture() -> None: + """Krea detection uses loaded tensor metadata rather than its filename.""" + + result = DiffusionModelMetadataInspector().inspect( + _ModelPatcher({"image_model": "krea2", "context_in_dim": 30_720}) + ) + + assert result == DiffusionModelMetadata("krea2", 30_720) + + +def test_metadata_inspector_rejects_unavailable_and_malformed_surfaces() -> None: + """Dynamic host values fail closed when structural metadata is unavailable.""" + + inspector = DiffusionModelMetadataInspector() + + assert inspector.inspect(object()) is None + assert inspector.inspect(_ModelPatcher("not a mapping")) is None diff --git a/tests/test_flux_runtime_loaders.py b/tests/test_flux_runtime_loaders.py index ceceef6..8cd7e3e 100644 --- a/tests/test_flux_runtime_loaders.py +++ b/tests/test_flux_runtime_loaders.py @@ -290,4 +290,5 @@ def _artifact() -> AutoModelArtifact: source_repo="example/auto", description="test encoder", sha256="0" * 64, + file_size_bytes=0, ) diff --git a/tests/test_krea2_artifacts.py b/tests/test_krea2_artifacts.py new file mode 100644 index 0000000..5deaa93 --- /dev/null +++ b/tests/test_krea2_artifacts.py @@ -0,0 +1,39 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Tests for trusted Krea 2 component metadata and choices.""" + +from __future__ import annotations + +from simple_syrup.runtime.krea2_artifacts import ( + KREA2_QWEN3_VL_4B_BF16, + KREA2_QWEN3_VL_4B_FP8, + KREA2_TEXT_ENCODER_ARTIFACTS, +) +from simple_syrup.runtime.qwen_artifacts import QWEN_IMAGE_VAE + + +def test_krea2_artifacts_are_checksum_and_revision_pinned() -> None: + """Every downloadable component uses an immutable URL and exact digest.""" + + artifacts = ( + KREA2_QWEN3_VL_4B_FP8, + KREA2_QWEN3_VL_4B_BF16, + QWEN_IMAGE_VAE, + ) + + for artifact in artifacts: + assert "/resolve/main/" not in artifact.source_url + assert "e5ea8b4dd7f38f348b138eb0fe29f92c0e367e96" in artifact.source_url + assert len(artifact.sha256) == 64 + assert set(artifact.sha256) <= set("0123456789abcdef") + + +def test_krea2_official_encoder_selections_map_to_downloadable_artifacts() -> None: + """Both public precision choices resolve to their exact official files.""" + + assert KREA2_TEXT_ENCODER_ARTIFACTS == { + "qwen3vl_4b_fp8_scaled.safetensors": KREA2_QWEN3_VL_4B_FP8, + "qwen3vl_4b_bf16.safetensors": KREA2_QWEN3_VL_4B_BF16, + } diff --git a/tests/test_krea2_loader_service.py b/tests/test_krea2_loader_service.py new file mode 100644 index 0000000..6638800 --- /dev/null +++ b/tests/test_krea2_loader_service.py @@ -0,0 +1,313 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Tests for Krea 2 component selection and runtime validation.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from types import ModuleType + +import pytest + +import simple_syrup.services.krea2_loader_service as service_module +from simple_syrup.runtime.auto_model_artifact import AutoModelArtifact +from simple_syrup.runtime.auto_model_resolver import AutoModelResolution +from simple_syrup.runtime.diffusion_model_metadata import DiffusionModelMetadata +from simple_syrup.runtime.krea2_artifacts import ( + KREA2_AUTO_TEXT_ENCODER, + KREA2_QWEN3_VL_4B_BF16, + KREA2_QWEN3_VL_4B_FP8, +) +from simple_syrup.runtime.model_downloads import ProgressReporter +from simple_syrup.runtime.qwen_artifacts import QWEN_IMAGE_VAE +from simple_syrup.services.krea2_loader_service import Krea2LoaderService + + +class _FolderPaths(ModuleType): + """Resolve deterministic manual text-encoder files.""" + + def __init__(self, root: Path) -> None: + """Store one fake models root.""" + + super().__init__("folder_paths") + self._root = root + + def get_full_path_or_raise(self, folder: str, name: str) -> str: + """Return a deterministic path for a manual selection.""" + + return str(self._root / folder / name) + + +class _DiffusionLoader: + """Record diffusion requests and return a stable model object.""" + + def __init__(self) -> None: + """Create an empty request log.""" + + self.calls: list[tuple[str, str]] = [] + self.model = object() + + def load(self, diffusion_model: str, weight_dtype: str) -> object: + """Record one model load.""" + + self.calls.append((diffusion_model, weight_dtype)) + return self.model + + +class _Inspector: + """Return configured loaded-model metadata.""" + + def __init__(self, image_model: str | None) -> None: + """Store the image-model identifier to expose.""" + + self.image_model = image_model + self.calls: list[object] = [] + + def inspect(self, model: object) -> DiffusionModelMetadata | None: + """Record inspection and return configured metadata.""" + + self.calls.append(model) + if self.image_model is None: + return None + return DiffusionModelMetadata(self.image_model, None) + + +class _ClipTypeSupport: + """Record or reject required Comfy CLIP types.""" + + def __init__(self, error: Exception | None = None) -> None: + """Configure an optional compatibility failure.""" + + self.error = error + self.calls: list[str] = [] + + def require(self, clip_type_name: str) -> None: + """Record the type and raise the configured failure.""" + + self.calls.append(clip_type_name) + if self.error is not None: + raise self.error + + +@dataclass +class _Resolver: + """Record automatic artifact resolutions.""" + + root: Path + calls: list[tuple[AutoModelArtifact, ProgressReporter | None]] = field( + default_factory=list + ) + + def resolve( + self, + artifact: AutoModelArtifact, + progress: ProgressReporter | None = None, + ) -> AutoModelResolution: + """Return a deterministic resolved path.""" + + self.calls.append((artifact, progress)) + return AutoModelResolution(self.root / artifact.filename, "downloaded") + + +@dataclass +class _TextEncoderLoader: + """Record Krea text-encoder adapter calls.""" + + calls: list[tuple[tuple[Path, ...], str, str]] = field(default_factory=list) + + def load( + self, + paths: tuple[Path, ...], + clip_type_name: str, + device: str, + ) -> object: + """Record and return a stable CLIP value.""" + + self.calls.append((paths, clip_type_name, device)) + return "clip" + + +@dataclass +class _VaeLoader: + """Record manual VAE selections.""" + + calls: list[str] = field(default_factory=list) + + def load_vae(self, vae_name: str) -> object: + """Record and return a stable manual VAE value.""" + + self.calls.append(vae_name) + return "manual vae" + + +class _Progress: + """Provide a progress object whose identity is asserted.""" + + def start(self, label: str, total: int | None) -> None: + """Accept a start update.""" + + def advance(self, current: int, total: int | None) -> None: + """Accept an absolute update.""" + + def finish(self) -> None: + """Accept completion.""" + + +@pytest.mark.parametrize( + ("selection", "artifact"), + ( + (KREA2_AUTO_TEXT_ENCODER, KREA2_QWEN3_VL_4B_FP8), + (KREA2_QWEN3_VL_4B_FP8.filename, KREA2_QWEN3_VL_4B_FP8), + (KREA2_QWEN3_VL_4B_BF16.filename, KREA2_QWEN3_VL_4B_BF16), + ), +) +def test_official_encoder_choices_resolve_downloadable_artifacts( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + selection: str, + artifact: AutoModelArtifact, +) -> None: + """Auto, explicit FP8, and explicit BF16 share verified download behavior.""" + + diffusion = _DiffusionLoader() + resolver = _Resolver(tmp_path / "resolved") + text_loader = _TextEncoderLoader() + support = _ClipTypeSupport() + progress = _Progress() + loaded_vae_paths: list[Path] = [] + + def load_automatic_vae(path: Path) -> object: + """Record the resolved VAE path and return a stable object.""" + + loaded_vae_paths.append(path) + return "auto vae" + + monkeypatch.setattr( + service_module, + "load_vae_path", + load_automatic_vae, + ) + service = Krea2LoaderService( + diffusion_loader=diffusion, + text_encoder_loader=text_loader, + model_inspector=_Inspector("krea2"), + clip_type_support=support, + resolver=resolver, + vae_loader=_VaeLoader(), + folder_paths_module=_FolderPaths(tmp_path), + ) + + result = service.load_models( + "renamed-model.safetensors", + "default", + selection, + "cpu", + "auto", + progress, + ) + + assert result == (diffusion.model, "clip", "auto vae") + assert resolver.calls == [(artifact, progress), (QWEN_IMAGE_VAE, progress)] + assert text_loader.calls == [ + ((tmp_path / "resolved" / artifact.filename,), "KREA2", "cpu") + ] + assert loaded_vae_paths == [tmp_path / "resolved" / QWEN_IMAGE_VAE.filename] + assert support.calls == ["KREA2"] + + +def test_manual_components_use_comfy_folder_and_vae_boundaries(tmp_path: Path) -> None: + """Manual escape hatches retain Krea validation and CLIP loading policy.""" + + resolver = _Resolver(tmp_path / "resolved") + text_loader = _TextEncoderLoader() + vae_loader = _VaeLoader() + service = Krea2LoaderService( + diffusion_loader=_DiffusionLoader(), + text_encoder_loader=text_loader, + model_inspector=_Inspector("krea2"), + clip_type_support=_ClipTypeSupport(), + resolver=resolver, + vae_loader=vae_loader, + folder_paths_module=_FolderPaths(tmp_path), + ) + + result = service.load_models( + "krea.safetensors", + "fp8_e4m3fn_fast", + "custom_qwen.safetensors", + "default", + "custom_vae.safetensors", + ) + + assert result[1:] == ("clip", "manual vae") + assert resolver.calls == [] + assert text_loader.calls == [ + ( + (tmp_path / "text_encoders" / "custom_qwen.safetensors",), + "KREA2", + "default", + ) + ] + assert vae_loader.calls == ["custom_vae.safetensors"] + + +@pytest.mark.parametrize("image_model", [None, "flux", "flux2", "anima"]) +def test_non_krea_models_fail_before_compatibility_checks_or_downloads( + tmp_path: Path, + image_model: str | None, +) -> None: + """Structural validation prevents expensive support downloads for wrong models.""" + + resolver = _Resolver(tmp_path) + support = _ClipTypeSupport() + service = Krea2LoaderService( + diffusion_loader=_DiffusionLoader(), + text_encoder_loader=_TextEncoderLoader(), + model_inspector=_Inspector(image_model), + clip_type_support=support, + resolver=resolver, + vae_loader=_VaeLoader(), + folder_paths_module=_FolderPaths(tmp_path), + ) + + with pytest.raises(ValueError, match="recognizes as Krea 2"): + service.load_models( + "wrong.safetensors", + "default", + "auto", + "default", + "auto", + ) + + assert support.calls == [] + assert resolver.calls == [] + + +def test_missing_host_krea_clip_type_fails_before_downloads(tmp_path: Path) -> None: + """An outdated ComfyUI reports an actionable error without downloading files.""" + + resolver = _Resolver(tmp_path) + support = _ClipTypeSupport(RuntimeError("Update ComfyUI")) + service = Krea2LoaderService( + diffusion_loader=_DiffusionLoader(), + text_encoder_loader=_TextEncoderLoader(), + model_inspector=_Inspector("krea2"), + clip_type_support=support, + resolver=resolver, + vae_loader=_VaeLoader(), + folder_paths_module=_FolderPaths(tmp_path), + ) + + with pytest.raises(RuntimeError, match="Update ComfyUI"): + service.load_models( + "krea.safetensors", + "default", + "auto", + "default", + "auto", + ) + + assert resolver.calls == [] diff --git a/tests/test_model_catalog.py b/tests/test_model_catalog.py index 7d336e7..2a3218f 100644 --- a/tests/test_model_catalog.py +++ b/tests/test_model_catalog.py @@ -8,10 +8,7 @@ from __future__ import annotations import pytest -from simple_syrup.runtime.anima_artifacts import ( - ANIMA_QWEN_TEXT_ENCODER, - ANIMA_QWEN_VAE, -) +from simple_syrup.runtime.anima_artifacts import ANIMA_QWEN_TEXT_ENCODER from simple_syrup.runtime.model_catalog import ( BERT_ENTRY, GROUNDING_DINO_ENTRIES, @@ -24,6 +21,7 @@ from simple_syrup.runtime.model_catalog import ( sam_choices, ultralytics_choices, ) +from simple_syrup.runtime.qwen_artifacts import QWEN_IMAGE_VAE def test_sam_catalog_exposes_layerstyle_compatible_models() -> None: @@ -162,10 +160,8 @@ def test_anima_catalog_has_trusted_auto_artifacts() -> None: ) assert len(ANIMA_QWEN_TEXT_ENCODER.sha256) == 64 - assert ANIMA_QWEN_VAE.folder_name == "vae" - assert ANIMA_QWEN_VAE.filename == "qwen_image_vae.safetensors" - assert ANIMA_QWEN_VAE.canonical_subfolder == "qwen" - assert ANIMA_QWEN_VAE.source_url.endswith( - "/split_files/vae/qwen_image_vae.safetensors" - ) - assert len(ANIMA_QWEN_VAE.sha256) == 64 + assert QWEN_IMAGE_VAE.folder_name == "vae" + assert QWEN_IMAGE_VAE.filename == "qwen_image_vae.safetensors" + assert QWEN_IMAGE_VAE.canonical_subfolder == "qwen" + assert QWEN_IMAGE_VAE.source_url.endswith("/vae/qwen_image_vae.safetensors") + assert len(QWEN_IMAGE_VAE.sha256) == 64 diff --git a/tests/test_registration.py b/tests/test_registration.py index c626794..bf2e6b9 100644 --- a/tests/test_registration.py +++ b/tests/test_registration.py @@ -63,6 +63,7 @@ BASE_NODE_IDS = [ "SimpleSyrup.SimpleLoadCheckpoint", "SimpleSyrup.SimpleLoadFlux", "SimpleSyrup.SimpleLoadFlux2", + "SimpleSyrup.SimpleLoadKrea2", "SimpleSyrup.SimpleVAEEncode", "SimpleSyrup.TagSEGSWithExternalLLM", "SimpleSyrup.TagSEGSWithWD14", diff --git a/tests/test_simple_load_krea2_node.py b/tests/test_simple_load_krea2_node.py new file mode 100644 index 0000000..2fb997c --- /dev/null +++ b/tests/test_simple_load_krea2_node.py @@ -0,0 +1,117 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Tests for the Simple Load Krea 2 v3 node contract.""" + +from __future__ import annotations + +import sys +from types import ModuleType + +import pytest + +from simple_syrup.nodes_v3.simple_load_krea2 import SimpleLoadKrea2V3 +from simple_syrup.runtime.krea2_artifacts import ( + KREA2_QWEN3_VL_4B_BF16, + KREA2_QWEN3_VL_4B_FP8, +) +from simple_syrup.runtime.model_downloads import ComfyProgressReporter + + +class _FolderPaths(ModuleType): + """Provide deterministic schema choices.""" + + def __init__(self) -> None: + """Create model lists including one official local encoder.""" + + super().__init__("folder_paths") + self._files = { + "diffusion_models": ["krea2_turbo.safetensors"], + "text_encoders": [ + "manual.safetensors", + KREA2_QWEN3_VL_4B_BF16.filename, + ], + "vae": ["manual_vae.safetensors"], + "vae_approx": [], + } + + def get_filename_list(self, folder_name: str) -> list[str]: + """Return configured files for one Comfy model folder.""" + + return self._files[folder_name] + + +def test_krea2_schema_keeps_only_model_selection_non_advanced( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The normal surface is one model choice with expert overrides collapsed.""" + + monkeypatch.setitem(sys.modules, "folder_paths", _FolderPaths()) + + schema = SimpleLoadKrea2V3.define_schema() + inputs = {item.id: item for item in schema.inputs} + + assert schema.node_id == "SimpleSyrup.SimpleLoadKrea2" + assert schema.display_name == "Simple Load Krea 2" + assert [item.id for item in schema.inputs] == [ + "diffusion_model", + "diffusion_weight_dtype", + "text_encoder", + "text_encoder_device", + "vae", + ] + assert [output.io_type for output in schema.outputs] == ["MODEL", "CLIP", "VAE"] + assert inputs["diffusion_model"].advanced is None + assert all(item.advanced is True for item in schema.inputs[1:]) + assert inputs["text_encoder"].options[:3] == [ + "auto", + KREA2_QWEN3_VL_4B_FP8.filename, + KREA2_QWEN3_VL_4B_BF16.filename, + ] + assert inputs["text_encoder"].options.count(KREA2_QWEN3_VL_4B_BF16.filename) == 1 + assert inputs["text_encoder"].default == "auto" + assert inputs["vae"].default == "auto" + assert "checksum-pinned" in schema.description + + +def test_krea2_execute_supplies_visible_comfy_progress( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Node execution forwards exact inputs and Comfy download progress.""" + + class _Service: + """Capture loader arguments and return fixed outputs.""" + + def __init__(self) -> None: + """Create an empty argument capture.""" + + self.arguments: dict[str, object] = {} + + def load_models(self, **arguments: object) -> tuple[str, str, str]: + """Record arguments and return deterministic outputs.""" + + self.arguments = arguments + return "model", "clip", "vae" + + service = _Service() + monkeypatch.setattr(SimpleLoadKrea2V3, "_service", service) + + result = SimpleLoadKrea2V3.execute( + "krea.safetensors", + "default", + KREA2_QWEN3_VL_4B_BF16.filename, + "cpu", + "auto", + ) + + assert result == ("model", "clip", "vae") + assert service.arguments == { + "diffusion_model": "krea.safetensors", + "diffusion_weight_dtype": "default", + "text_encoder": KREA2_QWEN3_VL_4B_BF16.filename, + "text_encoder_device": "cpu", + "vae": "auto", + "progress": service.arguments["progress"], + } + assert isinstance(service.arguments["progress"], ComfyProgressReporter)