update zerogpu

LHM/models/rendering/gs_renderer.py

@@ -829,14 +829,15 @@ class GS3DRenderer(nn.Module):
         # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
         screenspace_points = (
             torch.zeros_like(
-                gs.xyz, dtype=gs.xyz.dtype, requires_grad=False, device=self.device
+                gs.xyz, dtype=gs.xyz.dtype, requires_grad=True, device=self.device
             )
             + 0
         )
-        # try:
-        #     screenspace_points.retain_grad()
-        # except:
-        #     pass
+
+        try:
+            screenspace_points.retain_grad()
+        except:
+            pass
 
         bg_color = background_color
         # Set up rasterization configuration
@@ -855,7 +856,7 @@ class GS3DRenderer(nn.Module):
             sh_degree=self.sh_degree,
             campos=viewpoint_camera.camera_center,
             prefiltered=False,
-            debug=False,
+            debug=True,
         )
 
         rasterizer = GaussianRasterizer(raster_settings=raster_settings)
@@ -883,11 +884,13 @@ class GS3DRenderer(nn.Module):
             colors_precomp = None
             shs = gs.shs
 
+        # print(shs, colors_precomp)
+        # print(means3D.device, means2D.device, opacity.device, rotations.device, self.device)
+        # print(means3D.dtype, means2D.dtype, rotations.dtype, opacity.dtype)
+        # print(means3D.shape, means2D.shape, rotations.shape, opacity.shape)
         # Rasterize visible Gaussians to image, obtain their radii (on screen).
         # NOTE that dadong tries to regress rgb not shs
         # with torch.autocast(device_type=self.device.type, dtype=torch.float32):
-        print(means3D.device, means2D.device, colors_precomp.device, opacity.device, rotations.device, self.device)
-        print(means3D.dtype, means2D.dtype, colors_precomp.dtype)
         rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
             means3D=means3D,
             means2D=means2D,
@@ -937,7 +940,7 @@ class GS3DRenderer(nn.Module):
         #         ret["comp_mask"] = rendered_mask.permute(1, 2, 0)
 
         return ret
-    @torch.no_grad()
+    
     def animate_gs_model(
         self, gs_attr: GaussianAppOutput, query_points, smplx_data, debug=False
     ):
@@ -1066,7 +1069,7 @@ class GS3DRenderer(nn.Module):
                 gs_list.append(gs_copy)
 
         return gs_list, cano_gs_list
-    @torch.no_grad()
+    
     def forward_gs_attr(self, x, query_points, smplx_data, debug=False, x_fine=None):
         """
         x: [N, C] Float[Tensor, "Np Cp"],
@@ -1083,7 +1086,7 @@ class GS3DRenderer(nn.Module):
         gs_attr: GaussianAppOutput = self.gs_net(x, query_points, x_fine)
 
         return gs_attr
-    @torch.no_grad()
+    
     def get_query_points(self, smplx_data, device):
         with torch.no_grad():
             with torch.autocast(device_type=device.type, dtype=torch.float32):
@@ -1115,7 +1118,7 @@ class GS3DRenderer(nn.Module):
             pcl_embed.to(dtype=latent_feat.dtype), latent_feat, extra_info
         )
         return gs_feats
-    @torch.no_grad()
+    
     def query_latent_feat(
         self,
         positions: Float[Tensor, "*B N1 3"],
@@ -1140,7 +1143,7 @@ class GS3DRenderer(nn.Module):
             )
 
         return gs_feats, positions, smplx_data
-    @torch.no_grad()
+    
     def forward_single_batch(
         self,
         gs_list: list[GaussianModel],
@@ -1183,7 +1186,7 @@ class GS3DRenderer(nn.Module):
 
         return out
 
-    @torch.no_grad()
+    
     def forward_cano_batch(
         self,
         gs_list: list[GaussianModel],
@@ -1289,7 +1292,7 @@ class GS3DRenderer(nn.Module):
                     :, vidx : vidx + 1
                 ]  # e.g. body_pose: [1, N_v, 21, 3] -> [1, 1, 21, 3]
         return smpl_data_single_view
-    @torch.no_grad()
+    
     def forward_gs(
         self,
         gs_hidden_features: Float[Tensor, "B Np Cp"],
@@ -1397,7 +1400,7 @@ class GS3DRenderer(nn.Module):
         )  # [B, NV, H, W, 3] -> [B, NV, 1, H, W]
 
         return out
-    @torch.no_grad()
+    
     def forward(
         self,
         gs_hidden_features: Float[Tensor, "B Np Cp"],