From 3c80da7f87ee359b2d06f107cb3c0797079dfbbb Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Fri, 18 Sep 2026 11:35:35 -0700
Subject: [PATCH] Fix ace step VAE decode crashing on non bf16 GPUs. (#16405)

---
 comfy/sd.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/comfy/sd.py b/comfy/sd.py
index 230d4cf82..4380cffa5 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -1095,6 +1095,9 @@ class VAE:
         self.patcher = mp(self.first_stage_model, load_device=self.device, offload_device=offload_device, fast_disk=fast_disk)
 
         m, u = self.first_stage_model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic())
+        if not self.patcher.is_dynamic():
+            # Lazy parameters only exist after loading the state dict.
+            self.first_stage_model.to(self.vae_dtype)
         if len(m) > 0:
             logging.warning("Missing VAE keys {}".format(m))