Merge baf1e9f1fe6976e4be876457ea22e4c66fb9f962 into aebac221937b511d46fe601656acdc753435b849

This commit is contained in:
drhead 2025-08-03 20:29:20 +08:00 committed by GitHub
commit 39babfe942
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,6 +1,7 @@
import torch
import math
import comfy.utils
from comfy.model_management import device_supports_non_blocking
class CONDRegular:
@ -10,8 +11,13 @@ class CONDRegular:
def _copy_with(self, cond):
return self.__class__(cond)
def _pin_cond(self, device):
if self.cond.device == torch.device('cpu') and device_supports_non_blocking(device):
self.cond = self.cond.pin_memory(device)
def process_cond(self, batch_size, device, **kwargs):
return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device))
self._pin_cond(device)
return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond.to(device, non_blocking=device_supports_non_blocking(device)), batch_size))
def can_concat(self, other):
if self.cond.shape != other.cond.shape: