train to small denoising

Author: ostix360

audioenhancer/constants.py

@@ -3,11 +3,11 @@
 """
 
 SAMPLING_RATE = 44100
-MAX_AUDIO_LENGTH = 5
-BATCH_SIZE = 8
+MAX_AUDIO_LENGTH = 6
+BATCH_SIZE = 3
 EPOCH = 1
 LOGGING_STEPS = 10
-GRADIENT_ACCUMULATION_STEPS = 2
+GRADIENT_ACCUMULATION_STEPS = 3
 SAVE_STEPS = 100
 EVAL_STEPS = 100
 

audioenhancer/dataset/loader.py

@@ -62,7 +62,6 @@ def __init__(
             if os.path.isdir(f)
         ]
 
-        self.codecs += [audio_dir,]
 
         self._pad_length_input = 2 ** math.ceil(math.log2(max_duration * input_freq))
         self._pad_length_output = 2 ** math.ceil(math.log2(max_duration * output_freq))
@@ -174,22 +173,27 @@ def __getitem__(self, index: int) -> tuple:
             base_waveform
         )
 
-        noisy_waveform = torch.zeros_like(encoded_compressed_waveform)
-        for i in range(encoded_compressed_waveform.shape[1]):
-            if random.random() < 0.3:
-                noisy_waveform[:, i] = encoded_compressed_waveform[:, i]
-            else:
-                noisy_waveform[:, i] = encoded_base_waveform[:, i]
+        noise = encoded_base_waveform - encoded_compressed_waveform
+        noise_levels = []
+        noise_level = random.random()
+        encoded_compressed_waveform[:] = encoded_base_waveform[:] + noise[:] * noise_level
+        if noise_level <= 0.10:
+            t_noise_level = 0
+        else:
+            t_noise_level = noise_level - 0.10
+        encoded_base_waveform[:] = encoded_base_waveform[:] + noise[:] * t_noise_level
+        noise_levels.append(t_noise_level)
+
 
         # class_id = [0]
         # if "dac" in codec or "encodec" in codec or "opus" in codec or use_transform:
         #     class_id = [1]
 
-        # class_id = torch.tensor(class_id).cuda()
+        noise_levels = torch.tensor(noise_levels)
 
         return (
-            noisy_waveform,
+            encoded_compressed_waveform,
             encoded_base_waveform,
             base_waveform,
-            # class_id,
+            noise_levels,
         )

audioenhancer/inference.py

@@ -10,16 +10,22 @@
 from einops import rearrange
 
 from audioenhancer.model.audio_ae.model import mamba_model as model
+from audioenhancer.model.audio_ae.model import disc_model as disc_model
 
 
 class Inference:
     def __init__(self, model_path: str, sampling_rate: int):
         self.model = model
+        self.disc_model = disc_model
         self.device = torch.device("cuda")
         self.model = self.model.to(self.device)
         self.model.load_state_dict(torch.load(model_path))
         self.model.eval()
 
+        # self.disc_model = self.disc_model.to(self.device)
+        # self.disc_model.load_state_dict(torch.load("data/model/disc_model_1400.pt"))
+        # self.disc_model.eval()
+
         self._sampling_rate = sampling_rate
 
         autoencoder_path = dac.utils.download(model_type="44khz")
@@ -84,7 +90,14 @@ def inference(self, audio_path: str, chunk_duration: int = 5):
                 c, d = encoded.shape[1], encoded.shape[2]
                 encoded = rearrange(encoded, "b c d t -> b (t c) d")
 
-                pred = self.model(encoded, None)
+                noise_level = torch.tensor([10]).to(self.device)
+                for i in range(1, 15):
+                    pred, _ = self.model(encoded, None, False, None)
+                    # _, logits = self.disc_model(pred, None, True, None)
+                    noise_level -= 1
+                    print(f"Noise level: {noise_level.item()}")
+                    if noise_level == 0:
+                        break
 
                 pred = rearrange(pred, "b (t c) d -> b c d t", c=c, d=d)
                 pred = pred.squeeze(0)

audioenhancer/model/audio_ae/latent.py

@@ -1,6 +1,9 @@
 """This module contains all the process for the latent space of the audio autoencoder."""
+import random
+
 import torch
 from torch import nn
+from torch.nn import functional as F
 
 from audioenhancer.model.audio_ae.expert import Expert
 from audioenhancer.model.audio_ae.mamba import MambaBlock
@@ -76,7 +79,7 @@ class LatentProcessor(nn.Module):
     This module processes the latent space of the audio autoencoder.
     """
 
-    def __init__(self, in_dim: int, out_dim: int, latent_dim, num_layer, num_expert=1):
+    def __init__(self, in_dim: int, out_dim: int, latent_dim, num_layer, noise_grad=1):
         super().__init__()
         self.latent_dim = latent_dim
         self.num_layer = num_layer
@@ -85,9 +88,11 @@ def __init__(self, in_dim: int, out_dim: int, latent_dim, num_layer, num_expert=
         self.in_proj = nn.Linear(in_dim, latent_dim)
 
         self.out_proj = nn.Linear(latent_dim, out_dim)
-        self.num_expert = num_expert
 
         self.mambas = nn.ModuleList([MambaBlock(config) for _ in range(num_layer)])
+        self.unknow_noise = nn.Parameter(torch.randn(latent_dim))
+        self.noise_embed = nn.Embedding(noise_grad, latent_dim)
+        self.noise_head = nn.Linear(latent_dim, noise_grad)
         # self.pre_process = nn.Sequential(
         #     MambaBlock(config),
         #     MambaBlock(config),
@@ -104,14 +109,32 @@ def classify(self, x):
         x = self.pre_process(x)
         return self.classifier(x)
 
-    def forward(self, x, classes):
+    def forward(self, x, noise, gen_noise=False, noise_label=None):
+        bzs = x.size(0)
         h = self.in_proj(x)
+        if noise is not None and not gen_noise:
+            noise = self.noise_embed(noise).reshape(bzs, 1, -1)
+            h = torch.cat([h, noise], dim=1)
+            gen_noise = True
+        else:
+            noise = self.unknow_noise.reshape(1, 1, -1).repeat(bzs, 1, 1)
+            h = torch.cat([h, noise], dim=1)
+
         # h = self.pre_process(h)
         for mamba in self.mambas:
-            h = mamba(h)
+            h = mamba(h, gen_noise=gen_noise)
         # if classes is not None:
         #     return x * classes[:, None, None, 0] + self.out_proj(h) * classes[:, None, None, 1]
-        return self.out_proj(h)
+
+        logits = self.noise_head(h[:, -1])
+        h = h[:, :-1]
+        if noise_label is not None:
+            if not gen_noise:
+                return self.out_proj(h), 0
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), noise_label.view(-1))
+            return self.out_proj(h), loss
+
+        return self.out_proj(h), logits
 
     # expert
     # def forward(self, x, expert_id=None):

audioenhancer/model/audio_ae/mamba.py

@@ -212,7 +212,8 @@ def __init__(self, config):
         self.mixer_rev = MambaMixer(config)
         # self.mlp = MLP(config)
 
-    def forward(self, hidden_states):
+    def forward(self, hidden_states, gen_noise):
+        bzs, _ , h_dim = hidden_states.shape
         residual = hidden_states
         hidden_states = self.norm(hidden_states)
         if self.residual_in_fp32:
@@ -223,11 +224,17 @@ def forward(self, hidden_states):
             self.mixer_rev.to(torch.float32)
 
         out = self.mixer(hidden_states)
-        out_rev = self.mixer_rev(
-            hidden_states.flip(dims=(1,))
-        ).flip(dims=(1,))
-        hidden_states = out + out_rev
-
+        if gen_noise:
+            out_rev = self.mixer_rev(
+                hidden_states.flip(dims=(1,))[..., 1:, :]
+            ).flip(dims=(1,))
+            hidden_states = out + torch.cat(
+                [out_rev, torch.zeros([bzs, 1, h_dim]).to(device=out.device, dtype=out.dtype)],
+                dim=1
+            )
+        else:
+            out_rev = self.mixer_rev(hidden_states.flip(dims=(1,))).flip(dims=(1,))
+            hidden_states = out + out_rev
         if self.residual_in_fp32:
             hidden_states = hidden_states.to(original_dtype)
             residual = residual.to(original_dtype)

audioenhancer/model/audio_ae/model.py

@@ -16,6 +16,7 @@
 from x_transformers import ContinuousTransformerWrapper, Decoder, Encoder
 
 from audioenhancer.model.audio_ae.latent import LatentProcessor
+from audioenhancer.model.audio_ae.transformer import Transformer
 from audioenhancer.model.audio_ae.vdiffusion import CustomVDiffusion
 
 model = DiffusionModel(
@@ -197,10 +198,25 @@
     ),
 )
 
+transformer = Transformer(
+    in_dim=1024,
+    out_dim=1024,
+    latent_dim=1024,
+    num_layer=6,
+)
+
 mamba_model = LatentProcessor(
     in_dim=1024,
     out_dim=1024,
     latent_dim=2048,
     num_layer=6,
-    num_expert=1,
+    noise_grad=11,
+)
+
+disc_model = LatentProcessor(
+    in_dim=1024,
+    out_dim=1024,
+    latent_dim=2048,
+    num_layer=6,
+    noise_grad=11,
 )

audioenhancer/model/audio_ae/transformer.py

@@ -0,0 +1,53 @@
+import warnings
+
+import torch
+from torch import nn
+from x_transformers import XTransformer
+from xformers.factory import xFormerConfig, xFormer
+
+
+class Transformer(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, latent_dim, num_layer):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.latent_dim = latent_dim
+        self.num_layer = num_layer
+
+        transformer = XTransformer(
+            dim=latent_dim,
+            enc_depth=num_layer,
+            enc_heads=16,
+            enc_max_seq_len=0,
+            enc_attn_flash=True,
+            enc_num_tokens=256,
+            enc_cross_attend=False,
+            enc_ff_glu=True,
+            enc_rotary_pos_emb=True,
+            enc_use_scalenorm=True,
+            enc_zero_init_branch_output=True,
+            dec_num_tokens=256,
+            dec_depth=num_layer,
+            dec_heads=16,
+            dec_ff_glu=True,
+            dec_rotary_pos_emb=True,
+            dec_use_scalenorm=True,
+            dec_attn_flash=True,
+            dec_max_seq_len=0,
+            dec_zero_init_branch_output=True,
+        )
+        self.embed = nn.Parameter(torch.randn(latent_dim))
+        self.encoders = transformer.encoder.attn_layers
+        self.decoders = transformer.decoder.net.attn_layers
+        self.in_proj = nn.Linear(in_dim, latent_dim)
+        self.out_proj = nn.Linear(latent_dim, out_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.in_proj(x)
+
+        x = self.encoders(x)
+
+        h = self.embed.expand(x.shape)
+        h = self.decoders(h, context=x)
+        h = self.out_proj(h)
+        return x - h

scripts/gradio_demo.py

@@ -11,7 +11,7 @@
     "--model_path",
     type=str,
     required=False,
-    default="data/model/model_1000.pt",
+    default="data/model/model_3700.pt",
     help="The path to the model",
 )
 

scripts/inference.py

@@ -10,7 +10,7 @@
 parser = argparse.ArgumentParser()
 parser.add_argument(
     "--audio",
-    default="../media/works/dataset/opus/5700_part2.mp3",  # ../media/works/dataset/dac/5700_part2.mp3
+    default="../media/works/dataset/encodec/5700_part2.mp3",  # ../media/works/dataset/dac/5700_part2.mp3
     type=str,
     required=False,
     help="The path to the audio file to enhance",
@@ -20,7 +20,7 @@
     "--model_path",
     type=str,
     required=False,
-    default="data/model/model_300.pt",
+    default="data/model/model_200.pt",
     help="The path to the model",
 )