Refactor skin weight calculations to handle division by zero and ensure valid index access in Exporter and SAMPart3DDataset classes

src/data/exporter.py

@@ -290,11 +290,15 @@ class Exporter():
         vertex_group_reweight = skin[np.arange(skin.shape[0])[..., None], argsorted]
         if group_per_vertex == -1:
             group_per_vertex = vertex_group_reweight.shape[-1]
+        # Ensure we don't access more columns than available
+        max_groups = min(group_per_vertex, argsorted.shape[1])
         if not do_not_normalize:
-            vertex_group_reweight = vertex_group_reweight / vertex_group_reweight[..., :group_per_vertex].sum(axis=1)[...,None]
+            vertex_group_sum = vertex_group_reweight[..., :max_groups].sum(axis=1)[..., None]
+            vertex_group_sum = np.where(vertex_group_sum == 0, 1.0, vertex_group_sum)  # Avoid division by zero
+            vertex_group_reweight = vertex_group_reweight / vertex_group_sum
 
         for v, w in enumerate(skin):
-            for ii in range(group_per_vertex):
+            for ii in range(max_groups):
                 i = argsorted[v, ii]
                 if i >= J:
                     continue

src/inference/merge.py

@@ -262,7 +262,11 @@ def make_armature(
     
     argsorted = np.argsort(-skin, axis=1)
     vertex_group_reweight = skin[np.arange(skin.shape[0])[..., None], argsorted]
-    vertex_group_reweight = vertex_group_reweight / vertex_group_reweight[..., :group_per_vertex].sum(axis=1)[...,None]
+    # Handle division by zero and limit to group_per_vertex
+    max_groups = min(group_per_vertex, vertex_group_reweight.shape[1])
+    vertex_group_sum = vertex_group_reweight[..., :max_groups].sum(axis=1)[..., None]
+    vertex_group_sum = np.where(vertex_group_sum == 0, 1.0, vertex_group_sum)  # Avoid division by zero
+    vertex_group_reweight = vertex_group_reweight / vertex_group_sum
     vertex_group_reweight = np.nan_to_num(vertex_group_reweight)
     tree = cKDTree(vertices)
     for ob in objects:
@@ -286,7 +290,9 @@ def make_armature(
         _, index = tree.query(n_vertices)
 
         for v, co in enumerate(tqdm(n_vertices)):
-            for ii in range(group_per_vertex):
+            # Ensure we don't access more columns than available in argsorted
+            max_groups = min(group_per_vertex, argsorted.shape[1])
+            for ii in range(max_groups):
                 i = argsorted[index[v], ii]
                 if i >= len(names):
                     continue

src/model/michelangelo/models/modules/transformer_blocks.py

@@ -35,13 +35,13 @@ def init_linear(l, stddev):
         nn.init.constant_(l.bias, 0.0)
 
 def flash_attention(q, k, v):
-    with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-        out = F.scaled_dot_product_attention(q, k, v)
-        out = out.transpose(1, 2)
-        # print("use flash atten 2")   
+    # Use torch.nn.functional.scaled_dot_product_attention directly instead of deprecated context manager
+    q = q.transpose(1, 2)
+    k = k.transpose(1, 2)
+    v = v.transpose(1, 2)
+    out = F.scaled_dot_product_attention(q, k, v)
+    out = out.transpose(1, 2)
+    # print("use flash atten 2")   
 
     return out
 

src/model/pointcept/datasets/dataset_render_16views.py

@@ -387,13 +387,17 @@ class SAMPart3DDataset16Views(Dataset):
             if per_pixel_index.shape[-1] == 1:
                 per_pixel_mask = per_pixel_index.squeeze()
             else:
+                # Clamp random_index to valid range to prevent out of bounds error
+                random_index_clamped = torch.clamp(random_index.unsqueeze(-1), 0, per_pixel_index.shape[1] - 1)
                 per_pixel_mask = torch.gather(
-                    per_pixel_index, 1, random_index.unsqueeze(-1)
+                    per_pixel_index, 1, random_index_clamped
                 ).squeeze()
+                # Clamp the previous index to valid range as well
+                prev_index_clamped = torch.clamp(random_index.unsqueeze(-1) - 1, 0, per_pixel_index.shape[1] - 1)
                 per_pixel_mask_ = torch.gather(
                     per_pixel_index,
                     1,
-                    torch.max(random_index.unsqueeze(-1) - 1, torch.Tensor([0]).int()),
+                    prev_index_clamped,
                 ).squeeze()
 
             mask_id[i : i + npximg] = per_pixel_mask.to(self.device)

src/model/unirig_skin.py

@@ -356,7 +356,7 @@ class UniRigSkin(ModelSpec):
             ptv3_input = {
                 'coord': vertices.reshape(-1, 3),
                 'feat': feat.reshape(-1, 9),
-                'offset': torch.tensor(batch['offset']),
+                'offset': batch['offset'].detach().clone(),
                 'grid_size': self.grid_size,
             }
             if not self.training:
@@ -420,7 +420,7 @@ class UniRigSkin(ModelSpec):
                 input_features = attn_weight[i, :, :num_bones[i], :].reshape(-1, attn_weight.shape[-1])
                 
                 pred = self.skinweight_pred(input_features).reshape(cur_N, num_bones[i])
-                skin_pred[i, :, :num_bones[i]] = F.softmax(pred)
+                skin_pred[i, :, :num_bones[i]] = F.softmax(pred, dim=-1)
             skin_pred_list.append(skin_pred)
         skin_pred_list = torch.cat(skin_pred_list, dim=1)
         for i in range(B):
@@ -437,4 +437,5 @@ class UniRigSkin(ModelSpec):
             outputs = []
             for i in range(skin_pred.shape[0]):
                 outputs.append(skin_pred[i, :, :num_bones[i]])
+            
             return outputs

src/system/skin.py

@@ -118,7 +118,7 @@ class SkinWriter(BasePredictionWriter):
             J = num_bones[id]
             F = num_faces[id]
             o_vertices = vertices[id, :N]
-
+            
             _parents = parents_list[id]
             parents = []
             for i in range(J):
@@ -156,36 +156,34 @@ class SkinWriter(BasePredictionWriter):
             raw_data = RawSkin(skin=skin_pred, vertices=sampled_vertices[id], joints=joints[id, :J])
             if self.export_npz is not None:
                 raw_data.save(path=make_path(self.export_npz, 'npz'))
+
             if self.export_fbx is not None:
-                try:
-                    exporter = Exporter()
-                    names = RawData.load(path=os.path.join(paths[id], data_names[id])).names
-                    if names is None:
-                        names = [f"bone_{i}" for i in range(J)]
-                    if self.user_mode:
-                        if self.output_name is not None:
-                            path = self.output_name
-                        else:
-                            path = make_path(self.save_name, 'fbx', trim=True)
+                exporter = Exporter()
+                names = RawData.load(path=os.path.join(paths[id], data_names[id])).names
+                if names is None:
+                    names = [f"bone_{i}" for i in range(J)]
+                if self.user_mode:
+                    if self.output_name is not None:
+                        path = self.output_name
                     else:
-                        path = make_path(self.export_fbx, 'fbx')
-                    exporter._export_fbx(
-                        path=path,
-                        vertices=o_vertices,
-                        joints=joints[id, :J],
-                        skin=skin_resampled,
-                        parents=parents,
-                        names=names,
-                        faces=faces[id, :F],
-                        group_per_vertex=4,
-                        tails=tails[id, :J],
-                        use_extrude_bone=False,
-                        use_connect_unique_child=False,
-                        # do_not_normalize=True,
-                    )
-                except Exception as e:
-                    print(str(e))
-    
+                        path = make_path(self.save_name, 'fbx', trim=True)
+                else:
+                    path = make_path(self.export_fbx, 'fbx')
+                exporter._export_fbx(
+                    path=path,
+                    vertices=o_vertices,
+                    joints=joints[id, :J],
+                    skin=skin_resampled,
+                    parents=parents,
+                    names=names,
+                    faces=faces[id, :F],
+                    group_per_vertex=4,
+                    tails=tails[id, :J],
+                    use_extrude_bone=False,
+                    use_connect_unique_child=False,
+                    # do_not_normalize=True,
+                )
+
     def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
         self._epoch += 1