From 3c7fd499b67bfaef367a2e8313c8aa52663eb7a9 Mon Sep 17 00:00:00 2001 From: BobsBlazed <39865461+BobsBlazed@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:04:26 -0400 Subject: [PATCH] Add a Universal loader covering every supported architecture (#5) The FLUX and SDXL loaders know their architecture's block names up front, which does not scale across ComfyUI's ~99 model configs. The new Universal node discovers a model's block layout at runtime instead: stack names, stack sizes and their execution order all come from the loaded model, so pruned, distilled and brand-new architectures work without a code change. Concatenating the discovered stacks puts every block on one normalised depth axis, split into five buckets, alongside embeddings, output head and text encoder. Covers SD1.5/2/3, SDXL, FLUX, Chroma, AuraFlow, PixArt, HiDream, Qwen-Image, Wan, LTX-Video, Mochi, HunyuanVideo/DiT, Lumina and Cosmos. Verified against real ComfyUI: models built from ComfyUI's own configs on the meta device plus its *_to_diffusers key tables, 8000+ authentic state-dict keys across 11 architectures, none falling through to "Other Tensors". That pass found and fixed real gaps: SD3's negative "joint_blocks.-1" index, Lumina's noise_refiner/context_refiner stacks, AuraFlow's native double_layers naming, several conditioning embedders, and FLUX's ControlNet pos_embed_input. No widget was added, removed or reordered on the FLUX or SDXL nodes, so existing workflows are unaffected. Bumps version to 1.2.0. --- README.md | 86 ++++++++ __init__.py | 2 + bobs_blocks.py | 1 + bobs_lora_loader.py | 79 +++++-- bobs_universal.py | 402 +++++++++++++++++++++++++++++++++++ pyproject.toml | 4 +- tests/test_bobs_universal.py | 324 ++++++++++++++++++++++++++++ 7 files changed, 884 insertions(+), 14 deletions(-) create mode 100644 bobs_universal.py create mode 100644 tests/test_bobs_universal.py diff --git a/README.md b/README.md index 61e1288..0c38370 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ This node allows you to go beyond a single strength slider and specify different ## Features +- **Any Architecture**: A third **Universal** loader covers everything else ComfyUI supports — SD1.5, SD2, SD3/3.5, Chroma, AuraFlow, PixArt, HiDream, Qwen-Image, Wan, LTX-Video, Mochi, HunyuanVideo/DiT, Lumina, Cosmos and more. It discovers the model's block layout at runtime instead of using a hard-coded table, so new and pruned architectures work without an update. - **Dual Model Support**: Separate, optimized loaders for `SDXL` and `FLUX` models, each tailored to the architecture's specific blocks. - **Granular Block-Level Control**: Fine-tune the strength of a LoRA on different conceptual parts of the diffusion model. - **Intelligent Presets**: Comes with pre-configured presets for common use cases like `Character`, `Style`, `Concept`, `Detail & Texture` and `Fix Hands/Anatomy`. @@ -95,6 +96,50 @@ Ranges below are for the canonical FLUX.1 geometry (19 double-stream blocks, 38 | Output Blocks | `output_blocks.*` | | Other Tensors | `time_embed`, `label_emb`, `out.*` | +### Universal + +The Universal loader works on a single normalised **depth axis**. Almost every +diffusion backbone is one or more ordered stacks of repeated blocks: + +``` +UNet (SD1.5 / SDXL) input_blocks.N -> middle_block -> output_blocks.N +Dual-stream DiT (FLUX) double_blocks.N -> single_blocks.N +HiDream double_stream_blocks.N -> single_stream_blocks.N +MMDiT (SD3) joint_blocks.N +AuraFlow double_layers.N -> single_layers.N +Qwen-Image / LTX-Video transformer_blocks.N +Wan / Mochi / PixArt blocks.N +Lumina noise_refiner.N -> context_refiner.N -> layers.N +``` + +Those stacks are **discovered from the loaded model**, concatenated in execution +order, and every block gets a position from 0 to 1 along the result. That axis is +split into five buckets, so the same five sliders mean the same thing on a +19+38-block FLUX, a 60-block Qwen-Image and a 20-stage SDXL UNet. + +| Block | Covers | +|---|---| +| Text Encoder | Every text-encoder weight (CLIP / T5 / LLM) | +| Input & Embeddings | Patch, timestep, guidance and context embedders | +| Early Blocks (Composition) | First 20% of the stack | +| Early-Mid Blocks (Subject) | 20–40% | +| Mid Blocks (Concept & Style) | 40–60% | +| Late-Mid Blocks (Detail) | 60–80% | +| Late Blocks (Texture) | Final 20% | +| Output Head | Final projection back to latent space | +| Other Tensors | Anything unmatched (normally empty) | + +The `info` output names the detected architecture and the discovered stacks, e.g. + +``` +[UNIVERSAL] mylora.safetensors (preset: Style) +architecture: QwenImage transformer_blocks[60] (total 60) +``` + +Use the dedicated FLUX or SDXL loader when you want that architecture's named +blocks; use Universal for everything else, or when you want one node whose +sliders behave consistently across models. + ## Why Use Block-Weighted LoRA? A single LoRA file often contains training for multiple concepts (e.g. a character's face, their clothing, and the overall artistic style). A standard LoRA loader applies the LoRA with one uniform strength across the entire model. @@ -115,6 +160,47 @@ python -m unittest discover -s tests -v ## Changelog +### 1.2.0 + +**New: a Universal loader covering every architecture ComfyUI supports.** + +- **`Bobs LoRA Loader (Universal)`** handles SD1.5, SD2, SDXL, SD3/3.5, FLUX, + Chroma, AuraFlow, PixArt, HiDream, Qwen-Image, Wan, LTX-Video, Mochi, + HunyuanVideo/DiT, Lumina, Cosmos and anything else built as stacks of + repeated blocks. The block layout is *discovered from the loaded model* — + stack names, stack sizes and their execution order — rather than read from a + per-family table, so pruned, distilled and brand-new architectures work + without a code change. +- Weights are assigned along a normalised depth axis (Early → Late), plus + embeddings, output head and text encoder, so the same sliders mean the same + thing across very different models. +- The `info` output now reports the detected architecture and the discovered + stacks, e.g. `architecture: QwenImage transformer_blocks[60] (total 60)`. + The FLUX and SDXL loaders report their architecture too. + +**Verified against real ComfyUI.** The classification logic was run against +models built from ComfyUI's own configs and against its `*_to_diffusers` key +tables, covering SD1.5, SDXL, SD3, FLUX (full and pruned geometry), AuraFlow, +PixArt, LTX-Video, Lumina, Qwen-Image and Wan — 8,000+ authentic state-dict +keys, all classified with none falling through to `Other Tensors`. That pass +found and fixed several real gaps that synthetic fixtures had missed: + +- SD3 addresses its final block as `joint_blocks.-1`; negative indices are now + resolved against the stack size instead of failing to match. +- Lumina's `noise_refiner` / `context_refiner` stacks are recognised and + ordered ahead of the main `layers` stack. +- AuraFlow's native `double_layers` / `single_layers` names are handled, not + just the diffusers spelling. +- Conditioning embedders that previously fell through — PixArt's `ar_embedder`, + `csize_embedder` and `t_block`, LTX-Video's `adaln_single` and + `scale_shift_table`, Qwen-Image's `txt_norm`, Wan's `time_projection`, + AuraFlow's `cond_seq_linear` / `positional_encoding` — now land in + `Input & Embeddings`. +- FLUX's ControlNet `pos_embed_input` now maps to `Image Hint`. + +Existing FLUX and SDXL workflows are unaffected: no widget was added, removed or +reordered on those two nodes. + ### 1.1.0 **Block weighting now actually works.** This release fixes a defect that made the per-block sliders unreliable for every LoRA. diff --git a/__init__.py b/__init__.py index 066be93..2e2a432 100644 --- a/__init__.py +++ b/__init__.py @@ -7,6 +7,7 @@ from .bobs_lora_loader import ( NODE_DISPLAY_NAME_MAPPINGS, BobsLoraLoaderFlux, BobsLoraLoaderSdxl, + BobsLoraLoaderUniversal, ) __all__ = [ @@ -14,6 +15,7 @@ __all__ = [ "NODE_DISPLAY_NAME_MAPPINGS", "BobsLoraLoaderFlux", "BobsLoraLoaderSdxl", + "BobsLoraLoaderUniversal", ] logging.getLogger("BobsLoraLoader").info( diff --git a/bobs_blocks.py b/bobs_blocks.py index fa204a4..172d709 100644 --- a/bobs_blocks.py +++ b/bobs_blocks.py @@ -340,6 +340,7 @@ _FLUX_TOKEN_MAP: Sequence[Tuple[str, str]] = ( ("_context_embedder", FLUX_TEXT_CONDITIONING), # diffusers name for txt_in ("_img_in", FLUX_IMAGE_HINT), ("_x_embedder", FLUX_IMAGE_HINT), # diffusers name for img_in + ("_pos_embed_input", FLUX_IMAGE_HINT), # FLUX ControlNet hint input ("_final_layer", FLUX_FINAL), ("_proj_out", FLUX_FINAL), # diffusers name for final_layer ("_norm_out", FLUX_FINAL), diff --git a/bobs_lora_loader.py b/bobs_lora_loader.py index 8de16a5..58d75b6 100644 --- a/bobs_lora_loader.py +++ b/bobs_lora_loader.py @@ -37,6 +37,12 @@ from .bobs_blocks import ( flux_block_ranges, resolve_block_strengths, ) +from .bobs_universal import ( # registers the UNIVERSAL family on import + ALL_UNIVERSAL_BLOCKS, + UNIVERSAL_TOOLTIPS, + classify_universal_key, + discover_layout, +) try: # Added to ComfyUI in 2025; handles BFL-control / Wan-Fun / USO LoRAs. import comfy.lora_convert as _lora_convert @@ -84,6 +90,14 @@ def _build_key_maps(model, clip) -> Tuple[Dict[str, Any], set]: return key_map, unet_targets +def _architecture_name(model) -> str: + """Best-effort name of the loaded backbone, for the report.""" + try: + return type(model.model).__name__ + except Exception: # noqa: BLE001 + return "unknown" + + def _flux_geometry(model) -> Tuple[int, int]: """Read the double/single stream depths off the loaded FLUX model.""" try: @@ -101,9 +115,12 @@ def _format_report(tag: str, blocks: List[str], grouped: Dict[str, Dict[Any, Any]], strengths: Dict[str, float], - applied: Dict[str, int]) -> str: - lines = [f"[{tag}] {lora_name} (preset: {preset})", - f"{'block':<40} {'weight':>7} {'found':>7} {'applied':>8}"] + applied: Dict[str, int], + detail: str = "") -> str: + lines = [f"[{tag}] {lora_name} (preset: {preset})"] + if detail: + lines.append(detail) + lines.append(f"{'block':<40} {'weight':>7} {'found':>7} {'applied':>8}") for name in blocks: lines.append( f"{name:<40} {strengths.get(name, 0.0):>7.2f} " @@ -209,8 +226,8 @@ class _BobsLoraLoaderBase: # ------------------------------------------------------------ classifier -- - def _classifier(self, model): - """Return ``fn(model_state_dict_key) -> block name`` for this family.""" + def _classifier(self, model, unet_targets): + """Return ``(fn(model_key) -> block name, detail_line)`` for this family.""" raise NotImplementedError # ----------------------------------------------------------------- apply -- @@ -253,7 +270,7 @@ class _BobsLoraLoaderBase: # Group UNet patches per conceptual block; CLIP patches all share the # text-encoder block. - classify = self._classifier(model) + classify, detail = self._classifier(model, unet_targets) text_encoder_block = TEXT_ENCODER_BLOCK[self.FAMILY] grouped: Dict[str, Dict[Any, Any]] = {name: {} for name in self.BLOCKS} @@ -278,7 +295,7 @@ class _BobsLoraLoaderBase: # Any tensor in the file that ComfyUI could not place is reported by # comfy.lora.load_lora itself as a "lora key not loaded" warning. report = _format_report(tag, lora_name, preset, self.BLOCKS, - grouped, strengths, applied) + grouped, strengths, applied, detail) self.logger.info("\n%s", report) _explain_empty_blocks(tag, self.BLOCKS, grouped, strengths) @@ -303,9 +320,12 @@ class BobsLoraLoaderFlux(_BobsLoraLoaderBase): def INPUT_TYPES(cls): return cls._input_types() - def _classifier(self, model): - ranges = flux_block_ranges(*_flux_geometry(model)) - return lambda key: classify_flux_key(key, ranges) + def _classifier(self, model, unet_targets): + depth, depth_single = _flux_geometry(model) + ranges = flux_block_ranges(depth, depth_single) + detail = (f"architecture: {_architecture_name(model)} " + f"double_blocks[{depth}] -> single_blocks[{depth_single}]") + return (lambda key: classify_flux_key(key, ranges)), detail # -----------------------------------------------------------------------------# @@ -324,8 +344,41 @@ class BobsLoraLoaderSdxl(_BobsLoraLoaderBase): def INPUT_TYPES(cls): return cls._input_types() - def _classifier(self, model): - return classify_sdxl_key + def _classifier(self, model, unet_targets): + return classify_sdxl_key, f"architecture: {_architecture_name(model)}" + + +# -----------------------------------------------------------------------------# +# UNIVERSAL LOADER # +# -----------------------------------------------------------------------------# + +class BobsLoraLoaderUniversal(_BobsLoraLoaderBase): + FAMILY = "UNIVERSAL" + BLOCKS = ALL_UNIVERSAL_BLOCKS + DESCRIPTION = ( + "Block-weighted LoRA loading for any architecture ComfyUI supports — " + "SD1.5, SD2, SDXL, SD3/3.5, FLUX, Chroma, AuraFlow, PixArt, HiDream, " + "Qwen-Image, Wan, LTX-Video, Mochi, HunyuanVideo/DiT, Lumina, Cosmos and " + "others. The block stacks are discovered from the loaded model itself " + "rather than a hard-coded table, so new and pruned architectures work " + "without an update. Weights are assigned by depth: Early through Late " + "across the whole block stack, plus embeddings, output head and text " + "encoder." + ) + + @classmethod + def INPUT_TYPES(cls): + return cls._input_types() + + def _classifier(self, model, unet_targets): + layout = discover_layout(unet_targets) + detail = f"architecture: {_architecture_name(model)} {layout.describe()}" + if not layout: + self.logger.warning( + "[UNIVERSAL] No block stacks recognised in this model; only the " + "embedding, output-head and text-encoder weights can be targeted." + ) + return (lambda key: classify_universal_key(key, layout)), detail # -----------------------------------------------------------------------------# @@ -335,9 +388,11 @@ class BobsLoraLoaderSdxl(_BobsLoraLoaderBase): NODE_CLASS_MAPPINGS = { "BobsLoraLoaderFlux": BobsLoraLoaderFlux, "BobsLoraLoaderSdxl": BobsLoraLoaderSdxl, + "BobsLoraLoaderUniversal": BobsLoraLoaderUniversal, } NODE_DISPLAY_NAME_MAPPINGS = { "BobsLoraLoaderFlux": "Bobs LoRA Loader (FLUX)", "BobsLoraLoaderSdxl": "Bobs LoRA Loader (SDXL)", + "BobsLoraLoaderUniversal": "Bobs LoRA Loader (Universal)", } diff --git a/bobs_universal.py b/bobs_universal.py new file mode 100644 index 0000000..7b5e996 --- /dev/null +++ b/bobs_universal.py @@ -0,0 +1,402 @@ +""" +Architecture-agnostic block analysis for Bobs LoRA Loader. + +The FLUX and SDXL loaders in :mod:`bobs_blocks` know their architecture's block +names up front. That does not scale: ComfyUI ships close to a hundred model +configs, and new ones land regularly. Rather than maintain a table per family, +this module *discovers* a model's block layout from the model's own key set at +runtime, then buckets each key by how deep it sits in the execution order. + +The observation that makes this work: essentially every diffusion backbone is +one or more ordered stacks of repeated blocks, whatever they are named — + + UNet (SD1.5 / SDXL) input_blocks.N -> middle_block -> output_blocks.N + Dual-stream DiT (FLUX) double_blocks.N -> single_blocks.N + HiDream double_stream_blocks.N -> single_stream_blocks.N + MMDiT (SD3) joint_blocks.N + AuraFlow joint_transformer_blocks.N -> single_transformer_blocks.N + Qwen-Image / LTXV transformer_blocks.N + Wan / Mochi blocks.N + Lumina layers.N + +Concatenating those stacks in execution order gives every block a position on a +single 0..1 depth axis, which is then split into five buckets. Anything outside +a stack is an embedding/input layer, an output head, or genuinely unrecognised. + +Because the stack sizes come from the model rather than a constant, this adapts +to pruned, distilled and brand-new architectures without a code change. + +Imports nothing from ComfyUI or torch; see ``tests/test_bobs_universal.py``. +""" + +import re +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +from .bobs_blocks import ( + ALL_BLOCKS, + BLOCK_TOOLTIPS, + LORA_BLOCK_PRESETS, + TEXT_ENCODER_BLOCK, + normalize_key, +) + +# -----------------------------------------------------------------------------# +# BLOCK NAMES # +# -----------------------------------------------------------------------------# +# +# Widget values serialise positionally — only ever append to this list. + +UNIVERSAL_TEXT_ENCODER = "Text Encoder" +UNIVERSAL_INPUT = "Input & Embeddings" +UNIVERSAL_EARLY = "Early Blocks (Composition)" +UNIVERSAL_EARLY_MID = "Early-Mid Blocks (Subject)" +UNIVERSAL_MID = "Mid Blocks (Concept & Style)" +UNIVERSAL_LATE_MID = "Late-Mid Blocks (Detail)" +UNIVERSAL_LATE = "Late Blocks (Texture)" +UNIVERSAL_OUTPUT = "Output Head" +UNIVERSAL_OTHER = "Other Tensors" + +ALL_UNIVERSAL_BLOCKS: List[str] = [ + UNIVERSAL_TEXT_ENCODER, + UNIVERSAL_INPUT, + UNIVERSAL_EARLY, + UNIVERSAL_EARLY_MID, + UNIVERSAL_MID, + UNIVERSAL_LATE_MID, + UNIVERSAL_LATE, + UNIVERSAL_OUTPUT, + UNIVERSAL_OTHER, +] + +# Depth buckets, as upper-exclusive fractions of the concatenated block stacks. +_DEPTH_BUCKETS: Sequence[Tuple[float, str]] = ( + (0.2, UNIVERSAL_EARLY), + (0.4, UNIVERSAL_EARLY_MID), + (0.6, UNIVERSAL_MID), + (0.8, UNIVERSAL_LATE_MID), + (1.01, UNIVERSAL_LATE), +) + +UNIVERSAL_TOOLTIPS: Dict[str, str] = { + UNIVERSAL_TEXT_ENCODER: "All text-encoder weights (CLIP / T5 / LLM), whatever the model uses.", + UNIVERSAL_INPUT: "Patch, timestep, guidance and context embedding layers feeding the block stack.", + UNIVERSAL_EARLY: "First 20% of the block stack. Global composition and layout.", + UNIVERSAL_EARLY_MID: "20–40% of the block stack. Subject identity.", + UNIVERSAL_MID: "40–60% of the block stack. Concept and dominant style.", + UNIVERSAL_LATE_MID: "60–80% of the block stack. Detail generation.", + UNIVERSAL_LATE: "Final 20% of the block stack. Fine texture and surface.", + UNIVERSAL_OUTPUT: "Final projection back to latent space.", + UNIVERSAL_OTHER: "Tensors that matched no stack, embedding or head pattern.", +} + +# -----------------------------------------------------------------------------# +# STACK DISCOVERY # +# -----------------------------------------------------------------------------# + +# Execution order of the stack names we know about. A model only ever uses a +# few of these; the order here is what puts them on one axis correctly. +# Unknown stack names still work — they sort after these, alphabetically. +STACK_ORDER: Tuple[str, ...] = ( + # UNet + "input_blocks", + "middle_block", + "output_blocks", + # Dual-stream DiT: FLUX, HunyuanVideo, Chroma + "double_blocks", + "single_blocks", + # HiDream + "double_stream_blocks", + "single_stream_blocks", + # MMDiT: SD3 / SD3.5 + "joint_blocks", + # Lumina: input refiners run ahead of the main stack + "noise_refiner", + "context_refiner", + # AuraFlow (ComfyUI names them *_layers; the diffusers names also appear) + "double_layers", + "single_layers", + "joint_transformer_blocks", + "single_transformer_blocks", + # Qwen-Image, LTX-Video, PixArt, Cosmos + "transformer_blocks", + # Wan, Mochi, Kandinsky + "blocks", + # Lumina, generic transformer stacks + "layers", +) + +# Stack matching runs on the RAW dotted key, never the underscore-normalised +# one. Normalising would erase the separator that tells "mystery_blocks.0" +# (stack "mystery_blocks") apart from a stack literally named "blocks" — both +# collapse to "..._blocks_0". The dot boundary is the whole signal. +# +# No list of known names is needed to *match*: a repeated stack is always a +# module path component ending in "blocks"/"layers" followed by an index. +# STACK_ORDER below is used only to order the stacks once discovered. +# The name prefix is optional so a stack named exactly "blocks" or "layers" +# (Wan, Mochi, PixArt, Lumina) matches just as well as "single_stream_blocks". +# "refiner" is included for Lumina's noise_refiner/context_refiner stacks. +# The index may be negative: ComfyUI's SD3 key map addresses the final joint +# block as "joint_blocks.-1", which is resolved against the stack size below. +_STACK_RE = re.compile( + r"(?:^|\.)((?:[A-Za-z][A-Za-z0-9_]*)?(?:blocks|layers|refiner))\.(-?\d+)(?=\.|$)", + re.IGNORECASE, +) +# ``middle_block`` is a single module rather than an indexed list. +_MIDDLE_RE = re.compile(r"(?:^|\.)(middle_block)(?=\.|$)", re.IGNORECASE) + +# Tokens that identify a non-stack tensor. Checked in order, output before +# input, because some head names also contain an embedding-ish word. +_OUTPUT_TOKENS: Tuple[str, ...] = ( + "_final_layer", + "_final_linear", + "_proj_out", + "_norm_out", + "_unpatchify", + "_final_norm", + "_head_", +) +_OUTPUT_RE = re.compile(r"_out_\d+(?=_|$)") # SD "out.0" / "out.2" head + +_INPUT_TOKENS: Tuple[str, ...] = ( + "_img_in", + "_txt_in", + "_x_embedder", + "_patch_embed", + "_patchify", + "_time_in", + "_time_embed", + "_time_text_embed", + "_timestep_embed", + "_t_embedder", + "_guidance_in", + "_vector_in", + "_context_embedder", + "_caption_projection", + "_y_embedder", + "_label_emb", + "_pos_embed", + "_text_embedding", + "_condition_embedder", + "_img_emb", + "_register_tokens", + "_cap_embedder", + "_input_hint", + "_init_x_linear", # AuraFlow + "_cond_seq_linear", # AuraFlow text conditioning + "_positional_encoding", + "_modf", # AuraFlow final modulation + "_adaln_single", # LTX-Video / PixArt shared modulation + "_t_block", # PixArt timestep modulation + "_scale_shift_table", + "_txt_norm", # Qwen-Image text stream norm + "_time_projection", # Wan + "_rope", + "_freqs", + # Generic catch-all, last: covers *_embedder / *_embedding / *_embed names + # we have not enumerated (PixArt's ar_embedder and csize_embedder, and + # whatever the next architecture calls its conditioning embedders). + "_embed", +) + + +def _stack_match(model_key: str) -> Optional[Tuple[str, int]]: + """Return ``(stack_name, index)`` for the outermost stack in a raw key. + + The *leftmost* match is deliberate, and matters twice over: + + - an SDXL key ``input_blocks.4.1.transformer_blocks.0.attn1.to_q.weight`` + belongs to ``input_blocks``, not the attention module's inner + ``transformer_blocks``; + - ``middle_block.1.transformer_blocks.0...`` belongs to ``middle_block``, + which is why the two patterns are compared by position rather than + tried in a fixed order. + """ + stack = _STACK_RE.search(model_key) + middle = _MIDDLE_RE.search(model_key) + + if stack and middle: + return ((stack.group(1), int(stack.group(2))) if stack.start() < middle.start() + else ("middle_block", 0)) + if stack: + return stack.group(1), int(stack.group(2)) + if middle: + return "middle_block", 0 + return None + + +def _stack_sort_key(name: str) -> Tuple[int, str]: + try: + return (STACK_ORDER.index(name), "") + except ValueError: + return (len(STACK_ORDER), name) + + +class BlockLayout: + """The block-stack layout discovered from one model's key set.""" + + def __init__(self, stacks: Dict[str, int]): + #: stack name -> number of blocks, in execution order + self.stacks: Dict[str, int] = { + name: stacks[name] for name in sorted(stacks, key=_stack_sort_key) + } + self.offsets: Dict[str, int] = {} + running = 0 + for name, size in self.stacks.items(): + self.offsets[name] = running + running += size + self.total: int = running + + def __bool__(self) -> bool: + return self.total > 0 + + def describe(self) -> str: + if not self.stacks: + return "no block stacks detected" + parts = [f"{name}[{size}]" for name, size in self.stacks.items()] + return " -> ".join(parts) + f" (total {self.total})" + + def depth_fraction(self, stack: str, index: int) -> Optional[float]: + """Position of a block on the 0..1 depth axis, centred in its slot.""" + if not self.total or stack not in self.offsets: + return None + size = self.stacks[stack] + if index < 0: # Python-style: "joint_blocks.-1" is the last block. + index += size + # Guard against an index beyond what discovery saw. + index = min(max(index, 0), max(size - 1, 0)) + return (self.offsets[stack] + index + 0.5) / self.total + + +def discover_layout(model_keys: Iterable[str]) -> BlockLayout: + """Infer the block-stack layout from a model's state-dict key names.""" + sizes: Dict[str, int] = {} + for key in model_keys: + found = _stack_match(key) + if found is None: + continue + name, index = found + if index + 1 > sizes.get(name, 0): + sizes[name] = index + 1 + return BlockLayout(sizes) + + +# -----------------------------------------------------------------------------# +# CLASSIFICATION # +# -----------------------------------------------------------------------------# + +def classify_universal_key(model_key: str, layout: BlockLayout) -> str: + """Bucket a UNet/DiT state-dict key using a discovered :class:`BlockLayout`.""" + nk = normalize_key(model_key) + + found = _stack_match(model_key) + if found is not None: + fraction = layout.depth_fraction(*found) + if fraction is not None: + for upper, name in _DEPTH_BUCKETS: + if fraction < upper: + return name + return UNIVERSAL_LATE + + for token in _OUTPUT_TOKENS: + if token in nk: + return UNIVERSAL_OUTPUT + if _OUTPUT_RE.search(nk): + return UNIVERSAL_OUTPUT + + for token in _INPUT_TOKENS: + if token in nk: + return UNIVERSAL_INPUT + + return UNIVERSAL_OTHER + + +# -----------------------------------------------------------------------------# +# PRESETS # +# -----------------------------------------------------------------------------# +# +# Registered into the shared tables in bobs_blocks so resolve_block_strengths +# treats "UNIVERSAL" exactly like the two hand-tuned families. + +UNIVERSAL_PRESETS = { + "Custom": {}, + "Full (Normal LoRA)": { + "strength": 1.0, + "block_weights": {name: 1.0 for name in ALL_UNIVERSAL_BLOCKS}, + }, + "Character": { + "strength": 1.0, + "block_weights": { + UNIVERSAL_TEXT_ENCODER: 1.0, + UNIVERSAL_INPUT: 1.0, + UNIVERSAL_EARLY: 0.8, + UNIVERSAL_EARLY_MID: 1.0, + UNIVERSAL_MID: 1.0, + UNIVERSAL_LATE_MID: 0.2, + UNIVERSAL_LATE: 0.0, + UNIVERSAL_OUTPUT: 0.0, + UNIVERSAL_OTHER: 1.0, + }, + }, + "Style": { + "strength": 1.0, + "block_weights": { + UNIVERSAL_TEXT_ENCODER: 0.2, + UNIVERSAL_INPUT: 1.0, + UNIVERSAL_EARLY: 0.1, + UNIVERSAL_EARLY_MID: 0.0, + UNIVERSAL_MID: 0.5, + UNIVERSAL_LATE_MID: 1.0, + UNIVERSAL_LATE: 1.0, + UNIVERSAL_OUTPUT: 1.0, + UNIVERSAL_OTHER: 1.0, + }, + }, + "Concept": { + "strength": 1.0, + "block_weights": { + UNIVERSAL_TEXT_ENCODER: 1.0, + UNIVERSAL_INPUT: 1.0, + UNIVERSAL_EARLY: 1.0, + UNIVERSAL_EARLY_MID: 0.9, + UNIVERSAL_MID: 0.7, + UNIVERSAL_LATE_MID: 0.4, + UNIVERSAL_LATE: 0.2, + UNIVERSAL_OUTPUT: 0.0, + UNIVERSAL_OTHER: 1.0, + }, + }, + "Detail & Texture": { + "strength": 1.0, + "block_weights": { + UNIVERSAL_TEXT_ENCODER: 0.0, + UNIVERSAL_INPUT: 1.0, + UNIVERSAL_EARLY: 0.0, + UNIVERSAL_EARLY_MID: 0.0, + UNIVERSAL_MID: 0.2, + UNIVERSAL_LATE_MID: 1.0, + UNIVERSAL_LATE: 1.0, + UNIVERSAL_OUTPUT: 1.0, + UNIVERSAL_OTHER: 0.0, + }, + }, + "Fix Hands/Anatomy": { + "strength": 0.4, + "block_weights": { + UNIVERSAL_TEXT_ENCODER: 0.2, + UNIVERSAL_INPUT: 1.0, + UNIVERSAL_EARLY: 1.0, + UNIVERSAL_EARLY_MID: 0.5, + UNIVERSAL_MID: 0.0, + UNIVERSAL_LATE_MID: 0.0, + UNIVERSAL_LATE: 0.0, + UNIVERSAL_OUTPUT: 0.0, + UNIVERSAL_OTHER: 0.0, + }, + }, +} + +LORA_BLOCK_PRESETS["UNIVERSAL"] = UNIVERSAL_PRESETS +ALL_BLOCKS["UNIVERSAL"] = ALL_UNIVERSAL_BLOCKS +TEXT_ENCODER_BLOCK["UNIVERSAL"] = UNIVERSAL_TEXT_ENCODER +BLOCK_TOOLTIPS.update(UNIVERSAL_TOOLTIPS) diff --git a/pyproject.toml b/pyproject.toml index c738ac8..9f57908 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "Bobs_LoRA_Loader" -description = "A custom LoRA loader node for ComfyUI with advanced block-weighting controls for both SDXL and FLUX models. Features presets for common use-cases like 'Character' and 'Style', and a 'Custom' mode for fine-grained control over individual model blocks." -version = "1.1.0" +description = "Block-weighted LoRA loaders for ComfyUI. Dedicated SDXL and FLUX nodes plus a Universal node that discovers any supported architecture's block layout at runtime (SD1.5/2/3, Chroma, AuraFlow, PixArt, HiDream, Qwen-Image, Wan, LTX-Video, Mochi, HunyuanVideo, Lumina, Cosmos and more). Presets for common use-cases like 'Character' and 'Style', plus a 'Custom' mode for per-block control." +version = "1.2.0" license = {file = "LICENSE"} [project.urls] diff --git a/tests/test_bobs_universal.py b/tests/test_bobs_universal.py new file mode 100644 index 0000000..419800b --- /dev/null +++ b/tests/test_bobs_universal.py @@ -0,0 +1,324 @@ +"""Unit tests for the architecture-agnostic block analysis. + +Runs without ComfyUI or torch: + + python -m unittest discover -s tests -v +""" + +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +# The package modules use relative imports; expose them under a package name. +import importlib.util # noqa: E402 +import types # noqa: E402 + +_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +_pkg = types.ModuleType("bobs_pkg") +_pkg.__path__ = [_REPO] +sys.modules.setdefault("bobs_pkg", _pkg) + + +def _load(name): + spec = importlib.util.spec_from_file_location( + f"bobs_pkg.{name}", os.path.join(_REPO, f"{name}.py")) + module = importlib.util.module_from_spec(spec) + sys.modules[f"bobs_pkg.{name}"] = module + spec.loader.exec_module(module) + return module + + +bobs_blocks = _load("bobs_blocks") +bu = _load("bobs_universal") + + +# --------------------------------------------------------------------------- # +# Representative state-dict key sets, one per architecture family. +# Names follow ComfyUI's canonical "diffusion_model.*" naming. +# --------------------------------------------------------------------------- # + +def _unet_sd(n_in=9, n_out=9): + keys = ["diffusion_model.time_embed.0.weight", + "diffusion_model.label_emb.0.0.weight", + "diffusion_model.out.2.weight"] + for i in range(n_in): + keys.append(f"diffusion_model.input_blocks.{i}.1.transformer_blocks.0.attn1.to_q.weight") + keys.append("diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_k.weight") + for i in range(n_out): + keys.append(f"diffusion_model.output_blocks.{i}.1.transformer_blocks.0.attn1.to_v.weight") + return keys + + +def _flat_stack(stack, count, extra=()): + keys = [f"diffusion_model.{stack}.{i}.attn.to_q.weight" for i in range(count)] + return keys + [f"diffusion_model.{k}" for k in extra] + + +def _dual_stack(first, second, n1, n2, extra=()): + keys = [f"diffusion_model.{first}.{i}.attn.qkv.weight" for i in range(n1)] + keys += [f"diffusion_model.{second}.{i}.linear1.weight" for i in range(n2)] + return keys + [f"diffusion_model.{k}" for k in extra] + + +FAMILIES = { + "SD15": _unet_sd(12, 12), + "SDXL": _unet_sd(9, 9), + "FLUX": _dual_stack("double_blocks", "single_blocks", 19, 38, + extra=("img_in.weight", "txt_in.weight", + "time_in.in_layer.weight", "guidance_in.in_layer.weight", + "vector_in.in_layer.weight", "final_layer.linear.weight")), + "HiDream": _dual_stack("double_stream_blocks", "single_stream_blocks", 16, 32, + extra=("x_embedder.weight", "final_layer.linear.weight")), + "HunyuanVideo": _dual_stack("double_blocks", "single_blocks", 20, 40, + extra=("img_in.proj.weight", "final_layer.linear.weight")), + "SD3": _flat_stack("joint_blocks", 24, + extra=("x_embedder.proj.weight", "t_embedder.mlp.0.weight", + "y_embedder.mlp.0.weight", "context_embedder.weight", + "pos_embed", "final_layer.linear.weight")), + "AuraFlow": (_flat_stack("joint_transformer_blocks", 4) + + _flat_stack("single_transformer_blocks", 32) + + ["diffusion_model.init_x_linear.weight", + "diffusion_model.final_linear.weight"]), + "QwenImage": _flat_stack("transformer_blocks", 60, + extra=("img_in.weight", "txt_in.weight", + "time_text_embed.timestep_embedder.linear_1.weight", + "proj_out.weight")), + "LTXV": _flat_stack("transformer_blocks", 28, + extra=("patchify_proj.weight", "caption_projection.linear_1.weight", + "proj_out.weight")), + "PixArt": _flat_stack("blocks", 28, + extra=("x_embedder.proj.weight", "t_embedder.mlp.0.weight", + "final_layer.linear.weight")), + "Wan21": _flat_stack("blocks", 40, + extra=("patch_embedding.weight", "time_embedding.0.weight", + "text_embedding.0.weight", "head.head.weight")), + "Mochi": _flat_stack("blocks", 48, + extra=("x_embedder.proj.weight", "final_layer.linear.weight")), + "Lumina2": _flat_stack("layers", 26, + extra=("x_embedder.weight", "cap_embedder.1.weight", + "final_layer.linear.weight")), + "Cosmos": _flat_stack("blocks", 28, + extra=("x_embedder.proj.1.weight", "final_layer.linear.weight")), +} + + +class TestStackDiscovery(unittest.TestCase): + def test_discovers_expected_stacks_and_sizes(self): + cases = { + "SD15": {"input_blocks": 12, "middle_block": 1, "output_blocks": 12}, + "SDXL": {"input_blocks": 9, "middle_block": 1, "output_blocks": 9}, + "FLUX": {"double_blocks": 19, "single_blocks": 38}, + "HiDream": {"double_stream_blocks": 16, "single_stream_blocks": 32}, + "SD3": {"joint_blocks": 24}, + "AuraFlow": {"joint_transformer_blocks": 4, "single_transformer_blocks": 32}, + "QwenImage": {"transformer_blocks": 60}, + "Wan21": {"blocks": 40}, + "Lumina2": {"layers": 26}, + } + for family, expected in cases.items(): + layout = bu.discover_layout(FAMILIES[family]) + self.assertEqual(layout.stacks, expected, family) + + def test_stacks_are_ordered_by_execution_order(self): + layout = bu.discover_layout(FAMILIES["SD15"]) + self.assertEqual(list(layout.stacks), ["input_blocks", "middle_block", "output_blocks"]) + layout = bu.discover_layout(FAMILIES["FLUX"]) + self.assertEqual(list(layout.stacks), ["double_blocks", "single_blocks"]) + layout = bu.discover_layout(FAMILIES["AuraFlow"]) + self.assertEqual(list(layout.stacks), + ["joint_transformer_blocks", "single_transformer_blocks"]) + + def test_outermost_stack_wins_over_nested_one(self): + # SDXL nests transformer_blocks inside input_blocks; the outer one counts. + key = "diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_q.weight" + self.assertEqual(bu._stack_match(key), ("input_blocks", 4)) + + def test_longer_stack_name_wins(self): + for key, expected in [ + ("diffusion_model.single_transformer_blocks.3.x.weight", + ("single_transformer_blocks", 3)), + ("diffusion_model.single_blocks.3.x.weight", ("single_blocks", 3)), + ("diffusion_model.double_stream_blocks.2.x.weight", ("double_stream_blocks", 2)), + ("diffusion_model.blocks.7.x.weight", ("blocks", 7)), + ("diffusion_model.layers.7.x.weight", ("layers", 7)), + ]: + self.assertEqual(bu._stack_match(key), expected, key) + + def test_unknown_stack_name_still_discovered(self): + keys = [f"diffusion_model.mystery_blocks.{i}.attn.weight" for i in range(10)] + layout = bu.discover_layout(keys) + self.assertEqual(layout.stacks, {"mystery_blocks": 10}) + self.assertEqual(layout.total, 10) + + def test_empty_layout_is_falsy_and_described(self): + layout = bu.discover_layout(["diffusion_model.final_layer.linear.weight"]) + self.assertFalse(layout) + self.assertIn("no block stacks", layout.describe()) + + +class TestUniversalClassification(unittest.TestCase): + def test_every_family_covers_its_whole_stack_without_other(self): + """No key from any family should fall through to 'Other Tensors'.""" + for family, keys in FAMILIES.items(): + layout = bu.discover_layout(keys) + unclassified = [k for k in keys + if bu.classify_universal_key(k, layout) == bu.UNIVERSAL_OTHER] + self.assertEqual(unclassified, [], f"{family}: {unclassified[:5]}") + + def test_every_depth_bucket_is_reachable(self): + for family, keys in FAMILIES.items(): + layout = bu.discover_layout(keys) + if not layout: + continue + buckets = {bu.classify_universal_key(k, layout) for k in keys} + for expected in (bu.UNIVERSAL_EARLY, bu.UNIVERSAL_MID, bu.UNIVERSAL_LATE): + self.assertIn(expected, buckets, f"{family} missing {expected}") + + def test_depth_ordering_is_monotonic(self): + """Walking a flat stack front to back must never move backwards.""" + order = [bu.UNIVERSAL_EARLY, bu.UNIVERSAL_EARLY_MID, bu.UNIVERSAL_MID, + bu.UNIVERSAL_LATE_MID, bu.UNIVERSAL_LATE] + keys = _flat_stack("blocks", 40) + layout = bu.discover_layout(keys) + seen = [order.index(bu.classify_universal_key(k, layout)) for k in keys] + self.assertEqual(seen, sorted(seen)) + self.assertEqual(seen[0], 0) + self.assertEqual(seen[-1], len(order) - 1) + + def test_dual_stack_spans_the_whole_axis(self): + """FLUX double blocks sit early, single blocks run to the end.""" + layout = bu.discover_layout(FAMILIES["FLUX"]) + first = bu.classify_universal_key( + "diffusion_model.double_blocks.0.img_attn.qkv.weight", layout) + last = bu.classify_universal_key( + "diffusion_model.single_blocks.37.linear1.weight", layout) + self.assertEqual(first, bu.UNIVERSAL_EARLY) + self.assertEqual(last, bu.UNIVERSAL_LATE) + + def test_unet_stages_land_in_sensible_buckets(self): + layout = bu.discover_layout(FAMILIES["SDXL"]) + early = bu.classify_universal_key( + "diffusion_model.input_blocks.0.1.transformer_blocks.0.attn1.to_q.weight", layout) + late = bu.classify_universal_key( + "diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_v.weight", layout) + self.assertEqual(early, bu.UNIVERSAL_EARLY) + self.assertEqual(late, bu.UNIVERSAL_LATE) + + def test_embedding_and_head_tokens(self): + layout = bu.discover_layout(FAMILIES["FLUX"]) + inputs = ["diffusion_model.img_in.weight", "diffusion_model.txt_in.weight", + "diffusion_model.time_in.in_layer.weight", + "diffusion_model.x_embedder.proj.weight", + "diffusion_model.patch_embedding.weight", + "diffusion_model.caption_projection.linear_1.weight", + "diffusion_model.label_emb.0.0.weight"] + heads = ["diffusion_model.final_layer.linear.weight", + "diffusion_model.proj_out.weight", + "diffusion_model.norm_out.linear.weight", + "diffusion_model.out.2.weight", + "diffusion_model.final_linear.weight"] + for key in inputs: + self.assertEqual(bu.classify_universal_key(key, layout), + bu.UNIVERSAL_INPUT, key) + for key in heads: + self.assertEqual(bu.classify_universal_key(key, layout), + bu.UNIVERSAL_OUTPUT, key) + + def test_unrecognised_key_falls_through(self): + layout = bu.discover_layout(FAMILIES["FLUX"]) + self.assertEqual( + bu.classify_universal_key("diffusion_model.mystery.weight", layout), + bu.UNIVERSAL_OTHER) + + def test_classification_is_safe_with_an_empty_layout(self): + layout = bu.discover_layout([]) + self.assertEqual( + bu.classify_universal_key("diffusion_model.blocks.3.attn.weight", layout), + bu.UNIVERSAL_OTHER) + + def test_index_beyond_discovered_size_is_clamped(self): + layout = bu.discover_layout(_flat_stack("blocks", 10)) + result = bu.classify_universal_key("diffusion_model.blocks.99.attn.weight", layout) + self.assertEqual(result, bu.UNIVERSAL_LATE) + + +class TestRealWorldRegressions(unittest.TestCase): + """Cases found by running against real ComfyUI models, locked in here. + + Each of these silently landed in 'Other Tensors' before being fixed. + """ + + def test_sd3_negative_block_index(self): + # ComfyUI's SD3 key map addresses the final block as "joint_blocks.-1". + layout = bu.discover_layout(_flat_stack("joint_blocks", 24)) + key = "diffusion_model.joint_blocks.-1.context_block.adaLN_modulation.1.weight" + self.assertEqual(bu._stack_match(key), ("joint_blocks", -1)) + self.assertEqual(bu.classify_universal_key(key, layout), bu.UNIVERSAL_LATE) + + def test_lumina_refiner_stacks_are_discovered(self): + keys = (_flat_stack("noise_refiner", 2) + _flat_stack("context_refiner", 2) + + _flat_stack("layers", 32)) + layout = bu.discover_layout(keys) + self.assertEqual(layout.stacks, + {"noise_refiner": 2, "context_refiner": 2, "layers": 32}) + self.assertEqual(list(layout.stacks)[:2], ["noise_refiner", "context_refiner"]) + self.assertEqual( + [k for k in keys if bu.classify_universal_key(k, layout) == bu.UNIVERSAL_OTHER], + []) + + def test_auraflow_native_layer_stacks(self): + # ComfyUI names these *_layers, not the diffusers *_transformer_blocks. + keys = _flat_stack("double_layers", 4) + _flat_stack("single_layers", 32) + layout = bu.discover_layout(keys) + self.assertEqual(list(layout.stacks), ["double_layers", "single_layers"]) + + def test_conditioning_embedders_are_input_not_other(self): + layout = bu.discover_layout(_flat_stack("blocks", 28)) + for key in ("diffusion_model.ar_embedder.mlp.0.weight", # PixArt + "diffusion_model.csize_embedder.mlp.0.weight", # PixArt + "diffusion_model.t_block.1.weight", # PixArt + "diffusion_model.adaln_single.linear.weight", # LTX-Video + "diffusion_model.scale_shift_table", # LTX-Video + "diffusion_model.txt_norm.weight", # Qwen-Image + "diffusion_model.time_projection.1.weight", # Wan + "diffusion_model.cond_seq_linear.weight", # AuraFlow + "diffusion_model.positional_encoding", # AuraFlow + "diffusion_model.init_x_linear.weight"): # AuraFlow + self.assertEqual(bu.classify_universal_key(key, layout), + bu.UNIVERSAL_INPUT, key) + + def test_flux_controlnet_hint_input(self): + self.assertEqual( + bobs_blocks.classify_flux_key("diffusion_model.pos_embed_input.weight"), + bobs_blocks.FLUX_IMAGE_HINT) + + +class TestUniversalPresets(unittest.TestCase): + def test_registered_into_the_shared_tables(self): + self.assertIn("UNIVERSAL", bobs_blocks.LORA_BLOCK_PRESETS) + self.assertIs(bobs_blocks.ALL_BLOCKS["UNIVERSAL"], bu.ALL_UNIVERSAL_BLOCKS) + self.assertEqual(bobs_blocks.TEXT_ENCODER_BLOCK["UNIVERSAL"], + bu.UNIVERSAL_TEXT_ENCODER) + + def test_every_preset_covers_every_block(self): + for name, config in bu.UNIVERSAL_PRESETS.items(): + if name == "Custom": + self.assertEqual(config, {}) + continue + self.assertEqual(set(config["block_weights"]), set(bu.ALL_UNIVERSAL_BLOCKS), name) + + def test_resolve_block_strengths_handles_the_universal_family(self): + result = bobs_blocks.resolve_block_strengths("UNIVERSAL", "Full (Normal LoRA)", 0.7, {}) + self.assertEqual(set(result), set(bu.ALL_UNIVERSAL_BLOCKS)) + for value in result.values(): + self.assertAlmostEqual(value, 0.7) + + def test_every_block_has_a_tooltip(self): + for name in bu.ALL_UNIVERSAL_BLOCKS: + self.assertIn(name, bobs_blocks.BLOCK_TOOLTIPS) + + +if __name__ == "__main__": + unittest.main()