xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-26 14:57:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

copy conds to device before repeating

Browse Source
This commit is contained in:
drhead 2025-05-25 12:49:49 -04:00 committed by GitHub
parent 87f4673af9
commit baf1e9f1fe
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
comfy/conds.py
View File

@ -11,17 +11,13 @@ class CONDRegular:
def _copy_with(self, cond): def _copy_with(self, cond):
return self.__class__(cond) return self.__class__(cond)
def _pin_memory(self, cond): def _pin_cond(self, device):
if cond.device == torch.device('cpu'): if self.cond.device == torch.device('cpu') and device_supports_non_blocking(device):
return cond.pin_memory() self.cond = self.cond.pin_memory(device)
else:
return cond
def process_cond(self, batch_size, device, **kwargs): def process_cond(self, batch_size, device, **kwargs):
if device_supports_non_blocking(device): self._pin_cond(device)
return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=True)) return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond.to(device, non_blocking=device_supports_non_blocking(device)), batch_size))
else:
return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device))
def can_concat(self, other): def can_concat(self, other):
if self.cond.shape != other.cond.shape: if self.cond.shape != other.cond.shape: