v7/runtime/multiview.py

from __future__ import annotations
import math,torch
import torch.nn as nn

class PoseFourier(nn.Module):
    def __init__(self,bands=8,out=64):super().__init__();self.bands=bands;self.proj=nn.Linear(bands*4,out)
    def forward(self,azimuth_deg,elevation_deg):
        a=torch.deg2rad(azimuth_deg.float());e=torch.deg2rad(elevation_deg.float());freq=2**torch.arange(self.bands,device=a.device,dtype=a.dtype)
        z=torch.cat([torch.sin(a[:,None]*freq),torch.cos(a[:,None]*freq),torch.sin(e[:,None]*freq),torch.cos(e[:,None]*freq)],-1);return self.proj(z)

class MultiViewAggregator(nn.Module):
    """Order-aware but permutation-stable view fusion for AniGen condition tokens."""
    def __init__(self,dim=1024,pose_dim=64,heads=8):
        super().__init__();self.pose=PoseFourier(out=pose_dim);self.pose_proj=nn.Linear(pose_dim,dim);self.attn=nn.MultiheadAttention(dim,heads,batch_first=True);self.norm=nn.LayerNorm(dim)
    def forward(self,view_tokens,azimuth,elevation,view_mask=None):
        # view_tokens [B,V,T,C]
        b,v,t,c=view_tokens.shape;p=self.pose(azimuth.reshape(-1),elevation.reshape(-1)).view(b,v,-1);x=view_tokens+self.pose_proj(p)[:,:,None,:]
        # fuse each token position across views; output remains [B,T,C]
        q=x.mean(1);kv=x.transpose(1,2).reshape(b*t,v,c);qq=q.reshape(b*t,1,c);key_padding=None if view_mask is None else ~view_mask[:,None,:].expand(b,t,v).reshape(b*t,v)
        y,_=self.attn(qq,kv,kv,key_padding_mask=key_padding,need_weights=False);return self.norm(q+y.reshape(b,t,c))

DEFAULT_VIEWS={'front':(0.,0.),'left':(-90.,0.),'back':(180.,0.),'right':(90.,0.)}