395 lines
12 KiB
Python
395 lines
12 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
from contextlib import nullcontext
|
|
|
|
import torch
|
|
|
|
import comfy.lora
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
import folder_paths
|
|
from comfy.cli_args import args
|
|
|
|
|
|
DTYPES = {
|
|
"fp32": torch.float32,
|
|
"bf16": torch.bfloat16,
|
|
"fp16": torch.float16,
|
|
}
|
|
|
|
MATERIALISE_MODES = {
|
|
"match_live": (
|
|
"Use ModelPatcher.patch_weight_to_device(return_weight=True), matching "
|
|
"ComfyUI's live LoRA/patch computation, set_func handling, compute dtype, "
|
|
"and deterministic stochastic rounding."
|
|
),
|
|
"fp32_math": (
|
|
"Rebuild patches manually on CPU in FP32 before the final save cast. "
|
|
"This is the high-precision mathematical merge path from V1."
|
|
),
|
|
}
|
|
|
|
|
|
class PrecisionModelSaveV2:
|
|
"""
|
|
Saves a standalone diffusion model from a live ComfyUI MODEL/ModelPatcher.
|
|
|
|
Modes
|
|
-----
|
|
match_live:
|
|
Materialises each tensor through ComfyUI's own live patching path using
|
|
patch_weight_to_device(..., return_weight=True). This is intended to
|
|
preserve the behaviour of a live model with ordinary attached LoRAs and
|
|
model-merge patches as closely as possible.
|
|
|
|
fp32_math:
|
|
Reconstructs each patched tensor manually on CPU in float32, then casts
|
|
once to the selected output dtype. This preserves V1's high-precision
|
|
merge behaviour, but it may differ slightly from live inference because
|
|
live ComfyUI can use device-dependent LoRA compute dtypes, custom
|
|
set_func logic, and deterministic stochastic rounding.
|
|
|
|
Notes
|
|
-----
|
|
Forward hooks, injections, bypass LoRAs, and weight-wrapper patches may not
|
|
be representable as ordinary standalone checkpoint tensors. The node logs a
|
|
warning when it detects those structures.
|
|
"""
|
|
|
|
VERSION = "2"
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"filename_prefix": (
|
|
"STRING",
|
|
{"default": "diffusion_models/precision_merge_v2"},
|
|
),
|
|
"materialise_mode": (
|
|
["match_live", "fp32_math"],
|
|
{"default": "match_live"},
|
|
),
|
|
"save_dtype": (
|
|
["fp32", "bf16", "fp16"],
|
|
{"default": "bf16"},
|
|
),
|
|
"strip_diffusion_model_prefix": (
|
|
"BOOLEAN",
|
|
{"default": True},
|
|
),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("saved_path",)
|
|
FUNCTION = "save"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "model/merging"
|
|
|
|
@staticmethod
|
|
def _to_output_tensor(weight, output_dtype):
|
|
if not isinstance(weight, torch.Tensor):
|
|
raise TypeError(
|
|
f"Materialised weight is {type(weight).__name__}, not Tensor"
|
|
)
|
|
return (
|
|
weight.detach()
|
|
.to(device="cpu", dtype=output_dtype, copy=True)
|
|
.contiguous()
|
|
)
|
|
|
|
@classmethod
|
|
def _materialise_fp32_math(
|
|
cls,
|
|
key,
|
|
patch_description,
|
|
output_dtype,
|
|
):
|
|
"""
|
|
V1-compatible manual FP32 materialisation.
|
|
|
|
patch_description follows ModelPatcher.get_key_patches():
|
|
[(base_weight, convert_func), patch_1, patch_2, ...]
|
|
"""
|
|
if not patch_description:
|
|
raise RuntimeError(f"No patch description for {key}")
|
|
|
|
base_weight, convert_func = patch_description[0]
|
|
|
|
if not isinstance(base_weight, torch.Tensor):
|
|
raise TypeError(
|
|
f"{key}: unsupported base weight type "
|
|
f"{type(base_weight).__name__}"
|
|
)
|
|
|
|
weight = base_weight.detach().to(
|
|
device="cpu",
|
|
dtype=torch.float32,
|
|
copy=True,
|
|
)
|
|
|
|
if convert_func is not None:
|
|
try:
|
|
weight = convert_func(weight, inplace=True)
|
|
except TypeError:
|
|
weight = convert_func(weight)
|
|
|
|
patches = patch_description[1:]
|
|
if patches:
|
|
try:
|
|
weight = comfy.lora.calculate_weight(
|
|
patches,
|
|
weight,
|
|
key,
|
|
intermediate_dtype=torch.float32,
|
|
)
|
|
except TypeError:
|
|
# Compatibility fallback for older ComfyUI builds.
|
|
weight = comfy.lora.calculate_weight(
|
|
patches,
|
|
weight,
|
|
key,
|
|
)
|
|
|
|
return cls._to_output_tensor(weight, output_dtype)
|
|
|
|
@classmethod
|
|
def _materialise_match_live(
|
|
cls,
|
|
model,
|
|
key,
|
|
output_dtype,
|
|
):
|
|
"""
|
|
Materialise through ComfyUI's actual live ModelPatcher path.
|
|
|
|
return_weight=True prevents the call from replacing the parameter in
|
|
the live model while retaining ComfyUI's normal convert_func, set_func,
|
|
LoRA compute dtype, and deterministic stochastic-rounding behaviour.
|
|
"""
|
|
patch_method = getattr(model, "patch_weight_to_device", None)
|
|
if patch_method is None:
|
|
raise RuntimeError(
|
|
"This ComfyUI build does not expose "
|
|
"ModelPatcher.patch_weight_to_device(). "
|
|
"Use fp32_math mode or update ComfyUI."
|
|
)
|
|
|
|
load_device = getattr(model, "load_device", None)
|
|
if load_device is None:
|
|
raise RuntimeError(
|
|
f"{key}: MODEL has no load_device; cannot match live patching."
|
|
)
|
|
|
|
try:
|
|
weight = patch_method(
|
|
key,
|
|
device_to=load_device,
|
|
return_weight=True,
|
|
)
|
|
except TypeError as exc:
|
|
raise RuntimeError(
|
|
"Your ComfyUI build has an incompatible "
|
|
"patch_weight_to_device() signature. "
|
|
"Use fp32_math mode or update ComfyUI."
|
|
) from exc
|
|
|
|
return cls._to_output_tensor(weight, output_dtype)
|
|
|
|
@staticmethod
|
|
def _warn_about_non_bakeable_features(model):
|
|
warnings = []
|
|
|
|
hook_patches = getattr(model, "hook_patches", None)
|
|
if hook_patches:
|
|
warnings.append(
|
|
f"{len(hook_patches)} hook-patch group(s)"
|
|
)
|
|
|
|
wrapper_patches = getattr(model, "weight_wrapper_patches", None)
|
|
if wrapper_patches:
|
|
warnings.append(
|
|
f"{len(wrapper_patches)} weight-wrapper patch group(s)"
|
|
)
|
|
|
|
injections = getattr(model, "injections", None)
|
|
if injections:
|
|
warnings.append(
|
|
f"{len(injections)} injection group(s)"
|
|
)
|
|
|
|
if warnings:
|
|
logging.warning(
|
|
"[PrecisionModelSaveV2] Detected %s. These can modify the "
|
|
"forward pass and may not be fully bakeable into ordinary "
|
|
"checkpoint tensors.",
|
|
", ".join(warnings),
|
|
)
|
|
|
|
def save(
|
|
self,
|
|
model,
|
|
filename_prefix,
|
|
materialise_mode,
|
|
save_dtype,
|
|
strip_diffusion_model_prefix,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
if materialise_mode not in MATERIALISE_MODES:
|
|
raise ValueError(
|
|
f"Unknown materialise mode: {materialise_mode}"
|
|
)
|
|
|
|
output_dtype = DTYPES[save_dtype]
|
|
|
|
(
|
|
full_output_folder,
|
|
filename,
|
|
counter,
|
|
subfolder,
|
|
resolved_prefix,
|
|
) = folder_paths.get_save_image_path(
|
|
filename_prefix,
|
|
self.output_dir,
|
|
)
|
|
os.makedirs(full_output_folder, exist_ok=True)
|
|
|
|
output_name = f"{filename}_{counter:05}_.safetensors"
|
|
output_path = os.path.join(
|
|
full_output_folder,
|
|
output_name,
|
|
)
|
|
|
|
logging.info(
|
|
"[PrecisionModelSaveV2] Mode=%s; saving as %s to %s",
|
|
materialise_mode,
|
|
save_dtype,
|
|
output_path,
|
|
)
|
|
logging.info(
|
|
"[PrecisionModelSaveV2] %s",
|
|
MATERIALISE_MODES[materialise_mode],
|
|
)
|
|
|
|
self._warn_about_non_bakeable_features(model)
|
|
|
|
# Includes base diffusion tensors and ordinary live merge/LoRA patches.
|
|
patch_map = model.get_key_patches("diffusion_model.")
|
|
|
|
if not patch_map:
|
|
raise RuntimeError(
|
|
"No diffusion_model.* weights found. "
|
|
"This node expects a ComfyUI MODEL/ModelPatcher."
|
|
)
|
|
|
|
state_dict = {}
|
|
total = len(patch_map)
|
|
|
|
# Eject model injections while reading ordinary patch tensors, matching
|
|
# ComfyUI's own saving/state-dict behaviour where available.
|
|
use_ejected = getattr(model, "use_ejected", None)
|
|
context = use_ejected() if callable(use_ejected) else nullcontext()
|
|
|
|
with context:
|
|
for index, (
|
|
internal_key,
|
|
patch_description,
|
|
) in enumerate(
|
|
patch_map.items(),
|
|
start=1,
|
|
):
|
|
comfy.model_management.throw_exception_if_processing_interrupted()
|
|
|
|
save_key = internal_key
|
|
if (
|
|
strip_diffusion_model_prefix
|
|
and save_key.startswith("diffusion_model.")
|
|
):
|
|
save_key = save_key[
|
|
len("diffusion_model.") :
|
|
]
|
|
|
|
if materialise_mode == "match_live":
|
|
tensor = self._materialise_match_live(
|
|
model,
|
|
internal_key,
|
|
output_dtype,
|
|
)
|
|
else:
|
|
tensor = self._materialise_fp32_math(
|
|
internal_key,
|
|
patch_description,
|
|
output_dtype,
|
|
)
|
|
|
|
state_dict[save_key] = tensor
|
|
|
|
if (
|
|
index == 1
|
|
or index % 100 == 0
|
|
or index == total
|
|
):
|
|
logging.info(
|
|
"[PrecisionModelSaveV2] %d/%d tensors materialised",
|
|
index,
|
|
total,
|
|
)
|
|
|
|
metadata = {
|
|
"format": "pt",
|
|
"comfy_precision_model_save": self.VERSION,
|
|
"materialise_mode": materialise_mode,
|
|
"saved_dtype": save_dtype,
|
|
"stripped_diffusion_model_prefix": str(
|
|
bool(strip_diffusion_model_prefix)
|
|
).lower(),
|
|
}
|
|
|
|
if materialise_mode == "fp32_math":
|
|
metadata["materialise_dtype"] = "fp32"
|
|
else:
|
|
metadata["materialise_dtype"] = "comfy_live_path"
|
|
|
|
if not args.disable_metadata:
|
|
if prompt is not None:
|
|
metadata["prompt"] = json.dumps(prompt)
|
|
if extra_pnginfo is not None:
|
|
for key, value in extra_pnginfo.items():
|
|
metadata[key] = json.dumps(value)
|
|
|
|
comfy.utils.save_torch_file(
|
|
state_dict,
|
|
output_path,
|
|
metadata=metadata,
|
|
)
|
|
|
|
del state_dict
|
|
|
|
logging.info(
|
|
"[PrecisionModelSaveV2] Saved: %s",
|
|
output_path,
|
|
)
|
|
return (output_path,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PrecisionModelSaveV2": PrecisionModelSaveV2,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PrecisionModelSaveV2": (
|
|
"Precision Model Save V2 (Live Match / FP32 Math)"
|
|
),
|
|
}
|