From 9dc737e560fad298c783b82418479b68842d2f57 Mon Sep 17 00:00:00 2001 From: DocShotgun <126566557+DocShotgun@users.noreply.github.com> Date: Fri, 28 Feb 2025 10:59:29 -0800 Subject: [PATCH] Fix nightly torch check for fp16 accumulation --- nodes/model_optimization_nodes.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/nodes/model_optimization_nodes.py b/nodes/model_optimization_nodes.py index b97142a..15075f8 100644 --- a/nodes/model_optimization_nodes.py +++ b/nodes/model_optimization_nodes.py @@ -197,13 +197,14 @@ class DiffusionModelLoaderKJ(BaseLoaderKJ): model_options["dtype"] = torch.float8_e4m3fn model_options["fp8_optimizations"] = True - try: - if enable_fp16_accumulation: + if enable_fp16_accumulation: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): torch.backends.cuda.matmul.allow_fp16_accumulation = True else: + raise RuntimeError("Failed to set fp16 accumulation, this requires pytorch 2.7.0 nightly currently") + else: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): torch.backends.cuda.matmul.allow_fp16_accumulation = False - except: - raise RuntimeError("Failed to set fp16 accumulation, this requires pytorch 2.7.0 nightly currently") unet_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name) model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options) @@ -673,4 +674,4 @@ class TorchCompileCosmosModel: except: raise RuntimeError("Failed to compile model") - return (m, ) \ No newline at end of file + return (m, )