mul / modules /octree_decode.py
Upload folder using huggingface_hub
90ca518 verified 16 days ago
26.5 kB
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)