Upload folder using huggingface_hub

mulres_infer/.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.pyc
+ckpt/*.pth
+ckpt/*.ckpt
+outputs/

mulres_infer/ckpt/generator.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:edf562427ae7c583d5185890769e24266455346650a4292ab41d00a29eddf6a5
+size 3284020143

mulres_infer/ckpt/readme.md

@@ -0,0 +1,28 @@
+# 权重
+
+推理需要两个文件，放在本目录下：
+
+| 文件 | 说明 | 大小 |
+|---|---|---|
+| `generator.pth` | 低分辨率生成器（TSSR 三前向离散扩散），逐级生成到格 64 | 3.1 GB |
+| `subdivider.pth` | 高分辨率细分器（acc_uni），64 → 128 → 256 → 512 | 1.2 GB |
+
+两者都**只含 `model_state_dict`**，已剥离优化器/调度器状态，且 EMA 影子权重已合并进权重本体
+（推理的原始行为就是「加载 model_state_dict 后用 EMA 覆盖」，合并后逐位等价）。
+
+想自己从训练产出的 ckpt 转换：
+
+```bash
+python scripts/strip_ckpt.py <训练ckpt.pth> ckpt/generator.pth
+# 想要未经 EMA 平滑的原始权重: 加 --no-ema
+```
+
+## 可选
+
+`shapevae-256.ckpt`（Michelangelo 点云编码器）**推理不需要**——点云条件编码器的权重已包含在
+`generator.pth` 内，缺失时 `miche_conditioner` 会先构造未初始化的编码器再被权重覆盖，结果不受影响。
+只有从头训练才需要它。
+
+## 路径覆盖
+
+环境变量 `RAIL_CKPT` / `SUPERRES_CKPT` / `SHAPEVAE_CKPT`，或 `infer.py --ckpt / --superres-ckpt`。

mulres_infer/ckpt/subdivider.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55feaa1feaf4254f2ba9d94cc2c3def528bbb5408068a8191f636dbcd55e2bd2
+size 1258077735

mulres_infer/configs/generator.yaml

@@ -0,0 +1,151 @@
+# 几何 RoPE vcat 线(2026-08-05): geo 配方不变, 只把粗尺度(3/face,1/face)锚点从"面质心"
+# 换成"三顶点各占1/3通道拼接"(rope_mode="geo_vcat", 用户提议) —— 共质心三角形可分、
+# 顶点信息在粗尺度保留。9/face 尺度与 geo 逐位一致。从 geo ckpt_47 剥离 init 热启动。
+#   (token 位置 = 3D 锚点[顶点/边中点/面质心] + 面内类型, 序列顺序不参与注意力, 面片置换等变;
+#    设计动机: 跨面"同顶点=同相位"取代 z-major 排序邻近这个有损协调代理, 见 model_acc_fix3_altpc.__init__)。
+#   [vcat_1024] 放大版(2026-08-08 用户定): 1024宽/16头/[2,4,12]/ffn2816, 骨干624M+编码器~193M≈818M。
+#   init = vcat ckpt_18 的 Net2Net 扩展(script/expand_vcat_768_to_1024.py: 旧权重按组拷入/新维小随机/
+#   新谷底块6-11零门控恒等), strict 加载。
+#   [首启后] 本目录出现自己的 ckpt 后, 把 init_ckpt_path 置 null 切回 resume_latest(否则重启永远回 init)。
+
+project_name: "gobjaverse_fix3_octant_factorized_1_to_1024_geo_vcat_1024_pc1024_sfocal_tssr"
+output_folder_base: "./pretrained_model"
+
+data:
+  # 2026-09-10 上集群: 与 sfocal config 同步做前缀迁移(内容同源, /root/dataset -> 仓库内 dataset_lists/,
+  # mesh 实体在 /data/home/from-baidu-h100-jfs/)。仍是 08-21 用户定的 focal 时代旧数据, 未换数据集。
+  train_dataset_path: "./dataset_lists/under_20k_concat.txt"
+  eval_dataset_path: "./dataset_lists/unique_files_glb_under20000face_compress_glb_with_info_val.txt"
+  num_workers: 12
+
+data_processing:
+  overfit: false
+  y_up: true
+  n_discrete_size: 128
+  vertex_order: "zyx"
+  max_len: 600000             # focal 时代口径(60/face, 1万面上限)
+  min_len: 0
+  max_face_num: 10000         # focal 时代口径
+  min_face_num: 200
+  pc_num: 40960
+  total_pc_num: 50000
+  windowing:
+    window_size: -1
+    window_stride: 9000
+    split_method: "sliding_window"
+  train_augmentation:
+    shift_augment: false
+    scale_augment: true
+    dec_augment: false
+    rot_augment: false         # focal 时代口径
+    sample_augment: true
+  val_augmentation:
+    shift_augment: false
+    scale_augment: false
+    dec_augment: false
+    rot_augment: false
+    sample_augment: false
+
+tokenizer:
+  max_depth: 10
+  max_resolution: 1024
+  use_hourglass: true
+  max_children: 8
+  target_depth: [0,1,2,3,4,5,6]   # 用户定(08-19): 砍掉深度7(格256)提速, 训到格128
+
+model:
+  feature_dim: 1024
+  num_heads: 16               # head_dim 保持 64(flash+rope 槽位不变)
+  hourglass_depth: [2, 4, 12]  # 谷底 6->12; 新块 6-11 由扩展 init 零门控(起步恒等)
+  dropout: 0.0
+  encoder_freeze: false
+  conditioned_on_pc: true
+  max_seq_len: 150000
+  # ---- 几何 RoPE ----
+  rope_mode: "geo_vcat"       # 顶点显式锚(2026-08-05 用户提议): 粗尺度=三顶点各1/3通道拼接;
+                               # 9/face=双端点布局(F/M_ii=(v,v)复制, M_ij=(vi,vj)拆分, 治对角线中点混叠)
+  rope_geo_scale: 1.0          # 体素单位(不按 depth 归一, 理由见模型注释)
+  rope_theta_xyz: 100.0        # 邻格 1 rad 可分辨; depth7 全程 128 格最低频 2.0 rad < π
+  rope_theta_type: 10.0
+  rope_type_slots: 2           # 32 槽 = xyz 30 + type 2; vcat 粗尺度切 3顶点x9槽+type2+pad3
+  rope_theta_vcat: 300.0       # 3槽/轴的频率梯子(邻格1rad可分辨, 128格最低频2.85rad<π)
+  ffn_hidden: 2816             # 显式覆盖(8/3*1024=2730 非 256 倍数, 取 2816 对齐 Tensor Core)
+  pc_num_latents: 1024         # 点云编码器 latent 1+1024(query 平铺扩容, 注入方式不变: 每层 cross-attn 双云交替)
+  miche_sdpa: true             # miche 注意力 SDPA 化(省 ~20GB 瞬时峰, fpsq 线验证过, 数值等价)
+
+training:
+  skip_resume: false
+  # 冷启动: 从 hg246 ckpt_39 的剥离 init(只留 model+ema)起步, epoch 从 0 计。
+  # 首启跑出本目录 ckpt 后改: init_ckpt_path: null
+  init_ckpt_path: null         # ckpt_1(0.1894)已落盘(08-23), 切回 resume 语义; 冷启曾用 sfocal/tssr_init_from_ep24.pth
+  resume_latest: true
+  load_strict: false           # TSSR 新头 mode_emb/to_acc fresh(零初始化尾层, warm-start 友好)
+  reset_epoch: false           # 已切 resume 语义
+
+  epoches_num: 40
+  base_lr: 3.0e-5              # warm-start 大模型 + 新头 + 新目标(uniform/分类器), 小 lr 防骨干漂移; 稳定后可提
+  min_lr: 1.0e-6
+  warmup_ratio: 0.02
+
+  cond_drop_prob: 0.1
+  pc_target_drop_prob: 1.0    # 恒丢云B => 每层恒定注入同一朵全分辨率云A(用户最终定案: 经典focal式单云)
+  single_target_pc: false      # 已否决(08-21二次修正): 用户要的是全分辨率云A, 不是云B
+
+  nonempty_row_weight: 1.0     # 原始 focal 的正确配置(focal γ 即类平衡, 别叠 w2; 2.0是后来的w2密度fork)
+  focal_gamma: 2.0            # focal loss(同 sfocal)
+  # ----- TSSR 三前向超参(与旧 tssr 线一致) -----
+  tssr_p_mask: 0.5             # 每 step 走 mask 步的概率(逐 step 跨 rank 同步)
+  tssr_p_selfgen: 0.1          # no_mask 步用自生成脏输入的概率(治暴露偏差)
+  tssr_acc_loss_weight: 1.0    # 逐 token 正确性分类器 BCE 权重
+  tssr_acc_pos_weight: 1.0     # BCE 中「错」类上权
+  tssr_mask_logits_lo: 0.2     # sample_tssr 分类器 re-mask 阈值下限
+  tssr_mask_logits_hi: 0.9     # 上限
+
+  use_ema: true
+  ema_decay: 0.999
+  ema_start_step: 0
+  ema_update_every: 1
+
+  grad_checkpoint: true        # 1024x24块@90k闸 激活~2x于768线的69GB悬崖, 必开(激活省~10x/步时+30%)
+  save_ckpt_every: 1
+  save_ckpt_step_num: 3000     # 9h/epoch(9400步)且宿主机不稳: ~每1.8h一个中途档, 防整epoch丢失
+  batch_size: 6
+  diffusion_train_steps: 500
+  optimizer:
+    type: "AdamW"
+    weight_decay: 0.01
+    betas: [0.9, 0.95]
+  debug_fixed_mask: false
+  debug_fixed_t: None
+  debug_train_on_all_valid_tokens: false
+  dynamic_token_batch:
+    enabled: true
+    tokens_per_face: 9
+    budget_mode: "varlen"
+    # 2026-07-30 批点重定(零丢样档, 用户选定): 校准 ratio 后 est≈1.31x 实际, 预算 95k -> 批实际
+    # 位置中位 ~58k / P90 ~76k; 闸 72k 只兜长尾(全量模拟: 触闸批 2%, 丢 0.3% 样本/0.6% 位置)。
+    # 旧点 60k+旧ratio+55k 闸实际在静默丢 12% 样本/22% 位置(偏丢深 depth 大 mesh);
+    # tssr 式 340k 会丢 64% 位置 —— 预算永远要和 hard_max_positions 对齐着调, 别只动一个。
+    # 2026-07-30 三调(用户定 90k): 闸 72k 实测触发 ~14% 批(全在 72-82k 区间, 丢 ~1.4% 样本,
+    # 实际长度绕估计 ±20% 噪声推批过阈)。90k 下预算 95k(est≈1.31x实际→实际中位 ~72k)几乎不可能触闸,
+    # 只剩单样本巨 mesh 长尾。实测 bs14+闸72k 峰值 64.9GB, +18k 位置 ≈ +4-5GB -> ~69-70GB, 安全。
+    max_tokens_per_batch: 200000   # 用户定(2026-08-12): 容纳2万面样本@depth7(est~164k), 大mesh获得全深度覆盖
+    max_batch_size: 20
+    hard_max_positions: 250000     # 用户定: 预算-闸余量25%, 兜深样本组合的估计低估长尾
+    bucket_size: 2048
+    length_safety_factor: 1.05
+
+accelerator:
+  gradient_accumulation_steps: 1
+  mixed_precision: "bf16"
+
+evaluation:
+  generate_only_after_training: false
+  generate_val_mesh_every: 3   # 大 val mesh 逐个采样很慢, 降频防 epoch 末长尾(NCCL 超时已提1h兜底)
+  diffusion_sample_steps: 20
+  pc_num: 1024
+  cfg_scale: 1.5
+  cfg_schedule: "linear"
+  cfg_start: 1.0
+  best_of_n: 0
+  sample_temperature: 0.0

mulres_infer/configs/subdivider.yaml

@@ -0,0 +1,126 @@
+project_name: "gobjaverse_5k_diffusion_pc_fix_tokenizer_1_to_1024_hourglass_fix_pc_more_data_10k_acc_uni_fix_tokenizer"
+output_folder_base: "./pretrained_model"
+
+data:
+  train_dataset_path: "/root/dataset/under_20k_concat.txt"
+  eval_dataset_path: "/root/dataset/unique_files_glb_under20000face_compress_glb_with_info_val.txt"
+  num_workers: 12
+
+data_processing:
+  overfit: false
+  y_up: true
+  n_discrete_size: 128
+  vertex_order: "zyx"
+  max_len: 270000
+  min_len: 0
+  max_face_num: 13000
+  min_face_num: 200
+  pc_num: 40960
+  total_pc_num: 50000
+  windowing:
+    window_size: -1
+    window_stride: 9000
+    split_method: "sliding_window"
+  train_augmentation:
+    shift_augment: false
+    scale_augment: true
+    dec_augment: false
+    rot_augment: false
+    sample_augment: true
+  val_augmentation:
+    shift_augment: false
+    scale_augment: false
+    dec_augment: false
+    rot_augment: false
+    sample_augment: false
+
+tokenizer:
+  max_depth: 10
+  max_resolution: 1024
+  use_hourglass: true
+  max_children: 3
+  target_depth: [0,1,2,3,4,5,6,7,8,9,10]
+
+model:
+  codebook_path: "/root/mul-res/checkpoints/octree_codebook_meshgpt_shapenet_d5_6.npz"
+  use_hourglass: true
+  depth: 12
+  model_version: "v1"
+  feature_dim: 768
+  num_heads: 12
+  root_pred_depth: 3
+  context_embedding_dim: 256
+  num_classify: 100
+  hourglass_depth: [2, 2, 4]
+  input_embe_pos: "no"
+  root_pred_method: "index"
+  conditioned_on_pc: true
+  encoder_freeze: false
+  dropout: 0.0
+
+training:
+  # ---- fine-tune knobs ----
+  # checkpoint_394_..._loss.pth 是 pre-CFG (496 参数, 无 null_cond / EMA, 用它做
+  # init_ckpt_path 会触发 optimizer param-group size mismatch).
+  # 真正的 refine ckpt 是下面这个 (497 参数, 有 null_cond + ema_state_dict).
+  init_ckpt_path: "/root/junkai_lin/mul_res_diffusion/pretrained_model/gobjaverse_5k_diffusion_pc_fix_tokenizer_1_to_1024_hourglass_fix_pc_more_data_10k_acc_uni/checkpoint_19_epoch_0.1954_loss.pth"
+  load_strict: true           # refine ckpt 已含 null_cond, 严格加载更安全
+  reset_epoch: true           # epoch 从 0 重新计数
+  skip_resume: false
+
+  # ---- LR / 时长 ----
+  epoches_num: 20              # 由于 trainer 循环 off-by-one (range(1, epoches_num)), 实际只跑 5 epoch
+  base_lr: 5.0e-5             # 比原训练 (1e-4) 低 10x, 不毁掉已收敛的权重
+  min_lr: 1.0e-6
+  warmup_ratio: 0.02
+
+  # ---- CFG cond-drop ----
+  cond_drop_prob: 0.1         # 10% batch 把 PC 置空, 学习 null_cond
+
+  # ---- EMA (开启 + 快 decay) ----
+  use_ema: true
+  ema_decay: 0.99             # fine-tune 步数少, 必须用快 decay
+  ema_start_step: 0           # 从第一步就开始累积
+  ema_update_every: 1
+
+  # ---- 保留的训练参数 ----
+  save_ckpt_every: 1
+  save_ckpt_step_num: 10000
+  batch_size: 3
+  diffusion_train_steps: 500
+  optimizer:
+    type: "AdamW"
+    weight_decay: 0.01
+    betas: [0.9, 0.95]
+  loss_weights:
+    root_pred_weight: 0.2
+  root_method: "root_move_index"
+  root_embed: false
+  debug_fixed_mask: false
+  debug_fixed_batch: false
+  debug_fixed_t: None
+  debug_train_on_all_valid_tokens: false
+  dynamic_token_batch:
+    enabled: true
+    tokens_per_face: 9
+    budget_mode: "varlen"
+    max_tokens_per_batch: 117000
+    max_batch_size: 6
+    bucket_size: 2048
+    length_safety_factor: 1.05
+
+accelerator:
+  gradient_accumulation_steps: 4
+  mixed_precision: "bf16"
+
+evaluation:
+  generate_only_after_training: false
+  debug_root: false
+  generate_val_mesh_every: 1
+  completion_input_len: 18
+  pc_num: 1024
+  use_toppk: true
+  top_k: 8
+  top_p: 0.8
+  temperature: 1
+  diffusion_sample_steps: 40

mulres_infer/dataset/mesh_dataset.py

@@ -0,0 +1,940 @@
+import trimesh
+import networkx as nx
+import numpy as np
+import os, pickle
+import json
+import math
+import threading
+import torch
+from torch.utils.data import Dataset
+import glob
+from torch.utils.data import Sampler
+from tqdm import tqdm
+import numpy as np
+from utils.data_utils import undiscretize, discretize, normalize, clean_mesh
+import logging
+from concurrent.futures import ProcessPoolExecutor, as_completed
+import random
+import multiprocessing as mp
+from pathlib import Path
+import open3d
+from scipy.spatial import cKDTree
+import concurrent.futures
+from PIL import Image
+
+
+# ----------------------------------------------------------------------------
+# nvdiffrast on-the-fly normal-map renderer.
+#
+# Vendored from /root/junkai_lin/TRELLIS/dataset_toolkits/render_normals_gpu.py
+# so the dataset stays self-contained. When `image_cond.render_on_the_fly` is
+# enabled, we render the AUGMENTED mesh directly in the dataset's world frame
+# (Z-up, normalized to [-1, 1]) so the projection math in
+# `diffusion/img_cond_conditioner.py` lines up without needing the
+# `undo_dataset_y_up` swap or the 0.5 pc_render_scale that the TRELLIS
+# pre-renders required. Set those config flags accordingly when switching to
+# render-on-the-fly mode (see img_cond yaml).
+#
+# Render failures (e.g. degenerate mesh, CUDA OOM, zero-face mesh) raise; the
+# `except Exception` loop in __getitem__ then resamples a different mesh.
+# ----------------------------------------------------------------------------
+
+_NVDR_CTX = None
+_NVDR_LOCK = threading.Lock()
+
+
+def _get_nvdr_ctx():
+    """Lazily create one nvdiffrast CUDA context per process.
+
+    DataLoader workers run as separate processes; each one initializes its own
+    context the first time __getitem__ touches the renderer. The default
+    `fork` start method may conflict with CUDA already initialized in the
+    parent — if you hit `cuda: initialization error` here, switch the
+    DataLoader to `multiprocessing_context='spawn'`.
+    """
+    global _NVDR_CTX
+    if _NVDR_CTX is not None:
+        return _NVDR_CTX
+    with _NVDR_LOCK:
+        if _NVDR_CTX is None:
+            import nvdiffrast.torch as dr
+            if not torch.cuda.is_available():
+                raise RuntimeError(
+                    "image_cond.render_on_the_fly=true requires CUDA; no GPU visible to this worker"
+                )
+            device = torch.device("cuda", 0)
+            _NVDR_CTX = {"glctx": dr.RasterizeCudaContext(device=device), "device": device}
+    return _NVDR_CTX
+
+
+def _build_camera_for_render(yaw, pitch, radius, fov, aspect=1.0, near=0.01, far=100.0):
+    """Z-up world, origin-tracking camera. yaw/pitch/fov in radians.
+
+    Returns (view[4,4], proj[4,4], cam_pos[3], R_c2w[3,3]) as float64 numpy.
+    Matches `render_normals_gpu.build_camera` exactly so the pose convention
+    is consistent with the TRELLIS pre-render path.
+    """
+    cx = radius * math.cos(yaw) * math.cos(pitch)
+    cy = radius * math.sin(yaw) * math.cos(pitch)
+    cz = radius * math.sin(pitch)
+    cam_pos = np.array([cx, cy, cz], dtype=np.float64)
+
+    forward = -cam_pos / np.linalg.norm(cam_pos)
+    world_up = np.array([0.0, 0.0, 1.0])
+    right = np.cross(forward, world_up)
+    rn = np.linalg.norm(right)
+    if rn < 1e-9:
+        world_up = np.array([0.0, 1.0, 0.0])
+        right = np.cross(forward, world_up)
+        rn = np.linalg.norm(right)
+    right /= rn
+    up = np.cross(right, forward)
+    R_c2w = np.stack([right, up, -forward], axis=1)
+    R_w2c = R_c2w.T
+    t_w2c = -R_w2c @ cam_pos
+
+    view = np.eye(4, dtype=np.float64)
+    view[:3, :3] = R_w2c
+    view[:3, 3] = t_w2c
+
+    f = 1.0 / math.tan(fov / 2.0)
+    proj = np.zeros((4, 4), dtype=np.float64)
+    proj[0, 0] = f / aspect
+    # nvdiffrast image-space y is flipped vs OpenGL; negate so row 0 = top.
+    proj[1, 1] = -f
+    proj[2, 2] = -(far + near) / (far - near)
+    proj[2, 3] = -2.0 * far * near / (far - near)
+    proj[3, 2] = -1.0
+    return view, proj, cam_pos, R_c2w
+
+
+def _explode_to_flat_normals(verts, faces):
+    """Per-face geometric normals via mesh explosion (flat shading)."""
+    F = faces.shape[0]
+    v0 = verts[faces[:, 0]]
+    v1 = verts[faces[:, 1]]
+    v2 = verts[faces[:, 2]]
+    fn = np.cross(v1 - v0, v2 - v0)
+    fn_len = np.linalg.norm(fn, axis=1, keepdims=True)
+    fn = fn / np.where(fn_len < 1e-12, 1.0, fn_len)
+    new_verts = np.empty((F * 3, 3), dtype=verts.dtype)
+    new_verts[0::3] = v0
+    new_verts[1::3] = v1
+    new_verts[2::3] = v2
+    new_normals = np.repeat(fn, 3, axis=0).astype(verts.dtype)
+    new_faces = np.arange(F * 3, dtype=np.int32).reshape(F, 3)
+    return new_verts, new_faces, new_normals
+
+
+def _fit_radius_to_silhouette(verts_t, yaw, pitch, fov, init_radius, target_fill=0.9, iters=3):
+    """Shrink/grow camera radius so the mesh's max |NDC| equals target_fill."""
+    radius = float(init_radius)
+    ones = torch.ones((verts_t.shape[0], 1), device=verts_t.device, dtype=verts_t.dtype)
+    v_h = torch.cat([verts_t, ones], dim=1)
+    for _ in range(iters):
+        view, proj, _, _ = _build_camera_for_render(yaw, pitch, radius, fov)
+        mvp = (proj @ view).astype(np.float32)
+        mvp_t = torch.from_numpy(mvp).to(verts_t.device)
+        v_clip = v_h @ mvp_t.T
+        front = v_clip[:, 3] > 1e-3
+        if not front.any():
+            break
+        c = v_clip[front]
+        ndc_xy = c[:, :2] / c[:, 3:4].clamp_min(1e-6)
+        max_ndc = float(ndc_xy.abs().max().item())
+        if max_ndc < 1e-6:
+            break
+        radius *= max_ndc / target_fill
+        radius = max(0.5, min(radius, 50.0))
+    return radius
+
+
+def _render_normals_image(
+    verts_np,
+    faces_np,
+    yaw_deg=315.0,
+    pitch_deg=30.0,
+    fov_deg=40.0,
+    init_radius=4.0,
+    resolution=518,
+    auto_radius=True,
+    target_fill=0.9,
+):
+    """Render a normal-visualization image of `(verts, faces)` in the mesh's
+    own world frame (assumed Z-up).
+
+    Returns
+    -------
+    image_tensor : torch.float16, [3, H, W], values in [0, 1] (white background)
+    transform_matrix : torch.float32, [4, 4]  (camera-to-world, NeRF convention)
+    camera_angle_x : torch.float32, scalar    (horizontal FOV in radians)
+    """
+    import nvdiffrast.torch as dr
+
+    if faces_np is None or faces_np.shape[0] == 0:
+        raise ValueError("render: mesh has zero faces")
+
+    ctx = _get_nvdr_ctx()
+    device = ctx["device"]
+    glctx = ctx["glctx"]
+
+    yaw = math.radians(yaw_deg)
+    pitch = math.radians(pitch_deg)
+    fov = math.radians(fov_deg)
+
+    verts_np = np.ascontiguousarray(verts_np, dtype=np.float32)
+    faces_np = np.ascontiguousarray(faces_np, dtype=np.int32)
+    verts_np, faces_np, normals_np = _explode_to_flat_normals(verts_np, faces_np)
+
+    verts_t = torch.from_numpy(verts_np).to(device)
+    faces_t = torch.from_numpy(faces_np).to(device).to(torch.int32)
+    normals_t = torch.from_numpy(normals_np).to(device)
+    ones = torch.ones((verts_t.shape[0], 1), device=device, dtype=verts_t.dtype)
+    verts_h = torch.cat([verts_t, ones], dim=1)
+
+    radius = (
+        _fit_radius_to_silhouette(verts_t, yaw, pitch, fov, init_radius, target_fill)
+        if auto_radius
+        else float(init_radius)
+    )
+
+    view_M, proj_M, cam_pos, R_c2w = _build_camera_for_render(yaw, pitch, radius, fov)
+    mvp = (proj_M @ view_M).astype(np.float32)
+    mvp_t = torch.from_numpy(mvp).to(device)
+
+    clip = (verts_h @ mvp_t.T).unsqueeze(0).contiguous()  # (1, V, 4)
+    rast, _ = dr.rasterize(glctx, clip, faces_t, resolution=(resolution, resolution))
+    alpha = (rast[0, ..., 3] > 0).float()  # (H, W)
+    if float(alpha.sum().item()) < 4.0:
+        raise ValueError("render: empty silhouette (mesh out of frustum?)")
+
+    normal_img, _ = dr.interpolate(normals_t.unsqueeze(0), rast, faces_t)
+    normal_img = normal_img[0]  # (H, W, 3)
+    nlen = torch.linalg.norm(normal_img, dim=-1, keepdim=True).clamp_min(1e-8)
+    normal_img = normal_img / nlen
+    normal_img = normal_img * alpha.unsqueeze(-1)
+
+    # Visualize normals as RGB: (N * 0.5 + 0.5), white background.
+    vis = (normal_img * 0.5 + 0.5).clamp(0.0, 1.0)
+    a = alpha.unsqueeze(-1)
+    vis = vis * a + (1.0 - a)  # background = white
+    image_tensor = vis.permute(2, 0, 1).contiguous().to(dtype=torch.float16).cpu()
+
+    transform_matrix = torch.eye(4, dtype=torch.float32)
+    transform_matrix[:3, :3] = torch.from_numpy(R_c2w).float()
+    transform_matrix[:3, 3] = torch.from_numpy(cam_pos).float()
+    camera_angle_x = torch.tensor(float(fov), dtype=torch.float32)
+    return image_tensor, transform_matrix, camera_angle_x
+
+
+def _read_image_cond_cfg(opt):
+    """Return (enabled, render_root, image_resolution, render_on_the_fly,
+    render_cfg, render_manifest_path).
+
+    `render_on_the_fly=True` enables the nvdiffrast path. In that mode
+    `render_root` is unused (no disk-render filter), but it may still be set
+    in the yaml for backward compatibility.
+
+    `render_manifest_path` is the path to a precomputed list of valid render
+    subdir stems (one per line). When set and the file exists, the dataset
+    skips the slow per-line `os.path.exists` and uses O(1) set membership.
+    Build it with `script/build_render_manifest.py`.
+    """
+    cfg = getattr(opt, "image_cond", None)
+    if cfg is None or not getattr(cfg, "enabled", False):
+        return False, None, None, False, None, None
+    image_resolution = int(getattr(cfg, "image_resolution", 512))
+    render_on_the_fly = bool(getattr(cfg, "render_on_the_fly", False))
+    render_root = getattr(cfg, "render_root", None)
+    if not render_on_the_fly and render_root is None:
+        raise ValueError(
+            "image_cond.enabled=true but image_cond.render_root is not set "
+            "(and render_on_the_fly is false)"
+        )
+    render_cfg = getattr(cfg, "render", None)
+    render_manifest_path = getattr(cfg, "render_manifest_path", None)
+    return True, render_root, image_resolution, render_on_the_fly, render_cfg, render_manifest_path
+
+
+def _render_dir_for(file_path: str, render_root: str) -> str:
+    return os.path.join(render_root, Path(file_path).stem)
+
+
+def _has_render(file_path: str, render_root: str) -> bool:
+    d = _render_dir_for(file_path, render_root)
+    return os.path.exists(os.path.join(d, "transforms.json")) and os.path.exists(
+        os.path.join(d, "000.png")
+    )
+
+
+# Per-path LRU cache so the manifest is read once per process.
+_RENDER_MANIFEST_CACHE: dict = {}
+
+
+def _load_render_manifest(manifest_path):
+    """Load a manifest file into a set of valid render stems.
+
+    Returns None when `manifest_path` is None/empty or the file doesn't
+    exist, so the caller falls back to per-line `os.path.exists`. Build the
+    manifest with `script/build_render_manifest.py` whenever the render
+    contents change — it drops startup time from O(N stat) to O(N lookup).
+    """
+    if not manifest_path:
+        return None
+    if manifest_path in _RENDER_MANIFEST_CACHE:
+        return _RENDER_MANIFEST_CACHE[manifest_path]
+    if not os.path.isfile(manifest_path):
+        _RENDER_MANIFEST_CACHE[manifest_path] = None
+        return None
+    with open(manifest_path, "r") as f:
+        stems = {line.strip() for line in f if line.strip()}
+    _RENDER_MANIFEST_CACHE[manifest_path] = stems
+    return stems
+
+
+# TRELLIS Blender depth encoding (see dataset_toolkits/blender_script/render.py):
+#   pixel_val (uint16) = ((depth_world - depth.min) / (depth.max - depth.min)) * 65535
+# Background ends up clamped to 65535 (== depth.max output). We treat 65535 as
+# "no surface" so points landing on BG pixels are considered occluded (not
+# visible) in that view.
+_DEPTH_BG_VALUE = 65535
+
+
+def _load_depth_map(depth_png_path: str, depth_min: float, depth_max: float,
+                    target_resolution: int):
+    """Decode TRELLIS depth PNG to world-space depth.
+
+    Returns (depth[H,W] float32, valid_mask[H,W] bool).
+    `valid_mask` is False on background pixels (no surface hit).
+    """
+    arr = np.array(Image.open(depth_png_path))
+    if arr.dtype != np.uint16:
+        # 8-bit fallback — still treat the brightest value as background.
+        max_val = np.iinfo(arr.dtype).max
+        bg_val = max_val
+        norm = arr.astype(np.float32) / float(max_val)
+    else:
+        bg_val = _DEPTH_BG_VALUE
+        norm = arr.astype(np.float32) / 65535.0
+    valid = arr != bg_val
+    depth = depth_min + norm * (depth_max - depth_min)
+    if (arr.shape[0], arr.shape[1]) != (target_resolution, target_resolution):
+        # Resize via PIL (nearest) so depth values aren't blended across the
+        # silhouette — bilinear would corrupt the foreground/background edge.
+        depth_img = Image.fromarray(depth.astype(np.float32), mode="F").resize(
+            (target_resolution, target_resolution), Image.NEAREST
+        )
+        valid_img = Image.fromarray(valid.astype(np.uint8) * 255, mode="L").resize(
+            (target_resolution, target_resolution), Image.NEAREST
+        )
+        depth = np.asarray(depth_img, dtype=np.float32)
+        valid = np.asarray(valid_img, dtype=np.uint8) > 0
+    return depth, valid
+
+
+def _load_image_and_camera(file_path: str, render_root: str, image_resolution: int):
+    """Load BOTH rendered views + depth maps for back-projection.
+
+    Returns:
+        image_tensor       : [V, 3, H, W] float16 in [0, 1] (white bg)
+        transform_matrix   : [V, 4, 4]    float32   (camera-to-world)
+        camera_angle_x     : [V]          float32   (horizontal FOV, rad)
+        depth_map          : [V, H, W]    float32   (world-space depth)
+        depth_valid        : [V, H, W]    bool      (False on background)
+
+    Iterates every frame listed in transforms.json (the dataset stores two
+    views: 000 and 001). Missing depth PNGs raise — the caller will catch and
+    resample a different mesh.
+    """
+    d = _render_dir_for(file_path, render_root)
+    with open(os.path.join(d, "transforms.json"), "r") as f:
+        meta = json.load(f)
+
+    images = []
+    transforms = []
+    fovs = []
+    depths = []
+    valids = []
+    for frame in meta["frames"]:
+        rgb_path = os.path.join(d, frame["file_path"])
+        img = Image.open(rgb_path)
+        if img.mode == "RGBA":
+            bg = Image.new("RGB", img.size, (255, 255, 255))
+            bg.paste(img, mask=img.split()[-1])
+            img = bg
+        else:
+            img = img.convert("RGB")
+        if img.size != (image_resolution, image_resolution):
+            img = img.resize((image_resolution, image_resolution), Image.LANCZOS)
+        arr = np.asarray(img, dtype=np.float32) / 255.0
+        images.append(torch.from_numpy(arr).permute(2, 0, 1).to(dtype=torch.float16))
+
+        transforms.append(
+            torch.tensor(frame["transform_matrix"], dtype=torch.float32)
+        )
+        fovs.append(torch.tensor(float(frame["camera_angle_x"]), dtype=torch.float32))
+
+        depth_info = frame.get("depth")
+        if depth_info is None:
+            raise ValueError(
+                f"transforms.json frame {frame.get('file_path')} has no 'depth' entry; "
+                "back-projection requires depth bounds."
+            )
+        stem = os.path.splitext(frame["file_path"])[0]
+        depth_path = os.path.join(d, f"{stem}_depth.png")
+        depth_np, valid_np = _load_depth_map(
+            depth_path,
+            float(depth_info["min"]),
+            float(depth_info["max"]),
+            image_resolution,
+        )
+        depths.append(torch.from_numpy(depth_np))
+        valids.append(torch.from_numpy(valid_np))
+
+    image_tensor = torch.stack(images, dim=0)  # [V, 3, H, W]
+    transform_matrix = torch.stack(transforms, dim=0)  # [V, 4, 4]
+    camera_angle_x = torch.stack(fovs, dim=0)  # [V]
+    depth_map = torch.stack(depths, dim=0)  # [V, H, W]
+    depth_valid = torch.stack(valids, dim=0)  # [V, H, W]
+    return image_tensor, transform_matrix, camera_angle_x, depth_map, depth_valid
+def shift_vertices(vertices, x_lims=(-0.1, 0.1), y_lims=(-0.1, 0.1), z_lims=(-0.075, 0.075)):
+    # shift x, y, z
+    x = np.random.uniform(low=x_lims[0], high=x_lims[1], size=(1,))
+    y = np.random.uniform(low=y_lims[0], high=y_lims[1], size=(1,))
+    z = np.random.uniform(low=z_lims[0], high=z_lims[1], size=(1,))
+    x = max(min(x, 0.5 - vertices[:, 0].max()), -0.5 - vertices[:, 0].min())
+    y = max(min(y, 0.5 - vertices[:, 1].max()), -0.5 - vertices[:, 1].min())
+    z = max(min(z, 0.5 - vertices[:, 2].max()), -0.5 - vertices[:, 2].min())
+    vertices = np.stack([vertices[:, 0] + x, vertices[:, 1] + y, vertices[:, 2] + z], axis=-1)
+    return vertices
+
+def scale_vertices(vertices, x_lims=(0.75, 1.25), y_lims=(0.75, 1.25), z_lims=(0.75, 1.25)):
+    # scale x, y, z
+    x = np.random.uniform(low=x_lims[0], high=x_lims[1], size=(1,))
+    y = np.random.uniform(low=y_lims[0], high=y_lims[1], size=(1,))
+    z = np.random.uniform(low=z_lims[0], high=z_lims[1], size=(1,))
+    vertices = np.stack([vertices[:, 0] * x, vertices[:, 1] * y, vertices[:, 2] * z], axis=-1)
+    return vertices
+
+def normalize_vertices(vertices, eps: float = 1e-8):
+    bounds = np.array([vertices.min(axis=0), vertices.max(axis=0)])  # type: ignore
+    vertices = vertices - (bounds[0] + bounds[1])[None, :] / 2
+    # 数据中存在退化 mesh (bounds 范围为 0), 分母加 eps 避免除零产生 inf/NaN
+    scale = (bounds[1] - bounds[0]).max() + eps
+    vertices = vertices / scale
+    return vertices
+
+def agu_mesh(vertices: np.ndarray, shift_augment, scale_augment):
+    # 随机轴缩放
+    if scale_augment:
+        # lower, upper = scale_factor
+        # scale_axis = lower + np.random.rand(vertices.shape[-1]) * (upper - lower)
+        # vertices = vertices * scale_axis
+        vertices = scale_vertices(vertices)
+    vertices = normalize_vertices(vertices)
+    if shift_augment:
+        vertices = shift_vertices(vertices)
+    return vertices
+
+
+    
+def dec(mesh, target_face):
+    open3d_mesh = open3d.geometry.TriangleMesh(
+            vertices=open3d.utility.Vector3dVector(mesh.vertices),
+            triangles=open3d.utility.Vector3iVector(mesh.faces))
+    simple = open3d_mesh.simplify_quadric_decimation(target_face)
+    return trimesh.Trimesh(vertices=simple.vertices, faces=simple.triangles)
+
+def normalize(mesh, eps: float = 1e-8):
+    bounds = np.array([mesh.vertices.min(axis=0), mesh.vertices.max(axis=0)])
+    mesh.vertices = mesh.vertices - (bounds[0] + bounds[1])[None, :] / 2
+    # 数据中存在退化 mesh (bounds 范围为 0), 分母加 eps 避免除零产生 inf/NaN
+    scale = (bounds[1] - bounds[0]).max() + eps
+    mesh.vertices = 2 * (mesh.vertices / scale)
+    mesh.vertices = mesh.vertices.clip(-1.0, 1.0)  # 限制顶点坐标在[-1.0, 1.0]
+
+def merge_duplicate_vertices(mesh, tolerance=0):
+    """更可靠的顶点合并方法"""
+    vertices = mesh.vertices
+    faces = mesh.faces
+    
+    # 使用KD树查找相近的顶点
+    tree = cKDTree(vertices)
+    
+    # 查找在容差范围内的顶点对
+    pairs = tree.query_pairs(tolerance, output_type='ndarray')
+    
+    if len(pairs) == 0:
+        return mesh  # 没有需要合并的顶点
+    
+    # 创建顶点映射：将重复顶点映射到第一个顶点
+    vertex_map = {i: i for i in range(len(vertices))}
+    
+    for i, j in pairs:
+        # 总是映射到较小的索引
+        if i < j:
+            vertex_map[j] = i
+        else:
+            vertex_map[i] = j
+    
+    # 应用顶点映射
+    new_faces = []
+    for face in faces:
+        new_face = [vertex_map[v] for v in face]
+        new_faces.append(new_face)
+    
+    # 创建新的网格
+    new_mesh = trimesh.Trimesh(vertices=vertices, faces=new_faces)
+    
+    # 移除未引用的顶点
+    new_mesh.remove_unreferenced_vertices()
+    
+    return new_mesh
+
+def discrete_undiscrete_and_clean(mesh, n_discrete, clean=True):
+    vertices = discretize(mesh.vertices, num_discrete=n_discrete)
+    vertices = undiscretize(vertices, num_discrete=n_discrete)
+    dis_mesh = trimesh.Trimesh(vertices=vertices, faces=mesh.faces)
+    new_mesh = merge_duplicate_vertices(dis_mesh)
+    new_mesh.merge_vertices()
+    new_mesh.update_faces(new_mesh.nondegenerate_faces())
+    new_mesh.update_faces(new_mesh.unique_faces())
+    new_mesh.remove_unreferenced_vertices()
+    new_mesh.fix_normals()
+    return new_mesh
+
+# 点云加噪开关(环境变量,默认关=旧加噪;DataLoader worker 会继承 env)。
+#   PAPER_PC_NOISE=1 开启论文式加噪: 位置 σ_pos "up to" PC_NOISE_POS(逐样本采强度) + 法向 σ_normal=PC_NOISE_NORMAL
+#   (加噪后重归一化)。默认关 -> 走旧加噪(仅位置 +N(0,0.01)、50% 概率, 法向不加噪),其它训练零影响。
+#   幅度也可用 env PC_NOISE_POS / PC_NOISE_NORMAL 覆盖(默认 0.1 / 0.2)。只在 aug=True(训练)时生效,推理不加噪。
+def _paper_pc_noise_cfg():
+    return (os.environ.get("PAPER_PC_NOISE", "0") == "1",
+            float(os.environ.get("PC_NOISE_POS", "0.1")),
+            float(os.environ.get("PC_NOISE_NORMAL", "0.2")))
+
+
+def sample_pc(mesh, pc_num, total_pc_num = 50000, with_normal=True, aug=True, overfit=False):
+    if overfit:
+        np.random.seed(42)
+    if not with_normal:
+        points, _ = mesh.sample(pc_num, return_index=True)
+        return points
+
+    points, face_idx = mesh.sample(total_pc_num, return_index=True)
+    normals = mesh.face_normals[face_idx]
+    paper_noise, sigma_pos_max, sigma_normal = _paper_pc_noise_cfg()
+    if paper_noise:
+        # 论文式加噪: 位置 σ_pos<=sigma_pos_max + 法向 σ_normal(重归一化为单位向量)
+        if aug:
+            points = points.astype(np.float32)
+            normals = normals.astype(np.float32)
+            sigma_pos = np.random.uniform(0.0, sigma_pos_max)                  # "up to" -> 逐样本采噪声强度
+            points = points + np.random.randn(*points.shape).astype(np.float32) * sigma_pos
+            normals = normals + np.random.randn(*normals.shape).astype(np.float32) * sigma_normal
+            nlen = np.linalg.norm(normals, axis=-1, keepdims=True)
+            normals = normals / np.clip(nlen, 1e-8, None)                      # 加噪后法向重归一化为单位向量
+    else:
+        # 旧加噪(默认): 仅位置 +N(0,0.01)、50% 概率, 法向不加噪(保持既有训练行为不变)
+        if aug and random.random() < 0.5:
+            points = points + np.random.randn(*points.shape) * 0.01
+    pc_normal = np.concatenate([points, normals], axis=-1, dtype=np.float16)
+
+    # random sample point cloud
+    ind = np.random.choice(pc_normal.shape[0], pc_num, replace=False)
+    pc_normal = pc_normal[ind]
+    
+    return pc_normal
+
+def mul_encode_mesh(raw_mesh, gt_mesh, codec, tokenizer_target_depth, pc_num, total_pc_num, use_gt_pc, clean=True,
+                    normalize_input=True):
+    mesh = raw_mesh
+
+    # 归一化。normalize_input=False: 跳过对 raw_mesh 的归一化 —— 级联推理用(2026-08-22 修复):
+    # 级间 mesh 是上一级 tokenizer 解码产物, 天生在规范 [-1,1] 格坐标系; 按它自身 bbox 再归一化会在
+    # 预测带极值杂散面时把脚手架帧拉离 GT/点云帧(z 轴撕裂 bug, 见 infer_multilevel CASCADE_RENORM)。
+    # 训练路径 raw_mesh==gt_mesh(磁盘任意坐标系), 保持默认 True 不变。
+    if normalize_input:
+        normalize(mesh)
+    normalize(gt_mesh)
+    
+    if isinstance(tokenizer_target_depth, list):
+        selected_depth = random.choice(tokenizer_target_depth)
+    else:
+        selected_depth = tokenizer_target_depth
+    raw_dis_mesh = discrete_undiscrete_and_clean(gt_mesh, n_discrete=2**selected_depth, clean=clean)
+    gt_dis_mesh = discrete_undiscrete_and_clean(gt_mesh, n_discrete=2**(selected_depth+1), clean=clean)
+    # 采样点云
+    if use_gt_pc:
+        pc = sample_pc(gt_mesh, pc_num=pc_num, total_pc_num=total_pc_num, 
+                        with_normal=True, aug=False, overfit=False)
+    else:
+        pc = sample_pc(mesh, pc_num=pc_num, total_pc_num=total_pc_num, 
+                        with_normal=True, aug=False, overfit=False)
+    # raw_token, token = self.codec.encode_mesh(mesh, target_depth=self.tokenizer_target_depth)
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        
+        future = executor.submit(codec.encode_mesh, mesh, target_depth=selected_depth)
+        # 如果超过 2.0 秒没算完，这里会直接抛出 concurrent.futures.TimeoutError
+        raw_token, token, _ = future.result(timeout=20.0)
+        
+    n = len(token)
+
+    data_dict = {}
+    data_dict["token"] = torch.tensor(token, dtype=torch.int64)
+    data_dict["raw_token"] = torch.tensor(raw_token, dtype=torch.int64)
+    data_dict["raw_vertices"] = mesh.vertices
+    data_dict["raw_faces"] = mesh.faces
+    data_dict["raw_dis_mesh"] = raw_dis_mesh
+    data_dict["gt_dis_mesh"] = gt_dis_mesh
+
+    data_dict["start"] = 0
+    data_dict["end"] = len(data_dict["token"])
+    data_dict["pc"] = torch.from_numpy(pc)
+    data_dict["token_len"]  = torch.tensor([len(data_dict["token"])], dtype=torch.int64)
+
+    data_dict["selected_depth"] = selected_depth
+    return data_dict
+
+class MeshDataset(Dataset):
+    """
+    重构后的版本，初始化参数通过 opt 对象传入
+    """
+    def __init__(self, opt, data, codec, all_token_len, is_train):
+        # 将配置对象和核心数据存储到实例中
+        self.opt = opt
+        self.data = data
+        self.codec = codec
+        self.all_token_len = all_token_len
+        self.is_train = is_train
+        # 从 opt 对象中解包参数，保持原有代码逻辑不变
+        # 这样做的好处是 __getitem__ 等其他方法不需要修改
+        if is_train:
+            self.shift_augment = self.opt.data_processing.train_augmentation.shift_augment
+            self.scale_augment = self.opt.data_processing.train_augmentation.scale_augment
+            self.dec_augment = self.opt.data_processing.train_augmentation.dec_augment
+            self.rot_augment = self.opt.data_processing.train_augmentation.rot_augment
+            self.sample_augment = self.opt.data_processing.train_augmentation.sample_augment
+        else:
+            self.shift_augment = self.opt.data_processing.val_augmentation.shift_augment
+            self.scale_augment = self.opt.data_processing.val_augmentation.scale_augment
+            self.dec_augment = self.opt.data_processing.val_augmentation.dec_augment
+            self.rot_augment = self.opt.data_processing.val_augmentation.rot_augment
+            self.sample_augment = self.opt.data_processing.val_augmentation.sample_augment
+        self.n_discrete_size = self.opt.data_processing.n_discrete_size
+        self.vertex_order = self.opt.data_processing.vertex_order
+        self.window_size = self.opt.data_processing.windowing.window_size
+        self.window_stride = self.opt.data_processing.windowing.window_stride
+        self.split_method = self.opt.data_processing.windowing.split_method
+        self.max_length = self.opt.data_processing.max_len-20 # 假设max_length也在opt中
+        self.min_length = self.opt.data_processing.min_len # 假设max_length也在opt中
+        self.pc_num = self.opt.data_processing.pc_num
+        self.total_pc_num = self.opt.data_processing.total_pc_num
+        self.y_up = self.opt.data_processing.y_up
+        self.overfit = self.opt.data_processing.overfit if hasattr(self.opt.data_processing, 'overfit') else False
+        self.tokenizer_target_depth = self.opt.tokenizer.target_depth
+        # 2a 暴露偏差训练腐蚀(仅训练): 粗输入加伪面(过连接)+节点坐标抖动(位移),目标查 GT 不变。
+        # 仅 FastMeshCodec.encode_mesh 实现; 推理/eval 传 None -> 与原版逐位一致。
+        self.corrupt_cfg = None
+        _ct = getattr(getattr(self.opt, "training", None), "corrupt", None)
+        if is_train and _ct is not None and bool(getattr(_ct, "enabled", False)):
+            self.corrupt_cfg = {
+                "spurious_rate": float(getattr(_ct, "spurious_rate", 0.0)),
+                "spatial_radius": int(getattr(_ct, "spatial_radius", 2)),
+                "jitter_prob": float(getattr(_ct, "jitter_prob", 0.0)),
+                "jitter_max": int(getattr(_ct, "jitter_max", 1)),
+            }
+            print(f"  [2a Corrupt] ENABLED (is_train): {self.corrupt_cfg}")
+        (
+            self.image_cond_enabled,
+            self.image_cond_render_root,
+            self.image_cond_resolution,
+            self.image_cond_render_on_the_fly,
+            self.image_cond_render_cfg,
+            self.image_cond_render_manifest_path,
+        ) = _read_image_cond_cfg(opt)
+        # Default render parameters; can be overridden by `image_cond.render` in yaml.
+        # Defaults match render_normals_gpu.sh (yaw=315, pitch=30, FOV=40, auto_radius)
+        # but with init_radius=4.0 because our mesh frame is [-1, 1] (extent ~2) vs.
+        # the TRELLIS pre-render frame [-0.5, 0.5] (extent ~1).
+        rcfg = self.image_cond_render_cfg
+        self.render_yaw_deg = float(getattr(rcfg, "yaw_deg", 315.0)) if rcfg is not None else 315.0
+        self.render_pitch_deg = float(getattr(rcfg, "pitch_deg", 30.0)) if rcfg is not None else 30.0
+        self.render_fov_deg = float(getattr(rcfg, "fov_deg", 40.0)) if rcfg is not None else 40.0
+        self.render_init_radius = float(getattr(rcfg, "init_radius", 4.0)) if rcfg is not None else 4.0
+        self.render_auto_radius = bool(getattr(rcfg, "auto_radius", True)) if rcfg is not None else True
+        self.render_target_fill = float(getattr(rcfg, "target_fill", 0.9)) if rcfg is not None else 0.9
+        # If true (and is_train), randomly sample yaw in [0, 360) per call.
+        self.render_randomize_yaw = bool(getattr(rcfg, "randomize_yaw", False)) if rcfg is not None else False
+        print("\n" + "="*70)
+        print("="*70)
+        # 使用 f-string 格式化输出，将相关的配置项打印在同一行
+        print(f"  [Augmentation] Shift: {self.shift_augment}, Scale: {self.scale_augment}, Decimation: {self.dec_augment}, Rotation: {self.rot_augment}, Sample: {self.sample_augment}")
+        print(f"  [Data Specs]   Discretization Size: {self.n_discrete_size}, Vertex Order: '{self.vertex_order}'")
+        print(f"  [Windowing]    Size: {self.window_size}, Stride: {self.window_stride}, Method: '{self.split_method}'")
+        print(f"  [Sequence Len] Min: {self.min_length}, Max: {self.max_length}")
+        print(f"  [Point Cloud]  PC Num: {self.pc_num}, Total PC Num: {self.total_pc_num}")
+        print(f"  [Tokenizer]    Target Depth: {self.tokenizer_target_depth}")
+        if self.image_cond_enabled:
+            if self.image_cond_render_on_the_fly:
+                print(
+                    f"  [Image Cond]   Enabled | mode=render_on_the_fly | res={self.image_cond_resolution} "
+                    f"| yaw={self.render_yaw_deg} pitch={self.render_pitch_deg} fov={self.render_fov_deg} "
+                    f"radius_init={self.render_init_radius} auto_radius={self.render_auto_radius} "
+                    f"randomize_yaw={self.render_randomize_yaw}"
+                )
+            else:
+                print(f"  [Image Cond]   Enabled | mode=disk | render_root={self.image_cond_render_root} | res={self.image_cond_resolution}")
+        print("="*70 + "\n")
+        print(f"[MeshDataset] Created from {len(self.data)} entries")
+
+    def __len__(self):
+        return len(self.data)
+
+    def get_all_token_lenth(self):
+        return self.all_token_len
+
+    # __getitem__ 方法保持不变
+    def __getitem__(self, item):
+        iter_cnt = 0
+        if isinstance(item, tuple):
+            idx, start, end = item
+        else:
+            idx = item
+            start = 0
+            end = -1
+        while True:
+            file_path = self.data[idx]
+            iter_cnt += 1
+            if iter_cnt >= 4:
+                print(f'iter_cnt>=4: {iter_cnt}')
+            try:
+                mesh = trimesh.load(file_path, force="mesh")
+                # y轴朝上->z轴朝上
+                if self.y_up:
+                    mesh.vertices = mesh.vertices[:, [2, 0, 1]]
+
+                # 随机塌缩
+                if self.dec_augment:
+                    if mesh.faces.shape[0] >= self.min_length and random.random() < 0.5:
+                        target = np.random.randint(max(self.min_length, mesh.faces.shape[0] // 4), mesh.faces.shape[0])
+                        mesh = dec(mesh=mesh, target_face=target)
+                
+                # 随机拉伸和平移
+                if self.scale_augment or self.shift_augment:
+                    mesh.vertices = agu_mesh(mesh.vertices, self.shift_augment, self.scale_augment)
+                
+                # 随机旋转
+                if self.rot_augment:
+                    # azimuth = np.random.uniform(0, 360)
+                    azimuth = np.random.choice([0, 30, 60, 90, 120, 150, 180, 210, 240, 270, 300, 330], 1)[0]
+                    theta = np.radians(-azimuth)
+                    rotz = np.array([
+                        [np.cos(theta), -np.sin(theta), 0],
+                        [np.sin(theta), np.cos(theta), 0],
+                        [0, 0, 1]
+                    ])
+                    mesh.vertices = mesh.vertices @ rotz.T  # 绕z轴旋转
+                
+                # 归一化
+                normalize(mesh)
+
+                if isinstance(self.tokenizer_target_depth, list):
+                    selected_depth = random.choice(self.tokenizer_target_depth)
+                else:
+                    selected_depth = self.tokenizer_target_depth
+
+                # discrete_mesh = discrete_undiscrete_and_clean(mesh, n_discrete=2**selected_depth)
+
+                # # 采样点云
+                # pc = sample_pc(discrete_mesh, pc_num=self.pc_num, total_pc_num=self.total_pc_num,
+                #                with_normal=True, aug=self.sample_augment, overfit=self.overfit)
+                # 采样点云
+                pc = sample_pc(mesh, pc_num=self.pc_num, total_pc_num=self.total_pc_num,
+                               with_normal=True, aug=self.sample_augment, overfit=self.overfit)
+                if not np.isfinite(pc).all():
+                    raise ValueError("NaN/Inf in sampled point cloud")
+                # raw_token, token = self.codec.encode_mesh(mesh, target_depth=self.tokenizer_target_depth)
+                with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+
+                    future = executor.submit(self.codec.encode_mesh, mesh, target_depth=selected_depth,
+                                             corrupt=self.corrupt_cfg)
+                    # 超时会抛 concurrent.futures.TimeoutError(str() 为空)-> 换成带 depth/面数的明确消息
+                    try:
+                        raw_token, token, dis_mesh = future.result(timeout=10.0)
+                    except concurrent.futures.TimeoutError:
+                        raise TimeoutError(f"encode_mesh 超时(>10s, depth={selected_depth} res{2**(selected_depth+1)}, 输入面={len(mesh.faces)}) —— 多半粗分辨率过连接")
+
+                n = len(token)
+                if n < self.min_length:
+                    raise ValueError(f"Token length {n} is less than min_length {self.min_length}")
+
+                if n > self.max_length:
+                    raise ValueError(f"Token length {n} exceeds max_length {self.max_length}")
+
+                # Render the (augmented) mesh on-the-fly. Done inside the try
+                # block so any render failure resamples a different mesh.
+                rendered = None
+                if self.image_cond_enabled:
+                    if self.image_cond_render_on_the_fly:
+                        yaw_deg = (
+                            float(np.random.uniform(0.0, 360.0))
+                            if (self.render_randomize_yaw and self.is_train)
+                            else self.render_yaw_deg
+                        )
+                        rendered = _render_normals_image(
+                            np.asarray(mesh.vertices, dtype=np.float32),
+                            np.asarray(mesh.faces, dtype=np.int32),
+                            yaw_deg=yaw_deg,
+                            pitch_deg=self.render_pitch_deg,
+                            fov_deg=self.render_fov_deg,
+                            init_radius=self.render_init_radius,
+                            resolution=self.image_cond_resolution,
+                            auto_radius=self.render_auto_radius,
+                            target_fill=self.render_target_fill,
+                        )
+                    else:
+                        rendered = _load_image_and_camera(
+                            file_path,
+                            self.image_cond_render_root,
+                            self.image_cond_resolution,
+                        )
+
+                data_dict = {}
+                data_dict["token"] = torch.tensor(token, dtype=torch.int64)
+                data_dict["raw_token"] = torch.tensor(raw_token, dtype=torch.int64)
+                data_dict["raw_vertices"] = mesh.vertices
+                data_dict["raw_faces"] = mesh.faces
+                data_dict["dis_mesh"] = dis_mesh
+                del mesh
+                break
+            except Exception as e:
+                _reason = str(e) if str(e) else f"{type(e).__name__}(空消息)"
+                print(f'[dataiter wrong]: {file_path}, {_reason}')
+                idx = np.random.randint(0, len(self.data))
+                # start = 0
+                # end = -1
+
+        if len(data_dict["token"]) < self.window_size:
+            data_dict["start"] = 0
+            data_dict["end"] = len(data_dict["token"])
+        else:
+            if start == 0:
+                data_dict["start"] = 0
+                data_dict["end"] = self.window_size
+            # 随机截断
+            elif start == -1 and end == -1:
+                data_dict["start"] = -1
+                data_dict["end"] = -1
+            else:
+                data_dict["start"] = len(data_dict["token"]) - self.window_size
+                data_dict["end"] = len(data_dict["token"])
+        data_dict["pc"] = torch.from_numpy(pc)
+        data_dict["token_len"]  = torch.tensor([len(data_dict["token"])], dtype=torch.int64)
+        data_dict['name'] = Path(file_path).stem
+        data_dict["file_path"] = file_path
+        data_dict["selected_depth"] = selected_depth
+        if self.image_cond_enabled:
+            if self.image_cond_render_on_the_fly:
+                # Render-on-the-fly path still returns the legacy single-view tuple
+                # (no depth). Add a leading view-dim so downstream code is uniform.
+                image_tensor, transform_matrix, camera_angle_x = rendered
+                data_dict["image"] = image_tensor.unsqueeze(0)
+                data_dict["transform_matrix"] = transform_matrix.unsqueeze(0)
+                data_dict["camera_angle_x"] = camera_angle_x.unsqueeze(0)
+                # No depth available — emit dummies; the conditioner will skip
+                # occlusion testing when depth_valid is all False.
+                H = image_tensor.shape[-1]
+                data_dict["depth_map"] = torch.zeros(1, H, H, dtype=torch.float32)
+                data_dict["depth_valid"] = torch.zeros(1, H, H, dtype=torch.bool)
+            else:
+                (image_tensor, transform_matrix, camera_angle_x,
+                 depth_map, depth_valid) = rendered
+                data_dict["image"] = image_tensor
+                data_dict["transform_matrix"] = transform_matrix
+                data_dict["camera_angle_x"] = camera_angle_x
+                data_dict["depth_map"] = depth_map
+                data_dict["depth_valid"] = depth_valid
+        return data_dict
+
+    @classmethod
+    def load(cls, opt, path, codec, accelerator, is_train):
+        """
+        读取数据并使用 opt 对象来初始化 Dataset 实例。
+        """
+        all_token_len = []
+        data = []
+        max_length = opt.data_processing.max_face_num
+        min_len = opt.data_processing.min_face_num
+        (image_cond_enabled, render_root, _, render_on_the_fly,
+         _, render_manifest_path) = _read_image_cond_cfg(opt)
+        # Disk-render availability filter only applies when we read PNGs from
+        # disk; in render_on_the_fly mode the renderer makes its own image.
+        apply_disk_render_filter = image_cond_enabled and not render_on_the_fly
+        # Prefer the manifest when the disk-render filter is on. Path comes
+        # from `image_cond.render_manifest_path` in yaml; if absent or the
+        # file doesn't exist, fall back to per-line os.path.exists (slow).
+        render_manifest = None
+        if apply_disk_render_filter:
+            render_manifest = _load_render_manifest(render_manifest_path)
+            if accelerator.is_main_process:
+                if render_manifest is not None:
+                    print(f"[MeshDataset] Using render manifest ({len(render_manifest)} stems) "
+                          f"from {render_manifest_path}")
+                else:
+                    if render_manifest_path:
+                        print(f"[MeshDataset] WARNING: image_cond.render_manifest_path={render_manifest_path} "
+                              f"not found — falling back to per-line os.path.exists (slow). "
+                              f"Run script/build_render_manifest.py to fix.")
+                    else:
+                        print(f"[MeshDataset] WARNING: image_cond.render_manifest_path not set — "
+                              f"falling back to per-line os.path.exists (slow). "
+                              f"Run script/build_render_manifest.py and set the yaml key.")
+        dropped_no_render = 0
+        with open(path, 'r') as f:
+            for line in tqdm(f):
+                # 这段逻辑可能需要根据你的文件名格式调整
+                try:
+                    token_len = int(line.split("_f")[-1].split("_")[0])
+                except (IndexError, ValueError):
+                    try:
+                        token_len = int(line.split("_facenum_")[-1].split("_")[0])
+                    except (ValueError, IndexError):
+                        try:
+                            token_len = int(line.split("_facenum_")[-1].split(".")[0])
+                        except (ValueError, IndexError):
+                            # 如果文件名不包含token长度
+                            mesh = trimesh.load(line.strip(), force="mesh")
+                            token_len = mesh.faces.shape[0]
+                if token_len <= max_length and token_len>= min_len:
+                    file_path = line.strip()
+                    if apply_disk_render_filter:
+                        if render_manifest is not None:
+                            if Path(file_path).stem not in render_manifest:
+                                dropped_no_render += 1
+                                continue
+                        elif not _has_render(file_path, render_root):
+                            dropped_no_render += 1
+                            continue
+                    data.append(file_path)
+                    all_token_len.append(token_len)
+
+        if accelerator.is_main_process:
+            print(f"[MeshDataset] Loaded {len(data)} entries from {path}")
+            if apply_disk_render_filter:
+                print(
+                    f"[MeshDataset] Image-cond filter dropped {dropped_no_render} entries with no render under {render_root}"
+                )
+                if len(data) == 0:
+                    print(
+                        f"[MeshDataset] WARNING: no usable entries after image-cond filter — render root may be wrong."
+                    )
+            elif image_cond_enabled and render_on_the_fly:
+                print(f"[MeshDataset] Image-cond mode=render_on_the_fly — no disk-render filter applied.")
+
+        # 使用 opt 对象来创建类的实例
+        return cls(opt,
+                   data=data,
+                   codec=codec,
+                   all_token_len=all_token_len,
+                   is_train=is_train)
+

mulres_infer/dataset/mesh_dataset_altpc.py

@@ -0,0 +1,94 @@
+"""altpc 变体:pc_target 改用「编码→解码的 mesh」(octree token 解码的 GT@目标分辨率),
+而非 discrete_undiscrete 塌缩。理由:解码 mesh 与模型实际输出同一表示(含退化面结构),条件更一致。
+PCT_DEGEN=1 时, 额外在退化面(零面积, 即栏杆/薄结构边)上采点并注入 pc_target(法向置0)。
+直接 decode 父类已存的 d["token"](无需重编码)。
+"""
+import os
+import numpy as np
+import trimesh
+import torch
+
+from dataset.mesh_dataset import MeshDataset, sample_pc
+
+PCT_DEGEN = os.environ.get("PCT_DEGEN", "0") == "1"   # 是否在退化面上也采样注入
+_EPS = 1e-7
+
+
+def rail_edges(mesh):
+    """退化面(零面积)的有效栏杆边:取每个退化面最长边(对 (A,B,B),B-B 长 0,最长边即真实 A-B)。
+    返回 (A[D,3], B[D,3], 边长[D])。无退化面则 D=0。"""
+    af = np.asarray(mesh.area_faces)
+    deg = np.where(af < _EPS)[0]
+    if len(deg) == 0:
+        return np.zeros((0, 3)), np.zeros((0, 3)), np.zeros((0,))
+    v = np.asarray(mesh.vertices); f = np.asarray(mesh.faces)[deg]
+    ev = v[f]
+    pairs = [(0, 1), (1, 2), (0, 2)]
+    lens = np.stack([np.linalg.norm(ev[:, i] - ev[:, j], axis=1) for i, j in pairs], axis=1)  # [D,3]
+    pick = lens.argmax(1)
+    A = ev[np.arange(len(deg)), [pairs[p][0] for p in pick]]
+    B = ev[np.arange(len(deg)), [pairs[p][1] for p in pick]]
+    return A, B, lens[np.arange(len(deg)), pick]
+
+
+def sample_rail(A, B, lengths, n):
+    """沿栏杆边按边长加权采 n 个点(长边多采)。返回 [n,3]。"""
+    if n <= 0 or len(lengths) == 0 or lengths.sum() <= 0:
+        return np.zeros((0, 3))
+    counts = np.random.multinomial(n, lengths / lengths.sum())   # 每条边分到的点数, 和=n
+    idx = np.repeat(np.arange(len(lengths)), counts)             # [n] 每点所属边
+    t = np.random.rand(len(idx), 1).astype(np.float32)
+    return A[idx] + t * (B[idx] - A[idx])
+
+
+def degen_point_budget(mesh, pc_num, resolution):
+    """按「一致面密度」分配表面/栏杆点数, 不固定比例。
+    退化边当 1 体素宽薄条: 等效面积 = 边长 × 体素宽(mesh尺度/分辨率)。
+    ρ = pc_num/(表面积+栏杆等效面积); n_rail = round(ρ·栏杆等效面积)。
+    高分辨率体素小→栏杆相对少; 栏杆多的 mesh→栏杆点自然多。返回 (n_surf, n_rail, A, B, lengths)。"""
+    af = np.asarray(mesh.area_faces)
+    surf_area = float(af[af >= _EPS].sum())
+    A, B, lengths = rail_edges(mesh)
+    rail_len = float(lengths.sum())
+    if rail_len <= 0 or surf_area <= 0:
+        return pc_num, 0, A, B, lengths
+    voxel = float(np.max(mesh.extents)) / max(1, resolution)     # 1 体素的物理尺度
+    rail_eff_area = rail_len * voxel
+    rho = pc_num / (surf_area + rail_eff_area)
+    n_rail = int(round(rho * rail_eff_area))
+    n_rail = max(0, min(n_rail, pc_num))
+    return pc_num - n_rail, n_rail, A, B, lengths
+
+
+def sample_pc_with_degen(mesh, pc_num, total_pc_num, aug, overfit, resolution=64):
+    """表面采样 + 退化面边上采样, 按一致密度自适应分配(不固定比例)。栏杆点法向置 0。总数 = pc_num。"""
+    n_surf, n_rail, A, B, lengths = degen_point_budget(mesh, pc_num, resolution)
+    surf = sample_pc(mesh, pc_num=pc_num, total_pc_num=total_pc_num, with_normal=True, aug=aug, overfit=overfit)
+    if n_rail <= 0:
+        return surf
+    rail_xyz = sample_rail(A, B, lengths, n_rail)
+    if len(rail_xyz) == 0:
+        return surf
+    rail = np.concatenate([rail_xyz, np.zeros_like(rail_xyz)], axis=1).astype(surf.dtype)  # 法向 0
+    return np.concatenate([surf[:pc_num - len(rail)], rail], axis=0)
+
+
+class MeshDatasetAltPCDecoded(MeshDataset):
+    def __getitem__(self, item):
+        d = super().__getitem__(item)
+        sd = int(d["selected_depth"])
+        # 直接 decode 父类已存的 token -> octree 解码 GT@目标分辨率(与模型输出同表示, 含退化面)
+        try:
+            decoded = self.codec.decode_tokens(d["token"].tolist(), target_depth=sd)
+        except Exception:
+            decoded = None
+        if decoded is None or len(getattr(decoded, "faces", [])) < 4:
+            decoded = trimesh.Trimesh(vertices=d["raw_vertices"], faces=d["raw_faces"], process=False)
+        if PCT_DEGEN:
+            pc_t = sample_pc_with_degen(decoded, self.pc_num, self.total_pc_num, self.sample_augment,
+                                        self.overfit, resolution=2 ** (sd + 1))
+        else:
+            pc_t = sample_pc(decoded, pc_num=self.pc_num, total_pc_num=self.total_pc_num,
+                             with_normal=True, aug=self.sample_augment, overfit=self.overfit)
+        d["pc_target"] = torch.from_numpy(pc_t)
+        return d

mulres_infer/models/generator.py

@@ -0,0 +1,1507 @@
+"""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

mulres_infer/models/point_conditioner.py

@@ -0,0 +1,200 @@
+import torch
+from torch import nn
+from beartype import beartype
+from modules.miche.encode import load_model
+import os
+
+_PKG_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+# helper functions
+
+def exists(val):
+    return val is not None
+
+def default(*values):
+    for value in values:
+        if exists(value):
+            return value
+    return None
+
+
+# point-cloud encoder from Michelangelo
+@beartype
+class PointConditioner(torch.nn.Module):
+    def __init__(
+        self,
+        *,
+        dim_latent = None,
+        model_name = 'miche-256-feature',
+        cond_dim = 768,
+        feature_dim = 768,
+        freeze = True,
+        img_feat_dim: int = 0,
+        img_feat_extractor=None,
+        num_latents = None,   # None=预训练默认(256); >256 时把可学习 query 平铺+微噪扩容(预训练权重原样热启动)
+    ):
+        super().__init__()
+
+        # open-source version of miche
+        if model_name == 'miche-256-feature':
+            ckpt_path = os.environ.get("SHAPEVAE_CKPT",
+                os.path.join(_PKG_ROOT, "ckpt", "shapevae-256.ckpt"))
+            if not os.path.exists(ckpt_path):
+                ckpt_path=None
+                print('[WARNING] Michelangelo ckpt not exist, please check if you are training')
+            config_path = os.path.join(_PKG_ROOT, 'modules', 'miche', 'shapevae-256.yaml')
+
+            self.feature_dim = feature_dim    # embedding dimension
+            self.cond_length = 257     # length of embedding
+            self.point_encoder = load_model(ckpt_path=ckpt_path, config_path=config_path)
+            if num_latents:
+                self._expand_latents(int(num_latents))
+
+            # additional layers to connect miche and GPT
+            self.cond_head_proj = nn.Linear(cond_dim, self.feature_dim)
+            self.cond_proj = nn.Linear(cond_dim, self.feature_dim)
+
+        else:
+            raise NotImplementedError
+
+        # whether to finetuen point-cloud encoder
+        if freeze:
+            for parameter in self.point_encoder.parameters():
+                parameter.requires_grad = False
+
+        self.freeze = freeze
+        self.model_name = model_name
+        self.dim_latent = self.feature_dim
+
+        # ---------- image-conditioning extension ----------
+        self.img_feat_dim = int(img_feat_dim or 0)
+        self.img_feat_extractor = img_feat_extractor
+        if self.img_feat_dim > 0:
+            self._widen_input_proj(self.img_feat_dim)
+
+        self.register_buffer('_device_param', torch.tensor(0.), persistent = False)
+
+
+    @property
+    def device(self):
+        return next(self.buffers()).device
+
+    def _expand_latents(self, num_latents: int):
+        """把 perceiver 可学习 query 扩到 [1头 + num_latents body]: 预训练 query 是 257=1+256(行0=头token),
+        头保持 1 个, body 256 平铺复制 + 1% std 微噪打破对称(首份无噪)。attention/n_ctx 均动态读 shape,
+        无其他结构依赖(SDPA 补丁同样形状无关)。训练热启动时上层 ckpt(若已含扩容 query)会原样覆盖。"""
+        enc = self.point_encoder.model.shape_model.encoder
+        old_q = enc.query.data                          # [1+B, w]
+        body = old_q[1:]
+        if body.shape[0] == num_latents:
+            self.cond_length = 1 + num_latents
+            return
+        assert num_latents % body.shape[0] == 0, f"num_latents={num_latents} 须为 body={body.shape[0]} 的整数倍"
+        reps = num_latents // body.shape[0]
+        new_body = body.repeat(reps, 1)
+        noise = torch.randn_like(new_body) * body.std() * 0.01
+        noise[: body.shape[0]] = 0                      # 首份 body 保持原 query 精确不动
+        enc.query = nn.Parameter(torch.cat([old_q[:1], new_body + noise], dim=0))
+        enc.num_latents = 1 + num_latents
+        self.cond_length = 1 + num_latents
+        print(f"[PointConditioner] query 扩容 1+{body.shape[0]} -> 1+{num_latents} (body tile+1%噪, 首份无噪)")
+
+    def _widen_input_proj(self, img_feat_dim: int):
+        """Replace Michelangelo's input_proj Linear with a wider one and zero-init
+        the new input columns so loading the pretrained ckpt doesn't change init
+        behavior (image features get multiplied by zero on construction)."""
+        enc = self.point_encoder.model.shape_model.encoder
+        if getattr(enc, "img_feat_dim", 0) == img_feat_dim and enc.input_proj.in_features == (
+            enc.fourier_embedder.out_dim + enc.point_feats + img_feat_dim
+        ):
+            return  # already widened (e.g. resume from a widened ckpt)
+        old = enc.input_proj
+        new_in = enc.fourier_embedder.out_dim + enc.point_feats + img_feat_dim
+        if new_in == old.in_features:
+            enc.img_feat_dim = img_feat_dim
+            return
+        new_proj = nn.Linear(
+            new_in,
+            old.out_features,
+            bias=old.bias is not None,
+            device=old.weight.device,
+            dtype=old.weight.dtype,
+        )
+        with torch.no_grad():
+            new_proj.weight.zero_()
+            new_proj.weight[:, : old.in_features].copy_(old.weight)
+            if old.bias is not None:
+                new_proj.bias.copy_(old.bias)
+        enc.input_proj = new_proj
+        enc.img_feat_dim = img_feat_dim
+        if self.freeze:
+            # input_proj is part of the shape encoder, which we keep frozen by
+            # default; only the new image-feature columns ever produce nonzero
+            # gradients on top of frozen weights. If callers want to fine-tune
+            # the whole input_proj they can flip freeze=False.
+            for p in new_proj.parameters():
+                p.requires_grad = False
+
+
+    def embed_pc(
+        self,
+        pc_normal,
+        image=None,
+        transform_matrix=None,
+        camera_angle_x=None,
+        depth_map=None,
+        depth_valid=None,
+        img_feats=None,
+    ):
+        # encode point cloud to embeddings
+        if self.model_name == 'miche-256-feature':
+            if self.img_feat_dim > 0:
+                if img_feats is None:
+                    if self.img_feat_extractor is None:
+                        raise RuntimeError(
+                            "img_feat_dim>0 but no img_feats provided and no extractor configured."
+                        )
+                    if image is None or transform_matrix is None or camera_angle_x is None:
+                        raise ValueError(
+                            "image, transform_matrix, camera_angle_x are required when img_feat_extractor is set."
+                        )
+                    img_feats = self.img_feat_extractor(
+                        image, pc_normal[..., :3], transform_matrix, camera_angle_x,
+                        depth_map=depth_map, depth_valid=depth_valid,
+                    )
+                img_feats = img_feats.to(dtype=pc_normal.dtype)
+                point_feature = self.point_encoder.encode_latents(pc_normal, img_feats=img_feats)
+            else:
+                point_feature = self.point_encoder.encode_latents(pc_normal)
+            pc_embed_head = self.cond_head_proj(point_feature[:, 0:1])
+            pc_embed = self.cond_proj(point_feature[:, 1:])
+            pc_embed = torch.cat([pc_embed_head, pc_embed], dim=1)
+
+        return pc_embed
+
+
+    def forward(
+        self,
+        pc = None,
+        pc_embeds = None,
+        image=None,
+        transform_matrix=None,
+        camera_angle_x=None,
+        depth_map=None,
+        depth_valid=None,
+        img_feats=None,
+    ):
+        if pc_embeds is None:
+            target_dtype = next(self.buffers()).dtype
+            pc_embeds = self.embed_pc(
+                pc.to(target_dtype),
+                image=image,
+                transform_matrix=transform_matrix,
+                camera_angle_x=camera_angle_x,
+                depth_map=depth_map,
+                depth_valid=depth_valid,
+                img_feats=img_feats,
+            )
+
+        assert not torch.any(torch.isnan(pc_embeds)), 'NAN values in pc embedings'
+
+        return pc_embeds

mulres_infer/models/subdivider.py

@@ -0,0 +1,1188 @@
+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
+
+

mulres_infer/modules/attention.py

@@ -0,0 +1,376 @@
+"""注意力与 Transformer 原语(从 pure_iflame_with_condition 提取, 只保留 MLA 及其依赖)。
+
+MLA 是细分器 subdivider.DiscreteDiffusionTransformer 使用的多头潜在注意力块。
+原目录里的 MeshRes / MLCA / LinearDownsample / LinearUpsample / MultiheadCrossFlashrope
+都不在推理路径上, 已删除。
+"""
+
+import torch
+import torch.nn as nn
+import math
+import torch.nn.functional as F
+from torch import nn
+import time #cal_time new at 10.21
+import numpy as np
+from typing import List, Optional, Tuple
+from tqdm import tqdm
+from einops import rearrange, repeat, pack
+
+
+# ---- 原 rms_norm.py ----
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine=True, memory_efficient=False):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim))
+        else:
+            self.register_parameter('weight', None)
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        if self.weight is not None:
+            output = output * self.weight
+        return output
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'
+
+# ---- 原 multihead_flashdiff_2.py ----
+
+class Rotary(nn.Module):
+    def __init__(self, dim: int, max_seq_len=9*50000, dtype=torch.float16):
+        super().__init__()
+        # half-truncate RoPE by @YouJiacheng (w/ base freq tuning)
+        angular_freq = (1 / 1024) ** torch.linspace(0, 1, steps=dim//4, dtype=dtype)
+        angular_freq = torch.cat([angular_freq, angular_freq.new_zeros(dim//4)])
+        t = torch.arange(max_seq_len, dtype=dtype)
+        theta = torch.einsum("i,j -> ij", t, angular_freq)
+        self.register_buffer('cos', theta.cos(), persistent=False)
+        self.register_buffer('sin', theta.sin(), persistent=False)
+        self.max_seq_len = max_seq_len
+        self.dtype = dtype
+
+    def forward(self, x_BTHD, offset=0):
+        """
+        Args:
+            x_BTHD: Input tensor of shape [batch, seq_len, num_heads, dim]
+            offset: Position offset for the sequence
+        """
+        batch_size, seq_len = x_BTHD.size(0), x_BTHD.size(-3)
+        assert torch.max(offset + seq_len) <= self.max_seq_len, f"Sequence length {offset + seq_len} exceeds maximum length {self.max_seq_len}"
+
+        # Get the appropriate slice of cos/sin based on offset
+        positions = torch.arange(seq_len, device=x_BTHD.device).expand(batch_size, -1) + offset.unsqueeze(-1)
+        cos = self.cos[positions].unsqueeze(-2)
+        sin = self.sin[positions].unsqueeze(-2)
+
+        x1, x2 = x_BTHD.to(dtype=self.dtype).chunk(2, dim=-1)
+        y1 = x1 * cos - x2 * sin
+        y2 = x1 * sin + x2 * cos
+        return torch.cat((y1, y2), dim=-1)
+
+class MultiheadFlashrope(nn.Module):
+    def __init__(self, args, embed_dim, num_heads):
+        super().__init__()
+        self.args = args
+        kvcache_method, kv_cache_window_size, window_stride = args
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads  
+        self.num_kv_heads = num_heads
+        self.head_dim = embed_dim // num_heads 
+        self.scaling = self.head_dim ** -0.5
+        self.kv_cache_window_size = kv_cache_window_size
+        self.window_stride = window_stride
+        self.kvcache_method = kvcache_method
+        # KV cache related attributes
+        self.kv_cache_enabled = False
+        self.k_cache = None
+        self.v_cache = None
+        self.cache_pos = 0  # Track position in cache
+        
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.rotary = Rotary(self.head_dim)
+        self.sliding_window_cur_len = 0
+        # self.attention_scores = None
+
+    def forward(self, x, start_pos, context=None, attention_mask=None, causal=False):
+        """
+        Args:
+            x: [batch_size, tgt_len, embed_dim]
+            start_pos: [batch_size]
+            context: Optional context for cross-attention [batch_size, src_len, embed_dim]
+        """
+        bsz, tgt_len, embed_dim = x.size()
+        # Project inputs to q, k, v
+        q = self.q_proj(x)
+        k = self.k_proj(context if context is not None else x)
+        v = self.v_proj(context if context is not None else x)
+
+        # Reshape to [batch_size, seq_len, num_heads, head_dim]
+        q = q.view(bsz, -1, self.num_heads, self.head_dim)
+        k = k.view(bsz, -1, self.num_heads, self.head_dim)
+        v = v.view(bsz, -1, self.num_heads, self.head_dim)
+        q = self.rotary(q, offset=start_pos)
+        # Normal operation without cache
+        k = self.rotary(k, offset=start_pos)
+        if attention_mask is not None:
+            attn_output = self.attention_pytorch_optimized(q, k, v, mask=attention_mask)
+        else:
+            attn_output = self.attention_pytorch_optimized(q, k, v, causal=causal)
+
+        # Reshape output
+        attn_output = attn_output.contiguous().view(bsz, tgt_len, embed_dim)
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output
+    
+    def attention_pytorch_optimized(self, q, k, v, mask=None, causal=False):
+        b, seq_q, head_num, c = q.shape
+        seq_k = k.shape[1]
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+        if mask is not None:
+            atten_mask_npu = mask.unsqueeze(1).expand(b, head_num, seq_q, seq_k)
+            atten_mask_npu = torch.logical_not(atten_mask_npu)
+            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=atten_mask_npu)
+        else:
+            if causal:
+                attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+            else:
+                attn_output = F.scaled_dot_product_attention(q, k, v)
+
+        return attn_output.permute(0, 2, 1, 3)
+
+    def attention_pytorch(self, q, k, v, causal=False):
+        # [batch_size, num_heads, seq_len, head_dim]
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scaling
+
+        if causal:
+            seq_len = q.size(-2)
+            causal_mask = torch.triu(
+                torch.ones(seq_len, seq_len, dtype=torch.bool, device=q.device),
+                diagonal=1
+            )
+            attn_weights.masked_fill_(causal_mask, float('-inf'))
+  
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        
+  
+        attn_output = torch.matmul(attn_weights, v)
+        
+
+        return attn_output.permute(0, 2, 1, 3)
+
+    def empty_kv_cache(self, batch_size: int, dtype: torch.dtype):
+        """Initialize empty KV cache"""
+        self.kv_cache_enabled=True
+
+        device = next(self.parameters()).device
+        kv_cache_maxlen = self.kv_cache_window_size
+        # Initialize empty caches
+        k_cache = torch.zeros(
+            batch_size, 
+            kv_cache_maxlen, 
+            self.num_heads,
+            self.head_dim,
+            dtype=dtype,
+            device=device
+        )
+        self.k_cache=k_cache
+
+        v_cache = torch.zeros(
+            batch_size,
+            kv_cache_maxlen,
+            self.num_heads,
+            self.head_dim,
+            dtype=dtype,
+            device=device
+        )       
+        self.v_cache=v_cache
+        self.cache_pos = 0  # Reset cache position
+
+        element_size = {
+            torch.float32: 4,
+            torch.float16: 2,
+            torch.bfloat16: 2,
+        }[dtype]
+        k_cache_size = (batch_size * 
+                    kv_cache_maxlen * 
+                    self.num_heads * 
+                    self.head_dim * 
+                    element_size)
+        v_cache_size = k_cache_size 
+        total_cache_size = k_cache_size + v_cache_size
+        cache_size_gb = total_cache_size / (1024**3)
+        return cache_size_gb
+        
+
+    def reset_cache(self):
+        """Reset KV cache state"""
+        self.k_cache = None
+        self.v_cache = None
+        self.cache_pos = 0
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim, context_dim, n_heads, dropout=0.0):
+        super().__init__()
+        self.n_heads = n_heads
+        self.scale = (dim // n_heads) ** -0.5
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.kv_proj = nn.Linear(context_dim, 2 * dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, context):
+        # x: [batch, seq_len, dim], context: [batch, context_len, context_dim]
+        B, N, C = x.shape
+        _, M, _ = context.shape
+        H = self.n_heads
+
+        # Linear projections
+        q = self.q_proj(x).view(B, N, H, C // H)  # [B, H, seq_len, dim//H]
+        k, v = self.kv_proj(context).chunk(2, dim=-1)
+        k = k.view(B, M, H, C // H)  # [B, context_len, H, dim//H]
+        v = v.view(B, M, H, C // H)  # [B, context_len, H, dim//H]
+        
+        attn_output = self.attention_pytorch_optimized(q,k,v)
+        attn_output = attn_output.contiguous().view(B, N, C)
+        
+        # q = self.q_proj(x).view(B, N, H, C // H).transpose(1, 2)  # [B, H, seq_len, dim//H]
+        # k, v = self.kv_proj(context).chunk(2, dim=-1)
+        # k = k.view(B, M, H, C // H).transpose(1, 2)  # [B, H, context_len, dim//H]
+        # v = v.view(B, M, H, C // H).transpose(1, 2)  # [B, H, context_len, dim//H]
+        # # Scaled dot-product attention
+        # attn_weights = (q @ k.transpose(-2, -1)) * self.scale  # [B, H, seq_len, context_len]
+        # attn_weights = attn_weights.softmax(dim=-1)
+        # attn_output = (attn_weights @ v).transpose(1, 2).contiguous().view(B, N, C)  # [B, seq_len, dim]
+        out = self.out_proj(attn_output)
+        out = self.dropout(out)
+
+        return out
+    
+    def attention_pytorch_optimized(self, q, k, v, mask=None, causal=False):
+        b, seq_q, head_num, c = q.shape
+        seq_k = k.shape[1]
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+        if mask is not None:
+            atten_mask_npu = mask.unsqueeze(1).expand(b, head_num, seq_q, seq_k)
+            atten_mask_npu = torch.logical_not(atten_mask_npu)
+            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=atten_mask_npu)
+        else:
+            if causal:
+                attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+            else:
+                attn_output = F.scaled_dot_product_attention(q, k, v)
+
+        return attn_output.permute(0, 2, 1, 3)
+
+# ---- 原 transformer.py ----
+
+class SwiGLU(nn.Module):
+    def __init__(self, embed_dim):
+        super(SwiGLU, self).__init__()
+        self.proj1 = nn.Linear(embed_dim, int(8 / 3 * embed_dim))  # XWG
+        self.proj2 = nn.Linear(embed_dim, int(8 / 3 * embed_dim))  # XW1
+        self.proj3 = nn.Linear(int(8 / 3 * embed_dim), embed_dim)  # W2
+
+    def forward(self, x):
+        # Apply Swish (SiLU in PyTorch) to the first projection
+        x_proj1 = F.silu(self.proj1(x))  # Swish(XWG)
+        # Apply the second projection
+        x_proj2 = self.proj2(x)  # XW1
+        # Element-wise multiplication
+        x_glu = x_proj1 * x_proj2  # (swish(XWG) ⊙ XW1)
+        # Final projection back to the original dimension
+        output = self.proj3(x_glu)  # (swish(XWG) ⊙ XW1)W2
+        
+        return output
+
+class MLA(nn.Module):
+    def __init__(self, embed_dim, num_heads, args, cross_embed_dim, causal=True, dropout=0.):
+        """
+        Differential Transformer Block with optional causal self-attention or cross-attention
+        and automatic application of ROPE.
+
+        Args:
+            embed_dim (int): Embedding dimension.
+            num_heads (int): Number of attention heads.
+            args (Namespace): Arguments for attention settings.
+            causal (bool): Whether to use causal masking by default for self-attention.
+        """
+        super(MLA, self).__init__()
+        # Multihead differential attention module
+        self.attn = MultiheadFlashrope(args, embed_dim, num_heads)
+        self.causal = causal
+
+        self.feed_forward = SwiGLU(embed_dim)
+        self.norm1 = RMSNorm(embed_dim, eps=1e-5)
+        self.norm2 = RMSNorm(embed_dim, eps=1e-5)
+        self.attn_dropout = nn.Dropout(dropout)
+        self.ff_dropout = nn.Dropout(dropout)
+        self.residual_dropout = nn.Dropout(dropout) 
+        if cross_embed_dim != -1:
+            self.cross_norm = RMSNorm(embed_dim, eps=1e-5)
+            self.cross_attn = CrossAttention(embed_dim, cross_embed_dim, num_heads)
+
+    def forward(
+        self, x, start, conds=None, attention_mask=None, use_cache=False
+    ):
+        """
+        Args:
+            x: Input tensor
+            use_cache: Whether to use KV cache
+        """
+        # Enable/disable KV cache in attention module
+        self.attn.kv_cache_enabled = use_cache
+
+        # 第一个残差连接：注意力 + dropout
+        residual = x
+        x = self.norm1(x)
+        attn_out = self.attn(
+            x,
+            start_pos=start,
+            attention_mask=attention_mask,
+            causal=self.causal,
+        )
+        attn_out = self.attn_dropout(attn_out)
+        x = residual + attn_out
+        if conds is not None:
+            x_skip = x
+            x = self.cross_norm(x)
+            x = self.cross_attn(x, conds) + x_skip 
+
+        # 第二个残差连接：前馈网络 + dropout
+        residual = x
+        ff_out = self.ff_dropout((self.feed_forward(self.norm2(x))))
+        x = residual + ff_out
+
+        return x
+
+    def init_kv_cache(self, batch_size, dtype=torch.bfloat16):
+        """Initialize KV cache for this transformer block"""
+        return self.attn.empty_kv_cache(
+            batch_size=batch_size, dtype=dtype
+        )
+
+    def reset_kv_cache(self):
+        """Reset KV cache for this transformer block"""
+        self.attn.reset_cache()

mulres_infer/modules/fast_octree_level.py

@@ -0,0 +1,52 @@
+"""数组化 + Numba 的「单层」八叉树编码,等价替换
+  MatrixOctree(max_depth=target_depth+1).build + OctreeFaceTokenBuilder(mc=8).encode_level_tokens(target_depth)
+
+前提(per-row 栈已满足):max_depth = target_depth+1 -> 只有一次「叶子(td+1)->节点(td)」聚合,
+不需要多层 neighbor 传播;且 max_children=8 时 reduce_node_children_to_max 是 no-op(活跃 octant ≤ 8)。
+
+去掉 OctreeNode 对象 / dict-of-tuple / set-of-tuple:
+  * 叶子坐标量化 -> 边表(numpy) -> 节点坐标(leaf>>1)+ octant 局部索引(位运算, 向量化)
+  * 按节点聚合 intra 8×8 / 活跃掩码 F / 节点-节点邻接(向量化 scatter / np.unique)
+  * 三角形(3-clique)枚举走 Numba njit(CSR 两指针求交), 这是原 get_level_info 的热点
+  * inter 8×8 按 (节点,邻居) 分组填充
+输出 (raw_t, out_t) 与原实现「逐 slot 解码后网格一致」(顺序可不同;decode 与顺序无关)。
+"""
+import numpy as np
+from numba import njit
+
+from utils.data_utils import discretize
+
+_OCT_W = np.array([1, 2, 4], dtype=np.int64)   # octant 局部索引 = x + 2y + 4z
+
+
+@njit(cache=False)
+def _triangles_csr(indptr, indices):
+    """前向邻接 CSR(indices[indptr[u]:indptr[u+1]] = 与 u 相连且 >u 的节点, 升序)。
+    枚举 u<v<w 三角形, 返回 [[u,v,w],...]。两指针求 adj[u]∩adj[v]。"""
+    out_u = []
+    out_v = []
+    out_w = []
+    n = indptr.shape[0] - 1
+    for u in range(n):
+        su, eu = indptr[u], indptr[u + 1]
+        for pa in range(su, eu):
+            v = indices[pa]
+            sv, ev = indptr[v], indptr[v + 1]
+            # 交集: adj[u] 中 >v 的部分(即 pa+1..eu, 因为升序且后面都 >v) 与 adj[v](sv..ev)
+            i = pa + 1
+            j = sv
+            while i < eu and j < ev:
+                wi = indices[i]
+                wj = indices[j]
+                if wi == wj:
+                    out_u.append(u); out_v.append(v); out_w.append(wi)
+                    i += 1; j += 1
+                elif wi < wj:
+                    i += 1
+                else:
+                    j += 1
+    return (np.array(out_u, dtype=np.int64),
+            np.array(out_v, dtype=np.int64),
+            np.array(out_w, dtype=np.int64))
+
+

mulres_infer/modules/miche/LICENSE

@@ -0,0 +1,674 @@
+                    GNU GENERAL PUBLIC LICENSE
+                       Version 3, 29 June 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+                            Preamble
+
+  The GNU General Public License is a free, copyleft license for
+software and other kinds of works.
+
+  The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works.  By contrast,
+the GNU General Public License is intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users.  We, the Free Software Foundation, use the
+GNU General Public License for most of our software; it applies also to
+any other work released this way by its authors.  You can apply it to
+your programs, too.
+
+  When we speak of free software, we are referring to freedom, not
+price.  Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+  To protect your rights, we need to prevent others from denying you
+these rights or asking you to surrender the rights.  Therefore, you have
+certain responsibilities if you distribute copies of the software, or if
+you modify it: responsibilities to respect the freedom of others.
+
+  For example, if you distribute copies of such a program, whether
+gratis or for a fee, you must pass on to the recipients the same
+freedoms that you received.  You must make sure that they, too, receive
+or can get the source code.  And you must show them these terms so they
+know their rights.
+
+  Developers that use the GNU GPL protect your rights with two steps:
+(1) assert copyright on the software, and (2) offer you this License
+giving you legal permission to copy, distribute and/or modify it.
+
+  For the developers' and authors' protection, the GPL clearly explains
+that there is no warranty for this free software.  For both users' and
+authors' sake, the GPL requires that modified versions be marked as
+changed, so that their problems will not be attributed erroneously to
+authors of previous versions.
+
+  Some devices are designed to deny users access to install or run
+modified versions of the software inside them, although the manufacturer
+can do so.  This is fundamentally incompatible with the aim of
+protecting users' freedom to change the software.  The systematic
+pattern of such abuse occurs in the area of products for individuals to
+use, which is precisely where it is most unacceptable.  Therefore, we
+have designed this version of the GPL to prohibit the practice for those
+products.  If such problems arise substantially in other domains, we
+stand ready to extend this provision to those domains in future versions
+of the GPL, as needed to protect the freedom of users.
+
+  Finally, every program is threatened constantly by software patents.
+States should not allow patents to restrict development and use of
+software on general-purpose computers, but in those that do, we wish to
+avoid the special danger that patents applied to a free program could
+make it effectively proprietary.  To prevent this, the GPL assures that
+patents cannot be used to render the program non-free.
+
+  The precise terms and conditions for copying, distribution and
+modification follow.
+
+                       TERMS AND CONDITIONS
+
+  0. Definitions.
+
+  "This License" refers to version 3 of the GNU General Public License.
+
+  "Copyright" also means copyright-like laws that apply to other kinds of
+works, such as semiconductor masks.
+
+  "The Program" refers to any copyrightable work licensed under this
+License.  Each licensee is addressed as "you".  "Licensees" and
+"recipients" may be individuals or organizations.
+
+  To "modify" a work means to copy from or adapt all or part of the work
+in a fashion requiring copyright permission, other than the making of an
+exact copy.  The resulting work is called a "modified version" of the
+earlier work or a work "based on" the earlier work.
+
+  A "covered work" means either the unmodified Program or a work based
+on the Program.
+
+  To "propagate" a work means to do anything with it that, without
+permission, would make you directly or secondarily liable for
+infringement under applicable copyright law, except executing it on a
+computer or modifying a private copy.  Propagation includes copying,
+distribution (with or without modification), making available to the
+public, and in some countries other activities as well.
+
+  To "convey" a work means any kind of propagation that enables other
+parties to make or receive copies.  Mere interaction with a user through
+a computer network, with no transfer of a copy, is not conveying.
+
+  An interactive user interface displays "Appropriate Legal Notices"
+to the extent that it includes a convenient and prominently visible
+feature that (1) displays an appropriate copyright notice, and (2)
+tells the user that there is no warranty for the work (except to the
+extent that warranties are provided), that licensees may convey the
+work under this License, and how to view a copy of this License.  If
+the interface presents a list of user commands or options, such as a
+menu, a prominent item in the list meets this criterion.
+
+  1. Source Code.
+
+  The "source code" for a work means the preferred form of the work
+for making modifications to it.  "Object code" means any non-source
+form of a work.
+
+  A "Standard Interface" means an interface that either is an official
+standard defined by a recognized standards body, or, in the case of
+interfaces specified for a particular programming language, one that
+is widely used among developers working in that language.
+
+  The "System Libraries" of an executable work include anything, other
+than the work as a whole, that (a) is included in the normal form of
+packaging a Major Component, but which is not part of that Major
+Component, and (b) serves only to enable use of the work with that
+Major Component, or to implement a Standard Interface for which an
+implementation is available to the public in source code form.  A
+"Major Component", in this context, means a major essential component
+(kernel, window system, and so on) of the specific operating system
+(if any) on which the executable work runs, or a compiler used to
+produce the work, or an object code interpreter used to run it.
+
+  The "Corresponding Source" for a work in object code form means all
+the source code needed to generate, install, and (for an executable
+work) run the object code and to modify the work, including scripts to
+control those activities.  However, it does not include the work's
+System Libraries, or general-purpose tools or generally available free
+programs which are used unmodified in performing those activities but
+which are not part of the work.  For example, Corresponding Source
+includes interface definition files associated with source files for
+the work, and the source code for shared libraries and dynamically
+linked subprograms that the work is specifically designed to require,
+such as by intimate data communication or control flow between those
+subprograms and other parts of the work.
+
+  The Corresponding Source need not include anything that users
+can regenerate automatically from other parts of the Corresponding
+Source.
+
+  The Corresponding Source for a work in source code form is that
+same work.
+
+  2. Basic Permissions.
+
+  All rights granted under this License are granted for the term of
+copyright on the Program, and are irrevocable provided the stated
+conditions are met.  This License explicitly affirms your unlimited
+permission to run the unmodified Program.  The output from running a
+covered work is covered by this License only if the output, given its
+content, constitutes a covered work.  This License acknowledges your
+rights of fair use or other equivalent, as provided by copyright law.
+
+  You may make, run and propagate covered works that you do not
+convey, without conditions so long as your license otherwise remains
+in force.  You may convey covered works to others for the sole purpose
+of having them make modifications exclusively for you, or provide you
+with facilities for running those works, provided that you comply with
+the terms of this License in conveying all material for which you do
+not control copyright.  Those thus making or running the covered works
+for you must do so exclusively on your behalf, under your direction
+and control, on terms that prohibit them from making any copies of
+your copyrighted material outside their relationship with you.
+
+  Conveying under any other circumstances is permitted solely under
+the conditions stated below.  Sublicensing is not allowed; section 10
+makes it unnecessary.
+
+  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
+  No covered work shall be deemed part of an effective technological
+measure under any applicable law fulfilling obligations under article
+11 of the WIPO copyright treaty adopted on 20 December 1996, or
+similar laws prohibiting or restricting circumvention of such
+measures.
+
+  When you convey a covered work, you waive any legal power to forbid
+circumvention of technological measures to the extent such circumvention
+is effected by exercising rights under this License with respect to
+the covered work, and you disclaim any intention to limit operation or
+modification of the work as a means of enforcing, against the work's
+users, your or third parties' legal rights to forbid circumvention of
+technological measures.
+
+  4. Conveying Verbatim Copies.
+
+  You may convey verbatim copies of the Program's source code as you
+receive it, in any medium, provided that you conspicuously and
+appropriately publish on each copy an appropriate copyright notice;
+keep intact all notices stating that this License and any
+non-permissive terms added in accord with section 7 apply to the code;
+keep intact all notices of the absence of any warranty; and give all
+recipients a copy of this License along with the Program.
+
+  You may charge any price or no price for each copy that you convey,
+and you may offer support or warranty protection for a fee.
+
+  5. Conveying Modified Source Versions.
+
+  You may convey a work based on the Program, or the modifications to
+produce it from the Program, in the form of source code under the
+terms of section 4, provided that you also meet all of these conditions:
+
+    a) The work must carry prominent notices stating that you modified
+    it, and giving a relevant date.
+
+    b) The work must carry prominent notices stating that it is
+    released under this License and any conditions added under section
+    7.  This requirement modifies the requirement in section 4 to
+    "keep intact all notices".
+
+    c) You must license the entire work, as a whole, under this
+    License to anyone who comes into possession of a copy.  This
+    License will therefore apply, along with any applicable section 7
+    additional terms, to the whole of the work, and all its parts,
+    regardless of how they are packaged.  This License gives no
+    permission to license the work in any other way, but it does not
+    invalidate such permission if you have separately received it.
+
+    d) If the work has interactive user interfaces, each must display
+    Appropriate Legal Notices; however, if the Program has interactive
+    interfaces that do not display Appropriate Legal Notices, your
+    work need not make them do so.
+
+  A compilation of a covered work with other separate and independent
+works, which are not by their nature extensions of the covered work,
+and which are not combined with it such as to form a larger program,
+in or on a volume of a storage or distribution medium, is called an
+"aggregate" if the compilation and its resulting copyright are not
+used to limit the access or legal rights of the compilation's users
+beyond what the individual works permit.  Inclusion of a covered work
+in an aggregate does not cause this License to apply to the other
+parts of the aggregate.
+
+  6. Conveying Non-Source Forms.
+
+  You may convey a covered work in object code form under the terms
+of sections 4 and 5, provided that you also convey the
+machine-readable Corresponding Source under the terms of this License,
+in one of these ways:
+
+    a) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by the
+    Corresponding Source fixed on a durable physical medium
+    customarily used for software interchange.
+
+    b) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by a
+    written offer, valid for at least three years and valid for as
+    long as you offer spare parts or customer support for that product
+    model, to give anyone who possesses the object code either (1) a
+    copy of the Corresponding Source for all the software in the
+    product that is covered by this License, on a durable physical
+    medium customarily used for software interchange, for a price no
+    more than your reasonable cost of physically performing this
+    conveying of source, or (2) access to copy the
+    Corresponding Source from a network server at no charge.
+
+    c) Convey individual copies of the object code with a copy of the
+    written offer to provide the Corresponding Source.  This
+    alternative is allowed only occasionally and noncommercially, and
+    only if you received the object code with such an offer, in accord
+    with subsection 6b.
+
+    d) Convey the object code by offering access from a designated
+    place (gratis or for a charge), and offer equivalent access to the
+    Corresponding Source in the same way through the same place at no
+    further charge.  You need not require recipients to copy the
+    Corresponding Source along with the object code.  If the place to
+    copy the object code is a network server, the Corresponding Source
+    may be on a different server (operated by you or a third party)
+    that supports equivalent copying facilities, provided you maintain
+    clear directions next to the object code saying where to find the
+    Corresponding Source.  Regardless of what server hosts the
+    Corresponding Source, you remain obligated to ensure that it is
+    available for as long as needed to satisfy these requirements.
+
+    e) Convey the object code using peer-to-peer transmission, provided
+    you inform other peers where the object code and Corresponding
+    Source of the work are being offered to the general public at no
+    charge under subsection 6d.
+
+  A separable portion of the object code, whose source code is excluded
+from the Corresponding Source as a System Library, need not be
+included in conveying the object code work.
+
+  A "User Product" is either (1) a "consumer product", which means any
+tangible personal property which is normally used for personal, family,
+or household purposes, or (2) anything designed or sold for incorporation
+into a dwelling.  In determining whether a product is a consumer product,
+doubtful cases shall be resolved in favor of coverage.  For a particular
+product received by a particular user, "normally used" refers to a
+typical or common use of that class of product, regardless of the status
+of the particular user or of the way in which the particular user
+actually uses, or expects or is expected to use, the product.  A product
+is a consumer product regardless of whether the product has substantial
+commercial, industrial or non-consumer uses, unless such uses represent
+the only significant mode of use of the product.
+
+  "Installation Information" for a User Product means any methods,
+procedures, authorization keys, or other information required to install
+and execute modified versions of a covered work in that User Product from
+a modified version of its Corresponding Source.  The information must
+suffice to ensure that the continued functioning of the modified object
+code is in no case prevented or interfered with solely because
+modification has been made.
+
+  If you convey an object code work under this section in, or with, or
+specifically for use in, a User Product, and the conveying occurs as
+part of a transaction in which the right of possession and use of the
+User Product is transferred to the recipient in perpetuity or for a
+fixed term (regardless of how the transaction is characterized), the
+Corresponding Source conveyed under this section must be accompanied
+by the Installation Information.  But this requirement does not apply
+if neither you nor any third party retains the ability to install
+modified object code on the User Product (for example, the work has
+been installed in ROM).
+
+  The requirement to provide Installation Information does not include a
+requirement to continue to provide support service, warranty, or updates
+for a work that has been modified or installed by the recipient, or for
+the User Product in which it has been modified or installed.  Access to a
+network may be denied when the modification itself materially and
+adversely affects the operation of the network or violates the rules and
+protocols for communication across the network.
+
+  Corresponding Source conveyed, and Installation Information provided,
+in accord with this section must be in a format that is publicly
+documented (and with an implementation available to the public in
+source code form), and must require no special password or key for
+unpacking, reading or copying.
+
+  7. Additional Terms.
+
+  "Additional permissions" are terms that supplement the terms of this
+License by making exceptions from one or more of its conditions.
+Additional permissions that are applicable to the entire Program shall
+be treated as though they were included in this License, to the extent
+that they are valid under applicable law.  If additional permissions
+apply only to part of the Program, that part may be used separately
+under those permissions, but the entire Program remains governed by
+this License without regard to the additional permissions.
+
+  When you convey a copy of a covered work, you may at your option
+remove any additional permissions from that copy, or from any part of
+it.  (Additional permissions may be written to require their own
+removal in certain cases when you modify the work.)  You may place
+additional permissions on material, added by you to a covered work,
+for which you have or can give appropriate copyright permission.
+
+  Notwithstanding any other provision of this License, for material you
+add to a covered work, you may (if authorized by the copyright holders of
+that material) supplement the terms of this License with terms:
+
+    a) Disclaiming warranty or limiting liability differently from the
+    terms of sections 15 and 16 of this License; or
+
+    b) Requiring preservation of specified reasonable legal notices or
+    author attributions in that material or in the Appropriate Legal
+    Notices displayed by works containing it; or
+
+    c) Prohibiting misrepresentation of the origin of that material, or
+    requiring that modified versions of such material be marked in
+    reasonable ways as different from the original version; or
+
+    d) Limiting the use for publicity purposes of names of licensors or
+    authors of the material; or
+
+    e) Declining to grant rights under trademark law for use of some
+    trade names, trademarks, or service marks; or
+
+    f) Requiring indemnification of licensors and authors of that
+    material by anyone who conveys the material (or modified versions of
+    it) with contractual assumptions of liability to the recipient, for
+    any liability that these contractual assumptions directly impose on
+    those licensors and authors.
+
+  All other non-permissive additional terms are considered "further
+restrictions" within the meaning of section 10.  If the Program as you
+received it, or any part of it, contains a notice stating that it is
+governed by this License along with a term that is a further
+restriction, you may remove that term.  If a license document contains
+a further restriction but permits relicensing or conveying under this
+License, you may add to a covered work material governed by the terms
+of that license document, provided that the further restriction does
+not survive such relicensing or conveying.
+
+  If you add terms to a covered work in accord with this section, you
+must place, in the relevant source files, a statement of the
+additional terms that apply to those files, or a notice indicating
+where to find the applicable terms.
+
+  Additional terms, permissive or non-permissive, may be stated in the
+form of a separately written license, or stated as exceptions;
+the above requirements apply either way.
+
+  8. Termination.
+
+  You may not propagate or modify a covered work except as expressly
+provided under this License.  Any attempt otherwise to propagate or
+modify it is void, and will automatically terminate your rights under
+this License (including any patent licenses granted under the third
+paragraph of section 11).
+
+  However, if you cease all violation of this License, then your
+license from a particular copyright holder is reinstated (a)
+provisionally, unless and until the copyright holder explicitly and
+finally terminates your license, and (b) permanently, if the copyright
+holder fails to notify you of the violation by some reasonable means
+prior to 60 days after the cessation.
+
+  Moreover, your license from a particular copyright holder is
+reinstated permanently if the copyright holder notifies you of the
+violation by some reasonable means, this is the first time you have
+received notice of violation of this License (for any work) from that
+copyright holder, and you cure the violation prior to 30 days after
+your receipt of the notice.
+
+  Termination of your rights under this section does not terminate the
+licenses of parties who have received copies or rights from you under
+this License.  If your rights have been terminated and not permanently
+reinstated, you do not qualify to receive new licenses for the same
+material under section 10.
+
+  9. Acceptance Not Required for Having Copies.
+
+  You are not required to accept this License in order to receive or
+run a copy of the Program.  Ancillary propagation of a covered work
+occurring solely as a consequence of using peer-to-peer transmission
+to receive a copy likewise does not require acceptance.  However,
+nothing other than this License grants you permission to propagate or
+modify any covered work.  These actions infringe copyright if you do
+not accept this License.  Therefore, by modifying or propagating a
+covered work, you indicate your acceptance of this License to do so.
+
+  10. Automatic Licensing of Downstream Recipients.
+
+  Each time you convey a covered work, the recipient automatically
+receives a license from the original licensors, to run, modify and
+propagate that work, subject to this License.  You are not responsible
+for enforcing compliance by third parties with this License.
+
+  An "entity transaction" is a transaction transferring control of an
+organization, or substantially all assets of one, or subdividing an
+organization, or merging organizations.  If propagation of a covered
+work results from an entity transaction, each party to that
+transaction who receives a copy of the work also receives whatever
+licenses to the work the party's predecessor in interest had or could
+give under the previous paragraph, plus a right to possession of the
+Corresponding Source of the work from the predecessor in interest, if
+the predecessor has it or can get it with reasonable efforts.
+
+  You may not impose any further restrictions on the exercise of the
+rights granted or affirmed under this License.  For example, you may
+not impose a license fee, royalty, or other charge for exercise of
+rights granted under this License, and you may not initiate litigation
+(including a cross-claim or counterclaim in a lawsuit) alleging that
+any patent claim is infringed by making, using, selling, offering for
+sale, or importing the Program or any portion of it.
+
+  11. Patents.
+
+  A "contributor" is a copyright holder who authorizes use under this
+License of the Program or a work on which the Program is based.  The
+work thus licensed is called the contributor's "contributor version".
+
+  A contributor's "essential patent claims" are all patent claims
+owned or controlled by the contributor, whether already acquired or
+hereafter acquired, that would be infringed by some manner, permitted
+by this License, of making, using, or selling its contributor version,
+but do not include claims that would be infringed only as a
+consequence of further modification of the contributor version.  For
+purposes of this definition, "control" includes the right to grant
+patent sublicenses in a manner consistent with the requirements of
+this License.
+
+  Each contributor grants you a non-exclusive, worldwide, royalty-free
+patent license under the contributor's essential patent claims, to
+make, use, sell, offer for sale, import and otherwise run, modify and
+propagate the contents of its contributor version.
+
+  In the following three paragraphs, a "patent license" is any express
+agreement or commitment, however denominated, not to enforce a patent
+(such as an express permission to practice a patent or covenant not to
+sue for patent infringement).  To "grant" such a patent license to a
+party means to make such an agreement or commitment not to enforce a
+patent against the party.
+
+  If you convey a covered work, knowingly relying on a patent license,
+and the Corresponding Source of the work is not available for anyone
+to copy, free of charge and under the terms of this License, through a
+publicly available network server or other readily accessible means,
+then you must either (1) cause the Corresponding Source to be so
+available, or (2) arrange to deprive yourself of the benefit of the
+patent license for this particular work, or (3) arrange, in a manner
+consistent with the requirements of this License, to extend the patent
+license to downstream recipients.  "Knowingly relying" means you have
+actual knowledge that, but for the patent license, your conveying the
+covered work in a country, or your recipient's use of the covered work
+in a country, would infringe one or more identifiable patents in that
+country that you have reason to believe are valid.
+
+  If, pursuant to or in connection with a single transaction or
+arrangement, you convey, or propagate by procuring conveyance of, a
+covered work, and grant a patent license to some of the parties
+receiving the covered work authorizing them to use, propagate, modify
+or convey a specific copy of the covered work, then the patent license
+you grant is automatically extended to all recipients of the covered
+work and works based on it.
+
+  A patent license is "discriminatory" if it does not include within
+the scope of its coverage, prohibits the exercise of, or is
+conditioned on the non-exercise of one or more of the rights that are
+specifically granted under this License.  You may not convey a covered
+work if you are a party to an arrangement with a third party that is
+in the business of distributing software, under which you make payment
+to the third party based on the extent of your activity of conveying
+the work, and under which the third party grants, to any of the
+parties who would receive the covered work from you, a discriminatory
+patent license (a) in connection with copies of the covered work
+conveyed by you (or copies made from those copies), or (b) primarily
+for and in connection with specific products or compilations that
+contain the covered work, unless you entered into that arrangement,
+or that patent license was granted, prior to 28 March 2007.
+
+  Nothing in this License shall be construed as excluding or limiting
+any implied license or other defenses to infringement that may
+otherwise be available to you under applicable patent law.
+
+  12. No Surrender of Others' Freedom.
+
+  If conditions are imposed on you (whether by court order, agreement or
+otherwise) that contradict the conditions of this License, they do not
+excuse you from the conditions of this License.  If you cannot convey a
+covered work so as to satisfy simultaneously your obligations under this
+License and any other pertinent obligations, then as a consequence you may
+not convey it at all.  For example, if you agree to terms that obligate you
+to collect a royalty for further conveying from those to whom you convey
+the Program, the only way you could satisfy both those terms and this
+License would be to refrain entirely from conveying the Program.
+
+  13. Use with the GNU Affero General Public License.
+
+  Notwithstanding any other provision of this License, you have
+permission to link or combine any covered work with a work licensed
+under version 3 of the GNU Affero General Public License into a single
+combined work, and to convey the resulting work.  The terms of this
+License will continue to apply to the part which is the covered work,
+but the special requirements of the GNU Affero General Public License,
+section 13, concerning interaction through a network will apply to the
+combination as such.
+
+  14. Revised Versions of this License.
+
+  The Free Software Foundation may publish revised and/or new versions of
+the GNU General Public License from time to time.  Such new versions will
+be similar in spirit to the present version, but may differ in detail to
+address new problems or concerns.
+
+  Each version is given a distinguishing version number.  If the
+Program specifies that a certain numbered version of the GNU General
+Public License "or any later version" applies to it, you have the
+option of following the terms and conditions either of that numbered
+version or of any later version published by the Free Software
+Foundation.  If the Program does not specify a version number of the
+GNU General Public License, you may choose any version ever published
+by the Free Software Foundation.
+
+  If the Program specifies that a proxy can decide which future
+versions of the GNU General Public License can be used, that proxy's
+public statement of acceptance of a version permanently authorizes you
+to choose that version for the Program.
+
+  Later license versions may give you additional or different
+permissions.  However, no additional obligations are imposed on any
+author or copyright holder as a result of your choosing to follow a
+later version.
+
+  15. Disclaimer of Warranty.
+
+  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
+APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
+HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
+OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
+IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+  16. Limitation of Liability.
+
+  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
+WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
+THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
+GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
+USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
+DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
+PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
+EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGES.
+
+  17. Interpretation of Sections 15 and 16.
+
+  If the disclaimer of warranty and limitation of liability provided
+above cannot be given local legal effect according to their terms,
+reviewing courts shall apply local law that most closely approximates
+an absolute waiver of all civil liability in connection with the
+Program, unless a warranty or assumption of liability accompanies a
+copy of the Program in return for a fee.
+
+                     END OF TERMS AND CONDITIONS
+
+            How to Apply These Terms to Your New Programs
+
+  If you develop a new program, and you want it to be of the greatest
+possible use to the public, the best way to achieve this is to make it
+free software which everyone can redistribute and change under these terms.
+
+  To do so, attach the following notices to the program.  It is safest
+to attach them to the start of each source file to most effectively
+state the exclusion of warranty; and each file should have at least
+the "copyright" line and a pointer to where the full notice is found.
+
+    <one line to give the program's name and a brief idea of what it does.>
+    Copyright (C) <year>  <name of author>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+Also add information on how to contact you by electronic and paper mail.
+
+  If the program does terminal interaction, make it output a short
+notice like this when it starts in an interactive mode:
+
+    <program>  Copyright (C) <year>  <name of author>
+    This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
+    This is free software, and you are welcome to redistribute it
+    under certain conditions; type `show c' for details.
+
+The hypothetical commands `show w' and `show c' should show the appropriate
+parts of the General Public License.  Of course, your program's commands
+might be different; for a GUI interface, you would use an "about box".
+
+  You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU GPL, see
+<https://www.gnu.org/licenses/>.
+
+  The GNU General Public License does not permit incorporating your program
+into proprietary programs.  If your program is a subroutine library, you
+may consider it more useful to permit linking proprietary applications with
+the library.  If this is what you want to do, use the GNU Lesser General
+Public License instead of this License.  But first, please read
+<https://www.gnu.org/licenses/why-not-lgpl.html>.

mulres_infer/modules/miche/download_ckpt.py

@@ -0,0 +1,25 @@
+from huggingface_hub import hf_hub_download
+import os
+
+# 设置镜像（国内加速）
+# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
+
+def download_michelangelo_checkpoint():
+    """下载Michelangelo模型的checkpoint文件"""
+    try:
+        # 下载文件
+        local_path = hf_hub_download(
+            repo_id="Maikou/Michelangelo",
+            filename="checkpoints/aligned_shape_latents/shapevae-256.ckpt",
+            local_dir="./downloads",  # 本地保存目录
+            local_dir_use_symlinks=False,  # 不使用符号链接
+            resume_download=True,     # 支持断点续传
+        )
+        print(f"文件已下载到: {local_path}")
+        return local_path
+    except Exception as e:
+        print(f"下载失败: {e}")
+        return None
+
+if __name__ == "__main__":
+    download_michelangelo_checkpoint()

mulres_infer/modules/miche/encode.py

@@ -0,0 +1,84 @@
+import os
+# -*- coding: utf-8 -*-
+import argparse
+from omegaconf import OmegaConf
+import numpy as np
+import torch
+from .michelangelo.utils.misc import instantiate_from_config
+
+def load_surface(fp):
+    
+    with np.load(fp) as input_pc:
+        surface = input_pc['points']
+        normal = input_pc['normals']
+    
+    rng = np.random.default_rng()
+    ind = rng.choice(surface.shape[0], 4096, replace=False)
+    surface = torch.FloatTensor(surface[ind])
+    normal = torch.FloatTensor(normal[ind])
+    
+    surface = torch.cat([surface, normal], dim=-1).unsqueeze(0).cuda()
+    
+    return surface
+
+def reconstruction(args, model, bounds=(-1.25, -1.25, -1.25, 1.25, 1.25, 1.25), octree_depth=7, num_chunks=10000):
+
+    surface = load_surface(args.pointcloud_path)
+    # old_surface = surface.clone()
+
+    # surface[0,:,0]*=-1
+    # surface[0,:,1]*=-1
+    surface[0,:,2]*=-1
+
+    # encoding
+    shape_embed, shape_latents = model.model.encode_shape_embed(surface, return_latents=True)    
+    shape_zq, posterior = model.model.shape_model.encode_kl_embed(shape_latents)
+
+    # decoding
+    latents = model.model.shape_model.decode(shape_zq)
+    # geometric_func = partial(model.model.shape_model.query_geometry, latents=latents)
+    
+    return 0
+
+def load_model(ckpt_path=None, config_path=None):
+    _HERE = os.path.dirname(os.path.abspath(__file__))
+    config_path = config_path or os.path.join(_HERE, "shapevae-256.yaml")
+    
+    model_config = OmegaConf.load(config_path)
+    # print(model_config)
+    if hasattr(model_config, "model"):
+        model_config = model_config.model
+
+    original_load = torch.load
+    torch.load = lambda path, *args, **kwargs: original_load(path, map_location="cpu", weights_only=False)
+    try:
+        model = instantiate_from_config(model_config, ckpt_path=ckpt_path)
+        model = model.eval()
+    finally:
+        # 恢复原始的 torch.load
+        torch.load = original_load
+
+    return model
+if __name__ == "__main__":
+    '''
+    1. Reconstruct point cloud
+    2. Image-conditioned generation
+    3. Text-conditioned generation
+    '''
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, required=True)
+    parser.add_argument("--ckpt_path", type=str, required=True)
+    parser.add_argument("--pointcloud_path", type=str, default='./example_data/surface.npz', 
+                        help='Path to the input point cloud')
+    parser.add_argument("--image_path", type=str, help='Path to the input image')
+    parser.add_argument("--text", type=str, 
+                        help='Input text within a format: A 3D model of motorcar; Porsche 911.')
+    parser.add_argument("--output_dir", type=str, default='./output')
+    parser.add_argument("-s", "--seed", type=int, default=0)
+    args = parser.parse_args()
+    
+    print(f'-----------------------------------------------------------------------------')
+    print(f'>>> Output directory: {args.output_dir}')
+    print(f'-----------------------------------------------------------------------------')
+    
+    reconstruction(args, load_model(args))

mulres_infer/modules/miche/michelangelo/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

mulres_infer/modules/miche/michelangelo/graphics/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

mulres_infer/modules/miche/michelangelo/graphics/primitives/__init__.py

@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+
+from .volume import generate_dense_grid_points
+

mulres_infer/modules/miche/michelangelo/graphics/primitives/volume.py

@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+
+# produce dense points
+def generate_dense_grid_points(bbox_min: np.ndarray,
+                               bbox_max: np.ndarray,
+                               octree_depth: int,
+                               indexing: str = "ij"):
+    length = bbox_max - bbox_min
+    num_cells = np.exp2(octree_depth)
+    x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32)
+    y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32)
+    z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32)
+    [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing)
+    xyz = np.stack((xs, ys, zs), axis=-1)
+    xyz = xyz.reshape(-1, 3)
+    grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1]
+
+    return xyz, grid_size, length
+

mulres_infer/modules/miche/michelangelo/models/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

mulres_infer/modules/miche/michelangelo/models/modules/__init__.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+
+from .checkpoint import checkpoint

mulres_infer/modules/miche/michelangelo/models/modules/checkpoint.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+
+import torch
+from typing import Callable, Iterable, Sequence, Union
+
+
+def checkpoint(
+    func: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor]]],
+    inputs: Sequence[torch.Tensor],
+    params: Iterable[torch.Tensor],
+    flag: bool,
+    use_deepspeed: bool = False
+):
+    # Evaluate a function without caching intermediate activations, allowing for
+    # reduced memory at the expense of extra compute in the backward pass.
+    # :param func: the function to evaluate.
+    # :param inputs: the argument sequence to pass to `func`.
+    # :param params: a sequence of parameters `func` depends on but does not
+    #                explicitly take as arguments.
+    # :param flag: if False, disable gradient checkpointing.
+    # :param use_deepspeed: if True, use deepspeed
+    if flag:
+        if use_deepspeed:
+            import deepspeed
+            return deepspeed.checkpointing.checkpoint(func, *inputs)
+
+        args = tuple(inputs) + tuple(params)
+        return CheckpointFunction.apply(func, len(inputs), *args)
+    else:
+        return func(*inputs)
+
+
+class CheckpointFunction(torch.autograd.Function):
+    @staticmethod
+    @torch.cuda.amp.custom_fwd
+    def forward(ctx, run_function, length, *args):
+        ctx.run_function = run_function
+        ctx.input_tensors = list(args[:length])
+        ctx.input_params = list(args[length:])
+
+        with torch.no_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        return output_tensors
+
+    @staticmethod
+    @torch.cuda.amp.custom_bwd
+    def backward(ctx, *output_grads):
+        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
+        with torch.enable_grad():
+            # Fixes a bug where the first op in run_function modifies the
+            # Tensor storage in place, which is not allowed for detach()'d
+            # Tensors.
+            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+            output_tensors = ctx.run_function(*shallow_copies)
+        input_grads = torch.autograd.grad(
+            output_tensors,
+            ctx.input_tensors + ctx.input_params,
+            output_grads,
+            allow_unused=True,
+        )
+        del ctx.input_tensors
+        del ctx.input_params
+        del output_tensors
+        return (None, None) + input_grads

mulres_infer/modules/miche/michelangelo/models/modules/distributions.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+
+import torch
+import numpy as np
+from typing import Union, List
+
+
+class DiagonalGaussianDistribution(object):
+    # Gaussian distribution
+    def __init__(self, parameters: Union[torch.Tensor, List[torch.Tensor]], deterministic=False, feat_dim=1):
+        self.feat_dim = feat_dim
+        self.parameters = parameters
+
+        if isinstance(parameters, list):
+            self.mean = parameters[0]
+            self.logvar = parameters[1]
+        else:
+            self.mean, self.logvar = torch.chunk(parameters, 2, dim=feat_dim)
+
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean)
+
+    # sample from the guassian distribution
+    def sample(self):
+        x = self.mean + self.std * torch.randn_like(self.mean)
+        return x
+
+    def kl(self, other=None, dims=(1, 2, 3)):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        else:
+            if other is None:
+                return 0.5 * torch.mean(torch.pow(self.mean, 2)
+                                        + self.var - 1.0 - self.logvar,
+                                        dim=dims)
+            else:
+                return 0.5 * torch.mean(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var - 1.0 - self.logvar + other.logvar,
+                    dim=dims)
+
+    def nll(self, sample, dims=(1, 2, 3)):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims)
+
+    def mode(self):
+        return self.mean
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    # Compute the KL divergence between two gaussians.
+    # Shapes are automatically broadcasted, so batches can be compared to
+    # scalars, among other use cases.
+
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, torch.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for torch.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+            -1.0
+            + logvar2
+            - logvar1
+            + torch.exp(logvar1 - logvar2)
+            + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
+    )

mulres_infer/modules/miche/michelangelo/models/modules/embedder.py

@@ -0,0 +1,213 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+import torch
+import torch.nn as nn
+import math
+
+VALID_EMBED_TYPES = ["identity", "fourier", "hashgrid", "sphere_harmonic", "triplane_fourier"]
+
+
+class FourierEmbedder(nn.Module):
+    """The sin/cosine positional embedding. Given an input tensor `x` of shape [n_batch, ..., c_dim], it converts
+    each feature dimension of `x[..., i]` into:
+        [
+            sin(x[..., i]),
+            sin(f_1*x[..., i]),
+            sin(f_2*x[..., i]),
+            ...
+            sin(f_N * x[..., i]),
+            cos(x[..., i]),
+            cos(f_1*x[..., i]),
+            cos(f_2*x[..., i]),
+            ...
+            cos(f_N * x[..., i]),
+            x[..., i]     # only present if include_input is True.
+        ], here f_i is the frequency.
+
+    Denote the space is [0 / num_freqs, 1 / num_freqs, 2 / num_freqs, 3 / num_freqs, ..., (num_freqs - 1) / num_freqs].
+    If logspace is True, then the frequency f_i is [2^(0 / num_freqs), ..., 2^(i / num_freqs), ...];
+    Otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1)].
+
+    Args:
+        num_freqs (int): the number of frequencies, default is 6;
+        logspace (bool): If logspace is True, then the frequency f_i is [..., 2^(i / num_freqs), ...],
+            otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1)];
+        input_dim (int): the input dimension, default is 3;
+        include_input (bool): include the input tensor or not, default is True.
+
+    Attributes:
+        frequencies (torch.Tensor): If logspace is True, then the frequency f_i is [..., 2^(i / num_freqs), ...],
+                otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1);
+
+        out_dim (int): the embedding size, if include_input is True, it is input_dim * (num_freqs * 2 + 1),
+            otherwise, it is input_dim * num_freqs * 2.
+
+    """
+
+    def __init__(self,
+                 num_freqs: int = 6,
+                 logspace: bool = True,
+                 input_dim: int = 3,
+                 include_input: bool = True,
+                 include_pi: bool = True) -> None:
+
+        """The initialization"""
+
+        super().__init__()
+
+        if logspace:
+            frequencies = 2.0 ** torch.arange(
+                num_freqs,
+                dtype=torch.float32
+            )
+        else:
+            frequencies = torch.linspace(
+                1.0,
+                2.0 ** (num_freqs - 1),
+                num_freqs,
+                dtype=torch.float32
+            )
+
+        if include_pi:
+            frequencies *= torch.pi
+
+        self.register_buffer("frequencies", frequencies, persistent=False)
+        self.include_input = include_input
+        self.num_freqs = num_freqs
+
+        self.out_dim = self.get_dims(input_dim)
+
+    def get_dims(self, input_dim):
+        temp = 1 if self.include_input or self.num_freqs == 0 else 0
+        out_dim = input_dim * (self.num_freqs * 2 + temp)
+
+        return out_dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """ Forward process.
+
+        Args:
+            x: tensor of shape [..., dim]
+
+        Returns:
+            embedding: an embedding of `x` of shape [..., dim * (num_freqs * 2 + temp)]
+                where temp is 1 if include_input is True and 0 otherwise.
+        """
+
+        if self.num_freqs > 0:
+            embed = (x[..., None].contiguous() * self.frequencies).view(*x.shape[:-1], -1)
+            if self.include_input:
+                return torch.cat((x, embed.sin(), embed.cos()), dim=-1)
+            else:
+                return torch.cat((embed.sin(), embed.cos()), dim=-1)
+        else:
+            return x
+
+
+class LearnedFourierEmbedder(nn.Module):
+    """ following @crowsonkb "s lead with learned sinusoidal pos emb """
+    """ https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """
+
+    def __init__(self, in_channels, dim):
+        super().__init__()
+        assert (dim % 2) == 0
+        half_dim = dim // 2
+        per_channel_dim = half_dim // in_channels
+        self.weights = nn.Parameter(torch.randn(per_channel_dim))
+
+    def forward(self, x):
+        """
+
+        Args:
+            x (torch.FloatTensor): [..., c]
+
+        Returns:
+            x (torch.FloatTensor): [..., d]
+        """
+
+        # [b, t, c, 1] * [1, d] = [b, t, c, d] -> [b, t, c * d]
+        freqs = (x[..., None] * self.weights[None] * 2 * np.pi).view(*x.shape[:-1], -1)
+        fouriered = torch.cat((x, freqs.sin(), freqs.cos()), dim=-1)
+        return fouriered
+
+
+class TriplaneLearnedFourierEmbedder(nn.Module):
+    def __init__(self, in_channels, dim):
+        super().__init__()
+
+        self.yz_plane_embedder = LearnedFourierEmbedder(in_channels, dim)
+        self.xz_plane_embedder = LearnedFourierEmbedder(in_channels, dim)
+        self.xy_plane_embedder = LearnedFourierEmbedder(in_channels, dim)
+
+        self.out_dim = in_channels + dim
+
+    def forward(self, x):
+
+        yz_embed = self.yz_plane_embedder(x)
+        xz_embed = self.xz_plane_embedder(x)
+        xy_embed = self.xy_plane_embedder(x)
+
+        embed = yz_embed + xz_embed + xy_embed
+
+        return embed
+
+
+def sequential_pos_embed(num_len, embed_dim):
+    assert embed_dim % 2 == 0
+
+    pos = torch.arange(num_len, dtype=torch.float32)
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000 ** omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = torch.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = torch.sin(out)  # (M, D/2)
+    emb_cos = torch.cos(out)  # (M, D/2)
+
+    embeddings = torch.cat([emb_sin, emb_cos], dim=1)  # (M, D)
+
+    return embeddings
+
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """
+    Create sinusoidal timestep embeddings.
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].to(timesteps.dtype) * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+def get_embedder(embed_type="fourier", num_freqs=-1, input_dim=3, degree=4,
+                 num_levels=16, level_dim=2, per_level_scale=2, base_resolution=16,
+                 log2_hashmap_size=19, desired_resolution=None):
+    if embed_type == "identity" or (embed_type == "fourier" and num_freqs == -1):
+        return nn.Identity(), input_dim
+
+    elif embed_type == "fourier":
+        embedder_obj = FourierEmbedder(num_freqs=num_freqs, input_dim=input_dim,
+                                       logspace=True, include_input=True)
+        return embedder_obj, embedder_obj.out_dim
+
+    elif embed_type == "hashgrid":
+        raise NotImplementedError
+
+    elif embed_type == "sphere_harmonic":
+        raise NotImplementedError
+
+    else:
+        raise ValueError(f"{embed_type} is not valid. Currently only supprts {VALID_EMBED_TYPES}")

mulres_infer/modules/miche/michelangelo/models/modules/transformer_blocks.py

@@ -0,0 +1,309 @@
+# -*- coding: utf-8 -*-
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+from modules.miche.michelangelo.models.modules.checkpoint import checkpoint
+
+# Initialize linear layers with normal distribution weights and zero biases
+def init_linear(l, stddev):
+    nn.init.normal_(l.weight, std=stddev)
+    if l.bias is not None:
+        nn.init.constant_(l.bias, 0.0)
+
+# Multihead attention module
+class MultiheadAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int,  # Context size
+        width: int,  # Width of the input tensor
+        heads: int,  # Number of attention heads
+        init_scale: float,  # Initialization scale for weights
+        qkv_bias: bool,  # Whether to use bias in QKV layers
+        flash: bool = False  # Whether to use flash attention
+    ):
+        super().__init__()
+        self.n_ctx = n_ctx
+        self.width = width
+        self.heads = heads
+        self.c_qkv = nn.Linear(width, width * 3, bias=qkv_bias, device=device, dtype=dtype)
+        self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
+        self.attention = QKVMultiheadAttention(device=device, dtype=dtype, heads=heads, n_ctx=n_ctx, flash=flash)
+        init_linear(self.c_qkv, init_scale)
+        init_linear(self.c_proj, init_scale)
+
+    def forward(self, x):
+        x = self.c_qkv(x)
+        x = checkpoint(self.attention, (x,), (), True)
+        x = self.c_proj(x)
+        return x
+
+# QKV multihead attention module
+class QKVMultiheadAttention(nn.Module):
+    def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int, n_ctx: int, flash: bool = False):
+        super().__init__()
+        self.device = device
+        self.dtype = dtype
+        self.heads = heads
+        self.n_ctx = n_ctx
+        self.flash = flash
+
+    def forward(self, qkv):
+        bs, n_ctx, width = qkv.shape
+        attn_ch = width // self.heads // 3
+        scale = 1 / math.sqrt(math.sqrt(attn_ch))
+        qkv = qkv.view(bs, n_ctx, self.heads, -1)
+        q, k, v = torch.split(qkv, attn_ch, dim=-1)
+
+        if self.flash:
+            # print(1)
+            # # out = F.scaled_dot_product_attention(q, k, v)
+            # q = q.permute(0, 2, 1, 3)  # -> (b, h, t, c)
+            # k = k.permute(0, 2, 1, 3)  # -> (b, h, t, c)
+            # v = v.permute(0, 2, 1, 3)  # -> (b, h, t, c)
+
+            # # F.scaled_dot_product_attention 直接返回计算结果，无需手动 reshape
+            # out = F.scaled_dot_product_attention(q, k, v)  # -> (b, h, t, c)
+
+            # # 关键：将输出变回 (b, t, h, c) 并 reshape，以匹配 else 分支的输出
+            # out = out.permute(0, 2, 1, 3).reshape(bs, n_ctx, -1) # -> (b, t, h*c)
+            out = torch_npu.npu_fusion_attention(
+                q, k, v, self.heads, "BSND", keep_prob=1.0, scale=scale)[0].reshape(bs, n_ctx, -1)
+        else:
+            weight = torch.einsum(
+                "bthc,bshc->bhts", q * scale, k * scale
+            )  # More stable with f16 than dividing afterwards
+            wdtype = weight.dtype
+            weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
+            out = torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
+
+        return out
+
+# Residual attention block module
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        use_checkpoint: bool = False, 
+        n_ctx: int,  # Context size
+        width: int,  # Width of the input tensor
+        heads: int,  # Number of attention heads
+        init_scale: float,  # Initialization scale for weights
+        qkv_bias: bool,  # Whether to use bias in QKV layers
+        flash: bool = False  # Whether to use flash attention
+    ):
+        super().__init__()
+
+        self.use_checkpoint = use_checkpoint
+
+        self.attn = MultiheadAttention(
+            device=device,
+            dtype=dtype,
+            n_ctx=n_ctx,
+            width=width,
+            heads=heads,
+            init_scale=init_scale,
+            qkv_bias=qkv_bias,
+            flash=flash
+        )
+        self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
+        self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=init_scale)
+        self.ln_2 = nn.LayerNorm(width, device=device, dtype=dtype)
+
+    def _forward(self, x: torch.Tensor):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+    def forward(self, x: torch.Tensor):
+        return checkpoint(self._forward, (x,), self.parameters(), self.use_checkpoint)
+
+# Multihead cross attention module
+class MultiheadCrossAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_data: Optional[int] = None,
+        data_width: Optional[int] = None,
+        width: int,  # Width of the input tensor
+        heads: int,  # Number of attention heads
+        init_scale: float,  # Initialization scale for weights
+        qkv_bias: bool,  # Whether to use bias in QKV layers
+        flash: bool = False  # Whether to use flash attention
+    ):
+        super().__init__()
+        self.n_data = n_data
+        self.width = width
+        self.heads = heads
+        self.data_width = width if data_width is None else data_width
+        self.c_q = nn.Linear(width, width, bias=qkv_bias, device=device, dtype=dtype)
+        self.c_kv = nn.Linear(self.data_width, width * 2, bias=qkv_bias, device=device, dtype=dtype)
+        self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
+        self.attention = QKVMultiheadCrossAttention(
+            device=device, dtype=dtype, heads=heads, n_data=n_data, flash=flash
+        )
+        init_linear(self.c_q, init_scale)
+        init_linear(self.c_kv, init_scale)
+        init_linear(self.c_proj, init_scale)
+
+    def forward(self, x, data):
+        x = self.c_q(x)
+        data = self.c_kv(data)
+        x = checkpoint(self.attention, (x, data), (), True)
+        x = self.c_proj(x)
+        return x
+
+# QKV multihead cross attention module
+class QKVMultiheadCrossAttention(nn.Module):
+    def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int,
+                 flash: bool = False, n_data: Optional[int] = None):
+
+        super().__init__()
+        self.device = device
+        self.dtype = dtype
+        self.heads = heads
+        self.n_data = n_data
+        self.flash = flash
+
+    def forward(self, q, kv):
+        _, n_ctx, _ = q.shape
+        bs, n_data, width = kv.shape
+        attn_ch = width // self.heads // 2
+        scale = 1 / math.sqrt(math.sqrt(attn_ch))
+        q = q.view(bs, n_ctx, self.heads, -1)
+        kv = kv.view(bs, n_data, self.heads, -1)
+        k, v = torch.split(kv, attn_ch, dim=-1)
+
+        if self.flash:
+            # print(2)
+            # # out = F.scaled_dot_product_attention(q, k, v)
+            # q = q.permute(0, 2, 1, 3)  # -> (b, h, t, c)
+            # k = k.permute(0, 2, 1, 3)  # -> (b, h, t, c)
+            # v = v.permute(0, 2, 1, 3)  # -> (b, h, t, c)
+
+            # # F.scaled_dot_product_attention 直接返回计算结果，无需手动 reshape
+            # out = F.scaled_dot_product_attention(q, k, v)  # -> (b, h, t, c)
+
+            # # 关键：将输出变回 (b, t, h, c) 并 reshape，以匹配 else 分支的输出
+            # out = out.permute(0, 2, 1, 3).reshape(bs, n_ctx, -1) # -> (b, t, h*c)
+            out = torch_npu.npu_fusion_attention(
+                q, k, v, self.heads, "BSND", keep_prob=1.0, scale=scale)[0].reshape(bs, n_ctx, -1)
+        else:
+            weight = torch.einsum(
+                "bthc,bshc->bhts", q * scale, k * scale
+            )  # More stable with f16 than dividing afterwards
+            wdtype = weight.dtype
+            weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
+            out = torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
+
+        return out
+
+# Residual cross attention block module
+class ResidualCrossAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: Optional[torch.device],
+        dtype: Optional[torch.dtype],
+        n_data: Optional[int] = None,
+        data_width: Optional[int] = None,
+        width: int,  # Width of the input tensor
+        heads: int,  # Number of attention heads
+        init_scale: float,  # Initialization scale for weights
+        qkv_bias: bool,  # Whether to use bias in QKV layers
+        flash: bool = False  # Whether to use flash attention
+    ):
+        super().__init__()
+
+        if data_width is None:
+            data_width = width
+
+        self.attn = MultiheadCrossAttention(
+            device=device,
+            dtype=dtype,
+            n_data=n_data,
+            width=width,
+            heads=heads,
+            data_width=data_width,
+            init_scale=init_scale,
+            qkv_bias=qkv_bias,
+            flash=flash,
+        )
+        self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
+        self.ln_2 = nn.LayerNorm(data_width, device=device, dtype=dtype)
+        self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=init_scale)
+        self.ln_3 = nn.LayerNorm(width, device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor, data: torch.Tensor):
+        x = x + self.attn(self.ln_1(x), self.ln_2(data))
+        x = x + self.mlp(self.ln_3(x))
+        return x
+
+# MLP Module
+class MLP(nn.Module):
+    def __init__(self, *,
+                 device: Optional[torch.device],
+                 dtype: Optional[torch.dtype],
+                 width: int,
+                 init_scale: float):
+        super().__init__()
+        self.width = width
+        self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype)
+        self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype)
+        self.gelu = nn.GELU()
+        init_linear(self.c_fc, init_scale)
+        init_linear(self.c_proj, init_scale)
+
+    def forward(self, x):
+        return self.c_proj(self.gelu(self.c_fc(x)))
+
+# Transformer Module
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: Optional[torch.device],
+        dtype: Optional[torch.dtype],
+        layers: int,
+        use_checkpoint: bool = False,
+        n_ctx: int,  # Context size
+        width: int,  # Width of the input tensor
+        heads: int,  # Number of attention heads
+        init_scale: float,  # Initialization scale for weights
+        qkv_bias: bool,  # Whether to use bias in QKV layers
+        flash: bool = False  # Whether to use flash attention
+    ):
+        super().__init__()
+        self.n_ctx = n_ctx
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.ModuleList(
+            [
+                ResidualAttentionBlock(
+                    device=device,
+                    dtype=dtype,
+                    n_ctx=n_ctx,
+                    width=width,
+                    heads=heads,
+                    init_scale=init_scale,
+                    qkv_bias=qkv_bias,
+                    flash=flash,
+                    use_checkpoint=use_checkpoint
+                )
+                for _ in range(layers)
+            ]
+        )
+
+    def forward(self, x: torch.Tensor):
+        for block in self.resblocks:
+            x = block(x)
+        return x

mulres_infer/modules/miche/michelangelo/models/tsal/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

mulres_infer/modules/miche/michelangelo/models/tsal/asl_pl_module.py

@@ -0,0 +1,384 @@
+# -*- coding: utf-8 -*-
+
+from typing import List, Tuple, Dict, Optional
+from omegaconf import DictConfig
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.optim import lr_scheduler
+from typing import Union
+from functools import partial
+
+from modules.miche.michelangelo.utils import instantiate_from_config
+
+from .tsal_base import (
+    AlignedShapeAsLatentModule,
+    ShapeAsLatentModule,
+    Latent2MeshOutput,
+    AlignedMeshOutput
+)
+from modules.miche.michelangelo.models.tsal.inference_utils import extract_geometry
+import trimesh
+
+class AlignedShapeAsLatentPLModule(nn.Module):
+    def __init__(self, *,
+                 shape_module_cfg,
+                 aligned_module_cfg,
+                 loss_cfg,
+                 optimizer_cfg: Optional[DictConfig] = None,
+                 ckpt_path: Optional[str] = None,
+                 ignore_keys: Union[Tuple[str], List[str]] = ()):
+
+        super().__init__()
+
+        shape_model: ShapeAsLatentModule = instantiate_from_config(
+            shape_module_cfg, device=None, dtype=None
+        )
+        self.model: AlignedShapeAsLatentModule = instantiate_from_config(
+            aligned_module_cfg, shape_model=shape_model
+        )
+
+        self.loss = instantiate_from_config(loss_cfg)
+
+        self.optimizer_cfg = optimizer_cfg
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+    def set_shape_model_only(self):
+        self.model.set_shape_model_only()
+
+    @property
+    def latent_shape(self):
+        return self.model.shape_model.latent_shape
+
+    @property
+    def zero_rank(self):
+        if self._trainer:
+            zero_rank = self.trainer.local_rank == 0
+        else:
+            zero_rank = True
+
+        return zero_rank
+
+    def init_from_ckpt(self, path, ignore_keys=()):
+        state_dict = torch.load(path, map_location="cpu")["state_dict"]
+
+        keys = list(state_dict.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del state_dict[k]
+
+        missing, unexpected = self.load_state_dict(state_dict, strict=False)
+        print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
+        if len(missing) > 0:
+            print(f"Missing Keys: {missing}")
+            print(f"Unexpected Keys: {unexpected}")
+
+    def configure_optimizers(self) -> Tuple[List, List]:
+        lr = self.learning_rate
+
+        trainable_parameters = list(self.model.parameters())
+
+        if self.optimizer_cfg is None:
+            optimizers = [torch.optim.AdamW(trainable_parameters, lr=lr, betas=(0.9, 0.99), weight_decay=1e-3)]
+            schedulers = []
+        else:
+            optimizer = instantiate_from_config(self.optimizer_cfg.optimizer, params=trainable_parameters)
+            scheduler_func = instantiate_from_config(
+                self.optimizer_cfg.scheduler,
+                max_decay_steps=self.trainer.max_steps,
+                lr_max=lr
+            )
+            scheduler = {
+                "scheduler": lr_scheduler.LambdaLR(optimizer, lr_lambda=scheduler_func.schedule),
+                "interval": "step",
+                "frequency": 1
+            }
+            optimizers = [optimizer]
+            schedulers = [scheduler]
+
+        return optimizers, schedulers
+
+    def forward(self,
+                surface: torch.FloatTensor,
+                image: torch.FloatTensor,
+                text: torch.FloatTensor,
+                volume_queries: torch.FloatTensor):
+        # Args:
+        #     surface (torch.FloatTensor):
+        #     image (torch.FloatTensor):
+        #     text (torch.FloatTensor):
+        #     volume_queries (torch.FloatTensor):
+        # 
+        # Returns:
+
+        embed_outputs, shape_z = self.model(surface, image, text)
+
+        shape_zq, posterior = self.model.shape_model.encode_kl_embed(shape_z)
+        latents = self.model.shape_model.decode(shape_zq)
+        logits = self.model.shape_model.query_geometry(volume_queries, latents)
+
+        return embed_outputs, logits, posterior
+
+    def encode(self, surface: torch.FloatTensor, sample_posterior=True):
+
+        pc = surface[..., 0:3]
+        feats = surface[..., 3:6]
+
+        shape_embed, shape_zq, posterior = self.model.shape_model.encode(
+            pc=pc, feats=feats, sample_posterior=sample_posterior
+        )
+
+        return shape_zq
+
+    def encode_latents(self, surface: torch.FloatTensor, img_feats: Optional[torch.FloatTensor] = None):
+
+        pc = surface[..., 0:3]
+        feats = surface[..., 3:6]
+
+        shape_embed, shape_latents = self.model.shape_model.encode_latents(
+            pc=pc, feats=feats, img_feats=img_feats,
+        )
+        shape_embed = shape_embed.unsqueeze(1)
+        # 长度无关: latent 数可被 PointConditioner._expand_latents 扩容(如 1024), 只校验头 token
+        assert shape_embed.shape[1] == 1
+        cat_latents = torch.cat([shape_embed, shape_latents], dim=1)
+
+        return cat_latents
+
+    def recon(self, surface):
+        cat_latents = self.encode_latents(surface)
+        shape_latents = cat_latents[:, 1:]
+        shape_zq, posterior = self.model.shape_model.encode_kl_embed(shape_latents)
+
+        # decoding
+        latents = self.model.shape_model.decode(shape_zq)
+        geometric_func = partial(self.model.shape_model.query_geometry, latents=latents)
+
+        # reconstruction
+        mesh_v_f, has_surface = extract_geometry(
+            geometric_func=geometric_func,
+            device=surface.device,
+            batch_size=surface.shape[0],
+            bounds=(-1.25, -1.25, -1.25, 1.25, 1.25, 1.25),
+            octree_depth=7,
+            num_chunks=10000,
+        )
+        recon_mesh = trimesh.Trimesh(mesh_v_f[0][0], mesh_v_f[0][1])
+
+        return recon_mesh
+
+
+    def to_shape_latents(self, latents):
+
+        shape_zq, posterior = self.model.shape_model.encode_kl_embed(latents, sample_posterior = False)
+        return self.model.shape_model.decode(shape_zq)
+
+    def decode(self,
+               z_q,
+               bounds: Union[Tuple[float], List[float], float] = 1.1,
+               octree_depth: int = 7,
+               num_chunks: int = 10000) -> List[Latent2MeshOutput]:
+
+        latents = self.model.shape_model.decode(z_q)  # latents: [bs, num_latents, dim]
+        outputs = self.latent2mesh(latents, bounds=bounds, octree_depth=octree_depth, num_chunks=num_chunks)
+
+        return outputs
+
+    def training_step(self, batch: Dict[str, torch.FloatTensor],
+                      batch_idx: int, optimizer_idx: int = 0) -> torch.FloatTensor:
+        #Args:
+        #    batch (dict): the batch sample, and it contains:
+        #        - surface (torch.FloatTensor): [bs, n_surface, (3 + input_dim)]
+        #        - image (torch.FloatTensor): [bs, 3, 224, 224]
+        #        - text (torch.FloatTensor): [bs, num_templates, 77]
+        #        - geo_points (torch.FloatTensor): [bs, n_pts, (3 + 1)]
+        #
+        #    batch_idx (int):
+        #
+        #    optimizer_idx (int):
+        #
+        # Returns:
+        #    loss (torch.FloatTensor):
+
+        surface = batch["surface"]
+        image = batch["image"]
+        text = batch["text"]
+
+        volume_queries = batch["geo_points"][..., 0:3]
+        shape_labels = batch["geo_points"][..., -1]
+
+        embed_outputs, shape_logits, posteriors = self(surface, image, text, volume_queries)
+
+        aeloss, log_dict_ae = self.loss(
+            **embed_outputs,
+            posteriors=posteriors,
+            shape_logits=shape_logits,
+            shape_labels=shape_labels,
+            split="train"
+        )
+
+        self.log_dict(log_dict_ae, prog_bar=True, logger=True, batch_size=shape_logits.shape[0],
+                      sync_dist=False, rank_zero_only=True)
+
+        return aeloss
+
+    def validation_step(self, batch: Dict[str, torch.FloatTensor], batch_idx: int) -> torch.FloatTensor:
+
+        surface = batch["surface"]
+        image = batch["image"]
+        text = batch["text"]
+
+        volume_queries = batch["geo_points"][..., 0:3]
+        shape_labels = batch["geo_points"][..., -1]
+
+        embed_outputs, shape_logits, posteriors = self(surface, image, text, volume_queries)
+
+        aeloss, log_dict_ae = self.loss(
+            **embed_outputs,
+            posteriors=posteriors,
+            shape_logits=shape_logits,
+            shape_labels=shape_labels,
+            split="val"
+        )
+        self.log_dict(log_dict_ae, prog_bar=True, logger=True, batch_size=shape_logits.shape[0],
+                      sync_dist=False, rank_zero_only=True)
+
+        return aeloss
+
+    def visual_alignment(self,
+                         surface: torch.FloatTensor,
+                         image: torch.FloatTensor,
+                         text: torch.FloatTensor,
+                         description: Optional[List[str]] = None,
+                         bounds: Union[Tuple[float], List[float]] = (-1.25, -1.25, -1.25, 1.25, 1.25, 1.25),
+                         octree_depth: int = 7,
+                         num_chunks: int = 10000) -> List[AlignedMeshOutput]:
+
+        """
+
+        Args:
+            surface:
+            image:
+            text:
+            description:
+            bounds:
+            octree_depth:
+            num_chunks:
+
+        Returns:
+            mesh_outputs (List[AlignedMeshOutput]): the mesh outputs list.
+
+        """
+
+        outputs = []
+
+        device = surface.device
+        bs = surface.shape[0]
+
+        embed_outputs, shape_z = self.model(surface, image, text)
+
+        # calculate the similarity
+        image_embed = embed_outputs["image_embed"]
+        text_embed = embed_outputs["text_embed"]
+        shape_embed = embed_outputs["shape_embed"]
+
+        # normalized features
+        shape_embed = F.normalize(shape_embed, dim=-1, p=2)
+        text_embed = F.normalize(text_embed, dim=-1, p=2)
+        image_embed = F.normalize(image_embed, dim=-1, p=2)
+
+        # B x B
+        shape_text_similarity = (100.0 * shape_embed @ text_embed.T).softmax(dim=-1)
+
+        # B x B
+        shape_image_similarity = (100.0 * shape_embed @ image_embed.T).softmax(dim=-1)
+
+        # shape reconstruction
+        shape_zq, posterior = self.model.shape_model.encode_kl_embed(shape_z)
+        latents = self.model.shape_model.decode(shape_zq)
+        geometric_func = partial(self.model.shape_model.query_geometry, latents=latents)
+
+        # 2. decode geometry
+        mesh_v_f, has_surface = extract_geometry(
+            geometric_func=geometric_func,
+            device=device,
+            batch_size=bs,
+            bounds=bounds,
+            octree_depth=octree_depth,
+            num_chunks=num_chunks,
+            disable=not self.zero_rank
+        )
+
+        # 3. decode texture
+        for i, ((mesh_v, mesh_f), is_surface) in enumerate(zip(mesh_v_f, has_surface)):
+            if not is_surface:
+                outputs.append(None)
+                continue
+
+            out = AlignedMeshOutput()
+            out.mesh_v = mesh_v
+            out.mesh_f = mesh_f
+            out.surface = surface[i].cpu().numpy()
+            out.image = image[i].cpu().numpy()
+            if description is not None:
+                out.text = description[i]
+            out.shape_text_similarity = shape_text_similarity[i, i]
+            out.shape_image_similarity = shape_image_similarity[i, i]
+
+            outputs.append(out)
+
+        return outputs
+
+    def latent2mesh(self,
+                    latents: torch.FloatTensor,
+                    bounds: Union[Tuple[float], List[float], float] = 1.1,
+                    octree_depth: int = 7,
+                    num_chunks: int = 10000) -> List[Latent2MeshOutput]:
+
+        """
+
+        Args:
+            latents: [bs, num_latents, dim]
+            bounds:
+            octree_depth:
+            num_chunks:
+
+        Returns:
+            mesh_outputs (List[MeshOutput]): the mesh outputs list.
+
+        """
+
+        outputs = []
+
+        geometric_func = partial(self.model.shape_model.query_geometry, latents=latents)
+
+        # 2. decode geometry
+        device = latents.device
+        mesh_v_f, has_surface = extract_geometry(
+            geometric_func=geometric_func,
+            device=device,
+            batch_size=len(latents),
+            bounds=bounds,
+            octree_depth=octree_depth,
+            num_chunks=num_chunks,
+            disable=not self.zero_rank
+        )
+
+        # 3. decode texture
+        for i, ((mesh_v, mesh_f), is_surface) in enumerate(zip(mesh_v_f, has_surface)):
+            if not is_surface:
+                outputs.append(None)
+                continue
+
+            out = Latent2MeshOutput()
+            out.mesh_v = mesh_v
+            out.mesh_f = mesh_f
+
+            outputs.append(out)
+
+        return outputs

mulres_infer/modules/miche/michelangelo/models/tsal/clip_asl_module.py

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+
+import torch
+from torch import nn
+from einops import rearrange
+from transformers import CLIPModel
+
+from modules.miche.michelangelo.models.tsal.tsal_base import AlignedShapeAsLatentModule
+
+
+class CLIPAlignedShapeAsLatentModule(AlignedShapeAsLatentModule):
+
+    def __init__(self, *,
+                 shape_model,
+                 clip_model_version: str = "openai/clip-vit-large-patch14"):
+
+        super().__init__()
+
+        # self.clip_model: CLIPModel = CLIPModel.from_pretrained(clip_model_version)
+        # for params in self.clip_model.parameters():
+        #     params.requires_grad = False
+        self.clip_model = None
+        self.shape_model = shape_model
+        self.shape_projection = nn.Parameter(torch.empty(self.shape_model.width, self.shape_model.width))
+        # nn.init.normal_(self.shape_projection, std=self.shape_model.width ** -0.5)
+
+    def set_shape_model_only(self):
+        self.clip_model = None
+
+    def encode_shape_embed(self, surface, return_latents: bool = False):
+        """
+
+        Args:
+            surface (torch.FloatTensor): [bs, n, 3 + c]
+            return_latents (bool):
+
+        Returns:
+            x (torch.FloatTensor): [bs, projection_dim]
+            shape_latents (torch.FloatTensor): [bs, m, d]
+        """
+
+        pc = surface[..., 0:3]
+        feats = surface[..., 3:]
+
+        shape_embed, shape_latents = self.shape_model.encode_latents(pc, feats)
+        x = shape_embed @ self.shape_projection
+
+        if return_latents:
+            return x, shape_latents
+        else:
+            return x
+
+    def encode_image_embed(self, image):
+        """
+
+        Args:
+            image (torch.FloatTensor): [bs, 3, h, w]
+
+        Returns:
+            x (torch.FloatTensor): [bs, projection_dim]
+        """
+
+        x = self.clip_model.get_image_features(image)
+
+        return x
+
+    def encode_text_embed(self, text):
+        x = self.clip_model.get_text_features(text)
+        return x
+
+    def forward(self, surface, image, text):
+        """
+
+        Args:
+            surface (torch.FloatTensor):
+            image (torch.FloatTensor): [bs, 3, 224, 224]
+            text (torch.LongTensor): [bs, num_templates, 77]
+
+        Returns:
+            embed_outputs (dict): the embedding outputs, and it contains:
+                - image_embed (torch.FloatTensor):
+                - text_embed (torch.FloatTensor):
+                - shape_embed (torch.FloatTensor):
+                - logit_scale (float):
+        """
+
+        # # text embedding
+        # text_embed_all = []
+        # for i in range(text.shape[0]):
+        #     text_for_one_sample = text[i]
+        #     text_embed = self.encode_text_embed(text_for_one_sample)
+        #     text_embed = text_embed / text_embed.norm(dim=-1, keepdim=True)
+        #     text_embed = text_embed.mean(dim=0)
+        #     text_embed = text_embed / text_embed.norm(dim=-1, keepdim=True)
+        #     text_embed_all.append(text_embed)
+        # text_embed_all = torch.stack(text_embed_all)
+
+        b = text.shape[0]
+        text_tokens = rearrange(text, "b t l -> (b t) l")
+        text_embed = self.encode_text_embed(text_tokens)
+        text_embed = rearrange(text_embed, "(b t) d -> b t d", b=b)
+        text_embed = text_embed.mean(dim=1)
+        text_embed = text_embed / text_embed.norm(dim=-1, keepdim=True)
+
+        # image embedding
+        image_embed = self.encode_image_embed(image)
+
+        # shape embedding
+        shape_embed, shape_latents = self.encode_shape_embed(surface, return_latents=True)
+
+        embed_outputs = {
+            "image_embed": image_embed,
+            "text_embed": text_embed,
+            "shape_embed": shape_embed,
+            # "logit_scale": self.clip_model.logit_scale.exp()
+        }
+
+        return embed_outputs, shape_latents

mulres_infer/modules/miche/michelangelo/models/tsal/inference_utils.py

@@ -0,0 +1,76 @@
+# -*- coding: utf-8 -*-
+
+import torch
+from tqdm import tqdm
+from einops import repeat
+import numpy as np
+from typing import Callable, Tuple, List, Union, Optional
+from skimage import measure
+
+from modules.miche.michelangelo.graphics.primitives import generate_dense_grid_points
+
+
+@torch.no_grad()
+def extract_geometry(geometric_func: Callable,
+                     device: torch.device,
+                     batch_size: int = 1,
+                     bounds: Union[Tuple[float], List[float], float] = (-1.25, -1.25, -1.25, 1.25, 1.25, 1.25),
+                     octree_depth: int = 7,
+                     num_chunks: int = 10000,
+                     disable: bool = True):
+
+    # Args:
+    #     geometric_func:
+    #     device:
+    #     bounds:
+    #     octree_depth:
+    #     batch_size:
+    #     num_chunks:
+    #     disable:
+    # Returns:
+
+    if isinstance(bounds, float):
+        bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds]
+
+    bbox_min = np.array(bounds[0:3])
+    bbox_max = np.array(bounds[3:6])
+    bbox_size = bbox_max - bbox_min
+
+    xyz_samples, grid_size, length = generate_dense_grid_points(
+        bbox_min=bbox_min,
+        bbox_max=bbox_max,
+        octree_depth=octree_depth,
+        indexing="ij"
+    )
+    xyz_samples = torch.FloatTensor(xyz_samples)
+
+    batch_logits = []
+    for start in tqdm(range(0, xyz_samples.shape[0], num_chunks),
+                      desc="Implicit Function:", disable=disable, leave=False):
+        queries = xyz_samples[start: start + num_chunks, :].to(device)
+        batch_queries = repeat(queries, "p c -> b p c", b=batch_size)
+
+        logits = geometric_func(batch_queries)
+        batch_logits.append(logits.cpu())
+
+    grid_logits = torch.cat(batch_logits, dim=1).view((batch_size, grid_size[0], grid_size[1], grid_size[2])).numpy()
+
+    mesh_v_f = []
+    has_surface = np.zeros((batch_size,), dtype=np.bool_)
+    for i in range(batch_size):
+        try:
+            vertices, faces, normals, _ = measure.marching_cubes(grid_logits[i], 0, method="lewiner")
+            vertices = vertices / grid_size * bbox_size + bbox_min
+            # vertices[:, [0, 1]] = vertices[:, [1, 0]]
+            mesh_v_f.append((vertices.astype(np.float32), np.ascontiguousarray(faces)))
+            has_surface[i] = True
+
+        except ValueError:
+            mesh_v_f.append((None, None))
+            has_surface[i] = False
+
+        except RuntimeError:
+            mesh_v_f.append((None, None))
+            has_surface[i] = False
+
+    return mesh_v_f, has_surface

mulres_infer/modules/miche/michelangelo/models/tsal/loss.py

@@ -0,0 +1,130 @@
+# -*- coding: utf-8 -*-
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Optional
+
+from modules.miche.michelangelo.models.modules.distributions import DiagonalGaussianDistribution
+from modules.miche.michelangelo.utils import misc
+
+
+class ContrastKLNearFar(nn.Module):
+    def __init__(self,
+                 contrast_weight: float = 1.0,
+                 near_weight: float = 0.1,
+                 kl_weight: float = 1.0,
+                 num_near_samples: Optional[int] = None):
+
+        super().__init__()
+
+        self.labels = None
+        self.last_local_batch_size = None
+
+        self.contrast_weight = contrast_weight
+        self.near_weight = near_weight
+        self.kl_weight = kl_weight
+        self.num_near_samples = num_near_samples
+        self.geo_criterion = nn.BCEWithLogitsLoss()
+
+    def forward(self,
+                shape_embed: torch.FloatTensor,
+                text_embed: torch.FloatTensor,
+                image_embed: torch.FloatTensor,
+                logit_scale: torch.FloatTensor,
+                posteriors: Optional[DiagonalGaussianDistribution],
+                shape_logits: torch.FloatTensor,
+                shape_labels: torch.FloatTensor,
+                split: Optional[str] = "train", **kwargs):
+        
+        # shape_embed: torch.FloatTensor
+        # text_embed: torch.FloatTensor
+        # image_embed: torch.FloatTensor
+        # logit_scale: torch.FloatTensor
+        # posteriors: Optional[DiagonalGaussianDistribution]
+        # shape_logits: torch.FloatTensor
+        # shape_labels: torch.FloatTensor
+
+        local_batch_size = shape_embed.size(0)
+
+        if local_batch_size != self.last_local_batch_size:
+            self.labels = local_batch_size * misc.get_rank() + torch.arange(
+                local_batch_size, device=shape_embed.device
+            ).long()
+            self.last_local_batch_size = local_batch_size
+
+        # normalized features
+        shape_embed = F.normalize(shape_embed, dim=-1, p=2)
+        text_embed = F.normalize(text_embed, dim=-1, p=2)
+        image_embed = F.normalize(image_embed, dim=-1, p=2)
+
+        # gather features from all GPUs
+        shape_embed_all, text_embed_all, image_embed_all = misc.all_gather_batch(
+            [shape_embed, text_embed, image_embed]
+        )
+
+        # cosine similarity as logits
+        logits_per_shape_text = logit_scale * shape_embed @ text_embed_all.t()
+        logits_per_text_shape = logit_scale * text_embed @ shape_embed_all.t()
+        logits_per_shape_image = logit_scale * shape_embed @ image_embed_all.t()
+        logits_per_image_shape = logit_scale * image_embed @ shape_embed_all.t()
+        contrast_loss = (F.cross_entropy(logits_per_shape_text, self.labels) +
+                         F.cross_entropy(logits_per_text_shape, self.labels)) / 2 + \
+                        (F.cross_entropy(logits_per_shape_image, self.labels) +
+                         F.cross_entropy(logits_per_image_shape, self.labels)) / 2
+
+        # shape reconstruction
+        if self.num_near_samples is None:
+            num_vol = shape_logits.shape[1] // 2
+        else:
+            num_vol = shape_logits.shape[1] - self.num_near_samples
+
+        vol_logits = shape_logits[:, 0:num_vol]
+        vol_labels = shape_labels[:, 0:num_vol]
+
+        near_logits = shape_logits[:, num_vol:]
+        near_labels = shape_labels[:, num_vol:]
+
+        # occupancy loss
+        vol_bce = self.geo_criterion(vol_logits.float(), vol_labels.float())
+        near_bce = self.geo_criterion(near_logits.float(), near_labels.float())
+
+        if posteriors is None:
+            kl_loss = torch.tensor(0.0, dtype=vol_logits.dtype, device=vol_logits.device)
+        else:
+            kl_loss = posteriors.kl(dims=(1, 2))
+            kl_loss = torch.mean(kl_loss)
+
+        loss = vol_bce + near_bce * self.near_weight + kl_loss * self.kl_weight + contrast_loss * self.contrast_weight
+
+        # compute accuracy
+        with torch.no_grad():
+            pred = torch.argmax(logits_per_shape_text, dim=-1)
+            correct = pred.eq(self.labels).sum()
+            shape_text_acc = 100 * correct / local_batch_size
+
+            pred = torch.argmax(logits_per_shape_image, dim=-1)
+            correct = pred.eq(self.labels).sum()
+            shape_image_acc = 100 * correct / local_batch_size
+
+            preds = shape_logits >= 0
+            accuracy = (preds == shape_labels).float()
+            accuracy = accuracy.mean()
+
+            log = {
+                "{}/contrast".format(split): contrast_loss.clone().detach(),
+                "{}/near".format(split): near_bce.detach(),
+                "{}/far".format(split): vol_bce.detach(),
+                "{}/kl".format(split): kl_loss.detach(),
+                "{}/shape_text_acc".format(split): shape_text_acc,
+                "{}/shape_image_acc".format(split): shape_image_acc,
+                "{}/total_loss".format(split): loss.clone().detach(),
+                "{}/accuracy".format(split): accuracy,
+            }
+
+            if posteriors is not None:
+                log[f"{split}/mean"] = posteriors.mean.mean().detach()
+                log[f"{split}/std_mean"] = posteriors.std.mean().detach()
+                log[f"{split}/std_max"] = posteriors.std.max().detach()
+
+        return loss, log

mulres_infer/modules/miche/michelangelo/models/tsal/sal_perceiver.py

@@ -0,0 +1,459 @@
+# -*- coding: utf-8 -*-
+
+import torch
+import torch.nn as nn
+from typing import Optional
+from einops import repeat
+import math
+
+from modules.miche.michelangelo.models.modules import checkpoint
+from modules.miche.michelangelo.models.modules.embedder import FourierEmbedder
+from modules.miche.michelangelo.models.modules.distributions import DiagonalGaussianDistribution
+from modules.miche.michelangelo.models.modules.transformer_blocks import (
+    ResidualCrossAttentionBlock,
+    Transformer
+)
+
+from .tsal_base import ShapeAsLatentModule
+
+
+class CrossAttentionEncoder(nn.Module):
+
+    def __init__(self, *,
+                 device: Optional[torch.device],
+                 dtype: Optional[torch.dtype],
+                 num_latents: int,
+                 fourier_embedder: FourierEmbedder,
+                 point_feats: int,
+                 width: int,
+                 heads: int,
+                 layers: int,
+                 init_scale: float = 0.25,
+                 qkv_bias: bool = True,
+                 flash: bool = False,
+                 use_ln_post: bool = False,
+                 use_checkpoint: bool = False,
+                 img_feat_dim: int = 0):
+
+        super().__init__()
+
+        self.use_checkpoint = use_checkpoint
+        self.num_latents = num_latents
+
+        self.query = nn.Parameter(torch.randn((num_latents, width), device=device, dtype=dtype) * 0.02)
+
+        self.fourier_embedder = fourier_embedder
+        self.point_feats = point_feats
+        self.img_feat_dim = img_feat_dim
+        self.input_proj = nn.Linear(
+            self.fourier_embedder.out_dim + point_feats + img_feat_dim,
+            width,
+            device=device,
+            dtype=dtype,
+        )
+        self.cross_attn = ResidualCrossAttentionBlock(
+            device=device,
+            dtype=dtype,
+            width=width,
+            heads=heads,
+            init_scale=init_scale,
+            qkv_bias=qkv_bias,
+            flash=flash,
+        )
+
+        self.self_attn = Transformer(
+            device=device,
+            dtype=dtype,
+            n_ctx=num_latents,
+            width=width,
+            layers=layers,
+            heads=heads,
+            init_scale=init_scale,
+            qkv_bias=qkv_bias,
+            flash=flash,
+            use_checkpoint=False
+        )
+
+        if use_ln_post:
+            self.ln_post = nn.LayerNorm(width, dtype=dtype, device=device)
+        else:
+            self.ln_post = None
+
+    def _forward(self, pc, feats, img_feats=None):
+
+        # Args:
+        #     pc (torch.FloatTensor): [B, N, 3]
+        #     feats (torch.FloatTensor or None): [B, N, C]
+        #     img_feats (torch.FloatTensor or None): [B, N, img_feat_dim]
+
+        bs = pc.shape[0]
+
+        parts = [self.fourier_embedder(pc)]
+        if feats is not None:
+            parts.append(feats)
+        if self.img_feat_dim > 0:
+            if img_feats is None:
+                raise ValueError("img_feats is required when img_feat_dim > 0")
+            parts.append(img_feats)
+        data = self.input_proj(torch.cat(parts, dim=-1))
+
+        query = repeat(self.query, "m c -> b m c", b=bs)
+        latents = self.cross_attn(query, data)
+        latents = self.self_attn(latents)
+
+        if self.ln_post is not None:
+            latents = self.ln_post(latents)
+
+        return latents, pc
+
+    def forward(
+        self,
+        pc: torch.FloatTensor,
+        feats: Optional[torch.FloatTensor] = None,
+        img_feats: Optional[torch.FloatTensor] = None,
+    ):
+
+        # Args:
+        #     pc (torch.FloatTensor): [B, N, 3]
+        #     feats (torch.FloatTensor or None): [B, N, C]
+        #     img_feats (torch.FloatTensor or None): [B, N, img_feat_dim]
+
+        # NOTE: Michelangelo's custom CheckpointFunction.backward does
+        # `x.detach()` on every entry of ctx.input_tensors and cannot tolerate
+        # None values. So we MUST NOT pass `img_feats=None` into the checkpoint
+        # inputs tuple. Build the inputs based on which optional args are
+        # actually tensors.
+        inputs = [pc]
+        if feats is not None:
+            inputs.append(feats)
+        if img_feats is not None:
+            inputs.append(img_feats)
+
+        # _forward expects (pc, feats, img_feats). When some are None we cannot
+        # rely on positional defaults inside checkpoint, so wrap to restore them.
+        feats_is_none = feats is None
+        img_feats_is_none = img_feats is None
+
+        def _forward_wrapper(*args):
+            it = iter(args)
+            pc_arg = next(it)
+            feats_arg = None if feats_is_none else next(it)
+            img_feats_arg = None if img_feats_is_none else next(it)
+            return self._forward(pc_arg, feats_arg, img_feats_arg)
+
+        return checkpoint(_forward_wrapper, tuple(inputs), self.parameters(), self.use_checkpoint)
+
+
+class CrossAttentionDecoder(nn.Module):
+
+    def __init__(self, *,
+                 device: Optional[torch.device],
+                 dtype: Optional[torch.dtype],
+                 num_latents: int,
+                 out_channels: int,
+                 fourier_embedder: FourierEmbedder,
+                 width: int,
+                 heads: int,
+                 init_scale: float = 0.25,
+                 qkv_bias: bool = True,
+                 flash: bool = False,
+                 use_checkpoint: bool = False):
+
+        super().__init__()
+
+        self.use_checkpoint = use_checkpoint
+        self.fourier_embedder = fourier_embedder
+
+        self.query_proj = nn.Linear(self.fourier_embedder.out_dim, width, device=device, dtype=dtype)
+
+        self.cross_attn_decoder = ResidualCrossAttentionBlock(
+            device=device,
+            dtype=dtype,
+            n_data=num_latents,
+            width=width,
+            heads=heads,
+            init_scale=init_scale,
+            qkv_bias=qkv_bias,
+            flash=flash
+        )
+
+        self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
+        self.output_proj = nn.Linear(width, out_channels, device=device, dtype=dtype)
+
+    def _forward(self, queries: torch.FloatTensor, latents: torch.FloatTensor):
+        queries = self.query_proj(self.fourier_embedder(queries))
+        x = self.cross_attn_decoder(queries, latents)
+        x = self.ln_post(x)
+        x = self.output_proj(x)
+        return x
+
+    def forward(self, queries: torch.FloatTensor, latents: torch.FloatTensor):
+        return checkpoint(self._forward, (queries, latents), self.parameters(), self.use_checkpoint)
+
+
+class ShapeAsLatentPerceiver(ShapeAsLatentModule):
+    def __init__(self, *,
+                 device: Optional[torch.device],
+                 dtype: Optional[torch.dtype],
+                 num_latents: int,
+                 point_feats: int = 0,
+                 embed_dim: int = 0,
+                 num_freqs: int = 8,
+                 include_pi: bool = True,
+                 width: int,
+                 heads: int,
+                 num_encoder_layers: int,
+                 num_decoder_layers: int,
+                 init_scale: float = 0.25,
+                 qkv_bias: bool = True,
+                 flash: bool = False,
+                 use_ln_post: bool = False,
+                 use_checkpoint: bool = False,
+                 img_feat_dim: int = 0):
+
+        super().__init__()
+
+        self.use_checkpoint = use_checkpoint
+
+        self.num_latents = num_latents
+        self.fourier_embedder = FourierEmbedder(num_freqs=num_freqs, include_pi=include_pi)
+        self.img_feat_dim = img_feat_dim
+
+        init_scale = init_scale * math.sqrt(1.0 / width)
+        self.encoder = CrossAttentionEncoder(
+            device=device,
+            dtype=dtype,
+            fourier_embedder=self.fourier_embedder,
+            num_latents=num_latents,
+            point_feats=point_feats,
+            width=width,
+            heads=heads,
+            layers=num_encoder_layers,
+            init_scale=init_scale,
+            qkv_bias=qkv_bias,
+            flash=flash,
+            use_ln_post=use_ln_post,
+            use_checkpoint=use_checkpoint,
+            img_feat_dim=img_feat_dim,
+        )
+
+        self.embed_dim = embed_dim
+        if embed_dim > 0:
+            # VAE embed
+            self.pre_kl = nn.Linear(width, embed_dim * 2, device=device, dtype=dtype)
+            self.post_kl = nn.Linear(embed_dim, width, device=device, dtype=dtype)
+            self.latent_shape = (num_latents, embed_dim)
+        else:
+            self.latent_shape = (num_latents, width)
+
+        self.transformer = Transformer(
+            device=device,
+            dtype=dtype,
+            n_ctx=num_latents,
+            width=width,
+            layers=num_decoder_layers,
+            heads=heads,
+            init_scale=init_scale,
+            qkv_bias=qkv_bias,
+            flash=flash,
+            use_checkpoint=use_checkpoint
+        )
+
+        # geometry decoder
+        self.geo_decoder = CrossAttentionDecoder(
+            device=device,
+            dtype=dtype,
+            fourier_embedder=self.fourier_embedder,
+            out_channels=1,
+            num_latents=num_latents,
+            width=width,
+            heads=heads,
+            init_scale=init_scale,
+            qkv_bias=qkv_bias,
+            flash=flash,
+            use_checkpoint=use_checkpoint
+        )
+
+    def encode(self,
+               pc: torch.FloatTensor,
+               feats: Optional[torch.FloatTensor] = None,
+               sample_posterior: bool = True,
+               img_feats: Optional[torch.FloatTensor] = None):
+
+
+        # Args:
+        #     pc (torch.FloatTensor): [B, N, 3]
+        #     feats (torch.FloatTensor or None): [B, N, C]
+        #     sample_posterior (bool):
+
+        # Returns:
+        #     latents (torch.FloatTensor)
+        #     center_pos (torch.FloatTensor or None):
+        #     posterior (DiagonalGaussianDistribution or None):
+
+
+        latents, center_pos = self.encoder(pc, feats, img_feats=img_feats)
+
+        posterior = None
+        if self.embed_dim > 0:
+            moments = self.pre_kl(latents)
+            posterior = DiagonalGaussianDistribution(moments, feat_dim=-1)
+
+            if sample_posterior:
+                latents = posterior.sample()
+            else:
+                latents = posterior.mode()
+
+        return latents, center_pos, posterior
+
+    def decode(self, latents: torch.FloatTensor):
+        latents = self.post_kl(latents)
+        return self.transformer(latents)
+
+    def query_geometry(self, queries: torch.FloatTensor, latents: torch.FloatTensor):
+        logits = self.geo_decoder(queries, latents).squeeze(-1)
+        return logits
+
+    def forward(self,
+                pc: torch.FloatTensor,
+                feats: torch.FloatTensor,
+                volume_queries: torch.FloatTensor,
+                sample_posterior: bool = True):
+
+        # Args:
+        #     pc (torch.FloatTensor): [B, N, 3]
+        #     feats (torch.FloatTensor or None): [B, N, C]
+        #     volume_queries (torch.FloatTensor): [B, P, 3]
+        #     sample_posterior (bool):
+
+        # Returns:
+        #     logits (torch.FloatTensor): [B, P]
+        #     center_pos (torch.FloatTensor): [B, M, 3]
+        #     posterior (DiagonalGaussianDistribution or None).
+
+
+
+        latents, center_pos, posterior = self.encode(pc, feats, sample_posterior=sample_posterior)
+
+        latents = self.decode(latents)
+        logits = self.query_geometry(volume_queries, latents)
+
+        return logits, center_pos, posterior
+
+
+class AlignedShapeLatentPerceiver(ShapeAsLatentPerceiver):
+
+    def __init__(self, *,
+                 device: Optional[torch.device],
+                 dtype: Optional[torch.dtype],
+                 num_latents: int,
+                 point_feats: int = 0,
+                 embed_dim: int = 0,
+                 num_freqs: int = 8,
+                 include_pi: bool = True,
+                 width: int,
+                 heads: int,
+                 num_encoder_layers: int,
+                 num_decoder_layers: int,
+                 init_scale: float = 0.25,
+                 qkv_bias: bool = True,
+                 flash: bool = False,
+                 use_ln_post: bool = False,
+                 use_checkpoint: bool = False,
+                 img_feat_dim: int = 0):
+
+        super().__init__(
+            device=device,
+            dtype=dtype,
+            num_latents=1 + num_latents,
+            point_feats=point_feats,
+            embed_dim=embed_dim,
+            num_freqs=num_freqs,
+            include_pi=include_pi,
+            width=width,
+            heads=heads,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            init_scale=init_scale,
+            qkv_bias=qkv_bias,
+            flash=flash,
+            use_ln_post=use_ln_post,
+            use_checkpoint=use_checkpoint,
+            img_feat_dim=img_feat_dim,
+        )
+
+        self.width = width
+
+    def encode(self,
+               pc: torch.FloatTensor,
+               feats: Optional[torch.FloatTensor] = None,
+               sample_posterior: bool = True,
+               img_feats: Optional[torch.FloatTensor] = None):
+
+        # Args:
+        #     pc (torch.FloatTensor): [B, N, 3]
+        #     feats (torch.FloatTensor or None): [B, N, c]
+        #     sample_posterior (bool):
+
+        # Returns:
+        #     shape_embed (torch.FloatTensor)
+        #     kl_embed (torch.FloatTensor):
+        #     posterior (DiagonalGaussianDistribution or None):
+
+
+        shape_embed, latents = self.encode_latents(pc, feats, img_feats=img_feats)
+        kl_embed, posterior = self.encode_kl_embed(latents, sample_posterior)
+
+        return shape_embed, kl_embed, posterior
+
+    def encode_latents(self,
+                       pc: torch.FloatTensor,
+                       feats: Optional[torch.FloatTensor] = None,
+                       img_feats: Optional[torch.FloatTensor] = None):
+
+        x, _ = self.encoder(pc, feats, img_feats=img_feats)
+
+        shape_embed = x[:, 0]
+        latents = x[:, 1:]
+
+        return shape_embed, latents
+
+    def encode_kl_embed(self, latents: torch.FloatTensor, sample_posterior: bool = True):
+        posterior = None
+        if self.embed_dim > 0:
+            moments = self.pre_kl(latents)
+            posterior = DiagonalGaussianDistribution(moments, feat_dim=-1)
+
+            if sample_posterior:
+                kl_embed = posterior.sample()
+            else:
+                kl_embed = posterior.mode()
+        else:
+            kl_embed = latents
+
+        return kl_embed, posterior
+
+    def forward(self,
+                pc: torch.FloatTensor,
+                feats: torch.FloatTensor,
+                volume_queries: torch.FloatTensor,
+                sample_posterior: bool = True):
+ 
+        # Args:
+        #     pc (torch.FloatTensor): [B, N, 3]
+        #     feats (torch.FloatTensor or None): [B, N, C]
+        #     volume_queries (torch.FloatTensor): [B, P, 3]
+        #     sample_posterior (bool):
+
+        # Returns:
+        #     shape_embed (torch.FloatTensor): [B, projection_dim]
+        #     logits (torch.FloatTensor): [B, M]
+        #     posterior (DiagonalGaussianDistribution or None).
+
+
+        shape_embed, kl_embed, posterior = self.encode(pc, feats, sample_posterior=sample_posterior)
+
+        latents = self.decode(kl_embed)
+        logits = self.query_geometry(volume_queries, latents)
+
+        return shape_embed, logits, posterior

mulres_infer/modules/miche/michelangelo/models/tsal/tsal_base.py

@@ -0,0 +1,125 @@
+# -*- coding: utf-8 -*-
+
+import torch.nn as nn
+from typing import Tuple, List, Optional
+
+# Base class for output of Point to Mesh transformation
+class Point2MeshOutput(object):
+    def __init__(self):
+        self.mesh_v = None  # Vertices of the mesh
+        self.mesh_f = None  # Faces of the mesh
+        self.center = None  # Center of the mesh
+        self.pc = None  # Point cloud data
+
+
+# Base class for output of Latent to Mesh transformation
+class Latent2MeshOutput(object):
+    def __init__(self):
+        self.mesh_v = None  # Vertices of the mesh
+        self.mesh_f = None  # Faces of the mesh
+
+
+# Base class for output of Aligned Mesh transformation
+class AlignedMeshOutput(object):
+    def __init__(self):
+        self.mesh_v = None  # Vertices of the mesh
+        self.mesh_f = None  # Faces of the mesh
+        self.surface = None  # Surface data
+        self.image = None  # Aligned image data
+        self.text: Optional[str] = None  # Aligned text data
+        self.shape_text_similarity: Optional[float] = None  # Similarity between shape and text
+        self.shape_image_similarity: Optional[float] = None  # Similarity between shape and image
+
+
+# Base class for Shape as Latent with Point to Mesh transformation module
+class ShapeAsLatentPLModule(nn.Module):
+    latent_shape: Tuple[int]  # Shape of the latent space
+    
+    def encode(self, surface, *args, **kwargs):
+        raise NotImplementedError
+
+    def decode(self, z_q, *args, **kwargs):
+        raise NotImplementedError
+
+    def latent2mesh(self, latents, *args, **kwargs) -> List[Latent2MeshOutput]:
+        raise NotImplementedError
+
+    def point2mesh(self, *args, **kwargs) -> List[Point2MeshOutput]:
+        raise NotImplementedError
+
+
+# Base class for Shape as Latent module
+class ShapeAsLatentModule(nn.Module):
+    latent_shape: Tuple[int, int]  # Shape of the latent space
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def decode(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def query_geometry(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+# Base class for Aligned Shape as Latent with Point to Mesh transformation module
+class AlignedShapeAsLatentPLModule(nn.Module):
+    latent_shape: Tuple[int]  # Shape of the latent space
+
+    def set_shape_model_only(self):
+        raise NotImplementedError
+
+    def encode(self, surface, *args, **kwargs):
+        raise NotImplementedError
+
+    def decode(self, z_q, *args, **kwargs):
+        raise NotImplementedError
+
+    def latent2mesh(self, latents, *args, **kwargs) -> List[Latent2MeshOutput]:
+        raise NotImplementedError
+
+    def point2mesh(self, *args, **kwargs) -> List[Point2MeshOutput]:
+        raise NotImplementedError
+
+
+# Base class for Aligned Shape as Latent module
+class AlignedShapeAsLatentModule(nn.Module):
+    shape_model: ShapeAsLatentModule  # Shape model module
+    latent_shape: Tuple[int, int]  # Shape of the latent space
+
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def set_shape_model_only(self):
+        raise NotImplementedError
+
+    def encode_image_embed(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def encode_text_embed(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def encode_shape_embed(self, *args, **kwargs):
+        raise NotImplementedError
+
+# Base class for Textured Shape as Latent module
+class TexturedShapeAsLatentModule(nn.Module):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def decode(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def query_geometry(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def query_color(self, *args, **kwargs):
+        raise NotImplementedError

mulres_infer/modules/miche/michelangelo/utils/__init__.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+
+from .misc import instantiate_from_config

mulres_infer/modules/miche/michelangelo/utils/misc.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+
+import importlib
+
+import torch
+import torch.distributed as dist
+
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def get_obj_from_config(config):
+    if "target" not in config:
+        raise KeyError("Expected key `target` to instantiate.")
+
+    return get_obj_from_str(config["target"])
+
+
+def instantiate_from_config(config, **kwargs):
+    if "target" not in config:
+        raise KeyError("Expected key `target` to instantiate.")
+
+    cls = get_obj_from_str(config["target"])
+
+    params = config.get("params", dict())
+    # params.update(kwargs)
+    # instance = cls(**params)
+    kwargs.update(params)
+    instance = cls(**kwargs)
+
+    return instance
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def all_gather_batch(tensors):
+    """
+    Performs all_gather operation on the provided tensors.
+    """
+    # Queue the gathered tensors
+    world_size = get_world_size()
+    # There is no need for reduction in the single-proc case
+    if world_size == 1:
+        return tensors
+    tensor_list = []
+    output_tensor = []
+    for tensor in tensors:
+        tensor_all = [torch.ones_like(tensor) for _ in range(world_size)]
+        dist.all_gather(
+            tensor_all,
+            tensor,
+            async_op=False  # performance opt
+        )
+
+        tensor_list.append(tensor_all)
+
+    for tensor_all in tensor_list:
+        output_tensor.append(torch.cat(tensor_all, dim=0))
+    return output_tensor

mulres_infer/modules/miche/shapevae-256.yaml

@@ -0,0 +1,46 @@
+model:
+  target: modules.miche.michelangelo.models.tsal.asl_pl_module.AlignedShapeAsLatentPLModule
+  params:
+    shape_module_cfg:
+      target: modules.miche.michelangelo.models.tsal.sal_perceiver.AlignedShapeLatentPerceiver
+      params:
+        num_latents: 256
+        embed_dim: 64
+        point_feats: 3   # normal
+        num_freqs: 8
+        include_pi: false
+        heads: 12
+        width: 768
+        num_encoder_layers: 8
+        num_decoder_layers: 16
+        use_ln_post: true
+        init_scale: 0.25
+        qkv_bias: false
+        use_checkpoint: true
+    aligned_module_cfg:
+      target: modules.miche.michelangelo.models.tsal.clip_asl_module.CLIPAlignedShapeAsLatentModule
+      params:
+        clip_model_version: "./checkpoints/clip/clip-vit-large-patch14"
+
+    loss_cfg:
+      target: modules.miche.michelangelo.models.tsal.loss.ContrastKLNearFar
+      params:
+        contrast_weight: 0.1
+        near_weight: 0.1
+        kl_weight: 0.001
+
+    optimizer_cfg:
+      optimizer:
+        target: torch.optim.AdamW
+        params:
+          betas: [0.9, 0.99]
+          eps: 1.e-6
+          weight_decay: 1.e-2
+
+      scheduler:
+        target: modules.miche.michelangelo.utils.trainings.lr_scheduler.LambdaWarmUpCosineFactorScheduler
+        params:
+          warm_up_steps: 5000
+          f_start: 1.e-6
+          f_min: 1.e-3
+          f_max: 1.0

mulres_infer/modules/octree_codec.py

@@ -0,0 +1,981 @@
+"""fix3 tokenizer(已融合 numba 加速实现, 原 *_fast.py 并入本文件)。
+
+encode_mesh / decode_tokens 走 numpy 向量化 + numba 枚举, 比纯 python 版快约 5x。
+quantized_mesh 与 decode 输出和纯 python 版完全一致; token 的 face 块集合一致,
+仅少数平局块顺序不同(本实现补了 (f2.y,f2.x) 决定性平局键, 比原版更可复现)。
+
+原文件头:
+
+fix3: 连接矩阵改为「octant 绝对原位 8x8 + 因子分解(按行)」编码，保留 F。
+
+相对 fix2 的核心变化
+--------------------
+1. 不再按 F 排序、不再受 max_children 限制。每个连接矩阵(intra / inter)统一是
+   8x8 的「octant 绝对原位」矩阵：行/列 = octant 0..7 的固定语义。
+   → 天然 mc=8 无损，"超标 / 合并 / 退化分摊" 这套概念全部消失。
+2. 8x8 矩阵不再编码成单个 token(那是 2^64 词表)，而是按行拆成 8 个「行子码」，
+   每个行子码 = 该行 8 个 bit ∈ [0, 256)。intra / inter 各 8 个子码。
+3. **逻辑序列长度不变**：每个 face 仍是 18 个「位置」(9 coord + 3 F + 3 intra + 3 inter)，
+   只是 intra / inter 这 6 个位置每个携带 8 个子码(给因子分解头：同一位置 8 个并行 256-way 头)。
+   coord / F 仍是单子码位置。tokenizer.encode 返回 List[List[int]]：每个内层 list = 一个位置的子码。
+4. 保留 F(标识 active octant)。重建时用 F mask 掉非 active 行列(M &= F⊗F)，消除歧义。
+
+词表(极小)：pad(1) + coord(resolution) + F(256) + intra_row(256) + inter_row(256)。
+"""
+import numpy as np
+import trimesh
+from tqdm import tqdm
+import os
+import sys
+import glob
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from pathlib import Path
+from utils.data_utils import undiscretize, discretize, normalize, clean_mesh
+from utils.mesh_metric import compute_mesh_metrics
+from modules.fast_octree_level import _triangles_csr
+import itertools
+from dataclasses import dataclass
+
+NUM_OCTANTS = 8
+
+# 预计算 24 个节点(3 槽位 x 8 octant)中选 3 个的所有组合，用于向量化找三角形
+_TRI_INDICES = np.array(list(itertools.combinations(range(24), 3)))
+_IDX_I = _TRI_INDICES[:, 0]
+_IDX_J = _TRI_INDICES[:, 1]
+_IDX_K = _TRI_INDICES[:, 2]
+
+
+# ==========================================
+# Tokenizer：octant 原位 + 按行子码
+# ==========================================
+class DynamicMeshTokenizer:
+    """
+    encode 输出格式(fine)：List[List[int]]，每个内层 list 是一个「逻辑位置」的子码。
+      每 face 18 个位置： [v0x][v0y][v0z][v1x][v1y][v1z][v2x][v2y][v2z]  (9 个 coord, 各 1 子码)
+                          [F0][F1][F2]                                   (3 个 F,     各 1 子码)
+                          [M00 的 8 行][M11..][M22..]                    (3 个 intra, 各 8 子码)
+                          [M01 的 8 行][M02..][M12..]                    (3 个 inter, 各 8 子码)
+    coord / F 位置长度 1；intra / inter 位置长度 8(行子码)。逻辑序列长度 = 18/face，与 fix2 相同。
+    """
+    def __init__(self, resolution, max_children=NUM_OCTANTS):
+        self.resolution = resolution
+        self.max_children = NUM_OCTANTS  # octant 原位下固定为 8；保留字段仅为兼容 config
+
+        self.pad_id = 0
+        self.coord_offset = 1
+        self.f_offset = self.coord_offset + self.resolution
+        self.row_vocab = 1 << NUM_OCTANTS                       # 2^8 = 256
+        self.f_vocab = 1 << NUM_OCTANTS                         # F 也是 8bit
+        self.intra_row_offset = self.f_offset + self.f_vocab
+        self.inter_row_offset = self.intra_row_offset + self.row_vocab
+        self.vocab_size = self.inter_row_offset + self.row_vocab
+
+    # ---- bit <-> int ----
+    def _row_to_int(self, bits):
+        v = 0
+        for i, b in enumerate(bits):
+            if b:
+                v |= (1 << i)
+        return v
+
+    def _int_to_row(self, val):
+        return [(val >> i) & 1 for i in range(NUM_OCTANTS)]
+
+    def _mat_to_rows(self, M, offset):
+        return [self._row_to_int(M[r]) + offset for r in range(NUM_OCTANTS)]
+
+    def _rows_to_mat(self, rows, offset):
+        return [self._int_to_row(rows[r] - offset) for r in range(NUM_OCTANTS)]
+
+    # ------------------------------------
+    # 细分网格 (Fine Mesh) —— 扁平 1D int 布局，每 face 60 个 token：
+    #   9 coord + 3 F + 6 矩阵 x 8 行子码
+    #   顺序: [coord(9)]  [F0 F1 F2]  [M00 M11 M22 各 8 行]  [M01 M02 M12 各 8 行]
+    #   F: 每个 1 个 token；intra/inter: 每个矩阵按行拆成 8 个行子码(每行 8bit)
+    #   所有真实 token 都带 offset(>=1)，0 仅作 trailing pad，所以 F 不需要复制避坑。
+    # 与现有 dataset/collate(1D int, pad_id=0) 兼容；模型侧整理成 9 个 refine 位置/face
+    #   (F 位置 1 子码、M 位置 8 行子码 sum 成一个 embedding token)，逻辑序列 18/face 不变。
+    # ------------------------------------
+    TOKENS_PER_FACE = 9 + 3 + 6 * NUM_OCTANTS   # 60
+    REFINE_POS_PER_FACE = 9                       # 模型逻辑: 3 F + 3 intra + 3 inter
+
+    def encode(self, faces_data, use_hourglass=True):
+        seq = []
+        for face in faces_data:
+            v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12 = face
+            for v in (v0, v1, v2):
+                seq.extend([v[0] + self.coord_offset, v[1] + self.coord_offset, v[2] + self.coord_offset])
+            for F in (F0, F1, F2):
+                seq.append(self._row_to_int(F) + self.f_offset)        # F: 1 个 token
+            for M in (M00, M11, M22):
+                seq.extend(self._mat_to_rows(M, self.intra_row_offset))  # 各 8 行
+            for M in (M01, M02, M12):
+                seq.extend(self._mat_to_rows(M, self.inter_row_offset))  # 各 8 行
+        return seq
+
+    def decode(self, seq):
+        faces_data = []
+        B = self.TOKENS_PER_FACE  # 60
+        clean_seq = [t for t in seq if t != self.pad_id]
+        co = self.coord_offset
+        for i in range(0, len(clean_seq) - B + 1, B):
+            ch = clean_seq[i:i + B]
+            cv = [ch[k] - co for k in range(9)]
+            v0 = (cv[0], cv[1], cv[2]); v1 = (cv[3], cv[4], cv[5]); v2 = (cv[6], cv[7], cv[8])
+            F0 = self._int_to_row(ch[9] - self.f_offset)
+            F1 = self._int_to_row(ch[10] - self.f_offset)
+            F2 = self._int_to_row(ch[11] - self.f_offset)
+            M00 = self._rows_to_mat(ch[12:20], self.intra_row_offset)
+            M11 = self._rows_to_mat(ch[20:28], self.intra_row_offset)
+            M22 = self._rows_to_mat(ch[28:36], self.intra_row_offset)
+            M01 = self._rows_to_mat(ch[36:44], self.inter_row_offset)
+            M02 = self._rows_to_mat(ch[44:52], self.inter_row_offset)
+            M12 = self._rows_to_mat(ch[52:60], self.inter_row_offset)
+            faces_data.append((v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12))
+        return faces_data
+
+    # ------------------------------------
+    # 粗糙网格 (Raw Mesh)：仅顶点坐标，每 face 9 个单子码位置
+    # ------------------------------------
+    def encode_raw(self, raw_faces_data, use_hourglass=True):
+        seq = []
+        for face in raw_faces_data:
+            v0, v1, v2 = face
+            for v in (v0, v1, v2):
+                seq.append([v[0] + self.coord_offset])
+                seq.append([v[1] + self.coord_offset])
+                seq.append([v[2] + self.coord_offset])
+        return seq
+
+    def decode_raw(self, seq):
+        faces_data = []
+        clean_seq = [p for p in seq if not (len(p) == 1 and p[0] == self.pad_id)]
+        for i in range(0, len(clean_seq) - 9 + 1, 9):
+            chunk = clean_seq[i:i + 9]
+            if len(chunk) < 9:
+                break
+            cv = [chunk[k][0] - self.coord_offset for k in range(9)]
+            faces_data.append(((cv[0], cv[1], cv[2]), (cv[3], cv[4], cv[5]), (cv[6], cv[7], cv[8])))
+        return faces_data
+
+
+# ==========================================
+# Builder：直接构建 octant 原位 8x8 矩阵
+# ==========================================
+class OctreeFaceTokenBuilder:
+    def __init__(self, tree, max_children=NUM_OCTANTS):
+        self.tree = tree
+        self.max_children = NUM_OCTANTS  # octant 原位固定 8
+
+    def get_level_info(self, depth):
+        nodes_at_depth = {}
+
+        def collect_nodes(current, target_depth):
+            if current.depth == target_depth:
+                nodes_at_depth[current.coord] = current
+                return
+            for child in current.children:
+                if child:
+                    collect_nodes(child, target_depth)
+
+        collect_nodes(self.tree.root, depth)
+
+        coords = sorted(nodes_at_depth.keys(), key=lambda x: (x[2], x[1], x[0]))
+        coord_to_idx = {c: i for i, c in enumerate(coords)}
+
+        adj = {i: set() for i in range(len(coords))}
+        for i, c in enumerate(coords):
+            node = nodes_at_depth[c]
+            all_neighbors = node.intra_neighbors | node.inter_neighbors
+            for neighbor_coord in all_neighbors:
+                if neighbor_coord in coord_to_idx:
+                    j = coord_to_idx[neighbor_coord]
+                    if i < j:
+                        adj[i].add(j)
+
+        raw_faces_set = set()
+        covered_edges = set()
+
+        # 1. 3-clique (三角形)
+        for u in range(len(coords)):
+            for v in adj[u]:
+                common_neighbors = adj[u] & adj[v]
+                for w in common_neighbors:
+                    face_coords = tuple(sorted([coords[u], coords[v], coords[w]]))
+                    raw_faces_set.add(face_coords)
+                    covered_edges.add(tuple(sorted([face_coords[0], face_coords[1]])))
+                    covered_edges.add(tuple(sorted([face_coords[1], face_coords[2]])))
+                    covered_edges.add(tuple(sorted([face_coords[0], face_coords[2]])))
+
+        # 2. 未被三角形覆盖的独立边 -> 退化面片 (A, B, B)
+        for u in range(len(coords)):
+            for v in adj[u]:
+                edge = tuple(sorted([coords[u], coords[v]]))
+                if edge not in covered_edges:
+                    raw_faces_set.add((edge[0], edge[1], edge[1]))
+
+        # 3. 孤立顶点 -> (A, A, A)
+        for u in range(len(coords)):
+            has_edges = len(adj[u]) > 0 or any(u in adj[i] for i in range(u))
+            if not has_edges:
+                raw_faces_set.add((coords[u], coords[u], coords[u]))
+
+        raw_faces_at_depth = list(raw_faces_set)
+        raw_faces_at_depth.sort(key=lambda f: (f[0][2], f[0][1], f[0][0],
+                                               f[1][2], f[1][1], f[1][0],
+                                               f[2][2]))
+        return raw_faces_at_depth, nodes_at_depth
+
+    def _active_octants(self, coord, n, nodes_map):
+        """该节点哪些 octant 子节点存在(occupancy)。
+        用 children 占据而非"有连接"标记 —— 因为「有连接 ⟹ 必有子节点」，occupancy 是
+        "有连接"的超集，且正好等于 raw 的 get_level_info 收集的 depth+1 节点集；
+        这样「有几何但无边」的孤立子节点也会被 F 标记，fine 重构能补回孤立点，与 raw 完全一致。
+        M(连接矩阵)不变，故三角形/边重构不受影响。"""
+        return set(i for i in range(NUM_OCTANTS) if n.children[i] is not None)
+
+    def encode_level_tokens(self, depth):
+        raw_faces, nodes_map = self.get_level_info(depth)
+        empty = np.zeros((NUM_OCTANTS, NUM_OCTANTS), dtype=np.int8)
+
+        # F：每个节点的 active octant 掩码 (8bit)
+        node_F = {}
+        for coord, n in nodes_map.items():
+            active = self._active_octants(coord, n, nodes_map)
+            node_F[coord] = [1 if i in active else 0 for i in range(NUM_OCTANTS)]
+
+        intra_cache = {}
+
+        def intra8(coord):
+            if coord not in intra_cache:
+                M = (nodes_map[coord].intra_matrix != 0).astype(np.int8)  # 8x8 octant 原位、对称
+                intra_cache[coord] = M.tolist()
+            return intra_cache[coord]
+
+        def inter8(a, b):
+            if a == b:
+                return intra8(a)
+            na = nodes_map[a]
+            M = na.inter_matrices.get(b)
+            if M is None:
+                nb = nodes_map[b]
+                Mt = nb.inter_matrices.get(a)
+                M = Mt.T if Mt is not None else empty
+            return (np.asarray(M) != 0).astype(np.int8).tolist()
+
+        raw_tokens = []
+        output_tokens = []
+        for face in raw_faces:
+            v0_c, v1_c, v2_c = face
+            raw_tokens.append([v0_c, v1_c, v2_c])
+
+            F0, F1, F2 = node_F[v0_c], node_F[v1_c], node_F[v2_c]
+            M00, M11, M22 = intra8(v0_c), intra8(v1_c), intra8(v2_c)
+            M01 = inter8(v0_c, v1_c)
+            M02 = inter8(v0_c, v2_c)
+            M12 = inter8(v1_c, v2_c)
+
+            output_tokens.append((v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12))
+        return raw_tokens, output_tokens
+
+
+class OctreeNode:
+    def __init__(self, x, y, z, depth):
+        self.coord = (x, y, z)
+        self.depth = depth
+        self.children = [None] * 8
+        self.intra_neighbors = set()
+        self.inter_neighbors = set()
+        self.intra_matrix = np.zeros((8, 8), dtype=np.int8)
+        self.inter_matrices = {}
+
+
+class MatrixOctree:
+    def __init__(self, max_depth=7):
+        self.max_depth = max_depth
+        self.root = None
+
+    def _get_parent_coord(self, coord):
+        return (coord[0] >> 1, coord[1] >> 1, coord[2] >> 1)
+
+    def _get_local_index(self, coord):
+        return ((coord[2] & 1) << 2) | ((coord[1] & 1) << 1) | (coord[0] & 1)
+
+    def build(self, vertices, faces):
+        resolution = 2 ** self.max_depth
+        quantized = discretize(vertices, num_discrete=resolution)
+
+        leaf_nodes_map = {}
+        for i in range(len(vertices)):
+            coord = tuple(quantized[i])
+            if coord not in leaf_nodes_map:
+                leaf_nodes_map[coord] = OctreeNode(coord[0], coord[1], coord[2], self.max_depth)
+
+        for face in faces:
+            coords = [tuple(quantized[i]) for i in face]
+            for i in range(3):
+                u, v = coords[i], coords[(i + 1) % 3]
+                if u == v:
+                    continue
+                parent_u = self._get_parent_coord(u)
+                parent_v = self._get_parent_coord(v)
+                node_u = leaf_nodes_map[u]
+                node_v = leaf_nodes_map[v]
+                if parent_u == parent_v:
+                    node_u.intra_neighbors.add(v)
+                    node_v.intra_neighbors.add(u)
+                else:
+                    node_u.inter_neighbors.add(v)
+                    node_v.inter_neighbors.add(u)
+
+        current_layer_nodes = list(leaf_nodes_map.values())
+
+        for d in range(self.max_depth, 0, -1):
+            parent_layer_map = {}
+            for child in current_layer_nodes:
+                parent_coord = self._get_parent_coord(child.coord)
+                if parent_coord not in parent_layer_map:
+                    parent_layer_map[parent_coord] = OctreeNode(*parent_coord, d - 1)
+                parent_node = parent_layer_map[parent_coord]
+                idx = self._get_local_index(child.coord)
+                parent_node.children[idx] = child
+
+            for child in current_layer_nodes:
+                parent_coord = self._get_parent_coord(child.coord)
+                parent_node = parent_layer_map[parent_coord]
+                child_idx = self._get_local_index(child.coord)
+
+                for sibling_coord in child.intra_neighbors:
+                    sibling_idx = self._get_local_index(sibling_coord)
+                    parent_node.intra_matrix[child_idx, sibling_idx] = 1
+
+                for neighbor_coord in child.inter_neighbors:
+                    neighbor_parent_coord = self._get_parent_coord(neighbor_coord)
+                    if neighbor_parent_coord not in parent_node.inter_matrices:
+                        parent_node.inter_matrices[neighbor_parent_coord] = np.zeros((8, 8), dtype=np.int8)
+                    neighbor_child_idx = self._get_local_index(neighbor_coord)
+                    parent_node.inter_matrices[neighbor_parent_coord][child_idx, neighbor_child_idx] = 1
+
+                    grandparent_p = self._get_parent_coord(parent_coord)
+                    grandparent_n = self._get_parent_coord(neighbor_parent_coord)
+                    if grandparent_p == grandparent_n:
+                        parent_node.intra_neighbors.add(neighbor_parent_coord)
+                    else:
+                        parent_node.inter_neighbors.add(neighbor_parent_coord)
+
+            current_layer_nodes = list(parent_layer_map.values())
+        if current_layer_nodes:
+            self.root = current_layer_nodes[0]
+
+
+def reconstruct_next_level_optimized(output_tokens, resolution, max_children=NUM_OCTANTS):
+    """octant 绝对原位重建：M 已是 8x8、行列=octant；用 F mask 掉非 active 行列。"""
+    indices = np.arange(8)
+    offsets = np.stack([indices & 1, (indices >> 1) & 1, (indices >> 2) & 1], axis=1).astype(np.int32)
+
+    tri_blocks = []          # 逐 face 向量化收集的三角形 (N,3,3)，最后统一去重
+    edge_keys = set()        # 全局所有边: canonical (key_a, key_b)
+    tri_edge_keys = set()    # 被三角形覆盖的边(全局)
+    active_keys = set()      # 全局 active 节点 key
+    key_to_coord = {}        # key(int 三元组) -> np.int32 坐标(3,)
+
+    for f_token in output_tokens:
+        v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12 = f_token
+        fa = [np.asarray(F0, dtype=bool), np.asarray(F1, dtype=bool), np.asarray(F2, dtype=bool)]
+
+        def intra(F, M):
+            return np.asarray(M, dtype=bool) & np.outer(F, F)
+
+        def inter(Fa, Fb, M):
+            return np.asarray(M, dtype=bool) & np.outer(Fa, Fb)
+
+        in0 = intra(fa[0], M00); in1 = intra(fa[1], M11); in2 = intra(fa[2], M22)
+        m01 = inter(fa[0], fa[1], M01); m02 = inter(fa[0], fa[2], M02); m12 = inter(fa[1], fa[2], M12)
+
+        A = np.zeros((24, 24), dtype=bool)
+        A[0:8, 0:8] = in0 | in0.T
+        A[8:16, 8:16] = in1 | in1.T
+        A[16:24, 16:24] = in2 | in2.T
+        A[0:8, 8:16] = m01;  A[8:16, 0:8] = m01.T
+        A[0:8, 16:24] = m02; A[16:24, 0:8] = m02.T
+        A[8:16, 16:24] = m12; A[16:24, 8:16] = m12.T
+
+        # 24 个节点(slot p, octant c)的世界坐标 = v_parents[p]*2 + offset[c]
+        v_parents_arr = np.stack([np.array(v0_c[:3], dtype=np.int32),
+                                  np.array(v1_c[:3], dtype=np.int32),
+                                  np.array(v2_c[:3], dtype=np.int32)])
+        node_coords = v_parents_arr[np.arange(24) // 8] * 2 + offsets[np.arange(24) % 8]   # (24,3)
+        keys = [tuple(int(x) for x in node_coords[i]) for i in range(24)]
+        for i in range(24):
+            if keys[i] not in key_to_coord:
+                key_to_coord[keys[i]] = node_coords[i]
+
+        # 三角形(逐 face 向量化)；同时把三角形的边记入全局 covered 集合
+        valid_mask = A[_IDX_I, _IDX_J] & A[_IDX_J, _IDX_K] & A[_IDX_I, _IDX_K]
+        tris = _TRI_INDICES[valid_mask]
+        if len(tris) > 0:
+            p_idx = tris // 8
+            c_idx = tris % 8
+            tri_blocks.append(v_parents_arr[p_idx] * 2 + offsets[c_idx])
+            for i, j, k in tris:
+                ka, kb, kc = sorted((keys[int(i)], keys[int(j)], keys[int(k)]))
+                tri_edge_keys.add((ka, kb)); tri_edge_keys.add((kb, kc)); tri_edge_keys.add((ka, kc))
+
+        # 全局收集所有边
+        ei, ej = np.where(np.triu(A, 1))
+        for i, j in zip(ei.tolist(), ej.tolist()):
+            ki, kj = keys[i], keys[j]
+            edge_keys.add((ki, kj) if ki <= kj else (kj, ki))
+
+        # 全局收集 active 节点
+        for i in np.where(np.concatenate(fa))[0].tolist():
+            active_keys.add(keys[i])
+
+    all_blocks = list(tri_blocks)
+
+    # 未被任何三角形覆盖的边 -> 退化面 (A, B, B)   [全局判断，与 raw 的 get_level_info 一致]
+    for a, b in (edge_keys - tri_edge_keys):
+        all_blocks.append(np.stack([key_to_coord[a], key_to_coord[b], key_to_coord[b]])[None])
+
+    # 孤立的 active 节点(全局无任何边) -> 点 (A, A, A)
+    edge_nodes = set()
+    for a, b in edge_keys:
+        edge_nodes.add(a); edge_nodes.add(b)
+    for a in (active_keys - edge_nodes):
+        c = key_to_coord[a]
+        all_blocks.append(np.stack([c, c, c])[None])
+
+    if not all_blocks:
+        return np.array([]), np.array([])
+
+    flat_v = np.concatenate(all_blocks, axis=0).reshape(-1, 3)
+    V, inv = np.unique(flat_v, axis=0, return_inverse=True)
+    F_raw = inv.reshape(-1, 3)
+    F_sorted = np.sort(F_raw, axis=1)
+    _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True)
+    F = F_raw[unique_face_indices]
+    V = undiscretize(V, num_discrete=resolution)
+    return V, F
+
+
+def reconstruct_base_mesh(output_tokens, resolution):
+    all_face_coords = []
+    for f_token in output_tokens:
+        v0_c, v1_c, v2_c = f_token
+        c0 = np.array(v0_c[:3], dtype=np.int32)
+        c1 = np.array(v1_c[:3], dtype=np.int32)
+        c2 = np.array(v2_c[:3], dtype=np.int32)
+        all_face_coords.append(np.stack([c0, c1, c2]))
+
+    if not all_face_coords:
+        return np.array([]), np.array([])
+
+    flat_v = np.concatenate(all_face_coords, axis=0).reshape(-1, 3)
+    V, inv = np.unique(flat_v, axis=0, return_inverse=True)
+    F_raw = inv.reshape(-1, 3)
+    F_sorted = np.sort(F_raw, axis=1)
+    _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True)
+    F = F_raw[unique_face_indices]
+    V = undiscretize(V, num_discrete=resolution)
+    return V, F
+
+
+class MeshCodec:
+    def __init__(self, args):
+        self.max_depth = args.max_depth
+        self.max_resolution = args.max_resolution
+        self.use_hourglass = args.use_hourglass
+        self.max_children = NUM_OCTANTS  # octant 原位固定 8(忽略 config 的 max_children)
+
+        self.tokenizer = DynamicMeshTokenizer(resolution=self.max_resolution)
+
+    def encode_mesh(self, mesh: trimesh.Trimesh, target_depth: int, corrupt=None) -> tuple:
+        if isinstance(mesh, trimesh.Scene):
+            mesh = mesh.dump(concatenate=True)
+        V = np.asarray(mesh.vertices, dtype=np.float64)
+
+        res = 2 ** (target_depth + 1)
+        v_restored = undiscretize(discretize(V, num_discrete=res), num_discrete=res)
+        quantized_mesh = trimesh.Trimesh(vertices=v_restored, faces=mesh.faces, process=False)
+
+        # corrupt(仅训练时由 dataset 传入)= 2a 暴露偏差腐蚀; None 时与原版逐位一致(推理不受影响)
+        raw_seq, seq = _fast_fix3_level_tokens(V, mesh.faces, target_depth, self.tokenizer, corrupt=corrupt)
+        return raw_seq, seq, quantized_mesh
+
+
+    def decode_tokens(self, token_sequence: list, target_depth: int):
+        if os.environ.get("FIX3_FAST_DECODE", "1") != "1":
+            return super().decode_tokens(token_sequence, target_depth)
+        if token_sequence is None or len(token_sequence) == 0:
+            return None
+        return _fast_decode_tokens(token_sequence, target_depth, self.tokenizer)
+
+
+    def decode_raw_tokens(self, raw_token_sequence: list, target_depth: int) -> trimesh.Trimesh:
+        if not raw_token_sequence:
+            return None
+        resolution = 2 ** (target_depth)
+        raw_t = self.tokenizer.decode_raw(raw_token_sequence)
+        v, f = reconstruct_base_mesh(raw_t, resolution)
+        if len(v) == 0 or len(f) == 0:
+            return None
+        return trimesh.Trimesh(vertices=v, faces=f, process=False)
+
+    def get_model_config(self) -> dict:
+        t = self.tokenizer
+        return {
+            "total_vocab_size": t.vocab_size,
+            "coord_vocab_size": t.resolution,
+            "f_vocab_size": t.f_vocab,
+            "row_vocab_size": t.row_vocab,            # intra/inter 行子码各 256
+            "coord_offset": t.coord_offset,
+            "f_offset": t.f_offset,
+            "intra_row_offset": t.intra_row_offset,
+            "inter_row_offset": t.inter_row_offset,
+            "pad_token_id": t.pad_id,
+            "num_octants": NUM_OCTANTS,
+            "subcodes_per_position": NUM_OCTANTS,     # 每个 refine 位置 8 个子码槽(F 复制；intra/inter 为 8 行)
+            "tokens_per_face": DynamicMeshTokenizer.TOKENS_PER_FACE,   # 60 扁平 token/face (9+3+6*8)
+            "coord_per_face": 9,
+            "refine_positions_per_face": DynamicMeshTokenizer.REFINE_POS_PER_FACE,  # 9 (3F+3intra+3inter)
+            "input_num_categories": t.resolution + 1,
+            # 输出端：F / intra_row / inter_row 三类，每类 256；coord 为输入端条件
+            "output_num_categories": 1 + t.f_vocab + 2 * t.row_vocab,
+        }
+
+
+# ==========================================
+# 自测：encode -> decode -> reconstruct 往返
+# ==========================================
+if __name__ == "__main__":
+    # ==========================================
+    # 编解码 + 重构导出流程 (参照 fix2 的 __main__)
+    #   每个 mesh: 导出 GT；对每个 depth 做
+    #     encode -> decode_tokens(细分/Fine) / decode_raw_tokens(粗糙/Coarse) / 量化网格
+    #   并打印计时、token 数(60/face)与 token 级往返一致性。
+    # ==========================================
+    import time
+
+    # 1. 输入/输出文件夹 (按需修改)
+    INPUT_DIRECTORY = "/root/junkai_lin/mul_res_diffusion/demo/complex_test_data/files"
+    OUTPUT_DIRECTORY = "./output/fix3_octant_factorized2"
+    os.makedirs(OUTPUT_DIRECTORY, exist_ok=True)
+
+    @dataclass
+    class MeshCodecArgs:
+        max_depth: int = 10
+        max_resolution: int = 1024
+        use_hourglass: bool = True
+        max_children: int = 8      # fix3 octant 原位固定 8 (此字段实际被忽略)
+
+    # 2. 初始化 Codec
+    args = MeshCodecArgs()
+    codec = MeshCodec(args=args)
+    print(f"[*] Initialized fix3 Codec. Total Vocab Size: {codec.tokenizer.vocab_size} "
+          f"(tokens/face={DynamicMeshTokenizer.TOKENS_PER_FACE})")
+
+    # 3. 收集所有 .obj / .glb
+    input_dir_path = Path(INPUT_DIRECTORY)
+    mesh_files = sorted(input_dir_path.rglob("*.obj")) + sorted(input_dir_path.rglob("*.glb"))
+    if not mesh_files:
+        print(f"[!] No .obj or .glb files found in {INPUT_DIRECTORY}")
+
+    DEPTHS = [0, 1, 2, 3, 4, 5, 6]
+
+    # 4. 逐文件处理
+    for idx, file_path in enumerate(mesh_files, 1):
+        obj_name = file_path.stem
+        print(f"\n{'=' * 50}")
+        print(f"[{idx}/{len(mesh_files)}] Processing: {file_path.name}")
+        print(f"{'=' * 50}")
+        file_start_time = time.perf_counter()
+
+        try:
+            raw_mesh = trimesh.load(str(file_path), force="mesh")
+            normalize(raw_mesh)
+        except Exception as e:
+            print(f"[!] load/normalize failed: {e}")
+            continue
+
+        # 导出 GT
+        gt_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj")
+        raw_mesh.export(gt_path)
+        print(f"[*] Saved GT -> {gt_path}")
+
+        # 不同深度的编解码
+        for depth in DEPTHS:
+            print(f"\n--- Depth {depth} | res(fine)={2 ** (depth + 1)} | octant 原位 ---")
+
+            # Encoding
+            t0 = time.perf_counter()
+            raw_tokens, tokens, quantized_mesh = codec.encode_mesh(raw_mesh, target_depth=depth)
+            encode_time = time.perf_counter() - t0
+            nface = len(tokens) // DynamicMeshTokenizer.TOKENS_PER_FACE
+            print(f"[Timer] Encoding: {encode_time:.4f}s | flat_tokens={len(tokens)} "
+                  f"(nface={nface}, 模型逻辑位置={18 * nface})")
+
+            # token 级往返一致性 (encode->decode->re-encode)
+            same = (tokens == codec.tokenizer.encode(codec.tokenizer.decode(tokens)))
+            print(f"[Check] token 往返一致: {same}")
+
+            # Decoding (Fine)
+            if len(tokens) > 0:
+                t1 = time.perf_counter()
+                recon_mesh = codec.decode_tokens(tokens, target_depth=depth)
+                decode_fine_time = time.perf_counter() - t1
+                if recon_mesh is not None:
+                    out_path = os.path.join(
+                        OUTPUT_DIRECTORY, f"{obj_name}_recon_d{depth}_res{2 ** (depth + 1)}_fix3.obj")
+                    recon_mesh.export(out_path)
+                    print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s "
+                          f"-> V={len(recon_mesh.vertices)} F={len(recon_mesh.faces)}")
+                else:
+                    print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s -> empty")
+
+            # Decoding (Coarse / Raw)
+            if len(raw_tokens) > 0:
+                t2 = time.perf_counter()
+                raw_recon_mesh = codec.decode_raw_tokens(raw_tokens, target_depth=depth)
+                decode_coarse_time = time.perf_counter() - t2
+                if raw_recon_mesh is not None:
+                    raw_out_path = os.path.join(
+                        OUTPUT_DIRECTORY, f"{obj_name}_raw_recon_d{depth}_res{2 ** depth}.obj")
+                    raw_recon_mesh.export(raw_out_path)
+                    print(f"[Timer] Decoding (Coarse): {decode_coarse_time:.4f}s "
+                          f"-> V={len(raw_recon_mesh.vertices)} F={len(raw_recon_mesh.faces)}")
+
+            # 量化网格 (输入离散化后的 GT，用于对照重构质量上限)
+            if quantized_mesh is not None:
+                quant_out_path = os.path.join(
+                    OUTPUT_DIRECTORY, f"{obj_name}_dis_d{depth + 1}_res{2 ** (depth + 1)}.obj")
+                quantized_mesh.export(quant_out_path)
+
+        file_total_time = time.perf_counter() - file_start_time
+        print(f"\n[*] Finished {file_path.name} in {file_total_time:.2f}s")
+
+
+# ==========================================================================
+# encode_mesh / decode_tokens 的 numba 加速实现细节(原 *_fast.py 并入)
+# ==========================================================================
+
+_OCT_W = np.array([1, 2, 4], dtype=np.int64)   # octant 局部索引 = x + 2y + 4z
+
+def _corrupt_fid_nodes(fid, nodes, und, nN, target_depth, corrupt):
+    """2a 暴露偏差训练腐蚀(仅训练): 模拟级联实测的两类粗输入误差。
+      (a) 伪面: 在已存在的 GT 节点间加不该有的面(过连接)。下游打包按 fid 查表 ->
+          非边对的 inter 行自动全 0(empty),教模型"看 PC 把伪面 refine 成空"。
+      (b) 节点坐标抖动 ±jitter_max(位移): 只改 coord_part 的输入坐标,目标 F/intra/inter
+          仍按真节点身份查表不变,教模型"输入坐标偏了也 refine 对"。
+    目标无需重新导出 —— 全部沿用 fid 索引,与干净路径同一套向量化打包。返回 (fid_aug, nodes_coord)。
+    """
+    sp_rate = float(corrupt.get("spurious_rate", 0.0))
+    jp = float(corrupt.get("jitter_prob", 0.0))
+
+    if sp_rate > 0.0 and nN >= 3 and und.shape[0] > 0:
+        radius = int(corrupt.get("spatial_radius", 2))
+        n_add = int(sp_rate * fid.shape[0])
+        if n_add > 0:
+            coord2idx = {(int(c[0]), int(c[1]), int(c[2])): i for i, c in enumerate(nodes)}
+            existing = set(map(tuple, np.sort(fid, axis=1).tolist()))
+            ne = und.shape[0]
+            spurious, tries, maxtries = [], 0, n_add * 6
+            while len(spurious) < n_add and tries < maxtries:
+                tries += 1
+                ei = int(np.random.randint(ne))
+                ai, bi = int(und[ei, 0]), int(und[ei, 1])
+                base = nodes[ai] if np.random.random() < 0.5 else nodes[bi]
+                off = np.random.randint(-radius, radius + 1, size=3)
+                if not off.any():
+                    continue
+                ck = (int(base[0] + off[0]), int(base[1] + off[1]), int(base[2] + off[2]))
+                ci = coord2idx.get(ck)
+                if ci is None or ci == ai or ci == bi:
+                    continue
+                tri = tuple(sorted((ai, bi, ci)))
+                if tri in existing:
+                    continue
+                existing.add(tri)
+                spurious.append(tri)
+            if spurious:
+                fid = np.concatenate([fid, np.asarray(spurious, dtype=fid.dtype)], axis=0)
+
+    nodes_coord = nodes
+    if jp > 0.0:
+        jmax = int(corrupt.get("jitter_max", 1))
+        res_c = 2 ** target_depth
+        nodes_coord = nodes.copy()
+        mask = np.random.random(nN) < jp
+        k = int(mask.sum())
+        if k > 0:
+            noise = np.random.randint(-jmax, jmax + 1, size=(k, 3))
+            nodes_coord[mask] = np.clip(nodes_coord[mask] + noise, 0, res_c - 1)
+    return fid, nodes_coord
+
+def _fast_fix3_level_tokens(vertices, faces, target_depth, tok, corrupt=None):
+    """等价替换 MatrixOctree(max_depth=td+1).build + encode_level_tokens(td) + tokenizer.encode/encode_raw。
+    返回 (raw_token_sequence, token_sequence)，与原版逐位一致。"""
+    res = 2 ** (target_depth + 1)
+    Q = discretize(np.asarray(vertices, dtype=np.float64), num_discrete=res).astype(np.int64)
+    faces = np.asarray(faces, dtype=np.int64).reshape(-1, 3)
+
+    # ---- 节点(td) = 所有顶点叶子坐标的 parent；占据 F = 所有顶点的 octant ----
+    P_all = Q >> 1
+    i_all = (Q & 1) @ _OCT_W
+
+    # ---- 叶子边表(每面 3 条，跳过零长边 u==v，与原版一致) ----
+    if faces.size:
+        e = np.concatenate([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=0)
+        u, v = Q[e[:, 0]], Q[e[:, 1]]
+        keep = np.any(u != v, axis=1)
+        u, v = u[keep], v[keep]
+    else:
+        u = v = np.zeros((0, 3), dtype=np.int64)
+    pu, pv = u >> 1, v >> 1
+    iu = (u & 1) @ _OCT_W
+    iv = (v & 1) @ _OCT_W
+
+    # 节点表(np.unique 行字典序 = 坐标 (x,y,z) 升序 = 原版 face 内顶点 sorted() 同序)
+    nodes, inv = np.unique(np.concatenate([P_all, pu, pv], axis=0), axis=0, return_inverse=True)
+    nV, M = P_all.shape[0], pu.shape[0]
+    occ_idx = inv[:nV]
+    a, b = inv[nV:nV + M], inv[nV + M:]
+    nN = nodes.shape[0]
+    intra = a == b
+
+    # ---- F 行子码(8bit int)：所有顶点的 octant 占据 ----
+    Fbits = np.zeros(nN, dtype=np.int64)
+    np.bitwise_or.at(Fbits, occ_idx, np.int64(1) << i_all)
+
+    # ---- intra 行子码 [nN,8]：绝对原位、对称双向、无对角填充(与原版一致) ----
+    intra_rows = np.zeros((nN, 8), dtype=np.int64)
+    ai, iui, ivi = a[intra], iu[intra], iv[intra]
+    np.bitwise_or.at(intra_rows, (ai, iui), np.int64(1) << ivi)
+    np.bitwise_or.at(intra_rows, (ai, ivi), np.int64(1) << iui)
+
+    # ---- inter 行子码：有向 (a,b) 分组(双向都存，与原版 build 一致) ----
+    im = ~intra
+    ea, eb, eiu, eiv = a[im], b[im], iu[im], iv[im]
+    dkey = np.concatenate([ea * nN + eb, eb * nN + ea])
+    drow = np.concatenate([eiu, eiv])
+    dcol = np.concatenate([eiv, eiu])
+    ukeys, kinv = np.unique(dkey, return_inverse=True)
+    inter_rows = np.zeros((max(len(ukeys), 1), 8), dtype=np.int64)
+    if dkey.size:
+        np.bitwise_or.at(inter_rows, (kinv, drow), np.int64(1) << dcol)
+
+    # ---- 无向节点邻接 -> CSR -> numba 三角形枚举 ----
+    if ea.size:
+        und = np.unique(np.stack([np.minimum(ea, eb), np.maximum(ea, eb)], axis=1), axis=0)
+    else:
+        und = np.zeros((0, 2), dtype=np.int64)
+    deg = np.zeros(nN, dtype=np.int64)
+    if und.size:
+        np.add.at(deg, und[:, 0], 1)
+        np.add.at(deg, und[:, 1], 1)
+        order = np.lexsort((und[:, 1], und[:, 0]))
+        fu, fv = und[order, 0], und[order, 1]
+        indptr = np.zeros(nN + 1, dtype=np.int64)
+        np.add.at(indptr, fu + 1, 1)
+        indptr = np.cumsum(indptr)
+        tu, tv, tw = _triangles_csr(indptr, fv.astype(np.int64))
+    else:
+        tu = tv = tw = np.zeros(0, dtype=np.int64)
+
+    # ---- faces_ids：三角形 + 未覆盖边退化 (A,B,B) + 孤立点 (A,A,A)，集合与原版一致 ----
+    tris = np.stack([tu, tv, tw], axis=1) if tu.size else np.zeros((0, 3), dtype=np.int64)
+    covered = set()
+    for x, y, z in tris:
+        x, y, z = int(x), int(y), int(z)
+        covered.add((x, y)); covered.add((y, z)); covered.add((x, z))
+    extra = []
+    for x, y in und:
+        x, y = int(x), int(y)
+        if (x, y) not in covered:
+            extra.append((x, y, y))
+    for x in np.flatnonzero(deg == 0):
+        x = int(x)
+        extra.append((x, x, x))
+    fid = np.concatenate([tris, np.asarray(extra, dtype=np.int64).reshape(-1, 3)], axis=0)
+    if fid.shape[0] == 0:
+        return [], []
+
+    # ---- 原版排序键 + 决定性平局键 ----
+    c0, c1, c2 = nodes[fid[:, 0]], nodes[fid[:, 1]], nodes[fid[:, 2]]
+    order = np.lexsort((c2[:, 0], c2[:, 1], c2[:, 2],
+                        c1[:, 0], c1[:, 1], c1[:, 2],
+                        c0[:, 0], c0[:, 1], c0[:, 2]))   # 主键=c0.z(最后) ... 与原版 7 键一致+补齐平局
+    fid = fid[order]
+    nodes_coord = nodes
+    if corrupt is not None:
+        fid, nodes_coord = _corrupt_fid_nodes(fid, nodes, und, nN, target_depth, corrupt)
+    c0, c1, c2 = nodes_coord[fid[:, 0]], nodes_coord[fid[:, 1]], nodes_coord[fid[:, 2]]
+
+    # ---- 60/face 扁平 token 向量化打包 ----
+    # 注: F_part/intra_part/inter_part 下面仍用 fid(节点身份)查表 -> 腐蚀的伪面自动得正确目标
+    # (非边对 inter 全 0),坐标抖动只影响上面 coord_part,目标不变。
+    co, fo = tok.coord_offset, tok.f_offset
+    io, eo = tok.intra_row_offset, tok.inter_row_offset
+    nF = fid.shape[0]
+
+    coord_part = np.concatenate([c0, c1, c2], axis=1) + co                      # (nF,9)
+    F_part = np.stack([Fbits[fid[:, 0]], Fbits[fid[:, 1]], Fbits[fid[:, 2]]], axis=1) + fo
+    intra_part = np.concatenate([intra_rows[fid[:, 0]],
+                                 intra_rows[fid[:, 1]],
+                                 intra_rows[fid[:, 2]]], axis=1) + io           # (nF,24)
+
+    inter_blocks = []
+    for x, y in ((0, 1), (0, 2), (1, 2)):
+        A_, B_ = fid[:, x], fid[:, y]
+        rows = np.zeros((nF, 8), dtype=np.int64)
+        same = A_ == B_
+        rows[same] = intra_rows[A_[same]]                  # inter8(a,a)=intra8(a)，与原版一致
+        ns = np.flatnonzero(~same)
+        if ns.size:
+            dk = A_[ns] * nN + B_[ns]
+            pos = np.searchsorted(ukeys, dk)
+            pos_c = np.minimum(pos, max(len(ukeys) - 1, 0))
+            found = (ukeys[pos_c] == dk) if len(ukeys) else np.zeros(ns.size, dtype=bool)
+            rows[ns[found]] = inter_rows[pos_c[found]]     # 未找到 -> 保持全 0(原版 empty 矩阵)
+        inter_blocks.append(rows)
+    inter_part = np.concatenate(inter_blocks, axis=1) + eo                      # (nF,24)
+
+    token = np.concatenate([coord_part, F_part, intra_part, inter_part], axis=1)  # (nF,60)
+    token_sequence = token.ravel().tolist()
+    raw_token_sequence = coord_part.reshape(-1, 1).tolist()                       # 9 coord/face, 各 [t]
+    return raw_token_sequence, token_sequence
+
+_BIT_SHIFTS = np.arange(NUM_OCTANTS, dtype=np.int64)                 # _int_to_row: bit i = (v>>i)&1
+
+_NODE_PARENT = np.arange(24) // NUM_OCTANTS                          # 24 个节点 -> 哪个父顶点 slot
+
+_NODE_OCTANT = np.arange(24) % NUM_OCTANTS                           # 24 个节点 -> 哪个 octant
+
+_OCT_OFFSETS = np.stack([np.arange(8) & 1, (np.arange(8) >> 1) & 1, (np.arange(8) >> 2) & 1],
+                        axis=1).astype(np.int32)                     # octant -> (dx,dy,dz)
+
+_INTER_ROW_SLOT = np.array([0, 0, 1])
+
+_INTER_COL_SLOT = np.array([1, 2, 2])
+
+def _unpack_rows(values):
+    """int 行子码 -> bool 位。values (...,) -> (..., 8), 与 _int_to_row 同序(bit i = (v>>i)&1)。"""
+    return ((values[..., None].astype(np.int64) >> _BIT_SHIFTS) & 1).astype(bool)
+
+def _pack_keys(coords, resolution):
+    """(...,3) int 坐标 -> (...) int64 唯一 key。坐标范围 [0, resolution)。"""
+    c = coords.astype(np.int64)
+    return (c[..., 0] * resolution + c[..., 1]) * resolution + c[..., 2]
+
+def _fast_reconstruct_next_level(token_array, resolution, tokenizer):
+    """token_array: (nface, TOKENS_PER_FACE) int。返回 (V, F), 语义同 reconstruct_next_level_optimized。"""
+    from modules.octree_codec import _TRI_INDICES
+
+    nface = token_array.shape[0]
+    parents = (token_array[:, :9] - tokenizer.coord_offset).reshape(nface, 3, 3).astype(np.int32)
+    occupancy = _unpack_rows(token_array[:, 9:12] - tokenizer.f_offset)          # (nface,3,8) bool
+
+    def rows_to_mat(lo, hi, offset):
+        return _unpack_rows(token_array[:, lo:hi] - offset)                      # (nface,8,8) bool, [行,列]
+
+    intra = [rows_to_mat(12 + 8 * i, 20 + 8 * i, tokenizer.intra_row_offset) for i in range(3)]
+    inter = [rows_to_mat(36 + 8 * i, 44 + 8 * i, tokenizer.inter_row_offset) for i in range(3)]
+
+    # ---- 24x24 邻接: intra 块按 F&F 掩码后对称化, inter 块按 Fa&Fb 掩码 ----
+    adjacency = np.zeros((nface, 24, 24), dtype=bool)
+    for slot in range(3):
+        f = occupancy[:, slot]                                                   # (nface,8)
+        block = intra[slot] & (f[:, :, None] & f[:, None, :])
+        block = block | np.swapaxes(block, 1, 2)
+        adjacency[:, slot * 8:(slot + 1) * 8, slot * 8:(slot + 1) * 8] = block
+    for b in range(3):
+        ra, ca = _INTER_ROW_SLOT[b], _INTER_COL_SLOT[b]
+        block = inter[b] & (occupancy[:, ra][:, :, None] & occupancy[:, ca][:, None, :])
+        adjacency[:, ra * 8:(ra + 1) * 8, ca * 8:(ca + 1) * 8] = block
+        adjacency[:, ca * 8:(ca + 1) * 8, ra * 8:(ra + 1) * 8] = np.swapaxes(block, 1, 2)
+
+    # ---- 24 个节点的世界坐标与 key ----
+    node_coords = parents[:, _NODE_PARENT] * 2 + _OCT_OFFSETS[_NODE_OCTANT]      # (nface,24,3)
+    node_keys = _pack_keys(node_coords, resolution)                              # (nface,24)
+
+    # ---- 三角形: 三条边都在邻接里 ----
+    i_idx, j_idx, k_idx = _TRI_INDICES[:, 0], _TRI_INDICES[:, 1], _TRI_INDICES[:, 2]
+    valid = (adjacency[:, i_idx, j_idx] & adjacency[:, j_idx, k_idx] & adjacency[:, i_idx, k_idx])
+    face_of_tri, tri_of_tri = np.nonzero(valid)                                  # 各 (ntri,)
+    tri_nodes = _TRI_INDICES[tri_of_tri]                                         # (ntri,3) 节点编号
+    tri_coords = node_coords[face_of_tri[:, None], tri_nodes]                    # (ntri,3,3)
+    tri_keys = node_keys[face_of_tri[:, None], tri_nodes]                        # (ntri,3)
+
+    # ---- 全部边(上三角) ----
+    upper = np.triu(np.ones((24, 24), dtype=bool), 1)
+    face_of_edge, node_i, node_j = np.nonzero(adjacency & upper)
+    key_i = node_keys[face_of_edge, node_i]
+    key_j = node_keys[face_of_edge, node_j]
+    edge_pairs = np.unique(np.stack([np.minimum(key_i, key_j), np.maximum(key_i, key_j)], axis=1), axis=0)
+
+    # ---- 被三角形覆盖的边 ----
+    if len(tri_keys):
+        tri_sorted = np.sort(tri_keys, axis=1)
+        covered = np.concatenate([tri_sorted[:, [0, 1]], tri_sorted[:, [1, 2]], tri_sorted[:, [0, 2]]], axis=0)
+        covered = np.unique(covered, axis=0)
+    else:
+        covered = np.zeros((0, 2), dtype=np.int64)
+
+    blocks = [tri_coords] if len(tri_coords) else []
+
+    # ---- 未被任何三角形覆盖的边 -> 退化面 (A,B,B) ----
+    if len(edge_pairs):
+        edge_flat = edge_pairs[:, 0] * (resolution ** 3) + edge_pairs[:, 1]
+        covered_flat = (covered[:, 0] * (resolution ** 3) + covered[:, 1]) if len(covered) \
+            else np.zeros(0, dtype=np.int64)
+        open_edges = edge_pairs[~np.isin(edge_flat, covered_flat)]
+        if len(open_edges):
+            key_to_coord = _build_key_lookup(node_keys, node_coords)
+            ca = key_to_coord[open_edges[:, 0]]
+            cb = key_to_coord[open_edges[:, 1]]
+            blocks.append(np.stack([ca, cb, cb], axis=1))
+
+    # ---- 孤立的 active 节点(无任何边) -> 点 (A,A,A) ----
+    active_keys = np.unique(node_keys[occupancy.reshape(nface, 24)])
+    if len(active_keys):
+        edge_nodes = np.unique(edge_pairs.reshape(-1)) if len(edge_pairs) else np.zeros(0, dtype=np.int64)
+        isolated = active_keys[~np.isin(active_keys, edge_nodes)]
+        if len(isolated):
+            key_to_coord = _build_key_lookup(node_keys, node_coords)
+            c = key_to_coord[isolated]
+            blocks.append(np.stack([c, c, c], axis=1))
+
+    if not blocks:
+        return np.array([]), np.array([])
+
+    flat_v = np.concatenate(blocks, axis=0).reshape(-1, 3)
+    V, inverse = np.unique(flat_v, axis=0, return_inverse=True)
+    faces_raw = inverse.reshape(-1, 3)
+    _, unique_face_indices = np.unique(np.sort(faces_raw, axis=1), axis=0, return_index=True)
+    F = faces_raw[unique_face_indices]
+    return undiscretize(V, num_discrete=resolution), F
+
+def _build_key_lookup(node_keys, node_coords):
+    """key -> 坐标 的稠密查表(用 np.unique 的第一次出现), 供退化面/孤立点还原坐标。"""
+    flat_keys = node_keys.reshape(-1)
+    flat_coords = node_coords.reshape(-1, 3)
+    uniq, first = np.unique(flat_keys, return_index=True)
+    lookup = np.zeros((int(uniq.max()) + 1, 3), dtype=np.int32)
+    lookup[uniq] = flat_coords[first]
+    return lookup
+
+def _fast_decode_tokens(token_sequence, target_depth, tokenizer):
+    """decode_tokens 的向量化实现。返回 trimesh.Trimesh 或 None。"""
+    arr = np.asarray(token_sequence).reshape(-1)
+    arr = arr[arr != tokenizer.pad_id]                                # 与原版 clean_seq 一致
+    per_face = tokenizer.TOKENS_PER_FACE
+    nface = arr.shape[0] // per_face
+    if nface == 0:
+        return None
+    token_array = arr[:nface * per_face].reshape(nface, per_face).astype(np.int64)
+    V, F = _fast_reconstruct_next_level(token_array, 2 ** (target_depth + 1), tokenizer)
+    if len(V) == 0 or len(F) == 0:
+        return None
+    return trimesh.Trimesh(vertices=V, faces=F, process=False)
+
+
+FastMeshCodec = MeshCodec   # 兼容别名: 加速实现已并入 MeshCodec

mulres_infer/modules/octree_codec_acc.py

@@ -0,0 +1,857 @@
+import numpy as np
+import trimesh
+from tqdm import tqdm
+import os
+import sys
+import glob
+import multiprocessing
+import random
+from collections import Counter
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from pathlib import Path
+from utils.data_utils import undiscretize, discretize, normalize, clean_mesh
+from utils.mesh_metric import compute_mesh_metrics
+import itertools
+from dataclasses import dataclass
+# 预计算 24 个节点中选 3 个的所有可能组合 (共 2024 种)
+_TRI_INDICES = np.array(list(itertools.combinations(range(24), 3)))
+_IDX_I = _TRI_INDICES[:, 0]
+_IDX_J = _TRI_INDICES[:, 1]
+_IDX_K = _TRI_INDICES[:, 2]
+# ==========================================
+# 新增：动态网格 Tokenizer
+# ==========================================
+class DynamicMeshTokenizer:
+    def __init__(self, resolution, max_children):
+        self.resolution = resolution
+        self.max_children = max_children
+        self.m_bits = max_children * max_children
+        
+        # 1. 特殊控制 Token
+        self.pad_id = 0
+        
+        # 2. 计算各特征在词表中的偏移量 (Offsets)
+        self.coord_offset = 1
+        self.f_offset = self.coord_offset + self.resolution
+        self.m_intra_offset = self.f_offset + 256 # 1x8 对应 256 种可能
+        
+        self.m_vocab_size = 2 ** self.m_bits
+        self.m_inter_offset = self.m_intra_offset + self.m_vocab_size
+        
+        # 最终 Transformer 的总词表大小
+        self.vocab_size = self.m_inter_offset + self.m_vocab_size
+
+    def _list_to_int(self, bit_list):
+        res = 0
+        for i, bit in enumerate(bit_list):
+            if bit: res |= (1 << i)
+        return res
+
+    def _int_to_list(self, val, length):
+        return [(val >> i) & 1 for i in range(length)]
+        
+    def _m_to_int(self, M):
+        flat = [bit for row in M for bit in row]
+        return self._list_to_int(flat)
+        
+    def _int_to_m(self, val):
+        flat = self._int_to_list(val, self.m_bits)
+        return [flat[i*self.max_children : (i+1)*self.max_children] for i in range(self.max_children)]
+
+    # ------------------------------------
+    # 细分网格 (Fine Mesh) Tokenizer
+    # ------------------------------------
+    def encode(self, faces_data, use_hourglass=True):
+        """将细分几何特征列表编码为 1D Token 序列 (18 Tokens / Face)"""
+        seq = []
+            
+        for face in faces_data:
+            v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12 = face
+            
+            seq.extend([v0[0] + self.coord_offset, v0[1] + self.coord_offset, v0[2] + self.coord_offset])
+            seq.extend([v1[0] + self.coord_offset, v1[1] + self.coord_offset, v1[2] + self.coord_offset])
+            seq.extend([v2[0] + self.coord_offset, v2[1] + self.coord_offset, v2[2] + self.coord_offset])
+            
+            seq.extend([self._list_to_int(F0) + self.f_offset,
+                        self._list_to_int(F1) + self.f_offset,
+                        self._list_to_int(F2) + self.f_offset])
+            
+            seq.extend([self._m_to_int(M00) + self.m_intra_offset,
+                        self._m_to_int(M11) + self.m_intra_offset,
+                        self._m_to_int(M22) + self.m_intra_offset])
+            
+            seq.extend([self._m_to_int(M01) + self.m_inter_offset,
+                        self._m_to_int(M02) + self.m_inter_offset,
+                        self._m_to_int(M12) + self.m_inter_offset])
+        return seq
+
+    def decode(self, seq):
+        """将细分网格的 1D Token 序列解码回几何特征列表"""
+        faces_data = []
+        clean_seq = [t for t in seq if t not in ([self.pad_id])]
+        
+        for i in range(0, len(clean_seq), 18):
+            chunk = clean_seq[i:i+18]
+            if len(chunk) < 18:
+                break 
+                
+            v0 = (chunk[0] - self.coord_offset, chunk[1] - self.coord_offset, chunk[2] - self.coord_offset)
+            v1 = (chunk[3] - self.coord_offset, chunk[4] - self.coord_offset, chunk[5] - self.coord_offset)
+            v2 = (chunk[6] - self.coord_offset, chunk[7] - self.coord_offset, chunk[8] - self.coord_offset)
+            
+            F0 = self._int_to_list(chunk[9] - self.f_offset, 8)
+            F1 = self._int_to_list(chunk[10] - self.f_offset, 8)
+            F2 = self._int_to_list(chunk[11] - self.f_offset, 8)
+            
+            M00 = self._int_to_m(chunk[12] - self.m_intra_offset)
+            M11 = self._int_to_m(chunk[13] - self.m_intra_offset)
+            M22 = self._int_to_m(chunk[14] - self.m_intra_offset)
+            
+            M01 = self._int_to_m(chunk[15] - self.m_inter_offset)
+            M02 = self._int_to_m(chunk[16] - self.m_inter_offset)
+            M12 = self._int_to_m(chunk[17] - self.m_inter_offset)
+            
+            faces_data.append((v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12))
+            
+        return faces_data
+
+    # ------------------------------------
+    # [新增] 基础/粗糙网格 (Raw Mesh) Tokenizer
+    # ------------------------------------
+    def encode_raw(self, raw_faces_data, use_hourglass=True):
+        """将粗糙层级特征（仅含顶点）编码为 1D Token 序列 (9 Tokens / Face)"""
+        seq = []
+            
+        for face in raw_faces_data:
+            v0, v1, v2 = face
+            seq.extend([v0[0] + self.coord_offset, v0[1] + self.coord_offset, v0[2] + self.coord_offset])
+            seq.extend([v1[0] + self.coord_offset, v1[1] + self.coord_offset, v1[2] + self.coord_offset])
+            seq.extend([v2[0] + self.coord_offset, v2[1] + self.coord_offset, v2[2] + self.coord_offset])
+            
+        return seq
+
+    def decode_raw(self, seq):
+        """将粗糙层级的 1D Token 序列解码回几何特征列表 [(v0, v1, v2), ...]"""
+        faces_data = []
+        clean_seq = [t for t in seq if t not in ([self.pad_id])]
+        
+        # 每个面片固定转化为 9 个 Token
+        for i in range(0, len(clean_seq), 9):
+            chunk = clean_seq[i:i+9]
+            if len(chunk) < 9:
+                break 
+                
+            v0 = (chunk[0] - self.coord_offset, chunk[1] - self.coord_offset, chunk[2] - self.coord_offset)
+            v1 = (chunk[3] - self.coord_offset, chunk[4] - self.coord_offset, chunk[5] - self.coord_offset)
+            v2 = (chunk[6] - self.coord_offset, chunk[7] - self.coord_offset, chunk[8] - self.coord_offset)
+            
+            faces_data.append((v0, v1, v2))
+            
+        return faces_data
+
+
+# --- 以下为原有的类声明 (OctreeFaceTokenBuilder, OctreeNode, MatrixOctree 等保持不变) ---
+class OctreeFaceTokenBuilder:
+    def __init__(self, tree, max_children=4):
+        self.tree = tree
+        self.max_children = max_children 
+
+    def get_level_info(self, depth):
+        nodes_at_depth = {}
+        def collect_nodes(current, target_depth):
+            if current.depth == target_depth:
+                nodes_at_depth[current.coord] = current
+                return
+            for child in current.children:
+                if child: collect_nodes(child, target_depth)
+        
+        collect_nodes(self.tree.root, depth)
+        
+        coords = sorted(nodes_at_depth.keys(), key=lambda x: (x[2], x[1], x[0]))
+        coord_to_idx = {c: i for i, c in enumerate(coords)}
+        
+        adj = {i: set() for i in range(len(coords))}
+        for i, c in enumerate(coords):
+            node = nodes_at_depth[c]
+            all_neighbors = node.intra_neighbors | node.inter_neighbors
+            for neighbor_coord in all_neighbors:
+                if neighbor_coord in coord_to_idx:
+                    j = coord_to_idx[neighbor_coord]
+                    if i < j: 
+                        adj[i].add(j)
+
+        raw_faces_set = set()
+        covered_edges = set()
+
+        # 1. 提取所有互相连接的 3-clique (正常的三角形面片)
+        for u in range(len(coords)):
+            for v in adj[u]:
+                common_neighbors = adj[u] & adj[v]
+                for w in common_neighbors:
+                    # 排序坐标点以确保集合去重时的唯一性
+                    face_coords = tuple(sorted([coords[u], coords[v], coords[w]]))
+                    raw_faces_set.add(face_coords)
+                    
+                    # 记录已经被三角形覆盖的边
+                    covered_edges.add(tuple(sorted([face_coords[0], face_coords[1]])))
+                    covered_edges.add(tuple(sorted([face_coords[1], face_coords[2]])))
+                    covered_edges.add(tuple(sorted([face_coords[0], face_coords[2]])))
+
+        # 2. 提取未形成三角形的独立边 (线段结构，例如细长桌腿)
+        for u in range(len(coords)):
+            for v in adj[u]:
+                edge = tuple(sorted([coords[u], coords[v]]))
+                if edge not in covered_edges:
+                    # 将 1D 的线段转换为退化面片 (A, B, B) 以适配 Tokenizer 格式
+                    degenerate_face = (edge[0], edge[1], edge[1])
+                    raw_faces_set.add(degenerate_face)
+                    
+        # 3. (可选) 极端情况下处理没有任何连接的孤立顶点 (A, A, A)
+        for u in range(len(coords)):
+            has_edges = len(adj[u]) > 0 or any(u in adj[i] for i in range(u))
+            if not has_edges:
+                isolated_face = (coords[u], coords[u], coords[u])
+                raw_faces_set.add(isolated_face)
+
+        raw_faces_at_depth = list(raw_faces_set)
+        
+        # 保持原始排序逻辑以确保输出的 Determinism
+        raw_faces_at_depth.sort(key=lambda f: (f[0][2], f[0][1], f[0][0], 
+                                               f[1][2], f[1][1], f[1][0],
+                                               f[2][2]))
+        
+        return raw_faces_at_depth, nodes_at_depth
+
+    def reduce_node_children_to_max(self, nodes_map):
+        def bit_dist(a, b):
+            return bin(a ^ b).count('1')
+
+        # Pass 1: 基于原始图结构为所有节点一次性计算 merge plan，消除顺序依赖。
+        merge_plan = {}
+        for coord, n in nodes_map.items():
+            active_indices = set()
+            has_intra_row = n.intra_matrix.any(axis=1)
+            has_intra_col = n.intra_matrix.any(axis=0)
+            for i in range(8):
+                if has_intra_row[i] or has_intra_col[i]:
+                    active_indices.add(i)
+                    
+            for mat in n.inter_matrices.values():
+                has_inter_row = mat.any(axis=1)
+                for i in range(8):
+                    if has_inter_row[i]:
+                        active_indices.add(i)
+                        
+            for n_inter_coord in n.inter_neighbors:
+                if n_inter_coord in nodes_map:
+                    neighbor = nodes_map[n_inter_coord]
+                    if coord in neighbor.inter_matrices:
+                        has_in_col = neighbor.inter_matrices[coord].any(axis=0)
+                        for i in range(8):
+                            if has_in_col[i]:
+                                active_indices.add(i)
+
+            if len(active_indices) <= self.max_children:
+                merge_plan[coord] = None
+                continue
+
+            child_degrees = [0] * 8
+            intra_row_counts = np.count_nonzero(n.intra_matrix, axis=1)
+            intra_col_counts = np.count_nonzero(n.intra_matrix, axis=0)
+            for i in range(8):
+                child_degrees[i] += intra_row_counts[i] + intra_col_counts[i]
+                
+            out_inter_counts = np.zeros(8, dtype=int)
+            for mat in n.inter_matrices.values():
+                out_inter_counts += np.count_nonzero(mat, axis=1)
+
+            in_inter_counts = np.zeros(8, dtype=int)
+            for n_inter_coord in n.inter_neighbors:
+                if n_inter_coord in nodes_map:
+                    neighbor = nodes_map[n_inter_coord]
+                    if coord in neighbor.inter_matrices:
+                        in_inter_counts += np.count_nonzero(neighbor.inter_matrices[coord], axis=0)
+
+            for i in range(8):
+                child_degrees[i] += 2 * (out_inter_counts[i] + in_inter_counts[i])
+
+            active_list = sorted(list(active_indices))
+            scored = sorted(
+                active_list,
+                key=lambda idx: (-child_degrees[idx], idx)
+            )
+            keep_nodes = sorted(scored[:self.max_children])
+
+            mapping = {idx: idx for idx in range(8)}
+            for m_idx in active_list:
+                if m_idx in keep_nodes:
+                    continue
+                target_merge = min(keep_nodes, key=lambda x: (bit_dist(x, m_idx), x))
+                mapping[m_idx] = target_merge
+
+            merge_plan[coord] = mapping
+
+        # Pass 2: 根据全局 merge plan 同步重映射所有 intra/inter 矩阵。
+        for coord, n in nodes_map.items():
+            mapping = merge_plan.get(coord)
+            if mapping is not None:
+                old_intra = n.intra_matrix
+                new_intra = np.zeros_like(old_intra)
+                nz = np.argwhere(old_intra != 0)
+                for i, j in nz:
+                    new_intra[mapping[int(i)], mapping[int(j)]] = 1
+                np.fill_diagonal(new_intra, 0)
+                n.intra_matrix = new_intra
+
+            for neigh_coord, mat in list(n.inter_matrices.items()):
+                src_map = merge_plan.get(coord)
+                dst_map = merge_plan.get(neigh_coord)
+                if src_map is None and dst_map is None:
+                    continue
+
+                new_mat = np.zeros_like(mat)
+                nz = np.argwhere(mat != 0)
+                for i, j in nz:
+                    ni = src_map[int(i)] if src_map is not None else int(i)
+                    nj = dst_map[int(j)] if dst_map is not None else int(j)
+                    new_mat[ni, nj] = 1
+                n.inter_matrices[neigh_coord] = new_mat
+
+    def encode_level_tokens(self, depth):
+        raw_faces, nodes_map = self.get_level_info(depth)
+        
+        self.reduce_node_children_to_max(nodes_map)
+        
+        node_F = {}
+        node_active = {}
+        empty_mat = np.zeros((8, 8), dtype=np.int8)
+        
+        for coord, n in nodes_map.items():
+            active_indices = set()
+            has_intra_row = n.intra_matrix.any(axis=1)
+            has_intra_col = n.intra_matrix.any(axis=0)
+            for i in range(8):
+                if has_intra_row[i] or has_intra_col[i]:
+                    active_indices.add(i)
+                    
+            for mat in n.inter_matrices.values():
+                has_inter_row = mat.any(axis=1)
+                for i in range(8):
+                    if has_inter_row[i]:
+                        active_indices.add(i)
+                        
+            for n_inter_coord in n.inter_neighbors:
+                if n_inter_coord in nodes_map:
+                    neighbor = nodes_map[n_inter_coord]
+                    if n.coord in neighbor.inter_matrices:
+                        has_in_col = neighbor.inter_matrices[n.coord].any(axis=0)
+                        for i in range(8):
+                            if has_in_col[i]:
+                                active_indices.add(i)
+            
+            F = [0] * 8
+            for idx in active_indices: 
+                F[idx] = 1
+            node_F[coord] = F
+            node_active[coord] = sorted(list(active_indices))
+            
+        memo_intra = {}
+        memo_inter = {}
+        
+        def extract_NxN_intra(node_coord):
+            if node_coord in memo_intra:
+                return memo_intra[node_coord]
+                
+            active = node_active[node_coord][:self.max_children]
+            n = nodes_map[node_coord]
+            
+            if len(active) > 0:
+                M = np.zeros((self.max_children, self.max_children), dtype=np.int8)
+                sub = n.intra_matrix[np.ix_(active, active)]
+                M[:len(active), :len(active)] = sub
+                np.fill_diagonal(M[:len(active), :len(active)], 1)
+                res = M.tolist()
+            else:
+                res = np.zeros((self.max_children, self.max_children), dtype=np.int8).tolist()
+                
+            memo_intra[node_coord] = res
+            return res
+            
+        def extract_NxN_inter(coord_A, coord_B):
+            if coord_A == coord_B:
+                return extract_NxN_intra(coord_A)
+                
+            key = (coord_A, coord_B)
+            if key in memo_inter:
+                return memo_inter[key]
+                
+            active_A = node_active[coord_A][:self.max_children]
+            active_B = node_active[coord_B][:self.max_children]
+            nA = nodes_map[coord_A]
+            orig_M = nA.inter_matrices.get(coord_B, empty_mat)
+            
+            if len(active_A) > 0 and len(active_B) > 0:
+                M = np.zeros((self.max_children, self.max_children), dtype=np.int8)
+                M[:len(active_A), :len(active_B)] = orig_M[np.ix_(active_A, active_B)]
+                res = M.tolist()
+            else:
+                res = np.zeros((self.max_children, self.max_children), dtype=np.int8).tolist()
+                
+            memo_inter[key] = res
+            return res
+
+        raw_tokens = []
+        output_tokens = []
+        
+        for face in raw_faces:
+            v0_c, v1_c, v2_c = face
+            raw_tokens.append([v0_c, v1_c, v2_c])
+            
+            F0, F1, F2 = node_F[v0_c], node_F[v1_c], node_F[v2_c]
+            
+            M00 = extract_NxN_intra(v0_c)
+            M11 = extract_NxN_intra(v1_c)
+            M22 = extract_NxN_intra(v2_c)
+            
+            M01 = extract_NxN_inter(v0_c, v1_c)
+            M02 = extract_NxN_inter(v0_c, v2_c)
+            M12 = extract_NxN_inter(v1_c, v2_c)
+            
+            f_encoded = (
+                v0_c, v1_c, v2_c,
+                F0, F1, F2,
+                M00, M11, M22, M01, M02, M12
+            )
+            output_tokens.append(f_encoded)
+            
+        return raw_tokens, output_tokens
+
+class OctreeNode:
+    def __init__(self, x, y, z, depth):
+        self.coord = (x, y, z)
+        self.depth = depth
+        self.children = [None] * 8
+        self.intra_neighbors = set()
+        self.inter_neighbors = set()
+        self.intra_matrix = np.zeros((8, 8), dtype=np.int8)
+        self.inter_matrices = {} 
+
+class MatrixOctree:
+    def __init__(self, max_depth=7):
+        self.max_depth = max_depth
+        self.root = None
+
+    def _get_parent_coord(self, coord):
+        return (coord[0] >> 1, coord[1] >> 1, coord[2] >> 1)
+    
+    def _get_local_index(self, coord):
+        return ((coord[2] & 1) << 2) | ((coord[1] & 1) << 1) | (coord[0] & 1)
+
+    def build(self, vertices, faces):
+        resolution = 2 ** self.max_depth
+        quantized = discretize(vertices, num_discrete=resolution)
+        
+        leaf_nodes_map = {}
+        for i in range(len(vertices)):
+            coord = tuple(quantized[i])
+            if coord not in leaf_nodes_map:
+                leaf_nodes_map[coord] = OctreeNode(coord[0], coord[1], coord[2], self.max_depth)
+        
+        for face in faces:
+            coords = [tuple(quantized[i]) for i in face]
+            for i in range(3):
+                u, v = coords[i], coords[(i+1)%3]
+                if u == v: continue
+                
+                parent_u = self._get_parent_coord(u)
+                parent_v = self._get_parent_coord(v)
+                
+                node_u = leaf_nodes_map[u]
+                node_v = leaf_nodes_map[v]
+                
+                if parent_u == parent_v:
+                    node_u.intra_neighbors.add(v)
+                    node_v.intra_neighbors.add(u)
+                else:
+                    node_u.inter_neighbors.add(v)
+                    node_v.inter_neighbors.add(u)
+
+        current_layer_nodes = list(leaf_nodes_map.values())
+
+        for d in range(self.max_depth, 0, -1):
+            parent_layer_map = {}
+            for child in current_layer_nodes:
+                parent_coord = self._get_parent_coord(child.coord)
+                if parent_coord not in parent_layer_map:
+                    parent_layer_map[parent_coord] = OctreeNode(*parent_coord, d - 1)
+                
+                parent_node = parent_layer_map[parent_coord]
+                idx = self._get_local_index(child.coord)
+                parent_node.children[idx] = child
+
+            for child in current_layer_nodes:
+                parent_coord = self._get_parent_coord(child.coord)
+                parent_node = parent_layer_map[parent_coord]
+                child_idx = self._get_local_index(child.coord)
+                
+                for sibling_coord in child.intra_neighbors:
+                    sibling_idx = self._get_local_index(sibling_coord)
+                    parent_node.intra_matrix[child_idx, sibling_idx] = 1
+                
+                for neighbor_coord in child.inter_neighbors:
+                    neighbor_parent_coord = self._get_parent_coord(neighbor_coord)
+                    if neighbor_parent_coord not in parent_node.inter_matrices:
+                        parent_node.inter_matrices[neighbor_parent_coord] = np.zeros((8, 8), dtype=np.int8)
+                    
+                    neighbor_child_idx = self._get_local_index(neighbor_coord)
+                    parent_node.inter_matrices[neighbor_parent_coord][child_idx, neighbor_child_idx] = 1
+                    
+                    grandparent_p = self._get_parent_coord(parent_coord)
+                    grandparent_n = self._get_parent_coord(neighbor_parent_coord)
+                    
+                    if grandparent_p == grandparent_n:
+                        parent_node.intra_neighbors.add(neighbor_parent_coord)
+                    else:
+                        parent_node.inter_neighbors.add(neighbor_parent_coord)
+
+            current_layer_nodes = list(parent_layer_map.values())
+        if current_layer_nodes:
+            self.root = current_layer_nodes[0]
+
+
+def reconstruct_next_level_optimized(output_tokens, resolution, max_children=4):
+    all_face_coords = []
+    
+    indices = np.arange(8)
+    offsets = np.stack([indices & 1, (indices >> 1) & 1, (indices >> 2) & 1], axis=1).astype(np.int32)
+
+    def expand_intra(F, M_NxN):
+        M8x8 = np.zeros((8, 8), dtype=bool)
+        active = [i for i, val in enumerate(F) if val]
+        for i_n, i_o in enumerate(active):
+            for j_n, j_o in enumerate(active):
+                if i_n < max_children and j_n < max_children and M_NxN[i_n][j_n]:
+                    M8x8[i_o][j_o] = True
+        return M8x8
+        
+    def expand_inter(FA, FB, M_NxN):
+        M8x8 = np.zeros((8, 8), dtype=bool)
+        activeA = [i for i, val in enumerate(FA) if val]
+        activeB = [i for i, val in enumerate(FB) if val]
+        for i_n, i_o in enumerate(activeA):
+            for j_n, j_o in enumerate(activeB):
+                if i_n < max_children and j_n < max_children and M_NxN[i_n][j_n]:
+                    M8x8[i_o][j_o] = True
+        return M8x8
+
+    for f_token in output_tokens:
+        v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12 = f_token
+        
+        v_parents = [np.array(v[:3], dtype=np.int32) for v in [v0_c, v1_c, v2_c]]
+        
+        in0 = expand_intra(F0, M00)
+        in1 = expand_intra(F1, M11)
+        in2 = expand_intra(F2, M22)
+        
+        m01 = expand_inter(F0, F1, M01)
+        m02 = expand_inter(F0, F2, M02)
+        m12 = expand_inter(F1, F2, M12)
+
+        # in_mats = [in0, in1, in2]
+        # inter_edges = {
+        #     (0, 1): m01,
+        #     (1, 2): m12,
+        #     (0, 2): m02
+        # }
+
+        # def are_connected_opt(p_a, c_a, p_b, c_b):
+        #     # p_a 和 p_b 是 0, 1, 2 这三个顶点索引。
+        #     # 只要索引相同，说明是查询同一个顶点内部；索引不同，无脑查 inter_edges
+        #     if p_a == p_b:
+        #         return in_mats[p_a][c_a, c_b]
+        #     if p_a < p_b:
+        #         return inter_edges[(p_a, p_b)][c_a, c_b]
+        #     else:
+        #         return inter_edges[(p_b, p_a)][c_b, c_a]
+
+        # from itertools import combinations_with_replacement
+        # for pi, pj, pk in combinations_with_replacement(range(3), 3):
+        #     for i in range(8):
+        #         start_j = i + 1 if pi == pj else 0
+        #         for j in range(start_j, 8):
+        #             if not are_connected_opt(pi, i, pj, j): continue
+                    
+        #             start_k = j + 1 if pj == pk else 0
+        #             for k in range(start_k, 8):
+        #                 if are_connected_opt(pj, j, pk, k) and are_connected_opt(pk, k, pi, i):
+        #                     c0 = v_parents[pi] * 2 + offsets[i]
+        #                     c1 = v_parents[pj] * 2 + offsets[j]
+        #                     c2 = v_parents[pk] * 2 + offsets[k]
+        #                     all_face_coords.append(np.stack([c0, c1, c2]))
+        # 1. 构建 24x24 全局邻接矩阵 A
+        A = np.zeros((24, 24), dtype=bool)
+        
+        # 填充内部矩阵 (Intra)，使用 | .T 确保无向图的对称性
+        A[0:8, 0:8] = in0 | in0.T
+        A[8:16, 8:16] = in1 | in1.T
+        A[16:24, 16:24] = in2 | in2.T
+        
+        # 填充外部跨体素矩阵 (Inter)
+        A[0:8, 8:16] = m01
+        A[8:16, 0:8] = m01.T
+        A[0:8, 16:24] = m02
+        A[16:24, 0:8] = m02.T
+        A[8:16, 16:24] = m12
+        A[16:24, 8:16] = m12.T
+
+        # 2. 向量化寻找所有的 3-clique (三角形)
+        # 只需要一行代码，判断 i-j, j-k, i-k 是否同时连通
+        valid_mask = A[_IDX_I, _IDX_J] & A[_IDX_J, _IDX_K] & A[_IDX_I, _IDX_K]
+        valid_triangles = _TRI_INDICES[valid_mask] # 过滤出所有形成三角形的节点组合
+
+        # 3. 批量计算物理坐标
+        if len(valid_triangles) > 0:
+            # 将 0~23 的索引还原为父节点索引 (p_idx: 0~2) 和子节点索引 (c_idx: 0~7)
+            p_idx = valid_triangles // 8
+            c_idx = valid_triangles % 8
+            
+            v_parents_arr = np.stack(v_parents) # shape: (3, 3)
+            
+            # 批量算出所有符合条件的三角面的世界坐标，shape: (N, 3, 3)
+            coords = v_parents_arr[p_idx] * 2 + offsets[c_idx]
+            all_face_coords.append(coords)
+                            
+    if not all_face_coords:
+        return np.array([]), np.array([])
+    
+    flat_v = np.concatenate(all_face_coords, axis=0).reshape(-1, 3)
+    V, inv = np.unique(flat_v, axis=0, return_inverse=True)
+    F_raw = inv.reshape(-1, 3)
+
+    # mask_not_degenerate = (F_raw[:, 0] != F_raw[:, 1]) & \
+    #                       (F_raw[:, 1] != F_raw[:, 2]) & \
+    #                       (F_raw[:, 2] != F_raw[:, 0])
+    # F_valid = F_raw[mask_not_degenerate]
+    F_valid = F_raw
+
+    F_sorted = np.sort(F_valid, axis=1)
+    _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True)
+    F = F_valid[unique_face_indices]
+    V = undiscretize(V, num_discrete=resolution)
+    return V, F
+
+def reconstruct_base_mesh(output_tokens, resolution):
+    all_face_coords = []
+    for f_token in output_tokens:
+        v0_c, v1_c, v2_c = f_token
+        c0 = np.array(v0_c[:3], dtype=np.int32)
+        c1 = np.array(v1_c[:3], dtype=np.int32)
+        c2 = np.array(v2_c[:3], dtype=np.int32)
+        all_face_coords.append(np.stack([c0, c1, c2]))
+
+    if not all_face_coords:
+        return np.array([]), np.array([])
+    
+    flat_v = np.concatenate(all_face_coords, axis=0).reshape(-1, 3)
+    V, inv = np.unique(flat_v, axis=0, return_inverse=True)
+    F_raw = inv.reshape(-1, 3)
+
+    # mask_not_degenerate = (F_raw[:, 0] != F_raw[:, 1]) & \
+    #                       (F_raw[:, 1] != F_raw[:, 2]) & \
+    #                       (F_raw[:, 2] != F_raw[:, 0])
+    # F_valid = F_raw[mask_not_degenerate]
+    F_valid = F_raw
+
+    F_sorted = np.sort(F_valid, axis=1)
+    _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True)
+    F = F_valid[unique_face_indices]
+    V = undiscretize(V, num_discrete=resolution)
+    return V, F
+
+
+class MeshCodec:
+    def __init__(self, args):
+        self.max_depth = args.max_depth
+        self.max_resolution = args.max_resolution
+        self.use_hourglass = args.use_hourglass
+        self.max_children = args.max_children
+        
+        self.tokenizer = DynamicMeshTokenizer(resolution=self.max_resolution, max_children=self.max_children)
+
+    def encode_mesh(self, mesh: trimesh.Trimesh, target_depth: int) -> tuple:
+
+        mesh_copy = mesh.copy()
+        if isinstance(mesh_copy, trimesh.Scene):
+            mesh_copy = mesh_copy.dump(concatenate=True)
+        
+        v_restored = discretize(mesh_copy.vertices, num_discrete=2**(target_depth+1))
+        v_restored = undiscretize(v_restored, num_discrete=2**(target_depth+1))
+        
+        # 4. 生成新的量化网格
+        quantized_mesh = trimesh.Trimesh(vertices=v_restored, faces=mesh.faces, process=False)
+        
+        # tree = MatrixOctree(max_depth=self.max_depth)
+        tree = MatrixOctree(max_depth=target_depth+1)
+        tree.build(mesh_copy.vertices, mesh_copy.faces)
+
+        builder = OctreeFaceTokenBuilder(tree, max_children=self.max_children) 
+        
+        raw_t, out_t = builder.encode_level_tokens(depth=target_depth)
+        
+        token_sequence = self.tokenizer.encode(out_t, self.use_hourglass)
+        raw_token_sequence = self.tokenizer.encode_raw(raw_t, self.use_hourglass)
+        
+        return raw_token_sequence, token_sequence, quantized_mesh
+
+    def decode_tokens(self, token_sequence: list, target_depth: int) -> trimesh.Trimesh:
+        if not token_sequence:
+            return None
+
+        resolution = 2 ** (target_depth + 1)
+        out_t = self.tokenizer.decode(token_sequence)
+        v_next, f_next = reconstruct_next_level_optimized(out_t, resolution, max_children=self.max_children)
+        
+        if len(v_next) == 0 or len(f_next) == 0:
+            return None
+            
+        detokenized_mesh = trimesh.Trimesh(vertices=v_next, faces=f_next, process=False)
+        # trimesh.repair.fix_normals(detokenized_mesh)
+        # clean_mesh(detokenized_mesh)
+        return detokenized_mesh
+
+    def decode_raw_tokens(self, raw_token_sequence: list, target_depth:int) -> trimesh.Trimesh:
+        if not raw_token_sequence:
+            return None
+            
+        resolution = 2 ** (target_depth)
+        
+        # [修改] 先使用 tokenizer 从 1D 整数序列解码回 raw_t 结构
+        raw_t = self.tokenizer.decode_raw(raw_token_sequence)
+        
+        v,f = reconstruct_base_mesh(raw_t, resolution)
+        if len(v) == 0 or len(f) == 0:
+            return None
+            
+        detokenized_mesh = trimesh.Trimesh(vertices=v, faces=f, process=False)
+        # trimesh.repair.fix_normals(detokenized_mesh)
+        # clean_mesh(detokenized_mesh)
+        
+        return detokenized_mesh
+    
+    def get_model_config(self) -> dict:
+        return {
+            "total_vocab_size": self.tokenizer.vocab_size,       
+            "coord_vocab_size": self.tokenizer.resolution,       
+            "f_vocab_size": 256,                                 
+            "m_vocab_size": self.tokenizer.m_vocab_size,         
+            "coord_offset": self.tokenizer.coord_offset,
+            "f_offset": self.tokenizer.f_offset,
+            "m_intra_offset": self.tokenizer.m_intra_offset,
+            "m_inter_offset": self.tokenizer.m_inter_offset,
+            "pad_token_id": self.tokenizer.pad_id,
+            "max_children": self.max_children,
+            "input_num_categories": self.tokenizer.resolution+1,
+            "output_num_categories": 1 + 2*2**self.tokenizer.m_bits + 256,     
+        }
+
+# ==========================================
+# 测试重构与无损化流程
+# ==========================================
+if __name__ == "__main__":
+    # 1. 定义输入和输出文件夹
+    INPUT_DIRECTORY = "/root/junkai_lin/mul_res_diffusion/demo/complex_test_data/files"  # 【修改这里】替换为你的模型文件夹路径
+    OUTPUT_DIRECTORY = "./output/complex_child8"
+    os.makedirs(OUTPUT_DIRECTORY, exist_ok=True)
+    import time
+    @dataclass
+    class MeshCodecArgs:
+        max_depth: int = 10          
+        max_resolution: int = 1024   
+        use_hourglass: bool = True   
+        max_children: int = 8       
+        
+    # 2. 初始化 Codec (放在循环外，避免重复加载)
+    args = MeshCodecArgs()
+    codec = MeshCodec(args=args)
+    
+    print(f"[*] Initialized Codec. Total Vocab Size: {codec.tokenizer.vocab_size}")
+    
+    # 3. 获取所有 .obj 和 .glb 文件
+    input_dir_path = Path(INPUT_DIRECTORY)
+    # rglob 会递归查找子文件夹，如果只需要当前层级，可以把 rglob 改成 glob
+    mesh_files = list(input_dir_path.rglob("*.obj")) + list(input_dir_path.rglob("*.glb"))
+    
+    if not mesh_files:
+        print(f"[!] No .obj or .glb files found in {INPUT_DIRECTORY}")
+    
+    # 4. 遍历处理每个文件
+    for idx, file_path in enumerate(mesh_files, 1):
+        obj_name = file_path.stem # 获取不带后缀的文件名
+        print(f"\n{'='*50}")
+        print(f"[{idx}/{len(mesh_files)}] Processing: {file_path.name}")
+        print(f"{'='*50}")
+        file_start_time = time.perf_counter()
+        # try:
+        # 加载并归一化 Mesh
+        raw_mesh = trimesh.load(str(file_path), force="mesh")
+        normalize(raw_mesh)
+        # clean_mesh(raw_mesh)
+        
+        if raw_mesh:
+            out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj")
+            raw_mesh.export(out_path)
+            print(f"[*] Successfully saved GT to: {out_path}")
+        
+        # 开始不同深度 (Depth) 的推理
+        for idx, file_path in enumerate(mesh_files, 1):
+            obj_name = file_path.stem
+            print(f"\n{'='*50}")
+            print(f"[{idx}/{len(mesh_files)}] Processing: {file_path.name}")
+            print(f"{'='*50}")
+            
+            file_start_time = time.perf_counter()  # 记录单个文件处理开始
+
+            raw_mesh = trimesh.load(str(file_path), force="mesh")
+            normalize(raw_mesh)
+            
+            if raw_mesh:
+                out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj")
+                raw_mesh.export(out_path)
+            
+            # 开始不同深度 (Depth) 的推理
+            for depth in [0,1,2,3,4,5,6]:
+                print(f"\n--- Depth {depth} | N={args.max_children} ---")
+                
+                # --- 计时：Encoding ---
+                t0 = time.perf_counter()
+                raw_tokens, tokens, quantized_mesh = codec.encode_mesh(raw_mesh, target_depth=depth)
+                encode_time = time.perf_counter() - t0
+                print(f"[Timer] Encoding: {encode_time:.4f}s")
+                
+                # --- 计时：Decoding (Fine) ---
+                if len(tokens) > 0:
+                    t1 = time.perf_counter()
+                    recon_mesh = codec.decode_tokens(tokens, target_depth=depth)
+                    decode_fine_time = time.perf_counter() - t1
+                    print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s")
+                    
+                    # 导出保存
+                    out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_recon_d{depth}_res{2**(depth+1)}_N{args.max_children}.obj")
+                    recon_mesh.export(out_path)
+
+                # --- 计时：Decoding (Coarse/Raw) ---
+                if len(raw_tokens) > 0:
+                    t2 = time.perf_counter()
+                    raw_recon_mesh = codec.decode_raw_tokens(raw_tokens, target_depth=depth)
+                    decode_coarse_time = time.perf_counter() - t2
+                    print(f"[Timer] Decoding (Coarse): {decode_coarse_time:.4f}s")
+                    
+                    raw_out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_raw_recon_d{depth}_res{2**depth}.obj")
+                    raw_recon_mesh.export(raw_out_path)
+
+                if quantized_mesh:
+                    quant_out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_dis_d{depth+1}_res{2**(depth+1)}.obj")
+                    quantized_mesh.export(quant_out_path)
+
+            file_total_time = time.perf_counter() - file_start_time
+            print(f"\n[*] Finished {file_path.name} in {file_total_time:.2f}s")

mulres_infer/modules/octree_decode.py

@@ -0,0 +1,625 @@
+import gc
+import glob
+import math
+import os
+import random
+import re
+from pathlib import Path
+
+import numpy as np
+import torch
+import trimesh
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from tqdm import tqdm
+
+import importlib.util
+
+from utils.config import load_config
+from dataset.mesh_dataset import mul_encode_mesh, discrete_undiscrete_and_clean
+from models.generator import FactorizedDiffusionTransformer
+from modules.octree_codec import MeshCodec
+from utils.postprocessing import _fill_holes
+from utils.data_utils import normalize
+
+# fix3 编码格式的数据通路(与训练同一份实现, 已抽成独立模块)
+from modules.token_utils import split_fix3_token_list, subcodes_to_flat_token
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    os.environ["PYTHONHASHSEED"] = str(seed)
+
+def _depth_from_resolution(resolution):
+    if resolution <= 0:
+        raise ValueError(f"resolution must be positive, got {resolution}")
+    depth = int(round(math.log2(resolution)))
+    if (1 << depth) != resolution:
+        raise ValueError(f"resolution must be power of 2, got {resolution}")
+    return depth
+
+
+def read_dataset_paths(txt_path):
+    mesh_paths = []
+    with open(txt_path, "r") as f:
+        for line in f:
+            path = line.strip()
+            if path:
+                mesh_paths.append(path)
+    return mesh_paths
+
+
+def raw_stem_to_pair_key(raw_path):
+    stem = Path(raw_path).stem
+    # raw name example: 1(42)_y_f_4244 -> pair key: 1(42)
+    return re.sub(r"_y_f_\d+$", "", stem)
+
+
+def gt_stem_to_pair_key(gt_path):
+    return Path(gt_path).stem
+
+def get_face_count(file_path):
+    """
+    使用正则表达式从文件名中提取面片数量
+    支持格式: *_f{num}_v*.obj 或 *_facenum_{num}_*.glb
+    """
+    filename = Path(file_path).name
+    match = re.search(r'_f(\d+)_v|_facenum_(\d+)', filename)
+    if match:
+        # group(1) 对应 _f(\d+)_v 里的数字
+        # group(2) 对应 _facenum_(\d+) 里的数字
+        return int(match.group(1) or match.group(2))
+    return None
+
+# def build_raw_gt_pairs(raw_txt_path, gt_txt_path):
+#     raw_paths = read_dataset_paths(raw_txt_path)
+#     gt_paths = read_dataset_paths(gt_txt_path)
+
+#     # If two txt paths are the same source, pair each mesh with itself directly.
+#     if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve():
+#         return [(p, p) for p in raw_paths], []
+
+#     gt_by_key = {}
+#     for gt_path in gt_paths:
+#         gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path
+
+#     pairs = []
+#     missing_gt = []
+#     for raw_path in raw_paths:
+#         key = raw_stem_to_pair_key(raw_path)
+#         gt_path = gt_by_key.get(key)
+#         if gt_path is None:
+#             missing_gt.append(raw_path)
+#             continue
+#         pairs.append((raw_path, gt_path))
+
+#     return pairs, missing_gt
+def build_raw_gt_pairs(raw_txt_path, gt_txt_path, min_faces=None, max_faces=None):
+    raw_paths = read_dataset_paths(raw_txt_path)
+    gt_paths = read_dataset_paths(gt_txt_path)
+
+    # ===== 新增过滤逻辑 =====
+    filtered_raw_paths = []
+    for p in raw_paths:
+        if min_faces is None and max_faces is None:
+            filtered_raw_paths.append(p)
+            continue
+            
+        faces = get_face_count(p)
+        if faces is not None:
+            # 检查是否在指定范围内
+            if min_faces is not None and faces < min_faces:
+                continue
+            if max_faces is not None and faces > max_faces:
+                continue
+            filtered_raw_paths.append(p)
+        else:
+            # 如果文件名中没找到面片数特征，默认保留（你也可以改为 continue 丢弃它）
+            filtered_raw_paths.append(p)
+            
+    raw_paths = filtered_raw_paths
+    # ========================
+
+    # If two txt paths are the same source, pair each mesh with itself directly.
+    if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve():
+        return [(p, p) for p in raw_paths], []
+
+    gt_by_key = {}
+    for gt_path in gt_paths:
+        gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path
+
+    pairs = []
+    missing_gt = []
+    for raw_path in raw_paths:
+        key = raw_stem_to_pair_key(raw_path)
+        gt_path = gt_by_key.get(key)
+        if gt_path is None:
+            missing_gt.append(raw_path)
+            continue
+        pairs.append((raw_path, gt_path))
+
+    return pairs, missing_gt
+
+def save_gt_mesh(mesh, out_dir, name):
+    gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False)
+    gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]]
+    gt_mesh.export(os.path.join(out_dir, f"{name}_gt.obj"))
+
+def save_raw_gt_mesh(mesh, out_dir, name):
+    gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False)
+    gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]]
+    gt_mesh.export(os.path.join(out_dir, f"{name}_raw_gt.obj"))
+
+
+def save_stage_input_mesh(mesh, out_dir, name, stage_in_res, stage_out_res, stage_in_depth):
+    coarse_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False)
+    coarse_mesh.vertices = coarse_mesh.vertices[:, [1, 2, 0]]
+    suffix = f"_stage_{stage_in_res}to{stage_out_res}_input_coarse"
+    coarse_mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{stage_in_depth}.obj"))
+
+
+def save_stage_prediction(codec, coord_seq, gen_subcodes, codec_cfg, out_dir, name, target_depth, suffix, fix_holes, raw_dis_mesh=None, gt_dis_mesh=None, remove_degenerate=True):
+    # fix3 数据通路：coord(粗层条件, 含 offset) + 采样得到的 refine 子码(local[0,256)) ->
+    # subcodes_to_flat_token 还原成扁平 60/face token(含 offset) -> codec.decode_tokens 解成 mesh。
+    # 与 acc_v2_fix3.fix3_generate 完全一致；不再做旧 18/face 的 coord/refine 交错。
+    flat = subcodes_to_flat_token(coord_seq, gen_subcodes, codec_cfg)
+    if not flat:
+        return None
+
+    mesh = codec.decode_tokens(flat, target_depth=target_depth - 1)
+    if mesh is None:
+        return None
+
+    # 先剔除退化面(三顶点存在重合的零面积面)，再做 normals/holes 修复，避免退化面干扰后续几何修复。
+    if remove_degenerate:
+        mesh = remove_degenerate_faces(mesh)
+
+    trimesh.repair.fix_normals(mesh)
+    if fix_holes:
+        verts = torch.tensor(mesh.vertices, device="cuda", dtype=torch.float32)
+        faces = torch.tensor(mesh.faces, device="cuda", dtype=torch.int32)
+        new_verts, new_faces = _fill_holes(verts, faces, debug=False, verbose=False, only_fill_holes=True)
+        mesh = trimesh.Trimesh(vertices=new_verts.cpu().numpy(), faces=new_faces.cpu().numpy())
+    
+    return_mesh = mesh.copy()
+    mesh.vertices = mesh.vertices[:, [1, 2, 0]]
+    mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{target_depth}.obj"))
+
+    if raw_dis_mesh is not None:
+        raw_dis_mesh.vertices = raw_dis_mesh.vertices[:, [1, 2, 0]]
+        raw_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_raw_dis_depth{target_depth-1}.obj"))
+    if gt_dis_mesh is not None:
+        gt_dis_mesh.vertices = gt_dis_mesh.vertices[:, [1, 2, 0]]
+        gt_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_gt_dis_depth{target_depth}.obj"))
+    return return_mesh
+
+
+def _decode_subcodes_to_mesh(codec, coord_seq, gen_subcodes, codec_cfg, target_depth):
+    """与 save_stage_prediction 同款 decode(不 fix_normals/fix_holes/不导出)，用于 best-of-N 打分。"""
+    flat = subcodes_to_flat_token(coord_seq, gen_subcodes, codec_cfg)
+    if not flat:
+        return None
+    return codec.decode_tokens(flat, target_depth=target_depth - 1)
+
+
+def _pc_chamfer_score(mesh, pc_xyz, n_sample=4096):
+    """候选 mesh 表面采样点 与 条件点云 pc_xyz 的对称 chamfer(越小越贴合目标点云)。
+    两者同处归一化 frame(mul_encode_mesh 内 normalize + decode 反离散化)，可直接算。
+    退化面(零面积)不会被 sample_surface 采到 -> 天然只评有效曲面；无有效面返回 inf。"""
+    if mesh is None or len(getattr(mesh, "faces", [])) == 0 or len(mesh.vertices) == 0:
+        return float("inf")
+    try:
+        pts, _ = trimesh.sample.sample_surface(mesh, n_sample)
+        pts = np.asarray(pts)
+    except Exception:
+        pts = np.asarray(mesh.vertices)
+    if pts.shape[0] == 0:
+        return float("inf")
+    pts = torch.as_tensor(pts, dtype=torch.float32, device=pc_xyz.device)
+    ref = pc_xyz
+    if ref.shape[0] > n_sample:
+        idx = torch.randperm(ref.shape[0], device=ref.device)[:n_sample]
+        ref = ref[idx]
+    d = torch.cdist(pts, ref)                                   # [Np, Nr]
+    return float((d.min(1).values.mean() + d.min(0).values.mean()).item())
+
+
+def remove_degenerate_faces(mesh):
+    """删除退化面：三角形三个顶点中存在重合顶点(任意两顶点坐标相同 -> 零面积/退化)。
+    按顶点坐标判定(而非仅 index)，因为 octree 解码后不同 index 可能落到同一离散坐标。
+    返回过滤后的新 mesh；若没有退化面则原样返回；若全部退化则返回原 mesh(避免产出空 mesh)。"""
+    if mesh is None or len(getattr(mesh, "faces", [])) == 0 or len(getattr(mesh, "vertices", [])) == 0:
+        return mesh
+    faces = np.asarray(mesh.faces)
+    verts = np.asarray(mesh.vertices)
+    v0, v1, v2 = verts[faces[:, 0]], verts[faces[:, 1]], verts[faces[:, 2]]
+    # 三条边都非零长度(三顶点两两不重合)才保留
+    keep = (
+        np.any(v0 != v1, axis=1)
+        & np.any(v1 != v2, axis=1)
+        & np.any(v0 != v2, axis=1)
+    )
+    if keep.all():
+        return mesh
+    if not keep.any():
+        return mesh
+    return trimesh.Trimesh(vertices=verts.copy(), faces=faces[keep].copy(), process=False)
+
+
+def process_one_mesh(
+    raw_mesh_path,
+    gt_mesh_path,
+    model,
+    codec,
+    codec_cfg,
+    device,
+    pad_id,
+    config,
+    output_folder,
+    stage_resolutions,
+    fix_holes=True,
+    use_gt_pc=True,
+    remove_degenerate=True,
+):
+    import time
+    t_start = time.time()
+    import time
+    t_total_start = time.time()
+    
+    name = Path(raw_mesh_path).stem
+    t_start = time.time()
+    mesh = trimesh.load(raw_mesh_path, force="mesh")
+    mesh.vertices = mesh.vertices[:, [2, 0, 1]]
+    gt_mesh = trimesh.load(gt_mesh_path, force="mesh")
+    gt_mesh.vertices = gt_mesh.vertices[:, [2, 0, 1]]
+    normalize(gt_mesh)
+    normalize(mesh)
+    save_raw_gt_mesh(mesh, output_folder, name)
+    save_gt_mesh(gt_mesh, output_folder, name)
+    t_load = time.time()
+    print(f"[{name}] Load & normalize time: {t_load - t_start:.4f}s")
+
+    stage_depths = [_depth_from_resolution(x) for x in stage_resolutions]
+
+    # First-stage input uses codec encode-decode result at the minimum resolution.
+    min_stage_depth = stage_depths[0]
+    t_start = time.time()
+    init_raw_token, _, init_quantized_mesh = codec.encode_mesh(mesh, target_depth=min_stage_depth)
+    init_mesh = codec.decode_raw_tokens(init_raw_token, target_depth=min_stage_depth)
+    t_init = time.time()
+    print(f"[{name}] Init encode & decode (depth={min_stage_depth}): {t_init - t_start:.4f}s")
+    
+    if init_mesh is not None:
+        current_mesh = init_mesh
+    elif init_quantized_mesh is not None:
+        current_mesh = init_quantized_mesh
+    else:
+        current_mesh = mesh.copy()
+        
+    for stage_idx in range(len(stage_depths) - 1):
+        stage_start_time = time.time()
+        stage_in_depth = stage_depths[stage_idx]
+        stage_out_depth = stage_depths[stage_idx + 1]
+        stage_in_res = stage_resolutions[stage_idx]
+        stage_out_res = stage_resolutions[stage_idx + 1]
+
+        save_stage_input_mesh(
+            mesh=current_mesh,
+            out_dir=output_folder,
+            name=name,
+            stage_in_res=stage_in_res,
+            stage_out_res=stage_out_res,
+            stage_in_depth=stage_in_depth,
+        )
+        
+        t_start = time.time()
+        stage_data = mul_encode_mesh(
+            raw_mesh=current_mesh,
+            gt_mesh=gt_mesh,
+            codec=codec,
+            tokenizer_target_depth=stage_in_depth,
+            pc_num=config.data_processing.pc_num,
+            total_pc_num=config.data_processing.total_pc_num,
+            use_gt_pc=use_gt_pc,
+            clean=False
+        )
+        t_mul_encode = time.time()
+        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | mul_encode_mesh time: {t_mul_encode - t_start:.4f}s")
+
+        pc = stage_data["pc"].unsqueeze(0).to(device)
+
+        # fix3 数据通路(与 acc_v2_fix3.fix3_generate 一致)：
+        #   扁平 60/face token -> 取 coord(粗层条件) -> 因子分解采样 refine 子码。
+        # 这里只用 coord 作为条件，refine 子码全部由模型 sample 生成(GT refine 不参与推理)。
+        coord_list, _ = split_fix3_token_list([stage_data["token"]], pad_id, device, codec_cfg)
+        depth_t = torch.tensor([stage_in_depth], device=device, dtype=torch.long)
+
+        cfg_scale = float(getattr(config.evaluation, "cfg_scale", 1.0))
+        # CFG ramp：linear 时每级采样的最初几步用低 cfg(cfg_start)，按步线性升到 cfg_scale，
+        # 避免高 cfg 在早期步把没把握的位置置信度抬高、过早 commit。constant=旧行为。
+        cfg_schedule = str(getattr(config.evaluation, "cfg_schedule", "constant"))
+        cfg_start = float(getattr(config.evaluation, "cfg_start", 1.0))
+        best_of_n = int(getattr(config.evaluation, "best_of_n", 1))
+        # best-of-N 对抗级联累积误差：每级采 N 个候选 -> 各自 decode -> 选 chamfer-to-pc 最小者(最贴合
+        # 目标点云)作为本级输出喂下一级。点云 pc 是跨级不漂移的全局锚，用它选优能降低单级错误率、抑制误差累积。
+        # N=1 退化为单次贪心采样(确定性, 行为同前)；N>1 需 temperature>0 才能采出多样候选。
+        sample_temp = float(getattr(config.evaluation, "sample_temperature", 1.0)) if best_of_n > 1 else 0.0
+        pc_xyz = stage_data["pc"][:, :3].to(device).float()
+
+        t_start = time.time()
+        gen_subcodes, best_score = None, float("inf")
+        for cand in range(max(1, best_of_n)):
+            gs = model.sample(
+                coord_list,
+                depth_t,
+                pc=pc,
+                T_infer=config.evaluation.diffusion_sample_steps,
+                cfg_scale=cfg_scale,
+                temperature=sample_temp,
+                cfg_schedule=cfg_schedule,
+                cfg_start=cfg_start,
+            )
+            if best_of_n <= 1:
+                gen_subcodes = gs
+                break
+            cand_mesh = _decode_subcodes_to_mesh(codec, coord_list[0], gs[0], codec_cfg, stage_out_depth)
+            score = _pc_chamfer_score(cand_mesh, pc_xyz)
+            if score < best_score:
+                best_score, gen_subcodes = score, gs
+        torch.cuda.synchronize()
+        t_sample = time.time()
+        if best_of_n > 1:
+            print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | best-of-{best_of_n} (T={sample_temp}, cfg={cfg_scale}) "
+                  f"chamfer={best_score:.6f} | sample time: {t_sample - t_start:.4f}s")
+        else:
+            print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | factorized sample time: {t_sample - t_start:.4f}s")
+
+        t_start = time.time()
+        stage_suffix = f"_stage_{stage_in_res}to{stage_out_res}"
+        stage_mesh = save_stage_prediction(
+            codec=codec,
+            coord_seq=coord_list[0],
+            gen_subcodes=gen_subcodes[0],
+            codec_cfg=codec_cfg,
+            out_dir=output_folder,
+            name=name,
+            target_depth=stage_out_depth,
+            suffix=stage_suffix,
+            fix_holes=fix_holes,
+            raw_dis_mesh=stage_data.get("raw_dis_mesh"),
+            gt_dis_mesh=stage_data.get("gt_dis_mesh"),
+            remove_degenerate=remove_degenerate,
+        )
+        t_save_pred = time.time()
+        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | save_stage_prediction (decode + fix) time: {t_save_pred - t_start:.4f}s")
+        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | total stage time: {time.time() - stage_start_time:.4f}s\n")
+
+        if stage_mesh is None:
+            break
+
+        current_mesh = stage_mesh
+        
+    print(f"[{name}] Total process_one_mesh time: {time.time() - t_total_start:.4f}s")
+
+
+class EMA:
+    def __init__(self, model, decay=0.999):
+        self.decay = decay
+        self.shadow = {}
+        self.register(model)
+
+    def register(self, model):
+        self.shadow = {}
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                self.shadow[name] = param.data.clone()
+
+    def load_state_dict(self, state):
+        self.decay = state.get("decay", self.decay)
+        self.shadow = {k: v.clone() for k, v in state["shadow"].items()}
+
+    @torch.no_grad()
+    def copy_to(self, model):
+        """Swap shadow weights into the model in-place (for inference)."""
+        copied, missing = 0, 0
+        for name, param in model.named_parameters():
+            if name in self.shadow:
+                param.data.copy_(self.shadow[name].to(device=param.device, dtype=param.dtype))
+                copied += 1
+            else:
+                missing += 1
+        print(f"[EMA.copy_to] copied {copied} params into model ({missing} not present in shadow)")
+
+
+def run_inference_on_all_checkpoints(model, accelerator, config, device, pad_id, codec, codec_cfg, ema=None):
+    ckpt_dir = "ckpt_fix"
+    cfg_scale = float(getattr(config.evaluation, "cfg_scale", 1.0))   # 采样 CFG 强度(与细分循环 model.sample 一致)
+    output_folder = f"./pretrain_model_mul_more_data_1_to_64_eval_fix_easy2_cfg_{cfg_scale}_step_20"
+
+    # raw_data_txt = "dataset/bpt_res_list.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_retrain_gt/data.txt
+    # gt_data_txt = "dataset/trellis_clean_mesh.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_gt/data.txt
+    # raw_data_txt = "/root/dataset/unique_files_glb_under20000face_compress_glb_with_info_val.txt"
+    # gt_data_txt = "/root/dataset/unique_files_glb_under20000face_compress_glb_with_info_val.txt"
+    raw_data_txt = "/root/dataset/unique_files_glb_under6000face_2degree_30ratio_0.01_with_info_val.txt"
+    gt_data_txt = "/root/dataset/unique_files_glb_under6000face_2degree_30ratio_0.01_with_info_val.txt"
+    
+    # raw_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt"
+    # gt_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt"
+    stage_resolutions = [1,2,4,8,16,32,64]
+    # raw_data_txt = "demo/tripo 1.0/data.txt"
+    # gt_data_txt = "demo/tripo 1.0/data.txt"
+    
+    use_gt_pc=True
+    fix_holes=False
+    # 退化面删除开关：True=导出前剔除三顶点存在重合的零面积面。可在 config.evaluation.remove_degenerate_faces 覆盖。
+    remove_degenerate = bool(getattr(config.evaluation, "remove_degenerate_faces", True))
+    if accelerator.is_main_process:
+        print(f"Using raw data txt: {raw_data_txt}")
+        print(f"Using gt  data txt: {gt_data_txt}")
+        print(f"Searching checkpoints in: {ckpt_dir}")
+
+    # mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=config.data_processing.max_face_num, min_faces=config.data_processing.min_face_num)
+    mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=20000, min_faces=500)
+    if accelerator.is_main_process:
+        print(f"Matched raw/gt pairs: {len(mesh_pairs)}")
+        if missing_gt:
+            print(f"Raw without matched gt: {len(missing_gt)}")
+
+    search_pattern = os.path.join(ckpt_dir, "checkpoint_*_epoch_*_loss.pth")
+    ckpt_files = glob.glob(search_pattern)
+    if not ckpt_files:
+        if accelerator.is_main_process:
+            print("No checkpoints found!")
+        return
+
+    ckpt_data = []
+    for f in ckpt_files:
+        match = re.search(r"checkpoint_(\d+)_epoch_", f)
+        if match:
+            epoch = int(match.group(1))
+            ckpt_data.append((epoch, f))
+    ckpt_data.sort(key=lambda x: x[0])
+
+    # stage_resolutions = [32, 64, 128, 256]
+    
+
+    # stage_resolutions = [128, 256]
+
+    for epoch, ckpt_path in ckpt_data:
+        if accelerator.is_main_process:
+            print(f"\n[{'='*40}]")
+            print(f"Starting inference for epoch {epoch}")
+            print(f"Loading weights from: {ckpt_path}")
+
+        output_folder_ply = os.path.join(output_folder, f"{epoch}_epoch_val_generate_mul")
+        if accelerator.is_main_process:
+            os.makedirs(output_folder_ply, exist_ok=True)
+
+        checkpoint = torch.load(ckpt_path, map_location=device)
+        model_to_load = accelerator.unwrap_model(model)
+        # 加载的是用「新模型代码(AR 连接头 + factorized)」训练出的 ckpt，必须 strict=True：
+        # key/shape 一旦不匹配立即报错，避免 strict=False 把 AR 连接头静默留成随机初始化 -> 输出垃圾还不报警。
+        # (只有训练脚本 acc_v2_fix3 从旧 Linear 头 ckpt 暖启动时才用 load_strict=false，推理这里不允许。)
+        model_to_load.load_state_dict(checkpoint["model_state_dict"], strict=True)
+
+        # Auto-detect EMA: if the ckpt has an `ema_state_dict`, swap those
+        # shadow weights into the model for inference. Otherwise fall back to
+        # the raw `model_state_dict` we already loaded above. No reliance on
+        # config `use_ema` / a separate EMA object — just check the ckpt.
+        ema_state = checkpoint.get("ema_state_dict")
+        shadow = ema_state.get("shadow") if isinstance(ema_state, dict) else None
+        if shadow:
+            with torch.no_grad():
+                copied, skipped = 0, 0
+                for name, param in model_to_load.named_parameters():
+                    if name in shadow:
+                        param.data.copy_(shadow[name].to(device=param.device, dtype=param.dtype))
+                        copied += 1
+                    else:
+                        skipped += 1
+            if accelerator.is_main_process:
+                msg = f"[Epoch {epoch}] EMA shadow detected in ckpt -> applied to model ({copied} params"
+                if skipped:
+                    msg += f", {skipped} model params not in shadow"
+                print(msg + ")")
+        else:
+            if accelerator.is_main_process:
+                print(f"[Epoch {epoch}] no EMA in ckpt -> running with raw model_state_dict weights.")
+
+        with torch.no_grad():
+            # 用 unwrap 后的网络做推理：FactorizedDiffusionTransformer.sample 定义在裸模型上，
+            # DDP 包装层没有 .sample 方法。
+            model_to_load.eval()
+            pbar = tqdm(
+                enumerate(mesh_pairs),
+                total=len(mesh_pairs),
+                desc=f"Infer Epoch {epoch}",
+                disable=not accelerator.is_local_main_process,
+                dynamic_ncols=True,
+            )
+            for line_idx, (raw_mesh_path, gt_mesh_path) in pbar:
+                if line_idx % accelerator.num_processes != accelerator.process_index:
+                    continue
+                try:
+                    process_one_mesh(
+                        raw_mesh_path=raw_mesh_path,
+                        gt_mesh_path=gt_mesh_path,
+                        model=model_to_load,
+                        codec=codec,
+                        codec_cfg=codec_cfg,
+                        device=device,
+                        pad_id=pad_id,
+                        config=config,
+                        output_folder=output_folder_ply,
+                        stage_resolutions=stage_resolutions,
+                        fix_holes=fix_holes,
+                        use_gt_pc=use_gt_pc,
+                        remove_degenerate=remove_degenerate,
+                    )
+                except Exception as e:
+                    print(f"[infer error] raw={raw_mesh_path}, gt={gt_mesh_path}: {str(e)}")
+
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    if accelerator.is_main_process:
+        print("\nBatch inference complete for all checkpoints.")
+
+
+def main(config):
+    print("start init Accelerator...")
+    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+    accelerator = Accelerator(
+        gradient_accumulation_steps=config.accelerator.gradient_accumulation_steps,
+        mixed_precision=config.accelerator.mixed_precision,
+        kwargs_handlers=[ddp_kwargs],
+    )
+    device = accelerator.device
+
+    codec = MeshCodec(args=config.tokenizer)
+    codec_cfg = codec.get_model_config()
+    input_num_categories = codec_cfg["input_num_categories"]   # coord 词表 = resolution + 1 (含 offset)
+    pad_id = codec_cfg["pad_token_id"]
+
+    # fix3 因子分解模型(与 acc_v2_fix3 训练时一致的构造方式)。
+    # rotary 只需覆盖「模型位置」(9/face)，用 config.model.max_seq_len(回退到 max_len)。
+    model_max_seq_len = int(getattr(config.model, "max_seq_len", config.data_processing.max_len))
+    model = FactorizedDiffusionTransformer(
+        coord_vocab_size=input_num_categories,
+        max_seq_len=model_max_seq_len,
+        d_model=config.model.feature_dim,
+        nhead=config.model.num_heads,
+        depth=config.model.hourglass_depth,
+        dropout=config.model.dropout,
+        encoder_freeze=getattr(config.model, "encoder_freeze", True),
+        use_pc=True,
+        time_embed_scale=config.training.diffusion_train_steps,
+        nonempty_row_weight=float(getattr(config.training, "nonempty_row_weight", 1.0)),  # 与 acc_v2_fix3 构造一致(推理不算 loss, 仅对齐)
+    )
+
+    use_ema = getattr(config.training, "use_ema", False)
+    ema_decay = getattr(config.training, "ema_decay", 0.999)
+    ema = EMA(model, decay=ema_decay) if use_ema else None
+
+    model = accelerator.prepare(model)
+
+    run_inference_on_all_checkpoints(
+        model=model,
+        accelerator=accelerator,
+        config=config,
+        device=device,
+        pad_id=pad_id,
+        codec=codec,
+        codec_cfg=codec_cfg,
+        ema=ema,
+    )
+
+
+if __name__ == "__main__":
+    set_seed(42)
+    config = load_config()
+    main(config)

mulres_infer/modules/pipeline.py

@@ -0,0 +1,964 @@
+"""多级"从零生成": res2 -> 4 -> 8 -> ... -> 目标分辨率 的级联生成。
+
+基于 infer_remask_geom_altpc.py 的级联骨架, 但用 infer_single_layer.py 的干净基建 + 新 model.sample
+(align_fm / align_inter / align_per_step)。**不含 remask/requery**(按需求省去), 只保留:
+  - 从最粗分辨率 STAGE[0] 起(GT 八叉树解码作粗起点), 每一级把"上一级预测的 mesh"当输入喂下一级(级联累积);
+  - 每级 model.sample 细分一层, 可选末步跨面一致性对齐(F/intra/inter);
+  - 可选 fill_small_holes 补小洞(不考虑法向)。
+
+与单级(infer_single_layer)的区别: 单级每级喂 GT@in_depth(teacher-forced, 无累积); 本脚本喂上一级预测(真实从零级联)。
+
+⚠ v2 口径(2026-08-17, 实测校准): 所有文件名 = 内容的**真实体素格**(顶点格步实测 2/res)。
+  - fix3: 深度 d 的完整表示真实格 = 2^(d+1); 一个级联 stage = "在深度 d 上把占用补全成格 2^(d+1) 的 mesh",
+    分辨率翻倍发生在下一级重编码——所以 STAGE 每项即该级完成的格, fix3 终点 = 格 32(深度4)。
+  - acc_uni 细分器: 深度 D 真实格 = 2^D(与 fix3 差一档!), 细分终点 SUPERRES_TO=512 = 深度9。
+  - 生成产物 res{r} 与 GT 天花板 gt_res{r} 同格同名, 可直接叠加; 起点若是 raw 占用块命名 raw_res{r}。
+  - 历史目录三代命名: 旧标签(≤08-14, 生成侧=真实格/GT侧=一半), 08-14~17(生成侧虚高一倍), v2(全对齐);
+    新目录一律带 _v2 后缀防覆盖。塌缩自 08-17 起先按坐标焊接(WELD_DEC, 修 UV 接缝撕裂)。
+
+用法:
+  CUDA_VISIBLE_DEVICES=7 STAGE=2,4,8,16,32 DATASET=trellis NCASE=8 STEPS=20 \
+    GEOM_STAGES=8,16,32 SUPERRES=1 SUPERRES_TO=512 python infer_multilevel.py
+"""
+import os, sys, glob as _glob, re as _re, time, math
+from collections import defaultdict
+import numpy as np, torch, trimesh
+CONFIG = os.environ.get("CONFIG", "configs/generator.yaml")
+sys.argv = ["x", "--config", CONFIG]
+from utils.config import load_config
+from dataset.mesh_dataset import mul_encode_mesh, sample_pc, dec
+from dataset.mesh_dataset_altpc import sample_pc_with_degen
+from utils.data_utils import normalize
+from models.generator import REFINE_POS_PER_FACE, NUM_OCTANTS, MASK_ID, TYPE_F, TYPE_INTRA, TYPE_INTER
+from models.generator import FactorizedDiffusionTransformerAltPC as BaseModel
+from modules.octree_codec import MeshCodec as FastMeshCodec
+import utils.distance as metric
+from scipy.spatial import cKDTree
+import importlib.util as _ilu
+from utils.postprocessing import _fill_holes
+_v2_spec = _ilu.spec_from_file_location("fix3_v2", os.path.join(os.path.dirname(os.path.abspath(__file__)), "octree_decode.py"))
+fix3_v2 = _ilu.module_from_spec(_v2_spec); _v2_spec.loader.exec_module(fix3_v2)
+decode_subcodes_to_mesh = fix3_v2._decode_subcodes_to_mesh
+# acc_uni 高分辨率细分器(SUPERRES 用): model_acc + 旧 octree-codebook tokenizer(与 fix3 不同模型/tokenizer)
+from models.subdivider import DiscreteDiffusionTransformer, sample_refine_tokens_with_topology
+from modules.octree_codec_acc import MeshCodec as OctreeMeshCodec
+import pyvista as pv
+from tqdm import tqdm
+TOKENS_PER_FACE = 60
+STAGE = [int(x) for x in os.environ.get("STAGE", "2,4,8,16,32").split(",")]   # 级联阶梯: 每项 = 该级"完成"的真实体素格
+# (2026-08-17 v2 口径: 一级 = 在模型深度 d=log2(r)-1 上把占用补全成格 r 的完整 mesh; 翻倍发生在下一级重编码。
+#  fix3 深度 d 的表示真实格 = 2^(d+1); 生成产物文件 res{r} 与 GT 天花板 gt_res{r} 同格可直接对照)
+DATASET = os.environ.get("DATASET", "trellis")
+NCASE = int(os.environ.get("NCASE", "8")); CASE_STEM = os.environ.get("CASE_STEM", "")
+STEPS = int(os.environ.get("STEPS", "20")); PCT_DEG = os.environ.get("PCT_DEG", "1") == "1"
+ALIGN_FM = os.environ.get("ALIGN_FM", "0") == "1"
+ALIGN_INTER = os.environ.get("ALIGN_INTER", "0") == "1"
+ALIGN_PER_STEP = os.environ.get("ALIGN_PER_STEP", "0") == "1"
+FILL = os.environ.get("FILL", "0") == "1"; FILL_MAX_SIDES = int(os.environ.get("FILL_MAX_SIDES", "4"))
+FILL_MIN_RES = int(os.environ.get("FILL_MIN_RES", "0"))   # >0: FILL=1 的级联内补洞只在"产物真实格 >= 此值"的级做
+                                                         # (默认 0 = 历史行为 stage_i>=4, 在 STAGE=2..64 下即 res32/res64 两级)
+                                                         # 例: FILL_MIN_RES=64 -> 只在 res64 补, res32 不补
+# method=fill: 只用直接补洞(_fill_holes)做级联修复, 在这些"输出分辨率"阶段补洞后喂下一级(默认 8/16/32/64)
+FILL_STAGES = set(int(x) for x in os.environ.get("FILL_STAGES", "8,16,32").split(",") if x.strip())   # 按产物真实格
+SAVE = os.environ.get("SAVE", "1") == "1"
+SAVE_GT_LEVELS = os.environ.get("SAVE_GT_LEVELS", "1") == "1"   # 每 case 导出 GT 经 tokenizer 逐分辨率编码->解码的天花板(gt_res{N}.obj, 真实res命名, 含退化面), 用户要求所有多级生成默认带
+SAVE_PC_LEVELS = os.environ.get("SAVE_PC_LEVELS", "1") == "1"   # 每级导出实际注入的多级点云 pc_target_res{N}.ply(云B, 该级条件) + pc_full.ply(云A, 全分辨率GT云), 用户要求默认带
+EMPTY_BIAS = float(os.environ.get("EMPTY_BIAS", "0"))   # >0: 连接AR头空字节类(0=无连接)logit+bias -> 稀疏化控密度(移植自 infer_remask_geom_altpc, δ=0.5 历史面数比→1.10); requery 补洞期间自动关
+EB_ACTIVE = [True]
+QUAD_FILL = os.environ.get("QUAD_FILL", "0") == "1"     # 每级产物确定性微补洞: 开边图长度3/4简单环(4须双对角线缺失+bbox≤2格)补1-2三角。
+                                                        # 数据背书: GT天花板"无对角线四边洞"=0(20case×res8/16/32), 误补风险极低; 修在级间阻断洞传播
+QUAD_FILL_STATS = [0, 0]                                # [补的三角洞数, 补的四边洞数]
+SAVE_REMASK_VIS = os.environ.get("SAVE_REMASK_VIS", "0") == "1"   # geom 级导出可视化: *_remask_targets.ply(重采目标区=开边1-ring, 橙) / *_remask_changed.ply(重采后实际改变的面, 红)
+
+
+def _face_keys(mesh):
+    v = np.round(np.asarray(mesh.vertices), 6)
+    return [tuple(sorted(map(tuple, v[f]))) for f in np.asarray(mesh.faces)]
+
+
+def save_remask_vis(pre_mesh, post_mesh, out_prefix):
+    """pre=修补前解码, post=修补后。targets: pre 上开边顶点 1-ring 面(≈build_remask目标)橙色;
+    changed: post 上'pre 中不存在'的面红色。PLY 带面色, y-up 帧与 obj 一致。"""
+    GRAY, ORANGE, RED = [200, 200, 200, 255], [255, 140, 0, 255], [220, 30, 30, 255]
+    def yup(m):
+        mm = m.copy(); mm.vertices = mm.vertices[:, [1, 2, 0]]; return mm
+    # targets on pre
+    v = np.asarray(pre_mesh.vertices); f = np.asarray(pre_mesh.faces)
+    key = np.round(v, 6); uniq, inv = np.unique(key, axis=0, return_inverse=True)
+    f2 = inv[f]
+    fg = f2[(f2[:, 0] != f2[:, 1]) & (f2[:, 1] != f2[:, 2]) & (f2[:, 0] != f2[:, 2])]
+    e = np.sort(np.vstack([fg[:, [0, 1]], fg[:, [1, 2]], fg[:, [0, 2]]]), axis=1)
+    ue, cnt = np.unique(e, axis=0, return_counts=True)
+    open_verts = set(ue[cnt == 1].reshape(-1).tolist())
+    tmask = np.array([any(x in open_verts for x in row) for row in f2])
+    mp = yup(pre_mesh); mp.visual.face_colors = np.where(tmask[:, None], ORANGE, GRAY)
+    mp.export(out_prefix + "_remask_targets.ply")
+    # changed on post
+    pre_keys = set(_face_keys(pre_mesh))
+    cmask = np.array([k not in pre_keys for k in _face_keys(post_mesh)])
+    mq = yup(post_mesh); mq.visual.face_colors = np.where(cmask[:, None], RED, GRAY)
+    mq.export(out_prefix + "_remask_changed.ply")
+    print(f"  [REMASK_VIS] {os.path.basename(out_prefix)}: 目标面 {int(tmask.sum())}/{len(tmask)}, 改变面 {int(cmask.sum())}/{len(cmask)}")
+
+
+def quad_fill_mesh(mesh, res):
+    """长度3/4边界环确定性补洞(见 QUAD_FILL 注释)。返回(焊接后的)新 mesh。"""
+    v = np.asarray(mesh.vertices); f = np.asarray(mesh.faces)
+    key = np.round(v, 6); uniq, first_idx, inv = np.unique(key, axis=0, return_index=True, return_inverse=True)
+    f2 = inv[f]; f2 = f2[(f2[:, 0] != f2[:, 1]) & (f2[:, 1] != f2[:, 2]) & (f2[:, 0] != f2[:, 2])]   # 图只用正常面; 输出保留原mesh全部面(含退化载体)
+    e = np.sort(np.vstack([f2[:, [0, 1]], f2[:, [1, 2]], f2[:, [0, 2]]]), axis=1)
+    ue, cnt = np.unique(e, axis=0, return_counts=True)
+    all_edges = set(map(tuple, ue))
+    open_e = ue[cnt == 1]
+    from collections import defaultdict
+    adj = defaultdict(set)
+    for a, b in open_e:
+        adj[int(a)].add(int(b)); adj[int(b)].add(int(a))
+    new_faces = []; seen = set(); step = 2.0 / res
+    for s in list(adj):
+        if s in seen:
+            continue
+        comp = [s]; seen.add(s); st = [s]
+        while st:
+            x = st.pop()
+            for y in adj[x]:
+                if y not in seen:
+                    seen.add(y); comp.append(y); st.append(y)
+        n = len(comp); ne = sum(len(adj[x]) for x in comp) // 2
+        if not (ne == n and all(len(adj[x]) == 2 for x in comp)):
+            continue                                   # 非简单环不碰
+        if n == 3:
+            new_faces.append(tuple(comp)); QUAD_FILL_STATS[0] += 1
+        elif n == 4:
+            ring = [comp[0]]; prev = None
+            while len(ring) < 4:
+                nxt = [y for y in adj[ring[-1]] if y != prev and y in comp]
+                prev = ring[-1]; ring.append(nxt[0])
+            d1 = (min(ring[0], ring[2]), max(ring[0], ring[2]))
+            d2 = (min(ring[1], ring[3]), max(ring[1], ring[3]))
+            if d1 in all_edges or d2 in all_edges:
+                continue                               # 已有对角线(半填充/薄结构合法边界)不碰
+            if (uniq[comp].max(0) - uniq[comp].min(0)).max() > 2 * step + 1e-6:
+                continue                               # 跨区大环不碰(只修 raw face 内量级)
+            if np.linalg.norm(uniq[ring[0]] - uniq[ring[2]]) <= np.linalg.norm(uniq[ring[1]] - uniq[ring[3]]):
+                new_faces += [(ring[0], ring[1], ring[2]), (ring[0], ring[2], ring[3])]
+            else:
+                new_faces += [(ring[1], ring[2], ring[3]), (ring[1], ring[3], ring[0])]
+            QUAD_FILL_STATS[1] += 1
+    if not new_faces:
+        return mesh
+    nf = first_idx[np.array(new_faces)]   # 焊接索引 -> 原顶点索引; 原 mesh 零改动只追加
+    return trimesh.Trimesh(vertices=v, faces=np.vstack([f, nf]), process=False)
+DECIMATE_RATIO = float(os.environ.get("DECIMATE_RATIO", "0.5"))  # GT 预塌缩比例(pyvista decimate 目标削减比, 0=不塌缩)
+DEC_FILL = os.environ.get("DEC_FILL", "0") == "1"           # 塌缩后先补洞(TRELLIS _fill_holes only_fill_holes=pymeshfix)再作输入(条件PC/天花板/粗起点); 指标参照仍用未补洞塌缩GT, 与基线可比
+# 级间是否按预测 mesh 自身 bbox 重新归一化(2026-08-22 修复, 默认关): 级间 mesh 是上一级解码产物,
+# 天生在规范 [-1,1] 帧; 重归一化在预测带极值杂散面时把脚手架帧拉离 GT/点云帧 -> z 轴撕裂
+# (1(10)/1(101): res4 顶层 3-4% 杂散面 -> z中心+0.25/缩放4/3 -> 本体父层劈成 {-0.75,+0.25} 中间断层,
+#  res8 出现 -0.625 悬空板 + -0.375/-0.125 全空, 级联放大成面数腰斩)。=1 复现历史(≤08-22)行为。
+CASCADE_RENORM = os.environ.get("CASCADE_RENORM", "0") == "1"
+# 每级采样方法(可多选逗号): plain / consensus(align) / log / log_fmm / geom / geom_fmm / fill。ALIGN_* 仍对 plain 生效(向后兼容)。
+METHODS = [m.strip() for m in os.environ.get("METHODS", "plain").split(",") if m.strip()]
+MULTI_PC = os.environ.get("MULTI_PC", "1") == "1"           # 0=单点云(pc_target=None, focal 单PC ckpt 必须)
+SINGLE_TARGET_PC = os.environ.get("SINGLE_TARGET_PC", "0") == "1"   # 1=云B作唯一云每层恒注入(sfocal 线口径); 需 MULTI_PC=1 以构建云B
+NITER_GEOM = int(os.environ.get("NITER_GEOM", "3"))         # geom: 每级终点 remask+requery 轮数
+KREQ_GEOM = int(os.environ.get("KREQ_GEOM", "5"))           # geom: requery 每轮步数(5=甜点, 质量≈10 快~3x)
+MAXFACE_REMASK = int(os.environ.get("MAXFACE", "1200"))     # geom: 每轮最多重掩多少 face 的连接
+REMASK_UNION = os.environ.get("REMASK_UNION", "1") == "1"   # geom: 开边两父 token 取并集(1-ring)
+REMASK_LOOP_MAX = int(os.environ.get("REMASK_LOOP_MAX", "0"))   # >0: 只对边界环长<=K的开边做 remask(外科修小洞; 0=全部开边)
+REMASK_GROUPS = int(os.environ.get("REMASK_GROUPS", "1"))       # >1: 每轮候选面按 BFS 序分 G 批逐批 requery(后批看得到前批修复, 用户方案); 1=整轮一次性
+RUN_SEED = int(os.environ.get("RUN_SEED", "0"))                 # >0: 整体偏移随机种子, 用于"同配置多次独立采样各自保存"(看运行间方差); 0=与历史逐位一致
+BO_N = int(os.environ.get("BO_N", "1"))                         # >1: 级内 best-of-N——重采条件云B+换torch种子造候选(实测±40开边方差来源), 结构惩罚(开边+10*NM%)选优
+BO_STAGES = set(int(x) for x in os.environ.get("BO_STAGES", "32").split(",") if x.strip())   # 只在这些产物格做 best-of-N
+REQUERY_EMPTY_BIAS = float(os.environ.get("REQUERY_EMPTY_BIAS", "0"))   # requery(缺陷区重问)期的空行偏置: 负值="不许再答无连接"推向补上; 0=关(现状)
+GEOM_STAGE_PARAMS = {}                                          # 按产物格覆盖 remask 轮数x步数, 如 "32:6x10"(res32 用 6x10, 其余用方法默认); 扫参结论: res8/16 甜点 3x5, res32 甜点 6x10
+for _kv in (x for x in os.environ.get("GEOM_STAGE_PARAMS", "").split(",") if x.strip()):
+    _r, _p = _kv.split(":"); _a, _b = _p.lower().split("x")
+    GEOM_STAGE_PARAMS[int(_r)] = (int(_a), int(_b))
+REMASK_DEGEN = os.environ.get("REMASK_DEGEN", "0") == "1"   # geom: 额外纳入"有退化三角的 raw face"一起 remask
+REMASK_NM = os.environ.get("REMASK_NM", "0") == "1"         # geom: 非流形边(覆盖>2)的 1-ring face 也纳入 remask+requery。
+# 实测(val20@res32 A/B, ab_remask_nm/)净负: 开边131->171/非流形12.9->14.5%双变差, 调大MAXFACE无关(cap不触顶)、
+# 加迭代不救; 模型对 NM 区域重采仍给过连接答案, 反在原闭合处造新开边 -> 默认关, 仅留开关。
+GEOM_STAGES = set(int(x) for x in os.environ.get("GEOM_STAGES", "").split(",") if x.strip())  # geom: 只在这些"产物真实格"级做 remask(空=每级都做)
+GT_START_RES = int(os.environ.get("GT_START_RES", "0"))     # >0: 用 GT 编码->解码天花板@该真实分辨率作级联起点, 跳过更粗级(暴露偏差上界实验); 须在 STAGE 中
+EXT_START_DIR = os.environ.get("EXT_START_DIR", "")          # 非空: 外部预测 mesh 作级联起点, 文件 {case_stem}{EXT_START_SUFFIX}, 保存帧(与 gt_original.obj 同帧), 只跑有该文件的 case
+EXT_START_RES = int(os.environ.get("EXT_START_RES", "32"))   # 外部起点的真实分辨率(须在 STAGE 中)
+EXT_START_SUFFIX = os.environ.get("EXT_START_SUFFIX", "_pred.obj")
+assert not (GT_START_RES and EXT_START_DIR), "GT_START_RES 与 EXT_START_DIR 互斥"
+_START_RES = GT_START_RES or (EXT_START_RES if EXT_START_DIR else 0)
+# 起点(GT天花板/外部预测)已是"完成的格 _START_RES mesh" -> 只跑更细的级(v2 语义: 不再重做起点那一级)
+STAGE_RUN = [r for r in STAGE if r > _START_RES] if _START_RES else STAGE   # 实际执行的阶梯(每项=完成该格)
+OFF = np.stack([(np.arange(8) >> k) & 1 for k in range(3)], 1).astype(np.int64)      # [8,3]
+TRI = np.array([(i, j, k) for i in range(24) for j in range(i + 1, 24) for k in range(j + 1, 24)])  # C(24,3)
+POS = {(0, 0): 3, (1, 1): 4, (2, 2): 5, (0, 1): 6, (0, 2): 7, (1, 2): 8}
+# ---- SUPERRES: fix3 级联到 STAGE[-1] 后, 用 acc_uni 细分器继续细分到更高分辨率 ----
+SUPERRES = os.environ.get("SUPERRES", "0") == "1"                # 是否开启更高分辨率细分
+SUPERRES_TO = int(os.environ.get("SUPERRES_TO", "512"))          # 细分终点真实格(从 STAGE[-1] 倍增: 32->64->128->256->512)
+# acc_uni 与 fix3 口径不同: acc_uni 深度 D 的真实格 = 2^D(fix3 是 2^(d+1)); 细分用 ACC_UN._depth_from_resolution(log2)
+SUPERRES_STEPS = int(os.environ.get("SUPERRES_STEPS", "10")); SUPERRES_CFG = float(os.environ.get("SUPERRES_CFG", "1.0"))
+SUPERRES_FIX_HOLES = os.environ.get("SUPERRES_FIX_HOLES", "1") == "1"
+SR_FILL_IN = os.environ.get("SR_FILL_IN", "1") == "1"            # 细分入口先 pymeshfix 补洞(对齐旧 cross512_fillall: 粗输入进细分器前补洞)
+SUPERRES_CONFIG = os.environ.get("SUPERRES_CONFIG", "configs/subdivider.yaml")
+SUPERRES_CKPT = os.environ.get("SUPERRES_CKPT", os.path.join(os.path.dirname(os.path.abspath(__file__)), "ckpt", "subdivider.pth"))
+FINAL_RES = SUPERRES_TO if SUPERRES else STAGE[-1]
+NSAMP = int(os.environ.get("NSAMP", "1024")); TAU = 2.0 / FINAL_RES   # v2: FINAL_RES=512 真实格 -> TAU=2/512, 数值与历史协议一致; NSAMP=1024 严格对齐 metric_mul_dense_mesh_v2
+
+BASE_DIR = os.environ.get("BASE_DIR", "pretrained_model/gobjaverse_fix3_octant_factorized_1_to_1024_ce_altpc_decoded_deg_w1")
+_ckpt_cands = sorted(_glob.glob(os.path.join(BASE_DIR, "checkpoint_*.pth")))
+CKPT = os.environ.get("RAIL_CKPT") or (_ckpt_cands[-1] if _ckpt_cands else "")
+if not CKPT or not os.path.exists(CKPT):
+    raise SystemExit(f"未找到生成器权重: 用 scripts/infer.py --ckpt <path> 指定, 或设环境变量 RAIL_CKPT\n"
+                     f"  (已查找: {os.path.join(BASE_DIR, 'checkpoint_*.pth')})")
+_ep_m = _re.search(r"checkpoint_(\d+)_", os.path.basename(CKPT))
+EPOCH = _ep_m.group(1) if _ep_m else os.path.splitext(os.path.basename(CKPT))[0]   # 权重文件名不含 epoch 时用文件名作标签
+_RUN_TAG = f"ep{EPOCH}" if _ep_m else EPOCH
+TRELLIS = os.environ.get("TRELLIS_DIR", "")      # 可选: 内置数据集快捷方式; 一般用 DATASET_DIR 指定任意目录
+VAL = os.environ.get("VAL_LIST", "")             # 可选: DATASET=val 时的列表文件
+CAR = "1307d3253490441ab12f765c23fda2c5"
+TOYS4K = os.environ.get("TOYS4K_DIR", "")        # 可选: 域外玩具集
+_tag = ("_fm" if ALIGN_FM else "") + ("_inter" if ALIGN_INTER else "") + ("_perstep" if ALIGN_PER_STEP else "") + ("_fill" if FILL else "") + (f"_sr{SUPERRES_TO}" if SUPERRES else "") + (f"_gtstart{GT_START_RES}" if GT_START_RES else "") + (f"_ext{EXT_START_RES}" if EXT_START_DIR else "") + ("_decfill" if DEC_FILL else "") + "_v2"   # v2=2026-08-17 真实格命名+焊接塌缩口径, 防与历史目录同名覆盖
+OUT = os.path.join(os.environ.get("OUT_ROOT") or os.path.dirname(CKPT),
+                   f"multilevel_{_RUN_TAG}_res{FINAL_RES}_s{STEPS}_{DATASET}{_tag}") + os.environ.get("OUT_SUFFIX", "")
+os.makedirs(OUT, exist_ok=True)
+
+
+WELD_DEC = os.environ.get("WELD_DEC", "1") == "1"   # 塌缩前按坐标焊接顶点。修根本bug: GLB 的 UV 接缝重复顶点
+                                                     # 使 vtkQuadricDecimation 把接缝当边界、两侧独立塌缩 -> 沿接缝撕出
+                                                     # 数千条真裂缝(1(11): 3699条 -> 焊接后 0)。0=旧撕裂行为(复现历史用)
+def decimate_mesh(vertices, faces, simplify=True, simplify_ratio=0.9, verbose=True):
+    # Simplify
+    if simplify and simplify_ratio > 0:
+        if WELD_DEC:
+            _key = np.round(np.asarray(vertices, dtype=np.float64), 6)
+            _uniq, _inv = np.unique(_key, axis=0, return_inverse=True)
+            _f = _inv[np.asarray(faces)]
+            _f = _f[(_f[:, 0] != _f[:, 1]) & (_f[:, 1] != _f[:, 2]) & (_f[:, 0] != _f[:, 2])]
+            vertices, faces = _uniq, _f
+        mesh = pv.PolyData(vertices, np.concatenate([np.full((faces.shape[0], 1), 3), faces], axis=1))
+        mesh = mesh.decimate(simplify_ratio, progress_bar=verbose)
+        vertices, faces = mesh.points, mesh.faces.reshape(-1, 4)[:, 1:]
+        if verbose:
+            tqdm.write(f'After decimate: {vertices.shape[0]} vertices, {faces.shape[0]} faces')
+    return vertices, faces
+
+def depth_from_resolution(res):
+    """真实体素分辨率 -> 八叉树深度(res = 2^(depth+1), 与训练/单级脚本口径一致)。
+    旧口径(2026-08-14 前)是级别标签 res = 2^depth; 旧输出目录文件名按旧标签, 见文件头说明。"""
+    d = int(round(np.log2(res))) - 1
+    assert d >= 0 and 2 ** (d + 1) == res, f"res={res} 须为 >=2 的 2 的幂"
+    return d
+
+
+def mesh_open_edge_stats(mesh):
+    """网格拓扑统计。返回 (开边数, 面数, 非流形边占比%)。
+    先 merge_vertices 合并重合顶点再统计(否则补洞类输出的重复顶点会把开边/非流形虚高)。"""
+    if mesh is None or len(getattr(mesh, "faces", [])) == 0:
+        return -1, 0, 0.0
+    m = mesh.copy(); m.merge_vertices()
+    edges = np.sort(np.asarray(m.edges), axis=1)
+    _, counts = np.unique(edges, axis=0, return_counts=True)
+    return int((counts == 1).sum()), len(m.faces), float((counts > 2).mean() * 100)
+
+
+def fill_small_holes(mesh, max_sides=4):
+    """补 <=max_sides 边的小闭环洞(**不考虑法向**, 直接扇形三角化连上)。返回 (新mesh, 补洞数)。"""
+    if mesh is None or len(getattr(mesh, "faces", [])) == 0:
+        return mesh, 0
+    verts = np.asarray(mesh.vertices); faces = np.asarray(mesh.faces)
+    directed_edges = np.concatenate([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=0)
+    undirected_key = np.sort(directed_edges, axis=1)
+    _, inverse, count = np.unique(undirected_key, axis=0, return_inverse=True, return_counts=True)
+    boundary_edges = directed_edges[count[inverse] == 1]                 # 无向边只被 1 面用 = 边界
+    if len(boundary_edges) == 0:
+        return mesh, 0
+    next_vertex = {}                                                     # a -> [b,...] 边界后继
+    for a, b in boundary_edges:
+        next_vertex.setdefault(int(a), []).append(int(b))
+    new_faces = []; n_filled = 0; visited = set()
+    for start in list(next_vertex.keys()):
+        if start in visited:
+            continue
+        loop = [start]; cur = start; closed = False
+        for _ in range(max_sides + 1):
+            outs = next_vertex.get(cur)
+            if not outs or len(outs) != 1:                              # 断裂 / 非流形边界顶点 -> 放弃该环
+                break
+            nb = outs[0]
+            if nb == start:
+                closed = True; break
+            if nb in loop:                                             # 自交 -> 放弃
+                break
+            loop.append(nb); cur = nb
+        if closed and 3 <= len(loop) <= max_sides:
+            for v in loop:
+                visited.add(v)
+            v0 = loop[0]
+            for i in range(1, len(loop) - 1):                          # 扇形三角化(不考虑法向/绕序)
+                new_faces.append([v0, loop[i], loop[i + 1]])
+            n_filled += 1
+    if not new_faces:
+        return mesh, 0
+    faces_out = np.concatenate([faces, np.asarray(new_faces, dtype=faces.dtype)], axis=0)
+    return trimesh.Trimesh(vertices=verts, faces=faces_out, process=False), n_filled
+
+
+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[9*nface,8])。位置 [F0..2, M00 M11 M22, M01 M02 M12]。"""
+    f_offset = codec_cfg["f_offset"]; intra_offset = codec_cfg["intra_row_offset"]; inter_offset = codec_cfg["inter_row_offset"]
+    coord_list, subcode_list = [], []
+    for token in token_list:
+        token = _trim_trailing_pad(token.to(device=device, dtype=torch.long), pad_id)
+        if token.numel() % TOKENS_PER_FACE != 0:
+            raise ValueError(f"token 长度 {token.numel()} 不是 {TOKENS_PER_FACE} 的倍数")
+        token = token.view(-1, TOKENS_PER_FACE); nface = token.size(0)
+        coord = token[:, :9].reshape(-1)
+        occupancy = token[:, 9:12] - f_offset
+        intra = (token[:, 12:36] - intra_offset).view(nface, 3, NUM_OCTANTS)
+        inter = (token[:, 36:60] - inter_offset).view(nface, 3, NUM_OCTANTS)
+        subcode = torch.zeros(nface, REFINE_POS_PER_FACE, NUM_OCTANTS, dtype=torch.long, device=device)
+        subcode[:, 0:3, 0] = occupancy; subcode[:, 3:6, :] = intra; subcode[:, 6:9, :] = inter
+        subcode_list.append(subcode.reshape(-1, NUM_OCTANTS)); coord_list.append(coord)
+    return coord_list, subcode_list
+
+
+def _metric_normalize(mesh):
+    """metric_mul_dense_mesh_v2 协议: 每 mesh 独立归一化到 [-0.5, 0.5] 盒。"""
+    m = mesh.copy()
+    bounds = np.array([m.vertices.min(axis=0), m.vertices.max(axis=0)])
+    m.vertices = m.vertices - (bounds[0] + bounds[1])[None, :] / 2
+    m.vertices = m.vertices / (bounds[1] - bounds[0]).max()
+    m.vertices = m.vertices.clip(-0.5, 0.5)
+    return m
+
+
+def chamfer_hd_fscore(pred_mesh, ref_mesh):
+    """返回 (CD-L2*1e3, HD*1e3, F@2voxel, |NC|)。
+    CD/HD/|NC| 按 metric_mul_dense_mesh_v2 协议: 每 mesh 独立归一化到 [-0.5,0.5] 后采样(带法向),
+    |NC| = 双向最近邻法向 |点积| 均值。F-score 沿用原 [-1,1] 管线帧 + TAU=4/真实FINAL_RES(=旧 2/标签res, 阈值数值与历史可比)。"""
+    if pred_mesh is None or ref_mesh is None or len(getattr(pred_mesh, "faces", [])) == 0:
+        return -1., -1., -1., -1.
+    ref_pts = np.asarray(trimesh.sample.sample_surface(ref_mesh, NSAMP)[0])
+    pred_pts = np.asarray(trimesh.sample.sample_surface(pred_mesh, NSAMP)[0])
+    recall = (cKDTree(pred_pts).query(ref_pts)[0] < TAU).mean()
+    precision = (cKDTree(ref_pts).query(pred_pts)[0] < TAU).mean()
+    fscore = 200 * precision * recall / (precision + recall + 1e-9)
+    # ---- v2 协议: 独立归一化 [-0.5,0.5] + 带法向采样 ----
+    ref_m, pred_m = _metric_normalize(ref_mesh), _metric_normalize(pred_mesh)
+    ref_p, ref_fi = trimesh.sample.sample_surface(ref_m, NSAMP)
+    pred_p, pred_fi = trimesh.sample.sample_surface(pred_m, NSAMP)
+    ref_p, pred_p = np.asarray(ref_p), np.asarray(pred_p)
+    ref_n, pred_n = np.asarray(ref_m.face_normals[ref_fi]), np.asarray(pred_m.face_normals[pred_fi])
+    chamfer_l2, hausdorff, _ = metric.calculate_distances_numpy(ref_p, pred_p)
+    i1 = cKDTree(pred_p).query(ref_p)[1]                      # GT->Pred 最近邻
+    i2 = cKDTree(ref_p).query(pred_p)[1]                      # Pred->GT 最近邻
+    nc_abs = 0.5 * (np.abs(np.einsum('ij,ij->i', ref_n, pred_n[i1])).mean()
+                    + np.abs(np.einsum('ij,ij->i', pred_n, ref_n[i2])).mean())
+    return chamfer_l2 * 1e3, hausdorff * 1e3, fscore, nc_abs
+
+
+# ===== geom (方案A 几何 remask+requery, 移植自 infer_single_layer.py; requery 支持 consensus) =====
+def open_edges_and_faces(coord_seq, sub_seq, coff):
+    """coord 空间复现解码(全向量化, 无逐面 Python 循环) ->
+    (open_edges set, parent2faces dict, degen_faces set, nonmanifold_edges set)。
+    开边=覆盖数1, 非流形边=覆盖数>2(与 mesh_open_edge_stats 同判据)。与原逐面版逐位等价, 快 ~10x。"""
+    nface = coord_seq.numel() // 9
+    P = (coord_seq.view(nface, 9).reshape(nface, 3, 3).cpu().numpy() - coff).astype(np.int64)   # [nface,3(父),3(xyz)]
+    S = sub_seq.view(nface, 9, 8).cpu().numpy().astype(np.int64)                                 # [nface,9,8]
+    ar = np.arange(8)
+    fa = ((S[:, 0:3, 0][:, :, None] >> ar[None, None]) & 1).astype(bool)                         # [nface,3,8] 各父 F
+    A = np.zeros((nface, 24, 24), bool)                                                          # 24=3父x8子 邻接
+    for (pi, pj), pos in POS.items():
+        M = ((S[:, pos, :][:, :, None] >> ar[None, None]) & 1).astype(bool)                      # [nface,8(子a),8(bit b)]
+        M = M & fa[:, pi, :, None] & fa[:, pj, None, :]
+        if pi == pj:
+            A[:, pi * 8:pi * 8 + 8, pi * 8:pi * 8 + 8] = M | M.transpose(0, 2, 1)
+        else:
+            A[:, pi * 8:pi * 8 + 8, pj * 8:pj * 8 + 8] = M
+            A[:, pj * 8:pj * 8 + 8, pi * 8:pi * 8 + 8] = M.transpose(0, 2, 1)
+    vm = A[:, TRI[:, 0], TRI[:, 1]] & A[:, TRI[:, 1], TRI[:, 2]] & A[:, TRI[:, 0], TRI[:, 2]]     # [nface,nTRI] 3-clique
+    # parent2faces: 每个父坐标 -> 所属 face 集(所有 face, 与 tris 无关)
+    parent2faces = defaultdict(set)
+    Pl = P.reshape(-1, 3); fl = np.repeat(np.arange(nface), 3)
+    for i in range(Pl.shape[0]):
+        parent2faces[(int(Pl[i, 0]), int(Pl[i, 1]), int(Pl[i, 2]))].add(int(fl[i]))
+    fidx, tidx = np.nonzero(vm)
+    if fidx.size == 0:
+        return set(), parent2faces, set(), set()
+    pnode = TRI[tidx] // 8; cnode = TRI[tidx] % 8                                                 # [N,3] 父/子索引
+    Pf = P[fidx]                                                                                  # [N,3,3]
+    tri = np.take_along_axis(Pf, np.broadcast_to(pnode[:, :, None], (fidx.size, 3, 3)), axis=1) * 2 + OFF[cnode]  # [N,3,3]
+    cross = np.cross(tri[:, 1] - tri[:, 0], tri[:, 2] - tri[:, 0])                                # 面积 x2
+    degen_faces = set(np.unique(fidx[np.abs(cross).sum(1) == 0]).tolist())                        # 有零面积三角的 raw face
+    # 开边: 全局边覆盖计数==1。点->单调标量(保持字典序), 边->(小,大)标量对
+    B = int(tri.max()) + 1
+    scal = tri[:, :, 0] * B * B + tri[:, :, 1] * B + tri[:, :, 2]                                 # [N,3] 每顶点标量
+    # ---- 修复(2026-08-02, 同 infer_single_layer): 只在「唯一且非退化」三角形上数边覆盖 ----
+    # 同一子三角形可被多个父块重复编码, decode 后 mesh 合并成一个 -> 实为开边, 旧统计判"已覆盖",
+    # build_remask 永远不掩其直接父面; 重复覆盖同样虚增非流形计数。去重+去退化后与 mesh 拓扑对齐。
+    nondeg = np.abs(cross).sum(1) > 0
+    sig = np.sort(scal[nondeg], axis=1)
+    _, first_idx = np.unique(sig, axis=0, return_index=True)
+    scal_u = scal[nondeg][first_idx]                                                              # [Nu,3] 唯一非退化三角
+    sa = scal_u[:, [0, 1, 2]].reshape(-1); sb = scal_u[:, [1, 2, 0]].reshape(-1)                  # [3Nu] 三边端���
+    key = np.stack([np.minimum(sa, sb), np.maximum(sa, sb)], axis=1)                              # [3Nu,2] 无向边
+    uk, cnt = np.unique(key, axis=0, return_counts=True)
+    ok = uk[cnt == 1]                                                                             # 开边 (小,大) 标量对
+    nk = uk[cnt > 2]                                                                              # 非流形边 (覆盖>2)
+    dec = lambda s: (int(s // (B * B)), int((s // B) % B), int(s % B))
+    open_edges = {(dec(mn), dec(mx)) for mn, mx in ok}
+    nm_edges = {(dec(mn), dec(mx)) for mn, mx in nk}
+    return open_edges, parent2faces, degen_faces, nm_edges
+
+def build_remask(coord_seq, sub_seq, coff, total, cap=None, use_degen=False):
+    """几何缺陷驱动选 remask 目标: 触及缺陷边(开边 + REMASK_NM 时非流形边)的 face
+    (两端点 parent 的1-ring)的连接位置(3-8)置 MASK。
+    use_degen=True: 额外把"有退化三角(零面积)的 raw face"也纳入 remask(直接按 face)。"""
+    if cap is None: cap = MAXFACE_REMASK
+    oe, p2f, degen, nm = open_edges_and_faces(coord_seq, sub_seq, coff)
+    if REMASK_LOOP_MAX and oe:   # 只用"边界环长<=K"的开边作种子: 小洞外科修复, 薄结构大环/复杂边界不触发
+        from collections import defaultdict as _dd
+        _adj = _dd(set)
+        for (c1, c2) in oe:
+            _adj[c1].add(c2); _adj[c2].add(c1)
+        _seen = set(); _keep = set()
+        for _s in list(_adj):
+            if _s in _seen: continue
+            _comp = [_s]; _seen.add(_s); _st = [_s]
+            while _st:
+                _x = _st.pop()
+                for _y in _adj[_x]:
+                    if _y not in _seen: _seen.add(_y); _comp.append(_y); _st.append(_y)
+            _ne = sum(len(_adj[_x]) for _x in _comp) // 2
+            if _ne <= REMASK_LOOP_MAX:
+                _keep.update(_comp)
+        oe = {(c1, c2) for (c1, c2) in oe if c1 in _keep and c2 in _keep}
+    cand = defaultdict(int)                                  # face -> 触及缺陷数(开边 + 非流形 + 退化)
+    oe_touch = set()                                         # 触及开边的 face: cap 内永远优先(补洞不被非流形稀释)
+    for edges, is_oe in ((oe, True), (nm if REMASK_NM else set(), False)):   # 覆盖1 vs >2 判据互斥
+        for (c1, c2) in edges:
+            g1 = tuple(np.array(c1) // 2); g2 = tuple(np.array(c2) // 2)
+            fs = (p2f.get(g1, set()) | p2f.get(g2, set())) if REMASK_UNION else (p2f.get(g1, set()) & p2f.get(g2, set()))
+            for f in fs:
+                cand[f] += 1
+                if is_oe: oe_touch.add(f)
+    if use_degen:                                           # 退化面 raw face 直接进 remask
+        for f in degen:
+            cand[f] += 1
+    faces = sorted(cand, key=lambda x: (-(x in oe_touch), -cand[x]))[:cap]
+    rem = torch.zeros(total, dtype=torch.bool)
+    for f in faces:
+        rem[f * 9 + 3: f * 9 + 9] = True                    # 只重掩连接位置(3-8),保 F(0-2)+coord
+    return rem, len(oe) + (len(nm) if REMASK_NM else 0), faces
+
+
+@torch.no_grad()
+def requery(model, coord_list, subcodes, decided, depth, pc, K, cfg_scale, cfg_schedule, cfg_start, pc_target=None, consensus=False):
+    EB_ACTIVE[0] = False   # requery=补洞期: 关 empty_bias, 否则补洞也被推向空(移植自 infer_remask_geom_altpc)
+    try:
+        return _requery_impl(model, coord_list, subcodes, decided, depth, pc, K, cfg_scale, cfg_schedule, cfg_start, pc_target, consensus)
+    finally:
+        EB_ACTIVE[0] = True
+
+
+def _requery_impl(model, coord_list, subcodes, decided, depth, pc, K, cfg_scale, cfg_schedule, cfg_start, pc_target=None, consensus=False):
+    """冻结已 commit, 对被 remask 的位置按真实掩码比例设噪声级 t, K 步 in-distribution 重采。
+       consensus=True(geom_fmm): 重采时也做跨面共识——F+intra 概率平均(按共享父/顶点)、inter 8-subcode
+       众数投票(按共享父对/边), 组原子揭示; 末步对全体做一次组投影(F 字节众数/intra,inter 向量众数),
+       保证同顶点/同边一致(含之前冻结的 token)。等价于把 log_fmm 的共识折进 geom 的补洞 requery。"""
+    device = subcodes.device
+    lengths = torch.tensor([c.numel() for c in coord_list], device=device)
+    types = model.types_from_lengths(lengths)
+    total = int(lengths.sum().item())
+    subcodes = subcodes.clone(); decided = decided.clone()
+    subcodes[~decided] = MASK_ID
+    n0 = int((~decided).sum().item())
+    want_cfg = (float(cfg_scale) != 1.0 or float(cfg_start) != 1.0) and model.use_pc and (pc is not None)
+    pred = subcodes.clone()
+
+    if consensus:   # ---- 跨面共识分组 + 工具(F/intra 按共享父, inter 按共享父对) ----
+        C = 4096; nface = total // REFINE_POS_PER_FACE
+        cr = torch.cat([c.to(device) for c in coord_list]).view(nface, 3, 3)
+        def _pcode(p3): return p3[:, 0] * C * C + p3[:, 1] * C + p3[:, 2]
+        f_pos = (types == TYPE_F).nonzero(as_tuple=True)[0]
+        codef = _pcode(cr[f_pos // REFINE_POS_PER_FACE, f_pos % REFINE_POS_PER_FACE])
+        _uf, f_group = torch.unique(codef, return_inverse=True); n_fgroup = int(_uf.numel())
+        i_pos = (types == TYPE_INTRA).nonzero(as_tuple=True)[0]
+        codei = _pcode(cr[i_pos // REFINE_POS_PER_FACE, (i_pos % REFINE_POS_PER_FACE) - 3])
+        _ui, i_group = torch.unique(codei, return_inverse=True); n_igroup = int(_ui.numel())
+        allpos = torch.cat([f_pos, i_pos])                              # 统一顶点分组(揭示原子)
+        _uv, vid_all = torch.unique(torch.cat([codef, codei]), return_inverse=True); n_vgroup = int(_uv.numel())
+        is_fi = torch.zeros(total, dtype=torch.bool, device=device); is_fi[allpos] = True
+        tok2vid = torch.full((total,), -1, dtype=torch.long, device=device); tok2vid[allpos] = vid_all
+        e_pos = (types == TYPE_INTER).nonzero(as_tuple=True)[0]         # inter 按共享父对(边)
+        e_face = e_pos // REFINE_POS_PER_FACE; e_local = (e_pos % REFINE_POS_PER_FACE) - 6
+        PA = torch.tensor([0, 0, 1], device=device); PB = torch.tensor([1, 2, 2], device=device)
+        ca = _pcode(cr[e_face, PA[e_local]]); cb = _pcode(cr[e_face, PB[e_local]])
+        pair = torch.stack([torch.minimum(ca, cb), torch.maximum(ca, cb)], dim=1)
+        _ue, e_group = torch.unique(pair, dim=0, return_inverse=True); n_egroup = int(_ue.shape[0])
+        is_inter = torch.zeros(total, dtype=torch.bool, device=device); is_inter[e_pos] = True
+        tok2eid = torch.full((total,), -1, dtype=torch.long, device=device); tok2eid[e_pos] = e_group
+        def _gmode(group, vid, G):
+            U = int(vid.max().item()) + 1 if vid.numel() else 1
+            key, cc = torch.unique(group.long() * U + vid.long(), return_counts=True)
+            g, v = key // U, key % U
+            order = torch.argsort(g * (int(cc.max().item()) + 1) - cc); 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(G, dtype=torch.long, device=device); mode[g[first]] = v[first]
+            return mode[group]
+        def vote_vec(vecs, group, ngroup):
+            uv, vid = torch.unique(vecs, dim=0, return_inverse=True); mvid = _gmode(group, vid, ngroup)
+            return vid != mvid, uv[mvid]
+        def group_avg_argmax(prob, group, ngroup):
+            agg = torch.zeros(ngroup, prob.shape[-1], device=device, dtype=prob.dtype); agg.index_add_(0, group, prob)
+            agg = agg / torch.bincount(group, minlength=ngroup).clamp(min=1).unsqueeze(1)
+            am_g = agg.argmax(-1); return am_g[group], agg.gather(-1, am_g.unsqueeze(-1)).squeeze(-1)[group]
+        def group_rollout(head, hh, hh_u, group, ngroup, cfg_t):
+            feats = [hh]; feats_u = [hh_u] if hh_u is not None else None; preds, confs = [], []
+            for r in range(head.num_rows):
+                logit = head.mlps[r](torch.cat(feats, dim=-1))
+                if feats_u is not None:
+                    lu = head.mlps[r](torch.cat(feats_u, dim=-1)); logit = lu + cfg_t * (logit - lu)
+                am, pm = group_avg_argmax(logit.float().softmax(-1), group, ngroup); preds.append(am); confs.append(pm)
+                if r < head.num_rows - 1:
+                    e = head.row_emb[r](am).to(hh.dtype); feats.append(e)
+                    if feats_u is not None: feats_u.append(e)
+            return torch.stack(preds, 1), torch.stack(confs, 1)
+
+    _rq_cc = {}                                              # 条件跨步缓存(点云在 K 步重采内不变, 严格等价)
+    for step in range(K):
+        mf = float((~decided).sum().item()) / max(1, total)
+        tval = (2.0 / math.pi) * math.acos(min(1.0, max(0.0, mf)))    # 真实掩码比例 -> 噪声时间
+        cfg_t = float(cfg_scale) if cfg_schedule != "linear" else float(cfg_start) + (float(cfg_scale) - float(cfg_start)) * step / max(K - 1, 1)
+        use_cfg = want_cfg and (cfg_t != 1.0)
+        t = torch.full((len(coord_list),), tval, device=device, dtype=torch.float)
+        sub_list = list(torch.split(subcodes, lengths.tolist(), dim=0))
+        h, _ = model.forward_varlen(coord_list, sub_list, t, depth, pc, pc_target, cond_cache=_rq_cc)
+        h_u = model.forward_varlen(coord_list, sub_list, t, depth, None, None, cond_cache=_rq_cc)[0] if use_cfg else None
+        h = h.float()                                  # 头部权重 fp32; BF16_RESIDUAL=1 时需 cast(fp32 下 no-op)
+        if h_u is not None: h_u = h_u.float()
+        pred = subcodes.clone(); conf = torch.zeros(total, device=device)
+        m = types == TYPE_F
+        if m.any():
+            lg = model.f_head(h[m])
+            if use_cfg: lu = model.f_head(h_u[m]); lg = lu + cfg_t * (lg - lu)
+            if consensus:
+                am, pm = group_avg_argmax(lg.float().softmax(-1), f_group, n_fgroup)
+            else:
+                prob = lg.softmax(-1); pm, am = prob.max(-1)
+            pred[m, :] = am.unsqueeze(1).expand(-1, NUM_OCTANTS); conf[m] = pm.to(conf.dtype)
+        mI = types == TYPE_INTRA
+        if mI.any():
+            if consensus:
+                p, pmx = group_rollout(model.intra_head, h[mI], h_u[mI] if use_cfg else None, i_group, n_igroup, cfg_t)
+            else:
+                p, pmx = model.intra_head.rollout(h[mI], h_u[mI] if use_cfg else None, cfg_t, 0.0)
+            pred[i_pos if consensus else mI] = p; conf[i_pos if consensus else mI] = pmx.mean(1).to(conf.dtype)
+        mE = types == TYPE_INTER
+        if mE.any():
+            p, pmx = model.inter_head.rollout(h[mE], h_u[mE] if use_cfg else None, cfg_t, 0.0)
+            pred[mE] = p; conf[mE] = pmx.mean(1).to(conf.dtype)
+            if consensus and e_pos.numel():   # inter 8-subcode 众数投票 + 组均置信度
+                _, vE = vote_vec(pred[e_pos], e_group, n_egroup); pred[e_pos] = vE
+                cagg = torch.zeros(n_egroup, device=device); cagg.index_add_(0, e_group, conf[e_pos])
+                cagg = cagg / torch.bincount(e_group, minlength=n_egroup).clamp(min=1); conf[e_pos] = cagg[e_group].to(conf.dtype)
+        n_keep = int(round(n0 * (1 - (step + 1) / K)))
+        n_dec = max(0, (total - int(decided.sum().item())) - n_keep)
+        if n_dec > 0:
+            und = (~decided).nonzero(as_tuple=True)[0]
+            order = conf[und].argsort(descending=True); ch = und[order[:n_dec]]
+            if consensus:   # 组原子揭示: 选中某顶点/边 -> 该顶点全部 F+intra / 该边全部 inter 一起
+                take = torch.zeros(total, dtype=torch.bool, device=device); take[ch] = True
+                ch_fi = ch[is_fi[ch]]
+                if ch_fi.numel():
+                    sel_v = torch.zeros(n_vgroup, dtype=torch.bool, device=device); sel_v[tok2vid[ch_fi]] = True
+                    take[allpos[sel_v[vid_all]]] = True
+                ch_e = ch[is_inter[ch]]
+                if ch_e.numel():
+                    sel_e = torch.zeros(n_egroup, dtype=torch.bool, device=device); sel_e[tok2eid[ch_e]] = True
+                    take[e_pos[sel_e[e_group]]] = True
+                take &= ~decided; subcodes[take] = pred[take]; decided[take] = True
+            else:
+                subcodes[ch] = pred[ch]; decided[ch] = True
+        subcodes[~decided] = MASK_ID
+    subcodes[~decided] = pred[~decided]
+    if consensus:   # 末步全体组投影: 同顶点/同边强制一致(含之前冻结 token)
+        if f_pos.numel():
+            cnt = torch.bincount(f_group * 256 + subcodes[f_pos, 0].clamp(0, 255), minlength=n_fgroup * 256).view(n_fgroup, 256)
+            subcodes[f_pos, :] = cnt.argmax(1)[f_group].unsqueeze(1).expand(-1, NUM_OCTANTS)
+        if i_pos.numel(): subcodes[i_pos] = vote_vec(subcodes[i_pos], i_group, n_igroup)[1]
+        if e_pos.numel(): subcodes[e_pos] = vote_vec(subcodes[e_pos], e_group, n_egroup)[1]
+    return list(torch.split(subcodes, lengths.tolist(), dim=0))
+
+
+def main():
+    device = "cuda"; cfg = load_config()
+    codec = FastMeshCodec(args=cfg.tokenizer); codec_cfg = codec.get_model_config(); pad_id = codec_cfg["pad_token_id"]
+    coord_offset = int(codec_cfg["coord_offset"])            # geom: coord 空间复现解码用
+    pc_num = int(cfg.data_processing.pc_num); total_pc_num = int(cfg.data_processing.total_pc_num)
+    cfg_scale = float(os.environ.get("CFG_SCALE", cfg.evaluation.cfg_scale))
+    cfg_schedule = cfg.evaluation.cfg_schedule; cfg_start = float(cfg.evaluation.cfg_start)
+    # 模型变体: config model.variant="fpsq" -> FPS-query 1024 latent + 稀疏注入(与 trainer/infer_remask 同开关)
+    _MC, _vk = BaseModel, {}      # 发布版固定 altpc 系(fix3 -> focal -> altpc -> tssr)
+    _margs = dict(_vk,
+                  coord_vocab_size=codec_cfg["input_num_categories"], max_seq_len=int(cfg.model.max_seq_len),
+                  d_model=cfg.model.feature_dim, nhead=cfg.model.num_heads, depth=cfg.model.hourglass_depth,
+                  dropout=0.0, encoder_freeze=True, use_pc=True, time_embed_scale=cfg.training.diffusion_train_steps,
+                  nonempty_row_weight=1.0,
+                  rope_mode=getattr(cfg.model, "rope_mode", "flat"),   # geo/geo_vcat 线 config 带(零新增参数, 不传会静默装成 flat)
+                  rope_theta_vcat=float(getattr(cfg.model, "rope_theta_vcat", 300.0)),
+                  ffn_hidden=(int(cfg.model.ffn_hidden) if getattr(cfg.model, "ffn_hidden", None) else None),
+                  pc_num_latents=(int(getattr(cfg.model, "pc_num_latents", 0)) or None))
+
+    def _load_model(ckpt_path):
+        m = _MC(**_margs).to(device)
+        ck = torch.load(ckpt_path, map_location=device)
+        m.load_state_dict(ck["model_state_dict"], strict=True)
+        sh = (ck.get("ema_state_dict") or {}).get("shadow")
+        if sh:
+            for name, param in m.named_parameters():
+                if name in sh:
+                    param.data.copy_(sh[name].to(device=param.device, dtype=param.dtype))
+        m.eval()
+        if EMPTY_BIAS or REQUERY_EMPTY_BIAS:   # 连接 rollout 空字节类加 bias; 生成期用 EMPTY_BIAS, requery(缺陷重��)期用 REQUERY_EMPTY_BIAS
+            import types as _types
+            def _biased_rollout(self, h, h_u=None, cfg_scale=1.0, temperature=0.0):
+                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)
+                    _b = EMPTY_BIAS if EB_ACTIVE[0] else REQUERY_EMPTY_BIAS
+                    if _b:
+                        logit = logit.clone(); logit[:, 0] = logit[:, 0] + _b
+                    if temperature and temperature > 0:
+                        prob = (logit / temperature).softmax(-1); am = torch.multinomial(prob.float(), 1).squeeze(-1); 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)
+            m.intra_head.rollout = _types.MethodType(_biased_rollout, m.intra_head)
+            m.inter_head.rollout = _types.MethodType(_biased_rollout, m.inter_head)
+            print(f"[EMPTY_BIAS={EMPTY_BIAS}] 连接空字节类已加 bias -> 稀疏化(requery 期自动关)")
+        raw = m.forward_varlen
+        def _autocast_forward(*args, **kwargs):
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                out = raw(*args, **kwargs)
+            # 关键: geo 线的 f_head / AR 头是 fp32(rope fp32 精度), 且 model.sample() 在 autocast 外调用它们。
+            # 把 backbone 输出 h 上采回 fp32, 才能与下游 fp32 头匹配(否则 bf16×fp32 -> mat dtype 崩)。
+            if isinstance(out, tuple):
+                return tuple(o.float() if torch.is_tensor(o) else o for o in out)
+            return out.float() if torch.is_tensor(out) else out
+        m.forward_varlen = _autocast_forward
+        return m
+
+    model = _load_model(CKPT)
+    # ---- 按 in_depth 换模型(专家组装): STAGE_CKPTS="2:a.pth,3:b.pth,4:c.pth", 未列出的深度用主 CKPT ----
+    stage_models = {}
+    for _kv in (x for x in os.environ.get("STAGE_CKPTS", "").split(",") if x.strip()):
+        _d, _p = _kv.split(":", 1)
+        stage_models[int(_d)] = _load_model(_p)
+        print(f"[STAGE_CKPTS] in_depth={_d} <- {os.path.basename(_p)}")
+
+    # ---- acc_uni 高分辨率细分器(SUPERRES): 独立 config/codec/model, 载入一次 ----
+    run_superres = None
+    if SUPERRES:
+        _acc_spec = _ilu.spec_from_file_location("acc_un", os.path.join(os.path.dirname(os.path.abspath(__file__)), "subdivider_runtime.py"))
+        ACC_UN = _ilu.module_from_spec(_acc_spec); _acc_spec.loader.exec_module(ACC_UN)
+        _saved_argv = sys.argv; sys.argv = ["x", "--config", SUPERRES_CONFIG]
+        sr_cfg = load_config(); sys.argv = _saved_argv                          # 单独加载 acc_uni config
+        sr_codec = OctreeMeshCodec(args=sr_cfg.tokenizer); sr_cc = sr_codec.get_model_config()
+        sr_model = DiscreteDiffusionTransformer(
+            vocab_A=sr_cc["input_num_categories"], vocab_B=sr_cc["output_num_categories"],
+            max_seq_len=sr_cfg.data_processing.max_len, d_model=sr_cfg.model.feature_dim,
+            encoder_freeze=sr_cfg.model.encoder_freeze, nhead=sr_cfg.model.num_heads,
+            depth=sr_cfg.model.hourglass_depth, dropout=sr_cfg.model.dropout,
+            time_embed_scale=sr_cfg.training.diffusion_train_steps).to(device)
+        sr_ckpt = torch.load(SUPERRES_CKPT, map_location=device)
+        sr_model.load_state_dict(sr_ckpt["model_state_dict"], strict=False)
+        _sr_shadow = (sr_ckpt.get("ema_state_dict") or {}).get("shadow") if isinstance(sr_ckpt.get("ema_state_dict"), dict) else None
+        if _sr_shadow:
+            for name, param in sr_model.named_parameters():
+                if name in _sr_shadow: param.data.copy_(_sr_shadow[name].to(device=param.device, dtype=param.dtype))
+        sr_model.eval()
+        sr_token_shift = sr_cc["input_num_categories"] - 1
+        # 细分阶梯起点 = 级联实际终点(真实完成的格): 正常跑完 STAGE_RUN 时是 STAGE[-1];
+        # 但 GT_START_RES/EXT_START_RES 起点高于整条 STAGE 时 STAGE_RUN=[](级联被跳过,
+        # current 就是外部/GT 起点本身) -> 阶梯必须从 _START_RES 起, 不能想当然用 STAGE[-1]
+        # (否则会把已经是 res128 的 mesh 又拿去按 res32 重新离散化, 白丢 3 级细节)。
+        _sr_start = STAGE_RUN[-1] if STAGE_RUN else (_START_RES or STAGE[-1])
+        sr_ladder = [_sr_start]                                                 # 细分阶梯: 倍增到 SUPERRES_TO
+        while sr_ladder[-1] < SUPERRES_TO:
+            sr_ladder.append(sr_ladder[-1] * 2)
+        print(f"[SUPERRES] acc_uni 细分器 {os.path.basename(SUPERRES_CKPT)} | 阶梯 {sr_ladder} | steps={SUPERRES_STEPS} cfg={SUPERRES_CFG} fix_holes={int(SUPERRES_FIX_HOLES)}")
+
+        def run_superres(cur, gt_mesh, stem, case_dir, hill_holes=True):
+            """用 acc_uni 细分器把 cur 从 STAGE[-1] 继续细分到 SUPERRES_TO(每级 save_stage_prediction 自存)。"""
+            sr_depth = ACC_UN._depth_from_resolution   # acc_uni 真实格=2^D -> log2 正确(实测: decode@D 格步=2/2^D); 勿用 fix3 的 log2-1
+            if SR_FILL_IN and cur is not None and len(getattr(cur, "faces", [])):   # 入口补洞: 对齐旧 fillall(粗输入先补再细分)
+                _v = torch.as_tensor(np.asarray(cur.vertices), dtype=torch.float32, device=device)
+                _f = torch.as_tensor(np.asarray(cur.faces), dtype=torch.long, device=device)
+                _v, _f = _fill_holes(_v, _f, only_fill_holes=True)
+                cur = trimesh.Trimesh(vertices=(_v.detach().cpu().numpy() if torch.is_tensor(_v) else np.asarray(_v)),
+                                      faces=(_f.detach().cpu().numpy() if torch.is_tensor(_f) else np.asarray(_f)), process=False)
+            sr_bar = tqdm(range(len(sr_ladder) - 1), desc=f"{stem[:12]} 细分", leave=False)
+            for i in sr_bar:
+                in_res, out_res = sr_ladder[i], sr_ladder[i + 1]
+                sr_bar.set_postfix_str(f"res{in_res}->{out_res}")
+                in_d, out_d = sr_depth(in_res), sr_depth(out_res)
+                enc = mul_encode_mesh(raw_mesh=cur.copy(), gt_mesh=gt_mesh.copy(), codec=sr_codec, tokenizer_target_depth=in_d,
+                                      pc_num=pc_num, total_pc_num=total_pc_num, use_gt_pc=True, clean=False,
+                                      normalize_input=CASCADE_RENORM)
+                raw_tok = enc["raw_token"].unsqueeze(0).to(device); pc = enc["pc"].unsqueeze(0).to(device)
+                with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+                    refine = sample_refine_tokens_with_topology(sr_model, sr_model.mask_token_id, raw_tok, pc,
+                                                                use_gumbel=False, use_soft_consensus=False,
+                                                                consensus_start_step_ratio=0, cfg_scale=SUPERRES_CFG, T_infer=SUPERRES_STEPS)
+                
+                refine = ACC_UN.restore_non_pad_token(refine, shift=sr_token_shift, pad_id=sr_cc["pad_token_id"])
+                m = ACC_UN.save_stage_prediction(codec=sr_codec, raw_token=raw_tok, refine_token=refine,
+                                                 out_dir=case_dir, name=stem, target_depth=out_d,
+                                                 suffix=f"_sr_{in_res}to{out_res}", fix_holes=SUPERRES_FIX_HOLES,
+                                                 raw_dis_mesh=None, gt_dis_mesh=None)
+                if m is None:
+                    break
+                cur = m
+            return cur
+
+    def save_yup(mesh, name, out_dir):
+        if SAVE and mesh is not None and len(getattr(mesh, "faces", [])):
+            m = mesh.copy(); m.vertices = m.vertices[:, [1, 2, 0]]; m.export(os.path.join(out_dir, name))
+
+    if DATASET == "val":
+        paths = [line.split()[0] for line in open(VAL) if line.strip()]
+        car_path = next((p for p in paths if CAR in p), None)
+        cases = (([car_path] if car_path else []) + [p for p in paths if (not car_path or CAR not in p)])[:NCASE]
+    elif DATASET == "toys4k":
+        cases = sorted(_glob.glob(os.path.join(TOYS4K, "*.glb")))[:NCASE]
+    elif os.environ.get("DATASET_DIR"):   # 任意外部目录: DATASET 只作命名标签, 实际目录由 DATASET_DIR 给
+        _dd = os.environ["DATASET_DIR"]
+        cases = sorted(set(sum((_glob.glob(os.path.join(_dd, f"*.{_e}")) for _e in ("glb", "obj", "ply")), [])))[:NCASE]
+    else:
+        cases = sorted(_glob.glob(os.path.join(TRELLIS, "*.glb")))[:NCASE]
+    if CASE_STEM:   # 逗号分隔多个子串, 任一命中即保留(如 "1(10).,1(101)." 用点号防前缀误配)
+        _stems = [s for s in CASE_STEM.split(",") if s]
+        cases = [p for p in cases if any(s in p for s in _stems)]
+    if EXT_START_DIR:   # 只跑有外部起点文件的 case
+        cases = [p for p in cases if os.path.exists(os.path.join(EXT_START_DIR, os.path.splitext(os.path.basename(p))[0] + EXT_START_SUFFIX))]
+        print(f"[EXT_START] {EXT_START_DIR} 匹配到 {len(cases)} case (后缀 {EXT_START_SUFFIX}, res{EXT_START_RES})")
+
+    print(f"多级生成(真实res口径) STAGE={STAGE_RUN}{f' + SUPERRES->{SUPERRES_TO}' if SUPERRES else ''}{f' | GT起点@res{GT_START_RES}' if GT_START_RES else ''}{f' | 外部起点@res{EXT_START_RES}' if EXT_START_DIR else ''} | ep{EPOCH} STEPS={STEPS} | "
+          f"METHODS={METHODS} 多级PC={int(MULTI_PC)} ALIGN fm={int(ALIGN_FM)} inter={int(ALIGN_INTER)} FILL={int(FILL)} | {len(cases)} case -> {OUT}")
+    print(f"{'case':12s} {'method':9s} | 开边 面 非流形 | CD/HD/F@GT(CD/HD为v2协议[-0.5,0.5]帧) |NC| | 用时")
+    print("-" * 82)
+    results = {m: [] for m in METHODS}
+    for path in cases:
+        stem = os.path.splitext(os.path.basename(path))[0]
+        gt = trimesh.load(path, force="mesh"); gt.vertices = gt.vertices[:, [2, 0, 1]]; normalize(gt)
+        case_dir = os.path.join(OUT, stem); os.makedirs(case_dir, exist_ok=True)
+        save_yup(gt, "gt_original.obj", case_dir)
+        gt.vertices, gt.faces = decimate_mesh(gt.vertices, gt.faces, simplify=True, simplify_ratio=DECIMATE_RATIO, verbose=False)
+        save_yup(gt, "gt_original_dec.obj", case_dir)
+        gt_eval = gt.copy()                     # 指标参照: 固定为未补洞的塌缩GT, DEC_FILL 只改输入侧
+        if DEC_FILL and len(getattr(gt, "faces", [])):
+            _v = torch.as_tensor(np.asarray(gt.vertices), dtype=torch.float32, device=device)
+            _f = torch.as_tensor(np.asarray(gt.faces), dtype=torch.long, device=device)
+            _v, _f = _fill_holes(_v, _f, only_fill_holes=True)
+            gt = trimesh.Trimesh(vertices=(_v.detach().cpu().numpy() if torch.is_tensor(_v) else np.asarray(_v)),
+                                 faces=(_f.detach().cpu().numpy() if torch.is_tensor(_f) else np.asarray(_f)), process=False)
+            save_yup(gt, "gt_original_dec_fill.obj", case_dir)
+        # ---- 粗起点(所有方法共用) ----
+        if EXT_START_DIR:  # 外部预测起点(如 BPT res32 预测): 保存帧 -> 内部帧(与 gt 载入同款 [2,0,1] 交换)
+            coarse_start = trimesh.load(os.path.join(EXT_START_DIR, stem + EXT_START_SUFFIX), force="mesh", process=False)
+            coarse_start.vertices = np.asarray(coarse_start.vertices)[:, [2, 0, 1]]
+            assert len(getattr(coarse_start, "faces", [])), f"外部起点为空 mesh: {stem}"
+        elif GT_START_RES:   # GT 天花板起点: GT 编码->解码(含连接)@该分辨率, 测级联暴露偏差上界
+            _d0 = depth_from_resolution(GT_START_RES)
+            _enc0 = mul_encode_mesh(raw_mesh=gt.copy(), gt_mesh=gt.copy(), codec=codec, tokenizer_target_depth=_d0,
+                                    pc_num=pc_num, total_pc_num=total_pc_num, use_gt_pc=True, clean=False)
+            coarse_start = codec.decode_tokens(_enc0["token"].tolist(), target_depth=_d0)
+            assert coarse_start is not None and len(getattr(coarse_start, "faces", [])), f"GT 天花板@res{GT_START_RES} 解码为空: {stem}"
+        else:              # GT 八叉树解码到最粗级(只有占用的 raw 起点, 非完整 mesh -> 文件名带 raw_ 前缀)
+            raw_token, _, init_quant = codec.encode_mesh(gt.copy(), target_depth=depth_from_resolution(STAGE_RUN[0]))
+            coarse_start = codec.decode_raw_tokens(raw_token, target_depth=depth_from_resolution(STAGE_RUN[0]))
+            if coarse_start is None or len(getattr(coarse_start, "faces", [])) == 0:
+                coarse_start = init_quant if init_quant is not None else gt.copy()
+        save_yup(coarse_start, (f"res{_START_RES}.obj" if _START_RES else f"raw_res{STAGE_RUN[0]}.obj"), case_dir)
+        if SAVE_GT_LEVELS and SAVE:   # GT 逐分辨率天花板(tokenizer 量化口径), 供逐级对照
+            for _r in STAGE_RUN:
+                try:
+                    _d = depth_from_resolution(_r)
+                    _enc = mul_encode_mesh(raw_mesh=gt.copy(), gt_mesh=gt.copy(), codec=codec, tokenizer_target_depth=_d,
+                                           pc_num=pc_num, total_pc_num=total_pc_num, use_gt_pc=True, clean=False)
+                    _gm = codec.decode_tokens(_enc["token"].tolist(), target_depth=_d)
+                    save_yup(_gm, f"gt_res{_r}.obj", case_dir)
+                except Exception as _e:
+                    print(f"  [SAVE_GT_LEVELS] {stem} res{_r} 失败: {_e}")
+        for method in METHODS:
+            np.random.seed(777 + 7919 * RUN_SEED)   # 每方法同种子 -> 点云/退化采样一致, 干净对比; RUN_SEED 偏移用于多次独立采样
+            t_start = time.time()
+            current = coarse_start.copy()
+            base = method.split(":")[0]      # geom 可用 "geom_fmm:2x5" 覆盖 NITER=2/KREQ=5
+            ni_g, kq_g = NITER_GEOM, KREQ_GEOM
+            if ":" in method:
+                _s = method.split(":")[1].lower().split("x"); ni_g, kq_g = int(_s[0]), int(_s[1])
+            msave = method.replace(":", "_")  # 文件名去冒号
+            # ---- 级联: 每级把上一级产物重编码成占用, 在模型深度 d=log2(out_res)-1 补全成格 out_res 的完整 mesh ----
+            stage_bar = tqdm(range(len(STAGE_RUN)), desc=f"{stem[:12]} {method} 级联", leave=False)
+            for stage_i in stage_bar:
+                out_res = STAGE_RUN[stage_i]
+                d = depth_from_resolution(out_res)
+                stage_bar.set_postfix_str(f"->res{out_res}")
+                encoded = mul_encode_mesh(raw_mesh=current.copy(), gt_mesh=gt.copy(), codec=codec, tokenizer_target_depth=d,
+                                          pc_num=pc_num, total_pc_num=total_pc_num, use_gt_pc=True, clean=False,
+                                          normalize_input=CASCADE_RENORM)
+                coord_list, _ = split_fix3_token_list([encoded["token"]], pad_id, device, codec_cfg)
+                pc_full = encoded["pc"].unsqueeze(0).to(device)
+                if MULTI_PC:   # altpc pc_target: GT 天花板@d 解码 -> 退化面感知采点
+                    # resolution=out_res=2^(d+1) 严格对齐训练(mesh_dataset_altpc_decoded: decode@sd + 2^(sd+1))
+                    gt_encoded = mul_encode_mesh(raw_mesh=gt.copy(), gt_mesh=gt.copy(), codec=codec, tokenizer_target_depth=d,
+                                                 pc_num=pc_num, total_pc_num=total_pc_num, use_gt_pc=True, clean=False)
+                    gt_ceiling_mesh = codec.decode_tokens(gt_encoded["token"].tolist(), target_depth=d)
+                    pc_target_np = (sample_pc_with_degen(gt_ceiling_mesh, pc_num, total_pc_num, aug=False, overfit=False, resolution=out_res)
+                                    if PCT_DEG else sample_pc(gt_ceiling_mesh, pc_num=pc_num, total_pc_num=total_pc_num, with_normal=True, aug=False))
+                    pc_target = torch.from_numpy(pc_target_np).unsqueeze(0).to(device)
+                    if SAVE and SAVE_PC_LEVELS:   # 实际注入的云B(帧与 obj 一致); 多方法时同名覆盖(同种子, 云几乎一致)
+                        trimesh.PointCloud(pc_target_np[:, :3][:, [1, 2, 0]]).export(os.path.join(case_dir, f"pc_target_res{out_res}.ply"))
+                else:          # 单点云(focal 单PC ckpt): pc_target=None -> altpc 退化为单 PC
+                    pc_target = None
+                if SINGLE_TARGET_PC and pc_target is not None:   # sfocal: 云B进 pc 位作唯一云(每层恒注入), 云A弃用
+                    pc_full = pc_target
+                    pc_target = None
+                if SAVE and SAVE_PC_LEVELS and stage_i == 0 and not os.path.exists(os.path.join(case_dir, "pc_full.ply")):
+                    _pa = encoded["pc"].cpu().numpy()             # 云A: 全分辨率 GT 云(每级重采, 分布相同, 存首级一份)
+                    trimesh.PointCloud(_pa[:, :3][:, [1, 2, 0]]).export(os.path.join(case_dir, "pc_full.ply"))
+                torch.manual_seed(20000 + stage_i + 7919 * RUN_SEED)
+                depth_t = torch.tensor([d], device=device)
+                model_d = stage_models.get(d, model)   # 专家组装: 该深度有专属 ckpt 则用之
+                if base in ("log", "log_fmm"):        # 概率平均共识(log_fmm 再加 inter 众数投票)
+                    gen_subcodes = model_d.sample_log(coord_list, depth_t, pc=pc_full, pc_target=pc_target, T_infer=STEPS,
+                                                    cfg_scale=cfg_scale, temperature=0.0, cfg_schedule=cfg_schedule, cfg_start=cfg_start,
+                                                    inter_vote=(base == "log_fmm"))
+                elif base in ("geom", "geom_fmm", "geom_deg", "geom_fmm_deg"):    # plain 采样 + 终点几何 remask+requery
+                    # _fmm -> requery 内共识; _deg -> remask 额外纳入退化面(否则用全局 REMASK_DEGEN)
+                    use_cons = ("fmm" in base); use_deg = base.endswith("_deg") or REMASK_DEGEN
+                    ni_g, kq_g = GEOM_STAGE_PARAMS.get(out_res, (ni_g, kq_g))   # 按级覆盖 remask 预算
+                    _n_cand = BO_N if (BO_N > 1 and out_res in BO_STAGES) else 1
+                    _best = None                            # (score, gen_cur, rem_union)
+                    for _cand in range(_n_cand):
+                        _pt = pc_target
+                        if _cand > 0 and pc_target is not None:   # 候选多样性: 重采条件云B(实测±40方差来源)
+                            _ptn = (sample_pc_with_degen(gt_ceiling_mesh, pc_num, total_pc_num, aug=False, overfit=False, resolution=out_res)
+                                    if PCT_DEG else sample_pc(gt_ceiling_mesh, pc_num=pc_num, total_pc_num=total_pc_num, with_normal=True, aug=False))
+                            _pt = torch.from_numpy(_ptn).unsqueeze(0).to(device)
+                        torch.manual_seed(20000 + stage_i + 7919 * _cand)
+                        gen_list = model_d.sample(coord_list, depth_t, pc=pc_full, pc_target=_pt, T_infer=STEPS,
+                                                cfg_scale=cfg_scale, temperature=0.0, cfg_schedule=cfg_schedule, cfg_start=cfg_start,
+                                                align_fm=False, align_inter=False, align_per_step=False, align_debug=False)
+                        gen_cur = gen_list[0].to(device); total_tok = coord_list[0].numel()
+                        rem_union = set()                      # SAVE_REMASK_VIS: 各轮被重掩 token 的并集
+                        for _ in range(ni_g if (not GEOM_STAGES or out_res in GEOM_STAGES) else 0):
+                            rem, _, rfaces = build_remask(coord_list[0].cpu(), gen_cur.cpu(), coord_offset, total_tok, use_degen=use_deg)
+                            if int(rem.sum().item()) == 0:
+                                break
+                            rem_union.update(torch.nonzero(rem, as_tuple=False).reshape(-1).tolist() if rem.dtype == torch.bool else rem.reshape(-1).tolist())
+                            if REMASK_GROUPS > 1 and len(rfaces) > REMASK_GROUPS:   # 分批: BFS序(≈空间局部)分 G 批逐批重采
+                                for _ch in np.array_split(np.array(sorted(rfaces)), REMASK_GROUPS):
+                                    if len(_ch) == 0:
+                                        continue
+                                    _rm = torch.zeros(total_tok, dtype=torch.bool)
+                                    for _f in _ch:
+                                        _rm[int(_f) * 9 + 3: int(_f) * 9 + 9] = True
+                                    dec_b = torch.ones(total_tok, dtype=torch.bool, device=device); dec_b[_rm.to(device)] = False
+                                    gen_cur = requery(model_d, coord_list, gen_cur, dec_b, depth_t, pc_full, kq_g,
+                                                      cfg_scale, "constant", 1.0, pc_target=_pt, consensus=use_cons)[0]
+                            else:
+                                dec_b = torch.ones(total_tok, dtype=torch.bool, device=device); dec_b[rem.to(device)] = False
+                                gen_cur = requery(model_d, coord_list, gen_cur, dec_b, depth_t, pc_full, kq_g,
+                                                  cfg_scale, "constant", 1.0, pc_target=_pt, consensus=use_cons)[0]
+                        if _n_cand == 1:
+                            _best = (0.0, gen_cur, rem_union)
+                            break
+                        _cm = decode_subcodes_to_mesh(codec, coord_list[0], gen_cur, codec_cfg, d + 1)
+                        _oe, _nf, _nmp = mesh_open_edge_stats(_cm)
+                        _score = _oe + 10.0 * _nmp             # 结构惩罚: 开边 + 10*非流形%(bo4 同思路)
+                        if _best is None or _score < _best[0]:
+                            _best = (_score, gen_cur, rem_union)
+                    gen_cur, rem_union = _best[1], _best[2]
+                    gen_subcodes = [gen_cur]
+                else:                                    # plain / consensus(末步 align 众数表决)
+                    is_cons = (base == "consensus")
+                    gen_subcodes = model_d.sample(coord_list, depth_t, pc=pc_full, pc_target=pc_target, T_infer=STEPS,
+                                                cfg_scale=cfg_scale, temperature=0.0, cfg_schedule=cfg_schedule, cfg_start=cfg_start,
+                                                align_fm=(is_cons or ALIGN_FM), align_inter=(is_cons or ALIGN_INTER),
+                                                align_per_step=ALIGN_PER_STEP, align_debug=False)
+                stage_mesh = decode_subcodes_to_mesh(codec, coord_list[0], gen_subcodes[0], codec_cfg, d + 1)   # 产物=深度d完整mesh, 真实格 out_res(实测)
+                if SAVE_REMASK_VIS and base.startswith("geom") and SAVE:
+                    _pre = decode_subcodes_to_mesh(codec, coord_list[0], gen_list[0].to(device), codec_cfg, d + 1)
+                    if _pre is not None:
+                        save_yup(_pre, f"res{out_res}_{msave}_pre.obj", case_dir)   # 补洞(remask/requery)前的 plain 采样解码
+                    if _pre is not None and stage_mesh is not None:
+                        save_remask_vis(_pre, stage_mesh, os.path.join(case_dir, f"res{out_res}_{msave}"))
+                    if rem_union:   # remask token -> 输入侧 raw face block(每块 REFINE_POS_PER_FACE=9 token) 投影上色
+                        try:
+                            _rt, _, _ = codec.encode_mesh(current.copy(), target_depth=d)
+                            _raw = codec.decode_raw_tokens(_rt, target_depth=d)
+                            _nb = coord_list[0].numel() // REFINE_POS_PER_FACE
+                            if _raw is not None and len(_raw.faces) == _nb:
+                                _bm = np.zeros(_nb, bool)
+                                _bm[sorted({i // REFINE_POS_PER_FACE for i in rem_union})] = True
+                                _mr = _raw.copy(); _mr.vertices = _mr.vertices[:, [1, 2, 0]]
+                                _mr.visual.face_colors = np.where(_bm[:, None], [255, 140, 0, 255], [200, 200, 200, 255])
+                                _mr.export(os.path.join(case_dir, f"res{out_res}_{msave}_remask_rawface.ply"))
+                                print(f"  [REMASK_VIS] raw face 投影: 被重掩块 {int(_bm.sum())}/{_nb}")
+                            else:
+                                print(f"  [REMASK_VIS] raw face 数不匹配: raw {0 if _raw is None else len(_raw.faces)} vs 块 {_nb}, 跳过投影")
+                        except Exception as _e:
+                            print(f"  [REMASK_VIS] raw face 投影失败: {_e}")
+                if QUAD_FILL and stage_mesh is not None and len(getattr(stage_mesh, "faces", [])):
+                    stage_mesh = quad_fill_mesh(stage_mesh, out_res)
+                # 直接补洞: method=fill 在 FILL_STAGES 各级补洞喂下一级(纯几何, 不过模型); 或全局 FILL 开关(stage_i>=4)
+                _fill_hit = (out_res >= FILL_MIN_RES) if FILL_MIN_RES else (stage_i >= 4)
+                do_fill = (FILL and _fill_hit) or (base == "fill" and out_res in FILL_STAGES)
+                if do_fill and stage_mesh is not None and len(getattr(stage_mesh, "faces", [])):
+                    _v = torch.as_tensor(np.asarray(stage_mesh.vertices), dtype=torch.float32, device=device)
+                    _f = torch.as_tensor(np.asarray(stage_mesh.faces), dtype=torch.long, device=device)
+                    _v, _f = _fill_holes(_v, _f, only_fill_holes=True, keep_degenerate_faces=True)
+                    _vn = _v.detach().cpu().numpy() if torch.is_tensor(_v) else np.asarray(_v)
+                    _fn = _f.detach().cpu().numpy() if torch.is_tensor(_f) else np.asarray(_f)
+                    stage_mesh = trimesh.Trimesh(vertices=_vn, faces=_fn, process=False)
+                current = stage_mesh if stage_mesh is not None else current
+                save_yup(current, f"res{out_res}_{msave}.obj", case_dir)
+            if SUPERRES and run_superres is not None and current is not None and len(getattr(current, "faces", [])):
+                current = run_superres(current, gt, stem, case_dir)
+            save_yup(current, f"final_res{FINAL_RES}_{msave}.obj", case_dir)
+            elapsed = time.time() - t_start
+            open_edges, num_faces, nonmanifold = mesh_open_edge_stats(current)
+            n_degen = int((np.asarray(current.area_faces) <= 1e-12).sum()) if (current is not None and len(getattr(current, "faces", []))) else 0
+            cd, hd, fscore, nc = chamfer_hd_fscore(current, gt_eval)
+            print(f"{stem[:12]:12s} {method:12s} | {open_edges:5d} {num_faces:6d} {nonmanifold:4.1f} 退化{n_degen:4d} | {cd:5.1f}/{hd:5.0f}/{fscore:4.0f} |NC|{nc:.3f} | {elapsed:.1f}s")
+            results[method].append((open_edges, num_faces, nonmanifold, n_degen, cd, hd, fscore, nc, elapsed))
+    print("=" * 88)
+    for method in METHODS:
+        r = results[method]
+        if not r:
+            continue
+        mr = np.array(r, np.float64).mean(0)
+        print(f"{'均值':12s} {method:12s} | {mr[0]:5.0f} {mr[1]:6.0f} {mr[2]:4.1f} 退化{mr[3]:4.0f} | {mr[4]:5.1f}/{mr[5]:5.0f}/{mr[6]:4.0f} |NC|{mr[7]:.3f} | {mr[8]:.1f}s")
+
+
+if __name__ == "__main__":
+    main()
+
+# 运行(GPU7; 训练占 0-6):
+#   cd /root/junkai_lin/mul_res_diffusion && source /root/miniconda3/etc/profile.d/conda.sh && conda activate mul_res
+#   HG=pretrained_model/gobjaverse_fix3_octant_factorized_1_to_1024_ce_altpc_decoded_deg_w1_hg246
+#   CUDA_VISIBLE_DEVICES=7 STAGE=2,4,8,16,32 DATASET=trellis NCASE=8 STEPS=20 ALIGN_INTER=1 FILL=1 \
+#     BASE_DIR=$HG RAIL_CKPT=$HG/checkpoint_30_epoch_0.05907_loss.pth python infer_multilevel.py

mulres_infer/modules/subdivider_runtime.py

@@ -0,0 +1,575 @@
+import gc
+import glob
+import math
+import os
+import random
+import re
+from pathlib import Path
+
+import numpy as np
+import torch
+import trimesh
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from tqdm import tqdm
+
+from utils.config import load_config
+from dataset.mesh_dataset import mul_encode_mesh, discrete_undiscrete_and_clean
+from models.subdivider import DiscreteDiffusionTransformer, sample_refine_tokens_with_topology, sample_refine_tokens_with_topology_fm
+from modules.octree_codec_acc import MeshCodec
+from utils.postprocessing import _fill_holes
+from utils.data_utils import normalize
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    os.environ["PYTHONHASHSEED"] = str(seed)
+
+def _depth_from_resolution(resolution):
+    if resolution <= 0:
+        raise ValueError(f"resolution must be positive, got {resolution}")
+    depth = int(round(math.log2(resolution)))
+    if (1 << depth) != resolution:
+        raise ValueError(f"resolution must be power of 2, got {resolution}")
+    return depth
+
+
+def restore_non_pad_token(token: torch.Tensor, shift: int, pad_id: int = 0):
+    if shift == 0:
+        return token
+    restored = token.clone()
+    non_pad_mask = restored != pad_id
+    restored[non_pad_mask] = restored[non_pad_mask] + shift
+    return restored
+
+
+def read_dataset_paths(txt_path):
+    mesh_paths = []
+    with open(txt_path, "r") as f:
+        for line in f:
+            path = line.strip()
+            if path:
+                mesh_paths.append(path)
+    return mesh_paths
+
+
+def raw_stem_to_pair_key(raw_path):
+    stem = Path(raw_path).stem
+    # raw name example: 1(42)_y_f_4244 -> pair key: 1(42)
+    return re.sub(r"_y_f_\d+$", "", stem)
+
+
+def gt_stem_to_pair_key(gt_path):
+    return Path(gt_path).stem
+
+def get_face_count(file_path):
+    """
+    使用正则表达式从文件名中提取面片数量
+    支持格式: *_f{num}_v*.obj 或 *_facenum_{num}_*.glb
+    """
+    filename = Path(file_path).name
+    match = re.search(r'_f(\d+)_v|_facenum_(\d+)', filename)
+    if match:
+        # group(1) 对应 _f(\d+)_v 里的数字
+        # group(2) 对应 _facenum_(\d+) 里的数字
+        return int(match.group(1) or match.group(2))
+    return None
+
+# def build_raw_gt_pairs(raw_txt_path, gt_txt_path):
+#     raw_paths = read_dataset_paths(raw_txt_path)
+#     gt_paths = read_dataset_paths(gt_txt_path)
+
+#     # If two txt paths are the same source, pair each mesh with itself directly.
+#     if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve():
+#         return [(p, p) for p in raw_paths], []
+
+#     gt_by_key = {}
+#     for gt_path in gt_paths:
+#         gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path
+
+#     pairs = []
+#     missing_gt = []
+#     for raw_path in raw_paths:
+#         key = raw_stem_to_pair_key(raw_path)
+#         gt_path = gt_by_key.get(key)
+#         if gt_path is None:
+#             missing_gt.append(raw_path)
+#             continue
+#         pairs.append((raw_path, gt_path))
+
+#     return pairs, missing_gt
+def build_raw_gt_pairs(raw_txt_path, gt_txt_path, min_faces=None, max_faces=None):
+    raw_paths = read_dataset_paths(raw_txt_path)
+    gt_paths = read_dataset_paths(gt_txt_path)
+
+    # ===== 新增过滤逻辑 =====
+    filtered_raw_paths = []
+    for p in raw_paths:
+        if min_faces is None and max_faces is None:
+            filtered_raw_paths.append(p)
+            continue
+            
+        faces = get_face_count(p)
+        if faces is not None:
+            # 检查是否在指定范围内
+            if min_faces is not None and faces < min_faces:
+                continue
+            if max_faces is not None and faces > max_faces:
+                continue
+            filtered_raw_paths.append(p)
+        else:
+            # 如果文件名中没找到面片数特征，默认保留（你也可以改为 continue 丢弃它）
+            filtered_raw_paths.append(p)
+            
+    raw_paths = filtered_raw_paths
+    # ========================
+
+    # If two txt paths are the same source, pair each mesh with itself directly.
+    if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve():
+        return [(p, p) for p in raw_paths], []
+
+    gt_by_key = {}
+    for gt_path in gt_paths:
+        gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path
+
+    pairs = []
+    missing_gt = []
+    for raw_path in raw_paths:
+        key = raw_stem_to_pair_key(raw_path)
+        gt_path = gt_by_key.get(key)
+        if gt_path is None:
+            missing_gt.append(raw_path)
+            continue
+        pairs.append((raw_path, gt_path))
+
+    return pairs, missing_gt
+
+def save_gt_mesh(mesh, out_dir, name):
+    gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False)
+    gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]]
+    gt_mesh.export(os.path.join(out_dir, f"{name}_gt.obj"))
+
+def save_raw_gt_mesh(mesh, out_dir, name):
+    gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False)
+    gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]]
+    gt_mesh.export(os.path.join(out_dir, f"{name}_raw_gt.obj"))
+
+
+def save_stage_input_mesh(mesh, out_dir, name, stage_in_res, stage_out_res, stage_in_depth):
+    coarse_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False)
+    coarse_mesh.vertices = coarse_mesh.vertices[:, [1, 2, 0]]
+    suffix = f"_stage_{stage_in_res}to{stage_out_res}_input_coarse"
+    coarse_mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{stage_in_depth}.obj"))
+
+
+def save_stage_prediction(codec, raw_token, refine_token, out_dir, name, target_depth, suffix, fix_holes, raw_dis_mesh=None, gt_dis_mesh=None):
+    eos_mask = raw_token[0] == 0
+    if eos_mask.any():
+        raw_tokens = raw_token[0, : int(torch.where(eos_mask)[0][0])].detach().cpu().long().tolist()
+    else:
+        raw_tokens = raw_token[0].detach().cpu().long().tolist()
+
+    eos_mask_refine = refine_token[0] == 0
+    if eos_mask_refine.any():
+        refine_tokens = refine_token[0, : int(torch.where(eos_mask_refine)[0][0])].detach().cpu().long().tolist()
+    else:
+        refine_tokens = refine_token[0].detach().cpu().long().tolist()
+
+    shared_len = min(len(raw_tokens), len(refine_tokens))
+    block = (shared_len // 9) * 9
+    if block <= 0:
+        return None
+
+    raw_tokens = raw_tokens[:block]
+    refine_tokens = refine_tokens[:block]
+
+    interleaved = []
+    for i in range(0, block, 9):
+        interleaved.extend(raw_tokens[i : i + 9])
+        interleaved.extend(refine_tokens[i : i + 9])
+
+    mesh = codec.decode_tokens(interleaved, target_depth=target_depth-1)
+    if mesh is None:
+        return None
+    
+    trimesh.repair.fix_normals(mesh)
+    if fix_holes:
+        verts = torch.tensor(mesh.vertices, device="cuda", dtype=torch.float32)
+        faces = torch.tensor(mesh.faces, device="cuda", dtype=torch.int32)
+        new_verts, new_faces = _fill_holes(verts, faces, debug=False, verbose=False, only_fill_holes=True)
+        mesh = trimesh.Trimesh(vertices=new_verts.cpu().numpy(), faces=new_faces.cpu().numpy())
+    
+    return_mesh = mesh.copy()
+    mesh.vertices = mesh.vertices[:, [1, 2, 0]]
+    mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{target_depth}.obj"))
+
+    if raw_dis_mesh is not None:
+        raw_dis_mesh.vertices = raw_dis_mesh.vertices[:, [1, 2, 0]]
+        raw_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_raw_dis_depth{target_depth-1}.obj"))
+    if gt_dis_mesh is not None:
+        gt_dis_mesh.vertices = gt_dis_mesh.vertices[:, [1, 2, 0]]
+        gt_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_gt_dis_depth{target_depth}.obj"))
+    return return_mesh
+
+
+def process_one_mesh(
+    raw_mesh_path,
+    gt_mesh_path,
+    model,
+    mask_token_id,
+    codec,
+    device,
+    pad_id,
+    token_shift,
+    config,
+    output_folder,
+    stage_resolutions,
+    fix_holes=True,
+    use_gt_pc=True
+):
+    import time
+    t_start = time.time()
+    import time
+    t_total_start = time.time()
+    
+    name = Path(raw_mesh_path).stem
+    t_start = time.time()
+    mesh = trimesh.load(raw_mesh_path, force="mesh")
+    mesh.vertices = mesh.vertices[:, [2, 0, 1]]
+    gt_mesh = trimesh.load(gt_mesh_path, force="mesh")
+    gt_mesh.vertices = gt_mesh.vertices[:, [2, 0, 1]]
+    normalize(gt_mesh)
+    normalize(mesh)
+    save_raw_gt_mesh(mesh, output_folder, name)
+    save_gt_mesh(gt_mesh, output_folder, name)
+    t_load = time.time()
+    print(f"[{name}] Load & normalize time: {t_load - t_start:.4f}s")
+
+    stage_depths = [_depth_from_resolution(x) for x in stage_resolutions]
+
+    # First-stage input uses codec encode-decode result at the minimum resolution.
+    min_stage_depth = stage_depths[0]
+    t_start = time.time()
+    init_raw_token, _, init_quantized_mesh = codec.encode_mesh(mesh, target_depth=min_stage_depth)
+    init_mesh = codec.decode_raw_tokens(init_raw_token, target_depth=min_stage_depth)
+    t_init = time.time()
+    print(f"[{name}] Init encode & decode (depth={min_stage_depth}): {t_init - t_start:.4f}s")
+    
+    if init_mesh is not None:
+        current_mesh = init_mesh
+    elif init_quantized_mesh is not None:
+        current_mesh = init_quantized_mesh
+    else:
+        current_mesh = mesh.copy()
+        
+    for stage_idx in range(len(stage_depths) - 1):
+        stage_start_time = time.time()
+        stage_in_depth = stage_depths[stage_idx]
+        stage_out_depth = stage_depths[stage_idx + 1]
+        stage_in_res = stage_resolutions[stage_idx]
+        stage_out_res = stage_resolutions[stage_idx + 1]
+
+        save_stage_input_mesh(
+            mesh=current_mesh,
+            out_dir=output_folder,
+            name=name,
+            stage_in_res=stage_in_res,
+            stage_out_res=stage_out_res,
+            stage_in_depth=stage_in_depth,
+        )
+        
+        t_start = time.time()
+        stage_data = mul_encode_mesh(
+            raw_mesh=current_mesh,
+            gt_mesh=gt_mesh,
+            codec=codec,
+            tokenizer_target_depth=stage_in_depth,
+            pc_num=config.data_processing.pc_num,
+            total_pc_num=config.data_processing.total_pc_num,
+            use_gt_pc=use_gt_pc,
+            clean=False
+        )
+        t_mul_encode = time.time()
+        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | mul_encode_mesh time: {t_mul_encode - t_start:.4f}s")
+
+        raw_token = stage_data["raw_token"].unsqueeze(0).to(device)
+        pc = stage_data["pc"].unsqueeze(0).to(device)
+    
+        t_start = time.time()
+        # refine_token = sample_refine_tokens_with_topology(
+        #     model,
+        #     mask_token_id,
+        #     raw_token,
+        #     pc,
+        #     use_gumbel=False,
+        #     gumbel_temp_init=1.0,
+        #     use_soft_consensus=False,
+        #     consensus_start_step_ratio=0,
+        #     cfg_scale=1.0,
+        #     T_infer=config.evaluation.diffusion_sample_steps,
+        # )
+        refine_token = sample_refine_tokens_with_topology_fm(
+            model,
+            mask_token_id,
+            raw_token,
+            pc,
+            use_gumbel=False,
+            gumbel_temp_init=1.0,
+            use_soft_consensus=False,
+            consensus_start_step_ratio=0,
+            cfg_scale=1.0,
+            T_infer=config.evaluation.diffusion_sample_steps,
+        )
+        
+        torch.cuda.synchronize()
+        t_sample = time.time()
+        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | sample_refine_tokens_with_topology time: {t_sample - t_start:.4f}s")
+        
+        refine_token = restore_non_pad_token(refine_token, shift=token_shift, pad_id=pad_id)
+
+        t_start = time.time()
+        stage_suffix = f"_stage_{stage_in_res}to{stage_out_res}"
+        stage_mesh = save_stage_prediction(
+            codec=codec,
+            raw_token=raw_token,
+            refine_token=refine_token,
+            out_dir=output_folder,
+            name=name,
+            target_depth=stage_out_depth,
+            suffix=stage_suffix,
+            fix_holes=fix_holes,
+            raw_dis_mesh=stage_data.get("raw_dis_mesh"),
+            gt_dis_mesh=stage_data.get("gt_dis_mesh"),
+        )
+        t_save_pred = time.time()
+        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | save_stage_prediction (decode + fix) time: {t_save_pred - t_start:.4f}s")
+        print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | total stage time: {time.time() - stage_start_time:.4f}s\n")
+
+        if stage_mesh is None:
+            break
+
+        current_mesh = stage_mesh
+        
+    print(f"[{name}] Total process_one_mesh time: {time.time() - t_total_start:.4f}s")
+
+
+class EMA:
+    def __init__(self, model, decay=0.999):
+        self.decay = decay
+        self.shadow = {}
+        self.register(model)
+
+    def register(self, model):
+        self.shadow = {}
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                self.shadow[name] = param.data.clone()
+
+    def load_state_dict(self, state):
+        self.decay = state.get("decay", self.decay)
+        self.shadow = {k: v.clone() for k, v in state["shadow"].items()}
+
+    @torch.no_grad()
+    def copy_to(self, model):
+        """Swap shadow weights into the model in-place (for inference)."""
+        copied, missing = 0, 0
+        for name, param in model.named_parameters():
+            if name in self.shadow:
+                param.data.copy_(self.shadow[name].to(device=param.device, dtype=param.dtype))
+                copied += 1
+            else:
+                missing += 1
+        print(f"[EMA.copy_to] copied {copied} params into model ({missing} not present in shadow)")
+
+
+def run_inference_on_all_checkpoints(model, mask_token_id, accelerator, config, device, pad_id, codec, token_shift, ema=None):
+    ckpt_dir = "ckpt"
+    output_folder = "./pretrain_model_mul_more_data_64_to_256_eval_complex_test_no_fix_holes_no_clean_cfg_1.0_fm"
+
+    # raw_data_txt = "dataset/bpt_res_list.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_retrain_gt/data.txt
+    # gt_data_txt = "dataset/trellis_clean_mesh.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_gt/data.txt
+    raw_data_txt = "demo/complex_test_data/files/data.txt"
+    gt_data_txt = "demo/complex_test_data/files/data.txt"
+    # raw_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt"
+    # gt_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt"
+    stage_resolutions = [64, 128, 256]
+    # raw_data_txt = "demo/tripo 1.0/data.txt"
+    # gt_data_txt = "demo/tripo 1.0/data.txt"
+    
+    use_gt_pc=True
+    fix_holes=False
+    if accelerator.is_main_process:
+        print(f"Using raw data txt: {raw_data_txt}")
+        print(f"Using gt  data txt: {gt_data_txt}")
+        print(f"Searching checkpoints in: {ckpt_dir}")
+
+    # mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=config.data_processing.max_face_num, min_faces=config.data_processing.min_face_num)
+    mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=20000, min_faces=500)
+    if accelerator.is_main_process:
+        print(f"Matched raw/gt pairs: {len(mesh_pairs)}")
+        if missing_gt:
+            print(f"Raw without matched gt: {len(missing_gt)}")
+
+    search_pattern = os.path.join(ckpt_dir, "checkpoint_*_epoch_*_loss.pth")
+    ckpt_files = glob.glob(search_pattern)
+    if not ckpt_files:
+        if accelerator.is_main_process:
+            print("No checkpoints found!")
+        return
+
+    ckpt_data = []
+    for f in ckpt_files:
+        match = re.search(r"checkpoint_(\d+)_epoch_", f)
+        if match:
+            epoch = int(match.group(1))
+            ckpt_data.append((epoch, f))
+    ckpt_data.sort(key=lambda x: x[0])
+
+    # stage_resolutions = [32, 64, 128, 256]
+    
+
+    # stage_resolutions = [128, 256]
+
+    for epoch, ckpt_path in ckpt_data:
+        if accelerator.is_main_process:
+            print(f"\n[{'='*40}]")
+            print(f"Starting inference for epoch {epoch}")
+            print(f"Loading weights from: {ckpt_path}")
+
+        output_folder_ply = os.path.join(output_folder, f"{epoch}_epoch_val_generate_mul")
+        if accelerator.is_main_process:
+            os.makedirs(output_folder_ply, exist_ok=True)
+
+        checkpoint = torch.load(ckpt_path, map_location=device)
+        model_to_load = accelerator.unwrap_model(model)
+        # Use strict=False so we can load older ckpts (no null_cond) into the
+        # newer model class, and vice versa. Missing/unexpected keys are logged.
+        load_result = model_to_load.load_state_dict(checkpoint["model_state_dict"], strict=False)
+        missing = list(getattr(load_result, "missing_keys", []) or [])
+        unexpected = list(getattr(load_result, "unexpected_keys", []) or [])
+        if (missing or unexpected) and accelerator.is_main_process:
+            print(f"[load_state_dict] strict=False -> missing={len(missing)} unexpected={len(unexpected)}")
+            if missing:
+                print(f"  missing: {missing[:8]}{' ...' if len(missing) > 8 else ''}")
+            if unexpected:
+                print(f"  unexpected: {unexpected[:8]}{' ...' if len(unexpected) > 8 else ''}")
+
+        # Auto-detect EMA: if the ckpt has an `ema_state_dict`, swap those
+        # shadow weights into the model for inference. Otherwise fall back to
+        # the raw `model_state_dict` we already loaded above. No reliance on
+        # config `use_ema` / a separate EMA object — just check the ckpt.
+        ema_state = checkpoint.get("ema_state_dict")
+        shadow = ema_state.get("shadow") if isinstance(ema_state, dict) else None
+        if shadow:
+            with torch.no_grad():
+                copied, skipped = 0, 0
+                for name, param in model_to_load.named_parameters():
+                    if name in shadow:
+                        param.data.copy_(shadow[name].to(device=param.device, dtype=param.dtype))
+                        copied += 1
+                    else:
+                        skipped += 1
+            if accelerator.is_main_process:
+                msg = f"[Epoch {epoch}] EMA shadow detected in ckpt -> applied to model ({copied} params"
+                if skipped:
+                    msg += f", {skipped} model params not in shadow"
+                print(msg + ")")
+        else:
+            if accelerator.is_main_process:
+                print(f"[Epoch {epoch}] no EMA in ckpt -> running with raw model_state_dict weights.")
+
+        with torch.no_grad():
+            model.eval()
+            pbar = tqdm(
+                enumerate(mesh_pairs),
+                total=len(mesh_pairs),
+                desc=f"Infer Epoch {epoch}",
+                disable=not accelerator.is_local_main_process,
+                dynamic_ncols=True,
+            )
+            for line_idx, (raw_mesh_path, gt_mesh_path) in pbar:
+                if line_idx % accelerator.num_processes != accelerator.process_index:
+                    continue
+                try:
+                    process_one_mesh(
+                        raw_mesh_path=raw_mesh_path,
+                        gt_mesh_path=gt_mesh_path,
+                        model=model,
+                        mask_token_id=mask_token_id,
+                        codec=codec,
+                        device=device,
+                        pad_id=pad_id,
+                        token_shift=token_shift,
+                        config=config,
+                        output_folder=output_folder_ply,
+                        stage_resolutions=stage_resolutions,
+                        fix_holes=fix_holes,
+                        use_gt_pc=use_gt_pc,
+                    )
+                except Exception as e:
+                    print(f"[infer error] raw={raw_mesh_path}, gt={gt_mesh_path}: {str(e)}")
+
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    if accelerator.is_main_process:
+        print("\nBatch inference complete for all checkpoints.")
+
+
+def main(config):
+    print("start init Accelerator...")
+    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+    accelerator = Accelerator(
+        gradient_accumulation_steps=config.accelerator.gradient_accumulation_steps,
+        mixed_precision=config.accelerator.mixed_precision,
+        kwargs_handlers=[ddp_kwargs],
+    )
+    device = accelerator.device
+
+    codec = MeshCodec(args=config.tokenizer)
+    codec_config = codec.get_model_config()
+    input_num_categories = codec_config["input_num_categories"]
+    lm_num_categories = codec_config["output_num_categories"]
+    pad_id = codec_config["pad_token_id"]
+    token_shift = input_num_categories - 1
+
+    model = DiscreteDiffusionTransformer(
+        vocab_A=input_num_categories,
+        vocab_B=lm_num_categories,
+        max_seq_len=config.data_processing.max_len,
+        d_model=config.model.feature_dim,
+        encoder_freeze=config.model.encoder_freeze,
+        nhead=config.model.num_heads,
+        depth=config.model.hourglass_depth,
+        dropout=config.model.dropout,
+        time_embed_scale=config.training.diffusion_train_steps,
+    )
+    mask_token_id = model.mask_token_id
+
+    use_ema = getattr(config.training, "use_ema", False)
+    ema_decay = getattr(config.training, "ema_decay", 0.999)
+    ema = EMA(model, decay=ema_decay) if use_ema else None
+
+    model = accelerator.prepare(model)
+
+    run_inference_on_all_checkpoints(
+        model=model,
+        mask_token_id=mask_token_id,
+        accelerator=accelerator,
+        config=config,
+        device=device,
+        pad_id=pad_id,
+        codec=codec,
+        token_shift=token_shift,
+        ema=ema,
+    )
+
+
+if __name__ == "__main__":
+    set_seed(42)
+    config = load_config()
+    main(config)

mulres_infer/modules/token_utils.py

@@ -0,0 +1,52 @@
+"""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()

mulres_infer/readme.md

@@ -0,0 +1,122 @@
+# 多分辨率网格生成 · 推理
+
+点云条件下的**多分辨率八叉树网格生成**：低分辨率生成器逐级长出粗网格，高分辨率细分器把它推到 512 格。
+
+## 流程
+
+```
+输入 mesh ──采样点云──┐
+                      ▼
+  ┌───────── 生成器 (TSSR 三前向离散扩散) ─────────┐
+  │  格2 → 格4 → 格8 → 格16 → 格32 → 格64  每级20步 │
+  │                        └补洞┘  └补洞┘          │
+  └──────────────────────┬─────────────────────────┘
+                         ▼
+  ┌────────── 细分器 (acc_uni) ──────────┐
+  │  格64 → 格128 → 格256 → 格512  每级20步 │
+  └──────────────────────────────────────┘
+                         ▼
+               final_res512_plain.obj
+```
+
+**生成器**每级把上一级网格重新编码成八叉树 token，用离散扩散预测该级完整子码。
+TSSR 解码每步三次前向：mask 预测 → no-mask 精修 → 逐 token 分类器挑可疑位置 re-mask，
+避免 MaskGIT 式"一旦揭示即定死"的错误累积。
+
+**细分器**只做细分不做生成，粗输入上的破洞它补不出来——这是 res32/res64 要先补洞再喂下去的原因。
+
+## 目录
+
+```
+.
+├── scripts/            入口
+│   ├── infer.py            端到端推理
+│   ├── eval.py             评测(CD / HD / 法向一致性)
+│   └── strip_ckpt.py       权重瘦身: 剥离优化器、合并 EMA
+├── models/             模型
+│   ├── generator.py        生成器(骨干 → 双云注入 → TSSR)
+│   ├── subdivider.py       细分器
+│   └── point_conditioner.py 点云条件编码器
+├── modules/            算子与管线
+│   ├── pipeline.py         多级级联主管线
+│   ├── attention.py        MLA 等注意力原语
+│   ├── octree_codec.py     八叉树 tokenizer(numba 加速)
+│   ├── octree_codec_acc.py 细分器口径的 tokenizer
+│   ├── octree_decode.py    子码 → mesh 解码
+│   ├── token_utils.py      token 通路
+│   └── miche/              Michelangelo 点云编码器(第三方)
+├── dataset/            数据处理与点云采样
+├── utils/              配置、网格工具、补洞、指标
+├── configs/            generator.yaml / subdivider.yaml
+└── ckpt/               权重(见 ckpt/README.md)
+```
+
+## 安装
+
+```bash
+conda create -n mulres python=3.12 -y && conda activate mulres
+pip install -r requirements.txt
+# flash-attn 用与 torch/CUDA 匹配的预编译 wheel, 不要从源码构建
+```
+
+## 使用
+
+```bash
+# 全部 mesh
+python scripts/infer.py --data /path/to/meshes --out_dir outputs/run1
+
+# 只跑前 20 个
+python scripts/infer.py --data /path/to/meshes --out_dir outputs/run1 --ncase 20
+
+# 指定 case(文件名子串, 带尾点防前缀误配: "a." 不匹配 "a1.glb")
+python scripts/infer.py --data /path/to/meshes --out_dir outputs/run1 --case "model_a.,model_b."
+
+# 面数太密: 给连接头"空连接"类加 logit 偏置(0.3 约 -25% 面, 保真度几乎不变)
+python scripts/infer.py --data /path/to/meshes --out_dir outputs/run1 --empty_bias 0.3
+```
+
+每个 case 输出一个子目录：
+
+| 文件 | 说明 |
+|---|---|
+| `final_res512_plain.obj` | **最终结果** |
+| `res{2,4,8,16,32,64}_plain.obj` | 生成器各级产物(补洞后) |
+| `*_sr_64to128_depth7.obj` 等 | 细分器各级产物 |
+| `gt_original.obj` | 归一化后的原始输入(塌缩前) |
+| `gt_original_dec.obj` | 塌缩后的输入(级联起点与条件点云的来源) |
+| `gt_res{N}.obj` | 输入经 tokenizer 编解码的天花板, 用于逐级对照 |
+| `pc_full.ply` / `pc_target_res{N}.ply` | 注入的条件点云 |
+
+## 评测
+
+```bash
+# 预测目录需扁平: <case>_y.obj ; GT 目录: <case>.glb 或 .obj
+python scripts/eval.py /path/to/pred_flat --gt /path/to/gt_meshes
+```
+
+输出倒角距离(L1/L2)、豪斯多夫距离、法向一致性 |NC| 与 NC。
+
+> 解码不编码面的绕序，有向法向随机，所以看 **|NC|**（约 0.85），`NC (Real)` 接近 0 属预期。
+
+## 参数
+
+`scripts/infer.py` 固化了默认配方，更细的控制用环境变量（定义见 `modules/pipeline.py` 顶部）：
+
+| 变量 | 默认 | 说明 |
+|---|---|---|
+| `STAGE` | `2,4,8,16,32,64` | 生成器级联阶梯(产物真实体素格) |
+| `STEPS` | `20` | 生成器每级扩散步数 |
+| `SUPERRES_TO` / `SUPERRES_STEPS` | `512` / `20` | 细分终点与每级步数 |
+| `FILL` / `FILL_MIN_RES` | `1` / `0` | 级联内补洞; `FILL_MIN_RES=64` 则只在格 64 补 |
+| `SR_FILL_IN` / `SUPERRES_FIX_HOLES` | `1` / `1` | 细分入口 / 每级输出补洞 |
+| `DECIMATE_RATIO` | `0.5` | 输入预塌缩比(0=不塌缩) |
+| `EMPTY_BIAS` | `0` | >0 稀疏化控面数 |
+| `CASE_STEM` | — | 文件名子串过滤(在 `NCASE` 截断**之后**生效) |
+
+## 已知行为
+
+- **步数**：20 与 40 步保真度无差别；10 步会让开边 +27%、非流形 +34%。20 是性价比点。
+- **补洞是置换不是净赚**：开边 −16% 的代价是非流形 +12%。要纯净拓扑用 `--no_fill`。
+- **分辨率终点**：生成器自己走到格 64 比走到格 32 再交给细分器更好（非流形 −59%、CD −6%），
+  代价是开边变多——细分器会把粗输入的破洞一路放大。
+- **GPU 非确定性**：flash-attn / bf16 的归约顺序不固定，同配置重跑面数有约 3% 波动，单 case 开边波动更大。

mulres_infer/requirements.txt

@@ -0,0 +1,22 @@
+# 核心
+torch==2.7.0
+numpy
+scipy
+trimesh==5.1.0
+numba
+einops
+omegaconf
+transformers
+accelerate
+timm
+scikit-image
+pymeshfix
+pyvista
+tqdm
+
+# flash-attn 需与 torch/CUDA 匹配, 用官方预编译 wheel 安装, 不要从源码构建:
+#   pip install flash_attn-2.8.3+cu12torch2.7cxx11abiFALSE-cp312-cp312-linux_x86_64.whl
+
+# 可选: utils/postprocessing_utils.py 顶部导入, 但补洞路径实际不调用其功能
+# utils3d @ git+https://github.com/EasternJournalist/utils3d.git
+# nvdiffrast @ git+https://github.com/NVlabs/nvdiffrast.git

mulres_infer/scripts/eval.py

@@ -0,0 +1,429 @@
+
+import os, sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from utils.distance import calculate_distances_numpy, calculate_normal_consistency_score
+import os
+import numpy as np
+from tqdm import tqdm
+import trimesh
+import json
+import glob
+import multiprocessing # 导入多进程模块
+MIN_FACE_AMOUNT = 0
+MAX_FACE_AMOUNT = 100000
+# ===============================================================
+# 第一步：你的距离计算函数 (无变化)
+# ===============================================================
+pc_num = 1024
+
+def normalize(mesh, file_path):
+    # ... 现有代码 ...
+    bounds = np.array([mesh.vertices.min(axis=0), mesh.vertices.max(axis=0)])
+    mesh.vertices = mesh.vertices - (bounds[0] + bounds[1])[None, :] / 2
+    mesh.vertices = mesh.vertices / (bounds[1] - bounds[0]).max()
+    mesh.vertices = mesh.vertices.clip(-0.5, 0.5) # 限制顶点坐标在[-0.5, 0.5] 
+
+# ===============================================================
+# 第二步：读取 .obj 文件的函数 (无变化)
+# ===============================================================
+def sample_from_obj(file_path, alter=False):
+    # ... 现有代码 ...
+    """
+    从 .obj 文件中加载顶点坐标。
+    """
+    mesh = trimesh.load(file_path, force="mesh")
+    if alter:
+        mesh.vertices = mesh.vertices[:, [1,2,0]]
+    normalize(mesh, file_path)
+    points, face_idx = mesh.sample(pc_num, return_index=True)
+    normals = mesh.face_normals[face_idx]
+    return points,normals
+
+# ===============================================================
+# [修改] 法向量一致性计算
+# ===============================================================
+
+# ===============================================================
+# 第三步：[修改] 多进程工作函数
+# ===============================================================
+
+def process_common_file_worker(base_name, gt_folder, folder_paths):
+    """
+    [已修改] 这是一个工作进程执行的任务：
+    处理单个 base_name 对应的所有文件夹。
+    它会加载一次GT文件，然后遍历所有预测文件夹进行计算。
+    返回 L1, L2 CD 距离, |NC| 和 NC。
+    """
+    
+    # 1. 查找并加载 GT 文件
+    gt_filename = base_name + '.glb'
+    # ... 现有代码 ...
+    gt_filepath = os.path.join(gt_folder, gt_filename)
+
+    if not os.path.exists(gt_filepath):
+        gt_filename = base_name + '.obj'
+    # ... 现有代码 ...
+        gt_filepath = os.path.join(gt_folder, gt_filename)
+        if not os.path.exists(gt_filepath):
+            # 警告：在工作进程中打印 (print) 不是线程安全的，
+    # ... 现有代码 ...
+            # 但返回空列表是安全的，主进程会处理它。
+            # print(f"\n警告：共同文件 '{base_name}' 对应的GT文件 '{gt_filepath}' 不存在，已跳过。")
+            return [] # 返回空列表表示此文件处理失败
+    mesh = trimesh.load(gt_filepath, force="mesh")
+    if len(mesh.faces) < MIN_FACE_AMOUNT or len(mesh.faces) > MAX_FACE_AMOUNT:
+        return []
+        
+    try:
+        pcd_gt, normal_gt = sample_from_obj(gt_filepath)
+        if pcd_gt is None or pcd_gt.size == 0:
+    # ... 现有代码 ...
+            # print(f"\n警告：加载GT文件 {gt_filename} 失败或点云为空。已跳过。")
+            return []
+    except Exception as e:
+    # ... 现有代码 ...
+        # print(f"\n加载GT文件 {gt_filename} 时发生错误: {e}。已跳过。")
+        return []
+
+    # 2. 遍历所有文件夹，加载 Pred 文件并计算
+    results_for_this_file = []
+
+    for folder_path in folder_paths:
+        pred_filename = base_name + '_y.obj'
+    # ... 现有代码 ...
+        pred_filepath = os.path.join(folder_path, pred_filename)
+        
+        try:
+            pcd_pred, normal_pred = sample_from_obj(pred_filepath)
+            if pcd_pred is None or pcd_pred.size == 0:
+    # ... 现有代码 ...
+                # print(f"\n警告：加载文件 {pred_filepath} 失败或点云为空。已跳过。")
+                results_for_this_file.append((folder_path, None, None, None, None, None)) # [修改] 6 个 None
+                continue
+            
+            # [修改] 接收 L1 和 L2 CD
+            chamfer_l2, hausdorff, chamfer_l1 = calculate_distances_numpy(pcd_gt, pcd_pred)
+            # [修改] 接收两个 NC 分数
+            normal_error_abs, normal_error_real = calculate_normal_consistency_score(pcd_gt, normal_gt, pcd_pred, normal_pred)
+            
+            # [修改] 存储有效结果 (6个值)
+            results_for_this_file.append((folder_path, chamfer_l2, hausdorff, normal_error_abs, chamfer_l1, normal_error_real))
+            
+        except Exception as e:
+    # ... 现有代码 ...
+            # print(f"\n计算文件对 {pred_filepath} 时发生错误: {e}。已跳过。")
+            results_for_this_file.append((folder_path, None, None, None, None, None)) # [修改] 6 个 None
+    
+    # 3. 返回此 base_name 对应的所有结果
+    return results_for_this_file
+
+def starmap_helper(args):
+    # ... 现有代码 ...
+    """
+    [新增] 辅助函数，用于解包 pool.imap 的参数，使其能接受多个参数。
+    """
+    return process_common_file_worker(*args)
+
+# ===============================================================
+# 第四步：[修改] 评估函数，使用多进程
+# ===============================================================
+
+def evaluate_common_files(folder_paths, gt_folder, base_dir):
+    """
+    [已修改] 评估多个文件夹中 *共有* 的文件，并计算总平均指标和逐文件夹的平均指标。
+    现在使用 multiprocessing.Pool 来并行处理文件。
+    """
+    if not folder_paths:
+    # ... 现有代码 ...
+        print("错误：没有提供用于评估共同文件的文件夹。")
+        return None
+
+    print("\n" + "="*80)
+    # ... 现有代码 ...
+    print(" 正在计算共同文件指标 ".center(80, "="))
+    print(f"总共 {len(folder_paths)} 个文件夹参与比较。")
+
+    # 1. 查找每个文件夹中的文件基础名（不含 '_y.obj'）(无变化)
+    files_by_folder = {}
+    # ... 现有代码 ...
+    for folder_path in folder_paths:
+        try:
+            all_files = os.listdir(folder_path)
+    # ... 现有代码 ...
+            base_names = set([f.replace('_y.obj', '') for f in all_files
+                if (f.endswith('_y.obj') and not f.endswith('_gt_y.obj') and not f.endswith("input_y.obj"))]
+                )
+            files_by_folder[folder_path] = base_names
+    # ... 现有代码 ...
+            print(f" ‐ 文件夹 '{os.path.basename(folder_path)}' 找到 {len(base_names)} 个文件。")
+        except Exception as e:
+            print(f" ‐ 错误：无法读取文件夹 {folder_path}: {e}")
+    # ... 现有代码 ...
+            files_by_folder[folder_path] = set()
+
+    # 2. 找到所有文件夹中 *共有* 的文件基础名 (无变化)
+    common_base_names = set.intersection(*files_by_folder.values())
+    
+    if not common_base_names:
+    # ... 现有代码 ...
+        print("\n错误：在所有文件夹中没有找到任何共有的 `_y.obj` 文件。")
+        print("="*80)
+        return None
+
+    print(f"\n在所有文件夹中找到了 {len(common_base_names)} 个共同文件。将使用多进程处理...")
+
+    # 3. [修改] 初始化指标存储
+    global_chamfer_l2 = []
+    global_hausdorff = []
+    global_normal_abs = [] # [修改]
+    global_chamfer_l1 = []
+    global_normal_real = [] # [新增]
+    
+    per_folder_metrics = {path: {'chamfer_l2': [], 'hausdorff': [], 'normal_abs': [], 'chamfer_l1': [], 'normal_real': []} for path in folder_paths} # [修改]
+
+    # 4. [修改] 设置并运行多进程任务
+    
+    # 准备任务参数列表
+    # ... 现有代码 ...
+    task_args = [(base_name, gt_folder, folder_paths) for base_name in common_base_names]
+    all_results_list = []
+    
+    # 确定进程数，可以适当调整，例如 os.cpu_count()
+    # ... 现有代码 ...
+    # 留出2个核心给系统，防止卡死
+    num_processes = max(1, 60) 
+    print(f"启动 {num_processes} 个工作进程...")
+
+    # 创建进程池
+    # ... 现有代码 ...
+    with multiprocessing.Pool(processes=num_processes) as pool:
+        # 使用 pool.imap 来迭代任务，并用 tqdm 显示进度条
+        # imap 会按顺序返回结果，同时保持并行执行
+    # ... 现有代码 ...
+        for result in tqdm(pool.imap(starmap_helper, task_args), 
+                           total=len(common_base_names), 
+                           desc="评估共同文件"):
+            all_results_list.append(result)
+
+    # 5. [修改] 聚合所有进程返回的结果
+    print("\n正在聚合所有进程的结果...")
+    valid_pairs_processed = 0
+    for results_for_this_file in all_results_list:
+        if not results_for_this_file:
+    # ... 现有代码 ...
+            # GT 文件加载失败或跳过
+            continue
+            
+        # [修改] 解包 L1, L2, |NC|, NC
+        for (folder_path, chamfer_l2, hausdorff, normal_error_abs, chamfer_l1, normal_error_real) in results_for_this_file:
+            if chamfer_l2 is not None and hausdorff is not None and normal_error_abs is not None and chamfer_l1 is not None and normal_error_real is not None:
+                # 这是一个有效的结果
+                global_chamfer_l2.append(chamfer_l2)
+                global_hausdorff.append(hausdorff)
+                global_normal_abs.append(normal_error_abs) # [修改]
+                global_chamfer_l1.append(chamfer_l1)
+                global_normal_real.append(normal_error_real) # [新增]
+                
+                per_folder_metrics[folder_path]['chamfer_l2'].append(chamfer_l2)
+                per_folder_metrics[folder_path]['hausdorff'].append(hausdorff)
+                per_folder_metrics[folder_path]['normal_abs'].append(normal_error_abs) # [修改]
+                per_folder_metrics[folder_path]['chamfer_l1'].append(chamfer_l1)
+                per_folder_metrics[folder_path]['normal_real'].append(normal_error_real) # [新增]
+                valid_pairs_processed += 1
+            else:
+    # ... 现有代码 ...
+                # 预测文件加载失败或计算错误
+                # 可以在这里添加日志
+                pass
+    
+    print(f"聚合完成。总共处理了 {valid_pairs_processed} 个有效的 (GT, Pred) 文件对。")
+
+    # 6. [修改] 计算最终平均值
+    if not global_chamfer_l2: # [修改]
+        print("\n错误：未能成功计算任何共同文件的指标。")
+        print("="*80)
+        return None
+
+    # [修改] 计算总平均
+    global_avg_results = {
+        "evaluated_pairs_total": len(global_chamfer_l2),
+        "common_files_count": len(common_base_names),
+        "folders_count": len(folder_paths),
+        "avg_chamfer_distance_l2": np.mean(global_chamfer_l2),
+        "avg_hausdorff_distance": np.mean(global_hausdorff),
+        "avg_normal_consistency_abs": np.mean(global_normal_abs), # [修改]
+        "avg_chamfer_distance_l1": np.mean(global_chamfer_l1),
+        "avg_normal_consistency_real": np.mean(global_normal_real) # [新增]
+    }
+
+    # [修改] 计算逐文件夹平均
+    per_folder_avg_results = {}
+    for folder_path, metrics in per_folder_metrics.items():
+        folder_name = os.path.basename(folder_path)
+        if metrics['chamfer_l2']: # [修改]
+            per_folder_avg_results[folder_name] = {
+                "evaluated_pairs": len(metrics['chamfer_l2']),
+                "avg_chamfer_distance_l2": np.mean(metrics['chamfer_l2']),
+                "avg_hausdorff_distance": np.mean(metrics['hausdorff']),
+                "avg_normal_consistency_abs": np.mean(metrics['normal_abs']), # [修改]
+                "avg_chamfer_distance_l1": np.mean(metrics['chamfer_l1']),
+                "avg_normal_consistency_real": np.mean(metrics['normal_real']) # [新增]
+            }
+        else:
+            per_folder_avg_results[folder_name] = {
+                "evaluated_pairs": 0,
+                "avg_chamfer_distance_l2": 0,
+                "avg_hausdorff_distance": 0,
+                "avg_normal_consistency_abs": 0, # [修改]
+                "avg_chamfer_distance_l1": 0,
+                "avg_normal_consistency_real": 0 # [新增]
+            }
+            
+    # 7. [修改] 打印共同文件评估总结
+    print("\n--- 共同文件评估总结 ---")
+    print(f"\n全局平均 (基于 {len(global_chamfer_l2)} 个总比较对, {len(common_base_names)} 个文件 x {len(folder_paths)} 个文件夹):")
+    print(f" ‐ 平均倒角距离 (L2): {global_avg_results['avg_chamfer_distance_l2']:.6f}")
+    print(f" ‐ 平均倒角距离 (L1): {global_avg_results['avg_chamfer_distance_l1']:.6f}")
+    print(f" ‐ 平均豪斯多夫距离: {global_avg_results['avg_hausdorff_distance']:.6f}")
+    print(f" ‐ 平均法向量 |NC| (Abs): {global_avg_results['avg_normal_consistency_abs']:.6f}") # [修改]
+    print(f" ‐ 平均法向量 NC (Real): {global_avg_results['avg_normal_consistency_real']:.6f}") # [新增]
+    
+    print("\n逐文件夹平均 (仅基于共同文件):")
+    for folder_name, results in per_folder_avg_results.items():
+        print(f" ‐ 文件夹: {folder_name}")
+        # [修改] 打印所有6个指标
+        print(f" ‐ 倒角(L2): {results['avg_chamfer_distance_l2']:.6f}, 倒角(L1): {results['avg_chamfer_distance_l1']:.6f}, 豪斯多夫: {results['avg_hausdorff_distance']:.6f}, 法向量|NC|: {results['avg_normal_consistency_abs']:.6f}, 法向量NC: {results['avg_normal_consistency_real']:.6f} (共 {results['evaluated_pairs']} 对)")
+    
+    print("="*80)
+    
+    # 8. 保存结果 (无变化)
+    final_results = {
+    # ... 现有代码 ...
+        "global_average_on_common_files": global_avg_results,
+        "per_folder_average_on_common_files": per_folder_avg_results
+    }
+    
+    output_filepath = os.path.join(base_dir, "common_evaluation_results.json")
+    # ... 现有代码 ...
+    try:
+        with open(output_filepath, 'w', encoding='utf-8') as f:
+            json.dump(final_results, f, ensure_ascii=False, indent=4)
+    # ... 现有代码 ...
+        print(f"\n共同文件评估结果已保存到: {output_filepath}")
+    except Exception as e:
+        print(f"\n错误：保存共同文件结果时失败: {e}")
+    # ... 现有代码 ...
+        
+    return final_results
+
+
+# --- [主函数：run_evaluation_on_common_files_only] (打印部分已修改) ---
+def run_evaluation_on_common_files_only(target_folders, gt_folder, base_dir=None):
+    """
+    [无变化]
+    在手动指定的多个文件夹上 *仅* 运行 *共同文件* 评估。
+    """
+    
+    # 1. 验证输入的文件夹列表
+    # ... 现有代码 ...
+    if not isinstance(target_folders, list) or not target_folders:
+        print("错误：未提供任何文件夹用于评估。请检查 'target_folders' 列表。")
+        return
+
+    # 2. 检查路径有效性并过滤非文件夹
+    # ... 现有代码 ...
+    valid_folders = []
+    for f_path in target_folders:
+        if os.path.isdir(f_path):
+    # ... 现有代码 ...
+            valid_folders.append(f_path)
+        else:
+            print(f"警告：提供的路径 '{f_path}' 不是一个有效的文件夹，已跳过。")
+    
+    if not valid_folders:
+    # ... 现有代码 ...
+        print("错误：提供的列表中没有找到任何有效的文件夹。")
+        return
+
+    # 确保排序一致
+    # ... 现有代码 ...
+    valid_folders.sort()
+    
+    # 3. 确定用于保存汇总文件的 base_dir
+    if base_dir is None:
+    # ... 现有代码 ...
+        # 如果未指定，使用第一个有效文件夹的父目录
+        base_dir = os.path.dirname(valid_folders[0])
+        print(f"信息：未指定 base_dir，将使用 '{base_dir}' 存放共同文件评估结果。")
+    elif not os.path.isdir(base_dir):
+    # ... 现有代码 ...
+        print(f"警告：指定的 base_dir '{base_dir}' 不存在。将改用第一个文件夹的父目录。")
+        base_dir = os.path.dirname(valid_folders[0])
+    
+    print(f"\n找到了 {len(valid_folders)} 个有效的文件夹，将评估它们的共同文件：")
+    # ... 现有代码 ...
+    for folder in valid_folders:
+        print(f" - {folder}") # 打印完整路径
+
+    # 4. 第1部分：直接进行共同文件评估
+    # ... 现有代码 ...
+    print("\n" + "#"*80)
+    print(" 开始共同文件评估 (仅计算所有文件夹共有的文件) ".center(80, "#"))
+    print("#"*80)
+    
+    # [调用修改后的函数]
+    # ... 现有代码 ...
+    common_results = evaluate_common_files(valid_folders, gt_folder, base_dir)
+
+    # 5. [修改] 第2部分：打印最终的汇总报告
+    print("\n" + "#"*80)
+    print(" 最终汇总报告 ".center(80, "#"))
+    print("#"*80)
+
+    print("\n--- 共同文件评估结果 (仅基于共同文件) ---")
+    if common_results:
+        # 打印全局平均
+        global_avg = common_results.get("global_average_on_common_files", {})
+        if global_avg:
+            print(f"\n [总平均 (基于 {global_avg.get('common_files_count', 0)} 个共同文件)]")
+            print(f" ‐ 倒角(L2): {global_avg.get('avg_chamfer_distance_l2', 0):.6f}, "
+                  f"倒角(L1): {global_avg.get('avg_chamfer_distance_l1', 0):.6f}, "
+                  f"豪斯多夫: {global_avg.get('avg_hausdorff_distance', 0):.6f}, "
+                  f"法向量|NC|: {global_avg.get('avg_normal_consistency_abs', 0):.6f}, " # [修改]
+                  f"法向量NC: {global_avg.get('avg_normal_consistency_real', 0):.6f}") # [新增]
+        
+        # 打印逐文件夹平均
+        per_folder_avg = common_results.get("per_folder_average_on_common_files", {})
+        if per_folder_avg:
+            print("\n [逐文件夹平均 (仅共同文件)]")
+            for folder_name, results in per_folder_avg.items():
+                print(f" ‐ 文件夹: {folder_name}")
+                print(f" ‐ 倒角(L2): {results.get('avg_chamfer_distance_l2', 0):.6f}, "
+                      f"倒角(L1): {results.get('avg_chamfer_distance_l1', 0):.6f}, "
+                      f"豪斯多夫: {results.get('avg_hausdorff_distance', 0):.6f}, "
+                      f"法向量|NC|: {results.get('avg_normal_consistency_abs', 0):.6f}, " # [修改]
+                      f"法向量NC: {results.get('avg_normal_consistency_real', 0):.6f} " # [新增]
+                      f"(共 {results.get('evaluated_pairs', 0)} 对)")
+    else:
+        print(" (没有共同文件的评估结果)")
+        
+    print("\n" + "="*80)
+    print("所有评估任务已完成。")
+    print("="*80)
+
+
+# ===============================================================
+# 第五步：运行主程序 (无变化)
+# ===============================================================
+        
+if __name__ == "__main__":
+    import argparse
+    ap = argparse.ArgumentParser(description="评测: 预测目录(含 <case>_y.obj) vs GT 目录(<case>.glb/.obj)")
+    ap.add_argument("folders", nargs="+", help="一个或多个预测目录(只评估各目录共有的 case)")
+    ap.add_argument("--gt", required=True, help="Ground Truth 目录")
+    ap.add_argument("--out", default=None, help="汇总 JSON 的输出目录(默认第一个预测目录)")
+    a = ap.parse_args()
+    if not os.path.isdir(a.gt):
+        raise SystemExit(f"GT 目录不存在: {a.gt}")
+    run_evaluation_on_common_files_only(a.folders, a.gt, base_dir=a.out)