models/vae_v1.py

import os
import sys

triposg_root = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/TripoSG"
if triposg_root not in sys.path:
    sys.path.append(triposg_root)

import torch
import torch.nn as nn
import numpy as np
import trimesh

# 现在可以完美导入了
from triposg.models.autoencoders import TripoSGVAEModel


class TripoSGVaeWrapper(nn.Module):
    def __init__(self, weights_dir: str = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/TripoSG/pretrained_weights/TripoSG", device: str = "cuda", dtype: torch.dtype = torch.float16):
        """
        加载预训练且冻结的 TripoSG VAE 模型，集成内置的 Mesh 采样与归一化预处理。
        """
        super().__init__()
        self.device = device
        self.dtype = dtype
        
        # 1. 加载预训练 VAE
        assert os.path.exists(os.path.join(weights_dir, "vae")), f"Error: 找不到 VAE 子文件夹，请确认路径: {weights_dir}"
        self.vae = TripoSGVAEModel.from_pretrained(
            weights_dir,
            subfolder="vae",
        ).to(self.device, dtype=self.dtype)
        
        # 2. 彻底冻结所有 VAE 参数
        self.vae.eval()
        for param in self.vae.parameters():
            param.requires_grad = False

    def preprocess_mesh(self, vertices: torch.Tensor, faces: torch.Tensor, num_samples: int = 204800) -> torch.Tensor:
        # 转换并确保输入合法
        verts_np = vertices.detach().cpu().numpy()
        faces_np = faces.detach().cpu().numpy()
        
        # 修复1：使用 repair 选项去除退化面
        mesh = trimesh.Trimesh(vertices=verts_np, faces=faces_np, process=True) 
        if mesh.is_empty: return None

        # 修复2：更稳健的归一化
        # 增加一个 epsilon 避免除以过小的值
        max_edge = np.max(mesh.bounding_box.extents)
        scale = 2.0 / (max_edge + 1e-6) 
        normalized_vertices = (mesh.vertices - mesh.center_mass) * scale
        
        norm_mesh = trimesh.Trimesh(vertices=normalized_vertices, faces=mesh.faces, process=False)
        
        # 修复3：采样后进行 NaN 清洗
        surface_points, face_indices = trimesh.sample.sample_surface(norm_mesh, count=num_samples)
        surface_normals = norm_mesh.face_normals[face_indices]
        
        # 关键检查：如果采样结果有 NaN，强行替换为 0
        surface_points = np.nan_to_num(surface_points, nan=0.0)
        surface_normals = np.nan_to_num(surface_normals, nan=0.0)
        
        surface_tensor = torch.cat([torch.from_numpy(surface_points), torch.from_numpy(surface_normals)], dim=-1).float()
        
        # 修复4：为了安全，VAE 输入可以使用 float32，但在 VAE 内部使用半精度，
        # 或者确保输入范围严格限制在 [-1, 1] 之间
        return surface_tensor.unsqueeze(0).to(device=self.device, dtype=torch.float16)

    @torch.no_grad()
    def encode_mesh(self, vertices: np.ndarray, faces: np.ndarray, num_samples: int = 204800) -> torch.Tensor:
        """
        输入原始静态 Mesh 的顶点和面片，提取出压缩后的隐向量 (Latent Space)
        """
        surface_tensor = self.preprocess_mesh(vertices, faces, num_samples=num_samples)
        # 编码并采样分布
        latent_sample = self.vae.encode(surface_tensor).latent_dist.sample()
        return latent_sample

    @torch.no_grad()
    def decode_latent(self, latent_sample: torch.Tensor, sampled_points: torch.Tensor) -> torch.Tensor:
        """
        传入隐特征向量以及需要查询的 3D 空间坐标点，解码返回这些坐标点上的 SDF/Occupancy 预测值
        """
        return self.vae.decode(latent_sample, sampled_points=sampled_points).sample