diff --git a/README.md b/README.md new file mode 100644 index 0000000..7aafbc4 --- /dev/null +++ b/README.md @@ -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. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3aa636f --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .precision_model_save import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/precision_model_save.py b/precision_model_save.py new file mode 100644 index 0000000..b59a535 --- /dev/null +++ b/precision_model_save.py @@ -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)", +}