Skip to content

Commit

Permalink
Merge PR #163 from Kosinkadink/develop - SparseCtrl lowvram fix for n…
Browse files Browse the repository at this point in the history
…ewest ComfyUI

Make SparseCtrl work with newest ComfyUI updates
  • Loading branch information
Kosinkadink committed Aug 20, 2024
2 parents 85d4970 + d98cd8d commit 1b3cb5f
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 3 deletions.
12 changes: 10 additions & 2 deletions adv_control/control_sparsectrl.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,9 +114,13 @@ def __init__(self, *args, **kwargs):
self.model: SparseControlNet
super().__init__(*args, **kwargs)

def patch_model_lowvram(self, device_to=None, *args, **kwargs):
patched_model = super().patch_model_lowvram(device_to, *args, **kwargs)
def load(self, device_to=None, lowvram_model_memory=0, *args, **kwargs):
to_return = super().load(device_to=device_to, lowvram_model_memory=lowvram_model_memory, *args, **kwargs)
if lowvram_model_memory > 0:
self._patch_lowvram_extras(device_to=device_to)
return to_return

def _patch_lowvram_extras(self, device_to=None):
if self.model.motion_wrapper is not None:
# figure out the tensors (likely pe's) that should be cast to device besides just the named_modules
remaining_tensors = list(self.model.motion_wrapper.state_dict().keys())
Expand All @@ -134,6 +138,10 @@ def patch_model_lowvram(self, device_to=None, *args, **kwargs):
if device_to is not None:
comfy.utils.set_attr(self.model.motion_wrapper, key, comfy.utils.get_attr(self.model.motion_wrapper, key).to(device_to))

# NOTE: no longer called by ComfyUI, but here for backwards compatibility
def patch_model_lowvram(self, device_to=None, *args, **kwargs):
patched_model = super().patch_model_lowvram(device_to, *args, **kwargs)
self._patch_lowvram_extras(device_to=device_to)
return patched_model

def clone(self):
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
[project]
name = "comfyui-advanced-controlnet"
description = "Nodes for scheduling ControlNet strength across timesteps and batched latents, as well as applying custom weights and attention masks."
version = "1.2.0"
version = "1.2.1"
license = { file = "LICENSE" }
dependencies = []

Expand Down

0 comments on commit 1b3cb5f

Please sign in to comment.