Add files via upload
This commit is contained in:
l0stl0rd
@@ -0,0 +1,48 @@
|
||||
# ComfyUI Precision Model Save
|
||||
|
||||
A small diagnostic custom node for saving a **live merged MODEL** without using
|
||||
ComfyUI's normal `ModelSave` materialisation path.
|
||||
|
||||
## What it does
|
||||
|
||||
1. Reads the live patch descriptions from `MODEL.get_key_patches()`.
|
||||
2. Materialises every merged tensor on CPU in **FP32**.
|
||||
3. Casts only the final tensor to your selected file dtype:
|
||||
- `fp32`
|
||||
- `bf16`
|
||||
- `fp16`
|
||||
4. Saves a standalone `.safetensors` diffusion model.
|
||||
|
||||
The first test should be `fp32`. It will be roughly twice the size of BF16 and
|
||||
can require a large amount of system RAM while saving.
|
||||
|
||||
## Install
|
||||
|
||||
Extract the folder into:
|
||||
|
||||
`ComfyUI/custom_nodes/ComfyUI-Precision-Model-Save/`
|
||||
|
||||
Restart ComfyUI.
|
||||
|
||||
The node appears under:
|
||||
|
||||
`model/merging -> Precision Model Save (FP32 Materialise)`
|
||||
|
||||
## Suggested test
|
||||
|
||||
- Feed the exact live merge that produces the expected image into this node.
|
||||
- Select `fp32`.
|
||||
- Keep `strip_diffusion_model_prefix = true`.
|
||||
- Save, restart or fully unload models, then load the saved file with
|
||||
`Load Diffusion Model`.
|
||||
- Render the same prompt, seed, sampler, scheduler, guidance and steps.
|
||||
|
||||
If FP32 matches the live merge while BF16/FP16 drift, the culprit is output
|
||||
precision. If FP32 still differs, the issue is likely not merely the final
|
||||
save dtype.
|
||||
|
||||
## Important
|
||||
|
||||
This is an experimental diagnostic node, not a universal saver. It is designed
|
||||
for ordinary floating-point diffusion models and live ComfyUI model merges.
|
||||
Quantised tensor formats may need architecture-specific handling.
|
||||
@@ -0,0 +1,3 @@
|
||||
from .precision_model_save import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,215 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
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"},
|
||||
),
|
||||
"save_dtype": (["fp32", "bf16", "fp16"], {"default": "fp32"}),
|
||||
"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 _materialise_weight(key, patch_description, output_dtype):
|
||||
"""
|
||||
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]
|
||||
|
||||
# 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__}"
|
||||
)
|
||||
|
||||
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:
|
||||
# Current ComfyUI supports intermediate_dtype here. The fallback
|
||||
# keeps the node usable on somewhat older builds.
|
||||
try:
|
||||
weight = comfy.lora.calculate_weight(
|
||||
patches,
|
||||
weight,
|
||||
key,
|
||||
intermediate_dtype=torch.float32,
|
||||
)
|
||||
except TypeError:
|
||||
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 weight.detach().to(
|
||||
device="cpu", dtype=output_dtype, copy=False
|
||||
).contiguous()
|
||||
|
||||
def save(
|
||||
self,
|
||||
model,
|
||||
filename_prefix,
|
||||
save_dtype,
|
||||
strip_diffusion_model_prefix,
|
||||
prompt=None,
|
||||
extra_pnginfo=None,
|
||||
):
|
||||
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(
|
||||
"[PrecisionModelSave] Materialising model patches in FP32; "
|
||||
"saving as %s to %s",
|
||||
save_dtype,
|
||||
output_path,
|
||||
)
|
||||
|
||||
# This returns the base tensor plus every live merge/patch attached to
|
||||
# the MODEL object, rather than only model.model.state_dict().
|
||||
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)
|
||||
|
||||
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."):]
|
||||
|
||||
state_dict[save_key] = self._materialise_weight(
|
||||
internal_key,
|
||||
patch_description,
|
||||
output_dtype,
|
||||
)
|
||||
|
||||
if index == 1 or index % 100 == 0 or index == total:
|
||||
logging.info(
|
||||
"[PrecisionModelSave] %d/%d tensors materialised",
|
||||
index,
|
||||
total,
|
||||
)
|
||||
|
||||
metadata = {
|
||||
"format": "pt",
|
||||
"comfy_precision_model_save": "1",
|
||||
"materialise_dtype": "fp32",
|
||||
"saved_dtype": save_dtype,
|
||||
"stripped_diffusion_model_prefix": str(
|
||||
bool(strip_diffusion_model_prefix)
|
||||
).lower(),
|
||||
}
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
# 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)
|
||||
return (output_path,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PrecisionModelSave": PrecisionModelSave,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PrecisionModelSave": "Precision Model Save (FP32 Materialise)",
|
||||
}
|
||||
Reference in New Issue
Block a user