Skip to content

Commit

Permalink
fix: adjust to use new signature for cleanup_models
Browse files Browse the repository at this point in the history
  • Loading branch information
tazlin committed Dec 11, 2024
1 parent d6af90d commit d49a637
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions hordelib/comfy_horde.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@
_comfy_model_loading: types.ModuleType
_comfy_free_memory: Callable[[float, torch.device, list], None]
"""Will aggressively unload models from memory"""
_comfy_cleanup_models: Callable[[bool], None]
_comfy_cleanup_models: Callable[[None], None]
"""Will unload unused models from memory"""
_comfy_soft_empty_cache: Callable
"""Triggers comfyui and torch to empty their caches"""
Expand Down Expand Up @@ -944,7 +944,7 @@ def _run_pipeline(
if self.aggressive_unloading:
global _comfy_cleanup_models
logger.debug("Cleaning up models")
_comfy_cleanup_models(False)
_comfy_cleanup_models()
_comfy_soft_empty_cache()

stdio.replay()
Expand Down

0 comments on commit d49a637

Please sign in to comment.