967 lines
34 KiB
Python
967 lines
34 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, Iterable, Optional, Set, Tuple
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
|
|
import comfy.lora
|
|
import comfy.lora_convert
|
|
import comfy.model_management
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import folder_paths
|
|
|
|
|
|
NODE_CATEGORY = "model_merging/lora"
|
|
SAVE_DTYPES = {
|
|
"fp16": torch.float16,
|
|
"bf16": torch.bfloat16,
|
|
"fp32": torch.float32,
|
|
}
|
|
LAYER_SCOPE_OPTIONS = ["original_lora_only", "all_compatible_weights"]
|
|
RANK_MODE_OPTIONS = ["match_input_lora_when_possible", "fixed"]
|
|
OUTPUT_MODE_OPTIONS = ["standard", "full_diff_patch"]
|
|
COMPUTE_DEVICE_OPTIONS = ["auto", "cpu", "cuda"]
|
|
SVD_MODE_OPTIONS = ["auto", "full", "lowrank"]
|
|
CUDA_SVD_DRIVER_OPTIONS = ["auto", "gesvdj", "gesvda", "gesvd"]
|
|
|
|
|
|
@dataclass
|
|
class GroupStats:
|
|
processed: int = 0
|
|
written_lora: int = 0
|
|
written_diff: int = 0
|
|
written_bias: int = 0
|
|
skipped_zero: int = 0
|
|
skipped_missing: int = 0
|
|
skipped_shape: int = 0
|
|
skipped_nonfloat: int = 0
|
|
skipped_unsupported: int = 0
|
|
cuda_fallbacks: int = 0
|
|
errors: int = 0
|
|
|
|
def written_total(self) -> int:
|
|
return self.written_lora + self.written_diff + self.written_bias
|
|
|
|
|
|
@dataclass
|
|
class RebaseStats:
|
|
model: GroupStats
|
|
clip: GroupStats
|
|
loaded_patch_keys: int
|
|
selected_model_keys: int
|
|
selected_clip_keys: int
|
|
saved_relpath: str
|
|
saved_fullpath: str
|
|
compute_device: str
|
|
svd_mode: str
|
|
cuda_svd_driver: str
|
|
|
|
|
|
|
|
def _sanitize_save_relpath(save_relpath: str) -> str:
|
|
rel = (save_relpath or "").strip().replace("\\", "/")
|
|
if not rel:
|
|
raise ValueError("save_relpath cannot be empty.")
|
|
|
|
if rel.lower().endswith(".safetensors"):
|
|
rel = rel[: -len(".safetensors")]
|
|
|
|
rel = os.path.normpath(rel)
|
|
if rel in ("", "."):
|
|
raise ValueError("save_relpath cannot be empty.")
|
|
if os.path.isabs(rel):
|
|
raise ValueError("save_relpath must be relative to the ComfyUI loras folder, not absolute.")
|
|
if rel.startswith("..") or "/../" in rel.replace("\\", "/"):
|
|
raise ValueError("save_relpath cannot escape the ComfyUI loras folder.")
|
|
|
|
return rel.replace("\\", "/")
|
|
|
|
|
|
|
|
def _read_safetensors_metadata(filepath: str) -> Dict[str, str]:
|
|
try:
|
|
if not str(filepath).lower().endswith(".safetensors"):
|
|
return {}
|
|
with safe_open(filepath, framework="pt") as handle:
|
|
meta = handle.metadata()
|
|
return dict(meta) if meta else {}
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
|
|
def _first_lora_root() -> str:
|
|
lora_roots = folder_paths.get_folder_paths("loras")
|
|
if not lora_roots:
|
|
raise RuntimeError("No ComfyUI loras folder is registered.")
|
|
root = os.path.abspath(lora_roots[0])
|
|
os.makedirs(root, exist_ok=True)
|
|
return root
|
|
|
|
|
|
|
|
def _next_available_lora_path(save_relpath: str) -> Tuple[str, str]:
|
|
root = _first_lora_root()
|
|
rel = _sanitize_save_relpath(save_relpath)
|
|
|
|
rel_with_ext = rel if rel.lower().endswith(".safetensors") else f"{rel}.safetensors"
|
|
full_path = os.path.abspath(os.path.join(root, rel_with_ext))
|
|
parent = os.path.dirname(full_path)
|
|
|
|
if os.path.commonpath((root, parent)) != root:
|
|
raise ValueError("save_relpath cannot escape the ComfyUI loras folder.")
|
|
|
|
os.makedirs(parent, exist_ok=True)
|
|
|
|
if not os.path.exists(full_path):
|
|
return full_path, os.path.relpath(full_path, root).replace("\\", "/")
|
|
|
|
stem, ext = os.path.splitext(os.path.basename(full_path))
|
|
counter = 1
|
|
while True:
|
|
candidate = os.path.join(parent, f"{stem}_{counter:05d}{ext}")
|
|
if not os.path.exists(candidate):
|
|
return candidate, os.path.relpath(candidate, root).replace("\\", "/")
|
|
counter += 1
|
|
|
|
|
|
|
|
def _cast_for_save(tensor: torch.Tensor, save_dtype: str) -> torch.Tensor:
|
|
out_dtype = SAVE_DTYPES[save_dtype]
|
|
return tensor.detach().to(device="cpu", dtype=out_dtype).contiguous()
|
|
|
|
|
|
|
|
def _cast_scalar_for_save(value: float) -> torch.Tensor:
|
|
return torch.tensor(float(value), dtype=torch.float32)
|
|
|
|
|
|
|
|
def _get_model_candidate_keys(model_patcher: Any) -> Set[str]:
|
|
sd = model_patcher.model_state_dict()
|
|
return {
|
|
key
|
|
for key in sd.keys()
|
|
if key.startswith("diffusion_model.") and (key.endswith(".weight") or key.endswith(".bias"))
|
|
}
|
|
|
|
|
|
|
|
def _get_clip_candidate_keys(clip: Any) -> Set[str]:
|
|
if clip is None:
|
|
return set()
|
|
sd = clip.patcher.model_state_dict()
|
|
return {key for key in sd.keys() if key.endswith(".weight") or key.endswith(".bias")}
|
|
|
|
|
|
|
|
def _state_dict_shapes(state_dict: Dict[str, torch.Tensor]) -> Dict[str, Tuple[int, ...]]:
|
|
return {key: tuple(value.shape) for key, value in state_dict.items()}
|
|
|
|
|
|
|
|
def _maybe_interrupt() -> None:
|
|
if hasattr(comfy.model_management, "throw_exception_if_processing_interrupted"):
|
|
comfy.model_management.throw_exception_if_processing_interrupted()
|
|
|
|
|
|
|
|
def _resolve_compute_device(device_name: str) -> torch.device:
|
|
requested = (device_name or "auto").lower()
|
|
|
|
if requested == "cpu":
|
|
return torch.device("cpu")
|
|
|
|
if requested == "cuda":
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("compute_device='cuda' was requested, but CUDA is not available.")
|
|
return torch.device("cuda")
|
|
|
|
if hasattr(comfy.model_management, "get_torch_device"):
|
|
try:
|
|
comfy_device = comfy.model_management.get_torch_device()
|
|
if comfy_device is not None:
|
|
if isinstance(comfy_device, str):
|
|
comfy_device = torch.device(comfy_device)
|
|
if isinstance(comfy_device, torch.device):
|
|
if comfy_device.type != "cuda" or torch.cuda.is_available():
|
|
return comfy_device
|
|
except Exception:
|
|
pass
|
|
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda")
|
|
return torch.device("cpu")
|
|
|
|
|
|
|
|
def _effective_weight_from_patcher(model_patcher: Any, key: str, device: torch.device) -> torch.Tensor:
|
|
tensor = model_patcher.patch_weight_to_device(key, device_to=device, return_weight=True)
|
|
if tensor is None:
|
|
raise KeyError(f"Key {key} was not found in the model patcher.")
|
|
return tensor.detach().to(device=device, dtype=torch.float32).contiguous()
|
|
|
|
|
|
|
|
def _extract_actual_key(mapping_value: Any) -> Optional[str]:
|
|
if isinstance(mapping_value, str):
|
|
return mapping_value
|
|
if isinstance(mapping_value, (tuple, list)) and mapping_value:
|
|
head = mapping_value[0]
|
|
if isinstance(head, str):
|
|
return head
|
|
if isinstance(mapping_value, dict):
|
|
for candidate_key in ("target", "weight", "actual_key", "key"):
|
|
candidate = mapping_value.get(candidate_key)
|
|
actual = _extract_actual_key(candidate)
|
|
if actual:
|
|
return actual
|
|
return None
|
|
|
|
|
|
|
|
def _normalize_patch_target_keys(patch_keys: Iterable[Any]) -> Set[str]:
|
|
normalized: Set[str] = set()
|
|
for patch_key in patch_keys:
|
|
actual = _extract_actual_key(patch_key)
|
|
if actual:
|
|
normalized.add(actual)
|
|
return normalized
|
|
|
|
|
|
|
|
def _canonical_prefix_lookup(key_map: Dict[str, Any]) -> Dict[str, str]:
|
|
lookup: Dict[str, str] = {}
|
|
for source_prefix, mapping_value in key_map.items():
|
|
actual_key = _extract_actual_key(mapping_value)
|
|
if not actual_key:
|
|
continue
|
|
canonical = source_prefix
|
|
if isinstance(mapping_value, dict):
|
|
canonical = mapping_value.get("canonical_prefix", source_prefix)
|
|
lookup.setdefault(actual_key, canonical)
|
|
return lookup
|
|
|
|
|
|
|
|
def _strip_patch_suffix(prefix: str) -> str:
|
|
if prefix.endswith(".weight"):
|
|
return prefix[:-len(".weight")]
|
|
if prefix.endswith(".bias"):
|
|
return prefix[:-len(".bias")]
|
|
return prefix
|
|
|
|
|
|
|
|
def _logical_patch_prefix(actual_key: str, canonical_lookup: Dict[str, str]) -> str:
|
|
if actual_key.endswith(".bias"):
|
|
canonical = canonical_lookup.get(actual_key[:-len(".bias")] + ".weight")
|
|
if canonical:
|
|
return _strip_patch_suffix(canonical)
|
|
|
|
canonical = canonical_lookup.get(actual_key)
|
|
if canonical:
|
|
return _strip_patch_suffix(canonical)
|
|
|
|
if actual_key.endswith(".weight"):
|
|
return actual_key[:-len(".weight")]
|
|
if actual_key.endswith(".bias"):
|
|
return actual_key[:-len(".bias")]
|
|
return actual_key
|
|
|
|
|
|
|
|
def _infer_input_ranks(converted_lora: Dict[str, torch.Tensor], key_map: Dict[str, Any]) -> Dict[str, int]:
|
|
ranks: Dict[str, int] = {}
|
|
for source_prefix, mapping_value in key_map.items():
|
|
actual_key = _extract_actual_key(mapping_value)
|
|
if not actual_key:
|
|
continue
|
|
down_key = f"{source_prefix}.lora_down.weight"
|
|
if down_key not in converted_lora:
|
|
continue
|
|
try:
|
|
rank = int(converted_lora[down_key].shape[0])
|
|
except Exception:
|
|
continue
|
|
if actual_key not in ranks:
|
|
ranks[actual_key] = rank
|
|
return ranks
|
|
|
|
|
|
|
|
def _choose_svd_strategy(rank: int, matrix_shape: Tuple[int, int], svd_mode: str, device_type: str) -> str:
|
|
min_dim = min(matrix_shape)
|
|
if svd_mode == "full":
|
|
return "full"
|
|
if svd_mode == "lowrank":
|
|
return "lowrank"
|
|
|
|
# auto: keep CPU deterministic/exact by default.
|
|
# Use randomized low-rank mainly on CUDA, where small-rank extraction tends to benefit more.
|
|
if device_type == "cuda" and min_dim >= 512 and rank <= max(64, min_dim // 4):
|
|
return "lowrank"
|
|
return "full"
|
|
|
|
|
|
|
|
def _extract_lora_factors_impl(
|
|
diff: torch.Tensor,
|
|
rank: int,
|
|
clamp_quantile: float,
|
|
svd_mode: str,
|
|
cuda_svd_driver: str,
|
|
lowrank_oversample: int,
|
|
lowrank_niter: int,
|
|
) -> Optional[Tuple[torch.Tensor, torch.Tensor, int]]:
|
|
if diff.ndim == 2:
|
|
diff_2d = diff
|
|
out_dim, in_dim = diff.shape
|
|
conv2d = False
|
|
kernel_size = None
|
|
elif diff.ndim == 4:
|
|
out_dim, in_dim = diff.shape[:2]
|
|
kernel_size = diff.shape[2:4]
|
|
conv2d = True
|
|
if kernel_size == (1, 1):
|
|
diff_2d = diff.reshape(out_dim, in_dim)
|
|
else:
|
|
diff_2d = diff.flatten(start_dim=1)
|
|
else:
|
|
return None
|
|
|
|
if diff_2d.ndim != 2:
|
|
return None
|
|
|
|
rank = max(1, min(int(rank), diff_2d.shape[0], diff_2d.shape[1]))
|
|
if rank <= 0:
|
|
return None
|
|
|
|
matrix = diff_2d.float().contiguous()
|
|
strategy = _choose_svd_strategy(rank, (matrix.shape[0], matrix.shape[1]), svd_mode, matrix.device.type)
|
|
|
|
if strategy == "lowrank":
|
|
q = min(min(matrix.shape), rank + max(int(lowrank_oversample), max(8, rank // 5)))
|
|
U, S, V = torch.svd_lowrank(matrix, q=q, niter=max(0, int(lowrank_niter)))
|
|
U = U[:, :rank]
|
|
S = S[:rank]
|
|
V = V[:, :rank]
|
|
sqrt_s = S.sqrt()
|
|
up = U * sqrt_s.unsqueeze(0)
|
|
down = (V * sqrt_s.unsqueeze(0)).T
|
|
else:
|
|
svd_kwargs = {}
|
|
if matrix.is_cuda and cuda_svd_driver != "auto":
|
|
svd_kwargs["driver"] = cuda_svd_driver
|
|
U, S, Vh = torch.linalg.svd(matrix, full_matrices=False, **svd_kwargs)
|
|
U = U[:, :rank]
|
|
S = S[:rank]
|
|
Vh = Vh[:rank, :]
|
|
sqrt_s = S.sqrt()
|
|
up = U * sqrt_s.unsqueeze(0)
|
|
down = sqrt_s.unsqueeze(1) * Vh
|
|
|
|
if 0.0 < clamp_quantile < 1.0:
|
|
dist = torch.cat([up.flatten(), down.flatten()])
|
|
hi_val = torch.quantile(dist.abs(), clamp_quantile)
|
|
up = up.clamp(-hi_val, hi_val)
|
|
down = down.clamp(-hi_val, hi_val)
|
|
|
|
if conv2d:
|
|
if kernel_size == (1, 1):
|
|
up = up.reshape(out_dim, rank, 1, 1)
|
|
down = down.reshape(rank, in_dim, 1, 1)
|
|
else:
|
|
up = up.reshape(out_dim, rank, 1, 1)
|
|
down = down.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
|
|
|
|
return up, down, rank
|
|
|
|
|
|
|
|
def _extract_lora_factors(
|
|
diff: torch.Tensor,
|
|
rank: int,
|
|
clamp_quantile: float,
|
|
svd_mode: str,
|
|
cuda_svd_driver: str,
|
|
lowrank_oversample: int,
|
|
lowrank_niter: int,
|
|
) -> Tuple[Optional[Tuple[torch.Tensor, torch.Tensor, int]], bool]:
|
|
try:
|
|
return (
|
|
_extract_lora_factors_impl(
|
|
diff=diff,
|
|
rank=rank,
|
|
clamp_quantile=clamp_quantile,
|
|
svd_mode=svd_mode,
|
|
cuda_svd_driver=cuda_svd_driver,
|
|
lowrank_oversample=lowrank_oversample,
|
|
lowrank_niter=lowrank_niter,
|
|
),
|
|
False,
|
|
)
|
|
except (torch.cuda.OutOfMemoryError, RuntimeError) as exc:
|
|
if diff.is_cuda:
|
|
logging.warning("CUDA SVD failed; falling back to CPU for this layer: %s", exc)
|
|
try:
|
|
torch.cuda.empty_cache()
|
|
except Exception:
|
|
pass
|
|
return (
|
|
_extract_lora_factors_impl(
|
|
diff=diff.to(device="cpu"),
|
|
rank=rank,
|
|
clamp_quantile=clamp_quantile,
|
|
svd_mode=svd_mode,
|
|
cuda_svd_driver="auto",
|
|
lowrank_oversample=lowrank_oversample,
|
|
lowrank_niter=lowrank_niter,
|
|
),
|
|
True,
|
|
)
|
|
raise
|
|
|
|
|
|
|
|
def _store_delta(
|
|
output_sd: Dict[str, torch.Tensor],
|
|
actual_key: str,
|
|
canonical_lookup: Dict[str, str],
|
|
diff: torch.Tensor,
|
|
rank: int,
|
|
save_dtype: str,
|
|
output_mode: str,
|
|
bias_diff: bool,
|
|
clamp_quantile: float,
|
|
svd_mode: str,
|
|
cuda_svd_driver: str,
|
|
lowrank_oversample: int,
|
|
lowrank_niter: int,
|
|
) -> Tuple[str, bool]:
|
|
patch_prefix = _logical_patch_prefix(actual_key, canonical_lookup)
|
|
|
|
if actual_key.endswith(".bias"):
|
|
if not bias_diff:
|
|
return "skip_unsupported", False
|
|
output_sd[f"{patch_prefix}.diff_b"] = _cast_for_save(diff, save_dtype)
|
|
return "bias", False
|
|
|
|
if not actual_key.endswith(".weight"):
|
|
return "skip_unsupported", False
|
|
|
|
if output_mode == "full_diff_patch":
|
|
output_sd[f"{patch_prefix}.diff"] = _cast_for_save(diff, save_dtype)
|
|
return "diff", False
|
|
|
|
factors, fell_back_to_cpu = _extract_lora_factors(
|
|
diff=diff,
|
|
rank=rank,
|
|
clamp_quantile=clamp_quantile,
|
|
svd_mode=svd_mode,
|
|
cuda_svd_driver=cuda_svd_driver,
|
|
lowrank_oversample=lowrank_oversample,
|
|
lowrank_niter=lowrank_niter,
|
|
)
|
|
if factors is None:
|
|
output_sd[f"{patch_prefix}.diff"] = _cast_for_save(diff, save_dtype)
|
|
return "diff", fell_back_to_cpu
|
|
|
|
up, down, written_rank = factors
|
|
output_sd[f"{patch_prefix}.lora_up.weight"] = _cast_for_save(up, save_dtype)
|
|
output_sd[f"{patch_prefix}.lora_down.weight"] = _cast_for_save(down, save_dtype)
|
|
output_sd[f"{patch_prefix}.alpha"] = _cast_scalar_for_save(float(written_rank))
|
|
return "lora", fell_back_to_cpu
|
|
|
|
|
|
|
|
def _process_group(
|
|
stats: GroupStats,
|
|
output_sd: Dict[str, torch.Tensor],
|
|
target_patcher: Any,
|
|
finetuned_patcher: Any,
|
|
candidate_keys: Iterable[str],
|
|
canonical_lookup: Dict[str, str],
|
|
fixed_rank: int,
|
|
rank_mode: str,
|
|
input_ranks: Dict[str, int],
|
|
output_mode: str,
|
|
bias_diff: bool,
|
|
clamp_quantile: float,
|
|
save_dtype: str,
|
|
compute_device: torch.device,
|
|
svd_mode: str,
|
|
cuda_svd_driver: str,
|
|
lowrank_oversample: int,
|
|
lowrank_niter: int,
|
|
) -> None:
|
|
target_shapes = _state_dict_shapes(target_patcher.model_state_dict())
|
|
finetuned_shapes = _state_dict_shapes(finetuned_patcher.model_state_dict())
|
|
|
|
for key in sorted(candidate_keys):
|
|
_maybe_interrupt()
|
|
stats.processed += 1
|
|
|
|
if key not in target_shapes or key not in finetuned_shapes:
|
|
stats.skipped_missing += 1
|
|
continue
|
|
|
|
if target_shapes[key] != finetuned_shapes[key]:
|
|
stats.skipped_shape += 1
|
|
continue
|
|
|
|
try:
|
|
target_weight = _effective_weight_from_patcher(target_patcher, key, compute_device)
|
|
finetuned_weight = _effective_weight_from_patcher(finetuned_patcher, key, compute_device)
|
|
except Exception:
|
|
logging.exception("Failed to fetch effective weights for key: %s", key)
|
|
stats.errors += 1
|
|
continue
|
|
|
|
if not (target_weight.is_floating_point() and finetuned_weight.is_floating_point()):
|
|
stats.skipped_nonfloat += 1
|
|
del target_weight, finetuned_weight
|
|
continue
|
|
|
|
if target_weight.shape != finetuned_weight.shape:
|
|
stats.skipped_shape += 1
|
|
del target_weight, finetuned_weight
|
|
continue
|
|
|
|
diff = target_weight - finetuned_weight
|
|
if diff.numel() == 0 or float(diff.abs().max().item()) == 0.0:
|
|
stats.skipped_zero += 1
|
|
del target_weight, finetuned_weight, diff
|
|
continue
|
|
|
|
rank = int(input_ranks.get(key, fixed_rank)) if rank_mode == "match_input_lora_when_possible" else int(fixed_rank)
|
|
rank = max(1, rank)
|
|
|
|
try:
|
|
result, fell_back_to_cpu = _store_delta(
|
|
output_sd=output_sd,
|
|
actual_key=key,
|
|
canonical_lookup=canonical_lookup,
|
|
diff=diff,
|
|
rank=rank,
|
|
save_dtype=save_dtype,
|
|
output_mode=output_mode,
|
|
bias_diff=bias_diff,
|
|
clamp_quantile=clamp_quantile,
|
|
svd_mode=svd_mode,
|
|
cuda_svd_driver=cuda_svd_driver,
|
|
lowrank_oversample=lowrank_oversample,
|
|
lowrank_niter=lowrank_niter,
|
|
)
|
|
except Exception:
|
|
logging.exception("Failed to store rebased delta for key: %s", key)
|
|
stats.errors += 1
|
|
del target_weight, finetuned_weight, diff
|
|
continue
|
|
|
|
if fell_back_to_cpu:
|
|
stats.cuda_fallbacks += 1
|
|
|
|
if result == "lora":
|
|
stats.written_lora += 1
|
|
elif result == "diff":
|
|
stats.written_diff += 1
|
|
elif result == "bias":
|
|
stats.written_bias += 1
|
|
else:
|
|
stats.skipped_unsupported += 1
|
|
|
|
del target_weight, finetuned_weight, diff
|
|
|
|
|
|
|
|
def _build_report(
|
|
stats: RebaseStats,
|
|
base_lora_name: str,
|
|
strength_model: float,
|
|
strength_clip: float,
|
|
layer_scope: str,
|
|
rank_mode: str,
|
|
rank: int,
|
|
output_mode: str,
|
|
save_dtype: str,
|
|
) -> str:
|
|
lines = [
|
|
f"saved={stats.saved_relpath}",
|
|
f"source_lora={base_lora_name}",
|
|
f"strength_model={strength_model}",
|
|
f"strength_clip={strength_clip}",
|
|
f"layer_scope={layer_scope}",
|
|
f"rank_mode={rank_mode}",
|
|
f"rank={rank}",
|
|
f"output_mode={output_mode}",
|
|
f"save_dtype={save_dtype}",
|
|
f"compute_device={stats.compute_device}",
|
|
f"svd_mode={stats.svd_mode}",
|
|
f"cuda_svd_driver={stats.cuda_svd_driver}",
|
|
f"loaded_patch_keys={stats.loaded_patch_keys}",
|
|
f"selected_model_keys={stats.selected_model_keys}",
|
|
f"selected_clip_keys={stats.selected_clip_keys}",
|
|
(
|
|
"model="
|
|
f"processed:{stats.model.processed},"
|
|
f"written_lora:{stats.model.written_lora},"
|
|
f"written_diff:{stats.model.written_diff},"
|
|
f"written_bias:{stats.model.written_bias},"
|
|
f"skipped_zero:{stats.model.skipped_zero},"
|
|
f"skipped_missing:{stats.model.skipped_missing},"
|
|
f"skipped_shape:{stats.model.skipped_shape},"
|
|
f"skipped_nonfloat:{stats.model.skipped_nonfloat},"
|
|
f"skipped_unsupported:{stats.model.skipped_unsupported},"
|
|
f"cuda_fallbacks:{stats.model.cuda_fallbacks},"
|
|
f"errors:{stats.model.errors}"
|
|
),
|
|
(
|
|
"clip="
|
|
f"processed:{stats.clip.processed},"
|
|
f"written_lora:{stats.clip.written_lora},"
|
|
f"written_diff:{stats.clip.written_diff},"
|
|
f"written_bias:{stats.clip.written_bias},"
|
|
f"skipped_zero:{stats.clip.skipped_zero},"
|
|
f"skipped_missing:{stats.clip.skipped_missing},"
|
|
f"skipped_shape:{stats.clip.skipped_shape},"
|
|
f"skipped_nonfloat:{stats.clip.skipped_nonfloat},"
|
|
f"skipped_unsupported:{stats.clip.skipped_unsupported},"
|
|
f"cuda_fallbacks:{stats.clip.cuda_fallbacks},"
|
|
f"errors:{stats.clip.errors}"
|
|
),
|
|
]
|
|
return "\n".join(lines)
|
|
|
|
|
|
|
|
def _build_metadata(
|
|
saved_relpath: str,
|
|
base_lora_name: str,
|
|
strength_model: float,
|
|
strength_clip: float,
|
|
layer_scope: str,
|
|
rank_mode: str,
|
|
rank: int,
|
|
output_mode: str,
|
|
save_dtype: str,
|
|
bias_diff: bool,
|
|
clamp_quantile: float,
|
|
compute_device: str,
|
|
svd_mode: str,
|
|
cuda_svd_driver: str,
|
|
source_metadata: Dict[str, str],
|
|
report: str,
|
|
) -> Dict[str, str]:
|
|
metadata = {
|
|
"rebase.saved_relpath": str(saved_relpath),
|
|
"rebase.source_lora": str(base_lora_name),
|
|
"rebase.strength_model": str(strength_model),
|
|
"rebase.strength_clip": str(strength_clip),
|
|
"rebase.layer_scope": str(layer_scope),
|
|
"rebase.rank_mode": str(rank_mode),
|
|
"rebase.rank": str(rank),
|
|
"rebase.output_mode": str(output_mode),
|
|
"rebase.save_dtype": str(save_dtype),
|
|
"rebase.bias_diff": json.dumps(bool(bias_diff)),
|
|
"rebase.clamp_quantile": str(clamp_quantile),
|
|
"rebase.compute_device": str(compute_device),
|
|
"rebase.svd_mode": str(svd_mode),
|
|
"rebase.cuda_svd_driver": str(cuda_svd_driver),
|
|
"rebase.report": report,
|
|
}
|
|
|
|
for passthrough_key in ("prompt", "description", "tool"):
|
|
value = source_metadata.get(passthrough_key)
|
|
if value:
|
|
metadata[passthrough_key] = value
|
|
|
|
return metadata
|
|
|
|
|
|
class RebaseLoRAFromFinetune:
|
|
"""
|
|
Build a new LoRA/patch file relative to a finetuned model by first applying a
|
|
source LoRA onto the base model, then extracting the delta:
|
|
|
|
(base + source_lora) - finetuned
|
|
|
|
The node writes the output into the first registered ComfyUI loras folder.
|
|
"""
|
|
|
|
SEARCH_ALIASES = [
|
|
"rebase lora",
|
|
"convert lora to finetune",
|
|
"lora transplant",
|
|
"lora remap",
|
|
]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base_model": ("MODEL",),
|
|
"finetuned_model": ("MODEL",),
|
|
"base_lora_name": (folder_paths.get_filename_list("loras"), {"tooltip": "The source LoRA trained on top of the base model."}),
|
|
"save_relpath": ("STRING", {"default": "rebased/rebased_lora", "tooltip": "Relative path inside the ComfyUI loras folder. .safetensors is optional."}),
|
|
"rank": ("INT", {"default": 16, "min": 1, "max": 4096, "step": 1}),
|
|
"layer_scope": (LAYER_SCOPE_OPTIONS, {"tooltip": "original_lora_only keeps only layers touched by the source LoRA. all_compatible_weights tries to recreate the full target delta."}),
|
|
"rank_mode": (RANK_MODE_OPTIONS, {"tooltip": "Reuse the input LoRA layer rank when possible, or force one fixed rank for every extracted layer."}),
|
|
"output_mode": (OUTPUT_MODE_OPTIONS, {"tooltip": "standard writes LoRA factors where possible. full_diff_patch writes exact diff tensors instead."}),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
"strength_clip": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
"save_dtype": (list(SAVE_DTYPES.keys()), {"tooltip": "Tensor dtype used in the written safetensors file."}),
|
|
"bias_diff": ("BOOLEAN", {"default": True, "tooltip": "When enabled, bias deltas are written as .diff_b patches."}),
|
|
"clamp_quantile": ("FLOAT", {"default": 0.99, "min": 0.0, "max": 1.0, "step": 0.001, "tooltip": "Symmetric clamp used before writing LoRA factors. Set to 0 or 1 to disable."}),
|
|
"include_clip": ("BOOLEAN", {"default": True, "tooltip": "Also extract text-encoder deltas when clip inputs are connected."}),
|
|
"compute_device": (COMPUTE_DEVICE_OPTIONS, {"default": "auto", "tooltip": "Where to run subtraction and SVD. GPU mainly helps standard LoRA extraction; full_diff_patch is usually fine on CPU."}),
|
|
"svd_mode": (SVD_MODE_OPTIONS, {"default": "auto", "tooltip": "auto uses full SVD for most layers and randomized low-rank SVD only when the requested rank is much smaller than the matrix."}),
|
|
"cuda_svd_driver": (CUDA_SVD_DRIVER_OPTIONS, {"default": "auto", "tooltip": "CUDA-only torch.linalg.svd driver. Leave on auto unless you are tuning speed/accuracy."}),
|
|
"lowrank_oversample": ("INT", {"default": 16, "min": 0, "max": 512, "step": 1, "tooltip": "Extra rank budget used by randomized low-rank SVD in auto/lowrank mode."}),
|
|
"lowrank_niter": ("INT", {"default": 2, "min": 0, "max": 16, "step": 1, "tooltip": "Power iterations for randomized low-rank SVD. Higher is slower but usually more accurate."}),
|
|
},
|
|
"optional": {
|
|
"base_clip": ("CLIP",),
|
|
"finetuned_clip": ("CLIP",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("saved_lora", "report")
|
|
FUNCTION = "rebase_and_save"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = NODE_CATEGORY
|
|
DESCRIPTION = (
|
|
"Applies a source LoRA to the base MODEL/CLIP, subtracts the finetuned MODEL/CLIP, "
|
|
"then writes a new LoRA/patch file relative to the finetuned weights."
|
|
)
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, base_lora_name, save_relpath, **kwargs):
|
|
if folder_paths.get_full_path("loras", base_lora_name) is None:
|
|
return f"Invalid LoRA: {base_lora_name}"
|
|
try:
|
|
_sanitize_save_relpath(save_relpath)
|
|
except Exception as exc:
|
|
return str(exc)
|
|
return True
|
|
|
|
def rebase_and_save(
|
|
self,
|
|
base_model,
|
|
finetuned_model,
|
|
base_lora_name,
|
|
save_relpath,
|
|
rank,
|
|
layer_scope,
|
|
rank_mode,
|
|
output_mode,
|
|
strength_model,
|
|
strength_clip,
|
|
save_dtype,
|
|
bias_diff,
|
|
clamp_quantile,
|
|
include_clip,
|
|
compute_device,
|
|
svd_mode,
|
|
cuda_svd_driver,
|
|
lowrank_oversample,
|
|
lowrank_niter,
|
|
base_clip=None,
|
|
finetuned_clip=None,
|
|
):
|
|
if save_dtype not in SAVE_DTYPES:
|
|
raise ValueError(f"Unsupported save_dtype: {save_dtype}")
|
|
if layer_scope not in LAYER_SCOPE_OPTIONS:
|
|
raise ValueError(f"Unsupported layer_scope: {layer_scope}")
|
|
if rank_mode not in RANK_MODE_OPTIONS:
|
|
raise ValueError(f"Unsupported rank_mode: {rank_mode}")
|
|
if output_mode not in OUTPUT_MODE_OPTIONS:
|
|
raise ValueError(f"Unsupported output_mode: {output_mode}")
|
|
if compute_device not in COMPUTE_DEVICE_OPTIONS:
|
|
raise ValueError(f"Unsupported compute_device: {compute_device}")
|
|
if svd_mode not in SVD_MODE_OPTIONS:
|
|
raise ValueError(f"Unsupported svd_mode: {svd_mode}")
|
|
if cuda_svd_driver not in CUDA_SVD_DRIVER_OPTIONS:
|
|
raise ValueError(f"Unsupported cuda_svd_driver: {cuda_svd_driver}")
|
|
|
|
resolved_device = _resolve_compute_device(compute_device)
|
|
|
|
save_fullpath, saved_relpath = _next_available_lora_path(save_relpath)
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", base_lora_name)
|
|
source_metadata = _read_safetensors_metadata(lora_path)
|
|
lora_sd = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
converted_lora = comfy.lora_convert.convert_lora(lora_sd)
|
|
|
|
key_map: Dict[str, Any] = {}
|
|
key_map = comfy.lora.model_lora_keys_unet(base_model.model, key_map)
|
|
if include_clip and base_clip is not None:
|
|
key_map = comfy.lora.model_lora_keys_clip(base_clip.cond_stage_model, key_map)
|
|
|
|
loaded_patches = comfy.lora.load_lora(converted_lora, key_map)
|
|
if not loaded_patches:
|
|
raise RuntimeError("The selected LoRA did not load any keys onto the provided base model / clip inputs.")
|
|
|
|
target_model, target_clip = comfy.sd.load_lora_for_models(
|
|
base_model,
|
|
base_clip if include_clip else None,
|
|
lora_sd,
|
|
strength_model,
|
|
strength_clip,
|
|
)
|
|
|
|
if target_model is None:
|
|
raise RuntimeError("Applying the source LoRA to the base MODEL produced no model output.")
|
|
|
|
model_candidate_keys = _get_model_candidate_keys(target_model)
|
|
clip_candidate_keys = _get_clip_candidate_keys(target_clip) if include_clip and target_clip is not None and finetuned_clip is not None else set()
|
|
|
|
if layer_scope == "original_lora_only":
|
|
loaded_patch_target_keys = _normalize_patch_target_keys(loaded_patches.keys())
|
|
selected_model_keys = {key for key in loaded_patch_target_keys if key in model_candidate_keys}
|
|
selected_clip_keys = {key for key in loaded_patch_target_keys if key in clip_candidate_keys}
|
|
else:
|
|
finetuned_model_keys = _get_model_candidate_keys(finetuned_model)
|
|
selected_model_keys = model_candidate_keys & finetuned_model_keys
|
|
|
|
if include_clip and target_clip is not None and finetuned_clip is not None:
|
|
finetuned_clip_keys = _get_clip_candidate_keys(finetuned_clip)
|
|
selected_clip_keys = clip_candidate_keys & finetuned_clip_keys
|
|
else:
|
|
selected_clip_keys = set()
|
|
|
|
if not selected_model_keys and not selected_clip_keys:
|
|
raise RuntimeError(
|
|
"No compatible keys were selected for extraction. "
|
|
"Check that the base LoRA matches the base MODEL/CLIP and that the finetuned MODEL/CLIP share compatible keys."
|
|
)
|
|
|
|
input_ranks = _infer_input_ranks(converted_lora, key_map)
|
|
canonical_lookup = _canonical_prefix_lookup(key_map)
|
|
output_sd: Dict[str, torch.Tensor] = {}
|
|
|
|
model_stats = GroupStats()
|
|
clip_stats = GroupStats()
|
|
|
|
context = torch.inference_mode if hasattr(torch, "inference_mode") else torch.no_grad
|
|
with context():
|
|
_process_group(
|
|
stats=model_stats,
|
|
output_sd=output_sd,
|
|
target_patcher=target_model,
|
|
finetuned_patcher=finetuned_model,
|
|
candidate_keys=selected_model_keys,
|
|
canonical_lookup=canonical_lookup,
|
|
fixed_rank=int(rank),
|
|
rank_mode=rank_mode,
|
|
input_ranks=input_ranks,
|
|
output_mode=output_mode,
|
|
bias_diff=bias_diff,
|
|
clamp_quantile=float(clamp_quantile),
|
|
save_dtype=save_dtype,
|
|
compute_device=resolved_device,
|
|
svd_mode=svd_mode,
|
|
cuda_svd_driver=cuda_svd_driver,
|
|
lowrank_oversample=int(lowrank_oversample),
|
|
lowrank_niter=int(lowrank_niter),
|
|
)
|
|
|
|
if include_clip and target_clip is not None and finetuned_clip is not None and selected_clip_keys:
|
|
_process_group(
|
|
stats=clip_stats,
|
|
output_sd=output_sd,
|
|
target_patcher=target_clip.patcher,
|
|
finetuned_patcher=finetuned_clip.patcher,
|
|
candidate_keys=selected_clip_keys,
|
|
canonical_lookup=canonical_lookup,
|
|
fixed_rank=int(rank),
|
|
rank_mode=rank_mode,
|
|
input_ranks=input_ranks,
|
|
output_mode=output_mode,
|
|
bias_diff=bias_diff,
|
|
clamp_quantile=float(clamp_quantile),
|
|
save_dtype=save_dtype,
|
|
compute_device=resolved_device,
|
|
svd_mode=svd_mode,
|
|
cuda_svd_driver=cuda_svd_driver,
|
|
lowrank_oversample=int(lowrank_oversample),
|
|
lowrank_niter=int(lowrank_niter),
|
|
)
|
|
|
|
if not output_sd:
|
|
raise RuntimeError("No deltas were written. The computed difference was empty or all compatible keys were skipped.")
|
|
|
|
stats = RebaseStats(
|
|
model=model_stats,
|
|
clip=clip_stats,
|
|
loaded_patch_keys=len(loaded_patches),
|
|
selected_model_keys=len(selected_model_keys),
|
|
selected_clip_keys=len(selected_clip_keys),
|
|
saved_relpath=saved_relpath,
|
|
saved_fullpath=save_fullpath,
|
|
compute_device=str(resolved_device),
|
|
svd_mode=svd_mode,
|
|
cuda_svd_driver=cuda_svd_driver,
|
|
)
|
|
report = _build_report(
|
|
stats=stats,
|
|
base_lora_name=base_lora_name,
|
|
strength_model=float(strength_model),
|
|
strength_clip=float(strength_clip),
|
|
layer_scope=layer_scope,
|
|
rank_mode=rank_mode,
|
|
rank=int(rank),
|
|
output_mode=output_mode,
|
|
save_dtype=save_dtype,
|
|
)
|
|
metadata = _build_metadata(
|
|
saved_relpath=saved_relpath,
|
|
base_lora_name=base_lora_name,
|
|
strength_model=float(strength_model),
|
|
strength_clip=float(strength_clip),
|
|
layer_scope=layer_scope,
|
|
rank_mode=rank_mode,
|
|
rank=int(rank),
|
|
output_mode=output_mode,
|
|
save_dtype=save_dtype,
|
|
bias_diff=bool(bias_diff),
|
|
clamp_quantile=float(clamp_quantile),
|
|
compute_device=str(resolved_device),
|
|
svd_mode=svd_mode,
|
|
cuda_svd_driver=cuda_svd_driver,
|
|
source_metadata=source_metadata,
|
|
report=report,
|
|
)
|
|
|
|
comfy.utils.save_torch_file(output_sd, save_fullpath, metadata=metadata)
|
|
|
|
if hasattr(folder_paths, "filename_list_cache") and isinstance(folder_paths.filename_list_cache, dict):
|
|
folder_paths.filename_list_cache.pop("loras", None)
|
|
|
|
if resolved_device.type == "cuda":
|
|
try:
|
|
torch.cuda.empty_cache()
|
|
except Exception:
|
|
pass
|
|
|
|
if hasattr(comfy.model_management, "soft_empty_cache"):
|
|
try:
|
|
comfy.model_management.soft_empty_cache()
|
|
except Exception:
|
|
pass
|
|
|
|
return (saved_relpath, report)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RebaseLoRAFromFinetune": RebaseLoRAFromFinetune,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RebaseLoRAFromFinetune": "Rebase LoRA From Finetune",
|
|
}
|