diff --git a/__init__.py b/__init__.py index ebe1a1d9..ffa8d52d 100644 --- a/__init__.py +++ b/__init__.py @@ -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", } diff --git a/loaders.py b/loaders.py new file mode 100644 index 00000000..989a973b --- /dev/null +++ b/loaders.py @@ -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}") diff --git a/model_management.py b/model_management.py index 2685941a..fa7af3dc 100644 --- a/model_management.py +++ b/model_management.py @@ -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( diff --git a/tests/test_loaders.py b/tests/test_loaders.py new file mode 100644 index 00000000..05bc277f --- /dev/null +++ b/tests/test_loaders.py @@ -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") diff --git a/widgets.py b/widgets.py new file mode 100644 index 00000000..fe39357e --- /dev/null +++ b/widgets.py @@ -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