models/vae.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):
        super().__init__()
        self.device = device
        self.dtype = dtype
        
        self.vae = TripoSGVAEModel.from_pretrained(
            weights_dir, subfolder="vae",
        ).to(self.device).to(torch.float32) # 内部运算保持 float32
        
        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 = 51200, seed: int = 42) -> torch.Tensor:
        # 0. 基础清洗：防止输入数据包含 inf 或极大的离群点
        vertices = torch.clamp(vertices, -1e6, 1e6)
        vertices = torch.nan_to_num(vertices, nan=0.0, posinf=1e6, neginf=-1e6)
        
        verts_np = vertices.detach().cpu().numpy()
        faces_np = faces.detach().cpu().numpy()
        
        # 1. 基础修复 (保留原始形状)
        mesh = trimesh.Trimesh(vertices=verts_np, faces=faces_np, process=True)
        if mesh.is_empty: return None
        
        # 2. 原始比例归一化 (关键点：不要改变尺度比例)
        max_edge = np.max(mesh.bounding_box.extents)
        scale = 2.0 / (max_edge + 1e-8)
        # 使用原始中心位置，不要做任何裁剪
        normalized_vertices = (mesh.vertices - mesh.center_mass) * scale
        
        norm_mesh = trimesh.Trimesh(vertices=normalized_vertices, faces=mesh.faces, process=False)
        
        # 3. 采样 (使用 trimesh 原生 sampler)
        if seed is not None:
            np.random.seed(seed)
        surface_points, face_indices = trimesh.sample.sample_surface(norm_mesh, count=num_samples)
        surface_normals = norm_mesh.face_normals[face_indices]
        
        # 4. 最后防线：对采样结果进行 NaN 检查，如果极个别 NaN，用 0 替换，而不改变整体点云
        surface_points = np.nan_to_num(surface_points, nan=0.0)
        surface_normals = np.nan_to_num(surface_normals, nan=0.0)
        
        # 转换为 float32 保证 VAE 计算精度
        surface_tensor = torch.cat([torch.from_numpy(surface_points), torch.from_numpy(surface_normals)], dim=-1).float()
        
        return surface_tensor.unsqueeze(0).to(self.device)


    @torch.no_grad()
    def encode_mesh(
        self,
        vertices,
        faces,
        num_samples: int = 51200,
        surface_seed: int = 42,
        vae_seed: int = 42,
    ) -> torch.Tensor:
        surface_tensor = self.preprocess_mesh(
            vertices,
            faces,
            num_samples=num_samples,
            seed=surface_seed,
        )
        if surface_tensor is None:
            return None

        encoder_output = self.vae.encode(surface_tensor, seed=vae_seed)
        latent = encoder_output.latent_dist.mode()
        return latent.to(self.dtype)
    
    

    @torch.no_grad()
    def decode_latent(self, latent_sample: torch.Tensor, sampled_points: torch.Tensor) -> torch.Tensor:
        # 解码时保持 float32，避免计算溢出
        return self.vae.decode(latent_sample.to(torch.float32), sampled_points=sampled_points.to(torch.float32)).sample