brunvelop/ComfyUI

1import torch2 3class LatentRebatch:4    @classmethod5    def INPUT_TYPES(s):6        return {"required": { "latents": ("LATENT",),7                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),8                              }}9    RETURN_TYPES = ("LATENT",)10    INPUT_IS_LIST = True11    OUTPUT_IS_LIST = (True, )12 13    FUNCTION = "rebatch"14 15    CATEGORY = "latent/batch"16 17    @staticmethod18    def get_batch(latents, list_ind, offset):19        '''prepare a batch out of the list of latents'''20        samples = latents[list_ind]['samples']21        shape = samples.shape22        mask = latents[list_ind]['noise_mask'] if 'noise_mask' in latents[list_ind] else torch.ones((shape[0], 1, shape[2]*8, shape[3]*8), device='cpu')23        if mask.shape[-1] != shape[-1] * 8 or mask.shape[-2] != shape[-2]:24            torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[-2]*8, shape[-1]*8), mode="bilinear")25        if mask.shape[0] < samples.shape[0]:26            mask = mask.repeat((shape[0] - 1) // mask.shape[0] + 1, 1, 1, 1)[:shape[0]]27        if 'batch_index' in latents[list_ind]:28            batch_inds = latents[list_ind]['batch_index']29        else:30            batch_inds = [x+offset for x in range(shape[0])]31        return samples, mask, batch_inds32 33    @staticmethod34    def get_slices(indexable, num, batch_size):35        '''divides an indexable object into num slices of length batch_size, and a remainder'''36        slices = []37        for i in range(num):38            slices.append(indexable[i*batch_size:(i+1)*batch_size])39        if num * batch_size < len(indexable):40            return slices, indexable[num * batch_size:]41        else:42            return slices, None43    44    @staticmethod45    def slice_batch(batch, num, batch_size):46        result = [LatentRebatch.get_slices(x, num, batch_size) for x in batch]47        return list(zip(*result))48 49    @staticmethod50    def cat_batch(batch1, batch2):51        if batch1[0] is None:52            return batch253        result = [torch.cat((b1, b2)) if torch.is_tensor(b1) else b1 + b2 for b1, b2 in zip(batch1, batch2)]54        return result55 56    def rebatch(self, latents, batch_size):57        batch_size = batch_size[0]58 59        output_list = []60        current_batch = (None, None, None)61        processed = 062 63        for i in range(len(latents)):64            # fetch new entry of list65            #samples, masks, indices = self.get_batch(latents, i)66            next_batch = self.get_batch(latents, i, processed)67            processed += len(next_batch[2])68            # set to current if current is None69            if current_batch[0] is None:70                current_batch = next_batch71            # add previous to list if dimensions do not match72            elif next_batch[0].shape[-1] != current_batch[0].shape[-1] or next_batch[0].shape[-2] != current_batch[0].shape[-2]:73                sliced, _ = self.slice_batch(current_batch, 1, batch_size)74                output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]})75                current_batch = next_batch76            # cat if everything checks out77            else:78                current_batch = self.cat_batch(current_batch, next_batch)79 80            # add to list if dimensions gone above target batch size81            if current_batch[0].shape[0] > batch_size:82                num = current_batch[0].shape[0] // batch_size83                sliced, remainder = self.slice_batch(current_batch, num, batch_size)84                85                for i in range(num):86                    output_list.append({'samples': sliced[0][i], 'noise_mask': sliced[1][i], 'batch_index': sliced[2][i]})87 88                current_batch = remainder89 90        #add remainder91        if current_batch[0] is not None:92            sliced, _ = self.slice_batch(current_batch, 1, batch_size)93            output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]})94 95        #get rid of empty masks96        for s in output_list:97            if s['noise_mask'].mean() == 1.0:98                del s['noise_mask']99 100        return (output_list,)101 102NODE_CLASS_MAPPINGS = {103    "RebatchLatents": LatentRebatch,104}105 106NODE_DISPLAY_NAME_MAPPINGS = {107    "RebatchLatents": "Rebatch Latents",108}