+Small models

app.py

@@ -1,6 +1,8 @@
 import subprocess
 import sys
 import shlex
+
+from OtherMethods.Demucs.denoiser import enhance
 #update the system
 subprocess.check_call(["apt-get", "update"])
 subprocess.check_call([sys.executable,"-m","pip","install",
@@ -72,7 +74,10 @@ def mel_transform(audio, X_norm):
     return transform(audio, X_norm)
 
 def load_cleanmel(model_name):
-    model_config = f"./configs/cleanmel_offline.yaml"
+    if "S" in model_name:
+        model_config = f"./configs/cleanmel_offline_S.yaml"
+    else:
+        model_config = f"./configs/cleanmel_offline_L.yaml"
     model_config = yaml.safe_load(open(model_config, "r"))["model"]["arch"]["init_args"]
     cleanmel = CleanMel(**model_config)
     cleanmel.load_state_dict(torch.load(f"./ckpts/CleanMel/{model_name}.ckpt", map_location=DEVICE))
@@ -129,6 +134,20 @@ def enhance_cleanmel_L_mask(audio_path):
     y_hat = vocos(logMel_hat, X_norm).clamp(min=-1, max=1)
     return output(y_hat, logMel_hat)
 
+@spaces.GPU
+@torch.inference_mode()
+def enhance_cleanmel_S_mask(audio_path):
+    model = load_cleanmel("offline_CleanMel_S_mask").to(DEVICE)
+    vocos = load_vocos().to(DEVICE)
+    x = read_audio(audio_path).to(DEVICE)
+    X, X_norm = stft(x)
+    Y_hat = model(X, inference=True)
+    MRM_hat = torch.sigmoid(Y_hat)
+    Y_hat = get_mrm_pred(MRM_hat, x, X_norm)
+    logMel_hat = safe_log(Y_hat)
+    y_hat = vocos(logMel_hat, X_norm).clamp(min=-1, max=1)
+    return output(y_hat, logMel_hat)
+
 @spaces.GPU
 @torch.inference_mode()
 def enhance_cleanmel_L_map(audio_path):
@@ -140,6 +159,17 @@ def enhance_cleanmel_L_map(audio_path):
     y_hat = vocos(logMel_hat, X_norm).clamp(min=-1, max=1)
     return output(y_hat, logMel_hat)
 
+@spaces.GPU
+@torch.inference_mode()
+def enhance_cleanmel_S_map(audio_path):
+    model = load_cleanmel("offline_CleanMel_S_map").to(DEVICE)
+    vocos = load_vocos().to(DEVICE)
+    x = read_audio(audio_path).to(DEVICE)
+    X, X_norm = stft(x)
+    logMel_hat = model(X, inference=True)
+    y_hat = vocos(logMel_hat, X_norm).clamp(min=-1, max=1)
+    return output(y_hat, logMel_hat)
+
 def reset_everything():
     """Reset all components to initial state"""
     return None, None, None
@@ -153,8 +183,12 @@ with gr.Blocks(title="CleanMel Demo") as demo:
     with gr.Row():
         audio_input = gr.Audio(label="Input Audio", type="filepath", sources="upload")
         with gr.Column():
-            enhance_button_map = gr.Button("Enhance Audio (offline CleanMel_L_map)")
-            enhance_button_mask = gr.Button("Enhance Audio (offline CleanMel_L_mask)")
+            
+            enhance_button_map_S = gr.Button("Enhance Audio (offline CleanMel_S_map), 4 mins for 10-second audio")
+            enhance_button_mask_S = gr.Button("Enhance Audio (offline CleanMel_S_mask), 4 mins for 10-second audio")
+            
+            enhance_button_map_L = gr.Button("Enhance Audio (offline CleanMel_L_map), 10 mins for 10-second audio")
+            enhance_button_mask_L = gr.Button("Enhance Audio (offline CleanMel_L_mask), 10 mins for 10-second audio")
             clear_btn = gr.Button(
                 "🗑️ Clear All",
                 variant="secondary",
@@ -165,17 +199,30 @@ with gr.Blocks(title="CleanMel Demo") as demo:
     output_mel = gr.Image(label="Output LogMel Spectrogram", type="filepath", visible=True)
     output_np = gr.File(label="Enhanced LogMel Spec. (.npy)", type="filepath")
     
-    enhance_button_map.click(
+    enhance_button_map_L.click(
         enhance_cleanmel_L_map, 
         inputs=audio_input, 
         outputs=[output_audio, output_mel, output_np]
     )
     
-    enhance_button_mask.click(
+    enhance_button_mask_L.click(
         enhance_cleanmel_L_mask, 
         inputs=audio_input, 
         outputs=[output_audio, output_mel, output_np]
     )
+    
+    enhance_button_map_S.click(
+        enhance_cleanmel_S_map,
+        inputs=audio_input,
+        outputs=[output_audio, output_mel, output_np]
+    )
+    
+    enhance_button_mask_S.click(
+        enhance_cleanmel_S_mask,
+        inputs=audio_input,
+        outputs=[output_audio, output_mel, output_np]
+    )
+    
     clear_btn.click(
             fn=reset_everything,
             outputs=[output_audio, output_mel, output_np]

configs/cleanmel_offline_S.yaml

@@ -0,0 +1,67 @@
+seed_everything: 2
+
+trainer:
+  gradient_clip_val: 10
+  gradient_clip_algorithm: norm
+  devices: null
+  accelerator: gpu
+  strategy: ddp_find_unused_parameters_false
+  sync_batchnorm: false
+  precision: 32
+  num_sanity_val_steps: 3
+  deterministic: true
+  max_epochs: 100
+  log_every_n_steps: 40
+
+model:
+  arch:
+    class_path: model.arch.cleanmel.CleanMel
+    init_args:
+      dim_input: 2
+      dim_output: 1
+      n_layers: 8
+      dim_hidden: 96
+      layer_linear_freq: 1
+      f_kernel_size: 5
+      f_conv_groups: 8
+      n_freqs: 257
+      n_mels: 80
+      mamba_state: 16
+      mamba_conv_kernel: 4
+      online: false
+      sr: 16000
+      n_fft: 512
+  input_stft:
+    class_path: model.io.stft.InputSTFT
+    init_args:
+      n_fft: 512 
+      n_win: 512 
+      n_hop: 128
+      center: true 
+      normalize: false 
+      onesided: true 
+      online: false
+  target_stft:
+    class_path: model.io.stft.TargetMel
+    init_args:
+        sample_rate: 16000
+        n_fft: 512
+        n_win: 512
+        n_hop: 128
+        n_mels: 80
+        f_min: 0
+        f_max: 8000
+        power: 2
+        center: true
+        normalize: false
+        onesided: true
+        mel_norm: "slaney"
+        mel_scale: "slaney"
+        librosa_mel: true
+        online: false
+
+  optimizer: [AdamW, { lr: 0.001, weight_decay: 0.001}]
+  lr_scheduler: [ExponentialLR, { gamma: 0.99 }]
+  exp_name: exp
+  metrics: [DNSMOS]
+  log_eps: 1e-5