diff --git a/precision_model_save.py b/precision_model_save.py index b59a535..c0a2d1a 100644 --- a/precision_model_save.py +++ b/precision_model_save.py @@ -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)" + ), }