Skip to content

Commit 98898c4

Browse files
Fix quant cache for comfy node (#1187)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Improved model loading to allow specifying a checkpoint name during speedup operations. This provides more flexibility when selecting and accelerating models. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent dc2cd27 commit 98898c4

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

onediff_comfy_nodes/_nodes.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ class SpeedupMixin:
4747
def speedup(
4848
self,
4949
model,
50+
ckpt_name: Optional[str] = None,
5051
inplace: bool = False,
5152
custom_booster: Optional[BoosterScheduler] = None,
5253
booster_settings: Optional[BoosterSettings] = None,
@@ -77,7 +78,7 @@ def speedup(
7778
booster.settings = (
7879
self.booster_settings if booster_settings is None else booster_settings
7980
)
80-
return (booster(model, *args, **kwargs),)
81+
return (booster(model, ckpt_name, *args, **kwargs),)
8182

8283

8384
class ModelSpeedup(SpeedupMixin):
@@ -244,6 +245,7 @@ def onediff_load_checkpoint(
244245
modelpatcher, clip, vae = self.load_checkpoint(ckpt_name)
245246
modelpatcher = self.speedup(
246247
modelpatcher,
248+
ckpt_name=ckpt_name,
247249
inplace=True,
248250
custom_booster=custom_booster,
249251
booster_settings=self.unet_booster_settings,

0 commit comments

Comments
 (0)