(img_features, depth_features)
| 7 | return torch.cat((-x2, x1), dim=-1) |
| 8 | |
| 9 | def depth_rope_encoding(img_features, depth_features): |
| 10 | orig_dtype = img_features.dtype |
| 11 | depth_features = torch.cat(depth_features,dim=0) |
| 12 | depth_features = depth_features.reshape(depth_features.shape[0],-1) |
| 13 | B, L, dim = img_features.shape |
| 14 | assert dim % 2 == 0, "wrong dim" |
| 15 | |
| 16 | theta = 10000.0 |
| 17 | seqlen = 10000 |
| 18 | inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / (dim))) |
| 19 | seq = torch.arange(seqlen, device=inv_freq.device, dtype=inv_freq.dtype) |
| 20 | freqs = torch.outer(seq, inv_freq) |
| 21 | |
| 22 | depth_features = depth_features.clone() |
| 23 | depth_features = (depth_features * seqlen).to(torch.long) |
| 24 | |
| 25 | rotary_pos_emb = freqs[depth_features] |
| 26 | rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) |
| 27 | cos = rotary_pos_emb.cos() |
| 28 | sin = rotary_pos_emb.sin() |
| 29 | img_features = (img_features * cos) + (rotate_half(img_features) * sin) |
| 30 | img_features = img_features.to(orig_dtype) |
| 31 | |
| 32 | return img_features |
| 33 | |
| 34 | if __name__ == "__main__": |
| 35 | img_features = torch.rand(2, 10, 16) # Example image features |
no test coverage detected