feat(loaders): FunPack diffusion/CLIP/VAE loaders; LoRA loader works alone
Quantization + a per-model attention backend read from ComfyUI's own registry, so only installed ones are offered. The LoRA loader now takes its own list and scores wrapper prefixes against the model's key map, instead of applying a diffusers LoRA as a silent no-op.
This commit is contained in:
digital-garbage
+13
@@ -19,6 +19,7 @@ if __package__:
|
||||
FunPackStoryMemLastFrameExtractor,
|
||||
FunPackVideoStitch,
|
||||
)
|
||||
from .loaders import FunPackCLIPLoader, FunPackDiffusionModelLoader, FunPackVAELoader
|
||||
from .model_management import FunPackApplyLoraWeights, FunPackLoraLoader
|
||||
from .samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler
|
||||
from .templates import FunPackRefinementKeyLoader
|
||||
@@ -62,6 +63,12 @@ else:
|
||||
except Exception:
|
||||
FunPackApplyLoraWeights = None
|
||||
FunPackLoraLoader = None
|
||||
try:
|
||||
from loaders import FunPackCLIPLoader, FunPackDiffusionModelLoader, FunPackVAELoader
|
||||
except Exception:
|
||||
FunPackCLIPLoader = None
|
||||
FunPackDiffusionModelLoader = None
|
||||
FunPackVAELoader = None
|
||||
try:
|
||||
from samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler
|
||||
except Exception:
|
||||
@@ -96,6 +103,9 @@ NODE_CLASS_MAPPINGS = {
|
||||
"FunPackLTXAVSceneChainSampler": FunPackLTXAVSceneChainSampler,
|
||||
"FunPackApplyLoraWeights": FunPackApplyLoraWeights,
|
||||
"FunPackLoraLoader": FunPackLoraLoader,
|
||||
"FunPackDiffusionModelLoader": FunPackDiffusionModelLoader,
|
||||
"FunPackCLIPLoader": FunPackCLIPLoader,
|
||||
"FunPackVAELoader": FunPackVAELoader,
|
||||
"FunPackRefinementKeyLoader": FunPackRefinementKeyLoader,
|
||||
}
|
||||
NODE_CLASS_MAPPINGS = {name: cls for name, cls in NODE_CLASS_MAPPINGS.items() if cls is not None}
|
||||
@@ -121,6 +131,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FunPackLTXAVSceneChainSampler": "FunPack LTXAV Scene Chain Sampler",
|
||||
"FunPackApplyLoraWeights": "FunPack Apply LoRA Weights",
|
||||
"FunPackLoraLoader": "FunPack LoRA Loader",
|
||||
"FunPackDiffusionModelLoader": "FunPack Diffusion Model Loader",
|
||||
"FunPackCLIPLoader": "FunPack CLIP Loader",
|
||||
"FunPackVAELoader": "FunPack VAE Loader",
|
||||
"FunPackRefinementKeyLoader": "FunPack Refinement Key Loader",
|
||||
}
|
||||
|
||||
|
||||
+281
@@ -0,0 +1,281 @@
|
||||
"""FunPack's own model loaders.
|
||||
|
||||
The point is not to replace ComfyUI's loaders — it is to make the FunPack pipeline
|
||||
configurable by picking files and nothing else. One loader per model kind, every option
|
||||
that actually matters on the node (quantization, compute dtype, attention backend), and
|
||||
list inputs where a family needs more than one file.
|
||||
|
||||
Attention backends come from ComfyUI's own registry rather than a private sageattention
|
||||
import, so this node offers exactly what the machine can run — sage/sage3 when
|
||||
SageAttention is installed, flash when flash-attn is, xformers when enabled — and never
|
||||
lies about the rest.
|
||||
"""
|
||||
import logging
|
||||
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
import torch
|
||||
|
||||
try:
|
||||
from .widgets import field, list_widget, parse_rows
|
||||
except ImportError: # standalone tests import the modules directly
|
||||
from widgets import field, list_widget, parse_rows
|
||||
|
||||
WEIGHT_DTYPES = ["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2", "fp16", "bf16", "fp32"]
|
||||
COMPUTE_DTYPES = ["default", "fp16", "bf16", "fp32"]
|
||||
VAE_DTYPES = ["default", "fp16", "bf16", "fp32"]
|
||||
|
||||
_DTYPE_BY_NAME = {
|
||||
"fp16": torch.float16,
|
||||
"bf16": torch.bfloat16,
|
||||
"fp32": torch.float32,
|
||||
"fp8_e4m3fn": getattr(torch, "float8_e4m3fn", None),
|
||||
"fp8_e5m2": getattr(torch, "float8_e5m2", None),
|
||||
# fp8_e4m3fn_fast is 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 dtype choice names, or None for 'default'/unsupported builds."""
|
||||
return _DTYPE_BY_NAME.get(str(name or "default"))
|
||||
|
||||
|
||||
def weight_model_options(weight_dtype):
|
||||
"""comfy `model_options` for a weight dtype choice. Same 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():
|
||||
"""Attention backends this ComfyUI can actually run, newest-registry first.
|
||||
|
||||
'default' leaves the model on whatever ComfyUI itself selected (the --use-sage-attention
|
||||
/ --use-flash-attention launch flags), which is why it stays the default here.
|
||||
"""
|
||||
names = []
|
||||
try:
|
||||
from comfy.ldm.modules.attention import REGISTERED_ATTENTION_FUNCTIONS
|
||||
names = sorted(REGISTERED_ATTENTION_FUNCTIONS.keys())
|
||||
except Exception: # noqa: BLE001 - older ComfyUI without the registry
|
||||
pass
|
||||
return ["default"] + names
|
||||
|
||||
|
||||
def attention_override(name):
|
||||
"""A transformer_options override that routes every attention call to `name`.
|
||||
|
||||
ComfyUI wraps each attention implementation with `wrap_attn`, which hands the override
|
||||
the original function plus the call's arguments. Calling the wrapped replacement would
|
||||
re-enter that machinery, so the unwrapped function 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. Returns what it ended up as, or None if unsupported."""
|
||||
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)
|
||||
|
||||
|
||||
class FunPackDiffusionModelLoader:
|
||||
"""Loads a diffusion model with quantization and an attention backend chosen per node."""
|
||||
|
||||
CATEGORY = "FunPack/Loaders"
|
||||
RETURN_TYPES = ("MODEL", "STRING")
|
||||
RETURN_NAMES = ("MODEL", "status")
|
||||
FUNCTION = "load_model"
|
||||
DESCRIPTION = ("Diffusion model loader with weight/compute dtype and a per-model attention "
|
||||
"backend (SageAttention, FlashAttention, xformers — whichever are installed).")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model_name": (folder_paths.get_filename_list("diffusion_models"), {
|
||||
"tooltip": "The diffusion model file, from ComfyUI/models/diffusion_models."}),
|
||||
"weight_dtype": (WEIGHT_DTYPES, {
|
||||
"default": "default",
|
||||
"tooltip": "How weights are stored in VRAM. fp8_e4m3fn roughly halves the "
|
||||
"model's memory; _fast additionally runs fp8 matmuls, which is "
|
||||
"faster on Ada and newer and slightly less precise. 'default' "
|
||||
"keeps whatever the file already is."}),
|
||||
"compute_dtype": (COMPUTE_DTYPES, {
|
||||
"default": "default",
|
||||
"tooltip": "The dtype the model computes in. Setting this stops per-weight "
|
||||
"casting, so an fp8 model computes in bf16 without re-casting "
|
||||
"every layer."}),
|
||||
"attention": (attention_choices(), {
|
||||
"default": "default",
|
||||
"tooltip": "Attention backend for THIS model. 'default' uses whatever "
|
||||
"ComfyUI was launched with. Only backends installed on this "
|
||||
"machine are listed."}),
|
||||
},
|
||||
"optional": {
|
||||
"fp16_accumulation": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Accumulate fp16 matmuls in fp16. Faster on recent NVIDIA cards; "
|
||||
"needs a torch build that supports it, ignored otherwise."}),
|
||||
},
|
||||
}
|
||||
|
||||
def load_model(self, model_name, weight_dtype, compute_dtype, attention, fp16_accumulation=False):
|
||||
notes = [f"FunPack Diffusion Model Loader | {model_name}"]
|
||||
|
||||
accum = set_fp16_accumulation(fp16_accumulation)
|
||||
if fp16_accumulation and accum is None:
|
||||
notes.append("fp16_accumulation: unsupported by this torch build, ignored")
|
||||
elif accum is not None:
|
||||
notes.append(f"fp16_accumulation: {accum}")
|
||||
|
||||
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:
|
||||
raise RuntimeError(
|
||||
f"ERROR: could not detect a diffusion model in {model_name}. "
|
||||
"Loading a text encoder or 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 (model, "\n".join(notes))
|
||||
|
||||
|
||||
class FunPackCLIPLoader:
|
||||
"""Loads one or more text encoder files into a single CLIP."""
|
||||
|
||||
CATEGORY = "FunPack/Loaders"
|
||||
RETURN_TYPES = ("CLIP", "STRING")
|
||||
RETURN_NAMES = ("CLIP", "status")
|
||||
FUNCTION = "load_clip"
|
||||
DESCRIPTION = ("Text encoder loader. Add one slot for a single-file encoder (LTX-2.5's "
|
||||
"Gemma4), or two for an encoder plus its connector (LTX-2.3's Gemma3).")
|
||||
|
||||
@classmethod
|
||||
def clip_types(cls):
|
||||
try:
|
||||
return [t.name.lower() for t in comfy.sd.CLIPType]
|
||||
except Exception: # noqa: BLE001
|
||||
return ["ltxv"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
encoders = folder_paths.get_filename_list("text_encoders")
|
||||
types = cls.clip_types()
|
||||
return {
|
||||
"required": {
|
||||
"clip_list": list_widget(
|
||||
"text encoder",
|
||||
[field("clip_name", "combo", label="file", choices=encoders)],
|
||||
add_label="+ Add slot",
|
||||
tooltip="Text encoder files, in load order. One slot for LTX-2.5 (Gemma4); "
|
||||
"two for LTX-2.3 (Gemma3 + its connector)."),
|
||||
"type": (types, {
|
||||
"default": "ltxv" if "ltxv" in types else (types[0] if types else "ltxv"),
|
||||
"tooltip": "Which model family the encoder is for. LTX-2 (all point "
|
||||
"releases) is 'ltxv'; MiniMax H3 is 'minimax'."}),
|
||||
},
|
||||
"optional": {
|
||||
"device": (["default", "cpu"], {
|
||||
"default": "default",
|
||||
"tooltip": "'cpu' keeps the encoder off the GPU — slower prompts, more VRAM "
|
||||
"left for the diffusion model."}),
|
||||
},
|
||||
}
|
||||
|
||||
def load_clip(self, clip_list, type, device="default"):
|
||||
rows = parse_rows(clip_list, [field("clip_name", "combo")], key="clip_name")
|
||||
names = [row["clip_name"] for row in rows]
|
||||
if not names:
|
||||
raise RuntimeError("FunPack CLIP Loader: add 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", n) for n 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,
|
||||
)
|
||||
status = "FunPack CLIP Loader | type={} device={}\n{}".format(
|
||||
type, device, "\n".join(f" {i + 1}. {n}" for i, n in enumerate(names)))
|
||||
return (clip, status)
|
||||
|
||||
|
||||
class FunPackVAELoader:
|
||||
"""Loads one VAE. Add a second instance for the audio VAE."""
|
||||
|
||||
CATEGORY = "FunPack/Loaders"
|
||||
RETURN_TYPES = ("VAE", "STRING")
|
||||
RETURN_NAMES = ("VAE", "status")
|
||||
FUNCTION = "load_vae"
|
||||
DESCRIPTION = ("VAE loader with an explicit dtype. LTX-2 and MiniMax H3 both ship a video "
|
||||
"VAE and an audio VAE — use one of these per VAE.")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"vae_name": (folder_paths.get_filename_list("vae"), {
|
||||
"tooltip": "The VAE file, from ComfyUI/models/vae."}),
|
||||
"dtype": (VAE_DTYPES, {
|
||||
"default": "default",
|
||||
"tooltip": "MiniMax H3's VAEs must run in bf16 — fp16 produces garbage and "
|
||||
"fp32 decodes slowly. 'default' lets ComfyUI choose."}),
|
||||
},
|
||||
}
|
||||
|
||||
def load_vae(self, vae_name, dtype="default"):
|
||||
path = folder_paths.get_full_path_or_raise("vae", vae_name)
|
||||
state_dict, metadata = comfy.utils.load_torch_file(path, return_metadata=True)
|
||||
vae = comfy.sd.VAE(sd=state_dict, dtype=dtype_of(dtype), metadata=metadata)
|
||||
check = getattr(vae, "throw_exception_if_invalid", None)
|
||||
if callable(check):
|
||||
check()
|
||||
return (vae, f"FunPack VAE Loader | {vae_name} dtype={dtype}")
|
||||
+154
-23
@@ -15,6 +15,11 @@ import torch
|
||||
from aiohttp import web
|
||||
from server import PromptServer
|
||||
|
||||
try:
|
||||
from .widgets import field, list_widget, parse_rows
|
||||
except ImportError: # standalone tests import the modules directly
|
||||
from widgets import field, list_widget, parse_rows
|
||||
|
||||
LORA_TYPES = ["general", "action", "style", "quality", "character"]
|
||||
LORA_STACK_TYPE = "FUNPACK_LORA_STACK"
|
||||
TRANSFORMER_BLOCK_PATTERN = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.")
|
||||
@@ -183,6 +188,109 @@ def patch_energy(value):
|
||||
return 0.0
|
||||
|
||||
|
||||
def lora_row_fields():
|
||||
"""The columns of one LoRA row, shared by the list widget and its parser."""
|
||||
loras = ["None"] + folder_paths.get_filename_list("loras")
|
||||
return [
|
||||
field("lora", "combo", label="LoRA", choices=loras, default="None"),
|
||||
field("type", "combo", label="type", choices=LORA_TYPES, default="general",
|
||||
tooltip="What the LoRA is for. Drives per-block placement when per_block is on."),
|
||||
field("strength", "float", label="strength", default=1.0, min=-10.0, max=10.0, step=0.01),
|
||||
]
|
||||
|
||||
|
||||
def lora_list_input(tooltip):
|
||||
return list_widget("LoRA", lora_row_fields(), add_label="+ Add LoRA", tooltip=tooltip)
|
||||
|
||||
|
||||
def stack_from_lora_list(value, per_block=False):
|
||||
"""A LoRA stack straight from a list widget, with no refinement suggestions applied.
|
||||
|
||||
Same row shape FunPack Apply LoRA Weights emits, so the loader cannot tell the two
|
||||
apart — which is what lets a LoRA be used without wiring a second node for it.
|
||||
"""
|
||||
entries = []
|
||||
for row in parse_rows(value, lora_row_fields(), key="lora"):
|
||||
name = row["lora"]
|
||||
lora_type = normalize_lora_type(row["type"])
|
||||
entries.append({
|
||||
"slot": row["index"],
|
||||
"name": name,
|
||||
"type": lora_type,
|
||||
"id": lora_state_id(name, lora_type),
|
||||
"base_model_weight": row["strength"],
|
||||
"model_weight": row["strength"],
|
||||
"source": "list",
|
||||
})
|
||||
return {
|
||||
"version": 2,
|
||||
"refinement_key": "",
|
||||
"mode": "ltx2",
|
||||
"per_block": coerce_bool(per_block),
|
||||
"positive_prompt": "",
|
||||
"prompt_key": "",
|
||||
"loras": entries,
|
||||
}
|
||||
|
||||
|
||||
# Wrapper prefixes trainers put in front of otherwise-standard keys. ComfyUI's own key maps
|
||||
# already cover diffusion_model./lora_unet_/lycoris_ and, for LTX, the bare
|
||||
# transformer_blocks.* form -- these are the ones that arrive wrapped and match nothing at
|
||||
# all until the wrapper comes off.
|
||||
LORA_KEY_PREFIXES = (
|
||||
"transformer.",
|
||||
"diffusion_model.",
|
||||
"model.diffusion_model.",
|
||||
"base_model.model.",
|
||||
"lora_model.",
|
||||
"net.",
|
||||
)
|
||||
|
||||
|
||||
def _count_lora_matches(lora, key_map):
|
||||
try:
|
||||
return len(comfy.lora.load_lora(lora, key_map, log_missing=False))
|
||||
except TypeError: # ComfyUI without the log_missing argument
|
||||
return len(comfy.lora.load_lora(lora, key_map))
|
||||
|
||||
|
||||
def _strip_key_prefix(lora, prefix):
|
||||
if not any(k.startswith(prefix) for k in lora):
|
||||
return None
|
||||
return {(k[len(prefix):] if k.startswith(prefix) else k): v for k, v in lora.items()}
|
||||
|
||||
|
||||
def resolve_lora_patches(model, lora, clip=None):
|
||||
"""Match a LoRA against a model, trying known wrapper prefixes. -> (patches, note).
|
||||
|
||||
ComfyUI's converters handle the tensor-naming dialects (lora_up/lora_A/lora.up/PEFT
|
||||
defaults); what they do not handle is a whole state dict nested under a wrapper the
|
||||
model's key map never mentions, which is how a diffusers-trained LoRA ends up applying
|
||||
zero weights in silence. Every candidate is scored against the real key map and the one
|
||||
that matches the most weights wins, so a format nobody anticipated still lands as long
|
||||
as its keys are recognisable underneath.
|
||||
"""
|
||||
key_map = comfy.lora.model_lora_keys_unet(model.model, {})
|
||||
if clip is not None:
|
||||
key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)
|
||||
converted = comfy.lora_convert.convert_lora(lora)
|
||||
|
||||
best, best_count, best_note = converted, _count_lora_matches(converted, key_map), "as-is"
|
||||
for prefix in LORA_KEY_PREFIXES:
|
||||
variant = _strip_key_prefix(converted, prefix)
|
||||
if variant is None:
|
||||
continue
|
||||
count = _count_lora_matches(variant, key_map)
|
||||
if count > best_count:
|
||||
best, best_count, best_note = variant, count, f"stripped {prefix}"
|
||||
|
||||
patches = comfy.lora.load_lora(best, key_map)
|
||||
if not patches:
|
||||
logging.warning("[FunPack] LoRA matched no weights in this model - wrong model family?")
|
||||
return patches, "MATCHED NOTHING"
|
||||
return patches, f"keys={len(patches)} fmt={best_note}"
|
||||
|
||||
|
||||
class FunPackApplyLoraWeights:
|
||||
"""
|
||||
Builds a LoRA stack from user base weights, then applies prompt-specific
|
||||
@@ -201,7 +309,10 @@ class FunPackApplyLoraWeights:
|
||||
optional = FlexibleOptionalInputType(
|
||||
any_type,
|
||||
{
|
||||
"lora_list": ("STRING", {"default": "[]", "multiline": False}),
|
||||
# The canvas draws this with web/funpack_lora_weights.js; the funpack_list
|
||||
# spec is what lets the FunPack Editor draw the same rows instead of raw JSON.
|
||||
"lora_list": lora_list_input(
|
||||
"LoRAs whose weights this node looks up for the current prompt."),
|
||||
"refinement_key_input": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
@@ -458,10 +569,21 @@ class FunPackLoraLoader:
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"lora_stack": (LORA_STACK_TYPE,),
|
||||
"lora_list": lora_list_input(
|
||||
"LoRAs to apply, top to bottom. Enough on its own — a stack is only "
|
||||
"needed for prompt-specific trained weights."),
|
||||
},
|
||||
"optional": {
|
||||
"clip": ("CLIP",),
|
||||
"lora_stack": (LORA_STACK_TYPE, {
|
||||
"tooltip": "Optional stack from FunPack Apply LoRA Weights, carrying "
|
||||
"prompt-specific trained strengths. Its LoRAs are applied "
|
||||
"first, then this node's own list."}),
|
||||
"per_block": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Analyze each LoRA's block deltas and balance competing block "
|
||||
"strengths. LTX models only; a wired stack can switch this on "
|
||||
"by itself."}),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -512,21 +634,18 @@ class FunPackLoraLoader:
|
||||
)
|
||||
|
||||
def _load_model_lora_patches(self, model, lora, lora_cache_key, model_cache_key=None):
|
||||
"""(patches, format note). The note says how the LoRA's keys had to be read."""
|
||||
if lora_cache_key is None:
|
||||
key_map = comfy.lora.model_lora_keys_unet(model.model, {})
|
||||
converted_lora = comfy.lora_convert.convert_lora(lora)
|
||||
return comfy.lora.load_lora(converted_lora, key_map)
|
||||
return resolve_lora_patches(model, lora)
|
||||
|
||||
cache_key = (lora_cache_key, model_cache_key or self._model_cache_key(model))
|
||||
loaded = self._cache_get(self.model_patch_cache, cache_key)
|
||||
if loaded is not None:
|
||||
return loaded
|
||||
cached = self._cache_get(self.model_patch_cache, cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
key_map = comfy.lora.model_lora_keys_unet(model.model, {})
|
||||
converted_lora = comfy.lora_convert.convert_lora(lora)
|
||||
loaded = comfy.lora.load_lora(converted_lora, key_map)
|
||||
self._cache_put(self.model_patch_cache, cache_key, loaded, LORA_PATCH_CACHE_SIZE)
|
||||
return loaded
|
||||
result = resolve_lora_patches(model, lora)
|
||||
self._cache_put(self.model_patch_cache, cache_key, result, LORA_PATCH_CACHE_SIZE)
|
||||
return result
|
||||
|
||||
def _split_model_patches_by_block(self, loaded):
|
||||
global_patches = {}
|
||||
@@ -793,10 +912,10 @@ class FunPackLoraLoader:
|
||||
|
||||
def _load_lora_per_block(self, model, lora, model_weight):
|
||||
lora_cache_key = None
|
||||
loaded = self._load_model_lora_patches(model, lora, lora_cache_key)
|
||||
loaded, fmt = self._load_model_lora_patches(model, lora, lora_cache_key)
|
||||
profile = self._lora_block_profile({}, loaded, model_weight)
|
||||
if not profile:
|
||||
return None, "per-block fallback=global"
|
||||
return None, f"per-block fallback=global {fmt}"
|
||||
|
||||
new_model, non_block_count, _ = self._apply_model_patches(
|
||||
model,
|
||||
@@ -809,13 +928,19 @@ class FunPackLoraLoader:
|
||||
max_eff = model_weight * max(scales.values())
|
||||
status = (
|
||||
f"per-block blocks={len(scales)} non_block={non_block_count} "
|
||||
f"range={min_eff:.3f}..{max_eff:.3f}"
|
||||
f"range={min_eff:.3f}..{max_eff:.3f} {fmt}"
|
||||
)
|
||||
return new_model, status
|
||||
|
||||
def load_loras(self, model, lora_stack, clip=None):
|
||||
loras = lora_stack.get("loras", []) if isinstance(lora_stack, dict) else []
|
||||
per_block = coerce_bool(lora_stack.get("per_block", False)) if isinstance(lora_stack, dict) else False
|
||||
def load_loras(self, model, lora_list="[]", lora_stack=None, clip=None, per_block=False):
|
||||
# Both sources are honoured: a wired stack carries trained, prompt-specific weights,
|
||||
# the list is what the user typed here. Dropping either one silently would make a
|
||||
# filled-in field do nothing.
|
||||
stack = lora_stack if isinstance(lora_stack, dict) else {}
|
||||
own = stack_from_lora_list(lora_list, per_block)
|
||||
loras = list(stack.get("loras", [])) + own["loras"]
|
||||
per_block = coerce_bool(stack.get("per_block", False)) or coerce_bool(per_block)
|
||||
lora_stack = {**own, **stack, "per_block": per_block, "loras": loras}
|
||||
lines = [f"FunPack LoRA Loader | loading {len(loras)} LoRA(s)"]
|
||||
lines.append(f"Per-block application: {'enabled' if per_block else 'disabled'}")
|
||||
loaded_count = 0
|
||||
@@ -832,7 +957,7 @@ class FunPackLoraLoader:
|
||||
lora, lora_cache_key = self._load_lora_file(entry["name"])
|
||||
if self._per_block_supported(model, lora_stack, entry):
|
||||
patch_cache_key = (lora_cache_key, model_cache_key)
|
||||
loaded = self._load_model_lora_patches(model, lora, lora_cache_key, model_cache_key)
|
||||
loaded, fmt = self._load_model_lora_patches(model, lora, lora_cache_key, model_cache_key)
|
||||
profile = self._lora_block_profile(entry, loaded, model_weight, patch_cache_key)
|
||||
if profile is not None:
|
||||
item = {
|
||||
@@ -851,7 +976,7 @@ class FunPackLoraLoader:
|
||||
"loaded": loaded,
|
||||
"mode": "patch_global",
|
||||
"model_weight": model_weight,
|
||||
"status": "per-block fallback=global",
|
||||
"status": f"per-block fallback=global {fmt}",
|
||||
}
|
||||
)
|
||||
elif self._per_block_requested(entry, lora_stack):
|
||||
@@ -859,6 +984,7 @@ class FunPackLoraLoader:
|
||||
{
|
||||
"entry": entry,
|
||||
"lora": lora,
|
||||
"cache_key": lora_cache_key,
|
||||
"mode": "global",
|
||||
"model_weight": model_weight,
|
||||
"status": "per-block unsupported -> global",
|
||||
@@ -869,6 +995,7 @@ class FunPackLoraLoader:
|
||||
{
|
||||
"entry": entry,
|
||||
"lora": lora,
|
||||
"cache_key": lora_cache_key,
|
||||
"mode": "global",
|
||||
"model_weight": model_weight,
|
||||
"status": "global",
|
||||
@@ -893,8 +1020,12 @@ class FunPackLoraLoader:
|
||||
apply_status = item["status"]
|
||||
model, _, _ = self._apply_model_patches(model, item["loaded"], model_weight)
|
||||
else:
|
||||
apply_status = item["status"]
|
||||
model, clip = comfy.sd.load_lora_for_models(model, clip, item["lora"], model_weight, 0.0)
|
||||
# Same resolver as the per-block path: comfy's own converters, plus the
|
||||
# wrapper-prefix search, so an unusual format is not silently a no-op.
|
||||
loaded, fmt = self._load_model_lora_patches(
|
||||
model, item["lora"], item.get("cache_key"), model_cache_key)
|
||||
model, _, _ = self._apply_model_patches(model, loaded, model_weight)
|
||||
apply_status = f"{item['status']} {fmt}"
|
||||
|
||||
loaded_count += 1
|
||||
lines.append(
|
||||
|
||||
@@ -0,0 +1,201 @@
|
||||
"""FunPack's own loaders: dtype/attention plumbing, list inputs, LoRA format matching."""
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
import _comfy_stubs
|
||||
from _comfy_stubs import install_module
|
||||
|
||||
_comfy_stubs.install_all()
|
||||
|
||||
|
||||
class _CLIPType:
|
||||
LTXV = "LTXV"
|
||||
STABLE_DIFFUSION = "STABLE_DIFFUSION"
|
||||
|
||||
def __iter__(self):
|
||||
return iter([])
|
||||
|
||||
|
||||
install_module("comfy.sd", CLIPType=_CLIPType, VAE=object,
|
||||
load_clip=lambda **kw: kw, load_diffusion_model_state_dict=lambda *a, **k: None)
|
||||
install_module("comfy.lora_convert", convert_lora=lambda sd: sd)
|
||||
|
||||
|
||||
def _fake_load_lora(lora, to_load, log_missing=True):
|
||||
"""Only the shape that matters here: a base key matches when its lora_up/down exist."""
|
||||
return {to_load[x]: ("lora", x) for x in to_load if f"{x}.lora_up.weight" in lora}
|
||||
|
||||
|
||||
install_module("comfy.lora", load_lora=_fake_load_lora,
|
||||
model_lora_keys_unet=lambda model, key_map={}: dict(key_map),
|
||||
model_lora_keys_clip=lambda model, key_map={}: dict(key_map))
|
||||
|
||||
import loaders # noqa: E402
|
||||
import model_management as mm # noqa: E402
|
||||
import widgets # noqa: E402
|
||||
|
||||
FIELDS = [
|
||||
widgets.field("lora", "combo", choices=["None", "a.safetensors"], default="None"),
|
||||
widgets.field("type", "combo", choices=["general", "style"], default="general"),
|
||||
widgets.field("strength", "float", default=1.0),
|
||||
widgets.field("on", "boolean", default=True),
|
||||
]
|
||||
|
||||
|
||||
# ── list inputs ───────────────────────────────────────────────────────────────
|
||||
|
||||
def test_parse_rows_reads_the_json_string_a_widget_holds():
|
||||
rows = widgets.parse_rows('[{"lora": "a.safetensors", "strength": "0.5"}]', FIELDS, key="lora")
|
||||
assert rows == [{"index": 0, "lora": "a.safetensors", "type": "general",
|
||||
"strength": 0.5, "on": True}]
|
||||
|
||||
|
||||
def test_parse_rows_drops_unset_and_switched_off_rows_but_keeps_their_index():
|
||||
rows = widgets.parse_rows(
|
||||
[{"lora": "None"}, {"lora": "a.safetensors", "on": False}, {"lora": "a.safetensors"}],
|
||||
FIELDS, key="lora")
|
||||
assert [r["index"] for r in rows] == [2]
|
||||
|
||||
|
||||
def test_parse_rows_survives_a_value_that_is_not_a_list():
|
||||
assert widgets.parse_rows("not json", FIELDS) == []
|
||||
assert widgets.parse_rows(None, FIELDS) == []
|
||||
assert widgets.parse_rows({"lora": "a.safetensors"}, FIELDS, key="lora")[0]["lora"] == "a.safetensors"
|
||||
|
||||
|
||||
def test_list_widget_declares_its_row_shape_for_both_frontends():
|
||||
_, opts = widgets.list_widget("LoRA", FIELDS, add_label="+ Add LoRA")
|
||||
assert json.loads(opts["default"]) == []
|
||||
assert opts["funpack_list"]["add_label"] == "+ Add LoRA"
|
||||
assert [f["name"] for f in opts["funpack_list"]["fields"]] == ["lora", "type", "strength", "on"]
|
||||
|
||||
|
||||
# ── diffusion model loader ────────────────────────────────────────────────────
|
||||
|
||||
def test_weight_dtype_default_leaves_the_file_alone():
|
||||
assert loaders.weight_model_options("default") == {}
|
||||
|
||||
|
||||
def test_fp8_fast_asks_for_fp8_matmuls_as_well_as_fp8_storage():
|
||||
opts = loaders.weight_model_options("fp8_e4m3fn_fast")
|
||||
assert opts["fp8_optimizations"] is True
|
||||
assert opts["dtype"] is loaders.dtype_of("fp8_e4m3fn")
|
||||
|
||||
|
||||
def test_attention_choices_only_offers_backends_this_machine_registered(monkeypatch):
|
||||
attention = install_module("comfy.ldm.modules.attention")
|
||||
monkeypatch.setattr(attention, "REGISTERED_ATTENTION_FUNCTIONS", {"sage": lambda: None,
|
||||
"pytorch": lambda: None},
|
||||
raising=False)
|
||||
assert loaders.attention_choices() == ["default", "pytorch", "sage"]
|
||||
|
||||
|
||||
def test_attention_override_is_none_for_default_so_the_launch_flag_wins():
|
||||
assert loaders.attention_override("default") is None
|
||||
assert loaders.attention_override("") is None
|
||||
|
||||
|
||||
def test_attention_override_calls_the_backend_unwrapped(monkeypatch):
|
||||
attention = install_module("comfy.ldm.modules.attention")
|
||||
calls = []
|
||||
|
||||
def inner(*args, **kwargs):
|
||||
calls.append(args)
|
||||
return "sage-out"
|
||||
|
||||
wrapped = types.SimpleNamespace(__wrapped__=inner)
|
||||
monkeypatch.setattr(attention, "get_attention_function",
|
||||
lambda name, default=...: wrapped if name == "sage" else default,
|
||||
raising=False)
|
||||
|
||||
override = loaders.attention_override("sage")
|
||||
assert override(lambda *a, **k: "core-out", "q", "k", "v", 8) == "sage-out"
|
||||
assert calls == [("q", "k", "v", 8)]
|
||||
|
||||
|
||||
def test_attention_override_declines_a_backend_that_is_not_installed(monkeypatch):
|
||||
attention = install_module("comfy.ldm.modules.attention")
|
||||
monkeypatch.setattr(attention, "get_attention_function", lambda name, default=...: default,
|
||||
raising=False)
|
||||
assert loaders.attention_override("sage3") is None
|
||||
|
||||
|
||||
# ── LoRA: any format ──────────────────────────────────────────────────────────
|
||||
|
||||
class _FakeModel:
|
||||
def __init__(self, keys):
|
||||
self.model = types.SimpleNamespace(keys=keys)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ltx_model(monkeypatch):
|
||||
"""A model whose key map names the bare LTX form comfy produces."""
|
||||
key_map = {"transformer_blocks.0.attn1.to_q": "diffusion_model.transformer_blocks.0.attn1.to_q.weight"}
|
||||
monkeypatch.setattr(sys.modules["comfy.lora"], "model_lora_keys_unet",
|
||||
lambda model, km={}: dict(key_map))
|
||||
return _FakeModel(key_map)
|
||||
|
||||
|
||||
def _lora(*base_keys):
|
||||
out = {}
|
||||
for base in base_keys:
|
||||
out[f"{base}.lora_up.weight"] = 1
|
||||
out[f"{base}.lora_down.weight"] = 1
|
||||
return out
|
||||
|
||||
|
||||
def test_a_lora_already_in_the_models_own_naming_loads_as_is(ltx_model):
|
||||
patches, note = mm.resolve_lora_patches(ltx_model, _lora("transformer_blocks.0.attn1.to_q"))
|
||||
assert len(patches) == 1
|
||||
assert note == "keys=1 fmt=as-is"
|
||||
|
||||
|
||||
def test_a_diffusers_wrapped_lora_loads_once_the_wrapper_is_stripped(ltx_model):
|
||||
patches, note = mm.resolve_lora_patches(
|
||||
ltx_model, _lora("transformer.transformer_blocks.0.attn1.to_q"))
|
||||
assert len(patches) == 1
|
||||
assert note == "keys=1 fmt=stripped transformer."
|
||||
|
||||
|
||||
def test_a_peft_wrapped_lora_loads_too(ltx_model):
|
||||
patches, _ = mm.resolve_lora_patches(
|
||||
ltx_model, _lora("base_model.model.transformer_blocks.0.attn1.to_q"))
|
||||
assert len(patches) == 1
|
||||
|
||||
|
||||
def test_stripping_a_wrapper_keeps_the_keys_that_were_already_bare(ltx_model):
|
||||
# A half-converted file: one key bare, one still wrapped. Stripping must add the wrapped
|
||||
# one without losing the bare one — that is why the variant is SCORED, not just tried.
|
||||
ltx_model.model.keys["transformer_blocks.9.attn1.to_q"] = "diffusion_model.x.weight"
|
||||
patches, note = mm.resolve_lora_patches(ltx_model, _lora(
|
||||
"transformer_blocks.9.attn1.to_q", "transformer.transformer_blocks.0.attn1.to_q"))
|
||||
assert len(patches) == 2
|
||||
assert note == "keys=2 fmt=stripped transformer."
|
||||
|
||||
|
||||
def test_a_lora_for_a_different_model_family_says_so_instead_of_silently_doing_nothing(ltx_model):
|
||||
patches, note = mm.resolve_lora_patches(ltx_model, _lora("double_blocks.0.img_attn.qkv"))
|
||||
assert patches == {}
|
||||
assert note == "MATCHED NOTHING"
|
||||
|
||||
|
||||
# ── LoRA loader without a stack ───────────────────────────────────────────────
|
||||
|
||||
def test_the_lora_loader_no_longer_needs_a_second_node_wired_into_it():
|
||||
spec = mm.FunPackLoraLoader.INPUT_TYPES()
|
||||
assert "lora_list" in spec["required"]
|
||||
assert "lora_stack" in spec["optional"]
|
||||
assert "funpack_list" in spec["required"]["lora_list"][1]
|
||||
|
||||
|
||||
def test_a_list_becomes_the_same_stack_shape_apply_lora_weights_emits():
|
||||
stack = mm.stack_from_lora_list(
|
||||
'[{"lora": "a.safetensors", "type": "style", "strength": 0.8}, {"lora": "None"}]')
|
||||
assert [e["name"] for e in stack["loras"]] == ["a.safetensors"]
|
||||
entry = stack["loras"][0]
|
||||
assert entry["type"] == "style"
|
||||
assert entry["model_weight"] == entry["base_model_weight"] == 0.8
|
||||
assert entry["id"] == mm.lora_state_id("a.safetensors", "style")
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
"""Declarative list inputs: one STRING widget holding a JSON array of rows.
|
||||
|
||||
ComfyUI has no repeatable input, so a node that needs N of something (text encoders,
|
||||
LoRAs) carries ONE string widget whose value is a JSON array. The `funpack_list` entry
|
||||
in that widget's options describes a row, so both frontends can render it as rows with
|
||||
an Add button instead of raw JSON: the canvas widget in `web/funpack_list.js`, and the
|
||||
FunPack Editor's node page (`movie_editor/frontend/models.js`).
|
||||
|
||||
The node itself never has to care which frontend wrote the value — `parse_rows` accepts
|
||||
the JSON string, an already-parsed list, or a single row.
|
||||
"""
|
||||
import json
|
||||
|
||||
# Row field kinds a frontend must know how to draw.
|
||||
FIELD_KINDS = {"combo", "string", "int", "float", "boolean"}
|
||||
|
||||
|
||||
def field(name, kind, *, label=None, choices=None, default=None, min=None, max=None,
|
||||
step=None, tooltip=None, width=None):
|
||||
"""One column of a list row."""
|
||||
spec = {"name": name, "kind": kind, "label": label or name}
|
||||
if choices is not None:
|
||||
spec["choices"] = list(choices)
|
||||
if default is not None:
|
||||
spec["default"] = default
|
||||
for key, value in (("min", min), ("max", max), ("step", step),
|
||||
("tooltip", tooltip), ("width", width)):
|
||||
if value is not None:
|
||||
spec[key] = value
|
||||
return spec
|
||||
|
||||
|
||||
def list_widget(item, fields, *, default=None, add_label=None, max_rows=0, tooltip=None):
|
||||
"""An INPUT_TYPES entry for a list input. Returns the ("STRING", {...}) tuple."""
|
||||
spec = {"item": item, "fields": list(fields), "add_label": add_label or f"+ Add {item}"}
|
||||
if max_rows:
|
||||
spec["max_rows"] = int(max_rows)
|
||||
opts = {
|
||||
"default": json.dumps(default or []),
|
||||
"multiline": False,
|
||||
"funpack_list": spec,
|
||||
}
|
||||
if tooltip:
|
||||
opts["tooltip"] = tooltip
|
||||
return ("STRING", opts)
|
||||
|
||||
|
||||
def _coerce(value, kind, default=None):
|
||||
if kind == "boolean":
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return bool(value) if value is not None else bool(default)
|
||||
if kind in ("int", "float"):
|
||||
try:
|
||||
number = float(value)
|
||||
except (TypeError, ValueError):
|
||||
number = float(default or 0)
|
||||
if number != number or number in (float("inf"), float("-inf")):
|
||||
number = float(default or 0)
|
||||
return int(number) if kind == "int" else number
|
||||
if value is None:
|
||||
return default
|
||||
return value
|
||||
|
||||
|
||||
def parse_rows(value, fields, *, key=None, drop_empty=True):
|
||||
"""Normalize a list widget's value into a list of row dicts.
|
||||
|
||||
`key` names the field a row is worthless without (the file name, usually); rows whose
|
||||
key is missing, empty or "None" are dropped unless `drop_empty` is off. A row may also
|
||||
carry "on": False to be skipped without deleting it.
|
||||
"""
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = json.loads(value or "[]")
|
||||
except (TypeError, ValueError, json.JSONDecodeError):
|
||||
return []
|
||||
if isinstance(value, dict):
|
||||
value = [value]
|
||||
if not isinstance(value, list):
|
||||
return []
|
||||
|
||||
specs = {f["name"]: f for f in fields}
|
||||
if key is None and fields:
|
||||
key = fields[0]["name"]
|
||||
|
||||
rows = []
|
||||
for index, raw in enumerate(value):
|
||||
if not isinstance(raw, dict):
|
||||
continue
|
||||
if not _coerce(raw.get("on", True), "boolean", True):
|
||||
continue
|
||||
row = {"index": index}
|
||||
for name, spec in specs.items():
|
||||
row[name] = _coerce(raw.get(name), spec["kind"], spec.get("default"))
|
||||
if drop_empty and key:
|
||||
k = row.get(key)
|
||||
if k is None or (isinstance(k, str) and k.strip() in ("", "None")):
|
||||
continue
|
||||
rows.append(row)
|
||||
return rows
|
||||
Reference in New Issue
Block a user