File size: 73,644 Bytes
90ca518 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 | """多级"从零生成": 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
|