File size: 3,324 Bytes
90ca518 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """fix3 编码格式的数据通路工具 —— 从训练脚本抽出, 保证推理与训练用同一份实现。
split_fix3_token_list : 扁平 60/face token -> (coord_list[9*nface], subcode_list[9*nface,8] local)
subcodes_to_flat_token : (coord, 采样得到的 subcodes) -> fix3 扁平 60/face token(含 offset), 供 decode_tokens
"""
import torch
from models.generator import REFINE_POS_PER_FACE, NUM_OCTANTS
TOKENS_PER_FACE = 60 # fix3 扁平: 9 coord + 3 F + 6 矩阵 x 8 行
def _trim_trailing_pad(token, pad_id):
pad_pos = (token == pad_id).nonzero(as_tuple=True)[0]
return token[:pad_pos[0].item()] if pad_pos.numel() > 0 else token
def split_fix3_token_list(token_list, pad_id, device, codec_cfg):
"""扁平 60/face token -> (coord_list, subcode_list)。
coord_list[i] : [9*nface] 坐标(保留 coord_offset, 直接喂 raw_embedding)
subcode_list[i] : [9*nface, 8] refine 子码(去 offset, local[0,256))
位置顺序 [F0 F1 F2 M00 M11 M22 M01 M02 M12]
F 位置只用 slot0(其余槽填 0, 模型按类型只读 slot0)；intra/inter 用 8 行
"""
f_off = codec_cfg["f_offset"]; ia = codec_cfg["intra_row_offset"]; ie = codec_cfg["inter_row_offset"]
coord_list, sub_list = [], []
for tok in token_list:
tok = _trim_trailing_pad(tok.to(device=device, dtype=torch.long), pad_id)
if tok.numel() % TOKENS_PER_FACE != 0:
raise ValueError(f"token 长度 {tok.numel()} 不是 {TOKENS_PER_FACE} 的倍数")
tok = tok.view(-1, TOKENS_PER_FACE) # [nface, 60]
nface = tok.size(0)
coord = tok[:, :9].reshape(-1) # [9*nface]
Fv = tok[:, 9:12] - f_off # [nface, 3]
intra = (tok[:, 12:36] - ia).view(nface, 3, NUM_OCTANTS) # [nface,3,8]
inter = (tok[:, 36:60] - ie).view(nface, 3, NUM_OCTANTS) # [nface,3,8]
sub = torch.zeros(nface, REFINE_POS_PER_FACE, NUM_OCTANTS, dtype=torch.long, device=device)
sub[:, 0:3, 0] = Fv # F -> slot0
sub[:, 3:6, :] = intra
sub[:, 6:9, :] = inter
sub_list.append(sub.reshape(-1, NUM_OCTANTS)) # [9*nface, 8]
coord_list.append(coord)
return coord_list, sub_list
def subcodes_to_flat_token(coord_seq, sub_seq, codec_cfg):
"""采样得到的 (coord[9*nface], subcodes[9*nface,8] local) 还原成 fix3 扁平 60/face token(含 offset)，
供 codec.decode_tokens 使用。F 位置取 slot0。"""
f_off = codec_cfg["f_offset"]; ia = codec_cfg["intra_row_offset"]; ie = codec_cfg["inter_row_offset"]
nface = coord_seq.numel() // REFINE_POS_PER_FACE
coord = coord_seq.view(nface, REFINE_POS_PER_FACE) # [nface,9]
sub = sub_seq.view(nface, REFINE_POS_PER_FACE, NUM_OCTANTS) # [nface,9,8]
Fv = sub[:, 0:3, 0] + f_off # [nface,3] 取 slot0
intra = sub[:, 3:6, :] + ia # [nface,3,8]
inter = sub[:, 6:9, :] + ie
flat = torch.cat([coord, Fv, intra.reshape(nface, 24), inter.reshape(nface, 24)], dim=1) # [nface,60]
return flat.reshape(-1).tolist()
|