small update
This commit is contained in:
Terpentine
@@ -92,6 +92,7 @@ Console output keeps you informed at every step without overwhelming.
|
||||
- **Install:** Drop into `ComfyUI/custom_nodes/` and restart ComfyUI.
|
||||
- **Explore:** Drag [`assets/nodes.png`](assets/nodes.png) into the workflow area to see the suite in action.
|
||||
- **Experiment:** Connect a *Easy LoRA Merger* → *Easy LoRA Baker* pipeline, or use *Easy Checkpoint Studio* to shrink a 12GB checkpoint to FP8.
|
||||
- **Optional — GGUF export:** Install `pip install gguf` to enable GGUF output format in Easy Checkpoint Studio. GGUF files are loaded with the [ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF) custom node. All other Easy LoRA Merger nodes work without either dependency.
|
||||
- **Feedback:** Open an issue on GitHub — contributions and ideas welcome.
|
||||
|
||||
---
|
||||
@@ -128,8 +129,11 @@ Used by: **Easy Checkpoint Studio**
|
||||
| `int8` | Explicit INT8 quantization, 1 byte/param |
|
||||
| `int8_convrot` | INT8 with convolution rotation handling |
|
||||
| `svd_only` | SVD compression only, no dtype conversion |
|
||||
| `gguf_q8_0` / `gguf_q5_0` / `gguf_q4_0` | Block‑wise GGUF quantization — requires `gguf` Python package (`pip install gguf`). Output loads with [ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF) nodes. |
|
||||
|
||||
> 💡 **Tip:** For everyday use, just leave it on `auto`. The nodes will pick the best format for your hardware automatically.
|
||||
>
|
||||
> ⚠️ **Note:** The `gguf` package is **optional**. All nodes work without it — GGUF output is only available in Easy Checkpoint Studio when `gguf` is installed. If missing, GGUF options are hidden from the dropdown and the Checkpoint Studio loads gracefully without them.
|
||||
|
||||
---
|
||||
|
||||
|
||||
+19
-5
@@ -1,19 +1,33 @@
|
||||
"""
|
||||
Easy LoRA Merger — ComfyUI custom node pack.
|
||||
Re-exports all node class/display-name mappings and registers the SmartModelBaker.
|
||||
"""
|
||||
|
||||
from .easy_lora_merger import NODE_CLASS_MAPPINGS as ELM_CLASS_MAPPINGS
|
||||
from .easy_lora_merger import NODE_DISPLAY_NAME_MAPPINGS as ELM_DISPLAY_MAPPINGS
|
||||
from .baker_node import SmartModelBaker
|
||||
|
||||
# This line is the bridge that tells ComfyUI to load your JS files
|
||||
WEB_DIRECTORY = "./js"
|
||||
|
||||
# Merge mappings: start with existing, add SmartModelBaker
|
||||
# ── SmartModelBaker (guarded) ──────────────────────────────────────────────
|
||||
try:
|
||||
from .baker_node import SmartModelBaker
|
||||
_BAKER_AVAILABLE = True
|
||||
except ImportError as _e:
|
||||
SmartModelBaker = None # type: ignore[assignment]
|
||||
_BAKER_AVAILABLE = False
|
||||
print(f"⚠️ Easy LoRA Merger: SmartModelBaker not available ({_e})")
|
||||
|
||||
# Merge mappings: start with existing, add SmartModelBaker if available
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
**ELM_CLASS_MAPPINGS,
|
||||
"SmartModelBaker": SmartModelBaker,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
**ELM_DISPLAY_MAPPINGS,
|
||||
"SmartModelBaker": "🔥 Easy LoRA Baker",
|
||||
}
|
||||
|
||||
if _BAKER_AVAILABLE:
|
||||
NODE_CLASS_MAPPINGS["SmartModelBaker"] = SmartModelBaker
|
||||
NODE_DISPLAY_NAME_MAPPINGS["SmartModelBaker"] = "🔥 Easy LoRA Baker"
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 154 KiB |
@@ -97,10 +97,21 @@ PRECISION_STUDIO = PRECISION_EXTENDED + ["int8", "int8_convrot", "svd_only"]
|
||||
|
||||
DEVICE_OPTIONS = ["auto", "cuda", "cpu"]
|
||||
|
||||
# ── GGUF availability check ────────────────────────────────────────────────
|
||||
# If 'gguf' is not installed, hide GGUF format options from the UI dropdown
|
||||
# to prevent users from selecting options that will fail at runtime.
|
||||
try:
|
||||
import gguf # noqa: F401 — just checking availability
|
||||
_GGUF_AVAILABLE = True
|
||||
except ImportError:
|
||||
_GGUF_AVAILABLE = False
|
||||
|
||||
# GGUF output format options for Checkpoint Studio.
|
||||
# Each entry maps to a GGMLQuantizationType in engine/gguf_writer.py.
|
||||
# "safetensors" retains the existing output path (FP8/BF16/INT8 safetensors file).
|
||||
FORMAT_OPTIONS = ["safetensors", "gguf_q8_0", "gguf_q5_0", "gguf_q4_0"]
|
||||
FORMAT_OPTIONS = ["safetensors"] + (
|
||||
["gguf_q8_0", "gguf_q5_0", "gguf_q4_0"] if _GGUF_AVAILABLE else []
|
||||
)
|
||||
|
||||
|
||||
# ── Unified device/precision resolver ──
|
||||
|
||||
+95
-21
@@ -1,31 +1,21 @@
|
||||
"""
|
||||
Easy LoRA Merger - Node class mappings and display name mappings.
|
||||
This module acts as a re-export hub for all remaining nodes.
|
||||
|
||||
Each node import is individually guarded so that a failure in one
|
||||
module (e.g. missing optional dependency) does not cascade and
|
||||
unregister all nodes.
|
||||
"""
|
||||
|
||||
from .triple_merger_node import EasyLoRATripleMerger
|
||||
from .engine.lora_studio_converter import MusubiLoraConverter
|
||||
from .engine.musubi_checkpoint_studio import MusubiCheckpointStudio
|
||||
from .checkpoint_merger_node import EasyCheckpointMerger
|
||||
from .engine.lora_extractor import EasyLoRAExtractor
|
||||
from .engine.component_extractor import EasyComponentExtractor
|
||||
from .engine.component_combiner import EasyComponentCombiner
|
||||
from .engine.component_merger import EasyComponentMerger
|
||||
# ── Individual guarded imports ────────────────────────────────────────────
|
||||
# Each import is wrapped in try/except so one failure doesn't cascade.
|
||||
# Only successfully imported nodes are registered below.
|
||||
|
||||
# ==================== REGISTRATION ====================
|
||||
_node_registry = {}
|
||||
_node_display_names = {}
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"EasyLoRATripleMerger": EasyLoRATripleMerger,
|
||||
"MusubiLoraConverter": MusubiLoraConverter,
|
||||
"MusubiCheckpointStudio": MusubiCheckpointStudio,
|
||||
"EasyCheckpointMerger": EasyCheckpointMerger,
|
||||
"EasyLoRAExtractor": EasyLoRAExtractor,
|
||||
"EasyComponentExtractor": EasyComponentExtractor,
|
||||
"EasyComponentCombiner": EasyComponentCombiner,
|
||||
"EasyComponentMerger": EasyComponentMerger,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
# Mapping of node IDs to their display names (used if import succeeds)
|
||||
_DISPLAY_NAMES = {
|
||||
"EasyLoRATripleMerger": "🎨 Easy LoRA Merger",
|
||||
"MusubiLoraConverter": "🛡️ Easy LoRA Studio",
|
||||
"MusubiCheckpointStudio": "🛡️ Easy Checkpoint Studio",
|
||||
@@ -34,4 +24,88 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"EasyComponentExtractor": "🔧 Easy Component Extractor",
|
||||
"EasyComponentCombiner": "🧩 Easy Component Combiner",
|
||||
"EasyComponentMerger": "🧩 Easy Component Merger",
|
||||
"EasyTextDisplay": "📋 Easy Text Display",
|
||||
}
|
||||
|
||||
# ── EasyLoRATripleMerger ─────────────────────────────────────────────────
|
||||
try:
|
||||
from .triple_merger_node import EasyLoRATripleMerger
|
||||
_node_registry["EasyLoRATripleMerger"] = EasyLoRATripleMerger
|
||||
except ImportError as _e:
|
||||
print(f"⚠️ EasyLoRAMerger: EasyLoRATripleMerger not available ({_e})")
|
||||
|
||||
# ── MusubiLoraConverter ──────────────────────────────────────────────────
|
||||
try:
|
||||
from .engine.lora_studio_converter import MusubiLoraConverter
|
||||
_node_registry["MusubiLoraConverter"] = MusubiLoraConverter
|
||||
except ImportError as _e:
|
||||
print(f"⚠️ EasyLoRAMerger: MusubiLoraConverter not available ({_e})")
|
||||
|
||||
# ── MusubiCheckpointStudio ───────────────────────────────────────────────
|
||||
try:
|
||||
from .engine.musubi_checkpoint_studio import MusubiCheckpointStudio
|
||||
_node_registry["MusubiCheckpointStudio"] = MusubiCheckpointStudio
|
||||
except ImportError as _e:
|
||||
print(f"⚠️ EasyLoRAMerger: MusubiCheckpointStudio not available ({_e})")
|
||||
|
||||
# ── EasyCheckpointMerger ─────────────────────────────────────────────────
|
||||
try:
|
||||
from .checkpoint_merger_node import EasyCheckpointMerger
|
||||
_node_registry["EasyCheckpointMerger"] = EasyCheckpointMerger
|
||||
except ImportError as _e:
|
||||
print(f"⚠️ EasyLoRAMerger: EasyCheckpointMerger not available ({_e})")
|
||||
|
||||
# ── EasyLoRAExtractor ────────────────────────────────────────────────────
|
||||
try:
|
||||
from .engine.lora_extractor import EasyLoRAExtractor
|
||||
_node_registry["EasyLoRAExtractor"] = EasyLoRAExtractor
|
||||
except ImportError as _e:
|
||||
print(f"⚠️ EasyLoRAMerger: EasyLoRAExtractor not available ({_e})")
|
||||
|
||||
# ── EasyComponentExtractor ───────────────────────────────────────────────
|
||||
try:
|
||||
from .engine.component_extractor import EasyComponentExtractor
|
||||
_node_registry["EasyComponentExtractor"] = EasyComponentExtractor
|
||||
except ImportError as _e:
|
||||
print(f"⚠️ EasyLoRAMerger: EasyComponentExtractor not available ({_e})")
|
||||
|
||||
# ── EasyComponentCombiner ────────────────────────────────────────────────
|
||||
try:
|
||||
from .engine.component_combiner import EasyComponentCombiner
|
||||
_node_registry["EasyComponentCombiner"] = EasyComponentCombiner
|
||||
except ImportError as _e:
|
||||
print(f"⚠️ EasyLoRAMerger: EasyComponentCombiner not available ({_e})")
|
||||
|
||||
# ── EasyComponentMerger ──────────────────────────────────────────────────
|
||||
try:
|
||||
from .engine.component_merger import EasyComponentMerger
|
||||
_node_registry["EasyComponentMerger"] = EasyComponentMerger
|
||||
except ImportError as _e:
|
||||
print(f"⚠️ EasyLoRAMerger: EasyComponentMerger not available ({_e})")
|
||||
|
||||
# ── EasyTextDisplay ─────────────────────────────────────────────────────
|
||||
try:
|
||||
from .easy_text_display import EasyTextDisplay
|
||||
_node_registry["EasyTextDisplay"] = EasyTextDisplay
|
||||
except ImportError as _e:
|
||||
print(f"⚠️ EasyLoRAMerger: EasyTextDisplay not available ({_e})")
|
||||
|
||||
|
||||
# ==================== REGISTRATION ====================
|
||||
|
||||
NODE_CLASS_MAPPINGS = dict(_node_registry)
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
node_id: _DISPLAY_NAMES[node_id]
|
||||
for node_id in _node_registry
|
||||
if node_id in _DISPLAY_NAMES
|
||||
}
|
||||
|
||||
# Startup summary — users can see how many nodes loaded in the ComfyUI console
|
||||
_loaded = len(_node_registry)
|
||||
_total = len(_DISPLAY_NAMES)
|
||||
if _loaded == _total:
|
||||
print(f"✅ Easy LoRA Merger — all {_total} nodes loaded successfully.")
|
||||
else:
|
||||
_missing = _total - _loaded
|
||||
print(f"⚠️ Easy LoRA Merger — {_loaded}/{_total} nodes loaded ({_missing} unavailable).")
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
"""
|
||||
Easy Text Display — ComfyUI Node Definition
|
||||
|
||||
Generic text display node that renders any STRING input as visible text on the
|
||||
node body in the graph UI. Uses the same ``OUTPUT_NODE + {"ui": {"text": …}}``
|
||||
pattern as other ComfyUI text‑display nodes, so no external dependencies are
|
||||
required.
|
||||
|
||||
Use cases
|
||||
---------
|
||||
- View forensic reports from any EasyLoRAMerger node (merger, baker, studio, …)
|
||||
- Inspect output_path strings without opening the console
|
||||
- Debug any STRING wire in your workflow
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
|
||||
class EasyTextDisplay:
|
||||
"""
|
||||
Display a STRING value as visible text on the node body.
|
||||
|
||||
Inputs
|
||||
------
|
||||
text : STRING
|
||||
Any text content to display. Accepts multi‑line strings such as
|
||||
forensic reports, file paths, JSON dumps, etc.
|
||||
|
||||
Outputs
|
||||
-------
|
||||
text : STRING
|
||||
Passthrough of the input — allows chaining to other nodes.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": (
|
||||
"STRING",
|
||||
{
|
||||
"forceInput": True,
|
||||
"multiline": True,
|
||||
"tooltip": "Any text to display — forensic reports, paths, debug output…",
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
FUNCTION = "display"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "LoRA/Utils"
|
||||
|
||||
def display(self, text: str, unique_id=None, extra_pnginfo=None) -> tuple:
|
||||
"""
|
||||
Forward the text to the UI and pass it through the output socket.
|
||||
|
||||
The ``{"ui": {"text": …}}`` payload is picked up by the JavaScript
|
||||
extension which renders the content as a read‑only multiline widget
|
||||
on the node.
|
||||
"""
|
||||
# ── Persist the text into the workflow JSON so it survives saves ──
|
||||
if unique_id is not None and extra_pnginfo is not None:
|
||||
try:
|
||||
# ComfyUI may pass extra_pnginfo as a list [{"workflow": …}]
|
||||
# or as a bare dict {"workflow": …} depending on the version.
|
||||
workflow_container = (
|
||||
extra_pnginfo[0]
|
||||
if isinstance(extra_pnginfo, list)
|
||||
else extra_pnginfo
|
||||
)
|
||||
workflow = (
|
||||
workflow_container.get("workflow")
|
||||
if isinstance(workflow_container, dict)
|
||||
else None
|
||||
)
|
||||
if workflow is None:
|
||||
print(
|
||||
"⚠️ EasyTextDisplay: extra_pnginfo has no 'workflow' key "
|
||||
"— cannot persist text"
|
||||
)
|
||||
else:
|
||||
node = next(
|
||||
(
|
||||
x
|
||||
for x in workflow["nodes"]
|
||||
if str(x["id"]) == str(unique_id)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if node:
|
||||
node["widgets_values"] = [text]
|
||||
except Exception as exc:
|
||||
print(f"⚠️ EasyTextDisplay: failed to persist text ({exc})")
|
||||
|
||||
return {"ui": {"text": [text]}, "result": (text,)}
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Registry helpers (used by easy_lora_merger.py)
|
||||
# =====================================================================
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"EasyTextDisplay": EasyTextDisplay,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"EasyTextDisplay": "📋 Easy Text Display",
|
||||
}
|
||||
@@ -28,10 +28,18 @@ CHECKPOINT_ARCHITECTURES = {
|
||||
["net.blocks"],
|
||||
# Also detect by net.final_layer as secondary marker
|
||||
["net.final_layer"],
|
||||
# Detect bare blocks.* keys without net. prefix (common in fine-tuned
|
||||
# checkpoints). Uses AND logic: both "blocks." and "adaln_modulation"
|
||||
# must exist in the keyset. "blocks." alone would also match SDXL's
|
||||
# input_blocks.* / output_blocks.* and Flux's double_blocks.*, so
|
||||
# "adaln_modulation" narrows it to Anima-specific key patterns.
|
||||
["blocks.", "adaln_modulation"],
|
||||
],
|
||||
"strip_prefixes": [],
|
||||
"canonical_form": "with net. prefix — keys like net.blocks.0.adaln_modulation_cross_attn.1.weight",
|
||||
"description": "Anima / Anime Diffusion (FLOW-based model with net.blocks architecture)",
|
||||
# Strip "net." prefix so that net.blocks.* and bare blocks.* keys
|
||||
# normalize to the same canonical form for matching.
|
||||
"strip_prefixes": ["net."],
|
||||
"canonical_form": "bare blocks. prefix — keys like blocks.0.adaln_modulation_cross_attn.1.weight",
|
||||
"description": "Anima / Anime Diffusion (FLOW-based model with blocks architecture)",
|
||||
},
|
||||
"flux": {
|
||||
"detect": ["double_blocks", "single_blocks"],
|
||||
|
||||
@@ -94,7 +94,7 @@ GGUF_QTYPE_LABELS = {
|
||||
}
|
||||
|
||||
# File type for LlamaFileType metadata
|
||||
_GGUF_QTYPE_TO_FILE_TYPE = {
|
||||
_GGUF_QTYPE_TO_FILE_TYPE = {} if gguf is None else {
|
||||
gguf.GGMLQuantizationType.Q8_0: gguf.LlamaFileType.MOSTLY_Q8_0,
|
||||
gguf.GGMLQuantizationType.Q5_0: gguf.LlamaFileType.MOSTLY_Q5_0,
|
||||
gguf.GGMLQuantizationType.Q4_0: gguf.LlamaFileType.MOSTLY_Q4_0,
|
||||
|
||||
@@ -308,11 +308,15 @@ _ARCH_KEY_RULES: Dict[str, Dict] = {
|
||||
},
|
||||
# ── Anima ────────────────────────────────────────────────────────────
|
||||
# Anima: net.blocks.xxx → diffusion_model.blocks.xxx
|
||||
# After checkpoint_normalizer strips "net." prefix, keys are bare
|
||||
# blocks.xxx. The bare_paths rule catches these normalized keys.
|
||||
"anima": {
|
||||
"checkpoint_prefixes": [
|
||||
("net.blocks.", "diffusion_model.blocks."),
|
||||
],
|
||||
"bare_paths": [],
|
||||
"bare_paths": [
|
||||
"blocks.", # normalized bare blocks.* → diffusion_model.blocks.*
|
||||
],
|
||||
"lora_prefix": "diffusion_model.",
|
||||
},
|
||||
# ── Additional TE/CLIP prefix patterns ───────────────────────────────
|
||||
|
||||
+218
-128
@@ -24,6 +24,8 @@ from ..utils import (
|
||||
categorize_key,
|
||||
DeviceManager,
|
||||
ProgressTracker,
|
||||
comfyui_yield,
|
||||
cleanup_memory,
|
||||
)
|
||||
from .klein_normalizer import (
|
||||
detect_lora_format,
|
||||
@@ -193,9 +195,16 @@ class MusubiLoraConverter:
|
||||
key_str = str(cache_key).encode('utf-8')
|
||||
return hashlib.md5(key_str).hexdigest()
|
||||
|
||||
def _apply_svd_compression(self, normalized_sd, compression_mode, target_rank):
|
||||
def _apply_svd_compression(self, normalized_sd, compression_mode, target_rank, device=None):
|
||||
"""
|
||||
Apply SVD compression to normalized state dict.
|
||||
Apply SVD compression to normalized state dict using batched SVD.
|
||||
|
||||
LoRA layers within a model architecture typically share dimensions
|
||||
(e.g., all DiT blocks have QKV projections of the same size).
|
||||
We group pairs by (out_features, in_features) and call
|
||||
torch.linalg.svd once per group with a batched 3D tensor,
|
||||
dramatically reducing GPU kernel launch overhead.
|
||||
|
||||
Returns (compressed_sd, compression_info) where compression_info dict contains:
|
||||
- compression_mode
|
||||
- original_total_params
|
||||
@@ -235,137 +244,223 @@ class MusubiLoraConverter:
|
||||
layer_count = 0
|
||||
skipped_conv = 0
|
||||
|
||||
# Progress indicator — replace 10% console with ComfyUI ProgressBar
|
||||
total_pairs = len(pairs)
|
||||
with ProgressTracker(total=total_pairs, desc="[SVD] Processing layers") as svd_progress:
|
||||
# ================================================================
|
||||
# PHASE 1: Analyze all pairs, detect orientation, group by shape
|
||||
# ================================================================
|
||||
# Groups are keyed by (out_features, in_features) — pairs sharing
|
||||
# identical matrix dimensions are processed together in one batched
|
||||
# SVD call, which is dramatically faster than one-by-one on GPU.
|
||||
shape_groups = {} # {(out, in): [(dk, uk, base, transposed, rank), ...]}
|
||||
conv_pairs = [] # handled separately (no SVD)
|
||||
skipped_mismatch = [] # rank mismatch warnings
|
||||
|
||||
for i, (dk, uk, base) in enumerate(pairs):
|
||||
down = compressed_sd[dk]
|
||||
up = compressed_sd[uk]
|
||||
for dk, uk, base in pairs:
|
||||
down = compressed_sd[dk]
|
||||
up = compressed_sd[uk]
|
||||
|
||||
# Skip convolutional layers (spatial dimensions >2) to avoid shape errors
|
||||
if down.ndim > 2 or up.ndim > 2:
|
||||
skipped_conv += 1
|
||||
print(f"ℹ️ Skipping convolutional layer {base} (shape {down.shape}, {up.shape})")
|
||||
# Ensure tensors are contiguous for safetensors
|
||||
if not down.is_contiguous():
|
||||
compressed_sd[dk] = down.contiguous()
|
||||
if not up.is_contiguous():
|
||||
compressed_sd[uk] = up.contiguous()
|
||||
continue
|
||||
# Skip convolutional layers (spatial dimensions >2) to avoid shape errors
|
||||
if down.ndim > 2 or up.ndim > 2:
|
||||
conv_pairs.append((dk, uk, base))
|
||||
# Ensure tensors are contiguous for safetensors
|
||||
if not down.is_contiguous():
|
||||
compressed_sd[dk] = down.contiguous()
|
||||
if not up.is_contiguous():
|
||||
compressed_sd[uk] = up.contiguous()
|
||||
continue
|
||||
|
||||
# Detect orientation: down shape (out_features, rank) or (rank, out_features)
|
||||
transposed = False
|
||||
# Check both possible orientations (ambiguous if both conditions hold)
|
||||
if down.shape[1] == up.shape[0] and down.shape[0] == up.shape[1]:
|
||||
# Both orientations possible (square matrices). Decide based on typical LoRA rank being smaller dimension.
|
||||
# Standard orientation: rank = down.shape[1], out_features = down.shape[0], in_features = up.shape[1]
|
||||
# Transposed orientation: rank = down.shape[0], out_features = down.shape[1], in_features = up.shape[0]
|
||||
rank_standard = down.shape[1]
|
||||
rank_transposed = down.shape[0]
|
||||
# Choose orientation where rank is smaller than out_features and in_features (typical LoRA)
|
||||
# If both satisfy, prefer standard orientation.
|
||||
if rank_standard < down.shape[0] and rank_standard < up.shape[1]:
|
||||
# standard orientation is correct
|
||||
out_features, rank = down.shape
|
||||
rank2, in_features = up.shape
|
||||
else:
|
||||
# transposed orientation is correct
|
||||
down = down.T
|
||||
up = up.T
|
||||
out_features, rank = down.shape
|
||||
rank2, in_features = up.shape
|
||||
transposed = True
|
||||
elif down.shape[1] == up.shape[0]:
|
||||
# standard orientation: down (out_features, rank), up (rank, in_features)
|
||||
out_features, rank = down.shape
|
||||
rank2, in_features = up.shape
|
||||
elif down.shape[0] == up.shape[1]:
|
||||
# transposed orientation: down (rank, out_features), up (in_features, rank)
|
||||
# transpose to standard orientation for processing
|
||||
down = down.T
|
||||
up = up.T
|
||||
# Detect orientation: down shape (out_features, rank) or (rank, out_features)
|
||||
transposed = False
|
||||
if down.shape[1] == up.shape[0] and down.shape[0] == up.shape[1]:
|
||||
# Both orientations possible (square matrices). Decide based on typical LoRA rank being smaller dimension.
|
||||
rank_standard = down.shape[1]
|
||||
if rank_standard < down.shape[0] and rank_standard < up.shape[1]:
|
||||
out_features, rank = down.shape
|
||||
rank2, in_features = up.shape
|
||||
else:
|
||||
out_features, rank = down.T.shape
|
||||
rank2, in_features = up.T.shape
|
||||
transposed = True
|
||||
else:
|
||||
# mismatch, skip this pair with warning
|
||||
print(f"⚠️ Rank mismatch {down.shape} vs {up.shape} for pair {dk}, {uk}; skipping")
|
||||
continue
|
||||
assert rank == rank2, f"Rank mismatch {rank} != {rank2} for pair {dk}, {uk}"
|
||||
|
||||
# Compute product W = down @ up (out_features, in_features)
|
||||
W = down.float() @ up.float()
|
||||
|
||||
# SVD
|
||||
max_singular_values = min(out_features, in_features)
|
||||
if compression_mode in ('auto-fast', 'manual'):
|
||||
elif down.shape[1] == up.shape[0]:
|
||||
out_features, rank = down.shape
|
||||
rank2, in_features = up.shape
|
||||
elif down.shape[0] == up.shape[1]:
|
||||
out_features, rank = down.T.shape
|
||||
rank2, in_features = up.T.shape
|
||||
transposed = True
|
||||
else:
|
||||
skipped_mismatch.append((dk, uk, base))
|
||||
continue
|
||||
|
||||
assert rank == rank2, f"Rank mismatch {rank} != {rank2} for pair {dk}, {uk}"
|
||||
|
||||
# Group by (out_features, in_features) for batched SVD
|
||||
shape_key = (out_features, in_features)
|
||||
if shape_key not in shape_groups:
|
||||
shape_groups[shape_key] = []
|
||||
shape_groups[shape_key].append((dk, uk, base, transposed, rank))
|
||||
|
||||
if skipped_mismatch:
|
||||
for dk, uk, base in skipped_mismatch:
|
||||
print(f"⚠️ Rank mismatch for pair {dk}, {uk}; skipping")
|
||||
|
||||
# ================================================================
|
||||
# PHASE 2: Micro-batched SVD per shape group
|
||||
# ================================================================
|
||||
# Large shape groups (e.g., 105 identical QKV projections in DiT models)
|
||||
# produce W = down @ up matrices of size (out_features, in_features).
|
||||
# For 4096×4096 W, stacking ALL 105 into one W_batch would consume
|
||||
# ~7 GB + SVD intermediates — exceeding most consumer GPUs.
|
||||
#
|
||||
# Instead, process in micro-batches of SVD_MICRO_BATCH pairs each.
|
||||
# This caps peak VRAM while still amortizing kernel launch overhead.
|
||||
SVD_MICRO_BATCH = 8 # pairs per batched SVD call
|
||||
total_pairs = len(pairs)
|
||||
|
||||
with ProgressTracker(total=total_pairs, desc="[SVD] Processing layers") as svd_progress:
|
||||
|
||||
for (out_features, in_features), group in shape_groups.items():
|
||||
group_size = len(group)
|
||||
|
||||
# Process this shape group in micro-batches to cap peak VRAM
|
||||
for start_idx in range(0, group_size, SVD_MICRO_BATCH):
|
||||
micro_batch = group[start_idx:start_idx + SVD_MICRO_BATCH]
|
||||
mb_size = len(micro_batch)
|
||||
|
||||
# ── Step 1: Compute W matrices for this micro-batch ──
|
||||
W_list = []
|
||||
for dk, uk, base, transposed, rank in micro_batch:
|
||||
down = compressed_sd[dk]
|
||||
up = compressed_sd[uk]
|
||||
|
||||
# On-demand GPU transfer (lazy, per-pair)
|
||||
down_on_gpu = False
|
||||
up_on_gpu = False
|
||||
if device is not None and down.device != device:
|
||||
down = down.to(device=device)
|
||||
down_on_gpu = True
|
||||
if device is not None and up.device != device:
|
||||
up = up.to(device=device)
|
||||
up_on_gpu = True
|
||||
|
||||
if transposed:
|
||||
W = down.T.float() @ up.T.float()
|
||||
else:
|
||||
W = down.float() @ up.float()
|
||||
|
||||
# Free GPU copies of down/up immediately — no longer needed
|
||||
if down_on_gpu:
|
||||
del down
|
||||
if up_on_gpu:
|
||||
del up
|
||||
|
||||
W_list.append(W)
|
||||
|
||||
# ── Step 2: Batched SVD for this micro-batch ──
|
||||
if compression_mode == 'auto-fast':
|
||||
# randomized low-rank SVD (fast): compute up to original rank
|
||||
q = min(rank, max_singular_values)
|
||||
# Randomized SVD (fast): compute only top-q singular values.
|
||||
# For LoRA, W = down @ up has rank <= training rank.
|
||||
# q = rank + 10 oversampling ensures >99.9% energy capture.
|
||||
pair_rank = micro_batch[0][4] # original training rank
|
||||
q = min(pair_rank + 10, out_features, in_features)
|
||||
q = max(q, 1)
|
||||
if mb_size == 1:
|
||||
W = W_list[0]
|
||||
U, S, V = torch.svd_lowrank(W, q=q)
|
||||
Vh = V.mT # svd_lowrank returns V; convert to Vh convention
|
||||
U_batch, S_batch, Vh_batch = U.unsqueeze(0), S.unsqueeze(0), Vh.unsqueeze(0)
|
||||
del W
|
||||
else:
|
||||
W_batch = torch.stack(W_list, dim=0)
|
||||
del W_list
|
||||
U_batch, S_batch, V_batch = torch.svd_lowrank(W_batch, q=q)
|
||||
Vh_batch = V_batch.mT # convert to Vh convention
|
||||
del W_batch
|
||||
else:
|
||||
# manual target rank: ensure we capture at least target_rank singular values
|
||||
q = min(max(rank, target_rank), max_singular_values)
|
||||
q = max(q, 1)
|
||||
U, S, V = torch.svd_lowrank(W, q=q, niter=4)
|
||||
Vh = V.T # convert to (q, n) format consistent with linalg.svd
|
||||
else:
|
||||
# auto-full, original, or other (full SVD)
|
||||
U, S, Vh = torch.linalg.svd(W, full_matrices=False)
|
||||
|
||||
# Determine target rank
|
||||
if compression_mode == 'original':
|
||||
k = rank
|
||||
elif compression_mode == 'manual':
|
||||
k = min(target_rank, rank)
|
||||
k = max(k, 1)
|
||||
else: # auto-fast or auto-full
|
||||
# compute cumulative energy
|
||||
total_energy = torch.sum(S ** 2)
|
||||
cumulative = torch.cumsum(S ** 2, dim=0)
|
||||
# find smallest k where cumulative >= 0.95 * total_energy
|
||||
k = torch.searchsorted(cumulative, 0.95 * total_energy).item() + 1
|
||||
k = min(k, rank)
|
||||
k = max(k, 1)
|
||||
|
||||
# Truncate and re‑factorize
|
||||
sqrt_Sk = torch.sqrt(S[:k])
|
||||
down_new = (U[:, :k] * sqrt_Sk).contiguous() # (out_features, k)
|
||||
up_new = (sqrt_Sk[:, None] * Vh[:k, :]).contiguous() # (k, in_features)
|
||||
# Cast back to original dtype
|
||||
down_new = down_new.to(dtype=down.dtype)
|
||||
up_new = up_new.to(dtype=up.dtype)
|
||||
|
||||
# Update state dict
|
||||
if transposed:
|
||||
down_new = down_new.T.contiguous()
|
||||
up_new = up_new.T.contiguous()
|
||||
compressed_sd[dk] = down_new
|
||||
compressed_sd[uk] = up_new
|
||||
|
||||
# Update alpha if present
|
||||
alpha_key = f"{base}.alpha"
|
||||
if alpha_key in compressed_sd:
|
||||
alpha = compressed_sd[alpha_key]
|
||||
# Alpha scaling formula: new_alpha = old_alpha * (k / rank)
|
||||
# Alpha may be scalar tensor or float; ensure tensor
|
||||
alpha = alpha * (k / rank)
|
||||
compressed_sd[alpha_key] = alpha
|
||||
|
||||
# Statistics
|
||||
total_original_params += out_features * rank + rank * in_features
|
||||
total_compressed_params += out_features * k + k * in_features
|
||||
rank_sum += k
|
||||
layer_count += 1
|
||||
# Exact SVD (auto-full, original, manual)
|
||||
if mb_size == 1:
|
||||
W = W_list[0]
|
||||
U, S, Vh = torch.linalg.svd(W, full_matrices=False)
|
||||
U_batch, S_batch, Vh_batch = U.unsqueeze(0), S.unsqueeze(0), Vh.unsqueeze(0)
|
||||
del W # free single W immediately
|
||||
else:
|
||||
W_batch = torch.stack(W_list, dim=0) # (B, M, N)
|
||||
del W_list # free list before SVD to save memory
|
||||
U_batch, S_batch, Vh_batch = torch.linalg.svd(W_batch, full_matrices=False)
|
||||
del W_batch # free stacked W immediately after SVD
|
||||
|
||||
# ── Step 3: Process results (truncate and refactorize) ──
|
||||
for i, (dk, uk, base, transposed, rank) in enumerate(micro_batch):
|
||||
S = S_batch[i]
|
||||
U = U_batch[i]
|
||||
Vh = Vh_batch[i]
|
||||
|
||||
# Determine target rank
|
||||
if compression_mode == 'original':
|
||||
k = rank
|
||||
elif compression_mode == 'manual':
|
||||
k = min(target_rank, rank)
|
||||
k = max(k, 1)
|
||||
else: # auto-fast or auto-full (same 95% energy threshold)
|
||||
total_energy = torch.sum(S ** 2)
|
||||
cumulative = torch.cumsum(S ** 2, dim=0)
|
||||
k = torch.searchsorted(cumulative, 0.95 * total_energy).item() + 1
|
||||
k = min(k, rank)
|
||||
k = max(k, 1)
|
||||
|
||||
# Truncate and re-factorize
|
||||
sqrt_Sk = torch.sqrt(S[:k])
|
||||
down_new = (U[:, :k] * sqrt_Sk).contiguous() # (out_features, k)
|
||||
up_new = (sqrt_Sk[:, None] * Vh[:k, :]).contiguous() # (k, in_features)
|
||||
|
||||
# Cast back to original dtype
|
||||
down_dtype = compressed_sd[dk].dtype
|
||||
up_dtype = compressed_sd[uk].dtype
|
||||
down_new = down_new.to(dtype=down_dtype)
|
||||
up_new = up_new.to(dtype=up_dtype)
|
||||
|
||||
if transposed:
|
||||
down_new = down_new.T.contiguous()
|
||||
up_new = up_new.T.contiguous()
|
||||
|
||||
compressed_sd[dk] = down_new
|
||||
compressed_sd[uk] = up_new
|
||||
|
||||
# Update alpha if present
|
||||
alpha_key = f"{base}.alpha"
|
||||
if alpha_key in compressed_sd:
|
||||
alpha = compressed_sd[alpha_key]
|
||||
alpha = alpha * (k / rank)
|
||||
compressed_sd[alpha_key] = alpha
|
||||
|
||||
# Statistics
|
||||
total_original_params += out_features * rank + rank * in_features
|
||||
total_compressed_params += out_features * k + k * in_features
|
||||
rank_sum += k
|
||||
layer_count += 1
|
||||
svd_progress += 1
|
||||
|
||||
# Free SVD result tensors for this micro-batch
|
||||
del U_batch, S_batch, Vh_batch
|
||||
|
||||
# Yield to ComfyUI scheduler between micro-batches
|
||||
comfyui_yield()
|
||||
|
||||
# Clear GPU cache between shape groups
|
||||
cleanup_memory()
|
||||
|
||||
# Handle convolutional pairs (no SVD, just ensure contiguous)
|
||||
for dk, uk, base in conv_pairs:
|
||||
skipped_conv += 1
|
||||
print(f"ℹ️ Skipping convolutional layer {base} (shape {compressed_sd[dk].shape}, {compressed_sd[uk].shape})")
|
||||
svd_progress += 1
|
||||
|
||||
if skipped_conv > 0:
|
||||
print(f"ℹ️ Skipped {skipped_conv} convolutional layers (normal for SDXL/UNet).")
|
||||
|
||||
|
||||
# Compute size saved (assuming float32 = 4 bytes)
|
||||
size_saved_mb = (total_original_params - total_compressed_params) * 4 / (1024 * 1024)
|
||||
avg_final_rank = rank_sum / layer_count if layer_count > 0 else 0
|
||||
|
||||
|
||||
compression_info = {
|
||||
'compression_mode': compression_mode,
|
||||
'original_total_params': total_original_params,
|
||||
@@ -373,7 +468,7 @@ class MusubiLoraConverter:
|
||||
'size_saved_mb': size_saved_mb,
|
||||
'final_rank': avg_final_rank,
|
||||
}
|
||||
|
||||
|
||||
return compressed_sd, compression_info
|
||||
|
||||
RETURN_TYPES = ("LORA", "MODEL", "CLIP", "STRING", "STRING")
|
||||
@@ -429,21 +524,16 @@ class MusubiLoraConverter:
|
||||
"""Apply SVD compression if requested. Modifies info in-place.
|
||||
|
||||
Returns (normalized_sd, compression_info).
|
||||
Tensors are moved to the target device on-demand inside
|
||||
_apply_svd_compression (lazy, per-pair) instead of eagerly
|
||||
moving all tensors upfront — reduces PCIe traffic.
|
||||
"""
|
||||
target_device = DeviceManager.get_device(device)
|
||||
if compression_mode != "original" and target_device.type != "cpu":
|
||||
moved_count = 0
|
||||
for k in list(normalized_sd.keys()):
|
||||
if normalized_sd[k].device != target_device:
|
||||
normalized_sd[k] = normalized_sd[k].to(device=target_device)
|
||||
moved_count += 1
|
||||
if moved_count > 0:
|
||||
print(f"\U0001f527 Moved {moved_count} tensors to {target_device} for SVD compression")
|
||||
|
||||
if compression_mode != "original":
|
||||
print(f"\U0001f527 Applying SVD compression ({compression_mode})...")
|
||||
normalized_sd, compression_info = self._apply_svd_compression(
|
||||
normalized_sd, compression_mode, target_rank
|
||||
normalized_sd, compression_mode, target_rank, device=target_device
|
||||
)
|
||||
info.update(compression_info)
|
||||
print(f"\U0001f527 Compression applied. Final rank: {compression_info.get('final_rank', 'N/A'):.1f}, "
|
||||
|
||||
@@ -90,12 +90,22 @@ from .svd_quantizer import (
|
||||
)
|
||||
|
||||
# GGUF writer — block-wise quantization output for ComfyUI-GGUF
|
||||
from .gguf_writer import (
|
||||
GGUFSaveWriter,
|
||||
gguf_arch_from_arch,
|
||||
is_gguf_precision,
|
||||
GGUF_SUPPORTED_ARCHS,
|
||||
)
|
||||
# Guarded: if the 'gguf' package is not installed, GGUF features are disabled
|
||||
# but the rest of the node (precision conversion, stripping, SVD) still works.
|
||||
try:
|
||||
from .gguf_writer import (
|
||||
GGUFSaveWriter,
|
||||
gguf_arch_from_arch,
|
||||
is_gguf_precision,
|
||||
GGUF_SUPPORTED_ARCHS,
|
||||
)
|
||||
except ImportError:
|
||||
GGUFSaveWriter = None
|
||||
gguf_arch_from_arch = None # type: ignore[assignment]
|
||||
is_gguf_precision = lambda fmt: False # Always returns False — clean fallback
|
||||
GGUF_SUPPORTED_ARCHS: set = set()
|
||||
print("⚠️ Checkpoint Studio: GGUF output disabled — 'gguf' package not installed. "
|
||||
"Install it with: pip install gguf")
|
||||
|
||||
try:
|
||||
from .metadata_factory import finalize_metadata
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
import { app } from "../../scripts/app.js";
|
||||
import { ComfyWidgets } from "../../scripts/widgets.js";
|
||||
|
||||
console.log("[EasyLoRAMerger] JS module loaded — registering extensions", Date.now());
|
||||
|
||||
// Enhanced method descriptions
|
||||
const METHOD_DESCRIPTIONS = {
|
||||
@@ -210,3 +213,124 @@ app.registerExtension({
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// EasyTextDisplay — renders STRING input as a visible multiline widget
|
||||
//
|
||||
// Pattern follows proven ShowText (comfyui-custom-scripts) implementation:
|
||||
// 1. Determine if there's a hidden converted-widget; clear from there
|
||||
// 2. Create read-only widgets named "text_N" via ComfyWidgets
|
||||
// 3. On execution: onExecuted → populate(message.text)
|
||||
// 4. On workflow load: configure → store values, onConfigure → restore
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
app.registerExtension({
|
||||
name: "EasyLoRA.TextDisplay",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
||||
if (nodeData.name === "EasyTextDisplay") {
|
||||
console.log("[EasyLoRA.TextDisplay] Registering for EasyTextDisplay node");
|
||||
|
||||
/**
|
||||
* Populate the node with read-only display widgets
|
||||
* from the given text array.
|
||||
*
|
||||
* @param {string[]} text — e.g. ["hello world"]
|
||||
*/
|
||||
function populate(text) {
|
||||
try {
|
||||
if (this.widgets) {
|
||||
// On older frontend versions there is a hidden converted-widget
|
||||
const isConvertedWidget = +!!this.inputs?.[0]?.widget;
|
||||
for (let i = isConvertedWidget; i < this.widgets.length; i++) {
|
||||
this.widgets[i].onRemove?.();
|
||||
}
|
||||
this.widgets.length = isConvertedWidget;
|
||||
}
|
||||
|
||||
// Guard: if text is undefined/null, there's nothing to display
|
||||
if (text == null) {
|
||||
console.warn("[EasyLoRA.TextDisplay] populate called with null/undefined text");
|
||||
return;
|
||||
}
|
||||
|
||||
const v = [...text];
|
||||
if (!v[0]) {
|
||||
v.shift();
|
||||
}
|
||||
for (let list of v) {
|
||||
// Force list to be an array, not sure why sometimes it is/isn't
|
||||
if (!(list instanceof Array)) list = [list];
|
||||
for (const l of list) {
|
||||
const w = ComfyWidgets["STRING"](
|
||||
this,
|
||||
"text_" + (this.widgets?.length ?? 0),
|
||||
["STRING", { multiline: true }],
|
||||
app
|
||||
).widget;
|
||||
w.inputEl.readOnly = true;
|
||||
w.inputEl.style.opacity = 0.6;
|
||||
w.value = l;
|
||||
}
|
||||
}
|
||||
|
||||
// Auto-resize to fit content
|
||||
requestAnimationFrame(() => {
|
||||
const sz = this.computeSize();
|
||||
if (sz[0] < this.size[0]) sz[0] = this.size[0];
|
||||
if (sz[1] < this.size[1]) sz[1] = this.size[1];
|
||||
this.onResize?.(sz);
|
||||
app.graph.setDirtyCanvas(true, false);
|
||||
});
|
||||
} catch (err) {
|
||||
console.error("[EasyLoRA.TextDisplay] populate() crashed:", err);
|
||||
}
|
||||
}
|
||||
|
||||
// ── onExecuted: triggered when the node finishes running ──
|
||||
const onExecuted = nodeType.prototype.onExecuted;
|
||||
nodeType.prototype.onExecuted = function (message) {
|
||||
try {
|
||||
onExecuted?.apply(this, arguments);
|
||||
console.log("[EasyLoRA.TextDisplay] onExecuted received:", JSON.stringify(message));
|
||||
populate.call(this, message?.text);
|
||||
} catch (err) {
|
||||
console.error("[EasyLoRA.TextDisplay] onExecuted crashed:", err);
|
||||
}
|
||||
};
|
||||
|
||||
// ── Persist text across workflow saves/loads ──
|
||||
const VALUES = Symbol("textDisplayValues");
|
||||
const configure = nodeType.prototype.configure;
|
||||
nodeType.prototype.configure = function () {
|
||||
// Store unmodified widget values as they get removed on configure by new frontend
|
||||
this[VALUES] = arguments[0]?.widgets_values;
|
||||
return configure?.apply(this, arguments);
|
||||
};
|
||||
|
||||
const onConfigure = nodeType.prototype.onConfigure;
|
||||
nodeType.prototype.onConfigure = function () {
|
||||
try {
|
||||
onConfigure?.apply(this, arguments);
|
||||
const widgets_values = this[VALUES];
|
||||
if (widgets_values?.length) {
|
||||
console.log("[EasyLoRA.TextDisplay] onConfigure restoring:", JSON.stringify(widgets_values));
|
||||
// In newer frontend there seems to be a delay in creating the initial widget
|
||||
requestAnimationFrame(() => {
|
||||
try {
|
||||
populate.call(
|
||||
this,
|
||||
widgets_values.slice(
|
||||
+(widgets_values.length > 1 && this.inputs?.[0]?.widget)
|
||||
)
|
||||
);
|
||||
} catch (err) {
|
||||
console.error("[EasyLoRA.TextDisplay] onConfigure populate crashed:", err);
|
||||
}
|
||||
});
|
||||
}
|
||||
} catch (err) {
|
||||
console.error("[EasyLoRA.TextDisplay] onConfigure crashed:", err);
|
||||
}
|
||||
};
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
@@ -1563,7 +1563,7 @@ def compute_component_energy_ratios(
|
||||
rank = max(1, rank)
|
||||
|
||||
if original_sds[i] is not None:
|
||||
from ..engine.scale_utils import find_alpha_value
|
||||
from .engine.scale_utils import find_alpha_value
|
||||
alpha_value = find_alpha_value(original_sds[i], key, mapping=mappings[i])
|
||||
if not converted_flags[i] and alpha_value is not None:
|
||||
scale_factor = alpha_value / rank
|
||||
@@ -1668,7 +1668,7 @@ def compute_primary_driver_intensity_metric(
|
||||
|
||||
# Alpha/rank scaling for unconverted LoRAs (same logic as existing functions)
|
||||
if original_sds[i] is not None:
|
||||
from ..engine.scale_utils import find_alpha_value
|
||||
from .engine.scale_utils import find_alpha_value
|
||||
alpha_value = find_alpha_value(original_sds[i], key, mapping=mappings[i])
|
||||
if not converted_flags[i] and alpha_value is not None:
|
||||
scale_factor = alpha_value / rank
|
||||
|
||||
Reference in New Issue
Block a user