diff --git a/mulres_infer/.gitignore b/mulres_infer/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..7ad99397846dbb58dae8b9ff5fb37eabb93e0596 --- /dev/null +++ b/mulres_infer/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.pyc +ckpt/*.pth +ckpt/*.ckpt +outputs/ diff --git a/mulres_infer/ckpt/generator.pth b/mulres_infer/ckpt/generator.pth new file mode 100644 index 0000000000000000000000000000000000000000..4ada7cc8f2d97263e9b461dcd416af91fbc01eab --- /dev/null +++ b/mulres_infer/ckpt/generator.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:edf562427ae7c583d5185890769e24266455346650a4292ab41d00a29eddf6a5 +size 3284020143 diff --git a/mulres_infer/ckpt/readme.md b/mulres_infer/ckpt/readme.md new file mode 100644 index 0000000000000000000000000000000000000000..3c55d8761beafcc7c49a0169199436627155ca55 --- /dev/null +++ b/mulres_infer/ckpt/readme.md @@ -0,0 +1,28 @@ +# 权重 + +推理需要两个文件，放在本目录下： + +| 文件 | 说明 | 大小 | +|---|---|---| +| `generator.pth` | 低分辨率生成器（TSSR 三前向离散扩散），逐级生成到格 64 | 3.1 GB | +| `subdivider.pth` | 高分辨率细分器（acc_uni），64 → 128 → 256 → 512 | 1.2 GB | + +两者都**只含 `model_state_dict`**，已剥离优化器/调度器状态，且 EMA 影子权重已合并进权重本体 +（推理的原始行为就是「加载 model_state_dict 后用 EMA 覆盖」，合并后逐位等价）。 + +想自己从训练产出的 ckpt 转换： + +```bash +python scripts/strip_ckpt.py <训练ckpt.pth> ckpt/generator.pth +# 想要未经 EMA 平滑的原始权重: 加 --no-ema +``` + +## 可选 + +`shapevae-256.ckpt`（Michelangelo 点云编码器）**推理不需要**——点云条件编码器的权重已包含在 +`generator.pth` 内，缺失时 `miche_conditioner` 会先构造未初始化的编码器再被权重覆盖，结果不受影响。 +只有从头训练才需要它。 + +## 路径覆盖 + +环境变量 `RAIL_CKPT` / `SUPERRES_CKPT` / `SHAPEVAE_CKPT`，或 `infer.py --ckpt / --superres-ckpt`。 diff --git a/mulres_infer/ckpt/subdivider.pth b/mulres_infer/ckpt/subdivider.pth new file mode 100644 index 0000000000000000000000000000000000000000..4b2ef2401188c991202befbb806fa816c115cb08 --- /dev/null +++ b/mulres_infer/ckpt/subdivider.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:55feaa1feaf4254f2ba9d94cc2c3def528bbb5408068a8191f636dbcd55e2bd2 +size 1258077735 diff --git a/mulres_infer/configs/generator.yaml b/mulres_infer/configs/generator.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c69a48ba625cec0bb61d8c91641f3d52ec2a5888 --- /dev/null +++ b/mulres_infer/configs/generator.yaml @@ -0,0 +1,151 @@ +# 几何 RoPE vcat 线(2026-08-05): geo 配方不变, 只把粗尺度(3/face,1/face)锚点从"面质心" +# 换成"三顶点各占1/3通道拼接"(rope_mode="geo_vcat", 用户提议) —— 共质心三角形可分、 +# 顶点信息在粗尺度保留。9/face 尺度与 geo 逐位一致。从 geo ckpt_47 剥离 init 热启动。 +# (token 位置 = 3D 锚点[顶点/边中点/面质心] + 面内类型, 序列顺序不参与注意力, 面片置换等变; +# 设计动机: 跨面"同顶点=同相位"取代 z-major 排序邻近这个有损协调代理, 见 model_acc_fix3_altpc.__init__)。 +# [vcat_1024] 放大版(2026-08-08 用户定): 1024宽/16头/[2,4,12]/ffn2816, 骨干624M+编码器~193M≈818M。 +# init = vcat ckpt_18 的 Net2Net 扩展(script/expand_vcat_768_to_1024.py: 旧权重按组拷入/新维小随机/ +# 新谷底块6-11零门控恒等), strict 加载。 +# [首启后] 本目录出现自己的 ckpt 后, 把 init_ckpt_path 置 null 切回 resume_latest(否则重启永远回 init)。 + +project_name: "gobjaverse_fix3_octant_factorized_1_to_1024_geo_vcat_1024_pc1024_sfocal_tssr" +output_folder_base: "./pretrained_model" + +data: + # 2026-09-10 上集群: 与 sfocal config 同步做前缀迁移(内容同源, /root/dataset -> 仓库内 dataset_lists/, + # mesh 实体在 /data/home/from-baidu-h100-jfs/)。仍是 08-21 用户定的 focal 时代旧数据, 未换数据集。 + train_dataset_path: "./dataset_lists/under_20k_concat.txt" + eval_dataset_path: "./dataset_lists/unique_files_glb_under20000face_compress_glb_with_info_val.txt" + num_workers: 12 + +data_processing: + overfit: false + y_up: true + n_discrete_size: 128 + vertex_order: "zyx" + max_len: 600000 # focal 时代口径(60/face, 1万面上限) + min_len: 0 + max_face_num: 10000 # focal 时代口径 + min_face_num: 200 + pc_num: 40960 + total_pc_num: 50000 + windowing: + window_size: -1 + window_stride: 9000 + split_method: "sliding_window" + train_augmentation: + shift_augment: false + scale_augment: true + dec_augment: false + rot_augment: false # focal 时代口径 + sample_augment: true + val_augmentation: + shift_augment: false + scale_augment: false + dec_augment: false + rot_augment: false + sample_augment: false + +tokenizer: + max_depth: 10 + max_resolution: 1024 + use_hourglass: true + max_children: 8 + target_depth: [0,1,2,3,4,5,6] # 用户定(08-19): 砍掉深度7(格256)提速, 训到格128 + +model: + feature_dim: 1024 + num_heads: 16 # head_dim 保持 64(flash+rope 槽位不变) + hourglass_depth: [2, 4, 12] # 谷底 6->12; 新块 6-11 由扩展 init 零门控(起步恒等) + dropout: 0.0 + encoder_freeze: false + conditioned_on_pc: true + max_seq_len: 150000 + # ---- 几何 RoPE ---- + rope_mode: "geo_vcat" # 顶点显式锚(2026-08-05 用户提议): 粗尺度=三顶点各1/3通道拼接; + # 9/face=双端点布局(F/M_ii=(v,v)复制, M_ij=(vi,vj)拆分, 治对角线中点混叠) + rope_geo_scale: 1.0 # 体素单位(不按 depth 归一, 理由见模型注释) + rope_theta_xyz: 100.0 # 邻格 1 rad 可分辨; depth7 全程 128 格最低频 2.0 rad < π + rope_theta_type: 10.0 + rope_type_slots: 2 # 32 槽 = xyz 30 + type 2; vcat 粗尺度切 3顶点x9槽+type2+pad3 + rope_theta_vcat: 300.0 # 3槽/轴的频率梯子(邻格1rad可分辨, 128格最低频2.85rad<π) + ffn_hidden: 2816 # 显式覆盖(8/3*1024=2730 非 256 倍数, 取 2816 对齐 Tensor Core) + pc_num_latents: 1024 # 点云编码器 latent 1+1024(query 平铺扩容, 注入方式不变: 每层 cross-attn 双云交替) + miche_sdpa: true # miche 注意力 SDPA 化(省 ~20GB 瞬时峰, fpsq 线验证过, 数值等价) + +training: + skip_resume: false + # 冷启动: 从 hg246 ckpt_39 的剥离 init(只留 model+ema)起步, epoch 从 0 计。 + # 首启跑出本目录 ckpt 后改: init_ckpt_path: null + init_ckpt_path: null # ckpt_1(0.1894)已落盘(08-23), 切回 resume 语义; 冷启曾用 sfocal/tssr_init_from_ep24.pth + resume_latest: true + load_strict: false # TSSR 新头 mode_emb/to_acc fresh(零初始化尾层, warm-start 友好) + reset_epoch: false # 已切 resume 语义 + + epoches_num: 40 + base_lr: 3.0e-5 # warm-start 大模型 + 新头 + 新目标(uniform/分类器), 小 lr 防骨干漂移; 稳定后可提 + min_lr: 1.0e-6 + warmup_ratio: 0.02 + + cond_drop_prob: 0.1 + pc_target_drop_prob: 1.0 # 恒丢云B => 每层恒定注入同一朵全分辨率云A(用户最终定案: 经典focal式单云) + single_target_pc: false # 已否决(08-21二次修正): 用户要的是全分辨率云A, 不是云B + + nonempty_row_weight: 1.0 # 原始 focal 的正确配置(focal γ 即类平衡, 别叠 w2; 2.0是后来的w2密度fork) + focal_gamma: 2.0 # focal loss(同 sfocal) + # ----- TSSR 三前向超参(与旧 tssr 线一致) ----- + tssr_p_mask: 0.5 # 每 step 走 mask 步的概率(逐 step 跨 rank 同步) + tssr_p_selfgen: 0.1 # no_mask 步用自生成脏输入的概率(治暴露偏差) + tssr_acc_loss_weight: 1.0 # 逐 token 正确性分类器 BCE 权重 + tssr_acc_pos_weight: 1.0 # BCE 中「错」类上权 + tssr_mask_logits_lo: 0.2 # sample_tssr 分类器 re-mask 阈值下限 + tssr_mask_logits_hi: 0.9 # 上限 + + use_ema: true + ema_decay: 0.999 + ema_start_step: 0 + ema_update_every: 1 + + grad_checkpoint: true # 1024x24块@90k闸 激活~2x于768线的69GB悬崖, 必开(激活省~10x/步时+30%) + save_ckpt_every: 1 + save_ckpt_step_num: 3000 # 9h/epoch(9400步)且宿主机不稳: ~每1.8h一个中途档, 防整epoch丢失 + batch_size: 6 + diffusion_train_steps: 500 + optimizer: + type: "AdamW" + weight_decay: 0.01 + betas: [0.9, 0.95] + debug_fixed_mask: false + debug_fixed_t: None + debug_train_on_all_valid_tokens: false + dynamic_token_batch: + enabled: true + tokens_per_face: 9 + budget_mode: "varlen" + # 2026-07-30 批点重定(零丢样档, 用户选定): 校准 ratio 后 est≈1.31x 实际, 预算 95k -> 批实际 + # 位置中位 ~58k / P90 ~76k; 闸 72k 只兜长尾(全量模拟: 触闸批 2%, 丢 0.3% 样本/0.6% 位置)。 + # 旧点 60k+旧ratio+55k 闸实际在静默丢 12% 样本/22% 位置(偏丢深 depth 大 mesh); + # tssr 式 340k 会丢 64% 位置 —— 预算永远要和 hard_max_positions 对齐着调, 别只动一个。 + # 2026-07-30 三调(用户定 90k): 闸 72k 实测触发 ~14% 批(全在 72-82k 区间, 丢 ~1.4% 样本, + # 实际长度绕估计 ±20% 噪声推批过阈)。90k 下预算 95k(est≈1.31x实际→实际中位 ~72k)几乎不可能触闸, + # 只剩单样本巨 mesh 长尾。实测 bs14+闸72k 峰值 64.9GB, +18k 位置 ≈ +4-5GB -> ~69-70GB, 安全。 + max_tokens_per_batch: 200000 # 用户定(2026-08-12): 容纳2万面样本@depth7(est~164k), 大mesh获得全深度覆盖 + max_batch_size: 20 + hard_max_positions: 250000 # 用户定: 预算-闸余量25%, 兜深样本组合的估计低估长尾 + bucket_size: 2048 + length_safety_factor: 1.05 + +accelerator: + gradient_accumulation_steps: 1 + mixed_precision: "bf16" + +evaluation: + generate_only_after_training: false + generate_val_mesh_every: 3 # 大 val mesh 逐个采样很慢, 降频防 epoch 末长尾(NCCL 超时已提1h兜底) + diffusion_sample_steps: 20 + pc_num: 1024 + cfg_scale: 1.5 + cfg_schedule: "linear" + cfg_start: 1.0 + best_of_n: 0 + sample_temperature: 0.0 diff --git a/mulres_infer/configs/subdivider.yaml b/mulres_infer/configs/subdivider.yaml new file mode 100644 index 0000000000000000000000000000000000000000..edd9fdf6f3f9d7f3bdee24e25adb1b8d4afa175c --- /dev/null +++ b/mulres_infer/configs/subdivider.yaml @@ -0,0 +1,126 @@ +project_name: "gobjaverse_5k_diffusion_pc_fix_tokenizer_1_to_1024_hourglass_fix_pc_more_data_10k_acc_uni_fix_tokenizer" +output_folder_base: "./pretrained_model" + +data: + train_dataset_path: "/root/dataset/under_20k_concat.txt" + eval_dataset_path: "/root/dataset/unique_files_glb_under20000face_compress_glb_with_info_val.txt" + num_workers: 12 + +data_processing: + overfit: false + y_up: true + n_discrete_size: 128 + vertex_order: "zyx" + max_len: 270000 + min_len: 0 + max_face_num: 13000 + min_face_num: 200 + pc_num: 40960 + total_pc_num: 50000 + windowing: + window_size: -1 + window_stride: 9000 + split_method: "sliding_window" + train_augmentation: + shift_augment: false + scale_augment: true + dec_augment: false + rot_augment: false + sample_augment: true + val_augmentation: + shift_augment: false + scale_augment: false + dec_augment: false + rot_augment: false + sample_augment: false + +tokenizer: + max_depth: 10 + max_resolution: 1024 + use_hourglass: true + max_children: 3 + target_depth: [0,1,2,3,4,5,6,7,8,9,10] + +model: + codebook_path: "/root/mul-res/checkpoints/octree_codebook_meshgpt_shapenet_d5_6.npz" + use_hourglass: true + depth: 12 + model_version: "v1" + feature_dim: 768 + num_heads: 12 + root_pred_depth: 3 + context_embedding_dim: 256 + num_classify: 100 + hourglass_depth: [2, 2, 4] + input_embe_pos: "no" + root_pred_method: "index" + conditioned_on_pc: true + encoder_freeze: false + dropout: 0.0 + +training: + # ---- fine-tune knobs ---- + # checkpoint_394_..._loss.pth 是 pre-CFG (496 参数, 无 null_cond / EMA, 用它做 + # init_ckpt_path 会触发 optimizer param-group size mismatch). + # 真正的 refine ckpt 是下面这个 (497 参数, 有 null_cond + ema_state_dict). + init_ckpt_path: "/root/junkai_lin/mul_res_diffusion/pretrained_model/gobjaverse_5k_diffusion_pc_fix_tokenizer_1_to_1024_hourglass_fix_pc_more_data_10k_acc_uni/checkpoint_19_epoch_0.1954_loss.pth" + load_strict: true # refine ckpt 已含 null_cond, 严格加载更安全 + reset_epoch: true # epoch 从 0 重新计数 + skip_resume: false + + # ---- LR / 时长 ---- + epoches_num: 20 # 由于 trainer 循环 off-by-one (range(1, epoches_num)), 实际只跑 5 epoch + base_lr: 5.0e-5 # 比原训练 (1e-4) 低 10x, 不毁掉已收敛的权重 + min_lr: 1.0e-6 + warmup_ratio: 0.02 + + # ---- CFG cond-drop ---- + cond_drop_prob: 0.1 # 10% batch 把 PC 置空, 学习 null_cond + + # ---- EMA (开启 + 快 decay) ---- + use_ema: true + ema_decay: 0.99 # fine-tune 步数少, 必须用快 decay + ema_start_step: 0 # 从第一步就开始累积 + ema_update_every: 1 + + # ---- 保留的训练参数 ---- + save_ckpt_every: 1 + save_ckpt_step_num: 10000 + batch_size: 3 + diffusion_train_steps: 500 + optimizer: + type: "AdamW" + weight_decay: 0.01 + betas: [0.9, 0.95] + loss_weights: + root_pred_weight: 0.2 + root_method: "root_move_index" + root_embed: false + debug_fixed_mask: false + debug_fixed_batch: false + debug_fixed_t: None + debug_train_on_all_valid_tokens: false + dynamic_token_batch: + enabled: true + tokens_per_face: 9 + budget_mode: "varlen" + max_tokens_per_batch: 117000 + max_batch_size: 6 + bucket_size: 2048 + length_safety_factor: 1.05 + +accelerator: + gradient_accumulation_steps: 4 + mixed_precision: "bf16" + +evaluation: + generate_only_after_training: false + debug_root: false + generate_val_mesh_every: 1 + completion_input_len: 18 + pc_num: 1024 + use_toppk: true + top_k: 8 + top_p: 0.8 + temperature: 1 + diffusion_sample_steps: 40 diff --git a/mulres_infer/dataset/__init__.py b/mulres_infer/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/mulres_infer/dataset/mesh_dataset.py b/mulres_infer/dataset/mesh_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..32acd66ac3a2c3950244a82bd36a085ea61f6168 --- /dev/null +++ b/mulres_infer/dataset/mesh_dataset.py @@ -0,0 +1,940 @@ +import trimesh +import networkx as nx +import numpy as np +import os, pickle +import json +import math +import threading +import torch +from torch.utils.data import Dataset +import glob +from torch.utils.data import Sampler +from tqdm import tqdm +import numpy as np +from utils.data_utils import undiscretize, discretize, normalize, clean_mesh +import logging +from concurrent.futures import ProcessPoolExecutor, as_completed +import random +import multiprocessing as mp +from pathlib import Path +import open3d +from scipy.spatial import cKDTree +import concurrent.futures +from PIL import Image + + +# ---------------------------------------------------------------------------- +# nvdiffrast on-the-fly normal-map renderer. +# +# Vendored from /root/junkai_lin/TRELLIS/dataset_toolkits/render_normals_gpu.py +# so the dataset stays self-contained. When `image_cond.render_on_the_fly` is +# enabled, we render the AUGMENTED mesh directly in the dataset's world frame +# (Z-up, normalized to [-1, 1]) so the projection math in +# `diffusion/img_cond_conditioner.py` lines up without needing the +# `undo_dataset_y_up` swap or the 0.5 pc_render_scale that the TRELLIS +# pre-renders required. Set those config flags accordingly when switching to +# render-on-the-fly mode (see img_cond yaml). +# +# Render failures (e.g. degenerate mesh, CUDA OOM, zero-face mesh) raise; the +# `except Exception` loop in __getitem__ then resamples a different mesh. +# ---------------------------------------------------------------------------- + +_NVDR_CTX = None +_NVDR_LOCK = threading.Lock() + + +def _get_nvdr_ctx(): + """Lazily create one nvdiffrast CUDA context per process. + + DataLoader workers run as separate processes; each one initializes its own + context the first time __getitem__ touches the renderer. The default + `fork` start method may conflict with CUDA already initialized in the + parent — if you hit `cuda: initialization error` here, switch the + DataLoader to `multiprocessing_context='spawn'`. + """ + global _NVDR_CTX + if _NVDR_CTX is not None: + return _NVDR_CTX + with _NVDR_LOCK: + if _NVDR_CTX is None: + import nvdiffrast.torch as dr + if not torch.cuda.is_available(): + raise RuntimeError( + "image_cond.render_on_the_fly=true requires CUDA; no GPU visible to this worker" + ) + device = torch.device("cuda", 0) + _NVDR_CTX = {"glctx": dr.RasterizeCudaContext(device=device), "device": device} + return _NVDR_CTX + + +def _build_camera_for_render(yaw, pitch, radius, fov, aspect=1.0, near=0.01, far=100.0): + """Z-up world, origin-tracking camera. yaw/pitch/fov in radians. + + Returns (view[4,4], proj[4,4], cam_pos[3], R_c2w[3,3]) as float64 numpy. + Matches `render_normals_gpu.build_camera` exactly so the pose convention + is consistent with the TRELLIS pre-render path. + """ + cx = radius * math.cos(yaw) * math.cos(pitch) + cy = radius * math.sin(yaw) * math.cos(pitch) + cz = radius * math.sin(pitch) + cam_pos = np.array([cx, cy, cz], dtype=np.float64) + + forward = -cam_pos / np.linalg.norm(cam_pos) + world_up = np.array([0.0, 0.0, 1.0]) + right = np.cross(forward, world_up) + rn = np.linalg.norm(right) + if rn < 1e-9: + world_up = np.array([0.0, 1.0, 0.0]) + right = np.cross(forward, world_up) + rn = np.linalg.norm(right) + right /= rn + up = np.cross(right, forward) + R_c2w = np.stack([right, up, -forward], axis=1) + R_w2c = R_c2w.T + t_w2c = -R_w2c @ cam_pos + + view = np.eye(4, dtype=np.float64) + view[:3, :3] = R_w2c + view[:3, 3] = t_w2c + + f = 1.0 / math.tan(fov / 2.0) + proj = np.zeros((4, 4), dtype=np.float64) + proj[0, 0] = f / aspect + # nvdiffrast image-space y is flipped vs OpenGL; negate so row 0 = top. + proj[1, 1] = -f + proj[2, 2] = -(far + near) / (far - near) + proj[2, 3] = -2.0 * far * near / (far - near) + proj[3, 2] = -1.0 + return view, proj, cam_pos, R_c2w + + +def _explode_to_flat_normals(verts, faces): + """Per-face geometric normals via mesh explosion (flat shading).""" + F = faces.shape[0] + v0 = verts[faces[:, 0]] + v1 = verts[faces[:, 1]] + v2 = verts[faces[:, 2]] + fn = np.cross(v1 - v0, v2 - v0) + fn_len = np.linalg.norm(fn, axis=1, keepdims=True) + fn = fn / np.where(fn_len < 1e-12, 1.0, fn_len) + new_verts = np.empty((F * 3, 3), dtype=verts.dtype) + new_verts[0::3] = v0 + new_verts[1::3] = v1 + new_verts[2::3] = v2 + new_normals = np.repeat(fn, 3, axis=0).astype(verts.dtype) + new_faces = np.arange(F * 3, dtype=np.int32).reshape(F, 3) + return new_verts, new_faces, new_normals + + +def _fit_radius_to_silhouette(verts_t, yaw, pitch, fov, init_radius, target_fill=0.9, iters=3): + """Shrink/grow camera radius so the mesh's max |NDC| equals target_fill.""" + radius = float(init_radius) + ones = torch.ones((verts_t.shape[0], 1), device=verts_t.device, dtype=verts_t.dtype) + v_h = torch.cat([verts_t, ones], dim=1) + for _ in range(iters): + view, proj, _, _ = _build_camera_for_render(yaw, pitch, radius, fov) + mvp = (proj @ view).astype(np.float32) + mvp_t = torch.from_numpy(mvp).to(verts_t.device) + v_clip = v_h @ mvp_t.T + front = v_clip[:, 3] > 1e-3 + if not front.any(): + break + c = v_clip[front] + ndc_xy = c[:, :2] / c[:, 3:4].clamp_min(1e-6) + max_ndc = float(ndc_xy.abs().max().item()) + if max_ndc < 1e-6: + break + radius *= max_ndc / target_fill + radius = max(0.5, min(radius, 50.0)) + return radius + + +def _render_normals_image( + verts_np, + faces_np, + yaw_deg=315.0, + pitch_deg=30.0, + fov_deg=40.0, + init_radius=4.0, + resolution=518, + auto_radius=True, + target_fill=0.9, +): + """Render a normal-visualization image of `(verts, faces)` in the mesh's + own world frame (assumed Z-up). + + Returns + ------- + image_tensor : torch.float16, [3, H, W], values in [0, 1] (white background) + transform_matrix : torch.float32, [4, 4] (camera-to-world, NeRF convention) + camera_angle_x : torch.float32, scalar (horizontal FOV in radians) + """ + import nvdiffrast.torch as dr + + if faces_np is None or faces_np.shape[0] == 0: + raise ValueError("render: mesh has zero faces") + + ctx = _get_nvdr_ctx() + device = ctx["device"] + glctx = ctx["glctx"] + + yaw = math.radians(yaw_deg) + pitch = math.radians(pitch_deg) + fov = math.radians(fov_deg) + + verts_np = np.ascontiguousarray(verts_np, dtype=np.float32) + faces_np = np.ascontiguousarray(faces_np, dtype=np.int32) + verts_np, faces_np, normals_np = _explode_to_flat_normals(verts_np, faces_np) + + verts_t = torch.from_numpy(verts_np).to(device) + faces_t = torch.from_numpy(faces_np).to(device).to(torch.int32) + normals_t = torch.from_numpy(normals_np).to(device) + ones = torch.ones((verts_t.shape[0], 1), device=device, dtype=verts_t.dtype) + verts_h = torch.cat([verts_t, ones], dim=1) + + radius = ( + _fit_radius_to_silhouette(verts_t, yaw, pitch, fov, init_radius, target_fill) + if auto_radius + else float(init_radius) + ) + + view_M, proj_M, cam_pos, R_c2w = _build_camera_for_render(yaw, pitch, radius, fov) + mvp = (proj_M @ view_M).astype(np.float32) + mvp_t = torch.from_numpy(mvp).to(device) + + clip = (verts_h @ mvp_t.T).unsqueeze(0).contiguous() # (1, V, 4) + rast, _ = dr.rasterize(glctx, clip, faces_t, resolution=(resolution, resolution)) + alpha = (rast[0, ..., 3] > 0).float() # (H, W) + if float(alpha.sum().item()) < 4.0: + raise ValueError("render: empty silhouette (mesh out of frustum?)") + + normal_img, _ = dr.interpolate(normals_t.unsqueeze(0), rast, faces_t) + normal_img = normal_img[0] # (H, W, 3) + nlen = torch.linalg.norm(normal_img, dim=-1, keepdim=True).clamp_min(1e-8) + normal_img = normal_img / nlen + normal_img = normal_img * alpha.unsqueeze(-1) + + # Visualize normals as RGB: (N * 0.5 + 0.5), white background. + vis = (normal_img * 0.5 + 0.5).clamp(0.0, 1.0) + a = alpha.unsqueeze(-1) + vis = vis * a + (1.0 - a) # background = white + image_tensor = vis.permute(2, 0, 1).contiguous().to(dtype=torch.float16).cpu() + + transform_matrix = torch.eye(4, dtype=torch.float32) + transform_matrix[:3, :3] = torch.from_numpy(R_c2w).float() + transform_matrix[:3, 3] = torch.from_numpy(cam_pos).float() + camera_angle_x = torch.tensor(float(fov), dtype=torch.float32) + return image_tensor, transform_matrix, camera_angle_x + + +def _read_image_cond_cfg(opt): + """Return (enabled, render_root, image_resolution, render_on_the_fly, + render_cfg, render_manifest_path). + + `render_on_the_fly=True` enables the nvdiffrast path. In that mode + `render_root` is unused (no disk-render filter), but it may still be set + in the yaml for backward compatibility. + + `render_manifest_path` is the path to a precomputed list of valid render + subdir stems (one per line). When set and the file exists, the dataset + skips the slow per-line `os.path.exists` and uses O(1) set membership. + Build it with `script/build_render_manifest.py`. + """ + cfg = getattr(opt, "image_cond", None) + if cfg is None or not getattr(cfg, "enabled", False): + return False, None, None, False, None, None + image_resolution = int(getattr(cfg, "image_resolution", 512)) + render_on_the_fly = bool(getattr(cfg, "render_on_the_fly", False)) + render_root = getattr(cfg, "render_root", None) + if not render_on_the_fly and render_root is None: + raise ValueError( + "image_cond.enabled=true but image_cond.render_root is not set " + "(and render_on_the_fly is false)" + ) + render_cfg = getattr(cfg, "render", None) + render_manifest_path = getattr(cfg, "render_manifest_path", None) + return True, render_root, image_resolution, render_on_the_fly, render_cfg, render_manifest_path + + +def _render_dir_for(file_path: str, render_root: str) -> str: + return os.path.join(render_root, Path(file_path).stem) + + +def _has_render(file_path: str, render_root: str) -> bool: + d = _render_dir_for(file_path, render_root) + return os.path.exists(os.path.join(d, "transforms.json")) and os.path.exists( + os.path.join(d, "000.png") + ) + + +# Per-path LRU cache so the manifest is read once per process. +_RENDER_MANIFEST_CACHE: dict = {} + + +def _load_render_manifest(manifest_path): + """Load a manifest file into a set of valid render stems. + + Returns None when `manifest_path` is None/empty or the file doesn't + exist, so the caller falls back to per-line `os.path.exists`. Build the + manifest with `script/build_render_manifest.py` whenever the render + contents change — it drops startup time from O(N stat) to O(N lookup). + """ + if not manifest_path: + return None + if manifest_path in _RENDER_MANIFEST_CACHE: + return _RENDER_MANIFEST_CACHE[manifest_path] + if not os.path.isfile(manifest_path): + _RENDER_MANIFEST_CACHE[manifest_path] = None + return None + with open(manifest_path, "r") as f: + stems = {line.strip() for line in f if line.strip()} + _RENDER_MANIFEST_CACHE[manifest_path] = stems + return stems + + +# TRELLIS Blender depth encoding (see dataset_toolkits/blender_script/render.py): +# pixel_val (uint16) = ((depth_world - depth.min) / (depth.max - depth.min)) * 65535 +# Background ends up clamped to 65535 (== depth.max output). We treat 65535 as +# "no surface" so points landing on BG pixels are considered occluded (not +# visible) in that view. +_DEPTH_BG_VALUE = 65535 + + +def _load_depth_map(depth_png_path: str, depth_min: float, depth_max: float, + target_resolution: int): + """Decode TRELLIS depth PNG to world-space depth. + + Returns (depth[H,W] float32, valid_mask[H,W] bool). + `valid_mask` is False on background pixels (no surface hit). + """ + arr = np.array(Image.open(depth_png_path)) + if arr.dtype != np.uint16: + # 8-bit fallback — still treat the brightest value as background. + max_val = np.iinfo(arr.dtype).max + bg_val = max_val + norm = arr.astype(np.float32) / float(max_val) + else: + bg_val = _DEPTH_BG_VALUE + norm = arr.astype(np.float32) / 65535.0 + valid = arr != bg_val + depth = depth_min + norm * (depth_max - depth_min) + if (arr.shape[0], arr.shape[1]) != (target_resolution, target_resolution): + # Resize via PIL (nearest) so depth values aren't blended across the + # silhouette — bilinear would corrupt the foreground/background edge. + depth_img = Image.fromarray(depth.astype(np.float32), mode="F").resize( + (target_resolution, target_resolution), Image.NEAREST + ) + valid_img = Image.fromarray(valid.astype(np.uint8) * 255, mode="L").resize( + (target_resolution, target_resolution), Image.NEAREST + ) + depth = np.asarray(depth_img, dtype=np.float32) + valid = np.asarray(valid_img, dtype=np.uint8) > 0 + return depth, valid + + +def _load_image_and_camera(file_path: str, render_root: str, image_resolution: int): + """Load BOTH rendered views + depth maps for back-projection. + + Returns: + image_tensor : [V, 3, H, W] float16 in [0, 1] (white bg) + transform_matrix : [V, 4, 4] float32 (camera-to-world) + camera_angle_x : [V] float32 (horizontal FOV, rad) + depth_map : [V, H, W] float32 (world-space depth) + depth_valid : [V, H, W] bool (False on background) + + Iterates every frame listed in transforms.json (the dataset stores two + views: 000 and 001). Missing depth PNGs raise — the caller will catch and + resample a different mesh. + """ + d = _render_dir_for(file_path, render_root) + with open(os.path.join(d, "transforms.json"), "r") as f: + meta = json.load(f) + + images = [] + transforms = [] + fovs = [] + depths = [] + valids = [] + for frame in meta["frames"]: + rgb_path = os.path.join(d, frame["file_path"]) + img = Image.open(rgb_path) + if img.mode == "RGBA": + bg = Image.new("RGB", img.size, (255, 255, 255)) + bg.paste(img, mask=img.split()[-1]) + img = bg + else: + img = img.convert("RGB") + if img.size != (image_resolution, image_resolution): + img = img.resize((image_resolution, image_resolution), Image.LANCZOS) + arr = np.asarray(img, dtype=np.float32) / 255.0 + images.append(torch.from_numpy(arr).permute(2, 0, 1).to(dtype=torch.float16)) + + transforms.append( + torch.tensor(frame["transform_matrix"], dtype=torch.float32) + ) + fovs.append(torch.tensor(float(frame["camera_angle_x"]), dtype=torch.float32)) + + depth_info = frame.get("depth") + if depth_info is None: + raise ValueError( + f"transforms.json frame {frame.get('file_path')} has no 'depth' entry; " + "back-projection requires depth bounds." + ) + stem = os.path.splitext(frame["file_path"])[0] + depth_path = os.path.join(d, f"{stem}_depth.png") + depth_np, valid_np = _load_depth_map( + depth_path, + float(depth_info["min"]), + float(depth_info["max"]), + image_resolution, + ) + depths.append(torch.from_numpy(depth_np)) + valids.append(torch.from_numpy(valid_np)) + + image_tensor = torch.stack(images, dim=0) # [V, 3, H, W] + transform_matrix = torch.stack(transforms, dim=0) # [V, 4, 4] + camera_angle_x = torch.stack(fovs, dim=0) # [V] + depth_map = torch.stack(depths, dim=0) # [V, H, W] + depth_valid = torch.stack(valids, dim=0) # [V, H, W] + return image_tensor, transform_matrix, camera_angle_x, depth_map, depth_valid +def shift_vertices(vertices, x_lims=(-0.1, 0.1), y_lims=(-0.1, 0.1), z_lims=(-0.075, 0.075)): + # shift x, y, z + x = np.random.uniform(low=x_lims[0], high=x_lims[1], size=(1,)) + y = np.random.uniform(low=y_lims[0], high=y_lims[1], size=(1,)) + z = np.random.uniform(low=z_lims[0], high=z_lims[1], size=(1,)) + x = max(min(x, 0.5 - vertices[:, 0].max()), -0.5 - vertices[:, 0].min()) + y = max(min(y, 0.5 - vertices[:, 1].max()), -0.5 - vertices[:, 1].min()) + z = max(min(z, 0.5 - vertices[:, 2].max()), -0.5 - vertices[:, 2].min()) + vertices = np.stack([vertices[:, 0] + x, vertices[:, 1] + y, vertices[:, 2] + z], axis=-1) + return vertices + +def scale_vertices(vertices, x_lims=(0.75, 1.25), y_lims=(0.75, 1.25), z_lims=(0.75, 1.25)): + # scale x, y, z + x = np.random.uniform(low=x_lims[0], high=x_lims[1], size=(1,)) + y = np.random.uniform(low=y_lims[0], high=y_lims[1], size=(1,)) + z = np.random.uniform(low=z_lims[0], high=z_lims[1], size=(1,)) + vertices = np.stack([vertices[:, 0] * x, vertices[:, 1] * y, vertices[:, 2] * z], axis=-1) + return vertices + +def normalize_vertices(vertices, eps: float = 1e-8): + bounds = np.array([vertices.min(axis=0), vertices.max(axis=0)]) # type: ignore + vertices = vertices - (bounds[0] + bounds[1])[None, :] / 2 + # 数据中存在退化 mesh (bounds 范围为 0), 分母加 eps 避免除零产生 inf/NaN + scale = (bounds[1] - bounds[0]).max() + eps + vertices = vertices / scale + return vertices + +def agu_mesh(vertices: np.ndarray, shift_augment, scale_augment): + # 随机轴缩放 + if scale_augment: + # lower, upper = scale_factor + # scale_axis = lower + np.random.rand(vertices.shape[-1]) * (upper - lower) + # vertices = vertices * scale_axis + vertices = scale_vertices(vertices) + vertices = normalize_vertices(vertices) + if shift_augment: + vertices = shift_vertices(vertices) + return vertices + + + +def dec(mesh, target_face): + open3d_mesh = open3d.geometry.TriangleMesh( + vertices=open3d.utility.Vector3dVector(mesh.vertices), + triangles=open3d.utility.Vector3iVector(mesh.faces)) + simple = open3d_mesh.simplify_quadric_decimation(target_face) + return trimesh.Trimesh(vertices=simple.vertices, faces=simple.triangles) + +def normalize(mesh, eps: float = 1e-8): + bounds = np.array([mesh.vertices.min(axis=0), mesh.vertices.max(axis=0)]) + mesh.vertices = mesh.vertices - (bounds[0] + bounds[1])[None, :] / 2 + # 数据中存在退化 mesh (bounds 范围为 0), 分母加 eps 避免除零产生 inf/NaN + scale = (bounds[1] - bounds[0]).max() + eps + mesh.vertices = 2 * (mesh.vertices / scale) + mesh.vertices = mesh.vertices.clip(-1.0, 1.0) # 限制顶点坐标在[-1.0, 1.0] + +def merge_duplicate_vertices(mesh, tolerance=0): + """更可靠的顶点合并方法""" + vertices = mesh.vertices + faces = mesh.faces + + # 使用KD树查找相近的顶点 + tree = cKDTree(vertices) + + # 查找在容差范围内的顶点对 + pairs = tree.query_pairs(tolerance, output_type='ndarray') + + if len(pairs) == 0: + return mesh # 没有需要合并的顶点 + + # 创建顶点映射：将重复顶点映射到第一个顶点 + vertex_map = {i: i for i in range(len(vertices))} + + for i, j in pairs: + # 总是映射到较小的索引 + if i < j: + vertex_map[j] = i + else: + vertex_map[i] = j + + # 应用顶点映射 + new_faces = [] + for face in faces: + new_face = [vertex_map[v] for v in face] + new_faces.append(new_face) + + # 创建新的网格 + new_mesh = trimesh.Trimesh(vertices=vertices, faces=new_faces) + + # 移除未引用的顶点 + new_mesh.remove_unreferenced_vertices() + + return new_mesh + +def discrete_undiscrete_and_clean(mesh, n_discrete, clean=True): + vertices = discretize(mesh.vertices, num_discrete=n_discrete) + vertices = undiscretize(vertices, num_discrete=n_discrete) + dis_mesh = trimesh.Trimesh(vertices=vertices, faces=mesh.faces) + new_mesh = merge_duplicate_vertices(dis_mesh) + new_mesh.merge_vertices() + new_mesh.update_faces(new_mesh.nondegenerate_faces()) + new_mesh.update_faces(new_mesh.unique_faces()) + new_mesh.remove_unreferenced_vertices() + new_mesh.fix_normals() + return new_mesh + +# 点云加噪开关(环境变量,默认关=旧加噪;DataLoader worker 会继承 env)。 +# PAPER_PC_NOISE=1 开启论文式加噪: 位置 σ_pos "up to" PC_NOISE_POS(逐样本采强度) + 法向 σ_normal=PC_NOISE_NORMAL +# (加噪后重归一化)。默认关 -> 走旧加噪(仅位置 +N(0,0.01)、50% 概率, 法向不加噪),其它训练零影响。 +# 幅度也可用 env PC_NOISE_POS / PC_NOISE_NORMAL 覆盖(默认 0.1 / 0.2)。只在 aug=True(训练)时生效,推理不加噪。 +def _paper_pc_noise_cfg(): + return (os.environ.get("PAPER_PC_NOISE", "0") == "1", + float(os.environ.get("PC_NOISE_POS", "0.1")), + float(os.environ.get("PC_NOISE_NORMAL", "0.2"))) + + +def sample_pc(mesh, pc_num, total_pc_num = 50000, with_normal=True, aug=True, overfit=False): + if overfit: + np.random.seed(42) + if not with_normal: + points, _ = mesh.sample(pc_num, return_index=True) + return points + + points, face_idx = mesh.sample(total_pc_num, return_index=True) + normals = mesh.face_normals[face_idx] + paper_noise, sigma_pos_max, sigma_normal = _paper_pc_noise_cfg() + if paper_noise: + # 论文式加噪: 位置 σ_pos<=sigma_pos_max + 法向 σ_normal(重归一化为单位向量) + if aug: + points = points.astype(np.float32) + normals = normals.astype(np.float32) + sigma_pos = np.random.uniform(0.0, sigma_pos_max) # "up to" -> 逐样本采噪声强度 + points = points + np.random.randn(*points.shape).astype(np.float32) * sigma_pos + normals = normals + np.random.randn(*normals.shape).astype(np.float32) * sigma_normal + nlen = np.linalg.norm(normals, axis=-1, keepdims=True) + normals = normals / np.clip(nlen, 1e-8, None) # 加噪后法向重归一化为单位向量 + else: + # 旧加噪(默认): 仅位置 +N(0,0.01)、50% 概率, 法向不加噪(保持既有训练行为不变) + if aug and random.random() < 0.5: + points = points + np.random.randn(*points.shape) * 0.01 + pc_normal = np.concatenate([points, normals], axis=-1, dtype=np.float16) + + # random sample point cloud + ind = np.random.choice(pc_normal.shape[0], pc_num, replace=False) + pc_normal = pc_normal[ind] + + return pc_normal + +def mul_encode_mesh(raw_mesh, gt_mesh, codec, tokenizer_target_depth, pc_num, total_pc_num, use_gt_pc, clean=True, + normalize_input=True): + mesh = raw_mesh + + # 归一化。normalize_input=False: 跳过对 raw_mesh 的归一化 —— 级联推理用(2026-08-22 修复): + # 级间 mesh 是上一级 tokenizer 解码产物, 天生在规范 [-1,1] 格坐标系; 按它自身 bbox 再归一化会在 + # 预测带极值杂散面时把脚手架帧拉离 GT/点云帧(z 轴撕裂 bug, 见 infer_multilevel CASCADE_RENORM)。 + # 训练路径 raw_mesh==gt_mesh(磁盘任意坐标系), 保持默认 True 不变。 + if normalize_input: + normalize(mesh) + normalize(gt_mesh) + + if isinstance(tokenizer_target_depth, list): + selected_depth = random.choice(tokenizer_target_depth) + else: + selected_depth = tokenizer_target_depth + raw_dis_mesh = discrete_undiscrete_and_clean(gt_mesh, n_discrete=2**selected_depth, clean=clean) + gt_dis_mesh = discrete_undiscrete_and_clean(gt_mesh, n_discrete=2**(selected_depth+1), clean=clean) + # 采样点云 + if use_gt_pc: + pc = sample_pc(gt_mesh, pc_num=pc_num, total_pc_num=total_pc_num, + with_normal=True, aug=False, overfit=False) + else: + pc = sample_pc(mesh, pc_num=pc_num, total_pc_num=total_pc_num, + with_normal=True, aug=False, overfit=False) + # raw_token, token = self.codec.encode_mesh(mesh, target_depth=self.tokenizer_target_depth) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + + future = executor.submit(codec.encode_mesh, mesh, target_depth=selected_depth) + # 如果超过 2.0 秒没算完，这里会直接抛出 concurrent.futures.TimeoutError + raw_token, token, _ = future.result(timeout=20.0) + + n = len(token) + + data_dict = {} + data_dict["token"] = torch.tensor(token, dtype=torch.int64) + data_dict["raw_token"] = torch.tensor(raw_token, dtype=torch.int64) + data_dict["raw_vertices"] = mesh.vertices + data_dict["raw_faces"] = mesh.faces + data_dict["raw_dis_mesh"] = raw_dis_mesh + data_dict["gt_dis_mesh"] = gt_dis_mesh + + data_dict["start"] = 0 + data_dict["end"] = len(data_dict["token"]) + data_dict["pc"] = torch.from_numpy(pc) + data_dict["token_len"] = torch.tensor([len(data_dict["token"])], dtype=torch.int64) + + data_dict["selected_depth"] = selected_depth + return data_dict + +class MeshDataset(Dataset): + """ + 重构后的版本，初始化参数通过 opt 对象传入 + """ + def __init__(self, opt, data, codec, all_token_len, is_train): + # 将配置对象和核心数据存储到实例中 + self.opt = opt + self.data = data + self.codec = codec + self.all_token_len = all_token_len + self.is_train = is_train + # 从 opt 对象中解包参数，保持原有代码逻辑不变 + # 这样做的好处是 __getitem__ 等其他方法不需要修改 + if is_train: + self.shift_augment = self.opt.data_processing.train_augmentation.shift_augment + self.scale_augment = self.opt.data_processing.train_augmentation.scale_augment + self.dec_augment = self.opt.data_processing.train_augmentation.dec_augment + self.rot_augment = self.opt.data_processing.train_augmentation.rot_augment + self.sample_augment = self.opt.data_processing.train_augmentation.sample_augment + else: + self.shift_augment = self.opt.data_processing.val_augmentation.shift_augment + self.scale_augment = self.opt.data_processing.val_augmentation.scale_augment + self.dec_augment = self.opt.data_processing.val_augmentation.dec_augment + self.rot_augment = self.opt.data_processing.val_augmentation.rot_augment + self.sample_augment = self.opt.data_processing.val_augmentation.sample_augment + self.n_discrete_size = self.opt.data_processing.n_discrete_size + self.vertex_order = self.opt.data_processing.vertex_order + self.window_size = self.opt.data_processing.windowing.window_size + self.window_stride = self.opt.data_processing.windowing.window_stride + self.split_method = self.opt.data_processing.windowing.split_method + self.max_length = self.opt.data_processing.max_len-20 # 假设max_length也在opt中 + self.min_length = self.opt.data_processing.min_len # 假设max_length也在opt中 + self.pc_num = self.opt.data_processing.pc_num + self.total_pc_num = self.opt.data_processing.total_pc_num + self.y_up = self.opt.data_processing.y_up + self.overfit = self.opt.data_processing.overfit if hasattr(self.opt.data_processing, 'overfit') else False + self.tokenizer_target_depth = self.opt.tokenizer.target_depth + # 2a 暴露偏差训练腐蚀(仅训练): 粗输入加伪面(过连接)+节点坐标抖动(位移),目标查 GT 不变。 + # 仅 FastMeshCodec.encode_mesh 实现; 推理/eval 传 None -> 与原版逐位一致。 + self.corrupt_cfg = None + _ct = getattr(getattr(self.opt, "training", None), "corrupt", None) + if is_train and _ct is not None and bool(getattr(_ct, "enabled", False)): + self.corrupt_cfg = { + "spurious_rate": float(getattr(_ct, "spurious_rate", 0.0)), + "spatial_radius": int(getattr(_ct, "spatial_radius", 2)), + "jitter_prob": float(getattr(_ct, "jitter_prob", 0.0)), + "jitter_max": int(getattr(_ct, "jitter_max", 1)), + } + print(f" [2a Corrupt] ENABLED (is_train): {self.corrupt_cfg}") + ( + self.image_cond_enabled, + self.image_cond_render_root, + self.image_cond_resolution, + self.image_cond_render_on_the_fly, + self.image_cond_render_cfg, + self.image_cond_render_manifest_path, + ) = _read_image_cond_cfg(opt) + # Default render parameters; can be overridden by `image_cond.render` in yaml. + # Defaults match render_normals_gpu.sh (yaw=315, pitch=30, FOV=40, auto_radius) + # but with init_radius=4.0 because our mesh frame is [-1, 1] (extent ~2) vs. + # the TRELLIS pre-render frame [-0.5, 0.5] (extent ~1). + rcfg = self.image_cond_render_cfg + self.render_yaw_deg = float(getattr(rcfg, "yaw_deg", 315.0)) if rcfg is not None else 315.0 + self.render_pitch_deg = float(getattr(rcfg, "pitch_deg", 30.0)) if rcfg is not None else 30.0 + self.render_fov_deg = float(getattr(rcfg, "fov_deg", 40.0)) if rcfg is not None else 40.0 + self.render_init_radius = float(getattr(rcfg, "init_radius", 4.0)) if rcfg is not None else 4.0 + self.render_auto_radius = bool(getattr(rcfg, "auto_radius", True)) if rcfg is not None else True + self.render_target_fill = float(getattr(rcfg, "target_fill", 0.9)) if rcfg is not None else 0.9 + # If true (and is_train), randomly sample yaw in [0, 360) per call. + self.render_randomize_yaw = bool(getattr(rcfg, "randomize_yaw", False)) if rcfg is not None else False + print("\n" + "="*70) + print("="*70) + # 使用 f-string 格式化输出，将相关的配置项打印在同一行 + print(f" [Augmentation] Shift: {self.shift_augment}, Scale: {self.scale_augment}, Decimation: {self.dec_augment}, Rotation: {self.rot_augment}, Sample: {self.sample_augment}") + print(f" [Data Specs] Discretization Size: {self.n_discrete_size}, Vertex Order: '{self.vertex_order}'") + print(f" [Windowing] Size: {self.window_size}, Stride: {self.window_stride}, Method: '{self.split_method}'") + print(f" [Sequence Len] Min: {self.min_length}, Max: {self.max_length}") + print(f" [Point Cloud] PC Num: {self.pc_num}, Total PC Num: {self.total_pc_num}") + print(f" [Tokenizer] Target Depth: {self.tokenizer_target_depth}") + if self.image_cond_enabled: + if self.image_cond_render_on_the_fly: + print( + f" [Image Cond] Enabled | mode=render_on_the_fly | res={self.image_cond_resolution} " + f"| yaw={self.render_yaw_deg} pitch={self.render_pitch_deg} fov={self.render_fov_deg} " + f"radius_init={self.render_init_radius} auto_radius={self.render_auto_radius} " + f"randomize_yaw={self.render_randomize_yaw}" + ) + else: + print(f" [Image Cond] Enabled | mode=disk | render_root={self.image_cond_render_root} | res={self.image_cond_resolution}") + print("="*70 + "\n") + print(f"[MeshDataset] Created from {len(self.data)} entries") + + def __len__(self): + return len(self.data) + + def get_all_token_lenth(self): + return self.all_token_len + + # __getitem__ 方法保持不变 + def __getitem__(self, item): + iter_cnt = 0 + if isinstance(item, tuple): + idx, start, end = item + else: + idx = item + start = 0 + end = -1 + while True: + file_path = self.data[idx] + iter_cnt += 1 + if iter_cnt >= 4: + print(f'iter_cnt>=4: {iter_cnt}') + try: + mesh = trimesh.load(file_path, force="mesh") + # y轴朝上->z轴朝上 + if self.y_up: + mesh.vertices = mesh.vertices[:, [2, 0, 1]] + + # 随机塌缩 + if self.dec_augment: + if mesh.faces.shape[0] >= self.min_length and random.random() < 0.5: + target = np.random.randint(max(self.min_length, mesh.faces.shape[0] // 4), mesh.faces.shape[0]) + mesh = dec(mesh=mesh, target_face=target) + + # 随机拉伸和平移 + if self.scale_augment or self.shift_augment: + mesh.vertices = agu_mesh(mesh.vertices, self.shift_augment, self.scale_augment) + + # 随机旋转 + if self.rot_augment: + # azimuth = np.random.uniform(0, 360) + azimuth = np.random.choice([0, 30, 60, 90, 120, 150, 180, 210, 240, 270, 300, 330], 1)[0] + theta = np.radians(-azimuth) + rotz = np.array([ + [np.cos(theta), -np.sin(theta), 0], + [np.sin(theta), np.cos(theta), 0], + [0, 0, 1] + ]) + mesh.vertices = mesh.vertices @ rotz.T # 绕z轴旋转 + + # 归一化 + normalize(mesh) + + if isinstance(self.tokenizer_target_depth, list): + selected_depth = random.choice(self.tokenizer_target_depth) + else: + selected_depth = self.tokenizer_target_depth + + # discrete_mesh = discrete_undiscrete_and_clean(mesh, n_discrete=2**selected_depth) + + # # 采样点云 + # pc = sample_pc(discrete_mesh, pc_num=self.pc_num, total_pc_num=self.total_pc_num, + # with_normal=True, aug=self.sample_augment, overfit=self.overfit) + # 采样点云 + pc = sample_pc(mesh, pc_num=self.pc_num, total_pc_num=self.total_pc_num, + with_normal=True, aug=self.sample_augment, overfit=self.overfit) + if not np.isfinite(pc).all(): + raise ValueError("NaN/Inf in sampled point cloud") + # raw_token, token = self.codec.encode_mesh(mesh, target_depth=self.tokenizer_target_depth) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + + future = executor.submit(self.codec.encode_mesh, mesh, target_depth=selected_depth, + corrupt=self.corrupt_cfg) + # 超时会抛 concurrent.futures.TimeoutError(str() 为空)-> 换成带 depth/面数的明确消息 + try: + raw_token, token, dis_mesh = future.result(timeout=10.0) + except concurrent.futures.TimeoutError: + raise TimeoutError(f"encode_mesh 超时(>10s, depth={selected_depth} res{2**(selected_depth+1)}, 输入面={len(mesh.faces)}) —— 多半粗分辨率过连接") + + n = len(token) + if n < self.min_length: + raise ValueError(f"Token length {n} is less than min_length {self.min_length}") + + if n > self.max_length: + raise ValueError(f"Token length {n} exceeds max_length {self.max_length}") + + # Render the (augmented) mesh on-the-fly. Done inside the try + # block so any render failure resamples a different mesh. + rendered = None + if self.image_cond_enabled: + if self.image_cond_render_on_the_fly: + yaw_deg = ( + float(np.random.uniform(0.0, 360.0)) + if (self.render_randomize_yaw and self.is_train) + else self.render_yaw_deg + ) + rendered = _render_normals_image( + np.asarray(mesh.vertices, dtype=np.float32), + np.asarray(mesh.faces, dtype=np.int32), + yaw_deg=yaw_deg, + pitch_deg=self.render_pitch_deg, + fov_deg=self.render_fov_deg, + init_radius=self.render_init_radius, + resolution=self.image_cond_resolution, + auto_radius=self.render_auto_radius, + target_fill=self.render_target_fill, + ) + else: + rendered = _load_image_and_camera( + file_path, + self.image_cond_render_root, + self.image_cond_resolution, + ) + + data_dict = {} + data_dict["token"] = torch.tensor(token, dtype=torch.int64) + data_dict["raw_token"] = torch.tensor(raw_token, dtype=torch.int64) + data_dict["raw_vertices"] = mesh.vertices + data_dict["raw_faces"] = mesh.faces + data_dict["dis_mesh"] = dis_mesh + del mesh + break + except Exception as e: + _reason = str(e) if str(e) else f"{type(e).__name__}(空消息)" + print(f'[dataiter wrong]: {file_path}, {_reason}') + idx = np.random.randint(0, len(self.data)) + # start = 0 + # end = -1 + + if len(data_dict["token"]) < self.window_size: + data_dict["start"] = 0 + data_dict["end"] = len(data_dict["token"]) + else: + if start == 0: + data_dict["start"] = 0 + data_dict["end"] = self.window_size + # 随机截断 + elif start == -1 and end == -1: + data_dict["start"] = -1 + data_dict["end"] = -1 + else: + data_dict["start"] = len(data_dict["token"]) - self.window_size + data_dict["end"] = len(data_dict["token"]) + data_dict["pc"] = torch.from_numpy(pc) + data_dict["token_len"] = torch.tensor([len(data_dict["token"])], dtype=torch.int64) + data_dict['name'] = Path(file_path).stem + data_dict["file_path"] = file_path + data_dict["selected_depth"] = selected_depth + if self.image_cond_enabled: + if self.image_cond_render_on_the_fly: + # Render-on-the-fly path still returns the legacy single-view tuple + # (no depth). Add a leading view-dim so downstream code is uniform. + image_tensor, transform_matrix, camera_angle_x = rendered + data_dict["image"] = image_tensor.unsqueeze(0) + data_dict["transform_matrix"] = transform_matrix.unsqueeze(0) + data_dict["camera_angle_x"] = camera_angle_x.unsqueeze(0) + # No depth available — emit dummies; the conditioner will skip + # occlusion testing when depth_valid is all False. + H = image_tensor.shape[-1] + data_dict["depth_map"] = torch.zeros(1, H, H, dtype=torch.float32) + data_dict["depth_valid"] = torch.zeros(1, H, H, dtype=torch.bool) + else: + (image_tensor, transform_matrix, camera_angle_x, + depth_map, depth_valid) = rendered + data_dict["image"] = image_tensor + data_dict["transform_matrix"] = transform_matrix + data_dict["camera_angle_x"] = camera_angle_x + data_dict["depth_map"] = depth_map + data_dict["depth_valid"] = depth_valid + return data_dict + + @classmethod + def load(cls, opt, path, codec, accelerator, is_train): + """ + 读取数据并使用 opt 对象来初始化 Dataset 实例。 + """ + all_token_len = [] + data = [] + max_length = opt.data_processing.max_face_num + min_len = opt.data_processing.min_face_num + (image_cond_enabled, render_root, _, render_on_the_fly, + _, render_manifest_path) = _read_image_cond_cfg(opt) + # Disk-render availability filter only applies when we read PNGs from + # disk; in render_on_the_fly mode the renderer makes its own image. + apply_disk_render_filter = image_cond_enabled and not render_on_the_fly + # Prefer the manifest when the disk-render filter is on. Path comes + # from `image_cond.render_manifest_path` in yaml; if absent or the + # file doesn't exist, fall back to per-line os.path.exists (slow). + render_manifest = None + if apply_disk_render_filter: + render_manifest = _load_render_manifest(render_manifest_path) + if accelerator.is_main_process: + if render_manifest is not None: + print(f"[MeshDataset] Using render manifest ({len(render_manifest)} stems) " + f"from {render_manifest_path}") + else: + if render_manifest_path: + print(f"[MeshDataset] WARNING: image_cond.render_manifest_path={render_manifest_path} " + f"not found — falling back to per-line os.path.exists (slow). " + f"Run script/build_render_manifest.py to fix.") + else: + print(f"[MeshDataset] WARNING: image_cond.render_manifest_path not set — " + f"falling back to per-line os.path.exists (slow). " + f"Run script/build_render_manifest.py and set the yaml key.") + dropped_no_render = 0 + with open(path, 'r') as f: + for line in tqdm(f): + # 这段逻辑可能需要根据你的文件名格式调整 + try: + token_len = int(line.split("_f")[-1].split("_")[0]) + except (IndexError, ValueError): + try: + token_len = int(line.split("_facenum_")[-1].split("_")[0]) + except (ValueError, IndexError): + try: + token_len = int(line.split("_facenum_")[-1].split(".")[0]) + except (ValueError, IndexError): + # 如果文件名不包含token长度 + mesh = trimesh.load(line.strip(), force="mesh") + token_len = mesh.faces.shape[0] + if token_len <= max_length and token_len>= min_len: + file_path = line.strip() + if apply_disk_render_filter: + if render_manifest is not None: + if Path(file_path).stem not in render_manifest: + dropped_no_render += 1 + continue + elif not _has_render(file_path, render_root): + dropped_no_render += 1 + continue + data.append(file_path) + all_token_len.append(token_len) + + if accelerator.is_main_process: + print(f"[MeshDataset] Loaded {len(data)} entries from {path}") + if apply_disk_render_filter: + print( + f"[MeshDataset] Image-cond filter dropped {dropped_no_render} entries with no render under {render_root}" + ) + if len(data) == 0: + print( + f"[MeshDataset] WARNING: no usable entries after image-cond filter — render root may be wrong." + ) + elif image_cond_enabled and render_on_the_fly: + print(f"[MeshDataset] Image-cond mode=render_on_the_fly — no disk-render filter applied.") + + # 使用 opt 对象来创建类的实例 + return cls(opt, + data=data, + codec=codec, + all_token_len=all_token_len, + is_train=is_train) + diff --git a/mulres_infer/dataset/mesh_dataset_altpc.py b/mulres_infer/dataset/mesh_dataset_altpc.py new file mode 100644 index 0000000000000000000000000000000000000000..29da88b29f4dada8a56f835713e809ce32ce5922 --- /dev/null +++ b/mulres_infer/dataset/mesh_dataset_altpc.py @@ -0,0 +1,94 @@ +"""altpc 变体:pc_target 改用「编码→解码的 mesh」(octree token 解码的 GT@目标分辨率), +而非 discrete_undiscrete 塌缩。理由:解码 mesh 与模型实际输出同一表示(含退化面结构),条件更一致。 +PCT_DEGEN=1 时, 额外在退化面(零面积, 即栏杆/薄结构边)上采点并注入 pc_target(法向置0)。 +直接 decode 父类已存的 d["token"](无需重编码)。 +""" +import os +import numpy as np +import trimesh +import torch + +from dataset.mesh_dataset import MeshDataset, sample_pc + +PCT_DEGEN = os.environ.get("PCT_DEGEN", "0") == "1" # 是否在退化面上也采样注入 +_EPS = 1e-7 + + +def rail_edges(mesh): + """退化面(零面积)的有效栏杆边:取每个退化面最长边(对 (A,B,B),B-B 长 0,最长边即真实 A-B)。 + 返回 (A[D,3], B[D,3], 边长[D])。无退化面则 D=0。""" + af = np.asarray(mesh.area_faces) + deg = np.where(af < _EPS)[0] + if len(deg) == 0: + return np.zeros((0, 3)), np.zeros((0, 3)), np.zeros((0,)) + v = np.asarray(mesh.vertices); f = np.asarray(mesh.faces)[deg] + ev = v[f] + pairs = [(0, 1), (1, 2), (0, 2)] + lens = np.stack([np.linalg.norm(ev[:, i] - ev[:, j], axis=1) for i, j in pairs], axis=1) # [D,3] + pick = lens.argmax(1) + A = ev[np.arange(len(deg)), [pairs[p][0] for p in pick]] + B = ev[np.arange(len(deg)), [pairs[p][1] for p in pick]] + return A, B, lens[np.arange(len(deg)), pick] + + +def sample_rail(A, B, lengths, n): + """沿栏杆边按边长加权采 n 个点(长边多采)。返回 [n,3]。""" + if n <= 0 or len(lengths) == 0 or lengths.sum() <= 0: + return np.zeros((0, 3)) + counts = np.random.multinomial(n, lengths / lengths.sum()) # 每条边分到的点数, 和=n + idx = np.repeat(np.arange(len(lengths)), counts) # [n] 每点所属边 + t = np.random.rand(len(idx), 1).astype(np.float32) + return A[idx] + t * (B[idx] - A[idx]) + + +def degen_point_budget(mesh, pc_num, resolution): + """按「一致面密度」分配表面/栏杆点数, 不固定比例。 + 退化边当 1 体素宽薄条: 等效面积 = 边长 × 体素宽(mesh尺度/分辨率)。 + ρ = pc_num/(表面积+栏杆等效面积); n_rail = round(ρ·栏杆等效面积)。 + 高分辨率体素小→栏杆相对少; 栏杆多的 mesh→栏杆点自然多。返回 (n_surf, n_rail, A, B, lengths)。""" + af = np.asarray(mesh.area_faces) + surf_area = float(af[af >= _EPS].sum()) + A, B, lengths = rail_edges(mesh) + rail_len = float(lengths.sum()) + if rail_len <= 0 or surf_area <= 0: + return pc_num, 0, A, B, lengths + voxel = float(np.max(mesh.extents)) / max(1, resolution) # 1 体素的物理尺度 + rail_eff_area = rail_len * voxel + rho = pc_num / (surf_area + rail_eff_area) + n_rail = int(round(rho * rail_eff_area)) + n_rail = max(0, min(n_rail, pc_num)) + return pc_num - n_rail, n_rail, A, B, lengths + + +def sample_pc_with_degen(mesh, pc_num, total_pc_num, aug, overfit, resolution=64): + """表面采样 + 退化面边上采样, 按一致密度自适应分配(不固定比例)。栏杆点法向置 0。总数 = pc_num。""" + n_surf, n_rail, A, B, lengths = degen_point_budget(mesh, pc_num, resolution) + surf = sample_pc(mesh, pc_num=pc_num, total_pc_num=total_pc_num, with_normal=True, aug=aug, overfit=overfit) + if n_rail <= 0: + return surf + rail_xyz = sample_rail(A, B, lengths, n_rail) + if len(rail_xyz) == 0: + return surf + rail = np.concatenate([rail_xyz, np.zeros_like(rail_xyz)], axis=1).astype(surf.dtype) # 法向 0 + return np.concatenate([surf[:pc_num - len(rail)], rail], axis=0) + + +class MeshDatasetAltPCDecoded(MeshDataset): + def __getitem__(self, item): + d = super().__getitem__(item) + sd = int(d["selected_depth"]) + # 直接 decode 父类已存的 token -> octree 解码 GT@目标分辨率(与模型输出同表示, 含退化面) + try: + decoded = self.codec.decode_tokens(d["token"].tolist(), target_depth=sd) + except Exception: + decoded = None + if decoded is None or len(getattr(decoded, "faces", [])) < 4: + decoded = trimesh.Trimesh(vertices=d["raw_vertices"], faces=d["raw_faces"], process=False) + if PCT_DEGEN: + pc_t = sample_pc_with_degen(decoded, self.pc_num, self.total_pc_num, self.sample_augment, + self.overfit, resolution=2 ** (sd + 1)) + else: + pc_t = sample_pc(decoded, pc_num=self.pc_num, total_pc_num=self.total_pc_num, + with_normal=True, aug=self.sample_augment, overfit=self.overfit) + d["pc_target"] = torch.from_numpy(pc_t) + return d diff --git a/mulres_infer/models/__init__.py b/mulres_infer/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/mulres_infer/models/generator.py b/mulres_infer/models/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..5e7eeb3e5f6eca4dd317127c77a80e5fa8b95cf9 --- /dev/null +++ b/mulres_infer/models/generator.py @@ -0,0 +1,1507 @@ +"""fix3 生成器模型(已融合: 基类 + alt-PC 双云注入 + TSSR 三前向)。 + +类继承链(融合后仍保留, 因各层职责清晰): + FactorizedDiffusionTransformer 骨干: 因子分解 token / AdaLN cross-attn / 自回归行头 + └ FactorizedDiffusionTransformerAltPC 逐层交替双点云注入 + 几何 RoPE + sample() + └ FactorizedTSSRTransformer mode flag + 逐 token 正确性分类器 + sample_tssr() + +原先的 FactorizedDiffusionTransformerFocal 一层只重写训练用的 factorized_loss, 推理路径上是空壳, +已移除并让 AltPC 直接继承骨干。各类的纯训练方法(factorized_loss / factorized_accuracy / +sample_log / forward_mask / forward_nomask / zero_touch_conditional_params)也已剥离。 +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +from models.subdivider import ( + SinusoidalPosEmb, + ModernCrossAttnTransformerBlock, + ModernAttention, + RMSNorm, + SwiGLU, + RotaryEmbedding, + LinearDownsample, + LinearUpsample, + apply_rotary_pos_emb, + apply_rotary_pos_emb_varlen, + lengths_to_position_ids, + cosine_schedule, +) +from models.point_conditioner import PointConditioner +import os +import numpy as np +from models.subdivider import ( + apply_rotary_pos_emb_varlen, lengths_to_position_ids, cosine_schedule, +) + + +# ========================================================================== +# 原 model_acc_fix3.py +# ========================================================================== + +def modulate(x, shift, scale): + """adaLN 调制: x * (1 + scale) + shift。shift/scale 与 x 同形状(逐 token 广播)。""" + return x * (1 + scale) + shift +class AdaLNCrossAttnBlock(nn.Module): + """fix3 专用 block：在 self-attn 与 FFN 上加 adaLN 调制(条件 c = time + depth)。 + cross-attn 仍走点云条件(不调制)。结构与 ModernCrossAttnTransformerBlock 同名(self_attn/ + cross_attn/norm1-3/ffn)以便迁移 fix2 权重；adaLN_modulation 为**新增**且做「恒等初始化」 + (shift=0, scale=0, gate=1)，使加载迁移权重时 block == 旧 block，backbone 不被破坏。""" + + def __init__(self, d_model, nhead, dropout=0., qk_norm=False, ffn_hidden=None): + super().__init__() + # QK-norm 只加在 self-attn(instability 主源);cross-attn 的 K 是点云条件,保持默认不动 + self.self_attn = ModernAttention(d_model, nhead, dropout=dropout, qk_norm=qk_norm) + self.cross_attn = ModernAttention(d_model, nhead, dropout=dropout) + self.norm1 = RMSNorm(d_model) + self.norm2 = RMSNorm(d_model) + self.norm3 = RMSNorm(d_model) + self.ffn = SwiGLU(d_model, ffn_hidden) + self.dropout = nn.Dropout(dropout) + # 从条件 c 产生 self-attn / FFN 各自的 (shift, scale, gate)，共 6*d_model + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(d_model, 6 * d_model)) + # 恒等初始化：weight=0；bias 使 shift=0, scale=0, gate=1 -> block 退化为旧 block + nn.init.zeros_(self.adaLN_modulation[1].weight) + with torch.no_grad(): + b = torch.zeros(6 * d_model) + b[2 * d_model:3 * d_model] = 1.0 # gate_sa = 1 + b[5 * d_model:6 * d_model] = 1.0 # gate_ff = 1 + self.adaLN_modulation[1].bias.copy_(b) + + def forward_varlen(self, x, c, lengths, cond=None, cond_lengths=None, rope_apply_fn=None, + max_len=None, cond_max_len=None): + """c: [total, d_model] 逐 token 条件(time+depth)。 + max_len / cond_max_len: 透传给 attention 的 max_seqlen(python int), 省掉 GPU->CPU 同步; + 不传则由 attention 内部回退到 `.item()`(见 ModernAttention.forward_varlen 的说明)。 + + 注: 曾试过「c 只传 [b,d]、在块内 adaLN 后再 repeat_interleave」以省掉 768->6*768 Linear + 在同一行上重复算 total 次的冗余(每块 424 GFLOP@60k token) —— 实测输出逐位相同但**零加速** + (hourglass 已把 token 降到 1/3、1/9, 那 3.1 TFLOP 不是瓶颈), 且会让 torch.compile 报 + dtype 不匹配。故回退, 别再试。 + """ + # c 来自 time_mlp + depth_embedding, 是 fp32(nn.Embedding 不在 autocast 转型名单里)。 + # 对齐到 x 的 dtype: eager+autocast 下等价(Linear 本来就会转), 但 torch.compile 编译这个方法时 + # 会丢掉外层 autocast 上下文, 不显式对齐就会在后面的 matmul 上报 "Float and BFloat16"。 + c = c.to(x.dtype) + shift_sa, scale_sa, gate_sa, shift_ff, scale_ff, gate_ff = self.adaLN_modulation(c).chunk(6, dim=-1) + h = modulate(self.norm1(x), shift_sa, scale_sa) + sa = self.self_attn.forward_varlen(h, h, h, lengths, rope_apply_fn=rope_apply_fn, + max_len_q=max_len) + x = x + gate_sa * self.dropout(sa) + if cond is not None: + h = self.norm2(x) + # CFG drop 的步里 cond 来自 fp32 的 null_cond; 同 c 一样显式对齐(见上), 否则 compile 会崩 + ca = self.cross_attn.forward_varlen(h, cond.to(x.dtype), cond.to(x.dtype), lengths, + lengths_k=cond_lengths, + max_len_q=max_len, max_len_k=cond_max_len) + x = x + self.dropout(ca) + h = modulate(self.norm3(x), shift_ff, scale_ff) + ff = self.ffn(h) + x = x + gate_ff * self.dropout(ff) + return x +SUB_VOCAB = 256 # 每个子码(8bit)的取值数 +MASK_ID = SUB_VOCAB # 子码 MASK id = 256（embedding 表大小 = 257） +NUM_OCTANTS = 8 +TYPE_F, TYPE_INTRA, TYPE_INTER = 0, 1, 2 +REFINE_POS_PER_FACE = 9 # 3 F + 3 intra + 3 inter +class AutoregressiveRowHead(nn.Module): + """连接矩阵(intra/inter)的 8 行**自回归** MLP 头(参考 TreeMeshGPT 的 head_coord1/2/3)。 + + 取代原先「8 行并行独立预测」：行 r 的 256-way logits 由 + cat([hidden, emb(行0), ..., emb(行r-1)]) -> 3 层 MLP + 给出，从而**显式建模行间依赖**，让模型只产出能共现的合法矩阵(对称/可成 clique)， + 减少独立预测导致的「悬边→退化面/碎块」。 + * 训练: teacher forcing(回喂 GT 行 embedding)，一次 forward 出全 8 行 logits。 + * 采样: 逐行 argmax rollout(回喂已采行)；支持 CFG —— cond/uncond 在每行 logits 级合并， + 合并后采样得到的同一行回喂两路，保持自回归状态一致。 + 行 embedding 维度默认 = d_model//8(与输入端 row_embed_dim 一致，控参数量)；最后一行不作输入故只需 7 张表。 + """ + + def __init__(self, d_model, num_rows=NUM_OCTANTS, vocab=SUB_VOCAB, row_emb_dim=None): + super().__init__() + self.num_rows = num_rows + self.vocab = vocab + self.row_emb_dim = row_emb_dim or (d_model // num_rows) + self.mlps = nn.ModuleList([ + nn.Sequential( + nn.Linear(d_model + r * self.row_emb_dim, d_model), nn.ReLU(), + nn.Linear(d_model, d_model), nn.ReLU(), + nn.Linear(d_model, vocab), + ) for r in range(num_rows) + ]) + self.row_emb = nn.ModuleList([nn.Embedding(vocab, self.row_emb_dim) for _ in range(num_rows - 1)]) + + def forward_teacher(self, h, gt_rows): + """h [n,d]；gt_rows [n,num_rows] long(取值[0,vocab)) -> logits [n,num_rows,vocab]。teacher forcing。""" + feats = [h] + logits = [] + for r in range(self.num_rows): + logits.append(self.mlps[r](torch.cat(feats, dim=-1))) + if r < self.num_rows - 1: + feats.append(self.row_emb[r](gt_rows[:, r]).to(h.dtype)) + return torch.stack(logits, dim=1) + + def rollout(self, h, h_u=None, cfg_scale=1.0, temperature=0.0): + """逐行采样。h [n,d]；h_u 提供则做 CFG(logits=uncond+w*(cond-uncond))。 + temperature<=0: argmax 贪心(确定性)；>0: 按 (logit/temperature) 多项式采样(用于 best-of-N 产生多样候选)。 + 返回 (pred [n,num_rows] long, conf [n,num_rows] 每行所选类的概率)。""" + feats = [h] + feats_u = [h_u] if h_u is not None else None + preds, confs = [], [] + for r in range(self.num_rows): + logit = self.mlps[r](torch.cat(feats, dim=-1)) + if feats_u is not None: + lu = self.mlps[r](torch.cat(feats_u, dim=-1)) + logit = lu + cfg_scale * (logit - lu) + if temperature and temperature > 0: + prob = (logit / temperature).softmax(-1) + am = torch.multinomial(prob.float(), 1).squeeze(-1) # [n] + pm = prob.gather(-1, am.unsqueeze(-1)).squeeze(-1) # 所选类概率 + else: + prob = logit.softmax(-1) + pm, am = prob.max(-1) + preds.append(am); confs.append(pm) + if r < self.num_rows - 1: + e = self.row_emb[r](am).to(h.dtype) + feats.append(e) + if feats_u is not None: + feats_u.append(e) + return torch.stack(preds, dim=1), torch.stack(confs, dim=1) +class FactorizedDiffusionTransformer(nn.Module): + def __init__(self, + coord_vocab_size, # raw_embedding 大小 = resolution+1 (coord 含 offset, 0=pad) + max_seq_len, + d_model=768, + nhead=12, + depth=(4, 4, 8), + dropout=0., + encoder_freeze=True, + encoder_name='miche-256-feature', + use_pc=True, + time_embed_scale=100.0, + num_depths=16, # depth(层级) embedding 大小，覆盖 target_depth 0..max + nonempty_row_weight=1.0, # 连接矩阵「非空行」(目标!=0, 即真实连接) 的 loss 上权 + qk_norm=False, # self-attn 的 Q/K 加 per-head RMSNorm(稳定 attention);默认关 + ffn_hidden=None, # FFN 隐层维度显式覆盖(None=8/3*d_model; 1024宽用2816对齐TensorCore) + pc_num_latents=None, # 点云编码器 latent 数(None=预训练默认256; 1024=条件带宽x4, query平铺扩容) + time_embed_dim=None): # 正弦基底/time_mlp 隐层维度(None=d_model, 默认即可)。备注: + # SinusoidalPosEmb 的"列-频率/sin-cos"布局随 dim 变化, 宽度扩展 + # 热启动时旧 time_mlp 权重不可拷贝(语义错位会爆), 见扩展脚本。 + super().__init__() + self.d_model = d_model + self.mask_token_id = MASK_ID + self.num_depths = num_depths + # 连接(intra/inter)行子码中「非空行」(target!=0) 的 loss 权重。connectivity 矩阵 ~92% 是空行， + # 等权 CE 会让模型坍缩到「全预测空行」(逐子码 acc 虚高、生成 mesh 欠连接碎裂)。>1 上权真实连接行。 + self.nonempty_row_weight = float(nonempty_row_weight) + + # 1. Embeddings ---------------------------------------------------- + self.raw_embedding = nn.Embedding(coord_vocab_size, d_model) # coord(条件) + self.f_sub_embedding = nn.Embedding(SUB_VOCAB + 1, d_model) # F: 单子码 + # intra/inter: 8 个行子码各查表(+1 行给 MASK)，再 cat 后 Linear 投影到 d_model。 + # 用 cat 而非 sum：保留 8 行的「行序/位置语义」(行 r = octant r)，避免 sum 的置换不变与叠加丢失； + # 行 embedding 维度 = d_model // 8，cat 后恰好 = d_model，再过一层 Linear 学行间组合(含 intra 对称)。 + self.row_embed_dim = d_model // NUM_OCTANTS + self.intra_sub_embedding = nn.Embedding(SUB_VOCAB + 1, self.row_embed_dim) + self.inter_sub_embedding = nn.Embedding(SUB_VOCAB + 1, self.row_embed_dim) + self.intra_in_proj = nn.Linear(NUM_OCTANTS * self.row_embed_dim, d_model) + self.inter_in_proj = nn.Linear(NUM_OCTANTS * self.row_embed_dim, d_model) + + # 2. 条件: 时间步 + 层级(depth) —— 一起作为 adaLN 的条件 c ------------- + te = int(time_embed_dim) if time_embed_dim else d_model + self.time_mlp = nn.Sequential( + SinusoidalPosEmb(te, scale=time_embed_scale), + nn.Linear(te, te * 4), + nn.GELU(), + nn.Linear(te * 4, d_model), + ) + self.depth_embedding = nn.Embedding(num_depths, d_model) # 层级 embedding(随机初始化) + + # 3. 主干: 带 adaLN 的 block(条件 c = time + depth)；blocks 内 self_attn/cross_attn/ + # norm/ffn 与 ModernCrossAttnTransformerBlock 同名可迁移，adaLN_modulation 为新增(恒等初始化) + self.depth = list(depth) + _blk = lambda: AdaLNCrossAttnBlock(d_model, nhead, dropout=dropout, qk_norm=qk_norm, ffn_hidden=ffn_hidden) + self.enc_blocks_9 = nn.ModuleList([_blk() for _ in range(depth[0])]) + self.enc_blocks_3 = nn.ModuleList([_blk() for _ in range(depth[1])]) + self.bottleneck_blocks = nn.ModuleList([_blk() for _ in range(depth[2])]) + self.dec_blocks_3 = nn.ModuleList([_blk() for _ in range(depth[1])]) + self.dec_blocks_9 = nn.ModuleList([_blk() for _ in range(depth[0])]) + self.down_9_to_3 = LinearDownsample(d_model, shorten_factor=3) + self.down_3_to_1 = LinearDownsample(d_model, shorten_factor=3) + self.up_1_to_3 = LinearUpsample(d_model, shorten_factor=3) + self.up_3_to_9 = LinearUpsample(d_model, shorten_factor=3) + self.rotary_emb = RotaryEmbedding(dim=d_model // nhead, max_position_embeddings=max_seq_len) + + # 4. 因子分解输出头 ------------------------------------------------ + self.f_head = nn.Linear(d_model, SUB_VOCAB) # F: 1 个子码(单头) + # intra/inter: 8 行**自回归** MLP 头(行间依赖, 取代旧的 8 行并行独立 Linear) + self.intra_head = AutoregressiveRowHead(d_model, num_rows=NUM_OCTANTS, vocab=SUB_VOCAB) + self.inter_head = AutoregressiveRowHead(d_model, num_rows=NUM_OCTANTS, vocab=SUB_VOCAB) + + # 5. 点云条件 ------------------------------------------------------ + self.use_pc = use_pc + if use_pc: + print(f'[fix3] Point cloud encoder: {encoder_name} | freeze: {encoder_freeze}') + self.conditioner = PointConditioner(model_name=encoder_name, freeze=encoder_freeze, feature_dim=d_model, + num_latents=pc_num_latents) + self.null_cond_len = 257 + self.null_cond = nn.Parameter(torch.randn(self.null_cond_len, d_model) * 0.02) + + # ------------------------------------------------------------------ # + # 工具 + # ------------------------------------------------------------------ # + def types_from_lengths(self, lengths): + """由每个样本长度(=9*nface)推出 packed 的位置类型 [total]，0/1/2 = F/intra/inter。""" + parts = [] + for L in lengths.tolist(): + idx = torch.arange(L, device=lengths.device) + fp = idx % REFINE_POS_PER_FACE + t = torch.where(fp < 3, TYPE_F, + torch.where(fp < 6, TYPE_INTRA, TYPE_INTER)) + parts.append(t) + return torch.cat(parts).long() if parts else torch.empty(0, dtype=torch.long, device=lengths.device) + + def embed_refine(self, subcodes, types): + """subcodes [N,8] long, types [N] long -> [N,d]。 + F 位置: 单子码查表。intra/inter 位置: 8 行各查表 -> cat -> Linear(保留行序，避免 sum 丢失)。""" + N = subcodes.size(0) + out = self.f_sub_embedding.weight.new_zeros((N, self.d_model)) + is_f = types == TYPE_F + is_in = types == TYPE_INTRA + is_ix = types == TYPE_INTER + # autocast(bf16) 下 Linear 输出是 bf16，而 out 是 fp32；原地 index_put 要求 dtype 一致， + # 故统一 .to(out.dtype)（保持 refine_emb 与 raw_emb 一致的 fp32，backbone matmul 再由 autocast 转 bf16）。 + if is_f.any(): + out[is_f] = self.f_sub_embedding(subcodes[is_f, 0]).to(out.dtype) + if is_in.any(): + e = self.intra_sub_embedding(subcodes[is_in]) # [n, 8, row_embed_dim] + out[is_in] = self.intra_in_proj(e.reshape(e.size(0), -1)).to(out.dtype) # cat -> Linear -> [n, d] + if is_ix.any(): + e = self.inter_sub_embedding(subcodes[is_ix]) + out[is_ix] = self.inter_in_proj(e.reshape(e.size(0), -1)).to(out.dtype) + return out + + # ------------------------------------------------------------------ # + # 前向(varlen)：返回主干 hidden h [total, d] 与 packed types [total] + # ------------------------------------------------------------------ # + def forward_varlen(self, coord_list, subcode_list, t, depth, pc=None): + device = t.device + coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list] + subcode_list = [s.to(device=device, dtype=torch.long) for s in subcode_list] + lengths = torch.tensor([c.numel() for c in coord_list], device=device, dtype=torch.long) + if int(lengths.sum().item()) == 0: + raise ValueError("forward_varlen received an empty batch.") + if torch.any(lengths % REFINE_POS_PER_FACE != 0): + raise ValueError("每个序列长度必须是 9 的倍数(每 face 9 个 refine 位置)。") + + coord_packed = torch.cat(coord_list, dim=0) + subcode_packed = torch.cat(subcode_list, dim=0) + types = self.types_from_lengths(lengths) + + raw_emb = self.raw_embedding(coord_packed) + refine_emb = self.embed_refine(subcode_packed, types) + x = raw_emb + refine_emb # 不再加法注入 time；time/depth 走 adaLN + + # adaLN 条件 c = time_emb + depth_emb（逐样本 [b,d]）；各 hourglass 阶段按其长度展开到逐 token + depth = depth.to(device=device, dtype=torch.long) + c_sample = self.time_mlp(t) + self.depth_embedding(depth) # [b, d_model] + + def c_for(cur_lengths): + return c_sample.repeat_interleave(cur_lengths, dim=0) + + if self.use_pc: + if pc is None: + bsz = len(coord_list) + null = self.null_cond + cond = null.unsqueeze(0).expand(bsz, -1, -1).reshape(-1, null.size(-1)) + cond_lengths = torch.full((bsz,), self.null_cond_len, device=device, dtype=torch.long) + else: + pc = pc.to(device=device) + cond = self.conditioner(pc) + cond_lengths = torch.full((cond.size(0),), cond.size(1), device=device, dtype=torch.long) + cond = cond.reshape(-1, cond.size(-1)) + else: + cond = None + cond_lengths = None + + def rope_fn_for(cur_lengths): + position_ids = lengths_to_position_ids(cur_lengths) + cos, sin = self.rotary_emb.forward_positions(position_ids, dtype=x.dtype) + return lambda q, k: apply_rotary_pos_emb_varlen(q, k, cos, sin) + + rope_fn_9 = rope_fn_for(lengths) + c9 = c_for(lengths) + for block in self.enc_blocks_9: + x = block.forward_varlen(x, c9, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_9) + skip_9 = x + lengths_9 = lengths + + x, lengths = self.down_9_to_3.forward_varlen(x, lengths) + rope_fn_3 = rope_fn_for(lengths) + c3 = c_for(lengths) + for block in self.enc_blocks_3: + x = block.forward_varlen(x, c3, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_3) + skip_3 = x + lengths_3 = lengths + + x, lengths = self.down_3_to_1.forward_varlen(x, lengths) + rope_fn_1 = rope_fn_for(lengths) + c1 = c_for(lengths) + for block in self.bottleneck_blocks: + x = block.forward_varlen(x, c1, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_1) + + x, lengths = self.up_1_to_3.forward_varlen(x, lengths) + x = x + skip_3 + for block in self.dec_blocks_3: + x = block.forward_varlen(x, c3, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_3) + + x, lengths = self.up_3_to_9.forward_varlen(x, lengths) + x = x + skip_9 + for block in self.dec_blocks_9: + x = block.forward_varlen(x, c9, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_9) + + return x, types # hidden [total,d], types [total] + + def forward(self, coord_list, noisy_subcode_list, t, depth, + target_subcode_list=None, loss_mask=None, pc=None, return_extra_acc=False): + """训练入口：把因子分解头放进 forward 内部(经 factorized_loss)，保证 DDP 正确同步。 + depth: [b] long，每个样本的层级(target_depth)，与 time 一起作为 adaLN 条件。 + - 传 target/loss_mask -> 返回 (loss, acc)；return_extra_acc=True 时返回 (loss, acc, acc_nonempty) + (acc_nonempty=连接非空行准确率，见 factorized_accuracy)。默认 False 与旧接口兼容。 + - 不传 -> 返回 (hidden, types)（推理/调试） + """ + h, types = self.forward_varlen(coord_list, noisy_subcode_list, t, depth, pc) + if target_subcode_list is None: + return h, types + target = torch.cat([s.to(device=h.device, dtype=torch.long) for s in target_subcode_list], dim=0) + loss_mask = loss_mask.to(device=h.device) + loss = self.factorized_loss(h, target, types, loss_mask) + with torch.no_grad(): + acc, acc_ne = self.factorized_accuracy(h, target, types, loss_mask) + if return_extra_acc: + return loss, acc, acc_ne + return loss, acc + + # ------------------------------------------------------------------ # + # 因子分解 loss / accuracy + # ------------------------------------------------------------------ # + + + # ------------------------------------------------------------------ # + # 采样(MaskGIT 风格，因子分解、varlen、逐位置置信度解 mask) + # ------------------------------------------------------------------ # + @torch.no_grad() + def sample(self, coord_list, depth, pc=None, T_infer=10, cfg_scale=1.0, temperature=0.0, + cfg_schedule="constant", cfg_start=1.0): + """给定 coord(粗层条件)与 depth(层级)，从全 MASK 迭代生成 refine 子码。 + depth: [b] long(每样本层级)。返回 subcode_list(list[ [L,8] ])。 + cfg_scale: classifier-free guidance 强度，约定 logits = uncond + cfg_scale*(cond - uncond)。 + =1.0 即纯条件(无引导，跳过无条件前向)；>1 放大点云条件(锐化连接预测、对抗欠连接)。 + 仅在 use_pc 且 pc 提供时生效(无条件分支用 pc=None 走训练过的 null_cond)。 + temperature: <=0 argmax 贪心(确定性，默认，与旧行为一致)；>0 多项式随机采样。 + best-of-N 时需 >0(如 1.0) 才能采出多样候选，否则贪心 N 次结果完全相同。 + cfg_schedule: "constant"=全程恒定 cfg_scale(默认，旧行为)；"linear"=按步从 cfg_start + 线性升到 cfg_scale(MUSE 式 ramp)。最初几步决定全局结构，CFG 锐化会把条件 + 敏感(=分歧大、没把握)位置的置信度人为抬高、过早 commit 且不可回滚；低 cfg + 起步可避开这种早期过冲，引导力留给后期细化。 + cfg_start: ramp 起点(仅 linear 生效)。默认 1.0 = 第一步纯条件(并跳过该步无条件前向)。""" + self.eval() + device = next(self.parameters()).device + coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list] + depth = depth.to(device=device, dtype=torch.long) + lengths = torch.tensor([c.numel() for c in coord_list], device=device, dtype=torch.long) + types = self.types_from_lengths(lengths) + want_cfg = (float(cfg_scale) != 1.0 or float(cfg_start) != 1.0) and self.use_pc and (pc is not None) + + total = int(lengths.sum().item()) + subcodes = torch.full((total, NUM_OCTANTS), MASK_ID, dtype=torch.long, device=device) + decided = torch.zeros(total, dtype=torch.bool, device=device) + + for step in range(T_infer): + # 当步实际引导强度：linear 时 step0 = cfg_start，末步 = cfg_scale + if cfg_schedule == "linear": + cfg_t = float(cfg_start) + (float(cfg_scale) - float(cfg_start)) * step / max(T_infer - 1, 1) + else: + cfg_t = float(cfg_scale) + use_cfg = want_cfg and (cfg_t != 1.0) + + ratio = step / T_infer + t = torch.full((len(coord_list),), ratio, device=device, dtype=torch.float) + sub_list = list(torch.split(subcodes, lengths.tolist(), dim=0)) + h, _ = self.forward_varlen(coord_list, sub_list, t, depth, pc) + h_u = self.forward_varlen(coord_list, sub_list, t, depth, None)[0] if use_cfg else None + + # 逐类型预测子码 + 置信度(该位置有效子码的平均最大概率) + pred = subcodes.clone() + conf = torch.zeros(total, device=device) + # F: 单头(可 CFG) + mf = types == TYPE_F + if mf.any(): + logits = self.f_head(h[mf]) # [nm,256] + if use_cfg: + lu = self.f_head(h_u[mf]) + logits = lu + cfg_t * (logits - lu) + if temperature and temperature > 0: + prob = (logits / temperature).softmax(-1) + am = torch.multinomial(prob.float(), 1).squeeze(-1) # [nm] + pm = prob.gather(-1, am.unsqueeze(-1)).squeeze(-1) + else: + prob = logits.softmax(-1) + pm, am = prob.max(-1) # [nm],[nm] + pred[mf, :] = am.unsqueeze(1).expand(-1, NUM_OCTANTS) # F 复制到 8 槽 + conf[mf] = pm.to(conf.dtype) + # intra/inter: 8 行自回归 rollout(可 CFG, cond/uncond 在每行 logits 级合并) + for ty, head in ((TYPE_INTRA, self.intra_head), (TYPE_INTER, self.inter_head)): + m = types == ty + if not m.any(): + continue + p, pmaxs = head.rollout(h[m], h_u[m] if use_cfg else None, cfg_t, temperature) # [nm,8],[nm,8] + pred[m] = p + conf[m] = pmaxs.mean(dim=1).to(conf.dtype) # autocast 下 pmax 可能是 bf16 + + # 已决定的位置保持不变 + keep_ratio = cosine_schedule(torch.tensor(float(step + 1)), float(T_infer)).item() + n_keep_masked = int(keep_ratio * total) # 仍保持 mask 的目标数(随步数减少) + conf_for_select = conf.clone() + conf_for_select[decided] = float('inf') # 已定的不再参与"新解mask" + # 选出本步要确定的位置：未定中置信度最高的若干个 + n_to_decide = max(0, (total - n_keep_masked) - int(decided.sum().item())) + if n_to_decide > 0: + undecided_idx = (~decided).nonzero(as_tuple=True)[0] + order = conf[undecided_idx].argsort(descending=True) + chosen = undecided_idx[order[:n_to_decide]] + subcodes[chosen] = pred[chosen] + decided[chosen] = True + # 未决定位置回到 MASK，已决定保持 + subcodes[~decided] = MASK_ID + + # 收尾：还没定的全部按当前预测填上 + subcodes[~decided] = pred[~decided] + return list(torch.split(subcodes, lengths.tolist(), dim=0)) + +# ========================================================================== +# 原 model_acc_fix3_altpc.py +# ========================================================================== + +_BF16_RESIDUAL = os.environ.get("BF16_RESIDUAL", "1") == "1" +_STATIC_PAD_BUCKETS = tuple(sorted(int(x) for x in os.environ.get("STATIC_PAD_BUCKETS", "").split(",") if x.strip())) +_STATIC_PAD_SLOTS = int(os.environ.get("STATIC_PAD_SLOTS", "16")) +assert all(b % 9 == 0 for b in _STATIC_PAD_BUCKETS), "STATIC_PAD_BUCKETS 必须都是 9 的倍数" +class GeoRoPE: + """几何 RoPE(rope_mode="geo"/"geo_vcat")的相位计算器。纯计算、零参数(不进 state_dict)。 + token 的"位置" = 3D 锚点 + 面内类型, 序列顺序完全不参与注意力(模型对面片块置换严格等变)。 + + mode="geo": 锚点用聚合点(旧线 ckpt 在用, 勿删): + 9/face: F_i/M_ii 锚在父顶点 v_i, M_ij 锚在边中点 (v_i+v_j)/2; 类型 0/1/2=F/intra/inter + 3/face: 三个组 token 同锚面质心, 类型 = 组 id 0/1/2(与 9/face 同一条语义轴) + 1/face: 面质心, 类型 = 3 + head_dim/2 个频率槽切 [n,n,n,type_slots](n=(d/2-type_slots)//3, 余数补恒等相位), + 依序旋转 x/y/z/type(LATO RotaryPositionPhasesEmbedder 的 rotate_half 版)。 + mode="geo_vcat"(2026-08-05, 用户提议): 锚点显式编码顶点, **不用质心/中点等聚合点**。 + 3/face 与 1/face 粗尺度: 「三顶点各占 1/3 通道拼接」(共质心三角形可分): + xyz 30 槽切 3 组 x 每顶点 9 槽(x/y/z 各 3, 顶点序=面内 zyx 字典序), type 2 槽, 余 3 pad。 + 9/face: 「双端点布局」(治中点混叠: 平行四边形两对角线同中点会被 geo 判成同一条边): + 每 token 两锚点各占一半 xyz 槽(x/y/z 各 5), F_i/M_ii=(v_i,v_i) 复制, M_ij=(v_i,v_j) + (zyx 排序 => 端点对全局规范, 共享边跨面同相位); F 与其所属边的 inter 前半通道对齐。 + 频率梯子统一 theta_vcat=300(3 槽/轴最低频 2.85 rad、5 槽/轴 1.33 rad, 均 <π 无歧义, + 邻格 1 rad 可分辨)。与 geo ckpt 参数完全同构, 但两个尺度的注意力几何都变了 -> + 热启动 loss 先跳后收(同 flat->geo 的迁移形态)。 + + 坐标约定(两种模式共同): + 锚点用**原始体素坐标**(乘 geo_scale, 默认 1.0), 刻意不按 depth 归一: + 注意力只见同样本内的相对相位(每样本单一 depth), 而连接决策的几何是体素邻接 —— + 体素单位让"相邻一格"在所有 depth 下相位签名一致(θ_xyz=100: 邻格 1 rad 可分辨, + depth7 全程 128 格最低频 2.0 rad < π 无歧义)。 + coord 不减 offset: 常数平移让 q/k 同转一个角, 相对相位(注意力唯一可见量)不变。 + 纯位置编码零新增参数 -> 任意旧 ckpt strict 热启动(注意几何变了, loss 先跳后收)。 + """ + + def __init__(self, mode="flat", geo_scale=1.0, theta_xyz=100.0, + theta_type=10.0, type_slots=2, theta_vcat=300.0): + self.mode = mode + self.geo_scale = float(geo_scale) + self.theta_xyz = float(theta_xyz) + self.theta_type = float(theta_type) + self.type_slots = int(type_slots) + self.theta_vcat = float(theta_vcat) + + @property + def is_geo(self): + return str(self.mode).startswith("geo") + + @staticmethod + def _ladder(theta, n, device): + """几何频率梯子: freq_i = 1/theta^(i/n), i=0..n-1。""" + return 1.0 / (theta ** (torch.arange(n, device=device, dtype=torch.float32) / max(n, 1))) + + @staticmethod + def _cos_sin(rot, dtype): + """[N, d_half] 相位 -> ([N,1,head_dim] cos, sin): 频率复制两半配 rotate_half, + 布局与 RotaryEmbedding.forward_positions 一致, apply_rotary_pos_emb_varlen 原样可用。""" + emb = torch.cat((rot, rot), dim=-1) + return emb.cos().to(dtype).unsqueeze(1), emb.sin().to(dtype).unsqueeze(1) + + def scales(self, coord_packed, d_half, dtype): + """由 coord 流现算三个 hourglass 尺度的 (cos, sin), 返回 {9:(cos,sin), 3:..., 1:...}。 + 每面 9 个 coord 标量 = 3 父顶点 x 3 分量(面内坐标字典序); 样本长度都是 9 的倍数且 + face 块不跨样本, 所以可以全局 view。""" + device = coord_packed.device + V = coord_packed.view(-1, 3, 3).to(torch.float32) * self.geo_scale # [nF, 3顶点, xyz] + nF = V.shape[0] + freqs_type = self._ladder(self.theta_type, self.type_slots, device) + # 面内类型分量(与 token 排布一致): 9/face = F F F | intra x3 | inter x3; + # 3/face = 组 id 0/1/2(与 9/face 同一条语义轴); 1/face = 3 + type9 = torch.tensor([0., 0., 0., 1., 1., 1., 2., 2., 2.], device=device).repeat(nF) + type3 = torch.tensor([0., 1., 2.], device=device).repeat(nF) + type1 = torch.full((nF,), 3.0, device=device) + if getattr(self, "mode", "geo") == "geo_vcat": + return self._scales_vcat(V, type9, type3, type1, freqs_type, d_half, dtype, device) + return self._scales_geo(V, type9, type3, type1, freqs_type, d_half, dtype, device) + + def _scales_geo(self, V, type9, type3, type1, freqs_type, d_half, dtype, device): + """geo 布局: [x*n_ax | y*n_ax | z*n_ax | type*type_slots | pad]。""" + n_type = self.type_slots + n_ax = (d_half - n_type) // 3 + n_pad = d_half - n_type - 3 * n_ax + freqs_ax = self._ladder(self.theta_xyz, n_ax, device) + + nF = V.shape[0] + mid01 = (V[:, 0] + V[:, 1]) * 0.5 + mid02 = (V[:, 0] + V[:, 2]) * 0.5 + mid12 = (V[:, 1] + V[:, 2]) * 0.5 + centroid = V.mean(dim=1) # [nF,3] + # 槽位序 = [F0 F1 F2 M00 M11 M22 M01 M02 M12](与 token 排布一致) + anchors9 = torch.stack([V[:, 0], V[:, 1], V[:, 2], + V[:, 0], V[:, 1], V[:, 2], + mid01, mid02, mid12], dim=1) # [nF,9,3] + anchors3 = centroid[:, None, :].expand(nF, 3, 3) # down_9_to_3 按连续三元折叠成组 + anchors1 = centroid[:, None, :] + + def build(anchors, type_comp): + a = anchors.reshape(-1, 3) + parts = [a[:, 0:1] * freqs_ax, a[:, 1:2] * freqs_ax, a[:, 2:3] * freqs_ax, + type_comp[:, None] * freqs_type] + if n_pad: + parts.append(torch.zeros(a.shape[0], n_pad, device=device, dtype=torch.float32)) + return self._cos_sin(torch.cat(parts, dim=-1), dtype) + + return {9: build(anchors9, type9), 3: build(anchors3, type3), 1: build(anchors1, type1)} + + def _scales_vcat(self, V, type9, type3, type1, freqs_type, d_half, dtype, device): + """geo_vcat: 锚点全部为显式顶点(不算质心/中点)。""" + n_type = self.type_slots + + # ---- 粗尺度(3/face 与 1/face) = 三顶点拼接: 每顶点每轴 n_per_vert 槽, + # 布局 [v0: x,y,z | v1: x,y,z | v2: x,y,z | type | pad] + n_per_vert = (d_half - n_type) // 9 + n_pad_v = d_half - n_type - 9 * n_per_vert + freqs_vcat = self._ladder(self.theta_vcat, n_per_vert, device) + + def build_vcat(type_comp, reps): + parts_xyz = [V[:, vtx, ax:ax + 1] * freqs_vcat for vtx in range(3) for ax in range(3)] + rot_face = torch.cat(parts_xyz, dim=-1) # [nF, 9*n_per_vert] + rot = rot_face.repeat_interleave(reps, dim=0) if reps > 1 else rot_face + parts = [rot, type_comp[:, None] * freqs_type] + if n_pad_v: + parts.append(torch.zeros(rot.shape[0], n_pad_v, device=device, dtype=torch.float32)) + return self._cos_sin(torch.cat(parts, dim=-1), dtype) + + # ---- 9/face = 双端点布局: 每 token 两个锚点各占一半 xyz 槽 + # [P1: x/y/z 各 n_ax9 | P2: 同 | type n_type], 30=2x15 无 pad。 + # F_i/M_ii -> (v_i, v_i) 复制; M_ij -> (v_i, v_j), i 端点对全局规范, + # 共享边跨面同相位)。免费性质: F 锚 v_a 与 inter 锚 (v_a,v_b) 前半通道完全对齐 —— + # "顶点是边的端点"在相位空间直接可见。 + n_ax9 = (d_half - n_type) // 6 + n_pad9 = d_half - n_type - 6 * n_ax9 + freqs_9 = self._ladder(self.theta_vcat, n_ax9, device) + P1 = torch.stack([V[:, 0], V[:, 1], V[:, 2], V[:, 0], V[:, 1], V[:, 2], + V[:, 0], V[:, 0], V[:, 1]], dim=1).reshape(-1, 3) # [nF*9,3] + P2 = torch.stack([V[:, 0], V[:, 1], V[:, 2], V[:, 0], V[:, 1], V[:, 2], + V[:, 1], V[:, 2], V[:, 2]], dim=1).reshape(-1, 3) + parts9 = [P1[:, 0:1] * freqs_9, P1[:, 1:2] * freqs_9, P1[:, 2:3] * freqs_9, + P2[:, 0:1] * freqs_9, P2[:, 1:2] * freqs_9, P2[:, 2:3] * freqs_9, + type9[:, None] * freqs_type] + if n_pad9: + parts9.append(torch.zeros(P1.shape[0], n_pad9, device=device, dtype=torch.float32)) + scales9 = self._cos_sin(torch.cat(parts9, dim=-1), dtype) + return {9: scales9, 3: build_vcat(type3, 3), 1: build_vcat(type1, 1)} +_COORD_PACK_BASE = 4096 # 坐标打包基数(> 最大体素坐标), xyz -> 单一整数 code +def _pack_coord(p3): + """[.,3] 整数 xyz -> 单一整数 code(供 torch.unique 分组用)。""" + return p3[:, 0] * _COORD_PACK_BASE * _COORD_PACK_BASE + p3[:, 1] * _COORD_PACK_BASE + p3[:, 2] +class TokenGroup: + """一个跨面 token 组划分 + 组内表决。 + 字段: pos=该类型 token 的全局位置; group=每 token 的组 id(0..n_group-1); + face=每 token 所在面; code=父顶点打包坐标(仅 by_parent, 顶点原子揭示要用)。""" + + def __init__(self, pos, group, n_group, face, code=None): + self.pos, self.group, self.n_group, self.face, self.code = pos, group, n_group, face, code + + @classmethod + def by_parent(cls, types, face_coords, token_type, parent_offset): + """F/intra 按父顶点分组(同一顶点在不同面的 token 同组)。 + parent_offset: F=0(面内位置 0/1/2 -> 父 v0/v1/v2), intra=3(位置 3/4/5)。""" + pos = (types == token_type).nonzero(as_tuple=True)[0] + face = pos // REFINE_POS_PER_FACE + code = _pack_coord(face_coords[face, (pos % REFINE_POS_PER_FACE) - parent_offset]) + uniq, group = torch.unique(code, return_inverse=True) + return cls(pos, group, int(uniq.numel()), face, code) + + @classmethod + def by_edge(cls, types, face_coords): + """inter 按「无序父对(=共享边)」分组: 面内位置 6/7/8 -> 父对 (v0,v1)/(v0,v2)/(v1,v2)。 + 父对取 (小,大) => 同一条边跨面同组, 朝向天然规范(无 M vs M^T 问题)。""" + device = face_coords.device + pos = (types == TYPE_INTER).nonzero(as_tuple=True)[0] + face = pos // REFINE_POS_PER_FACE + local = (pos % REFINE_POS_PER_FACE) - 6 # 0/1/2 -> M01/M02/M12 + pa = torch.tensor([0, 0, 1], device=device) + pb = torch.tensor([1, 2, 2], device=device) + code_a = _pack_coord(face_coords[face, pa[local]]) + code_b = _pack_coord(face_coords[face, pb[local]]) + pair = torch.stack([torch.minimum(code_a, code_b), torch.maximum(code_a, code_b)], dim=1) + uniq, group = torch.unique(pair, dim=0, return_inverse=True) + return cls(pos, group, int(uniq.shape[0]), face) + + @staticmethod + def _mode_of(group, vid, n_group): + """每组内 vid 的众数(sort+segment, 免 python 循环), 广播回每个成员。""" + n_vid = int(vid.max().item()) + 1 if vid.numel() else 1 + key, cnt = torch.unique(group.long() * n_vid + vid.long(), return_counts=True) + g, v = key // n_vid, key % n_vid + order = torch.argsort(g * (int(cnt.max().item()) + 1) - cnt) # 组内按计数降序 + g, v = g[order], v[order] + first = torch.ones_like(g, dtype=torch.bool) + if g.numel() > 1: + first[1:] = g[1:] != g[:-1] + mode = torch.zeros(n_group, dtype=torch.long, device=vid.device) + mode[g[first]] = v[first] + return mode[group] + + def vote_bytes(self, vals): + """F 通道: 每组对占用字节(0..255)计数表决, 返回广播回每 token 的众数字节。""" + cnt = torch.bincount(self.group * 256 + vals.clamp(0, 255), + minlength=self.n_group * 256).view(self.n_group, 256) + return cnt.argmax(1)[self.group] + + def vote_vec(self, vecs): + """intra/inter 通道: 8-subcode 向量整体表决。返回 (该 token 是否≠众数, 广播回的众数向量)。""" + uniq, vid = torch.unique(vecs, dim=0, return_inverse=True) + mode_vid = self._mode_of(self.group, vid, self.n_group) + return vid != mode_vid, uniq[mode_vid] +class _DualCloudCond: + """一次前向内的点云条件: A=全GT云, B=目标分辨率塌缩云, 按全局 block 序号交替提供 cross-attn KV + (奇数层->A, 偶数层->B; B 未提供时两套相同 = 原单点云行为; 云为 None 走 null_cond = CFG 无条件分支)。 + 每套 cond = (packed KV [sum(len),d], 每样本 KV 长度, KV max_len)。 + fake_lens 桶 pad 时在 conditioner **之后**补零行(不浪费 Michelangelo 前向)。 + cache: 推理跨步复用编码(采样各步点云不变, miche 逐步重算纯浪费; 数值严格等价), 训练传 None。""" + + def __init__(self, model, coord_list, fake_lens, device, pc, pc_target, cache=None): + self._model = model + self._n_seq = len(coord_list) # 含 fake 序列(null_cond 也要按批展开) + self._fake_lens = fake_lens + self._device = device + key = ("null" if pc is None else "cond", pc_target is None) + if cache is not None and key in cache: + self.cond_a, self.cond_b = cache[key] + return + self.cond_a = self._encode(pc) # 奇数层(全GT) + if model.use_pc and (pc_target is not None): + self.cond_b = self._encode(pc_target) # 偶数层(目标分辨率塌缩) + else: + self.cond_b = self.cond_a # 未提供 -> 两套相同(原行为) + if cache is not None: + cache[key] = (self.cond_a, self.cond_b) + + def _encode(self, cloud): + model, device = self._model, self._device + if not model.use_pc: + return None, None, None + if cloud is None: # CFG 无条件分支 + null = model.null_cond + kv = null.unsqueeze(0).expand(self._n_seq, -1, -1).reshape(-1, null.size(-1)) + kv_len = torch.full((self._n_seq,), model.null_cond_len, device=device, dtype=torch.long) + return kv, kv_len, int(model.null_cond_len) + emb = model.conditioner(cloud.to(device=device)) + if self._fake_lens: # 假序列的 cross-attn KV = 零行 + emb = torch.cat([emb, emb.new_zeros(len(self._fake_lens), emb.size(1), emb.size(2))], dim=0) + kv_len = torch.full((emb.size(0),), emb.size(1), device=device, dtype=torch.long) + return emb.reshape(-1, emb.size(-1)), kv_len, int(emb.size(1)) + + def for_layer(self, gi): + """gi: 0-based 全局 block 序号; 层号=gi+1 -> 奇数层(gi 偶)用 A, 偶数层(gi 奇)用 B。""" + return self.cond_a if gi % 2 == 0 else self.cond_b +class _ReMDM: + """ReMDM(arXiv 2503.00307, 训练免费): 步首以 σ_t 重掩已定 token, 本步前向重新预测; + 揭示额度按 n_keep_masked 自动补足 -> 掩码数轨迹严格贴原 cosine 调度(边际保持的计数空间等价)。 + σ_max=min(1, m_s/|C|); cap: σ=min(η,σ_max); rescale: σ=η·σ_max; conf: cap 的量+定向重掩 + 提交置信度最低者。env 门控(sample 调用时读取), REMDM 为空 = 与原行为逐位一致。""" + + def __init__(self): + self.mode = os.environ.get("REMDM", "").strip().lower() + self.eta = float(os.environ.get("REMDM_ETA", "0.05")) + self.t_start = float(os.environ.get("REMDM_TSTART", "0.0")) + self.t_end = float(os.environ.get("REMDM_TEND", "1.0")) + + def maybe_remask(self, step, T_infer, total, decided, subcodes, commit_conf, device): + """就地重掩(decided/subcodes 原位改写)。""" + if not (self.mode and step > 0 and self.t_start <= step / T_infer <= self.t_end): + return + n_committed = int(decided.sum().item()) + if n_committed == 0: + return + n_mask_target = int(cosine_schedule(torch.tensor(float(step + 1)), float(T_infer)).item() * total) + sigma_max = min(1.0, n_mask_target / n_committed) + sigma = self.eta * sigma_max if self.mode == "rescale" else min(self.eta, sigma_max) + n_remask = int(sigma * n_committed) + if n_remask <= 0: + return + committed_idx = decided.nonzero(as_tuple=True)[0] + if self.mode == "conf": + remask_idx = committed_idx[commit_conf[committed_idx].argsort()[:n_remask]] # 最低提交置信优先 + else: + remask_idx = committed_idx[torch.randperm(n_committed, device=device)[:n_remask]] + decided[remask_idx] = False + subcodes[remask_idx] = MASK_ID +class FactorizedDiffusionTransformerAltPC(FactorizedDiffusionTransformer): + def __init__(self, *args, rope_mode="flat", rope_geo_scale=1.0, + rope_theta_xyz=100.0, rope_theta_type=10.0, rope_type_slots=2, + rope_theta_vcat=300.0, **kwargs): + super().__init__(*args, **kwargs) + # 几何 RoPE 封装在 GeoRoPE(纯计算无参数, 不影响 state_dict)。 + # rope_mode="flat"(默认): 原行为, 逐 token 序列位置 0..L-1, 与所有旧 ckpt 完全兼容; + # "geo"/"geo_vcat" 的设计与 theta 约束见 GeoRoPE docstring。 + self.geo_rope = GeoRoPE(mode=rope_mode, geo_scale=rope_geo_scale, + theta_xyz=rope_theta_xyz, theta_type=rope_theta_type, + type_slots=rope_type_slots, theta_vcat=rope_theta_vcat) + + # rope_* 属性代理到 GeoRoPE(单一事实源): 外部读写 self.rope_mode 等保持兼容 + # (含子类/脚本事后改 rope_mode 的用法, 如 tssr 的 self.rope_mode = "hier")。 + rope_mode = property(lambda self: self.geo_rope.mode, + lambda self, v: setattr(self.geo_rope, "mode", v)) + rope_geo_scale = property(lambda self: self.geo_rope.geo_scale, + lambda self, v: setattr(self.geo_rope, "geo_scale", float(v))) + rope_theta_xyz = property(lambda self: self.geo_rope.theta_xyz, + lambda self, v: setattr(self.geo_rope, "theta_xyz", float(v))) + rope_theta_type = property(lambda self: self.geo_rope.theta_type, + lambda self, v: setattr(self.geo_rope, "theta_type", float(v))) + rope_type_slots = property(lambda self: self.geo_rope.type_slots, + lambda self, v: setattr(self.geo_rope, "type_slots", int(v))) + rope_theta_vcat = property(lambda self: self.geo_rope.theta_vcat, + lambda self, v: setattr(self.geo_rope, "theta_vcat", float(v))) + + def _geo_rope_scales(self, coord_packed, dtype): + """geo/geo_vcat: 三个 hourglass 尺度的 (cos,sin)。方法名/签名保留(fpsq 子类在调用)。""" + return self.geo_rope.scales(coord_packed, self.rotary_emb.inv_freq.numel(), dtype) + + @staticmethod + def _static_pad_plan(total, b_real): + """返回假序列长度列表(都是 9 的倍数且 >=9): padded 总数落到桶位、批槽位补满 _STATIC_PAD_SLOTS。 + 超过最大桶位时返回 None(不 pad, 触发一次形状特化编译; 配 hard_max_positions 应当罕见)。""" + n_fake = max(_STATIC_PAD_SLOTS - b_real, 1) + min_pad = 9 * n_fake + target = next((B for B in _STATIC_PAD_BUCKETS if B >= total + min_pad), None) + if target is None: + return None + fake = [9] * n_fake + fake[0] += target - total - min_pad # 桶位/总数都是 9 的倍数 -> 余量也是 + return fake + + def forward_varlen(self, coord_list, subcode_list, t, depth, pc=None, pc_target=None, + cond_cache=None): + device = t.device + coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list] + subcode_list = [s.to(device=device, dtype=torch.long) for s in subcode_list] + # 长度在 python 端本来就知道 -> 一并算出各尺度 max_seqlen 传给 attention, + # 省掉每个 flash 调用的 int(lengths.max().item()) GPU->CPU 同步(见 model_acc_fix3_tssr)。 + lengths_py = [int(c.numel()) for c in coord_list] + if sum(lengths_py) == 0: + raise ValueError("forward_varlen received an empty batch.") + if any(n % REFINE_POS_PER_FACE != 0 for n in lengths_py): + raise ValueError("每个序列长度必须是 9 的倍数(每 face 9 个 refine 位置)。") + + # ---- 静态桶化 padding(见文件顶部说明; STATIC_PAD_BUCKETS 为空 = 原行为) ---- + real_total, real_b = sum(lengths_py), len(lengths_py) + fake_lens = self._static_pad_plan(real_total, real_b) if _STATIC_PAD_BUCKETS else None + if fake_lens: + sub_tail = tuple(subcode_list[0].shape[1:]) + for L in fake_lens: + coord_list.append(torch.zeros(L, dtype=torch.long, device=device)) + subcode_list.append(torch.zeros((L,) + sub_tail, dtype=torch.long, device=device)) + lengths_py = lengths_py + fake_lens + lengths = torch.tensor(lengths_py, device=device, dtype=torch.long) + if fake_lens: + # pad 模式: max_len 传桶位常数(= padded 总数, 合法上界), 保证编译形状/guard 全静态 + max_len_9 = sum(lengths_py) + else: + max_len_9 = max(lengths_py) + max_len_3 = max_len_9 // 3 # hourglass 整除, 断言在上面 + max_len_1 = max_len_9 // 9 + + coord_packed = torch.cat(coord_list, dim=0) + subcode_packed = torch.cat(subcode_list, dim=0) + types = self.types_from_lengths(lengths) + + raw_emb = self.raw_embedding(coord_packed) + refine_emb = self.embed_refine(subcode_packed, types) + x = raw_emb + refine_emb # 不再加法注入 time；time/depth 走 adaLN + # 统一残差流 dtype(理由见文件顶部 _BF16_RESIDUAL 注释); c_sample 刻意保持 fp32(adaLN 惯例) + if _BF16_RESIDUAL and torch.is_autocast_enabled(): + x = x.to(torch.get_autocast_dtype("cuda")) + + # adaLN 条件 c = time_emb + depth_emb(逐样本 [b,d]); 各 hourglass 阶段按其长度展开到逐 token + depth = depth.to(device=device, dtype=torch.long) + c_sample = self.time_mlp(t) + self.depth_embedding(depth) # [b, d_model] + if fake_lens: # 假序列的 adaLN 条件 = 零行 + c_sample = torch.cat([c_sample, c_sample.new_zeros(len(fake_lens), c_sample.size(-1))], dim=0) + + def c_per_token(cur_lengths): + return c_sample.repeat_interleave(cur_lengths, dim=0) + + # 点云条件: 双云编码 + 逐层交替选择(见 _DualCloudCond) + cond = _DualCloudCond(self, coord_list, fake_lens, device, pc, pc_target, cond_cache) + + # geo 模式: 三个尺度的 cos/sin 一次算好(锚点只依赖 coord/depth, 与序列位置无关) + geo_scales = (self._geo_rope_scales(coord_packed, x.dtype) + if str(getattr(self, "rope_mode", "flat")).startswith("geo") else None) + + def rope_fn_for(cur_lengths, tokens_per_face=9): + if geo_scales is not None: + cos, sin = geo_scales[tokens_per_face] + else: + position_ids = lengths_to_position_ids(cur_lengths) + cos, sin = self.rotary_emb.forward_positions(position_ids, dtype=x.dtype) + return lambda q, k: apply_rotary_pos_emb_varlen(q, k, cos, sin) + + # ---- hourglass 五段: enc9 -> enc3 -> bottleneck -> dec3 -> dec9 ---- + gi = 0 # 全局 block 序号(横跨五段, 决定该层注入云A还是云B) + + def run_stage(x, blocks, c_stage, stage_lengths, rope_fn, stage_max_len): + nonlocal gi + for block in blocks: + cond_kv, cond_kv_len, cond_kv_max = cond.for_layer(gi) + gi += 1 + x = block.forward_varlen(x, c_stage, stage_lengths, cond_kv, cond_kv_len, + rope_apply_fn=rope_fn, max_len=stage_max_len, + cond_max_len=cond_kv_max) + return x + + rope_fn_9, c9 = rope_fn_for(lengths, 9), c_per_token(lengths) + x = run_stage(x, self.enc_blocks_9, c9, lengths, rope_fn_9, max_len_9) + skip_9 = x + + x, lengths = self.down_9_to_3.forward_varlen(x, lengths) + rope_fn_3, c3 = rope_fn_for(lengths, 3), c_per_token(lengths) + x = run_stage(x, self.enc_blocks_3, c3, lengths, rope_fn_3, max_len_3) + skip_3 = x + + x, lengths = self.down_3_to_1.forward_varlen(x, lengths) + rope_fn_1, c1 = rope_fn_for(lengths, 1), c_per_token(lengths) + x = run_stage(x, self.bottleneck_blocks, c1, lengths, rope_fn_1, max_len_1) + + x, lengths = self.up_1_to_3.forward_varlen(x, lengths) + x = run_stage(x + skip_3, self.dec_blocks_3, c3, lengths, rope_fn_3, max_len_3) + + x, lengths = self.up_3_to_9.forward_varlen(x, lengths) + x = run_stage(x + skip_9, self.dec_blocks_9, c9, lengths, rope_fn_9, max_len_9) + + if fake_lens: # 切掉假位置(不进头/loss) + x, types = x[:real_total], types[:real_total] + return x, types # hidden [total,d], types [total] + + def forward(self, coord_list, noisy_subcode_list, t, depth, + target_subcode_list=None, loss_mask=None, pc=None, pc_target=None, return_extra_acc=False): + """训练入口, 与基类同, 只多传 pc_target(逐层交替的第二套点云)。""" + h, types = self.forward_varlen(coord_list, noisy_subcode_list, t, depth, pc, pc_target) + if target_subcode_list is None: + return h, types + target = torch.cat([s.to(device=h.device, dtype=torch.long) for s in target_subcode_list], dim=0) + loss_mask = loss_mask.to(device=h.device) + loss = self.factorized_loss(h, target, types, loss_mask) # focal(继承) + with torch.no_grad(): + acc, acc_ne = self.factorized_accuracy(h, target, types, loss_mask) + if return_extra_acc: + return loss, acc, acc_ne + return loss, acc + + def _predict_heads(self, h, h_u, types, subcodes, cfg_t, temperature): + """三个头对全部位置出预测 + 置信度: F=256 类单分类(整块 9 位广播同值), + intra/inter=8 行自回归 rollout。h_u 非 None 时做 logits 级 CFG。返回 (pred, conf)。""" + pred = subcodes.clone() + conf = torch.zeros(types.numel(), device=h.device) + is_f = types == TYPE_F + if is_f.any(): + logits = self.f_head(h[is_f]) + if h_u is not None: + logits_u = self.f_head(h_u[is_f]) + logits = logits_u + cfg_t * (logits - logits_u) + if temperature and temperature > 0: + prob = (logits / temperature).softmax(-1) + choice = torch.multinomial(prob.float(), 1).squeeze(-1) + choice_prob = prob.gather(-1, choice.unsqueeze(-1)).squeeze(-1) + else: + prob = logits.softmax(-1) + choice_prob, choice = prob.max(-1) + pred[is_f, :] = choice.unsqueeze(1).expand(-1, NUM_OCTANTS) + conf[is_f] = choice_prob.to(conf.dtype) + for token_type, head in ((TYPE_INTRA, self.intra_head), (TYPE_INTER, self.inter_head)): + sel = types == token_type + if not sel.any(): + continue + row_pred, row_conf = head.rollout(h[sel], h_u[sel] if h_u is not None else None, cfg_t, temperature) + pred[sel] = row_pred + conf[sel] = row_conf.mean(dim=1).to(conf.dtype) + return pred, conf + + @torch.no_grad() + def sample(self, coord_list, depth, pc=None, pc_target=None, T_infer=10, cfg_scale=1.0, + temperature=0.0, cfg_schedule="constant", cfg_start=1.0, + align_fm=False, align_inter=False, align_per_step=False, align_debug=True, + diag_vertex=False): + """前向多传 pc_target(逐层交替)。CFG 无条件分支两套点云都走 null_cond。 + + 跨面一致性对齐(GT 侧这些量对共享结构 100% 一致, 见 check_tokenizer_consistency): + align_fm=True : 对齐 F(占用字节, 按"共享父"分组) + intra(M_inner 8-subcode 矩阵, 按"共享父")。 + align_inter=True : 对齐 inter(M_outer 8-subcode 矩阵, 按"共享父对"=共享边分组)。 + 如面 (v0,v1,v2) 与 (v0,v1,v3) 共享边 (v0,v1) -> 其 (v0,v1) 的 inter 应一致。 + align_per_step : False(默认, 推荐)= 只在收尾把各启用通道投影到表决值(不扰动去噪, 不造洞); + True = 每步额外把"不一致的面"整块 9 token 重掩重预测(实测打乱揭示节奏 -> + 反而多洞, 仅供对照)。 + 全 False -> 与原行为逐位一致。 + + env 门控(为空/未设 = 与原行为逐位一致): + REMDM=cap|rescale|conf (+REMDM_ETA/TSTART/TEND): 步首重掩已定 token(见 _ReMDM)。 + WINDOW_REVEAL=1(仅 batch=1): BFS 前缀窗内 token 才可提交(块因果 变体A), 注意力/预测不变。 + """ + self.eval() + device = next(self.parameters()).device + coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list] + depth = depth.to(device=device, dtype=torch.long) + lengths = torch.tensor([c.numel() for c in coord_list], device=device, dtype=torch.long) + types = self.types_from_lengths(lengths) + want_cfg = (float(cfg_scale) != 1.0 or float(cfg_start) != 1.0) and self.use_pc and (pc is not None) + + total = int(lengths.sum().item()) + nface = total // REFINE_POS_PER_FACE + + align_any = align_fm or align_inter + if align_any or diag_vertex: + face_coords = torch.cat(coord_list).view(nface, 3, 3) # [nface, 3父, xyz] + + # ---- 顶点组诊断: 同一父(顶点)在不同面的 F/intra token, 记录 commit step, 末尾统计一致性 ---- + if diag_vertex: + diag_f = TokenGroup.by_parent(types, face_coords, TYPE_F, 0) + diag_i = TokenGroup.by_parent(types, face_coords, TYPE_INTRA, 3) + commit_step = torch.full((total,), int(T_infer), dtype=torch.long, device=device) + + # ---- 跨面对齐分组: F/intra 按"共享父", inter 按"共享父对"(=共享边) ---- + if align_fm: + grp_f = TokenGroup.by_parent(types, face_coords, TYPE_F, 0) + grp_i = TokenGroup.by_parent(types, face_coords, TYPE_INTRA, 3) + if align_inter: + grp_e = TokenGroup.by_edge(types, face_coords) + + subcodes = torch.full((total, NUM_OCTANTS), MASK_ID, dtype=torch.long, device=device) + decided = torch.zeros(total, dtype=torch.bool, device=device) + pred = subcodes.clone() + + remdm = _ReMDM() # env 门控, 未设=无操作 + window_reveal = os.environ.get("WINDOW_REVEAL", "0") == "1" and len(lengths) == 1 + commit_conf = torch.zeros(total, device=device) # 每 token 提交时的置信度(ReMDM-conf 用) + cond_cache = {} # 点云条件跨步缓存(严格等价, 见 _DualCloudCond) + + for step in range(T_infer): + remdm.maybe_remask(step, T_infer, total, decided, subcodes, commit_conf, device) + + if cfg_schedule == "linear": + cfg_t = float(cfg_start) + (float(cfg_scale) - float(cfg_start)) * step / max(T_infer - 1, 1) + else: + cfg_t = float(cfg_scale) + use_cfg = want_cfg and (cfg_t != 1.0) + + # ---- 前向(有条件 + 可选无条件 CFG 分支) -> 三头全量预测 ---- + t = torch.full((len(coord_list),), step / T_infer, device=device, dtype=torch.float) + sub_list = list(torch.split(subcodes, lengths.tolist(), dim=0)) + h, _ = self.forward_varlen(coord_list, sub_list, t, depth, pc, pc_target, cond_cache=cond_cache) + h_u = self.forward_varlen(coord_list, sub_list, t, depth, None, None, cond_cache=cond_cache)[0] if use_cfg else None + h = h.float() # 头部权重 fp32; bf16 残差模式下需 cast(fp32 下 no-op) + if h_u is not None: + h_u = h_u.float() + pred, conf = self._predict_heads(h, h_u, types, subcodes, cfg_t, temperature) + + # ---- 每步重掩(align_per_step): 跨面不一致的面整块 9 token 重掩(本步不揭示+撤销已定)。 + # 检测 F/intra(按共享父)与 inter(按共享父对)的不一致; 任一通道不一致的面都重掩。实测会造洞, 默认关。 + if align_per_step and align_any: + force_face = torch.zeros(nface, dtype=torch.bool, device=device) + if align_fm: + if grp_f.pos.numel(): + bad = pred[grp_f.pos, 0] != grp_f.vote_bytes(pred[grp_f.pos, 0]) + force_face[grp_f.face[bad]] = True + if grp_i.pos.numel(): + dis_intra, _ = grp_i.vote_vec(pred[grp_i.pos]) + force_face[grp_i.face[dis_intra]] = True + if align_inter and grp_e.pos.numel(): + dis_inter, _ = grp_e.vote_vec(pred[grp_e.pos]) + force_face[grp_e.face[dis_inter]] = True + remask = force_face.repeat_interleave(REFINE_POS_PER_FACE) + conf[remask] = -float('inf') + decided = decided & (~remask) + + # ---- 按 cosine 调度揭示 top-k 置信 token ---- + keep_ratio = cosine_schedule(torch.tensor(float(step + 1)), float(T_infer)).item() + n_keep_masked = int(keep_ratio * total) + n_to_decide = max(0, (total - n_keep_masked) - int(decided.sum().item())) + if n_to_decide > 0: + undecided_idx = (~decided).nonzero(as_tuple=True)[0] + if window_reveal: + # 块序前沿窗口(推理期块因果 变体A): 仅允许 BFS 前缀窗内的 token 本步被提交。 + # 注意力/预测不变(coord 脚手架全局可见), 只限制提交资格; 窗内不够本步预算时自适应扩窗。 + n_blocks = total // 9 + win_blocks = max(1, int(np.ceil(n_blocks * (step + 1) / float(T_infer)))) + while True: + in_window = undecided_idx[undecided_idx < win_blocks * 9] + if len(in_window) >= n_to_decide or win_blocks >= n_blocks: + break + win_blocks += max(1, (n_to_decide - len(in_window) + 8) // 9) + undecided_idx = undecided_idx[undecided_idx < win_blocks * 9] + order = conf[undecided_idx].argsort(descending=True) + chosen = undecided_idx[order[:n_to_decide]] + subcodes[chosen] = pred[chosen] + decided[chosen] = True + commit_conf[chosen] = conf[chosen] + if diag_vertex: + commit_step[chosen] = step + subcodes[~decided] = MASK_ID + + subcodes[~decided] = pred[~decided] + if diag_vertex: + self._diag_vertex_report(subcodes, commit_step, diag_f.pos, diag_f.group, + diag_i.pos, diag_i.group, int(T_infer)) + + # ---- 末步: 把各启用通道强制对齐到表决值(不扰动去噪, 推荐的对齐方式)。 + # align_debug=True: 打印每通道"与表决值不一致(即被改动)的 token 数 / 涉及面数"便于 debug。 + def report(tag, disagree, face_of): + n_tok = int(disagree.sum().item()) + n_all = disagree.numel() + n_face = int(torch.unique(face_of[disagree]).numel()) if n_tok else 0 + print(f"[align_debug] {tag:5s}: 不一致 {n_tok}/{n_all} token ({100 * n_tok / max(n_all, 1):.1f}%), 涉及 {n_face} 面") + + if align_fm: + if grp_f.pos.numel(): + voted_f = grp_f.vote_bytes(subcodes[grp_f.pos, 0]) + if align_debug: + report("F", subcodes[grp_f.pos, 0] != voted_f, grp_f.face) + subcodes[grp_f.pos, :] = voted_f.unsqueeze(1).expand(-1, NUM_OCTANTS) + if grp_i.pos.numel(): + dis_intra, voted_intra = grp_i.vote_vec(subcodes[grp_i.pos]) + if align_debug: + report("intra", dis_intra, grp_i.face) + subcodes[grp_i.pos, :] = voted_intra + if align_inter and grp_e.pos.numel(): + dis_inter, voted_inter = grp_e.vote_vec(subcodes[grp_e.pos]) + if align_debug: + report("inter", dis_inter, grp_e.face) + subcodes[grp_e.pos, :] = voted_inter + return list(torch.split(subcodes, lengths.tolist(), dim=0)) + + + @torch.no_grad() + def _diag_vertex_report(self, subcodes, commit_step, f_pos, f_grp, i_pos, i_grp, T): + """诊断: 同一父(顶点)在不同面的 F/intra token 是否一致, 以及不一致的时序来源。 + 关键问题——不一致是(a)同步揭示时就分裂, 还是(b)已有同组 token commit(可见上下文)后仍预测错?""" + def analyze(pos, grp, tag, is_vec): + if pos.numel() == 0: + print(f"[diag {tag}] 无 token") + return + if is_vec: # intra: 8-subcode 向量 -> 唯一 id + _, vid = torch.unique(subcodes[pos], dim=0, return_inverse=True) + vals = vid.cpu().numpy() + else: # F: 占用字节 + vals = subcodes[pos, 0].clamp(0, 255).cpu().numpy() + groups = grp.cpu().numpy() + steps = commit_step[pos].cpu().numpy() + order = np.argsort(groups, kind="stable") + groups, vals, steps = groups[order], vals[order], steps[order] + segs = np.split(np.arange(len(groups)), np.flatnonzero(np.diff(groups)) + 1) + n_multi = n_incons = 0 + post_total = post_dis = 0 # 后揭示事件(同组已有更早 commit): 与已定值 一致/不一致 + coreveal = postcommit = both = 0 # 不一致组来源分类 + spreads = [] + for seg in segs: + if len(seg) < 2: + continue + n_multi += 1 + v, s = vals[seg], steps[seg] + inconsistent = len(set(v.tolist())) > 1 + first_step = int(s.min()) + early_vals = v[s == first_step] + ref = np.bincount(early_vals).argmax() # 最早 step 的众数值作参照 + coreveal_split = len(set(early_vals.tolist())) > 1 + later = v[s > first_step] + post_total += len(later) + post_dis += int((later != ref).sum()) + has_post = bool((later != ref).any()) + if inconsistent: + n_incons += 1 + spreads.append(int(s.max() - s.min())) + if coreveal_split and has_post: + both += 1 + elif coreveal_split: + coreveal += 1 + else: + postcommit += 1 + pct_dis = 100 * post_dis / max(post_total, 1) + med = int(np.median(spreads)) if spreads else 0 + print(f"[diag {tag}] 多面顶点组 {n_multi} | 最终不一致 {n_incons} ({100*n_incons/max(n_multi,1):.1f}%)") + print(f" 后揭示token(同组已有commit) {post_total}: 不一致 {post_dis} ({pct_dis:.1f}%) <- 看到已定值仍预测错的比例") + print(f" 不一致组来源: 同步揭示分裂 {coreveal} / 后揭示才分裂 {postcommit} / 两者 {both} | commit step 跨度中位 {med}/{T}") + + analyze(f_pos, f_grp, "F", False) + analyze(i_pos, i_grp, "intra", True) + +# ========================================================================== +# 原 model_acc_fix3_tssr.py +# ========================================================================== + +_BF16_RESIDUAL = os.environ.get("BF16_RESIDUAL", "1") == "1" +class FactorizedTSSRTransformer(FactorizedDiffusionTransformerAltPC): + def __init__(self, *args, mode_embed_scale=1000.0, rope_mode="flat", + rope_theta_face=10000.0, rope_theta_group=10000.0, rope_theta_member=10000.0, **kwargs): + super().__init__(*args, **kwargs) + # ---- 层级 RoPE 开关(移植自 Tencent-TSSR MultiLevelRoPE, 适配本仓库 rotate_half 约定) ---- + # rope_mode="flat"(默认): 原行为, 逐 token 平位置 0..L-1, 与所有旧 ckpt 完全兼容。 + # rope_mode="hier" : head 维的 d/2 个频率槽切三段, 分别用三个结构分量旋转: + # face 段: p // tokens_per_face (第几个面; 长程) + # group 段: 组内三元组编号 (9/face 时 (p%9)//3: 0=F,1=intra,2=inter) + # member 段: 三元组内成员 (9/face 时 p%3: 顶点/边槽位, F0↔M00 同相位) + # hourglass 三个尺度(9/3/1 token per face)各有对应分解(对齐 TSSR 的 token/vertex/face 三级)。 + # 纯位置编码无参数 -> 任意旧 ckpt 可直接热启动, 但注意几何变了, loss 会先跳再收。 + self.rope_mode = rope_mode + self.rope_thetas = (float(rope_theta_face), float(rope_theta_group), float(rope_theta_member)) + d = self.d_model + # 模式 flag 嵌入:标量 mode∈{0,1} -> 正弦 -> MLP -> [b,d](仿 TSSR DiscreteValueEmbeding)。 + # scale 大让 0 与 1 的正弦特征清晰可分;末层零初始化使初始 mode token≈0(warm-start 不破坏骨干)。 + self.mode_emb = nn.Sequential( + SinusoidalPosEmb(d, scale=mode_embed_scale), + nn.Linear(d, d), nn.SiLU(), + nn.Linear(d, d), + ) + nn.init.zeros_(self.mode_emb[-1].weight) + nn.init.zeros_(self.mode_emb[-1].bias) + # 逐 token 正确性分类器(detach 骨干) + self.to_acc = nn.Sequential( + nn.Linear(d, d // 2), nn.SiLU(), + nn.Linear(d // 2, d // 4), nn.SiLU(), + nn.Linear(d // 4, 1), + ) + + def _hier_cos_sin(self, position_ids, tokens_per_face, dtype): + """层级 RoPE 的 cos/sin, 形状/布局与 RotaryEmbedding.forward_positions 一致([total, head_dim], + 频率复制两半以配 rotate_half), 因此 apply_rotary_pos_emb_varlen 原样可用。 + d/2 个频率槽按 TSSR 方式切 [n, n, d/2-2n](n=d//6), 依序分给 face/group/member 分量。""" + device = position_ids.device + d_half = self.rotary_emb.inv_freq.numel() # = head_dim // 2 + n = d_half // 3 + chunk_sizes = (n, n, d_half - 2 * n) + p = position_ids + if tokens_per_face == 9: + comps = (p // 9, (p % 9) // 3, p % 3) + elif tokens_per_face == 3: + comps = (p // 3, p % 3, torch.zeros_like(p)) + elif tokens_per_face == 1: + comps = (p, torch.zeros_like(p), torch.zeros_like(p)) + else: + raise ValueError(f"tokens_per_face 只支持 9/3/1, 收到 {tokens_per_face}") + rotations = [] + for size, theta, comp in zip(chunk_sizes, self.rope_thetas, comps): + freqs = 1.0 / (theta ** (torch.arange(size, device=device, dtype=torch.float32) / max(size, 1))) + rotations.append(comp.to(torch.float32)[:, None] * freqs[None, :]) + rot = torch.cat(rotations, dim=-1) # [total, d_half] + emb = torch.cat((rot, rot), dim=-1) # [total, head_dim], 同 rotate_half 布局 + # [total, 1, head_dim]: 中间加 head 广播维, 与 forward_positions 的返回形状一致 + # (q/k 是 [total, nhead, d_head], 少了这一维会把 total 广播到 nhead 上直接报形状错) + return emb.cos().to(dtype).unsqueeze(1), emb.sin().to(dtype).unsqueeze(1) + + # ------------------------------------------------------------------ # + # forward_varlen:复制 AltPC 版,改动 = ①mode token 拼点云 context ②rope_mode=hier 时层级 RoPE + # ------------------------------------------------------------------ # + def forward_varlen(self, coord_list, subcode_list, t, depth, pc=None, pc_target=None, mode=None): + device = t.device + coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list] + subcode_list = [s.to(device=device, dtype=torch.long) for s in subcode_list] + # 长度在 python 端本来就知道 -> 一并算出各尺度的 max_seqlen, 传给 attention 省掉 + # 每次 flash 调用的 int(lengths.max().item()) 同步(一次前向 72 次, 见 ModernAttention 说明)。 + lengths_py = [int(c.numel()) for c in coord_list] + if sum(lengths_py) == 0: + raise ValueError("forward_varlen received an empty batch.") + if any(n % REFINE_POS_PER_FACE != 0 for n in lengths_py): + raise ValueError("每个序列长度必须是 9 的倍数(每 face 9 个 refine 位置)。") + lengths = torch.tensor(lengths_py, device=device, dtype=torch.long) + max_len_9 = max(lengths_py) + max_len_3 = max_len_9 // 3 # hourglass: lengths // 3, 整除有断言保证 + max_len_1 = max_len_9 // 9 + + coord_packed = torch.cat(coord_list, dim=0) + subcode_packed = torch.cat(subcode_list, dim=0) + types = self.types_from_lengths(lengths) + + raw_emb = self.raw_embedding(coord_packed) + refine_emb = self.embed_refine(subcode_packed, types) + x = raw_emb + refine_emb + + # nn.Embedding 不在 autocast 的转型名单里 -> 上面的 x 是 fp32。不显式降下来的话: + # ① enc_blocks_9 / dec_blocks_9(经 skip_9 带回)这 4 个块的残差流全程 fp32 —— 而它们 + # 正是 token 数最多的一档(hourglass 把其余降到 1/3、1/9), RMSNorm/modulate/残差加 + # 在最大的激活上按 fp32 跑, 内存流量白翻一倍; + # ② rope_fn_9 是按 `dtype=x.dtype` 建 cos/sin 的, x 是 fp32 就把 bf16 的 Q/K 提升回 fp32, + # 导致 can_use_flash_attn 判否、掉进 SDPA 逐序列 python 循环(见 apply_rotary_pos_emb_varlen)。 + # 统一降到 autocast dtype, 让整条残差流和 RoPE 都是 bf16。 + if _BF16_RESIDUAL and torch.is_autocast_enabled(): + x = x.to(torch.get_autocast_dtype("cuda")) + + depth = depth.to(device=device, dtype=torch.long) + # c_sample **刻意保持 fp32**(按混精惯例): 它只有 [b, d_model], 保高精度的成本可忽略, 而它经 + # adaLN 变成 shift/scale/gate 逐元素乘加到所有激活上 —— 误差是系统性的、不会被平均掉。 + # DiT / SD3 / Flux 都把时间嵌入与调制留在 fp32。对比: 点云 latent 是 cross-attn 的 K/V, + # 紧接着就进 k_proj/v_proj 的 matmul(autocast 必然转 bf16), 且 flash-attn 不接受 fp32, 故 bf16。 + c_sample = self.time_mlp(t) + self.depth_embedding(depth) # [b, d_model], fp32 + + def c_for(cur_lengths): + return c_sample.repeat_interleave(cur_lengths, dim=0) + + bsz = len(coord_list) + # 模式 token(每样本一个),拼到点云 context 尾部(TSSR 式 flag 注入) + mode_tok = None + if mode is not None: + if torch.is_tensor(mode): + mode_vec = mode.to(device=device, dtype=torch.float).reshape(-1) + else: + mode_vec = torch.full((bsz,), float(mode), device=device) + mode_tok = self.mode_emb(mode_vec).unsqueeze(1) # [b,1,d] + + # ---- 点云条件:两套(全GT / 目标分辨率塌缩),逐层交替; 各拼上 mode token ---- + def _encode_cond(_pc): + if not self.use_pc: + return None, None, None + if _pc is None: + emb = self.null_cond.unsqueeze(0).expand(bsz, -1, -1) # [b,257,d] + else: + _pc = _pc.to(device=device) + emb = self.conditioner(_pc) # [b,257,d] + if mode_tok is not None: + emb = torch.cat([emb, mode_tok.to(emb.dtype)], dim=1) # [b,258,d] + ctx_len = int(emb.size(1)) # 每样本 context 长度恒为 257(+1 = mode token) + cl = torch.full((emb.size(0),), ctx_len, device=device, dtype=torch.long) + return emb.reshape(-1, emb.size(-1)), cl, ctx_len + + cond_a, cond_a_len, cond_max_len = _encode_cond(pc) # 奇数层(全GT) + if self.use_pc and (pc_target is not None): + cond_b, cond_b_len, _ = _encode_cond(pc_target) # 偶数层(目标分辨率塌缩) + else: + cond_b, cond_b_len = cond_a, cond_a_len # 未提供 -> 两套相同 + + def _cond_for(gi): + return (cond_a, cond_a_len) if (gi % 2 == 0) else (cond_b, cond_b_len) + + # geo/geo_vcat(2026-08-23, 接 sfocal 大模型 warm-start): 三尺度 cos/sin 由 coord 现算, + # 复用父类 GeoRoPE(锚点/频率梯子与主线逐位一致); hier/flat 走原路径不变。 + geo_scales = (self._geo_rope_scales(coord_packed, x.dtype) + if str(getattr(self, "rope_mode", "flat")).startswith("geo") else None) + + def rope_fn_for(cur_lengths, tokens_per_face=9): + if geo_scales is not None: + cos, sin = geo_scales[tokens_per_face] + elif self.rope_mode == "hier": + position_ids = lengths_to_position_ids(cur_lengths) + cos, sin = self._hier_cos_sin(position_ids, tokens_per_face, x.dtype) + else: + position_ids = lengths_to_position_ids(cur_lengths) + cos, sin = self.rotary_emb.forward_positions(position_ids, dtype=x.dtype) + return lambda q, k: apply_rotary_pos_emb_varlen(q, k, cos, sin) + + gi = 0 + rope_fn_9 = rope_fn_for(lengths, 9) + c9 = c_for(lengths) + for block in self.enc_blocks_9: + cnd, cnl = _cond_for(gi); gi += 1 + x = block.forward_varlen(x, c9, lengths, cnd, cnl, rope_apply_fn=rope_fn_9, + max_len=max_len_9, cond_max_len=cond_max_len) + skip_9 = x + + x, lengths = self.down_9_to_3.forward_varlen(x, lengths) + rope_fn_3 = rope_fn_for(lengths, 3) + c3 = c_for(lengths) + for block in self.enc_blocks_3: + cnd, cnl = _cond_for(gi); gi += 1 + x = block.forward_varlen(x, c3, lengths, cnd, cnl, rope_apply_fn=rope_fn_3, + max_len=max_len_3, cond_max_len=cond_max_len) + skip_3 = x + + x, lengths = self.down_3_to_1.forward_varlen(x, lengths) + rope_fn_1 = rope_fn_for(lengths, 1) + c1 = c_for(lengths) + for block in self.bottleneck_blocks: + cnd, cnl = _cond_for(gi); gi += 1 + x = block.forward_varlen(x, c1, lengths, cnd, cnl, rope_apply_fn=rope_fn_1, + max_len=max_len_1, cond_max_len=cond_max_len) + + x, lengths = self.up_1_to_3.forward_varlen(x, lengths) + x = x + skip_3 + for block in self.dec_blocks_3: + cnd, cnl = _cond_for(gi); gi += 1 + x = block.forward_varlen(x, c3, lengths, cnd, cnl, rope_apply_fn=rope_fn_3, + max_len=max_len_3, cond_max_len=cond_max_len) + + x, lengths = self.up_3_to_9.forward_varlen(x, lengths) + x = x + skip_9 + for block in self.dec_blocks_9: + cnd, cnl = _cond_for(gi); gi += 1 + x = block.forward_varlen(x, c9, lengths, cnd, cnl, rope_apply_fn=rope_fn_9, + max_len=max_len_9, cond_max_len=cond_max_len) + + return x, types + + # ------------------------------------------------------------------ # + # 分类器工具 + # ------------------------------------------------------------------ # + def to_acc_logit(self, h): + """逐 token 正确性 logit [total];detach 骨干,分类器梯度不回流。""" + return self.to_acc(h.detach()).squeeze(-1) + + @torch.no_grad() + def _pos_correct(self, h, target, types): + """逐位置正确性 label [total] bool:该 refine 位置的头 argmax(teacher-forced)是否==GT。 + F 比 slot0;intra/inter 比 8 行全等。""" + total = h.size(0) + correct = torch.zeros(total, dtype=torch.bool, device=h.device) + mf = types == TYPE_F + if mf.any(): + pred = self.f_head(h[mf]).argmax(-1) + correct[mf] = pred == target[mf, 0] + for ty, head in ((TYPE_INTRA, self.intra_head), (TYPE_INTER, self.inter_head)): + m = types == ty + if not m.any(): + continue + tgt = target[m] + pred = head.forward_teacher(h[m], tgt).argmax(-1) # [n,8] teacher-forced + correct[m] = (pred == tgt).all(dim=1) + return correct + + # ------------------------------------------------------------------ # + # 训练入口 + # ------------------------------------------------------------------ # + def forward(self, coord_list, noisy_subcode_list, t, depth, target_subcode_list=None, + loss_mask=None, pc=None, pc_target=None, task="mask", acc_wrong_weight=1.0): + """DDP 单一入口(必须走 model(...) 才能触发梯度同步;自定义方法会绕过 DDP)。 + task="mask" -> (loss, acc, acc_ne) 吸收腐蚀 CE(仅 masked 位置) + task="nomask" -> (loss_nomask, loss_acc, frac_correct, cls_acc) uniform 去噪 CE(全位置)+ 分类器 BCE + """ + if task == "mask": + return self.forward_mask(coord_list, noisy_subcode_list, t, depth, + target_subcode_list, loss_mask, pc, pc_target) + elif task == "nomask": + return self.forward_nomask(coord_list, noisy_subcode_list, t, depth, + target_subcode_list, pc, pc_target, acc_wrong_weight) + raise ValueError(f"unknown task {task!r}") + + + + + # ------------------------------------------------------------------ # + # 采样工具 + # ------------------------------------------------------------------ # + def _predict_fill(self, h, h_u, types, cfg_t, use_cfg, temperature): + """由 hidden 出各类型头 argmax(可 CFG/温度),填满 subcodes [total,8]。""" + total = h.size(0) + pred = torch.zeros((total, NUM_OCTANTS), dtype=torch.long, device=h.device) + mf = types == TYPE_F + if mf.any(): + logits = self.f_head(h[mf]) + if use_cfg: + lu = self.f_head(h_u[mf]); logits = lu + cfg_t * (logits - lu) + if temperature and temperature > 0: + prob = (logits / temperature).softmax(-1) + am = torch.multinomial(prob.float(), 1).squeeze(-1) + else: + am = logits.argmax(-1) + pred[mf, :] = am.unsqueeze(1).expand(-1, NUM_OCTANTS) + for ty, head in ((TYPE_INTRA, self.intra_head), (TYPE_INTER, self.inter_head)): + m = types == ty + if not m.any(): + continue + p, _ = head.rollout(h[m], h_u[m] if use_cfg else None, cfg_t, temperature) + pred[m] = p + return pred + + @torch.no_grad() + def predict_subcodes(self, coord_list, sub_list, t, depth, pc=None, pc_target=None, mode=1, + cfg_scale=1.0, temperature=0.0): + """一次前向 + 头 argmax 得填满的 subcodes [total,8](供自生成脏输入 / 单步预测)。""" + lengths = torch.tensor([c.numel() for c in coord_list], device=t.device, dtype=torch.long) + types = self.types_from_lengths(lengths) + use_cfg = (float(cfg_scale) != 1.0) and self.use_pc and (pc is not None) + h, _ = self.forward_varlen(coord_list, sub_list, t, depth, pc, pc_target, mode=mode) + h_u = self.forward_varlen(coord_list, sub_list, t, depth, None, None, mode=mode)[0] if use_cfg else None + return self._predict_fill(h, h_u, types, float(cfg_scale), use_cfg, temperature) + + @torch.no_grad() + def sample_tssr(self, coord_list, depth, pc=None, pc_target=None, T_infer=10, + cfg_scale=1.0, cfg_schedule="constant", cfg_start=1.0, + mask_logits_lo=0.2, mask_logits_hi=0.9, temperature=0.0, + generator=None, return_stats=False): + """TSSR 三前向解码。返回 list[[L,8]] local 子码(drop-in 父类 decode 路径)。 + 每步:① mask 预测(mode=1)填满 -> ② no-mask 精修(mode=0)+ 同前向出置信 -> ③ re-mask 判错位置。 + 末步不再 re-mask。""" + self.eval() + device = next(self.parameters()).device + coord_list = [c.to(device=device, dtype=torch.long) for c in coord_list] + depth = depth.to(device=device, dtype=torch.long) + lengths = torch.tensor([c.numel() for c in coord_list], device=device, dtype=torch.long) + types = self.types_from_lengths(lengths) + want_cfg = (float(cfg_scale) != 1.0 or float(cfg_start) != 1.0) and self.use_pc and (pc is not None) + total = int(lengths.sum().item()) + + subcodes = torch.full((total, NUM_OCTANTS), MASK_ID, dtype=torch.long, device=device) + sched = torch.linspace(float(mask_logits_lo), float(mask_logits_hi), max(T_infer, 1), device=device) + n_remask_hist = [] + + for step in range(T_infer): + if cfg_schedule == "linear": + cfg_t = float(cfg_start) + (float(cfg_scale) - float(cfg_start)) * step / max(T_infer - 1, 1) + else: + cfg_t = float(cfg_scale) + use_cfg = want_cfg and (cfg_t != 1.0) + t = torch.full((len(coord_list),), step / T_infer, device=device, dtype=torch.float) + + # ① mask 预测(mode=1) + sub_list = list(torch.split(subcodes, lengths.tolist(), dim=0)) + h1, _ = self.forward_varlen(coord_list, sub_list, t, depth, pc, pc_target, mode=1) + h1u = self.forward_varlen(coord_list, sub_list, t, depth, None, None, mode=1)[0] if use_cfg else None + x_t = self._predict_fill(h1, h1u, types, cfg_t, use_cfg, temperature) + + # ② no-mask 精修(mode=0)+ 复用该前向出逐 token 置信 + sub_list2 = list(torch.split(x_t, lengths.tolist(), dim=0)) + h2, _ = self.forward_varlen(coord_list, sub_list2, t, depth, pc, pc_target, mode=0) + h2u = self.forward_varlen(coord_list, sub_list2, t, depth, None, None, mode=0)[0] if use_cfg else None + x_t = self._predict_fill(h2, h2u, types, cfg_t, use_cfg, temperature) + acc = torch.sigmoid(self.to_acc(h2).squeeze(-1)).float() # [total] ∈[0,1] + + # ③ re-mask 判错位置(末步不做) + if step < T_infer - 1: + mask_prob = ((1.0 - acc) * sched[step]).clamp(1e-6, 1.0 - 1e-6) + if generator is not None: + r = torch.rand(total, device=device, generator=generator) + else: + r = torch.rand(total, device=device) + remask = r < mask_prob + subcodes = x_t.clone() + subcodes[remask] = MASK_ID + n_remask_hist.append(int(remask.sum().item())) + else: + subcodes = x_t + + out = list(torch.split(subcodes, lengths.tolist(), dim=0)) + if return_stats: + return out, {"n_remask_per_step": n_remask_hist, "total_pos": total} + return out diff --git a/mulres_infer/models/point_conditioner.py b/mulres_infer/models/point_conditioner.py new file mode 100644 index 0000000000000000000000000000000000000000..2ca5f355d62f3e7d5e2db52b6fbd0326116650ec --- /dev/null +++ b/mulres_infer/models/point_conditioner.py @@ -0,0 +1,200 @@ +import torch +from torch import nn +from beartype import beartype +from modules.miche.encode import load_model +import os + +_PKG_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +# helper functions + +def exists(val): + return val is not None + +def default(*values): + for value in values: + if exists(value): + return value + return None + + +# point-cloud encoder from Michelangelo +@beartype +class PointConditioner(torch.nn.Module): + def __init__( + self, + *, + dim_latent = None, + model_name = 'miche-256-feature', + cond_dim = 768, + feature_dim = 768, + freeze = True, + img_feat_dim: int = 0, + img_feat_extractor=None, + num_latents = None, # None=预训练默认(256); >256 时把可学习 query 平铺+微噪扩容(预训练权重原样热启动) + ): + super().__init__() + + # open-source version of miche + if model_name == 'miche-256-feature': + ckpt_path = os.environ.get("SHAPEVAE_CKPT", + os.path.join(_PKG_ROOT, "ckpt", "shapevae-256.ckpt")) + if not os.path.exists(ckpt_path): + ckpt_path=None + print('[WARNING] Michelangelo ckpt not exist, please check if you are training') + config_path = os.path.join(_PKG_ROOT, 'modules', 'miche', 'shapevae-256.yaml') + + self.feature_dim = feature_dim # embedding dimension + self.cond_length = 257 # length of embedding + self.point_encoder = load_model(ckpt_path=ckpt_path, config_path=config_path) + if num_latents: + self._expand_latents(int(num_latents)) + + # additional layers to connect miche and GPT + self.cond_head_proj = nn.Linear(cond_dim, self.feature_dim) + self.cond_proj = nn.Linear(cond_dim, self.feature_dim) + + else: + raise NotImplementedError + + # whether to finetuen point-cloud encoder + if freeze: + for parameter in self.point_encoder.parameters(): + parameter.requires_grad = False + + self.freeze = freeze + self.model_name = model_name + self.dim_latent = self.feature_dim + + # ---------- image-conditioning extension ---------- + self.img_feat_dim = int(img_feat_dim or 0) + self.img_feat_extractor = img_feat_extractor + if self.img_feat_dim > 0: + self._widen_input_proj(self.img_feat_dim) + + self.register_buffer('_device_param', torch.tensor(0.), persistent = False) + + + @property + def device(self): + return next(self.buffers()).device + + def _expand_latents(self, num_latents: int): + """把 perceiver 可学习 query 扩到 [1头 + num_latents body]: 预训练 query 是 257=1+256(行0=头token), + 头保持 1 个, body 256 平铺复制 + 1% std 微噪打破对称(首份无噪)。attention/n_ctx 均动态读 shape, + 无其他结构依赖(SDPA 补丁同样形状无关)。训练热启动时上层 ckpt(若已含扩容 query)会原样覆盖。""" + enc = self.point_encoder.model.shape_model.encoder + old_q = enc.query.data # [1+B, w] + body = old_q[1:] + if body.shape[0] == num_latents: + self.cond_length = 1 + num_latents + return + assert num_latents % body.shape[0] == 0, f"num_latents={num_latents} 须为 body={body.shape[0]} 的整数倍" + reps = num_latents // body.shape[0] + new_body = body.repeat(reps, 1) + noise = torch.randn_like(new_body) * body.std() * 0.01 + noise[: body.shape[0]] = 0 # 首份 body 保持原 query 精确不动 + enc.query = nn.Parameter(torch.cat([old_q[:1], new_body + noise], dim=0)) + enc.num_latents = 1 + num_latents + self.cond_length = 1 + num_latents + print(f"[PointConditioner] query 扩容 1+{body.shape[0]} -> 1+{num_latents} (body tile+1%噪, 首份无噪)") + + def _widen_input_proj(self, img_feat_dim: int): + """Replace Michelangelo's input_proj Linear with a wider one and zero-init + the new input columns so loading the pretrained ckpt doesn't change init + behavior (image features get multiplied by zero on construction).""" + enc = self.point_encoder.model.shape_model.encoder + if getattr(enc, "img_feat_dim", 0) == img_feat_dim and enc.input_proj.in_features == ( + enc.fourier_embedder.out_dim + enc.point_feats + img_feat_dim + ): + return # already widened (e.g. resume from a widened ckpt) + old = enc.input_proj + new_in = enc.fourier_embedder.out_dim + enc.point_feats + img_feat_dim + if new_in == old.in_features: + enc.img_feat_dim = img_feat_dim + return + new_proj = nn.Linear( + new_in, + old.out_features, + bias=old.bias is not None, + device=old.weight.device, + dtype=old.weight.dtype, + ) + with torch.no_grad(): + new_proj.weight.zero_() + new_proj.weight[:, : old.in_features].copy_(old.weight) + if old.bias is not None: + new_proj.bias.copy_(old.bias) + enc.input_proj = new_proj + enc.img_feat_dim = img_feat_dim + if self.freeze: + # input_proj is part of the shape encoder, which we keep frozen by + # default; only the new image-feature columns ever produce nonzero + # gradients on top of frozen weights. If callers want to fine-tune + # the whole input_proj they can flip freeze=False. + for p in new_proj.parameters(): + p.requires_grad = False + + + def embed_pc( + self, + pc_normal, + image=None, + transform_matrix=None, + camera_angle_x=None, + depth_map=None, + depth_valid=None, + img_feats=None, + ): + # encode point cloud to embeddings + if self.model_name == 'miche-256-feature': + if self.img_feat_dim > 0: + if img_feats is None: + if self.img_feat_extractor is None: + raise RuntimeError( + "img_feat_dim>0 but no img_feats provided and no extractor configured." + ) + if image is None or transform_matrix is None or camera_angle_x is None: + raise ValueError( + "image, transform_matrix, camera_angle_x are required when img_feat_extractor is set." + ) + img_feats = self.img_feat_extractor( + image, pc_normal[..., :3], transform_matrix, camera_angle_x, + depth_map=depth_map, depth_valid=depth_valid, + ) + img_feats = img_feats.to(dtype=pc_normal.dtype) + point_feature = self.point_encoder.encode_latents(pc_normal, img_feats=img_feats) + else: + point_feature = self.point_encoder.encode_latents(pc_normal) + pc_embed_head = self.cond_head_proj(point_feature[:, 0:1]) + pc_embed = self.cond_proj(point_feature[:, 1:]) + pc_embed = torch.cat([pc_embed_head, pc_embed], dim=1) + + return pc_embed + + + def forward( + self, + pc = None, + pc_embeds = None, + image=None, + transform_matrix=None, + camera_angle_x=None, + depth_map=None, + depth_valid=None, + img_feats=None, + ): + if pc_embeds is None: + target_dtype = next(self.buffers()).dtype + pc_embeds = self.embed_pc( + pc.to(target_dtype), + image=image, + transform_matrix=transform_matrix, + camera_angle_x=camera_angle_x, + depth_map=depth_map, + depth_valid=depth_valid, + img_feats=img_feats, + ) + + assert not torch.any(torch.isnan(pc_embeds)), 'NAN values in pc embedings' + + return pc_embeds diff --git a/mulres_infer/models/subdivider.py b/mulres_infer/models/subdivider.py new file mode 100644 index 0000000000000000000000000000000000000000..c3a4f40c6e672fe35c32a801d1ea8ae4e3495c9d --- /dev/null +++ b/mulres_infer/models/subdivider.py @@ -0,0 +1,1188 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +from einops import rearrange +from modules.attention import MLA +from .point_conditioner import PointConditioner + +try: + from flash_attn import flash_attn_varlen_func +except Exception: + try: + from flash_attn.flash_attn_interface import flash_attn_varlen_func + except Exception: + flash_attn_varlen_func = None + +class SinusoidalPosEmb(nn.Module): + def __init__(self, dim, scale=100.0): + super().__init__() + self.dim = dim + self.scale = scale + + def forward(self, x): + device = x.device + dtype = x.dtype + x = x.float() * self.scale + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb.to(dtype=dtype) + +class RMSNorm(nn.Module): + """现代 LLM 标配：去均值的 RMSNorm""" + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + # 只需要学习一个缩放参数，去掉 bias + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + # 计算均方根 (Root Mean Square) + norm_x = torch.mean(x.pow(2), dim=-1, keepdim=True) + x_normed = x * torch.rsqrt(norm_x + self.eps) + return self.weight * x_normed + +class SwiGLU(nn.Module): + """现代 LLM 标配：SwiGLU 激活的门控 FFN""" + def __init__(self, d_model, hidden_dim=None): + super().__init__() + # 将原先 4 * d_model 的维度缩小到 8/3，保持参数量对等; hidden_dim 可显式覆盖 + # (如 1024 宽度时 8/3 给 2730 非 256 倍数, 显式取 2816 对齐 Tensor Core) + hidden_dim = int((8 / 3) * d_model) if hidden_dim is None else int(hidden_dim) + # 提示：在实际工程中，通常会把 hidden_dim 向上取整到 128 或 256 的倍数，以最大化 Tensor Core 效率 + + self.w1 = nn.Linear(d_model, hidden_dim, bias=False) + self.w2 = nn.Linear(d_model, hidden_dim, bias=False) + self.w3 = nn.Linear(hidden_dim, d_model, bias=False) + + def forward(self, x): + # 公式: Swish(xW1) ⊗ (xW2) -> W3 + # PyTorch 中 F.silu 就是 Swish + return self.w3(F.silu(self.w1(x)) * self.w2(x)) + +class ModernAttention(nn.Module): + def __init__(self, d_model, nhead, dropout=0., qk_norm=False): + super().__init__() + assert d_model % nhead == 0, "d_model 必须能被 nhead 整除" + self.nhead = nhead + self.d_head = d_model // nhead + + # 现代架构通常在 Attention 的投影层中也丢弃 bias + self.q_proj = nn.Linear(d_model, d_model, bias=False) + self.k_proj = nn.Linear(d_model, d_model, bias=False) + self.v_proj = nn.Linear(d_model, d_model, bias=False) + self.out_proj = nn.Linear(d_model, d_model, bias=False) + self.dropout = dropout + # QK-norm: 逐 head 对 Q/K 在 head_dim 上做 RMSNorm(在 RoPE 之前),稳定 attention logits、 + # 防幅度爆炸。默认关(不改变既有模型);flow 等新模型可开。新增参数 -> 旧 ckpt load_strict=False。 + self.qk_norm = qk_norm + if qk_norm: + self.q_norm = RMSNorm(self.d_head) + self.k_norm = RMSNorm(self.d_head) + + def forward(self, q, k, v, mask=None, is_causal=False, rope_apply_fn=None): + B, L_q, _ = q.size() + _, L_k, _ = k.size() + + # 投影并调整形状为 [Batch, Head, SeqLen, HeadDim] + Q = self.q_proj(q).view(B, L_q, self.nhead, self.d_head).transpose(1, 2) + K = self.k_proj(k).view(B, L_k, self.nhead, self.d_head).transpose(1, 2) + V = self.v_proj(v).view(B, L_k, self.nhead, self.d_head).transpose(1, 2) + + if self.qk_norm: # QK-norm(RoPE 前),保持 dtype 供 SDPA/flash + Q = self.q_norm(Q).to(Q.dtype); K = self.k_norm(K).to(K.dtype) + + # 注入旋转位置编码 (RoPE) + if rope_apply_fn is not None: + Q, K = rope_apply_fn(Q, K) + + # 使用 PyTorch 2.0+ 内置的 SDPA，它会自动调用 Flash Attention 节省巨量显存 + out = F.scaled_dot_product_attention( + Q, K, V, + attn_mask=mask, + dropout_p=self.dropout if self.training else 0.0, + is_causal=is_causal + ) + + # 恢复形状并输出 + out = out.transpose(1, 2).contiguous().view(B, L_q, -1) + return self.out_proj(out) + + def forward_varlen(self, q, k, v, lengths_q, lengths_k=None, is_causal=False, rope_apply_fn=None, + max_len_q=None, max_len_k=None): + """max_len_q / max_len_k: flash 需要的 max_seqlen。**传 python int 可避免 GPU->CPU 同步**。 + + 不传时回退到 `int(lengths.max().item())` —— 那是一次 cudaStreamSynchronize: CPU 必须等 + GPU 把已排队的活干完才能读到值, 于是无法提前排队 kernel。本模型一次前向有 36 次 attention + 调用(18 self + 18 cross), 即 72 次同步, 是「kernel 又多又小、GPU 有空隙」的直接成因, + 也让 torch.compile 在这里 graph break(数据依赖的 unbacked symint)。 + 调用方本来就知道这些长度(lengths 是从 [c.numel() for c in coord_list] 建的), 应当传进来。 + """ + if lengths_k is None: + lengths_k = lengths_q + if max_len_k is None: + max_len_k = max_len_q + + total_q = q.size(0) + total_k = k.size(0) + Q = self.q_proj(q).view(total_q, self.nhead, self.d_head) + K = self.k_proj(k).view(total_k, self.nhead, self.d_head) + V = self.v_proj(v).view(total_k, self.nhead, self.d_head) + + if self.qk_norm: # QK-norm(RoPE 前),.to 保持 dtype 供 flash-attn + Q = self.q_norm(Q).to(Q.dtype); K = self.k_norm(K).to(K.dtype) + + if rope_apply_fn is not None: + Q, K = rope_apply_fn(Q, K) + + dropout_p = self.dropout if self.training else 0.0 + if can_use_flash_attn(Q): + out = flash_attn_varlen_func( + Q.contiguous(), + K.contiguous(), + V.contiguous(), + lengths_to_cu_seqlens(lengths_q), + lengths_to_cu_seqlens(lengths_k), + int(lengths_q.max().item()) if max_len_q is None else int(max_len_q), + int(lengths_k.max().item()) if max_len_k is None else int(max_len_k), + dropout_p=dropout_p, + causal=is_causal, + ) + else: + out_chunks = [] + q_start = 0 + k_start = 0 + for q_len, k_len in zip(lengths_q.tolist(), lengths_k.tolist()): + q_end = q_start + q_len + k_end = k_start + k_len + q_i = Q[q_start:q_end].transpose(0, 1).unsqueeze(0) + k_i = K[k_start:k_end].transpose(0, 1).unsqueeze(0) + v_i = V[k_start:k_end].transpose(0, 1).unsqueeze(0) + out_i = F.scaled_dot_product_attention( + q_i, + k_i, + v_i, + dropout_p=dropout_p, + is_causal=is_causal, + ) + out_chunks.append(out_i.squeeze(0).transpose(0, 1)) + q_start = q_end + k_start = k_end + out = torch.cat(out_chunks, dim=0) + + return self.out_proj(out.contiguous().view(total_q, -1)) + +# class CrossAttnTransformerBlock(nn.Module): +# def __init__(self, d_model, nhead, dropout=0.): +# super().__init__() +# self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True) +# self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True) + +# self.norm1 = nn.LayerNorm(d_model) +# self.norm2 = nn.LayerNorm(d_model) +# self.norm3 = nn.LayerNorm(d_model) + +# self.ffn = nn.Sequential( +# nn.Linear(d_model, d_model * 4), +# nn.GELU(), +# nn.Dropout(dropout), +# nn.Linear(d_model * 4, d_model), +# ) +# self.dropout = nn.Dropout(dropout) + +# def forward(self, x, cond): +# # 1) token self-attention +# h = self.norm1(x) +# sa_out, _ = self.self_attn(h, h, h, need_weights=False) +# x = x + self.dropout(sa_out) + +# # 2) cross-attention: query from current tokens, key/value from condition tokens +# h = self.norm2(x) +# ca_out, _ = self.cross_attn(h, cond, cond, need_weights=False) +# x = x + self.dropout(ca_out) + +# # 3) feed-forward +# h = self.norm3(x) +# ff_out = self.ffn(h) +# x = x + self.dropout(ff_out) +# return x +class ModernCrossAttnTransformerBlock(nn.Module): + def __init__(self, d_model, nhead, dropout=0.): + super().__init__() + + # 替换为支持 Flash Attention 的现代 Attention + self.self_attn = ModernAttention(d_model, nhead, dropout=dropout) + self.cross_attn = ModernAttention(d_model, nhead, dropout=dropout) + + # 替换 LayerNorm 为 RMSNorm + self.norm1 = RMSNorm(d_model) + self.norm2 = RMSNorm(d_model) + self.norm3 = RMSNorm(d_model) + + # 替换传统 FFN 为 SwiGLU + self.ffn = SwiGLU(d_model) + + self.dropout = nn.Dropout(dropout) + + def forward(self, x, cond, rope_apply_fn=None, causal_mask=False): + # 1) token self-attention + # 你的 Pre-Norm 逻辑保持不变，这是正确的 + h = self.norm1(x) + # 注意：在这里把 RoPE 传进去处理 Query 和 Key + sa_out = self.self_attn(h, h, h, is_causal=causal_mask, rope_apply_fn=rope_apply_fn) + x = x + self.dropout(sa_out) + + # 2) cross-attention: query from current tokens, key/value from condition tokens + if cond is not None: + h = self.norm2(x) + # Cross Attention 一般不使用 RoPE，因为输入源不同 + ca_out = self.cross_attn(h, cond, cond) + x = x + self.dropout(ca_out) + + # 3) feed-forward (SwiGLU) + h = self.norm3(x) + ff_out = self.ffn(h) + x = x + self.dropout(ff_out) + + return x + + def forward_varlen(self, x, lengths, cond=None, cond_lengths=None, rope_apply_fn=None, causal_mask=False): + h = self.norm1(x) + sa_out = self.self_attn.forward_varlen( + h, + h, + h, + lengths, + is_causal=causal_mask, + rope_apply_fn=rope_apply_fn, + ) + x = x + self.dropout(sa_out) + + if cond is not None: + h = self.norm2(x) + ca_out = self.cross_attn.forward_varlen( + h, + cond, + cond, + lengths, + lengths_k=cond_lengths, + ) + x = x + self.dropout(ca_out) + + h = self.norm3(x) + ff_out = self.ffn(h) + x = x + self.dropout(ff_out) + return x + +def rotate_half(x): + """ + 旋转特征的一半维度。 + 对于向量 [x1, x2, ..., x_d/2, x_d/2+1, ..., x_d]， + 将其转换为 [-x_d/2+1, ..., -x_d, x1, x2, ..., x_d/2]。 + 这在数学上等价于复数平面上的乘以 i。 + """ + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + +def apply_rotary_pos_emb(q, k, cos, sin): + """ + 将预计算好的 cos 和 sin 注入到 Query 和 Key 中。 + q, k 形状期望为: (Batch, Num_Heads, Seq_Len, Head_Dim) + cos, sin 形状期望为: (1, 1, Seq_Len, Head_Dim) + """ + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + +def apply_rotary_pos_emb_varlen(q, k, cos, sin): + """RoPE。**cos/sin 先转成 q/k 的 dtype**，保证输出 dtype 不被 cos/sin 带跑。 + + 为什么必须这么写：cos/sin 由 `rotary_emb.forward_positions(..., dtype=x.dtype)` 生成，而 + 9 尺度的 rope_fn 是在 `x = raw_emb + refine_emb` 之后立刻构造的 —— `nn.Embedding` 不在 autocast + 的转型名单里，那时 x 还是 **fp32**，于是 cos/sin 是 fp32。bf16 的 Q/K 乘上 fp32 的 cos/sin 会 + 类型提升成 fp32，`can_use_flash_attn` 判否，直接掉进 `ModernAttention.forward_varlen` 里 + **SDPA 的逐序列 python 循环**。命中的恰好是 enc_blocks_9 / dec_blocks_9 这 4 层, + 也就是跑在满 token 长度(其余层已被 hourglass 降到 1/3、1/9)上最贵的 4 个注意力。 + 3 尺度与 bottleneck 的 rope_fn 因为构造时 x 已是 bf16，本来就走 flash —— 即同一网络里 + 4 层 fp32 慢路径 + 14 层 bf16 快路径，不一致。统一到 bf16。 + """ + cos_q, sin_q = cos.to(q.dtype), sin.to(q.dtype) + q_embed = (q * cos_q) + (rotate_half(q) * sin_q) + cos_k, sin_k = (cos_q, sin_q) if k.dtype == q.dtype else (cos.to(k.dtype), sin.to(k.dtype)) + k_embed = (k * cos_k) + (rotate_half(k) * sin_k) + return q_embed, k_embed + +def lengths_to_cu_seqlens(lengths): + lengths = lengths.to(dtype=torch.int32) + prefix = torch.empty(lengths.numel() + 1, device=lengths.device, dtype=torch.int32) + prefix[0] = 0 + prefix[1:] = torch.cumsum(lengths, dim=0, dtype=torch.int32) + return prefix.contiguous() + +def lengths_to_position_ids(lengths): + total = int(lengths.sum().item()) + if total == 0: + return torch.empty(0, device=lengths.device, dtype=torch.long) + starts = torch.cumsum(lengths, dim=0) - lengths + return torch.arange(total, device=lengths.device, dtype=torch.long) - torch.repeat_interleave(starts, lengths) + +def can_use_flash_attn(x): + return ( + flash_attn_varlen_func is not None + and x.is_cuda + and x.dtype in (torch.float16, torch.bfloat16) + ) + +class RotaryEmbedding(nn.Module): + """ + 预计算和缓存 RoPE 的 cos 和 sin 矩阵。 + """ + def __init__(self, dim, max_position_embeddings=90000, base=10000.0, device=None): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + + # 计算旋转频率: theta_i = 10000^(-2(i-1)/d) + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # 预先构建缓存 + self._set_cos_sin_cache(seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32) + + # freqs 形状: (seq_len, dim/2) + freqs = torch.outer(t, self.inv_freq) + + # 将频率在特征维度上复制一份，形状变为: (seq_len, dim) + # 前半部分和后半部分的频率一致，为了和 rotate_half 的逻辑对齐 + emb = torch.cat((freqs, freqs), dim=-1) + + # 增加 Batch 和 Head 的广播维度: (1, 1, Seq_Len, Head_Dim) + self.register_buffer("cos_cached", emb.cos().to(dtype).unsqueeze(0).unsqueeze(0), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype).unsqueeze(0).unsqueeze(0), persistent=False) + + def forward(self, x, seq_len=None): + """ + x 形状: (Batch, Num_Heads, Seq_Len, Head_Dim) + """ + # 如果当前序列长度超过了预计算的缓存，则动态扩展缓存 + if seq_len > self.max_seq_len_cached: + self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) + + # 截取当前需要的序列长度返回 + return ( + self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype), + self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype), + ) + + def forward_positions(self, position_ids, dtype=None): + if dtype is None: + dtype = self.cos_cached.dtype + if position_ids.numel() == 0: + empty = torch.empty(0, 1, self.dim, device=position_ids.device, dtype=dtype) + return empty, empty + seq_len = int(position_ids.max().item()) + 1 + if seq_len > self.max_seq_len_cached: + self._set_cos_sin_cache(seq_len=seq_len, device=position_ids.device, dtype=dtype) + cos = self.cos_cached[0, 0, position_ids, :].to(device=position_ids.device, dtype=dtype) + sin = self.sin_cached[0, 0, position_ids, :].to(device=position_ids.device, dtype=dtype) + return cos.unsqueeze(1), sin.unsqueeze(1) + +class LinearDownsample(nn.Module): + def __init__(self, dim, shorten_factor): + super().__init__() + self.proj = nn.Linear(dim * shorten_factor, dim) + self.shorten_factor = shorten_factor + + def forward(self, x): + x = rearrange(x, 'b (n s) d -> b n (s d)', s=self.shorten_factor) + return self.proj(x) + + def forward_varlen(self, x, lengths): + if torch.any(lengths % self.shorten_factor != 0): + raise ValueError(f"All sequence lengths must be divisible by {self.shorten_factor} for hourglass downsample.") + x = rearrange(x, '(n s) d -> n (s d)', s=self.shorten_factor) + return self.proj(x), lengths // self.shorten_factor + +class LinearUpsample(nn.Module): + def __init__(self, dim, shorten_factor): + super().__init__() + self.proj = nn.Linear(dim, dim * shorten_factor) + self.shorten_factor = shorten_factor + + def forward(self, x): + x = self.proj(x) + return rearrange(x, 'b n (s d) -> b (n s) d', s=self.shorten_factor) + + def forward_varlen(self, x, lengths): + x = self.proj(x) + x = rearrange(x, 'n (s d) -> (n s) d', s=self.shorten_factor) + return x, lengths * self.shorten_factor + +class DiscreteDiffusionTransformer(nn.Module): + def __init__(self, + vocab_A, + vocab_B, + max_seq_len, + d_model=768, + nhead=12, + depth = [4,4,8], + dropout=0., + encoder_freeze=True, + encoder_name='miche-256-feature', + use_pc=True, + time_embed_scale=100.0): + super().__init__() + self.vocab_A = vocab_A + self.vocab_B = vocab_B + self.d_model = d_model + + # MASK Token 的 ID，我们将其设为 B（因为真实的 refine_token 是 0 到 B-1） + self.mask_token_id = vocab_B + + # 1. 词嵌入 (Embeddings) + self.raw_embedding = nn.Embedding(vocab_A, d_model) + # refine_token 的词表大小是 B + 1 (包含了 MASK token) + self.refine_embedding = nn.Embedding(vocab_B + 1, d_model) + + # 2. 位置编码 (Positional Encoding) + # self.pos_embedding = nn.Embedding(max_seq_len, d_model) + + # 3. 时间步编码 (Timestep Embedding) + self.time_mlp = nn.Sequential( + SinusoidalPosEmb(d_model, scale=time_embed_scale), + nn.Linear(d_model, d_model * 4), + nn.GELU(), + nn.Linear(d_model * 4, d_model) + ) + + # 4. 主干网络 (Denoising Backbone) + # U-Net/Hourglass 主干: + # enc(9) -> down -> enc(3) -> down -> bottleneck(1) -> up+skip -> dec(3) -> up+skip -> dec(9) + self.depth = depth + + self.enc_blocks_9 = nn.ModuleList( + [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[0])] + ) + self.enc_blocks_3 = nn.ModuleList( + [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[1])] + ) + self.bottleneck_blocks = nn.ModuleList( + [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[2])] + ) + self.dec_blocks_3 = nn.ModuleList( + [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[1])] + ) + self.dec_blocks_9 = nn.ModuleList( + [ModernCrossAttnTransformerBlock(d_model, nhead, dropout=dropout) for _ in range(depth[0])] + ) + + self.down_9_to_3 = LinearDownsample(d_model, shorten_factor=3) + self.down_3_to_1 = LinearDownsample(d_model, shorten_factor=3) + self.up_1_to_3 = LinearUpsample(d_model, shorten_factor=3) + self.up_3_to_9 = LinearUpsample(d_model, shorten_factor=3) + + self.rotary_emb = RotaryEmbedding(dim=d_model//nhead, max_position_embeddings=max_seq_len) + # 5. 输出预测头 (只预测真实的 B 个 Token) + self.lm_head = nn.Linear(d_model, vocab_B) + self.use_pc = use_pc + if use_pc: + print(f'Point cloud encoder: {encoder_name} | freeze: {encoder_freeze}') + self.conditioner = PointConditioner(model_name=encoder_name, freeze=encoder_freeze, feature_dim=d_model) + # Learnable null condition for classifier-free guidance. + # Replaces self.conditioner(pc) when pc is None, so the model can learn + # an "unconditional" representation during cond-drop fine-tuning. + self.null_cond_len = 257 # matches PointConditioner.cond_length + self.null_cond = nn.Parameter(torch.randn(self.null_cond_len, d_model) * 0.02) + + def forward(self, raw_tokens, noisy_refine_tokens, t, pc=None): + """ + raw_tokens: [b, n], 取值 0~A-1 + noisy_refine_tokens: [b, n], 取值 0~B (B是MASK) + t: [b], 当前的时间步 + pc: [b, n, 6], 点云数据 + """ + if isinstance(raw_tokens, (list, tuple)): + return self.forward_varlen(raw_tokens, noisy_refine_tokens, t, pc) + + b, n = raw_tokens.shape + device = raw_tokens.device + + # 获取 Embedding + raw_emb = self.raw_embedding(raw_tokens) # [b, n, d_model] + refine_emb = self.refine_embedding(noisy_refine_tokens) # [b, n, d_model] + + # 获取位置和时间 Embedding + time_emb = self.time_mlp(t).unsqueeze(1) # [b, 1, d_model] + + # x 作为被去噪序列，cond 作为每层 cross-attn 的条件序列 + x = raw_emb + refine_emb + time_emb + + if self.use_pc: + if pc is None: + # Classifier-free guidance / cond-drop: use learnable null condition. + cond = self.null_cond.unsqueeze(0).expand(b, -1, -1) + else: + cond = self.conditioner(pc) # [b, 257, d_model] + else: + cond = None + + if n % 9 != 0: + raise ValueError(f"Sequence length {n} is not compatible with 9->3->1 hourglass compression.") + + dummy_tensor = torch.empty(1, device=x.device, dtype=torch.float32) + + def rope_fn_for(seq_len): + cos, sin = self.rotary_emb(dummy_tensor, seq_len=seq_len) + return lambda q, k: apply_rotary_pos_emb(q, k, cos, sin) + + # Encoder @ length 9 + rope_fn_9 = rope_fn_for(x.shape[1]) + for block in self.enc_blocks_9: + x = block(x, cond, rope_apply_fn=rope_fn_9) + skip_9 = x + + # Downsample 9 -> 3 + x = self.down_9_to_3(x) + + # Encoder @ length 3 + rope_fn_3 = rope_fn_for(x.shape[1]) + for block in self.enc_blocks_3: + x = block(x, cond, rope_apply_fn=rope_fn_3) + skip_3 = x + + # Downsample 3 -> 1 (or in general: n/3 -> n/9) + x = self.down_3_to_1(x) + + # Bottleneck + rope_fn_1 = rope_fn_for(x.shape[1]) + for block in self.bottleneck_blocks: + x = block(x, cond, rope_apply_fn=rope_fn_1) + + # Upsample 1 -> 3 and long skip connection + x = self.up_1_to_3(x) + x = x + skip_3 + + # Decoder @ length 3 + for block in self.dec_blocks_3: + x = block(x, cond, rope_apply_fn=rope_fn_3) + + # Upsample 3 -> 9 and long skip connection + x = self.up_3_to_9(x) + x = x + skip_9 + + # Decoder @ length 9 + for block in self.dec_blocks_9: + x = block(x, cond, rope_apply_fn=rope_fn_9) + + # 预测 Logits + logits = self.lm_head(x) # [b, n, vocab_B] + return logits + + def forward_varlen(self, raw_tokens, noisy_refine_tokens, t, pc=None): + device = t.device + raw_tokens = [seq.to(device=device, dtype=torch.long) for seq in raw_tokens] + noisy_refine_tokens = [seq.to(device=device, dtype=torch.long) for seq in noisy_refine_tokens] + lengths = torch.tensor([seq.numel() for seq in raw_tokens], device=device, dtype=torch.long) + if len(raw_tokens) == 0 or int(lengths.sum().item()) == 0: + raise ValueError("forward_varlen received an empty batch.") + if any(a.numel() != b.numel() for a, b in zip(raw_tokens, noisy_refine_tokens)): + raise ValueError("raw token lengths and refine token lengths must match per sample.") + if torch.any(lengths % 9 != 0): + bad = lengths[lengths % 9 != 0].detach().cpu().tolist() + raise ValueError(f"Varlen hourglass requires each sequence length to be divisible by 9, got {bad}.") + + raw_packed = torch.cat(raw_tokens, dim=0) + refine_packed = torch.cat(noisy_refine_tokens, dim=0) + raw_emb = self.raw_embedding(raw_packed) + refine_emb = self.refine_embedding(refine_packed) + time_emb = self.time_mlp(t).repeat_interleave(lengths, dim=0) + x = raw_emb + refine_emb + time_emb + + if self.use_pc: + if pc is None: + # Null condition path for CFG / cond-drop. Replicate the learned + # null_cond per sample and flatten to match the varlen layout. + bsz = len(raw_tokens) + null = self.null_cond # [257, d_model] + cond = null.unsqueeze(0).expand(bsz, -1, -1).reshape(-1, null.size(-1)) + cond_lengths = torch.full((bsz,), self.null_cond_len, device=device, dtype=torch.long) + else: + pc = pc.to(device=device) + cond = self.conditioner(pc) + cond_lengths = torch.full((cond.size(0),), cond.size(1), device=device, dtype=torch.long) + cond = cond.reshape(-1, cond.size(-1)) + else: + cond = None + cond_lengths = None + + def rope_fn_for(cur_lengths): + position_ids = lengths_to_position_ids(cur_lengths) + cos, sin = self.rotary_emb.forward_positions(position_ids, dtype=x.dtype) + return lambda q, k: apply_rotary_pos_emb_varlen(q, k, cos, sin) + + rope_fn_9 = rope_fn_for(lengths) + for block in self.enc_blocks_9: + x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_9) + skip_9 = x + lengths_9 = lengths + + x, lengths = self.down_9_to_3.forward_varlen(x, lengths) + rope_fn_3 = rope_fn_for(lengths) + for block in self.enc_blocks_3: + x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_3) + skip_3 = x + lengths_3 = lengths + + x, lengths = self.down_3_to_1.forward_varlen(x, lengths) + rope_fn_1 = rope_fn_for(lengths) + for block in self.bottleneck_blocks: + x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_1) + + x, lengths = self.up_1_to_3.forward_varlen(x, lengths) + if not torch.equal(lengths, lengths_3): + raise RuntimeError("Internal length mismatch after 1->3 upsample.") + x = x + skip_3 + for block in self.dec_blocks_3: + x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_3) + + x, lengths = self.up_3_to_9.forward_varlen(x, lengths) + if not torch.equal(lengths, lengths_9): + raise RuntimeError("Internal length mismatch after 3->9 upsample.") + x = x + skip_9 + for block in self.dec_blocks_9: + x = block.forward_varlen(x, lengths, cond, cond_lengths, rope_apply_fn=rope_fn_9) + + return self.lm_head(x) + +def cosine_schedule(t, T): + """ + 余弦调度函数。 + t: 当前进度步数 (0 表示起始/全噪，T 表示结束/真实数据) + 返回: 需要 mask 的比例 (t=0时返回1.0，t=T时返回0.0) + """ + # 确保 t 是浮点数以进行安全除法 + return torch.cos((t / T) * (math.pi / 2)) + + return torch.cos((t / T) * (math.pi / 2)) + + +def precompute_topology_tensors(raw_faces, tokens_per_face=9): + """ + 预处理：完全向量化的拓扑索引构建（在去噪循环外执行一次即可） + 返回的是一维的张量，方便在 GPU 上做极速的 Scatter/Gather 计算 + """ + B, num_faces, _, _ = raw_faces.shape + device = raw_faces.device + + # 转为 Python list 加速哈希构建 + raw_faces_list = raw_faces.cpu().tolist() + + token_indices_list = [] + v_indices_list = [] + v_shared_count = 0 + + for b in range(B): + vertex_dict = {} + for f in range(num_faces): + for v_idx in range(3): + coord = tuple(raw_faces_list[b][f][v_idx]) + + # 计算展平后的【全局绝对 Token 索引】 + global_token_idx = b * (num_faces * tokens_per_face) + f * tokens_per_face + v_idx + + if coord not in vertex_dict: + vertex_dict[coord] = [] + vertex_dict[coord].append(global_token_idx) + + # 仅保留被共享的顶点 + for indices in vertex_dict.values(): + if len(indices) > 1: + token_indices_list.extend(indices) + v_indices_list.extend([v_shared_count] * len(indices)) # 记录属于哪个全局共享顶点 ID + v_shared_count += 1 + + # 如果没有任何共享顶点，直接返回空 + if len(token_indices_list) == 0: + return None, None, None, 0 + + # [M] 个共享 Token 的绝对索引 + token_idx = torch.tensor(token_indices_list, dtype=torch.long, device=device) + # [M] 每个 Token 对应的全局顶点 ID + v_idx = torch.tensor(v_indices_list, dtype=torch.long, device=device) + # [M] 每个 Token 对应的绝对面片(Face)索引 + face_idx = token_idx // tokens_per_face + + return token_idx, v_idx, face_idx, v_shared_count + +def sample_refine_tokens_with_topology( + model, + mask_token_id, + raw_tokens, + pc, + T_infer=10, + T_train=100, + tokens_per_face=9, + use_gumbel=False, + gumbel_temp_init=1.0, + use_soft_consensus=False, + consensus_start_step_ratio=0.0, + cfg_scale=0.0, + null_pc=None, + diag_stats=None, + forbidden_token_ids=(0,), + forbid_mask_token=True, + empty_token_ids=(), + empty_token_bias=0.0, +): + """ + 极致优化的拓扑一致性采样代码（零 Python Loop） + + empty_token_ids / empty_token_bias: 对"空/背景"占用 token(实测 258,770 各~33%)的 logits + 减去 empty_token_bias(负偏置)。>0 -> 空 token 更不易被选 -> 更密的细分(更多面)。默认 no-op。 + + 优化 flag: + use_gumbel : 重 mask 时给 confidence 加温度退火的 Gumbel 噪声 (MaskGIT 标配) + gumbel_temp_init : 初始温度，随 step 线性退火到 0 + use_soft_consensus : 共享顶点用概率分布聚合代替硬多数表决 + consensus_start_step_ratio: 在 step / T_infer 达到该值前不做 topology consensus + cfg_scale : Classifier-free guidance 强度. 0 = 不开启 (单 forward). + >0 时每步多一次 forward (null_pc 作 uncond). + logits = logits_c + cfg_scale * (logits_c - logits_u) + null_pc : (B, N, C) 作为 uncond 输入. 若为 None, 用 zeros_like(pc). + 注: 当前模型未训练 cond-drop, 用 null pc 是 OOD 行为, 效果不保证. + diag_stats : 可选 dict, 把每步的诊断信息追加进去 (供测试用) + forbidden_token_ids : 禁止模型"决定"采样到的 token id 集合. 这些 id 的 logits + 在 softmax 前被设为 -inf, 后续 argmax / consensus 全部 + 自动避开. 默认 (0,) 即 pad/EOS——堵住"模型中段吐出 0 + -> save_stage_prediction 截断 -> mesh 只剩前半段"这条 bug. + forbid_mask_token : 同时也禁止预测 mask_token_id (默认 True). mask 是占位符, + 模型不应把它当最终答案. + """ + model.eval() + b, n = raw_tokens.shape + device = raw_tokens.device + num_faces = n // tokens_per_face + + # 收集要屏蔽的 token id; 真正裁剪并搬到 GPU 留到拿到 logits 后再做 + # (mask_token_id 通常等于 vocab_B, 不在输出 logits 的维度内, 直接索引会越界) + _forbidden_raw = list(forbidden_token_ids) if forbidden_token_ids is not None else [] + if forbid_mask_token: + _forbidden_raw.append(int(mask_token_id)) + _forbidden_raw = sorted(set(int(x) for x in _forbidden_raw)) + forbidden_t = None # 在首次 forward 后用 vocab_size 过滤构建 + + # --- 1. 预处理 (仅执行一次，耗时极短) --- + raw_faces = raw_tokens.view(b, num_faces, 3, 3) + token_idx, v_idx, face_idx, v_shared_count = precompute_topology_tensors(raw_faces, tokens_per_face) + + # 初始化 + current_tokens = torch.full((b, n), mask_token_id, dtype=torch.long, device=device) + unmasked_flags = torch.zeros((b, n), dtype=torch.bool, device=device) + vocab_size = None # 将在第一次 forward 后获取 + + def cosine_schedule(t, max_val): + return torch.cos(t * torch.pi * 0.5) * max_val + + # 预备 null_pc (CFG) + # 三种模式: + # - null_pc 是 tensor : 用该 tensor 作为 uncond 输入(老模型可用 zeros_like(pc)) + # - null_pc == "learned" : 调用 model(..., pc=None), 走模型的 self.null_cond 路径 + # (只有 cond-drop fine-tune 过的模型才会有意义) + # - null_pc is None (默认): 用 zeros_like(pc), 行为同老版 + use_learned_null = (isinstance(null_pc, str) and null_pc.lower() == "learned") + if cfg_scale > 0.0 and not use_learned_null: + if null_pc is None: + null_pc = torch.zeros_like(pc) + null_pc = null_pc.to(device=device, dtype=pc.dtype) + + for step in range(T_infer): + current_progress_ratio = step / T_infer + t_model = torch.full((b,), current_progress_ratio, device=device, dtype=torch.float) + + # [1] 前向预测 (可选 CFG: 双 forward 线性外推) + logits_c = model(raw_tokens, current_tokens, t_model, pc) # [b, n, vocab_size] + if cfg_scale > 0.0: + if use_learned_null: + logits_u = model(raw_tokens, current_tokens, t_model, None) + else: + logits_u = model(raw_tokens, current_tokens, t_model, null_pc) + logits = logits_c + cfg_scale * (logits_c - logits_u) + else: + logits = logits_c + if vocab_size is None: + vocab_size = logits.shape[-1] + # 现在知道真实输出 vocab 大小, 把越界的 id (例如 mask_token_id=vocab_B) + # 剔除掉后再构建 forbidden_t. 否则 logits[..., id] 会触发 CUDA assert. + _valid = [x for x in _forbidden_raw if 0 <= x < vocab_size] + if len(_valid) > 0: + forbidden_t = torch.tensor(_valid, dtype=torch.long, device=device) + + # 禁止采样到 pad / EOS. softmax 前置 -inf, 后续 argmax 与 + # consensus 聚合 (软/硬) 都因 probs[forbidden]=0 自动避开. + if forbidden_t is not None: + logits[..., forbidden_t] = float('-inf') + + # 空行/空占用 token 负偏置: 减小其 logits -> 更不易被选 -> 更密的细分 + if empty_token_bias != 0.0 and len(empty_token_ids) > 0: + _et = torch.tensor([int(i) for i in empty_token_ids if 0 <= int(i) < vocab_size], + dtype=torch.long, device=device) + if _et.numel() > 0: + logits[..., _et] = logits[..., _et] - float(empty_token_bias) + + probs = F.softmax(logits, dim=-1) + max_probs, pred_tokens = probs.max(dim=-1) # [b, n] + + force_mask = torch.zeros((b, n), dtype=torch.bool, device=device) + target_vals = None # 最后一步若开启 consensus 会用到 + + # --------------------------------------------------------- + # [2] 拓扑共识 (硬多数表决 / 软概率聚合) + # 通过 consensus_start_step_ratio 控制何时开启 + # --------------------------------------------------------- + apply_consensus = (v_shared_count > 0) and (current_progress_ratio >= consensus_start_step_ratio) + if apply_consensus: + if use_soft_consensus: + # 软共识: 直接 scatter_add 共享 token 的概率分布 + shared_probs = probs.view(-1, vocab_size)[token_idx] # [M, V] + agg_probs = torch.zeros((v_shared_count, vocab_size), dtype=shared_probs.dtype, device=device) + v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size) + agg_probs.scatter_add_(0, v_idx_expanded, shared_probs) + + mode_vals = agg_probs.argmax(dim=-1) + target_vals = mode_vals[v_idx] + + preds_flat = pred_tokens.view(-1) + shared_vals = preds_flat[token_idx] + is_disagree = (shared_vals != target_vals) + else: + # 硬多数表决 + preds_flat = pred_tokens.view(-1) + shared_vals = preds_flat[token_idx] + + one_hot_vals = F.one_hot(shared_vals, num_classes=vocab_size) # [M, V] + votes = torch.zeros((v_shared_count, vocab_size), dtype=torch.long, device=device) + v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size) + votes.scatter_add_(0, v_idx_expanded, one_hot_vals) + + mode_vals = votes.argmax(dim=-1) + target_vals = mode_vals[v_idx] + is_disagree = (shared_vals != target_vals) + + if is_disagree.any(): + disagree_faces = face_idx[is_disagree] + force_mask_faces = torch.zeros(b * num_faces, dtype=torch.bool, device=device) + force_mask_faces[disagree_faces] = True + force_mask_flat = force_mask_faces.repeat_interleave(tokens_per_face) + force_mask = force_mask_flat.view(b, n) + + # --------------------------------------------------------- + # [3] 置信度干预与 Mask 更新 + # --------------------------------------------------------- + # 注意: 这里先克隆，避免 Gumbel 噪声破坏掉之后我们对 unmasked/force 的 ±inf 标记 + confidence = max_probs.clone() + if use_gumbel: + # 温度随进度线性退火到 0 + temperature = gumbel_temp_init * (1.0 - current_progress_ratio) + if temperature > 0: + uniform = torch.rand_like(confidence).clamp_(1e-9, 1.0 - 1e-9) + gumbel = -torch.log(-torch.log(uniform)) + confidence = confidence + temperature * gumbel + + # 已确认的位置不应被再 mask + confidence[unmasked_flags] = float('inf') + + if apply_consensus: + confidence[force_mask] = -float('inf') + unmasked_flags = unmasked_flags & (~force_mask) + + # 诊断信息 + if diag_stats is not None: + diag_stats.setdefault('force_count', []).append(int(force_mask.sum().item())) + diag_stats.setdefault('avg_conf', []).append(float(max_probs.mean().item())) + + if step == T_infer - 1: + # 最后一步：用共识算出的最终值硬覆盖所有可能残余的不一致 + if apply_consensus and target_vals is not None: + preds_flat = pred_tokens.view(-1) + preds_flat[token_idx] = target_vals + pred_tokens = preds_flat.view(b, n) + + current_tokens = torch.where(unmasked_flags, current_tokens, pred_tokens) + + else: + next_progress_ratio = (step + 1) / T_infer + ratio = cosine_schedule(torch.tensor(next_progress_ratio), 1.0).item() + num_to_mask = int(ratio * n) + + max_forced_in_batch = int(force_mask.sum(dim=-1).max().item()) if apply_consensus else 0 + actual_k = min(max(num_to_mask, max_forced_in_batch), n) + + _, mask_indices = torch.topk(-confidence, k=actual_k, dim=-1) + + new_mask = torch.zeros_like(unmasked_flags) + new_mask.scatter_(1, mask_indices, True) + + # 兜底: force_mask 必入 new_mask (Gumbel 噪声不会让 -inf 漏选, 但保留以防边界条件) + new_mask = new_mask | force_mask + + current_tokens = torch.where(new_mask, mask_token_id, pred_tokens) + unmasked_flags = ~new_mask + + return current_tokens + + +def precompute_topology_tensors_fm(raw_faces, tokens_per_face=9): + """precompute_topology_tensors 的扩展版: 除了 F(refine 位置 0/1/2 = 3 个顶点/父的占用, + 按顶点 coord 分组), 额外把 intra(M_inner, refine 位置 3/4/5, 分别对应父 0/1/2 的内部连接 + 矩阵)也按其父 coord 分组。F 与 intra 用不同"通道标签"分开归组, 绝不互相表决。 + + token 布局(与 fix3 同构, 见 tokenizer decode 的 chunk[9..17]): + refine[0,1,2] = F0,F1,F2(父 0/1/2 占用) -> 按父 coord 分组(每父冗余, 应一致) + refine[3,4,5] = M00,M11,M22(父 0/1/2 intra)-> 按父 coord 分组(每父冗余, 应一致) + refine[6,7,8] = M01,M02,M12(inter, 每面本地)-> 不分组、不表决 + + 返回 (token_idx, v_idx, face_idx, v_shared_count), 结构与原版完全一致, 可直接喂给 + 采样循环(F 组的众数落在 F 值域, intra 组的众数落在 intra 值域, 天然不串)。 + """ + B, num_faces, _, _ = raw_faces.shape + device = raw_faces.device + raw_faces_list = raw_faces.cpu().tolist() + + # 每个"冗余通道" = (refine 位置, 该位置对应的父编号, 通道标签)。 + # F: 位置 p -> 父 p; intra: 位置 3+p -> 父 p。inter(6/7/8)不在其中。 + channels = [(0, 0, "F"), (1, 1, "F"), (2, 2, "F"), + (3, 0, "I"), (4, 1, "I"), (5, 2, "I")] + + token_indices_list = [] + v_indices_list = [] + v_shared_count = 0 + for b in range(B): + group_dict = {} + for f in range(num_faces): + for pos, parent, tag in channels: + coord = tuple(raw_faces_list[b][f][parent]) # 该父的坐标 + key = (tag, coord) # F / intra 分开归组 + global_token_idx = b * (num_faces * tokens_per_face) + f * tokens_per_face + pos + group_dict.setdefault(key, []).append(global_token_idx) + for indices in group_dict.values(): + if len(indices) > 1: # 仅保留被共享(>1 面引用)的 + token_indices_list.extend(indices) + v_indices_list.extend([v_shared_count] * len(indices)) + v_shared_count += 1 + + if len(token_indices_list) == 0: + return None, None, None, 0 + token_idx = torch.tensor(token_indices_list, dtype=torch.long, device=device) + v_idx = torch.tensor(v_indices_list, dtype=torch.long, device=device) + face_idx = token_idx // tokens_per_face + return token_idx, v_idx, face_idx, v_shared_count + + +def sample_refine_tokens_with_topology_fm( + model, + mask_token_id, + raw_tokens, + pc, + T_infer=10, + T_train=100, + tokens_per_face=9, + use_gumbel=False, + gumbel_temp_init=1.0, + use_soft_consensus=False, + consensus_start_step_ratio=0.0, + cfg_scale=0.0, + null_pc=None, + diag_stats=None, + forbidden_token_ids=(0,), + forbid_mask_token=True, +): + """sample_refine_tokens_with_topology 的"F + M_inner(intra)双对齐"版。 + + 与原版唯一区别: 拓扑分组用 precompute_topology_tensors_fm —— + 原版只把 refine 位置 0/1/2(F, 顶点占用)按共享顶点分组表决; + 本版额外把 3/4/5(intra, 每父内部连接矩阵)也按共享父分组表决。 + 因此每步的共识检测/末步硬覆盖会同时作用在 F 与 intra 上, 任一不一致的面整块 9 token + 重掩重预测。inter(6/7/8)每面本地、无冗余, 仍不表决(但会随所在面被连坐重掩)。 + + 其余采样逻辑(CFG 双前向、forbidden、Gumbel、cosine 揭示)与原版逐行一致。 + 参数含义见 sample_refine_tokens_with_topology。 + """ + model.eval() + b, n = raw_tokens.shape + device = raw_tokens.device + num_faces = n // tokens_per_face + + _forbidden_raw = list(forbidden_token_ids) if forbidden_token_ids is not None else [] + if forbid_mask_token: + _forbidden_raw.append(int(mask_token_id)) + _forbidden_raw = sorted(set(int(x) for x in _forbidden_raw)) + forbidden_t = None + + # --- 1. 预处理: F + intra 双通道拓扑分组(唯一改动) --- + raw_faces = raw_tokens.view(b, num_faces, 3, 3) + token_idx, v_idx, face_idx, v_shared_count = precompute_topology_tensors_fm(raw_faces, tokens_per_face) + + current_tokens = torch.full((b, n), mask_token_id, dtype=torch.long, device=device) + unmasked_flags = torch.zeros((b, n), dtype=torch.bool, device=device) + vocab_size = None + + def cosine_schedule(t, max_val): + return torch.cos(t * torch.pi * 0.5) * max_val + + use_learned_null = (isinstance(null_pc, str) and null_pc.lower() == "learned") + if cfg_scale > 0.0 and not use_learned_null: + if null_pc is None: + null_pc = torch.zeros_like(pc) + null_pc = null_pc.to(device=device, dtype=pc.dtype) + + for step in range(T_infer): + current_progress_ratio = step / T_infer + t_model = torch.full((b,), current_progress_ratio, device=device, dtype=torch.float) + + # [1] 前向预测 (可选 CFG) + logits_c = model(raw_tokens, current_tokens, t_model, pc) + if cfg_scale > 0.0: + if use_learned_null: + logits_u = model(raw_tokens, current_tokens, t_model, None) + else: + logits_u = model(raw_tokens, current_tokens, t_model, null_pc) + logits = logits_c + cfg_scale * (logits_c - logits_u) + else: + logits = logits_c + if vocab_size is None: + vocab_size = logits.shape[-1] + _valid = [x for x in _forbidden_raw if 0 <= x < vocab_size] + if len(_valid) > 0: + forbidden_t = torch.tensor(_valid, dtype=torch.long, device=device) + + if forbidden_t is not None: + logits[..., forbidden_t] = float('-inf') + + probs = F.softmax(logits, dim=-1) + max_probs, pred_tokens = probs.max(dim=-1) + + force_mask = torch.zeros((b, n), dtype=torch.bool, device=device) + target_vals = None + + # [2] 拓扑共识 (F + intra 同时, 因两者在 v_idx 里是不同组, 自动分别表决) + apply_consensus = (v_shared_count > 0) and (current_progress_ratio >= consensus_start_step_ratio) + if apply_consensus: + if use_soft_consensus: + shared_probs = probs.view(-1, vocab_size)[token_idx] + agg_probs = torch.zeros((v_shared_count, vocab_size), dtype=shared_probs.dtype, device=device) + v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size) + agg_probs.scatter_add_(0, v_idx_expanded, shared_probs) + mode_vals = agg_probs.argmax(dim=-1) + target_vals = mode_vals[v_idx] + preds_flat = pred_tokens.view(-1) + shared_vals = preds_flat[token_idx] + is_disagree = (shared_vals != target_vals) + else: + preds_flat = pred_tokens.view(-1) + shared_vals = preds_flat[token_idx] + one_hot_vals = F.one_hot(shared_vals, num_classes=vocab_size) + votes = torch.zeros((v_shared_count, vocab_size), dtype=torch.long, device=device) + v_idx_expanded = v_idx.unsqueeze(1).expand(-1, vocab_size) + votes.scatter_add_(0, v_idx_expanded, one_hot_vals) + mode_vals = votes.argmax(dim=-1) + target_vals = mode_vals[v_idx] + is_disagree = (shared_vals != target_vals) + + if is_disagree.any(): + disagree_faces = face_idx[is_disagree] + force_mask_faces = torch.zeros(b * num_faces, dtype=torch.bool, device=device) + force_mask_faces[disagree_faces] = True + force_mask_flat = force_mask_faces.repeat_interleave(tokens_per_face) + force_mask = force_mask_flat.view(b, n) + + # [3] 置信度干预与 Mask 更新 + confidence = max_probs.clone() + if use_gumbel: + temperature = gumbel_temp_init * (1.0 - current_progress_ratio) + if temperature > 0: + uniform = torch.rand_like(confidence).clamp_(1e-9, 1.0 - 1e-9) + gumbel = -torch.log(-torch.log(uniform)) + confidence = confidence + temperature * gumbel + + confidence[unmasked_flags] = float('inf') + if apply_consensus: + confidence[force_mask] = -float('inf') + unmasked_flags = unmasked_flags & (~force_mask) + + if diag_stats is not None: + diag_stats.setdefault('force_count', []).append(int(force_mask.sum().item())) + diag_stats.setdefault('avg_conf', []).append(float(max_probs.mean().item())) + + if step == T_infer - 1: + if apply_consensus and target_vals is not None: + preds_flat = pred_tokens.view(-1) + preds_flat[token_idx] = target_vals + pred_tokens = preds_flat.view(b, n) + current_tokens = torch.where(unmasked_flags, current_tokens, pred_tokens) + else: + next_progress_ratio = (step + 1) / T_infer + ratio = cosine_schedule(torch.tensor(next_progress_ratio), 1.0).item() + num_to_mask = int(ratio * n) + max_forced_in_batch = int(force_mask.sum(dim=-1).max().item()) if apply_consensus else 0 + actual_k = min(max(num_to_mask, max_forced_in_batch), n) + _, mask_indices = torch.topk(-confidence, k=actual_k, dim=-1) + new_mask = torch.zeros_like(unmasked_flags) + new_mask.scatter_(1, mask_indices, True) + new_mask = new_mask | force_mask + current_tokens = torch.where(new_mask, mask_token_id, pred_tokens) + unmasked_flags = ~new_mask + + return current_tokens + + + + +def precompute_topology_tensors_fmi(raw_faces, tokens_per_face=9): + """F + intra + inter 三通道拓扑分组(逐步共识用)。 + F (refine 0/1/2): 按单父 coord 分组(每父冗余, 引用同父的面应一致) + intra(refine 3/4/5): 按单父 coord 分组(每父冗余) + inter(refine 6/7/8): 按"无序父对 coord"分组 —— 两个面若恰好共享 2 顶点(=共享一条粗边) + 就共享同一父对, 其 inter 矩阵应一致; tokenizer 每面 3 父排序+按 + (小,大)取 inter, 同一父对在任何面里 token 值相同(朝向天然规范)。 + 三通道用不同 key 前缀("F"/"I"/"O")互不表决。返回结构与 precompute_topology_tensors 完全一致, + 可直接喂采样循环(各通道众数落各自值域, 不串)。 + """ + B, num_faces, _, _ = raw_faces.shape + device = raw_faces.device + raw_faces_list = raw_faces.cpu().tolist() + single_ch = [(0, 0, "F"), (1, 1, "F"), (2, 2, "F"), (3, 0, "I"), (4, 1, "I"), (5, 2, "I")] + pair_ch = [(6, 0, 1), (7, 0, 2), (8, 1, 2)] # inter 位置 -> 父对 + token_indices_list, v_indices_list, v_shared_count = [], [], 0 + for b in range(B): + gd = {} + for f in range(num_faces): + fc = raw_faces_list[b][f] + base = b * (num_faces * tokens_per_face) + f * tokens_per_face + for pos, parent, tag in single_ch: + gd.setdefault((tag, tuple(fc[parent])), []).append(base + pos) + for pos, a, bp in pair_ch: + ca, cb = tuple(fc[a]), tuple(fc[bp]) + pair = (ca, cb) if ca <= cb else (cb, ca) + gd.setdefault(("O", pair), []).append(base + pos) + for indices in gd.values(): + if len(indices) > 1: + token_indices_list.extend(indices) + v_indices_list.extend([v_shared_count] * len(indices)) + v_shared_count += 1 + if len(token_indices_list) == 0: + return None, None, None, 0 + token_idx = torch.tensor(token_indices_list, dtype=torch.long, device=device) + v_idx = torch.tensor(v_indices_list, dtype=torch.long, device=device) + face_idx = token_idx // tokens_per_face + return token_idx, v_idx, face_idx, v_shared_count + + diff --git a/mulres_infer/modules/__init__.py b/mulres_infer/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/mulres_infer/modules/attention.py b/mulres_infer/modules/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..0b915a1e0b86d8db95b1e208308177b142685074 --- /dev/null +++ b/mulres_infer/modules/attention.py @@ -0,0 +1,376 @@ +"""注意力与 Transformer 原语(从 pure_iflame_with_condition 提取, 只保留 MLA 及其依赖)。 + +MLA 是细分器 subdivider.DiscreteDiffusionTransformer 使用的多头潜在注意力块。 +原目录里的 MeshRes / MLCA / LinearDownsample / LinearUpsample / MultiheadCrossFlashrope +都不在推理路径上, 已删除。 +""" + +import torch +import torch.nn as nn +import math +import torch.nn.functional as F +from torch import nn +import time #cal_time new at 10.21 +import numpy as np +from typing import List, Optional, Tuple +from tqdm import tqdm +from einops import rearrange, repeat, pack + + +# ---- 原 rms_norm.py ---- + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine=True, memory_efficient=False): + super().__init__() + self.dim = dim + self.eps = eps + self.elementwise_affine = elementwise_affine + if self.elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim)) + else: + self.register_parameter('weight', None) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()).type_as(x) + if self.weight is not None: + output = output * self.weight + return output + + def extra_repr(self) -> str: + return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}' + +# ---- 原 multihead_flashdiff_2.py ---- + +class Rotary(nn.Module): + def __init__(self, dim: int, max_seq_len=9*50000, dtype=torch.float16): + super().__init__() + # half-truncate RoPE by @YouJiacheng (w/ base freq tuning) + angular_freq = (1 / 1024) ** torch.linspace(0, 1, steps=dim//4, dtype=dtype) + angular_freq = torch.cat([angular_freq, angular_freq.new_zeros(dim//4)]) + t = torch.arange(max_seq_len, dtype=dtype) + theta = torch.einsum("i,j -> ij", t, angular_freq) + self.register_buffer('cos', theta.cos(), persistent=False) + self.register_buffer('sin', theta.sin(), persistent=False) + self.max_seq_len = max_seq_len + self.dtype = dtype + + def forward(self, x_BTHD, offset=0): + """ + Args: + x_BTHD: Input tensor of shape [batch, seq_len, num_heads, dim] + offset: Position offset for the sequence + """ + batch_size, seq_len = x_BTHD.size(0), x_BTHD.size(-3) + assert torch.max(offset + seq_len) <= self.max_seq_len, f"Sequence length {offset + seq_len} exceeds maximum length {self.max_seq_len}" + + # Get the appropriate slice of cos/sin based on offset + positions = torch.arange(seq_len, device=x_BTHD.device).expand(batch_size, -1) + offset.unsqueeze(-1) + cos = self.cos[positions].unsqueeze(-2) + sin = self.sin[positions].unsqueeze(-2) + + x1, x2 = x_BTHD.to(dtype=self.dtype).chunk(2, dim=-1) + y1 = x1 * cos - x2 * sin + y2 = x1 * sin + x2 * cos + return torch.cat((y1, y2), dim=-1) + +class MultiheadFlashrope(nn.Module): + def __init__(self, args, embed_dim, num_heads): + super().__init__() + self.args = args + kvcache_method, kv_cache_window_size, window_stride = args + self.embed_dim = embed_dim + self.num_heads = num_heads + self.num_kv_heads = num_heads + self.head_dim = embed_dim // num_heads + self.scaling = self.head_dim ** -0.5 + self.kv_cache_window_size = kv_cache_window_size + self.window_stride = window_stride + self.kvcache_method = kvcache_method + # KV cache related attributes + self.kv_cache_enabled = False + self.k_cache = None + self.v_cache = None + self.cache_pos = 0 # Track position in cache + + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.rotary = Rotary(self.head_dim) + self.sliding_window_cur_len = 0 + # self.attention_scores = None + + def forward(self, x, start_pos, context=None, attention_mask=None, causal=False): + """ + Args: + x: [batch_size, tgt_len, embed_dim] + start_pos: [batch_size] + context: Optional context for cross-attention [batch_size, src_len, embed_dim] + """ + bsz, tgt_len, embed_dim = x.size() + # Project inputs to q, k, v + q = self.q_proj(x) + k = self.k_proj(context if context is not None else x) + v = self.v_proj(context if context is not None else x) + + # Reshape to [batch_size, seq_len, num_heads, head_dim] + q = q.view(bsz, -1, self.num_heads, self.head_dim) + k = k.view(bsz, -1, self.num_heads, self.head_dim) + v = v.view(bsz, -1, self.num_heads, self.head_dim) + q = self.rotary(q, offset=start_pos) + # Normal operation without cache + k = self.rotary(k, offset=start_pos) + if attention_mask is not None: + attn_output = self.attention_pytorch_optimized(q, k, v, mask=attention_mask) + else: + attn_output = self.attention_pytorch_optimized(q, k, v, causal=causal) + + # Reshape output + attn_output = attn_output.contiguous().view(bsz, tgt_len, embed_dim) + attn_output = self.out_proj(attn_output) + + return attn_output + + def attention_pytorch_optimized(self, q, k, v, mask=None, causal=False): + b, seq_q, head_num, c = q.shape + seq_k = k.shape[1] + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + if mask is not None: + atten_mask_npu = mask.unsqueeze(1).expand(b, head_num, seq_q, seq_k) + atten_mask_npu = torch.logical_not(atten_mask_npu) + attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=atten_mask_npu) + else: + if causal: + attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True) + else: + attn_output = F.scaled_dot_product_attention(q, k, v) + + return attn_output.permute(0, 2, 1, 3) + + def attention_pytorch(self, q, k, v, causal=False): + # [batch_size, num_heads, seq_len, head_dim] + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + + attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scaling + + if causal: + seq_len = q.size(-2) + causal_mask = torch.triu( + torch.ones(seq_len, seq_len, dtype=torch.bool, device=q.device), + diagonal=1 + ) + attn_weights.masked_fill_(causal_mask, float('-inf')) + + attn_weights = F.softmax(attn_weights, dim=-1) + + + attn_output = torch.matmul(attn_weights, v) + + + return attn_output.permute(0, 2, 1, 3) + + def empty_kv_cache(self, batch_size: int, dtype: torch.dtype): + """Initialize empty KV cache""" + self.kv_cache_enabled=True + + device = next(self.parameters()).device + kv_cache_maxlen = self.kv_cache_window_size + # Initialize empty caches + k_cache = torch.zeros( + batch_size, + kv_cache_maxlen, + self.num_heads, + self.head_dim, + dtype=dtype, + device=device + ) + self.k_cache=k_cache + + v_cache = torch.zeros( + batch_size, + kv_cache_maxlen, + self.num_heads, + self.head_dim, + dtype=dtype, + device=device + ) + self.v_cache=v_cache + self.cache_pos = 0 # Reset cache position + + element_size = { + torch.float32: 4, + torch.float16: 2, + torch.bfloat16: 2, + }[dtype] + k_cache_size = (batch_size * + kv_cache_maxlen * + self.num_heads * + self.head_dim * + element_size) + v_cache_size = k_cache_size + total_cache_size = k_cache_size + v_cache_size + cache_size_gb = total_cache_size / (1024**3) + return cache_size_gb + + + def reset_cache(self): + """Reset KV cache state""" + self.k_cache = None + self.v_cache = None + self.cache_pos = 0 + +class CrossAttention(nn.Module): + def __init__(self, dim, context_dim, n_heads, dropout=0.0): + super().__init__() + self.n_heads = n_heads + self.scale = (dim // n_heads) ** -0.5 + + self.q_proj = nn.Linear(dim, dim, bias=False) + self.kv_proj = nn.Linear(context_dim, 2 * dim, bias=False) + self.out_proj = nn.Linear(dim, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, context): + # x: [batch, seq_len, dim], context: [batch, context_len, context_dim] + B, N, C = x.shape + _, M, _ = context.shape + H = self.n_heads + + # Linear projections + q = self.q_proj(x).view(B, N, H, C // H) # [B, H, seq_len, dim//H] + k, v = self.kv_proj(context).chunk(2, dim=-1) + k = k.view(B, M, H, C // H) # [B, context_len, H, dim//H] + v = v.view(B, M, H, C // H) # [B, context_len, H, dim//H] + + attn_output = self.attention_pytorch_optimized(q,k,v) + attn_output = attn_output.contiguous().view(B, N, C) + + # q = self.q_proj(x).view(B, N, H, C // H).transpose(1, 2) # [B, H, seq_len, dim//H] + # k, v = self.kv_proj(context).chunk(2, dim=-1) + # k = k.view(B, M, H, C // H).transpose(1, 2) # [B, H, context_len, dim//H] + # v = v.view(B, M, H, C // H).transpose(1, 2) # [B, H, context_len, dim//H] + # # Scaled dot-product attention + # attn_weights = (q @ k.transpose(-2, -1)) * self.scale # [B, H, seq_len, context_len] + # attn_weights = attn_weights.softmax(dim=-1) + # attn_output = (attn_weights @ v).transpose(1, 2).contiguous().view(B, N, C) # [B, seq_len, dim] + out = self.out_proj(attn_output) + out = self.dropout(out) + + return out + + def attention_pytorch_optimized(self, q, k, v, mask=None, causal=False): + b, seq_q, head_num, c = q.shape + seq_k = k.shape[1] + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + if mask is not None: + atten_mask_npu = mask.unsqueeze(1).expand(b, head_num, seq_q, seq_k) + atten_mask_npu = torch.logical_not(atten_mask_npu) + attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=atten_mask_npu) + else: + if causal: + attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True) + else: + attn_output = F.scaled_dot_product_attention(q, k, v) + + return attn_output.permute(0, 2, 1, 3) + +# ---- 原 transformer.py ---- + +class SwiGLU(nn.Module): + def __init__(self, embed_dim): + super(SwiGLU, self).__init__() + self.proj1 = nn.Linear(embed_dim, int(8 / 3 * embed_dim)) # XWG + self.proj2 = nn.Linear(embed_dim, int(8 / 3 * embed_dim)) # XW1 + self.proj3 = nn.Linear(int(8 / 3 * embed_dim), embed_dim) # W2 + + def forward(self, x): + # Apply Swish (SiLU in PyTorch) to the first projection + x_proj1 = F.silu(self.proj1(x)) # Swish(XWG) + # Apply the second projection + x_proj2 = self.proj2(x) # XW1 + # Element-wise multiplication + x_glu = x_proj1 * x_proj2 # (swish(XWG) ⊙ XW1) + # Final projection back to the original dimension + output = self.proj3(x_glu) # (swish(XWG) ⊙ XW1)W2 + + return output + +class MLA(nn.Module): + def __init__(self, embed_dim, num_heads, args, cross_embed_dim, causal=True, dropout=0.): + """ + Differential Transformer Block with optional causal self-attention or cross-attention + and automatic application of ROPE. + + Args: + embed_dim (int): Embedding dimension. + num_heads (int): Number of attention heads. + args (Namespace): Arguments for attention settings. + causal (bool): Whether to use causal masking by default for self-attention. + """ + super(MLA, self).__init__() + # Multihead differential attention module + self.attn = MultiheadFlashrope(args, embed_dim, num_heads) + self.causal = causal + + self.feed_forward = SwiGLU(embed_dim) + self.norm1 = RMSNorm(embed_dim, eps=1e-5) + self.norm2 = RMSNorm(embed_dim, eps=1e-5) + self.attn_dropout = nn.Dropout(dropout) + self.ff_dropout = nn.Dropout(dropout) + self.residual_dropout = nn.Dropout(dropout) + if cross_embed_dim != -1: + self.cross_norm = RMSNorm(embed_dim, eps=1e-5) + self.cross_attn = CrossAttention(embed_dim, cross_embed_dim, num_heads) + + def forward( + self, x, start, conds=None, attention_mask=None, use_cache=False + ): + """ + Args: + x: Input tensor + use_cache: Whether to use KV cache + """ + # Enable/disable KV cache in attention module + self.attn.kv_cache_enabled = use_cache + + # 第一个残差连接：注意力 + dropout + residual = x + x = self.norm1(x) + attn_out = self.attn( + x, + start_pos=start, + attention_mask=attention_mask, + causal=self.causal, + ) + attn_out = self.attn_dropout(attn_out) + x = residual + attn_out + if conds is not None: + x_skip = x + x = self.cross_norm(x) + x = self.cross_attn(x, conds) + x_skip + + # 第二个残差连接：前馈网络 + dropout + residual = x + ff_out = self.ff_dropout((self.feed_forward(self.norm2(x)))) + x = residual + ff_out + + return x + + def init_kv_cache(self, batch_size, dtype=torch.bfloat16): + """Initialize KV cache for this transformer block""" + return self.attn.empty_kv_cache( + batch_size=batch_size, dtype=dtype + ) + + def reset_kv_cache(self): + """Reset KV cache for this transformer block""" + self.attn.reset_cache() diff --git a/mulres_infer/modules/fast_octree_level.py b/mulres_infer/modules/fast_octree_level.py new file mode 100644 index 0000000000000000000000000000000000000000..8e5d4e75f2766f3041ad2922faa1f05d6e04e130 --- /dev/null +++ b/mulres_infer/modules/fast_octree_level.py @@ -0,0 +1,52 @@ +"""数组化 + Numba 的「单层」八叉树编码,等价替换 + MatrixOctree(max_depth=target_depth+1).build + OctreeFaceTokenBuilder(mc=8).encode_level_tokens(target_depth) + +前提(per-row 栈已满足):max_depth = target_depth+1 -> 只有一次「叶子(td+1)->节点(td)」聚合, +不需要多层 neighbor 传播;且 max_children=8 时 reduce_node_children_to_max 是 no-op(活跃 octant ≤ 8)。 + +去掉 OctreeNode 对象 / dict-of-tuple / set-of-tuple: + * 叶子坐标量化 -> 边表(numpy) -> 节点坐标(leaf>>1)+ octant 局部索引(位运算, 向量化) + * 按节点聚合 intra 8×8 / 活跃掩码 F / 节点-节点邻接(向量化 scatter / np.unique) + * 三角形(3-clique)枚举走 Numba njit(CSR 两指针求交), 这是原 get_level_info 的热点 + * inter 8×8 按 (节点,邻居) 分组填充 +输出 (raw_t, out_t) 与原实现「逐 slot 解码后网格一致」(顺序可不同;decode 与顺序无关)。 +""" +import numpy as np +from numba import njit + +from utils.data_utils import discretize + +_OCT_W = np.array([1, 2, 4], dtype=np.int64) # octant 局部索引 = x + 2y + 4z + + +@njit(cache=False) +def _triangles_csr(indptr, indices): + """前向邻接 CSR(indices[indptr[u]:indptr[u+1]] = 与 u 相连且 >u 的节点, 升序)。 + 枚举 uv 的部分(即 pa+1..eu, 因为升序且后面都 >v) 与 adj[v](sv..ev) + i = pa + 1 + j = sv + while i < eu and j < ev: + wi = indices[i] + wj = indices[j] + if wi == wj: + out_u.append(u); out_v.append(v); out_w.append(wi) + i += 1; j += 1 + elif wi < wj: + i += 1 + else: + j += 1 + return (np.array(out_u, dtype=np.int64), + np.array(out_v, dtype=np.int64), + np.array(out_w, dtype=np.int64)) + + diff --git a/mulres_infer/modules/miche/LICENSE b/mulres_infer/modules/miche/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..3877ae0a7ff6f94ac222fd704e112723db776114 --- /dev/null +++ b/mulres_infer/modules/miche/LICENSE @@ -0,0 +1,674 @@ + GNU GENERAL PUBLIC LICENSE + Version 3, 29 June 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU General Public License is a free, copyleft license for +software and other kinds of works. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +the GNU General Public License is intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. We, the Free Software Foundation, use the +GNU General Public License for most of our software; it applies also to +any other work released this way by its authors. You can apply it to +your programs, too. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + To protect your rights, we need to prevent others from denying you +these rights or asking you to surrender the rights. Therefore, you have +certain responsibilities if you distribute copies of the software, or if +you modify it: responsibilities to respect the freedom of others. + + For example, if you distribute copies of such a program, whether +gratis or for a fee, you must pass on to the recipients the same +freedoms that you received. You must make sure that they, too, receive +or can get the source code. And you must show them these terms so they +know their rights. + + Developers that use the GNU GPL protect your rights with two steps: +(1) assert copyright on the software, and (2) offer you this License +giving you legal permission to copy, distribute and/or modify it. + + For the developers' and authors' protection, the GPL clearly explains +that there is no warranty for this free software. For both users' and +authors' sake, the GPL requires that modified versions be marked as +changed, so that their problems will not be attributed erroneously to +authors of previous versions. + + Some devices are designed to deny users access to install or run +modified versions of the software inside them, although the manufacturer +can do so. This is fundamentally incompatible with the aim of +protecting users' freedom to change the software. The systematic +pattern of such abuse occurs in the area of products for individuals to +use, which is precisely where it is most unacceptable. Therefore, we +have designed this version of the GPL to prohibit the practice for those +products. If such problems arise substantially in other domains, we +stand ready to extend this provision to those domains in future versions +of the GPL, as needed to protect the freedom of users. + + Finally, every program is threatened constantly by software patents. +States should not allow patents to restrict development and use of +software on general-purpose computers, but in those that do, we wish to +avoid the special danger that patents applied to a free program could +make it effectively proprietary. To prevent this, the GPL assures that +patents cannot be used to render the program non-free. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Use with the GNU Affero General Public License. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU Affero General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the special requirements of the GNU Affero General Public License, +section 13, concerning interaction through a network will apply to the +combination as such. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU General Public License from time to time. Such new versions will +be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If the program does terminal interaction, make it output a short +notice like this when it starts in an interactive mode: + + Copyright (C) + This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'. + This is free software, and you are welcome to redistribute it + under certain conditions; type `show c' for details. + +The hypothetical commands `show w' and `show c' should show the appropriate +parts of the General Public License. Of course, your program's commands +might be different; for a GUI interface, you would use an "about box". + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU GPL, see +. + + The GNU General Public License does not permit incorporating your program +into proprietary programs. If your program is a subroutine library, you +may consider it more useful to permit linking proprietary applications with +the library. If this is what you want to do, use the GNU Lesser General +Public License instead of this License. But first, please read +. diff --git a/mulres_infer/modules/miche/__init__.py b/mulres_infer/modules/miche/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/mulres_infer/modules/miche/download_ckpt.py b/mulres_infer/modules/miche/download_ckpt.py new file mode 100644 index 0000000000000000000000000000000000000000..17ea0f81bd0d54bb9e4c7879357da736cdbab4ef --- /dev/null +++ b/mulres_infer/modules/miche/download_ckpt.py @@ -0,0 +1,25 @@ +from huggingface_hub import hf_hub_download +import os + +# 设置镜像（国内加速） +# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' + +def download_michelangelo_checkpoint(): + """下载Michelangelo模型的checkpoint文件""" + try: + # 下载文件 + local_path = hf_hub_download( + repo_id="Maikou/Michelangelo", + filename="checkpoints/aligned_shape_latents/shapevae-256.ckpt", + local_dir="./downloads", # 本地保存目录 + local_dir_use_symlinks=False, # 不使用符号链接 + resume_download=True, # 支持断点续传 + ) + print(f"文件已下载到: {local_path}") + return local_path + except Exception as e: + print(f"下载失败: {e}") + return None + +if __name__ == "__main__": + download_michelangelo_checkpoint() \ No newline at end of file diff --git a/mulres_infer/modules/miche/encode.py b/mulres_infer/modules/miche/encode.py new file mode 100644 index 0000000000000000000000000000000000000000..b730925a4902f84fa41cbcb30dd30b00f9f4b5c2 --- /dev/null +++ b/mulres_infer/modules/miche/encode.py @@ -0,0 +1,84 @@ +import os +# -*- coding: utf-8 -*- +import argparse +from omegaconf import OmegaConf +import numpy as np +import torch +from .michelangelo.utils.misc import instantiate_from_config + +def load_surface(fp): + + with np.load(fp) as input_pc: + surface = input_pc['points'] + normal = input_pc['normals'] + + rng = np.random.default_rng() + ind = rng.choice(surface.shape[0], 4096, replace=False) + surface = torch.FloatTensor(surface[ind]) + normal = torch.FloatTensor(normal[ind]) + + surface = torch.cat([surface, normal], dim=-1).unsqueeze(0).cuda() + + return surface + +def reconstruction(args, model, bounds=(-1.25, -1.25, -1.25, 1.25, 1.25, 1.25), octree_depth=7, num_chunks=10000): + + surface = load_surface(args.pointcloud_path) + # old_surface = surface.clone() + + # surface[0,:,0]*=-1 + # surface[0,:,1]*=-1 + surface[0,:,2]*=-1 + + # encoding + shape_embed, shape_latents = model.model.encode_shape_embed(surface, return_latents=True) + shape_zq, posterior = model.model.shape_model.encode_kl_embed(shape_latents) + + # decoding + latents = model.model.shape_model.decode(shape_zq) + # geometric_func = partial(model.model.shape_model.query_geometry, latents=latents) + + return 0 + +def load_model(ckpt_path=None, config_path=None): + _HERE = os.path.dirname(os.path.abspath(__file__)) + config_path = config_path or os.path.join(_HERE, "shapevae-256.yaml") + + model_config = OmegaConf.load(config_path) + # print(model_config) + if hasattr(model_config, "model"): + model_config = model_config.model + + original_load = torch.load + torch.load = lambda path, *args, **kwargs: original_load(path, map_location="cpu", weights_only=False) + try: + model = instantiate_from_config(model_config, ckpt_path=ckpt_path) + model = model.eval() + finally: + # 恢复原始的 torch.load + torch.load = original_load + + return model +if __name__ == "__main__": + ''' + 1. Reconstruct point cloud + 2. Image-conditioned generation + 3. Text-conditioned generation + ''' + parser = argparse.ArgumentParser() + parser.add_argument("--config_path", type=str, required=True) + parser.add_argument("--ckpt_path", type=str, required=True) + parser.add_argument("--pointcloud_path", type=str, default='./example_data/surface.npz', + help='Path to the input point cloud') + parser.add_argument("--image_path", type=str, help='Path to the input image') + parser.add_argument("--text", type=str, + help='Input text within a format: A 3D model of motorcar; Porsche 911.') + parser.add_argument("--output_dir", type=str, default='./output') + parser.add_argument("-s", "--seed", type=int, default=0) + args = parser.parse_args() + + print(f'-----------------------------------------------------------------------------') + print(f'>>> Output directory: {args.output_dir}') + print(f'-----------------------------------------------------------------------------') + + reconstruction(args, load_model(args)) diff --git a/mulres_infer/modules/miche/michelangelo/__init__.py b/mulres_infer/modules/miche/michelangelo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4c48b5acbf4ae09a163f0f8fa01a91341695211f --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/__init__.py @@ -0,0 +1 @@ +# -*- coding: utf-8 -*- diff --git a/mulres_infer/modules/miche/michelangelo/graphics/__init__.py b/mulres_infer/modules/miche/michelangelo/graphics/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4c48b5acbf4ae09a163f0f8fa01a91341695211f --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/graphics/__init__.py @@ -0,0 +1 @@ +# -*- coding: utf-8 -*- diff --git a/mulres_infer/modules/miche/michelangelo/graphics/primitives/__init__.py b/mulres_infer/modules/miche/michelangelo/graphics/primitives/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8dd3e711be7a89dd96422803b9dc61d9c55b0815 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/graphics/primitives/__init__.py @@ -0,0 +1,4 @@ +# -*- coding: utf-8 -*- + +from .volume import generate_dense_grid_points + diff --git a/mulres_infer/modules/miche/michelangelo/graphics/primitives/volume.py b/mulres_infer/modules/miche/michelangelo/graphics/primitives/volume.py new file mode 100644 index 0000000000000000000000000000000000000000..31450efd0914646e7109d3de5f17fa30b4eed510 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/graphics/primitives/volume.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +import numpy as np + +# produce dense points +def generate_dense_grid_points(bbox_min: np.ndarray, + bbox_max: np.ndarray, + octree_depth: int, + indexing: str = "ij"): + length = bbox_max - bbox_min + num_cells = np.exp2(octree_depth) + x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32) + y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32) + z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32) + [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing) + xyz = np.stack((xs, ys, zs), axis=-1) + xyz = xyz.reshape(-1, 3) + grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1] + + return xyz, grid_size, length + diff --git a/mulres_infer/modules/miche/michelangelo/models/__init__.py b/mulres_infer/modules/miche/michelangelo/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4c48b5acbf4ae09a163f0f8fa01a91341695211f --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/__init__.py @@ -0,0 +1 @@ +# -*- coding: utf-8 -*- diff --git a/mulres_infer/modules/miche/michelangelo/models/modules/__init__.py b/mulres_infer/modules/miche/michelangelo/models/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0d0994729df700371c3f262f778a1dad44c87646 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/modules/__init__.py @@ -0,0 +1,3 @@ +# -*- coding: utf-8 -*- + +from .checkpoint import checkpoint diff --git a/mulres_infer/modules/miche/michelangelo/models/modules/checkpoint.py b/mulres_infer/modules/miche/michelangelo/models/modules/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..c5d13ba0980e9e43c9d756a3f838eb05af919c62 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/modules/checkpoint.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- + +import torch +from typing import Callable, Iterable, Sequence, Union + + +def checkpoint( + func: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor]]], + inputs: Sequence[torch.Tensor], + params: Iterable[torch.Tensor], + flag: bool, + use_deepspeed: bool = False +): + # Evaluate a function without caching intermediate activations, allowing for + # reduced memory at the expense of extra compute in the backward pass. + # :param func: the function to evaluate. + # :param inputs: the argument sequence to pass to `func`. + # :param params: a sequence of parameters `func` depends on but does not + # explicitly take as arguments. + # :param flag: if False, disable gradient checkpointing. + # :param use_deepspeed: if True, use deepspeed + if flag: + if use_deepspeed: + import deepspeed + return deepspeed.checkpointing.checkpoint(func, *inputs) + + args = tuple(inputs) + tuple(params) + return CheckpointFunction.apply(func, len(inputs), *args) + else: + return func(*inputs) + + +class CheckpointFunction(torch.autograd.Function): + @staticmethod + @torch.cuda.amp.custom_fwd + def forward(ctx, run_function, length, *args): + ctx.run_function = run_function + ctx.input_tensors = list(args[:length]) + ctx.input_params = list(args[length:]) + + with torch.no_grad(): + output_tensors = ctx.run_function(*ctx.input_tensors) + return output_tensors + + @staticmethod + @torch.cuda.amp.custom_bwd + def backward(ctx, *output_grads): + ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] + with torch.enable_grad(): + # Fixes a bug where the first op in run_function modifies the + # Tensor storage in place, which is not allowed for detach()'d + # Tensors. + shallow_copies = [x.view_as(x) for x in ctx.input_tensors] + output_tensors = ctx.run_function(*shallow_copies) + input_grads = torch.autograd.grad( + output_tensors, + ctx.input_tensors + ctx.input_params, + output_grads, + allow_unused=True, + ) + del ctx.input_tensors + del ctx.input_params + del output_tensors + return (None, None) + input_grads diff --git a/mulres_infer/modules/miche/michelangelo/models/modules/distributions.py b/mulres_infer/modules/miche/michelangelo/models/modules/distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..87ed88986ea9760e6611030efad74eabd7bff31f --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/modules/distributions.py @@ -0,0 +1,83 @@ +# -*- coding: utf-8 -*- + +import torch +import numpy as np +from typing import Union, List + + +class DiagonalGaussianDistribution(object): + # Gaussian distribution + def __init__(self, parameters: Union[torch.Tensor, List[torch.Tensor]], deterministic=False, feat_dim=1): + self.feat_dim = feat_dim + self.parameters = parameters + + if isinstance(parameters, list): + self.mean = parameters[0] + self.logvar = parameters[1] + else: + self.mean, self.logvar = torch.chunk(parameters, 2, dim=feat_dim) + + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean) + + # sample from the guassian distribution + def sample(self): + x = self.mean + self.std * torch.randn_like(self.mean) + return x + + def kl(self, other=None, dims=(1, 2, 3)): + if self.deterministic: + return torch.Tensor([0.]) + else: + if other is None: + return 0.5 * torch.mean(torch.pow(self.mean, 2) + + self.var - 1.0 - self.logvar, + dim=dims) + else: + return 0.5 * torch.mean( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var - 1.0 - self.logvar + other.logvar, + dim=dims) + + def nll(self, sample, dims=(1, 2, 3)): + if self.deterministic: + return torch.Tensor([0.]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims) + + def mode(self): + return self.mean + + +def normal_kl(mean1, logvar1, mean2, logvar2): + # Compute the KL divergence between two gaussians. + # Shapes are automatically broadcasted, so batches can be compared to + # scalars, among other use cases. + + tensor = None + for obj in (mean1, logvar1, mean2, logvar2): + if isinstance(obj, torch.Tensor): + tensor = obj + break + assert tensor is not None, "at least one argument must be a Tensor" + + # Force variances to be Tensors. Broadcasting helps convert scalars to + # Tensors, but it does not work for torch.exp(). + logvar1, logvar2 = [ + x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor) + for x in (logvar1, logvar2) + ] + + return 0.5 * ( + -1.0 + + logvar2 + - logvar1 + + torch.exp(logvar1 - logvar2) + + ((mean1 - mean2) ** 2) * torch.exp(-logvar2) + ) diff --git a/mulres_infer/modules/miche/michelangelo/models/modules/embedder.py b/mulres_infer/modules/miche/michelangelo/models/modules/embedder.py new file mode 100644 index 0000000000000000000000000000000000000000..1ed654feb5956b3ceb1d3312f1050bc3d0428386 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/modules/embedder.py @@ -0,0 +1,213 @@ +# -*- coding: utf-8 -*- + +import numpy as np +import torch +import torch.nn as nn +import math + +VALID_EMBED_TYPES = ["identity", "fourier", "hashgrid", "sphere_harmonic", "triplane_fourier"] + + +class FourierEmbedder(nn.Module): + """The sin/cosine positional embedding. Given an input tensor `x` of shape [n_batch, ..., c_dim], it converts + each feature dimension of `x[..., i]` into: + [ + sin(x[..., i]), + sin(f_1*x[..., i]), + sin(f_2*x[..., i]), + ... + sin(f_N * x[..., i]), + cos(x[..., i]), + cos(f_1*x[..., i]), + cos(f_2*x[..., i]), + ... + cos(f_N * x[..., i]), + x[..., i] # only present if include_input is True. + ], here f_i is the frequency. + + Denote the space is [0 / num_freqs, 1 / num_freqs, 2 / num_freqs, 3 / num_freqs, ..., (num_freqs - 1) / num_freqs]. + If logspace is True, then the frequency f_i is [2^(0 / num_freqs), ..., 2^(i / num_freqs), ...]; + Otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1)]. + + Args: + num_freqs (int): the number of frequencies, default is 6; + logspace (bool): If logspace is True, then the frequency f_i is [..., 2^(i / num_freqs), ...], + otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1)]; + input_dim (int): the input dimension, default is 3; + include_input (bool): include the input tensor or not, default is True. + + Attributes: + frequencies (torch.Tensor): If logspace is True, then the frequency f_i is [..., 2^(i / num_freqs), ...], + otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1); + + out_dim (int): the embedding size, if include_input is True, it is input_dim * (num_freqs * 2 + 1), + otherwise, it is input_dim * num_freqs * 2. + + """ + + def __init__(self, + num_freqs: int = 6, + logspace: bool = True, + input_dim: int = 3, + include_input: bool = True, + include_pi: bool = True) -> None: + + """The initialization""" + + super().__init__() + + if logspace: + frequencies = 2.0 ** torch.arange( + num_freqs, + dtype=torch.float32 + ) + else: + frequencies = torch.linspace( + 1.0, + 2.0 ** (num_freqs - 1), + num_freqs, + dtype=torch.float32 + ) + + if include_pi: + frequencies *= torch.pi + + self.register_buffer("frequencies", frequencies, persistent=False) + self.include_input = include_input + self.num_freqs = num_freqs + + self.out_dim = self.get_dims(input_dim) + + def get_dims(self, input_dim): + temp = 1 if self.include_input or self.num_freqs == 0 else 0 + out_dim = input_dim * (self.num_freqs * 2 + temp) + + return out_dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ Forward process. + + Args: + x: tensor of shape [..., dim] + + Returns: + embedding: an embedding of `x` of shape [..., dim * (num_freqs * 2 + temp)] + where temp is 1 if include_input is True and 0 otherwise. + """ + + if self.num_freqs > 0: + embed = (x[..., None].contiguous() * self.frequencies).view(*x.shape[:-1], -1) + if self.include_input: + return torch.cat((x, embed.sin(), embed.cos()), dim=-1) + else: + return torch.cat((embed.sin(), embed.cos()), dim=-1) + else: + return x + + +class LearnedFourierEmbedder(nn.Module): + """ following @crowsonkb "s lead with learned sinusoidal pos emb """ + """ https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """ + + def __init__(self, in_channels, dim): + super().__init__() + assert (dim % 2) == 0 + half_dim = dim // 2 + per_channel_dim = half_dim // in_channels + self.weights = nn.Parameter(torch.randn(per_channel_dim)) + + def forward(self, x): + """ + + Args: + x (torch.FloatTensor): [..., c] + + Returns: + x (torch.FloatTensor): [..., d] + """ + + # [b, t, c, 1] * [1, d] = [b, t, c, d] -> [b, t, c * d] + freqs = (x[..., None] * self.weights[None] * 2 * np.pi).view(*x.shape[:-1], -1) + fouriered = torch.cat((x, freqs.sin(), freqs.cos()), dim=-1) + return fouriered + + +class TriplaneLearnedFourierEmbedder(nn.Module): + def __init__(self, in_channels, dim): + super().__init__() + + self.yz_plane_embedder = LearnedFourierEmbedder(in_channels, dim) + self.xz_plane_embedder = LearnedFourierEmbedder(in_channels, dim) + self.xy_plane_embedder = LearnedFourierEmbedder(in_channels, dim) + + self.out_dim = in_channels + dim + + def forward(self, x): + + yz_embed = self.yz_plane_embedder(x) + xz_embed = self.xz_plane_embedder(x) + xy_embed = self.xy_plane_embedder(x) + + embed = yz_embed + xz_embed + xy_embed + + return embed + + +def sequential_pos_embed(num_len, embed_dim): + assert embed_dim % 2 == 0 + + pos = torch.arange(num_len, dtype=torch.float32) + omega = torch.arange(embed_dim // 2, dtype=torch.float32) + omega /= embed_dim / 2. + omega = 1. / 10000 ** omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = torch.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = torch.sin(out) # (M, D/2) + emb_cos = torch.cos(out) # (M, D/2) + + embeddings = torch.cat([emb_sin, emb_cos], dim=1) # (M, D) + + return embeddings + + +def timestep_embedding(timesteps, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + :param timesteps: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an [N x dim] Tensor of positional embeddings. + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half + ).to(device=timesteps.device) + args = timesteps[:, None].to(timesteps.dtype) * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +def get_embedder(embed_type="fourier", num_freqs=-1, input_dim=3, degree=4, + num_levels=16, level_dim=2, per_level_scale=2, base_resolution=16, + log2_hashmap_size=19, desired_resolution=None): + if embed_type == "identity" or (embed_type == "fourier" and num_freqs == -1): + return nn.Identity(), input_dim + + elif embed_type == "fourier": + embedder_obj = FourierEmbedder(num_freqs=num_freqs, input_dim=input_dim, + logspace=True, include_input=True) + return embedder_obj, embedder_obj.out_dim + + elif embed_type == "hashgrid": + raise NotImplementedError + + elif embed_type == "sphere_harmonic": + raise NotImplementedError + + else: + raise ValueError(f"{embed_type} is not valid. Currently only supprts {VALID_EMBED_TYPES}") diff --git a/mulres_infer/modules/miche/michelangelo/models/modules/transformer_blocks.py b/mulres_infer/modules/miche/michelangelo/models/modules/transformer_blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..3d1b489baf9df765bad4e6b3b4c47235b40cfdd8 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/modules/transformer_blocks.py @@ -0,0 +1,309 @@ +# -*- coding: utf-8 -*- + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Optional +from modules.miche.michelangelo.models.modules.checkpoint import checkpoint + +# Initialize linear layers with normal distribution weights and zero biases +def init_linear(l, stddev): + nn.init.normal_(l.weight, std=stddev) + if l.bias is not None: + nn.init.constant_(l.bias, 0.0) + +# Multihead attention module +class MultiheadAttention(nn.Module): + def __init__( + self, + *, + device: torch.device, + dtype: torch.dtype, + n_ctx: int, # Context size + width: int, # Width of the input tensor + heads: int, # Number of attention heads + init_scale: float, # Initialization scale for weights + qkv_bias: bool, # Whether to use bias in QKV layers + flash: bool = False # Whether to use flash attention + ): + super().__init__() + self.n_ctx = n_ctx + self.width = width + self.heads = heads + self.c_qkv = nn.Linear(width, width * 3, bias=qkv_bias, device=device, dtype=dtype) + self.c_proj = nn.Linear(width, width, device=device, dtype=dtype) + self.attention = QKVMultiheadAttention(device=device, dtype=dtype, heads=heads, n_ctx=n_ctx, flash=flash) + init_linear(self.c_qkv, init_scale) + init_linear(self.c_proj, init_scale) + + def forward(self, x): + x = self.c_qkv(x) + x = checkpoint(self.attention, (x,), (), True) + x = self.c_proj(x) + return x + +# QKV multihead attention module +class QKVMultiheadAttention(nn.Module): + def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int, n_ctx: int, flash: bool = False): + super().__init__() + self.device = device + self.dtype = dtype + self.heads = heads + self.n_ctx = n_ctx + self.flash = flash + + def forward(self, qkv): + bs, n_ctx, width = qkv.shape + attn_ch = width // self.heads // 3 + scale = 1 / math.sqrt(math.sqrt(attn_ch)) + qkv = qkv.view(bs, n_ctx, self.heads, -1) + q, k, v = torch.split(qkv, attn_ch, dim=-1) + + if self.flash: + # print(1) + # # out = F.scaled_dot_product_attention(q, k, v) + # q = q.permute(0, 2, 1, 3) # -> (b, h, t, c) + # k = k.permute(0, 2, 1, 3) # -> (b, h, t, c) + # v = v.permute(0, 2, 1, 3) # -> (b, h, t, c) + + # # F.scaled_dot_product_attention 直接返回计算结果，无需手动 reshape + # out = F.scaled_dot_product_attention(q, k, v) # -> (b, h, t, c) + + # # 关键：将输出变回 (b, t, h, c) 并 reshape，以匹配 else 分支的输出 + # out = out.permute(0, 2, 1, 3).reshape(bs, n_ctx, -1) # -> (b, t, h*c) + out = torch_npu.npu_fusion_attention( + q, k, v, self.heads, "BSND", keep_prob=1.0, scale=scale)[0].reshape(bs, n_ctx, -1) + else: + weight = torch.einsum( + "bthc,bshc->bhts", q * scale, k * scale + ) # More stable with f16 than dividing afterwards + wdtype = weight.dtype + weight = torch.softmax(weight.float(), dim=-1).type(wdtype) + out = torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1) + + return out + +# Residual attention block module +class ResidualAttentionBlock(nn.Module): + def __init__( + self, + *, + device: torch.device, + dtype: torch.dtype, + use_checkpoint: bool = False, + n_ctx: int, # Context size + width: int, # Width of the input tensor + heads: int, # Number of attention heads + init_scale: float, # Initialization scale for weights + qkv_bias: bool, # Whether to use bias in QKV layers + flash: bool = False # Whether to use flash attention + ): + super().__init__() + + self.use_checkpoint = use_checkpoint + + self.attn = MultiheadAttention( + device=device, + dtype=dtype, + n_ctx=n_ctx, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash + ) + self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype) + self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=init_scale) + self.ln_2 = nn.LayerNorm(width, device=device, dtype=dtype) + + def _forward(self, x: torch.Tensor): + x = x + self.attn(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + + def forward(self, x: torch.Tensor): + return checkpoint(self._forward, (x,), self.parameters(), self.use_checkpoint) + +# Multihead cross attention module +class MultiheadCrossAttention(nn.Module): + def __init__( + self, + *, + device: torch.device, + dtype: torch.dtype, + n_data: Optional[int] = None, + data_width: Optional[int] = None, + width: int, # Width of the input tensor + heads: int, # Number of attention heads + init_scale: float, # Initialization scale for weights + qkv_bias: bool, # Whether to use bias in QKV layers + flash: bool = False # Whether to use flash attention + ): + super().__init__() + self.n_data = n_data + self.width = width + self.heads = heads + self.data_width = width if data_width is None else data_width + self.c_q = nn.Linear(width, width, bias=qkv_bias, device=device, dtype=dtype) + self.c_kv = nn.Linear(self.data_width, width * 2, bias=qkv_bias, device=device, dtype=dtype) + self.c_proj = nn.Linear(width, width, device=device, dtype=dtype) + self.attention = QKVMultiheadCrossAttention( + device=device, dtype=dtype, heads=heads, n_data=n_data, flash=flash + ) + init_linear(self.c_q, init_scale) + init_linear(self.c_kv, init_scale) + init_linear(self.c_proj, init_scale) + + def forward(self, x, data): + x = self.c_q(x) + data = self.c_kv(data) + x = checkpoint(self.attention, (x, data), (), True) + x = self.c_proj(x) + return x + +# QKV multihead cross attention module +class QKVMultiheadCrossAttention(nn.Module): + def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int, + flash: bool = False, n_data: Optional[int] = None): + + super().__init__() + self.device = device + self.dtype = dtype + self.heads = heads + self.n_data = n_data + self.flash = flash + + def forward(self, q, kv): + _, n_ctx, _ = q.shape + bs, n_data, width = kv.shape + attn_ch = width // self.heads // 2 + scale = 1 / math.sqrt(math.sqrt(attn_ch)) + q = q.view(bs, n_ctx, self.heads, -1) + kv = kv.view(bs, n_data, self.heads, -1) + k, v = torch.split(kv, attn_ch, dim=-1) + + if self.flash: + # print(2) + # # out = F.scaled_dot_product_attention(q, k, v) + # q = q.permute(0, 2, 1, 3) # -> (b, h, t, c) + # k = k.permute(0, 2, 1, 3) # -> (b, h, t, c) + # v = v.permute(0, 2, 1, 3) # -> (b, h, t, c) + + # # F.scaled_dot_product_attention 直接返回计算结果，无需手动 reshape + # out = F.scaled_dot_product_attention(q, k, v) # -> (b, h, t, c) + + # # 关键：将输出变回 (b, t, h, c) 并 reshape，以匹配 else 分支的输出 + # out = out.permute(0, 2, 1, 3).reshape(bs, n_ctx, -1) # -> (b, t, h*c) + out = torch_npu.npu_fusion_attention( + q, k, v, self.heads, "BSND", keep_prob=1.0, scale=scale)[0].reshape(bs, n_ctx, -1) + else: + weight = torch.einsum( + "bthc,bshc->bhts", q * scale, k * scale + ) # More stable with f16 than dividing afterwards + wdtype = weight.dtype + weight = torch.softmax(weight.float(), dim=-1).type(wdtype) + out = torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1) + + return out + +# Residual cross attention block module +class ResidualCrossAttentionBlock(nn.Module): + def __init__( + self, + *, + device: Optional[torch.device], + dtype: Optional[torch.dtype], + n_data: Optional[int] = None, + data_width: Optional[int] = None, + width: int, # Width of the input tensor + heads: int, # Number of attention heads + init_scale: float, # Initialization scale for weights + qkv_bias: bool, # Whether to use bias in QKV layers + flash: bool = False # Whether to use flash attention + ): + super().__init__() + + if data_width is None: + data_width = width + + self.attn = MultiheadCrossAttention( + device=device, + dtype=dtype, + n_data=n_data, + width=width, + heads=heads, + data_width=data_width, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash, + ) + self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype) + self.ln_2 = nn.LayerNorm(data_width, device=device, dtype=dtype) + self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=init_scale) + self.ln_3 = nn.LayerNorm(width, device=device, dtype=dtype) + + def forward(self, x: torch.Tensor, data: torch.Tensor): + x = x + self.attn(self.ln_1(x), self.ln_2(data)) + x = x + self.mlp(self.ln_3(x)) + return x + +# MLP Module +class MLP(nn.Module): + def __init__(self, *, + device: Optional[torch.device], + dtype: Optional[torch.dtype], + width: int, + init_scale: float): + super().__init__() + self.width = width + self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype) + self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype) + self.gelu = nn.GELU() + init_linear(self.c_fc, init_scale) + init_linear(self.c_proj, init_scale) + + def forward(self, x): + return self.c_proj(self.gelu(self.c_fc(x))) + +# Transformer Module +class Transformer(nn.Module): + def __init__( + self, + *, + device: Optional[torch.device], + dtype: Optional[torch.dtype], + layers: int, + use_checkpoint: bool = False, + n_ctx: int, # Context size + width: int, # Width of the input tensor + heads: int, # Number of attention heads + init_scale: float, # Initialization scale for weights + qkv_bias: bool, # Whether to use bias in QKV layers + flash: bool = False # Whether to use flash attention + ): + super().__init__() + self.n_ctx = n_ctx + self.width = width + self.layers = layers + self.resblocks = nn.ModuleList( + [ + ResidualAttentionBlock( + device=device, + dtype=dtype, + n_ctx=n_ctx, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash, + use_checkpoint=use_checkpoint + ) + for _ in range(layers) + ] + ) + + def forward(self, x: torch.Tensor): + for block in self.resblocks: + x = block(x) + return x diff --git a/mulres_infer/modules/miche/michelangelo/models/tsal/__init__.py b/mulres_infer/modules/miche/michelangelo/models/tsal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4c48b5acbf4ae09a163f0f8fa01a91341695211f --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/tsal/__init__.py @@ -0,0 +1 @@ +# -*- coding: utf-8 -*- diff --git a/mulres_infer/modules/miche/michelangelo/models/tsal/asl_pl_module.py b/mulres_infer/modules/miche/michelangelo/models/tsal/asl_pl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..a03d2331d0592d6a12334abcabfa5fac05f84231 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/tsal/asl_pl_module.py @@ -0,0 +1,384 @@ +# -*- coding: utf-8 -*- + +from typing import List, Tuple, Dict, Optional +from omegaconf import DictConfig + +import torch +import torch.nn.functional as F +from torch import nn +from torch.optim import lr_scheduler +from typing import Union +from functools import partial + +from modules.miche.michelangelo.utils import instantiate_from_config + +from .tsal_base import ( + AlignedShapeAsLatentModule, + ShapeAsLatentModule, + Latent2MeshOutput, + AlignedMeshOutput +) +from modules.miche.michelangelo.models.tsal.inference_utils import extract_geometry +import trimesh + +class AlignedShapeAsLatentPLModule(nn.Module): + def __init__(self, *, + shape_module_cfg, + aligned_module_cfg, + loss_cfg, + optimizer_cfg: Optional[DictConfig] = None, + ckpt_path: Optional[str] = None, + ignore_keys: Union[Tuple[str], List[str]] = ()): + + super().__init__() + + shape_model: ShapeAsLatentModule = instantiate_from_config( + shape_module_cfg, device=None, dtype=None + ) + self.model: AlignedShapeAsLatentModule = instantiate_from_config( + aligned_module_cfg, shape_model=shape_model + ) + + self.loss = instantiate_from_config(loss_cfg) + + self.optimizer_cfg = optimizer_cfg + + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) + + def set_shape_model_only(self): + self.model.set_shape_model_only() + + @property + def latent_shape(self): + return self.model.shape_model.latent_shape + + @property + def zero_rank(self): + if self._trainer: + zero_rank = self.trainer.local_rank == 0 + else: + zero_rank = True + + return zero_rank + + def init_from_ckpt(self, path, ignore_keys=()): + state_dict = torch.load(path, map_location="cpu")["state_dict"] + + keys = list(state_dict.keys()) + for k in keys: + for ik in ignore_keys: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del state_dict[k] + + missing, unexpected = self.load_state_dict(state_dict, strict=False) + print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") + if len(missing) > 0: + print(f"Missing Keys: {missing}") + print(f"Unexpected Keys: {unexpected}") + + def configure_optimizers(self) -> Tuple[List, List]: + lr = self.learning_rate + + trainable_parameters = list(self.model.parameters()) + + if self.optimizer_cfg is None: + optimizers = [torch.optim.AdamW(trainable_parameters, lr=lr, betas=(0.9, 0.99), weight_decay=1e-3)] + schedulers = [] + else: + optimizer = instantiate_from_config(self.optimizer_cfg.optimizer, params=trainable_parameters) + scheduler_func = instantiate_from_config( + self.optimizer_cfg.scheduler, + max_decay_steps=self.trainer.max_steps, + lr_max=lr + ) + scheduler = { + "scheduler": lr_scheduler.LambdaLR(optimizer, lr_lambda=scheduler_func.schedule), + "interval": "step", + "frequency": 1 + } + optimizers = [optimizer] + schedulers = [scheduler] + + return optimizers, schedulers + + def forward(self, + surface: torch.FloatTensor, + image: torch.FloatTensor, + text: torch.FloatTensor, + volume_queries: torch.FloatTensor): + # Args: + # surface (torch.FloatTensor): + # image (torch.FloatTensor): + # text (torch.FloatTensor): + # volume_queries (torch.FloatTensor): + # + # Returns: + + embed_outputs, shape_z = self.model(surface, image, text) + + shape_zq, posterior = self.model.shape_model.encode_kl_embed(shape_z) + latents = self.model.shape_model.decode(shape_zq) + logits = self.model.shape_model.query_geometry(volume_queries, latents) + + return embed_outputs, logits, posterior + + def encode(self, surface: torch.FloatTensor, sample_posterior=True): + + pc = surface[..., 0:3] + feats = surface[..., 3:6] + + shape_embed, shape_zq, posterior = self.model.shape_model.encode( + pc=pc, feats=feats, sample_posterior=sample_posterior + ) + + return shape_zq + + def encode_latents(self, surface: torch.FloatTensor, img_feats: Optional[torch.FloatTensor] = None): + + pc = surface[..., 0:3] + feats = surface[..., 3:6] + + shape_embed, shape_latents = self.model.shape_model.encode_latents( + pc=pc, feats=feats, img_feats=img_feats, + ) + shape_embed = shape_embed.unsqueeze(1) + # 长度无关: latent 数可被 PointConditioner._expand_latents 扩容(如 1024), 只校验头 token + assert shape_embed.shape[1] == 1 + cat_latents = torch.cat([shape_embed, shape_latents], dim=1) + + return cat_latents + + def recon(self, surface): + cat_latents = self.encode_latents(surface) + shape_latents = cat_latents[:, 1:] + shape_zq, posterior = self.model.shape_model.encode_kl_embed(shape_latents) + + # decoding + latents = self.model.shape_model.decode(shape_zq) + geometric_func = partial(self.model.shape_model.query_geometry, latents=latents) + + # reconstruction + mesh_v_f, has_surface = extract_geometry( + geometric_func=geometric_func, + device=surface.device, + batch_size=surface.shape[0], + bounds=(-1.25, -1.25, -1.25, 1.25, 1.25, 1.25), + octree_depth=7, + num_chunks=10000, + ) + recon_mesh = trimesh.Trimesh(mesh_v_f[0][0], mesh_v_f[0][1]) + + return recon_mesh + + + def to_shape_latents(self, latents): + + shape_zq, posterior = self.model.shape_model.encode_kl_embed(latents, sample_posterior = False) + return self.model.shape_model.decode(shape_zq) + + def decode(self, + z_q, + bounds: Union[Tuple[float], List[float], float] = 1.1, + octree_depth: int = 7, + num_chunks: int = 10000) -> List[Latent2MeshOutput]: + + latents = self.model.shape_model.decode(z_q) # latents: [bs, num_latents, dim] + outputs = self.latent2mesh(latents, bounds=bounds, octree_depth=octree_depth, num_chunks=num_chunks) + + return outputs + + def training_step(self, batch: Dict[str, torch.FloatTensor], + batch_idx: int, optimizer_idx: int = 0) -> torch.FloatTensor: + #Args: + # batch (dict): the batch sample, and it contains: + # - surface (torch.FloatTensor): [bs, n_surface, (3 + input_dim)] + # - image (torch.FloatTensor): [bs, 3, 224, 224] + # - text (torch.FloatTensor): [bs, num_templates, 77] + # - geo_points (torch.FloatTensor): [bs, n_pts, (3 + 1)] + # + # batch_idx (int): + # + # optimizer_idx (int): + # + # Returns: + # loss (torch.FloatTensor): + + surface = batch["surface"] + image = batch["image"] + text = batch["text"] + + volume_queries = batch["geo_points"][..., 0:3] + shape_labels = batch["geo_points"][..., -1] + + embed_outputs, shape_logits, posteriors = self(surface, image, text, volume_queries) + + aeloss, log_dict_ae = self.loss( + **embed_outputs, + posteriors=posteriors, + shape_logits=shape_logits, + shape_labels=shape_labels, + split="train" + ) + + self.log_dict(log_dict_ae, prog_bar=True, logger=True, batch_size=shape_logits.shape[0], + sync_dist=False, rank_zero_only=True) + + return aeloss + + def validation_step(self, batch: Dict[str, torch.FloatTensor], batch_idx: int) -> torch.FloatTensor: + + surface = batch["surface"] + image = batch["image"] + text = batch["text"] + + volume_queries = batch["geo_points"][..., 0:3] + shape_labels = batch["geo_points"][..., -1] + + embed_outputs, shape_logits, posteriors = self(surface, image, text, volume_queries) + + aeloss, log_dict_ae = self.loss( + **embed_outputs, + posteriors=posteriors, + shape_logits=shape_logits, + shape_labels=shape_labels, + split="val" + ) + self.log_dict(log_dict_ae, prog_bar=True, logger=True, batch_size=shape_logits.shape[0], + sync_dist=False, rank_zero_only=True) + + return aeloss + + def visual_alignment(self, + surface: torch.FloatTensor, + image: torch.FloatTensor, + text: torch.FloatTensor, + description: Optional[List[str]] = None, + bounds: Union[Tuple[float], List[float]] = (-1.25, -1.25, -1.25, 1.25, 1.25, 1.25), + octree_depth: int = 7, + num_chunks: int = 10000) -> List[AlignedMeshOutput]: + + """ + + Args: + surface: + image: + text: + description: + bounds: + octree_depth: + num_chunks: + + Returns: + mesh_outputs (List[AlignedMeshOutput]): the mesh outputs list. + + """ + + outputs = [] + + device = surface.device + bs = surface.shape[0] + + embed_outputs, shape_z = self.model(surface, image, text) + + # calculate the similarity + image_embed = embed_outputs["image_embed"] + text_embed = embed_outputs["text_embed"] + shape_embed = embed_outputs["shape_embed"] + + # normalized features + shape_embed = F.normalize(shape_embed, dim=-1, p=2) + text_embed = F.normalize(text_embed, dim=-1, p=2) + image_embed = F.normalize(image_embed, dim=-1, p=2) + + # B x B + shape_text_similarity = (100.0 * shape_embed @ text_embed.T).softmax(dim=-1) + + # B x B + shape_image_similarity = (100.0 * shape_embed @ image_embed.T).softmax(dim=-1) + + # shape reconstruction + shape_zq, posterior = self.model.shape_model.encode_kl_embed(shape_z) + latents = self.model.shape_model.decode(shape_zq) + geometric_func = partial(self.model.shape_model.query_geometry, latents=latents) + + # 2. decode geometry + mesh_v_f, has_surface = extract_geometry( + geometric_func=geometric_func, + device=device, + batch_size=bs, + bounds=bounds, + octree_depth=octree_depth, + num_chunks=num_chunks, + disable=not self.zero_rank + ) + + # 3. decode texture + for i, ((mesh_v, mesh_f), is_surface) in enumerate(zip(mesh_v_f, has_surface)): + if not is_surface: + outputs.append(None) + continue + + out = AlignedMeshOutput() + out.mesh_v = mesh_v + out.mesh_f = mesh_f + out.surface = surface[i].cpu().numpy() + out.image = image[i].cpu().numpy() + if description is not None: + out.text = description[i] + out.shape_text_similarity = shape_text_similarity[i, i] + out.shape_image_similarity = shape_image_similarity[i, i] + + outputs.append(out) + + return outputs + + def latent2mesh(self, + latents: torch.FloatTensor, + bounds: Union[Tuple[float], List[float], float] = 1.1, + octree_depth: int = 7, + num_chunks: int = 10000) -> List[Latent2MeshOutput]: + + """ + + Args: + latents: [bs, num_latents, dim] + bounds: + octree_depth: + num_chunks: + + Returns: + mesh_outputs (List[MeshOutput]): the mesh outputs list. + + """ + + outputs = [] + + geometric_func = partial(self.model.shape_model.query_geometry, latents=latents) + + # 2. decode geometry + device = latents.device + mesh_v_f, has_surface = extract_geometry( + geometric_func=geometric_func, + device=device, + batch_size=len(latents), + bounds=bounds, + octree_depth=octree_depth, + num_chunks=num_chunks, + disable=not self.zero_rank + ) + + # 3. decode texture + for i, ((mesh_v, mesh_f), is_surface) in enumerate(zip(mesh_v_f, has_surface)): + if not is_surface: + outputs.append(None) + continue + + out = Latent2MeshOutput() + out.mesh_v = mesh_v + out.mesh_f = mesh_f + + outputs.append(out) + + return outputs diff --git a/mulres_infer/modules/miche/michelangelo/models/tsal/clip_asl_module.py b/mulres_infer/modules/miche/michelangelo/models/tsal/clip_asl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..45ceabd5c7e8725a1969cd699bee64977cd1981e --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/tsal/clip_asl_module.py @@ -0,0 +1,118 @@ +# -*- coding: utf-8 -*- + +import torch +from torch import nn +from einops import rearrange +from transformers import CLIPModel + +from modules.miche.michelangelo.models.tsal.tsal_base import AlignedShapeAsLatentModule + + +class CLIPAlignedShapeAsLatentModule(AlignedShapeAsLatentModule): + + def __init__(self, *, + shape_model, + clip_model_version: str = "openai/clip-vit-large-patch14"): + + super().__init__() + + # self.clip_model: CLIPModel = CLIPModel.from_pretrained(clip_model_version) + # for params in self.clip_model.parameters(): + # params.requires_grad = False + self.clip_model = None + self.shape_model = shape_model + self.shape_projection = nn.Parameter(torch.empty(self.shape_model.width, self.shape_model.width)) + # nn.init.normal_(self.shape_projection, std=self.shape_model.width ** -0.5) + + def set_shape_model_only(self): + self.clip_model = None + + def encode_shape_embed(self, surface, return_latents: bool = False): + """ + + Args: + surface (torch.FloatTensor): [bs, n, 3 + c] + return_latents (bool): + + Returns: + x (torch.FloatTensor): [bs, projection_dim] + shape_latents (torch.FloatTensor): [bs, m, d] + """ + + pc = surface[..., 0:3] + feats = surface[..., 3:] + + shape_embed, shape_latents = self.shape_model.encode_latents(pc, feats) + x = shape_embed @ self.shape_projection + + if return_latents: + return x, shape_latents + else: + return x + + def encode_image_embed(self, image): + """ + + Args: + image (torch.FloatTensor): [bs, 3, h, w] + + Returns: + x (torch.FloatTensor): [bs, projection_dim] + """ + + x = self.clip_model.get_image_features(image) + + return x + + def encode_text_embed(self, text): + x = self.clip_model.get_text_features(text) + return x + + def forward(self, surface, image, text): + """ + + Args: + surface (torch.FloatTensor): + image (torch.FloatTensor): [bs, 3, 224, 224] + text (torch.LongTensor): [bs, num_templates, 77] + + Returns: + embed_outputs (dict): the embedding outputs, and it contains: + - image_embed (torch.FloatTensor): + - text_embed (torch.FloatTensor): + - shape_embed (torch.FloatTensor): + - logit_scale (float): + """ + + # # text embedding + # text_embed_all = [] + # for i in range(text.shape[0]): + # text_for_one_sample = text[i] + # text_embed = self.encode_text_embed(text_for_one_sample) + # text_embed = text_embed / text_embed.norm(dim=-1, keepdim=True) + # text_embed = text_embed.mean(dim=0) + # text_embed = text_embed / text_embed.norm(dim=-1, keepdim=True) + # text_embed_all.append(text_embed) + # text_embed_all = torch.stack(text_embed_all) + + b = text.shape[0] + text_tokens = rearrange(text, "b t l -> (b t) l") + text_embed = self.encode_text_embed(text_tokens) + text_embed = rearrange(text_embed, "(b t) d -> b t d", b=b) + text_embed = text_embed.mean(dim=1) + text_embed = text_embed / text_embed.norm(dim=-1, keepdim=True) + + # image embedding + image_embed = self.encode_image_embed(image) + + # shape embedding + shape_embed, shape_latents = self.encode_shape_embed(surface, return_latents=True) + + embed_outputs = { + "image_embed": image_embed, + "text_embed": text_embed, + "shape_embed": shape_embed, + # "logit_scale": self.clip_model.logit_scale.exp() + } + + return embed_outputs, shape_latents diff --git a/mulres_infer/modules/miche/michelangelo/models/tsal/inference_utils.py b/mulres_infer/modules/miche/michelangelo/models/tsal/inference_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e6e0a9c2910adfc8cd1f174b6bbf217043e880df --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/tsal/inference_utils.py @@ -0,0 +1,76 @@ +# -*- coding: utf-8 -*- + +import torch +from tqdm import tqdm +from einops import repeat +import numpy as np +from typing import Callable, Tuple, List, Union, Optional +from skimage import measure + +from modules.miche.michelangelo.graphics.primitives import generate_dense_grid_points + + +@torch.no_grad() +def extract_geometry(geometric_func: Callable, + device: torch.device, + batch_size: int = 1, + bounds: Union[Tuple[float], List[float], float] = (-1.25, -1.25, -1.25, 1.25, 1.25, 1.25), + octree_depth: int = 7, + num_chunks: int = 10000, + disable: bool = True): + + # Args: + # geometric_func: + # device: + # bounds: + # octree_depth: + # batch_size: + # num_chunks: + # disable: + # Returns: + + if isinstance(bounds, float): + bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] + + bbox_min = np.array(bounds[0:3]) + bbox_max = np.array(bounds[3:6]) + bbox_size = bbox_max - bbox_min + + xyz_samples, grid_size, length = generate_dense_grid_points( + bbox_min=bbox_min, + bbox_max=bbox_max, + octree_depth=octree_depth, + indexing="ij" + ) + xyz_samples = torch.FloatTensor(xyz_samples) + + batch_logits = [] + for start in tqdm(range(0, xyz_samples.shape[0], num_chunks), + desc="Implicit Function:", disable=disable, leave=False): + queries = xyz_samples[start: start + num_chunks, :].to(device) + batch_queries = repeat(queries, "p c -> b p c", b=batch_size) + + logits = geometric_func(batch_queries) + batch_logits.append(logits.cpu()) + + grid_logits = torch.cat(batch_logits, dim=1).view((batch_size, grid_size[0], grid_size[1], grid_size[2])).numpy() + + mesh_v_f = [] + has_surface = np.zeros((batch_size,), dtype=np.bool_) + for i in range(batch_size): + try: + vertices, faces, normals, _ = measure.marching_cubes(grid_logits[i], 0, method="lewiner") + vertices = vertices / grid_size * bbox_size + bbox_min + # vertices[:, [0, 1]] = vertices[:, [1, 0]] + mesh_v_f.append((vertices.astype(np.float32), np.ascontiguousarray(faces))) + has_surface[i] = True + + except ValueError: + mesh_v_f.append((None, None)) + has_surface[i] = False + + except RuntimeError: + mesh_v_f.append((None, None)) + has_surface[i] = False + + return mesh_v_f, has_surface diff --git a/mulres_infer/modules/miche/michelangelo/models/tsal/loss.py b/mulres_infer/modules/miche/michelangelo/models/tsal/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..43ce85f0abb7d97c4e8fd1164dcf4fca287d7386 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/tsal/loss.py @@ -0,0 +1,130 @@ +# -*- coding: utf-8 -*- +import torch +import torch.nn as nn +import torch.nn.functional as F + +from typing import Optional + +from modules.miche.michelangelo.models.modules.distributions import DiagonalGaussianDistribution +from modules.miche.michelangelo.utils import misc + + +class ContrastKLNearFar(nn.Module): + def __init__(self, + contrast_weight: float = 1.0, + near_weight: float = 0.1, + kl_weight: float = 1.0, + num_near_samples: Optional[int] = None): + + super().__init__() + + self.labels = None + self.last_local_batch_size = None + + self.contrast_weight = contrast_weight + self.near_weight = near_weight + self.kl_weight = kl_weight + self.num_near_samples = num_near_samples + self.geo_criterion = nn.BCEWithLogitsLoss() + + def forward(self, + shape_embed: torch.FloatTensor, + text_embed: torch.FloatTensor, + image_embed: torch.FloatTensor, + logit_scale: torch.FloatTensor, + posteriors: Optional[DiagonalGaussianDistribution], + shape_logits: torch.FloatTensor, + shape_labels: torch.FloatTensor, + split: Optional[str] = "train", **kwargs): + + # shape_embed: torch.FloatTensor + # text_embed: torch.FloatTensor + # image_embed: torch.FloatTensor + # logit_scale: torch.FloatTensor + # posteriors: Optional[DiagonalGaussianDistribution] + # shape_logits: torch.FloatTensor + # shape_labels: torch.FloatTensor + + local_batch_size = shape_embed.size(0) + + if local_batch_size != self.last_local_batch_size: + self.labels = local_batch_size * misc.get_rank() + torch.arange( + local_batch_size, device=shape_embed.device + ).long() + self.last_local_batch_size = local_batch_size + + # normalized features + shape_embed = F.normalize(shape_embed, dim=-1, p=2) + text_embed = F.normalize(text_embed, dim=-1, p=2) + image_embed = F.normalize(image_embed, dim=-1, p=2) + + # gather features from all GPUs + shape_embed_all, text_embed_all, image_embed_all = misc.all_gather_batch( + [shape_embed, text_embed, image_embed] + ) + + # cosine similarity as logits + logits_per_shape_text = logit_scale * shape_embed @ text_embed_all.t() + logits_per_text_shape = logit_scale * text_embed @ shape_embed_all.t() + logits_per_shape_image = logit_scale * shape_embed @ image_embed_all.t() + logits_per_image_shape = logit_scale * image_embed @ shape_embed_all.t() + contrast_loss = (F.cross_entropy(logits_per_shape_text, self.labels) + + F.cross_entropy(logits_per_text_shape, self.labels)) / 2 + \ + (F.cross_entropy(logits_per_shape_image, self.labels) + + F.cross_entropy(logits_per_image_shape, self.labels)) / 2 + + # shape reconstruction + if self.num_near_samples is None: + num_vol = shape_logits.shape[1] // 2 + else: + num_vol = shape_logits.shape[1] - self.num_near_samples + + vol_logits = shape_logits[:, 0:num_vol] + vol_labels = shape_labels[:, 0:num_vol] + + near_logits = shape_logits[:, num_vol:] + near_labels = shape_labels[:, num_vol:] + + # occupancy loss + vol_bce = self.geo_criterion(vol_logits.float(), vol_labels.float()) + near_bce = self.geo_criterion(near_logits.float(), near_labels.float()) + + if posteriors is None: + kl_loss = torch.tensor(0.0, dtype=vol_logits.dtype, device=vol_logits.device) + else: + kl_loss = posteriors.kl(dims=(1, 2)) + kl_loss = torch.mean(kl_loss) + + loss = vol_bce + near_bce * self.near_weight + kl_loss * self.kl_weight + contrast_loss * self.contrast_weight + + # compute accuracy + with torch.no_grad(): + pred = torch.argmax(logits_per_shape_text, dim=-1) + correct = pred.eq(self.labels).sum() + shape_text_acc = 100 * correct / local_batch_size + + pred = torch.argmax(logits_per_shape_image, dim=-1) + correct = pred.eq(self.labels).sum() + shape_image_acc = 100 * correct / local_batch_size + + preds = shape_logits >= 0 + accuracy = (preds == shape_labels).float() + accuracy = accuracy.mean() + + log = { + "{}/contrast".format(split): contrast_loss.clone().detach(), + "{}/near".format(split): near_bce.detach(), + "{}/far".format(split): vol_bce.detach(), + "{}/kl".format(split): kl_loss.detach(), + "{}/shape_text_acc".format(split): shape_text_acc, + "{}/shape_image_acc".format(split): shape_image_acc, + "{}/total_loss".format(split): loss.clone().detach(), + "{}/accuracy".format(split): accuracy, + } + + if posteriors is not None: + log[f"{split}/mean"] = posteriors.mean.mean().detach() + log[f"{split}/std_mean"] = posteriors.std.mean().detach() + log[f"{split}/std_max"] = posteriors.std.max().detach() + + return loss, log diff --git a/mulres_infer/modules/miche/michelangelo/models/tsal/sal_perceiver.py b/mulres_infer/modules/miche/michelangelo/models/tsal/sal_perceiver.py new file mode 100644 index 0000000000000000000000000000000000000000..a2610867a567ede957bf99a77fe83147a86e8042 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/tsal/sal_perceiver.py @@ -0,0 +1,459 @@ +# -*- coding: utf-8 -*- + +import torch +import torch.nn as nn +from typing import Optional +from einops import repeat +import math + +from modules.miche.michelangelo.models.modules import checkpoint +from modules.miche.michelangelo.models.modules.embedder import FourierEmbedder +from modules.miche.michelangelo.models.modules.distributions import DiagonalGaussianDistribution +from modules.miche.michelangelo.models.modules.transformer_blocks import ( + ResidualCrossAttentionBlock, + Transformer +) + +from .tsal_base import ShapeAsLatentModule + + +class CrossAttentionEncoder(nn.Module): + + def __init__(self, *, + device: Optional[torch.device], + dtype: Optional[torch.dtype], + num_latents: int, + fourier_embedder: FourierEmbedder, + point_feats: int, + width: int, + heads: int, + layers: int, + init_scale: float = 0.25, + qkv_bias: bool = True, + flash: bool = False, + use_ln_post: bool = False, + use_checkpoint: bool = False, + img_feat_dim: int = 0): + + super().__init__() + + self.use_checkpoint = use_checkpoint + self.num_latents = num_latents + + self.query = nn.Parameter(torch.randn((num_latents, width), device=device, dtype=dtype) * 0.02) + + self.fourier_embedder = fourier_embedder + self.point_feats = point_feats + self.img_feat_dim = img_feat_dim + self.input_proj = nn.Linear( + self.fourier_embedder.out_dim + point_feats + img_feat_dim, + width, + device=device, + dtype=dtype, + ) + self.cross_attn = ResidualCrossAttentionBlock( + device=device, + dtype=dtype, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash, + ) + + self.self_attn = Transformer( + device=device, + dtype=dtype, + n_ctx=num_latents, + width=width, + layers=layers, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash, + use_checkpoint=False + ) + + if use_ln_post: + self.ln_post = nn.LayerNorm(width, dtype=dtype, device=device) + else: + self.ln_post = None + + def _forward(self, pc, feats, img_feats=None): + + # Args: + # pc (torch.FloatTensor): [B, N, 3] + # feats (torch.FloatTensor or None): [B, N, C] + # img_feats (torch.FloatTensor or None): [B, N, img_feat_dim] + + bs = pc.shape[0] + + parts = [self.fourier_embedder(pc)] + if feats is not None: + parts.append(feats) + if self.img_feat_dim > 0: + if img_feats is None: + raise ValueError("img_feats is required when img_feat_dim > 0") + parts.append(img_feats) + data = self.input_proj(torch.cat(parts, dim=-1)) + + query = repeat(self.query, "m c -> b m c", b=bs) + latents = self.cross_attn(query, data) + latents = self.self_attn(latents) + + if self.ln_post is not None: + latents = self.ln_post(latents) + + return latents, pc + + def forward( + self, + pc: torch.FloatTensor, + feats: Optional[torch.FloatTensor] = None, + img_feats: Optional[torch.FloatTensor] = None, + ): + + # Args: + # pc (torch.FloatTensor): [B, N, 3] + # feats (torch.FloatTensor or None): [B, N, C] + # img_feats (torch.FloatTensor or None): [B, N, img_feat_dim] + + # NOTE: Michelangelo's custom CheckpointFunction.backward does + # `x.detach()` on every entry of ctx.input_tensors and cannot tolerate + # None values. So we MUST NOT pass `img_feats=None` into the checkpoint + # inputs tuple. Build the inputs based on which optional args are + # actually tensors. + inputs = [pc] + if feats is not None: + inputs.append(feats) + if img_feats is not None: + inputs.append(img_feats) + + # _forward expects (pc, feats, img_feats). When some are None we cannot + # rely on positional defaults inside checkpoint, so wrap to restore them. + feats_is_none = feats is None + img_feats_is_none = img_feats is None + + def _forward_wrapper(*args): + it = iter(args) + pc_arg = next(it) + feats_arg = None if feats_is_none else next(it) + img_feats_arg = None if img_feats_is_none else next(it) + return self._forward(pc_arg, feats_arg, img_feats_arg) + + return checkpoint(_forward_wrapper, tuple(inputs), self.parameters(), self.use_checkpoint) + + +class CrossAttentionDecoder(nn.Module): + + def __init__(self, *, + device: Optional[torch.device], + dtype: Optional[torch.dtype], + num_latents: int, + out_channels: int, + fourier_embedder: FourierEmbedder, + width: int, + heads: int, + init_scale: float = 0.25, + qkv_bias: bool = True, + flash: bool = False, + use_checkpoint: bool = False): + + super().__init__() + + self.use_checkpoint = use_checkpoint + self.fourier_embedder = fourier_embedder + + self.query_proj = nn.Linear(self.fourier_embedder.out_dim, width, device=device, dtype=dtype) + + self.cross_attn_decoder = ResidualCrossAttentionBlock( + device=device, + dtype=dtype, + n_data=num_latents, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash + ) + + self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype) + self.output_proj = nn.Linear(width, out_channels, device=device, dtype=dtype) + + def _forward(self, queries: torch.FloatTensor, latents: torch.FloatTensor): + queries = self.query_proj(self.fourier_embedder(queries)) + x = self.cross_attn_decoder(queries, latents) + x = self.ln_post(x) + x = self.output_proj(x) + return x + + def forward(self, queries: torch.FloatTensor, latents: torch.FloatTensor): + return checkpoint(self._forward, (queries, latents), self.parameters(), self.use_checkpoint) + + +class ShapeAsLatentPerceiver(ShapeAsLatentModule): + def __init__(self, *, + device: Optional[torch.device], + dtype: Optional[torch.dtype], + num_latents: int, + point_feats: int = 0, + embed_dim: int = 0, + num_freqs: int = 8, + include_pi: bool = True, + width: int, + heads: int, + num_encoder_layers: int, + num_decoder_layers: int, + init_scale: float = 0.25, + qkv_bias: bool = True, + flash: bool = False, + use_ln_post: bool = False, + use_checkpoint: bool = False, + img_feat_dim: int = 0): + + super().__init__() + + self.use_checkpoint = use_checkpoint + + self.num_latents = num_latents + self.fourier_embedder = FourierEmbedder(num_freqs=num_freqs, include_pi=include_pi) + self.img_feat_dim = img_feat_dim + + init_scale = init_scale * math.sqrt(1.0 / width) + self.encoder = CrossAttentionEncoder( + device=device, + dtype=dtype, + fourier_embedder=self.fourier_embedder, + num_latents=num_latents, + point_feats=point_feats, + width=width, + heads=heads, + layers=num_encoder_layers, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash, + use_ln_post=use_ln_post, + use_checkpoint=use_checkpoint, + img_feat_dim=img_feat_dim, + ) + + self.embed_dim = embed_dim + if embed_dim > 0: + # VAE embed + self.pre_kl = nn.Linear(width, embed_dim * 2, device=device, dtype=dtype) + self.post_kl = nn.Linear(embed_dim, width, device=device, dtype=dtype) + self.latent_shape = (num_latents, embed_dim) + else: + self.latent_shape = (num_latents, width) + + self.transformer = Transformer( + device=device, + dtype=dtype, + n_ctx=num_latents, + width=width, + layers=num_decoder_layers, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash, + use_checkpoint=use_checkpoint + ) + + # geometry decoder + self.geo_decoder = CrossAttentionDecoder( + device=device, + dtype=dtype, + fourier_embedder=self.fourier_embedder, + out_channels=1, + num_latents=num_latents, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash, + use_checkpoint=use_checkpoint + ) + + def encode(self, + pc: torch.FloatTensor, + feats: Optional[torch.FloatTensor] = None, + sample_posterior: bool = True, + img_feats: Optional[torch.FloatTensor] = None): + + + # Args: + # pc (torch.FloatTensor): [B, N, 3] + # feats (torch.FloatTensor or None): [B, N, C] + # sample_posterior (bool): + + # Returns: + # latents (torch.FloatTensor) + # center_pos (torch.FloatTensor or None): + # posterior (DiagonalGaussianDistribution or None): + + + latents, center_pos = self.encoder(pc, feats, img_feats=img_feats) + + posterior = None + if self.embed_dim > 0: + moments = self.pre_kl(latents) + posterior = DiagonalGaussianDistribution(moments, feat_dim=-1) + + if sample_posterior: + latents = posterior.sample() + else: + latents = posterior.mode() + + return latents, center_pos, posterior + + def decode(self, latents: torch.FloatTensor): + latents = self.post_kl(latents) + return self.transformer(latents) + + def query_geometry(self, queries: torch.FloatTensor, latents: torch.FloatTensor): + logits = self.geo_decoder(queries, latents).squeeze(-1) + return logits + + def forward(self, + pc: torch.FloatTensor, + feats: torch.FloatTensor, + volume_queries: torch.FloatTensor, + sample_posterior: bool = True): + + # Args: + # pc (torch.FloatTensor): [B, N, 3] + # feats (torch.FloatTensor or None): [B, N, C] + # volume_queries (torch.FloatTensor): [B, P, 3] + # sample_posterior (bool): + + # Returns: + # logits (torch.FloatTensor): [B, P] + # center_pos (torch.FloatTensor): [B, M, 3] + # posterior (DiagonalGaussianDistribution or None). + + + + latents, center_pos, posterior = self.encode(pc, feats, sample_posterior=sample_posterior) + + latents = self.decode(latents) + logits = self.query_geometry(volume_queries, latents) + + return logits, center_pos, posterior + + +class AlignedShapeLatentPerceiver(ShapeAsLatentPerceiver): + + def __init__(self, *, + device: Optional[torch.device], + dtype: Optional[torch.dtype], + num_latents: int, + point_feats: int = 0, + embed_dim: int = 0, + num_freqs: int = 8, + include_pi: bool = True, + width: int, + heads: int, + num_encoder_layers: int, + num_decoder_layers: int, + init_scale: float = 0.25, + qkv_bias: bool = True, + flash: bool = False, + use_ln_post: bool = False, + use_checkpoint: bool = False, + img_feat_dim: int = 0): + + super().__init__( + device=device, + dtype=dtype, + num_latents=1 + num_latents, + point_feats=point_feats, + embed_dim=embed_dim, + num_freqs=num_freqs, + include_pi=include_pi, + width=width, + heads=heads, + num_encoder_layers=num_encoder_layers, + num_decoder_layers=num_decoder_layers, + init_scale=init_scale, + qkv_bias=qkv_bias, + flash=flash, + use_ln_post=use_ln_post, + use_checkpoint=use_checkpoint, + img_feat_dim=img_feat_dim, + ) + + self.width = width + + def encode(self, + pc: torch.FloatTensor, + feats: Optional[torch.FloatTensor] = None, + sample_posterior: bool = True, + img_feats: Optional[torch.FloatTensor] = None): + + # Args: + # pc (torch.FloatTensor): [B, N, 3] + # feats (torch.FloatTensor or None): [B, N, c] + # sample_posterior (bool): + + # Returns: + # shape_embed (torch.FloatTensor) + # kl_embed (torch.FloatTensor): + # posterior (DiagonalGaussianDistribution or None): + + + shape_embed, latents = self.encode_latents(pc, feats, img_feats=img_feats) + kl_embed, posterior = self.encode_kl_embed(latents, sample_posterior) + + return shape_embed, kl_embed, posterior + + def encode_latents(self, + pc: torch.FloatTensor, + feats: Optional[torch.FloatTensor] = None, + img_feats: Optional[torch.FloatTensor] = None): + + x, _ = self.encoder(pc, feats, img_feats=img_feats) + + shape_embed = x[:, 0] + latents = x[:, 1:] + + return shape_embed, latents + + def encode_kl_embed(self, latents: torch.FloatTensor, sample_posterior: bool = True): + posterior = None + if self.embed_dim > 0: + moments = self.pre_kl(latents) + posterior = DiagonalGaussianDistribution(moments, feat_dim=-1) + + if sample_posterior: + kl_embed = posterior.sample() + else: + kl_embed = posterior.mode() + else: + kl_embed = latents + + return kl_embed, posterior + + def forward(self, + pc: torch.FloatTensor, + feats: torch.FloatTensor, + volume_queries: torch.FloatTensor, + sample_posterior: bool = True): + + # Args: + # pc (torch.FloatTensor): [B, N, 3] + # feats (torch.FloatTensor or None): [B, N, C] + # volume_queries (torch.FloatTensor): [B, P, 3] + # sample_posterior (bool): + + # Returns: + # shape_embed (torch.FloatTensor): [B, projection_dim] + # logits (torch.FloatTensor): [B, M] + # posterior (DiagonalGaussianDistribution or None). + + + shape_embed, kl_embed, posterior = self.encode(pc, feats, sample_posterior=sample_posterior) + + latents = self.decode(kl_embed) + logits = self.query_geometry(volume_queries, latents) + + return shape_embed, logits, posterior diff --git a/mulres_infer/modules/miche/michelangelo/models/tsal/tsal_base.py b/mulres_infer/modules/miche/michelangelo/models/tsal/tsal_base.py new file mode 100644 index 0000000000000000000000000000000000000000..eae7275b4c3cb26d902c6ef754bb67d2f13fc712 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/models/tsal/tsal_base.py @@ -0,0 +1,125 @@ +# -*- coding: utf-8 -*- + +import torch.nn as nn +from typing import Tuple, List, Optional + +# Base class for output of Point to Mesh transformation +class Point2MeshOutput(object): + def __init__(self): + self.mesh_v = None # Vertices of the mesh + self.mesh_f = None # Faces of the mesh + self.center = None # Center of the mesh + self.pc = None # Point cloud data + + +# Base class for output of Latent to Mesh transformation +class Latent2MeshOutput(object): + def __init__(self): + self.mesh_v = None # Vertices of the mesh + self.mesh_f = None # Faces of the mesh + + +# Base class for output of Aligned Mesh transformation +class AlignedMeshOutput(object): + def __init__(self): + self.mesh_v = None # Vertices of the mesh + self.mesh_f = None # Faces of the mesh + self.surface = None # Surface data + self.image = None # Aligned image data + self.text: Optional[str] = None # Aligned text data + self.shape_text_similarity: Optional[float] = None # Similarity between shape and text + self.shape_image_similarity: Optional[float] = None # Similarity between shape and image + + +# Base class for Shape as Latent with Point to Mesh transformation module +class ShapeAsLatentPLModule(nn.Module): + latent_shape: Tuple[int] # Shape of the latent space + + def encode(self, surface, *args, **kwargs): + raise NotImplementedError + + def decode(self, z_q, *args, **kwargs): + raise NotImplementedError + + def latent2mesh(self, latents, *args, **kwargs) -> List[Latent2MeshOutput]: + raise NotImplementedError + + def point2mesh(self, *args, **kwargs) -> List[Point2MeshOutput]: + raise NotImplementedError + + +# Base class for Shape as Latent module +class ShapeAsLatentModule(nn.Module): + latent_shape: Tuple[int, int] # Shape of the latent space + + def __init__(self, *args, **kwargs): + super().__init__() + + def encode(self, *args, **kwargs): + raise NotImplementedError + + def decode(self, *args, **kwargs): + raise NotImplementedError + + def query_geometry(self, *args, **kwargs): + raise NotImplementedError + + +# Base class for Aligned Shape as Latent with Point to Mesh transformation module +class AlignedShapeAsLatentPLModule(nn.Module): + latent_shape: Tuple[int] # Shape of the latent space + + def set_shape_model_only(self): + raise NotImplementedError + + def encode(self, surface, *args, **kwargs): + raise NotImplementedError + + def decode(self, z_q, *args, **kwargs): + raise NotImplementedError + + def latent2mesh(self, latents, *args, **kwargs) -> List[Latent2MeshOutput]: + raise NotImplementedError + + def point2mesh(self, *args, **kwargs) -> List[Point2MeshOutput]: + raise NotImplementedError + + +# Base class for Aligned Shape as Latent module +class AlignedShapeAsLatentModule(nn.Module): + shape_model: ShapeAsLatentModule # Shape model module + latent_shape: Tuple[int, int] # Shape of the latent space + + + def __init__(self, *args, **kwargs): + super().__init__() + + def set_shape_model_only(self): + raise NotImplementedError + + def encode_image_embed(self, *args, **kwargs): + raise NotImplementedError + + def encode_text_embed(self, *args, **kwargs): + raise NotImplementedError + + def encode_shape_embed(self, *args, **kwargs): + raise NotImplementedError + +# Base class for Textured Shape as Latent module +class TexturedShapeAsLatentModule(nn.Module): + + def __init__(self, *args, **kwargs): + super().__init__() + + def encode(self, *args, **kwargs): + raise NotImplementedError + + def decode(self, *args, **kwargs): + raise NotImplementedError + + def query_geometry(self, *args, **kwargs): + raise NotImplementedError + + def query_color(self, *args, **kwargs): + raise NotImplementedError diff --git a/mulres_infer/modules/miche/michelangelo/utils/__init__.py b/mulres_infer/modules/miche/michelangelo/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..439f659e569036bbc4ffb4654f6e7edfde9b36c9 --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/utils/__init__.py @@ -0,0 +1,3 @@ +# -*- coding: utf-8 -*- + +from .misc import instantiate_from_config diff --git a/mulres_infer/modules/miche/michelangelo/utils/misc.py b/mulres_infer/modules/miche/michelangelo/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..a4326dc04a370803cd658121853568e8305e54ea --- /dev/null +++ b/mulres_infer/modules/miche/michelangelo/utils/misc.py @@ -0,0 +1,83 @@ +# -*- coding: utf-8 -*- + +import importlib + +import torch +import torch.distributed as dist + + + +def get_obj_from_str(string, reload=False): + module, cls = string.rsplit(".", 1) + if reload: + module_imp = importlib.import_module(module) + importlib.reload(module_imp) + return getattr(importlib.import_module(module, package=None), cls) + + +def get_obj_from_config(config): + if "target" not in config: + raise KeyError("Expected key `target` to instantiate.") + + return get_obj_from_str(config["target"]) + + +def instantiate_from_config(config, **kwargs): + if "target" not in config: + raise KeyError("Expected key `target` to instantiate.") + + cls = get_obj_from_str(config["target"]) + + params = config.get("params", dict()) + # params.update(kwargs) + # instance = cls(**params) + kwargs.update(params) + instance = cls(**kwargs) + + return instance + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def all_gather_batch(tensors): + """ + Performs all_gather operation on the provided tensors. + """ + # Queue the gathered tensors + world_size = get_world_size() + # There is no need for reduction in the single-proc case + if world_size == 1: + return tensors + tensor_list = [] + output_tensor = [] + for tensor in tensors: + tensor_all = [torch.ones_like(tensor) for _ in range(world_size)] + dist.all_gather( + tensor_all, + tensor, + async_op=False # performance opt + ) + + tensor_list.append(tensor_all) + + for tensor_all in tensor_list: + output_tensor.append(torch.cat(tensor_all, dim=0)) + return output_tensor diff --git a/mulres_infer/modules/miche/shapevae-256.yaml b/mulres_infer/modules/miche/shapevae-256.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8fe7ed0def1f7265eec38163cb5b881081f26e09 --- /dev/null +++ b/mulres_infer/modules/miche/shapevae-256.yaml @@ -0,0 +1,46 @@ +model: + target: modules.miche.michelangelo.models.tsal.asl_pl_module.AlignedShapeAsLatentPLModule + params: + shape_module_cfg: + target: modules.miche.michelangelo.models.tsal.sal_perceiver.AlignedShapeLatentPerceiver + params: + num_latents: 256 + embed_dim: 64 + point_feats: 3 # normal + num_freqs: 8 + include_pi: false + heads: 12 + width: 768 + num_encoder_layers: 8 + num_decoder_layers: 16 + use_ln_post: true + init_scale: 0.25 + qkv_bias: false + use_checkpoint: true + aligned_module_cfg: + target: modules.miche.michelangelo.models.tsal.clip_asl_module.CLIPAlignedShapeAsLatentModule + params: + clip_model_version: "./checkpoints/clip/clip-vit-large-patch14" + + loss_cfg: + target: modules.miche.michelangelo.models.tsal.loss.ContrastKLNearFar + params: + contrast_weight: 0.1 + near_weight: 0.1 + kl_weight: 0.001 + + optimizer_cfg: + optimizer: + target: torch.optim.AdamW + params: + betas: [0.9, 0.99] + eps: 1.e-6 + weight_decay: 1.e-2 + + scheduler: + target: modules.miche.michelangelo.utils.trainings.lr_scheduler.LambdaWarmUpCosineFactorScheduler + params: + warm_up_steps: 5000 + f_start: 1.e-6 + f_min: 1.e-3 + f_max: 1.0 diff --git a/mulres_infer/modules/octree_codec.py b/mulres_infer/modules/octree_codec.py new file mode 100644 index 0000000000000000000000000000000000000000..93d818eb53b4ca052d45c0f4c2f429217dafb9f1 --- /dev/null +++ b/mulres_infer/modules/octree_codec.py @@ -0,0 +1,981 @@ +"""fix3 tokenizer(已融合 numba 加速实现, 原 *_fast.py 并入本文件)。 + +encode_mesh / decode_tokens 走 numpy 向量化 + numba 枚举, 比纯 python 版快约 5x。 +quantized_mesh 与 decode 输出和纯 python 版完全一致; token 的 face 块集合一致, +仅少数平局块顺序不同(本实现补了 (f2.y,f2.x) 决定性平局键, 比原版更可复现)。 + +原文件头: + +fix3: 连接矩阵改为「octant 绝对原位 8x8 + 因子分解(按行)」编码，保留 F。 + +相对 fix2 的核心变化 +-------------------- +1. 不再按 F 排序、不再受 max_children 限制。每个连接矩阵(intra / inter)统一是 + 8x8 的「octant 绝对原位」矩阵：行/列 = octant 0..7 的固定语义。 + → 天然 mc=8 无损，"超标 / 合并 / 退化分摊" 这套概念全部消失。 +2. 8x8 矩阵不再编码成单个 token(那是 2^64 词表)，而是按行拆成 8 个「行子码」， + 每个行子码 = 该行 8 个 bit ∈ [0, 256)。intra / inter 各 8 个子码。 +3. **逻辑序列长度不变**：每个 face 仍是 18 个「位置」(9 coord + 3 F + 3 intra + 3 inter)， + 只是 intra / inter 这 6 个位置每个携带 8 个子码(给因子分解头：同一位置 8 个并行 256-way 头)。 + coord / F 仍是单子码位置。tokenizer.encode 返回 List[List[int]]：每个内层 list = 一个位置的子码。 +4. 保留 F(标识 active octant)。重建时用 F mask 掉非 active 行列(M &= F⊗F)，消除歧义。 + +词表(极小)：pad(1) + coord(resolution) + F(256) + intra_row(256) + inter_row(256)。 +""" +import numpy as np +import trimesh +from tqdm import tqdm +import os +import sys +import glob +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from pathlib import Path +from utils.data_utils import undiscretize, discretize, normalize, clean_mesh +from utils.mesh_metric import compute_mesh_metrics +from modules.fast_octree_level import _triangles_csr +import itertools +from dataclasses import dataclass + +NUM_OCTANTS = 8 + +# 预计算 24 个节点(3 槽位 x 8 octant)中选 3 个的所有组合，用于向量化找三角形 +_TRI_INDICES = np.array(list(itertools.combinations(range(24), 3))) +_IDX_I = _TRI_INDICES[:, 0] +_IDX_J = _TRI_INDICES[:, 1] +_IDX_K = _TRI_INDICES[:, 2] + + +# ========================================== +# Tokenizer：octant 原位 + 按行子码 +# ========================================== +class DynamicMeshTokenizer: + """ + encode 输出格式(fine)：List[List[int]]，每个内层 list 是一个「逻辑位置」的子码。 + 每 face 18 个位置： [v0x][v0y][v0z][v1x][v1y][v1z][v2x][v2y][v2z] (9 个 coord, 各 1 子码) + [F0][F1][F2] (3 个 F, 各 1 子码) + [M00 的 8 行][M11..][M22..] (3 个 intra, 各 8 子码) + [M01 的 8 行][M02..][M12..] (3 个 inter, 各 8 子码) + coord / F 位置长度 1；intra / inter 位置长度 8(行子码)。逻辑序列长度 = 18/face，与 fix2 相同。 + """ + def __init__(self, resolution, max_children=NUM_OCTANTS): + self.resolution = resolution + self.max_children = NUM_OCTANTS # octant 原位下固定为 8；保留字段仅为兼容 config + + self.pad_id = 0 + self.coord_offset = 1 + self.f_offset = self.coord_offset + self.resolution + self.row_vocab = 1 << NUM_OCTANTS # 2^8 = 256 + self.f_vocab = 1 << NUM_OCTANTS # F 也是 8bit + self.intra_row_offset = self.f_offset + self.f_vocab + self.inter_row_offset = self.intra_row_offset + self.row_vocab + self.vocab_size = self.inter_row_offset + self.row_vocab + + # ---- bit <-> int ---- + def _row_to_int(self, bits): + v = 0 + for i, b in enumerate(bits): + if b: + v |= (1 << i) + return v + + def _int_to_row(self, val): + return [(val >> i) & 1 for i in range(NUM_OCTANTS)] + + def _mat_to_rows(self, M, offset): + return [self._row_to_int(M[r]) + offset for r in range(NUM_OCTANTS)] + + def _rows_to_mat(self, rows, offset): + return [self._int_to_row(rows[r] - offset) for r in range(NUM_OCTANTS)] + + # ------------------------------------ + # 细分网格 (Fine Mesh) —— 扁平 1D int 布局，每 face 60 个 token： + # 9 coord + 3 F + 6 矩阵 x 8 行子码 + # 顺序: [coord(9)] [F0 F1 F2] [M00 M11 M22 各 8 行] [M01 M02 M12 各 8 行] + # F: 每个 1 个 token；intra/inter: 每个矩阵按行拆成 8 个行子码(每行 8bit) + # 所有真实 token 都带 offset(>=1)，0 仅作 trailing pad，所以 F 不需要复制避坑。 + # 与现有 dataset/collate(1D int, pad_id=0) 兼容；模型侧整理成 9 个 refine 位置/face + # (F 位置 1 子码、M 位置 8 行子码 sum 成一个 embedding token)，逻辑序列 18/face 不变。 + # ------------------------------------ + TOKENS_PER_FACE = 9 + 3 + 6 * NUM_OCTANTS # 60 + REFINE_POS_PER_FACE = 9 # 模型逻辑: 3 F + 3 intra + 3 inter + + def encode(self, faces_data, use_hourglass=True): + seq = [] + for face in faces_data: + v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12 = face + for v in (v0, v1, v2): + seq.extend([v[0] + self.coord_offset, v[1] + self.coord_offset, v[2] + self.coord_offset]) + for F in (F0, F1, F2): + seq.append(self._row_to_int(F) + self.f_offset) # F: 1 个 token + for M in (M00, M11, M22): + seq.extend(self._mat_to_rows(M, self.intra_row_offset)) # 各 8 行 + for M in (M01, M02, M12): + seq.extend(self._mat_to_rows(M, self.inter_row_offset)) # 各 8 行 + return seq + + def decode(self, seq): + faces_data = [] + B = self.TOKENS_PER_FACE # 60 + clean_seq = [t for t in seq if t != self.pad_id] + co = self.coord_offset + for i in range(0, len(clean_seq) - B + 1, B): + ch = clean_seq[i:i + B] + cv = [ch[k] - co for k in range(9)] + v0 = (cv[0], cv[1], cv[2]); v1 = (cv[3], cv[4], cv[5]); v2 = (cv[6], cv[7], cv[8]) + F0 = self._int_to_row(ch[9] - self.f_offset) + F1 = self._int_to_row(ch[10] - self.f_offset) + F2 = self._int_to_row(ch[11] - self.f_offset) + M00 = self._rows_to_mat(ch[12:20], self.intra_row_offset) + M11 = self._rows_to_mat(ch[20:28], self.intra_row_offset) + M22 = self._rows_to_mat(ch[28:36], self.intra_row_offset) + M01 = self._rows_to_mat(ch[36:44], self.inter_row_offset) + M02 = self._rows_to_mat(ch[44:52], self.inter_row_offset) + M12 = self._rows_to_mat(ch[52:60], self.inter_row_offset) + faces_data.append((v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12)) + return faces_data + + # ------------------------------------ + # 粗糙网格 (Raw Mesh)：仅顶点坐标，每 face 9 个单子码位置 + # ------------------------------------ + def encode_raw(self, raw_faces_data, use_hourglass=True): + seq = [] + for face in raw_faces_data: + v0, v1, v2 = face + for v in (v0, v1, v2): + seq.append([v[0] + self.coord_offset]) + seq.append([v[1] + self.coord_offset]) + seq.append([v[2] + self.coord_offset]) + return seq + + def decode_raw(self, seq): + faces_data = [] + clean_seq = [p for p in seq if not (len(p) == 1 and p[0] == self.pad_id)] + for i in range(0, len(clean_seq) - 9 + 1, 9): + chunk = clean_seq[i:i + 9] + if len(chunk) < 9: + break + cv = [chunk[k][0] - self.coord_offset for k in range(9)] + faces_data.append(((cv[0], cv[1], cv[2]), (cv[3], cv[4], cv[5]), (cv[6], cv[7], cv[8]))) + return faces_data + + +# ========================================== +# Builder：直接构建 octant 原位 8x8 矩阵 +# ========================================== +class OctreeFaceTokenBuilder: + def __init__(self, tree, max_children=NUM_OCTANTS): + self.tree = tree + self.max_children = NUM_OCTANTS # octant 原位固定 8 + + def get_level_info(self, depth): + nodes_at_depth = {} + + def collect_nodes(current, target_depth): + if current.depth == target_depth: + nodes_at_depth[current.coord] = current + return + for child in current.children: + if child: + collect_nodes(child, target_depth) + + collect_nodes(self.tree.root, depth) + + coords = sorted(nodes_at_depth.keys(), key=lambda x: (x[2], x[1], x[0])) + coord_to_idx = {c: i for i, c in enumerate(coords)} + + adj = {i: set() for i in range(len(coords))} + for i, c in enumerate(coords): + node = nodes_at_depth[c] + all_neighbors = node.intra_neighbors | node.inter_neighbors + for neighbor_coord in all_neighbors: + if neighbor_coord in coord_to_idx: + j = coord_to_idx[neighbor_coord] + if i < j: + adj[i].add(j) + + raw_faces_set = set() + covered_edges = set() + + # 1. 3-clique (三角形) + for u in range(len(coords)): + for v in adj[u]: + common_neighbors = adj[u] & adj[v] + for w in common_neighbors: + face_coords = tuple(sorted([coords[u], coords[v], coords[w]])) + raw_faces_set.add(face_coords) + covered_edges.add(tuple(sorted([face_coords[0], face_coords[1]]))) + covered_edges.add(tuple(sorted([face_coords[1], face_coords[2]]))) + covered_edges.add(tuple(sorted([face_coords[0], face_coords[2]]))) + + # 2. 未被三角形覆盖的独立边 -> 退化面片 (A, B, B) + for u in range(len(coords)): + for v in adj[u]: + edge = tuple(sorted([coords[u], coords[v]])) + if edge not in covered_edges: + raw_faces_set.add((edge[0], edge[1], edge[1])) + + # 3. 孤立顶点 -> (A, A, A) + for u in range(len(coords)): + has_edges = len(adj[u]) > 0 or any(u in adj[i] for i in range(u)) + if not has_edges: + raw_faces_set.add((coords[u], coords[u], coords[u])) + + raw_faces_at_depth = list(raw_faces_set) + raw_faces_at_depth.sort(key=lambda f: (f[0][2], f[0][1], f[0][0], + f[1][2], f[1][1], f[1][0], + f[2][2])) + return raw_faces_at_depth, nodes_at_depth + + def _active_octants(self, coord, n, nodes_map): + """该节点哪些 octant 子节点存在(occupancy)。 + 用 children 占据而非"有连接"标记 —— 因为「有连接 ⟹ 必有子节点」，occupancy 是 + "有连接"的超集，且正好等于 raw 的 get_level_info 收集的 depth+1 节点集； + 这样「有几何但无边」的孤立子节点也会被 F 标记，fine 重构能补回孤立点，与 raw 完全一致。 + M(连接矩阵)不变，故三角形/边重构不受影响。""" + return set(i for i in range(NUM_OCTANTS) if n.children[i] is not None) + + def encode_level_tokens(self, depth): + raw_faces, nodes_map = self.get_level_info(depth) + empty = np.zeros((NUM_OCTANTS, NUM_OCTANTS), dtype=np.int8) + + # F：每个节点的 active octant 掩码 (8bit) + node_F = {} + for coord, n in nodes_map.items(): + active = self._active_octants(coord, n, nodes_map) + node_F[coord] = [1 if i in active else 0 for i in range(NUM_OCTANTS)] + + intra_cache = {} + + def intra8(coord): + if coord not in intra_cache: + M = (nodes_map[coord].intra_matrix != 0).astype(np.int8) # 8x8 octant 原位、对称 + intra_cache[coord] = M.tolist() + return intra_cache[coord] + + def inter8(a, b): + if a == b: + return intra8(a) + na = nodes_map[a] + M = na.inter_matrices.get(b) + if M is None: + nb = nodes_map[b] + Mt = nb.inter_matrices.get(a) + M = Mt.T if Mt is not None else empty + return (np.asarray(M) != 0).astype(np.int8).tolist() + + raw_tokens = [] + output_tokens = [] + for face in raw_faces: + v0_c, v1_c, v2_c = face + raw_tokens.append([v0_c, v1_c, v2_c]) + + F0, F1, F2 = node_F[v0_c], node_F[v1_c], node_F[v2_c] + M00, M11, M22 = intra8(v0_c), intra8(v1_c), intra8(v2_c) + M01 = inter8(v0_c, v1_c) + M02 = inter8(v0_c, v2_c) + M12 = inter8(v1_c, v2_c) + + output_tokens.append((v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12)) + return raw_tokens, output_tokens + + +class OctreeNode: + def __init__(self, x, y, z, depth): + self.coord = (x, y, z) + self.depth = depth + self.children = [None] * 8 + self.intra_neighbors = set() + self.inter_neighbors = set() + self.intra_matrix = np.zeros((8, 8), dtype=np.int8) + self.inter_matrices = {} + + +class MatrixOctree: + def __init__(self, max_depth=7): + self.max_depth = max_depth + self.root = None + + def _get_parent_coord(self, coord): + return (coord[0] >> 1, coord[1] >> 1, coord[2] >> 1) + + def _get_local_index(self, coord): + return ((coord[2] & 1) << 2) | ((coord[1] & 1) << 1) | (coord[0] & 1) + + def build(self, vertices, faces): + resolution = 2 ** self.max_depth + quantized = discretize(vertices, num_discrete=resolution) + + leaf_nodes_map = {} + for i in range(len(vertices)): + coord = tuple(quantized[i]) + if coord not in leaf_nodes_map: + leaf_nodes_map[coord] = OctreeNode(coord[0], coord[1], coord[2], self.max_depth) + + for face in faces: + coords = [tuple(quantized[i]) for i in face] + for i in range(3): + u, v = coords[i], coords[(i + 1) % 3] + if u == v: + continue + parent_u = self._get_parent_coord(u) + parent_v = self._get_parent_coord(v) + node_u = leaf_nodes_map[u] + node_v = leaf_nodes_map[v] + if parent_u == parent_v: + node_u.intra_neighbors.add(v) + node_v.intra_neighbors.add(u) + else: + node_u.inter_neighbors.add(v) + node_v.inter_neighbors.add(u) + + current_layer_nodes = list(leaf_nodes_map.values()) + + for d in range(self.max_depth, 0, -1): + parent_layer_map = {} + for child in current_layer_nodes: + parent_coord = self._get_parent_coord(child.coord) + if parent_coord not in parent_layer_map: + parent_layer_map[parent_coord] = OctreeNode(*parent_coord, d - 1) + parent_node = parent_layer_map[parent_coord] + idx = self._get_local_index(child.coord) + parent_node.children[idx] = child + + for child in current_layer_nodes: + parent_coord = self._get_parent_coord(child.coord) + parent_node = parent_layer_map[parent_coord] + child_idx = self._get_local_index(child.coord) + + for sibling_coord in child.intra_neighbors: + sibling_idx = self._get_local_index(sibling_coord) + parent_node.intra_matrix[child_idx, sibling_idx] = 1 + + for neighbor_coord in child.inter_neighbors: + neighbor_parent_coord = self._get_parent_coord(neighbor_coord) + if neighbor_parent_coord not in parent_node.inter_matrices: + parent_node.inter_matrices[neighbor_parent_coord] = np.zeros((8, 8), dtype=np.int8) + neighbor_child_idx = self._get_local_index(neighbor_coord) + parent_node.inter_matrices[neighbor_parent_coord][child_idx, neighbor_child_idx] = 1 + + grandparent_p = self._get_parent_coord(parent_coord) + grandparent_n = self._get_parent_coord(neighbor_parent_coord) + if grandparent_p == grandparent_n: + parent_node.intra_neighbors.add(neighbor_parent_coord) + else: + parent_node.inter_neighbors.add(neighbor_parent_coord) + + current_layer_nodes = list(parent_layer_map.values()) + if current_layer_nodes: + self.root = current_layer_nodes[0] + + +def reconstruct_next_level_optimized(output_tokens, resolution, max_children=NUM_OCTANTS): + """octant 绝对原位重建：M 已是 8x8、行列=octant；用 F mask 掉非 active 行列。""" + indices = np.arange(8) + offsets = np.stack([indices & 1, (indices >> 1) & 1, (indices >> 2) & 1], axis=1).astype(np.int32) + + tri_blocks = [] # 逐 face 向量化收集的三角形 (N,3,3)，最后统一去重 + edge_keys = set() # 全局所有边: canonical (key_a, key_b) + tri_edge_keys = set() # 被三角形覆盖的边(全局) + active_keys = set() # 全局 active 节点 key + key_to_coord = {} # key(int 三元组) -> np.int32 坐标(3,) + + for f_token in output_tokens: + v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12 = f_token + fa = [np.asarray(F0, dtype=bool), np.asarray(F1, dtype=bool), np.asarray(F2, dtype=bool)] + + def intra(F, M): + return np.asarray(M, dtype=bool) & np.outer(F, F) + + def inter(Fa, Fb, M): + return np.asarray(M, dtype=bool) & np.outer(Fa, Fb) + + in0 = intra(fa[0], M00); in1 = intra(fa[1], M11); in2 = intra(fa[2], M22) + m01 = inter(fa[0], fa[1], M01); m02 = inter(fa[0], fa[2], M02); m12 = inter(fa[1], fa[2], M12) + + A = np.zeros((24, 24), dtype=bool) + A[0:8, 0:8] = in0 | in0.T + A[8:16, 8:16] = in1 | in1.T + A[16:24, 16:24] = in2 | in2.T + A[0:8, 8:16] = m01; A[8:16, 0:8] = m01.T + A[0:8, 16:24] = m02; A[16:24, 0:8] = m02.T + A[8:16, 16:24] = m12; A[16:24, 8:16] = m12.T + + # 24 个节点(slot p, octant c)的世界坐标 = v_parents[p]*2 + offset[c] + v_parents_arr = np.stack([np.array(v0_c[:3], dtype=np.int32), + np.array(v1_c[:3], dtype=np.int32), + np.array(v2_c[:3], dtype=np.int32)]) + node_coords = v_parents_arr[np.arange(24) // 8] * 2 + offsets[np.arange(24) % 8] # (24,3) + keys = [tuple(int(x) for x in node_coords[i]) for i in range(24)] + for i in range(24): + if keys[i] not in key_to_coord: + key_to_coord[keys[i]] = node_coords[i] + + # 三角形(逐 face 向量化)；同时把三角形的边记入全局 covered 集合 + valid_mask = A[_IDX_I, _IDX_J] & A[_IDX_J, _IDX_K] & A[_IDX_I, _IDX_K] + tris = _TRI_INDICES[valid_mask] + if len(tris) > 0: + p_idx = tris // 8 + c_idx = tris % 8 + tri_blocks.append(v_parents_arr[p_idx] * 2 + offsets[c_idx]) + for i, j, k in tris: + ka, kb, kc = sorted((keys[int(i)], keys[int(j)], keys[int(k)])) + tri_edge_keys.add((ka, kb)); tri_edge_keys.add((kb, kc)); tri_edge_keys.add((ka, kc)) + + # 全局收集所有边 + ei, ej = np.where(np.triu(A, 1)) + for i, j in zip(ei.tolist(), ej.tolist()): + ki, kj = keys[i], keys[j] + edge_keys.add((ki, kj) if ki <= kj else (kj, ki)) + + # 全局收集 active 节点 + for i in np.where(np.concatenate(fa))[0].tolist(): + active_keys.add(keys[i]) + + all_blocks = list(tri_blocks) + + # 未被任何三角形覆盖的边 -> 退化面 (A, B, B) [全局判断，与 raw 的 get_level_info 一致] + for a, b in (edge_keys - tri_edge_keys): + all_blocks.append(np.stack([key_to_coord[a], key_to_coord[b], key_to_coord[b]])[None]) + + # 孤立的 active 节点(全局无任何边) -> 点 (A, A, A) + edge_nodes = set() + for a, b in edge_keys: + edge_nodes.add(a); edge_nodes.add(b) + for a in (active_keys - edge_nodes): + c = key_to_coord[a] + all_blocks.append(np.stack([c, c, c])[None]) + + if not all_blocks: + return np.array([]), np.array([]) + + flat_v = np.concatenate(all_blocks, axis=0).reshape(-1, 3) + V, inv = np.unique(flat_v, axis=0, return_inverse=True) + F_raw = inv.reshape(-1, 3) + F_sorted = np.sort(F_raw, axis=1) + _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True) + F = F_raw[unique_face_indices] + V = undiscretize(V, num_discrete=resolution) + return V, F + + +def reconstruct_base_mesh(output_tokens, resolution): + all_face_coords = [] + for f_token in output_tokens: + v0_c, v1_c, v2_c = f_token + c0 = np.array(v0_c[:3], dtype=np.int32) + c1 = np.array(v1_c[:3], dtype=np.int32) + c2 = np.array(v2_c[:3], dtype=np.int32) + all_face_coords.append(np.stack([c0, c1, c2])) + + if not all_face_coords: + return np.array([]), np.array([]) + + flat_v = np.concatenate(all_face_coords, axis=0).reshape(-1, 3) + V, inv = np.unique(flat_v, axis=0, return_inverse=True) + F_raw = inv.reshape(-1, 3) + F_sorted = np.sort(F_raw, axis=1) + _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True) + F = F_raw[unique_face_indices] + V = undiscretize(V, num_discrete=resolution) + return V, F + + +class MeshCodec: + def __init__(self, args): + self.max_depth = args.max_depth + self.max_resolution = args.max_resolution + self.use_hourglass = args.use_hourglass + self.max_children = NUM_OCTANTS # octant 原位固定 8(忽略 config 的 max_children) + + self.tokenizer = DynamicMeshTokenizer(resolution=self.max_resolution) + + def encode_mesh(self, mesh: trimesh.Trimesh, target_depth: int, corrupt=None) -> tuple: + if isinstance(mesh, trimesh.Scene): + mesh = mesh.dump(concatenate=True) + V = np.asarray(mesh.vertices, dtype=np.float64) + + res = 2 ** (target_depth + 1) + v_restored = undiscretize(discretize(V, num_discrete=res), num_discrete=res) + quantized_mesh = trimesh.Trimesh(vertices=v_restored, faces=mesh.faces, process=False) + + # corrupt(仅训练时由 dataset 传入)= 2a 暴露偏差腐蚀; None 时与原版逐位一致(推理不受影响) + raw_seq, seq = _fast_fix3_level_tokens(V, mesh.faces, target_depth, self.tokenizer, corrupt=corrupt) + return raw_seq, seq, quantized_mesh + + + def decode_tokens(self, token_sequence: list, target_depth: int): + if os.environ.get("FIX3_FAST_DECODE", "1") != "1": + return super().decode_tokens(token_sequence, target_depth) + if token_sequence is None or len(token_sequence) == 0: + return None + return _fast_decode_tokens(token_sequence, target_depth, self.tokenizer) + + + def decode_raw_tokens(self, raw_token_sequence: list, target_depth: int) -> trimesh.Trimesh: + if not raw_token_sequence: + return None + resolution = 2 ** (target_depth) + raw_t = self.tokenizer.decode_raw(raw_token_sequence) + v, f = reconstruct_base_mesh(raw_t, resolution) + if len(v) == 0 or len(f) == 0: + return None + return trimesh.Trimesh(vertices=v, faces=f, process=False) + + def get_model_config(self) -> dict: + t = self.tokenizer + return { + "total_vocab_size": t.vocab_size, + "coord_vocab_size": t.resolution, + "f_vocab_size": t.f_vocab, + "row_vocab_size": t.row_vocab, # intra/inter 行子码各 256 + "coord_offset": t.coord_offset, + "f_offset": t.f_offset, + "intra_row_offset": t.intra_row_offset, + "inter_row_offset": t.inter_row_offset, + "pad_token_id": t.pad_id, + "num_octants": NUM_OCTANTS, + "subcodes_per_position": NUM_OCTANTS, # 每个 refine 位置 8 个子码槽(F 复制；intra/inter 为 8 行) + "tokens_per_face": DynamicMeshTokenizer.TOKENS_PER_FACE, # 60 扁平 token/face (9+3+6*8) + "coord_per_face": 9, + "refine_positions_per_face": DynamicMeshTokenizer.REFINE_POS_PER_FACE, # 9 (3F+3intra+3inter) + "input_num_categories": t.resolution + 1, + # 输出端：F / intra_row / inter_row 三类，每类 256；coord 为输入端条件 + "output_num_categories": 1 + t.f_vocab + 2 * t.row_vocab, + } + + +# ========================================== +# 自测：encode -> decode -> reconstruct 往返 +# ========================================== +if __name__ == "__main__": + # ========================================== + # 编解码 + 重构导出流程 (参照 fix2 的 __main__) + # 每个 mesh: 导出 GT；对每个 depth 做 + # encode -> decode_tokens(细分/Fine) / decode_raw_tokens(粗糙/Coarse) / 量化网格 + # 并打印计时、token 数(60/face)与 token 级往返一致性。 + # ========================================== + import time + + # 1. 输入/输出文件夹 (按需修改) + INPUT_DIRECTORY = "/root/junkai_lin/mul_res_diffusion/demo/complex_test_data/files" + OUTPUT_DIRECTORY = "./output/fix3_octant_factorized2" + os.makedirs(OUTPUT_DIRECTORY, exist_ok=True) + + @dataclass + class MeshCodecArgs: + max_depth: int = 10 + max_resolution: int = 1024 + use_hourglass: bool = True + max_children: int = 8 # fix3 octant 原位固定 8 (此字段实际被忽略) + + # 2. 初始化 Codec + args = MeshCodecArgs() + codec = MeshCodec(args=args) + print(f"[*] Initialized fix3 Codec. Total Vocab Size: {codec.tokenizer.vocab_size} " + f"(tokens/face={DynamicMeshTokenizer.TOKENS_PER_FACE})") + + # 3. 收集所有 .obj / .glb + input_dir_path = Path(INPUT_DIRECTORY) + mesh_files = sorted(input_dir_path.rglob("*.obj")) + sorted(input_dir_path.rglob("*.glb")) + if not mesh_files: + print(f"[!] No .obj or .glb files found in {INPUT_DIRECTORY}") + + DEPTHS = [0, 1, 2, 3, 4, 5, 6] + + # 4. 逐文件处理 + for idx, file_path in enumerate(mesh_files, 1): + obj_name = file_path.stem + print(f"\n{'=' * 50}") + print(f"[{idx}/{len(mesh_files)}] Processing: {file_path.name}") + print(f"{'=' * 50}") + file_start_time = time.perf_counter() + + try: + raw_mesh = trimesh.load(str(file_path), force="mesh") + normalize(raw_mesh) + except Exception as e: + print(f"[!] load/normalize failed: {e}") + continue + + # 导出 GT + gt_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj") + raw_mesh.export(gt_path) + print(f"[*] Saved GT -> {gt_path}") + + # 不同深度的编解码 + for depth in DEPTHS: + print(f"\n--- Depth {depth} | res(fine)={2 ** (depth + 1)} | octant 原位 ---") + + # Encoding + t0 = time.perf_counter() + raw_tokens, tokens, quantized_mesh = codec.encode_mesh(raw_mesh, target_depth=depth) + encode_time = time.perf_counter() - t0 + nface = len(tokens) // DynamicMeshTokenizer.TOKENS_PER_FACE + print(f"[Timer] Encoding: {encode_time:.4f}s | flat_tokens={len(tokens)} " + f"(nface={nface}, 模型逻辑位置={18 * nface})") + + # token 级往返一致性 (encode->decode->re-encode) + same = (tokens == codec.tokenizer.encode(codec.tokenizer.decode(tokens))) + print(f"[Check] token 往返一致: {same}") + + # Decoding (Fine) + if len(tokens) > 0: + t1 = time.perf_counter() + recon_mesh = codec.decode_tokens(tokens, target_depth=depth) + decode_fine_time = time.perf_counter() - t1 + if recon_mesh is not None: + out_path = os.path.join( + OUTPUT_DIRECTORY, f"{obj_name}_recon_d{depth}_res{2 ** (depth + 1)}_fix3.obj") + recon_mesh.export(out_path) + print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s " + f"-> V={len(recon_mesh.vertices)} F={len(recon_mesh.faces)}") + else: + print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s -> empty") + + # Decoding (Coarse / Raw) + if len(raw_tokens) > 0: + t2 = time.perf_counter() + raw_recon_mesh = codec.decode_raw_tokens(raw_tokens, target_depth=depth) + decode_coarse_time = time.perf_counter() - t2 + if raw_recon_mesh is not None: + raw_out_path = os.path.join( + OUTPUT_DIRECTORY, f"{obj_name}_raw_recon_d{depth}_res{2 ** depth}.obj") + raw_recon_mesh.export(raw_out_path) + print(f"[Timer] Decoding (Coarse): {decode_coarse_time:.4f}s " + f"-> V={len(raw_recon_mesh.vertices)} F={len(raw_recon_mesh.faces)}") + + # 量化网格 (输入离散化后的 GT，用于对照重构质量上限) + if quantized_mesh is not None: + quant_out_path = os.path.join( + OUTPUT_DIRECTORY, f"{obj_name}_dis_d{depth + 1}_res{2 ** (depth + 1)}.obj") + quantized_mesh.export(quant_out_path) + + file_total_time = time.perf_counter() - file_start_time + print(f"\n[*] Finished {file_path.name} in {file_total_time:.2f}s") + + +# ========================================================================== +# encode_mesh / decode_tokens 的 numba 加速实现细节(原 *_fast.py 并入) +# ========================================================================== + +_OCT_W = np.array([1, 2, 4], dtype=np.int64) # octant 局部索引 = x + 2y + 4z + +def _corrupt_fid_nodes(fid, nodes, und, nN, target_depth, corrupt): + """2a 暴露偏差训练腐蚀(仅训练): 模拟级联实测的两类粗输入误差。 + (a) 伪面: 在已存在的 GT 节点间加不该有的面(过连接)。下游打包按 fid 查表 -> + 非边对的 inter 行自动全 0(empty),教模型"看 PC 把伪面 refine 成空"。 + (b) 节点坐标抖动 ±jitter_max(位移): 只改 coord_part 的输入坐标,目标 F/intra/inter + 仍按真节点身份查表不变,教模型"输入坐标偏了也 refine 对"。 + 目标无需重新导出 —— 全部沿用 fid 索引,与干净路径同一套向量化打包。返回 (fid_aug, nodes_coord)。 + """ + sp_rate = float(corrupt.get("spurious_rate", 0.0)) + jp = float(corrupt.get("jitter_prob", 0.0)) + + if sp_rate > 0.0 and nN >= 3 and und.shape[0] > 0: + radius = int(corrupt.get("spatial_radius", 2)) + n_add = int(sp_rate * fid.shape[0]) + if n_add > 0: + coord2idx = {(int(c[0]), int(c[1]), int(c[2])): i for i, c in enumerate(nodes)} + existing = set(map(tuple, np.sort(fid, axis=1).tolist())) + ne = und.shape[0] + spurious, tries, maxtries = [], 0, n_add * 6 + while len(spurious) < n_add and tries < maxtries: + tries += 1 + ei = int(np.random.randint(ne)) + ai, bi = int(und[ei, 0]), int(und[ei, 1]) + base = nodes[ai] if np.random.random() < 0.5 else nodes[bi] + off = np.random.randint(-radius, radius + 1, size=3) + if not off.any(): + continue + ck = (int(base[0] + off[0]), int(base[1] + off[1]), int(base[2] + off[2])) + ci = coord2idx.get(ck) + if ci is None or ci == ai or ci == bi: + continue + tri = tuple(sorted((ai, bi, ci))) + if tri in existing: + continue + existing.add(tri) + spurious.append(tri) + if spurious: + fid = np.concatenate([fid, np.asarray(spurious, dtype=fid.dtype)], axis=0) + + nodes_coord = nodes + if jp > 0.0: + jmax = int(corrupt.get("jitter_max", 1)) + res_c = 2 ** target_depth + nodes_coord = nodes.copy() + mask = np.random.random(nN) < jp + k = int(mask.sum()) + if k > 0: + noise = np.random.randint(-jmax, jmax + 1, size=(k, 3)) + nodes_coord[mask] = np.clip(nodes_coord[mask] + noise, 0, res_c - 1) + return fid, nodes_coord + +def _fast_fix3_level_tokens(vertices, faces, target_depth, tok, corrupt=None): + """等价替换 MatrixOctree(max_depth=td+1).build + encode_level_tokens(td) + tokenizer.encode/encode_raw。 + 返回 (raw_token_sequence, token_sequence)，与原版逐位一致。""" + res = 2 ** (target_depth + 1) + Q = discretize(np.asarray(vertices, dtype=np.float64), num_discrete=res).astype(np.int64) + faces = np.asarray(faces, dtype=np.int64).reshape(-1, 3) + + # ---- 节点(td) = 所有顶点叶子坐标的 parent；占据 F = 所有顶点的 octant ---- + P_all = Q >> 1 + i_all = (Q & 1) @ _OCT_W + + # ---- 叶子边表(每面 3 条，跳过零长边 u==v，与原版一致) ---- + if faces.size: + e = np.concatenate([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=0) + u, v = Q[e[:, 0]], Q[e[:, 1]] + keep = np.any(u != v, axis=1) + u, v = u[keep], v[keep] + else: + u = v = np.zeros((0, 3), dtype=np.int64) + pu, pv = u >> 1, v >> 1 + iu = (u & 1) @ _OCT_W + iv = (v & 1) @ _OCT_W + + # 节点表(np.unique 行字典序 = 坐标 (x,y,z) 升序 = 原版 face 内顶点 sorted() 同序) + nodes, inv = np.unique(np.concatenate([P_all, pu, pv], axis=0), axis=0, return_inverse=True) + nV, M = P_all.shape[0], pu.shape[0] + occ_idx = inv[:nV] + a, b = inv[nV:nV + M], inv[nV + M:] + nN = nodes.shape[0] + intra = a == b + + # ---- F 行子码(8bit int)：所有顶点的 octant 占据 ---- + Fbits = np.zeros(nN, dtype=np.int64) + np.bitwise_or.at(Fbits, occ_idx, np.int64(1) << i_all) + + # ---- intra 行子码 [nN,8]：绝对原位、对称双向、无对角填充(与原版一致) ---- + intra_rows = np.zeros((nN, 8), dtype=np.int64) + ai, iui, ivi = a[intra], iu[intra], iv[intra] + np.bitwise_or.at(intra_rows, (ai, iui), np.int64(1) << ivi) + np.bitwise_or.at(intra_rows, (ai, ivi), np.int64(1) << iui) + + # ---- inter 行子码：有向 (a,b) 分组(双向都存，与原版 build 一致) ---- + im = ~intra + ea, eb, eiu, eiv = a[im], b[im], iu[im], iv[im] + dkey = np.concatenate([ea * nN + eb, eb * nN + ea]) + drow = np.concatenate([eiu, eiv]) + dcol = np.concatenate([eiv, eiu]) + ukeys, kinv = np.unique(dkey, return_inverse=True) + inter_rows = np.zeros((max(len(ukeys), 1), 8), dtype=np.int64) + if dkey.size: + np.bitwise_or.at(inter_rows, (kinv, drow), np.int64(1) << dcol) + + # ---- 无向节点邻接 -> CSR -> numba 三角形枚举 ---- + if ea.size: + und = np.unique(np.stack([np.minimum(ea, eb), np.maximum(ea, eb)], axis=1), axis=0) + else: + und = np.zeros((0, 2), dtype=np.int64) + deg = np.zeros(nN, dtype=np.int64) + if und.size: + np.add.at(deg, und[:, 0], 1) + np.add.at(deg, und[:, 1], 1) + order = np.lexsort((und[:, 1], und[:, 0])) + fu, fv = und[order, 0], und[order, 1] + indptr = np.zeros(nN + 1, dtype=np.int64) + np.add.at(indptr, fu + 1, 1) + indptr = np.cumsum(indptr) + tu, tv, tw = _triangles_csr(indptr, fv.astype(np.int64)) + else: + tu = tv = tw = np.zeros(0, dtype=np.int64) + + # ---- faces_ids：三角形 + 未覆盖边退化 (A,B,B) + 孤立点 (A,A,A)，集合与原版一致 ---- + tris = np.stack([tu, tv, tw], axis=1) if tu.size else np.zeros((0, 3), dtype=np.int64) + covered = set() + for x, y, z in tris: + x, y, z = int(x), int(y), int(z) + covered.add((x, y)); covered.add((y, z)); covered.add((x, z)) + extra = [] + for x, y in und: + x, y = int(x), int(y) + if (x, y) not in covered: + extra.append((x, y, y)) + for x in np.flatnonzero(deg == 0): + x = int(x) + extra.append((x, x, x)) + fid = np.concatenate([tris, np.asarray(extra, dtype=np.int64).reshape(-1, 3)], axis=0) + if fid.shape[0] == 0: + return [], [] + + # ---- 原版排序键 + 决定性平局键 ---- + c0, c1, c2 = nodes[fid[:, 0]], nodes[fid[:, 1]], nodes[fid[:, 2]] + order = np.lexsort((c2[:, 0], c2[:, 1], c2[:, 2], + c1[:, 0], c1[:, 1], c1[:, 2], + c0[:, 0], c0[:, 1], c0[:, 2])) # 主键=c0.z(最后) ... 与原版 7 键一致+补齐平局 + fid = fid[order] + nodes_coord = nodes + if corrupt is not None: + fid, nodes_coord = _corrupt_fid_nodes(fid, nodes, und, nN, target_depth, corrupt) + c0, c1, c2 = nodes_coord[fid[:, 0]], nodes_coord[fid[:, 1]], nodes_coord[fid[:, 2]] + + # ---- 60/face 扁平 token 向量化打包 ---- + # 注: F_part/intra_part/inter_part 下面仍用 fid(节点身份)查表 -> 腐蚀的伪面自动得正确目标 + # (非边对 inter 全 0),坐标抖动只影响上面 coord_part,目标不变。 + co, fo = tok.coord_offset, tok.f_offset + io, eo = tok.intra_row_offset, tok.inter_row_offset + nF = fid.shape[0] + + coord_part = np.concatenate([c0, c1, c2], axis=1) + co # (nF,9) + F_part = np.stack([Fbits[fid[:, 0]], Fbits[fid[:, 1]], Fbits[fid[:, 2]]], axis=1) + fo + intra_part = np.concatenate([intra_rows[fid[:, 0]], + intra_rows[fid[:, 1]], + intra_rows[fid[:, 2]]], axis=1) + io # (nF,24) + + inter_blocks = [] + for x, y in ((0, 1), (0, 2), (1, 2)): + A_, B_ = fid[:, x], fid[:, y] + rows = np.zeros((nF, 8), dtype=np.int64) + same = A_ == B_ + rows[same] = intra_rows[A_[same]] # inter8(a,a)=intra8(a)，与原版一致 + ns = np.flatnonzero(~same) + if ns.size: + dk = A_[ns] * nN + B_[ns] + pos = np.searchsorted(ukeys, dk) + pos_c = np.minimum(pos, max(len(ukeys) - 1, 0)) + found = (ukeys[pos_c] == dk) if len(ukeys) else np.zeros(ns.size, dtype=bool) + rows[ns[found]] = inter_rows[pos_c[found]] # 未找到 -> 保持全 0(原版 empty 矩阵) + inter_blocks.append(rows) + inter_part = np.concatenate(inter_blocks, axis=1) + eo # (nF,24) + + token = np.concatenate([coord_part, F_part, intra_part, inter_part], axis=1) # (nF,60) + token_sequence = token.ravel().tolist() + raw_token_sequence = coord_part.reshape(-1, 1).tolist() # 9 coord/face, 各 [t] + return raw_token_sequence, token_sequence + +_BIT_SHIFTS = np.arange(NUM_OCTANTS, dtype=np.int64) # _int_to_row: bit i = (v>>i)&1 + +_NODE_PARENT = np.arange(24) // NUM_OCTANTS # 24 个节点 -> 哪个父顶点 slot + +_NODE_OCTANT = np.arange(24) % NUM_OCTANTS # 24 个节点 -> 哪个 octant + +_OCT_OFFSETS = np.stack([np.arange(8) & 1, (np.arange(8) >> 1) & 1, (np.arange(8) >> 2) & 1], + axis=1).astype(np.int32) # octant -> (dx,dy,dz) + +_INTER_ROW_SLOT = np.array([0, 0, 1]) + +_INTER_COL_SLOT = np.array([1, 2, 2]) + +def _unpack_rows(values): + """int 行子码 -> bool 位。values (...,) -> (..., 8), 与 _int_to_row 同序(bit i = (v>>i)&1)。""" + return ((values[..., None].astype(np.int64) >> _BIT_SHIFTS) & 1).astype(bool) + +def _pack_keys(coords, resolution): + """(...,3) int 坐标 -> (...) int64 唯一 key。坐标范围 [0, resolution)。""" + c = coords.astype(np.int64) + return (c[..., 0] * resolution + c[..., 1]) * resolution + c[..., 2] + +def _fast_reconstruct_next_level(token_array, resolution, tokenizer): + """token_array: (nface, TOKENS_PER_FACE) int。返回 (V, F), 语义同 reconstruct_next_level_optimized。""" + from modules.octree_codec import _TRI_INDICES + + nface = token_array.shape[0] + parents = (token_array[:, :9] - tokenizer.coord_offset).reshape(nface, 3, 3).astype(np.int32) + occupancy = _unpack_rows(token_array[:, 9:12] - tokenizer.f_offset) # (nface,3,8) bool + + def rows_to_mat(lo, hi, offset): + return _unpack_rows(token_array[:, lo:hi] - offset) # (nface,8,8) bool, [行,列] + + intra = [rows_to_mat(12 + 8 * i, 20 + 8 * i, tokenizer.intra_row_offset) for i in range(3)] + inter = [rows_to_mat(36 + 8 * i, 44 + 8 * i, tokenizer.inter_row_offset) for i in range(3)] + + # ---- 24x24 邻接: intra 块按 F&F 掩码后对称化, inter 块按 Fa&Fb 掩码 ---- + adjacency = np.zeros((nface, 24, 24), dtype=bool) + for slot in range(3): + f = occupancy[:, slot] # (nface,8) + block = intra[slot] & (f[:, :, None] & f[:, None, :]) + block = block | np.swapaxes(block, 1, 2) + adjacency[:, slot * 8:(slot + 1) * 8, slot * 8:(slot + 1) * 8] = block + for b in range(3): + ra, ca = _INTER_ROW_SLOT[b], _INTER_COL_SLOT[b] + block = inter[b] & (occupancy[:, ra][:, :, None] & occupancy[:, ca][:, None, :]) + adjacency[:, ra * 8:(ra + 1) * 8, ca * 8:(ca + 1) * 8] = block + adjacency[:, ca * 8:(ca + 1) * 8, ra * 8:(ra + 1) * 8] = np.swapaxes(block, 1, 2) + + # ---- 24 个节点的世界坐标与 key ---- + node_coords = parents[:, _NODE_PARENT] * 2 + _OCT_OFFSETS[_NODE_OCTANT] # (nface,24,3) + node_keys = _pack_keys(node_coords, resolution) # (nface,24) + + # ---- 三角形: 三条边都在邻接里 ---- + i_idx, j_idx, k_idx = _TRI_INDICES[:, 0], _TRI_INDICES[:, 1], _TRI_INDICES[:, 2] + valid = (adjacency[:, i_idx, j_idx] & adjacency[:, j_idx, k_idx] & adjacency[:, i_idx, k_idx]) + face_of_tri, tri_of_tri = np.nonzero(valid) # 各 (ntri,) + tri_nodes = _TRI_INDICES[tri_of_tri] # (ntri,3) 节点编号 + tri_coords = node_coords[face_of_tri[:, None], tri_nodes] # (ntri,3,3) + tri_keys = node_keys[face_of_tri[:, None], tri_nodes] # (ntri,3) + + # ---- 全部边(上三角) ---- + upper = np.triu(np.ones((24, 24), dtype=bool), 1) + face_of_edge, node_i, node_j = np.nonzero(adjacency & upper) + key_i = node_keys[face_of_edge, node_i] + key_j = node_keys[face_of_edge, node_j] + edge_pairs = np.unique(np.stack([np.minimum(key_i, key_j), np.maximum(key_i, key_j)], axis=1), axis=0) + + # ---- 被三角形覆盖的边 ---- + if len(tri_keys): + tri_sorted = np.sort(tri_keys, axis=1) + covered = np.concatenate([tri_sorted[:, [0, 1]], tri_sorted[:, [1, 2]], tri_sorted[:, [0, 2]]], axis=0) + covered = np.unique(covered, axis=0) + else: + covered = np.zeros((0, 2), dtype=np.int64) + + blocks = [tri_coords] if len(tri_coords) else [] + + # ---- 未被任何三角形覆盖的边 -> 退化面 (A,B,B) ---- + if len(edge_pairs): + edge_flat = edge_pairs[:, 0] * (resolution ** 3) + edge_pairs[:, 1] + covered_flat = (covered[:, 0] * (resolution ** 3) + covered[:, 1]) if len(covered) \ + else np.zeros(0, dtype=np.int64) + open_edges = edge_pairs[~np.isin(edge_flat, covered_flat)] + if len(open_edges): + key_to_coord = _build_key_lookup(node_keys, node_coords) + ca = key_to_coord[open_edges[:, 0]] + cb = key_to_coord[open_edges[:, 1]] + blocks.append(np.stack([ca, cb, cb], axis=1)) + + # ---- 孤立的 active 节点(无任何边) -> 点 (A,A,A) ---- + active_keys = np.unique(node_keys[occupancy.reshape(nface, 24)]) + if len(active_keys): + edge_nodes = np.unique(edge_pairs.reshape(-1)) if len(edge_pairs) else np.zeros(0, dtype=np.int64) + isolated = active_keys[~np.isin(active_keys, edge_nodes)] + if len(isolated): + key_to_coord = _build_key_lookup(node_keys, node_coords) + c = key_to_coord[isolated] + blocks.append(np.stack([c, c, c], axis=1)) + + if not blocks: + return np.array([]), np.array([]) + + flat_v = np.concatenate(blocks, axis=0).reshape(-1, 3) + V, inverse = np.unique(flat_v, axis=0, return_inverse=True) + faces_raw = inverse.reshape(-1, 3) + _, unique_face_indices = np.unique(np.sort(faces_raw, axis=1), axis=0, return_index=True) + F = faces_raw[unique_face_indices] + return undiscretize(V, num_discrete=resolution), F + +def _build_key_lookup(node_keys, node_coords): + """key -> 坐标 的稠密查表(用 np.unique 的第一次出现), 供退化面/孤立点还原坐标。""" + flat_keys = node_keys.reshape(-1) + flat_coords = node_coords.reshape(-1, 3) + uniq, first = np.unique(flat_keys, return_index=True) + lookup = np.zeros((int(uniq.max()) + 1, 3), dtype=np.int32) + lookup[uniq] = flat_coords[first] + return lookup + +def _fast_decode_tokens(token_sequence, target_depth, tokenizer): + """decode_tokens 的向量化实现。返回 trimesh.Trimesh 或 None。""" + arr = np.asarray(token_sequence).reshape(-1) + arr = arr[arr != tokenizer.pad_id] # 与原版 clean_seq 一致 + per_face = tokenizer.TOKENS_PER_FACE + nface = arr.shape[0] // per_face + if nface == 0: + return None + token_array = arr[:nface * per_face].reshape(nface, per_face).astype(np.int64) + V, F = _fast_reconstruct_next_level(token_array, 2 ** (target_depth + 1), tokenizer) + if len(V) == 0 or len(F) == 0: + return None + return trimesh.Trimesh(vertices=V, faces=F, process=False) + + +FastMeshCodec = MeshCodec # 兼容别名: 加速实现已并入 MeshCodec diff --git a/mulres_infer/modules/octree_codec_acc.py b/mulres_infer/modules/octree_codec_acc.py new file mode 100644 index 0000000000000000000000000000000000000000..4fca5c711add04c2561935ee5993416e192be833 --- /dev/null +++ b/mulres_infer/modules/octree_codec_acc.py @@ -0,0 +1,857 @@ +import numpy as np +import trimesh +from tqdm import tqdm +import os +import sys +import glob +import multiprocessing +import random +from collections import Counter +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from pathlib import Path +from utils.data_utils import undiscretize, discretize, normalize, clean_mesh +from utils.mesh_metric import compute_mesh_metrics +import itertools +from dataclasses import dataclass +# 预计算 24 个节点中选 3 个的所有可能组合 (共 2024 种) +_TRI_INDICES = np.array(list(itertools.combinations(range(24), 3))) +_IDX_I = _TRI_INDICES[:, 0] +_IDX_J = _TRI_INDICES[:, 1] +_IDX_K = _TRI_INDICES[:, 2] +# ========================================== +# 新增：动态网格 Tokenizer +# ========================================== +class DynamicMeshTokenizer: + def __init__(self, resolution, max_children): + self.resolution = resolution + self.max_children = max_children + self.m_bits = max_children * max_children + + # 1. 特殊控制 Token + self.pad_id = 0 + + # 2. 计算各特征在词表中的偏移量 (Offsets) + self.coord_offset = 1 + self.f_offset = self.coord_offset + self.resolution + self.m_intra_offset = self.f_offset + 256 # 1x8 对应 256 种可能 + + self.m_vocab_size = 2 ** self.m_bits + self.m_inter_offset = self.m_intra_offset + self.m_vocab_size + + # 最终 Transformer 的总词表大小 + self.vocab_size = self.m_inter_offset + self.m_vocab_size + + def _list_to_int(self, bit_list): + res = 0 + for i, bit in enumerate(bit_list): + if bit: res |= (1 << i) + return res + + def _int_to_list(self, val, length): + return [(val >> i) & 1 for i in range(length)] + + def _m_to_int(self, M): + flat = [bit for row in M for bit in row] + return self._list_to_int(flat) + + def _int_to_m(self, val): + flat = self._int_to_list(val, self.m_bits) + return [flat[i*self.max_children : (i+1)*self.max_children] for i in range(self.max_children)] + + # ------------------------------------ + # 细分网格 (Fine Mesh) Tokenizer + # ------------------------------------ + def encode(self, faces_data, use_hourglass=True): + """将细分几何特征列表编码为 1D Token 序列 (18 Tokens / Face)""" + seq = [] + + for face in faces_data: + v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12 = face + + seq.extend([v0[0] + self.coord_offset, v0[1] + self.coord_offset, v0[2] + self.coord_offset]) + seq.extend([v1[0] + self.coord_offset, v1[1] + self.coord_offset, v1[2] + self.coord_offset]) + seq.extend([v2[0] + self.coord_offset, v2[1] + self.coord_offset, v2[2] + self.coord_offset]) + + seq.extend([self._list_to_int(F0) + self.f_offset, + self._list_to_int(F1) + self.f_offset, + self._list_to_int(F2) + self.f_offset]) + + seq.extend([self._m_to_int(M00) + self.m_intra_offset, + self._m_to_int(M11) + self.m_intra_offset, + self._m_to_int(M22) + self.m_intra_offset]) + + seq.extend([self._m_to_int(M01) + self.m_inter_offset, + self._m_to_int(M02) + self.m_inter_offset, + self._m_to_int(M12) + self.m_inter_offset]) + return seq + + def decode(self, seq): + """将细分网格的 1D Token 序列解码回几何特征列表""" + faces_data = [] + clean_seq = [t for t in seq if t not in ([self.pad_id])] + + for i in range(0, len(clean_seq), 18): + chunk = clean_seq[i:i+18] + if len(chunk) < 18: + break + + v0 = (chunk[0] - self.coord_offset, chunk[1] - self.coord_offset, chunk[2] - self.coord_offset) + v1 = (chunk[3] - self.coord_offset, chunk[4] - self.coord_offset, chunk[5] - self.coord_offset) + v2 = (chunk[6] - self.coord_offset, chunk[7] - self.coord_offset, chunk[8] - self.coord_offset) + + F0 = self._int_to_list(chunk[9] - self.f_offset, 8) + F1 = self._int_to_list(chunk[10] - self.f_offset, 8) + F2 = self._int_to_list(chunk[11] - self.f_offset, 8) + + M00 = self._int_to_m(chunk[12] - self.m_intra_offset) + M11 = self._int_to_m(chunk[13] - self.m_intra_offset) + M22 = self._int_to_m(chunk[14] - self.m_intra_offset) + + M01 = self._int_to_m(chunk[15] - self.m_inter_offset) + M02 = self._int_to_m(chunk[16] - self.m_inter_offset) + M12 = self._int_to_m(chunk[17] - self.m_inter_offset) + + faces_data.append((v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12)) + + return faces_data + + # ------------------------------------ + # [新增] 基础/粗糙网格 (Raw Mesh) Tokenizer + # ------------------------------------ + def encode_raw(self, raw_faces_data, use_hourglass=True): + """将粗糙层级特征（仅含顶点）编码为 1D Token 序列 (9 Tokens / Face)""" + seq = [] + + for face in raw_faces_data: + v0, v1, v2 = face + seq.extend([v0[0] + self.coord_offset, v0[1] + self.coord_offset, v0[2] + self.coord_offset]) + seq.extend([v1[0] + self.coord_offset, v1[1] + self.coord_offset, v1[2] + self.coord_offset]) + seq.extend([v2[0] + self.coord_offset, v2[1] + self.coord_offset, v2[2] + self.coord_offset]) + + return seq + + def decode_raw(self, seq): + """将粗糙层级的 1D Token 序列解码回几何特征列表 [(v0, v1, v2), ...]""" + faces_data = [] + clean_seq = [t for t in seq if t not in ([self.pad_id])] + + # 每个面片固定转化为 9 个 Token + for i in range(0, len(clean_seq), 9): + chunk = clean_seq[i:i+9] + if len(chunk) < 9: + break + + v0 = (chunk[0] - self.coord_offset, chunk[1] - self.coord_offset, chunk[2] - self.coord_offset) + v1 = (chunk[3] - self.coord_offset, chunk[4] - self.coord_offset, chunk[5] - self.coord_offset) + v2 = (chunk[6] - self.coord_offset, chunk[7] - self.coord_offset, chunk[8] - self.coord_offset) + + faces_data.append((v0, v1, v2)) + + return faces_data + + +# --- 以下为原有的类声明 (OctreeFaceTokenBuilder, OctreeNode, MatrixOctree 等保持不变) --- +class OctreeFaceTokenBuilder: + def __init__(self, tree, max_children=4): + self.tree = tree + self.max_children = max_children + + def get_level_info(self, depth): + nodes_at_depth = {} + def collect_nodes(current, target_depth): + if current.depth == target_depth: + nodes_at_depth[current.coord] = current + return + for child in current.children: + if child: collect_nodes(child, target_depth) + + collect_nodes(self.tree.root, depth) + + coords = sorted(nodes_at_depth.keys(), key=lambda x: (x[2], x[1], x[0])) + coord_to_idx = {c: i for i, c in enumerate(coords)} + + adj = {i: set() for i in range(len(coords))} + for i, c in enumerate(coords): + node = nodes_at_depth[c] + all_neighbors = node.intra_neighbors | node.inter_neighbors + for neighbor_coord in all_neighbors: + if neighbor_coord in coord_to_idx: + j = coord_to_idx[neighbor_coord] + if i < j: + adj[i].add(j) + + raw_faces_set = set() + covered_edges = set() + + # 1. 提取所有互相连接的 3-clique (正常的三角形面片) + for u in range(len(coords)): + for v in adj[u]: + common_neighbors = adj[u] & adj[v] + for w in common_neighbors: + # 排序坐标点以确保集合去重时的唯一性 + face_coords = tuple(sorted([coords[u], coords[v], coords[w]])) + raw_faces_set.add(face_coords) + + # 记录已经被三角形覆盖的边 + covered_edges.add(tuple(sorted([face_coords[0], face_coords[1]]))) + covered_edges.add(tuple(sorted([face_coords[1], face_coords[2]]))) + covered_edges.add(tuple(sorted([face_coords[0], face_coords[2]]))) + + # 2. 提取未形成三角形的独立边 (线段结构，例如细长桌腿) + for u in range(len(coords)): + for v in adj[u]: + edge = tuple(sorted([coords[u], coords[v]])) + if edge not in covered_edges: + # 将 1D 的线段转换为退化面片 (A, B, B) 以适配 Tokenizer 格式 + degenerate_face = (edge[0], edge[1], edge[1]) + raw_faces_set.add(degenerate_face) + + # 3. (可选) 极端情况下处理没有任何连接的孤立顶点 (A, A, A) + for u in range(len(coords)): + has_edges = len(adj[u]) > 0 or any(u in adj[i] for i in range(u)) + if not has_edges: + isolated_face = (coords[u], coords[u], coords[u]) + raw_faces_set.add(isolated_face) + + raw_faces_at_depth = list(raw_faces_set) + + # 保持原始排序逻辑以确保输出的 Determinism + raw_faces_at_depth.sort(key=lambda f: (f[0][2], f[0][1], f[0][0], + f[1][2], f[1][1], f[1][0], + f[2][2])) + + return raw_faces_at_depth, nodes_at_depth + + def reduce_node_children_to_max(self, nodes_map): + def bit_dist(a, b): + return bin(a ^ b).count('1') + + # Pass 1: 基于原始图结构为所有节点一次性计算 merge plan，消除顺序依赖。 + merge_plan = {} + for coord, n in nodes_map.items(): + active_indices = set() + has_intra_row = n.intra_matrix.any(axis=1) + has_intra_col = n.intra_matrix.any(axis=0) + for i in range(8): + if has_intra_row[i] or has_intra_col[i]: + active_indices.add(i) + + for mat in n.inter_matrices.values(): + has_inter_row = mat.any(axis=1) + for i in range(8): + if has_inter_row[i]: + active_indices.add(i) + + for n_inter_coord in n.inter_neighbors: + if n_inter_coord in nodes_map: + neighbor = nodes_map[n_inter_coord] + if coord in neighbor.inter_matrices: + has_in_col = neighbor.inter_matrices[coord].any(axis=0) + for i in range(8): + if has_in_col[i]: + active_indices.add(i) + + if len(active_indices) <= self.max_children: + merge_plan[coord] = None + continue + + child_degrees = [0] * 8 + intra_row_counts = np.count_nonzero(n.intra_matrix, axis=1) + intra_col_counts = np.count_nonzero(n.intra_matrix, axis=0) + for i in range(8): + child_degrees[i] += intra_row_counts[i] + intra_col_counts[i] + + out_inter_counts = np.zeros(8, dtype=int) + for mat in n.inter_matrices.values(): + out_inter_counts += np.count_nonzero(mat, axis=1) + + in_inter_counts = np.zeros(8, dtype=int) + for n_inter_coord in n.inter_neighbors: + if n_inter_coord in nodes_map: + neighbor = nodes_map[n_inter_coord] + if coord in neighbor.inter_matrices: + in_inter_counts += np.count_nonzero(neighbor.inter_matrices[coord], axis=0) + + for i in range(8): + child_degrees[i] += 2 * (out_inter_counts[i] + in_inter_counts[i]) + + active_list = sorted(list(active_indices)) + scored = sorted( + active_list, + key=lambda idx: (-child_degrees[idx], idx) + ) + keep_nodes = sorted(scored[:self.max_children]) + + mapping = {idx: idx for idx in range(8)} + for m_idx in active_list: + if m_idx in keep_nodes: + continue + target_merge = min(keep_nodes, key=lambda x: (bit_dist(x, m_idx), x)) + mapping[m_idx] = target_merge + + merge_plan[coord] = mapping + + # Pass 2: 根据全局 merge plan 同步重映射所有 intra/inter 矩阵。 + for coord, n in nodes_map.items(): + mapping = merge_plan.get(coord) + if mapping is not None: + old_intra = n.intra_matrix + new_intra = np.zeros_like(old_intra) + nz = np.argwhere(old_intra != 0) + for i, j in nz: + new_intra[mapping[int(i)], mapping[int(j)]] = 1 + np.fill_diagonal(new_intra, 0) + n.intra_matrix = new_intra + + for neigh_coord, mat in list(n.inter_matrices.items()): + src_map = merge_plan.get(coord) + dst_map = merge_plan.get(neigh_coord) + if src_map is None and dst_map is None: + continue + + new_mat = np.zeros_like(mat) + nz = np.argwhere(mat != 0) + for i, j in nz: + ni = src_map[int(i)] if src_map is not None else int(i) + nj = dst_map[int(j)] if dst_map is not None else int(j) + new_mat[ni, nj] = 1 + n.inter_matrices[neigh_coord] = new_mat + + def encode_level_tokens(self, depth): + raw_faces, nodes_map = self.get_level_info(depth) + + self.reduce_node_children_to_max(nodes_map) + + node_F = {} + node_active = {} + empty_mat = np.zeros((8, 8), dtype=np.int8) + + for coord, n in nodes_map.items(): + active_indices = set() + has_intra_row = n.intra_matrix.any(axis=1) + has_intra_col = n.intra_matrix.any(axis=0) + for i in range(8): + if has_intra_row[i] or has_intra_col[i]: + active_indices.add(i) + + for mat in n.inter_matrices.values(): + has_inter_row = mat.any(axis=1) + for i in range(8): + if has_inter_row[i]: + active_indices.add(i) + + for n_inter_coord in n.inter_neighbors: + if n_inter_coord in nodes_map: + neighbor = nodes_map[n_inter_coord] + if n.coord in neighbor.inter_matrices: + has_in_col = neighbor.inter_matrices[n.coord].any(axis=0) + for i in range(8): + if has_in_col[i]: + active_indices.add(i) + + F = [0] * 8 + for idx in active_indices: + F[idx] = 1 + node_F[coord] = F + node_active[coord] = sorted(list(active_indices)) + + memo_intra = {} + memo_inter = {} + + def extract_NxN_intra(node_coord): + if node_coord in memo_intra: + return memo_intra[node_coord] + + active = node_active[node_coord][:self.max_children] + n = nodes_map[node_coord] + + if len(active) > 0: + M = np.zeros((self.max_children, self.max_children), dtype=np.int8) + sub = n.intra_matrix[np.ix_(active, active)] + M[:len(active), :len(active)] = sub + np.fill_diagonal(M[:len(active), :len(active)], 1) + res = M.tolist() + else: + res = np.zeros((self.max_children, self.max_children), dtype=np.int8).tolist() + + memo_intra[node_coord] = res + return res + + def extract_NxN_inter(coord_A, coord_B): + if coord_A == coord_B: + return extract_NxN_intra(coord_A) + + key = (coord_A, coord_B) + if key in memo_inter: + return memo_inter[key] + + active_A = node_active[coord_A][:self.max_children] + active_B = node_active[coord_B][:self.max_children] + nA = nodes_map[coord_A] + orig_M = nA.inter_matrices.get(coord_B, empty_mat) + + if len(active_A) > 0 and len(active_B) > 0: + M = np.zeros((self.max_children, self.max_children), dtype=np.int8) + M[:len(active_A), :len(active_B)] = orig_M[np.ix_(active_A, active_B)] + res = M.tolist() + else: + res = np.zeros((self.max_children, self.max_children), dtype=np.int8).tolist() + + memo_inter[key] = res + return res + + raw_tokens = [] + output_tokens = [] + + for face in raw_faces: + v0_c, v1_c, v2_c = face + raw_tokens.append([v0_c, v1_c, v2_c]) + + F0, F1, F2 = node_F[v0_c], node_F[v1_c], node_F[v2_c] + + M00 = extract_NxN_intra(v0_c) + M11 = extract_NxN_intra(v1_c) + M22 = extract_NxN_intra(v2_c) + + M01 = extract_NxN_inter(v0_c, v1_c) + M02 = extract_NxN_inter(v0_c, v2_c) + M12 = extract_NxN_inter(v1_c, v2_c) + + f_encoded = ( + v0_c, v1_c, v2_c, + F0, F1, F2, + M00, M11, M22, M01, M02, M12 + ) + output_tokens.append(f_encoded) + + return raw_tokens, output_tokens + +class OctreeNode: + def __init__(self, x, y, z, depth): + self.coord = (x, y, z) + self.depth = depth + self.children = [None] * 8 + self.intra_neighbors = set() + self.inter_neighbors = set() + self.intra_matrix = np.zeros((8, 8), dtype=np.int8) + self.inter_matrices = {} + +class MatrixOctree: + def __init__(self, max_depth=7): + self.max_depth = max_depth + self.root = None + + def _get_parent_coord(self, coord): + return (coord[0] >> 1, coord[1] >> 1, coord[2] >> 1) + + def _get_local_index(self, coord): + return ((coord[2] & 1) << 2) | ((coord[1] & 1) << 1) | (coord[0] & 1) + + def build(self, vertices, faces): + resolution = 2 ** self.max_depth + quantized = discretize(vertices, num_discrete=resolution) + + leaf_nodes_map = {} + for i in range(len(vertices)): + coord = tuple(quantized[i]) + if coord not in leaf_nodes_map: + leaf_nodes_map[coord] = OctreeNode(coord[0], coord[1], coord[2], self.max_depth) + + for face in faces: + coords = [tuple(quantized[i]) for i in face] + for i in range(3): + u, v = coords[i], coords[(i+1)%3] + if u == v: continue + + parent_u = self._get_parent_coord(u) + parent_v = self._get_parent_coord(v) + + node_u = leaf_nodes_map[u] + node_v = leaf_nodes_map[v] + + if parent_u == parent_v: + node_u.intra_neighbors.add(v) + node_v.intra_neighbors.add(u) + else: + node_u.inter_neighbors.add(v) + node_v.inter_neighbors.add(u) + + current_layer_nodes = list(leaf_nodes_map.values()) + + for d in range(self.max_depth, 0, -1): + parent_layer_map = {} + for child in current_layer_nodes: + parent_coord = self._get_parent_coord(child.coord) + if parent_coord not in parent_layer_map: + parent_layer_map[parent_coord] = OctreeNode(*parent_coord, d - 1) + + parent_node = parent_layer_map[parent_coord] + idx = self._get_local_index(child.coord) + parent_node.children[idx] = child + + for child in current_layer_nodes: + parent_coord = self._get_parent_coord(child.coord) + parent_node = parent_layer_map[parent_coord] + child_idx = self._get_local_index(child.coord) + + for sibling_coord in child.intra_neighbors: + sibling_idx = self._get_local_index(sibling_coord) + parent_node.intra_matrix[child_idx, sibling_idx] = 1 + + for neighbor_coord in child.inter_neighbors: + neighbor_parent_coord = self._get_parent_coord(neighbor_coord) + if neighbor_parent_coord not in parent_node.inter_matrices: + parent_node.inter_matrices[neighbor_parent_coord] = np.zeros((8, 8), dtype=np.int8) + + neighbor_child_idx = self._get_local_index(neighbor_coord) + parent_node.inter_matrices[neighbor_parent_coord][child_idx, neighbor_child_idx] = 1 + + grandparent_p = self._get_parent_coord(parent_coord) + grandparent_n = self._get_parent_coord(neighbor_parent_coord) + + if grandparent_p == grandparent_n: + parent_node.intra_neighbors.add(neighbor_parent_coord) + else: + parent_node.inter_neighbors.add(neighbor_parent_coord) + + current_layer_nodes = list(parent_layer_map.values()) + if current_layer_nodes: + self.root = current_layer_nodes[0] + + +def reconstruct_next_level_optimized(output_tokens, resolution, max_children=4): + all_face_coords = [] + + indices = np.arange(8) + offsets = np.stack([indices & 1, (indices >> 1) & 1, (indices >> 2) & 1], axis=1).astype(np.int32) + + def expand_intra(F, M_NxN): + M8x8 = np.zeros((8, 8), dtype=bool) + active = [i for i, val in enumerate(F) if val] + for i_n, i_o in enumerate(active): + for j_n, j_o in enumerate(active): + if i_n < max_children and j_n < max_children and M_NxN[i_n][j_n]: + M8x8[i_o][j_o] = True + return M8x8 + + def expand_inter(FA, FB, M_NxN): + M8x8 = np.zeros((8, 8), dtype=bool) + activeA = [i for i, val in enumerate(FA) if val] + activeB = [i for i, val in enumerate(FB) if val] + for i_n, i_o in enumerate(activeA): + for j_n, j_o in enumerate(activeB): + if i_n < max_children and j_n < max_children and M_NxN[i_n][j_n]: + M8x8[i_o][j_o] = True + return M8x8 + + for f_token in output_tokens: + v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12 = f_token + + v_parents = [np.array(v[:3], dtype=np.int32) for v in [v0_c, v1_c, v2_c]] + + in0 = expand_intra(F0, M00) + in1 = expand_intra(F1, M11) + in2 = expand_intra(F2, M22) + + m01 = expand_inter(F0, F1, M01) + m02 = expand_inter(F0, F2, M02) + m12 = expand_inter(F1, F2, M12) + + # in_mats = [in0, in1, in2] + # inter_edges = { + # (0, 1): m01, + # (1, 2): m12, + # (0, 2): m02 + # } + + # def are_connected_opt(p_a, c_a, p_b, c_b): + # # p_a 和 p_b 是 0, 1, 2 这三个顶点索引。 + # # 只要索引相同，说明是查询同一个顶点内部；索引不同，无脑查 inter_edges + # if p_a == p_b: + # return in_mats[p_a][c_a, c_b] + # if p_a < p_b: + # return inter_edges[(p_a, p_b)][c_a, c_b] + # else: + # return inter_edges[(p_b, p_a)][c_b, c_a] + + # from itertools import combinations_with_replacement + # for pi, pj, pk in combinations_with_replacement(range(3), 3): + # for i in range(8): + # start_j = i + 1 if pi == pj else 0 + # for j in range(start_j, 8): + # if not are_connected_opt(pi, i, pj, j): continue + + # start_k = j + 1 if pj == pk else 0 + # for k in range(start_k, 8): + # if are_connected_opt(pj, j, pk, k) and are_connected_opt(pk, k, pi, i): + # c0 = v_parents[pi] * 2 + offsets[i] + # c1 = v_parents[pj] * 2 + offsets[j] + # c2 = v_parents[pk] * 2 + offsets[k] + # all_face_coords.append(np.stack([c0, c1, c2])) + # 1. 构建 24x24 全局邻接矩阵 A + A = np.zeros((24, 24), dtype=bool) + + # 填充内部矩阵 (Intra)，使用 | .T 确保无向图的对称性 + A[0:8, 0:8] = in0 | in0.T + A[8:16, 8:16] = in1 | in1.T + A[16:24, 16:24] = in2 | in2.T + + # 填充外部跨体素矩阵 (Inter) + A[0:8, 8:16] = m01 + A[8:16, 0:8] = m01.T + A[0:8, 16:24] = m02 + A[16:24, 0:8] = m02.T + A[8:16, 16:24] = m12 + A[16:24, 8:16] = m12.T + + # 2. 向量化寻找所有的 3-clique (三角形) + # 只需要一行代码，判断 i-j, j-k, i-k 是否同时连通 + valid_mask = A[_IDX_I, _IDX_J] & A[_IDX_J, _IDX_K] & A[_IDX_I, _IDX_K] + valid_triangles = _TRI_INDICES[valid_mask] # 过滤出所有形成三角形的节点组合 + + # 3. 批量计算物理坐标 + if len(valid_triangles) > 0: + # 将 0~23 的索引还原为父节点索引 (p_idx: 0~2) 和子节点索引 (c_idx: 0~7) + p_idx = valid_triangles // 8 + c_idx = valid_triangles % 8 + + v_parents_arr = np.stack(v_parents) # shape: (3, 3) + + # 批量算出所有符合条件的三角面的世界坐标，shape: (N, 3, 3) + coords = v_parents_arr[p_idx] * 2 + offsets[c_idx] + all_face_coords.append(coords) + + if not all_face_coords: + return np.array([]), np.array([]) + + flat_v = np.concatenate(all_face_coords, axis=0).reshape(-1, 3) + V, inv = np.unique(flat_v, axis=0, return_inverse=True) + F_raw = inv.reshape(-1, 3) + + # mask_not_degenerate = (F_raw[:, 0] != F_raw[:, 1]) & \ + # (F_raw[:, 1] != F_raw[:, 2]) & \ + # (F_raw[:, 2] != F_raw[:, 0]) + # F_valid = F_raw[mask_not_degenerate] + F_valid = F_raw + + F_sorted = np.sort(F_valid, axis=1) + _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True) + F = F_valid[unique_face_indices] + V = undiscretize(V, num_discrete=resolution) + return V, F + +def reconstruct_base_mesh(output_tokens, resolution): + all_face_coords = [] + for f_token in output_tokens: + v0_c, v1_c, v2_c = f_token + c0 = np.array(v0_c[:3], dtype=np.int32) + c1 = np.array(v1_c[:3], dtype=np.int32) + c2 = np.array(v2_c[:3], dtype=np.int32) + all_face_coords.append(np.stack([c0, c1, c2])) + + if not all_face_coords: + return np.array([]), np.array([]) + + flat_v = np.concatenate(all_face_coords, axis=0).reshape(-1, 3) + V, inv = np.unique(flat_v, axis=0, return_inverse=True) + F_raw = inv.reshape(-1, 3) + + # mask_not_degenerate = (F_raw[:, 0] != F_raw[:, 1]) & \ + # (F_raw[:, 1] != F_raw[:, 2]) & \ + # (F_raw[:, 2] != F_raw[:, 0]) + # F_valid = F_raw[mask_not_degenerate] + F_valid = F_raw + + F_sorted = np.sort(F_valid, axis=1) + _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True) + F = F_valid[unique_face_indices] + V = undiscretize(V, num_discrete=resolution) + return V, F + + +class MeshCodec: + def __init__(self, args): + self.max_depth = args.max_depth + self.max_resolution = args.max_resolution + self.use_hourglass = args.use_hourglass + self.max_children = args.max_children + + self.tokenizer = DynamicMeshTokenizer(resolution=self.max_resolution, max_children=self.max_children) + + def encode_mesh(self, mesh: trimesh.Trimesh, target_depth: int) -> tuple: + + mesh_copy = mesh.copy() + if isinstance(mesh_copy, trimesh.Scene): + mesh_copy = mesh_copy.dump(concatenate=True) + + v_restored = discretize(mesh_copy.vertices, num_discrete=2**(target_depth+1)) + v_restored = undiscretize(v_restored, num_discrete=2**(target_depth+1)) + + # 4. 生成新的量化网格 + quantized_mesh = trimesh.Trimesh(vertices=v_restored, faces=mesh.faces, process=False) + + # tree = MatrixOctree(max_depth=self.max_depth) + tree = MatrixOctree(max_depth=target_depth+1) + tree.build(mesh_copy.vertices, mesh_copy.faces) + + builder = OctreeFaceTokenBuilder(tree, max_children=self.max_children) + + raw_t, out_t = builder.encode_level_tokens(depth=target_depth) + + token_sequence = self.tokenizer.encode(out_t, self.use_hourglass) + raw_token_sequence = self.tokenizer.encode_raw(raw_t, self.use_hourglass) + + return raw_token_sequence, token_sequence, quantized_mesh + + def decode_tokens(self, token_sequence: list, target_depth: int) -> trimesh.Trimesh: + if not token_sequence: + return None + + resolution = 2 ** (target_depth + 1) + out_t = self.tokenizer.decode(token_sequence) + v_next, f_next = reconstruct_next_level_optimized(out_t, resolution, max_children=self.max_children) + + if len(v_next) == 0 or len(f_next) == 0: + return None + + detokenized_mesh = trimesh.Trimesh(vertices=v_next, faces=f_next, process=False) + # trimesh.repair.fix_normals(detokenized_mesh) + # clean_mesh(detokenized_mesh) + return detokenized_mesh + + def decode_raw_tokens(self, raw_token_sequence: list, target_depth:int) -> trimesh.Trimesh: + if not raw_token_sequence: + return None + + resolution = 2 ** (target_depth) + + # [修改] 先使用 tokenizer 从 1D 整数序列解码回 raw_t 结构 + raw_t = self.tokenizer.decode_raw(raw_token_sequence) + + v,f = reconstruct_base_mesh(raw_t, resolution) + if len(v) == 0 or len(f) == 0: + return None + + detokenized_mesh = trimesh.Trimesh(vertices=v, faces=f, process=False) + # trimesh.repair.fix_normals(detokenized_mesh) + # clean_mesh(detokenized_mesh) + + return detokenized_mesh + + def get_model_config(self) -> dict: + return { + "total_vocab_size": self.tokenizer.vocab_size, + "coord_vocab_size": self.tokenizer.resolution, + "f_vocab_size": 256, + "m_vocab_size": self.tokenizer.m_vocab_size, + "coord_offset": self.tokenizer.coord_offset, + "f_offset": self.tokenizer.f_offset, + "m_intra_offset": self.tokenizer.m_intra_offset, + "m_inter_offset": self.tokenizer.m_inter_offset, + "pad_token_id": self.tokenizer.pad_id, + "max_children": self.max_children, + "input_num_categories": self.tokenizer.resolution+1, + "output_num_categories": 1 + 2*2**self.tokenizer.m_bits + 256, + } + +# ========================================== +# 测试重构与无损化流程 +# ========================================== +if __name__ == "__main__": + # 1. 定义输入和输出文件夹 + INPUT_DIRECTORY = "/root/junkai_lin/mul_res_diffusion/demo/complex_test_data/files" # 【修改这里】替换为你的模型文件夹路径 + OUTPUT_DIRECTORY = "./output/complex_child8" + os.makedirs(OUTPUT_DIRECTORY, exist_ok=True) + import time + @dataclass + class MeshCodecArgs: + max_depth: int = 10 + max_resolution: int = 1024 + use_hourglass: bool = True + max_children: int = 8 + + # 2. 初始化 Codec (放在循环外，避免重复加载) + args = MeshCodecArgs() + codec = MeshCodec(args=args) + + print(f"[*] Initialized Codec. Total Vocab Size: {codec.tokenizer.vocab_size}") + + # 3. 获取所有 .obj 和 .glb 文件 + input_dir_path = Path(INPUT_DIRECTORY) + # rglob 会递归查找子文件夹，如果只需要当前层级，可以把 rglob 改成 glob + mesh_files = list(input_dir_path.rglob("*.obj")) + list(input_dir_path.rglob("*.glb")) + + if not mesh_files: + print(f"[!] No .obj or .glb files found in {INPUT_DIRECTORY}") + + # 4. 遍历处理每个文件 + for idx, file_path in enumerate(mesh_files, 1): + obj_name = file_path.stem # 获取不带后缀的文件名 + print(f"\n{'='*50}") + print(f"[{idx}/{len(mesh_files)}] Processing: {file_path.name}") + print(f"{'='*50}") + file_start_time = time.perf_counter() + # try: + # 加载并归一化 Mesh + raw_mesh = trimesh.load(str(file_path), force="mesh") + normalize(raw_mesh) + # clean_mesh(raw_mesh) + + if raw_mesh: + out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj") + raw_mesh.export(out_path) + print(f"[*] Successfully saved GT to: {out_path}") + + # 开始不同深度 (Depth) 的推理 + for idx, file_path in enumerate(mesh_files, 1): + obj_name = file_path.stem + print(f"\n{'='*50}") + print(f"[{idx}/{len(mesh_files)}] Processing: {file_path.name}") + print(f"{'='*50}") + + file_start_time = time.perf_counter() # 记录单个文件处理开始 + + raw_mesh = trimesh.load(str(file_path), force="mesh") + normalize(raw_mesh) + + if raw_mesh: + out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj") + raw_mesh.export(out_path) + + # 开始不同深度 (Depth) 的推理 + for depth in [0,1,2,3,4,5,6]: + print(f"\n--- Depth {depth} | N={args.max_children} ---") + + # --- 计时：Encoding --- + t0 = time.perf_counter() + raw_tokens, tokens, quantized_mesh = codec.encode_mesh(raw_mesh, target_depth=depth) + encode_time = time.perf_counter() - t0 + print(f"[Timer] Encoding: {encode_time:.4f}s") + + # --- 计时：Decoding (Fine) --- + if len(tokens) > 0: + t1 = time.perf_counter() + recon_mesh = codec.decode_tokens(tokens, target_depth=depth) + decode_fine_time = time.perf_counter() - t1 + print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s") + + # 导出保存 + out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_recon_d{depth}_res{2**(depth+1)}_N{args.max_children}.obj") + recon_mesh.export(out_path) + + # --- 计时：Decoding (Coarse/Raw) --- + if len(raw_tokens) > 0: + t2 = time.perf_counter() + raw_recon_mesh = codec.decode_raw_tokens(raw_tokens, target_depth=depth) + decode_coarse_time = time.perf_counter() - t2 + print(f"[Timer] Decoding (Coarse): {decode_coarse_time:.4f}s") + + raw_out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_raw_recon_d{depth}_res{2**depth}.obj") + raw_recon_mesh.export(raw_out_path) + + if quantized_mesh: + quant_out_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_dis_d{depth+1}_res{2**(depth+1)}.obj") + quantized_mesh.export(quant_out_path) + + file_total_time = time.perf_counter() - file_start_time + print(f"\n[*] Finished {file_path.name} in {file_total_time:.2f}s") \ No newline at end of file diff --git a/mulres_infer/modules/octree_decode.py b/mulres_infer/modules/octree_decode.py new file mode 100644 index 0000000000000000000000000000000000000000..dd77bda50a4115550fdb69ed469ca5118ed2c58c --- /dev/null +++ b/mulres_infer/modules/octree_decode.py @@ -0,0 +1,625 @@ +import gc +import glob +import math +import os +import random +import re +from pathlib import Path + +import numpy as np +import torch +import trimesh +from accelerate import Accelerator, DistributedDataParallelKwargs +from tqdm import tqdm + +import importlib.util + +from utils.config import load_config +from dataset.mesh_dataset import mul_encode_mesh, discrete_undiscrete_and_clean +from models.generator import FactorizedDiffusionTransformer +from modules.octree_codec import MeshCodec +from utils.postprocessing import _fill_holes +from utils.data_utils import normalize + +# fix3 编码格式的数据通路(与训练同一份实现, 已抽成独立模块) +from modules.token_utils import split_fix3_token_list, subcodes_to_flat_token +def set_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ["PYTHONHASHSEED"] = str(seed) + +def _depth_from_resolution(resolution): + if resolution <= 0: + raise ValueError(f"resolution must be positive, got {resolution}") + depth = int(round(math.log2(resolution))) + if (1 << depth) != resolution: + raise ValueError(f"resolution must be power of 2, got {resolution}") + return depth + + +def read_dataset_paths(txt_path): + mesh_paths = [] + with open(txt_path, "r") as f: + for line in f: + path = line.strip() + if path: + mesh_paths.append(path) + return mesh_paths + + +def raw_stem_to_pair_key(raw_path): + stem = Path(raw_path).stem + # raw name example: 1(42)_y_f_4244 -> pair key: 1(42) + return re.sub(r"_y_f_\d+$", "", stem) + + +def gt_stem_to_pair_key(gt_path): + return Path(gt_path).stem + +def get_face_count(file_path): + """ + 使用正则表达式从文件名中提取面片数量 + 支持格式: *_f{num}_v*.obj 或 *_facenum_{num}_*.glb + """ + filename = Path(file_path).name + match = re.search(r'_f(\d+)_v|_facenum_(\d+)', filename) + if match: + # group(1) 对应 _f(\d+)_v 里的数字 + # group(2) 对应 _facenum_(\d+) 里的数字 + return int(match.group(1) or match.group(2)) + return None + +# def build_raw_gt_pairs(raw_txt_path, gt_txt_path): +# raw_paths = read_dataset_paths(raw_txt_path) +# gt_paths = read_dataset_paths(gt_txt_path) + +# # If two txt paths are the same source, pair each mesh with itself directly. +# if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve(): +# return [(p, p) for p in raw_paths], [] + +# gt_by_key = {} +# for gt_path in gt_paths: +# gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path + +# pairs = [] +# missing_gt = [] +# for raw_path in raw_paths: +# key = raw_stem_to_pair_key(raw_path) +# gt_path = gt_by_key.get(key) +# if gt_path is None: +# missing_gt.append(raw_path) +# continue +# pairs.append((raw_path, gt_path)) + +# return pairs, missing_gt +def build_raw_gt_pairs(raw_txt_path, gt_txt_path, min_faces=None, max_faces=None): + raw_paths = read_dataset_paths(raw_txt_path) + gt_paths = read_dataset_paths(gt_txt_path) + + # ===== 新增过滤逻辑 ===== + filtered_raw_paths = [] + for p in raw_paths: + if min_faces is None and max_faces is None: + filtered_raw_paths.append(p) + continue + + faces = get_face_count(p) + if faces is not None: + # 检查是否在指定范围内 + if min_faces is not None and faces < min_faces: + continue + if max_faces is not None and faces > max_faces: + continue + filtered_raw_paths.append(p) + else: + # 如果文件名中没找到面片数特征，默认保留（你也可以改为 continue 丢弃它） + filtered_raw_paths.append(p) + + raw_paths = filtered_raw_paths + # ======================== + + # If two txt paths are the same source, pair each mesh with itself directly. + if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve(): + return [(p, p) for p in raw_paths], [] + + gt_by_key = {} + for gt_path in gt_paths: + gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path + + pairs = [] + missing_gt = [] + for raw_path in raw_paths: + key = raw_stem_to_pair_key(raw_path) + gt_path = gt_by_key.get(key) + if gt_path is None: + missing_gt.append(raw_path) + continue + pairs.append((raw_path, gt_path)) + + return pairs, missing_gt + +def save_gt_mesh(mesh, out_dir, name): + gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False) + gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]] + gt_mesh.export(os.path.join(out_dir, f"{name}_gt.obj")) + +def save_raw_gt_mesh(mesh, out_dir, name): + gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False) + gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]] + gt_mesh.export(os.path.join(out_dir, f"{name}_raw_gt.obj")) + + +def save_stage_input_mesh(mesh, out_dir, name, stage_in_res, stage_out_res, stage_in_depth): + coarse_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False) + coarse_mesh.vertices = coarse_mesh.vertices[:, [1, 2, 0]] + suffix = f"_stage_{stage_in_res}to{stage_out_res}_input_coarse" + coarse_mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{stage_in_depth}.obj")) + + +def save_stage_prediction(codec, coord_seq, gen_subcodes, codec_cfg, out_dir, name, target_depth, suffix, fix_holes, raw_dis_mesh=None, gt_dis_mesh=None, remove_degenerate=True): + # fix3 数据通路：coord(粗层条件, 含 offset) + 采样得到的 refine 子码(local[0,256)) -> + # subcodes_to_flat_token 还原成扁平 60/face token(含 offset) -> codec.decode_tokens 解成 mesh。 + # 与 acc_v2_fix3.fix3_generate 完全一致；不再做旧 18/face 的 coord/refine 交错。 + flat = subcodes_to_flat_token(coord_seq, gen_subcodes, codec_cfg) + if not flat: + return None + + mesh = codec.decode_tokens(flat, target_depth=target_depth - 1) + if mesh is None: + return None + + # 先剔除退化面(三顶点存在重合的零面积面)，再做 normals/holes 修复，避免退化面干扰后续几何修复。 + if remove_degenerate: + mesh = remove_degenerate_faces(mesh) + + trimesh.repair.fix_normals(mesh) + if fix_holes: + verts = torch.tensor(mesh.vertices, device="cuda", dtype=torch.float32) + faces = torch.tensor(mesh.faces, device="cuda", dtype=torch.int32) + new_verts, new_faces = _fill_holes(verts, faces, debug=False, verbose=False, only_fill_holes=True) + mesh = trimesh.Trimesh(vertices=new_verts.cpu().numpy(), faces=new_faces.cpu().numpy()) + + return_mesh = mesh.copy() + mesh.vertices = mesh.vertices[:, [1, 2, 0]] + mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{target_depth}.obj")) + + if raw_dis_mesh is not None: + raw_dis_mesh.vertices = raw_dis_mesh.vertices[:, [1, 2, 0]] + raw_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_raw_dis_depth{target_depth-1}.obj")) + if gt_dis_mesh is not None: + gt_dis_mesh.vertices = gt_dis_mesh.vertices[:, [1, 2, 0]] + gt_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_gt_dis_depth{target_depth}.obj")) + return return_mesh + + +def _decode_subcodes_to_mesh(codec, coord_seq, gen_subcodes, codec_cfg, target_depth): + """与 save_stage_prediction 同款 decode(不 fix_normals/fix_holes/不导出)，用于 best-of-N 打分。""" + flat = subcodes_to_flat_token(coord_seq, gen_subcodes, codec_cfg) + if not flat: + return None + return codec.decode_tokens(flat, target_depth=target_depth - 1) + + +def _pc_chamfer_score(mesh, pc_xyz, n_sample=4096): + """候选 mesh 表面采样点 与 条件点云 pc_xyz 的对称 chamfer(越小越贴合目标点云)。 + 两者同处归一化 frame(mul_encode_mesh 内 normalize + decode 反离散化)，可直接算。 + 退化面(零面积)不会被 sample_surface 采到 -> 天然只评有效曲面；无有效面返回 inf。""" + if mesh is None or len(getattr(mesh, "faces", [])) == 0 or len(mesh.vertices) == 0: + return float("inf") + try: + pts, _ = trimesh.sample.sample_surface(mesh, n_sample) + pts = np.asarray(pts) + except Exception: + pts = np.asarray(mesh.vertices) + if pts.shape[0] == 0: + return float("inf") + pts = torch.as_tensor(pts, dtype=torch.float32, device=pc_xyz.device) + ref = pc_xyz + if ref.shape[0] > n_sample: + idx = torch.randperm(ref.shape[0], device=ref.device)[:n_sample] + ref = ref[idx] + d = torch.cdist(pts, ref) # [Np, Nr] + return float((d.min(1).values.mean() + d.min(0).values.mean()).item()) + + +def remove_degenerate_faces(mesh): + """删除退化面：三角形三个顶点中存在重合顶点(任意两顶点坐标相同 -> 零面积/退化)。 + 按顶点坐标判定(而非仅 index)，因为 octree 解码后不同 index 可能落到同一离散坐标。 + 返回过滤后的新 mesh；若没有退化面则原样返回；若全部退化则返回原 mesh(避免产出空 mesh)。""" + if mesh is None or len(getattr(mesh, "faces", [])) == 0 or len(getattr(mesh, "vertices", [])) == 0: + return mesh + faces = np.asarray(mesh.faces) + verts = np.asarray(mesh.vertices) + v0, v1, v2 = verts[faces[:, 0]], verts[faces[:, 1]], verts[faces[:, 2]] + # 三条边都非零长度(三顶点两两不重合)才保留 + keep = ( + np.any(v0 != v1, axis=1) + & np.any(v1 != v2, axis=1) + & np.any(v0 != v2, axis=1) + ) + if keep.all(): + return mesh + if not keep.any(): + return mesh + return trimesh.Trimesh(vertices=verts.copy(), faces=faces[keep].copy(), process=False) + + +def process_one_mesh( + raw_mesh_path, + gt_mesh_path, + model, + codec, + codec_cfg, + device, + pad_id, + config, + output_folder, + stage_resolutions, + fix_holes=True, + use_gt_pc=True, + remove_degenerate=True, +): + import time + t_start = time.time() + import time + t_total_start = time.time() + + name = Path(raw_mesh_path).stem + t_start = time.time() + mesh = trimesh.load(raw_mesh_path, force="mesh") + mesh.vertices = mesh.vertices[:, [2, 0, 1]] + gt_mesh = trimesh.load(gt_mesh_path, force="mesh") + gt_mesh.vertices = gt_mesh.vertices[:, [2, 0, 1]] + normalize(gt_mesh) + normalize(mesh) + save_raw_gt_mesh(mesh, output_folder, name) + save_gt_mesh(gt_mesh, output_folder, name) + t_load = time.time() + print(f"[{name}] Load & normalize time: {t_load - t_start:.4f}s") + + stage_depths = [_depth_from_resolution(x) for x in stage_resolutions] + + # First-stage input uses codec encode-decode result at the minimum resolution. + min_stage_depth = stage_depths[0] + t_start = time.time() + init_raw_token, _, init_quantized_mesh = codec.encode_mesh(mesh, target_depth=min_stage_depth) + init_mesh = codec.decode_raw_tokens(init_raw_token, target_depth=min_stage_depth) + t_init = time.time() + print(f"[{name}] Init encode & decode (depth={min_stage_depth}): {t_init - t_start:.4f}s") + + if init_mesh is not None: + current_mesh = init_mesh + elif init_quantized_mesh is not None: + current_mesh = init_quantized_mesh + else: + current_mesh = mesh.copy() + + for stage_idx in range(len(stage_depths) - 1): + stage_start_time = time.time() + stage_in_depth = stage_depths[stage_idx] + stage_out_depth = stage_depths[stage_idx + 1] + stage_in_res = stage_resolutions[stage_idx] + stage_out_res = stage_resolutions[stage_idx + 1] + + save_stage_input_mesh( + mesh=current_mesh, + out_dir=output_folder, + name=name, + stage_in_res=stage_in_res, + stage_out_res=stage_out_res, + stage_in_depth=stage_in_depth, + ) + + t_start = time.time() + stage_data = mul_encode_mesh( + raw_mesh=current_mesh, + gt_mesh=gt_mesh, + codec=codec, + tokenizer_target_depth=stage_in_depth, + pc_num=config.data_processing.pc_num, + total_pc_num=config.data_processing.total_pc_num, + use_gt_pc=use_gt_pc, + clean=False + ) + t_mul_encode = time.time() + print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | mul_encode_mesh time: {t_mul_encode - t_start:.4f}s") + + pc = stage_data["pc"].unsqueeze(0).to(device) + + # fix3 数据通路(与 acc_v2_fix3.fix3_generate 一致)： + # 扁平 60/face token -> 取 coord(粗层条件) -> 因子分解采样 refine 子码。 + # 这里只用 coord 作为条件，refine 子码全部由模型 sample 生成(GT refine 不参与推理)。 + coord_list, _ = split_fix3_token_list([stage_data["token"]], pad_id, device, codec_cfg) + depth_t = torch.tensor([stage_in_depth], device=device, dtype=torch.long) + + cfg_scale = float(getattr(config.evaluation, "cfg_scale", 1.0)) + # CFG ramp：linear 时每级采样的最初几步用低 cfg(cfg_start)，按步线性升到 cfg_scale， + # 避免高 cfg 在早期步把没把握的位置置信度抬高、过早 commit。constant=旧行为。 + cfg_schedule = str(getattr(config.evaluation, "cfg_schedule", "constant")) + cfg_start = float(getattr(config.evaluation, "cfg_start", 1.0)) + best_of_n = int(getattr(config.evaluation, "best_of_n", 1)) + # best-of-N 对抗级联累积误差：每级采 N 个候选 -> 各自 decode -> 选 chamfer-to-pc 最小者(最贴合 + # 目标点云)作为本级输出喂下一级。点云 pc 是跨级不漂移的全局锚，用它选优能降低单级错误率、抑制误差累积。 + # N=1 退化为单次贪心采样(确定性, 行为同前)；N>1 需 temperature>0 才能采出多样候选。 + sample_temp = float(getattr(config.evaluation, "sample_temperature", 1.0)) if best_of_n > 1 else 0.0 + pc_xyz = stage_data["pc"][:, :3].to(device).float() + + t_start = time.time() + gen_subcodes, best_score = None, float("inf") + for cand in range(max(1, best_of_n)): + gs = model.sample( + coord_list, + depth_t, + pc=pc, + T_infer=config.evaluation.diffusion_sample_steps, + cfg_scale=cfg_scale, + temperature=sample_temp, + cfg_schedule=cfg_schedule, + cfg_start=cfg_start, + ) + if best_of_n <= 1: + gen_subcodes = gs + break + cand_mesh = _decode_subcodes_to_mesh(codec, coord_list[0], gs[0], codec_cfg, stage_out_depth) + score = _pc_chamfer_score(cand_mesh, pc_xyz) + if score < best_score: + best_score, gen_subcodes = score, gs + torch.cuda.synchronize() + t_sample = time.time() + if best_of_n > 1: + print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | best-of-{best_of_n} (T={sample_temp}, cfg={cfg_scale}) " + f"chamfer={best_score:.6f} | sample time: {t_sample - t_start:.4f}s") + else: + print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | factorized sample time: {t_sample - t_start:.4f}s") + + t_start = time.time() + stage_suffix = f"_stage_{stage_in_res}to{stage_out_res}" + stage_mesh = save_stage_prediction( + codec=codec, + coord_seq=coord_list[0], + gen_subcodes=gen_subcodes[0], + codec_cfg=codec_cfg, + out_dir=output_folder, + name=name, + target_depth=stage_out_depth, + suffix=stage_suffix, + fix_holes=fix_holes, + raw_dis_mesh=stage_data.get("raw_dis_mesh"), + gt_dis_mesh=stage_data.get("gt_dis_mesh"), + remove_degenerate=remove_degenerate, + ) + t_save_pred = time.time() + print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | save_stage_prediction (decode + fix) time: {t_save_pred - t_start:.4f}s") + print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | total stage time: {time.time() - stage_start_time:.4f}s\n") + + if stage_mesh is None: + break + + current_mesh = stage_mesh + + print(f"[{name}] Total process_one_mesh time: {time.time() - t_total_start:.4f}s") + + +class EMA: + def __init__(self, model, decay=0.999): + self.decay = decay + self.shadow = {} + self.register(model) + + def register(self, model): + self.shadow = {} + for name, param in model.named_parameters(): + if param.requires_grad: + self.shadow[name] = param.data.clone() + + def load_state_dict(self, state): + self.decay = state.get("decay", self.decay) + self.shadow = {k: v.clone() for k, v in state["shadow"].items()} + + @torch.no_grad() + def copy_to(self, model): + """Swap shadow weights into the model in-place (for inference).""" + copied, missing = 0, 0 + for name, param in model.named_parameters(): + if name in self.shadow: + param.data.copy_(self.shadow[name].to(device=param.device, dtype=param.dtype)) + copied += 1 + else: + missing += 1 + print(f"[EMA.copy_to] copied {copied} params into model ({missing} not present in shadow)") + + +def run_inference_on_all_checkpoints(model, accelerator, config, device, pad_id, codec, codec_cfg, ema=None): + ckpt_dir = "ckpt_fix" + cfg_scale = float(getattr(config.evaluation, "cfg_scale", 1.0)) # 采样 CFG 强度(与细分循环 model.sample 一致) + output_folder = f"./pretrain_model_mul_more_data_1_to_64_eval_fix_easy2_cfg_{cfg_scale}_step_20" + + # raw_data_txt = "dataset/bpt_res_list.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_retrain_gt/data.txt + # gt_data_txt = "dataset/trellis_clean_mesh.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_gt/data.txt + # raw_data_txt = "/root/dataset/unique_files_glb_under20000face_compress_glb_with_info_val.txt" + # gt_data_txt = "/root/dataset/unique_files_glb_under20000face_compress_glb_with_info_val.txt" + raw_data_txt = "/root/dataset/unique_files_glb_under6000face_2degree_30ratio_0.01_with_info_val.txt" + gt_data_txt = "/root/dataset/unique_files_glb_under6000face_2degree_30ratio_0.01_with_info_val.txt" + + # raw_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt" + # gt_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt" + stage_resolutions = [1,2,4,8,16,32,64] + # raw_data_txt = "demo/tripo 1.0/data.txt" + # gt_data_txt = "demo/tripo 1.0/data.txt" + + use_gt_pc=True + fix_holes=False + # 退化面删除开关：True=导出前剔除三顶点存在重合的零面积面。可在 config.evaluation.remove_degenerate_faces 覆盖。 + remove_degenerate = bool(getattr(config.evaluation, "remove_degenerate_faces", True)) + if accelerator.is_main_process: + print(f"Using raw data txt: {raw_data_txt}") + print(f"Using gt data txt: {gt_data_txt}") + print(f"Searching checkpoints in: {ckpt_dir}") + + # mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=config.data_processing.max_face_num, min_faces=config.data_processing.min_face_num) + mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=20000, min_faces=500) + if accelerator.is_main_process: + print(f"Matched raw/gt pairs: {len(mesh_pairs)}") + if missing_gt: + print(f"Raw without matched gt: {len(missing_gt)}") + + search_pattern = os.path.join(ckpt_dir, "checkpoint_*_epoch_*_loss.pth") + ckpt_files = glob.glob(search_pattern) + if not ckpt_files: + if accelerator.is_main_process: + print("No checkpoints found!") + return + + ckpt_data = [] + for f in ckpt_files: + match = re.search(r"checkpoint_(\d+)_epoch_", f) + if match: + epoch = int(match.group(1)) + ckpt_data.append((epoch, f)) + ckpt_data.sort(key=lambda x: x[0]) + + # stage_resolutions = [32, 64, 128, 256] + + + # stage_resolutions = [128, 256] + + for epoch, ckpt_path in ckpt_data: + if accelerator.is_main_process: + print(f"\n[{'='*40}]") + print(f"Starting inference for epoch {epoch}") + print(f"Loading weights from: {ckpt_path}") + + output_folder_ply = os.path.join(output_folder, f"{epoch}_epoch_val_generate_mul") + if accelerator.is_main_process: + os.makedirs(output_folder_ply, exist_ok=True) + + checkpoint = torch.load(ckpt_path, map_location=device) + model_to_load = accelerator.unwrap_model(model) + # 加载的是用「新模型代码(AR 连接头 + factorized)」训练出的 ckpt，必须 strict=True： + # key/shape 一旦不匹配立即报错，避免 strict=False 把 AR 连接头静默留成随机初始化 -> 输出垃圾还不报警。 + # (只有训练脚本 acc_v2_fix3 从旧 Linear 头 ckpt 暖启动时才用 load_strict=false，推理这里不允许。) + model_to_load.load_state_dict(checkpoint["model_state_dict"], strict=True) + + # Auto-detect EMA: if the ckpt has an `ema_state_dict`, swap those + # shadow weights into the model for inference. Otherwise fall back to + # the raw `model_state_dict` we already loaded above. No reliance on + # config `use_ema` / a separate EMA object — just check the ckpt. + ema_state = checkpoint.get("ema_state_dict") + shadow = ema_state.get("shadow") if isinstance(ema_state, dict) else None + if shadow: + with torch.no_grad(): + copied, skipped = 0, 0 + for name, param in model_to_load.named_parameters(): + if name in shadow: + param.data.copy_(shadow[name].to(device=param.device, dtype=param.dtype)) + copied += 1 + else: + skipped += 1 + if accelerator.is_main_process: + msg = f"[Epoch {epoch}] EMA shadow detected in ckpt -> applied to model ({copied} params" + if skipped: + msg += f", {skipped} model params not in shadow" + print(msg + ")") + else: + if accelerator.is_main_process: + print(f"[Epoch {epoch}] no EMA in ckpt -> running with raw model_state_dict weights.") + + with torch.no_grad(): + # 用 unwrap 后的网络做推理：FactorizedDiffusionTransformer.sample 定义在裸模型上， + # DDP 包装层没有 .sample 方法。 + model_to_load.eval() + pbar = tqdm( + enumerate(mesh_pairs), + total=len(mesh_pairs), + desc=f"Infer Epoch {epoch}", + disable=not accelerator.is_local_main_process, + dynamic_ncols=True, + ) + for line_idx, (raw_mesh_path, gt_mesh_path) in pbar: + if line_idx % accelerator.num_processes != accelerator.process_index: + continue + try: + process_one_mesh( + raw_mesh_path=raw_mesh_path, + gt_mesh_path=gt_mesh_path, + model=model_to_load, + codec=codec, + codec_cfg=codec_cfg, + device=device, + pad_id=pad_id, + config=config, + output_folder=output_folder_ply, + stage_resolutions=stage_resolutions, + fix_holes=fix_holes, + use_gt_pc=use_gt_pc, + remove_degenerate=remove_degenerate, + ) + except Exception as e: + print(f"[infer error] raw={raw_mesh_path}, gt={gt_mesh_path}: {str(e)}") + + torch.cuda.empty_cache() + gc.collect() + + if accelerator.is_main_process: + print("\nBatch inference complete for all checkpoints.") + + +def main(config): + print("start init Accelerator...") + ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator( + gradient_accumulation_steps=config.accelerator.gradient_accumulation_steps, + mixed_precision=config.accelerator.mixed_precision, + kwargs_handlers=[ddp_kwargs], + ) + device = accelerator.device + + codec = MeshCodec(args=config.tokenizer) + codec_cfg = codec.get_model_config() + input_num_categories = codec_cfg["input_num_categories"] # coord 词表 = resolution + 1 (含 offset) + pad_id = codec_cfg["pad_token_id"] + + # fix3 因子分解模型(与 acc_v2_fix3 训练时一致的构造方式)。 + # rotary 只需覆盖「模型位置」(9/face)，用 config.model.max_seq_len(回退到 max_len)。 + model_max_seq_len = int(getattr(config.model, "max_seq_len", config.data_processing.max_len)) + model = FactorizedDiffusionTransformer( + coord_vocab_size=input_num_categories, + max_seq_len=model_max_seq_len, + d_model=config.model.feature_dim, + nhead=config.model.num_heads, + depth=config.model.hourglass_depth, + dropout=config.model.dropout, + encoder_freeze=getattr(config.model, "encoder_freeze", True), + use_pc=True, + time_embed_scale=config.training.diffusion_train_steps, + nonempty_row_weight=float(getattr(config.training, "nonempty_row_weight", 1.0)), # 与 acc_v2_fix3 构造一致(推理不算 loss, 仅对齐) + ) + + use_ema = getattr(config.training, "use_ema", False) + ema_decay = getattr(config.training, "ema_decay", 0.999) + ema = EMA(model, decay=ema_decay) if use_ema else None + + model = accelerator.prepare(model) + + run_inference_on_all_checkpoints( + model=model, + accelerator=accelerator, + config=config, + device=device, + pad_id=pad_id, + codec=codec, + codec_cfg=codec_cfg, + ema=ema, + ) + + +if __name__ == "__main__": + set_seed(42) + config = load_config() + main(config) diff --git a/mulres_infer/modules/pipeline.py b/mulres_infer/modules/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..0ead7599dc91b4210fe73ded334873b918f08933 --- /dev/null +++ b/mulres_infer/modules/pipeline.py @@ -0,0 +1,964 @@ +"""多级"从零生成": res2 -> 4 -> 8 -> ... -> 目标分辨率 的级联生成。 + +基于 infer_remask_geom_altpc.py 的级联骨架, 但用 infer_single_layer.py 的干净基建 + 新 model.sample +(align_fm / align_inter / align_per_step)。**不含 remask/requery**(按需求省去), 只保留: + - 从最粗分辨率 STAGE[0] 起(GT 八叉树解码作粗起点), 每一级把"上一级预测的 mesh"当输入喂下一级(级联累积); + - 每级 model.sample 细分一层, 可选末步跨面一致性对齐(F/intra/inter); + - 可选 fill_small_holes 补小洞(不考虑法向)。 + +与单级(infer_single_layer)的区别: 单级每级喂 GT@in_depth(teacher-forced, 无累积); 本脚本喂上一级预测(真实从零级联)。 + +⚠ v2 口径(2026-08-17, 实测校准): 所有文件名 = 内容的**真实体素格**(顶点格步实测 2/res)。 + - fix3: 深度 d 的完整表示真实格 = 2^(d+1); 一个级联 stage = "在深度 d 上把占用补全成格 2^(d+1) 的 mesh", + 分辨率翻倍发生在下一级重编码——所以 STAGE 每项即该级完成的格, fix3 终点 = 格 32(深度4)。 + - acc_uni 细分器: 深度 D 真实格 = 2^D(与 fix3 差一档!), 细分终点 SUPERRES_TO=512 = 深度9。 + - 生成产物 res{r} 与 GT 天花板 gt_res{r} 同格同名, 可直接叠加; 起点若是 raw 占用块命名 raw_res{r}。 + - 历史目录三代命名: 旧标签(≤08-14, 生成侧=真实格/GT侧=一半), 08-14~17(生成侧虚高一倍), v2(全对齐); + 新目录一律带 _v2 后缀防覆盖。塌缩自 08-17 起先按坐标焊接(WELD_DEC, 修 UV 接缝撕裂)。 + +用法: + CUDA_VISIBLE_DEVICES=7 STAGE=2,4,8,16,32 DATASET=trellis NCASE=8 STEPS=20 \ + GEOM_STAGES=8,16,32 SUPERRES=1 SUPERRES_TO=512 python infer_multilevel.py +""" +import os, sys, glob as _glob, re as _re, time, math +from collections import defaultdict +import numpy as np, torch, trimesh +CONFIG = os.environ.get("CONFIG", "configs/generator.yaml") +sys.argv = ["x", "--config", CONFIG] +from utils.config import load_config +from dataset.mesh_dataset import mul_encode_mesh, sample_pc, dec +from dataset.mesh_dataset_altpc import sample_pc_with_degen +from utils.data_utils import normalize +from models.generator import REFINE_POS_PER_FACE, NUM_OCTANTS, MASK_ID, TYPE_F, TYPE_INTRA, TYPE_INTER +from models.generator import FactorizedDiffusionTransformerAltPC as BaseModel +from modules.octree_codec import MeshCodec as FastMeshCodec +import utils.distance as metric +from scipy.spatial import cKDTree +import importlib.util as _ilu +from utils.postprocessing import _fill_holes +_v2_spec = _ilu.spec_from_file_location("fix3_v2", os.path.join(os.path.dirname(os.path.abspath(__file__)), "octree_decode.py")) +fix3_v2 = _ilu.module_from_spec(_v2_spec); _v2_spec.loader.exec_module(fix3_v2) +decode_subcodes_to_mesh = fix3_v2._decode_subcodes_to_mesh +# acc_uni 高分辨率细分器(SUPERRES 用): model_acc + 旧 octree-codebook tokenizer(与 fix3 不同模型/tokenizer) +from models.subdivider import DiscreteDiffusionTransformer, sample_refine_tokens_with_topology +from modules.octree_codec_acc import MeshCodec as OctreeMeshCodec +import pyvista as pv +from tqdm import tqdm +TOKENS_PER_FACE = 60 +STAGE = [int(x) for x in os.environ.get("STAGE", "2,4,8,16,32").split(",")] # 级联阶梯: 每项 = 该级"完成"的真实体素格 +# (2026-08-17 v2 口径: 一级 = 在模型深度 d=log2(r)-1 上把占用补全成格 r 的完整 mesh; 翻倍发生在下一级重编码。 +# fix3 深度 d 的表示真实格 = 2^(d+1); 生成产物文件 res{r} 与 GT 天花板 gt_res{r} 同格可直接对照) +DATASET = os.environ.get("DATASET", "trellis") +NCASE = int(os.environ.get("NCASE", "8")); CASE_STEM = os.environ.get("CASE_STEM", "") +STEPS = int(os.environ.get("STEPS", "20")); PCT_DEG = os.environ.get("PCT_DEG", "1") == "1" +ALIGN_FM = os.environ.get("ALIGN_FM", "0") == "1" +ALIGN_INTER = os.environ.get("ALIGN_INTER", "0") == "1" +ALIGN_PER_STEP = os.environ.get("ALIGN_PER_STEP", "0") == "1" +FILL = os.environ.get("FILL", "0") == "1"; FILL_MAX_SIDES = int(os.environ.get("FILL_MAX_SIDES", "4")) +FILL_MIN_RES = int(os.environ.get("FILL_MIN_RES", "0")) # >0: FILL=1 的级联内补洞只在"产物真实格 >= 此值"的级做 + # (默认 0 = 历史行为 stage_i>=4, 在 STAGE=2..64 下即 res32/res64 两级) + # 例: FILL_MIN_RES=64 -> 只在 res64 补, res32 不补 +# method=fill: 只用直接补洞(_fill_holes)做级联修复, 在这些"输出分辨率"阶段补洞后喂下一级(默认 8/16/32/64) +FILL_STAGES = set(int(x) for x in os.environ.get("FILL_STAGES", "8,16,32").split(",") if x.strip()) # 按产物真实格 +SAVE = os.environ.get("SAVE", "1") == "1" +SAVE_GT_LEVELS = os.environ.get("SAVE_GT_LEVELS", "1") == "1" # 每 case 导出 GT 经 tokenizer 逐分辨率编码->解码的天花板(gt_res{N}.obj, 真实res命名, 含退化面), 用户要求所有多级生成默认带 +SAVE_PC_LEVELS = os.environ.get("SAVE_PC_LEVELS", "1") == "1" # 每级导出实际注入的多级点云 pc_target_res{N}.ply(云B, 该级条件) + pc_full.ply(云A, 全分辨率GT云), 用户要求默认带 +EMPTY_BIAS = float(os.environ.get("EMPTY_BIAS", "0")) # >0: 连接AR头空字节类(0=无连接)logit+bias -> 稀疏化控密度(移植自 infer_remask_geom_altpc, δ=0.5 历史面数比→1.10); requery 补洞期间自动关 +EB_ACTIVE = [True] +QUAD_FILL = os.environ.get("QUAD_FILL", "0") == "1" # 每级产物确定性微补洞: 开边图长度3/4简单环(4须双对角线缺失+bbox≤2格)补1-2三角。 + # 数据背书: GT天花板"无对角线四边洞"=0(20case×res8/16/32), 误补风险极低; 修在级间阻断洞传播 +QUAD_FILL_STATS = [0, 0] # [补的三角洞数, 补的四边洞数] +SAVE_REMASK_VIS = os.environ.get("SAVE_REMASK_VIS", "0") == "1" # geom 级导出可视化: *_remask_targets.ply(重采目标区=开边1-ring, 橙) / *_remask_changed.ply(重采后实际改变的面, 红) + + +def _face_keys(mesh): + v = np.round(np.asarray(mesh.vertices), 6) + return [tuple(sorted(map(tuple, v[f]))) for f in np.asarray(mesh.faces)] + + +def save_remask_vis(pre_mesh, post_mesh, out_prefix): + """pre=修补前解码, post=修补后。targets: pre 上开边顶点 1-ring 面(≈build_remask目标)橙色; + changed: post 上'pre 中不存在'的面红色。PLY 带面色, y-up 帧与 obj 一致。""" + GRAY, ORANGE, RED = [200, 200, 200, 255], [255, 140, 0, 255], [220, 30, 30, 255] + def yup(m): + mm = m.copy(); mm.vertices = mm.vertices[:, [1, 2, 0]]; return mm + # targets on pre + v = np.asarray(pre_mesh.vertices); f = np.asarray(pre_mesh.faces) + key = np.round(v, 6); uniq, inv = np.unique(key, axis=0, return_inverse=True) + f2 = inv[f] + fg = f2[(f2[:, 0] != f2[:, 1]) & (f2[:, 1] != f2[:, 2]) & (f2[:, 0] != f2[:, 2])] + e = np.sort(np.vstack([fg[:, [0, 1]], fg[:, [1, 2]], fg[:, [0, 2]]]), axis=1) + ue, cnt = np.unique(e, axis=0, return_counts=True) + open_verts = set(ue[cnt == 1].reshape(-1).tolist()) + tmask = np.array([any(x in open_verts for x in row) for row in f2]) + mp = yup(pre_mesh); mp.visual.face_colors = np.where(tmask[:, None], ORANGE, GRAY) + mp.export(out_prefix + "_remask_targets.ply") + # changed on post + pre_keys = set(_face_keys(pre_mesh)) + cmask = np.array([k not in pre_keys for k in _face_keys(post_mesh)]) + mq = yup(post_mesh); mq.visual.face_colors = np.where(cmask[:, None], RED, GRAY) + mq.export(out_prefix + "_remask_changed.ply") + print(f" [REMASK_VIS] {os.path.basename(out_prefix)}: 目标面 {int(tmask.sum())}/{len(tmask)}, 改变面 {int(cmask.sum())}/{len(cmask)}") + + +def quad_fill_mesh(mesh, res): + """长度3/4边界环确定性补洞(见 QUAD_FILL 注释)。返回(焊接后的)新 mesh。""" + v = np.asarray(mesh.vertices); f = np.asarray(mesh.faces) + key = np.round(v, 6); uniq, first_idx, inv = np.unique(key, axis=0, return_index=True, return_inverse=True) + f2 = inv[f]; f2 = f2[(f2[:, 0] != f2[:, 1]) & (f2[:, 1] != f2[:, 2]) & (f2[:, 0] != f2[:, 2])] # 图只用正常面; 输出保留原mesh全部面(含退化载体) + e = np.sort(np.vstack([f2[:, [0, 1]], f2[:, [1, 2]], f2[:, [0, 2]]]), axis=1) + ue, cnt = np.unique(e, axis=0, return_counts=True) + all_edges = set(map(tuple, ue)) + open_e = ue[cnt == 1] + from collections import defaultdict + adj = defaultdict(set) + for a, b in open_e: + adj[int(a)].add(int(b)); adj[int(b)].add(int(a)) + new_faces = []; seen = set(); step = 2.0 / res + for s in list(adj): + if s in seen: + continue + comp = [s]; seen.add(s); st = [s] + while st: + x = st.pop() + for y in adj[x]: + if y not in seen: + seen.add(y); comp.append(y); st.append(y) + n = len(comp); ne = sum(len(adj[x]) for x in comp) // 2 + if not (ne == n and all(len(adj[x]) == 2 for x in comp)): + continue # 非简单环不碰 + if n == 3: + new_faces.append(tuple(comp)); QUAD_FILL_STATS[0] += 1 + elif n == 4: + ring = [comp[0]]; prev = None + while len(ring) < 4: + nxt = [y for y in adj[ring[-1]] if y != prev and y in comp] + prev = ring[-1]; ring.append(nxt[0]) + d1 = (min(ring[0], ring[2]), max(ring[0], ring[2])) + d2 = (min(ring[1], ring[3]), max(ring[1], ring[3])) + if d1 in all_edges or d2 in all_edges: + continue # 已有对角线(半填充/薄结构合法边界)不碰 + if (uniq[comp].max(0) - uniq[comp].min(0)).max() > 2 * step + 1e-6: + continue # 跨区大环不碰(只修 raw face 内量级) + if np.linalg.norm(uniq[ring[0]] - uniq[ring[2]]) <= np.linalg.norm(uniq[ring[1]] - uniq[ring[3]]): + new_faces += [(ring[0], ring[1], ring[2]), (ring[0], ring[2], ring[3])] + else: + new_faces += [(ring[1], ring[2], ring[3]), (ring[1], ring[3], ring[0])] + QUAD_FILL_STATS[1] += 1 + if not new_faces: + return mesh + nf = first_idx[np.array(new_faces)] # 焊接索引 -> 原顶点索引; 原 mesh 零改动只追加 + return trimesh.Trimesh(vertices=v, faces=np.vstack([f, nf]), process=False) +DECIMATE_RATIO = float(os.environ.get("DECIMATE_RATIO", "0.5")) # GT 预塌缩比例(pyvista decimate 目标削减比, 0=不塌缩) +DEC_FILL = os.environ.get("DEC_FILL", "0") == "1" # 塌缩后先补洞(TRELLIS _fill_holes only_fill_holes=pymeshfix)再作输入(条件PC/天花板/粗起点); 指标参照仍用未补洞塌缩GT, 与基线可比 +# 级间是否按预测 mesh 自身 bbox 重新归一化(2026-08-22 修复, 默认关): 级间 mesh 是上一级解码产物, +# 天生在规范 [-1,1] 帧; 重归一化在预测带极值杂散面时把脚手架帧拉离 GT/点云帧 -> z 轴撕裂 +# (1(10)/1(101): res4 顶层 3-4% 杂散面 -> z中心+0.25/缩放4/3 -> 本体父层劈成 {-0.75,+0.25} 中间断层, +# res8 出现 -0.625 悬空板 + -0.375/-0.125 全空, 级联放大成面数腰斩)。=1 复现历史(≤08-22)行为。 +CASCADE_RENORM = os.environ.get("CASCADE_RENORM", "0") == "1" +# 每级采样方法(可多选逗号): plain / consensus(align) / log / log_fmm / geom / geom_fmm / fill。ALIGN_* 仍对 plain 生效(向后兼容)。 +METHODS = [m.strip() for m in os.environ.get("METHODS", "plain").split(",") if m.strip()] +MULTI_PC = os.environ.get("MULTI_PC", "1") == "1" # 0=单点云(pc_target=None, focal 单PC ckpt 必须) +SINGLE_TARGET_PC = os.environ.get("SINGLE_TARGET_PC", "0") == "1" # 1=云B作唯一云每层恒注入(sfocal 线口径); 需 MULTI_PC=1 以构建云B +NITER_GEOM = int(os.environ.get("NITER_GEOM", "3")) # geom: 每级终点 remask+requery 轮数 +KREQ_GEOM = int(os.environ.get("KREQ_GEOM", "5")) # geom: requery 每轮步数(5=甜点, 质量≈10 快~3x) +MAXFACE_REMASK = int(os.environ.get("MAXFACE", "1200")) # geom: 每轮最多重掩多少 face 的连接 +REMASK_UNION = os.environ.get("REMASK_UNION", "1") == "1" # geom: 开边两父 token 取并集(1-ring) +REMASK_LOOP_MAX = int(os.environ.get("REMASK_LOOP_MAX", "0")) # >0: 只对边界环长<=K的开边做 remask(外科修小洞; 0=全部开边) +REMASK_GROUPS = int(os.environ.get("REMASK_GROUPS", "1")) # >1: 每轮候选面按 BFS 序分 G 批逐批 requery(后批看得到前批修复, 用户方案); 1=整轮一次性 +RUN_SEED = int(os.environ.get("RUN_SEED", "0")) # >0: 整体偏移随机种子, 用于"同配置多次独立采样各自保存"(看运行间方差); 0=与历史逐位一致 +BO_N = int(os.environ.get("BO_N", "1")) # >1: 级内 best-of-N——重采条件云B+换torch种子造候选(实测±40开边方差来源), 结构惩罚(开边+10*NM%)选优 +BO_STAGES = set(int(x) for x in os.environ.get("BO_STAGES", "32").split(",") if x.strip()) # 只在这些产物格做 best-of-N +REQUERY_EMPTY_BIAS = float(os.environ.get("REQUERY_EMPTY_BIAS", "0")) # requery(缺陷区重问)期的空行偏置: 负值="不许再答无连接"推向补上; 0=关(现状) +GEOM_STAGE_PARAMS = {} # 按产物格覆盖 remask 轮数x步数, 如 "32:6x10"(res32 用 6x10, 其余用方法默认); 扫参结论: res8/16 甜点 3x5, res32 甜点 6x10 +for _kv in (x for x in os.environ.get("GEOM_STAGE_PARAMS", "").split(",") if x.strip()): + _r, _p = _kv.split(":"); _a, _b = _p.lower().split("x") + GEOM_STAGE_PARAMS[int(_r)] = (int(_a), int(_b)) +REMASK_DEGEN = os.environ.get("REMASK_DEGEN", "0") == "1" # geom: 额外纳入"有退化三角的 raw face"一起 remask +REMASK_NM = os.environ.get("REMASK_NM", "0") == "1" # geom: 非流形边(覆盖>2)的 1-ring face 也纳入 remask+requery。 +# 实测(val20@res32 A/B, ab_remask_nm/)净负: 开边131->171/非流形12.9->14.5%双变差, 调大MAXFACE无关(cap不触顶)、 +# 加迭代不救; 模型对 NM 区域重采仍给过连接答案, 反在原闭合处造新开边 -> 默认关, 仅留开关。 +GEOM_STAGES = set(int(x) for x in os.environ.get("GEOM_STAGES", "").split(",") if x.strip()) # geom: 只在这些"产物真实格"级做 remask(空=每级都做) +GT_START_RES = int(os.environ.get("GT_START_RES", "0")) # >0: 用 GT 编码->解码天花板@该真实分辨率作级联起点, 跳过更粗级(暴露偏差上界实验); 须在 STAGE 中 +EXT_START_DIR = os.environ.get("EXT_START_DIR", "") # 非空: 外部预测 mesh 作级联起点, 文件 {case_stem}{EXT_START_SUFFIX}, 保存帧(与 gt_original.obj 同帧), 只跑有该文件的 case +EXT_START_RES = int(os.environ.get("EXT_START_RES", "32")) # 外部起点的真实分辨率(须在 STAGE 中) +EXT_START_SUFFIX = os.environ.get("EXT_START_SUFFIX", "_pred.obj") +assert not (GT_START_RES and EXT_START_DIR), "GT_START_RES 与 EXT_START_DIR 互斥" +_START_RES = GT_START_RES or (EXT_START_RES if EXT_START_DIR else 0) +# 起点(GT天花板/外部预测)已是"完成的格 _START_RES mesh" -> 只跑更细的级(v2 语义: 不再重做起点那一级) +STAGE_RUN = [r for r in STAGE if r > _START_RES] if _START_RES else STAGE # 实际执行的阶梯(每项=完成该格) +OFF = np.stack([(np.arange(8) >> k) & 1 for k in range(3)], 1).astype(np.int64) # [8,3] +TRI = np.array([(i, j, k) for i in range(24) for j in range(i + 1, 24) for k in range(j + 1, 24)]) # C(24,3) +POS = {(0, 0): 3, (1, 1): 4, (2, 2): 5, (0, 1): 6, (0, 2): 7, (1, 2): 8} +# ---- SUPERRES: fix3 级联到 STAGE[-1] 后, 用 acc_uni 细分器继续细分到更高分辨率 ---- +SUPERRES = os.environ.get("SUPERRES", "0") == "1" # 是否开启更高分辨率细分 +SUPERRES_TO = int(os.environ.get("SUPERRES_TO", "512")) # 细分终点真实格(从 STAGE[-1] 倍增: 32->64->128->256->512) +# acc_uni 与 fix3 口径不同: acc_uni 深度 D 的真实格 = 2^D(fix3 是 2^(d+1)); 细分用 ACC_UN._depth_from_resolution(log2) +SUPERRES_STEPS = int(os.environ.get("SUPERRES_STEPS", "10")); SUPERRES_CFG = float(os.environ.get("SUPERRES_CFG", "1.0")) +SUPERRES_FIX_HOLES = os.environ.get("SUPERRES_FIX_HOLES", "1") == "1" +SR_FILL_IN = os.environ.get("SR_FILL_IN", "1") == "1" # 细分入口先 pymeshfix 补洞(对齐旧 cross512_fillall: 粗输入进细分器前补洞) +SUPERRES_CONFIG = os.environ.get("SUPERRES_CONFIG", "configs/subdivider.yaml") +SUPERRES_CKPT = os.environ.get("SUPERRES_CKPT", os.path.join(os.path.dirname(os.path.abspath(__file__)), "ckpt", "subdivider.pth")) +FINAL_RES = SUPERRES_TO if SUPERRES else STAGE[-1] +NSAMP = int(os.environ.get("NSAMP", "1024")); TAU = 2.0 / FINAL_RES # v2: FINAL_RES=512 真实格 -> TAU=2/512, 数值与历史协议一致; NSAMP=1024 严格对齐 metric_mul_dense_mesh_v2 + +BASE_DIR = os.environ.get("BASE_DIR", "pretrained_model/gobjaverse_fix3_octant_factorized_1_to_1024_ce_altpc_decoded_deg_w1") +_ckpt_cands = sorted(_glob.glob(os.path.join(BASE_DIR, "checkpoint_*.pth"))) +CKPT = os.environ.get("RAIL_CKPT") or (_ckpt_cands[-1] if _ckpt_cands else "") +if not CKPT or not os.path.exists(CKPT): + raise SystemExit(f"未找到生成器权重: 用 scripts/infer.py --ckpt 指定, 或设环境变量 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 diff --git a/mulres_infer/modules/subdivider_runtime.py b/mulres_infer/modules/subdivider_runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..6b1a315818d74989ce3399584b5cd69fa1a28199 --- /dev/null +++ b/mulres_infer/modules/subdivider_runtime.py @@ -0,0 +1,575 @@ +import gc +import glob +import math +import os +import random +import re +from pathlib import Path + +import numpy as np +import torch +import trimesh +from accelerate import Accelerator, DistributedDataParallelKwargs +from tqdm import tqdm + +from utils.config import load_config +from dataset.mesh_dataset import mul_encode_mesh, discrete_undiscrete_and_clean +from models.subdivider import DiscreteDiffusionTransformer, sample_refine_tokens_with_topology, sample_refine_tokens_with_topology_fm +from modules.octree_codec_acc import MeshCodec +from utils.postprocessing import _fill_holes +from utils.data_utils import normalize +def set_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ["PYTHONHASHSEED"] = str(seed) + +def _depth_from_resolution(resolution): + if resolution <= 0: + raise ValueError(f"resolution must be positive, got {resolution}") + depth = int(round(math.log2(resolution))) + if (1 << depth) != resolution: + raise ValueError(f"resolution must be power of 2, got {resolution}") + return depth + + +def restore_non_pad_token(token: torch.Tensor, shift: int, pad_id: int = 0): + if shift == 0: + return token + restored = token.clone() + non_pad_mask = restored != pad_id + restored[non_pad_mask] = restored[non_pad_mask] + shift + return restored + + +def read_dataset_paths(txt_path): + mesh_paths = [] + with open(txt_path, "r") as f: + for line in f: + path = line.strip() + if path: + mesh_paths.append(path) + return mesh_paths + + +def raw_stem_to_pair_key(raw_path): + stem = Path(raw_path).stem + # raw name example: 1(42)_y_f_4244 -> pair key: 1(42) + return re.sub(r"_y_f_\d+$", "", stem) + + +def gt_stem_to_pair_key(gt_path): + return Path(gt_path).stem + +def get_face_count(file_path): + """ + 使用正则表达式从文件名中提取面片数量 + 支持格式: *_f{num}_v*.obj 或 *_facenum_{num}_*.glb + """ + filename = Path(file_path).name + match = re.search(r'_f(\d+)_v|_facenum_(\d+)', filename) + if match: + # group(1) 对应 _f(\d+)_v 里的数字 + # group(2) 对应 _facenum_(\d+) 里的数字 + return int(match.group(1) or match.group(2)) + return None + +# def build_raw_gt_pairs(raw_txt_path, gt_txt_path): +# raw_paths = read_dataset_paths(raw_txt_path) +# gt_paths = read_dataset_paths(gt_txt_path) + +# # If two txt paths are the same source, pair each mesh with itself directly. +# if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve(): +# return [(p, p) for p in raw_paths], [] + +# gt_by_key = {} +# for gt_path in gt_paths: +# gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path + +# pairs = [] +# missing_gt = [] +# for raw_path in raw_paths: +# key = raw_stem_to_pair_key(raw_path) +# gt_path = gt_by_key.get(key) +# if gt_path is None: +# missing_gt.append(raw_path) +# continue +# pairs.append((raw_path, gt_path)) + +# return pairs, missing_gt +def build_raw_gt_pairs(raw_txt_path, gt_txt_path, min_faces=None, max_faces=None): + raw_paths = read_dataset_paths(raw_txt_path) + gt_paths = read_dataset_paths(gt_txt_path) + + # ===== 新增过滤逻辑 ===== + filtered_raw_paths = [] + for p in raw_paths: + if min_faces is None and max_faces is None: + filtered_raw_paths.append(p) + continue + + faces = get_face_count(p) + if faces is not None: + # 检查是否在指定范围内 + if min_faces is not None and faces < min_faces: + continue + if max_faces is not None and faces > max_faces: + continue + filtered_raw_paths.append(p) + else: + # 如果文件名中没找到面片数特征，默认保留（你也可以改为 continue 丢弃它） + filtered_raw_paths.append(p) + + raw_paths = filtered_raw_paths + # ======================== + + # If two txt paths are the same source, pair each mesh with itself directly. + if Path(raw_txt_path).resolve() == Path(gt_txt_path).resolve(): + return [(p, p) for p in raw_paths], [] + + gt_by_key = {} + for gt_path in gt_paths: + gt_by_key[gt_stem_to_pair_key(gt_path)] = gt_path + + pairs = [] + missing_gt = [] + for raw_path in raw_paths: + key = raw_stem_to_pair_key(raw_path) + gt_path = gt_by_key.get(key) + if gt_path is None: + missing_gt.append(raw_path) + continue + pairs.append((raw_path, gt_path)) + + return pairs, missing_gt + +def save_gt_mesh(mesh, out_dir, name): + gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False) + gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]] + gt_mesh.export(os.path.join(out_dir, f"{name}_gt.obj")) + +def save_raw_gt_mesh(mesh, out_dir, name): + gt_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False) + gt_mesh.vertices = gt_mesh.vertices[:, [1, 2, 0]] + gt_mesh.export(os.path.join(out_dir, f"{name}_raw_gt.obj")) + + +def save_stage_input_mesh(mesh, out_dir, name, stage_in_res, stage_out_res, stage_in_depth): + coarse_mesh = trimesh.Trimesh(vertices=mesh.vertices.copy(), faces=mesh.faces.copy(), process=False) + coarse_mesh.vertices = coarse_mesh.vertices[:, [1, 2, 0]] + suffix = f"_stage_{stage_in_res}to{stage_out_res}_input_coarse" + coarse_mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{stage_in_depth}.obj")) + + +def save_stage_prediction(codec, raw_token, refine_token, out_dir, name, target_depth, suffix, fix_holes, raw_dis_mesh=None, gt_dis_mesh=None): + eos_mask = raw_token[0] == 0 + if eos_mask.any(): + raw_tokens = raw_token[0, : int(torch.where(eos_mask)[0][0])].detach().cpu().long().tolist() + else: + raw_tokens = raw_token[0].detach().cpu().long().tolist() + + eos_mask_refine = refine_token[0] == 0 + if eos_mask_refine.any(): + refine_tokens = refine_token[0, : int(torch.where(eos_mask_refine)[0][0])].detach().cpu().long().tolist() + else: + refine_tokens = refine_token[0].detach().cpu().long().tolist() + + shared_len = min(len(raw_tokens), len(refine_tokens)) + block = (shared_len // 9) * 9 + if block <= 0: + return None + + raw_tokens = raw_tokens[:block] + refine_tokens = refine_tokens[:block] + + interleaved = [] + for i in range(0, block, 9): + interleaved.extend(raw_tokens[i : i + 9]) + interleaved.extend(refine_tokens[i : i + 9]) + + mesh = codec.decode_tokens(interleaved, target_depth=target_depth-1) + if mesh is None: + return None + + trimesh.repair.fix_normals(mesh) + if fix_holes: + verts = torch.tensor(mesh.vertices, device="cuda", dtype=torch.float32) + faces = torch.tensor(mesh.faces, device="cuda", dtype=torch.int32) + new_verts, new_faces = _fill_holes(verts, faces, debug=False, verbose=False, only_fill_holes=True) + mesh = trimesh.Trimesh(vertices=new_verts.cpu().numpy(), faces=new_faces.cpu().numpy()) + + return_mesh = mesh.copy() + mesh.vertices = mesh.vertices[:, [1, 2, 0]] + mesh.export(os.path.join(out_dir, f"{name}{suffix}_depth{target_depth}.obj")) + + if raw_dis_mesh is not None: + raw_dis_mesh.vertices = raw_dis_mesh.vertices[:, [1, 2, 0]] + raw_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_raw_dis_depth{target_depth-1}.obj")) + if gt_dis_mesh is not None: + gt_dis_mesh.vertices = gt_dis_mesh.vertices[:, [1, 2, 0]] + gt_dis_mesh.export(os.path.join(out_dir, f"{name}{suffix}_gt_dis_depth{target_depth}.obj")) + return return_mesh + + +def process_one_mesh( + raw_mesh_path, + gt_mesh_path, + model, + mask_token_id, + codec, + device, + pad_id, + token_shift, + config, + output_folder, + stage_resolutions, + fix_holes=True, + use_gt_pc=True +): + import time + t_start = time.time() + import time + t_total_start = time.time() + + name = Path(raw_mesh_path).stem + t_start = time.time() + mesh = trimesh.load(raw_mesh_path, force="mesh") + mesh.vertices = mesh.vertices[:, [2, 0, 1]] + gt_mesh = trimesh.load(gt_mesh_path, force="mesh") + gt_mesh.vertices = gt_mesh.vertices[:, [2, 0, 1]] + normalize(gt_mesh) + normalize(mesh) + save_raw_gt_mesh(mesh, output_folder, name) + save_gt_mesh(gt_mesh, output_folder, name) + t_load = time.time() + print(f"[{name}] Load & normalize time: {t_load - t_start:.4f}s") + + stage_depths = [_depth_from_resolution(x) for x in stage_resolutions] + + # First-stage input uses codec encode-decode result at the minimum resolution. + min_stage_depth = stage_depths[0] + t_start = time.time() + init_raw_token, _, init_quantized_mesh = codec.encode_mesh(mesh, target_depth=min_stage_depth) + init_mesh = codec.decode_raw_tokens(init_raw_token, target_depth=min_stage_depth) + t_init = time.time() + print(f"[{name}] Init encode & decode (depth={min_stage_depth}): {t_init - t_start:.4f}s") + + if init_mesh is not None: + current_mesh = init_mesh + elif init_quantized_mesh is not None: + current_mesh = init_quantized_mesh + else: + current_mesh = mesh.copy() + + for stage_idx in range(len(stage_depths) - 1): + stage_start_time = time.time() + stage_in_depth = stage_depths[stage_idx] + stage_out_depth = stage_depths[stage_idx + 1] + stage_in_res = stage_resolutions[stage_idx] + stage_out_res = stage_resolutions[stage_idx + 1] + + save_stage_input_mesh( + mesh=current_mesh, + out_dir=output_folder, + name=name, + stage_in_res=stage_in_res, + stage_out_res=stage_out_res, + stage_in_depth=stage_in_depth, + ) + + t_start = time.time() + stage_data = mul_encode_mesh( + raw_mesh=current_mesh, + gt_mesh=gt_mesh, + codec=codec, + tokenizer_target_depth=stage_in_depth, + pc_num=config.data_processing.pc_num, + total_pc_num=config.data_processing.total_pc_num, + use_gt_pc=use_gt_pc, + clean=False + ) + t_mul_encode = time.time() + print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | mul_encode_mesh time: {t_mul_encode - t_start:.4f}s") + + raw_token = stage_data["raw_token"].unsqueeze(0).to(device) + pc = stage_data["pc"].unsqueeze(0).to(device) + + t_start = time.time() + # refine_token = sample_refine_tokens_with_topology( + # model, + # mask_token_id, + # raw_token, + # pc, + # use_gumbel=False, + # gumbel_temp_init=1.0, + # use_soft_consensus=False, + # consensus_start_step_ratio=0, + # cfg_scale=1.0, + # T_infer=config.evaluation.diffusion_sample_steps, + # ) + refine_token = sample_refine_tokens_with_topology_fm( + model, + mask_token_id, + raw_token, + pc, + use_gumbel=False, + gumbel_temp_init=1.0, + use_soft_consensus=False, + consensus_start_step_ratio=0, + cfg_scale=1.0, + T_infer=config.evaluation.diffusion_sample_steps, + ) + + torch.cuda.synchronize() + t_sample = time.time() + print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | sample_refine_tokens_with_topology time: {t_sample - t_start:.4f}s") + + refine_token = restore_non_pad_token(refine_token, shift=token_shift, pad_id=pad_id) + + t_start = time.time() + stage_suffix = f"_stage_{stage_in_res}to{stage_out_res}" + stage_mesh = save_stage_prediction( + codec=codec, + raw_token=raw_token, + refine_token=refine_token, + out_dir=output_folder, + name=name, + target_depth=stage_out_depth, + suffix=stage_suffix, + fix_holes=fix_holes, + raw_dis_mesh=stage_data.get("raw_dis_mesh"), + gt_dis_mesh=stage_data.get("gt_dis_mesh"), + ) + t_save_pred = time.time() + print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | save_stage_prediction (decode + fix) time: {t_save_pred - t_start:.4f}s") + print(f"[{name}] Stage {stage_in_res}->{stage_out_res} | total stage time: {time.time() - stage_start_time:.4f}s\n") + + if stage_mesh is None: + break + + current_mesh = stage_mesh + + print(f"[{name}] Total process_one_mesh time: {time.time() - t_total_start:.4f}s") + + +class EMA: + def __init__(self, model, decay=0.999): + self.decay = decay + self.shadow = {} + self.register(model) + + def register(self, model): + self.shadow = {} + for name, param in model.named_parameters(): + if param.requires_grad: + self.shadow[name] = param.data.clone() + + def load_state_dict(self, state): + self.decay = state.get("decay", self.decay) + self.shadow = {k: v.clone() for k, v in state["shadow"].items()} + + @torch.no_grad() + def copy_to(self, model): + """Swap shadow weights into the model in-place (for inference).""" + copied, missing = 0, 0 + for name, param in model.named_parameters(): + if name in self.shadow: + param.data.copy_(self.shadow[name].to(device=param.device, dtype=param.dtype)) + copied += 1 + else: + missing += 1 + print(f"[EMA.copy_to] copied {copied} params into model ({missing} not present in shadow)") + + +def run_inference_on_all_checkpoints(model, mask_token_id, accelerator, config, device, pad_id, codec, token_shift, ema=None): + ckpt_dir = "ckpt" + output_folder = "./pretrain_model_mul_more_data_64_to_256_eval_complex_test_no_fix_holes_no_clean_cfg_1.0_fm" + + # raw_data_txt = "dataset/bpt_res_list.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_retrain_gt/data.txt + # gt_data_txt = "dataset/trellis_clean_mesh.txt" /root/junkai_lin/mul_res_diffusion/demo/bpt_gt/data.txt + raw_data_txt = "demo/complex_test_data/files/data.txt" + gt_data_txt = "demo/complex_test_data/files/data.txt" + # raw_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt" + # gt_data_txt = "/root/junkai_lin/mul_res_diffusion/demo/hun_yuan/data.txt" + stage_resolutions = [64, 128, 256] + # raw_data_txt = "demo/tripo 1.0/data.txt" + # gt_data_txt = "demo/tripo 1.0/data.txt" + + use_gt_pc=True + fix_holes=False + if accelerator.is_main_process: + print(f"Using raw data txt: {raw_data_txt}") + print(f"Using gt data txt: {gt_data_txt}") + print(f"Searching checkpoints in: {ckpt_dir}") + + # mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=config.data_processing.max_face_num, min_faces=config.data_processing.min_face_num) + mesh_pairs, missing_gt = build_raw_gt_pairs(raw_data_txt, gt_data_txt, max_faces=20000, min_faces=500) + if accelerator.is_main_process: + print(f"Matched raw/gt pairs: {len(mesh_pairs)}") + if missing_gt: + print(f"Raw without matched gt: {len(missing_gt)}") + + search_pattern = os.path.join(ckpt_dir, "checkpoint_*_epoch_*_loss.pth") + ckpt_files = glob.glob(search_pattern) + if not ckpt_files: + if accelerator.is_main_process: + print("No checkpoints found!") + return + + ckpt_data = [] + for f in ckpt_files: + match = re.search(r"checkpoint_(\d+)_epoch_", f) + if match: + epoch = int(match.group(1)) + ckpt_data.append((epoch, f)) + ckpt_data.sort(key=lambda x: x[0]) + + # stage_resolutions = [32, 64, 128, 256] + + + # stage_resolutions = [128, 256] + + for epoch, ckpt_path in ckpt_data: + if accelerator.is_main_process: + print(f"\n[{'='*40}]") + print(f"Starting inference for epoch {epoch}") + print(f"Loading weights from: {ckpt_path}") + + output_folder_ply = os.path.join(output_folder, f"{epoch}_epoch_val_generate_mul") + if accelerator.is_main_process: + os.makedirs(output_folder_ply, exist_ok=True) + + checkpoint = torch.load(ckpt_path, map_location=device) + model_to_load = accelerator.unwrap_model(model) + # Use strict=False so we can load older ckpts (no null_cond) into the + # newer model class, and vice versa. Missing/unexpected keys are logged. + load_result = model_to_load.load_state_dict(checkpoint["model_state_dict"], strict=False) + missing = list(getattr(load_result, "missing_keys", []) or []) + unexpected = list(getattr(load_result, "unexpected_keys", []) or []) + if (missing or unexpected) and accelerator.is_main_process: + print(f"[load_state_dict] strict=False -> missing={len(missing)} unexpected={len(unexpected)}") + if missing: + print(f" missing: {missing[:8]}{' ...' if len(missing) > 8 else ''}") + if unexpected: + print(f" unexpected: {unexpected[:8]}{' ...' if len(unexpected) > 8 else ''}") + + # Auto-detect EMA: if the ckpt has an `ema_state_dict`, swap those + # shadow weights into the model for inference. Otherwise fall back to + # the raw `model_state_dict` we already loaded above. No reliance on + # config `use_ema` / a separate EMA object — just check the ckpt. + ema_state = checkpoint.get("ema_state_dict") + shadow = ema_state.get("shadow") if isinstance(ema_state, dict) else None + if shadow: + with torch.no_grad(): + copied, skipped = 0, 0 + for name, param in model_to_load.named_parameters(): + if name in shadow: + param.data.copy_(shadow[name].to(device=param.device, dtype=param.dtype)) + copied += 1 + else: + skipped += 1 + if accelerator.is_main_process: + msg = f"[Epoch {epoch}] EMA shadow detected in ckpt -> applied to model ({copied} params" + if skipped: + msg += f", {skipped} model params not in shadow" + print(msg + ")") + else: + if accelerator.is_main_process: + print(f"[Epoch {epoch}] no EMA in ckpt -> running with raw model_state_dict weights.") + + with torch.no_grad(): + model.eval() + pbar = tqdm( + enumerate(mesh_pairs), + total=len(mesh_pairs), + desc=f"Infer Epoch {epoch}", + disable=not accelerator.is_local_main_process, + dynamic_ncols=True, + ) + for line_idx, (raw_mesh_path, gt_mesh_path) in pbar: + if line_idx % accelerator.num_processes != accelerator.process_index: + continue + try: + process_one_mesh( + raw_mesh_path=raw_mesh_path, + gt_mesh_path=gt_mesh_path, + model=model, + mask_token_id=mask_token_id, + codec=codec, + device=device, + pad_id=pad_id, + token_shift=token_shift, + config=config, + output_folder=output_folder_ply, + stage_resolutions=stage_resolutions, + fix_holes=fix_holes, + use_gt_pc=use_gt_pc, + ) + except Exception as e: + print(f"[infer error] raw={raw_mesh_path}, gt={gt_mesh_path}: {str(e)}") + + torch.cuda.empty_cache() + gc.collect() + + if accelerator.is_main_process: + print("\nBatch inference complete for all checkpoints.") + + +def main(config): + print("start init Accelerator...") + ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator( + gradient_accumulation_steps=config.accelerator.gradient_accumulation_steps, + mixed_precision=config.accelerator.mixed_precision, + kwargs_handlers=[ddp_kwargs], + ) + device = accelerator.device + + codec = MeshCodec(args=config.tokenizer) + codec_config = codec.get_model_config() + input_num_categories = codec_config["input_num_categories"] + lm_num_categories = codec_config["output_num_categories"] + pad_id = codec_config["pad_token_id"] + token_shift = input_num_categories - 1 + + model = DiscreteDiffusionTransformer( + vocab_A=input_num_categories, + vocab_B=lm_num_categories, + max_seq_len=config.data_processing.max_len, + d_model=config.model.feature_dim, + encoder_freeze=config.model.encoder_freeze, + nhead=config.model.num_heads, + depth=config.model.hourglass_depth, + dropout=config.model.dropout, + time_embed_scale=config.training.diffusion_train_steps, + ) + mask_token_id = model.mask_token_id + + use_ema = getattr(config.training, "use_ema", False) + ema_decay = getattr(config.training, "ema_decay", 0.999) + ema = EMA(model, decay=ema_decay) if use_ema else None + + model = accelerator.prepare(model) + + run_inference_on_all_checkpoints( + model=model, + mask_token_id=mask_token_id, + accelerator=accelerator, + config=config, + device=device, + pad_id=pad_id, + codec=codec, + token_shift=token_shift, + ema=ema, + ) + + +if __name__ == "__main__": + set_seed(42) + config = load_config() + main(config) diff --git a/mulres_infer/modules/token_utils.py b/mulres_infer/modules/token_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7da583e88aa59957b0591fb9e79c7a5f00168967 --- /dev/null +++ b/mulres_infer/modules/token_utils.py @@ -0,0 +1,52 @@ +"""fix3 编码格式的数据通路工具 —— 从训练脚本抽出, 保证推理与训练用同一份实现。 + + split_fix3_token_list : 扁平 60/face token -> (coord_list[9*nface], subcode_list[9*nface,8] local) + subcodes_to_flat_token : (coord, 采样得到的 subcodes) -> fix3 扁平 60/face token(含 offset), 供 decode_tokens +""" +import torch + +from models.generator import REFINE_POS_PER_FACE, NUM_OCTANTS + +TOKENS_PER_FACE = 60 # fix3 扁平: 9 coord + 3 F + 6 矩阵 x 8 行 + +def _trim_trailing_pad(token, pad_id): + pad_pos = (token == pad_id).nonzero(as_tuple=True)[0] + return token[:pad_pos[0].item()] if pad_pos.numel() > 0 else token +def split_fix3_token_list(token_list, pad_id, device, codec_cfg): + """扁平 60/face token -> (coord_list, subcode_list)。 + coord_list[i] : [9*nface] 坐标(保留 coord_offset, 直接喂 raw_embedding) + subcode_list[i] : [9*nface, 8] refine 子码(去 offset, local[0,256)) + 位置顺序 [F0 F1 F2 M00 M11 M22 M01 M02 M12] + F 位置只用 slot0(其余槽填 0, 模型按类型只读 slot0)；intra/inter 用 8 行 + """ + f_off = codec_cfg["f_offset"]; ia = codec_cfg["intra_row_offset"]; ie = codec_cfg["inter_row_offset"] + coord_list, sub_list = [], [] + for tok in token_list: + tok = _trim_trailing_pad(tok.to(device=device, dtype=torch.long), pad_id) + if tok.numel() % TOKENS_PER_FACE != 0: + raise ValueError(f"token 长度 {tok.numel()} 不是 {TOKENS_PER_FACE} 的倍数") + tok = tok.view(-1, TOKENS_PER_FACE) # [nface, 60] + nface = tok.size(0) + coord = tok[:, :9].reshape(-1) # [9*nface] + Fv = tok[:, 9:12] - f_off # [nface, 3] + intra = (tok[:, 12:36] - ia).view(nface, 3, NUM_OCTANTS) # [nface,3,8] + inter = (tok[:, 36:60] - ie).view(nface, 3, NUM_OCTANTS) # [nface,3,8] + sub = torch.zeros(nface, REFINE_POS_PER_FACE, NUM_OCTANTS, dtype=torch.long, device=device) + sub[:, 0:3, 0] = Fv # F -> slot0 + sub[:, 3:6, :] = intra + sub[:, 6:9, :] = inter + sub_list.append(sub.reshape(-1, NUM_OCTANTS)) # [9*nface, 8] + coord_list.append(coord) + return coord_list, sub_list +def subcodes_to_flat_token(coord_seq, sub_seq, codec_cfg): + """采样得到的 (coord[9*nface], subcodes[9*nface,8] local) 还原成 fix3 扁平 60/face token(含 offset)， + 供 codec.decode_tokens 使用。F 位置取 slot0。""" + f_off = codec_cfg["f_offset"]; ia = codec_cfg["intra_row_offset"]; ie = codec_cfg["inter_row_offset"] + nface = coord_seq.numel() // REFINE_POS_PER_FACE + coord = coord_seq.view(nface, REFINE_POS_PER_FACE) # [nface,9] + sub = sub_seq.view(nface, REFINE_POS_PER_FACE, NUM_OCTANTS) # [nface,9,8] + Fv = sub[:, 0:3, 0] + f_off # [nface,3] 取 slot0 + intra = sub[:, 3:6, :] + ia # [nface,3,8] + inter = sub[:, 6:9, :] + ie + flat = torch.cat([coord, Fv, intra.reshape(nface, 24), inter.reshape(nface, 24)], dim=1) # [nface,60] + return flat.reshape(-1).tolist() diff --git a/mulres_infer/readme.md b/mulres_infer/readme.md new file mode 100644 index 0000000000000000000000000000000000000000..e458e295db72933c16d3e1419235ee79b1231aae --- /dev/null +++ b/mulres_infer/readme.md @@ -0,0 +1,122 @@ +# 多分辨率网格生成 · 推理 + +点云条件下的**多分辨率八叉树网格生成**：低分辨率生成器逐级长出粗网格，高分辨率细分器把它推到 512 格。 + +## 流程 + +``` +输入 mesh ──采样点云──┐ + ▼ + ┌───────── 生成器 (TSSR 三前向离散扩散) ─────────┐ + │ 格2 → 格4 → 格8 → 格16 → 格32 → 格64 每级20步 │ + │ └补洞┘ └补洞┘ │ + └──────────────────────┬─────────────────────────┘ + ▼ + ┌────────── 细分器 (acc_uni) ──────────┐ + │ 格64 → 格128 → 格256 → 格512 每级20步 │ + └──────────────────────────────────────┘ + ▼ + final_res512_plain.obj +``` + +**生成器**每级把上一级网格重新编码成八叉树 token，用离散扩散预测该级完整子码。 +TSSR 解码每步三次前向：mask 预测 → no-mask 精修 → 逐 token 分类器挑可疑位置 re-mask， +避免 MaskGIT 式"一旦揭示即定死"的错误累积。 + +**细分器**只做细分不做生成，粗输入上的破洞它补不出来——这是 res32/res64 要先补洞再喂下去的原因。 + +## 目录 + +``` +. +├── scripts/ 入口 +│ ├── infer.py 端到端推理 +│ ├── eval.py 评测(CD / HD / 法向一致性) +│ └── strip_ckpt.py 权重瘦身: 剥离优化器、合并 EMA +├── models/ 模型 +│ ├── generator.py 生成器(骨干 → 双云注入 → TSSR) +│ ├── subdivider.py 细分器 +│ └── point_conditioner.py 点云条件编码器 +├── modules/ 算子与管线 +│ ├── pipeline.py 多级级联主管线 +│ ├── attention.py MLA 等注意力原语 +│ ├── octree_codec.py 八叉树 tokenizer(numba 加速) +│ ├── octree_codec_acc.py 细分器口径的 tokenizer +│ ├── octree_decode.py 子码 → mesh 解码 +│ ├── token_utils.py token 通路 +│ └── miche/ Michelangelo 点云编码器(第三方) +├── dataset/ 数据处理与点云采样 +├── utils/ 配置、网格工具、补洞、指标 +├── configs/ generator.yaml / subdivider.yaml +└── ckpt/ 权重(见 ckpt/README.md) +``` + +## 安装 + +```bash +conda create -n mulres python=3.12 -y && conda activate mulres +pip install -r requirements.txt +# flash-attn 用与 torch/CUDA 匹配的预编译 wheel, 不要从源码构建 +``` + +## 使用 + +```bash +# 全部 mesh +python scripts/infer.py --data /path/to/meshes --out_dir outputs/run1 + +# 只跑前 20 个 +python scripts/infer.py --data /path/to/meshes --out_dir outputs/run1 --ncase 20 + +# 指定 case(文件名子串, 带尾点防前缀误配: "a." 不匹配 "a1.glb") +python scripts/infer.py --data /path/to/meshes --out_dir outputs/run1 --case "model_a.,model_b." + +# 面数太密: 给连接头"空连接"类加 logit 偏置(0.3 约 -25% 面, 保真度几乎不变) +python scripts/infer.py --data /path/to/meshes --out_dir outputs/run1 --empty_bias 0.3 +``` + +每个 case 输出一个子目录： + +| 文件 | 说明 | +|---|---| +| `final_res512_plain.obj` | **最终结果** | +| `res{2,4,8,16,32,64}_plain.obj` | 生成器各级产物(补洞后) | +| `*_sr_64to128_depth7.obj` 等 | 细分器各级产物 | +| `gt_original.obj` | 归一化后的原始输入(塌缩前) | +| `gt_original_dec.obj` | 塌缩后的输入(级联起点与条件点云的来源) | +| `gt_res{N}.obj` | 输入经 tokenizer 编解码的天花板, 用于逐级对照 | +| `pc_full.ply` / `pc_target_res{N}.ply` | 注入的条件点云 | + +## 评测 + +```bash +# 预测目录需扁平: _y.obj ; GT 目录: .glb 或 .obj +python scripts/eval.py /path/to/pred_flat --gt /path/to/gt_meshes +``` + +输出倒角距离(L1/L2)、豪斯多夫距离、法向一致性 |NC| 与 NC。 + +> 解码不编码面的绕序，有向法向随机，所以看 **|NC|**（约 0.85），`NC (Real)` 接近 0 属预期。 + +## 参数 + +`scripts/infer.py` 固化了默认配方，更细的控制用环境变量（定义见 `modules/pipeline.py` 顶部）： + +| 变量 | 默认 | 说明 | +|---|---|---| +| `STAGE` | `2,4,8,16,32,64` | 生成器级联阶梯(产物真实体素格) | +| `STEPS` | `20` | 生成器每级扩散步数 | +| `SUPERRES_TO` / `SUPERRES_STEPS` | `512` / `20` | 细分终点与每级步数 | +| `FILL` / `FILL_MIN_RES` | `1` / `0` | 级联内补洞; `FILL_MIN_RES=64` 则只在格 64 补 | +| `SR_FILL_IN` / `SUPERRES_FIX_HOLES` | `1` / `1` | 细分入口 / 每级输出补洞 | +| `DECIMATE_RATIO` | `0.5` | 输入预塌缩比(0=不塌缩) | +| `EMPTY_BIAS` | `0` | >0 稀疏化控面数 | +| `CASE_STEM` | — | 文件名子串过滤(在 `NCASE` 截断**之后**生效) | + +## 已知行为 + +- **步数**：20 与 40 步保真度无差别；10 步会让开边 +27%、非流形 +34%。20 是性价比点。 +- **补洞是置换不是净赚**：开边 −16% 的代价是非流形 +12%。要纯净拓扑用 `--no_fill`。 +- **分辨率终点**：生成器自己走到格 64 比走到格 32 再交给细分器更好（非流形 −59%、CD −6%）， + 代价是开边变多——细分器会把粗输入的破洞一路放大。 +- **GPU 非确定性**：flash-attn / bf16 的归约顺序不固定，同配置重跑面数有约 3% 波动，单 case 开边波动更大。 diff --git a/mulres_infer/requirements.txt b/mulres_infer/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..d47908067842766250f5293d9153778cf6161934 --- /dev/null +++ b/mulres_infer/requirements.txt @@ -0,0 +1,22 @@ +# 核心 +torch==2.7.0 +numpy +scipy +trimesh==5.1.0 +numba +einops +omegaconf +transformers +accelerate +timm +scikit-image +pymeshfix +pyvista +tqdm + +# flash-attn 需与 torch/CUDA 匹配, 用官方预编译 wheel 安装, 不要从源码构建: +# pip install flash_attn-2.8.3+cu12torch2.7cxx11abiFALSE-cp312-cp312-linux_x86_64.whl + +# 可选: utils/postprocessing_utils.py 顶部导入, 但补洞路径实际不调用其功能 +# utils3d @ git+https://github.com/EasternJournalist/utils3d.git +# nvdiffrast @ git+https://github.com/NVlabs/nvdiffrast.git diff --git a/mulres_infer/scripts/__init__.py b/mulres_infer/scripts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/mulres_infer/scripts/eval.py b/mulres_infer/scripts/eval.py new file mode 100644 index 0000000000000000000000000000000000000000..82bbf3b29ddbcb26a8b331c630ec388f4403a383 --- /dev/null +++ b/mulres_infer/scripts/eval.py @@ -0,0 +1,429 @@ + +import os, sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from utils.distance import calculate_distances_numpy, calculate_normal_consistency_score +import os +import numpy as np +from tqdm import tqdm +import trimesh +import json +import glob +import multiprocessing # 导入多进程模块 +MIN_FACE_AMOUNT = 0 +MAX_FACE_AMOUNT = 100000 +# =============================================================== +# 第一步：你的距离计算函数 (无变化) +# =============================================================== +pc_num = 1024 + +def normalize(mesh, file_path): + # ... 现有代码 ... + bounds = np.array([mesh.vertices.min(axis=0), mesh.vertices.max(axis=0)]) + mesh.vertices = mesh.vertices - (bounds[0] + bounds[1])[None, :] / 2 + mesh.vertices = mesh.vertices / (bounds[1] - bounds[0]).max() + mesh.vertices = mesh.vertices.clip(-0.5, 0.5) # 限制顶点坐标在[-0.5, 0.5] + +# =============================================================== +# 第二步：读取 .obj 文件的函数 (无变化) +# =============================================================== +def sample_from_obj(file_path, alter=False): + # ... 现有代码 ... + """ + 从 .obj 文件中加载顶点坐标。 + """ + mesh = trimesh.load(file_path, force="mesh") + if alter: + mesh.vertices = mesh.vertices[:, [1,2,0]] + normalize(mesh, file_path) + points, face_idx = mesh.sample(pc_num, return_index=True) + normals = mesh.face_normals[face_idx] + return points,normals + +# =============================================================== +# [修改] 法向量一致性计算 +# =============================================================== + +# =============================================================== +# 第三步：[修改] 多进程工作函数 +# =============================================================== + +def process_common_file_worker(base_name, gt_folder, folder_paths): + """ + [已修改] 这是一个工作进程执行的任务： + 处理单个 base_name 对应的所有文件夹。 + 它会加载一次GT文件，然后遍历所有预测文件夹进行计算。 + 返回 L1, L2 CD 距离, |NC| 和 NC。 + """ + + # 1. 查找并加载 GT 文件 + gt_filename = base_name + '.glb' + # ... 现有代码 ... + gt_filepath = os.path.join(gt_folder, gt_filename) + + if not os.path.exists(gt_filepath): + gt_filename = base_name + '.obj' + # ... 现有代码 ... + gt_filepath = os.path.join(gt_folder, gt_filename) + if not os.path.exists(gt_filepath): + # 警告：在工作进程中打印 (print) 不是线程安全的， + # ... 现有代码 ... + # 但返回空列表是安全的，主进程会处理它。 + # print(f"\n警告：共同文件 '{base_name}' 对应的GT文件 '{gt_filepath}' 不存在，已跳过。") + return [] # 返回空列表表示此文件处理失败 + mesh = trimesh.load(gt_filepath, force="mesh") + if len(mesh.faces) < MIN_FACE_AMOUNT or len(mesh.faces) > MAX_FACE_AMOUNT: + return [] + + try: + pcd_gt, normal_gt = sample_from_obj(gt_filepath) + if pcd_gt is None or pcd_gt.size == 0: + # ... 现有代码 ... + # print(f"\n警告：加载GT文件 {gt_filename} 失败或点云为空。已跳过。") + return [] + except Exception as e: + # ... 现有代码 ... + # print(f"\n加载GT文件 {gt_filename} 时发生错误: {e}。已跳过。") + return [] + + # 2. 遍历所有文件夹，加载 Pred 文件并计算 + results_for_this_file = [] + + for folder_path in folder_paths: + pred_filename = base_name + '_y.obj' + # ... 现有代码 ... + pred_filepath = os.path.join(folder_path, pred_filename) + + try: + pcd_pred, normal_pred = sample_from_obj(pred_filepath) + if pcd_pred is None or pcd_pred.size == 0: + # ... 现有代码 ... + # print(f"\n警告：加载文件 {pred_filepath} 失败或点云为空。已跳过。") + results_for_this_file.append((folder_path, None, None, None, None, None)) # [修改] 6 个 None + continue + + # [修改] 接收 L1 和 L2 CD + chamfer_l2, hausdorff, chamfer_l1 = calculate_distances_numpy(pcd_gt, pcd_pred) + # [修改] 接收两个 NC 分数 + normal_error_abs, normal_error_real = calculate_normal_consistency_score(pcd_gt, normal_gt, pcd_pred, normal_pred) + + # [修改] 存储有效结果 (6个值) + results_for_this_file.append((folder_path, chamfer_l2, hausdorff, normal_error_abs, chamfer_l1, normal_error_real)) + + except Exception as e: + # ... 现有代码 ... + # print(f"\n计算文件对 {pred_filepath} 时发生错误: {e}。已跳过。") + results_for_this_file.append((folder_path, None, None, None, None, None)) # [修改] 6 个 None + + # 3. 返回此 base_name 对应的所有结果 + return results_for_this_file + +def starmap_helper(args): + # ... 现有代码 ... + """ + [新增] 辅助函数，用于解包 pool.imap 的参数，使其能接受多个参数。 + """ + return process_common_file_worker(*args) + +# =============================================================== +# 第四步：[修改] 评估函数，使用多进程 +# =============================================================== + +def evaluate_common_files(folder_paths, gt_folder, base_dir): + """ + [已修改] 评估多个文件夹中 *共有* 的文件，并计算总平均指标和逐文件夹的平均指标。 + 现在使用 multiprocessing.Pool 来并行处理文件。 + """ + if not folder_paths: + # ... 现有代码 ... + print("错误：没有提供用于评估共同文件的文件夹。") + return None + + print("\n" + "="*80) + # ... 现有代码 ... + print(" 正在计算共同文件指标 ".center(80, "=")) + print(f"总共 {len(folder_paths)} 个文件夹参与比较。") + + # 1. 查找每个文件夹中的文件基础名（不含 '_y.obj'）(无变化) + files_by_folder = {} + # ... 现有代码 ... + for folder_path in folder_paths: + try: + all_files = os.listdir(folder_path) + # ... 现有代码 ... + base_names = set([f.replace('_y.obj', '') for f in all_files + if (f.endswith('_y.obj') and not f.endswith('_gt_y.obj') and not f.endswith("input_y.obj"))] + ) + files_by_folder[folder_path] = base_names + # ... 现有代码 ... + print(f" ‐ 文件夹 '{os.path.basename(folder_path)}' 找到 {len(base_names)} 个文件。") + except Exception as e: + print(f" ‐ 错误：无法读取文件夹 {folder_path}: {e}") + # ... 现有代码 ... + files_by_folder[folder_path] = set() + + # 2. 找到所有文件夹中 *共有* 的文件基础名 (无变化) + common_base_names = set.intersection(*files_by_folder.values()) + + if not common_base_names: + # ... 现有代码 ... + print("\n错误：在所有文件夹中没有找到任何共有的 `_y.obj` 文件。") + print("="*80) + return None + + print(f"\n在所有文件夹中找到了 {len(common_base_names)} 个共同文件。将使用多进程处理...") + + # 3. [修改] 初始化指标存储 + global_chamfer_l2 = [] + global_hausdorff = [] + global_normal_abs = [] # [修改] + global_chamfer_l1 = [] + global_normal_real = [] # [新增] + + per_folder_metrics = {path: {'chamfer_l2': [], 'hausdorff': [], 'normal_abs': [], 'chamfer_l1': [], 'normal_real': []} for path in folder_paths} # [修改] + + # 4. [修改] 设置并运行多进程任务 + + # 准备任务参数列表 + # ... 现有代码 ... + task_args = [(base_name, gt_folder, folder_paths) for base_name in common_base_names] + all_results_list = [] + + # 确定进程数，可以适当调整，例如 os.cpu_count() + # ... 现有代码 ... + # 留出2个核心给系统，防止卡死 + num_processes = max(1, 60) + print(f"启动 {num_processes} 个工作进程...") + + # 创建进程池 + # ... 现有代码 ... + with multiprocessing.Pool(processes=num_processes) as pool: + # 使用 pool.imap 来迭代任务，并用 tqdm 显示进度条 + # imap 会按顺序返回结果，同时保持并行执行 + # ... 现有代码 ... + for result in tqdm(pool.imap(starmap_helper, task_args), + total=len(common_base_names), + desc="评估共同文件"): + all_results_list.append(result) + + # 5. [修改] 聚合所有进程返回的结果 + print("\n正在聚合所有进程的结果...") + valid_pairs_processed = 0 + for results_for_this_file in all_results_list: + if not results_for_this_file: + # ... 现有代码 ... + # GT 文件加载失败或跳过 + continue + + # [修改] 解包 L1, L2, |NC|, NC + for (folder_path, chamfer_l2, hausdorff, normal_error_abs, chamfer_l1, normal_error_real) in results_for_this_file: + if chamfer_l2 is not None and hausdorff is not None and normal_error_abs is not None and chamfer_l1 is not None and normal_error_real is not None: + # 这是一个有效的结果 + global_chamfer_l2.append(chamfer_l2) + global_hausdorff.append(hausdorff) + global_normal_abs.append(normal_error_abs) # [修改] + global_chamfer_l1.append(chamfer_l1) + global_normal_real.append(normal_error_real) # [新增] + + per_folder_metrics[folder_path]['chamfer_l2'].append(chamfer_l2) + per_folder_metrics[folder_path]['hausdorff'].append(hausdorff) + per_folder_metrics[folder_path]['normal_abs'].append(normal_error_abs) # [修改] + per_folder_metrics[folder_path]['chamfer_l1'].append(chamfer_l1) + per_folder_metrics[folder_path]['normal_real'].append(normal_error_real) # [新增] + valid_pairs_processed += 1 + else: + # ... 现有代码 ... + # 预测文件加载失败或计算错误 + # 可以在这里添加日志 + pass + + print(f"聚合完成。总共处理了 {valid_pairs_processed} 个有效的 (GT, Pred) 文件对。") + + # 6. [修改] 计算最终平均值 + if not global_chamfer_l2: # [修改] + print("\n错误：未能成功计算任何共同文件的指标。") + print("="*80) + return None + + # [修改] 计算总平均 + global_avg_results = { + "evaluated_pairs_total": len(global_chamfer_l2), + "common_files_count": len(common_base_names), + "folders_count": len(folder_paths), + "avg_chamfer_distance_l2": np.mean(global_chamfer_l2), + "avg_hausdorff_distance": np.mean(global_hausdorff), + "avg_normal_consistency_abs": np.mean(global_normal_abs), # [修改] + "avg_chamfer_distance_l1": np.mean(global_chamfer_l1), + "avg_normal_consistency_real": np.mean(global_normal_real) # [新增] + } + + # [修改] 计算逐文件夹平均 + per_folder_avg_results = {} + for folder_path, metrics in per_folder_metrics.items(): + folder_name = os.path.basename(folder_path) + if metrics['chamfer_l2']: # [修改] + per_folder_avg_results[folder_name] = { + "evaluated_pairs": len(metrics['chamfer_l2']), + "avg_chamfer_distance_l2": np.mean(metrics['chamfer_l2']), + "avg_hausdorff_distance": np.mean(metrics['hausdorff']), + "avg_normal_consistency_abs": np.mean(metrics['normal_abs']), # [修改] + "avg_chamfer_distance_l1": np.mean(metrics['chamfer_l1']), + "avg_normal_consistency_real": np.mean(metrics['normal_real']) # [新增] + } + else: + per_folder_avg_results[folder_name] = { + "evaluated_pairs": 0, + "avg_chamfer_distance_l2": 0, + "avg_hausdorff_distance": 0, + "avg_normal_consistency_abs": 0, # [修改] + "avg_chamfer_distance_l1": 0, + "avg_normal_consistency_real": 0 # [新增] + } + + # 7. [修改] 打印共同文件评估总结 + print("\n--- 共同文件评估总结 ---") + print(f"\n全局平均 (基于 {len(global_chamfer_l2)} 个总比较对, {len(common_base_names)} 个文件 x {len(folder_paths)} 个文件夹):") + print(f" ‐ 平均倒角距离 (L2): {global_avg_results['avg_chamfer_distance_l2']:.6f}") + print(f" ‐ 平均倒角距离 (L1): {global_avg_results['avg_chamfer_distance_l1']:.6f}") + print(f" ‐ 平均豪斯多夫距离: {global_avg_results['avg_hausdorff_distance']:.6f}") + print(f" ‐ 平均法向量 |NC| (Abs): {global_avg_results['avg_normal_consistency_abs']:.6f}") # [修改] + print(f" ‐ 平均法向量 NC (Real): {global_avg_results['avg_normal_consistency_real']:.6f}") # [新增] + + print("\n逐文件夹平均 (仅基于共同文件):") + for folder_name, results in per_folder_avg_results.items(): + print(f" ‐ 文件夹: {folder_name}") + # [修改] 打印所有6个指标 + print(f" ‐ 倒角(L2): {results['avg_chamfer_distance_l2']:.6f}, 倒角(L1): {results['avg_chamfer_distance_l1']:.6f}, 豪斯多夫: {results['avg_hausdorff_distance']:.6f}, 法向量|NC|: {results['avg_normal_consistency_abs']:.6f}, 法向量NC: {results['avg_normal_consistency_real']:.6f} (共 {results['evaluated_pairs']} 对)") + + print("="*80) + + # 8. 保存结果 (无变化) + final_results = { + # ... 现有代码 ... + "global_average_on_common_files": global_avg_results, + "per_folder_average_on_common_files": per_folder_avg_results + } + + output_filepath = os.path.join(base_dir, "common_evaluation_results.json") + # ... 现有代码 ... + try: + with open(output_filepath, 'w', encoding='utf-8') as f: + json.dump(final_results, f, ensure_ascii=False, indent=4) + # ... 现有代码 ... + print(f"\n共同文件评估结果已保存到: {output_filepath}") + except Exception as e: + print(f"\n错误：保存共同文件结果时失败: {e}") + # ... 现有代码 ... + + return final_results + + +# --- [主函数：run_evaluation_on_common_files_only] (打印部分已修改) --- +def run_evaluation_on_common_files_only(target_folders, gt_folder, base_dir=None): + """ + [无变化] + 在手动指定的多个文件夹上 *仅* 运行 *共同文件* 评估。 + """ + + # 1. 验证输入的文件夹列表 + # ... 现有代码 ... + if not isinstance(target_folders, list) or not target_folders: + print("错误：未提供任何文件夹用于评估。请检查 'target_folders' 列表。") + return + + # 2. 检查路径有效性并过滤非文件夹 + # ... 现有代码 ... + valid_folders = [] + for f_path in target_folders: + if os.path.isdir(f_path): + # ... 现有代码 ... + valid_folders.append(f_path) + else: + print(f"警告：提供的路径 '{f_path}' 不是一个有效的文件夹，已跳过。") + + if not valid_folders: + # ... 现有代码 ... + print("错误：提供的列表中没有找到任何有效的文件夹。") + return + + # 确保排序一致 + # ... 现有代码 ... + valid_folders.sort() + + # 3. 确定用于保存汇总文件的 base_dir + if base_dir is None: + # ... 现有代码 ... + # 如果未指定，使用第一个有效文件夹的父目录 + base_dir = os.path.dirname(valid_folders[0]) + print(f"信息：未指定 base_dir，将使用 '{base_dir}' 存放共同文件评估结果。") + elif not os.path.isdir(base_dir): + # ... 现有代码 ... + print(f"警告：指定的 base_dir '{base_dir}' 不存在。将改用第一个文件夹的父目录。") + base_dir = os.path.dirname(valid_folders[0]) + + print(f"\n找到了 {len(valid_folders)} 个有效的文件夹，将评估它们的共同文件：") + # ... 现有代码 ... + for folder in valid_folders: + print(f" - {folder}") # 打印完整路径 + + # 4. 第1部分：直接进行共同文件评估 + # ... 现有代码 ... + print("\n" + "#"*80) + print(" 开始共同文件评估 (仅计算所有文件夹共有的文件) ".center(80, "#")) + print("#"*80) + + # [调用修改后的函数] + # ... 现有代码 ... + common_results = evaluate_common_files(valid_folders, gt_folder, base_dir) + + # 5. [修改] 第2部分：打印最终的汇总报告 + print("\n" + "#"*80) + print(" 最终汇总报告 ".center(80, "#")) + print("#"*80) + + print("\n--- 共同文件评估结果 (仅基于共同文件) ---") + if common_results: + # 打印全局平均 + global_avg = common_results.get("global_average_on_common_files", {}) + if global_avg: + print(f"\n [总平均 (基于 {global_avg.get('common_files_count', 0)} 个共同文件)]") + print(f" ‐ 倒角(L2): {global_avg.get('avg_chamfer_distance_l2', 0):.6f}, " + f"倒角(L1): {global_avg.get('avg_chamfer_distance_l1', 0):.6f}, " + f"豪斯多夫: {global_avg.get('avg_hausdorff_distance', 0):.6f}, " + f"法向量|NC|: {global_avg.get('avg_normal_consistency_abs', 0):.6f}, " # [修改] + f"法向量NC: {global_avg.get('avg_normal_consistency_real', 0):.6f}") # [新增] + + # 打印逐文件夹平均 + per_folder_avg = common_results.get("per_folder_average_on_common_files", {}) + if per_folder_avg: + print("\n [逐文件夹平均 (仅共同文件)]") + for folder_name, results in per_folder_avg.items(): + print(f" ‐ 文件夹: {folder_name}") + print(f" ‐ 倒角(L2): {results.get('avg_chamfer_distance_l2', 0):.6f}, " + f"倒角(L1): {results.get('avg_chamfer_distance_l1', 0):.6f}, " + f"豪斯多夫: {results.get('avg_hausdorff_distance', 0):.6f}, " + f"法向量|NC|: {results.get('avg_normal_consistency_abs', 0):.6f}, " # [修改] + f"法向量NC: {results.get('avg_normal_consistency_real', 0):.6f} " # [新增] + f"(共 {results.get('evaluated_pairs', 0)} 对)") + else: + print(" (没有共同文件的评估结果)") + + print("\n" + "="*80) + print("所有评估任务已完成。") + print("="*80) + + +# =============================================================== +# 第五步：运行主程序 (无变化) +# =============================================================== + +if __name__ == "__main__": + import argparse + ap = argparse.ArgumentParser(description="评测: 预测目录(含 _y.obj) vs GT 目录(.glb/.obj)") + ap.add_argument("folders", nargs="+", help="一个或多个预测目录(只评估各目录共有的 case)") + ap.add_argument("--gt", required=True, help="Ground Truth 目录") + ap.add_argument("--out", default=None, help="汇总 JSON 的输出目录(默认第一个预测目录)") + a = ap.parse_args() + if not os.path.isdir(a.gt): + raise SystemExit(f"GT 目录不存在: {a.gt}") + run_evaluation_on_common_files_only(a.folders, a.gt, base_dir=a.out) diff --git a/mulres_infer/scripts/infer.py b/mulres_infer/scripts/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..abbbae4a66549451f23a8f4cb41af3a50b8b2633 --- /dev/null +++ b/mulres_infer/scripts/infer.py @@ -0,0 +1,92 @@ +"""多分辨率网格生成 —— 端到端推理。 + + 低分辨率生成器(TSSR 三前向离散扩散) 逐级生成到格 64, 每级 20 步; + res32 / res64 两级产物各补一次洞再喂下一级; + 高分辨率细分器继续 64 -> 128 -> 256 -> 512, 每级 20 步, 全程补洞。 + +用法: + python scripts/infer.py --data --out_dir outputs/run1 + python scripts/infer.py --data --out_dir outputs/run1 --ncase 20 --empty_bias 0.3 +""" +import argparse +import os +import sys + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, ROOT) + +DEFAULTS = dict( + STAGE="2,4,8,16,32,64", # 生成器级联阶梯(产物真实体素格) + STEPS="20", # 生成器每级扩散步数 + SUPERRES="1", SUPERRES_TO="512", SUPERRES_STEPS="20", # 细分器 64->512, 每级步数 + FILL="1", # 级联内补洞(res32/res64, 见 FILL_MIN_RES) + SR_FILL_IN="1", SUPERRES_FIX_HOLES="1", # 细分入口 + 每级输出补洞 + DECIMATE_RATIO="0.5", # 输入预塌缩比(0=不塌缩) + METHODS="plain", MULTI_PC="0", + SAVE_GT_LEVELS="1", SAVE_PC_LEVELS="1", + CONFIG=os.path.join(ROOT, "configs", "generator.yaml"), + SUPERRES_CONFIG=os.path.join(ROOT, "configs", "subdivider.yaml"), +) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--data", required=True, help="输入 mesh 目录(glb/obj/ply 混合)") + ap.add_argument("--out_dir", default=os.path.join(ROOT, "outputs"), help="输出根目录") + ap.add_argument("--ckpt", default=os.path.join(ROOT, "ckpt", "generator.pth"), help="生成器权重") + ap.add_argument("--subdivider_ckpt", default=os.path.join(ROOT, "ckpt", "subdivider.pth"), help="细分器权重") + ap.add_argument("--out_suffix", default="", help="输出子目录后缀") + ap.add_argument("--ncase", type=int, default=0, help="只跑前 N 个(文件名字母序); 0=全部") + ap.add_argument("--case", default="", help="只跑指定 case, 逗号分隔文件名子串(建议带尾点: 'a.' 不会匹配 'a1.glb')") + ap.add_argument("--gpu", default=None, help="CUDA_VISIBLE_DEVICES") + ap.add_argument("--steps", type=int, default=None, help="覆盖每级步数(默认 20)") + ap.add_argument("--no_fill", action="store_true", help="关闭所有补洞") + ap.add_argument("--fill_min_res", type=int, default=0, help=">0: 级联内补洞只在产物格>=该值的级做(默认 0=res32 起)") + ap.add_argument("--empty_bias", type=float, default=0.0, help=">0: 连接头空类 logit 偏置, 稀疏化控面数(0.3 约 -25%% 面)") + a = ap.parse_args() + + for k, v in DEFAULTS.items(): + os.environ.setdefault(k, v) + os.environ["RAIL_CKPT"] = os.path.abspath(a.ckpt) + os.environ["SUPERRES_CKPT"] = os.path.abspath(a.subdivider_ckpt) + os.environ["BASE_DIR"] = os.path.dirname(os.path.abspath(a.ckpt)) + os.environ["OUT_ROOT"] = os.path.abspath(a.out_dir) + os.environ["DATASET_DIR"] = os.path.abspath(a.data) + os.environ.setdefault("DATASET", "custom") + os.environ["NCASE"] = str(a.ncase if a.ncase > 0 else 10 ** 9) + if a.case: os.environ["CASE_STEM"] = a.case + if a.out_suffix: os.environ["OUT_SUFFIX"] = a.out_suffix + if a.gpu: os.environ["CUDA_VISIBLE_DEVICES"] = a.gpu + if a.steps: os.environ["STEPS"] = os.environ["SUPERRES_STEPS"] = str(a.steps) + if a.no_fill: os.environ.update(FILL="0", SR_FILL_IN="0", SUPERRES_FIX_HOLES="0") + if a.fill_min_res: os.environ["FILL_MIN_RES"] = str(a.fill_min_res) + if a.empty_bias: os.environ["EMPTY_BIAS"] = str(a.empty_bias) + + os.chdir(ROOT) + import torch + import modules.pipeline as P + from models.generator import FactorizedTSSRTransformer + + class TSSRAdapter(FactorizedTSSRTransformer): + """把主管线的 .sample() 转发到三前向 sample_tssr; forward 缺省补 mode=1(mask 步)。""" + + @torch.no_grad() + def sample(self, coord_list, depth, pc=None, pc_target=None, T_infer=10, cfg_scale=1.0, + temperature=0.0, cfg_schedule="constant", cfg_start=1.0, **_ignored): + return self.sample_tssr(coord_list, depth, pc=pc, pc_target=pc_target, T_infer=T_infer, + cfg_scale=cfg_scale, cfg_schedule=cfg_schedule, cfg_start=cfg_start, + mask_logits_lo=float(os.environ.get("MASK_LO", "0.2")), + mask_logits_hi=float(os.environ.get("MASK_HI", "0.9"))) + + def forward_varlen(self, *args, **kwargs): + kwargs.pop("cond_cache", None) + if kwargs.get("mode", None) is None and len(args) < 7: + kwargs["mode"] = 1 + return super().forward_varlen(*args, **kwargs) + + P.BaseModel = TSSRAdapter + P.main() + + +if __name__ == "__main__": + main() diff --git a/mulres_infer/scripts/strip_ckpt.py b/mulres_infer/scripts/strip_ckpt.py new file mode 100644 index 0000000000000000000000000000000000000000..1305b591811339488d4ceb8b11b2f8d269b830bc --- /dev/null +++ b/mulres_infer/scripts/strip_ckpt.py @@ -0,0 +1,46 @@ +"""权重瘦身: 只保留推理需要的部分。 + +推理的加载逻辑是「先 load model_state_dict, 再用 ema_state_dict.shadow 逐参数覆盖」, +所以把 EMA 影子权重合并进 model_state_dict 后只存一份, 推理结果逐位等价。 + +丢弃: optimizer_state_dict(最大头, AdamW 的 m/v 约 2x 参数量)、scheduler_state_dict、训练元信息。 + +用法: python tools_strip_ckpt.py <输入.pth> <输出.pth> [--no-ema] +""" +import os, sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import argparse, os, torch + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("src"); ap.add_argument("dst") + ap.add_argument("--no-ema", action="store_true", help="不合并 EMA, 保留原始训练权重") + a = ap.parse_args() + + ck = torch.load(a.src, map_location="cpu", weights_only=False) + print(f"输入 {a.src} ({os.path.getsize(a.src)/1e9:.2f} GB)") + print(" 顶层 key:", list(ck.keys())) + msd = ck["model_state_dict"] + shadow = (ck.get("ema_state_dict") or {}).get("shadow") if isinstance(ck.get("ema_state_dict"), dict) else None + + merged = 0 + if shadow and not a.no_ema: + for k in msd: + if k in shadow: + msd[k] = shadow[k].to(msd[k].dtype); merged += 1 + print(f" 已合并 EMA 影子权重: {merged}/{len(msd)} 个张量") + elif a.no_ema: + print(" --no-ema: 保留原始训练权重") + else: + print(" ckpt 内无 EMA, 直接使用 model_state_dict") + + out = {"model_state_dict": msd} + torch.save(out, a.dst) + print(f"输出 {a.dst} ({os.path.getsize(a.dst)/1e9:.2f} GB, " + f"省 {100*(1-os.path.getsize(a.dst)/os.path.getsize(a.src)):.0f}%)") + + +if __name__ == "__main__": + main() diff --git a/mulres_infer/utils/__init__.py b/mulres_infer/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/mulres_infer/utils/config.py b/mulres_infer/utils/config.py new file mode 100644 index 0000000000000000000000000000000000000000..f4282537dd5cefef5e9e74c8bc2c7a9bab2969cf --- /dev/null +++ b/mulres_infer/utils/config.py @@ -0,0 +1,37 @@ +# 配置加载: YAML -> 递归 DictToObject(点号访问) +import yaml +import argparse + +class DictToObject: + """将字典递归地转换为对象，允许使用点符号访问""" + def __init__(self, data): + for key, value in data.items(): + if isinstance(value, dict): + setattr(self, key, DictToObject(value)) + else: + setattr(self, key, value) + + def __repr__(self): + return f'{self.__dict__}' + +def load_config(): + """ + 解析命令行参数以获取配置文件路径， + 加载 YAML 文件并返回一个配置对象。 + """ + parser = argparse.ArgumentParser(description="Load training configuration from a YAML file.") + parser.add_argument('--config', type=str, default='./configs/generator.yaml', + help='Path to the configuration YAML file.') + + args = parser.parse_args() + print(f"config path :{args.config}") + + try: + with open(args.config, 'r') as f: + config_dict = yaml.safe_load(f) + print(f"Configuration loaded from {args.config}") + except FileNotFoundError: + print(f"Error: Configuration file not found at {args.config}") + exit(1) + + return DictToObject(config_dict) diff --git a/mulres_infer/utils/data_utils.py b/mulres_infer/utils/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9167b56756c5575d86c48d91fcd61153d68e878b --- /dev/null +++ b/mulres_infer/utils/data_utils.py @@ -0,0 +1,38 @@ +import numpy as np +import trimesh +def discretize( + t, + continuous_range = (-1.0, 1.0), + num_discrete: int = 128 +): + lo, hi = continuous_range + t = (t - lo) / (hi - lo) + t *= num_discrete + t -= 0.5 + + return t.round().astype(np.int32).clip(min = 0, max = num_discrete - 1) + +def undiscretize( + t, + continuous_range = (-1.0, 1.0), + num_discrete: int = 128 +): + lo, hi = continuous_range + t = t + 0.5 + t /= num_discrete + return t * (hi - lo) + lo + +def normalize(mesh, eps: float = 1e-8): + bounds = np.array([mesh.vertices.min(axis=0), mesh.vertices.max(axis=0)]) + mesh.vertices = mesh.vertices - (bounds[0] + bounds[1])[None, :] / 2 + # 数据中存在退化 mesh (bounds 范围为 0), 分母加 eps 避免除零产生 inf/NaN + scale = (bounds[1] - bounds[0]).max() + eps + mesh.vertices = 2 * (mesh.vertices / scale) + mesh.vertices = mesh.vertices.clip(-1.0, 1.0) # 限制顶点坐标在[-1.0, 1.0] + + +def clean_mesh(mesh): + mesh.merge_vertices() + mesh.update_faces(mesh.nondegenerate_faces()) + mesh.update_faces(mesh.unique_faces()) + mesh.remove_unreferenced_vertices() \ No newline at end of file diff --git a/mulres_infer/utils/distance.py b/mulres_infer/utils/distance.py new file mode 100644 index 0000000000000000000000000000000000000000..736939b515beddd68310388f28d7372078c782cb --- /dev/null +++ b/mulres_infer/utils/distance.py @@ -0,0 +1,88 @@ +"""点云距离与法向一致性指标(倒角 L1/L2、豪斯多夫、|NC| / NC)。 + +modules.pipeline 与 scripts.eval 共用同一份实现, 保证推理内打印的指标与离线评测同口径。 +""" +import numpy as np +from scipy.spatial import cKDTree + + +def calculate_distances_numpy(pcd1, pcd2): + """ + [已修改] + 使用 NumPy 高效计算 L2倒角距离、L1倒角距离 和 豪斯多夫距离。 + + 参数: + pcd1: np.ndarray, 形状为 (N, 3)，第一个点云。 + pcd2: np.ndarray, 形状为 (M, 3)，第二个点云。 + + 返回: + (float, float, float): (L2倒角距离, 豪斯多夫距离, L1倒角距离) + """ + # --- 计算所有点对之间的差异 --- + pcd1_expanded = pcd1[:, np.newaxis, :] + pcd2_expanded = pcd2[np.newaxis, :, :] + diff = pcd1_expanded - pcd2_expanded + + # --- L2 (欧氏) 距离计算 --- + dist_matrix_sq = np.sum(diff**2, axis=2) + + # --- L2 倒角距离 (Chamfer Distance L2) --- + min_dists_1_to_2_l2 = np.sqrt(np.min(dist_matrix_sq, axis=1)) + min_dists_2_to_1_l2 = np.sqrt(np.min(dist_matrix_sq, axis=0)) + chamfer_dist_l2 = np.mean(min_dists_1_to_2_l2) + np.mean(min_dists_2_to_1_l2) + + # --- L2 豪斯多夫距离 (Hausdorff Distance) --- + hausdorff_dist = max(np.max(min_dists_1_to_2_l2), np.max(min_dists_2_to_1_l2)) + + # --- L1 (曼哈顿) 距离计算 --- + dist_matrix_l1 = np.sum(np.abs(diff), axis=2) + + # --- L1 倒角距离 (Chamfer Distance L1) --- + min_dists_1_to_2_l1 = np.min(dist_matrix_l1, axis=1) + min_dists_2_to_1_l1 = np.min(dist_matrix_l1, axis=0) + chamfer_dist_l1 = np.mean(min_dists_1_to_2_l1) + np.mean(min_dists_2_to_1_l1) + + return chamfer_dist_l2, hausdorff_dist, chamfer_dist_l1 + + +def calculate_normal_consistency_score(pcd_gt, normal_gt, pcd_pred, normal_pred): + """ + 【修改版】根据提供的公式计算法向量一致性得分 (Normal Consistency Score)。 + 返回两个值： + 1. |NC| (绝对值 NC): np.mean( |n1·n2| ) + 2. NC (真实值 NC): np.mean( n1·n2 ) + """ + + def _calculate_one_way_score(source_pcd, source_normals, target_pcd, target_normals): + """ + 辅助函数：计算从 source 到 target 的单向 NC 得分。 + 返回 (mean_abs_dot, mean_real_dot) + """ + # 1. 计算 source 到 target 的所有点对的距离矩阵 + diff = source_pcd[:, np.newaxis, :] - target_pcd[np.newaxis, :, :] + dist_sq = np.sum(diff**2, axis=2) + + # 2. 为每个 source 点找到 target 中的最近邻点的索引 + indices = np.argmin(dist_sq, axis=1) + + # 3. 获取这些最近邻点对应的法向量 + corresponding_normals = target_normals[indices] + + # 4. 计算 source 法向量与对应 target 法向量的点积 + dot_products = np.einsum('ij,ij->i', source_normals, corresponding_normals) + + # 5. [修改] 返回点积绝对值的平均值 和 真实点积的平均值 + return np.mean(np.abs(dot_products)), np.mean(dot_products) + + # --- Part 1: 计算从 GT -> Pred 的得分 --- + score_gt_to_pred_abs, score_gt_to_pred_real = _calculate_one_way_score(pcd_gt, normal_gt, pcd_pred, normal_pred) + + # --- Part 2: 计算从 Pred -> GT 的得分 --- + score_pred_to_gt_abs, score_pred_to_gt_real = _calculate_one_way_score(pcd_pred, normal_pred, pcd_gt, normal_gt) + + # --- Final Score: 结合两部分 --- + final_score_abs = (score_gt_to_pred_abs + score_pred_to_gt_abs) / 2.0 + final_score_real = (score_gt_to_pred_real + score_pred_to_gt_real) / 2.0 + + # [修改] 返回两个分数 + return final_score_abs, final_score_real diff --git a/mulres_infer/utils/mesh_metric.py b/mulres_infer/utils/mesh_metric.py new file mode 100644 index 0000000000000000000000000000000000000000..f1321f2b54471704ec1fdbb63d6d5b0f477dc7f6 --- /dev/null +++ b/mulres_infer/utils/mesh_metric.py @@ -0,0 +1,91 @@ +import numpy as np +import trimesh +from scipy.spatial import cKDTree +pc_num = 10000 +def calculate_distances_numpy(pcd1, pcd2): + """ + 使用 NumPy 高效计算倒角距离和豪斯多夫距离。 + + 参数: + pcd1: np.ndarray, 形状为 (N, 3)，第一个点云。 + pcd2: np.ndarray, 形状为 (M, 3)，第二个点云。 + + 返回: + (float, float): (倒角距离, 豪斯多夫距离) + """ + # --- 计算所有点对之间的平方欧氏距离矩阵 --- + pcd1_expanded = pcd1[:, np.newaxis, :] + pcd2_expanded = pcd2[np.newaxis, :, :] + diff = pcd1_expanded - pcd2_expanded + dist_matrix_sq = np.sum(diff**2, axis=2) + + # --- 计算倒角距离 (Chamfer Distance) --- + min_dists_1_to_2 = np.sqrt(np.min(dist_matrix_sq, axis=1)) + min_dists_2_to_1 = np.sqrt(np.min(dist_matrix_sq, axis=0)) + chamfer_dist = np.mean(min_dists_1_to_2) + np.mean(min_dists_2_to_1) + + # --- 计算豪斯多夫距离 (Hausdorff Distance) --- + hausdorff_dist = max(np.max(min_dists_1_to_2), np.max(min_dists_2_to_1)) + + return chamfer_dist, hausdorff_dist + +def normalize(mesh): + bounds = np.array([mesh.vertices.min(axis=0), mesh.vertices.max(axis=0)]) + mesh.vertices = mesh.vertices - (bounds[0] + bounds[1])[None, :] / 2 + mesh.vertices = mesh.vertices / (bounds[1] - bounds[0]).max() + mesh.vertices = mesh.vertices.clip(-0.5, 0.5) # 限制顶点坐标在[-0.5, 0.5] + +def sample_from_obj(mesh): + """ + 从 .obj 文件中加载顶点坐标。 + + 参数: + file_path: str, .obj 文件的路径。 + + 返回: + np.ndarray: 形状为 (V, 3) 的顶点数组，V 是顶点数量。 + """ + normalize(mesh) + points, face_idx = mesh.sample(pc_num, return_index=True) + normals = mesh.face_normals[face_idx] + return points,normals + +def compute_mesh_metrics(gt_mesh, pred_mesh): + pcd_gt, normal_gt = sample_from_obj(gt_mesh) + pcd_pred, normal_pred = sample_from_obj(pred_mesh) + cd, hd = calculate_distances_numpy(pcd_gt, pcd_pred) + return cd, hd + +def compute_mesh_metrics_new(mesh_gt, mesh_recon, num_samples=100000): + """计算 Chamfer Distance 和 Hausdorff Distance""" + if mesh_gt is None or mesh_recon is None: + return None, None + + # 表面采样点云 + p_gt, _ = trimesh.sample.sample_surface(mesh_gt, num_samples) + p_recon, _ = trimesh.sample.sample_surface(mesh_recon, num_samples) + + # 构建 KD-Tree + tree_gt = cKDTree(p_gt) + tree_recon = cKDTree(p_recon) + + # 双向最近邻搜索 + d_recon_to_gt, _ = tree_gt.query(p_recon) + d_gt_to_recon, _ = tree_recon.query(p_gt) + + # CD: $\frac{1}{N} \sum \min\|p-q\| + \frac{1}{M} \sum \min\|q-p\|$ + cd = np.mean(d_recon_to_gt) + np.mean(d_gt_to_recon) + # HD: $\max( \max \min\|p-q\|, \max \min\|q-p\| )$ + hd = max(np.max(d_recon_to_gt), np.max(d_gt_to_recon)) + + return cd, hd + +if __name__ == "__main__": + gt_path = "octree_results_fix/res_128.obj" + pred_path = "tree_export_results/tree_level_128.obj" + gt_mesh = trimesh.load(gt_path) + pred_mesh = trimesh.load(pred_path) + normalize(gt_mesh) + normalize(pred_mesh) + cd, hd = compute_mesh_metrics_new(gt_mesh, pred_mesh) + print(cd, hd) \ No newline at end of file diff --git a/mulres_infer/utils/postprocessing.py b/mulres_infer/utils/postprocessing.py new file mode 100644 index 0000000000000000000000000000000000000000..3496d65f51e746618695aaa4b32624d4b4744570 --- /dev/null +++ b/mulres_infer/utils/postprocessing.py @@ -0,0 +1,1428 @@ +from typing import * +import numpy as np +import torch +import utils3d +import nvdiffrast.torch as dr +from tqdm import tqdm +import trimesh +import trimesh.visual +import pyvista as pv +from pymeshfix import _meshfix +import igraph +PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53] +import itertools +import os as _os, sys as _sys, contextlib as _contextlib + +MESHFIX_QUIET = _os.environ.get("MESHFIX_QUIET", "1") == "1" # 屏蔽 PyMeshFix/TMesh C++ 层日志(Loading..%, WARNING-, % done) + + +@_contextlib.contextmanager +def _suppress_c_output(enabled=True): + """临时把 fd1/fd2 重定向到 /dev/null: TMesh 的日志是 C 层 printf, Python 的 sys.stdout 拦不住。""" + if not enabled: + yield; return + _sys.stdout.flush(); _sys.stderr.flush() + saved_out, saved_err = _os.dup(1), _os.dup(2) + devnull = _os.open(_os.devnull, _os.O_WRONLY) + try: + _os.dup2(devnull, 1); _os.dup2(devnull, 2) + yield + finally: + _os.dup2(saved_out, 1); _os.dup2(saved_err, 2) + _os.close(saved_out); _os.close(saved_err); _os.close(devnull) + +def radical_inverse(base, n): + val = 0 + inv_base = 1.0 / base + inv_base_n = inv_base + while n > 0: + digit = n % base + val += digit * inv_base_n + n //= base + inv_base_n *= inv_base + return val + +def halton_sequence(dim, n): + return [radical_inverse(PRIMES[dim], n) for dim in range(dim)] + +def hammersley_sequence(dim, n, num_samples): + return [n / num_samples] + halton_sequence(dim - 1, n) + +def sphere_hammersley_sequence(n, num_samples, offset=(0, 0), remap=False): + u, v = hammersley_sequence(2, n, num_samples) + u += offset[0] / num_samples + v += offset[1] + if remap: + u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3 + theta = np.arccos(1 - 2 * u) - np.pi / 2 + phi = v * 2 * np.pi + return [phi, theta] + +@torch.no_grad() +# def _fill_holes( +# verts, +# faces, +# max_hole_size=0.04, +# max_hole_nbe=32, +# resolution=128, +# num_views=500, +# debug=False, +# verbose=False, +# only_fill_holes=False, +# ): +# """ +# Rasterize a mesh from multiple views and remove invisible faces. +# Also includes postprocessing to: +# 1. Remove connected components that are have low visibility. +# 2. Mincut to remove faces at the inner side of the mesh connected to the outer side with a small hole. + +# Args: +# verts (torch.Tensor): Vertices of the mesh. Shape (V, 3). +# faces (torch.Tensor): Faces of the mesh. Shape (F, 3). +# max_hole_size (float): Maximum area of a hole to fill. +# resolution (int): Resolution of the rasterization. +# num_views (int): Number of views to rasterize the mesh. +# verbose (bool): Whether to print progress. +# """ +# if only_fill_holes: +# mesh = _meshfix.PyTMesh() +# mesh.load_array(verts.cpu().numpy(), faces.cpu().numpy()) +# mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True) +# verts, faces = mesh.return_arrays() +# verts = torch.tensor(verts, device='cuda', dtype=torch.float32) +# faces = torch.tensor(faces, device='cuda', dtype=torch.int32) +# if verbose: +# tqdm.write('Skip mincut and run meshfix hole filling only') +# return verts, faces +# # Construct cameras +# yaws = [] +# pitchs = [] +# for i in range(num_views): +# y, p = sphere_hammersley_sequence(i, num_views) +# yaws.append(y) +# pitchs.append(p) +# yaws = torch.tensor(yaws).cuda() +# pitchs = torch.tensor(pitchs).cuda() +# radius = 2.0 +# fov = torch.deg2rad(torch.tensor(40)).cuda() +# projection = utils3d.torch.perspective_from_fov_xy(fov, fov, 1, 3) +# views = [] +# for (yaw, pitch) in zip(yaws, pitchs): +# orig = torch.tensor([ +# torch.sin(yaw) * torch.cos(pitch), +# torch.cos(yaw) * torch.cos(pitch), +# torch.sin(pitch), +# ]).cuda().float() * radius +# view = utils3d.torch.view_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda()) +# views.append(view) +# views = torch.stack(views, dim=0) + +# # Rasterize +# visblity = torch.zeros(faces.shape[0], dtype=torch.int32, device=verts.device) +# rastctx = utils3d.torch.RastContext(backend='cuda') +# for i in tqdm(range(views.shape[0]), total=views.shape[0], disable=not verbose, desc='Rasterizing'): +# view = views[i] +# buffers = utils3d.torch.rasterize_triangle_faces( +# rastctx, verts[None], faces, resolution, resolution, view=view, projection=projection +# ) +# face_id = buffers['face_id'][0][buffers['mask'][0] > 0.95] - 1 +# face_id = torch.unique(face_id).long() +# visblity[face_id] += 1 +# visblity = visblity.float() / num_views + + + +# # Mincut +# ## construct outer faces +# edges, face2edge, edge_degrees = utils3d.torch.compute_edges(faces) +# boundary_edge_indices = torch.nonzero(edge_degrees == 1).reshape(-1) +# connected_components = utils3d.torch.compute_connected_components(faces, edges, face2edge) +# outer_face_indices = torch.zeros(faces.shape[0], dtype=torch.bool, device=faces.device) +# for i in range(len(connected_components)): +# outer_face_indices[connected_components[i]] = visblity[connected_components[i]] > min(max(visblity[connected_components[i]].quantile(0.75).item(), 0.25), 0.5) +# outer_face_indices = outer_face_indices.nonzero().reshape(-1) + +# ## construct inner faces +# inner_face_indices = torch.nonzero(visblity == 0).reshape(-1) +# if verbose: +# tqdm.write(f'Found {inner_face_indices.shape[0]} invisible faces') +# if inner_face_indices.shape[0] == 0: +# return verts, faces + +# ## Construct dual graph (faces as nodes, edges as edges) +# dual_edges, dual_edge2edge = utils3d.torch.compute_dual_graph(face2edge) +# dual_edge2edge = edges[dual_edge2edge] +# dual_edges_weights = torch.norm(verts[dual_edge2edge[:, 0]] - verts[dual_edge2edge[:, 1]], dim=1) +# if verbose: +# tqdm.write(f'Dual graph: {dual_edges.shape[0]} edges') + +# ## solve mincut problem +# ### construct main graph +# g = igraph.Graph() +# g.add_vertices(faces.shape[0]) +# g.add_edges(dual_edges.cpu().numpy()) +# g.es['weight'] = dual_edges_weights.cpu().numpy() + +# ### source and target +# g.add_vertex('s') +# g.add_vertex('t') + +# ### connect invisible faces to source +# g.add_edges([(f, 's') for f in inner_face_indices], attributes={'weight': torch.ones(inner_face_indices.shape[0], dtype=torch.float32).cpu().numpy()}) + +# ### connect outer faces to target +# g.add_edges([(f, 't') for f in outer_face_indices], attributes={'weight': torch.ones(outer_face_indices.shape[0], dtype=torch.float32).cpu().numpy()}) + +# ### solve mincut +# cut = g.mincut('s', 't', (np.array(g.es['weight']) * 1000).tolist()) +# remove_face_indices = torch.tensor([v for v in cut.partition[0] if v < faces.shape[0]], dtype=torch.long, device=faces.device) +# if verbose: +# tqdm.write(f'Mincut solved, start checking the cut') + +# ### check if the cut is valid with each connected component +# to_remove_cc = utils3d.torch.compute_connected_components(faces[remove_face_indices]) +# if debug: +# tqdm.write(f'Number of connected components of the cut: {len(to_remove_cc)}') +# valid_remove_cc = [] +# cutting_edges = [] +# for cc in to_remove_cc: +# #### check if the connected component has low visibility +# visblity_median = visblity[remove_face_indices[cc]].median() +# if debug: +# tqdm.write(f'visblity_median: {visblity_median}') +# if visblity_median > 0.25: +# continue + +# #### check if the cuting loop is small enough +# cc_edge_indices, cc_edges_degree = torch.unique(face2edge[remove_face_indices[cc]], return_counts=True) +# cc_boundary_edge_indices = cc_edge_indices[cc_edges_degree == 1] +# cc_new_boundary_edge_indices = cc_boundary_edge_indices[~torch.isin(cc_boundary_edge_indices, boundary_edge_indices)] +# if len(cc_new_boundary_edge_indices) > 0: +# cc_new_boundary_edge_cc = utils3d.torch.compute_edge_connected_components(edges[cc_new_boundary_edge_indices]) +# cc_new_boundary_edges_cc_center = [verts[edges[cc_new_boundary_edge_indices[edge_cc]]].mean(dim=1).mean(dim=0) for edge_cc in cc_new_boundary_edge_cc] +# cc_new_boundary_edges_cc_area = [] +# for i, edge_cc in enumerate(cc_new_boundary_edge_cc): +# _e1 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 0]] - cc_new_boundary_edges_cc_center[i] +# _e2 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 1]] - cc_new_boundary_edges_cc_center[i] +# cc_new_boundary_edges_cc_area.append(torch.norm(torch.cross(_e1, _e2, dim=-1), dim=1).sum() * 0.5) +# if debug: +# cutting_edges.append(cc_new_boundary_edge_indices) +# tqdm.write(f'Area of the cutting loop: {cc_new_boundary_edges_cc_area}') +# if any([l > max_hole_size for l in cc_new_boundary_edges_cc_area]): +# continue + +# valid_remove_cc.append(cc) + +# if debug: +# face_v = verts[faces].mean(dim=1).cpu().numpy() +# vis_dual_edges = dual_edges.cpu().numpy() +# vis_colors = np.zeros((faces.shape[0], 3), dtype=np.uint8) +# vis_colors[inner_face_indices.cpu().numpy()] = [0, 0, 255] +# vis_colors[outer_face_indices.cpu().numpy()] = [0, 255, 0] +# vis_colors[remove_face_indices.cpu().numpy()] = [255, 0, 255] +# if len(valid_remove_cc) > 0: +# vis_colors[remove_face_indices[torch.cat(valid_remove_cc)].cpu().numpy()] = [255, 0, 0] +# utils3d.io.write_ply('dbg_dual.ply', face_v, edges=vis_dual_edges, vertex_colors=vis_colors) + +# vis_verts = verts.cpu().numpy() +# vis_edges = edges[torch.cat(cutting_edges)].cpu().numpy() +# utils3d.io.write_ply('dbg_cut.ply', vis_verts, edges=vis_edges) + + +# if len(valid_remove_cc) > 0: +# remove_face_indices = remove_face_indices[torch.cat(valid_remove_cc)] +# mask = torch.ones(faces.shape[0], dtype=torch.bool, device=faces.device) +# mask[remove_face_indices] = 0 +# faces = faces[mask] +# faces, verts = utils3d.torch.remove_unreferenced_vertices(faces, verts) +# if verbose: +# tqdm.write(f'Removed {(~mask).sum()} faces by mincut') +# else: +# if verbose: +# tqdm.write(f'Removed 0 faces by mincut') + +# mesh = _meshfix.PyTMesh() +# mesh.load_array(verts.cpu().numpy(), faces.cpu().numpy()) +# mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True) +# verts, faces = mesh.return_arrays() +# verts, faces = torch.tensor(verts, device='cuda', dtype=torch.float32), torch.tensor(faces, device='cuda', dtype=torch.int32) + +# return verts, faces + +def _detect_degenerate_faces( + verts: torch.Tensor, + faces: torch.Tensor, + area_eps: float = 1e-12, +) -> torch.Tensor: + """ + 检测退化三角面。 + + 判定条件： + 1. 一个面中存在重复顶点索引； + 2. 三角形二倍面积小于等于 area_eps； + 3. 三角形顶点坐标中存在 NaN 或 Inf。 + + Args: + verts: + [V, 3] 顶点。 + faces: + [F, 3] 三角面。 + area_eps: + 二倍面积阈值。 + + Returns: + [F] bool mask，True 表示退化面。 + """ + if faces.numel() == 0: + return torch.zeros( + faces.shape[0], + dtype=torch.bool, + device=faces.device, + ) + + faces_long = faces.long() + + repeated_vertex = ( + (faces_long[:, 0] == faces_long[:, 1]) + | (faces_long[:, 1] == faces_long[:, 2]) + | (faces_long[:, 2] == faces_long[:, 0]) + ) + + triangles = verts[faces_long] + + invalid_position = ~torch.isfinite(triangles).all( + dim=-1 + ).all(dim=-1) + + edge_1 = triangles[:, 1] - triangles[:, 0] + edge_2 = triangles[:, 2] - triangles[:, 0] + + double_area = torch.linalg.vector_norm( + torch.cross(edge_1, edge_2, dim=-1), + dim=-1, + ) + + zero_area = double_area <= area_eps + + return repeated_vertex | zero_area | invalid_position + + +def _find_matching_vertex( + position: np.ndarray, + vertices_list: list, + vertex_buckets: dict, + tolerance: float, + unavailable_indices: set, +): + """ + 在已有顶点中查找与 position 接近的顶点。 + + 会搜索当前量化单元及其相邻的 26 个单元，避免浮点量化边界问题。 + """ + base_key = np.round( + position / tolerance + ).astype(np.int64) + + best_index = None + best_distance = np.inf + + for offset in itertools.product( + (-1, 0, 1), + repeat=3, + ): + key = tuple( + base_key + np.asarray(offset, dtype=np.int64) + ) + + for candidate in vertex_buckets.get(key, []): + if candidate in unavailable_indices: + continue + + distance = np.linalg.norm( + vertices_list[candidate] - position + ) + + if ( + distance <= tolerance * 2.0 + and distance < best_distance + ): + best_distance = distance + best_index = candidate + + return best_index + + +def _restore_degenerate_faces( + vertices: np.ndarray, + faces: np.ndarray, + original_vertices: np.ndarray, + preserved_faces: np.ndarray, + vertex_match_eps: float = 1e-7, +): + """ + 将预先保存的退化面恢复到处理后的 mesh 中。 + + PyMeshFix 可能重排或删除顶点，因此不能直接使用原 face index。 + 本函数根据顶点坐标重新建立映射。 + + 同一个原始顶点 ID 会映射到同一个新顶点 ID； + 不同但坐标相同的原始顶点会尽量保持为不同顶点。 + + Args: + vertices: + 处理后的 [V, 3] 顶点。 + faces: + 处理后的 [F, 3] 面。 + original_vertices: + 处理前的原始顶点。 + preserved_faces: + 需要恢复的原始退化面索引。 + vertex_match_eps: + 顶点坐标匹配的相对容差。 + + Returns: + output_vertices, output_faces + """ + vertices = np.asarray(vertices) + faces = np.asarray(faces) + original_vertices = np.asarray(original_vertices) + preserved_faces = np.asarray(preserved_faces) + + if preserved_faces.size == 0: + return vertices, faces + + if vertices.ndim != 2: + vertices = vertices.reshape(-1, 3) + + if faces.size == 0: + faces = np.empty( + (0, 3), + dtype=preserved_faces.dtype, + ) + else: + faces = faces.reshape(-1, 3) + + if len(vertices) > 0: + bbox_diagonal = np.linalg.norm( + vertices.max(axis=0) + - vertices.min(axis=0) + ) + elif len(original_vertices) > 0: + bbox_diagonal = np.linalg.norm( + original_vertices.max(axis=0) + - original_vertices.min(axis=0) + ) + else: + bbox_diagonal = 1.0 + + tolerance = max( + float(vertex_match_eps) + * max(float(bbox_diagonal), 1.0), + 1e-12, + ) + + vertices_list = [ + vertex.copy() + for vertex in vertices + ] + + vertex_buckets = {} + + for vertex_index, vertex in enumerate(vertices_list): + key = tuple( + np.round( + vertex / tolerance + ).astype(np.int64) + ) + + vertex_buckets.setdefault( + key, + [], + ).append(vertex_index) + + original_vertex_ids = np.unique( + preserved_faces.reshape(-1) + ) + + old_to_new = {} + + # 防止两个不同的原始顶点 ID 因为坐标相同而被错误合并。 + claimed_existing_indices = set() + + for old_vertex_id in original_vertex_ids: + old_vertex_id = int(old_vertex_id) + position = original_vertices[old_vertex_id] + + matched_index = _find_matching_vertex( + position=position, + vertices_list=vertices_list, + vertex_buckets=vertex_buckets, + tolerance=tolerance, + unavailable_indices=claimed_existing_indices, + ) + + if matched_index is None: + matched_index = len(vertices_list) + vertices_list.append(position.copy()) + + key = tuple( + np.round( + position / tolerance + ).astype(np.int64) + ) + + vertex_buckets.setdefault( + key, + [], + ).append(matched_index) + else: + claimed_existing_indices.add(matched_index) + + old_to_new[old_vertex_id] = matched_index + + restored_faces = np.empty_like( + preserved_faces, + dtype=np.int64, + ) + + for column in range(3): + restored_faces[:, column] = np.fromiter( + ( + old_to_new[int(vertex_id)] + for vertex_id + in preserved_faces[:, column] + ), + dtype=np.int64, + count=preserved_faces.shape[0], + ) + + output_vertices = np.asarray( + vertices_list, + dtype=vertices.dtype, + ) + + restored_faces = restored_faces.astype( + faces.dtype, + copy=False, + ) + + output_faces = np.concatenate( + [faces, restored_faces], + axis=0, + ) + + return output_vertices, output_faces + + +def _run_meshfix_hole_filling( + verts: torch.Tensor, + faces: torch.Tensor, + max_hole_nbe: int, +): + """ + 调用 PyMeshFix 进行小边界补洞。 + + 注意：PyMeshFix 可能自行清理退化面。 + 需要保留的退化面应在调用本函数前暂时移除， + 并在本函数执行结束后恢复。 + """ + if faces.numel() == 0: + return verts, faces + + device = verts.device + verts_dtype = verts.dtype + faces_dtype = faces.dtype + + with _suppress_c_output(MESHFIX_QUIET): + mesh = _meshfix.PyTMesh() + + mesh.load_array( + verts.detach().cpu().numpy(), + faces.detach().cpu().numpy(), + ) + + mesh.fill_small_boundaries( + nbe=max_hole_nbe, + refine=True, + ) + + output_verts, output_faces = mesh.return_arrays() + + output_verts = torch.as_tensor( + output_verts, + device=device, + dtype=verts_dtype, + ) + + output_faces = torch.as_tensor( + output_faces, + device=device, + dtype=faces_dtype, + ) + + return output_verts, output_faces + + +def _fill_holes( + verts, + faces, + max_hole_size=0.04, + max_hole_nbe=32, + resolution=128, + num_views=500, + debug=False, + verbose=False, + only_fill_holes=False, + keep_degenerate_faces=False, + degenerate_area_eps=1e-12, + vertex_match_eps=1e-7, +): + """ + Rasterize a mesh from multiple views and remove invisible faces. + + 后处理包括： + 1. 删除低可见性的连通区域； + 2. 使用 mincut 删除通过小孔与外侧连接的内部面； + 3. 使用 PyMeshFix 填补小边界洞。 + + Args: + verts: + torch.Tensor，[V, 3] 顶点坐标。 + faces: + torch.Tensor，[F, 3] 三角面。 + max_hole_size: + mincut 新产生切割环允许的最大面积。 + max_hole_nbe: + PyMeshFix 填补边界时允许的最大边数量。 + resolution: + 光栅化分辨率。 + num_views: + 多视角光栅化视角数量。 + debug: + 是否输出 debug 文件。 + verbose: + 是否输出进度信息。 + only_fill_holes: + True 时跳过光栅化和 mincut，只运行 MeshFix 补洞。 + keep_degenerate_faces: + False： + 保持原始行为。退化面可能在 mincut 或 PyMeshFix 中消失。 + + True： + 处理前保存所有退化面，并暂时将它们从处理网格中移除； + 补洞结束后再将其恢复。 + degenerate_area_eps: + 判断零面积面的二倍面积阈值。 + vertex_match_eps: + 恢复退化面时的顶点坐标匹配容差。 + + Returns: + verts: + 处理后的顶点。 + faces: + 处理后的三角面。 + + 注意： + keep_degenerate_faces=True 时，退化面会被追加到输出 faces 末尾， + 所以输出 face 顺序不保证与输入完全一致。 + """ + if verts is None or faces is None: + return verts, faces + + if faces.numel() == 0: + return verts, faces + + original_device = verts.device + original_verts_dtype = verts.dtype + original_faces_dtype = faces.dtype + + original_vertices_numpy = None + preserved_degenerate_faces_numpy = None + + # ============================================================ + # 保存退化面，并将其从后续处理网格中暂时移除 + # ============================================================ + + if keep_degenerate_faces: + degenerate_mask = _detect_degenerate_faces( + verts=verts, + faces=faces, + area_eps=degenerate_area_eps, + ) + + num_degenerate = int( + degenerate_mask.sum().item() + ) + + if verbose: + tqdm.write( + f'Found {num_degenerate} degenerate input faces' + ) + + if num_degenerate > 0: + original_vertices_numpy = ( + verts.detach().cpu().numpy().copy() + ) + + preserved_degenerate_faces_numpy = ( + faces[degenerate_mask] + .detach() + .cpu() + .numpy() + .copy() + ) + + faces = faces[~degenerate_mask] + + def finalize_output( + output_verts: torch.Tensor, + output_faces: torch.Tensor, + ): + """ + 在返回前恢复之前保存的退化面。 + """ + if ( + not keep_degenerate_faces + or preserved_degenerate_faces_numpy is None + or preserved_degenerate_faces_numpy.size == 0 + ): + return output_verts, output_faces + + output_vertices_numpy = ( + output_verts.detach().cpu().numpy() + ) + + output_faces_numpy = ( + output_faces.detach().cpu().numpy() + ) + + ( + output_vertices_numpy, + output_faces_numpy, + ) = _restore_degenerate_faces( + vertices=output_vertices_numpy, + faces=output_faces_numpy, + original_vertices=original_vertices_numpy, + preserved_faces=preserved_degenerate_faces_numpy, + vertex_match_eps=vertex_match_eps, + ) + + output_verts = torch.as_tensor( + output_vertices_numpy, + device=original_device, + dtype=original_verts_dtype, + ) + + output_faces = torch.as_tensor( + output_faces_numpy, + device=original_device, + dtype=original_faces_dtype, + ) + + if verbose: + tqdm.write( + f'Restored ' + f'{len(preserved_degenerate_faces_numpy)} ' + f'degenerate faces' + ) + + return output_verts, output_faces + + # 如果所有输入面都是退化面，没有可补洞的有效表面 + if faces.numel() == 0: + if keep_degenerate_faces: + return finalize_output( + verts[:0], + faces.reshape(0, 3), + ) + + return verts, faces + + # ============================================================ + # 只运行 PyMeshFix 补洞 + # ============================================================ + + if only_fill_holes: + verts, faces = _run_meshfix_hole_filling( + verts=verts, + faces=faces, + max_hole_nbe=max_hole_nbe, + ) + + if verbose: + tqdm.write( + 'Skip mincut and run MeshFix hole filling only' + ) + + return finalize_output(verts, faces) + + device = verts.device + + # ============================================================ + # Construct cameras + # ============================================================ + + yaws = [] + pitches = [] + + for view_index in range(num_views): + yaw, pitch = sphere_hammersley_sequence( + view_index, + num_views, + ) + + yaws.append(yaw) + pitches.append(pitch) + + yaws = torch.as_tensor( + yaws, + device=device, + dtype=torch.float32, + ) + + pitches = torch.as_tensor( + pitches, + device=device, + dtype=torch.float32, + ) + + radius = 2.0 + + fov = torch.deg2rad( + torch.tensor( + 40.0, + device=device, + dtype=torch.float32, + ) + ) + + projection = utils3d.torch.perspective_from_fov_xy( + fov, + fov, + 1, + 3, + ) + + target = torch.zeros( + 3, + device=device, + dtype=torch.float32, + ) + + up = torch.tensor( + [0.0, 0.0, 1.0], + device=device, + dtype=torch.float32, + ) + + views = [] + + for yaw, pitch in zip(yaws, pitches): + origin = torch.stack( + [ + torch.sin(yaw) * torch.cos(pitch), + torch.cos(yaw) * torch.cos(pitch), + torch.sin(pitch), + ], + dim=0, + ).float() * radius + + view = utils3d.torch.view_look_at( + origin, + target, + up, + ) + + views.append(view) + + views = torch.stack( + views, + dim=0, + ) + + # ============================================================ + # Rasterize + # ============================================================ + + visibility = torch.zeros( + faces.shape[0], + dtype=torch.int32, + device=device, + ) + + rastctx = utils3d.torch.RastContext( + backend='cuda' + ) + + for view_index in tqdm( + range(views.shape[0]), + total=views.shape[0], + disable=not verbose, + desc='Rasterizing', + ): + view = views[view_index] + + buffers = utils3d.torch.rasterize_triangle_faces( + rastctx, + verts[None], + faces, + resolution, + resolution, + view=view, + projection=projection, + ) + + valid_pixels = buffers['mask'][0] > 0.95 + + face_id = ( + buffers['face_id'][0][valid_pixels] - 1 + ) + + if face_id.numel() == 0: + continue + + face_id = torch.unique( + face_id + ).long() + + # 避免异常背景 ID + face_id = face_id[ + (face_id >= 0) + & (face_id < faces.shape[0]) + ] + + if face_id.numel() > 0: + visibility[face_id] += 1 + + visibility = ( + visibility.float() + / float(num_views) + ) + + # ============================================================ + # Mincut: construct outer faces + # ============================================================ + + ( + edges, + face2edge, + edge_degrees, + ) = utils3d.torch.compute_edges(faces) + + boundary_edge_indices = torch.nonzero( + edge_degrees == 1, + as_tuple=False, + ).reshape(-1) + + connected_components = ( + utils3d.torch.compute_connected_components( + faces, + edges, + face2edge, + ) + ) + + outer_face_mask = torch.zeros( + faces.shape[0], + dtype=torch.bool, + device=faces.device, + ) + + for component in connected_components: + component_visibility = visibility[component] + + if component_visibility.numel() == 0: + continue + + visibility_threshold = min( + max( + component_visibility + .quantile(0.75) + .item(), + 0.25, + ), + 0.5, + ) + + outer_face_mask[component] = ( + component_visibility + > visibility_threshold + ) + + outer_face_indices = torch.nonzero( + outer_face_mask, + as_tuple=False, + ).reshape(-1) + + # ============================================================ + # Construct inner faces + # ============================================================ + + inner_face_indices = torch.nonzero( + visibility == 0, + as_tuple=False, + ).reshape(-1) + + if verbose: + tqdm.write( + f'Found {inner_face_indices.shape[0]} ' + f'invisible faces' + ) + + # 保持原函数行为： + # 若不存在不可见面，则不执行 mincut 和最终 MeshFix。 + if inner_face_indices.shape[0] == 0: + return finalize_output(verts, faces) + + # ============================================================ + # Construct dual graph + # ============================================================ + + ( + dual_edges, + dual_edge2edge, + ) = utils3d.torch.compute_dual_graph( + face2edge + ) + + dual_edge_vertices = edges[ + dual_edge2edge + ] + + dual_edges_weights = torch.linalg.vector_norm( + verts[dual_edge_vertices[:, 0]] + - verts[dual_edge_vertices[:, 1]], + dim=1, + ) + + if verbose: + tqdm.write( + f'Dual graph: {dual_edges.shape[0]} edges' + ) + + # ============================================================ + # Solve mincut + # ============================================================ + + graph = igraph.Graph() + + graph.add_vertices( + faces.shape[0] + ) + + if dual_edges.numel() > 0: + graph.add_edges( + dual_edges.detach().cpu().numpy() + ) + + graph.es['weight'] = ( + dual_edges_weights + .detach() + .cpu() + .numpy() + .tolist() + ) + else: + graph.es['weight'] = [] + + graph.add_vertex('s') + graph.add_vertex('t') + + # Connect invisible faces to source + if inner_face_indices.numel() > 0: + graph.add_edges( + [ + (int(face_index), 's') + for face_index + in inner_face_indices.tolist() + ], + attributes={ + 'weight': np.ones( + inner_face_indices.shape[0], + dtype=np.float32, + ) + }, + ) + + # Connect outer faces to target + if outer_face_indices.numel() > 0: + graph.add_edges( + [ + (int(face_index), 't') + for face_index + in outer_face_indices.tolist() + ], + attributes={ + 'weight': np.ones( + outer_face_indices.shape[0], + dtype=np.float32, + ) + }, + ) + + mincut_weights = ( + np.asarray( + graph.es['weight'], + dtype=np.float64, + ) + * 1000.0 + ).tolist() + + cut = graph.mincut( + 's', + 't', + mincut_weights, + ) + + remove_face_indices = torch.tensor( + [ + vertex_id + for vertex_id in cut.partition[0] + if vertex_id < faces.shape[0] + ], + dtype=torch.long, + device=faces.device, + ) + + if verbose: + tqdm.write( + 'Mincut solved, start checking the cut' + ) + + # ============================================================ + # Validate each connected component selected by mincut + # ============================================================ + + if remove_face_indices.numel() > 0: + to_remove_components = ( + utils3d.torch.compute_connected_components( + faces[remove_face_indices] + ) + ) + else: + to_remove_components = [] + + if debug: + tqdm.write( + 'Number of connected components of the cut: ' + f'{len(to_remove_components)}' + ) + + valid_remove_components = [] + cutting_edges = [] + + for component in to_remove_components: + selected_face_indices = remove_face_indices[ + component + ] + + component_visibility_median = visibility[ + selected_face_indices + ].median() + + if debug: + tqdm.write( + 'visibility_median: ' + f'{component_visibility_median.item()}' + ) + + # 可见性过高，不删除 + if component_visibility_median > 0.25: + continue + + # 检查切割形成的新边界环是否足够小 + ( + component_edge_indices, + component_edge_counts, + ) = torch.unique( + face2edge[selected_face_indices], + return_counts=True, + ) + + component_boundary_edge_indices = ( + component_edge_indices[ + component_edge_counts == 1 + ] + ) + + component_new_boundary_edge_indices = ( + component_boundary_edge_indices[ + ~torch.isin( + component_boundary_edge_indices, + boundary_edge_indices, + ) + ] + ) + + if component_new_boundary_edge_indices.numel() > 0: + new_boundary_components = ( + utils3d.torch + .compute_edge_connected_components( + edges[ + component_new_boundary_edge_indices + ] + ) + ) + + component_areas = [] + + for edge_component in new_boundary_components: + edge_indices = ( + component_new_boundary_edge_indices[ + edge_component + ] + ) + + edge_vertex_indices = edges[ + edge_indices + ] + + edge_center = verts[ + edge_vertex_indices + ].mean(dim=1).mean(dim=0) + + edge_vector_1 = ( + verts[edge_vertex_indices[:, 0]] + - edge_center + ) + + edge_vector_2 = ( + verts[edge_vertex_indices[:, 1]] + - edge_center + ) + + loop_area = ( + torch.linalg.vector_norm( + torch.cross( + edge_vector_1, + edge_vector_2, + dim=-1, + ), + dim=1, + ).sum() + * 0.5 + ) + + component_areas.append(loop_area) + + if debug: + cutting_edges.append( + component_new_boundary_edge_indices + ) + + tqdm.write( + 'Area of the cutting loop: ' + f'{[float(area.item()) for area in component_areas]}' + ) + + if any( + float(area.item()) > max_hole_size + for area in component_areas + ): + continue + + valid_remove_components.append( + component + ) + + # ============================================================ + # Debug outputs + # ============================================================ + + if debug: + face_centers = ( + verts[faces] + .mean(dim=1) + .detach() + .cpu() + .numpy() + ) + + visualization_dual_edges = ( + dual_edges.detach().cpu().numpy() + ) + + visualization_colors = np.zeros( + (faces.shape[0], 3), + dtype=np.uint8, + ) + + visualization_colors[ + inner_face_indices.detach().cpu().numpy() + ] = [0, 0, 255] + + visualization_colors[ + outer_face_indices.detach().cpu().numpy() + ] = [0, 255, 0] + + visualization_colors[ + remove_face_indices.detach().cpu().numpy() + ] = [255, 0, 255] + + if len(valid_remove_components) > 0: + valid_component_indices = torch.cat( + valid_remove_components + ) + + visualization_colors[ + remove_face_indices[ + valid_component_indices + ].detach().cpu().numpy() + ] = [255, 0, 0] + + utils3d.io.write_ply( + 'dbg_dual.ply', + face_centers, + edges=visualization_dual_edges, + vertex_colors=visualization_colors, + ) + + if len(cutting_edges) > 0: + debug_cutting_edge_indices = torch.cat( + cutting_edges + ) + + visualization_vertices = ( + verts.detach().cpu().numpy() + ) + + visualization_edges = ( + edges[debug_cutting_edge_indices] + .detach() + .cpu() + .numpy() + ) + + utils3d.io.write_ply( + 'dbg_cut.ply', + visualization_vertices, + edges=visualization_edges, + ) + + # ============================================================ + # Remove validated mincut components + # ============================================================ + + if len(valid_remove_components) > 0: + valid_component_indices = torch.cat( + valid_remove_components + ) + + remove_face_indices = remove_face_indices[ + valid_component_indices + ] + + if remove_face_indices.numel() > 0: + keep_face_mask = torch.ones( + faces.shape[0], + dtype=torch.bool, + device=faces.device, + ) + + keep_face_mask[ + remove_face_indices + ] = False + + faces = faces[ + keep_face_mask + ] + + # 该函数只删除不再被任何面引用的顶点， + # 不会主动删除仍存在的面。 + ( + faces, + verts, + ) = utils3d.torch.remove_unreferenced_vertices( + faces, + verts, + ) + + if verbose: + tqdm.write( + f'Removed ' + f'{int((~keep_face_mask).sum().item())} ' + f'faces by mincut' + ) + elif verbose: + tqdm.write( + 'Removed 0 faces by mincut' + ) + elif verbose: + tqdm.write( + 'Removed 0 faces by mincut' + ) + + # ============================================================ + # MeshFix hole filling + # ============================================================ + + verts, faces = _run_meshfix_hole_filling( + verts=verts, + faces=faces, + max_hole_nbe=max_hole_nbe, + ) + + # 恢复提前保存的退化面 + return finalize_output( + verts, + faces, + ) + +if __name__ == '__main__': + import trimesh + input_mesh_path = '/root/junkai_lin/mul_res_diffusion/pretrained_model/gobjaverse_fix3_octant_factorized_1_to_1024_ce_altpc_decoded_deg_w1_hg246/single_layer_ep28_res32to64_s20/000-002_21538150cde042bd87a0b0a75935c5a1_facenum_8572/pred_consensus.obj' + output_mesh_path = '/root/junkai_lin/mul_res_diffusion/pretrained_model/gobjaverse_fix3_octant_factorized_1_to_1024_ce_altpc_decoded_deg_w1_hg246/single_layer_ep28_res32to64_s20/000-002_21538150cde042bd87a0b0a75935c5a1_facenum_8572/pred_consensus_fix2.obj' + mesh = trimesh.load(input_mesh_path, force='mesh') + verts = torch.tensor(mesh.vertices, device='cuda', dtype=torch.float32) + faces = torch.tensor(mesh.faces, device='cuda', dtype=torch.int32) + new_verts, new_faces = _fill_holes(verts, faces, debug=True, verbose=True, only_fill_holes=True, keep_degenerate_faces=True) + new_mesh = trimesh.Trimesh(vertices=new_verts.cpu().numpy(), faces=new_faces.cpu().numpy()) + new_mesh.export(output_mesh_path) + # from pathlib import Path + + # import torch + # import trimesh + + # input_dir = Path("/root/junkai_lin/experiment/render/result_control/t-flow-ab") + # output_dir = Path("/root/junkai_lin/experiment/render/result_control/t-flow-ab_fix_holes") + # output_dir.mkdir(parents=True, exist_ok=True) + + # obj_paths = sorted(input_dir.glob("*.obj")) + + # print(f"Found {len(obj_paths)} obj files in {input_dir}") + + # for obj_path in obj_paths: + # print(f"\nProcessing: {obj_path.name}") + + # try: + # mesh = trimesh.load(str(obj_path), force="mesh") + + # verts = torch.tensor(mesh.vertices, device="cuda", dtype=torch.float32) + # faces = torch.tensor(mesh.faces, device="cuda", dtype=torch.int32) + + # new_verts, new_faces = _fill_holes( + # verts, + # faces, + # debug=True, + # verbose=True, + # only_fill_holes=True, + # max_hole_nbe=64, + # resolution=256, + # max_hole_size=0.12, + # num_views=1000 + # ) + + # new_mesh = trimesh.Trimesh( + # vertices=new_verts.detach().cpu().numpy(), + # faces=new_faces.detach().cpu().numpy(), + # process=False, + # ) + + # out_path = output_dir / f"{obj_path.stem}_fix2.obj" + # new_mesh.export(str(out_path)) + + # print(f"Saved: {out_path}") + + # except Exception as e: + # print(f"Failed: {obj_path.name}") + # print(f"Error: {e}") \ No newline at end of file