diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..b591e62dcbe280caf1fd6d090c3184364634bbcf 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,68 @@ saved_model/**/* 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
+assets/TIGER.png filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample1/dialog.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample1/effect.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample1/mixture.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample1/music.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample2/dialog.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample2/effect.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample2/mixture.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample2/music.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample3/dialog.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample3/effect.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample3/mixture.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr-demo/sample3/music.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/dnr.png filter=lfs diff=lfs merge=lfs -text
+assets/efficiency.png filter=lfs diff=lfs merge=lfs -text
+assets/logo.png filter=lfs diff=lfs merge=lfs -text
+assets/result.png filter=lfs diff=lfs merge=lfs -text
+assets/sample1/GroundTruth/mix.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample1/GroundTruth/s1.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample1/GroundTruth/s2.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample1/TFGNet/s1.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample1/TFGNet/s2.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample1/TIGER/s1.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample1/TIGER/s2.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample1/spec/TFGNet_s1.png filter=lfs diff=lfs merge=lfs -text
+assets/sample1/spec/TFGNet_s2.png filter=lfs diff=lfs merge=lfs -text
+assets/sample1/spec/TIGER_s1.png filter=lfs diff=lfs merge=lfs -text
+assets/sample1/spec/TIGER_s2.png filter=lfs diff=lfs merge=lfs -text
+assets/sample1/spec/ground_truth_s1.png filter=lfs diff=lfs merge=lfs -text
+assets/sample1/spec/ground_truth_s2.png filter=lfs diff=lfs merge=lfs -text
+assets/sample2/GroundTruth/mix.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample2/GroundTruth/s1.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample2/GroundTruth/s2.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample2/TFGNet/s1.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample2/TFGNet/s2.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample2/TIGER/s1.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample2/TIGER/s2.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample2/spec/TFGNet_s1.png filter=lfs diff=lfs merge=lfs -text
+assets/sample2/spec/TFGNet_s2.png filter=lfs diff=lfs merge=lfs -text
+assets/sample2/spec/TIGER_s1.png filter=lfs diff=lfs merge=lfs -text
+assets/sample2/spec/TIGER_s2.png filter=lfs diff=lfs merge=lfs -text
+assets/sample2/spec/ground_truth_s1.png filter=lfs diff=lfs merge=lfs -text
+assets/sample2/spec/ground_truth_s2.png filter=lfs diff=lfs merge=lfs -text
+assets/sample3/GroundTruth/mix.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample3/GroundTruth/s1.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample3/GroundTruth/s2.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample3/TFGNet/s1.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample3/TFGNet/s2.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample3/TIGER/s1.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample3/TIGER/s2.wav filter=lfs diff=lfs merge=lfs -text
+assets/sample3/spec/TFGNet_s1.png filter=lfs diff=lfs merge=lfs -text
+assets/sample3/spec/TFGNet_s2.png filter=lfs diff=lfs merge=lfs -text
+assets/sample3/spec/TIGER_s1.png filter=lfs diff=lfs merge=lfs -text
+assets/sample3/spec/TIGER_s2.png filter=lfs diff=lfs merge=lfs -text
+assets/sample3/spec/ground_truth_s1.png filter=lfs diff=lfs merge=lfs -text
+assets/sample3/spec/ground_truth_s2.png filter=lfs diff=lfs merge=lfs -text
+test/mix.wav filter=lfs diff=lfs merge=lfs -text
+test/s1.wav filter=lfs diff=lfs merge=lfs -text
+test/s2.wav filter=lfs diff=lfs merge=lfs -text
+test/spk1.wav filter=lfs diff=lfs merge=lfs -text
+test/spk2.wav filter=lfs diff=lfs merge=lfs -text
+test/test_mixture_466.wav filter=lfs diff=lfs merge=lfs -text
+test/test_target_dialog_466.wav filter=lfs diff=lfs merge=lfs -text
+test/test_target_effect_466.wav filter=lfs diff=lfs merge=lfs -text
+test/test_target_music_466.wav filter=lfs diff=lfs merge=lfs -text
diff --git a/DataPreProcess/preprocess_lrs2_audio.py b/DataPreProcess/preprocess_lrs2_audio.py
new file mode 100644
index 0000000000000000000000000000000000000000..2b51deeb067dae98ba2d860af933991ade5245be
--- /dev/null
+++ b/DataPreProcess/preprocess_lrs2_audio.py
@@ -0,0 +1,71 @@
+import argparse
+import json
+import os
+import soundfile as sf
+from tqdm import tqdm
+
+
+def get_mouth_path(in_mouth_dir, wav_file, spk, data_type):
+ wav_file = wav_file.replace(".wav", "").split("_")
+ if spk == "s1":
+ file_path = os.path.join(
+ in_mouth_dir, "{}_{}.npz".format(wav_file[0], wav_file[1])
+ )
+ else:
+ file_path = os.path.join(
+ in_mouth_dir, "{}_{}.npz".format(wav_file[3], wav_file[4])
+ )
+ return file_path
+
+
+def preprocess_one_dir(in_data_dir, out_dir, data_type, spk):
+ """Create .json file for one condition."""
+ file_infos = []
+ in_dir = os.path.abspath(os.path.join(in_data_dir, data_type, spk))
+ wav_list = os.listdir(in_dir)
+ wav_list.sort()
+ for wav_file in tqdm(wav_list):
+ if not wav_file.endswith(".wav"):
+ continue
+ wav_path = os.path.join(in_dir, wav_file)
+ samples = sf.SoundFile(wav_path)
+ if spk == "mix":
+ file_infos.append((wav_path, len(samples)))
+ else:
+ file_infos.append(
+ (
+ wav_path,
+ # get_mouth_path(os.path.join(in_data_dir, data_type, 'mouths'), wav_file, spk, data_type),
+ len(samples),
+ )
+ )
+ if not os.path.exists(os.path.join(out_dir, data_type)):
+ os.makedirs(os.path.join(out_dir, data_type))
+ with open(os.path.join(out_dir, data_type, spk + ".json"), "w") as f:
+ json.dump(file_infos, f, indent=4)
+
+
+def preprocess_lrs2_audio(inp_args):
+ """Create .json files for all conditions."""
+ speaker_list = ["mix", "s1", "s2"]
+ for data_type in ["tr", "cv", "tt"]:
+ for spk in speaker_list:
+ preprocess_one_dir(
+ inp_args.in_dir, inp_args.out_dir, data_type, spk,
+ )
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser("LRS2 audio data preprocessing")
+ parser.add_argument(
+ "--in_dir",
+ type=str,
+ default=None,
+ help="Directory path of audio including tr, cv and tt",
+ )
+ parser.add_argument(
+ "--out_dir", type=str, default=None, help="Directory path to put output files"
+ )
+ args = parser.parse_args()
+ print(args)
+ preprocess_lrs2_audio(args)
diff --git a/DataPreProcess/process_echoset.py b/DataPreProcess/process_echoset.py
new file mode 100644
index 0000000000000000000000000000000000000000..a384042e7e415dce72281f2dd4b65109b873053b
--- /dev/null
+++ b/DataPreProcess/process_echoset.py
@@ -0,0 +1,75 @@
+import argparse
+import json
+import os
+import soundfile as sf
+from tqdm import tqdm
+from rich import print
+
+
+def preprocess_one_dir(in_data_dir, out_dir, data_type):
+ """Create .json file for one condition."""
+ mix_infos = []
+ s1_infos = []
+ s2_infos = []
+ in_dir = os.path.abspath(os.path.join(in_data_dir, data_type))
+ print("Process {} set...".format(data_type))
+ for root, dirs, files in os.walk(in_dir):
+ for file in files:
+ if file.endswith(".wav") and file.startswith("mix"):
+ file_path = os.path.join(root, file)
+ audio, _ = sf.read(file_path)
+ mix_infos.append((
+ file_path,
+ len(audio),
+ ))
+
+ file_path = file_path.replace("mix", "spk1_reverb")
+ audio, _ = sf.read(file_path)
+ s1_infos.append((
+ file_path,
+ len(audio),
+ ))
+
+ file_path = file_path.replace("spk1_reverb", "spk2_reverb")
+ audio, _ = sf.read(file_path)
+ s2_infos.append((
+ file_path,
+ len(audio),
+ ))
+ print("Process num: {}".format(len(mix_infos)), end="\r")
+
+ if not os.path.exists(os.path.join(out_dir, data_type)):
+ os.makedirs(os.path.join(out_dir, data_type))
+ with open(os.path.join(out_dir, data_type, "mix.json"), "w") as f:
+ json.dump(mix_infos, f, indent=4)
+
+ with open(os.path.join(out_dir, data_type, "s1.json"), "w") as f:
+ json.dump(s1_infos, f, indent=4)
+
+ with open(os.path.join(out_dir, data_type, "s2.json"), "w") as f:
+ json.dump(s2_infos, f, indent=4)
+
+
+def preprocess_lrs2_audio(inp_args):
+ """Create .json files for all conditions."""
+ for data_type in ["train", "val", "test"]:
+ preprocess_one_dir(
+ inp_args.in_dir, inp_args.out_dir, data_type
+ )
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser("LRS2 audio data preprocessing")
+ parser.add_argument(
+ "--in_dir",
+ type=str,
+ default=None,
+ help="Directory path of audio including tr, cv and tt",
+ )
+ parser.add_argument(
+ "--out_dir", type=str, default=None, help="Directory path to put output files"
+ )
+ args = parser.parse_args()
+ print(args)
+ preprocess_lrs2_audio(args)
+
diff --git a/DataPreProcess/process_librimix.py b/DataPreProcess/process_librimix.py
new file mode 100644
index 0000000000000000000000000000000000000000..f5ce29411890fd2e84717f5cb1e3b1462787e2f9
--- /dev/null
+++ b/DataPreProcess/process_librimix.py
@@ -0,0 +1,57 @@
+import argparse
+import json
+import os
+import soundfile as sf
+from tqdm import tqdm
+
+
+def preprocess_one_dir(in_data_dir, out_dir, data_type, spk):
+ """Create .json file for one condition."""
+ file_infos = []
+ in_dir = os.path.abspath(os.path.join(in_data_dir, data_type, spk))
+ wav_list = os.listdir(in_dir)
+ wav_list.sort()
+ for wav_file in tqdm(wav_list):
+ if not wav_file.endswith(".wav"):
+ continue
+ wav_path = os.path.join(in_dir, wav_file)
+ samples = sf.SoundFile(wav_path)
+ if spk == "mix":
+ file_infos.append((wav_path, len(samples)))
+ else:
+ file_infos.append(
+ (
+ wav_path,
+ len(samples),
+ )
+ )
+ if not os.path.exists(os.path.join(out_dir, data_type)):
+ os.makedirs(os.path.join(out_dir, data_type))
+ with open(os.path.join(out_dir, data_type, spk + ".json"), "w") as f:
+ json.dump(file_infos, f, indent=4)
+
+
+def preprocess_librimix_audio(inp_args):
+ """Create .json files for all conditions."""
+ speaker_list = ["mix_both", "s1", "s2"]
+ for data_type in ["train-100", "dev", "test"]:
+ for spk in speaker_list:
+ preprocess_one_dir(
+ inp_args.in_dir, inp_args.out_dir, data_type, spk,
+ )
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser("Librimix audio data preprocessing")
+ parser.add_argument(
+ "--in_dir",
+ type=str,
+ default=None,
+ help="Directory path of audio including tr, cv and tt",
+ )
+ parser.add_argument(
+ "--out_dir", type=str, default=None, help="Directory path to put output files"
+ )
+ args = parser.parse_args()
+ print(args)
+ preprocess_librimix_audio(args)
\ No newline at end of file
diff --git a/ORIGINAL_README.md b/ORIGINAL_README.md
new file mode 100644
index 0000000000000000000000000000000000000000..2c0fb6c20818d396bcc6ac97eb42d8b14238bd6e
--- /dev/null
+++ b/ORIGINAL_README.md
@@ -0,0 +1,95 @@
+
+
+TIGER: Time-frequency Interleaved Gain Extraction and Reconstruction for Efficient Speech Separation
+
+ Mohan Xu* , Kai Li* , Guo Chen, Xiaolin Hu Tsinghua University, Beijing, China * Equal contribution📜 ICLR 2025 | 🎶 Demo | 🤗 Dataset
+
+
+
+
+
+
+> TIGER is a lightweight model for speech separation which effectively extracts key acoustic features through frequency band-split, multi-scale and full-frequency-frame modeling.
+
+## 💥 News
+
+- **[2025-01-23]** We release the code and pre-trained model of TIGER! 🚀
+- **[2025-01-23]** We release the TIGER model and the EchoSet dataset! 🚀
+
+## 📜 Abstract
+
+In this paper, we propose a speech separation model with significantly reduced parameter size and computational cost: Time-Frequency Interleaved Gain Extraction and Reconstruction Network (TIGER). TIGER leverages prior knowledge to divide frequency bands and applies compression on frequency information. We employ a multi-scale selective attention (MSA) module to extract contextual features, while introducing a full-frequency-frame attention (F^3A) module to capture both temporal and frequency contextual information. Additionally, to more realistically evaluate the performance of speech separation models in complex acoustic environments, we introduce a novel dataset called EchoSet. This dataset includes noise and more realistic reverberation (e.g., considering object occlusions and material properties), with speech from two speakers overlapping at random proportions. Experimental results demonstrated that TIGER significantly outperformed state-of-the-art (SOTA) model TF-GridNet on the EchoSet dataset in both inference speed and separation quality, while reducing the number of parameters by 94.3% and the MACs by 95.3%. These results indicate that by utilizing frequency band-split and interleaved modeling structures, TIGER achieves a substantial reduction in parameters and computational costs while maintaining high performance. Notably, TIGER is the first speech separation model with fewer than 1 million parameters that achieves performance close to the SOTA model.
+
+## TIGER
+
+Overall pipeline of the model architecture of TIGER and its modules.
+
+
+
+## Results
+
+Performance comparisons of TIGER and other existing separation models on ***Libri2Mix, LRS2-2Mix, and EchoSet***. Bold indicates optimal performance, and italics indicate suboptimal performance.
+
+
+
+Efficiency comparisons of TIGER and other models.
+
+
+
+Comparison of performance and efficiency of cinematic sound separation models on DnR. '*' means the result comes from the original paper of DnR.
+
+
+
+## 📦 Installation
+
+```bash
+git clone https://github.com/JusperLee/TIGER.git
+cd TIGER
+pip install -r requirements.txt
+```
+
+## 🚀 Quick Start
+
+### Test with Pre-trained Model
+
+```bash
+# Test using speech
+python inference_speech.py --audio_path test/mix.wav
+
+# Test using DnR
+python inference_dnr.py --audio_path test/test_mixture_466.wav
+```
+
+### Train with EchoSet
+
+```bash
+python audio_train.py --conf_dir configs/tiger.yml
+```
+
+### Evaluate with EchoSet
+
+```bash
+python audio_test.py --conf_dir configs/tiger.yml
+```
+
+## 📖 Citation
+
+```bibtex
+@article{xu2024tiger,
+ title={TIGER: Time-frequency Interleaved Gain Extraction and Reconstruction for Efficient Speech Separation},
+ author={Xu, Mohan and Li, Kai and Chen, Guo and Hu, Xiaolin},
+ journal={arXiv preprint arXiv:2410.01469},
+ year={2024}
+}
+```
+
+## 📧 Contact
+
+If you have any questions, please feel free to contact us via `tsinghua.kaili@gmail.com`.
diff --git a/assets/TIGER.png b/assets/TIGER.png
new file mode 100644
index 0000000000000000000000000000000000000000..21458a248c047762f08f03caf28fd5fe1c5bfac1
--- /dev/null
+++ b/assets/TIGER.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce1139d21f7b363f2210dc91f4ff6fce1aaebb4cb7432462096d53e76dfd806f
+size 1502140
diff --git a/assets/dnr-demo/sample1/dialog.mp4 b/assets/dnr-demo/sample1/dialog.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..0589d35e1752dc4a91fcec76e3fd5fef3a5b1eeb
--- /dev/null
+++ b/assets/dnr-demo/sample1/dialog.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f01fa3aeddcca598793810ead2ecf0140f82ebb43e96800578ec52af71f6e35
+size 2810085
diff --git a/assets/dnr-demo/sample1/effect.mp4 b/assets/dnr-demo/sample1/effect.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..c4988833ab761672be742e3ea42ab3ac30f191bb
--- /dev/null
+++ b/assets/dnr-demo/sample1/effect.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7cbf696ebb01eda225c9ebf77856677afe3ccd983b0e2bc45547ee96cd074fb
+size 2939960
diff --git a/assets/dnr-demo/sample1/mixture.mp4 b/assets/dnr-demo/sample1/mixture.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..3f0caaadda63bb806802a03a2e184b772a323a80
--- /dev/null
+++ b/assets/dnr-demo/sample1/mixture.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a66b7b61dc5a1e4639321b6255a23cc0e504a1492b9396c7e287d28d39627f35
+size 2918024
diff --git a/assets/dnr-demo/sample1/music.mp4 b/assets/dnr-demo/sample1/music.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..2d4c20cbd4bf72a53b4e89db0cdd4c90133a05cf
--- /dev/null
+++ b/assets/dnr-demo/sample1/music.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:413d85706158662a0c8852612cb4536be9b62fb4717cd4cbbc7c6a868194a0ca
+size 2884375
diff --git a/assets/dnr-demo/sample2/dialog.mp4 b/assets/dnr-demo/sample2/dialog.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..c445f5774d8e01213747928d915061653c81bc93
--- /dev/null
+++ b/assets/dnr-demo/sample2/dialog.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2677b02f921915a8bafcd9149a297f54657cafbc3464bec833743a6c7ad7779c
+size 1718587
diff --git a/assets/dnr-demo/sample2/effect.mp4 b/assets/dnr-demo/sample2/effect.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..4c45bba6884c07487dad674c5ec3ccddf1c66eb1
--- /dev/null
+++ b/assets/dnr-demo/sample2/effect.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:784a90eadcb0cae93d83ce739d95c06163423d8742448b7ffd373676a2f51a01
+size 1752110
diff --git a/assets/dnr-demo/sample2/mixture.mp4 b/assets/dnr-demo/sample2/mixture.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..15f0c6a76fbf408dd1a687bd116e501f6936a928
--- /dev/null
+++ b/assets/dnr-demo/sample2/mixture.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ab05854c7d9f4bbee331c9e8655b4bdadb4723050332baba027a06b33391fe36
+size 1781103
diff --git a/assets/dnr-demo/sample2/music.mp4 b/assets/dnr-demo/sample2/music.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..320d3f3a7dfb017c9a3dc17e4bebabd03bf36dae
--- /dev/null
+++ b/assets/dnr-demo/sample2/music.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a0b92919ad64ee5e11b14b8a8ab58361fc05c36ba4fe0873d6616f0d9eeead39
+size 1778868
diff --git a/assets/dnr-demo/sample3/dialog.mp4 b/assets/dnr-demo/sample3/dialog.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..da900ea201281156cd09c2139a4d627c8d69f023
--- /dev/null
+++ b/assets/dnr-demo/sample3/dialog.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc1e9c2336a61630f7f48561d21fc51c5f68731cf2a247e640b2ca8a5a576da2
+size 3988060
diff --git a/assets/dnr-demo/sample3/effect.mp4 b/assets/dnr-demo/sample3/effect.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..8bcb8140cdd6a2fd1b178414508dca16d85dbbcf
--- /dev/null
+++ b/assets/dnr-demo/sample3/effect.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efaf9714d91b8caf53b7f7b4f3e39a811d6e20393fc48bcbc5d23930ec66742a
+size 4035872
diff --git a/assets/dnr-demo/sample3/mixture.mp4 b/assets/dnr-demo/sample3/mixture.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..6e318d706e0337a3a232a898b8e291237190387a
--- /dev/null
+++ b/assets/dnr-demo/sample3/mixture.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ffb8dee7b106e5c7d1da78608b1992b93bb764825bdc9c22c286f6c04b020fa5
+size 4053761
diff --git a/assets/dnr-demo/sample3/music.mp4 b/assets/dnr-demo/sample3/music.mp4
new file mode 100644
index 0000000000000000000000000000000000000000..8b347505bc68e442047bdb1a2fcdebd13f799338
--- /dev/null
+++ b/assets/dnr-demo/sample3/music.mp4
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90b24a2e5aaf1a4c880c83f68d25179ebb119905658a8f28e8cdc3d33cd24f10
+size 4062344
diff --git a/assets/dnr.png b/assets/dnr.png
new file mode 100644
index 0000000000000000000000000000000000000000..ac81432da61acc90ab2b687a9cd0cd788f0d1cf3
--- /dev/null
+++ b/assets/dnr.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fce2988567b8637f68dce9d25aa17d17b98fc9b6f0ce4f80ea78144b3eeda80
+size 194016
diff --git a/assets/efficiency.png b/assets/efficiency.png
new file mode 100644
index 0000000000000000000000000000000000000000..fec7fa6b39d910b84343e05ef7f33a323d3bc8e2
--- /dev/null
+++ b/assets/efficiency.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ccc5b69e2ccbb48b03f44fdbbd88b99366b124307f84b78538996a5ba91d171
+size 405116
diff --git a/assets/logo.png b/assets/logo.png
new file mode 100644
index 0000000000000000000000000000000000000000..d22794ec0ac8766f5e40b0615f3ce606d1421cb7
--- /dev/null
+++ b/assets/logo.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e93c060d66ca24d7dcd8f30fab9546e0eea7b09027d8af6b67f72c76f66a88f2
+size 486031
diff --git a/assets/result.png b/assets/result.png
new file mode 100644
index 0000000000000000000000000000000000000000..4047b5bdc4da3b2bc49fde1a1730426e3da4ffb9
--- /dev/null
+++ b/assets/result.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd2d93eec2d4f7c44151a231b830fc687fba3bfd1081a220ac9ab8460eeffcf3
+size 469267
diff --git a/assets/sample1/GroundTruth/mix.wav b/assets/sample1/GroundTruth/mix.wav
new file mode 100644
index 0000000000000000000000000000000000000000..eb063704154aff42d7a52bc909f5a04ac315ac3c
--- /dev/null
+++ b/assets/sample1/GroundTruth/mix.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96096869c977120503b2ef4b53429a4bce510b344d6a274c3371457be854f38e
+size 192044
diff --git a/assets/sample1/GroundTruth/s1.wav b/assets/sample1/GroundTruth/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..ce773338cc58908074a345f7bef11c4e6b49ef24
--- /dev/null
+++ b/assets/sample1/GroundTruth/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6128c53d4d57c61eae0b1d7030a155b91db64fb6ca5ec72c982c4f3b6ac3e071
+size 192044
diff --git a/assets/sample1/GroundTruth/s2.wav b/assets/sample1/GroundTruth/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..b4c142414daedfb0cbba44923768936d0e17b384
--- /dev/null
+++ b/assets/sample1/GroundTruth/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:342f9c15fc0b0445ff82b1026551b33fbafe3ee89c060990ed402ad4ad2a62ee
+size 192044
diff --git a/assets/sample1/TFGNet/s1.wav b/assets/sample1/TFGNet/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..bd5bccd81064a97e67fcef30dabe24302dbfbf9c
--- /dev/null
+++ b/assets/sample1/TFGNet/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8676fdb6412ed3123d025e145abdd86bbaf2956e55cbc3be5034857ee03b562f
+size 384080
diff --git a/assets/sample1/TFGNet/s2.wav b/assets/sample1/TFGNet/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..e94d0d8da0001c4a33b4c36503be5e5d223fd7ee
--- /dev/null
+++ b/assets/sample1/TFGNet/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac533785786cac0cacc9f5dde2a32c90c7b057da759be4837b44aa7ddf586b94
+size 384080
diff --git a/assets/sample1/TIGER/s1.wav b/assets/sample1/TIGER/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..745df78928513718302157293d6ef5ed69a00d6f
--- /dev/null
+++ b/assets/sample1/TIGER/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8ec094c7c7d9262f734b651735ae209eeaf70f1f7d0e18085ab2db330cdc53f
+size 384080
diff --git a/assets/sample1/TIGER/s2.wav b/assets/sample1/TIGER/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..02f697df0a89fd6d89332d3e74c46dbe2efd86b2
--- /dev/null
+++ b/assets/sample1/TIGER/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92837bf98498e4f7f23d5f918cf042f60f05c534eb1cc4fd27fadce170e38132
+size 384080
diff --git a/assets/sample1/spec/TFGNet_s1.png b/assets/sample1/spec/TFGNet_s1.png
new file mode 100644
index 0000000000000000000000000000000000000000..a0c3c1c621ce9519265a0058e999e9e6c390fd9a
--- /dev/null
+++ b/assets/sample1/spec/TFGNet_s1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f629a31eebf229ed4aca06fdd15a2da516e2095967cba8590194a54d8719e10b
+size 774751
diff --git a/assets/sample1/spec/TFGNet_s2.png b/assets/sample1/spec/TFGNet_s2.png
new file mode 100644
index 0000000000000000000000000000000000000000..ca3f6fc9f5fd0c88a102153577f620aad7e7dc7e
--- /dev/null
+++ b/assets/sample1/spec/TFGNet_s2.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e9ab62f73ad6721f1b9edd570f364158f95e6ee0e446a4aebd9325aaea284d2
+size 767865
diff --git a/assets/sample1/spec/TIGER_s1.png b/assets/sample1/spec/TIGER_s1.png
new file mode 100644
index 0000000000000000000000000000000000000000..680dc9ee965f198b0eabb0c6cc6871c3e8279714
--- /dev/null
+++ b/assets/sample1/spec/TIGER_s1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34eac91edfbf96ce850c6474e60a5add51151579bc423d63d27593b07df64c9e
+size 771902
diff --git a/assets/sample1/spec/TIGER_s2.png b/assets/sample1/spec/TIGER_s2.png
new file mode 100644
index 0000000000000000000000000000000000000000..ed03dfa6de727dd286c2b6479dfe09cb03672a55
--- /dev/null
+++ b/assets/sample1/spec/TIGER_s2.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:019899feacd4df56a66f4052352e5e4252ac1e20fbcab106099f3696ef5b23ac
+size 792674
diff --git a/assets/sample1/spec/ground_truth_s1.png b/assets/sample1/spec/ground_truth_s1.png
new file mode 100644
index 0000000000000000000000000000000000000000..30deb410116f9ea5571dea52ff3935bc8dcfee2f
--- /dev/null
+++ b/assets/sample1/spec/ground_truth_s1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87fab6fb1e5cb7415e106b4aed57a108e8160cee893cd9c0ece3ecab2ca5bded
+size 619171
diff --git a/assets/sample1/spec/ground_truth_s2.png b/assets/sample1/spec/ground_truth_s2.png
new file mode 100644
index 0000000000000000000000000000000000000000..07444ce0fd1df8f4a671729d58b7891536170a56
--- /dev/null
+++ b/assets/sample1/spec/ground_truth_s2.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c160511e75a0c0e3998e634490f71db2f6a403b408ec2af8312b614aba6ad12
+size 685215
diff --git a/assets/sample2/GroundTruth/mix.wav b/assets/sample2/GroundTruth/mix.wav
new file mode 100644
index 0000000000000000000000000000000000000000..bcd6f7213bde7e9215008b6d8b4cae1412c585fb
--- /dev/null
+++ b/assets/sample2/GroundTruth/mix.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26f24e81d600dab059662afec2d0ef6197875ce8e233e554c423d435527560a8
+size 192044
diff --git a/assets/sample2/GroundTruth/s1.wav b/assets/sample2/GroundTruth/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..446b30d3e874c0e28b383489cef785e3cfec7a83
--- /dev/null
+++ b/assets/sample2/GroundTruth/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bbfcb21a93c1c1896800a5d057f8e9c497c27c92a36233c90920c6a2e1831ee2
+size 192044
diff --git a/assets/sample2/GroundTruth/s2.wav b/assets/sample2/GroundTruth/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..28b023dbeeb3d179bca5e732d935bb07f981ffb0
--- /dev/null
+++ b/assets/sample2/GroundTruth/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad6579a1ee21ac0d9fb416f9d79a10d65f52f7f5a51759355d259141e966e33e
+size 192044
diff --git a/assets/sample2/TFGNet/s1.wav b/assets/sample2/TFGNet/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..b7bfbbca5fa3745fb072779e8a59a7a2202ed601
--- /dev/null
+++ b/assets/sample2/TFGNet/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:401a5efe1d3feda8a1c34a4f887bd353114166561d6aa81e28097fc68933004b
+size 384080
diff --git a/assets/sample2/TFGNet/s2.wav b/assets/sample2/TFGNet/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..3d8f0796cae5ddc7c3efe84f0bc0e44cc6b10d96
--- /dev/null
+++ b/assets/sample2/TFGNet/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59c8a509d19b47254ff37a7f15b0194a49bd2b154ad2725ac7d07efe9ff174bd
+size 384080
diff --git a/assets/sample2/TIGER/s1.wav b/assets/sample2/TIGER/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..a8c61eff87eea39b0b0b9a668838e618ba6237c4
--- /dev/null
+++ b/assets/sample2/TIGER/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e804f6d6c6d63f3756427e8c6832a9a28a0bbbcb808e2f7faf14acdc5adc5d8
+size 384080
diff --git a/assets/sample2/TIGER/s2.wav b/assets/sample2/TIGER/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..f4ec2adebc434bd870c4116f2014035be4da9194
--- /dev/null
+++ b/assets/sample2/TIGER/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45f92e18ce5371fcc3d07a7010e406efad0092090a3003e9252f901e12e353ba
+size 384080
diff --git a/assets/sample2/spec/TFGNet_s1.png b/assets/sample2/spec/TFGNet_s1.png
new file mode 100644
index 0000000000000000000000000000000000000000..1e8fc2d5355eaaf0288c1606d38eecc5361131c1
--- /dev/null
+++ b/assets/sample2/spec/TFGNet_s1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b95c0a00dcd44f794747fc86d9f916031994e9f21728a31ecac97c2e941fc522
+size 434257
diff --git a/assets/sample2/spec/TFGNet_s2.png b/assets/sample2/spec/TFGNet_s2.png
new file mode 100644
index 0000000000000000000000000000000000000000..81d5ca8904760add9bac16fcdd8de0b91638f9a6
--- /dev/null
+++ b/assets/sample2/spec/TFGNet_s2.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bdd22e56556506cec95009049243cdf1286acd2dfd27e073b8ee75319d9c2bf
+size 727285
diff --git a/assets/sample2/spec/TIGER_s1.png b/assets/sample2/spec/TIGER_s1.png
new file mode 100644
index 0000000000000000000000000000000000000000..b1d6bd11d93e25b4f929d3394826ad641a374910
--- /dev/null
+++ b/assets/sample2/spec/TIGER_s1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f09dba6957f58fe02dbb4557d036d4edae00d101b1e96edfc76de27e44f4ba5a
+size 565372
diff --git a/assets/sample2/spec/TIGER_s2.png b/assets/sample2/spec/TIGER_s2.png
new file mode 100644
index 0000000000000000000000000000000000000000..188f6d3f09329e87f0ca784e4f970a56f572ac61
--- /dev/null
+++ b/assets/sample2/spec/TIGER_s2.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba061680e3a36ceebe158c40f51db15cd2c322de33c41adaaf8c0358b624dd08
+size 582429
diff --git a/assets/sample2/spec/ground_truth_s1.png b/assets/sample2/spec/ground_truth_s1.png
new file mode 100644
index 0000000000000000000000000000000000000000..41ed0fe6dee762c53314c65926fa71907cc4a41e
--- /dev/null
+++ b/assets/sample2/spec/ground_truth_s1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11216b3f40d513c5297b54da5dc989f96305efac6c915f893c3524d6c0734622
+size 534150
diff --git a/assets/sample2/spec/ground_truth_s2.png b/assets/sample2/spec/ground_truth_s2.png
new file mode 100644
index 0000000000000000000000000000000000000000..931b9de96e5b23528e68ac32caa0a71ad0b5558c
--- /dev/null
+++ b/assets/sample2/spec/ground_truth_s2.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:def1f97d15ef8c61d8e8a2e3144e861dcacd6c945e19375be3a2ca3749f15653
+size 584546
diff --git a/assets/sample3/GroundTruth/mix.wav b/assets/sample3/GroundTruth/mix.wav
new file mode 100644
index 0000000000000000000000000000000000000000..b01655e1b58ec92473bd96a096d14eee872376a9
--- /dev/null
+++ b/assets/sample3/GroundTruth/mix.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92be7d03f4ca97e3ee88272490838fe77d64207fb0d326cd70391da061498a8a
+size 192044
diff --git a/assets/sample3/GroundTruth/s1.wav b/assets/sample3/GroundTruth/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..59dbff1b2a3b4ef4ce57bdf4763bfb152c73c5d3
--- /dev/null
+++ b/assets/sample3/GroundTruth/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b4071c7f390ce9e9a3ccfcf58afaec1e5b237aa150c3c87cf46a3c5a361e2f7
+size 192044
diff --git a/assets/sample3/GroundTruth/s2.wav b/assets/sample3/GroundTruth/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..82e669cafb571aea92459f3f830e6dfc3bda19dc
--- /dev/null
+++ b/assets/sample3/GroundTruth/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48354c78aa7af0099c3bad5df18b587bc4ee897bd1ea00ac6f3db77c69055711
+size 192044
diff --git a/assets/sample3/TFGNet/s1.wav b/assets/sample3/TFGNet/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..bbbd1692e80eb4300fa1868d819550d58c1b8de6
--- /dev/null
+++ b/assets/sample3/TFGNet/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a669856a8a7ca1c4439a2cf1b88cfa5830a031449ef7950653120fd6ce0eeeca
+size 384080
diff --git a/assets/sample3/TFGNet/s2.wav b/assets/sample3/TFGNet/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..085bfc4614e658db90a01a70778b8c783494bfa1
--- /dev/null
+++ b/assets/sample3/TFGNet/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e409459991ef9b50a36267b6ea0cefa47263b25a3c6a6bb7d04a99681a06abaf
+size 384080
diff --git a/assets/sample3/TIGER/s1.wav b/assets/sample3/TIGER/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..8f6bd186b0a1b1b073aeaa7e952d3435bca97846
--- /dev/null
+++ b/assets/sample3/TIGER/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6784a1d0a0b2e7a1d69a523ba824b5493465301727fe79e0a99d0fc6f52d0ef
+size 384080
diff --git a/assets/sample3/TIGER/s2.wav b/assets/sample3/TIGER/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..cf2510e2f8487e8f409469bad1d537931c8ad084
--- /dev/null
+++ b/assets/sample3/TIGER/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:188edc3e4fac51b7494e5d7e06f103c250d185876e3cc71af908f613829de97a
+size 384080
diff --git a/assets/sample3/spec/TFGNet_s1.png b/assets/sample3/spec/TFGNet_s1.png
new file mode 100644
index 0000000000000000000000000000000000000000..9b14b798b1dc01ad60a000bd85bd602a9eed88e1
--- /dev/null
+++ b/assets/sample3/spec/TFGNet_s1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6cc1f43f5f52fa9faf2e2485a6196b5959f9679f996b3522d310e63a7c3e55f7
+size 631623
diff --git a/assets/sample3/spec/TFGNet_s2.png b/assets/sample3/spec/TFGNet_s2.png
new file mode 100644
index 0000000000000000000000000000000000000000..a4fa8cc3f3cb125bbbef5c07658abbcfa3d47c86
--- /dev/null
+++ b/assets/sample3/spec/TFGNet_s2.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83d0c746152c5808b33e18d04de0561cd503552b66299676e534eff24f0d0eea
+size 541215
diff --git a/assets/sample3/spec/TIGER_s1.png b/assets/sample3/spec/TIGER_s1.png
new file mode 100644
index 0000000000000000000000000000000000000000..a2babb155ef321359b40a667c2eee81f6b2af90e
--- /dev/null
+++ b/assets/sample3/spec/TIGER_s1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3913b89117845e2b1d0a794c4203eee01558bc6a1702cee4ca01c7d8e6ba104
+size 550682
diff --git a/assets/sample3/spec/TIGER_s2.png b/assets/sample3/spec/TIGER_s2.png
new file mode 100644
index 0000000000000000000000000000000000000000..408cab28c1f42232211437d243c229b87ef4a6f2
--- /dev/null
+++ b/assets/sample3/spec/TIGER_s2.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d1e4d09989f758371536f84c2576ac72a31914f4446fd589b37040a1302b40e
+size 660406
diff --git a/assets/sample3/spec/ground_truth_s1.png b/assets/sample3/spec/ground_truth_s1.png
new file mode 100644
index 0000000000000000000000000000000000000000..6c2e9341adb25258ecda566fa3045c843e5f8807
--- /dev/null
+++ b/assets/sample3/spec/ground_truth_s1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f170d33ffc66d730693cc071273d9aef3e58be850f9bd8b479cc93ced99d0cb
+size 548248
diff --git a/assets/sample3/spec/ground_truth_s2.png b/assets/sample3/spec/ground_truth_s2.png
new file mode 100644
index 0000000000000000000000000000000000000000..dc8e323920bd84fac4269dbc8f41598ea1c4f2b7
--- /dev/null
+++ b/assets/sample3/spec/ground_truth_s2.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80942ea5e1742fb77f653625540f36e08821c333631d5b9a6f1d37d570b68a8b
+size 686893
diff --git a/audio_test.py b/audio_test.py
new file mode 100644
index 0000000000000000000000000000000000000000..2602981634f91d4f03746f21927ae58407fdbdd4
--- /dev/null
+++ b/audio_test.py
@@ -0,0 +1,116 @@
+import os
+import random
+from typing import Union
+import soundfile as sf
+import torch
+import yaml
+import json
+import argparse
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from pprint import pprint
+from scipy.io import wavfile
+import warnings
+import torchaudio
+warnings.filterwarnings("ignore")
+import look2hear.models
+import look2hear.datas
+from look2hear.metrics import MetricsTracker
+from look2hear.utils import tensors_to_device, RichProgressBarTheme, MyMetricsTextColumn, BatchesProcessedColumn
+
+from rich.progress import (
+ BarColumn,
+ Progress,
+ TextColumn,
+ TimeRemainingColumn,
+ TransferSpeedColumn,
+)
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--conf_dir",
+ default="local/mixit_conf.yml",
+ help="Full path to save best validation model")
+
+
+compute_metrics = ["si_sdr", "sdr"]
+os.environ['CUDA_VISIBLE_DEVICES'] = "8"
+
+def main(config):
+ metricscolumn = MyMetricsTextColumn(style=RichProgressBarTheme.metrics)
+ progress = Progress(
+ TextColumn("[bold blue]Testing", justify="right"),
+ BarColumn(bar_width=None),
+ "•",
+ BatchesProcessedColumn(style=RichProgressBarTheme.batch_progress),
+ "•",
+ TransferSpeedColumn(),
+ "•",
+ TimeRemainingColumn(),
+ "•",
+ metricscolumn
+ )
+ # import pdb; pdb.set_trace()
+ config["train_conf"]["main_args"]["exp_dir"] = os.path.join(
+ os.getcwd(), "Experiments", "checkpoint", config["train_conf"]["exp"]["exp_name"]
+ )
+ model_path = os.path.join(config["train_conf"]["main_args"]["exp_dir"], "best_model.pth")
+ # import pdb; pdb.set_trace()
+ # conf["train_conf"]["masknet"].update({"n_src": 2})
+ model = getattr(look2hear.models, config["train_conf"]["audionet"]["audionet_name"]).from_pretrain(
+ model_path,
+ sample_rate=config["train_conf"]["datamodule"]["data_config"]["sample_rate"],
+ **config["train_conf"]["audionet"]["audionet_config"],
+ )
+ if config["train_conf"]["training"]["gpus"]:
+ device = "cuda"
+ model.to(device)
+ model_device = next(model.parameters()).device
+ datamodule: object = getattr(look2hear.datas, config["train_conf"]["datamodule"]["data_name"])(
+ **config["train_conf"]["datamodule"]["data_config"]
+ )
+ datamodule.setup()
+ _, _ , test_set = datamodule.make_sets
+
+ # Randomly choose the indexes of sentences to save.
+ ex_save_dir = os.path.join(config["train_conf"]["main_args"]["exp_dir"], "results/")
+ os.makedirs(ex_save_dir, exist_ok=True)
+ metrics = MetricsTracker(
+ save_file=os.path.join(ex_save_dir, "metrics.csv"))
+ torch.no_grad().__enter__()
+ with progress:
+ for idx in progress.track(range(len(test_set))):
+ if idx == 825:
+ # Forward the network on the mixture.
+ mix, sources, key = tensors_to_device(test_set[idx],
+ device=model_device)
+ est_sources = model(mix[None])
+ mix_np = mix
+ sources_np = sources
+ est_sources_np = est_sources.squeeze(0)
+ # metrics(mix=mix_np,
+ # clean=sources_np,
+ # estimate=est_sources_np,
+ # key=key)
+ save_dir = os.path.join("./result/TIGER", "idx{}".format(idx))
+ # est_sources_np = normalize_tensor_wav(est_sources_np)
+ for i in range(est_sources_np.shape[0]):
+ os.makedirs(os.path.join(save_dir, "s{}/".format(i + 1)), exist_ok=True)
+ # torchaudio.save(os.path.join(save_dir, "s{}/".format(i + 1)) + key, est_sources_np[i].unsqueeze(0).cpu(), 16000)
+ torchaudio.save(os.path.join(save_dir, "s{}/".format(i + 1)) + key.split("/")[-1], est_sources_np[i].unsqueeze(0).cpu(), 16000)
+ # if idx % 50 == 0:
+ # metricscolumn.update(metrics.update())
+ metrics.final()
+
+
+if __name__ == "__main__":
+ args = parser.parse_args()
+ arg_dic = dict(vars(args))
+
+ # Load training config
+ with open(args.conf_dir, "rb") as f:
+ train_conf = yaml.safe_load(f)
+ arg_dic["train_conf"] = train_conf
+ # print(arg_dic)
+ main(arg_dic)
diff --git a/audio_train.py b/audio_train.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e53b9d714e459598edd1730d6add4421f6eb47b
--- /dev/null
+++ b/audio_train.py
@@ -0,0 +1,202 @@
+import os
+import sys
+import torch
+from torch import Tensor
+import argparse
+import json
+import look2hear.datas
+import look2hear.models
+import look2hear.system
+import look2hear.losses
+import look2hear.metrics
+import look2hear.utils
+from look2hear.system import make_optimizer
+from dataclasses import dataclass
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, RichProgressBar
+from pytorch_lightning.callbacks.progress.rich_progress import *
+from rich.console import Console
+from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.loggers.wandb import WandbLogger
+from pytorch_lightning.strategies.ddp import DDPStrategy
+from rich import print, reconfigure
+from collections.abc import MutableMapping
+from look2hear.utils import print_only, MyRichProgressBar, RichProgressBarTheme
+
+import warnings
+
+warnings.filterwarnings("ignore")
+
+import wandb
+wandb.login()
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+ "--conf_dir",
+ default="local/conf.yml",
+ help="Full path to save best validation model",
+)
+
+def main(config):
+ print_only(
+ "Instantiating datamodule <{}>".format(config["datamodule"]["data_name"])
+ )
+ datamodule: object = getattr(look2hear.datas, config["datamodule"]["data_name"])(
+ **config["datamodule"]["data_config"]
+ )
+ datamodule.setup()
+
+ train_loader, val_loader, test_loader = datamodule.make_loader
+
+ # Define model and optimizer
+ print_only(
+ "Instantiating AudioNet <{}>".format(config["audionet"]["audionet_name"])
+ )
+ model = getattr(look2hear.models, config["audionet"]["audionet_name"])(
+ sample_rate=config["datamodule"]["data_config"]["sample_rate"],
+ **config["audionet"]["audionet_config"],
+ )
+ # import pdb; pdb.set_trace()
+ print_only("Instantiating Optimizer <{}>".format(config["optimizer"]["optim_name"]))
+ optimizer = make_optimizer(model.parameters(), **config["optimizer"])
+
+ # Define scheduler
+ scheduler = None
+ if config["scheduler"]["sche_name"]:
+ print_only(
+ "Instantiating Scheduler <{}>".format(config["scheduler"]["sche_name"])
+ )
+ if config["scheduler"]["sche_name"] != "DPTNetScheduler":
+ scheduler = getattr(torch.optim.lr_scheduler, config["scheduler"]["sche_name"])(
+ optimizer=optimizer, **config["scheduler"]["sche_config"]
+ )
+ else:
+ scheduler = {
+ "scheduler": getattr(look2hear.system.schedulers, config["scheduler"]["sche_name"])(
+ optimizer, len(train_loader) // config["datamodule"]["data_config"]["batch_size"], 64
+ ),
+ "interval": "step",
+ }
+
+ # Just after instantiating, save the args. Easy loading in the future.
+ config["main_args"]["exp_dir"] = os.path.join(
+ os.getcwd(), "Experiments", "checkpoint", config["exp"]["exp_name"]
+ )
+ exp_dir = config["main_args"]["exp_dir"]
+ os.makedirs(exp_dir, exist_ok=True)
+ conf_path = os.path.join(exp_dir, "conf.yml")
+ with open(conf_path, "w") as outfile:
+ yaml.safe_dump(config, outfile)
+
+ # Define Loss function.
+ print_only(
+ "Instantiating Loss, Train <{}>, Val <{}>".format(
+ config["loss"]["train"]["sdr_type"], config["loss"]["val"]["sdr_type"]
+ )
+ )
+ loss_func = {
+ "train": getattr(look2hear.losses, config["loss"]["train"]["loss_func"])(
+ getattr(look2hear.losses, config["loss"]["train"]["sdr_type"]),
+ **config["loss"]["train"]["config"],
+ ),
+ "val": getattr(look2hear.losses, config["loss"]["val"]["loss_func"])(
+ getattr(look2hear.losses, config["loss"]["val"]["sdr_type"]),
+ **config["loss"]["val"]["config"],
+ ),
+ }
+
+ print_only("Instantiating System <{}>".format(config["training"]["system"]))
+ system = getattr(look2hear.system, config["training"]["system"])(
+ audio_model=model,
+ loss_func=loss_func,
+ optimizer=optimizer,
+ train_loader=train_loader,
+ val_loader=val_loader,
+ test_loader=test_loader,
+ scheduler=scheduler,
+ config=config,
+ )
+
+ # Define callbacks
+ print_only("Instantiating ModelCheckpoint")
+ callbacks = []
+ checkpoint_dir = os.path.join(exp_dir)
+ checkpoint = ModelCheckpoint(
+ checkpoint_dir,
+ filename="{epoch}",
+ monitor="val_loss/dataloader_idx_0",
+ mode="min",
+ save_top_k=5,
+ verbose=True,
+ save_last=True,
+ )
+ callbacks.append(checkpoint)
+
+ if config["training"]["early_stop"]:
+ print_only("Instantiating EarlyStopping")
+ callbacks.append(EarlyStopping(**config["training"]["early_stop"]))
+ callbacks.append(MyRichProgressBar(theme=RichProgressBarTheme()))
+
+ # Don't ask GPU if they are not available.
+ gpus = config["training"]["gpus"] if torch.cuda.is_available() else None
+ distributed_backend = "cuda" if torch.cuda.is_available() else None
+
+ # default logger used by trainer
+ logger_dir = os.path.join(os.getcwd(), "Experiments", "tensorboard_logs")
+ os.makedirs(os.path.join(logger_dir, config["exp"]["exp_name"]), exist_ok=True)
+ # comet_logger = TensorBoardLogger(logger_dir, name=config["exp"]["exp_name"])
+ comet_logger = WandbLogger(
+ name=config["exp"]["exp_name"],
+ save_dir=os.path.join(logger_dir, config["exp"]["exp_name"]),
+ project="Real-work-dataset",
+ # offline=True
+ )
+
+ trainer = pl.Trainer(
+ max_epochs=config["training"]["epochs"],
+ callbacks=callbacks,
+ default_root_dir=exp_dir,
+ devices=gpus,
+ accelerator=distributed_backend,
+ strategy=DDPStrategy(find_unused_parameters=True),
+ limit_train_batches=1.0, # Useful for fast experiment
+ gradient_clip_val=5.0,
+ logger=comet_logger,
+ sync_batchnorm=True,
+ # precision="bf16-mixed",
+ # num_sanity_val_steps=0,
+ # sync_batchnorm=True,
+ # fast_dev_run=True,
+ )
+ trainer.fit(system)
+ print_only("Finished Training")
+ best_k = {k: v.item() for k, v in checkpoint.best_k_models.items()}
+ with open(os.path.join(exp_dir, "best_k_models.json"), "w") as f:
+ json.dump(best_k, f, indent=0)
+
+ state_dict = torch.load(checkpoint.best_model_path)
+ system.load_state_dict(state_dict=state_dict["state_dict"])
+ system.cpu()
+
+ to_save = system.audio_model.serialize()
+ torch.save(to_save, os.path.join(exp_dir, "best_model.pth"))
+
+
+if __name__ == "__main__":
+ import yaml
+ from pprint import pprint
+ from look2hear.utils.parser_utils import (
+ prepare_parser_from_dict,
+ parse_args_as_dict,
+ )
+
+ args = parser.parse_args()
+ with open(args.conf_dir) as f:
+ def_conf = yaml.safe_load(f)
+ parser = prepare_parser_from_dict(def_conf, parser=parser)
+
+ arg_dic, plain_args = parse_args_as_dict(parser, return_plain_args=True)
+ # pprint(arg_dic)
+ main(arg_dic)
diff --git a/configs/tiger-large.yml b/configs/tiger-large.yml
new file mode 100644
index 0000000000000000000000000000000000000000..61be3a25962433e4333e307933632c589e6e4d5b
--- /dev/null
+++ b/configs/tiger-large.yml
@@ -0,0 +1,71 @@
+# Network config
+audionet:
+ audionet_name: TIGER
+ audionet_config:
+ out_channels: 128
+ in_channels: 256
+ num_blocks: 8
+ upsampling_depth: 5
+ win: 640
+ stride: 160
+ num_sources: 2
+
+# Loss config
+loss:
+ train:
+ loss_func: PITLossWrapper
+ sdr_type: pairwise_neg_snr
+ config:
+ pit_from: pw_mtx
+ threshold_byloss: false
+ val:
+ loss_func: PITLossWrapper
+ sdr_type: pairwise_neg_sisdr
+ config:
+ pit_from: pw_mtx
+ threshold_byloss: false
+
+# Training config
+training:
+ system: AudioLightningModule
+ gpus: [2,3,4,5,6,7,8,9]
+ parallel: ddp
+ epochs: 500
+ early_stop:
+ monitor: val_loss/dataloader_idx_0
+ mode: min
+ patience: 20
+ verbose: true
+ SpeedAug: False
+
+# Optim config
+optimizer:
+ optim_name: adam
+ lr: 0.001
+ weight_decay: 0
+
+# Sche config
+scheduler:
+ sche_name: ReduceLROnPlateau
+ sche_config:
+ patience: 10
+ factor: 0.5
+
+# Data config
+datamodule:
+ data_name: EchoSetDataModule
+ data_config:
+ train_dir: DataPreProcess/EchoSet/train
+ valid_dir: DataPreProcess/EchoSet/val
+ test_dir: DataPreProcess/EchoSet/test
+ n_src: 2
+ sample_rate: 16000
+ segment: 3.0
+ normalize_audio: false
+ batch_size: 1
+ num_workers: 8
+ pin_memory: true
+ persistent_workers: false
+
+exp:
+ exp_name: TIGER-EchoSet
\ No newline at end of file
diff --git a/configs/tiger-small.yml b/configs/tiger-small.yml
new file mode 100644
index 0000000000000000000000000000000000000000..7b9e6c22cbe6bfb6a9faa206b34f1e47e668dc62
--- /dev/null
+++ b/configs/tiger-small.yml
@@ -0,0 +1,71 @@
+# Network config
+audionet:
+ audionet_name: TIGER
+ audionet_config:
+ out_channels: 128
+ in_channels: 256
+ num_blocks: 4
+ upsampling_depth: 5
+ win: 640
+ stride: 160
+ num_sources: 2
+
+# Loss config
+loss:
+ train:
+ loss_func: PITLossWrapper
+ sdr_type: pairwise_neg_snr
+ config:
+ pit_from: pw_mtx
+ threshold_byloss: false
+ val:
+ loss_func: PITLossWrapper
+ sdr_type: pairwise_neg_sisdr
+ config:
+ pit_from: pw_mtx
+ threshold_byloss: false
+
+# Training config
+training:
+ system: AudioLightningModule
+ gpus: [2,3,4,5,6,7,8,9]
+ parallel: ddp
+ epochs: 500
+ early_stop:
+ monitor: val_loss/dataloader_idx_0
+ mode: min
+ patience: 20
+ verbose: true
+ SpeedAug: False
+
+# Optim config
+optimizer:
+ optim_name: adam
+ lr: 0.001
+ weight_decay: 0
+
+# Sche config
+scheduler:
+ sche_name: ReduceLROnPlateau
+ sche_config:
+ patience: 10
+ factor: 0.5
+
+# Data config
+datamodule:
+ data_name: EchoSetDataModule
+ data_config:
+ train_dir: DataPreProcess/EchoSet/train
+ valid_dir: DataPreProcess/EchoSet/val
+ test_dir: DataPreProcess/EchoSet/test
+ n_src: 2
+ sample_rate: 16000
+ segment: 3.0
+ normalize_audio: false
+ batch_size: 1
+ num_workers: 8
+ pin_memory: true
+ persistent_workers: false
+
+exp:
+ exp_name: TIGER-EchoSet
\ No newline at end of file
diff --git a/evaluated_mac_params.py b/evaluated_mac_params.py
new file mode 100644
index 0000000000000000000000000000000000000000..3972bd6c0c41ab6b6316449359dd560e9f3c1d76
--- /dev/null
+++ b/evaluated_mac_params.py
@@ -0,0 +1,53 @@
+import os
+import argparse
+import json
+import time
+import torch
+import torch.nn as nn
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from look2hear.utils.parser_utils import prepare_parser_from_dict, parse_args_as_dict
+import look2hear.models
+import yaml
+from ptflops import get_model_complexity_info
+from rich import print
+
+def check_parameters(net):
+ """
+ Returns module parameters. Mb
+ """
+ parameters = sum(param.numel() for param in net.parameters())
+ return parameters / 10 ** 6
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+ "--exp_dir", default="exp/tmp", help="Full path to save best validation model"
+)
+
+with open("configs/tiger.yml") as f:
+ def_conf = yaml.safe_load(f)
+parser = prepare_parser_from_dict(def_conf, parser=parser)
+
+arg_dic, plain_args = parse_args_as_dict(parser, return_plain_args=True)
+audiomodel = getattr(look2hear.models, arg_dic["audionet"]["audionet_name"])(
+ sample_rate=arg_dic["datamodule"]["data_config"]["sample_rate"],
+ **arg_dic["audionet"]["audionet_config"]
+)
+# 配置GPU为mps
+device = torch.device("mps")
+a = torch.randn(1, 1, 16000).to(device)
+total_macs = 0
+total_params = 0
+model = audiomodel.to(device)
+with torch.no_grad():
+ macs, params = get_model_complexity_info(
+ model, (16000,), as_strings=False, print_per_layer_stat=True, verbose=False
+ )
+print(model(a).shape)
+total_macs += macs
+total_params += params
+print("MACs: ", total_macs / 10.0 ** 9)
+print("Params: ", total_params / 10.0 ** 6)
\ No newline at end of file
diff --git a/index.html b/index.html
new file mode 100644
index 0000000000000000000000000000000000000000..7f4c035f90a5d5e99585fadf759e77f2bf8a7542
--- /dev/null
+++ b/index.html
@@ -0,0 +1,791 @@
+
+
+
+
+
+
+
+ TIGER: Time-frequency Interleaved Gain Extraction and Reconstruction for Efficient Speech Separation
+
+
+
+
+
+
+ TIGER: Time-frequency Interleaved Gain Extraction and Reconstruction for Efficient
+ Speech Separation
+ ━ ICLR 2025 ━
+
+
+
+
+
+ Mohan Xu*
+
+
+
+
+ Kai Li*
+
+
+
+ Guo Chen
+
+
+
+
+ Xiaolin Hu
+
+
+
+
+
+
+
+
+
+ TIGER is a lightweight model
+ for speech separation which
+ effectively extracts key acoustic features through frequency band-split, multi-scale
+ and full-frequency-frame modeling.
+
+
+
+ Abstract
+
+
+
+ In recent years, much speech separation research has focused primarily on improving model performance.
+ However, for low-latency speech processing systems, computational efficiency is equally critical.
+ In this paper, we propose a speech separation model with significantly reduced parameter size and
+ computational cost:
+ Time-Frequency Interleaved Gain Extraction and Reconstruction Network (TIGER) .
+ TIGER leverages prior knowledge to divide frequency
+ bands and applies compression on frequency information.
+ We employ a multi-scale selective attention (MSA) module to extract contextual features,
+ while introducing a full-frequency-frame attention (F^3A) module to capture both temporal and
+ frequency contextual information.
+ Additionally, to more realistically evaluate the performance of speech separation models in complex
+ acoustic environments,
+ we introduce a novel dataset called EchoSet . This dataset includes noise and more realistic
+ reverberation
+ (e.g., considering object occlusions and material properties), with speech from two speakers overlapping
+ at random proportions.
+ Experimental results demonstrated that TIGER significantly outperformed
+ state-of-the-art (SOTA) model TF-GridNet
+ on the EchoSet dataset in both inference speed and separation quality, while reducing the number of
+ parameters by 94.3 %
+ and the MACs by 95.3 %. These results indicate that by utilizing frequency band-split and
+ interleaved modeling structures,
+ TIGER achieves a substantial reduction in parameters and computational costs while maintaining high
+ performance.
+ Notably, TIGER is the first speech separation model with fewer than 1 million parameters
+ that achieves performance close to the SOTA model.
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ Overall pipeline of the model architecture of TIGER and
+ its
+ modules .
+
+
+
+
+
+
+ Results on Speech Separation
+
+
+
+ Performance comparisons of TIGER and other existing separation models on Libri2Mix, LRS2-2Mix, and EchoSet .
+ Bold indicates optimal performance, and italics indicate suboptimal performance.
+
+
+
+
+ Efficiency comparisons of TIGER and other models.
+
+
+
+ Results on Cinematic Sound Separation
+
+
+ Comparison of performance and efficiency of cinematic sound
+ separation models on DnR .
+ ‘*’ means the result comes from the original paper of DnR.
+
+
+
+ Speech Separation Demo
+
+
+ EchoSet Sample I
+
+
+
+ Mixture
+
+
+
+
+
+ Your browser does
+ not support the audio element.
+
+
+
+
+
+
+ Ground Truth
+ TF-GridNet
+ TIGER
+
+
+
+
+
+ Your browser does
+ not support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+
+
+ Your browser does
+ not support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+
+
+
+
+ EchoSet Sample II
+
+
+
+ Mixture
+
+
+
+
+
+ Your browser does
+ not support the audio element.
+
+
+
+
+
+
+ Ground Truth
+ TF-GridNet
+ TIGER
+
+
+
+
+
+ Your browser does
+ not support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+
+
+ Your browser does
+ not support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+
+
+
+
+ EchoSet Sample III
+
+
+
+ Mixture
+
+
+
+
+
+ Your browser does
+ not support the audio element.
+
+
+
+
+
+
+ Ground Truth
+ TF-GridNet
+ TIGER
+
+
+
+
+
+ Your browser does
+ not support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+
+
+ Your browser does
+ not support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+ Your browser does not
+ support the audio element.
+
+
+
+
+
+ Cinematic Sound Demo
+
+
+
+
+
+ Deadpool
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ Dialog
+ Music
+ Effect
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ The Wandering Earth
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ Dialog
+ Music
+ Effect
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ Captain America
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ Dialog
+ Music
+ Effect
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ ▶ cslikai.cn's
+ clustrmaps 🌎 .
+
+
+ `_
+
+ Args:
+ d (MutableMapping): Dictionary to be flattened.
+ parent_key (str): String to use as a prefix to all subsequent keys.
+ sep (str): String to use as a separator between two key levels.
+
+ Returns:
+ dict: Single-level dictionary, flattened.
+ """
+ items = []
+ for k, v in d.items():
+ new_key = parent_key + sep + k if parent_key else k
+ if isinstance(v, MutableMapping):
+ items.extend(flatten_dict(v, new_key, sep=sep).items())
+ else:
+ items.append((new_key, v))
+ return dict(items)
+
+
+class AudioLightningModule(pl.LightningModule):
+ def __init__(
+ self,
+ audio_model=None,
+ video_model=None,
+ optimizer=None,
+ loss_func=None,
+ train_loader=None,
+ val_loader=None,
+ test_loader=None,
+ scheduler=None,
+ config=None,
+ ):
+ super().__init__()
+ self.audio_model = audio_model
+ self.video_model = video_model
+ self.optimizer = optimizer
+ self.loss_func = loss_func
+ self.train_loader = train_loader
+ self.val_loader = val_loader
+ self.test_loader = test_loader
+ self.scheduler = scheduler
+ self.config = {} if config is None else config
+ # Speed Aug
+ self.speedperturb = SpeedPerturb(
+ self.config["datamodule"]["data_config"]["sample_rate"],
+ speeds=[95, 100, 105],
+ perturb_prob=1.0
+ )
+ # Save lightning"s AttributeDict under self.hparams
+ self.default_monitor = "val_loss/dataloader_idx_0"
+ self.save_hyperparameters(self.config_to_hparams(self.config))
+ # self.print(self.audio_model)
+ self.validation_step_outputs = []
+ self.test_step_outputs = []
+
+
+ def forward(self, wav, mouth=None):
+ """Applies forward pass of the model.
+
+ Returns:
+ :class:`torch.Tensor`
+ """
+ return self.audio_model(wav)
+
+ def training_step(self, batch, batch_nb):
+ mixtures, targets, _ = batch
+
+ new_targets = []
+ min_len = -1
+ if self.config["training"]["SpeedAug"] == True:
+ with torch.no_grad():
+ for i in range(targets.shape[1]):
+ new_target = self.speedperturb(targets[:, i, :])
+ new_targets.append(new_target)
+ if i == 0:
+ min_len = new_target.shape[-1]
+ else:
+ if new_target.shape[-1] < min_len:
+ min_len = new_target.shape[-1]
+
+ targets = torch.zeros(
+ targets.shape[0],
+ targets.shape[1],
+ min_len,
+ device=targets.device,
+ dtype=torch.float,
+ )
+ for i, new_target in enumerate(new_targets):
+ targets[:, i, :] = new_targets[i][:, 0:min_len]
+
+ mixtures = targets.sum(1)
+ # print(mixtures.shape)
+ est_sources = self(mixtures)
+ loss = self.loss_func["train"](est_sources, targets)
+
+ self.log(
+ "train_loss",
+ loss,
+ on_epoch=True,
+ prog_bar=True,
+ sync_dist=True,
+ logger=True,
+ )
+
+ return {"loss": loss}
+
+
+ def validation_step(self, batch, batch_nb, dataloader_idx):
+ # cal val loss
+ if dataloader_idx == 0:
+ mixtures, targets, _ = batch
+ # print(mixtures.shape)
+ est_sources = self(mixtures)
+ loss = self.loss_func["val"](est_sources, targets)
+ self.log(
+ "val_loss",
+ loss,
+ on_epoch=True,
+ prog_bar=True,
+ sync_dist=True,
+ logger=True,
+ )
+
+ self.validation_step_outputs.append(loss)
+
+ return {"val_loss": loss}
+
+ # cal test loss
+ if (self.trainer.current_epoch) % 10 == 0 and dataloader_idx == 1:
+ mixtures, targets, _ = batch
+ # print(mixtures.shape)
+ est_sources = self(mixtures)
+ tloss = self.loss_func["val"](est_sources, targets)
+ self.log(
+ "test_loss",
+ tloss,
+ on_epoch=True,
+ prog_bar=True,
+ sync_dist=True,
+ logger=True,
+ )
+ self.test_step_outputs.append(tloss)
+ return {"test_loss": tloss}
+
+ def on_validation_epoch_end(self):
+ # val
+ avg_loss = torch.stack(self.validation_step_outputs).mean()
+ val_loss = torch.mean(self.all_gather(avg_loss))
+ self.log(
+ "lr",
+ self.optimizer.param_groups[0]["lr"],
+ on_epoch=True,
+ prog_bar=True,
+ sync_dist=True,
+ )
+ self.logger.experiment.log(
+ {"learning_rate": self.optimizer.param_groups[0]["lr"], "epoch": self.current_epoch}
+ )
+ self.logger.experiment.log(
+ {"val_pit_sisnr": -val_loss, "epoch": self.current_epoch}
+ )
+
+ # test
+ if (self.trainer.current_epoch) % 10 == 0:
+ avg_loss = torch.stack(self.test_step_outputs).mean()
+ test_loss = torch.mean(self.all_gather(avg_loss))
+ self.logger.experiment.log(
+ {"test_pit_sisnr": -test_loss, "epoch": self.current_epoch}
+ )
+ self.validation_step_outputs.clear() # free memory
+ self.test_step_outputs.clear() # free memory
+
+ def configure_optimizers(self):
+ """Initialize optimizers, batch-wise and epoch-wise schedulers."""
+ if self.scheduler is None:
+ return self.optimizer
+
+ if not isinstance(self.scheduler, (list, tuple)):
+ self.scheduler = [self.scheduler] # support multiple schedulers
+
+ epoch_schedulers = []
+ for sched in self.scheduler:
+ if not isinstance(sched, dict):
+ if isinstance(sched, ReduceLROnPlateau):
+ sched = {"scheduler": sched, "monitor": self.default_monitor}
+ epoch_schedulers.append(sched)
+ else:
+ sched.setdefault("monitor", self.default_monitor)
+ sched.setdefault("frequency", 1)
+ # Backward compat
+ if sched["interval"] == "batch":
+ sched["interval"] = "step"
+ assert sched["interval"] in [
+ "epoch",
+ "step",
+ ], "Scheduler interval should be either step or epoch"
+ epoch_schedulers.append(sched)
+ return [self.optimizer], epoch_schedulers
+
+ # def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
+ # if metric is None:
+ # scheduler.step()
+ # else:
+ # scheduler.step(metric)
+
+ def train_dataloader(self):
+ """Training dataloader"""
+ return self.train_loader
+
+ def val_dataloader(self):
+ """Validation dataloader"""
+ return [self.val_loader, self.test_loader]
+
+ def on_save_checkpoint(self, checkpoint):
+ """Overwrite if you want to save more things in the checkpoint."""
+ checkpoint["training_config"] = self.config
+ return checkpoint
+
+ @staticmethod
+ def config_to_hparams(dic):
+ """Sanitizes the config dict to be handled correctly by torch
+ SummaryWriter. It flatten the config dict, converts ``None`` to
+ ``"None"`` and any list and tuple into torch.Tensors.
+
+ Args:
+ dic (dict): Dictionary to be transformed.
+
+ Returns:
+ dict: Transformed dictionary.
+ """
+ dic = flatten_dict(dic)
+ for k, v in dic.items():
+ if v is None:
+ dic[k] = str(v)
+ elif isinstance(v, (list, tuple)):
+ dic[k] = torch.tensor(v)
+ return dic
diff --git a/look2hear/system/audio_litmodule_multidecoder.py b/look2hear/system/audio_litmodule_multidecoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b89663e2a7b58e8050397c06653bb1306ddafb3
--- /dev/null
+++ b/look2hear/system/audio_litmodule_multidecoder.py
@@ -0,0 +1,258 @@
+import torch
+import pytorch_lightning as pl
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from collections.abc import MutableMapping
+# from speechbrain.processing.speech_augmentation import SpeedPerturb
+
+def flatten_dict(d, parent_key="", sep="_"):
+ """Flattens a dictionary into a single-level dictionary while preserving
+ parent keys. Taken from
+ `SO `_
+
+ Args:
+ d (MutableMapping): Dictionary to be flattened.
+ parent_key (str): String to use as a prefix to all subsequent keys.
+ sep (str): String to use as a separator between two key levels.
+
+ Returns:
+ dict: Single-level dictionary, flattened.
+ """
+ items = []
+ for k, v in d.items():
+ new_key = parent_key + sep + k if parent_key else k
+ if isinstance(v, MutableMapping):
+ items.extend(flatten_dict(v, new_key, sep=sep).items())
+ else:
+ items.append((new_key, v))
+ return dict(items)
+
+
+class AudioLightningModuleMultiDecoder(pl.LightningModule):
+ def __init__(
+ self,
+ audio_model=None,
+ video_model=None,
+ optimizer=None,
+ loss_func=None,
+ train_loader=None,
+ val_loader=None,
+ test_loader=None,
+ scheduler=None,
+ config=None,
+ ):
+ super().__init__()
+ self.audio_model = audio_model
+ self.video_model = video_model
+ self.optimizer = optimizer
+ self.loss_func = loss_func
+ self.train_loader = train_loader
+ self.val_loader = val_loader
+ self.test_loader = test_loader
+ self.scheduler = scheduler
+ self.config = {} if config is None else config
+ # Speed Aug
+ # self.speedperturb = SpeedPerturb(
+ # self.config["datamodule"]["data_config"]["sample_rate"],
+ # speeds=[95, 100, 105],
+ # perturb_prob=1.0
+ # )
+ # Save lightning"s AttributeDict under self.hparams
+ self.default_monitor = "val_loss/dataloader_idx_0"
+ self.save_hyperparameters(self.config_to_hparams(self.config))
+ # self.print(self.audio_model)
+ self.validation_step_outputs = []
+ self.test_step_outputs = []
+
+
+ def forward(self, wav, mouth=None):
+ """Applies forward pass of the model.
+
+ Returns:
+ :class:`torch.Tensor`
+ """
+ if self.training:
+ return self.audio_model(wav)
+ else:
+ ests, num_blocks = self.audio_model(wav)
+ assert ests.shape[0] % num_blocks == 0
+ batch_size = int(ests.shape[0] / num_blocks)
+ return ests[-batch_size:]
+
+ def training_step(self, batch, batch_nb):
+ mixtures, targets, _ = batch
+
+ new_targets = []
+ min_len = -1
+ # if self.config["training"]["SpeedAug"] == True:
+ # with torch.no_grad():
+ # for i in range(targets.shape[1]):
+ # new_target = self.speedperturb(targets[:, i, :])
+ # new_targets.append(new_target)
+ # if i == 0:
+ # min_len = new_target.shape[-1]
+ # else:
+ # if new_target.shape[-1] < min_len:
+ # min_len = new_target.shape[-1]
+
+ # targets = torch.zeros(
+ # targets.shape[0],
+ # targets.shape[1],
+ # min_len,
+ # device=targets.device,
+ # dtype=torch.float,
+ # )
+ # for i, new_target in enumerate(new_targets):
+ # targets[:, i, :] = new_targets[i][:, 0:min_len]
+
+ # mixtures = targets.sum(1)
+ # print(mixtures.shape)
+ est_sources, num_blocks = self(mixtures)
+ assert est_sources.shape[0] % num_blocks == 0
+ batch_size = int(est_sources.shape[0] / num_blocks)
+ ests_sources_each_block = []
+ for i in range(num_blocks):
+ ests_sources_each_block.append(est_sources[i * batch_size : (i + 1) * batch_size])
+ loss = self.loss_func["train"](ests_sources_each_block, targets)
+
+ self.log(
+ "train_loss",
+ loss,
+ on_epoch=True,
+ prog_bar=True,
+ sync_dist=True,
+ logger=True,
+ )
+
+ return {"loss": loss}
+
+
+ def validation_step(self, batch, batch_nb, dataloader_idx):
+ # cal val loss
+ if dataloader_idx == 0:
+ mixtures, targets, _ = batch
+ # print(mixtures.shape)
+ est_sources = self(mixtures)
+ loss = self.loss_func["val"](est_sources, targets)
+ self.log(
+ "val_loss",
+ loss,
+ on_epoch=True,
+ prog_bar=True,
+ sync_dist=True,
+ logger=True,
+ )
+
+ self.validation_step_outputs.append(loss)
+
+ return {"val_loss": loss}
+
+ # cal test loss
+ if (self.trainer.current_epoch) % 10 == 0 and dataloader_idx == 1:
+ mixtures, targets, _ = batch
+ # print(mixtures.shape)
+ est_sources = self(mixtures)
+ tloss = self.loss_func["val"](est_sources, targets)
+ self.log(
+ "test_loss",
+ tloss,
+ on_epoch=True,
+ prog_bar=True,
+ sync_dist=True,
+ logger=True,
+ )
+ self.test_step_outputs.append(tloss)
+ return {"test_loss": tloss}
+
+ def on_validation_epoch_end(self):
+ # val
+ avg_loss = torch.stack(self.validation_step_outputs).mean()
+ val_loss = torch.mean(self.all_gather(avg_loss))
+ self.log(
+ "lr",
+ self.optimizer.param_groups[0]["lr"],
+ on_epoch=True,
+ prog_bar=True,
+ sync_dist=True,
+ )
+ self.logger.experiment.log(
+ {"learning_rate": self.optimizer.param_groups[0]["lr"], "epoch": self.current_epoch}
+ )
+ self.logger.experiment.log(
+ {"val_pit_sisnr": -val_loss, "epoch": self.current_epoch}
+ )
+
+ # test
+ if (self.trainer.current_epoch) % 10 == 0:
+ avg_loss = torch.stack(self.test_step_outputs).mean()
+ test_loss = torch.mean(self.all_gather(avg_loss))
+ self.logger.experiment.log(
+ {"test_pit_sisnr": -test_loss, "epoch": self.current_epoch}
+ )
+ self.validation_step_outputs.clear() # free memory
+ self.test_step_outputs.clear() # free memory
+
+ def configure_optimizers(self):
+ """Initialize optimizers, batch-wise and epoch-wise schedulers."""
+ if self.scheduler is None:
+ return self.optimizer
+
+ if not isinstance(self.scheduler, (list, tuple)):
+ self.scheduler = [self.scheduler] # support multiple schedulers
+
+ epoch_schedulers = []
+ for sched in self.scheduler:
+ if not isinstance(sched, dict):
+ if isinstance(sched, ReduceLROnPlateau):
+ sched = {"scheduler": sched, "monitor": self.default_monitor}
+ epoch_schedulers.append(sched)
+ else:
+ sched.setdefault("monitor", self.default_monitor)
+ sched.setdefault("frequency", 1)
+ # Backward compat
+ if sched["interval"] == "batch":
+ sched["interval"] = "step"
+ assert sched["interval"] in [
+ "epoch",
+ "step",
+ ], "Scheduler interval should be either step or epoch"
+ epoch_schedulers.append(sched)
+ return [self.optimizer], epoch_schedulers
+
+ # def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
+ # if metric is None:
+ # scheduler.step()
+ # else:
+ # scheduler.step(metric)
+
+ def train_dataloader(self):
+ """Training dataloader"""
+ return self.train_loader
+
+ def val_dataloader(self):
+ """Validation dataloader"""
+ return [self.val_loader, self.test_loader]
+
+ def on_save_checkpoint(self, checkpoint):
+ """Overwrite if you want to save more things in the checkpoint."""
+ checkpoint["training_config"] = self.config
+ return checkpoint
+
+ @staticmethod
+ def config_to_hparams(dic):
+ """Sanitizes the config dict to be handled correctly by torch
+ SummaryWriter. It flatten the config dict, converts ``None`` to
+ ``"None"`` and any list and tuple into torch.Tensors.
+
+ Args:
+ dic (dict): Dictionary to be transformed.
+
+ Returns:
+ dict: Transformed dictionary.
+ """
+ dic = flatten_dict(dic)
+ for k, v in dic.items():
+ if v is None:
+ dic[k] = str(v)
+ elif isinstance(v, (list, tuple)):
+ dic[k] = torch.tensor(v)
+ return dic
diff --git a/look2hear/system/optimizers.py b/look2hear/system/optimizers.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca3293e5f1cdf319fd1b39bb7c8649816f8d5d76
--- /dev/null
+++ b/look2hear/system/optimizers.py
@@ -0,0 +1,106 @@
+from torch.optim.optimizer import Optimizer
+from torch.optim import Adam, RMSprop, SGD, Adadelta, Adagrad, Adamax, AdamW, ASGD
+from torch_optimizer import (
+ AccSGD,
+ AdaBound,
+ AdaMod,
+ DiffGrad,
+ Lamb,
+ NovoGrad,
+ PID,
+ QHAdam,
+ QHM,
+ RAdam,
+ SGDW,
+ Yogi,
+ Ranger,
+ RangerQH,
+ RangerVA,
+)
+
+
+__all__ = [
+ "AccSGD",
+ "AdaBound",
+ "AdaMod",
+ "DiffGrad",
+ "Lamb",
+ "NovoGrad",
+ "PID",
+ "QHAdam",
+ "QHM",
+ "RAdam",
+ "SGDW",
+ "Yogi",
+ "Ranger",
+ "RangerQH",
+ "RangerVA",
+ "Adam",
+ "RMSprop",
+ "SGD",
+ "Adadelta",
+ "Adagrad",
+ "Adamax",
+ "AdamW",
+ "ASGD",
+ "make_optimizer",
+ "get",
+]
+
+
+def make_optimizer(params, optim_name="adam", **kwargs):
+ """
+
+ Args:
+ params (iterable): Output of `nn.Module.parameters()`.
+ optimizer (str or :class:`torch.optim.Optimizer`): Identifier understood
+ by :func:`~.get`.
+ **kwargs (dict): keyword arguments for the optimizer.
+
+ Returns:
+ torch.optim.Optimizer
+ Examples
+ >>> from torch import nn
+ >>> model = nn.Sequential(nn.Linear(10, 10))
+ >>> optimizer = make_optimizer(model.parameters(), optimizer='sgd',
+ >>> lr=1e-3)
+ """
+ return get(optim_name)(params, **kwargs)
+
+
+def register_optimizer(custom_opt):
+ """Register a custom opt, gettable with `optimzers.get`.
+
+ Args:
+ custom_opt: Custom optimizer to register.
+
+ """
+ if (
+ custom_opt.__name__ in globals().keys()
+ or custom_opt.__name__.lower() in globals().keys()
+ ):
+ raise ValueError(
+ f"Activation {custom_opt.__name__} already exists. Choose another name."
+ )
+ globals().update({custom_opt.__name__: custom_opt})
+
+
+def get(identifier):
+ """Returns an optimizer function from a string. Returns its input if it
+ is callable (already a :class:`torch.optim.Optimizer` for example).
+
+ Args:
+ identifier (str or Callable): the optimizer identifier.
+
+ Returns:
+ :class:`torch.optim.Optimizer` or None
+ """
+ if isinstance(identifier, Optimizer):
+ return identifier
+ elif isinstance(identifier, str):
+ to_get = {k.lower(): v for k, v in globals().items()}
+ cls = to_get.get(identifier.lower())
+ if cls is None:
+ raise ValueError(f"Could not interpret optimizer : {str(identifier)}")
+ return cls
+ raise ValueError(f"Could not interpret optimizer : {str(identifier)}")
diff --git a/look2hear/system/schedulers.py b/look2hear/system/schedulers.py
new file mode 100644
index 0000000000000000000000000000000000000000..085dcb7f1b21f479373b6f7bc7aeb5cef33b7731
--- /dev/null
+++ b/look2hear/system/schedulers.py
@@ -0,0 +1,115 @@
+import torch
+from torch.optim.optimizer import Optimizer
+import pytorch_lightning as pl
+
+
+class BaseScheduler(object):
+ """Base class for the step-wise scheduler logic.
+
+ Args:
+ optimizer (Optimize): Optimizer instance to apply lr schedule on.
+
+ Subclass this and overwrite ``_get_lr`` to write your own step-wise scheduler.
+ """
+
+ def __init__(self, optimizer):
+ self.optimizer = optimizer
+ self.step_num = 0
+
+ def zero_grad(self):
+ self.optimizer.zero_grad()
+
+ def _get_lr(self):
+ raise NotImplementedError
+
+ def _set_lr(self, lr):
+ for param_group in self.optimizer.param_groups:
+ param_group["lr"] = lr
+
+ def step(self, metrics=None, epoch=None):
+ """Update step-wise learning rate before optimizer.step."""
+ self.step_num += 1
+ lr = self._get_lr()
+ self._set_lr(lr)
+
+ def load_state_dict(self, state_dict):
+ self.__dict__.update(state_dict)
+
+ def state_dict(self):
+ return {key: value for key, value in self.__dict__.items() if key != "optimizer"}
+
+ def as_tensor(self, start=0, stop=100_000):
+ """Returns the scheduler values from start to stop."""
+ lr_list = []
+ for _ in range(start, stop):
+ self.step_num += 1
+ lr_list.append(self._get_lr())
+ self.step_num = 0
+ return torch.tensor(lr_list)
+
+ def plot(self, start=0, stop=100_000): # noqa
+ """Plot the scheduler values from start to stop."""
+ import matplotlib.pyplot as plt
+
+ all_lr = self.as_tensor(start=start, stop=stop)
+ plt.plot(all_lr.numpy())
+ plt.show()
+
+class DPTNetScheduler(BaseScheduler):
+ """Dual Path Transformer Scheduler used in [1]
+
+ Args:
+ optimizer (Optimizer): Optimizer instance to apply lr schedule on.
+ steps_per_epoch (int): Number of steps per epoch.
+ d_model(int): The number of units in the layer output.
+ warmup_steps (int): The number of steps in the warmup stage of training.
+ noam_scale (float): Linear increase rate in first phase.
+ exp_max (float): Max learning rate in second phase.
+ exp_base (float): Exp learning rate base in second phase.
+
+ Schedule:
+ This scheduler increases the learning rate linearly for the first
+ ``warmup_steps``, and then decay it by 0.98 for every two epochs.
+
+ References
+ [1]: Jingjing Chen et al. "Dual-Path Transformer Network: Direct Context-
+ Aware Modeling for End-to-End Monaural Speech Separation" Interspeech 2020.
+ """
+
+ def __init__(
+ self,
+ optimizer,
+ steps_per_epoch,
+ d_model,
+ warmup_steps=4000,
+ noam_scale=1.0,
+ exp_max=0.0004,
+ exp_base=0.98,
+ ):
+ super().__init__(optimizer)
+ self.noam_scale = noam_scale
+ self.d_model = d_model
+ self.warmup_steps = warmup_steps
+ self.exp_max = exp_max
+ self.exp_base = exp_base
+ self.steps_per_epoch = steps_per_epoch
+ self.epoch = 0
+
+ def _get_lr(self):
+ if self.step_num % self.steps_per_epoch == 0:
+ self.epoch += 1
+
+ if self.step_num > self.warmup_steps:
+ # exp decaying
+ lr = self.exp_max * (self.exp_base ** ((self.epoch - 1) // 2))
+ else:
+ # noam
+ lr = (
+ self.noam_scale
+ * self.d_model ** (-0.5)
+ * min(self.step_num ** (-0.5), self.step_num * self.warmup_steps ** (-1.5))
+ )
+ return lr
+
+# Backward compat
+_BaseScheduler = BaseScheduler
\ No newline at end of file
diff --git a/look2hear/utils/__init__.py b/look2hear/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9d90195f22f7a5eb37ca18d9b17bfc2f9126e5c
--- /dev/null
+++ b/look2hear/utils/__init__.py
@@ -0,0 +1,41 @@
+from .stft import STFT
+from .torch_utils import pad_x_to_y, shape_reconstructed, tensors_to_device
+from .parser_utils import (
+ prepare_parser_from_dict,
+ parse_args_as_dict,
+ str_int_float,
+ str2bool,
+ str2bool_arg,
+ isfloat,
+ isint,
+)
+from .lightning_utils import print_only, RichProgressBarTheme, MyRichProgressBar, BatchesProcessedColumn, MyMetricsTextColumn
+from .complex_utils import is_complex, is_torch_complex_tensor, new_complex_like
+from .get_layer_from_string import get_layer
+from .inversible_interface import InversibleInterface
+from .nets_utils import make_pad_mask
+
+__all__ = [
+ "STFT",
+ "pad_x_to_y",
+ "shape_reconstructed",
+ "tensors_to_device",
+ "prepare_parser_from_dict",
+ "parse_args_as_dict",
+ "str_int_float",
+ "str2bool",
+ "str2bool_arg",
+ "isfloat",
+ "isint",
+ "print_only",
+ "RichProgressBarTheme",
+ "MyRichProgressBar",
+ "BatchesProcessedColumn",
+ "MyMetricsTextColumn",
+ "is_complex",
+ "is_torch_complex_tensor",
+ "new_complex_like",
+ "get_layer",
+ "InversibleInterface",
+ "make_pad_mask",
+]
diff --git a/look2hear/utils/complex_utils.py b/look2hear/utils/complex_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..2e353197bfad9a62731a3265fd43bf628840b54d
--- /dev/null
+++ b/look2hear/utils/complex_utils.py
@@ -0,0 +1,191 @@
+"""Beamformer module."""
+from typing import Sequence, Tuple, Union
+
+import torch
+from packaging.version import parse as V
+from torch_complex import functional as FC
+from torch_complex.tensor import ComplexTensor
+
+EPS = torch.finfo(torch.double).eps
+is_torch_1_8_plus = V(torch.__version__) >= V("1.8.0")
+is_torch_1_9_plus = V(torch.__version__) >= V("1.9.0")
+
+
+def new_complex_like(
+ ref: Union[torch.Tensor, ComplexTensor],
+ real_imag: Tuple[torch.Tensor, torch.Tensor],
+):
+ if isinstance(ref, ComplexTensor):
+ return ComplexTensor(*real_imag)
+ elif is_torch_complex_tensor(ref):
+ return torch.complex(*real_imag)
+ else:
+ raise ValueError(
+ "Please update your PyTorch version to 1.9+ for complex support."
+ )
+
+
+def is_torch_complex_tensor(c):
+ return (
+ not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c)
+ )
+
+
+def is_complex(c):
+ return isinstance(c, ComplexTensor) or is_torch_complex_tensor(c)
+
+
+def to_double(c):
+ if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c):
+ return c.to(dtype=torch.complex128)
+ else:
+ return c.double()
+
+
+def to_float(c):
+ if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c):
+ return c.to(dtype=torch.complex64)
+ else:
+ return c.float()
+
+
+def cat(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs):
+ if not isinstance(seq, (list, tuple)):
+ raise TypeError(
+ "cat(): argument 'tensors' (position 1) must be tuple of Tensors, "
+ "not Tensor"
+ )
+ if isinstance(seq[0], ComplexTensor):
+ return FC.cat(seq, *args, **kwargs)
+ else:
+ return torch.cat(seq, *args, **kwargs)
+
+
+def complex_norm(
+ c: Union[torch.Tensor, ComplexTensor], dim=-1, keepdim=False
+) -> torch.Tensor:
+ if not is_complex(c):
+ raise TypeError("Input is not a complex tensor.")
+ if is_torch_complex_tensor(c):
+ return torch.norm(c, dim=dim, keepdim=keepdim)
+ else:
+ if dim is None:
+ return torch.sqrt((c.real**2 + c.imag**2).sum() + EPS)
+ else:
+ return torch.sqrt(
+ (c.real**2 + c.imag**2).sum(dim=dim, keepdim=keepdim) + EPS
+ )
+
+
+def einsum(equation, *operands):
+ # NOTE: Do not mix ComplexTensor and torch.complex in the input!
+ # NOTE (wangyou): Until PyTorch 1.9.0, torch.einsum does not support
+ # mixed input with complex and real tensors.
+ if len(operands) == 1:
+ if isinstance(operands[0], (tuple, list)):
+ operands = operands[0]
+ complex_module = FC if isinstance(operands[0], ComplexTensor) else torch
+ return complex_module.einsum(equation, *operands)
+ elif len(operands) != 2:
+ op0 = operands[0]
+ same_type = all(op.dtype == op0.dtype for op in operands[1:])
+ if same_type:
+ _einsum = FC.einsum if isinstance(op0, ComplexTensor) else torch.einsum
+ return _einsum(equation, *operands)
+ else:
+ raise ValueError("0 or More than 2 operands are not supported.")
+ a, b = operands
+ if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
+ return FC.einsum(equation, a, b)
+ elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
+ if not torch.is_complex(a):
+ o_real = torch.einsum(equation, a, b.real)
+ o_imag = torch.einsum(equation, a, b.imag)
+ return torch.complex(o_real, o_imag)
+ elif not torch.is_complex(b):
+ o_real = torch.einsum(equation, a.real, b)
+ o_imag = torch.einsum(equation, a.imag, b)
+ return torch.complex(o_real, o_imag)
+ else:
+ return torch.einsum(equation, a, b)
+ else:
+ return torch.einsum(equation, a, b)
+
+
+def inverse(
+ c: Union[torch.Tensor, ComplexTensor]
+) -> Union[torch.Tensor, ComplexTensor]:
+ if isinstance(c, ComplexTensor):
+ return c.inverse2()
+ else:
+ return c.inverse()
+
+
+def matmul(
+ a: Union[torch.Tensor, ComplexTensor], b: Union[torch.Tensor, ComplexTensor]
+) -> Union[torch.Tensor, ComplexTensor]:
+ # NOTE: Do not mix ComplexTensor and torch.complex in the input!
+ # NOTE (wangyou): Until PyTorch 1.9.0, torch.matmul does not support
+ # multiplication between complex and real tensors.
+ if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
+ return FC.matmul(a, b)
+ elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
+ if not torch.is_complex(a):
+ o_real = torch.matmul(a, b.real)
+ o_imag = torch.matmul(a, b.imag)
+ return torch.complex(o_real, o_imag)
+ elif not torch.is_complex(b):
+ o_real = torch.matmul(a.real, b)
+ o_imag = torch.matmul(a.imag, b)
+ return torch.complex(o_real, o_imag)
+ else:
+ return torch.matmul(a, b)
+ else:
+ return torch.matmul(a, b)
+
+
+def trace(a: Union[torch.Tensor, ComplexTensor]):
+ # NOTE (wangyou): until PyTorch 1.9.0, torch.trace does not
+ # support bacth processing. Use FC.trace() as fallback.
+ return FC.trace(a)
+
+
+def reverse(a: Union[torch.Tensor, ComplexTensor], dim=0):
+ if isinstance(a, ComplexTensor):
+ return FC.reverse(a, dim=dim)
+ else:
+ return torch.flip(a, dims=(dim,))
+
+
+def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]):
+ """Solve the linear equation ax = b."""
+ # NOTE: Do not mix ComplexTensor and torch.complex in the input!
+ # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support
+ # mixed input with complex and real tensors.
+ if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
+ if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
+ return FC.solve(b, a, return_LU=False)
+ else:
+ return matmul(inverse(a), b)
+ elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
+ if torch.is_complex(a) and torch.is_complex(b):
+ return torch.linalg.solve(a, b)
+ else:
+ return matmul(inverse(a), b)
+ else:
+ if is_torch_1_8_plus:
+ return torch.linalg.solve(a, b)
+ else:
+ return torch.solve(b, a)[0]
+
+
+def stack(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs):
+ if not isinstance(seq, (list, tuple)):
+ raise TypeError(
+ "stack(): argument 'tensors' (position 1) must be tuple of Tensors, "
+ "not Tensor"
+ )
+ if isinstance(seq[0], ComplexTensor):
+ return FC.stack(seq, *args, **kwargs)
+ else:
+ return torch.stack(seq, *args, **kwargs)
\ No newline at end of file
diff --git a/look2hear/utils/get_layer_from_string.py b/look2hear/utils/get_layer_from_string.py
new file mode 100644
index 0000000000000000000000000000000000000000..3987a8df65fa6fd6a69f74d2d3aa2bc06971512e
--- /dev/null
+++ b/look2hear/utils/get_layer_from_string.py
@@ -0,0 +1,43 @@
+import difflib
+
+import torch
+
+
+def get_layer(l_name, library=torch.nn):
+ """Return layer object handler from library e.g. from torch.nn
+
+ E.g. if l_name=="elu", returns torch.nn.ELU.
+
+ Args:
+ l_name (string): Case insensitive name for layer in library (e.g. .'elu').
+ library (module): Name of library/module where to search for object handler
+ with l_name e.g. "torch.nn".
+
+ Returns:
+ layer_handler (object): handler for the requested layer e.g. (torch.nn.ELU)
+
+ """
+
+ all_torch_layers = [x for x in dir(torch.nn)]
+ match = [x for x in all_torch_layers if l_name.lower() == x.lower()]
+ if len(match) == 0:
+ close_matches = difflib.get_close_matches(
+ l_name, [x.lower() for x in all_torch_layers]
+ )
+ raise NotImplementedError(
+ "Layer with name {} not found in {}.\n Closest matches: {}".format(
+ l_name, str(library), close_matches
+ )
+ )
+ elif len(match) > 1:
+ close_matches = difflib.get_close_matches(
+ l_name, [x.lower() for x in all_torch_layers]
+ )
+ raise NotImplementedError(
+ "Multiple matchs for layer with name {} not found in {}.\n "
+ "All matches: {}".format(l_name, str(library), close_matches)
+ )
+ else:
+ # valid
+ layer_handler = getattr(library, match[0])
+ return layer_handler
\ No newline at end of file
diff --git a/look2hear/utils/inversible_interface.py b/look2hear/utils/inversible_interface.py
new file mode 100644
index 0000000000000000000000000000000000000000..919107ea0c1c8b9e35b957ab3109cf11faa6067f
--- /dev/null
+++ b/look2hear/utils/inversible_interface.py
@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+from typing import Tuple
+
+import torch
+
+
+class InversibleInterface(ABC):
+ @abstractmethod
+ def inverse(
+ self, input: torch.Tensor, input_lengths: torch.Tensor = None
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ # return output, output_lengths
+ raise NotImplementedError
\ No newline at end of file
diff --git a/look2hear/utils/lightning_utils.py b/look2hear/utils/lightning_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..29f2f73cc3ca6c4587c63ed5c3ef051609a0cb1a
--- /dev/null
+++ b/look2hear/utils/lightning_utils.py
@@ -0,0 +1,104 @@
+from rich import print
+from dataclasses import dataclass
+from pytorch_lightning.utilities import rank_zero_only
+from typing import Union
+from pytorch_lightning.callbacks.progress.rich_progress import *
+from rich.console import Console, RenderableType
+from rich.progress_bar import ProgressBar
+from rich.style import Style
+from rich.text import Text
+from rich.progress import (
+ BarColumn,
+ DownloadColumn,
+ Progress,
+ TaskID,
+ TextColumn,
+ TimeRemainingColumn,
+ TransferSpeedColumn,
+ ProgressColumn
+)
+from rich import print, reconfigure
+
+@rank_zero_only
+def print_only(message: str):
+ print(message)
+
+@dataclass
+class RichProgressBarTheme:
+ """Styles to associate to different base components.
+
+ Args:
+ description: Style for the progress bar description. For eg., Epoch x, Testing, etc.
+ progress_bar: Style for the bar in progress.
+ progress_bar_finished: Style for the finished progress bar.
+ progress_bar_pulse: Style for the progress bar when `IterableDataset` is being processed.
+ batch_progress: Style for the progress tracker (i.e 10/50 batches completed).
+ time: Style for the processed time and estimate time remaining.
+ processing_speed: Style for the speed of the batches being processed.
+ metrics: Style for the metrics
+
+ https://rich.readthedocs.io/en/stable/style.html
+ """
+
+ description: Union[str, Style] = "#FF4500"
+ progress_bar: Union[str, Style] = "#f92672"
+ progress_bar_finished: Union[str, Style] = "#b7cc8a"
+ progress_bar_pulse: Union[str, Style] = "#f92672"
+ batch_progress: Union[str, Style] = "#fc608a"
+ time: Union[str, Style] = "#45ada2"
+ processing_speed: Union[str, Style] = "#DC143C"
+ metrics: Union[str, Style] = "#228B22"
+
+class BatchesProcessedColumn(ProgressColumn):
+ def __init__(self, style: Union[str, Style]):
+ self.style = style
+ super().__init__()
+
+ def render(self, task) -> RenderableType:
+ total = task.total if task.total != float("inf") else "--"
+ return Text(f"{int(task.completed)}/{int(total)}", style=self.style)
+
+class MyMetricsTextColumn(ProgressColumn):
+ """A column containing text."""
+
+ def __init__(self, style):
+ self._tasks = {}
+ self._current_task_id = 0
+ self._metrics = {}
+ self._style = style
+ super().__init__()
+
+ def update(self, metrics):
+ # Called when metrics are ready to be rendered.
+ # This is to prevent render from causing deadlock issues by requesting metrics
+ # in separate threads.
+ self._metrics = metrics
+
+ def render(self, task) -> Text:
+ text = ""
+ for k, v in self._metrics.items():
+ text += f"{k}: {round(v, 3) if isinstance(v, float) else v} "
+ return Text(text, justify="left", style=self._style)
+
+class MyRichProgressBar(RichProgressBar):
+ """A progress bar prints metrics at the end of each epoch
+ """
+
+ def _init_progress(self, trainer):
+ if self.is_enabled and (self.progress is None or self._progress_stopped):
+ self._reset_progress_bar_ids()
+ reconfigure(**self._console_kwargs)
+ # file = open("Look2Hear/Experiments/run_logs/EdgeFRCNN-Noncausal.log", 'w')
+ self._console: Console = Console(force_terminal=True)
+ self._console.clear_live()
+ self._metric_component = MetricsTextColumn(trainer, self.theme.metrics)
+ self.progress = CustomProgress(
+ *self.configure_columns(trainer),
+ self._metric_component,
+ auto_refresh=False,
+ disable=self.is_disabled,
+ console=self._console,
+ )
+ self.progress.start()
+ # progress has started
+ self._progress_stopped = False
\ No newline at end of file
diff --git a/look2hear/utils/nets_utils.py b/look2hear/utils/nets_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e7a677023d87e6370352454994d12e458437e2c2
--- /dev/null
+++ b/look2hear/utils/nets_utils.py
@@ -0,0 +1,503 @@
+# -*- coding: utf-8 -*-
+
+"""Network related utility tools."""
+
+import logging
+from typing import Dict
+
+import numpy as np
+import torch
+
+
+def to_device(m, x):
+ """Send tensor into the device of the module.
+
+ Args:
+ m (torch.nn.Module): Torch module.
+ x (Tensor): Torch tensor.
+
+ Returns:
+ Tensor: Torch tensor located in the same place as torch module.
+
+ """
+ if isinstance(m, torch.nn.Module):
+ device = next(m.parameters()).device
+ elif isinstance(m, torch.Tensor):
+ device = m.device
+ else:
+ raise TypeError(
+ "Expected torch.nn.Module or torch.tensor, " f"bot got: {type(m)}"
+ )
+ return x.to(device)
+
+
+def pad_list(xs, pad_value):
+ """Perform padding for the list of tensors.
+
+ Args:
+ xs (List): List of Tensors [(T_1, `*`), (T_2, `*`), ..., (T_B, `*`)].
+ pad_value (float): Value for padding.
+
+ Returns:
+ Tensor: Padded tensor (B, Tmax, `*`).
+
+ Examples:
+ >>> x = [torch.ones(4), torch.ones(2), torch.ones(1)]
+ >>> x
+ [tensor([1., 1., 1., 1.]), tensor([1., 1.]), tensor([1.])]
+ >>> pad_list(x, 0)
+ tensor([[1., 1., 1., 1.],
+ [1., 1., 0., 0.],
+ [1., 0., 0., 0.]])
+
+ """
+ n_batch = len(xs)
+ max_len = max(x.size(0) for x in xs)
+ pad = xs[0].new(n_batch, max_len, *xs[0].size()[1:]).fill_(pad_value)
+
+ for i in range(n_batch):
+ pad[i, : xs[i].size(0)] = xs[i]
+
+ return pad
+
+
+def make_pad_mask(lengths, xs=None, length_dim=-1, maxlen=None):
+ """Make mask tensor containing indices of padded part.
+
+ Args:
+ lengths (LongTensor or List): Batch of lengths (B,).
+ xs (Tensor, optional): The reference tensor.
+ If set, masks will be the same shape as this tensor.
+ length_dim (int, optional): Dimension indicator of the above tensor.
+ See the example.
+
+ Returns:
+ Tensor: Mask tensor containing indices of padded part.
+ dtype=torch.uint8 in PyTorch 1.2-
+ dtype=torch.bool in PyTorch 1.2+ (including 1.2)
+
+ Examples:
+ With only lengths.
+
+ >>> lengths = [5, 3, 2]
+ >>> make_pad_mask(lengths)
+ masks = [[0, 0, 0, 0 ,0],
+ [0, 0, 0, 1, 1],
+ [0, 0, 1, 1, 1]]
+
+ With the reference tensor.
+
+ >>> xs = torch.zeros((3, 2, 4))
+ >>> make_pad_mask(lengths, xs)
+ tensor([[[0, 0, 0, 0],
+ [0, 0, 0, 0]],
+ [[0, 0, 0, 1],
+ [0, 0, 0, 1]],
+ [[0, 0, 1, 1],
+ [0, 0, 1, 1]]], dtype=torch.uint8)
+ >>> xs = torch.zeros((3, 2, 6))
+ >>> make_pad_mask(lengths, xs)
+ tensor([[[0, 0, 0, 0, 0, 1],
+ [0, 0, 0, 0, 0, 1]],
+ [[0, 0, 0, 1, 1, 1],
+ [0, 0, 0, 1, 1, 1]],
+ [[0, 0, 1, 1, 1, 1],
+ [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8)
+
+ With the reference tensor and dimension indicator.
+
+ >>> xs = torch.zeros((3, 6, 6))
+ >>> make_pad_mask(lengths, xs, 1)
+ tensor([[[0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [1, 1, 1, 1, 1, 1]],
+ [[0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1]],
+ [[0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1]]], dtype=torch.uint8)
+ >>> make_pad_mask(lengths, xs, 2)
+ tensor([[[0, 0, 0, 0, 0, 1],
+ [0, 0, 0, 0, 0, 1],
+ [0, 0, 0, 0, 0, 1],
+ [0, 0, 0, 0, 0, 1],
+ [0, 0, 0, 0, 0, 1],
+ [0, 0, 0, 0, 0, 1]],
+ [[0, 0, 0, 1, 1, 1],
+ [0, 0, 0, 1, 1, 1],
+ [0, 0, 0, 1, 1, 1],
+ [0, 0, 0, 1, 1, 1],
+ [0, 0, 0, 1, 1, 1],
+ [0, 0, 0, 1, 1, 1]],
+ [[0, 0, 1, 1, 1, 1],
+ [0, 0, 1, 1, 1, 1],
+ [0, 0, 1, 1, 1, 1],
+ [0, 0, 1, 1, 1, 1],
+ [0, 0, 1, 1, 1, 1],
+ [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8)
+
+ """
+ if length_dim == 0:
+ raise ValueError("length_dim cannot be 0: {}".format(length_dim))
+
+ if not isinstance(lengths, list):
+ lengths = lengths.long().tolist()
+
+ bs = int(len(lengths))
+ if maxlen is None:
+ if xs is None:
+ maxlen = int(max(lengths))
+ else:
+ maxlen = xs.size(length_dim)
+ else:
+ assert xs is None
+ assert maxlen >= int(max(lengths))
+
+ seq_range = torch.arange(0, maxlen, dtype=torch.int64)
+ seq_range_expand = seq_range.unsqueeze(0).expand(bs, maxlen)
+ seq_length_expand = seq_range_expand.new(lengths).unsqueeze(-1)
+ mask = seq_range_expand >= seq_length_expand
+
+ if xs is not None:
+ assert xs.size(0) == bs, (xs.size(0), bs)
+
+ if length_dim < 0:
+ length_dim = xs.dim() + length_dim
+ # ind = (:, None, ..., None, :, , None, ..., None)
+ ind = tuple(
+ slice(None) if i in (0, length_dim) else None for i in range(xs.dim())
+ )
+ mask = mask[ind].expand_as(xs).to(xs.device)
+ return mask
+
+
+def make_non_pad_mask(lengths, xs=None, length_dim=-1):
+ """Make mask tensor containing indices of non-padded part.
+
+ Args:
+ lengths (LongTensor or List): Batch of lengths (B,).
+ xs (Tensor, optional): The reference tensor.
+ If set, masks will be the same shape as this tensor.
+ length_dim (int, optional): Dimension indicator of the above tensor.
+ See the example.
+
+ Returns:
+ ByteTensor: mask tensor containing indices of padded part.
+ dtype=torch.uint8 in PyTorch 1.2-
+ dtype=torch.bool in PyTorch 1.2+ (including 1.2)
+
+ Examples:
+ With only lengths.
+
+ >>> lengths = [5, 3, 2]
+ >>> make_non_pad_mask(lengths)
+ masks = [[1, 1, 1, 1 ,1],
+ [1, 1, 1, 0, 0],
+ [1, 1, 0, 0, 0]]
+
+ With the reference tensor.
+
+ >>> xs = torch.zeros((3, 2, 4))
+ >>> make_non_pad_mask(lengths, xs)
+ tensor([[[1, 1, 1, 1],
+ [1, 1, 1, 1]],
+ [[1, 1, 1, 0],
+ [1, 1, 1, 0]],
+ [[1, 1, 0, 0],
+ [1, 1, 0, 0]]], dtype=torch.uint8)
+ >>> xs = torch.zeros((3, 2, 6))
+ >>> make_non_pad_mask(lengths, xs)
+ tensor([[[1, 1, 1, 1, 1, 0],
+ [1, 1, 1, 1, 1, 0]],
+ [[1, 1, 1, 0, 0, 0],
+ [1, 1, 1, 0, 0, 0]],
+ [[1, 1, 0, 0, 0, 0],
+ [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8)
+
+ With the reference tensor and dimension indicator.
+
+ >>> xs = torch.zeros((3, 6, 6))
+ >>> make_non_pad_mask(lengths, xs, 1)
+ tensor([[[1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [0, 0, 0, 0, 0, 0]],
+ [[1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0]],
+ [[1, 1, 1, 1, 1, 1],
+ [1, 1, 1, 1, 1, 1],
+ [0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0]]], dtype=torch.uint8)
+ >>> make_non_pad_mask(lengths, xs, 2)
+ tensor([[[1, 1, 1, 1, 1, 0],
+ [1, 1, 1, 1, 1, 0],
+ [1, 1, 1, 1, 1, 0],
+ [1, 1, 1, 1, 1, 0],
+ [1, 1, 1, 1, 1, 0],
+ [1, 1, 1, 1, 1, 0]],
+ [[1, 1, 1, 0, 0, 0],
+ [1, 1, 1, 0, 0, 0],
+ [1, 1, 1, 0, 0, 0],
+ [1, 1, 1, 0, 0, 0],
+ [1, 1, 1, 0, 0, 0],
+ [1, 1, 1, 0, 0, 0]],
+ [[1, 1, 0, 0, 0, 0],
+ [1, 1, 0, 0, 0, 0],
+ [1, 1, 0, 0, 0, 0],
+ [1, 1, 0, 0, 0, 0],
+ [1, 1, 0, 0, 0, 0],
+ [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8)
+
+ """
+ return ~make_pad_mask(lengths, xs, length_dim)
+
+
+def mask_by_length(xs, lengths, fill=0):
+ """Mask tensor according to length.
+
+ Args:
+ xs (Tensor): Batch of input tensor (B, `*`).
+ lengths (LongTensor or List): Batch of lengths (B,).
+ fill (int or float): Value to fill masked part.
+
+ Returns:
+ Tensor: Batch of masked input tensor (B, `*`).
+
+ Examples:
+ >>> x = torch.arange(5).repeat(3, 1) + 1
+ >>> x
+ tensor([[1, 2, 3, 4, 5],
+ [1, 2, 3, 4, 5],
+ [1, 2, 3, 4, 5]])
+ >>> lengths = [5, 3, 2]
+ >>> mask_by_length(x, lengths)
+ tensor([[1, 2, 3, 4, 5],
+ [1, 2, 3, 0, 0],
+ [1, 2, 0, 0, 0]])
+
+ """
+ assert xs.size(0) == len(lengths)
+ ret = xs.data.new(*xs.size()).fill_(fill)
+ for i, l in enumerate(lengths):
+ ret[i, :l] = xs[i, :l]
+ return ret
+
+
+def th_accuracy(pad_outputs, pad_targets, ignore_label):
+ """Calculate accuracy.
+
+ Args:
+ pad_outputs (Tensor): Prediction tensors (B * Lmax, D).
+ pad_targets (LongTensor): Target label tensors (B, Lmax, D).
+ ignore_label (int): Ignore label id.
+
+ Returns:
+ float: Accuracy value (0.0 - 1.0).
+
+ """
+ pad_pred = pad_outputs.view(
+ pad_targets.size(0), pad_targets.size(1), pad_outputs.size(1)
+ ).argmax(2)
+ mask = pad_targets != ignore_label
+ numerator = torch.sum(
+ pad_pred.masked_select(mask) == pad_targets.masked_select(mask)
+ )
+ denominator = torch.sum(mask)
+ return float(numerator) / float(denominator)
+
+
+def to_torch_tensor(x):
+ """Change to torch.Tensor or ComplexTensor from numpy.ndarray.
+
+ Args:
+ x: Inputs. It should be one of numpy.ndarray, Tensor, ComplexTensor, and dict.
+
+ Returns:
+ Tensor or ComplexTensor: Type converted inputs.
+
+ Examples:
+ >>> xs = np.ones(3, dtype=np.float32)
+ >>> xs = to_torch_tensor(xs)
+ tensor([1., 1., 1.])
+ >>> xs = torch.ones(3, 4, 5)
+ >>> assert to_torch_tensor(xs) is xs
+ >>> xs = {'real': xs, 'imag': xs}
+ >>> to_torch_tensor(xs)
+ ComplexTensor(
+ Real:
+ tensor([1., 1., 1.])
+ Imag;
+ tensor([1., 1., 1.])
+ )
+
+ """
+ # If numpy, change to torch tensor
+ if isinstance(x, np.ndarray):
+ if x.dtype.kind == "c":
+ # Dynamically importing because torch_complex requires python3
+ from torch_complex.tensor import ComplexTensor
+
+ return ComplexTensor(x)
+ else:
+ return torch.from_numpy(x)
+
+ # If {'real': ..., 'imag': ...}, convert to ComplexTensor
+ elif isinstance(x, dict):
+ # Dynamically importing because torch_complex requires python3
+ from torch_complex.tensor import ComplexTensor
+
+ if "real" not in x or "imag" not in x:
+ raise ValueError("has 'real' and 'imag' keys: {}".format(list(x)))
+ # Relative importing because of using python3 syntax
+ return ComplexTensor(x["real"], x["imag"])
+
+ # If torch.Tensor, as it is
+ elif isinstance(x, torch.Tensor):
+ return x
+
+ else:
+ error = (
+ "x must be numpy.ndarray, torch.Tensor or a dict like "
+ "{{'real': torch.Tensor, 'imag': torch.Tensor}}, "
+ "but got {}".format(type(x))
+ )
+ try:
+ from torch_complex.tensor import ComplexTensor
+ except Exception:
+ # If PY2
+ raise ValueError(error)
+ else:
+ # If PY3
+ if isinstance(x, ComplexTensor):
+ return x
+ else:
+ raise ValueError(error)
+
+
+def get_subsample(train_args, mode, arch):
+ """Parse the subsampling factors from the args for the specified `mode` and `arch`.
+
+ Args:
+ train_args: argument Namespace containing options.
+ mode: one of ('asr', 'mt', 'st')
+ arch: one of ('rnn', 'rnn-t', 'rnn_mix', 'rnn_mulenc', 'transformer')
+
+ Returns:
+ np.ndarray / List[np.ndarray]: subsampling factors.
+ """
+ if arch == "transformer":
+ return np.array([1])
+
+ elif mode == "mt" and arch == "rnn":
+ # +1 means input (+1) and layers outputs (train_args.elayer)
+ subsample = np.ones(train_args.elayers + 1, dtype=np.int64)
+ logging.warning("Subsampling is not performed for machine translation.")
+ logging.info("subsample: " + " ".join([str(x) for x in subsample]))
+ return subsample
+
+ elif (
+ (mode == "asr" and arch in ("rnn", "rnn-t"))
+ or (mode == "mt" and arch == "rnn")
+ or (mode == "st" and arch == "rnn")
+ ):
+ subsample = np.ones(train_args.elayers + 1, dtype=np.int64)
+ if train_args.etype.endswith("p") and not train_args.etype.startswith("vgg"):
+ ss = train_args.subsample.split("_")
+ for j in range(min(train_args.elayers + 1, len(ss))):
+ subsample[j] = int(ss[j])
+ else:
+ logging.warning(
+ "Subsampling is not performed for vgg*. "
+ "It is performed in max pooling layers at CNN."
+ )
+ logging.info("subsample: " + " ".join([str(x) for x in subsample]))
+ return subsample
+
+ elif mode == "asr" and arch == "rnn_mix":
+ subsample = np.ones(
+ train_args.elayers_sd + train_args.elayers + 1, dtype=np.int64
+ )
+ if train_args.etype.endswith("p") and not train_args.etype.startswith("vgg"):
+ ss = train_args.subsample.split("_")
+ for j in range(
+ min(train_args.elayers_sd + train_args.elayers + 1, len(ss))
+ ):
+ subsample[j] = int(ss[j])
+ else:
+ logging.warning(
+ "Subsampling is not performed for vgg*. "
+ "It is performed in max pooling layers at CNN."
+ )
+ logging.info("subsample: " + " ".join([str(x) for x in subsample]))
+ return subsample
+
+ elif mode == "asr" and arch == "rnn_mulenc":
+ subsample_list = []
+ for idx in range(train_args.num_encs):
+ subsample = np.ones(train_args.elayers[idx] + 1, dtype=np.int64)
+ if train_args.etype[idx].endswith("p") and not train_args.etype[
+ idx
+ ].startswith("vgg"):
+ ss = train_args.subsample[idx].split("_")
+ for j in range(min(train_args.elayers[idx] + 1, len(ss))):
+ subsample[j] = int(ss[j])
+ else:
+ logging.warning(
+ "Encoder %d: Subsampling is not performed for vgg*. "
+ "It is performed in max pooling layers at CNN.",
+ idx + 1,
+ )
+ logging.info("subsample: " + " ".join([str(x) for x in subsample]))
+ subsample_list.append(subsample)
+ return subsample_list
+
+ else:
+ raise ValueError("Invalid options: mode={}, arch={}".format(mode, arch))
+
+
+def rename_state_dict(
+ old_prefix: str, new_prefix: str, state_dict: Dict[str, torch.Tensor]
+):
+ """Replace keys of old prefix with new prefix in state dict."""
+ # need this list not to break the dict iterator
+ old_keys = [k for k in state_dict if k.startswith(old_prefix)]
+ if len(old_keys) > 0:
+ logging.warning(f"Rename: {old_prefix} -> {new_prefix}")
+ for k in old_keys:
+ v = state_dict.pop(k)
+ new_k = k.replace(old_prefix, new_prefix)
+ state_dict[new_k] = v
+
+
+def get_activation(act):
+ """Return activation function."""
+ # Lazy load to avoid unused import
+ from espnet.nets.pytorch_backend.conformer.swish import Swish
+
+ activation_funcs = {
+ "hardtanh": torch.nn.Hardtanh,
+ "tanh": torch.nn.Tanh,
+ "relu": torch.nn.ReLU,
+ "selu": torch.nn.SELU,
+ "swish": Swish,
+ }
+
+ return activation_funcs[act]()
\ No newline at end of file
diff --git a/look2hear/utils/parser_utils.py b/look2hear/utils/parser_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..6cec1eaad6e74a24fd94be0c34ff54b001b22628
--- /dev/null
+++ b/look2hear/utils/parser_utils.py
@@ -0,0 +1,151 @@
+import sys
+import argparse
+
+
+def prepare_parser_from_dict(dic, parser=None):
+ """Prepare an argparser from a dictionary.
+
+ Args:
+ dic (dict): Two-level config dictionary with unique bottom-level keys.
+ parser (argparse.ArgumentParser, optional): If a parser already
+ exists, add the keys from the dictionary on the top of it.
+
+ Returns:
+ argparse.ArgumentParser:
+ Parser instance with groups corresponding to the first level keys
+ and arguments corresponding to the second level keys with default
+ values given by the values.
+ """
+
+ def standardized_entry_type(value):
+ """If the default value is None, replace NoneType by str_int_float.
+ If the default value is boolean, look for boolean strings."""
+ if value is None:
+ return str_int_float
+ if isinstance(str2bool(value), bool):
+ return str2bool_arg
+ return type(value)
+
+ if parser is None:
+ parser = argparse.ArgumentParser()
+ for k in dic.keys():
+ group = parser.add_argument_group(k)
+ if isinstance(dic[k], list):
+ entry_type = standardized_entry_type(dic[k])
+ group.add_argument("--" + k, default=dic[k], type=entry_type)
+ elif isinstance(dic[k], dict):
+ for kk in dic[k].keys():
+ entry_type = standardized_entry_type(dic[k][kk])
+ group.add_argument("--" + kk, default=dic[k][kk], type=entry_type)
+ elif isinstance(dic[k], str):
+ entry_type = standardized_entry_type(dic[k])
+ group.add_argument("--" + k, default=dic[k], type=entry_type)
+ return parser
+
+
+def str_int_float(value):
+ """Type to convert strings to int, float (in this order) if possible.
+
+ Args:
+ value (str): Value to convert.
+
+ Returns:
+ int, float, str: Converted value.
+ """
+ if isint(value):
+ return int(value)
+ if isfloat(value):
+ return float(value)
+ elif isinstance(value, str):
+ return value
+
+
+def str2bool(value):
+ """Type to convert strings to Boolean (returns input if not boolean)"""
+ if not isinstance(value, str):
+ return value
+ if value.lower() in ("yes", "true", "y", "1"):
+ return True
+ elif value.lower() in ("no", "false", "n", "0"):
+ return False
+ else:
+ return value
+
+
+def str2bool_arg(value):
+ """Argparse type to convert strings to Boolean"""
+ value = str2bool(value)
+ if isinstance(value, bool):
+ return value
+ raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+def isfloat(value):
+ """Computes whether `value` can be cast to a float.
+
+ Args:
+ value (str): Value to check.
+
+ Returns:
+ bool: Whether `value` can be cast to a float.
+
+ """
+ try:
+ float(value)
+ return True
+ except ValueError:
+ return False
+
+
+def isint(value):
+ """Computes whether `value` can be cast to an int
+
+ Args:
+ value (str): Value to check.
+
+ Returns:
+ bool: Whether `value` can be cast to an int.
+
+ """
+ try:
+ int(value)
+ return True
+ except ValueError:
+ return False
+
+
+def parse_args_as_dict(parser, return_plain_args=False, args=None):
+ """Get a dict of dicts out of process `parser.parse_args()`
+
+ Top-level keys corresponding to groups and bottom-level keys corresponding
+ to arguments. Under `'main_args'`, the arguments which don't belong to a
+ argparse group (i.e main arguments defined before parsing from a dict) can
+ be found.
+
+ Args:
+ parser (argparse.ArgumentParser): ArgumentParser instance containing
+ groups. Output of `prepare_parser_from_dict`.
+ return_plain_args (bool): Whether to return the output or
+ `parser.parse_args()`.
+ args (list): List of arguments as read from the command line.
+ Used for unit testing.
+
+ Returns:
+ dict:
+ Dictionary of dictionaries containing the arguments. Optionally the
+ direct output `parser.parse_args()`.
+ """
+ args = parser.parse_args(args=args)
+ args_dic = {}
+ for group in parser._action_groups:
+ group_dict = {a.dest: getattr(args, a.dest, None) for a in group._group_actions}
+ args_dic[group.title] = group_dict
+ if sys.version_info.minor == 10:
+ args_dic["main_args"] = args_dic["positional arguments"]
+ del args_dic["positional arguments"]
+ else:
+ args_dic["main_args"] = args_dic["optional arguments"]
+ del args_dic["optional arguments"]
+ if return_plain_args:
+ return args_dic, args
+ return args_dic
diff --git a/look2hear/utils/separator.py b/look2hear/utils/separator.py
new file mode 100644
index 0000000000000000000000000000000000000000..f38077bad5ef49d7af77e9aff4bc1476730036af
--- /dev/null
+++ b/look2hear/utils/separator.py
@@ -0,0 +1,69 @@
+import os
+import warnings
+import torch
+import numpy as np
+import soundfile as sf
+
+
+def get_device(tensor_or_module, default=None):
+ if hasattr(tensor_or_module, "device"):
+ return tensor_or_module.device
+ elif hasattr(tensor_or_module, "parameters"):
+ return next(tensor_or_module.parameters()).device
+ elif default is None:
+ raise TypeError(
+ f"Don't know how to get device of {type(tensor_or_module)} object"
+ )
+ else:
+ return torch.device(default)
+
+
+class Separator:
+ def forward_wav(self, wav, **kwargs):
+ raise NotImplementedError
+
+ def sample_rate(self):
+ raise NotImplementedError
+
+
+def separate(model, wav, **kwargs):
+ if isinstance(wav, np.ndarray):
+ return numpy_separate(model, wav, **kwargs)
+ elif isinstance(wav, torch.Tensor):
+ return torch_separate(model, wav, **kwargs)
+ else:
+ raise ValueError(
+ f"Only support filenames, numpy arrays and torch tensors, received {type(wav)}"
+ )
+
+
+@torch.no_grad()
+def torch_separate(model: Separator, wav: torch.Tensor, **kwargs) -> torch.Tensor:
+ """Core logic of `separate`."""
+ if model.in_channels is not None and wav.shape[-2] != model.in_channels:
+ raise RuntimeError(
+ f"Model supports {model.in_channels}-channel inputs but found audio with {wav.shape[-2]} channels."
+ f"Please match the number of channels."
+ )
+ # Handle device placement
+ input_device = get_device(wav, default="cpu")
+ model_device = get_device(model, default="cpu")
+ wav = wav.to(model_device)
+ # Forward
+ separate_func = getattr(model, "forward_wav", model)
+ out_wavs = separate_func(wav, **kwargs)
+
+ # FIXME: for now this is the best we can do.
+ out_wavs *= wav.abs().sum() / (out_wavs.abs().sum())
+
+ # Back to input device (and numpy if necessary)
+ out_wavs = out_wavs.to(input_device)
+ return out_wavs
+
+
+def numpy_separate(model: Separator, wav: np.ndarray, **kwargs) -> np.ndarray:
+ """Numpy interface to `separate`."""
+ wav = torch.from_numpy(wav)
+ out_wavs = torch_separate(model, wav, **kwargs)
+ out_wavs = out_wavs.data.numpy()
+ return out_wavs
diff --git a/look2hear/utils/stft.py b/look2hear/utils/stft.py
new file mode 100644
index 0000000000000000000000000000000000000000..0bef618093aa219f59cbd967f4ed134d20b0837e
--- /dev/null
+++ b/look2hear/utils/stft.py
@@ -0,0 +1,797 @@
+# Copyright 2019 Jian Wu
+# License: Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+
+import math
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as tf
+import librosa.filters as filters
+
+from typing import Optional, Tuple
+from distutils.version import LooseVersion
+
+EPSILON = float(np.finfo(np.float32).eps)
+TORCH_VERSION = th.__version__
+
+if TORCH_VERSION >= LooseVersion("1.7"):
+ from torch.fft import fft as fft_func
+else:
+ pass
+
+
+def export_jit(transform: nn.Module) -> nn.Module:
+ """
+ Export transform module for inference
+ """
+ export_out = [module for module in transform if module.exportable()]
+ return nn.Sequential(*export_out)
+
+
+def init_window(wnd: str, frame_len: int, device: th.device = "cpu") -> th.Tensor:
+ """
+ Return window coefficient
+ Args:
+ wnd: window name
+ frame_len: length of the frame
+ """
+
+ def sqrthann(frame_len, periodic=True):
+ return th.hann_window(frame_len, periodic=periodic) ** 0.5
+
+ if wnd not in ["bartlett", "hann", "hamm", "blackman", "rect", "sqrthann"]:
+ raise RuntimeError(f"Unknown window type: {wnd}")
+
+ wnd_tpl = {
+ "sqrthann": sqrthann,
+ "hann": th.hann_window,
+ "hamm": th.hamming_window,
+ "blackman": th.blackman_window,
+ "bartlett": th.bartlett_window,
+ "rect": th.ones,
+ }
+ if wnd != "rect":
+ # match with librosa
+ c = wnd_tpl[wnd](frame_len, periodic=True)
+ else:
+ c = wnd_tpl[wnd](frame_len)
+ return c.to(device)
+
+
+def init_kernel(
+ frame_len: int,
+ frame_hop: int,
+ window: th.Tensor,
+ round_pow_of_two: bool = True,
+ normalized: bool = False,
+ inverse: bool = False,
+ mode: str = "librosa",
+) -> Tuple[th.Tensor, th.Tensor]:
+ """
+ Return STFT kernels
+ Args:
+ frame_len: length of the frame
+ frame_hop: hop size between frames
+ window: window tensor
+ round_pow_of_two: if true, choose round(#power_of_two) as the FFT size
+ normalized: return normalized DFT matrix
+ inverse: return iDFT matrix
+ mode: framing mode (librosa or kaldi)
+ """
+ if mode not in ["librosa", "kaldi"]:
+ raise ValueError(f"Unsupported mode: {mode}")
+ # FFT size: B
+ if round_pow_of_two or mode == "kaldi":
+ fft_size = 2 ** math.ceil(math.log2(frame_len))
+ else:
+ fft_size = frame_len
+ # center padding window if needed
+ if mode == "librosa" and fft_size != frame_len:
+ lpad = (fft_size - frame_len) // 2
+ window = tf.pad(window, (lpad, fft_size - frame_len - lpad))
+ if normalized:
+ # make K^H * K = I
+ S = fft_size ** 0.5
+ else:
+ S = 1
+ # W x B x 2
+ if TORCH_VERSION >= LooseVersion("1.7"):
+ K = fft_func(th.eye(fft_size) / S, dim=-1)
+ K = th.stack([K.real, K.imag], dim=-1)
+ else:
+ I = th.stack([th.eye(fft_size), th.zeros(fft_size, fft_size)], dim=-1)
+ K = th.fft(I / S, 1)
+ if mode == "kaldi":
+ K = K[:frame_len]
+ if inverse and not normalized:
+ # to make K^H * K = I
+ K = K / fft_size
+ # 2 x B x W
+ K = th.transpose(K, 0, 2)
+ # 2B x 1 x W
+ K = th.reshape(K, (fft_size * 2, 1, K.shape[-1]))
+ return K.to(window.device), window
+
+
+def mel_filter(
+ frame_len: int,
+ round_pow_of_two: bool = True,
+ num_bins: Optional[int] = None,
+ sr: int = 16000,
+ num_mels: int = 80,
+ fmin: float = 0.0,
+ fmax: Optional[float] = None,
+ norm: bool = False,
+) -> th.Tensor:
+ """
+ Return mel filter coefficients
+ Args:
+ frame_len: length of the frame
+ round_pow_of_two: if true, choose round(#power_of_two) as the FFT size
+ num_bins: number of the frequency bins produced by STFT
+ num_mels: number of the mel bands
+ fmin: lowest frequency (in Hz)
+ fmax: highest frequency (in Hz)
+ norm: normalize the mel filter coefficients
+ """
+ # FFT points
+ if num_bins is None:
+ N = 2 ** math.ceil(math.log2(frame_len)) if round_pow_of_two else frame_len
+ else:
+ N = (num_bins - 1) * 2
+ # fmin & fmax
+ freq_upper = sr // 2
+ if fmax is None:
+ fmax = freq_upper
+ else:
+ fmax = min(fmax + freq_upper if fmax < 0 else fmax, freq_upper)
+ fmin = max(0, fmin)
+ # mel filter coefficients
+ mel = filters.mel(
+ sr,
+ N,
+ n_mels=num_mels,
+ fmax=fmax,
+ fmin=fmin,
+ htk=True,
+ norm="slaney" if norm else None,
+ )
+ # num_mels x (N // 2 + 1)
+ return th.tensor(mel, dtype=th.float32)
+
+
+def speed_perturb_filter(
+ src_sr: int, dst_sr: int, cutoff_ratio: float = 0.95, num_zeros: int = 64
+) -> th.Tensor:
+ """
+ Return speed perturb filters, reference:
+ https://github.com/danpovey/filtering/blob/master/lilfilter/resampler.py
+ Args:
+ src_sr: sample rate of the source signal
+ dst_sr: sample rate of the target signal
+ Return:
+ weight (Tensor): coefficients of the filter
+ """
+ if src_sr == dst_sr:
+ raise ValueError(f"src_sr should not be equal to dst_sr: {src_sr}/{dst_sr}")
+ gcd = math.gcd(src_sr, dst_sr)
+ src_sr = src_sr // gcd
+ dst_sr = dst_sr // gcd
+ if src_sr == 1 or dst_sr == 1:
+ raise ValueError("do not support integer downsample/upsample")
+ zeros_per_block = min(src_sr, dst_sr) * cutoff_ratio
+ padding = 1 + int(num_zeros / zeros_per_block)
+ # dst_sr x src_sr x K
+ times = (
+ np.arange(dst_sr)[:, None, None] / float(dst_sr)
+ - np.arange(src_sr)[None, :, None] / float(src_sr)
+ - np.arange(2 * padding + 1)[None, None, :]
+ + padding
+ )
+ window = np.heaviside(1 - np.abs(times / padding), 0.0) * (
+ 0.5 + 0.5 * np.cos(times / padding * math.pi)
+ )
+ weight = np.sinc(times * zeros_per_block) * window * zeros_per_block / float(src_sr)
+ return th.tensor(weight, dtype=th.float32)
+
+
+def splice_feature(
+ feats: th.Tensor, lctx: int = 1, rctx: int = 1, op: str = "cat"
+) -> th.Tensor:
+ """
+ Splice feature
+ Args:
+ feats (Tensor): N x ... x T x F, original feature
+ lctx: left context
+ rctx: right context
+ op: operator on feature context
+ Return:
+ splice (Tensor): feature with context padded
+ """
+ if lctx + rctx == 0:
+ return feats
+ if op not in ["cat", "stack"]:
+ raise ValueError(f"Unknown op for feature splicing: {op}")
+ # [N x ... x T x F, ...]
+ ctx = []
+ T = feats.shape[-2]
+ for c in range(-lctx, rctx + 1):
+ idx = th.arange(c, c + T, device=feats.device, dtype=th.int64)
+ idx = th.clamp(idx, min=0, max=T - 1)
+ ctx.append(th.index_select(feats, -2, idx))
+ if op == "cat":
+ # N x ... x T x FD
+ splice = th.cat(ctx, -1)
+ else:
+ # N x ... x T x F x D
+ splice = th.stack(ctx, -1)
+ return splice
+
+
+def _forward_stft(
+ wav: th.Tensor,
+ kernel: th.Tensor,
+ window: th.Tensor,
+ return_polar: bool = False,
+ pre_emphasis: float = 0,
+ frame_hop: int = 256,
+ onesided: bool = False,
+ center: bool = False,
+ eps: float = EPSILON,
+) -> th.Tensor:
+ """
+ STFT function implemented by conv1d (not efficient, but we don't care during training)
+ Args:
+ wav (Tensor): N x (C) x S
+ kernel (Tensor): STFT transform kernels, from init_kernel(...)
+ return_polar: return [magnitude; phase] Tensor or [real; imag] Tensor
+ pre_emphasis: factor of preemphasis
+ frame_hop: frame hop size in number samples
+ onesided: return half FFT bins
+ center: if true, we assumed to have centered frames
+ Return:
+ transform (Tensor): STFT transform results
+ """
+ wav_dim = wav.dim()
+ if wav_dim not in [2, 3]:
+ raise RuntimeError(f"STFT expect 2D/3D tensor, but got {wav_dim:d}D")
+ # if N x S, reshape N x 1 x S
+ # else: reshape NC x 1 x S
+ N, S = wav.shape[0], wav.shape[-1]
+ wav = wav.view(-1, 1, S)
+ # NC x 1 x S+2P
+ if center:
+ pad = kernel.shape[-1] // 2
+ # NOTE: match with librosa
+ wav = tf.pad(wav, (pad, pad), mode="reflect")
+ # STFT
+ kernel = kernel * window
+ if pre_emphasis > 0:
+ # NC x W x T
+ frames = tf.unfold(
+ wav[:, None], (1, kernel.shape[-1]), stride=frame_hop, padding=0
+ )
+ # follow Kaldi's Preemphasize
+ frames[:, 1:] = frames[:, 1:] - pre_emphasis * frames[:, :-1]
+ frames[:, 0] *= 1 - pre_emphasis
+ # 1 x 2B x W, NC x W x T, NC x 2B x T
+ packed = th.matmul(kernel[:, 0][None, ...], frames)
+ else:
+ packed = tf.conv1d(wav, kernel, stride=frame_hop, padding=0)
+ # NC x 2B x T => N x C x 2B x T
+ if wav_dim == 3:
+ packed = packed.view(N, -1, packed.shape[-2], packed.shape[-1])
+ # N x (C) x B x T
+ real, imag = th.chunk(packed, 2, dim=-2)
+ # N x (C) x B/2+1 x T
+ if onesided:
+ num_bins = kernel.shape[0] // 4 + 1
+ real = real[..., :num_bins, :]
+ imag = imag[..., :num_bins, :]
+ if return_polar:
+ mag = (real ** 2 + imag ** 2 + eps) ** 0.5
+ pha = th.atan2(imag, real)
+ return th.stack([mag, pha], dim=-1)
+ else:
+ return th.stack([real, imag], dim=-1)
+
+
+def _inverse_stft(
+ transform: th.Tensor,
+ kernel: th.Tensor,
+ window: th.Tensor,
+ return_polar: bool = False,
+ frame_hop: int = 256,
+ onesided: bool = False,
+ center: bool = False,
+ eps: float = EPSILON,
+) -> th.Tensor:
+ """
+ iSTFT function implemented by conv1d
+ Args:
+ transform (Tensor): STFT transform results
+ kernel (Tensor): STFT transform kernels, from init_kernel(...)
+ return_polar (bool): keep same with the one in _forward_stft
+ frame_hop: frame hop size in number samples
+ onesided: return half FFT bins
+ center: used in _forward_stft
+ Return:
+ wav (Tensor), N x S
+ """
+ # (N) x F x T x 2
+ transform_dim = transform.dim()
+ # if F x T x 2, reshape 1 x F x T x 2
+ if transform_dim == 3:
+ transform = th.unsqueeze(transform, 0)
+ if transform_dim != 4:
+ raise RuntimeError(f"Expect 4D tensor, but got {transform_dim}D")
+
+ if return_polar:
+ real = transform[..., 0] * th.cos(transform[..., 1])
+ imag = transform[..., 0] * th.sin(transform[..., 1])
+ else:
+ real, imag = transform[..., 0], transform[..., 1]
+
+ if onesided:
+ # [self.num_bins - 2, ..., 1]
+ reverse = range(kernel.shape[0] // 4 - 1, 0, -1)
+ # extend matrix: N x B x T
+ real = th.cat([real, real[:, reverse]], 1)
+ imag = th.cat([imag, -imag[:, reverse]], 1)
+ # pack: N x 2B x T
+ packed = th.cat([real, imag], dim=1)
+ # N x 1 x T
+ wav = tf.conv_transpose1d(packed, kernel * window, stride=frame_hop, padding=0)
+ # normalized audio samples
+ # refer: https://github.com/pytorch/audio/blob/2ebbbf511fb1e6c47b59fd32ad7e66023fa0dff1/torchaudio/functional.py#L171
+ num_frames = packed.shape[-1]
+ win_length = window.shape[0]
+ # W x T
+ win = th.repeat_interleave(window[..., None] ** 2, num_frames, dim=-1)
+ # Do OLA on windows
+ # v1)
+ I = th.eye(win_length, device=win.device)[:, None]
+ denorm = tf.conv_transpose1d(win[None, ...], I, stride=frame_hop, padding=0)
+ # v2)
+ # num_samples = (num_frames - 1) * frame_hop + win_length
+ # denorm = tf.fold(win[None, ...], (num_samples, 1), (win_length, 1),
+ # stride=frame_hop)[..., 0]
+ if center:
+ pad = kernel.shape[-1] // 2
+ wav = wav[..., pad:-pad]
+ denorm = denorm[..., pad:-pad]
+ wav = wav / (denorm + eps)
+ # N x S
+ return wav.squeeze(1)
+
+
+def _pytorch_stft(
+ wav: th.Tensor,
+ frame_len: int,
+ frame_hop: int,
+ n_fft: int = 512,
+ return_polar: bool = False,
+ window: str = "sqrthann",
+ normalized: bool = False,
+ onesided: bool = True,
+ center: bool = False,
+ eps: float = EPSILON,
+) -> th.Tensor:
+ """
+ Wrapper of PyTorch STFT function
+ Args:
+ wav (Tensor): source audio signal
+ frame_len: length of the frame
+ frame_hop: hop size between frames
+ n_fft: number of the FFT size
+ return_polar: return the results in polar coordinate
+ window: window tensor
+ center: same definition with the parameter in librosa.stft
+ normalized: use normalized DFT kernel
+ onesided: output onesided STFT
+ Return:
+ transform (Tensor), STFT transform results
+ """
+ if TORCH_VERSION < LooseVersion("1.7"):
+ raise RuntimeError("Can not use this function as TORCH_VERSION < 1.7")
+ wav_dim = wav.dim()
+ if wav_dim not in [2, 3]:
+ raise RuntimeError(f"STFT expect 2D/3D tensor, but got {wav_dim:d}D")
+ # if N x C x S, reshape NC x S
+ wav = wav.view(-1, wav.shape[-1])
+ # STFT: N x F x T x 2
+ stft = th.stft(
+ wav,
+ n_fft,
+ hop_length=frame_hop,
+ win_length=window.shape[-1],
+ window=window,
+ center=center,
+ normalized=normalized,
+ onesided=onesided,
+ return_complex=False,
+ )
+ if wav_dim == 3:
+ N, F, T, _ = stft.shape
+ stft = stft.view(N, -1, F, T, 2)
+ # N x (C) x F x T x 2
+ if not return_polar:
+ return stft
+ # N x (C) x F x T
+ real, imag = stft[..., 0], stft[..., 1]
+ mag = (real ** 2 + imag ** 2 + eps) ** 0.5
+ pha = th.atan2(imag, real)
+ return th.stack([mag, pha], dim=-1)
+
+
+def _pytorch_istft(
+ transform: th.Tensor,
+ frame_len: int,
+ frame_hop: int,
+ window: th.Tensor,
+ n_fft: int = 512,
+ return_polar: bool = False,
+ normalized: bool = False,
+ onesided: bool = True,
+ center: bool = False,
+ eps: float = EPSILON,
+) -> th.Tensor:
+ """
+ Wrapper of PyTorch iSTFT function
+ Args:
+ transform (Tensor): results of STFT
+ frame_len: length of the frame
+ frame_hop: hop size between frames
+ window: window tensor
+ n_fft: number of the FFT size
+ return_polar: keep same with _pytorch_stft
+ center: same definition with the parameter in librosa.stft
+ normalized: use normalized DFT kernel
+ onesided: output onesided STFT
+ Return:
+ wav (Tensor): synthetic audio
+ """
+ if TORCH_VERSION < LooseVersion("1.7"):
+ raise RuntimeError("Can not use this function as TORCH_VERSION < 1.7")
+
+ transform_dim = transform.dim()
+ # if F x T x 2, reshape 1 x F x T x 2
+ if transform_dim == 3:
+ transform = th.unsqueeze(transform, 0)
+ if transform_dim != 4:
+ raise RuntimeError(f"Expect 4D tensor, but got {transform_dim}D")
+
+ if return_polar:
+ real = transform[..., 0] * th.cos(transform[..., 1])
+ imag = transform[..., 0] * th.sin(transform[..., 1])
+ transform = th.stack([real, imag], -1)
+ # stft is a complex tensor of PyTorch
+ stft = th.view_as_complex(transform)
+ # (N) x S
+ wav = th.istft(
+ stft,
+ n_fft,
+ hop_length=frame_hop,
+ win_length=window.shape[-1],
+ window=window,
+ center=center,
+ normalized=normalized,
+ onesided=onesided,
+ return_complex=False,
+ )
+ return wav
+
+
+def forward_stft(
+ wav: th.Tensor,
+ frame_len: int,
+ frame_hop: int,
+ window: str = "sqrthann",
+ round_pow_of_two: bool = True,
+ return_polar: bool = False,
+ pre_emphasis: float = 0,
+ normalized: bool = False,
+ onesided: bool = True,
+ center: bool = False,
+ mode: str = "librosa",
+ eps: float = EPSILON,
+) -> th.Tensor:
+ """
+ STFT function implementation, equals to STFT layer
+ Args:
+ wav: source audio signal
+ frame_len: length of the frame
+ frame_hop: hop size between frames
+ return_polar: return [magnitude; phase] Tensor or [real; imag] Tensor
+ window: window name
+ center: center flag (similar with that in librosa.stft)
+ round_pow_of_two: if true, choose round(#power_of_two) as the FFT size
+ pre_emphasis: factor of preemphasis
+ normalized: use normalized DFT kernel
+ onesided: output onesided STFT
+ inverse: using iDFT kernel (for iSTFT)
+ mode: STFT mode, "kaldi" or "librosa" or "torch"
+ Return:
+ transform: results of STFT
+ """
+ window = init_window(window, frame_len, device=wav.device)
+ if mode == "torch":
+ n_fft = 2 ** math.ceil(math.log2(frame_len)) if round_pow_of_two else frame_len
+ return _pytorch_stft(
+ wav,
+ frame_len,
+ frame_hop,
+ n_fft=n_fft,
+ return_polar=return_polar,
+ window=window,
+ normalized=normalized,
+ onesided=onesided,
+ center=center,
+ eps=eps,
+ )
+ else:
+ kernel, window = init_kernel(
+ frame_len,
+ frame_hop,
+ window=window,
+ round_pow_of_two=round_pow_of_two,
+ normalized=normalized,
+ inverse=False,
+ mode=mode,
+ )
+ return _forward_stft(
+ wav,
+ kernel,
+ window,
+ return_polar=return_polar,
+ frame_hop=frame_hop,
+ pre_emphasis=pre_emphasis,
+ onesided=onesided,
+ center=center,
+ eps=eps,
+ )
+
+
+def inverse_stft(
+ transform: th.Tensor,
+ frame_len: int,
+ frame_hop: int,
+ return_polar: bool = False,
+ window: str = "sqrthann",
+ round_pow_of_two: bool = True,
+ normalized: bool = False,
+ onesided: bool = True,
+ center: bool = False,
+ mode: str = "librosa",
+ eps: float = EPSILON,
+) -> th.Tensor:
+ """
+ iSTFT function implementation, equals to iSTFT layer
+ Args:
+ transform: results of STFT
+ frame_len: length of the frame
+ frame_hop: hop size between frames
+ return_polar: keep same with function forward_stft(...)
+ window: window name
+ center: center flag (similar with that in librosa.stft)
+ round_pow_of_two: if true, choose round(#power_of_two) as the FFT size
+ normalized: use normalized DFT kernel
+ onesided: output onesided STFT
+ mode: STFT mode, "kaldi" or "librosa" or "torch"
+ Return:
+ wav: synthetic signals
+ """
+ window = init_window(window, frame_len, device=transform.device)
+ if mode == "torch":
+ n_fft = 2 ** math.ceil(math.log2(frame_len)) if round_pow_of_two else frame_len
+ return _pytorch_istft(
+ transform,
+ frame_len,
+ frame_hop,
+ n_fft=n_fft,
+ return_polar=return_polar,
+ window=window,
+ normalized=normalized,
+ onesided=onesided,
+ center=center,
+ eps=eps,
+ )
+ else:
+ kernel, window = init_kernel(
+ frame_len,
+ frame_hop,
+ window,
+ round_pow_of_two=round_pow_of_two,
+ normalized=normalized,
+ inverse=True,
+ mode=mode,
+ )
+ return _inverse_stft(
+ transform,
+ kernel,
+ window,
+ return_polar=return_polar,
+ frame_hop=frame_hop,
+ onesided=onesided,
+ center=center,
+ eps=eps,
+ )
+
+
+class STFTBase(nn.Module):
+ """
+ Base layer for (i)STFT
+ Args:
+ frame_len: length of the frame
+ frame_hop: hop size between frames
+ window: window name
+ center: center flag (similar with that in librosa.stft)
+ round_pow_of_two: if true, choose round(#power_of_two) as the FFT size
+ normalized: use normalized DFT kernel
+ pre_emphasis: factor of preemphasis
+ mode: STFT mode, "kaldi" or "librosa" or "torch"
+ onesided: output onesided STFT
+ inverse: using iDFT kernel (for iSTFT)
+ """
+
+ def __init__(
+ self,
+ frame_len: int,
+ frame_hop: int,
+ window: str = "sqrthann",
+ round_pow_of_two: bool = True,
+ normalized: bool = False,
+ pre_emphasis: float = 0,
+ onesided: bool = True,
+ inverse: bool = False,
+ center: bool = False,
+ mode: str = "librosa",
+ ) -> None:
+ super(STFTBase, self).__init__()
+ if mode != "torch":
+ K, w = init_kernel(
+ frame_len,
+ frame_hop,
+ init_window(window, frame_len),
+ round_pow_of_two=round_pow_of_two,
+ normalized=normalized,
+ inverse=inverse,
+ mode=mode,
+ )
+ self.K = nn.Parameter(K, requires_grad=False)
+ self.w = nn.Parameter(w, requires_grad=False)
+ self.num_bins = self.K.shape[0] // 4 + 1
+ self.pre_emphasis = pre_emphasis
+ self.win_length = self.K.shape[2]
+ else:
+ self.K = None
+ w = init_window(window, frame_len)
+ self.w = nn.Parameter(w, requires_grad=False)
+ fft_size = (
+ 2 ** math.ceil(math.log2(frame_len)) if round_pow_of_two else frame_len
+ )
+ self.num_bins = fft_size // 2 + 1
+ self.pre_emphasis = 0
+ self.win_length = fft_size
+ self.frame_len = frame_len
+ self.frame_hop = frame_hop
+ self.window = window
+ self.normalized = normalized
+ self.onesided = onesided
+ self.center = center
+ self.mode = mode
+
+ def num_frames(self, wav_len: th.Tensor) -> th.Tensor:
+ """
+ Compute number of the frames
+ """
+ assert th.sum(wav_len <= self.win_length) == 0
+ if self.center:
+ wav_len += self.win_length
+ return (
+ th.div(wav_len - self.win_length, self.frame_hop, rounding_mode="trunc") + 1
+ )
+
+ def extra_repr(self) -> str:
+ str_repr = (
+ f"num_bins={self.num_bins}, win_length={self.win_length}, "
+ + f"stride={self.frame_hop}, window={self.window}, "
+ + f"center={self.center}, mode={self.mode}"
+ )
+ if not self.onesided:
+ str_repr += f", onesided={self.onesided}"
+ if self.pre_emphasis > 0:
+ str_repr += f", pre_emphasis={self.pre_emphasis}"
+ if self.normalized:
+ str_repr += f", normalized={self.normalized}"
+ return str_repr
+
+
+class STFT(STFTBase):
+ """
+ Short-time Fourier Transform as a Layer
+ """
+
+ def __init__(self, *args, **kwargs):
+ super(STFT, self).__init__(*args, inverse=False, **kwargs)
+
+ def forward(
+ self, wav: th.Tensor, return_polar: bool = False, eps: float = EPSILON
+ ) -> th.Tensor:
+ """
+ Accept (single or multiple channel) raw waveform and output magnitude and phase
+ Args
+ wav (Tensor) input signal, N x (C) x S
+ Return
+ transform (Tensor), N x (C) x F x T x 2
+ """
+ if self.mode == "torch":
+ return _pytorch_stft(
+ wav,
+ self.frame_len,
+ self.frame_hop,
+ n_fft=(self.num_bins - 1) * 2,
+ return_polar=return_polar,
+ window=self.w,
+ normalized=self.normalized,
+ onesided=self.onesided,
+ center=self.center,
+ eps=eps,
+ )
+ else:
+ return _forward_stft(
+ wav,
+ self.K,
+ self.w,
+ return_polar=return_polar,
+ frame_hop=self.frame_hop,
+ pre_emphasis=self.pre_emphasis,
+ onesided=self.onesided,
+ center=self.center,
+ eps=eps,
+ )
+
+
+class iSTFT(STFTBase):
+ """
+ Inverse Short-time Fourier Transform as a Layer
+ """
+
+ def __init__(self, *args, **kwargs):
+ super(iSTFT, self).__init__(*args, inverse=True, **kwargs)
+
+ def forward(
+ self, transform: th.Tensor, return_polar: bool = False, eps: float = EPSILON
+ ) -> th.Tensor:
+ """
+ Accept phase & magnitude and output raw waveform
+ Args
+ transform (Tensor): STFT output, N x F x T x 2
+ Return
+ s (Tensor): N x S
+ """
+ if self.mode == "torch":
+ return _pytorch_istft(
+ transform,
+ self.frame_len,
+ self.frame_hop,
+ n_fft=(self.num_bins - 1) * 2,
+ return_polar=return_polar,
+ window=self.w,
+ normalized=self.normalized,
+ onesided=self.onesided,
+ center=self.center,
+ eps=eps,
+ )
+ else:
+ return _inverse_stft(
+ transform,
+ self.K,
+ self.w,
+ return_polar=return_polar,
+ frame_hop=self.frame_hop,
+ onesided=self.onesided,
+ center=self.center,
+ eps=eps,
+ )
diff --git a/look2hear/utils/torch_utils.py b/look2hear/utils/torch_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c3ec70df9e7b6c89f9beb4f5aa1cb8f9ac4cf90
--- /dev/null
+++ b/look2hear/utils/torch_utils.py
@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+
+
+def pad_x_to_y(x, y, axis: int = -1):
+ if axis != -1:
+ raise NotImplementedError
+ inp_len = y.shape[axis]
+ output_len = x.shape[axis]
+ return nn.functional.pad(x, [0, inp_len - output_len])
+
+
+def shape_reconstructed(reconstructed, size):
+ if len(size) == 1:
+ return reconstructed.squeeze(0)
+ return reconstructed
+
+
+def tensors_to_device(tensors, device):
+ """Transfer tensor, dict or list of tensors to device.
+
+ Args:
+ tensors (:class:`torch.Tensor`): May be a single, a list or a
+ dictionary of tensors.
+ device (:class: `torch.device`): the device where to place the tensors.
+
+ Returns:
+ Union [:class:`torch.Tensor`, list, tuple, dict]:
+ Same as input but transferred to device.
+ Goes through lists and dicts and transfers the torch.Tensor to
+ device. Leaves the rest untouched.
+ """
+ if isinstance(tensors, torch.Tensor):
+ return tensors.to(device)
+ elif isinstance(tensors, (list, tuple)):
+ return [tensors_to_device(tens, device) for tens in tensors]
+ elif isinstance(tensors, dict):
+ for key in tensors.keys():
+ tensors[key] = tensors_to_device(tensors[key], device)
+ return tensors
+ else:
+ return tensors
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..3bef70d53061bf077728444136b9e52697ec2ab9
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,114 @@
+accelerated-scan==0.2.0
+aiohappyeyeballs==2.4.3
+aiohttp==3.11.7
+aiosignal==1.3.1
+antlr4-python3-runtime==4.9.3
+attrs==24.2.0
+audioread==3.0.1
+av==13.1.0
+blessed==1.20.0
+blessings==1.7
+certifi==2024.8.30
+cffi==1.17.1
+charset-normalizer==3.4.0
+click==8.1.7
+coloredlogs==15.0.1
+conda-pack==0.8.1
+ctranslate2==4.5.0
+datasets==3.2.0
+decorator==5.1.1
+dill==0.3.8
+docker-pycreds==0.4.0
+einops==0.8.0
+fast_bss_eval==0.1.4
+filelock==3.16.1
+flatbuffers==24.3.25
+frozenlist==1.5.0
+fsspec==2024.9.0
+gitdb==4.0.11
+GitPython==3.1.43
+gpustat==1.1.1
+huggingface-hub==0.26.2
+humanfriendly==10.0
+hydra-core==1.3.2
+HyperPyYAML==1.2.2
+idna==3.10
+Jinja2==3.1.4
+joblib==1.4.2
+lazy_loader==0.4
+librosa==0.10.2.post1
+lightning-utilities==0.11.9
+llvmlite==0.43.0
+markdown-it-py==3.0.0
+MarkupSafe==3.0.2
+mdurl==0.1.2
+modelscope==1.20.1
+mpmath==1.3.0
+msgpack==1.1.0
+multidict==6.1.0
+multiprocess==0.70.16
+networkx==3.4.2
+ninja==1.11.1.3
+numba==0.60.0
+numpy==2.0.2
+omegaconf==2.3.0
+onnxruntime==1.20.1
+packaging==24.2
+pandas==2.2.3
+pesq==0.0.4
+pillow==11.0.0
+pip-system-certs==4.0
+platformdirs==4.3.6
+pooch==1.8.2
+propcache==0.2.0
+protobuf==5.28.3
+psutil==6.1.0
+ptflops==0.7.4
+pyarrow==18.1.0
+pycparser==2.22
+Pygments==2.18.0
+pystoi==0.4.1
+python-dateutil==2.9.0.post0
+pytorch-lightning==2.0.2
+pytorch-ranger==0.1.1
+pytz==2024.2
+PyYAML==6.0.2
+regex==2024.11.6
+requests==2.32.3
+rich==13.9.4
+rotary-embedding-torch==0.8.5
+ruamel.yaml==0.18.6
+ruamel.yaml.clib==0.2.12
+safetensors==0.4.5
+scikit-learn==1.5.2
+scipy==1.14.1
+sentencepiece==0.2.0
+sentry-sdk==2.19.0
+setproctitle==1.3.4
+six==1.16.0
+smmap==5.0.1
+soundfile==0.12.1
+soxr==0.5.0.post1
+speechbrain==0.5.14
+sympy==1.13.1
+threadpoolctl==3.5.0
+tokenizers==0.21.0
+torch==2.5.1
+torch-complex==0.4.4
+torch-mir-eval==0.4
+torch-optimizer==0.3.0
+torchaudio==2.5.1
+torchmetrics==1.6.0
+torchvision==0.20.1
+tqdm==4.67.0
+transformers==4.47.1
+triton==3.1.0
+typeguard==2.13.3
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+wandb==0.18.7
+wcwidth==0.2.13
+wrapt==1.17.0
+xxhash==3.5.0
+yarl==1.18.0
diff --git a/test/mix.wav b/test/mix.wav
new file mode 100644
index 0000000000000000000000000000000000000000..eb063704154aff42d7a52bc909f5a04ac315ac3c
--- /dev/null
+++ b/test/mix.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96096869c977120503b2ef4b53429a4bce510b344d6a274c3371457be854f38e
+size 192044
diff --git a/test/s1.wav b/test/s1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..50ceb11a9e40dd8af0371e21f0d4c9e597bfaf62
--- /dev/null
+++ b/test/s1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2bc86a053f8fd42703f5a402d6c10d71b78121854993896ae2f204646b251cd4
+size 384080
diff --git a/test/s2.wav b/test/s2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..f24d83d5cb77a350720303969356274f9a37b17a
--- /dev/null
+++ b/test/s2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a93d50682bd48181e0fc96f6bda41e923f8426402a29fb3f53f33919af6c381f
+size 384080
diff --git a/test/spk1.wav b/test/spk1.wav
new file mode 100644
index 0000000000000000000000000000000000000000..8163478833cf217a263577c5840d4d1031a459af
--- /dev/null
+++ b/test/spk1.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19d94b8621ca6111b1f55fd663d357f202d50feeb4d480435e9285bcc0dcf5f7
+size 384080
diff --git a/test/spk2.wav b/test/spk2.wav
new file mode 100644
index 0000000000000000000000000000000000000000..a2b5258a69e599754aa0a72775e72a806fb5e011
--- /dev/null
+++ b/test/spk2.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aba7da1920c8b736385378ea8e8b4b69e803eb44be353703e4babe60f600c4f0
+size 384080
diff --git a/test/test_mixture_466.wav b/test/test_mixture_466.wav
new file mode 100644
index 0000000000000000000000000000000000000000..c4b4e6f731bbacb07363640c5f26e2a63de3aeae
--- /dev/null
+++ b/test/test_mixture_466.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27d8a4c46ea934c264af5ec04f8db4f0a3b5d59481d1071b1aa8086ea3a230fd
+size 5292044
diff --git a/test/test_target_dialog_466.wav b/test/test_target_dialog_466.wav
new file mode 100644
index 0000000000000000000000000000000000000000..2ec68e47368662a6ff4aa78391ac3ae89af6a930
--- /dev/null
+++ b/test/test_target_dialog_466.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca08ef210dcdd4e705789cce4a5de300dbc9cee0c5498a810506cb7a4fdbcf74
+size 10584080
diff --git a/test/test_target_effect_466.wav b/test/test_target_effect_466.wav
new file mode 100644
index 0000000000000000000000000000000000000000..dd8969ae932de2aaab18d4b72ef2c87cf911b3c1
--- /dev/null
+++ b/test/test_target_effect_466.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2ef06b888f17ec1e13562a2a03d419ba7e955f2ba0634b60f7cf1eabca46ba0
+size 10584080
diff --git a/test/test_target_music_466.wav b/test/test_target_music_466.wav
new file mode 100644
index 0000000000000000000000000000000000000000..7867202d9971f917e1c523c4ec445e7632bf5df2
--- /dev/null
+++ b/test/test_target_music_466.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cdb26e3ab09a25694d8654f047fde930a9a6e04a0746f4a42e628cf80207479f
+size 10584080
+
+
+
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
