Update modeling_wavtokenizer.py

modeling_wavtokenizer.py

@@ -39,28 +39,27 @@ def convert_audio(wav, sr, target_sr, target_channels=1):
 
 
 # =============================================================================
-# Weight-Normalized Conv1d (matching checkpoint's weight_g/weight_v structure)
+# Weight-Normalized Conv1d (using parametrizations API to match checkpoint)
 # =============================================================================
 
 class WNConv1d(nn.Module):
-    """Weight-normalized Conv1d matching checkpoint structure with weight_g/weight_v."""
+    """Weight-normalized Conv1d using parametrizations API to match checkpoint structure."""
     def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
         super().__init__()
-        self.conv = nn.utils.weight_norm(
-            nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
-        )
+        conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
+        # Use parametrizations API (PyTorch 2.0+) to match checkpoint naming
+        self.conv = nn.utils.parametrizations.weight_norm(conv)
     
     def forward(self, x):
         return self.conv(x)
 
 
 class WNConvTranspose1d(nn.Module):
-    """Weight-normalized ConvTranspose1d."""
+    """Weight-normalized ConvTranspose1d using parametrizations API."""
     def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True):
         super().__init__()
-        self.convtr = nn.utils.weight_norm(
-            nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride, padding, output_padding, groups, bias)
-        )
+        convtr = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride, padding, output_padding, groups, bias)
+        self.convtr = nn.utils.parametrizations.weight_norm(convtr)
     
     def forward(self, x):
         return self.convtr(x)