From 159995e440b95efb95d9b95a4bb8d0cce2f84868 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Wed, 4 Feb 2026 16:47:23 +0200 Subject: [PATCH] Fix for newer ComfyUI version --- nodes_v2.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/nodes_v2.py b/nodes_v2.py index d90ef95..ea61ed9 100644 --- a/nodes_v2.py +++ b/nodes_v2.py @@ -902,7 +902,10 @@ high_vram: uses Accelerate to load weights to GPU, slightly faster model loading try: print(f"Attempting to load SDXL model from node inputs") mm.load_model_gpu(model) - sdxl_state_dict = model.model.state_dict_for_saving(None, vae.get_sd(), None) + try: + sdxl_state_dict = model.model.state_dict_for_saving(model.model.diffusion_model.state_dict(), vae_state_dict=vae.get_sd()) + except: + sdxl_state_dict = model.model.state_dict_for_saving(None, vae.get_sd(), None) if is_accelerate_available: for key in sdxl_state_dict: set_module_tensor_to_device(self.model, key, device=device, dtype=dtype, value=sdxl_state_dict[key]) @@ -1085,7 +1088,10 @@ high_vram: uses Accelerate to load weights to GPU, slightly faster model loading try: print(f"Attempting to load SDXL model from node inputs") mm.load_model_gpu(model) - sdxl_state_dict = model.model.state_dict_for_saving(None, vae.get_sd(), None) + try: + sdxl_state_dict = model.model.state_dict_for_saving(model.model.diffusion_model.state_dict(), vae_state_dict=vae.get_sd()) + except: + sdxl_state_dict = model.model.state_dict_for_saving(None, vae.get_sd(), None) if is_accelerate_available: for key in sdxl_state_dict: set_module_tensor_to_device(self.model, key, device=device, dtype=dtype, value=sdxl_state_dict[key])