import math import random from typing import Literal, List, Dict import torch import torch.nn as nn import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence from diffusers.models.embeddings import get_fourier_embeds_from_boundingbox from diffusers.utils import logging from diffusers import ModelMixin from diffusers.configuration_utils import ConfigMixin, register_to_config from transformers import T5TokenizerFast, T5EncoderModel HOI_N_MAX = 4 BOX_N_MAX = 12 # 4x3 logger = logging.get_logger(__name__) # pylint: disable=invalid-name class GroundingInput: # Padding @staticmethod def pad_to(x, pad_shape, value=0): pad_size = list(pad_shape) pad_size[0] = pad_shape[0] - x.shape[0] if pad_size[0] > 0: pad = torch.full(pad_size, value, dtype=x.dtype, device=x.device) return torch.cat([x, pad], dim=0) return x @staticmethod @torch.no_grad() def _encode_prompt_with_t5( text_encoder: T5EncoderModel, tokenizer: T5TokenizerFast, max_sequence_length=512, prompt=None, num_images_per_prompt=1, device=None, text_input_ids=None, padding: Literal["max_length", "do_not_pad"] = "max_length", ): prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) if tokenizer is not None: text_inputs = tokenizer( prompt, padding=padding, max_length=max_sequence_length, truncation=True, return_length=False, return_overflowing_tokens=False, return_tensors="pt", ) text_input_ids = text_inputs.input_ids else: if text_input_ids is None: raise ValueError("text_input_ids must be provided when the tokenizer is not specified") prompt_embeds = text_encoder(text_input_ids.to(device))[0] if hasattr(text_encoder, "module"): dtype = text_encoder.module.dtype else: dtype = text_encoder.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) _, seq_len, _ = prompt_embeds.shape # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) return prompt_embeds @classmethod def get_hoi_seq_len(cls, max_hoi_seq_len, total_hois): # Determine hoi sequence length # we maintain about total sequence = 6144 ? # this trigger OOM during training # or we should use 512 * 3 * 3 = 4608 if total_hois <= 3: hoi_seq_len = 512 max_hois = 3 elif total_hois <= 6: hoi_seq_len = 256 max_hois = 6 elif total_hois <= 12: hoi_seq_len = 128 max_hois = 12 elif total_hois <= 24: hoi_seq_len = 64 max_hois = 24 elif total_hois <= 48: hoi_seq_len = 32 max_hois = 48 elif total_hois <= 96: hoi_seq_len = 16 max_hois = 96 elif total_hois <= 192: hoi_seq_len = 8 max_hois = 192 else: hoi_seq_len = 8 max_hois = 192 logger.warning(f"Number of HOIs ({total_hois}) exceeds the maximum limit of 192. Truncateing to 192.") hoi_seq_len = min(hoi_seq_len, max_hoi_seq_len) return hoi_seq_len, max_hois @classmethod def get_box_seq_len(cls, max_hoi_seq_len, total_boxes): hoi_seq_len, max_hois = cls.get_hoi_seq_len(max_hoi_seq_len, math.ceil(total_boxes / 3)) return hoi_seq_len, max_hois * 3 @classmethod def get_rope_ids(cls, g_text_ids, img_width: int = 64, img_height: int = 64, cond_width: int = 64, cond_height: int = 64): # Avoid in-place modification of the input tensor that may be needed for gradient computation max_img_dim = max(img_height, cond_height, img_width, cond_width) slot_ids = g_text_ids[:, 1] updated_cols = (slot_ids + max_img_dim).unsqueeze(1) g_text_ids = g_text_ids.clone() g_text_ids[:, 0] = 0 # set frame ids = 0 g_text_ids[:, 1:] = updated_cols return g_text_ids @classmethod def get_prior(cls, sx: float, sy: float, h: int, w: int, device='cpu', dtype=torch.float32): ys = (torch.arange(h, device=device, dtype=dtype) + 0.5) / h xs = (torch.arange(w, device=device, dtype=dtype) + 0.5) / w Y, X = torch.meshgrid(ys, xs, indexing="ij") cx, cy = 0.5, 0.5 eps = 1.0 / max(h, w) # avoid near-zero std sx_ = max(float(sx), eps) sy_ = max(float(sy), eps) prior = torch.exp(-(((X - cx) ** 2) / (2 * sx_ ** 2) + ((Y - cy) ** 2) / (2 * sy_ ** 2))) prior /= prior.max() # normalize to max=1 return prior @classmethod def preprocess_arbitrary_masks(cls, arbitrary_mask, img_height, img_width): # resize the arbitrary mask to img_height and img_width if arbitrary_mask is None: return None if not isinstance(arbitrary_mask, torch.Tensor): arbitrary_mask = torch.tensor(arbitrary_mask, dtype=torch.float32) # resize bool mask pytorch arbitrary_mask = arbitrary_mask.unsqueeze(0).unsqueeze(0).to(dtype=torch.float32) # [1, 1, H, W] arbitrary_mask = F.interpolate(arbitrary_mask, size=(img_height, img_width), mode="bilinear", align_corners=False) arbitrary_mask = arbitrary_mask.squeeze(0).squeeze(0).to(dtype=torch.bool) return arbitrary_mask @classmethod def get_union_masks(cls, subject_mask, object_mask): if subject_mask is None or object_mask is None: return None if subject_mask.shape != object_mask.shape: raise ValueError(f"Shape mismatch: subject_mask {subject_mask.shape} and object_mask {object_mask.shape} must have the same shape.") subject_mask = subject_mask.to(torch.bool) object_mask = object_mask.to(torch.bool) return subject_mask | object_mask @classmethod def prepare_arbitrary_masks(cls, arbitrary_masks: List[List[torch.Tensor]], g_text_ids: torch.Tensor, img_height: int = 64, img_width: int = 64, hoi_seq_len: int = 64): """ Prepare arbitrary masks for each box in the batch. Args: arbitrary_masks (List[List[torch.Tensor]]): List of batch samples, each containing a list of masks. \ Each mask is a tensor of shape [img_tokens_size] or None. List must be in [B, N * M * T]. boxes (torch.Tensor): Tensor of shape [B, N_max * M * T_max, 4] with box coordinates. img_height (int): Height of the image, default is 64. img_width (int): Width of the image, default is 64. hoi_seq_len (int): Maximum sequence length for HOI text encoding. Returns: List[List[torch.Tensor]]: Processed arbitrary masks with the same structure as input. """ seq_len, _ = g_text_ids.shape batch_size = len(arbitrary_masks) for i in range(len(arbitrary_masks)): for j in range(len(arbitrary_masks[i])): for k in range(len(arbitrary_masks[i][j])): if arbitrary_masks[i][j][k] is not None: arbitrary_masks[i][j][k] = cls.preprocess_arbitrary_masks(arbitrary_masks[i][j][k], img_height, img_width) processed_masks = [] for i in range(batch_size): sample_masks = [] for j in range(seq_len): _, slot_id, role_id = g_text_ids[j] if role_id < 2: mask = arbitrary_masks[i][slot_id][role_id] if arbitrary_masks[i][slot_id][role_id] is not None else None elif role_id == 2: # for action, it is the intersect of subject and object mask = cls.get_union_masks(arbitrary_masks[i][slot_id][0], arbitrary_masks[i][slot_id][1]) else: raise ValueError(f"Invalid role_id {role_id} at batch {i}, index {j}") sample_masks.append(mask.flatten()) # flatten to [img_tokens_size] processed_masks.append(sample_masks) return processed_masks @classmethod def prepare_attention_mask(cls, out_text_ids: torch.Tensor, out_boxes: torch.Tensor, img_tokens_size: int = 4096, txt_tokens_size: int = 512, img_width: int = 64, img_height: int = 64, cond_tokens_size: int = 4096, cond_width: int = 64, cond_height: int = 64, arbitrary_masks: List[List[torch.Tensor]] = None, use_union_action_mask: bool = True): """ Given input out_text_ids: [B, N_max * M * T_max, 3], and out_boxes: [B, N_max * M * T_max, 4] returns attention mask for the grounding encoder, where M could be 1 for object and 3 for HOI. Args: out_text_ids (torch.Tensor): Tensor of shape [B, N_max * M * T_max, 3] with text ids. out_boxes (torch.Tensor): Tensor of shape [B, N_max * M * T_max, 4] with box coordinates. img_tokens_size (int): Size of image tokens, default is 4096. txt_tokens_size (int): Size of text tokens, default is 512. img_width (int): Width of the image, default is 64. img_height (int): Height of the image, default is 64. cond_tokens_size (int): Size of condition tokens, default is 4096. cond_width (int): Width of the condition, default is 64. cond_height (int): Height of the condition, default is 64. arbitrary_masks (List[List[torch.Tensor]]): Optional list of arbitrary masks to apply. \ Tensor is in shape [img_tokens_size] or None. List must be in [B, N * M * T]. Returns: torch.Tensor: Attention mask of shape [B, N_max * M * T_max, N_max * M * T_max]. """ # assert shape of out_text_ids and out_boxes if out_text_ids.shape[0] != out_boxes.shape[1]: raise ValueError(f"Shape mismatch: out_text_ids {out_text_ids.shape} and out_boxes {out_boxes.shape} must have the same sequence length.") assert img_tokens_size == img_width * img_height, \ f"Image tokens size {img_tokens_size} must equal width {img_width} * height {img_height} = {img_width * img_height}" assert cond_tokens_size == cond_width * cond_height, \ f"Condition tokens size {cond_tokens_size} must equal width {cond_width} * height {cond_height} = {cond_width * cond_height}" batch_size, seq_len, _ = out_boxes.shape mask_shape = seq_len + txt_tokens_size + img_tokens_size + cond_tokens_size all_img_tokens_size = img_tokens_size + cond_tokens_size attention_mask = torch.zeros(batch_size, mask_shape, mask_shape, dtype=torch.bool) # set image tokens attention mask to 1, last img_tokens_size tokens attention_mask[:, -all_img_tokens_size:, -all_img_tokens_size:] = 1 # set text tokens attention mask to 1, first txt_tokens_size tokens attention_mask[:, seq_len:seq_len+txt_tokens_size, seq_len:seq_len+txt_tokens_size] = 1 # set the cross attention mask for text tokens and image tokens attention_mask[:, seq_len:seq_len+txt_tokens_size, -all_img_tokens_size:] = 1 attention_mask[:, -all_img_tokens_size:, seq_len:seq_len+txt_tokens_size] = 1 # check if a token is valid, this can be obtained from box coordinates, it should be dropped if it is negative # we set attention mask of invalid one to False (at the end of this method) valid_seq = (out_boxes >= 0).all(dim=2).cpu() vq = valid_seq.unsqueeze(2) vk = valid_seq.unsqueeze(1) valid_seq = torch.ones([batch_size, seq_len, seq_len], dtype=torch.bool, device=attention_mask.device) & vq & vk # the text_ids could be in the form of: # tensor([[1, 0, 0], # [1, 0, 1], # [1, 0, 2], # [1, 1, 0], # [1, 1, 1], # [1, 1, 2]]), # where the first dimension is the batch size, and the second dimension is the sequence length. # for the element with same second element, we set their attention mask to 1 for i in range(batch_size): for j in range(seq_len): # Only compare with the seq_len tokens, and assign to the correct slice if out_text_ids[j, 0] == 0: # if the first element is 0, it is a empty token and not valid raise ValueError(f"Invalid token at batch {i}, index {j}: {out_text_ids[j]}") attention_mask[i, j, :seq_len] = ( out_text_ids[j, 1] == out_text_ids[:, 1] ) # based on the out_boxes at the same index, set the attention mask, the boxes are in the form of: # [x1, y1, x2, y2] box = out_boxes[i, j] # verify box are valid and make sure both width and height not negative if (box >= 0).all() and (box[2] - box[0]) >= 0 and (box[3] - box[1]) >= 0: if arbitrary_masks is not None and arbitrary_masks[i][j] is not None: box_attn_mask = arbitrary_masks[i][j] if box_attn_mask.numel() != img_tokens_size: raise ValueError(f"Arbitrary mask at batch {i}, index {j} has incorrect size {box_attn_mask.numel()}, expected {img_tokens_size}") elif box.sum() < 1e-6: # 2e-4 is min res for 64x64, 1e-6 is almost zero, here we want set zero box (randomly dropped box) with all attended box_attn_mask = torch.ones(img_tokens_size, dtype=torch.bool) elif use_union_action_mask and out_text_ids[j, 2] == 2: # for action, we use the union of subject and object # we assume both direction of attention is same, thus we take from one only subject_index = (out_text_ids[:, 2] == 0) & (out_text_ids[:, 1] == out_text_ids[j, 1]) subject_index = subject_index.to(device=attention_mask.device) object_index = (out_text_ids[:, 2] == 1) & (out_text_ids[:, 1] == out_text_ids[j, 1]) object_index = object_index.to(device=attention_mask.device) if img_tokens_size == all_img_tokens_size: subject_attn_mask = attention_mask[i, -img_tokens_size:, :seq_len][:, subject_index] object_attn_mask = attention_mask[i, -img_tokens_size:, :seq_len][:, object_index] else: subject_attn_mask = attention_mask[i, -all_img_tokens_size:-all_img_tokens_size+img_tokens_size, :seq_len][:, subject_index] object_attn_mask = attention_mask[i, -all_img_tokens_size:-all_img_tokens_size+img_tokens_size, :seq_len][:, object_index] if subject_attn_mask.numel() == 0 or object_attn_mask.numel() == 0: box_attn_mask = torch.ones(img_tokens_size, dtype=torch.bool) else: subject_attn_mask = subject_attn_mask.any(dim=1) object_attn_mask = object_attn_mask.any(dim=1) box_attn_mask = subject_attn_mask | object_attn_mask else: box_attn_mask = torch.zeros(img_tokens_size, dtype=torch.bool) box_attn_mask = box_attn_mask.reshape(img_height, img_width) x1_idx = int(box[0] * img_width) y1_idx = int(box[1] * img_height) x2_idx = int(box[2] * img_width) y2_idx = int(box[3] * img_height) # Make the end indices inclusive, but clamp to image size x2_idx = min(x2_idx, img_width - 1) y2_idx = min(y2_idx, img_height - 1) # Add 1 to end indices for inclusive slicing box_attn_mask[y1_idx:y2_idx+1, x1_idx:x2_idx+1] = 1 # flatten the box attention mask to match the img_tokens_size box_attn_mask = box_attn_mask.flatten() # set the attention mask for the box tokens if img_tokens_size == all_img_tokens_size: attention_mask[i, -img_tokens_size:, j] = box_attn_mask attention_mask[i, j, -img_tokens_size:] = box_attn_mask else: attention_mask[i, j, -all_img_tokens_size:-all_img_tokens_size+img_tokens_size] = box_attn_mask attention_mask[i, -all_img_tokens_size:-all_img_tokens_size+img_tokens_size, j] = box_attn_mask # For HOI, prevent S to attend to O and vice versa roles = out_text_ids[:, 2] is_S = roles == 0 is_O = roles == 1 # is_A = roles == 2 forbid_SO = (is_S[:, None] & is_O[None, :]) | (is_O[:, None] & is_S[None, :]) forbid = forbid_SO forbid = forbid.to(device=attention_mask.device) attention_mask[:, :seq_len, :seq_len] &= ~forbid # set invalid one to False # this seems to cause NaN, because for some invalid query, all its key now become 0 # we can fix it with minimal self-attention via a diagonal mask eye = torch.eye(seq_len, device=attention_mask.device, dtype=torch.bool)[None] attention_mask[:, :seq_len, :seq_len] = attention_mask[:, :seq_len, :seq_len].bool() & valid_seq[:, :seq_len, :seq_len] | eye negative_mask = attention_mask[:, seq_len:, seq_len:] return attention_mask, negative_mask @classmethod def prepare_train_input(cls, tokenizer, text_encoder, boxes=None, hois=None, objects=None, random_drop_boxes: float = 0.0, random_drop_hois: float = 0.0, hoi_seq_len: int = 64): """ Prepares input for the grounding encoder during training. Args: tokenizer (T5TokenizerFast): Tokenizer for encoding text. boxes (List[List[List[float]]]): List of batch samples, each containing a list of boxes, each box as [x1, y1, x2, y2]. hois (List[List[dict]]): List of batch samples, each containing a list of HOI labels. objects (List[List[dict]]): List of batch samples, each containing a list of object labels. random_drop_boxes (float): Probability of randomly dropping boxes during training. random_drop_hois (float): Probability of randomly dropping HOIs during training. hoi_seq_len (int): Maximum sequence length for HOI text encoding. max_box (int): Maximum number of boxes to consider. max_hoi (int): Maximum number of HOIs to consider. Returns: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - out_embeds: Tensor of shape [B, T_max, D] with text embeddings. - out_boxes: Tensor of shape [B, T_max, 4] with box coordinates. - out_text_ids: Tensor of shape [B, T_max, 3] with text ids. """ # If no boxes, hois, or objects are provided, return None if boxes[0] is None and hois[0] is None and objects[0] is None: return None, None, None if random.random() < random_drop_boxes: # replace boxes with zeros, maintaining the shape, # the shape is [B, N, 1, 4] or [B, N, 2, 4] for HOI # for each B x tensor(N,1,4) or B x tensor(N,2,4) # print(boxes) # print(boxes[0].shape) for sample in boxes: if isinstance(sample, torch.Tensor): sample.fill_(0.0) else: raise ValueError("boxes should be a list of tensors or None, got: {}. HOI: {} BOX:{}".format( type(sample), objects[0] is None, boxes[0] is not None)) # print(f"Randomly dropping boxes, replaced with zeros. hoi:{objects[0] is None} box:{objects[0] is not None}") # print(boxes) # HOI case: hois provided, objects are None if boxes[0] is not None and hois[0] is not None and objects[0] is None: if random.random() < random_drop_hois: box_labels = cls.obtain_only_box_labels_from_hoi(hois) # reshape boxes from hoi to independent boxes, to reduce it into Box Generation task. # Flatten boxes from [B, N, M, 4] to [B, N*M, 4] for each sample, then unsqueeze to [B, N*M, 1, 4] boxes = [ sample.reshape(-1, 4) if isinstance(sample, torch.Tensor) else torch.tensor(sample, dtype=torch.float32).reshape(-1, 4) for sample in boxes ] # boxes = [[torch.tensor(box, dtype=torch.float32) for box in sample] for sample in boxes] # keep as list of lists for variable N boxes = [sample.unsqueeze(1) for sample in boxes] # [B][N*M, 1, 4] _hoi_seq_len, _max_box = cls.get_box_seq_len(hoi_seq_len, total_boxes=max(len(h) for h in box_labels)) processed_boxes, box_prompt_embeds = cls.preprocess_box(boxes, box_labels, tokenizer, text_encoder, _hoi_seq_len, max_box=_max_box) return cls.prepare_box(box_prompt_embeds, processed_boxes, _max_box) else: box_labels, hoi_labels = cls.obtain_box_hoi_labels(hois) _hoi_seq_len, _max_hoi = cls.get_hoi_seq_len(hoi_seq_len, max(len(h) for h in hoi_labels) if hoi_labels is not None else 0) processed_boxes, box_prompt_embeds = cls.preprocess_hoi(boxes, box_labels, hoi_labels, tokenizer, text_encoder, _hoi_seq_len, max_hoi=_max_hoi) return cls.prepare_hoi(box_prompt_embeds, processed_boxes, _max_hoi) # If boxes and objects are provided elif boxes[0] is not None and objects[0] is not None and hois[0] is None: box_labels = cls.obtain_box_labels(objects) _hoi_seq_len, _max_box = cls.get_box_seq_len(hoi_seq_len, total_boxes=max(len(h) for h in box_labels)) processed_boxes, box_prompt_embeds = cls.preprocess_box(boxes, box_labels, tokenizer, text_encoder, _hoi_seq_len, max_box=_max_box) return cls.prepare_box(box_prompt_embeds, processed_boxes, _max_box) else: raise ValueError(f"Unexpected case of boxes={'None' if boxes is None else 'Not None'}," f" objects={'None' if objects is None else 'Not None'}," f" hois={'None' if hois is None else 'Not None'}.") @classmethod def prepare_mixed_pipeline_input(cls, tokenizer, text_encoder, mix_boxes=None, mix_box_labels=None, mix_hoi_labels=None, hoi_seq_len: int = 64, max_box: int | None = None): """ Prepares input for the grounding encoder pipeline, accepting arbitary modality. Args: tokenizer (T5TokenizerFast): Tokenizer for encoding text. text_encoder (T5EncoderModel): Text encoder model for encoding tokenized text. mix_boxes (List[List[List[float]]]): List of batch samples with mixed boxes or None. mix_box_labels (List[List[str]]): List of batch samples with mixed box labels or None. mix_hoi_labels (List[List[str]]): List of batch samples with mixed HOI labels or None. hoi_seq_len (int): Maximum sequence length for HOI text encoding. max_box (int): Maximum number of boxes to consider. max_hoi (int): Maximum number of HOIs to consider. """ if mix_boxes is None and mix_box_labels is None and mix_hoi_labels is None: return None, None, None processed_boxes, box_prompt_embeds = cls.preprocess_mixed(mix_boxes, mix_box_labels, mix_hoi_labels, tokenizer, text_encoder, hoi_seq_len, max_box=max_box) out_embeds, out_boxes, out_text_ids = cls.prepare_mixed(box_prompt_embeds, processed_boxes, max_box=max_box) return out_embeds, out_boxes, out_text_ids @classmethod def preprocess_mixed(cls, mix_boxes, mix_box_labels, mix_hoi_labels, tokenizer, text_encoder, hoi_seq_len, max_box=None): """ Preprocesses mixed boxes and labels for the grounding encoder. Args: mix_boxes (List[List[List[float]]]): List of batch samples with mixed boxes or None. The shape must be [B, N, M] x [4 or None] mix_box_labels (List[List[str]]): List of batch samples with mixed box labels or None. The shape must be [B, N, M] x (str or None) mix_hoi_labels (List[List[str]]): List of batch samples with mixed HOI labels or None. The shape must be [B, N, 1] x (str or None). If n-th box_labels has M=3, then this must be str. tokenizer (T5TokenizerFast): Tokenizer for encoding text. text_encoder (T5EncoderModel): Text encoder model for encoding tokenized text. hoi_seq_len (int): Maximum sequence length for HOI text encoding. max_box (int): Maximum number of boxes to consider. """ assert len(mix_boxes) == len(mix_box_labels) == len(mix_hoi_labels), \ f"Batch size mismatch: mix_boxes {len(mix_boxes)}, mix_box_labels {len(mix_box_labels)}, mix_hoi_labels {len(mix_hoi_labels)}" for b in range(len(mix_boxes)): assert len(mix_boxes[b]) == len(mix_box_labels[b]) == len(mix_hoi_labels[b]), \ f"HOI instance number mismatch at index {b}: mix_boxes {len(mix_boxes[b])}, mix_box_labels {len(mix_box_labels[b])}, mix_hoi_labels {len(mix_hoi_labels[b])}" for n in range(len(mix_boxes[b])): assert len(mix_boxes[b][n]) == len(mix_box_labels[b][n]), \ f"Role (subject/object/action) mismatch at index {b},{n}: mix_boxes {len(mix_boxes[b][n])}, mix_box_labels {len(mix_box_labels[b][n])}" if len(mix_boxes[b][n]) == 2: assert isinstance(mix_hoi_labels[b][n], str), \ f"HOI label must be str when box_labels has 2 roles at index {b},{n}: got {type(mix_hoi_labels[b][n])}" else: assert mix_hoi_labels[b][n] is None, \ f"HOI label must be None when box_labels has not 2 roles at index {b},{n}: got {mix_hoi_labels[b][n]}" assert len(mix_boxes[b][n]) in [1, 2], \ f"Number of roles (subject/object/action) must be 1 or 2 for Object or HOI instance at index {b},{n}: got {len(mix_boxes[b][n])}" B = len(mix_boxes) processed_boxes = [] box_prompt_embeds = [] for b in range(B): N = len(mix_boxes[b]) box_list = [] box_prompt_list = [] for n in range(N): M = len(mix_boxes[b][n]) # M is 1 for object and 2 for HOI for m, box in enumerate(mix_boxes[b][n]): if box is None: mix_boxes[b][n][m] = [0.0, 0.0, 0.0, 0.0] if M == 2: # get_action_boxes input is [B, N, 2, 4] _boxes = torch.tensor(mix_boxes[b][n], dtype=torch.float32).unsqueeze(0).unsqueeze(0) # [1, 1, 2, 4] action_box = cls.get_enclosing_action_boxes(_boxes)[0, 0].tolist() # [4] mix_boxes[b][n].append(action_box) mix_box_labels[b][n].append(mix_hoi_labels[b][n]) embs = cls._encode_prompt_with_t5( text_encoder, tokenizer, device=text_encoder.device, prompt=mix_box_labels[b][n], padding="max_length", max_sequence_length=hoi_seq_len ) # [chunk, T, D] box_list.append(mix_boxes[b][n]) # proc_box is a list of length B=1 box_prompt_list.append(embs) # box_prompt is a list of length B=1 # Concatenate all boxes and prompts for this batch item processed_boxes.append(box_list) # box_list: [N*M, 4] box_prompt_embeds.append(box_prompt_list) # box_prompt_list: [N*M, T, D] # processed_boxes: [B, N*M, 4] # box_prompt_embeds: [B, N*M, T, D] return processed_boxes, box_prompt_embeds @classmethod def prepare_mixed(cls, box_prompt_embeds, processed_boxes, max_box=None): """ Prepares mixed box embeddings and coordinates for the grounding encoder. Args: box_prompt_embeds (List[List[torch.Tensor]]): List of batch samples, each containing a list of box prompt embeddings. processed_boxes (List[List[List[float]]]): List of batch samples, each containing a list of boxes, each box as [x1, y1, x2, y2]. max_box (int): Maximum number of boxes to consider. Returns: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - out_embeds: Tensor of shape [B, T_max, D] with text embeddings. - out_boxes: Tensor of shape [B, T_max, 4] with box coordinates. - out_text_ids: Tensor of shape [B, T_max, 3] with text ids. """ B = len(processed_boxes) # if B is > 1, # when N is variable, then we must set all to N_max, and pad the missing ones # when M is variable (1 or 3), we must set all to 3, and pad the missing ones # assert B == 1, f"Batch size must be 1 for mixed input, got {B=}" T = box_prompt_embeds[0][0][0].shape[0] # Assuming all have the same T D = box_prompt_embeds[0][0][0].shape[-1] # Assuming all have the same D device = box_prompt_embeds[0][0][0].device if B > 1: N_max = min(max_box, max(len(b) for b in processed_boxes)) out_embeds = torch.zeros((B, N_max * 3 * T, D), dtype=torch.float32, device=device) out_boxes = torch.full((B, N_max * 3 * T, 4), -1.0, dtype=torch.float32, device=device) _M = 3 ones = torch.full((N_max, _M, T, 1), 8, dtype=torch.long, device=device) # it was 1, now 8 ns = torch.arange(N_max, device=device).view(N_max, 1, 1, 1).expand(N_max, _M, T, 1) ms = torch.arange(_M, device=device).view(1, _M, 1, 1).expand(N_max, _M, T, 1) ids = torch.cat([ones, ns, ms], dim=-1) # [N_max, M, T_max, 3] for i in range(B): N = min(len(processed_boxes[i]), N_max) for n in range(N): M = len(processed_boxes[i][n]) # M is 1 for object and 3 for HOI for m in range(M): out_embeds[i, (n * _M + m) * T : (n * _M + m + 1) * T] = box_prompt_embeds[i][n][m] out_boxes[i, (n * _M + m) * T : (n * _M + m + 1) * T] = torch.tensor(processed_boxes[i][n][m], dtype=torch.float32, device=device).unsqueeze(0).expand(T, 4) # out_embeds: [B, N_max*M*T, D] # out_boxes: [B, N_max*M*T, 4] # out_text_ids: [N_max*M*T, 3] out_embeds = out_embeds out_boxes = out_boxes out_text_ids = ids.view(N_max * _M * T, 3) elif B == 1: # this we could handle variable N and M N = len(processed_boxes[0]) NM = sum(len(b) for b in processed_boxes[0]) # N * M out_embeds = torch.zeros((1, NM * T, D), dtype=torch.float32, device=device) out_boxes = torch.full((1, NM * T, 4), -1.0, dtype=torch.float32, device=device) out_text_ids = torch.zeros((NM * T, 3), dtype=torch.long, device=device) out_text_ids[:, 0] = 8 # it was 1, now 8 nm = 0 for n in range(N): M = len(processed_boxes[0][n]) # M is 1 for object and 3 for HOI for m in range(M): out_embeds[0, nm * T : (nm + 1) * T] = box_prompt_embeds[0][n][m] out_boxes[0, nm * T : (nm + 1) * T] = torch.tensor(processed_boxes[0][n][m], dtype=torch.float32, device=device).unsqueeze(0).expand(T, 4) out_text_ids[nm * T : (nm + 1) * T, 1] = n out_text_ids[nm * T : (nm + 1) * T, 2] = m nm += 1 # out_embeds: [B, N*M*T, D] # out_boxes: [B, N*M*T, 4] # out_text_ids: [N*M*T, 3] else: raise ValueError(f"Batch size must not be 0, got {B=}") return out_embeds, out_boxes, out_text_ids @classmethod def prepare_pipeline_input(cls, tokenizer, text_encoder, boxes=None, box_labels=None, hoi_labels=None, hoi_seq_len: int = 64, max_hoi: int | None = None, max_box: int | None = None): """ Prepares input for the grounding encoder pipeline. Deterministic, no random drop of boxes and hois for inference. Args: tokenizer (T5TokenizerFast): Tokenizer for encoding text. text_encoder (T5EncoderModel): Text encoder model for encoding tokenized text. boxes (List[List[List[float]]]): List of batch samples, each containing a list of boxes, each box as [x1, y1, x2, y2]. box_labels (List[List[str]]): List of batch samples, each containing a list of labels for boxes. hoi_labels (List[List[str]]): List of batch samples, each containing a list of HOI labels. Returns: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - out_embeds: Tensor of shape [B, T_max, D] with text embeddings. - out_boxes: Tensor of shape [B, T_max, 4] with box coordinates. - out_text_ids: Tensor of shape [B, T_max, 3] with text ids. Raises: ValueError: If boxes, box_labels, and hoi_labels combo are unexpected. """ if boxes is None and box_labels is None and hoi_labels is None: return None, None, None # boxes is [[sub1, obj2], ...] if hoi_labels is not None and box_labels is not None and boxes is not None: max_hoi = max_hoi if max_hoi is not None else HOI_N_MAX processed_boxes, box_prompt_embeds = cls.preprocess_hoi(boxes, box_labels, hoi_labels, tokenizer, text_encoder, hoi_seq_len, limit_max_hoi=False, max_hoi=max_hoi) return cls.prepare_hoi(box_prompt_embeds, processed_boxes, max_hoi=max_hoi) elif hoi_labels is not None and box_labels is not None and boxes is None: # Create dummy subject and object boxes of zeros for each HOI in each batch max_hoi = max_hoi if max_hoi is not None else HOI_N_MAX boxes = [ [ [0.0, 0.0, 0.0, 0.0] for _ in range(len(hois) * 2) ] for hois in hoi_labels ] processed_boxes, box_prompt_embeds = cls.preprocess_hoi(boxes, box_labels, hoi_labels, tokenizer, text_encoder, hoi_seq_len, limit_max_hoi=False, max_hoi=max_hoi) return cls.prepare_hoi(box_prompt_embeds, processed_boxes, max_hoi=max_hoi) elif boxes is not None and box_labels is not None and hoi_labels is None: max_box = max_box if max_box is not None else BOX_N_MAX processed_boxes, box_prompt_embeds = cls.preprocess_box(boxes, box_labels, tokenizer, text_encoder, hoi_seq_len, limit_max_box=False, max_box=max_box) return cls.prepare_box(box_prompt_embeds, processed_boxes, max_box=max_box) elif boxes is None and box_labels is not None and hoi_labels is None: max_box = max_box if max_box is not None else BOX_N_MAX boxes = [ [ [0.0, 0.0, 0.0, 0.0] for _ in range(len(labels)) ] for labels in box_labels ] processed_boxes, box_prompt_embeds = cls.preprocess_box(boxes, box_labels, tokenizer, text_encoder, hoi_seq_len, limit_max_box=False, max_box=max_box) return cls.prepare_box(box_prompt_embeds, processed_boxes, max_box=max_box) else: raise ValueError(f"Unexpected case of boxes={'None' if boxes is None else 'Not None'}," f" box_labels={'None' if box_labels is None else 'Not None'}," f" hoi_labels={'None' if hoi_labels is None else 'Not None'}.") @classmethod def obtain_only_box_labels_from_hoi(cls, batch_hois: List[List[Dict]]): """ Extracts only box labels from a batch of HOIs. This is used when we randomly drop HOI labels. Args: batch_hois (List[List[Dict]]): List of batch samples, each containing a list of HOI dictionaries. Returns: List[List[str]]: List of box labels for each batch sample. """ box_labels = [] for batch in batch_hois: batch_box_labels = [] for hoi in batch: batch_box_labels.extend([hoi['subject']]) batch_box_labels.extend([hoi['object']]) box_labels.append(batch_box_labels) return box_labels @classmethod def obtain_box_hoi_labels(cls, batch_hois: List[List[Dict]]): box_labels = [] hoi_labels = [] for batch in batch_hois: batch_box_labels = [] batch_hoi_labels = [] for hoi in batch: batch_box_labels.extend([hoi['subject'], hoi['object']]) batch_hoi_labels.extend([hoi['action']]) box_labels.append(batch_box_labels) hoi_labels.append(batch_hoi_labels) return box_labels, hoi_labels @classmethod def obtain_box_labels(cls, batch_objs: List[List[Dict]]): box_labels = [] for batch in batch_objs: batch_box_labels = [] for obj in batch: batch_box_labels.extend([obj['phrases']]) box_labels.append(batch_box_labels) return box_labels @classmethod def preprocess_box(cls, boxes, box_labels, tokenizer, text_encoder, hoi_seq_len=64, limit_max_box=True, max_box=BOX_N_MAX): """ Prepares box data for the grounding encoder. Args: boxes (List[List[List[float]]]): List of batch samples, each containing a list of boxes, each box as [x1, y1, x2, y2]. box_labels (List[List[str]]): List of batch samples, each containing a list of labels for boxes. tokenizer: Tokenizer object used to tokenize box labels. text_encoder: Text encoder model used to encode tokenized labels. Returns: processed_boxes (List[List[List[List[float]]]]): List of batch samples, each containing a list of boxes, each box wrapped in a list as [[x1, y1, x2, y2]]. box_prompt_embeds (List[List[List[Tensor]]]): List of batch samples, each containing a list of boxes, each box as a list containing a tensor of shape [T', D] for token embeddings. """ B = len(boxes) processed_boxes = [] box_prompt_embeds = [] N = max_box if limit_max_box else max(len(b) for b in boxes) # number of boxes per sample, max 12 for b in range(B): token_budget = 9 * 512 chunk_size = max(1, token_budget // hoi_seq_len) all_embs = [] box_texts = box_labels[b][:N] _local_N = len(box_texts) for start in range(0, len(box_texts), chunk_size): chunk = box_texts[start:start + chunk_size] chunk_embs = cls._encode_prompt_with_t5( text_encoder, tokenizer, device=text_encoder.device, prompt=chunk, padding="max_length", max_sequence_length=hoi_seq_len ) # [chunk, T, D] all_embs.append(chunk_embs) embs = torch.cat(all_embs, dim=0) # [N*M, T, D] embs_ = embs.reshape(_local_N, 1, hoi_seq_len, -1) box_prompt_embeds.append(embs_) # this has to be [B, N, 1, T, D] if isinstance(boxes[b], torch.Tensor) and boxes[b].ndim == 3: # already in [N, 1, 4] processed_boxes.append(boxes[b][:N]) # act like no-op else: processed_boxes.append([[box] for box in boxes[b][:N]]) # processed_boxes should be [B, N ,1, 4] return processed_boxes, box_prompt_embeds @classmethod def prepare_box(cls, box_prompt_embeds, boxes, max_box=BOX_N_MAX): """ Prepares box features for the grounding encoder. boxes: [B, N, 1, 4] where last dim is [x1, y1, x2, y2] box_prompt_embeds: Tensor where each is a sample in Tensor [B, N, M, T, D] representing token embeddings for each box. Outputs: - out_embeds: [B, N_max * 1 * T, D] - out_boxes: [B, N_max * 1 * T_max, 4] - out_text_ids: [B, N_max * 1 * T_max, 3] """ B = len(boxes) N_max, M = max_box, 1 # max boxes per sample, 1 box per interaction N = min(N_max, max(len(b) for b in boxes)) # number of boxes per sample T = box_prompt_embeds[0][0][0].shape[0] # 64 or cfg.model.hoi_max_seq_len D = box_prompt_embeds[0][0][0].shape[-1] # 4096 device = box_prompt_embeds[0][0][0].device # Allocate outputs out_embeds = torch.zeros((B, N, M, T, D), device=device) out_boxes = torch.full((B, N, M, T, 4), -1.0, device=device) # Text ids for boxes, [B, N, M, T, 3] ones = torch.full((N, M, T, 1), 8, dtype=torch.long, device=device) # it was 1, now 8 ns = torch.arange(N, device=device).view(N, 1, 1, 1).expand(N, M, T, 1) ms = torch.arange(M, device=device).view(1, M, 1, 1).expand(N, M, T, 1) ids = torch.cat([ones, ns, ms], dim=-1) # [N, M, T, 3] ids = ids.unsqueeze(0).expand(B, -1, -1, -1, -1) # [B, N, M, T, 3] for i in range(B): N_i = min(len(box_prompt_embeds[i]), N) if isinstance(boxes[i], torch.Tensor): boxes_tensor = boxes[i][:N_i] # N, 4 else: boxes_tensor = torch.tensor(boxes[i][:N_i], dtype=torch.float32, device=device) # [N, 4] for n in range(N_i): for m in range(M): emb = box_prompt_embeds[i][n][m] # [T', D] out_embeds[i, n, m] = emb # boxes[i][n, m, :] is [4], expand to [T, 4] out_boxes[i, n, m] = boxes_tensor[n, m].expand(T, 4) # Padding happens automatically # Reshape to [B, N_max * M * T_max, ...] out_embeds = out_embeds.view(B, N * M * T, D) out_boxes = out_boxes.view(B, N * M * T, 4) out_text_ids = ids.view(B, N * M * T, 3)[0] return out_embeds, out_boxes, out_text_ids @classmethod def preprocess_hoi(cls, boxes, box_labels, hoi_labels, tokenizer, text_encoder, hoi_seq_len=64, limit_max_hoi=True, max_hoi=HOI_N_MAX): """ Prepares HOI data for the grounding encoder. Args: boxes (List[List[float]]): List of lists of lists of floats, where each innermost list contains [x1, y1, x2, y2] for subject and object boxes. box_labels (List[List[str]]): List of lists of strings, each inner list contains labels for boxes per interaction. hoi_labels (List[List[str]]): List of lists of strings, each inner list contains HOI labels per interaction. tokenizer: Tokenizer object used to tokenize box and HOI labels. text_encoder: Text encoder model used to encode tokenized labels. Returns: boxes (List[List[List[float]]]): Tensor of shape [B, N, 3, 4] where last dim is [x1, y1, x2, y2]. box_prompt_embeds (List[List[List[Tensor]]]): List of lists of lists of Tensors, where each Tensor is [3, D] for each box. """ # boxes: [B, N*2, 4] -> [B, N, 2, 4] # box_labels: List[List[str]], hoi_labels: List[List[str]] # For each interaction, create a dict with subject, object, action B = len(boxes) M = 3 if isinstance(boxes[0], torch.Tensor) and boxes[0].ndim == 3: # if already in [N, 2, 4], reshape back to [N*2, 4] boxes = [b.reshape(-1, 4) for b in boxes] processed_boxes = [] box_prompt_embeds = [] for b in range(B): N = min(len(boxes[b]) // 2, max_hoi) if limit_max_hoi else len(boxes[b]) // 2 # number of interactions, max 4 boxes_b = [] hoi_texts = [] for n in range(N): hoi_texts.extend([box_labels[b][n*2], box_labels[b][n*2+1], hoi_labels[b][n]]) # subject, object, action subject_box = boxes[b][n*2] object_box = boxes[b][n*2+1] boxes_b.append([subject_box, object_box]) # hoi_texts may be large; batch them to save GPU memory. # At hoi_seq_len=512 we can handle 4 HOIs -> token budget = 4 * 512 = 2048 "token-units". token_budget = 9 * 512 chunk_size = max(1, token_budget // hoi_seq_len) all_embs = [] for start in range(0, len(hoi_texts), chunk_size): chunk = hoi_texts[start:start + chunk_size] chunk_embs = cls._encode_prompt_with_t5( text_encoder, tokenizer, device=text_encoder.device, prompt=chunk, padding="max_length", max_sequence_length=hoi_seq_len ) # [chunk, T, D] all_embs.append(chunk_embs) embs = torch.cat(all_embs, dim=0) # [N*M, T, D] embs_ = embs.reshape(N, M, hoi_seq_len, -1) processed_boxes.append(boxes_b) box_prompt_embeds.append(embs_) return processed_boxes, box_prompt_embeds @classmethod def prepare_hoi(cls, box_prompt_embeds, boxes, max_hoi=HOI_N_MAX): """ Enforces: - max 8 interactions per sample - 3 boxes per interaction (subject, object, action) - 10 tokens per box Outputs all tensors with shape [B, 10*3*10, ...] Prepares box features for the grounding encoder. boxes: [B, N, 2, 4] where last dim is [x1, y1, x2, y2] and dim=2 indexes subject and object. box_prompt_embeds: List[List[List[Tensor]]] where each Tensor is [T', D] for each box. Outputs: - out_embeds: [B, N_max * M * T_max, D] - out_boxes: [B, N_max * M * T_max, 4] - out_text_ids: [B, N_max * M * T_max, 3 """ B = len(boxes) N_max, M = max_hoi, 3 # interactions, boxes N_max = min(N_max, max(len(b) for b in boxes)) # number of boxes per sample T = box_prompt_embeds[0][0][0].shape[0] # T_max, e.g. 64 D = box_prompt_embeds[0][0][0].shape[-1] device = box_prompt_embeds[0][0][0].device # Allocate outputs out_embeds = torch.zeros((B, N_max, M, T, D), device=device) out_boxes = torch.full((B, N_max, M, T, 4), -1.0, device=device) # out_text_ids = torch.full((B, N_max, M, T_max, 3), -1, dtype=torch.long, device=boxes[0].device) ones = torch.full((N_max, M, T, 1), 8, dtype=torch.long, device=device) # it was 1, now 8 ns = torch.arange(N_max, device=device).view(N_max, 1, 1, 1).expand(N_max, M, T, 1) ms = torch.arange(M, device=device).view(1, M, 1, 1).expand(N_max, M, T, 1) ids = torch.cat([ones, ns, ms], dim=-1) # [N_max, M, T_max, 3] ids = ids.unsqueeze(0).expand(B, -1, -1, -1, -1) # [B, N_max, M, T_max, 3] for i in range(B): N = min(len(box_prompt_embeds[i]), N_max) if isinstance(boxes[i], torch.Tensor): so_boxes = boxes[i][:N] elif isinstance(boxes[i], list) and isinstance(boxes[i][0], list): # List[List[Tensor]], B, N, 2 is nested list of [4] tensor # Convert to tensor [N, 2, 4], assuming boxes[i] is a list of lists of tensors [4] so_boxes = torch.stack( [torch.stack([torch.as_tensor(b, dtype=torch.float32) for b in pair], dim=0) for pair in boxes[i]], dim=0 ) else: so_boxes = torch.tensor(boxes[i][:N], dtype=torch.float32, device=device) # [N, 2, 4] action_boxes = cls.get_enclosing_action_boxes(so_boxes.unsqueeze(0)).squeeze(0) # [N, 4] # Union soa_boxes = torch.cat([so_boxes, action_boxes.unsqueeze(1)], dim=1) # [N, 3, 4] for n in range(N): for m in range(M): emb = box_prompt_embeds[i][n][m] # [T', D] # Fill the valid parts out_embeds[i, n, m] = emb out_boxes[i, n, m] = soa_boxes[n, m].expand(T, 4) # Reshape to [B, 300, ...] out_embeds = out_embeds.view(B, N_max * M * T, D) out_boxes = out_boxes.view(B, N_max * M * T, 4) out_text_ids = ids.view(B, N_max * M * T, 3)[0] return out_embeds, out_boxes, out_text_ids @classmethod def get_action_boxes(cls, boxes): """ Compute action boxes using the 'between' operation. Args: boxes (torch.Tensor): Tensor of shape [B, N, 2, 4], where the last dimension is [x1, y1, x2, y2] and dim=2 indexes subject and object. Returns: torch.Tensor: Action boxes of shape [B, N, 4] """ subj_boxes = boxes[:, :, 0, :] # [B, N, 4] obj_boxes = boxes[:, :, 1, :] # [B, N, 4] all_x = torch.cat([subj_boxes[:, :, 0::2], obj_boxes[:, :, 0::2]], dim=-1) # x1, x2 all_y = torch.cat([subj_boxes[:, :, 1::2], obj_boxes[:, :, 1::2]], dim=-1) # y1, y2 all_x, _ = all_x.sort(dim=-1) all_y, _ = all_y.sort(dim=-1) # return [x1, y1, x2, y2] between boxes return torch.stack([all_x[:, :, 1], all_y[:, :, 1], all_x[:, :, 2], all_y[:, :, 2]], dim=-1) @classmethod def get_enclosing_action_boxes(cls, boxes): """ Compute enclosing action boxes. Args: boxes (torch.Tensor): Tensor of shape [B, N, 2, 4], where the last dimension is [x1, y1, x2, y2] and dim=2 indexes subject and object. Returns: torch.Tensor: Union action boxes of shape [B, N, 4] """ subj_boxes = boxes[:, :, 0, :] # [B, N, 4] obj_boxes = boxes[:, :, 1, :] # [B, N, 4] x1 = torch.min(subj_boxes[:, :, 0], obj_boxes[:, :, 0]) y1 = torch.min(subj_boxes[:, :, 1], obj_boxes[:, :, 1]) x2 = torch.max(subj_boxes[:, :, 2], obj_boxes[:, :, 2]) y2 = torch.max(subj_boxes[:, :, 3], obj_boxes[:, :, 3]) return torch.stack([x1, y1, x2, y2], dim=-1) class GroundingEncoder(ModelMixin, ConfigMixin): @register_to_config def __init__(self, hidden_size=512, text_encoder_dim=4096, max_hoi_seq=32, pos_embed_dim=32, role_embed_dim=32, fourier_freq=32, init_logit=-5.0, n_roles=3, role_std=0.02, mlp_out_std=3e-4): super(GroundingEncoder, self).__init__() self.text_encoder_dim = text_encoder_dim self.role_emb = nn.Embedding(n_roles, role_embed_dim) nn.init.kaiming_normal_(self.role_emb.weight, nonlinearity="linear") self.pos_embed_dim = pos_embed_dim self.max_hoi_seq = max_hoi_seq self.fourier_freq = fourier_freq self.fourier_dim = fourier_freq * 2 * 4 position = torch.arange(max_hoi_seq).unsqueeze(1) div_term = torch.exp(torch.arange(0, pos_embed_dim, 2) * (-math.log(10000.0) / pos_embed_dim)) pe = torch.zeros(max_hoi_seq, pos_embed_dim) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer("pe", pe) self.norm = nn.LayerNorm(text_encoder_dim) self.mlp = nn.Sequential( nn.Linear(text_encoder_dim + self.fourier_dim + pos_embed_dim + role_embed_dim, hidden_size), nn.SiLU(), nn.Linear(hidden_size, text_encoder_dim), ) self.gate = nn.Parameter(torch.tensor(0.0, dtype=torch.float32)) def forward(self, x, role_ids: torch.Tensor, idx_ids: torch.Tensor, boxes: torch.Tensor): # x = prompt_embeds: [B, T, D] # role_ids: [hoi_seq_len] # idx_ids: [hoi_seq_len] B, _, _ = x.shape h = self.norm(x) role_pe = self.role_emb(role_ids) role_pe = role_pe.unsqueeze(0).expand(B, -1, -1) if self.training: idx_ids = (idx_ids + torch.randint(0, self.max_hoi_seq, (B, 1), device=x.device) ) % self.max_hoi_seq idx_pe = self.pe[idx_ids] else: idx_pe = self.pe[idx_ids].unsqueeze(0).expand(B, -1, -1) idx_pe = idx_pe - idx_pe.mean(-1, keepdim=True) boxes_features = get_fourier_embeds_from_boundingbox(self.fourier_freq, boxes) # [B, T, D=256] h = torch.cat([h, boxes_features, idx_pe, role_pe], dim=-1) x = x + self.gate.tanh() * self.mlp(h) # [B, T, D=4096] return x if __name__ == "__main__": # Example usage boxes = [[[0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6]], [[0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6]],] box_labels = [["person", "dog"], ["person", "cat", "person", "dog", "person", "cat", "person", "dog", "person", "cat", "person", "dog", "person", "cat", "person", "dog",]] hoi_labels = [["walking"], ["walking", "running", "walking", "running", "walking", "running", "walking", "running",]] tokenizer = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", subfolder="tokenizer_2") text_encoder = T5EncoderModel.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", subfolder="text_encoder_2") # processed_boxes, box_prompt_embeds = GroundingInput.preprocess_hoi(boxes, box_labels, hoi_labels, tokenizer, text_encoder) # # print(processed_boxes) # print(f"processed_boxes: {len(processed_boxes)}x{len(processed_boxes[0])}x{len(processed_boxes[0][0])}x4") # print(f"processed_boxes: {len(processed_boxes)}x{len(processed_boxes[1])}x{len(processed_boxes[1][0])}xD") # # print(box_prompt_embeds) # print(f"box_prompt_embeds: {len(box_prompt_embeds)}x{len(box_prompt_embeds[0])}x{len(box_prompt_embeds[0][0])}xD") # print(f"box_prompt_embeds: {len(box_prompt_embeds)}x{len(box_prompt_embeds[1])}x{len(box_prompt_embeds[1][0])}xD") # # Prepare HOI input # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_hoi(box_prompt_embeds, processed_boxes) # print(f"out_embeds: {out_embeds.shape}") # [B, N_max * M * T_max, D] # print(f"out_boxes: {out_boxes.shape}") # [B, N_max * M * T_max, 4] # print(f"out_text_ids: {out_text_ids.shape}") # [B, N_max * M * T_max, 3] # # Vanilla T2I # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_pipeline_input(tokenizer, text_encoder, # boxes=None, box_labels=None, hoi_labels=None) # print(f"out_embeds: {out_embeds}") # None # print(f"out_boxes: {out_boxes}") # None # print(f"out_text_ids: {out_text_ids}") # None # # HOI text control # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_pipeline_input(tokenizer, text_encoder, # boxes=None, box_labels=box_labels, hoi_labels=hoi_labels) # print(f"out_embeds: {out_embeds.shape}") # print(f"out_boxes: {out_boxes.shape}") # print(f"out_text_ids: {out_text_ids.shape}") # # HOI box control # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_pipeline_input(tokenizer, text_encoder, # boxes=boxes, box_labels=box_labels, hoi_labels=hoi_labels) # print(f"out_embeds: {out_embeds.shape}") # print(f"out_boxes: {out_boxes.shape}") # None # print(f"out_text_ids: {out_text_ids.shape}") # None # # Box control with no HOI labels # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_pipeline_input(tokenizer, text_encoder, # boxes=boxes, box_labels=box_labels, hoi_labels=None) # print(f"out_embeds: {out_embeds.shape}") # [B, N_max * M * T_max, D] # print(f"out_boxes: {out_boxes.shape}") # [B, N_max * M * T_max, 4] # print(f"out_text_ids: {out_text_ids.shape}") # [B, N_max * M * T_max, 3] # # Box control with no boxes # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_pipeline_input(tokenizer, text_encoder, # boxes=None, box_labels=box_labels, hoi_labels=None) # print(f"out_embeds: {out_embeds.shape}") # [B, N_max * M * T_max, D] # print(f"out_boxes: {out_boxes.shape}") # [B, N_max * M * T_max, 4] # print(f"out_text_ids: {out_text_ids.shape}") # [B, N_max * M * T_max, 3] #### NEW TEST #### # processed_boxes, box_prompt_embeds = GroundingInput.preprocess_box(boxes, box_labels, tokenizer, text_encoder) # print(f"processed_boxes: {len(processed_boxes)}x{len(processed_boxes[0])}x{len(processed_boxes[0][0])}x4") # print(f"box_prompt_embeds: {len(box_prompt_embeds)}x{len(box_prompt_embeds[0])}x{len(box_prompt_embeds[0][0])}xD") # D is the embedding dimension # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_box(box_prompt_embeds, processed_boxes) # print(f"out_embeds: {out_embeds.shape}") # [B, N_max * M * T_max, D] # print(f"out_boxes: {out_boxes.shape}") # [B, N_max * M * T_max, 4] # print(f"out_text_ids: {out_text_ids.shape}") # [B, N_max * M * T_max, 3] processed_boxes, box_prompt_embeds = GroundingInput.preprocess_hoi(boxes, box_labels, hoi_labels, tokenizer, text_encoder) out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_hoi(box_prompt_embeds, processed_boxes) attn_mask, _ = GroundingInput.prepare_attention_mask(out_text_ids, out_boxes, use_union_action_mask=True) pass #### TEST RANDOM DROP HOI #### # from datasets import load_from_disk # # ds = load_from_disk("data/synthesis_edits_kontext_9") # ds = load_from_disk("data/hicodet_kontext_dataset") # ds.set_format(type='torch', columns=['hois', 'boxes']) # # sample = ds[0::512] # batched for hoi edits datasets # sample = ds[0::16000] # batched for hicodet_kontext_dataset # boxes = sample['boxes'] # hois = sample['hois'] # objects = [None] # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_train_input(tokenizer, text_encoder, # boxes=boxes, hois=hois, objects=objects, # random_drop_boxes=0, random_drop_hois=1) # for b in range(len(boxes)): # n = len(boxes[b]) # assert (boxes[b] == out_boxes[b, ::64][:n*2].reshape(boxes[b].shape)).all() # boxes are [B, N, 1, 4], out_boxes is [B, T_max, 4] #### TEST LIMIT MAX HOI #### # boxes = [[[0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6]], # [[0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], # [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], # [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6], [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6]],] # box_labels = [["person", "dog"], # ["person", "cat", "person", "dog", # "person", "cat", "person", "dog", # "person", "cat", "person", "dog",]] # hoi_labels = [["walking"], # ["walking", "running", # "walking", "running", # "walking", "running",]] # processed_boxes, box_prompt_embeds = GroundingInput.preprocess_hoi(boxes, box_labels, hoi_labels, tokenizer, text_encoder) # assert len(processed_boxes[1]) == HOI_N_MAX # #### TEST LIMIT MAX BOX #### # processed_boxes, box_prompt_embeds = GroundingInput.preprocess_box(boxes, box_labels, tokenizer, text_encoder) # assert len(processed_boxes[1]) == BOX_N_MAX # #### TEST FIXED EMBED for HOI ROLE #### # processed_boxes, box_prompt_embeds = GroundingInput.preprocess_hoi(boxes, box_labels, hoi_labels, tokenizer, text_encoder) # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_hoi(box_prompt_embeds, processed_boxes) # grounding_encoder = GroundingEncoder() # mlp_prompt_embeds = grounding_encoder(out_embeds, out_text_ids[:, 2]) #### TEST FIXED EMBED for BOX ROLE #### # processed_boxes, box_prompt_embeds = GroundingInput.preprocess_box(boxes, box_labels, tokenizer, text_encoder) # out_embeds, out_boxes, out_text_ids = GroundingInput.prepare_box(box_prompt_embeds, processed_boxes) # new_text_ids = GroundingInput.get_rope_ids(out_text_ids, img_width=64, img_height=48, cond_width=64, cond_height=48) # grounding_encoder = GroundingEncoder() # grounding_encoder.eval() # mlp_prompt_embeds = grounding_encoder(out_embeds, out_text_ids[:, 2], out_text_ids[:, 1]) # grounding_encoder.train() # mlp_prompt_embeds = grounding_encoder(out_embeds, out_text_ids[:, 2], out_text_ids[:, 1]) grounding_encoder = GroundingEncoder() grounding_encoder.eval() mlp_prompt_embeds = grounding_encoder(out_embeds, out_text_ids[:, 2], out_text_ids[:, 1], boxes=out_boxes) # mix_boxes: [B, N, M, 4] # mix_boxes = [ # [ # b=0 # [ # n=0 # [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6] # ], # [ # n=1 # None # ], # ], # [ # b=1 # [ # n=0 # None, None # ], # [ # n=1 # [0, 0, 0.2, 0.2] # ], # ] # ] # # mix_box_labels: [B, N, M] # mix_box_labels = [ # [ # b=0 # ["person", "dog"], # n=0 # ["cat"] # n=1 # ], # [ # b=1 # ["person", "dog"], # ["watermelon"] # ] # ] # # mix_box_labels: [B, N] # mix_hoi_labels = [ # [ # b=0 # "walking", None, # ], # [ # b=1 # "hugging", None, # ], # ] # mix boxes: [B=1, N, M, 4] mix_boxes = [ [ # b=0 [ # n=0 [0, 0, 0.2, 0.2], [0.4, 0.4, 0.6, 0.6] ], [ # n=1 None ], ] ] # mix_box_labels: [B, N, M] mix_box_labels = [ [ # b=0 ["person", "dog"], # n=0 ["cat"] # n=1 ] ] # mix_box_labels: [B, N] mix_hoi_labels = [ [ # b=0 "walking", None, ] ] GroundingInput.prepare_mixed_pipeline_input(tokenizer=tokenizer, text_encoder=text_encoder, mix_boxes=mix_boxes, mix_box_labels=mix_box_labels, mix_hoi_labels=mix_hoi_labels, hoi_seq_len=64, max_box=9) pass