models/generator.py

"""fix3 生成器模型(已融合: 基类 + alt-PC 双云注入 + TSSR 三前向)。

类继承链(融合后仍保留, 因各层职责清晰):
    FactorizedDiffusionTransformer      骨干: 因子分解 token / AdaLN cross-attn / 自回归行头
      └ FactorizedDiffusionTransformerAltPC   逐层交替双点云注入 + 几何 RoPE + sample()
          └ FactorizedTSSRTransformer         mode flag + 逐 token 正确性分类器 + sample_tssr()

原先的 FactorizedDiffusionTransformerFocal 一层只重写训练用的 factorized_loss, 推理路径上是空壳,
已移除并让 AltPC 直接继承骨干。各类的纯训练方法(factorized_loss / factorized_accuracy /
sample_log / forward_mask / forward_nomask / zero_touch_conditional_params)也已剥离。
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from models.subdivider import (
    SinusoidalPosEmb,
    ModernCrossAttnTransformerBlock,
    ModernAttention,
    RMSNorm,
    SwiGLU,
    RotaryEmbedding,
    LinearDownsample,
    LinearUpsample,
    apply_rotary_pos_emb,
    apply_rotary_pos_emb_varlen,
    lengths_to_position_ids,
    cosine_schedule,
)
from models.point_conditioner import PointConditioner
import os
import numpy as np
from models.subdivider import (
    apply_rotary_pos_emb_varlen, lengths_to_position_ids, cosine_schedule,
)


# ==========================================================================
# 原 model_acc_fix3.py
# ==========================================================================

def modulate(x, shift, scale):
    """adaLN 调制: x * (1 + scale) + shift。shift/scale 与 x 同形状(逐 token 广播)。"""
    return x * (1 + scale) + shift
class AdaLNCrossAttnBlock(nn.Module):
    """fix3 专用 block：在 self-attn 与 FFN 上加 adaLN 调制(条件 c = time + depth)。
    cross-attn 仍走点云条件(不调制)。结构与 ModernCrossAttnTransformerBlock 同名(self_attn/
    cross_attn/norm1-3/ffn)以便迁移 fix2 权重；adaLN_modulation 为**新增**且做「恒等初始化」
    (shift=0, scale=0, gate=1)，使加载迁移权重时 block == 旧 block，backbone 不被破坏。"""

    def __init__(self, d_model, nhead, dropout=0., qk_norm=False, ffn_hidden=None):
        super().__init__()
        # QK-norm 只加在 self-attn(instability 主源);cross-attn 的 K 是点云条件,保持默认不动
        self.self_attn = ModernAttention(d_model, nhead, dropout=dropout, qk_norm=qk_norm)
        self.cross_attn = ModernAttention(d_model, nhead, dropout=dropout)
        self.norm1 = RMSNorm(d_model)
        self.norm2 = RMSNorm(d_model)
        self.norm3 = RMSNorm(d_model)
        self.ffn = SwiGLU(d_model, ffn_hidden)
        self.dropout = nn.Dropout(dropout)
        # 从条件 c 产生 self-attn / FFN 各自的 (shift, scale, gate)，共 6*d_model
        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(d_model, 6 * d_model))
        # 恒等初始化：weight=0；bias 使 shift=0, scale=0, gate=1 -> block 退化为旧 block
        nn.init.zeros_(self.adaLN_modulation[1].weight)
        with torch.no_grad():
            b = torch.zeros(6 * d_model)
            b[2 * d_model:3 * d_model] = 1.0   # gate_sa = 1
            b[5 * d_model:6 * d_model] = 1.0   # gate_ff = 1
            self.adaLN_modulation[1].bias.copy_(b)

    def forward_varlen(self, x, c, lengths, cond=None, cond_lengths=None, rope_apply_fn=None,
                       max_len=None, cond_max_len=None):
        """c: [total, d_model] 逐 token 条件(time+depth)。
        max_len / cond_max_len: 透传给 attention 的 max_seqlen(python int), 省掉 GPU->CPU 同步;
        不传则由 attention 内部回退到 `.item()`(见 ModernAttention.forward_varlen 的说明)。

        注: 曾试过「c 只传 [b,d]、在块内 adaLN 后再 repeat_interleave」以省掉 768->6*768 Linear
        在同一行上重复算 total 次的冗余(每块 424 GFLOP@60k token) —— 实测输出逐位相同但**零加速**
        (hourglass 已把 token 降到 1/3、1/9, 那 3.1 TFLOP 不是瓶颈), 且会让 torch.compile 报
        dtype 不匹配。故回退, 别再试。
        """
        # c 来自 time_mlp + depth_embedding, 是 fp32(nn.Embedding 不在 autocast 转型名单里)。
        # 对齐到 x 的 dtype: eager+autocast 下等价(Linear 本来就会转), 但 torch.compile 编译这个方法时
        # 会丢掉外层 autocast 上下文, 不显式对齐就会在后面的 matmul 上报 "Float and BFloat16"。
        c = c.to(x.dtype)
        shift_sa, scale_sa, gate_sa, shift_ff, scale_ff, gate_ff = self.adaLN_modulation(c).chunk(6, dim=-1)
        h = modulate(self.norm1(x), shift_sa, scale_sa)
        sa = self.self_attn.forward_varlen(h, h, h, lengths, rope_apply_fn=rope_apply_fn,
                                           max_len_q=max_len)
        x = x + gate_sa * self.dropout(sa)
        if cond is not None:
            h = self.norm2(x)
            # CFG drop 的步里 cond 来自 fp32 的 null_cond; 同 c 一样显式对齐(见上), 否则 compile 会崩
            ca = self.cross_attn.forward_varlen(h, cond.to(x.dtype), cond.to(x.dtype), lengths,
                                                lengths_k=cond_lengths,
                                                max_len_q=max_len, max_len_k=cond_max_len)
            x = x + self.dropout(ca)
        h = modulate(self.norm3(x), shift_ff, scale_ff)
        ff = self.ffn(h)
        x = x + gate_ff * self.dropout(ff)
        return x
SUB_VOCAB = 256          # 每个子码(8bit)的取值数
MASK_ID = SUB_VOCAB      # 子码 MASK id = 256（embedding 表大小 = 257）
NUM_OCTANTS = 8
TYPE_F, TYPE_INTRA, TYPE_INTER = 0, 1, 2
REFINE_POS_PER_FACE = 9  # 3 F + 3 intra + 3 inter
class AutoregressiveRowHead(nn.Module):
    """连接矩阵(intra/inter)的 8 行**自回归** MLP 头(参考 TreeMeshGPT 的 head_coord1/2/3)。

    取代原先「8 行并行独立预测」：行 r 的 256-way logits 由
        cat([hidden, emb(行0), ..., emb(行r-1)])  -> 3 层 MLP
    给出，从而**显式建模行间依赖**，让模型只产出能共现的合法矩阵(对称/可成 clique)，
    减少独立预测导致的「悬边→退化面/碎块」。
      * 训练: teacher forcing(回喂 GT 行 embedding)，一次 forward 出全 8 行 logits。
      * 采样: 逐行 argmax rollout(回喂已采行)；支持 CFG —— cond/uncond 在每行 logits 级合并，
              合并后采样得到的同一行回喂两路，保持自回归状态一致。
    行 embedding 维度默认 = d_model//8(与输入端 row_embed_dim 一致，控参数量)；最后一行不作输入故只需 7 张表。
    """

    def __init__(self, d_model, num_rows=NUM_OCTANTS, vocab=SUB_VOCAB, row_emb_dim=None):
        super().__init__()
        self.num_rows = num_rows
        self.vocab = vocab
        self.row_emb_dim = row_emb_dim or (d_model // num_rows)
        self.mlps = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model + r * self.row_emb_dim, d_model), nn.ReLU(),
                nn.Linear(d_model, d_model), nn.ReLU(),
                nn.Linear(d_model, vocab),
            ) for r in range(num_rows)
        ])
        self.row_emb = nn.ModuleList([nn.Embedding(vocab, self.row_emb_dim) for _ in range(num_rows - 1)])

    def forward_teacher(self, h, gt_rows):
        """h [n,d]；gt_rows [n,num_rows] long(取值[0,vocab)) -> logits [n,num_rows,vocab]。teacher forcing。"""
        feats = [h]
        logits = []
        for r in range(self.num_rows):
            logits.append(self.mlps[r](torch.cat(feats, dim=-1)))
            if r < self.num_rows - 1:
                feats.append(self.row_emb[r](gt_rows[:, r]).to(h.dtype))
        return torch.stack(logits, dim=1)

    def rollout(self, h, h_u=None, cfg_scale=1.0, temperature=0.0):
        """逐行采样。h [n,d]；h_u 提供则做 CFG(logits=uncond+w*(cond-uncond))。
        temperature<=0: argmax 贪心(确定性)；>0: 按 (logit/temperature) 多项式采样(用于 best-of-N 产生多样候选)。
        返回 (pred [n,num_rows] long, conf [n,num_rows] 每行所选类的概率)。"""
        feats = [h]
        feats_u = [h_u] if h_u is not None else None
        preds, confs = [], []
        for r in range(self.num_rows):
            logit = self.mlps[r](torch.cat(feats, dim=-1))
            if feats_u is not None:
                lu = self.mlps[r](torch.cat(feats_u, dim=-1))
                logit = lu + cfg_scale * (logit - lu)
            if temperature and temperature > 0:
                prob = (logit / temperature).softmax(-1)
                am = torch.multinomial(prob.float(), 1).squeeze(-1)              # [n]
                pm = prob.gather(-1, am.unsqueeze(-1)).squeeze(-1)               # 所选类概率
            else:
                prob = logit.softmax(-1)
                pm, am = prob.max(-1)
            preds.append(am); confs.append(pm)
            if r < self.num_rows - 1:
                e = self.row_emb[r](am).to(h.dtype)
                feats.append(e)
                if feats_u is not None:
                    feats_u.append(e)
        return torch.stack(preds, dim=1), torch.stack(confs, dim=1)
class FactorizedDiffusionTransformer(nn.Module):
    def __init__(self,
                 coord_vocab_size,         # raw_embedding 大小 = resolution+1 (coord 含 offset, 0=pad)
                 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,
                 num_depths=16,            # depth(层级) embedding 大小，覆盖 target_depth 0..max
                 nonempty_row_weight=1.0,  # 连接矩阵「非空行」(目标!=0, 即真实连接) 的 loss 上权
                 qk_norm=False,            # self-attn 的 Q/K 加 per-head RMSNorm(稳定 attention);默认关
                 ffn_hidden=None,          # FFN 隐层维度显式覆盖(None=8/3*d_model; 1024宽用2816对齐TensorCore)
                 pc_num_latents=None,      # 点云编码器 latent 数(None=预训练默认256; 1024=条件带宽x4, query平铺扩容)
                 time_embed_dim=None):     # 正弦基底/time_mlp 隐层维度(None=d_model, 默认即可)。备注:
                                           # SinusoidalPosEmb 的"列-频率/sin-cos"布局随 dim 变化, 宽度扩展
                                           # 热启动时旧 time_mlp 权重不可拷贝(语义错位会爆), 见扩展脚本。
        super().__init__()
        self.d_model = d_model
        self.mask_token_id = MASK_ID
        self.num_depths = num_depths
        # 连接(intra/inter)行子码中「非空行」(target!=0) 的 loss 权重。connectivity 矩阵 ~92% 是空行，
        # 等权 CE 会让模型坍缩到「全预测空行」(逐子码 acc 虚高、生成 mesh 欠连接碎裂)。>1 上权真实连接行。
        self.nonempty_row_weight = float(nonempty_row_weight)

        # 1. Embeddings ----------------------------------------------------
        self.raw_embedding = nn.Embedding(coord_vocab_size, d_model)        # coord(条件)
        self.f_sub_embedding = nn.Embedding(SUB_VOCAB + 1, d_model)         # F: 单子码
        # intra/inter: 8 个行子码各查表(+1 行给 MASK)，再 cat 后 Linear 投影到 d_model。
        # 用 cat 而非 sum：保留 8 行的「行序/位置语义」(行 r = octant r)，避免 sum 的置换不变与叠加丢失；
        # 行 embedding 维度 = d_model // 8，cat 后恰好 = d_model，再过一层 Linear 学行间组合(含 intra 对称)。
        self.row_embed_dim = d_model // NUM_OCTANTS
        self.intra_sub_embedding = nn.Embedding(SUB_VOCAB + 1, self.row_embed_dim)
        self.inter_sub_embedding = nn.Embedding(SUB_VOCAB + 1, self.row_embed_dim)
        self.intra_in_proj = nn.Linear(NUM_OCTANTS * self.row_embed_dim, d_model)
        self.inter_in_proj = nn.Linear(NUM_OCTANTS * self.row_embed_dim, d_model)

        # 2. 条件: 时间步 + 层级(depth) —— 一起作为 adaLN 的条件 c -------------
        te = int(time_embed_dim) if time_embed_dim else d_model
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(te, scale=time_embed_scale),
            nn.Linear(te, te * 4),
            nn.GELU(),
            nn.Linear(te * 4, d_model),
        )
        self.depth_embedding = nn.Embedding(num_depths, d_model)   # 层级 embedding(随机初始化)

        # 3. 主干: 带 adaLN 的 block(条件 c = time + depth)；blocks 内 self_attn/cross_attn/
        #    norm/ffn 与 ModernCrossAttnTransformerBlock 同名可迁移，adaLN_modulation 为新增(恒等初始化)
        self.depth = list(depth)
        _blk = lambda: AdaLNCrossAttnBlock(d_model, nhead, dropout=dropout, qk_norm=qk_norm, ffn_hidden=ffn_hidden)
        self.enc_blocks_9 = nn.ModuleList([_blk() for _ in range(depth[0])])
        self.enc_blocks_3 = nn.ModuleList([_blk() for _ in range(depth[1])])
        self.bottleneck_blocks = nn.ModuleList([_blk() for _ in range(depth[2])])
        self.dec_blocks_3 = nn.ModuleList([_blk() for _ in range(depth[1])])
        self.dec_blocks_9 = nn.ModuleList([_blk() 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)

        # 4. 因子分解输出头 ------------------------------------------------
        self.f_head = nn.Linear(d_model, SUB_VOCAB)                          # F: 1 个子码(单头)
        # intra/inter: 8 行**自回归** MLP 头(行间依赖, 取代旧的 8 行并行独立 Linear)
        self.intra_head = AutoregressiveRowHead(d_model, num_rows=NUM_OCTANTS, vocab=SUB_VOCAB)
        self.inter_head = AutoregressiveRowHead(d_model, num_rows=NUM_OCTANTS, vocab=SUB_VOCAB)

        # 5. 点云条件 ------------------------------------------------------
        self.use_pc = use_pc
        if use_pc:
            print(f'[fix3] Point cloud encoder: {encoder_name} | freeze: {encoder_freeze}')
            self.conditioner = PointConditioner(model_name=encoder_name, freeze=encoder_freeze, feature_dim=d_model,
                                                num_latents=pc_num_latents)
            self.null_cond_len = 257
            self.null_cond = nn.Parameter(torch.randn(self.null_cond_len, d_model) * 0.02)

    # ------------------------------------------------------------------ #
    # 工具
    # ------------------------------------------------------------------ #
    def types_from_lengths(self, lengths):
        """由每个样本长度(=9*nface)推出 packed 的位置类型 [total]，0/1/2 = F/intra/inter。"""
        parts = []
        for L in lengths.tolist():
            idx = torch.arange(L, device=lengths.device)
            fp = idx % REFINE_POS_PER_FACE
            t = torch.where(fp < 3, TYPE_F,
                            torch.where(fp < 6, TYPE_INTRA, TYPE_INTER))
            parts.append(t)
        return torch.cat(parts).long() if parts else torch.empty(0, dtype=torch.long, device=lengths.device)

    def embed_refine(self, subcodes, types):
        """subcodes [N,8] long, types [N] long -> [N,d]。
        F 位置: 单子码查表。intra/inter 位置: 8 行各查表 -> cat -> Linear(保留行序，避免 sum 丢失)。"""
        N = subcodes.size(0)
        out = self.f_sub_embedding.weight.new_zeros((N, self.d_model))
        is_f = types == TYPE_F
        is_in = types == TYPE_INTRA
        is_ix = types == TYPE_INTER
        # autocast(bf16) 下 Linear 输出是 bf16，而 out 是 fp32；原地 index_put 要求 dtype 一致，
        # 故统一 .to(out.dtype)（保持 refine_emb 与 raw_emb 一致的 fp32，backbone matmul 再由 autocast 转 bf16）。
        if is_f.any():
            out[is_f] = self.f_sub_embedding(subcodes[is_f, 0]).to(out.dtype)
        if is_in.any():
            e = self.intra_sub_embedding(subcodes[is_in])              # [n, 8, row_embed_dim]
            out[is_in] = self.intra_in_proj(e.reshape(e.size(0), -1)).to(out.dtype)  # cat -> Linear -> [n, d]
        if is_ix.any():
            e = self.inter_sub_embedding(subcodes[is_ix])
            out[is_ix] = self.inter_in_proj(e.reshape(e.size(0), -1)).to(out.dtype)
        return out

    # ------------------------------------------------------------------ #
    # 前向(varlen)：返回主干 hidden h [total, d] 与 packed types [total]
    # ------------------------------------------------------------------ #
    def forward_varlen(self, coord_list, subcode_list, t, depth, pc=None):
        device = t.device
        coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list]
        subcode_list = [s.to(device=device, dtype=torch.long) for s in subcode_list]
        lengths = torch.tensor([c.numel() for c in coord_list], device=device, dtype=torch.long)
        if int(lengths.sum().item()) == 0:
            raise ValueError("forward_varlen received an empty batch.")
        if torch.any(lengths % REFINE_POS_PER_FACE != 0):
            raise ValueError("每个序列长度必须是 9 的倍数(每 face 9 个 refine 位置)。")

        coord_packed = torch.cat(coord_list, dim=0)
        subcode_packed = torch.cat(subcode_list, dim=0)
        types = self.types_from_lengths(lengths)

        raw_emb = self.raw_embedding(coord_packed)
        refine_emb = self.embed_refine(subcode_packed, types)
        x = raw_emb + refine_emb                         # 不再加法注入 time；time/depth 走 adaLN

        # adaLN 条件 c = time_emb + depth_emb（逐样本 [b,d]）；各 hourglass 阶段按其长度展开到逐 token
        depth = depth.to(device=device, dtype=torch.long)
        c_sample = self.time_mlp(t) + self.depth_embedding(depth)   # [b, d_model]

        def c_for(cur_lengths):
            return c_sample.repeat_interleave(cur_lengths, dim=0)

        if self.use_pc:
            if pc is None:
                bsz = len(coord_list)
                null = self.null_cond
                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)
        c9 = c_for(lengths)
        for block in self.enc_blocks_9:
            x = block.forward_varlen(x, c9, 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)
        c3 = c_for(lengths)
        for block in self.enc_blocks_3:
            x = block.forward_varlen(x, c3, 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)
        c1 = c_for(lengths)
        for block in self.bottleneck_blocks:
            x = block.forward_varlen(x, c1, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_1)

        x, lengths = self.up_1_to_3.forward_varlen(x, lengths)
        x = x + skip_3
        for block in self.dec_blocks_3:
            x = block.forward_varlen(x, c3, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_3)

        x, lengths = self.up_3_to_9.forward_varlen(x, lengths)
        x = x + skip_9
        for block in self.dec_blocks_9:
            x = block.forward_varlen(x, c9, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_9)

        return x, types  # hidden [total,d], types [total]

    def forward(self, coord_list, noisy_subcode_list, t, depth,
                target_subcode_list=None, loss_mask=None, pc=None, return_extra_acc=False):
        """训练入口：把因子分解头放进 forward 内部(经 factorized_loss)，保证 DDP 正确同步。
        depth: [b] long，每个样本的层级(target_depth)，与 time 一起作为 adaLN 条件。
        - 传 target/loss_mask -> 返回 (loss, acc)；return_extra_acc=True 时返回 (loss, acc, acc_nonempty)
          (acc_nonempty=连接非空行准确率，见 factorized_accuracy)。默认 False 与旧接口兼容。
        - 不传            -> 返回 (hidden, types)（推理/调试）
        """
        h, types = self.forward_varlen(coord_list, noisy_subcode_list, t, depth, pc)
        if target_subcode_list is None:
            return h, types
        target = torch.cat([s.to(device=h.device, dtype=torch.long) for s in target_subcode_list], dim=0)
        loss_mask = loss_mask.to(device=h.device)
        loss = self.factorized_loss(h, target, types, loss_mask)
        with torch.no_grad():
            acc, acc_ne = self.factorized_accuracy(h, target, types, loss_mask)
        if return_extra_acc:
            return loss, acc, acc_ne
        return loss, acc

    # ------------------------------------------------------------------ #
    # 因子分解 loss / accuracy
    # ------------------------------------------------------------------ #


    # ------------------------------------------------------------------ #
    # 采样(MaskGIT 风格，因子分解、varlen、逐位置置信度解 mask)
    # ------------------------------------------------------------------ #
    @torch.no_grad()
    def sample(self, coord_list, depth, pc=None, T_infer=10, cfg_scale=1.0, temperature=0.0,
               cfg_schedule="constant", cfg_start=1.0):
        """给定 coord(粗层条件)与 depth(层级)，从全 MASK 迭代生成 refine 子码。
        depth: [b] long(每样本层级)。返回 subcode_list(list[ [L,8] ])。
        cfg_scale: classifier-free guidance 强度，约定 logits = uncond + cfg_scale*(cond - uncond)。
                   =1.0 即纯条件(无引导，跳过无条件前向)；>1 放大点云条件(锐化连接预测、对抗欠连接)。
                   仅在 use_pc 且 pc 提供时生效(无条件分支用 pc=None 走训练过的 null_cond)。
        temperature: <=0 argmax 贪心(确定性，默认，与旧行为一致)；>0 多项式随机采样。
                   best-of-N 时需 >0(如 1.0) 才能采出多样候选，否则贪心 N 次结果完全相同。
        cfg_schedule: "constant"=全程恒定 cfg_scale(默认，旧行为)；"linear"=按步从 cfg_start
                   线性升到 cfg_scale(MUSE 式 ramp)。最初几步决定全局结构，CFG 锐化会把条件
                   敏感(=分歧大、没把握)位置的置信度人为抬高、过早 commit 且不可回滚；低 cfg
                   起步可避开这种早期过冲，引导力留给后期细化。
        cfg_start: ramp 起点(仅 linear 生效)。默认 1.0 = 第一步纯条件(并跳过该步无条件前向)。"""
        self.eval()
        device = next(self.parameters()).device
        coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list]
        depth = depth.to(device=device, dtype=torch.long)
        lengths = torch.tensor([c.numel() for c in coord_list], device=device, dtype=torch.long)
        types = self.types_from_lengths(lengths)
        want_cfg = (float(cfg_scale) != 1.0 or float(cfg_start) != 1.0) and self.use_pc and (pc is not None)

        total = int(lengths.sum().item())
        subcodes = torch.full((total, NUM_OCTANTS), MASK_ID, dtype=torch.long, device=device)
        decided = torch.zeros(total, dtype=torch.bool, device=device)

        for step in range(T_infer):
            # 当步实际引导强度：linear 时 step0 = cfg_start，末步 = cfg_scale
            if cfg_schedule == "linear":
                cfg_t = float(cfg_start) + (float(cfg_scale) - float(cfg_start)) * step / max(T_infer - 1, 1)
            else:
                cfg_t = float(cfg_scale)
            use_cfg = want_cfg and (cfg_t != 1.0)

            ratio = step / T_infer
            t = torch.full((len(coord_list),), ratio, device=device, dtype=torch.float)
            sub_list = list(torch.split(subcodes, lengths.tolist(), dim=0))
            h, _ = self.forward_varlen(coord_list, sub_list, t, depth, pc)
            h_u = self.forward_varlen(coord_list, sub_list, t, depth, None)[0] if use_cfg else None

            # 逐类型预测子码 + 置信度(该位置有效子码的平均最大概率)
            pred = subcodes.clone()
            conf = torch.zeros(total, device=device)
            # F: 单头(可 CFG)
            mf = types == TYPE_F
            if mf.any():
                logits = self.f_head(h[mf])                              # [nm,256]
                if use_cfg:
                    lu = self.f_head(h_u[mf])
                    logits = lu + cfg_t * (logits - lu)
                if temperature and temperature > 0:
                    prob = (logits / temperature).softmax(-1)
                    am = torch.multinomial(prob.float(), 1).squeeze(-1)              # [nm]
                    pm = prob.gather(-1, am.unsqueeze(-1)).squeeze(-1)
                else:
                    prob = logits.softmax(-1)
                    pm, am = prob.max(-1)                                # [nm],[nm]
                pred[mf, :] = am.unsqueeze(1).expand(-1, NUM_OCTANTS)    # F 复制到 8 槽
                conf[mf] = pm.to(conf.dtype)
            # intra/inter: 8 行自回归 rollout(可 CFG, cond/uncond 在每行 logits 级合并)
            for ty, head in ((TYPE_INTRA, self.intra_head), (TYPE_INTER, self.inter_head)):
                m = types == ty
                if not m.any():
                    continue
                p, pmaxs = head.rollout(h[m], h_u[m] if use_cfg else None, cfg_t, temperature)  # [nm,8],[nm,8]
                pred[m] = p
                conf[m] = pmaxs.mean(dim=1).to(conf.dtype)   # autocast 下 pmax 可能是 bf16

            # 已决定的位置保持不变
            keep_ratio = cosine_schedule(torch.tensor(float(step + 1)), float(T_infer)).item()
            n_keep_masked = int(keep_ratio * total)  # 仍保持 mask 的目标数(随步数减少)
            conf_for_select = conf.clone()
            conf_for_select[decided] = float('inf')   # 已定的不再参与"新解mask"
            # 选出本步要确定的位置：未定中置信度最高的若干个
            n_to_decide = max(0, (total - n_keep_masked) - int(decided.sum().item()))
            if n_to_decide > 0:
                undecided_idx = (~decided).nonzero(as_tuple=True)[0]
                order = conf[undecided_idx].argsort(descending=True)
                chosen = undecided_idx[order[:n_to_decide]]
                subcodes[chosen] = pred[chosen]
                decided[chosen] = True
            # 未决定位置回到 MASK，已决定保持
            subcodes[~decided] = MASK_ID

        # 收尾：还没定的全部按当前预测填上
        subcodes[~decided] = pred[~decided]
        return list(torch.split(subcodes, lengths.tolist(), dim=0))

# ==========================================================================
# 原 model_acc_fix3_altpc.py
# ==========================================================================

_BF16_RESIDUAL = os.environ.get("BF16_RESIDUAL", "1") == "1"
_STATIC_PAD_BUCKETS = tuple(sorted(int(x) for x in os.environ.get("STATIC_PAD_BUCKETS", "").split(",") if x.strip()))
_STATIC_PAD_SLOTS = int(os.environ.get("STATIC_PAD_SLOTS", "16"))
assert all(b % 9 == 0 for b in _STATIC_PAD_BUCKETS), "STATIC_PAD_BUCKETS 必须都是 9 的倍数"
class GeoRoPE:
    """几何 RoPE(rope_mode="geo"/"geo_vcat")的相位计算器。纯计算、零参数(不进 state_dict)。
    token 的"位置" = 3D 锚点 + 面内类型, 序列顺序完全不参与注意力(模型对面片块置换严格等变)。

    mode="geo": 锚点用聚合点(旧线 ckpt 在用, 勿删):
        9/face: F_i/M_ii 锚在父顶点 v_i, M_ij 锚在边中点 (v_i+v_j)/2; 类型 0/1/2=F/intra/inter
        3/face: 三个组 token 同锚面质心, 类型 = 组 id 0/1/2(与 9/face 同一条语义轴)
        1/face: 面质心, 类型 = 3
      head_dim/2 个频率槽切 [n,n,n,type_slots](n=(d/2-type_slots)//3, 余数补恒等相位),
      依序旋转 x/y/z/type(LATO RotaryPositionPhasesEmbedder 的 rotate_half 版)。
    mode="geo_vcat"(2026-08-05, 用户提议): 锚点显式编码顶点, **不用质心/中点等聚合点**。
        3/face 与 1/face 粗尺度: 「三顶点各占 1/3 通道拼接」(共质心三角形可分):
          xyz 30 槽切 3 组 x 每顶点 9 槽(x/y/z 各 3, 顶点序=面内 zyx 字典序), type 2 槽, 余 3 pad。
        9/face: 「双端点布局」(治中点混叠: 平行四边形两对角线同中点会被 geo 判成同一条边):
          每 token 两锚点各占一半 xyz 槽(x/y/z 各 5), F_i/M_ii=(v_i,v_i) 复制, M_ij=(v_i,v_j)
          (zyx 排序 => 端点对全局规范, 共享边跨面同相位); F 与其所属边的 inter 前半通道对齐。
      频率梯子统一 theta_vcat=300(3 槽/轴最低频 2.85 rad、5 槽/轴 1.33 rad, 均 <π 无歧义,
      邻格 1 rad 可分辨)。与 geo ckpt 参数完全同构, 但两个尺度的注意力几何都变了 ->
      热启动 loss 先跳后收(同 flat->geo 的迁移形态)。

    坐标约定(两种模式共同):
      锚点用**原始体素坐标**(乘 geo_scale, 默认 1.0), 刻意不按 depth 归一:
        注意力只见同样本内的相对相位(每样本单一 depth), 而连接决策的几何是体素邻接 ——
        体素单位让"相邻一格"在所有 depth 下相位签名一致(θ_xyz=100: 邻格 1 rad 可分辨,
        depth7 全程 128 格最低频 2.0 rad < π 无歧义)。
      coord 不减 offset: 常数平移让 q/k 同转一个角, 相对相位(注意力唯一可见量)不变。
      纯位置编码零新增参数 -> 任意旧 ckpt strict 热启动(注意几何变了, loss 先跳后收)。
    """

    def __init__(self, mode="flat", geo_scale=1.0, theta_xyz=100.0,
                 theta_type=10.0, type_slots=2, theta_vcat=300.0):
        self.mode = mode
        self.geo_scale = float(geo_scale)
        self.theta_xyz = float(theta_xyz)
        self.theta_type = float(theta_type)
        self.type_slots = int(type_slots)
        self.theta_vcat = float(theta_vcat)

    @property
    def is_geo(self):
        return str(self.mode).startswith("geo")

    @staticmethod
    def _ladder(theta, n, device):
        """几何频率梯子: freq_i = 1/theta^(i/n), i=0..n-1。"""
        return 1.0 / (theta ** (torch.arange(n, device=device, dtype=torch.float32) / max(n, 1)))

    @staticmethod
    def _cos_sin(rot, dtype):
        """[N, d_half] 相位 -> ([N,1,head_dim] cos, sin): 频率复制两半配 rotate_half,
        布局与 RotaryEmbedding.forward_positions 一致, apply_rotary_pos_emb_varlen 原样可用。"""
        emb = torch.cat((rot, rot), dim=-1)
        return emb.cos().to(dtype).unsqueeze(1), emb.sin().to(dtype).unsqueeze(1)

    def scales(self, coord_packed, d_half, dtype):
        """由 coord 流现算三个 hourglass 尺度的 (cos, sin), 返回 {9:(cos,sin), 3:..., 1:...}。
        每面 9 个 coord 标量 = 3 父顶点 x 3 分量(面内坐标字典序); 样本长度都是 9 的倍数且
        face 块不跨样本, 所以可以全局 view。"""
        device = coord_packed.device
        V = coord_packed.view(-1, 3, 3).to(torch.float32) * self.geo_scale   # [nF, 3顶点, xyz]
        nF = V.shape[0]
        freqs_type = self._ladder(self.theta_type, self.type_slots, device)
        # 面内类型分量(与 token 排布一致): 9/face = F F F | intra x3 | inter x3;
        # 3/face = 组 id 0/1/2(与 9/face 同一条语义轴); 1/face = 3
        type9 = torch.tensor([0., 0., 0., 1., 1., 1., 2., 2., 2.], device=device).repeat(nF)
        type3 = torch.tensor([0., 1., 2.], device=device).repeat(nF)
        type1 = torch.full((nF,), 3.0, device=device)
        if getattr(self, "mode", "geo") == "geo_vcat":
            return self._scales_vcat(V, type9, type3, type1, freqs_type, d_half, dtype, device)
        return self._scales_geo(V, type9, type3, type1, freqs_type, d_half, dtype, device)

    def _scales_geo(self, V, type9, type3, type1, freqs_type, d_half, dtype, device):
        """geo 布局: [x*n_ax | y*n_ax | z*n_ax | type*type_slots | pad]。"""
        n_type = self.type_slots
        n_ax = (d_half - n_type) // 3
        n_pad = d_half - n_type - 3 * n_ax
        freqs_ax = self._ladder(self.theta_xyz, n_ax, device)

        nF = V.shape[0]
        mid01 = (V[:, 0] + V[:, 1]) * 0.5
        mid02 = (V[:, 0] + V[:, 2]) * 0.5
        mid12 = (V[:, 1] + V[:, 2]) * 0.5
        centroid = V.mean(dim=1)                                        # [nF,3]
        # 槽位序 = [F0 F1 F2 M00 M11 M22 M01 M02 M12](与 token 排布一致)
        anchors9 = torch.stack([V[:, 0], V[:, 1], V[:, 2],
                                V[:, 0], V[:, 1], V[:, 2],
                                mid01, mid02, mid12], dim=1)            # [nF,9,3]
        anchors3 = centroid[:, None, :].expand(nF, 3, 3)                # down_9_to_3 按连续三元折叠成组
        anchors1 = centroid[:, None, :]

        def build(anchors, type_comp):
            a = anchors.reshape(-1, 3)
            parts = [a[:, 0:1] * freqs_ax, a[:, 1:2] * freqs_ax, a[:, 2:3] * freqs_ax,
                     type_comp[:, None] * freqs_type]
            if n_pad:
                parts.append(torch.zeros(a.shape[0], n_pad, device=device, dtype=torch.float32))
            return self._cos_sin(torch.cat(parts, dim=-1), dtype)

        return {9: build(anchors9, type9), 3: build(anchors3, type3), 1: build(anchors1, type1)}

    def _scales_vcat(self, V, type9, type3, type1, freqs_type, d_half, dtype, device):
        """geo_vcat: 锚点全部为显式顶点(不算质心/中点)。"""
        n_type = self.type_slots

        # ---- 粗尺度(3/face 与 1/face) = 三顶点拼接: 每顶点每轴 n_per_vert 槽,
        # 布局 [v0: x,y,z | v1: x,y,z | v2: x,y,z | type | pad]
        n_per_vert = (d_half - n_type) // 9
        n_pad_v = d_half - n_type - 9 * n_per_vert
        freqs_vcat = self._ladder(self.theta_vcat, n_per_vert, device)

        def build_vcat(type_comp, reps):
            parts_xyz = [V[:, vtx, ax:ax + 1] * freqs_vcat for vtx in range(3) for ax in range(3)]
            rot_face = torch.cat(parts_xyz, dim=-1)                     # [nF, 9*n_per_vert]
            rot = rot_face.repeat_interleave(reps, dim=0) if reps > 1 else rot_face
            parts = [rot, type_comp[:, None] * freqs_type]
            if n_pad_v:
                parts.append(torch.zeros(rot.shape[0], n_pad_v, device=device, dtype=torch.float32))
            return self._cos_sin(torch.cat(parts, dim=-1), dtype)

        # ---- 9/face = 双端点布局: 每 token 两个锚点各占一半 xyz 槽
        # [P1: x/y/z 各 n_ax9 | P2: 同 | type n_type], 30=2x15 无 pad。
        # F_i/M_ii -> (v_i, v_i) 复制; M_ij -> (v_i, v_j), i<j(面内 zyx 排序 => 端点对全局规范,
        # 共享边跨面同相位)。免费性质: F 锚 v_a 与 inter 锚 (v_a,v_b) 前半通道完全对齐 ——
        # "顶点是边的端点"在相位空间直接可见。
        n_ax9 = (d_half - n_type) // 6
        n_pad9 = d_half - n_type - 6 * n_ax9
        freqs_9 = self._ladder(self.theta_vcat, n_ax9, device)
        P1 = torch.stack([V[:, 0], V[:, 1], V[:, 2], V[:, 0], V[:, 1], V[:, 2],
                          V[:, 0], V[:, 0], V[:, 1]], dim=1).reshape(-1, 3)   # [nF*9,3]
        P2 = torch.stack([V[:, 0], V[:, 1], V[:, 2], V[:, 0], V[:, 1], V[:, 2],
                          V[:, 1], V[:, 2], V[:, 2]], dim=1).reshape(-1, 3)
        parts9 = [P1[:, 0:1] * freqs_9, P1[:, 1:2] * freqs_9, P1[:, 2:3] * freqs_9,
                  P2[:, 0:1] * freqs_9, P2[:, 1:2] * freqs_9, P2[:, 2:3] * freqs_9,
                  type9[:, None] * freqs_type]
        if n_pad9:
            parts9.append(torch.zeros(P1.shape[0], n_pad9, device=device, dtype=torch.float32))
        scales9 = self._cos_sin(torch.cat(parts9, dim=-1), dtype)
        return {9: scales9, 3: build_vcat(type3, 3), 1: build_vcat(type1, 1)}
_COORD_PACK_BASE = 4096   # 坐标打包基数(> 最大体素坐标), xyz -> 单一整数 code
def _pack_coord(p3):
    """[.,3] 整数 xyz -> 单一整数 code(供 torch.unique 分组用)。"""
    return p3[:, 0] * _COORD_PACK_BASE * _COORD_PACK_BASE + p3[:, 1] * _COORD_PACK_BASE + p3[:, 2]
class TokenGroup:
    """一个跨面 token 组划分 + 组内表决。
    字段: pos=该类型 token 的全局位置; group=每 token 的组 id(0..n_group-1);
          face=每 token 所在面; code=父顶点打包坐标(仅 by_parent, 顶点原子揭示要用)。"""

    def __init__(self, pos, group, n_group, face, code=None):
        self.pos, self.group, self.n_group, self.face, self.code = pos, group, n_group, face, code

    @classmethod
    def by_parent(cls, types, face_coords, token_type, parent_offset):
        """F/intra 按父顶点分组(同一顶点在不同面的 token 同组)。
        parent_offset: F=0(面内位置 0/1/2 -> 父 v0/v1/v2), intra=3(位置 3/4/5)。"""
        pos = (types == token_type).nonzero(as_tuple=True)[0]
        face = pos // REFINE_POS_PER_FACE
        code = _pack_coord(face_coords[face, (pos % REFINE_POS_PER_FACE) - parent_offset])
        uniq, group = torch.unique(code, return_inverse=True)
        return cls(pos, group, int(uniq.numel()), face, code)

    @classmethod
    def by_edge(cls, types, face_coords):
        """inter 按「无序父对(=共享边)」分组: 面内位置 6/7/8 -> 父对 (v0,v1)/(v0,v2)/(v1,v2)。
        父对取 (小,大) => 同一条边跨面同组, 朝向天然规范(无 M vs M^T 问题)。"""
        device = face_coords.device
        pos = (types == TYPE_INTER).nonzero(as_tuple=True)[0]
        face = pos // REFINE_POS_PER_FACE
        local = (pos % REFINE_POS_PER_FACE) - 6                     # 0/1/2 -> M01/M02/M12
        pa = torch.tensor([0, 0, 1], device=device)
        pb = torch.tensor([1, 2, 2], device=device)
        code_a = _pack_coord(face_coords[face, pa[local]])
        code_b = _pack_coord(face_coords[face, pb[local]])
        pair = torch.stack([torch.minimum(code_a, code_b), torch.maximum(code_a, code_b)], dim=1)
        uniq, group = torch.unique(pair, dim=0, return_inverse=True)
        return cls(pos, group, int(uniq.shape[0]), face)

    @staticmethod
    def _mode_of(group, vid, n_group):
        """每组内 vid 的众数(sort+segment, 免 python 循环), 广播回每个成员。"""
        n_vid = int(vid.max().item()) + 1 if vid.numel() else 1
        key, cnt = torch.unique(group.long() * n_vid + vid.long(), return_counts=True)
        g, v = key // n_vid, key % n_vid
        order = torch.argsort(g * (int(cnt.max().item()) + 1) - cnt)   # 组内按计数降序
        g, v = g[order], v[order]
        first = torch.ones_like(g, dtype=torch.bool)
        if g.numel() > 1:
            first[1:] = g[1:] != g[:-1]
        mode = torch.zeros(n_group, dtype=torch.long, device=vid.device)
        mode[g[first]] = v[first]
        return mode[group]

    def vote_bytes(self, vals):
        """F 通道: 每组对占用字节(0..255)计数表决, 返回广播回每 token 的众数字节。"""
        cnt = torch.bincount(self.group * 256 + vals.clamp(0, 255),
                             minlength=self.n_group * 256).view(self.n_group, 256)
        return cnt.argmax(1)[self.group]

    def vote_vec(self, vecs):
        """intra/inter 通道: 8-subcode 向量整体表决。返回 (该 token 是否≠众数, 广播回的众数向量)。"""
        uniq, vid = torch.unique(vecs, dim=0, return_inverse=True)
        mode_vid = self._mode_of(self.group, vid, self.n_group)
        return vid != mode_vid, uniq[mode_vid]
class _DualCloudCond:
    """一次前向内的点云条件: A=全GT云, B=目标分辨率塌缩云, 按全局 block 序号交替提供 cross-attn KV
    (奇数层->A, 偶数层->B; B 未提供时两套相同 = 原单点云行为; 云为 None 走 null_cond = CFG 无条件分支)。
    每套 cond = (packed KV [sum(len),d], 每样本 KV 长度, KV max_len)。
    fake_lens 桶 pad 时在 conditioner **之后**补零行(不浪费 Michelangelo 前向)。
    cache: 推理跨步复用编码(采样各步点云不变, miche 逐步重算纯浪费; 数值严格等价), 训练传 None。"""

    def __init__(self, model, coord_list, fake_lens, device, pc, pc_target, cache=None):
        self._model = model
        self._n_seq = len(coord_list)          # 含 fake 序列(null_cond 也要按批展开)
        self._fake_lens = fake_lens
        self._device = device
        key = ("null" if pc is None else "cond", pc_target is None)
        if cache is not None and key in cache:
            self.cond_a, self.cond_b = cache[key]
            return
        self.cond_a = self._encode(pc)                               # 奇数层(全GT)
        if model.use_pc and (pc_target is not None):
            self.cond_b = self._encode(pc_target)                    # 偶数层(目标分辨率塌缩)
        else:
            self.cond_b = self.cond_a                                # 未提供 -> 两套相同(原行为)
        if cache is not None:
            cache[key] = (self.cond_a, self.cond_b)

    def _encode(self, cloud):
        model, device = self._model, self._device
        if not model.use_pc:
            return None, None, None
        if cloud is None:                                            # CFG 无条件分支
            null = model.null_cond
            kv = null.unsqueeze(0).expand(self._n_seq, -1, -1).reshape(-1, null.size(-1))
            kv_len = torch.full((self._n_seq,), model.null_cond_len, device=device, dtype=torch.long)
            return kv, kv_len, int(model.null_cond_len)
        emb = model.conditioner(cloud.to(device=device))
        if self._fake_lens:                                          # 假序列的 cross-attn KV = 零行
            emb = torch.cat([emb, emb.new_zeros(len(self._fake_lens), emb.size(1), emb.size(2))], dim=0)
        kv_len = torch.full((emb.size(0),), emb.size(1), device=device, dtype=torch.long)
        return emb.reshape(-1, emb.size(-1)), kv_len, int(emb.size(1))

    def for_layer(self, gi):
        """gi: 0-based 全局 block 序号; 层号=gi+1 -> 奇数层(gi 偶)用 A, 偶数层(gi 奇)用 B。"""
        return self.cond_a if gi % 2 == 0 else self.cond_b
class _ReMDM:
    """ReMDM(arXiv 2503.00307, 训练免费): 步首以 σ_t 重掩已定 token, 本步前向重新预测;
    揭示额度按 n_keep_masked 自动补足 -> 掩码数轨迹严格贴原 cosine 调度(边际保持的计数空间等价)。
    σ_max=min(1, m_s/|C|); cap: σ=min(η,σ_max); rescale: σ=η·σ_max; conf: cap 的量+定向重掩
    提交置信度最低者。env 门控(sample 调用时读取), REMDM 为空 = 与原行为逐位一致。"""

    def __init__(self):
        self.mode = os.environ.get("REMDM", "").strip().lower()
        self.eta = float(os.environ.get("REMDM_ETA", "0.05"))
        self.t_start = float(os.environ.get("REMDM_TSTART", "0.0"))
        self.t_end = float(os.environ.get("REMDM_TEND", "1.0"))

    def maybe_remask(self, step, T_infer, total, decided, subcodes, commit_conf, device):
        """就地重掩(decided/subcodes 原位改写)。"""
        if not (self.mode and step > 0 and self.t_start <= step / T_infer <= self.t_end):
            return
        n_committed = int(decided.sum().item())
        if n_committed == 0:
            return
        n_mask_target = int(cosine_schedule(torch.tensor(float(step + 1)), float(T_infer)).item() * total)
        sigma_max = min(1.0, n_mask_target / n_committed)
        sigma = self.eta * sigma_max if self.mode == "rescale" else min(self.eta, sigma_max)
        n_remask = int(sigma * n_committed)
        if n_remask <= 0:
            return
        committed_idx = decided.nonzero(as_tuple=True)[0]
        if self.mode == "conf":
            remask_idx = committed_idx[commit_conf[committed_idx].argsort()[:n_remask]]   # 最低提交置信优先
        else:
            remask_idx = committed_idx[torch.randperm(n_committed, device=device)[:n_remask]]
        decided[remask_idx] = False
        subcodes[remask_idx] = MASK_ID
class FactorizedDiffusionTransformerAltPC(FactorizedDiffusionTransformer):
    def __init__(self, *args, rope_mode="flat", rope_geo_scale=1.0,
                 rope_theta_xyz=100.0, rope_theta_type=10.0, rope_type_slots=2,
                 rope_theta_vcat=300.0, **kwargs):
        super().__init__(*args, **kwargs)
        # 几何 RoPE 封装在 GeoRoPE(纯计算无参数, 不影响 state_dict)。
        # rope_mode="flat"(默认): 原行为, 逐 token 序列位置 0..L-1, 与所有旧 ckpt 完全兼容;
        # "geo"/"geo_vcat" 的设计与 theta 约束见 GeoRoPE docstring。
        self.geo_rope = GeoRoPE(mode=rope_mode, geo_scale=rope_geo_scale,
                                theta_xyz=rope_theta_xyz, theta_type=rope_theta_type,
                                type_slots=rope_type_slots, theta_vcat=rope_theta_vcat)

    # rope_* 属性代理到 GeoRoPE(单一事实源): 外部读写 self.rope_mode 等保持兼容
    # (含子类/脚本事后改 rope_mode 的用法, 如 tssr 的 self.rope_mode = "hier")。
    rope_mode = property(lambda self: self.geo_rope.mode,
                         lambda self, v: setattr(self.geo_rope, "mode", v))
    rope_geo_scale = property(lambda self: self.geo_rope.geo_scale,
                              lambda self, v: setattr(self.geo_rope, "geo_scale", float(v)))
    rope_theta_xyz = property(lambda self: self.geo_rope.theta_xyz,
                              lambda self, v: setattr(self.geo_rope, "theta_xyz", float(v)))
    rope_theta_type = property(lambda self: self.geo_rope.theta_type,
                               lambda self, v: setattr(self.geo_rope, "theta_type", float(v)))
    rope_type_slots = property(lambda self: self.geo_rope.type_slots,
                               lambda self, v: setattr(self.geo_rope, "type_slots", int(v)))
    rope_theta_vcat = property(lambda self: self.geo_rope.theta_vcat,
                               lambda self, v: setattr(self.geo_rope, "theta_vcat", float(v)))

    def _geo_rope_scales(self, coord_packed, dtype):
        """geo/geo_vcat: 三个 hourglass 尺度的 (cos,sin)。方法名/签名保留(fpsq 子类在调用)。"""
        return self.geo_rope.scales(coord_packed, self.rotary_emb.inv_freq.numel(), dtype)

    @staticmethod
    def _static_pad_plan(total, b_real):
        """返回假序列长度列表(都是 9 的倍数且 >=9): padded 总数落到桶位、批槽位补满 _STATIC_PAD_SLOTS。
        超过最大桶位时返回 None(不 pad, 触发一次形状特化编译; 配 hard_max_positions 应当罕见)。"""
        n_fake = max(_STATIC_PAD_SLOTS - b_real, 1)
        min_pad = 9 * n_fake
        target = next((B for B in _STATIC_PAD_BUCKETS if B >= total + min_pad), None)
        if target is None:
            return None
        fake = [9] * n_fake
        fake[0] += target - total - min_pad          # 桶位/总数都是 9 的倍数 -> 余量也是
        return fake

    def forward_varlen(self, coord_list, subcode_list, t, depth, pc=None, pc_target=None,
                       cond_cache=None):
        device = t.device
        coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list]
        subcode_list = [s.to(device=device, dtype=torch.long) for s in subcode_list]
        # 长度在 python 端本来就知道 -> 一并算出各尺度 max_seqlen 传给 attention,
        # 省掉每个 flash 调用的 int(lengths.max().item()) GPU->CPU 同步(见 model_acc_fix3_tssr)。
        lengths_py = [int(c.numel()) for c in coord_list]
        if sum(lengths_py) == 0:
            raise ValueError("forward_varlen received an empty batch.")
        if any(n % REFINE_POS_PER_FACE != 0 for n in lengths_py):
            raise ValueError("每个序列长度必须是 9 的倍数(每 face 9 个 refine 位置)。")

        # ---- 静态桶化 padding(见文件顶部说明; STATIC_PAD_BUCKETS 为空 = 原行为) ----
        real_total, real_b = sum(lengths_py), len(lengths_py)
        fake_lens = self._static_pad_plan(real_total, real_b) if _STATIC_PAD_BUCKETS else None
        if fake_lens:
            sub_tail = tuple(subcode_list[0].shape[1:])
            for L in fake_lens:
                coord_list.append(torch.zeros(L, dtype=torch.long, device=device))
                subcode_list.append(torch.zeros((L,) + sub_tail, dtype=torch.long, device=device))
            lengths_py = lengths_py + fake_lens
        lengths = torch.tensor(lengths_py, device=device, dtype=torch.long)
        if fake_lens:
            # pad 模式: max_len 传桶位常数(= padded 总数, 合法上界), 保证编译形状/guard 全静态
            max_len_9 = sum(lengths_py)
        else:
            max_len_9 = max(lengths_py)
        max_len_3 = max_len_9 // 3          # hourglass 整除, 断言在上面
        max_len_1 = max_len_9 // 9

        coord_packed = torch.cat(coord_list, dim=0)
        subcode_packed = torch.cat(subcode_list, dim=0)
        types = self.types_from_lengths(lengths)

        raw_emb = self.raw_embedding(coord_packed)
        refine_emb = self.embed_refine(subcode_packed, types)
        x = raw_emb + refine_emb                         # 不再加法注入 time；time/depth 走 adaLN
        # 统一残差流 dtype(理由见文件顶部 _BF16_RESIDUAL 注释); c_sample 刻意保持 fp32(adaLN 惯例)
        if _BF16_RESIDUAL and torch.is_autocast_enabled():
            x = x.to(torch.get_autocast_dtype("cuda"))

        # adaLN 条件 c = time_emb + depth_emb(逐样本 [b,d]); 各 hourglass 阶段按其长度展开到逐 token
        depth = depth.to(device=device, dtype=torch.long)
        c_sample = self.time_mlp(t) + self.depth_embedding(depth)   # [b, d_model]
        if fake_lens:                                                # 假序列的 adaLN 条件 = 零行
            c_sample = torch.cat([c_sample, c_sample.new_zeros(len(fake_lens), c_sample.size(-1))], dim=0)

        def c_per_token(cur_lengths):
            return c_sample.repeat_interleave(cur_lengths, dim=0)

        # 点云条件: 双云编码 + 逐层交替选择(见 _DualCloudCond)
        cond = _DualCloudCond(self, coord_list, fake_lens, device, pc, pc_target, cond_cache)

        # geo 模式: 三个尺度的 cos/sin 一次算好(锚点只依赖 coord/depth, 与序列位置无关)
        geo_scales = (self._geo_rope_scales(coord_packed, x.dtype)
                      if str(getattr(self, "rope_mode", "flat")).startswith("geo") else None)

        def rope_fn_for(cur_lengths, tokens_per_face=9):
            if geo_scales is not None:
                cos, sin = geo_scales[tokens_per_face]
            else:
                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)

        # ---- hourglass 五段: enc9 -> enc3 -> bottleneck -> dec3 -> dec9 ----
        gi = 0   # 全局 block 序号(横跨五段, 决定该层注入云A还是云B)

        def run_stage(x, blocks, c_stage, stage_lengths, rope_fn, stage_max_len):
            nonlocal gi
            for block in blocks:
                cond_kv, cond_kv_len, cond_kv_max = cond.for_layer(gi)
                gi += 1
                x = block.forward_varlen(x, c_stage, stage_lengths, cond_kv, cond_kv_len,
                                         rope_apply_fn=rope_fn, max_len=stage_max_len,
                                         cond_max_len=cond_kv_max)
            return x

        rope_fn_9, c9 = rope_fn_for(lengths, 9), c_per_token(lengths)
        x = run_stage(x, self.enc_blocks_9, c9, lengths, rope_fn_9, max_len_9)
        skip_9 = x

        x, lengths = self.down_9_to_3.forward_varlen(x, lengths)
        rope_fn_3, c3 = rope_fn_for(lengths, 3), c_per_token(lengths)
        x = run_stage(x, self.enc_blocks_3, c3, lengths, rope_fn_3, max_len_3)
        skip_3 = x

        x, lengths = self.down_3_to_1.forward_varlen(x, lengths)
        rope_fn_1, c1 = rope_fn_for(lengths, 1), c_per_token(lengths)
        x = run_stage(x, self.bottleneck_blocks, c1, lengths, rope_fn_1, max_len_1)

        x, lengths = self.up_1_to_3.forward_varlen(x, lengths)
        x = run_stage(x + skip_3, self.dec_blocks_3, c3, lengths, rope_fn_3, max_len_3)

        x, lengths = self.up_3_to_9.forward_varlen(x, lengths)
        x = run_stage(x + skip_9, self.dec_blocks_9, c9, lengths, rope_fn_9, max_len_9)

        if fake_lens:                                                # 切掉假位置(不进头/loss)
            x, types = x[:real_total], types[:real_total]
        return x, types  # hidden [total,d], types [total]

    def forward(self, coord_list, noisy_subcode_list, t, depth,
                target_subcode_list=None, loss_mask=None, pc=None, pc_target=None, return_extra_acc=False):
        """训练入口, 与基类同, 只多传 pc_target(逐层交替的第二套点云)。"""
        h, types = self.forward_varlen(coord_list, noisy_subcode_list, t, depth, pc, pc_target)
        if target_subcode_list is None:
            return h, types
        target = torch.cat([s.to(device=h.device, dtype=torch.long) for s in target_subcode_list], dim=0)
        loss_mask = loss_mask.to(device=h.device)
        loss = self.factorized_loss(h, target, types, loss_mask)   # focal(继承)
        with torch.no_grad():
            acc, acc_ne = self.factorized_accuracy(h, target, types, loss_mask)
        if return_extra_acc:
            return loss, acc, acc_ne
        return loss, acc

    def _predict_heads(self, h, h_u, types, subcodes, cfg_t, temperature):
        """三个头对全部位置出预测 + 置信度: F=256 类单分类(整块 9 位广播同值),
        intra/inter=8 行自回归 rollout。h_u 非 None 时做 logits 级 CFG。返回 (pred, conf)。"""
        pred = subcodes.clone()
        conf = torch.zeros(types.numel(), device=h.device)
        is_f = types == TYPE_F
        if is_f.any():
            logits = self.f_head(h[is_f])
            if h_u is not None:
                logits_u = self.f_head(h_u[is_f])
                logits = logits_u + cfg_t * (logits - logits_u)
            if temperature and temperature > 0:
                prob = (logits / temperature).softmax(-1)
                choice = torch.multinomial(prob.float(), 1).squeeze(-1)
                choice_prob = prob.gather(-1, choice.unsqueeze(-1)).squeeze(-1)
            else:
                prob = logits.softmax(-1)
                choice_prob, choice = prob.max(-1)
            pred[is_f, :] = choice.unsqueeze(1).expand(-1, NUM_OCTANTS)
            conf[is_f] = choice_prob.to(conf.dtype)
        for token_type, head in ((TYPE_INTRA, self.intra_head), (TYPE_INTER, self.inter_head)):
            sel = types == token_type
            if not sel.any():
                continue
            row_pred, row_conf = head.rollout(h[sel], h_u[sel] if h_u is not None else None, cfg_t, temperature)
            pred[sel] = row_pred
            conf[sel] = row_conf.mean(dim=1).to(conf.dtype)
        return pred, conf

    @torch.no_grad()
    def sample(self, coord_list, depth, pc=None, pc_target=None, T_infer=10, cfg_scale=1.0,
               temperature=0.0, cfg_schedule="constant", cfg_start=1.0,
               align_fm=False, align_inter=False, align_per_step=False, align_debug=True,
               diag_vertex=False):
        """前向多传 pc_target(逐层交替)。CFG 无条件分支两套点云都走 null_cond。

        跨面一致性对齐(GT 侧这些量对共享结构 100% 一致, 见 check_tokenizer_consistency):
          align_fm=True    : 对齐 F(占用字节, 按"共享父"分组) + intra(M_inner 8-subcode 矩阵, 按"共享父")。
          align_inter=True : 对齐 inter(M_outer 8-subcode 矩阵, 按"共享父对"=共享边分组)。
                             如面 (v0,v1,v2) 与 (v0,v1,v3) 共享边 (v0,v1) -> 其 (v0,v1) 的 inter 应一致。
          align_per_step   : False(默认, 推荐)= 只在收尾把各启用通道投影到表决值(不扰动去噪, 不造洞);
                             True = 每步额外把"不一致的面"整块 9 token 重掩重预测(实测打乱揭示节奏 ->
                             反而多洞, 仅供对照)。
        全 False -> 与原行为逐位一致。

        env 门控(为空/未设 = 与原行为逐位一致):
          REMDM=cap|rescale|conf (+REMDM_ETA/TSTART/TEND): 步首重掩已定 token(见 _ReMDM)。
          WINDOW_REVEAL=1(仅 batch=1): BFS 前缀窗内 token 才可提交(块因果 变体A), 注意力/预测不变。
        """
        self.eval()
        device = next(self.parameters()).device
        coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list]
        depth = depth.to(device=device, dtype=torch.long)
        lengths = torch.tensor([c.numel() for c in coord_list], device=device, dtype=torch.long)
        types = self.types_from_lengths(lengths)
        want_cfg = (float(cfg_scale) != 1.0 or float(cfg_start) != 1.0) and self.use_pc and (pc is not None)

        total = int(lengths.sum().item())
        nface = total // REFINE_POS_PER_FACE

        align_any = align_fm or align_inter
        if align_any or diag_vertex:
            face_coords = torch.cat(coord_list).view(nface, 3, 3)   # [nface, 3父, xyz]

        # ---- 顶点组诊断: 同一父(顶点)在不同面的 F/intra token, 记录 commit step, 末尾统计一致性 ----
        if diag_vertex:
            diag_f = TokenGroup.by_parent(types, face_coords, TYPE_F, 0)
            diag_i = TokenGroup.by_parent(types, face_coords, TYPE_INTRA, 3)
            commit_step = torch.full((total,), int(T_infer), dtype=torch.long, device=device)

        # ---- 跨面对齐分组: F/intra 按"共享父", inter 按"共享父对"(=共享边) ----
        if align_fm:
            grp_f = TokenGroup.by_parent(types, face_coords, TYPE_F, 0)
            grp_i = TokenGroup.by_parent(types, face_coords, TYPE_INTRA, 3)
        if align_inter:
            grp_e = TokenGroup.by_edge(types, face_coords)

        subcodes = torch.full((total, NUM_OCTANTS), MASK_ID, dtype=torch.long, device=device)
        decided = torch.zeros(total, dtype=torch.bool, device=device)
        pred = subcodes.clone()

        remdm = _ReMDM()                                    # env 门控, 未设=无操作
        window_reveal = os.environ.get("WINDOW_REVEAL", "0") == "1" and len(lengths) == 1
        commit_conf = torch.zeros(total, device=device)     # 每 token 提交时的置信度(ReMDM-conf 用)
        cond_cache = {}                                     # 点云条件跨步缓存(严格等价, 见 _DualCloudCond)

        for step in range(T_infer):
            remdm.maybe_remask(step, T_infer, total, decided, subcodes, commit_conf, device)

            if cfg_schedule == "linear":
                cfg_t = float(cfg_start) + (float(cfg_scale) - float(cfg_start)) * step / max(T_infer - 1, 1)
            else:
                cfg_t = float(cfg_scale)
            use_cfg = want_cfg and (cfg_t != 1.0)

            # ---- 前向(有条件 + 可选无条件 CFG 分支) -> 三头全量预测 ----
            t = torch.full((len(coord_list),), step / T_infer, device=device, dtype=torch.float)
            sub_list = list(torch.split(subcodes, lengths.tolist(), dim=0))
            h, _ = self.forward_varlen(coord_list, sub_list, t, depth, pc, pc_target, cond_cache=cond_cache)
            h_u = self.forward_varlen(coord_list, sub_list, t, depth, None, None, cond_cache=cond_cache)[0] if use_cfg else None
            h = h.float()                                    # 头部权重 fp32; bf16 残差模式下需 cast(fp32 下 no-op)
            if h_u is not None:
                h_u = h_u.float()
            pred, conf = self._predict_heads(h, h_u, types, subcodes, cfg_t, temperature)

            # ---- 每步重掩(align_per_step): 跨面不一致的面整块 9 token 重掩(本步不揭示+撤销已定)。
            # 检测 F/intra(按共享父)与 inter(按共享父对)的不一致; 任一通道不一致的面都重掩。实测会造洞, 默认关。
            if align_per_step and align_any:
                force_face = torch.zeros(nface, dtype=torch.bool, device=device)
                if align_fm:
                    if grp_f.pos.numel():
                        bad = pred[grp_f.pos, 0] != grp_f.vote_bytes(pred[grp_f.pos, 0])
                        force_face[grp_f.face[bad]] = True
                    if grp_i.pos.numel():
                        dis_intra, _ = grp_i.vote_vec(pred[grp_i.pos])
                        force_face[grp_i.face[dis_intra]] = True
                if align_inter and grp_e.pos.numel():
                    dis_inter, _ = grp_e.vote_vec(pred[grp_e.pos])
                    force_face[grp_e.face[dis_inter]] = True
                remask = force_face.repeat_interleave(REFINE_POS_PER_FACE)
                conf[remask] = -float('inf')
                decided = decided & (~remask)

            # ---- 按 cosine 调度揭示 top-k 置信 token ----
            keep_ratio = cosine_schedule(torch.tensor(float(step + 1)), float(T_infer)).item()
            n_keep_masked = int(keep_ratio * total)
            n_to_decide = max(0, (total - n_keep_masked) - int(decided.sum().item()))
            if n_to_decide > 0:
                undecided_idx = (~decided).nonzero(as_tuple=True)[0]
                if window_reveal:
                    # 块序前沿窗口(推理期块因果 变体A): 仅允许 BFS 前缀窗内的 token 本步被提交。
                    # 注意力/预测不变(coord 脚手架全局可见), 只限制提交资格; 窗内不够本步预算时自适应扩窗。
                    n_blocks = total // 9
                    win_blocks = max(1, int(np.ceil(n_blocks * (step + 1) / float(T_infer))))
                    while True:
                        in_window = undecided_idx[undecided_idx < win_blocks * 9]
                        if len(in_window) >= n_to_decide or win_blocks >= n_blocks:
                            break
                        win_blocks += max(1, (n_to_decide - len(in_window) + 8) // 9)
                    undecided_idx = undecided_idx[undecided_idx < win_blocks * 9]
                order = conf[undecided_idx].argsort(descending=True)
                chosen = undecided_idx[order[:n_to_decide]]
                subcodes[chosen] = pred[chosen]
                decided[chosen] = True
                commit_conf[chosen] = conf[chosen]
                if diag_vertex:
                    commit_step[chosen] = step
            subcodes[~decided] = MASK_ID

        subcodes[~decided] = pred[~decided]
        if diag_vertex:
            self._diag_vertex_report(subcodes, commit_step, diag_f.pos, diag_f.group,
                                     diag_i.pos, diag_i.group, int(T_infer))

        # ---- 末步: 把各启用通道强制对齐到表决值(不扰动去噪, 推荐的对齐方式)。
        # align_debug=True: 打印每通道"与表决值不一致(即被改动)的 token 数 / 涉及面数"便于 debug。
        def report(tag, disagree, face_of):
            n_tok = int(disagree.sum().item())
            n_all = disagree.numel()
            n_face = int(torch.unique(face_of[disagree]).numel()) if n_tok else 0
            print(f"[align_debug] {tag:5s}: 不一致 {n_tok}/{n_all} token ({100 * n_tok / max(n_all, 1):.1f}%), 涉及 {n_face} 面")

        if align_fm:
            if grp_f.pos.numel():
                voted_f = grp_f.vote_bytes(subcodes[grp_f.pos, 0])
                if align_debug:
                    report("F", subcodes[grp_f.pos, 0] != voted_f, grp_f.face)
                subcodes[grp_f.pos, :] = voted_f.unsqueeze(1).expand(-1, NUM_OCTANTS)
            if grp_i.pos.numel():
                dis_intra, voted_intra = grp_i.vote_vec(subcodes[grp_i.pos])
                if align_debug:
                    report("intra", dis_intra, grp_i.face)
                subcodes[grp_i.pos, :] = voted_intra
        if align_inter and grp_e.pos.numel():
            dis_inter, voted_inter = grp_e.vote_vec(subcodes[grp_e.pos])
            if align_debug:
                report("inter", dis_inter, grp_e.face)
            subcodes[grp_e.pos, :] = voted_inter
        return list(torch.split(subcodes, lengths.tolist(), dim=0))


    @torch.no_grad()
    def _diag_vertex_report(self, subcodes, commit_step, f_pos, f_grp, i_pos, i_grp, T):
        """诊断: 同一父(顶点)在不同面的 F/intra token 是否一致, 以及不一致的时序来源。
        关键问题——不一致是(a)同步揭示时就分裂, 还是(b)已有同组 token commit(可见上下文)后仍预测错?"""
        def analyze(pos, grp, tag, is_vec):
            if pos.numel() == 0:
                print(f"[diag {tag}] 无 token")
                return
            if is_vec:   # intra: 8-subcode 向量 -> 唯一 id
                _, vid = torch.unique(subcodes[pos], dim=0, return_inverse=True)
                vals = vid.cpu().numpy()
            else:        # F: 占用字节
                vals = subcodes[pos, 0].clamp(0, 255).cpu().numpy()
            groups = grp.cpu().numpy()
            steps = commit_step[pos].cpu().numpy()
            order = np.argsort(groups, kind="stable")
            groups, vals, steps = groups[order], vals[order], steps[order]
            segs = np.split(np.arange(len(groups)), np.flatnonzero(np.diff(groups)) + 1)
            n_multi = n_incons = 0
            post_total = post_dis = 0                     # 后揭示事件(同组已有更早 commit): 与已定值 一致/不一致
            coreveal = postcommit = both = 0              # 不一致组来源分类
            spreads = []
            for seg in segs:
                if len(seg) < 2:
                    continue
                n_multi += 1
                v, s = vals[seg], steps[seg]
                inconsistent = len(set(v.tolist())) > 1
                first_step = int(s.min())
                early_vals = v[s == first_step]
                ref = np.bincount(early_vals).argmax()    # 最早 step 的众数值作参照
                coreveal_split = len(set(early_vals.tolist())) > 1
                later = v[s > first_step]
                post_total += len(later)
                post_dis += int((later != ref).sum())
                has_post = bool((later != ref).any())
                if inconsistent:
                    n_incons += 1
                    spreads.append(int(s.max() - s.min()))
                    if coreveal_split and has_post:
                        both += 1
                    elif coreveal_split:
                        coreveal += 1
                    else:
                        postcommit += 1
            pct_dis = 100 * post_dis / max(post_total, 1)
            med = int(np.median(spreads)) if spreads else 0
            print(f"[diag {tag}] 多面顶点组 {n_multi} | 最终不一致 {n_incons} ({100*n_incons/max(n_multi,1):.1f}%)")
            print(f"          后揭示token(同组已有commit) {post_total}: 不一致 {post_dis} ({pct_dis:.1f}%) <- 看到已定值仍预测错的比例")
            print(f"          不一致组来源: 同步揭示分裂 {coreveal} / 后揭示才分裂 {postcommit} / 两者 {both} | commit step 跨度中位 {med}/{T}")

        analyze(f_pos, f_grp, "F", False)
        analyze(i_pos, i_grp, "intra", True)

# ==========================================================================
# 原 model_acc_fix3_tssr.py
# ==========================================================================

_BF16_RESIDUAL = os.environ.get("BF16_RESIDUAL", "1") == "1"
class FactorizedTSSRTransformer(FactorizedDiffusionTransformerAltPC):
    def __init__(self, *args, mode_embed_scale=1000.0, rope_mode="flat",
                 rope_theta_face=10000.0, rope_theta_group=10000.0, rope_theta_member=10000.0, **kwargs):
        super().__init__(*args, **kwargs)
        # ---- 层级 RoPE 开关(移植自 Tencent-TSSR MultiLevelRoPE, 适配本仓库 rotate_half 约定) ----
        # rope_mode="flat"(默认): 原行为, 逐 token 平位置 0..L-1, 与所有旧 ckpt 完全兼容。
        # rope_mode="hier"    : head 维的 d/2 个频率槽切三段, 分别用三个结构分量旋转:
        #     face   段: p // tokens_per_face          (第几个面; 长程)
        #     group  段: 组内三元组编号                 (9/face 时 (p%9)//3: 0=F,1=intra,2=inter)
        #     member 段: 三元组内成员                   (9/face 时 p%3: 顶点/边槽位, F0↔M00 同相位)
        #   hourglass 三个尺度(9/3/1 token per face)各有对应分解(对齐 TSSR 的 token/vertex/face 三级)。
        #   纯位置编码无参数 -> 任意旧 ckpt 可直接热启动, 但注意几何变了, loss 会先跳再收。
        self.rope_mode = rope_mode
        self.rope_thetas = (float(rope_theta_face), float(rope_theta_group), float(rope_theta_member))
        d = self.d_model
        # 模式 flag 嵌入:标量 mode∈{0,1} -> 正弦 -> MLP -> [b,d](仿 TSSR DiscreteValueEmbeding)。
        # scale 大让 0 与 1 的正弦特征清晰可分;末层零初始化使初始 mode token≈0(warm-start 不破坏骨干)。
        self.mode_emb = nn.Sequential(
            SinusoidalPosEmb(d, scale=mode_embed_scale),
            nn.Linear(d, d), nn.SiLU(),
            nn.Linear(d, d),
        )
        nn.init.zeros_(self.mode_emb[-1].weight)
        nn.init.zeros_(self.mode_emb[-1].bias)
        # 逐 token 正确性分类器(detach 骨干)
        self.to_acc = nn.Sequential(
            nn.Linear(d, d // 2), nn.SiLU(),
            nn.Linear(d // 2, d // 4), nn.SiLU(),
            nn.Linear(d // 4, 1),
        )

    def _hier_cos_sin(self, position_ids, tokens_per_face, dtype):
        """层级 RoPE 的 cos/sin, 形状/布局与 RotaryEmbedding.forward_positions 一致([total, head_dim],
        频率复制两半以配 rotate_half), 因此 apply_rotary_pos_emb_varlen 原样可用。
        d/2 个频率槽按 TSSR 方式切 [n, n, d/2-2n](n=d//6), 依序分给 face/group/member 分量。"""
        device = position_ids.device
        d_half = self.rotary_emb.inv_freq.numel()                 # = head_dim // 2
        n = d_half // 3
        chunk_sizes = (n, n, d_half - 2 * n)
        p = position_ids
        if tokens_per_face == 9:
            comps = (p // 9, (p % 9) // 3, p % 3)
        elif tokens_per_face == 3:
            comps = (p // 3, p % 3, torch.zeros_like(p))
        elif tokens_per_face == 1:
            comps = (p, torch.zeros_like(p), torch.zeros_like(p))
        else:
            raise ValueError(f"tokens_per_face 只支持 9/3/1, 收到 {tokens_per_face}")
        rotations = []
        for size, theta, comp in zip(chunk_sizes, self.rope_thetas, comps):
            freqs = 1.0 / (theta ** (torch.arange(size, device=device, dtype=torch.float32) / max(size, 1)))
            rotations.append(comp.to(torch.float32)[:, None] * freqs[None, :])
        rot = torch.cat(rotations, dim=-1)                        # [total, d_half]
        emb = torch.cat((rot, rot), dim=-1)                       # [total, head_dim], 同 rotate_half 布局
        # [total, 1, head_dim]: 中间加 head 广播维, 与 forward_positions 的返回形状一致
        # (q/k 是 [total, nhead, d_head], 少了这一维会把 total 广播到 nhead 上直接报形状错)
        return emb.cos().to(dtype).unsqueeze(1), emb.sin().to(dtype).unsqueeze(1)

    # ------------------------------------------------------------------ #
    # forward_varlen:复制 AltPC 版,改动 = ①mode token 拼点云 context ②rope_mode=hier 时层级 RoPE
    # ------------------------------------------------------------------ #
    def forward_varlen(self, coord_list, subcode_list, t, depth, pc=None, pc_target=None, mode=None):
        device = t.device
        coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list]
        subcode_list = [s.to(device=device, dtype=torch.long) for s in subcode_list]
        # 长度在 python 端本来就知道 -> 一并算出各尺度的 max_seqlen, 传给 attention 省掉
        # 每次 flash 调用的 int(lengths.max().item()) 同步(一次前向 72 次, 见 ModernAttention 说明)。
        lengths_py = [int(c.numel()) for c in coord_list]
        if sum(lengths_py) == 0:
            raise ValueError("forward_varlen received an empty batch.")
        if any(n % REFINE_POS_PER_FACE != 0 for n in lengths_py):
            raise ValueError("每个序列长度必须是 9 的倍数(每 face 9 个 refine 位置)。")
        lengths = torch.tensor(lengths_py, device=device, dtype=torch.long)
        max_len_9 = max(lengths_py)
        max_len_3 = max_len_9 // 3          # hourglass: lengths // 3, 整除有断言保证
        max_len_1 = max_len_9 // 9

        coord_packed = torch.cat(coord_list, dim=0)
        subcode_packed = torch.cat(subcode_list, dim=0)
        types = self.types_from_lengths(lengths)

        raw_emb = self.raw_embedding(coord_packed)
        refine_emb = self.embed_refine(subcode_packed, types)
        x = raw_emb + refine_emb

        # nn.Embedding 不在 autocast 的转型名单里 -> 上面的 x 是 fp32。不显式降下来的话:
        #   ① enc_blocks_9 / dec_blocks_9(经 skip_9 带回)这 4 个块的残差流全程 fp32 —— 而它们
        #      正是 token 数最多的一档(hourglass 把其余降到 1/3、1/9), RMSNorm/modulate/残差加
        #      在最大的激活上按 fp32 跑, 内存流量白翻一倍;
        #   ② rope_fn_9 是按 `dtype=x.dtype` 建 cos/sin 的, x 是 fp32 就把 bf16 的 Q/K 提升回 fp32,
        #      导致 can_use_flash_attn 判否、掉进 SDPA 逐序列 python 循环(见 apply_rotary_pos_emb_varlen)。
        # 统一降到 autocast dtype, 让整条残差流和 RoPE 都是 bf16。
        if _BF16_RESIDUAL and torch.is_autocast_enabled():
            x = x.to(torch.get_autocast_dtype("cuda"))

        depth = depth.to(device=device, dtype=torch.long)
        # c_sample **刻意保持 fp32**(按混精惯例): 它只有 [b, d_model], 保高精度的成本可忽略, 而它经
        # adaLN 变成 shift/scale/gate 逐元素乘加到所有激活上 —— 误差是系统性的、不会被平均掉。
        # DiT / SD3 / Flux 都把时间嵌入与调制留在 fp32。对比: 点云 latent 是 cross-attn 的 K/V,
        # 紧接着就进 k_proj/v_proj 的 matmul(autocast 必然转 bf16), 且 flash-attn 不接受 fp32, 故 bf16。
        c_sample = self.time_mlp(t) + self.depth_embedding(depth)   # [b, d_model], fp32

        def c_for(cur_lengths):
            return c_sample.repeat_interleave(cur_lengths, dim=0)

        bsz = len(coord_list)
        # 模式 token(每样本一个),拼到点云 context 尾部(TSSR 式 flag 注入)
        mode_tok = None
        if mode is not None:
            if torch.is_tensor(mode):
                mode_vec = mode.to(device=device, dtype=torch.float).reshape(-1)
            else:
                mode_vec = torch.full((bsz,), float(mode), device=device)
            mode_tok = self.mode_emb(mode_vec).unsqueeze(1)          # [b,1,d]

        # ---- 点云条件:两套(全GT / 目标分辨率塌缩),逐层交替; 各拼上 mode token ----
        def _encode_cond(_pc):
            if not self.use_pc:
                return None, None, None
            if _pc is None:
                emb = self.null_cond.unsqueeze(0).expand(bsz, -1, -1)     # [b,257,d]
            else:
                _pc = _pc.to(device=device)
                emb = self.conditioner(_pc)                              # [b,257,d]
            if mode_tok is not None:
                emb = torch.cat([emb, mode_tok.to(emb.dtype)], dim=1)     # [b,258,d]
            ctx_len = int(emb.size(1))                  # 每样本 context 长度恒为 257(+1 = mode token)
            cl = torch.full((emb.size(0),), ctx_len, device=device, dtype=torch.long)
            return emb.reshape(-1, emb.size(-1)), cl, ctx_len

        cond_a, cond_a_len, cond_max_len = _encode_cond(pc)         # 奇数层(全GT)
        if self.use_pc and (pc_target is not None):
            cond_b, cond_b_len, _ = _encode_cond(pc_target)         # 偶数层(目标分辨率塌缩)
        else:
            cond_b, cond_b_len = cond_a, cond_a_len                # 未提供 -> 两套相同

        def _cond_for(gi):
            return (cond_a, cond_a_len) if (gi % 2 == 0) else (cond_b, cond_b_len)

        # geo/geo_vcat(2026-08-23, 接 sfocal 大模型 warm-start): 三尺度 cos/sin 由 coord 现算,
        # 复用父类 GeoRoPE(锚点/频率梯子与主线逐位一致); hier/flat 走原路径不变。
        geo_scales = (self._geo_rope_scales(coord_packed, x.dtype)
                      if str(getattr(self, "rope_mode", "flat")).startswith("geo") else None)

        def rope_fn_for(cur_lengths, tokens_per_face=9):
            if geo_scales is not None:
                cos, sin = geo_scales[tokens_per_face]
            elif self.rope_mode == "hier":
                position_ids = lengths_to_position_ids(cur_lengths)
                cos, sin = self._hier_cos_sin(position_ids, tokens_per_face, x.dtype)
            else:
                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)

        gi = 0
        rope_fn_9 = rope_fn_for(lengths, 9)
        c9 = c_for(lengths)
        for block in self.enc_blocks_9:
            cnd, cnl = _cond_for(gi); gi += 1
            x = block.forward_varlen(x, c9, lengths, cnd, cnl, rope_apply_fn=rope_fn_9,
                                     max_len=max_len_9, cond_max_len=cond_max_len)
        skip_9 = x

        x, lengths = self.down_9_to_3.forward_varlen(x, lengths)
        rope_fn_3 = rope_fn_for(lengths, 3)
        c3 = c_for(lengths)
        for block in self.enc_blocks_3:
            cnd, cnl = _cond_for(gi); gi += 1
            x = block.forward_varlen(x, c3, lengths, cnd, cnl, rope_apply_fn=rope_fn_3,
                                     max_len=max_len_3, cond_max_len=cond_max_len)
        skip_3 = x

        x, lengths = self.down_3_to_1.forward_varlen(x, lengths)
        rope_fn_1 = rope_fn_for(lengths, 1)
        c1 = c_for(lengths)
        for block in self.bottleneck_blocks:
            cnd, cnl = _cond_for(gi); gi += 1
            x = block.forward_varlen(x, c1, lengths, cnd, cnl, rope_apply_fn=rope_fn_1,
                                     max_len=max_len_1, cond_max_len=cond_max_len)

        x, lengths = self.up_1_to_3.forward_varlen(x, lengths)
        x = x + skip_3
        for block in self.dec_blocks_3:
            cnd, cnl = _cond_for(gi); gi += 1
            x = block.forward_varlen(x, c3, lengths, cnd, cnl, rope_apply_fn=rope_fn_3,
                                     max_len=max_len_3, cond_max_len=cond_max_len)

        x, lengths = self.up_3_to_9.forward_varlen(x, lengths)
        x = x + skip_9
        for block in self.dec_blocks_9:
            cnd, cnl = _cond_for(gi); gi += 1
            x = block.forward_varlen(x, c9, lengths, cnd, cnl, rope_apply_fn=rope_fn_9,
                                     max_len=max_len_9, cond_max_len=cond_max_len)

        return x, types

    # ------------------------------------------------------------------ #
    # 分类器工具
    # ------------------------------------------------------------------ #
    def to_acc_logit(self, h):
        """逐 token 正确性 logit [total];detach 骨干,分类器梯度不回流。"""
        return self.to_acc(h.detach()).squeeze(-1)

    @torch.no_grad()
    def _pos_correct(self, h, target, types):
        """逐位置正确性 label [total] bool:该 refine 位置的头 argmax(teacher-forced)是否==GT。
        F 比 slot0;intra/inter 比 8 行全等。"""
        total = h.size(0)
        correct = torch.zeros(total, dtype=torch.bool, device=h.device)
        mf = types == TYPE_F
        if mf.any():
            pred = self.f_head(h[mf]).argmax(-1)
            correct[mf] = pred == target[mf, 0]
        for ty, head in ((TYPE_INTRA, self.intra_head), (TYPE_INTER, self.inter_head)):
            m = types == ty
            if not m.any():
                continue
            tgt = target[m]
            pred = head.forward_teacher(h[m], tgt).argmax(-1)   # [n,8] teacher-forced
            correct[m] = (pred == tgt).all(dim=1)
        return correct

    # ------------------------------------------------------------------ #
    # 训练入口
    # ------------------------------------------------------------------ #
    def forward(self, coord_list, noisy_subcode_list, t, depth, target_subcode_list=None,
                loss_mask=None, pc=None, pc_target=None, task="mask", acc_wrong_weight=1.0):
        """DDP 单一入口(必须走 model(...) 才能触发梯度同步;自定义方法会绕过 DDP)。
        task="mask"   -> (loss, acc, acc_ne)          吸收腐蚀 CE(仅 masked 位置)
        task="nomask" -> (loss_nomask, loss_acc, frac_correct, cls_acc)  uniform 去噪 CE(全位置)+ 分类器 BCE
        """
        if task == "mask":
            return self.forward_mask(coord_list, noisy_subcode_list, t, depth,
                                     target_subcode_list, loss_mask, pc, pc_target)
        elif task == "nomask":
            return self.forward_nomask(coord_list, noisy_subcode_list, t, depth,
                                       target_subcode_list, pc, pc_target, acc_wrong_weight)
        raise ValueError(f"unknown task {task!r}")




    # ------------------------------------------------------------------ #
    # 采样工具
    # ------------------------------------------------------------------ #
    def _predict_fill(self, h, h_u, types, cfg_t, use_cfg, temperature):
        """由 hidden 出各类型头 argmax(可 CFG/温度),填满 subcodes [total,8]。"""
        total = h.size(0)
        pred = torch.zeros((total, NUM_OCTANTS), dtype=torch.long, device=h.device)
        mf = types == TYPE_F
        if mf.any():
            logits = self.f_head(h[mf])
            if use_cfg:
                lu = self.f_head(h_u[mf]); logits = lu + cfg_t * (logits - lu)
            if temperature and temperature > 0:
                prob = (logits / temperature).softmax(-1)
                am = torch.multinomial(prob.float(), 1).squeeze(-1)
            else:
                am = logits.argmax(-1)
            pred[mf, :] = am.unsqueeze(1).expand(-1, NUM_OCTANTS)
        for ty, head in ((TYPE_INTRA, self.intra_head), (TYPE_INTER, self.inter_head)):
            m = types == ty
            if not m.any():
                continue
            p, _ = head.rollout(h[m], h_u[m] if use_cfg else None, cfg_t, temperature)
            pred[m] = p
        return pred

    @torch.no_grad()
    def predict_subcodes(self, coord_list, sub_list, t, depth, pc=None, pc_target=None, mode=1,
                         cfg_scale=1.0, temperature=0.0):
        """一次前向 + 头 argmax 得填满的 subcodes [total,8](供自生成脏输入 / 单步预测)。"""
        lengths = torch.tensor([c.numel() for c in coord_list], device=t.device, dtype=torch.long)
        types = self.types_from_lengths(lengths)
        use_cfg = (float(cfg_scale) != 1.0) and self.use_pc and (pc is not None)
        h, _ = self.forward_varlen(coord_list, sub_list, t, depth, pc, pc_target, mode=mode)
        h_u = self.forward_varlen(coord_list, sub_list, t, depth, None, None, mode=mode)[0] if use_cfg else None
        return self._predict_fill(h, h_u, types, float(cfg_scale), use_cfg, temperature)

    @torch.no_grad()
    def sample_tssr(self, coord_list, depth, pc=None, pc_target=None, T_infer=10,
                    cfg_scale=1.0, cfg_schedule="constant", cfg_start=1.0,
                    mask_logits_lo=0.2, mask_logits_hi=0.9, temperature=0.0,
                    generator=None, return_stats=False):
        """TSSR 三前向解码。返回 list[[L,8]] local 子码(drop-in 父类 decode 路径)。
        每步:① mask 预测(mode=1)填满 -> ② no-mask 精修(mode=0)+ 同前向出置信 -> ③ re-mask 判错位置。
        末步不再 re-mask。"""
        self.eval()
        device = next(self.parameters()).device
        coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list]
        depth = depth.to(device=device, dtype=torch.long)
        lengths = torch.tensor([c.numel() for c in coord_list], device=device, dtype=torch.long)
        types = self.types_from_lengths(lengths)
        want_cfg = (float(cfg_scale) != 1.0 or float(cfg_start) != 1.0) and self.use_pc and (pc is not None)
        total = int(lengths.sum().item())

        subcodes = torch.full((total, NUM_OCTANTS), MASK_ID, dtype=torch.long, device=device)
        sched = torch.linspace(float(mask_logits_lo), float(mask_logits_hi), max(T_infer, 1), device=device)
        n_remask_hist = []

        for step in range(T_infer):
            if cfg_schedule == "linear":
                cfg_t = float(cfg_start) + (float(cfg_scale) - float(cfg_start)) * step / max(T_infer - 1, 1)
            else:
                cfg_t = float(cfg_scale)
            use_cfg = want_cfg and (cfg_t != 1.0)
            t = torch.full((len(coord_list),), step / T_infer, device=device, dtype=torch.float)

            # ① mask 预测(mode=1)
            sub_list = list(torch.split(subcodes, lengths.tolist(), dim=0))
            h1, _ = self.forward_varlen(coord_list, sub_list, t, depth, pc, pc_target, mode=1)
            h1u = self.forward_varlen(coord_list, sub_list, t, depth, None, None, mode=1)[0] if use_cfg else None
            x_t = self._predict_fill(h1, h1u, types, cfg_t, use_cfg, temperature)

            # ② no-mask 精修(mode=0)+ 复用该前向出逐 token 置信
            sub_list2 = list(torch.split(x_t, lengths.tolist(), dim=0))
            h2, _ = self.forward_varlen(coord_list, sub_list2, t, depth, pc, pc_target, mode=0)
            h2u = self.forward_varlen(coord_list, sub_list2, t, depth, None, None, mode=0)[0] if use_cfg else None
            x_t = self._predict_fill(h2, h2u, types, cfg_t, use_cfg, temperature)
            acc = torch.sigmoid(self.to_acc(h2).squeeze(-1)).float()      # [total] ∈[0,1]

            # ③ re-mask 判错位置(末步不做)
            if step < T_infer - 1:
                mask_prob = ((1.0 - acc) * sched[step]).clamp(1e-6, 1.0 - 1e-6)
                if generator is not None:
                    r = torch.rand(total, device=device, generator=generator)
                else:
                    r = torch.rand(total, device=device)
                remask = r < mask_prob
                subcodes = x_t.clone()
                subcodes[remask] = MASK_ID
                n_remask_hist.append(int(remask.sum().item()))
            else:
                subcodes = x_t

        out = list(torch.split(subcodes, lengths.tolist(), dim=0))
        if return_stats:
            return out, {"n_remask_per_step": n_remask_hist, "total_pos": total}
        return out