RoPE for video tokens with 3D structure. Args: embed_dim: (`int`): The embedding dimension size, corresponding to hidden_size_head. crops_coords (`Tuple[int]`): The top-left and bottom-right coordinates of the crop. grid_size (`Tuple[int]`): The grid siz
(
embed_dim, crops_coords, grid_size, temporal_size, theta: int = 10000, use_real: bool = True
)
| 443 | |
| 444 | |
| 445 | def get_3d_rotary_pos_embed( |
| 446 | embed_dim, crops_coords, grid_size, temporal_size, theta: int = 10000, use_real: bool = True |
| 447 | ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: |
| 448 | """ |
| 449 | RoPE for video tokens with 3D structure. |
| 450 | |
| 451 | Args: |
| 452 | embed_dim: (`int`): |
| 453 | The embedding dimension size, corresponding to hidden_size_head. |
| 454 | crops_coords (`Tuple[int]`): |
| 455 | The top-left and bottom-right coordinates of the crop. |
| 456 | grid_size (`Tuple[int]`): |
| 457 | The grid size of the spatial positional embedding (height, width). |
| 458 | temporal_size (`int`): |
| 459 | The size of the temporal dimension. |
| 460 | theta (`float`): |
| 461 | Scaling factor for frequency computation. |
| 462 | |
| 463 | Returns: |
| 464 | `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`. |
| 465 | """ |
| 466 | if use_real is not True: |
| 467 | raise ValueError(" `use_real = False` is not currently supported for get_3d_rotary_pos_embed") |
| 468 | start, stop = crops_coords |
| 469 | grid_size_h, grid_size_w = grid_size |
| 470 | grid_h = np.linspace(start[0], stop[0], grid_size_h, endpoint=False, dtype=np.float32) |
| 471 | grid_w = np.linspace(start[1], stop[1], grid_size_w, endpoint=False, dtype=np.float32) |
| 472 | grid_t = np.linspace(0, temporal_size, temporal_size, endpoint=False, dtype=np.float32) |
| 473 | |
| 474 | # Compute dimensions for each axis |
| 475 | dim_t = embed_dim // 4 |
| 476 | dim_h = embed_dim // 8 * 3 |
| 477 | dim_w = embed_dim // 8 * 3 |
| 478 | |
| 479 | # Temporal frequencies |
| 480 | freqs_t = get_1d_rotary_pos_embed(dim_t, grid_t, use_real=True) |
| 481 | # Spatial frequencies for height and width |
| 482 | freqs_h = get_1d_rotary_pos_embed(dim_h, grid_h, use_real=True) |
| 483 | freqs_w = get_1d_rotary_pos_embed(dim_w, grid_w, use_real=True) |
| 484 | |
| 485 | # BroadCast and concatenate temporal and spaial frequencie (height and width) into a 3d tensor |
| 486 | def combine_time_height_width(freqs_t, freqs_h, freqs_w): |
| 487 | freqs_t = freqs_t[:, None, None, :].expand( |
| 488 | -1, grid_size_h, grid_size_w, -1 |
| 489 | ) # temporal_size, grid_size_h, grid_size_w, dim_t |
| 490 | freqs_h = freqs_h[None, :, None, :].expand( |
| 491 | temporal_size, -1, grid_size_w, -1 |
| 492 | ) # temporal_size, grid_size_h, grid_size_2, dim_h |
| 493 | freqs_w = freqs_w[None, None, :, :].expand( |
| 494 | temporal_size, grid_size_h, -1, -1 |
| 495 | ) # temporal_size, grid_size_h, grid_size_2, dim_w |
| 496 | |
| 497 | freqs = torch.cat( |
| 498 | [freqs_t, freqs_h, freqs_w], dim=-1 |
| 499 | ) # temporal_size, grid_size_h, grid_size_w, (dim_t + dim_h + dim_w) |
| 500 | freqs = freqs.view( |
| 501 | temporal_size * grid_size_h * grid_size_w, -1 |
| 502 | ) # (temporal_size * grid_size_h * grid_size_w), (dim_t + dim_h + dim_w) |
no test coverage detected