wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-03-12 12:47:45 +08:00
Code Issues Actions 17 Packages Projects Releases Wiki Activity

Allow model patches to have a cleanup function. (#12878)

Browse Source 
The function gets called after sampling is finished.
This commit is contained in:
comfyanonymous 2026-03-10 17:09:12 -07:00 committed by GitHub
parent 8086468d2a
commit 3ad36d6be6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 22 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
comfy/model_patcher.py
View File

@ -599,6 +599,27 @@ class ModelPatcher:
return models return models
def model_patches_call_function(self, function_name="cleanup", arguments={}):
to = self.model_options["transformer_options"]
if "patches" in to:
patches = to["patches"]
for name in patches:
patch_list = patches[name]
for i in range(len(patch_list)):
if hasattr(patch_list[i], function_name):
getattr(patch_list[i], function_name)(**arguments)
if "patches_replace" in to:
patches = to["patches_replace"]
for name in patches:
patch_list = patches[name]
for k in patch_list:
if hasattr(patch_list[k], function_name):
getattr(patch_list[k], function_name)(**arguments)
if "model_function_wrapper" in self.model_options:
wrap_func = self.model_options["model_function_wrapper"]
if hasattr(wrap_func, function_name):
getattr(wrap_func, function_name)(**arguments)
def model_dtype(self): def model_dtype(self):
if hasattr(self.model, "get_dtype"): if hasattr(self.model, "get_dtype"):
return self.model.get_dtype() return self.model.get_dtype()
@ -1062,6 +1083,7 @@ class ModelPatcher:
return comfy.lora.calculate_weight(patches, weight, key, intermediate_dtype=intermediate_dtype) return comfy.lora.calculate_weight(patches, weight, key, intermediate_dtype=intermediate_dtype)
def cleanup(self): def cleanup(self):
self.model_patches_call_function(function_name="cleanup")
self.clean_hooks() self.clean_hooks()
if hasattr(self.model, "current_patcher"): if hasattr(self.model, "current_patcher"):
self.model.current_patcher = None self.model.current_patcher = None