init

.gitignore

@@ -0,0 +1,3 @@
+.git
+.vscode
+**/__pycache__

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,2 @@
+# ComfyUI-InstanceDiffusion
+ComfyUI nodes to use InstanceDiffusion

__init__.py

@@ -0,0 +1,17 @@
+from .nodes.instance_diffusion_tracking_prompt_node import InstanceDiffusionTrackingPromptNode
+
+NODE_CLASS_MAPPINGS = {
+    "LoadInstanceDiffusion": None,
+    "LoadInstancePositionNetModel": None,
+    "LoadInstanceScaleUModel": None,
+    "LoadInstanceAttentionModel": None,
+    "InstanceDiffusionTrackingPrompt": InstanceDiffusionTrackingPromptNode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LoadInstanceDiffusion": "Load Instance Diffusion",
+    "LoadInstancePositionNetModel": "Load Instance PositionNet Model",
+    "LoadInstanceScaleUModel": "Load Instance ScaleU Model",
+    "LoadInstanceAttentionModel": "Load Instance Attention Model",
+    "InstanceDiffusionTrackingPrompt": "Instance Diffusion Tracking Prompt"
+}

conditioning/embeddings.py

@@ -0,0 +1,203 @@
+import torch
+import numpy as np
+from skimage.transform import resize
+from ..utils.decode_item import binary_mask_to_polygon, sample_uniform_sparse_points
+
+
+N_SCRIBBLE_POINTS = 20
+N_POLYGON_POINTS = 256
+N_MAX_OBJECTS = 30
+
+
+def get_point_from_box(bbox):
+  x0, y0, x1, y1 = bbox[0], bbox[1], bbox[2], bbox[3]
+  return [(x0 + x1) / 2.0, (y0 + y1) / 2.0]
+
+
+def get_empty_binary_mask(img_width, img_height):
+  return np.zeros((img_width, img_height, 1))
+
+
+def sample_random_points_from_mask(mask, k=N_SCRIBBLE_POINTS):
+  mask = mask[:, :, 0]
+  # Find the coordinates of non-zero pixels in the binary mask
+  nonzero_coords = np.transpose(np.nonzero(mask))
+
+  # Randomly sample 'k' points
+  # return all zeros if there is no non-zero pixel
+  if len(nonzero_coords) == 0:
+    xy_points = [0 for _ in range(k * 2)]
+    return xy_points
+
+  # randomly sample with replacement if there are not enough non-zero pixels
+  if len(nonzero_coords) < k and len(nonzero_coords) > 0:
+    random_indices = np.random.choice(len(nonzero_coords), k, replace=True)
+  # randomly sample withiout replacement if there are enough non-zero pixels
+  else:
+    random_indices = np.random.choice(
+        len(nonzero_coords), k, replace=False)
+  sampled_points = nonzero_coords[random_indices]
+
+  # order the points by their distance to (0, 0)
+  # center = np.array([mask.shape[0] // 2, mask.shape[1] // 2])
+  center = np.array([0, 0])
+  sampled_points = sorted(sampled_points, key=lambda x: np.linalg.norm(
+      np.array(x) - center))  # np.linalg.norm
+
+  # concatenate x and y coordinates and return them as a list
+  # [x1,y1,x2,y2,...,x_k,y_k]
+  xy_points = []
+  for x in sampled_points:
+    xy_points.append(float(x[1]))
+    xy_points.append(float(x[0]))
+  return xy_points
+
+
+def convert_points(points, img_width, img_height):
+  # convert polygons/scribbless' coordinates to the relative values (0, 1)
+  for i in range(len(points)):
+    if i % 2 == 0:
+      points[i] = min(points[i] / img_width, 1.0)
+    else:
+      points[i] = min(points[i] / img_height, 1.0)
+  return points
+
+
+def sample_sparse_points_from_mask(mask, k=256):
+  n_points = k
+  n_polygons = n_points // 2  # half points should be sampled from the polygons
+  mask = mask[:, :, 0]
+  # sample sparse points from the polygons (boundary)
+  polygons = binary_mask_to_polygon(mask, tolerance=0.0)
+  # concatenate polygons to a single list
+  polygons_single = []
+  for polygon in polygons:
+    polygons_single += polygon
+  if len(polygons_single) != 0:
+    # uniformly sample points from the polygon
+    polygons_single = np.array(polygons_single).reshape(-1, 2)
+    indexes = np.linspace(0, polygons_single.shape[0] - 1, n_polygons)
+    indexes = list([int(i) for i in indexes])
+
+    polygons_single = polygons_single[indexes]
+    sampled_polygons = [(x[0], x[1]) for x in polygons_single]
+  else:
+    return [0 for _ in range(256 * 2)]
+
+  # sample sparse points from the mask
+  n_inside_points = n_points - len(sampled_polygons)
+  inside_points = sample_uniform_sparse_points(mask, n_inside_points)
+
+  # combine inside_points and sampled_polygons
+  xy_points = inside_points + sampled_polygons
+
+  # order the points by their distance to (0, 0)
+  center = np.array([0, 0])
+  xy_points = sorted(xy_points, key=lambda x: np.linalg.norm(
+      np.array(x) - center))  # np.linalg.norm
+
+  # return the sampled points
+  sampled_points = []
+  for x in xy_points:
+    sampled_points.append(x[0])
+    sampled_points.append(x[1])
+  return sampled_points
+
+
+# [x0, y0, x1, y1]
+def get_grounding_input_from_coords(coords, img_width, img_height):
+  x0, y0, x1, y1, coord_width, coord_height = coords
+  location = [x0 / coord_width, y0 / coord_height,
+              x1 / coord_width, y1 / coord_height]
+
+  point = get_point_from_box(location)
+  binary_mask = get_empty_binary_mask(img_width, img_height)
+
+  scribble = sample_random_points_from_mask(binary_mask, k=N_SCRIBBLE_POINTS)
+  scribble = convert_points(scribble, img_width, img_height)
+
+  polygon = sample_sparse_points_from_mask(binary_mask, k=N_POLYGON_POINTS)
+  polygon = convert_points(polygon, img_width, img_height)
+
+  segment = resize(binary_mask.astype(np.float32),
+                   (img_width, img_height)).squeeze()
+  # segment = np.stack(segment).astype(np.float32).squeeze() if len(segment) > 0 else segment
+
+  return dict(
+      polygon=polygon,
+      scribble=scribble,
+      segment=segment,
+      box=location,
+      point=point,
+  )
+
+
+def create_zero_input_tensors(n_frames, img_width, img_height):
+  masks = torch.zeros(n_frames, N_MAX_OBJECTS)
+  text_masks = torch.zeros(n_frames, N_MAX_OBJECTS)
+  text_embeddings = torch.zeros(n_frames, N_MAX_OBJECTS, 768)
+  box_embeddings = torch.zeros(n_frames, N_MAX_OBJECTS, 4)
+  polygon_embeddings = torch.zeros(
+      n_frames, N_MAX_OBJECTS, N_POLYGON_POINTS * 2)
+  scribble_embeddings = torch.zeros(
+      n_frames, N_MAX_OBJECTS, N_SCRIBBLE_POINTS * 2)
+  segment_embeddings = torch.zeros(
+      n_frames, N_MAX_OBJECTS, img_width, img_height)  # TODO: width height order
+  point_embeddings = torch.zeros(n_frames, N_MAX_OBJECTS, 2)
+
+  return dict(
+      masks=masks,
+      text_masks=text_masks,
+      prompts=text_embeddings,
+      boxes=box_embeddings,
+      polygons=polygon_embeddings,
+      scribbles=scribble_embeddings,
+      segments=segment_embeddings,
+      points=point_embeddings
+  )
+
+
+def get_attn_mask(img_size=64):
+  return torch.zeros(N_MAX_OBJECTS, img_size, img_size)
+
+
+def prepare_embeddings(conds, latent_shape, idxs, use_masked_att=False):
+  batch_size, _, latent_height, latent_width = latent_shape
+  embeddings = create_zero_input_tensors(
+    batch_size, latent_width, latent_height)
+  if use_masked_att:
+    embeddings['att_masks'] = torch.zeros(
+        idx, N_MAX_OBJECTS, latent_width, latent_height)
+
+  for idx in idxs:
+    for cond_idx, cond in enumerate(conds):
+      if cond['positions'][idx] is None:
+        continue
+
+      grounding = get_grounding_input_from_coords(
+          cond['positions'][idx], latent_width, latent_height)
+      embeddings['masks'][idx][cond_idx] = 1
+      embeddings['text_masks'][idx][cond_idx] = 1
+      embeddings['prompts'][idx][cond_idx] = cond['cond_pooled']
+      embeddings['boxes'][idx][cond_idx] = torch.tensor(
+          grounding['box'])
+      embeddings['polygons'][idx][cond_idx] = torch.tensor(
+          grounding['polygon'])
+      embeddings['scribbles'][idx][cond_idx] = torch.tensor(
+          grounding['scribble'])
+      embeddings['segments'][idx][cond_idx] = torch.tensor(
+          grounding['segment'])
+      embeddings['points'][idx][cond_idx] = torch.tensor(
+          grounding['point'])
+
+      if use_masked_att:
+        box = grounding['box']
+        x1, y1, x2, y2 = int(np.round(box[0] * latent_width)), int(np.round(box[1] * latent_height)), int(
+            np.round(box[2] * latent_width)), int(np.round(box[3] * latent_height))
+        embeddings['att_masks'][idx][cond_idx][x1:x2, y1:y2] = 1
+
+  return embeddings
+
+
+def get_model_inputs():
+  return

conditioning/instance_conditioning.py

@@ -0,0 +1,63 @@
+import torch
+
+import comfy.model_management
+
+from .embeddings import prepare_embeddings
+
+
+class InstanceConditioning:
+  """
+  This class masquerades as both a ControlNet and Gligen in order to trigger setup
+  """
+
+  def __init__(self, fusers, positionnet, fusers_batch_size):
+    self.fusers_list = fusers['model_list']
+    self.positionnet = positionnet['model']
+    self.fusers_batch_size = fusers_batch_size
+    self.conds = []
+
+    # ControlNet hacks
+    self.load_device = comfy.model_management.get_torch_device()
+
+    # Gligen hacks
+    self.model = self
+
+    # AnimateDiff hacks
+    self.sub_idxs = None
+    self.full_latent_length = None
+    self.context_length = None
+
+  def get_fusers_patch(self, objs, drop_box_mask):
+    return
+
+  def set_position(self, latent_shape, _, device):
+    # Called in samplers by gligen cond to return middle attention patch
+    batch_size = latent_shape[0]
+    idxs = self.sub_idxs if self.sub_idxs is not None else list(
+      range(batch_size))
+
+    # TODO conds
+    embeddings = prepare_embeddings(self.conds, latent_shape, idxs, True)
+    objs, drop_box_mask = self.positionnet.to(device)(embeddings)
+    fusers_patch = self.get_fusers_patch(objs, drop_box_mask, device)
+    return fusers_patch
+
+  def add_conds(self, conds):
+    self.conds.extend(conds)
+
+  def get_models(self) -> list[torch.nn.Module]:
+    # Used to get models for loading/offloading
+    # TODO need to make this work for ControlNet and Gligen
+    return [(None, model) for model in [*self.fusers_list, self.positionnet]]
+
+  def inference_memory_requirements(self, dtype) -> int:
+    # Used to calculate memory requirements by ControlNet
+    return 0
+
+  def is_clone(self, other):
+    return other == self
+
+  # def __getattribute__(self, __name: str) -> Any:
+  #   if __name == 'gligen':
+  #     return 'something'  # TODO
+  #   return super().__getattribute__(__name)

constants.py

@@ -0,0 +1,5 @@
+
+INSTANCE_MODELS_DIR = "instance_models"
+INSTANCE_FUSERS_DIR = "fuser_models"
+INSTANCE_POSITIONNET_DIR = "positionnet_models"
+INSTANCE_SCALEU_DIR = "scaleu_models"

models/prepare_fusers.py

@@ -0,0 +1,30 @@
+import torch
+
+from ..modules.attention import GatedSelfAttentionDense
+
+
+def prepare_fusers(fusers_ckpt) -> list[torch.nn.Module]:
+  fusers_list = []
+  for key in fusers_ckpt['input_blocks']:
+    fusers_ckpt['input_blocks'][key]['params']['query_dim'] = fusers_ckpt['input_blocks'][key]['params']['n_heads'] * \
+      fusers_ckpt['input_blocks'][key]['params']['d_head']
+    fuser = GatedSelfAttentionDense(
+      **fusers_ckpt['input_blocks'][key]['params'])
+    fuser.load_state_dict(fusers_ckpt['input_blocks'][key]['state'])
+    fusers_list.append(fuser)
+
+  fusers_ckpt['middle_block']['1']['params']['query_dim'] = fusers_ckpt['middle_block']['1']['params']['n_heads'] * \
+      fusers_ckpt['middle_block']['1']['params']['d_head']
+  fuser = GatedSelfAttentionDense(**fusers_ckpt['middle_block']['1']['params'])
+  fuser.load_state_dict(fusers_ckpt['middle_block']['1']['state'])
+  fusers_list.append(fuser)
+
+  for key in fusers_ckpt['output_blocks']:
+    fusers_ckpt['output_blocks'][key]['params']['query_dim'] = fusers_ckpt['output_blocks'][key]['params']['n_heads'] * \
+      fusers_ckpt['output_blocks'][key]['params']['d_head']
+    fuser = GatedSelfAttentionDense(
+      **fusers_ckpt['output_blocks'][key]['params'])
+    fuser.load_state_dict(fusers_ckpt['output_blocks'][key]['state'])
+    fusers_list.append(fuser)
+
+  return fusers_list

models/prepare_positionnet.py

@@ -0,0 +1,26 @@
+import torch
+
+from ..modules.text_grounding_net import UniFusion
+
+
+def get_positionnet_default_params():
+  return {
+      "in_dim": 768,
+      "mid_dim": 3072,
+      "out_dim": 768,
+      "test_drop_boxes": False,
+      "test_drop_masks": True,
+      "test_drop_points": False,
+      "test_drop_scribbles": True,
+      "train_add_boxes": True,
+      "train_add_masks": True,
+      "train_add_points": True,
+      "train_add_scribbles": True,
+      "use_seperate_tokenizer": True,
+  }
+
+
+def prepare_positionnet(checkpoint, params) -> torch.nn.Module:
+  model = UniFusion(**params)
+  model.load_state_dict(checkpoint)
+  return model

models/prepare_scaleu.py

@@ -0,0 +1,22 @@
+import torch
+
+from ..modules.scaleu import ScaleU
+
+
+def get_scaleu_patch(scaleu_nets):
+  def scaleu_patch(h, hsp, transformer_options):
+    _, idx = transformer_options['block']
+    sk = scaleu_nets[idx](h, hsp)
+    return sk
+
+  return scaleu_patch
+
+
+def prepare_scaleu_nets(scaleu_ckpt) -> torch.nn.Module:
+  scaleu_nets = []
+  for i in range(12):
+    ckpt = scaleu_ckpt[f'{i}']
+    scaleu = ScaleU(True, len(ckpt['scaleu_b']), len(ckpt['scaleu_s']))
+    scaleu.load_state_dict(ckpt)
+    scaleu_nets.append(scaleu)
+  return scaleu_nets

modules/attention.py

@@ -0,0 +1,209 @@
+import torch
+import numpy as np
+from torch import nn
+from inspect import isfunction
+import torch.nn.functional as F
+
+
+def exists(val):
+  return val is not None
+
+
+def default(val, d):
+  if exists(val):
+    return val
+  return d() if isfunction(d) else d
+
+
+# feedforward
+class GEGLU(nn.Module):
+  def __init__(self, dim_in, dim_out):
+    super().__init__()
+    self.proj = nn.Linear(dim_in, dim_out * 2)
+
+  def forward(self, x):
+    x, gate = self.proj(x).chunk(2, dim=-1)
+    return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+  def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
+    super().__init__()
+    inner_dim = int(dim * mult)
+    dim_out = default(dim_out, dim)
+    project_in = nn.Sequential(
+        nn.Linear(dim, inner_dim),
+        nn.GELU()
+    ) if not glu else GEGLU(dim, inner_dim)
+
+    self.net = nn.Sequential(
+        project_in,
+        nn.Dropout(dropout),
+        nn.Linear(inner_dim, dim_out)
+    )
+
+  def forward(self, x):
+    return self.net(x)
+
+
+class SelfAttention(nn.Module):
+  def __init__(self, query_dim, heads=8, dim_head=64, dropout=0.):
+    super().__init__()
+    inner_dim = dim_head * heads
+    self.scale = dim_head ** -0.5
+    self.heads = heads
+
+    self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+    self.to_k = nn.Linear(query_dim, inner_dim, bias=False)
+    self.to_v = nn.Linear(query_dim, inner_dim, bias=False)
+
+    self.to_out = nn.Sequential(
+        nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
+
+  def forward(self, x, instance_options={}):
+    grounding_input = instance_options.get('grounding_input', None)
+    drop_box_mask = instance_options.get('drop_box_mask', None)
+    q = self.to_q(x)  # B*N*(H*C)
+    k = self.to_k(x)  # B*N*(H*C)
+    v = self.to_v(x)  # B*N*(H*C)
+
+    B, N, HC = q.shape
+    H = self.heads
+    C = HC // H
+
+    q = q.view(B, N, H, C).permute(
+        0, 2, 1, 3).reshape(B * H, N, C)  # (B*H)*N*C
+    k = k.view(B, N, H, C).permute(
+        0, 2, 1, 3).reshape(B * H, N, C)  # (B*H)*N*C
+    v = v.view(B, N, H, C).permute(
+        0, 2, 1, 3).reshape(B * H, N, C)  # (B*H)*N*C
+
+    sim = torch.einsum('b i c, b j c -> b i j', q, k) * \
+        self.scale  # (B*H)*N*N
+
+    if grounding_input is not None:
+      # att_masks: B*(n_objs*4)*sqrt(w_h)*sqrt(w_h)
+      n_objs = grounding_input["att_masks"].shape[1]
+
+      att_masks_reshape = None
+      if N - n_objs * 4 == 64 * 64:
+        att_masks_reshape = grounding_input["att_masks"]
+        att_key = "att_masks_selfAtt64"
+
+      # if the mask is all zeros or box&mask are dropped, we do not need to mask the self-attention
+      if att_masks_reshape is not None and torch.sum(att_masks_reshape) > 0.0 and not drop_box_mask:
+        # the attention mask is shared across all blocks for more efficient forwarding
+        # we only need to get the mask for the first block.
+        if att_key in grounding_input:
+          att_masks_ = grounding_input[att_key]
+        else:
+          w_h = att_masks_reshape.shape[2] * \
+              att_masks_reshape.shape[3]
+          att_masks_ = torch.ones(B, 1, N, N).type(
+              att_masks_reshape.dtype).to(att_masks_reshape.device)
+          # mask the self-attention between visual tokens
+          # all object patches do not exchange information with each other
+          # but object patches share information with all background regions
+          att_masks_reshape_v = att_masks_reshape.view(
+              B * n_objs, w_h, 1)
+
+          # get the self-attention between att_masks_reshape with batch matrix multiplication
+          #  1. att_masks_reshape_v: (B*n_objs)*w_h*1
+          #  2. att_masks_reshape_v.permute(0,1,3,2): (B*n_objs)*1*w_h
+          #  3. torch.bmm(att_masks_reshape_v, att_masks_reshape_v.permute(0,2,1)): (B*n_objs)*w_h*w_h
+          self_att_all = torch.bmm(
+              att_masks_reshape_v, att_masks_reshape_v.permute(0, 2, 1))
+
+          # sum the values alone the n_objs dimension
+          #  1. self_att_all: B*n_objs*w_h*w_h
+          #  2. self_att_all.sum(dim=1): B*w_h*w_h
+          self_att_ind_objs = self_att_all.view(
+              B, n_objs, w_h, w_h).sum(dim=1)
+
+          # sum att_masks_reshape alone the n_objs dimension and measure the self-attention with batch matrix multiplication
+          #  1. att_masks_reshape_all: B*n_objs*w_h*1
+          #  2. att_masks_reshape_all.sum(dim=1): B*w_h*1
+          #  3. set non-zero values to 1: B*w_h*1
+          #  4. use batch matrix multiplication to measure the self-attention: B*w_h*w_h
+          att_masks_reshape_all = att_masks_reshape_v.view(
+              B, n_objs, w_h, 1).sum(dim=1)
+          att_masks_reshape_all[att_masks_reshape_all >= 1.0] = 1.0
+          self_att_all_objs = torch.bmm(
+              att_masks_reshape_all, att_masks_reshape_all.permute(0, 2, 1))
+
+          # get the masks for avoiding information leakage between object patches
+          visual_token_masks = self_att_all_objs + self_att_ind_objs
+
+          # avoid the communications between objects and background
+          # objects-background can not communicate
+          visual_token_masks[self_att_ind_objs < 1.0] = 0.0
+          visual_token_masks[self_att_ind_objs >=
+                             1.0] = 1.0  # binay mask
+
+          att_masks_[:, :, :w_h, :w_h] = visual_token_masks.view(
+              B, 1, w_h, w_h)
+
+          # mask the self-attention between grounded tokens and the visual tokens
+          att_masks_reshape = att_masks_reshape.view(
+              B, 1, n_objs, w_h)
+          # the order of inputs are [box, point, scribble, mask]. only box and mask need to have masked self-attention
+          att_masks_[:, :, w_h:, :w_h] = att_masks_reshape.repeat(
+              1, 1, 4, 1)
+          att_masks_[:, :, w_h + n_objs:w_h + n_objs * 3, :w_h] = 1
+          att_masks_[:, :, :w_h, w_h:] = att_masks_reshape.permute(
+              0, 1, 3, 2).repeat(1, 1, 1, 4)
+          att_masks_[:, :, :w_h, w_h + n_objs:w_h + n_objs * 3] = 1
+
+          # add 1e-9 along the diagonal to avoid nan
+          diagonal_epsilon = torch.eye(N).view(1, 1, N, N).type(
+              att_masks_.dtype).to(att_masks_.device) * 1e-9
+          att_masks_ = att_masks_ + diagonal_epsilon
+
+          # save the masks for later blocks
+          grounding_input['next_att_mask'] = att_masks_
+        # the larger the threshold is, the smaller area the grounded info can influence
+        # if threshold = 1.01, all grounded info has no influence on the rest of the patches.
+        sim = sim.view(B, H, N, N).masked_fill(
+            att_masks_ <= 0.0, -np.inf).view(B * H, N, N)  # -np.inf
+
+    attn = sim.softmax(dim=-1)  # (B*H)*N*N
+
+    out = torch.einsum('b i j, b j c -> b i c', attn, v)  # (B*H)*N*C
+    out = out.view(B, H, N, C).permute(
+        0, 2, 1, 3).reshape(B, N, (H * C))  # B*N*(H*C)
+    return self.to_out(out)
+
+
+class GatedSelfAttentionDense(nn.Module):
+  def __init__(self, query_dim, context_dim, n_heads, d_head):
+    super().__init__()
+    # we need a linear projection since we need cat visual feature and obj feature
+    self.linear = nn.Linear(context_dim, query_dim)
+
+    self.attn = SelfAttention(
+        query_dim=query_dim, heads=n_heads, dim_head=d_head)
+    self.ff = FeedForward(query_dim, glu=True)
+
+    self.norm1 = nn.LayerNorm(query_dim)
+    self.norm2 = nn.LayerNorm(query_dim)
+
+    self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
+    self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
+
+    # this can be useful: we can externally change magnitude of tanh(alpha)
+    # for example, when it is set to 0, then the entire model is same as original one
+    self.scale = 1
+
+  def forward(self, x, instance_options={}):
+    objs = instance_options['objs']
+
+    N_visual = x.shape[1]
+    objs = self.linear(objs)
+    attention_output = self.attn(self.norm1(
+        torch.cat([x, objs], dim=1)), instance_options=instance_options)
+    x = x + self.scale * \
+        torch.tanh(self.alpha_attn) * attention_output[:, 0:N_visual, :]
+    x = x + self.scale * \
+        torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
+
+    return x.to(torch.float16)

modules/convnext.py

@@ -0,0 +1,217 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the original repo.
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import trunc_normal_, DropPath
+from timm.models.registry import register_model
+
+
+class Block(nn.Module):
+  r""" ConvNeXt Block. There are two equivalent implementations:
+  (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+  (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+  We use (2) as we find it slightly faster in PyTorch
+
+  Args:
+      dim (int): Number of input channels.
+      drop_path (float): Stochastic depth rate. Default: 0.0
+      layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+  """
+
+  def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
+    super().__init__()
+    self.dwconv = nn.Conv2d(dim, dim, kernel_size=7,
+                            padding=3, groups=dim)  # depthwise conv
+    self.norm = LayerNorm(dim, eps=1e-6)
+    # pointwise/1x1 convs, implemented with linear layers
+    self.pwconv1 = nn.Linear(dim, 4 * dim)
+    self.act = nn.GELU()
+    self.pwconv2 = nn.Linear(4 * dim, dim)
+    self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
+                              requires_grad=True) if layer_scale_init_value > 0 else None
+    self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+  def forward(self, x):
+    input = x
+    x = self.dwconv(x)
+    x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+    x = self.norm(x)
+    x = self.pwconv1(x)
+    x = self.act(x)
+    x = self.pwconv2(x)
+    if self.gamma is not None:
+      x = self.gamma * x
+    x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+    x = input + self.drop_path(x)
+    return x
+
+
+class ConvNeXt(nn.Module):
+  r""" ConvNeXt
+      A PyTorch impl of : `A ConvNet for the 2020s`  -
+        https://arxiv.org/pdf/2201.03545.pdf
+
+  Args:
+      in_chans (int): Number of input image channels. Default: 3
+      num_classes (int): Number of classes for classification head. Default: 1000
+      depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+      dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+      drop_path_rate (float): Stochastic depth rate. Default: 0.
+      layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+      head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+  """
+
+  def __init__(self, in_chans=3, num_classes=1000,
+               depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0.,
+               layer_scale_init_value=1e-6, head_init_scale=1.,
+               ):
+    super().__init__()
+
+    # stem and 3 intermediate downsampling conv layers
+    self.downsample_layers = nn.ModuleList()
+    stem = nn.Sequential(
+        nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+        LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
+    )
+    self.downsample_layers.append(stem)
+    for i in range(3):
+      downsample_layer = nn.Sequential(
+          LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+          nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+      )
+      self.downsample_layers.append(downsample_layer)
+
+    # 4 feature resolution stages, each consisting of multiple residual blocks
+    self.stages = nn.ModuleList()
+    dp_rates = [x.item()
+                for x in torch.linspace(0, drop_path_rate, sum(depths))]
+    cur = 0
+    for i in range(4):
+      stage = nn.Sequential(
+          *[Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                  layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])]
+      )
+      self.stages.append(stage)
+      cur += depths[i]
+
+  def _init_weights(self, m):
+    if isinstance(m, (nn.Conv2d, nn.Linear)):
+      trunc_normal_(m.weight, std=.02)
+      nn.init.constant_(m.bias, 0)
+
+  def forward_features(self, x):
+    for i in range(4):
+      x = self.downsample_layers[i](x)
+      x = self.stages[i](x)
+    return x
+
+  def forward(self, x):
+    x = self.forward_features(x)
+    # x = self.head(x)
+    return x
+
+
+class LayerNorm(nn.Module):
+  r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
+  The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
+  shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
+  with shape (batch_size, channels, height, width).
+  """
+
+  def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+    super().__init__()
+    self.weight = nn.Parameter(torch.ones(normalized_shape))
+    self.bias = nn.Parameter(torch.zeros(normalized_shape))
+    self.eps = eps
+    self.data_format = data_format
+    if self.data_format not in ["channels_last", "channels_first"]:
+      raise NotImplementedError
+    self.normalized_shape = (normalized_shape, )
+
+  def forward(self, x):
+    if self.data_format == "channels_last":
+      return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+    elif self.data_format == "channels_first":
+      u = x.mean(1, keepdim=True)
+      s = (x - u).pow(2).mean(1, keepdim=True)
+      x = (x - u) / torch.sqrt(s + self.eps)
+      x = self.weight[:, None, None] * x + self.bias[:, None, None]
+      return x
+
+
+model_urls = {
+    "convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
+    "convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
+    "convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
+    "convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
+    "convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
+    "convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
+    "convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
+    "convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
+    "convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
+}
+
+
+@register_model
+def convnext_tiny(pretrained=False, in_22k=False, **kwargs):
+  model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+  if pretrained:
+    url = model_urls['convnext_tiny_22k'] if in_22k else model_urls['convnext_tiny_1k']
+    checkpoint = torch.hub.load_state_dict_from_url(
+      url=url, map_location="cpu", check_hash=True)
+    # we remove classifer head
+    model.load_state_dict(checkpoint["model"], strict=False)
+  return model
+
+
+@register_model
+def convnext_small(pretrained=False, in_22k=False, **kwargs):
+  model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
+  if pretrained:
+    url = model_urls['convnext_small_22k'] if in_22k else model_urls['convnext_small_1k']
+    checkpoint = torch.hub.load_state_dict_from_url(
+      url=url, map_location="cpu")
+    model.load_state_dict(checkpoint["model"])
+  return model
+
+
+@register_model
+def convnext_base(pretrained=False, in_22k=False, **kwargs):
+  model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
+  if pretrained:
+    url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
+    checkpoint = torch.hub.load_state_dict_from_url(
+      url=url, map_location="cpu")
+    model.load_state_dict(checkpoint["model"])
+  return model
+
+
+@register_model
+def convnext_large(pretrained=False, in_22k=False, **kwargs):
+  model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
+  if pretrained:
+    url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']
+    checkpoint = torch.hub.load_state_dict_from_url(
+      url=url, map_location="cpu")
+    model.load_state_dict(checkpoint["model"])
+  return model
+
+
+@register_model
+def convnext_xlarge(pretrained=False, in_22k=False, **kwargs):
+  model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
+  if pretrained:
+    assert in_22k, "only ImageNet-22K pre-trained ConvNeXt-XL is available; please set in_22k=True"
+    url = model_urls['convnext_xlarge_22k']
+    checkpoint = torch.hub.load_state_dict_from_url(
+      url=url, map_location="cpu")
+    model.load_state_dict(checkpoint["model"])
+  return model

modules/scaleu.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+import torch.fft as fft
+
+
+def Fourier_filter(x_in, threshold, scale):
+    x = x_in
+    B, C, H, W = x.shape
+
+    # Non-power of 2 images must be float32
+    if (W & (W - 1)) != 0 or (H & (H - 1)) != 0:
+        x = x.to(dtype=torch.float32)
+
+    # FFT
+    x_freq = fft.fftn(x, dim=(-2, -1))
+    x_freq = fft.fftshift(x_freq, dim=(-2, -1))
+
+    B, C, H, W = x_freq.shape
+    mask = torch.ones((B, C, H, W), device=x.device)
+
+    crow, ccol = H // 2, W // 2
+    mask[..., crow - threshold : crow + threshold, ccol - threshold : ccol + threshold] = scale
+    x_freq = x_freq * mask
+
+    # IFFT
+    x_freq = fft.ifftshift(x_freq, dim=(-2, -1))
+    x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real
+
+    return x_filtered.to(dtype=x_in.dtype)
+
+
+class ScaleU(nn.Module):
+    def __init__(self, enable_se_scaleu=True, b_size=1280, s_size=1):
+        super(ScaleU, self).__init__()
+        self.scaleu_b = nn.Parameter(torch.zeros(b_size))
+        self.scaleu_s = nn.Parameter(torch.zeros(s_size))
+        self.enable_se_scaleu = enable_se_scaleu
+
+    def forward(self, h, hs_, transformer_options={}):
+      h = h.to(torch.float32)
+      hs_ = hs_.to(torch.float32)
+      b = torch.tanh(self.scaleu_b) + 1
+      s = torch.tanh(self.scaleu_s) + 1
+      if self.enable_se_scaleu:
+          hidden_mean = h.mean(1).unsqueeze(1) # B,1,H,W 
+          B = hidden_mean.shape[0]
+          hidden_max, _ = torch.max(hidden_mean.view(B, -1), dim=-1, keepdim=True) # B,1
+          hidden_min, _ = torch.min(hidden_mean.view(B, -1), dim=-1, keepdim=True) # B,1
+          # duplicate the hidden_mean dimension 1 to C
+          hidden_mean = (hidden_mean - hidden_min.unsqueeze(2).unsqueeze(3)) / (hidden_max - hidden_min).unsqueeze(2).unsqueeze(3) # B,1,H,W
+          b = torch.einsum('c,bchw->bchw', b-1, hidden_mean) + 1.0 # B,C,H,W
+          h = torch.einsum('bchw,bchw->bchw', h, b)
+      else:      
+          h = torch.einsum('bchw,c->bchw', h, b)
+      
+      hs_ = Fourier_filter(hs_, threshold=1, scale=s)
+      return h.to(torch.float16), hs_.to(torch.float16)

modules/text_grounding_net.py

@@ -0,0 +1,351 @@
+import torch
+import torch.nn as nn
+from .util import FourierEmbedder
+from .convnext import convnext_tiny
+
+
+class UniFusion(nn.Module):
+  def __init__(self, in_dim, out_dim, mid_dim=3072, fourier_freqs=8,
+               train_add_boxes=True, train_add_points=True, train_add_scribbles=True, train_add_masks=True,
+               test_drop_boxes=False, test_drop_points=False, test_drop_scribbles=True, test_drop_masks=False,
+               use_seperate_tokenizer=True):
+    super().__init__()
+    self.in_dim = in_dim
+    self.out_dim = out_dim
+    self.mid_dim = mid_dim
+
+    # InstanceDiffusion hyper-parameters
+    self.n_scribble_points = 20
+    self.n_polygon_points = 256
+    fourier_freqs = 16
+    fourier_freqs_polygons = 16
+    self.add_boxes = train_add_boxes
+    self.add_points = train_add_points
+    self.add_scribbles = train_add_scribbles
+    self.add_masks = train_add_masks
+    self.use_seperate_tokenizer = use_seperate_tokenizer
+
+    # Use instance masks as additional model inputs for mask conditioned image generation
+    self.use_segs = True if self.add_masks else False
+
+    if self.use_segs:
+      in_dim = 30
+      self.resize_input = 512
+      self.down_factor = 64  # determined by the convnext backbone
+      # from num_sem to 3 channels
+      self.in_conv = nn.Conv2d(in_dim, 3, 3, 1, 1)
+      self.convnext_tiny_backbone = convnext_tiny(pretrained=True)
+      self.num_tokens = (self.resize_input // self.down_factor) ** 2
+      self.convnext_feature_dim = 3072
+      self.pos_embedding = nn.Parameter(torch.empty(
+        1, self.num_tokens, self.convnext_feature_dim).normal_(std=0.02))  # from BERT
+
+    self.test_drop_boxes = test_drop_boxes
+    self.test_drop_points = test_drop_points
+    self.test_drop_scribbles = test_drop_scribbles
+    self.test_drop_masks = test_drop_masks
+    self.test_drop_segs = test_drop_masks
+
+    self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
+    self.fourier_embedder_polygons = FourierEmbedder(
+      num_freqs=fourier_freqs_polygons)
+    input_dim = self.in_dim
+    input_dim_list = []
+    if self.add_boxes:
+      # 2: sin and cos; 4: (x1,y1) and (x2,y2)
+      self.position_dim = fourier_freqs * 2 * 4
+      input_dim += self.position_dim
+      input_dim_list.append(self.in_dim + self.position_dim)
+    if self.add_points:
+      self.point_dim = fourier_freqs * 2 * 2  # 2: sin and cos; 2: (x,y)
+      input_dim += self.point_dim
+      input_dim_list.append(self.in_dim + self.point_dim)
+    if self.add_scribbles:
+      self.scribble_dim = fourier_freqs_polygons * 2 * \
+        self.n_scribble_points * 2  # 2: sin and cos; 2: (x,y)
+      input_dim += self.scribble_dim
+      input_dim_list.append(self.in_dim + self.scribble_dim)
+    if self.add_masks:
+      self.polygon_dim = fourier_freqs_polygons * 2 * \
+        self.n_polygon_points * 2  # 2: sin and cos; 2: (x,y)
+      input_dim += self.polygon_dim
+      input_dim_list.append(self.in_dim + self.polygon_dim)
+      if self.use_segs:
+        input_dim += self.convnext_feature_dim
+        input_dim_list.append(self.convnext_feature_dim)
+
+    if self.use_seperate_tokenizer:
+      self.linears_list = nn.ModuleList([])
+      for idx, input_dim_ in enumerate(input_dim_list):
+        mid_dim = self.mid_dim
+        self.linears_list.append(nn.Sequential(
+            nn.Linear(input_dim_, mid_dim),
+            nn.SiLU(),
+            nn.Linear(mid_dim, mid_dim),
+            nn.SiLU(),
+            nn.Linear(mid_dim, out_dim),
+        ))
+    else:
+      self.linears = nn.Sequential(
+          nn.Linear(input_dim, self.mid_dim),
+          nn.SiLU(),
+          nn.Linear(self.mid_dim, self.mid_dim),
+          nn.SiLU(),
+          nn.Linear(self.mid_dim, out_dim),
+      )
+
+    self.null_positive_feature = torch.nn.Parameter(
+      torch.zeros([self.in_dim]))  # text
+    if self.add_boxes:
+      self.null_position_feature = torch.nn.Parameter(
+        torch.zeros([self.position_dim]))
+    if self.add_points:
+      self.null_point_feature = torch.nn.Parameter(
+        torch.zeros([self.point_dim]))
+    if self.add_scribbles:
+      self.null_scribble_feature = torch.nn.Parameter(
+        torch.zeros([self.scribble_dim]))
+    if self.add_masks:
+      self.null_polygon_feature = torch.nn.Parameter(
+        torch.zeros([self.polygon_dim]))
+      if self.use_segs:
+        self.null_seg_feature = torch.nn.Parameter(
+          torch.zeros([self.convnext_feature_dim]))
+
+  def reset_dropout_test(self):
+    # drop_box = True
+    # drop_point = False
+    # drop_scribble = True
+    # drop_polygons = True
+    # drop_segs = True
+    drop_box = self.test_drop_boxes
+    drop_point = self.test_drop_points
+    drop_scribble = self.test_drop_scribbles
+    drop_polygons = self.test_drop_masks
+    drop_segs = self.test_drop_masks
+
+    return drop_point, drop_box, drop_scribble, drop_polygons, drop_segs
+
+  def reset_dropout(self):
+    drop_box = False
+    drop_point = False
+    drop_scribble = False
+    drop_polygons = False
+    drop_segs = False
+    return drop_point, drop_box, drop_scribble, drop_polygons, drop_segs
+
+  def reset_dropout_train(self, drop_point, drop_box, drop_scribble, drop_polygons, drop_segs):
+    if not drop_polygons:
+      drop_box = False
+      drop_point = False
+    if not drop_box or not drop_polygons:
+      drop_point = False
+
+    # keep point only for 10% of the time
+    keep_point_only_ratio = 0.1
+    keep_point_only = torch.rand(1).item() < keep_point_only_ratio
+    if keep_point_only:
+      drop_point = False
+      drop_box = True
+      drop_scribble = True
+      drop_polygons = True
+      drop_segs = True
+
+    # keep scribble only for 0% of the time
+    keep_scribble_only_ratio = 0.0
+    keep_scribble_only = torch.rand(
+      1).item() < keep_scribble_only_ratio and not drop_scribble
+    if keep_scribble_only:
+      drop_point = True
+      drop_box = True
+      drop_scribble = False
+      drop_polygons = True
+      drop_segs = True
+
+    # keep mask only for 0% of the time
+    keep_mask_only_ratio = 0.0
+    keep_mask_only = torch.rand(
+      1).item() < keep_mask_only_ratio and not drop_polygons
+    if keep_mask_only:
+      drop_point = True
+      drop_box = True
+      drop_scribble = True
+      drop_polygons = False
+      drop_segs = False
+
+    # keep seg only for 10% of the time
+    keep_seg_only_ratio = 0.1  # default 0.1
+    keep_seg_only = torch.rand(
+      1).item() < keep_seg_only_ratio and not drop_segs
+    if keep_seg_only:
+      drop_point = False
+      drop_box = False
+      drop_scribble = True
+      drop_polygons = False
+      drop_segs = False
+
+    # keep box only for 0% of the time
+    keep_box_only_ratio = 0.0  # default 0.0
+    keep_box_only = torch.rand(1).item() < keep_box_only_ratio and not drop_box
+    if keep_box_only:
+      drop_point = True
+      drop_box = False
+      drop_scribble = True
+      drop_polygons = True
+      drop_segs = True
+
+    return drop_point, drop_box, drop_scribble, drop_polygons, drop_segs
+
+  def forward(self, embeddings):
+    boxes = embeddings['boxes']
+    masks = embeddings['masks']
+    positive_embeddings = embeddings['prompts']
+    scribbles = embeddings['scribbles']
+    polygons = embeddings['polygons']
+    segs = embeddings['segments']
+    points = embeddings['points']
+
+    B, N, _ = boxes.shape
+    masks = masks.unsqueeze(-1)
+
+    drop_point, drop_box, drop_scribble, drop_polygons, drop_segs = self.reset_dropout()
+    # randomly drop boxes or points embeddings.
+    if self.add_boxes:
+      drop_box_ratio = 0.1
+      drop_box = torch.rand(1).item() < drop_box_ratio
+    if self.add_points:
+      drop_point_ratio = 0.1
+      drop_point = torch.rand(1).item() < drop_point_ratio
+    if self.add_scribbles:
+      drop_scribble_ratio = 0.1
+      drop_scribble = torch.rand(1).item() < drop_scribble_ratio
+    if self.add_masks:
+      drop_polygon_ratio = 0.1
+      drop_polygons = torch.rand(1).item() < drop_polygon_ratio
+      drop_segs = drop_polygons
+
+    # Not training, always keep both boxes and points
+    if not self.training:
+      drop_point, drop_box, drop_scribble, drop_polygons, drop_segs = self.reset_dropout_test()
+    else:
+      drop_point, drop_box, drop_scribble, drop_polygons, drop_segs = self.reset_dropout_train(
+        drop_point, drop_box, drop_scribble, drop_polygons, drop_segs)
+
+    # set drop_box to False if all other inputs are dropped
+    if drop_point and drop_box and drop_scribble and drop_polygons and drop_segs:
+      drop_box = False
+
+    # embedding position (it may includes padding as placeholder)
+    if self.add_boxes:
+      xyxy_embedding = self.fourier_embedder(
+        boxes)  # B*N*4 --> B*N*C (C=8*2*4)
+    if self.add_points:
+      if points is None:  # we can always get a point using a box
+        points = (boxes[:, :, :2] + boxes[:, :, 2:]) / 2.0
+      point_embedding = self.fourier_embedder(points)  # B*N*2 --> B*N*(8*2*2)
+    if self.add_scribbles:
+      scribble_embedding = self.fourier_embedder_polygons(
+        scribbles)  # B*N*20 --> B*N*(8*20*2)
+    if self.add_masks:
+      polygon_embedding = self.fourier_embedder_polygons(
+        polygons)  # B*N*128 --> B*N*(16*128*2)
+    if self.use_segs:
+      segs = torch.nn.functional.interpolate(
+        segs, self.resize_input, mode="nearest")
+      segs_feature = self.in_conv(segs)
+      segs_feature = self.convnext_tiny_backbone(segs_feature)
+      segs_feature = segs_feature.reshape(B, -1, self.num_tokens)
+      segs_feature = segs_feature.permute(0, 2, 1)
+
+    # learnable null embedding
+    positive_null = self.null_positive_feature.view(1, 1, -1)
+    if self.add_boxes:
+      xyxy_null = self.null_position_feature.view(1, 1, -1)
+    if self.add_points:
+      point_null = self.null_point_feature.view(1, 1, -1)
+    if self.add_scribbles:
+      scribble_null = self.null_scribble_feature.view(1, 1, -1)
+    if self.add_masks:
+      polygon_null = self.null_polygon_feature.view(1, 1, -1)
+    if self.use_segs:
+      seg_null = self.null_seg_feature.view(1, 1, -1)
+      seg_null = seg_null.repeat(B, self.num_tokens, 1)
+
+    # replace padding with learnable null embedding
+    positive_embeddings = positive_embeddings * \
+        masks + (1 - masks) * positive_null
+    if self.use_seperate_tokenizer:
+      embeddings_list = []
+    if self.add_boxes:
+      # replace padding with learnable null embedding for boxes
+      xyxy_masks = torch.zeros_like(masks).to(
+        masks.device) if drop_box else masks.detach().clone()
+      xyxy_embedding = xyxy_embedding * \
+          xyxy_masks + (1 - xyxy_masks) * xyxy_null
+      if self.use_seperate_tokenizer:
+        embeddings_list.append(xyxy_embedding)
+    if self.add_points:
+      # replace padding with learnable null embedding for points
+      point_masks = torch.zeros_like(masks).to(
+        boxes.device) if drop_point else masks.detach().clone()
+      point_embedding = point_embedding * \
+          point_masks + (1 - point_masks) * point_null
+      if self.use_seperate_tokenizer:
+        embeddings_list.append(point_embedding)
+    if self.add_scribbles:
+      # sum along the batch dimension and check if all scribbles are 0s
+      # replace padding with learnable null embedding for scribbles
+      # scribble_embedding: torch.Size([bs, n_objs, 640]); masks_scribble: torch.Size([bs, n_objs, 1]); scribble_null: torch.Size([1, 1, 640])
+      masks_scribble = torch.zeros_like(masks).to(masks.device) if drop_scribble else (
+        (torch.sum(scribbles, dim=-1).unsqueeze(-1) + masks.detach().clone()) > 0).float()
+      scribble_embedding = scribble_embedding * \
+          masks_scribble + (1 - masks_scribble) * scribble_null
+      if self.use_seperate_tokenizer:
+        embeddings_list.append(scribble_embedding)
+    if self.add_masks:
+      masks_polygons = torch.zeros_like(masks).to(masks.device) if drop_polygons else (
+        (torch.sum(polygons, dim=-1).unsqueeze(-1) + masks.detach().clone()) > 0).float()
+      assert torch.sum(scribbles, dim=-1).unsqueeze(-1).size() == masks.size()
+      polygon_embedding = polygon_embedding * \
+          masks_polygons + (1 - masks_polygons) * polygon_null
+      if self.use_seperate_tokenizer:
+        embeddings_list.append(polygon_embedding)
+    if self.use_segs:
+      # mask replacing
+      masks_segs = torch.zeros(masks.shape[0]).to(masks.device) if drop_segs else (
+        torch.sum(segs, dim=(1, 2, 3)) > 0).float()
+      masks_segs = masks_segs.view(-1, 1, 1)
+      assert masks_segs.size()[0] == masks.shape[0]
+      seg_embedding = segs_feature * masks_segs
+      seg_embedding = seg_embedding + (1 - masks_segs) * seg_null
+      # add pos
+      seg_embedding = seg_embedding + self.pos_embedding
+      if self.use_seperate_tokenizer:
+        embeddings_list.append(seg_embedding)
+
+    inputs = [positive_embeddings]
+    if self.use_seperate_tokenizer:
+      objs = []
+      # forward all types of embeddings using the corresponding linear layers
+      for i, (linears, layout_embeddings) in enumerate(zip(self.linears_list, embeddings_list)):
+        if i == len(embeddings_list) - 1 and self.use_segs:
+          objs.append(linears(layout_embeddings))
+        else:
+          objs.append(
+            linears(torch.cat([positive_embeddings, layout_embeddings], dim=-1)))
+      objs = torch.cat(objs, dim=1)
+    else:
+      # NOTE: orders should the same for training and testing
+      if self.add_boxes:
+        inputs.append(xyxy_embedding)
+      if self.add_points:
+        inputs.append(point_embedding)
+      if self.add_scribbles:
+        inputs.append(scribble_embedding)
+      if self.add_masks:
+        inputs.append(polygon_embedding)
+
+      objs = self.linears(torch.cat(inputs, dim=-1))
+      assert objs.shape == torch.Size([B, N, self.out_dim])
+    drop_box_mask = True if drop_box and drop_polygons else False
+    return objs, drop_box_mask

modules/text_grounding_tokenizer_input.py

@@ -0,0 +1,96 @@
+import torch
+
+
+class GroundingNetInput:
+  def __init__(self):
+    self.set = False
+    self.return_att_masks = False
+    self.image_size = 64
+    self.return_att_masks32 = False
+
+  def prepare(self, batch, image_size=64, device=None, dtype=None, return_att_masks=False):
+    """
+    batch should be the output from dataset.
+    Please define here how to process the batch and prepare the 
+    input only for the ground tokenizer. 
+    """
+    output = {}
+    self.set = True
+    self.return_att_masks = return_att_masks
+
+    boxes = batch['boxes']
+    masks = batch['masks']
+    positive_embeddings = batch["prompts"]
+
+    if self.return_att_masks:
+      assert 'att_masks' in batch
+      att_masks = batch['att_masks']
+
+    scribbles = batch['scribbles']
+    polygons = batch['polygons']
+    self.dim_scribbles = scribbles.shape[-1]
+    self.dim_polygons = polygons.shape[-1]
+    # NOTE: New Seg
+    segs = batch["segments"]
+    self.dim_segs = segs.shape[-1]
+    points = batch["points"]
+
+    self.batch, self.max_box, self.in_dim = positive_embeddings.shape
+    self.device = positive_embeddings.device
+    self.dtype = positive_embeddings.dtype
+
+    output = {
+        "boxes": boxes,
+        "masks": masks,
+        "prompts": positive_embeddings,
+    }
+    output["scribbles"] = scribbles
+    output["polygons"] = polygons
+    output["segments"] = segs
+    output["points"] = points
+
+    if self.return_att_masks:
+      output['att_masks'] = att_masks
+    return output
+
+  def get_null_input(self, batch=None, latent_width=64, latent_height=64, device=None, dtype=None):
+    """
+    Guidance for training (drop) or inference, 
+    please define the null input for the grounding tokenizer 
+    """
+
+    assert self.set, "not set yet, cannot call this funcion"
+    batch = self.batch if batch is None else batch
+    device = self.device if device is None else device
+    dtype = self.dtype if dtype is None else dtype
+
+    boxes = torch.zeros(batch, self.max_box, 4,).type(dtype).to(device)
+    masks = torch.zeros(batch, self.max_box).type(dtype).to(device)
+    # NOTE: New Seg
+    segs = torch.zeros(batch, self.max_box, self.dim_segs,
+                       self.dim_segs).type(dtype).to(device)
+
+    scribbles = torch.zeros(batch, self.max_box,
+                            self.dim_scribbles).type(dtype).to(device)
+    polygons = torch.zeros(batch, self.max_box,
+                           self.dim_polygons).type(dtype).to(device)
+    points = torch.zeros(batch, self.max_box, 2).type(dtype).to(device)
+
+    positive_embeddings = torch.zeros(
+      batch, self.max_box, self.in_dim).type(dtype).to(device)
+
+    output = {
+        "boxes": boxes,
+        "masks": masks,
+        "prompts": positive_embeddings,
+    }
+    output["scribbles"] = scribbles
+    output["polygons"] = polygons
+    output["segments"] = segs
+    output["points"] = points
+
+    if self.return_att_masks:
+      att_masks = torch.zeros(batch, self.max_box, latent_width, latent_height).type(
+        dtype).to(device)  # TODO Order width/height
+      output['att_masks'] = att_masks
+    return output

modules/util.py

@@ -0,0 +1,19 @@
+import torch
+
+
+class FourierEmbedder():
+    def __init__(self, num_freqs=64, temperature=100):
+
+        self.num_freqs = num_freqs
+        self.temperature = temperature
+        self.freq_bands = temperature ** ( torch.arange(num_freqs) / num_freqs )  
+
+    @ torch.no_grad()
+    def __call__(self, x, cat_dim=-1):
+        "x: arbitrary shape of tensor. dim: cat dim"
+        out = []
+        for freq in self.freq_bands:
+            out.append( torch.sin( freq*x ) )
+            out.append( torch.cos( freq*x ) )
+        return torch.cat(out, cat_dim)
+

nodes/apply_scaleu_model_node.py

@@ -0,0 +1,30 @@
+from .. import constants as constants
+from ..models.prepare_scaleu import get_scaleu_patch
+
+
+class ApplyScaleUModel:
+  @classmethod
+  def INPUT_TYPES(s):
+    return {"required": {
+        "model": ("MODEL",),
+        "scaleu": ("SCALEU",),
+    }}
+
+  RETURN_TYPES = ("MODEL",)
+  FUNCTION = "apply"
+
+  CATEGORY = "instance"
+
+  def apply(self, model, scaleu):
+    # Validate patches dict is setup correctly
+    if 'patches' not in model or model['patches'] is None:
+      model['patches'] = {}
+
+    if 'output_block_patch' not in model['patches']:
+      model['patches']['output_block_patch'] = []
+
+    # Add scaleu patch to model patches
+    scaleu_nets = scaleu['model_list']
+    model['patches']['output_block_patch'].append(
+      get_scaleu_patch(scaleu_nets))
+    return (model,)

nodes/instance_diffusion_tracking_prompt_node.py

@@ -0,0 +1,61 @@
+from ..utils.prompt_utils import extract_prompts
+from ..conditioning.instance_conditioning import InstanceConditioning
+
+
+class InstanceDiffusionTrackingPromptNode:
+  @classmethod
+  def INPUT_TYPES(s):
+    return {"required": {"conditioning": ("CONDITIONING", ),
+                         "clip": ("CLIP", ),
+                         "tracking": ("TRACKING", ),
+                         "positionnet": ("POSITIONNET", ),
+                         "fusers": ("FUSERS", ),
+                         "fusers_batch_size": ("INT", {"default": 5, "min": 0, "max": 10000, "step": 8}),
+                         "text": ("STRING", {"multiline": True}),
+                         }}
+  RETURN_TYPES = ("CONDITIONING",)
+  FUNCTION = "append"
+
+  CATEGORY = "instance/conditioning"
+
+  def append(self, conditioning, clip, tracking, fusers, positionnet, fusers_batch_size, text):
+    # Get prompts and their class id and trakcer id
+    prompt_pairs = extract_prompts(text)
+
+    # Go through prompt pairs, encode prompts, and join with positions from tracking
+    position_conds = []
+    for tracker_id, class_id, prompt in prompt_pairs:
+      _, cond_pooled = clip.encode_from_tokens(
+        clip.tokenize(prompt), return_pooled=True)
+      # A tracker_id of -1 means that it is prompting all instances of a single class
+      if tracker_id != -1:
+        position_cond = {'cond_pooled': cond_pooled, 'positions':
+                         tracking[class_id][tracker_id]}
+        position_conds.append(position_cond)
+      else:
+        for tracker_id in tracking[class_id]:
+          position_cond = {'cond_pooled': cond_pooled,
+                           'positions': tracking[class_id][tracker_id]}
+          position_conds.append(position_cond)
+
+    # Add prompts+embeddings to the input conditionings
+    cond_out = []
+    for t in conditioning:
+      n = [t[0], t[1].copy()]
+      cond = n[1]
+      prev = []
+      has_instance = 'instance_diffusion' in cond
+      instance_conditioning = conditioning['instance_diffusion'] if has_instance else InstanceConditioning(
+        fusers, positionnet, fusers_batch_size)
+      cond['instance_diffusion'] = instance_conditioning
+      instance_conditioning.add_conds(position_conds)
+
+      cond['instance_diffusion'] = {
+        'conds': prev + position_conds,
+      }
+
+      cond['gligen'] = ('position', instance_conditioning, None)
+
+      cond_out.append(n)
+
+    return (cond_out,)

nodes/load_instance_fusers_node.py

@@ -0,0 +1,24 @@
+from .. import constants as constants
+from ..utils.model_utils import get_model_dir, load_checkpoint
+from ..models.prepare_fusers import prepare_fusers
+
+
+class LoadInstanceFusersNode:
+  @classmethod
+  def INPUT_TYPES(s):
+    return {"required": {
+        "model_path": (get_model_dir(constants.INSTANCE_FUSERS_DIR)),
+    }}
+
+  RETURN_TYPES = ("FUSERS",)
+  FUNCTION = "load_model"
+
+  CATEGORY = "instance/loaders"
+
+  def load_model(self, model_path: str):
+    checkpoint = load_checkpoint(model_path)
+    fusers_list = prepare_fusers(checkpoint)
+    fusers = {
+      'model_list': fusers_list
+    }
+    return (fusers,)

nodes/load_instance_positionnet_node.py

@@ -0,0 +1,25 @@
+from .. import constants as constants
+from ..utils.model_utils import get_model_dir, load_checkpoint
+from ..models.prepare_positionnet import prepare_positionnet, get_positionnet_default_params
+
+
+class LoadInstancePositionNetNode:
+  @classmethod
+  def INPUT_TYPES(s):
+    return {"required": {
+        "model_path": (get_model_dir(constants.INSTANCE_POSITIONNET_DIR)),
+    }}
+
+  RETURN_TYPES = ("POSITIONNET",)
+  FUNCTION = "load_model"
+
+  CATEGORY = "instance/loaders"
+
+  def load_model(self, model_path: str):
+    checkpoint = load_checkpoint(model_path)
+    params = get_positionnet_default_params()
+    model = prepare_positionnet(checkpoint, params)
+    positionnet = {
+      'model': model,
+    }
+    return (positionnet,)

nodes/load_instance_scaleu_node.py

@@ -0,0 +1,24 @@
+from .. import constants as constants
+from ..utils.model_utils import get_model_dir, load_checkpoint
+from ..models.prepare_scaleu import prepare_scaleu_nets
+
+
+class LoadInstanceScaleUNode:
+  @classmethod
+  def INPUT_TYPES(s):
+    return {"required": {
+        "model_path": (get_model_dir(constants.INSTANCE_SCALEU_DIR)),
+    }}
+
+  RETURN_TYPES = ("SCALEU",)
+  FUNCTION = "load_model"
+
+  CATEGORY = "instance/loaders"
+
+  def load_model(self, model_path: str):
+    checkpoint = load_checkpoint(model_path)
+    scaleu_list = prepare_scaleu_nets(checkpoint)
+    scaleu = {
+      'model_list': scaleu_list
+    }
+    return (scaleu,)

requirements.txt

@@ -0,0 +1 @@
+torch

utils/decode_item.py

@@ -0,0 +1,422 @@
+# Directly taken from InstanceDiffusion repo
+import torch
+import random
+import base64
+import numpy as np
+from io import BytesIO
+from collections import Counter
+from PIL import Image, ImageDraw
+import base64
+from skimage import measure
+
+
+# import nltk
+# from nltk.corpus import stopwords
+
+def decode_base64_to_pillow(image_b64):
+  return Image.open(BytesIO(base64.b64decode(image_b64))).convert('RGB')
+
+
+def decode_tensor_from_string(arr_str, use_tensor=True):
+  arr = np.frombuffer(base64.b64decode(arr_str), dtype='float32')
+  if use_tensor:
+    arr = torch.from_numpy(arr)
+  return arr
+
+
+def close_contour(contour):
+  if not np.array_equal(contour[0], contour[-1]):
+    contour = np.vstack((contour, contour[0]))
+  return contour
+
+# convert binay mask to polygon format
+
+
+def binary_mask_to_polygon(binary_mask, tolerance=0):
+  """Converts a binary mask to COCO polygon representation
+
+  Args:
+      binary_mask: a 2D binary numpy array where '1's represent the object
+      tolerance: Maximum distance from original points of polygon to approximated
+          polygonal chain. If tolerance is 0, the original coordinate array is returned.
+
+  """
+  polygons = []
+  # pad mask to close contours of shapes which start and end at an edge
+  padded_binary_mask = np.pad(
+    binary_mask, pad_width=1, mode='constant', constant_values=0)
+  contours = measure.find_contours(padded_binary_mask, tolerance)
+  polygons = []
+  # print(contours)
+  for contour in contours:
+    contour = close_contour(contour)
+    contour = measure.approximate_polygon(contour, tolerance)
+    if len(contour) < 3:
+      continue
+    contour = np.flip(contour, axis=1)
+    segmentation = contour.ravel().tolist()
+    # after padding and subtracting 1 we may get -0.5 points in our segmentation
+    segmentation = [0 if i < 0 else i for i in segmentation]
+    polygons.append(segmentation)
+
+  return polygons
+
+
+def sample_random_points_from_mask(mask, k):
+  mask = mask[:, :, 0]
+  # Find the coordinates of non-zero pixels in the binary mask
+  nonzero_coords = np.transpose(np.nonzero(mask))
+
+  # Randomly sample 'k' points
+  # return all zeros if there is no non-zero pixel
+  if len(nonzero_coords) == 0:
+    xy_points = [0 for _ in range(k * 2)]
+    return xy_points
+
+  # randomly sample with replacement if there are not enough non-zero pixels
+  if len(nonzero_coords) < k and len(nonzero_coords) > 0:
+    random_indices = np.random.choice(len(nonzero_coords), k, replace=True)
+  # randomly sample withiout replacement if there are enough non-zero pixels
+  else:
+    random_indices = np.random.choice(len(nonzero_coords), k, replace=False)
+  sampled_points = nonzero_coords[random_indices]
+
+  # order the points by their distance to (0, 0)
+  # center = np.array([mask.shape[0] // 2, mask.shape[1] // 2])
+  center = np.array([0, 0])
+  sampled_points = sorted(sampled_points, key=lambda x: np.linalg.norm(
+    np.array(x) - center))  # np.linalg.norm
+
+  # concatenate x and y coordinates and return them as a list
+  # [x1,y1,x2,y2,...,x_k,y_k]
+  xy_points = []
+  for x in sampled_points:
+    xy_points.append(float(x[1]))
+    xy_points.append(float(x[0]))
+  return xy_points
+
+# convert numpy array of bool mask to float mask
+
+
+def binary_mask_to_int(binary_mask):
+  return binary_mask.astype(np.int32)
+
+# uniformly sample points from the mask
+
+
+def sample_sparse_points(binary_mask, k, return_2d=False):
+  # Find the coordinates of non-zero pixels in the binary mask
+  nonzero_coords = np.array(np.nonzero(binary_mask))
+  if len(nonzero_coords) == 0:
+    xy_points = [0 for _ in range(k * 2)]
+    return xy_points
+
+  # Calculate the total number of non-zero pixels
+  num_nonzero_pixels = len(nonzero_coords)
+
+  xy_points = []
+  if k >= num_nonzero_pixels:
+    for x in nonzero_coords:
+      xy_points.append(float(x[1]))
+      xy_points.append(float(x[0]))
+    for _ in range(k - num_nonzero_pixels):
+      xy_points.append(nonzero_coords[-1][1])
+      xy_points.append(nonzero_coords[-1][0])
+    return nonzero_coords
+
+  # Calculate the number of points to sample in each dimension
+  num_points_per_dim = int(np.sqrt(k))
+
+  # Calculate the step size to ensure equal spacing
+  step_size = max(1, num_nonzero_pixels // (num_points_per_dim ** 2))
+
+  # Sample points with equal spacing
+  sampled_points = nonzero_coords[::step_size][:k]
+  if return_2d:
+    sampled_points = [(x[1], x[0]) for x in sampled_points]
+  else:
+    for x in sampled_points:
+      xy_points.append(float(x[1]))
+      xy_points.append(float(x[0]))
+    return xy_points
+
+
+def sample_uniform_sparse_points(binary_mask, k):
+  # binary_mask = binary_mask[:,:,0]
+  # Step 1: Get the indices of '1' values in the binary mask
+  foreground_indices = np.argwhere(binary_mask == 1)
+
+  if len(foreground_indices) == 0:
+    return []
+
+  selected_points = []
+  if len(foreground_indices) < k:
+    # randomly sample with replacement if there are not enough non-zero pixels
+    for i in range(k):
+      random_point = random.choice(foreground_indices)
+      selected_points.append((random_point[1], random_point[0]))
+  else:
+    # rank the points by their distance to the mean of the foreground_indices
+    center = np.mean(foreground_indices, axis=0)
+    # print(center)
+    foreground_indices = sorted(
+      foreground_indices, key=lambda x: np.linalg.norm(x - center))  # np.linalg.norm
+    # Calculate the number of points to select from each segment
+    points_per_segment = len(foreground_indices) // k
+
+    # Step 2: Randomly select one point from each segment
+    # print(k)
+    for i in range(k):
+      segment_points = foreground_indices[i *
+                                          points_per_segment: (i + 1) * points_per_segment]
+      # choose the middle point in each segment
+      random_point = segment_points[len(segment_points) // 2]
+      # random_point = random.choice(segment_points)
+      selected_points.append((random_point[1], random_point[0]))
+
+  return selected_points
+
+
+def sample_sparse_points_from_mask(mask, k):
+  n_points = k
+  n_polygons = n_points // 2  # half points should be sampled from the polygons
+  mask = mask[:, :, 0]
+  # sample sparse points from the polygons (boundary)
+  polygons = binary_mask_to_polygon(mask, tolerance=0.0)
+  # concatenate polygons to a single list
+  polygons_single = []
+  for polygon in polygons:
+    polygons_single += polygon
+  if len(polygons_single) != 0:
+    # uniformly sample points from the polygon
+    polygons_single = np.array(polygons_single).reshape(-1, 2)
+    indexes = np.linspace(0, polygons_single.shape[0] - 1, n_polygons)
+    indexes = list([int(i) for i in indexes])
+
+    polygons_single = polygons_single[indexes]
+    sampled_polygons = [(x[0], x[1]) for x in polygons_single]
+  else:
+    return None
+
+  # sample sparse points from the mask
+  n_inside_points = n_points - len(sampled_polygons)
+  inside_points = sample_uniform_sparse_points(mask, n_inside_points)
+
+  # combine inside_points and sampled_polygons
+  xy_points = inside_points + sampled_polygons
+
+  # order the points by their distance to (0, 0)
+  center = np.array([0, 0])
+  xy_points = sorted(xy_points, key=lambda x: np.linalg.norm(
+    np.array(x) - center))  # np.linalg.norm
+
+  # return the sampled points
+  sampled_points = []
+  for x in xy_points:
+    sampled_points.append(x[0])
+    sampled_points.append(x[1])
+  return sampled_points
+
+
+def get_polygons_from_mask(mask, tolerance=0, n_polygon_points=256):
+  mask = binary_mask_to_int(mask)
+  return_polygons = True
+  if return_polygons:
+    # convert float mask to polygons
+    polygons = binary_mask_to_polygon(mask[:, :, 0], tolerance=tolerance)
+
+    # return all zeros if there is no polygon
+    if len(polygons) == 0:
+      polygons = [0 for _ in range(n_polygon_points * 2)]
+      return polygons
+
+    # concatenate polygons to a single list
+    polygon = []
+    for p in polygons:
+      polygon += p
+
+    # uniformly sample points the polygon
+    polygon = np.array(polygon).reshape(-1, 2)
+    indexes = np.linspace(0, polygon.shape[0] - 1, n_polygon_points)
+    indexes = [int(i) for i in indexes]
+    polygon = polygon[indexes].reshape(-1)
+
+    return polygon
+  else:
+    sampled_points = sample_sparse_points(mask, n_polygon_points)
+    return sampled_points
+
+
+def decode_item(item):
+  # convert string to dict
+  if "image" in item and isinstance(item['image'], Image.Image):
+    return item
+
+  item['image'] = decode_base64_to_pillow(item['image'])
+  segs = []
+  for anno in item['annos']:
+    anno['image_embedding_before'] = decode_tensor_from_string(
+      anno['image_embedding_before'])
+    anno['text_embedding_before'] = decode_tensor_from_string(
+      anno['text_embedding_before'])
+    anno['image_embedding_after'] = decode_tensor_from_string(
+      anno['image_embedding_after'])
+    anno['text_embedding_after'] = decode_tensor_from_string(
+      anno['text_embedding_after'])
+    if "blip_clip_embeddings" in anno:
+      anno['blip_clip_embeddings'] = decode_tensor_from_string(
+        anno['blip_clip_embeddings'])
+    if 'mask' in anno:
+      # sample k random points from the mask
+      n_scribble_points = 20
+      rle = anno['mask']
+      binary_mask = decodeToBinaryMask(rle)
+      segs.append(binary_mask)
+      if "scribbles" in anno:
+        anno['scribbles'] = anno["scribbles"]
+      else:
+        anno['scribbles'] = sample_random_points_from_mask(
+          binary_mask, n_scribble_points)
+      # convert mask to polygon
+      n_polygon_points = 256
+      polygons = sample_sparse_points_from_mask(
+        binary_mask, k=n_polygon_points)
+      if polygons != None:
+        anno['polygons'] = polygons
+      else:
+        anno['polygons'] = [0 for _ in range(n_polygon_points * 2)]
+  if len(segs) > 0:
+    item['segs'] = np.stack(segs).astype(np.float32).squeeze()
+  return item
+
+
+def check_unique(images, fields):
+  for field in fields:
+    temp_list = []
+    for img_info in images:
+      temp_list.append(img_info[field])
+    assert len(set(temp_list)) == len(temp_list), field
+
+
+def clean_data(data):
+  for data_info in data:
+    data_info.pop("original_img_id", None)
+    data_info.pop("original_id", None)
+    # sentence id for each image (multiple sentences for one image)
+    data_info.pop("sentence_id", None)
+    data_info.pop("dataset_name", None)
+    data_info.pop("data_source", None)
+    data_info["data_id"] = data_info.pop("id")
+
+
+def clean_annotations(annotations):
+  for anno_info in annotations:
+    anno_info.pop("iscrowd", None)
+    anno_info.pop("category_id", None)
+    anno_info.pop("area", None)
+    anno_info["data_id"] = anno_info.pop("image_id")
+
+
+def draw_box(img, boxes):
+  draw = ImageDraw.Draw(img)
+  for box in boxes:
+    draw.rectangle([box[0], box[1], box[2], box[3]],
+                   outline="red", width=2)  # x0 y0 x1 y1
+  return img
+
+
+def xyhw2xyxy(box):
+  x0, y0, w, h = box
+  return [x0, y0, x0 + w, y0 + h]
+
+
+def make_a_sentence_count_nums(obj_names):
+  # count the number of duplicated strings in the list
+  # ["dog", "dog", "cat"]
+  obj_names = dict(Counter(obj_names))
+  # {'dog': 2, 'cat': 1}
+  caption = ""
+  for item in obj_names:
+    caption += str(obj_names[item]) + " " + item + ", "
+  return caption[:-2]
+
+
+def make_a_sentence(obj_names, clean=False):
+
+  if clean:
+    obj_names = [name[:-6] if ("-other" in name)
+                 else name for name in obj_names]
+
+  caption = ""
+  tokens_positive = []
+  for obj_name in obj_names:
+    start_len = len(caption)
+    caption += obj_name
+    end_len = len(caption)
+    caption += ", "
+    tokens_positive.append(
+        # in real caption, positive tokens can be disjoint, thus using list of list
+        [[start_len, end_len]]
+    )
+  caption = caption[:-2]  # remove last ", "
+
+  return caption  # , tokens_positive
+
+
+def mask_for_random_drop_text_or_image_feature(masks, random_drop_embedding):
+  """
+  input masks tell how many valid grounding tokens for this image
+  e.g., 1,1,1,1,0,0,0,0,0,0...
+
+  If random_drop_embedding=both.  we will random drop either image or
+  text feature for each token, 
+  but we always make sure there is at least one feature used. 
+  In other words, the following masks are not valid 
+  (because for the second obj, no feature at all):
+  image: 1,0,1,1,0,0,0,0,0
+  text:  1,0,0,0,0,0,0,0,0
+
+  if random_drop_embedding=image. we will random drop image feature 
+  and always keep the text one.  
+
+  """
+  N = masks.shape[0]
+
+  if random_drop_embedding == 'both':
+    temp_mask = torch.ones(2, N)
+    for i in range(N):
+      if random.uniform(0, 1) < 0.5:  # else keep both features
+        # randomly choose to drop image or text feature
+        idx = random.sample([0, 1], 1)[0]
+        temp_mask[idx, i] = 0
+    image_masks = temp_mask[0] * masks
+    text_masks = temp_mask[1] * masks
+
+  if random_drop_embedding == 'image':
+    image_masks = masks * (torch.rand(N) > 0.5) * 1
+    text_masks = masks
+
+  return image_masks, text_masks
+
+
+def project(x, projection_matrix):
+  """
+  x (Batch*768) should be the penultimate feature of CLIP (before projection)
+  projection_matrix (768*768) is the CLIP projection matrix, which should be weight.data of Linear layer 
+  defined in CLIP (out_dim, in_dim), thus we need to apply transpose below.  
+  this function will return the CLIP feature (without normalziation)
+  """
+  return x @ torch.transpose(projection_matrix, 0, 1)
+
+
+def inv_project(y, projection_matrix):
+  """
+  y (Batch*768) should be the CLIP feature (after projection)
+  projection_matrix (768*768) is the CLIP projection matrix, which should be weight.data of Linear layer 
+  defined in CLIP (out_dim, in_dim).  
+  this function will return the CLIP penultimate feature. 
+
+  Note: to make sure getting the correct penultimate feature, the input y should not be normalized. 
+  If it is normalized, then the result will be scaled by CLIP feature norm, which is unknown.   
+  """
+  return y @ torch.transpose(torch.linalg.inv(projection_matrix), 0, 1)

utils/model_utils.py

@@ -0,0 +1,18 @@
+import os
+import folder_paths
+import torch
+
+import comfy.utils
+
+from . import constants as constants
+
+
+def get_model_dir(model_dir) -> list[str]:
+  path = os.path.join(constants.INSTANCE_MODELS_DIR,
+                      model_dir)
+  return folder_paths.get_filename_list(path)
+
+
+def load_checkpoint(checkpoint_path):
+  checkpoint = comfy.utils.load_torch_file(checkpoint_path, safe_load=True)
+  return checkpoint

utils/prompt_utils.py

@@ -0,0 +1,15 @@
+import re
+
+
+def extract_prompts(input_string):
+  # Define the new regex pattern to include negative numbers
+  pattern = r"\"(-?\d+)\.([^\"]+)\":\s*\"([^\"]+)\""
+
+  # Find all matches using the pattern
+  matches = re.findall(pattern, input_string)
+
+  # Convert matches to a list of tuples (number, name, text)
+  result = [(int(number), name.strip(), text)
+            for number, name, text in matches]
+
+  return result