Hugging Face's logo

MayMeHappy
/
mul

Model card Files
xet
 Community
mul / modules /pipeline.py
MayMeHappy's picture
MayMeHappy
Upload folder using huggingface_hub
90ca518 verified
Raw History Blame Contribute Delete
73.6 kB
"""多级"从零生成": 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