124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Load checkpoints with optional external VAE replacement."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
from collections.abc import Iterable
|
|
from types import ModuleType
|
|
from typing import Any, Protocol, cast, runtime_checkable
|
|
|
|
from .vae_loader import VaeLoaderService
|
|
|
|
USE_CHECKPOINT_VAE_CHOICE = "Use Checkpoint VAE"
|
|
CLIP_SKIP_DEFAULT = False
|
|
CLIP_SKIP_LAYER = -2
|
|
|
|
|
|
class CheckpointLoaderService:
|
|
"""Load a checkpoint and optionally adjust its CLIP and VAE outputs."""
|
|
|
|
def __init__(
|
|
self,
|
|
folder_paths_module: ModuleType | None = None,
|
|
vae_loader: VaeLoaderBoundary | None = None,
|
|
) -> None:
|
|
"""Create a checkpoint loader with injectable runtime boundaries."""
|
|
|
|
self._folder_paths_module = folder_paths_module
|
|
self._vae_loader = vae_loader or VaeLoaderService(folder_paths_module)
|
|
|
|
def load_checkpoint(
|
|
self,
|
|
ckpt_name: str,
|
|
vae_name: str,
|
|
clip_skip: bool = CLIP_SKIP_DEFAULT,
|
|
) -> tuple[object, object, object]:
|
|
"""Return MODEL, CLIP, and selected VAE objects."""
|
|
|
|
_validate_clip_skip(clip_skip)
|
|
folder_paths = self._folder_paths()
|
|
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
|
comfy_sd = _comfy_sd()
|
|
loaded = tuple(
|
|
cast(
|
|
Iterable[object],
|
|
comfy_sd.load_checkpoint_guess_config(
|
|
ckpt_path,
|
|
output_vae=True,
|
|
output_clip=True,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
),
|
|
)
|
|
)
|
|
model = loaded[0]
|
|
clip = loaded[1]
|
|
checkpoint_vae = loaded[2]
|
|
selected_clip = _selected_clip(clip, clip_skip)
|
|
|
|
if vae_name == USE_CHECKPOINT_VAE_CHOICE:
|
|
return model, selected_clip, checkpoint_vae
|
|
|
|
return model, selected_clip, self._vae_loader.load_vae(vae_name)
|
|
|
|
def _folder_paths(self) -> ModuleType:
|
|
"""Return the ComfyUI folder_paths module."""
|
|
|
|
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 VaeLoaderBoundary(Protocol):
|
|
"""VAE loader interface required by the checkpoint loader."""
|
|
|
|
def load_vae(self, vae_name: str) -> object:
|
|
"""Load the named external VAE."""
|
|
|
|
|
|
@runtime_checkable
|
|
class ClipLayerBoundary(Protocol):
|
|
"""CLIP interface required to apply the ComfyUI clip-skip layer."""
|
|
|
|
def clone(self) -> ClipLayerBoundary:
|
|
"""Return an independent CLIP object."""
|
|
|
|
def clip_layer(self, layer_idx: int) -> None:
|
|
"""Set the CLIP layer index used during prompt encoding."""
|
|
|
|
|
|
def _validate_clip_skip(clip_skip: object) -> None:
|
|
"""Reject non-boolean clip-skip selections from runtime callers."""
|
|
|
|
if not isinstance(clip_skip, bool):
|
|
raise TypeError("clip_skip must be a boolean.")
|
|
|
|
|
|
def _selected_clip(clip: object, clip_skip: bool) -> object:
|
|
"""Return the loaded CLIP or a cloned CLIP with clip skip applied."""
|
|
|
|
if not clip_skip:
|
|
return clip
|
|
|
|
if not isinstance(clip, ClipLayerBoundary):
|
|
raise TypeError(
|
|
"clip_skip requires a CLIP object with clone() and clip_layer()."
|
|
)
|
|
|
|
selected_clip = clip.clone()
|
|
selected_clip.clip_layer(CLIP_SKIP_LAYER)
|
|
return selected_clip
|
|
|
|
|
|
def _comfy_sd() -> Any:
|
|
"""Import ComfyUI's stable diffusion loading module lazily."""
|
|
|
|
return importlib.import_module("comfy.sd")
|