update 1.0.0.post2

DCNv4_op/DCNv4/__init__.py

@@ -0,0 +1,2 @@
+from .functions import DCNv4Function, FlashDeformAttnFunction
+from .modules import DCNv4, FlashDeformAttn

DCNv4_op/DCNv4/functions/flash_deform_attn_func.py

@@ -57,7 +57,7 @@ def findspec_bwd(B, Q, G, C):
         d_stride = 2
     else:
         d_stride = 1
-    
+
     ms = factors(B*Q)
     multiplier = 1
     for m in ms:
@@ -70,7 +70,7 @@ class FlashDeformAttnFunction(Function):
     @staticmethod
     @torch.autocast("cuda", enabled=True, dtype=torch.float16)
     def forward(
-        ctx, value, value_spatial_shapes, value_level_start_index, 
+        ctx, value, value_spatial_shapes, value_level_start_index,
         sampling_loc_attn, im2col_step, K=8
     ):
 

DCNv4_op/DCNv4/modules/dcnv4.py

@@ -139,13 +139,15 @@ class DCNv4(nn.Module):
             self.remove_center
             )
 
-        x = x.view(N, L, -1)
         if self.center_feature_scale:
             center_feature_scale = self.center_feature_scale_module(
                 x, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias)
             center_feature_scale = center_feature_scale[..., None].repeat(
                 1, 1, 1, 1, self.channels // self.group).flatten(-2)
             x = x * (1 - center_feature_scale) + x_proj * center_feature_scale
+
+        x = x.view(N, L, -1)
+
         if not self.without_pointwise:
             x = self.output_proj(x)
         return x

DCNv4_op/DCNv4/modules/flash_deform_attn.py

@@ -136,6 +136,7 @@ class FlashDeformAttn(nn.Module):
             sampling_locations,
             attention_weights,
             self.im2col_step,
+            self.n_points           
         )
         output = self.output_proj(output)
         return output