""" Easy LoRA Baker — ComfyUI Node Definition Bakes a LoRA (or merged LoRA) into a full model checkpoint at the tensor level, producing ComfyUI MODEL+CLIP+VAE objects directly. Lazy assembly via mmap means near-zero RAM overhead for preserved keys — no temp files needed. Features: - VAE passthrough & output - Per-component scaling (weight_unet, weight_te, weight_clip, weight_vae) - Lazy assembly: baked keys in write cache, preserved keys mmap'd from file - Metadata modes: none / preserve_a / merge_basic - Expanded precision: auto / float32 / bfloat16 / float16 / fp8_e4m3fn / fp8_e5m2 """ import torch import json import time _T0 = time.time() from pathlib import Path from collections import Counter import folder_paths import comfy.sd from .utils import ( load_checkpoint_with_metadata, load_lora_with_metadata, read_safetensors_header_only, cleanup_memory, NodeCache, save_safetensors_file, save_safetensors_stream, load_state_dict_as_model_objects, get_available_ram, get_combined_model_list, resolve_model_path, ) from .config import PRECISION_EXTENDED, DEVICE_OPTIONS, DevicePrecisionConfig from .engine.baking_processor import SmartBakingProcessor # =================================================================== # FP8 metadata helpers — imported from shared fp8_quantizer module # (One Pattern Per Concept — working_principles.md) # =================================================================== from .engine.fp8_quantizer import ( build_fp8_quantization_metadata as _build_fp8_quantization_metadata, strip_input_scale_keys as _strip_input_scale_keys, ) # =================================================================== # Auto-precision helper: detect checkpoint's native weight dtype # =================================================================== def _detect_native_ckpt_dtype(ckpt_header: dict) -> "torch.dtype | None": """Detect the dominant weight dtype from the safetensors header. Scans all tensor entries in the header, skipping scale-factor keys (.weight_scale, .input_scale), and finds the most common floating-point dtype. Only float types (float16, bfloat16, float32, float8_e4m3fn, float8_e5m2) are considered — int/bool tensors are excluded since they are not precision targets for weight computation. Returns: The most common torch.dtype among float weight tensors, or None if no float tensors are found. """ from .engine.key_utils import SAFETENSORS_DTYPE_MAP as _SAFETENSORS_DTYPE_MAP _FLOAT_DTYPES = { torch.float16, torch.bfloat16, torch.float32, torch.float8_e4m3fn, torch.float8_e5m2, } dtype_counter: Counter = Counter() for key, info in ckpt_header.items(): if not isinstance(info, dict): continue if key.endswith((".weight_scale", ".input_scale", ".comfy_quant")): continue dtype_str = info.get("dtype", "") if not dtype_str: continue dtype = _SAFETENSORS_DTYPE_MAP.get(dtype_str) if dtype is not None and dtype in _FLOAT_DTYPES: dtype_counter[dtype] += 1 if dtype_counter: return dtype_counter.most_common(1)[0][0] return None class SmartModelBaker: """ ComfyUI node: Bakes a LoRA into a checkpoint at the tensor level. Inputs: - checkpoint: Pick from available checkpoints dropdown - lora_data: LORA type from Triple Merger or LoRA-Only Merger (optional) - lora_name: Alternate single LoRA from dropdown (optional) - baking_method: linear / impact_weighted / orthogonal - strength: LoRA strength 0.0–2.0 - weight_unet: Per-component weight for U-Net keys (0.0–2.0) - weight_te: Per-component weight for Text Encoder keys (0.0–2.0) - weight_clip: Per-component weight for CLIP Vision keys (0.0–2.0) - weight_vae: Per-component weight for VAE keys (0.0–2.0) - metadata_mode: "none" / "preserve_a" / "merge_basic" - precision: auto / float32 / bfloat16 / float16 / fp8_e4m3fn / fp8_e5m2 - save_trigger: Toggle to bake and save to disk permanently. When True, MODEL+CLIP+VAE are lazy-loaded from the saved file for low RAM usage. When False, baked result is kept in memory with automatic RAM Guard fallback. - save_folder: Custom output folder - filename: Output filename (.safetensors added automatically) - device: auto / cuda / cpu Outputs: - model: MODEL - clip: CLIP - vae: VAE - output: STRING path to saved baked checkpoint (empty in preview mode) - forensic_report: STRING Human-readable forensic report (Lora Studio style) """ @classmethod def INPUT_TYPES(cls): checkpoints = get_combined_model_list() loras = folder_paths.get_filename_list("loras") return { "required": {}, "optional": { # ── SECTION 1: INPUTS ─────────────────────────────────── "checkpoint": (checkpoints,), # Wired LORA overrides dropdown — when lora_data is connected, lora_name is ignored "lora_data": ("LORA", { "tooltip": "LORA output from Triple Merger or LoRA-Only Merger. When connected, overrides lora_name dropdown." }), "lora_name": (["None"] + loras, { "tooltip": "Pick a single LoRA from dropdown. Used only when lora_data is not connected." }), # ── SECTION 3: SETTINGS ───────────────────────────────── "baking_method": ( ["linear", "impact_weighted", "orthogonal"], {"default": "linear"}, ), "strength": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "LoRA strength multiplier" }), "weight_unet": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "Per-component weight for U-Net (diffusion model) keys" }), "weight_te": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "Per-component weight for Text Encoder keys" }), "weight_clip": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "Per-component weight for CLIP Vision keys" }), "weight_vae": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "Per-component weight for VAE keys" }), # ── SECTION 5: HARDWARE ───────────────────────────────── "device": (DEVICE_OPTIONS, {"default": "auto"}), "precision": (PRECISION_EXTENDED, {"default": "auto"}), "batch_size": ("INT", { "default": 64, "min": 1, "max": 256, "step": 8, "tooltip": "Number of keys to process per batch. Larger = faster but more VRAM. " "memory_guard() runs between batches to prevent OOM. " "DeviceManager.suggest_batch_size() can auto-tune based on VRAM." }), # ── SECTION 6: OUTPUT ─────────────────────────────────── "save_trigger": ("BOOLEAN", { "default": False, "tooltip": "When True: bake LoRA into checkpoint AND save as permanent .safetensors file (low RAM: saves first, then lazy-loads from file). When False: bake in-memory only (preview mode) with automatic RAM Guard fallback — connect MODEL+CLIP+VAE outputs downstream to test results before committing." }), "filename": ("STRING", { "default": "baked_checkpoint", "multiline": False, "tooltip": "Output filename (.safetensors added automatically in disk mode)" }), "save_folder": ("STRING", { "default": "", "multiline": False, "tooltip": "Custom output folder (leave empty for default checkpoints dir)" }), "metadata_mode": ( ["preserve_a", "preserve_b", "none", "merge_basic"], {"default": "preserve_a", "tooltip": "'none'=baking only, 'preserve_a'=original priority, 'preserve_b'=second-source (fallback to A), 'merge_basic'=baking priority"} ), }, "hidden": { "node_id": "UNIQUE_ID", }, } RETURN_TYPES = ("MODEL", "CLIP", "VAE", "STRING", "STRING") RETURN_NAMES = ("model", "clip", "vae", "output_path", "forensic_report") FUNCTION = "bake" CATEGORY = "LoRA/Baking" @classmethod def IS_CHANGED(cls, **kwargs): """Cache implementation for ComfyUI change detection.""" # Inline data guard: wired lora_data can't be cheaply hashed if kwargs.get('lora_data') is not None: return float("nan") return NodeCache.is_changed( cls.__name__, checkpoint=kwargs.get('checkpoint', ''), lora_name=kwargs.get('lora_name', 'None'), baking_method=kwargs.get('baking_method', 'linear'), strength=kwargs.get('strength', 1.0), weight_unet=kwargs.get('weight_unet', 1.0), weight_te=kwargs.get('weight_te', 1.0), weight_clip=kwargs.get('weight_clip', 1.0), weight_vae=kwargs.get('weight_vae', 1.0), metadata_mode=kwargs.get('metadata_mode', 'preserve_a'), save_trigger=kwargs.get('save_trigger', False), device=kwargs.get('device', 'auto'), precision=kwargs.get('precision', 'auto'), batch_size=kwargs.get('batch_size', 64), ) def bake( self, checkpoint: str, baking_method: str = "linear", strength: float = 1.0, weight_unet: float = 1.0, weight_te: float = 1.0, weight_clip: float = 1.0, weight_vae: float = 1.0, batch_size: int = 64, metadata_mode: str = "preserve_a", save_trigger: bool = False, lora_data=None, lora_name: str = "None", save_folder: str = "", filename: str = "baked_checkpoint", device: str = "auto", precision: str = "auto", **kwargs, ): """ Execute the baking pipeline with per-component scaling, RAM mode with automatic RAM Guard fallback, VAE output, and metadata control. Returns (model, clip, vae, output, forensic_report). """ print("\n" + "=" * 60) print("🔥 Easy LoRA Baker — Baking LoRA into Checkpoint") print("=" * 60) print(f" Checkpoint: {checkpoint}") print(f" Method: {baking_method}, Strength: {strength}") print(f" 🧱 Weights — UNet:{weight_unet} TE:{weight_te} CLIP:{weight_clip} VAE:{weight_vae}") print(f" 💾 Preview: RAM mode with RAM Guard, Metadata: {metadata_mode}") # ══ Runtime precision guard — protect against stale workflow JSONs ══ if precision not in PRECISION_EXTENDED: print(f" ⚠️ Precision '{precision}' is no longer available, falling back to 'auto'") precision = "auto" # Default return values save_path = "" model_out = None clip_out = None vae_out = None forensic_report = "" # Validate inputs if not checkpoint: error_msg = "❌ No checkpoint selected — please choose a checkpoint from the dropdown" print(error_msg) forensic_report = json.dumps({"error": error_msg}, indent=2) return (model_out, clip_out, vae_out, save_path, forensic_report) if lora_data is None and (lora_name == "None" or not lora_name): error_msg = "❌ Either lora_data or lora_name must be provided" print(error_msg) forensic_report = json.dumps({"error": error_msg}, indent=2) return (model_out, clip_out, vae_out, save_path, forensic_report) # ------------------------------------------------------------------ # Step 1: Load checkpoint # ------------------------------------------------------------------ print("\n📥 Loading checkpoint...") try: ckpt_path = resolve_model_path(checkpoint) if not ckpt_path: # Try direct path ckpt_path = checkpoint ckpt_path = Path(ckpt_path) # RAM Guard (informational): Log file size for diagnostics. # With mmap-based lazy loading, the checkpoint file is never fully # loaded into RAM. Only individual tensors are loaded on demand via # safe_open (mmap). A 17 GiB checkpoint uses ~50 KB for the header. ckpt_file_size = ckpt_path.stat().st_size print(f" 📦 Checkpoint file: {ckpt_file_size / (1024**3):.2f} GB on disk") # Use mmap-based lazy loading: read only the safetensors header JSON # (~50 KB), then wrap the file in a dict-like lazy mapping. # Tensors are loaded on demand when ckpt_sd[key] is first accessed. ckpt_header, ckpt_metadata = read_safetensors_header_only(ckpt_path) # 🔥 FP8 DETECTION: Check if any checkpoint tensor uses FP8 dtype. # FP8 e4m3fn has only 3 mantissa bits — step size ~0.125 at magnitude # 1.0, which is 64× coarser than a typical LoRA delta (~0.002). # We dequantize to bfloat16 for all baking computation, then # quantize once at final output (if the user selected fp8 precision). # Safetensors header uses short dtype strings: "F8_E4M3" for fp8_e4m3fn, # "F8_E5M2" for fp8_e5m2 (see _SAFETENSORS_DTYPE_MAP in musubi_checkpoint_studio.py). # Check for the "F8_" prefix which is unique to FP8 dtypes. detected_fp8 = any( info.get('dtype', '').startswith('F8_') for info in ckpt_header.values() if isinstance(info, dict) ) if detected_fp8: print(f" 🚀 FP8 checkpoint detected — dequantizing to bfloat16 for baking precision") # Strip stale quantization metadata that tells ComfyUI to use # its slow mixed-precision path. Since we dequantize to bf16, # this metadata is stale and would trigger incorrect behavior. # Secondary cleanup also happens in _build_metadata (Fix 6). fp8_meta_keys = [k for k in list(ckpt_metadata.keys()) if 'quantization' in k.lower()] for k in fp8_meta_keys: del ckpt_metadata[k] if fp8_meta_keys: print(f" 🧹 Stripped {len(fp8_meta_keys)} quantization metadata keys") from .engine.musubi_checkpoint_studio import MusubiCheckpointStudio # 🔥 Pass target_dtype=bf16 for FP8 checkpoints so that __getitem__ # applies per-channel scale factors during fp8→bf16 dequantization. # Without this, _safe_bake_add receives raw fp8 tensors and its naive # base.to(dtype=compute_dtype) cast ignores the per-channel scale # factors, producing completely wrong weight values (pure noise). ckpt_sd = MusubiCheckpointStudio._LazyCheckpointMapping( ckpt_path, ckpt_metadata, target_dtype=torch.bfloat16 if detected_fp8 else None, ) print(f" ✅ Lazily mapped checkpoint: {len(ckpt_header)} keys via mmap (0 bytes loaded)") except Exception as e: error_msg = f"❌ Failed to load checkpoint: {e}" print(error_msg) forensic_report = json.dumps({"error": error_msg}, indent=2) return (model_out, clip_out, vae_out, save_path, forensic_report) # ------------------------------------------------------------------ # Step 2: Load LoRA # ------------------------------------------------------------------ print("\n📥 Loading LoRA...") lora_sd = None lora_source = "unknown" try: if lora_data is not None: # LORA type from Triple Merger or LoRA-Only Merger # lora_data is a tuple: (state_dict, weight, weight) if isinstance(lora_data, tuple) and len(lora_data) >= 1: lora_sd = lora_data[0] elif isinstance(lora_data, dict): lora_sd = lora_data lora_source = "lora_data (from merger)" print(f" ✅ Using lora_data from upstream node ({len(lora_sd)} tensors)") # Apply ComfyUI's convert_lora for non-standard LoRA formats # (BFL control, WanFun, USO — standard SD1.5 Kohya passes through) # Must match ComfyUI's ordering: convert_lora BEFORE key matching # See: comfy/sd.py:85 → comfy.lora_convert.convert_lora(lora) try: from comfy.lora_convert import convert_lora lora_sd = convert_lora(lora_sd) print(f" 🔄 Applied convert_lora ({len(lora_sd)} tensors)") except ImportError: pass # Fallback: comfy not available in this context if lora_sd is None and lora_name != "None" and lora_name: # Load from dropdown pick lora_path = folder_paths.get_full_path("loras", lora_name) if lora_path: lora_sd, _ = load_lora_with_metadata(Path(lora_path)) lora_source = lora_name print(f" ✅ Loaded LoRA: {lora_name} ({len(lora_sd)} tensors)") if lora_sd is None: error_msg = "❌ Failed to load LoRA data" print(error_msg) forensic_report = json.dumps({"error": error_msg}, indent=2) return (model_out, clip_out, vae_out, save_path, forensic_report) # Filter out non-tensor metadata from LoRA state dict lora_sd = {k: v for k, v in lora_sd.items() if isinstance(v, torch.Tensor)} except Exception as e: error_msg = f"❌ Failed to load LoRA: {e}" print(error_msg) forensic_report = json.dumps({"error": error_msg}, indent=2) return (model_out, clip_out, vae_out, save_path, forensic_report) # ------------------------------------------------------------------ # Step 3-9: Run baking pipeline # ------------------------------------------------------------------ print("\n🔨 Running baking pipeline...") # 🎯 Auto-precision: detect checkpoint native dtype and use it if precision == "auto": native_dtype = _detect_native_ckpt_dtype(ckpt_header) if native_dtype is not None: _DTYPE_TO_PRECISION = { torch.float8_e4m3fn: "fp8_e4m3fn", torch.float8_e5m2: "fp8_e5m2", torch.bfloat16: "bfloat16", torch.float16: "float16", torch.float32: "float32", } mapped = _DTYPE_TO_PRECISION.get(native_dtype) if mapped: precision = mapped print(f" 🎯 Auto precision: detected {native_dtype} → using {mapped}") try: dpc = DevicePrecisionConfig(device_type=device, precision=precision) processor = SmartBakingProcessor(device_precision=dpc) output_sd, metadata, forensic_report, _ = processor.bake( ckpt_sd=ckpt_sd, ckpt_header=ckpt_header, lora_sd=lora_sd, baking_method=baking_method, batch_size=batch_size, strength=strength, lora_source=lora_source, checkpoint_name=checkpoint, original_metadata=ckpt_metadata, weight_unet=weight_unet, weight_te=weight_te, weight_clip=weight_clip, weight_vae=weight_vae, metadata_mode=metadata_mode, detected_fp8=detected_fp8, ) print(f" 🕐 [t={time.time()-_T0:.1f}s] Baking pipeline returned to baker_node — output_sd type: {type(output_sd).__name__}") except Exception as e: error_msg = f"❌ Baking pipeline failed: {e}" print(error_msg) import traceback traceback.print_exc() forensic_report = json.dumps({"error": error_msg}, indent=2) return (model_out, clip_out, vae_out, save_path, forensic_report) # ------------------------------------------------------------------ # Step 10-11: Output as MODEL+CLIP+VAE with conditional save # ------------------------------------------------------------------ print(f"\n💾 save_trigger={save_trigger}") # ── Measure pre-load RAM for post-load delta ── _pre_load_ram = get_available_ram() if save_trigger: # === Save Mode: permanent disk save + lazy load from file === from .engine.musubi_checkpoint_studio import MusubiCheckpointStudio # 🎯 FP8 SAVE: Build _quantization_metadata for MixedPrecisionOps. # This runs before the if-else dispatch so ALL save paths benefit. metadata = _build_fp8_quantization_metadata(output_sd, metadata, label="save") if isinstance(output_sd, MusubiCheckpointStudio._LazyCheckpointMapping): # ── Lazy assembly mode (Fix 1): baked keys in write cache ── # output_sd.filepath points to the ORIGINAL checkpoint, NOT the # baked output. Must materialize all keys to a new file. print("Saving baked checkpoint to disk (lazy assembly — materializing to file)...") from .utils import _get_output_path output_path = _get_output_path( save_folder, filename, "baked_checkpoint", "checkpoints" ) output_path.parent.mkdir(parents=True, exist_ok=True) # Materialize all keys (write cache + lazy reads from original file) all_tensors = dict(output_sd.items()) # 🎯 FP8 SAVE: Strip .input_scale (MixedPrecisionOps only needs .weight_scale) _strip_input_scale_keys(all_tensors) save_safetensors_stream(all_tensors, output_path, metadata=metadata) del all_tensors save_path = str(output_path) file_size_mb = output_path.stat().st_size / (1024 * 1024) print(f" ✅ Saved: {output_path} ({file_size_mb:.1f} MB)") # Fix 2: Close the lazy mapping permanently — no more reads from original output_sd.permanent_close() cleanup_memory() else: # ── Dict mode: use existing save_safetensors_file path ── print("Saving baked checkpoint to disk permanently...") output_path, save_path = save_safetensors_file( output_sd, save_folder, filename, metadata=metadata, folder_type="checkpoints", default_name="baked_checkpoint", ) if output_path is None: # save_safetensors_file already printed the error forensic_report = json.dumps({"error": save_path}, indent=2) return (None, None, None, save_path, forensic_report) # Free in-memory state dict BEFORE loading model objects (low RAM) print(" 🧹 Freeing in-memory state dict...") cleanup_memory(output_sd) # ── Common model loading for all save modes (Fix 5) ── print(f"\n📦 Loading baked checkpoint as MODEL+CLIP+VAE (lazy from file)...") lazy_mapping = None try: lazy_mapping = MusubiCheckpointStudio._LazyCheckpointMapping( output_path, metadata, target_dtype=None, # 🎯 MixedPrecisionOps handles dequant via _quantization_metadata ) model_out, clip_out, vae_out = load_state_dict_as_model_objects( lazy_mapping, metadata=metadata, ) print(f" ✅ Loaded MODEL + CLIP + VAE from saved file (lazy, low RAM)") except Exception as e: print(f" ⚠️ Failed to load baked checkpoint as model objects: {e}") print(f" ℹ️ The checkpoint was saved successfully at: {save_path}") finally: # Fix 5: Close the mapping permanently after model objects are built. # This prevents __del__ from reopening the file handle during GC, # eliminating post-prompt file I/O on stale lazy mappings. if lazy_mapping is not None: lazy_mapping.permanent_close() else: # === Preview Mode: load as MODEL+CLIP+VAE directly === # output_sd is a lazy mapping (from _assemble_output_lazy) or a real dict # (from _assemble_output for non-lazy mode). Lazy mappings have near-zero # RAM overhead — no temp file fallback needed. print("Preview mode — loading as MODEL+CLIP+VAE...") try: # 🎯 FP8 PREVIEW: Use MixedPrecisionOps instead of eager dequant. # Unified detection: works for BOTH lazy mappings (has _user_fp8_mode flag) # and real dicts (check for FP8 tensor dtype). _has_fp8 = getattr(output_sd, '_user_fp8_mode', False) or any( hasattr(v, 'dtype') and v.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for v in output_sd.values() ) if _has_fp8: # Step 1: Pop .input_scale keys (MixedPrecisionOps # only needs .weight_scale — .input_scale would appear as # 'unet unexpected' keys) _strip_input_scale_keys(output_sd) # Step 2: Build _quantization_metadata from the mapping. metadata = _build_fp8_quantization_metadata(output_sd, metadata, label="preview") print(f" 🕐 [t={time.time()-_T0:.1f}s] Starting model loading — load_state_dict_as_model_objects") model_out, clip_out, vae_out = load_state_dict_as_model_objects( output_sd, metadata=metadata, ) mode_label = "lazy (mmap)" if hasattr(output_sd, 'filepath') else "in-memory" print(f" ✅ Loaded MODEL + CLIP + VAE from baked state dict ({mode_label})") print(f" 🕐 [t={time.time()-_T0:.1f}s] Model loading complete — MODEL+CLIP+VAE ready") # Free state dict after loading to recover memory cleanup_memory(output_sd) # Close lazy mapping to prevent stale file handles if hasattr(output_sd, 'filepath'): output_sd.permanent_close() except Exception as e: error_msg = f"❌ In-memory model load failed: {e}" print(error_msg) import traceback traceback.print_exc() forensic_report = json.dumps({"error": error_msg}, indent=2) return (None, None, None, "", forensic_report) save_path = "" # No permanent save path in preview mode # SSD Fix: skip gc.collect() at end of bake — GC already ran naturally # during the pipeline. Calling gc.collect() here triggers __del__ on # stale objects at unpredictable times, causing unexpected file I/O. cleanup_memory(skip_gc=True) # ── Measure post-load RAM and print delta ── _post_load_ram = get_available_ram() if _pre_load_ram is not None and _post_load_ram is not None: _delta_gb = (_pre_load_ram - _post_load_ram) / (1024**3) print(f"📊 RAM: Pre-load {_pre_load_ram / (1024**3):.2f} GB → " f"Post-load {_post_load_ram / (1024**3):.2f} GB " f"(delta: {_delta_gb:+.2f} GB)") elif _post_load_ram is not None: print(f"📊 Post-load available RAM: {_post_load_ram / (1024**3):.2f} GB") print("\n" + "=" * 60) print("✅ Easy LoRA Baker complete!") if save_trigger: print(f" 💾 Saved to: {save_path}") if model_out is not None: print(f" 📦 Loaded MODEL + CLIP + VAE from saved file (lazy)") else: print(f" 🧠 Preview mode — no disk save") if model_out is not None: print(f" 📦 Loaded baked MODEL + CLIP + VAE in memory") print(f" ℹ️ Set save_trigger=True to save permanently to disk") print("=" * 60) print(f" 🕐 [t={time.time()-_T0:.1f}s] baker_node.bake returning") return (model_out, clip_out, vae_out, save_path, forensic_report)