Function omni_attn_mask_naive

show-o2/models/omni_attention.py:56–70

(B, LEN, modalities, device, inverted=True)

def omni_attn_mask_naive(B, LEN, modalities, device, inverted=True):
    attention_mask = torch.tril(torch.ones((B, 1, LEN, LEN), dtype=torch.long)).to(device)
    for b in range(B):
        modality_batch = modalities[b]
        for offset, length in modality_batch:
            attention_mask[b, :, offset:offset + length, offset:offset + length] = 1

    if inverted:
        inverted_attention_mask = 1 - attention_mask
        inverted_attention_mask = inverted_attention_mask.masked_fill(
            inverted_attention_mask.to(torch.bool), torch.iinfo(torch.long).min
        )
        return inverted_attention_mask
    else:
        return attention_mask


def full_attn_mask_naive(B, LEN, device, inverted=True):