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

formatting and ADE support

Browse Source
This commit is contained in:
logtd
2024-02-23 14:54:01 -07:00
parent ff935165a8
commit f269a7aafb
19 changed files with 1317 additions and 1277 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2 -1
View File
@@ -1,4 +1,5 @@
.git
.vscode
**/__pycache__
**/*.ckpt
**/*.ckpt
.DS_Store
conditioning/fusers_patch.py
+59 -16
View File
@@ -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)
conditioning/instance_conditioning.py
+48 -67
View File
@@ -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
model_helpers/prepare_fusers.py
+21 -20
View File
@@ -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
model_helpers/prepare_positionnet.py
+17 -17
View File
@@ -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
model_helpers/prepare_scaleu.py
+12 -12
View File
@@ -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
modules/attention.py
+155 -155
View File
@@ -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)
modules/convnext.py
+153 -149
View File
@@ -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
modules/scaleu.py
+23 -19
View File
@@ -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)
modules/text_grounding_net.py
+339 -336
View File
@@ -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
modules/text_grounding_tokenizer_input.py
+78 -78
View File
@@ -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
modules/util.py
+3 -4
View File
@@ -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)
nodes/apply_scaleu_model_node.py
+25 -24
View File
@@ -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,)
nodes/load_instance_fusers_node.py
+16 -15
View File
@@ -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,)
nodes/load_instance_positionnet_node.py
+17 -17
View File
@@ -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,)
nodes/load_instance_scaleu_node.py
+16 -15
View File
@@ -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,)
utils/decode_item.py
+316 -315
View File
@@ -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)
utils/model_utils.py
+9 -9
View File
@@ -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
utils/prompt_utils.py
+8 -8
View File
@@ -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