298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-13 15:03:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge d175cbd315 into 22ad513c72

Browse Source
This commit is contained in:
drhead 2025-04-11 09:46:27 -04:00 committed by GitHub
commit 1b5933c829
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
comfy/conds.py
View File

@ -1,6 +1,7 @@
import torch import torch
import math import math
import comfy.utils import comfy.utils
from comfy.model_management import device_should_use_non_blocking
class CONDRegular: class CONDRegular:
@ -10,8 +11,17 @@ 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):
if cond.device == torch.device('cpu'):
return cond.pin_memory()
else:
return cond
def process_cond(self, batch_size, device, **kwargs): def process_cond(self, batch_size, device, **kwargs):
return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device)) if device_should_use_non_blocking(device):
return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=True))
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: