import torch import torch.nn as nn import torch.nn.functional as F import math from einops import rearrange from modules.attention import MLA from .point_conditioner import PointConditioner try: from flash_attn import flash_attn_varlen_func except Exception: try: from flash_attn.flash_attn_interface import flash_attn_varlen_func except Exception: flash_attn_varlen_func = None class SinusoidalPosEmb(nn.Module): def __init__(self, dim, scale=100.0): super().__init__() self.dim = dim self.scale = scale def forward(self, x): device = x.device dtype = x.dtype x = x.float() * self.scale half_dim = self.dim // 2 emb = math.log(10000) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, device=device) * -emb) emb = x[:, None] * emb[None, :] emb = torch.cat((emb.sin(), emb.cos()), dim=-1) return emb.to(dtype=dtype) class RMSNorm(nn.Module): """现代 LLM 标配：去均值的 RMSNorm""" def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps # 只需要学习一个缩放参数，去掉 bias self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): # 计算均方根 (Root Mean Square) norm_x = torch.mean(x.pow(2), dim=-1, keepdim=True) x_normed = x * torch.rsqrt(norm_x + self.eps) return self.weight * x_normed class SwiGLU(nn.Module): """现代 LLM 标配：SwiGLU 激活的门控 FFN""" def __init__(self, d_model, hidden_dim=None): super().__init__() # 将原先 4 * d_model 的维度缩小到 8/3，保持参数量对等; hidden_dim 可显式覆盖 # (如 1024 宽度时 8/3 给 2730 非 256 倍数, 显式取 2816 对齐 Tensor Core) hidden_dim = int((8 / 3) * d_model) if hidden_dim is None else int(hidden_dim) # 提示：在实际工程中，通常会把 hidden_dim 向上取整到 128 或 256 的倍数，以最大化 Tensor Core 效率 self.w1 = nn.Linear(d_model, hidden_dim, bias=False) self.w2 = nn.Linear(d_model, hidden_dim, bias=False) self.w3 = nn.Linear(hidden_dim, d_model, bias=False) def forward(self, x): # 公式: Swish(xW1) ⊗ (xW2) -> W3 # PyTorch 中 F.silu 就是 Swish return self.w3(F.silu(self.w1(x)) * self.w2(x)) class ModernAttention(nn.Module): def __init__(self, d_model, nhead, dropout=0., qk_norm=False): super().__init__() assert d_model % nhead == 0, "d_model 必须能被 nhead 整除" self.nhead = nhead self.d_head = d_model // nhead # 现代架构通常在 Attention 的投影层中也丢弃 bias 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.out_proj = nn.Linear(d_model, d_model, bias=False) self.dropout = dropout # QK-norm: 逐 head 对 Q/K 在 head_dim 上做 RMSNorm(在 RoPE 之前),稳定 attention logits、 # 防幅度爆炸。默认关(不改变既有模型);flow 等新模型可开。新增参数 -> 旧 ckpt load_strict=False。 self.qk_norm = qk_norm if qk_norm: self.q_norm = RMSNorm(self.d_head) self.k_norm = RMSNorm(self.d_head) def forward(self, q, k, v, mask=None, is_causal=False, rope_apply_fn=None): B, L_q, _ = q.size() _, L_k, _ = k.size() # 投影并调整形状为 [Batch, Head, SeqLen, HeadDim] Q = self.q_proj(q).view(B, L_q, self.nhead, self.d_head).transpose(1, 2) K = self.k_proj(k).view(B, L_k, self.nhead, self.d_head).transpose(1, 2) V = self.v_proj(v).view(B, L_k, self.nhead, self.d_head).transpose(1, 2) if self.qk_norm: # QK-norm(RoPE 前),保持 dtype 供 SDPA/flash Q = self.q_norm(Q).to(Q.dtype); K = self.k_norm(K).to(K.dtype) # 注入旋转位置编码 (RoPE) if rope_apply_fn is not None: Q, K = rope_apply_fn(Q, K) # 使用 PyTorch 2.0+ 内置的 SDPA，它会自动调用 Flash Attention 节省巨量显存 out = F.scaled_dot_product_attention( Q, K, V, attn_mask=mask, dropout_p=self.dropout if self.training else 0.0, is_causal=is_causal ) # 恢复形状并输出 out = out.transpose(1, 2).contiguous().view(B, L_q, -1) return self.out_proj(out) def forward_varlen(self, q, k, v, lengths_q, lengths_k=None, is_causal=False, rope_apply_fn=None, max_len_q=None, max_len_k=None): """max_len_q / max_len_k: flash 需要的 max_seqlen。**传 python int 可避免 GPU->CPU 同步**。 不传时回退到 `int(lengths.max().item())` —— 那是一次 cudaStreamSynchronize: CPU 必须等 GPU 把已排队的活干完才能读到值, 于是无法提前排队 kernel。本模型一次前向有 36 次 attention 调用(18 self + 18 cross), 即 72 次同步, 是「kernel 又多又小、GPU 有空隙」的直接成因, 也让 torch.compile 在这里 graph break(数据依赖的 unbacked symint)。 调用方本来就知道这些长度(lengths 是从 [c.numel() for c in coord_list] 建的), 应当传进来。 """ if lengths_k is None: lengths_k = lengths_q if max_len_k is None: max_len_k = max_len_q total_q = q.size(0) total_k = k.size(0) Q = self.q_proj(q).view(total_q, self.nhead, self.d_head) K = self.k_proj(k).view(total_k, self.nhead, self.d_head) V = self.v_proj(v).view(total_k, self.nhead, self.d_head) if self.qk_norm: # QK-norm(RoPE 前),.to 保持 dtype 供 flash-attn Q = self.q_norm(Q).to(Q.dtype); K = self.k_norm(K).to(K.dtype) if rope_apply_fn is not None: Q, K = rope_apply_fn(Q, K) dropout_p = self.dropout if self.training else 0.0 if can_use_flash_attn(Q): out = flash_attn_varlen_func( Q.contiguous(), K.contiguous(), V.contiguous(), lengths_to_cu_seqlens(lengths_q), lengths_to_cu_seqlens(lengths_k), int(lengths_q.max().item()) if max_len_q is None else int(max_len_q), int(lengths_k.max().item()) if max_len_k is None else int(max_len_k), dropout_p=dropout_p, causal=is_causal, ) else: out_chunks = [] q_start = 0 k_start = 0 for q_len, k_len in zip(lengths_q.tolist(), lengths_k.tolist()): q_end = q_start + q_len k_end = k_start + k_len q_i = Q[q_start:q_end].transpose(0, 1).unsqueeze(0) k_i = K[k_start:k_end].transpose(0, 1).unsqueeze(0) v_i = V[k_start:k_end].transpose(0, 1).unsqueeze(0) out_i = F.scaled_dot_product_attention( q_i, k_i, v_i, dropout_p=dropout_p, is_causal=is_causal, ) out_chunks.append(out_i.squeeze(0).transpose(0, 1)) q_start = q_end k_start = k_end out = torch.cat(out_chunks, dim=0) return self.out_proj(out.contiguous().view(total_q, -1)) # class CrossAttnTransformerBlock(nn.Module): # def __init__(self, d_model, nhead, dropout=0.): # super().__init__() # self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True) # self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True) # self.norm1 = nn.LayerNorm(d_model) # self.norm2 = nn.LayerNorm(d_model) # self.norm3 = nn.LayerNorm(d_model) # self.ffn = nn.Sequential( # nn.Linear(d_model, d_model * 4), # nn.GELU(), # nn.Dropout(dropout), # nn.Linear(d_model * 4, d_model), # ) # self.dropout = nn.Dropout(dropout) # def forward(self, x, cond): # # 1) token self-attention # h = self.norm1(x) # sa_out, _ = self.self_attn(h, h, h, need_weights=False) # x = x + self.dropout(sa_out) # # 2) cross-attention: query from current tokens, key/value from condition tokens # h = self.norm2(x) # ca_out, _ = self.cross_attn(h, cond, cond, need_weights=False) # x = x + self.dropout(ca_out) # # 3) feed-forward # h = self.norm3(x) # ff_out = self.ffn(h) # x = x + self.dropout(ff_out) # return x class ModernCrossAttnTransformerBlock(nn.Module): def __init__(self, d_model, nhead, dropout=0.): super().__init__() # 替换为支持 Flash Attention 的现代 Attention self.self_attn = ModernAttention(d_model, nhead, dropout=dropout) self.cross_attn = ModernAttention(d_model, nhead, dropout=dropout) # 替换 LayerNorm 为 RMSNorm self.norm1 = RMSNorm(d_model) self.norm2 = RMSNorm(d_model) self.norm3 = RMSNorm(d_model) # 替换传统 FFN 为 SwiGLU self.ffn = SwiGLU(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, cond, rope_apply_fn=None, causal_mask=False): # 1) token self-attention # 你的 Pre-Norm 逻辑保持不变，这是正确的 h = self.norm1(x) # 注意：在这里把 RoPE 传进去处理 Query 和 Key sa_out = self.self_attn(h, h, h, is_causal=causal_mask, rope_apply_fn=rope_apply_fn) x = x + self.dropout(sa_out) # 2) cross-attention: query from current tokens, key/value from condition tokens if cond is not None: h = self.norm2(x) # Cross Attention 一般不使用 RoPE，因为输入源不同 ca_out = self.cross_attn(h, cond, cond) x = x + self.dropout(ca_out) # 3) feed-forward (SwiGLU) h = self.norm3(x) ff_out = self.ffn(h) x = x + self.dropout(ff_out) return x def forward_varlen(self, x, lengths, cond=None, cond_lengths=None, rope_apply_fn=None, causal_mask=False): h = self.norm1(x) sa_out = self.self_attn.forward_varlen( h, h, h, lengths, is_causal=causal_mask, rope_apply_fn=rope_apply_fn, ) x = x + self.dropout(sa_out) if cond is not None: h = self.norm2(x) ca_out = self.cross_attn.forward_varlen( h, cond, cond, lengths, lengths_k=cond_lengths, ) x = x + self.dropout(ca_out) h = self.norm3(x) ff_out = self.ffn(h) x = x + self.dropout(ff_out) return x def rotate_half(x): """ 旋转特征的一半维度。 对于向量 [x1, x2, ..., x_d/2, x_d/2+1, ..., x_d]， 将其转换为 [-x_d/2+1, ..., -x_d, x1, x2, ..., x_d/2]。 这在数学上等价于复数平面上的乘以 i。 """ x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin): """ 将预计算好的 cos 和 sin 注入到 Query 和 Key 中。 q, k 形状期望为: (Batch, Num_Heads, Seq_Len, Head_Dim) cos, sin 形状期望为: (1, 1, Seq_Len, Head_Dim) """ q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed def apply_rotary_pos_emb_varlen(q, k, cos, sin): """RoPE。**cos/sin 先转成 q/k 的 dtype**，保证输出 dtype 不被 cos/sin 带跑。 为什么必须这么写：cos/sin 由 `rotary_emb.forward_positions(..., dtype=x.dtype)` 生成，而 9 尺度的 rope_fn 是在 `x = raw_emb + refine_emb` 之后立刻构造的 —— `nn.Embedding` 不在 autocast 的转型名单里，那时 x 还是 **fp32**，于是 cos/sin 是 fp32。bf16 的 Q/K 乘上 fp32 的 cos/sin 会 类型提升成 fp32，`can_use_flash_attn` 判否，直接掉进 `ModernAttention.forward_varlen` 里 **SDPA 的逐序列 python 循环**。命中的恰好是 enc_blocks_9 / dec_blocks_9 这 4 层, 也就是跑在满 token 长度(其余层已被 hourglass 降到 1/3、1/9)上最贵的 4 个注意力。 3 尺度与 bottleneck 的 rope_fn 因为构造时 x 已是 bf16，本来就走 flash —— 即同一网络里 4 层 fp32 慢路径 + 14 层 bf16 快路径，不一致。统一到 bf16。 """ cos_q, sin_q = cos.to(q.dtype), sin.to(q.dtype) q_embed = (q * cos_q) + (rotate_half(q) * sin_q) cos_k, sin_k = (cos_q, sin_q) if k.dtype == q.dtype else (cos.to(k.dtype), sin.to(k.dtype)) k_embed = (k * cos_k) + (rotate_half(k) * sin_k) return q_embed, k_embed def lengths_to_cu_seqlens(lengths): lengths = lengths.to(dtype=torch.int32) prefix = torch.empty(lengths.numel() + 1, device=lengths.device, dtype=torch.int32) prefix[0] = 0 prefix[1:] = torch.cumsum(lengths, dim=0, dtype=torch.int32) return prefix.contiguous() def lengths_to_position_ids(lengths): total = int(lengths.sum().item()) if total == 0: return torch.empty(0, device=lengths.device, dtype=torch.long) starts = torch.cumsum(lengths, dim=0) - lengths return torch.arange(total, device=lengths.device, dtype=torch.long) - torch.repeat_interleave(starts, lengths) def can_use_flash_attn(x): return ( flash_attn_varlen_func is not None and x.is_cuda and x.dtype in (torch.float16, torch.bfloat16) ) class RotaryEmbedding(nn.Module): """ 预计算和缓存 RoPE 的 cos 和 sin 矩阵。 """ def __init__(self, dim, max_position_embeddings=90000, base=10000.0, device=None): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base # 计算旋转频率: theta_i = 10000^(-2(i-1)/d) inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) # 预先构建缓存 self._set_cos_sin_cache(seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32) # freqs 形状: (seq_len, dim/2) freqs = torch.outer(t, self.inv_freq) # 将频率在特征维度上复制一份，形状变为: (seq_len, dim) # 前半部分和后半部分的频率一致，为了和 rotate_half 的逻辑对齐 emb = torch.cat((freqs, freqs), dim=-1) # 增加 Batch 和 Head 的广播维度: (1, 1, Seq_Len, Head_Dim) self.register_buffer("cos_cached", emb.cos().to(dtype).unsqueeze(0).unsqueeze(0), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype).unsqueeze(0).unsqueeze(0), persistent=False) def forward(self, x, seq_len=None): """ x 形状: (Batch, Num_Heads, Seq_Len, Head_Dim) """ # 如果当前序列长度超过了预计算的缓存，则动态扩展缓存 if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) # 截取当前需要的序列长度返回 return ( self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype), self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype), ) def forward_positions(self, position_ids, dtype=None): if dtype is None: dtype = self.cos_cached.dtype if position_ids.numel() == 0: empty = torch.empty(0, 1, self.dim, device=position_ids.device, dtype=dtype) return empty, empty seq_len = int(position_ids.max().item()) + 1 if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len=seq_len, device=position_ids.device, dtype=dtype) cos = self.cos_cached[0, 0, position_ids, :].to(device=position_ids.device, dtype=dtype) sin = self.sin_cached[0, 0, position_ids, :].to(device=position_ids.device, dtype=dtype) return cos.unsqueeze(1), sin.unsqueeze(1) class LinearDownsample(nn.Module): def __init__(self, dim, shorten_factor): super().__init__() self.proj = nn.Linear(dim * shorten_factor, dim) self.shorten_factor = shorten_factor def forward(self, x): x = rearrange(x, 'b (n s) d -> b n (s d)', s=self.shorten_factor) return self.proj(x) def forward_varlen(self, x, lengths): if torch.any(lengths % self.shorten_factor != 0): raise ValueError(f"All sequence lengths must be divisible by {self.shorten_factor} for hourglass downsample.") x = rearrange(x, '(n s) d -> n (s d)', s=self.shorten_factor) return self.proj(x), lengths // self.shorten_factor class LinearUpsample(nn.Module): def __init__(self, dim, shorten_factor): super().__init__() self.proj = nn.Linear(dim, dim * shorten_factor) self.shorten_factor = shorten_factor def forward(self, x): x = self.proj(x) return rearrange(x, 'b n (s d) -> b (n s) d', s=self.shorten_factor) def forward_varlen(self, x, lengths): x = self.proj(x) x = rearrange(x, 'n (s d) -> (n s) d', s=self.shorten_factor) return x, lengths * self.shorten_factor class DiscreteDiffusionTransformer(nn.Module): def __init__(self, vocab_A, vocab_B, max_seq_len, d_model=768, nhead=12, depth = [4,4,8], dropout=0., encoder_freeze=True, encoder_name='miche-256-feature', use_pc=True, time_embed_scale=100.0): super().__init__() self.vocab_A = vocab_A self.vocab_B = vocab_B self.d_model = d_model # MASK Token 的 ID，我们将其设为 B（因为真实的 refine_token 是 0 到 B-1） self.mask_token_id = vocab_B # 1. 词嵌入 (Embeddings) self.raw_embedding = nn.Embedding(vocab_A, d_model) # refine_token 的词表大小是 B + 1 (包含了 MASK token) self.refine_embedding = nn.Embedding(vocab_B + 1, d_model) # 2. 位置编码 (Positional Encoding) # self.pos_embedding = nn.Embedding(max_seq_len, d_model) # 3. 时间步编码 (Timestep Embedding) self.time_mlp = nn.Sequential( SinusoidalPosEmb(d_model, scale=time_embed_scale), nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Linear(d_model * 4, d_model) ) # 4. 主干网络 (Denoising Backbone) # U-Net/Hourglass 主干: # enc(9) -> down -> enc(3) -> down -> bottleneck(1) -> up+skip -> dec(3) -> up+skip -> dec(9) self.depth = depth self.enc_blocks_9 = nn.ModuleList( [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[0])] ) self.enc_blocks_3 = nn.ModuleList( [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[1])] ) self.bottleneck_blocks = nn.ModuleList( [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[2])] ) self.dec_blocks_3 = nn.ModuleList( [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[1])] ) self.dec_blocks_9 = nn.ModuleList( [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[0])] ) self.down_9_to_3 = LinearDownsample(d_model, shorten_factor=3) self.down_3_to_1 = LinearDownsample(d_model, shorten_factor=3) self.up_1_to_3 = LinearUpsample(d_model, shorten_factor=3) self.up_3_to_9 = LinearUpsample(d_model, shorten_factor=3) self.rotary_emb = RotaryEmbedding(dim=d_model//nhead, max_position_embeddings=max_seq_len) # 5. 输出预测头 (只预测真实的 B 个 Token) self.lm_head = nn.Linear(d_model, vocab_B) self.use_pc = use_pc if use_pc: print(f'Point cloud encoder: {encoder_name} | freeze: {encoder_freeze}') self.conditioner = PointConditioner(model_name=encoder_name, freeze=encoder_freeze, feature_dim=d_model) # Learnable null condition for classifier-free guidance. # Replaces self.conditioner(pc) when pc is None, so the model can learn # an "unconditional" representation during cond-drop fine-tuning. self.null_cond_len = 257 # matches PointConditioner.cond_length self.null_cond = nn.Parameter(torch.randn(self.null_cond_len, d_model) * 0.02) def forward(self, raw_tokens, noisy_refine_tokens, t, pc=None): """ raw_tokens: [b, n], 取值 0~A-1 noisy_refine_tokens: [b, n], 取值 0~B (B是MASK) t: [b], 当前的时间步 pc: [b, n, 6], 点云数据 """ if isinstance(raw_tokens, (list, tuple)): return self.forward_varlen(raw_tokens, noisy_refine_tokens, t, pc) b, n = raw_tokens.shape device = raw_tokens.device # 获取 Embedding raw_emb = self.raw_embedding(raw_tokens) # [b, n, d_model] refine_emb = self.refine_embedding(noisy_refine_tokens) # [b, n, d_model] # 获取位置和时间 Embedding time_emb = self.time_mlp(t).unsqueeze(1) # [b, 1, d_model] # x 作为被去噪序列，cond 作为每层 cross-attn 的条件序列 x = raw_emb + refine_emb + time_emb if self.use_pc: if pc is None: # Classifier-free guidance / cond-drop: use learnable null condition. cond = self.null_cond.unsqueeze(0).expand(b, -1, -1) else: cond = self.conditioner(pc) # [b, 257, d_model] else: cond = None if n % 9 != 0: raise ValueError(f"Sequence length {n} is not compatible with 9->3->1 hourglass compression.") dummy_tensor = torch.empty(1, device=x.device, dtype=torch.float32) def rope_fn_for(seq_len): cos, sin = self.rotary_emb(dummy_tensor, seq_len=seq_len) return lambda q, k: apply_rotary_pos_emb(q, k, cos, sin) # Encoder @ length 9 rope_fn_9 = rope_fn_for(x.shape[1]) for block in self.enc_blocks_9: x = block(x, cond, rope_apply_fn=rope_fn_9) skip_9 = x # Downsample 9 -> 3 x = self.down_9_to_3(x) # Encoder @ length 3 rope_fn_3 = rope_fn_for(x.shape[1]) for block in self.enc_blocks_3: x = block(x, cond, rope_apply_fn=rope_fn_3) skip_3 = x # Downsample 3 -> 1 (or in general: n/3 -> n/9) x = self.down_3_to_1(x) # Bottleneck rope_fn_1 = rope_fn_for(x.shape[1]) for block in self.bottleneck_blocks: x = block(x, cond, rope_apply_fn=rope_fn_1) # Upsample 1 -> 3 and long skip connection x = self.up_1_to_3(x) x = x + skip_3 # Decoder @ length 3 for block in self.dec_blocks_3: x = block(x, cond, rope_apply_fn=rope_fn_3) # Upsample 3 -> 9 and long skip connection x = self.up_3_to_9(x) x = x + skip_9 # Decoder @ length 9 for block in self.dec_blocks_9: x = block(x, cond, rope_apply_fn=rope_fn_9) # 预测 Logits logits = self.lm_head(x) # [b, n, vocab_B] return logits def forward_varlen(self, raw_tokens, noisy_refine_tokens, t, pc=None): device = t.device raw_tokens = [seq.to(device=device, dtype=torch.long) for seq in raw_tokens] noisy_refine_tokens = [seq.to(device=device, dtype=torch.long) for seq in noisy_refine_tokens] lengths = torch.tensor([seq.numel() for seq in raw_tokens], device=device, dtype=torch.long) if len(raw_tokens) == 0 or int(lengths.sum().item()) == 0: raise ValueError("forward_varlen received an empty batch.") if any(a.numel() != b.numel() for a, b in zip(raw_tokens, noisy_refine_tokens)): raise ValueError("raw token lengths and refine token lengths must match per sample.") if torch.any(lengths % 9 != 0): bad = lengths[lengths % 9 != 0].detach().cpu().tolist() raise ValueError(f"Varlen hourglass requires each sequence length to be divisible by 9, got {bad}.") raw_packed = torch.cat(raw_tokens, dim=0) refine_packed = torch.cat(noisy_refine_tokens, dim=0) raw_emb = self.raw_embedding(raw_packed) refine_emb = self.refine_embedding(refine_packed) time_emb = self.time_mlp(t).repeat_interleave(lengths, dim=0) x = raw_emb + refine_emb + time_emb if self.use_pc: if pc is None: # Null condition path for CFG / cond-drop. Replicate the learned # null_cond per sample and flatten to match the varlen layout. bsz = len(raw_tokens) null = self.null_cond # [257, d_model] cond = null.unsqueeze(0).expand(bsz, -1, -1).reshape(-1, null.size(-1)) cond_lengths = torch.full((bsz,), self.null_cond_len, device=device, dtype=torch.long) else: pc = pc.to(device=device) cond = self.conditioner(pc) cond_lengths = torch.full((cond.size(0),), cond.size(1), device=device, dtype=torch.long) cond = cond.reshape(-1, cond.size(-1)) else: cond = None cond_lengths = None def rope_fn_for(cur_lengths): position_ids = lengths_to_position_ids(cur_lengths) cos, sin = self.rotary_emb.forward_positions(position_ids, dtype=x.dtype) return lambda q, k: apply_rotary_pos_emb_varlen(q, k, cos, sin) rope_fn_9 = rope_fn_for(lengths) for block in self.enc_blocks_9: x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_9) skip_9 = x lengths_9 = lengths x, lengths = self.down_9_to_3.forward_varlen(x, lengths) rope_fn_3 = rope_fn_for(lengths) for block in self.enc_blocks_3: x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_3) skip_3 = x lengths_3 = lengths x, lengths = self.down_3_to_1.forward_varlen(x, lengths) rope_fn_1 = rope_fn_for(lengths) for block in self.bottleneck_blocks: x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_1) x, lengths = self.up_1_to_3.forward_varlen(x, lengths) if not torch.equal(lengths, lengths_3): raise RuntimeError("Internal length mismatch after 1->3 upsample.") x = x + skip_3 for block in self.dec_blocks_3: x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_3) x, lengths = self.up_3_to_9.forward_varlen(x, lengths) if not torch.equal(lengths, lengths_9): raise RuntimeError("Internal length mismatch after 3->9 upsample.") x = x + skip_9 for block in self.dec_blocks_9: x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_9) return self.lm_head(x) def cosine_schedule(t, T): """ 余弦调度函数。 t: 当前进度步数 (0 表示起始/全噪，T 表示结束/真实数据) 返回: 需要 mask 的比例 (t=0时返回1.0，t=T时返回0.0) """ # 确保 t 是浮点数以进行安全除法 return torch.cos((t / T) * (math.pi / 2)) return torch.cos((t / T) * (math.pi / 2)) def precompute_topology_tensors(raw_faces, tokens_per_face=9): """ 预处理：完全向量化的拓扑索引构建（在去噪循环外执行一次即可） 返回的是一维的张量，方便在 GPU 上做极速的 Scatter/Gather 计算 """ B, num_faces, _, _ = raw_faces.shape device = raw_faces.device # 转为 Python list 加速哈希构建 raw_faces_list = raw_faces.cpu().tolist() token_indices_list = [] v_indices_list = [] v_shared_count = 0 for b in range(B): vertex_dict = {} for f in range(num_faces): for v_idx in range(3): coord = tuple(raw_faces_list[b][f][v_idx]) # 计算展平后的【全局绝对 Token 索引】 global_token_idx = b * (num_faces * tokens_per_face) + f * tokens_per_face + v_idx if coord not in vertex_dict: vertex_dict[coord] = [] vertex_dict[coord].append(global_token_idx) # 仅保留被共享的顶点 for indices in vertex_dict.values(): if len(indices) > 1: token_indices_list.extend(indices) v_indices_list.extend([v_shared_count] * len(indices)) # 记录属于哪个全局共享顶点 ID v_shared_count += 1 # 如果没有任何共享顶点，直接返回空 if len(token_indices_list) == 0: return None, None, None, 0 # [M] 个共享 Token 的绝对索引 token_idx = torch.tensor(token_indices_list, dtype=torch.long, device=device) # [M] 每个 Token 对应的全局顶点 ID v_idx = torch.tensor(v_indices_list, dtype=torch.long, device=device) # [M] 每个 Token 对应的绝对面片(Face)索引 face_idx = token_idx // tokens_per_face return token_idx, v_idx, face_idx, v_shared_count def sample_refine_tokens_with_topology( model, mask_token_id, raw_tokens, pc, T_infer=10, T_train=100, tokens_per_face=9, use_gumbel=False, gumbel_temp_init=1.0, use_soft_consensus=False, consensus_start_step_ratio=0.0, cfg_scale=0.0, null_pc=None, diag_stats=None, forbidden_token_ids=(0,), forbid_mask_token=True, empty_token_ids=(), empty_token_bias=0.0, ): """ 极致优化的拓扑一致性采样代码（零 Python Loop） empty_token_ids / empty_token_bias: 对"空/背景"占用 token(实测 258,770 各~33%)的 logits 减去 empty_token_bias(负偏置)。>0 -> 空 token 更不易被选 -> 更密的细分(更多面)。默认 no-op。 优化 flag: use_gumbel : 重 mask 时给 confidence 加温度退火的 Gumbel 噪声 (MaskGIT 标配) gumbel_temp_init : 初始温度，随 step 线性退火到 0 use_soft_consensus : 共享顶点用概率分布聚合代替硬多数表决 consensus_start_step_ratio: 在 step / T_infer 达到该值前不做 topology consensus cfg_scale : Classifier-free guidance 强度. 0 = 不开启 (单 forward). >0 时每步多一次 forward (null_pc 作 uncond). logits = logits_c + cfg_scale * (logits_c - logits_u) null_pc : (B, N, C) 作为 uncond 输入. 若为 None, 用 zeros_like(pc). 注: 当前模型未训练 cond-drop, 用 null pc 是 OOD 行为, 效果不保证. diag_stats : 可选 dict, 把每步的诊断信息追加进去 (供测试用) forbidden_token_ids : 禁止模型"决定"采样到的 token id 集合. 这些 id 的 logits 在 softmax 前被设为 -inf, 后续 argmax / consensus 全部 自动避开. 默认 (0,) 即 pad/EOS——堵住"模型中段吐出 0 -> save_stage_prediction 截断 -> mesh 只剩前半段"这条 bug. forbid_mask_token : 同时也禁止预测 mask_token_id (默认 True). mask 是占位符, 模型不应把它当最终答案. """ model.eval() b, n = raw_tokens.shape device = raw_tokens.device num_faces = n // tokens_per_face # 收集要屏蔽的 token id; 真正裁剪并搬到 GPU 留到拿到 logits 后再做 # (mask_token_id 通常等于 vocab_B, 不在输出 logits 的维度内, 直接索引会越界) _forbidden_raw = list(forbidden_token_ids) if forbidden_token_ids is not None else [] if forbid_mask_token: _forbidden_raw.append(int(mask_token_id)) _forbidden_raw = sorted(set(int(x) for x in _forbidden_raw)) forbidden_t = None # 在首次 forward 后用 vocab_size 过滤构建 # --- 1. 预处理 (仅执行一次，耗时极短) --- raw_faces = raw_tokens.view(b, num_faces, 3, 3) token_idx, v_idx, face_idx, v_shared_count = precompute_topology_tensors(raw_faces, tokens_per_face) # 初始化 current_tokens = torch.full((b, n), mask_token_id, dtype=torch.long, device=device) unmasked_flags = torch.zeros((b, n), dtype=torch.bool, device=device) vocab_size = None # 将在第一次 forward 后获取 def cosine_schedule(t, max_val): return torch.cos(t * torch.pi * 0.5) * max_val # 预备 null_pc (CFG) # 三种模式: # - null_pc 是 tensor : 用该 tensor 作为 uncond 输入(老模型可用 zeros_like(pc)) # - null_pc == "learned" : 调用 model(..., pc=None), 走模型的 self.null_cond 路径 # (只有 cond-drop fine-tune 过的模型才会有意义) # - null_pc is None (默认): 用 zeros_like(pc), 行为同老版 use_learned_null = (isinstance(null_pc, str) and null_pc.lower() == "learned") if cfg_scale > 0.0 and not use_learned_null: if null_pc is None: null_pc = torch.zeros_like(pc) null_pc = null_pc.to(device=device, dtype=pc.dtype) for step in range(T_infer): current_progress_ratio = step / T_infer t_model = torch.full((b,), current_progress_ratio, device=device, dtype=torch.float) # [1] 前向预测 (可选 CFG: 双 forward 线性外推) logits_c = model(raw_tokens, current_tokens, t_model, pc) # [b, n, vocab_size] if cfg_scale > 0.0: if use_learned_null: logits_u = model(raw_tokens, current_tokens, t_model, None) else: logits_u = model(raw_tokens, current_tokens, t_model, null_pc) logits = logits_c + cfg_scale * (logits_c - logits_u) else: logits = logits_c if vocab_size is None: vocab_size = logits.shape[-1] # 现在知道真实输出 vocab 大小, 把越界的 id (例如 mask_token_id=vocab_B) # 剔除掉后再构建 forbidden_t. 否则 logits[..., id] 会触发 CUDA assert. _valid = [x for x in _forbidden_raw if 0 <= x < vocab_size] if len(_valid) > 0: forbidden_t = torch.tensor(_valid, dtype=torch.long, device=device) # 禁止采样到 pad / EOS. softmax 前置 -inf, 后续 argmax 与 # consensus 聚合 (软/硬) 都因 probs[forbidden]=0 自动避开. if forbidden_t is not None: logits[..., forbidden_t] = float('-inf') # 空行/空占用 token 负偏置: 减小其 logits -> 更不易被选 -> 更密的细分 if empty_token_bias != 0.0 and len(empty_token_ids) > 0: _et = torch.tensor([int(i) for i in empty_token_ids if 0 <= int(i) < vocab_size], dtype=torch.long, device=device) if _et.numel() > 0: logits[..., _et] = logits[..., _et] - float(empty_token_bias) probs = F.softmax(logits, dim=-1) max_probs, pred_tokens = probs.max(dim=-1) # [b, n] force_mask = torch.zeros((b, n), dtype=torch.bool, device=device) target_vals = None # 最后一步若开启 consensus 会用到 # --------------------------------------------------------- # [2] 拓扑共识 (硬多数表决 / 软概率聚合) # 通过 consensus_start_step_ratio 控制何时开启 # --------------------------------------------------------- apply_consensus = (v_shared_count > 0) and (current_progress_ratio >= consensus_start_step_ratio) if apply_consensus: if use_soft_consensus: # 软共识: 直接 scatter_add 共享 token 的概率分布 shared_probs = probs.view(-1, vocab_size)[token_idx] # [M, V] agg_probs = torch.zeros((v_shared_count, vocab_size), dtype=shared_probs.dtype, device=device) v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size) agg_probs.scatter_add_(0, v_idx_expanded, shared_probs) mode_vals = agg_probs.argmax(dim=-1) target_vals = mode_vals[v_idx] preds_flat = pred_tokens.view(-1) shared_vals = preds_flat[token_idx] is_disagree = (shared_vals != target_vals) else: # 硬多数表决 preds_flat = pred_tokens.view(-1) shared_vals = preds_flat[token_idx] one_hot_vals = F.one_hot(shared_vals, num_classes=vocab_size) # [M, V] votes = torch.zeros((v_shared_count, vocab_size), dtype=torch.long, device=device) v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size) votes.scatter_add_(0, v_idx_expanded, one_hot_vals) mode_vals = votes.argmax(dim=-1) target_vals = mode_vals[v_idx] is_disagree = (shared_vals != target_vals) if is_disagree.any(): disagree_faces = face_idx[is_disagree] force_mask_faces = torch.zeros(b * num_faces, dtype=torch.bool, device=device) force_mask_faces[disagree_faces] = True force_mask_flat = force_mask_faces.repeat_interleave(tokens_per_face) force_mask = force_mask_flat.view(b, n) # --------------------------------------------------------- # [3] 置信度干预与 Mask 更新 # --------------------------------------------------------- # 注意: 这里先克隆，避免 Gumbel 噪声破坏掉之后我们对 unmasked/force 的 ±inf 标记 confidence = max_probs.clone() if use_gumbel: # 温度随进度线性退火到 0 temperature = gumbel_temp_init * (1.0 - current_progress_ratio) if temperature > 0: uniform = torch.rand_like(confidence).clamp_(1e-9, 1.0 - 1e-9) gumbel = -torch.log(-torch.log(uniform)) confidence = confidence + temperature * gumbel # 已确认的位置不应被再 mask confidence[unmasked_flags] = float('inf') if apply_consensus: confidence[force_mask] = -float('inf') unmasked_flags = unmasked_flags & (~force_mask) # 诊断信息 if diag_stats is not None: diag_stats.setdefault('force_count', []).append(int(force_mask.sum().item())) diag_stats.setdefault('avg_conf', []).append(float(max_probs.mean().item())) if step == T_infer - 1: # 最后一步：用共识算出的最终值硬覆盖所有可能残余的不一致 if apply_consensus and target_vals is not None: preds_flat = pred_tokens.view(-1) preds_flat[token_idx] = target_vals pred_tokens = preds_flat.view(b, n) current_tokens = torch.where(unmasked_flags, current_tokens, pred_tokens) else: next_progress_ratio = (step + 1) / T_infer ratio = cosine_schedule(torch.tensor(next_progress_ratio), 1.0).item() num_to_mask = int(ratio * n) max_forced_in_batch = int(force_mask.sum(dim=-1).max().item()) if apply_consensus else 0 actual_k = min(max(num_to_mask, max_forced_in_batch), n) _, mask_indices = torch.topk(-confidence, k=actual_k, dim=-1) new_mask = torch.zeros_like(unmasked_flags) new_mask.scatter_(1, mask_indices, True) # 兜底: force_mask 必入 new_mask (Gumbel 噪声不会让 -inf 漏选, 但保留以防边界条件) new_mask = new_mask | force_mask current_tokens = torch.where(new_mask, mask_token_id, pred_tokens) unmasked_flags = ~new_mask return current_tokens def precompute_topology_tensors_fm(raw_faces, tokens_per_face=9): """precompute_topology_tensors 的扩展版: 除了 F(refine 位置 0/1/2 = 3 个顶点/父的占用, 按顶点 coord 分组), 额外把 intra(M_inner, refine 位置 3/4/5, 分别对应父 0/1/2 的内部连接 矩阵)也按其父 coord 分组。F 与 intra 用不同"通道标签"分开归组, 绝不互相表决。 token 布局(与 fix3 同构, 见 tokenizer decode 的 chunk[9..17]): refine[0,1,2] = F0,F1,F2(父 0/1/2 占用) -> 按父 coord 分组(每父冗余, 应一致) refine[3,4,5] = M00,M11,M22(父 0/1/2 intra)-> 按父 coord 分组(每父冗余, 应一致) refine[6,7,8] = M01,M02,M12(inter, 每面本地)-> 不分组、不表决 返回 (token_idx, v_idx, face_idx, v_shared_count), 结构与原版完全一致, 可直接喂给 采样循环(F 组的众数落在 F 值域, intra 组的众数落在 intra 值域, 天然不串)。 """ B, num_faces, _, _ = raw_faces.shape device = raw_faces.device raw_faces_list = raw_faces.cpu().tolist() # 每个"冗余通道" = (refine 位置, 该位置对应的父编号, 通道标签)。 # F: 位置 p -> 父 p; intra: 位置 3+p -> 父 p。inter(6/7/8)不在其中。 channels = [(0, 0, "F"), (1, 1, "F"), (2, 2, "F"), (3, 0, "I"), (4, 1, "I"), (5, 2, "I")] token_indices_list = [] v_indices_list = [] v_shared_count = 0 for b in range(B): group_dict = {} for f in range(num_faces): for pos, parent, tag in channels: coord = tuple(raw_faces_list[b][f][parent]) # 该父的坐标 key = (tag, coord) # F / intra 分开归组 global_token_idx = b * (num_faces * tokens_per_face) + f * tokens_per_face + pos group_dict.setdefault(key, []).append(global_token_idx) for indices in group_dict.values(): if len(indices) > 1: # 仅保留被共享(>1 面引用)的 token_indices_list.extend(indices) v_indices_list.extend([v_shared_count] * len(indices)) v_shared_count += 1 if len(token_indices_list) == 0: return None, None, None, 0 token_idx = torch.tensor(token_indices_list, dtype=torch.long, device=device) v_idx = torch.tensor(v_indices_list, dtype=torch.long, device=device) face_idx = token_idx // tokens_per_face return token_idx, v_idx, face_idx, v_shared_count def sample_refine_tokens_with_topology_fm( model, mask_token_id, raw_tokens, pc, T_infer=10, T_train=100, tokens_per_face=9, use_gumbel=False, gumbel_temp_init=1.0, use_soft_consensus=False, consensus_start_step_ratio=0.0, cfg_scale=0.0, null_pc=None, diag_stats=None, forbidden_token_ids=(0,), forbid_mask_token=True, ): """sample_refine_tokens_with_topology 的"F + M_inner(intra)双对齐"版。 与原版唯一区别: 拓扑分组用 precompute_topology_tensors_fm —— 原版只把 refine 位置 0/1/2(F, 顶点占用)按共享顶点分组表决; 本版额外把 3/4/5(intra, 每父内部连接矩阵)也按共享父分组表决。 因此每步的共识检测/末步硬覆盖会同时作用在 F 与 intra 上, 任一不一致的面整块 9 token 重掩重预测。inter(6/7/8)每面本地、无冗余, 仍不表决(但会随所在面被连坐重掩)。 其余采样逻辑(CFG 双前向、forbidden、Gumbel、cosine 揭示)与原版逐行一致。 参数含义见 sample_refine_tokens_with_topology。 """ model.eval() b, n = raw_tokens.shape device = raw_tokens.device num_faces = n // tokens_per_face _forbidden_raw = list(forbidden_token_ids) if forbidden_token_ids is not None else [] if forbid_mask_token: _forbidden_raw.append(int(mask_token_id)) _forbidden_raw = sorted(set(int(x) for x in _forbidden_raw)) forbidden_t = None # --- 1. 预处理: F + intra 双通道拓扑分组(唯一改动) --- raw_faces = raw_tokens.view(b, num_faces, 3, 3) token_idx, v_idx, face_idx, v_shared_count = precompute_topology_tensors_fm(raw_faces, tokens_per_face) current_tokens = torch.full((b, n), mask_token_id, dtype=torch.long, device=device) unmasked_flags = torch.zeros((b, n), dtype=torch.bool, device=device) vocab_size = None def cosine_schedule(t, max_val): return torch.cos(t * torch.pi * 0.5) * max_val use_learned_null = (isinstance(null_pc, str) and null_pc.lower() == "learned") if cfg_scale > 0.0 and not use_learned_null: if null_pc is None: null_pc = torch.zeros_like(pc) null_pc = null_pc.to(device=device, dtype=pc.dtype) for step in range(T_infer): current_progress_ratio = step / T_infer t_model = torch.full((b,), current_progress_ratio, device=device, dtype=torch.float) # [1] 前向预测 (可选 CFG) logits_c = model(raw_tokens, current_tokens, t_model, pc) if cfg_scale > 0.0: if use_learned_null: logits_u = model(raw_tokens, current_tokens, t_model, None) else: logits_u = model(raw_tokens, current_tokens, t_model, null_pc) logits = logits_c + cfg_scale * (logits_c - logits_u) else: logits = logits_c if vocab_size is None: vocab_size = logits.shape[-1] _valid = [x for x in _forbidden_raw if 0 <= x < vocab_size] if len(_valid) > 0: forbidden_t = torch.tensor(_valid, dtype=torch.long, device=device) if forbidden_t is not None: logits[..., forbidden_t] = float('-inf') probs = F.softmax(logits, dim=-1) max_probs, pred_tokens = probs.max(dim=-1) force_mask = torch.zeros((b, n), dtype=torch.bool, device=device) target_vals = None # [2] 拓扑共识 (F + intra 同时, 因两者在 v_idx 里是不同组, 自动分别表决) apply_consensus = (v_shared_count > 0) and (current_progress_ratio >= consensus_start_step_ratio) if apply_consensus: if use_soft_consensus: shared_probs = probs.view(-1, vocab_size)[token_idx] agg_probs = torch.zeros((v_shared_count, vocab_size), dtype=shared_probs.dtype, device=device) v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size) agg_probs.scatter_add_(0, v_idx_expanded, shared_probs) mode_vals = agg_probs.argmax(dim=-1) target_vals = mode_vals[v_idx] preds_flat = pred_tokens.view(-1) shared_vals = preds_flat[token_idx] is_disagree = (shared_vals != target_vals) else: preds_flat = pred_tokens.view(-1) shared_vals = preds_flat[token_idx] one_hot_vals = F.one_hot(shared_vals, num_classes=vocab_size) votes = torch.zeros((v_shared_count, vocab_size), dtype=torch.long, device=device) v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size) votes.scatter_add_(0, v_idx_expanded, one_hot_vals) mode_vals = votes.argmax(dim=-1) target_vals = mode_vals[v_idx] is_disagree = (shared_vals != target_vals) if is_disagree.any(): disagree_faces = face_idx[is_disagree] force_mask_faces = torch.zeros(b * num_faces, dtype=torch.bool, device=device) force_mask_faces[disagree_faces] = True force_mask_flat = force_mask_faces.repeat_interleave(tokens_per_face) force_mask = force_mask_flat.view(b, n) # [3] 置信度干预与 Mask 更新 confidence = max_probs.clone() if use_gumbel: temperature = gumbel_temp_init * (1.0 - current_progress_ratio) if temperature > 0: uniform = torch.rand_like(confidence).clamp_(1e-9, 1.0 - 1e-9) gumbel = -torch.log(-torch.log(uniform)) confidence = confidence + temperature * gumbel confidence[unmasked_flags] = float('inf') if apply_consensus: confidence[force_mask] = -float('inf') unmasked_flags = unmasked_flags & (~force_mask) if diag_stats is not None: diag_stats.setdefault('force_count', []).append(int(force_mask.sum().item())) diag_stats.setdefault('avg_conf', []).append(float(max_probs.mean().item())) if step == T_infer - 1: if apply_consensus and target_vals is not None: preds_flat = pred_tokens.view(-1) preds_flat[token_idx] = target_vals pred_tokens = preds_flat.view(b, n) current_tokens = torch.where(unmasked_flags, current_tokens, pred_tokens) else: next_progress_ratio = (step + 1) / T_infer ratio = cosine_schedule(torch.tensor(next_progress_ratio), 1.0).item() num_to_mask = int(ratio * n) max_forced_in_batch = int(force_mask.sum(dim=-1).max().item()) if apply_consensus else 0 actual_k = min(max(num_to_mask, max_forced_in_batch), n) _, mask_indices = torch.topk(-confidence, k=actual_k, dim=-1) new_mask = torch.zeros_like(unmasked_flags) new_mask.scatter_(1, mask_indices, True) new_mask = new_mask | force_mask current_tokens = torch.where(new_mask, mask_token_id, pred_tokens) unmasked_flags = ~new_mask return current_tokens def precompute_topology_tensors_fmi(raw_faces, tokens_per_face=9): """F + intra + inter 三通道拓扑分组(逐步共识用)。 F (refine 0/1/2): 按单父 coord 分组(每父冗余, 引用同父的面应一致) intra(refine 3/4/5): 按单父 coord 分组(每父冗余) inter(refine 6/7/8): 按"无序父对 coord"分组 —— 两个面若恰好共享 2 顶点(=共享一条粗边) 就共享同一父对, 其 inter 矩阵应一致; tokenizer 每面 3 父排序+按 (小,大)取 inter, 同一父对在任何面里 token 值相同(朝向天然规范)。 三通道用不同 key 前缀("F"/"I"/"O")互不表决。返回结构与 precompute_topology_tensors 完全一致, 可直接喂采样循环(各通道众数落各自值域, 不串)。 """ B, num_faces, _, _ = raw_faces.shape device = raw_faces.device raw_faces_list = raw_faces.cpu().tolist() single_ch = [(0, 0, "F"), (1, 1, "F"), (2, 2, "F"), (3, 0, "I"), (4, 1, "I"), (5, 2, "I")] pair_ch = [(6, 0, 1), (7, 0, 2), (8, 1, 2)] # inter 位置 -> 父对 token_indices_list, v_indices_list, v_shared_count = [], [], 0 for b in range(B): gd = {} for f in range(num_faces): fc = raw_faces_list[b][f] base = b * (num_faces * tokens_per_face) + f * tokens_per_face for pos, parent, tag in single_ch: gd.setdefault((tag, tuple(fc[parent])), []).append(base + pos) for pos, a, bp in pair_ch: ca, cb = tuple(fc[a]), tuple(fc[bp]) pair = (ca, cb) if ca <= cb else (cb, ca) gd.setdefault(("O", pair), []).append(base + pos) for indices in gd.values(): if len(indices) > 1: token_indices_list.extend(indices) v_indices_list.extend([v_shared_count] * len(indices)) v_shared_count += 1 if len(token_indices_list) == 0: return None, None, None, 0 token_idx = torch.tensor(token_indices_list, dtype=torch.long, device=device) v_idx = torch.tensor(v_indices_list, dtype=torch.long, device=device) face_idx = token_idx // tokens_per_face return token_idx, v_idx, face_idx, v_shared_count