update

models/DiMeR/models/DiMeR.py

@@ -17,6 +17,7 @@ import torch
 import torch.nn as nn
 import nvdiffrast.torch as dr
 from einops import rearrange, repeat
+import spaces
 
 from .encoder.dino_wrapper import DinoWrapper
 from .decoder.transformer import TriplaneTransformer
@@ -31,32 +32,8 @@ from PIL import Image
 import torchvision.utils as vutils
 from models.DiMeR.utils import render
 from models.DiMeR.utils.render_utils import rotate_x, rotate_y
-import open3d as o3d
 from .geometry.rep_3d import mesh
 
-def save_sdf_as_ply(positions, sdf_values, filename, number_samples=100000):
-    # 归一化 SDF 值到 [0, 1] 之间
-    sdf_normalized = (sdf_values - sdf_values.min()) / (sdf_values.max() - sdf_values.min())
-    sdf_normalized = sdf_normalized.squeeze().cpu().numpy()
-
-    if positions.shape[0] > number_samples:
-        indices = np.random.choice(positions.shape[0], number_samples, replace=False)
-        positions = positions[indices]
-        sdf_normalized = sdf_normalized[indices]
-
-    # 创建颜色映射 (从蓝色到红色)
-    colors = np.zeros((sdf_normalized.shape[0], 3), dtype=np.float64)
-    colors[:, 0] = sdf_normalized  # Red channel
-    colors[:, 2] = 1.0 - sdf_normalized  # Blue channel
-
-    # 创建 Open3D 点云对象
-    point_cloud = o3d.geometry.PointCloud()
-    point_cloud.points = o3d.utility.Vector3dVector(positions.cpu().numpy())
-    point_cloud.colors = o3d.utility.Vector3dVector(colors)
-
-    # 保存为 PLY 文件
-    o3d.io.write_point_cloud(filename, point_cloud)
-
 class DiMeR(nn.Module):
     """
     Full model of the large reconstruction model.
@@ -106,7 +83,7 @@ class DiMeR(nn.Module):
             samples_per_ray=rendering_samples_per_ray,
             use_deformation_weight=use_deformation_weight,
         )
-
+    @spaces.GPU
     def init_flexicubes_geometry(self, device, fovy=50.0):
         camera = PerspectiveCamera(fovy=fovy, device=device)
         renderer = NeuralRender(device, camera_model=camera)