File size: 1,618 Bytes
27fe41c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | 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.)}
|