update app

app.py

@@ -5,29 +5,34 @@ from pathlib import Path
 from denoisers.SpectralGating import SpectralGating
 from huggingface_hub import hf_hub_download
 from denoisers.demucs import Demucs
-import hydra
-from omegaconf import DictConfig
 import torch
+import torchaudio
+import yaml
 
 
 
-@hydra.main(version_base=None, config_path="conf", config_name="config")
-def run_app(cfg: DictConfig):
-    model = Demucs(cfg['model'])
-    model_path = hf_hub_download(repo_id="BorisovMaksim/demucs", filename="Demucs_original_sr_epoch3.pt")
-    checkpoint = torch.load(model_path)
+def run_app(model_filename, config_filename):
+    model_path = hf_hub_download(repo_id="BorisovMaksim/demucs", filename=model_filename)
+    config_path = hf_hub_download(repo_id="BorisovMaksim/demucs", filename=config_filename)
+    with open(config_path, 'r') as f:
+        config = yaml.safe_load(f)
+    model = Demucs(config['demucs'])
+    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
     model.load_state_dict(checkpoint['model_state_dict'])
 
     def denoising_transform(audio):
-        src_path = Path(__file__).parent.resolve() / Path("cache_wav/original/{}.wav".format(str(uuid.uuid4())))
-        tgt_path = Path(__file__).parent.resolve() / Path("cache_wav/denoised/{}.wav".format(str(uuid.uuid4())))
+        # Path(__file__).parent.resolve()
+        src_path = Path("cache_wav/original/{}.wav".format(str(uuid.uuid4())))
+        tgt_path = Path("cache_wav/denoised/{}.wav".format(str(uuid.uuid4())))
         src_path.parent.mkdir(exist_ok=True, parents=True)
         tgt_path.parent.mkdir(exist_ok=True, parents=True)
         (ffmpeg.input(audio)
          .output(src_path.as_posix(), acodec='pcm_s16le', ac=1, ar=22050)
          .run()
          )
-        model.predict(audio, tgt_path)
+        wav, rate = torchaudio.load(audio)
+        reduced_noise = model.predict(wav)
+        torchaudio.save(tgt_path, reduced_noise, rate)
         return tgt_path
 
     demo = gr.Interface(
@@ -45,5 +50,7 @@ def run_app(cfg: DictConfig):
     demo.launch()
 
 if __name__ == "__main__":
-    run_app()
+    model_filename = "original_sr/Demucs_original_sr_epoch3.pt"
+    config_filename = "original_sr/config.yaml"
+    run_app(model_filename, config_filename)
 

denoisers/demucs.py

@@ -83,3 +83,8 @@ class Demucs(torch.nn.Module):
             x = decoder(x + outs[i])
         x = pad_cut_batch_audio(x, model_input.shape)
         return x
+
+    def predict(self, wav):
+        prediction = self.forward(torch.reshape(wav, (1, 1, -1)))
+        return prediction.detach()[0]
+