From 047e491505bc7d7310776a4f5db26b7522e25b81 Mon Sep 17 00:00:00 2001 From: logtd Date: Thu, 22 Feb 2024 21:46:05 -0700 Subject: [PATCH] init --- .gitignore | 3 + LICENSE | 201 +++++++++ README.md | 2 + __init__.py | 17 + conditioning/embeddings.py | 203 +++++++++ conditioning/fusers_patch.py | 0 conditioning/instance_conditioning.py | 63 +++ constants.py | 5 + models/prepare_fusers.py | 30 ++ models/prepare_positionnet.py | 26 ++ models/prepare_scaleu.py | 22 + modules/attention.py | 209 +++++++++ modules/convnext.py | 217 +++++++++ modules/scaleu.py | 57 +++ modules/text_grounding_net.py | 351 +++++++++++++++ modules/text_grounding_tokenizer_input.py | 96 ++++ modules/util.py | 19 + nodes/apply_scaleu_model_node.py | 30 ++ ...instance_diffusion_tracking_prompt_node.py | 61 +++ nodes/load_instance_fusers_node.py | 24 + nodes/load_instance_positionnet_node.py | 25 ++ nodes/load_instance_scaleu_node.py | 24 + requirements.txt | 1 + utils/decode_item.py | 422 ++++++++++++++++++ utils/model_utils.py | 18 + utils/prompt_utils.py | 15 + 26 files changed, 2141 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 README.md create mode 100644 __init__.py create mode 100644 conditioning/embeddings.py create mode 100644 conditioning/fusers_patch.py create mode 100644 conditioning/instance_conditioning.py create mode 100644 constants.py create mode 100644 models/prepare_fusers.py create mode 100644 models/prepare_positionnet.py create mode 100644 models/prepare_scaleu.py create mode 100644 modules/attention.py create mode 100644 modules/convnext.py create mode 100644 modules/scaleu.py create mode 100644 modules/text_grounding_net.py create mode 100644 modules/text_grounding_tokenizer_input.py create mode 100644 modules/util.py create mode 100644 nodes/apply_scaleu_model_node.py create mode 100644 nodes/instance_diffusion_tracking_prompt_node.py create mode 100644 nodes/load_instance_fusers_node.py create mode 100644 nodes/load_instance_positionnet_node.py create mode 100644 nodes/load_instance_scaleu_node.py create mode 100644 requirements.txt create mode 100644 utils/decode_item.py create mode 100644 utils/model_utils.py create mode 100644 utils/prompt_utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bbbe7d6 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +.git +.vscode +**/__pycache__ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..2ea9a9a --- /dev/null +++ b/README.md @@ -0,0 +1,2 @@ +# ComfyUI-InstanceDiffusion +ComfyUI nodes to use InstanceDiffusion diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..7a4e73e --- /dev/null +++ b/__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" +} diff --git a/conditioning/embeddings.py b/conditioning/embeddings.py new file mode 100644 index 0000000..aa0373d --- /dev/null +++ b/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 diff --git a/conditioning/fusers_patch.py b/conditioning/fusers_patch.py new file mode 100644 index 0000000..e69de29 diff --git a/conditioning/instance_conditioning.py b/conditioning/instance_conditioning.py new file mode 100644 index 0000000..13eb977 --- /dev/null +++ b/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) diff --git a/constants.py b/constants.py new file mode 100644 index 0000000..fd5ac12 --- /dev/null +++ b/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" diff --git a/models/prepare_fusers.py b/models/prepare_fusers.py new file mode 100644 index 0000000..b3bb20b --- /dev/null +++ b/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 diff --git a/models/prepare_positionnet.py b/models/prepare_positionnet.py new file mode 100644 index 0000000..41f3bf2 --- /dev/null +++ b/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 diff --git a/models/prepare_scaleu.py b/models/prepare_scaleu.py new file mode 100644 index 0000000..96ed207 --- /dev/null +++ b/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 diff --git a/modules/attention.py b/modules/attention.py new file mode 100644 index 0000000..3311e06 --- /dev/null +++ b/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) diff --git a/modules/convnext.py b/modules/convnext.py new file mode 100644 index 0000000..480f1a0 --- /dev/null +++ b/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 diff --git a/modules/scaleu.py b/modules/scaleu.py new file mode 100644 index 0000000..99da6d1 --- /dev/null +++ b/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) diff --git a/modules/text_grounding_net.py b/modules/text_grounding_net.py new file mode 100644 index 0000000..07b119a --- /dev/null +++ b/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 diff --git a/modules/text_grounding_tokenizer_input.py b/modules/text_grounding_tokenizer_input.py new file mode 100644 index 0000000..33dda7d --- /dev/null +++ b/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 diff --git a/modules/util.py b/modules/util.py new file mode 100644 index 0000000..05ec855 --- /dev/null +++ b/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) + diff --git a/nodes/apply_scaleu_model_node.py b/nodes/apply_scaleu_model_node.py new file mode 100644 index 0000000..7d9f26e --- /dev/null +++ b/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,) diff --git a/nodes/instance_diffusion_tracking_prompt_node.py b/nodes/instance_diffusion_tracking_prompt_node.py new file mode 100644 index 0000000..030d02e --- /dev/null +++ b/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,) diff --git a/nodes/load_instance_fusers_node.py b/nodes/load_instance_fusers_node.py new file mode 100644 index 0000000..9386800 --- /dev/null +++ b/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,) diff --git a/nodes/load_instance_positionnet_node.py b/nodes/load_instance_positionnet_node.py new file mode 100644 index 0000000..d22a9b9 --- /dev/null +++ b/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,) diff --git a/nodes/load_instance_scaleu_node.py b/nodes/load_instance_scaleu_node.py new file mode 100644 index 0000000..fe0bdff --- /dev/null +++ b/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,) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..08ed5ee --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +torch \ No newline at end of file diff --git a/utils/decode_item.py b/utils/decode_item.py new file mode 100644 index 0000000..f57f192 --- /dev/null +++ b/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) diff --git a/utils/model_utils.py b/utils/model_utils.py new file mode 100644 index 0000000..69138a5 --- /dev/null +++ b/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 diff --git a/utils/prompt_utils.py b/utils/prompt_utils.py new file mode 100644 index 0000000..195cb9a --- /dev/null +++ b/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