Fix issues found in code review

README.md

@@ -160,6 +160,36 @@ python -m unittest discover -s tests -v
 
 ## Changelog
 
+### 1.2.1
+
+Code-review fixes. No behaviour change for a correctly-paired LoRA and model.
+
+-   **The FLUX and SDXL classifiers now anchor on the module path**, matching
+    the Universal loader. They previously matched on an underscore-normalised
+    key, which erases the separator that tells an *outermost* block stack from
+    one nested inside a block. An SDXL key such as
+    `input_blocks.4.1.transformer_blocks.0.attn1.to_q.weight` was read as
+    double-stream block 0, so pointing the FLUX node at an SDXL model filed
+    every tensor under "Early Downsampling". Real FLUX and SDXL keys are
+    unaffected — this only ever misfired across families.
+-   **Mismatched loader and model now warn** instead of failing silently. If
+    90% or more of a LoRA's UNet tensors land in one block, the `info` output
+    and the log say so and point at the Universal loader. Measured: a real SDXL
+    key set through the FLUX node warns at 99%, a FLUX key set through the SDXL
+    node at 100%, while all six correctly-paired combinations stay silent.
+-   **Families register explicitly** via `register_family()` rather than by
+    mutating `bobs_blocks`'s tables at import time, removing an import-order
+    dependency that also made the test suite fragile.
+-   **Tooltips are per family.** FLUX/SDXL and Universal each define a block
+    called "Text Encoder"; the Universal wording used to overwrite the other.
+-   **`Output Head` detection tightened.** The `_head_` substring also matched
+    `multi_head_attention` and `head_dim_proj`; it is now anchored to a real
+    head module.
+-   An unknown `preset` name (only reachable from hand-edited workflow JSON)
+    logs a warning instead of silently behaving as `Custom`.
+-   Test suite grown to 71, including a new `tests/test_node_behaviour.py`
+    covering node registration, widget order, the report and the new warning.
+
 ### 1.2.0
 
 **New: a Universal loader covering every architecture ComfyUI supports.**

bobs_blocks.py

@@ -98,6 +98,11 @@
 }
 BLOCK_TOOLTIPS[SDXL_TEXT_ENCODER] = BLOCK_TOOLTIPS[FLUX_TEXT_ENCODER]
 
+#: Per-family tooltip overrides, keyed family -> block -> text. Families share
+#: some block *names* ("Text Encoder", "Other Tensors") but want different help
+#: text, so a single flat dict cannot hold both.
+FAMILY_TOOLTIPS: Dict[str, Dict[str, str]] = {}
+
 # -----------------------------------------------------------------------------#
 #                             PRESET  DEFINITIONS                              #
 # -----------------------------------------------------------------------------#
@@ -278,6 +283,35 @@
 
 ALL_BLOCKS = {"FLUX": ALL_FLUX_BLOCKS, "SDXL": ALL_SDXL_BLOCKS}
 
+#: Families defined in this module. Additional families register themselves via
+#: :func:`register_family`; keep this constant for code that needs the built-in
+#: two regardless of what else has been imported.
+BUILTIN_FAMILIES = ("FLUX", "SDXL")
+
+
+def register_family(family: str,
+                    blocks: List[str],
+                    presets: Dict[str, Dict],
+                    text_encoder_block: str,
+                    tooltips: Dict[str, str] = None) -> None:
+    """Register an additional block family (see :mod:`bobs_universal`).
+
+    Kept explicit rather than done by assignment at import time: the tables
+    below are module state, and a silent import-order dependency makes both
+    behaviour and tests fragile. Tooltips are stored per family so that two
+    families sharing a block name (both call one "Text Encoder") do not
+    overwrite each other's help text.
+    """
+    ALL_BLOCKS[family] = blocks
+    LORA_BLOCK_PRESETS[family] = presets
+    TEXT_ENCODER_BLOCK[family] = text_encoder_block
+    FAMILY_TOOLTIPS[family] = dict(tooltips or {})
+
+
+def tooltip_for(family: str, block: str) -> str:
+    """Help text for a block, preferring the family's own wording."""
+    return FAMILY_TOOLTIPS.get(family, {}).get(block) or BLOCK_TOOLTIPS.get(block, block)
+
 # -----------------------------------------------------------------------------#
 #                              KEY NORMALISATION                               #
 # -----------------------------------------------------------------------------#
@@ -322,10 +356,27 @@ def normalize_key(key: str) -> str:
     (1.0, FLUX_LATE_UP),
 )
 
-# ``single`` must be tested before ``double``: "single_transformer_blocks_3"
-# would otherwise also satisfy the double-stream pattern.
-_RE_SINGLE_BLOCK = re.compile(r"_(?:single_blocks|single_transformer_blocks)_(\d+)(?:_|$)")
-_RE_DOUBLE_BLOCK = re.compile(r"_(?:double_blocks|transformer_blocks)_(\d+)(?:_|$)")
+# Stack matching runs on the RAW dotted key, never the underscore-normalised
+# one. Normalising erases the separator that distinguishes an *outermost* stack
+# from one nested inside a block: an SDXL key
+# "input_blocks.4.1.transformer_blocks.0.attn1.to_q.weight" normalises to
+# "..._input_blocks_4_1_transformer_blocks_0_..." where a naive search matches
+# the inner transformer_blocks and reads the key as double-stream block 0.
+# Anchoring on "." keeps the two apart, and makes the leftmost match the
+# outermost stack.
+#
+# ``single`` is still tested before ``double`` so that
+# "single_transformer_blocks.3" cannot be read as "transformer_blocks.3".
+#
+# Anchoring on "." alone is not enough — the nested transformer_blocks is also
+# dot-preceded. FLUX's stacks are always *top level*, so the pattern is anchored
+# to the start of the module path (after an optional known prefix). SDXL's
+# nested stack then cannot match, because "input_blocks" sits in front of it.
+_PREFIX = r"^(?:model\.)?(?:diffusion_model\.|transformer\.)?"
+_RE_SINGLE_BLOCK = re.compile(
+    _PREFIX + r"(?:single_blocks|single_transformer_blocks)\.(-?\d+)(?=\.|$)")
+_RE_DOUBLE_BLOCK = re.compile(
+    _PREFIX + r"(?:double_blocks|transformer_blocks)\.(-?\d+)(?=\.|$)")
 
 # Head / tail tokens, most specific first.
 _FLUX_TOKEN_MAP: Sequence[Tuple[str, str]] = (
@@ -375,19 +426,31 @@ def _bucket(index: int, boundaries: Sequence[Tuple[int, str]], fallback: str) ->
     return fallback
 
 
+def _resolve_index(index: int, boundaries: Sequence[Tuple[int, str]]) -> int:
+    """Resolve a Python-style negative block index against the stack size."""
+    if index >= 0:
+        return index
+    size = boundaries[-1][0] if boundaries else 0
+    return max(index + size, 0)
+
+
 def classify_flux_key(model_key: str, ranges=None) -> str:
     """Map a FLUX UNet state-dict key to one of :data:`ALL_FLUX_BLOCKS`."""
     double_bounds, single_bounds = ranges or flux_block_ranges()
-    nk = normalize_key(model_key)
 
-    match = _RE_SINGLE_BLOCK.search(nk)
+    # Block stacks match on the raw dotted key (see _RE_SINGLE_BLOCK); the
+    # token checks below match on the normalised one.
+    match = _RE_SINGLE_BLOCK.search(model_key)
     if match:
-        return _bucket(int(match.group(1)), single_bounds, FLUX_LATE_UP)
+        return _bucket(_resolve_index(int(match.group(1)), single_bounds),
+                       single_bounds, FLUX_LATE_UP)
 
-    match = _RE_DOUBLE_BLOCK.search(nk)
+    match = _RE_DOUBLE_BLOCK.search(model_key)
     if match:
-        return _bucket(int(match.group(1)), double_bounds, FLUX_CORE)
+        return _bucket(_resolve_index(int(match.group(1)), double_bounds),
+                       double_bounds, FLUX_CORE)
 
+    nk = normalize_key(model_key)
     for token, block in _FLUX_TOKEN_MAP:
         if token in nk:
             return block
@@ -399,15 +462,23 @@ def classify_flux_key(model_key: str, ranges=None) -> str:
 #                             SDXL  CLASSIFICATION                             #
 # -----------------------------------------------------------------------------#
 
+# Anchored on "." for the same reason as the FLUX patterns above: the stage is
+# the outermost path component, not whatever is nested inside a block.
+_RE_SDXL_STAGE = re.compile(
+    _PREFIX + r"(input_blocks|middle_block|output_blocks)(?=\.|$)")
+
+_SDXL_STAGE_BLOCKS = {
+    "input_blocks": SDXL_INPUT_BLOCKS,
+    "middle_block": SDXL_MIDDLE_BLOCK,
+    "output_blocks": SDXL_OUTPUT_BLOCKS,
+}
+
+
 def classify_sdxl_key(model_key: str) -> str:
     """Map an SDXL UNet state-dict key to one of :data:`ALL_SDXL_BLOCKS`."""
-    nk = normalize_key(model_key)
-    if "_input_blocks_" in nk:
-        return SDXL_INPUT_BLOCKS
-    if "_middle_block_" in nk:
-        return SDXL_MIDDLE_BLOCK
-    if "_output_blocks_" in nk:
-        return SDXL_OUTPUT_BLOCKS
+    match = _RE_SDXL_STAGE.search(model_key)
+    if match:
+        return _SDXL_STAGE_BLOCKS[match.group(1)]
     return SDXL_OTHER
 
 

bobs_lora_loader.py

@@ -27,19 +27,21 @@
 from .bobs_blocks import (
     ALL_FLUX_BLOCKS,
     ALL_SDXL_BLOCKS,
-    BLOCK_TOOLTIPS,
     FLUX_DEFAULT_DEPTH,
     FLUX_DEFAULT_DEPTH_SINGLE,
+    FLUX_OTHER,
     LORA_BLOCK_PRESETS,
+    SDXL_OTHER,
     TEXT_ENCODER_BLOCK,
     classify_flux_key,
     classify_sdxl_key,
     flux_block_ranges,
     resolve_block_strengths,
+    tooltip_for,
 )
 from .bobs_universal import (  # registers the UNIVERSAL family on import
     ALL_UNIVERSAL_BLOCKS,
-    UNIVERSAL_TOOLTIPS,
+    UNIVERSAL_OTHER,
     classify_universal_key,
     discover_layout,
 )
@@ -54,6 +56,11 @@
 MAX_STRENGTH = 5.0
 MAX_BLOCK_WEIGHT = 2.0
 
+#: Share of UNet tensors in a single block above which the family looks wrong.
+#: A correctly-paired LoRA spreads across blocks; the worst real case measured
+#: was 60/200 (30%) in one block, so 0.9 leaves a wide margin.
+_MISMATCH_SHARE = 0.9
+
 
 # -----------------------------------------------------------------------------#
 #                               HELPER FUNCTIONS                               #
@@ -132,6 +139,51 @@ def _format_report(tag: str,
     return "\n".join(lines)
 
 
+def _mismatch_warning(tag: str,
+                      blocks: List[str],
+                      grouped: Dict[str, Dict[Any, Any]],
+                      text_encoder_block: str,
+                      catch_all_block: str) -> Optional[str]:
+    """Detect a LoRA/model pairing that the family classifier cannot resolve.
+
+    Picking the wrong loader for a model is silent otherwise: the keys still
+    match (ComfyUI's key map is built from the model, so the LoRA loads), they
+    just all land in whichever single bucket the classifier happens to reach.
+    Two shapes of that are worth calling out:
+
+    - everything piled into the family's catch-all, which then applies at that
+      slider's weight with no block weighting at all;
+    - everything piled into one *named* block, which means the sliders are
+      lying about what they control.
+    """
+    unet_counts = {name: len(patches) for name, patches in grouped.items()
+                   if name != text_encoder_block and patches}
+    total = sum(unet_counts.values())
+    if total < 8:  # too few patches for the distribution to mean anything
+        return None
+
+    # A threshold rather than "all of them": cross-family key names overlap just
+    # enough to scatter a few tensors. Feeding a real SDXL key set through the
+    # FLUX classifier leaves 98.7% in the catch-all and sends the rest to
+    # "Final Output Layer" via the shared proj_out token, which an exact-equality
+    # test would wave through.
+    def _share(name: str) -> float:
+        return unet_counts.get(name, 0) / total
+
+    if _share(catch_all_block) >= _MISMATCH_SHARE:
+        return (f"WARNING: {_share(catch_all_block):.0%} of UNet tensors landed in "
+                f"'{catch_all_block}'. This LoRA does not look like a {tag} LoRA, or the "
+                f"model is not a {tag} model — the per-block sliders are barely doing "
+                f"anything. Try the Universal loader.")
+
+    dominant, count = max(unet_counts.items(), key=lambda kv: kv[1])
+    if count / total >= _MISMATCH_SHARE and len(blocks) > 2:
+        return (f"WARNING: {count / total:.0%} of UNet tensors landed in '{dominant}'. "
+                f"That usually means this model is not a {tag} model, so the block "
+                f"sliders do not map to what their names say. Try the Universal loader.")
+    return None
+
+
 def _explain_empty_blocks(tag: str,
                           blocks: List[str],
                           grouped: Dict[str, Dict[Any, Any]],
@@ -155,6 +207,8 @@ class _BobsLoraLoaderBase:
 
     FAMILY = "FLUX"
     BLOCKS: List[str] = ALL_FLUX_BLOCKS
+    #: Bucket that receives anything the family classifier could not place.
+    CATCH_ALL_BLOCK = FLUX_OTHER
 
     RETURN_TYPES = ("MODEL", "CLIP", "STRING")
     RETURN_NAMES = ("MODEL", "CLIP", "info")
@@ -193,7 +247,7 @@ def _input_types(cls) -> Dict[str, Any]:
         for block in cls.BLOCKS:
             required[block] = ("FLOAT", {
                 "default": 1.0, "min": -MAX_BLOCK_WEIGHT, "max": MAX_BLOCK_WEIGHT, "step": 0.05,
-                "tooltip": BLOCK_TOOLTIPS.get(block, block),
+                "tooltip": tooltip_for(cls.FAMILY, block),
             })
         return {
             "required": required,
@@ -253,6 +307,9 @@ def apply_lora(self, model, lora_name, strength, preset, clip=None, **kwargs):
             self.logger.error(message)
             return (model, clip, message)
 
+        if preset not in LORA_BLOCK_PRESETS[self.FAMILY]:
+            self.logger.warning(
+                "[%s] unknown preset %r; falling back to the slider values.", tag, preset)
         strengths = resolve_block_strengths(self.FAMILY, preset, strength, kwargs)
 
         key_map, unet_targets = _build_key_maps(model, clip)
@@ -296,6 +353,13 @@ def apply_lora(self, model, lora_name, strength, preset, clip=None, **kwargs):
         # comfy.lora.load_lora itself as a "lora key not loaded" warning.
         report = _format_report(tag, lora_name, preset, self.BLOCKS,
                                 grouped, strengths, applied, detail)
+
+        mismatch = _mismatch_warning(tag, self.BLOCKS, grouped,
+                                     text_encoder_block, self.CATCH_ALL_BLOCK)
+        if mismatch:
+            self.logger.warning("[%s] %s", tag, mismatch)
+            report = f"{report}\n{mismatch}"
+
         self.logger.info("\n%s", report)
         _explain_empty_blocks(tag, self.BLOCKS, grouped, strengths)
 
@@ -335,6 +399,7 @@ def _classifier(self, model, unet_targets):
 class BobsLoraLoaderSdxl(_BobsLoraLoaderBase):
     FAMILY = "SDXL"
     BLOCKS = ALL_SDXL_BLOCKS
+    CATCH_ALL_BLOCK = SDXL_OTHER
     DESCRIPTION = (
         "Applies a LoRA to an SDXL model with independent weights for the text "
         "encoder and the UNet input / middle / output stages."
@@ -355,6 +420,7 @@ def _classifier(self, model, unet_targets):
 class BobsLoraLoaderUniversal(_BobsLoraLoaderBase):
     FAMILY = "UNIVERSAL"
     BLOCKS = ALL_UNIVERSAL_BLOCKS
+    CATCH_ALL_BLOCK = UNIVERSAL_OTHER
     DESCRIPTION = (
         "Block-weighted LoRA loading for any architecture ComfyUI supports — "
         "SD1.5, SD2, SDXL, SD3/3.5, FLUX, Chroma, AuraFlow, PixArt, HiDream, "

bobs_universal.py

@@ -32,13 +32,7 @@
 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,
-)
+from .bobs_blocks import normalize_key, register_family
 
 # -----------------------------------------------------------------------------#
 #                                 BLOCK NAMES                                  #
@@ -74,7 +68,7 @@
     (0.4, UNIVERSAL_EARLY_MID),
     (0.6, UNIVERSAL_MID),
     (0.8, UNIVERSAL_LATE_MID),
-    (1.01, UNIVERSAL_LATE),
+    (1.0, UNIVERSAL_LATE),
 )
 
 UNIVERSAL_TOOLTIPS: Dict[str, str] = {
@@ -154,9 +148,12 @@
     "_norm_out",
     "_unpatchify",
     "_final_norm",
-    "_head_",
 )
-_OUTPUT_RE = re.compile(r"_out_\d+(?=_|$)")  # SD "out.0" / "out.2" head
+# Anchored patterns rather than loose substrings. A bare "_head_" token also
+# matched "multi_head_attention" and "head_dim_proj"; these require the segment
+# to actually be a head module ("head.head.weight", "out.0", "out.2").
+_OUTPUT_RE = re.compile(
+    r"(?:^|\.)(?:out\.\d+|head(?:\.head)?)(?=\.|$)", re.IGNORECASE)
 
 _INPUT_TOKENS: Tuple[str, ...] = (
     "_img_in",
@@ -301,7 +298,7 @@ def classify_universal_key(model_key: str, layout: BlockLayout) -> str:
     for token in _OUTPUT_TOKENS:
         if token in nk:
             return UNIVERSAL_OUTPUT
-    if _OUTPUT_RE.search(nk):
+    if _OUTPUT_RE.search(model_key):  # dot-anchored: raw key, not normalised
         return UNIVERSAL_OUTPUT
 
     for token in _INPUT_TOKENS:
@@ -396,7 +393,12 @@ def classify_universal_key(model_key: str, layout: BlockLayout) -> str:
     },
 }
 
-LORA_BLOCK_PRESETS["UNIVERSAL"] = UNIVERSAL_PRESETS
-ALL_BLOCKS["UNIVERSAL"] = ALL_UNIVERSAL_BLOCKS
-TEXT_ENCODER_BLOCK["UNIVERSAL"] = UNIVERSAL_TEXT_ENCODER
-BLOCK_TOOLTIPS.update(UNIVERSAL_TOOLTIPS)
+UNIVERSAL_FAMILY = "UNIVERSAL"
+
+register_family(
+    UNIVERSAL_FAMILY,
+    blocks=ALL_UNIVERSAL_BLOCKS,
+    presets=UNIVERSAL_PRESETS,
+    text_encoder_block=UNIVERSAL_TEXT_ENCODER,
+    tooltips=UNIVERSAL_TOOLTIPS,
+)