loaders: the four FunPack loaders, as V3 nodes
Ported from v4, minus two mechanisms that are no longer needed. The LoRA loader takes one LoRA and chains, so the FlexibleOptionalInputType dict subclass that lied to ComfyUI's validator has nothing left to do. The CLIP loader uses numbered slots like core's DualCLIPLoader instead of a JSON array in a string widget. Autogrow looked like the right answer for the encoder list and is not: it sets force_input on any widget input, so a list of combos renders as sockets and the file picker stops being a picker. Every core use of it is a connection type. GGUF and SLA are deliberately not here. Both are capabilities rather than loader options, and both belong in modules of their own. The tests follow what the loaders PRODUCE rather than stopping at the call: the attention override is carried into comfy's real dispatch and checked against every backend this machine registers, because an override that only fails at the first attention call fails at sampling, not at load. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Opus 5
Claude Opus 5
parent
eaae196c1f
commit
ab9cbdb556
+38
@@ -0,0 +1,38 @@
|
||||
"""Shared test setup for the whole tree.
|
||||
|
||||
Kept at the root so `modules/**/tests` get the same fixtures as `core/tests`
|
||||
without each one re-deriving the path to ComfyUI.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parent
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(ROOT))
|
||||
|
||||
|
||||
def _comfyui_root():
|
||||
candidates = [os.environ.get("COMFYUI_ROOT"), str(Path.home() / "Documents" / "ComfyUI")]
|
||||
for candidate in candidates:
|
||||
if candidate and (Path(candidate) / "comfy" / "supported_models.py").is_file():
|
||||
if candidate not in sys.path:
|
||||
sys.path.append(candidate)
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def comfyui():
|
||||
"""ComfyUI's source tree on sys.path, or skip.
|
||||
|
||||
Everything that touches a node schema needs it, because a schema is only
|
||||
meaningful in terms of comfy_api's own types.
|
||||
"""
|
||||
root = _comfyui_root()
|
||||
if root is None:
|
||||
pytest.skip("no ComfyUI source on this machine")
|
||||
return root
|
||||
@@ -43,6 +43,15 @@ def test_every_shipped_node_is_a_real_comfy_node(shipped, comfyui):
|
||||
# breaks graphs that already exist, so this test is meant to FAIL on a rename and
|
||||
# make that a decision rather than an accident.
|
||||
FROZEN = {
|
||||
"FunPackCLIPLoader": [
|
||||
"clip_name1", "clip_name2", "clip_name3", "clip_name4", "type", "device",
|
||||
],
|
||||
"FunPackDiffusionModelLoader": [
|
||||
"model_name", "weight_dtype", "compute_dtype", "attention", "fp16_accumulation",
|
||||
],
|
||||
"FunPackLoraLoader": [
|
||||
"model", "lora_name", "strength_model", "clip", "strength_clip",
|
||||
],
|
||||
"FunPackVAELoader": ["vae_name", "dtype"],
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Text encoder loader. A node-only module."""
|
||||
|
||||
from .nodes import FunPackCLIPLoader
|
||||
|
||||
ID = "loader_clip"
|
||||
TITLE = "CLIP loader"
|
||||
STAGE = "load"
|
||||
STATUS = "proven"
|
||||
|
||||
NODES = [FunPackCLIPLoader]
|
||||
@@ -0,0 +1,113 @@
|
||||
"""The text encoder loader.
|
||||
|
||||
One node, up to four files, one CLIP. Unlike the LoRA loader this cannot be a
|
||||
chain: the files are combined into a single encoder, so they have to arrive
|
||||
together. One slot for LTX-2.5 (Gemma4); two for LTX-2.3 (Gemma3 plus its
|
||||
connector).
|
||||
|
||||
The slots are numbered combos, which is what core's DualCLIPLoader and
|
||||
TripleCLIPLoader use. Two alternatives were tried and rejected:
|
||||
|
||||
* **ComfyUI's autogrow input** looks like the right answer and is not.
|
||||
`_AutogrowTemplate.__init__` sets `force_input = True` on any WidgetInput, so a
|
||||
Combo template becomes a row of SOCKETS instead of dropdowns -- the file picker
|
||||
stops being a picker. Every core usage of autogrow is a connection type (Image,
|
||||
Audio, Splat), never a Combo.
|
||||
* **v4's `funpack_list`**, a JSON array in a STRING widget, needed a hand-written
|
||||
canvas widget to draw it, and without one it renders as raw JSON. The file it
|
||||
documented (`web/funpack_list.js`) did not even exist, so on the stock canvas
|
||||
the list was raw JSON in v4 too.
|
||||
|
||||
The cost of numbered slots is a fixed ceiling. Four is above anything shipping
|
||||
(LTX-2.3 needs two), and raising it later only adds optional inputs, which is
|
||||
backwards-compatible with every saved graph.
|
||||
"""
|
||||
|
||||
import comfy.sd
|
||||
import folder_paths
|
||||
import torch
|
||||
from comfy_api.latest import io
|
||||
|
||||
MAX_ENCODERS = 4
|
||||
|
||||
|
||||
def clip_types():
|
||||
try:
|
||||
return [t.name.lower() for t in comfy.sd.CLIPType]
|
||||
except Exception: # noqa: BLE001
|
||||
return ["stable_diffusion"]
|
||||
|
||||
|
||||
def encoder_files():
|
||||
return folder_paths.get_filename_list("text_encoders")
|
||||
|
||||
|
||||
class FunPackCLIPLoader(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
types = clip_types()
|
||||
files = encoder_files()
|
||||
default_type = "ltxv" if "ltxv" in types else (types[0] if types else "stable_diffusion")
|
||||
|
||||
slots = [io.Combo.Input("clip_name1", options=files,
|
||||
tooltip="Text encoder file. The first slot is the encoder itself.")]
|
||||
# Optional, and "" means "no file here" -- an empty slot must not be a
|
||||
# validation failure, because most families only need one.
|
||||
slots += [
|
||||
io.Combo.Input(f"clip_name{n}", options=[""] + files, default="", optional=True,
|
||||
tooltip="Another encoder file, loaded after the ones above. "
|
||||
"Leave empty if the family needs no more.")
|
||||
for n in range(2, MAX_ENCODERS + 1)
|
||||
]
|
||||
|
||||
return io.Schema(
|
||||
node_id="FunPackCLIPLoader",
|
||||
display_name="FunPack CLIP Loader",
|
||||
category="FunPack/Loaders",
|
||||
description="Load one or more text encoder files into a single CLIP.",
|
||||
inputs=[
|
||||
*slots,
|
||||
io.Combo.Input("type", options=types, default=default_type,
|
||||
tooltip="Which family the encoder is for. LTX-2 is 'ltxv'; "
|
||||
"MiniMax H3 is 'minimax'."),
|
||||
io.Combo.Input("device", options=["default", "cpu"], default="default",
|
||||
optional=True,
|
||||
tooltip="'cpu' keeps the encoder off the GPU: slower prompts, "
|
||||
"more VRAM for the diffusion model."),
|
||||
],
|
||||
outputs=[
|
||||
io.Clip.Output(display_name="clip"),
|
||||
io.String.Output(display_name="status"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, clip_name1: str, type: str, device: str = "default", **slots) -> io.NodeOutput:
|
||||
# Order matters to load_clip, so the slots are read in slot order rather
|
||||
# than in whatever order kwargs arrive.
|
||||
names = [clip_name1] + [slots.get(f"clip_name{n}") for n in range(2, MAX_ENCODERS + 1)]
|
||||
names = [name for name in names if name]
|
||||
if not names:
|
||||
raise RuntimeError("FunPack CLIP Loader: pick at least one text encoder file.")
|
||||
|
||||
clip_type = getattr(comfy.sd.CLIPType, str(type).upper(),
|
||||
comfy.sd.CLIPType.STABLE_DIFFUSION)
|
||||
model_options = {}
|
||||
if device == "cpu":
|
||||
model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")
|
||||
|
||||
paths = [folder_paths.get_full_path_or_raise("text_encoders", name) for name in names]
|
||||
clip = comfy.sd.load_clip(
|
||||
ckpt_paths=paths,
|
||||
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
||||
clip_type=clip_type,
|
||||
model_options=model_options,
|
||||
)
|
||||
if clip is None:
|
||||
# comfy returns None for an unrecognised encoder rather than raising,
|
||||
# and a None CLIP fails at encode time, far from the mistake.
|
||||
raise RuntimeError(
|
||||
f"Could not load a text encoder from {', '.join(names)} as type {type!r}.")
|
||||
|
||||
listing = "\n".join(f" {i + 1}. {name}" for i, name in enumerate(names))
|
||||
return io.NodeOutput(clip, f"FunPack CLIP Loader | type={type} device={device}\n{listing}")
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Shared loading options.
|
||||
|
||||
Not a module: it announces nothing and ships no node. `registry.scan` only walks
|
||||
two levels of DIRECTORIES, so a plain file at the domain level is a helper the
|
||||
loaders import and nothing else.
|
||||
|
||||
These options are on the loader node rather than in a settings payload, and the
|
||||
line is worth stating: a loader widget describes the FILE and how it is read,
|
||||
which the loader cannot be handed by anyone else. Feature knobs -- strengths,
|
||||
sigma windows, anything a modifier owns -- are settings, and never appear here.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
|
||||
# Named exactly as core's UNETLoader names them where they overlap, so a value
|
||||
# copied from a stock workflow means the same thing here.
|
||||
WEIGHT_DTYPES = ["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2", "fp16", "bf16", "fp32"]
|
||||
COMPUTE_DTYPES = ["default", "fp16", "bf16", "fp32"]
|
||||
VAE_DTYPES = ["bfloat16", "float16", "float32"]
|
||||
|
||||
_BY_NAME = {
|
||||
"fp16": torch.float16,
|
||||
"bf16": torch.bfloat16,
|
||||
"fp32": torch.float32,
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float32": torch.float32,
|
||||
"fp8_e4m3fn": getattr(torch, "float8_e4m3fn", None),
|
||||
"fp8_e5m2": getattr(torch, "float8_e5m2", None),
|
||||
# fp8_e4m3fn plus comfy's fp8 matmul path; see weight_model_options.
|
||||
"fp8_e4m3fn_fast": getattr(torch, "float8_e4m3fn", None),
|
||||
}
|
||||
|
||||
|
||||
def dtype_of(name):
|
||||
"""The torch dtype a choice names, or None for 'default'/an unsupported build."""
|
||||
return _BY_NAME.get(str(name or "default"))
|
||||
|
||||
|
||||
def weight_model_options(weight_dtype):
|
||||
"""comfy `model_options` for a weight dtype. The mapping core's UNETLoader uses."""
|
||||
options = {}
|
||||
dtype = dtype_of(weight_dtype)
|
||||
if dtype is not None:
|
||||
options["dtype"] = dtype
|
||||
if weight_dtype == "fp8_e4m3fn_fast":
|
||||
options["fp8_optimizations"] = True
|
||||
return options
|
||||
|
||||
|
||||
def attention_choices():
|
||||
"""Backends this ComfyUI can actually run.
|
||||
|
||||
Read from ComfyUI's own registry rather than a private import, so the list
|
||||
offers exactly what the machine has -- sage where SageAttention is installed,
|
||||
flash where flash-attn is -- and never lies about the rest. 'default' leaves
|
||||
whatever ComfyUI itself selected from its launch flags.
|
||||
"""
|
||||
names = []
|
||||
try:
|
||||
from comfy.ldm.modules.attention import REGISTERED_ATTENTION_FUNCTIONS
|
||||
names = sorted(REGISTERED_ATTENTION_FUNCTIONS.keys())
|
||||
except Exception: # noqa: BLE001 -- older ComfyUI
|
||||
pass
|
||||
return ["default"] + names
|
||||
|
||||
|
||||
def attention_override(name):
|
||||
"""A transformer_options override routing every attention call to `name`.
|
||||
|
||||
ComfyUI wraps each implementation with `wrap_attn`, which hands the override
|
||||
the original function plus the call's arguments -- so calling the WRAPPED
|
||||
replacement would re-enter that machinery. The unwrapped one is what gets
|
||||
called.
|
||||
"""
|
||||
if not name or name == "default":
|
||||
return None
|
||||
try:
|
||||
from comfy.ldm.modules.attention import get_attention_function
|
||||
except ImportError:
|
||||
logging.warning("[FunPack] this ComfyUI has no attention registry; leaving attention alone")
|
||||
return None
|
||||
chosen = get_attention_function(name, None)
|
||||
if chosen is None:
|
||||
logging.warning("[FunPack] attention backend %r is not available; leaving attention alone", name)
|
||||
return None
|
||||
inner = getattr(chosen, "__wrapped__", chosen)
|
||||
|
||||
def override(_func, *args, **kwargs):
|
||||
return inner(*args, **kwargs)
|
||||
|
||||
return override
|
||||
|
||||
|
||||
def set_fp16_accumulation(enabled):
|
||||
"""torch's fp16 accumulation switch, or None where the build has no such knob."""
|
||||
matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
|
||||
if matmul is None or not hasattr(matmul, "allow_fp16_accumulation"):
|
||||
return None
|
||||
matmul.allow_fp16_accumulation = bool(enabled)
|
||||
return bool(enabled)
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Diffusion model loader. A node-only module."""
|
||||
|
||||
from .nodes import FunPackDiffusionModelLoader
|
||||
|
||||
ID = "loader_diffusion_model"
|
||||
TITLE = "Diffusion model loader"
|
||||
STAGE = "load"
|
||||
STATUS = "proven"
|
||||
|
||||
NODES = [FunPackDiffusionModelLoader]
|
||||
@@ -0,0 +1,86 @@
|
||||
"""The diffusion model loader.
|
||||
|
||||
Picking a file and how to read it. Every option here changes what the loaded
|
||||
model IS, which is why they are node widgets and not settings: nobody downstream
|
||||
can supply them, and a saved workflow has to record them.
|
||||
"""
|
||||
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
from comfy_api.latest import io
|
||||
|
||||
from ..common import (COMPUTE_DTYPES, WEIGHT_DTYPES, attention_choices,
|
||||
attention_override, dtype_of, set_fp16_accumulation,
|
||||
weight_model_options)
|
||||
|
||||
|
||||
class FunPackDiffusionModelLoader(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
node_id="FunPackDiffusionModelLoader",
|
||||
display_name="FunPack Diffusion Model Loader",
|
||||
category="FunPack/Loaders",
|
||||
description="Load a diffusion model, choosing precision and attention backend.",
|
||||
inputs=[
|
||||
io.Combo.Input("model_name", options=folder_paths.get_filename_list("diffusion_models")),
|
||||
io.Combo.Input("weight_dtype", options=WEIGHT_DTYPES, default="default",
|
||||
tooltip="How the weights are stored. 'default' keeps whatever "
|
||||
"the file already is."),
|
||||
io.Combo.Input("compute_dtype", options=COMPUTE_DTYPES, default="default",
|
||||
tooltip="What the maths runs in, independently of storage."),
|
||||
io.Combo.Input("attention", options=attention_choices(), default="default",
|
||||
tooltip="Attention backend. Only what this machine can run is "
|
||||
"listed; 'default' keeps ComfyUI's launch choice."),
|
||||
io.Boolean.Input("fp16_accumulation", default=False, optional=True,
|
||||
tooltip="Faster fp16 matmuls where the torch build supports it."),
|
||||
],
|
||||
outputs=[
|
||||
io.Model.Output(display_name="model"),
|
||||
io.String.Output(display_name="status"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_name: str, weight_dtype: str, compute_dtype: str,
|
||||
attention: str, fp16_accumulation: bool = False) -> io.NodeOutput:
|
||||
notes = [f"FunPack Diffusion Model Loader | {model_name}"]
|
||||
|
||||
accumulation = set_fp16_accumulation(fp16_accumulation)
|
||||
if fp16_accumulation and accumulation is None:
|
||||
notes.append("fp16_accumulation: unsupported by this torch build, ignored")
|
||||
elif accumulation is not None:
|
||||
notes.append(f"fp16_accumulation: {accumulation}")
|
||||
|
||||
model_options = weight_model_options(weight_dtype)
|
||||
notes.append(f"weight dtype: {weight_dtype}")
|
||||
|
||||
path = folder_paths.get_full_path_or_raise("diffusion_models", model_name)
|
||||
state_dict, metadata = comfy.utils.load_torch_file(path, return_metadata=True)
|
||||
model = comfy.sd.load_diffusion_model_state_dict(
|
||||
state_dict, model_options=model_options, metadata=metadata)
|
||||
if model is None:
|
||||
# comfy returns None rather than raising, and a None reaching the
|
||||
# sampler fails much further away than the mistake.
|
||||
raise RuntimeError(
|
||||
f"Could not detect a diffusion model in {model_name}. Loading a text "
|
||||
f"encoder or a VAE file here is the usual cause.")
|
||||
|
||||
dtype = dtype_of(compute_dtype)
|
||||
if dtype is not None and hasattr(model, "set_model_compute_dtype"):
|
||||
model.set_model_compute_dtype(dtype)
|
||||
model.force_cast_weights = False
|
||||
notes.append(f"compute dtype: {compute_dtype}")
|
||||
elif dtype is not None:
|
||||
notes.append("compute dtype: unsupported by this ComfyUI, ignored")
|
||||
|
||||
override = attention_override(attention)
|
||||
if override is not None:
|
||||
model.model_options.setdefault("transformer_options", {})[
|
||||
"optimized_attention_override"] = override
|
||||
notes.append(f"attention: {attention}")
|
||||
else:
|
||||
notes.append("attention: default (as launched)")
|
||||
|
||||
return io.NodeOutput(model, "\n".join(notes))
|
||||
@@ -0,0 +1,10 @@
|
||||
"""LoRA loader. A node-only module."""
|
||||
|
||||
from .nodes import FunPackLoraLoader
|
||||
|
||||
ID = "loader_lora"
|
||||
TITLE = "LoRA loader"
|
||||
STAGE = "load"
|
||||
STATUS = "proven"
|
||||
|
||||
NODES = [FunPackLoraLoader]
|
||||
@@ -0,0 +1,58 @@
|
||||
"""The LoRA loader.
|
||||
|
||||
One LoRA per node, chained for more. v4 carried N LoRAs in a single node behind a
|
||||
JSON list widget and a `FlexibleOptionalInputType` -- a dict subclass that lied to
|
||||
ComfyUI's validator so arbitrary keys would pass. Chaining needs neither: it is
|
||||
what ComfyUI's own LoraLoader does, every strength is a real socket, and the graph
|
||||
shows how many LoRAs are actually applied instead of hiding the count in a string.
|
||||
"""
|
||||
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
from comfy_api.latest import io
|
||||
|
||||
|
||||
class FunPackLoraLoader(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
node_id="FunPackLoraLoader",
|
||||
display_name="FunPack LoRA Loader",
|
||||
category="FunPack/Loaders",
|
||||
description="Apply one LoRA to a model, and optionally its text encoder. "
|
||||
"Chain these for more than one.",
|
||||
inputs=[
|
||||
io.Model.Input("model"),
|
||||
io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras")),
|
||||
io.Float.Input("strength_model", default=1.0, min=-100.0, max=100.0, step=0.01),
|
||||
io.Clip.Input("clip", optional=True,
|
||||
tooltip="Leave unwired for a model-only LoRA."),
|
||||
io.Float.Input("strength_clip", default=1.0, min=-100.0, max=100.0, step=0.01,
|
||||
optional=True),
|
||||
],
|
||||
outputs=[
|
||||
io.Model.Output(display_name="model"),
|
||||
io.Clip.Output(display_name="clip"),
|
||||
io.String.Output(display_name="status"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model, lora_name: str, strength_model: float,
|
||||
clip=None, strength_clip: float = 1.0) -> io.NodeOutput:
|
||||
# Doing nothing is a valid request, but it must return the SAME objects
|
||||
# rather than clones: a clone here would silently drop patches a later
|
||||
# loader in the chain applied to the original.
|
||||
if strength_model == 0 and (clip is None or strength_clip == 0):
|
||||
return io.NodeOutput(model, clip, f"{lora_name}: both strengths are 0, not applied")
|
||||
|
||||
path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
||||
lora = comfy.utils.load_torch_file(path, safe_load=True)
|
||||
patched_model, patched_clip = comfy.sd.load_lora_for_models(
|
||||
model, clip, lora, strength_model, strength_clip)
|
||||
|
||||
applied = [f"model={strength_model}"]
|
||||
applied.append(f"clip={strength_clip}" if clip is not None else "clip=not wired")
|
||||
return io.NodeOutput(patched_model, patched_clip,
|
||||
f"FunPack LoRA Loader | {lora_name} ({', '.join(applied)})")
|
||||
@@ -0,0 +1,119 @@
|
||||
"""The attention override, carried all the way to a real attention call.
|
||||
|
||||
This is the seam the loader is most likely to get silently wrong: a bad override
|
||||
loads perfectly and then fails at the first attention call, which is sampling.
|
||||
None of it needs model weights, so there is no excuse for not running it.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def torch_and_attention(comfyui):
|
||||
import torch
|
||||
import comfy.ldm.modules.attention as attention
|
||||
return torch, attention
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def qkv(torch_and_attention):
|
||||
torch, _ = torch_and_attention
|
||||
torch.manual_seed(0)
|
||||
heads, dim = 2, 16
|
||||
shape = (1, 8, heads * dim)
|
||||
return torch.randn(shape), torch.randn(shape), torch.randn(shape), heads
|
||||
|
||||
|
||||
def test_every_backend_this_machine_offers_survives_the_override(torch_and_attention, qkv):
|
||||
"""The override is handed the WRAPPED function plus the wrapper's own kwargs
|
||||
(`_inside_attn_wrapper` among them). If the unwrapped backend chokes on any
|
||||
of that, choosing a backend breaks sampling and nothing else."""
|
||||
torch, attention = torch_and_attention
|
||||
from modules.loaders.common import attention_override
|
||||
|
||||
q, k, v, heads = qkv
|
||||
reference = attention.get_attention_function("pytorch")(q, k, v, heads=heads)
|
||||
|
||||
for name in sorted(attention.REGISTERED_ATTENTION_FUNCTIONS):
|
||||
override = attention_override(name)
|
||||
assert override is not None, f"{name} is registered but produced no override"
|
||||
out = override(attention.get_attention_function(name), q, k, v, heads=heads,
|
||||
transformer_options={}, _inside_attn_wrapper=True)
|
||||
assert out.shape == reference.shape, f"{name} changed the output shape"
|
||||
|
||||
|
||||
def test_an_installed_override_is_actually_reached_by_comfys_dispatch(torch_and_attention, qkv):
|
||||
"""Proves the KEY the loader writes is the key comfy reads. A typo here is
|
||||
invisible: the model loads, sampling runs, and the chosen backend is simply
|
||||
never used."""
|
||||
torch, attention = torch_and_attention
|
||||
from modules.loaders.common import attention_override
|
||||
|
||||
q, k, v, heads = qkv
|
||||
inner = attention_override("pytorch")
|
||||
reached = []
|
||||
|
||||
def spy(func, *args, **kwargs):
|
||||
reached.append(True)
|
||||
return inner(func, *args, **kwargs)
|
||||
|
||||
# Exactly the dict the diffusion loader writes into model_options.
|
||||
transformer_options = {"optimized_attention_override": spy}
|
||||
out = attention.optimized_attention(q, k, v, heads=heads,
|
||||
transformer_options=transformer_options)
|
||||
|
||||
assert reached, "the override was installed where nothing reads it"
|
||||
plain = attention.optimized_attention(q, k, v, heads=heads, transformer_options={})
|
||||
assert torch.allclose(out, plain, atol=1e-5), (
|
||||
"routing through the override changed the result")
|
||||
|
||||
|
||||
def test_the_loader_writes_the_override_where_dispatch_looks(comfyui, monkeypatch):
|
||||
"""The loader's own installation step, without loading a model: the key name
|
||||
and its nesting under model_options are what this pins."""
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
from comfy.ldm.modules import attention as attn
|
||||
from modules.loaders.diffusion_model import nodes
|
||||
|
||||
class FakeModel:
|
||||
def __init__(self):
|
||||
self.model_options = {}
|
||||
|
||||
def wrapped(*args, **kwargs):
|
||||
raise AssertionError("called the wrapped function, re-entering wrap_attn")
|
||||
|
||||
wrapped.__wrapped__ = lambda *a, **k: "attended"
|
||||
monkeypatch.setattr(attn, "get_attention_function", lambda name, default: wrapped)
|
||||
monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/fake/{name}")
|
||||
monkeypatch.setattr(comfy.utils, "load_torch_file", lambda p, **kw: ({}, None))
|
||||
monkeypatch.setattr(comfy.sd, "load_diffusion_model_state_dict",
|
||||
lambda sd, **kw: FakeModel())
|
||||
|
||||
out = nodes.FunPackDiffusionModelLoader.execute(
|
||||
model_name="m.safetensors", weight_dtype="default",
|
||||
compute_dtype="default", attention="pytorch")
|
||||
|
||||
model = out.result[0]
|
||||
installed = model.model_options["transformer_options"]["optimized_attention_override"]
|
||||
assert callable(installed)
|
||||
assert installed(wrapped, "q") == "attended"
|
||||
|
||||
|
||||
def test_a_model_comfy_could_not_identify_fails_at_the_loader(comfyui, monkeypatch):
|
||||
"""comfy returns None instead of raising, and a None model reaches the
|
||||
sampler before anything complains."""
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
from modules.loaders.diffusion_model import nodes
|
||||
|
||||
monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/fake/{name}")
|
||||
monkeypatch.setattr(comfy.utils, "load_torch_file", lambda p, **kw: ({}, None))
|
||||
monkeypatch.setattr(comfy.sd, "load_diffusion_model_state_dict", lambda sd, **kw: None)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Could not detect a diffusion model"):
|
||||
nodes.FunPackDiffusionModelLoader.execute(
|
||||
model_name="actually_a_vae.safetensors", weight_dtype="default",
|
||||
compute_dtype="default", attention="default")
|
||||
@@ -0,0 +1,156 @@
|
||||
"""What the loaders actually do with what they are given.
|
||||
|
||||
Real weights cannot be loaded here (this machine's model files are zero-byte
|
||||
placeholders), so the file-reading boundary is stubbed and everything on OUR side
|
||||
of it runs for real: slot ordering, empty slots, no-op strengths, and where the
|
||||
attention override is installed.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _needs_comfy(comfyui):
|
||||
"""Every test here imports a node module, which imports comfy."""
|
||||
|
||||
|
||||
def test_empty_encoder_slots_are_skipped_not_looked_up(monkeypatch):
|
||||
"""An empty optional slot means "no file here". Passing "" through to
|
||||
get_full_path_or_raise turns a normal one-encoder setup into a hard error."""
|
||||
import comfy.sd
|
||||
import folder_paths
|
||||
from modules.loaders.clip import nodes
|
||||
|
||||
asked = []
|
||||
monkeypatch.setattr(folder_paths, "get_full_path_or_raise",
|
||||
lambda kind, name: asked.append(name) or f"/fake/{name}")
|
||||
monkeypatch.setattr(folder_paths, "get_folder_paths", lambda kind: [])
|
||||
monkeypatch.setattr(comfy.sd, "load_clip", lambda **kw: object())
|
||||
|
||||
nodes.FunPackCLIPLoader.execute(
|
||||
clip_name1="encoder.safetensors", type="ltxv",
|
||||
clip_name2="", clip_name3=None, clip_name4="")
|
||||
|
||||
assert asked == ["encoder.safetensors"]
|
||||
|
||||
|
||||
def test_encoder_slots_load_in_slot_order(monkeypatch):
|
||||
"""load_clip cares about order: the encoder comes before its connector.
|
||||
Reading **kwargs in arrival order would make that arbitrary."""
|
||||
import comfy.sd
|
||||
import folder_paths
|
||||
from modules.loaders.clip import nodes
|
||||
|
||||
seen = {}
|
||||
monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/fake/{name}")
|
||||
monkeypatch.setattr(folder_paths, "get_folder_paths", lambda kind: [])
|
||||
monkeypatch.setattr(comfy.sd, "load_clip",
|
||||
lambda **kw: seen.update(kw) or object())
|
||||
|
||||
# Deliberately out of order in the call, which is how kwargs can arrive.
|
||||
nodes.FunPackCLIPLoader.execute(
|
||||
clip_name3="third.safetensors", clip_name2="second.safetensors",
|
||||
clip_name1="first.safetensors", type="ltxv")
|
||||
|
||||
assert seen["ckpt_paths"] == ["/fake/first.safetensors",
|
||||
"/fake/second.safetensors",
|
||||
"/fake/third.safetensors"]
|
||||
|
||||
|
||||
def test_an_unrecognised_encoder_fails_here_not_at_encode_time(monkeypatch):
|
||||
"""comfy returns None rather than raising, and a None CLIP fails much later,
|
||||
a long way from the wrong file that caused it."""
|
||||
import comfy.sd
|
||||
import folder_paths
|
||||
from modules.loaders.clip import nodes
|
||||
|
||||
monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/fake/{name}")
|
||||
monkeypatch.setattr(folder_paths, "get_folder_paths", lambda kind: [])
|
||||
monkeypatch.setattr(comfy.sd, "load_clip", lambda **kw: None)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Could not load a text encoder"):
|
||||
nodes.FunPackCLIPLoader.execute(clip_name1="wrong.safetensors", type="ltxv")
|
||||
|
||||
|
||||
def test_no_encoder_at_all_is_refused(monkeypatch):
|
||||
from modules.loaders.clip import nodes
|
||||
with pytest.raises(RuntimeError, match="at least one"):
|
||||
nodes.FunPackCLIPLoader.execute(clip_name1="", type="ltxv")
|
||||
|
||||
|
||||
def test_a_lora_at_zero_strength_passes_the_originals_through(monkeypatch):
|
||||
"""Not clones. A clone here silently discards patches an earlier loader in
|
||||
the chain applied to the same objects."""
|
||||
import comfy.sd
|
||||
from modules.loaders.lora import nodes
|
||||
|
||||
def _explode(*a, **k):
|
||||
raise AssertionError("a zero-strength LoRA still read the file")
|
||||
|
||||
monkeypatch.setattr(comfy.sd, "load_lora_for_models", _explode)
|
||||
|
||||
model, clip = object(), object()
|
||||
out = nodes.FunPackLoraLoader.execute(
|
||||
model=model, lora_name="x.safetensors", strength_model=0.0,
|
||||
clip=clip, strength_clip=0.0)
|
||||
|
||||
assert out.result[0] is model and out.result[1] is clip
|
||||
|
||||
|
||||
def test_a_model_only_lora_still_applies_at_zero_clip_strength(monkeypatch):
|
||||
"""clip is unwired, so strength_clip is meaningless -- but strength_model is
|
||||
not, and skipping the whole thing would silently drop the LoRA."""
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
from modules.loaders.lora import nodes
|
||||
|
||||
called = {}
|
||||
monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/fake/{name}")
|
||||
monkeypatch.setattr(comfy.utils, "load_torch_file", lambda p, **kw: {"fake": 1})
|
||||
monkeypatch.setattr(comfy.sd, "load_lora_for_models",
|
||||
lambda m, c, l, sm, sc: called.update(sm=sm, sc=sc) or ("patched", None))
|
||||
|
||||
out = nodes.FunPackLoraLoader.execute(
|
||||
model=object(), lora_name="x.safetensors", strength_model=0.8, clip=None)
|
||||
|
||||
assert called["sm"] == 0.8
|
||||
assert out.result[0] == "patched"
|
||||
|
||||
|
||||
def test_the_attention_override_calls_the_unwrapped_backend(monkeypatch):
|
||||
"""ComfyUI wraps each backend with wrap_attn and hands the override the
|
||||
ORIGINAL function. Calling the wrapped one re-enters that machinery."""
|
||||
from comfy.ldm.modules import attention as attn
|
||||
from modules.loaders import common
|
||||
|
||||
calls = []
|
||||
|
||||
def wrapped(*args, **kwargs):
|
||||
calls.append("wrapped")
|
||||
|
||||
def unwrapped(*args, **kwargs):
|
||||
calls.append("unwrapped")
|
||||
return "result"
|
||||
|
||||
wrapped.__wrapped__ = unwrapped
|
||||
monkeypatch.setattr(attn, "get_attention_function", lambda name, default: wrapped)
|
||||
|
||||
override = common.attention_override("anything")
|
||||
assert override(wrapped, "q", "k", "v") == "result"
|
||||
assert calls == ["unwrapped"], "the override re-entered wrap_attn"
|
||||
|
||||
|
||||
def test_default_attention_installs_nothing(monkeypatch):
|
||||
"""'default' must leave ComfyUI's own launch-flag choice alone."""
|
||||
from modules.loaders import common
|
||||
assert common.attention_override("default") is None
|
||||
assert common.attention_override("") is None
|
||||
assert common.attention_override(None) is None
|
||||
|
||||
|
||||
def test_an_unavailable_backend_does_not_silently_become_a_broken_override(monkeypatch):
|
||||
from comfy.ldm.modules import attention as attn
|
||||
from modules.loaders import common
|
||||
monkeypatch.setattr(attn, "get_attention_function", lambda name, default: None)
|
||||
assert common.attention_override("sage") is None
|
||||
@@ -0,0 +1,60 @@
|
||||
"""What the loaders present to the graph, and to the person wiring it.
|
||||
|
||||
These are seam tests, not unit tests: a loader that "loads" is worth nothing if
|
||||
what it offers cannot be filled in, or if what it produces is not what the next
|
||||
stage reads. No model weights are needed for any of this.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def schemas(comfyui):
|
||||
from modules.loaders.clip.nodes import FunPackCLIPLoader
|
||||
from modules.loaders.diffusion_model.nodes import FunPackDiffusionModelLoader
|
||||
from modules.loaders.lora.nodes import FunPackLoraLoader
|
||||
from modules.loaders.vae.nodes import FunPackVAELoader
|
||||
return [FunPackCLIPLoader, FunPackDiffusionModelLoader,
|
||||
FunPackLoraLoader, FunPackVAELoader]
|
||||
|
||||
|
||||
def test_a_file_picker_is_a_widget_not_a_socket(schemas):
|
||||
"""The failure this exists for: ComfyUI's autogrow input sets
|
||||
force_input=True on any WidgetInput, so a list of Combos renders as a row of
|
||||
sockets and the file picker stops being a picker. A dropdown nobody can drop
|
||||
down is not a loader."""
|
||||
for node in schemas:
|
||||
spec = node.INPUT_TYPES()
|
||||
for section in ("required", "optional"):
|
||||
for name, (kind, options) in (spec.get(section) or {}).items():
|
||||
if kind != "COMBO":
|
||||
continue
|
||||
assert not options.get("forceInput"), (
|
||||
f"{node.__name__}.{name} is a COMBO forced to be a socket; "
|
||||
f"nobody can pick a file from it"
|
||||
)
|
||||
|
||||
|
||||
def test_every_combo_offers_something_to_pick(schemas):
|
||||
for node in schemas:
|
||||
spec = node.INPUT_TYPES()
|
||||
for section in ("required", "optional"):
|
||||
for name, (kind, options) in (spec.get(section) or {}).items():
|
||||
if kind != "COMBO":
|
||||
continue
|
||||
assert "options" in options, f"{node.__name__}.{name} has no options at all"
|
||||
|
||||
|
||||
def test_a_default_is_actually_one_of_the_options(schemas):
|
||||
"""A default outside its own list makes ComfyUI refuse to queue the whole
|
||||
prompt with 'Value not in list' -- and the node looks configured."""
|
||||
for node in schemas:
|
||||
spec = node.INPUT_TYPES()
|
||||
for section in ("required", "optional"):
|
||||
for name, (kind, options) in (spec.get(section) or {}).items():
|
||||
if kind != "COMBO" or "default" not in options:
|
||||
continue
|
||||
assert options["default"] in options["options"], (
|
||||
f"{node.__name__}.{name} defaults to {options['default']!r}, "
|
||||
f"which is not one of its options"
|
||||
)
|
||||
Reference in New Issue
Block a user