775 lines
38 KiB
Python
775 lines
38 KiB
Python
"""
|
|
Easy LoRA Baker — Core tensor-level LoRA baking engine (orchestration layer).
|
|
|
|
Rewritten as a thin orchestration core that imports specialized functions
|
|
from 5 sub-modules:
|
|
- baking_processor_constants: Registry data (prefix mappings, patterns)
|
|
- baking_processor_delta: Delta reconstruction, alpha/rank scaling
|
|
- baking_processor_baking: Bake methods, shape alignment, assembly
|
|
- baking_processor_matching: Key matching, reverse map, strategy cascade
|
|
- baking_processor_sd15: SD1.5 auto-scale channel safety only
|
|
(old SD1.5 matching removed — bridge handles it)
|
|
|
|
The SmartBakingProcessor class delegates all method calls to module-level
|
|
functions assigned as class methods. External API (baker_node.py import)
|
|
is unchanged: SmartBakingProcessor still lives in engine/baking_processor.py.
|
|
"""
|
|
|
|
import torch
|
|
import time
|
|
_T0 = time.time()
|
|
import statistics
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple, Set, Any
|
|
from collections import Counter
|
|
|
|
from ..utils import (
|
|
DeviceManager,
|
|
ProgressTracker,
|
|
categorize_checkpoint_key,
|
|
comfyui_yield,
|
|
memory_optimized_merge,
|
|
get_available_ram,
|
|
cleanup_memory,
|
|
check_ram_guard,
|
|
)
|
|
from ..config import DevicePrecisionConfig
|
|
|
|
# =====================================================================
|
|
# Imports from sub-modules
|
|
# =====================================================================
|
|
|
|
# Delta reconstruction — lazy and eager
|
|
from .baking_processor_delta import (
|
|
_reconstruct_lora_delta,
|
|
_LazyDeltaDict,
|
|
_MatchedDeltas,
|
|
_group_lora_pairs,
|
|
_ShapeInfo,
|
|
)
|
|
|
|
# FP8 quantizer — scale-then-cast, BF16 preservation, companion scales
|
|
from .fp8_quantizer import (
|
|
quantize_to_fp8,
|
|
compute_weight_scale,
|
|
should_preserve_bf16,
|
|
quantize_weight_to_fp8_with_scales,
|
|
)
|
|
|
|
# Baking methods and assembly
|
|
from .baking_processor_baking import (
|
|
_check_shape_compatible,
|
|
_align_delta_shape,
|
|
_align_delta_safe,
|
|
_orthogonal_projection_2d,
|
|
_orthogonal_projection,
|
|
_get_component_weight,
|
|
bake_linear,
|
|
bake_impact_weighted,
|
|
bake_orthogonal,
|
|
_assemble_output,
|
|
_assemble_output_lazy,
|
|
_sanitize_or_revert,
|
|
_sanitize_baked,
|
|
_build_metadata,
|
|
_build_forensic_report,
|
|
)
|
|
|
|
# Key matching
|
|
from .baking_processor_matching import (
|
|
_lora_key_to_checkpoint_key,
|
|
_try_underscore_to_dot_conversion,
|
|
_try_te_conditioner_prefix,
|
|
_build_reverse_key_map,
|
|
_is_spatial_depth_block,
|
|
_extract_block_id,
|
|
_build_shape_block_index, # NEW: O(1) shape+block index for Strategy 4
|
|
_build_shape_index, # NEW: O(1) shape index for Strategy 5
|
|
_global_search_by_shape_block,
|
|
_categorize_lora_key,
|
|
_create_effect_bar,
|
|
_find_matching_keys,
|
|
)
|
|
|
|
# SD1.5 Diffusers Bridge
|
|
from .diffusers_bridge import (
|
|
detect_native_sd15_checkpoint,
|
|
detect_diffusers_sd15_lora,
|
|
normalize_diffusers_preserving,
|
|
build_diffusers_reverse_map,
|
|
match_diffusers_deltas,
|
|
)
|
|
|
|
|
|
# =====================================================================
|
|
# SmartBakingProcessor — Orchestration Core
|
|
# =====================================================================
|
|
|
|
class SmartBakingProcessor:
|
|
"""
|
|
Performs tensor-level LoRA baking into full model checkpoints.
|
|
|
|
Workflow:
|
|
1. Load checkpoint state dict (with metadata)
|
|
2. Load LoRA state dict (from lora_data tuple or temp file)
|
|
3. Apply alpha/rank scaling correction to LoRA tensors
|
|
4. Reconstruct LoRA deltas from A/B pairs: delta = B @ A
|
|
5. Match LoRA keys to checkpoint keys
|
|
6. Apply selected baking method (linear/impact_weighted/orthogonal)
|
|
7. Apply safety features (preservation mask, TE toggle)
|
|
8. Assemble final state dict with VAE pass-through
|
|
9. Embed baking metadata history
|
|
10. Save as .safetensors file
|
|
|
|
All method implementations are delegated to module-level functions
|
|
from the baking_processor_* sub-modules.
|
|
"""
|
|
|
|
# ------------------------------------------------------------------
|
|
# Static method assignments (pure functions, no self)
|
|
# ------------------------------------------------------------------
|
|
_reconstruct_lora_delta = staticmethod(_reconstruct_lora_delta)
|
|
|
|
_check_shape_compatible = staticmethod(_check_shape_compatible)
|
|
_align_delta_shape = staticmethod(_align_delta_shape)
|
|
_align_delta_safe = staticmethod(_align_delta_safe)
|
|
_orthogonal_projection_2d = staticmethod(_orthogonal_projection_2d)
|
|
_orthogonal_projection = staticmethod(_orthogonal_projection)
|
|
_get_component_weight = staticmethod(_get_component_weight)
|
|
|
|
_assemble_output = staticmethod(_assemble_output)
|
|
_assemble_output_lazy = staticmethod(_assemble_output_lazy)
|
|
_sanitize_or_revert = staticmethod(_sanitize_or_revert)
|
|
_sanitize_baked = staticmethod(_sanitize_baked)
|
|
_build_metadata = staticmethod(_build_metadata)
|
|
_build_forensic_report = staticmethod(_build_forensic_report)
|
|
|
|
_lora_key_to_checkpoint_key = staticmethod(_lora_key_to_checkpoint_key)
|
|
_try_underscore_to_dot_conversion = staticmethod(_try_underscore_to_dot_conversion)
|
|
_try_te_conditioner_prefix = staticmethod(_try_te_conditioner_prefix)
|
|
_build_reverse_key_map = staticmethod(_build_reverse_key_map)
|
|
_is_spatial_depth_block = staticmethod(_is_spatial_depth_block)
|
|
_extract_block_id = staticmethod(_extract_block_id)
|
|
_global_search_by_shape_block = staticmethod(_global_search_by_shape_block)
|
|
_categorize_lora_key = staticmethod(_categorize_lora_key)
|
|
_create_effect_bar = staticmethod(_create_effect_bar)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Instance method assignments (functions that take self)
|
|
# ------------------------------------------------------------------
|
|
bake_linear = bake_linear
|
|
bake_impact_weighted = bake_impact_weighted
|
|
bake_orthogonal = bake_orthogonal
|
|
_find_matching_keys = _find_matching_keys
|
|
|
|
# ------------------------------------------------------------------
|
|
# Constructor
|
|
# ------------------------------------------------------------------
|
|
|
|
def __init__(self, device: str = "auto", precision: str = "auto",
|
|
verbose: bool = False,
|
|
device_precision: Optional['DevicePrecisionConfig'] = None):
|
|
if device_precision is not None:
|
|
self.device = device_precision.device
|
|
self.dtype = device_precision.dtype
|
|
else:
|
|
self.device = DeviceManager.get_device(device)
|
|
self.dtype = DeviceManager.get_dtype(precision, self.device)
|
|
self._verbose = verbose
|
|
|
|
# ------------------------------------------------------------------
|
|
# Early-abort helper: unifies the two nearly-identical error returns
|
|
# ------------------------------------------------------------------
|
|
|
|
def _abort_bake(
|
|
self,
|
|
error_msg: str,
|
|
ckpt_sd: Dict[str, torch.Tensor],
|
|
baking_method: str,
|
|
strength: float,
|
|
lora_source: str,
|
|
checkpoint_name: str,
|
|
original_metadata: Optional[Dict[str, str]],
|
|
weight_unet: float,
|
|
weight_te: float,
|
|
weight_clip: float,
|
|
weight_vae: float,
|
|
metadata_mode: str,
|
|
) -> Tuple[Dict[str, torch.Tensor], Dict[str, str], str, Dict[str, Any]]:
|
|
"""Early-abort return with error profile — preserves metadata/forensic shape."""
|
|
print(f"❌ {error_msg}")
|
|
# CRITICAL: When ckpt_sd is a _LazyCheckpointMapping, dict(ckpt_sd) would
|
|
# trigger .items() which loads ALL tensors via mmap (17 GiB for Klein 9B).
|
|
# Return the lazy mapping as-is (no tensor data loaded) or an empty dict.
|
|
if hasattr(ckpt_sd, 'filepath'):
|
|
output_for_abort: Dict[str, torch.Tensor] = {} # type: ignore[assignment]
|
|
else:
|
|
output_for_abort = dict(ckpt_sd)
|
|
return (
|
|
output_for_abort,
|
|
self._build_metadata(baking_method, strength,
|
|
lora_source, checkpoint_name,
|
|
0, len(ckpt_sd), original_metadata,
|
|
weight_unet, weight_te, weight_clip, weight_vae,
|
|
metadata_mode),
|
|
self._build_forensic_report(baking_method, strength,
|
|
lora_source, checkpoint_name,
|
|
0, len(ckpt_sd), len(ckpt_sd),
|
|
weight_unet=weight_unet, weight_te=weight_te,
|
|
weight_clip=weight_clip, weight_vae=weight_vae,
|
|
component_breakdown=getattr(self, '_component_breakdown', None),
|
|
delta_analysis=getattr(self, '_delta_analysis', None)),
|
|
{"error": error_msg},
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Main Entry Point: bake
|
|
# ------------------------------------------------------------------
|
|
|
|
def bake(
|
|
self,
|
|
ckpt_sd: Dict[str, torch.Tensor],
|
|
lora_sd: Dict[str, torch.Tensor],
|
|
ckpt_header: Optional[Dict[str, Any]] = None,
|
|
baking_method: str = "linear",
|
|
strength: float = 1.0,
|
|
energy_concentration: float = 0.80,
|
|
lora_source: str = "unknown",
|
|
checkpoint_name: str = "unknown",
|
|
original_metadata: Optional[Dict[str, str]] = None,
|
|
weight_unet: float = 1.0,
|
|
weight_te: float = 1.0,
|
|
weight_clip: float = 1.0,
|
|
weight_vae: float = 1.0,
|
|
metadata_mode: str = "preserve_a",
|
|
batch_size: int = 64,
|
|
detected_fp8: bool = False,
|
|
) -> Tuple[
|
|
Dict[str, torch.Tensor],
|
|
Dict[str, str],
|
|
str,
|
|
Dict[str, Any],
|
|
]:
|
|
"""
|
|
Execute the full baking pipeline with per-component scaling.
|
|
|
|
Args:
|
|
ckpt_sd: Source checkpoint state dict (can be _LazyCheckpointMapping).
|
|
lora_sd: LoRA state dict to bake in.
|
|
ckpt_header: Optional safetensors header dict for mmap-based
|
|
matching and lazy assembly. When provided, the
|
|
pipeline avoids loading tensor data for key matching
|
|
and assembles output as a lazy mapping.
|
|
baking_method: "linear", "impact_weighted", or "orthogonal".
|
|
strength: Overall baking strength multiplier.
|
|
energy_concentration: Energy threshold for impact-weighted method.
|
|
lora_source: Description of LoRA source for metadata.
|
|
checkpoint_name: Source checkpoint filename for metadata.
|
|
original_metadata: Original checkpoint .safetensors metadata dict.
|
|
weight_unet: Per-component weight for U-Net (diffusion model) keys.
|
|
weight_te: Per-component weight for Text Encoder keys.
|
|
weight_clip: Per-component weight for CLIP Vision keys.
|
|
weight_vae: Per-component weight for VAE keys.
|
|
metadata_mode: "none" (baking only), "preserve_a" (original priority),
|
|
or "merge_basic" (baking priority).
|
|
detected_fp8: When True, the checkpoint has FP8 native weights.
|
|
Computation is performed in bfloat16 (via Fix 1 in
|
|
_safe_bake_add) and stale FP8 scale keys are stripped
|
|
from the output. Preserved (non-baked) keys are
|
|
auto-converted from FP8 to bfloat16 on access.
|
|
|
|
Returns:
|
|
output_sd: Final assembled state dict with baked weights
|
|
metadata: Metadata dict for the .safetensors header
|
|
forensic_report: Human-readable forensic report string (Lora Studio style)
|
|
impact_profile: Dict with impact analysis details
|
|
"""
|
|
# Detect if ckpt_sd is a lazy mapping (has .filepath attribute)
|
|
# _LazyCheckpointMapping from musubi_checkpoint_studio.py stores the
|
|
# original file path, enabling lazy assembly without a separate path arg.
|
|
self._lazy_mode = hasattr(ckpt_sd, 'filepath')
|
|
self._ckpt_path = getattr(ckpt_sd, 'filepath', None)
|
|
# 🔥 FP8: Store detection flag for assembly pass-through
|
|
self._detected_fp8 = detected_fp8
|
|
|
|
with memory_optimized_merge():
|
|
# Step 2a: SD1.5 Diffusers Bridge detection
|
|
use_diffusers_bridge = False
|
|
if detect_native_sd15_checkpoint(ckpt_sd) and detect_diffusers_sd15_lora(lora_sd):
|
|
print(" 🔄 SD1.5 Diffusers Bridge ACTIVE")
|
|
use_diffusers_bridge = True
|
|
|
|
comfyui_yield()
|
|
|
|
print("\n" + "=" * 50)
|
|
print("🔥 Easy LoRA Baker — Baking Pipeline")
|
|
precision_info = (f"(compute uses float16 fallback)"
|
|
if self.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
|
|
else "")
|
|
print(f" 🎯 Precision: {self.dtype} {precision_info}")
|
|
print("=" * 50)
|
|
|
|
# =====================================================================
|
|
# PIPELINE PROGRESS: Track outer bake phases
|
|
# =====================================================================
|
|
# The inner baking methods (linear/impact_weighted/orthogonal) and
|
|
# key matching each have their own ProgressTracker. This outer tracker
|
|
# covers the phases BETWEEN those: normalization, delta reconstruction,
|
|
# statistics, assembly, and metadata.
|
|
PIPELINE_PHASES = 5
|
|
pipeline_progress = ProgressTracker(total=PIPELINE_PHASES, desc="Bake pipeline")
|
|
pipeline_progress.update(0)
|
|
|
|
# Step 2b: Normalize LoRA keys to master format
|
|
# (handles all formats: SDXL Kohya, ComfyUI native, Musubi, Flux Klein, Z-Image)
|
|
if use_diffusers_bridge:
|
|
# Preserve Diffusers block structure (no down_blocks→input_blocks conversion)
|
|
lora_sd, lora_key_map = normalize_diffusers_preserving(lora_sd)
|
|
print(f" 🔄 Normalized {len(lora_sd)} keys (Diffusers-preserving)")
|
|
else:
|
|
from .identity_normalizer import identity_normalize
|
|
lora_sd, lora_key_map = identity_normalize(lora_sd)
|
|
print(f" 🔄 Normalized {len(lora_sd)} keys to master format")
|
|
|
|
pipeline_progress += 1 # Phase 1: Normalization done
|
|
|
|
# Step 3: Group LoRA pairs for on-demand delta reconstruction
|
|
# Instead of reconstructing all ≈112 deltas upfront (≈9.6 GB), store only
|
|
# the A/B component tensors (≈14 MB) and reconstruct each delta on demand
|
|
# during matching (shape inference) and baking (.pop() → one at a time).
|
|
#
|
|
# GPU acceleration: if CUDA is available, the A/B tensors are moved to GPU
|
|
# during _reconstruct_one() inside _LazyDeltaDict.reconstruct(), keeping the
|
|
# fast matmul path while avoiding VRAM accumulation since only one delta
|
|
# is reconstructed at a time.
|
|
device_for_reconstruction = self.device if 'cuda' in str(self.device) else None
|
|
lora_pairs, lora_standalone = _group_lora_pairs(lora_sd)
|
|
lora_deltas = _LazyDeltaDict(lora_pairs, lora_standalone, device=device_for_reconstruction)
|
|
|
|
if not lora_deltas:
|
|
return self._abort_bake(
|
|
"No LoRA deltas found — aborting bake",
|
|
ckpt_sd, baking_method, strength,
|
|
lora_source, checkpoint_name, original_metadata,
|
|
weight_unet, weight_te, weight_clip, weight_vae,
|
|
metadata_mode,
|
|
)
|
|
|
|
# Step 5: Log all LoRA delta keys before matching (debug aid)
|
|
# Models ComfyUI's approach at comfy/sd.py:78-106 where convert_lora+load_lora
|
|
# produce the patch_dict; we log pre-matching deltas to trace key coverage.
|
|
print(f"\n 📋 LoRA delta keys ({len(lora_deltas)} total):")
|
|
# 🔥 FIX: Use _categorize_lora_key() instead of naive 'lora_unet'/'lora_te' substring
|
|
# checks. After Diffusers→ComfyUI conversion, keys use diffusion_model.* prefix
|
|
# (not lora_unet_*), so the old substring check always returned 0 UNet keys.
|
|
# _categorize_lora_key() recognizes all known normalized formats:
|
|
# unet: diffusion_model.*, input_blocks.*, output_blocks.*, middle_block.*
|
|
# te: text_model.*, te.*, lora_te_*, transformer.text_model.*
|
|
# clip: clip*
|
|
# vae: vae*, first_stage*
|
|
comp_counts: Counter = Counter()
|
|
for k in lora_deltas:
|
|
comp_counts[self._categorize_lora_key(k)] += 1
|
|
unet_count = comp_counts.get('unet', 0)
|
|
te_count = comp_counts.get('te', 0)
|
|
clip_count = comp_counts.get('clip', 0)
|
|
vae_count = comp_counts.get('vae', 0)
|
|
other_count = len(lora_deltas) - unet_count - te_count - clip_count - vae_count
|
|
print(f" UNet: {unet_count} keys, TE: {te_count} keys, "
|
|
f"CLIP: {clip_count} keys, VAE: {vae_count} keys, Other: {other_count} keys")
|
|
if self._verbose:
|
|
for i, (base, delta) in enumerate(sorted(lora_deltas.items())):
|
|
if isinstance(delta, _ShapeInfo):
|
|
norm_str = "N/A (lazy)"
|
|
else:
|
|
norm_str = f"{torch.norm(delta.float()).item():.4f}"
|
|
print(f" [{i:>3d}] {str(base):>70s} shape={str(list(delta.shape)):>20s} norm={norm_str}")
|
|
|
|
comfyui_yield()
|
|
pipeline_progress += 1 # Phase 2: Delta reconstruction done
|
|
|
|
# Step 5: Match keys
|
|
if use_diffusers_bridge:
|
|
# Build Diffusers-aware reverse map for 1:1 exact matching
|
|
bridge_reverse_map = build_diffusers_reverse_map(ckpt_sd)
|
|
matched_deltas = match_diffusers_deltas(ckpt_sd, lora_deltas, bridge_reverse_map)
|
|
match_stats = {"diffusers_bridge_1:1": len(matched_deltas)}
|
|
print(f"\n 📊 Match strategy distribution ({len(matched_deltas)} total):")
|
|
print(f" diffusers_bridge_1:1: {len(matched_deltas)} keys")
|
|
else:
|
|
matched_deltas, match_stats = self._find_matching_keys(
|
|
ckpt_sd, lora_deltas, return_stats=True,
|
|
ckpt_header=ckpt_header,
|
|
)
|
|
# Log match strategy distribution
|
|
if match_stats:
|
|
print(f"\n 📊 Match strategy distribution ({sum(match_stats.values())} total):")
|
|
for strategy, count in sorted(match_stats.items(), key=lambda x: -x[1]):
|
|
print(f" {strategy}: {count} keys")
|
|
# NOTE: matched_deltas stay on CPU — moved to GPU on-demand per-key
|
|
# inside _prepare_bake_key / each bake method's inner loop.
|
|
if not matched_deltas:
|
|
return self._abort_bake(
|
|
"No keys matched between LoRA and checkpoint — aborting bake",
|
|
ckpt_sd, baking_method, strength,
|
|
lora_source, checkpoint_name, original_metadata,
|
|
weight_unet, weight_te, weight_clip, weight_vae,
|
|
metadata_mode,
|
|
)
|
|
|
|
pipeline_progress += 1 # Phase 3: Key matching done
|
|
|
|
# Step 6: Apply baking method with per-component scaling
|
|
BAKE_DISPATCH = {
|
|
"linear": self.bake_linear,
|
|
"impact_weighted": self.bake_impact_weighted,
|
|
"orthogonal": self.bake_orthogonal,
|
|
}
|
|
|
|
bake_fn = BAKE_DISPATCH.get(baking_method)
|
|
if bake_fn is None:
|
|
print(f"⚠️ Unknown baking method '{baking_method}' — falling back to linear")
|
|
bake_fn = self.bake_linear
|
|
baking_method = "linear"
|
|
|
|
# ── RAM Guard: per-batch GPU memory guard ──
|
|
if matched_deltas and ckpt_header is not None:
|
|
sample_delta = next(iter(matched_deltas.values()))
|
|
per_key_bytes = sample_delta.numel() * sample_delta.element_size() * 3
|
|
batch_peak = per_key_bytes * batch_size
|
|
_avail_ram = get_available_ram()
|
|
if _avail_ram is not None:
|
|
_ram_gb = _avail_ram / (1024**3)
|
|
_peak_gb = batch_peak / (1024**3)
|
|
print(f"🧠 RAM Guard (baking): Estimated peak ~{_peak_gb:.2f} GB/batch, "
|
|
f"Available RAM ~{_ram_gb:.2f} GB, "
|
|
f"85% threshold ~{_ram_gb * 0.85:.2f} GB")
|
|
if batch_peak > _avail_ram * 0.85:
|
|
old_bs = batch_size
|
|
batch_size = 1
|
|
print(f" ⚠️ Batch size {old_bs} exceeds RAM threshold — "
|
|
f"falling back to sequential (batch_size=1)")
|
|
else:
|
|
print(f"✅ RAM Guard (baking): Batch size {batch_size} within safe limits")
|
|
else:
|
|
print(f"⚠️ RAM Guard (baking): Cannot detect available RAM — proceeding")
|
|
|
|
# Initialize output dict (always dict mode — streaming path removed)
|
|
output_dict: Dict[str, torch.Tensor] = {}
|
|
|
|
if baking_method == "impact_weighted":
|
|
delta_norms, delta_ratios = bake_fn(
|
|
ckpt_sd, matched_deltas, output_dict, strength, energy_concentration,
|
|
weight_unet, weight_te, weight_clip, weight_vae,
|
|
batch_size=batch_size,
|
|
)
|
|
else:
|
|
delta_norms, delta_ratios = bake_fn(
|
|
ckpt_sd, matched_deltas, output_dict, strength,
|
|
weight_unet, weight_te, weight_clip, weight_vae,
|
|
batch_size=batch_size,
|
|
)
|
|
|
|
# Capture matched count before deletion
|
|
_matched_count = len(matched_deltas) if matched_deltas is not None else 0
|
|
|
|
# matched_deltas and lora_deltas are no longer needed after baking.
|
|
# In lazy mode, _MatchedDeltas is a thin wrapper (few KB) and
|
|
# _LazyDeltaDict holds only component tensors (≈14 MB for 112 pairs).
|
|
# Delete eagerly to free residual memory (Fix #6).
|
|
del matched_deltas, lora_deltas
|
|
cleanup_memory()
|
|
comfyui_yield()
|
|
|
|
|
|
# Alias output_dict as baked for minimal downstream changes
|
|
baked = output_dict
|
|
|
|
# =====================================================================
|
|
# COMPONENT BAKED COUNTING: Count how many keys per component were baked
|
|
# =====================================================================
|
|
if hasattr(self, '_component_breakdown') and self._component_breakdown:
|
|
_comp_map = {
|
|
'model_diffusion': 'unet',
|
|
'te': 'te',
|
|
'clip_vision': 'clip',
|
|
'vae': 'vae',
|
|
}
|
|
for ckpt_key in baked:
|
|
ckpt_comp = categorize_checkpoint_key(ckpt_key)
|
|
mapped = _comp_map.get(ckpt_comp, 'unet')
|
|
if mapped in self._component_breakdown:
|
|
self._component_breakdown[mapped].setdefault("baked", 0)
|
|
self._component_breakdown[mapped]["baked"] += 1
|
|
# Ensure all components have a "baked" field
|
|
for comp, data in self._component_breakdown.items():
|
|
data.setdefault("baked", 0)
|
|
|
|
# =====================================================================
|
|
# BAKE DROPOUT REPORT: Keys matched but not baked
|
|
# =====================================================================
|
|
if hasattr(self, '_bake_dropped') and self._bake_dropped:
|
|
drop_by_reason: Dict[str, int] = {}
|
|
for key, reason in self._bake_dropped.items():
|
|
# Categorize by reason type (first word before colon)
|
|
reason_cat = reason.split(":")[0].strip() if ":" in reason else reason
|
|
drop_by_reason[reason_cat] = drop_by_reason.get(reason_cat, 0) + 1
|
|
total_dropped = len(self._bake_dropped)
|
|
print(f"\n ⚠️ BAKE DROPOUT: {total_dropped} matched keys were NOT baked:")
|
|
for reason, count in sorted(drop_by_reason.items(), key=lambda x: -x[1]):
|
|
print(f" {reason}: {count} keys")
|
|
if self._verbose:
|
|
for key, reason in sorted(self._bake_dropped.items()):
|
|
print(f" - {key}: {reason}")
|
|
|
|
# =====================================================================
|
|
# DELTA STATISTICS: Use norms/ratios collected during bake loop
|
|
# =====================================================================
|
|
print("\n" + "-" * 50)
|
|
print("📊 DELTA STATISTICS (aggregate across all baked keys):")
|
|
|
|
if delta_norms:
|
|
mean_dn = statistics.mean(delta_norms)
|
|
max_dn = max(delta_norms)
|
|
mean_ratio = statistics.mean(delta_ratios)
|
|
max_ratio = max(delta_ratios)
|
|
median_ratio = statistics.median(delta_ratios)
|
|
# Manual stddev — Python 3.12's statistics.stdev has an internal
|
|
# Fraction conversion issue that raises:
|
|
# AttributeError: 'float' object has no attribute 'numerator'
|
|
if len(delta_ratios) > 1:
|
|
mean_r = statistics.mean(delta_ratios)
|
|
variance = sum((x - mean_r) ** 2 for x in delta_ratios) / (len(delta_ratios) - 1)
|
|
std_ratio = variance ** 0.5
|
|
else:
|
|
std_ratio = 0.0
|
|
mean_ratio_percent = round(mean_ratio * 100, 2)
|
|
max_ratio_percent = round(max_ratio * 100, 2)
|
|
self._delta_analysis = {
|
|
"layers_analyzed": len(delta_norms),
|
|
"mean_ratio_percent": mean_ratio_percent,
|
|
"median_ratio_percent": round(median_ratio * 100, 2),
|
|
"max_ratio_percent": max_ratio_percent,
|
|
"std_ratio_percent": round(std_ratio * 100, 2),
|
|
"mean_delta_norm": round(mean_dn, 4),
|
|
"max_delta_norm": round(max_dn, 4),
|
|
"effect_visual": self._create_effect_bar(mean_ratio_percent, max_ratio_percent),
|
|
}
|
|
print(f" Layers analyzed: {len(delta_norms)}")
|
|
print(f" Mean delta norm: {mean_dn:.4e}")
|
|
print(f" Max delta norm: {max_dn:.4e}")
|
|
print(f" Mean ratio (%): {mean_ratio*100:.2f}%")
|
|
print(f" Median ratio (%): {median_ratio*100:.2f}%")
|
|
print(f" Max ratio (%): {max_ratio*100:.2f}%")
|
|
print(f" Std ratio (%): {std_ratio*100:.2f}%")
|
|
else:
|
|
self._delta_analysis = {
|
|
"layers_analyzed": 0, "mean_ratio_percent": 0.0,
|
|
"median_ratio_percent": 0.0, "max_ratio_percent": 0.0,
|
|
"std_ratio_percent": 0.0, "mean_delta_norm": 0.0, "max_delta_norm": 0.0,
|
|
"effect_visual": self._create_effect_bar(0.0, 0.0),
|
|
}
|
|
print("-" * 50)
|
|
pipeline_progress += 1 # Phase 4: Baking + statistics done
|
|
|
|
# =====================================================================
|
|
# 🔍 BAKE VERIFICATION: Sample-check baked weights differ from original
|
|
# =====================================================================
|
|
if baked and ckpt_sd:
|
|
_verify_keys = list(baked.keys())[:3]
|
|
_all_differ = True
|
|
for _vk in _verify_keys:
|
|
if _vk in ckpt_sd:
|
|
_orig = ckpt_sd[_vk]
|
|
_baked_t = baked[_vk]
|
|
if isinstance(_orig, torch.Tensor) and isinstance(_baked_t, torch.Tensor):
|
|
if _orig.shape == _baked_t.shape:
|
|
_diff = (_baked_t.to(dtype=torch.float32) - _orig.to(dtype=torch.float32)).abs().max().item()
|
|
_dtype_match = "✅" if _baked_t.dtype == _orig.dtype else f"⚠️ dtype: {_baked_t.dtype}≠{_orig.dtype}"
|
|
print(f" 🔍 Bake verify '{_vk[:60]}...': |diff|_max={_diff:.6e} {_dtype_match}")
|
|
if _diff == 0.0:
|
|
_all_differ = False
|
|
else:
|
|
print(f" 🔍 Bake verify '{_vk[:60]}...': shape mismatch {_baked_t.shape}≠{_orig.shape}")
|
|
_all_differ = False
|
|
else:
|
|
print(f" 🔍 Bake verify '{_vk[:60]}...': non-tensor type")
|
|
_all_differ = False
|
|
if _all_differ and _verify_keys:
|
|
print(f" ✅ Bake verification passed — sampled weights differ from originals")
|
|
elif _verify_keys:
|
|
print(f" ⚠️ Bake verification: some sampled weights match originals (delta=0)")
|
|
else:
|
|
print(f" ⚠️ Bake verification: no baked keys to verify")
|
|
|
|
# NaN/Inf sanitization
|
|
# 🔥 REVERT-TO-ORIGINAL: Pass ckpt_sd so any key with persistent NaN/Inf
|
|
# reverts to its original checkpoint weight instead of being replaced
|
|
# with zeros (which cause black/dark patches in generated images).
|
|
baked = self._sanitize_baked(baked, original_ckpt_sd=ckpt_sd)
|
|
|
|
# Close mmap handles if ckpt_sd is a lazy mapping (Fix #4).
|
|
# After sanitize, we no longer need to read from the original checkpoint.
|
|
# The lazy assembly step creates a new mapping from self._ckpt_path instead.
|
|
if hasattr(ckpt_sd, 'close'):
|
|
ckpt_sd.close()
|
|
|
|
# ── Measure pre-assembly RAM ──
|
|
_pre_assembly_ram = get_available_ram()
|
|
print(f" 🕐 [t={time.time()-_T0:.1f}s] Starting output assembly (_assemble_output_lazy)")
|
|
|
|
# Step 8: Assemble output — always lazy when checkpoint is mmap'd.
|
|
# _assemble_output_lazy creates a _LazyCheckpointMapping with baked keys
|
|
# in the write cache. Non-baked keys stay mmap'd from the original file —
|
|
# zero bytes copied. Compatible with both preview mode
|
|
# (load_state_dict_as_model_objects) and save mode
|
|
# (materialize via dict(items()) then save_safetensors_stream).
|
|
if self._lazy_mode:
|
|
if self._ckpt_path is not None and ckpt_header is not None:
|
|
_user_fp8 = (
|
|
self.dtype
|
|
if self.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
|
|
else None
|
|
)
|
|
output_sd = _assemble_output_lazy(
|
|
self._ckpt_path, ckpt_header, baked,
|
|
detected_fp8=self._detected_fp8,
|
|
user_fp8_dtype=_user_fp8,
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
"Lazy assembly requires both ckpt_path and ckpt_header. "
|
|
"Cannot fall back to eager assembly — would load 17+ GiB into RAM."
|
|
)
|
|
else:
|
|
output_sd = self._assemble_output(
|
|
ckpt_sd, baked,
|
|
fp8_dequantized=self._detected_fp8,
|
|
target_dtype=self.dtype,
|
|
)
|
|
|
|
# Free the baked dict — its CPU tensors are now held by output_sd
|
|
# (or the lazy mapping's write cache). Reclaim the dict wrapper + any
|
|
# stale references eagerly so cleanup_memory() can free residual GPU pages.
|
|
# NOTE: Keep `baked` alive until all downstream uses (impact_profile,
|
|
# metadata, forensic report, final print) are done.
|
|
# Convert per-key count to local variables first for safe deletion later.
|
|
_baked_count = len(baked)
|
|
cleanup_memory()
|
|
comfyui_yield()
|
|
|
|
# ── Measure post-assembly RAM and print delta ──
|
|
_post_assembly_ram = get_available_ram()
|
|
if _pre_assembly_ram is not None and _post_assembly_ram is not None:
|
|
_delta_gb = (_pre_assembly_ram - _post_assembly_ram) / (1024**3)
|
|
print(f"📊 RAM: Pre-assembly {_pre_assembly_ram / (1024**3):.2f} GB → "
|
|
f"Post-assembly {_post_assembly_ram / (1024**3):.2f} GB "
|
|
f"(delta: {_delta_gb:+.2f} GB)")
|
|
elif _post_assembly_ram is not None:
|
|
print(f"📊 Post-assembly available RAM: {_post_assembly_ram / (1024**3):.2f} GB")
|
|
cleanup_memory()
|
|
print(f" 🕐 [t={time.time()-_T0:.1f}s] Output assembly complete — {len(output_sd)} keys in output_sd")
|
|
|
|
# =====================================================================
|
|
# 🔍 DIAGNOSTIC: Validate baked keys match checkpoint key format
|
|
# =====================================================================
|
|
# After the Section 4d/4e fix in _build_reverse_key_map(), all Flux Klein
|
|
# LoRA keys match via reverse_map (Strategy 0) instead of shape_block
|
|
# (Strategy 4). This diagnostic verifies:
|
|
# 1. Baked keys use checkpoint-native prefix
|
|
# - Vanilla Flux: transformer.*
|
|
# - Klein Flux: double_blocks.* / single_blocks.* (bare)
|
|
# 2. The strategy distribution shows 0 shape_block keys for Flux
|
|
# 3. ComfyUI's "left over keys" in Flux model loading is NORMAL —
|
|
# ComfyUI internally strips prefixes during loading
|
|
if baked and ckpt_sd:
|
|
# Check if this is a Flux checkpoint
|
|
_ckpt_sample_keys = list(ckpt_sd.keys())[:5]
|
|
_is_flux = any(
|
|
k.startswith(('transformer.', 'double_blocks.', 'single_blocks.'))
|
|
for k in _ckpt_sample_keys
|
|
)
|
|
if _is_flux:
|
|
# Detect checkpoint key prefix style
|
|
_has_transformer = any(k.startswith('transformer.') for k in _ckpt_sample_keys)
|
|
_has_bare_blocks = any(
|
|
k.startswith(('double_blocks.', 'single_blocks.'))
|
|
for k in _ckpt_sample_keys
|
|
)
|
|
_expected_prefix = 'transformer.' if _has_transformer else 'double_blocks./single_blocks.'
|
|
|
|
# Verify baked keys use correct prefix
|
|
if _has_transformer:
|
|
_all_proper_prefix = all(k.startswith('transformer.') for k in baked.keys())
|
|
else:
|
|
_all_proper_prefix = all(
|
|
k.startswith(('double_blocks.', 'single_blocks.'))
|
|
for k in baked.keys()
|
|
)
|
|
|
|
if _all_proper_prefix:
|
|
print(f" ✅ Baked key format: all {len(baked)} keys use "
|
|
f"'{_expected_prefix}' prefix — correct for Flux model loader")
|
|
print(f" (ComfyUI's Flux loader strips prefixes internally;")
|
|
print(f" 'left over keys' in logs is NORMAL behavior — weights ARE consumed)")
|
|
else:
|
|
_wrong_prefix = [k for k in baked.keys()
|
|
if _has_transformer and not k.startswith('transformer.')
|
|
or not _has_transformer and not k.startswith(('double_blocks.', 'single_blocks.'))]
|
|
print(f" ⚠️ {len(_wrong_prefix)} baked keys with unexpected prefix:")
|
|
for _nk in _wrong_prefix[:3]:
|
|
print(f" - {_nk}")
|
|
print(f" These may NOT be consumed by ComfyUI's Flux model loader!")
|
|
|
|
# Build impact profile (use captured lengths — baked still alive here)
|
|
impact_profile = {
|
|
"baking_method": baking_method,
|
|
"matched_keys": _matched_count,
|
|
"baked_keys": _baked_count,
|
|
"preserved_keys": len(ckpt_sd) - _baked_count,
|
|
"total_keys": len(ckpt_sd),
|
|
"strength": strength,
|
|
"weight_unet": weight_unet,
|
|
"weight_te": weight_te,
|
|
"weight_clip": weight_clip,
|
|
"weight_vae": weight_vae,
|
|
}
|
|
|
|
# Step 9: Build metadata
|
|
metadata = self._build_metadata(
|
|
baking_method, strength,
|
|
lora_source, checkpoint_name,
|
|
_baked_count, len(ckpt_sd) - _baked_count,
|
|
original_metadata,
|
|
weight_unet, weight_te, weight_clip, weight_vae,
|
|
metadata_mode,
|
|
)
|
|
|
|
# Build forensic report
|
|
forensic_report = self._build_forensic_report(
|
|
baking_method, strength,
|
|
lora_source, checkpoint_name,
|
|
_baked_count, len(ckpt_sd) - _baked_count,
|
|
len(ckpt_sd), impact_profile,
|
|
weight_unet=weight_unet, weight_te=weight_te,
|
|
weight_clip=weight_clip, weight_vae=weight_vae,
|
|
component_breakdown=getattr(self, '_component_breakdown', None),
|
|
delta_analysis=getattr(self, '_delta_analysis', None),
|
|
)
|
|
|
|
print("=" * 50)
|
|
print("✅ Baking pipeline complete!")
|
|
print(f" Baked {_baked_count} keys, preserved {len(ckpt_sd) - _baked_count} keys")
|
|
print(f" 🧱 Weights — UNet:{weight_unet} TE:{weight_te} CLIP:{weight_clip} VAE:{weight_vae}")
|
|
print(f" 📋 Metadata mode: {metadata_mode}")
|
|
print("=" * 50)
|
|
pipeline_progress += 1 # Phase 5: Assembly + metadata done
|
|
pipeline_progress.complete()
|
|
|
|
# Free the baked dict now that all uses are done.
|
|
del baked
|
|
cleanup_memory()
|
|
|
|
print(f" 🕐 [t={time.time()-_T0:.1f}s] SmartBakingProcessor.bake returning")
|
|
return output_sd, metadata, forensic_report, impact_profile
|