From 08a077c8ef811283a3eaaa5d70ef9890c682c1dd Mon Sep 17 00:00:00 2001 From: LUCAS M MORGAN Date: Fri, 29 May 2026 20:57:41 -0500 Subject: [PATCH] Add model torch settings patch node --- README.md | 14 +++++++++ pyproject.toml | 2 +- utility_nodes.py | 77 ++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 92 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 2ed888e..75f004d 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,20 @@ Inputs: The CPU methods use `color-matcher`. The `reinhard_lab_gpu` method uses ComfyUI's Kornia dependency for Lab-space matching. +### `Model Patch Torch Settings` + +Patches a `MODEL` with ComfyUI model callbacks for PyTorch backend settings. + +Inputs: + +- `model`: model to clone and patch +- `enable_fp16_accumulation`: toggles + `torch.backends.cuda.matmul.allow_fp16_accumulation` while the model runs + +When enabled, the patched model turns full FP16 accumulation on before model +execution and turns it back off during cleanup. When disabled, it forces the +setting off before model execution. + ### `LM Studio Unified (URL + API Key)` Sends text, an image, or both to an LM Studio OpenAI-compatible chat endpoint. diff --git a/pyproject.toml b/pyproject.toml index d29b4db..228fc1c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "enviral-design-node-pack" -version = "0.1.10" +version = "0.1.11" description = "Small, dependency-light ComfyUI utility nodes from Enviral Design." readme = "README.md" requires-python = ">=3.10" diff --git a/utility_nodes.py b/utility_nodes.py index 855bbaf..219f3e1 100644 --- a/utility_nodes.py +++ b/utility_nodes.py @@ -478,3 +478,80 @@ an optional Kornia Lab-space Reinhard path for GPU-friendly batch work. NODE_CLASS_MAPPINGS["EnviralColorMatchV2"] = EnviralColorMatchV2 NODE_DISPLAY_NAME_MAPPINGS["EnviralColorMatchV2"] = "Enviral Color Match V2" + + +class EnviralModelPatchTorchSettings: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "enable_fp16_accumulation": ( + "BOOLEAN", + { + "default": False, + "tooltip": ( + "Toggles torch.backends.cuda.matmul.allow_fp16_accumulation " + "while the patched model runs." + ), + }, + ), + }, + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "patch" + CATEGORY = "EnviralDesign/model" + DESCRIPTION = """ +Adds model callbacks that enable or disable PyTorch full FP16 accumulation for +CUDA FP16 matmuls while the model runs. +""" + EXPERIMENTAL = True + + @staticmethod + def _ensure_supported(): + if not hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): + raise RuntimeError( + "torch.backends.cuda.matmul.allow_fp16_accumulation is not " + "available in this PyTorch build." + ) + + @staticmethod + def _set_fp16_accumulation(enabled): + logging.info( + "Patching torch settings: " + "torch.backends.cuda.matmul.allow_fp16_accumulation = %s", + enabled, + ) + torch.backends.cuda.matmul.allow_fp16_accumulation = enabled + + def patch(self, model, enable_fp16_accumulation): + from comfy.patcher_extension import CallbacksMP + + self._ensure_supported() + if not hasattr(model, "clone") or not hasattr(model, "add_callback"): + raise RuntimeError( + "Model Patch Torch Settings requires a ComfyUI MODEL patcher " + "with clone() and add_callback() support." + ) + + model_clone = model.clone() + + def enable_callback(model_patcher): + self._set_fp16_accumulation(True) + + def disable_callback(model_patcher): + self._set_fp16_accumulation(False) + + if enable_fp16_accumulation: + model_clone.add_callback(CallbacksMP.ON_PRE_RUN, enable_callback) + model_clone.add_callback(CallbacksMP.ON_CLEANUP, disable_callback) + else: + model_clone.add_callback(CallbacksMP.ON_PRE_RUN, disable_callback) + + return (model_clone,) + + +NODE_CLASS_MAPPINGS["EnviralModelPatchTorchSettings"] = EnviralModelPatchTorchSettings +NODE_DISPLAY_NAME_MAPPINGS["EnviralModelPatchTorchSettings"] = "Model Patch Torch Settings"