new

models/lrm/models/geometry/render/neural_render.py

@@ -7,6 +7,7 @@
 # license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
 
 import torch
+import spaces
 import torch.nn.functional as F
 import nvdiffrast.torch as dr
 from . import Renderer
@@ -66,7 +67,7 @@ def compute_vertex_normal(v_pos, t_pos_idx):
 
     return v_nrm
 
-
+@spaces.GPU
 class NeuralRender(Renderer):
     def __init__(self, device='cuda', camera_model=None):
         super(NeuralRender, self).__init__()
@@ -146,6 +147,7 @@ class NeuralRender(Renderer):
     #  - Single light
     #  - Single material
     # ==============================================================================================
+    @spaces.GPU
     def render_layer(
             self,
             rast,
@@ -189,6 +191,7 @@ class NeuralRender(Renderer):
 
         return gb_pos, gb_normal
 
+    @spaces.GPU
     def render_mesh(
             self,
             mesh_v_pos_bxnx3,
@@ -245,6 +248,7 @@ class NeuralRender(Renderer):
         # normal = torch.lerp(torch.zeros_like(normal), (normal + 1.0) / 2.0, hard_mask.float())      # black background
         return ori_mesh_feature, antialias_mask, hard_mask, rast, v_pos_clip, mask_pyramid, depth, normal, gb_normal
     
+    @spaces.GPU
     def render_mesh_light(
             self,
             mesh_v_pos_bxnx3,