From 0c3fdccf2f271fb7c44f6ea6e9f4ee234795f2c5 Mon Sep 17 00:00:00 2001 From: Matthijs Hollemans Date: Wed, 14 Jun 2023 18:57:23 +0200 Subject: [PATCH] [WIP] add EnCodec model (#23655) * boilerplate stuff * messing around with the feature extractor * fix feature extractor * unit tests for feature extractor * rename speech to audio * quick-and-dirty import of Meta's code * import weights (sort of) * cleaning up * more cleaning up * move encoder/decoder args into config * cleanup model * rename EnCodec -> Encodec * RVQ parameters in config * add slow test * add lstm init and test_init * Add save & load * finish EncodecModel * remove decoder_input_values as they are ont used anywhere (not removed from doc yet) * fix test feature extraction model name * Add better slow test * Fix tests * some fixup and cleaning * Improve further * cleaning up quantizer * fix up conversion script * test don't pass, _encode_fram does not work * update tests with output per encode and decode * more cleanup * rename _codebook * remove old config cruft * ratios & hop_length * use ModuleList instead of Sequential * clean up resnet block * update types * update tests * fixup * quick cleanup * fix padding * more styl,ing * add patrick feedback * fix copies * fixup * fix lstm * fix shape issues * fixup * rename conv layers * fixup * fix decoding * small conv refactoring * remove norm_params * simplify conv layers * rename conv layers * stuff * Clean up * Add padding logic use padding mask small conv refactoring remove norm_params simplify conv layers rename conv layers stuff add batched test update Clean up merge and update for padding fix padding fixup * clean up more * clean up more * More clean ups * cleanup convolutions * typo * fix typos * fixup * build PR doc? * start refactoring docstring * fix don't pad when no strid and chunk * update docstring * update docstring * nits * update going to lunch * update config and model * fix broken testse (becaue of the config changes) * fix scale computation * fixu[ * only return dict if speciefied or if config returns it * remove todos * update defaults in config * update conversion script * fix doctest * more docstring + fixup * nits on batched_tests * more nits * Apply suggestions from code review Co-authored-by: Patrick von Platen * update basxed on review * fix update * updaet tests * Apply suggestions from code review Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> * fixup * add overlap and chunl_length_s * cleanup feature extraction * teste edge cases truncation and padding * correct processor values * update config encodec, nits * fix tests * fixup * fix 24Hz test * elle tests are green * fix fixup * Apply suggestions from code review * revert readme changes * fixup * add example * use facebook checkpoints * fix typo * no pipeline tests * use slef.pad everywhere we can * Apply suggestions from code review Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * update based on review * update * update mdx * fix bug and tests * fixup * fix doctest * remove comment * more nits * add more coverage for `test_truncation_and_padding` * fixup * add last test * fix text * nits * Update tests/models/encodec/test_modeling_encodec.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * take care of the last comments * typo * fix test * nits * fixup * Update src/transformers/models/encodec/feature_extraction_encodec.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> --------- Co-authored-by: Patrick von Platen Co-authored-by: arthur.zucker@gmail.com Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> --- README.md | 1 + README_es.md | 1 + README_hd.md | 1 + README_ja.md | 1 + README_ko.md | 1 + README_zh-hans.md | 1 + README_zh-hant.md | 1 + docs/source/en/_toctree.yml | 2 + docs/source/en/index.mdx | 2 + docs/source/en/model_doc/encodec.mdx | 59 ++ src/transformers/__init__.py | 22 + src/transformers/models/__init__.py | 1 + .../models/auto/configuration_auto.py | 3 + .../models/auto/feature_extraction_auto.py | 1 + src/transformers/models/auto/modeling_auto.py | 1 + src/transformers/models/encodec/__init__.py | 65 ++ .../models/encodec/configuration_encodec.py | 189 ++++ .../convert_encodec_checkpoint_to_pytorch.py | 352 ++++++++ .../encodec/feature_extraction_encodec.py | 206 +++++ .../models/encodec/modeling_encodec.py | 808 ++++++++++++++++++ src/transformers/utils/dummy_pt_objects.py | 17 + .../utils/dummy_speech_objects.py | 7 + tests/models/encodec/__init__.py | 0 .../test_feature_extraction_encodec.py | 255 ++++++ tests/models/encodec/test_modeling_encodec.py | 571 +++++++++++++ utils/check_config_attributes.py | 2 + utils/documentation_tests.txt | 2 + 27 files changed, 2572 insertions(+) create mode 100644 docs/source/en/model_doc/encodec.mdx create mode 100644 src/transformers/models/encodec/__init__.py create mode 100644 src/transformers/models/encodec/configuration_encodec.py create mode 100644 src/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py create mode 100644 src/transformers/models/encodec/feature_extraction_encodec.py create mode 100644 src/transformers/models/encodec/modeling_encodec.py create mode 100644 tests/models/encodec/__init__.py create mode 100644 tests/models/encodec/test_feature_extraction_encodec.py create mode 100644 tests/models/encodec/test_modeling_encodec.py diff --git a/README.md b/README.md index b5ce283be4..7d6ea27878 100644 --- a/README.md +++ b/README.md @@ -346,6 +346,7 @@ Current number of checkpoints: ![](https://img.shields.io/endpoint?url=https://h 1. **[EfficientFormer](https://huggingface.co/docs/transformers/model_doc/efficientformer)** (from Snap Research) released with the paper [EfficientFormer: Vision Transformers at MobileNetSpeed](https://arxiv.org/abs/2206.01191) by Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, Jian Ren. 1. **[EfficientNet](https://huggingface.co/docs/transformers/model_doc/efficientnet)** (from Google Brain) released with the paper [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) by Mingxing Tan, Quoc V. Le. 1. **[ELECTRA](https://huggingface.co/docs/transformers/model_doc/electra)** (from Google Research/Stanford University) released with the paper [ELECTRA: Pre-training text encoders as discriminators rather than generators](https://arxiv.org/abs/2003.10555) by Kevin Clark, Minh-Thang Luong, Quoc V. Le, Christopher D. Manning. +1. **[EnCodec](https://huggingface.co/docs/transformers/main/model_doc/encodec)** (from Meta AI) released with the paper [High Fidelity Neural Audio Compression](https://arxiv.org/abs/2210.13438) by Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi. 1. **[EncoderDecoder](https://huggingface.co/docs/transformers/model_doc/encoder-decoder)** (from Google Research) released with the paper [Leveraging Pre-trained Checkpoints for Sequence Generation Tasks](https://arxiv.org/abs/1907.12461) by Sascha Rothe, Shashi Narayan, Aliaksei Severyn. 1. **[ERNIE](https://huggingface.co/docs/transformers/model_doc/ernie)** (from Baidu) released with the paper [ERNIE: Enhanced Representation through Knowledge Integration](https://arxiv.org/abs/1904.09223) by Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Xuyi Chen, Han Zhang, Xin Tian, Danxiang Zhu, Hao Tian, Hua Wu. 1. **[ErnieM](https://huggingface.co/docs/transformers/model_doc/ernie_m)** (from Baidu) released with the paper [ERNIE-M: Enhanced Multilingual Representation by Aligning Cross-lingual Semantics with Monolingual Corpora](https://arxiv.org/abs/2012.15674) by Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang. diff --git a/README_es.md b/README_es.md index 32f955dbda..76270b2c7c 100644 --- a/README_es.md +++ b/README_es.md @@ -321,6 +321,7 @@ Número actual de puntos de control: ![](https://img.shields.io/endpoint?url=htt 1. **[EfficientFormer](https://huggingface.co/docs/transformers/model_doc/efficientformer)** (from Snap Research) released with the paper [EfficientFormer: Vision Transformers at MobileNetSpeed](https://arxiv.org/abs/2206.01191) by Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, Jian Ren. 1. **[EfficientNet](https://huggingface.co/docs/transformers/model_doc/efficientnet)** (from Google Brain) released with the paper [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) by Mingxing Tan, Quoc V. Le. 1. **[ELECTRA](https://huggingface.co/docs/transformers/model_doc/electra)** (from Google Research/Stanford University) released with the paper [ELECTRA: Pre-training text encoders as discriminators rather than generators](https://arxiv.org/abs/2003.10555) by Kevin Clark, Minh-Thang Luong, Quoc V. Le, Christopher D. Manning. +1. **[EnCodec](https://huggingface.co/docs/transformers/main/model_doc/encodec)** (from Meta AI) released with the paper [High Fidelity Neural Audio Compression](https://arxiv.org/abs/2210.13438) by Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi. 1. **[EncoderDecoder](https://huggingface.co/docs/transformers/model_doc/encoder-decoder)** (from Google Research) released with the paper [Leveraging Pre-trained Checkpoints for Sequence Generation Tasks](https://arxiv.org/abs/1907.12461) by Sascha Rothe, Shashi Narayan, Aliaksei Severyn. 1. **[ERNIE](https://huggingface.co/docs/transformers/model_doc/ernie)** (from Baidu) released with the paper [ERNIE: Enhanced Representation through Knowledge Integration](https://arxiv.org/abs/1904.09223) by Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Xuyi Chen, Han Zhang, Xin Tian, Danxiang Zhu, Hao Tian, Hua Wu. 1. **[ErnieM](https://huggingface.co/docs/transformers/model_doc/ernie_m)** (from Baidu) released with the paper [ERNIE-M: Enhanced Multilingual Representation by Aligning Cross-lingual Semantics with Monolingual Corpora](https://arxiv.org/abs/2012.15674) by Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang. diff --git a/README_hd.md b/README_hd.md index 2195e7cc63..02fdadf751 100644 --- a/README_hd.md +++ b/README_hd.md @@ -293,6 +293,7 @@ conda install -c huggingface transformers 1. **[EfficientFormer](https://huggingface.co/docs/transformers/model_doc/efficientformer)** (from Snap Research) released with the paper [EfficientFormer: Vision Transformers at MobileNetSpeed](https://arxiv.org/abs/2206.01191) by Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, Jian Ren. 1. **[EfficientNet](https://huggingface.co/docs/transformers/model_doc/efficientnet)** (from Google Brain) released with the paper [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) by Mingxing Tan, Quoc V. Le. 1. **[ELECTRA](https://huggingface.co/docs/transformers/model_doc/electra)** (Google रिसर्च/स्टैनफोर्ड यूनिवर्सिटी से) साथ में दिया गया पेपर [इलेक्ट्रा: जेनरेटर के बजाय भेदभाव करने वाले के रूप में टेक्स्ट एन्कोडर्स का पूर्व-प्रशिक्षण] (https://arxiv.org/abs/2003.10555) केविन क्लार्क, मिन्ह-थांग लुओंग, क्वोक वी. ले, क्रिस्टोफर डी. मैनिंग द्वारा पोस्ट किया गया। +1. **[EnCodec](https://huggingface.co/docs/transformers/main/model_doc/encodec)** (Meta AI से) Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi. द्वाराअनुसंधान पत्र [High Fidelity Neural Audio Compression](https://arxiv.org/abs/2210.13438) के साथ जारी किया गया 1. **[EncoderDecoder](https://huggingface.co/docs/transformers/model_doc/encoder-decoder)** (Google रिसर्च से) साथ में दिया गया पेपर [सीक्वेंस जेनरेशन टास्क के लिए प्री-ट्रेंड चेकपॉइंट का इस्तेमाल करना](https:/ /arxiv.org/abs/1907.12461) साशा रोठे, शशि नारायण, अलियाक्सि सेवेरिन द्वारा। 1. **[ERNIE](https://huggingface.co/docs/transformers/model_doc/ernie)**(Baidu से) साथ देने वाला पेपर [ERNIE: एन्हांस्ड रिप्रेजेंटेशन थ्रू नॉलेज इंटीग्रेशन](https://arxiv.org/abs/1904.09223) यू सन, शुओहुआन वांग, युकुन ली, शिकुन फेंग, ज़ुई चेन, हान झांग, शिन तियान, डैनक्सियांग झू, हाओ तियान, हुआ वू द्वारा पोस्ट किया गया। 1. **[ErnieM](https://huggingface.co/docs/transformers/model_doc/ernie_m)** (Baidu से) Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang. द्वाराअनुसंधान पत्र [ERNIE-M: Enhanced Multilingual Representation by Aligning Cross-lingual Semantics with Monolingual Corpora](https://arxiv.org/abs/2012.15674) के साथ जारी किया गया diff --git a/README_ja.md b/README_ja.md index 2331bb06cf..2116fee408 100644 --- a/README_ja.md +++ b/README_ja.md @@ -355,6 +355,7 @@ Flax、PyTorch、TensorFlowをcondaでインストールする方法は、それ 1. **[EfficientFormer](https://huggingface.co/docs/transformers/model_doc/efficientformer)** (Snap Research から) Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, Jian Ren. から公開された研究論文 [EfficientFormer: Vision Transformers at MobileNetSpeed](https://arxiv.org/abs/2206.01191) 1. **[EfficientNet](https://huggingface.co/docs/transformers/model_doc/efficientnet)** (from Google Brain) released with the paper [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) by Mingxing Tan, Quoc V. Le. 1. **[ELECTRA](https://huggingface.co/docs/transformers/model_doc/electra)** (Google Research/Stanford University から) Kevin Clark, Minh-Thang Luong, Quoc V. Le, Christopher D. Manning から公開された研究論文: [ELECTRA: Pre-training text encoders as discriminators rather than generators](https://arxiv.org/abs/2003.10555) +1. **[EnCodec](https://huggingface.co/docs/transformers/main/model_doc/encodec)** (Meta AI から) Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi. から公開された研究論文 [High Fidelity Neural Audio Compression](https://arxiv.org/abs/2210.13438) 1. **[EncoderDecoder](https://huggingface.co/docs/transformers/model_doc/encoder-decoder)** (Google Research から) Sascha Rothe, Shashi Narayan, Aliaksei Severyn から公開された研究論文: [Leveraging Pre-trained Checkpoints for Sequence Generation Tasks](https://arxiv.org/abs/1907.12461) 1. **[ERNIE](https://huggingface.co/docs/transformers/model_doc/ernie)** (Baidu から) Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Xuyi Chen, Han Zhang, Xin Tian, Danxiang Zhu, Hao Tian, Hua Wu から公開された研究論文: [ERNIE: Enhanced Representation through Knowledge Integration](https://arxiv.org/abs/1904.09223) 1. **[ErnieM](https://huggingface.co/docs/transformers/model_doc/ernie_m)** (Baidu から) Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang. から公開された研究論文 [ERNIE-M: Enhanced Multilingual Representation by Aligning Cross-lingual Semantics with Monolingual Corpora](https://arxiv.org/abs/2012.15674) diff --git a/README_ko.md b/README_ko.md index b1fffac022..34ffed005b 100644 --- a/README_ko.md +++ b/README_ko.md @@ -270,6 +270,7 @@ Flax, PyTorch, TensorFlow 설치 페이지에서 이들을 conda로 설치하는 1. **[EfficientFormer](https://huggingface.co/docs/transformers/model_doc/efficientformer)** (from Snap Research) released with the paper [EfficientFormer: Vision Transformers at MobileNetSpeed](https://arxiv.org/abs/2206.01191) by Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, Jian Ren. 1. **[EfficientNet](https://huggingface.co/docs/transformers/model_doc/efficientnet)** (from Google Brain) released with the paper [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) by Mingxing Tan, Quoc V. Le. 1. **[ELECTRA](https://huggingface.co/docs/transformers/model_doc/electra)** (Google Research/Stanford University 에서) Kevin Clark, Minh-Thang Luong, Quoc V. Le, Christopher D. Manning 의 [ELECTRA: Pre-training text encoders as discriminators rather than generators](https://arxiv.org/abs/2003.10555) 논문과 함께 발표했습니다. +1. **[EnCodec](https://huggingface.co/docs/transformers/main/model_doc/encodec)** (Meta AI 에서 제공)은 Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi.의 [High Fidelity Neural Audio Compression](https://arxiv.org/abs/2210.13438)논문과 함께 발표했습니다. 1. **[EncoderDecoder](https://huggingface.co/docs/transformers/model_doc/encoder-decoder)** (Google Research 에서) Sascha Rothe, Shashi Narayan, Aliaksei Severyn 의 [Leveraging Pre-trained Checkpoints for Sequence Generation Tasks](https://arxiv.org/abs/1907.12461) 논문과 함께 발표했습니다. 1. **[ERNIE](https://huggingface.co/docs/transformers/model_doc/ernie)** (Baidu 에서) Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Xuyi Chen, Han Zhang, Xin Tian, Danxiang Zhu, Hao Tian, Hua Wu 의 [ERNIE: Enhanced Representation through Knowledge Integration](https://arxiv.org/abs/1904.09223) 논문과 함께 발표했습니다. 1. **[ErnieM](https://huggingface.co/docs/transformers/model_doc/ernie_m)** (Baidu 에서 제공)은 Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang.의 [ERNIE-M: Enhanced Multilingual Representation by Aligning Cross-lingual Semantics with Monolingual Corpora](https://arxiv.org/abs/2012.15674)논문과 함께 발표했습니다. diff --git a/README_zh-hans.md b/README_zh-hans.md index a463c171e5..469fd86760 100644 --- a/README_zh-hans.md +++ b/README_zh-hans.md @@ -294,6 +294,7 @@ conda install -c huggingface transformers 1. **[EfficientFormer](https://huggingface.co/docs/transformers/model_doc/efficientformer)** (来自 Snap Research) 伴随论文 [EfficientFormer: Vision Transformers at MobileNetSpeed](https://arxiv.org/abs/2206.01191) 由 Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, Jian Ren 发布。 1. **[EfficientNet](https://huggingface.co/docs/transformers/model_doc/efficientnet)** (from Google Brain) released with the paper [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) by Mingxing Tan, Quoc V. Le. 1. **[ELECTRA](https://huggingface.co/docs/transformers/model_doc/electra)** (来自 Google Research/Stanford University) 伴随论文 [ELECTRA: Pre-training text encoders as discriminators rather than generators](https://arxiv.org/abs/2003.10555) 由 Kevin Clark, Minh-Thang Luong, Quoc V. Le, Christopher D. Manning 发布。 +1. **[EnCodec](https://huggingface.co/docs/transformers/main/model_doc/encodec)** (来自 Meta AI) 伴随论文 [High Fidelity Neural Audio Compression](https://arxiv.org/abs/2210.13438) 由 Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi 发布。 1. **[EncoderDecoder](https://huggingface.co/docs/transformers/model_doc/encoder-decoder)** (来自 Google Research) 伴随论文 [Leveraging Pre-trained Checkpoints for Sequence Generation Tasks](https://arxiv.org/abs/1907.12461) 由 Sascha Rothe, Shashi Narayan, Aliaksei Severyn 发布。 1. **[ERNIE](https://huggingface.co/docs/transformers/model_doc/ernie)** (来自 Baidu) 伴随论文 [ERNIE: Enhanced Representation through Knowledge Integration](https://arxiv.org/abs/1904.09223) by Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Xuyi Chen, Han Zhang, Xin Tian, Danxiang Zhu, Hao Tian, Hua Wu 发布。 1. **[ErnieM](https://huggingface.co/docs/transformers/model_doc/ernie_m)** (来自 Baidu) 伴随论文 [ERNIE-M: Enhanced Multilingual Representation by Aligning Cross-lingual Semantics with Monolingual Corpora](https://arxiv.org/abs/2012.15674) 由 Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang 发布。 diff --git a/README_zh-hant.md b/README_zh-hant.md index 7d28c925a7..56b4bfad6d 100644 --- a/README_zh-hant.md +++ b/README_zh-hant.md @@ -306,6 +306,7 @@ conda install -c huggingface transformers 1. **[EfficientFormer](https://huggingface.co/docs/transformers/model_doc/efficientformer)** (from Snap Research) released with the paper [EfficientFormer: Vision Transformers at MobileNetSpeed](https://arxiv.org/abs/2206.01191) by Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, Jian Ren. 1. **[EfficientNet](https://huggingface.co/docs/transformers/model_doc/efficientnet)** (from Google Brain) released with the paper [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) by Mingxing Tan, Quoc V. Le. 1. **[ELECTRA](https://huggingface.co/docs/transformers/model_doc/electra)** (from Google Research/Stanford University) released with the paper [ELECTRA: Pre-training text encoders as discriminators rather than generators](https://arxiv.org/abs/2003.10555) by Kevin Clark, Minh-Thang Luong, Quoc V. Le, Christopher D. Manning. +1. **[EnCodec](https://huggingface.co/docs/transformers/main/model_doc/encodec)** (from Meta AI) released with the paper [High Fidelity Neural Audio Compression](https://arxiv.org/abs/2210.13438) by Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi. 1. **[EncoderDecoder](https://huggingface.co/docs/transformers/model_doc/encoder-decoder)** (from Google Research) released with the paper [Leveraging Pre-trained Checkpoints for Sequence Generation Tasks](https://arxiv.org/abs/1907.12461) by Sascha Rothe, Shashi Narayan, Aliaksei Severyn. 1. **[ERNIE](https://huggingface.co/docs/transformers/model_doc/ernie)** (from Baidu) released with the paper [ERNIE: Enhanced Representation through Knowledge Integration](https://arxiv.org/abs/1904.09223) by Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Xuyi Chen, Han Zhang, Xin Tian, Danxiang Zhu, Hao Tian, Hua Wu. 1. **[ErnieM](https://huggingface.co/docs/transformers/model_doc/ernie_m)** (from Baidu) released with the paper [ERNIE-M: Enhanced Multilingual Representation by Aligning Cross-lingual Semantics with Monolingual Corpora](https://arxiv.org/abs/2012.15674) by Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang. diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 60f98607bf..51ef7287d8 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -541,6 +541,8 @@ title: Audio Spectrogram Transformer - local: model_doc/clap title: CLAP + - local: model_doc/encodec + title: EnCodec - local: model_doc/hubert title: Hubert - local: model_doc/mctct diff --git a/docs/source/en/index.mdx b/docs/source/en/index.mdx index 6075c0ecf3..1690b4cb7c 100644 --- a/docs/source/en/index.mdx +++ b/docs/source/en/index.mdx @@ -107,6 +107,7 @@ The documentation is organized into five sections: 1. **[EfficientFormer](model_doc/efficientformer)** (from Snap Research) released with the paper [EfficientFormer: Vision Transformers at MobileNetSpeed](https://arxiv.org/abs/2206.01191) by Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, Jian Ren. 1. **[EfficientNet](model_doc/efficientnet)** (from Google Brain) released with the paper [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) by Mingxing Tan, Quoc V. Le. 1. **[ELECTRA](model_doc/electra)** (from Google Research/Stanford University) released with the paper [ELECTRA: Pre-training text encoders as discriminators rather than generators](https://arxiv.org/abs/2003.10555) by Kevin Clark, Minh-Thang Luong, Quoc V. Le, Christopher D. Manning. +1. **[EnCodec](model_doc/encodec)** (from Meta AI) released with the paper [High Fidelity Neural Audio Compression](https://arxiv.org/abs/2210.13438) by Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi. 1. **[EncoderDecoder](model_doc/encoder-decoder)** (from Google Research) released with the paper [Leveraging Pre-trained Checkpoints for Sequence Generation Tasks](https://arxiv.org/abs/1907.12461) by Sascha Rothe, Shashi Narayan, Aliaksei Severyn. 1. **[ERNIE](model_doc/ernie)** (from Baidu) released with the paper [ERNIE: Enhanced Representation through Knowledge Integration](https://arxiv.org/abs/1904.09223) by Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Xuyi Chen, Han Zhang, Xin Tian, Danxiang Zhu, Hao Tian, Hua Wu. 1. **[ErnieM](model_doc/ernie_m)** (from Baidu) released with the paper [ERNIE-M: Enhanced Multilingual Representation by Aligning Cross-lingual Semantics with Monolingual Corpora](https://arxiv.org/abs/2012.15674) by Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang. @@ -318,6 +319,7 @@ Flax), PyTorch, and/or TensorFlow. | EfficientFormer | ❌ | ❌ | ✅ | ✅ | ❌ | | EfficientNet | ❌ | ❌ | ✅ | ❌ | ❌ | | ELECTRA | ✅ | ✅ | ✅ | ✅ | ✅ | +| EnCodec | ❌ | ❌ | ✅ | ❌ | ❌ | | Encoder decoder | ❌ | ❌ | ✅ | ✅ | ✅ | | ERNIE | ❌ | ❌ | ✅ | ❌ | ❌ | | ErnieM | ✅ | ❌ | ✅ | ❌ | ❌ | diff --git a/docs/source/en/model_doc/encodec.mdx b/docs/source/en/model_doc/encodec.mdx new file mode 100644 index 0000000000..f98156db1d --- /dev/null +++ b/docs/source/en/model_doc/encodec.mdx @@ -0,0 +1,59 @@ + + +# EnCodec + +## Overview + +The EnCodec neural codec model was proposed in [High Fidelity Neural Audio Compression](https://arxiv.org/abs/2210.13438) by Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi. + +The abstract from the paper is the following: + +*We introduce a state-of-the-art real-time, high-fidelity, audio codec leveraging neural networks. It consists in a streaming encoder-decoder architecture with quantized latent space trained in an end-to-end fashion. We simplify and speed-up the training by using a single multiscale spectrogram adversary that efficiently reduces artifacts and produce high-quality samples. We introduce a novel loss balancer mechanism to stabilize training: the weight of a loss now defines the fraction of the overall gradient it should represent, thus decoupling the choice of this hyper-parameter from the typical scale of the loss. Finally, we study how lightweight Transformer models can be used to further compress the obtained representation by up to 40%, while staying faster than real time. We provide a detailed description of the key design choices of the proposed model including: training objective, architectural changes and a study of various perceptual loss functions. We present an extensive subjective evaluation (MUSHRA tests) together with an ablation study for a range of bandwidths and audio domains, including speech, noisy-reverberant speech, and music. Our approach is superior to the baselines methods across all evaluated settings, considering both 24 kHz monophonic and 48 kHz stereophonic audio.* + +This model was contributed by [Matthijs](https://huggingface.co/Matthijs), [Patrick Von Platen](https://huggingface.co/patrickvonplaten) and [Arthur Zucker](https://huggingface.co/ArthurZ). +The original code can be found [here](https://github.com/facebookresearch/encodec). +Here is a quick example of how to encode and decode an audio using this model: + +```python +>>> from datasets import load_dataset, Audio +>>> from transformers import EncodecModel, AutoProcessor +>>> librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + +>>> model = EncodecModel.from_pretrained("facebook/encodec_24khz") +>>> processor = AutoProcessor.from_pretrained("facebook/encodec_24khz") +>>> librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate)) +>>> audio_sample = librispeech_dummy[-1]["audio"]["array"] +>>> inputs = processor(raw_audio=audio_sample, sampling_rate=processor.sampling_rate, return_tensors="pt") + +>>> encoder_outputs = model.encode(inputs["input_values"], inputs["padding_mask"]) +>>> audio_values = model.decode(encoder_outputs.audio_codes, encoder_outputs.audio_scales, inputs["padding_mask"])[0] +>>> # or the equivalent with a forward pass +>>> audio_values = model(inputs["input_values"], inputs["padding_mask"]).audio_values +``` + + +## EncodecConfig + +[[autodoc]] EncodecConfig + +## EncodecFeatureExtractor + +[[autodoc]] EncodecFeatureExtractor + - __call__ + +## EncodecModel + +[[autodoc]] EncodecModel + - decode + - encode + - forward diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index c2966423da..f4e53099b1 100644 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -281,6 +281,10 @@ _import_structure = { "models.efficientformer": ["EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "EfficientFormerConfig"], "models.efficientnet": ["EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP", "EfficientNetConfig"], "models.electra": ["ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP", "ElectraConfig", "ElectraTokenizer"], + "models.encodec": [ + "ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP", + "EncodecConfig", + ], "models.encoder_decoder": ["EncoderDecoderConfig"], "models.ernie": [ "ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP", @@ -825,6 +829,7 @@ except OptionalDependencyNotAvailable: ] else: _import_structure["models.audio_spectrogram_transformer"].append("ASTFeatureExtractor") + _import_structure["models.encodec"].append("EncodecFeatureExtractor") _import_structure["models.mctct"].append("MCTCTFeatureExtractor") _import_structure["models.speech_to_text"].append("Speech2TextFeatureExtractor") _import_structure["models.speecht5"].append("SpeechT5FeatureExtractor") @@ -1568,6 +1573,13 @@ else: "load_tf_weights_in_electra", ] ) + _import_structure["models.encodec"].extend( + [ + "ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST", + "EncodecModel", + "EncodecPreTrainedModel", + ] + ) _import_structure["models.encoder_decoder"].append("EncoderDecoderModel") _import_structure["models.ernie"].extend( [ @@ -4100,6 +4112,10 @@ if TYPE_CHECKING: from .models.efficientformer import EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, EfficientFormerConfig from .models.efficientnet import EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP, EfficientNetConfig from .models.electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraTokenizer + from .models.encodec import ( + ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP, + EncodecConfig, + ) from .models.encoder_decoder import EncoderDecoderConfig from .models.ernie import ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieConfig from .models.ernie_m import ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieMConfig @@ -4598,6 +4614,7 @@ if TYPE_CHECKING: from .utils.dummy_speech_objects import * else: from .models.audio_spectrogram_transformer import ASTFeatureExtractor + from .models.encodec import EncodecFeatureExtractor from .models.mctct import MCTCTFeatureExtractor from .models.speech_to_text import Speech2TextFeatureExtractor from .models.speecht5 import SpeechT5FeatureExtractor @@ -5210,6 +5227,11 @@ if TYPE_CHECKING: ElectraPreTrainedModel, load_tf_weights_in_electra, ) + from .models.encodec import ( + ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST, + EncodecModel, + EncodecPreTrainedModel, + ) from .models.encoder_decoder import EncoderDecoderModel from .models.ernie import ( ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST, diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py index 0ee88e0ba3..c410596985 100644 --- a/src/transformers/models/__init__.py +++ b/src/transformers/models/__init__.py @@ -72,6 +72,7 @@ from . import ( efficientformer, efficientnet, electra, + encodec, encoder_decoder, ernie, ernie_m, diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index 6445f62acb..f2c1c88299 100755 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -80,6 +80,7 @@ CONFIG_MAPPING_NAMES = OrderedDict( ("efficientformer", "EfficientFormerConfig"), ("efficientnet", "EfficientNetConfig"), ("electra", "ElectraConfig"), + ("encodec", "EncodecConfig"), ("encoder-decoder", "EncoderDecoderConfig"), ("ernie", "ErnieConfig"), ("ernie_m", "ErnieMConfig"), @@ -273,6 +274,7 @@ CONFIG_ARCHIVE_MAP_MAPPING_NAMES = OrderedDict( ("efficientformer", "EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("efficientnet", "EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("electra", "ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP"), + ("encodec", "ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("ernie", "ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("ernie_m", "ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("esm", "ESM_PRETRAINED_CONFIG_ARCHIVE_MAP"), @@ -461,6 +463,7 @@ MODEL_NAMES_MAPPING = OrderedDict( ("efficientformer", "EfficientFormer"), ("efficientnet", "EfficientNet"), ("electra", "ELECTRA"), + ("encodec", "EnCodec"), ("encoder-decoder", "Encoder decoder"), ("ernie", "ERNIE"), ("ernie_m", "ErnieM"), diff --git a/src/transformers/models/auto/feature_extraction_auto.py b/src/transformers/models/auto/feature_extraction_auto.py index 681a59e0b2..b100b62ac0 100644 --- a/src/transformers/models/auto/feature_extraction_auto.py +++ b/src/transformers/models/auto/feature_extraction_auto.py @@ -54,6 +54,7 @@ FEATURE_EXTRACTOR_MAPPING_NAMES = OrderedDict( ("dinat", "ViTFeatureExtractor"), ("donut-swin", "DonutFeatureExtractor"), ("dpt", "DPTFeatureExtractor"), + ("encodec", "EncodecFeatureExtractor"), ("flava", "FlavaFeatureExtractor"), ("glpn", "GLPNFeatureExtractor"), ("groupvit", "CLIPFeatureExtractor"), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 2c23da3c55..4c675ff050 100755 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -79,6 +79,7 @@ MODEL_MAPPING_NAMES = OrderedDict( ("efficientformer", "EfficientFormerModel"), ("efficientnet", "EfficientNetModel"), ("electra", "ElectraModel"), + ("encodec", "EncodecModel"), ("ernie", "ErnieModel"), ("ernie_m", "ErnieMModel"), ("esm", "EsmModel"), diff --git a/src/transformers/models/encodec/__init__.py b/src/transformers/models/encodec/__init__.py new file mode 100644 index 0000000000..d3d9488968 --- /dev/null +++ b/src/transformers/models/encodec/__init__.py @@ -0,0 +1,65 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import ( + OptionalDependencyNotAvailable, + _LazyModule, + is_torch_available, +) + + +_import_structure = { + "configuration_encodec": [ + "ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP", + "EncodecConfig", + ], + "feature_extraction_encodec": ["EncodecFeatureExtractor"], +} + +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["modeling_encodec"] = [ + "ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST", + "EncodecModel", + "EncodecPreTrainedModel", + ] + +if TYPE_CHECKING: + from .configuration_encodec import ( + ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP, + EncodecConfig, + ) + from .feature_extraction_encodec import EncodecFeatureExtractor + + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .modeling_encodec import ( + ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST, + EncodecModel, + EncodecPreTrainedModel, + ) + +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/src/transformers/models/encodec/configuration_encodec.py b/src/transformers/models/encodec/configuration_encodec.py new file mode 100644 index 0000000000..9ea2cfee94 --- /dev/null +++ b/src/transformers/models/encodec/configuration_encodec.py @@ -0,0 +1,189 @@ +# coding=utf-8 +# Copyright 2023 Meta Platforms, Inc. and affiliates, and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" EnCodec model configuration""" + + +import math +from typing import Optional + +import numpy as np + +from ...configuration_utils import PretrainedConfig +from ...utils import logging + + +logger = logging.get_logger(__name__) + +ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP = { + "facebook/encodec_24khz": "https://huggingface.co/facebook/encodec_24khz/resolve/main/config.json", + "facebook/encodec_48khz": "https://huggingface.co/facebook/encodec_48khz/resolve/main/config.json", +} + + +class EncodecConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of an [`EncodecModel`]. It is used to instantiate a + Encodec model according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of the + [facebook/encodec_24khz](https://huggingface.co/facebook/encodec_24khz) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + target_bandwidths (`List[float]`, *optional*, defaults to `[1.5, 3.0, 6.0, 12.0, 24.0]`): + The range of diffent bandwiths the model can encode audio with. + sampling_rate (`int`, *optional*, defaults to 24000): + The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz). + audio_channels (`int`, *optional*, defaults to 1): + Number of channels in the audio data. Either 1 for mono or 2 for stereo. + normalize (`bool`, *optional*, defaults to `False`): + Whether the audio shall be normalized when passed. + chunk_length_s (`float`, *optional*): + If defined the audio is pre-processed into chunks of lengths `chunk_length_s` and then encoded. + overlap (`float`, *optional*): + Defines the overlap between each chunk. It is used to compute the `chunk_stride` using the following + formulae : `int((1.0 - self.overlap) * self.chunk_length)`. + hidden_size (`int`, *optional*, defaults to 128): + Intermediate representation dimension. + num_filters (`int`, *optional*, defaults to 32): + Number of convolution kernels of first `EncodecConv1d` down sampling layer. + num_residual_layers (`int`, *optional*, defaults to 1): + Number of residual layers. + upsampling_ratios (`Sequence[int]` , *optional*, defaults to `[8, 5, 4, 2]`): + Kernel size and stride ratios. The encoder uses downsampling ratios instead of upsampling ratios, hence it + will use the ratios in the reverse order to the ones specified here that must match the decoder order. + norm_type (`str`, *optional*, defaults to `"weight_norm"`): + Normalization method. Should be in `["weight_norm", "time_group_norm"]` + kernel_size (`int`, *optional*, defaults to 7): + Kernel size for the initial convolution. + last_kernel_size (`int`, *optional*, defaults to 7): + Kernel size for the last convolution layer. + residual_kernel_size (`int`, *optional*, defaults to 3): + Kernel size for the residual layers. + dilation_growth_rate (`int`, *optional*, defaults to 2): + How much to increase the dilation with each layer. + use_causal_conv (`bool`, *optional*, defaults to `True`): + Whether to use fully causal convolution. + pad_mode (`str`, *optional*, defaults to `"reflect"`): + Padding mode for the convolutions. + compress (`int`, *optional*, defaults to 2): + Reduced dimensionality in residual branches (from Demucs v3). + num_lstm_layers (`int`, *optional*, defaults to 2): + Number of LSTM layers at the end of the encoder. + trim_right_ratio (`float`, *optional*, defaults to 1.0): + Ratio for trimming at the right of the transposed convolution under the `use_causal_conv = True` setup. If + equal to 1.0, it means that all the trimming is done at the right. + codebook_size (`int`, *optional*, defaults to 1024): + Number of discret codes that make up VQVAE. + codebook_dim (`int`, *optional*): + Dimension of the codebook vectors. If not defined, uses `hidden_size`. + + Example: + + ```python + >>> from transformers import EncodecModel, EncodecConfig + + >>> # Initializing a "facebook/encodec_24khz" style configuration + >>> configuration = EncodecConfig() + + >>> # Initializing a model (with random weights) from the "facebook/encodec_24khz" style configuration + >>> model = EncodecModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + model_type = "encodec" + + def __init__( + self, + target_bandwidths=[1.5, 3.0, 6.0, 12.0, 24.0], + sampling_rate=24_000, + audio_channels=1, + normalize=False, + chunk_length_s=None, + overlap=None, + hidden_size=128, + num_filters=32, + num_residual_layers=1, + upsampling_ratios=[8, 5, 4, 2], + norm_type="weight_norm", + kernel_size=7, + last_kernel_size=7, + residual_kernel_size=3, + dilation_growth_rate=2, + use_causal_conv=True, + pad_mode="reflect", + compress=2, + num_lstm_layers=2, + trim_right_ratio=1.0, + codebook_size=1024, + codebook_dim=None, + **kwargs, + ): + self.target_bandwidths = target_bandwidths + self.sampling_rate = sampling_rate + self.audio_channels = audio_channels + self.normalize = normalize + self.chunk_length_s = chunk_length_s + self.overlap = overlap + self.hidden_size = hidden_size + self.num_filters = num_filters + self.num_residual_layers = num_residual_layers + self.upsampling_ratios = upsampling_ratios + self.norm_type = norm_type + self.kernel_size = kernel_size + self.last_kernel_size = last_kernel_size + self.residual_kernel_size = residual_kernel_size + self.dilation_growth_rate = dilation_growth_rate + self.use_causal_conv = use_causal_conv + self.pad_mode = pad_mode + self.compress = compress + self.num_lstm_layers = num_lstm_layers + self.trim_right_ratio = trim_right_ratio + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim if codebook_dim is not None else hidden_size + + if self.norm_type not in ["weight_norm", "time_group_norm"]: + raise ValueError( + f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}' + ) + + super().__init__(**kwargs) + + # This is a property because you might want to change the chunk_length_s on the fly + @property + def chunk_length(self) -> Optional[int]: + if self.chunk_length_s is None: + return None + else: + return int(self.chunk_length_s * self.sampling_rate) + + # This is a property because you might want to change the chunk_length_s on the fly + @property + def chunk_stride(self) -> Optional[int]: + if self.chunk_length_s is None or self.overlap is None: + return None + else: + return max(1, int((1.0 - self.overlap) * self.chunk_length)) + + @property + def frame_rate(self) -> int: + hop_length = np.prod(self.upsampling_ratios) + return math.ceil(self.sampling_rate / hop_length) + + @property + def num_quantizers(self) -> int: + return int(1000 * self.target_bandwidths[-1] // (self.frame_rate * 10)) diff --git a/src/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py b/src/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py new file mode 100644 index 0000000000..cd7ead3d72 --- /dev/null +++ b/src/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py @@ -0,0 +1,352 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert EnCodec checkpoints.""" + +import argparse + +import torch + +from transformers import ( + EncodecConfig, + EncodecFeatureExtractor, + EncodecModel, + logging, +) + + +# checkpoints downloaded from: +# https://dl.fbaipublicfiles.com/encodec/v0/encodec_24khz-d7cc33bc.th +# https://dl.fbaipublicfiles.com/encodec/v0/encodec_48khz-7e698e3e.th + + +logging.set_verbosity_info() +logger = logging.get_logger("transformers.models.encodec") + +MAPPING_QUANTIZER = { + "quantizer.vq.layers.*._codebook.inited": "quantizer.layers.*.codebook.inited", + "quantizer.vq.layers.*._codebook.cluster_size": "quantizer.layers.*.codebook.cluster_size", + "quantizer.vq.layers.*._codebook.embed": "quantizer.layers.*.codebook.embed", + "quantizer.vq.layers.*._codebook.embed_avg": "quantizer.layers.*.codebook.embed_avg", +} +MAPPING_ENCODER = { + "encoder.model.0.conv.conv": "encoder.layers.0.conv", + "encoder.model.1.block.1.conv.conv": "encoder.layers.1.block.1.conv", + "encoder.model.1.block.3.conv.conv": "encoder.layers.1.block.3.conv", + "encoder.model.1.shortcut.conv.conv": "encoder.layers.1.shortcut.conv", + "encoder.model.3.conv.conv": "encoder.layers.3.conv", + "encoder.model.4.block.1.conv.conv": "encoder.layers.4.block.1.conv", + "encoder.model.4.block.3.conv.conv": "encoder.layers.4.block.3.conv", + "encoder.model.4.shortcut.conv.conv": "encoder.layers.4.shortcut.conv", + "encoder.model.6.conv.conv": "encoder.layers.6.conv", + "encoder.model.7.block.1.conv.conv": "encoder.layers.7.block.1.conv", + "encoder.model.7.block.3.conv.conv": "encoder.layers.7.block.3.conv", + "encoder.model.7.shortcut.conv.conv": "encoder.layers.7.shortcut.conv", + "encoder.model.9.conv.conv": "encoder.layers.9.conv", + "encoder.model.10.block.1.conv.conv": "encoder.layers.10.block.1.conv", + "encoder.model.10.block.3.conv.conv": "encoder.layers.10.block.3.conv", + "encoder.model.10.shortcut.conv.conv": "encoder.layers.10.shortcut.conv", + "encoder.model.12.conv.conv": "encoder.layers.12.conv", + "encoder.model.13.lstm": "encoder.layers.13.lstm", + "encoder.model.15.conv.conv": "encoder.layers.15.conv", +} +MAPPING_ENCODER_48K = { + "encoder.model.0.conv.norm": "encoder.layers.0.norm", + "encoder.model.1.block.1.conv.norm": "encoder.layers.1.block.1.norm", + "encoder.model.1.block.3.conv.norm": "encoder.layers.1.block.3.norm", + "encoder.model.1.shortcut.conv.norm": "encoder.layers.1.shortcut.norm", + "encoder.model.3.conv.norm": "encoder.layers.3.norm", + "encoder.model.4.block.1.conv.norm": "encoder.layers.4.block.1.norm", + "encoder.model.4.block.3.conv.norm": "encoder.layers.4.block.3.norm", + "encoder.model.4.shortcut.conv.norm": "encoder.layers.4.shortcut.norm", + "encoder.model.6.conv.norm": "encoder.layers.6.norm", + "encoder.model.7.block.1.conv.norm": "encoder.layers.7.block.1.norm", + "encoder.model.7.block.3.conv.norm": "encoder.layers.7.block.3.norm", + "encoder.model.7.shortcut.conv.norm": "encoder.layers.7.shortcut.norm", + "encoder.model.9.conv.norm": "encoder.layers.9.norm", + "encoder.model.10.block.1.conv.norm": "encoder.layers.10.block.1.norm", + "encoder.model.10.block.3.conv.norm": "encoder.layers.10.block.3.norm", + "encoder.model.10.shortcut.conv.norm": "encoder.layers.10.shortcut.norm", + "encoder.model.12.conv.norm": "encoder.layers.12.norm", + "encoder.model.15.conv.norm": "encoder.layers.15.norm", +} +MAPPING_DECODER = { + "decoder.model.0.conv.conv": "decoder.layers.0.conv", + "decoder.model.1.lstm": "decoder.layers.1.lstm", + "decoder.model.3.convtr.convtr": "decoder.layers.3.conv", + "decoder.model.4.block.1.conv.conv": "decoder.layers.4.block.1.conv", + "decoder.model.4.block.3.conv.conv": "decoder.layers.4.block.3.conv", + "decoder.model.4.shortcut.conv.conv": "decoder.layers.4.shortcut.conv", + "decoder.model.6.convtr.convtr": "decoder.layers.6.conv", + "decoder.model.7.block.1.conv.conv": "decoder.layers.7.block.1.conv", + "decoder.model.7.block.3.conv.conv": "decoder.layers.7.block.3.conv", + "decoder.model.7.shortcut.conv.conv": "decoder.layers.7.shortcut.conv", + "decoder.model.9.convtr.convtr": "decoder.layers.9.conv", + "decoder.model.10.block.1.conv.conv": "decoder.layers.10.block.1.conv", + "decoder.model.10.block.3.conv.conv": "decoder.layers.10.block.3.conv", + "decoder.model.10.shortcut.conv.conv": "decoder.layers.10.shortcut.conv", + "decoder.model.12.convtr.convtr": "decoder.layers.12.conv", + "decoder.model.13.block.1.conv.conv": "decoder.layers.13.block.1.conv", + "decoder.model.13.block.3.conv.conv": "decoder.layers.13.block.3.conv", + "decoder.model.13.shortcut.conv.conv": "decoder.layers.13.shortcut.conv", + "decoder.model.15.conv.conv": "decoder.layers.15.conv", +} +MAPPING_DECODER_48K = { + "decoder.model.0.conv.norm": "decoder.layers.0.norm", + "decoder.model.3.convtr.norm": "decoder.layers.3.norm", + "decoder.model.4.block.1.conv.norm": "decoder.layers.4.block.1.norm", + "decoder.model.4.block.3.conv.norm": "decoder.layers.4.block.3.norm", + "decoder.model.4.shortcut.conv.norm": "decoder.layers.4.shortcut.norm", + "decoder.model.6.convtr.norm": "decoder.layers.6.norm", + "decoder.model.7.block.1.conv.norm": "decoder.layers.7.block.1.norm", + "decoder.model.7.block.3.conv.norm": "decoder.layers.7.block.3.norm", + "decoder.model.7.shortcut.conv.norm": "decoder.layers.7.shortcut.norm", + "decoder.model.9.convtr.norm": "decoder.layers.9.norm", + "decoder.model.10.block.1.conv.norm": "decoder.layers.10.block.1.norm", + "decoder.model.10.block.3.conv.norm": "decoder.layers.10.block.3.norm", + "decoder.model.10.shortcut.conv.norm": "decoder.layers.10.shortcut.norm", + "decoder.model.12.convtr.norm": "decoder.layers.12.norm", + "decoder.model.13.block.1.conv.norm": "decoder.layers.13.block.1.norm", + "decoder.model.13.block.3.conv.norm": "decoder.layers.13.block.3.norm", + "decoder.model.13.shortcut.conv.norm": "decoder.layers.13.shortcut.norm", + "decoder.model.15.conv.norm": "decoder.layers.15.norm", +} +MAPPING_24K = { + **MAPPING_QUANTIZER, + **MAPPING_ENCODER, + **MAPPING_DECODER, +} +MAPPING_48K = { + **MAPPING_QUANTIZER, + **MAPPING_ENCODER, + **MAPPING_ENCODER_48K, + **MAPPING_DECODER, + **MAPPING_DECODER_48K, +} +TOP_LEVEL_KEYS = [] +IGNORE_KEYS = [] + + +def set_recursively(hf_pointer, key, value, full_name, weight_type): + for attribute in key.split("."): + hf_pointer = getattr(hf_pointer, attribute) + + if weight_type is not None: + hf_shape = getattr(hf_pointer, weight_type).shape + else: + hf_shape = hf_pointer.shape + + if hf_shape != value.shape: + raise ValueError( + f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be" + f" {value.shape} for {full_name}" + ) + + if weight_type == "weight": + hf_pointer.weight.data = value + elif weight_type == "weight_g": + hf_pointer.weight_g.data = value + elif weight_type == "weight_v": + hf_pointer.weight_v.data = value + elif weight_type == "bias": + hf_pointer.bias.data = value + elif weight_type == "running_mean": + hf_pointer.running_mean.data = value + elif weight_type == "running_var": + hf_pointer.running_var.data = value + elif weight_type == "num_batches_tracked": + hf_pointer.num_batches_tracked.data = value + elif weight_type == "weight_ih_l0": + hf_pointer.weight_ih_l0.data = value + elif weight_type == "weight_hh_l0": + hf_pointer.weight_hh_l0.data = value + elif weight_type == "bias_ih_l0": + hf_pointer.bias_ih_l0.data = value + elif weight_type == "bias_hh_l0": + hf_pointer.bias_hh_l0.data = value + elif weight_type == "weight_ih_l1": + hf_pointer.weight_ih_l1.data = value + elif weight_type == "weight_hh_l1": + hf_pointer.weight_hh_l1.data = value + elif weight_type == "bias_ih_l1": + hf_pointer.bias_ih_l1.data = value + elif weight_type == "bias_hh_l1": + hf_pointer.bias_hh_l1.data = value + else: + hf_pointer.data = value + + logger.info(f"{key + ('.' + weight_type if weight_type is not None else '')} was initialized from {full_name}.") + + +def should_ignore(name, ignore_keys): + for key in ignore_keys: + if key.endswith(".*"): + if name.startswith(key[:-1]): + return True + elif ".*." in key: + prefix, suffix = key.split(".*.") + if prefix in name and suffix in name: + return True + elif key in name: + return True + return False + + +def recursively_load_weights(orig_dict, hf_model, model_name): + unused_weights = [] + + if model_name == "encodec_24khz": + MAPPING = MAPPING_24K + elif model_name == "encodec_48khz": + MAPPING = MAPPING_48K + else: + raise ValueError(f"Unsupported model: {model_name}") + + for name, value in orig_dict.items(): + if should_ignore(name, IGNORE_KEYS): + logger.info(f"{name} was ignored") + continue + + is_used = False + for key, mapped_key in MAPPING.items(): + if "*" in key: + prefix, suffix = key.split(".*.") + if prefix in name and suffix in name: + key = suffix + + if key in name: + # HACK otherwise .embed gets initialized with .embed_avg too + if key.endswith("embed") and name.endswith("embed_avg"): + continue + + is_used = True + if "*" in mapped_key: + layer_index = name.split(key)[0].split(".")[-2] + mapped_key = mapped_key.replace("*", layer_index) + if "weight_g" in name: + weight_type = "weight_g" + elif "weight_v" in name: + weight_type = "weight_v" + elif "weight_ih_l0" in name: + weight_type = "weight_ih_l0" + elif "weight_hh_l0" in name: + weight_type = "weight_hh_l0" + elif "bias_ih_l0" in name: + weight_type = "bias_ih_l0" + elif "bias_hh_l0" in name: + weight_type = "bias_hh_l0" + elif "weight_ih_l1" in name: + weight_type = "weight_ih_l1" + elif "weight_hh_l1" in name: + weight_type = "weight_hh_l1" + elif "bias_ih_l1" in name: + weight_type = "bias_ih_l1" + elif "bias_hh_l1" in name: + weight_type = "bias_hh_l1" + elif "bias" in name: + weight_type = "bias" + elif "weight" in name: + weight_type = "weight" + elif "running_mean" in name: + weight_type = "running_mean" + elif "running_var" in name: + weight_type = "running_var" + elif "num_batches_tracked" in name: + weight_type = "num_batches_tracked" + else: + weight_type = None + set_recursively(hf_model, mapped_key, value, name, weight_type) + continue + if not is_used: + unused_weights.append(name) + + logger.warning(f"Unused weights: {unused_weights}") + + +@torch.no_grad() +def convert_checkpoint( + model_name, + checkpoint_path, + pytorch_dump_folder_path, + config_path=None, + repo_id=None, +): + """ + Copy/paste/tweak model's weights to transformers design. + """ + if config_path is not None: + config = EncodecConfig.from_pretrained(config_path) + else: + config = EncodecConfig() + + if model_name == "encodec_24khz": + pass # config is already correct + elif model_name == "encodec_48khz": + config.upsampling_ratios = [8, 5, 4, 2] + config.target_bandwidths = [3.0, 6.0, 12.0, 24.0] + config.sampling_rate = 48_000 + config.audio_channels = 2 + config.use_causal_conv = False + config.norm_type = "time_group_norm" + config.normalize = True + config.chunk_length_s = 1.0 + config.overlap = 0.01 + else: + raise ValueError(f"Unknown model name: {model_name}") + + model = EncodecModel(config) + + feature_extractor = EncodecFeatureExtractor( + feature_size=config.audio_channels, + sampling_rate=config.sampling_rate, + chunk_length_s=config.chunk_length_s, + overlap=config.overlap, + ) + feature_extractor.save_pretrained(pytorch_dump_folder_path) + + original_checkpoint = torch.load(checkpoint_path) + recursively_load_weights(original_checkpoint, model, model_name) + model.save_pretrained(pytorch_dump_folder_path) + + if repo_id: + print("Pushing to the hub...") + feature_extractor.push_to_hub(repo_id) + model.push_to_hub(repo_id) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + default="encodec_24khz", + type=str, + help="The model to convert. Should be one of 'encodec_24khz', 'encodec_48khz'.", + ) + parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint") + parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert") + parser.add_argument( + "--pytorch_dump_folder_path", required=True, default=None, type=str, help="Path to the output PyTorch model." + ) + parser.add_argument( + "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub." + ) + + args = parser.parse_args() + convert_checkpoint( + args.model, + args.checkpoint_path, + args.pytorch_dump_folder_path, + args.config_path, + args.push_to_hub, + ) diff --git a/src/transformers/models/encodec/feature_extraction_encodec.py b/src/transformers/models/encodec/feature_extraction_encodec.py new file mode 100644 index 0000000000..6f7536a52e --- /dev/null +++ b/src/transformers/models/encodec/feature_extraction_encodec.py @@ -0,0 +1,206 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Feature extractor class for EnCodec.""" + +from typing import List, Optional, Union + +import numpy as np + +from ...feature_extraction_sequence_utils import SequenceFeatureExtractor +from ...feature_extraction_utils import BatchFeature +from ...utils import PaddingStrategy, TensorType, logging + + +logger = logging.get_logger(__name__) + + +class EncodecFeatureExtractor(SequenceFeatureExtractor): + r""" + Constructs an EnCodec feature extractor. + + This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains + most of the main methods. Users should refer to this superclass for more information regarding those methods. + + Instantiating a feature extractor with the defaults will yield a similar configuration to that of the + [facebook/encodec_24khz](https://huggingface.co/facebook/encodec_24khz) architecture. + + Args: + feature_size (`int`, *optional*, defaults to 1): + The feature dimension of the extracted features. Use 1 for mono, 2 for stereo. + sampling_rate (`int`, *optional*, defaults to 24000): + The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz). + padding_value (`float`, *optional*, defaults to 0.0): + The value that is used to fill the padding values. + chunk_length_s (`float`, *optional*): + If defined the audio is pre-processed into chunks of lengths `chunk_length_s` and then encoded. + overlap (`float`, *optional*): + Defines the overlap between each chunk. It is used to compute the `chunk_stride` using the following + formulae : `int((1.0 - self.overlap) * self.chunk_length)`. + """ + + model_input_names = ["input_values", "padding_mask"] + + def __init__( + self, + feature_size: int = 1, + sampling_rate: int = 24000, + padding_value: float = 0.0, + chunk_length_s: float = None, + overlap: float = None, + **kwargs, + ): + super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs) + self.chunk_length_s = chunk_length_s + self.overlap = overlap + + # This is a property because you might want to change the chunk_length_s on the fly + @property + def chunk_length(self) -> Optional[int]: + if self.chunk_length_s is None: + return None + else: + return int(self.chunk_length_s * self.sampling_rate) + + # This is a property because you might want to change the chunk_length_s on the fly + @property + def chunk_stride(self) -> Optional[int]: + if self.chunk_length_s is None or self.overlap is None: + return None + else: + return max(1, int((1.0 - self.overlap) * self.chunk_length)) + + def __call__( + self, + raw_audio: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]], + padding: Optional[Union[bool, str, PaddingStrategy]] = None, + truncation: Optional[bool] = False, + max_length: Optional[int] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + sampling_rate: Optional[int] = None, + ) -> BatchFeature: + """ + Main method to featurize and prepare for the model one or several sequence(s). + + Args: + raw_audio (`np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[List[float]]`): + The sequence or batch of sequences to be processed. Each sequence can be a numpy array, a list of float + values, a list of numpy arrays or a list of list of float values. The numpy array must be of shape + `(num_samples,)` for mono audio (`feature_size = 1`), or `(2, num_samples)` for stereo audio + (`feature_size = 2`). + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, *optional*, defaults to `False`): + Activates truncation to cut input sequences longer than `max_length` to `max_length`. + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see above). + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + sampling_rate (`int`, *optional*): + The sampling rate at which the `audio` input was sampled. It is strongly recommended to pass + `sampling_rate` at the forward call to prevent silent errors. + """ + if sampling_rate is not None: + if sampling_rate != self.sampling_rate: + raise ValueError( + f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of" + f" {self.sampling_rate}. Please make sure that the provided audio input was sampled with" + f" {self.sampling_rate} and not {sampling_rate}." + ) + else: + logger.warning( + "It is strongly recommended to pass the `sampling_rate` argument to this function. " + "Failing to do so can result in silent errors that might be hard to debug." + ) + + if padding and truncation: + raise ValueError("Both padding and truncation were set. Make sure you only set one.") + elif padding is None: + # by default let's pad the inputs + padding = True + + is_batched = bool( + isinstance(raw_audio, (list, tuple)) and (isinstance(raw_audio[0], (np.ndarray, tuple, list))) + ) + + if is_batched: + raw_audio = [np.asarray(audio, dtype=np.float32).T for audio in raw_audio] + elif not is_batched and not isinstance(raw_audio, np.ndarray): + raw_audio = np.asarray(raw_audio, dtype=np.float32) + elif isinstance(raw_audio, np.ndarray) and raw_audio.dtype is np.dtype(np.float64): + raw_audio = raw_audio.astype(np.float32) + + # always return batch + if not is_batched: + raw_audio = [np.asarray(raw_audio).T] + + # verify inputs are valid + for idx, example in enumerate(raw_audio): + if example.ndim > 2: + raise ValueError(f"Expected input shape (channels, length) but got shape {example.shape}") + if self.feature_size == 1 and example.ndim != 1: + raise ValueError(f"Expected mono audio but example has {example.shape[-1]} channels") + if self.feature_size == 2 and example.shape[-1] != 2: + raise ValueError(f"Expected stereo audio but example has {example.shape[-1]} channels") + + padded_inputs = None + input_values = BatchFeature({"input_values": raw_audio}) + if self.chunk_stride is not None and self.chunk_length is not None and max_length is None: + if truncation: + max_length = min(array.shape[0] for array in raw_audio) + nb_step = int(np.floor(max_length / self.chunk_stride)) + max_length = (nb_step - 1) * self.chunk_stride + self.chunk_length + elif padding: + max_length = max(array.shape[0] for array in raw_audio) + nb_step = int(np.ceil(max_length / self.chunk_stride)) + max_length = (nb_step - 1) * self.chunk_stride + self.chunk_length + padding = "max_length" + else: + padded_inputs = input_values + + # normal padding on batch + if padded_inputs is None: + padded_inputs = self.pad( + input_values, + max_length=max_length, + truncation=truncation, + padding=padding, + return_attention_mask=padding, + ) + if padding: + padded_inputs["padding_mask"] = padded_inputs.pop("attention_mask") + + input_values = [] + for example in padded_inputs.pop("input_values"): + if self.feature_size == 1: + example = example[..., None] + input_values.append(example.T) + + padded_inputs["input_values"] = input_values + if return_tensors is not None: + padded_inputs = padded_inputs.convert_to_tensors(return_tensors) + + return padded_inputs diff --git a/src/transformers/models/encodec/modeling_encodec.py b/src/transformers/models/encodec/modeling_encodec.py new file mode 100644 index 0000000000..ad1f6a0ee8 --- /dev/null +++ b/src/transformers/models/encodec/modeling_encodec.py @@ -0,0 +1,808 @@ +# coding=utf-8 +# Copyright 2023 Meta Platforms, Inc. and affiliates, and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" PyTorch EnCodec model.""" + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn + +from ...modeling_utils import PreTrainedModel +from ...utils import ( + ModelOutput, + add_start_docstrings, + add_start_docstrings_to_model_forward, + logging, + replace_return_docstrings, +) +from .configuration_encodec import EncodecConfig + + +logger = logging.get_logger(__name__) + + +# General docstring +_CONFIG_FOR_DOC = "EncodecConfig" + + +ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "facebook/encodec_24khz", + "facebook/encodec_48khz", + # See all EnCodec models at https://huggingface.co/models?filter=encodec +] + + +@dataclass +class EncodecOutput(ModelOutput): + """ + Args: + audio_codes (`torch.FloatTensor` of shape `(batch_size, nb_chunks, chunk_length)`, *optional*): + Discret code embeddings computed using `model.encode`. + audio_values (`torch.FlaotTensor` of shape `(batch_size, sequence_length)`, *optional*) + Decoded audio values, obtained using the decoder part of Encodec. + """ + + audio_codes: torch.FloatTensor = None + audio_values: torch.FloatTensor = None + + +@dataclass +class EncodecEncoderOutput(ModelOutput): + """ + Args: + audio_codes (`torch.FloatTensor` of shape `(batch_size, nb_chunks, chunk_length)`, *optional*): + Discret code embeddings computed using `model.encode`. + audio_scales (`torch.Tensor` of shape `(batch_size, nb_chunks)`, *optional*): + Scaling factor for each `audio_codes` input. This is used to unscale each chunk of audio when decoding. + """ + + audio_codes: torch.FloatTensor = None + audio_scales: torch.FloatTensor = None + + +@dataclass +class EncodecDecoderOutput(ModelOutput): + """ + Args: + audio_values (`torch.FloatTensor` of shape `(batch_size, segment_length)`, *optional*): + Decoded audio values, obtained using the decoder part of Encodec. + """ + + audio_values: torch.FloatTensor = None + + +class EncodecConv1d(nn.Module): + """Conv1d with asymmetric or causal padding and normalization.""" + + def __init__( + self, config, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, dilation: int = 1 + ): + super().__init__() + self.causal = config.use_causal_conv + self.pad_mode = config.pad_mode + self.norm_type = config.norm_type + + if self.norm_type not in ["weight_norm", "time_group_norm"]: + raise ValueError( + f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}' + ) + + # warn user on unusual setup between dilation and stride + if stride > 1 and dilation > 1: + logger.warning( + "EncodecConv1d has been initialized with stride > 1 and dilation > 1" + f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})." + ) + + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, dilation=dilation) + if self.norm_type == "weight_norm": + self.conv = nn.utils.weight_norm(self.conv) + elif self.norm_type == "time_group_norm": + self.norm = nn.GroupNorm(1, out_channels) + + @staticmethod + def _get_extra_padding_for_conv1d( + hidden_states: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0 + ) -> int: + """See `pad_for_conv1d`.""" + length = hidden_states.shape[-1] + n_frames = (length - kernel_size + padding_total) / stride + 1 + ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) + return ideal_length - length + + @staticmethod + def _pad1d(hidden_states: torch.Tensor, paddings: Tuple[int, int], mode: str = "zero", value: float = 0.0): + """Tiny wrapper around torch.nn.functional.pad, just to allow for reflect padding on small input. + If this is the case, we insert extra 0 padding to the right before the reflection happens. + """ + length = hidden_states.shape[-1] + padding_left, padding_right = paddings + if not mode == "reflect": + return nn.functional.pad(hidden_states, paddings, mode, value) + + max_pad = max(padding_left, padding_right) + extra_pad = 0 + if length <= max_pad: + extra_pad = max_pad - length + 1 + hidden_states = nn.functional.pad(hidden_states, (0, extra_pad)) + padded = nn.functional.pad(hidden_states, paddings, mode, value) + end = padded.shape[-1] - extra_pad + return padded[..., :end] + + def forward(self, hidden_states): + kernel_size = self.conv.kernel_size[0] + stride = self.conv.stride[0] + dilation = self.conv.dilation[0] + kernel_size = (kernel_size - 1) * dilation + 1 # effective kernel size with dilations + padding_total = kernel_size - stride + extra_padding = self._get_extra_padding_for_conv1d(hidden_states, kernel_size, stride, padding_total) + + if self.causal: + # Left padding for causal + hidden_states = self._pad1d(hidden_states, (padding_total, extra_padding), mode=self.pad_mode) + else: + # Asymmetric padding required for odd strides + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + hidden_states = self._pad1d( + hidden_states, (padding_left, padding_right + extra_padding), mode=self.pad_mode + ) + + hidden_states = self.conv(hidden_states) + + if self.norm_type == "time_group_norm": + hidden_states = self.norm(hidden_states) + + return hidden_states + + +class EncodecConvTranspose1d(nn.Module): + """ConvTranspose1d with asymmetric or causal padding and normalization.""" + + def __init__(self, config, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1): + super().__init__() + self.causal = config.use_causal_conv + self.trim_right_ratio = config.trim_right_ratio + self.norm_type = config.norm_type + if self.norm_type not in ["weight_norm", "time_group_norm"]: + raise ValueError( + f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}' + ) + + self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride) + if config.norm_type == "weight_norm": + self.conv = nn.utils.weight_norm(self.conv) + elif config.norm_type == "time_group_norm": + self.norm = nn.GroupNorm(1, out_channels) + + if not (self.causal or self.trim_right_ratio == 1.0): + raise ValueError("`trim_right_ratio` != 1.0 only makes sense for causal convolutions") + + def forward(self, hidden_states): + kernel_size = self.conv.kernel_size[0] + stride = self.conv.stride[0] + padding_total = kernel_size - stride + + hidden_states = self.conv(hidden_states) + + if self.norm_type == "time_group_norm": + hidden_states = self.norm(hidden_states) + + # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be + # removed at the very end, when keeping only the right length for the output, + # as removing it here would require also passing the length at the matching layer + # in the encoder. + if self.causal: + # Trim the padding on the right according to the specified ratio + # if trim_right_ratio = 1.0, trim everything from right + padding_right = math.ceil(padding_total * self.trim_right_ratio) + else: + # Asymmetric padding required for odd strides + padding_right = padding_total // 2 + + padding_left = padding_total - padding_right + + # unpad + end = hidden_states.shape[-1] - padding_right + hidden_states = hidden_states[..., padding_left:end] + return hidden_states + + +class EncodecLSTM(nn.Module): + """ + LSTM without worrying about the hidden state, nor the layout of the data. Expects input as convolutional layout. + """ + + def __init__(self, config, dimension): + super().__init__() + self.lstm = nn.LSTM(dimension, dimension, config.num_lstm_layers) + + def forward(self, hidden_states): + hidden_states = hidden_states.permute(2, 0, 1) + hidden_states = self.lstm(hidden_states)[0] + hidden_states + hidden_states = hidden_states.permute(1, 2, 0) + return hidden_states + + +class EncodecResnetBlock(nn.Module): + """ + Residual block from SEANet model as used by EnCodec. + """ + + def __init__(self, config: EncodecConfig, dim: int, dilations: List[int]): + super().__init__() + kernel_sizes = (config.residual_kernel_size, 1) + if len(kernel_sizes) != len(dilations): + raise ValueError("Number of kernel sizes should match number of dilations") + + hidden = dim // config.compress + block = [] + for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)): + in_chs = dim if i == 0 else hidden + out_chs = dim if i == len(kernel_sizes) - 1 else hidden + block += [nn.ELU()] + block += [EncodecConv1d(config, in_chs, out_chs, kernel_size, dilation=dilation)] + self.block = nn.ModuleList(block) + + self.shortcut = EncodecConv1d(config, dim, dim, kernel_size=1) + + def forward(self, hidden_states): + residual = hidden_states + for layer in self.block: + hidden_states = layer(hidden_states) + + return self.shortcut(residual) + hidden_states + + +class EncodecEncoder(nn.Module): + """SEANet encoder as used by EnCodec.""" + + def __init__(self, config: EncodecConfig): + super().__init__() + model = [EncodecConv1d(config, config.audio_channels, config.num_filters, config.kernel_size)] + scaling = 1 + + # Downsample to raw audio scale + for ratio in reversed(config.upsampling_ratios): + current_scale = scaling * config.num_filters + # Add residual layers + for j in range(config.num_residual_layers): + model += [EncodecResnetBlock(config, current_scale, [config.dilation_growth_rate**j, 1])] + # Add downsampling layers + model += [nn.ELU()] + model += [EncodecConv1d(config, current_scale, current_scale * 2, kernel_size=ratio * 2, stride=ratio)] + scaling *= 2 + + model += [EncodecLSTM(config, scaling * config.num_filters)] + model += [nn.ELU()] + model += [EncodecConv1d(config, scaling * config.num_filters, config.hidden_size, config.last_kernel_size)] + + self.layers = nn.ModuleList(model) + + def forward(self, hidden_states): + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +class EncodecDecoder(nn.Module): + """SEANet decoder as used by EnCodec.""" + + def __init__(self, config: EncodecConfig): + super().__init__() + scaling = int(2 ** len(config.upsampling_ratios)) + model = [EncodecConv1d(config, config.hidden_size, scaling * config.num_filters, config.kernel_size)] + + model += [EncodecLSTM(config, scaling * config.num_filters)] + + # Upsample to raw audio scale + for ratio in config.upsampling_ratios: + current_scale = scaling * config.num_filters + # Add upsampling layers + model += [nn.ELU()] + model += [ + EncodecConvTranspose1d(config, current_scale, current_scale // 2, kernel_size=ratio * 2, stride=ratio) + ] + # Add residual layers + for j in range(config.num_residual_layers): + model += [EncodecResnetBlock(config, current_scale // 2, (config.dilation_growth_rate**j, 1))] + scaling //= 2 + + # Add final layers + model += [nn.ELU()] + model += [EncodecConv1d(config, config.num_filters, config.audio_channels, config.last_kernel_size)] + self.layers = nn.ModuleList(model) + + def forward(self, hidden_states): + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +class EncodecEuclideanCodebook(nn.Module): + """Codebook with Euclidean distance.""" + + def __init__(self, config: EncodecConfig): + super().__init__() + embed = torch.zeros(config.codebook_size, config.codebook_dim) + + self.codebook_size = config.codebook_size + + self.register_buffer("inited", torch.Tensor([True])) + self.register_buffer("cluster_size", torch.zeros(config.codebook_size)) + self.register_buffer("embed", embed) + self.register_buffer("embed_avg", embed.clone()) + + def quantize(self, hidden_states): + embed = self.embed.t() + scaled_states = hidden_states.pow(2).sum(1, keepdim=True) + dist = -(scaled_states - 2 * hidden_states @ embed + embed.pow(2).sum(0, keepdim=True)) + embed_ind = dist.max(dim=-1).indices + return embed_ind + + def encode(self, hidden_states): + shape = hidden_states.shape + # pre-process + hidden_states = hidden_states.reshape((-1, shape[-1])) + # quantize + embed_ind = self.quantize(hidden_states) + # post-process + embed_ind = embed_ind.view(*shape[:-1]) + return embed_ind + + def decode(self, embed_ind): + quantize = nn.functional.embedding(embed_ind, self.embed) + return quantize + + +class EncodecVectorQuantization(nn.Module): + """ + Vector quantization implementation. Currently supports only euclidean distance. + """ + + def __init__(self, config: EncodecConfig): + super().__init__() + self.codebook = EncodecEuclideanCodebook(config) + + def encode(self, hidden_states): + hidden_states = hidden_states.permute(0, 2, 1) + embed_in = self.codebook.encode(hidden_states) + return embed_in + + def decode(self, embed_ind): + quantize = self.codebook.decode(embed_ind) + quantize = quantize.permute(0, 2, 1) + return quantize + + +class EncodecResidualVectorQuantizer(nn.Module): + """Residual Vector Quantizer.""" + + def __init__(self, config: EncodecConfig): + super().__init__() + self.codebook_size = config.codebook_size + self.frame_rate = config.frame_rate + self.num_quantizers = config.num_quantizers + self.layers = nn.ModuleList([EncodecVectorQuantization(config) for _ in range(config.num_quantizers)]) + + def get_num_quantizers_for_bandwidth(self, bandwidth: Optional[float] = None) -> int: + """Return num_quantizers based on specified target bandwidth.""" + bw_per_q = math.log2(self.codebook_size) * self.frame_rate + num_quantizers = self.num_quantizers + if bandwidth is not None and bandwidth > 0.0: + num_quantizers = int(max(1, math.floor(bandwidth * 1000 / bw_per_q))) + return num_quantizers + + def encode(self, embeddings: torch.Tensor, bandwidth: Optional[float] = None) -> torch.Tensor: + """ + Encode a given input tensor with the specified frame rate at the given bandwidth. The RVQ encode method sets + the appropriate number of quantizers to use and returns indices for each quantizer. + """ + num_quantizers = self.get_num_quantizers_for_bandwidth(bandwidth) + residual = embeddings + all_indices = [] + for layer in self.layers[:num_quantizers]: + indices = layer.encode(residual) + quantized = layer.decode(indices) + residual = residual - quantized + all_indices.append(indices) + out_indices = torch.stack(all_indices) + return out_indices + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """Decode the given codes to the quantized representation.""" + quantized_out = torch.tensor(0.0, device=codes.device) + for i, indices in enumerate(codes): + layer = self.layers[i] + quantized = layer.decode(indices) + quantized_out = quantized_out + quantized + return quantized_out + + +class EncodecPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = EncodecConfig + base_model_prefix = "encodec" + main_input_name = "input_values" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + elif isinstance(module, nn.Conv1d): + nn.init.kaiming_normal_(module.weight) + if module.bias is not None: + k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0])) + nn.init.uniform_(module.bias, a=-k, b=k) + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, nn.LSTM): + for name, param in module.named_parameters(): + if "weight" in name: + nn.init.xavier_uniform_(param) + elif "bias" in name: + nn.init.constant_(param, 0.0) + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (EncodecEncoder, EncodecDecoder)): + module.gradient_checkpointing = value + + +ENCODEC_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`EncodecConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +ENCODEC_INPUTS_DOCSTRING = r""" + Args: + input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Raw audio input converted to Float and padded to the approriate length in order to be encoded using chunks + of length self.chunk_length and a stride of `config.chunk_stride`. + padding_mask (`torch.BoolTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Mask to avoid computing scaling factors on padding token indices (can we avoid computing conv on these+). + Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + + + `padding_mask` should always be passed, unless the input was truncated or not padded. This is because in + order to process tensors effectively, the input audio should be padded so that `input_length % stride = + step` with `step = chunk_length-stride`. This ensures that all chunks are of the same shape + + + + bandwidth (`float`, *optional*): + The target bandwidth. Must be one of `config.target_bandwidths`. If `None`, uses the smallest possible + bandwidth. bandwidth is represented as a thousandth of what it is, e.g. 6kbps bandwidth is represented as + `bandwidth == 6.0` + audio_codes (`torch.FloatTensor` of shape `(batch_size, nb_chunks, chunk_length)`, *optional*): + Discret code embeddings computed using `model.encode`. + audio_scales (`torch.Tensor` of shape `(batch_size, nb_chunks)`, *optional*): + Scaling factor for each `audio_codes` input. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The EnCodec neural audio codec model.", + ENCODEC_START_DOCSTRING, +) +class EncodecModel(EncodecPreTrainedModel): + def __init__(self, config: EncodecConfig): + super().__init__(config) + self.config = config + + self.encoder = EncodecEncoder(config) + self.decoder = EncodecDecoder(config) + + self.quantizer = EncodecResidualVectorQuantizer(config) + + self.bits_per_codebook = int(math.log2(self.config.codebook_size)) + if 2**self.bits_per_codebook != self.config.codebook_size: + raise ValueError("The codebook_size must be a power of 2.") + + # Initialize weights and apply final processing + self.post_init() + + def get_encoder(self): + return self.encoder + + def get_decoder(self): + return self.decoder + + def _encode_frame( + self, input_values: torch.Tensor, bandwidth: float, padding_mask: int + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Encodes the given input using the underlying VQVAE. If `config.normalize` is set to `True` the input is first + normalized. The padding mask is required to compute the correct scale. + """ + length = input_values.shape[-1] + duration = length / self.config.sampling_rate + + if self.config.chunk_length_s is not None and duration > 1e-5 + self.config.chunk_length_s: + raise RuntimeError(f"Duration of frame ({duration}) is longer than chunk {self.config.chunk_length_s}") + + scale = None + if self.config.normalize: + # if the padding is non zero + input_values = input_values * padding_mask + mono = torch.sum(input_values, 1, keepdim=True) / input_values.shape[1] + scale = mono.pow(2).mean(dim=-1, keepdim=True).sqrt() + 1e-8 + input_values = input_values / scale + + embeddings = self.encoder(input_values) + codes = self.quantizer.encode(embeddings, bandwidth) + codes = codes.transpose(0, 1) + return codes, scale + + def encode( + self, + input_values: torch.Tensor, + padding_mask: torch.Tensor = None, + bandwidth: Optional[float] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor, Optional[torch.Tensor]], EncodecEncoderOutput]: + """ + Encodes the input audio waveform into discrete codes. + + Args: + input_values (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`): + Float values of the input audio waveform. + padding_mask (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`): + Padding mask used to pad the `input_values`. + bandwidth (`float`, *optional*): + The target bandwidth. Must be one of `config.target_bandwidths`. If `None`, uses the smallest possible + bandwidth. bandwidth is represented as a thousandth of what it is, e.g. 6kbps bandwidth is represented + as bandwidth == 6.0 + + Returns: + A list of frames containing the discrete encoded codes for the input audio waveform, along with rescaling + factors for each chunk when `normalize` is True. Each frames is a tuple `(codebook, scale)`, with + `codebook` of shape `[batch_size, num_codebooks, frames]`. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if bandwidth is None: + bandwidth = self.config.target_bandwidths[0] + if bandwidth not in self.config.target_bandwidths: + raise ValueError( + f"This model doesn't support the bandwidth {bandwidth}. " + f"Select one of {self.config.target_bandwidths}." + ) + + _, channels, input_length = input_values.shape + + if channels < 1 or channels > 2: + raise ValueError(f"Number of audio channels must be 1 or 2, but got {channels}") + + chunk_length = self.config.chunk_length + if chunk_length is None: + chunk_length = input_length + stride = input_length + else: + stride = self.config.chunk_stride + + if padding_mask is None: + padding_mask = torch.ones_like(input_values).bool() + + encoded_frames = [] + scales = [] + + step = chunk_length - stride + if (input_length % stride) - step != 0: + raise ValueError( + "The input length is not properly padded for batched chunked decoding. Make sure to pad the input correctly." + ) + + for offset in range(0, input_length - step, stride): + mask = padding_mask[..., offset : offset + chunk_length].bool() + frame = input_values[:, :, offset : offset + chunk_length] + encoded_frame, scale = self._encode_frame(frame, bandwidth, mask) + encoded_frames.append(encoded_frame) + scales.append(scale) + + encoded_frames = torch.stack(encoded_frames) + + if not return_dict: + return (encoded_frames, scales) + + return EncodecEncoderOutput(encoded_frames, scales) + + @staticmethod + def _linear_overlap_add(frames: List[torch.Tensor], stride: int): + # Generic overlap add, with linear fade-in/fade-out, supporting complex scenario + # e.g., more than 2 frames per position. + # The core idea is to use a weight function that is a triangle, + # with a maximum value at the middle of the chunk. + # We use this weighting when summing the frames, and divide by the sum of weights + # for each positions at the end. Thus: + # - if a frame is the only one to cover a position, the weighting is a no-op. + # - if 2 frames cover a position: + # ... ... + # / \/ \ + # / /\ \ + # S T , i.e. S offset of second frame starts, T end of first frame. + # Then the weight function for each one is: (t - S), (T - t), with `t` a given offset. + # After the final normalization, the weight of the second frame at position `t` is + # (t - S) / (t - S + (T - t)) = (t - S) / (T - S), which is exactly what we want. + # + # - if more than 2 frames overlap at a given point, we hope that by induction + # something sensible happens. + if len(frames) == 0: + raise ValueError("`frames` cannot be an empty list.") + + device = frames[0].device + dtype = frames[0].dtype + shape = frames[0].shape[:-1] + total_size = stride * (len(frames) - 1) + frames[-1].shape[-1] + + frame_length = frames[0].shape[-1] + time_vec = torch.linspace(0, 1, frame_length + 2, device=device, dtype=dtype)[1:-1] + weight = 0.5 - (time_vec - 0.5).abs() + + sum_weight = torch.zeros(total_size, device=device, dtype=dtype) + out = torch.zeros(*shape, total_size, device=device, dtype=dtype) + offset: int = 0 + + for frame in frames: + frame_length = frame.shape[-1] + out[..., offset : offset + frame_length] += weight[:frame_length] * frame + sum_weight[offset : offset + frame_length] += weight[:frame_length] + offset += stride + + if sum_weight.min() == 0: + raise ValueError(f"`sum_weight` minimum element must be bigger than zero: {sum_weight}`") + + return out / sum_weight + + def _decode_frame(self, codes: torch.Tensor, scale: Optional[torch.Tensor] = None) -> torch.Tensor: + codes = codes.transpose(0, 1) + embeddings = self.quantizer.decode(codes) + outputs = self.decoder(embeddings) + if scale is not None: + outputs = outputs * scale.view(-1, 1, 1) + return outputs + + def decode( + self, + audio_codes: torch.Tensor, + audio_scales: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], EncodecDecoderOutput]: + """ + Decodes the given frames into an output audio waveform. + + Note that the output might be a bit bigger than the input. In that case, any extra steps at the end can be + trimmed. + + Args: + audio_codes (`torch.FloatTensor` of shape `(batch_size, nb_chunks, chunk_length)`, *optional*): + Discret code embeddings computed using `model.encode`. + audio_scales (`torch.Tensor` of shape `(batch_size, nb_chunks)`, *optional*): + Scaling factor for each `audio_codes` input. + padding_mask (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`): + Padding mask used to pad the `input_values`. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + + """ + return_dict = return_dict or self.config.return_dict + + chunk_length = self.config.chunk_length + if chunk_length is None: + if len(audio_codes) != 1: + raise ValueError(f"Expected one frame, got {len(audio_codes)}") + audio_values = self._decode_frame(audio_codes[0], audio_scales[0]) + else: + decoded_frames = [] + + for frame, scale in zip(audio_codes, audio_scales): + frames = self._decode_frame(frame, scale) + decoded_frames.append(frames) + + audio_values = self._linear_overlap_add(decoded_frames, self.config.chunk_stride or 1) + + # truncate based on padding mask + if padding_mask is not None and padding_mask.shape[-1] < audio_values.shape[-1]: + audio_values = audio_values[..., : padding_mask.shape[-1]] + + if not return_dict: + return (audio_values,) + return EncodecDecoderOutput(audio_values) + + @add_start_docstrings_to_model_forward(ENCODEC_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=EncodecOutput, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_values: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + bandwidth: Optional[float] = None, + audio_codes: Optional[torch.Tensor] = None, + audio_scales: Optional[torch.Tensor] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], EncodecOutput]: + r""" + Returns: + + Examples: + + ```python + >>> from datasets import load_dataset + >>> from transformers import AutoProcessor, EncodecModel + + >>> dataset = load_dataset("ashraq/esc50") + >>> audio_sample = dataset["train"]["audio"][0]["array"] + + >>> model_id = "facebook/encodec_24khz" + >>> model = EncodecModel.from_pretrained(model_id) + >>> processor = AutoProcessor.from_pretrained(model_id) + + >>> inputs = processor(raw_audio=audio_sample, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> audio_codes = outputs.audio_codes + >>> audio_values = outputs.audio_values + ```""" + return_dict = return_dict or self.config.return_dict + + if padding_mask is None: + padding_mask = torch.ones_like(input_values).bool() + + if audio_codes is not None and audio_scales is None: + raise ValueError("You specified `audio_codes` but did not specify the `audio_scales`") + + if audio_scales is not None and audio_codes is None: + raise ValueError("You specified `audio_scales` but did not specify the `audio_codes`") + + if audio_scales is None and audio_codes is None: + audio_codes, audio_scales = self.encode(input_values, padding_mask, bandwidth, False) + + audio_values = self.decode(audio_codes, audio_scales, padding_mask, return_dict=return_dict)[0] + if not return_dict: + return (audio_codes, audio_values) + + return EncodecOutput(audio_codes=audio_codes, audio_values=audio_values) diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py index 1eba8c0ca7..11e763e2a0 100644 --- a/src/transformers/utils/dummy_pt_objects.py +++ b/src/transformers/utils/dummy_pt_objects.py @@ -2655,6 +2655,23 @@ def load_tf_weights_in_electra(*args, **kwargs): requires_backends(load_tf_weights_in_electra, ["torch"]) +ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST = None + + +class EncodecModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EncodecPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + class EncoderDecoderModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/transformers/utils/dummy_speech_objects.py b/src/transformers/utils/dummy_speech_objects.py index f85182c8bc..6c59158858 100644 --- a/src/transformers/utils/dummy_speech_objects.py +++ b/src/transformers/utils/dummy_speech_objects.py @@ -9,6 +9,13 @@ class ASTFeatureExtractor(metaclass=DummyObject): requires_backends(self, ["speech"]) +class EncodecFeatureExtractor(metaclass=DummyObject): + _backends = ["speech"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["speech"]) + + class MCTCTFeatureExtractor(metaclass=DummyObject): _backends = ["speech"] diff --git a/tests/models/encodec/__init__.py b/tests/models/encodec/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/models/encodec/test_feature_extraction_encodec.py b/tests/models/encodec/test_feature_extraction_encodec.py new file mode 100644 index 0000000000..95639fcda5 --- /dev/null +++ b/tests/models/encodec/test_feature_extraction_encodec.py @@ -0,0 +1,255 @@ +# coding=utf-8 +# Copyright 2021-2023 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the EnCodec feature extractor.""" + +import itertools +import random +import unittest + +import numpy as np + +from transformers import is_speech_available +from transformers.testing_utils import require_torch +from transformers.utils.import_utils import is_torch_available + +from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin + + +if is_speech_available(): + from transformers import EncodecFeatureExtractor + +if is_torch_available(): + import torch + + +global_rng = random.Random() + + +def floats_list(shape, scale=1.0, rng=None, name=None): + """Creates a random float32 tensor""" + if rng is None: + rng = global_rng + + values = [] + for batch_idx in range(shape[0]): + values.append([]) + for _ in range(shape[1]): + values[-1].append(rng.random() * scale) + + return values + + +@require_torch +class EnCodecFeatureExtractionTester(unittest.TestCase): + def __init__( + self, + parent, + batch_size=7, + min_seq_length=400, + max_seq_length=2000, + feature_size=1, + padding_value=0.0, + sampling_rate=24000, + return_attention_mask=True, + ): + self.parent = parent + self.batch_size = batch_size + self.min_seq_length = min_seq_length + self.max_seq_length = max_seq_length + self.seq_length_diff = (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1) + self.feature_size = feature_size + self.padding_value = padding_value + self.sampling_rate = sampling_rate + self.return_attention_mask = return_attention_mask + + def prepare_feat_extract_dict(self): + return { + "feature_size": self.feature_size, + "padding_value": self.padding_value, + "sampling_rate": self.sampling_rate, + "return_attention_mask": self.return_attention_mask, + } + + def prepare_inputs_for_common(self, equal_length=False, numpify=False): + def _flatten(list_of_lists): + return list(itertools.chain(*list_of_lists)) + + if equal_length: + audio_inputs = floats_list((self.batch_size, self.max_seq_length)) + else: + # make sure that inputs increase in size + audio_inputs = [ + _flatten(floats_list((x, self.feature_size))) + for x in range(self.min_seq_length, self.max_seq_length, self.seq_length_diff) + ] + + if numpify: + audio_inputs = [np.asarray(x) for x in audio_inputs] + + return audio_inputs + + +@require_torch +class EnCodecFeatureExtractionTest(SequenceFeatureExtractionTestMixin, unittest.TestCase): + feature_extraction_class = EncodecFeatureExtractor if is_speech_available() else None + + def setUp(self): + self.feat_extract_tester = EnCodecFeatureExtractionTester(self) + + def test_call(self): + # Tests that all call wrap to encode_plus and batch_encode_plus + feat_extract = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()) + # create three inputs of length 800, 1000, and 1200 + audio_inputs = [floats_list((1, x))[0] for x in range(800, 1400, 200)] + np_audio_inputs = [np.asarray(audio_input) for audio_input in audio_inputs] + + # Test not batched input + encoded_sequences_1 = feat_extract(audio_inputs[0], return_tensors="np").input_values + encoded_sequences_2 = feat_extract(np_audio_inputs[0], return_tensors="np").input_values + self.assertTrue(np.allclose(encoded_sequences_1, encoded_sequences_2, atol=1e-3)) + + # Test batched + encoded_sequences_1 = feat_extract(audio_inputs, padding=True, return_tensors="np").input_values + encoded_sequences_2 = feat_extract(np_audio_inputs, padding=True, return_tensors="np").input_values + for enc_seq_1, enc_seq_2 in zip(encoded_sequences_1, encoded_sequences_2): + self.assertTrue(np.allclose(enc_seq_1, enc_seq_2, atol=1e-3)) + + def test_double_precision_pad(self): + feature_extractor = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()) + np_audio_inputs = np.random.rand(100).astype(np.float64) + py_audio_inputs = np_audio_inputs.tolist() + + for inputs in [py_audio_inputs, np_audio_inputs]: + np_processed = feature_extractor.pad([{"input_values": inputs}], return_tensors="np") + self.assertTrue(np_processed.input_values.dtype == np.float32) + pt_processed = feature_extractor.pad([{"input_values": inputs}], return_tensors="pt") + self.assertTrue(pt_processed.input_values.dtype == torch.float32) + + def _load_datasamples(self, num_samples): + from datasets import load_dataset + + ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + # automatic decoding with librispeech + audio_samples = ds.sort("id").select(range(num_samples))[:num_samples]["audio"] + + return [x["array"] for x in audio_samples] + + def test_integration(self): + # fmt: off + EXPECTED_INPUT_VALUES = torch.tensor( + [2.3804e-03, 2.0752e-03, 1.9836e-03, 2.1057e-03, 1.6174e-03, + 3.0518e-04, 9.1553e-05, 3.3569e-04, 9.7656e-04, 1.8311e-03, + 2.0142e-03, 2.1057e-03, 1.7395e-03, 4.5776e-04, -3.9673e-04, + 4.5776e-04, 1.0071e-03, 9.1553e-05, 4.8828e-04, 1.1597e-03, + 7.3242e-04, 9.4604e-04, 1.8005e-03, 1.8311e-03, 8.8501e-04, + 4.2725e-04, 4.8828e-04, 7.3242e-04, 1.0986e-03, 2.1057e-03] + ) + # fmt: on + input_audio = self._load_datasamples(1) + feature_extractor = EncodecFeatureExtractor() + input_values = feature_extractor(input_audio, return_tensors="pt").input_values + self.assertEquals(input_values.shape, (1, 1, 93680)) + self.assertTrue(torch.allclose(input_values[0, 0, :30], EXPECTED_INPUT_VALUES, atol=1e-6)) + + def test_integration_stereo(self): + # fmt: off + EXPECTED_INPUT_VALUES = torch.tensor( + [2.3804e-03, 2.0752e-03, 1.9836e-03, 2.1057e-03, 1.6174e-03, + 3.0518e-04, 9.1553e-05, 3.3569e-04, 9.7656e-04, 1.8311e-03, + 2.0142e-03, 2.1057e-03, 1.7395e-03, 4.5776e-04, -3.9673e-04, + 4.5776e-04, 1.0071e-03, 9.1553e-05, 4.8828e-04, 1.1597e-03, + 7.3242e-04, 9.4604e-04, 1.8005e-03, 1.8311e-03, 8.8501e-04, + 4.2725e-04, 4.8828e-04, 7.3242e-04, 1.0986e-03, 2.1057e-03] + ) + # fmt: on + input_audio = self._load_datasamples(1) + input_audio = [np.tile(input_audio[0][None], reps=(2, 1))] + input_audio[0][1] *= 0.5 + feature_extractor = EncodecFeatureExtractor(feature_size=2) + input_values = feature_extractor(input_audio, return_tensors="pt").input_values + self.assertEquals(input_values.shape, (1, 2, 93680)) + self.assertTrue(torch.allclose(input_values[0, 0, :30], EXPECTED_INPUT_VALUES, atol=1e-6)) + self.assertTrue(torch.allclose(input_values[0, 1, :30], EXPECTED_INPUT_VALUES * 0.5, atol=1e-6)) + + def test_truncation_and_padding(self): + input_audio = self._load_datasamples(2) + # would be easier if the stride was like + feature_extractor = EncodecFeatureExtractor(feature_size=1, chunk_length_s=1, overlap=0.01) + + # pad and trunc raise an error ? + with self.assertRaisesRegex( + ValueError, + "^Both padding and truncation were set. Make sure you only set one.$", + ): + truncated_outputs = feature_extractor( + input_audio, padding="max_length", truncation=True, return_tensors="pt" + ).input_values + + # truncate to chunk + truncated_outputs = feature_extractor(input_audio, truncation=True, return_tensors="pt").input_values + self.assertEquals(truncated_outputs.shape, (2, 1, 71520)) # 2 chunks + + # force truncate to max_length + truncated_outputs = feature_extractor( + input_audio, truncation=True, max_length=48000, return_tensors="pt" + ).input_values + self.assertEquals(truncated_outputs.shape, (2, 1, 48000)) + + # pad to chunk + padded_outputs = feature_extractor(input_audio, padding=True, return_tensors="pt").input_values + self.assertEquals(padded_outputs.shape, (2, 1, 95280)) + + # pad to chunk + truncated_outputs = feature_extractor(input_audio, return_tensors="pt").input_values + self.assertEquals(truncated_outputs.shape, (2, 1, 95280)) + + # force pad to max length + truncated_outputs = feature_extractor( + input_audio, padding="max_length", max_length=100000, return_tensors="pt" + ).input_values + self.assertEquals(truncated_outputs.shape, (2, 1, 100000)) + + # force no pad + with self.assertRaisesRegex( + ValueError, + "^Unable to create tensor, you should probably activate padding with 'padding=True' to have batched tensors with the same length.$", + ): + truncated_outputs = feature_extractor(input_audio, padding=False, return_tensors="pt").input_values + + truncated_outputs = feature_extractor(input_audio[0], padding=False, return_tensors="pt").input_values + self.assertEquals(truncated_outputs.shape, (1, 1, 93680)) + + # no pad if no chunk_length_s + feature_extractor.chunk_length_s = None + with self.assertRaisesRegex( + ValueError, + "^Unable to create tensor, you should probably activate padding with 'padding=True' to have batched tensors with the same length.$", + ): + truncated_outputs = feature_extractor(input_audio, padding=False, return_tensors="pt").input_values + + truncated_outputs = feature_extractor(input_audio[0], padding=False, return_tensors="pt").input_values + self.assertEquals(truncated_outputs.shape, (1, 1, 93680)) + + # no pad if no overlap + feature_extractor.chunk_length_s = 2 + feature_extractor.overlap = None + with self.assertRaisesRegex( + ValueError, + "^Unable to create tensor, you should probably activate padding with 'padding=True' to have batched tensors with the same length.$", + ): + truncated_outputs = feature_extractor(input_audio, padding=False, return_tensors="pt").input_values + + truncated_outputs = feature_extractor(input_audio[0], padding=False, return_tensors="pt").input_values + self.assertEquals(truncated_outputs.shape, (1, 1, 93680)) diff --git a/tests/models/encodec/test_modeling_encodec.py b/tests/models/encodec/test_modeling_encodec.py new file mode 100644 index 0000000000..23b2114a5d --- /dev/null +++ b/tests/models/encodec/test_modeling_encodec.py @@ -0,0 +1,571 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" Testing suite for the PyTorch Encodec model. """ + +import copy +import inspect +import os +import tempfile +import unittest +from typing import Dict, List, Tuple + +import numpy as np +from datasets import Audio, load_dataset + +from transformers import AutoProcessor, EncodecConfig +from transformers.testing_utils import ( + is_torch_available, + require_torch, + slow, + torch_device, +) + +from ...test_configuration_common import ConfigTester +from ...test_modeling_common import ( + ModelTesterMixin, + _config_zero_init, + floats_tensor, +) +from ...test_pipeline_mixin import PipelineTesterMixin + + +if is_torch_available(): + import torch + + from transformers import EncodecModel + + +def prepare_inputs_dict( + config, + input_ids=None, + input_values=None, + decoder_input_ids=None, + attention_mask=None, + decoder_attention_mask=None, + head_mask=None, + decoder_head_mask=None, + cross_attn_head_mask=None, +): + if input_ids is not None: + encoder_dict = {"input_ids": input_ids} + else: + encoder_dict = {"input_values": input_values} + + decoder_dict = {"decoder_input_ids": decoder_input_ids} if decoder_input_ids is not None else {} + + return {**encoder_dict, **decoder_dict} + + +@require_torch +class EncodecModelTester: + def __init__( + self, + parent, + batch_size=13, + num_channels=2, + is_training=False, + num_hidden_layers=4, + intermediate_size=40, + ): + self.parent = parent + self.batch_size = batch_size + self.num_channels = num_channels + self.is_training = is_training + + self.num_hidden_layers = num_hidden_layers + self.intermediate_size = intermediate_size + + def prepare_config_and_inputs(self): + input_values = floats_tensor([self.batch_size, self.num_channels, self.intermediate_size], scale=1.0) + config = self.get_config() + inputs_dict = {"input_values": input_values} + return config, inputs_dict + + def prepare_config_and_inputs_for_common(self): + config, inputs_dict = self.prepare_config_and_inputs() + return config, inputs_dict + + def get_config(self): + return EncodecConfig(audio_channels=self.num_channels, chunk_in_sec=None) + + def create_and_check_model_forward(self, config, inputs_dict): + model = EncodecModel(config=config).to(torch_device).eval() + + input_values = inputs_dict["input_values"] + result = model(input_values) + self.parent.assertEqual( + result.audio_values.shape, (self.batch_size, self.num_channels, self.intermediate_size) + ) + + +@require_torch +class EncodecModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): + all_model_classes = (EncodecModel,) if is_torch_available() else () + is_encoder_decoder = True + test_pruning = False + test_headmasking = False + test_resize_embeddings = False + pipeline_model_mapping = {} + input_name = "input_values" + + def _prepare_for_class(self, inputs_dict, model_class, return_labels=False): + # model does not have attention and does not support returning hidden states + inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels) + if "output_attentions" in inputs_dict: + inputs_dict.pop("output_attentions") + if "output_hidden_states" in inputs_dict: + inputs_dict.pop("output_hidden_states") + return inputs_dict + + def setUp(self): + self.model_tester = EncodecModelTester(self) + self.config_tester = ConfigTester( + self, config_class=EncodecConfig, hidden_size=37, common_properties=[], has_text_modality=False + ) + + def test_config(self): + self.config_tester.run_common_tests() + + def test_model_forward(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_model_forward(*config_and_inputs) + + def test_forward_signature(self): + config, _ = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + signature = inspect.signature(model.forward) + # signature.parameters is an OrderedDict => so arg_names order is deterministic + arg_names = [*signature.parameters.keys()] + + expected_arg_names = ["input_values", "padding_mask", "bandwidth"] + self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names) + + @unittest.skip("The EncodecModel is not transformers based, thus it does not have `inputs_embeds` logics") + def test_inputs_embeds(self): + pass + + @unittest.skip("The EncodecModel is not transformers based, thus it does not have `inputs_embeds` logics") + def test_model_common_attributes(self): + pass + + @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `attention` logic") + def test_retain_grad_hidden_states_attentions(self): + pass + + @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `attention` logic") + def test_torchscript_output_attentions(self): + pass + + @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `hidden_states` logic") + def test_torchscript_output_hidden_state(self): + pass + + def _create_and_check_torchscript(self, config, inputs_dict): + if not self.test_torchscript: + return + + configs_no_init = _config_zero_init(config) # To be sure we have no Nan + configs_no_init.torchscript = True + configs_no_init.return_dict = False + for model_class in self.all_model_classes: + model = model_class(config=configs_no_init) + model.to(torch_device) + model.eval() + inputs = self._prepare_for_class(inputs_dict, model_class) + + main_input_name = model_class.main_input_name + + try: + main_input = inputs[main_input_name] + model(main_input) + traced_model = torch.jit.trace(model, main_input) + except RuntimeError: + self.fail("Couldn't trace module.") + + with tempfile.TemporaryDirectory() as tmp_dir_name: + pt_file_name = os.path.join(tmp_dir_name, "traced_model.pt") + + try: + torch.jit.save(traced_model, pt_file_name) + except Exception: + self.fail("Couldn't save module.") + + try: + loaded_model = torch.jit.load(pt_file_name) + except Exception: + self.fail("Couldn't load module.") + + model.to(torch_device) + model.eval() + + loaded_model.to(torch_device) + loaded_model.eval() + + model_state_dict = model.state_dict() + loaded_model_state_dict = loaded_model.state_dict() + + non_persistent_buffers = {} + for key in loaded_model_state_dict.keys(): + if key not in model_state_dict.keys(): + non_persistent_buffers[key] = loaded_model_state_dict[key] + + loaded_model_state_dict = { + key: value for key, value in loaded_model_state_dict.items() if key not in non_persistent_buffers + } + + self.assertEqual(set(model_state_dict.keys()), set(loaded_model_state_dict.keys())) + + model_buffers = list(model.buffers()) + for non_persistent_buffer in non_persistent_buffers.values(): + found_buffer = False + for i, model_buffer in enumerate(model_buffers): + if torch.equal(non_persistent_buffer, model_buffer): + found_buffer = True + break + + self.assertTrue(found_buffer) + model_buffers.pop(i) + + models_equal = True + for layer_name, p1 in model_state_dict.items(): + if layer_name in loaded_model_state_dict: + p2 = loaded_model_state_dict[layer_name] + if p1.data.ne(p2.data).sum() > 0: + models_equal = False + + self.assertTrue(models_equal) + + # Avoid memory leak. Without this, each call increase RAM usage by ~20MB. + # (Even with this call, there are still memory leak by ~0.04MB) + self.clear_torch_jit_class_registry() + + @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `attention` logic") + def test_attention_outputs(self): + pass + + def test_feed_forward_chunking(self): + (original_config, inputs_dict) = self.model_tester.prepare_config_and_inputs_for_common() + for model_class in self.all_model_classes: + torch.manual_seed(0) + config = copy.deepcopy(original_config) + config.chunk_length_s = None + config.overlap = None + config.sampling_rate = 10 + + model = model_class(config) + model.to(torch_device) + model.eval() + inputs = self._prepare_for_class(inputs_dict, model_class) + inputs["input_values"] = inputs["input_values"].repeat(1, 1, 10) + + hidden_states_no_chunk = model(**inputs)[0] + + torch.manual_seed(0) + config.chunk_length_s = 1 + config.overlap = 0 + config.sampling_rate = 10 + + model = model_class(config) + model.to(torch_device) + model.eval() + + hidden_states_with_chunk = model(**inputs)[0] + self.assertTrue(torch.allclose(hidden_states_no_chunk, hidden_states_with_chunk, atol=1e-3)) + + @unittest.skip("The EncodecModel is not transformers based, thus it does not have the usual `hidden_states` logic") + def test_hidden_states_output(self): + pass + + def test_determinism(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + def check_determinism(first, second): + # outputs are not tensors but list (since each sequence don't have the same frame_length) + out_1 = first.cpu().numpy() + out_2 = second.cpu().numpy() + out_1 = out_1[~np.isnan(out_1)] + out_2 = out_2[~np.isnan(out_2)] + max_diff = np.amax(np.abs(out_1 - out_2)) + self.assertLessEqual(max_diff, 1e-5) + + for model_class in self.all_model_classes: + model = model_class(config) + model.to(torch_device) + model.eval() + with torch.no_grad(): + first = model(**self._prepare_for_class(inputs_dict, model_class))[0] + second = model(**self._prepare_for_class(inputs_dict, model_class))[0] + + if isinstance(first, tuple) and isinstance(second, tuple): + for tensor1, tensor2 in zip(first, second): + check_determinism(tensor1, tensor2) + else: + check_determinism(first, second) + + def test_model_outputs_equivalence(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + def set_nan_tensor_to_zero(t): + t[t != t] = 0 + return t + + def check_equivalence(model, tuple_inputs, dict_inputs, additional_kwargs={}): + with torch.no_grad(): + tuple_output = model(**tuple_inputs, return_dict=False, **additional_kwargs) + dict_output = model(**dict_inputs, return_dict=True, **additional_kwargs) + + def recursive_check(tuple_object, dict_object): + if isinstance(tuple_object, (List, Tuple)): + for tuple_iterable_value, dict_iterable_value in zip(tuple_object, dict_object): + recursive_check(tuple_iterable_value, dict_iterable_value) + elif isinstance(tuple_object, Dict): + for tuple_iterable_value, dict_iterable_value in zip( + tuple_object.values(), dict_object.values() + ): + recursive_check(tuple_iterable_value, dict_iterable_value) + elif tuple_object is None: + return + else: + self.assertTrue( + torch.allclose( + set_nan_tensor_to_zero(tuple_object), set_nan_tensor_to_zero(dict_object), atol=1e-5 + ), + msg=( + "Tuple and dict output are not equal. Difference:" + f" {torch.max(torch.abs(tuple_object - dict_object))}. Tuple has `nan`:" + f" {torch.isnan(tuple_object).any()} and `inf`: {torch.isinf(tuple_object)}. Dict has" + f" `nan`: {torch.isnan(dict_object).any()} and `inf`: {torch.isinf(dict_object)}." + ), + ) + + recursive_check(tuple_output, dict_output) + + for model_class in self.all_model_classes: + model = model_class(config) + model.to(torch_device) + model.eval() + + tuple_inputs = self._prepare_for_class(inputs_dict, model_class) + dict_inputs = self._prepare_for_class(inputs_dict, model_class) + check_equivalence(model, tuple_inputs, dict_inputs) + + def test_initialization(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + configs_no_init = _config_zero_init(config) + for model_class in self.all_model_classes: + model = model_class(config=configs_no_init) + for name, param in model.named_parameters(): + uniform_init_parms = ["conv"] + ignore_init = ["lstm"] + if param.requires_grad: + if any([x in name for x in uniform_init_parms]): + self.assertTrue( + -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0, + msg=f"Parameter {name} of model {model_class} seems not properly initialized", + ) + elif not any([x in name for x in ignore_init]): + self.assertIn( + ((param.data.mean() * 1e9).round() / 1e9).item(), + [0.0, 1.0], + msg=f"Parameter {name} of model {model_class} seems not properly initialized", + ) + + +def normalize(arr): + norm = np.linalg.norm(arr) + normalized_arr = arr / norm + return normalized_arr + + +def compute_rmse(arr1, arr2): + arr1_normalized = normalize(arr1) + arr2_normalized = normalize(arr2) + return np.sqrt(((arr1_normalized - arr2_normalized) ** 2).mean()) + + +@slow +@require_torch +class EncodecIntegrationTest(unittest.TestCase): + def test_integration_24kHz(self): + expected_rmse = { + "1.5": 0.0025, + "24.0": 0.0015, + } + expected_codesums = { + "1.5": [367184], + "24.0": [6648961], + } + librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + model_id = "facebook/encodec_24khz" + + model = EncodecModel.from_pretrained(model_id).to(torch_device) + processor = AutoProcessor.from_pretrained(model_id) + + librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate)) + audio_sample = librispeech_dummy[-1]["audio"]["array"] + + inputs = processor( + raw_audio=audio_sample, + sampling_rate=processor.sampling_rate, + return_tensors="pt", + ).to(torch_device) + + for bandwidth, expected_rmse in expected_rmse.items(): + with torch.no_grad(): + # use max bandwith for best possible reconstruction + encoder_outputs = model.encode(inputs["input_values"], bandwidth=float(bandwidth)) + + audio_code_sums = [a[0].sum().cpu().item() for a in encoder_outputs[0]] + + # make sure audio encoded codes are correct + self.assertListEqual(audio_code_sums, expected_codesums[bandwidth]) + + audio_codes, scales = encoder_outputs.to_tuple() + input_values_dec = model.decode(audio_codes, scales, inputs["padding_mask"])[0] + input_values_enc_dec = model( + inputs["input_values"], inputs["padding_mask"], bandwidth=float(bandwidth) + )[-1] + + # make sure forward and decode gives same result + self.assertTrue(torch.allclose(input_values_dec, input_values_enc_dec, atol=1e-3)) + + # make sure shape matches + self.assertTrue(inputs["input_values"].shape == input_values_enc_dec.shape) + + arr = inputs["input_values"][0].cpu().numpy() + arr_enc_dec = input_values_enc_dec[0].cpu().numpy() + + # make sure audios are more or less equal + # the RMSE of two random gaussian noise vectors with ~N(0, 1) is around 1.0 + rmse = compute_rmse(arr, arr_enc_dec) + self.assertTrue(rmse < expected_rmse) + + def test_integration_48kHz(self): + expected_rmse = { + "3.0": 0.001, + "24.0": 0.0005, + } + expected_codesums = { + "3.0": [142174, 147901, 154090, 178965, 161879], + "24.0": [1561048, 1284593, 1278330, 1487220, 1659404], + } + librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + model_id = "facebook/encodec_48khz" + + model = EncodecModel.from_pretrained(model_id).to(torch_device) + model = model.eval() + processor = AutoProcessor.from_pretrained(model_id) + + librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate)) + audio_sample = librispeech_dummy[-1]["audio"]["array"] + + # transform mono to stereo + audio_sample = np.array([audio_sample, audio_sample]) + + inputs = processor(raw_audio=audio_sample, sampling_rate=processor.sampling_rate, return_tensors="pt").to( + torch_device + ) + + for bandwidth, expected_rmse in expected_rmse.items(): + with torch.no_grad(): + # use max bandwith for best possible reconstruction + encoder_outputs = model.encode( + inputs["input_values"], inputs["padding_mask"], bandwidth=float(bandwidth), return_dict=False + ) + audio_code_sums = [a[0].sum().cpu().item() for a in encoder_outputs[0]] + + # make sure audio encoded codes are correct + self.assertListEqual(audio_code_sums, expected_codesums[bandwidth]) + audio_codes, scales = encoder_outputs + input_values_dec = model.decode(audio_codes, scales, inputs["padding_mask"])[0] + input_values_enc_dec = model( + inputs["input_values"], inputs["padding_mask"], bandwidth=float(bandwidth) + )[-1] + + # make sure forward and decode gives same result + self.assertTrue(torch.allclose(input_values_dec, input_values_enc_dec, atol=1e-3)) + + # make sure shape matches + self.assertTrue(inputs["input_values"].shape == input_values_enc_dec.shape) + + arr = inputs["input_values"][0].cpu().numpy() + arr_enc_dec = input_values_enc_dec[0].cpu().numpy() + + # make sure audios are more or less equal + # the RMSE of two random gaussian noise vectors with ~N(0, 1) is around 1.0 + rmse = compute_rmse(arr, arr_enc_dec) + self.assertTrue(rmse < expected_rmse) + + def test_batch_48kHz(self): + expected_rmse = { + "3.0": 0.001, + "24.0": 0.0005, + } + expected_codesums = { + "3.0": [ + [71689, 78549, 75644, 88889, 73100, 82509, 71449, 82835], + [84427, 82356, 75809, 52509, 80137, 87672, 87436, 70456], + ], + "24.0": [ + [71689, 78549, 75644, 88889, 73100, 82509, 71449, 82835], + [84427, 82356, 75809, 52509, 80137, 87672, 87436, 70456], + ], + } + librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + model_id = "facebook/encodec_48khz" + + model = EncodecModel.from_pretrained(model_id).to(torch_device) + processor = AutoProcessor.from_pretrained(model_id, chunk_length_s=1, overlap=0.01) + + librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate)) + + audio_samples = [ + np.array([audio_sample["array"], audio_sample["array"]]) + for audio_sample in librispeech_dummy[-2:]["audio"] + ] + + inputs = processor(raw_audio=audio_samples, sampling_rate=processor.sampling_rate, return_tensors="pt") + input_values = inputs["input_values"].to(torch_device) + for bandwidth, expected_rmse in expected_rmse.items(): + with torch.no_grad(): + # use max bandwith for best possible reconstruction + encoder_outputs = model.encode(input_values, bandwidth=float(bandwidth), return_dict=False) + audio_code_sums_0 = [a[0][0].sum().cpu().item() for a in encoder_outputs[0]] + audio_code_sums_1 = [a[0][1].sum().cpu().item() for a in encoder_outputs[0]] + + # make sure audio encoded codes are correct + self.assertListEqual(audio_code_sums_0, expected_codesums[bandwidth][0]) + self.assertListEqual(audio_code_sums_1, expected_codesums[bandwidth][1]) + + audio_codes, scales = encoder_outputs + input_values_dec = model.decode(audio_codes, scales)[0] + input_values_enc_dec = model(input_values, bandwidth=float(bandwidth))[-1] + + # make sure forward and decode gives same result + self.assertTrue(torch.allclose(input_values_dec, input_values_enc_dec, atol=1e-3)) + + # make sure shape matches + self.assertTrue(input_values.shape == input_values_enc_dec.shape) + + arr = input_values[0].cpu().numpy() + arr_enc_dec = input_values_enc_dec[0].cpu().numpy() + + # make sure audios are more or less equal + # the RMSE of two random gaussian noise vectors with ~N(0, 1) is around 1.0 + rmse = compute_rmse(arr, arr_enc_dec) + self.assertTrue(rmse < expected_rmse) diff --git a/utils/check_config_attributes.py b/utils/check_config_attributes.py index 02c3d2276f..63b1bacbbe 100644 --- a/utils/check_config_attributes.py +++ b/utils/check_config_attributes.py @@ -31,6 +31,8 @@ transformers = direct_transformers_import(PATH_TO_TRANSFORMERS) CONFIG_MAPPING = transformers.models.auto.configuration_auto.CONFIG_MAPPING SPECIAL_CASES_TO_ALLOW = { + # used to compute the property `self.chunk_length` + "EncodecConfig": ["overlap"], # used as `self.bert_model = BertModel(config, ...)` "DPRConfig": True, # not used in modeling files, but it's an important information diff --git a/utils/documentation_tests.txt b/utils/documentation_tests.txt index 5cdded64bb..63eb4a0002 100644 --- a/utils/documentation_tests.txt +++ b/utils/documentation_tests.txt @@ -83,6 +83,7 @@ src/transformers/models/electra/configuration_electra.py src/transformers/models/electra/modeling_electra.py src/transformers/models/electra/modeling_tf_electra.py src/transformers/models/efficientformer/modeling_tf_efficientformer.py +src/transformers/models/encodec/modeling_encodec.py src/transformers/models/ernie/configuration_ernie.py src/transformers/models/ernie_m/configuration_ernie_m.py src/transformers/models/ernie_m/modeling_ernie_m.py @@ -395,6 +396,7 @@ src/transformers/models/deit/feature_extraction_deit.py src/transformers/models/detr/feature_extraction_detr.py src/transformers/models/donut/feature_extraction_donut.py src/transformers/models/dpt/feature_extraction_dpt.py +src/transformers/models/encodec/feature_extraction_encodec.py src/transformers/models/flava/feature_extraction_flava.py src/transformers/models/glpn/feature_extraction_glpn.py src/transformers/models/imagegpt/feature_extraction_imagegpt.py