diff --git a/README.md b/README.md index 85c2bf2..711a378 100644 --- a/README.md +++ b/README.md @@ -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. --- diff --git a/__init__.py b/__init__.py index 9ff73a9..7a2137b 100644 --- a/__init__.py +++ b/__init__.py @@ -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'] diff --git a/assets/saving.jpg b/assets/saving.jpg new file mode 100644 index 0000000..39eef72 Binary files /dev/null and b/assets/saving.jpg differ diff --git a/config.py b/config.py index b410cfc..0e6b3e1 100644 --- a/config.py +++ b/config.py @@ -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 ── diff --git a/easy_lora_merger.py b/easy_lora_merger.py index 91a11f0..371f827 100644 --- a/easy_lora_merger.py +++ b/easy_lora_merger.py @@ -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).") diff --git a/easy_text_display.py b/easy_text_display.py new file mode 100644 index 0000000..fdec684 --- /dev/null +++ b/easy_text_display.py @@ -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", +} diff --git a/engine/checkpoint_normalizer.py b/engine/checkpoint_normalizer.py index 5f639aa..7a3876e 100644 --- a/engine/checkpoint_normalizer.py +++ b/engine/checkpoint_normalizer.py @@ -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"], diff --git a/engine/gguf_writer.py b/engine/gguf_writer.py index 056993b..8f97e65 100644 --- a/engine/gguf_writer.py +++ b/engine/gguf_writer.py @@ -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, diff --git a/engine/key_mapper.py b/engine/key_mapper.py index 631094e..ab024fc 100644 --- a/engine/key_mapper.py +++ b/engine/key_mapper.py @@ -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 ─────────────────────────────── diff --git a/engine/lora_studio_converter.py b/engine/lora_studio_converter.py index 2f7cbaa..6bb4c2b 100644 --- a/engine/lora_studio_converter.py +++ b/engine/lora_studio_converter.py @@ -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}, " diff --git a/engine/musubi_checkpoint_studio.py b/engine/musubi_checkpoint_studio.py index 4dc2a9e..9a4c7e3 100644 --- a/engine/musubi_checkpoint_studio.py +++ b/engine/musubi_checkpoint_studio.py @@ -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 diff --git a/js/easy_lora_merger.js b/js/easy_lora_merger.js index 94c5053..c0dcd55 100644 --- a/js/easy_lora_merger.js +++ b/js/easy_lora_merger.js @@ -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); + } + }; + } + }, +}); diff --git a/utils.py b/utils.py index e1c94df..9256788 100644 --- a/utils.py +++ b/utils.py @@ -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