diff --git a/.gitignore b/.gitignore index dd0c0c3..7c0ed73 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ .git .vscode **/__pycache__ -**/*.ckpt \ No newline at end of file +**/*.ckpt +.DS_Store diff --git a/conditioning/fusers_patch.py b/conditioning/fusers_patch.py index 2bc8297..a58f023 100644 --- a/conditioning/fusers_patch.py +++ b/conditioning/fusers_patch.py @@ -1,21 +1,64 @@ import torch +from .embeddings import prepare_embeddings + + +block_map = { + 'input': { + 1: 0, + 2: 1, + 4: 2, + 5: 3, + 7: 4, + 8: 5 + }, + 'middle': { + 0: 6, + }, + 'output': { + 3: 7, + 4: 8, + 5: 9, + 6: 10, + 7: 11, + 8: 12, + 9: 13, + 10: 14, + 11: 15 + } +} + class FusersPatch(torch.nn.Module): - def __init__(self, fusers_list, objs, drop_box_mask, device, fusers_batch_size=5): - super(FusersPatch, self).__init__() - self.fusers_list = fusers_list - self.transformer_options = {'objs': objs, 'drop_box_mask': drop_box_mask} - self.device = device - self.fusers_batch_size = fusers_batch_size - # call_idx traces where in the UNet this is being called from - # a hack around lack of Comfy support - self.call_idx = 0 + def __init__(self, conds, fusers_list, fusers_batch_size, positionnet, latent_shape, idxs, device): + super(FusersPatch, self).__init__() + self.conds = conds + self.fusers_list = fusers_list + self.fusers_batch_size = fusers_batch_size + self.positionnet = positionnet + self.latent_shape = latent_shape + self.idxs = idxs + self.device = device - @torch.no_grad() - def forward(self, x, extra_options): - fuser = self.fusers_list[self.call_idx % - len(self.fusers_list)] - attn = fuser(x, self.transformer_options) - self.call_idx += 1 - return attn.to(torch.float16) + def _get_position_objs(self, extra_options): + idxs = None + if extra_options is not None: + if 'ad_params' in extra_options: + idxs = extra_options['ad_params']['sub_idxs'] + elif 'sub_idxs' in extra_options: + idxs = extra_options['sub_idxs'] + + embeddings = prepare_embeddings( + self.conds, self.latent_shape, idxs, True) + for key in embeddings: + embeddings[key] = embeddings[key].to(self.device) + objs, drop_box_mask = self.positionnet(embeddings) + return {'objs': objs, 'drop_box_mask': drop_box_mask} + + @torch.no_grad() + def forward(self, x, extra_options): + block, idx = extra_options['block'] + fuser_idx = block_map[block][idx] + fuser = self.fusers_list[fuser_idx] + attn = fuser(x, self._get_position_objs(extra_options)) + return attn.to(torch.float16) diff --git a/conditioning/instance_conditioning.py b/conditioning/instance_conditioning.py index a3a3f1b..efa77c8 100644 --- a/conditioning/instance_conditioning.py +++ b/conditioning/instance_conditioning.py @@ -2,86 +2,67 @@ import torch import comfy.model_management -from .embeddings import prepare_embeddings from .fusers_patch import FusersPatch +from .embeddings import prepare_embeddings class InstanceConditioning: - """ - This class masquerades as both a ControlNet and Gligen in order to trigger setup - """ + """ + This class masquerades as 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 = [] - self.current_device = comfy.model_management.intermediate_device() + 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 = [] + self.current_device = comfy.model_management.intermediate_device() - # ControlNet hacks - self.load_device = comfy.model_management.get_torch_device() - self.offload_device = comfy.model_management.intermediate_device() + # Gligen hacks + self.model = self + self.load_device = comfy.model_management.get_torch_device() + self.offload_device = comfy.model_management.intermediate_device() - # Gligen hacks - self.model = self + def get_fusers_patch(self, latent_shape, idxs, device): + return FusersPatch(self.conds, self.fusers_list, self.fusers_batch_size, self.positionnet, latent_shape, idxs, device) - # AnimateDiff hacks - self.sub_idxs = None - self.full_latent_length = None - self.context_length = None + def set_position(self, latent_shape, _, device): + # Called in samplers by gligen cond to return middle attention patch + batch_size = latent_shape[0] + idxs = list(range(batch_size)) + fusers_patch = self.get_fusers_patch(latent_shape, idxs, device) + return fusers_patch - def get_fusers_patch(self, objs, drop_box_mask, device): - return FusersPatch(self.fusers_list, objs, drop_box_mask, device, self.fusers_batch_size) + def add_conds(self, conds): + self.conds.extend(conds) - 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)) + def get_models(self) -> list[torch.nn.Module]: + # Used to get models for loading/offloading + return [(None, model) for model in [*self.fusers_list, self.positionnet]] - embeddings = prepare_embeddings(self.conds, latent_shape, idxs, True) - for key in embeddings: - embeddings[key] = embeddings[key].to(device) - objs, drop_box_mask = self.positionnet.to(device)(embeddings) - fusers_patch = self.get_fusers_patch(objs, drop_box_mask, device) - return fusers_patch + def inference_memory_requirements(self, dtype) -> int: + # Used to calculate memory requirements by ControlNet + return 0 - def add_conds(self, conds): - self.conds.extend(conds) + def is_clone(self, other): + return other == self - 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 model_size(self): + return 0 - def inference_memory_requirements(self, dtype) -> int: - # Used to calculate memory requirements by ControlNet - return 0 + def model_patches_to(self, device_or_dtype): + if device_or_dtype == torch.float16 or device_or_dtype == torch.float32: + return + if device_or_dtype is None: + return + self.positionnet = self.positionnet.to(device_or_dtype) + for i, fuser in enumerate(self.fusers_list): + self.fusers_list[i] = fuser.to(device_or_dtype) - def is_clone(self, other): - return other == self + def model_dtype(self): + # Fusers requires float32 so we will ignore this + return torch.float32 - def model_size(self): - return 0 - - def model_patches_to(self, device_or_dtype): - if device_or_dtype == torch.float16 or device_or_dtype == torch.float32: - return - if device_or_dtype is None: - return - self.positionnet = self.positionnet.to(device_or_dtype) - for i, fuser in enumerate(self.fusers_list): - self.fusers_list[i] = fuser.to(device_or_dtype) - - def model_dtype(self): - # Fusers requires float32 so we will ignore this - return torch.float32 - - def patch_model(self, device_to): - # TODO what is this for? - return - - # def __getattribute__(self, __name: str) -> Any: - # if __name == 'gligen': - # return 'something' # TODO - # return super().__getattribute__(__name) + def patch_model(self, device_to): + # TODO what is this for? + return diff --git a/model_helpers/prepare_fusers.py b/model_helpers/prepare_fusers.py index b3bb20b..2d2858a 100644 --- a/model_helpers/prepare_fusers.py +++ b/model_helpers/prepare_fusers.py @@ -4,27 +4,28 @@ 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'] + 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['input_blocks'][key]['params']) - fuser.load_state_dict(fusers_ckpt['input_blocks'][key]['state']) + **fusers_ckpt['middle_block']['1']['params']) + fuser.load_state_dict(fusers_ckpt['middle_block']['1']['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) - 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 + return fusers_list diff --git a/model_helpers/prepare_positionnet.py b/model_helpers/prepare_positionnet.py index 41f3bf2..5c05830 100644 --- a/model_helpers/prepare_positionnet.py +++ b/model_helpers/prepare_positionnet.py @@ -4,23 +4,23 @@ 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, - } + 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 + model = UniFusion(**params) + model.load_state_dict(checkpoint) + return model diff --git a/model_helpers/prepare_scaleu.py b/model_helpers/prepare_scaleu.py index 96ed207..eebc3d0 100644 --- a/model_helpers/prepare_scaleu.py +++ b/model_helpers/prepare_scaleu.py @@ -4,19 +4,19 @@ 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 + def scaleu_patch(h, hsp, transformer_options): + _, idx = transformer_options['block'] + sk = scaleu_nets[idx](h, hsp) + return sk - return scaleu_patch + 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 + 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 index 3311e06..bf39a9e 100644 --- a/modules/attention.py +++ b/modules/attention.py @@ -6,204 +6,204 @@ import torch.nn.functional as F def exists(val): - return val is not None + return val is not None def default(val, d): - if exists(val): - return val - return d() if isfunction(d) else 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 __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) + 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) + 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) - ) + self.net = nn.Sequential( + project_in, + nn.Dropout(dropout), + nn.Linear(inner_dim, dim_out) + ) - def forward(self, x): - return self.net(x) + 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 + 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_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)) + 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) + 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 + 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 + 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 + 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] + 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" + 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) + # 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)) + # 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 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)) + # 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 + # 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 + # 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) + 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 + # 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 + # 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 + # 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 + 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) + 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) + 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.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.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.))) + 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 + # 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'] + 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)) + 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) + return x.to(torch.float16) diff --git a/modules/convnext.py b/modules/convnext.py index 480f1a0..54e058e 100644 --- a/modules/convnext.py +++ b/modules/convnext.py @@ -14,137 +14,138 @@ 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 + 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. - """ + 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 __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) + 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 + 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 + 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. - """ + 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__() + 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) + # 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] + # 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 _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_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 + 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). - """ + 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 __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 + 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 = { @@ -162,56 +163,59 @@ model_urls = { @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 + 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 + 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 + 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 + 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 + 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 index 99da6d1..37e9d03 100644 --- a/modules/scaleu.py +++ b/modules/scaleu.py @@ -19,7 +19,8 @@ def Fourier_filter(x_in, threshold, scale): 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 + mask[..., crow - threshold: crow + threshold, + ccol - threshold: ccol + threshold] = scale x_freq = x_freq * mask # IFFT @@ -37,21 +38,24 @@ class ScaleU(nn.Module): 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) + 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 index 07b119a..ba77d7f 100644 --- a/modules/text_grounding_net.py +++ b/modules/text_grounding_net.py @@ -5,347 +5,350 @@ 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 + 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 + # 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 + # 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 + 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.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) + 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)) + 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: - 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) + 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), + ) - 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 + 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 index 33dda7d..fdf05bd 100644 --- a/modules/text_grounding_tokenizer_input.py +++ b/modules/text_grounding_tokenizer_input.py @@ -2,95 +2,95 @@ 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 __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 + 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"] + 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'] + 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"] + 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 + 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 + 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 + 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 - """ + 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 + 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) + 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) + 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) + 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 + 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 + 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 index 05ec855..cad0375 100644 --- a/modules/util.py +++ b/modules/util.py @@ -6,14 +6,13 @@ class FourierEmbedder(): self.num_freqs = num_freqs self.temperature = temperature - self.freq_bands = temperature ** ( torch.arange(num_freqs) / num_freqs ) + 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 ) ) + 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 index 0aa9f5b..9a1f6e1 100644 --- a/nodes/apply_scaleu_model_node.py +++ b/nodes/apply_scaleu_model_node.py @@ -5,32 +5,33 @@ from ..model_helpers.prepare_scaleu import get_scaleu_patch class ApplyScaleUModelNode: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "model": ("MODEL",), - "scaleu": ("SCALEU",), - }} + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "scaleu": ("SCALEU",), + }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "apply" + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply" - CATEGORY = "instance" + CATEGORY = "instance" - def apply(self, model, scaleu): - # Validate patches dict is setup correctly - transformer_options = model.model_options['transformer_options'] - if 'patches' not in transformer_options: - transformer_options['patches'] = {} + def apply(self, model, scaleu): + # Validate patches dict is setup correctly + transformer_options = model.model_options['transformer_options'] + if 'patches' not in transformer_options: + transformer_options['patches'] = {} - if 'output_block_patch' not in transformer_options['patches']: - transformer_options['patches']['output_block_patch'] = [] + if 'output_block_patch' not in transformer_options['patches']: + transformer_options['patches']['output_block_patch'] = [] - # Add scaleu patch to model patches - scaleu_nets = scaleu['model_list'] - # TODO make this load in KSampler - for i, scaleu in enumerate(scaleu_nets): - scaleu_nets[i] = scaleu.to(comfy.model_management.get_torch_device()) - transformer_options['patches']['output_block_patch'].append( - get_scaleu_patch(scaleu_nets)) - return (model,) + # Add scaleu patch to model patches + scaleu_nets = scaleu['model_list'] + # TODO make this load in KSampler + for i, scaleu in enumerate(scaleu_nets): + scaleu_nets[i] = scaleu.to( + comfy.model_management.get_torch_device()) + transformer_options['patches']['output_block_patch'].append( + get_scaleu_patch(scaleu_nets)) + return (model,) diff --git a/nodes/load_instance_fusers_node.py b/nodes/load_instance_fusers_node.py index ccf253b..612b216 100644 --- a/nodes/load_instance_fusers_node.py +++ b/nodes/load_instance_fusers_node.py @@ -4,21 +4,22 @@ from ..model_helpers.prepare_fusers import prepare_fusers class LoadInstanceFusersNode: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "model_filename": (get_model_list(constants.INSTANCE_FUSERS_DIR),), - }} + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model_filename": (get_model_list(constants.INSTANCE_FUSERS_DIR),), + }} - RETURN_TYPES = ("FUSERS",) - FUNCTION = "load_model" + RETURN_TYPES = ("FUSERS",) + FUNCTION = "load_model" - CATEGORY = "instance/loaders" + CATEGORY = "instance/loaders" - def load_model(self, model_filename: str): - checkpoint = load_checkpoint(constants.INSTANCE_FUSERS_DIR, model_filename) - fusers_list = prepare_fusers(checkpoint) - fusers = { - 'model_list': fusers_list - } - return (fusers,) + def load_model(self, model_filename: str): + checkpoint = load_checkpoint( + constants.INSTANCE_FUSERS_DIR, model_filename) + 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 index ec55db6..5ee8b9d 100644 --- a/nodes/load_instance_positionnet_node.py +++ b/nodes/load_instance_positionnet_node.py @@ -4,23 +4,23 @@ from ..model_helpers.prepare_positionnet import prepare_positionnet, get_positio class LoadInstancePositionNetNode: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "model_filename": (get_model_list(constants.INSTANCE_POSITIONNET_DIR),), - }} + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model_filename": (get_model_list(constants.INSTANCE_POSITIONNET_DIR),), + }} - RETURN_TYPES = ("POSITIONNET",) - FUNCTION = "load_model" + RETURN_TYPES = ("POSITIONNET",) + FUNCTION = "load_model" - CATEGORY = "instance/loaders" + CATEGORY = "instance/loaders" - def load_model(self, model_filename: str): - checkpoint = load_checkpoint( - constants.INSTANCE_POSITIONNET_DIR, model_filename) - params = get_positionnet_default_params() - model = prepare_positionnet(checkpoint, params) - positionnet = { - 'model': model, - } - return (positionnet,) + def load_model(self, model_filename: str): + checkpoint = load_checkpoint( + constants.INSTANCE_POSITIONNET_DIR, model_filename) + 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 index 7548756..9fde343 100644 --- a/nodes/load_instance_scaleu_node.py +++ b/nodes/load_instance_scaleu_node.py @@ -4,21 +4,22 @@ from ..model_helpers.prepare_scaleu import prepare_scaleu_nets class LoadInstanceScaleUNode: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "model_filename": (get_model_list(constants.INSTANCE_SCALEU_DIR),), - }} + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model_filename": (get_model_list(constants.INSTANCE_SCALEU_DIR),), + }} - RETURN_TYPES = ("SCALEU",) - FUNCTION = "load_model" + RETURN_TYPES = ("SCALEU",) + FUNCTION = "load_model" - CATEGORY = "instance/loaders" + CATEGORY = "instance/loaders" - def load_model(self, model_filename: str): - checkpoint = load_checkpoint(constants.INSTANCE_SCALEU_DIR, model_filename) - scaleu_list = prepare_scaleu_nets(checkpoint) - scaleu = { - 'model_list': scaleu_list - } - return (scaleu,) + def load_model(self, model_filename: str): + checkpoint = load_checkpoint( + constants.INSTANCE_SCALEU_DIR, model_filename) + scaleu_list = prepare_scaleu_nets(checkpoint) + scaleu = { + 'model_list': scaleu_list + } + return (scaleu,) diff --git a/utils/decode_item.py b/utils/decode_item.py index f57f192..9972e23 100644 --- a/utils/decode_item.py +++ b/utils/decode_item.py @@ -14,409 +14,410 @@ from skimage import measure # from nltk.corpus import stopwords def decode_base64_to_pillow(image_b64): - return Image.open(BytesIO(base64.b64decode(image_b64))).convert('RGB') + 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 + 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 + 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 + """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. + 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) + """ + 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 + 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)) + 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)] + # 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 - # 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) + 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 + # 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) + # 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 + 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 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)) + # 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 + # 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) + # 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 [] + 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 + 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])) + # 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 + 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 - + 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 - polygon = [] - for p in polygons: - polygon += p + 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]) - # 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) + polygons_single = polygons_single[indexes] + sampled_polygons = [(x[0], x[1]) for x in polygons_single] + else: + return None - return polygon - else: - sampled_points = sample_sparse_points(mask, n_polygon_points) + # 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 decode_item(item): - # convert string to dict - if "image" in item and isinstance(item['image'], Image.Image): - return item +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) - 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 + # 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 + 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") + 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") + 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 + 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] + 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] + # 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] + 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 ", " + 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 + 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... + """ + 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=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. + if random_drop_embedding=image. we will random drop image feature + and always keep the text one. - """ - N = masks.shape[0] + """ + 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 == '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 + if random_drop_embedding == 'image': + image_masks = masks * (torch.rand(N) > 0.5) * 1 + text_masks = masks - return image_masks, text_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) + """ + 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. + """ + 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) + 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 index 37c3e61..5e9fe30 100644 --- a/utils/model_utils.py +++ b/utils/model_utils.py @@ -7,18 +7,18 @@ from .. import constants def get_model_dir(model_dir): - root_path = folder_paths.folder_names_and_paths['custom_nodes'][0][0] - path = os.path.join(root_path, 'ComfyUI-InstanceDiffusion', constants.INSTANCE_MODELS_DIR, - model_dir) - return path + root_path = folder_paths.folder_names_and_paths['custom_nodes'][0][0] + path = os.path.join(root_path, 'ComfyUI-InstanceDiffusion', constants.INSTANCE_MODELS_DIR, + model_dir) + return path def get_model_list(model_dir) -> list[str]: - path = get_model_dir(model_dir) - return os.listdir(path) + path = get_model_dir(model_dir) + return os.listdir(path) def load_checkpoint(model_dir, filename): - checkpoint_path = os.path.join(get_model_dir(model_dir), filename) - checkpoint = comfy.utils.load_torch_file(checkpoint_path, safe_load=True) - return checkpoint + checkpoint_path = os.path.join(get_model_dir(model_dir), filename) + 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 index 195cb9a..bbd9293 100644 --- a/utils/prompt_utils.py +++ b/utils/prompt_utils.py @@ -2,14 +2,14 @@ import re def extract_prompts(input_string): - # Define the new regex pattern to include negative numbers - pattern = r"\"(-?\d+)\.([^\"]+)\":\s*\"([^\"]+)\"" + # 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) + # 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] + # Convert matches to a list of tuples (number, name, text) + result = [(int(number), name.strip(), text) + for number, name, text in matches] - return result + return result