formatting and ADE support

2024-02-23 14:54:01 -07:00
parent ff935165a8
commit f269a7aafb
19 changed files with 1317 additions and 1277 deletions
@@ -1,21 +1,64 @@
 import torch

+from .embeddings import prepare_embeddings
+
+
+block_map = {
+    'input': {
+        1: 0,
+        2: 1,
+        4: 2,
+        5: 3,
+        7: 4,
+        8: 5
+    },
+    'middle': {
+        0: 6,
+    },
+    'output': {
+        3: 7,
+        4: 8,
+        5: 9,
+        6: 10,
+        7: 11,
+        8: 12,
+        9: 13,
+        10: 14,
+        11: 15
+    }
+}
+

 class FusersPatch(torch.nn.Module):
-  def __init__(self, fusers_list, objs, drop_box_mask, device, fusers_batch_size=5):
-    super(FusersPatch, self).__init__()
-    self.fusers_list = fusers_list
-    self.transformer_options = {'objs': objs, 'drop_box_mask': drop_box_mask}
-    self.device = device
-    self.fusers_batch_size = fusers_batch_size
-    # call_idx traces where in the UNet this is being called from
-    # a hack around lack of Comfy support
-    self.call_idx = 0
+    def __init__(self, conds, fusers_list, fusers_batch_size, positionnet, latent_shape, idxs, device):
+        super(FusersPatch, self).__init__()
+        self.conds = conds
+        self.fusers_list = fusers_list
+        self.fusers_batch_size = fusers_batch_size
+        self.positionnet = positionnet
+        self.latent_shape = latent_shape
+        self.idxs = idxs
+        self.device = device

-  @torch.no_grad()
-  def forward(self, x, extra_options):
-    fuser = self.fusers_list[self.call_idx %
-                             len(self.fusers_list)]
-    attn = fuser(x, self.transformer_options)
-    self.call_idx += 1
-    return attn.to(torch.float16)
+    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']
+
+        embeddings = prepare_embeddings(
+            self.conds, self.latent_shape, idxs, True)
+        for key in embeddings:
+            embeddings[key] = embeddings[key].to(self.device)
+        objs, drop_box_mask = self.positionnet(embeddings)
+        return {'objs': objs, 'drop_box_mask': drop_box_mask}
+
+    @torch.no_grad()
+    def forward(self, x, extra_options):
+        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)