From 92326b68ba16144c0e7dcd2cdbc2aa9334e2f815 Mon Sep 17 00:00:00 2001
From: lixiang007666 <88304454@qq.com>
Date: Thu, 8 May 2025 15:46:24 +0800
Subject: [PATCH] Fix quant cache for comfy node

---
 onediff_comfy_nodes/_nodes.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/onediff_comfy_nodes/_nodes.py b/onediff_comfy_nodes/_nodes.py
index 74a864477..6a2622b9f 100644
--- a/onediff_comfy_nodes/_nodes.py
+++ b/onediff_comfy_nodes/_nodes.py
@@ -47,6 +47,7 @@ class SpeedupMixin:
     def speedup(
         self,
         model,
+        ckpt_name: Optional[str] = None,
         inplace: bool = False,
         custom_booster: Optional[BoosterScheduler] = None,
         booster_settings: Optional[BoosterSettings] = None,
@@ -77,7 +78,7 @@ def speedup(
         booster.settings = (
             self.booster_settings if booster_settings is None else booster_settings
         )
-        return (booster(model, *args, **kwargs),)
+        return (booster(model, ckpt_name, *args, **kwargs),)
 
 
 class ModelSpeedup(SpeedupMixin):
@@ -244,6 +245,7 @@ def onediff_load_checkpoint(
         modelpatcher, clip, vae = self.load_checkpoint(ckpt_name)
         modelpatcher = self.speedup(
             modelpatcher,
+            ckpt_name=ckpt_name,
             inplace=True,
             custom_booster=custom_booster,
             booster_settings=self.unet_booster_settings,