adde revision

.gitattributes

@@ -1,35 +1,6 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
+*.json filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.txt filter=lfs diff=lfs merge=lfs -text
+*.pt* filter=lfs diff=lfs merge=lfs -text
+*.ckpt* filter=lfs diff=lfs merge=lfs -text
+*.pl filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,163 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+Temp_Audios/
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+*.ini
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

Languages/ben/G_100000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a8c098eab2e5e378fc52bec57683839bbc641b2241033dab17174f6e37db29a4
+size 145512166

Languages/ben/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec453d50be1976ddb3b501c7dc09128cdf360b83e5b0cf7dec68dc5caa460f49
+size 1887

Languages/ben/vocab.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7085f1a1f6040b4da0ac55bb3ff91b77229d1ed14f7d86df2b23676a1a2cb81b
+size 268

Languages/ell/G_100000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75bfa237f0fe859b34c4340bc7dccd944678cf9984bce5b5a82e2c90ca268db8
+size 145504497

Languages/ell/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec453d50be1976ddb3b501c7dc09128cdf360b83e5b0cf7dec68dc5caa460f49
+size 1887

Languages/ell/vocab.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c53d89f446eba9d061b510e31900d235ef0e021e44a790978dae5a4350a4013
+size 164

Languages/fra/G_100000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63725b5a9201548b2247af02bd69a059335bddf52c1b858dbe38a43a40478bd7
+size 145489135

Languages/fra/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec453d50be1976ddb3b501c7dc09128cdf360b83e5b0cf7dec68dc5caa460f49
+size 1887

Languages/fra/vocab.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b57f0f246b488fe914508d82a8607e1aea357beb0f801069b39bfeb3a4c0d47
+size 104

Languages/guj/G_100000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:427ac3c74f61be494b389cae7d771311d0bcf576f4e2f1b22f257539e26e323a
+size 145501427

Languages/guj/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec453d50be1976ddb3b501c7dc09128cdf360b83e5b0cf7dec68dc5caa460f49
+size 1887

Languages/guj/vocab.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:611d4c5d7ba4bce727c1154277aea43df7a534e22e523877d1885a36727d63c3
+size 232

Languages/hin/G_100000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9f1d5e47edd7368ff40ff5673ddfc606ea713e785420d26c2da396b555458d3b
+size 145510619

Languages/hin/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec453d50be1976ddb3b501c7dc09128cdf360b83e5b0cf7dec68dc5caa460f49
+size 1887

Languages/hin/vocab.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eea03474615c78a1c42d1299c345b6865421c07485544ac3361bff472e5005ac
+size 266

Languages/nld/G_100000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b09e9917b07f06dd911045c8fc8738594b4c4d65c55223c46335093a4904816
+size 145486855

Languages/nld/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec453d50be1976ddb3b501c7dc09128cdf360b83e5b0cf7dec68dc5caa460f49
+size 1887

Languages/nld/vocab.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d7e65d89b8be768ac4a1e53643aebfe42830b82910bfed87f13904e2c5292a4
+size 94

Languages/pol/G_100000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d6f4a9de92a6eb15bca8cb01826d8a9938ab6fb2c04a1c13a06d1d170c88ba6
+size 145490647

Languages/pol/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec453d50be1976ddb3b501c7dc09128cdf360b83e5b0cf7dec68dc5caa460f49
+size 1887

Languages/pol/vocab.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5514b17eb0cd17950849e3f2f2f22a6c7c2d18f2729f5b2fbfc2f2e5f035dc4a
+size 103

app.py

@@ -0,0 +1,317 @@
+
+# load the libraries for the application
+# -------------------------------------------
+import os
+import re
+import nltk
+import torch
+import librosa
+import tempfile
+import subprocess
+
+import gradio as gr
+
+from scipy.io import wavfile
+from nnet import utils, commons
+from transformers import pipeline
+from scipy.io.wavfile import write
+from faster_whisper import WhisperModel
+from nnet.models import SynthesizerTrn as vitsTRN
+from nnet.models_vc import SynthesizerTrn as freeTRN
+from nnet.mel_processing import mel_spectrogram_torch
+from configurations.get_constants import constantConfig
+
+from speaker_encoder.voice_encoder import SpeakerEncoder
+
+from df.enhance import enhance, init_df, load_audio, save_audio
+from configurations.get_hyperparameters import hyperparameterConfig
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+
+nltk.download('punkt')
+from nltk.tokenize import sent_tokenize
+
+# making the FreeVC function
+# ---------------------------------
+class FreeVCModel:
+    def __init__(self, config, ptfile, speaker_model, wavLM_model, device='cpu'):
+        self.hps = utils.get_hparams_from_file(config)
+
+        self.net_g = freeTRN(
+            self.hps.data.filter_length // 2 + 1,
+            self.hps.train.segment_size // self.hps.data.hop_length,
+            **self.hps.model
+        ).to(hyperparameters.device)
+        _ = self.net_g.eval()
+        _ = utils.load_checkpoint(ptfile, self.net_g, None, True)
+        
+        self.cmodel = utils.get_cmodel(device, wavLM_model)
+
+        if self.hps.model.use_spk:
+            self.smodel = SpeakerEncoder(speaker_model)
+
+    def convert(self, src, tgt):
+        fs_src, src_audio = src
+        fs_tgt, tgt_audio = tgt
+
+        src = f"{constants.temp_audio_folder}/src.wav"
+        tgt = f"{constants.temp_audio_folder}/tgt.wav"
+        out = f"{constants.temp_audio_folder}/cnvr.wav"
+        with torch.no_grad():
+            wavfile.write(tgt, fs_tgt, tgt_audio)
+            wav_tgt, _ = librosa.load(tgt, sr=self.hps.data.sampling_rate)
+            wav_tgt, _ = librosa.effects.trim(wav_tgt, top_db=20)
+            if self.hps.model.use_spk:
+                g_tgt = self.smodel.embed_utterance(wav_tgt)
+                g_tgt = torch.from_numpy(g_tgt).unsqueeze(0).to(hyperparameters.device.type)
+            else:
+                wav_tgt = torch.from_numpy(wav_tgt).unsqueeze(0).to(hyperparameters.device.type)
+                mel_tgt = mel_spectrogram_torch(
+                    wav_tgt,
+                    self.hps.data.filter_length,
+                    self.hps.data.n_mel_channels,
+                    self.hps.data.sampling_rate,
+                    self.hps.data.hop_length,
+                    self.hps.data.win_length,
+                    self.hps.data.mel_fmin,
+                    self.hps.data.mel_fmax,
+                )
+            wavfile.write(src, fs_src, src_audio)
+            wav_src, _ = librosa.load(src, sr=self.hps.data.sampling_rate)
+            wav_src = torch.from_numpy(wav_src).unsqueeze(0).to(hyperparameters.device.type)
+            c = utils.get_content(self.cmodel, wav_src)
+
+            if self.hps.model.use_spk:
+                audio = self.net_g.infer(c, g=g_tgt)
+            else:
+                audio = self.net_g.infer(c, mel=mel_tgt)
+            audio = audio[0][0].data.cpu().float().numpy()
+            write(out, 24000, audio)
+
+            return out
+
+# load the system configurations
+constants       = constantConfig()
+hyperparameters = hyperparameterConfig()
+
+# load the models
+model, df_state, _  = init_df(hyperparameters.voice_enhacing_model, config_allow_defaults=True) # voice enhancing model
+stt_model           = WhisperModel(hyperparameters.stt_model, device=hyperparameters.device.type, compute_type="float32") #speech to text model
+
+trans_model     = AutoModelForSeq2SeqLM.from_pretrained(constants.model_name_dict[hyperparameters.nllb_model], torch_dtype=torch.bfloat16).to(hyperparameters.device)
+trans_tokenizer = AutoTokenizer.from_pretrained(constants.model_name_dict[hyperparameters.nllb_model])
+
+modelConvertSpeech = FreeVCModel(config=hyperparameters.text2speech_config, ptfile=hyperparameters.text2speech_model,
+                                 speaker_model=hyperparameters.text2speech_encoder, wavLM_model=hyperparameters.wavlm_model,
+                                 device=hyperparameters.device.type)
+
+# download the language model if doesn't existing
+# ----------------------------------------------------
+def download(lang, lang_directory):
+    
+    if not os.path.exists(f"{lang_directory}/{lang}"):
+        cmd = ";".join([
+                f"wget {constants.language_download_web}/{lang}.tar.gz -O {lang_directory}/{lang}.tar.gz",
+                f"tar zxvf {lang_directory}/{lang}.tar.gz -C {lang_directory}"
+        ])
+        subprocess.check_output(cmd, shell=True)
+    try:
+        os.remove(f"{lang_directory}/{lang}.tar.gz")
+    except:
+        pass
+    return f"{lang_directory}/{lang}"
+
+def preprocess_char(text, lang=None):
+    """
+    Special treatement of characters in certain languages
+    """
+    if lang == 'ron':
+        text = text.replace("ț", "ţ")
+    return text
+
+def preprocess_text(txt, text_mapper, hps, uroman_dir=None, lang=None):
+    txt = preprocess_char(txt, lang=lang)
+    is_uroman = hps.data.training_files.split('.')[-1] == 'uroman'
+    if is_uroman:
+        txt  = text_mapper.uromanize(txt, f'{uroman_dir}/bin/uroman.pl')
+            
+    txt = txt.lower()
+    txt = text_mapper.filter_oov(txt)
+    return txt
+
+def detect_language(text,LID):
+    predictions = LID.predict(text)
+    detected_lang_code = predictions[0][0].replace("__label__", "")
+    return detected_lang_code
+
+# text to speech
+class TextMapper(object):
+    def __init__(self, vocab_file):
+        self.symbols = [x.replace("\n", "") for x in open(vocab_file, encoding="utf-8").readlines()]
+        self.SPACE_ID = self.symbols.index(" ")
+        self._symbol_to_id = {s: i for i, s in enumerate(self.symbols)}
+        self._id_to_symbol = {i: s for i, s in enumerate(self.symbols)}
+
+    def text_to_sequence(self, text, cleaner_names):
+        '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+        Args:
+        text: string to convert to a sequence
+        cleaner_names: names of the cleaner functions to run the text through
+        Returns:
+        List of integers corresponding to the symbols in the text
+        '''
+        sequence = []
+        clean_text = text.strip()
+        for symbol in clean_text:
+            symbol_id = self._symbol_to_id[symbol]
+            sequence += [symbol_id]
+        return sequence
+    
+    def uromanize(self, text, uroman_pl):
+        with tempfile.NamedTemporaryFile() as tf, \
+             tempfile.NamedTemporaryFile() as tf2:
+            with open(tf.name, "w") as f:
+                f.write("\n".join([text]))
+            cmd = f"perl " + uroman_pl
+            cmd += f" -l xxx "
+            cmd +=  f" < {tf.name} > {tf2.name}"
+            os.system(cmd)
+            outtexts = []
+            with open(tf2.name) as f:
+                for line in f:
+                    line =  re.sub(r"\s+", " ", line).strip()
+                    outtexts.append(line)
+            outtext = outtexts[0]
+        return outtext
+
+    def get_text(self, text, hps):
+        text_norm = self.text_to_sequence(text, hps.data.text_cleaners)
+        if hps.data.add_blank:
+            text_norm = commons.intersperse(text_norm, 0)
+        text_norm = torch.LongTensor(text_norm)
+        return text_norm
+    
+    def filter_oov(self, text):
+        val_chars = self._symbol_to_id
+        txt_filt = "".join(list(filter(lambda x: x in val_chars, text)))
+        return txt_filt
+
+def speech_to_text(audio_file):
+    try:
+        fs, audio = audio_file
+        wavfile.write(constants.input_speech_file, fs, audio)
+        audio0, _ = load_audio(constants.input_speech_file, sr=df_state.sr())
+
+        # Enhance the SNR of the audio
+        enhanced = enhance(model, df_state, audio0)
+        save_audio(constants.enhanced_speech_file, enhanced, df_state.sr())
+
+        segments, info = stt_model.transcribe(constants.enhanced_speech_file)
+
+        speech_text = ''
+        for segment in segments:
+            speech_text = f'{speech_text}{segment.text}'
+        try:
+            source_lang_nllb = [k for k, v in constants.flores_codes_to_tts_codes.items() if v[:2] == info.language][0]
+        except:
+            source_lang_nllb = 'language cant be determined, select manually'
+        
+        # text translation
+        return speech_text, gr.Dropdown.update(value=source_lang_nllb)
+    except:
+        return '', gr.Dropdown.update(value='English')
+
+# Text tp speech
+def text_to_speech(text, target_lang):
+    txt = text
+
+    # LANG = get_target_tts_lang(target_lang)
+    LANG = constants.flores_codes_to_tts_codes[target_lang]
+    ckpt_dir = download(LANG, lang_directory=constants.language_directory)
+
+    vocab_file  = f"{ckpt_dir}/{constants.language_vocab_text}"
+    config_file = f"{ckpt_dir}/{constants.language_vocab_configuration}"
+    hps = utils.get_hparams_from_file(config_file)
+    text_mapper = TextMapper(vocab_file)
+    net_g = vitsTRN(
+        len(text_mapper.symbols),
+        hps.data.filter_length // 2 + 1,
+        hps.train.segment_size // hps.data.hop_length,
+        **hps.model)
+    net_g.to(hyperparameters.device)
+    _ = net_g.eval()
+
+    g_pth = f"{ckpt_dir}/{constants.language_vocab_model}"
+
+    _ = utils.load_checkpoint(g_pth, net_g, None)
+
+    txt = preprocess_text(txt, text_mapper, hps, lang=LANG, uroman_dir=constants.uroman_directory)
+    stn_tst = text_mapper.get_text(txt, hps)
+    with torch.no_grad():
+        x_tst = stn_tst.unsqueeze(0).to(hyperparameters.device)
+        x_tst_lengths = torch.LongTensor([stn_tst.size(0)]).to(hyperparameters.device)
+        hyp = net_g.infer(
+            x_tst, x_tst_lengths, noise_scale=.667,
+            noise_scale_w=0.8, length_scale=1.0
+        )[0][0,0].cpu().float().numpy()
+
+    return hps.data.sampling_rate, hyp
+
+def translation(audio, text, source_lang_nllb, target_code_nllb, output_type, sentence_mode):
+    target_code    = constants.flores_codes[target_code_nllb]
+    translator     = pipeline('translation', model=trans_model, tokenizer=trans_tokenizer, src_lang=source_lang_nllb, tgt_lang=target_code, device=hyperparameters.device)
+
+    # output = translator(text, max_length=400)[0]['translation_text']
+    if sentence_mode == "Sentence-wise":
+        sentences = sent_tokenize(text)
+        translated_sentences = []
+        for sentence in sentences:
+            translated_sentence = translator(sentence, max_length=400)[0]['translation_text']
+            translated_sentences.append(translated_sentence)
+        output = ' '.join(translated_sentences)
+    else:
+        output = translator(text, max_length=1024)[0]['translation_text']
+
+    # get the text to speech 
+    fs_out, audio_out = text_to_speech(output, target_code_nllb)
+
+    if output_type == 'own voice':
+        out_file = modelConvertSpeech.convert((fs_out, audio_out), audio)
+        return output, out_file
+
+    wavfile.write(constants.text2speech_wavfile, fs_out, audio_out)
+    return output, constants.text2speech_wavfile
+
+with gr.Blocks(title = "Octopus Translation App") as octopus_translator:
+    with gr.Row():
+        audio_file = gr.Audio(source="microphone")
+    
+    with gr.Row():
+        input_text  = gr.Textbox(label="Input text")
+        source_language     = gr.Dropdown(list(constants.flores_codes.keys()), value='English', label='Source (Autoselected)', interactive=True)
+    
+    with gr.Row():
+        output_text = gr.Textbox(label='Translated text')
+        target_language  = gr.Dropdown(list(constants.flores_codes.keys()), value='German', label='Target', interactive=True)
+    
+    
+    with gr.Row():
+        output_speech = gr.Audio(label='Translated speech')
+        translate_button = gr.Button('Translate')
+
+    
+    with gr.Row():
+        enhance_audio       = gr.Radio(['yes', 'no'], value='yes', label='Enhance input voice', interactive=True)
+        input_type          = gr.Radio(['Whole text', 'Sentence-wise'],value='Sentence-wise', label="Translation Mode", interactive=True)
+        output_audio_type   = gr.Radio(['standard speaker', 'voice transfer'], value='voice transfer', label='Enhance output voice', interactive=True)
+
+    audio_file.change(speech_to_text,
+                      inputs=[audio_file],
+                      outputs=[input_text, source_language])
+
+    translate_button.click(translation,
+                           inputs=[audio_file, input_text, 
+                                   source_language, target_language, 
+                                   output_audio_type, input_type],
+                           outputs=[output_text, output_speech])
+
+octopus_translator.launch(share=False)

aux_files/uroman.pl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ceece2c05343e8bc3b1a7cdc8cecd530af94a7928013c0e4224fd5c729fb29a
+size 5347

configurations/get_constants.py

@@ -0,0 +1,176 @@
+import os
+
+class constantConfig:
+    def __init__(self):
+        self.flores_codes={'Acehnese (Arabic script)': 'ace_Arab', 
+                           'Acehnese (Latin script)': 'ace_Latn', 
+                           'Mesopotamian Arabic': 'acm_Arab', 
+                           'Ta’izzi-Adeni Arabic': 'acq_Arab', 
+                           'Tunisian Arabic': 'aeb_Arab', 
+                           'Afrikaans': 'afr_Latn', 
+                           'South Levantine Arabic': 'ajp_Arab', 
+                           'Akan': 'aka_Latn', 
+                           'Amharic': 'amh_Ethi', 
+                           'North Levantine Arabic': 'apc_Arab', 
+                           'Modern Standard Arabic': 'arb_Arab', 
+                           'Modern Standard Arabic (Romanized)': 'arb_Latn', 
+                           'Najdi Arabic': 'ars_Arab', 
+                           'Moroccan Arabic': 'ary_Arab', 
+                           'Egyptian Arabic': 'arz_Arab', 
+                           'Assamese': 'asm_Beng', 
+                           'Asturian': 'ast_Latn', 
+                           'Awadhi': 'awa_Deva', 
+                           'Central Aymara': 'ayr_Latn', 
+                           'South Azerbaijani': 'azb_Arab', 
+                           'North Azerbaijani': 'azj_Latn', 
+                           'Bashkir': 'bak_Cyrl', 
+                           'Bambara': 'bam_Latn', 
+                           'Balinese': 'ban_Latn', 
+                           'Belarusian': 'bel_Cyrl', 
+                           'Bemba': 'bem_Latn', 
+                           'Bengali': 'ben_Beng', 
+                           'Bhojpuri': 'bho_Deva', 
+                           'Banjar (Arabic script)': 'bjn_Arab', 
+                           'Banjar (Latin script)': 'bjn_Latn', 
+                           'Standard Tibetan': 'bod_Tibt', 
+                           'Bosnian': 'bos_Latn', 
+                           'Buginese': 'bug_Latn', 
+                           'Bulgarian': 'bul_Cyrl', 
+                           'Catalan': 'cat_Latn', 
+                           'Cebuano': 'ceb_Latn', 
+                           'Czech': 'ces_Latn', 
+                           'Chokwe': 'cjk_Latn', 
+                           'Central Kurdish': 'ckb_Arab', 
+                           'Crimean Tatar': 'crh_Latn', 
+                           'Welsh': 'cym_Latn', 
+                           'Danish': 'dan_Latn', 
+                           'German': 'deu_Latn', 
+                           'Southwestern Dinka': 'dik_Latn', 
+                           'Dyula': 'dyu_Latn', 
+                           'Dzongkha': 'dzo_Tibt', 
+                           'Greek': 'ell_Grek', 
+                           'English': 'eng_Latn', 
+                           'Esperanto': 'epo_Latn', 
+                           'Estonian': 'est_Latn', 
+                           'Basque': 'eus_Latn', 
+                           'Ewe': 'ewe_Latn', 
+                           'Faroese': 'fao_Latn', 
+                           'Fijian': 'fij_Latn', 
+                           'Finnish': 'fin_Latn', 
+                           'Fon': 'fon_Latn', 
+                           'French': 'fra_Latn', 
+                           'Friulian': 'fur_Latn', 
+                           'Nigerian Fulfulde': 'fuv_Latn', 
+                           'Scottish Gaelic': 'gla_Latn', 
+                           'Irish': 'gle_Latn', 
+                           'Galician': 'glg_Latn', 
+                           'Guarani': 'grn_Latn', 
+                           'Gujarati': 'guj_Gujr', 
+                           'Haitian Creole': 'hat_Latn', 
+                           'Hausa': 'hau_Latn', 
+                           'Hebrew': 'heb_Hebr', 
+                           'Hindi': 'hin_Deva', 
+                           'Chhattisgarhi': 'hne_Deva', 
+                           'Croatian': 'hrv_Latn', 
+                           'Hungarian': 'hun_Latn', 
+                           'Armenian': 'hye_Armn', 
+                           'Igbo': 'ibo_Latn', 
+                           'Ilocano': 'ilo_Latn', 
+                           'Indonesian': 'ind_Latn', 
+                           'Icelandic': 'isl_Latn', 
+                           'Italian': 'ita_Latn', 
+                           'Javanese': 'jav_Latn', 
+                           'Japanese': 'jpn_Jpan', 
+                           'Kabyle': 'kab_Latn', 
+                           'Jingpho': 'kac_Latn', 
+                           'Kamba': 'kam_Latn', 
+                           'Kannada': 'kan_Knda', 
+                           'Kashmiri (Arabic script)': 'kas_Arab', 
+                           'Kashmiri (Devanagari script)': 'kas_Deva', 
+                           'Georgian': 'kat_Geor', 
+                           'Central Kanuri (Arabic script)': 'knc_Arab', 
+                           'Central Kanuri (Latin script)': 'knc_Latn', 
+                           'Kazakh': 'kaz_Cyrl', 
+                           'Kabiyè': 'kbp_Latn', 
+                           'Kabuverdianu': 'kea_Latn', 
+                           'Khmer': 'khm_Khmr', 
+                           'Kikuyu': 'kik_Latn', 
+                           'Kinyarwanda': 'kin_Latn', 'Kyrgyz': 'kir_Cyrl', 'Kimbundu': 'kmb_Latn', 
+                           'Northern Kurdish': 'kmr_Latn', 'Kikongo': 'kon_Latn', 
+                           'Korean': 'kor_Hang', 'Lao': 'lao_Laoo', 'Ligurian': 'lij_Latn', 
+                           'Limburgish': 'lim_Latn', 'Lingala': 'lin_Latn', 'Lithuanian': 'lit_Latn', 'Lombard': 'lmo_Latn', 
+                           'Latgalian': 'ltg_Latn', 'Luxembourgish': 'ltz_Latn', 'Luba-Kasai': 'lua_Latn', 'Ganda': 'lug_Latn', 
+                           'Luo': 'luo_Latn', 'Mizo': 'lus_Latn', 'Standard Latvian': 'lvs_Latn', 'Magahi': 'mag_Deva', 
+                           'Maithili': 'mai_Deva', 'Malayalam': 'mal_Mlym', 'Marathi': 'mar_Deva', 
+                           'Minangkabau (Arabic script)': 'min_Arab', 'Minangkabau (Latin script)': 'min_Latn', 
+                           'Macedonian': 'mkd_Cyrl', 'Plateau Malagasy': 'plt_Latn', 'Maltese': 'mlt_Latn', 
+                           'Meitei (Bengali script)': 'mni_Beng', 'Halh Mongolian': 'khk_Cyrl', 'Mossi': 'mos_Latn', 
+                           'Maori': 'mri_Latn', 'Burmese': 'mya_Mymr', 'Dutch': 'nld_Latn', 'Norwegian Nynorsk': 'nno_Latn', 
+                           'Norwegian Bokmål': 'nob_Latn', 'Nepali': 'npi_Deva', 'Northern Sotho': 'nso_Latn', 
+                           'Nuer': 'nus_Latn', 
+                           'Nyanja': 'nya_Latn', 'Occitan': 'oci_Latn', 'West Central Oromo': 'gaz_Latn', 'Odia': 'ory_Orya', 
+                           'Pangasinan': 'pag_Latn', 'Eastern Panjabi': 'pan_Guru', 'Papiamento': 'pap_Latn', 
+                           'Western Persian': 'pes_Arab', 
+                           'Polish': 'pol_Latn', 'Portuguese': 'por_Latn', 'Dari': 'prs_Arab', 'Southern Pashto': 'pbt_Arab', 
+                           'Ayacucho Quechua': 'quy_Latn', 'Romanian': 'ron_Latn', 'Rundi': 'run_Latn', 'Russian': 'rus_Cyrl', 
+                           'Sango': 'sag_Latn', 'Sanskrit': 'san_Deva', 'Santali': 'sat_Olck', 'Sicilian': 'scn_Latn', 
+                           'Shan': 'shn_Mymr', 
+                           'Sinhala': 'sin_Sinh', 'Slovak': 'slk_Latn', 'Slovenian': 'slv_Latn', 'Samoan': 'smo_Latn', 
+                           'Shona': 'sna_Latn', 
+                           'Sindhi': 'snd_Arab', 'Somali': 'som_Latn', 'Southern Sotho': 'sot_Latn', 'Spanish': 'spa_Latn', 
+                           'Tosk Albanian': 'als_Latn', 'Sardinian': 'srd_Latn', 'Serbian': 'srp_Cyrl', 'Swati': 'ssw_Latn', 
+                           'Sundanese': 'sun_Latn', 'Swedish': 'swe_Latn', 'Swahili': 'swh_Latn', 'Silesian': 'szl_Latn', 
+                           'Tamil': 'tam_Taml', 'Tatar': 'tat_Cyrl', 'Telugu': 'tel_Telu', 'Tajik': 'tgk_Cyrl', 
+                           'Tagalog': 'tgl_Latn', 
+                           'Thai': 'tha_Thai', 'Tigrinya': 'tir_Ethi', 'Tamasheq (Latin script)': 'taq_Latn', 
+                           'Tamasheq (Tifinagh script)': 'taq_Tfng', 
+                           'Tok Pisin': 'tpi_Latn', 'Tswana': 'tsn_Latn', 'Tsonga': 'tso_Latn', 'Turkmen': 'tuk_Latn', 'Tumbuka': 'tum_Latn', 
+                           'Turkish': 'tur_Latn', 'Twi': 'twi_Latn', 'Central Atlas Tamazight': 'tzm_Tfng', 
+                           'Uyghur': 'uig_Arab', 
+                           'Ukrainian': 'ukr_Cyrl', 'Umbundu': 'umb_Latn', 'Urdu': 'urd_Arab', 'Northern Uzbek': 'uzn_Latn', 
+                           'Venetian': 'vec_Latn', 
+                           'Vietnamese': 'vie_Latn', 'Waray': 'war_Latn', 'Wolof': 'wol_Latn', 'Xhosa': 'xho_Latn', 
+                           'Eastern Yiddish': 'ydd_Hebr', 
+                           'Yoruba': 'yor_Latn', 'Yue Chinese': 'yue_Hant', 'Chinese (Simplified)': 'zho_Hans', 
+                           'Chinese (Traditional)': 'zho_Hant', 
+                           'Standard Malay': 'zsm_Latn', 'Zulu': 'zul_Latn'}
+        
+        self.model_name_dict = {'0.6B': 'facebook/nllb-200-distilled-600M',
+                                '1.3B': 'facebook/nllb-200-distilled-1.3B',
+                                '3.3B': 'facebook/nllb-200-3.3B',
+                                }
+        
+        self.whisper_codes_to_flores_codes = {"de" : self.flores_codes['German'], 
+                                              "en" : self.flores_codes['English'],
+                                              "pl" : self.flores_codes['Polish'],
+                                              "hi" : self.flores_codes['Hindi']
+                                              }
+        
+        self.flores_codes_to_tts_codes = {'Acehnese': 'ace', 'Mesopotamian Arabic': 'acm', 'Ta’izzi-Adeni Arabic': 'acq', 'Tunisian Arabic': 'aeb', 'Afrikaans': 'afr', 'South Levantine Arabic': 'ajp', 'Akan': 'aka', 'Amharic': 'amh', 'North Levantine Arabic': 'apc', 'Modern Standard Arabic': 'arb', 'Najdi Arabic': 'ars', 'Moroccan Arabic': 'ary', 'Egyptian Arabic': 'arz', 'Assamese': 'asm', 'Asturian': 'ast', 'Awadhi': 'awa', 'Central Aymara': 'ayr', 'South Azerbaijani': 'azb', 'North Azerbaijani': 'azj', 'Bashkir': 'bak', 'Bambara': 'bam', 'Balinese': 'ban', 'Belarusian': 'bel', 'Bemba': 'bem', 'Bengali': 'ben', 'Bhojpuri': 'bho', 'Banjar': 'bjn', 'Standard Tibetan': 'bod', 'Bosnian': 'bos', 'Buginese': 'bug', 'Bulgarian': 'bul', 'Catalan': 'cat', 'Cebuano': 'ceb', 'Czech': 'ces', 'Chokwe': 'cjk', 'Central Kurdish': 'ckb', 'Crimean Tatar': 'crh', 'Welsh': 'cym', 'Danish': 'dan', 'German': 'deu', 'Southwestern Dinka': 'dik', 'Dyula': 'dyu', 'Dzongkha': 'dzo', 'Greek': 'ell', 'English': 'eng', 'Esperanto': 'epo', 'Estonian': 'est', 'Basque': 'eus', 'Ewe': 'ewe', 'Faroese': 'fao', 'Fijian': 'fij', 'Finnish': 'fin', 'Fon': 'fon', 'French': 'fra', 'Friulian': 'fur', 'Nigerian Fulfulde': 'fuv', 'Scottish Gaelic': 'gla', 'Irish': 'gle', 'Galician': 'glg', 'Guarani': 'grn', 'Gujarati': 'guj', 'Haitian Creole': 'hat', 'Hausa': 'hau', 'Hebrew': 'heb', 'Hindi': 'hin', 'Chhattisgarhi': 'hne', 'Croatian': 'hrv', 'Hungarian': 'hun', 'Armenian': 'hye', 'Igbo': 'ibo', 'Ilocano': 'ilo', 'Indonesian': 'ind', 'Icelandic': 'isl', 'Italian': 'ita', 'Javanese': 'jav', 'Japanese': 'jpn', 'Kabyle': 'kab', 'Jingpho': 'kac', 'Kamba': 'kam', 'Kannada': 'kan', 'Kashmiri': 'kas', 'Georgian': 'kat', 'Central Kanuri': 'knc', 'Kazakh': 'kaz', 'Kabiyè': 'kbp', 'Kabuverdianu': 'kea', 'Khmer': 'khm', 'Kikuyu': 'kik', 'Kinyarwanda': 'kin', 'Kyrgyz': 'kir', 'Kimbundu': 'kmb', 'Northern Kurdish': 'kmr', 'Kikongo': 'kon', 'Korean': 'kor', 'Lao': 'lao', 'Ligurian': 'lij', 'Limburgish': 'lim', 'Lingala': 'lin', 'Lithuanian': 'lit', 'Lombard': 'lmo', 'Latgalian': 'ltg', 'Luxembourgish': 'ltz', 'Luba-Kasai': 'lua', 'Ganda': 'lug', 'Luo': 'luo', 'Mizo': 'lus', 'Standard Latvian': 'lvs', 'Magahi': 'mag', 'Maithili': 'mai', 'Malayalam': 'mal', 'Marathi': 'mar', 'Minangkabau': 'min', 'Macedonian': 'mkd', 'Plateau Malagasy': 'plt', 'Maltese': 'mlt', 'Meitei': 'mni', 'Halh Mongolian': 'khk', 'Mossi': 'mos', 'Maori': 'mri', 'Burmese': 'mya', 'Dutch': 'nld', 'Norwegian Nynorsk': 'nno', 'Norwegian Bokmål': 'nob', 'Nepali': 'npi', 'Northern Sotho': 'nso', 'Nuer': 'nus', 'Nyanja': 'nya', 'Occitan': 'oci', 'West Central Oromo': 'gaz', 'Odia': 'ory', 'Pangasinan': 'pag', 'Eastern Panjabi': 'pan', 'Papiamento': 'pap', 'Western Persian': 'pes', 'Polish': 'pol', 'Portuguese': 'por', 'Dari': 'prs', 'Southern Pashto': 'pbt', 'Ayacucho Quechua': 'quy', 'Romanian': 'ron', 'Rundi': 'run', 'Russian': 'rus', 'Sango': 'sag', 'Sanskrit': 'san', 'Santali': 'sat', 'Sicilian': 'scn', 'Shan': 'shn', 'Sinhala': 'sin', 'Slovak': 'slk', 'Slovenian': 'slv', 'Samoan': 'smo', 'Shona': 'sna', 'Sindhi': 'snd', 'Somali': 'som', 'Southern Sotho': 'sot', 'Spanish': 'spa', 'Tosk Albanian': 'als', 'Sardinian': 'srd', 'Serbian': 'srp', 'Swati': 'ssw', 'Sundanese': 'sun', 'Swedish': 'swe', 'Swahili': 'swh', 'Silesian': 'szl', 'Tamil': 'tam', 'Tatar': 'tat', 'Telugu': 'tel', 'Tajik': 'tgk', 'Tagalog': 'tgl', 'Thai': 'tha', 'Tigrinya': 'tir', 'Tamasheq': 'taq', 'Tok Pisin': 'tpi', 'Tswana': 'tsn', 'Tsonga': 'tso', 'Turkmen': 'tuk', 'Tumbuka': 'tum', 'Turkish': 'tur', 'Twi': 'twi', 'Central Atlas Tamazight': 'tzm', 'Uyghur': 'uig', 'Ukrainian': 'ukr', 'Umbundu': 'umb', 'Urdu': 'urd', 'Northern Uzbek': 'uzn', 'Venetian': 'vec', 'Vietnamese': 'vie', 'Waray': 'war', 'Wolof': 'wol', 'Xhosa': 'xho', 'Eastern Yiddish': 'ydd', 'Yoruba': 'yor', 'Yue Chinese': 'yue', 'Chinese': 'zho', 'Standard Malay': 'zsm', 'Zulu': 'zul'}
+
+        self.language_directory = 'Languages'
+        self.uroman_directory   = 'aux_files'
+
+        self.language_download_web          = 'https://dl.fbaipublicfiles.com/mms/tts'
+        self.language_vocab_text            = "vocab.txt"
+        self.language_vocab_configuration   = "config.json"
+        self.language_vocab_model           = "G_100000.pth"
+
+        # creating the audio files temporary
+        # ---------------------------------------
+        self.temp_audio_folder = 'Temp_Audios'
+
+        self.text2speech_wavfile    = f'{self.temp_audio_folder}/text2speech.wav'
+        self.enhanced_speech_file   = f"{self.temp_audio_folder}/enhanced.mp3"
+        self.input_speech_file      = f'{self.temp_audio_folder}/output.wav'
+
+
+        try:
+            os.makedirs(self.language_directory)
+        except:
+            pass
+        
+        try:
+            os.makedirs(self.temp_audio_folder)
+        except:
+            pass

configurations/get_hyperparameters.py

@@ -0,0 +1,19 @@
+import torch
+
+class hyperparameterConfig:
+    def __init__(self):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        self.stt_model  = "large-v2"
+        self.nllb_model = '1.3B'
+
+        # text to speech model
+        self.text2speech_model   = 'model_weights/voiceover/freevc-24.pth'
+        self.text2speech_config  = 'model_weights/voiceover/freevc-24.json'
+        self.text2speech_encoder = 'model_weights/voice_enhance/ckpt/pretrained_bak_5805000.pt'
+
+        # voice enhancing model
+        self.voice_enhacing_model = 'model_weights/voice_enhance'
+
+        # loading the wavlm model
+        self.wavlm_model = 'model_weights/wavlm_models/WavLM-Large.pt'

df/__init__.py

@@ -0,0 +1,3 @@
+from .config import config
+
+__all__ = ["config"]

df/checkpoint.py

@@ -0,0 +1,213 @@
+import glob
+import os
+import re
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from loguru import logger
+from torch import nn
+
+from df.config import Csv, config
+from df.model import init_model
+from df.utils import check_finite_module
+from libdf import DF
+
+
+def get_epoch(cp) -> int:
+    return int(os.path.basename(cp).split(".")[0].split("_")[-1])
+
+
+def load_model(
+    cp_dir: Optional[str],
+    df_state: DF,
+    jit: bool = False,
+    mask_only: bool = False,
+    train_df_only: bool = False,
+    extension: str = "ckpt",
+    epoch: Union[str, int, None] = "latest",
+) -> Tuple[nn.Module, int]:
+    if mask_only and train_df_only:
+        raise ValueError("Only one of `mask_only` `train_df_only` can be enabled")
+    model = init_model(df_state, run_df=mask_only is False, train_mask=train_df_only is False)
+    if jit:
+        model = torch.jit.script(model)
+    blacklist: List[str] = config("CP_BLACKLIST", [], Csv(), save=False, section="train")  # type: ignore
+    if cp_dir is not None:
+        epoch = read_cp(
+            model, "model", cp_dir, blacklist=blacklist, extension=extension, epoch=epoch
+        )
+        epoch = 0 if epoch is None else epoch
+    else:
+        epoch = 0
+    return model, epoch
+
+
+def read_cp(
+    obj: Union[torch.optim.Optimizer, nn.Module],
+    name: str,
+    dirname: str,
+    epoch: Union[str, int, None] = "latest",
+    extension="ckpt",
+    blacklist=[],
+    log: bool = True,
+):
+    checkpoints = []
+    if isinstance(epoch, str):
+        assert epoch in ("best", "latest")
+    if epoch == "best":
+        checkpoints = glob.glob(os.path.join(dirname, f"{name}*.{extension}.best"))
+        if len(checkpoints) == 0:
+            logger.warning("Could not find `best` checkpoint. Checking for default...")
+    if len(checkpoints) == 0:
+        checkpoints = glob.glob(os.path.join(dirname, f"{name}*.{extension}"))
+        checkpoints += glob.glob(os.path.join(dirname, f"{name}*.{extension}.best"))
+    if len(checkpoints) == 0:
+        return None
+    if isinstance(epoch, int):
+        latest = next((x for x in checkpoints if get_epoch(x) == epoch), None)
+        if latest is None:
+            logger.error(f"Could not find checkpoint of epoch {epoch}")
+            exit(1)
+    else:
+        latest = max(checkpoints, key=get_epoch)
+        epoch = get_epoch(latest)
+    if log:
+        logger.info("Found checkpoint {} with epoch {}".format(latest, epoch))
+    latest = torch.load(latest, map_location="cpu")
+    latest = {k.replace("clc", "df"): v for k, v in latest.items()}
+    if blacklist:
+        reg = re.compile("".join(f"({b})|" for b in blacklist)[:-1])
+        len_before = len(latest)
+        latest = {k: v for k, v in latest.items() if reg.search(k) is None}
+        if len(latest) < len_before:
+            logger.info("Filtered checkpoint modules: {}".format(blacklist))
+    if isinstance(obj, nn.Module):
+        while True:
+            try:
+                missing, unexpected = obj.load_state_dict(latest, strict=False)
+            except RuntimeError as e:
+                e_str = str(e)
+                logger.warning(e_str)
+                if "size mismatch" in e_str:
+                    latest = {k: v for k, v in latest.items() if k not in e_str}
+                    continue
+                raise e
+            break
+        for key in missing:
+            logger.warning(f"Missing key: '{key}'")
+        for key in unexpected:
+            if key.endswith(".h0"):
+                continue
+            logger.warning(f"Unexpected key: {key}")
+        return epoch
+    obj.load_state_dict(latest)
+
+
+def write_cp(
+    obj: Union[torch.optim.Optimizer, nn.Module],
+    name: str,
+    dirname: str,
+    epoch: int,
+    extension="ckpt",
+    metric: Optional[float] = None,
+    cmp="min",
+):
+    check_finite_module(obj)
+    n_keep = config("n_checkpoint_history", default=3, cast=int, section="train")
+    n_keep_best = config("n_best_checkpoint_history", default=5, cast=int, section="train")
+    if metric is not None:
+        assert cmp in ("min", "max")
+        metric = float(metric)  # Make sure it is not an integer
+        # Each line contains a previous best with entries: (epoch, metric)
+        with open(os.path.join(dirname, ".best"), "a+") as prev_best_f:
+            prev_best_f.seek(0)  # "a+" creates a file in read/write mode without truncating
+            lines = prev_best_f.readlines()
+            if len(lines) == 0:
+                prev_best = float("inf" if cmp == "min" else "-inf")
+            else:
+                prev_best = float(lines[-1].strip().split(" ")[1])
+            cmp = "__lt__" if cmp == "min" else "__gt__"
+            if getattr(metric, cmp)(prev_best):
+                logger.info(f"Saving new best checkpoint at epoch {epoch} with metric: {metric}")
+                prev_best_f.seek(0, os.SEEK_END)
+                np.savetxt(prev_best_f, np.array([[float(epoch), metric]]))
+                cp_name = os.path.join(dirname, f"{name}_{epoch}.{extension}.best")
+                torch.save(obj.state_dict(), cp_name)
+                cleanup(name, dirname, extension + ".best", nkeep=n_keep_best)
+    cp_name = os.path.join(dirname, f"{name}_{epoch}.{extension}")
+    logger.info(f"Writing checkpoint {cp_name} with epoch {epoch}")
+    torch.save(obj.state_dict(), cp_name)
+    cleanup(name, dirname, extension, nkeep=n_keep)
+
+
+def cleanup(name: str, dirname: str, extension: str, nkeep=5):
+    if nkeep < 0:
+        return
+    checkpoints = glob.glob(os.path.join(dirname, f"{name}*.{extension}"))
+    if len(checkpoints) == 0:
+        return
+    checkpoints = sorted(checkpoints, key=get_epoch, reverse=True)
+    for cp in checkpoints[nkeep:]:
+        logger.debug("Removing old checkpoint: {}".format(cp))
+        os.remove(cp)
+
+
+def check_patience(
+    dirname: str, max_patience: int, new_metric: float, cmp: str = "min", raise_: bool = True
+):
+    cmp = "__lt__" if cmp == "min" else "__gt__"
+    new_metric = float(new_metric)  # Make sure it is not an integer
+    prev_patience, prev_metric = read_patience(dirname)
+    if prev_patience is None or getattr(new_metric, cmp)(prev_metric):
+        # We have a better new_metric, reset patience
+        write_patience(dirname, 0, new_metric)
+    else:
+        # We don't have a better metric, decrement patience
+        new_patience = prev_patience + 1
+        write_patience(dirname, new_patience, prev_metric)
+        if new_patience >= max_patience:
+            if raise_:
+                raise ValueError(
+                    f"No improvements on validation metric ({new_metric}) for {max_patience} epochs. "
+                    "Stopping."
+                )
+            else:
+                return False
+    return True
+
+
+def read_patience(dirname: str) -> Tuple[Optional[int], float]:
+    fn = os.path.join(dirname, ".patience")
+    if not os.path.isfile(fn):
+        return None, 0.0
+    patience, metric = np.loadtxt(fn)
+    return int(patience), float(metric)
+
+
+def write_patience(dirname: str, new_patience: int, metric: float):
+    return np.savetxt(os.path.join(dirname, ".patience"), [new_patience, metric])
+
+
+def test_check_patience():
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as d:
+        check_patience(d, 3, 1.0)
+        check_patience(d, 3, 1.0)
+        check_patience(d, 3, 1.0)
+        assert check_patience(d, 3, 1.0, raise_=False) is False
+
+    with tempfile.TemporaryDirectory() as d:
+        check_patience(d, 3, 1.0)
+        check_patience(d, 3, 0.9)
+        check_patience(d, 3, 1.0)
+        check_patience(d, 3, 1.0)
+        assert check_patience(d, 3, 1.0, raise_=False) is False
+
+    with tempfile.TemporaryDirectory() as d:
+        check_patience(d, 3, 1.0, cmp="max")
+        check_patience(d, 3, 1.9, cmp="max")
+        check_patience(d, 3, 1.0, cmp="max")
+        check_patience(d, 3, 1.0, cmp="max")
+        assert check_patience(d, 3, 1.0, cmp="max", raise_=False) is False

df/config.py

@@ -0,0 +1,266 @@
+import os
+import string
+from configparser import ConfigParser
+from shlex import shlex
+from typing import Any, List, Optional, Tuple, Type, TypeVar, Union
+
+from loguru import logger
+
+T = TypeVar("T")
+
+
+class DfParams:
+    def __init__(self):
+        # Sampling rate used for training
+        self.sr: int = config("SR", cast=int, default=48_000, section="DF")
+        # FFT size in samples
+        self.fft_size: int = config("FFT_SIZE", cast=int, default=960, section="DF")
+        # STFT Hop size in samples
+        self.hop_size: int = config("HOP_SIZE", cast=int, default=480, section="DF")
+        # Number of ERB bands
+        self.nb_erb: int = config("NB_ERB", cast=int, default=32, section="DF")
+        # Number of deep filtering bins; DF is applied from 0th to nb_df-th frequency bins
+        self.nb_df: int = config("NB_DF", cast=int, default=96, section="DF")
+        # Normalization decay factor; used for complex and erb features
+        self.norm_tau: float = config("NORM_TAU", 1, float, section="DF")
+        # Local SNR minimum value, ground truth will be truncated
+        self.lsnr_max: int = config("LSNR_MAX", 35, int, section="DF")
+        # Local SNR maximum value, ground truth will be truncated
+        self.lsnr_min: int = config("LSNR_MIN", -15, int, section="DF")
+        # Minimum number of frequency bins per ERB band
+        self.min_nb_freqs = config("MIN_NB_ERB_FREQS", 2, int, section="DF")
+        # Deep Filtering order
+        self.df_order: int = config("DF_ORDER", cast=int, default=5, section="DF")
+        # Deep Filtering look-ahead
+        self.df_lookahead: int = config("DF_LOOKAHEAD", cast=int, default=0, section="DF")
+        # Pad mode. By default, padding will be handled on the input side:
+        # - `input`, which pads the input features passed to the model
+        # - `output`, which pads the output spectrogram corresponding to `df_lookahead`
+        self.pad_mode: str = config("PAD_MODE", default="input_specf", section="DF")
+
+
+class Config:
+    """Adopted from python-decouple"""
+
+    DEFAULT_SECTION = "settings"
+
+    def __init__(self):
+        self.parser: ConfigParser = None  # type: ignore
+        self.path = ""
+        self.modified = False
+        self.allow_defaults = True
+
+    def load(
+        self, path: Optional[str], config_must_exist=False, allow_defaults=True, allow_reload=False
+    ):
+        self.allow_defaults = allow_defaults
+        if self.parser is not None and not allow_reload:
+            raise ValueError("Config already loaded")
+        self.parser = ConfigParser()
+        self.path = path
+        if path is not None and os.path.isfile(path):
+            with open(path) as f:
+                self.parser.read_file(f)
+        else:
+            if config_must_exist:
+                raise ValueError(f"No config file found at '{path}'.")
+        if not self.parser.has_section(self.DEFAULT_SECTION):
+            self.parser.add_section(self.DEFAULT_SECTION)
+        self._fix_clc()
+        self._fix_df()
+
+    def use_defaults(self):
+        self.load(path=None, config_must_exist=False)
+
+    def save(self, path: str):
+        if not self.modified:
+            logger.debug("Config not modified. No need to overwrite on disk.")
+            return
+        if self.parser is None:
+            self.parser = ConfigParser()
+        for section in self.parser.sections():
+            if len(self.parser[section]) == 0:
+                self.parser.remove_section(section)
+        with open(path, mode="w") as f:
+            self.parser.write(f)
+
+    def tostr(self, value, cast):
+        if isinstance(cast, Csv) and isinstance(value, (tuple, list)):
+            return "".join(str(v) + cast.delimiter for v in value)[:-1]
+        return str(value)
+
+    def set(self, option: str, value: T, cast: Type[T], section: Optional[str] = None) -> T:
+        section = self.DEFAULT_SECTION if section is None else section
+        section = section.lower()
+        if not self.parser.has_section(section):
+            self.parser.add_section(section)
+        if self.parser.has_option(section, option):
+            if value == self.cast(self.parser.get(section, option), cast):
+                return value
+        self.modified = True
+        self.parser.set(section, option, self.tostr(value, cast))
+        return value
+
+    def __call__(
+        self,
+        option: str,
+        default: Any = None,
+        cast: Type[T] = str,
+        save: bool = True,
+        section: Optional[str] = None,
+    ) -> T:
+        # Get value either from an ENV or from the .ini file
+        section = self.DEFAULT_SECTION if section is None else section
+        value = None
+        if self.parser is None:
+            raise ValueError("No configuration loaded")
+        if not self.parser.has_section(section.lower()):
+            self.parser.add_section(section.lower())
+        if option in os.environ:
+            value = os.environ[option]
+            if save:
+                self.parser.set(section, option, self.tostr(value, cast))
+        elif self.parser.has_option(section, option):
+            value = self.read_from_section(section, option, default, cast=cast, save=save)
+        elif self.parser.has_option(section.lower(), option):
+            value = self.read_from_section(section.lower(), option, default, cast=cast, save=save)
+        elif self.parser.has_option(self.DEFAULT_SECTION, option):
+            logger.warning(
+                f"Couldn't find option {option} in section {section}. "
+                "Falling back to default settings section."
+            )
+            value = self.read_from_section(self.DEFAULT_SECTION, option, cast=cast, save=save)
+        elif default is None:
+            raise ValueError("Value {} not found.".format(option))
+        elif not self.allow_defaults and save:
+            raise ValueError(f"Value '{option}' not found in config (defaults not allowed).")
+        else:
+            value = default
+            if save:
+                self.set(option, value, cast, section)
+        return self.cast(value, cast)
+
+    def cast(self, value, cast):
+        # Do the casting to get the correct type
+        if cast is bool:
+            value = str(value).lower()
+            if value in {"true", "yes", "y", "on", "1"}:
+                return True  # type: ignore
+            elif value in {"false", "no", "n", "off", "0"}:
+                return False  # type: ignore
+            raise ValueError("Parse error")
+        return cast(value)
+
+    def get(self, option: str, cast: Type[T] = str, section: Optional[str] = None) -> T:
+        section = self.DEFAULT_SECTION if section is None else section
+        if not self.parser.has_section(section):
+            raise KeyError(section)
+        if not self.parser.has_option(section, option):
+            raise KeyError(option)
+        return self.cast(self.parser.get(section, option), cast)
+
+    def read_from_section(
+        self, section: str, option: str, default: Any = None, cast: Type = str, save: bool = True
+    ) -> str:
+        value = self.parser.get(section, option)
+        if not save:
+            # Set to default or remove to not read it at trainig start again
+            if default is None:
+                self.parser.remove_option(section, option)
+            elif not self.allow_defaults:
+                raise ValueError(f"Value '{option}' not found in config (defaults not allowed).")
+            else:
+                self.parser.set(section, option, self.tostr(default, cast))
+        elif section.lower() != section:
+            self.parser.set(section.lower(), option, self.tostr(value, cast))
+            self.parser.remove_option(section, option)
+            self.modified = True
+        return value
+
+    def overwrite(self, section: str, option: str, value: Any):
+        if not self.parser.has_section(section):
+            return ValueError(f"Section not found: '{section}'")
+        if not self.parser.has_option(section, option):
+            return ValueError(f"Option not found '{option}' in section '{section}'")
+        self.modified = True
+        cast = type(value)
+        return self.parser.set(section, option, self.tostr(value, cast))
+
+    def _fix_df(self):
+        """Renaming of some groups/options for compatibility with old models."""
+        if self.parser.has_section("deepfilternet") and self.parser.has_section("df"):
+            sec_deepfilternet = self.parser["deepfilternet"]
+            sec_df = self.parser["df"]
+            if "df_order" in sec_deepfilternet:
+                sec_df["df_order"] = sec_deepfilternet["df_order"]
+                del sec_deepfilternet["df_order"]
+            if "df_lookahead" in sec_deepfilternet:
+                sec_df["df_lookahead"] = sec_deepfilternet["df_lookahead"]
+                del sec_deepfilternet["df_lookahead"]
+
+    def _fix_clc(self):
+        """Renaming of some groups/options for compatibility with old models."""
+        if (
+            not self.parser.has_section("deepfilternet")
+            and self.parser.has_section("train")
+            and self.parser.get("train", "model") == "convgru5"
+        ):
+            self.overwrite("train", "model", "deepfilternet")
+            self.parser.add_section("deepfilternet")
+            self.parser["deepfilternet"] = self.parser["convgru"]
+            del self.parser["convgru"]
+        if not self.parser.has_section("df") and self.parser.has_section("clc"):
+            self.parser["df"] = self.parser["clc"]
+            del self.parser["clc"]
+        for section in self.parser.sections():
+            for k, v in self.parser[section].items():
+                if "clc" in k.lower():
+                    self.parser.set(section, k.lower().replace("clc", "df"), v)
+                    del self.parser[section][k]
+
+    def __repr__(self):
+        msg = ""
+        for section in self.parser.sections():
+            msg += f"{section}:\n"
+            for k, v in self.parser[section].items():
+                msg += f"  {k}: {v}\n"
+        return msg
+
+
+config = Config()
+
+
+class Csv(object):
+    """
+    Produces a csv parser that return a list of transformed elements. From python-decouple.
+    """
+
+    def __init__(
+        self, cast: Type[T] = str, delimiter=",", strip=string.whitespace, post_process=list
+    ):
+        """
+        Parameters:
+        cast -- callable that transforms the item just before it's added to the list.
+        delimiter -- string of delimiters chars passed to shlex.
+        strip -- string of non-relevant characters to be passed to str.strip after the split.
+        post_process -- callable to post process all casted values. Default is `list`.
+        """
+        self.cast: Type[T] = cast
+        self.delimiter = delimiter
+        self.strip = strip
+        self.post_process = post_process
+
+    def __call__(self, value: Union[str, Tuple[T], List[T]]) -> List[T]:
+        """The actual transformation"""
+        if isinstance(value, (tuple, list)):
+            # if default value is a list
+            value = "".join(str(v) + self.delimiter for v in value)[:-1]
+
+        def transform(s):
+            return self.cast(s.strip(self.strip))
+
+        splitter = shlex(value, posix=True)
+        splitter.whitespace = self.delimiter
+        splitter.whitespace_split = True
+
+        return self.post_process(transform(s) for s in splitter)

df/deepfilternet2.py

@@ -0,0 +1,453 @@
+from functools import partial
+from typing import Final, List, Optional, Tuple, Union
+
+import torch
+from loguru import logger
+from torch import Tensor, nn
+
+from df.config import Csv, DfParams, config
+from df.modules import (
+    Conv2dNormAct,
+    ConvTranspose2dNormAct,
+    DfOp,
+    GroupedGRU,
+    GroupedLinear,
+    GroupedLinearEinsum,
+    Mask,
+    SqueezedGRU,
+    erb_fb,
+    get_device,
+)
+from df.multiframe import MF_METHODS, MultiFrameModule
+from libdf import DF
+
+
+class ModelParams(DfParams):
+    section = "deepfilternet"
+
+    def __init__(self):
+        super().__init__()
+        self.conv_lookahead: int = config(
+            "CONV_LOOKAHEAD", cast=int, default=0, section=self.section
+        )
+        self.conv_ch: int = config("CONV_CH", cast=int, default=16, section=self.section)
+        self.conv_depthwise: bool = config(
+            "CONV_DEPTHWISE", cast=bool, default=True, section=self.section
+        )
+        self.convt_depthwise: bool = config(
+            "CONVT_DEPTHWISE", cast=bool, default=True, section=self.section
+        )
+        self.conv_kernel: List[int] = config(
+            "CONV_KERNEL", cast=Csv(int), default=(1, 3), section=self.section  # type: ignore
+        )
+        self.conv_kernel_inp: List[int] = config(
+            "CONV_KERNEL_INP", cast=Csv(int), default=(3, 3), section=self.section  # type: ignore
+        )
+        self.emb_hidden_dim: int = config(
+            "EMB_HIDDEN_DIM", cast=int, default=256, section=self.section
+        )
+        self.emb_num_layers: int = config(
+            "EMB_NUM_LAYERS", cast=int, default=2, section=self.section
+        )
+        self.df_hidden_dim: int = config(
+            "DF_HIDDEN_DIM", cast=int, default=256, section=self.section
+        )
+        self.df_gru_skip: str = config("DF_GRU_SKIP", default="none", section=self.section)
+        self.df_output_layer: str = config(
+            "DF_OUTPUT_LAYER", default="linear", section=self.section
+        )
+        self.df_pathway_kernel_size_t: int = config(
+            "DF_PATHWAY_KERNEL_SIZE_T", cast=int, default=1, section=self.section
+        )
+        self.enc_concat: bool = config("ENC_CONCAT", cast=bool, default=False, section=self.section)
+        self.df_num_layers: int = config("DF_NUM_LAYERS", cast=int, default=3, section=self.section)
+        self.df_n_iter: int = config("DF_N_ITER", cast=int, default=2, section=self.section)
+        self.gru_type: str = config("GRU_TYPE", default="grouped", section=self.section)
+        self.gru_groups: int = config("GRU_GROUPS", cast=int, default=1, section=self.section)
+        self.lin_groups: int = config("LINEAR_GROUPS", cast=int, default=1, section=self.section)
+        self.group_shuffle: bool = config(
+            "GROUP_SHUFFLE", cast=bool, default=True, section=self.section
+        )
+        self.dfop_method: str = config("DFOP_METHOD", cast=str, default="df", section=self.section)
+        self.mask_pf: bool = config("MASK_PF", cast=bool, default=False, section=self.section)
+
+
+def init_model(df_state: Optional[DF] = None, run_df: bool = True, train_mask: bool = True):
+    p = ModelParams()
+    if df_state is None:
+        df_state = DF(sr=p.sr, fft_size=p.fft_size, hop_size=p.hop_size, nb_bands=p.nb_erb)
+    erb = erb_fb(df_state.erb_widths(), p.sr, inverse=False)
+    erb_inverse = erb_fb(df_state.erb_widths(), p.sr, inverse=True)
+    model = DfNet(erb, erb_inverse, run_df, train_mask)
+    return model.to(device=get_device())
+
+
+class Add(nn.Module):
+    def forward(self, a, b):
+        return a + b
+
+
+class Concat(nn.Module):
+    def forward(self, a, b):
+        return torch.cat((a, b), dim=-1)
+
+
+class Encoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        p = ModelParams()
+        assert p.nb_erb % 4 == 0, "erb_bins should be divisible by 4"
+
+        self.erb_conv0 = Conv2dNormAct(
+            1, p.conv_ch, kernel_size=p.conv_kernel_inp, bias=False, separable=True
+        )
+        conv_layer = partial(
+            Conv2dNormAct,
+            in_ch=p.conv_ch,
+            out_ch=p.conv_ch,
+            kernel_size=p.conv_kernel,
+            bias=False,
+            separable=True,
+        )
+        self.erb_conv1 = conv_layer(fstride=2)
+        self.erb_conv2 = conv_layer(fstride=2)
+        self.erb_conv3 = conv_layer(fstride=1)
+        self.df_conv0 = Conv2dNormAct(
+            2, p.conv_ch, kernel_size=p.conv_kernel_inp, bias=False, separable=True
+        )
+        self.df_conv1 = conv_layer(fstride=2)
+        self.erb_bins = p.nb_erb
+        self.emb_in_dim = p.conv_ch * p.nb_erb // 4
+        self.emb_out_dim = p.emb_hidden_dim
+        if p.gru_type == "grouped":
+            self.df_fc_emb = GroupedLinear(
+                p.conv_ch * p.nb_df // 2, self.emb_in_dim, groups=p.lin_groups
+            )
+        else:
+            df_fc_emb = GroupedLinearEinsum(
+                p.conv_ch * p.nb_df // 2, self.emb_in_dim, groups=p.lin_groups
+            )
+            self.df_fc_emb = nn.Sequential(df_fc_emb, nn.ReLU(inplace=True))
+        if p.enc_concat:
+            self.emb_in_dim *= 2
+            self.combine = Concat()
+        else:
+            self.combine = Add()
+        self.emb_out_dim = p.emb_hidden_dim
+        self.emb_n_layers = p.emb_num_layers
+        assert p.gru_type in ("grouped", "squeeze"), f"But got {p.gru_type}"
+        if p.gru_type == "grouped":
+            self.emb_gru = GroupedGRU(
+                self.emb_in_dim,
+                self.emb_out_dim,
+                num_layers=1,
+                batch_first=True,
+                groups=p.gru_groups,
+                shuffle=p.group_shuffle,
+                add_outputs=True,
+            )
+        else:
+            self.emb_gru = SqueezedGRU(
+                self.emb_in_dim,
+                self.emb_out_dim,
+                num_layers=1,
+                batch_first=True,
+                linear_groups=p.lin_groups,
+                linear_act_layer=partial(nn.ReLU, inplace=True),
+            )
+        self.lsnr_fc = nn.Sequential(nn.Linear(self.emb_out_dim, 1), nn.Sigmoid())
+        self.lsnr_scale = p.lsnr_max - p.lsnr_min
+        self.lsnr_offset = p.lsnr_min
+
+    def forward(
+        self, feat_erb: Tensor, feat_spec: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+        # Encodes erb; erb should be in dB scale + normalized; Fe are number of erb bands.
+        # erb: [B, 1, T, Fe]
+        # spec: [B, 2, T, Fc]
+        # b, _, t, _ = feat_erb.shape
+        e0 = self.erb_conv0(feat_erb)  # [B, C, T, F]
+        e1 = self.erb_conv1(e0)  # [B, C*2, T, F/2]
+        e2 = self.erb_conv2(e1)  # [B, C*4, T, F/4]
+        e3 = self.erb_conv3(e2)  # [B, C*4, T, F/4]
+        c0 = self.df_conv0(feat_spec)  # [B, C, T, Fc]
+        c1 = self.df_conv1(c0)  # [B, C*2, T, Fc]
+        cemb = c1.permute(0, 2, 3, 1).flatten(2)  # [B, T, -1]
+        cemb = self.df_fc_emb(cemb)  # [T, B, C * F/4]
+        emb = e3.permute(0, 2, 3, 1).flatten(2)  # [B, T, C * F/4]
+        emb = self.combine(emb, cemb)
+        emb, _ = self.emb_gru(emb)  # [B, T, -1]
+        lsnr = self.lsnr_fc(emb) * self.lsnr_scale + self.lsnr_offset
+        return e0, e1, e2, e3, emb, c0, lsnr
+
+
+class ErbDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        p = ModelParams()
+        assert p.nb_erb % 8 == 0, "erb_bins should be divisible by 8"
+
+        self.emb_out_dim = p.emb_hidden_dim
+
+        if p.gru_type == "grouped":
+            self.emb_gru = GroupedGRU(
+                p.conv_ch * p.nb_erb // 4,  # For compat
+                self.emb_out_dim,
+                num_layers=p.emb_num_layers - 1,
+                batch_first=True,
+                groups=p.gru_groups,
+                shuffle=p.group_shuffle,
+                add_outputs=True,
+            )
+            # SqueezedGRU uses GroupedLinearEinsum, so let's use it here as well
+            fc_emb = GroupedLinear(
+                p.emb_hidden_dim,
+                p.conv_ch * p.nb_erb // 4,
+                groups=p.lin_groups,
+                shuffle=p.group_shuffle,
+            )
+            self.fc_emb = nn.Sequential(fc_emb, nn.ReLU(inplace=True))
+        else:
+            self.emb_gru = SqueezedGRU(
+                self.emb_out_dim,
+                self.emb_out_dim,
+                output_size=p.conv_ch * p.nb_erb // 4,
+                num_layers=p.emb_num_layers - 1,
+                batch_first=True,
+                gru_skip_op=nn.Identity,
+                linear_groups=p.lin_groups,
+                linear_act_layer=partial(nn.ReLU, inplace=True),
+            )
+            self.fc_emb = nn.Identity()
+        tconv_layer = partial(
+            ConvTranspose2dNormAct,
+            kernel_size=p.conv_kernel,
+            bias=False,
+            separable=True,
+        )
+        conv_layer = partial(
+            Conv2dNormAct,
+            bias=False,
+            separable=True,
+        )
+        # convt: TransposedConvolution, convp: Pathway (encoder to decoder) convolutions
+        self.conv3p = conv_layer(p.conv_ch, p.conv_ch, kernel_size=1)
+        self.convt3 = conv_layer(p.conv_ch, p.conv_ch, kernel_size=p.conv_kernel)
+        self.conv2p = conv_layer(p.conv_ch, p.conv_ch, kernel_size=1)
+        self.convt2 = tconv_layer(p.conv_ch, p.conv_ch, fstride=2)
+        self.conv1p = conv_layer(p.conv_ch, p.conv_ch, kernel_size=1)
+        self.convt1 = tconv_layer(p.conv_ch, p.conv_ch, fstride=2)
+        self.conv0p = conv_layer(p.conv_ch, p.conv_ch, kernel_size=1)
+        self.conv0_out = conv_layer(
+            p.conv_ch, 1, kernel_size=p.conv_kernel, activation_layer=nn.Sigmoid
+        )
+
+    def forward(self, emb, e3, e2, e1, e0) -> Tensor:
+        # Estimates erb mask
+        b, _, t, f8 = e3.shape
+        emb, _ = self.emb_gru(emb)
+        emb = self.fc_emb(emb)
+        emb = emb.view(b, t, f8, -1).permute(0, 3, 1, 2)  # [B, C*8, T, F/8]
+        e3 = self.convt3(self.conv3p(e3) + emb)  # [B, C*4, T, F/4]
+        e2 = self.convt2(self.conv2p(e2) + e3)  # [B, C*2, T, F/2]
+        e1 = self.convt1(self.conv1p(e1) + e2)  # [B, C, T, F]
+        m = self.conv0_out(self.conv0p(e0) + e1)  # [B, 1, T, F]
+        return m
+
+
+class DfOutputReshapeMF(nn.Module):
+    """Coefficients output reshape for multiframe/MultiFrameModule
+
+    Requires input of shape B, C, T, F, 2.
+    """
+
+    def __init__(self, df_order: int, df_bins: int):
+        super().__init__()
+        self.df_order = df_order
+        self.df_bins = df_bins
+
+    def forward(self, coefs: Tensor) -> Tensor:
+        # [B, T, F, O*2] -> [B, O, T, F, 2]
+        coefs = coefs.view(*coefs.shape[:-1], -1, 2)
+        coefs = coefs.permute(0, 3, 1, 2, 4)
+        return coefs
+
+
+class DfDecoder(nn.Module):
+    def __init__(self, out_channels: int = -1):
+        super().__init__()
+        p = ModelParams()
+        layer_width = p.conv_ch
+        self.emb_dim = p.emb_hidden_dim
+
+        self.df_n_hidden = p.df_hidden_dim
+        self.df_n_layers = p.df_num_layers
+        self.df_order = p.df_order
+        self.df_bins = p.nb_df
+        self.gru_groups = p.gru_groups
+        self.df_out_ch = out_channels if out_channels > 0 else p.df_order * 2
+
+        conv_layer = partial(Conv2dNormAct, separable=True, bias=False)
+        kt = p.df_pathway_kernel_size_t
+        self.df_convp = conv_layer(layer_width, self.df_out_ch, fstride=1, kernel_size=(kt, 1))
+        if p.gru_type == "grouped":
+            self.df_gru = GroupedGRU(
+                p.emb_hidden_dim,
+                p.df_hidden_dim,
+                num_layers=self.df_n_layers,
+                batch_first=True,
+                groups=p.gru_groups,
+                shuffle=p.group_shuffle,
+                add_outputs=True,
+            )
+        else:
+            self.df_gru = SqueezedGRU(
+                p.emb_hidden_dim,
+                p.df_hidden_dim,
+                num_layers=self.df_n_layers,
+                batch_first=True,
+                gru_skip_op=nn.Identity,
+                linear_act_layer=partial(nn.ReLU, inplace=True),
+            )
+        p.df_gru_skip = p.df_gru_skip.lower()
+        assert p.df_gru_skip in ("none", "identity", "groupedlinear")
+        self.df_skip: Optional[nn.Module]
+        if p.df_gru_skip == "none":
+            self.df_skip = None
+        elif p.df_gru_skip == "identity":
+            assert p.emb_hidden_dim == p.df_hidden_dim, "Dimensions do not match"
+            self.df_skip = nn.Identity()
+        elif p.df_gru_skip == "groupedlinear":
+            self.df_skip = GroupedLinearEinsum(
+                p.emb_hidden_dim, p.df_hidden_dim, groups=p.lin_groups
+            )
+        else:
+            raise NotImplementedError()
+        assert p.df_output_layer in ("linear", "groupedlinear")
+        self.df_out: nn.Module
+        out_dim = self.df_bins * self.df_out_ch
+        if p.df_output_layer == "linear":
+            df_out = nn.Linear(self.df_n_hidden, out_dim)
+        elif p.df_output_layer == "groupedlinear":
+            df_out = GroupedLinearEinsum(self.df_n_hidden, out_dim, groups=p.lin_groups)
+        else:
+            raise NotImplementedError
+        self.df_out = nn.Sequential(df_out, nn.Tanh())
+        self.df_fc_a = nn.Sequential(nn.Linear(self.df_n_hidden, 1), nn.Sigmoid())
+        self.out_transform = DfOutputReshapeMF(self.df_order, self.df_bins)
+
+    def forward(self, emb: Tensor, c0: Tensor) -> Tuple[Tensor, Tensor]:
+        b, t, _ = emb.shape
+        c, _ = self.df_gru(emb)  # [B, T, H], H: df_n_hidden
+        if self.df_skip is not None:
+            c += self.df_skip(emb)
+        c0 = self.df_convp(c0).permute(0, 2, 3, 1)  # [B, T, F, O*2], channels_last
+        alpha = self.df_fc_a(c)  # [B, T, 1]
+        c = self.df_out(c)  # [B, T, F*O*2], O: df_order
+        c = c.view(b, t, self.df_bins, self.df_out_ch) + c0  # [B, T, F, O*2]
+        c = self.out_transform(c)
+        return c, alpha
+
+
+class DfNet(nn.Module):
+    run_df: Final[bool]
+    pad_specf: Final[bool]
+
+    def __init__(
+        self,
+        erb_fb: Tensor,
+        erb_inv_fb: Tensor,
+        run_df: bool = True,
+        train_mask: bool = True,
+    ):
+        super().__init__()
+        p = ModelParams()
+        layer_width = p.conv_ch
+        assert p.nb_erb % 8 == 0, "erb_bins should be divisible by 8"
+        self.df_lookahead = p.df_lookahead if p.pad_mode == "model" else 0
+        self.nb_df = p.nb_df
+        self.freq_bins: int = p.fft_size // 2 + 1
+        self.emb_dim: int = layer_width * p.nb_erb
+        self.erb_bins: int = p.nb_erb
+        if p.conv_lookahead > 0 and p.pad_mode.startswith("input"):
+            self.pad_feat = nn.ConstantPad2d((0, 0, -p.conv_lookahead, p.conv_lookahead), 0.0)
+        else:
+            self.pad_feat = nn.Identity()
+        self.pad_specf = p.pad_mode.endswith("specf")
+        if p.df_lookahead > 0 and self.pad_specf:
+            self.pad_spec = nn.ConstantPad3d((0, 0, 0, 0, -p.df_lookahead, p.df_lookahead), 0.0)
+        else:
+            self.pad_spec = nn.Identity()
+        if (p.conv_lookahead > 0 or p.df_lookahead > 0) and p.pad_mode.startswith("output"):
+            assert p.conv_lookahead == p.df_lookahead
+            pad = (0, 0, 0, 0, -p.conv_lookahead, p.conv_lookahead)
+            self.pad_out = nn.ConstantPad3d(pad, 0.0)
+        else:
+            self.pad_out = nn.Identity()
+        self.register_buffer("erb_fb", erb_fb)
+        self.enc = Encoder()
+        self.erb_dec = ErbDecoder()
+        self.mask = Mask(erb_inv_fb, post_filter=p.mask_pf)
+
+        self.df_order = p.df_order
+        self.df_bins = p.nb_df
+        self.df_op: Union[DfOp, MultiFrameModule]
+        if p.dfop_method == "real_unfold":
+            raise ValueError("RealUnfold DF OP is now unsupported.")
+        assert p.df_output_layer != "linear", "Must be used with `groupedlinear`"
+        self.df_op = MF_METHODS[p.dfop_method](
+            num_freqs=p.nb_df, frame_size=p.df_order, lookahead=self.df_lookahead
+        )
+        n_ch_out = self.df_op.num_channels()
+        self.df_dec = DfDecoder(out_channels=n_ch_out)
+
+        self.run_df = run_df
+        if not run_df:
+            logger.warning("Runing without DF")
+        self.train_mask = train_mask
+        assert p.df_n_iter == 1
+
+    def forward(
+        self,
+        spec: Tensor,
+        feat_erb: Tensor,
+        feat_spec: Tensor,  # Not used, take spec modified by mask instead
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+        """Forward method of DeepFilterNet2.
+
+        Args:
+            spec (Tensor): Spectrum of shape [B, 1, T, F, 2]
+            feat_erb (Tensor): ERB features of shape [B, 1, T, E]
+            feat_spec (Tensor): Complex spectrogram features of shape [B, 1, T, F']
+
+        Returns:
+            spec (Tensor): Enhanced spectrum of shape [B, 1, T, F, 2]
+            m (Tensor): ERB mask estimate of shape [B, 1, T, E]
+            lsnr (Tensor): Local SNR estimate of shape [B, T, 1]
+        """
+        feat_spec = feat_spec.squeeze(1).permute(0, 3, 1, 2)
+
+        feat_erb = self.pad_feat(feat_erb)
+        feat_spec = self.pad_feat(feat_spec)
+        e0, e1, e2, e3, emb, c0, lsnr = self.enc(feat_erb, feat_spec)
+        m = self.erb_dec(emb, e3, e2, e1, e0)
+
+        m = self.pad_out(m.unsqueeze(-1)).squeeze(-1)
+        spec = self.mask(spec, m)
+
+        if self.run_df:
+            df_coefs, df_alpha = self.df_dec(emb, c0)
+            df_coefs = self.pad_out(df_coefs)
+
+            if self.pad_specf:
+                # Only pad the lower part of the spectrum.
+                spec_f = self.pad_spec(spec)
+                spec_f = self.df_op(spec_f, df_coefs)
+                spec[..., : self.nb_df, :] = spec_f[..., : self.nb_df, :]
+            else:
+                spec = self.pad_spec(spec)
+                spec = self.df_op(spec, df_coefs)
+        else:
+            df_alpha = torch.zeros(())
+
+        return spec, m, lsnr, df_alpha

df/enhance.py

@@ -0,0 +1,333 @@
+import argparse
+import os
+import time
+import warnings
+from typing import Optional, Tuple, Union
+
+import torch
+import torchaudio as ta
+from loguru import logger
+from numpy import ndarray
+from torch import Tensor, nn
+from torch.nn import functional as F
+from torchaudio.backend.common import AudioMetaData
+
+import df
+from df import config
+from df.checkpoint import load_model as load_model_cp
+from df.logger import init_logger, warn_once
+from df.model import ModelParams
+from df.modules import get_device
+from df.utils import as_complex, as_real, get_norm_alpha, resample
+from libdf import DF, erb, erb_norm, unit_norm
+
+
+def main(args):
+    model, df_state, suffix = init_df(
+        args.model_base_dir,
+        post_filter=args.pf,
+        log_level=args.log_level,
+        config_allow_defaults=True,
+        epoch=args.epoch,
+    )
+    if args.output_dir is None:
+        args.output_dir = "."
+    elif not os.path.isdir(args.output_dir):
+        os.mkdir(args.output_dir)
+    df_sr = ModelParams().sr
+    n_samples = len(args.noisy_audio_files)
+    for i, file in enumerate(args.noisy_audio_files):
+        progress = (i + 1) / n_samples * 100
+        audio, meta = load_audio(file, df_sr)
+        t0 = time.time()
+        audio = enhance(
+            model, df_state, audio, pad=args.compensate_delay, atten_lim_db=args.atten_lim
+        )
+        t1 = time.time()
+        t_audio = audio.shape[-1] / df_sr
+        t = t1 - t0
+        rtf = t / t_audio
+        fn = os.path.basename(file)
+        p_str = f"{progress:2.0f}% | " if n_samples > 1 else ""
+        logger.info(f"{p_str}Enhanced noisy audio file '{fn}' in {t:.1f}s (RT factor: {rtf:.3f})")
+        audio = resample(audio, df_sr, meta.sample_rate)
+        save_audio(
+            file, audio, sr=meta.sample_rate, output_dir=args.output_dir, suffix=suffix, log=False
+        )
+
+
+def init_df(
+    model_base_dir: Optional[str] = None,
+    post_filter: bool = False,
+    log_level: str = "INFO",
+    log_file: Optional[str] = "enhance.log",
+    config_allow_defaults: bool = False,
+    epoch: Union[str, int, None] = "best",
+    default_model: str = "DeepFilterNet2",
+) -> Tuple[nn.Module, DF, str]:
+    """Initializes and loads config, model and deep filtering state.
+
+    Args:
+        model_base_dir (str): Path to the model directory containing checkpoint and config. If None,
+            load the pretrained DeepFilterNet2 model.
+        post_filter (bool): Enable post filter for some minor, extra noise reduction.
+        log_level (str): Control amount of logging. Defaults to `INFO`.
+        log_file (str): Optional log file name. None disables it. Defaults to `enhance.log`.
+        config_allow_defaults (bool): Whether to allow initializing new config values with defaults.
+        epoch (str): Checkpoint epoch to load. Options are `best`, `latest`, `<int>`, and `none`.
+            `none` disables checkpoint loading. Defaults to `best`.
+
+    Returns:
+        model (nn.Modules): Intialized model, moved to GPU if available.
+        df_state (DF): Deep filtering state for stft/istft/erb
+        suffix (str): Suffix based on the model name. This can be used for saving the enhanced
+            audio.
+    """
+    try:
+        from icecream import ic, install
+
+        ic.configureOutput(includeContext=True)
+        install()
+    except ImportError:
+        pass
+    use_default_model = False
+    if model_base_dir == "DeepFilterNet":
+        default_model = "DeepFilterNet"
+        use_default_model = True
+    elif model_base_dir == "DeepFilterNet2":
+        use_default_model = True
+    if model_base_dir is None or use_default_model:
+        use_default_model = True
+        model_base_dir = os.path.relpath(
+            os.path.join(
+                os.path.dirname(df.__file__), os.pardir, "pretrained_models", default_model
+            )
+        )
+    if not os.path.isdir(model_base_dir):
+        raise NotADirectoryError("Base directory not found at {}".format(model_base_dir))
+    log_file = os.path.join(model_base_dir, log_file) if log_file is not None else None
+    init_logger(file=log_file, level=log_level, model=model_base_dir)
+    if use_default_model:
+        logger.info(f"Using {default_model} model at {model_base_dir}")
+    config.load(
+        os.path.join(model_base_dir, "config.ini"),
+        config_must_exist=True,
+        allow_defaults=config_allow_defaults,
+        allow_reload=True,
+    )
+    if post_filter:
+        config.set("mask_pf", True, bool, ModelParams().section)
+        logger.info("Running with post-filter")
+    p = ModelParams()
+    df_state = DF(
+        sr=p.sr,
+        fft_size=p.fft_size,
+        hop_size=p.hop_size,
+        nb_bands=p.nb_erb,
+        min_nb_erb_freqs=p.min_nb_freqs,
+    )
+    checkpoint_dir = os.path.join(model_base_dir, "checkpoints")
+    load_cp = epoch is not None and not (isinstance(epoch, str) and epoch.lower() == "none")
+    if not load_cp:
+        checkpoint_dir = None
+    try:
+        mask_only = config.get("mask_only", cast=bool, section="train")
+    except KeyError:
+        mask_only = False
+    model, epoch = load_model_cp(checkpoint_dir, df_state, epoch=epoch, mask_only=mask_only)
+    if (epoch is None or epoch == 0) and load_cp:
+        logger.error("Could not find a checkpoint")
+        exit(1)
+    logger.debug(f"Loaded checkpoint from epoch {epoch}")
+    model = model.to(get_device())
+    # Set suffix to model name
+    suffix = os.path.basename(os.path.abspath(model_base_dir))
+    if post_filter:
+        suffix += "_pf"
+    logger.info("Model loaded")
+    return model, df_state, suffix
+
+
+def df_features(audio: Tensor, df: DF, nb_df: int, device=None) -> Tuple[Tensor, Tensor, Tensor]:
+    spec = df.analysis(audio.numpy())  # [C, Tf] -> [C, Tf, F]
+    a = get_norm_alpha(False)
+    erb_fb = df.erb_widths()
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        erb_feat = torch.as_tensor(erb_norm(erb(spec, erb_fb), a)).unsqueeze(1)
+    spec_feat = as_real(torch.as_tensor(unit_norm(spec[..., :nb_df], a)).unsqueeze(1))
+    spec = as_real(torch.as_tensor(spec).unsqueeze(1))
+    if device is not None:
+        spec = spec.to(device)
+        erb_feat = erb_feat.to(device)
+        spec_feat = spec_feat.to(device)
+    return spec, erb_feat, spec_feat
+
+
+def load_audio(
+    file: str, sr: Optional[int], verbose=True, **kwargs
+) -> Tuple[Tensor, AudioMetaData]:
+    """Loads an audio file using torchaudio.
+
+    Args:
+        file (str): Path to an audio file.
+        sr (int): Optionally resample audio to specified target sampling rate.
+        **kwargs: Passed to torchaudio.load(). Depends on the backend. The resample method
+            may be set via `method` which is passed to `resample()`.
+
+    Returns:
+        audio (Tensor): Audio tensor of shape [C, T], if channels_first=True (default).
+        info (AudioMetaData): Meta data of the original audio file. Contains the original sr.
+    """
+    ikwargs = {}
+    if "format" in kwargs:
+        ikwargs["format"] = kwargs["format"]
+    rkwargs = {}
+    if "method" in kwargs:
+        rkwargs["method"] = kwargs.pop("method")
+    info: AudioMetaData = ta.info(file, **ikwargs)
+    audio, orig_sr = ta.load(file, **kwargs)
+    if sr is not None and orig_sr != sr:
+        if verbose:
+            warn_once(
+                f"Audio sampling rate does not match model sampling rate ({orig_sr}, {sr}). "
+                "Resampling..."
+            )
+        audio = resample(audio, orig_sr, sr, **rkwargs)
+    return audio, info
+
+
+def save_audio(
+    file: str,
+    audio: Union[Tensor, ndarray],
+    sr: int,
+    output_dir: Optional[str] = None,
+    suffix: Optional[str] = None,
+    log: bool = False,
+    dtype=torch.int16,
+):
+    outpath = file
+    if suffix is not None:
+        file, ext = os.path.splitext(file)
+        outpath = file + f"_{suffix}" + ext
+    if output_dir is not None:
+        outpath = os.path.join(output_dir, os.path.basename(outpath))
+    if log:
+        logger.info(f"Saving audio file '{outpath}'")
+    audio = torch.as_tensor(audio)
+    if audio.ndim == 1:
+        audio.unsqueeze_(0)
+    if dtype == torch.int16 and audio.dtype != torch.int16:
+        audio = (audio * (1 << 15)).to(torch.int16)
+    if dtype == torch.float32 and audio.dtype != torch.float32:
+        audio = audio.to(torch.float32) / (1 << 15)
+    ta.save(outpath, audio, sr)
+
+
+@torch.no_grad()
+def enhance(
+    model: nn.Module, df_state: DF, audio: Tensor, pad=False, atten_lim_db: Optional[float] = None
+):
+    model.eval()
+    bs = audio.shape[0]
+    if hasattr(model, "reset_h0"):
+        model.reset_h0(batch_size=bs, device=get_device())
+    orig_len = audio.shape[-1]
+    n_fft, hop = 0, 0
+    if pad:
+        n_fft, hop = df_state.fft_size(), df_state.hop_size()
+        # Pad audio to compensate for the delay due to the real-time STFT implementation
+        audio = F.pad(audio, (0, n_fft))
+    nb_df = getattr(model, "nb_df", getattr(model, "df_bins", ModelParams().nb_df))
+    spec, erb_feat, spec_feat = df_features(audio, df_state, nb_df, device=get_device())
+    enhanced = model(spec, erb_feat, spec_feat)[0].cpu()
+    enhanced = as_complex(enhanced.squeeze(1))
+    if atten_lim_db is not None and abs(atten_lim_db) > 0:
+        lim = 10 ** (-abs(atten_lim_db) / 20)
+        enhanced = as_complex(spec.squeeze(1)) * lim + enhanced * (1 - lim)
+    audio = torch.as_tensor(df_state.synthesis(enhanced.numpy()))
+    if pad:
+        # The frame size is equal to p.hop_size. Given a new frame, the STFT loop requires e.g.
+        # ceil((n_fft-hop)/hop). I.e. for 50% overlap, then hop=n_fft//2
+        # requires 1 additional frame lookahead; 75% requires 3 additional frames lookahead.
+        # Thus, the STFT/ISTFT loop introduces an algorithmic delay of n_fft - hop.
+        assert n_fft % hop == 0  # This is only tested for 50% and 75% overlap
+        d = n_fft - hop
+        audio = audio[:, d : orig_len + d]
+    return audio
+
+
+def parse_epoch_type(value: str) -> Union[int, str]:
+    try:
+        return int(value)
+    except ValueError:
+        assert value in ("best", "latest")
+        return value
+
+
+def setup_df_argument_parser(default_log_level: str = "INFO") -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model-base-dir",
+        "-m",
+        type=str,
+        default=None,
+        help="Model directory containing checkpoints and config. "
+        "To load a pretrained model, you may just provide the model name, e.g. `DeepFilterNet`. "
+        "By default, the pretrained DeepFilterNet2 model is loaded.",
+    )
+    parser.add_argument(
+        "--pf",
+        help="Post-filter that slightly over-attenuates very noisy sections.",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        default=None,
+        help="Directory in which the enhanced audio files will be stored.",
+    )
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default=default_log_level,
+        help="Logger verbosity. Can be one of (debug, info, error, none)",
+    )
+    parser.add_argument("--debug", "-d", action="store_const", const="DEBUG", dest="log_level")
+    parser.add_argument(
+        "--epoch",
+        "-e",
+        default="best",
+        type=parse_epoch_type,
+        help="Epoch for checkpoint loading. Can be one of ['best', 'latest', <int>].",
+    )
+    return parser
+
+
+def run():
+    parser = setup_df_argument_parser()
+    parser.add_argument(
+        "--compensate-delay",
+        "-D",
+        action="store_true",
+        help="Add some paddig to compensate the delay introduced by the real-time STFT/ISTFT implementation.",
+    )
+    parser.add_argument(
+        "--atten-lim",
+        "-a",
+        type=int,
+        default=None,
+        help="Attenuation limit in dB by mixing the enhanced signal with the noisy signal.",
+    )
+    parser.add_argument(
+        "noisy_audio_files",
+        type=str,
+        nargs="+",
+        help="List of noise files to mix with the clean speech file.",
+    )
+    main(parser.parse_args())
+
+
+if __name__ == "__main__":
+    run()

df/logger.py

@@ -0,0 +1,212 @@
+import os
+import sys
+import warnings
+from collections import defaultdict
+from copy import deepcopy
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import torch
+from loguru import logger
+from torch.types import Number
+
+from df.modules import GroupedLinearEinsum
+from df.utils import get_branch_name, get_commit_hash, get_device, get_host
+
+_logger_initialized = False
+WARN_ONCE_NO = logger.level("WARNING").no + 1
+DEPRECATED_NO = logger.level("WARNING").no + 2
+
+
+def init_logger(file: Optional[str] = None, level: str = "INFO", model: Optional[str] = None):
+    global _logger_initialized, _duplicate_filter
+    if _logger_initialized:
+        logger.debug("Logger already initialized.")
+    else:
+        logger.remove()
+        level = level.upper()
+        if level.lower() != "none":
+            log_format = Formatter(debug=logger.level(level).no <= logger.level("DEBUG").no).format
+            logger.add(
+                sys.stdout,
+                level=level,
+                format=log_format,
+                filter=lambda r: r["level"].no not in {WARN_ONCE_NO, DEPRECATED_NO},
+            )
+            if file is not None:
+                logger.add(
+                    file,
+                    level=level,
+                    format=log_format,
+                    filter=lambda r: r["level"].no != WARN_ONCE_NO,
+                )
+
+            logger.info(f"Running on torch {torch.__version__}")
+            logger.info(f"Running on host {get_host()}")
+            commit = get_commit_hash()
+            if commit is not None:
+                logger.info(f"Git commit: {commit}, branch: {get_branch_name()}")
+            if (jobid := os.getenv("SLURM_JOB_ID")) is not None:
+                logger.info(f"Slurm jobid: {jobid}")
+            logger.level("WARNONCE", no=WARN_ONCE_NO, color="<yellow><bold>")
+            logger.add(
+                sys.stderr,
+                level=max(logger.level(level).no, WARN_ONCE_NO),
+                format=log_format,
+                filter=lambda r: r["level"].no == WARN_ONCE_NO and _duplicate_filter(r),
+            )
+            logger.level("DEPRECATED", no=DEPRECATED_NO, color="<yellow><bold>")
+            logger.add(
+                sys.stderr,
+                level=max(logger.level(level).no, DEPRECATED_NO),
+                format=log_format,
+                filter=lambda r: r["level"].no == DEPRECATED_NO and _duplicate_filter(r),
+            )
+    if model is not None:
+        logger.info("Loading model settings of {}", os.path.basename(model.rstrip("/")))
+    _logger_initialized = True
+
+
+def warn_once(message, *args, **kwargs):
+    logger.log("WARNONCE", message, *args, **kwargs)
+
+
+def log_deprecated(message, *args, **kwargs):
+    logger.log("DEPRECATED", message, *args, **kwargs)
+
+
+class Formatter:
+    def __init__(self, debug=False):
+        if debug:
+            self.fmt = (
+                "<green>{time:YYYY-MM-DD HH:mm:ss}</green>"
+                " | <level>{level: <8}</level>"
+                " | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan>"
+                " | <level>{message}</level>"
+            )
+        else:
+            self.fmt = (
+                "<green>{time:YYYY-MM-DD HH:mm:ss}</green>"
+                " | <level>{level: <8}</level>"
+                " | <cyan>DF</cyan>"
+                " | <level>{message}</level>"
+            )
+        self.fmt += "\n{exception}"
+
+    def format(self, record):
+        if record["level"].no == WARN_ONCE_NO:
+            return self.fmt.replace("{level: <8}", "WARNING ")
+        return self.fmt
+
+
+def _metrics_key(k_: Tuple[str, float]):
+    k0 = k_[0]
+    ks = k0.split("_")
+    if len(ks) > 2:
+        try:
+            return int(ks[-1])
+        except ValueError:
+            return 1000
+    elif k0 == "loss":
+        return -999
+    elif "loss" in k0.lower():
+        return -998
+    elif k0 == "lr":
+        return 998
+    elif k0 == "wd":
+        return 999
+    else:
+        return -101
+
+
+def log_metrics(prefix: str, metrics: Dict[str, Number], level="INFO"):
+    msg = ""
+    stages = defaultdict(str)
+    loss_msg = ""
+    for n, v in sorted(metrics.items(), key=_metrics_key):
+        if abs(v) > 1e-3:
+            m = f" | {n}: {v:.5f}"
+        else:
+            m = f" | {n}: {v:.3E}"
+        if "stage" in n:
+            s = n.split("stage_")[1].split("_snr")[0]
+            stages[s] += m.replace(f"stage_{s}_", "")
+        elif ("valid" in prefix or "test" in prefix) and "loss" in n.lower():
+            loss_msg += m
+        else:
+            msg += m
+    for s, msg_s in stages.items():
+        logger.log(level, f"{prefix} | stage {s}" + msg_s)
+    if len(stages) == 0:
+        logger.log(level, prefix + msg)
+    if len(loss_msg) > 0:
+        logger.log(level, prefix + loss_msg)
+
+
+class DuplicateFilter:
+    """
+    Filters away duplicate log messages.
+    Modified version of: https://stackoverflow.com/a/60462619
+    """
+
+    def __init__(self):
+        self.msgs = set()
+
+    def __call__(self, record) -> bool:
+        k = f"{record['level']}{record['message']}"
+        if k in self.msgs:
+            return False
+        else:
+            self.msgs.add(k)
+            return True
+
+
+_duplicate_filter = DuplicateFilter()
+
+
+def log_model_summary(model: torch.nn.Module, verbose=False):
+    try:
+        import ptflops
+    except ImportError:
+        logger.debug("Failed to import ptflops. Cannot print model summary.")
+        return
+
+    from df.model import ModelParams
+
+    # Generate input of 1 second audio
+    # Necessary inputs are:
+    #   spec: [B, 1, T, F, 2], F: freq bin
+    #   feat_erb: [B, 1, T, E], E: ERB bands
+    #   feat_spec: [B, 2, T, C*2], C: Complex features
+    p = ModelParams()
+    b = 1
+    t = p.sr // p.hop_size
+    device = get_device()
+    spec = torch.randn([b, 1, t, p.fft_size // 2 + 1, 2]).to(device)
+    feat_erb = torch.randn([b, 1, t, p.nb_erb]).to(device)
+    feat_spec = torch.randn([b, 1, t, p.nb_df, 2]).to(device)
+
+    warnings.filterwarnings("ignore", "RNN module weights", category=UserWarning, module="torch")
+    macs, params = ptflops.get_model_complexity_info(
+        deepcopy(model),
+        (t,),
+        input_constructor=lambda _: {"spec": spec, "feat_erb": feat_erb, "feat_spec": feat_spec},
+        as_strings=False,
+        print_per_layer_stat=verbose,
+        verbose=verbose,
+        custom_modules_hooks={
+            GroupedLinearEinsum: grouped_linear_flops_counter_hook,
+        },
+    )
+    logger.info(f"Model complexity: {params/1e6:.3f}M #Params, {macs/1e6:.1f}M MACS")
+
+
+def grouped_linear_flops_counter_hook(module: GroupedLinearEinsum, input, output):
+    # input: ([B, T, I],)
+    # output: [B, T, H]
+    input = input[0]  # [B, T, I]
+    output_last_dim = module.weight.shape[-1]
+    input = input.unflatten(-1, (module.groups, module.ws))  # [B, T, G, I/G]
+    # GroupedLinear calculates "...gi,...gih->...gh"
+    weight_flops = np.prod(input.shape) * output_last_dim
+    module.__flops__ += int(weight_flops)  # type: ignore

df/model.py

@@ -0,0 +1,24 @@
+from importlib import import_module
+
+import torch
+from loguru import logger
+
+from df.config import DfParams, config
+
+
+class ModelParams(DfParams):
+    def __init__(self):
+        self.__model = config("MODEL", default="deepfilternet", section="train")
+        self.__params = getattr(import_module("df." + self.__model), "ModelParams")()
+
+    def __getattr__(self, attr: str):
+        return getattr(self.__params, attr)
+
+
+def init_model(*args, **kwargs):
+    """Initialize the model specified in the config."""
+    model = config("MODEL", default="deepfilternet", section="train")
+    logger.info(f"Initializing model `{model}`")
+    model = getattr(import_module("df." + model), "init_model")(*args, **kwargs)
+    model.to(memory_format=torch.channels_last)
+    return model

df/modules.py

@@ -0,0 +1,956 @@
+import math
+from collections import OrderedDict
+from typing import Callable, Iterable, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+from torch.nn import init
+from torch.nn.parameter import Parameter
+from typing_extensions import Final
+
+from df.model import ModelParams
+from df.utils import as_complex, as_real, get_device, get_norm_alpha
+from libdf import unit_norm_init
+
+
+class Conv2dNormAct(nn.Sequential):
+    def __init__(
+        self,
+        in_ch: int,
+        out_ch: int,
+        kernel_size: Union[int, Iterable[int]],
+        fstride: int = 1,
+        dilation: int = 1,
+        fpad: bool = True,
+        bias: bool = True,
+        separable: bool = False,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+    ):
+        """Causal Conv2d by delaying the signal for any lookahead.
+
+        Expected input format: [B, C, T, F]
+        """
+        lookahead = 0  # This needs to be handled on the input feature side
+        # Padding on time axis
+        kernel_size = (
+            (kernel_size, kernel_size) if isinstance(kernel_size, int) else tuple(kernel_size)
+        )
+        if fpad:
+            fpad_ = kernel_size[1] // 2 + dilation - 1
+        else:
+            fpad_ = 0
+        pad = (0, 0, kernel_size[0] - 1 - lookahead, lookahead)
+        layers = []
+        if any(x > 0 for x in pad):
+            layers.append(nn.ConstantPad2d(pad, 0.0))
+        groups = math.gcd(in_ch, out_ch) if separable else 1
+        if groups == 1:
+            separable = False
+        if max(kernel_size) == 1:
+            separable = False
+        layers.append(
+            nn.Conv2d(
+                in_ch,
+                out_ch,
+                kernel_size=kernel_size,
+                padding=(0, fpad_),
+                stride=(1, fstride),  # Stride over time is always 1
+                dilation=(1, dilation),  # Same for dilation
+                groups=groups,
+                bias=bias,
+            )
+        )
+        if separable:
+            layers.append(nn.Conv2d(out_ch, out_ch, kernel_size=1, bias=False))
+        if norm_layer is not None:
+            layers.append(norm_layer(out_ch))
+        if activation_layer is not None:
+            layers.append(activation_layer())
+        super().__init__(*layers)
+
+
+class ConvTranspose2dNormAct(nn.Sequential):
+    def __init__(
+        self,
+        in_ch: int,
+        out_ch: int,
+        kernel_size: Union[int, Tuple[int, int]],
+        fstride: int = 1,
+        dilation: int = 1,
+        fpad: bool = True,
+        bias: bool = True,
+        separable: bool = False,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+    ):
+        """Causal ConvTranspose2d.
+
+        Expected input format: [B, C, T, F]
+        """
+        # Padding on time axis, with lookahead = 0
+        lookahead = 0  # This needs to be handled on the input feature side
+        kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
+        if fpad:
+            fpad_ = kernel_size[1] // 2
+        else:
+            fpad_ = 0
+        pad = (0, 0, kernel_size[0] - 1 - lookahead, lookahead)
+        layers = []
+        if any(x > 0 for x in pad):
+            layers.append(nn.ConstantPad2d(pad, 0.0))
+        groups = math.gcd(in_ch, out_ch) if separable else 1
+        if groups == 1:
+            separable = False
+        layers.append(
+            nn.ConvTranspose2d(
+                in_ch,
+                out_ch,
+                kernel_size=kernel_size,
+                padding=(kernel_size[0] - 1, fpad_ + dilation - 1),
+                output_padding=(0, fpad_),
+                stride=(1, fstride),  # Stride over time is always 1
+                dilation=(1, dilation),
+                groups=groups,
+                bias=bias,
+            )
+        )
+        if separable:
+            layers.append(nn.Conv2d(out_ch, out_ch, kernel_size=1, bias=False))
+        if norm_layer is not None:
+            layers.append(norm_layer(out_ch))
+        if activation_layer is not None:
+            layers.append(activation_layer())
+        super().__init__(*layers)
+
+
+def convkxf(
+    in_ch: int,
+    out_ch: Optional[int] = None,
+    k: int = 1,
+    f: int = 3,
+    fstride: int = 2,
+    lookahead: int = 0,
+    batch_norm: bool = False,
+    act: nn.Module = nn.ReLU(inplace=True),
+    mode="normal",  # must be "normal", "transposed" or "upsample"
+    depthwise: bool = True,
+    complex_in: bool = False,
+):
+    bias = batch_norm is False
+    assert f % 2 == 1
+    stride = 1 if f == 1 else (1, fstride)
+    if out_ch is None:
+        out_ch = in_ch * 2 if mode == "normal" else in_ch // 2
+    fpad = (f - 1) // 2
+    convpad = (0, fpad)
+    modules = []
+    # Manually pad for time axis kernel to not introduce delay
+    pad = (0, 0, k - 1 - lookahead, lookahead)
+    if any(p > 0 for p in pad):
+        modules.append(("pad", nn.ConstantPad2d(pad, 0.0)))
+    if depthwise:
+        groups = min(in_ch, out_ch)
+    else:
+        groups = 1
+    if in_ch % groups != 0 or out_ch % groups != 0:
+        groups = 1
+    if complex_in and groups % 2 == 0:
+        groups //= 2
+    convkwargs = {
+        "in_channels": in_ch,
+        "out_channels": out_ch,
+        "kernel_size": (k, f),
+        "stride": stride,
+        "groups": groups,
+        "bias": bias,
+    }
+    if mode == "normal":
+        modules.append(("sconv", nn.Conv2d(padding=convpad, **convkwargs)))
+    elif mode == "transposed":
+        # Since pytorch's transposed conv padding does not correspond to the actual padding but
+        # rather the padding that was used in the encoder conv, we need to set time axis padding
+        # according to k. E.g., this disables the padding for k=2:
+        #     dilation - (k - 1) - padding
+        #   = 1        - (2 - 1) - 1 = 0; => padding = fpad (=1 for k=2)
+        padding = (k - 1, fpad)
+        modules.append(
+            ("sconvt", nn.ConvTranspose2d(padding=padding, output_padding=convpad, **convkwargs))
+        )
+    elif mode == "upsample":
+        modules.append(("upsample", FreqUpsample(fstride)))
+        convkwargs["stride"] = 1
+        modules.append(("sconv", nn.Conv2d(padding=convpad, **convkwargs)))
+    else:
+        raise NotImplementedError()
+    if groups > 1:
+        modules.append(("1x1conv", nn.Conv2d(out_ch, out_ch, 1, bias=False)))
+    if batch_norm:
+        modules.append(("norm", nn.BatchNorm2d(out_ch)))
+    modules.append(("act", act))
+    return nn.Sequential(OrderedDict(modules))
+
+
+class FreqUpsample(nn.Module):
+    def __init__(self, factor: int, mode="nearest"):
+        super().__init__()
+        self.f = float(factor)
+        self.mode = mode
+
+    def forward(self, x: Tensor) -> Tensor:
+        return F.interpolate(x, scale_factor=[1.0, self.f], mode=self.mode)
+
+
+def erb_fb(widths: np.ndarray, sr: int, normalized: bool = True, inverse: bool = False) -> Tensor:
+    n_freqs = int(np.sum(widths))
+    all_freqs = torch.linspace(0, sr // 2, n_freqs + 1)[:-1]
+
+    b_pts = np.cumsum([0] + widths.tolist()).astype(int)[:-1]
+
+    fb = torch.zeros((all_freqs.shape[0], b_pts.shape[0]))
+    for i, (b, w) in enumerate(zip(b_pts.tolist(), widths.tolist())):
+        fb[b : b + w, i] = 1
+    # Normalize to constant energy per resulting band
+    if inverse:
+        fb = fb.t()
+        if not normalized:
+            fb /= fb.sum(dim=1, keepdim=True)
+    else:
+        if normalized:
+            fb /= fb.sum(dim=0)
+    return fb.to(device=get_device())
+
+
+class Mask(nn.Module):
+    def __init__(self, erb_inv_fb: Tensor, post_filter: bool = False, eps: float = 1e-12):
+        super().__init__()
+        self.erb_inv_fb: Tensor
+        self.register_buffer("erb_inv_fb", erb_inv_fb)
+        self.clamp_tensor = torch.__version__ > "1.9.0" or torch.__version__ == "1.9.0"
+        self.post_filter = post_filter
+        self.eps = eps
+
+    def pf(self, mask: Tensor, beta: float = 0.02) -> Tensor:
+        """Post-Filter proposed by Valin et al. [1].
+
+        Args:
+            mask (Tensor): Real valued mask, typically of shape [B, C, T, F].
+            beta: Global gain factor.
+        Refs:
+            [1]: Valin et al.: A Perceptually-Motivated Approach for Low-Complexity, Real-Time Enhancement of Fullband Speech.
+        """
+        mask_sin = mask * torch.sin(np.pi * mask / 2)
+        mask_pf = (1 + beta) * mask / (1 + beta * mask.div(mask_sin.clamp_min(self.eps)).pow(2))
+        return mask_pf
+
+    def forward(self, spec: Tensor, mask: Tensor, atten_lim: Optional[Tensor] = None) -> Tensor:
+        # spec (real) [B, 1, T, F, 2], F: freq_bins
+        # mask (real): [B, 1, T, Fe], Fe: erb_bins
+        # atten_lim: [B]
+        if not self.training and self.post_filter:
+            mask = self.pf(mask)
+        if atten_lim is not None:
+            # dB to amplitude
+            atten_lim = 10 ** (-atten_lim / 20)
+            # Greater equal (__ge__) not implemented for TorchVersion.
+            if self.clamp_tensor:
+                # Supported by torch >= 1.9
+                mask = mask.clamp(min=atten_lim.view(-1, 1, 1, 1))
+            else:
+                m_out = []
+                for i in range(atten_lim.shape[0]):
+                    m_out.append(mask[i].clamp_min(atten_lim[i].item()))
+                mask = torch.stack(m_out, dim=0)
+        mask = mask.matmul(self.erb_inv_fb)  # [B, 1, T, F]
+        return spec * mask.unsqueeze(4)
+
+
+class ExponentialUnitNorm(nn.Module):
+    """Unit norm for a complex spectrogram.
+
+    This should match the rust code:
+    ```rust
+        for (x, s) in xs.iter_mut().zip(state.iter_mut()) {
+            *s = x.norm() * (1. - alpha) + *s * alpha;
+            *x /= s.sqrt();
+        }
+    ```
+    """
+
+    alpha: Final[float]
+    eps: Final[float]
+
+    def __init__(self, alpha: float, num_freq_bins: int, eps: float = 1e-14):
+        super().__init__()
+        self.alpha = alpha
+        self.eps = eps
+        self.init_state: Tensor
+        s = torch.from_numpy(unit_norm_init(num_freq_bins)).view(1, 1, num_freq_bins, 1)
+        self.register_buffer("init_state", s)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # x: [B, C, T, F, 2]
+        b, c, t, f, _ = x.shape
+        x_abs = x.square().sum(dim=-1, keepdim=True).clamp_min(self.eps).sqrt()
+        state = self.init_state.clone().expand(b, c, f, 1)
+        out_states: List[Tensor] = []
+        for t in range(t):
+            state = x_abs[:, :, t] * (1 - self.alpha) + state * self.alpha
+            out_states.append(state)
+        return x / torch.stack(out_states, 2).sqrt()
+
+
+class DfOp(nn.Module):
+    df_order: Final[int]
+    df_bins: Final[int]
+    df_lookahead: Final[int]
+    freq_bins: Final[int]
+
+    def __init__(
+        self,
+        df_bins: int,
+        df_order: int = 5,
+        df_lookahead: int = 0,
+        method: str = "complex_strided",
+        freq_bins: int = 0,
+    ):
+        super().__init__()
+        self.df_order = df_order
+        self.df_bins = df_bins
+        self.df_lookahead = df_lookahead
+        self.freq_bins = freq_bins
+        self.set_forward(method)
+
+    def set_forward(self, method: str):
+        # All forward methods should be mathematically similar.
+        # DeepFilterNet results are obtained with 'real_unfold'.
+        forward_methods = {
+            "real_loop": self.forward_real_loop,
+            "real_strided": self.forward_real_strided,
+            "real_unfold": self.forward_real_unfold,
+            "complex_strided": self.forward_complex_strided,
+            "real_one_step": self.forward_real_no_pad_one_step,
+            "real_hidden_state_loop": self.forward_real_hidden_state_loop,
+        }
+        if method not in forward_methods.keys():
+            raise NotImplementedError(f"`method` must be one of {forward_methods.keys()}")
+        if method == "real_hidden_state_loop":
+            assert self.freq_bins >= self.df_bins
+            self.spec_buf: Tensor
+            # Currently only designed for batch size of 1
+            self.register_buffer(
+                "spec_buf", torch.zeros(1, 1, self.df_order, self.freq_bins, 2), persistent=False
+            )
+        self.forward = forward_methods[method]
+
+    def forward_real_loop(
+        self, spec: Tensor, coefs: Tensor, alpha: Optional[Tensor] = None
+    ) -> Tensor:
+        # Version 0: Manual loop over df_order, maybe best for onnx export?
+        b, _, t, _, _ = spec.shape
+        f = self.df_bins
+        padded = spec_pad(
+            spec[..., : self.df_bins, :].squeeze(1), self.df_order, self.df_lookahead, dim=-3
+        )
+
+        spec_f = torch.zeros((b, t, f, 2), device=spec.device)
+        for i in range(self.df_order):
+            spec_f[..., 0] += padded[:, i : i + t, ..., 0] * coefs[:, :, i, :, 0]
+            spec_f[..., 0] -= padded[:, i : i + t, ..., 1] * coefs[:, :, i, :, 1]
+            spec_f[..., 1] += padded[:, i : i + t, ..., 1] * coefs[:, :, i, :, 0]
+            spec_f[..., 1] += padded[:, i : i + t, ..., 0] * coefs[:, :, i, :, 1]
+        return assign_df(spec, spec_f.unsqueeze(1), self.df_bins, alpha)
+
+    def forward_real_strided(
+        self, spec: Tensor, coefs: Tensor, alpha: Optional[Tensor] = None
+    ) -> Tensor:
+        # Version1: Use as_strided instead of unfold
+        # spec (real) [B, 1, T, F, 2], O: df_order
+        # coefs (real) [B, T, O, F, 2]
+        # alpha (real) [B, T, 1]
+        padded = as_strided(
+            spec[..., : self.df_bins, :].squeeze(1), self.df_order, self.df_lookahead, dim=-3
+        )
+        # Complex numbers are not supported by onnx
+        re = padded[..., 0] * coefs[..., 0]
+        re -= padded[..., 1] * coefs[..., 1]
+        im = padded[..., 1] * coefs[..., 0]
+        im += padded[..., 0] * coefs[..., 1]
+        spec_f = torch.stack((re, im), -1).sum(2)
+        return assign_df(spec, spec_f.unsqueeze(1), self.df_bins, alpha)
+
+    def forward_real_unfold(
+        self, spec: Tensor, coefs: Tensor, alpha: Optional[Tensor] = None
+    ) -> Tensor:
+        # Version2: Unfold
+        # spec (real) [B, 1, T, F, 2], O: df_order
+        # coefs (real) [B, T, O, F, 2]
+        # alpha (real) [B, T, 1]
+        padded = spec_pad(
+            spec[..., : self.df_bins, :].squeeze(1), self.df_order, self.df_lookahead, dim=-3
+        )
+        padded = padded.unfold(dimension=1, size=self.df_order, step=1)  # [B, T, F, 2, O]
+        padded = padded.permute(0, 1, 4, 2, 3)
+        spec_f = torch.empty_like(padded)
+        spec_f[..., 0] = padded[..., 0] * coefs[..., 0]  # re1
+        spec_f[..., 0] -= padded[..., 1] * coefs[..., 1]  # re2
+        spec_f[..., 1] = padded[..., 1] * coefs[..., 0]  # im1
+        spec_f[..., 1] += padded[..., 0] * coefs[..., 1]  # im2
+        spec_f = spec_f.sum(dim=2)
+        return assign_df(spec, spec_f.unsqueeze(1), self.df_bins, alpha)
+
+    def forward_complex_strided(
+        self, spec: Tensor, coefs: Tensor, alpha: Optional[Tensor] = None
+    ) -> Tensor:
+        # Version3: Complex strided; definatly nicest, no permute, no indexing, but complex gradient
+        # spec (real) [B, 1, T, F, 2], O: df_order
+        # coefs (real) [B, T, O, F, 2]
+        # alpha (real) [B, T, 1]
+        padded = as_strided(
+            spec[..., : self.df_bins, :].squeeze(1), self.df_order, self.df_lookahead, dim=-3
+        )
+        spec_f = torch.sum(torch.view_as_complex(padded) * torch.view_as_complex(coefs), dim=2)
+        spec_f = torch.view_as_real(spec_f)
+        return assign_df(spec, spec_f.unsqueeze(1), self.df_bins, alpha)
+
+    def forward_real_no_pad_one_step(
+        self, spec: Tensor, coefs: Tensor, alpha: Optional[Tensor] = None
+    ) -> Tensor:
+        # Version4: Only viable for onnx handling. `spec` needs external (ring-)buffer handling.
+        # Thus, time steps `t` must be equal to `df_order`.
+
+        # spec (real) [B, 1, O, F', 2]
+        # coefs (real) [B, 1, O, F, 2]
+        assert (
+            spec.shape[2] == self.df_order
+        ), "This forward method needs spectrogram buffer with `df_order` time steps as input"
+        assert coefs.shape[1] == 1, "This forward method is only valid for 1 time step"
+        sre, sim = spec[..., : self.df_bins, :].split(1, -1)
+        cre, cim = coefs.split(1, -1)
+        outr = torch.sum(sre * cre - sim * cim, dim=2).squeeze(-1)
+        outi = torch.sum(sre * cim + sim * cre, dim=2).squeeze(-1)
+        spec_f = torch.stack((outr, outi), dim=-1)
+        return assign_df(
+            spec[:, :, self.df_order - self.df_lookahead - 1],
+            spec_f.unsqueeze(1),
+            self.df_bins,
+            alpha,
+        )
+
+    def forward_real_hidden_state_loop(self, spec: Tensor, coefs: Tensor, alpha: Tensor) -> Tensor:
+        # Version5: Designed for onnx export. `spec` buffer handling is done via a torch buffer.
+
+        # spec (real) [B, 1, T, F', 2]
+        # coefs (real) [B, T, O, F, 2]
+        b, _, t, _, _ = spec.shape
+        spec_out = torch.empty((b, 1, t, self.freq_bins, 2), device=spec.device)
+        for t in range(spec.shape[2]):
+            self.spec_buf = self.spec_buf.roll(-1, dims=2)
+            self.spec_buf[:, :, -1] = spec[:, :, t]
+            sre, sim = self.spec_buf[..., : self.df_bins, :].split(1, -1)
+            cre, cim = coefs[:, t : t + 1].split(1, -1)
+            outr = torch.sum(sre * cre - sim * cim, dim=2).squeeze(-1)
+            outi = torch.sum(sre * cim + sim * cre, dim=2).squeeze(-1)
+            spec_f = torch.stack((outr, outi), dim=-1)
+            spec_out[:, :, t] = assign_df(
+                self.spec_buf[:, :, self.df_order - self.df_lookahead - 1].unsqueeze(2),
+                spec_f.unsqueeze(1),
+                self.df_bins,
+                alpha[:, t],
+            ).squeeze(2)
+        return spec_out
+
+
+def assign_df(spec: Tensor, spec_f: Tensor, df_bins: int, alpha: Optional[Tensor]):
+    spec_out = spec.clone()
+    if alpha is not None:
+        b = spec.shape[0]
+        alpha = alpha.view(b, 1, -1, 1, 1)
+        spec_out[..., :df_bins, :] = spec_f * alpha + spec[..., :df_bins, :] * (1 - alpha)
+    else:
+        spec_out[..., :df_bins, :] = spec_f
+    return spec_out
+
+
+def spec_pad(x: Tensor, window_size: int, lookahead: int, dim: int = 0) -> Tensor:
+    pad = [0] * x.dim() * 2
+    if dim >= 0:
+        pad[(x.dim() - dim - 1) * 2] = window_size - lookahead - 1
+        pad[(x.dim() - dim - 1) * 2 + 1] = lookahead
+    else:
+        pad[(-dim - 1) * 2] = window_size - lookahead - 1
+        pad[(-dim - 1) * 2 + 1] = lookahead
+    return F.pad(x, pad)
+
+
+def as_strided(x: Tensor, window_size: int, lookahead: int, step: int = 1, dim: int = 0) -> Tensor:
+    shape = list(x.shape)
+    shape.insert(dim + 1, window_size)
+    x = spec_pad(x, window_size, lookahead, dim=dim)
+    # torch.fx workaround
+    step = 1
+    stride = [x.stride(0), x.stride(1), x.stride(2), x.stride(3)]
+    stride.insert(dim, stride[dim] * step)
+    return torch.as_strided(x, shape, stride)
+
+
+class GroupedGRULayer(nn.Module):
+    input_size: Final[int]
+    hidden_size: Final[int]
+    out_size: Final[int]
+    bidirectional: Final[bool]
+    num_directions: Final[int]
+    groups: Final[int]
+    batch_first: Final[bool]
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        groups: int,
+        batch_first: bool = True,
+        bias: bool = True,
+        dropout: float = 0,
+        bidirectional: bool = False,
+    ):
+        super().__init__()
+        assert input_size % groups == 0
+        assert hidden_size % groups == 0
+        kwargs = {
+            "bias": bias,
+            "batch_first": batch_first,
+            "dropout": dropout,
+            "bidirectional": bidirectional,
+        }
+        self.input_size = input_size // groups
+        self.hidden_size = hidden_size // groups
+        self.out_size = hidden_size
+        self.bidirectional = bidirectional
+        self.num_directions = 2 if bidirectional else 1
+        self.groups = groups
+        self.batch_first = batch_first
+        assert (self.hidden_size % groups) == 0, "Hidden size must be divisible by groups"
+        self.layers = nn.ModuleList(
+            (nn.GRU(self.input_size, self.hidden_size, **kwargs) for _ in range(groups))
+        )
+
+    def flatten_parameters(self):
+        for layer in self.layers:
+            layer.flatten_parameters()
+
+    def get_h0(self, batch_size: int = 1, device: torch.device = torch.device("cpu")):
+        return torch.zeros(
+            self.groups * self.num_directions,
+            batch_size,
+            self.hidden_size,
+            device=device,
+        )
+
+    def forward(self, input: Tensor, h0: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
+        # input shape: [B, T, I] if batch_first else [T, B, I], B: batch_size, I: input_size
+        # state shape: [G*D, B, H], where G: groups, D: num_directions, H: hidden_size
+
+        if h0 is None:
+            dim0, dim1 = input.shape[:2]
+            bs = dim0 if self.batch_first else dim1
+            h0 = self.get_h0(bs, device=input.device)
+        outputs: List[Tensor] = []
+        outstates: List[Tensor] = []
+        for i, layer in enumerate(self.layers):
+            o, s = layer(
+                input[..., i * self.input_size : (i + 1) * self.input_size],
+                h0[i * self.num_directions : (i + 1) * self.num_directions].detach(),
+            )
+            outputs.append(o)
+            outstates.append(s)
+        output = torch.cat(outputs, dim=-1)
+        h = torch.cat(outstates, dim=0)
+        return output, h
+
+
+class GroupedGRU(nn.Module):
+    groups: Final[int]
+    num_layers: Final[int]
+    batch_first: Final[bool]
+    hidden_size: Final[int]
+    bidirectional: Final[bool]
+    num_directions: Final[int]
+    shuffle: Final[bool]
+    add_outputs: Final[bool]
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int = 1,
+        groups: int = 4,
+        bias: bool = True,
+        batch_first: bool = True,
+        dropout: float = 0,
+        bidirectional: bool = False,
+        shuffle: bool = True,
+        add_outputs: bool = False,
+    ):
+        super().__init__()
+        kwargs = {
+            "groups": groups,
+            "bias": bias,
+            "batch_first": batch_first,
+            "dropout": dropout,
+            "bidirectional": bidirectional,
+        }
+        assert input_size % groups == 0
+        assert hidden_size % groups == 0
+        assert num_layers > 0
+        self.input_size = input_size
+        self.groups = groups
+        self.num_layers = num_layers
+        self.batch_first = batch_first
+        self.hidden_size = hidden_size // groups
+        self.bidirectional = bidirectional
+        self.num_directions = 2 if bidirectional else 1
+        if groups == 1:
+            shuffle = False  # Fully connected, no need to shuffle
+        self.shuffle = shuffle
+        self.add_outputs = add_outputs
+        self.grus: List[GroupedGRULayer] = nn.ModuleList()  # type: ignore
+        self.grus.append(GroupedGRULayer(input_size, hidden_size, **kwargs))
+        for _ in range(1, num_layers):
+            self.grus.append(GroupedGRULayer(hidden_size, hidden_size, **kwargs))
+        self.flatten_parameters()
+
+    def flatten_parameters(self):
+        for gru in self.grus:
+            gru.flatten_parameters()
+
+    def get_h0(self, batch_size: int, device: torch.device = torch.device("cpu")) -> Tensor:
+        return torch.zeros(
+            (self.num_layers * self.groups * self.num_directions, batch_size, self.hidden_size),
+            device=device,
+        )
+
+    def forward(self, input: Tensor, state: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
+        dim0, dim1, _ = input.shape
+        b = dim0 if self.batch_first else dim1
+        if state is None:
+            state = self.get_h0(b, input.device)
+        output = torch.zeros(
+            dim0, dim1, self.hidden_size * self.num_directions * self.groups, device=input.device
+        )
+        outstates = []
+        h = self.groups * self.num_directions
+        for i, gru in enumerate(self.grus):
+            input, s = gru(input, state[i * h : (i + 1) * h])
+            outstates.append(s)
+            if self.shuffle and i < self.num_layers - 1:
+                input = (
+                    input.view(dim0, dim1, -1, self.groups).transpose(2, 3).reshape(dim0, dim1, -1)
+                )
+            if self.add_outputs:
+                output += input
+            else:
+                output = input
+        outstate = torch.cat(outstates, dim=0)
+        return output, outstate
+
+
+class SqueezedGRU(nn.Module):
+    input_size: Final[int]
+    hidden_size: Final[int]
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        output_size: Optional[int] = None,
+        num_layers: int = 1,
+        linear_groups: int = 8,
+        batch_first: bool = True,
+        gru_skip_op: Optional[Callable[..., torch.nn.Module]] = None,
+        linear_act_layer: Callable[..., torch.nn.Module] = nn.Identity,
+    ):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.linear_in = nn.Sequential(
+            GroupedLinearEinsum(input_size, hidden_size, linear_groups), linear_act_layer()
+        )
+        self.gru = nn.GRU(hidden_size, hidden_size, num_layers=num_layers, batch_first=batch_first)
+        self.gru_skip = gru_skip_op() if gru_skip_op is not None else None
+        if output_size is not None:
+            self.linear_out = nn.Sequential(
+                GroupedLinearEinsum(hidden_size, output_size, linear_groups), linear_act_layer()
+            )
+        else:
+            self.linear_out = nn.Identity()
+
+    def forward(self, input: Tensor, h=None) -> Tuple[Tensor, Tensor]:
+        input = self.linear_in(input)
+        x, h = self.gru(input, h)
+        if self.gru_skip is not None:
+            x = x + self.gru_skip(input)
+        x = self.linear_out(x)
+        return x, h
+
+
+class GroupedLinearEinsum(nn.Module):
+    input_size: Final[int]
+    hidden_size: Final[int]
+    groups: Final[int]
+
+    def __init__(self, input_size: int, hidden_size: int, groups: int = 1):
+        super().__init__()
+        # self.weight: Tensor
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.groups = groups
+        assert input_size % groups == 0
+        self.ws = input_size // groups
+        self.register_parameter(
+            "weight",
+            Parameter(
+                torch.zeros(groups, input_size // groups, hidden_size // groups), requires_grad=True
+            ),
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))  # type: ignore
+
+    def forward(self, x: Tensor) -> Tensor:
+        # x: [..., I]
+        x = x.unflatten(-1, (self.groups, self.ws))  # [..., G, I/G]
+        x = torch.einsum("...gi,...gih->...gh", x, self.weight)  # [..., G, H/G]
+        x = x.flatten(2, 3)  # [B, T, H]
+        return x
+
+
+class GroupedLinear(nn.Module):
+    input_size: Final[int]
+    hidden_size: Final[int]
+    groups: Final[int]
+    shuffle: Final[bool]
+
+    def __init__(self, input_size: int, hidden_size: int, groups: int = 1, shuffle: bool = True):
+        super().__init__()
+        assert input_size % groups == 0
+        assert hidden_size % groups == 0
+        self.groups = groups
+        self.input_size = input_size // groups
+        self.hidden_size = hidden_size // groups
+        if groups == 1:
+            shuffle = False
+        self.shuffle = shuffle
+        self.layers = nn.ModuleList(
+            nn.Linear(self.input_size, self.hidden_size) for _ in range(groups)
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        outputs: List[Tensor] = []
+        for i, layer in enumerate(self.layers):
+            outputs.append(layer(x[..., i * self.input_size : (i + 1) * self.input_size]))
+        output = torch.cat(outputs, dim=-1)
+        if self.shuffle:
+            orig_shape = output.shape
+            output = (
+                output.view(-1, self.hidden_size, self.groups).transpose(-1, -2).reshape(orig_shape)
+            )
+        return output
+
+
+class LocalSnrTarget(nn.Module):
+    def __init__(
+        self, ws: int = 20, db: bool = True, ws_ns: Optional[int] = None, target_snr_range=None
+    ):
+        super().__init__()
+        self.ws = self.calc_ws(ws)
+        self.ws_ns = self.ws * 2 if ws_ns is None else self.calc_ws(ws_ns)
+        self.db = db
+        self.range = target_snr_range
+
+    def calc_ws(self, ws_ms: int) -> int:
+        # Calculates windows size in stft domain given a window size in ms
+        p = ModelParams()
+        ws = ws_ms - p.fft_size / p.sr * 1000  # length ms of an fft_window
+        ws = 1 + ws / (p.hop_size / p.sr * 1000)  # consider hop_size
+        return max(int(round(ws)), 1)
+
+    def forward(self, clean: Tensor, noise: Tensor, max_bin: Optional[int] = None) -> Tensor:
+        # clean: [B, 1, T, F]
+        # out: [B, T']
+        if max_bin is not None:
+            clean = as_complex(clean[..., :max_bin])
+            noise = as_complex(noise[..., :max_bin])
+        return (
+            local_snr(clean, noise, window_size=self.ws, db=self.db, window_size_ns=self.ws_ns)[0]
+            .clamp(self.range[0], self.range[1])
+            .squeeze(1)
+        )
+
+
+def _local_energy(x: Tensor, ws: int, device: torch.device) -> Tensor:
+    if (ws % 2) == 0:
+        ws += 1
+    ws_half = ws // 2
+    x = F.pad(x.pow(2).sum(-1).sum(-1), (ws_half, ws_half, 0, 0))
+    w = torch.hann_window(ws, device=device, dtype=x.dtype)
+    x = x.unfold(-1, size=ws, step=1) * w
+    return torch.sum(x, dim=-1).div(ws)
+
+
+def local_snr(
+    clean: Tensor,
+    noise: Tensor,
+    window_size: int,
+    db: bool = False,
+    window_size_ns: Optional[int] = None,
+    eps: float = 1e-12,
+) -> Tuple[Tensor, Tensor, Tensor]:
+    # clean shape: [B, C, T, F]
+    clean = as_real(clean)
+    noise = as_real(noise)
+    assert clean.dim() == 5
+
+    E_speech = _local_energy(clean, window_size, clean.device)
+    window_size_ns = window_size if window_size_ns is None else window_size_ns
+    E_noise = _local_energy(noise, window_size_ns, clean.device)
+
+    snr = E_speech / E_noise.clamp_min(eps)
+    if db:
+        snr = snr.clamp_min(eps).log10().mul(10)
+    return snr, E_speech, E_noise
+
+
+def test_grouped_gru():
+    from icecream import ic
+
+    g = 2  # groups
+    h = 4  # hidden_size
+    i = 2  # input_size
+    b = 1  # batch_size
+    t = 5  # time_steps
+    m = GroupedGRULayer(i, h, g, batch_first=True)
+    ic(m)
+    input = torch.randn((b, t, i))
+    h0 = m.get_h0(b)
+    assert list(h0.shape) == [g, b, h // g]
+    out, hout = m(input, h0)
+
+    # Should be exportable as raw nn.Module
+    torch.onnx.export(
+        m, (input, h0), "out/grouped.onnx", example_outputs=(out, hout), opset_version=13
+    )
+    # Should be exportable as traced
+    m = torch.jit.trace(m, (input, h0))
+    torch.onnx.export(
+        m, (input, h0), "out/grouped.onnx", example_outputs=(out, hout), opset_version=13
+    )
+    # and as scripted module
+    m = torch.jit.script(m)
+    torch.onnx.export(
+        m, (input, h0), "out/grouped.onnx", example_outputs=(out, hout), opset_version=13
+    )
+
+    # now grouped gru
+    num = 2
+    m = GroupedGRU(i, h, num, g, batch_first=True, shuffle=True)
+    ic(m)
+    h0 = m.get_h0(b)
+    assert list(h0.shape) == [num * g, b, h // g]
+    out, hout = m(input, h0)
+
+    # Should be exportable as traced
+    m = torch.jit.trace(m, (input, h0))
+    torch.onnx.export(
+        m, (input, h0), "out/grouped.onnx", example_outputs=(out, hout), opset_version=13
+    )
+    # and scripted module
+    m = torch.jit.script(m)
+    torch.onnx.export(
+        m, (input, h0), "out/grouped.onnx", example_outputs=(out, hout), opset_version=13
+    )
+
+
+def test_erb():
+    import libdf
+    from df.config import config
+
+    config.use_defaults()
+    p = ModelParams()
+    n_freq = p.fft_size // 2 + 1
+    df_state = libdf.DF(sr=p.sr, fft_size=p.fft_size, hop_size=p.hop_size, nb_bands=p.nb_erb)
+    erb = erb_fb(df_state.erb_widths(), p.sr)
+    erb_inverse = erb_fb(df_state.erb_widths(), p.sr, inverse=True)
+    input = torch.randn((1, 1, 1, n_freq), dtype=torch.complex64)
+    input_abs = input.abs().square()
+    erb_widths = df_state.erb_widths()
+    df_erb = torch.from_numpy(libdf.erb(input.numpy(), erb_widths, False))
+    py_erb = torch.matmul(input_abs, erb)
+    assert torch.allclose(df_erb, py_erb)
+    df_out = torch.from_numpy(libdf.erb_inv(df_erb.numpy(), erb_widths))
+    py_out = torch.matmul(py_erb, erb_inverse)
+    assert torch.allclose(df_out, py_out)
+
+
+def test_unit_norm():
+    from df.config import config
+    from libdf import unit_norm
+
+    config.use_defaults()
+    p = ModelParams()
+    b = 2
+    F = p.nb_df
+    t = 100
+    spec = torch.randn(b, 1, t, F, 2)
+    alpha = get_norm_alpha(log=False)
+    # Expects complex input of shape [C, T, F]
+    norm_lib = torch.as_tensor(unit_norm(torch.view_as_complex(spec).squeeze(1).numpy(), alpha))
+    m = ExponentialUnitNorm(alpha, F)
+    norm_torch = torch.view_as_complex(m(spec).squeeze(1))
+    assert torch.allclose(norm_lib.real, norm_torch.real)
+    assert torch.allclose(norm_lib.imag, norm_torch.imag)
+    assert torch.allclose(norm_lib.abs(), norm_torch.abs())
+
+
+def test_dfop():
+    from df.config import config
+
+    config.use_defaults()
+    p = ModelParams()
+    f = p.nb_df
+    F = f * 2
+    o = p.df_order
+    d = p.df_lookahead
+    t = 100
+    spec = torch.randn(1, 1, t, F, 2)
+    coefs = torch.randn(1, t, o, f, 2)
+    alpha = torch.randn(1, t, 1)
+    dfop = DfOp(df_bins=p.nb_df)
+    dfop.set_forward("real_loop")
+    out1 = dfop(spec, coefs, alpha)
+    dfop.set_forward("real_strided")
+    out2 = dfop(spec, coefs, alpha)
+    dfop.set_forward("real_unfold")
+    out3 = dfop(spec, coefs, alpha)
+    dfop.set_forward("complex_strided")
+    out4 = dfop(spec, coefs, alpha)
+    torch.testing.assert_allclose(out1, out2)
+    torch.testing.assert_allclose(out1, out3)
+    torch.testing.assert_allclose(out1, out4)
+    # This forward method requires external padding/lookahead as well as spectrogram buffer
+    # handling, i.e. via a ring buffer. Could be used in real time usage.
+    dfop.set_forward("real_one_step")
+    spec_padded = spec_pad(spec, o, d, dim=-3)
+    out5 = torch.zeros_like(out1)
+    for i in range(t):
+        out5[:, :, i] = dfop(
+            spec_padded[:, :, i : i + o], coefs[:, i].unsqueeze(1), alpha[:, i].unsqueeze(1)
+        )
+    torch.testing.assert_allclose(out1, out5)
+    # Forward method that does the padding/lookahead handling using an internal hidden state.
+    dfop.freq_bins = F
+    dfop.set_forward("real_hidden_state_loop")
+    out6 = dfop(spec, coefs, alpha)
+    torch.testing.assert_allclose(out1, out6)

df/multiframe.py

@@ -0,0 +1,329 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Final
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+
+class MultiFrameModule(nn.Module, ABC):
+    """Multi-frame speech enhancement modules.
+
+    Signal model and notation:
+        Noisy: `x = s + n`
+        Enhanced: `y = f(x)`
+        Objective: `min ||s - y||`
+
+        PSD: Power spectral density, notated eg. as `Rxx` for noisy PSD.
+        IFC: Inter-frame correlation vector: PSD*u, u: selection vector. Notated as `rxx`
+    """
+
+    num_freqs: Final[int]
+    frame_size: Final[int]
+    need_unfold: Final[bool]
+
+    def __init__(self, num_freqs: int, frame_size: int, lookahead: int = 0):
+        """Multi-Frame filtering module.
+
+        Args:
+            num_freqs (int): Number of frequency bins used for filtering.
+            frame_size (int): Frame size in FD domain.
+            lookahead (int): Lookahead, may be used to select the output time step. Note: This
+                module does not add additional padding according to lookahead!
+        """
+        super().__init__()
+        self.num_freqs = num_freqs
+        self.frame_size = frame_size
+        self.pad = nn.ConstantPad2d((0, 0, frame_size - 1, 0), 0.0)
+        self.need_unfold = frame_size > 1
+        self.lookahead = lookahead
+
+    def spec_unfold(self, spec: Tensor):
+        """Pads and unfolds the spectrogram according to frame_size.
+
+        Args:
+            spec (complex Tensor): Spectrogram of shape [B, C, T, F]
+        Returns:
+            spec (Tensor): Unfolded spectrogram of shape [B, C, T, F, N], where N: frame_size.
+        """
+        if self.need_unfold:
+            return self.pad(spec).unfold(2, self.frame_size, 1)
+        return spec.unsqueeze(-1)
+
+    def forward(self, spec: Tensor, coefs: Tensor):
+        """Pads and unfolds the spectrogram and forwards to impl.
+
+        Args:
+            spec (Tensor): Spectrogram of shape [B, C, T, F, 2]
+            coefs (Tensor): Spectrogram of shape [B, C, T, F, 2]
+        """
+        spec_u = self.spec_unfold(torch.view_as_complex(spec))
+        coefs = torch.view_as_complex(coefs)
+        spec_f = spec_u.narrow(-2, 0, self.num_freqs)
+        spec_f = self.forward_impl(spec_f, coefs)
+        if self.training:
+            spec = spec.clone()
+        spec[..., : self.num_freqs, :] = torch.view_as_real(spec_f)
+        return spec
+
+    @abstractmethod
+    def forward_impl(self, spec: Tensor, coefs: Tensor) -> Tensor:
+        """Forward impl taking complex spectrogram and coefficients.
+
+        Args:
+            spec (complex Tensor): Spectrogram of shape [B, C1, T, F, N]
+            coefs (complex Tensor): Coefficients [B, C2, T, F]
+
+        Returns:
+            spec (complex Tensor): Enhanced spectrogram of shape [B, C1, T, F]
+        """
+        ...
+
+    @abstractmethod
+    def num_channels(self) -> int:
+        """Return the number of required channels.
+
+        If multiple inputs are required, then all these should be combined in one Tensor containing
+        the summed channels.
+        """
+        ...
+
+
+def psd(x: Tensor, n: int) -> Tensor:
+    """Compute the PSD correlation matrix Rxx for a spectrogram.
+
+    That is, `X*conj(X)`, where `*` is the outer product.
+
+    Args:
+        x (complex Tensor): Spectrogram of shape [B, C, T, F]. Will be unfolded with `n` steps over
+            the time axis.
+
+    Returns:
+        Rxx (complex Tensor): Correlation matrix of shape [B, C, T, F, N, N]
+    """
+    x = F.pad(x, (0, 0, n - 1, 0)).unfold(-2, n, 1)
+    return torch.einsum("...n,...m->...mn", x, x.conj())
+
+
+def df(spec: Tensor, coefs: Tensor) -> Tensor:
+    """Deep filter implemenation using `torch.einsum`. Requires unfolded spectrogram.
+
+    Args:
+        spec (complex Tensor): Spectrogram of shape [B, C, T, F, N]
+        coefs (complex Tensor): Spectrogram of shape [B, C, N, T, F]
+
+    Returns:
+        spec (complex Tensor): Spectrogram of shape [B, C, T, F]
+    """
+    return torch.einsum("...tfn,...ntf->...tf", spec, coefs)
+
+
+class CRM(MultiFrameModule):
+    """Complex ratio mask."""
+
+    def __init__(self, num_freqs: int, frame_size: int = 1, lookahead: int = 0):
+        assert frame_size == 1 and lookahead == 0, (frame_size, lookahead)
+        super().__init__(num_freqs, 1)
+
+    def forward_impl(self, spec: Tensor, coefs: Tensor):
+        return spec.squeeze(-1).mul(coefs)
+
+    def num_channels(self):
+        return 2
+
+
+class DF(MultiFrameModule):
+    conj: Final[bool]
+    """Deep Filtering."""
+
+    def __init__(self, num_freqs: int, frame_size: int, lookahead: int = 0, conj: bool = False):
+        super().__init__(num_freqs, frame_size, lookahead)
+        self.conj = conj
+
+    def forward_impl(self, spec: Tensor, coefs: Tensor):
+        coefs = coefs.view(coefs.shape[0], -1, self.frame_size, *coefs.shape[2:])
+        if self.conj:
+            coefs = coefs.conj()
+        return df(spec, coefs)
+
+    def num_channels(self):
+        return self.frame_size * 2
+
+
+class MfWf(MultiFrameModule):
+    """Multi-frame Wiener filter base module."""
+
+    def __init__(self, num_freqs: int, frame_size: int, lookahead: int = 0):
+        """Multi-frame Wiener Filter.
+
+        Several implementation methods are available resulting in different number of required input
+        coefficient channels.
+
+        Methods:
+            psd_ifc: Predict PSD `Rxx` and IFC `rss`.
+            df: Use deep filtering to predict speech and noisy spectrograms. These will be used for
+                PSD calculation for Wiener filtering. Alias: `df_sx`
+            c: Directly predict Wiener filter coefficients. Computation same as deep filtering.
+
+        """
+        super().__init__(num_freqs, frame_size, lookahead=0)
+        self.idx = -lookahead
+
+    def num_channels(self):
+        return self.num_channels
+
+    @staticmethod
+    def solve(Rxx, rss, diag_eps: float = 1e-8, eps: float = 1e-7) -> Tensor:
+        return torch.einsum(
+            "...nm,...m->...n", torch.inverse(_tik_reg(Rxx, diag_eps, eps)), rss
+        )  # [T, F, N]
+
+    @abstractmethod
+    def mfwf(self, spec: Tensor, coefs: Tensor) -> Tensor:
+        """Multi-frame Wiener filter impl taking complex spectrogram and coefficients.
+
+        Coefficients may be split into multiple parts w.g. for multiple DF coefs or PSDs.
+
+        Args:
+            spec (complex Tensor): Spectrogram of shape [B, C1, T, F, N]
+            coefs (complex Tensor): Coefficients [B, C2, T, F]
+
+        Returns:
+            c (complex Tensor): MfWf coefs of shape [B, C1, T, F, N]
+        """
+        ...
+
+    def forward_impl(self, spec: Tensor, coefs: Tensor) -> Tensor:
+        coefs = self.mfwf(spec, coefs)
+        return self.apply_coefs(spec, coefs)
+
+    @staticmethod
+    def apply_coefs(spec: Tensor, coefs: Tensor) -> Tensor:
+        # spec: [B, C, T, F, N]
+        # coefs: [B, C, T, F, N]
+        return torch.einsum("...n,...n->...", spec, coefs)
+
+
+class MfWfDf(MfWf):
+    eps_diag: Final[float]
+
+    def __init__(
+        self,
+        num_freqs: int,
+        frame_size: int,
+        lookahead: int = 0,
+        eps_diag: float = 1e-7,
+        eps: float = 1e-7,
+    ):
+        super().__init__(num_freqs, frame_size, lookahead)
+        self.eps_diag = eps_diag
+        self.eps = eps
+
+    def num_channels(self):
+        # frame_size/df_order * 2 (x/s) * 2 (re/im)
+        return self.frame_size * 4
+
+    def mfwf(self, spec: Tensor, coefs: Tensor) -> Tensor:
+        coefs.chunk
+        df_s, df_x = torch.chunk(coefs, 2, 1)  # [B, C, T, F, N]
+        df_s = df_s.unflatten(1, (-1, self.frame_size))
+        df_x = df_x.unflatten(1, (-1, self.frame_size))
+        spec_s = df(spec, df_s)  # [B, C, T, F]
+        spec_x = df(spec, df_x)
+        Rss = psd(spec_s, self.frame_size)  # [B, C, T, F, N. N]
+        Rxx = psd(spec_x, self.frame_size)
+        rss = Rss[..., -1]  # TODO: use -1 or self.idx?
+        c = self.solve(Rxx, rss, self.eps_diag, self.eps)  # [B, C, T, F, N]
+        return c
+
+
+class MfWfPsd(MfWf):
+    """Multi-frame Wiener filter by predicting noisy PSD `Rxx` and speech IFC `rss`."""
+
+    def num_channels(self):
+        # (Rxx + rss) * 2 (re/im)
+        return (self.frame_size**2 + self.frame_size) * 2
+
+    def mfwf(self, spec: Tensor, coefs: Tensor) -> Tensor:  # type: ignore
+        Rxx, rss = torch.split(coefs.movedim(1, -1), [self.frame_size**2, self.frame_size], -1)
+        c = self.solve(Rxx.unflatten(-1, (self.frame_size, self.frame_size)), rss)
+        return c
+
+
+class MfWfC(MfWf):
+    """Multi-frame Wiener filter by directly predicting the MfWf coefficients."""
+
+    def num_channels(self):
+        # mfwf coefs * 2 (re/im)
+        return self.frame_size * 2
+
+    def mfwf(self, spec: Tensor, coefs: Tensor) -> Tensor:  # type: ignore
+        coefs = coefs.unflatten(1, (-1, self.frame_size)).permute(
+            0, 1, 3, 4, 2
+        )  # [B, C*N, T, F] -> [B, C, T, F, N]
+        return coefs
+
+
+class MvdrSouden(MultiFrameModule):
+    def __init__(self, num_freqs: int, frame_size: int, lookahead: int = 0):
+        super().__init__(num_freqs, frame_size, lookahead)
+
+
+class MvdrEvd(MultiFrameModule):
+    def __init__(self, num_freqs: int, frame_size: int, lookahead: int = 0):
+        super().__init__(num_freqs, frame_size, lookahead)
+
+
+class MvdrRtfPower(MultiFrameModule):
+    def __init__(self, num_freqs: int, frame_size: int, lookahead: int = 0):
+        super().__init__(num_freqs, frame_size, lookahead)
+
+
+MF_METHODS: Dict[str, MultiFrameModule] = {
+    "crm": CRM,
+    "df": DF,
+    "mfwf_df": MfWfDf,
+    "mfwf_df_sx": MfWfDf,
+    "mfwf_psd": MfWfPsd,
+    "mfwf_psd_ifc": MfWfPsd,
+    "mfwf_c": MfWfC,
+}
+
+
+# From torchaudio
+def _compute_mat_trace(input: torch.Tensor, dim1: int = -1, dim2: int = -2) -> torch.Tensor:
+    r"""Compute the trace of a Tensor along ``dim1`` and ``dim2`` dimensions.
+    Args:
+        input (torch.Tensor): Tensor of dimension `(..., channel, channel)`
+        dim1 (int, optional): the first dimension of the diagonal matrix
+            (Default: -1)
+        dim2 (int, optional): the second dimension of the diagonal matrix
+            (Default: -2)
+    Returns:
+        Tensor: trace of the input Tensor
+    """
+    assert input.ndim >= 2, "The dimension of the tensor must be at least 2."
+    assert (
+        input.shape[dim1] == input.shape[dim2]
+    ), "The size of ``dim1`` and ``dim2`` must be the same."
+    input = torch.diagonal(input, 0, dim1=dim1, dim2=dim2)
+    return input.sum(dim=-1)
+
+
+def _tik_reg(mat: torch.Tensor, reg: float = 1e-7, eps: float = 1e-8) -> torch.Tensor:
+    """Perform Tikhonov regularization (only modifying real part).
+    Args:
+        mat (torch.Tensor): input matrix (..., channel, channel)
+        reg (float, optional): regularization factor (Default: 1e-8)
+        eps (float, optional): a value to avoid the correlation matrix is all-zero (Default: ``1e-8``)
+    Returns:
+        Tensor: regularized matrix (..., channel, channel)
+    """
+    # Add eps
+    C = mat.size(-1)
+    eye = torch.eye(C, dtype=mat.dtype, device=mat.device)
+    epsilon = _compute_mat_trace(mat).real[..., None, None] * reg
+    # in case that correlation_matrix is all-zero
+    epsilon = epsilon + eps
+    mat = mat + epsilon * eye[..., :, :]
+    return mat

df/utils.py

@@ -0,0 +1,230 @@
+import collections
+import math
+import os
+import random
+import subprocess
+from socket import gethostname
+from typing import Any, Dict, Set, Tuple, Union
+
+import numpy as np
+import torch
+from loguru import logger
+from torch import Tensor
+#from torch._six import string_classes
+from torch.autograd import Function
+from torch.types import Number
+
+from df.config import config
+from df.model import ModelParams
+
+try:
+    from torchaudio.functional import resample as ta_resample
+except ImportError:
+    from torchaudio.compliance.kaldi import resample_waveform as ta_resample  # type: ignore
+
+
+def get_resample_params(method: str) -> Dict[str, Any]:
+    params = {
+        "sinc_fast": {"resampling_method": "sinc_interpolation", "lowpass_filter_width": 16},
+        "sinc_best": {"resampling_method": "sinc_interpolation", "lowpass_filter_width": 64},
+        "kaiser_fast": {
+            "resampling_method": "kaiser_window",
+            "lowpass_filter_width": 16,
+            "rolloff": 0.85,
+            "beta": 8.555504641634386,
+        },
+        "kaiser_best": {
+            "resampling_method": "kaiser_window",
+            "lowpass_filter_width": 16,
+            "rolloff": 0.9475937167399596,
+            "beta": 14.769656459379492,
+        },
+    }
+    assert method in params.keys(), f"method must be one of {list(params.keys())}"
+    return params[method]
+
+
+def resample(audio: Tensor, orig_sr: int, new_sr: int, method="sinc_fast"):
+    params = get_resample_params(method)
+    return ta_resample(audio, orig_sr, new_sr, **params)
+
+
+def get_device():
+    s = config("DEVICE", default="", section="train")
+    if s == "":
+        if torch.cuda.is_available():
+            DEVICE = torch.device("cuda:0")
+        else:
+            DEVICE = torch.device("cpu")
+    else:
+        DEVICE = torch.device(s)
+    return DEVICE
+
+
+def as_complex(x: Tensor):
+    if torch.is_complex(x):
+        return x
+    if x.shape[-1] != 2:
+        raise ValueError(f"Last dimension need to be of length 2 (re + im), but got {x.shape}")
+    if x.stride(-1) != 1:
+        x = x.contiguous()
+    return torch.view_as_complex(x)
+
+
+def as_real(x: Tensor):
+    if torch.is_complex(x):
+        return torch.view_as_real(x)
+    return x
+
+
+class angle_re_im(Function):
+    """Similar to torch.angle but robustify the gradient for zero magnitude."""
+
+    @staticmethod
+    def forward(ctx, re: Tensor, im: Tensor):
+        ctx.save_for_backward(re, im)
+        return torch.atan2(im, re)
+
+    @staticmethod
+    def backward(ctx, grad: Tensor) -> Tuple[Tensor, Tensor]:
+        re, im = ctx.saved_tensors
+        grad_inv = grad / (re.square() + im.square()).clamp_min_(1e-10)
+        return -im * grad_inv, re * grad_inv
+
+
+class angle(Function):
+    """Similar to torch.angle but robustify the gradient for zero magnitude."""
+
+    @staticmethod
+    def forward(ctx, x: Tensor):
+        ctx.save_for_backward(x)
+        return torch.atan2(x.imag, x.real)
+
+    @staticmethod
+    def backward(ctx, grad: Tensor):
+        (x,) = ctx.saved_tensors
+        grad_inv = grad / (x.real.square() + x.imag.square()).clamp_min_(1e-10)
+        return torch.view_as_complex(torch.stack((-x.imag * grad_inv, x.real * grad_inv), dim=-1))
+
+
+def check_finite_module(obj, name="Module", _raise=True) -> Set[str]:
+    out: Set[str] = set()
+    if isinstance(obj, torch.nn.Module):
+        for name, child in obj.named_children():
+            out = out | check_finite_module(child, name)
+        for name, param in obj.named_parameters():
+            out = out | check_finite_module(param, name)
+        for name, buf in obj.named_buffers():
+            out = out | check_finite_module(buf, name)
+    if _raise and len(out) > 0:
+        raise ValueError(f"{name} not finite during checkpoint writing including: {out}")
+    return out
+
+
+def make_np(x: Union[Tensor, np.ndarray, Number]) -> np.ndarray:
+    """Transforms Tensor to numpy.
+    Args:
+      x: An instance of torch tensor or caffe blob name
+
+    Returns:
+        numpy.array: Numpy array
+    """
+    if isinstance(x, np.ndarray):
+        return x
+    if np.isscalar(x):
+        return np.array([x])
+    if isinstance(x, Tensor):
+        return x.detach().cpu().numpy()
+    raise NotImplementedError(
+        "Got {}, but numpy array, scalar, or torch tensor are expected.".format(type(x))
+    )
+
+
+def get_norm_alpha(log: bool = True) -> float:
+    p = ModelParams()
+    a_ = _calculate_norm_alpha(sr=p.sr, hop_size=p.hop_size, tau=p.norm_tau)
+    precision = 3
+    a = 1.0
+    while a >= 1.0:
+        a = round(a_, precision)
+        precision += 1
+    if log:
+        logger.info(f"Running with normalization window alpha = '{a}'")
+    return a
+
+
+def _calculate_norm_alpha(sr: int, hop_size: int, tau: float):
+    """Exponential decay factor alpha for a given tau (decay window size [s])."""
+    dt = hop_size / sr
+    return math.exp(-dt / tau)
+
+
+def check_manual_seed(seed: int = None):
+    """If manual seed is not specified, choose a random one and communicate it to the user."""
+    seed = seed or random.randint(1, 10000)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)
+    return seed
+
+
+def get_git_root():
+    git_local_dir = os.path.dirname(os.path.abspath(__file__))
+    args = ["git", "-C", git_local_dir, "rev-parse", "--show-toplevel"]
+    return subprocess.check_output(args).strip().decode()
+
+
+def get_commit_hash():
+    """Returns the current git commit."""
+    try:
+        git_dir = get_git_root()
+        args = ["git", "-C", git_dir, "rev-parse", "--short", "--verify", "HEAD"]
+        commit = subprocess.check_output(args).strip().decode()
+    except subprocess.CalledProcessError:
+        # probably not in git repo
+        commit = None
+    return commit
+
+
+def get_host() -> str:
+    return gethostname()
+
+
+def get_branch_name():
+    try:
+        git_dir = os.path.dirname(os.path.abspath(__file__))
+        args = ["git", "-C", git_dir, "rev-parse", "--abbrev-ref", "HEAD"]
+        branch = subprocess.check_output(args).strip().decode()
+    except subprocess.CalledProcessError:
+        # probably not in git repo
+        branch = None
+    return branch
+
+
+# from pytorch/ignite:
+def apply_to_tensor(input_, func):
+    """Apply a function on a tensor or mapping, or sequence of tensors."""
+    if isinstance(input_, torch.nn.Module):
+        return [apply_to_tensor(c, func) for c in input_.children()]
+    elif isinstance(input_, torch.nn.Parameter):
+        return func(input_.data)
+    elif isinstance(input_, Tensor):
+        return func(input_)
+    elif isinstance(input_, str):
+        return input_
+    elif isinstance(input_, collections.Mapping):
+        return {k: apply_to_tensor(sample, func) for k, sample in input_.items()}
+    elif isinstance(input_, collections.Iterable):
+        return [apply_to_tensor(sample, func) for sample in input_]
+    elif input_ is None:
+        return input_
+    else:
+        return input_
+
+
+def detach_hidden(hidden: Any) -> Any:
+    """Cut backpropagation graph.
+    Auxillary function to cut the backpropagation graph by detaching the hidden
+    vector.
+    """
+    return apply_to_tensor(hidden, Tensor.detach)

libdf/__init__.py

@@ -0,0 +1,3 @@
+from .libdf import *
+
+__doc__ = libdf.__doc__

libdf/__init__.pyi

@@ -0,0 +1,57 @@
+from typing import List, Optional, Union
+
+from numpy import ndarray
+
+class DF:
+    def __init__(
+        self,
+        sr: int,
+        fft_size: int,
+        hop_size: int,
+        nb_bands: int,
+        min_nb_erb_freqs: Optional[int] = 1,
+    ):
+        """DeepFilter state used for analysis and synthesis.
+
+        Args:
+            sr (int): Sampling rate.
+            fft_size (int): Window length used for the Fast Fourier transform.
+            hop_size (int): Hop size between two analysis windows. Also called frame size.
+            nb_bands (int): Number of ERB bands.
+            min_nb_erb_freqs (int): Minimum number of frequency bands per ERB band. Defaults to 1.
+        """
+        ...
+    def analysis(self, input: ndarray) -> ndarray:
+        """Analysis of a time-domain signal.
+
+        Args:
+            input (ndarray): 2D real-valued array of shape [C, T].
+        Output:
+            output (ndarray): 3D complex-valued array of shape [C, T', F], where F is the `fft_size`,
+                and T' the original time T divided by `hop_size`.
+        """
+        ...
+    def synthesis(self, input: ndarray) -> ndarray:
+        """Synthesis of a frequency-domain signal.
+
+        Args:
+            input (ndarray): 3D complex-valued array of shape [C, T, F].
+        Output:
+            output (ndarray): 2D real-valued array of shape [C, T].
+        """
+        ...
+    def erb_widths(self) -> ndarray: ...
+    def fft_window(self) -> ndarray: ...
+    def sr(self) -> int: ...
+    def fft_size(self) -> int: ...
+    def hop_size(self) -> int: ...
+    def nb_erb(self) -> int: ...
+    def reset(self) -> None: ...
+
+def erb(
+    input: ndarray, erb_fb: Union[ndarray, List[int]], db: Optional[bool] = None
+) -> ndarray: ...
+def erb_inv(input: ndarray, erb_fb: Union[ndarray, List[int]]) -> ndarray: ...
+def erb_norm(erb: ndarray, alpha: float, state: Optional[ndarray] = None) -> ndarray: ...
+def unit_norm(spec: ndarray, alpha: float, state: Optional[ndarray] = None) -> ndarray: ...
+def unit_norm_init(num_freq_bins: int) -> ndarray: ...

model_weights/voice_enhance/checkpoints/model_96.ckpt.best

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb5eccb429e675bb4ec5ec9e280f048bfff9787b40bd3eb835fd11509eb14a3e
+size 9397209

model_weights/voice_enhance/ckpt/pretrained_bak_5805000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc7ff82ef75becd495aab2ede3a8220da393a717f178ae9534df355a6173bbca
+size 17090379

model_weights/voice_enhance/ckpt/pretrained_bak_5805000.pt.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7dfd48d0da24db35ee4a653d0d36a4104cb26873050a5c3584675eee21937621
+size 69

model_weights/voiceover/freevc-24.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:872360b61e6bbe09bec29810e7ad0d16318e379f6195a7ff3b06e50efb08ad31
+size 1264

model_weights/voiceover/freevc-24.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b39a86fefbc9ec6e30be8d26ee2a6aa5ffe6d235f6ab15773d01cdf348e5b20
+size 472644351

model_weights/wavlm_models/WavLM-Large.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6fb4b3c3e6aa567f0a997b30855859cb81528ee8078802af439f7b2da0bf100f
+size 1261965425

model_weights/wavlm_models/WavLM-Large.pt.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9836bca8ab0e9d0b4797aa78f41b367800d26cfd25ade7b1edcb35bc3c171e4
+size 52

nnet/attentions.py

@@ -0,0 +1,300 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from nnet import commons
+from nnet.modules import LayerNorm
+   
+
+class Encoder(nn.Module):
+  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs):
+    super().__init__()
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.window_size = window_size
+
+    self.drop = nn.Dropout(p_dropout)
+    self.attn_layers = nn.ModuleList()
+    self.norm_layers_1 = nn.ModuleList()
+    self.ffn_layers = nn.ModuleList()
+    self.norm_layers_2 = nn.ModuleList()
+    for i in range(self.n_layers):
+      self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
+      self.norm_layers_1.append(LayerNorm(hidden_channels))
+      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
+      self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+  def forward(self, x, x_mask):
+    attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+    x = x * x_mask
+    for i in range(self.n_layers):
+      y = self.attn_layers[i](x, x, attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_1[i](x + y)
+
+      y = self.ffn_layers[i](x, x_mask)
+      y = self.drop(y)
+      x = self.norm_layers_2[i](x + y)
+    x = x * x_mask
+    return x
+
+
+class Decoder(nn.Module):
+  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., proximal_bias=False, proximal_init=True, **kwargs):
+    super().__init__()
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.proximal_bias = proximal_bias
+    self.proximal_init = proximal_init
+
+    self.drop = nn.Dropout(p_dropout)
+    self.self_attn_layers = nn.ModuleList()
+    self.norm_layers_0 = nn.ModuleList()
+    self.encdec_attn_layers = nn.ModuleList()
+    self.norm_layers_1 = nn.ModuleList()
+    self.ffn_layers = nn.ModuleList()
+    self.norm_layers_2 = nn.ModuleList()
+    for i in range(self.n_layers):
+      self.self_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias, proximal_init=proximal_init))
+      self.norm_layers_0.append(LayerNorm(hidden_channels))
+      self.encdec_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout))
+      self.norm_layers_1.append(LayerNorm(hidden_channels))
+      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True))
+      self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+  def forward(self, x, x_mask, h, h_mask):
+    """
+    x: decoder input
+    h: encoder output
+    """
+    self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype)
+    encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+    x = x * x_mask
+    for i in range(self.n_layers):
+      y = self.self_attn_layers[i](x, x, self_attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_0[i](x + y)
+
+      y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_1[i](x + y)
+      
+      y = self.ffn_layers[i](x, x_mask)
+      y = self.drop(y)
+      x = self.norm_layers_2[i](x + y)
+    x = x * x_mask
+    return x
+
+
+class MultiHeadAttention(nn.Module):
+  def __init__(self, channels, out_channels, n_heads, p_dropout=0., window_size=None, heads_share=True, block_length=None, proximal_bias=False, proximal_init=False):
+    super().__init__()
+    assert channels % n_heads == 0
+
+    self.channels = channels
+    self.out_channels = out_channels
+    self.n_heads = n_heads
+    self.p_dropout = p_dropout
+    self.window_size = window_size
+    self.heads_share = heads_share
+    self.block_length = block_length
+    self.proximal_bias = proximal_bias
+    self.proximal_init = proximal_init
+    self.attn = None
+
+    self.k_channels = channels // n_heads
+    self.conv_q = nn.Conv1d(channels, channels, 1)
+    self.conv_k = nn.Conv1d(channels, channels, 1)
+    self.conv_v = nn.Conv1d(channels, channels, 1)
+    self.conv_o = nn.Conv1d(channels, out_channels, 1)
+    self.drop = nn.Dropout(p_dropout)
+
+    if window_size is not None:
+      n_heads_rel = 1 if heads_share else n_heads
+      rel_stddev = self.k_channels**-0.5
+      self.emb_rel_k = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
+      self.emb_rel_v = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
+
+    nn.init.xavier_uniform_(self.conv_q.weight)
+    nn.init.xavier_uniform_(self.conv_k.weight)
+    nn.init.xavier_uniform_(self.conv_v.weight)
+    if proximal_init:
+      with torch.no_grad():
+        self.conv_k.weight.copy_(self.conv_q.weight)
+        self.conv_k.bias.copy_(self.conv_q.bias)
+      
+  def forward(self, x, c, attn_mask=None):
+    q = self.conv_q(x)
+    k = self.conv_k(c)
+    v = self.conv_v(c)
+    
+    x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+    x = self.conv_o(x)
+    return x
+
+  def attention(self, query, key, value, mask=None):
+    # reshape [b, d, t] -> [b, n_h, t, d_k]
+    b, d, t_s, t_t = (*key.size(), query.size(2))
+    query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+    key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+    value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+    scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+    if self.window_size is not None:
+      assert t_s == t_t, "Relative attention is only available for self-attention."
+      key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+      rel_logits = self._matmul_with_relative_keys(query /math.sqrt(self.k_channels), key_relative_embeddings)
+      scores_local = self._relative_position_to_absolute_position(rel_logits)
+      scores = scores + scores_local
+    if self.proximal_bias:
+      assert t_s == t_t, "Proximal bias is only available for self-attention."
+      scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
+    if mask is not None:
+      scores = scores.masked_fill(mask == 0, -1e4)
+      if self.block_length is not None:
+        assert t_s == t_t, "Local attention is only available for self-attention."
+        block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length)
+        scores = scores.masked_fill(block_mask == 0, -1e4)
+    p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s]
+    p_attn = self.drop(p_attn)
+    output = torch.matmul(p_attn, value)
+    if self.window_size is not None:
+      relative_weights = self._absolute_position_to_relative_position(p_attn)
+      value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
+      output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
+    output = output.transpose(2, 3).contiguous().view(b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t]
+    return output, p_attn
+
+  def _matmul_with_relative_values(self, x, y):
+    """
+    x: [b, h, l, m]
+    y: [h or 1, m, d]
+    ret: [b, h, l, d]
+    """
+    ret = torch.matmul(x, y.unsqueeze(0))
+    return ret
+
+  def _matmul_with_relative_keys(self, x, y):
+    """
+    x: [b, h, l, d]
+    y: [h or 1, m, d]
+    ret: [b, h, l, m]
+    """
+    ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+    return ret
+
+  def _get_relative_embeddings(self, relative_embeddings, length):
+    max_relative_position = 2 * self.window_size + 1
+    # Pad first before slice to avoid using cond ops.
+    pad_length = max(length - (self.window_size + 1), 0)
+    slice_start_position = max((self.window_size + 1) - length, 0)
+    slice_end_position = slice_start_position + 2 * length - 1
+    if pad_length > 0:
+      padded_relative_embeddings = F.pad(
+          relative_embeddings,
+          commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
+    else:
+      padded_relative_embeddings = relative_embeddings
+    used_relative_embeddings = padded_relative_embeddings[:,slice_start_position:slice_end_position]
+    return used_relative_embeddings
+
+  def _relative_position_to_absolute_position(self, x):
+    """
+    x: [b, h, l, 2*l-1]
+    ret: [b, h, l, l]
+    """
+    batch, heads, length, _ = x.size()
+    # Concat columns of pad to shift from relative to absolute indexing.
+    x = F.pad(x, commons.convert_pad_shape([[0,0],[0,0],[0,0],[0,1]]))
+
+    # Concat extra elements so to add up to shape (len+1, 2*len-1).
+    x_flat = x.view([batch, heads, length * 2 * length])
+    x_flat = F.pad(x_flat, commons.convert_pad_shape([[0,0],[0,0],[0,length-1]]))
+
+    # Reshape and slice out the padded elements.
+    x_final = x_flat.view([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:]
+    return x_final
+
+  def _absolute_position_to_relative_position(self, x):
+    """
+    x: [b, h, l, l]
+    ret: [b, h, l, 2*l-1]
+    """
+    batch, heads, length, _ = x.size()
+    # padd along column
+    x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]]))
+    x_flat = x.view([batch, heads, length**2 + length*(length -1)])
+    # add 0's in the beginning that will skew the elements after reshape
+    x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
+    x_final = x_flat.view([batch, heads, length, 2*length])[:,:,:,1:]
+    return x_final
+
+  def _attention_bias_proximal(self, length):
+    """Bias for self-attention to encourage attention to close positions.
+    Args:
+      length: an integer scalar.
+    Returns:
+      a Tensor with shape [1, 1, length, length]
+    """
+    r = torch.arange(length, dtype=torch.float32)
+    diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+    return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+  def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.filter_channels = filter_channels
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.activation = activation
+    self.causal = causal
+
+    if causal:
+      self.padding = self._causal_padding
+    else:
+      self.padding = self._same_padding
+
+    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+    self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+    self.drop = nn.Dropout(p_dropout)
+
+  def forward(self, x, x_mask):
+    x = self.conv_1(self.padding(x * x_mask))
+    if self.activation == "gelu":
+      x = x * torch.sigmoid(1.702 * x)
+    else:
+      x = torch.relu(x)
+    x = self.drop(x)
+    x = self.conv_2(self.padding(x * x_mask))
+    return x * x_mask
+  
+  def _causal_padding(self, x):
+    if self.kernel_size == 1:
+      return x
+    pad_l = self.kernel_size - 1
+    pad_r = 0
+    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+    x = F.pad(x, commons.convert_pad_shape(padding))
+    return x
+
+  def _same_padding(self, x):
+    if self.kernel_size == 1:
+      return x
+    pad_l = (self.kernel_size - 1) // 2
+    pad_r = self.kernel_size // 2
+    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+    x = F.pad(x, commons.convert_pad_shape(padding))
+    return x