nodes/logtd-ComfyUI-InstanceDiffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

batching

Browse Source
This commit is contained in:
logtd
2024-02-23 21:42:42 -07:00
parent 210e36fbb0
commit 8ea7395006
3 changed files with 30 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
conditioning/fusers_patch.py
+27 -10
View File
@@ -40,14 +40,7 @@ class FusersPatch(torch.nn.Module):
self.idxs = idxs
self.device = device
def _get_position_objs(self, extra_options):
idxs = None
if extra_options is not None:
if 'ad_params' in extra_options:
idxs = extra_options['ad_params']['sub_idxs']
elif 'sub_idxs' in extra_options:
idxs = extra_options['sub_idxs']
def _get_position_objs(self, idxs):
embeddings = prepare_embeddings(
self.conds, self.latent_shape, idxs, True)
for key in embeddings:
@@ -55,10 +48,34 @@ class FusersPatch(torch.nn.Module):
objs, drop_box_mask = self.positionnet(embeddings)
return {'objs': objs, 'drop_box_mask': drop_box_mask}
def _get_idxs(self, extra_options):
idxs = None
if extra_options is not None:
if 'ad_params' in extra_options:
idxs = extra_options['ad_params']['sub_idxs']
elif 'sub_idxs' in extra_options:
idxs = extra_options['sub_idxs']
return idxs
@torch.no_grad()
def forward(self, x, extra_options):
print('--------', x.shape)
block, idx = extra_options['block']
fuser_idx = block_map[block][idx]
fuser = self.fusers_list[fuser_idx]
attn = fuser(x, self._get_position_objs(extra_options))
return attn.to(torch.float16)
attn_total = []
idxs = self._get_idxs(extra_options)
print('idxs', idxs)
for i in range(0, len(x), self.fusers_batch_size):
batch_idxs = idxs[i:i+self.fusers_batch_size]
x_idxs = list(range(i, i+self.fusers_batch_size))
# print('batch_idxs', batch_idxs, x_idxs)
# print('x[x_idxs].shape', x[x_idxs].shape)
# print('doing attn')
attn = fuser(x[x_idxs], self._get_position_objs(batch_idxs))
# print('attn.shape', attn.shape)
attn_total.append(attn)
attn_total = torch.cat(attn_total)
print(attn_total.shape)
return attn_total.to(torch.float16)