493 lines
17 KiB
Python
493 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
ComfyUI-Krea2-Model-Merge
|
|
=========================
|
|
|
|
Krea 2 の標準 MODEL を直接受け取る簡易マージ／保存ノード。
|
|
|
|
接続:
|
|
[Load Diffusion Model] -- MODEL --\
|
|
[Krea2 Model Merge Simple] -- MODEL
|
|
[Load Diffusion Model] -- MODEL --/
|
|
|
|
[Krea2 Model Merge Simple] -- MODEL --> [Krea2 Model Save]
|
|
|
|
専用ローダーや独自型は使用しない。
|
|
量子化モデルの state_dict に現れる仮想キー
|
|
(weight_scale / comfy_quant など)をマージ対象から除外し、
|
|
標準 ModelPatcher の遅延パッチ機構を維持したまま MODEL を返す。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import os
|
|
import shutil
|
|
import struct
|
|
import time
|
|
import uuid
|
|
from typing import Any, Dict, Iterable, List, Mapping, Sequence, Tuple
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
|
|
import folder_paths
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
from comfy.model_patcher import get_key_weight
|
|
|
|
|
|
CATEGORY = "Krea2 Model Merge"
|
|
|
|
_QUANT_AUX_SUFFIXES = (
|
|
".weight_scale",
|
|
".weight_scale_2",
|
|
".input_scale",
|
|
".pre_quant_scale",
|
|
".comfy_quant",
|
|
".scaled_fp8",
|
|
)
|
|
|
|
_SAVE_DTYPES = ["bf16", "fp16", "fp32", "fp8_scaled"]
|
|
_DTYPE_MAP = {
|
|
"bf16": torch.bfloat16,
|
|
"fp16": torch.float16,
|
|
"fp32": torch.float32,
|
|
}
|
|
_ST_DTYPE = {
|
|
torch.float64: "F64",
|
|
torch.float32: "F32",
|
|
torch.float16: "F16",
|
|
torch.bfloat16: "BF16",
|
|
torch.int64: "I64",
|
|
torch.int32: "I32",
|
|
torch.int16: "I16",
|
|
torch.int8: "I8",
|
|
torch.uint8: "U8",
|
|
torch.bool: "BOOL",
|
|
torch.float8_e4m3fn: "F8_E4M3",
|
|
torch.float8_e5m2: "F8_E5M2",
|
|
}
|
|
_ITEMSIZE = {
|
|
"F64": 8,
|
|
"F32": 4,
|
|
"F16": 2,
|
|
"BF16": 2,
|
|
"I64": 8,
|
|
"I32": 4,
|
|
"I16": 2,
|
|
"I8": 1,
|
|
"U8": 1,
|
|
"BOOL": 1,
|
|
"F8_E4M3": 1,
|
|
"F8_E5M2": 1,
|
|
}
|
|
_FP8_MAX = float(torch.finfo(torch.float8_e4m3fn).max)
|
|
_FP8_CONF = {"format": "float8_e4m3fn"}
|
|
_FP8_CONF_BYTES = json.dumps(_FP8_CONF, separators=(",", ":")).encode("utf-8")
|
|
|
|
|
|
def _is_quant_aux_key(key: str) -> bool:
|
|
return key.endswith(_QUANT_AUX_SUFFIXES)
|
|
|
|
|
|
def _is_krea2_state_dict_key(key: str) -> bool:
|
|
"""Krea2 の主要構造。厳密な拒否判定には使わず、警告用に使う。"""
|
|
k = key[len("diffusion_model."):] if key.startswith("diffusion_model.") else key
|
|
return (
|
|
k.startswith("first.")
|
|
or k.startswith("tmlp.")
|
|
or k.startswith("txtmlp.")
|
|
or k.startswith("tproj.")
|
|
or k.startswith("txtfusion.")
|
|
or k.startswith("blocks.")
|
|
or k.startswith("last.")
|
|
)
|
|
|
|
|
|
def _actual_key_patches(model, filter_prefix: str = "diffusion_model.") -> Tuple[Dict[str, list], List[str]]:
|
|
"""
|
|
ModelPatcher.get_key_patches() と同じ形式を返すが、量子化 state_dict の
|
|
仮想補助キーを除外し、実在しない属性への getattr を避ける。
|
|
"""
|
|
model_sd = model.model_state_dict(filter_prefix)
|
|
patches: Dict[str, list] = {}
|
|
skipped: List[str] = []
|
|
|
|
backup = getattr(model, "backup", {})
|
|
hook_backup = getattr(model, "hook_backup", {})
|
|
model_patches = getattr(model, "patches", {})
|
|
|
|
for key in model_sd.keys():
|
|
if _is_quant_aux_key(key):
|
|
skipped.append(key)
|
|
continue
|
|
|
|
try:
|
|
weight, _set_func, convert_func = get_key_weight(model.model, key)
|
|
except (AttributeError, KeyError, TypeError):
|
|
# state_dict にだけ存在する仮想キーや、ModelPatcher から直接取得できない項目。
|
|
skipped.append(key)
|
|
continue
|
|
|
|
if weight is None:
|
|
skipped.append(key)
|
|
continue
|
|
|
|
bk = backup.get(key)
|
|
hbk = hook_backup.get(key)
|
|
if bk is not None:
|
|
weight = bk.weight
|
|
if hbk is not None:
|
|
weight = hbk[0]
|
|
if convert_func is None:
|
|
convert_func = lambda value, **kwargs: value
|
|
|
|
base = [(weight, convert_func)]
|
|
if key in model_patches:
|
|
base += model_patches[key]
|
|
patches[key] = base
|
|
|
|
return patches, skipped
|
|
|
|
|
|
def _shape_of_actual_key(model, key: str) -> Tuple[int, ...] | None:
|
|
try:
|
|
weight, _set_func, _convert_func = get_key_weight(model.model, key)
|
|
except (AttributeError, KeyError, TypeError):
|
|
return None
|
|
return tuple(int(v) for v in weight.shape) if weight is not None else None
|
|
|
|
|
|
def _validate_merge_shapes(model1, key_patches: Mapping[str, list]) -> None:
|
|
mismatches = []
|
|
model1_sd = model1.model_state_dict("diffusion_model.")
|
|
for key, value in key_patches.items():
|
|
if key not in model1_sd:
|
|
continue
|
|
shape1 = _shape_of_actual_key(model1, key)
|
|
try:
|
|
shape2 = tuple(int(v) for v in value[0][0].shape)
|
|
except Exception:
|
|
shape2 = None
|
|
if shape1 is not None and shape2 is not None and shape1 != shape2:
|
|
mismatches.append((key, shape1, shape2))
|
|
if len(mismatches) >= 8:
|
|
break
|
|
|
|
if mismatches:
|
|
detail = "\n".join(f" {k}: {a} != {b}" for k, a, b in mismatches)
|
|
raise ValueError(
|
|
"モデル構造またはテンソル形状が一致しないためマージできません。\n" + detail
|
|
)
|
|
|
|
|
|
def _merge_models(model1, model2, ratio: float):
|
|
ratio = float(ratio)
|
|
merged = model1.clone()
|
|
key_patches, skipped = _actual_key_patches(model2)
|
|
if not key_patches:
|
|
raise ValueError(
|
|
"model2 からマージ可能な diffusion_model の重みを取得できませんでした。"
|
|
)
|
|
|
|
_validate_merge_shapes(model1, key_patches)
|
|
|
|
accepted = 0
|
|
for key, patch in key_patches.items():
|
|
added = merged.add_patches({key: patch}, 1.0 - ratio, ratio)
|
|
if added:
|
|
accepted += 1
|
|
|
|
if accepted == 0:
|
|
raise ValueError("2つのモデルに共通するマージ可能な重みがありません。")
|
|
|
|
krea_keys = sum(1 for key in key_patches if _is_krea2_state_dict_key(key))
|
|
print(
|
|
f"[Krea2ModelMerge] merged={accepted}, ratio={ratio:.4f}, "
|
|
f"krea2_keys={krea_keys}, skipped_quant_aux={len(skipped)}"
|
|
)
|
|
return merged
|
|
|
|
|
|
def _named_internal_tensors(model) -> List[Tuple[str, torch.Tensor]]:
|
|
"""実在する parameter / buffer のみを列挙し、量子化仮想キーを含めない。"""
|
|
tensors: Dict[str, torch.Tensor] = {}
|
|
for key, value in model.model.named_parameters(recurse=True):
|
|
if key.startswith("diffusion_model."):
|
|
tensors[key] = value
|
|
for key, value in model.model.named_buffers(recurse=True):
|
|
if key.startswith("diffusion_model.") and key not in tensors:
|
|
tensors[key] = value
|
|
return sorted(tensors.items(), key=lambda item: item[0])
|
|
|
|
|
|
def _external_key(internal_key: str) -> str:
|
|
prefix = "diffusion_model."
|
|
return internal_key[len(prefix):] if internal_key.startswith(prefix) else internal_key
|
|
|
|
|
|
def _dequantize_tensor(tensor: torch.Tensor) -> torch.Tensor:
|
|
# ComfyUI QuantizedTensor
|
|
if hasattr(tensor, "dequantize"):
|
|
try:
|
|
tensor = tensor.dequantize()
|
|
except Exception:
|
|
pass
|
|
if tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
|
tensor = tensor.to(torch.bfloat16)
|
|
return tensor
|
|
|
|
|
|
def _materialize_weight(model, key: str) -> torch.Tensor:
|
|
"""パッチ適用後の1テンソルだけを実体化する。"""
|
|
device = getattr(model, "load_device", None) or comfy.model_management.get_torch_device()
|
|
with torch.inference_mode():
|
|
value = model.patch_weight_to_device(
|
|
key,
|
|
device_to=device,
|
|
return_weight=True,
|
|
force_cast=True,
|
|
)
|
|
value = _dequantize_tensor(value)
|
|
return value.detach().to("cpu").contiguous()
|
|
|
|
|
|
def _product(shape: Sequence[int]) -> int:
|
|
out = 1
|
|
for value in shape:
|
|
out *= int(value)
|
|
return out
|
|
|
|
|
|
def _write_tensor_bytes(file_obj, tensor: torch.Tensor, chunk_size: int = 64 * 1024 * 1024) -> None:
|
|
tensor = tensor.detach().to("cpu").contiguous()
|
|
raw = tensor.reshape(-1).view(torch.uint8).numpy()
|
|
view = memoryview(raw)
|
|
for start in range(0, view.nbytes, chunk_size):
|
|
file_obj.write(view[start:start + chunk_size])
|
|
|
|
|
|
def _is_fp8_eligible(key: str, tensor: torch.Tensor) -> bool:
|
|
return key.endswith(".weight") and tensor.ndim == 2 and tensor.is_floating_point()
|
|
|
|
|
|
def _quantize_fp8_scaled(tensor: torch.Tensor, key: str) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
tensor = tensor.to(torch.float32)
|
|
max_abs = torch.amax(torch.abs(tensor)).to(torch.float32)
|
|
max_value = float(max_abs.item())
|
|
if not math.isfinite(max_value):
|
|
raise ValueError(f"NaNまたはInfを含むためFP8量子化できません: {key}")
|
|
scale_value = max_value / _FP8_MAX if max_value > 0.0 else 1.0
|
|
scale = torch.tensor(scale_value, dtype=torch.float32)
|
|
quantized = torch.clamp(tensor / scale, -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn)
|
|
return quantized.contiguous(), scale
|
|
|
|
|
|
def _safe_output_path(filename_prefix: str) -> str:
|
|
paths = folder_paths.get_folder_paths("diffusion_models")
|
|
if not paths:
|
|
raise ValueError("models/diffusion_models フォルダが見つかりません。")
|
|
output_dir = paths[0]
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
name = os.path.basename((filename_prefix or "").strip()) or "krea2_merged"
|
|
if not name.lower().endswith(".safetensors"):
|
|
name += ".safetensors"
|
|
path = os.path.join(output_dir, name)
|
|
if not os.path.exists(path):
|
|
return path
|
|
|
|
stem = name[:-len(".safetensors")]
|
|
index = 1
|
|
while True:
|
|
candidate = os.path.join(output_dir, f"{stem}_{index:02d}.safetensors")
|
|
if not os.path.exists(candidate):
|
|
return candidate
|
|
index += 1
|
|
|
|
|
|
def _build_save_entries(model, save_dtype: str) -> Tuple[List[dict], List[Tuple[str, torch.Tensor]]]:
|
|
named = _named_internal_tensors(model)
|
|
if not named:
|
|
raise ValueError("保存可能な diffusion_model の実テンソルが見つかりません。")
|
|
|
|
entries: List[dict] = []
|
|
for internal_key, tensor in named:
|
|
out_key = _external_key(internal_key)
|
|
shape = [int(v) for v in tensor.shape]
|
|
if save_dtype == "fp8_scaled" and _is_fp8_eligible(out_key, tensor):
|
|
entries.append({"name": out_key, "dtype": "F8_E4M3", "shape": shape, "kind": "fp8", "source": internal_key})
|
|
prefix = out_key[:-len("weight")]
|
|
entries.append({"name": prefix + "weight_scale", "dtype": "F32", "shape": [], "kind": "fp8_scale", "source": internal_key})
|
|
entries.append({"name": prefix + "comfy_quant", "dtype": "U8", "shape": [len(_FP8_CONF_BYTES)], "kind": "fp8_conf", "source": internal_key})
|
|
else:
|
|
if tensor.is_floating_point() or tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
|
dtype = _ST_DTYPE[torch.bfloat16 if save_dtype == "fp8_scaled" else _DTYPE_MAP[save_dtype]]
|
|
else:
|
|
dtype = _ST_DTYPE.get(tensor.dtype)
|
|
if dtype is None:
|
|
raise ValueError(f"保存未対応dtype: {tensor.dtype} ({out_key})")
|
|
entries.append({"name": out_key, "dtype": dtype, "shape": shape, "kind": "tensor", "source": internal_key})
|
|
return entries, named
|
|
|
|
|
|
def _header_bytes(entries: Sequence[dict], metadata: Mapping[str, str]) -> Tuple[bytes, int]:
|
|
header: Dict[str, Any] = {"__metadata__": dict(metadata)}
|
|
offset = 0
|
|
for entry in entries:
|
|
size = _product(entry["shape"]) * _ITEMSIZE[entry["dtype"]]
|
|
header[entry["name"]] = {
|
|
"dtype": entry["dtype"],
|
|
"shape": entry["shape"],
|
|
"data_offsets": [offset, offset + size],
|
|
}
|
|
offset += size
|
|
encoded = json.dumps(header, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
|
encoded += b" " * ((-len(encoded)) % 8)
|
|
return encoded, offset
|
|
|
|
|
|
def _save_model(model, filename_prefix: str, save_dtype: str) -> str:
|
|
if save_dtype not in _SAVE_DTYPES:
|
|
raise ValueError(f"save_dtype が不正です: {save_dtype}")
|
|
|
|
entries, named = _build_save_entries(model, save_dtype)
|
|
entry_groups: Dict[str, List[dict]] = {}
|
|
for entry in entries:
|
|
entry_groups.setdefault(entry["source"], []).append(entry)
|
|
|
|
metadata = {
|
|
"merge_tool": "ComfyUI-Krea2-Model-Merge",
|
|
"save_dtype": save_dtype,
|
|
"created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
|
}
|
|
if save_dtype == "fp8_scaled":
|
|
layers = {
|
|
entry["name"][:-len(".weight")]: {"format": "float8_e4m3fn"}
|
|
for entry in entries
|
|
if entry["kind"] == "fp8"
|
|
}
|
|
metadata["_quantization_metadata"] = json.dumps(
|
|
{"format_version": "1.0", "layers": layers},
|
|
separators=(",", ":"),
|
|
)
|
|
|
|
header, data_size = _header_bytes(entries, metadata)
|
|
output_path = _safe_output_path(filename_prefix)
|
|
free_space = shutil.disk_usage(os.path.dirname(output_path)).free
|
|
expected_size = 8 + len(header) + data_size
|
|
if free_space < expected_size + 256 * 1024 * 1024:
|
|
raise OSError(f"保存先の空き容量が不足しています。必要量目安: {expected_size / (1024**3):.2f} GB")
|
|
|
|
temp_path = output_path + f".tmp-{uuid.uuid4().hex}"
|
|
progress = comfy.utils.ProgressBar(len(named))
|
|
|
|
try:
|
|
with open(temp_path, "wb", buffering=0) as file_obj:
|
|
file_obj.write(struct.pack("<Q", len(header)))
|
|
file_obj.write(header)
|
|
|
|
for internal_key, _placeholder in named:
|
|
value = _materialize_weight(model, internal_key)
|
|
group = entry_groups[internal_key]
|
|
if group[0]["kind"] == "fp8":
|
|
out_key = group[0]["name"]
|
|
quantized, scale = _quantize_fp8_scaled(value, out_key)
|
|
_write_tensor_bytes(file_obj, quantized)
|
|
_write_tensor_bytes(file_obj, scale)
|
|
_write_tensor_bytes(file_obj, torch.tensor(list(_FP8_CONF_BYTES), dtype=torch.uint8))
|
|
del quantized, scale
|
|
else:
|
|
target_dtype = {
|
|
"F64": torch.float64,
|
|
"F32": torch.float32,
|
|
"F16": torch.float16,
|
|
"BF16": torch.bfloat16,
|
|
"I64": torch.int64,
|
|
"I32": torch.int32,
|
|
"I16": torch.int16,
|
|
"I8": torch.int8,
|
|
"U8": torch.uint8,
|
|
"BOOL": torch.bool,
|
|
}[group[0]["dtype"]]
|
|
value = value.to(target_dtype)
|
|
_write_tensor_bytes(file_obj, value)
|
|
del value
|
|
progress.update(1)
|
|
|
|
file_obj.flush()
|
|
os.fsync(file_obj.fileno())
|
|
|
|
if os.path.getsize(temp_path) != expected_size:
|
|
raise IOError("保存後のファイルサイズが想定値と一致しません。")
|
|
|
|
with safe_open(temp_path, framework="pt", device="cpu") as check:
|
|
if len(check.keys()) != len(entries):
|
|
raise IOError("保存後のテンソル数が一致しません。")
|
|
|
|
os.replace(temp_path, output_path)
|
|
except Exception:
|
|
try:
|
|
if os.path.exists(temp_path):
|
|
os.remove(temp_path)
|
|
except Exception:
|
|
pass
|
|
raise
|
|
|
|
print(f"[Krea2ModelSave] saved={output_path} dtype={save_dtype}")
|
|
return output_path
|
|
|
|
|
|
class Krea2ModelMergeSimple:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model1": ("MODEL",),
|
|
"model2": ("MODEL",),
|
|
"ratio": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "merge"
|
|
CATEGORY = CATEGORY
|
|
|
|
def merge(self, model1, model2, ratio):
|
|
return (_merge_models(model1, model2, ratio),)
|
|
|
|
|
|
class Krea2ModelSave:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"filename_prefix": ("STRING", {"default": "krea2_merged"}),
|
|
"save_dtype": (_SAVE_DTYPES, {"default": "bf16"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("saved_path",)
|
|
FUNCTION = "save"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = CATEGORY
|
|
|
|
def save(self, model, filename_prefix, save_dtype):
|
|
return (_save_model(model, filename_prefix, save_dtype),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Krea2ModelMergeSimple": Krea2ModelMergeSimple,
|
|
"Krea2ModelSave": Krea2ModelSave,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Krea2ModelMergeSimple": "Krea2 Model Merge Simple",
|
|
"Krea2ModelSave": "Krea2 Model Save",
|
|
}
|