Update precision_model_save.py
This commit is contained in:
Tiwaz
+239
-60
@@ -1,6 +1,7 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from contextlib import nullcontext
|
||||
|
||||
import torch
|
||||
|
||||
@@ -17,16 +18,46 @@ DTYPES = {
|
||||
"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 PrecisionModelSave:
|
||||
"""
|
||||
Materialises a ComfyUI MODEL's weight patches in float32, then saves the
|
||||
resulting standalone diffusion model in a chosen output dtype.
|
||||
|
||||
This intentionally does not call comfy.sd.save_checkpoint(), so it avoids
|
||||
the normal ModelSave materialisation/casting path. It is intended as a
|
||||
diagnostic saver for live model merges.
|
||||
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()
|
||||
@@ -38,9 +69,16 @@ class PrecisionModelSave:
|
||||
"model": ("MODEL",),
|
||||
"filename_prefix": (
|
||||
"STRING",
|
||||
{"default": "diffusion_models/precision_merge"},
|
||||
{"default": "diffusion_models/precision_merge_v2"},
|
||||
),
|
||||
"materialise_mode": (
|
||||
["match_live", "fp32_math"],
|
||||
{"default": "match_live"},
|
||||
),
|
||||
"save_dtype": (
|
||||
["fp32", "bf16", "fp16"],
|
||||
{"default": "bf16"},
|
||||
),
|
||||
"save_dtype": (["fp32", "bf16", "fp16"], {"default": "fp32"}),
|
||||
"strip_diffusion_model_prefix": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
@@ -59,25 +97,45 @@ class PrecisionModelSave:
|
||||
CATEGORY = "model/merging"
|
||||
|
||||
@staticmethod
|
||||
def _materialise_weight(key, patch_description, output_dtype):
|
||||
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, ...]
|
||||
[(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]
|
||||
|
||||
# Work on CPU in FP32 so interpolation/extrapolation is not rounded to
|
||||
# the source checkpoint dtype before saving.
|
||||
if not isinstance(base_weight, torch.Tensor):
|
||||
raise TypeError(
|
||||
f"{key}: unsupported base weight type {type(base_weight).__name__}"
|
||||
f"{key}: unsupported base weight type "
|
||||
f"{type(base_weight).__name__}"
|
||||
)
|
||||
|
||||
weight = base_weight.detach().to(
|
||||
device="cpu", dtype=torch.float32, copy=True
|
||||
device="cpu",
|
||||
dtype=torch.float32,
|
||||
copy=True,
|
||||
)
|
||||
|
||||
if convert_func is not None:
|
||||
@@ -88,8 +146,6 @@ class PrecisionModelSave:
|
||||
|
||||
patches = patch_description[1:]
|
||||
if patches:
|
||||
# Current ComfyUI supports intermediate_dtype here. The fallback
|
||||
# keeps the node usable on somewhat older builds.
|
||||
try:
|
||||
weight = comfy.lora.calculate_weight(
|
||||
patches,
|
||||
@@ -98,48 +154,137 @@ class PrecisionModelSave:
|
||||
intermediate_dtype=torch.float32,
|
||||
)
|
||||
except TypeError:
|
||||
weight = comfy.lora.calculate_weight(patches, weight, key)
|
||||
# Compatibility fallback for older ComfyUI builds.
|
||||
weight = comfy.lora.calculate_weight(
|
||||
patches,
|
||||
weight,
|
||||
key,
|
||||
)
|
||||
|
||||
if not isinstance(weight, torch.Tensor):
|
||||
raise TypeError(
|
||||
f"{key}: materialised weight is {type(weight).__name__}, not Tensor"
|
||||
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."
|
||||
)
|
||||
|
||||
return weight.detach().to(
|
||||
device="cpu", dtype=output_dtype, copy=False
|
||||
).contiguous()
|
||||
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,
|
||||
)
|
||||
(
|
||||
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)
|
||||
output_path = os.path.join(
|
||||
full_output_folder,
|
||||
output_name,
|
||||
)
|
||||
|
||||
logging.info(
|
||||
"[PrecisionModelSave] Materialising model patches in FP32; "
|
||||
"saving as %s to %s",
|
||||
"[PrecisionModelSaveV2] Mode=%s; saving as %s to %s",
|
||||
materialise_mode,
|
||||
save_dtype,
|
||||
output_path,
|
||||
)
|
||||
logging.info(
|
||||
"[PrecisionModelSaveV2] %s",
|
||||
MATERIALISE_MODES[materialise_mode],
|
||||
)
|
||||
|
||||
# This returns the base tensor plus every live merge/patch attached to
|
||||
# the MODEL object, rather than only model.model.state_dict().
|
||||
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:
|
||||
@@ -151,40 +296,71 @@ class PrecisionModelSave:
|
||||
state_dict = {}
|
||||
total = len(patch_map)
|
||||
|
||||
for index, (internal_key, patch_description) in enumerate(
|
||||
patch_map.items(), start=1
|
||||
):
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
# 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()
|
||||
|
||||
save_key = internal_key
|
||||
if strip_diffusion_model_prefix and save_key.startswith(
|
||||
"diffusion_model."
|
||||
):
|
||||
save_key = save_key[len("diffusion_model."):]
|
||||
|
||||
state_dict[save_key] = self._materialise_weight(
|
||||
with context:
|
||||
for index, (
|
||||
internal_key,
|
||||
patch_description,
|
||||
output_dtype,
|
||||
)
|
||||
) in enumerate(
|
||||
patch_map.items(),
|
||||
start=1,
|
||||
):
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
|
||||
if index == 1 or index % 100 == 0 or index == total:
|
||||
logging.info(
|
||||
"[PrecisionModelSave] %d/%d tensors materialised",
|
||||
index,
|
||||
total,
|
||||
)
|
||||
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": "1",
|
||||
"materialise_dtype": "fp32",
|
||||
"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)
|
||||
@@ -198,18 +374,21 @@ class PrecisionModelSave:
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
# Release references promptly; large FP32 dictionaries can consume a
|
||||
# lot of system RAM until Python's next collection.
|
||||
del state_dict
|
||||
|
||||
logging.info("[PrecisionModelSave] Saved: %s", output_path)
|
||||
logging.info(
|
||||
"[PrecisionModelSaveV2] Saved: %s",
|
||||
output_path,
|
||||
)
|
||||
return (output_path,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PrecisionModelSave": PrecisionModelSave,
|
||||
"PrecisionModelSaveV2": PrecisionModelSaveV2,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PrecisionModelSave": "Precision Model Save (FP32 Materialise)",
|
||||
"PrecisionModelSaveV2": (
|
||||
"Precision Model Save V2 (Live Match / FP32 Math)"
|
||||
),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user