AR / models /auto_v4.py

Upload folder using huggingface_hub

394918c verified 4 months ago

13.8 kB

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=2048):
        """
        旋转位置编码 (RoPE)
        dim: 必须是 head_dim (token_dim // num_heads)，通常为偶数
        max_seq_len: 预设的最大序列长度
        """
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        
        # 计算旋转频率项
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=True)
        
        # 初始化缓存
        self._set_cos_sin_cache(max_seq_len, device=torch.device("cpu"))
        

    def _set_cos_sin_cache(self, seq_len, device):
        self.max_seq_len = seq_len
        t = torch.arange(seq_len, device=device, dtype=torch.float32)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(device))
        emb = torch.cat((freqs, freqs), dim=-1)
        
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)

    def _rotate_half(self, x):
        x1 = x[..., :x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2:]
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, x, seq_len):
        if seq_len > self.max_seq_len or self.cos_cached.device != x.device:
            self._set_cos_sin_cache(seq_len=max(seq_len, self.max_seq_len), device=x.device)
            
        cos = self.cos_cached[..., :seq_len, :].to(dtype=x.dtype)
        sin = self.sin_cached[..., :seq_len, :].to(dtype=x.dtype)
        return (x * cos) + (self._rotate_half(x) * sin)


class TransformerLayer(nn.Module):
    """标准的自注意力时序层 (Self-Attention)"""
    def __init__(self, token_dim, num_heads, dim_feedforward, dropout):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads
        
        self.qkv = nn.Linear(token_dim, token_dim * 3)
        self.out_proj = nn.Linear(token_dim, token_dim)
        
        self.norm1 = nn.LayerNorm(token_dim)
        self.norm2 = nn.LayerNorm(token_dim)
        
        self.ffn = nn.Sequential(
            nn.Linear(token_dim, dim_feedforward),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_feedforward, token_dim),
            nn.Dropout(dropout)
        )
        self.dropout = dropout

    def forward(self, x, rope=None, mask=None):
        residual = x
        x = self.norm1(x)
        B_N, T, _ = x.shape
        
        qkv = self.qkv(x).chunk(3, dim=-1)
        q, k, v = map(
            lambda t: t.view(B_N, T, self.num_heads, self.head_dim).transpose(1, 2).contiguous(), 
            qkv
        )
        
        if rope is not None:
            q = rope(q, seq_len=T)
            k = rope(k, seq_len=T)
        
        x = F.scaled_dot_product_attention(
            q, k, v, 
            attn_mask=mask, 
            dropout_p=self.dropout if self.training else 0.0
        )
        
        x = x.transpose(1, 2).reshape(B_N, T, -1).contiguous()
        x = self.out_proj(x) + residual
        
        x = x + self.ffn(self.norm2(x))
        return x


class CrossAttentionLayer(nn.Module):
    """专门用于时序特征与文本条件对齐的交叉注意力层 (Cross-Attention)"""
    def __init__(self, token_dim, text_dim, num_heads, dim_feedforward, dropout):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads
        
        self.q_proj = nn.Linear(token_dim, token_dim)
        self.kv_proj = nn.Linear(text_dim, token_dim * 2)
        self.out_proj = nn.Linear(token_dim, token_dim)
        
        self.norm1 = nn.LayerNorm(token_dim)
        self.norm2 = nn.LayerNorm(token_dim)
        
        self.ffn = nn.Sequential(
            nn.Linear(token_dim, dim_feedforward),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_feedforward, token_dim),
            nn.Dropout(dropout)
        )
        self.dropout = dropout

    def forward(self, x, text_embed):
        # x 形状: (B_N, T, token_dim)
        # text_embed 形状: (B_N, L_text, text_dim) -> 外部需保证已对齐至 B_N 轴
        residual = x
        x = self.norm1(x)
        B_N, T, _ = x.shape
        
        q = self.q_proj(x).view(B_N, T, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        
        k_v = self.kv_proj(text_embed)
        k, v = k_v.chunk(2, dim=-1)
        k = k.view(B_N, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(B_N, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
        
        # 文本条件通常不需要时序的因果 Mask
        x = F.scaled_dot_product_attention(
            q, k, v, 
            attn_mask=None, 
            dropout_p=self.dropout if self.training else 0.0
        )
        
        x = x.transpose(1, 2).reshape(B_N, T, -1).contiguous()
        x = self.out_proj(x) + residual
        
        x = x + self.ffn(self.norm2(x))
        return x


class TemporalTrajectoryPredictor(nn.Module):
    def __init__(
        self,
        num_tokens=512,
        token_dim=768,
        text_dim=64,      # 新增 text_dim 输入参数
        max_seq_len=64,
        num_blocks=5,       # 重构为大块数量
        layers_per_block=4, # 每个大块内部自注意力层数
        num_heads=8,
        dim_feedforward=2048,
        dropout=0.05
    ):
        super().__init__()
        self.num_tokens = num_tokens
        self.token_dim = token_dim
        self.max_seq_len = max_seq_len
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads

        self.rope = RotaryEmbedding(self.head_dim, max_seq_len)
        
        # 动态构建层序列
        self.layers = nn.ModuleList()
        for _ in range(num_blocks):
            # 添加 4 层 Self-Attention 层
            for _ in range(layers_per_block):
                self.layers.append(
                    TransformerLayer(token_dim, num_heads, dim_feedforward, dropout)
                )
            # 紧接着添加 1 层 Cross-Attention 层
            self.layers.append(
                CrossAttentionLayer(token_dim, text_dim, num_heads, dim_feedforward, dropout)
            )
        
        self.output_projection = nn.Linear(token_dim, token_dim)
        self._init_weights()

    def _init_weights(self):
        nn.init.xavier_uniform_(self.output_projection.weight)
        if self.output_projection.bias is not None:
            nn.init.constant_(self.output_projection.bias, 0)

    def forward(self, history_frames, text_embed, mask=None, predict_next_only=False):
        """
        Args:
            history_frames: 4D 历史帧输入，形状为 (B, t, N, dim)
            text_embed: 文本条件，形状为 (B, 1, text_dim) 或 (B, L_text, text_dim)
            mask: 可选。自定义注意力掩码
            predict_next_only: 是否只预测未来下一帧
        """
        target_dtype = self.output_projection.weight.dtype 
        
        history_frames = history_frames.to(dtype=target_dtype)
        text_embed = text_embed.to(dtype=target_dtype)

        B, t, N, dim = history_frames.shape
        
        # 空间维度 N 转入 Batch 轴，暴露出纯粹的时间序列维度 t: (B * N, t, dim)
        x = history_frames.transpose(1, 2).reshape(B * N, t, dim).contiguous()
        
        # 同步将 text_embed 扩展到与 x 相同的 Batch 维度 (B * N)
        # text_embed 原始: (B, L, text_dim) -> 扩展后: (B, N, L, text_dim) -> 重塑为: (B * N, L, text_dim)
        L_text = text_embed.shape[1]
        text_dim_size = text_embed.shape[-1]
        text_embed_expanded = text_embed.unsqueeze(1).expand(-1, N, -1, -1).reshape(B * N, L_text, text_dim_size).contiguous()
        
        if mask is None:
            mask = torch.ones((t, t), device=x.device, dtype=torch.bool).tril()
            mask = mask.unsqueeze(0).unsqueeze(1)
        elif mask.dtype != torch.bool:
            mask = (mask < 0) if mask.dtype.is_floating_point else mask.bool()

        # 遍历执行所有层
        for layer in self.layers:
            if isinstance(layer, TransformerLayer):
                # 自注意力机制：传入 rope 和时序因果掩码 mask
                x = layer(x, rope=self.rope, mask=mask)
            elif isinstance(layer, CrossAttentionLayer):
                # 交叉注意力机制：传入对齐后的 text_embed_expanded
                x = layer(x, text_embed=text_embed_expanded)
            
        if predict_next_only:
            x_last = x[:, -1:, :]  
            next_frame = self.output_projection(x_last)  
            return next_frame.reshape(B, N, 1, dim).transpose(1, 2).contiguous()
        else:
            all_preds = self.output_projection(x)  
            return all_preds.reshape(B, N, t, dim).transpose(1, 2).contiguous()


class SpatialGraphRefiner(nn.Module):
    def __init__(
        self,
        num_tokens=512,
        token_dim=768,
        text_dim=1024,
        max_graph_dist=10,
        num_blocks=3,
        num_heads=8,
        dim_feedforward=2048,
        dropout=0.05
    ):
        super().__init__()
        self.num_tokens = num_tokens
        self.token_dim = token_dim
        self.max_graph_dist = max_graph_dist
        self.num_heads = num_heads
        self.head_dim = token_dim // num_heads

        distances = torch.arange(max_graph_dist + 1, dtype=torch.float32)
        scales = torch.linspace(0.1, 1.0, num_heads)
        self.register_buffer("graph_bias_table", -1.0 * (distances.unsqueeze(1) * scales.unsqueeze(0)))

        self.blocks = nn.ModuleList([
            nn.ModuleDict({
                "qkv1": nn.Linear(token_dim, token_dim * 3),
                "out1": nn.Linear(token_dim, token_dim),
                "ln1": nn.LayerNorm(token_dim),
                "qkv2": nn.Linear(token_dim, token_dim * 3),
                "out2": nn.Linear(token_dim, token_dim),
                "ln2": nn.LayerNorm(token_dim),
                "q_ca": nn.Linear(token_dim, token_dim),
                "kv_ca": nn.Linear(text_dim, token_dim * 2),
                "out_ca": nn.Linear(token_dim, token_dim),
                "ln3": nn.LayerNorm(token_dim),
                "ffn": nn.Sequential(
                    nn.Linear(token_dim, dim_feedforward),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(dim_feedforward, token_dim),
                    nn.Dropout(dropout)
                ),
                "ln4": nn.LayerNorm(token_dim)
            }) for _ in range(num_blocks)
        ])

        # self.delta_proj = nn.Linear(token_dim, token_dim)
        # nn.init.zeros_(self.delta_proj.weight)
        # nn.init.zeros_(self.delta_proj.bias)

    def _compute_graph_bias(self, shortest_path_matrix, B, t):
        N = shortest_path_matrix.shape[-1]
        dist_clipped = torch.clamp(shortest_path_matrix, 0, self.max_graph_dist).long()
        bias = F.embedding(dist_clipped, self.graph_bias_table) 
        return bias.unsqueeze(1).expand(-1, t, -1, -1, -1).permute(0, 1, 4, 2, 3).reshape(B * t, self.num_heads, N, N)

    def _sdpa_attention(self, q, k, v, bias=None):
        return F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.05 if self.training else 0.0)

    def forward(self, next_frame_init, text_embed, shortest_path_matrix):
        B, T, N, dim = next_frame_init.shape  
        graph_attn_bias = self._compute_graph_bias(shortest_path_matrix, B, T)
        
        x = next_frame_init.reshape(B * T, N, dim)
        text_embed_expanded = text_embed.unsqueeze(1).expand(-1, T, -1, -1).reshape(B * T, -1, text_embed.shape[-1])
        
        for b in self.blocks:
            # 1. Pre-LN SA1
            residual = x
            x_norm = b["ln1"](x)
            q, k, v = b["qkv1"](x_norm).chunk(3, dim=-1)
            q = q.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            k = k.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            v = v.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            attn = self._sdpa_attention(q, k, v, bias=graph_attn_bias)
            x = residual + b["out1"](attn.transpose(1, 2).reshape(B * T, N, dim))
            
            # 2. Pre-LN SA2
            residual = x
            x_norm = b["ln2"](x)
            q, k, v = b["qkv2"](x_norm).chunk(3, dim=-1)
            q = q.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            k = k.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            v = v.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            attn = self._sdpa_attention(q, k, v, bias=graph_attn_bias)
            x = residual + b["out2"](attn.transpose(1, 2).reshape(B * T, N, dim))
            
            # 3. Pre-LN CA
            residual = x
            x_norm = b["ln3"](x)
            q = b["q_ca"](x_norm).view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
            k_v = b["kv_ca"](text_embed_expanded)
            k, v = k_v.chunk(2, dim=-1)
            k = k.view(B * T, -1, self.num_heads, self.head_dim).transpose(1, 2)
            v = v.view(B * T, -1, self.num_heads, self.head_dim).transpose(1, 2)
            attn = self._sdpa_attention(q, k, v)
            x = residual + b["out_ca"](attn.transpose(1, 2).reshape(B * T, N, dim))
            
            # 4. Pre-LN FFN
            x = x + b["ffn"](b["ln4"](x))
            
        # delta = self.delta_proj(x).reshape(B, T, N, dim)
        return x.reshape(B, T, N, dim)

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	import math

	class RotaryEmbedding(nn.Module):
	def __init__(self, dim, max_seq_len=2048):
	"""
	旋转位置编码 (RoPE)
	dim: 必须是 head_dim (token_dim // num_heads)，通常为偶数
	max_seq_len: 预设的最大序列长度
	"""
	super().__init__()
	self.dim = dim
	self.max_seq_len = max_seq_len

	# 计算旋转频率项
	inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
	self.register_buffer("inv_freq", inv_freq, persistent=True)

	# 初始化缓存
	self._set_cos_sin_cache(max_seq_len, device=torch.device("cpu"))


	def _set_cos_sin_cache(self, seq_len, device):
	self.max_seq_len = seq_len
	t = torch.arange(seq_len, device=device, dtype=torch.float32)
	freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(device))
	emb = torch.cat((freqs, freqs), dim=-1)

	self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
	self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)

	def _rotate_half(self, x):
	x1 = x[..., :x.shape[-1] // 2]
	x2 = x[..., x.shape[-1] // 2:]
	return torch.cat((-x2, x1), dim=-1)

	def forward(self, x, seq_len):
	if seq_len > self.max_seq_len or self.cos_cached.device != x.device:
	self._set_cos_sin_cache(seq_len=max(seq_len, self.max_seq_len), device=x.device)

	cos = self.cos_cached[..., :seq_len, :].to(dtype=x.dtype)
	sin = self.sin_cached[..., :seq_len, :].to(dtype=x.dtype)
	return (x * cos) + (self._rotate_half(x) * sin)


	class TransformerLayer(nn.Module):
	"""标准的自注意力时序层 (Self-Attention)"""
	def __init__(self, token_dim, num_heads, dim_feedforward, dropout):
	super().__init__()
	self.num_heads = num_heads
	self.head_dim = token_dim // num_heads

	self.qkv = nn.Linear(token_dim, token_dim * 3)
	self.out_proj = nn.Linear(token_dim, token_dim)

	self.norm1 = nn.LayerNorm(token_dim)
	self.norm2 = nn.LayerNorm(token_dim)

	self.ffn = nn.Sequential(
	nn.Linear(token_dim, dim_feedforward),
	nn.GELU(),
	nn.Dropout(dropout),
	nn.Linear(dim_feedforward, token_dim),
	nn.Dropout(dropout)
	)
	self.dropout = dropout

	def forward(self, x, rope=None, mask=None):
	residual = x
	x = self.norm1(x)
	B_N, T, _ = x.shape

	qkv = self.qkv(x).chunk(3, dim=-1)
	q, k, v = map(
	lambda t: t.view(B_N, T, self.num_heads, self.head_dim).transpose(1, 2).contiguous(),
	qkv
	)

	if rope is not None:
	q = rope(q, seq_len=T)
	k = rope(k, seq_len=T)

	x = F.scaled_dot_product_attention(
	q, k, v,
	attn_mask=mask,
	dropout_p=self.dropout if self.training else 0.0
	)

	x = x.transpose(1, 2).reshape(B_N, T, -1).contiguous()
	x = self.out_proj(x) + residual

	x = x + self.ffn(self.norm2(x))
	return x


	class CrossAttentionLayer(nn.Module):
	"""专门用于时序特征与文本条件对齐的交叉注意力层 (Cross-Attention)"""
	def __init__(self, token_dim, text_dim, num_heads, dim_feedforward, dropout):
	super().__init__()
	self.num_heads = num_heads
	self.head_dim = token_dim // num_heads

	self.q_proj = nn.Linear(token_dim, token_dim)
	self.kv_proj = nn.Linear(text_dim, token_dim * 2)
	self.out_proj = nn.Linear(token_dim, token_dim)

	self.norm1 = nn.LayerNorm(token_dim)
	self.norm2 = nn.LayerNorm(token_dim)

	self.ffn = nn.Sequential(
	nn.Linear(token_dim, dim_feedforward),
	nn.GELU(),
	nn.Dropout(dropout),
	nn.Linear(dim_feedforward, token_dim),
	nn.Dropout(dropout)
	)
	self.dropout = dropout

	def forward(self, x, text_embed):
	# x 形状: (B_N, T, token_dim)
	# text_embed 形状: (B_N, L_text, text_dim) -> 外部需保证已对齐至 B_N 轴
	residual = x
	x = self.norm1(x)
	B_N, T, _ = x.shape

	q = self.q_proj(x).view(B_N, T, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

	k_v = self.kv_proj(text_embed)
	k, v = k_v.chunk(2, dim=-1)
	k = k.view(B_N, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
	v = v.view(B_N, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

	# 文本条件通常不需要时序的因果 Mask
	x = F.scaled_dot_product_attention(
	q, k, v,
	attn_mask=None,
	dropout_p=self.dropout if self.training else 0.0
	)

	x = x.transpose(1, 2).reshape(B_N, T, -1).contiguous()
	x = self.out_proj(x) + residual

	x = x + self.ffn(self.norm2(x))
	return x


	class TemporalTrajectoryPredictor(nn.Module):
	def __init__(
	self,
	num_tokens=512,
	token_dim=768,
	text_dim=64, # 新增 text_dim 输入参数
	max_seq_len=64,
	num_blocks=5, # 重构为大块数量
	layers_per_block=4, # 每个大块内部自注意力层数
	num_heads=8,
	dim_feedforward=2048,
	dropout=0.05
	):
	super().__init__()
	self.num_tokens = num_tokens
	self.token_dim = token_dim
	self.max_seq_len = max_seq_len
	self.num_heads = num_heads
	self.head_dim = token_dim // num_heads

	self.rope = RotaryEmbedding(self.head_dim, max_seq_len)

	# 动态构建层序列
	self.layers = nn.ModuleList()
	for _ in range(num_blocks):
	# 添加 4 层 Self-Attention 层
	for _ in range(layers_per_block):
	self.layers.append(
	TransformerLayer(token_dim, num_heads, dim_feedforward, dropout)
	)
	# 紧接着添加 1 层 Cross-Attention 层
	self.layers.append(
	CrossAttentionLayer(token_dim, text_dim, num_heads, dim_feedforward, dropout)
	)

	self.output_projection = nn.Linear(token_dim, token_dim)
	self._init_weights()

	def _init_weights(self):
	nn.init.xavier_uniform_(self.output_projection.weight)
	if self.output_projection.bias is not None:
	nn.init.constant_(self.output_projection.bias, 0)

	def forward(self, history_frames, text_embed, mask=None, predict_next_only=False):
	"""
	Args:
	history_frames: 4D 历史帧输入，形状为 (B, t, N, dim)
	text_embed: 文本条件，形状为 (B, 1, text_dim) 或 (B, L_text, text_dim)
	mask: 可选。自定义注意力掩码
	predict_next_only: 是否只预测未来下一帧
	"""
	target_dtype = self.output_projection.weight.dtype

	history_frames = history_frames.to(dtype=target_dtype)
	text_embed = text_embed.to(dtype=target_dtype)

	B, t, N, dim = history_frames.shape

	# 空间维度 N 转入 Batch 轴，暴露出纯粹的时间序列维度 t: (B * N, t, dim)
	x = history_frames.transpose(1, 2).reshape(B * N, t, dim).contiguous()

	# 同步将 text_embed 扩展到与 x 相同的 Batch 维度 (B * N)
	# text_embed 原始: (B, L, text_dim) -> 扩展后: (B, N, L, text_dim) -> 重塑为: (B * N, L, text_dim)
	L_text = text_embed.shape[1]
	text_dim_size = text_embed.shape[-1]
	text_embed_expanded = text_embed.unsqueeze(1).expand(-1, N, -1, -1).reshape(B * N, L_text, text_dim_size).contiguous()

	if mask is None:
	mask = torch.ones((t, t), device=x.device, dtype=torch.bool).tril()
	mask = mask.unsqueeze(0).unsqueeze(1)
	elif mask.dtype != torch.bool:
	mask = (mask < 0) if mask.dtype.is_floating_point else mask.bool()

	# 遍历执行所有层
	for layer in self.layers:
	if isinstance(layer, TransformerLayer):
	# 自注意力机制：传入 rope 和时序因果掩码 mask
	x = layer(x, rope=self.rope, mask=mask)
	elif isinstance(layer, CrossAttentionLayer):
	# 交叉注意力机制：传入对齐后的 text_embed_expanded
	x = layer(x, text_embed=text_embed_expanded)

	if predict_next_only:
	x_last = x[:, -1:, :]
	next_frame = self.output_projection(x_last)
	return next_frame.reshape(B, N, 1, dim).transpose(1, 2).contiguous()
	else:
	all_preds = self.output_projection(x)
	return all_preds.reshape(B, N, t, dim).transpose(1, 2).contiguous()


	class SpatialGraphRefiner(nn.Module):
	def __init__(
	self,
	num_tokens=512,
	token_dim=768,
	text_dim=1024,
	max_graph_dist=10,
	num_blocks=3,
	num_heads=8,
	dim_feedforward=2048,
	dropout=0.05
	):
	super().__init__()
	self.num_tokens = num_tokens
	self.token_dim = token_dim
	self.max_graph_dist = max_graph_dist
	self.num_heads = num_heads
	self.head_dim = token_dim // num_heads

	distances = torch.arange(max_graph_dist + 1, dtype=torch.float32)
	scales = torch.linspace(0.1, 1.0, num_heads)
	self.register_buffer("graph_bias_table", -1.0 * (distances.unsqueeze(1) * scales.unsqueeze(0)))

	self.blocks = nn.ModuleList([
	nn.ModuleDict({
	"qkv1": nn.Linear(token_dim, token_dim * 3),
	"out1": nn.Linear(token_dim, token_dim),
	"ln1": nn.LayerNorm(token_dim),
	"qkv2": nn.Linear(token_dim, token_dim * 3),
	"out2": nn.Linear(token_dim, token_dim),
	"ln2": nn.LayerNorm(token_dim),
	"q_ca": nn.Linear(token_dim, token_dim),
	"kv_ca": nn.Linear(text_dim, token_dim * 2),
	"out_ca": nn.Linear(token_dim, token_dim),
	"ln3": nn.LayerNorm(token_dim),
	"ffn": nn.Sequential(
	nn.Linear(token_dim, dim_feedforward),
	nn.GELU(),
	nn.Dropout(dropout),
	nn.Linear(dim_feedforward, token_dim),
	nn.Dropout(dropout)
	),
	"ln4": nn.LayerNorm(token_dim)
	}) for _ in range(num_blocks)
	])

	# self.delta_proj = nn.Linear(token_dim, token_dim)
	# nn.init.zeros_(self.delta_proj.weight)
	# nn.init.zeros_(self.delta_proj.bias)

	def _compute_graph_bias(self, shortest_path_matrix, B, t):
	N = shortest_path_matrix.shape[-1]
	dist_clipped = torch.clamp(shortest_path_matrix, 0, self.max_graph_dist).long()
	bias = F.embedding(dist_clipped, self.graph_bias_table)
	return bias.unsqueeze(1).expand(-1, t, -1, -1, -1).permute(0, 1, 4, 2, 3).reshape(B * t, self.num_heads, N, N)

	def _sdpa_attention(self, q, k, v, bias=None):
	return F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.05 if self.training else 0.0)

	def forward(self, next_frame_init, text_embed, shortest_path_matrix):
	B, T, N, dim = next_frame_init.shape
	graph_attn_bias = self._compute_graph_bias(shortest_path_matrix, B, T)

	x = next_frame_init.reshape(B * T, N, dim)
	text_embed_expanded = text_embed.unsqueeze(1).expand(-1, T, -1, -1).reshape(B * T, -1, text_embed.shape[-1])

	for b in self.blocks:
	# 1. Pre-LN SA1
	residual = x
	x_norm = b["ln1"](x)
	q, k, v = b["qkv1"](x_norm).chunk(3, dim=-1)
	q = q.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
	k = k.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
	v = v.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
	attn = self._sdpa_attention(q, k, v, bias=graph_attn_bias)
	x = residual + b["out1"](attn.transpose(1, 2).reshape(B * T, N, dim))

	# 2. Pre-LN SA2
	residual = x
	x_norm = b["ln2"](x)
	q, k, v = b["qkv2"](x_norm).chunk(3, dim=-1)
	q = q.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
	k = k.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
	v = v.view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
	attn = self._sdpa_attention(q, k, v, bias=graph_attn_bias)
	x = residual + b["out2"](attn.transpose(1, 2).reshape(B * T, N, dim))

	# 3. Pre-LN CA
	residual = x
	x_norm = b["ln3"](x)
	q = b["q_ca"](x_norm).view(B * T, N, self.num_heads, self.head_dim).transpose(1, 2)
	k_v = b["kv_ca"](text_embed_expanded)
	k, v = k_v.chunk(2, dim=-1)
	k = k.view(B * T, -1, self.num_heads, self.head_dim).transpose(1, 2)
	v = v.view(B * T, -1, self.num_heads, self.head_dim).transpose(1, 2)
	attn = self._sdpa_attention(q, k, v)
	x = residual + b["out_ca"](attn.transpose(1, 2).reshape(B * T, N, dim))

	# 4. Pre-LN FFN
	x = x + b["ffn"](b["ln4"](x))

	# delta = self.delta_proj(x).reshape(B, T, N, dim)
	return x.reshape(B, T, N, dim)