Buffer

third_party/MMAudio/mmaudio/ext/autoencoder/vae.py

@@ -75,11 +75,16 @@ class VAE(nn.Module):
         super().__init__()
 
         if data_dim == 80:
-            self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_80D, dtype=torch.float32))
-            self.data_std = nn.Buffer(torch.tensor(DATA_STD_80D, dtype=torch.float32))
+            # self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_80D, dtype=torch.float32))
+            # self.data_std = nn.Buffer(torch.tensor(DATA_STD_80D, dtype=torch.float32))
+            self.register_buffer("data_mean", torch.tensor(DATA_MEAN_80D, dtype=torch.float32))
+            self.register_buffer("data_std", torch.tensor(DATA_STD_80D, dtype=torch.float32))
+
         elif data_dim == 128:
-            self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32))
-            self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32))
+            # self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32))
+            # self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32))
+            self.register_buffer("data_mean", torch.tensor(DATA_MEAN_128D, dtype=torch.float32))
+            self.register_buffer("data_std", torch.tensor(DATA_STD_128D, dtype=torch.float32))
 
         self.data_mean = self.data_mean.view(1, -1, 1)
         self.data_std = self.data_std.view(1, -1, 1)

third_party/MMAudio/mmaudio/model/embeddings.py

@@ -21,10 +21,17 @@ class TimestepEmbedder(nn.Module):
         assert dim % 2 == 0, 'dim must be even.'
 
         with torch.autocast('cuda', enabled=False):
-            self.freqs = nn.Buffer(
+            # self.freqs = nn.Buffer(
+            #     1.0 / (10000**(torch.arange(0, frequency_embedding_size, 2, dtype=torch.float32) /
+            #                    frequency_embedding_size)),
+            #     persistent=False)
+            self.register_buffer(
+                "freqs",
                 1.0 / (10000**(torch.arange(0, frequency_embedding_size, 2, dtype=torch.float32) /
-                               frequency_embedding_size)),
-                persistent=False)
+                            frequency_embedding_size)),
+                persistent=False
+            )
+
             freq_scale = 10000 / max_period
             self.freqs = freq_scale * self.freqs
 

third_party/MMAudio/mmaudio/model/networks.py

@@ -166,8 +166,11 @@ class MMAudio(nn.Module):
                                           self._clip_seq_len,
                                           device=self.device)
 
-        self.latent_rot = nn.Buffer(latent_rot, persistent=False)
-        self.clip_rot = nn.Buffer(clip_rot, persistent=False)
+        # self.latent_rot = nn.Buffer(latent_rot, persistent=False)
+        # self.clip_rot = nn.Buffer(clip_rot, persistent=False)
+        self.register_buffer("latent_rot", latent_rot, persistent=False)
+        self.register_buffer("clip_rot", clip_rot, persistent=False)
+
 
     def update_seq_lengths(self, latent_seq_len: int, clip_seq_len: int, sync_seq_len: int) -> None:
         self._latent_seq_len = latent_seq_len