Spaces:
Running on Zero
Running on Zero
Download src/model/encoder/unimatch/mv_transformer.py from hugging-apps/codecsplat: direct link, hf CLI and curl.
- Browser
- Download file 22.2 kB
-
https://huggingface.co/spaces/hugging-apps/codecsplat/resolve/main/src/model/encoder/unimatch/mv_transformer.py
- Command line
-
hf download hf://spaces/hugging-apps/codecsplat/src/model/encoder/unimatch/mv_transformer.py
-
curl -L -o mv_transformer.py https://huggingface.co/spaces/hugging-apps/codecsplat/resolve/main/src/model/encoder/unimatch/mv_transformer.py
22.2 kB
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from .utils import split_feature, merge_splits | |
| def gather_view_tensor(tensor, view_indices): | |
| view_indices = view_indices.to(device=tensor.device, dtype=torch.long) | |
| batch_indices = torch.arange(tensor.size(0), device=tensor.device) | |
| batch_indices = batch_indices.reshape( | |
| tensor.size(0), *([1] * (view_indices.dim() - 1)) | |
| ) | |
| return tensor[batch_indices, view_indices] | |
| def non_self_view_indices(num_views, device): | |
| view_ids = torch.arange(num_views, device=device) | |
| return torch.stack( | |
| [torch.cat((view_ids[:i], view_ids[i + 1 :])) for i in range(num_views)], | |
| dim=0, | |
| ) | |
| def target_view_indices(batch_size, num_views, nn_matrix, device): | |
| if nn_matrix is not None: | |
| return nn_matrix[..., 1:].to(device=device, dtype=torch.long) | |
| return non_self_view_indices(num_views, device).unsqueeze(0).expand( | |
| batch_size, -1, -1 | |
| ) | |
| def flatten_batch_view(tensor): | |
| # Batch-major flatten: [B, V, ...] -> [B*V, ...]. | |
| return tensor.flatten(0, 1) | |
| def flatten_view_batch(tensor): | |
| # View-major flatten: [B, V, ...] -> [V*B, ...]. | |
| return tensor.transpose(0, 1).reshape( | |
| tensor.size(0) * tensor.size(1), *tensor.shape[2:] | |
| ) | |
| def unflatten_view_batch(tensor, batch_size, num_views): | |
| # Inverse of flatten_view_batch. | |
| return tensor.reshape(num_views, batch_size, *tensor.shape[1:]).transpose(0, 1) | |
| def single_head_full_attention(q, k, v): | |
| # q, k, v: [B, L, C] | |
| assert q.dim() == k.dim() == v.dim() == 3 | |
| return F.scaled_dot_product_attention( | |
| q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1) | |
| ).squeeze(1) | |
| def generate_shift_window_attn_mask( | |
| input_resolution, | |
| window_size_h, | |
| window_size_w, | |
| shift_size_h, | |
| shift_size_w, | |
| device=torch.device("cuda"), | |
| ): | |
| # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py | |
| # Build the shifted-window attention mask used by Swin-style local attention. | |
| h, w = input_resolution | |
| img_mask = torch.zeros((1, h, w, 1)).to(device) # 1 H W 1 | |
| h_slices = ( | |
| slice(0, -window_size_h), | |
| slice(-window_size_h, -shift_size_h), | |
| slice(-shift_size_h, None), | |
| ) | |
| w_slices = ( | |
| slice(0, -window_size_w), | |
| slice(-window_size_w, -shift_size_w), | |
| slice(-shift_size_w, None), | |
| ) | |
| cnt = 0 | |
| for h in h_slices: | |
| for w in w_slices: | |
| img_mask[:, h, w, :] = cnt | |
| cnt += 1 | |
| mask_windows = split_feature( | |
| img_mask, num_splits=input_resolution[-1] // window_size_w, channel_last=True | |
| ) | |
| mask_windows = mask_windows.view(-1, window_size_h * window_size_w) | |
| attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) | |
| attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( | |
| attn_mask == 0, float(0.0) | |
| ) | |
| return attn_mask | |
| def single_head_split_window_attention( | |
| q, | |
| k, | |
| v, | |
| num_splits=1, | |
| with_shift=False, | |
| h=None, | |
| w=None, | |
| attn_mask=None, | |
| ): | |
| # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py | |
| # q, k, v: [B, L, C] for self-attention or pairwise cross-attention. | |
| # For multi-view cross-attention, q: [B, L, C], k/v: [B, M, L, C]. | |
| # Multi-view cross-attention, where M is the number of target views per query. | |
| if not (q.dim() == k.dim() == v.dim() == 3): | |
| assert k.dim() == v.dim() == 4 | |
| assert h is not None and w is not None | |
| assert q.size(1) == h * w | |
| m = k.size(1) | |
| b, _, c = q.size() | |
| b_new = b * num_splits * num_splits | |
| window_size_h = h // num_splits | |
| window_size_w = w // num_splits | |
| q = q.view(b, h, w, c) # [B, H, W, C] | |
| k = k.view(b, m, h, w, c) # [B, M, H, W, C] | |
| v = v.view(b, m, h, w, c) | |
| if with_shift: | |
| assert attn_mask is not None # compute once | |
| shift_size_h = window_size_h // 2 | |
| shift_size_w = window_size_w // 2 | |
| q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) | |
| k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(2, 3)) | |
| v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(2, 3)) | |
| q = split_feature( | |
| q, num_splits=num_splits, channel_last=True | |
| ) # [B*K*K, H/K, W/K, C] | |
| k = split_feature( | |
| k.permute(0, 2, 3, 4, 1).reshape(b, h, w, -1), | |
| num_splits=num_splits, | |
| channel_last=True, | |
| ) # [B*K*K, H/K, W/K, C*M] | |
| v = split_feature( | |
| v.permute(0, 2, 3, 4, 1).reshape(b, h, w, -1), | |
| num_splits=num_splits, | |
| channel_last=True, | |
| ) # [B*K*K, H/K, W/K, C*M] | |
| k = ( | |
| k.view(b_new, h // num_splits, w // num_splits, c, m) | |
| .permute(0, 1, 2, 4, 3) | |
| .reshape(b_new, -1, c) | |
| ) # [B*K*K, H/K*W/K*M, C] | |
| v = ( | |
| v.view(b_new, h // num_splits, w // num_splits, c, m) | |
| .permute(0, 1, 2, 4, 3) | |
| .reshape(b_new, -1, c) | |
| ) # [B*K*K, H/K*W/K*M, C] | |
| attn_bias = None | |
| if with_shift: | |
| attn_bias = attn_mask.repeat(b, 1, m).unsqueeze(1).to(q.dtype) | |
| out = F.scaled_dot_product_attention( | |
| q.view(b_new, -1, c).unsqueeze(1), | |
| k.unsqueeze(1), | |
| v.unsqueeze(1), | |
| attn_mask=attn_bias, | |
| ).squeeze(1) | |
| out = merge_splits( | |
| out.view(b_new, h // num_splits, w // num_splits, c), | |
| num_splits=num_splits, | |
| channel_last=True, | |
| ) # [B, H, W, C] | |
| # shift back | |
| if with_shift: | |
| out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2)) | |
| out = out.view(b, -1, c) | |
| else: | |
| # Self-attention or pairwise cross-attention. | |
| assert q.dim() == k.dim() == v.dim() == 3 | |
| assert h is not None and w is not None | |
| assert q.size(1) == h * w | |
| b, _, c = q.size() | |
| b_new = b * num_splits * num_splits | |
| window_size_h = h // num_splits | |
| window_size_w = w // num_splits | |
| q = q.view(b, h, w, c) # [B, H, W, C] | |
| k = k.view(b, h, w, c) | |
| v = v.view(b, h, w, c) | |
| if with_shift: | |
| assert attn_mask is not None # compute once | |
| shift_size_h = window_size_h // 2 | |
| shift_size_w = window_size_w // 2 | |
| q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) | |
| k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) | |
| v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) | |
| q = split_feature( | |
| q, num_splits=num_splits, channel_last=True | |
| ) # [B*K*K, H/K, W/K, C] | |
| k = split_feature(k, num_splits=num_splits, channel_last=True) | |
| v = split_feature(v, num_splits=num_splits, channel_last=True) | |
| attn_bias = None | |
| if with_shift: | |
| attn_bias = attn_mask.repeat(b, 1, 1).unsqueeze(1).to(q.dtype) | |
| out = F.scaled_dot_product_attention( | |
| q.view(b_new, -1, c).unsqueeze(1), | |
| k.view(b_new, -1, c).unsqueeze(1), | |
| v.view(b_new, -1, c).unsqueeze(1), | |
| attn_mask=attn_bias, | |
| ).squeeze(1) | |
| out = merge_splits( | |
| out.view(b_new, h // num_splits, w // num_splits, c), | |
| num_splits=num_splits, | |
| channel_last=True, | |
| ) # [B, H, W, C] | |
| # shift back | |
| if with_shift: | |
| out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2)) | |
| out = out.view(b, -1, c) | |
| return out | |
| def multi_head_split_window_attention( | |
| q, | |
| k, | |
| v, | |
| num_splits=1, | |
| with_shift=False, | |
| h=None, | |
| w=None, | |
| attn_mask=None, | |
| num_head=1, | |
| ): | |
| """Windowed multi-head scaled dot-product attention. | |
| Args: | |
| q: [B, L, C] | |
| k: [B, S, C] | |
| v: [B, S, C] | |
| Returns: | |
| out: [B, L, C] | |
| """ | |
| assert h is not None and w is not None | |
| assert q.size(1) == h * w | |
| b, _, c = q.size() | |
| b_new = b * num_splits * num_splits | |
| window_size_h = h // num_splits | |
| window_size_w = w // num_splits | |
| q = q.view(b, h, w, c) # [B, H, W, C] | |
| k = k.view(b, h, w, c) | |
| v = v.view(b, h, w, c) | |
| assert c % num_head == 0 | |
| if with_shift: | |
| assert attn_mask is not None # compute once | |
| shift_size_h = window_size_h // 2 | |
| shift_size_w = window_size_w // 2 | |
| q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) | |
| k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) | |
| v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) | |
| q = split_feature(q, num_splits=num_splits) # [B*K*K, H/K, W/K, C] | |
| k = split_feature(k, num_splits=num_splits) | |
| v = split_feature(v, num_splits=num_splits) | |
| # Split channels into attention heads for PyTorch SDPA. | |
| head_dim = c // num_head | |
| q = q.view(b_new, -1, num_head, head_dim).transpose(1, 2) | |
| k = k.view(b_new, -1, num_head, head_dim).transpose(1, 2) | |
| v = v.view(b_new, -1, num_head, head_dim).transpose(1, 2) | |
| attn_bias = None | |
| if with_shift: | |
| attn_bias = attn_mask.unsqueeze(1).repeat(b, num_head, 1, 1).to(q.dtype) | |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias) | |
| out = merge_splits( | |
| out.transpose(1, 2).reshape(b_new, h // num_splits, w // num_splits, c), | |
| num_splits=num_splits, | |
| ) # [B, H, W, C] | |
| # shift back | |
| if with_shift: | |
| out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2)) | |
| out = out.view(b, -1, c) | |
| return out | |
| class TransformerLayer(nn.Module): | |
| def __init__( | |
| self, | |
| d_model=256, | |
| nhead=1, | |
| attention_type="swin", | |
| no_ffn=False, | |
| ffn_dim_expansion=4, | |
| with_shift=False, | |
| add_per_view_attn=False, | |
| **kwargs, | |
| ): | |
| super(TransformerLayer, self).__init__() | |
| self.dim = d_model | |
| self.nhead = nhead | |
| self.attention_type = attention_type | |
| self.no_ffn = no_ffn | |
| self.add_per_view_attn = add_per_view_attn | |
| self.with_shift = with_shift | |
| # Linear projections used by both single-head and multi-head attention. | |
| self.q_proj = nn.Linear(d_model, d_model, bias=False) | |
| self.k_proj = nn.Linear(d_model, d_model, bias=False) | |
| self.v_proj = nn.Linear(d_model, d_model, bias=False) | |
| self.merge = nn.Linear(d_model, d_model, bias=False) | |
| self.norm1 = nn.LayerNorm(d_model) | |
| # no ffn after self-attn, with ffn after cross-attn | |
| if not self.no_ffn: | |
| in_channels = d_model * 2 | |
| self.mlp = nn.Sequential( | |
| nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False), | |
| nn.GELU(), | |
| nn.Linear(in_channels * ffn_dim_expansion, d_model, bias=False), | |
| ) | |
| self.norm2 = nn.LayerNorm(d_model) | |
| def forward( | |
| self, | |
| source, | |
| target, | |
| height=None, | |
| width=None, | |
| shifted_window_attn_mask=None, | |
| attn_num_splits=None, | |
| **kwargs, | |
| ): | |
| if "attn_type" in kwargs: | |
| attn_type = kwargs["attn_type"] | |
| else: | |
| attn_type = self.attention_type | |
| # source: [B, L, C] | |
| # target: [B, L, C] for pairwise attention or [B, M, L, C] for multi-view. | |
| query, key, value = source, target, target | |
| # Project tokens before dispatching to local/global attention. | |
| query = self.q_proj(query) # [B, L, C] | |
| key = self.k_proj(key) # [B, L, C] or [B, M, L, C] | |
| value = self.v_proj(value) # [B, L, C] or [B, M, L, C] | |
| if attn_type == "swin" and attn_num_splits > 1: | |
| if self.nhead > 1: | |
| message = multi_head_split_window_attention( | |
| query, | |
| key, | |
| value, | |
| num_splits=attn_num_splits, | |
| with_shift=self.with_shift, | |
| h=height, | |
| w=width, | |
| attn_mask=shifted_window_attn_mask, | |
| num_head=self.nhead, | |
| ) | |
| else: | |
| if self.add_per_view_attn: | |
| assert query.dim() == 3 and key.dim() == 4 and value.dim() == 4 | |
| b, l, c = query.size() | |
| query = query.unsqueeze(1).expand( | |
| -1, key.size(1), -1, -1 | |
| ).reshape(-1, l, c) # [B*M, L, C] | |
| key = key.reshape(-1, l, c) | |
| value = value.reshape(-1, l, c) | |
| message = single_head_split_window_attention( | |
| query, | |
| key, | |
| value, | |
| num_splits=attn_num_splits, | |
| with_shift=self.with_shift, | |
| h=height, | |
| w=width, | |
| attn_mask=shifted_window_attn_mask, | |
| ) | |
| # Sum per-target-view messages back to [B, L, C]. | |
| message = message.view(b, -1, l, c).sum(1) | |
| else: | |
| message = single_head_split_window_attention( | |
| query, | |
| key, | |
| value, | |
| num_splits=attn_num_splits, | |
| with_shift=self.with_shift, | |
| h=height, | |
| w=width, | |
| attn_mask=shifted_window_attn_mask, | |
| ) | |
| else: | |
| message = single_head_full_attention(query, key, value) # [B, L, C] | |
| message = self.merge(message) # [B, L, C] | |
| message = self.norm1(message) | |
| if not self.no_ffn: | |
| message = self.mlp(torch.cat([source, message], dim=-1)) | |
| message = self.norm2(message) | |
| return source + message | |
| class TransformerBlock(nn.Module): | |
| """self attention + cross attention + FFN""" | |
| def __init__( | |
| self, | |
| d_model=256, | |
| nhead=1, | |
| attention_type="swin", | |
| ffn_dim_expansion=4, | |
| with_shift=False, | |
| add_per_view_attn=False, | |
| no_cross_attn=False, | |
| **kwargs, | |
| ): | |
| super(TransformerBlock, self).__init__() | |
| self.no_cross_attn = no_cross_attn | |
| if no_cross_attn: | |
| self.self_attn = TransformerLayer( | |
| d_model=d_model, | |
| nhead=nhead, | |
| attention_type=attention_type, | |
| ffn_dim_expansion=ffn_dim_expansion, | |
| with_shift=with_shift, | |
| add_per_view_attn=add_per_view_attn, | |
| ) | |
| else: | |
| self.self_attn = TransformerLayer( | |
| d_model=d_model, | |
| nhead=nhead, | |
| attention_type=attention_type, | |
| no_ffn=True, | |
| ffn_dim_expansion=ffn_dim_expansion, | |
| with_shift=with_shift, | |
| ) | |
| self.cross_attn_ffn = TransformerLayer( | |
| d_model=d_model, | |
| nhead=nhead, | |
| attention_type=attention_type, | |
| ffn_dim_expansion=ffn_dim_expansion, | |
| with_shift=with_shift, | |
| add_per_view_attn=add_per_view_attn, | |
| ) | |
| def forward( | |
| self, | |
| source, | |
| target, | |
| height=None, | |
| width=None, | |
| shifted_window_attn_mask=None, | |
| attn_num_splits=None, | |
| **kwargs, | |
| ): | |
| # source, target: [B, L, C] | |
| # self attention | |
| source = self.self_attn( | |
| source, | |
| source, | |
| height=height, | |
| width=width, | |
| shifted_window_attn_mask=shifted_window_attn_mask, | |
| attn_num_splits=attn_num_splits, | |
| **kwargs, | |
| ) | |
| if self.no_cross_attn: | |
| return source | |
| # cross attention and ffn | |
| source = self.cross_attn_ffn( | |
| source, | |
| target, | |
| height=height, | |
| width=width, | |
| shifted_window_attn_mask=shifted_window_attn_mask, | |
| attn_num_splits=attn_num_splits, | |
| **kwargs, | |
| ) | |
| return source | |
| def batch_features(features, nn_matrix=None): | |
| # Construct transformer inputs from tensor views. | |
| # features: [B, V, C, H, W] or [B, V, H*W, C]. | |
| # Returns view-major query/target tensors: [V*B, ...] and [V*B, M, ...]. | |
| assert features.dim() in (4, 5) | |
| batch_size, num_views = features.shape[:2] | |
| indices = target_view_indices( | |
| batch_size, num_views, nn_matrix, features.device | |
| ) | |
| q = flatten_view_batch(features) | |
| kv = gather_view_tensor(features, indices) | |
| kv = flatten_view_batch(kv) | |
| return q, kv | |
| class MultiViewFeatureTransformer(nn.Module): | |
| def __init__( | |
| self, | |
| num_layers=6, | |
| d_model=128, | |
| nhead=1, | |
| attention_type="swin", | |
| ffn_dim_expansion=4, | |
| add_per_view_attn=False, | |
| no_cross_attn=False, | |
| **kwargs, | |
| ): | |
| super(MultiViewFeatureTransformer, self).__init__() | |
| self.attention_type = attention_type | |
| self.d_model = d_model | |
| self.nhead = nhead | |
| self.layers = nn.ModuleList( | |
| [ | |
| TransformerBlock( | |
| d_model=d_model, | |
| nhead=nhead, | |
| attention_type=attention_type, | |
| ffn_dim_expansion=ffn_dim_expansion, | |
| with_shift=( | |
| True if attention_type == "swin" and i % 2 == 1 else False | |
| ), | |
| add_per_view_attn=add_per_view_attn, | |
| no_cross_attn=no_cross_attn, | |
| ) | |
| for i in range(num_layers) | |
| ] | |
| ) | |
| for p in self.parameters(): | |
| if p.dim() > 1: | |
| nn.init.xavier_uniform_(p) | |
| # zero init layers beyond 6 | |
| if num_layers > 6: | |
| for i in range(6, num_layers): | |
| self.layers[i].self_attn.norm1.weight.data.zero_() | |
| self.layers[i].self_attn.norm1.bias.data.zero_() | |
| self.layers[i].cross_attn_ffn.norm2.weight.data.zero_() | |
| self.layers[i].cross_attn_ffn.norm2.bias.data.zero_() | |
| def forward( | |
| self, | |
| multi_view_features, | |
| attn_num_splits=None, | |
| **kwargs, | |
| ): | |
| nn_matrix = kwargs.pop("nn_matrix", None) | |
| # multi_view_features: [B, V, C, H, W]. | |
| assert multi_view_features.dim() == 5 | |
| b, num_views, c, h, w = multi_view_features.shape | |
| assert self.d_model == c | |
| if self.attention_type == "swin" and attn_num_splits > 1: | |
| # global and refine use different number of splits | |
| window_size_h = h // attn_num_splits | |
| window_size_w = w // attn_num_splits | |
| # compute attn mask once | |
| shifted_window_attn_mask = generate_shift_window_attn_mask( | |
| input_resolution=(h, w), | |
| window_size_h=window_size_h, | |
| window_size_w=window_size_w, | |
| shift_size_h=window_size_h // 2, | |
| shift_size_w=window_size_w // 2, | |
| device=multi_view_features.device, | |
| ) # [K*K, H/K*W/K, H/K*W/K] | |
| else: | |
| shifted_window_attn_mask = None | |
| # View-major query/target tensors: [V*B, C, H, W], [V*B, M, C, H, W]. | |
| concat0, concat1 = batch_features(multi_view_features, nn_matrix=nn_matrix) | |
| concat0 = concat0.reshape(num_views * b, c, -1).permute( | |
| 0, 2, 1 | |
| ) # [V*B, H*W, C] | |
| c1_v = num_views - 1 if nn_matrix is None else nn_matrix.shape[-1] - 1 | |
| concat1 = concat1.reshape(num_views * b, c1_v, c, -1).permute( | |
| 0, 1, 3, 2 | |
| ) # [V*B, M, H*W, C] | |
| for i, layer in enumerate(self.layers): | |
| concat0 = layer( | |
| concat0, | |
| concat1, | |
| height=h, | |
| width=w, | |
| shifted_window_attn_mask=shifted_window_attn_mask, | |
| attn_num_splits=attn_num_splits, | |
| ) | |
| if i < len(self.layers) - 1: | |
| # Rebuild target-view groups after each layer. | |
| features = unflatten_view_batch(concat0, b, num_views) | |
| concat0, concat1 = batch_features(features, nn_matrix=nn_matrix) | |
| features = unflatten_view_batch(concat0, b, num_views) | |
| features = features.reshape(b, num_views, h, w, c) | |
| features = features.permute(0, 1, 4, 2, 3).contiguous() | |
| return features | |
| def batch_features_camera_parameters( | |
| features, | |
| intrinsics, | |
| extrinsics, | |
| nn_matrix=None, | |
| ): | |
| # Construct batch-major inputs for plane-sweep warping. | |
| # features: [B, V, C, H, W]. | |
| # intrinsics/extrinsics: [B, V, 3, 3] and [B, V, 4, 4]. | |
| # Returns [B*V, ...] query tensors and [B*V, M, ...] target tensors. | |
| assert features.dim() == 5 | |
| assert intrinsics.shape[-2:] == (3, 3) | |
| assert extrinsics.shape[-2:] == (4, 4) | |
| batch_size, num_views, c, h, w = features.shape | |
| indices = target_view_indices( | |
| batch_size, num_views, nn_matrix, features.device | |
| ) | |
| num_selected_views = indices.size(-1) | |
| q = features | |
| q_intrinsics = intrinsics | |
| q_extrinsics = extrinsics | |
| kv = gather_view_tensor(features, indices) | |
| kv_intrinsics = gather_view_tensor(intrinsics, indices) | |
| kv_extrinsics = gather_view_tensor(extrinsics, indices) | |
| q = flatten_batch_view(q) # [BV, C, H, W] | |
| q_intrinsics = flatten_batch_view(q_intrinsics) # [BV, 3, 3] | |
| q_extrinsics = flatten_batch_view(q_extrinsics) # [BV, 4, 4] | |
| kv = kv.reshape(-1, num_selected_views, c, h, w) # [BV, V-1, C, H, W] | |
| kv_intrinsics = kv_intrinsics.reshape(-1, num_selected_views, 3, 3) | |
| kv_extrinsics = kv_extrinsics.reshape(-1, num_selected_views, 4, 4) | |
| return q, q_intrinsics, q_extrinsics, kv, kv_intrinsics, kv_extrinsics | |