"""多级"从零生成": 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 指定, 或设环境变量 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