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.
This commit is contained in:
BobsBlazed
@@ -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.
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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),
|
||||
|
||||
+67
-12
@@ -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)",
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
+2
-2
@@ -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]
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user