fred-dev/comfy_ui_ali

1import torch2import math3import comfy.utils4 5 6class CONDRegular:7    def __init__(self, cond):8        self.cond = cond9 10    def _copy_with(self, cond):11        return self.__class__(cond)12 13    def process_cond(self, batch_size, device, **kwargs):14        return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device))15 16    def can_concat(self, other):17        if self.cond.shape != other.cond.shape:18            return False19        return True20 21    def concat(self, others):22        conds = [self.cond]23        for x in others:24            conds.append(x.cond)25        return torch.cat(conds)26 27class CONDNoiseShape(CONDRegular):28    def process_cond(self, batch_size, device, area, **kwargs):29        data = self.cond30        if area is not None:31            dims = len(area) // 232            for i in range(dims):33                data = data.narrow(i + 2, area[i + dims], area[i])34 35        return self._copy_with(comfy.utils.repeat_to_batch_size(data, batch_size).to(device))36 37 38class CONDCrossAttn(CONDRegular):39    def can_concat(self, other):40        s1 = self.cond.shape41        s2 = other.cond.shape42        if s1 != s2:43            if s1[0] != s2[0] or s1[2] != s2[2]: #these 2 cases should not happen44                return False45 46            mult_min = math.lcm(s1[1], s2[1])47            diff = mult_min // min(s1[1], s2[1])48            if diff > 4: #arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much49                return False50        return True51 52    def concat(self, others):53        conds = [self.cond]54        crossattn_max_len = self.cond.shape[1]55        for x in others:56            c = x.cond57            crossattn_max_len = math.lcm(crossattn_max_len, c.shape[1])58            conds.append(c)59 60        out = []61        for c in conds:62            if c.shape[1] < crossattn_max_len:63                c = c.repeat(1, crossattn_max_len // c.shape[1], 1) #padding with repeat doesn't change result64            out.append(c)65        return torch.cat(out)66 67class CONDConstant(CONDRegular):68    def __init__(self, cond):69        self.cond = cond70 71    def process_cond(self, batch_size, device, **kwargs):72        return self._copy_with(self.cond)73 74    def can_concat(self, other):75        if self.cond != other.cond:76            return False77        return True78 79    def concat(self, others):80        return self.cond81