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

make conds use non-blocking transfers

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

9
comfy/conds.py
View File

@ -1,6 +1,7 @@
import torch
import math
import comfy.utils
from comfy.model_management import device_should_use_non_blocking
class CONDRegular:
@ -10,8 +11,14 @@ class CONDRegular:
def _copy_with(self, 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):
return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device))
return self._copy_with(comfy.utils.repeat_to_batch_size(self._pin_memory(self.cond), batch_size).to(device, non_blocking=device_should_use_non_blocking(device))
def can_concat(self, other):
if self.cond.shape != other.cond.shape: