[ { "path": ".github/actions/audiocraft_build/action.yml", "content": "name: audiocraft_build\ndescription: 'Build audiocraft env.'\nruns:\n using: \"composite\"\n steps:\n - uses: actions/setup-python@v2\n with:\n python-version: 3.9\n - uses: actions/cache@v3\n id: cache\n with:\n path: env\n key: audiocraft_env-${{ hashFiles('**/requirements.txt') }}\n\n - if: ${{ steps.cache.outputs.cache-hit != 'true' }}\n name: Install dependencies\n shell: bash\n run: |\n sudo apt-get update\n sudo apt-get install libsndfile1-dev ffmpeg\n python3 -m venv env\n . env/bin/activate\n python -m pip install --upgrade pip\n pip install 'numpy<2' torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0\n pip install xformers==0.0.22.post7\n pip install -e '.[dev,wm]'\n - name: System Dependencies\n shell: bash\n run: |\n sudo apt-get update\n sudo apt-get install libsndfile1-dev ffmpeg\n" }, { "path": ".github/workflows/audiocraft_docs.yml", "content": "name: audiocraft_docs\non:\n push:\n branches: [ main ]\n\njobs:\n run_docs:\n name: Run docs\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v2\n - uses: ./.github/actions/audiocraft_build\n - name: Config git\n run: |\n git config --global user.email \"defossez@fb.com\"\n git config --global user.name \"Alexandre Défossez (autodoc)\"\n\n - name: Reset branch\n run: |\n git branch -f gh-docs main\n git checkout gh-docs\n\n - name: Make docs\n run: |\n . env/bin/activate\n make api_docs\n git add -f api_docs\n git commit -m api_docs\n\n - name: Push branch\n run: |\n git push -f -u origin gh-docs\n" }, { "path": ".github/workflows/audiocraft_linter.yml", "content": "name: audiocraft_linter\non:\n push:\n branches: [ main ]\n pull_request:\n branches: [ main, audiocraft_pub_main ]\n\njobs:\n run_linter:\n name: Run linter\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v2\n - uses: ./.github/actions/audiocraft_build\n - run: |\n . env/bin/activate\n make linter\n" }, { "path": ".github/workflows/audiocraft_tests.yml", "content": "name: audiocraft_tests\non:\n push:\n branches: [ main ]\n pull_request:\n branches: [ main, audiocraft_pub_main ]\n\njobs:\n run_tests:\n name: Run tests\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v2\n - uses: ./.github/actions/audiocraft_build\n - name: Run unit tests\n run: |\n . env/bin/activate\n make tests\n - name: Run integration tests\n run: |\n . env/bin/activate\n make tests_integ\n" }, { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# macOS dir files\n.DS_Store\n\n# Distribution / packaging\n.Python\nenv/\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\n*.egg-info/\n.installed.cfg\n*.egg\n.ipynb_checkpoints\n\n# Tests and linter\n.pytest_cache/\n.mypy_cache/\n.coverage\n\n# docs\n/api_docs\n\n# dotenv\n.env\n.envrc\n\n# virtualenv\n.venv\nvenv/\nENV/\n\n# egs with manifest files\negs/*\n!egs/example\n# local datasets\ndataset/*\n!dataset/example\n\n# personal notebooks & scripts\n*/local_scripts\n*/notes\n.vscode/\n/notebooks\n/local_scripts\n/notes\n" }, { "path": "CHANGELOG.md", "content": "# Changelog\n\nAll notable changes to this project will be documented in this file.\n\nThe format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).\n\n## [1.4.0a2] - 2025-01-14\n\nAdd training and inference code for JASCO (https://arxiv.org/abs/2406.10970) along with the [hf checkpoints](https://huggingface.co/facebook/jasco-chords-drums-melody-1B).\n\n## [1.4.0a1] - 2024-06-03\n\nAdding new metric PesqMetric ([Perceptual Evaluation of Speech Quality](https://doi.org/10.5281/zenodo.6549559))\n\nAdding multiple audio augmentation functions: generating pink noises, up-/downsampling, low-/highpass filtering, banpass filtering, smoothing, duck masking, boosting. All are wrapped in the `audiocraft.utils.audio_effects.AudioEffects` and can be called with the API `audiocraft.utils.audio_effects.select_audio_effects`.\n\nAdd training code for AudioSeal (https://arxiv.org/abs/2401.17264) along with the [hf checkpoints]( https://huggingface.co/facebook/audioseal).\n\n## [1.3.0] - 2024-05-02\n\nAdding the MAGNeT model (https://arxiv.org/abs/2401.04577) along with hf checkpoints and a gradio demo app.\n\nTypo fixes.\n\nFixing setup.py to install only audiocraft, not the unit tests and scripts.\n\nFix FSDP support with PyTorch 2.1.0. \n\n## [1.2.0] - 2024-01-11\n\nAdding stereo models.\n\nFixed the commitment loss, which was until now only applied to the first RVQ layer.\n\nRemoved compression model state from the LM checkpoints, for consistency, it\nshould always be loaded from the original `compression_model_checkpoint`.\n\n\n## [1.1.0] - 2023-11-06\n\nNot using torchaudio anymore when writing audio files, relying instead directly on the commandline ffmpeg. Also not using it anymore for reading audio files, for similar reasons.\n\nFixed DAC support with non default number of codebooks.\n\nFixed bug when `two_step_cfg` was overriden when calling `generate()`.\n\nFixed samples being always prompted with audio, rather than having both prompted and unprompted.\n\n**Backward incompatible change:** A `torch.no_grad` around the computation of the conditioning made its way in the public release.\n\tThe released models were trained without this. Those impact linear layers applied to the output of the T5 or melody conditioners.\n\tWe removed it, so you might need to retrain models.\n\n**Backward incompatible change:** Fixing wrong sample rate in CLAP (WARNING if you trained model with CLAP before).\n\n**Backward incompatible change:** Renamed VALLEPattern to CoarseFirstPattern, as it was wrongly named. Probably no one\n\tretrained a model with this pattern, so hopefully this won't impact you!\n\n\n## [1.0.0] - 2023-09-07\n\nMajor revision, added training code for EnCodec, AudioGen, MusicGen, and MultiBandDiffusion.\nAdded pretrained model for AudioGen and MultiBandDiffusion.\n\n## [0.0.2] - 2023-08-01\n\nImproved demo, fixed top p (thanks @jnordberg).\n\nCompressor tanh on output to avoid clipping with some style (especially piano).\nNow repeating the conditioning periodically if it is too short.\n\nMore options when launching Gradio app locally (thanks @ashleykleynhans).\n\nTesting out PyTorch 2.0 memory efficient attention.\n\nAdded extended generation (infinite length) by slowly moving the windows.\nNote that other implementations exist: https://github.com/camenduru/MusicGen-colab.\n\n## [0.0.1] - 2023-06-09\n\nInitial release, with model evaluation only.\n" }, { "path": "CODE_OF_CONDUCT.md", "content": "# Code of Conduct\n\n## Our Pledge\n\nIn the interest of fostering an open and welcoming environment, we as\ncontributors and maintainers pledge to make participation in our project and\nour community a harassment-free experience for everyone, regardless of age, body\nsize, disability, ethnicity, sex characteristics, gender identity and expression,\nlevel of experience, education, socio-economic status, nationality, personal\nappearance, race, religion, or sexual identity and orientation.\n\n## Our Standards\n\nExamples of behavior that contributes to creating a positive environment\ninclude:\n\n* Using welcoming and inclusive language\n* Being respectful of differing viewpoints and experiences\n* Gracefully accepting constructive criticism\n* Focusing on what is best for the community\n* Showing empathy towards other community members\n\nExamples of unacceptable behavior by participants include:\n\n* The use of sexualized language or imagery and unwelcome sexual attention or\nadvances\n* Trolling, insulting/derogatory comments, and personal or political attacks\n* Public or private harassment\n* Publishing others' private information, such as a physical or electronic\naddress, without explicit permission\n* Other conduct which could reasonably be considered inappropriate in a\nprofessional setting\n\n## Our Responsibilities\n\nProject maintainers are responsible for clarifying the standards of acceptable\nbehavior and are expected to take appropriate and fair corrective action in\nresponse to any instances of unacceptable behavior.\n\nProject maintainers have the right and responsibility to remove, edit, or\nreject comments, commits, code, wiki edits, issues, and other contributions\nthat are not aligned to this Code of Conduct, or to ban temporarily or\npermanently any contributor for other behaviors that they deem inappropriate,\nthreatening, offensive, or harmful.\n\n## Scope\n\nThis Code of Conduct applies within all project spaces, and it also applies when\nan individual is representing the project or its community in public spaces.\nExamples of representing a project or community include using an official\nproject e-mail address, posting via an official social media account, or acting\nas an appointed representative at an online or offline event. Representation of\na project may be further defined and clarified by project maintainers.\n\nThis Code of Conduct also applies outside the project spaces when there is a\nreasonable belief that an individual's behavior may have a negative impact on\nthe project or its community.\n\n## Enforcement\n\nInstances of abusive, harassing, or otherwise unacceptable behavior may be\nreported by contacting the project team at . All\ncomplaints will be reviewed and investigated and will result in a response that\nis deemed necessary and appropriate to the circumstances. The project team is\nobligated to maintain confidentiality with regard to the reporter of an incident.\nFurther details of specific enforcement policies may be posted separately.\n\nProject maintainers who do not follow or enforce the Code of Conduct in good\nfaith may face temporary or permanent repercussions as determined by other\nmembers of the project's leadership.\n\n## Attribution\n\nThis Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,\navailable at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html\n\n[homepage]: https://www.contributor-covenant.org\n\nFor answers to common questions about this code of conduct, see\nhttps://www.contributor-covenant.org/faq\n" }, { "path": "CONTRIBUTING.md", "content": "# Contributing to AudioCraft\n\nWe want to make contributing to this project as easy and transparent as\npossible.\n\n## Pull Requests\n\nAudioCraft is the implementation of a research paper.\nTherefore, we do not plan on accepting many pull requests for new features.\nWe certainly welcome them for bug fixes.\n\n1. Fork the repo and create your branch from `main`.\n2. If you've added code that should be tested, add tests.\n3. If you've changed APIs, update the documentation.\n4. Ensure the test suite passes.\n5. Make sure your code lints.\n6. If you haven't already, complete the Contributor License Agreement (\"CLA\").\n\n## Contributor License Agreement (\"CLA\")\nIn order to accept your pull request, we need you to submit a CLA. You only need\nto do this once to work on any of Meta's open source projects.\n\nComplete your CLA here: \n\n## Issues\nWe use GitHub issues to track public bugs. Please ensure your description is\nclear and has sufficient instructions to be able to reproduce the issue.\n\nMeta has a [bounty program](https://www.facebook.com/whitehat/) for the safe\ndisclosure of security bugs. In those cases, please go through the process\noutlined on that page and do not file a public issue.\n\n## License\nBy contributing to encodec, you agree that your contributions will be licensed\nunder the LICENSE file in the root directory of this source tree.\n" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) Meta Platforms, Inc. and affiliates.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "LICENSE_weights", "content": "Attribution-NonCommercial 4.0 International\n\n=======================================================================\n\nCreative Commons Corporation (\"Creative Commons\") is not a law firm and\ndoes not provide legal services or legal advice. Distribution of\nCreative Commons public licenses does not create a lawyer-client or\nother relationship. Creative Commons makes its licenses and related\ninformation available on an \"as-is\" basis. Creative Commons gives no\nwarranties regarding its licenses, any material licensed under their\nterms and conditions, or any related information. Creative Commons\ndisclaims all liability for damages resulting from their use to the\nfullest extent possible.\n\nUsing Creative Commons Public Licenses\n\nCreative Commons public licenses provide a standard set of terms and\nconditions that creators and other rights holders may use to share\noriginal works of authorship and other material subject to copyright\nand certain other rights specified in the public license below. The\nfollowing considerations are for informational purposes only, are not\nexhaustive, and do not form part of our licenses.\n\n Considerations for licensors: Our public licenses are\n intended for use by those authorized to give the public\n permission to use material in ways otherwise restricted by\n copyright and certain other rights. Our licenses are\n irrevocable. Licensors should read and understand the terms\n and conditions of the license they choose before applying it.\n Licensors should also secure all rights necessary before\n applying our licenses so that the public can reuse the\n material as expected. Licensors should clearly mark any\n material not subject to the license. This includes other CC-\n licensed material, or material used under an exception or\n limitation to copyright. More considerations for licensors:\n\twiki.creativecommons.org/Considerations_for_licensors\n\n Considerations for the public: By using one of our public\n licenses, a licensor grants the public permission to use the\n licensed material under specified terms and conditions. If\n the licensor's permission is not necessary for any reason--for\n example, because of any applicable exception or limitation to\n copyright--then that use is not regulated by the license. Our\n licenses grant only permissions under copyright and certain\n other rights that a licensor has authority to grant. Use of\n the licensed material may still be restricted for other\n reasons, including because others have copyright or other\n rights in the material. A licensor may make special requests,\n such as asking that all changes be marked or described.\n Although not required by our licenses, you are encouraged to\n respect those requests where reasonable. More_considerations\n for the public: \n\twiki.creativecommons.org/Considerations_for_licensees\n\n=======================================================================\n\nCreative Commons Attribution-NonCommercial 4.0 International Public\nLicense\n\nBy exercising the Licensed Rights (defined below), You accept and agree\nto be bound by the terms and conditions of this Creative Commons\nAttribution-NonCommercial 4.0 International Public License (\"Public\nLicense\"). To the extent this Public License may be interpreted as a\ncontract, You are granted the Licensed Rights in consideration of Your\nacceptance of these terms and conditions, and the Licensor grants You\nsuch rights in consideration of benefits the Licensor receives from\nmaking the Licensed Material available under these terms and\nconditions.\n\nSection 1 -- Definitions.\n\n a. Adapted Material means material subject to Copyright and Similar\n Rights that is derived from or based upon the Licensed Material\n and in which the Licensed Material is translated, altered,\n arranged, transformed, or otherwise modified in a manner requiring\n permission under the Copyright and Similar Rights held by the\n Licensor. For purposes of this Public License, where the Licensed\n Material is a musical work, performance, or sound recording,\n Adapted Material is always produced where the Licensed Material is\n synched in timed relation with a moving image.\n\n b. Adapter's License means the license You apply to Your Copyright\n and Similar Rights in Your contributions to Adapted Material in\n accordance with the terms and conditions of this Public License.\n\n c. Copyright and Similar Rights means copyright and/or similar rights\n closely related to copyright including, without limitation,\n performance, broadcast, sound recording, and Sui Generis Database\n Rights, without regard to how the rights are labeled or\n categorized. For purposes of this Public License, the rights\n specified in Section 2(b)(1)-(2) are not Copyright and Similar\n Rights.\n d. Effective Technological Measures means those measures that, in the\n absence of proper authority, may not be circumvented under laws\n fulfilling obligations under Article 11 of the WIPO Copyright\n Treaty adopted on December 20, 1996, and/or similar international\n agreements.\n\n e. Exceptions and Limitations means fair use, fair dealing, and/or\n any other exception or limitation to Copyright and Similar Rights\n that applies to Your use of the Licensed Material.\n\n f. Licensed Material means the artistic or literary work, database,\n or other material to which the Licensor applied this Public\n License.\n\n g. Licensed Rights means the rights granted to You subject to the\n terms and conditions of this Public License, which are limited to\n all Copyright and Similar Rights that apply to Your use of the\n Licensed Material and that the Licensor has authority to license.\n\n h. Licensor means the individual(s) or entity(ies) granting rights\n under this Public License.\n\n i. NonCommercial means not primarily intended for or directed towards\n commercial advantage or monetary compensation. For purposes of\n this Public License, the exchange of the Licensed Material for\n other material subject to Copyright and Similar Rights by digital\n file-sharing or similar means is NonCommercial provided there is\n no payment of monetary compensation in connection with the\n exchange.\n\n j. Share means to provide material to the public by any means or\n process that requires permission under the Licensed Rights, such\n as reproduction, public display, public performance, distribution,\n dissemination, communication, or importation, and to make material\n available to the public including in ways that members of the\n public may access the material from a place and at a time\n individually chosen by them.\n\n k. Sui Generis Database Rights means rights other than copyright\n resulting from Directive 96/9/EC of the European Parliament and of\n the Council of 11 March 1996 on the legal protection of databases,\n as amended and/or succeeded, as well as other essentially\n equivalent rights anywhere in the world.\n\n l. You means the individual or entity exercising the Licensed Rights\n under this Public License. Your has a corresponding meaning.\n\nSection 2 -- Scope.\n\n a. License grant.\n\n 1. Subject to the terms and conditions of this Public License,\n the Licensor hereby grants You a worldwide, royalty-free,\n non-sublicensable, non-exclusive, irrevocable license to\n exercise the Licensed Rights in the Licensed Material to:\n\n a. reproduce and Share the Licensed Material, in whole or\n in part, for NonCommercial purposes only; and\n\n b. produce, reproduce, and Share Adapted Material for\n NonCommercial purposes only.\n\n 2. Exceptions and Limitations. For the avoidance of doubt, where\n Exceptions and Limitations apply to Your use, this Public\n License does not apply, and You do not need to comply with\n its terms and conditions.\n\n 3. Term. The term of this Public License is specified in Section\n 6(a).\n\n 4. Media and formats; technical modifications allowed. The\n Licensor authorizes You to exercise the Licensed Rights in\n all media and formats whether now known or hereafter created,\n and to make technical modifications necessary to do so. The\n Licensor waives and/or agrees not to assert any right or\n authority to forbid You from making technical modifications\n necessary to exercise the Licensed Rights, including\n technical modifications necessary to circumvent Effective\n Technological Measures. For purposes of this Public License,\n simply making modifications authorized by this Section 2(a)\n (4) never produces Adapted Material.\n\n 5. Downstream recipients.\n\n a. Offer from the Licensor -- Licensed Material. Every\n recipient of the Licensed Material automatically\n receives an offer from the Licensor to exercise the\n Licensed Rights under the terms and conditions of this\n Public License.\n\n b. No downstream restrictions. You may not offer or impose\n any additional or different terms or conditions on, or\n apply any Effective Technological Measures to, the\n Licensed Material if doing so restricts exercise of the\n Licensed Rights by any recipient of the Licensed\n Material.\n\n 6. No endorsement. Nothing in this Public License constitutes or\n may be construed as permission to assert or imply that You\n are, or that Your use of the Licensed Material is, connected\n with, or sponsored, endorsed, or granted official status by,\n the Licensor or others designated to receive attribution as\n provided in Section 3(a)(1)(A)(i).\n\n b. Other rights.\n\n 1. Moral rights, such as the right of integrity, are not\n licensed under this Public License, nor are publicity,\n privacy, and/or other similar personality rights; however, to\n the extent possible, the Licensor waives and/or agrees not to\n assert any such rights held by the Licensor to the limited\n extent necessary to allow You to exercise the Licensed\n Rights, but not otherwise.\n\n 2. Patent and trademark rights are not licensed under this\n Public License.\n\n 3. To the extent possible, the Licensor waives any right to\n collect royalties from You for the exercise of the Licensed\n Rights, whether directly or through a collecting society\n under any voluntary or waivable statutory or compulsory\n licensing scheme. In all other cases the Licensor expressly\n reserves any right to collect such royalties, including when\n the Licensed Material is used other than for NonCommercial\n purposes.\n\nSection 3 -- License Conditions.\n\nYour exercise of the Licensed Rights is expressly made subject to the\nfollowing conditions.\n\n a. Attribution.\n\n 1. If You Share the Licensed Material (including in modified\n form), You must:\n\n a. retain the following if it is supplied by the Licensor\n with the Licensed Material:\n\n i. identification of the creator(s) of the Licensed\n Material and any others designated to receive\n attribution, in any reasonable manner requested by\n the Licensor (including by pseudonym if\n designated);\n\n ii. a copyright notice;\n\n iii. a notice that refers to this Public License;\n\n iv. a notice that refers to the disclaimer of\n warranties;\n\n v. a URI or hyperlink to the Licensed Material to the\n extent reasonably practicable;\n\n b. indicate if You modified the Licensed Material and\n retain an indication of any previous modifications; and\n\n c. indicate the Licensed Material is licensed under this\n Public License, and include the text of, or the URI or\n hyperlink to, this Public License.\n\n 2. You may satisfy the conditions in Section 3(a)(1) in any\n reasonable manner based on the medium, means, and context in\n which You Share the Licensed Material. For example, it may be\n reasonable to satisfy the conditions by providing a URI or\n hyperlink to a resource that includes the required\n information.\n\n 3. If requested by the Licensor, You must remove any of the\n information required by Section 3(a)(1)(A) to the extent\n reasonably practicable.\n\n 4. If You Share Adapted Material You produce, the Adapter's\n License You apply must not prevent recipients of the Adapted\n Material from complying with this Public License.\n\nSection 4 -- Sui Generis Database Rights.\n\nWhere the Licensed Rights include Sui Generis Database Rights that\napply to Your use of the Licensed Material:\n\n a. for the avoidance of doubt, Section 2(a)(1) grants You the right\n to extract, reuse, reproduce, and Share all or a substantial\n portion of the contents of the database for NonCommercial purposes\n only;\n\n b. if You include all or a substantial portion of the database\n contents in a database in which You have Sui Generis Database\n Rights, then the database in which You have Sui Generis Database\n Rights (but not its individual contents) is Adapted Material; and\n\n c. You must comply with the conditions in Section 3(a) if You Share\n all or a substantial portion of the contents of the database.\n\nFor the avoidance of doubt, this Section 4 supplements and does not\nreplace Your obligations under this Public License where the Licensed\nRights include other Copyright and Similar Rights.\n\nSection 5 -- Disclaimer of Warranties and Limitation of Liability.\n\n a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE\n EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS\n AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF\n ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,\n IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,\n WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR\n PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,\n ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT\n KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT\n ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.\n\n b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE\n TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,\n NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,\n INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,\n COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR\n USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN\n ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR\n DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR\n IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.\n\n c. The disclaimer of warranties and limitation of liability provided\n above shall be interpreted in a manner that, to the extent\n possible, most closely approximates an absolute disclaimer and\n waiver of all liability.\n\nSection 6 -- Term and Termination.\n\n a. This Public License applies for the term of the Copyright and\n Similar Rights licensed here. However, if You fail to comply with\n this Public License, then Your rights under this Public License\n terminate automatically.\n\n b. Where Your right to use the Licensed Material has terminated under\n Section 6(a), it reinstates:\n\n 1. automatically as of the date the violation is cured, provided\n it is cured within 30 days of Your discovery of the\n violation; or\n\n 2. upon express reinstatement by the Licensor.\n\n For the avoidance of doubt, this Section 6(b) does not affect any\n right the Licensor may have to seek remedies for Your violations\n of this Public License.\n\n c. For the avoidance of doubt, the Licensor may also offer the\n Licensed Material under separate terms or conditions or stop\n distributing the Licensed Material at any time; however, doing so\n will not terminate this Public License.\n\n d. Sections 1, 5, 6, 7, and 8 survive termination of this Public\n License.\n\nSection 7 -- Other Terms and Conditions.\n\n a. The Licensor shall not be bound by any additional or different\n terms or conditions communicated by You unless expressly agreed.\n\n b. Any arrangements, understandings, or agreements regarding the\n Licensed Material not stated herein are separate from and\n independent of the terms and conditions of this Public License.\n\nSection 8 -- Interpretation.\n\n a. For the avoidance of doubt, this Public License does not, and\n shall not be interpreted to, reduce, limit, restrict, or impose\n conditions on any use of the Licensed Material that could lawfully\n be made without permission under this Public License.\n\n b. To the extent possible, if any provision of this Public License is\n deemed unenforceable, it shall be automatically reformed to the\n minimum extent necessary to make it enforceable. If the provision\n cannot be reformed, it shall be severed from this Public License\n without affecting the enforceability of the remaining terms and\n conditions.\n\n c. No term or condition of this Public License will be waived and no\n failure to comply consented to unless expressly agreed to by the\n Licensor.\n\n d. Nothing in this Public License constitutes or may be interpreted\n as a limitation upon, or waiver of, any privileges and immunities\n that apply to the Licensor or You, including from the legal\n processes of any jurisdiction or authority.\n\n=======================================================================\n\nCreative Commons is not a party to its public\nlicenses. Notwithstanding, Creative Commons may elect to apply one of\nits public licenses to material it publishes and in those instances\nwill be considered the “Licensor.” The text of the Creative Commons\npublic licenses is dedicated to the public domain under the CC0 Public\nDomain Dedication. Except for the limited purpose of indicating that\nmaterial is shared under a Creative Commons public license or as\notherwise permitted by the Creative Commons policies published at\ncreativecommons.org/policies, Creative Commons does not authorize the\nuse of the trademark \"Creative Commons\" or any other trademark or logo\nof Creative Commons without its prior written consent including,\nwithout limitation, in connection with any unauthorized modifications\nto any of its public licenses or any other arrangements,\nunderstandings, or agreements concerning use of licensed material. For\nthe avoidance of doubt, this paragraph does not form part of the\npublic licenses.\n\nCreative Commons may be contacted at creativecommons.org.\n" }, { "path": "MANIFEST.in", "content": "include Makefile\ninclude LICENSE\ninclude LICENSE_weights\ninclude *.md\ninclude *.ini\ninclude requirements.txt\ninclude audiocraft/py.typed\ninclude assets/*.mp3\ninclude datasets/*.mp3\nrecursive-include config *.yaml\nrecursive-include demos *.py\nrecursive-include demos *.ipynb\nrecursive-include scripts *.py\nrecursive-include model_cards *.md\nrecursive-include docs *.md\n" }, { "path": "Makefile", "content": "INTEG=AUDIOCRAFT_DORA_DIR=\"/tmp/magma_$(USER)\" python3 -m dora -v run --clear device=cpu dataset.num_workers=0 optim.epochs=1 \\\n\tdataset.train.num_samples=10 dataset.valid.num_samples=10 \\\n\tdataset.evaluate.num_samples=10 dataset.generate.num_samples=2 sample_rate=16000 \\\n\tlogging.level=DEBUG\nINTEG_COMPRESSION = $(INTEG) solver=compression/debug rvq.n_q=2 rvq.bins=48 checkpoint.save_last=true # SIG is 5091833e\nINTEG_MUSICGEN = $(INTEG) solver=musicgen/debug dset=audio/example compression_model_checkpoint=//sig/5091833e \\\n\ttransformer_lm.n_q=2 transformer_lm.card=48 transformer_lm.dim=16 checkpoint.save_last=false # Using compression model from 5091833e\nINTEG_AUDIOGEN = $(INTEG) solver=audiogen/debug dset=audio/example compression_model_checkpoint=//sig/5091833e \\\n\ttransformer_lm.n_q=2 transformer_lm.card=48 transformer_lm.dim=16 checkpoint.save_last=false # Using compression model from 5091833e\nINTEG_MBD = $(INTEG) solver=diffusion/debug dset=audio/example \\\n\tcheckpoint.save_last=false # Using compression model from 616d7b3c\nINTEG_WATERMARK = AUDIOCRAFT_DORA_DIR=\"/tmp/wm_$(USER)\" dora run device=cpu dataset.num_workers=0 optim.epochs=1 \\\n dataset.train.num_samples=10 dataset.valid.num_samples=10 dataset.evaluate.num_samples=10 dataset.generate.num_samples=10 \\\n\tlogging.level=DEBUG solver=watermark/robustness checkpoint.save_last=false dset=audio/example \n\ndefault: linter tests\n\ninstall:\n\tpip install -U pip\n\tpip install -U -e '.[dev]'\n\nlinter:\n\tflake8 audiocraft && mypy audiocraft\n\tflake8 tests && mypy tests\n\ntests:\n\tcoverage run -m pytest tests\n\tcoverage report\n\ntests_integ:\n\t$(INTEG_COMPRESSION)\n\t$(INTEG_MBD)\n\t$(INTEG_MUSICGEN)\n\t$(INTEG_AUDIOGEN)\n\t$(INTEG_WATERMARK)\n\n\napi_docs:\n\tpdoc3 --html -o api_docs -f audiocraft\n\ndist:\n\tpython setup.py sdist\n\n.PHONY: linter tests api_docs dist\n" }, { "path": "README.md", "content": "# AudioCraft\n![docs badge](https://github.com/facebookresearch/audiocraft/workflows/audiocraft_docs/badge.svg)\n![linter badge](https://github.com/facebookresearch/audiocraft/workflows/audiocraft_linter/badge.svg)\n![tests badge](https://github.com/facebookresearch/audiocraft/workflows/audiocraft_tests/badge.svg)\n\nAudioCraft is a PyTorch library for deep learning research on audio generation. AudioCraft contains inference and training code\nfor two state-of-the-art AI generative models producing high-quality audio: AudioGen and MusicGen.\n\n\n## Installation\nAudioCraft requires Python 3.9, PyTorch 2.1.0. To install AudioCraft, you can run the following:\n\n```shell\n# Best to make sure you have torch installed first, in particular before installing xformers.\n# Don't run this if you already have PyTorch installed.\npython -m pip install 'torch==2.1.0'\n# You might need the following before trying to install the packages\npython -m pip install setuptools wheel\n# Then proceed to one of the following\npython -m pip install -U audiocraft # stable release\npython -m pip install -U git+https://git@github.com/facebookresearch/audiocraft#egg=audiocraft # bleeding edge\npython -m pip install -e . # or if you cloned the repo locally (mandatory if you want to train).\npython -m pip install -e '.[wm]' # if you want to train a watermarking model\n```\n\nWe also recommend having `ffmpeg` installed, either through your system or Anaconda:\n```bash\nsudo apt-get install ffmpeg\n# Or if you are using Anaconda or Miniconda\nconda install \"ffmpeg<5\" -c conda-forge\n```\n\n## Models\n\nAt the moment, AudioCraft contains the training code and inference code for:\n* [MusicGen](./docs/MUSICGEN.md): A state-of-the-art controllable text-to-music model.\n* [AudioGen](./docs/AUDIOGEN.md): A state-of-the-art text-to-sound model.\n* [EnCodec](./docs/ENCODEC.md): A state-of-the-art high fidelity neural audio codec.\n* [Multi Band Diffusion](./docs/MBD.md): An EnCodec compatible decoder using diffusion.\n* [MAGNeT](./docs/MAGNET.md): A state-of-the-art non-autoregressive model for text-to-music and text-to-sound.\n* [AudioSeal](./docs/WATERMARKING.md): A state-of-the-art audio watermarking.\n* [MusicGen Style](./docs/MUSICGEN_STYLE.md): A state-of-the-art text-and-style-to-music model.\n* [JASCO](./docs/JASCO.md): \"High quality text-to-music model conditioned on chords, melodies and drum tracks\"\n\n\n## Training code\n\nAudioCraft contains PyTorch components for deep learning research in audio and training pipelines for the developed models.\nFor a general introduction of AudioCraft design principles and instructions to develop your own training pipeline, refer to\nthe [AudioCraft training documentation](./docs/TRAINING.md).\n\nFor reproducing existing work and using the developed training pipelines, refer to the instructions for each specific model\nthat provides pointers to configuration, example grids and model/task-specific information and FAQ.\n\n\n## API documentation\n\nWe provide some [API documentation](https://facebookresearch.github.io/audiocraft/api_docs/audiocraft/index.html) for AudioCraft.\n\n\n## FAQ\n\n#### Is the training code available?\n\nYes! We provide the training code for [EnCodec](./docs/ENCODEC.md), [MusicGen](./docs/MUSICGEN.md),[Multi Band Diffusion](./docs/MBD.md) and [JASCO](./docs/JASCO.md).\n\n#### Where are the models stored?\n\nHugging Face stored the model in a specific location, which can be overridden by setting the `AUDIOCRAFT_CACHE_DIR` environment variable for the AudioCraft models.\nIn order to change the cache location of the other Hugging Face models, please check out the [Hugging Face Transformers documentation for the cache setup](https://huggingface.co/docs/transformers/installation#cache-setup).\nFinally, if you use a model that relies on Demucs (e.g. `musicgen-melody`) and want to change the download location for Demucs, refer to the [Torch Hub documentation](https://pytorch.org/docs/stable/hub.html#where-are-my-downloaded-models-saved).\n\n\n## License\n* The code in this repository is released under the MIT license as found in the [LICENSE file](LICENSE).\n* The models weights in this repository are released under the CC-BY-NC 4.0 license as found in the [LICENSE_weights file](LICENSE_weights).\n\n\n## Citation\n\nFor the general framework of AudioCraft, please cite the following.\n```\n@inproceedings{copet2023simple,\n title={Simple and Controllable Music Generation},\n author={Jade Copet and Felix Kreuk and Itai Gat and Tal Remez and David Kant and Gabriel Synnaeve and Yossi Adi and Alexandre Défossez},\n booktitle={Thirty-seventh Conference on Neural Information Processing Systems},\n year={2023},\n}\n```\n\nWhen referring to a specific model, please cite as mentioned in the model specific README, e.g\n[./docs/MUSICGEN.md](./docs/MUSICGEN.md), [./docs/AUDIOGEN.md](./docs/AUDIOGEN.md), etc.\n" }, { "path": "audiocraft/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"\nAudioCraft is a general framework for training audio generative models.\nAt the moment we provide the training code for:\n\n- [MusicGen](https://arxiv.org/abs/2306.05284), a state-of-the-art\n text-to-music and melody+text autoregressive generative model.\n For the solver, see `audiocraft.solvers.musicgen.MusicGenSolver`, and for the model,\n `audiocraft.models.musicgen.MusicGen`.\n- [AudioGen](https://arxiv.org/abs/2209.15352), a state-of-the-art\n text-to-general-audio generative model.\n- [EnCodec](https://arxiv.org/abs/2210.13438), efficient and high fidelity\n neural audio codec which provides an excellent tokenizer for autoregressive language models.\n See `audiocraft.solvers.compression.CompressionSolver`, and `audiocraft.models.encodec.EncodecModel`.\n- [MultiBandDiffusion](TODO), alternative diffusion-based decoder compatible with EnCodec that\n improves the perceived quality and reduces the artifacts coming from adversarial decoders.\n- [JASCO](https://arxiv.org/abs/2406.10970) Joint Audio and Symbolic Conditioning for Temporally Controlled\n Text-to-Music Generation.\n\"\"\"\n\n# flake8: noqa\nfrom . import data, modules, models\n\n__version__ = '1.4.0a2'\n" }, { "path": "audiocraft/adversarial/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Adversarial losses and discriminator architectures.\"\"\"\n\n# flake8: noqa\nfrom .discriminators import (\n MultiPeriodDiscriminator,\n MultiScaleDiscriminator,\n MultiScaleSTFTDiscriminator\n)\nfrom .losses import (\n AdversarialLoss,\n AdvLossType,\n get_adv_criterion,\n get_fake_criterion,\n get_real_criterion,\n FeatLossType,\n FeatureMatchingLoss\n)\n" }, { "path": "audiocraft/adversarial/discriminators/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# flake8: noqa\nfrom .mpd import MultiPeriodDiscriminator\nfrom .msd import MultiScaleDiscriminator\nfrom .msstftd import MultiScaleSTFTDiscriminator\n" }, { "path": "audiocraft/adversarial/discriminators/base.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom abc import ABC, abstractmethod\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\n\n\nFeatureMapType = tp.List[torch.Tensor]\nLogitsType = torch.Tensor\nMultiDiscriminatorOutputType = tp.Tuple[tp.List[LogitsType], tp.List[FeatureMapType]]\n\n\nclass MultiDiscriminator(ABC, nn.Module):\n \"\"\"Base implementation for discriminators composed of sub-discriminators acting at different scales.\n \"\"\"\n def __init__(self):\n super().__init__()\n\n @abstractmethod\n def forward(self, x: torch.Tensor) -> MultiDiscriminatorOutputType:\n ...\n\n @property\n @abstractmethod\n def num_discriminators(self) -> int:\n \"\"\"Number of discriminators.\n \"\"\"\n ...\n" }, { "path": "audiocraft/adversarial/discriminators/mpd.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom ...modules import NormConv2d\nfrom .base import MultiDiscriminator, MultiDiscriminatorOutputType\n\n\ndef get_padding(kernel_size: int, dilation: int = 1) -> int:\n return int((kernel_size * dilation - dilation) / 2)\n\n\nclass PeriodDiscriminator(nn.Module):\n \"\"\"Period sub-discriminator.\n\n Args:\n period (int): Period between samples of audio.\n in_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n n_layers (int): Number of convolutional layers.\n kernel_sizes (list of int): Kernel sizes for convolutions.\n stride (int): Stride for convolutions.\n filters (int): Initial number of filters in convolutions.\n filters_scale (int): Multiplier of number of filters as we increase depth.\n max_filters (int): Maximum number of filters.\n norm (str): Normalization method.\n activation (str): Activation function.\n activation_params (dict): Parameters to provide to the activation function.\n \"\"\"\n def __init__(self, period: int, in_channels: int = 1, out_channels: int = 1,\n n_layers: int = 5, kernel_sizes: tp.List[int] = [5, 3], stride: int = 3,\n filters: int = 8, filters_scale: int = 4, max_filters: int = 1024,\n norm: str = 'weight_norm', activation: str = 'LeakyReLU',\n activation_params: dict = {'negative_slope': 0.2}):\n super().__init__()\n self.period = period\n self.n_layers = n_layers\n self.activation = getattr(torch.nn, activation)(**activation_params)\n self.convs = nn.ModuleList()\n in_chs = in_channels\n for i in range(self.n_layers):\n out_chs = min(filters * (filters_scale ** (i + 1)), max_filters)\n eff_stride = 1 if i == self.n_layers - 1 else stride\n self.convs.append(NormConv2d(in_chs, out_chs, kernel_size=(kernel_sizes[0], 1), stride=(eff_stride, 1),\n padding=((kernel_sizes[0] - 1) // 2, 0), norm=norm))\n in_chs = out_chs\n self.conv_post = NormConv2d(in_chs, out_channels, kernel_size=(kernel_sizes[1], 1), stride=1,\n padding=((kernel_sizes[1] - 1) // 2, 0), norm=norm)\n\n def forward(self, x: torch.Tensor):\n fmap = []\n # 1d to 2d\n b, c, t = x.shape\n if t % self.period != 0: # pad first\n n_pad = self.period - (t % self.period)\n x = F.pad(x, (0, n_pad), 'reflect')\n t = t + n_pad\n x = x.view(b, c, t // self.period, self.period)\n\n for conv in self.convs:\n x = conv(x)\n x = self.activation(x)\n fmap.append(x)\n x = self.conv_post(x)\n fmap.append(x)\n # x = torch.flatten(x, 1, -1)\n\n return x, fmap\n\n\nclass MultiPeriodDiscriminator(MultiDiscriminator):\n \"\"\"Multi-Period (MPD) Discriminator.\n\n Args:\n in_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n periods (Sequence[int]): Periods between samples of audio for the sub-discriminators.\n **kwargs: Additional args for `PeriodDiscriminator`\n \"\"\"\n def __init__(self, in_channels: int = 1, out_channels: int = 1,\n periods: tp.Sequence[int] = [2, 3, 5, 7, 11], **kwargs):\n super().__init__()\n self.discriminators = nn.ModuleList([\n PeriodDiscriminator(p, in_channels, out_channels, **kwargs) for p in periods\n ])\n\n @property\n def num_discriminators(self):\n return len(self.discriminators)\n\n def forward(self, x: torch.Tensor) -> MultiDiscriminatorOutputType:\n logits = []\n fmaps = []\n for disc in self.discriminators:\n logit, fmap = disc(x)\n logits.append(logit)\n fmaps.append(fmap)\n return logits, fmaps\n" }, { "path": "audiocraft/adversarial/discriminators/msd.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nfrom ...modules import NormConv1d\nfrom .base import MultiDiscriminator, MultiDiscriminatorOutputType\n\n\nclass ScaleDiscriminator(nn.Module):\n \"\"\"Waveform sub-discriminator.\n\n Args:\n in_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n kernel_sizes (Sequence[int]): Kernel sizes for first and last convolutions.\n filters (int): Number of initial filters for convolutions.\n max_filters (int): Maximum number of filters.\n downsample_scales (Sequence[int]): Scale for downsampling implemented as strided convolutions.\n inner_kernel_sizes (Sequence[int] or None): Kernel sizes for inner convolutions.\n groups (Sequence[int] or None): Groups for inner convolutions.\n strides (Sequence[int] or None): Strides for inner convolutions.\n paddings (Sequence[int] or None): Paddings for inner convolutions.\n norm (str): Normalization method.\n activation (str): Activation function.\n activation_params (dict): Parameters to provide to the activation function.\n pad (str): Padding for initial convolution.\n pad_params (dict): Parameters to provide to the padding module.\n \"\"\"\n def __init__(self, in_channels=1, out_channels=1, kernel_sizes: tp.Sequence[int] = [5, 3],\n filters: int = 16, max_filters: int = 1024, downsample_scales: tp.Sequence[int] = [4, 4, 4, 4],\n inner_kernel_sizes: tp.Optional[tp.Sequence[int]] = None, groups: tp.Optional[tp.Sequence[int]] = None,\n strides: tp.Optional[tp.Sequence[int]] = None, paddings: tp.Optional[tp.Sequence[int]] = None,\n norm: str = 'weight_norm', activation: str = 'LeakyReLU',\n activation_params: dict = {'negative_slope': 0.2}, pad: str = 'ReflectionPad1d',\n pad_params: dict = {}):\n super().__init__()\n assert len(kernel_sizes) == 2\n assert kernel_sizes[0] % 2 == 1\n assert kernel_sizes[1] % 2 == 1\n assert (inner_kernel_sizes is None or len(inner_kernel_sizes) == len(downsample_scales))\n assert (groups is None or len(groups) == len(downsample_scales))\n assert (strides is None or len(strides) == len(downsample_scales))\n assert (paddings is None or len(paddings) == len(downsample_scales))\n self.activation = getattr(torch.nn, activation)(**activation_params)\n self.convs = nn.ModuleList()\n self.convs.append(\n nn.Sequential(\n getattr(torch.nn, pad)((np.prod(kernel_sizes) - 1) // 2, **pad_params),\n NormConv1d(in_channels, filters, kernel_size=np.prod(kernel_sizes), stride=1, norm=norm)\n )\n )\n\n in_chs = filters\n for i, downsample_scale in enumerate(downsample_scales):\n out_chs = min(in_chs * downsample_scale, max_filters)\n default_kernel_size = downsample_scale * 10 + 1\n default_stride = downsample_scale\n default_padding = (default_kernel_size - 1) // 2\n default_groups = in_chs // 4\n self.convs.append(\n NormConv1d(in_chs, out_chs,\n kernel_size=inner_kernel_sizes[i] if inner_kernel_sizes else default_kernel_size,\n stride=strides[i] if strides else default_stride,\n groups=groups[i] if groups else default_groups,\n padding=paddings[i] if paddings else default_padding,\n norm=norm))\n in_chs = out_chs\n\n out_chs = min(in_chs * 2, max_filters)\n self.convs.append(NormConv1d(in_chs, out_chs, kernel_size=kernel_sizes[0], stride=1,\n padding=(kernel_sizes[0] - 1) // 2, norm=norm))\n self.conv_post = NormConv1d(out_chs, out_channels, kernel_size=kernel_sizes[1], stride=1,\n padding=(kernel_sizes[1] - 1) // 2, norm=norm)\n\n def forward(self, x: torch.Tensor):\n fmap = []\n for layer in self.convs:\n x = layer(x)\n x = self.activation(x)\n fmap.append(x)\n x = self.conv_post(x)\n fmap.append(x)\n # x = torch.flatten(x, 1, -1)\n return x, fmap\n\n\nclass MultiScaleDiscriminator(MultiDiscriminator):\n \"\"\"Multi-Scale (MSD) Discriminator,\n\n Args:\n in_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n downsample_factor (int): Downsampling factor between the different scales.\n scale_norms (Sequence[str]): Normalization for each sub-discriminator.\n **kwargs: Additional args for ScaleDiscriminator.\n \"\"\"\n def __init__(self, in_channels: int = 1, out_channels: int = 1, downsample_factor: int = 2,\n scale_norms: tp.Sequence[str] = ['weight_norm', 'weight_norm', 'weight_norm'], **kwargs):\n super().__init__()\n self.discriminators = nn.ModuleList([\n ScaleDiscriminator(in_channels, out_channels, norm=norm, **kwargs) for norm in scale_norms\n ])\n self.downsample = nn.AvgPool1d(downsample_factor * 2, downsample_factor, padding=downsample_factor)\n\n @property\n def num_discriminators(self):\n return len(self.discriminators)\n\n def forward(self, x: torch.Tensor) -> MultiDiscriminatorOutputType:\n logits = []\n fmaps = []\n for i, disc in enumerate(self.discriminators):\n if i != 0:\n self.downsample(x)\n logit, fmap = disc(x)\n logits.append(logit)\n fmaps.append(fmap)\n return logits, fmaps\n" }, { "path": "audiocraft/adversarial/discriminators/msstftd.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nimport torchaudio\nimport torch\nfrom torch import nn\nfrom einops import rearrange\n\nfrom ...modules import NormConv2d\nfrom .base import MultiDiscriminator, MultiDiscriminatorOutputType\n\n\ndef get_2d_padding(kernel_size: tp.Tuple[int, int], dilation: tp.Tuple[int, int] = (1, 1)):\n return (((kernel_size[0] - 1) * dilation[0]) // 2, ((kernel_size[1] - 1) * dilation[1]) // 2)\n\n\nclass DiscriminatorSTFT(nn.Module):\n \"\"\"STFT sub-discriminator.\n\n Args:\n filters (int): Number of filters in convolutions.\n in_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n n_fft (int): Size of FFT for each scale.\n hop_length (int): Length of hop between STFT windows for each scale.\n kernel_size (tuple of int): Inner Conv2d kernel sizes.\n stride (tuple of int): Inner Conv2d strides.\n dilations (list of int): Inner Conv2d dilation on the time dimension.\n win_length (int): Window size for each scale.\n normalized (bool): Whether to normalize by magnitude after stft.\n norm (str): Normalization method.\n activation (str): Activation function.\n activation_params (dict): Parameters to provide to the activation function.\n growth (int): Growth factor for the filters.\n \"\"\"\n def __init__(self, filters: int, in_channels: int = 1, out_channels: int = 1,\n n_fft: int = 1024, hop_length: int = 256, win_length: int = 1024, max_filters: int = 1024,\n filters_scale: int = 1, kernel_size: tp.Tuple[int, int] = (3, 9), dilations: tp.List = [1, 2, 4],\n stride: tp.Tuple[int, int] = (1, 2), normalized: bool = True, norm: str = 'weight_norm',\n activation: str = 'LeakyReLU', activation_params: dict = {'negative_slope': 0.2}):\n super().__init__()\n assert len(kernel_size) == 2\n assert len(stride) == 2\n self.filters = filters\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.n_fft = n_fft\n self.hop_length = hop_length\n self.win_length = win_length\n self.normalized = normalized\n self.activation = getattr(torch.nn, activation)(**activation_params)\n self.spec_transform = torchaudio.transforms.Spectrogram(\n n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, window_fn=torch.hann_window,\n normalized=self.normalized, center=False, pad_mode=None, power=None)\n spec_channels = 2 * self.in_channels\n self.convs = nn.ModuleList()\n self.convs.append(\n NormConv2d(spec_channels, self.filters, kernel_size=kernel_size, padding=get_2d_padding(kernel_size))\n )\n in_chs = min(filters_scale * self.filters, max_filters)\n for i, dilation in enumerate(dilations):\n out_chs = min((filters_scale ** (i + 1)) * self.filters, max_filters)\n self.convs.append(NormConv2d(in_chs, out_chs, kernel_size=kernel_size, stride=stride,\n dilation=(dilation, 1), padding=get_2d_padding(kernel_size, (dilation, 1)),\n norm=norm))\n in_chs = out_chs\n out_chs = min((filters_scale ** (len(dilations) + 1)) * self.filters, max_filters)\n self.convs.append(NormConv2d(in_chs, out_chs, kernel_size=(kernel_size[0], kernel_size[0]),\n padding=get_2d_padding((kernel_size[0], kernel_size[0])),\n norm=norm))\n self.conv_post = NormConv2d(out_chs, self.out_channels,\n kernel_size=(kernel_size[0], kernel_size[0]),\n padding=get_2d_padding((kernel_size[0], kernel_size[0])),\n norm=norm)\n\n def forward(self, x: torch.Tensor):\n fmap = []\n z = self.spec_transform(x) # [B, 2, Freq, Frames, 2]\n z = torch.cat([z.real, z.imag], dim=1)\n z = rearrange(z, 'b c w t -> b c t w')\n for i, layer in enumerate(self.convs):\n z = layer(z)\n z = self.activation(z)\n fmap.append(z)\n z = self.conv_post(z)\n return z, fmap\n\n\nclass MultiScaleSTFTDiscriminator(MultiDiscriminator):\n \"\"\"Multi-Scale STFT (MS-STFT) discriminator.\n\n Args:\n filters (int): Number of filters in convolutions.\n in_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n sep_channels (bool): Separate channels to distinct samples for stereo support.\n n_ffts (Sequence[int]): Size of FFT for each scale.\n hop_lengths (Sequence[int]): Length of hop between STFT windows for each scale.\n win_lengths (Sequence[int]): Window size for each scale.\n **kwargs: Additional args for STFTDiscriminator.\n \"\"\"\n def __init__(self, filters: int, in_channels: int = 1, out_channels: int = 1, sep_channels: bool = False,\n n_ffts: tp.List[int] = [1024, 2048, 512], hop_lengths: tp.List[int] = [256, 512, 128],\n win_lengths: tp.List[int] = [1024, 2048, 512], **kwargs):\n super().__init__()\n assert len(n_ffts) == len(hop_lengths) == len(win_lengths)\n self.sep_channels = sep_channels\n self.discriminators = nn.ModuleList([\n DiscriminatorSTFT(filters, in_channels=in_channels, out_channels=out_channels,\n n_fft=n_ffts[i], win_length=win_lengths[i], hop_length=hop_lengths[i], **kwargs)\n for i in range(len(n_ffts))\n ])\n\n @property\n def num_discriminators(self):\n return len(self.discriminators)\n\n def _separate_channels(self, x: torch.Tensor) -> torch.Tensor:\n B, C, T = x.shape\n return x.view(-1, 1, T)\n\n def forward(self, x: torch.Tensor) -> MultiDiscriminatorOutputType:\n logits = []\n fmaps = []\n for disc in self.discriminators:\n logit, fmap = disc(x)\n logits.append(logit)\n fmaps.append(fmap)\n return logits, fmaps\n" }, { "path": "audiocraft/adversarial/losses.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nUtility module to handle adversarial losses without requiring to mess up the main training loop.\n\"\"\"\n\nimport typing as tp\n\nimport flashy\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nADVERSARIAL_LOSSES = ['mse', 'hinge', 'hinge2']\n\n\nAdvLossType = tp.Union[nn.Module, tp.Callable[[torch.Tensor], torch.Tensor]]\nFeatLossType = tp.Union[nn.Module, tp.Callable[[torch.Tensor, torch.Tensor], torch.Tensor]]\n\n\nclass AdversarialLoss(nn.Module):\n \"\"\"Adversary training wrapper.\n\n Args:\n adversary (nn.Module): The adversary module will be used to estimate the logits given the fake and real samples.\n We assume here the adversary output is ``Tuple[List[torch.Tensor], List[List[torch.Tensor]]]``\n where the first item is a list of logits and the second item is a list of feature maps.\n optimizer (torch.optim.Optimizer): Optimizer used for training the given module.\n loss (AdvLossType): Loss function for generator training.\n loss_real (AdvLossType): Loss function for adversarial training on logits from real samples.\n loss_fake (AdvLossType): Loss function for adversarial training on logits from fake samples.\n loss_feat (FeatLossType): Feature matching loss function for generator training.\n normalize (bool): Whether to normalize by number of sub-discriminators.\n\n Example of usage:\n adv_loss = AdversarialLoss(adversaries, optimizer, loss, loss_real, loss_fake)\n for real in loader:\n noise = torch.randn(...)\n fake = model(noise)\n adv_loss.train_adv(fake, real)\n loss, _ = adv_loss(fake, real)\n loss.backward()\n \"\"\"\n def __init__(self,\n adversary: nn.Module,\n optimizer: torch.optim.Optimizer,\n loss: AdvLossType,\n loss_real: AdvLossType,\n loss_fake: AdvLossType,\n loss_feat: tp.Optional[FeatLossType] = None,\n normalize: bool = True):\n super().__init__()\n self.adversary: nn.Module = adversary\n flashy.distrib.broadcast_model(self.adversary)\n self.optimizer = optimizer\n self.loss = loss\n self.loss_real = loss_real\n self.loss_fake = loss_fake\n self.loss_feat = loss_feat\n self.normalize = normalize\n\n def _save_to_state_dict(self, destination, prefix, keep_vars):\n # Add the optimizer state dict inside our own.\n super()._save_to_state_dict(destination, prefix, keep_vars)\n destination[prefix + 'optimizer'] = self.optimizer.state_dict()\n return destination\n\n def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):\n # Load optimizer state.\n self.optimizer.load_state_dict(state_dict.pop(prefix + 'optimizer'))\n super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)\n\n def get_adversary_pred(self, x):\n \"\"\"Run adversary model, validating expected output format.\"\"\"\n logits, fmaps = self.adversary(x)\n assert isinstance(logits, list) and all([isinstance(t, torch.Tensor) for t in logits]), \\\n f'Expecting a list of tensors as logits but {type(logits)} found.'\n assert isinstance(fmaps, list), f'Expecting a list of features maps but {type(fmaps)} found.'\n for fmap in fmaps:\n assert isinstance(fmap, list) and all([isinstance(f, torch.Tensor) for f in fmap]), \\\n f'Expecting a list of tensors as feature maps but {type(fmap)} found.'\n return logits, fmaps\n\n def train_adv(self, fake: torch.Tensor, real: torch.Tensor) -> torch.Tensor:\n \"\"\"Train the adversary with the given fake and real example.\n\n We assume the adversary output is the following format: Tuple[List[torch.Tensor], List[List[torch.Tensor]]].\n The first item being the logits and second item being a list of feature maps for each sub-discriminator.\n\n This will automatically synchronize gradients (with `flashy.distrib.eager_sync_model`)\n and call the optimizer.\n \"\"\"\n loss = torch.tensor(0., device=fake.device)\n all_logits_fake_is_fake, _ = self.get_adversary_pred(fake.detach())\n all_logits_real_is_fake, _ = self.get_adversary_pred(real.detach())\n n_sub_adversaries = len(all_logits_fake_is_fake)\n for logit_fake_is_fake, logit_real_is_fake in zip(all_logits_fake_is_fake, all_logits_real_is_fake):\n loss += self.loss_fake(logit_fake_is_fake) + self.loss_real(logit_real_is_fake)\n\n if self.normalize:\n loss /= n_sub_adversaries\n\n self.optimizer.zero_grad()\n with flashy.distrib.eager_sync_model(self.adversary):\n loss.backward()\n self.optimizer.step()\n\n return loss\n\n def forward(self, fake: torch.Tensor, real: torch.Tensor) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Return the loss for the generator, i.e. trying to fool the adversary,\n and feature matching loss if provided.\n \"\"\"\n adv = torch.tensor(0., device=fake.device)\n feat = torch.tensor(0., device=fake.device)\n with flashy.utils.readonly(self.adversary):\n all_logits_fake_is_fake, all_fmap_fake = self.get_adversary_pred(fake)\n all_logits_real_is_fake, all_fmap_real = self.get_adversary_pred(real)\n n_sub_adversaries = len(all_logits_fake_is_fake)\n for logit_fake_is_fake in all_logits_fake_is_fake:\n adv += self.loss(logit_fake_is_fake)\n if self.loss_feat:\n for fmap_fake, fmap_real in zip(all_fmap_fake, all_fmap_real):\n feat += self.loss_feat(fmap_fake, fmap_real)\n\n if self.normalize:\n adv /= n_sub_adversaries\n feat /= n_sub_adversaries\n\n return adv, feat\n\n\ndef get_adv_criterion(loss_type: str) -> tp.Callable:\n assert loss_type in ADVERSARIAL_LOSSES\n if loss_type == 'mse':\n return mse_loss\n elif loss_type == 'hinge':\n return hinge_loss\n elif loss_type == 'hinge2':\n return hinge2_loss\n raise ValueError('Unsupported loss')\n\n\ndef get_fake_criterion(loss_type: str) -> tp.Callable:\n assert loss_type in ADVERSARIAL_LOSSES\n if loss_type == 'mse':\n return mse_fake_loss\n elif loss_type in ['hinge', 'hinge2']:\n return hinge_fake_loss\n raise ValueError('Unsupported loss')\n\n\ndef get_real_criterion(loss_type: str) -> tp.Callable:\n assert loss_type in ADVERSARIAL_LOSSES\n if loss_type == 'mse':\n return mse_real_loss\n elif loss_type in ['hinge', 'hinge2']:\n return hinge_real_loss\n raise ValueError('Unsupported loss')\n\n\ndef mse_real_loss(x: torch.Tensor) -> torch.Tensor:\n return F.mse_loss(x, torch.tensor(1., device=x.device).expand_as(x))\n\n\ndef mse_fake_loss(x: torch.Tensor) -> torch.Tensor:\n return F.mse_loss(x, torch.tensor(0., device=x.device).expand_as(x))\n\n\ndef hinge_real_loss(x: torch.Tensor) -> torch.Tensor:\n return -torch.mean(torch.min(x - 1, torch.tensor(0., device=x.device).expand_as(x)))\n\n\ndef hinge_fake_loss(x: torch.Tensor) -> torch.Tensor:\n return -torch.mean(torch.min(-x - 1, torch.tensor(0., device=x.device).expand_as(x)))\n\n\ndef mse_loss(x: torch.Tensor) -> torch.Tensor:\n if x.numel() == 0:\n return torch.tensor([0.0], device=x.device)\n return F.mse_loss(x, torch.tensor(1., device=x.device).expand_as(x))\n\n\ndef hinge_loss(x: torch.Tensor) -> torch.Tensor:\n if x.numel() == 0:\n return torch.tensor([0.0], device=x.device)\n return -x.mean()\n\n\ndef hinge2_loss(x: torch.Tensor) -> torch.Tensor:\n if x.numel() == 0:\n return torch.tensor([0.0])\n return -torch.mean(torch.min(x - 1, torch.tensor(0., device=x.device).expand_as(x)))\n\n\nclass FeatureMatchingLoss(nn.Module):\n \"\"\"Feature matching loss for adversarial training.\n\n Args:\n loss (nn.Module): Loss to use for feature matching (default=torch.nn.L1).\n normalize (bool): Whether to normalize the loss.\n by number of feature maps.\n \"\"\"\n def __init__(self, loss: nn.Module = torch.nn.L1Loss(), normalize: bool = True):\n super().__init__()\n self.loss = loss\n self.normalize = normalize\n\n def forward(self, fmap_fake: tp.List[torch.Tensor], fmap_real: tp.List[torch.Tensor]) -> torch.Tensor:\n assert len(fmap_fake) == len(fmap_real) and len(fmap_fake) > 0\n feat_loss = torch.tensor(0., device=fmap_fake[0].device)\n feat_scale = torch.tensor(0., device=fmap_fake[0].device)\n n_fmaps = 0\n for (feat_fake, feat_real) in zip(fmap_fake, fmap_real):\n assert feat_fake.shape == feat_real.shape\n n_fmaps += 1\n feat_loss += self.loss(feat_fake, feat_real)\n feat_scale += torch.mean(torch.abs(feat_real))\n\n if self.normalize:\n feat_loss /= n_fmaps\n\n return feat_loss\n" }, { "path": "audiocraft/data/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Audio loading and writing support. Datasets for raw audio\nor also including some metadata.\"\"\"\n\n# flake8: noqa\nfrom . import audio, audio_dataset, info_audio_dataset, music_dataset, sound_dataset, jasco_dataset\n" }, { "path": "audiocraft/data/audio.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nAudio IO methods are defined in this module (info, read, write),\nWe rely on av library for faster read when possible, otherwise on torchaudio.\n\"\"\"\n\nfrom dataclasses import dataclass\nfrom pathlib import Path\nimport logging\nimport typing as tp\n\nimport numpy as np\nimport soundfile\nimport torch\nfrom torch.nn import functional as F\n\nimport av\nimport subprocess as sp\n\nfrom .audio_utils import f32_pcm, normalize_audio\n\n\n_av_initialized = False\n\n\ndef _init_av():\n global _av_initialized\n if _av_initialized:\n return\n logger = logging.getLogger('libav.mp3')\n logger.setLevel(logging.ERROR)\n _av_initialized = True\n\n\n@dataclass(frozen=True)\nclass AudioFileInfo:\n sample_rate: int\n duration: float\n channels: int\n\n\ndef _av_info(filepath: tp.Union[str, Path]) -> AudioFileInfo:\n _init_av()\n with av.open(str(filepath)) as af:\n stream = af.streams.audio[0]\n sample_rate = stream.codec_context.sample_rate\n duration = float(stream.duration * stream.time_base)\n channels = stream.channels\n return AudioFileInfo(sample_rate, duration, channels)\n\n\ndef _soundfile_info(filepath: tp.Union[str, Path]) -> AudioFileInfo:\n info = soundfile.info(filepath)\n return AudioFileInfo(info.samplerate, info.duration, info.channels)\n\n\ndef audio_info(filepath: tp.Union[str, Path]) -> AudioFileInfo:\n # torchaudio no longer returns useful duration informations for some formats like mp3s.\n filepath = Path(filepath)\n if filepath.suffix in ['.flac', '.ogg']: # TODO: Validate .ogg can be safely read with av_info\n # ffmpeg has some weird issue with flac.\n return _soundfile_info(filepath)\n else:\n return _av_info(filepath)\n\n\ndef _av_read(filepath: tp.Union[str, Path], seek_time: float = 0, duration: float = -1.) -> tp.Tuple[torch.Tensor, int]:\n \"\"\"FFMPEG-based audio file reading using PyAV bindings.\n Soundfile cannot read mp3 and av_read is more efficient than torchaudio.\n\n Args:\n filepath (str or Path): Path to audio file to read.\n seek_time (float): Time at which to start reading in the file.\n duration (float): Duration to read from the file. If set to -1, the whole file is read.\n Returns:\n tuple of torch.Tensor, int: Tuple containing audio data and sample rate\n \"\"\"\n _init_av()\n with av.open(str(filepath)) as af:\n stream = af.streams.audio[0]\n sr = stream.codec_context.sample_rate\n num_frames = int(sr * duration) if duration >= 0 else -1\n frame_offset = int(sr * seek_time)\n # we need a small negative offset otherwise we get some edge artifact\n # from the mp3 decoder.\n af.seek(int(max(0, (seek_time - 0.1)) / stream.time_base), stream=stream)\n frames = []\n length = 0\n for frame in af.decode(streams=stream.index):\n current_offset = int(frame.rate * frame.pts * frame.time_base)\n strip = max(0, frame_offset - current_offset)\n buf = torch.from_numpy(frame.to_ndarray())\n if buf.shape[0] != stream.channels:\n buf = buf.view(-1, stream.channels).t()\n buf = buf[:, strip:]\n frames.append(buf)\n length += buf.shape[1]\n if num_frames > 0 and length >= num_frames:\n break\n assert frames\n # If the above assert fails, it is likely because we seeked past the end of file point,\n # in which case ffmpeg returns a single frame with only zeros, and a weird timestamp.\n # This will need proper debugging, in due time.\n wav = torch.cat(frames, dim=1)\n assert wav.shape[0] == stream.channels\n if num_frames > 0:\n wav = wav[:, :num_frames]\n return f32_pcm(wav), sr\n\n\ndef audio_read(filepath: tp.Union[str, Path], seek_time: float = 0.,\n duration: float = -1.0, pad: bool = False) -> tp.Tuple[torch.Tensor, int]:\n \"\"\"Read audio by picking the most appropriate backend tool based on the audio format.\n\n Args:\n filepath (str or Path): Path to audio file to read.\n seek_time (float): Time at which to start reading in the file.\n duration (float): Duration to read from the file. If set to -1, the whole file is read.\n pad (bool): Pad output audio if not reaching expected duration.\n Returns:\n tuple of torch.Tensor, int: Tuple containing audio data and sample rate.\n \"\"\"\n fp = Path(filepath)\n if fp.suffix in ['.flac', '.ogg']: # TODO: check if we can safely use av_read for .ogg\n # There is some bug with ffmpeg and reading flac\n info = _soundfile_info(filepath)\n frames = -1 if duration <= 0 else int(duration * info.sample_rate)\n frame_offset = int(seek_time * info.sample_rate)\n wav, sr = soundfile.read(filepath, start=frame_offset, frames=frames, dtype=np.float32)\n assert info.sample_rate == sr, f\"Mismatch of sample rates {info.sample_rate} {sr}\"\n wav = torch.from_numpy(wav).t().contiguous()\n if len(wav.shape) == 1:\n wav = torch.unsqueeze(wav, 0)\n else:\n wav, sr = _av_read(filepath, seek_time, duration)\n if pad and duration > 0:\n expected_frames = int(duration * sr)\n wav = F.pad(wav, (0, expected_frames - wav.shape[-1]))\n return wav, sr\n\n\ndef _piping_to_ffmpeg(out_path: tp.Union[str, Path], wav: torch.Tensor, sample_rate: int, flags: tp.List[str]):\n # ffmpeg is always installed and torchaudio is a bit unstable lately, so let's bypass it entirely.\n assert wav.dim() == 2, wav.shape\n command = [\n 'ffmpeg',\n '-loglevel', 'error',\n '-y', '-f', 'f32le', '-ar', str(sample_rate), '-ac', str(wav.shape[0]),\n '-i', '-'] + flags + [str(out_path)]\n input_ = f32_pcm(wav).t().detach().cpu().numpy().tobytes()\n sp.run(command, input=input_, check=True)\n\n\ndef audio_write(stem_name: tp.Union[str, Path],\n wav: torch.Tensor, sample_rate: int,\n format: str = 'wav', mp3_rate: int = 320, ogg_rate: tp.Optional[int] = None,\n normalize: bool = True, strategy: str = 'peak', peak_clip_headroom_db: float = 1,\n rms_headroom_db: float = 18, loudness_headroom_db: float = 14,\n loudness_compressor: bool = False,\n log_clipping: bool = True, make_parent_dir: bool = True,\n add_suffix: bool = True) -> Path:\n \"\"\"Convenience function for saving audio to disk. Returns the filename the audio was written to.\n\n Args:\n stem_name (str or Path): Filename without extension which will be added automatically.\n wav (torch.Tensor): Audio data to save.\n sample_rate (int): Sample rate of audio data.\n format (str): Either \"wav\", \"mp3\", \"ogg\", or \"flac\".\n mp3_rate (int): kbps when using mp3s.\n ogg_rate (int): kbps when using ogg/vorbis. If not provided, let ffmpeg decide for itself.\n normalize (bool): if `True` (default), normalizes according to the prescribed\n strategy (see after). If `False`, the strategy is only used in case clipping\n would happen.\n strategy (str): Can be either 'clip', 'peak', or 'rms'. Default is 'peak',\n i.e. audio is normalized by its largest value. RMS normalizes by root-mean-square\n with extra headroom to avoid clipping. 'clip' just clips.\n peak_clip_headroom_db (float): Headroom in dB when doing 'peak' or 'clip' strategy.\n rms_headroom_db (float): Headroom in dB when doing 'rms' strategy. This must be much larger\n than the `peak_clip` one to avoid further clipping.\n loudness_headroom_db (float): Target loudness for loudness normalization.\n loudness_compressor (bool): Uses tanh for soft clipping when strategy is 'loudness'.\n when strategy is 'loudness' log_clipping (bool): If True, basic logging on stderr when clipping still\n occurs despite strategy (only for 'rms').\n make_parent_dir (bool): Make parent directory if it doesn't exist.\n Returns:\n Path: Path of the saved audio.\n \"\"\"\n assert wav.dtype.is_floating_point, \"wav is not floating point\"\n if wav.dim() == 1:\n wav = wav[None]\n elif wav.dim() > 2:\n raise ValueError(\"Input wav should be at most 2 dimension.\")\n assert wav.isfinite().all()\n wav = normalize_audio(wav, normalize, strategy, peak_clip_headroom_db,\n rms_headroom_db, loudness_headroom_db, loudness_compressor,\n log_clipping=log_clipping, sample_rate=sample_rate,\n stem_name=str(stem_name))\n if format == 'mp3':\n suffix = '.mp3'\n flags = ['-f', 'mp3', '-c:a', 'libmp3lame', '-b:a', f'{mp3_rate}k']\n elif format == 'wav':\n suffix = '.wav'\n flags = ['-f', 'wav', '-c:a', 'pcm_s16le']\n elif format == 'ogg':\n suffix = '.ogg'\n flags = ['-f', 'ogg', '-c:a', 'libvorbis']\n if ogg_rate is not None:\n flags += ['-b:a', f'{ogg_rate}k']\n elif format == 'flac':\n suffix = '.flac'\n flags = ['-f', 'flac']\n else:\n raise RuntimeError(f\"Invalid format {format}. Only wav or mp3 are supported.\")\n if not add_suffix:\n suffix = ''\n path = Path(str(stem_name) + suffix)\n if make_parent_dir:\n path.parent.mkdir(exist_ok=True, parents=True)\n try:\n _piping_to_ffmpeg(path, wav, sample_rate, flags)\n except Exception:\n if path.exists():\n # we do not want to leave half written files around.\n path.unlink()\n raise\n return path\n\n\ndef get_spec(y, sr=16000, n_fft=4096, hop_length=128, dur=8) -> np.ndarray:\n \"\"\"Get the mel-spectrogram from the raw audio.\n\n Args:\n y (numpy array): raw input\n sr (int): Sampling rate\n n_fft (int): Number of samples per FFT. Default is 2048.\n hop_length (int): Number of samples between successive frames. Default is 512.\n dur (float): Maxium duration to get the spectrograms\n Returns:\n spectro histogram as a numpy array\n \"\"\"\n import librosa\n import librosa.display\n\n spectrogram = librosa.feature.melspectrogram(\n y=y, sr=sr, n_fft=n_fft, hop_length=hop_length\n )\n spectrogram_db = librosa.power_to_db(spectrogram, ref=np.max)\n return spectrogram_db\n\n\ndef save_spectrograms(\n ys: tp.List[np.ndarray],\n sr: int,\n path: str,\n names: tp.List[str],\n n_fft: int = 4096,\n hop_length: int = 128,\n dur: float = 8.0,\n):\n \"\"\"Plot a spectrogram for an audio file.\n\n Args:\n ys: List of audio spectrograms\n sr (int): Sampling rate of the audio file. Default is 22050 Hz.\n path (str): Path to the plot file.\n names: name of each spectrogram plot\n n_fft (int): Number of samples per FFT. Default is 2048.\n hop_length (int): Number of samples between successive frames. Default is 512.\n dur (float): Maxium duration to plot the spectrograms\n\n Returns:\n None (plots the spectrogram using matplotlib)\n \"\"\"\n import matplotlib as mpl # type: ignore\n import matplotlib.pyplot as plt # type: ignore\n import librosa.display\n\n if not names:\n names = [\"Ground Truth\", \"Audio Watermarked\", \"Watermark\"]\n ys = [wav[: int(dur * sr)] for wav in ys] # crop\n assert len(names) == len(\n ys\n ), f\"There are {len(ys)} wavs but {len(names)} names ({names})\"\n\n # Set matplotlib stuff\n BIGGER_SIZE = 10\n SMALLER_SIZE = 8\n linewidth = 234.8775 # linewidth in pt\n\n plt.rc(\"font\", size=BIGGER_SIZE, family=\"serif\") # controls default text sizes\n plt.rcParams[\"font.family\"] = \"DeJavu Serif\"\n plt.rcParams[\"font.serif\"] = [\"Times New Roman\"]\n\n plt.rc(\"axes\", titlesize=BIGGER_SIZE) # fontsize of the axes title\n plt.rc(\"axes\", labelsize=BIGGER_SIZE) # fontsize of the x and y labels\n plt.rc(\"xtick\", labelsize=BIGGER_SIZE) # fontsize of the tick labels\n plt.rc(\"ytick\", labelsize=SMALLER_SIZE) # fontsize of the tick labels\n plt.rc(\"legend\", fontsize=BIGGER_SIZE) # legend fontsize\n plt.rc(\"figure\", titlesize=BIGGER_SIZE)\n height = 1.6 * linewidth / 72.0\n fig, ax = plt.subplots(\n nrows=len(ys),\n ncols=1,\n sharex=True,\n figsize=(linewidth / 72.0, height),\n )\n fig.tight_layout()\n\n # Plot the spectrogram\n\n for i, ysi in enumerate(ys):\n spectrogram_db = get_spec(ysi, sr=sr, n_fft=n_fft, hop_length=hop_length)\n if i == 0:\n cax = fig.add_axes(\n [\n ax[0].get_position().x1 + 0.01, # type: ignore\n ax[-1].get_position().y0,\n 0.02,\n ax[0].get_position().y1 - ax[-1].get_position().y0,\n ]\n )\n fig.colorbar(\n mpl.cm.ScalarMappable(\n norm=mpl.colors.Normalize(\n np.min(spectrogram_db), np.max(spectrogram_db)\n ),\n cmap=\"magma\",\n ),\n ax=ax,\n orientation=\"vertical\",\n format=\"%+2.0f dB\",\n cax=cax,\n )\n librosa.display.specshow(\n spectrogram_db,\n sr=sr,\n hop_length=hop_length,\n x_axis=\"time\",\n y_axis=\"mel\",\n ax=ax[i],\n )\n ax[i].set(title=names[i])\n ax[i].yaxis.set_label_text(None)\n ax[i].label_outer()\n fig.savefig(path, bbox_inches=\"tight\")\n plt.close()\n" }, { "path": "audiocraft/data/audio_dataset.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"AudioDataset support. In order to handle a larger number of files\nwithout having to scan again the folders, we precompute some metadata\n(filename, sample rate, duration), and use that to efficiently sample audio segments.\n\"\"\"\nimport argparse\nimport copy\nfrom concurrent.futures import ThreadPoolExecutor, Future\nfrom dataclasses import dataclass, fields\nfrom contextlib import ExitStack\nfrom functools import lru_cache\nimport gzip\nimport json\nimport logging\nimport os\nfrom pathlib import Path\nimport random\nimport sys\nimport typing as tp\n\nimport torch\nimport torch.nn.functional as F\n\nfrom .audio import audio_read, audio_info\nfrom .audio_utils import convert_audio\nfrom .zip import PathInZip\n\ntry:\n import dora\nexcept ImportError:\n dora = None # type: ignore\n\n\n@dataclass(order=True)\nclass BaseInfo:\n\n @classmethod\n def _dict2fields(cls, dictionary: dict):\n return {\n field.name: dictionary[field.name]\n for field in fields(cls) if field.name in dictionary\n }\n\n @classmethod\n def from_dict(cls, dictionary: dict):\n _dictionary = cls._dict2fields(dictionary)\n return cls(**_dictionary)\n\n def to_dict(self):\n return {\n field.name: self.__getattribute__(field.name)\n for field in fields(self)\n }\n\n\n@dataclass(order=True)\nclass AudioMeta(BaseInfo):\n path: str\n duration: float\n sample_rate: int\n amplitude: tp.Optional[float] = None\n weight: tp.Optional[float] = None\n # info_path is used to load additional information about the audio file that is stored in zip files.\n info_path: tp.Optional[PathInZip] = None\n\n @classmethod\n def from_dict(cls, dictionary: dict):\n base = cls._dict2fields(dictionary)\n if 'info_path' in base and base['info_path'] is not None:\n base['info_path'] = PathInZip(base['info_path'])\n return cls(**base)\n\n def to_dict(self):\n d = super().to_dict()\n if d['info_path'] is not None:\n d['info_path'] = str(d['info_path'])\n return d\n\n\n@dataclass(order=True)\nclass SegmentInfo(BaseInfo):\n meta: AudioMeta\n seek_time: float\n # The following values are given once the audio is processed, e.g.\n # at the target sample rate and target number of channels.\n n_frames: int # actual number of frames without padding\n total_frames: int # total number of frames, padding included\n sample_rate: int # actual sample rate\n channels: int # number of audio channels.\n\n\nDEFAULT_EXTS = ['.wav', '.mp3', '.flac', '.ogg', '.m4a']\n\nlogger = logging.getLogger(__name__)\n\n\ndef _get_audio_meta(file_path: str, minimal: bool = True) -> AudioMeta:\n \"\"\"AudioMeta from a path to an audio file.\n\n Args:\n file_path (str): Resolved path of valid audio file.\n minimal (bool): Whether to only load the minimal set of metadata (takes longer if not).\n Returns:\n AudioMeta: Audio file path and its metadata.\n \"\"\"\n info = audio_info(file_path)\n amplitude: tp.Optional[float] = None\n if not minimal:\n wav, sr = audio_read(file_path)\n amplitude = wav.abs().max().item()\n return AudioMeta(file_path, info.duration, info.sample_rate, amplitude)\n\n\ndef _resolve_audio_meta(m: AudioMeta, fast: bool = True) -> AudioMeta:\n \"\"\"If Dora is available as a dependency, try to resolve potential relative paths\n in list of AudioMeta. This method is expected to be used when loading meta from file.\n\n Args:\n m (AudioMeta): Audio meta to resolve.\n fast (bool): If True, uses a really fast check for determining if a file\n is already absolute or not. Only valid on Linux/Mac.\n Returns:\n AudioMeta: Audio meta with resolved path.\n \"\"\"\n def is_abs(m):\n if fast:\n return str(m)[0] == '/'\n else:\n os.path.isabs(str(m))\n\n if not dora:\n return m\n\n if not is_abs(m.path):\n m.path = dora.git_save.to_absolute_path(m.path)\n if m.info_path is not None and not is_abs(m.info_path.zip_path):\n m.info_path.zip_path = dora.git_save.to_absolute_path(m.path)\n return m\n\n\ndef find_audio_files(path: tp.Union[Path, str],\n exts: tp.List[str] = DEFAULT_EXTS,\n resolve: bool = True,\n minimal: bool = True,\n progress: bool = False,\n workers: int = 0) -> tp.List[AudioMeta]:\n \"\"\"Build a list of AudioMeta from a given path,\n collecting relevant audio files and fetching meta info.\n\n Args:\n path (str or Path): Path to folder containing audio files.\n exts (list of str): List of file extensions to consider for audio files.\n minimal (bool): Whether to only load the minimal set of metadata (takes longer if not).\n progress (bool): Whether to log progress on audio files collection.\n workers (int): number of parallel workers, if 0, use only the current thread.\n Returns:\n list of AudioMeta: List of audio file path and its metadata.\n \"\"\"\n audio_files = []\n futures: tp.List[Future] = []\n pool: tp.Optional[ThreadPoolExecutor] = None\n with ExitStack() as stack:\n if workers > 0:\n pool = ThreadPoolExecutor(workers)\n stack.enter_context(pool)\n\n if progress:\n print(\"Finding audio files...\")\n for root, folders, files in os.walk(path, followlinks=True):\n for file in files:\n full_path = Path(root) / file\n if full_path.suffix.lower() in exts:\n audio_files.append(full_path)\n if pool is not None:\n futures.append(pool.submit(_get_audio_meta, str(audio_files[-1]), minimal))\n if progress:\n print(format(len(audio_files), \" 8d\"), end='\\r', file=sys.stderr)\n\n if progress:\n print(\"Getting audio metadata...\")\n meta: tp.List[AudioMeta] = []\n for idx, file_path in enumerate(audio_files):\n try:\n if pool is None:\n m = _get_audio_meta(str(file_path), minimal)\n else:\n m = futures[idx].result()\n if resolve:\n m = _resolve_audio_meta(m)\n except Exception as err:\n print(\"Error with\", str(file_path), err, file=sys.stderr)\n continue\n meta.append(m)\n if progress:\n print(format((1 + idx) / len(audio_files), \" 3.1%\"), end='\\r', file=sys.stderr)\n meta.sort()\n return meta\n\n\ndef load_audio_meta(path: tp.Union[str, Path],\n resolve: bool = True, fast: bool = True) -> tp.List[AudioMeta]:\n \"\"\"Load list of AudioMeta from an optionally compressed json file.\n\n Args:\n path (str or Path): Path to JSON file.\n resolve (bool): Whether to resolve the path from AudioMeta (default=True).\n fast (bool): activates some tricks to make things faster.\n Returns:\n list of AudioMeta: List of audio file path and its total duration.\n \"\"\"\n open_fn = gzip.open if str(path).lower().endswith('.gz') else open\n with open_fn(path, 'rb') as fp: # type: ignore\n lines = fp.readlines()\n meta = []\n for line in lines:\n d = json.loads(line)\n m = AudioMeta.from_dict(d)\n if resolve:\n m = _resolve_audio_meta(m, fast=fast)\n meta.append(m)\n return meta\n\n\ndef save_audio_meta(path: tp.Union[str, Path], meta: tp.List[AudioMeta]):\n \"\"\"Save the audio metadata to the file pointer as json.\n\n Args:\n path (str or Path): Path to JSON file.\n metadata (list of BaseAudioMeta): List of audio meta to save.\n \"\"\"\n Path(path).parent.mkdir(exist_ok=True, parents=True)\n open_fn = gzip.open if str(path).lower().endswith('.gz') else open\n with open_fn(path, 'wb') as fp: # type: ignore\n for m in meta:\n json_str = json.dumps(m.to_dict()) + '\\n'\n json_bytes = json_str.encode('utf-8')\n fp.write(json_bytes)\n\n\nclass AudioDataset:\n \"\"\"Base audio dataset.\n\n The dataset takes a list of AudioMeta and create a dataset composed of segments of audio\n and potentially additional information, by creating random segments from the list of audio\n files referenced in the metadata and applying minimal data pre-processing such as resampling,\n mixing of channels, padding, etc.\n\n If no segment_duration value is provided, the AudioDataset will return the full wav for each\n audio file. Otherwise, it will randomly sample audio files and create a segment of the specified\n duration, applying padding if required.\n\n By default, only the torch Tensor corresponding to the waveform is returned. Setting return_info=True\n allows to return a tuple containing the torch Tensor and additional metadata on the segment and the\n original audio meta.\n\n Note that you can call `start_epoch(epoch)` in order to get\n a deterministic \"randomization\" for `shuffle=True`.\n For a given epoch and dataset index, this will always return the same extract.\n You can get back some diversity by setting the `shuffle_seed` param.\n\n Args:\n meta (list of AudioMeta): List of audio files metadata.\n segment_duration (float, optional): Optional segment duration of audio to load.\n If not specified, the dataset will load the full audio segment from the file.\n shuffle (bool): Set to `True` to have the data reshuffled at every epoch.\n sample_rate (int): Target sample rate of the loaded audio samples.\n channels (int): Target number of channels of the loaded audio samples.\n sample_on_duration (bool): Set to `True` to sample segments with probability\n dependent on audio file duration. This is only used if `segment_duration` is provided.\n sample_on_weight (bool): Set to `True` to sample segments using the `weight` entry of\n `AudioMeta`. If `sample_on_duration` is also True, the actual weight will be the product\n of the file duration and file weight. This is only used if `segment_duration` is provided.\n min_segment_ratio (float): Minimum segment ratio to use when the audio file\n is shorter than the desired segment.\n max_read_retry (int): Maximum number of retries to sample an audio segment from the dataset.\n return_info (bool): Whether to return the wav only or return wav along with segment info and metadata.\n min_audio_duration (float, optional): Minimum audio file duration, in seconds, if provided\n audio shorter than this will be filtered out.\n max_audio_duration (float, optional): Maximal audio file duration in seconds, if provided\n audio longer than this will be filtered out.\n shuffle_seed (int): can be used to further randomize\n load_wav (bool): if False, skip loading the wav but returns a tensor of 0\n with the expected segment_duration (which must be provided if load_wav is False).\n permutation_on_files (bool): only if `sample_on_weight` and `sample_on_duration`\n are False. Will ensure a permutation on files when going through the dataset.\n In that case the epoch number must be provided in order for the model\n to continue the permutation across epochs. In that case, it is assumed\n that `num_samples = total_batch_size * num_updates_per_epoch`, with\n `total_batch_size` the overall batch size accounting for all gpus.\n \"\"\"\n def __init__(self,\n meta: tp.List[AudioMeta],\n segment_duration: tp.Optional[float] = None,\n shuffle: bool = True,\n num_samples: int = 10_000,\n sample_rate: int = 48_000,\n channels: int = 2,\n pad: bool = True,\n sample_on_duration: bool = True,\n sample_on_weight: bool = True,\n min_segment_ratio: float = 0.5,\n max_read_retry: int = 10,\n return_info: bool = False,\n min_audio_duration: tp.Optional[float] = None,\n max_audio_duration: tp.Optional[float] = None,\n shuffle_seed: int = 0,\n load_wav: bool = True,\n permutation_on_files: bool = False,\n ):\n assert len(meta) > 0, \"No audio meta provided to AudioDataset. Please check loading of audio meta.\"\n assert segment_duration is None or segment_duration > 0\n assert segment_duration is None or min_segment_ratio >= 0\n self.segment_duration = segment_duration\n self.min_segment_ratio = min_segment_ratio\n self.max_audio_duration = max_audio_duration\n self.min_audio_duration = min_audio_duration\n if self.min_audio_duration is not None and self.max_audio_duration is not None:\n assert self.min_audio_duration <= self.max_audio_duration\n self.meta: tp.List[AudioMeta] = self._filter_duration(meta)\n assert len(self.meta) # Fail fast if all data has been filtered.\n self.total_duration = sum(d.duration for d in self.meta)\n\n if segment_duration is None:\n num_samples = len(self.meta)\n self.num_samples = num_samples\n self.shuffle = shuffle\n self.sample_rate = sample_rate\n self.channels = channels\n self.pad = pad\n self.sample_on_weight = sample_on_weight\n self.sample_on_duration = sample_on_duration\n self.sampling_probabilities = self._get_sampling_probabilities()\n self.max_read_retry = max_read_retry\n self.return_info = return_info\n self.shuffle_seed = shuffle_seed\n self.current_epoch: tp.Optional[int] = None\n self.load_wav = load_wav\n if not load_wav:\n assert segment_duration is not None\n self.permutation_on_files = permutation_on_files\n if permutation_on_files:\n assert not self.sample_on_duration\n assert not self.sample_on_weight\n assert self.shuffle\n\n def start_epoch(self, epoch: int):\n self.current_epoch = epoch\n\n def __len__(self):\n return self.num_samples\n\n def _get_sampling_probabilities(self, normalized: bool = True):\n \"\"\"Return the sampling probabilities for each file inside `self.meta`.\"\"\"\n scores: tp.List[float] = []\n for file_meta in self.meta:\n score = 1.\n if self.sample_on_weight and file_meta.weight is not None:\n score *= file_meta.weight\n if self.sample_on_duration:\n score *= file_meta.duration\n scores.append(score)\n probabilities = torch.tensor(scores)\n if normalized:\n probabilities /= probabilities.sum()\n return probabilities\n\n @staticmethod\n @lru_cache(16)\n def _get_file_permutation(num_files: int, permutation_index: int, base_seed: int):\n # Used to keep the most recent files permutation in memory implicitely.\n # will work unless someone is using a lot of Datasets in parallel.\n rng = torch.Generator()\n rng.manual_seed(base_seed + permutation_index)\n return torch.randperm(num_files, generator=rng)\n\n def sample_file(self, index: int, rng: torch.Generator) -> AudioMeta:\n \"\"\"Sample a given file from `self.meta`. Can be overridden in subclasses.\n This is only called if `segment_duration` is not None.\n\n You must use the provided random number generator `rng` for reproducibility.\n You can further make use of the index accessed.\n \"\"\"\n if self.permutation_on_files:\n assert self.current_epoch is not None\n total_index = self.current_epoch * len(self) + index\n permutation_index = total_index // len(self.meta)\n relative_index = total_index % len(self.meta)\n permutation = AudioDataset._get_file_permutation(\n len(self.meta), permutation_index, self.shuffle_seed)\n file_index = permutation[relative_index]\n return self.meta[file_index]\n\n if not self.sample_on_weight and not self.sample_on_duration:\n file_index = int(torch.randint(len(self.sampling_probabilities), (1,), generator=rng).item())\n else:\n file_index = int(torch.multinomial(self.sampling_probabilities, 1, generator=rng).item())\n\n return self.meta[file_index]\n\n def _audio_read(self, path: str, seek_time: float = 0, duration: float = -1):\n # Override this method in subclass if needed.\n if self.load_wav:\n return audio_read(path, seek_time, duration, pad=False)\n else:\n assert self.segment_duration is not None\n n_frames = int(self.sample_rate * self.segment_duration)\n return torch.zeros(self.channels, n_frames), self.sample_rate\n\n def __getitem__(self, index: int) -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, SegmentInfo]]:\n if self.segment_duration is None:\n file_meta = self.meta[index]\n out, sr = audio_read(file_meta.path)\n out = convert_audio(out, sr, self.sample_rate, self.channels)\n n_frames = out.shape[-1]\n segment_info = SegmentInfo(file_meta, seek_time=0., n_frames=n_frames, total_frames=n_frames,\n sample_rate=self.sample_rate, channels=out.shape[0])\n else:\n rng = torch.Generator()\n if self.shuffle:\n # We use index, plus extra randomness, either totally random if we don't know the epoch.\n # otherwise we make use of the epoch number and optional shuffle_seed.\n if self.current_epoch is None:\n rng.manual_seed(index + self.num_samples * random.randint(0, 2**24))\n else:\n rng.manual_seed(index + self.num_samples * (self.current_epoch + self.shuffle_seed))\n else:\n # We only use index\n rng.manual_seed(index)\n\n for retry in range(self.max_read_retry):\n file_meta = self.sample_file(index, rng)\n # We add some variance in the file position even if audio file is smaller than segment\n # without ending up with empty segments\n max_seek = max(0, file_meta.duration - self.segment_duration * self.min_segment_ratio)\n seek_time = torch.rand(1, generator=rng).item() * max_seek\n try:\n out, sr = audio_read(file_meta.path, seek_time, self.segment_duration, pad=False)\n out = convert_audio(out, sr, self.sample_rate, self.channels)\n n_frames = out.shape[-1]\n target_frames = int(self.segment_duration * self.sample_rate)\n if self.pad:\n out = F.pad(out, (0, target_frames - n_frames))\n segment_info = SegmentInfo(file_meta, seek_time, n_frames=n_frames, total_frames=target_frames,\n sample_rate=self.sample_rate, channels=out.shape[0])\n except Exception as exc:\n logger.warning(\"Error opening file %s: %r\", file_meta.path, exc)\n if retry == self.max_read_retry - 1:\n raise\n else:\n break\n\n if self.return_info:\n # Returns the wav and additional information on the wave segment\n return out, segment_info\n else:\n return out\n\n def collater(self, samples):\n \"\"\"The collater function has to be provided to the dataloader\n if AudioDataset has return_info=True in order to properly collate\n the samples of a batch.\n \"\"\"\n if self.segment_duration is None and len(samples) > 1:\n assert self.pad, \"Must allow padding when batching examples of different durations.\"\n\n # In this case the audio reaching the collater is of variable length as segment_duration=None.\n to_pad = self.segment_duration is None and self.pad\n if to_pad:\n max_len = max([wav.shape[-1] for wav, _ in samples])\n\n def _pad_wav(wav):\n return F.pad(wav, (0, max_len - wav.shape[-1]))\n\n if self.return_info:\n if len(samples) > 0:\n assert len(samples[0]) == 2\n assert isinstance(samples[0][0], torch.Tensor)\n assert isinstance(samples[0][1], SegmentInfo)\n\n wavs = [wav for wav, _ in samples]\n segment_infos = [copy.deepcopy(info) for _, info in samples]\n\n if to_pad:\n # Each wav could be of a different duration as they are not segmented.\n for i in range(len(samples)):\n # Determines the total length of the signal with padding, so we update here as we pad.\n segment_infos[i].total_frames = max_len\n wavs[i] = _pad_wav(wavs[i])\n\n wav = torch.stack(wavs)\n return wav, segment_infos\n else:\n assert isinstance(samples[0], torch.Tensor)\n if to_pad:\n samples = [_pad_wav(s) for s in samples]\n return torch.stack(samples)\n\n def _filter_duration(self, meta: tp.List[AudioMeta]) -> tp.List[AudioMeta]:\n \"\"\"Filters out audio files with audio durations that will not allow to sample examples from them.\"\"\"\n orig_len = len(meta)\n\n # Filter data that is too short.\n if self.min_audio_duration is not None:\n meta = [m for m in meta if m.duration >= self.min_audio_duration]\n\n # Filter data that is too long.\n if self.max_audio_duration is not None:\n meta = [m for m in meta if m.duration <= self.max_audio_duration]\n\n filtered_len = len(meta)\n removed_percentage = 100*(1-float(filtered_len)/orig_len)\n msg = 'Removed %.2f percent of the data because it was too short or too long.' % removed_percentage\n if removed_percentage < 10:\n logging.debug(msg)\n else:\n logging.warning(msg)\n return meta\n\n @classmethod\n def from_meta(cls, root: tp.Union[str, Path], **kwargs):\n \"\"\"Instantiate AudioDataset from a path to a directory containing a manifest as a jsonl file.\n\n Args:\n root (str or Path): Path to root folder containing audio files.\n kwargs: Additional keyword arguments for the AudioDataset.\n \"\"\"\n root = Path(root)\n if root.is_dir():\n if (root / 'data.jsonl').exists():\n root = root / 'data.jsonl'\n elif (root / 'data.jsonl.gz').exists():\n root = root / 'data.jsonl.gz'\n else:\n raise ValueError(\"Don't know where to read metadata from in the dir. \"\n \"Expecting either a data.jsonl or data.jsonl.gz file but none found.\")\n meta = load_audio_meta(root)\n return cls(meta, **kwargs)\n\n @classmethod\n def from_path(cls, root: tp.Union[str, Path], minimal_meta: bool = True,\n exts: tp.List[str] = DEFAULT_EXTS, **kwargs):\n \"\"\"Instantiate AudioDataset from a path containing (possibly nested) audio files.\n\n Args:\n root (str or Path): Path to root folder containing audio files.\n minimal_meta (bool): Whether to only load minimal metadata or not.\n exts (list of str): Extensions for audio files.\n kwargs: Additional keyword arguments for the AudioDataset.\n \"\"\"\n root = Path(root)\n if root.is_file():\n meta = load_audio_meta(root, resolve=True)\n else:\n meta = find_audio_files(root, exts, minimal=minimal_meta, resolve=True)\n return cls(meta, **kwargs)\n\n\ndef main():\n logging.basicConfig(stream=sys.stderr, level=logging.INFO)\n parser = argparse.ArgumentParser(\n prog='audio_dataset',\n description='Generate .jsonl files by scanning a folder.')\n parser.add_argument('root', help='Root folder with all the audio files')\n parser.add_argument('output_meta_file',\n help='Output file to store the metadata, ')\n parser.add_argument('--complete',\n action='store_false', dest='minimal', default=True,\n help='Retrieve all metadata, even the one that are expansive '\n 'to compute (e.g. normalization).')\n parser.add_argument('--resolve',\n action='store_true', default=False,\n help='Resolve the paths to be absolute and with no symlinks.')\n parser.add_argument('--workers',\n default=10, type=int,\n help='Number of workers.')\n args = parser.parse_args()\n meta = find_audio_files(args.root, DEFAULT_EXTS, progress=True,\n resolve=args.resolve, minimal=args.minimal, workers=args.workers)\n save_audio_meta(args.output_meta_file, meta)\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "audiocraft/data/audio_utils.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Various utilities for audio convertion (pcm format, sample rate and channels),\nand volume normalization.\"\"\"\nimport io\nimport logging\nimport re\nimport sys\nimport typing as tp\n\nimport julius\nimport torch\nimport torchaudio\n\nlogger = logging.getLogger(__name__)\n\n\ndef convert_audio_channels(wav: torch.Tensor, channels: int = 2) -> torch.Tensor:\n \"\"\"Convert audio to the given number of channels.\n\n Args:\n wav (torch.Tensor): Audio wave of shape [B, C, T].\n channels (int): Expected number of channels as output.\n Returns:\n torch.Tensor: Downmixed or unchanged audio wave [B, C, T].\n \"\"\"\n *shape, src_channels, length = wav.shape\n if src_channels == channels:\n pass\n elif channels == 1:\n # Case 1:\n # The caller asked 1-channel audio, and the stream has multiple\n # channels, downmix all channels.\n wav = wav.mean(dim=-2, keepdim=True)\n elif src_channels == 1:\n # Case 2:\n # The caller asked for multiple channels, but the input file has\n # a single channel, replicate the audio over all channels.\n wav = wav.expand(*shape, channels, length)\n elif src_channels >= channels:\n # Case 3:\n # The caller asked for multiple channels, and the input file has\n # more channels than requested. In that case return the first channels.\n wav = wav[..., :channels, :]\n else:\n # Case 4: What is a reasonable choice here?\n raise ValueError('The audio file has less channels than requested but is not mono.')\n return wav\n\n\ndef convert_audio(wav: torch.Tensor, from_rate: float,\n to_rate: float, to_channels: int) -> torch.Tensor:\n \"\"\"Convert audio to new sample rate and number of audio channels.\"\"\"\n wav = julius.resample_frac(wav, int(from_rate), int(to_rate))\n wav = convert_audio_channels(wav, to_channels)\n return wav\n\n\ndef normalize_loudness(wav: torch.Tensor, sample_rate: int, loudness_headroom_db: float = 14,\n loudness_compressor: bool = False, energy_floor: float = 2e-3):\n \"\"\"Normalize an input signal to a user loudness in dB LKFS.\n Audio loudness is defined according to the ITU-R BS.1770-4 recommendation.\n\n Args:\n wav (torch.Tensor): Input multichannel audio data.\n sample_rate (int): Sample rate.\n loudness_headroom_db (float): Target loudness of the output in dB LUFS.\n loudness_compressor (bool): Uses tanh for soft clipping.\n energy_floor (float): anything below that RMS level will not be rescaled.\n Returns:\n torch.Tensor: Loudness normalized output data.\n \"\"\"\n energy = wav.pow(2).mean().sqrt().item()\n if energy < energy_floor:\n return wav\n transform = torchaudio.transforms.Loudness(sample_rate)\n input_loudness_db = transform(wav).item()\n # calculate the gain needed to scale to the desired loudness level\n delta_loudness = -loudness_headroom_db - input_loudness_db\n gain = 10.0 ** (delta_loudness / 20.0)\n output = gain * wav\n if loudness_compressor:\n output = torch.tanh(output)\n assert output.isfinite().all(), (input_loudness_db, wav.pow(2).mean().sqrt())\n return output\n\n\ndef _clip_wav(wav: torch.Tensor, log_clipping: bool = False, stem_name: tp.Optional[str] = None) -> None:\n \"\"\"\n Utility function to clip the audio with logging if specified.\n \"\"\"\n max_scale = wav.abs().max()\n if log_clipping and max_scale > 1:\n clamp_prob = (wav.abs() > 1).float().mean().item()\n print(f\"CLIPPING {stem_name or ''} happening with proba (a bit of clipping is okay):\",\n clamp_prob, \"maximum scale: \", max_scale.item(), file=sys.stderr)\n wav.clamp_(-1, 1)\n\n\ndef normalize_audio(wav: torch.Tensor, normalize: bool = True,\n strategy: str = 'peak', peak_clip_headroom_db: float = 1,\n rms_headroom_db: float = 18, loudness_headroom_db: float = 14,\n loudness_compressor: bool = False, log_clipping: bool = False,\n sample_rate: tp.Optional[int] = None,\n stem_name: tp.Optional[str] = None) -> torch.Tensor:\n \"\"\"Normalize the audio according to the prescribed strategy (see after).\n\n Args:\n wav (torch.Tensor): Audio data.\n normalize (bool): if `True` (default), normalizes according to the prescribed\n strategy (see after). If `False`, the strategy is only used in case clipping\n would happen.\n strategy (str): Can be either 'clip', 'peak', or 'rms'. Default is 'peak',\n i.e. audio is normalized by its largest value. RMS normalizes by root-mean-square\n with extra headroom to avoid clipping. 'clip' just clips.\n peak_clip_headroom_db (float): Headroom in dB when doing 'peak' or 'clip' strategy.\n rms_headroom_db (float): Headroom in dB when doing 'rms' strategy. This must be much larger\n than the `peak_clip` one to avoid further clipping.\n loudness_headroom_db (float): Target loudness for loudness normalization.\n loudness_compressor (bool): If True, uses tanh based soft clipping.\n log_clipping (bool): If True, basic logging on stderr when clipping still\n occurs despite strategy (only for 'rms').\n sample_rate (int): Sample rate for the audio data (required for loudness).\n stem_name (str, optional): Stem name for clipping logging.\n Returns:\n torch.Tensor: Normalized audio.\n \"\"\"\n scale_peak = 10 ** (-peak_clip_headroom_db / 20)\n scale_rms = 10 ** (-rms_headroom_db / 20)\n if strategy == 'peak':\n rescaling = (scale_peak / wav.abs().max())\n if normalize or rescaling < 1:\n wav = wav * rescaling\n elif strategy == 'clip':\n wav = wav.clamp(-scale_peak, scale_peak)\n elif strategy == 'rms':\n mono = wav.mean(dim=0)\n rescaling = scale_rms / mono.pow(2).mean().sqrt()\n if normalize or rescaling < 1:\n wav = wav * rescaling\n _clip_wav(wav, log_clipping=log_clipping, stem_name=stem_name)\n elif strategy == 'loudness':\n assert sample_rate is not None, \"Loudness normalization requires sample rate.\"\n wav = normalize_loudness(wav, sample_rate, loudness_headroom_db, loudness_compressor)\n _clip_wav(wav, log_clipping=log_clipping, stem_name=stem_name)\n else:\n assert wav.abs().max() < 1\n assert strategy == '' or strategy == 'none', f\"Unexpected strategy: '{strategy}'\"\n return wav\n\n\ndef f32_pcm(wav: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Convert audio to float 32 bits PCM format.\n Args:\n wav (torch.tensor): Input wav tensor\n Returns:\n same wav in float32 PCM format\n \"\"\"\n if wav.dtype.is_floating_point:\n return wav\n elif wav.dtype == torch.int16:\n return wav.float() / 2**15\n elif wav.dtype == torch.int32:\n return wav.float() / 2**31\n raise ValueError(f\"Unsupported wav dtype: {wav.dtype}\")\n\n\ndef i16_pcm(wav: torch.Tensor) -> torch.Tensor:\n \"\"\"Convert audio to int 16 bits PCM format.\n\n ..Warning:: There exist many formula for doing this conversion. None are perfect\n due to the asymmetry of the int16 range. One either have possible clipping, DC offset,\n or inconsistencies with f32_pcm. If the given wav doesn't have enough headroom,\n it is possible that `i16_pcm(f32_pcm)) != Identity`.\n Args:\n wav (torch.tensor): Input wav tensor\n Returns:\n same wav in float16 PCM format\n \"\"\"\n if wav.dtype.is_floating_point:\n assert wav.abs().max() <= 1\n candidate = (wav * 2 ** 15).round()\n if candidate.max() >= 2 ** 15: # clipping would occur\n candidate = (wav * (2 ** 15 - 1)).round()\n return candidate.short()\n else:\n assert wav.dtype == torch.int16\n return wav\n\n\ndef compress(wav: torch.Tensor, sr: int,\n target_format: tp.Literal[\"mp3\", \"ogg\", \"flac\"] = \"mp3\",\n bitrate: str = \"128k\") -> tp.Tuple[torch.Tensor, int]:\n \"\"\"Convert audio wave form to a specified lossy format: mp3, ogg, flac\n\n Args:\n wav (torch.Tensor): Input wav tensor.\n sr (int): Sampling rate.\n target_format (str): Compression format (e.g., 'mp3').\n bitrate (str): Bitrate for compression.\n\n Returns:\n Tuple of compressed WAV tensor and sampling rate.\n \"\"\"\n\n # Extract the bit rate from string (e.g., '128k')\n match = re.search(r\"\\d+(\\.\\d+)?\", str(bitrate))\n parsed_bitrate = float(match.group()) if match else None\n assert parsed_bitrate, f\"Invalid bitrate specified (got {parsed_bitrate})\"\n try:\n # Create a virtual file instead of saving to disk\n buffer = io.BytesIO()\n\n torchaudio.save(\n buffer, wav, sr, format=target_format, bits_per_sample=parsed_bitrate,\n )\n # Move to the beginning of the file\n buffer.seek(0)\n compressed_wav, sr = torchaudio.load(buffer)\n return compressed_wav, sr\n\n except RuntimeError:\n logger.warning(\n f\"compression failed skipping compression: {format} {parsed_bitrate}\"\n )\n return wav, sr\n\n\ndef get_mp3(wav_tensor: torch.Tensor, sr: int, bitrate: str = \"128k\") -> torch.Tensor:\n \"\"\"Convert a batch of audio files to MP3 format, maintaining the original shape.\n\n This function takes a batch of audio files represented as a PyTorch tensor, converts\n them to MP3 format using the specified bitrate, and returns the batch in the same\n shape as the input.\n\n Args:\n wav_tensor (torch.Tensor): Batch of audio files represented as a tensor.\n Shape should be (batch_size, channels, length).\n sr (int): Sampling rate of the audio.\n bitrate (str): Bitrate for MP3 conversion, default is '128k'.\n\n Returns:\n torch.Tensor: Batch of audio files converted to MP3 format, with the same\n shape as the input tensor.\n \"\"\"\n device = wav_tensor.device\n batch_size, channels, original_length = wav_tensor.shape\n\n # Flatten tensor for conversion and move to CPU\n wav_tensor_flat = wav_tensor.view(1, -1).cpu()\n\n # Convert to MP3 format with specified bitrate\n wav_tensor_flat, _ = compress(wav_tensor_flat, sr, bitrate=bitrate)\n\n # Reshape back to original batch format and trim or pad if necessary\n wav_tensor = wav_tensor_flat.view(batch_size, channels, -1)\n compressed_length = wav_tensor.shape[-1]\n if compressed_length > original_length:\n wav_tensor = wav_tensor[:, :, :original_length] # Trim excess frames\n elif compressed_length < original_length:\n padding = torch.zeros(\n batch_size, channels, original_length - compressed_length, device=device\n )\n wav_tensor = torch.cat((wav_tensor, padding), dim=-1) # Pad with zeros\n\n # Move tensor back to the original device\n return wav_tensor.to(device)\n\n\ndef get_aac(\n wav_tensor: torch.Tensor,\n sr: int,\n bitrate: str = \"128k\",\n lowpass_freq: tp.Optional[int] = None,\n) -> torch.Tensor:\n \"\"\"Converts a batch of audio tensors to AAC format and then back to tensors.\n\n This function first saves the input tensor batch as WAV files, then uses FFmpeg to convert\n these WAV files to AAC format. Finally, it loads the AAC files back into tensors.\n\n Args:\n wav_tensor (torch.Tensor): A batch of audio files represented as a tensor.\n Shape should be (batch_size, channels, length).\n sr (int): Sampling rate of the audio.\n bitrate (str): Bitrate for AAC conversion, default is '128k'.\n lowpass_freq (Optional[int]): Frequency for a low-pass filter. If None, no filter is applied.\n\n Returns:\n torch.Tensor: Batch of audio files converted to AAC and back, with the same\n shape as the input tensor.\n \"\"\"\n import tempfile\n import subprocess\n\n device = wav_tensor.device\n batch_size, channels, original_length = wav_tensor.shape\n\n # Parse the bitrate value from the string\n match = re.search(r\"\\d+(\\.\\d+)?\", bitrate)\n parsed_bitrate = (\n match.group() if match else \"128\"\n ) # Default to 128 if parsing fails\n\n # Flatten tensor for conversion and move to CPU\n wav_tensor_flat = wav_tensor.view(1, -1).cpu()\n\n with tempfile.NamedTemporaryFile(\n suffix=\".wav\"\n ) as f_in, tempfile.NamedTemporaryFile(suffix=\".aac\") as f_out:\n input_path, output_path = f_in.name, f_out.name\n\n # Save the tensor as a WAV file\n torchaudio.save(input_path, wav_tensor_flat, sr, backend=\"ffmpeg\")\n\n # Prepare FFmpeg command for AAC conversion\n command = [\n \"ffmpeg\",\n \"-y\",\n \"-i\",\n input_path,\n \"-ar\",\n str(sr),\n \"-b:a\",\n f\"{parsed_bitrate}k\",\n \"-c:a\",\n \"aac\",\n ]\n if lowpass_freq is not None:\n command += [\"-cutoff\", str(lowpass_freq)]\n command.append(output_path)\n\n try:\n # Run FFmpeg and suppress output\n subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)\n\n # Load the AAC audio back into a tensor\n aac_tensor, _ = torchaudio.load(output_path, backend=\"ffmpeg\")\n except Exception as exc:\n raise RuntimeError(\n \"Failed to run command \" \".join(command)} \"\n \"(Often this means ffmpeg is not installed or the encoder is not supported, \"\n \"make sure you installed an older version ffmpeg<5)\"\n ) from exc\n\n original_length_flat = batch_size * channels * original_length\n compressed_length_flat = aac_tensor.shape[-1]\n\n # Trim excess frames\n if compressed_length_flat > original_length_flat:\n aac_tensor = aac_tensor[:, :original_length_flat]\n\n # Pad the shortedn frames\n elif compressed_length_flat < original_length_flat:\n padding = torch.zeros(\n 1, original_length_flat - compressed_length_flat, device=device\n )\n aac_tensor = torch.cat((aac_tensor, padding), dim=-1)\n\n # Reshape and adjust length to match original tensor\n wav_tensor = aac_tensor.view(batch_size, channels, -1)\n compressed_length = wav_tensor.shape[-1]\n\n assert compressed_length == original_length, (\n \"AAC-compressed audio does not have the same frames as original one. \"\n \"One reason can be ffmpeg is not installed and used as proper backed \"\n \"for torchaudio, or the AAC encoder is not correct. Run \"\n \"`torchaudio.utils.ffmpeg_utils.get_audio_encoders()` and make sure we see entry for\"\n \"AAC in the output.\"\n )\n return wav_tensor.to(device)\n" }, { "path": "audiocraft/data/info_audio_dataset.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Base classes for the datasets that also provide non-audio metadata,\ne.g. description, text transcription etc.\n\"\"\"\nfrom dataclasses import dataclass\nimport logging\nimport math\nimport re\nimport typing as tp\n\nimport torch\n\nfrom .audio_dataset import AudioDataset, AudioMeta\nfrom ..environment import AudioCraftEnvironment\nfrom ..modules.conditioners import SegmentWithAttributes, ConditioningAttributes\n\n\nlogger = logging.getLogger(__name__)\n\n\ndef _clusterify_meta(meta: AudioMeta) -> AudioMeta:\n \"\"\"Monkey-patch meta to match cluster specificities.\"\"\"\n meta.path = AudioCraftEnvironment.apply_dataset_mappers(meta.path)\n if meta.info_path is not None:\n meta.info_path.zip_path = AudioCraftEnvironment.apply_dataset_mappers(meta.info_path.zip_path)\n return meta\n\n\ndef clusterify_all_meta(meta: tp.List[AudioMeta]) -> tp.List[AudioMeta]:\n \"\"\"Monkey-patch all meta to match cluster specificities.\"\"\"\n return [_clusterify_meta(m) for m in meta]\n\n\n@dataclass\nclass AudioInfo(SegmentWithAttributes):\n \"\"\"Dummy SegmentInfo with empty attributes.\n\n The InfoAudioDataset is expected to return metadata that inherits\n from SegmentWithAttributes class and can return conditioning attributes.\n\n This basically guarantees all datasets will be compatible with current\n solver that contain conditioners requiring this.\n \"\"\"\n audio_tokens: tp.Optional[torch.Tensor] = None # populated when using cached batch for training a LM.\n\n def to_condition_attributes(self) -> ConditioningAttributes:\n return ConditioningAttributes()\n\n\nclass InfoAudioDataset(AudioDataset):\n \"\"\"AudioDataset that always returns metadata as SegmentWithAttributes along with the audio waveform.\n\n See `audiocraft.data.audio_dataset.AudioDataset` for initialization arguments.\n \"\"\"\n def __init__(self, meta: tp.List[AudioMeta], **kwargs):\n super().__init__(clusterify_all_meta(meta), **kwargs)\n\n def __getitem__(self, index: int) -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, SegmentWithAttributes]]:\n if not self.return_info:\n wav = super().__getitem__(index)\n assert isinstance(wav, torch.Tensor)\n return wav\n wav, meta = super().__getitem__(index)\n return wav, AudioInfo(**meta.to_dict())\n\n\ndef get_keyword_or_keyword_list(value: tp.Optional[str]) -> tp.Union[tp.Optional[str], tp.Optional[tp.List[str]]]:\n \"\"\"Preprocess a single keyword or possible a list of keywords.\"\"\"\n if isinstance(value, list):\n return get_keyword_list(value)\n else:\n return get_keyword(value)\n\n\ndef get_string(value: tp.Optional[str]) -> tp.Optional[str]:\n \"\"\"Preprocess a single keyword.\"\"\"\n if value is None or (not isinstance(value, str)) or len(value) == 0 or value == 'None':\n return None\n else:\n return value.strip()\n\n\ndef get_keyword(value: tp.Optional[str]) -> tp.Optional[str]:\n \"\"\"Preprocess a single keyword.\"\"\"\n if value is None or (not isinstance(value, str)) or len(value) == 0 or value == 'None':\n return None\n else:\n return value.strip().lower()\n\n\ndef get_keyword_list(values: tp.Union[str, tp.List[str]]) -> tp.Optional[tp.List[str]]:\n \"\"\"Preprocess a list of keywords.\"\"\"\n if isinstance(values, str):\n values = [v.strip() for v in re.split(r'[,\\s]', values)]\n elif isinstance(values, float) and math.isnan(values):\n values = []\n if not isinstance(values, list):\n logger.debug(f\"Unexpected keyword list {values}\")\n values = [str(values)]\n\n kws = [get_keyword(v) for v in values]\n kw_list = [k for k in kws if k is not None]\n if len(kw_list) == 0:\n return None\n else:\n return kw_list\n" }, { "path": "audiocraft/data/jasco_dataset.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\nimport bisect\nimport pickle\nimport math\nimport os\nimport torch\nimport typing as tp\nfrom pathlib import Path\nfrom dataclasses import dataclass, fields\nfrom ..utils.utils import construct_frame_chords\nfrom .music_dataset import MusicDataset, MusicInfo\nfrom .audio_dataset import load_audio_meta\nfrom ..modules.conditioners import (ConditioningAttributes, SymbolicCondition)\nimport librosa\nimport numpy as np\n\n\n@dataclass\nclass JascoInfo(MusicInfo):\n \"\"\"\n A data class extending MusicInfo for JASCO. The following attributes are added:\n Attributes:\n frame_chords (Optional[list]): A list of chords associated with frames in the music piece.\n \"\"\"\n chords: tp.Optional[SymbolicCondition] = None\n melody: tp.Optional[SymbolicCondition] = None\n\n def to_condition_attributes(self) -> ConditioningAttributes:\n out = ConditioningAttributes()\n for _field in fields(self):\n key, value = _field.name, getattr(self, _field.name)\n if key == 'self_wav':\n out.wav[key] = value\n elif key in {'chords', 'melody'}:\n out.symbolic[key] = value\n elif key == 'joint_embed':\n for embed_attribute, embed_cond in value.items():\n out.joint_embed[embed_attribute] = embed_cond\n else:\n if isinstance(value, list):\n value = ' '.join(value)\n out.text[key] = value\n return out\n\n\nclass MelodyData:\n\n SALIENCE_MODEL_EXPECTED_SAMPLE_RATE = 22050\n SALIENCE_MODEL_EXPECTED_HOP_SIZE = 256\n\n def __init__(self,\n latent_fr: int,\n segment_duration: float,\n melody_fr: int = 86,\n melody_salience_dim: int = 53,\n chroma_root: tp.Optional[str] = None,\n override_cache: bool = False,\n do_argmax: bool = True):\n \"\"\"Module to load salience matrix for a given info.\n\n Args:\n latent_fr (int): latent frame rate to match (interpolates model frame rate accordingly).\n segment_duration (float): expected segment duration.\n melody_fr (int, optional): extracted salience frame rate. Defaults to 86.\n melody_salience_dim (int, optional): salience dim. Defaults to 53.\n chroma_root (str, optional): path to root containing salience cache. Defaults to None.\n override_cache (bool, optional): rewrite cache. Defaults to False.\n do_argmax (bool, optional): argmax the melody matrix. Defaults to True.\n \"\"\"\n\n self.segment_duration = segment_duration\n self.melody_fr = melody_fr\n self.latent_fr = latent_fr\n self.melody_salience_dim = melody_salience_dim\n self.do_argmax = do_argmax\n self.tgt_chunk_len = int(latent_fr * segment_duration)\n\n self.null_op = False\n if chroma_root is None:\n self.null_op = True\n elif not os.path.exists(f\"{chroma_root}/cache.pkl\") or override_cache:\n self.tracks = []\n for file in librosa.util.find_files(chroma_root, ext='txt'):\n with open(file, 'r') as f:\n lines = f.readlines()\n for line in lines:\n self.tracks.append(line.strip())\n\n # go over tracks and add the corresponding saliency file to self.saliency_files\n self.saliency_files = []\n for track in self.tracks:\n # saliency file name\n salience_file = f\"{chroma_root}/{track.split('/')[-1].split('.')[0]}_multif0_salience.npz\"\n assert os.path.exists(salience_file), f\"File {salience_file} does not exist\"\n self.saliency_files.append(salience_file)\n\n self.trk2idx = {trk.split('/')[-1].split('.')[0]: i for i, trk in enumerate(self.tracks)}\n torch.save({'tracks': self.tracks,\n 'saliency_files': self.saliency_files,\n 'trk2idx': self.trk2idx}, f\"{chroma_root}/cache.pkl\")\n else:\n tmp = torch.load(f\"{chroma_root}/cache.pkl\")\n self.tracks = tmp['tracks']\n self.saliency_files = tmp['saliency_files']\n self.trk2idx = tmp['trk2idx']\n self.model_frame_rate = int(self.SALIENCE_MODEL_EXPECTED_SAMPLE_RATE / self.SALIENCE_MODEL_EXPECTED_HOP_SIZE)\n\n def load_saliency_from_saliency_dict(self,\n saliency_dict: tp.Dict[str, tp.Any],\n offset: float) -> torch.Tensor:\n \"\"\"\n construct the salience matrix and perform linear interpolation w.r.t the temporal axis to match the expected\n frame rate.\n \"\"\"\n # get saliency map for the segment\n saliency_dict_ = {}\n l, r = int(offset * self.model_frame_rate), int((offset + self.segment_duration) * self.model_frame_rate)\n saliency_dict_['salience'] = saliency_dict['salience'][:, l: r].T\n saliency_dict_['times'] = saliency_dict['times'][l: r] - offset\n saliency_dict_['freqs'] = saliency_dict['freqs']\n\n saliency_dict_['salience'] = torch.Tensor(saliency_dict_['salience']).float().permute(1, 0) # C, T\n if saliency_dict_['salience'].shape[-1] <= int(self.model_frame_rate) / self.latent_fr: # empty chroma\n saliency_dict_['salience'] = torch.zeros((saliency_dict_['salience'].shape[0], self.tgt_chunk_len))\n else:\n salience = torch.nn.functional.interpolate(saliency_dict_['salience'].unsqueeze(0),\n scale_factor=self.latent_fr/int(self.model_frame_rate),\n mode='linear').squeeze(0)\n if salience.shape[-1] < self.tgt_chunk_len:\n salience = torch.nn.functional.pad(salience,\n (0, self.tgt_chunk_len - salience.shape[-1]),\n mode='constant',\n value=0)\n elif salience.shape[-1] > self.tgt_chunk_len:\n salience = salience[..., :self.tgt_chunk_len]\n saliency_dict_['salience'] = salience\n\n salience = saliency_dict_['salience']\n if self.do_argmax:\n binary_mask = torch.zeros_like(salience)\n binary_mask[torch.argmax(salience, dim=0), torch.arange(salience.shape[-1])] = 1\n binary_mask *= (salience != 0).float()\n salience = binary_mask\n return salience\n\n def get_null_salience(self) -> torch.Tensor:\n return torch.zeros((self.melody_salience_dim, self.tgt_chunk_len))\n\n def __call__(self, x: MusicInfo) -> torch.Tensor:\n \"\"\"Reads salience matrix from memory, shifted by seek time\n\n Args:\n x (MusicInfo): Music info of a single sample\n\n Returns:\n torch.Tensor: salience matrix matching the target info\n \"\"\"\n fname: str = x.meta.path.split(\"/\")[-1].split(\".\")[0] if x.meta.path is not None else \"\"\n if x.meta.path is None or x.meta.path == \"\" or fname not in self.trk2idx:\n salience = self.get_null_salience()\n else:\n assert fname in self.trk2idx, f\"Track {fname} not found in the cache\"\n idx = self.trk2idx[fname]\n saliency_dict = np.load(self.saliency_files[idx], allow_pickle=True)\n salience = self.load_saliency_from_saliency_dict(saliency_dict, x.seek_time)\n return salience\n\n\nclass JascoDataset(MusicDataset):\n \"\"\"JASCO dataset is a MusicDataset with jasco-related symbolic data (chords, melody).\n\n Args:\n chords_card (int): The cardinality of the chords, default is 194.\n compression_model_framerate (int): The framerate for the compression model, default is 50.\n\n See `audiocraft.data.info_audio_dataset.MusicDataset` for full initialization arguments.\n \"\"\"\n @classmethod\n def from_meta(cls, root: tp.Union[str, Path], **kwargs):\n \"\"\"Instantiate AudioDataset from a path to a directory containing a manifest as a jsonl file.\n\n Args:\n root (str or Path): Path to root folder containing audio files.\n kwargs: Additional keyword arguments for the AudioDataset.\n \"\"\"\n root = Path(root)\n # a directory is given\n if root.is_dir():\n if (root / 'data.jsonl').exists():\n meta_json = root / 'data.jsonl'\n elif (root / 'data.jsonl.gz').exists():\n meta_json = root / 'data.jsonl.gz'\n else:\n raise ValueError(\"Don't know where to read metadata from in the dir. \"\n \"Expecting either a data.jsonl or data.jsonl.gz file but none found.\")\n # jsonl file was specified\n else:\n assert root.exists() and root.suffix == '.jsonl', \\\n \"Either specified path not exist or it is not a jsonl format\"\n meta_json = root\n root = root.parent\n meta = load_audio_meta(meta_json)\n kwargs['root'] = root\n return cls(meta, **kwargs)\n\n def __init__(self, *args,\n chords_card: int = 194,\n compression_model_framerate: float = 50.,\n melody_kwargs: tp.Optional[tp.Dict[str, tp.Any]] = {},\n **kwargs):\n \"\"\"Dataset class for text-to-music generation with temporal controls as in\n (JASCO)[https://arxiv.org/pdf/2406.10970]\n\n Args:\n chords_card (int, optional): Number of chord ebeddings. Defaults to 194.\n compression_model_framerate (float, optional): Expected frame rate of the resulted latent. Defaults to 50.\n melody_kwargs (tp.Optional[tp.Dict[str, tp.Any]], optional): See MelodyData class. Defaults to {}.\n \"\"\"\n root = kwargs.pop('root')\n super().__init__(*args, **kwargs)\n\n chords_mapping_path = root / 'chord_to_index_mapping.pkl'\n chords_path = root / 'chords_per_track.pkl'\n self.mapping_dict = pickle.load(open(chords_mapping_path, \"rb\")) if \\\n os.path.exists(chords_mapping_path) else None\n\n self.chords_per_track = pickle.load(open(chords_path, \"rb\")) if \\\n os.path.exists(chords_path) else None\n\n self.compression_model_framerate = compression_model_framerate\n self.null_chord_idx = chords_card\n\n self.melody_module = MelodyData(**melody_kwargs) # type: ignore\n\n def _get_relevant_sublist(self, chords, timestamp):\n \"\"\"\n Returns the sublist of chords within the specified timestamp and segment length.\n\n Args:\n chords (list): A sorted list of tuples containing (time changed, chord).\n timestamp (float): The timestamp at which to start the sublist.\n\n Returns:\n list: A list of chords within the specified timestamp and segment length.\n \"\"\"\n end_time = timestamp + self.segment_duration\n\n # Use binary search to find the starting index of the relevant sublist\n start_index = bisect.bisect_left(chords, (timestamp,))\n\n if start_index != 0:\n prev_chord = chords[start_index - 1]\n else:\n prev_chord = (0.0, \"N\")\n\n relevant_chords = []\n\n for time_changed, chord in chords[start_index:]:\n if time_changed >= end_time:\n break\n relevant_chords.append((time_changed, chord))\n\n return relevant_chords, prev_chord\n\n def _get_chords(self, music_info: MusicInfo, effective_segment_dur: float) -> torch.Tensor:\n if self.chords_per_track is None:\n # use null chord when there's no chords in dataset\n seq_len = math.ceil(self.compression_model_framerate * effective_segment_dur)\n return torch.ones(seq_len, dtype=int) * self.null_chord_idx # type: ignore\n\n fr = self.compression_model_framerate\n\n idx = music_info.meta.path.split(\"/\")[-1].split(\".\")[0]\n chords = self.chords_per_track[idx]\n\n min_timestamp = music_info.seek_time\n\n chords = [(item[1], item[0]) for item in chords]\n chords, prev_chord = self._get_relevant_sublist(\n chords, min_timestamp\n )\n\n iter_min_timestamp = int(min_timestamp * fr) + 1\n\n frame_chords = construct_frame_chords(\n iter_min_timestamp, chords, self.mapping_dict, prev_chord[1], # type: ignore\n fr, self.segment_duration # type: ignore\n )\n\n return torch.tensor(frame_chords)\n\n def __getitem__(self, index):\n wav, music_info = super().__getitem__(index)\n assert not wav.isinfinite().any(), f\"inf detected in wav file: {music_info}\"\n wav = wav.float()\n\n # downcast music info to jasco info\n jasco_info = JascoInfo({k: v for k, v in music_info.__dict__.items()})\n\n # get chords\n effective_segment_dur = (wav.shape[-1] / self.sample_rate) if \\\n self.segment_duration is None else self.segment_duration\n frame_chords = self._get_chords(music_info, effective_segment_dur)\n jasco_info.chords = SymbolicCondition(frame_chords=frame_chords)\n\n # get melody\n jasco_info.melody = SymbolicCondition(melody=self.melody_module(music_info))\n return wav, jasco_info\n" }, { "path": "audiocraft/data/music_dataset.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Dataset of music tracks with rich metadata.\n\"\"\"\nfrom dataclasses import dataclass, field, fields, replace\nimport gzip\nimport json\nimport logging\nfrom pathlib import Path\nimport random\nimport typing as tp\n\nimport torch\n\nfrom .info_audio_dataset import (\n InfoAudioDataset,\n AudioInfo,\n get_keyword_list,\n get_keyword,\n get_string\n)\nfrom ..modules.conditioners import (\n ConditioningAttributes,\n JointEmbedCondition,\n WavCondition,\n)\nfrom ..utils.utils import warn_once\n\n\nlogger = logging.getLogger(__name__)\n\n\n@dataclass\nclass MusicInfo(AudioInfo):\n \"\"\"Segment info augmented with music metadata.\n \"\"\"\n # music-specific metadata\n title: tp.Optional[str] = None\n artist: tp.Optional[str] = None # anonymized artist id, used to ensure no overlap between splits\n key: tp.Optional[str] = None\n bpm: tp.Optional[float] = None\n genre: tp.Optional[str] = None\n moods: tp.Optional[list] = None\n keywords: tp.Optional[list] = None\n description: tp.Optional[str] = None\n name: tp.Optional[str] = None\n instrument: tp.Optional[str] = None\n # original wav accompanying the metadata\n self_wav: tp.Optional[WavCondition] = None\n # dict mapping attributes names to tuple of wav, text and metadata\n joint_embed: tp.Dict[str, JointEmbedCondition] = field(default_factory=dict)\n\n @property\n def has_music_meta(self) -> bool:\n return self.name is not None\n\n def to_condition_attributes(self) -> ConditioningAttributes:\n out = ConditioningAttributes()\n for _field in fields(self):\n key, value = _field.name, getattr(self, _field.name)\n if key == 'self_wav':\n out.wav[key] = value\n elif key == 'joint_embed':\n for embed_attribute, embed_cond in value.items():\n out.joint_embed[embed_attribute] = embed_cond\n else:\n if isinstance(value, list):\n value = ' '.join(value)\n out.text[key] = value\n return out\n\n @staticmethod\n def attribute_getter(attribute):\n if attribute == 'bpm':\n preprocess_func = get_bpm\n elif attribute == 'key':\n preprocess_func = get_musical_key\n elif attribute in ['moods', 'keywords']:\n preprocess_func = get_keyword_list\n elif attribute in ['genre', 'name', 'instrument']:\n preprocess_func = get_keyword\n elif attribute in ['title', 'artist', 'description']:\n preprocess_func = get_string\n else:\n preprocess_func = None\n return preprocess_func\n\n @classmethod\n def from_dict(cls, dictionary: dict, fields_required: bool = False):\n _dictionary: tp.Dict[str, tp.Any] = {}\n\n # allow a subset of attributes to not be loaded from the dictionary\n # these attributes may be populated later\n post_init_attributes = ['self_wav', 'joint_embed']\n optional_fields = ['keywords']\n\n for _field in fields(cls):\n if _field.name in post_init_attributes:\n continue\n elif _field.name not in dictionary:\n if fields_required and _field.name not in optional_fields:\n raise KeyError(f\"Unexpected missing key: {_field.name}\")\n else:\n preprocess_func: tp.Optional[tp.Callable] = cls.attribute_getter(_field.name)\n value = dictionary[_field.name]\n if preprocess_func:\n value = preprocess_func(value)\n _dictionary[_field.name] = value\n return cls(**_dictionary)\n\n\ndef augment_music_info_description(music_info: MusicInfo, merge_text_p: float = 0.,\n drop_desc_p: float = 0., drop_other_p: float = 0.) -> MusicInfo:\n \"\"\"Augment MusicInfo description with additional metadata fields and potential dropout.\n Additional textual attributes are added given probability 'merge_text_conditions_p' and\n the original textual description is dropped from the augmented description given probability drop_desc_p.\n\n Args:\n music_info (MusicInfo): The music metadata to augment.\n merge_text_p (float): Probability of merging additional metadata to the description.\n If provided value is 0, then no merging is performed.\n drop_desc_p (float): Probability of dropping the original description on text merge.\n if provided value is 0, then no drop out is performed.\n drop_other_p (float): Probability of dropping the other fields used for text augmentation.\n Returns:\n MusicInfo: The MusicInfo with augmented textual description.\n \"\"\"\n def is_valid_field(field_name: str, field_value: tp.Any) -> bool:\n valid_field_name = field_name in ['key', 'bpm', 'genre', 'moods', 'instrument', 'keywords']\n valid_field_value = field_value is not None and isinstance(field_value, (int, float, str, list))\n keep_field = random.uniform(0, 1) < drop_other_p\n return valid_field_name and valid_field_value and keep_field\n\n def process_value(v: tp.Any) -> str:\n if isinstance(v, (int, float, str)):\n return str(v)\n if isinstance(v, list):\n return \", \".join(v)\n else:\n raise ValueError(f\"Unknown type for text value! ({type(v), v})\")\n\n description = music_info.description\n\n metadata_text = \"\"\n if random.uniform(0, 1) < merge_text_p:\n meta_pairs = [f'{_field.name}: {process_value(getattr(music_info, _field.name))}'\n for _field in fields(music_info) if is_valid_field(_field.name, getattr(music_info, _field.name))]\n random.shuffle(meta_pairs)\n metadata_text = \". \".join(meta_pairs)\n description = description if not random.uniform(0, 1) < drop_desc_p else None\n logger.debug(f\"Applying text augmentation on MMI info. description: {description}, metadata: {metadata_text}\")\n\n if description is None:\n description = metadata_text if len(metadata_text) > 1 else None\n else:\n description = \". \".join([description.rstrip('.'), metadata_text])\n description = description.strip() if description else None\n\n music_info = replace(music_info)\n music_info.description = description\n return music_info\n\n\nclass Paraphraser:\n def __init__(self, paraphrase_source: tp.Union[str, Path], paraphrase_p: float = 0.):\n self.paraphrase_p = paraphrase_p\n open_fn = gzip.open if str(paraphrase_source).lower().endswith('.gz') else open\n with open_fn(paraphrase_source, 'rb') as f: # type: ignore\n self.paraphrase_source = json.loads(f.read())\n logger.info(f\"loaded paraphrasing source from: {paraphrase_source}\")\n\n def sample_paraphrase(self, audio_path: str, description: str):\n if random.random() >= self.paraphrase_p:\n return description\n info_path = Path(audio_path).with_suffix('.json')\n if info_path not in self.paraphrase_source:\n warn_once(logger, f\"{info_path} not in paraphrase source!\")\n return description\n new_desc = random.choice(self.paraphrase_source[info_path])\n logger.debug(f\"{description} -> {new_desc}\")\n return new_desc\n\n\nclass MusicDataset(InfoAudioDataset):\n \"\"\"Music dataset is an AudioDataset with music-related metadata.\n\n Args:\n info_fields_required (bool): Whether to enforce having required fields.\n merge_text_p (float): Probability of merging additional metadata to the description.\n drop_desc_p (float): Probability of dropping the original description on text merge.\n drop_other_p (float): Probability of dropping the other fields used for text augmentation.\n joint_embed_attributes (list[str]): A list of attributes for which joint embedding metadata is returned.\n paraphrase_source (str, optional): Path to the .json or .json.gz file containing the\n paraphrases for the description. The json should be a dict with keys are the\n original info path (e.g. track_path.json) and each value is a list of possible\n paraphrased.\n paraphrase_p (float): probability of taking a paraphrase.\n\n See `audiocraft.data.info_audio_dataset.InfoAudioDataset` for full initialization arguments.\n \"\"\"\n def __init__(self, *args, info_fields_required: bool = True,\n merge_text_p: float = 0., drop_desc_p: float = 0., drop_other_p: float = 0.,\n joint_embed_attributes: tp.List[str] = [],\n paraphrase_source: tp.Optional[str] = None, paraphrase_p: float = 0,\n **kwargs):\n kwargs['return_info'] = True # We require the info for each song of the dataset.\n super().__init__(*args, **kwargs)\n self.info_fields_required = info_fields_required\n self.merge_text_p = merge_text_p\n self.drop_desc_p = drop_desc_p\n self.drop_other_p = drop_other_p\n self.joint_embed_attributes = joint_embed_attributes\n self.paraphraser = None\n if paraphrase_source is not None:\n self.paraphraser = Paraphraser(paraphrase_source, paraphrase_p)\n\n def __getitem__(self, index):\n wav, info = super().__getitem__(index)\n info_data = info.to_dict()\n music_info_path = Path(info.meta.path).with_suffix('.json')\n\n if Path(music_info_path).exists():\n with open(music_info_path, 'r') as json_file:\n music_data = json.load(json_file)\n music_data.update(info_data)\n music_info = MusicInfo.from_dict(music_data, fields_required=self.info_fields_required)\n if self.paraphraser is not None:\n music_info.description = self.paraphraser.sample(music_info.meta.path, music_info.description)\n if self.merge_text_p:\n music_info = augment_music_info_description(\n music_info, self.merge_text_p, self.drop_desc_p, self.drop_other_p)\n else:\n music_info = MusicInfo.from_dict(info_data, fields_required=False)\n\n music_info.self_wav = WavCondition(\n wav=wav[None], length=torch.tensor([info.n_frames]),\n sample_rate=[info.sample_rate], path=[info.meta.path], seek_time=[info.seek_time])\n\n for att in self.joint_embed_attributes:\n att_value = getattr(music_info, att)\n joint_embed_cond = JointEmbedCondition(\n wav[None], [att_value], torch.tensor([info.n_frames]),\n sample_rate=[info.sample_rate], path=[info.meta.path], seek_time=[info.seek_time])\n music_info.joint_embed[att] = joint_embed_cond\n\n return wav, music_info\n\n\ndef get_musical_key(value: tp.Optional[str]) -> tp.Optional[str]:\n \"\"\"Preprocess key keywords, discarding them if there are multiple key defined.\"\"\"\n if value is None or (not isinstance(value, str)) or len(value) == 0 or value == 'None':\n return None\n elif ',' in value:\n # For now, we discard when multiple keys are defined separated with comas\n return None\n else:\n return value.strip().lower()\n\n\ndef get_bpm(value: tp.Optional[str]) -> tp.Optional[float]:\n \"\"\"Preprocess to a float.\"\"\"\n if value is None:\n return None\n try:\n return float(value)\n except ValueError:\n return None\n" }, { "path": "audiocraft/data/sound_dataset.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Dataset of audio with a simple description.\n\"\"\"\n\nfrom dataclasses import dataclass, fields, replace\nimport json\nfrom pathlib import Path\nimport random\nimport typing as tp\n\nimport numpy as np\nimport torch\n\nfrom .info_audio_dataset import (\n InfoAudioDataset,\n get_keyword_or_keyword_list\n)\nfrom ..modules.conditioners import (\n ConditioningAttributes,\n SegmentWithAttributes,\n WavCondition,\n)\n\n\nEPS = torch.finfo(torch.float32).eps\nTARGET_LEVEL_LOWER = -35\nTARGET_LEVEL_UPPER = -15\n\n\n@dataclass\nclass SoundInfo(SegmentWithAttributes):\n \"\"\"Segment info augmented with Sound metadata.\n \"\"\"\n description: tp.Optional[str] = None\n self_wav: tp.Optional[torch.Tensor] = None\n\n @property\n def has_sound_meta(self) -> bool:\n return self.description is not None\n\n def to_condition_attributes(self) -> ConditioningAttributes:\n out = ConditioningAttributes()\n\n for _field in fields(self):\n key, value = _field.name, getattr(self, _field.name)\n if key == 'self_wav':\n out.wav[key] = value\n else:\n out.text[key] = value\n return out\n\n @staticmethod\n def attribute_getter(attribute):\n if attribute == 'description':\n preprocess_func = get_keyword_or_keyword_list\n else:\n preprocess_func = None\n return preprocess_func\n\n @classmethod\n def from_dict(cls, dictionary: dict, fields_required: bool = False):\n _dictionary: tp.Dict[str, tp.Any] = {}\n\n # allow a subset of attributes to not be loaded from the dictionary\n # these attributes may be populated later\n post_init_attributes = ['self_wav']\n\n for _field in fields(cls):\n if _field.name in post_init_attributes:\n continue\n elif _field.name not in dictionary:\n if fields_required:\n raise KeyError(f\"Unexpected missing key: {_field.name}\")\n else:\n preprocess_func: tp.Optional[tp.Callable] = cls.attribute_getter(_field.name)\n value = dictionary[_field.name]\n if preprocess_func:\n value = preprocess_func(value)\n _dictionary[_field.name] = value\n return cls(**_dictionary)\n\n\nclass SoundDataset(InfoAudioDataset):\n \"\"\"Sound audio dataset: Audio dataset with environmental sound-specific metadata.\n\n Args:\n info_fields_required (bool): Whether all the mandatory metadata fields should be in the loaded metadata.\n external_metadata_source (tp.Optional[str]): Folder containing JSON metadata for the corresponding dataset.\n The metadata files contained in this folder are expected to match the stem of the audio file with\n a json extension.\n aug_p (float): Probability of performing audio mixing augmentation on the batch.\n mix_p (float): Proportion of batch items that are mixed together when applying audio mixing augmentation.\n mix_snr_low (int): Lowerbound for SNR value sampled for mixing augmentation.\n mix_snr_high (int): Upperbound for SNR value sampled for mixing augmentation.\n mix_min_overlap (float): Minimum overlap between audio files when performing mixing augmentation.\n kwargs: Additional arguments for AudioDataset.\n\n See `audiocraft.data.info_audio_dataset.InfoAudioDataset` for full initialization arguments.\n \"\"\"\n def __init__(\n self,\n *args,\n info_fields_required: bool = True,\n external_metadata_source: tp.Optional[str] = None,\n aug_p: float = 0.,\n mix_p: float = 0.,\n mix_snr_low: int = -5,\n mix_snr_high: int = 5,\n mix_min_overlap: float = 0.5,\n **kwargs\n ):\n kwargs['return_info'] = True # We require the info for each song of the dataset.\n super().__init__(*args, **kwargs)\n self.info_fields_required = info_fields_required\n self.external_metadata_source = external_metadata_source\n self.aug_p = aug_p\n self.mix_p = mix_p\n if self.aug_p > 0:\n assert self.mix_p > 0, \"Expecting some mixing proportion mix_p if aug_p > 0\"\n assert self.channels == 1, \"SoundDataset with audio mixing considers only monophonic audio\"\n self.mix_snr_low = mix_snr_low\n self.mix_snr_high = mix_snr_high\n self.mix_min_overlap = mix_min_overlap\n\n def _get_info_path(self, path: tp.Union[str, Path]) -> Path:\n \"\"\"Get path of JSON with metadata (description, etc.).\n If there exists a JSON with the same name as 'path.name', then it will be used.\n Else, such JSON will be searched for in an external json source folder if it exists.\n \"\"\"\n info_path = Path(path).with_suffix('.json')\n if Path(info_path).exists():\n return info_path\n elif self.external_metadata_source and (Path(self.external_metadata_source) / info_path.name).exists():\n return Path(self.external_metadata_source) / info_path.name\n else:\n raise Exception(f\"Unable to find a metadata JSON for path: {path}\")\n\n def __getitem__(self, index):\n wav, info = super().__getitem__(index)\n info_data = info.to_dict()\n info_path = self._get_info_path(info.meta.path)\n if Path(info_path).exists():\n with open(info_path, 'r') as json_file:\n sound_data = json.load(json_file)\n sound_data.update(info_data)\n sound_info = SoundInfo.from_dict(sound_data, fields_required=self.info_fields_required)\n # if there are multiple descriptions, sample one randomly\n if isinstance(sound_info.description, list):\n sound_info.description = random.choice(sound_info.description)\n else:\n sound_info = SoundInfo.from_dict(info_data, fields_required=False)\n\n sound_info.self_wav = WavCondition(\n wav=wav[None], length=torch.tensor([info.n_frames]),\n sample_rate=[sound_info.sample_rate], path=[info.meta.path], seek_time=[info.seek_time])\n\n return wav, sound_info\n\n def collater(self, samples):\n # when training, audio mixing is performed in the collate function\n wav, sound_info = super().collater(samples) # SoundDataset always returns infos\n if self.aug_p > 0:\n wav, sound_info = mix_samples(wav, sound_info, self.aug_p, self.mix_p,\n snr_low=self.mix_snr_low, snr_high=self.mix_snr_high,\n min_overlap=self.mix_min_overlap)\n return wav, sound_info\n\n\ndef rms_f(x: torch.Tensor) -> torch.Tensor:\n return (x ** 2).mean(1).pow(0.5)\n\n\ndef normalize(audio: torch.Tensor, target_level: int = -25) -> torch.Tensor:\n \"\"\"Normalize the signal to the target level.\"\"\"\n rms = rms_f(audio)\n scalar = 10 ** (target_level / 20) / (rms + EPS)\n audio = audio * scalar.unsqueeze(1)\n return audio\n\n\ndef is_clipped(audio: torch.Tensor, clipping_threshold: float = 0.99) -> torch.Tensor:\n return (abs(audio) > clipping_threshold).any(1)\n\n\ndef mix_pair(src: torch.Tensor, dst: torch.Tensor, min_overlap: float) -> torch.Tensor:\n start = random.randint(0, int(src.shape[1] * (1 - min_overlap)))\n remainder = src.shape[1] - start\n if dst.shape[1] > remainder:\n src[:, start:] = src[:, start:] + dst[:, :remainder]\n else:\n src[:, start:start+dst.shape[1]] = src[:, start:start+dst.shape[1]] + dst\n return src\n\n\ndef snr_mixer(clean: torch.Tensor, noise: torch.Tensor, snr: int, min_overlap: float,\n target_level: int = -25, clipping_threshold: float = 0.99) -> torch.Tensor:\n \"\"\"Function to mix clean speech and noise at various SNR levels.\n\n Args:\n clean (torch.Tensor): Clean audio source to mix, of shape [B, T].\n noise (torch.Tensor): Noise audio source to mix, of shape [B, T].\n snr (int): SNR level when mixing.\n min_overlap (float): Minimum overlap between the two mixed sources.\n target_level (int): Gain level in dB.\n clipping_threshold (float): Threshold for clipping the audio.\n Returns:\n torch.Tensor: The mixed audio, of shape [B, T].\n \"\"\"\n if clean.shape[1] > noise.shape[1]:\n noise = torch.nn.functional.pad(noise, (0, clean.shape[1] - noise.shape[1]))\n else:\n noise = noise[:, :clean.shape[1]]\n\n # normalizing to -25 dB FS\n clean = clean / (clean.max(1)[0].abs().unsqueeze(1) + EPS)\n clean = normalize(clean, target_level)\n rmsclean = rms_f(clean)\n\n noise = noise / (noise.max(1)[0].abs().unsqueeze(1) + EPS)\n noise = normalize(noise, target_level)\n rmsnoise = rms_f(noise)\n\n # set the noise level for a given SNR\n noisescalar = (rmsclean / (10 ** (snr / 20)) / (rmsnoise + EPS)).unsqueeze(1)\n noisenewlevel = noise * noisescalar\n\n # mix noise and clean speech\n noisyspeech = mix_pair(clean, noisenewlevel, min_overlap)\n\n # randomly select RMS value between -15 dBFS and -35 dBFS and normalize noisyspeech with that value\n # there is a chance of clipping that might happen with very less probability, which is not a major issue.\n noisy_rms_level = np.random.randint(TARGET_LEVEL_LOWER, TARGET_LEVEL_UPPER)\n rmsnoisy = rms_f(noisyspeech)\n scalarnoisy = (10 ** (noisy_rms_level / 20) / (rmsnoisy + EPS)).unsqueeze(1)\n noisyspeech = noisyspeech * scalarnoisy\n clean = clean * scalarnoisy\n noisenewlevel = noisenewlevel * scalarnoisy\n\n # final check to see if there are any amplitudes exceeding +/- 1. If so, normalize all the signals accordingly\n clipped = is_clipped(noisyspeech)\n if clipped.any():\n noisyspeech_maxamplevel = noisyspeech[clipped].max(1)[0].abs().unsqueeze(1) / (clipping_threshold - EPS)\n noisyspeech[clipped] = noisyspeech[clipped] / noisyspeech_maxamplevel\n\n return noisyspeech\n\n\ndef snr_mix(src: torch.Tensor, dst: torch.Tensor, snr_low: int, snr_high: int, min_overlap: float):\n if snr_low == snr_high:\n snr = snr_low\n else:\n snr = np.random.randint(snr_low, snr_high)\n mix = snr_mixer(src, dst, snr, min_overlap)\n return mix\n\n\ndef mix_text(src_text: str, dst_text: str):\n \"\"\"Mix text from different sources by concatenating them.\"\"\"\n if src_text == dst_text:\n return src_text\n return src_text + \" \" + dst_text\n\n\ndef mix_samples(wavs: torch.Tensor, infos: tp.List[SoundInfo], aug_p: float, mix_p: float,\n snr_low: int, snr_high: int, min_overlap: float):\n \"\"\"Mix samples within a batch, summing the waveforms and concatenating the text infos.\n\n Args:\n wavs (torch.Tensor): Audio tensors of shape [B, C, T].\n infos (list[SoundInfo]): List of SoundInfo items corresponding to the audio.\n aug_p (float): Augmentation probability.\n mix_p (float): Proportion of items in the batch to mix (and merge) together.\n snr_low (int): Lowerbound for sampling SNR.\n snr_high (int): Upperbound for sampling SNR.\n min_overlap (float): Minimum overlap between mixed samples.\n Returns:\n tuple[torch.Tensor, list[SoundInfo]]: A tuple containing the mixed wavs\n and mixed SoundInfo for the given batch.\n \"\"\"\n # no mixing to perform within the batch\n if mix_p == 0:\n return wavs, infos\n\n if random.uniform(0, 1) < aug_p:\n # perform all augmentations on waveforms as [B, T]\n # randomly picking pairs of audio to mix\n assert wavs.size(1) == 1, f\"Mix samples requires monophonic audio but C={wavs.size(1)}\"\n wavs = wavs.mean(dim=1, keepdim=False)\n B, T = wavs.shape\n k = int(mix_p * B)\n mixed_sources_idx = torch.randperm(B)[:k]\n mixed_targets_idx = torch.randperm(B)[:k]\n aug_wavs = snr_mix(\n wavs[mixed_sources_idx],\n wavs[mixed_targets_idx],\n snr_low,\n snr_high,\n min_overlap,\n )\n # mixing textual descriptions in metadata\n descriptions = [info.description for info in infos]\n aug_infos = []\n for i, j in zip(mixed_sources_idx, mixed_targets_idx):\n text = mix_text(descriptions[i], descriptions[j])\n m = replace(infos[i])\n m.description = text\n aug_infos.append(m)\n\n # back to [B, C, T]\n aug_wavs = aug_wavs.unsqueeze(1)\n assert aug_wavs.shape[0] > 0, \"Samples mixing returned empty batch.\"\n assert aug_wavs.dim() == 3, f\"Returned wav should be [B, C, T] but dim = {aug_wavs.dim()}\"\n assert aug_wavs.shape[0] == len(aug_infos), \"Mismatch between number of wavs and infos in the batch\"\n\n return aug_wavs, aug_infos # [B, C, T]\n else:\n # randomly pick samples in the batch to match\n # the batch size when performing audio mixing\n B, C, T = wavs.shape\n k = int(mix_p * B)\n wav_idx = torch.randperm(B)[:k]\n wavs = wavs[wav_idx]\n infos = [infos[i] for i in wav_idx]\n assert wavs.shape[0] == len(infos), \"Mismatch between number of wavs and infos in the batch\"\n\n return wavs, infos # [B, C, T]\n" }, { "path": "audiocraft/data/zip.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Utility for reading some info from inside a zip file.\n\"\"\"\n\nimport typing\nimport zipfile\n\nfrom dataclasses import dataclass\nfrom functools import lru_cache\nfrom typing_extensions import Literal\n\n\nDEFAULT_SIZE = 32\nMODE = Literal['r', 'w', 'x', 'a']\n\n\n@dataclass(order=True)\nclass PathInZip:\n \"\"\"Hold a path of file within a zip file.\n\n Args:\n path (str): The convention is :.\n Let's assume there is a zip file /some/location/foo.zip\n and inside of it is a json file located at /data/file1.json,\n Then we expect path = \"/some/location/foo.zip:/data/file1.json\".\n \"\"\"\n\n INFO_PATH_SEP = ':'\n zip_path: str\n file_path: str\n\n def __init__(self, path: str) -> None:\n split_path = path.split(self.INFO_PATH_SEP)\n assert len(split_path) == 2\n self.zip_path, self.file_path = split_path\n\n @classmethod\n def from_paths(cls, zip_path: str, file_path: str):\n return cls(zip_path + cls.INFO_PATH_SEP + file_path)\n\n def __str__(self) -> str:\n return self.zip_path + self.INFO_PATH_SEP + self.file_path\n\n\ndef _open_zip(path: str, mode: MODE = 'r'):\n return zipfile.ZipFile(path, mode)\n\n\n_cached_open_zip = lru_cache(DEFAULT_SIZE)(_open_zip)\n\n\ndef set_zip_cache_size(max_size: int):\n \"\"\"Sets the maximal LRU caching for zip file opening.\n\n Args:\n max_size (int): the maximal LRU cache.\n \"\"\"\n global _cached_open_zip\n _cached_open_zip = lru_cache(max_size)(_open_zip)\n\n\ndef open_file_in_zip(path_in_zip: PathInZip, mode: str = 'r') -> typing.IO:\n \"\"\"Opens a file stored inside a zip and returns a file-like object.\n\n Args:\n path_in_zip (PathInZip): A PathInZip object representing the file to return a file-like object of.\n mode (str): The mode in which to open the file with.\n Returns:\n A file-like object for PathInZip.\n \"\"\"\n zf = _cached_open_zip(path_in_zip.zip_path)\n return zf.open(path_in_zip.file_path)\n" }, { "path": "audiocraft/environment.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nProvides cluster and tools configuration across clusters (slurm, dora, utilities).\n\"\"\"\n\nimport logging\nimport os\nfrom pathlib import Path\nimport re\nimport typing as tp\n\nimport omegaconf\n\nfrom .utils.cluster import _guess_cluster_type\n\n\nlogger = logging.getLogger(__name__)\n\n\nclass AudioCraftEnvironment:\n \"\"\"Environment configuration for teams and clusters.\n\n AudioCraftEnvironment picks compute cluster settings (slurm, dora) from the current running environment\n or declared variable and the loaded team configuration. Additionally, the AudioCraftEnvironment\n provides pointers to a reference folder resolved automatically across clusters that is shared across team members,\n allowing to share sigs or other files to run jobs. Finally, it provides dataset mappers to automatically\n map dataset file paths to new locations across clusters, allowing to use the same manifest of files across cluters.\n\n The cluster type is identified automatically and base configuration file is read from config/teams.yaml.\n Use the following environment variables to specify the cluster, team or configuration:\n\n AUDIOCRAFT_CLUSTER (optional): Cluster type to enforce. Useful if the cluster type\n cannot be inferred automatically.\n AUDIOCRAFT_CONFIG (optional): Path to yaml config holding the teams configuration.\n If not set, configuration is read from config/teams.yaml.\n AUDIOCRAFT_TEAM (optional): Name of the team. Recommended to set to your own team.\n Cluster configuration are shared across teams to match compute allocation,\n specify your cluster configuration in the configuration file under a key mapping\n your team name.\n \"\"\"\n _instance = None\n DEFAULT_TEAM = \"default\"\n\n def __init__(self) -> None:\n \"\"\"Loads configuration.\"\"\"\n self.team: str = os.getenv(\"AUDIOCRAFT_TEAM\", self.DEFAULT_TEAM)\n cluster_type = _guess_cluster_type()\n cluster = os.getenv(\n \"AUDIOCRAFT_CLUSTER\", cluster_type.value\n )\n logger.info(\"Detecting cluster type %s\", cluster_type)\n\n self.cluster: str = cluster\n\n config_path = os.getenv(\n \"AUDIOCRAFT_CONFIG\",\n Path(__file__)\n .parent.parent.joinpath(\"config/teams\", self.team)\n .with_suffix(\".yaml\"),\n )\n self.config = omegaconf.OmegaConf.load(config_path)\n self._dataset_mappers = []\n cluster_config = self._get_cluster_config()\n if \"dataset_mappers\" in cluster_config:\n for pattern, repl in cluster_config[\"dataset_mappers\"].items():\n regex = re.compile(pattern)\n self._dataset_mappers.append((regex, repl))\n\n def _get_cluster_config(self) -> omegaconf.DictConfig:\n assert isinstance(self.config, omegaconf.DictConfig)\n return self.config[self.cluster]\n\n @classmethod\n def instance(cls):\n if cls._instance is None:\n cls._instance = cls()\n return cls._instance\n\n @classmethod\n def reset(cls):\n \"\"\"Clears the environment and forces a reload on next invocation.\"\"\"\n cls._instance = None\n\n @classmethod\n def get_team(cls) -> str:\n \"\"\"Gets the selected team as dictated by the AUDIOCRAFT_TEAM env var.\n If not defined, defaults to \"labs\".\n \"\"\"\n return cls.instance().team\n\n @classmethod\n def get_cluster(cls) -> str:\n \"\"\"Gets the detected cluster.\n This value can be overridden by the AUDIOCRAFT_CLUSTER env var.\n \"\"\"\n return cls.instance().cluster\n\n @classmethod\n def get_dora_dir(cls) -> Path:\n \"\"\"Gets the path to the dora directory for the current team and cluster.\n Value is overridden by the AUDIOCRAFT_DORA_DIR env var.\n \"\"\"\n cluster_config = cls.instance()._get_cluster_config()\n dora_dir = os.getenv(\"AUDIOCRAFT_DORA_DIR\", cluster_config[\"dora_dir\"])\n logger.warning(f\"Dora directory: {dora_dir}\")\n return Path(dora_dir)\n\n @classmethod\n def get_reference_dir(cls) -> Path:\n \"\"\"Gets the path to the reference directory for the current team and cluster.\n Value is overridden by the AUDIOCRAFT_REFERENCE_DIR env var.\n \"\"\"\n cluster_config = cls.instance()._get_cluster_config()\n return Path(os.getenv(\"AUDIOCRAFT_REFERENCE_DIR\", cluster_config[\"reference_dir\"]))\n\n @classmethod\n def get_slurm_exclude(cls) -> tp.Optional[str]:\n \"\"\"Get the list of nodes to exclude for that cluster.\"\"\"\n cluster_config = cls.instance()._get_cluster_config()\n return cluster_config.get(\"slurm_exclude\")\n\n @classmethod\n def get_slurm_partitions(cls, partition_types: tp.Optional[tp.List[str]] = None) -> str:\n \"\"\"Gets the requested partitions for the current team and cluster as a comma-separated string.\n\n Args:\n partition_types (list[str], optional): partition types to retrieve. Values must be\n from ['global', 'team']. If not provided, the global partition is returned.\n \"\"\"\n if not partition_types:\n partition_types = [\"global\"]\n\n cluster_config = cls.instance()._get_cluster_config()\n partitions = [\n cluster_config[\"partitions\"][partition_type]\n for partition_type in partition_types\n ]\n return \",\".join(partitions)\n\n @classmethod\n def resolve_reference_path(cls, path: tp.Union[str, Path]) -> Path:\n \"\"\"Converts reference placeholder in path with configured reference dir to resolve paths.\n\n Args:\n path (str or Path): Path to resolve.\n Returns:\n Path: Resolved path.\n \"\"\"\n path = str(path)\n\n if path.startswith(\"//reference\"):\n reference_dir = cls.get_reference_dir()\n logger.warn(f\"Reference directory: {reference_dir}\")\n assert (\n reference_dir.exists() and reference_dir.is_dir()\n ), f\"Reference directory does not exist: {reference_dir}.\"\n path = re.sub(\"^//reference\", str(reference_dir), path)\n\n return Path(path)\n\n @classmethod\n def apply_dataset_mappers(cls, path: str) -> str:\n \"\"\"Applies dataset mapping regex rules as defined in the configuration.\n If no rules are defined, the path is returned as-is.\n \"\"\"\n instance = cls.instance()\n\n for pattern, repl in instance._dataset_mappers:\n path = pattern.sub(repl, path)\n\n return path\n" }, { "path": "audiocraft/grids/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Dora Grids.\"\"\"\n" }, { "path": "audiocraft/grids/_base_explorers.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom abc import ABC, abstractmethod\nimport time\nimport typing as tp\nfrom dora import Explorer\nimport treetable as tt\n\n\ndef get_sheep_ping(sheep) -> tp.Optional[str]:\n \"\"\"Return the amount of time since the Sheep made some update\n to its log. Returns a str using the relevant time unit.\"\"\"\n ping = None\n if sheep.log is not None and sheep.log.exists():\n delta = time.time() - sheep.log.stat().st_mtime\n if delta > 3600 * 24:\n ping = f'{delta / (3600 * 24):.1f}d'\n elif delta > 3600:\n ping = f'{delta / (3600):.1f}h'\n elif delta > 60:\n ping = f'{delta / 60:.1f}m'\n else:\n ping = f'{delta:.1f}s'\n return ping\n\n\nclass BaseExplorer(ABC, Explorer):\n \"\"\"Base explorer for AudioCraft grids.\n\n All task specific solvers are expected to implement the `get_grid_metrics`\n method to specify logic about metrics to display for a given task.\n\n If additional stages are used, the child explorer must define how to handle\n these new stages in the `process_history` and `process_sheep` methods.\n \"\"\"\n def stages(self):\n return [\"train\", \"valid\", \"evaluate\"]\n\n def get_grid_meta(self):\n \"\"\"Returns the list of Meta information to display for each XP/job.\n \"\"\"\n return [\n tt.leaf(\"index\", align=\">\"),\n tt.leaf(\"name\", wrap=140),\n tt.leaf(\"state\"),\n tt.leaf(\"sig\", align=\">\"),\n tt.leaf(\"sid\", align=\"<\"),\n ]\n\n @abstractmethod\n def get_grid_metrics(self):\n \"\"\"Return the metrics that should be displayed in the tracking table.\n \"\"\"\n ...\n\n def process_sheep(self, sheep, history):\n train = {\n \"epoch\": len(history),\n }\n parts = {\"train\": train}\n for metrics in history:\n for key, sub in metrics.items():\n part = parts.get(key, {})\n if 'duration' in sub:\n # Convert to minutes for readability.\n sub['duration'] = sub['duration'] / 60.\n part.update(sub)\n parts[key] = part\n ping = get_sheep_ping(sheep)\n if ping is not None:\n for name in self.stages():\n if name not in parts:\n parts[name] = {}\n # Add the ping to each part for convenience.\n parts[name]['ping'] = ping\n return parts\n" }, { "path": "audiocraft/grids/audiogen/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"AudioGen grids.\"\"\"\n" }, { "path": "audiocraft/grids/audiogen/audiogen_base_16khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom ..musicgen._explorers import LMExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@LMExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=64, partition=partitions)\n launcher.bind_(solver='audiogen/audiogen_base_16khz')\n # replace this by the desired environmental sound dataset\n launcher.bind_(dset='internal/sounds_16khz')\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n medium = {'model/lm/model_scale': 'medium'}\n\n launcher.bind_(fsdp)\n launcher(medium)\n" }, { "path": "audiocraft/grids/audiogen/audiogen_pretrained_16khz_eval.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nEvaluation with objective metrics for the pretrained AudioGen models.\nThis grid takes signature from the training grid and runs evaluation-only stage.\n\nWhen running the grid for the first time, please use:\nREGEN=1 dora grid audiogen.audiogen_pretrained_16khz_eval\nand re-use the REGEN=1 option when the grid is changed to force regenerating it.\n\nNote that you need the proper metrics external libraries setup to use all\nthe objective metrics activated in this grid. Refer to the README for more information.\n\"\"\"\n\nimport os\n\nfrom ..musicgen._explorers import GenerationEvalExplorer\nfrom ...environment import AudioCraftEnvironment\nfrom ... import train\n\n\ndef eval(launcher, batch_size: int = 32):\n opts = {\n 'dset': 'audio/audiocaps_16khz',\n 'solver/audiogen/evaluation': 'objective_eval',\n 'execute_only': 'evaluate',\n '+dataset.evaluate.batch_size': batch_size,\n '+metrics.fad.tf.batch_size': 32,\n }\n # binary for FAD computation: replace this path with your own path\n metrics_opts = {\n 'metrics.fad.tf.bin': '/data/home/jadecopet/local/usr/opt/google-research'\n }\n opt1 = {'generate.lm.use_sampling': True, 'generate.lm.top_k': 250, 'generate.lm.top_p': 0.}\n opt2 = {'transformer_lm.two_step_cfg': True}\n\n sub = launcher.bind(opts)\n sub.bind_(metrics_opts)\n\n # base objective metrics\n sub(opt1, opt2)\n\n\n@GenerationEvalExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=4, partition=partitions)\n\n if 'REGEN' not in os.environ:\n folder = train.main.dora.dir / 'grids' / __name__.split('.', 2)[-1]\n with launcher.job_array():\n for sig in folder.iterdir():\n if not sig.is_symlink():\n continue\n xp = train.main.get_xp_from_sig(sig.name)\n launcher(xp.argv)\n return\n\n audiogen_base = launcher.bind(solver=\"audiogen/audiogen_base_16khz\")\n audiogen_base.bind_({'autocast': False, 'fsdp.use': True})\n\n audiogen_base_medium = audiogen_base.bind({'continue_from': '//pretrained/facebook/audiogen-medium'})\n audiogen_base_medium.bind_({'model/lm/model_scale': 'medium'})\n eval(audiogen_base_medium, batch_size=128)\n" }, { "path": "audiocraft/grids/compression/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"EnCodec grids.\"\"\"\n" }, { "path": "audiocraft/grids/compression/_explorers.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport treetable as tt\n\nfrom .._base_explorers import BaseExplorer\n\n\nclass CompressionExplorer(BaseExplorer):\n eval_metrics = [\"sisnr\", \"visqol\"]\n\n def stages(self):\n return [\"train\", \"valid\", \"evaluate\"]\n\n def get_grid_meta(self):\n \"\"\"Returns the list of Meta information to display for each XP/job.\n \"\"\"\n return [\n tt.leaf(\"index\", align=\">\"),\n tt.leaf(\"name\", wrap=140),\n tt.leaf(\"state\"),\n tt.leaf(\"sig\", align=\">\"),\n ]\n\n def get_grid_metrics(self):\n \"\"\"Return the metrics that should be displayed in the tracking table.\n \"\"\"\n return [\n tt.group(\n \"train\",\n [\n tt.leaf(\"epoch\"),\n tt.leaf(\"bandwidth\", \".2f\"),\n tt.leaf(\"adv\", \".4f\"),\n tt.leaf(\"d_loss\", \".4f\"),\n ],\n align=\">\",\n ),\n tt.group(\n \"valid\",\n [\n tt.leaf(\"bandwidth\", \".2f\"),\n tt.leaf(\"adv\", \".4f\"),\n tt.leaf(\"msspec\", \".4f\"),\n tt.leaf(\"sisnr\", \".2f\"),\n ],\n align=\">\",\n ),\n tt.group(\n \"evaluate\", [tt.leaf(name, \".3f\") for name in self.eval_metrics], align=\">\"\n ),\n ]\n" }, { "path": "audiocraft/grids/compression/debug.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nGrid search file, simply list all the exp you want in `explorer`.\nAny new exp added there will be scheduled.\nYou can cancel and experiment by commenting its line.\n\nThis grid is a minimal example for debugging compression task\nand how to override parameters directly in a grid.\nLearn more about dora grids: https://github.com/facebookresearch/dora\n\"\"\"\n\nfrom ._explorers import CompressionExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@CompressionExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=2, partition=partitions)\n launcher.bind_(solver='compression/debug')\n\n with launcher.job_array():\n # base debug task using config from solver=compression/debug\n launcher()\n # we can override parameters in the grid to launch additional xps\n launcher({'rvq.bins': 2048, 'rvq.n_q': 4})\n" }, { "path": "audiocraft/grids/compression/encodec_audiogen_16khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nGrid search file, simply list all the exp you want in `explorer`.\nAny new exp added there will be scheduled.\nYou can cancel and experiment by commenting its line.\n\nThis grid shows how to train the new AudioGen EnCodec model at 16 kHz.\n\"\"\"\n\nfrom ._explorers import CompressionExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@CompressionExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=8, partition=partitions)\n # use configuration for AudioGen's EnCodec model trained on monophonic audio sampled at 16 kHz\n # AudioGen's EnCodec is trained with a total stride of 320 leading to a frame rate of 50 hz\n launcher.bind_(solver='compression/encodec_audiogen_16khz')\n # replace this by the desired sound dataset\n launcher.bind_(dset='internal/sounds_16khz')\n # launch xp\n launcher()\n" }, { "path": "audiocraft/grids/compression/encodec_base_24khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nGrid search file, simply list all the exp you want in `explorer`.\nAny new exp added there will be scheduled.\nYou can cancel and experiment by commenting its line.\n\nThis grid shows how to train a base causal EnCodec model at 24 kHz.\n\"\"\"\n\nfrom ._explorers import CompressionExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@CompressionExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=8, partition=partitions)\n # base causal EnCodec trained on monophonic audio sampled at 24 kHz\n launcher.bind_(solver='compression/encodec_base_24khz')\n # replace this by the desired dataset\n launcher.bind_(dset='audio/example')\n # launch xp\n launcher()\n" }, { "path": "audiocraft/grids/compression/encodec_musicgen_32khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nGrid search file, simply list all the exp you want in `explorer`.\nAny new exp added there will be scheduled.\nYou can cancel and experiment by commenting its line.\n\nThis grid shows how to train a MusicGen EnCodec model at 32 kHz.\n\"\"\"\n\nfrom ._explorers import CompressionExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@CompressionExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=8, partition=partitions)\n # use configuration for MusicGen's EnCodec model trained on monophonic audio sampled at 32 kHz\n # MusicGen's EnCodec is trained with a total stride of 640 leading to a frame rate of 50 hz\n launcher.bind_(solver='compression/encodec_musicgen_32khz')\n # replace this by the desired music dataset\n launcher.bind_(dset='internal/music_400k_32khz')\n # launch xp\n launcher()\n launcher({\n 'metrics.visqol.bin': '/data/home/jadecopet/local/usr/opt/visqol',\n 'label': 'visqol',\n 'evaluate.metrics.visqol': True\n })\n" }, { "path": "audiocraft/grids/diffusion/4_bands_base_32khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nTraining of the 4 diffusion models described in\n\"From Discrete Tokens to High-Fidelity Audio Using Multi-Band Diffusion\"\n(paper link).\n\"\"\"\n\nfrom ._explorers import DiffusionExplorer\n\n\n@DiffusionExplorer\ndef explorer(launcher):\n launcher.slurm_(gpus=4, partition='learnfair')\n\n launcher.bind_({'solver': 'diffusion/default',\n 'dset': 'internal/music_10k_32khz'})\n\n with launcher.job_array():\n launcher({'filter.use': True, 'filter.idx_band': 0, \"processor.use\": False, 'processor.power_std': 0.4})\n launcher({'filter.use': True, 'filter.idx_band': 1, \"processor.use\": False, 'processor.power_std': 0.4})\n launcher({'filter.use': True, 'filter.idx_band': 2, \"processor.use\": True, 'processor.power_std': 0.4})\n launcher({'filter.use': True, 'filter.idx_band': 3, \"processor.use\": True, 'processor.power_std': 0.75})\n" }, { "path": "audiocraft/grids/diffusion/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Diffusion grids.\"\"\"\n" }, { "path": "audiocraft/grids/diffusion/_explorers.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport treetable as tt\n\nfrom .._base_explorers import BaseExplorer\n\n\nclass DiffusionExplorer(BaseExplorer):\n eval_metrics = [\"sisnr\", \"visqol\"]\n\n def stages(self):\n return [\"train\", \"valid\", \"valid_ema\", \"evaluate\", \"evaluate_ema\"]\n\n def get_grid_meta(self):\n \"\"\"Returns the list of Meta information to display for each XP/job.\n \"\"\"\n return [\n tt.leaf(\"index\", align=\">\"),\n tt.leaf(\"name\", wrap=140),\n tt.leaf(\"state\"),\n tt.leaf(\"sig\", align=\">\"),\n ]\n\n def get_grid_metrics(self):\n \"\"\"Return the metrics that should be displayed in the tracking table.\n \"\"\"\n return [\n tt.group(\n \"train\",\n [\n tt.leaf(\"epoch\"),\n tt.leaf(\"loss\", \".3%\"),\n ],\n align=\">\",\n ),\n tt.group(\n \"valid\",\n [\n tt.leaf(\"loss\", \".3%\"),\n # tt.leaf(\"loss_0\", \".3%\"),\n ],\n align=\">\",\n ),\n tt.group(\n \"valid_ema\",\n [\n tt.leaf(\"loss\", \".3%\"),\n # tt.leaf(\"loss_0\", \".3%\"),\n ],\n align=\">\",\n ),\n tt.group(\n \"evaluate\", [tt.leaf(\"rvm\", \".4f\"), tt.leaf(\"rvm_0\", \".4f\"),\n tt.leaf(\"rvm_1\", \".4f\"), tt.leaf(\"rvm_2\", \".4f\"),\n tt.leaf(\"rvm_3\", \".4f\"), ], align=\">\"\n ),\n tt.group(\n \"evaluate_ema\", [tt.leaf(\"rvm\", \".4f\"), tt.leaf(\"rvm_0\", \".4f\"),\n tt.leaf(\"rvm_1\", \".4f\"), tt.leaf(\"rvm_2\", \".4f\"),\n tt.leaf(\"rvm_3\", \".4f\")], align=\">\"\n ),\n ]\n" }, { "path": "audiocraft/grids/magnet/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"MAGNeT grids.\"\"\"\n" }, { "path": "audiocraft/grids/magnet/audio_magnet_16khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom ..musicgen._explorers import LMExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@LMExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=32, partition=partitions)\n launcher.bind_(solver='magnet/audio_magnet_16khz')\n # replace this by the desired environmental sound dataset\n launcher.bind_(dset='internal/sounds_16khz')\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n medium = {'model/lm/model_scale': 'medium'}\n\n # Small model (300M)\n launcher.slurm_(gpus=32).bind_(label='32gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub()\n\n # Medium model (1.5B)\n launcher.slurm_(gpus=64).bind_(label='64gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub(medium, fsdp)\n" }, { "path": "audiocraft/grids/magnet/audio_magnet_pretrained_16khz_eval.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nEvaluation with objective metrics for the pretrained audio-MAGNeT models.\nThis grid takes signature from the training grid and runs evaluation-only stage.\n\nWhen running the grid for the first time, please use:\nREGEN=1 dora grid magnet.audio_magnet_pretrained_16khz_eval\nand re-use the REGEN=1 option when the grid is changed to force regenerating it.\n\nNote that you need the proper metrics external libraries setup to use all\nthe objective metrics activated in this grid. Refer to the README for more information.\n\"\"\"\n\nimport os\n\nfrom ..musicgen._explorers import GenerationEvalExplorer\nfrom ...environment import AudioCraftEnvironment\nfrom ... import train\n\n\ndef eval(launcher, batch_size: int = 32):\n opts = {\n 'dset': 'audio/audiocaps_16khz',\n 'solver/audiogen/evaluation': 'objective_eval',\n 'execute_only': 'evaluate',\n '+dataset.evaluate.batch_size': batch_size,\n '+metrics.fad.tf.batch_size': 32,\n }\n # binary for FAD computation: replace this path with your own path\n metrics_opts = {\n 'metrics.fad.tf.bin': '/data/home/jadecopet/local/usr/opt/google-research'\n }\n\n sub = launcher.bind(opts)\n sub.bind_(metrics_opts)\n\n # base objective metrics\n sub()\n\n\n@GenerationEvalExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=4, partition=partitions)\n\n if 'REGEN' not in os.environ:\n folder = train.main.dora.dir / 'grids' / __name__.split('.', 2)[-1]\n with launcher.job_array():\n for sig in folder.iterdir():\n if not sig.is_symlink():\n continue\n xp = train.main.get_xp_from_sig(sig.name)\n launcher(xp.argv)\n return\n\n with launcher.job_array():\n audio_magnet = launcher.bind(solver=\"magnet/audio_magnet_16khz\")\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n\n # Small audio-MAGNeT model (300M)\n audio_magnet_small = audio_magnet.bind({'continue_from': '//pretrained/facebook/audio-magnet-small'})\n eval(audio_magnet_small, batch_size=128)\n\n # Medium audio-MAGNeT model (1.5B)\n audio_magnet_medium = audio_magnet.bind({'continue_from': '//pretrained/facebook/audio-magnet-medium'})\n audio_magnet_medium.bind_({'model/lm/model_scale': 'medium'})\n audio_magnet_medium.bind_(fsdp)\n eval(audio_magnet_medium, batch_size=128)\n" }, { "path": "audiocraft/grids/magnet/magnet_32khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom ..musicgen._explorers import LMExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@LMExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=32, partition=partitions)\n launcher.bind_(solver='magnet/magnet_32khz')\n # replace this by the desired music dataset\n launcher.bind_(dset='internal/music_400k_32khz')\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n medium = {'model/lm/model_scale': 'medium'}\n adam = {'optim.optimizer': 'adamw', 'optim.lr': 1e-4}\n segdur_10secs = {'dataset.segment_duration': 10,\n 'dataset.batch_size': 576,\n 'generate.lm.decoding_steps': [20, 10, 10, 10]}\n\n # Small models (300M)\n launcher.slurm_(gpus=32).bind_(label='32gpus')\n with launcher.job_array():\n # 30 seconds\n sub = launcher.bind()\n sub()\n\n # 10 seconds\n sub = launcher.bind()\n sub(segdur_10secs)\n\n # Medium models (1.5B)\n launcher.bind_(fsdp)\n launcher.slurm_(gpus=64).bind_(label='64gpus')\n with launcher.job_array():\n # 30 seconds\n sub = launcher.bind()\n sub(medium, adam)\n\n # 10 seconds\n sub = launcher.bind()\n sub(segdur_10secs)\n" }, { "path": "audiocraft/grids/magnet/magnet_pretrained_32khz_eval.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nEvaluation with objective metrics for the pretrained MAGNeT models.\nThis grid takes signature from the training grid and runs evaluation-only stage.\n\nWhen running the grid for the first time, please use:\nREGEN=1 dora grid magnet.magnet_pretrained_32khz_eval\nand re-use the REGEN=1 option when the grid is changed to force regenerating it.\n\nNote that you need the proper metrics external libraries setup to use all\nthe objective metrics activated in this grid. Refer to the README for more information.\n\"\"\"\n\nimport os\n\nfrom ..musicgen._explorers import GenerationEvalExplorer\nfrom ...environment import AudioCraftEnvironment\nfrom ... import train\n\n\ndef eval(launcher, batch_size: int = 32):\n opts = {\n 'dset': 'audio/musiccaps_32khz',\n 'solver/musicgen/evaluation': 'objective_eval',\n 'execute_only': 'evaluate',\n '+dataset.evaluate.batch_size': batch_size,\n '+metrics.fad.tf.batch_size': 16,\n }\n # binary for FAD computation: replace this path with your own path\n metrics_opts = {\n 'metrics.fad.tf.bin': '/data/home/jadecopet/local/usr/opt/google-research'\n }\n\n sub = launcher.bind(opts)\n sub.bind_(metrics_opts)\n\n # base objective metrics\n sub()\n\n\n@GenerationEvalExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=4, partition=partitions)\n\n if 'REGEN' not in os.environ:\n folder = train.main.dora.dir / 'grids' / __name__.split('.', 2)[-1]\n with launcher.job_array():\n for sig in folder.iterdir():\n if not sig.is_symlink():\n continue\n xp = train.main.get_xp_from_sig(sig.name)\n launcher(xp.argv)\n return\n\n with launcher.job_array():\n magnet = launcher.bind(solver=\"magnet/magnet_32khz\")\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n\n segdur_10secs = {'dataset.segment_duration': 10,\n 'generate.lm.decoding_steps': [20, 10, 10, 10]}\n\n # 10-second magnet models\n magnet_small_10secs = magnet.bind({'continue_from': '//pretrained/facebook/magnet-small-10secs'})\n magnet_small_10secs.bind_(segdur_10secs)\n eval(magnet_small_10secs, batch_size=128)\n\n magnet_medium_10secs = magnet.bind({'continue_from': '//pretrained/facebook/magnet-medium-10secs'})\n magnet_medium_10secs.bind_(segdur_10secs)\n magnet_medium_10secs.bind_({'model/lm/model_scale': 'medium'})\n magnet_medium_10secs.bind_(fsdp)\n eval(magnet_medium_10secs, batch_size=128)\n\n # 30-second magnet models\n magnet_small_30secs = magnet.bind({'continue_from': '//pretrained/facebook/magnet-small-30secs'})\n eval(magnet_small_30secs, batch_size=128)\n\n magnet_medium_30secs = magnet.bind({'continue_from': '//pretrained/facebook/magnet-medium-30secs'})\n magnet_medium_30secs.bind_({'model/lm/model_scale': 'medium'})\n magnet_medium_30secs.bind_(fsdp)\n eval(magnet_medium_30secs, batch_size=128)\n" }, { "path": "audiocraft/grids/musicgen/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"MusicGen grids.\"\"\"\n" }, { "path": "audiocraft/grids/musicgen/_explorers.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nimport treetable as tt\n\nfrom .._base_explorers import BaseExplorer\n\n\nclass LMExplorer(BaseExplorer):\n eval_metrics: tp.List[str] = []\n\n def stages(self) -> tp.List[str]:\n return ['train', 'valid']\n\n def get_grid_metrics(self):\n \"\"\"Return the metrics that should be displayed in the tracking table.\"\"\"\n return [\n tt.group(\n 'train',\n [\n tt.leaf('epoch'),\n tt.leaf('duration', '.1f'), # duration in minutes\n tt.leaf('ping'),\n tt.leaf('ce', '.4f'), # cross entropy\n tt.leaf(\"ppl\", '.3f'), # perplexity\n ],\n align='>',\n ),\n tt.group(\n 'valid',\n [\n tt.leaf('ce', '.4f'),\n tt.leaf('ppl', '.3f'),\n tt.leaf('best_ppl', '.3f'),\n ],\n align='>',\n ),\n ]\n\n def process_sheep(self, sheep, history):\n parts = super().process_sheep(sheep, history)\n\n track_by = {'ppl': 'lower'} # values should be in ['lower', 'higher']\n best_metrics = {k: (1 if v == 'lower' else -1) * float('inf') for k, v in track_by.items()}\n\n def comparator(mode, a, b):\n return a < b if mode == 'lower' else a > b\n\n for metrics in history:\n for key, sub in metrics.items():\n for metric in track_by:\n # for the validation set, keep track of best metrics (ppl in this example)\n # this is so we can conveniently compare metrics between runs in the grid\n if key == 'valid' and metric in sub and comparator(\n track_by[metric], sub[metric], best_metrics[metric]\n ):\n best_metrics[metric] = sub[metric]\n\n if 'valid' in parts:\n parts['valid'].update({f'best_{k}': v for k, v in best_metrics.items()})\n return parts\n\n\nclass GenerationEvalExplorer(BaseExplorer):\n eval_metrics: tp.List[str] = []\n\n def stages(self) -> tp.List[str]:\n return ['evaluate']\n\n def get_grid_metrics(self):\n \"\"\"Return the metrics that should be displayed in the tracking table.\"\"\"\n return [\n tt.group(\n 'evaluate',\n [\n tt.leaf('epoch', '.3f'),\n tt.leaf('duration', '.1f'),\n tt.leaf('ping'),\n tt.leaf('ce', '.4f'),\n tt.leaf('ppl', '.3f'),\n tt.leaf('fad', '.3f'),\n tt.leaf('kld', '.3f'),\n tt.leaf('text_consistency', '.3f'),\n tt.leaf('chroma_cosine', '.3f'),\n ],\n align='>',\n ),\n ]\n" }, { "path": "audiocraft/grids/musicgen/musicgen_base_32khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom ._explorers import LMExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@LMExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=32, partition=partitions)\n launcher.bind_(solver='musicgen/musicgen_base_32khz')\n # replace this by the desired music dataset\n launcher.bind_(dset='internal/music_400k_32khz')\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n medium = {'model/lm/model_scale': 'medium'}\n large = {'model/lm/model_scale': 'large'}\n\n cfg_low = {'classifier_free_guidance.training_dropout': 0.2}\n wd_low = {'conditioners.description.t5.word_dropout': 0.2}\n\n adam = {'optim.optimizer': 'adamw', 'optim.lr': 1e-4}\n\n launcher.bind_(fsdp)\n\n launcher.slurm_(gpus=32).bind_(label='32gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub()\n\n launcher.slurm_(gpus=64).bind_(label='64gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub(medium, adam)\n\n launcher.slurm_(gpus=96).bind_(label='96gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub(large, cfg_low, wd_low, adam, {'optim.max_norm': 3})\n" }, { "path": "audiocraft/grids/musicgen/musicgen_base_cached_32khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom ._explorers import LMExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@LMExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=32, partition=partitions)\n launcher.bind_(solver='musicgen/musicgen_base_32khz')\n # replace this by the desired music dataset\n launcher.bind_(dset='internal/music_400k_32khz')\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n medium = {'model/lm/model_scale': 'medium'}\n large = {'model/lm/model_scale': 'large'}\n\n cfg_low = {'classifier_free_guidance.training_dropout': 0.2}\n wd_low = {'conditioners.description.t5.word_dropout': 0.2}\n\n adam = {'optim.optimizer': 'adamw', 'optim.lr': 1e-4}\n\n # BEGINNING OF CACHE WRITING JOBS.\n cache_write = {\n 'cache.path': '/fsx-codegen/defossez/cache/interleave_stereo_nv_32k',\n 'cache.write': True,\n 'generate.every': 500,\n 'evaluate.every': 500,\n 'logging.log_updates': 50,\n }\n\n cache_sub = launcher.bind({'model/lm/model_scale': 'xsmall', 'conditioner': 'none'})\n cache_sub.bind_({'deadlock.use': True})\n cache_sub.slurm_(gpus=8)\n with launcher.job_array():\n num_shards = 10 # total number of jobs running in parallel.\n for shard in range(0, num_shards):\n launcher(cache_write, {'cache.write_num_shards': num_shards, 'cache.write_shard': shard})\n\n # REMOVE THE FOLLOWING RETURN STATEMENT ONCE THE ABOVE JOBS ARE DONE,\n # OR SUFFICIENTLY AHEAD.\n return\n\n cache = {\n 'cache.path': '/fsx-codegen/defossez/cache/interleave_stereo_nv_32k',\n }\n launcher.bind_(fsdp, cache)\n\n launcher.slurm_(gpus=32).bind_(label='32gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub()\n\n launcher.slurm_(gpus=64).bind_(label='64gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub(medium, adam)\n\n launcher.slurm_(gpus=96).bind_(label='96gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub(large, cfg_low, wd_low, adam, {'optim.max_norm': 3})\n" }, { "path": "audiocraft/grids/musicgen/musicgen_clapemb_32khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom ._explorers import LMExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@LMExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=32, partition=partitions)\n launcher.bind_(solver='musicgen/musicgen_base_32khz')\n # replace this by the desired music dataset\n launcher.bind_(dset='internal/music_400k_32khz')\n launcher.bind_(conditioner='clapemb2music')\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n cache_path = {'conditioners.description.clap.cache_path':\n '/fsx-audio-craft-llm/jadecopet/experiments/audiocraft/caches/clap_embed_music'}\n text_wav_training_opt = {'conditioners.description.clap.text_p': 0.5}\n\n launcher.bind_(fsdp)\n\n launcher.slurm_(gpus=32).bind_(label='32gpus')\n with launcher.job_array():\n launcher()\n launcher(text_wav_training_opt)\n launcher(cache_path)\n launcher(cache_path, text_wav_training_opt)\n" }, { "path": "audiocraft/grids/musicgen/musicgen_melody_32khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom ._explorers import LMExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@LMExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=32, partition=partitions)\n launcher.bind_(solver='musicgen/musicgen_melody_32khz')\n # replace this by the desired music dataset\n launcher.bind_(dset='internal/music_400k_32khz')\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n medium = {'model/lm/model_scale': 'medium'}\n large = {'model/lm/model_scale': 'large'}\n\n cfg_low = {'classifier_free_guidance.training_dropout': 0.2}\n wd_low = {'conditioners.description.t5.word_dropout': 0.2}\n\n adam = {'optim.optimizer': 'adamw', 'optim.lr': 1e-4}\n\n cache_path = {'conditioners.self_wav.chroma_stem.cache_path':\n '/fsx-audio-craft-llm/jadecopet/experiments/audiocraft/caches/chroma_stem'}\n\n # CACHE GENERATION JOBS\n n_cache_gen_jobs = 4\n gen_sub = launcher.slurm(gpus=1)\n gen_sub.bind_(\n cache_path, {\n # the cache is always computed over the whole file, so duration doesn't matter here.\n 'dataset.segment_duration': 2.,\n 'dataset.batch_size': 8,\n 'dataset.train.permutation_on_files': True, # try to not repeat files.\n 'optim.epochs': 10,\n 'model/lm/model_scale': 'xsmall',\n\n })\n with gen_sub.job_array():\n for gen_job in range(n_cache_gen_jobs):\n gen_sub({'dataset.train.shuffle_seed': gen_job})\n\n # ACTUAL TRAINING JOBS.\n launcher.bind_(fsdp)\n\n launcher.slurm_(gpus=32).bind_(label='32gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub()\n sub(cache_path)\n\n launcher.slurm_(gpus=64).bind_(label='64gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub(medium, adam)\n\n launcher.slurm_(gpus=96).bind_(label='96gpus')\n with launcher.job_array():\n sub = launcher.bind()\n sub(large, cfg_low, wd_low, adam, {'optim.max_norm': 3})\n" }, { "path": "audiocraft/grids/musicgen/musicgen_pretrained_32khz_eval.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nEvaluation with objective metrics for the pretrained MusicGen models.\nThis grid takes signature from the training grid and runs evaluation-only stage.\n\nWhen running the grid for the first time, please use:\nREGEN=1 dora grid musicgen.musicgen_pretrained_32khz_eval\nand re-use the REGEN=1 option when the grid is changed to force regenerating it.\n\nNote that you need the proper metrics external libraries setup to use all\nthe objective metrics activated in this grid. Refer to the README for more information.\n\"\"\"\n\nimport os\n\nfrom ._explorers import GenerationEvalExplorer\nfrom ...environment import AudioCraftEnvironment\nfrom ... import train\n\n\ndef eval(launcher, batch_size: int = 32, eval_melody: bool = False):\n opts = {\n 'dset': 'audio/musiccaps_32khz',\n 'solver/musicgen/evaluation': 'objective_eval',\n 'execute_only': 'evaluate',\n '+dataset.evaluate.batch_size': batch_size,\n '+metrics.fad.tf.batch_size': 16,\n }\n # chroma-specific evaluation\n chroma_opts = {\n 'dset': 'internal/music_400k_32khz',\n 'dataset.evaluate.segment_duration': 30,\n 'dataset.evaluate.num_samples': 1000,\n 'evaluate.metrics.chroma_cosine': True,\n 'evaluate.metrics.fad': False,\n 'evaluate.metrics.kld': False,\n 'evaluate.metrics.text_consistency': False,\n }\n # binary for FAD computation: replace this path with your own path\n metrics_opts = {\n 'metrics.fad.tf.bin': '/data/home/jadecopet/local/usr/opt/google-research'\n }\n opt1 = {'generate.lm.use_sampling': True, 'generate.lm.top_k': 250, 'generate.lm.top_p': 0.}\n opt2 = {'transformer_lm.two_step_cfg': True}\n\n sub = launcher.bind(opts)\n sub.bind_(metrics_opts)\n\n # base objective metrics\n sub(opt1, opt2)\n\n if eval_melody:\n # chroma-specific metrics\n sub(opt1, opt2, chroma_opts)\n\n\n@GenerationEvalExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=4, partition=partitions)\n\n if 'REGEN' not in os.environ:\n folder = train.main.dora.dir / 'grids' / __name__.split('.', 2)[-1]\n with launcher.job_array():\n for sig in folder.iterdir():\n if not sig.is_symlink():\n continue\n xp = train.main.get_xp_from_sig(sig.name)\n launcher(xp.argv)\n return\n\n with launcher.job_array():\n musicgen_base = launcher.bind(solver=\"musicgen/musicgen_base_32khz\")\n musicgen_base.bind_({'autocast': False, 'fsdp.use': True})\n\n # base musicgen models\n musicgen_base_small = musicgen_base.bind({'continue_from': '//pretrained/facebook/musicgen-small'})\n eval(musicgen_base_small, batch_size=128)\n\n musicgen_base_medium = musicgen_base.bind({'continue_from': '//pretrained/facebook/musicgen-medium'})\n musicgen_base_medium.bind_({'model/lm/model_scale': 'medium'})\n eval(musicgen_base_medium, batch_size=128)\n\n musicgen_base_large = musicgen_base.bind({'continue_from': '//pretrained/facebook/musicgen-large'})\n musicgen_base_large.bind_({'model/lm/model_scale': 'large'})\n eval(musicgen_base_large, batch_size=128)\n\n # melody musicgen model\n musicgen_melody = launcher.bind(solver=\"musicgen/musicgen_melody_32khz\")\n musicgen_melody.bind_({'autocast': False, 'fsdp.use': True})\n\n musicgen_melody_medium = musicgen_melody.bind({'continue_from': '//pretrained/facebook/musicgen-melody'})\n musicgen_melody_medium.bind_({'model/lm/model_scale': 'medium'})\n eval(musicgen_melody_medium, batch_size=128, eval_melody=True)\n" }, { "path": "audiocraft/grids/musicgen/musicgen_stereo_finetune_32khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom pathlib import Path\nfrom ._explorers import LMExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@LMExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=32, partition=partitions)\n launcher.bind_(solver='musicgen/musicgen_base_32khz')\n # replace this by the desired music dataset, which needs to be stereo\n launcher.bind_(dset='audio/example')\n\n fsdp = {'autocast': False, 'fsdp.use': True}\n medium = {'model/lm/model_scale': 'medium'}\n large = {'model/lm/model_scale': 'large'}\n\n cfg_low = {'classifier_free_guidance.training_dropout': 0.2}\n wd_low = {'conditioners.description.t5.word_dropout': 0.2}\n\n adam = {'optim.optimizer': 'adamw', 'optim.lr': 1e-4}\n\n stereo = {\n 'codebooks_pattern.delay.delays': [0, 0, 1, 1, 2, 2, 3, 3],\n 'transformer_lm.n_q': 8,\n 'interleave_stereo_codebooks.use': True,\n 'channels': 2,\n }\n\n # You must follow the instructions in docs/MUSICGEN.md about the creation\n # of the proper fine tuning checkpoints. We will assume they are stored under\n # ~/checkpoints/{mode_name}.\n\n checkpoints = Path.home() / 'checkpoints'\n\n launcher.bind_(fsdp, stereo, {'optim.epochs': 100})\n\n launcher.slurm_(gpus=32).bind_(label='32gpus')\n with launcher.job_array():\n sub = launcher.bind({'continue_from': str(checkpoints / 'stereo_finetune_musicgen-small.th')})\n sub()\n\n launcher.slurm_(gpus=64).bind_(label='64gpus')\n with launcher.job_array():\n sub = launcher.bind({'continue_from': str(checkpoints / 'stereo_finetune_musicgen-medium.th')})\n sub(medium, adam)\n\n launcher.slurm_(gpus=96).bind_(label='96gpus')\n with launcher.job_array():\n sub = launcher.bind({'continue_from': str(checkpoints / 'stereo_finetune_musicgen-large.th')})\n sub(large, cfg_low, wd_low, adam, {'optim.max_norm': 3})\n" }, { "path": "audiocraft/grids/musicgen/musicgen_style_32khz.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\nfrom ._explorers import LMExplorer\nfrom ...environment import AudioCraftEnvironment\n\n\n@LMExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(gpus=64, partition=partitions, constraint='volta32gb').bind_(label='64gpus')\n launcher.bind_(dset='internal/music_400k_32khz')\n\n sub = launcher.bind_({'solver': 'musicgen/musicgen_style_32khz',\n 'autocast': False,\n 'fsdp.use': True,\n 'model/lm/model_scale': 'medium',\n 'optim.optimizer': 'adamw',\n 'optim.lr': 1e-4,\n 'generate.every': 25,\n 'dataset.generate.num_samples': 64,\n })\n sub()\n" }, { "path": "audiocraft/grids/watermarking/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"watermarking grids.\"\"\"\n" }, { "path": "audiocraft/grids/watermarking/_explorers.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport treetable as tt\n\nfrom .._base_explorers import BaseExplorer\n\n\nclass WatermarkingMbExplorer(BaseExplorer):\n eval_metrics = [\"acc\", \"bit_acc\", \"visqol\", \"fnr\", \"fpr\", \"sisnr\"]\n\n def stages(self):\n return [\"train\", \"valid\", \"valid_ema\", \"evaluate\", \"evaluate_ema\"]\n\n def get_grid_meta(self):\n \"\"\"Returns the list of Meta information to display for each XP/job.\"\"\"\n return [\n tt.leaf(\"index\", align=\">\"),\n tt.leaf(\"name\", wrap=140),\n tt.leaf(\"state\"),\n tt.leaf(\"sig\", align=\">\"),\n ]\n\n def get_grid_metrics(self):\n \"\"\"Return the metrics that should be displayed in the tracking table.\"\"\"\n return [\n tt.group(\n \"train\",\n [\n tt.leaf(\"epoch\"),\n tt.leaf(\"sisnr\", \".3%\"),\n tt.leaf(\"wm_detection_identity\", \".3%\"),\n tt.leaf(\"wm_mb_identity\", \".3%\"),\n ],\n align=\">\",\n ),\n tt.group(\n \"valid\",\n [\n tt.leaf(\"sisnr\", \".3%\"),\n tt.leaf(\"wm_detection_identity\", \".3%\"),\n tt.leaf(\"wm_mb_identity\", \".3%\"),\n # tt.leaf(\"loss_0\", \".3%\"),\n ],\n align=\">\",\n ),\n tt.group(\n \"evaluate\",\n [\n tt.leaf(\"aug_identity_acc\", \".4f\"),\n tt.leaf(\"aug_identity_fnr\", \".4f\"),\n tt.leaf(\"aug_identity_fpr\", \".4f\"),\n tt.leaf(\"aug_identity_bit_acc\", \".4f\"),\n tt.leaf(\"pesq\", \".4f\"),\n tt.leaf(\"all_aug_acc\", \".4f\"),\n tt.leaf(\"localization_acc_padding\", \".4f\"),\n ],\n align=\">\",\n ),\n ]\n\n\nclass WatermarkingExplorer(BaseExplorer):\n eval_metrics = [\"acc\", \"visqol\", \"fnr\", \"fpr\", \"sisnr\"]\n\n def stages(self):\n return [\"train\", \"valid\", \"valid_ema\", \"evaluate\", \"evaluate_ema\"]\n\n def get_grid_meta(self):\n \"\"\"Returns the list of Meta information to display for each XP/job.\"\"\"\n return [\n tt.leaf(\"index\", align=\">\"),\n tt.leaf(\"name\", wrap=140),\n tt.leaf(\"state\"),\n tt.leaf(\"sig\", align=\">\"),\n ]\n\n def get_grid_metrics(self):\n \"\"\"Return the metrics that should be displayed in the tracking table.\"\"\"\n return [\n tt.group(\n \"train\",\n [\n tt.leaf(\"epoch\"),\n tt.leaf(\"sisnr\", \".3f\"),\n tt.leaf(\"wm_detection_identity\"),\n ],\n align=\">\",\n ),\n tt.group(\n \"valid\",\n [\n tt.leaf(\"sisnr\", \".3f\"),\n tt.leaf(\"wm_detection_identity\"),\n # tt.leaf(\"loss_0\", \".3%\"),\n ],\n align=\">\",\n ),\n tt.group(\n \"evaluate\",\n [\n tt.leaf(\"aug_identity_acc\", \".4f\"),\n tt.leaf(\"aug_identity_fnr\", \".4f\"),\n tt.leaf(\"aug_identity_fpr\", \".4f\"),\n tt.leaf(\"pesq\", \".4f\"),\n tt.leaf(\"all_aug_acc\", \".4f\"),\n tt.leaf(\"localization_acc_padding\", \".4f\"),\n\n ],\n align=\">\",\n ),\n ]\n" }, { "path": "audiocraft/grids/watermarking/audioseal.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n#\n\"\"\"\ndora grid watermarking.audioseal --clear\n\"\"\"\nfrom audiocraft.environment import AudioCraftEnvironment\nfrom ._explorers import WatermarkingExplorer\n\n\n@WatermarkingExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(\n gpus=8,\n partition=partitions,\n constraint=\"volta32gb\",\n )\n launcher.bind_(\n {\n \"solver\": \"watermark/robustness\",\n \"dset\": \"audio/example\",\n }\n )\n launcher.bind_(label=\"audioseal\")\n\n with launcher.job_array():\n launcher()\n" }, { "path": "audiocraft/grids/watermarking/kbits.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\ndora grid watermarking.kbits --clear\n\"\"\"\nimport os\nfrom audiocraft.environment import AudioCraftEnvironment\nfrom ._explorers import WatermarkingMbExplorer\n\n\n@WatermarkingMbExplorer\ndef explorer(launcher):\n partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global'])\n launcher.slurm_(\n gpus=8,\n partition=partitions,\n constraint=\"volta32gb\",\n )\n launcher.bind_(\n {\n \"solver\": \"watermark/robustness\",\n \"dset\": os.getenv(\"AUDIOCRAFT_DSET\", \"audio/example\"),\n \"dataset.batch_size\": 16,\n # optim\n \"optim.epochs\": 300,\n \"schedule\": {\n \"lr_scheduler\": \"cosine\",\n \"cosine\": {\n \"warmup\": 4000,\n \"lr_min_ratio\": 0.0,\n \"cycle_length\": 1.0,\n },\n },\n # crop and padding\n \"crop\": {\n \"prob\": 0.4,\n \"shuffle_prob\": 0.2,\n \"pad_prob\": 0.2,\n \"size\": 0.5,\n \"max_n_windows\": 5,\n },\n # augmentations\n \"select_aug_mode\": 'use_eval',\n \"aug_weights.updownresample\": 0.1,\n \"aug_weights.speed\": 0.1,\n \"aug_weights.echo\": 0.1,\n \"aug_weights.pink_noise\": 0.1,\n \"aug_weights.lowpass_filter\": 0.1,\n \"aug_weights.highpass_filter\": 0.1,\n \"aug_weights.bandpass_filter\": 0.1,\n \"aug_weights.smooth\": 0.1,\n \"aug_weights.boost_audio\": 0.1,\n \"aug_weights.duck_audio\": 0.1,\n \"aug_weights.mp3_compression\": 0.1,\n \"aug_weights.encodec\": 0.1,\n \"aug_weights.identity\": 1.0,\n # multi-bit\n \"audioseal.nbits\": 16,\n \"detector.output_dim\": 32,\n \"wm_mb.loss_type\": \"bce\",\n \"wm_mb.temperature\": 0.1,\n # losses\n \"losses\": { # encodec loss + tf = 10\n \"adv\": 4.0,\n \"feat\": 4.0,\n \"l1\": 0.1,\n \"mel\": 0.0,\n \"msspec\": 2.0,\n \"sisnr\": 0.0,\n \"tf_loudnessratio\": 10.0,\n },\n \"losses.wm_detection\": 1.0,\n \"losses.wm_mb\": 1.0,\n }\n )\n launcher.bind_(label=\"kbits16\")\n\n lrs = [5e-5]\n seeds = [1, 2, 3, 4]\n\n with launcher.job_array():\n for lr in lrs:\n for seed in seeds:\n launcher({\n \"optim.lr\": lr,\n \"seed\": seed,\n })\n" }, { "path": "audiocraft/losses/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Loss related classes and functions. In particular the loss balancer from\nEnCodec, and the usual spectral losses.\"\"\"\n\n# flake8: noqa\nfrom .balancer import Balancer\nfrom .sisnr import SISNR\nfrom .stftloss import (\n LogSTFTMagnitudeLoss,\n MRSTFTLoss,\n SpectralConvergenceLoss,\n STFTLoss\n)\nfrom .specloss import (\n MelSpectrogramL1Loss,\n MultiScaleMelSpectrogramLoss,\n)\n\nfrom .wmloss import (\n WMDetectionLoss,\n WMMbLoss\n)\n\nfrom .loudnessloss import TFLoudnessRatio\n" }, { "path": "audiocraft/losses/balancer.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nimport flashy\nimport torch\nfrom torch import autograd\n\n\nclass Balancer:\n \"\"\"Loss balancer.\n\n The loss balancer combines losses together to compute gradients for the backward.\n Given `y = f(...)`, and a number of losses `l1(y, ...)`, `l2(y, ...)`, with `...`\n not having any dependence on `f`, the balancer can efficiently normalize the partial gradients\n `d l1 / d y`, `d l2 / dy` before summing them in order to achieve a desired ratio between\n the losses. For instance if `weights = {'l1': 2, 'l2': 1}`, 66% of the gradient\n going into `f(...)` will come from `l1` on average, and 33% from `l2`. This allows for an easy\n interpration of the weights even if the intrisic scale of `l1`, `l2` ... is unknown.\n\n Noting `g1 = d l1 / dy`, etc., the balanced gradient `G` will be\n (with `avg` an exponential moving average over the updates),\n\n G = sum_i total_norm * g_i / avg(||g_i||) * w_i / sum(w_i)\n\n If `balance_grads` is False, this is deactivated, and instead the gradient will just be the\n standard sum of the partial gradients with the given weights.\n\n A call to the backward method of the balancer will compute the the partial gradients,\n combining all the losses and potentially rescaling the gradients,\n which can help stabilize the training and reason about multiple losses with varying scales.\n The obtained gradient with respect to `y` is then back-propagated to `f(...)`.\n\n Expected usage:\n\n weights = {'loss_a': 1, 'loss_b': 4}\n balancer = Balancer(weights, ...)\n losses: dict = {}\n losses['loss_a'] = compute_loss_a(x, y)\n losses['loss_b'] = compute_loss_b(x, y)\n if model.training():\n effective_loss = balancer.backward(losses, x)\n\n Args:\n weights (dict[str, float]): Weight coefficient for each loss. The balancer expect the losses keys\n from the backward method to match the weights keys to assign weight to each of the provided loss.\n balance_grads (bool): Whether to rescale gradients so that weights reflect the fraction of the\n overall gradient, rather than a constant multiplier.\n total_norm (float): Reference norm when rescaling gradients, ignored otherwise.\n emay_decay (float): EMA decay for averaging the norms.\n per_batch_item (bool): Whether to compute the averaged norm per batch item or not. This only holds\n when rescaling the gradients.\n epsilon (float): Epsilon value for numerical stability.\n monitor (bool): If True, stores in `self.metrics` the relative ratio between the norm of the gradients\n coming from each loss, when calling `backward()`.\n \"\"\"\n def __init__(self, weights: tp.Dict[str, float], balance_grads: bool = True, total_norm: float = 1.,\n ema_decay: float = 0.999, per_batch_item: bool = True, epsilon: float = 1e-12,\n monitor: bool = False):\n self.weights = weights\n self.per_batch_item = per_batch_item\n self.total_norm = total_norm or 1.\n self.averager = flashy.averager(ema_decay or 1.)\n self.epsilon = epsilon\n self.monitor = monitor\n self.balance_grads = balance_grads\n self._metrics: tp.Dict[str, tp.Any] = {}\n\n @property\n def metrics(self):\n return self._metrics\n\n def backward(self, losses: tp.Dict[str, torch.Tensor], input: torch.Tensor) -> torch.Tensor:\n \"\"\"Compute the backward and return the effective train loss, e.g. the loss obtained from\n computing the effective weights. If `balance_grads` is True, the effective weights\n are the one that needs to be applied to each gradient to respect the desired relative\n scale of gradients coming from each loss.\n\n Args:\n losses (Dict[str, torch.Tensor]): dictionary with the same keys as `self.weights`.\n input (torch.Tensor): the input of the losses, typically the output of the model.\n This should be the single point of dependence between the losses\n and the model being trained.\n \"\"\"\n norms = {}\n grads = {}\n for name, loss in losses.items():\n # Compute partial derivative of the less with respect to the input.\n grad, = autograd.grad(loss, [input], retain_graph=True)\n if self.per_batch_item:\n # We do not average the gradient over the batch dimension.\n dims = tuple(range(1, grad.dim()))\n norm = grad.norm(dim=dims, p=2).mean()\n else:\n norm = grad.norm(p=2)\n norms[name] = norm\n grads[name] = grad\n\n count = 1\n if self.per_batch_item:\n count = len(grad)\n # Average norms across workers. Theoretically we should average the\n # squared norm, then take the sqrt, but it worked fine like that.\n avg_norms = flashy.distrib.average_metrics(self.averager(norms), count)\n # We approximate the total norm of the gradient as the sums of the norms.\n # Obviously this can be very incorrect if all gradients are aligned, but it works fine.\n total = sum(avg_norms.values())\n\n self._metrics = {}\n if self.monitor:\n # Store the ratio of the total gradient represented by each loss.\n for k, v in avg_norms.items():\n self._metrics[f'ratio_{k}'] = v / total\n\n total_weights = sum([self.weights[k] for k in avg_norms])\n assert total_weights > 0.\n desired_ratios = {k: w / total_weights for k, w in self.weights.items()}\n\n out_grad = torch.zeros_like(input)\n effective_loss = torch.tensor(0., device=input.device, dtype=input.dtype)\n for name, avg_norm in avg_norms.items():\n if self.balance_grads:\n # g_balanced = g / avg(||g||) * total_norm * desired_ratio\n scale = desired_ratios[name] * self.total_norm / (self.epsilon + avg_norm)\n else:\n # We just do regular weighted sum of the gradients.\n scale = self.weights[name]\n out_grad.add_(grads[name], alpha=scale)\n effective_loss += scale * losses[name].detach()\n # Send the computed partial derivative with respect to the output of the model to the model.\n input.backward(out_grad)\n return effective_loss\n" }, { "path": "audiocraft/losses/loudnessloss.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nimport typing as tp\n\nimport julius\nimport torch\nimport torchaudio\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torchaudio.functional.filtering import highpass_biquad, treble_biquad\n\n\ndef basic_loudness(waveform: torch.Tensor, sample_rate: int) -> torch.Tensor:\n \"\"\"This is a simpler loudness function that is more stable.\n Args:\n waveform(torch.Tensor): audio waveform of dimension `(..., channels, time)`\n sample_rate (int): sampling rate of the waveform\n Returns:\n loudness loss as a scalar\n \"\"\"\n\n if waveform.size(-2) > 5:\n raise ValueError(\"Only up to 5 channels are supported.\")\n eps = torch.finfo(torch.float32).eps\n gate_duration = 0.4\n overlap = 0.75\n gate_samples = int(round(gate_duration * sample_rate))\n step = int(round(gate_samples * (1 - overlap)))\n\n # Apply K-weighting\n waveform = treble_biquad(waveform, sample_rate, 4.0, 1500.0, 1 / math.sqrt(2))\n waveform = highpass_biquad(waveform, sample_rate, 38.0, 0.5)\n\n # Compute the energy for each block\n energy = torch.square(waveform).unfold(-1, gate_samples, step)\n energy = torch.mean(energy, dim=-1)\n\n # Compute channel-weighted summation\n g = torch.tensor([1.0, 1.0, 1.0, 1.41, 1.41], dtype=waveform.dtype, device=waveform.device)\n g = g[: energy.size(-2)]\n\n energy_weighted = torch.sum(g.unsqueeze(-1) * energy, dim=-2)\n # loudness with epsilon for stability. Not as much precision in the very low loudness sections\n loudness = -0.691 + 10 * torch.log10(energy_weighted + eps)\n return loudness\n\n\ndef _unfold(a: torch.Tensor, kernel_size: int, stride: int) -> torch.Tensor:\n \"\"\"Given input of size [*OT, T], output Tensor of size [*OT, F, K]\n with K the kernel size, by extracting frames with the given stride.\n This will pad the input so that `F = ceil(T / K)`.\n see https://github.com/pytorch/pytorch/issues/60466\n \"\"\"\n *shape, length = a.shape\n n_frames = math.ceil(length / stride)\n tgt_length = (n_frames - 1) * stride + kernel_size\n a = F.pad(a, (0, tgt_length - length))\n strides = list(a.stride())\n assert strides[-1] == 1, \"data should be contiguous\"\n strides = strides[:-1] + [stride, 1]\n return a.as_strided([*shape, n_frames, kernel_size], strides)\n\n\nclass FLoudnessRatio(nn.Module):\n \"\"\"FSNR loss.\n\n Input should be [B, C, T], output is scalar.\n\n Args:\n sample_rate (int): Sample rate.\n segment (float or None): Evaluate on chunks of that many seconds. If None, evaluate on\n entire audio only.\n overlap (float): Overlap between chunks, i.e. 0.5 = 50 % overlap.\n epsilon (float): Epsilon value for numerical stability.\n n_bands (int): number of mel scale bands that we include\n \"\"\"\n def __init__(\n self,\n sample_rate: int = 16000,\n segment: tp.Optional[float] = 20,\n overlap: float = 0.5,\n epsilon: float = torch.finfo(torch.float32).eps,\n n_bands: int = 0,\n ):\n super().__init__()\n self.sample_rate = sample_rate\n self.segment = segment\n self.overlap = overlap\n self.epsilon = epsilon\n if n_bands == 0:\n self.filter = None\n else:\n self.filter = julius.SplitBands(sample_rate=sample_rate, n_bands=n_bands)\n self.loudness = torchaudio.transforms.Loudness(sample_rate)\n\n def forward(self, out_sig: torch.Tensor, ref_sig: torch.Tensor) -> torch.Tensor:\n B, C, T = ref_sig.shape\n assert ref_sig.shape == out_sig.shape\n assert self.filter is not None\n bands_ref = self.filter(ref_sig)\n bands_out = self.filter(out_sig)\n l_noise = self.loudness(bands_ref - bands_out)\n l_ref = self.loudness(bands_ref)\n l_ratio = (l_noise - l_ref).view(-1, B)\n loss = torch.nn.functional.softmax(l_ratio, dim=0) * l_ratio\n return loss.sum()\n\n\nclass TLoudnessRatio(nn.Module):\n \"\"\"TSNR loss.\n\n Input should be [B, C, T], output is scalar.\n\n Args:\n sample_rate (int): Sample rate.\n segment (float or None): Evaluate on chunks of that many seconds. If None, evaluate on\n entire audio only.\n overlap (float): Overlap between chunks, i.e. 0.5 = 50 % overlap.\n \"\"\"\n def __init__(\n self,\n sample_rate: int = 16000,\n segment: float = 0.5,\n overlap: float = 0.5,\n ):\n super().__init__()\n self.sample_rate = sample_rate\n self.segment = segment\n self.overlap = overlap\n self.loudness = torchaudio.transforms.Loudness(sample_rate)\n\n def forward(self, out_sig: torch.Tensor, ref_sig: torch.Tensor) -> torch.Tensor:\n B, C, T = ref_sig.shape\n assert ref_sig.shape == out_sig.shape\n assert C == 1\n\n frame = int(self.segment * self.sample_rate)\n stride = int(frame * (1 - self.overlap))\n gt = _unfold(ref_sig, frame, stride).view(-1, 1, frame)\n est = _unfold(out_sig, frame, stride).view(-1, 1, frame)\n l_noise = self.loudness(gt - est) # watermark\n l_ref = self.loudness(gt) # ground truth\n l_ratio = (l_noise - l_ref).view(-1, B)\n loss = torch.nn.functional.softmax(l_ratio, dim=0) * l_ratio\n return loss.sum()\n\n\nclass TFLoudnessRatio(nn.Module):\n \"\"\"TF-loudness ratio loss.\n\n Input should be [B, C, T], output is scalar.\n\n Args:\n sample_rate (int): Sample rate.\n segment (float or None): Evaluate on chunks of that many seconds. If None, evaluate on\n entire audio only.\n overlap (float): Overlap between chunks, i.e. 0.5 = 50 % overlap.\n n_bands (int): number of bands to separate\n temperature (float): temperature of the softmax step\n \"\"\"\n def __init__(\n self,\n sample_rate: int = 16000,\n segment: float = 0.5,\n overlap: float = 0.5,\n n_bands: int = 0,\n clip_min: float = -100,\n temperature: float = 1.0,\n ):\n super().__init__()\n self.sample_rate = sample_rate\n self.segment = segment\n self.overlap = overlap\n self.clip_min = clip_min\n self.temperature = temperature\n if n_bands == 0:\n self.filter = None\n else:\n self.n_bands = n_bands\n self.filter = julius.SplitBands(sample_rate=sample_rate, n_bands=n_bands)\n\n def forward(self, out_sig: torch.Tensor, ref_sig: torch.Tensor) -> torch.Tensor:\n B, C, T = ref_sig.shape\n\n assert ref_sig.shape == out_sig.shape\n assert C == 1\n assert self.filter is not None\n\n bands_ref = self.filter(ref_sig).view(B * self.n_bands, 1, -1)\n bands_out = self.filter(out_sig).view(B * self.n_bands, 1, -1)\n frame = int(self.segment * self.sample_rate)\n stride = int(frame * (1 - self.overlap))\n gt = _unfold(bands_ref, frame, stride).squeeze(1).contiguous().view(-1, 1, frame)\n est = _unfold(bands_out, frame, stride).squeeze(1).contiguous().view(-1, 1, frame)\n l_noise = basic_loudness(est - gt, sample_rate=self.sample_rate) # watermark\n l_ref = basic_loudness(gt, sample_rate=self.sample_rate) # ground truth\n l_ratio = (l_noise - l_ref).view(-1, B)\n loss = torch.nn.functional.softmax(l_ratio / self.temperature, dim=0) * l_ratio\n return loss.mean()\n" }, { "path": "audiocraft/losses/sisnr.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nimport typing as tp\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n\ndef _unfold(a: torch.Tensor, kernel_size: int, stride: int) -> torch.Tensor:\n \"\"\"Given input of size [*OT, T], output Tensor of size [*OT, F, K]\n with K the kernel size, by extracting frames with the given stride.\n This will pad the input so that `F = ceil(T / K)`.\n see https://github.com/pytorch/pytorch/issues/60466\n \"\"\"\n *shape, length = a.shape\n n_frames = math.ceil(length / stride)\n tgt_length = (n_frames - 1) * stride + kernel_size\n a = F.pad(a, (0, tgt_length - length))\n strides = list(a.stride())\n assert strides[-1] == 1, \"data should be contiguous\"\n strides = strides[:-1] + [stride, 1]\n return a.as_strided([*shape, n_frames, kernel_size], strides)\n\n\ndef _center(x: torch.Tensor) -> torch.Tensor:\n return x - x.mean(-1, True)\n\n\ndef _norm2(x: torch.Tensor) -> torch.Tensor:\n return x.pow(2).sum(-1, True)\n\n\nclass SISNR(nn.Module):\n \"\"\"SISNR loss.\n\n Input should be [B, C, T], output is scalar.\n\n ..Warning:: This function returns the opposite of the SI-SNR (e.g. `-1 * regular_SI_SNR`).\n Consequently, lower scores are better in terms of reconstruction quality,\n in particular, it should be negative if training goes well. This done this way so\n that this module can also be used as a loss function for training model.\n\n Args:\n sample_rate (int): Sample rate.\n segment (float or None): Evaluate on chunks of that many seconds. If None, evaluate on\n entire audio only.\n overlap (float): Overlap between chunks, i.e. 0.5 = 50 % overlap.\n epsilon (float): Epsilon value for numerical stability.\n \"\"\"\n def __init__(\n self,\n sample_rate: int = 16000,\n segment: tp.Optional[float] = 20,\n overlap: float = 0.5,\n epsilon: float = torch.finfo(torch.float32).eps,\n ):\n super().__init__()\n self.sample_rate = sample_rate\n self.segment = segment\n self.overlap = overlap\n self.epsilon = epsilon\n\n def forward(self, out_sig: torch.Tensor, ref_sig: torch.Tensor) -> torch.Tensor:\n B, C, T = ref_sig.shape\n assert ref_sig.shape == out_sig.shape\n\n if self.segment is None:\n frame = T\n stride = T\n else:\n frame = int(self.segment * self.sample_rate)\n stride = int(frame * (1 - self.overlap))\n\n epsilon = self.epsilon * frame # make epsilon prop to frame size.\n\n gt = _unfold(ref_sig, frame, stride)\n est = _unfold(out_sig, frame, stride)\n if self.segment is None:\n assert gt.shape[-1] == 1\n\n gt = _center(gt)\n est = _center(est)\n dot = torch.einsum(\"bcft,bcft->bcf\", gt, est)\n\n proj = dot[:, :, :, None] * gt / (epsilon + _norm2(gt))\n noise = est - proj\n\n sisnr = 10 * (\n torch.log10(epsilon + _norm2(proj)) - torch.log10(epsilon + _norm2(noise))\n )\n return -1 * sisnr[..., 0].mean()\n" }, { "path": "audiocraft/losses/specloss.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nimport numpy as np\nfrom torchaudio.transforms import MelSpectrogram\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom ..modules import pad_for_conv1d\n\n\nclass MelSpectrogramWrapper(nn.Module):\n \"\"\"Wrapper around MelSpectrogram torchaudio transform providing proper padding\n and additional post-processing including log scaling.\n\n Args:\n n_mels (int): Number of mel bins.\n n_fft (int): Number of fft.\n hop_length (int): Hop size.\n win_length (int): Window length.\n n_mels (int): Number of mel bins.\n sample_rate (int): Sample rate.\n f_min (float or None): Minimum frequency.\n f_max (float or None): Maximum frequency.\n log (bool): Whether to scale with log.\n normalized (bool): Whether to normalize the melspectrogram.\n floor_level (float): Floor level based on human perception (default=1e-5).\n \"\"\"\n def __init__(self, n_fft: int = 1024, hop_length: int = 256, win_length: tp.Optional[int] = None,\n n_mels: int = 80, sample_rate: float = 22050, f_min: float = 0.0, f_max: tp.Optional[float] = None,\n log: bool = True, normalized: bool = False, floor_level: float = 1e-5):\n super().__init__()\n self.n_fft = n_fft\n hop_length = int(hop_length)\n self.hop_length = hop_length\n self.mel_transform = MelSpectrogram(n_mels=n_mels, sample_rate=sample_rate, n_fft=n_fft, hop_length=hop_length,\n win_length=win_length, f_min=f_min, f_max=f_max, normalized=normalized,\n window_fn=torch.hann_window, center=False)\n self.floor_level = floor_level\n self.log = log\n\n def forward(self, x):\n p = int((self.n_fft - self.hop_length) // 2)\n if len(x.shape) == 2:\n x = x.unsqueeze(1)\n x = F.pad(x, (p, p), \"reflect\")\n # Make sure that all the frames are full.\n # The combination of `pad_for_conv1d` and the above padding\n # will make the output of size ceil(T / hop).\n x = pad_for_conv1d(x, self.n_fft, self.hop_length)\n self.mel_transform.to(x.device)\n mel_spec = self.mel_transform(x)\n B, C, freqs, frame = mel_spec.shape\n if self.log:\n mel_spec = torch.log10(self.floor_level + mel_spec)\n return mel_spec.reshape(B, C * freqs, frame)\n\n\nclass MelSpectrogramL1Loss(torch.nn.Module):\n \"\"\"L1 Loss on MelSpectrogram.\n\n Args:\n sample_rate (int): Sample rate.\n n_fft (int): Number of fft.\n hop_length (int): Hop size.\n win_length (int): Window length.\n n_mels (int): Number of mel bins.\n f_min (float or None): Minimum frequency.\n f_max (float or None): Maximum frequency.\n log (bool): Whether to scale with log.\n normalized (bool): Whether to normalize the melspectrogram.\n floor_level (float): Floor level value based on human perception (default=1e-5).\n \"\"\"\n def __init__(self, sample_rate: int, n_fft: int = 1024, hop_length: int = 256, win_length: int = 1024,\n n_mels: int = 80, f_min: float = 0.0, f_max: tp.Optional[float] = None,\n log: bool = True, normalized: bool = False, floor_level: float = 1e-5):\n super().__init__()\n self.l1 = torch.nn.L1Loss()\n self.melspec = MelSpectrogramWrapper(n_fft=n_fft, hop_length=hop_length, win_length=win_length,\n n_mels=n_mels, sample_rate=sample_rate, f_min=f_min, f_max=f_max,\n log=log, normalized=normalized, floor_level=floor_level)\n\n def forward(self, x, y):\n self.melspec.to(x.device)\n s_x = self.melspec(x)\n s_y = self.melspec(y)\n return self.l1(s_x, s_y)\n\n\nclass MultiScaleMelSpectrogramLoss(nn.Module):\n \"\"\"Multi-Scale spectrogram loss (msspec).\n\n Args:\n sample_rate (int): Sample rate.\n range_start (int): Power of 2 to use for the first scale.\n range_stop (int): Power of 2 to use for the last scale.\n n_mels (int): Number of mel bins.\n f_min (float): Minimum frequency.\n f_max (float or None): Maximum frequency.\n normalized (bool): Whether to normalize the melspectrogram.\n alphas (bool): Whether to use alphas as coefficients or not.\n floor_level (float): Floor level value based on human perception (default=1e-5).\n \"\"\"\n def __init__(self, sample_rate: int, range_start: int = 6, range_end: int = 11,\n n_mels: int = 64, f_min: float = 0.0, f_max: tp.Optional[float] = None,\n normalized: bool = False, alphas: bool = True, floor_level: float = 1e-5):\n super().__init__()\n l1s = list()\n l2s = list()\n self.alphas = list()\n self.total = 0\n self.normalized = normalized\n for i in range(range_start, range_end):\n l1s.append(\n MelSpectrogramWrapper(n_fft=2 ** i, hop_length=(2 ** i) / 4, win_length=2 ** i,\n n_mels=n_mels, sample_rate=sample_rate, f_min=f_min, f_max=f_max,\n log=False, normalized=normalized, floor_level=floor_level))\n l2s.append(\n MelSpectrogramWrapper(n_fft=2 ** i, hop_length=(2 ** i) / 4, win_length=2 ** i,\n n_mels=n_mels, sample_rate=sample_rate, f_min=f_min, f_max=f_max,\n log=True, normalized=normalized, floor_level=floor_level))\n if alphas:\n self.alphas.append(np.sqrt(2 ** i - 1))\n else:\n self.alphas.append(1)\n self.total += self.alphas[-1] + 1\n\n self.l1s = nn.ModuleList(l1s)\n self.l2s = nn.ModuleList(l2s)\n\n def forward(self, x, y):\n loss = 0.0\n self.l1s.to(x.device)\n self.l2s.to(x.device)\n for i in range(len(self.alphas)):\n s_x_1 = self.l1s[i](x)\n s_y_1 = self.l1s[i](y)\n s_x_2 = self.l2s[i](x)\n s_y_2 = self.l2s[i](y)\n loss += F.l1_loss(s_x_1, s_y_1) + self.alphas[i] * F.mse_loss(s_x_2, s_y_2)\n if self.normalized:\n loss = loss / self.total\n return loss\n" }, { "path": "audiocraft/losses/stftloss.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n# Adapted from MIT code under the original license\n# Copyright 2019 Tomoki Hayashi\n# MIT License (https://opensource.org/licenses/MIT)\nimport typing as tp\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n\n# TODO: Replace with torchaudio.STFT?\ndef _stft(x: torch.Tensor, fft_size: int, hop_length: int, win_length: int,\n window: tp.Optional[torch.Tensor], normalized: bool) -> torch.Tensor:\n \"\"\"Perform STFT and convert to magnitude spectrogram.\n\n Args:\n x: Input signal tensor (B, C, T).\n fft_size (int): FFT size.\n hop_length (int): Hop size.\n win_length (int): Window length.\n window (torch.Tensor or None): Window function type.\n normalized (bool): Whether to normalize the STFT or not.\n\n Returns:\n torch.Tensor: Magnitude spectrogram (B, C, #frames, fft_size // 2 + 1).\n \"\"\"\n B, C, T = x.shape\n x_stft = torch.stft(\n x.view(-1, T), fft_size, hop_length, win_length, window,\n normalized=normalized, return_complex=True,\n )\n x_stft = x_stft.view(B, C, *x_stft.shape[1:])\n real = x_stft.real\n imag = x_stft.imag\n\n # NOTE(kan-bayashi): clamp is needed to avoid nan or inf\n return torch.sqrt(torch.clamp(real ** 2 + imag ** 2, min=1e-7)).transpose(2, 1)\n\n\nclass SpectralConvergenceLoss(nn.Module):\n \"\"\"Spectral convergence loss.\n \"\"\"\n def __init__(self, epsilon: float = torch.finfo(torch.float32).eps):\n super().__init__()\n self.epsilon = epsilon\n\n def forward(self, x_mag: torch.Tensor, y_mag: torch.Tensor):\n \"\"\"Calculate forward propagation.\n\n Args:\n x_mag: Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).\n y_mag: Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).\n Returns:\n torch.Tensor: Spectral convergence loss value.\n \"\"\"\n return torch.norm(y_mag - x_mag, p=\"fro\") / (torch.norm(y_mag, p=\"fro\") + self.epsilon)\n\n\nclass LogSTFTMagnitudeLoss(nn.Module):\n \"\"\"Log STFT magnitude loss.\n\n Args:\n epsilon (float): Epsilon value for numerical stability.\n \"\"\"\n def __init__(self, epsilon: float = torch.finfo(torch.float32).eps):\n super().__init__()\n self.epsilon = epsilon\n\n def forward(self, x_mag: torch.Tensor, y_mag: torch.Tensor):\n \"\"\"Calculate forward propagation.\n\n Args:\n x_mag (torch.Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).\n y_mag (torch.Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).\n Returns:\n torch.Tensor: Log STFT magnitude loss value.\n \"\"\"\n return F.l1_loss(torch.log(self.epsilon + y_mag), torch.log(self.epsilon + x_mag))\n\n\nclass STFTLosses(nn.Module):\n \"\"\"STFT losses.\n\n Args:\n n_fft (int): Size of FFT.\n hop_length (int): Hop length.\n win_length (int): Window length.\n window (str): Window function type.\n normalized (bool): Whether to use normalized STFT or not.\n epsilon (float): Epsilon for numerical stability.\n \"\"\"\n def __init__(self, n_fft: int = 1024, hop_length: int = 120, win_length: int = 600,\n window: str = \"hann_window\", normalized: bool = False,\n epsilon: float = torch.finfo(torch.float32).eps):\n super().__init__()\n self.n_fft = n_fft\n self.hop_length = hop_length\n self.win_length = win_length\n self.normalized = normalized\n self.register_buffer(\"window\", getattr(torch, window)(win_length))\n self.spectral_convergenge_loss = SpectralConvergenceLoss(epsilon)\n self.log_stft_magnitude_loss = LogSTFTMagnitudeLoss(epsilon)\n\n def forward(self, x: torch.Tensor, y: torch.Tensor) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Calculate forward propagation.\n\n Args:\n x (torch.Tensor): Predicted signal (B, T).\n y (torch.Tensor): Groundtruth signal (B, T).\n Returns:\n torch.Tensor: Spectral convergence loss value.\n torch.Tensor: Log STFT magnitude loss value.\n \"\"\"\n x_mag = _stft(x, self.n_fft, self.hop_length,\n self.win_length, self.window, self.normalized) # type: ignore\n y_mag = _stft(y, self.n_fft, self.hop_length,\n self.win_length, self.window, self.normalized) # type: ignore\n sc_loss = self.spectral_convergenge_loss(x_mag, y_mag)\n mag_loss = self.log_stft_magnitude_loss(x_mag, y_mag)\n\n return sc_loss, mag_loss\n\n\nclass STFTLoss(nn.Module):\n \"\"\"Single Resolution STFT loss.\n\n Args:\n n_fft (int): Nb of FFT.\n hop_length (int): Hop length.\n win_length (int): Window length.\n window (str): Window function type.\n normalized (bool): Whether to use normalized STFT or not.\n epsilon (float): Epsilon for numerical stability.\n factor_sc (float): Coefficient for the spectral loss.\n factor_mag (float): Coefficient for the magnitude loss.\n \"\"\"\n def __init__(self, n_fft: int = 1024, hop_length: int = 120, win_length: int = 600,\n window: str = \"hann_window\", normalized: bool = False,\n factor_sc: float = 0.1, factor_mag: float = 0.1,\n epsilon: float = torch.finfo(torch.float32).eps):\n super().__init__()\n self.loss = STFTLosses(n_fft, hop_length, win_length, window, normalized, epsilon)\n self.factor_sc = factor_sc\n self.factor_mag = factor_mag\n\n def forward(self, x: torch.Tensor, y: torch.Tensor) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Calculate forward propagation.\n\n Args:\n x (torch.Tensor): Predicted signal (B, T).\n y (torch.Tensor): Groundtruth signal (B, T).\n Returns:\n torch.Tensor: Single resolution STFT loss.\n \"\"\"\n sc_loss, mag_loss = self.loss(x, y)\n return self.factor_sc * sc_loss + self.factor_mag * mag_loss\n\n\nclass MRSTFTLoss(nn.Module):\n \"\"\"Multi resolution STFT loss.\n\n Args:\n n_ffts (Sequence[int]): Sequence of FFT sizes.\n hop_lengths (Sequence[int]): Sequence of hop sizes.\n win_lengths (Sequence[int]): Sequence of window lengths.\n window (str): Window function type.\n factor_sc (float): Coefficient for the spectral loss.\n factor_mag (float): Coefficient for the magnitude loss.\n normalized (bool): Whether to use normalized STFT or not.\n epsilon (float): Epsilon for numerical stability.\n \"\"\"\n def __init__(self, n_ffts: tp.Sequence[int] = [1024, 2048, 512], hop_lengths: tp.Sequence[int] = [120, 240, 50],\n win_lengths: tp.Sequence[int] = [600, 1200, 240], window: str = \"hann_window\",\n factor_sc: float = 0.1, factor_mag: float = 0.1,\n normalized: bool = False, epsilon: float = torch.finfo(torch.float32).eps):\n super().__init__()\n assert len(n_ffts) == len(hop_lengths) == len(win_lengths)\n self.stft_losses = torch.nn.ModuleList()\n for fs, ss, wl in zip(n_ffts, hop_lengths, win_lengths):\n self.stft_losses += [STFTLosses(fs, ss, wl, window, normalized, epsilon)]\n self.factor_sc = factor_sc\n self.factor_mag = factor_mag\n\n def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:\n \"\"\"Calculate forward propagation.\n\n Args:\n x (torch.Tensor): Predicted signal (B, T).\n y (torch.Tensor): Groundtruth signal (B, T).\n Returns:\n torch.Tensor: Multi resolution STFT loss.\n \"\"\"\n sc_loss = torch.Tensor([0.0])\n mag_loss = torch.Tensor([0.0])\n for f in self.stft_losses:\n sc_l, mag_l = f(x, y)\n sc_loss += sc_l\n mag_loss += mag_l\n sc_loss /= len(self.stft_losses)\n mag_loss /= len(self.stft_losses)\n\n return self.factor_sc * sc_loss + self.factor_mag * mag_loss\n" }, { "path": "audiocraft/losses/wmloss.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Literal\n\nimport torch\nimport torch.nn as nn\n\n\nclass WMDetectionLoss(nn.Module):\n \"\"\"Compute the detection loss\"\"\"\n def __init__(self, p_weight: float = 1.0, n_weight: float = 1.0) -> None:\n super().__init__()\n self.criterion = nn.NLLLoss()\n self.p_weight = p_weight\n self.n_weight = n_weight\n\n def forward(self, positive, negative, mask, message=None):\n\n positive = positive[:, :2, :] # b 2+nbits t -> b 2 t\n negative = negative[:, :2, :] # b 2+nbits t -> b 2 t\n\n # dimensionality of positive [bsz, classes=2, time_steps]\n # correct classes for pos = [bsz, time_steps] where all values = 1 for positive\n classes_shape = positive[\n :, 0, :\n ] # same as positive or negative but dropping dim=1\n pos_correct_classes = torch.ones_like(classes_shape, dtype=int)\n neg_correct_classes = torch.zeros_like(classes_shape, dtype=int)\n\n # taking log because network outputs softmax\n # NLLLoss expects a logsoftmax input\n positive = torch.log(positive)\n negative = torch.log(negative)\n\n if not torch.all(mask == 1):\n # pos_correct_classes [bsz, timesteps] mask [bsz, 1, timesptes]\n # mask is applied to the watermark, this basically flips the tgt class from 1 (positive)\n # to 0 (negative) in the correct places\n pos_correct_classes = pos_correct_classes * mask[:, 0, :].to(int)\n loss_p = self.p_weight * self.criterion(positive, pos_correct_classes)\n # no need for negative class loss here since some of the watermark\n # is masked to negative\n return loss_p\n\n else:\n loss_p = self.p_weight * self.criterion(positive, pos_correct_classes)\n loss_n = self.n_weight * self.criterion(negative, neg_correct_classes)\n return loss_p + loss_n\n\n\nclass WMMbLoss(nn.Module):\n def __init__(self, temperature: float, loss_type: Literal[\"bce\", \"mse\"]) -> None:\n \"\"\"\n Compute the masked sample-level detection loss\n (https://arxiv.org/pdf/2401.17264)\n\n Args:\n temperature: temperature for loss computation\n loss_type: bce or mse between outputs and original message\n \"\"\"\n super().__init__()\n self.bce_with_logits = (\n nn.BCEWithLogitsLoss()\n ) # same as Softmax + NLLLoss, but when only 1 output unit\n self.mse = nn.MSELoss()\n self.loss_type = loss_type\n self.temperature = temperature\n\n def forward(self, positive, negative, mask, message):\n \"\"\"\n Compute decoding loss\n Args:\n positive: outputs on watermarked samples [bsz, 2+nbits, time_steps]\n negative: outputs on not watermarked samples [bsz, 2+nbits, time_steps]\n mask: watermark mask [bsz, 1, time_steps]\n message: original message [bsz, nbits] or None\n \"\"\"\n # # no use of negative at the moment\n # negative = negative[:, 2:, :] # b 2+nbits t -> b nbits t\n # negative = torch.masked_select(negative, mask)\n if message.size(0) == 0:\n return torch.tensor(0.0)\n positive = positive[:, 2:, :] # b 2+nbits t -> b nbits t\n assert (\n positive.shape[-2] == message.shape[1]\n ), \"in decoding loss: \\\n enc and dec don't share nbits, are you using multi-bit?\"\n\n # cut last dim of positive to keep only where mask is 1\n new_shape = [*positive.shape[:-1], -1] # b nbits -1\n positive = torch.masked_select(positive, mask == 1).reshape(new_shape)\n\n message = message.unsqueeze(-1).repeat(1, 1, positive.shape[2]) # b k -> b k t\n if self.loss_type == \"bce\":\n # in this case similar to temperature in softmax\n loss = self.bce_with_logits(positive / self.temperature, message.float())\n elif self.loss_type == \"mse\":\n loss = self.mse(positive / self.temperature, message.float())\n\n return loss\n" }, { "path": "audiocraft/metrics/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Metrics like CLAP score, FAD, KLD, Visqol, Chroma similarity, etc.\n\"\"\"\n# flake8: noqa\nfrom .clap_consistency import CLAPTextConsistencyMetric, TextConsistencyMetric\nfrom .chroma_cosinesim import ChromaCosineSimilarityMetric\nfrom .fad import FrechetAudioDistanceMetric\nfrom .kld import KLDivergenceMetric, PasstKLDivergenceMetric\nfrom .rvm import RelativeVolumeMel\nfrom .visqol import ViSQOL\n" }, { "path": "audiocraft/metrics/chroma_cosinesim.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\nimport torchmetrics\n\nfrom ..data.audio_utils import convert_audio\nfrom ..modules.chroma import ChromaExtractor\n\n\nclass ChromaCosineSimilarityMetric(torchmetrics.Metric):\n \"\"\"Chroma cosine similarity metric.\n\n This metric extracts a chromagram for a reference waveform and\n a generated waveform and compares each frame using the cosine similarity\n function. The output is the mean cosine similarity.\n\n Args:\n sample_rate (int): Sample rate used by the chroma extractor.\n n_chroma (int): Number of chroma used by the chroma extractor.\n radix2_exp (int): Exponent for the chroma extractor.\n argmax (bool): Whether the chroma extractor uses argmax.\n eps (float): Epsilon for cosine similarity computation.\n \"\"\"\n def __init__(self, sample_rate: int, n_chroma: int, radix2_exp: int, argmax: bool, eps: float = 1e-8):\n super().__init__()\n self.chroma_sample_rate = sample_rate\n self.n_chroma = n_chroma\n self.eps = eps\n self.chroma_extractor = ChromaExtractor(sample_rate=self.chroma_sample_rate, n_chroma=self.n_chroma,\n radix2_exp=radix2_exp, argmax=argmax)\n self.add_state(\"cosine_sum\", default=torch.tensor(0.), dist_reduce_fx=\"sum\")\n self.add_state(\"weight\", default=torch.tensor(0.), dist_reduce_fx=\"sum\")\n\n def update(self, preds: torch.Tensor, targets: torch.Tensor,\n sizes: torch.Tensor, sample_rates: torch.Tensor) -> None:\n \"\"\"Compute cosine similarity between chromagrams and accumulate scores over the dataset.\"\"\"\n if preds.size(0) == 0:\n return\n\n assert preds.shape == targets.shape, (\n f\"Preds and target shapes mismatch: preds={preds.shape}, targets={targets.shape}\")\n assert preds.size(0) == sizes.size(0), (\n f\"Number of items in preds ({preds.shape}) mismatch \",\n f\"with sizes ({sizes.shape})\")\n assert preds.size(0) == sample_rates.size(0), (\n f\"Number of items in preds ({preds.shape}) mismatch \",\n f\"with sample_rates ({sample_rates.shape})\")\n assert torch.all(sample_rates == sample_rates[0].item()), \"All sample rates are not the same in the batch\"\n\n device = self.weight.device\n preds, targets = preds.to(device), targets.to(device) # type: ignore\n sample_rate = sample_rates[0].item()\n preds = convert_audio(preds, from_rate=sample_rate, to_rate=self.chroma_sample_rate, to_channels=1)\n targets = convert_audio(targets, from_rate=sample_rate, to_rate=self.chroma_sample_rate, to_channels=1)\n gt_chroma = self.chroma_extractor(targets)\n gen_chroma = self.chroma_extractor(preds)\n chroma_lens = (sizes / self.chroma_extractor.winhop).ceil().int()\n for i in range(len(gt_chroma)):\n t = int(chroma_lens[i].item())\n cosine_sim = torch.nn.functional.cosine_similarity(\n gt_chroma[i, :t], gen_chroma[i, :t], dim=1, eps=self.eps)\n self.cosine_sum += cosine_sim.sum(dim=0) # type: ignore\n self.weight += torch.tensor(t) # type: ignore\n\n def compute(self) -> float:\n \"\"\"Computes the average cosine similarty across all generated/target chromagrams pairs.\"\"\"\n assert self.weight.item() > 0, \"Unable to compute with total number of comparisons <= 0\" # type: ignore\n return (self.cosine_sum / self.weight).item() # type: ignore\n" }, { "path": "audiocraft/metrics/clap_consistency.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom pathlib import Path\nimport typing as tp\n\nimport torch\nimport torchmetrics\nfrom transformers import RobertaTokenizer # type: ignore\n\nfrom ..data.audio_utils import convert_audio\nfrom ..environment import AudioCraftEnvironment\nfrom ..utils.utils import load_clap_state_dict\n\ntry:\n import laion_clap # type: ignore\nexcept ImportError:\n laion_clap = None\n\n\nclass TextConsistencyMetric(torchmetrics.Metric):\n \"\"\"Text consistency metric measuring consistency between audio and text pairs.\"\"\"\n\n def update(self, audio: torch.Tensor, text: tp.List[str], sizes: torch.Tensor, sample_rates: torch.Tensor) -> None:\n raise NotImplementedError(\"implement how to update the metric from the audio and text pairs.\")\n\n def compute(self):\n raise NotImplementedError(\"implement how to compute the final metric score.\")\n\n\nclass CLAPTextConsistencyMetric(TextConsistencyMetric):\n \"\"\"Text consistency metric relying on Contrastive Language-Audio Pretraining (CLAP).\n\n This metric is similar to the MuLan Cycle Consistency from MusicLM (https://arxiv.org/pdf/2301.11325.pdf)\n or the CLAP score used in Make-An-Audio (https://arxiv.org/pdf/2301.12661v1.pdf).\n\n As a joint audio-text embedding model, a pretrained CLAP model can be used to quantify the\n similarity between audio-text pairs. We compute the CLAP embeddings from the text descriptions as\n well as the generated audio based on them, and define the MCC metric as the average cosine similarity\n between these embeddings.\n\n Model implementation & pre-trained checkpoints: https://github.com/LAION-AI/CLAP\n \"\"\"\n def __init__(self, model_path: tp.Union[str, Path], model_arch: str = 'HTSAT-tiny', enable_fusion: bool = False):\n super().__init__()\n if laion_clap is None:\n raise ImportError(\"Please install CLAP to compute text consistency: 'pip install laion_clap'\")\n self.add_state(\"cosine_sum\", default=torch.tensor(0.), dist_reduce_fx=\"sum\")\n self.add_state(\"weight\", default=torch.tensor(0.), dist_reduce_fx=\"sum\")\n self._initialize_model(model_path, model_arch, enable_fusion)\n\n def _initialize_model(self, model_path: tp.Union[str, Path], model_arch: str, enable_fusion: bool):\n model_path = AudioCraftEnvironment.resolve_reference_path(model_path)\n self.tokenize = RobertaTokenizer.from_pretrained('roberta-base')\n self.model = laion_clap.CLAP_Module(enable_fusion=enable_fusion, amodel=model_arch)\n self.model_sample_rate = 48_000\n load_clap_state_dict(self.model, model_path)\n self.model.eval()\n\n def _tokenizer(self, texts: tp.Union[str, tp.List[str]]) -> dict:\n # we use the default params from CLAP module here as well\n return self.tokenize(texts, padding=\"max_length\", truncation=True, max_length=77, return_tensors=\"pt\")\n\n def update(self, audio: torch.Tensor, text: tp.List[str], sizes: torch.Tensor, sample_rates: torch.Tensor) -> None:\n \"\"\"Compute cosine similarity between audio and text pairs and accumulate scores over the dataset.\"\"\"\n assert audio.size(0) == len(text), \"Number of audio and text samples should match\"\n assert torch.all(sample_rates == sample_rates[0].item()), \"All items in batch should have the same sample rate\"\n sample_rate = int(sample_rates[0].item())\n # convert audio batch to 48kHz monophonic audio with no channel dimension: [B, C, T] -> [B, T]\n audio = convert_audio(audio, from_rate=sample_rate, to_rate=self.model_sample_rate, to_channels=1).mean(dim=1)\n audio_embeddings = self.model.get_audio_embedding_from_data(audio, use_tensor=True)\n text_embeddings = self.model.get_text_embedding(text, tokenizer=self._tokenizer, use_tensor=True)\n # cosine similarity between the text and the audio embedding\n cosine_sim = torch.nn.functional.cosine_similarity(audio_embeddings, text_embeddings, dim=1, eps=1e-8)\n self.cosine_sum += cosine_sim.sum(dim=0)\n self.weight += torch.tensor(cosine_sim.size(0))\n\n def compute(self):\n \"\"\"Computes the average cosine similarty across all audio/text pairs.\"\"\"\n assert self.weight.item() > 0, \"Unable to compute with total number of comparisons <= 0\" # type: ignore\n return (self.cosine_sum / self.weight).item() # type: ignore\n" }, { "path": "audiocraft/metrics/fad.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nfrom pathlib import Path\nimport os\nimport subprocess\nimport tempfile\nimport typing as tp\n\nfrom audiocraft.data.audio import audio_write\nfrom audiocraft.data.audio_utils import convert_audio\nimport flashy\nimport torch\nimport torchmetrics\n\nfrom ..environment import AudioCraftEnvironment\n\n\nlogger = logging.getLogger(__name__)\n\nVGGISH_SAMPLE_RATE = 16_000\nVGGISH_CHANNELS = 1\n\n\nclass FrechetAudioDistanceMetric(torchmetrics.Metric):\n \"\"\"Fréchet Audio Distance computation based on official TensorFlow implementation from Google Research.\n\n From: D.C. Dowson & B.V. Landau The Fréchet distance between\n multivariate normal distributions\n https://doi.org/10.1016/0047-259X(82)90077-X\n The Fréchet distance between two multivariate gaussians,\n `X ~ N(mu_x, sigma_x)` and `Y ~ N(mu_y, sigma_y)`, is `d^2`.\n d^2 = (mu_x - mu_y)^2 + Tr(sigma_x + sigma_y - 2 * sqrt(sigma_x*sigma_y))\n = (mu_x - mu_y)^2 + Tr(sigma_x) + Tr(sigma_y)\n - 2 * Tr(sqrt(sigma_x*sigma_y)))\n\n To use this FAD computation metric, you need to have the proper Frechet Audio Distance tool setup\n from: https://github.com/google-research/google-research/tree/master/frechet_audio_distance\n We provide the below instructions as reference but we do not guarantee for further support\n in frechet_audio_distance installation. This was tested with python 3.10, cuda 11.8, tensorflow 2.12.0.\n\n We recommend installing the frechet_audio_distance library in a dedicated env (e.g. conda).\n\n 1. Get the code and models following the repository instructions. We used the steps below:\n git clone git@github.com:google-research/google-research.git\n git clone git@github.com:tensorflow/models.git\n mkdir google-research/tensorflow_models\n touch google-research/tensorflow_models/__init__.py\n cp -r models/research/audioset google-research/tensorflow_models/\n touch google-research/tensorflow_models/audioset/__init__.py\n echo \"from .vggish import mel_features, vggish_params, vggish_slim\" > \\\n google-research/tensorflow_models/audioset/__init__.py\n # we can now remove the tensorflow models repository\n # rm -r models\n cd google-research\n Follow the instructions to download the vggish checkpoint. AudioCraft base configuration\n assumes it is placed in the AudioCraft reference dir.\n\n Note that we operate the following changes for the code to work with TensorFlow 2.X and python 3:\n - Update xrange for range in:\n https://github.com/google-research/google-research/blob/master/frechet_audio_distance/audioset_model.py\n - Update `tf_record = tf.python_io.tf_record_iterator(filename).next()` to\n `tf_record = tf.python_io.tf_record_iterator(filename).__next__()` in\n https://github.com/google-research/google-research/blob/master/frechet_audio_distance/fad_utils.py\n - Update `import vggish_params as params` to `from . import vggish_params as params` in:\n https://github.com/tensorflow/models/blob/master/research/audioset/vggish/vggish_slim.py\n - Add flag to provide a given batch size for running the AudioSet model in:\n https://github.com/google-research/google-research/blob/master/frechet_audio_distance/create_embeddings_main.py\n ```\n flags.DEFINE_integer('batch_size', 64,\n 'Number of samples in the batch for AudioSet model.')\n ```\n Ensure you pass the flag to the create_embeddings_beam.create_pipeline function, adding:\n `batch_size=FLAGS.batch_size` to the provided parameters.\n\n 2. Follow instructions for the library installation and a valid TensorFlow installation\n ```\n # e.g. instructions from: https://www.tensorflow.org/install/pip\n conda install -c conda-forge cudatoolkit=11.8.0\n python3 -m pip install nvidia-cudnn-cu11==8.6.0.163 tensorflow==2.12.*\n mkdir -p $CONDA_PREFIX/etc/conda/activate.d\n echo 'CUDNN_PATH=$(dirname $(python -c \"import nvidia.cudnn;print(nvidia.cudnn.__file__)\"))' \\\n >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh\n echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib' \\\n >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh\n source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh\n # Verify install: on a machine with GPU device\n python3 -c \"import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))\"\n ```\n\n Now install frechet_audio_distance required dependencies:\n ```\n # We assume we already have TensorFlow installed from the above steps\n pip install apache-beam numpy scipy tf_slim\n ```\n\n Finally, follow remaining library instructions to ensure you have a working frechet_audio_distance setup\n (you may want to specify --model_ckpt flag pointing to the model's path).\n\n 3. AudioCraft's FrechetAudioDistanceMetric requires 2 environment variables pointing to the python executable\n and Tensorflow library path from the above installation steps:\n export TF_PYTHON_EXE=\"\"\n export TF_LIBRARY_PATH=\"\"\n\n e.g. assuming we have installed everything in a dedicated conda env\n with python 3.10 that is currently active:\n export TF_PYTHON_EXE=\"$CONDA_PREFIX/bin/python\"\n export TF_LIBRARY_PATH=\"$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/cudnn/lib\"\n\n Finally you may want to export the following variable:\n export TF_FORCE_GPU_ALLOW_GROWTH=true\n See: https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth\n\n You can save those environment variables in your training conda env, when currently active:\n `$CONDA_PREFIX/etc/conda/activate.d/env_vars.sh`\n e.g. assuming the env with TensorFlow and frechet_audio_distance install is named ac_eval,\n and the training conda env is named audiocraft:\n ```\n # activate training env\n conda activate audiocraft\n # get path to all envs\n CONDA_ENV_DIR=$(dirname $CONDA_PREFIX)\n # export pointers to evaluation env for using TensorFlow in FrechetAudioDistanceMetric\n touch $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh\n echo 'export TF_PYTHON_EXE=\"$CONDA_ENV_DIR/ac_eval/bin/python\"' >> \\\n $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh\n echo 'export TF_LIBRARY_PATH=\"$CONDA_ENV_DIR/ac_eval/lib/python3.10/site-packages/nvidia/cudnn/lib\"' >> \\\n $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh\n # optionally:\n echo 'export TF_FORCE_GPU_ALLOW_GROWTH=true' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh\n # you may need to reactivate the audiocraft env for this to take effect\n ```\n\n Args:\n bin (Path or str): Path to installed frechet audio distance code.\n model_path (Path or str): Path to Tensorflow checkpoint for the model\n used to compute statistics over the embedding beams.\n format (str): Audio format used to save files.\n log_folder (Path or str, optional): Path where to write process logs.\n \"\"\"\n def __init__(self, bin: tp.Union[Path, str], model_path: tp.Union[Path, str],\n format: str = \"wav\", batch_size: tp.Optional[int] = None,\n log_folder: tp.Optional[tp.Union[Path, str]] = None):\n super().__init__()\n self.model_sample_rate = VGGISH_SAMPLE_RATE\n self.model_channels = VGGISH_CHANNELS\n self.model_path = AudioCraftEnvironment.resolve_reference_path(model_path)\n assert Path(self.model_path).exists(), f\"Could not find provided model checkpoint path at: {self.model_path}\"\n self.format = format\n self.batch_size = batch_size\n self.bin = bin\n self.tf_env = {\"PYTHONPATH\": str(self.bin)}\n self.python_path = os.environ.get('TF_PYTHON_EXE') or 'python'\n logger.info(\"Python exe for TF is %s\", self.python_path)\n if 'TF_LIBRARY_PATH' in os.environ:\n self.tf_env['LD_LIBRARY_PATH'] = os.environ['TF_LIBRARY_PATH']\n if 'TF_FORCE_GPU_ALLOW_GROWTH' in os.environ:\n self.tf_env['TF_FORCE_GPU_ALLOW_GROWTH'] = os.environ['TF_FORCE_GPU_ALLOW_GROWTH']\n logger.info(\"Env for TF is %r\", self.tf_env)\n self.reset(log_folder)\n self.add_state(\"total_files\", default=torch.tensor(0.), dist_reduce_fx=\"sum\")\n\n def reset(self, log_folder: tp.Optional[tp.Union[Path, str]] = None):\n \"\"\"Reset torchmetrics.Metrics state.\"\"\"\n log_folder = Path(log_folder or tempfile.mkdtemp())\n self.tmp_dir = log_folder / 'fad'\n self.tmp_dir.mkdir(exist_ok=True)\n self.samples_tests_dir = self.tmp_dir / 'tests'\n self.samples_tests_dir.mkdir(exist_ok=True)\n self.samples_background_dir = self.tmp_dir / 'background'\n self.samples_background_dir.mkdir(exist_ok=True)\n self.manifest_tests = self.tmp_dir / 'files_tests.cvs'\n self.manifest_background = self.tmp_dir / 'files_background.cvs'\n self.stats_tests_dir = self.tmp_dir / 'stats_tests'\n self.stats_background_dir = self.tmp_dir / 'stats_background'\n self.counter = 0\n\n def update(self, preds: torch.Tensor, targets: torch.Tensor,\n sizes: torch.Tensor, sample_rates: torch.Tensor,\n stems: tp.Optional[tp.List[str]] = None):\n \"\"\"Update torchmetrics.Metrics by saving the audio and updating the manifest file.\"\"\"\n assert preds.shape == targets.shape, f\"preds={preds.shape} != targets={targets.shape}\"\n num_samples = preds.shape[0]\n assert num_samples == sizes.size(0) and num_samples == sample_rates.size(0)\n assert stems is None or num_samples == len(set(stems))\n for i in range(num_samples):\n self.total_files += 1 # type: ignore\n self.counter += 1\n wav_len = int(sizes[i].item())\n sample_rate = int(sample_rates[i].item())\n pred_wav = preds[i]\n target_wav = targets[i]\n pred_wav = pred_wav[..., :wav_len]\n target_wav = target_wav[..., :wav_len]\n stem_name = stems[i] if stems is not None else f'sample_{self.counter}_{flashy.distrib.rank()}'\n # dump audio files\n try:\n pred_wav = convert_audio(\n pred_wav.unsqueeze(0), from_rate=sample_rate,\n to_rate=self.model_sample_rate, to_channels=1).squeeze(0)\n audio_write(\n self.samples_tests_dir / stem_name, pred_wav, sample_rate=self.model_sample_rate,\n format=self.format, strategy=\"peak\")\n except Exception as e:\n logger.error(f\"Exception occured when saving tests files for FAD computation: {repr(e)} - {e}\")\n try:\n # for the ground truth audio, we enforce the 'peak' strategy to avoid modifying\n # the original audio when writing it\n target_wav = convert_audio(\n target_wav.unsqueeze(0), from_rate=sample_rate,\n to_rate=self.model_sample_rate, to_channels=1).squeeze(0)\n audio_write(\n self.samples_background_dir / stem_name, target_wav, sample_rate=self.model_sample_rate,\n format=self.format, strategy=\"peak\")\n except Exception as e:\n logger.error(f\"Exception occured when saving background files for FAD computation: {repr(e)} - {e}\")\n\n def _get_samples_name(self, is_background: bool):\n return 'background' if is_background else 'tests'\n\n def _create_embedding_beams(self, is_background: bool, gpu_index: tp.Optional[int] = None):\n if is_background:\n input_samples_dir = self.samples_background_dir\n input_filename = self.manifest_background\n stats_name = self.stats_background_dir\n else:\n input_samples_dir = self.samples_tests_dir\n input_filename = self.manifest_tests\n stats_name = self.stats_tests_dir\n beams_name = self._get_samples_name(is_background)\n log_file = self.tmp_dir / f'fad_logs_create_beams_{beams_name}.log'\n\n logger.info(f\"Scanning samples folder to fetch list of files: {input_samples_dir}\")\n with open(input_filename, \"w\") as fout:\n for path in Path(input_samples_dir).glob(f\"*.{self.format}\"):\n fout.write(f\"{str(path)}\\n\")\n\n cmd = [\n self.python_path, \"-m\",\n \"frechet_audio_distance.create_embeddings_main\",\n \"--model_ckpt\", f\"{self.model_path}\",\n \"--input_files\", f\"{str(input_filename)}\",\n \"--stats\", f\"{str(stats_name)}\",\n ]\n if self.batch_size is not None:\n cmd += [\"--batch_size\", str(self.batch_size)]\n logger.info(f\"Launching frechet_audio_distance embeddings main method: {' '.join(cmd)} on {beams_name}\")\n env = os.environ\n if gpu_index is not None:\n env[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_index)\n process = subprocess.Popen(\n cmd, stdout=open(log_file, \"w\"), env={**env, **self.tf_env}, stderr=subprocess.STDOUT)\n return process, log_file\n\n def _compute_fad_score(self, gpu_index: tp.Optional[int] = None):\n cmd = [\n self.python_path, \"-m\", \"frechet_audio_distance.compute_fad\",\n \"--test_stats\", f\"{str(self.stats_tests_dir)}\",\n \"--background_stats\", f\"{str(self.stats_background_dir)}\",\n ]\n logger.info(f\"Launching frechet_audio_distance compute fad method: {' '.join(cmd)}\")\n env = os.environ\n if gpu_index is not None:\n env[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_index)\n result = subprocess.run(cmd, env={**env, **self.tf_env}, capture_output=True)\n if result.returncode:\n logger.error(\n \"Error with FAD computation from stats: \\n %s \\n %s\",\n result.stdout.decode(), result.stderr.decode()\n )\n raise RuntimeError(\"Error while executing FAD computation from stats\")\n try:\n # result is \"FAD: (d+).(d+)\" hence we remove the prefix with (d+) being one digit or more\n fad_score = float(result.stdout[4:])\n return fad_score\n except Exception as e:\n raise RuntimeError(f\"Error parsing FAD score from command stdout: {e}\")\n\n def _log_process_result(self, returncode: int, log_file: tp.Union[Path, str], is_background: bool) -> None:\n beams_name = self._get_samples_name(is_background)\n if returncode:\n with open(log_file, \"r\") as f:\n error_log = f.read()\n logger.error(error_log)\n os._exit(1)\n else:\n logger.info(f\"Successfully computed embedding beams on {beams_name} samples.\")\n\n def _parallel_create_embedding_beams(self, num_of_gpus: int):\n assert num_of_gpus > 0\n logger.info(\"Creating embeddings beams in a parallel manner on different GPUs\")\n tests_beams_process, tests_beams_log_file = self._create_embedding_beams(is_background=False, gpu_index=0)\n bg_beams_process, bg_beams_log_file = self._create_embedding_beams(is_background=True, gpu_index=1)\n tests_beams_code = tests_beams_process.wait()\n bg_beams_code = bg_beams_process.wait()\n self._log_process_result(tests_beams_code, tests_beams_log_file, is_background=False)\n self._log_process_result(bg_beams_code, bg_beams_log_file, is_background=True)\n\n def _sequential_create_embedding_beams(self):\n logger.info(\"Creating embeddings beams in a sequential manner\")\n tests_beams_process, tests_beams_log_file = self._create_embedding_beams(is_background=False)\n tests_beams_code = tests_beams_process.wait()\n self._log_process_result(tests_beams_code, tests_beams_log_file, is_background=False)\n bg_beams_process, bg_beams_log_file = self._create_embedding_beams(is_background=True)\n bg_beams_code = bg_beams_process.wait()\n self._log_process_result(bg_beams_code, bg_beams_log_file, is_background=True)\n\n @flashy.distrib.rank_zero_only\n def _local_compute_frechet_audio_distance(self):\n \"\"\"Compute Frechet Audio Distance score calling TensorFlow API.\"\"\"\n num_of_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0\n if num_of_gpus > 1:\n self._parallel_create_embedding_beams(num_of_gpus)\n else:\n self._sequential_create_embedding_beams()\n fad_score = self._compute_fad_score(gpu_index=0)\n return fad_score\n\n def compute(self) -> float:\n \"\"\"Compute metrics.\"\"\"\n assert self.total_files.item() > 0, \"No files dumped for FAD computation!\" # type: ignore\n fad_score = self._local_compute_frechet_audio_distance()\n logger.warning(f\"FAD score = {fad_score}\")\n fad_score = flashy.distrib.broadcast_object(fad_score, src=0)\n return fad_score\n" }, { "path": "audiocraft/metrics/kld.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport contextlib\nfrom functools import partial\nimport logging\nimport os\nimport typing as tp\n\nimport torch\nimport torchmetrics\n\nfrom ..data.audio_utils import convert_audio\n\n\nlogger = logging.getLogger(__name__)\n\n\nclass _patch_passt_stft:\n \"\"\"Decorator to patch torch.stft in PaSST.\"\"\"\n def __init__(self):\n self.old_stft = torch.stft\n\n def __enter__(self):\n # return_complex is a mandatory parameter in latest torch versions\n # torch is throwing RuntimeErrors when not set\n torch.stft = partial(torch.stft, return_complex=False)\n\n def __exit__(self, *exc):\n torch.stft = self.old_stft\n\n\ndef kl_divergence(pred_probs: torch.Tensor, target_probs: torch.Tensor, epsilon: float = 1e-6) -> torch.Tensor:\n \"\"\"Computes the elementwise KL-Divergence loss between probability distributions\n from generated samples and target samples.\n\n Args:\n pred_probs (torch.Tensor): Probabilities for each label obtained\n from a classifier on generated audio. Expected shape is [B, num_classes].\n target_probs (torch.Tensor): Probabilities for each label obtained\n from a classifier on target audio. Expected shape is [B, num_classes].\n epsilon (float): Epsilon value.\n Returns:\n kld (torch.Tensor): KLD loss between each generated sample and target pair.\n \"\"\"\n kl_div = torch.nn.functional.kl_div((pred_probs + epsilon).log(), target_probs, reduction=\"none\")\n return kl_div.sum(-1)\n\n\nclass KLDivergenceMetric(torchmetrics.Metric):\n \"\"\"Base implementation for KL Divergence metric.\n\n The KL divergence is measured between probability distributions\n of class predictions returned by a pre-trained audio classification model.\n When the KL-divergence is low, the generated audio is expected to\n have similar acoustic characteristics as the reference audio,\n according to the classifier.\n \"\"\"\n def __init__(self):\n super().__init__()\n self.add_state(\"kld_pq_sum\", default=torch.tensor(0.), dist_reduce_fx=\"sum\")\n self.add_state(\"kld_qp_sum\", default=torch.tensor(0.), dist_reduce_fx=\"sum\")\n self.add_state(\"kld_all_sum\", default=torch.tensor(0.), dist_reduce_fx=\"sum\")\n self.add_state(\"weight\", default=torch.tensor(0), dist_reduce_fx=\"sum\")\n\n def _get_label_distribution(self, x: torch.Tensor, sizes: torch.Tensor,\n sample_rates: torch.Tensor) -> tp.Optional[torch.Tensor]:\n \"\"\"Get model output given provided input tensor.\n\n Args:\n x (torch.Tensor): Input audio tensor of shape [B, C, T].\n sizes (torch.Tensor): Actual audio sample length, of shape [B].\n sample_rates (torch.Tensor): Actual audio sample rate, of shape [B].\n Returns:\n probs (torch.Tensor): Probabilities over labels, of shape [B, num_classes].\n \"\"\"\n raise NotImplementedError(\"implement method to extract label distributions from the model.\")\n\n def update(self, preds: torch.Tensor, targets: torch.Tensor,\n sizes: torch.Tensor, sample_rates: torch.Tensor) -> None:\n \"\"\"Calculates running KL-Divergence loss between batches of audio\n preds (generated) and target (ground-truth)\n Args:\n preds (torch.Tensor): Audio samples to evaluate, of shape [B, C, T].\n targets (torch.Tensor): Target samples to compare against, of shape [B, C, T].\n sizes (torch.Tensor): Actual audio sample length, of shape [B].\n sample_rates (torch.Tensor): Actual audio sample rate, of shape [B].\n \"\"\"\n assert preds.shape == targets.shape\n assert preds.size(0) > 0, \"Cannot update the loss with empty tensors\"\n preds_probs = self._get_label_distribution(preds, sizes, sample_rates)\n targets_probs = self._get_label_distribution(targets, sizes, sample_rates)\n if preds_probs is not None and targets_probs is not None:\n assert preds_probs.shape == targets_probs.shape\n kld_scores = kl_divergence(preds_probs, targets_probs)\n assert not torch.isnan(kld_scores).any(), \"kld_scores contains NaN value(s)!\"\n self.kld_pq_sum += torch.sum(kld_scores)\n kld_qp_scores = kl_divergence(targets_probs, preds_probs)\n self.kld_qp_sum += torch.sum(kld_qp_scores)\n self.weight += torch.tensor(kld_scores.size(0))\n\n def compute(self) -> dict:\n \"\"\"Computes KL-Divergence across all evaluated pred/target pairs.\"\"\"\n weight: float = float(self.weight.item()) # type: ignore\n assert weight > 0, \"Unable to compute with total number of comparisons <= 0\"\n logger.info(f\"Computing KL divergence on a total of {weight} samples\")\n kld_pq = self.kld_pq_sum.item() / weight # type: ignore\n kld_qp = self.kld_qp_sum.item() / weight # type: ignore\n kld_both = kld_pq + kld_qp\n return {'kld': kld_pq, 'kld_pq': kld_pq, 'kld_qp': kld_qp, 'kld_both': kld_both}\n\n\nclass PasstKLDivergenceMetric(KLDivergenceMetric):\n \"\"\"KL-Divergence metric based on pre-trained PASST classifier on AudioSet.\n\n From: PaSST: Efficient Training of Audio Transformers with Patchout\n Paper: https://arxiv.org/abs/2110.05069\n Implementation: https://github.com/kkoutini/PaSST\n\n Follow instructions from the github repo:\n ```\n pip install 'git+https://github.com/kkoutini/passt_hear21@0.0.19#egg=hear21passt'\n ```\n\n Args:\n pretrained_length (float, optional): Audio duration used for the pretrained model.\n \"\"\"\n def __init__(self, pretrained_length: tp.Optional[float] = None):\n super().__init__()\n self._initialize_model(pretrained_length)\n\n def _initialize_model(self, pretrained_length: tp.Optional[float] = None):\n \"\"\"Initialize underlying PaSST audio classifier.\"\"\"\n model, sr, max_frames, min_frames = self._load_base_model(pretrained_length)\n self.min_input_frames = min_frames\n self.max_input_frames = max_frames\n self.model_sample_rate = sr\n self.model = model\n self.model.eval()\n self.model.to(self.device)\n\n def _load_base_model(self, pretrained_length: tp.Optional[float]):\n \"\"\"Load pretrained model from PaSST.\"\"\"\n try:\n if pretrained_length == 30:\n from hear21passt.base30sec import get_basic_model # type: ignore\n max_duration = 30\n elif pretrained_length == 20:\n from hear21passt.base20sec import get_basic_model # type: ignore\n max_duration = 20\n else:\n from hear21passt.base import get_basic_model # type: ignore\n # Original PASST was trained on AudioSet with 10s-long audio samples\n max_duration = 10\n min_duration = 0.15\n min_duration = 0.15\n except ModuleNotFoundError:\n raise ModuleNotFoundError(\n \"Please install hear21passt to compute KL divergence: \",\n \"pip install 'git+https://github.com/kkoutini/passt_hear21@0.0.19#egg=hear21passt'\"\n )\n model_sample_rate = 32_000\n max_input_frames = int(max_duration * model_sample_rate)\n min_input_frames = int(min_duration * model_sample_rate)\n with open(os.devnull, 'w') as f, contextlib.redirect_stdout(f):\n model = get_basic_model(mode='logits')\n return model, model_sample_rate, max_input_frames, min_input_frames\n\n def _process_audio(self, wav: torch.Tensor, sample_rate: int, wav_len: int) -> tp.List[torch.Tensor]:\n \"\"\"Process audio to feed to the pretrained model.\"\"\"\n wav = wav.unsqueeze(0)\n wav = wav[..., :wav_len]\n wav = convert_audio(wav, from_rate=sample_rate, to_rate=self.model_sample_rate, to_channels=1)\n wav = wav.squeeze(0)\n # we don't pad but return a list of audio segments as this otherwise affects the KLD computation\n segments = torch.split(wav, self.max_input_frames, dim=-1)\n valid_segments = []\n for s in segments:\n # ignoring too small segments that are breaking the model inference\n if s.size(-1) > self.min_input_frames:\n valid_segments.append(s)\n return [s[None] for s in valid_segments]\n\n def _get_model_preds(self, wav: torch.Tensor) -> torch.Tensor:\n \"\"\"Run the pretrained model and get the predictions.\"\"\"\n assert wav.dim() == 3, f\"Unexpected number of dims for preprocessed wav: {wav.shape}\"\n wav = wav.mean(dim=1)\n # PaSST is printing a lot of garbage that we are not interested in\n with open(os.devnull, \"w\") as f, contextlib.redirect_stdout(f):\n with torch.no_grad(), _patch_passt_stft():\n logits = self.model(wav.to(self.device))\n probs = torch.softmax(logits, dim=-1)\n return probs\n\n def _get_label_distribution(self, x: torch.Tensor, sizes: torch.Tensor,\n sample_rates: torch.Tensor) -> tp.Optional[torch.Tensor]:\n \"\"\"Get model output given provided input tensor.\n\n Args:\n x (torch.Tensor): Input audio tensor of shape [B, C, T].\n sizes (torch.Tensor): Actual audio sample length, of shape [B].\n sample_rates (torch.Tensor): Actual audio sample rate, of shape [B].\n Returns:\n probs (torch.Tensor, optional): Probabilities over labels, of shape [B, num_classes].\n \"\"\"\n all_probs: tp.List[torch.Tensor] = []\n for i, wav in enumerate(x):\n sample_rate = int(sample_rates[i].item())\n wav_len = int(sizes[i].item())\n wav_segments = self._process_audio(wav, sample_rate, wav_len)\n for segment in wav_segments:\n probs = self._get_model_preds(segment).mean(dim=0)\n all_probs.append(probs)\n if len(all_probs) > 0:\n return torch.stack(all_probs, dim=0)\n else:\n return None\n" }, { "path": "audiocraft/metrics/miou.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\n\n\ndef calculate_miou(y_pred: torch.Tensor, y_true: torch.Tensor) -> float:\n \"\"\"\n Calculate the mean Intersection over Union (mIoU) between two binary tensors using PyTorch.\n\n Args:\n y_pred (torch.Tensor): Predicted binary tensor of shape [bsz, frames].\n y_true (torch.Tensor): Ground truth binary tensor of shape [bsz, frames].\n\n Returns:\n float: The mean Intersection over Union (mIoU) score.\n\n Reference:\n The Intersection over Union (IoU) metric is commonly used in computer vision.\n For more information, refer to the following paper:\n \"SegNet: A Deep Convolutional Encoder-Decoder Architecture for Image Segmentation\"\n by Vijay Badrinarayanan, Alex Kendall, Roberto Cipolla\n \"\"\"\n # Ensure y_pred and y_true have the same shape\n if y_pred.shape != y_true.shape:\n raise ValueError(\"Input tensors must have the same shape\")\n\n # converting predictions to binary vector\n y_pred = y_pred > 0.5\n # Compute the intersection and union\n intersection = torch.logical_and(y_pred, y_true)\n union = torch.logical_or(y_pred, y_true)\n\n # Compute IoU for each sample in the batch\n iou_per_sample = torch.sum(intersection, dim=1) / torch.sum(union, dim=1)\n # Calculate mIoU by taking the mean across the batch\n miou = torch.mean(iou_per_sample).item()\n\n return miou\n" }, { "path": "audiocraft/metrics/pesq.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport julius\nimport pesq\n\nimport torch\nimport torchmetrics\n\n\nclass PesqMetric(torchmetrics.Metric):\n \"\"\"Metric for Perceptual Evaluation of Speech Quality.\n (https://doi.org/10.5281/zenodo.6549559)\n\n \"\"\"\n\n sum_pesq: torch.Tensor\n total: torch.Tensor\n\n def __init__(self, sample_rate: int):\n super().__init__()\n self.sr = sample_rate\n\n self.add_state(\"sum_pesq\", default=torch.tensor(0.0), dist_reduce_fx=\"sum\")\n self.add_state(\"total\", default=torch.tensor(0), dist_reduce_fx=\"sum\")\n\n def update(self, preds: torch.Tensor, targets: torch.Tensor):\n if self.sr != 16000:\n preds = julius.resample_frac(preds, self.sr, 16000)\n targets = julius.resample_frac(targets, self.sr, 16000)\n for ii in range(preds.size(0)):\n try:\n self.sum_pesq += pesq.pesq(\n 16000, targets[ii, 0].detach().cpu().numpy(), preds[ii, 0].detach().cpu().numpy()\n )\n self.total += 1\n except (\n pesq.NoUtterancesError\n ): # this error can append when the sample don't contain speech\n pass\n\n def compute(self) -> torch.Tensor:\n return (\n self.sum_pesq / self.total\n if (self.total != 0).item()\n else torch.tensor(0.0)\n )\n" }, { "path": "audiocraft/metrics/rvm.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\nimport torch\nfrom torch import nn\nimport torchaudio\n\n\ndef db_to_scale(volume: tp.Union[float, torch.Tensor]):\n return 10 ** (volume / 20)\n\n\ndef scale_to_db(scale: torch.Tensor, min_volume: float = -120):\n min_scale = db_to_scale(min_volume)\n return 20 * torch.log10(scale.clamp(min=min_scale))\n\n\nclass RelativeVolumeMel(nn.Module):\n \"\"\"Relative volume melspectrogram measure.\n\n Computes a measure of distance over two mel spectrogram that is interpretable in terms\n of decibels. Given `x_ref` and `x_est` two waveforms of shape `[*, T]`, it will\n first renormalize both by the ground truth of `x_ref`.\n\n ..Warning:: This class returns the volume of the distortion at the spectrogram level,\n e.g. low negative values reflects lower distortion levels. For a SNR (like reported\n in the MultiBandDiffusion paper), just take `-rvm`.\n\n Then it computes the mel spectrogram `z_ref` and `z_est` and compute volume of the difference\n relative to the volume of `z_ref` for each time-frequency bin. It further adds some limits, e.g.\n clamping the values between -25 and 25 dB (controlled by `min_relative_volume` and `max_relative_volume`)\n with the goal of avoiding the loss being dominated by parts where the reference is almost silent.\n Indeed, volumes in dB can take unbounded values both towards -oo and +oo, which can make the final\n average metric harder to interpret. Besides, anything below -30 dB of attenuation would sound extremely\n good (for a neural network output, although sound engineers typically aim for much lower attenuations).\n Similarly, anything above +30 dB would just be completely missing the target, and there is no point\n in measuring by exactly how much it missed it. -25, 25 is a more conservative range, but also more\n in line with what neural nets currently can achieve.\n\n For instance, a Relative Volume Mel (RVM) score of -10 dB means that on average, the delta between\n the target and reference mel-spec is 10 dB lower than the reference mel-spec value.\n\n The metric can be aggregated over a given frequency band in order have different insights for\n different region of the spectrum. `num_aggregated_bands` controls the number of bands.\n\n ..Warning:: While this function is optimized for interpretability, nothing was done to ensure it\n is numerically stable when computing its gradient. We thus advise against using it as a training loss.\n\n Args:\n sample_rate (int): Sample rate of the input audio.\n n_mels (int): Number of mel bands to use.\n n_fft (int): Number of frequency bins for the STFT.\n hop_length (int): Hop length of the STFT and the mel-spectrogram.\n min_relative_volume (float): The error `z_ref - z_est` volume is given relative to\n the volume of `z_ref`. If error is smaller than -25 dB of `z_ref`, then it is clamped.\n max_relative_volume (float): Same as `min_relative_volume` but clamping if the error is larger than that.\n max_initial_gain (float): When rescaling the audio at the very beginning, we will limit the gain\n to that amount, to avoid rescaling near silence. Given in dB.\n min_activity_volume (float): When computing the reference level from `z_ref`, will clamp low volume\n bins to that amount. This is effectively our \"zero\" level for the reference mel-spectrogram,\n and anything below that will be considered equally.\n num_aggregated_bands (int): Number of bands to keep when computing the average RVM value.\n For instance, a value of 3 would give 3 scores, roughly for low, mid and high freqs.\n \"\"\"\n def __init__(self, sample_rate: int = 24000, n_mels: int = 80, n_fft: int = 512,\n hop_length: int = 128, min_relative_volume: float = -25,\n max_relative_volume: float = 25, max_initial_gain: float = 25,\n min_activity_volume: float = -25,\n num_aggregated_bands: int = 4) -> None:\n super().__init__()\n self.melspec = torchaudio.transforms.MelSpectrogram(\n n_mels=n_mels, n_fft=n_fft, hop_length=hop_length,\n normalized=True, sample_rate=sample_rate, power=2)\n self.min_relative_volume = min_relative_volume\n self.max_relative_volume = max_relative_volume\n self.max_initial_gain = max_initial_gain\n self.min_activity_volume = min_activity_volume\n self.num_aggregated_bands = num_aggregated_bands\n\n def forward(self, estimate: torch.Tensor, ground_truth: torch.Tensor) -> tp.Dict[str, torch.Tensor]:\n \"\"\"Compute RVM metric between estimate and reference samples.\n\n Args:\n estimate (torch.Tensor): Estimate sample.\n ground_truth (torch.Tensor): Reference sample.\n\n Returns:\n dict[str, torch.Tensor]: Metrics with keys `rvm` for the overall average, and `rvm_{k}`\n for the RVM over the k-th band (k=0..num_aggregated_bands - 1).\n \"\"\"\n min_scale = db_to_scale(-self.max_initial_gain)\n std = ground_truth.pow(2).mean().sqrt().clamp(min=min_scale)\n z_gt = self.melspec(ground_truth / std).sqrt()\n z_est = self.melspec(estimate / std).sqrt()\n\n delta = z_gt - z_est\n ref_db = scale_to_db(z_gt, self.min_activity_volume)\n delta_db = scale_to_db(delta.abs(), min_volume=-120)\n relative_db = (delta_db - ref_db).clamp(self.min_relative_volume, self.max_relative_volume)\n dims = list(range(relative_db.dim()))\n dims.remove(dims[-2])\n losses_per_band = relative_db.mean(dim=dims)\n aggregated = [chunk.mean() for chunk in losses_per_band.chunk(self.num_aggregated_bands, dim=0)]\n metrics = {f'rvm_{index}': value for index, value in enumerate(aggregated)}\n metrics['rvm'] = losses_per_band.mean()\n return metrics\n" }, { "path": "audiocraft/metrics/visqol.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport csv\nimport json\nimport logging\nfrom pathlib import Path\nimport tempfile\nimport typing as tp\nimport subprocess\nimport shutil\n\nimport torch\nimport torchaudio\n\nlogger = logging.getLogger(__name__)\n\n\nclass ViSQOL:\n \"\"\"ViSQOL wrapper to run ViSQOL from Python using a pre-installed binary.\n\n To learn more about ViSQOL and how to build ViSQOL binary using bazel, please refer to the\n instructions available in the open source repository: https://github.com/google/visqol\n\n ViSQOL is capable of running in two modes:\n\n Audio Mode:\n When running in audio mode, input signals must have a 48kHz sample rate. Input should be resampled to 48kHz.\n Input signals can be multi-channel, but they will be down-mixed to mono for performing the comparison.\n Audio mode uses support vector regression, with the maximum range at ~4.75.\n\n Speech Mode:\n When running in speech mode, ViSQOL uses a wideband model. It therefore expects input sample rates of 16kHz.\n Input should be resampled to 16kHz.\n As part of the speech mode processing, a root mean square implementation for voice activity detection\n is performed on the reference signal to determine what parts of the signal have voice activity and\n should therefore be included in the comparison. The signal is normalized before performing the voice\n activity detection.\n Input signals can be multi-channel, but they will be down-mixed to mono for performing the comparison.\n Speech mode is scaled to have a maximum MOS of 5.0 to match previous version behavior.\n\n For more details, check the guidelines: https://github.com/google/visqol#general-guidelines-for-input\n\n Args:\n visqol_bin (str): Path to the ViSQOL binary.\n mode (str): ViSQOL computation mode, expecting \"audio\" or \"speech\".\n model (str): Name of the model to use for similarity to quality model.\n debug (bool): Whether to also get debug metrics from ViSQOL or not.\n \"\"\"\n SAMPLE_RATES_MODES = {\"audio\": 48_000, \"speech\": 16_000}\n ALLOWED_SAMPLE_RATES = frozenset(SAMPLE_RATES_MODES.values())\n\n def __init__(self, bin: tp.Union[Path, str], mode: str = \"audio\",\n model: str = \"libsvm_nu_svr_model.txt\", debug: bool = False):\n assert bin is not None and Path(bin).exists(), f\"Could not find ViSQOL binary in specified path: {bin}\"\n self.visqol_bin = str(bin)\n self.visqol_mode = mode\n self.target_sr = self._get_target_sr(self.visqol_mode)\n self.model = model\n self.debug = debug\n assert Path(self.visqol_model).exists(), \\\n f\"Could not find the specified model in ViSQOL install: {self.visqol_model}\"\n\n def _get_target_sr(self, mode: str) -> int:\n # returns target sampling rate for the corresponding ViSQOL mode.\n if mode not in ViSQOL.SAMPLE_RATES_MODES:\n raise ValueError(\n f\"Unsupported mode! Allowed are: {', '.join(ViSQOL.SAMPLE_RATES_MODES.keys())}\"\n )\n return ViSQOL.SAMPLE_RATES_MODES[mode]\n\n def _prepare_files(\n self, ref_sig: torch.Tensor, deg_sig: torch.Tensor, sr: int, target_sr: int, pad_with_silence: bool = False\n ):\n # prepare files for ViSQOL evaluation.\n assert target_sr in ViSQOL.ALLOWED_SAMPLE_RATES\n assert len(ref_sig) == len(deg_sig), (\n \"Expects same number of ref and degraded inputs\",\n f\" but ref len {len(ref_sig)} != deg len {len(deg_sig)}\"\n )\n # resample audio if needed\n if sr != target_sr:\n transform = torchaudio.transforms.Resample(sr, target_sr)\n pad = int(0.5 * target_sr)\n rs_ref = []\n rs_deg = []\n for i in range(len(ref_sig)):\n rs_ref_i = transform(ref_sig[i])\n rs_deg_i = transform(deg_sig[i])\n if pad_with_silence:\n rs_ref_i = torch.nn.functional.pad(rs_ref_i, (pad, pad), mode='constant', value=0)\n rs_deg_i = torch.nn.functional.pad(rs_deg_i, (pad, pad), mode='constant', value=0)\n rs_ref.append(rs_ref_i)\n rs_deg.append(rs_deg_i)\n ref_sig = torch.stack(rs_ref)\n deg_sig = torch.stack(rs_deg)\n # save audio chunks to tmp dir and create csv\n tmp_dir = Path(tempfile.mkdtemp())\n try:\n tmp_input_csv_path = tmp_dir / \"input.csv\"\n tmp_results_csv_path = tmp_dir / \"results.csv\"\n tmp_debug_json_path = tmp_dir / \"debug.json\"\n with open(tmp_input_csv_path, \"w\") as csv_file:\n csv_writer = csv.writer(csv_file)\n csv_writer.writerow([\"reference\", \"degraded\"])\n for i in range(len(ref_sig)):\n tmp_ref_filename = tmp_dir / f\"ref_{i}.wav\"\n tmp_deg_filename = tmp_dir / f\"deg_{i}.wav\"\n torchaudio.save(\n tmp_ref_filename,\n torch.clamp(ref_sig[i], min=-0.99, max=0.99),\n sample_rate=target_sr,\n bits_per_sample=16,\n encoding=\"PCM_S\"\n )\n torchaudio.save(\n tmp_deg_filename,\n torch.clamp(deg_sig[i], min=-0.99, max=0.99),\n sample_rate=target_sr,\n bits_per_sample=16,\n encoding=\"PCM_S\"\n )\n csv_writer.writerow([str(tmp_ref_filename), str(tmp_deg_filename)])\n return tmp_dir, tmp_input_csv_path, tmp_results_csv_path, tmp_debug_json_path\n except Exception as e:\n logger.error(\"Exception occurred when preparing files for ViSQOL: %s\", e)\n return tmp_dir, None, None, None\n\n def _flush_files(self, tmp_dir: tp.Union[Path, str]):\n # flush tmp files used to compute ViSQOL.\n shutil.rmtree(str(tmp_dir))\n\n def _collect_moslqo_score(self, results_csv_path: tp.Union[Path, str]) -> float:\n # collect results for each evaluated pair and return averaged moslqo score.\n with open(results_csv_path, \"r\") as csv_file:\n reader = csv.DictReader(csv_file)\n moslqo_scores = [float(row[\"moslqo\"]) for row in reader]\n if len(moslqo_scores) > 0:\n return sum(moslqo_scores) / len(moslqo_scores)\n else:\n return 0.0\n\n def _collect_debug_data(self, debug_json_path: tp.Union[Path, str]) -> dict:\n # collect debug data for the visqol inference.\n with open(debug_json_path, \"r\") as f:\n data = json.load(f)\n return data\n\n @property\n def visqol_model(self):\n return f'{self.visqol_bin}/model/{self.model}'\n\n def _run_visqol(\n self,\n input_csv_path: tp.Union[Path, str],\n results_csv_path: tp.Union[Path, str],\n debug_csv_path: tp.Optional[tp.Union[Path, str]],\n ):\n input_csv_path = str(input_csv_path)\n results_csv_path = str(results_csv_path)\n debug_csv_path = str(debug_csv_path)\n cmd = [\n f'{self.visqol_bin}/bazel-bin/visqol',\n '--batch_input_csv', f'{input_csv_path}',\n '--results_csv', f'{results_csv_path}'\n ]\n if debug_csv_path is not None:\n cmd += ['--output_debug', f'{debug_csv_path}']\n if self.visqol_mode == \"speech\":\n cmd += ['--use_speech_mode']\n cmd += ['--similarity_to_quality_model', f'{self.visqol_model}']\n result = subprocess.run(cmd, capture_output=True)\n if result.returncode:\n logger.error(\"Error with visqol: \\n %s \\n %s\", result.stdout.decode(), result.stderr.decode())\n raise RuntimeError(\"Error while executing visqol\")\n result.check_returncode()\n\n def __call__(\n self,\n ref_sig: torch.Tensor,\n deg_sig: torch.Tensor,\n sr: int,\n pad_with_silence: bool = False,\n ):\n \"\"\"Calculate the ViSQOL metric for a pair of audio signals at a given sample rate.\n Args:\n ref_sig (torch.Tensor): Reference signals as [B, C, T].\n deg_sig (torch.Tensor): Degraded signals as [B, C, T].\n sr (int): Sample rate of the two audio signals.\n pad_with_silence (bool): Whether to pad the file with silences as recommended\n in visqol guidelines (see: https://github.com/google/visqol#general-guidelines-for-input).\n Returns:\n float: The ViSQOL score or mean score for the batch.\n \"\"\"\n logger.debug(f\"Calculating visqol with mode={self.visqol_mode} on {len(ref_sig)} samples\")\n tmp_dir, input_csv, results_csv, debug_json = self._prepare_files(\n ref_sig, deg_sig, sr, self.target_sr, pad_with_silence\n )\n try:\n if input_csv and results_csv:\n self._run_visqol(\n input_csv,\n results_csv,\n debug_json if self.debug else None,\n )\n mosqol = self._collect_moslqo_score(results_csv)\n return mosqol\n else:\n raise RuntimeError(\"Something unexpected happened when running VISQOL!\")\n except Exception as e:\n logger.error(\"Exception occurred when running ViSQOL: %s\", e)\n finally:\n self._flush_files(tmp_dir)\n" }, { "path": "audiocraft/models/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"\nModels for EnCodec, AudioGen, MusicGen, as well as the generic LMModel.\n\"\"\"\n# flake8: noqa\nfrom . import builders, loaders\nfrom .encodec import (\n CompressionModel, EncodecModel, DAC,\n HFEncodecModel, HFEncodecCompressionModel)\nfrom .audiogen import AudioGen\nfrom .lm import LMModel\nfrom .lm_magnet import MagnetLMModel\nfrom .flow_matching import FlowMatchingModel\nfrom .multibanddiffusion import MultiBandDiffusion\nfrom .musicgen import MusicGen\nfrom .magnet import MAGNeT\nfrom .unet import DiffusionUnet\nfrom .watermark import WMModel\nfrom .jasco import JASCO\n" }, { "path": "audiocraft/models/audiogen.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nMain model for using AudioGen. This will combine all the required components\nand provide easy access to the generation API.\n\"\"\"\n\nimport typing as tp\n\nimport torch\n\nfrom .encodec import CompressionModel\nfrom .genmodel import BaseGenModel\nfrom .lm import LMModel\nfrom .builders import get_debug_compression_model, get_debug_lm_model\nfrom .loaders import load_compression_model, load_lm_model\n\n\nclass AudioGen(BaseGenModel):\n \"\"\"AudioGen main model with convenient generation API.\n\n Args:\n name (str): name of the model.\n compression_model (CompressionModel): Compression model\n used to map audio to invertible discrete representations.\n lm (LMModel): Language model over discrete representations.\n max_duration (float, optional): maximum duration the model can produce,\n otherwise, inferred from the training params.\n \"\"\"\n def __init__(self, name: str, compression_model: CompressionModel, lm: LMModel,\n max_duration: tp.Optional[float] = None):\n super().__init__(name, compression_model, lm, max_duration)\n self.set_generation_params(duration=5) # default duration\n\n @staticmethod\n def get_pretrained(name: str = 'facebook/audiogen-medium', device=None):\n \"\"\"Return pretrained model, we provide a single model for now:\n - facebook/audiogen-medium (1.5B), text to sound,\n # see: https://huggingface.co/facebook/audiogen-medium\n \"\"\"\n if device is None:\n if torch.cuda.device_count():\n device = 'cuda'\n else:\n device = 'cpu'\n\n if name == 'debug':\n # used only for unit tests\n compression_model = get_debug_compression_model(device, sample_rate=16000)\n lm = get_debug_lm_model(device)\n return AudioGen(name, compression_model, lm, max_duration=10)\n\n compression_model = load_compression_model(name, device=device)\n lm = load_lm_model(name, device=device)\n assert 'self_wav' not in lm.condition_provider.conditioners, \\\n \"AudioGen do not support waveform conditioning for now\"\n return AudioGen(name, compression_model, lm)\n\n def set_generation_params(self, use_sampling: bool = True, top_k: int = 250,\n top_p: float = 0.0, temperature: float = 1.0,\n duration: float = 10.0, cfg_coef: float = 3.0,\n two_step_cfg: bool = False, extend_stride: float = 2):\n \"\"\"Set the generation parameters for AudioGen.\n\n Args:\n use_sampling (bool, optional): Use sampling if True, else do argmax decoding. Defaults to True.\n top_k (int, optional): top_k used for sampling. Defaults to 250.\n top_p (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.0.\n temperature (float, optional): Softmax temperature parameter. Defaults to 1.0.\n duration (float, optional): Duration of the generated waveform. Defaults to 10.0.\n cfg_coef (float, optional): Coefficient used for classifier free guidance. Defaults to 3.0.\n two_step_cfg (bool, optional): If True, performs 2 forward for Classifier Free Guidance,\n instead of batching together the two. This has some impact on how things\n are padded but seems to have little impact in practice.\n extend_stride: when doing extended generation (i.e. more than 10 seconds), by how much\n should we extend the audio each time. Larger values will mean less context is\n preserved, and shorter value will require extra computations.\n \"\"\"\n assert extend_stride < self.max_duration, \"Cannot stride by more than max generation duration.\"\n self.extend_stride = extend_stride\n self.duration = duration\n self.generation_params = {\n 'use_sampling': use_sampling,\n 'temp': temperature,\n 'top_k': top_k,\n 'top_p': top_p,\n 'cfg_coef': cfg_coef,\n 'two_step_cfg': two_step_cfg,\n }\n" }, { "path": "audiocraft/models/builders.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nAll the functions to build the relevant models and modules\nfrom the Hydra config.\n\"\"\"\n\nimport typing as tp\n\nimport omegaconf\nimport torch\n\nimport audiocraft\n\nfrom .. import quantization as qt\nfrom ..modules.codebooks_patterns import (CoarseFirstPattern,\n CodebooksPatternProvider,\n DelayedPatternProvider,\n MusicLMPattern,\n ParallelPatternProvider,\n UnrolledPatternProvider)\nfrom ..modules.conditioners import (BaseConditioner, ChromaStemConditioner,\n CLAPEmbeddingConditioner,\n ConditionFuser, JascoCondConst,\n ConditioningProvider, LUTConditioner,\n T5Conditioner, StyleConditioner)\nfrom ..modules.jasco_conditioners import (JascoConditioningProvider, ChordsEmbConditioner,\n DrumsConditioner, MelodyConditioner)\nfrom ..modules.diffusion_schedule import MultiBandProcessor, SampleProcessor\nfrom ..utils.utils import dict_from_config\nfrom .encodec import (CompressionModel, EncodecModel,\n InterleaveStereoCompressionModel)\nfrom .lm import LMModel\nfrom .lm_magnet import MagnetLMModel\nfrom .flow_matching import FlowMatchingModel\nfrom .unet import DiffusionUnet\nfrom .watermark import WMModel\n\n\ndef get_quantizer(\n quantizer: str, cfg: omegaconf.DictConfig, dimension: int\n) -> qt.BaseQuantizer:\n klass = {\"no_quant\": qt.DummyQuantizer, \"rvq\": qt.ResidualVectorQuantizer}[\n quantizer\n ]\n kwargs = dict_from_config(getattr(cfg, quantizer))\n if quantizer != \"no_quant\":\n kwargs[\"dimension\"] = dimension\n return klass(**kwargs)\n\n\ndef get_encodec_autoencoder(encoder_name: str, cfg: omegaconf.DictConfig):\n if encoder_name == \"seanet\":\n kwargs = dict_from_config(getattr(cfg, \"seanet\"))\n encoder_override_kwargs = kwargs.pop(\"encoder\")\n decoder_override_kwargs = kwargs.pop(\"decoder\")\n encoder_kwargs = {**kwargs, **encoder_override_kwargs}\n decoder_kwargs = {**kwargs, **decoder_override_kwargs}\n encoder = audiocraft.modules.SEANetEncoder(**encoder_kwargs)\n decoder = audiocraft.modules.SEANetDecoder(**decoder_kwargs)\n return encoder, decoder\n else:\n raise KeyError(f\"Unexpected compression model {cfg.compression_model}\")\n\n\ndef get_compression_model(cfg: omegaconf.DictConfig) -> CompressionModel:\n \"\"\"Instantiate a compression model.\"\"\"\n if cfg.compression_model == \"encodec\":\n kwargs = dict_from_config(getattr(cfg, \"encodec\"))\n encoder_name = kwargs.pop(\"autoencoder\")\n quantizer_name = kwargs.pop(\"quantizer\")\n encoder, decoder = get_encodec_autoencoder(encoder_name, cfg)\n quantizer = get_quantizer(quantizer_name, cfg, encoder.dimension)\n frame_rate = kwargs[\"sample_rate\"] // encoder.hop_length\n renormalize = kwargs.pop(\"renormalize\", False)\n # deprecated params\n kwargs.pop(\"renorm\", None)\n return EncodecModel(\n encoder,\n decoder,\n quantizer,\n frame_rate=frame_rate,\n renormalize=renormalize,\n **kwargs,\n ).to(cfg.device)\n else:\n raise KeyError(f\"Unexpected compression model {cfg.compression_model}\")\n\n\ndef get_jasco_model(cfg: omegaconf.DictConfig,\n compression_model: tp.Optional[CompressionModel] = None) -> FlowMatchingModel:\n kwargs = dict_from_config(getattr(cfg, \"transformer_lm\"))\n attribute_dropout = dict_from_config(getattr(cfg, \"attribute_dropout\"))\n cls_free_guidance = dict_from_config(getattr(cfg, \"classifier_free_guidance\"))\n cfg_prob = cls_free_guidance[\"training_dropout\"]\n cfg_coef = cls_free_guidance[\"inference_coef\"]\n fuser = get_condition_fuser(cfg)\n condition_provider = get_conditioner_provider(kwargs[\"dim\"], cfg).to(cfg.device)\n if JascoCondConst.DRM.value in condition_provider.conditioners: # use self_wav for drums\n assert compression_model is not None\n\n # use compression model for drums conditioning\n condition_provider.conditioners.self_wav.compression_model = compression_model\n condition_provider.conditioners.self_wav.compression_model.requires_grad_(False)\n\n # downcast to jasco conditioning provider\n seq_len = cfg.compression_model_framerate * cfg.dataset.segment_duration\n chords_card = cfg.conditioners.chords.chords_emb.card if JascoCondConst.CRD.value in cfg.conditioners else -1\n condition_provider = JascoConditioningProvider(device=condition_provider.device,\n conditioners=condition_provider.conditioners,\n chords_card=chords_card,\n sequence_length=seq_len)\n\n if len(fuser.fuse2cond[\"cross\"]) > 0: # enforce cross-att programmatically\n kwargs[\"cross_attention\"] = True\n\n kwargs.pop(\"n_q\", None)\n kwargs.pop(\"card\", None)\n\n return FlowMatchingModel(\n condition_provider=condition_provider,\n fuser=fuser,\n cfg_dropout=cfg_prob,\n cfg_coef=cfg_coef,\n attribute_dropout=attribute_dropout,\n dtype=getattr(torch, cfg.dtype),\n device=cfg.device,\n **kwargs,\n ).to(cfg.device)\n\n\ndef get_lm_model(cfg: omegaconf.DictConfig) -> LMModel:\n \"\"\"Instantiate a transformer LM.\"\"\"\n if cfg.lm_model in [\"transformer_lm\", \"transformer_lm_magnet\"]:\n kwargs = dict_from_config(getattr(cfg, \"transformer_lm\"))\n n_q = kwargs[\"n_q\"]\n q_modeling = kwargs.pop(\"q_modeling\", None)\n codebooks_pattern_cfg = getattr(cfg, \"codebooks_pattern\")\n attribute_dropout = dict_from_config(getattr(cfg, \"attribute_dropout\"))\n cls_free_guidance = dict_from_config(getattr(cfg, \"classifier_free_guidance\"))\n cfg_prob, cfg_coef = (\n cls_free_guidance[\"training_dropout\"],\n cls_free_guidance[\"inference_coef\"],\n )\n fuser = get_condition_fuser(cfg)\n condition_provider = get_conditioner_provider(kwargs[\"dim\"], cfg).to(cfg.device)\n if len(fuser.fuse2cond[\"cross\"]) > 0: # enforce cross-att programmatically\n kwargs[\"cross_attention\"] = True\n if codebooks_pattern_cfg.modeling is None:\n assert (\n q_modeling is not None\n ), \"LM model should either have a codebook pattern defined or transformer_lm.q_modeling\"\n codebooks_pattern_cfg = omegaconf.OmegaConf.create(\n {\"modeling\": q_modeling, \"delay\": {\"delays\": list(range(n_q))}}\n )\n\n pattern_provider = get_codebooks_pattern_provider(n_q, codebooks_pattern_cfg)\n lm_class = MagnetLMModel if cfg.lm_model == \"transformer_lm_magnet\" else LMModel\n return lm_class(\n pattern_provider=pattern_provider,\n condition_provider=condition_provider,\n fuser=fuser,\n cfg_dropout=cfg_prob,\n cfg_coef=cfg_coef,\n attribute_dropout=attribute_dropout,\n dtype=getattr(torch, cfg.dtype),\n device=cfg.device,\n **kwargs,\n ).to(cfg.device)\n else:\n raise KeyError(f\"Unexpected LM model {cfg.lm_model}\")\n\n\ndef get_conditioner_provider(\n output_dim: int, cfg: omegaconf.DictConfig\n) -> ConditioningProvider:\n \"\"\"Instantiate a conditioning model.\"\"\"\n device = cfg.device\n duration = cfg.dataset.segment_duration\n cfg = getattr(cfg, \"conditioners\")\n dict_cfg = {} if cfg is None else dict_from_config(cfg)\n conditioners: tp.Dict[str, BaseConditioner] = {}\n condition_provider_args = dict_cfg.pop(\"args\", {})\n condition_provider_args.pop(\"merge_text_conditions_p\", None)\n condition_provider_args.pop(\"drop_desc_p\", None)\n\n for cond, cond_cfg in dict_cfg.items():\n model_type = cond_cfg[\"model\"]\n model_args = cond_cfg[model_type]\n if model_type == \"t5\":\n conditioners[str(cond)] = T5Conditioner(\n output_dim=output_dim, device=device, **model_args\n )\n elif model_type == \"lut\":\n conditioners[str(cond)] = LUTConditioner(\n output_dim=output_dim, **model_args\n )\n elif model_type == \"chroma_stem\":\n conditioners[str(cond)] = ChromaStemConditioner(\n output_dim=output_dim, duration=duration, device=device, **model_args\n )\n elif model_type in {\"chords_emb\", \"drum_latents\", \"melody\"}:\n conditioners_classes = {\"chords_emb\": ChordsEmbConditioner,\n \"drum_latents\": DrumsConditioner,\n \"melody\": MelodyConditioner}\n conditioner_class = conditioners_classes[model_type]\n conditioners[str(cond)] = conditioner_class(device=device, **model_args)\n elif model_type == \"clap\":\n conditioners[str(cond)] = CLAPEmbeddingConditioner(\n output_dim=output_dim, device=device, **model_args\n )\n elif model_type == 'style':\n conditioners[str(cond)] = StyleConditioner(\n output_dim=output_dim,\n device=device,\n **model_args\n )\n else:\n raise ValueError(f\"Unrecognized conditioning model: {model_type}\")\n conditioner = ConditioningProvider(\n conditioners, device=device, **condition_provider_args\n )\n return conditioner\n\n\ndef get_condition_fuser(cfg: omegaconf.DictConfig) -> ConditionFuser:\n \"\"\"Instantiate a condition fuser object.\"\"\"\n fuser_cfg = getattr(cfg, \"fuser\")\n fuser_methods = [\"sum\", \"cross\", \"prepend\", \"ignore\", \"input_interpolate\"]\n fuse2cond = {k: fuser_cfg[k] for k in fuser_methods if k in fuser_cfg}\n kwargs = {k: v for k, v in fuser_cfg.items() if k not in fuser_methods}\n fuser = ConditionFuser(fuse2cond=fuse2cond, **kwargs)\n return fuser\n\n\ndef get_codebooks_pattern_provider(\n n_q: int, cfg: omegaconf.DictConfig\n) -> CodebooksPatternProvider:\n \"\"\"Instantiate a codebooks pattern provider object.\"\"\"\n pattern_providers = {\n \"parallel\": ParallelPatternProvider,\n \"delay\": DelayedPatternProvider,\n \"unroll\": UnrolledPatternProvider,\n \"coarse_first\": CoarseFirstPattern,\n \"musiclm\": MusicLMPattern,\n }\n name = cfg.modeling\n kwargs = dict_from_config(cfg.get(name)) if hasattr(cfg, name) else {}\n klass = pattern_providers[name]\n return klass(n_q, **kwargs)\n\n\ndef get_debug_compression_model(device=\"cpu\", sample_rate: int = 32000):\n \"\"\"Instantiate a debug compression model to be used for unit tests.\"\"\"\n assert sample_rate in [\n 16000,\n 32000,\n ], \"unsupported sample rate for debug compression model\"\n model_ratios = {\n 16000: [10, 8, 8], # 25 Hz at 16kHz\n 32000: [10, 8, 16], # 25 Hz at 32kHz\n }\n ratios: tp.List[int] = model_ratios[sample_rate]\n frame_rate = 25\n seanet_kwargs: dict = {\n \"n_filters\": 4,\n \"n_residual_layers\": 1,\n \"dimension\": 32,\n \"ratios\": ratios,\n }\n encoder = audiocraft.modules.SEANetEncoder(**seanet_kwargs)\n decoder = audiocraft.modules.SEANetDecoder(**seanet_kwargs)\n quantizer = qt.ResidualVectorQuantizer(dimension=32, bins=400, n_q=4)\n init_x = torch.randn(8, 32, 128)\n quantizer(init_x, 1) # initialize kmeans etc.\n compression_model = EncodecModel(\n encoder,\n decoder,\n quantizer,\n frame_rate=frame_rate,\n sample_rate=sample_rate,\n channels=1,\n ).to(device)\n return compression_model.eval()\n\n\ndef get_diffusion_model(cfg: omegaconf.DictConfig):\n # TODO Find a way to infer the channels from dset\n channels = cfg.channels\n num_steps = cfg.schedule.num_steps\n return DiffusionUnet(chin=channels, num_steps=num_steps, **cfg.diffusion_unet)\n\n\ndef get_processor(cfg, sample_rate: int = 24000):\n sample_processor = SampleProcessor()\n if cfg.use:\n kw = dict(cfg)\n kw.pop(\"use\")\n kw.pop(\"name\")\n if cfg.name == \"multi_band_processor\":\n sample_processor = MultiBandProcessor(sample_rate=sample_rate, **kw)\n return sample_processor\n\n\ndef get_debug_lm_model(device=\"cpu\"):\n \"\"\"Instantiate a debug LM to be used for unit tests.\"\"\"\n pattern = DelayedPatternProvider(n_q=4)\n dim = 16\n providers = {\n \"description\": LUTConditioner(\n n_bins=128, dim=dim, output_dim=dim, tokenizer=\"whitespace\"\n ),\n }\n condition_provider = ConditioningProvider(providers)\n fuser = ConditionFuser(\n {\"cross\": [\"description\"], \"prepend\": [], \"sum\": [], \"input_interpolate\": []}\n )\n lm = LMModel(\n pattern,\n condition_provider,\n fuser,\n n_q=4,\n card=400,\n dim=dim,\n num_heads=4,\n custom=True,\n num_layers=2,\n cross_attention=True,\n causal=True,\n )\n return lm.to(device).eval()\n\n\ndef get_wrapped_compression_model(\n compression_model: CompressionModel, cfg: omegaconf.DictConfig\n) -> CompressionModel:\n if hasattr(cfg, \"interleave_stereo_codebooks\"):\n if cfg.interleave_stereo_codebooks.use:\n kwargs = dict_from_config(cfg.interleave_stereo_codebooks)\n kwargs.pop(\"use\")\n compression_model = InterleaveStereoCompressionModel(\n compression_model, **kwargs\n )\n if hasattr(cfg, \"compression_model_n_q\"):\n if cfg.compression_model_n_q is not None:\n compression_model.set_num_codebooks(cfg.compression_model_n_q)\n return compression_model\n\n\ndef get_watermark_model(cfg: omegaconf.DictConfig) -> WMModel:\n \"\"\"Build a WMModel based by audioseal. This requires audioseal to be installed\"\"\"\n import audioseal\n\n from .watermark import AudioSeal\n\n # Builder encoder and decoder directly using audiocraft API to avoid cyclic import\n assert hasattr(\n cfg, \"seanet\"\n ), \"Missing required `seanet` parameters in AudioSeal config\"\n encoder, decoder = get_encodec_autoencoder(\"seanet\", cfg)\n\n # Build message processor\n kwargs = (\n dict_from_config(getattr(cfg, \"audioseal\")) if hasattr(cfg, \"audioseal\") else {}\n )\n nbits = kwargs.get(\"nbits\", 0)\n hidden_size = getattr(cfg.seanet, \"dimension\", 128)\n msg_processor = audioseal.MsgProcessor(nbits, hidden_size=hidden_size)\n\n # Build detector using audioseal API\n def _get_audioseal_detector():\n # We don't need encoder and decoder params from seanet, remove them\n seanet_cfg = dict_from_config(cfg.seanet)\n seanet_cfg.pop(\"encoder\")\n seanet_cfg.pop(\"decoder\")\n detector_cfg = dict_from_config(cfg.detector)\n\n typed_seanet_cfg = audioseal.builder.SEANetConfig(**seanet_cfg)\n typed_detector_cfg = audioseal.builder.DetectorConfig(**detector_cfg)\n _cfg = audioseal.builder.AudioSealDetectorConfig(\n nbits=nbits, seanet=typed_seanet_cfg, detector=typed_detector_cfg\n )\n return audioseal.builder.create_detector(_cfg)\n\n detector = _get_audioseal_detector()\n generator = audioseal.AudioSealWM(\n encoder=encoder, decoder=decoder, msg_processor=msg_processor\n )\n model = AudioSeal(generator=generator, detector=detector, nbits=nbits)\n\n device = torch.device(getattr(cfg, \"device\", \"cpu\"))\n dtype = getattr(torch, getattr(cfg, \"dtype\", \"float32\"))\n return model.to(device=device, dtype=dtype)\n" }, { "path": "audiocraft/models/encodec.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Compression models or wrapper around existing models.\nAlso defines the main interface that a model must follow to be usable as an audio tokenizer.\n\"\"\"\n\nfrom abc import ABC, abstractmethod\nimport logging\nimport math\nfrom pathlib import Path\nimport typing as tp\n\nfrom einops import rearrange\nimport numpy as np\nimport torch\nfrom torch import nn\nfrom transformers import EncodecModel as HFEncodecModel\n\nfrom .. import quantization as qt\n\n\nlogger = logging.getLogger()\n\n\nclass CompressionModel(ABC, nn.Module):\n \"\"\"Base API for all compression models that aim at being used as audio tokenizers\n with a language model.\n \"\"\"\n\n @abstractmethod\n def forward(self, x: torch.Tensor) -> qt.QuantizedResult:\n ...\n\n @abstractmethod\n def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:\n \"\"\"See `EncodecModel.encode`.\"\"\"\n ...\n\n @abstractmethod\n def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):\n \"\"\"See `EncodecModel.decode`.\"\"\"\n ...\n\n @abstractmethod\n def decode_latent(self, codes: torch.Tensor):\n \"\"\"Decode from the discrete codes to continuous latent space.\"\"\"\n ...\n\n @property\n @abstractmethod\n def channels(self) -> int:\n ...\n\n @property\n @abstractmethod\n def frame_rate(self) -> float:\n ...\n\n @property\n @abstractmethod\n def sample_rate(self) -> int:\n ...\n\n @property\n @abstractmethod\n def cardinality(self) -> int:\n ...\n\n @property\n @abstractmethod\n def num_codebooks(self) -> int:\n ...\n\n @property\n @abstractmethod\n def total_codebooks(self) -> int:\n ...\n\n @abstractmethod\n def set_num_codebooks(self, n: int):\n \"\"\"Set the active number of codebooks used by the quantizer.\"\"\"\n ...\n\n @staticmethod\n def get_pretrained(\n name: str, device: tp.Union[torch.device, str] = 'cpu'\n ) -> 'CompressionModel':\n \"\"\"Instantiate a CompressionModel from a given pretrained model.\n\n Args:\n name (Path or str): name of the pretrained model. See after.\n device (torch.device or str): Device on which the model is loaded.\n\n Pretrained models:\n - dac_44khz (https://github.com/descriptinc/descript-audio-codec)\n - dac_24khz (same)\n - facebook/encodec_24khz (https://huggingface.co/facebook/encodec_24khz)\n - facebook/encodec_32khz (https://huggingface.co/facebook/encodec_32khz)\n - your own model on Hugging Face. Export instructions to come...\n \"\"\"\n\n from . import builders, loaders\n model: CompressionModel\n if name in ['dac_44khz', 'dac_24khz']:\n model_type = name.split('_')[1]\n logger.info(\"Getting pretrained compression model from DAC %s\", model_type)\n model = DAC(model_type)\n elif name in ['debug_compression_model']:\n logger.info(\"Getting pretrained compression model for debug\")\n model = builders.get_debug_compression_model()\n elif Path(name).exists():\n # We assume here if the path exists that it is in fact an AC checkpoint\n # that was exported using `audiocraft.utils.export` functions.\n model = loaders.load_compression_model(name, device=device)\n else:\n logger.info(\"Getting pretrained compression model from HF %s\", name)\n hf_model = HFEncodecModel.from_pretrained(name)\n model = HFEncodecCompressionModel(hf_model).to(device)\n return model.to(device).eval()\n\n\nclass EncodecModel(CompressionModel):\n \"\"\"Encodec model operating on the raw waveform.\n\n Args:\n encoder (nn.Module): Encoder network.\n decoder (nn.Module): Decoder network.\n quantizer (qt.BaseQuantizer): Quantizer network.\n frame_rate (int): Frame rate for the latent representation.\n sample_rate (int): Audio sample rate.\n channels (int): Number of audio channels.\n causal (bool): Whether to use a causal version of the model.\n renormalize (bool): Whether to renormalize the audio before running the model.\n \"\"\"\n # we need assignment to override the property in the abstract class,\n # I couldn't find a better way...\n frame_rate: float = 0\n sample_rate: int = 0\n channels: int = 0\n\n def __init__(self,\n encoder: nn.Module,\n decoder: nn.Module,\n quantizer: qt.BaseQuantizer,\n frame_rate: int,\n sample_rate: int,\n channels: int,\n causal: bool = False,\n renormalize: bool = False):\n super().__init__()\n self.encoder = encoder\n self.decoder = decoder\n self.quantizer = quantizer\n self.frame_rate = frame_rate\n self.sample_rate = sample_rate\n self.channels = channels\n self.renormalize = renormalize\n self.causal = causal\n if self.causal:\n # we force disabling here to avoid handling linear overlap of segments\n # as supported in original EnCodec codebase.\n assert not self.renormalize, 'Causal model does not support renormalize'\n\n @property\n def total_codebooks(self):\n \"\"\"Total number of quantizer codebooks available.\"\"\"\n return self.quantizer.total_codebooks\n\n @property\n def num_codebooks(self):\n \"\"\"Active number of codebooks used by the quantizer.\"\"\"\n return self.quantizer.num_codebooks\n\n def set_num_codebooks(self, n: int):\n \"\"\"Set the active number of codebooks used by the quantizer.\"\"\"\n self.quantizer.set_num_codebooks(n)\n\n @property\n def cardinality(self):\n \"\"\"Cardinality of each codebook.\"\"\"\n return self.quantizer.bins\n\n def preprocess(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:\n scale: tp.Optional[torch.Tensor]\n if self.renormalize:\n mono = x.mean(dim=1, keepdim=True)\n volume = mono.pow(2).mean(dim=2, keepdim=True).sqrt()\n scale = 1e-8 + volume\n x = x / scale\n scale = scale.view(-1, 1)\n else:\n scale = None\n return x, scale\n\n def postprocess(self,\n x: torch.Tensor,\n scale: tp.Optional[torch.Tensor] = None) -> torch.Tensor:\n if scale is not None:\n assert self.renormalize\n x = x * scale.view(-1, 1, 1)\n return x\n\n def forward(self, x: torch.Tensor) -> qt.QuantizedResult:\n assert x.dim() == 3\n length = x.shape[-1]\n x, scale = self.preprocess(x)\n\n emb = self.encoder(x)\n q_res = self.quantizer(emb, self.frame_rate)\n out = self.decoder(q_res.x)\n\n # remove extra padding added by the encoder and decoder\n assert out.shape[-1] >= length, (out.shape[-1], length)\n out = out[..., :length]\n\n q_res.x = self.postprocess(out, scale)\n\n return q_res\n\n def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:\n \"\"\"Encode the given input tensor to quantized representation along with scale parameter.\n\n Args:\n x (torch.Tensor): Float tensor of shape [B, C, T]\n\n Returns:\n codes, scale (tuple of torch.Tensor, torch.Tensor): Tuple composed of:\n codes: a float tensor of shape [B, K, T] with K the number of codebooks used and T the timestep.\n scale: a float tensor containing the scale for audio renormalization.\n \"\"\"\n assert x.dim() == 3\n x, scale = self.preprocess(x)\n emb = self.encoder(x)\n codes = self.quantizer.encode(emb)\n return codes, scale\n\n def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):\n \"\"\"Decode the given codes to a reconstructed representation, using the scale to perform\n audio denormalization if needed.\n\n Args:\n codes (torch.Tensor): Int tensor of shape [B, K, T]\n scale (torch.Tensor, optional): Float tensor containing the scale value.\n\n Returns:\n out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio.\n \"\"\"\n emb = self.decode_latent(codes)\n out = self.decoder(emb)\n out = self.postprocess(out, scale)\n # out contains extra padding added by the encoder and decoder\n return out\n\n def decode_latent(self, codes: torch.Tensor):\n \"\"\"Decode from the discrete codes to continuous latent space.\"\"\"\n return self.quantizer.decode(codes)\n\n\nclass DAC(CompressionModel):\n def __init__(self, model_type: str = \"44khz\"):\n super().__init__()\n try:\n import dac.utils\n except ImportError:\n raise RuntimeError(\"Could not import dac, make sure it is installed, \"\n \"please run `pip install descript-audio-codec`\")\n self.model = dac.utils.load_model(model_type=model_type)\n self.n_quantizers = self.total_codebooks\n self.model.eval()\n\n def forward(self, x: torch.Tensor) -> qt.QuantizedResult:\n # We don't support training with this.\n raise NotImplementedError(\"Forward and training with DAC not supported.\")\n\n def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:\n codes = self.model.encode(x, self.n_quantizers)[1]\n return codes[:, :self.n_quantizers], None\n\n def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):\n assert scale is None\n z_q = self.decode_latent(codes)\n return self.model.decode(z_q)\n\n def decode_latent(self, codes: torch.Tensor):\n \"\"\"Decode from the discrete codes to continuous latent space.\"\"\"\n return self.model.quantizer.from_codes(codes)[0]\n\n @property\n def channels(self) -> int:\n return 1\n\n @property\n def frame_rate(self) -> float:\n return self.model.sample_rate / self.model.hop_length\n\n @property\n def sample_rate(self) -> int:\n return self.model.sample_rate\n\n @property\n def cardinality(self) -> int:\n return self.model.codebook_size\n\n @property\n def num_codebooks(self) -> int:\n return self.n_quantizers\n\n @property\n def total_codebooks(self) -> int:\n return self.model.n_codebooks\n\n def set_num_codebooks(self, n: int):\n \"\"\"Set the active number of codebooks used by the quantizer.\n \"\"\"\n assert n >= 1\n assert n <= self.total_codebooks\n self.n_quantizers = n\n\n\nclass HFEncodecCompressionModel(CompressionModel):\n \"\"\"Wrapper around HuggingFace Encodec.\n \"\"\"\n def __init__(self, model: HFEncodecModel):\n super().__init__()\n self.model = model\n bws = self.model.config.target_bandwidths\n num_codebooks = [\n bw * 1000 / (self.frame_rate * math.log2(self.cardinality))\n for bw in bws\n ]\n deltas = [nc - int(nc) for nc in num_codebooks]\n # Checking we didn't do some bad maths and we indeed have integers!\n assert all(deltas) <= 1e-3, deltas\n self.possible_num_codebooks = [int(nc) for nc in num_codebooks]\n self.set_num_codebooks(max(self.possible_num_codebooks))\n\n def forward(self, x: torch.Tensor) -> qt.QuantizedResult:\n # We don't support training with this.\n raise NotImplementedError(\"Forward and training with HF EncodecModel not supported.\")\n\n def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:\n bandwidth_index = self.possible_num_codebooks.index(self.num_codebooks)\n bandwidth = self.model.config.target_bandwidths[bandwidth_index]\n res = self.model.encode(x, None, bandwidth)\n assert len(res[0]) == 1\n assert len(res[1]) == 1\n return res[0][0], res[1][0]\n\n def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):\n if scale is None:\n scales = [None] # type: ignore\n else:\n scales = scale # type: ignore\n res = self.model.decode(codes[None], scales)\n return res[0]\n\n def decode_latent(self, codes: torch.Tensor):\n \"\"\"Decode from the discrete codes to continuous latent space.\"\"\"\n return self.model.quantizer.decode(codes.transpose(0, 1))\n\n @property\n def channels(self) -> int:\n return self.model.config.audio_channels\n\n @property\n def frame_rate(self) -> float:\n hop_length = int(np.prod(self.model.config.upsampling_ratios))\n return self.sample_rate / hop_length\n\n @property\n def sample_rate(self) -> int:\n return self.model.config.sampling_rate\n\n @property\n def cardinality(self) -> int:\n return self.model.config.codebook_size\n\n @property\n def num_codebooks(self) -> int:\n return self._num_codebooks\n\n @property\n def total_codebooks(self) -> int:\n return max(self.possible_num_codebooks)\n\n def set_num_codebooks(self, n: int):\n \"\"\"Set the active number of codebooks used by the quantizer.\n \"\"\"\n if n not in self.possible_num_codebooks:\n raise ValueError(f\"Allowed values for num codebooks: {self.possible_num_codebooks}\")\n self._num_codebooks = n\n\n\nclass InterleaveStereoCompressionModel(CompressionModel):\n \"\"\"Wraps a CompressionModel to support stereo inputs. The wrapped model\n will be applied independently to the left and right channels, and both codebooks\n will be interleaved. If the wrapped model returns a representation `[B, K ,T]` per\n channel, then the output will be `[B, K * 2, T]` or `[B, K, T * 2]` depending on\n `per_timestep`.\n\n Args:\n model (CompressionModel): Compression model to wrap.\n per_timestep (bool): Whether to interleave on the timestep dimension\n or on the codebooks dimension.\n \"\"\"\n def __init__(self, model: CompressionModel, per_timestep: bool = False):\n super().__init__()\n self.model = model\n self.per_timestep = per_timestep\n assert self.model.channels == 1, \"Wrapped model is expected to be for monophonic audio\"\n\n @property\n def total_codebooks(self):\n return self.model.total_codebooks\n\n @property\n def num_codebooks(self):\n \"\"\"Active number of codebooks used by the quantizer.\n\n ..Warning:: this reports the number of codebooks after the interleaving\n of the codebooks!\n \"\"\"\n return self.model.num_codebooks if self.per_timestep else self.model.num_codebooks * 2\n\n def set_num_codebooks(self, n: int):\n \"\"\"Set the active number of codebooks used by the quantizer.\n\n ..Warning:: this sets the number of codebooks before the interleaving!\n \"\"\"\n self.model.set_num_codebooks(n)\n\n @property\n def num_virtual_steps(self) -> float:\n \"\"\"Return the number of virtual steps, e.g. one real step\n will be split into that many steps.\n \"\"\"\n return 2 if self.per_timestep else 1\n\n @property\n def frame_rate(self) -> float:\n return self.model.frame_rate * self.num_virtual_steps\n\n @property\n def sample_rate(self) -> int:\n return self.model.sample_rate\n\n @property\n def channels(self) -> int:\n return 2\n\n @property\n def cardinality(self):\n \"\"\"Cardinality of each codebook.\n \"\"\"\n return self.model.cardinality\n\n def forward(self, x: torch.Tensor) -> qt.QuantizedResult:\n raise NotImplementedError(\"Not supported, use encode and decode.\")\n\n def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:\n B, C, T = x.shape\n assert C == self.channels, f\"Expecting stereo audio but audio num channels is {C}\"\n\n indices_c0, scales_c0 = self.model.encode(x[:, 0, ...].unsqueeze(1))\n indices_c1, scales_c1 = self.model.encode(x[:, 1, ...].unsqueeze(1))\n indices = torch.stack([indices_c0, indices_c1], dim=0)\n scales: tp.Optional[torch.Tensor] = None\n if scales_c0 is not None and scales_c1 is not None:\n scales = torch.stack([scales_c0, scales_c1], dim=1)\n\n if self.per_timestep:\n indices = rearrange(indices, 'c b k t -> b k (t c)', c=2)\n else:\n indices = rearrange(indices, 'c b k t -> b (k c) t', c=2)\n\n return (indices, scales)\n\n def get_left_right_codes(self, codes: torch.Tensor) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n if self.per_timestep:\n codes = rearrange(codes, 'b k (t c) -> c b k t', c=2)\n else:\n codes = rearrange(codes, 'b (k c) t -> c b k t', c=2)\n return codes[0], codes[1]\n\n def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):\n B, K, T = codes.shape\n assert T % self.num_virtual_steps == 0, \"Provided codes' number of timesteps does not match\"\n assert K == self.num_codebooks, \"Provided codes' number of codebooks does not match\"\n\n scale_c0, scale_c1 = None, None\n if scale is not None:\n assert scale.size(0) == B and scale.size(1) == 2, f\"Scale has unexpected shape: {scale.shape}\"\n scale_c0 = scale[0, ...]\n scale_c1 = scale[1, ...]\n\n codes_c0, codes_c1 = self.get_left_right_codes(codes)\n audio_c0 = self.model.decode(codes_c0, scale_c0)\n audio_c1 = self.model.decode(codes_c1, scale_c1)\n return torch.cat([audio_c0, audio_c1], dim=1)\n\n def decode_latent(self, codes: torch.Tensor):\n \"\"\"Decode from the discrete codes to continuous latent space.\"\"\"\n raise NotImplementedError(\"Not supported by interleaved stereo wrapped models.\")\n" }, { "path": "audiocraft/models/flow_matching.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom dataclasses import dataclass\nfrom functools import partial\nimport logging\nimport math\nimport typing as tp\nimport torch\nfrom torch import nn\nfrom torchdiffeq import odeint # type: ignore\nfrom ..modules.streaming import StreamingModule\nfrom ..modules.transformer import create_norm_fn, StreamingTransformerLayer\nfrom ..modules.unet_transformer import UnetTransformer\nfrom ..modules.conditioners import (\n ConditionFuser,\n ClassifierFreeGuidanceDropout,\n AttributeDropout,\n ConditioningAttributes,\n JascoCondConst\n)\nfrom ..modules.jasco_conditioners import JascoConditioningProvider\nfrom ..modules.activations import get_activation_fn\n\nfrom .lm import ConditionTensors, init_layer\n\n\nlogger = logging.getLogger(__name__)\n\n\n@dataclass\nclass FMOutput:\n latents: torch.Tensor # [B, T, D]\n mask: torch.Tensor # [B, T]\n\n\nclass CFGTerm:\n \"\"\"\n Base class for Multi Source Classifier-Free Guidance (CFG) terms. This class represents a term in the CFG process,\n which is used to guide the generation process by adjusting the influence of different conditions.\n Attributes:\n conditions (dict): A dictionary of conditions that influence the generation process.\n weight (float): The weight of the CFG term, determining its influence on the generation.\n \"\"\"\n def __init__(self, conditions, weight):\n self.conditions = conditions\n self.weight = weight\n\n def drop_irrelevant_conds(self, conditions):\n \"\"\"\n Drops irrelevant conditions from the CFG term. This method should be implemented by subclasses.\n Args:\n conditions (dict): The conditions to be filtered.\n Raises:\n NotImplementedError: If the method is not implemented in a subclass.\n \"\"\"\n raise NotImplementedError(\"No base implementation for setting generation params.\")\n\n\nclass AllCFGTerm(CFGTerm):\n \"\"\"\n A CFG term that retains all conditions. This class does not drop any condition.\n \"\"\"\n def __init__(self, conditions, weight):\n super().__init__(conditions, weight)\n self.drop_irrelevant_conds()\n\n def drop_irrelevant_conds(self):\n pass\n\n\nclass NullCFGTerm(CFGTerm):\n \"\"\"\n A CFG term that drops all conditions, effectively nullifying their influence.\n \"\"\"\n def __init__(self, conditions, weight):\n super().__init__(conditions, weight)\n self.drop_irrelevant_conds()\n\n def drop_irrelevant_conds(self):\n \"\"\"\n Drops all conditions by applying a dropout with probability 1.0, effectively nullifying their influence.\n \"\"\"\n self.conditions = ClassifierFreeGuidanceDropout(p=1.0)(\n samples=self.conditions,\n cond_types=[\"wav\", \"text\", \"symbolic\"])\n\n\nclass TextCFGTerm(CFGTerm):\n \"\"\"\n A CFG term that selectively drops conditions based on specified dropout probabilities for different types\n of conditions, such as 'symbolic' and 'wav'.\n \"\"\"\n def __init__(self, conditions, weight, model_att_dropout):\n \"\"\"\n Initializes a TextCFGTerm with specified conditions, weight, and model attention dropout configuration.\n Args:\n conditions (dict): The conditions to be used in the CFG process.\n weight (float): The weight of the CFG term.\n model_att_dropout (object): The attribute dropouts used by the model.\n \"\"\"\n super().__init__(conditions, weight)\n if 'symbolic' in model_att_dropout.p:\n self.drop_symbolics = {k: 1.0 for k in model_att_dropout.p['symbolic'].keys()}\n else:\n self.drop_symbolics = {}\n if 'wav' in model_att_dropout.p:\n self.drop_wav = {k: 1.0 for k in model_att_dropout.p['wav'].keys()}\n else:\n self.drop_wav = {}\n self.drop_irrelevant_conds()\n\n def drop_irrelevant_conds(self):\n self.conditions = AttributeDropout({'symbolic': self.drop_symbolics,\n 'wav': self.drop_wav})(self.conditions) # drop temporal conds\n\n\nclass FlowMatchingModel(StreamingModule):\n \"\"\"\n A flow matching model inherits from StreamingModule.\n This model uses a transformer architecture to process and fuse conditions, applying learned embeddings and\n transformations and predicts multi-source guided vector fields.\n Attributes:\n condition_provider (JascoConditioningProvider): Provider for conditioning attributes.\n fuser (ConditionFuser): Fuser for combining multiple conditions.\n dim (int): Dimensionality of the model's main features.\n num_heads (int): Number of attention heads in the transformer.\n flow_dim (int): Dimensionality of the flow features.\n chords_dim (int): Dimensionality for chord embeddings, if used.\n drums_dim (int): Dimensionality for drums embeddings, if used.\n melody_dim (int): Dimensionality for melody embeddings, if used.\n hidden_scale (int): Scaling factor for the dimensionality of the feedforward network in the transformer.\n norm (str): Type of normalization to use ('layer_norm' or other supported types).\n norm_first (bool): Whether to apply normalization before other operations in the transformer layers.\n bias_proj (bool): Whether to include bias in the projection layers.\n weight_init (Optional[str]): Method for initializing weights.\n depthwise_init (Optional[str]): Method for initializing depthwise convolutional layers.\n zero_bias_init (bool): Whether to initialize biases to zero.\n cfg_dropout (float): Dropout rate for configuration settings.\n cfg_coef (float): Coefficient for configuration influence.\n attribute_dropout (Dict[str, Dict[str, float]]): Dropout rates for specific attributes.\n time_embedding_dim (int): Dimensionality of time embeddings.\n **kwargs: Additional keyword arguments for the transformer.\n Methods:\n __init__: Initializes the model with the specified attributes and configuration.\n \"\"\"\n def __init__(self, condition_provider: JascoConditioningProvider,\n fuser: ConditionFuser,\n dim: int = 128,\n num_heads: int = 8,\n flow_dim: int = 128,\n chords_dim: int = 0,\n drums_dim: int = 0,\n melody_dim: int = 0,\n hidden_scale: int = 4,\n norm: str = 'layer_norm',\n norm_first: bool = False,\n bias_proj: bool = True,\n weight_init: tp.Optional[str] = None,\n depthwise_init: tp.Optional[str] = None,\n zero_bias_init: bool = False,\n cfg_dropout: float = 0,\n cfg_coef: float = 1.0,\n attribute_dropout: tp.Dict[str, tp.Dict[str, float]] = {},\n time_embedding_dim: int = 128,\n **kwargs):\n super().__init__()\n self.cfg_coef = cfg_coef\n\n self.cfg_dropout = ClassifierFreeGuidanceDropout(p=cfg_dropout)\n self.att_dropout = AttributeDropout(p=attribute_dropout)\n self.condition_provider = condition_provider\n self.fuser = fuser\n self.dim = dim # transformer dim\n self.flow_dim = flow_dim\n self.chords_dim = chords_dim\n self.emb = nn.Linear(flow_dim + chords_dim + drums_dim + melody_dim, dim, bias=False)\n if 'activation' in kwargs:\n kwargs['activation'] = get_activation_fn(kwargs['activation'])\n\n self.transformer = UnetTransformer(\n d_model=dim, num_heads=num_heads, dim_feedforward=int(hidden_scale * dim),\n norm=norm, norm_first=norm_first,\n layer_class=StreamingTransformerLayer,\n **kwargs)\n self.out_norm: tp.Optional[nn.Module] = None\n if norm_first:\n self.out_norm = create_norm_fn(norm, dim)\n self.linear = nn.Linear(dim, flow_dim, bias=bias_proj)\n self._init_weights(weight_init, depthwise_init, zero_bias_init)\n self._fsdp: tp.Optional[nn.Module]\n self.__dict__['_fsdp'] = None\n\n # init time parameter embedding\n self.d_temb1 = time_embedding_dim\n self.d_temb2 = 4 * time_embedding_dim\n self.temb = nn.Module()\n self.temb.dense = nn.ModuleList([\n torch.nn.Linear(self.d_temb1,\n self.d_temb2),\n torch.nn.Linear(self.d_temb2,\n self.d_temb2),\n ])\n self.temb_proj = nn.Linear(self.d_temb2, dim)\n\n def _get_timestep_embedding(self, timesteps, embedding_dim):\n \"\"\"\n #######################################################################################################\n TAKEN FROM: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/model.py\n #######################################################################################################\n This matches the implementation in Denoising Diffusion Probabilistic Models:\n From Fairseq.\n Build sinusoidal embeddings.\n This matches the implementation in tensor2tensor, but differs slightly\n from the description in Section 3.5 of \"Attention Is All You Need\".\n \"\"\"\n assert len(timesteps.shape) == 1\n\n half_dim = embedding_dim // 2\n emb = math.log(10000) / (half_dim - 1)\n emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)\n emb = emb.to(device=timesteps.device)\n emb = timesteps.float()[:, None] * emb[None, :]\n emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)\n if embedding_dim % 2 == 1: # zero pad\n emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))\n return emb\n\n def _embed_time_parameter(self, t: torch.Tensor):\n \"\"\"\n #######################################################################################################\n TAKEN FROM: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/model.py\n #######################################################################################################\n \"\"\"\n temb = self._get_timestep_embedding(t.flatten(), self.d_temb1)\n temb = self.temb.dense[0](temb)\n temb = temb * torch.sigmoid(temb) # swish activation\n temb = self.temb.dense[1](temb)\n return temb\n\n def _init_weights(self, weight_init: tp.Optional[str], depthwise_init: tp.Optional[str], zero_bias_init: bool):\n \"\"\"Initialization of the transformer module weights.\n\n Args:\n weight_init (str, optional): Weight initialization strategy. See ``get_init_fn`` for valid options.\n depthwise_init (str, optional): Depthwise initialization strategy. The following options are valid:\n 'current' where the depth corresponds to the current layer index or 'global' where the total number\n of layer is used as depth. If not set, no depthwise initialization strategy is used.\n zero_bias_init (bool): Whether to initialize bias to zero or not.\n \"\"\"\n assert depthwise_init is None or depthwise_init in ['current', 'global']\n assert depthwise_init is None or weight_init is not None, \\\n \"If 'depthwise_init' is defined, a 'weight_init' method should be provided.\"\n assert not zero_bias_init or weight_init is not None, \\\n \"If 'zero_bias_init', a 'weight_init' method should be provided\"\n\n if weight_init is None:\n return\n\n init_layer(self.emb, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init)\n\n for layer_idx, tr_layer in enumerate(self.transformer.layers):\n depth = None\n if depthwise_init == 'current':\n depth = layer_idx + 1\n elif depthwise_init == 'global':\n depth = len(self.transformer.layers)\n init_fn = partial(init_layer, method=weight_init, init_depth=depth, zero_bias_init=zero_bias_init)\n tr_layer.apply(init_fn)\n\n init_layer(self.linear, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init)\n\n def _align_seq_length(self,\n cond: torch.Tensor,\n seq_len: int = 500):\n # trim if needed\n cond = cond[:, :seq_len, :]\n\n # pad if needed\n B, T, C = cond.shape\n if T < seq_len:\n cond = torch.cat((cond, torch.zeros((B, seq_len - T, C), dtype=cond.dtype, device=cond.device)), dim=1)\n\n return cond\n\n def forward(self,\n latents: torch.Tensor,\n t: torch.Tensor,\n conditions: tp.List[ConditioningAttributes],\n condition_tensors: tp.Optional[ConditionTensors] = None) -> torch.Tensor:\n \"\"\"Apply flow matching forward pass on latents and conditions.\n Given a tensor of noisy latents of shape [B, T, D] with D the flow dim and T the sequence steps,\n and a time parameter tensor t, return the vector field with shape [B, T, D].\n\n Args:\n latents (torch.Tensor): noisy latents.\n conditions (list of ConditioningAttributes): Conditions to use when modeling\n the given codes. Note that when evaluating multiple time with the same conditioning\n you should pre-compute those and pass them as `condition_tensors`.\n condition_tensors (dict[str, ConditionType], optional): Pre-computed conditioning\n tensors, see `conditions`.\n Returns:\n torch.Tensor: estimated vector field v_theta.\n \"\"\"\n assert condition_tensors is not None, \"FlowMatchingModel require pre-calculation of condition tensors\"\n assert not conditions, \"Shouldn't pass unprocessed conditions to FlowMatchingModel.\"\n\n B, T, D = latents.shape\n x = latents\n\n # concat temporal conditions on the feature dimension\n temporal_conds = JascoCondConst.ALL.value\n for cond in temporal_conds:\n if cond not in condition_tensors:\n continue\n c = self._align_seq_length(condition_tensors[cond][0], seq_len=T)\n x = torch.concat((x, c), dim=-1)\n\n # project to transformer dimension\n input_ = self.emb(x)\n\n input_, cross_attention_input = self.fuser(input_, condition_tensors)\n\n # embed time parameter\n t_embs = self._embed_time_parameter(t)\n\n # add it to cross_attention_input\n cross_attention_input = cross_attention_input + self.temb_proj(t_embs[:, None, :])\n\n out = self.transformer(input_, cross_attention_src=cross_attention_input)\n\n if self.out_norm:\n out = self.out_norm(out)\n v_theta = self.linear(out) # [B, T, D]\n\n # remove the prefix from the model outputs\n if len(self.fuser.fuse2cond['prepend']) > 0:\n v_theta = v_theta[:, :, -T:]\n\n return v_theta # [B, T, D]\n\n def _multi_source_cfg_preprocess(self,\n conditions: tp.List[ConditioningAttributes],\n cfg_coef_all: float,\n cfg_coef_txt: float,\n min_weight: float = 1e-6):\n \"\"\"\n Preprocesses the CFG terms for multi-source conditional generation.\n Args:\n conditions (list): A list of conditions to be applied.\n cfg_coef_all (float): The coefficient for all conditions.\n cfg_coef_txt (float): The coefficient for text conditions.\n min_weight (float): The minimal absolute weight for calculating a CFG term.\n Returns:\n tuple: A tuple containing condition_tensors and cfg_terms.\n condition_tensors is a dictionary or ConditionTensors object with tokenized conditions.\n cfg_terms is a list of CFGTerm objects with weights adjusted based on the coefficients.\n \"\"\"\n condition_tensors: tp.Optional[ConditionTensors]\n cfg_terms = []\n if conditions:\n # conditional terms\n cfg_terms = [AllCFGTerm(conditions=conditions, weight=cfg_coef_all),\n TextCFGTerm(conditions=conditions, weight=cfg_coef_txt,\n model_att_dropout=self.att_dropout)]\n\n # add null term\n cfg_terms.append(NullCFGTerm(conditions=conditions, weight=1 - sum([ct.weight for ct in cfg_terms])))\n\n # remove terms with negligible weight\n for ct in cfg_terms:\n if abs(ct.weight) < min_weight:\n cfg_terms.remove(ct)\n\n conds: tp.List[ConditioningAttributes] = sum([ct.conditions for ct in cfg_terms], [])\n tokenized = self.condition_provider.tokenize(conds)\n condition_tensors = self.condition_provider(tokenized)\n else:\n condition_tensors = {}\n\n return condition_tensors, cfg_terms\n\n def estimated_vector_field(self, z, t, condition_tensors=None, cfg_terms=[]):\n \"\"\"\n Estimates the vector field for the given latent variables and time parameter,\n conditioned on the provided conditions.\n Args:\n z (Tensor): The latent variables.\n t (float): The time variable.\n condition_tensors (ConditionTensors, optional): The condition tensors. Defaults to None.\n cfg_terms (list, optional): The list of CFG terms. Defaults to an empty list.\n Returns:\n Tensor: The estimated vector field.\n \"\"\"\n if len(cfg_terms) > 1:\n z = z.repeat(len(cfg_terms), 1, 1) # duplicate noisy latents for multi-source CFG\n v_thetas = self(latents=z, t=t, conditions=[], condition_tensors=condition_tensors)\n return self._multi_source_cfg_postprocess(v_thetas, cfg_terms)\n\n def _multi_source_cfg_postprocess(self, v_thetas, cfg_terms):\n \"\"\"\n Postprocesses the vector fields generated for each CFG term to combine them into a single vector field.\n Multi source guidance occurs here.\n Args:\n v_thetas (Tensor): The vector fields for each CFG term.\n cfg_terms (list): The CFG terms used.\n Returns:\n Tensor: The combined vector field.\n \"\"\"\n if len(cfg_terms) <= 1:\n return v_thetas\n v_theta_per_term = v_thetas.chunk(len(cfg_terms))\n return sum([ct.weight * term_vf for ct, term_vf in zip(cfg_terms, v_theta_per_term)])\n\n @torch.no_grad()\n def generate(self,\n prompt: tp.Optional[torch.Tensor] = None,\n conditions: tp.List[ConditioningAttributes] = [],\n num_samples: tp.Optional[int] = None,\n max_gen_len: int = 256,\n callback: tp.Optional[tp.Callable[[int, int], None]] = None,\n cfg_coef_all: float = 3.0,\n cfg_coef_txt: float = 1.0,\n euler: bool = False,\n euler_steps: int = 100,\n ode_rtol: float = 1e-5,\n ode_atol: float = 1e-5,\n ) -> torch.Tensor:\n \"\"\"\n Generate audio latents given a prompt or unconditionally. This method supports both Euler integration\n and adaptive ODE solving to generate sequences based on the specified conditions and configuration coefficients.\n\n Args:\n prompt (torch.Tensor, optional): Initial prompt to condition the generation. defaults to None\n conditions (List[ConditioningAttributes]): List of conditioning attributes - text, symbolic or audio.\n num_samples (int, optional): Number of samples to generate.\n If None, it is inferred from the number of conditions.\n max_gen_len (int): Maximum length of the generated sequence.\n callback (Callable[[int, int], None], optional): Callback function to monitor the generation process.\n cfg_coef_all (float): Coefficient for the fully conditional CFG term.\n cfg_coef_txt (float): Coefficient for text CFG term.\n euler (bool): If True, use Euler integration, otherwise use adaptive ODE solver.\n euler_steps (int): Number of Euler steps to perform if Euler integration is used.\n ode_rtol (float): ODE solver rtol threshold.\n ode_atol (float): ODE solver atol threshold.\n\n Returns:\n torch.Tensor: Generated latents, shaped as (num_samples, max_gen_len, feature_dim).\n \"\"\"\n\n assert not self.training, \"generation shouldn't be used in training mode.\"\n first_param = next(iter(self.parameters()))\n device = first_param.device\n\n # Checking all input shapes are consistent.\n possible_num_samples = []\n if num_samples is not None:\n possible_num_samples.append(num_samples)\n elif prompt is not None:\n possible_num_samples.append(prompt.shape[0])\n elif conditions:\n possible_num_samples.append(len(conditions))\n else:\n possible_num_samples.append(1)\n assert [x == possible_num_samples[0] for x in possible_num_samples], \"Inconsistent inputs shapes\"\n num_samples = possible_num_samples[0]\n\n condition_tensors, cfg_terms = self._multi_source_cfg_preprocess(conditions, cfg_coef_all, cfg_coef_txt)\n\n # flow matching inference\n B, T, D = num_samples, max_gen_len, self.flow_dim\n\n z_0 = torch.randn((B, T, D), device=device)\n\n if euler:\n # vanilla Euler intergration\n dt = (1 / euler_steps)\n z = z_0\n t = torch.zeros((1, ), device=device)\n for _ in range(euler_steps):\n v_theta = self.estimated_vector_field(z, t,\n condition_tensors=condition_tensors,\n cfg_terms=cfg_terms)\n z = z + dt * v_theta\n t = t + dt\n z_1 = z\n else:\n # solve with dynamic ode integrator (dopri5)\n t = torch.tensor([0, 1.0 - 1e-5], device=device)\n num_evals = 0\n\n # define ode vector field function\n def inner_ode_func(t, z):\n nonlocal num_evals\n num_evals += 1\n if callback is not None:\n ESTIMATED_ODE_SOLVER_STEPS = 300\n callback(num_evals, ESTIMATED_ODE_SOLVER_STEPS)\n return self.estimated_vector_field(z, t,\n condition_tensors=condition_tensors,\n cfg_terms=cfg_terms)\n\n ode_opts: dict = {\"options\": {}}\n z = odeint(\n inner_ode_func,\n z_0,\n t,\n **{\"atol\": ode_atol, \"rtol\": ode_rtol, **ode_opts},\n )\n logger.info(\"Generated in %d steps\", num_evals)\n z_1 = z[-1]\n\n return z_1\n" }, { "path": "audiocraft/models/genmodel.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nBase implementation for audio generative models. This base implementation\ncombines all the required components to run inference with pretrained audio\ngenerative models. It can be easily inherited by downstream model classes to\nprovide easy access to the generation API.\n\"\"\"\n\nfrom abc import ABC, abstractmethod\nimport typing as tp\n\nimport omegaconf\nimport torch\n\nfrom .encodec import CompressionModel\nfrom .lm import LMModel\nfrom .builders import get_wrapped_compression_model\nfrom ..data.audio_utils import convert_audio\nfrom ..modules.conditioners import ConditioningAttributes\nfrom ..utils.autocast import TorchAutocast\n\n\nclass BaseGenModel(ABC):\n \"\"\"Base generative model with convenient generation API.\n\n Args:\n name (str): name of the model.\n compression_model (CompressionModel): Compression model\n used to map audio to invertible discrete representations.\n lm (LMModel): Language model over discrete representations.\n max_duration (float, optional): maximum duration the model can produce,\n otherwise, inferred from the training params.\n \"\"\"\n def __init__(self, name: str, compression_model: CompressionModel, lm: LMModel,\n max_duration: tp.Optional[float] = None):\n self.name = name\n self.compression_model = compression_model\n self.lm = lm\n self.cfg: tp.Optional[omegaconf.DictConfig] = None\n # Just to be safe, let's put everything in eval mode.\n self.compression_model.eval()\n self.lm.eval()\n\n if hasattr(lm, 'cfg'):\n cfg = lm.cfg\n assert isinstance(cfg, omegaconf.DictConfig)\n self.cfg = cfg\n\n if self.cfg is not None:\n self.compression_model = get_wrapped_compression_model(self.compression_model, self.cfg)\n\n if max_duration is None:\n if self.cfg is not None:\n max_duration = lm.cfg.dataset.segment_duration # type: ignore\n else:\n raise ValueError(\"You must provide max_duration when building directly your GenModel\")\n assert max_duration is not None\n\n self.max_duration: float = max_duration\n self.duration = self.max_duration\n\n # self.extend_stride is the length of audio extension when generating samples longer\n # than self.max_duration. NOTE: the derived class must set self.extend_stride to a\n # positive float value when generating with self.duration > self.max_duration.\n self.extend_stride: tp.Optional[float] = None\n self.device = next(iter(lm.parameters())).device\n self.generation_params: dict = {}\n self._progress_callback: tp.Optional[tp.Callable[[int, int], None]] = None\n if self.device.type == 'cpu':\n self.autocast = TorchAutocast(enabled=False)\n else:\n self.autocast = TorchAutocast(\n enabled=True, device_type=self.device.type, dtype=torch.float16)\n\n @property\n def frame_rate(self) -> float:\n \"\"\"Roughly the number of AR steps per seconds.\"\"\"\n return self.compression_model.frame_rate\n\n @property\n def sample_rate(self) -> int:\n \"\"\"Sample rate of the generated audio.\"\"\"\n return self.compression_model.sample_rate\n\n @property\n def audio_channels(self) -> int:\n \"\"\"Audio channels of the generated audio.\"\"\"\n return self.compression_model.channels\n\n def set_custom_progress_callback(self, progress_callback: tp.Optional[tp.Callable[[int, int], None]] = None):\n \"\"\"Override the default progress callback.\"\"\"\n self._progress_callback = progress_callback\n\n @abstractmethod\n def set_generation_params(self, *args, **kwargs):\n \"\"\"Set the generation parameters.\"\"\"\n raise NotImplementedError(\"No base implementation for setting generation params.\")\n\n @staticmethod\n @abstractmethod\n def get_pretrained(name: str, device=None):\n raise NotImplementedError(\"No base implementation for getting pretrained model\")\n\n @torch.no_grad()\n def _prepare_tokens_and_attributes(\n self,\n descriptions: tp.Sequence[tp.Optional[str]],\n prompt: tp.Optional[torch.Tensor],\n ) -> tp.Tuple[tp.List[ConditioningAttributes], tp.Optional[torch.Tensor]]:\n \"\"\"Prepare model inputs.\n\n Args:\n descriptions (list of str): A list of strings used as text conditioning.\n prompt (torch.Tensor): A batch of waveforms used for continuation.\n \"\"\"\n attributes = [\n ConditioningAttributes(text={'description': description})\n for description in descriptions]\n\n if prompt is not None:\n if descriptions is not None:\n assert len(descriptions) == len(prompt), \"Prompt and nb. descriptions doesn't match\"\n prompt = prompt.to(self.device)\n prompt_tokens, scale = self.compression_model.encode(prompt)\n assert scale is None\n else:\n prompt_tokens = None\n return attributes, prompt_tokens\n\n def generate_unconditional(self, num_samples: int, progress: bool = False,\n return_tokens: bool = False) -> tp.Union[torch.Tensor,\n tp.Tuple[torch.Tensor, torch.Tensor]]:\n \"\"\"Generate samples in an unconditional manner.\n\n Args:\n num_samples (int): Number of samples to be generated.\n progress (bool, optional): Flag to display progress of the generation process. Defaults to False.\n \"\"\"\n descriptions: tp.List[tp.Optional[str]] = [None] * num_samples\n attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions, None)\n tokens = self._generate_tokens(attributes, prompt_tokens, progress)\n if return_tokens:\n return self.generate_audio(tokens), tokens\n return self.generate_audio(tokens)\n\n def generate(self, descriptions: tp.List[str], progress: bool = False, return_tokens: bool = False) \\\n -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]:\n \"\"\"Generate samples conditioned on text.\n\n Args:\n descriptions (list of str): A list of strings used as text conditioning.\n progress (bool, optional): Flag to display progress of the generation process. Defaults to False.\n \"\"\"\n attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions, None)\n assert prompt_tokens is None\n tokens = self._generate_tokens(attributes, prompt_tokens, progress)\n if return_tokens:\n return self.generate_audio(tokens), tokens\n return self.generate_audio(tokens)\n\n def generate_continuation(self, prompt: torch.Tensor, prompt_sample_rate: int,\n descriptions: tp.Optional[tp.List[tp.Optional[str]]] = None,\n progress: bool = False, return_tokens: bool = False) \\\n -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]:\n \"\"\"Generate samples conditioned on audio prompts and an optional text description.\n\n Args:\n prompt (torch.Tensor): A batch of waveforms used for continuation.\n Prompt should be [B, C, T], or [C, T] if only one sample is generated.\n prompt_sample_rate (int): Sampling rate of the given audio waveforms.\n descriptions (list of str, optional): A list of strings used as text conditioning. Defaults to None.\n progress (bool, optional): Flag to display progress of the generation process. Defaults to False.\n \"\"\"\n if prompt.dim() == 2:\n prompt = prompt[None]\n if prompt.dim() != 3:\n raise ValueError(\"prompt should have 3 dimensions: [B, C, T] (C = 1).\")\n prompt = convert_audio(prompt, prompt_sample_rate, self.sample_rate, self.audio_channels)\n if descriptions is None:\n descriptions = [None] * len(prompt)\n attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions, prompt)\n assert prompt_tokens is not None\n tokens = self._generate_tokens(attributes, prompt_tokens, progress)\n if return_tokens:\n return self.generate_audio(tokens), tokens\n return self.generate_audio(tokens)\n\n def _generate_tokens(self, attributes: tp.List[ConditioningAttributes],\n prompt_tokens: tp.Optional[torch.Tensor], progress: bool = False) -> torch.Tensor:\n \"\"\"Generate discrete audio tokens given audio prompt and/or conditions.\n\n Args:\n attributes (list of ConditioningAttributes): Conditions used for generation (here text).\n prompt_tokens (torch.Tensor, optional): Audio prompt used for continuation.\n progress (bool, optional): Flag to display progress of the generation process. Defaults to False.\n Returns:\n torch.Tensor: Generated audio, of shape [B, C, T], T is defined by the generation params.\n \"\"\"\n total_gen_len = int(self.duration * self.frame_rate)\n max_prompt_len = int(min(self.duration, self.max_duration) * self.frame_rate)\n current_gen_offset: int = 0\n\n def _progress_callback(generated_tokens: int, tokens_to_generate: int):\n generated_tokens += current_gen_offset\n if self._progress_callback is not None:\n # Note that total_gen_len might be quite wrong depending on the\n # codebook pattern used, but with delay it is almost accurate.\n self._progress_callback(generated_tokens, tokens_to_generate)\n else:\n print(f'{generated_tokens: 6d} / {tokens_to_generate: 6d}', end='\\r')\n\n if prompt_tokens is not None:\n assert max_prompt_len >= prompt_tokens.shape[-1], \\\n \"Prompt is longer than audio to generate\"\n\n callback = None\n if progress:\n callback = _progress_callback\n\n if self.duration <= self.max_duration:\n # generate by sampling from LM, simple case.\n with self.autocast:\n gen_tokens = self.lm.generate(\n prompt_tokens, attributes,\n callback=callback, max_gen_len=total_gen_len, **self.generation_params)\n\n else:\n assert self.extend_stride is not None, \"Stride should be defined to generate beyond max_duration\"\n assert self.extend_stride < self.max_duration, \"Cannot stride by more than max generation duration.\"\n all_tokens = []\n if prompt_tokens is None:\n prompt_length = 0\n else:\n all_tokens.append(prompt_tokens)\n prompt_length = prompt_tokens.shape[-1]\n\n stride_tokens = int(self.frame_rate * self.extend_stride)\n while current_gen_offset + prompt_length < total_gen_len:\n time_offset = current_gen_offset / self.frame_rate\n chunk_duration = min(self.duration - time_offset, self.max_duration)\n max_gen_len = int(chunk_duration * self.frame_rate)\n with self.autocast:\n gen_tokens = self.lm.generate(\n prompt_tokens, attributes,\n callback=callback, max_gen_len=max_gen_len, **self.generation_params)\n if prompt_tokens is None:\n all_tokens.append(gen_tokens)\n else:\n all_tokens.append(gen_tokens[:, :, prompt_tokens.shape[-1]:])\n prompt_tokens = gen_tokens[:, :, stride_tokens:]\n prompt_length = prompt_tokens.shape[-1]\n current_gen_offset += stride_tokens\n\n gen_tokens = torch.cat(all_tokens, dim=-1)\n return gen_tokens\n\n def generate_audio(self, gen_tokens: torch.Tensor) -> torch.Tensor:\n \"\"\"Generate Audio from tokens.\"\"\"\n assert gen_tokens.dim() == 3\n with torch.no_grad():\n gen_audio = self.compression_model.decode(gen_tokens, None)\n return gen_audio\n" }, { "path": "audiocraft/models/jasco.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nMain model for using JASCO. This will combine all the required components\nand provide easy access to the generation API.\n\"\"\"\nimport os\nimport math\nimport pickle\nimport torch\nimport typing as tp\n\nfrom audiocraft.utils.utils import construct_frame_chords\nfrom .genmodel import BaseGenModel\nfrom .loaders import load_compression_model, load_jasco_model\nfrom ..data.audio_utils import convert_audio\nfrom ..modules.conditioners import WavCondition, ConditioningAttributes, SymbolicCondition, JascoCondConst\n\n\nclass JASCO(BaseGenModel):\n \"\"\"JASCO main model with convenient generation API.\n Args:\n chords_mapping_path: path to chords to index mapping pickle\n kwargs - See MusicGen class.\n \"\"\"\n def __init__(self, chords_mapping_path='assets/chord_to_index_mapping.pkl', **kwargs):\n super().__init__(**kwargs)\n # JASCO operates over a fixed sequence length defined in it's config.\n self.duration = self.lm.cfg.dataset.segment_duration\n\n # load chord2index mapping of Chordino (https://github.com/ohollo/chord-extractor)\n assert os.path.exists(chords_mapping_path)\n self.chords_mapping = pickle.load(open(chords_mapping_path, \"rb\"))\n\n # set generation parameters\n self.set_generation_params()\n\n @staticmethod\n def get_pretrained(name: str = 'facebook/jasco-chords-drums-400M', device=None,\n chords_mapping_path='assets/chord_to_index_mapping.pkl'):\n \"\"\"Return pretrained model, we provide 2 models:\n 1. facebook/jasco-chords-drums-400M: 10s music generation conditioned on\n text, chords and drums, 400M parameters.\n 2. facebook/jasco-chords-drums-1B: 10s music generation conditioned on\n text, chords and drums, 1B parameters.\n \"\"\"\n if device is None:\n if torch.cuda.device_count():\n device = 'cuda'\n else:\n device = 'cpu'\n\n compression_model = load_compression_model(name, device=device)\n lm = load_jasco_model(name, compression_model, device=device)\n\n kwargs = {'name': name,\n 'compression_model': compression_model,\n 'lm': lm,\n 'chords_mapping_path': chords_mapping_path}\n return JASCO(**kwargs)\n\n def set_generation_params(self,\n cfg_coef_all: float = 5.0,\n cfg_coef_txt: float = 0.0,\n **kwargs):\n \"\"\"Set the generation parameters for JASCO.\n\n Args:\n cfg_coef_all (float, optional): Coefficient used in multi-source classifier free guidance -\n all conditions term. Defaults to 5.0.\n cfg_coef_txt (float, optional): Coefficient used in multi-source classifier free guidance -\n text condition term. Defaults to 0.0.\n\n \"\"\"\n self.generation_params = {\n 'cfg_coef_all': cfg_coef_all,\n 'cfg_coef_txt': cfg_coef_txt\n }\n self.generation_params.update(kwargs)\n\n def _unnormalized_latents(self, latents: torch.Tensor) -> torch.Tensor:\n \"\"\"Unnormalize latents, shifting back to EnCodec's expected mean, std\"\"\"\n assert self.cfg is not None\n scaled = latents * self.cfg.compression_model_latent_std\n return scaled + self.cfg.compression_model_latent_mean\n\n def generate_audio(self, gen_latents: torch.Tensor) -> torch.Tensor:\n \"\"\"Decode audio from generated latents\"\"\"\n assert gen_latents.dim() == 3 # [B, T, C]\n\n # unnormalize latents\n gen_latents = self._unnormalized_latents(gen_latents)\n return self.compression_model.model.decoder(gen_latents.permute(0, 2, 1))\n\n def _generate_tokens(self, attributes: tp.List[ConditioningAttributes],\n prompt_tokens: tp.Optional[torch.Tensor], progress: bool = False) -> torch.Tensor:\n \"\"\"Generate continuous audio latents given conditions.\n\n Args:\n attributes (list of ConditioningAttributes): Conditions used for generation (here text).\n prompt_tokens (torch.Tensor, optional): Audio prompt used for continuation.\n progress (bool, optional): Flag to display progress of the generation process. Defaults to False.\n Returns:\n torch.Tensor: Generated latents, of shape [B, T, C].\n \"\"\"\n total_gen_len = int(self.duration * self.frame_rate)\n max_prompt_len = int(min(self.duration, self.max_duration) * self.frame_rate)\n\n def _progress_callback(ode_steps: int, max_ode_steps: int):\n ode_steps += 1\n if self._progress_callback is not None:\n # Note that total_gen_len might be quite wrong depending on the\n # codebook pattern used, but with delay it is almost accurate.\n self._progress_callback(ode_steps, max_ode_steps)\n else:\n print(f'{ode_steps: 6d} / {max_ode_steps: 6d}', end='\\r')\n\n if prompt_tokens is not None:\n assert max_prompt_len >= prompt_tokens.shape[-1], \\\n \"Prompt is longer than audio to generate\"\n\n callback = None\n if progress:\n callback = _progress_callback\n\n # generate by sampling from the LM\n with self.autocast:\n total_gen_len = math.ceil(self.duration * self.compression_model.frame_rate)\n return self.lm.generate(\n prompt_tokens, attributes,\n callback=callback, max_gen_len=total_gen_len, **self.generation_params)\n\n def _prepare_chord_conditions(\n self,\n attributes: tp.List[ConditioningAttributes],\n chords: tp.Optional[tp.List[tp.Tuple[str, float]]],\n ) -> tp.List[ConditioningAttributes]:\n \"\"\"\n Prepares chord conditions by translating symbolic chord progressions into a sequence of integers.\n This method updates the ConditioningAttributes with per-frame chords information.\n Args:\n attributes (List[ConditioningAttributes]):\n The initial attributes and optional tensor data.\n chords (List[Tuple[str, float]]):\n A list of tuples containing chord labels and their start times.\n Returns:\n List[ConditioningAttributes]:\n The updated attributes with frame chords integrated, alongside the original optional tensor data.\n \"\"\"\n if chords is None or chords == []:\n for att in attributes:\n att.symbolic[JascoCondConst.CRD.value] = SymbolicCondition(frame_chords=-1 *\n torch.ones(1, dtype=torch.int32))\n return attributes\n\n # flip from (chord, start_time) to (start_time, chord)\n chords_time_first: tp.List[tuple[float, str]] = [(item[1], item[0]) for item in chords]\n\n # translate symbolic chord progression into a sequence of ints\n frame_chords = construct_frame_chords(min_timestamp=0,\n chord_changes=chords_time_first,\n mapping_dict=self.chords_mapping,\n prev_chord='',\n frame_rate=self.compression_model.frame_rate,\n segment_duration=self.duration)\n # update the attribute objects\n for att in attributes:\n att.symbolic[JascoCondConst.CRD.value] = SymbolicCondition(frame_chords=torch.tensor(frame_chords))\n return attributes\n\n @torch.no_grad()\n def _prepare_drums_conditions(self,\n attributes:\n tp.List[ConditioningAttributes],\n drums_wav: tp.Optional[torch.Tensor],\n ):\n # prepare drums cond\n for attr in attributes:\n if drums_wav is None:\n attr.wav[JascoCondConst.DRM.value] = WavCondition(\n torch.zeros((1, 1, 1), device=self.device),\n torch.tensor([0], device=self.device),\n sample_rate=[self.sample_rate],\n path=[None])\n else:\n if JascoCondConst.DRM.value not in self.lm.condition_provider.conditioners:\n raise RuntimeError(\"This model doesn't support drums conditioning. \")\n\n expected_length = self.lm.cfg.dataset.segment_duration * self.sample_rate\n # trim if needed\n drums_wav = drums_wav[..., :expected_length]\n\n # pad if needed\n if drums_wav.shape[-1] < expected_length:\n diff = expected_length - drums_wav.shape[-1]\n diff_zeros = torch.zeros((drums_wav.shape[0], drums_wav.shape[1], diff),\n device=drums_wav.device, dtype=drums_wav.dtype)\n drums_wav = torch.cat((drums_wav, diff_zeros), dim=-1)\n\n attr.wav[JascoCondConst.DRM.value] = WavCondition(\n drums_wav.to(device=self.device),\n torch.tensor([drums_wav.shape[-1]], device=self.device),\n sample_rate=[self.sample_rate],\n path=[None],\n )\n\n return attributes\n\n @torch.no_grad()\n def _prepare_melody_conditions(\n self,\n attributes: tp.List[ConditioningAttributes],\n melody: tp.Optional[torch.Tensor],\n expected_length: int,\n melody_bins: int = 53,\n ) -> tp.List[ConditioningAttributes]:\n \"\"\"\n Prepares melody conditions by subtituting with pre-computed salience matrix.\n This method updates the ConditioningAttributes with per-frame chords information.\n Args:\n attributes (List[ConditioningAttributes]):\n The initial attributes and optional tensor data.\n chords (List[Tuple[str, float]]):\n A list of tuples containing chord labels and their start times.\n Returns:\n List[ConditioningAttributes]:\n The updated attributes with frame chords integrated, alongside the original optional tensor data.\n \"\"\"\n for attr in attributes:\n if melody is None:\n melody = torch.zeros((melody_bins, expected_length))\n attr.symbolic[JascoCondConst.MLD.value] = SymbolicCondition(melody=melody)\n return attributes\n\n @torch.no_grad()\n def _prepare_temporal_conditions(\n self,\n attributes: tp.List[ConditioningAttributes],\n expected_length: int,\n chords: tp.Optional[tp.List[tp.Tuple[str, float]]],\n drums_wav: tp.Optional[torch.Tensor],\n salience_matrix: tp.Optional[torch.Tensor],\n melody_bins: int = 53,\n ) -> tp.List[ConditioningAttributes]:\n \"\"\"\n Prepares temporal conditions (chords, drums).\n Args:\n attributes (List[ConditioningAttributes]): The initial attributes and optional tensor data.\n expected_length (int): The expected number of generated frames.\n chords (List[Tuple[str, float]]): A list of tuples containing chord labels and their start times.\n drums_wav (List[Tuple[str, float]]): tensor of extracted drums wav.\n salience_matrix (List[Tuple[str, float]]): melody matrix.\n melody_bins (int): number of melody bins the model was trained with, only relevant if trained with melody.\n Returns:\n List[ConditioningAttributes]:\n The updated attributes after processing chord conditions.\n \"\"\"\n attributes = self._prepare_chord_conditions(attributes=attributes, chords=chords)\n attributes = self._prepare_drums_conditions(attributes=attributes, drums_wav=drums_wav)\n attributes = self._prepare_melody_conditions(attributes=attributes, melody=salience_matrix,\n expected_length=expected_length, melody_bins=melody_bins)\n return attributes\n\n @torch.no_grad()\n def generate_music(\n self, descriptions: tp.List[str],\n drums_wav: tp.Optional[torch.Tensor] = None,\n drums_sample_rate: int = 32000,\n chords: tp.Optional[tp.List[tp.Tuple[str, float]]] = None,\n melody_salience_matrix: tp.Optional[torch.Tensor] = None,\n iopaint_wav: tp.Optional[torch.Tensor] = None,\n segment_duration: float = 10.0,\n frame_rate: float = 50.0,\n melody_bins: int = 53,\n progress: bool = False, return_latents: bool = False) \\\n -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]:\n \"\"\"Generate samples conditioned on text and temporal conditions (chords, melody, drums).\n\n Args:\n descriptions (list of str): A list of strings used as text conditioning.\n chords (list of (str, float) tuples): Chord progression represented as chord, start time (sec), e.g.:\n [(\"C\", 0.0), (\"F\", 4.0), (\"G\", 6.0), (\"C\", 8.0)]\n melody_salience_matrix (torch.Tensor, optional): melody saliency matrix. Default=None.\n iopaint_wav (torch.Tensor, optional): in/out=painting waveform. Default=None.\n segment_duration (float): the segment duration the model was trained on. Default=None.\n frame_rate (float): the frame_rate model was trained on. Default=None.\n melody_bins (int): number of melody bins the model was trained with, only relevant if trained with melody.\n progress (bool, optional): Flag to display progress of the generation process. Defaults to False.\n \"\"\"\n\n if drums_wav is not None:\n if drums_wav.dim() == 2:\n drums_wav = drums_wav[None]\n assert drums_wav.dim() == 3, \"drums wav should have a shape [B, C, T].\"\n drums_wav = convert_audio(drums_wav, drums_sample_rate, self.sample_rate, self.audio_channels)\n\n cond_attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions=descriptions,\n prompt=None)\n\n # prepare temporal conds (symbolic / audio)\n jasco_attributes = self._prepare_temporal_conditions(attributes=cond_attributes,\n expected_length=int(segment_duration * frame_rate),\n chords=chords,\n drums_wav=drums_wav,\n salience_matrix=melody_salience_matrix,\n melody_bins=melody_bins)\n assert prompt_tokens is None\n tokens = self._generate_tokens(jasco_attributes, prompt_tokens, progress)\n if return_latents:\n return self.generate_audio(tokens), tokens\n return self.generate_audio(tokens)\n\n @torch.no_grad()\n def generate(self, descriptions: tp.List[str], progress: bool = False, return_latents: bool = False) \\\n -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]:\n \"\"\"Generate samples conditioned on text.\n\n Args:\n descriptions (list of str): A list of strings used as text conditioning.\n progress (bool, optional): Flag to display progress of the generation process. Defaults to False.\n \"\"\"\n return self.generate_music(descriptions=descriptions, progress=progress, return_latents=return_latents)\n" }, { "path": "audiocraft/models/lm.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom dataclasses import dataclass\nfrom functools import partial\nimport logging\nimport math\nimport typing as tp\n\nimport torch\nfrom torch import nn\n\nfrom ..utils import utils\nfrom ..modules.streaming import StreamingModule, State\nfrom ..modules.transformer import StreamingTransformer, create_norm_fn\nfrom ..modules.conditioners import (\n ConditionFuser,\n ClassifierFreeGuidanceDropout,\n AttributeDropout,\n ConditioningProvider,\n ConditioningAttributes,\n ConditionType,\n _drop_description_condition\n)\nfrom ..modules.codebooks_patterns import CodebooksPatternProvider\nfrom ..modules.activations import get_activation_fn\n\n\nlogger = logging.getLogger(__name__)\nConditionTensors = tp.Dict[str, ConditionType]\nCFGConditions = tp.Union[ConditionTensors, tp.Tuple[ConditionTensors, ConditionTensors]]\n\n\ndef get_init_fn(method: str, input_dim: int, init_depth: tp.Optional[int] = None):\n \"\"\"LM layer initialization.\n Inspired from xlformers: https://github.com/fairinternal/xlformers\n\n Args:\n method (str): Method name for init function. Valid options are:\n 'gaussian', 'uniform'.\n input_dim (int): Input dimension of the initialized module.\n init_depth (int, optional): Optional init depth value used to rescale\n the standard deviation if defined.\n \"\"\"\n # Compute std\n std = 1 / math.sqrt(input_dim)\n # Rescale with depth\n if init_depth is not None:\n std = std / math.sqrt(2 * init_depth)\n\n if method == 'gaussian':\n return partial(\n torch.nn.init.trunc_normal_, mean=0.0, std=std, a=-3 * std, b=3 * std\n )\n elif method == 'uniform':\n bound = math.sqrt(3) * std # ensure the standard deviation is `std`\n return partial(torch.nn.init.uniform_, a=-bound, b=bound)\n else:\n raise ValueError(\"Unsupported layer initialization method\")\n\n\ndef init_layer(m: nn.Module,\n method: str,\n init_depth: tp.Optional[int] = None,\n zero_bias_init: bool = False):\n \"\"\"Wrapper around ``get_init_fn`` for proper initialization of LM modules.\n\n Args:\n m (nn.Module): Module to initialize.\n method (str): Method name for the init function.\n init_depth (int, optional): Optional init depth value used to rescale\n the standard deviation if defined.\n zero_bias_init (bool): Whether to initialize the bias to 0 or not.\n \"\"\"\n if isinstance(m, nn.Linear):\n init_fn = get_init_fn(method, m.in_features, init_depth=init_depth)\n if m.weight.device.type == 'cpu' and m.weight.dtype == torch.float16:\n weight = m.weight.float()\n init_fn(weight)\n m.weight.data[:] = weight.half()\n else:\n init_fn(m.weight)\n if zero_bias_init and m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Embedding):\n init_fn = get_init_fn(method, m.embedding_dim, init_depth=None)\n if m.weight.device.type == 'cpu' and m.weight.dtype == torch.float16:\n weight = m.weight.float()\n init_fn(weight)\n m.weight.data[:] = weight.half()\n else:\n init_fn(m.weight)\n\n\nclass ScaledEmbedding(nn.Embedding):\n \"\"\"Boost learning rate for embeddings (with `scale`).\n \"\"\"\n def __init__(self, *args, lr=None, **kwargs):\n super().__init__(*args, **kwargs)\n self.lr = lr\n\n def make_optim_group(self):\n group = {\"params\": list(self.parameters())}\n if self.lr is not None:\n group[\"lr\"] = self.lr\n return group\n\n\n@dataclass\nclass LMOutput:\n # The logits are already re-aligned with the input codes\n # hence no extra shift is required, e.g. when computing CE\n logits: torch.Tensor # [B, K, T, card]\n mask: torch.Tensor # [B, K, T]\n\n\nclass LMModel(StreamingModule):\n \"\"\"Transformer-based language model on multiple streams of codes.\n\n Args:\n pattern_provider (CodebooksPatternProvider): Pattern provider for codebook interleaving.\n condition_provider (MusicConditioningProvider): Conditioning provider from metadata.\n fuser (ConditionFuser): Fuser handling the fusing of conditions with language model input.\n n_q (int): Number of parallel streams to model.\n card (int): Cardinality, vocabulary size.\n dim (int): Dimension of the transformer encoder.\n num_heads (int): Number of heads for the transformer encoder.\n hidden_scale (int): Scale for hidden feed forward dimension of the transformer encoder.\n norm (str): Normalization method.\n norm_first (bool): Use pre-norm instead of post-norm.\n emb_lr (float, optional): Embedding-specific learning rate.\n bias_proj (bool): Use bias for output projections.\n weight_init (str, optional): Method for weight initialization.\n depthwise_init (str, optional): Method for depthwise weight initialization.\n zero_bias_init (bool): If true and bias in Linears, initialize bias to zeros.\n cfg_dropout (float): Classifier-free guidance dropout.\n cfg_coef (float): Classifier-free guidance coefficient.\n attribute_dropout (dict): Attribute dropout probabilities.\n two_step_cfg (bool): Whether to run classifier free-guidance with 2 distinct steps.\n **kwargs: Additional parameters for the transformer encoder.\n \"\"\"\n def __init__(self, pattern_provider: CodebooksPatternProvider, condition_provider: ConditioningProvider,\n fuser: ConditionFuser, n_q: int = 8, card: int = 1024, dim: int = 128, num_heads: int = 8,\n hidden_scale: int = 4, norm: str = 'layer_norm', norm_first: bool = False,\n emb_lr: tp.Optional[float] = None, bias_proj: bool = True,\n weight_init: tp.Optional[str] = None, depthwise_init: tp.Optional[str] = None,\n zero_bias_init: bool = False, cfg_dropout: float = 0, cfg_coef: float = 1.0,\n attribute_dropout: tp.Dict[str, tp.Dict[str, float]] = {}, two_step_cfg: bool = False,\n **kwargs):\n super().__init__()\n self.cfg_coef = cfg_coef\n self.cfg_dropout = ClassifierFreeGuidanceDropout(p=cfg_dropout)\n self.att_dropout = AttributeDropout(p=attribute_dropout)\n self.condition_provider = condition_provider\n self.fuser = fuser\n self.card = card\n embed_dim = self.card + 1\n self.n_q = n_q\n self.dim = dim\n self.pattern_provider = pattern_provider\n self.two_step_cfg = two_step_cfg\n self.emb = nn.ModuleList([ScaledEmbedding(embed_dim, dim, lr=emb_lr) for _ in range(n_q)])\n if 'activation' in kwargs:\n kwargs['activation'] = get_activation_fn(kwargs['activation'])\n self.transformer = StreamingTransformer(\n d_model=dim, num_heads=num_heads, dim_feedforward=int(hidden_scale * dim),\n norm=norm, norm_first=norm_first, **kwargs)\n self.out_norm: tp.Optional[nn.Module] = None\n if norm_first:\n self.out_norm = create_norm_fn(norm, dim)\n self.linears = nn.ModuleList([nn.Linear(dim, self.card, bias=bias_proj) for _ in range(n_q)])\n self._init_weights(weight_init, depthwise_init, zero_bias_init)\n self._fsdp: tp.Optional[nn.Module]\n self.__dict__['_fsdp'] = None\n\n def _init_weights(self, weight_init: tp.Optional[str], depthwise_init: tp.Optional[str], zero_bias_init: bool):\n \"\"\"Initialization of the transformer module weights.\n\n Args:\n weight_init (str, optional): Weight initialization strategy. See ``get_init_fn`` for valid options.\n depthwise_init (str, optional): Depthwise initialization strategy. The following options are valid:\n 'current' where the depth corresponds to the current layer index or 'global' where the total number\n of layer is used as depth. If not set, no depthwise initialization strategy is used.\n zero_bias_init (bool): Whether to initialize bias to zero or not.\n \"\"\"\n assert depthwise_init is None or depthwise_init in ['current', 'global']\n assert depthwise_init is None or weight_init is not None, \\\n \"If 'depthwise_init' is defined, a 'weight_init' method should be provided.\"\n assert not zero_bias_init or weight_init is not None, \\\n \"If 'zero_bias_init', a 'weight_init' method should be provided\"\n\n if weight_init is None:\n return\n\n for emb_layer in self.emb:\n init_layer(emb_layer, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init)\n\n for layer_idx, tr_layer in enumerate(self.transformer.layers):\n depth = None\n if depthwise_init == 'current':\n depth = layer_idx + 1\n elif depthwise_init == 'global':\n depth = len(self.transformer.layers)\n init_fn = partial(init_layer, method=weight_init, init_depth=depth, zero_bias_init=zero_bias_init)\n tr_layer.apply(init_fn)\n\n for linear in self.linears:\n init_layer(linear, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init)\n\n @property\n def special_token_id(self) -> int:\n return self.card\n\n @property\n def num_codebooks(self) -> int:\n return self.n_q\n\n def forward(self, sequence: torch.Tensor,\n conditions: tp.List[ConditioningAttributes],\n condition_tensors: tp.Optional[ConditionTensors] = None,\n stage: int = -1) -> torch.Tensor:\n \"\"\"Apply language model on sequence and conditions.\n Given a tensor of sequence of shape [B, K, S] with K the number of codebooks and\n S the sequence steps, return the logits with shape [B, card, K, S].\n\n Args:\n indices (torch.Tensor): Indices of the codes to model.\n conditions (list of ConditioningAttributes): Conditions to use when modeling\n the given codes. Note that when evaluating multiple time with the same conditioning\n you should pre-compute those and pass them as `condition_tensors`.\n condition_tensors (dict[str, ConditionType], optional): Pre-computed conditioning\n tensors, see `conditions`.\n stage (int): The codebook level that is being predicted. Relevant for MAGNeT\n in which prediction is done in a codebook-by-codebook manner.\n Takes values in range(n_q), and ignored by default.\n Returns:\n torch.Tensor: Logits.\n \"\"\"\n B, K, S = sequence.shape\n assert K == self.num_codebooks, \"Sequence shape must match the specified number of codebooks\"\n input_ = sum([self.emb[k](sequence[:, k]) for k in range(K)])\n if condition_tensors is None:\n assert not self._is_streaming, \"Conditions tensors should be precomputed when streaming.\"\n # apply dropout modules\n conditions = self.cfg_dropout(conditions)\n conditions = self.att_dropout(conditions)\n tokenized = self.condition_provider.tokenize(conditions)\n # encode conditions and fuse, both have a streaming cache to not recompute when generating.\n condition_tensors = self.condition_provider(tokenized)\n else:\n assert not conditions, \"Shouldn't pass both conditions and condition_tensors.\"\n\n input_, cross_attention_input = self.fuser(input_, condition_tensors)\n\n out = self.transformer(input_, cross_attention_src=cross_attention_input,\n src_mask=(self.attn_mask_per_stage[stage] if stage >= 0 else None)) # type: ignore\n if self.out_norm:\n out = self.out_norm(out)\n logits = torch.stack([self.linears[k](out) for k in range(K)], dim=1) # [B, K, S, card]\n\n # remove the prefix from the model outputs\n if len(self.fuser.fuse2cond['prepend']) > 0:\n logits = logits[:, :, -S:]\n\n return logits # [B, K, S, card]\n\n def compute_predictions(\n self, codes: torch.Tensor,\n conditions: tp.List[ConditioningAttributes],\n condition_tensors: tp.Optional[ConditionTensors] = None,\n stage: int = -1,\n keep_only_valid_steps: bool = True) -> LMOutput:\n \"\"\"Given an input tensor of codes [B, K, T] and list of conditions, runs the model\n forward using the specified codes interleaving pattern.\n\n Args:\n codes (torch.Tensor): Input codes of shape [B, K, T] with B the batch size,\n K the number of codebooks and T the number of timesteps.\n conditions (list of ConditioningAttributes): conditionings to use when modeling\n the given codes. Note that when evaluating multiple time with the same conditioning\n you should pre-compute those and pass them as `condition_tensors`.\n condition_tensors (dict[str, ConditionType], optional): pre-computed conditioning\n tensors, see `conditions`.\n stage (int): The codebook level that is being predicted. Relevant for MAGNeT\n in which prediction is done in a codebook-by-codebook manner.\n Takes values in range(n_q), and ignored by default.\n keep_only_valid_steps (bool): Build a sequence from the pattern up to valid (= fully defined) steps.\n Steps that are beyond valid steps will be replaced by the special_token in that case.\n Returns:\n LMOutput: Language model outputs\n logits (torch.Tensor) of shape [B, K, T, card] corresponding to the provided codes,\n i.e. the first item corresponds to logits to predict the first code, meaning that\n no additional shifting of codes and logits is required.\n mask (torch.Tensor) of shape [B, K, T], mask over valid and invalid positions.\n Given the specified interleaving strategies, parts of the logits and codes should\n not be considered as valid predictions because of invalid context.\n \"\"\"\n B, K, T = codes.shape\n codes = codes.contiguous()\n # map codes [B, K, T] into pattern sequence [B, K, S] using special_token_id for masked tokens\n pattern = self.pattern_provider.get_pattern(T)\n sequence_codes, sequence_indexes, sequence_mask = pattern.build_pattern_sequence(\n codes, self.special_token_id, keep_only_valid_steps=keep_only_valid_steps,\n )\n\n # apply model on pattern sequence\n model = self if self._fsdp is None else self._fsdp\n logits = model(sequence_codes, conditions, condition_tensors, stage=stage) # [B, K, S, card]\n # map back the logits on pattern sequence to logits on original codes: [B, K, S, card] -> [B, K, T, card]\n # and provide the corresponding mask over invalid positions of tokens\n logits = logits.permute(0, 3, 1, 2) # [B, card, K, S]\n # note: we use nans as special token to make it obvious if we feed unexpected logits\n logits, logits_indexes, logits_mask = pattern.revert_pattern_logits(\n logits, float('nan'), keep_only_valid_steps=keep_only_valid_steps\n )\n logits = logits.permute(0, 2, 3, 1) # [B, K, T, card]\n logits_mask = logits_mask[None, :, :].expand(B, -1, -1) # [K, T] -> [B, K, T]\n return LMOutput(logits, logits_mask)\n\n def _sample_next_token(self,\n sequence: torch.Tensor,\n cfg_conditions: CFGConditions,\n unconditional_state: State,\n use_sampling: bool = False,\n temp: float = 1.0,\n top_k: int = 0,\n top_p: float = 0.0,\n cfg_coef: tp.Optional[float] = None,\n cfg_coef_beta: tp.Optional[float] = None,\n two_step_cfg: tp.Optional[bool] = None) -> torch.Tensor:\n \"\"\"Sample next token from the model given a sequence and a set of conditions. The model supports\n multiple sampling strategies (greedy sampling, softmax, top-k, top-p...).\n\n Args:\n sequence (torch.Tensor): Current sequence of shape [B, K, S]\n with K corresponding to the number of codebooks and S the number of sequence steps.\n S = 1 in streaming mode, except for the first step that contains a bigger prompt.\n condition_tensors (dict[str, ConditionType): Set of conditions. If CFG is used,\n should be twice the batch size, being the concatenation of the conditions + null conditions.\n use_sampling (bool): Whether to use a sampling strategy or not.\n temp (float): Sampling temperature.\n top_k (int): K for \"top-k\" sampling.\n top_p (float): P for \"top-p\" sampling.\n cfg_coef (float, optional): classifier free guidance coefficient\n cfg_coef_beta (float, optional): If None, simple classifier free guidance is used with cfg_coef.\n If not None, we apply double classifier free guidance as introduced in MusicGen-Style\n in paragraph 4.3 (https://arxiv.org/pdf/2407.12563). This beta coefficient is meant to\n push the text condition more than the style condition in the case where both text and style\n conditions are being used.\n two_step_cfg (bool): Whether to run classifier free-guidance with 2 distinct steps.\n\n Returns:\n next_token (torch.Tensor): Next token tensor of shape [B, K, 1].\n \"\"\"\n B = sequence.shape[0]\n cfg_coef = self.cfg_coef if cfg_coef is None else cfg_coef\n model = self if self._fsdp is None else self._fsdp\n two_step_cfg = self.two_step_cfg if two_step_cfg is None else two_step_cfg\n if cfg_coef_beta is not None:\n assert isinstance(cfg_conditions, dict)\n condition_tensors = cfg_conditions\n if condition_tensors:\n # Preparing for CFG, predicting conditional text and style, conditional style\n # and unconditional\n sequence = torch.cat([sequence, sequence, sequence], dim=0)\n all_logits = model(\n sequence,\n conditions=[], condition_tensors=condition_tensors)\n if condition_tensors:\n cond_logits, wav_logits, uncond_logits = all_logits.split(B, dim=0) # [B, K, T, card]\n logits = uncond_logits + cfg_coef * (\n wav_logits + cfg_coef_beta * (cond_logits - wav_logits) - uncond_logits\n )\n\n elif two_step_cfg and cfg_conditions != {}:\n assert isinstance(cfg_conditions, tuple), type(cfg_conditions)\n condition_tensors, null_condition_tensors = cfg_conditions\n cond_logits = model(sequence, conditions=[], condition_tensors=condition_tensors)\n state = self.get_streaming_state()\n self.set_streaming_state(unconditional_state)\n uncond_logits = model(sequence, conditions=[], condition_tensors=null_condition_tensors)\n unconditional_state.update(self.get_streaming_state())\n self.set_streaming_state(state)\n logits = uncond_logits + (cond_logits - uncond_logits) * self.cfg_coef\n else:\n assert isinstance(cfg_conditions, dict)\n condition_tensors = cfg_conditions\n if condition_tensors:\n # Preparing for CFG, predicting both conditional and unconditional logits.\n sequence = torch.cat([sequence, sequence], dim=0)\n all_logits = model(\n sequence,\n conditions=[], condition_tensors=condition_tensors)\n if condition_tensors:\n cond_logits, uncond_logits = all_logits.split(B, dim=0) # [B, K, T, card]\n logits = uncond_logits + (cond_logits - uncond_logits) * cfg_coef\n else:\n logits = all_logits\n\n logits = logits.permute(0, 1, 3, 2) # [B, K, card, T]\n logits = logits[..., -1] # [B x K x card]\n\n # Apply softmax for sampling if temp > 0. Else, do greedy sampling to avoid zero division error.\n if use_sampling and temp > 0.0:\n probs = torch.softmax(logits / temp, dim=-1)\n if top_p > 0.0:\n next_token = utils.sample_top_p(probs, p=top_p)\n elif top_k > 0:\n next_token = utils.sample_top_k(probs, k=top_k)\n else:\n next_token = utils.multinomial(probs, num_samples=1)\n else:\n next_token = torch.argmax(logits, dim=-1, keepdim=True)\n\n return next_token\n\n @torch.no_grad()\n def generate(self,\n prompt: tp.Optional[torch.Tensor] = None,\n conditions: tp.List[ConditioningAttributes] = [],\n num_samples: tp.Optional[int] = None,\n max_gen_len: int = 256,\n use_sampling: bool = True,\n temp: float = 1.0,\n top_k: int = 250,\n top_p: float = 0.0,\n cfg_coef: tp.Optional[float] = None,\n cfg_coef_beta: tp.Optional[float] = None,\n two_step_cfg: tp.Optional[bool] = None,\n remove_prompts: bool = False,\n check: bool = False,\n callback: tp.Optional[tp.Callable[[int, int], None]] = None,\n ) -> torch.Tensor:\n \"\"\"Generate tokens sampling from the model given a prompt or unconditionally. Generation can\n be performed in a greedy fashion or using sampling with top K and top P strategies.\n\n Args:\n prompt (torch.Tensor, optional): Prompt tokens of shape [B, K, T].\n conditions (list of ConditioningAttributes, optional): List of conditions.\n num_samples (int, optional): Number of samples to generate when no prompt and no conditions are given.\n max_gen_len (int): Maximum generation length.\n use_sampling (bool): Whether to use a sampling strategy or not.\n temp (float): Sampling temperature.\n top_k (int): K for \"top-k\" sampling.\n top_p (float): P for \"top-p\" sampling.\n cfg_coef (float, optional): Classifier-free guidance coefficient.\n cfg_coef_beta (float, optional): If None, simple classifier free guidance is used with cfg_coef.\n If not None, we apply double classifier free guidance as introduced in MusicGen-Style\n in paragraph 4.3 (https://arxiv.org/pdf/2407.12563). This beta coefficient is meant to\n push the text condition more than the style condition in the case where both text and style\n conditions are being used.\n two_step_cfg (bool, optional): Whether to perform classifier-free guidance with two steps generation.\n remove_prompts (bool): Whether to remove prompts from generation or not.\n check (bool): Whether to apply further checks on generated sequence.\n callback (Callback, optional): Callback function to report generation progress.\n Returns:\n torch.Tensor: Generated tokens.\n \"\"\"\n assert not self.training, \"generation shouldn't be used in training mode.\"\n first_param = next(iter(self.parameters()))\n device = first_param.device\n\n # Checking all input shapes are consistent.\n possible_num_samples = []\n if num_samples is not None:\n possible_num_samples.append(num_samples)\n elif prompt is not None:\n possible_num_samples.append(prompt.shape[0])\n elif conditions:\n possible_num_samples.append(len(conditions))\n else:\n possible_num_samples.append(1)\n assert [x == possible_num_samples[0] for x in possible_num_samples], \"Inconsistent inputs shapes\"\n num_samples = possible_num_samples[0]\n\n # below we create set of conditions: one conditional and one unconditional\n # to do that we merge the regular condition together with the null condition\n # we then do 1 forward pass instead of 2.\n # the reason for that is two-fold:\n # 1. it is about x2 faster than doing 2 forward passes\n # 2. avoid the streaming API treating the 2 passes as part of different time steps\n # We also support doing two different passes, in particular to ensure that\n # the padding structure is exactly the same between train and test.\n # With a batch size of 1, this can be slower though.\n cfg_conditions: CFGConditions\n cfg_conditions = {}\n if cfg_coef_beta is not None:\n if conditions:\n wav_conditions = _drop_description_condition(conditions)\n null_conditions = ClassifierFreeGuidanceDropout(p=1.0)(conditions)\n conditions = conditions + wav_conditions + null_conditions\n tokenized = self.condition_provider.tokenize(conditions)\n cfg_conditions = self.condition_provider(tokenized)\n elif conditions:\n two_step_cfg = self.two_step_cfg if two_step_cfg is None else two_step_cfg\n if conditions:\n null_conditions = ClassifierFreeGuidanceDropout(p=1.0)(conditions)\n if two_step_cfg:\n cfg_conditions = (\n self.condition_provider(self.condition_provider.tokenize(conditions)),\n self.condition_provider(self.condition_provider.tokenize(null_conditions)),\n )\n else:\n conditions = conditions + null_conditions\n tokenized = self.condition_provider.tokenize(conditions)\n cfg_conditions = self.condition_provider(tokenized)\n else:\n cfg_conditions = {}\n\n if prompt is None:\n assert num_samples > 0\n prompt = torch.zeros((num_samples, self.num_codebooks, 0), dtype=torch.long, device=device)\n\n B, K, T = prompt.shape\n start_offset = T\n assert start_offset < max_gen_len\n\n pattern = self.pattern_provider.get_pattern(max_gen_len)\n # this token is used as default value for codes that are not generated yet\n unknown_token = -1\n\n # we generate codes up to the max_gen_len that will be mapped to the pattern sequence\n gen_codes = torch.full((B, K, max_gen_len), unknown_token, dtype=torch.long, device=device)\n # filling the gen_codes with the prompt if needed\n gen_codes[..., :start_offset] = prompt\n # create the gen_sequence with proper interleaving from the pattern: [B, K, S]\n gen_sequence, indexes, mask = pattern.build_pattern_sequence(gen_codes, self.special_token_id)\n # retrieve the start_offset in the sequence:\n # it is the first sequence step that contains the `start_offset` timestep\n start_offset_sequence = pattern.get_first_step_with_timesteps(start_offset)\n assert start_offset_sequence is not None\n\n with self.streaming():\n unconditional_state = self.get_streaming_state()\n prev_offset = 0\n gen_sequence_len = gen_sequence.shape[-1] # gen_sequence shape is [B, K, S]\n for offset in range(start_offset_sequence, gen_sequence_len):\n # get current sequence (note that the streaming API is providing the caching over previous offsets)\n curr_sequence = gen_sequence[..., prev_offset:offset]\n curr_mask = mask[None, ..., prev_offset:offset].expand(B, -1, -1)\n if check:\n # check coherence between mask and sequence\n assert (curr_sequence == torch.where(curr_mask, curr_sequence, self.special_token_id)).all()\n # should never happen as gen_sequence is filled progressively\n assert not (curr_sequence == unknown_token).any()\n # sample next token from the model, next token shape is [B, K, 1]\n next_token = self._sample_next_token(\n curr_sequence, cfg_conditions, unconditional_state, use_sampling, temp, top_k, top_p,\n cfg_coef=cfg_coef, cfg_coef_beta=cfg_coef_beta, two_step_cfg=two_step_cfg)\n # ensure the tokens that should be masked are properly set to special_token_id\n # as the model never output special_token_id\n valid_mask = mask[..., offset:offset+1].expand(B, -1, -1)\n next_token[~valid_mask] = self.special_token_id\n # ensure we don't overwrite prompt tokens, we only write over unknown tokens\n # (then mask tokens should be left as is as well, which is correct)\n gen_sequence[..., offset:offset+1] = torch.where(\n gen_sequence[..., offset:offset+1] == unknown_token,\n next_token, gen_sequence[..., offset:offset+1]\n )\n prev_offset = offset\n if callback is not None:\n callback(1 + offset - start_offset_sequence, gen_sequence_len - start_offset_sequence)\n unconditional_state.clear()\n\n # ensure sequence has been entirely filled\n assert not (gen_sequence == unknown_token).any()\n # ensure gen_sequence pattern and mask are matching\n # which means the gen_sequence is valid according to the pattern\n assert (\n gen_sequence == torch.where(mask[None, ...].expand(B, -1, -1), gen_sequence, self.special_token_id)\n ).all()\n # get back the codes, trimming the prompt if needed and cutting potentially incomplete timesteps\n out_codes, out_indexes, out_mask = pattern.revert_pattern_sequence(gen_sequence, special_token=unknown_token)\n\n # sanity checks over the returned codes and corresponding masks\n assert (out_codes[..., :max_gen_len] != unknown_token).all()\n assert (out_mask[..., :max_gen_len] == 1).all()\n\n out_start_offset = start_offset if remove_prompts else 0\n out_codes = out_codes[..., out_start_offset:max_gen_len]\n\n # ensure the returned codes are all valid\n assert (out_codes >= 0).all() and (out_codes <= self.card).all()\n return out_codes\n" }, { "path": "audiocraft/models/lm_magnet.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nimport math\nimport typing as tp\nimport torch\nimport numpy as np\n\nfrom ..utils import utils\nfrom ..modules.conditioners import (\n ClassifierFreeGuidanceDropout,\n ConditioningAttributes,\n ConditionType,\n)\nfrom .lm import LMModel\n\nlogger = logging.getLogger(__name__)\nConditionTensors = tp.Dict[str, ConditionType]\nCFGConditions = tp.Union[ConditionTensors, tp.Tuple[ConditionTensors, ConditionTensors]]\n\n\nclass MagnetLMModel(LMModel):\n \"\"\"Transformer-based, non-autoregressive model, operates on multiple streams of audio tokens (MAGNeT).\n Args:\n subcodes_context (int): The number of timesteps attended in the self-attention blocks of codebooks > 0.\n When set to -1, attention is unrestricted and all timesteps are attended. Defaults to 5.\n compression_model_framerate (int): frame rate of the audio tokenizer.\n segment_duration (int): Sample length in seconds.\n span_len (int): Determines the length of masking spans. This is the minimal length of consecutive masked tokens,\n for both training and inference. Defaults to 3.\n **kwargs: Additional parameters for the LMModel.\n \"\"\"\n def __init__(self, subcodes_context: int = 5, compression_model_framerate: int = 50,\n segment_duration: int = 10, span_len: int = 3, **kwargs):\n super().__init__(**kwargs)\n self.causal = kwargs['causal']\n self.subcodes_context = subcodes_context\n self.span_len = span_len\n self._build_attn_masks(compression_model_framerate=compression_model_framerate,\n segment_duration=segment_duration,\n num_heads=kwargs['num_heads'],\n device=kwargs['device'], dtype=kwargs['dtype'])\n\n def restricted_context_attn_mask(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:\n \"\"\"Creates a restricted attention mask (local attention map) where the context\n is determined by self.subcodes_context.\n Args:\n seq_len (int): token sequence length.\n device (torch.device): device of the output tensor.\n dtype (torch.dtype): data type of the output tensor.\n Returns:\n torch.Tensor: The restricted attention mask.\n \"\"\"\n # Return a context restricted non-causal att mask\n queries_pos = torch.arange(seq_len, device=device).view(-1, 1)\n keys_pos = torch.arange(seq_len, device=device).view(1, -1)\n\n delta = queries_pos - keys_pos\n valid = torch.abs(delta) <= self.subcodes_context\n return torch.where(\n valid,\n torch.zeros([], device=device, dtype=dtype),\n torch.full([], float('-inf'), device=device, dtype=dtype))\n\n def _stage_attn_mask(self, stage: int, seq_len: int, num_heads: int,\n device: torch.device, dtype: torch.dtype) -> tp.Optional[torch.Tensor]:\n \"\"\"Creates a restricted attention mask given the stage (codebook index).\n Args:\n stage (int): The codebook index. Takes values in [0, n_q].\n seq_len (int): Token sequence length.\n num_heads (int): Num transformer attention heads.\n device (torch.device): device of the output tensor.\n dtype (torch.dtype): data type of the output tensor.\n Returns:\n torch.Tensor: Either a restricted attention mask or None if stage attention is unrestricted.\n \"\"\"\n sa_mask = None\n\n if stage > 0 and self.subcodes_context > -1:\n # parallel - non-causal - with restricted subcodes context\n sa_mask = self.restricted_context_attn_mask(seq_len, device=device, dtype=dtype)\n\n if sa_mask is not None:\n # Repeat for each attention head\n sa_mask = sa_mask.repeat((1, num_heads, 1, 1))\n\n # align8 to enable memory efficient attention\n MEMORY_EFFICIENT_ATTN_ALIGN_FACTOR = 8\n seq_len_aligned = \\\n int(np.ceil(seq_len / MEMORY_EFFICIENT_ATTN_ALIGN_FACTOR)) * MEMORY_EFFICIENT_ATTN_ALIGN_FACTOR\n\n sa_mask_aligned = torch.zeros((1, num_heads, seq_len_aligned, seq_len_aligned), device=device, dtype=dtype)\n sa_mask_aligned[..., :seq_len, :seq_len] = sa_mask\n sa_mask = sa_mask_aligned\n\n return sa_mask\n\n def _build_attn_masks(self, compression_model_framerate: int, segment_duration: int, num_heads: int,\n device: torch.device, dtype: torch.dtype):\n \"\"\"Construct attention mask per stage. For each of the RVQ codebook levels in the [0, n_q] range,\n either a local attention map or None would be stored as an entry in the self.attn_mask_per_stage list.\n Args:\n compression_model_framerate (int): The frame rate of the tokenizer.\n segment_duration (int): Sample length in seconds.\n num_heads (int): Num transformer attention heads.\n device (torch.device): device of the output tensor.\n dtype (torch.dtype): data type of the output tensor.\n \"\"\"\n seq_len = compression_model_framerate * segment_duration\n self.attn_mask_per_stage = [self._stage_attn_mask(stage, seq_len, num_heads,\n device, dtype) for stage in range(self.n_q)]\n\n @torch.no_grad()\n def generate(self,\n prompt: tp.Optional[torch.Tensor] = None,\n conditions: tp.List[ConditioningAttributes] = [],\n num_samples: tp.Optional[int] = None,\n max_gen_len: int = 256,\n use_sampling: bool = True,\n temp: float = 1.0,\n top_k: int = 250,\n top_p: float = 0.0,\n cfg_coef: tp.Optional[float] = None,\n cfg_coef_beta: tp.Optional[float] = None,\n two_step_cfg: tp.Optional[bool] = None,\n remove_prompts: bool = False,\n check: bool = False,\n callback: tp.Optional[tp.Callable[[int, int], None]] = None,\n **kwargs) -> torch.Tensor:\n\n assert cfg_coef is None, \"Unsupported in MAGNeT. Use max_cfg_coef,min_cfg_coef instead.\"\n assert two_step_cfg is None, \"MAGNeT currently doesn't support two step classifier-free-guidance.\"\n assert remove_prompts is False, \"MAGNeT currently doesn't support the remove_prompts arg.\"\n assert check is False, \"MAGNeT currently doesn't support the check arg.\"\n assert cfg_coef_beta is None, \"MAGNeT currently doesn't support the cfg_coef_beta arg.\"\n # Call the MAGNeT-specific generation method\n return self._generate_magnet(prompt=prompt,\n conditions=conditions,\n num_samples=num_samples,\n max_gen_len=max_gen_len,\n use_sampling=use_sampling,\n temp=temp,\n top_k=top_k,\n top_p=top_p,\n callback=callback, **kwargs)\n\n @torch.no_grad()\n def _generate_magnet(self,\n prompt: tp.Optional[torch.Tensor] = None,\n conditions: tp.List[ConditioningAttributes] = [],\n num_samples: tp.Optional[int] = None,\n max_gen_len: int = 256,\n use_sampling: bool = True,\n temp: float = 3.0,\n top_k: int = 0,\n top_p: float = 0.9,\n callback: tp.Optional[tp.Callable[[int, int], None]] = None,\n max_cfg_coef: float = 10.0,\n min_cfg_coef: float = 1.0,\n decoding_steps: tp.List[int] = [20, 10, 10, 10],\n anneal_temp: bool = True,\n span_scoring='max',\n span_arrangement='nonoverlap') -> torch.Tensor:\n \"\"\"Generate audio tokens given textual conditions, and optionally given audio prompts,\n by running MAGNeT's iterative decoding algorithm for each of the n_q RVQ levels.\n Args:\n prompt (torch.Tensor): Prompt tokens of shape [B, K, T].\n conditions (list of ConditioningAttributes): List of conditions.\n num_samples (int): Number of samples to generate when no prompt and no conditions are given.\n max_gen_len (int): Maximum generation length.\n use_sampling (bool): Whether to use a sampling strategy or not.\n temp (float): Initial sampling temperature.\n top_k (int): k for \"top-k\" sampling.\n top_p (float): p for \"top-p\" sampling.\n callback (Callback): Callback function to report generation progress.\n max_clsfg_coef (float): Initial coefficient used for classifier free guidance.\n min_clsfg_coef (float): Final coefficient used for classifier free guidance.\n decoding_steps (list of n_q ints): The number of iterative decoding steps,\n for each of the n_q RVQ codebooks.\n anneal_temp (bool): When set to True, softmax temperature will be linearly decayed to zero, at each stage.\n span_scoring (str): Use the maximum probability of each span ('max')\n or the product of probabilities ('prod').\n span_arrangement (str): Use either non-overlapping spans ('nonoverlap') or overlapping spans ('stride1').\n in the masking scheme.\n Returns:\n torch.Tensor: Generated tokens.\n \"\"\"\n assert not self.training, \"generation shouldn't be used in training mode.\"\n first_param = next(iter(self.parameters()))\n device = first_param.device\n\n # Checking all input shapes are consistent.\n possible_num_samples = []\n if num_samples is not None:\n possible_num_samples.append(num_samples)\n elif prompt is not None:\n possible_num_samples.append(prompt.shape[0])\n elif conditions:\n possible_num_samples.append(len(conditions))\n else:\n possible_num_samples.append(1)\n assert [x == possible_num_samples[0] for x in possible_num_samples], \"Inconsistent inputs shapes\"\n num_samples = possible_num_samples[0]\n\n # below we create set of conditions: one conditional and one unconditional\n # to do that we merge the regular condition together with the null condition\n # we then do 1 forward pass instead of 2.\n cfg_conditions: tp.Optional[ConditionTensors]\n if conditions:\n null_conditions = ClassifierFreeGuidanceDropout(p=1.0)(conditions)\n conditions = conditions + null_conditions\n tokenized = self.condition_provider.tokenize(conditions)\n cfg_conditions = self.condition_provider(tokenized)\n else:\n cfg_conditions = {}\n\n if prompt is None:\n assert num_samples > 0\n prompt = torch.zeros((num_samples, self.num_codebooks, 0), dtype=torch.long, device=device)\n\n B, K, prompt_length = prompt.shape\n start_offset = prompt_length\n assert start_offset < max_gen_len\n\n mask_id = self.special_token_id\n\n # we generate codes with a fixed sequence length\n shape = (B, K, max_gen_len)\n\n gen_codes = torch.full(shape, mask_id, dtype=torch.long, device=device)\n # filling the gen_codes with the prompt if needed\n gen_codes[..., :start_offset] = prompt\n # create the gen_sequence with proper interleaving from the pattern: [B, K, S]\n gen_sequence = gen_codes\n\n curr_step = 0\n for stage, n_steps in zip(range(self.n_q), decoding_steps):\n gen_sequence, curr_step = self._generate_stage(gen_sequence,\n cfg_conditions,\n stage=stage,\n device=device,\n prompt_length=prompt_length,\n prompt=prompt,\n temp=temp,\n max_cfg_coef=max_cfg_coef,\n min_cfg_coef=min_cfg_coef,\n top_k=top_k,\n top_p=top_p,\n timesteps=n_steps,\n anneal_temp=anneal_temp,\n span_scoring=span_scoring,\n use_sampling=use_sampling,\n span_arrangement=span_arrangement,\n curr_step=curr_step,\n total_steps=sum(decoding_steps),\n callback=callback)\n\n return gen_sequence\n\n @torch.no_grad()\n def _generate_stage(self,\n gen_sequence: torch.Tensor,\n condition_tensors: tp.Optional[ConditionTensors],\n stage: int,\n device: torch.device,\n prompt_length: int = 0,\n prompt: tp.Optional[torch.Tensor] = None,\n use_sampling: bool = True,\n temp: float = 3.0,\n max_cfg_coef: float = 10.0,\n min_cfg_coef: float = 1.0,\n top_k: int = 0,\n top_p: float = 0.0,\n timesteps: int = 10,\n anneal_temp: bool = True,\n span_scoring: str = 'max',\n span_arrangement: str = 'nonoverlap',\n curr_step: int = 0,\n total_steps: int = 0,\n callback: tp.Optional[tp.Callable[[int, int], None]] = None) -> tp.Tuple[torch.Tensor, int]:\n \"\"\"Generate audio tokens of a single RVQ level (stage), given the previously generated stages,\n and the textual conditions.\n Args:\n gen_sequence (torch.Tensor): Previously generated tokens.\n condition_tensors (tp.Optional[ConditionTensors]): pre-computed conditioning tensors.\n stage (int): RVQ level to generate.\n device (torch.device): device of the output tensor.\n prompt_length (int): Temporal length of the audio prompt.\n prompt (torch.Tensor): Prompt tokens of shape [B, K, T].\n use_sampling (bool): Whether to use a sampling strategy or not.\n temp (float): Initial sampling temperature.\n max_clsfg_coef (float): Initial coefficient used for classifier free guidance.\n min_clsfg_coef (float): Final coefficient used for classifier free guidance.\n top_k (int): k for \"top-k\" sampling.\n top_p (float): p for \"top-p\" sampling.\n timesteps (int): Number of iterative decoding steps.\n anneal_temp (bool): When set to True, softmax temperature will be linearly decayed to zero, at each stage.\n span_scoring (str): Use the maximum probability of each span ('max')\n or the product of probabilities ('prod').\n span_arrangement (str): Use either non-overlapping spans ('nonoverlap') or overlapping spans ('stride1').\n in the masking scheme.\n curr_step (int): Global iterative decoding step counter.\n total_steps (int): Total decoding steps.\n callback (Callback): Callback function to report generation progress.\n Returns:\n tuple(torch.Tensor, int): Generated tokens and the current decoding step counter.\n \"\"\"\n B, K, T = gen_sequence.shape\n shape = (B, 1, T) # generating a single codebook per stage\n\n mask_id = self.special_token_id\n stage_gen_seq = torch.full(shape, mask_id, dtype=torch.long, device=device)\n\n assert span_arrangement == 'nonoverlap' or span_arrangement == 'stride1'\n chunk_masking = self.span_len > 1 and span_arrangement == 'nonoverlap'\n\n DONT_REMASK_ME_SCORE = -1e4\n\n model = self if self._fsdp is None else self._fsdp\n\n if chunk_masking:\n # span-wise scores\n n_chunks = T // self.span_len\n if T % self.span_len != 0:\n # trim sequence ending to achieve a multiple of span_len\n T = self.span_len * n_chunks\n gen_sequence = gen_sequence[..., :T]\n stage_gen_seq = stage_gen_seq[..., :T]\n\n chunked_shape = (B, 1, n_chunks)\n n_prompt_chunks = prompt_length // self.span_len\n scores = torch.zeros(chunked_shape, dtype=torch.float32, device=device)\n scores[..., :n_prompt_chunks] = DONT_REMASK_ME_SCORE\n num_chunks_to_gen = n_chunks - n_prompt_chunks\n else:\n # token-wise scores\n scores = torch.zeros(shape, dtype=torch.float32, device=device)\n scores[..., :prompt_length] = DONT_REMASK_ME_SCORE\n gen_T = T - prompt_length\n\n # run MAGNeT iterative decoding for \"timesteps\" iterations\n for timestep, steps_left in zip(torch.linspace(0, 1, timesteps, device=device), reversed(range(timesteps))):\n\n mask_p = torch.cos(timestep * math.pi * 0.5)\n\n if chunk_masking:\n num_masked = max(int((mask_p * num_chunks_to_gen).item()), 1)\n else:\n num_masked = max(int((mask_p * gen_T).item()), 1)\n\n # masking\n run_lps_masking = (span_arrangement == 'stride1') and self.span_len > 1\n if run_lps_masking:\n # masking of the k least probable overlapping (stride 1) spans\n mask = torch.concat((\n [self._least_probable_span_masking(scores[[i], :, :], num_masked).to(device)\n for i in range(B)]), dim=0)\n stage_gen_seq[mask] = mask_id\n else:\n # masking of the k least probable non-overlapping spans\n masked = scores.topk(num_masked, dim=-1).indices\n if chunk_masking:\n chunks_mask = torch.full(chunked_shape, False, dtype=torch.bool, device=device)\n chunks_mask = chunks_mask.scatter(2, masked, True)\n mask = torch.repeat_interleave(chunks_mask, self.span_len, dim=-1)\n stage_gen_seq[mask] = mask_id\n else:\n stage_gen_seq = stage_gen_seq.scatter(2, masked, mask_id)\n\n if prompt is not None:\n stage_gen_seq[..., :prompt_length] = prompt[:, stage, :].unsqueeze(1)\n\n gen_sequence[:, [stage], :] = stage_gen_seq\n if condition_tensors:\n # duplicate input for classifier free guidance\n sequence = torch.cat([gen_sequence, gen_sequence], dim=0)\n\n all_logits = model(sequence, [], condition_tensors, stage=stage)\n\n if condition_tensors:\n # classifier free guidance with annealing\n cond_logits, uncond_logits = all_logits.split(B, dim=0) # [B, K, T, card]\n clsfg_coef = float(mask_p) * max_cfg_coef + (1 - float(mask_p)) * min_cfg_coef\n logits = uncond_logits + (cond_logits - uncond_logits) * clsfg_coef\n else:\n logits = all_logits\n\n # temperature annealing - linear\n t = temp * (steps_left / timesteps) if anneal_temp else temp\n\n # sampling\n logits = logits[:, stage, :, :].unsqueeze(1)\n probs = torch.softmax(logits / max(t, 1e-2), dim=-1)\n if use_sampling:\n if top_p > 0.0:\n sampled_tokens = utils.sample_top_p(probs, p=top_p)\n elif top_k > 0:\n sampled_tokens = utils.sample_top_k(probs, k=top_k)\n else:\n sampled_tokens = utils.multinomial(probs, num_samples=1)\n else:\n sampled_tokens = torch.argmax(logits, dim=-1, keepdim=True)\n\n # place mask_id token in each of the masked positions\n mask = stage_gen_seq == mask_id\n stage_gen_seq = torch.where(mask, sampled_tokens[..., 0], stage_gen_seq)\n gen_sequence[:, [stage], :] = stage_gen_seq\n\n # get probs of sampled tokens\n sampled_probs = torch.gather(probs, 3, sampled_tokens)[..., 0]\n\n # span scoring\n if chunk_masking:\n if span_scoring == 'max':\n # max in linear space\n scores = 1 - torch.max(sampled_probs.reshape((B, 1, n_chunks, -1)), dim=-1)[0]\n elif span_scoring == 'prod':\n # prod in log space\n scores = torch.sum(-torch.log(sampled_probs).reshape((B, 1, n_chunks, -1)), dim=-1)\n else:\n raise NotImplementedError\n else:\n # prod in log space for lps masking (stride1)\n scores = -torch.log(sampled_probs)\n\n # Fix unmasked tokens by placing inf probs (-inf scores)\n if chunk_masking:\n scores = scores.masked_fill(~chunks_mask, DONT_REMASK_ME_SCORE)\n else:\n scores = scores.masked_fill(~mask, DONT_REMASK_ME_SCORE)\n\n if callback is not None:\n curr_step += 1\n callback(curr_step, total_steps)\n\n return gen_sequence, curr_step\n\n def _construct_spans_mask(self, span_starts: torch.Tensor, T: int, device: torch.device) -> torch.Tensor:\n \"\"\"Build a [1x1xT] boolean mask consists of overlapping spans of True values, where\n span_starts defines the initial index of each span, and the span length is\n defined by self.span_len.\n Args:\n span_starts (torch.Tensor): Boolean mask determines the temporal location of each span start.\n T (int): Sequence length.\n device (torch.device): device of the output tensor.\n Returns:\n torch.Tensor: Spans mask of shape [1x1xT]\n \"\"\"\n mask = torch.full((1, 1, T), False, device=device)\n mask[:, :, span_starts] = True\n shifted_mask = mask.clone()\n for _ in range(self.span_len - 1):\n shifted_mask = torch.concat((torch.full((1, 1, 1), False, device=device), shifted_mask[:, :, :-1]), dim=-1)\n mask = torch.logical_or(mask, shifted_mask)\n return mask\n\n def _least_probable_span_masking(self, scores: torch.Tensor, num_masked_trg: int) -> torch.Tensor:\n \"\"\"Construct a [1x1xT] boolean mask, consists of the u least probable spans,\n where the token probability is determined by -scores, and the total\n number of masked tokens is as closest as possible to num_masked_trg.\n Find u using binary search.\n Args:\n scores (torch.Tensor): Per token score [-log(prob)]\n num_masked_trg: int: The desired amount of tokens to be masked.\n Returns:\n torch.Tensor: Spans mask of shape [1x1xT]\n \"\"\"\n T = scores.shape[-1]\n device = scores.device\n scores_unfolded = scores.unfold(2, self.span_len, 1)\n # Span score is the product of probs (sum in log space)\n span_scores = scores_unfolded.sum(dim=-1)\n spans_by_scores = torch.argsort(span_scores[0, 0], descending=True)\n\n num_masked_trg = max(num_masked_trg, self.span_len)\n\n # Binary search for u - the number least probable overlapping masked spans s.t.\n # the total masking rate is the closest to num_masked_trg / T.\n min_u = num_masked_trg // self.span_len\n max_u = num_masked_trg - self.span_len + 1\n mid = round(0.5 * (min_u + max_u))\n\n if mid == min_u or mid == max_u:\n return self._construct_spans_mask(spans_by_scores[:mid], T, device)\n\n while mid > min_u and mid < max_u:\n mask = self._construct_spans_mask(spans_by_scores[:mid], T, device)\n n_masked = mask.sum()\n if n_masked > num_masked_trg:\n max_u = mid\n mid = round(0.5 * (min_u + max_u))\n else:\n min_u = mid\n mid = round(0.5 * (min_u + max_u))\n\n return mask\n" }, { "path": "audiocraft/models/loaders.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nUtility functions to load from the checkpoints.\nEach checkpoint is a torch.saved dict with the following keys:\n- 'xp.cfg': the hydra config as dumped during training. This should be used\n to rebuild the object using the audiocraft.models.builders functions,\n- 'model_best_state': a readily loadable best state for the model, including\n the conditioner. The model obtained from `xp.cfg` should be compatible\n with this state dict. In the case of a LM, the encodec model would not be\n bundled along but instead provided separately.\n\nThose functions also support loading from a remote location with the Torch Hub API.\nThey also support overriding some parameters, in particular the device and dtype\nof the returned model.\n\"\"\"\n\nfrom pathlib import Path\nfrom huggingface_hub import hf_hub_download\nimport typing as tp\nimport os\n\nfrom omegaconf import OmegaConf, DictConfig\nimport torch\n\nimport audiocraft\n\nfrom . import builders\nfrom .encodec import CompressionModel\n\n\ndef get_audiocraft_cache_dir() -> tp.Optional[str]:\n return os.environ.get('AUDIOCRAFT_CACHE_DIR', None)\n\n\ndef _get_state_dict(\n file_or_url_or_id: tp.Union[Path, str],\n filename: tp.Optional[str] = None,\n device='cpu',\n cache_dir: tp.Optional[str] = None,\n):\n if cache_dir is None:\n cache_dir = get_audiocraft_cache_dir()\n # Return the state dict either from a file or url\n file_or_url_or_id = str(file_or_url_or_id)\n assert isinstance(file_or_url_or_id, str)\n\n if os.path.isfile(file_or_url_or_id):\n return torch.load(file_or_url_or_id, map_location=device)\n\n if os.path.isdir(file_or_url_or_id):\n file = f\"{file_or_url_or_id}/{filename}\"\n return torch.load(file, map_location=device)\n\n elif file_or_url_or_id.startswith('https://'):\n return torch.hub.load_state_dict_from_url(file_or_url_or_id, map_location=device, check_hash=True)\n\n else:\n assert filename is not None, \"filename needs to be defined if using HF checkpoints\"\n file = hf_hub_download(\n repo_id=file_or_url_or_id,\n filename=filename,\n cache_dir=cache_dir,\n library_name=\"audiocraft\",\n library_version=audiocraft.__version__,\n )\n return torch.load(file, map_location=device)\n\n\ndef load_compression_model_ckpt(file_or_url_or_id: tp.Union[Path, str], cache_dir: tp.Optional[str] = None):\n return _get_state_dict(file_or_url_or_id, filename=\"compression_state_dict.bin\", cache_dir=cache_dir)\n\n\ndef load_compression_model(\n file_or_url_or_id: tp.Union[Path, str],\n device=\"cpu\",\n cache_dir: tp.Optional[str] = None,\n):\n pkg = load_compression_model_ckpt(file_or_url_or_id, cache_dir=cache_dir)\n if 'pretrained' in pkg:\n return CompressionModel.get_pretrained(pkg['pretrained'], device=device)\n cfg = OmegaConf.create(pkg['xp.cfg'])\n cfg.device = str(device)\n model = builders.get_compression_model(cfg)\n model.load_state_dict(pkg[\"best_state\"])\n model.eval()\n return model\n\n\ndef load_lm_model_ckpt(file_or_url_or_id: tp.Union[Path, str], cache_dir: tp.Optional[str] = None):\n return _get_state_dict(file_or_url_or_id, filename=\"state_dict.bin\", cache_dir=cache_dir)\n\n\ndef _delete_param(cfg: DictConfig, full_name: str):\n parts = full_name.split('.')\n for part in parts[:-1]:\n if part in cfg:\n cfg = cfg[part]\n else:\n return\n OmegaConf.set_struct(cfg, False)\n if parts[-1] in cfg:\n del cfg[parts[-1]]\n OmegaConf.set_struct(cfg, True)\n\n\ndef load_lm_model(file_or_url_or_id: tp.Union[Path, str], device='cpu', cache_dir: tp.Optional[str] = None):\n pkg = load_lm_model_ckpt(file_or_url_or_id, cache_dir=cache_dir)\n cfg = OmegaConf.create(pkg['xp.cfg'])\n cfg.device = str(device)\n if cfg.device == 'cpu':\n cfg.dtype = 'float32'\n else:\n cfg.dtype = 'float16'\n _delete_param(cfg, 'conditioners.self_wav.chroma_stem.cache_path')\n _delete_param(cfg, 'conditioners.args.merge_text_conditions_p')\n _delete_param(cfg, 'conditioners.args.drop_desc_p')\n model = builders.get_lm_model(cfg)\n model.load_state_dict(pkg['best_state'])\n model.eval()\n model.cfg = cfg\n return model\n\n\ndef load_lm_model_magnet(file_or_url_or_id: tp.Union[Path, str], compression_model_frame_rate: int,\n device='cpu', cache_dir: tp.Optional[str] = None):\n pkg = load_lm_model_ckpt(file_or_url_or_id, cache_dir=cache_dir)\n cfg = OmegaConf.create(pkg['xp.cfg'])\n cfg.device = str(device)\n if cfg.device == 'cpu':\n cfg.dtype = 'float32'\n else:\n cfg.dtype = 'float16'\n _delete_param(cfg, 'conditioners.args.merge_text_conditions_p')\n _delete_param(cfg, 'conditioners.args.drop_desc_p')\n\n cfg.transformer_lm.compression_model_framerate = compression_model_frame_rate\n cfg.transformer_lm.segment_duration = cfg.dataset.segment_duration\n cfg.transformer_lm.span_len = cfg.masking.span_len\n\n # MAGNeT models v1 support only xformers backend.\n from audiocraft.modules.transformer import set_efficient_attention_backend\n\n if cfg.transformer_lm.memory_efficient:\n set_efficient_attention_backend(\"xformers\")\n\n model = builders.get_lm_model(cfg)\n model.load_state_dict(pkg['best_state'])\n model.eval()\n model.cfg = cfg\n return model\n\n\ndef load_jasco_model(file_or_url_or_id: tp.Union[Path, str],\n compression_model: CompressionModel,\n device='cpu', cache_dir: tp.Optional[str] = None):\n pkg = load_lm_model_ckpt(file_or_url_or_id, cache_dir=cache_dir)\n cfg = OmegaConf.create(pkg['xp.cfg'])\n cfg.device = str(device)\n if cfg.device == 'cpu':\n cfg.dtype = 'float32'\n else:\n cfg.dtype = 'float16'\n model = builders.get_jasco_model(cfg, compression_model)\n model.load_state_dict(pkg['best_state'])\n model.eval()\n model.cfg = cfg\n return model\n\n\ndef load_mbd_ckpt(file_or_url_or_id: tp.Union[Path, str],\n filename: tp.Optional[str] = None,\n cache_dir: tp.Optional[str] = None):\n return _get_state_dict(file_or_url_or_id, filename=filename, cache_dir=cache_dir)\n\n\ndef load_diffusion_models(file_or_url_or_id: tp.Union[Path, str],\n device='cpu',\n filename: tp.Optional[str] = None,\n cache_dir: tp.Optional[str] = None):\n pkg = load_mbd_ckpt(file_or_url_or_id, filename=filename, cache_dir=cache_dir)\n models = []\n processors = []\n cfgs = []\n sample_rate = pkg['sample_rate']\n for i in range(pkg['n_bands']):\n cfg = pkg[i]['cfg']\n model = builders.get_diffusion_model(cfg)\n model_dict = pkg[i]['model_state']\n model.load_state_dict(model_dict)\n model.to(device)\n processor = builders.get_processor(cfg=cfg.processor, sample_rate=sample_rate)\n processor_dict = pkg[i]['processor_state']\n processor.load_state_dict(processor_dict)\n processor.to(device)\n models.append(model)\n processors.append(processor)\n cfgs.append(cfg)\n return models, processors, cfgs\n\n\ndef load_audioseal_models(\n file_or_url_or_id: tp.Union[Path, str],\n device=\"cpu\",\n filename: tp.Optional[str] = None,\n cache_dir: tp.Optional[str] = None,\n):\n\n detector_ckpt = _get_state_dict(\n file_or_url_or_id,\n filename=f\"detector_{filename}.pth\",\n device=device,\n cache_dir=cache_dir,\n )\n assert (\n \"model\" in detector_ckpt\n ), f\"No model state dict found in {file_or_url_or_id}/detector_{filename}.pth\"\n detector_state = detector_ckpt[\"model\"]\n\n generator_ckpt = _get_state_dict(\n file_or_url_or_id,\n filename=f\"generator_{filename}.pth\",\n device=device,\n cache_dir=cache_dir,\n )\n assert (\n \"model\" in generator_ckpt\n ), f\"No model state dict found in {file_or_url_or_id}/generator_{filename}.pth\"\n generator_state = generator_ckpt[\"model\"]\n\n def load_model_config():\n if Path(file_or_url_or_id).joinpath(f\"{filename}.yaml\").is_file():\n return OmegaConf.load(Path(file_or_url_or_id).joinpath(f\"{filename}.yaml\"))\n elif file_or_url_or_id.startswith(\"https://\"):\n import requests # type: ignore\n\n resp = requests.get(f\"{file_or_url_or_id}/{filename}.yaml\")\n return OmegaConf.create(resp.text)\n else:\n file = hf_hub_download(\n repo_id=file_or_url_or_id,\n filename=f\"{filename}.yaml\",\n cache_dir=cache_dir,\n library_name=\"audiocraft\",\n library_version=audiocraft.__version__,\n )\n return OmegaConf.load(file)\n\n try:\n cfg = load_model_config()\n except Exception as exc: # noqa\n cfg_fp = (\n Path(__file__)\n .parents[2]\n .joinpath(\"config\", \"model\", \"watermark\", \"default.yaml\")\n )\n cfg = OmegaConf.load(cfg_fp)\n\n OmegaConf.resolve(cfg)\n model = builders.get_watermark_model(cfg)\n\n model.generator.load_state_dict(generator_state)\n model.detector.load_state_dict(detector_state)\n return model.to(device)\n" }, { "path": "audiocraft/models/magnet.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nMain model for using MAGNeT. This will combine all the required components\nand provide easy access to the generation API.\n\"\"\"\nimport typing as tp\nimport torch\n\nfrom .genmodel import BaseGenModel\nfrom .loaders import load_compression_model, load_lm_model_magnet\n\n\nclass MAGNeT(BaseGenModel):\n \"\"\"MAGNeT main model with convenient generation API.\n Args:\n See MusicGen class.\n \"\"\"\n def __init__(self, **kwargs):\n super().__init__(**kwargs)\n # MAGNeT operates over a fixed sequence length defined in it's config.\n self.duration = self.lm.cfg.dataset.segment_duration\n self.set_generation_params()\n\n @staticmethod\n def get_pretrained(name: str = 'facebook/magnet-small-10secs', device=None):\n \"\"\"Return pretrained model, we provide six models:\n - facebook/magnet-small-10secs (300M), text to music, 10-second audio samples.\n # see: https://huggingface.co/facebook/magnet-small-10secs\n - facebook/magnet-medium-10secs (1.5B), text to music, 10-second audio samples.\n # see: https://huggingface.co/facebook/magnet-medium-10secs\n - facebook/magnet-small-30secs (300M), text to music, 30-second audio samples.\n # see: https://huggingface.co/facebook/magnet-small-30secs\n - facebook/magnet-medium-30secs (1.5B), text to music, 30-second audio samples.\n # see: https://huggingface.co/facebook/magnet-medium-30secs\n - facebook/audio-magnet-small (300M), text to sound-effect (10-second samples).\n # see: https://huggingface.co/facebook/audio-magnet-small\n - facebook/audio-magnet-medium (1.5B), text to sound-effect (10-second samples).\n # see: https://huggingface.co/facebook/audio-magnet-medium\n \"\"\"\n if device is None:\n if torch.cuda.device_count():\n device = 'cuda'\n else:\n device = 'cpu'\n\n compression_model = load_compression_model(name, device=device)\n lm = load_lm_model_magnet(name, compression_model_frame_rate=int(compression_model.frame_rate), device=device)\n\n if 'self_wav' in lm.condition_provider.conditioners:\n lm.condition_provider.conditioners['self_wav'].match_len_on_eval = True\n\n kwargs = {'name': name, 'compression_model': compression_model, 'lm': lm}\n return MAGNeT(**kwargs)\n\n def set_generation_params(self, use_sampling: bool = True, top_k: int = 0,\n top_p: float = 0.9, temperature: float = 3.0,\n max_cfg_coef: float = 10.0, min_cfg_coef: float = 1.0,\n decoding_steps: tp.List[int] = [20, 10, 10, 10],\n span_arrangement: str = 'nonoverlap'):\n \"\"\"Set the generation parameters for MAGNeT.\n\n Args:\n use_sampling (bool, optional): Use sampling if True, else do argmax decoding. Defaults to True.\n top_k (int, optional): top_k used for sampling. Defaults to 0.\n top_p (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.9.\n temperature (float, optional): Initial softmax temperature parameter. Defaults to 3.0.\n max_cfg_coef (float, optional): Coefficient used for classifier free guidance. Defaults to 10.0.\n min_cfg_coef (float, optional): End coefficient of classifier free guidance annealing. Defaults to 1.0.\n decoding_steps (list of n_q ints, optional): The number of iterative decoding steps,\n for each of the n_q RVQ codebooks.\n span_arrangement (str, optional): Use either non-overlapping spans ('nonoverlap')\n or overlapping spans ('stride1') in the masking scheme.\n \"\"\"\n self.generation_params = {\n 'use_sampling': use_sampling,\n 'temp': temperature,\n 'top_k': top_k,\n 'top_p': top_p,\n 'max_cfg_coef': max_cfg_coef,\n 'min_cfg_coef': min_cfg_coef,\n 'decoding_steps': [int(s) for s in decoding_steps],\n 'span_arrangement': span_arrangement\n }\n" }, { "path": "audiocraft/models/multibanddiffusion.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nMulti Band Diffusion models as described in\n\"From Discrete Tokens to High-Fidelity Audio Using Multi-Band Diffusion\"\n(paper link).\n\"\"\"\n\nimport typing as tp\n\nimport torch\nimport julius\n\nfrom .unet import DiffusionUnet\nfrom ..modules.diffusion_schedule import NoiseSchedule\nfrom .encodec import CompressionModel\nfrom ..solvers.compression import CompressionSolver\nfrom .loaders import load_compression_model, load_diffusion_models\n\n\nclass DiffusionProcess:\n \"\"\"Sampling for a diffusion Model.\n\n Args:\n model (DiffusionUnet): Diffusion U-Net model.\n noise_schedule (NoiseSchedule): Noise schedule for diffusion process.\n \"\"\"\n def __init__(self, model: DiffusionUnet, noise_schedule: NoiseSchedule) -> None:\n self.model = model\n self.schedule = noise_schedule\n\n def generate(self, condition: torch.Tensor, initial_noise: torch.Tensor,\n step_list: tp.Optional[tp.List[int]] = None):\n \"\"\"Perform one diffusion process to generate one of the bands.\n\n Args:\n condition (torch.Tensor): The embeddings from the compression model.\n initial_noise (torch.Tensor): The initial noise to start the process.\n \"\"\"\n return self.schedule.generate_subsampled(model=self.model, initial=initial_noise, step_list=step_list,\n condition=condition)\n\n\nclass MultiBandDiffusion:\n \"\"\"Sample from multiple diffusion models.\n\n Args:\n DPs (list of DiffusionProcess): Diffusion processes.\n codec_model (CompressionModel): Underlying compression model used to obtain discrete tokens.\n \"\"\"\n def __init__(self, DPs: tp.List[DiffusionProcess], codec_model: CompressionModel) -> None:\n self.DPs = DPs\n self.codec_model = codec_model\n self.device = next(self.codec_model.parameters()).device\n\n @property\n def sample_rate(self) -> int:\n return self.codec_model.sample_rate\n\n @staticmethod\n def get_mbd_musicgen(device=None):\n \"\"\"Load our diffusion models trained for MusicGen.\"\"\"\n if device is None:\n device = 'cuda' if torch.cuda.is_available() else 'cpu'\n path = 'facebook/multiband-diffusion'\n filename = 'mbd_musicgen_32khz.th'\n name = 'facebook/musicgen-small'\n codec_model = load_compression_model(name, device=device)\n models, processors, cfgs = load_diffusion_models(path, filename=filename, device=device)\n DPs = []\n for i in range(len(models)):\n schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)\n DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))\n return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)\n\n @staticmethod\n def get_mbd_24khz(bw: float = 3.0,\n device: tp.Optional[tp.Union[torch.device, str]] = None,\n n_q: tp.Optional[int] = None):\n \"\"\"Get the pretrained Models for MultibandDiffusion.\n\n Args:\n bw (float): Bandwidth of the compression model.\n device (torch.device or str, optional): Device on which the models are loaded.\n n_q (int, optional): Number of quantizers to use within the compression model.\n \"\"\"\n if device is None:\n device = 'cuda' if torch.cuda.is_available() else 'cpu'\n assert bw in [1.5, 3.0, 6.0], f\"bandwidth {bw} not available\"\n if n_q is not None:\n assert n_q in [2, 4, 8]\n assert {1.5: 2, 3.0: 4, 6.0: 8}[bw] == n_q, \\\n f\"bandwidth and number of codebooks missmatch to use n_q = {n_q} bw should be {n_q * (1.5 / 2)}\"\n n_q = {1.5: 2, 3.0: 4, 6.0: 8}[bw]\n codec_model = CompressionSolver.model_from_checkpoint(\n '//pretrained/facebook/encodec_24khz', device=device)\n codec_model.set_num_codebooks(n_q)\n codec_model = codec_model.to(device)\n path = 'facebook/multiband-diffusion'\n filename = f'mbd_comp_{n_q}.pt'\n models, processors, cfgs = load_diffusion_models(path, filename=filename, device=device)\n DPs = []\n for i in range(len(models)):\n schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)\n DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))\n return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)\n\n @torch.no_grad()\n def get_condition(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:\n \"\"\"Get the conditioning (i.e. latent representations of the compression model) from a waveform.\n Args:\n wav (torch.Tensor): The audio that we want to extract the conditioning from.\n sample_rate (int): Sample rate of the audio.\"\"\"\n if sample_rate != self.sample_rate:\n wav = julius.resample_frac(wav, sample_rate, self.sample_rate)\n codes, scale = self.codec_model.encode(wav)\n assert scale is None, \"Scaled compression models not supported.\"\n emb = self.get_emb(codes)\n return emb\n\n @torch.no_grad()\n def get_emb(self, codes: torch.Tensor):\n \"\"\"Get latent representation from the discrete codes.\n Args:\n codes (torch.Tensor): Discrete tokens.\"\"\"\n emb = self.codec_model.decode_latent(codes)\n return emb\n\n def generate(self, emb: torch.Tensor, size: tp.Optional[torch.Size] = None,\n step_list: tp.Optional[tp.List[int]] = None):\n \"\"\"Generate waveform audio from the latent embeddings of the compression model.\n Args:\n emb (torch.Tensor): Conditioning embeddings\n size (None, torch.Size): Size of the output\n if None this is computed from the typical upsampling of the model.\n step_list (list[int], optional): list of Markov chain steps, defaults to 50 linearly spaced step.\n \"\"\"\n if size is None:\n upsampling = int(self.codec_model.sample_rate / self.codec_model.frame_rate)\n size = torch.Size([emb.size(0), self.codec_model.channels, emb.size(-1) * upsampling])\n assert size[0] == emb.size(0)\n out = torch.zeros(size).to(self.device)\n for DP in self.DPs:\n out += DP.generate(condition=emb, step_list=step_list, initial_noise=torch.randn_like(out))\n return out\n\n def re_eq(self, wav: torch.Tensor, ref: torch.Tensor, n_bands: int = 32, strictness: float = 1):\n \"\"\"Match the eq to the encodec output by matching the standard deviation of some frequency bands.\n Args:\n wav (torch.Tensor): Audio to equalize.\n ref (torch.Tensor): Reference audio from which we match the spectrogram.\n n_bands (int): Number of bands of the eq.\n strictness (float): How strict the matching. 0 is no matching, 1 is exact matching.\n \"\"\"\n split = julius.SplitBands(n_bands=n_bands, sample_rate=self.codec_model.sample_rate).to(wav.device)\n bands = split(wav)\n bands_ref = split(ref)\n out = torch.zeros_like(ref)\n for i in range(n_bands):\n out += bands[i] * (bands_ref[i].std() / bands[i].std()) ** strictness\n return out\n\n def regenerate(self, wav: torch.Tensor, sample_rate: int):\n \"\"\"Regenerate a waveform through compression and diffusion regeneration.\n Args:\n wav (torch.Tensor): Original 'ground truth' audio.\n sample_rate (int): Sample rate of the input (and output) wav.\n \"\"\"\n if sample_rate != self.codec_model.sample_rate:\n wav = julius.resample_frac(wav, sample_rate, self.codec_model.sample_rate)\n emb = self.get_condition(wav, sample_rate=self.codec_model.sample_rate)\n size = wav.size()\n out = self.generate(emb, size=size)\n if sample_rate != self.codec_model.sample_rate:\n out = julius.resample_frac(out, self.codec_model.sample_rate, sample_rate)\n return out\n\n def tokens_to_wav(self, tokens: torch.Tensor, n_bands: int = 32):\n \"\"\"Generate Waveform audio with diffusion from the discrete codes.\n Args:\n tokens (torch.Tensor): Discrete codes.\n n_bands (int): Bands for the eq matching.\n \"\"\"\n wav_encodec = self.codec_model.decode(tokens)\n condition = self.get_emb(tokens)\n wav_diffusion = self.generate(emb=condition, size=wav_encodec.size())\n return self.re_eq(wav=wav_diffusion, ref=wav_encodec, n_bands=n_bands)\n" }, { "path": "audiocraft/models/musicgen.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nMain model for using MusicGen. This will combine all the required components\nand provide easy access to the generation API.\n\"\"\"\n\nimport typing as tp\nimport warnings\n\nimport torch\n\nfrom .encodec import CompressionModel\nfrom .genmodel import BaseGenModel\nfrom .lm import LMModel\nfrom .builders import get_debug_compression_model, get_debug_lm_model\nfrom .loaders import load_compression_model, load_lm_model\nfrom ..data.audio_utils import convert_audio\nfrom ..modules.conditioners import ConditioningAttributes, WavCondition, StyleConditioner\n\n\nMelodyList = tp.List[tp.Optional[torch.Tensor]]\nMelodyType = tp.Union[torch.Tensor, MelodyList]\n\n\n# backward compatible names mapping\n_HF_MODEL_CHECKPOINTS_MAP = {\n \"small\": \"facebook/musicgen-small\",\n \"medium\": \"facebook/musicgen-medium\",\n \"large\": \"facebook/musicgen-large\",\n \"melody\": \"facebook/musicgen-melody\",\n \"style\": \"facebook/musicgen-style\",\n}\n\n\nclass MusicGen(BaseGenModel):\n \"\"\"MusicGen main model with convenient generation API.\n\n Args:\n name (str): name of the model.\n compression_model (CompressionModel): Compression model\n used to map audio to invertible discrete representations.\n lm (LMModel): Language model over discrete representations.\n max_duration (float, optional): maximum duration the model can produce,\n otherwise, inferred from the training params.\n \"\"\"\n def __init__(self, name: str, compression_model: CompressionModel, lm: LMModel,\n max_duration: tp.Optional[float] = None):\n super().__init__(name, compression_model, lm, max_duration)\n self.set_generation_params(duration=15) # default duration\n\n @staticmethod\n def get_pretrained(name: str = 'facebook/musicgen-melody', device=None):\n \"\"\"Return pretrained model, we provide four models:\n - facebook/musicgen-small (300M), text to music,\n # see: https://huggingface.co/facebook/musicgen-small\n - facebook/musicgen-medium (1.5B), text to music,\n # see: https://huggingface.co/facebook/musicgen-medium\n - facebook/musicgen-melody (1.5B) text to music and text+melody to music,\n # see: https://huggingface.co/facebook/musicgen-melody\n - facebook/musicgen-large (3.3B), text to music,\n # see: https://huggingface.co/facebook/musicgen-large\n - facebook/musicgen-style (1.5 B), text and style to music,\n # see: https://huggingface.co/facebook/musicgen-style\n \"\"\"\n if device is None:\n if torch.cuda.device_count():\n device = 'cuda'\n else:\n device = 'cpu'\n\n if name == 'debug':\n # used only for unit tests\n compression_model = get_debug_compression_model(device)\n lm = get_debug_lm_model(device)\n return MusicGen(name, compression_model, lm, max_duration=30)\n\n if name in _HF_MODEL_CHECKPOINTS_MAP:\n warnings.warn(\n \"MusicGen pretrained model relying on deprecated checkpoint mapping. \" +\n f\"Please use full pre-trained id instead: facebook/musicgen-{name}\")\n name = _HF_MODEL_CHECKPOINTS_MAP[name]\n\n lm = load_lm_model(name, device=device)\n compression_model = load_compression_model(name, device=device)\n if 'self_wav' in lm.condition_provider.conditioners:\n lm.condition_provider.conditioners['self_wav'].match_len_on_eval = True\n lm.condition_provider.conditioners['self_wav']._use_masking = False\n\n return MusicGen(name, compression_model, lm)\n\n def set_generation_params(self, use_sampling: bool = True, top_k: int = 250,\n top_p: float = 0.0, temperature: float = 1.0,\n duration: float = 30.0, cfg_coef: float = 3.0,\n cfg_coef_beta: tp.Optional[float] = None,\n two_step_cfg: bool = False, extend_stride: float = 18,):\n \"\"\"Set the generation parameters for MusicGen.\n\n Args:\n use_sampling (bool, optional): Use sampling if True, else do argmax decoding. Defaults to True.\n top_k (int, optional): top_k used for sampling. Defaults to 250.\n top_p (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.0.\n temperature (float, optional): Softmax temperature parameter. Defaults to 1.0.\n duration (float, optional): Duration of the generated waveform. Defaults to 30.0.\n cfg_coef (float, optional): Coefficient used for classifier free guidance. Defaults to 3.0.\n cfg_coef_beta (float, optional): beta coefficient in double classifier free guidance.\n Should be only used for MusicGen melody if we want to push the text condition more than\n the audio conditioning. See paragraph 4.3 in https://arxiv.org/pdf/2407.12563 to understand\n double CFG.\n two_step_cfg (bool, optional): If True, performs 2 forward for Classifier Free Guidance,\n instead of batching together the two. This has some impact on how things\n are padded but seems to have little impact in practice.\n extend_stride: when doing extended generation (i.e. more than 30 seconds), by how much\n should we extend the audio each time. Larger values will mean less context is\n preserved, and shorter value will require extra computations.\n \"\"\"\n assert extend_stride < self.max_duration, \"Cannot stride by more than max generation duration.\"\n self.extend_stride = extend_stride\n self.duration = duration\n self.generation_params = {\n 'use_sampling': use_sampling,\n 'temp': temperature,\n 'top_k': top_k,\n 'top_p': top_p,\n 'cfg_coef': cfg_coef,\n 'two_step_cfg': two_step_cfg,\n 'cfg_coef_beta': cfg_coef_beta,\n }\n\n def set_style_conditioner_params(self, eval_q: int = 3, excerpt_length: float = 3.0,\n ds_factor: tp.Optional[int] = None,\n encodec_n_q: tp.Optional[int] = None) -> None:\n \"\"\"Set the parameters of the style conditioner\n Args:\n eval_q (int): the number of residual quantization streams used to quantize the style condition\n the smaller it is, the narrower is the information bottleneck\n excerpt_length (float): the excerpt length in seconds that is extracted from the audio\n conditioning\n ds_factor: (int): the downsampling factor used to downsample the style tokens before\n using them as a prefix\n encodec_n_q: (int, optional): if encodec is used as a feature extractor, sets the number\n of streams that is used to extract features\n \"\"\"\n assert isinstance(self.lm.condition_provider.conditioners.self_wav, StyleConditioner), \\\n \"Only use this function if you model is MusicGen-Style\"\n self.lm.condition_provider.conditioners.self_wav.set_params(eval_q=eval_q,\n excerpt_length=excerpt_length,\n ds_factor=ds_factor,\n encodec_n_q=encodec_n_q)\n\n def generate_with_chroma(self, descriptions: tp.List[str], melody_wavs: MelodyType,\n melody_sample_rate: int, progress: bool = False,\n return_tokens: bool = False) -> tp.Union[torch.Tensor,\n tp.Tuple[torch.Tensor, torch.Tensor]]:\n \"\"\"Generate samples conditioned on text and melody.\n\n Args:\n descriptions (list of str): A list of strings used as text conditioning.\n melody_wavs: (torch.Tensor or list of Tensor): A batch of waveforms used as\n melody conditioning. Should have shape [B, C, T] with B matching the description length,\n C=1 or 2. It can be [C, T] if there is a single description. It can also be\n a list of [C, T] tensors.\n melody_sample_rate: (int): Sample rate of the melody waveforms.\n progress (bool, optional): Flag to display progress of the generation process. Defaults to False.\n \"\"\"\n if isinstance(melody_wavs, torch.Tensor):\n if melody_wavs.dim() == 2:\n melody_wavs = melody_wavs[None]\n if melody_wavs.dim() != 3:\n raise ValueError(\"Melody wavs should have a shape [B, C, T].\")\n melody_wavs = list(melody_wavs)\n else:\n for melody in melody_wavs:\n if melody is not None:\n assert melody.dim() == 2, \"One melody in the list has the wrong number of dims.\"\n\n melody_wavs = [\n convert_audio(wav, melody_sample_rate, self.sample_rate, self.audio_channels)\n if wav is not None else None\n for wav in melody_wavs]\n attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions=descriptions, prompt=None,\n melody_wavs=melody_wavs)\n assert prompt_tokens is None\n tokens = self._generate_tokens(attributes, prompt_tokens, progress)\n if return_tokens:\n return self.generate_audio(tokens), tokens\n return self.generate_audio(tokens)\n\n @torch.no_grad()\n def _prepare_tokens_and_attributes(\n self,\n descriptions: tp.Sequence[tp.Optional[str]],\n prompt: tp.Optional[torch.Tensor],\n melody_wavs: tp.Optional[MelodyList] = None,\n ) -> tp.Tuple[tp.List[ConditioningAttributes], tp.Optional[torch.Tensor]]:\n \"\"\"Prepare model inputs.\n\n Args:\n descriptions (list of str): A list of strings used as text conditioning.\n prompt (torch.Tensor): A batch of waveforms used for continuation.\n melody_wavs (torch.Tensor, optional): A batch of waveforms\n used as melody conditioning. Defaults to None.\n \"\"\"\n attributes = [\n ConditioningAttributes(text={'description': description})\n for description in descriptions]\n\n if melody_wavs is None:\n for attr in attributes:\n attr.wav['self_wav'] = WavCondition(\n torch.zeros((1, 1, 1), device=self.device),\n torch.tensor([0], device=self.device),\n sample_rate=[self.sample_rate],\n path=[None])\n else:\n if 'self_wav' not in self.lm.condition_provider.conditioners:\n raise RuntimeError(\"This model doesn't support melody conditioning. \"\n \"Use the `melody` model.\")\n assert len(melody_wavs) == len(descriptions), \\\n f\"number of melody wavs must match number of descriptions! \" \\\n f\"got melody len={len(melody_wavs)}, and descriptions len={len(descriptions)}\"\n for attr, melody in zip(attributes, melody_wavs):\n if melody is None:\n attr.wav['self_wav'] = WavCondition(\n torch.zeros((1, 1, 1), device=self.device),\n torch.tensor([0], device=self.device),\n sample_rate=[self.sample_rate],\n path=[None])\n else:\n attr.wav['self_wav'] = WavCondition(\n melody[None].to(device=self.device),\n torch.tensor([melody.shape[-1]], device=self.device),\n sample_rate=[self.sample_rate],\n path=[None],\n )\n\n if prompt is not None:\n if descriptions is not None:\n assert len(descriptions) == len(prompt), \"Prompt and nb. descriptions doesn't match\"\n prompt = prompt.to(self.device)\n prompt_tokens, scale = self.compression_model.encode(prompt)\n assert scale is None\n else:\n prompt_tokens = None\n return attributes, prompt_tokens\n\n def _generate_tokens(self, attributes: tp.List[ConditioningAttributes],\n prompt_tokens: tp.Optional[torch.Tensor], progress: bool = False) -> torch.Tensor:\n \"\"\"Generate discrete audio tokens given audio prompt and/or conditions.\n\n Args:\n attributes (list of ConditioningAttributes): Conditions used for generation (text/melody).\n prompt_tokens (torch.Tensor, optional): Audio prompt used for continuation.\n progress (bool, optional): Flag to display progress of the generation process. Defaults to False.\n Returns:\n torch.Tensor: Generated audio, of shape [B, C, T], T is defined by the generation params.\n \"\"\"\n total_gen_len = int(self.duration * self.frame_rate)\n max_prompt_len = int(min(self.duration, self.max_duration) * self.frame_rate)\n current_gen_offset: int = 0\n\n def _progress_callback(generated_tokens: int, tokens_to_generate: int):\n generated_tokens += current_gen_offset\n if self._progress_callback is not None:\n # Note that total_gen_len might be quite wrong depending on the\n # codebook pattern used, but with delay it is almost accurate.\n self._progress_callback(generated_tokens, tokens_to_generate)\n else:\n print(f'{generated_tokens: 6d} / {tokens_to_generate: 6d}', end='\\r')\n\n if prompt_tokens is not None:\n assert max_prompt_len >= prompt_tokens.shape[-1], \\\n \"Prompt is longer than audio to generate\"\n\n callback = None\n if progress:\n callback = _progress_callback\n\n if self.duration <= self.max_duration:\n # generate by sampling from LM, simple case.\n with self.autocast:\n gen_tokens = self.lm.generate(\n prompt_tokens, attributes,\n callback=callback, max_gen_len=total_gen_len, **self.generation_params)\n\n else:\n # now this gets a bit messier, we need to handle prompts,\n # melody conditioning etc.\n ref_wavs = [attr.wav['self_wav'] for attr in attributes]\n all_tokens = []\n if prompt_tokens is None:\n prompt_length = 0\n else:\n all_tokens.append(prompt_tokens)\n prompt_length = prompt_tokens.shape[-1]\n\n assert self.extend_stride is not None, \"Stride should be defined to generate beyond max_duration\"\n assert self.extend_stride < self.max_duration, \"Cannot stride by more than max generation duration.\"\n stride_tokens = int(self.frame_rate * self.extend_stride)\n\n while current_gen_offset + prompt_length < total_gen_len:\n time_offset = current_gen_offset / self.frame_rate\n chunk_duration = min(self.duration - time_offset, self.max_duration)\n max_gen_len = int(chunk_duration * self.frame_rate)\n for attr, ref_wav in zip(attributes, ref_wavs):\n wav_length = ref_wav.length.item()\n if wav_length == 0:\n continue\n # We will extend the wav periodically if it not long enough.\n # we have to do it here rather than in conditioners.py as otherwise\n # we wouldn't have the full wav.\n initial_position = int(time_offset * self.sample_rate)\n wav_target_length = int(self.max_duration * self.sample_rate)\n positions = torch.arange(initial_position,\n initial_position + wav_target_length, device=self.device)\n attr.wav['self_wav'] = WavCondition(\n ref_wav[0][..., positions % wav_length],\n torch.full_like(ref_wav[1], wav_target_length),\n [self.sample_rate] * ref_wav[0].size(0),\n [None], [0.])\n with self.autocast:\n gen_tokens = self.lm.generate(\n prompt_tokens, attributes,\n callback=callback, max_gen_len=max_gen_len, **self.generation_params)\n if prompt_tokens is None:\n all_tokens.append(gen_tokens)\n else:\n all_tokens.append(gen_tokens[:, :, prompt_tokens.shape[-1]:])\n prompt_tokens = gen_tokens[:, :, stride_tokens:]\n prompt_length = prompt_tokens.shape[-1]\n current_gen_offset += stride_tokens\n\n gen_tokens = torch.cat(all_tokens, dim=-1)\n return gen_tokens\n" }, { "path": "audiocraft/models/unet.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nPytorch Unet Module used for diffusion.\n\"\"\"\n\nfrom dataclasses import dataclass\nimport typing as tp\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom audiocraft.modules.transformer import StreamingTransformer, create_sin_embedding\n\n\n@dataclass\nclass Output:\n sample: torch.Tensor\n\n\ndef get_model(cfg, channels: int, side: int, num_steps: int):\n if cfg.model == 'unet':\n return DiffusionUnet(\n chin=channels, num_steps=num_steps, **cfg.diffusion_unet)\n else:\n raise RuntimeError('Not Implemented')\n\n\nclass ResBlock(nn.Module):\n def __init__(self, channels: int, kernel: int = 3, norm_groups: int = 4,\n dilation: int = 1, activation: tp.Type[nn.Module] = nn.ReLU,\n dropout: float = 0.):\n super().__init__()\n stride = 1\n padding = dilation * (kernel - stride) // 2\n Conv = nn.Conv1d\n Drop = nn.Dropout1d\n self.norm1 = nn.GroupNorm(norm_groups, channels)\n self.conv1 = Conv(channels, channels, kernel, 1, padding, dilation=dilation)\n self.activation1 = activation()\n self.dropout1 = Drop(dropout)\n\n self.norm2 = nn.GroupNorm(norm_groups, channels)\n self.conv2 = Conv(channels, channels, kernel, 1, padding, dilation=dilation)\n self.activation2 = activation()\n self.dropout2 = Drop(dropout)\n\n def forward(self, x):\n h = self.dropout1(self.conv1(self.activation1(self.norm1(x))))\n h = self.dropout2(self.conv2(self.activation2(self.norm2(h))))\n return x + h\n\n\nclass DecoderLayer(nn.Module):\n def __init__(self, chin: int, chout: int, kernel: int = 4, stride: int = 2,\n norm_groups: int = 4, res_blocks: int = 1, activation: tp.Type[nn.Module] = nn.ReLU,\n dropout: float = 0.):\n super().__init__()\n padding = (kernel - stride) // 2\n self.res_blocks = nn.Sequential(\n *[ResBlock(chin, norm_groups=norm_groups, dilation=2**idx, dropout=dropout)\n for idx in range(res_blocks)])\n self.norm = nn.GroupNorm(norm_groups, chin)\n ConvTr = nn.ConvTranspose1d\n self.convtr = ConvTr(chin, chout, kernel, stride, padding, bias=False)\n self.activation = activation()\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n x = self.res_blocks(x)\n x = self.norm(x)\n x = self.activation(x)\n x = self.convtr(x)\n return x\n\n\nclass EncoderLayer(nn.Module):\n def __init__(self, chin: int, chout: int, kernel: int = 4, stride: int = 2,\n norm_groups: int = 4, res_blocks: int = 1, activation: tp.Type[nn.Module] = nn.ReLU,\n dropout: float = 0.):\n super().__init__()\n padding = (kernel - stride) // 2\n Conv = nn.Conv1d\n self.conv = Conv(chin, chout, kernel, stride, padding, bias=False)\n self.norm = nn.GroupNorm(norm_groups, chout)\n self.activation = activation()\n self.res_blocks = nn.Sequential(\n *[ResBlock(chout, norm_groups=norm_groups, dilation=2**idx, dropout=dropout)\n for idx in range(res_blocks)])\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n B, C, T = x.shape\n stride, = self.conv.stride\n pad = (stride - (T % stride)) % stride\n x = F.pad(x, (0, pad))\n\n x = self.conv(x)\n x = self.norm(x)\n x = self.activation(x)\n x = self.res_blocks(x)\n return x\n\n\nclass BLSTM(nn.Module):\n \"\"\"BiLSTM with same hidden units as input dim.\n \"\"\"\n def __init__(self, dim, layers=2):\n super().__init__()\n self.lstm = nn.LSTM(bidirectional=True, num_layers=layers, hidden_size=dim, input_size=dim)\n self.linear = nn.Linear(2 * dim, dim)\n\n def forward(self, x):\n x = x.permute(2, 0, 1)\n x = self.lstm(x)[0]\n x = self.linear(x)\n x = x.permute(1, 2, 0)\n return x\n\n\nclass DiffusionUnet(nn.Module):\n def __init__(self, chin: int = 3, hidden: int = 24, depth: int = 3, growth: float = 2.,\n max_channels: int = 10_000, num_steps: int = 1000, emb_all_layers=False, cross_attention: bool = False,\n bilstm: bool = False, transformer: bool = False,\n codec_dim: tp.Optional[int] = None, **kwargs):\n super().__init__()\n self.encoders = nn.ModuleList()\n self.decoders = nn.ModuleList()\n self.embeddings: tp.Optional[nn.ModuleList] = None\n self.embedding = nn.Embedding(num_steps, hidden)\n if emb_all_layers:\n self.embeddings = nn.ModuleList()\n self.condition_embedding: tp.Optional[nn.Module] = None\n for d in range(depth):\n encoder = EncoderLayer(chin, hidden, **kwargs)\n decoder = DecoderLayer(hidden, chin, **kwargs)\n self.encoders.append(encoder)\n self.decoders.insert(0, decoder)\n if emb_all_layers and d > 0:\n assert self.embeddings is not None\n self.embeddings.append(nn.Embedding(num_steps, hidden))\n chin = hidden\n hidden = min(int(chin * growth), max_channels)\n self.bilstm: tp.Optional[nn.Module]\n if bilstm:\n self.bilstm = BLSTM(chin)\n else:\n self.bilstm = None\n self.use_transformer = transformer\n self.cross_attention = False\n if transformer:\n self.cross_attention = cross_attention\n self.transformer = StreamingTransformer(chin, 8, 6, bias_ff=False, bias_attn=False,\n cross_attention=cross_attention)\n\n self.use_codec = False\n if codec_dim is not None:\n self.conv_codec = nn.Conv1d(codec_dim, chin, 1)\n self.use_codec = True\n\n def forward(self, x: torch.Tensor, step: tp.Union[int, torch.Tensor], condition: tp.Optional[torch.Tensor] = None):\n skips = []\n bs = x.size(0)\n z = x\n view_args = [1]\n if type(step) is torch.Tensor:\n step_tensor = step\n else:\n step_tensor = torch.tensor([step], device=x.device, dtype=torch.long).expand(bs)\n\n for idx, encoder in enumerate(self.encoders):\n z = encoder(z)\n if idx == 0:\n z = z + self.embedding(step_tensor).view(bs, -1, *view_args).expand_as(z)\n elif self.embeddings is not None:\n z = z + self.embeddings[idx - 1](step_tensor).view(bs, -1, *view_args).expand_as(z)\n\n skips.append(z)\n\n if self.use_codec: # insert condition in the bottleneck\n assert condition is not None, \"Model defined for conditionnal generation\"\n condition_emb = self.conv_codec(condition) # reshape to the bottleneck dim\n assert condition_emb.size(-1) <= 2 * z.size(-1), \\\n f\"You are downsampling the conditionning with factor >=2 : {condition_emb.size(-1)=} and {z.size(-1)=}\"\n if not self.cross_attention:\n\n condition_emb = torch.nn.functional.interpolate(condition_emb, z.size(-1))\n assert z.size() == condition_emb.size()\n z += condition_emb\n cross_attention_src = None\n else:\n cross_attention_src = condition_emb.permute(0, 2, 1) # B, T, C\n B, T, C = cross_attention_src.shape\n positions = torch.arange(T, device=x.device).view(1, -1, 1)\n pos_emb = create_sin_embedding(positions, C, max_period=10_000, dtype=cross_attention_src.dtype)\n cross_attention_src = cross_attention_src + pos_emb\n if self.use_transformer:\n z = self.transformer(z.permute(0, 2, 1), cross_attention_src=cross_attention_src).permute(0, 2, 1)\n else:\n if self.bilstm is None:\n z = torch.zeros_like(z)\n else:\n z = self.bilstm(z)\n\n for decoder in self.decoders:\n s = skips.pop(-1)\n z = z[:, :, :s.shape[2]]\n z = z + s\n z = decoder(z)\n\n z = z[:, :, :x.shape[2]]\n return Output(z)\n" }, { "path": "audiocraft/models/watermark.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\nimport typing as tp\nfrom abc import ABC, abstractmethod\n\nimport torch\nimport torch.nn as nn\n\nfrom audiocraft.models.loaders import load_audioseal_models\n\n\nclass WMModel(ABC, nn.Module):\n \"\"\"\n A wrapper interface to different watermarking models for\n training or evaluation purporses\n \"\"\"\n\n @abstractmethod\n def get_watermark(\n self,\n x: torch.Tensor,\n message: tp.Optional[torch.Tensor] = None,\n sample_rate: int = 16_000,\n ) -> torch.Tensor:\n \"\"\"Get the watermark from an audio tensor and a message.\n If the input message is None, a random message of\n n bits {0,1} will be generated\n \"\"\"\n\n @abstractmethod\n def detect_watermark(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"Detect the watermarks from the audio signal\n\n Args:\n x: Audio signal, size batch x frames\n\n Returns:\n tensor of size (B, 2+n, frames) where:\n Detection results of shape (B, 2, frames)\n Message decoding results of shape (B, n, frames)\n \"\"\"\n\n\nclass AudioSeal(WMModel):\n \"\"\"Wrap Audioseal (https://github.com/facebookresearch/audioseal) for the\n training and evaluation. The generator and detector are jointly trained\n \"\"\"\n\n def __init__(\n self,\n generator: nn.Module,\n detector: nn.Module,\n nbits: int = 0,\n ):\n super().__init__()\n self.generator = generator # type: ignore\n self.detector = detector # type: ignore\n\n # Allow to re-train an n-bit model with new 0-bit message\n self.nbits = nbits if nbits else self.generator.msg_processor.nbits\n\n def get_watermark(\n self,\n x: torch.Tensor,\n message: tp.Optional[torch.Tensor] = None,\n sample_rate: int = 16_000,\n ) -> torch.Tensor:\n return self.generator.get_watermark(x, message=message, sample_rate=sample_rate)\n\n def detect_watermark(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Detect the watermarks from the audio signal. The first two units of the output\n are used for detection, the rest is used to decode the message. If the audio is\n not watermarked, the message will be random.\n\n Args:\n x: Audio signal, size batch x frames\n Returns\n torch.Tensor: Detection + decoding results of shape (B, 2+nbits, T).\n \"\"\"\n\n # Getting the direct decoded message from the detector\n result = self.detector.detector(x) # b x 2+nbits\n # hardcode softmax on 2 first units used for detection\n result[:, :2, :] = torch.softmax(result[:, :2, :], dim=1)\n return result\n\n def forward( # generator\n self,\n x: torch.Tensor,\n message: tp.Optional[torch.Tensor] = None,\n sample_rate: int = 16_000,\n alpha: float = 1.0,\n ) -> torch.Tensor:\n \"\"\"Apply the watermarking to the audio signal x with a tune-down ratio (default 1.0)\"\"\"\n wm = self.get_watermark(x, message)\n return x + alpha * wm\n\n @staticmethod\n def get_pretrained(name=\"base\", device=None) -> WMModel:\n if device is None:\n if torch.cuda.device_count():\n device = \"cuda\"\n else:\n device = \"cpu\"\n return load_audioseal_models(\"facebook/audioseal\", filename=name, device=device)\n" }, { "path": "audiocraft/modules/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Modules used for building the models.\"\"\"\n\n# flake8: noqa\nfrom .conv import (\n NormConv1d,\n NormConv2d,\n NormConvTranspose1d,\n NormConvTranspose2d,\n StreamableConv1d,\n StreamableConvTranspose1d,\n pad_for_conv1d,\n pad1d,\n unpad1d,\n)\nfrom .lstm import StreamableLSTM\nfrom .seanet import SEANetEncoder, SEANetDecoder\nfrom .transformer import StreamingTransformer" }, { "path": "audiocraft/modules/activations.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\nimport torch.nn as nn\nfrom torch import Tensor\nfrom typing import Union, Callable\n\n\nclass CustomGLU(nn.Module):\n \"\"\"Custom Gated Linear Unit activation.\n Applies a modified gated linear unit :math:`a * f(b)` where :math:`a` is the first half\n of the input matrices, :math:`b` is the second half, and :math:`f` is a provided activation\n function (i.e. sigmoid, swish, etc.).\n\n Args:\n activation (nn.Module): The custom activation to apply in the Gated Linear Unit\n dim (int): the dimension on which to split the input. Default: -1\n\n Shape:\n - Input: :math:`(\\ast_1, N, \\ast_2)` where `*` means, any number of additional\n dimensions\n - Output: :math:`(\\ast_1, M, \\ast_2)` where :math:`M=N/2`\n\n Examples::\n >>> m = CustomGLU(nn.Sigmoid())\n >>> input = torch.randn(4, 2)\n >>> output = m(input)\n \"\"\"\n def __init__(self, activation: nn.Module, dim: int = -1):\n super(CustomGLU, self).__init__()\n self.dim = dim\n self.activation = activation\n\n def forward(self, x: Tensor):\n assert x.shape[self.dim] % 2 == 0 # M = N / 2\n a, b = torch.chunk(x, 2, dim=self.dim)\n return a * self.activation(b)\n\n\nclass SwiGLU(CustomGLU):\n \"\"\"SiLU Gated Linear Unit activation.\n Applies SiLU Gated Linear Unit :math:`a * SiLU(b)` where :math:`a` is\n the first half of the input matrices, :math:`b` is the second half.\n\n Args:\n dim (int): the dimension on which to split the input. Default: -1\n \"\"\"\n def __init__(self, dim: int = -1):\n super(SwiGLU, self).__init__(nn.SiLU(), dim)\n\n\nclass GeGLU(CustomGLU):\n \"\"\"GeLU Gated Linear Unit activation.\n Applies GeLU Gated Linear Unit :math:`a * GELU(b)` where :math:`a` is\n the first half of the input matrices, :math:`b` is the second half.\n\n Args:\n dim (int): the dimension on which to split the input. Default: -1\n \"\"\"\n def __init__(self, dim: int = -1):\n super(GeGLU, self).__init__(nn.GELU(), dim)\n\n\nclass ReGLU(CustomGLU):\n \"\"\"ReLU Gated Linear Unit activation.\n Applies ReLU Gated Linear Unit :math:`a * ReLU(b)` where :math:`a` is\n the first half of the input matrices, :math:`b` is the second half.\n\n Args:\n dim (int): the dimension on which to split the input. Default: -1\n \"\"\"\n def __init__(self, dim: int = -1):\n super(ReGLU, self).__init__(nn.ReLU(), dim)\n\n\ndef get_activation_fn(\n activation: Union[str, Callable[[Tensor], Tensor]]\n) -> Union[str, Callable[[Tensor], Tensor]]:\n \"\"\"Helper function to map an activation string to the activation class.\n If the supplied activation is not a string that is recognized, the activation is passed back.\n\n Args:\n activation (str, or Callable[[Tensor], Tensor]): Activation to check\n \"\"\"\n if isinstance(activation, str):\n if activation == \"reglu\":\n return ReGLU()\n elif activation == \"geglu\":\n return GeGLU()\n elif activation == \"swiglu\":\n return SwiGLU()\n return activation\n" }, { "path": "audiocraft/modules/chroma.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\nimport typing as tp\n\nfrom einops import rearrange\nfrom librosa import filters\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nimport torchaudio\n\n\nclass ChromaExtractor(nn.Module):\n \"\"\"Chroma extraction and quantization.\n\n Args:\n sample_rate (int): Sample rate for the chroma extraction.\n n_chroma (int): Number of chroma bins for the chroma extraction.\n radix2_exp (int): Size of stft window for the chroma extraction (power of 2, e.g. 12 -> 2^12).\n nfft (int, optional): Number of FFT.\n winlen (int, optional): Window length.\n winhop (int, optional): Window hop size.\n argmax (bool, optional): Whether to use argmax. Defaults to False.\n norm (float, optional): Norm for chroma normalization. Defaults to inf.\n \"\"\"\n def __init__(self, sample_rate: int, n_chroma: int = 12, radix2_exp: int = 12, nfft: tp.Optional[int] = None,\n winlen: tp.Optional[int] = None, winhop: tp.Optional[int] = None, argmax: bool = False,\n norm: float = torch.inf):\n super().__init__()\n self.winlen = winlen or 2 ** radix2_exp\n self.nfft = nfft or self.winlen\n self.winhop = winhop or (self.winlen // 4)\n self.sample_rate = sample_rate\n self.n_chroma = n_chroma\n self.norm = norm\n self.argmax = argmax\n self.register_buffer('fbanks', torch.from_numpy(filters.chroma(sr=sample_rate, n_fft=self.nfft, tuning=0,\n n_chroma=self.n_chroma)), persistent=False)\n self.spec = torchaudio.transforms.Spectrogram(n_fft=self.nfft, win_length=self.winlen,\n hop_length=self.winhop, power=2, center=True,\n pad=0, normalized=True)\n\n def forward(self, wav: torch.Tensor) -> torch.Tensor:\n T = wav.shape[-1]\n # in case we are getting a wav that was dropped out (nullified)\n # from the conditioner, make sure wav length is no less that nfft\n if T < self.nfft:\n pad = self.nfft - T\n r = 0 if pad % 2 == 0 else 1\n wav = F.pad(wav, (pad // 2, pad // 2 + r), 'constant', 0)\n assert wav.shape[-1] == self.nfft, f\"expected len {self.nfft} but got {wav.shape[-1]}\"\n\n spec = self.spec(wav).squeeze(1)\n raw_chroma = torch.einsum('cf,...ft->...ct', self.fbanks, spec)\n norm_chroma = torch.nn.functional.normalize(raw_chroma, p=self.norm, dim=-2, eps=1e-6)\n norm_chroma = rearrange(norm_chroma, 'b d t -> b t d')\n\n if self.argmax:\n idx = norm_chroma.argmax(-1, keepdim=True)\n norm_chroma[:] = 0\n norm_chroma.scatter_(dim=-1, index=idx, value=1)\n\n return norm_chroma\n" }, { "path": "audiocraft/modules/codebooks_patterns.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom collections import namedtuple\nfrom dataclasses import dataclass\nfrom functools import lru_cache\nimport logging\nimport typing as tp\n\nfrom abc import ABC, abstractmethod\nimport torch\n\nLayoutCoord = namedtuple('LayoutCoord', ['t', 'q']) # (timestep, codebook index)\nPatternLayout = tp.List[tp.List[LayoutCoord]] # Sequence of coordinates\nlogger = logging.getLogger(__name__)\n\n\n@dataclass\nclass Pattern:\n \"\"\"Base implementation of a pattern over a sequence with multiple codebooks.\n\n The codebook pattern consists in a layout, defining for each sequence step\n the list of coordinates of each codebook timestep in the resulting interleaved sequence.\n The first item of the pattern is always an empty list in order to properly insert a special token\n to start with. For convenience, we also keep track of ``n_q`` the number of codebooks used for the pattern\n and ``timesteps`` the number of timesteps corresponding to the original sequence.\n\n The pattern provides convenient methods to build and revert interleaved sequences from it:\n ``build_pattern_sequence`` maps a given a dense input tensor of multi-codebook sequence from [B, K, T]\n to the interleaved sequence of shape [B, K, S] applying the pattern, with B being the batch size,\n K being the number of codebooks, T the number of original timesteps and S the number of sequence steps\n for the output sequence. The unfilled positions are replaced with a special token and the built sequence\n is returned along with a mask indicating valid tokens.\n ``revert_pattern_sequence`` maps back an interleaved sequence of shape [B, K, S] to the original alignment\n of codebooks across timesteps to an output tensor of shape [B, K, T], using again a special token and a mask\n to fill and specify invalid positions if needed.\n See the dedicated methods for more details.\n \"\"\"\n # Pattern layout, for each sequence step, we have a list of coordinates\n # corresponding to the original codebook timestep and position.\n # The first list is always an empty list in order to properly insert\n # a special token to start with.\n layout: PatternLayout\n timesteps: int\n n_q: int\n\n def __post_init__(self):\n assert len(self.layout) > 0\n self._validate_layout()\n self._build_reverted_sequence_scatter_indexes = lru_cache(100)(self._build_reverted_sequence_scatter_indexes)\n self._build_pattern_sequence_scatter_indexes = lru_cache(100)(self._build_pattern_sequence_scatter_indexes)\n logger.info(\"New pattern, time steps: %d, sequence steps: %d\", self.timesteps, len(self.layout))\n\n def _validate_layout(self):\n \"\"\"Runs checks on the layout to ensure a valid pattern is defined.\n A pattern is considered invalid if:\n - Multiple timesteps for a same codebook are defined in the same sequence step\n - The timesteps for a given codebook are not in ascending order as we advance in the sequence\n (this would mean that we have future timesteps before past timesteps).\n \"\"\"\n q_timesteps = {q: 0 for q in range(self.n_q)}\n for s, seq_coords in enumerate(self.layout):\n if len(seq_coords) > 0:\n qs = set()\n for coord in seq_coords:\n qs.add(coord.q)\n last_q_timestep = q_timesteps[coord.q]\n assert coord.t >= last_q_timestep, \\\n f\"Past timesteps are found in the sequence for codebook = {coord.q} at step {s}\"\n q_timesteps[coord.q] = coord.t\n # each sequence step contains at max 1 coordinate per codebook\n assert len(qs) == len(seq_coords), \\\n f\"Multiple entries for a same codebook are found at step {s}\"\n\n @property\n def num_sequence_steps(self):\n return len(self.layout) - 1\n\n @property\n def max_delay(self):\n max_t_in_seq_coords = 0\n for seq_coords in self.layout[1:]:\n for coords in seq_coords:\n max_t_in_seq_coords = max(max_t_in_seq_coords, coords.t + 1)\n return max_t_in_seq_coords - self.timesteps\n\n @property\n def valid_layout(self):\n valid_step = len(self.layout) - self.max_delay\n return self.layout[:valid_step]\n\n def starts_with_special_token(self):\n return self.layout[0] == []\n\n def get_sequence_coords_with_timestep(self, t: int, q: tp.Optional[int] = None):\n \"\"\"Get codebook coordinates in the layout that corresponds to the specified timestep t\n and optionally to the codebook q. Coordinates are returned as a tuple with the sequence step\n and the actual codebook coordinates.\n \"\"\"\n assert t <= self.timesteps, \"provided timesteps is greater than the pattern's number of timesteps\"\n if q is not None:\n assert q <= self.n_q, \"provided number of codebooks is greater than the pattern's number of codebooks\"\n coords = []\n for s, seq_codes in enumerate(self.layout):\n for code in seq_codes:\n if code.t == t and (q is None or code.q == q):\n coords.append((s, code))\n return coords\n\n def get_steps_with_timestep(self, t: int, q: tp.Optional[int] = None) -> tp.List[int]:\n return [step for step, coords in self.get_sequence_coords_with_timestep(t, q)]\n\n def get_first_step_with_timesteps(self, t: int, q: tp.Optional[int] = None) -> tp.Optional[int]:\n steps_with_timesteps = self.get_steps_with_timestep(t, q)\n return steps_with_timesteps[0] if len(steps_with_timesteps) > 0 else None\n\n def _build_pattern_sequence_scatter_indexes(self, timesteps: int, n_q: int, keep_only_valid_steps: bool,\n device: tp.Union[torch.device, str] = 'cpu'):\n \"\"\"Build scatter indexes corresponding to the pattern, up to the provided sequence_steps.\n\n Args:\n timesteps (int): Maximum number of timesteps steps to consider.\n keep_only_valid_steps (bool): Restrict the pattern layout to match only valid steps.\n device (torch.device or str): Device for created tensors.\n Returns:\n indexes (torch.Tensor): Indexes corresponding to the sequence, of shape [K, S].\n mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes, of shape [K, S].\n \"\"\"\n assert n_q == self.n_q, f\"invalid number of codebooks for the sequence and the pattern: {n_q} != {self.n_q}\"\n assert timesteps <= self.timesteps, \"invalid number of timesteps used to build the sequence from the pattern\"\n # use the proper layout based on whether we limit ourselves to valid steps only or not,\n # note that using the valid_layout will result in a truncated sequence up to the valid steps\n ref_layout = self.valid_layout if keep_only_valid_steps else self.layout\n # single item indexing being super slow with pytorch vs. numpy, so we use numpy here\n indexes = torch.zeros(n_q, len(ref_layout), dtype=torch.long).numpy()\n mask = torch.zeros(n_q, len(ref_layout), dtype=torch.bool).numpy()\n # fill indexes with last sequence step value that will correspond to our special token\n # the last value is n_q * timesteps as we have flattened z and append special token as the last token\n # which will correspond to the index: n_q * timesteps\n indexes[:] = n_q * timesteps\n # iterate over the pattern and fill scattered indexes and mask\n for s, sequence_coords in enumerate(ref_layout):\n for coords in sequence_coords:\n if coords.t < timesteps:\n indexes[coords.q, s] = coords.t + coords.q * timesteps\n mask[coords.q, s] = 1\n indexes = torch.from_numpy(indexes).to(device)\n mask = torch.from_numpy(mask).to(device)\n return indexes, mask\n\n def build_pattern_sequence(self, z: torch.Tensor, special_token: int, keep_only_valid_steps: bool = False):\n \"\"\"Build sequence corresponding to the pattern from the input tensor z.\n The sequence is built using up to sequence_steps if specified, and non-pattern\n coordinates are filled with the special token.\n\n Args:\n z (torch.Tensor): Input tensor of multi-codebooks sequence, of shape [B, K, T].\n special_token (int): Special token used to fill non-pattern coordinates in the new sequence.\n keep_only_valid_steps (bool): Build a sequence from the pattern up to valid (= fully defined) steps.\n Steps that are beyond valid steps will be replaced by the special_token in that case.\n Returns:\n values (torch.Tensor): Interleaved sequence matching the pattern, of shape [B, K, S] with S\n corresponding either to the sequence_steps if provided, otherwise to the length of the pattern.\n indexes (torch.Tensor): Indexes corresponding to the interleaved sequence, of shape [K, S].\n mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, S].\n \"\"\"\n B, K, T = z.shape\n indexes, mask = self._build_pattern_sequence_scatter_indexes(\n T, K, keep_only_valid_steps=keep_only_valid_steps, device=str(z.device)\n )\n z = z.view(B, -1)\n # we append the special token as the last index of our flattened z tensor\n z = torch.cat([z, torch.zeros_like(z[:, :1]) + special_token], dim=1)\n values = z[:, indexes.view(-1)]\n values = values.view(B, K, indexes.shape[-1])\n return values, indexes, mask\n\n def _build_reverted_sequence_scatter_indexes(self, sequence_steps: int, n_q: int,\n keep_only_valid_steps: bool = False,\n is_model_output: bool = False,\n device: tp.Union[torch.device, str] = 'cpu'):\n \"\"\"Builds scatter indexes required to retrieve the original multi-codebook sequence\n from interleaving pattern.\n\n Args:\n sequence_steps (int): Sequence steps.\n n_q (int): Number of codebooks.\n keep_only_valid_steps (bool): Build a sequence from the pattern up to valid (= fully defined) steps.\n Steps that are beyond valid steps will be replaced by the special_token in that case.\n is_model_output (bool): Whether to keep the sequence item corresponding to initial special token or not.\n device (torch.device or str): Device for created tensors.\n Returns:\n indexes (torch.Tensor): Indexes for reconstructing the output, of shape [K, T].\n mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, T].\n \"\"\"\n ref_layout = self.valid_layout if keep_only_valid_steps else self.layout\n # TODO(jade): Do we want to further truncate to only valid timesteps here as well?\n timesteps = self.timesteps\n assert n_q == self.n_q, f\"invalid number of codebooks for the sequence and the pattern: {n_q} != {self.n_q}\"\n assert sequence_steps <= len(ref_layout), \\\n f\"sequence to revert is longer than the defined pattern: {sequence_steps} > {len(ref_layout)}\"\n\n # ensure we take the appropriate indexes to keep the model output from the first special token as well\n if is_model_output and self.starts_with_special_token():\n ref_layout = ref_layout[1:]\n\n # single item indexing being super slow with pytorch vs. numpy, so we use numpy here\n indexes = torch.zeros(n_q, timesteps, dtype=torch.long).numpy()\n mask = torch.zeros(n_q, timesteps, dtype=torch.bool).numpy()\n # fill indexes with last sequence step value that will correspond to our special token\n indexes[:] = n_q * sequence_steps\n for s, sequence_codes in enumerate(ref_layout):\n if s < sequence_steps:\n for code in sequence_codes:\n if code.t < timesteps:\n indexes[code.q, code.t] = s + code.q * sequence_steps\n mask[code.q, code.t] = 1\n indexes = torch.from_numpy(indexes).to(device)\n mask = torch.from_numpy(mask).to(device)\n return indexes, mask\n\n def revert_pattern_sequence(self, s: torch.Tensor, special_token: int, keep_only_valid_steps: bool = False):\n \"\"\"Revert a sequence built from the pattern back to the original multi-codebook sequence without interleaving.\n The sequence is reverted using up to timesteps if specified, and non-pattern coordinates\n are filled with the special token.\n\n Args:\n s (torch.Tensor): Interleaved sequence tensor obtained from the pattern, of shape [B, K, S].\n special_token (int or float): Special token used to fill non-pattern coordinates in the new sequence.\n Returns:\n values (torch.Tensor): Interleaved sequence matching the pattern, of shape [B, K, T] with T\n corresponding either to the timesteps if provided, or the total timesteps in pattern otherwise.\n indexes (torch.Tensor): Indexes corresponding to the interleaved sequence, of shape [K, T].\n mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, T].\n \"\"\"\n B, K, S = s.shape\n indexes, mask = self._build_reverted_sequence_scatter_indexes(\n S, K, keep_only_valid_steps, is_model_output=False, device=str(s.device)\n )\n s = s.view(B, -1)\n # we append the special token as the last index of our flattened z tensor\n s = torch.cat([s, torch.zeros_like(s[:, :1]) + special_token], dim=1)\n values = s[:, indexes.view(-1)]\n values = values.view(B, K, indexes.shape[-1])\n return values, indexes, mask\n\n def revert_pattern_logits(self, logits: torch.Tensor, special_token: float, keep_only_valid_steps: bool = False):\n \"\"\"Revert model logits obtained on a sequence built from the pattern\n back to a tensor matching the original sequence.\n\n This method is similar to ``revert_pattern_sequence`` with the following specificities:\n 1. It is designed to work with the extra cardinality dimension\n 2. We return the logits for the first sequence item that matches the special_token and\n which matching target in the original sequence is the first item of the sequence,\n while we skip the last logits as there is no matching target\n \"\"\"\n B, card, K, S = logits.shape\n indexes, mask = self._build_reverted_sequence_scatter_indexes(\n S, K, keep_only_valid_steps, is_model_output=True, device=logits.device\n )\n logits = logits.reshape(B, card, -1)\n # we append the special token as the last index of our flattened z tensor\n logits = torch.cat([logits, torch.zeros_like(logits[:, :, :1]) + special_token], dim=-1) # [B, card, K x S]\n values = logits[:, :, indexes.view(-1)]\n values = values.view(B, card, K, indexes.shape[-1])\n return values, indexes, mask\n\n\nclass CodebooksPatternProvider(ABC):\n \"\"\"Abstraction around providing pattern for interleaving codebooks.\n\n The CodebooksPatternProvider abstraction allows to implement various strategies to\n define interleaving pattern of sequences composed of multiple codebooks. For a given\n number of codebooks `n_q`, the pattern provider can generate a specified pattern\n corresponding to a sequence of `T` timesteps with `n_q` parallel codebooks. This pattern\n can be used to construct a new sequence from the original codes respecting the specified\n pattern. The pattern is defined as a list of list of code coordinates, code coordinate\n being a tuple with the original timestep and codebook to build the new sequence.\n Note that all patterns must start with an empty list that is then used to insert a first\n sequence step of special tokens in the newly generated sequence.\n\n Args:\n n_q (int): number of codebooks.\n cached (bool): if True, patterns for a given length are cached. In general\n that should be true for efficiency reason to avoid synchronization points.\n \"\"\"\n def __init__(self, n_q: int, cached: bool = True):\n assert n_q > 0\n self.n_q = n_q\n self.get_pattern = lru_cache(100)(self.get_pattern) # type: ignore\n\n @abstractmethod\n def get_pattern(self, timesteps: int) -> Pattern:\n \"\"\"Builds pattern with specific interleaving between codebooks.\n\n Args:\n timesteps (int): Total number of timesteps.\n \"\"\"\n raise NotImplementedError()\n\n\nclass DelayedPatternProvider(CodebooksPatternProvider):\n \"\"\"Provider for delayed pattern across delayed codebooks.\n Codebooks are delayed in the sequence and sequence steps will contain codebooks\n from different timesteps.\n\n Example:\n Taking timesteps=4 and n_q=3, delays=None, the multi-codebook sequence:\n [[1, 2, 3, 4],\n [1, 2, 3, 4],\n [1, 2, 3, 4]]\n The resulting sequence obtained from the returned pattern is:\n [[S, 1, 2, 3, 4],\n [S, S, 1, 2, 3],\n [S, S, S, 1, 2]]\n (with S being a special token)\n\n Args:\n n_q (int): Number of codebooks.\n delays (list of int, optional): Delay for each of the codebooks.\n If delays not defined, each codebook is delayed by 1 compared to the previous one.\n flatten_first (int): Flatten the first N timesteps.\n empty_initial (int): Prepend with N empty list of coordinates.\n \"\"\"\n def __init__(self, n_q: int, delays: tp.Optional[tp.List[int]] = None,\n flatten_first: int = 0, empty_initial: int = 0):\n super().__init__(n_q)\n if delays is None:\n delays = list(range(n_q))\n self.delays = delays\n self.flatten_first = flatten_first\n self.empty_initial = empty_initial\n assert len(self.delays) == self.n_q\n assert sorted(self.delays) == self.delays\n\n def get_pattern(self, timesteps: int) -> Pattern:\n omit_special_token = self.empty_initial < 0\n out: PatternLayout = [] if omit_special_token else [[]]\n max_delay = max(self.delays)\n if self.empty_initial:\n out += [[] for _ in range(self.empty_initial)]\n if self.flatten_first:\n for t in range(min(timesteps, self.flatten_first)):\n for q in range(self.n_q):\n out.append([LayoutCoord(t, q)])\n for t in range(self.flatten_first, timesteps + max_delay):\n v = []\n for q, delay in enumerate(self.delays):\n t_for_q = t - delay\n if t_for_q >= self.flatten_first:\n v.append(LayoutCoord(t_for_q, q))\n out.append(v)\n return Pattern(out, n_q=self.n_q, timesteps=timesteps)\n\n\nclass ParallelPatternProvider(DelayedPatternProvider):\n \"\"\"Provider for parallel pattern across codebooks.\n This pattern provider is a special case of the delayed pattern with actually no delay,\n hence delays=repeat(0, n_q).\n\n Args:\n n_q (int): Number of codebooks.\n empty_initial (int): Prepend with N empty list of coordinates.\n \"\"\"\n def __init__(self, n_q: int, empty_initial: int = 0):\n super().__init__(n_q, [0] * n_q, empty_initial=empty_initial)\n\n\nclass UnrolledPatternProvider(CodebooksPatternProvider):\n \"\"\"Provider for unrolling codebooks pattern.\n This pattern provider enables to represent the codebook flattened completely or only to some extend\n while also specifying a given delay between the flattened codebooks representation, allowing to\n unroll the codebooks in the sequence.\n\n Example:\n 1. Flattening of the codebooks.\n By default, the pattern provider will fully flatten the codebooks such as flattening=range(n_q),\n taking n_q = 3 and timesteps = 4:\n [[1, 2, 3, 4],\n [1, 2, 3, 4],\n [1, 2, 3, 4]]\n will result into:\n [[S, S, 1, S, S, 2, S, S, 3, S, S, 4],\n [S, 1, S, S, 2, S, S, 3, S, S, 4, S],\n [1, S, S, 2, S, S, 3, S, S, 4, S, S]]\n 2. Partial flattening of the codebooks. The ``flattening`` parameter allows to specify the inner step\n for each of the codebook, allowing to define which codebook to flatten (or keep in parallel), for example\n taking n_q = 3, timesteps = 4 and flattening = [0, 1, 1]:\n [[1, 2, 3, 4],\n [1, 2, 3, 4],\n [1, 2, 3, 4]]\n will result into:\n [[S, 1, S, S, 2, S, S, 3, S, S, 4, S],\n [S, 1, S, S, 2, S, S, 3, S, S, 4, S],\n [1, S, S, 2, S, S, 3, S, S, 4, S, S]]\n 3. Flattening with delay. The ``delay`` parameter allows to further unroll the sequence of codebooks\n allowing to specify the delay per codebook. Note that the delay between codebooks flattened to the\n same inner timestep should be coherent. For example, taking n_q = 3, timesteps = 4, flattening = [0, 1, 1]\n and delays = [0, 3, 3]:\n [[1, 2, 3, 4],\n [1, 2, 3, 4],\n [1, 2, 3, 4]]\n will result into:\n [[S, S, S, 1, S, 2, S, 3, S, 4],\n [S, S, S, 1, S, 2, S, 3, S, 4],\n [1, 2, 3, S, 4, S, 5, S, 6, S]]\n\n Args:\n n_q (int): Number of codebooks.\n flattening (list of int, optional): Flattening schema over the codebooks. If not defined,\n the codebooks will be flattened to 1 codebook per step, meaning that the sequence will\n have n_q extra steps for each timestep.\n delays (list of int, optional): Delay for each of the codebooks. If not defined,\n no delay is added and therefore will default to [0] * ``n_q``.\n Note that two codebooks that will be flattened to the same inner step\n should have the same delay, otherwise the pattern is considered as invalid.\n \"\"\"\n FlattenedCodebook = namedtuple('FlattenedCodebook', ['codebooks', 'delay'])\n\n def __init__(self, n_q: int, flattening: tp.Optional[tp.List[int]] = None,\n delays: tp.Optional[tp.List[int]] = None):\n super().__init__(n_q)\n if flattening is None:\n flattening = list(range(n_q))\n if delays is None:\n delays = [0] * n_q\n assert len(flattening) == n_q\n assert len(delays) == n_q\n assert sorted(flattening) == flattening\n assert sorted(delays) == delays\n self._flattened_codebooks = self._build_flattened_codebooks(delays, flattening)\n self.max_delay = max(delays)\n\n def _build_flattened_codebooks(self, delays: tp.List[int], flattening: tp.List[int]):\n \"\"\"Build a flattened codebooks representation as a dictionary of inner step\n and the actual codebook indices corresponding to the flattened codebook. For convenience, we\n also store the delay associated to the flattened codebook to avoid maintaining an extra mapping.\n \"\"\"\n flattened_codebooks: dict = {}\n for q, (inner_step, delay) in enumerate(zip(flattening, delays)):\n if inner_step not in flattened_codebooks:\n flat_codebook = UnrolledPatternProvider.FlattenedCodebook(codebooks=[q], delay=delay)\n else:\n flat_codebook = flattened_codebooks[inner_step]\n assert flat_codebook.delay == delay, (\n \"Delay and flattening between codebooks is inconsistent: \",\n \"two codebooks flattened to the same position should have the same delay.\"\n )\n flat_codebook.codebooks.append(q)\n flattened_codebooks[inner_step] = flat_codebook\n return flattened_codebooks\n\n @property\n def _num_inner_steps(self):\n \"\"\"Number of inner steps to unroll between timesteps in order to flatten the codebooks.\n \"\"\"\n return max([inner_step for inner_step in self._flattened_codebooks.keys()]) + 1\n\n def num_virtual_steps(self, timesteps: int) -> int:\n return timesteps * self._num_inner_steps + 1\n\n def get_pattern(self, timesteps: int) -> Pattern:\n \"\"\"Builds pattern for delay across codebooks.\n\n Args:\n timesteps (int): Total number of timesteps.\n \"\"\"\n # the PatternLayout is built as a tuple of sequence position and list of coordinates\n # so that it can be reordered properly given the required delay between codebooks of given timesteps\n indexed_out: list = [(-1, [])]\n max_timesteps = timesteps + self.max_delay\n for t in range(max_timesteps):\n # for each timestep, we unroll the flattened codebooks,\n # emitting the sequence step with the corresponding delay\n for step in range(self._num_inner_steps):\n if step in self._flattened_codebooks:\n # we have codebooks at this virtual step to emit\n step_codebooks = self._flattened_codebooks[step]\n t_for_q = t + step_codebooks.delay\n coords = [LayoutCoord(t, q) for q in step_codebooks.codebooks]\n if t_for_q < max_timesteps and t < max_timesteps:\n indexed_out.append((t_for_q, coords))\n else:\n # there is no codebook in this virtual step so we emit an empty list\n indexed_out.append((t, []))\n out = [coords for _, coords in sorted(indexed_out)]\n return Pattern(out, n_q=self.n_q, timesteps=timesteps)\n\n\nclass CoarseFirstPattern(CodebooksPatternProvider):\n \"\"\"First generates all the codebooks #1 (e.g. coarser), then the remaining ones,\n potentially with delays.\n\n ..Warning:: You must always generate the full training duration at test time, for instance,\n 30 seconds, as otherwise, the fine codebooks will start being generated in an unexpected\n location. This is due to the non causality of the remaining codebooks with respect to\n the first ones.\n\n Args:\n n_q (int): Number of codebooks.\n delays (list of int, optional): Delay for each of the codebooks.\n If delays not defined, each codebook is delayed by 1 compared to the previous one.\n \"\"\"\n def __init__(self, n_q: int, delays: tp.Optional[tp.List[int]] = None):\n super().__init__(n_q)\n if delays is None:\n delays = [0] * (n_q - 1)\n self.delays = delays\n assert len(self.delays) == self.n_q - 1\n assert sorted(self.delays) == self.delays\n\n def get_pattern(self, timesteps: int) -> Pattern:\n out: PatternLayout = [[]]\n for t in range(timesteps):\n out.append([LayoutCoord(t, 0)])\n max_delay = max(self.delays)\n for t in range(timesteps + max_delay):\n v = []\n for q, delay in enumerate(self.delays):\n t_for_q = t - delay\n if t_for_q >= 0:\n v.append(LayoutCoord(t_for_q, q + 1))\n out.append(v)\n return Pattern(out, n_q=self.n_q, timesteps=timesteps)\n\n\nclass MusicLMPattern(CodebooksPatternProvider):\n \"\"\"Almost MusicLM style pattern. This is equivalent to full flattening\n but in a different order.\n\n Args:\n n_q (int): Number of codebooks.\n group_by (int): Number of codebooks to group together.\n \"\"\"\n def __init__(self, n_q: int, group_by: int = 2):\n super().__init__(n_q)\n self.group_by = group_by\n\n def get_pattern(self, timesteps: int) -> Pattern:\n out: PatternLayout = [[]]\n for offset in range(0, self.n_q, self.group_by):\n for t in range(timesteps):\n for q in range(offset, offset + self.group_by):\n out.append([LayoutCoord(t, q)])\n return Pattern(out, n_q=self.n_q, timesteps=timesteps)\n" }, { "path": "audiocraft/modules/conditioners.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom collections import defaultdict\nfrom copy import deepcopy\nfrom dataclasses import dataclass, field\nfrom itertools import chain\nimport logging\nimport math\nfrom pathlib import Path\nimport random\nimport re\nimport typing as tp\nimport warnings\nimport einops\nimport flashy\nfrom num2words import num2words\nimport spacy\nfrom transformers import RobertaTokenizer, T5EncoderModel, T5Tokenizer # type: ignore\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.nn.utils.rnn import pad_sequence\nfrom enum import Enum\nfrom .chroma import ChromaExtractor\nfrom .streaming import StreamingModule\nfrom .transformer import create_sin_embedding, StreamingTransformer\nfrom ..data.audio import audio_read\nfrom ..data.audio_dataset import SegmentInfo\nfrom ..data.audio_utils import convert_audio\nfrom ..environment import AudioCraftEnvironment\nfrom ..quantization import ResidualVectorQuantizer\nfrom ..utils.autocast import TorchAutocast\nfrom ..utils.cache import EmbeddingCache\nfrom ..utils.utils import collate, hash_trick, length_to_mask, load_clap_state_dict, warn_once\n\n\nlogger = logging.getLogger(__name__)\nTextCondition = tp.Optional[str] # a text condition can be a string or None (if doesn't exist)\nConditionType = tp.Tuple[torch.Tensor, torch.Tensor] # condition, mask\n\n\nclass JascoCondConst(Enum):\n DRM = 'self_wav'\n CRD = 'chords'\n MLD = 'melody'\n SYM = {'chords', 'melody'}\n LAT = {'self_wav'}\n ALL = ['chords', 'self_wav', 'melody'] # order matters\n\n\nclass WavCondition(tp.NamedTuple):\n wav: torch.Tensor\n length: torch.Tensor\n sample_rate: tp.List[int]\n path: tp.List[tp.Optional[str]] = []\n seek_time: tp.List[tp.Optional[float]] = []\n\n\nclass JointEmbedCondition(tp.NamedTuple):\n wav: torch.Tensor\n text: tp.List[tp.Optional[str]]\n length: torch.Tensor\n sample_rate: tp.List[int]\n path: tp.List[tp.Optional[str]] = []\n seek_time: tp.List[tp.Optional[float]] = []\n\n\nclass SymbolicCondition(tp.NamedTuple):\n frame_chords: tp.Optional[torch.Tensor] = None\n melody: tp.Optional[torch.Tensor] = None\n\n\n@dataclass\nclass ConditioningAttributes:\n text: tp.Dict[str, tp.Optional[str]] = field(default_factory=dict)\n wav: tp.Dict[str, WavCondition] = field(default_factory=dict)\n joint_embed: tp.Dict[str, JointEmbedCondition] = field(default_factory=dict)\n symbolic: tp.Dict[str, SymbolicCondition] = field(default_factory=dict)\n\n def __getitem__(self, item):\n return getattr(self, item)\n\n @property\n def text_attributes(self):\n return self.text.keys()\n\n @property\n def wav_attributes(self):\n return self.wav.keys()\n\n @property\n def joint_embed_attributes(self):\n return self.joint_embed.keys()\n\n @property\n def symbolic_attributes(self):\n return self.symbolic.keys()\n\n @property\n def attributes(self):\n return {\n \"text\": self.text_attributes,\n \"wav\": self.wav_attributes,\n \"joint_embed\": self.joint_embed_attributes,\n \"symbolic\": self.symbolic_attributes,\n }\n\n def to_flat_dict(self):\n return {\n **{f\"text.{k}\": v for k, v in self.text.items()},\n **{f\"wav.{k}\": v for k, v in self.wav.items()},\n **{f\"joint_embed.{k}\": v for k, v in self.joint_embed.items()},\n **{f\"symbolic.{k}\": v for k, v in self.symbolic.items()}\n }\n\n @classmethod\n def from_flat_dict(cls, x):\n out = cls()\n for k, v in x.items():\n kind, att = k.split(\".\")\n out[kind][att] = v\n return out\n\n\nclass SegmentWithAttributes(SegmentInfo):\n \"\"\"Base class for all dataclasses that are used for conditioning.\n All child classes should implement `to_condition_attributes` that converts\n the existing attributes to a dataclass of type ConditioningAttributes.\n \"\"\"\n def to_condition_attributes(self) -> ConditioningAttributes:\n raise NotImplementedError()\n\n\ndef nullify_condition(condition: ConditionType, dim: int = 1):\n \"\"\"Transform an input condition to a null condition.\n The way it is done by converting it to a single zero vector similarly\n to how it is done inside WhiteSpaceTokenizer and NoopTokenizer.\n\n Args:\n condition (ConditionType): A tuple of condition and mask (tuple[torch.Tensor, torch.Tensor])\n dim (int): The dimension that will be truncated (should be the time dimension)\n WARNING!: dim should not be the batch dimension!\n Returns:\n ConditionType: A tuple of null condition and mask\n \"\"\"\n assert dim != 0, \"dim cannot be the batch dimension!\"\n assert isinstance(condition, tuple) and \\\n isinstance(condition[0], torch.Tensor) and \\\n isinstance(condition[1], torch.Tensor), \"'nullify_condition' got an unexpected input type!\"\n cond, mask = condition\n B = cond.shape[0]\n last_dim = cond.dim() - 1\n out = cond.transpose(dim, last_dim)\n out = 0. * out[..., :1]\n out = out.transpose(dim, last_dim)\n mask = torch.zeros((B, 1), device=out.device).int()\n assert cond.dim() == out.dim()\n return out, mask\n\n\ndef nullify_wav(cond: WavCondition) -> WavCondition:\n \"\"\"Transform a WavCondition to a nullified WavCondition.\n It replaces the wav by a null tensor, forces its length to 0, and replaces metadata by dummy attributes.\n\n Args:\n cond (WavCondition): Wav condition with wav, tensor of shape [B, T].\n Returns:\n WavCondition: Nullified wav condition.\n \"\"\"\n null_wav, _ = nullify_condition((cond.wav, torch.zeros_like(cond.wav)), dim=cond.wav.dim() - 1)\n return WavCondition(\n wav=null_wav,\n length=torch.tensor([0] * cond.wav.shape[0], device=cond.wav.device),\n sample_rate=cond.sample_rate,\n path=[None] * cond.wav.shape[0],\n seek_time=[None] * cond.wav.shape[0],\n )\n\n\ndef nullify_joint_embed(embed: JointEmbedCondition) -> JointEmbedCondition:\n \"\"\"Nullify the joint embedding condition by replacing it by a null tensor, forcing its length to 0,\n and replacing metadata by dummy attributes.\n\n Args:\n cond (JointEmbedCondition): Joint embedding condition with wav and text, wav tensor of shape [B, C, T].\n \"\"\"\n null_wav, _ = nullify_condition((embed.wav, torch.zeros_like(embed.wav)), dim=embed.wav.dim() - 1)\n return JointEmbedCondition(\n wav=null_wav, text=[None] * len(embed.text),\n length=torch.LongTensor([0]).to(embed.wav.device),\n sample_rate=embed.sample_rate,\n path=[None] * embed.wav.shape[0],\n seek_time=[0] * embed.wav.shape[0],\n )\n\n\ndef nullify_chords(sym_cond: SymbolicCondition, null_chord_idx: int = 194) -> SymbolicCondition:\n \"\"\"Nullify the symbolic condition by setting all frame chords to a specified null chord index.\n Args:\n sym_cond (SymbolicCondition): The symbolic condition containing frame chords to be nullified.\n null_chord_idx (int, optional): The index to use for nullifying the chords. Defaults to 194 (Chordino).\n Returns:\n SymbolicCondition: A new symbolic condition with all frame chords set to the null chord index.\n \"\"\"\n return SymbolicCondition(frame_chords=torch.ones_like(sym_cond.frame_chords) * null_chord_idx) # type: ignore\n\n\ndef nullify_melody(sym_cond: SymbolicCondition) -> SymbolicCondition:\n \"\"\"Nullify the symbolic condition by replacing the melody matrix with zeros matrix.\n Args:\n sym_cond (SymbolicCondition): The symbolic condition containing frame chords to be nullified.\n null_chord_idx (int, optional): The index to use for nullifying the chords. Defaults to 194 (Chordino).\n Returns:\n SymbolicCondition: A new symbolic condition with all frame chords set to the null chord index.\n \"\"\"\n return SymbolicCondition(melody=torch.zeros_like(sym_cond.melody)) # type: ignore\n\n\ndef _drop_description_condition(conditions: tp.List[ConditioningAttributes]) -> tp.List[ConditioningAttributes]:\n \"\"\"Drop the text condition but keep the wav conditon on a list of ConditioningAttributes.\n This is useful to calculate l_style in the double classifier free guidance formula.\n See paragraph 4.3 in https://arxiv.org/pdf/2407.12563\n\n Args:\n conditions (tp.List[ConditioningAttributes]): List of conditions.\n \"\"\"\n # We assert that description and self_wav are in the conditions\n for condition in conditions:\n assert 'description' in condition.text.keys()\n assert 'self_wav' in condition.wav.keys()\n return AttributeDropout(p={'text': {'description': 1.0},\n 'wav': {'self_wav': 0.0}})(conditions)\n\n\nclass Tokenizer:\n \"\"\"Base tokenizer implementation\n (in case we want to introduce more advances tokenizers in the future).\n \"\"\"\n def __call__(self, texts: tp.List[tp.Optional[str]]) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n raise NotImplementedError()\n\n\nclass WhiteSpaceTokenizer(Tokenizer):\n \"\"\"This tokenizer should be used for natural language descriptions.\n For example:\n [\"he didn't, know he's going home.\", 'shorter sentence'] =>\n [[78, 62, 31, 4, 78, 25, 19, 34],\n [59, 77, 0, 0, 0, 0, 0, 0]]\n \"\"\"\n PUNCTUATION = \"?:!.,;\"\n\n def __init__(self, n_bins: int, pad_idx: int = 0, language: str = \"en_core_web_sm\",\n lemma: bool = True, stopwords: bool = True) -> None:\n self.n_bins = n_bins\n self.pad_idx = pad_idx\n self.lemma = lemma\n self.stopwords = stopwords\n try:\n self.nlp = spacy.load(language)\n except IOError:\n spacy.cli.download(language) # type: ignore\n self.nlp = spacy.load(language)\n\n @tp.no_type_check\n def __call__(self, texts: tp.List[tp.Optional[str]],\n return_text: bool = False) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Take a list of strings and convert them to a tensor of indices.\n\n Args:\n texts (list[str]): List of strings.\n return_text (bool, optional): Whether to return text as additional tuple item. Defaults to False.\n Returns:\n tuple[torch.Tensor, torch.Tensor]:\n - Indices of words in the LUT.\n - And a mask indicating where the padding tokens are\n \"\"\"\n output, lengths = [], []\n texts = deepcopy(texts)\n for i, text in enumerate(texts):\n # if current sample doesn't have a certain attribute, replace with pad token\n if text is None:\n output.append(torch.Tensor([self.pad_idx]))\n lengths.append(0)\n continue\n\n # convert numbers to words\n text = re.sub(r\"(\\d+)\", lambda x: num2words(int(x.group(0))), text) # type: ignore\n # normalize text\n text = self.nlp(text) # type: ignore\n # remove stopwords\n if self.stopwords:\n text = [w for w in text if not w.is_stop] # type: ignore\n # remove punctuation\n text = [w for w in text if w.text not in self.PUNCTUATION] # type: ignore\n # lemmatize if needed\n text = [getattr(t, \"lemma_\" if self.lemma else \"text\") for t in text] # type: ignore\n\n texts[i] = \" \".join(text)\n lengths.append(len(text))\n # convert to tensor\n tokens = torch.Tensor([hash_trick(w, self.n_bins) for w in text])\n output.append(tokens)\n\n mask = length_to_mask(torch.IntTensor(lengths)).int()\n padded_output = pad_sequence(output, padding_value=self.pad_idx).int().t()\n if return_text:\n return padded_output, mask, texts # type: ignore\n return padded_output, mask\n\n\nclass NoopTokenizer(Tokenizer):\n \"\"\"This tokenizer should be used for global conditioners such as: artist, genre, key, etc.\n The difference between this and WhiteSpaceTokenizer is that NoopTokenizer does not split\n strings, so \"Jeff Buckley\" will get it's own index. Whereas WhiteSpaceTokenizer will\n split it to [\"Jeff\", \"Buckley\"] and return an index per word.\n\n For example:\n [\"Queen\", \"ABBA\", \"Jeff Buckley\"] => [43, 55, 101]\n [\"Metal\", \"Rock\", \"Classical\"] => [0, 223, 51]\n \"\"\"\n def __init__(self, n_bins: int, pad_idx: int = 0):\n self.n_bins = n_bins\n self.pad_idx = pad_idx\n\n def __call__(self, texts: tp.List[tp.Optional[str]]) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n output, lengths = [], []\n for text in texts:\n # if current sample doesn't have a certain attribute, replace with pad token\n if text is None:\n output.append(self.pad_idx)\n lengths.append(0)\n else:\n output.append(hash_trick(text, self.n_bins))\n lengths.append(1)\n\n tokens = torch.LongTensor(output).unsqueeze(1)\n mask = length_to_mask(torch.IntTensor(lengths)).int()\n return tokens, mask\n\n\nclass BaseConditioner(nn.Module):\n \"\"\"Base model for all conditioner modules.\n We allow the output dim to be different than the hidden dim for two reasons:\n 1) keep our LUTs small when the vocab is large;\n 2) make all condition dims consistent.\n\n Args:\n dim (int): Hidden dim of the model.\n output_dim (int): Output dim of the conditioner.\n \"\"\"\n def __init__(self, dim: int, output_dim: int):\n super().__init__()\n self.dim = dim\n self.output_dim = output_dim\n if self.output_dim > -1: # omit projection when output_dim <= 0\n self.output_proj = nn.Linear(dim, output_dim)\n\n def tokenize(self, *args, **kwargs) -> tp.Any:\n \"\"\"Should be any part of the processing that will lead to a synchronization\n point, e.g. BPE tokenization with transfer to the GPU.\n\n The returned value will be saved and return later when calling forward().\n \"\"\"\n raise NotImplementedError()\n\n def forward(self, inputs: tp.Any) -> ConditionType:\n \"\"\"Gets input that should be used as conditioning (e.g, genre, description or a waveform).\n Outputs a ConditionType, after the input data was embedded as a dense vector.\n\n Returns:\n ConditionType:\n - A tensor of size [B, T, D] where B is the batch size, T is the length of the\n output embedding and D is the dimension of the embedding.\n - And a mask indicating where the padding tokens.\n \"\"\"\n raise NotImplementedError()\n\n\nclass TextConditioner(BaseConditioner):\n ...\n\n\nclass LUTConditioner(TextConditioner):\n \"\"\"Lookup table TextConditioner.\n\n Args:\n n_bins (int): Number of bins.\n dim (int): Hidden dim of the model (text-encoder/LUT).\n output_dim (int): Output dim of the conditioner.\n tokenizer (str): Name of the tokenizer.\n pad_idx (int, optional): Index for padding token. Defaults to 0.\n \"\"\"\n def __init__(self, n_bins: int, dim: int, output_dim: int, tokenizer: str, pad_idx: int = 0):\n super().__init__(dim, output_dim)\n self.embed = nn.Embedding(n_bins, dim)\n self.tokenizer: Tokenizer\n if tokenizer == 'whitespace':\n self.tokenizer = WhiteSpaceTokenizer(n_bins, pad_idx=pad_idx)\n elif tokenizer == 'noop':\n self.tokenizer = NoopTokenizer(n_bins, pad_idx=pad_idx)\n else:\n raise ValueError(f\"unrecognized tokenizer `{tokenizer}`.\")\n\n def tokenize(self, x: tp.List[tp.Optional[str]]) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n device = self.embed.weight.device\n tokens, mask = self.tokenizer(x)\n tokens, mask = tokens.to(device), mask.to(device)\n return tokens, mask\n\n def forward(self, inputs: tp.Tuple[torch.Tensor, torch.Tensor]) -> ConditionType:\n tokens, mask = inputs\n embeds = self.embed(tokens)\n embeds = self.output_proj(embeds)\n embeds = (embeds * mask.unsqueeze(-1))\n return embeds, mask\n\n\nclass T5Conditioner(TextConditioner):\n \"\"\"T5-based TextConditioner.\n\n Args:\n name (str): Name of the T5 model.\n output_dim (int): Output dim of the conditioner.\n finetune (bool): Whether to fine-tune T5 at train time.\n device (str): Device for T5 Conditioner.\n autocast_dtype (tp.Optional[str], optional): Autocast dtype.\n word_dropout (float, optional): Word dropout probability.\n normalize_text (bool, optional): Whether to apply text normalization.\n \"\"\"\n MODELS = [\"t5-small\", \"t5-base\", \"t5-large\", \"t5-3b\", \"t5-11b\",\n \"google/flan-t5-small\", \"google/flan-t5-base\", \"google/flan-t5-large\",\n \"google/flan-t5-xl\", \"google/flan-t5-xxl\"]\n MODELS_DIMS = {\n \"t5-small\": 512,\n \"t5-base\": 768,\n \"t5-large\": 1024,\n \"t5-3b\": 1024,\n \"t5-11b\": 1024,\n \"google/flan-t5-small\": 512,\n \"google/flan-t5-base\": 768,\n \"google/flan-t5-large\": 1024,\n \"google/flan-t5-3b\": 1024,\n \"google/flan-t5-11b\": 1024,\n }\n\n def __init__(self, name: str, output_dim: int, finetune: bool, device: str,\n autocast_dtype: tp.Optional[str] = 'float32', word_dropout: float = 0.,\n normalize_text: bool = False):\n assert name in self.MODELS, f\"Unrecognized t5 model name (should in {self.MODELS})\"\n super().__init__(self.MODELS_DIMS[name], output_dim)\n self.device = device\n self.name = name\n self.finetune = finetune\n self.word_dropout = word_dropout\n if autocast_dtype is None or self.device == 'cpu':\n self.autocast = TorchAutocast(enabled=False)\n if self.device != 'cpu':\n logger.warning(\"T5 has no autocast, this might lead to NaN\")\n else:\n dtype = getattr(torch, autocast_dtype)\n assert isinstance(dtype, torch.dtype)\n logger.info(f\"T5 will be evaluated with autocast as {autocast_dtype}\")\n self.autocast = TorchAutocast(enabled=True, device_type=self.device, dtype=dtype)\n # Let's disable logging temporarily because T5 will vomit some errors otherwise.\n # thanks https://gist.github.com/simon-weber/7853144\n previous_level = logging.root.manager.disable\n logging.disable(logging.ERROR)\n with warnings.catch_warnings():\n warnings.simplefilter(\"ignore\")\n try:\n self.t5_tokenizer = T5Tokenizer.from_pretrained(name)\n t5 = T5EncoderModel.from_pretrained(name).train(mode=finetune)\n finally:\n logging.disable(previous_level)\n if finetune:\n self.t5 = t5\n else:\n # this makes sure that the t5 models is not part\n # of the saved checkpoint\n self.__dict__['t5'] = t5.to(device)\n\n self.normalize_text = normalize_text\n if normalize_text:\n self.text_normalizer = WhiteSpaceTokenizer(1, lemma=True, stopwords=True)\n\n def tokenize(self, x: tp.List[tp.Optional[str]]) -> tp.Dict[str, torch.Tensor]:\n # if current sample doesn't have a certain attribute, replace with empty string\n entries: tp.List[str] = [xi if xi is not None else \"\" for xi in x]\n if self.normalize_text:\n _, _, entries = self.text_normalizer(entries, return_text=True)\n if self.word_dropout > 0. and self.training:\n new_entries = []\n for entry in entries:\n words = [word for word in entry.split(\" \") if random.random() >= self.word_dropout]\n new_entries.append(\" \".join(words))\n entries = new_entries\n\n empty_idx = torch.LongTensor([i for i, xi in enumerate(entries) if xi == \"\"])\n\n inputs = self.t5_tokenizer(entries, return_tensors='pt', padding=True).to(self.device)\n mask = inputs['attention_mask']\n mask[empty_idx, :] = 0 # zero-out index where the input is non-existant\n return inputs\n\n def forward(self, inputs: tp.Dict[str, torch.Tensor]) -> ConditionType:\n mask = inputs['attention_mask']\n with torch.set_grad_enabled(self.finetune), self.autocast:\n embeds = self.t5(**inputs).last_hidden_state\n embeds = self.output_proj(embeds.to(self.output_proj.weight))\n embeds = (embeds * mask.unsqueeze(-1))\n return embeds, mask\n\n\nclass WaveformConditioner(BaseConditioner):\n \"\"\"Base class for all conditioners that take a waveform as input.\n Classes that inherit must implement `_get_wav_embedding` that outputs\n a continuous tensor, and `_downsampling_factor` that returns the down-sampling\n factor of the embedding model.\n\n Args:\n dim (int): The internal representation dimension.\n output_dim (int): Output dimension.\n device (tp.Union[torch.device, str]): Device.\n \"\"\"\n def __init__(self, dim: int, output_dim: int, device: tp.Union[torch.device, str]):\n super().__init__(dim, output_dim)\n self.device = device\n # if False no masking is done, used in ChromaStemConditioner when completing by periodicity a sample.\n self._use_masking = True\n\n def tokenize(self, x: WavCondition) -> WavCondition:\n wav, length, sample_rate, path, seek_time = x\n assert length is not None\n return WavCondition(wav.to(self.device), length.to(self.device), sample_rate, path, seek_time)\n\n def _get_wav_embedding(self, x: WavCondition) -> torch.Tensor:\n \"\"\"Gets as input a WavCondition and returns a dense embedding.\"\"\"\n raise NotImplementedError()\n\n def _downsampling_factor(self):\n \"\"\"Returns the downsampling factor of the embedding model.\"\"\"\n raise NotImplementedError()\n\n def forward(self, x: WavCondition) -> ConditionType:\n \"\"\"Extract condition embedding and mask from a waveform and its metadata.\n Args:\n x (WavCondition): Waveform condition containing raw waveform and metadata.\n Returns:\n ConditionType: a dense vector representing the conditioning along with its mask\n \"\"\"\n wav, lengths, *_ = x\n with torch.no_grad():\n embeds = self._get_wav_embedding(x)\n if hasattr(self, 'output_proj'):\n embeds = embeds.to(self.output_proj.weight)\n embeds = self.output_proj(embeds)\n\n if lengths is not None and self._use_masking:\n lengths = lengths / self._downsampling_factor()\n mask = length_to_mask(lengths, max_len=embeds.shape[1]).int() # type: ignore\n else:\n mask = torch.ones_like(embeds[..., 0])\n embeds = (embeds * mask.unsqueeze(-1))\n return embeds, mask\n\n\nclass ChromaStemConditioner(WaveformConditioner):\n \"\"\"Chroma conditioner based on stems.\n The ChromaStemConditioner uses DEMUCS to first filter out drums and bass, as\n the drums and bass often dominate the chroma leading to the chroma features\n not containing information about the melody.\n\n Args:\n output_dim (int): Output dimension for the conditioner.\n sample_rate (int): Sample rate for the chroma extractor.\n n_chroma (int): Number of chroma bins for the chroma extractor.\n radix2_exp (int): Size of stft window for the chroma extractor (power of 2, e.g. 12 -> 2^12).\n duration (int): duration used during training. This is later used for correct padding\n in case we are using chroma as prefix.\n match_len_on_eval (bool, optional): if True then all chromas are padded to the training\n duration. Defaults to False.\n eval_wavs (str, optional): path to a dataset manifest with waveform, this waveforms are used as\n conditions during eval (for cases where we don't want to leak test conditions like MusicCaps).\n Defaults to None.\n n_eval_wavs (int, optional): limits the number of waveforms used for conditioning. Defaults to 0.\n device (tp.Union[torch.device, str], optional): Device for the conditioner.\n **kwargs: Additional parameters for the chroma extractor.\n \"\"\"\n def __init__(self, output_dim: int, sample_rate: int, n_chroma: int, radix2_exp: int,\n duration: float, match_len_on_eval: bool = True, eval_wavs: tp.Optional[str] = None,\n n_eval_wavs: int = 0, cache_path: tp.Optional[tp.Union[str, Path]] = None,\n device: tp.Union[torch.device, str] = 'cpu', **kwargs):\n from demucs import pretrained\n super().__init__(dim=n_chroma, output_dim=output_dim, device=device)\n self.autocast = TorchAutocast(enabled=device != 'cpu', device_type=self.device, dtype=torch.float32)\n self.sample_rate = sample_rate\n self.match_len_on_eval = match_len_on_eval\n if match_len_on_eval:\n self._use_masking = False\n self.duration = duration\n self.__dict__['demucs'] = pretrained.get_model('htdemucs').to(device)\n stem_sources: list = self.demucs.sources # type: ignore\n self.stem_indices = torch.LongTensor([stem_sources.index('vocals'), stem_sources.index('other')]).to(device)\n self.chroma = ChromaExtractor(sample_rate=sample_rate, n_chroma=n_chroma,\n radix2_exp=radix2_exp, **kwargs).to(device)\n self.chroma_len = self._get_chroma_len()\n self.eval_wavs: tp.Optional[torch.Tensor] = self._load_eval_wavs(eval_wavs, n_eval_wavs)\n self.cache = None\n if cache_path is not None:\n self.cache = EmbeddingCache(Path(cache_path) / 'wav', self.device,\n compute_embed_fn=self._get_full_chroma_for_cache,\n extract_embed_fn=self._extract_chroma_chunk)\n\n def _downsampling_factor(self) -> int:\n return self.chroma.winhop\n\n def _load_eval_wavs(self, path: tp.Optional[str], num_samples: int) -> tp.Optional[torch.Tensor]:\n \"\"\"Load pre-defined waveforms from a json.\n These waveforms will be used for chroma extraction during evaluation.\n This is done to make the evaluation on MusicCaps fair (we shouldn't see the chromas of MusicCaps).\n \"\"\"\n if path is None:\n return None\n\n logger.info(f\"Loading evaluation wavs from {path}\")\n from audiocraft.data.audio_dataset import AudioDataset\n dataset: AudioDataset = AudioDataset.from_meta(\n path, segment_duration=self.duration, min_audio_duration=self.duration,\n sample_rate=self.sample_rate, channels=1)\n\n if len(dataset) > 0:\n eval_wavs = dataset.collater([dataset[i] for i in range(num_samples)]).to(self.device)\n logger.info(f\"Using {len(eval_wavs)} evaluation wavs for chroma-stem conditioner\")\n return eval_wavs\n else:\n raise ValueError(\"Could not find evaluation wavs, check lengths of wavs\")\n\n def reset_eval_wavs(self, eval_wavs: tp.Optional[torch.Tensor]) -> None:\n self.eval_wavs = eval_wavs\n\n def has_eval_wavs(self) -> bool:\n return self.eval_wavs is not None\n\n def _sample_eval_wavs(self, num_samples: int) -> torch.Tensor:\n \"\"\"Sample wavs from a predefined list.\"\"\"\n assert self.eval_wavs is not None, \"Cannot sample eval wavs as no eval wavs provided.\"\n total_eval_wavs = len(self.eval_wavs)\n out = self.eval_wavs\n if num_samples > total_eval_wavs:\n out = self.eval_wavs.repeat(num_samples // total_eval_wavs + 1, 1, 1)\n return out[torch.randperm(len(out))][:num_samples]\n\n def _get_chroma_len(self) -> int:\n \"\"\"Get length of chroma during training.\"\"\"\n dummy_wav = torch.zeros((1, int(self.sample_rate * self.duration)), device=self.device)\n dummy_chr = self.chroma(dummy_wav)\n return dummy_chr.shape[1]\n\n @torch.no_grad()\n def _get_stemmed_wav(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:\n \"\"\"Get parts of the wav that holds the melody, extracting the main stems from the wav.\"\"\"\n from demucs.apply import apply_model\n from demucs.audio import convert_audio\n with self.autocast:\n wav = convert_audio(\n wav, sample_rate, self.demucs.samplerate, self.demucs.audio_channels) # type: ignore\n stems = apply_model(self.demucs, wav, device=self.device) # type: ignore\n stems = stems[:, self.stem_indices] # extract relevant stems for melody conditioning\n mix_wav = stems.sum(1) # merge extracted stems to single waveform\n mix_wav = convert_audio(mix_wav, self.demucs.samplerate, self.sample_rate, 1) # type: ignore\n return mix_wav\n\n @torch.no_grad()\n def _extract_chroma(self, wav: torch.Tensor) -> torch.Tensor:\n \"\"\"Extract chroma features from the waveform.\"\"\"\n with self.autocast:\n return self.chroma(wav)\n\n @torch.no_grad()\n def _compute_wav_embedding(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:\n \"\"\"Compute wav embedding, applying stem and chroma extraction.\"\"\"\n # avoid 0-size tensors when we are working with null conds\n if wav.shape[-1] == 1:\n return self._extract_chroma(wav)\n stems = self._get_stemmed_wav(wav, sample_rate)\n chroma = self._extract_chroma(stems)\n return chroma\n\n @torch.no_grad()\n def _get_full_chroma_for_cache(self, path: tp.Union[str, Path], x: WavCondition, idx: int) -> torch.Tensor:\n \"\"\"Extract chroma from the whole audio waveform at the given path.\"\"\"\n wav, sr = audio_read(path)\n wav = wav[None].to(self.device)\n wav = convert_audio(wav, sr, self.sample_rate, to_channels=1)\n chroma = self._compute_wav_embedding(wav, self.sample_rate)[0]\n return chroma\n\n def _extract_chroma_chunk(self, full_chroma: torch.Tensor, x: WavCondition, idx: int) -> torch.Tensor:\n \"\"\"Extract a chunk of chroma from the full chroma derived from the full waveform.\"\"\"\n wav_length = x.wav.shape[-1]\n seek_time = x.seek_time[idx]\n assert seek_time is not None, (\n \"WavCondition seek_time is required \"\n \"when extracting chroma chunks from pre-computed chroma.\")\n full_chroma = full_chroma.float()\n frame_rate = self.sample_rate / self._downsampling_factor()\n target_length = int(frame_rate * wav_length / self.sample_rate)\n index = int(frame_rate * seek_time)\n out = full_chroma[index: index + target_length]\n out = F.pad(out[None], (0, 0, 0, target_length - out.shape[0]))[0]\n return out.to(self.device)\n\n @torch.no_grad()\n def _get_wav_embedding(self, x: WavCondition) -> torch.Tensor:\n \"\"\"Get the wav embedding from the WavCondition.\n The conditioner will either extract the embedding on-the-fly computing it from the condition wav directly\n or will rely on the embedding cache to load the pre-computed embedding if relevant.\n \"\"\"\n sampled_wav: tp.Optional[torch.Tensor] = None\n if not self.training and self.eval_wavs is not None:\n warn_once(logger, \"Using precomputed evaluation wavs!\")\n sampled_wav = self._sample_eval_wavs(len(x.wav))\n\n no_undefined_paths = all(p is not None for p in x.path)\n no_nullified_cond = x.wav.shape[-1] > 1\n if sampled_wav is not None:\n chroma = self._compute_wav_embedding(sampled_wav, self.sample_rate)\n elif self.cache is not None and no_undefined_paths and no_nullified_cond:\n paths = [Path(p) for p in x.path if p is not None]\n chroma = self.cache.get_embed_from_cache(paths, x)\n else:\n assert all(sr == x.sample_rate[0] for sr in x.sample_rate), \"All sample rates in batch should be equal.\"\n chroma = self._compute_wav_embedding(x.wav, x.sample_rate[0])\n\n if self.match_len_on_eval:\n B, T, C = chroma.shape\n if T > self.chroma_len:\n chroma = chroma[:, :self.chroma_len]\n logger.debug(f\"Chroma was truncated to match length! ({T} -> {chroma.shape[1]})\")\n elif T < self.chroma_len:\n n_repeat = int(math.ceil(self.chroma_len / T))\n chroma = chroma.repeat(1, n_repeat, 1)\n chroma = chroma[:, :self.chroma_len]\n logger.debug(f\"Chroma was repeated to match length! ({T} -> {chroma.shape[1]})\")\n\n return chroma\n\n def tokenize(self, x: WavCondition) -> WavCondition:\n \"\"\"Apply WavConditioner tokenization and populate cache if needed.\"\"\"\n x = super().tokenize(x)\n no_undefined_paths = all(p is not None for p in x.path)\n if self.cache is not None and no_undefined_paths:\n paths = [Path(p) for p in x.path if p is not None]\n self.cache.populate_embed_cache(paths, x)\n return x\n\n\nclass FeatureExtractor(WaveformConditioner):\n \"\"\"\n Feature Extractor used for the style conditioner of the paper AUDIO CONDITIONING\n FOR MUSIC GENERATION VIA DISCRETE BOTTLENECK FEATURES.\n\n Given a waveform, we extract an excerpt of defined length randomly subsampled.\n Then, we feed this excerpt to a feature extractor.\n\n Args:\n model_name (str): 'encodec' or 'mert'.\n sample_rate (str): sample rate of the input audio. (32000)\n encodec_checkpoint (str): if encodec is used as a feature extractor, checkpoint\n of the model. ('//pretrained/facebook/encodec_32khz' is the default)\n encodec_n_q (int): if encodec is used as a feature extractor it sets the number of\n quantization streams used in it.\n length (float): length in seconds of the random subsampled excerpt that is used\n for conditioning.\n dim (int): The internal representation dimension.\n output_dim (int): Output dimension for the conditioner.\n device (tp.Union[torch.device, str], optional): Device for the conditioner.\n compute_mask (bool): whether to mask the tokens corresponding to the subsampled\n excerpt in the computation of the music language model cross-entropy loss.\n use_middle_of_segment (bool): if True, always take the middle of the input\n instead of a random subsampled excerpt.\n ds_rate_compression (int): downsampling parameter of the compression model used\n for the music language model. (640 for encodec_32khz)\n num_codebooks_lm (int): the number of codebooks used by the music language model.\n \"\"\"\n def __init__(\n self, model_name: str,\n sample_rate: int, encodec_checkpoint: str, encodec_n_q: int, length: float,\n dim: int, output_dim: int, device: tp.Union[torch.device, str],\n compute_mask: bool = True,\n use_middle_of_segment: bool = False, ds_rate_compression: int = 640,\n num_codebooks_lm: int = 4\n ):\n assert model_name in ['encodec', 'mert']\n if model_name == 'encodec':\n from ..solvers.compression import CompressionSolver\n feat_extractor = CompressionSolver.model_from_checkpoint(encodec_checkpoint, device)\n elif model_name == 'mert':\n from transformers import AutoModel\n feat_extractor = AutoModel.from_pretrained(\"m-a-p/MERT-v1-95M\", trust_remote_code=True)\n super().__init__(\n dim=dim,\n output_dim=output_dim,\n device=device\n )\n self.sample_rate = sample_rate\n self.compute_mask = compute_mask\n self.feat_extractor: nn.Module\n self.embed: tp.Union[nn.ModuleList, nn.Linear]\n if model_name == 'encodec':\n self.__dict__[\"feat_extractor\"] = feat_extractor.to(device)\n self.encodec_n_q = encodec_n_q\n self.embed = nn.ModuleList([nn.Embedding(feat_extractor.cardinality, dim) for _ in range(encodec_n_q)])\n if model_name == 'mert':\n self.__dict__[\"feat_extractor\"] = feat_extractor.eval().to(device)\n self.embed = nn.Linear(768, dim) # hardcoded\n self.length_subwav = int(length * sample_rate)\n self.ds_rate_compression = ds_rate_compression\n self.model_name = model_name\n self.use_middle_of_segment = use_middle_of_segment\n self.num_codebooks_lm = num_codebooks_lm\n\n def _get_wav_embedding(self, x: WavCondition) -> torch.Tensor:\n if x.wav.shape[-1] == 1:\n self.temp_mask = None\n return torch.zeros(x.wav.shape[0], 1, self.dim, device=self.device)\n else:\n with torch.no_grad():\n if self.use_middle_of_segment:\n start = int((x.wav.shape[-1] - self.length_subwav) / 2)\n wav = x.wav[:, :, start:start+self.length_subwav]\n else:\n start = random.randint(0, x.wav.shape[-1] - self.length_subwav)\n wav = x.wav[:, :, start:start+self.length_subwav]\n if self.compute_mask:\n self.temp_mask = self._get_mask_wav(x, start)\n if self.model_name == 'encodec':\n tokens = self.feat_extractor.encode(wav)[0] # type: ignore\n elif self.model_name == 'mert':\n wav = convert_audio(wav, from_rate=x.sample_rate[0], to_rate=24000, to_channels=1)\n embeds = self.feat_extractor(wav.squeeze(-2)).last_hidden_state\n if self.model_name == 'encodec':\n tokens = tokens[:, :self.encodec_n_q]\n embeds = sum([self.embed[k](tokens[:, k]) for k in range(self.encodec_n_q)]) # type: ignore\n else:\n embeds = self.embed(embeds)\n\n return embeds # [B, T, dim]\n\n def _downsampling_factor(self):\n if self.model_name == 'encodec':\n return self.sample_rate / self.feat_extractor.frame_rate\n elif self.model_name == 'mert':\n return self.sample_rate / 75\n\n def _get_mask_wav(self, x: WavCondition, start: int) -> tp.Union[torch.Tensor, None]:\n if x.wav.shape[-1] == 1:\n return None\n total_length = int(x.wav.shape[-1] / self.ds_rate_compression)\n mask_length = int(self.length_subwav / self.ds_rate_compression)\n start = int(start / self.ds_rate_compression)\n mask = torch.ones(x.wav.shape[0], self.num_codebooks_lm,\n total_length, device=self.device, dtype=torch.bool)\n mask[:, :, start:start+mask_length] = 0\n return mask\n\n\nclass StyleConditioner(FeatureExtractor):\n \"\"\"Conditioner from the paper AUDIO CONDITIONING FOR MUSIC GENERATION VIA\n DISCRETE BOTTLENECK FEATURES.\n Given an audio input, it is passed through a Feature Extractor and a\n transformer encoder. Then it is quantized through RVQ.\n\n Args:\n transformer_scale (str): size of the transformer. See in the __init__ to have more infos.\n ds_factor (int): the downsampling factor applied to the representation after quantization.\n encodec_n_q (int): if encodec is used as a feature extractor it sets the number of\n quantization streams used in it.\n n_q_out (int): the number of quantization streams used for the RVQ. If increased, there\n is more information passing as a conditioning.\n eval_q (int): the number of quantization streams used for the RVQ at evaluation time.\n q_dropout (bool): if True, at training time, a random number of stream is sampled\n at each step in the interval [1, n_q_out].\n bins (int): the codebook size used for each quantization stream.\n varying_lengths (List[float]): list of the min and max duration in seconds for the\n randomly subsampled excerpt at training time. For each step a length is sampled\n in this interval.\n batch_norm (bool): use of batch normalization after the transformer. Stabilizes the\n training.\n rvq_threshold_ema_dead_code (float): threshold for dropping dead codes in the\n RVQ.\n \"\"\"\n def __init__(self, transformer_scale: str = 'default', ds_factor: int = 15, encodec_n_q: int = 4,\n n_q_out: int = 6, eval_q: int = 3, q_dropout: bool = True, bins: int = 1024,\n varying_lengths: tp.List[float] = [1.5, 4.5],\n batch_norm: bool = True, rvq_threshold_ema_dead_code: float = 0.1,\n **kwargs):\n tr_args: tp.Dict[str, tp.Any]\n if transformer_scale == 'xsmall':\n tr_args = {'d_model': 256, 'num_heads': 8, 'num_layers': 4}\n elif transformer_scale == 'large':\n tr_args = {'d_model': 1024, 'num_heads': 16, 'num_layers': 24}\n elif transformer_scale == 'default':\n tr_args = {'d_model': 512, 'num_heads': 8, 'num_layers': 8}\n elif transformer_scale == 'none':\n tr_args = {'d_model': 512}\n tr_args.update({\n 'memory_efficient': True, 'activation': 'gelu',\n 'norm_first': True, 'causal': False, 'layer_scale': None,\n 'bias_ff': False, 'bias_attn': False,\n })\n dim = tr_args['d_model']\n super().__init__(dim=dim, encodec_n_q=encodec_n_q, **kwargs)\n\n self.ds_factor = ds_factor\n if transformer_scale == 'none':\n self.transformer = None\n else:\n self.transformer = StreamingTransformer(dim_feedforward=int(4 * dim), **tr_args)\n self.n_q_out = n_q_out\n self.eval_q = eval_q\n self.rvq = None\n if n_q_out > 0:\n self.rvq = ResidualVectorQuantizer(dim, n_q=n_q_out, q_dropout=q_dropout, bins=bins,\n threshold_ema_dead_code=rvq_threshold_ema_dead_code)\n self.autocast = TorchAutocast(enabled=self.device != 'cpu', device_type=self.device, dtype=torch.float32)\n self.varying_lengths = varying_lengths\n self.batch_norm = None\n if batch_norm:\n self.batch_norm = nn.BatchNorm1d(dim, affine=False)\n self.mask = None\n\n def _get_wav_embedding(self, wav: WavCondition) -> torch.Tensor:\n with self.autocast:\n # Sample the length of the excerpts\n if self.varying_lengths and self.training:\n assert len(self.varying_lengths) == 2\n length = random.uniform(self.varying_lengths[0], self.varying_lengths[1])\n self.length_subwav = int(length * self.sample_rate)\n z1 = super()._get_wav_embedding(wav)\n if self.compute_mask:\n self.mask = self.temp_mask # type: ignore\n self.temp_mask = None\n\n if self.transformer is not None:\n out1 = self.transformer(z1)\n else:\n out1 = z1\n if self.batch_norm:\n out1 = self.batch_norm(out1.transpose(1, 2)).transpose(1, 2)\n # Apply quantization\n if self.rvq:\n if self.training:\n self.rvq.set_num_codebooks(self.n_q_out)\n else:\n self.rvq.set_num_codebooks(self.eval_q)\n out1 = self.rvq(out1.transpose(1, 2), frame_rate=1.)\n if self.training:\n flashy.distrib.average_tensors(self.rvq.buffers())\n out1 = out1.x.transpose(1, 2)\n # Apply fix downsample\n out1 = out1[:, ::self.ds_factor]\n\n return out1\n\n def set_params(self, eval_q: int = 3,\n excerpt_length: float = 3.0,\n ds_factor: tp.Optional[int] = None, encodec_n_q: tp.Optional[int] = None):\n \"\"\"Modify the parameters of the SSL or introduce new parameters to add noise to\n the conditioning or to downsample it\n\n Args:\n eval_q (int): number of codebooks used when evaluating the model\n excerpt_length (float): the length of the excerpts used to condition the model\n \"\"\"\n self.eval_q = eval_q\n self.length_subwav = int(excerpt_length * self.sample_rate)\n if ds_factor is not None:\n self.ds_factor = ds_factor\n if encodec_n_q is not None:\n self.encodec_n_q = encodec_n_q\n\n def _downsampling_factor(self):\n df = super()._downsampling_factor()\n return df * self.ds_factor\n\n def forward(self, x: WavCondition) -> ConditionType:\n wav, lengths, *_ = x\n\n embeds = self._get_wav_embedding(x)\n embeds = embeds.to(self.output_proj.weight)\n embeds = self.output_proj(embeds)\n\n lengths = lengths / self._downsampling_factor()\n mask = length_to_mask(lengths, max_len=embeds.shape[1]).int() # type: ignore\n\n embeds = (embeds * mask.unsqueeze(2).to(self.device))\n\n return embeds, mask\n\n\nclass JointEmbeddingConditioner(BaseConditioner):\n \"\"\"Joint embedding conditioning supporting both audio or text conditioning.\n\n Args:\n dim (int): Dimension.\n output_dim (int): Output dimension.\n device (str): Device.\n attribute (str): Attribute used by the conditioner.\n autocast_dtype (str): Autocast for the conditioner.\n quantize (bool): Whether to quantize the CLAP embedding.\n n_q (int): Number of residual quantizers (used if quantize is true).\n bins (int): Quantizers' codebooks size (used if quantize is true).\n kwargs: Additional parameters for residual vector quantizer.\n \"\"\"\n def __init__(self, dim: int, output_dim: int, device: str, attribute: str,\n autocast_dtype: tp.Optional[str] = 'float32', quantize: bool = True,\n n_q: int = 12, bins: int = 1024, **kwargs):\n super().__init__(dim=dim, output_dim=output_dim)\n self.device = device\n self.attribute = attribute\n if autocast_dtype is None or device == 'cpu':\n self.autocast = TorchAutocast(enabled=False)\n logger.warning(\"JointEmbeddingConditioner has no autocast, this might lead to NaN.\")\n else:\n dtype = getattr(torch, autocast_dtype)\n assert isinstance(dtype, torch.dtype)\n logger.info(f\"JointEmbeddingConditioner will be evaluated with autocast as {autocast_dtype}.\")\n self.autocast = TorchAutocast(enabled=True, device_type=self.device, dtype=dtype)\n # residual vector quantizer to discretize the conditioned embedding\n self.quantizer: tp.Optional[ResidualVectorQuantizer] = None\n if quantize:\n self.quantizer = ResidualVectorQuantizer(dim, n_q=n_q, bins=bins, **kwargs)\n\n def _get_embed(self, x: JointEmbedCondition) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Get joint embedding in latent space from the inputs.\n\n Returns:\n tuple[torch.Tensor, torch.Tensor]: Tensor for the latent embedding\n and corresponding empty indexes.\n \"\"\"\n raise NotImplementedError()\n\n def forward(self, x: JointEmbedCondition) -> ConditionType:\n with self.autocast:\n embed, empty_idx = self._get_embed(x)\n if self.quantizer is not None:\n embed = embed.view(-1, self.dim, 1)\n q_res = self.quantizer(embed, frame_rate=1)\n out_embed = q_res.x.view(-1, self.dim)\n else:\n out_embed = embed\n out_embed = self.output_proj(out_embed).view(-1, 1, self.output_dim)\n mask = torch.ones(*out_embed.shape[:2], device=out_embed.device)\n mask[empty_idx, :] = 0 # zero-out index where the input is non-existant\n out_embed = (out_embed * mask.unsqueeze(-1))\n return out_embed, mask\n\n def tokenize(self, x: JointEmbedCondition) -> JointEmbedCondition:\n return x\n\n\nclass CLAPEmbeddingConditioner(JointEmbeddingConditioner):\n \"\"\"Joint Embedding conditioner based on pre-trained CLAP model.\n\n This CLAP-based conditioner supports a caching mechanism\n over the computed embeddings for faster training.\n\n Args:\n dim (int): Dimension.\n output_dim (int): Output dimension.\n device (str): Device.\n attribute (str): Attribute used by the conditioner.\n quantize (bool): Whether to quantize the CLAP embedding.\n n_q (int): Number of residual quantizers (used if quantize is true).\n bins (int): Quantizers' codebooks size (used if quantize is true).\n checkpoint (str): Path to CLAP checkpoint.\n model_arch (str): CLAP model architecture.\n enable_fusion (bool): Enable fusion for CLAP model.\n sample_rate (int): Sample rate used by CLAP model.\n max_audio_length (float): Maximum audio length for CLAP model.\n audio_stride (float): Stride to use for getting a CLAP embedding on the full sequence.\n normalize (bool): Whether to normalize the CLAP embedding.\n text_p (float): Probability of using text representation instead of audio at train time.\n batch_size (Optional[int]): Batch size for CLAP embedding computation.\n autocast_dtype (str): Autocast for the conditioner.\n cache_path (Optional[str]): Path for pre-computed embeddings caching.\n kwargs: Additional parameters for residual vector quantizer.\n \"\"\"\n def __init__(self, dim: int, output_dim: int, device: str, attribute: str,\n quantize: bool, n_q: int, bins: int, checkpoint: tp.Union[str, Path], model_arch: str,\n enable_fusion: bool, sample_rate: int, max_audio_length: int, audio_stride: int,\n normalize: bool, text_p: bool, batch_size: tp.Optional[int] = None,\n autocast_dtype: tp.Optional[str] = 'float32', cache_path: tp.Optional[str] = None, **kwargs):\n try:\n import laion_clap # type: ignore\n except ImportError:\n raise ImportError(\"Please install CLAP to use the CLAPEmbeddingConditioner: 'pip install laion_clap'\")\n warnings.warn(\"Sample rate for CLAP conditioner was fixed in version v1.1.0, (from 44.1 to 48 kHz). \"\n \"Please retrain all models.\")\n checkpoint = AudioCraftEnvironment.resolve_reference_path(checkpoint)\n clap_tokenize = RobertaTokenizer.from_pretrained('roberta-base')\n clap_model = laion_clap.CLAP_Module(enable_fusion=enable_fusion, amodel=model_arch)\n load_clap_state_dict(clap_model, checkpoint)\n clap_model.eval()\n clap_model.to(device)\n super().__init__(dim=dim, output_dim=output_dim, device=device, attribute=attribute,\n autocast_dtype=autocast_dtype, quantize=quantize, n_q=n_q, bins=bins,\n **kwargs)\n self.checkpoint = checkpoint\n self.enable_fusion = enable_fusion\n self.model_arch = model_arch\n self.clap: laion_clap.CLAP_Module\n self.clap_tokenize: RobertaTokenizer\n self.clap_sample_rate = sample_rate\n self.clap_max_frames = int(self.clap_sample_rate * max_audio_length)\n self.clap_stride = int(self.clap_sample_rate * audio_stride)\n self.batch_size = batch_size or 1\n self.normalize = normalize\n self.text_p = text_p\n self.__dict__['clap_tokenize'] = clap_tokenize\n self.__dict__['clap'] = clap_model\n self.wav_cache, self.text_cache = None, None\n if cache_path is not None:\n self.wav_cache = EmbeddingCache(Path(cache_path) / 'wav', self.device,\n compute_embed_fn=self._get_wav_embedding_for_cache,\n extract_embed_fn=self._extract_wav_embedding_chunk)\n self.text_cache = EmbeddingCache(Path(cache_path) / 'text', self.device,\n compute_embed_fn=self._get_text_embedding_for_cache)\n\n def _tokenizer(self, texts: tp.Union[str, tp.List[str]]) -> dict:\n # we use the default params from CLAP module here as well\n return self.clap_tokenize(texts, padding=\"max_length\", truncation=True, max_length=77, return_tensors=\"pt\")\n\n def _compute_text_embedding(self, text: tp.List[str]) -> torch.Tensor:\n \"\"\"Compute text embedding from CLAP model on a given a batch of text.\n\n Args:\n text (list[str]): List of text for the batch, with B items.\n Returns:\n torch.Tensor: CLAP embedding derived from text, of shape [B, 1, D], with D the CLAP embedding dimension.\n \"\"\"\n with torch.no_grad():\n embed = self.clap.get_text_embedding(text, tokenizer=self._tokenizer, use_tensor=True)\n return embed.view(embed.size(0), 1, embed.size(-1))\n\n def _get_text_embedding_for_cache(self, path: tp.Union[Path, str],\n x: JointEmbedCondition, idx: int) -> torch.Tensor:\n \"\"\"Get text embedding function for the cache.\"\"\"\n text = x.text[idx]\n text = text if text is not None else \"\"\n return self._compute_text_embedding([text])[0]\n\n def _preprocess_wav(self, wav: torch.Tensor, length: torch.Tensor, sample_rates: tp.List[int]) -> torch.Tensor:\n \"\"\"Preprocess wav to expected format by CLAP model.\n\n Args:\n wav (torch.Tensor): Audio wav, of shape [B, C, T].\n length (torch.Tensor): Actual length of the audio for each item in the batch, of shape [B].\n sample_rates (list[int]): Sample rates for each sample in the batch\n Returns:\n torch.Tensor: Audio wav of shape [B, T].\n \"\"\"\n assert wav.dim() == 3, \"Expecting wav to be [B, C, T]\"\n if sample_rates is not None:\n _wav = []\n for i, audio in enumerate(wav):\n sr = sample_rates[i]\n audio = convert_audio(audio, from_rate=sr, to_rate=self.clap_sample_rate, to_channels=1)\n _wav.append(audio)\n wav = torch.stack(_wav, dim=0)\n wav = wav.mean(dim=1)\n return wav\n\n def _compute_wav_embedding(self, wav: torch.Tensor, length: torch.Tensor,\n sample_rates: tp.List[int], reduce_mean: bool = False) -> torch.Tensor:\n \"\"\"Compute audio wave embedding from CLAP model.\n\n Since CLAP operates on a fixed sequence length audio inputs and we need to process longer audio sequences,\n we calculate the wav embeddings on `clap_max_frames` windows with `clap_stride`-second stride and\n average the resulting embeddings.\n\n Args:\n wav (torch.Tensor): Audio wav, of shape [B, C, T].\n length (torch.Tensor): Actual length of the audio for each item in the batch, of shape [B].\n sample_rates (list[int]): Sample rates for each sample in the batch.\n reduce_mean (bool): Whether to get the average tensor.\n Returns:\n torch.Tensor: Audio embedding of shape [B, F, D], F being the number of chunks, D the dimension.\n \"\"\"\n with torch.no_grad():\n wav = self._preprocess_wav(wav, length, sample_rates)\n B, T = wav.shape\n if T >= self.clap_max_frames:\n wav = wav.unfold(-1, self.clap_max_frames, self.clap_stride) # [B, F, T]\n else:\n wav = wav.view(-1, 1, T) # [B, F, T] with F=1\n wav = einops.rearrange(wav, 'b f t -> (b f) t')\n embed_list = []\n for i in range(0, wav.size(0), self.batch_size):\n _wav = wav[i:i+self.batch_size, ...]\n _embed = self.clap.get_audio_embedding_from_data(_wav, use_tensor=True)\n embed_list.append(_embed)\n embed = torch.cat(embed_list, dim=0)\n embed = einops.rearrange(embed, '(b f) d -> b f d', b=B)\n if reduce_mean:\n embed = embed.mean(dim=1, keepdim=True)\n return embed # [B, F, D] with F=1 if reduce_mean is True\n\n def _get_wav_embedding_for_cache(self, path: tp.Union[str, Path],\n x: JointEmbedCondition, idx: int) -> torch.Tensor:\n \"\"\"Compute audio wave embedding for the cache.\n The embedding is computed on a given audio read from file.\n\n Args:\n path (str or Path): Path to the full audio file.\n Returns:\n torch.Tensor: Single-item tensor of shape [F, D], F being the number of chunks, D the dimension.\n \"\"\"\n wav, sr = audio_read(path) # [C, T]\n wav = wav.unsqueeze(0).to(self.device) # [1, C, T]\n wav_len = torch.LongTensor([wav.shape[-1]]).to(self.device)\n embed = self._compute_wav_embedding(wav, wav_len, [sr], reduce_mean=False) # [B, F, D]\n return embed.squeeze(0) # [F, D]\n\n def _extract_wav_embedding_chunk(self, full_embed: torch.Tensor, x: JointEmbedCondition, idx: int) -> torch.Tensor:\n \"\"\"Extract the chunk of embedding matching the seek_time and length from the full CLAP audio embedding.\n\n Args:\n full_embed (torch.Tensor): CLAP embedding computed on the full wave, of shape [F, D].\n x (JointEmbedCondition): Joint embedding condition for the full batch.\n idx (int): Index considered for the given embedding to extract.\n Returns:\n torch.Tensor: Wav embedding averaged on sliding window, of shape [1, D].\n \"\"\"\n sample_rate = x.sample_rate[idx]\n seek_time = x.seek_time[idx]\n seek_time = 0. if seek_time is None else seek_time\n clap_stride = int(self.clap_stride / self.clap_sample_rate) * sample_rate\n end_seek_time = seek_time + self.clap_max_frames / self.clap_sample_rate\n start_offset = int(seek_time * sample_rate // clap_stride)\n end_offset = int(end_seek_time * sample_rate // clap_stride)\n wav_embed = full_embed[start_offset:end_offset, ...]\n wav_embed = wav_embed.mean(dim=0, keepdim=True)\n return wav_embed.to(self.device) # [F, D]\n\n def _get_text_embedding(self, x: JointEmbedCondition) -> torch.Tensor:\n \"\"\"Get CLAP embedding from a batch of text descriptions.\"\"\"\n no_nullified_cond = x.wav.shape[-1] > 1 # we don't want to read from cache when condition dropout\n if self.text_cache is not None and no_nullified_cond:\n assert all(p is not None for p in x.path), \"Cache requires all JointEmbedCondition paths to be provided\"\n paths = [Path(p) for p in x.path if p is not None]\n embed = self.text_cache.get_embed_from_cache(paths, x)\n else:\n text = [xi if xi is not None else \"\" for xi in x.text]\n embed = self._compute_text_embedding(text)\n if self.normalize:\n embed = torch.nn.functional.normalize(embed, p=2.0, dim=-1)\n return embed\n\n def _get_wav_embedding(self, x: JointEmbedCondition) -> torch.Tensor:\n \"\"\"Get CLAP embedding from a batch of audio tensors (and corresponding sample rates).\"\"\"\n no_undefined_paths = all(p is not None for p in x.path)\n no_nullified_cond = x.wav.shape[-1] > 1 # we don't want to read from cache when condition dropout\n if self.wav_cache is not None and no_undefined_paths and no_nullified_cond:\n paths = [Path(p) for p in x.path if p is not None]\n embed = self.wav_cache.get_embed_from_cache(paths, x)\n else:\n embed = self._compute_wav_embedding(x.wav, x.length, x.sample_rate, reduce_mean=True)\n if self.normalize:\n embed = torch.nn.functional.normalize(embed, p=2.0, dim=-1)\n return embed\n\n def tokenize(self, x: JointEmbedCondition) -> JointEmbedCondition:\n # Trying to limit as much as possible sync points when the cache is warm.\n no_undefined_paths = all(p is not None for p in x.path)\n if self.wav_cache is not None and no_undefined_paths:\n assert all([p is not None for p in x.path]), \"Cache requires all JointEmbedCondition paths to be provided\"\n paths = [Path(p) for p in x.path if p is not None]\n self.wav_cache.populate_embed_cache(paths, x)\n if self.text_cache is not None and no_undefined_paths:\n assert all([p is not None for p in x.path]), \"Cache requires all JointEmbedCondition paths to be provided\"\n paths = [Path(p) for p in x.path if p is not None]\n self.text_cache.populate_embed_cache(paths, x)\n return x\n\n def _get_embed(self, x: JointEmbedCondition) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Extract shared latent representation from either the wav or the text using CLAP.\"\"\"\n # decide whether to use text embedding at train time or not\n use_text_embed = random.random() < self.text_p\n if self.training and not use_text_embed:\n embed = self._get_wav_embedding(x)\n empty_idx = torch.LongTensor([]) # we assume we always have the audio wav\n else:\n embed = self._get_text_embedding(x)\n empty_idx = torch.LongTensor([i for i, xi in enumerate(x.text) if xi is None or xi == \"\"])\n return embed, empty_idx\n\n\ndef dropout_symbolic_conditions(sample: ConditioningAttributes,\n condition: str, null_chord_idx: int = 194) -> ConditioningAttributes:\n \"\"\"\n Applies dropout to symbolic conditions within the sample based on the specified condition by setting the condition\n value to a null index.\n Args:\n sample (ConditioningAttributes): The sample containing symbolic attributes to potentially dropout.\n condition (str): The specific condition within the symbolic attributes to apply dropout.\n null_chord_idx (int, optional): The index used to represent a null chord. Defaults to 194.\n Returns:\n ConditioningAttributes: The modified sample with dropout applied to the specified condition.\n Raises:\n ValueError: If the specified condition is not present in the sample's symbolic attributes.\n \"\"\"\n if sample.symbolic == {} or sample.symbolic[JascoCondConst.CRD.value].frame_chords.shape[-1] <= 1: # type: ignore\n # nothing to drop\n return sample\n\n if condition not in getattr(sample, 'symbolic'):\n raise ValueError(\n \"dropout_symbolic_condition received an unexpected condition!\"\n f\" expected {sample.symbolic.keys()}\"\n f\" but got '{condition}'!\"\n )\n\n if condition == JascoCondConst.CRD.value:\n sample.symbolic[condition] = nullify_chords(sample.symbolic[condition], null_chord_idx=null_chord_idx)\n elif condition == JascoCondConst.MLD.value:\n sample.symbolic[condition] = nullify_melody(sample.symbolic[condition])\n\n return sample\n\n\ndef dropout_condition(sample: ConditioningAttributes,\n condition_type: str, condition: str,\n **kwargs) -> ConditioningAttributes:\n \"\"\"Utility function for nullifying an attribute inside an ConditioningAttributes object.\n If the condition is of type \"wav\", then nullify it using `nullify_condition` function.\n If the condition is of any other type, set its value to None.\n Works in-place.\n \"\"\"\n if condition_type not in ['text', 'wav', 'joint_embed', 'symbolic']:\n raise ValueError(\n \"dropout_condition got an unexpected condition type!\"\n f\" expected 'text', 'wav' or 'joint_embed' but got '{condition_type}'\"\n )\n\n if condition not in getattr(sample, condition_type):\n raise ValueError(\n \"dropout_condition received an unexpected condition!\"\n f\" expected wav={sample.wav.keys()} and text={sample.text.keys()}\"\n f\" but got '{condition}' of type '{condition_type}'!\"\n )\n\n if condition_type == 'wav':\n wav_cond = sample.wav[condition]\n sample.wav[condition] = nullify_wav(wav_cond)\n elif condition_type == 'joint_embed':\n embed = sample.joint_embed[condition]\n sample.joint_embed[condition] = nullify_joint_embed(embed)\n elif condition_type == 'symbolic':\n sample = dropout_symbolic_conditions(sample=sample, condition=condition, **kwargs)\n else:\n sample.text[condition] = None\n\n return sample\n\n\nclass DropoutModule(nn.Module):\n \"\"\"Base module for all dropout modules.\"\"\"\n def __init__(self, seed: int = 1234):\n super().__init__()\n self.rng = torch.Generator()\n self.rng.manual_seed(seed)\n\n\nclass AttributeDropout(DropoutModule):\n \"\"\"Dropout with a given probability per attribute.\n This is different from the behavior of ClassifierFreeGuidanceDropout as this allows for attributes\n to be dropped out separately. For example, \"artist\" can be dropped while \"genre\" remains.\n This is in contrast to ClassifierFreeGuidanceDropout where if \"artist\" is dropped \"genre\"\n must also be dropped.\n\n Args:\n p (tp.Dict[str, float]): A dict mapping between attributes and dropout probability. For example:\n ...\n \"genre\": 0.1,\n \"artist\": 0.5,\n \"wav\": 0.25,\n ...\n active_on_eval (bool, optional): Whether the dropout is active at eval. Default to False.\n seed (int, optional): Random seed.\n \"\"\"\n def __init__(self, p: tp.Dict[str, tp.Dict[str, float]], active_on_eval: bool = False, seed: int = 1234):\n super().__init__(seed=seed)\n self.active_on_eval = active_on_eval\n # construct dict that return the values from p otherwise 0\n self.p = {}\n for condition_type, probs in p.items():\n self.p[condition_type] = defaultdict(lambda: 0, probs)\n\n def forward(self, samples: tp.List[ConditioningAttributes]) -> tp.List[ConditioningAttributes]:\n \"\"\"\n Args:\n samples (list[ConditioningAttributes]): List of conditions.\n Returns:\n list[ConditioningAttributes]: List of conditions after certain attributes were set to None.\n \"\"\"\n if not self.training and not self.active_on_eval:\n return samples\n\n samples = deepcopy(samples)\n for condition_type, ps in self.p.items(): # for condition types [text, wav, symbolic]\n for condition, p in ps.items(): # for attributes of each type (e.g., [artist, genre])\n if torch.rand(1, generator=self.rng).item() < p:\n for sample in samples:\n dropout_condition(sample, condition_type, condition)\n return samples\n\n def __repr__(self):\n return f\"AttributeDropout({dict(self.p)})\"\n\n\nclass ClassifierFreeGuidanceDropout(DropoutModule):\n \"\"\"Classifier Free Guidance dropout.\n All attributes are dropped with the same probability.\n\n Args:\n p (float): Probability to apply condition dropout during training.\n seed (int): Random seed.\n \"\"\"\n def __init__(self, p: float, seed: int = 1234):\n super().__init__(seed=seed)\n self.p = p\n\n def forward(self, samples: tp.List[ConditioningAttributes],\n cond_types: tp.List[str] = [\"wav\", \"text\"],\n **kwargs) -> tp.List[ConditioningAttributes]:\n \"\"\"\n Args:\n samples (list[ConditioningAttributes]): List of conditions.\n Returns:\n list[ConditioningAttributes]: List of conditions after all attributes were set to None.\n \"\"\"\n if not self.training:\n return samples\n\n # decide on which attributes to drop in a batched fashion\n drop = torch.rand(1, generator=self.rng).item() < self.p\n if not drop:\n return samples\n\n # nullify conditions of all attributes\n samples = deepcopy(samples)\n for condition_type in cond_types:\n for sample in samples:\n for condition in sample.attributes[condition_type]:\n dropout_condition(sample, condition_type, condition,\n **kwargs)\n return samples\n\n def __repr__(self):\n return f\"ClassifierFreeGuidanceDropout(p={self.p})\"\n\n\nclass ConditioningProvider(nn.Module):\n \"\"\"Prepare and provide conditions given all the supported conditioners.\n\n Args:\n conditioners (dict): Dictionary of conditioners.\n device (torch.device or str, optional): Device for conditioners and output condition types.\n \"\"\"\n def __init__(self, conditioners: tp.Dict[str, BaseConditioner], device: tp.Union[torch.device, str] = \"cpu\"):\n super().__init__()\n self.device = device\n self.conditioners = nn.ModuleDict(conditioners)\n\n @property\n def joint_embed_conditions(self):\n return [m.attribute for m in self.conditioners.values() if isinstance(m, JointEmbeddingConditioner)]\n\n @property\n def has_joint_embed_conditions(self):\n return len(self.joint_embed_conditions) > 0\n\n @property\n def text_conditions(self):\n return [k for k, v in self.conditioners.items() if isinstance(v, TextConditioner)]\n\n @property\n def wav_conditions(self):\n return [k for k, v in self.conditioners.items() if isinstance(v, WaveformConditioner)]\n\n @property\n def has_wav_condition(self):\n return len(self.wav_conditions) > 0\n\n def tokenize(self, inputs: tp.List[ConditioningAttributes]) -> tp.Dict[str, tp.Any]:\n \"\"\"Match attributes/wavs with existing conditioners in self, and compute tokenize them accordingly.\n This should be called before starting any real GPU work to avoid synchronization points.\n This will return a dict matching conditioner names to their arbitrary tokenized representations.\n\n Args:\n inputs (list[ConditioningAttributes]): List of ConditioningAttributes objects containing\n text and wav conditions.\n \"\"\"\n assert all([isinstance(x, ConditioningAttributes) for x in inputs]), (\n \"Got unexpected types input for conditioner! should be tp.List[ConditioningAttributes]\",\n f\" but types were {set([type(x) for x in inputs])}\"\n )\n\n output = {}\n text = self._collate_text(inputs)\n wavs = self._collate_wavs(inputs)\n joint_embeds = self._collate_joint_embeds(inputs)\n\n assert set(text.keys() | wavs.keys() | joint_embeds.keys()).issubset(set(self.conditioners.keys())), (\n f\"Got an unexpected attribute! Expected {self.conditioners.keys()}, \",\n f\"got {text.keys(), wavs.keys(), joint_embeds.keys()}\"\n )\n\n for attribute, batch in chain(text.items(), wavs.items(), joint_embeds.items()):\n output[attribute] = self.conditioners[attribute].tokenize(batch)\n return output\n\n def forward(self, tokenized: tp.Dict[str, tp.Any]) -> tp.Dict[str, ConditionType]:\n \"\"\"Compute pairs of `(embedding, mask)` using the configured conditioners and the tokenized representations.\n The output is for example:\n {\n \"genre\": (torch.Tensor([B, 1, D_genre]), torch.Tensor([B, 1])),\n \"description\": (torch.Tensor([B, T_desc, D_desc]), torch.Tensor([B, T_desc])),\n ...\n }\n\n Args:\n tokenized (dict): Dict of tokenized representations as returned by `tokenize()`.\n \"\"\"\n output = {}\n for attribute, inputs in tokenized.items():\n condition, mask = self.conditioners[attribute](inputs)\n output[attribute] = (condition, mask)\n return output\n\n def _collate_text(self, samples: tp.List[ConditioningAttributes]) -> tp.Dict[str, tp.List[tp.Optional[str]]]:\n \"\"\"Given a list of ConditioningAttributes objects, compile a dictionary where the keys\n are the attributes and the values are the aggregated input per attribute.\n For example:\n Input:\n [\n ConditioningAttributes(text={\"genre\": \"Rock\", \"description\": \"A rock song with a guitar solo\"}, wav=...),\n ConditioningAttributes(text={\"genre\": \"Hip-hop\", \"description\": \"A hip-hop verse\"}, wav=...),\n ]\n Output:\n {\n \"genre\": [\"Rock\", \"Hip-hop\"],\n \"description\": [\"A rock song with a guitar solo\", \"A hip-hop verse\"]\n }\n\n Args:\n samples (list of ConditioningAttributes): List of ConditioningAttributes samples.\n Returns:\n dict[str, list[str, optional]]: A dictionary mapping an attribute name to text batch.\n \"\"\"\n out: tp.Dict[str, tp.List[tp.Optional[str]]] = defaultdict(list)\n texts = [x.text for x in samples]\n for text in texts:\n for condition in self.text_conditions:\n out[condition].append(text[condition])\n return out\n\n def _collate_wavs(self, samples: tp.List[ConditioningAttributes]) -> tp.Dict[str, WavCondition]:\n \"\"\"Generate a dict where the keys are attributes by which we fetch similar wavs,\n and the values are Tensors of wavs according to said attributes.\n\n *Note*: by the time the samples reach this function, each sample should have some waveform\n inside the \"wav\" attribute. It should be either:\n 1. A real waveform\n 2. A null waveform due to the sample having no similar waveforms (nullified by the dataset)\n 3. A null waveform due to it being dropped in a dropout module (nullified by dropout)\n\n Args:\n samples (list of ConditioningAttributes): List of ConditioningAttributes samples.\n Returns:\n dict[str, WavCondition]: A dictionary mapping an attribute name to wavs.\n \"\"\"\n wavs = defaultdict(list)\n lengths = defaultdict(list)\n sample_rates = defaultdict(list)\n paths = defaultdict(list)\n seek_times = defaultdict(list)\n out: tp.Dict[str, WavCondition] = {}\n\n for sample in samples:\n for attribute in self.wav_conditions:\n wav, length, sample_rate, path, seek_time = sample.wav[attribute]\n assert wav.dim() == 3, f\"Got wav with dim={wav.dim()}, but expected 3 [1, C, T]\"\n assert wav.size(0) == 1, f\"Got wav [B, C, T] with shape={wav.shape}, but expected B == 1\"\n # mono-channel conditioning\n wav = wav.mean(1, keepdim=True) # [1, 1, T]\n wavs[attribute].append(wav.flatten()) # [T]\n lengths[attribute].append(length)\n sample_rates[attribute].extend(sample_rate)\n paths[attribute].extend(path)\n seek_times[attribute].extend(seek_time)\n\n # stack all wavs to a single tensor\n for attribute in self.wav_conditions:\n stacked_wav, _ = collate(wavs[attribute], dim=0)\n out[attribute] = WavCondition(\n stacked_wav.unsqueeze(1), torch.cat(lengths[attribute]), sample_rates[attribute],\n paths[attribute], seek_times[attribute])\n\n return out\n\n def _collate_joint_embeds(self, samples: tp.List[ConditioningAttributes]) -> tp.Dict[str, JointEmbedCondition]:\n \"\"\"Generate a dict where the keys are attributes by which we compute joint embeddings,\n and the values are Tensors of pre-computed embeddings and the corresponding text attributes.\n\n Args:\n samples (list[ConditioningAttributes]): List of ConditioningAttributes samples.\n Returns:\n A dictionary mapping an attribute name to joint embeddings.\n \"\"\"\n texts = defaultdict(list)\n wavs = defaultdict(list)\n lengths = defaultdict(list)\n sample_rates = defaultdict(list)\n paths = defaultdict(list)\n seek_times = defaultdict(list)\n channels: int = 0\n\n out = {}\n for sample in samples:\n for attribute in self.joint_embed_conditions:\n wav, text, length, sample_rate, path, seek_time = sample.joint_embed[attribute]\n assert wav.dim() == 3\n if channels == 0:\n channels = wav.size(1)\n else:\n assert channels == wav.size(1), \"not all audio has same number of channels in batch\"\n assert wav.size(0) == 1, \"Expecting single-wav batch in the collate method\"\n wav = einops.rearrange(wav, \"b c t -> (b c t)\") # [1, C, T] => [C * T]\n wavs[attribute].append(wav)\n texts[attribute].extend(text)\n lengths[attribute].append(length)\n sample_rates[attribute].extend(sample_rate)\n paths[attribute].extend(path)\n seek_times[attribute].extend(seek_time)\n\n for attribute in self.joint_embed_conditions:\n stacked_texts = texts[attribute]\n stacked_paths = paths[attribute]\n stacked_seek_times = seek_times[attribute]\n stacked_wavs = pad_sequence(wavs[attribute]).to(self.device)\n stacked_wavs = einops.rearrange(stacked_wavs, \"(c t) b -> b c t\", c=channels)\n stacked_sample_rates = sample_rates[attribute]\n stacked_lengths = torch.cat(lengths[attribute]).to(self.device)\n assert stacked_lengths.size(0) == stacked_wavs.size(0)\n assert len(stacked_sample_rates) == stacked_wavs.size(0)\n assert len(stacked_texts) == stacked_wavs.size(0)\n out[attribute] = JointEmbedCondition(\n text=stacked_texts, wav=stacked_wavs,\n length=stacked_lengths, sample_rate=stacked_sample_rates,\n path=stacked_paths, seek_time=stacked_seek_times)\n\n return out\n\n\nclass ConditionFuser(StreamingModule):\n \"\"\"Condition fuser handles the logic to combine the different conditions\n to the actual model input.\n\n Args:\n fuse2cond (tp.Dict[str, str]): A dictionary that says how to fuse\n each condition. For example:\n {\n \"prepend\": [\"description\"],\n \"sum\": [\"genre\", \"bpm\"],\n \"cross\": [\"description\"],\n }\n cross_attention_pos_emb (bool, optional): Use positional embeddings in cross attention.\n cross_attention_pos_emb_scale (int): Scale for positional embeddings in cross attention if used.\n \"\"\"\n FUSING_METHODS = [\"sum\", \"prepend\", \"cross\", \"ignore\", \"input_interpolate\"]\n\n def __init__(self, fuse2cond: tp.Dict[str, tp.List[str]], cross_attention_pos_emb: bool = False,\n cross_attention_pos_emb_scale: float = 1.0):\n super().__init__()\n assert all(\n [k in self.FUSING_METHODS for k in fuse2cond.keys()]\n ), f\"Got invalid fuse method, allowed methods: {self.FUSING_METHODS}\"\n self.cross_attention_pos_emb = cross_attention_pos_emb\n self.cross_attention_pos_emb_scale = cross_attention_pos_emb_scale\n self.fuse2cond: tp.Dict[str, tp.List[str]] = fuse2cond\n self.cond2fuse: tp.Dict[str, str] = {}\n for fuse_method, conditions in fuse2cond.items():\n for condition in conditions:\n self.cond2fuse[condition] = fuse_method\n\n def forward(\n self,\n input: torch.Tensor,\n conditions: tp.Dict[str, ConditionType]\n ) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:\n \"\"\"Fuse the conditions to the provided model input.\n\n Args:\n input (torch.Tensor): Transformer input.\n conditions (dict[str, ConditionType]): Dict of conditions.\n Returns:\n tuple[torch.Tensor, torch.Tensor]: The first tensor is the transformer input\n after the conditions have been fused. The second output tensor is the tensor\n used for cross-attention or None if no cross attention inputs exist.\n \"\"\"\n B, T, _ = input.shape\n\n if 'offsets' in self._streaming_state:\n first_step = False\n offsets = self._streaming_state['offsets']\n else:\n first_step = True\n offsets = torch.zeros(input.shape[0], dtype=torch.long, device=input.device)\n\n assert set(conditions.keys()).issubset(set(self.cond2fuse.keys())), \\\n f\"given conditions contain unknown attributes for fuser, \" \\\n f\"expected {self.cond2fuse.keys()}, got {conditions.keys()}\"\n cross_attention_output = None\n for cond_type, (cond, cond_mask) in conditions.items():\n op = self.cond2fuse[cond_type]\n if op == 'sum':\n input += cond\n elif op == 'input_interpolate':\n cond = einops.rearrange(cond, \"b t d -> b d t\")\n cond = F.interpolate(cond, size=input.shape[1])\n input += einops.rearrange(cond, \"b d t -> b t d\")\n elif op == 'prepend':\n if first_step:\n input = torch.cat([cond, input], dim=1)\n elif op == 'cross':\n if cross_attention_output is not None:\n cross_attention_output = torch.cat([cross_attention_output, cond], dim=1)\n else:\n cross_attention_output = cond\n elif op == 'ignore':\n continue\n else:\n raise ValueError(f\"unknown op ({op})\")\n\n if self.cross_attention_pos_emb and cross_attention_output is not None:\n positions = torch.arange(\n cross_attention_output.shape[1],\n device=cross_attention_output.device\n ).view(1, -1, 1)\n pos_emb = create_sin_embedding(positions, cross_attention_output.shape[-1])\n cross_attention_output = cross_attention_output + self.cross_attention_pos_emb_scale * pos_emb\n\n if self._is_streaming:\n self._streaming_state['offsets'] = offsets + T\n\n return input, cross_attention_output\n" }, { "path": "audiocraft/modules/conv.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nimport typing as tp\nimport warnings\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.nn.utils import spectral_norm, weight_norm\n\n\nCONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',\n 'time_group_norm'])\n\n\ndef apply_parametrization_norm(module: nn.Module, norm: str = 'none'):\n assert norm in CONV_NORMALIZATIONS\n if norm == 'weight_norm':\n return weight_norm(module)\n elif norm == 'spectral_norm':\n return spectral_norm(module)\n else:\n # We already check was in CONV_NORMALIZATION, so any other choice\n # doesn't need reparametrization.\n return module\n\n\ndef get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs):\n \"\"\"Return the proper normalization module. If causal is True, this will ensure the returned\n module is causal, or return an error if the normalization doesn't support causal evaluation.\n \"\"\"\n assert norm in CONV_NORMALIZATIONS\n if norm == 'time_group_norm':\n if causal:\n raise ValueError(\"GroupNorm doesn't support causal evaluation.\")\n assert isinstance(module, nn.modules.conv._ConvNd)\n return nn.GroupNorm(1, module.out_channels, **norm_kwargs)\n else:\n return nn.Identity()\n\n\ndef get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,\n padding_total: int = 0) -> int:\n \"\"\"See `pad_for_conv1d`.\"\"\"\n length = x.shape[-1]\n n_frames = (length - kernel_size + padding_total) / stride + 1\n ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)\n return ideal_length - length\n\n\ndef pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):\n \"\"\"Pad for a convolution to make sure that the last window is full.\n Extra padding is added at the end. This is required to ensure that we can rebuild\n an output of the same length, as otherwise, even with padding, some time steps\n might get removed.\n For instance, with total padding = 4, kernel size = 4, stride = 2:\n 0 0 1 2 3 4 5 0 0 # (0s are padding)\n 1 2 3 # (output frames of a convolution, last 0 is never used)\n 0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding)\n 1 2 3 4 # once you removed padding, we are missing one time step !\n \"\"\"\n extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)\n return F.pad(x, (0, extra_padding))\n\n\ndef pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'constant', value: float = 0.):\n \"\"\"Tiny wrapper around F.pad, just to allow for reflect padding on small input.\n If this is the case, we insert extra 0 padding to the right before the reflection happen.\n \"\"\"\n length = x.shape[-1]\n padding_left, padding_right = paddings\n assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)\n if mode == 'reflect':\n max_pad = max(padding_left, padding_right)\n extra_pad = 0\n if length <= max_pad:\n extra_pad = max_pad - length + 1\n x = F.pad(x, (0, extra_pad))\n padded = F.pad(x, paddings, mode, value)\n end = padded.shape[-1] - extra_pad\n return padded[..., :end]\n else:\n return F.pad(x, paddings, mode, value)\n\n\ndef unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):\n \"\"\"Remove padding from x, handling properly zero padding. Only for 1d!\"\"\"\n padding_left, padding_right = paddings\n assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)\n assert (padding_left + padding_right) <= x.shape[-1]\n end = x.shape[-1] - padding_right\n return x[..., padding_left: end]\n\n\nclass NormConv1d(nn.Module):\n \"\"\"Wrapper around Conv1d and normalization applied to this conv\n to provide a uniform interface across normalization approaches.\n \"\"\"\n def __init__(self, *args, causal: bool = False, norm: str = 'none',\n norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):\n super().__init__()\n self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)\n self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)\n self.norm_type = norm\n\n def forward(self, x):\n x = self.conv(x)\n x = self.norm(x)\n return x\n\n\nclass NormConv2d(nn.Module):\n \"\"\"Wrapper around Conv2d and normalization applied to this conv\n to provide a uniform interface across normalization approaches.\n \"\"\"\n def __init__(self, *args, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):\n super().__init__()\n self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)\n self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)\n self.norm_type = norm\n\n def forward(self, x):\n x = self.conv(x)\n x = self.norm(x)\n return x\n\n\nclass NormConvTranspose1d(nn.Module):\n \"\"\"Wrapper around ConvTranspose1d and normalization applied to this conv\n to provide a uniform interface across normalization approaches.\n \"\"\"\n def __init__(self, *args, causal: bool = False, norm: str = 'none',\n norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):\n super().__init__()\n self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)\n self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)\n self.norm_type = norm\n\n def forward(self, x):\n x = self.convtr(x)\n x = self.norm(x)\n return x\n\n\nclass NormConvTranspose2d(nn.Module):\n \"\"\"Wrapper around ConvTranspose2d and normalization applied to this conv\n to provide a uniform interface across normalization approaches.\n \"\"\"\n def __init__(self, *args, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):\n super().__init__()\n self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm)\n self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)\n\n def forward(self, x):\n x = self.convtr(x)\n x = self.norm(x)\n return x\n\n\nclass StreamableConv1d(nn.Module):\n \"\"\"Conv1d with some builtin handling of asymmetric or causal padding\n and normalization.\n \"\"\"\n def __init__(self, in_channels: int, out_channels: int,\n kernel_size: int, stride: int = 1, dilation: int = 1,\n groups: int = 1, bias: bool = True, causal: bool = False,\n norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},\n pad_mode: str = 'reflect'):\n super().__init__()\n # warn user on unusual setup between dilation and stride\n if stride > 1 and dilation > 1:\n warnings.warn(\"StreamableConv1d has been initialized with stride > 1 and dilation > 1\"\n f\" (kernel_size={kernel_size} stride={stride}, dilation={dilation}).\")\n self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,\n dilation=dilation, groups=groups, bias=bias, causal=causal,\n norm=norm, norm_kwargs=norm_kwargs)\n self.causal = causal\n self.pad_mode = pad_mode\n\n def forward(self, x):\n B, C, T = x.shape\n kernel_size = self.conv.conv.kernel_size[0]\n stride = self.conv.conv.stride[0]\n dilation = self.conv.conv.dilation[0]\n kernel_size = (kernel_size - 1) * dilation + 1 # effective kernel size with dilations\n padding_total = kernel_size - stride\n extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)\n if self.causal:\n # Left padding for causal\n x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)\n else:\n # Asymmetric padding required for odd strides\n padding_right = padding_total // 2\n padding_left = padding_total - padding_right\n x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)\n return self.conv(x)\n\n\nclass StreamableConvTranspose1d(nn.Module):\n \"\"\"ConvTranspose1d with some builtin handling of asymmetric or causal padding\n and normalization.\n \"\"\"\n def __init__(self, in_channels: int, out_channels: int,\n kernel_size: int, stride: int = 1, causal: bool = False,\n norm: str = 'none', trim_right_ratio: float = 1.,\n norm_kwargs: tp.Dict[str, tp.Any] = {}):\n super().__init__()\n self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,\n causal=causal, norm=norm, norm_kwargs=norm_kwargs)\n self.causal = causal\n self.trim_right_ratio = trim_right_ratio\n assert self.causal or self.trim_right_ratio == 1., \\\n \"`trim_right_ratio` != 1.0 only makes sense for causal convolutions\"\n assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.\n\n def forward(self, x):\n kernel_size = self.convtr.convtr.kernel_size[0]\n stride = self.convtr.convtr.stride[0]\n padding_total = kernel_size - stride\n\n y = self.convtr(x)\n\n # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be\n # removed at the very end, when keeping only the right length for the output,\n # as removing it here would require also passing the length at the matching layer\n # in the encoder.\n if self.causal:\n # Trim the padding on the right according to the specified ratio\n # if trim_right_ratio = 1.0, trim everything from right\n padding_right = math.ceil(padding_total * self.trim_right_ratio)\n padding_left = padding_total - padding_right\n y = unpad1d(y, (padding_left, padding_right))\n else:\n # Asymmetric padding required for odd strides\n padding_right = padding_total // 2\n padding_left = padding_total - padding_right\n y = unpad1d(y, (padding_left, padding_right))\n return y\n" }, { "path": "audiocraft/modules/diffusion_schedule.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nFunctions for Noise Schedule, defines diffusion process, reverse process and data processor.\n\"\"\"\n\nfrom collections import namedtuple\nimport random\nimport typing as tp\nimport julius\nimport torch\n\nTrainingItem = namedtuple(\"TrainingItem\", \"noisy noise step\")\n\n\ndef betas_from_alpha_bar(alpha_bar):\n alphas = torch.cat([torch.Tensor([alpha_bar[0]]), alpha_bar[1:]/alpha_bar[:-1]])\n return 1 - alphas\n\n\nclass SampleProcessor(torch.nn.Module):\n def project_sample(self, x: torch.Tensor):\n \"\"\"Project the original sample to the 'space' where the diffusion will happen.\"\"\"\n return x\n\n def return_sample(self, z: torch.Tensor):\n \"\"\"Project back from diffusion space to the actual sample space.\"\"\"\n return z\n\n\nclass MultiBandProcessor(SampleProcessor):\n \"\"\"\n MultiBand sample processor. The input audio is splitted across\n frequency bands evenly distributed in mel-scale.\n\n Each band will be rescaled to match the power distribution\n of Gaussian noise in that band, using online metrics\n computed on the first few samples.\n\n Args:\n n_bands (int): Number of mel-bands to split the signal over.\n sample_rate (int): Sample rate of the audio.\n num_samples (int): Number of samples to use to fit the rescaling\n for each band. The processor won't be stable\n until it has seen that many samples.\n power_std (float or list/tensor): The rescaling factor computed to match the\n power of Gaussian noise in each band is taken to\n that power, i.e. `1.` means full correction of the energy\n in each band, and values less than `1` means only partial\n correction. Can be used to balance the relative importance\n of low vs. high freq in typical audio signals.\n \"\"\"\n def __init__(self, n_bands: int = 8, sample_rate: float = 24_000,\n num_samples: int = 10_000, power_std: tp.Union[float, tp.List[float], torch.Tensor] = 1.):\n super().__init__()\n self.n_bands = n_bands\n self.split_bands = julius.SplitBands(sample_rate, n_bands=n_bands)\n self.num_samples = num_samples\n self.power_std = power_std\n if isinstance(power_std, list):\n assert len(power_std) == n_bands\n power_std = torch.tensor(power_std)\n self.register_buffer('counts', torch.zeros(1))\n self.register_buffer('sum_x', torch.zeros(n_bands))\n self.register_buffer('sum_x2', torch.zeros(n_bands))\n self.register_buffer('sum_target_x2', torch.zeros(n_bands))\n self.counts: torch.Tensor\n self.sum_x: torch.Tensor\n self.sum_x2: torch.Tensor\n self.sum_target_x2: torch.Tensor\n\n @property\n def mean(self):\n mean = self.sum_x / self.counts\n return mean\n\n @property\n def std(self):\n std = (self.sum_x2 / self.counts - self.mean**2).clamp(min=0).sqrt()\n return std\n\n @property\n def target_std(self):\n target_std = self.sum_target_x2 / self.counts\n return target_std\n\n def project_sample(self, x: torch.Tensor):\n assert x.dim() == 3\n bands = self.split_bands(x)\n if self.counts.item() < self.num_samples:\n ref_bands = self.split_bands(torch.randn_like(x))\n self.counts += len(x)\n self.sum_x += bands.mean(dim=(2, 3)).sum(dim=1)\n self.sum_x2 += bands.pow(2).mean(dim=(2, 3)).sum(dim=1)\n self.sum_target_x2 += ref_bands.pow(2).mean(dim=(2, 3)).sum(dim=1)\n rescale = (self.target_std / self.std.clamp(min=1e-12)) ** self.power_std # same output size\n bands = (bands - self.mean.view(-1, 1, 1, 1)) * rescale.view(-1, 1, 1, 1)\n return bands.sum(dim=0)\n\n def return_sample(self, x: torch.Tensor):\n assert x.dim() == 3\n bands = self.split_bands(x)\n rescale = (self.std / self.target_std) ** self.power_std\n bands = bands * rescale.view(-1, 1, 1, 1) + self.mean.view(-1, 1, 1, 1)\n return bands.sum(dim=0)\n\n\nclass NoiseSchedule:\n \"\"\"Noise schedule for diffusion.\n\n Args:\n beta_t0 (float): Variance of the first diffusion step.\n beta_t1 (float): Variance of the last diffusion step.\n beta_exp (float): Power schedule exponent\n num_steps (int): Number of diffusion step.\n variance (str): choice of the sigma value for the denoising eq. Choices: \"beta\" or \"beta_tilde\"\n clip (float): clipping value for the denoising steps\n rescale (float): rescaling value to avoid vanishing signals unused by default (i.e 1)\n repartition (str): shape of the schedule only power schedule is supported\n sample_processor (SampleProcessor): Module that normalize data to match better the gaussian distribution\n noise_scale (float): Scaling factor for the noise\n \"\"\"\n def __init__(self, beta_t0: float = 1e-4, beta_t1: float = 0.02, num_steps: int = 1000, variance: str = 'beta',\n clip: float = 5., rescale: float = 1., device='cuda', beta_exp: float = 1,\n repartition: str = \"power\", alpha_sigmoid: dict = {}, n_bands: tp.Optional[int] = None,\n sample_processor: SampleProcessor = SampleProcessor(), noise_scale: float = 1.0, **kwargs):\n\n self.beta_t0 = beta_t0\n self.beta_t1 = beta_t1\n self.variance = variance\n self.num_steps = num_steps\n self.clip = clip\n self.sample_processor = sample_processor\n self.rescale = rescale\n self.n_bands = n_bands\n self.noise_scale = noise_scale\n assert n_bands is None\n if repartition == \"power\":\n self.betas = torch.linspace(beta_t0 ** (1 / beta_exp), beta_t1 ** (1 / beta_exp), num_steps,\n device=device, dtype=torch.float) ** beta_exp\n else:\n raise RuntimeError('Not implemented')\n self.rng = random.Random(1234)\n\n def get_beta(self, step: tp.Union[int, torch.Tensor]):\n if self.n_bands is None:\n return self.betas[step]\n else:\n return self.betas[:, step] # [n_bands, len(step)]\n\n def get_initial_noise(self, x: torch.Tensor):\n if self.n_bands is None:\n return torch.randn_like(x)\n return torch.randn((x.size(0), self.n_bands, x.size(2)))\n\n def get_alpha_bar(self, step: tp.Optional[tp.Union[int, torch.Tensor]] = None) -> torch.Tensor:\n \"\"\"Return 'alpha_bar', either for a given step, or as a tensor with its value for each step.\"\"\"\n if step is None:\n return (1 - self.betas).cumprod(dim=-1) # works for simgle and multi bands\n if type(step) is int:\n return (1 - self.betas[:step + 1]).prod()\n else:\n return (1 - self.betas).cumprod(dim=0)[step].view(-1, 1, 1)\n\n def get_training_item(self, x: torch.Tensor, tensor_step: bool = False) -> TrainingItem:\n \"\"\"Create a noisy data item for diffusion model training:\n\n Args:\n x (torch.Tensor): clean audio data torch.tensor(bs, 1, T)\n tensor_step (bool): If tensor_step = false, only one step t is sample,\n the whole batch is diffused to the same step and t is int.\n If tensor_step = true, t is a tensor of size (x.size(0),)\n every element of the batch is diffused to a independently sampled.\n \"\"\"\n step: tp.Union[int, torch.Tensor]\n if tensor_step:\n bs = x.size(0)\n step = torch.randint(0, self.num_steps, size=(bs,), device=x.device)\n else:\n step = self.rng.randrange(self.num_steps)\n alpha_bar = self.get_alpha_bar(step) # [batch_size, n_bands, 1]\n\n x = self.sample_processor.project_sample(x)\n noise = torch.randn_like(x)\n noisy = (alpha_bar.sqrt() / self.rescale) * x + (1 - alpha_bar).sqrt() * noise * self.noise_scale\n return TrainingItem(noisy, noise, step)\n\n def generate(self, model: torch.nn.Module, initial: tp.Optional[torch.Tensor] = None,\n condition: tp.Optional[torch.Tensor] = None, return_list: bool = False):\n \"\"\"Full ddpm reverse process.\n\n Args:\n model (nn.Module): Diffusion model.\n initial (tensor): Initial Noise.\n condition (tensor): Input conditionning Tensor (e.g. encodec compressed representation).\n return_list (bool): Whether to return the whole process or only the sampled point.\n \"\"\"\n alpha_bar = self.get_alpha_bar(step=self.num_steps - 1)\n current = initial\n iterates = [initial]\n for step in range(self.num_steps)[::-1]:\n with torch.no_grad():\n estimate = model(current, step, condition=condition).sample\n alpha = 1 - self.betas[step]\n previous = (current - (1 - alpha) / (1 - alpha_bar).sqrt() * estimate) / alpha.sqrt()\n previous_alpha_bar = self.get_alpha_bar(step=step - 1)\n if step == 0:\n sigma2 = 0\n elif self.variance == 'beta':\n sigma2 = 1 - alpha\n elif self.variance == 'beta_tilde':\n sigma2 = (1 - previous_alpha_bar) / (1 - alpha_bar) * (1 - alpha)\n elif self.variance == 'none':\n sigma2 = 0\n else:\n raise ValueError(f'Invalid variance type {self.variance}')\n\n if sigma2 > 0:\n previous += sigma2**0.5 * torch.randn_like(previous) * self.noise_scale\n if self.clip:\n previous = previous.clamp(-self.clip, self.clip)\n current = previous\n alpha_bar = previous_alpha_bar\n if step == 0:\n previous *= self.rescale\n if return_list:\n iterates.append(previous.cpu())\n\n if return_list:\n return iterates\n else:\n return self.sample_processor.return_sample(previous)\n\n def generate_subsampled(self, model: torch.nn.Module, initial: torch.Tensor, step_list: tp.Optional[list] = None,\n condition: tp.Optional[torch.Tensor] = None, return_list: bool = False):\n \"\"\"Reverse process that only goes through Markov chain states in step_list.\"\"\"\n if step_list is None:\n step_list = list(range(1000))[::-50] + [0]\n alpha_bar = self.get_alpha_bar(step=self.num_steps - 1)\n alpha_bars_subsampled = (1 - self.betas).cumprod(dim=0)[list(reversed(step_list))].cpu()\n betas_subsampled = betas_from_alpha_bar(alpha_bars_subsampled)\n current = initial * self.noise_scale\n iterates = [current]\n for idx, step in enumerate(step_list[:-1]):\n with torch.no_grad():\n estimate = model(current, step, condition=condition).sample * self.noise_scale\n alpha = 1 - betas_subsampled[-1 - idx]\n previous = (current - (1 - alpha) / (1 - alpha_bar).sqrt() * estimate) / alpha.sqrt()\n previous_alpha_bar = self.get_alpha_bar(step_list[idx + 1])\n if step == step_list[-2]:\n sigma2 = 0\n previous_alpha_bar = torch.tensor(1.0)\n else:\n sigma2 = (1 - previous_alpha_bar) / (1 - alpha_bar) * (1 - alpha)\n if sigma2 > 0:\n previous += sigma2**0.5 * torch.randn_like(previous) * self.noise_scale\n if self.clip:\n previous = previous.clamp(-self.clip, self.clip)\n current = previous\n alpha_bar = previous_alpha_bar\n if step == 0:\n previous *= self.rescale\n if return_list:\n iterates.append(previous.cpu())\n if return_list:\n return iterates\n else:\n return self.sample_processor.return_sample(previous)\n" }, { "path": "audiocraft/modules/jasco_conditioners.py", "content": "import torch\nimport typing as tp\nfrom itertools import chain\nfrom pathlib import Path\nfrom torch import nn\nfrom .conditioners import (ConditioningAttributes, BaseConditioner, ConditionType,\n ConditioningProvider, JascoCondConst,\n WaveformConditioner, WavCondition, SymbolicCondition)\nfrom ..data.audio import audio_read\nfrom ..data.audio_utils import convert_audio\nfrom ..utils.autocast import TorchAutocast\nfrom ..utils.cache import EmbeddingCache\n\n\nclass MelodyConditioner(BaseConditioner):\n \"\"\"\n A conditioner that handles melody conditioning from pre-computed salience matrix.\n Attributes:\n card (int): The cardinality of the melody matrix.\n out_dim (int): The dimensionality of the output projection.\n device (Union[torch.device, str]): The device on which the embeddings are stored.\n \"\"\"\n def __init__(self, card: int, out_dim: int, device: tp.Union[torch.device, str] = 'cpu', **kwargs):\n super().__init__(dim=card, output_dim=out_dim)\n self.device = device\n\n def tokenize(self, x: SymbolicCondition) -> SymbolicCondition:\n return SymbolicCondition(melody=x.melody.to(self.device)) # type: ignore\n\n def forward(self, x: SymbolicCondition) -> ConditionType:\n embeds = self.output_proj(x.melody.permute(0, 2, 1)) # type: ignore\n mask = torch.ones_like(embeds[..., 0])\n return embeds, mask\n\n\nclass ChordsEmbConditioner(BaseConditioner):\n \"\"\"\n A conditioner that embeds chord symbols into a continuous vector space.\n Attributes:\n card (int): The cardinality of the chord vocabulary.\n out_dim (int): The dimensionality of the output embeddings.\n device (Union[torch.device, str]): The device on which the embeddings are stored.\n \"\"\"\n def __init__(self, card: int, out_dim: int, device: tp.Union[torch.device, str] = 'cpu', **kwargs):\n vocab_size = card + 1 # card + 1 - for null chord used during dropout\n super().__init__(dim=vocab_size, output_dim=-1) # out_dim=-1 to avoid another projection\n self.emb = nn.Embedding(vocab_size, out_dim, device=device)\n self.device = device\n\n def tokenize(self, x: SymbolicCondition) -> SymbolicCondition:\n return SymbolicCondition(frame_chords=x.frame_chords.to(self.device)) # type: ignore\n\n def forward(self, x: SymbolicCondition) -> ConditionType:\n embeds = self.emb(x.frame_chords)\n mask = torch.ones_like(embeds[..., 0])\n return embeds, mask\n\n\nclass DrumsConditioner(WaveformConditioner):\n def __init__(self, out_dim: int, sample_rate: int, blurring_factor: int = 3,\n cache_path: tp.Optional[tp.Union[str, Path]] = None,\n compression_model_latent_dim: int = 128,\n compression_model_framerate: float = 50,\n segment_duration: float = 10.0,\n device: tp.Union[torch.device, str] = 'cpu',\n **kwargs):\n \"\"\"Drum condition conditioner\n\n Args:\n out_dim (int): _description_\n sample_rate (int): _description_\n blurring_factor (int, optional): _description_. Defaults to 3.\n cache_path (tp.Optional[tp.Union[str, Path]], optional): path to precomputed cache. Defaults to None.\n compression_model_latent_dim (int, optional): latent dimensino. Defaults to 128.\n compression_model_framerate (float, optional): frame rate of the representation model. Defaults to 50.\n segment_duration (float, optional): duration in sec for each audio segment. Defaults to 10.0.\n device (tp.Union[torch.device, str], optional): device. Defaults to 'cpu'.\n \"\"\"\n from demucs import pretrained\n self.sample_rate = sample_rate\n self.__dict__['demucs'] = pretrained.get_model('htdemucs').to(device)\n stem_sources: list = self.demucs.sources # type: ignore\n self.stem_idx = stem_sources.index('drums')\n self.compression_model = None\n self.latent_dim = compression_model_latent_dim\n super().__init__(dim=self.latent_dim, output_dim=out_dim, device=device)\n self.autocast = TorchAutocast(enabled=device != 'cpu', device_type=self.device, dtype=torch.float32)\n self._use_masking = False\n self.blurring_factor = blurring_factor\n self.seq_len = int(segment_duration * compression_model_framerate)\n self.cache = None # If you wish to train with EmbeddingCache, call self.create_embedding_cache(cache_path)\n\n def create_embedding_cache(self, cache_path):\n if cache_path is not None:\n self.cache = EmbeddingCache(Path(cache_path) / 'wav', self.device,\n compute_embed_fn=self._calc_coarse_drum_codes_for_cache,\n extract_embed_fn=self._load_drum_codes_chunk)\n\n @torch.no_grad()\n def _get_drums_stem(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:\n \"\"\"Get parts of the wav that holds the drums, extracting the main stems from the wav.\"\"\"\n from demucs.apply import apply_model\n from demucs.audio import convert_audio\n with self.autocast:\n wav = convert_audio(\n wav, sample_rate, self.demucs.samplerate, self.demucs.audio_channels) # type: ignore\n stems = apply_model(self.demucs, wav, device=self.device)\n drum_stem = stems[:, self.stem_idx] # extract relevant stems for drums conditioning\n return convert_audio(drum_stem, self.demucs.samplerate, self.sample_rate, 1) # type: ignore\n\n def _temporal_blur(self, z: torch.Tensor):\n # z: (B, T, C)\n B, T, C = z.shape\n if T % self.blurring_factor != 0:\n # pad with reflect for T % self.temporal_blurring on the right in dim=1\n pad_val = self.blurring_factor - T % self.blurring_factor\n z = torch.nn.functional.pad(z, (0, 0, 0, pad_val), mode='reflect')\n z = z.reshape(B, -1, self.blurring_factor, C).sum(dim=2) / self.blurring_factor\n z = z.unsqueeze(2).repeat(1, 1, self.blurring_factor, 1).reshape(B, -1, C)\n z = z[:, :T]\n assert z.shape == (B, T, C)\n return z\n\n @torch.no_grad()\n def _extract_coarse_drum_codes(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:\n assert self.compression_model is not None\n\n # stem separation of drums\n drums = self._get_drums_stem(wav, sample_rate)\n\n # continuous encoding with compression model\n latents = self.compression_model.model.encoder(drums)\n\n # quantization to coarsest codebook\n coarsest_quantizer = self.compression_model.model.quantizer.layers[0]\n drums = coarsest_quantizer.encode(latents).to(torch.int16)\n return drums\n\n @torch.no_grad()\n def _calc_coarse_drum_codes_for_cache(self, path: tp.Union[str, Path],\n x: WavCondition, idx: int,\n max_duration_to_process: float = 600) -> torch.Tensor:\n \"\"\"Extract blurred drum latents from the whole audio waveform at the given path.\"\"\"\n wav, sr = audio_read(path)\n wav = wav[None].to(self.device)\n wav = convert_audio(wav, sr, self.sample_rate, to_channels=1)\n\n max_frames_to_process = int(max_duration_to_process * self.sample_rate)\n if wav.shape[-1] > max_frames_to_process:\n # process very long tracks in chunks\n start = 0\n codes = []\n while start < wav.shape[-1] - 1:\n wav_chunk = wav[..., start: start + max_frames_to_process]\n codes.append(self._extract_coarse_drum_codes(wav_chunk, self.sample_rate)[0])\n start += max_frames_to_process\n return torch.cat(codes)\n\n return self._extract_coarse_drum_codes(wav, self.sample_rate)[0]\n\n def _load_drum_codes_chunk(self, full_coarse_drum_codes: torch.Tensor, x: WavCondition, idx: int) -> torch.Tensor:\n \"\"\"Extract a chunk of coarse drum codes from the full coarse drum codes derived from the full waveform.\"\"\"\n wav_length = x.wav.shape[-1]\n seek_time = x.seek_time[idx]\n assert seek_time is not None, (\n \"WavCondition seek_time is required \"\n \"when extracting chunks from pre-computed drum codes.\")\n assert self.compression_model is not None\n frame_rate = self.compression_model.frame_rate\n target_length = int(frame_rate * wav_length / self.sample_rate)\n target_length = max(target_length, self.seq_len)\n index = int(frame_rate * seek_time)\n out = full_coarse_drum_codes[index: index + target_length]\n # pad\n out = torch.cat((out, torch.zeros(target_length - out.shape[0], dtype=out.dtype, device=out.device)))\n return out.to(self.device)\n\n @torch.no_grad()\n def _get_wav_embedding(self, x: WavCondition) -> torch.Tensor:\n bs = x.wav.shape[0]\n if x.wav.shape[-1] <= 1:\n # null condition\n return torch.zeros((bs, self.seq_len, self.latent_dim), device=x.wav.device, dtype=x.wav.dtype)\n\n # extract coarse drum codes\n no_undefined_paths = all(p is not None for p in x.path)\n no_nullified_cond = x.wav.shape[-1] > 1\n if self.cache is not None and no_undefined_paths and no_nullified_cond:\n paths = [Path(p) for p in x.path if p is not None]\n codes = self.cache.get_embed_from_cache(paths, x)\n else:\n assert all(sr == x.sample_rate[0] for sr in x.sample_rate), \"All sample rates in batch should be equal.\"\n codes = self._extract_coarse_drum_codes(x.wav, x.sample_rate[0])\n\n assert self.compression_model is not None\n # decode back to the continuous representation of compression model\n codes = codes.unsqueeze(1).permute(1, 0, 2) # (B, T) -> (1, B, T)\n codes = codes.to(torch.int64)\n latents = self.compression_model.model.quantizer.decode(codes)\n\n latents = latents.permute(0, 2, 1) # [B, C, T] -> [B, T, C]\n\n # temporal blurring\n return self._temporal_blur(latents)\n\n def tokenize(self, x: WavCondition) -> WavCondition:\n \"\"\"Apply WavConditioner tokenization and populate cache if needed.\"\"\"\n x = super().tokenize(x)\n no_undefined_paths = all(p is not None for p in x.path)\n if self.cache is not None and no_undefined_paths:\n paths = [Path(p) for p in x.path if p is not None]\n self.cache.populate_embed_cache(paths, x)\n return x\n\n\nclass JascoConditioningProvider(ConditioningProvider):\n \"\"\"\n A cond-provider that manages and tokenizes various types of conditioning attributes for Jasco models.\n Attributes:\n chords_card (int): The cardinality of the chord vocabulary.\n sequence_length (int): The length of the sequence for padding purposes.\n melody_dim (int): The dimensionality of the melody matrix.\n \"\"\"\n def __init__(self, *args,\n chords_card: int = 194,\n sequence_length: int = 500,\n melody_dim: int = 53, **kwargs):\n self.null_chord = chords_card\n self.sequence_len = sequence_length\n self.melody_dim = melody_dim\n super().__init__(*args, **kwargs)\n\n def tokenize(self, inputs: tp.List[ConditioningAttributes]) -> tp.Dict[str, tp.Any]:\n \"\"\"Match attributes/wavs with existing conditioners in self, and compute tokenize them accordingly.\n This should be called before starting any real GPU work to avoid synchronization points.\n This will return a dict matching conditioner names to their arbitrary tokenized representations.\n\n Args:\n inputs (list[ConditioningAttributes]): List of ConditioningAttributes objects containing\n text and wav conditions.\n \"\"\"\n assert all([isinstance(x, ConditioningAttributes) for x in inputs]), (\n \"Got unexpected types input for conditioner! should be tp.List[ConditioningAttributes]\",\n f\" but types were {set([type(x) for x in inputs])}\"\n )\n\n output = {}\n text = self._collate_text(inputs)\n wavs = self._collate_wavs(inputs)\n\n symbolic = self._collate_symbolic(inputs, self.conditioners.keys())\n\n assert set(text.keys() | wavs.keys() | symbolic.keys()).issubset(set(self.conditioners.keys())), (\n f\"Got an unexpected attribute! Expected {self.conditioners.keys()}, \",\n f\"got {text.keys(), wavs.keys(), symbolic.keys()}\"\n )\n\n for attribute, batch in chain(text.items(), wavs.items(), symbolic.items()):\n output[attribute] = self.conditioners[attribute].tokenize(batch)\n return output\n\n def _collate_symbolic(self, samples: tp.List[ConditioningAttributes],\n conditioner_keys: tp.Set) -> tp.Dict[str, SymbolicCondition]:\n output = {}\n\n # collate if symbolic cond exists\n if any(x in conditioner_keys for x in JascoCondConst.SYM.value):\n\n for s in samples:\n # hydrate with null chord if chords not exist - for inference support\n if (s.symbolic == {} or\n s.symbolic[JascoCondConst.CRD.value].frame_chords is None or\n s.symbolic[JascoCondConst.CRD.value].frame_chords.shape[-1] <= 1): # type: ignore\n # no chords conditioning - fill with null chord token\n s.symbolic[JascoCondConst.CRD.value] = SymbolicCondition(\n frame_chords=torch.ones(self.sequence_len, dtype=torch.int32) * self.null_chord)\n\n if (s.symbolic == {} or\n s.symbolic[JascoCondConst.MLD.value].melody is None or\n s.symbolic[JascoCondConst.MLD.value].melody.shape[-1] <= 1): # type: ignore\n # no chords conditioning - fill with null chord token\n s.symbolic[JascoCondConst.MLD.value] = SymbolicCondition(\n melody=torch.zeros((self.melody_dim, self.sequence_len)))\n\n if JascoCondConst.CRD.value in conditioner_keys:\n # pad to max\n max_seq_len = max(\n [s.symbolic[JascoCondConst.CRD.value].frame_chords.shape[-1] for s in samples]) # type: ignore\n padded_chords = [\n torch.cat((x.symbolic[JascoCondConst.CRD.value].frame_chords, # type: ignore\n torch.ones(max_seq_len -\n x.symbolic[JascoCondConst.CRD.value].frame_chords.shape[-1], # type: ignore\n dtype=torch.int32) * self.null_chord))\n for x in samples\n ]\n output[JascoCondConst.CRD.value] = SymbolicCondition(frame_chords=torch.stack(padded_chords))\n if JascoCondConst.MLD.value in conditioner_keys:\n melodies = torch.stack([x.symbolic[JascoCondConst.MLD.value].melody for x in samples]) # type: ignore\n output[JascoCondConst.MLD.value] = SymbolicCondition(melody=melodies)\n return output\n" }, { "path": "audiocraft/modules/lstm.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom torch import nn\n\n\nclass StreamableLSTM(nn.Module):\n \"\"\"LSTM without worrying about the hidden state, nor the layout of the data.\n Expects input as convolutional layout.\n \"\"\"\n def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):\n super().__init__()\n self.skip = skip\n self.lstm = nn.LSTM(dimension, dimension, num_layers)\n\n def forward(self, x):\n x = x.permute(2, 0, 1)\n y, _ = self.lstm(x)\n if self.skip:\n y = y + x\n y = y.permute(1, 2, 0)\n return y\n" }, { "path": "audiocraft/modules/rope.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nfrom torch import nn\nimport torch\n\n\nclass XPos(nn.Module):\n \"\"\"Length-extrapolatable positional embedding (xPos) from [Sun et al 2022](https://arxiv.org/abs/2212.10554v1).\n This applies an exponential decay to the RoPE rotation matrix.\n\n Args:\n dim (int): Embedding dimension.\n smoothing (float): Smoothing factor applied to the decay rates.\n base_scale (int): Base decay rate, given in terms of scaling time.\n device (torch.device, optional): Device on which to initialize the module.\n dtype (torch.dtype): dtype to use to generate the embedding.\n \"\"\"\n def __init__(self, dim: int, smoothing: float = 0.4, base_scale: int = 512,\n device=None, dtype: torch.dtype = torch.float32):\n super().__init__()\n assert dim % 2 == 0\n assert dtype in [torch.float64, torch.float32]\n self.dtype = dtype\n self.base_scale = base_scale\n\n half_dim = dim // 2\n adim = torch.arange(half_dim, device=device, dtype=dtype)\n decay_rates = (adim / half_dim + smoothing) / (1.0 + smoothing)\n self.register_buffer(\"decay_rates\", decay_rates)\n self.decay: tp.Optional[torch.Tensor] = None\n\n def get_decay(self, start: int, end: int):\n \"\"\"Create complex decay tensor, cache values for fast computation.\"\"\"\n if self.decay is None or end > self.decay.shape[0]:\n assert isinstance(self.decay_rates, torch.Tensor) # Satisfy type checker.\n idx = torch.arange(end, device=self.decay_rates.device, dtype=self.dtype)\n power = idx / self.base_scale\n scale = self.decay_rates ** power.unsqueeze(-1)\n self.decay = torch.polar(scale, torch.zeros_like(scale))\n return self.decay[start:end] # [T, C/2]\n\n\nclass RotaryEmbedding(nn.Module):\n \"\"\"Rotary positional embedding (RoPE) from [Su et al 2022](https://arxiv.org/abs/2104.09864).\n\n Args:\n dim (int): Embedding dimension (twice the number of frequencies).\n max_period (float): Maximum period of the rotation frequencies.\n xpos (bool): Use xPos, applies an exponential decay to rotation matrix.\n scale (float): Scale of positional embedding, set to 0 to deactivate.\n device (torch.device, optional): Device on which to initialize the module.\n dtype (torch.dtype): dtype to use to generate the embedding.\n \"\"\"\n def __init__(self, dim: int, max_period: float = 10000.0, xpos: bool = False,\n scale: float = 1.0, device=None, dtype: torch.dtype = torch.float32):\n super().__init__()\n assert dim % 2 == 0\n self.scale = scale\n assert dtype in [torch.float64, torch.float32]\n self.dtype = dtype\n\n adim = torch.arange(0, dim, 2, device=device, dtype=dtype)[: (dim // 2)]\n frequencies = 1.0 / (max_period ** (adim / dim))\n self.register_buffer(\"frequencies\", frequencies)\n self.rotation: tp.Optional[torch.Tensor] = None\n\n self.xpos = XPos(dim, device=device, dtype=dtype) if xpos else None\n\n def get_rotation(self, start: int, end: int):\n \"\"\"Create complex rotation tensor, cache values for fast computation.\"\"\"\n if self.rotation is None or end > self.rotation.shape[0]:\n assert isinstance(self.frequencies, torch.Tensor) # Satisfy type checker.\n idx = torch.arange(end, device=self.frequencies.device, dtype=self.dtype)\n angles = torch.outer(idx, self.frequencies)\n self.rotation = torch.polar(torch.ones_like(angles), angles)\n return self.rotation[start:end]\n\n def rotate(self, x: torch.Tensor, start: int = 0, time_dim: int = 1, invert_decay: bool = False):\n \"\"\"Apply rope rotation to query or key tensor.\"\"\"\n T = x.shape[time_dim]\n target_shape = [1] * x.dim()\n target_shape[time_dim] = T\n target_shape[-1] = -1\n rotation = self.get_rotation(start, start + T).view(target_shape)\n\n if self.xpos:\n decay = self.xpos.get_decay(start, start + T).view(target_shape)\n else:\n decay = 1.0\n\n if invert_decay:\n decay = decay ** -1\n\n x_complex = torch.view_as_complex(x.to(self.dtype).reshape(*x.shape[:-1], -1, 2))\n scaled_rotation = (rotation * decay) * self.scale + (1.0 - self.scale)\n x_out = torch.view_as_real(x_complex * scaled_rotation).view_as(x)\n\n return x_out.type_as(x)\n\n def rotate_qk(self, query: torch.Tensor, key: torch.Tensor, start: int = 0, time_dim: int = 1):\n \"\"\" Apply rope rotation to both query and key tensors.\n Supports streaming mode, in which query and key are not expected to have the same shape.\n In streaming mode, key will be of length [P + C] with P the cached past timesteps, but\n query will be [C] (typically C == 1).\n\n Args:\n query (torch.Tensor): Query to rotate.\n key (torch.Tensor): Key to rotate.\n start (int): Start index of the sequence for time offset.\n time_dim (int): which dimension represent the time steps.\n \"\"\"\n query_timesteps = query.shape[time_dim]\n key_timesteps = key.shape[time_dim]\n streaming_offset = key_timesteps - query_timesteps\n\n query_out = self.rotate(query, start + streaming_offset, time_dim)\n key_out = self.rotate(key, start, time_dim, invert_decay=True)\n\n return query_out, key_out\n" }, { "path": "audiocraft/modules/seanet.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nimport numpy as np\nimport torch.nn as nn\n\nfrom .conv import StreamableConv1d, StreamableConvTranspose1d\nfrom .lstm import StreamableLSTM\n\n\nclass SEANetResnetBlock(nn.Module):\n \"\"\"Residual block from SEANet model.\n\n Args:\n dim (int): Dimension of the input/output.\n kernel_sizes (list): List of kernel sizes for the convolutions.\n dilations (list): List of dilations for the convolutions.\n activation (str): Activation function.\n activation_params (dict): Parameters to provide to the activation function.\n norm (str): Normalization method.\n norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.\n causal (bool): Whether to use fully causal convolution.\n pad_mode (str): Padding mode for the convolutions.\n compress (int): Reduced dimensionality in residual branches (from Demucs v3).\n true_skip (bool): Whether to use true skip connection or a simple\n (streamable) convolution as the skip connection.\n \"\"\"\n def __init__(self, dim: int, kernel_sizes: tp.List[int] = [3, 1], dilations: tp.List[int] = [1, 1],\n activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},\n norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, causal: bool = False,\n pad_mode: str = 'reflect', compress: int = 2, true_skip: bool = True):\n super().__init__()\n assert len(kernel_sizes) == len(dilations), 'Number of kernel sizes should match number of dilations'\n act = getattr(nn, activation)\n hidden = dim // compress\n block = []\n for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):\n in_chs = dim if i == 0 else hidden\n out_chs = dim if i == len(kernel_sizes) - 1 else hidden\n block += [\n act(**activation_params),\n StreamableConv1d(in_chs, out_chs, kernel_size=kernel_size, dilation=dilation,\n norm=norm, norm_kwargs=norm_params,\n causal=causal, pad_mode=pad_mode),\n ]\n self.block = nn.Sequential(*block)\n self.shortcut: nn.Module\n if true_skip:\n self.shortcut = nn.Identity()\n else:\n self.shortcut = StreamableConv1d(dim, dim, kernel_size=1, norm=norm, norm_kwargs=norm_params,\n causal=causal, pad_mode=pad_mode)\n\n def forward(self, x):\n return self.shortcut(x) + self.block(x)\n\n\nclass SEANetEncoder(nn.Module):\n \"\"\"SEANet encoder.\n\n Args:\n channels (int): Audio channels.\n dimension (int): Intermediate representation dimension.\n n_filters (int): Base width for the model.\n n_residual_layers (int): nb of residual layers.\n ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of\n upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here\n that must match the decoder order. We use the decoder order as some models may only employ the decoder.\n activation (str): Activation function.\n activation_params (dict): Parameters to provide to the activation function.\n norm (str): Normalization method.\n norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.\n kernel_size (int): Kernel size for the initial convolution.\n last_kernel_size (int): Kernel size for the initial convolution.\n residual_kernel_size (int): Kernel size for the residual layers.\n dilation_base (int): How much to increase the dilation with each layer.\n causal (bool): Whether to use fully causal convolution.\n pad_mode (str): Padding mode for the convolutions.\n true_skip (bool): Whether to use true skip connection or a simple\n (streamable) convolution as the skip connection in the residual network blocks.\n compress (int): Reduced dimensionality in residual branches (from Demucs v3).\n lstm (int): Number of LSTM layers at the end of the encoder.\n disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.\n For the encoder, it corresponds to the N first blocks.\n \"\"\"\n def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 3,\n ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},\n norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,\n last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,\n pad_mode: str = 'reflect', true_skip: bool = True, compress: int = 2, lstm: int = 0,\n disable_norm_outer_blocks: int = 0):\n super().__init__()\n self.channels = channels\n self.dimension = dimension\n self.n_filters = n_filters\n self.ratios = list(reversed(ratios))\n del ratios\n self.n_residual_layers = n_residual_layers\n self.hop_length = np.prod(self.ratios)\n self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks\n self.disable_norm_outer_blocks = disable_norm_outer_blocks\n assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \\\n \"Number of blocks for which to disable norm is invalid.\" \\\n \"It should be lower or equal to the actual number of blocks in the network and greater or equal to 0.\"\n\n act = getattr(nn, activation)\n mult = 1\n model: tp.List[nn.Module] = [\n StreamableConv1d(channels, mult * n_filters, kernel_size,\n norm='none' if self.disable_norm_outer_blocks >= 1 else norm,\n norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)\n ]\n # Downsample to raw audio scale\n for i, ratio in enumerate(self.ratios):\n block_norm = 'none' if self.disable_norm_outer_blocks >= i + 2 else norm\n # Add residual layers\n for j in range(n_residual_layers):\n model += [\n SEANetResnetBlock(mult * n_filters, kernel_sizes=[residual_kernel_size, 1],\n dilations=[dilation_base ** j, 1],\n norm=block_norm, norm_params=norm_params,\n activation=activation, activation_params=activation_params,\n causal=causal, pad_mode=pad_mode, compress=compress, true_skip=true_skip)]\n\n # Add downsampling layers\n model += [\n act(**activation_params),\n StreamableConv1d(mult * n_filters, mult * n_filters * 2,\n kernel_size=ratio * 2, stride=ratio,\n norm=block_norm, norm_kwargs=norm_params,\n causal=causal, pad_mode=pad_mode),\n ]\n mult *= 2\n\n if lstm:\n model += [StreamableLSTM(mult * n_filters, num_layers=lstm)]\n\n model += [\n act(**activation_params),\n StreamableConv1d(mult * n_filters, dimension, last_kernel_size,\n norm='none' if self.disable_norm_outer_blocks == self.n_blocks else norm,\n norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)\n ]\n\n self.model = nn.Sequential(*model)\n\n def forward(self, x):\n return self.model(x)\n\n\nclass SEANetDecoder(nn.Module):\n \"\"\"SEANet decoder.\n\n Args:\n channels (int): Audio channels.\n dimension (int): Intermediate representation dimension.\n n_filters (int): Base width for the model.\n n_residual_layers (int): nb of residual layers.\n ratios (Sequence[int]): kernel size and stride ratios.\n activation (str): Activation function.\n activation_params (dict): Parameters to provide to the activation function.\n final_activation (str): Final activation function after all convolutions.\n final_activation_params (dict): Parameters to provide to the activation function.\n norm (str): Normalization method.\n norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.\n kernel_size (int): Kernel size for the initial convolution.\n last_kernel_size (int): Kernel size for the initial convolution.\n residual_kernel_size (int): Kernel size for the residual layers.\n dilation_base (int): How much to increase the dilation with each layer.\n causal (bool): Whether to use fully causal convolution.\n pad_mode (str): Padding mode for the convolutions.\n true_skip (bool): Whether to use true skip connection or a simple.\n (streamable) convolution as the skip connection in the residual network blocks.\n compress (int): Reduced dimensionality in residual branches (from Demucs v3).\n lstm (int): Number of LSTM layers at the end of the encoder.\n disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.\n For the decoder, it corresponds to the N last blocks.\n trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup.\n If equal to 1.0, it means that all the trimming is done at the right.\n \"\"\"\n def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 3,\n ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},\n final_activation: tp.Optional[str] = None, final_activation_params: tp.Optional[dict] = None,\n norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,\n last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,\n pad_mode: str = 'reflect', true_skip: bool = True, compress: int = 2, lstm: int = 0,\n disable_norm_outer_blocks: int = 0, trim_right_ratio: float = 1.0):\n super().__init__()\n self.dimension = dimension\n self.channels = channels\n self.n_filters = n_filters\n self.ratios = ratios\n del ratios\n self.n_residual_layers = n_residual_layers\n self.hop_length = np.prod(self.ratios)\n self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks\n self.disable_norm_outer_blocks = disable_norm_outer_blocks\n assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \\\n \"Number of blocks for which to disable norm is invalid.\" \\\n \"It should be lower or equal to the actual number of blocks in the network and greater or equal to 0.\"\n\n act = getattr(nn, activation)\n mult = int(2 ** len(self.ratios))\n model: tp.List[nn.Module] = [\n StreamableConv1d(dimension, mult * n_filters, kernel_size,\n norm='none' if self.disable_norm_outer_blocks == self.n_blocks else norm,\n norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)\n ]\n\n if lstm:\n model += [StreamableLSTM(mult * n_filters, num_layers=lstm)]\n\n # Upsample to raw audio scale\n for i, ratio in enumerate(self.ratios):\n block_norm = 'none' if self.disable_norm_outer_blocks >= self.n_blocks - (i + 1) else norm\n # Add upsampling layers\n model += [\n act(**activation_params),\n StreamableConvTranspose1d(mult * n_filters, mult * n_filters // 2,\n kernel_size=ratio * 2, stride=ratio,\n norm=block_norm, norm_kwargs=norm_params,\n causal=causal, trim_right_ratio=trim_right_ratio),\n ]\n # Add residual layers\n for j in range(n_residual_layers):\n model += [\n SEANetResnetBlock(mult * n_filters // 2, kernel_sizes=[residual_kernel_size, 1],\n dilations=[dilation_base ** j, 1],\n activation=activation, activation_params=activation_params,\n norm=block_norm, norm_params=norm_params, causal=causal,\n pad_mode=pad_mode, compress=compress, true_skip=true_skip)]\n\n mult //= 2\n\n # Add final layers\n model += [\n act(**activation_params),\n StreamableConv1d(n_filters, channels, last_kernel_size,\n norm='none' if self.disable_norm_outer_blocks >= 1 else norm,\n norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)\n ]\n # Add optional final activation to decoder (eg. tanh)\n if final_activation is not None:\n final_act = getattr(nn, final_activation)\n final_activation_params = final_activation_params or {}\n model += [\n final_act(**final_activation_params)\n ]\n self.model = nn.Sequential(*model)\n\n def forward(self, z):\n y = self.model(z)\n return y\n" }, { "path": "audiocraft/modules/streaming.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nStreaming module API that should be implemented by all Streaming components,\n\"\"\"\n\nfrom contextlib import contextmanager\nimport typing as tp\nfrom torch import nn\nimport torch\n\n\nState = tp.Dict[str, torch.Tensor]\n\n\nclass StreamingModule(nn.Module):\n \"\"\"Common API for streaming components.\n\n Each streaming component has a streaming state, which is just a dict[str, Tensor].\n By convention, the first dim of each tensor must be the batch size.\n Don't use dots in the key names, as this would clash with submodules\n (like in state_dict).\n\n If `self._is_streaming` is True, the component should use and remember\n the proper state inside `self._streaming_state`.\n\n To set a streaming component in streaming state, use\n\n with module.streaming():\n ...\n\n This will automatically reset the streaming state when exiting the context manager.\n This also automatically propagates to all streaming children module.\n\n Some module might also implement the `StreamingModule.flush` method, although\n this one is trickier, as all parents module must be StreamingModule and implement\n it as well for it to work properly. See `StreamingSequential` after.\n \"\"\"\n def __init__(self) -> None:\n super().__init__()\n self._streaming_state: State = {}\n self._is_streaming = False\n\n def _apply_named_streaming(self, fn: tp.Any):\n for name, module in self.named_modules():\n if isinstance(module, StreamingModule):\n fn(name, module)\n\n def _set_streaming(self, streaming: bool):\n def _set_streaming(name, module):\n module._is_streaming = streaming\n self._apply_named_streaming(_set_streaming)\n\n @contextmanager\n def streaming(self):\n \"\"\"Context manager to enter streaming mode. Reset streaming state on exit.\"\"\"\n self._set_streaming(True)\n try:\n yield\n finally:\n self._set_streaming(False)\n self.reset_streaming()\n\n def reset_streaming(self):\n \"\"\"Reset the streaming state.\"\"\"\n def _reset(name: str, module: StreamingModule):\n module._streaming_state.clear()\n\n self._apply_named_streaming(_reset)\n\n def get_streaming_state(self) -> State:\n \"\"\"Return the streaming state, including that of sub-modules.\"\"\"\n state: State = {}\n\n def _add(name: str, module: StreamingModule):\n if name:\n name += \".\"\n for key, value in module._streaming_state.items():\n state[name + key] = value\n\n self._apply_named_streaming(_add)\n return state\n\n def set_streaming_state(self, state: State):\n \"\"\"Set the streaming state, including that of sub-modules.\"\"\"\n state = dict(state)\n\n def _set(name: str, module: StreamingModule):\n if name:\n name += \".\"\n module._streaming_state.clear()\n for key, value in list(state.items()):\n # complexity is not ideal here, but probably fine.\n if key.startswith(name):\n local_key = key[len(name):]\n if '.' not in local_key:\n module._streaming_state[local_key] = value\n del state[key]\n\n self._apply_named_streaming(_set)\n assert len(state) == 0, list(state.keys())\n\n def flush(self, x: tp.Optional[torch.Tensor] = None):\n \"\"\"Flush any remaining outputs that were waiting for completion.\n Typically, for convolutions, this will add the final padding\n and process the last buffer.\n\n This should take an optional argument `x`, which will be provided\n if a module before this one in the streaming pipeline has already\n spitted out a flushed out buffer.\n \"\"\"\n if x is None:\n return None\n else:\n return self(x)\n\n\nclass StreamingSequential(StreamingModule, nn.Sequential):\n \"\"\"A streaming compatible alternative of `nn.Sequential`.\n \"\"\"\n def flush(self, x: tp.Optional[torch.Tensor] = None):\n for module in self:\n if isinstance(module, StreamingModule):\n x = module.flush(x)\n elif x is not None:\n x = module(x)\n return x\n" }, { "path": "audiocraft/modules/transformer.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nTransformer model, with streaming support, xformer attention support\nand easy causal attention with a potentially finite receptive field.\n\nSee `StreamingTransformer` for more information.\n\nUnlike regular PyTorch Transformer, we make the hard choice that batches are first.\n\"\"\"\n\nimport typing as tp\n\nfrom einops import rearrange\nimport torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\nfrom torch.utils.checkpoint import checkpoint as torch_checkpoint\nfrom xformers import ops\n\nfrom .rope import RotaryEmbedding\nfrom .streaming import StreamingModule\n\n_efficient_attention_backend: str = 'torch'\n\n\ndef set_efficient_attention_backend(backend: str = 'torch'):\n # Using torch by default, it seems a bit faster on older P100 GPUs (~20% faster).\n global _efficient_attention_backend\n assert _efficient_attention_backend in ['xformers', 'torch']\n _efficient_attention_backend = backend\n\n\ndef _get_attention_time_dimension(memory_efficient: bool) -> int:\n if _efficient_attention_backend == 'torch' and memory_efficient:\n return 2\n else:\n return 1\n\n\ndef _is_profiled() -> bool:\n # Return true if we are currently running with a xformers profiler activated.\n try:\n from xformers.profiler import profiler\n except ImportError:\n return False\n return profiler._Profiler._CURRENT_PROFILER is not None\n\n\ndef create_norm_fn(norm_type: str, dim: int, **kwargs) -> nn.Module:\n \"\"\"Create normalization module for transformer encoder layer.\n\n Args:\n norm_type (str): Normalization method.\n dim (int): Dimension of the normalized layer.\n **kwargs (dict): Additional parameters for normalization layer.\n Returns:\n nn.Module: Normalization module.\n \"\"\"\n if norm_type == 'layer_norm':\n return nn.LayerNorm(dim, eps=1e-5, **kwargs)\n else:\n raise ValueError(f\"Unknown norm type: {norm_type}\")\n\n\ndef create_sin_embedding(positions: torch.Tensor, dim: int, max_period: float = 10000,\n dtype: torch.dtype = torch.float32) -> torch.Tensor:\n \"\"\"Create sinusoidal positional embedding, with shape `[B, T, C]`.\n\n Args:\n positions (torch.Tensor): LongTensor of positions.\n dim (int): Dimension of the embedding.\n max_period (float): Maximum period of the cosine/sine functions.\n dtype (torch.dtype or str): dtype to use to generate the embedding.\n Returns:\n torch.Tensor: Sinusoidal positional embedding.\n \"\"\"\n # We aim for BTC format\n assert dim % 2 == 0\n half_dim = dim // 2\n positions = positions.to(dtype)\n adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1)\n max_period_tensor = torch.full([], max_period, device=positions.device, dtype=dtype) # avoid sync point\n phase = positions / (max_period_tensor ** (adim / (half_dim - 1)))\n return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)\n\n\ndef expand_repeated_kv(x: torch.Tensor, n_rep: int, memory_efficient: bool) -> torch.Tensor:\n \"\"\"torch.repeat_interleave(x, dim=2, repeats=n_rep) from xlformers.\"\"\"\n if n_rep == 1:\n return x\n if _efficient_attention_backend == 'torch' and memory_efficient:\n bs, n_kv_heads, slen, head_dim = x.shape\n return (\n x[:, :, None, :, :]\n .expand(bs, n_kv_heads, n_rep, slen, head_dim)\n .reshape(bs, n_kv_heads * n_rep, slen, head_dim)\n )\n else:\n bs, slen, n_kv_heads, head_dim = x.shape\n return (\n x[:, :, :, None, :]\n .expand(bs, slen, n_kv_heads, n_rep, head_dim)\n .reshape(bs, slen, n_kv_heads * n_rep, head_dim)\n )\n\n\nclass LayerScale(nn.Module):\n \"\"\"Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf).\n This rescales diagonally the residual outputs close to 0, with a learnt scale.\n\n Args:\n channels (int): Number of channels.\n init (float): Initial scale.\n channel_last (bool): If True, expect `[*, C]` shaped tensors, otherwise, `[*, C, T]`.\n device (torch.device or str, optional): Device on which to initialize the module.\n dtype (torch.dtype, optional): dtype to use to initialize the module.\n \"\"\"\n def __init__(self, channels: int, init: float = 1e-4, channel_last: bool = True,\n device=None, dtype=None):\n super().__init__()\n self.channel_last = channel_last\n self.scale = nn.Parameter(\n torch.full((channels,), init,\n requires_grad=True, device=device, dtype=dtype))\n\n def forward(self, x: torch.Tensor):\n if self.channel_last:\n return self.scale * x\n else:\n return self.scale[:, None] * x\n\n\nclass StreamingMultiheadAttention(StreamingModule):\n \"\"\"Similar to `nn.MultiheadAttention` but with support for streaming, causal evaluation.\n\n Args:\n embed_dim (int): Dimension to project to.\n num_heads (int): Number of heads.\n dropout (float): Dropout level.\n bias (bool): Use bias in projections.\n causal (bool): Causal mask applied automatically.\n past_context (int, optional): Receptive field for the causal mask, infinite if None.\n custom (bool): Use custom MHA implementation, for testing / benchmarking.\n memory_efficient (bool): Use xformers based memory efficient attention.\n attention_as_float32 (bool): Perform the attention as float32\n (especially important with memory_efficient as autocast won't do this automatically).\n rope (`RotaryEmbedding`, optional): Rope embedding to use.\n cross_attention: Should be true when used as a cross attention.\n All keys and values must be available at once, streaming is only for the queries.\n Cannot be used with `causal` or `rope` (as it wouldn't make sens to\n interpret the time steps in the keys relative to those in the queries).\n safe_streaming (bool): Bug fix, will go away with xformers update.\n qk_layer_norm (bool): Layer normalization applied to queries and keys before dot product.\n kv_repeat (int): If > 1, will repeat keys and queries multiple times (need to divide num_heads).\n This will lead to faster decoding time on A100 or other GPUs with tensorcore.\n device (torch.device, optional): Device on which to initialize.\n dtype (torch.dtype, optional): dtype to use.\n \"\"\"\n def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0, bias: bool = True,\n causal: bool = False, past_context: tp.Optional[int] = None, custom: bool = False,\n memory_efficient: bool = False, attention_as_float32: bool = False,\n rope: tp.Optional[RotaryEmbedding] = None, cross_attention: bool = False,\n safe_streaming: bool = True, qk_layer_norm: bool = False, kv_repeat: int = 1,\n device=None, dtype=None):\n super().__init__()\n factory_kwargs = {'device': device, 'dtype': dtype}\n if past_context is not None:\n assert causal\n\n self.embed_dim = embed_dim\n self.causal = causal\n self.past_context = past_context\n self.memory_efficient = memory_efficient\n self.attention_as_float32 = attention_as_float32\n self.rope = rope\n self.cross_attention = cross_attention\n self.safe_streaming = safe_streaming\n self.num_heads = num_heads\n self.dropout = dropout\n self.kv_repeat = kv_repeat\n if cross_attention:\n assert not causal, \"Causal cannot work with cross attention.\"\n assert rope is None, \"Rope cannot work with cross attention.\"\n\n if memory_efficient:\n _verify_xformers_memory_efficient_compat()\n\n self.custom = _is_custom(custom, memory_efficient)\n if self.custom:\n out_dim = embed_dim\n assert num_heads % kv_repeat == 0\n assert not cross_attention or kv_repeat == 1\n num_kv = num_heads // kv_repeat\n kv_dim = (embed_dim // num_heads) * num_kv\n out_dim += 2 * kv_dim\n in_proj = nn.Linear(embed_dim, out_dim, bias=bias, **factory_kwargs)\n # We try to follow the default PyTorch MHA convention, to easily compare results.\n self.in_proj_weight = in_proj.weight\n self.in_proj_bias = in_proj.bias\n if bias:\n self.in_proj_bias.data.zero_() # Following Pytorch convention\n self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias, **factory_kwargs)\n if bias:\n self.out_proj.bias.data.zero_()\n else:\n assert not qk_layer_norm\n assert kv_repeat == 1\n self.mha = nn.MultiheadAttention(\n embed_dim, num_heads, dropout=dropout, bias=bias, batch_first=True,\n **factory_kwargs)\n self.qk_layer_norm = qk_layer_norm\n if qk_layer_norm:\n assert self.custom\n assert kv_repeat == 1\n ln_dim = embed_dim\n self.q_layer_norm = nn.LayerNorm(ln_dim)\n self.k_layer_norm = nn.LayerNorm(ln_dim)\n\n def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):\n if not self.custom:\n # Support compat with regular MHA\n keys = [n for n, _ in self.mha.named_parameters()]\n for key in keys:\n if prefix + key in state_dict:\n state_dict[prefix + \"mha.\" + key] = state_dict.pop(prefix + key)\n super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)\n\n def _get_mask(self, current_steps: int, device: torch.device, dtype: torch.dtype):\n # Return a causal mask, accounting for potentially stored past keys/values\n # We actually return a bias for the attention score, as this has the same\n # convention both in the builtin MHA in Pytorch, and Xformers functions.\n time_dim = _get_attention_time_dimension(self.memory_efficient)\n if self.memory_efficient:\n from xformers.ops import LowerTriangularMask\n if current_steps == 1:\n # If we only have one step, then we do not need a mask.\n return None\n elif 'past_keys' in self._streaming_state:\n raise RuntimeError(\"Not supported at the moment\")\n else:\n # Then we can safely use a lower triangular mask\n return LowerTriangularMask()\n if self._streaming_state:\n past_keys = self._streaming_state['past_keys']\n past_steps = past_keys.shape[time_dim]\n else:\n past_steps = 0\n\n queries_pos = torch.arange(\n past_steps, current_steps + past_steps, device=device).view(-1, 1)\n keys_pos = torch.arange(past_steps + current_steps, device=device).view(1, -1)\n delta = queries_pos - keys_pos\n valid = delta >= 0\n if self.past_context is not None:\n valid &= (delta <= self.past_context)\n return torch.where(\n valid,\n torch.zeros([], device=device, dtype=dtype),\n torch.full([], float('-inf'), device=device, dtype=dtype))\n\n def _complete_kv(self, k, v):\n time_dim = _get_attention_time_dimension(self.memory_efficient)\n if self.cross_attention:\n # With cross attention we assume all keys and values\n # are already available, and streaming is with respect\n # to the queries only.\n return k, v\n # Complete the key/value pair using the streaming state.\n if self._streaming_state:\n pk = self._streaming_state['past_keys']\n nk = torch.cat([pk, k], dim=time_dim)\n if v is k:\n nv = nk\n else:\n pv = self._streaming_state['past_values']\n nv = torch.cat([pv, v], dim=time_dim)\n else:\n nk = k\n nv = v\n\n assert nk.shape[time_dim] == nv.shape[time_dim]\n offset = 0\n if self.past_context is not None:\n offset = max(0, nk.shape[time_dim] - self.past_context)\n if self._is_streaming:\n self._streaming_state['past_keys'] = nk[:, offset:]\n if v is not k:\n self._streaming_state['past_values'] = nv[:, offset:]\n if 'offset' in self._streaming_state:\n self._streaming_state['offset'] += offset\n else:\n self._streaming_state['offset'] = torch.tensor(0)\n return nk, nv\n\n def _apply_rope(self, query: torch.Tensor, key: torch.Tensor):\n time_dim = _get_attention_time_dimension(self.memory_efficient)\n # Apply rope embeddings to query and key tensors.\n assert self.rope is not None\n if 'past_keys' in self._streaming_state:\n past_keys_offset = self._streaming_state['past_keys'].shape[1]\n else:\n past_keys_offset = 0\n if 'offset' in self._streaming_state:\n past_context_offset = int(self._streaming_state['offset'].item())\n else:\n past_context_offset = 0\n streaming_offset = past_context_offset + past_keys_offset\n return self.rope.rotate_qk(query, key, start=streaming_offset, time_dim=time_dim)\n\n def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,\n key_padding_mask=None, need_weights=False, attn_mask=None,\n average_attn_weights=True, is_causal=False):\n assert not is_causal, (\"New param added in torch 2.0.1 not supported, \"\n \"use the causal args in the constructor.\")\n\n time_dim = _get_attention_time_dimension(self.memory_efficient)\n if time_dim == 2:\n layout = \"b h t d\"\n else:\n layout = \"b t h d\"\n dtype = query.dtype\n if self._is_streaming:\n assert self.causal or self.cross_attention, \\\n \"Streaming only available for causal or cross attention\"\n\n custom_attn_mask = attn_mask is not None\n\n if self.causal:\n assert attn_mask is None\n # At the moment we specialize only for the self-attention case.\n assert query.shape[1] == key.shape[1], \"Causal only for same length query / key / value\"\n assert value.shape[1] == key.shape[1], \"Causal only for same length query / key / value\"\n attn_mask = self._get_mask(query.shape[1], query.device, query.dtype)\n\n if self.custom:\n # custom implementation\n assert need_weights is False\n assert key_padding_mask is None\n if self.cross_attention:\n # Different queries, keys, values, we have to spit manually the weights\n # before applying the linear.\n dim = self.in_proj_weight.shape[0] // 3\n if self.in_proj_bias is None:\n bias_q, bias_k, bias_v = None, None, None\n else:\n bias_q = self.in_proj_bias[:dim]\n bias_k = self.in_proj_bias[dim: 2 * dim]\n bias_v = self.in_proj_bias[2 * dim:]\n q = nn.functional.linear(query, self.in_proj_weight[:dim], bias_q)\n # todo: when streaming, we could actually save k, v and check the shape actually match.\n k = nn.functional.linear(key, self.in_proj_weight[dim: 2 * dim], bias_k)\n v = nn.functional.linear(value, self.in_proj_weight[2 * dim:], bias_v)\n if self.qk_layer_norm is True:\n q = self.q_layer_norm(q)\n k = self.k_layer_norm(k)\n q, k, v = [rearrange(x, f\"b t (h d) -> {layout}\", h=self.num_heads) for x in [q, k, v]]\n else:\n if not _is_profiled():\n # profiling breaks that propertysomehow.\n assert query is key, \"specialized implementation\"\n assert value is key, \"specialized implementation\"\n projected = nn.functional.linear(query, self.in_proj_weight, self.in_proj_bias)\n if self.kv_repeat == 1:\n if time_dim == 2:\n bound_layout = \"b h p t d\"\n else:\n bound_layout = \"b t p h d\"\n packed = rearrange(projected, f\"b t (p h d) -> {bound_layout}\", p=3, h=self.num_heads)\n q, k, v = ops.unbind(packed, dim=2)\n else:\n embed_dim = self.embed_dim\n per_head_dim = (embed_dim // self.num_heads)\n kv_heads = self.num_heads // self.kv_repeat\n q = projected[:, :, :embed_dim]\n start = embed_dim\n end = start + per_head_dim * kv_heads\n k = projected[:, :, start: end]\n v = projected[:, :, end:]\n q = rearrange(q, f\"b t (h d) -> {layout}\", h=self.num_heads)\n k = rearrange(k, f\"b t (h d) -> {layout}\", h=kv_heads)\n v = rearrange(v, f\"b t (h d) -> {layout}\", h=kv_heads)\n\n if self.qk_layer_norm is True:\n assert self.kv_repeat == 1\n q, k = [rearrange(x, f\"{layout} -> b t (h d)\") for x in [q, k]]\n q = self.q_layer_norm(q)\n k = self.k_layer_norm(k)\n q, k = [rearrange(x, f\"b t (h d) -> {layout}\", h=self.num_heads) for x in [q, k]]\n if self.rope:\n q, k = self._apply_rope(q, k)\n k, v = self._complete_kv(k, v)\n if self.kv_repeat > 1:\n k = expand_repeated_kv(k, self.kv_repeat, self.memory_efficient)\n v = expand_repeated_kv(v, self.kv_repeat, self.memory_efficient)\n if self.attention_as_float32:\n q, k, v = [x.float() for x in [q, k, v]]\n if self.memory_efficient:\n if custom_attn_mask:\n # When using a custom attn mask:\n # Move to query's device, repeat for each sample, remove align8 padding\n seq_len = query.shape[1]\n attn_mask = attn_mask.to(q.dtype)\n attn_mask = attn_mask.repeat((q.shape[0], 1, 1, 1))\n attn_mask = attn_mask[..., :seq_len, :seq_len]\n\n p = self.dropout if self.training else 0\n if _efficient_attention_backend == 'torch':\n x = torch.nn.functional.scaled_dot_product_attention(\n q, k, v, is_causal=attn_mask is not None, dropout_p=p)\n else:\n x = ops.memory_efficient_attention(q, k, v, attn_mask, p=p)\n else:\n # We include the dot product as float32, for consistency\n # with the other implementations that include that step\n # as part of the attention. Note that when using `autocast`,\n # the einsums would be done as bfloat16, but the softmax\n # would be done as bfloat16, so `attention_as_float32` will\n # extend a bit the range of operations done in float32,\n # although this should make no difference.\n q = q / q.shape[-1] ** 0.5\n key_layout = layout.replace('t', 'k')\n query_layout = layout\n if self._is_streaming and self.safe_streaming and q.device.type == 'cuda':\n with torch.autocast(device_type=q.device.type, dtype=torch.float32):\n pre_w = torch.einsum(f\"{query_layout},{key_layout}-> b h t k\", q, k)\n else:\n pre_w = torch.einsum(f\"{query_layout},{key_layout}-> b h t k\", q, k)\n if attn_mask is not None:\n pre_w = pre_w + attn_mask\n w = torch.softmax(pre_w, dim=-1)\n w = F.dropout(w, self.dropout, training=self.training).to(v)\n # Key and value have the same format.\n x = torch.einsum(f\"b h t k, {key_layout} -> {layout}\", w, v)\n x = x.to(dtype)\n x = rearrange(x, f\"{layout} -> b t (h d)\", h=self.num_heads)\n x = self.out_proj(x)\n else:\n key, value = self._complete_kv(key, value)\n if self.attention_as_float32:\n query, key, value = [x.float() for x in [query, key, value]]\n x, _ = self.mha(\n query, key, value, key_padding_mask,\n need_weights, attn_mask, average_attn_weights)\n x = x.to(dtype)\n\n return x, None\n\n\nclass StreamingTransformerLayer(nn.TransformerEncoderLayer):\n \"\"\"TransformerLayer with Streaming / Causal support.\n This also integrates cross_attention, when passing `cross_attention=True`,\n rather than having two separate classes like in PyTorch.\n\n Args:\n d_model (int): Dimension of the data.\n num_heads (int): Number of heads.\n dim_feedforward (int): Intermediate dimension of FF module.\n dropout (float): Dropout both for MHA and FF.\n bias_ff (bool): Use bias for FF.\n bias_attn (bool): Use bias for MHA.\n causal (bool): Causal mask applied automatically.\n past_context (int, optional): Receptive field for the causal mask, infinite if None.\n custom (bool): Use custom MHA implementation, for testing / benchmarking.\n memory_efficient (bool): Use xformers based memory efficient attention.\n attention_as_float32 (bool): Perform the attention as float32\n (especially important with memory_efficient as autocast won't do this automatically).\n qk_layer_norm (bool): Layer normalization applied to queries and keys before dot product in attention.\n qk_layer_norm_cross (bool): Same for the cross attention.\n cross_attention (bool): If True, expect to get secondary input for cross-attention.\n Cross attention will use the default MHA, as it typically won't require\n special treatment.\n layer_scale (float, optional): If not None, LayerScale will be used with\n the given value as initial scale.\n rope (`RotaryEmbedding`, optional): Rope embedding to use.\n attention_dropout (float, optional): If not None, separate the value of the dimension dropout\n in FFN and of the attention dropout.\n kv_repeat (int): If > 1, will repeat keys and queries multiple times (need to divide num_heads).\n This will lead to faster decoding time on A100 or other GPUs with tensorcore.\n device (torch.device, optional): Device on which to initialize.\n dtype (torch.dtype, optional): dtype to use.\n **kwargs: See `nn.TransformerEncoderLayer`.\n \"\"\"\n def __init__(self, d_model: int, num_heads: int, dim_feedforward: int = 2048, dropout: float = 0.1,\n bias_ff: bool = True, bias_attn: bool = True, causal: bool = False,\n past_context: tp.Optional[int] = None, custom: bool = False,\n memory_efficient: bool = False, attention_as_float32: bool = False,\n qk_layer_norm: bool = False, qk_layer_norm_cross: bool = False,\n cross_attention: bool = False, layer_scale: tp.Optional[float] = None,\n rope: tp.Optional[RotaryEmbedding] = None, attention_dropout: tp.Optional[float] = None,\n kv_repeat: int = 1, norm: str = 'layer_norm', device=None, dtype=None, **kwargs):\n super().__init__(d_model, num_heads, dim_feedforward, dropout,\n device=device, dtype=dtype, batch_first=True, **kwargs)\n factory_kwargs = {'device': device, 'dtype': dtype}\n # Redefine self_attn to our streaming multi-head attention\n attn_kwargs: tp.Dict[str, tp.Any] = {\n 'embed_dim': d_model,\n 'num_heads': num_heads,\n 'dropout': dropout if attention_dropout is None else attention_dropout,\n 'bias': bias_attn,\n 'custom': custom,\n 'memory_efficient': memory_efficient,\n 'attention_as_float32': attention_as_float32,\n }\n self.self_attn: StreamingMultiheadAttention = StreamingMultiheadAttention(\n causal=causal, past_context=past_context, rope=rope, qk_layer_norm=qk_layer_norm,\n kv_repeat=kv_repeat, **attn_kwargs, **factory_kwargs) # type: ignore\n # Redefine feedforward layers to expose bias parameter\n self.linear1 = nn.Linear(d_model, dim_feedforward, bias=bias_ff, **factory_kwargs)\n self.linear2 = nn.Linear(dim_feedforward, d_model, bias=bias_ff, **factory_kwargs)\n\n self.layer_scale_1: nn.Module\n self.layer_scale_2: nn.Module\n if layer_scale is None:\n self.layer_scale_1 = nn.Identity()\n self.layer_scale_2 = nn.Identity()\n else:\n self.layer_scale_1 = LayerScale(d_model, layer_scale, **factory_kwargs)\n self.layer_scale_2 = LayerScale(d_model, layer_scale, **factory_kwargs)\n\n self.cross_attention: tp.Optional[nn.Module] = None\n if cross_attention:\n self.cross_attention = StreamingMultiheadAttention(\n cross_attention=True, qk_layer_norm=qk_layer_norm_cross,\n **attn_kwargs, **factory_kwargs)\n # Norm and dropout\n self.dropout_cross = nn.Dropout(dropout)\n # eps value matching that used in PyTorch reference implementation.\n self.norm_cross = nn.LayerNorm(d_model, eps=1e-5, **factory_kwargs)\n self.layer_scale_cross: nn.Module\n if layer_scale is None:\n self.layer_scale_cross = nn.Identity()\n else:\n self.layer_scale_cross = LayerScale(d_model, layer_scale, **factory_kwargs)\n self.norm1 = create_norm_fn(norm, d_model, **factory_kwargs) # type: ignore\n self.norm2 = create_norm_fn(norm, d_model, **factory_kwargs) # type: ignore\n\n def _cross_attention_block(self, src: torch.Tensor,\n cross_attention_src: torch.Tensor) -> torch.Tensor:\n assert self.cross_attention is not None\n # queries are from src, keys and values from cross_attention_src.\n x = self.cross_attention(\n src, cross_attention_src, cross_attention_src, need_weights=False)[0]\n return self.dropout_cross(x) # type: ignore\n\n def forward(self, src: torch.Tensor, src_mask: tp.Optional[torch.Tensor] = None, # type: ignore\n src_key_padding_mask: tp.Optional[torch.Tensor] = None,\n cross_attention_src: tp.Optional[torch.Tensor] = None):\n if self.cross_attention is None:\n assert cross_attention_src is None\n else:\n assert cross_attention_src is not None\n x = src\n if self.norm_first:\n x = x + self.layer_scale_1(\n self._sa_block(self.norm1(x), src_mask, src_key_padding_mask))\n if cross_attention_src is not None:\n x = x + self.layer_scale_cross(\n self._cross_attention_block(\n self.norm_cross(x), cross_attention_src))\n x = x + self.layer_scale_2(self._ff_block(self.norm2(x)))\n else:\n x = self.norm1(x + self.layer_scale_1(\n self._sa_block(x, src_mask, src_key_padding_mask)))\n if cross_attention_src is not None:\n x = self.norm_cross(\n x + self.layer_scale_cross(\n self._cross_attention_block(src, cross_attention_src)))\n x = self.norm2(x + self.layer_scale_2(self._ff_block(x)))\n return x\n\n\nclass StreamingTransformer(StreamingModule):\n \"\"\"Transformer with Streaming / Causal support.\n\n Args:\n d_model (int): Dimension of the data.\n num_heads (int): Number of heads.\n dim_feedforward (int): Intermediate dimension of FF module.\n dropout (float): Dropout both for MHA and FF.\n bias_ff (bool): Use bias for FF.\n bias_attn (bool): Use bias for MHA.\n causal (bool): Causal mask applied automatically.\n past_context (int, optional): Receptive field for the causal mask, infinite if None.\n custom (bool): Use custom MHA implementation, for testing / benchmarking.\n memory_efficient (bool): Use xformers based memory efficient attention.\n attention_as_float32 (bool): Perform the attention as float32\n (especially important with memory_efficient as autocast won't do this automatically).\n cross_attention (bool): If True, expect to get secondary input for cross-attention.\n layer_scale (float, optional): If not None, LayerScale will be used\n with the given value as initial scale.\n positional_embedding (str): Positional embedding strategy (sin, rope, or sin_rope).\n max_period (float): Maximum period of the time embedding.\n positional_scale (float): Scale of positional embedding, set to 0 to deactivate.\n xpos (bool): Apply xpos exponential decay to positional embedding (rope only).\n lr (float, optional): learning rate override through the `make_optim_group` API.\n weight_decay (float, optional): Weight_decay override through the `make_optim_group` API.\n layer_class: (subclass of `StreamingTransformerLayer): class to use\n to initialize the layers, allowing further customization outside of AudioCraft.\n checkpointing (str): Checkpointing strategy to reduce memory usage.\n No checkpointing if set to 'none'. Per layer checkpointing using PyTorch\n if set to 'torch' (entire layer checkpointed, i.e. linears are evaluated twice,\n minimal memory usage, but maximal runtime). Finally, `xformers_default` provide\n a policy for opting-out some operations of the checkpointing like\n linear layers and attention, providing a middle ground between speed and memory.\n device (torch.device, optional): Device on which to initialize.\n dtype (torch.dtype, optional): dtype to use.\n **kwargs: See `nn.TransformerEncoderLayer`.\n \"\"\"\n def __init__(self, d_model: int, num_heads: int, num_layers: int, dim_feedforward: int = 2048,\n dropout: float = 0.1, bias_ff: bool = True, bias_attn: bool = True,\n causal: bool = False, past_context: tp.Optional[int] = None,\n custom: bool = False, memory_efficient: bool = False, attention_as_float32: bool = False,\n cross_attention: bool = False, layer_scale: tp.Optional[float] = None,\n positional_embedding: str = 'sin', max_period: float = 10_000, positional_scale: float = 1.,\n xpos: bool = False, lr: tp.Optional[float] = None, weight_decay: tp.Optional[float] = None,\n layer_class: tp.Type[StreamingTransformerLayer] = StreamingTransformerLayer,\n checkpointing: str = 'none', device=None, dtype=None, **kwargs):\n super().__init__()\n assert d_model % num_heads == 0\n\n self.positional_embedding = positional_embedding\n self.max_period = max_period\n self.positional_scale = positional_scale\n self.weight_decay = weight_decay\n self.lr = lr\n\n assert positional_embedding in ['sin', 'rope', 'sin_rope']\n self.rope: tp.Optional[RotaryEmbedding] = None\n if self.positional_embedding in ['rope', 'sin_rope']:\n assert _is_custom(custom, memory_efficient)\n self.rope = RotaryEmbedding(d_model // num_heads, max_period=max_period,\n xpos=xpos, scale=positional_scale, device=device)\n\n self.checkpointing = checkpointing\n\n assert checkpointing in ['none', 'torch', 'xformers_default', 'xformers_mm']\n if self.checkpointing.startswith('xformers'):\n _verify_xformers_internal_compat()\n\n self.layers = nn.ModuleList()\n for idx in range(num_layers):\n self.layers.append(\n layer_class(\n d_model=d_model, num_heads=num_heads, dim_feedforward=dim_feedforward,\n dropout=dropout, bias_ff=bias_ff, bias_attn=bias_attn,\n causal=causal, past_context=past_context, custom=custom,\n memory_efficient=memory_efficient, attention_as_float32=attention_as_float32,\n cross_attention=cross_attention, layer_scale=layer_scale, rope=self.rope,\n device=device, dtype=dtype, **kwargs))\n\n if self.checkpointing != 'none':\n for layer in self.layers:\n # see audiocraft/optim/fsdp.py, magic signal to indicate this requires fixing the\n # backward hook inside of FSDP...\n layer._magma_checkpointed = True # type: ignore\n\n def _apply_layer(self, layer, *args, **kwargs):\n method = self.checkpointing\n if method == 'none':\n return layer(*args, **kwargs)\n elif method == 'torch':\n return torch_checkpoint(layer, *args, use_reentrant=False, **kwargs)\n elif method.startswith('xformers'):\n from xformers.checkpoint_fairinternal import checkpoint, _get_default_policy\n if method == 'xformers_default':\n # those operations will be saved, and not recomputed.\n # According to Francisco we can get smarter policies but this is a good start.\n allow_list = [\n \"xformers.efficient_attention_forward_cutlass.default\",\n \"xformers_flash.flash_fwd.default\",\n \"aten.addmm.default\",\n \"aten.mm.default\",\n ]\n elif method == 'xformers_mm':\n # those operations will be saved, and not recomputed.\n # According to Francisco we can get smarter policies but this is a good start.\n allow_list = [\n \"aten.addmm.default\",\n \"aten.mm.default\",\n ]\n else:\n raise ValueError(f\"xformers checkpointing xformers policy {method} is not known.\")\n policy_fn = _get_default_policy(allow_list)\n return checkpoint(layer, *args, policy_fn=policy_fn, **kwargs)\n else:\n raise ValueError(f\"Checkpointing method {method} is unknown.\")\n\n def forward(self, x: torch.Tensor, *args, **kwargs):\n B, T, C = x.shape\n\n if 'offsets' in self._streaming_state:\n offsets = self._streaming_state['offsets']\n else:\n offsets = torch.zeros(B, dtype=torch.long, device=x.device)\n\n if self.positional_embedding in ['sin', 'sin_rope']:\n positions = torch.arange(T, device=x.device).view(1, -1, 1)\n positions = positions + offsets.view(-1, 1, 1)\n pos_emb = create_sin_embedding(positions, C, max_period=self.max_period, dtype=x.dtype)\n x = x + self.positional_scale * pos_emb\n\n for layer in self.layers:\n x = self._apply_layer(layer, x, *args, **kwargs)\n\n if self._is_streaming:\n self._streaming_state['offsets'] = offsets + T\n\n return x\n\n def make_optim_group(self):\n group = {\"params\": list(self.parameters())}\n if self.lr is not None:\n group[\"lr\"] = self.lr\n if self.weight_decay is not None:\n group[\"weight_decay\"] = self.weight_decay\n return group\n\n\n# special attention related function\n\ndef _verify_xformers_memory_efficient_compat():\n try:\n from xformers.ops import memory_efficient_attention, LowerTriangularMask # noqa\n except ImportError:\n raise ImportError(\n \"xformers is not installed. Please install it and try again.\\n\"\n \"To install on AWS and Azure, run \\n\"\n \"FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST='8.0'\\\\\\n\"\n \"pip install -U git+https://git@github.com/fairinternal/xformers.git#egg=xformers\\n\"\n \"To install on FAIR Cluster, run \\n\"\n \"FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST='6.0;7.0'\\\\\\n\"\n \"pip install -U git+https://git@github.com/fairinternal/xformers.git#egg=xformers\\n\")\n\n\ndef _verify_xformers_internal_compat():\n try:\n from xformers.checkpoint_fairinternal import checkpoint, _get_default_policy # noqa\n except ImportError:\n raise ImportError(\n \"Francisco's fairinternal xformers is not installed. Please install it and try again.\\n\"\n \"To install on AWS and Azure, run \\n\"\n \"FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST='8.0'\\\\\\n\"\n \"pip install -U git+https://git@github.com/fairinternal/xformers.git#egg=xformers\\n\"\n \"To install on FAIR Cluster, run \\n\"\n \"FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST='6.0;7.0'\\\\\\n\"\n \"pip install -U git+https://git@github.com/fairinternal/xformers.git#egg=xformers\\n\")\n\n\ndef _is_custom(custom: bool, memory_efficient: bool):\n return custom or memory_efficient\n" }, { "path": "audiocraft/modules/unet_transformer.py", "content": "import torch\nimport typing as tp\nfrom .transformer import StreamingTransformer, create_sin_embedding\n\n\nclass UnetTransformer(StreamingTransformer):\n \"\"\"U-net Transformer for processing sequences with optional skip connections.\n This transformer architecture incorporates U-net style skip connections\n between layers, which can be optionally enabled. It inherits from a\n StreamingTransformer.\n\n Args:\n d_model (int): Dimension of the model, typically the number of expected features in the input.\n num_layers (int): Total number of layers in the transformer.\n skip_connections (bool, optional): Flag to determine whether skip connections should be used.\n Defaults to False.\n layer_dropout_p (float, Optional): if given, defined bernoulli prob. to drop a skip connection (in training).\n **kwargs: Additional keyword arguments inherited from `nn.StreamingTransformer`.\n \"\"\"\n def __init__(self, d_model: int, num_layers: int, skip_connections: bool = False,\n layer_dropout_p: tp.Optional[float] = None, **kwargs):\n super().__init__(d_model=d_model,\n num_layers=num_layers,\n **kwargs)\n self.skip_connect = skip_connections\n if self.skip_connect:\n self.skip_projections = torch.nn.ModuleList([torch.nn.Linear(d_model * 2, d_model)\n for _ in range(num_layers // 2)])\n self.num_layers = num_layers\n self.layer_drop_p = max(min(layer_dropout_p, 1.), 0.) if layer_dropout_p is not None else 0.0\n\n def forward(self, x: torch.Tensor, *args, **kwargs):\n B, T, C = x.shape\n\n if 'offsets' in self._streaming_state:\n offsets = self._streaming_state['offsets']\n else:\n offsets = torch.zeros(B, dtype=torch.long, device=x.device)\n\n if self.positional_embedding in ['sin', 'sin_rope']:\n positions = torch.arange(T, device=x.device).view(1, -1, 1)\n positions = positions + offsets.view(-1, 1, 1)\n pos_emb = create_sin_embedding(positions, C, max_period=self.max_period, dtype=x.dtype)\n x = x + self.positional_scale * pos_emb\n\n skip_connections: tp.List[torch.Tensor] = []\n\n for i, layer in enumerate(self.layers):\n if self.skip_connect and i >= self.num_layers // 2:\n\n # in the second half of the layers, add residual connection\n # and linearly project the concatenated features back to d_model\n x = torch.cat([x, skip_connections.pop()], dim=-1)\n x = self.skip_projections[i % len(self.skip_projections)](x)\n\n x = self._apply_layer(layer, x, *args, **kwargs)\n\n if self.skip_connect and i < self.num_layers // 2:\n if self.training and torch.rand(1,) < self.layer_drop_p: # drop skip\n skip_connections.append(torch.zeros_like(x))\n else:\n skip_connections.append(x)\n\n if self._is_streaming:\n self._streaming_state['offsets'] = offsets + T\n\n return x\n" }, { "path": "audiocraft/modules/watermark.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\nimport random\n\nimport torch\n\n\ndef pad(\n x_wm: torch.Tensor, central: bool = False\n) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Pad a watermarked signal at the begining and the end\n\n Args:\n x_wm (torch.Tensor) : watermarked audio\n central (bool): Whether to mask the middle of the wave (around 34%) or the two tails\n (beginning and ending frames)\n\n Returns:\n padded (torch.Tensor): padded signal\n true_predictions(torch.Tensor): A binary mask where 1 represents\n watermarked and 0 represents non-watermarked.\"\"\"\n # keep at leat 34% of watermarked signal\n max_start = int(0.33 * x_wm.size(-1))\n min_end = int(0.66 * x_wm.size(-1))\n starts = torch.randint(0, max_start, size=(x_wm.size(0),))\n ends = torch.randint(min_end, x_wm.size(-1), size=(x_wm.size(0),))\n mask = torch.zeros_like(x_wm)\n for i in range(x_wm.size(0)):\n mask[i, :, starts[i]: ends[i]] = 1\n if central:\n mask = 1 - mask\n padded = x_wm * mask\n true_predictions = torch.cat([1 - mask, mask], dim=1)\n return padded, true_predictions\n\n\ndef mix(\n x: torch.Tensor, x_wm: torch.Tensor, window_size: float = 0.5, shuffle: bool = False\n) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"\n Mixes a window of the non-watermarked audio signal 'x' into the watermarked audio signal 'x_wm'.\n\n This function takes two tensors of shape [batch, channels, frames], copies a window of 'x' with the specified\n 'window_size' into 'x_wm', and returns a new tensor that is a mix between the watermarked (1 - mix_percent %)\n and non-watermarked audio (mix_percent %).\n\n Args:\n x (torch.Tensor): The non-watermarked audio signal tensor.\n x_wm (torch.Tensor): The watermarked audio signal tensor.\n window_size (float, optional): The percentage of 'x' to copy into 'x_wm' (between 0 and 1).\n shuffle (bool): whether or no keep the mix from the same batch element\n\n Returns:\n tuple: A tuple containing two tensors:\n - mixed_tensor (torch.Tensor): The resulting mixed audio signal tensor.\n - mask (torch.Tensor): A binary mask where 1 represents watermarked and 0 represents non-watermarked.\n\n Raises:\n AssertionError: If 'window_size' is not between 0 and 1.\n \"\"\"\n assert 0 < window_size <= 1, \"window_size should be between 0 and 1\"\n\n # Calculate the maximum starting point for the window\n max_start_point = x.shape[-1] - int(window_size * x.shape[-1])\n\n # Generate a random starting point within the adjusted valid range\n start_point = random.randint(0, max_start_point)\n\n # Calculate the window size in frames\n total_frames = x.shape[-1]\n window_frames = int(window_size * total_frames)\n\n # Create a mask tensor to identify watermarked and non-watermarked portions\n # it outputs two classes to match the detector output shape of [bsz, 2, frames]\n # Copy the random window from 'x' to 'x_wm'\n mixed = x_wm.detach().clone()\n\n true_predictions = torch.cat(\n [torch.zeros_like(mixed), torch.ones_like(mixed)], dim=1\n )\n # non-watermark class correct labels.\n true_predictions[:, 0, start_point: start_point + window_frames] = 1.0\n # watermarked class correct labels\n true_predictions[:, 1, start_point: start_point + window_frames] = 0.0\n\n if shuffle:\n # Take the middle part from a random element of the batch\n shuffle_idx = torch.randint(0, x.size(0), (x.size(0),))\n mixed[:, :, start_point: start_point + window_frames] = x[shuffle_idx][\n :, :, start_point: start_point + window_frames\n ]\n else:\n mixed[:, :, start_point: start_point + window_frames] = x[\n :, :, start_point: start_point + window_frames\n ]\n\n return mixed, true_predictions\n" }, { "path": "audiocraft/optim/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Optimization stuff. In particular, optimizers (DAdaptAdam), schedulers\nand Exponential Moving Average.\n\"\"\"\n\n# flake8: noqa\nfrom .cosine_lr_scheduler import CosineLRScheduler\nfrom .dadam import DAdaptAdam\nfrom .inverse_sqrt_lr_scheduler import InverseSquareRootLRScheduler\nfrom .linear_warmup_lr_scheduler import LinearWarmupLRScheduler\nfrom .polynomial_decay_lr_scheduler import PolynomialDecayLRScheduler\nfrom .ema import ModuleDictEMA\n" }, { "path": "audiocraft/optim/cosine_lr_scheduler.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\n\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler\n\n\nclass CosineLRScheduler(_LRScheduler):\n \"\"\"Cosine LR scheduler.\n\n Args:\n optimizer (Optimizer): Torch optimizer.\n warmup_steps (int): Number of warmup steps.\n total_steps (int): Total number of steps.\n lr_min_ratio (float): Minimum learning rate.\n cycle_length (float): Cycle length.\n \"\"\"\n def __init__(self, optimizer: Optimizer, total_steps: int, warmup_steps: int,\n lr_min_ratio: float = 0.0, cycle_length: float = 1.0):\n self.warmup_steps = warmup_steps\n assert self.warmup_steps >= 0\n self.total_steps = total_steps\n assert self.total_steps >= 0\n self.lr_min_ratio = lr_min_ratio\n self.cycle_length = cycle_length\n super().__init__(optimizer)\n\n def _get_sched_lr(self, lr: float, step: int):\n if step < self.warmup_steps:\n lr_ratio = step / self.warmup_steps\n lr = lr_ratio * lr\n elif step <= self.total_steps:\n s = (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)\n lr_ratio = self.lr_min_ratio + 0.5 * (1 - self.lr_min_ratio) * \\\n (1. + math.cos(math.pi * s / self.cycle_length))\n lr = lr_ratio * lr\n else:\n lr_ratio = self.lr_min_ratio\n lr = lr_ratio * lr\n return lr\n\n def get_lr(self):\n return [self._get_sched_lr(lr, self.last_epoch) for lr in self.base_lrs]\n" }, { "path": "audiocraft/optim/dadam.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nfrom typing import Any\n\nimport torch\nimport torch.optim\nimport torch.distributed as dist\n\n\nlogger = logging.getLogger(__name__)\n_params_t = Any\n\n\ndef to_real(x):\n if torch.is_complex(x):\n return x.real\n else:\n return x\n\n\nclass DAdaptAdam(torch.optim.Optimizer):\n \"\"\"Adam with D-Adaptation automatic step-sizes.\n Leave LR set to 1 unless you encounter instability.\n\n Args:\n params (iterable):\n Iterable of parameters to optimize or dicts defining parameter groups.\n lr (float):\n Learning rate adjustment parameter. Increases or decreases the D-adapted learning rate.\n betas (tuple[float, float], optional): coefficients used for computing\n running averages of gradient and its square (default: (0.9, 0.999))\n momentum (float):\n Momentum value in the range [0,1) (default: 0.9).\n eps (float):\n Term added to the denominator outside of the root operation to improve numerical stability. (default: 1e-8).\n weight_decay (float):\n Weight decay, i.e. a L2 penalty (default: 0).\n log_every (int):\n Log using print every k steps, default 0 (no logging).\n decouple (boolean):\n Use AdamW style decoupled weight decay\n d0 (float):\n Initial D estimate for D-adaptation (default 1e-6). Rarely needs changing.\n growth_rate (float):\n prevent the D estimate from growing faster than this multiplicative rate.\n Default is inf, for unrestricted. Values like 1.02 give a kind of learning\n rate warmup effect.\n fsdp_in_use (bool):\n If you're using sharded parameters, this should be set to True. The optimizer\n will attempt to auto-detect this, but if you're using an implementation other\n than PyTorch's builtin version, the auto-detection won't work.\n \"\"\"\n def __init__(self, params, lr=1.0,\n betas=(0.9, 0.999),\n eps=1e-8,\n weight_decay=0,\n log_every=0,\n decouple=True,\n d0=1e-6,\n growth_rate=float('inf')):\n if not 0.0 < d0:\n raise ValueError(\"Invalid d0 value: {}\".format(d0))\n if not 0.0 < lr:\n raise ValueError(\"Invalid learning rate: {}\".format(lr))\n if not 0.0 < eps:\n raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n if not 0.0 <= betas[0] < 1.0:\n raise ValueError(\"Invalid beta parameter at index 0: {}\".format(betas[0]))\n if not 0.0 <= betas[1] < 1.0:\n raise ValueError(\"Invalid beta parameter at index 1: {}\".format(betas[1]))\n\n if decouple:\n logger.info(\"Using decoupled weight decay\")\n\n from .fsdp import is_fsdp_used\n fsdp_in_use = is_fsdp_used()\n defaults = dict(lr=lr, betas=betas, eps=eps,\n weight_decay=weight_decay,\n d=d0,\n k=0,\n gsq_weighted=0.0,\n log_every=log_every,\n decouple=decouple,\n growth_rate=growth_rate,\n fsdp_in_use=fsdp_in_use)\n\n super().__init__(params, defaults)\n\n @property\n def supports_memory_efficient_fp16(self):\n return False\n\n @property\n def supports_flat_params(self):\n return True\n\n def step(self, closure=None):\n \"\"\"Performs a single optimization step.\n\n Args:\n closure (callable, optional): A closure that reevaluates the model\n and returns the loss.\n \"\"\"\n loss = None\n if closure is not None:\n loss = closure()\n\n g_sq = 0.0\n sksq_weighted = 0.0\n sk_l1 = 0.0\n\n lr = max(group['lr'] for group in self.param_groups)\n\n group = self.param_groups[0]\n gsq_weighted = group['gsq_weighted']\n d = group['d']\n dlr = d*lr\n\n growth_rate = group['growth_rate']\n decouple = group['decouple']\n fsdp_in_use = group['fsdp_in_use']\n log_every = group['log_every']\n\n beta1, beta2 = group['betas']\n\n for group in self.param_groups:\n group_lr = group['lr']\n decay = group['weight_decay']\n k = group['k']\n eps = group['eps']\n\n if group_lr not in [lr, 0.0]:\n raise RuntimeError(\"Setting different lr values in different parameter \"\n \"groups is only supported for values of 0\")\n\n for p in group['params']:\n if p.grad is None:\n continue\n if hasattr(p, \"_fsdp_flattened\"):\n fsdp_in_use = True\n grad = p.grad.data\n\n # Apply weight decay (coupled variant)\n if decay != 0 and not decouple:\n grad.add_(p.data, alpha=decay)\n\n state = self.state[p]\n\n # State initialization\n if 'step' not in state:\n state['step'] = 0\n state['s'] = torch.zeros_like(p.data, memory_format=torch.preserve_format).detach()\n # Exponential moving average of gradient values\n state['exp_avg'] = torch.zeros_like(p.data, memory_format=torch.preserve_format).detach()\n # Exponential moving average of squared gradient values\n state['exp_avg_sq'] = torch.zeros_like(\n to_real(p.data), memory_format=torch.preserve_format).detach()\n\n exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n\n grad_grad = to_real(grad * grad.conj())\n\n # Adam EMA updates\n if group_lr > 0:\n exp_avg.mul_(beta1).add_(grad, alpha=dlr*(1-beta1))\n exp_avg_sq.mul_(beta2).add_(grad_grad, alpha=1-beta2)\n\n denom = exp_avg_sq.sqrt().add_(eps)\n\n g_sq += grad_grad.div_(denom).sum().item()\n\n s = state['s']\n s.mul_(beta2).add_(grad, alpha=dlr*(1-beta2))\n sksq_weighted += to_real(s * s.conj()).div_(denom).sum().item()\n sk_l1 += s.abs().sum().item()\n\n ######\n\n gsq_weighted = beta2*gsq_weighted + g_sq*(dlr**2)*(1-beta2)\n d_hat = d\n\n # if we have not done any progres, return\n # if we have any gradients available, will have sk_l1 > 0 (unless \\|g\\|=0)\n if sk_l1 == 0:\n return loss\n\n if lr > 0.0:\n if fsdp_in_use:\n dist_tensor = torch.zeros(3, device='cuda')\n dist_tensor[0] = sksq_weighted\n dist_tensor[1] = gsq_weighted\n dist_tensor[2] = sk_l1\n dist.all_reduce(dist_tensor, op=dist.ReduceOp.SUM)\n global_sksq_weighted = dist_tensor[0]\n global_gsq_weighted = dist_tensor[1]\n global_sk_l1 = dist_tensor[2]\n else:\n global_sksq_weighted = sksq_weighted\n global_gsq_weighted = gsq_weighted\n global_sk_l1 = sk_l1\n\n d_hat = (global_sksq_weighted/(1-beta2) - global_gsq_weighted)/global_sk_l1\n d = max(d, min(d_hat, d*growth_rate))\n\n if log_every > 0 and k % log_every == 0:\n logger.info(\n f\"(k={k}) dlr: {dlr:1.1e} d_hat: {d_hat:1.1e}, d: {d:1.8}. \"\n f\"sksq_weighted={global_sksq_weighted:1.1e} gsq_weighted={global_gsq_weighted:1.1e} \"\n f\"sk_l1={global_sk_l1:1.1e}{' (FSDP)' if fsdp_in_use else ''}\")\n\n for group in self.param_groups:\n group['gsq_weighted'] = gsq_weighted\n group['d'] = d\n\n group_lr = group['lr']\n decay = group['weight_decay']\n k = group['k']\n eps = group['eps']\n\n for p in group['params']:\n if p.grad is None:\n continue\n grad = p.grad.data\n\n state = self.state[p]\n\n exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n\n state['step'] += 1\n\n denom = exp_avg_sq.sqrt().add_(eps)\n denom = denom.type(p.type())\n\n # Apply weight decay (decoupled variant)\n if decay != 0 and decouple and group_lr > 0:\n p.data.add_(p.data, alpha=-decay * dlr)\n\n # Take step\n p.data.addcdiv_(exp_avg, denom, value=-1)\n\n group['k'] = k + 1\n\n return loss\n" }, { "path": "audiocraft/optim/ema.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# ModelEMA implementation is taken from\n# https://github.com/facebookresearch/demucs\n\nfrom collections import defaultdict\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\n\n\ndef _get_all_non_persistent_buffers_set(module: nn.Module, root: str = \"\") -> set:\n names: set = set()\n for (name, sub_module) in module.named_modules():\n if name == '':\n buffer_names = module._non_persistent_buffers_set\n buffer_names = {f\"{root}.{buff_name}\" if len(root) > 0 else buff_name\n for buff_name in buffer_names}\n names.update(buffer_names)\n else:\n sub_name = f\"{root}.{name}\" if len(root) > 0 else name\n sub_buffer_names = _get_all_non_persistent_buffers_set(sub_module, sub_name)\n names.update(sub_buffer_names)\n return names\n\n\ndef _get_named_tensors(module: nn.Module):\n non_persistent_buffers_set = _get_all_non_persistent_buffers_set(module)\n named_buffers = [(name, buffer) for (name, buffer) in module.named_buffers()\n if name not in non_persistent_buffers_set]\n named_parameters = list(module.named_parameters())\n return named_parameters + named_buffers\n\n\nclass ModuleDictEMA:\n \"\"\"Exponential Moving Average over a nn.ModuleDict.\n\n You can switch to the EMA weights temporarily.\n \"\"\"\n def __init__(self, module_dict: nn.ModuleDict, decay: float = 0.999,\n unbias: bool = True, device: tp.Union[torch.device, str] = 'cpu'):\n self.decay = decay\n self.module_dict = module_dict\n self.state: dict = defaultdict(dict)\n self.count = 0\n self.device = device\n self.unbias = unbias\n self._init()\n\n def _init(self):\n for module_name, module in self.module_dict.items():\n for key, val in _get_named_tensors(module):\n if not val.is_floating_point():\n continue\n device = self.device or val.device\n if key not in self.state[module_name]:\n self.state[module_name][key] = val.detach().to(device, copy=True)\n\n def step(self):\n if self.unbias:\n self.count = self.count * self.decay + 1\n w = 1 / self.count\n else:\n w = 1 - self.decay\n for module_name, module in self.module_dict.items():\n for key, val in _get_named_tensors(module):\n if not val.is_floating_point():\n continue\n device = self.device or val.device\n self.state[module_name][key].mul_(1 - w)\n self.state[module_name][key].add_(val.detach().to(device), alpha=w)\n\n def state_dict(self):\n return {'state': self.state, 'count': self.count}\n\n def load_state_dict(self, state):\n self.count = state['count']\n for module_name, module in state['state'].items():\n for key, val in module.items():\n self.state[module_name][key].copy_(val)\n" }, { "path": "audiocraft/optim/fsdp.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nWrapper around FSDP for more convenient use in the training loops.\n\"\"\"\n\nfrom contextlib import contextmanager\nimport typing as tp\nimport dora\nimport torch\n\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp import (\n MixedPrecision, ShardingStrategy, FullStateDictConfig, StateDictType)\nfrom torch.distributed._shard.sharded_tensor.api import ShardedTensor\n\n\ndef is_fsdp_used() -> bool:\n \"\"\"Return whether we are using FSDP.\"\"\"\n # A bit of a hack but should work from anywhere.\n if dora.is_xp():\n cfg = dora.get_xp().cfg\n if hasattr(cfg, 'fsdp'):\n return cfg.fsdp.use\n return False\n\n\ndef is_sharded_tensor(x: tp.Any) -> bool:\n return isinstance(x, ShardedTensor)\n\n\n@contextmanager\ndef switch_to_full_state_dict(models: tp.List[FSDP]):\n # Another bug in FSDP makes it that we cannot use the `state_dict_type` API,\n # so let's do thing manually.\n for model in models:\n FSDP.set_state_dict_type( # type: ignore\n model, StateDictType.FULL_STATE_DICT,\n FullStateDictConfig(offload_to_cpu=True, rank0_only=True))\n try:\n yield\n finally:\n for model in models:\n FSDP.set_state_dict_type(model, StateDictType.LOCAL_STATE_DICT) # type: ignore\n\n\ndef wrap_with_fsdp(cfg, model: torch.nn.Module,\n block_classes: tp.Optional[tp.Set[tp.Type]] = None) -> FSDP:\n \"\"\"Wraps a model with FSDP.\"\"\"\n # Some of the typing is disabled until this gets integrated\n # into the stable version of PyTorch.\n from torch.distributed.fsdp.wrap import ModuleWrapPolicy # type: ignore\n\n # we import this here to prevent circular import.\n from ..modules.transformer import StreamingTransformerLayer\n from ..modules.conditioners import ConditioningProvider\n\n _fix_post_backward_hook()\n\n assert cfg.use\n sharding_strategy_dict = {\n \"no_shard\": ShardingStrategy.NO_SHARD,\n \"shard_grad_op\": ShardingStrategy.SHARD_GRAD_OP,\n \"full_shard\": ShardingStrategy.FULL_SHARD,\n }\n\n dtype_dict = {\n \"float32\": torch.float32,\n \"float16\": torch.float16,\n \"bfloat16\": torch.bfloat16,\n }\n\n mixed_precision_config = MixedPrecision(\n param_dtype=dtype_dict[cfg.param_dtype],\n reduce_dtype=dtype_dict[cfg.reduce_dtype],\n buffer_dtype=dtype_dict[cfg.buffer_dtype],\n )\n\n sharding_strategy_config = sharding_strategy_dict[cfg.sharding_strategy]\n # The following is going to require being a bit smart\n # when doing LM, because this would flush the weights for every time step\n # during generation. One possiblity is to use hybrid sharding:\n # See: https://pytorch.org/docs/master/fsdp.html#torch.distributed.fsdp.ShardingStrategy\n assert sharding_strategy_config != ShardingStrategy.FULL_SHARD, \\\n \"Not supported at the moment, requires a bit more work.\"\n\n local_rank = dora.distrib.get_distrib_spec().local_rank\n assert local_rank < torch.cuda.device_count(), \"Please upgrade Dora!\"\n\n auto_wrap_policy = None\n if block_classes is None:\n block_classes = {StreamingTransformerLayer, ConditioningProvider}\n if cfg.per_block:\n auto_wrap_policy = ModuleWrapPolicy(block_classes)\n wrapped = _FSDPFixStateDict(\n model,\n sharding_strategy=sharding_strategy_config,\n mixed_precision=mixed_precision_config,\n device_id=local_rank,\n sync_module_states=True,\n use_orig_params=True,\n auto_wrap_policy=auto_wrap_policy,\n ) # type: ignore\n FSDP.set_state_dict_type(wrapped, StateDictType.LOCAL_STATE_DICT) # type: ignore\n\n # Let the wrapped model know about the wrapping!\n # We use __dict__ to avoid it going into the state dict.\n # This is a bit dirty, but needed during generation, as otherwise\n # the wrapped model would call itself and bypass FSDP.\n for module in FSDP.fsdp_modules(wrapped):\n original = module._fsdp_wrapped_module\n original.__dict__['_fsdp'] = module\n return wrapped\n\n\ndef purge_fsdp(model: FSDP):\n \"\"\"Purge the FSDP cached shard inside the model. This should\n allow setting the best state or switching to the EMA.\n \"\"\"\n from torch.distributed.fsdp._runtime_utils import _reshard # type: ignore\n for module in FSDP.fsdp_modules(model):\n if hasattr(module, \"_handles\"):\n # support for FSDP with torch<2.1.0\n handles = module._handles\n if not handles:\n continue\n handle = handles[0]\n unsharded_flat_param = handle._get_padded_unsharded_flat_param()\n storage_size: int = unsharded_flat_param._typed_storage()._size() # type: ignore\n if storage_size == 0:\n continue\n true_list = [True for h in handles]\n _reshard(module, handles, true_list)\n else:\n handle = module._handle\n if not handle:\n continue\n unsharded_flat_param = handle._get_padded_unsharded_flat_param()\n storage_size: int = unsharded_flat_param._typed_storage()._size() # type: ignore\n if storage_size == 0:\n continue\n _reshard(module, handle, True)\n\n\nclass _FSDPFixStateDict(FSDP):\n @staticmethod\n def _name_without_fsdp_prefix(name: str) -> str:\n from torch.distributed.fsdp._common_utils import FSDP_WRAPPED_MODULE # type: ignore\n parts = name.split('.')\n new_parts = [part for part in parts if part != FSDP_WRAPPED_MODULE]\n return '.'.join(new_parts)\n\n def state_dict(self, *args, **kwargs) -> tp.Dict[str, tp.Any]: # type: ignore\n state = dict(super().state_dict(*args, **kwargs))\n for key, value in list(state.items()):\n if is_sharded_tensor(value):\n del state[key]\n return state\n\n def load_state_dict(self, state: tp.Dict[str, tp.Any]): # type: ignore\n if self._state_dict_type is StateDictType.FULL_STATE_DICT:\n super().load_state_dict(state)\n purge_fsdp(self)\n return\n # Fix FSDP load state dict in all situation.\n # Use this only with LOCAL_STATE_DICT !!!\n current_state = dict(super().state_dict())\n for key, value in state.items():\n key = _FSDPFixStateDict._name_without_fsdp_prefix(key)\n if key not in current_state:\n # Emulate strict loading manually.\n raise RuntimeError(f\"Unknown state key {key}\")\n current_state[key].copy_(value)\n\n # Purging cached weights from previous forward.\n purge_fsdp(self)\n\n\n_hook_fixed = False\n\n\ndef _fix_post_backward_hook():\n global _hook_fixed\n if _hook_fixed:\n return\n _hook_fixed = True\n\n from torch.distributed.fsdp import _runtime_utils\n from torch.distributed.fsdp._common_utils import TrainingState, HandleTrainingState\n old_hook = _runtime_utils._post_backward_hook\n\n def _post_backward_hook(state, handle, *args, **kwargs):\n checkpointed = getattr(state._fsdp_wrapped_module, '_audiocraft_checkpointed', False)\n if checkpointed:\n # there will be one more forward in the backward with checkpointing and that will\n # massively confuse FSDP, so we have to make it think everything\n # is going according to the plan.\n state.training_state = TrainingState.FORWARD_BACKWARD\n handle._training_state = HandleTrainingState.BACKWARD_PRE\n old_hook(state, handle, *args, **kwargs)\n\n _runtime_utils._post_backward_hook = _post_backward_hook\n" }, { "path": "audiocraft/optim/inverse_sqrt_lr_scheduler.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler\n\n\nclass InverseSquareRootLRScheduler(_LRScheduler):\n \"\"\"Inverse square root LR scheduler.\n\n Args:\n optimizer (Optimizer): Torch optimizer.\n warmup_steps (int): Number of warmup steps.\n warmup_init_lr (tp.Optional[float]): Initial learning rate\n during warmup phase. When not set, use the provided learning rate.\n \"\"\"\n def __init__(self, optimizer: Optimizer, warmup_steps: int, warmup_init_lr: tp.Optional[float] = 0):\n self.warmup_steps = warmup_steps\n self.warmup_init_lr = warmup_init_lr\n super().__init__(optimizer)\n\n def _get_sched_lr(self, lr: float, step: int):\n if step < self.warmup_steps:\n warmup_init_lr = self.warmup_init_lr or 0\n lr_step = (lr - warmup_init_lr) / self.warmup_steps\n lr = warmup_init_lr + step * lr_step\n else:\n decay_factor = lr * self.warmup_steps**0.5\n lr = decay_factor * step**-0.5\n return lr\n\n def get_lr(self):\n return [self._get_sched_lr(base_lr, self._step_count) for base_lr in self.base_lrs]\n" }, { "path": "audiocraft/optim/linear_warmup_lr_scheduler.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler\n\n\nclass LinearWarmupLRScheduler(_LRScheduler):\n \"\"\"Inverse square root LR scheduler.\n\n Args:\n optimizer (Optimizer): Torch optimizer.\n warmup_steps (int): Number of warmup steps.\n warmup_init_lr (tp.Optional[float]): Initial learning rate\n during warmup phase. When not set, use the provided learning rate.\n \"\"\"\n def __init__(self, optimizer: Optimizer, warmup_steps: int, warmup_init_lr: tp.Optional[float] = 0):\n self.warmup_steps = warmup_steps\n self.warmup_init_lr = warmup_init_lr\n super().__init__(optimizer)\n\n def _get_sched_lr(self, lr: float, step: int):\n if step < self.warmup_steps:\n warmup_init_lr = self.warmup_init_lr or 0\n lr_step = (lr - warmup_init_lr) / self.warmup_steps\n lr = warmup_init_lr + step * lr_step\n return lr\n\n def get_lr(self):\n return [self._get_sched_lr(base_lr, self.last_epoch) for base_lr in self.base_lrs]\n" }, { "path": "audiocraft/optim/polynomial_decay_lr_scheduler.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler\n\n\nclass PolynomialDecayLRScheduler(_LRScheduler):\n \"\"\"Polynomial decay LR scheduler.\n\n Args:\n optimizer (Optimizer): Torch optimizer.\n warmup_steps (int): Number of warmup steps.\n total_steps (int): Total number of steps.\n end_lr (float): Final learning rate to achieve over total number of steps.\n zero_lr_warmup_steps (int): Number of steps with a learning rate of value 0.\n power (float): Decay exponent.\n \"\"\"\n def __init__(self, optimizer: Optimizer, warmup_steps: int, total_steps: int,\n end_lr: float = 0., zero_lr_warmup_steps: int = 0, power: float = 1.):\n self.warmup_steps = warmup_steps\n self.total_steps = total_steps\n self.end_lr = end_lr\n self.zero_lr_warmup_steps = zero_lr_warmup_steps\n self.power = power\n super().__init__(optimizer)\n\n def _get_sched_lr(self, lr: float, step: int):\n if self.zero_lr_warmup_steps > 0 and step <= self.zero_lr_warmup_steps:\n lr = 0\n elif self.warmup_steps > 0 and step <= self.warmup_steps + self.zero_lr_warmup_steps:\n lr_ratio = (step - self.zero_lr_warmup_steps) / float(self.warmup_steps)\n lr = lr_ratio * lr\n elif step >= self.total_steps:\n lr = self.end_lr\n else:\n total_warmup_steps = self.warmup_steps + self.zero_lr_warmup_steps\n lr_range = lr - self.end_lr\n pct_remaining = 1 - (step - total_warmup_steps) / (self.total_steps - total_warmup_steps)\n lr = lr_range * pct_remaining ** self.power + self.end_lr\n return lr\n\n def get_lr(self):\n return [self._get_sched_lr(base_lr, self.last_epoch) for base_lr in self.base_lrs]\n" }, { "path": "audiocraft/py.typed", "content": "" }, { "path": "audiocraft/quantization/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"RVQ.\"\"\"\n# flake8: noqa\nfrom .vq import ResidualVectorQuantizer\nfrom .base import BaseQuantizer, DummyQuantizer, QuantizedResult\n" }, { "path": "audiocraft/quantization/base.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nBase class for all quantizers.\n\"\"\"\n\nfrom dataclasses import dataclass, field\nimport typing as tp\n\nimport torch\nfrom torch import nn\n\n\n@dataclass\nclass QuantizedResult:\n x: torch.Tensor\n codes: torch.Tensor\n bandwidth: torch.Tensor # bandwidth in kb/s used, per batch item.\n penalty: tp.Optional[torch.Tensor] = None\n metrics: dict = field(default_factory=dict)\n\n\nclass BaseQuantizer(nn.Module):\n \"\"\"Base class for quantizers.\n \"\"\"\n\n def forward(self, x: torch.Tensor, frame_rate: int) -> QuantizedResult:\n \"\"\"\n Given input tensor x, returns first the quantized (or approximately quantized)\n representation along with quantized codes, bandwidth, and any penalty term for the loss.\n Finally, this returns a dict of metrics to update logging etc.\n Frame rate must be passed so that the bandwidth is properly computed.\n \"\"\"\n raise NotImplementedError()\n\n def encode(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"Encode a given input tensor with the specified sample rate at the given bandwidth.\"\"\"\n raise NotImplementedError()\n\n def decode(self, codes: torch.Tensor) -> torch.Tensor:\n \"\"\"Decode the given codes to the quantized representation.\"\"\"\n raise NotImplementedError()\n\n @property\n def total_codebooks(self):\n \"\"\"Total number of codebooks.\"\"\"\n raise NotImplementedError()\n\n @property\n def num_codebooks(self):\n \"\"\"Number of active codebooks.\"\"\"\n raise NotImplementedError()\n\n def set_num_codebooks(self, n: int):\n \"\"\"Set the number of active codebooks.\"\"\"\n raise NotImplementedError()\n\n\nclass DummyQuantizer(BaseQuantizer):\n \"\"\"Fake quantizer that actually does not perform any quantization.\n \"\"\"\n def __init__(self):\n super().__init__()\n\n def forward(self, x: torch.Tensor, frame_rate: int):\n q = x.unsqueeze(1)\n return QuantizedResult(x, q, torch.tensor(q.numel() * 32 * frame_rate / 1000 / len(x)).to(x))\n\n def encode(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"Encode a given input tensor with the specified sample rate at the given bandwidth.\n In the case of the DummyQuantizer, the codes are actually identical\n to the input and resulting quantized representation as no quantization is done.\n \"\"\"\n return x.unsqueeze(1)\n\n def decode(self, codes: torch.Tensor) -> torch.Tensor:\n \"\"\"Decode the given codes to the quantized representation.\n In the case of the DummyQuantizer, the codes are actually identical\n to the input and resulting quantized representation as no quantization is done.\n \"\"\"\n return codes.squeeze(1)\n\n @property\n def total_codebooks(self):\n \"\"\"Total number of codebooks.\"\"\"\n return 1\n\n @property\n def num_codebooks(self):\n \"\"\"Total number of codebooks.\"\"\"\n return self.total_codebooks\n\n def set_num_codebooks(self, n: int):\n \"\"\"Set the number of active codebooks.\"\"\"\n raise AttributeError(\"Cannot override the number of codebooks for the dummy quantizer\")\n" }, { "path": "audiocraft/quantization/core_vq.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nfrom einops import rearrange, repeat\nimport flashy\nimport torch\nfrom torch import nn, einsum\nimport torch.nn.functional as F\n\n\ndef exists(val: tp.Optional[tp.Any]) -> bool:\n return val is not None\n\n\ndef default(val: tp.Any, d: tp.Any) -> tp.Any:\n return val if exists(val) else d\n\n\ndef l2norm(t):\n return F.normalize(t, p=2, dim=-1)\n\n\ndef ema_inplace(moving_avg, new, decay: float):\n moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))\n\n\ndef laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):\n return (x + epsilon) / (x.sum() + n_categories * epsilon)\n\n\ndef uniform_init(*shape: int):\n t = torch.empty(shape)\n nn.init.kaiming_uniform_(t)\n return t\n\n\ndef sample_vectors(samples, num: int):\n num_samples, device = samples.shape[0], samples.device\n\n if num_samples >= num:\n indices = torch.randperm(num_samples, device=device)[:num]\n else:\n indices = torch.randint(0, num_samples, (num,), device=device)\n\n return samples[indices]\n\n\ndef kmeans(samples, num_clusters: int, num_iters: int = 10):\n dim, dtype = samples.shape[-1], samples.dtype\n\n means = sample_vectors(samples, num_clusters)\n\n for _ in range(num_iters):\n diffs = rearrange(samples, \"n d -> n () d\") - rearrange(\n means, \"c d -> () c d\"\n )\n dists = -(diffs ** 2).sum(dim=-1)\n\n buckets = dists.max(dim=-1).indices\n bins = torch.bincount(buckets, minlength=num_clusters)\n zero_mask = bins == 0\n bins_min_clamped = bins.masked_fill(zero_mask, 1)\n\n new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)\n new_means.scatter_add_(0, repeat(buckets, \"n -> n d\", d=dim), samples)\n new_means = new_means / bins_min_clamped[..., None]\n\n means = torch.where(zero_mask[..., None], means, new_means)\n\n return means, bins\n\n\ndef orthogonal_loss_fn(t):\n # eq (2) from https://arxiv.org/abs/2112.00384\n n = t.shape[0]\n normed_codes = l2norm(t)\n identity = torch.eye(n, device=t.device)\n cosine_sim = einsum(\"i d, j d -> i j\", normed_codes, normed_codes)\n return ((cosine_sim - identity) ** 2).sum() / (n ** 2)\n\n\nclass EuclideanCodebook(nn.Module):\n \"\"\"Codebook with Euclidean distance.\n\n Args:\n dim (int): Dimension.\n codebook_size (int): Codebook size.\n kmeans_init (bool): Whether to use k-means to initialize the codebooks.\n If set to true, run the k-means algorithm on the first training batch and use\n the learned centroids as initialization.\n kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.\n decay (float): Decay for exponential moving average over the codebooks.\n epsilon (float): Epsilon value for numerical stability.\n threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes\n that have an exponential moving average cluster size less than the specified threshold with\n randomly selected vector from the current batch.\n \"\"\"\n def __init__(\n self,\n dim: int,\n codebook_size: int,\n kmeans_init: int = False,\n kmeans_iters: int = 10,\n decay: float = 0.8,\n epsilon: float = 1e-5,\n threshold_ema_dead_code: float = 2.,\n ):\n super().__init__()\n self.decay = decay\n init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = uniform_init if not kmeans_init else torch.zeros\n embed = init_fn(codebook_size, dim)\n\n self.codebook_size = codebook_size\n\n self.kmeans_iters = kmeans_iters\n self.epsilon = epsilon\n self.threshold_ema_dead_code = threshold_ema_dead_code\n\n self.register_buffer(\"inited\", torch.Tensor([not kmeans_init]))\n self.register_buffer(\"cluster_size\", torch.zeros(codebook_size))\n self.register_buffer(\"embed\", embed)\n self.register_buffer(\"embed_avg\", embed.clone())\n\n @torch.jit.ignore\n def init_embed_(self, data):\n if self.inited:\n return\n\n embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)\n self.embed.data.copy_(embed)\n self.embed_avg.data.copy_(embed.clone())\n self.cluster_size.data.copy_(cluster_size)\n self.inited.data.copy_(torch.Tensor([True]))\n # Make sure all buffers across workers are in sync after initialization\n flashy.distrib.broadcast_tensors(self.buffers())\n\n def replace_(self, samples, mask):\n modified_codebook = torch.where(\n mask[..., None], sample_vectors(samples, self.codebook_size), self.embed\n )\n self.embed.data.copy_(modified_codebook)\n\n def expire_codes_(self, batch_samples):\n if self.threshold_ema_dead_code == 0:\n return\n\n expired_codes = self.cluster_size < self.threshold_ema_dead_code\n if not torch.any(expired_codes):\n return\n\n batch_samples = rearrange(batch_samples, \"... d -> (...) d\")\n self.replace_(batch_samples, mask=expired_codes)\n flashy.distrib.broadcast_tensors(self.buffers())\n\n def preprocess(self, x):\n x = rearrange(x, \"... d -> (...) d\")\n return x\n\n def quantize(self, x):\n embed = self.embed.t()\n dist = -(\n x.pow(2).sum(1, keepdim=True)\n - 2 * x @ embed\n + embed.pow(2).sum(0, keepdim=True)\n )\n embed_ind = dist.max(dim=-1).indices\n return embed_ind\n\n def postprocess_emb(self, embed_ind, shape):\n return embed_ind.view(*shape[:-1])\n\n def dequantize(self, embed_ind):\n quantize = F.embedding(embed_ind, self.embed)\n return quantize\n\n def encode(self, x):\n shape = x.shape\n # pre-process\n x = self.preprocess(x)\n # quantize\n embed_ind = self.quantize(x)\n # post-process\n embed_ind = self.postprocess_emb(embed_ind, shape)\n return embed_ind\n\n def decode(self, embed_ind):\n quantize = self.dequantize(embed_ind)\n return quantize\n\n def forward(self, x):\n shape, dtype = x.shape, x.dtype\n x = self.preprocess(x)\n self.init_embed_(x)\n\n embed_ind = self.quantize(x)\n embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)\n embed_ind = self.postprocess_emb(embed_ind, shape)\n quantize = self.dequantize(embed_ind)\n\n if self.training:\n # We do the expiry of code at that point as buffers are in sync\n # and all the workers will take the same decision.\n self.expire_codes_(x)\n ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)\n embed_sum = x.t() @ embed_onehot\n ema_inplace(self.embed_avg, embed_sum.t(), self.decay)\n cluster_size = (\n laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon)\n * self.cluster_size.sum()\n )\n embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)\n self.embed.data.copy_(embed_normalized)\n\n return quantize, embed_ind\n\n\nclass VectorQuantization(nn.Module):\n \"\"\"Vector quantization implementation.\n Currently supports only euclidean distance.\n\n Args:\n dim (int): Dimension\n codebook_size (int): Codebook size\n codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.\n decay (float): Decay for exponential moving average over the codebooks.\n epsilon (float): Epsilon value for numerical stability.\n kmeans_init (bool): Whether to use kmeans to initialize the codebooks.\n kmeans_iters (int): Number of iterations used for kmeans initialization.\n channels_last (bool): Channels are the last dimension in the input tensors.\n commitment_weight (float): Weight for commitment loss.\n orthogonal_reg_weight (float): Orthogonal regularization weights.\n orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.\n orthogonal_reg_max_codes (optional int): Maximum number of codes to consider\n for orthogonal regularization.\n threshold_ema_dead_code (float): Threshold for dead code expiration. Replace any codes\n that have an exponential moving average cluster size less than the specified threshold with\n randomly selected vector from the current batch.\n \"\"\"\n def __init__(\n self,\n dim: int,\n codebook_size: int,\n codebook_dim: tp.Optional[int] = None,\n decay: float = 0.8,\n epsilon: float = 1e-5,\n kmeans_init: bool = False,\n kmeans_iters: int = 10,\n threshold_ema_dead_code: float = 2.,\n channels_last: bool = False,\n commitment_weight: float = 1.,\n orthogonal_reg_weight: float = 0.0,\n orthogonal_reg_active_codes_only: bool = False,\n orthogonal_reg_max_codes: tp.Optional[int] = None,\n ):\n super().__init__()\n _codebook_dim: int = default(codebook_dim, dim)\n\n requires_projection = _codebook_dim != dim\n self.project_in = (nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity())\n self.project_out = (nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity())\n\n self.epsilon = epsilon\n self.commitment_weight = commitment_weight\n\n self.orthogonal_reg_weight = orthogonal_reg_weight\n self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only\n self.orthogonal_reg_max_codes = orthogonal_reg_max_codes\n\n self._codebook = EuclideanCodebook(dim=_codebook_dim, codebook_size=codebook_size,\n kmeans_init=kmeans_init, kmeans_iters=kmeans_iters,\n decay=decay, epsilon=epsilon,\n threshold_ema_dead_code=threshold_ema_dead_code)\n self.codebook_size = codebook_size\n\n self.channels_last = channels_last\n\n @property\n def codebook(self):\n return self._codebook.embed\n\n @property\n def inited(self):\n return self._codebook.inited\n\n def _preprocess(self, x):\n if not self.channels_last:\n x = rearrange(x, \"b d n -> b n d\")\n return x\n\n def _postprocess(self, quantize):\n if not self.channels_last:\n quantize = rearrange(quantize, \"b n d -> b d n\")\n return quantize\n\n def encode(self, x):\n x = self._preprocess(x)\n x = self.project_in(x)\n embed_in = self._codebook.encode(x)\n return embed_in\n\n def decode(self, embed_ind):\n quantize = self._codebook.decode(embed_ind)\n quantize = self.project_out(quantize)\n quantize = self._postprocess(quantize)\n return quantize\n\n def forward(self, x):\n device = x.device\n x = self._preprocess(x)\n\n x = self.project_in(x)\n quantize, embed_ind = self._codebook(x)\n\n if self.training:\n quantize = x + (quantize - x).detach()\n\n loss = torch.tensor([0.0], device=device, requires_grad=self.training)\n\n if self.training:\n if self.commitment_weight > 0:\n commit_loss = F.mse_loss(quantize.detach(), x)\n loss = loss + commit_loss * self.commitment_weight\n\n if self.orthogonal_reg_weight > 0:\n codebook = self.codebook\n\n if self.orthogonal_reg_active_codes_only:\n # only calculate orthogonal loss for the activated codes for this batch\n unique_code_ids = torch.unique(embed_ind)\n codebook = codebook[unique_code_ids]\n\n num_codes = codebook.shape[0]\n if exists(self.orthogonal_reg_max_codes) and num_codes > self.orthogonal_reg_max_codes:\n rand_ids = torch.randperm(num_codes, device=device)[:self.orthogonal_reg_max_codes]\n codebook = codebook[rand_ids]\n\n orthogonal_reg_loss = orthogonal_loss_fn(codebook)\n loss = loss + orthogonal_reg_loss * self.orthogonal_reg_weight\n\n quantize = self.project_out(quantize)\n quantize = self._postprocess(quantize)\n\n return quantize, embed_ind, loss\n\n\nclass ResidualVectorQuantization(nn.Module):\n \"\"\"Residual vector quantization implementation.\n\n Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf\n \"\"\"\n def __init__(self, *, num_quantizers, **kwargs):\n super().__init__()\n self.layers = nn.ModuleList(\n [VectorQuantization(**kwargs) for _ in range(num_quantizers)]\n )\n\n def forward(self, x, n_q: tp.Optional[int] = None):\n quantized_out = 0.0\n residual = x\n\n all_losses = []\n all_indices = []\n\n n_q = n_q or len(self.layers)\n\n for i, layer in enumerate(self.layers[:n_q]):\n quantized, indices, loss = layer(residual)\n quantized = quantized.detach()\n residual = residual - quantized\n quantized_out = quantized_out + quantized\n all_indices.append(indices)\n all_losses.append(loss)\n\n if self.training:\n # Solving subtle bug with STE and RVQ: https://github.com/facebookresearch/encodec/issues/25\n quantized_out = x + (quantized_out - x).detach()\n\n out_losses, out_indices = map(torch.stack, (all_losses, all_indices))\n return quantized_out, out_indices, out_losses\n\n def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor:\n residual = x\n all_indices = []\n n_q = n_q or len(self.layers)\n for layer in self.layers[:n_q]:\n indices = layer.encode(residual)\n quantized = layer.decode(indices)\n residual = residual - quantized\n all_indices.append(indices)\n out_indices = torch.stack(all_indices)\n return out_indices\n\n def decode(self, q_indices: torch.Tensor) -> torch.Tensor:\n quantized_out = torch.tensor(0.0, device=q_indices.device)\n for i, indices in enumerate(q_indices):\n layer = self.layers[i]\n quantized = layer.decode(indices)\n quantized_out = quantized_out + quantized\n return quantized_out\n" }, { "path": "audiocraft/quantization/vq.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nimport typing as tp\n\nimport torch\n\nfrom .base import BaseQuantizer, QuantizedResult\nfrom .core_vq import ResidualVectorQuantization\n\n\nclass ResidualVectorQuantizer(BaseQuantizer):\n \"\"\"Residual Vector Quantizer.\n\n Args:\n dimension (int): Dimension of the codebooks.\n n_q (int): Number of residual vector quantizers used.\n q_dropout (bool): Random quantizer drop out at train time.\n bins (int): Codebook size.\n decay (float): Decay for exponential moving average over the codebooks.\n kmeans_init (bool): Whether to use kmeans to initialize the codebooks.\n kmeans_iters (int): Number of iterations used for kmeans initialization.\n threshold_ema_dead_code (float): Threshold for dead code expiration. Replace any codes\n that have an exponential moving average cluster size less than the specified threshold with\n randomly selected vector from the current batch.\n orthogonal_reg_weight (float): Orthogonal regularization weights.\n orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.\n orthogonal_reg_max_codes (optional int): Maximum number of codes to consider.\n for orthogonal regularization.\n \"\"\"\n def __init__(\n self,\n dimension: int = 256,\n n_q: int = 8,\n q_dropout: bool = False,\n bins: int = 1024,\n decay: float = 0.99,\n kmeans_init: bool = True,\n kmeans_iters: int = 10,\n threshold_ema_dead_code: float = 2.,\n orthogonal_reg_weight: float = 0.0,\n orthogonal_reg_active_codes_only: bool = False,\n orthogonal_reg_max_codes: tp.Optional[int] = None,\n ):\n super().__init__()\n self.max_n_q = n_q\n self.n_q = n_q\n self.q_dropout = q_dropout\n self.dimension = dimension\n self.bins = bins\n self.decay = decay\n self.kmeans_init = kmeans_init\n self.kmeans_iters = kmeans_iters\n self.threshold_ema_dead_code = threshold_ema_dead_code\n self.orthogonal_reg_weight = orthogonal_reg_weight\n self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only\n self.orthogonal_reg_max_codes = orthogonal_reg_max_codes\n self.vq = ResidualVectorQuantization(\n dim=self.dimension,\n codebook_size=self.bins,\n num_quantizers=self.n_q,\n decay=self.decay,\n kmeans_init=self.kmeans_init,\n kmeans_iters=self.kmeans_iters,\n threshold_ema_dead_code=self.threshold_ema_dead_code,\n orthogonal_reg_weight=self.orthogonal_reg_weight,\n orthogonal_reg_active_codes_only=self.orthogonal_reg_active_codes_only,\n orthogonal_reg_max_codes=self.orthogonal_reg_max_codes,\n channels_last=False\n )\n\n def forward(self, x: torch.Tensor, frame_rate: int):\n n_q = self.n_q\n if self.training and self.q_dropout:\n n_q = int(torch.randint(1, self.n_q + 1, (1,)).item())\n bw_per_q = math.log2(self.bins) * frame_rate / 1000\n quantized, codes, commit_loss = self.vq(x, n_q=n_q)\n codes = codes.transpose(0, 1)\n # codes is [B, K, T], with T frames, K nb of codebooks.\n bw = torch.tensor(n_q * bw_per_q).to(x)\n return QuantizedResult(quantized, codes, bw, penalty=torch.mean(commit_loss))\n\n def encode(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"Encode a given input tensor with the specified frame rate at the given bandwidth.\n The RVQ encode method sets the appropriate number of quantizer to use\n and returns indices for each quantizer.\n \"\"\"\n n_q = self.n_q\n codes = self.vq.encode(x, n_q=n_q)\n codes = codes.transpose(0, 1)\n # codes is [B, K, T], with T frames, K nb of codebooks.\n return codes\n\n def decode(self, codes: torch.Tensor) -> torch.Tensor:\n \"\"\"Decode the given codes to the quantized representation.\"\"\"\n # codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T].\n codes = codes.transpose(0, 1)\n quantized = self.vq.decode(codes)\n return quantized\n\n @property\n def total_codebooks(self):\n return self.max_n_q\n\n @property\n def num_codebooks(self):\n return self.n_q\n\n def set_num_codebooks(self, n: int):\n assert n > 0 and n <= self.max_n_q\n self.n_q = n\n" }, { "path": "audiocraft/solvers/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"\nSolvers. A Solver is a training recipe, combining the dataloaders, models,\noptimizer, losses etc into a single convenient object.\n\"\"\"\n\n# flake8: noqa\nfrom .audiogen import AudioGenSolver\nfrom .builders import get_solver\nfrom .base import StandardSolver\nfrom .compression import CompressionSolver\nfrom .musicgen import MusicGenSolver\nfrom .diffusion import DiffusionSolver\n" }, { "path": "audiocraft/solvers/audiogen.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom . import builders, musicgen\n\n\nclass AudioGenSolver(musicgen.MusicGenSolver):\n \"\"\"Solver for AudioGen re-implementation training task.\n\n Note that this implementation does not strictly follows\n the method proposed in https://arxiv.org/abs/2209.15352\n but is derived from MusicGen's training pipeline.\n\n More information can be found in the AudioGen model card.\n \"\"\"\n DATASET_TYPE: builders.DatasetType = builders.DatasetType.SOUND\n" }, { "path": "audiocraft/solvers/base.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom abc import ABC, abstractmethod\nfrom contextlib import contextmanager\nfrom pathlib import Path\nimport typing as tp\n\nimport flashy\nimport omegaconf\nimport torch\nfrom torch import nn\n\nfrom .. import optim\nfrom ..optim import fsdp\nfrom ..utils import checkpoint\nfrom ..utils.autocast import TorchAutocast\nfrom ..utils.best_state import BestStateDictManager\nfrom ..utils.deadlock import DeadlockDetect\nfrom ..utils.profiler import Profiler\nfrom ..utils.utils import copy_state, dict_from_config, model_hash, with_rank_rng\n\n\nclass StandardSolver(ABC, flashy.BaseSolver):\n \"\"\"Standard solver for AudioCraft.\n\n The standard solver implements a base training loop with the following stages:\n train, valid, evaluate and generate that are expected to be all defined for\n solvers in AudioCraft. It also provides a nice default management of Dora history replay,\n checkpoint management across epoch, and logging configuration.\n\n AudioCraft solvers must inherit from the StandardSolver and define the methods\n associated to each stage as well as the show, build_model and build_dataloaders methods.\n \"\"\"\n def __init__(self, cfg: omegaconf.DictConfig):\n super().__init__()\n self.logger.info(f\"Instantiating solver {self.__class__.__name__} for XP {self.xp.sig}\")\n self.logger.info(f\"All XP logs are stored in {self.xp.folder}\")\n self.cfg = cfg\n self.device = cfg.device\n self.model: nn.Module\n self._continue_best_source_keys = ['best_state', 'fsdp_best_state']\n self._fsdp_modules: tp.List[fsdp.FSDP] = []\n self._ema_sources: nn.ModuleDict = nn.ModuleDict()\n self.ema: tp.Optional[optim.ModuleDictEMA] = None\n self.dataloaders: tp.Dict[str, torch.utils.data.DataLoader] = dict()\n self._log_updates = self.cfg.logging.get('log_updates', 10)\n if self.cfg.logging.log_tensorboard:\n self.init_tensorboard(**self.cfg.get('tensorboard'))\n if self.cfg.logging.log_wandb and self:\n self.init_wandb(**self.cfg.get('wandb'))\n # keep a copy of the best performing state for stateful objects\n # used for evaluation and generation stages\n dtype_best: tp.Optional[torch.dtype] = None\n if self.cfg.fsdp.use:\n dtype_best = getattr(torch, self.cfg.fsdp.param_dtype) # type: ignore\n assert isinstance(dtype_best, torch.dtype)\n elif self.cfg.autocast:\n dtype_best = getattr(torch, self.cfg.autocast_dtype) # type: ignore\n assert isinstance(dtype_best, torch.dtype)\n self.best_state: BestStateDictManager = BestStateDictManager(dtype=dtype_best)\n # Hacky support for keeping a copy of the full best state in rank0.\n self.fsdp_best_state: tp.Dict[str, tp.Any] = {}\n self.register_stateful('best_state', 'fsdp_best_state') # register best_state object to keep it in state_dict\n self._new_best_state: bool = False # should save a new checkpoint\n # instantiate datasets and appropriate number of updates per epoch\n self.build_dataloaders()\n if self.cfg.execute_only is None:\n assert 'train' in self.dataloaders, \"The train dataset split must be provided.\"\n assert 'valid' in self.dataloaders, \"The valid dataset split must be provided.\"\n self.train_updates_per_epoch = len(self.dataloaders['train']) if 'train' in self.dataloaders else 0\n if self.cfg.optim.updates_per_epoch:\n self.train_updates_per_epoch = self.cfg.optim.updates_per_epoch\n self.total_updates = self.train_updates_per_epoch * self.cfg.optim.epochs\n # instantiate model & exponential moving average on the model\n self.build_model()\n self.logger.info(\"Model hash: %s\", model_hash(self.model))\n assert 'model' in self.stateful.sources, \\\n \"Please register the model to stateful with self.register_stateful('model') in build_model.\"\n self.profiler = Profiler(self.model, **self.cfg.profiler)\n self.initialize_ema()\n self.register_stateful('ema')\n assert self.ema is None or 'ema' in self.stateful.sources, \\\n \"Please register the ema to stateful with self.register_stateful('ema') in build_model.\"\n self.deadlock_detect = DeadlockDetect(**self.cfg.deadlock)\n # basic statistics on the trained model\n model_size = sum(p.numel() for p in self.model.parameters() if p.requires_grad) / 1e6\n # one copy of grad, one copy of momentum, one copy of denominator and model weights.\n # and 4 bytes for each float!\n mem_usage = model_size * 4 * 4 / 1000\n self.logger.info(\"Model size: %.2f M params\", model_size)\n self.logger.info(\"Base memory usage, with model, grad and optim: %.2f GB\", mem_usage)\n\n @property\n def autocast(self):\n \"\"\"Convenient autocast (or not) using the solver configuration.\"\"\"\n return TorchAutocast(enabled=self.cfg.autocast, device_type=self.device, dtype=self.autocast_dtype)\n\n def _get_state_source(self, name) -> flashy.state.StateDictSource:\n # Internal utility to get a state source from the solver\n return self.stateful.sources[name]\n\n @property\n def best_metric_name(self) -> tp.Optional[str]:\n \"\"\"Metric name used to identify the best state. This metric should be stored in the metrics\n used on the stage for best state identification (most likely, `valid`). If None, then\n no best state is saved.\n \"\"\"\n return None\n\n def register_best_state(self, *args: str):\n \"\"\"Register state sources in `BestStateDictManager` to keep their best states along with their\n latest states. The best state will be used at evaluation stages instead of the latest states.\n\n Shortcut around `BestStateDictManager.register` method. You can pass any number of\n attribute, included nested attributes and those will be included into the checkpoints\n and automatically restored when `BaseSolver.restore` is called.\n \"\"\"\n for name in args:\n state_source = self._get_state_source(name)\n assert name in self.stateful.sources, \"Registered states in best should be registered in stateful first!\"\n self.best_state.register(name, state_source)\n\n def register_ema(self, *args: str):\n \"\"\"Register state sources for exponential moving average.\n\n The registered sources are used to instantiate a ModuleDictEMA instance.\n The ModuleDictEMA keeps a `nn.ModuleDict` module that is updated when self.ema.step() is called\n and swapped with the original state sources with self.swap_ema_state() method.\n\n Usage:\n self.register_ema('model')\n \"\"\"\n assert self.ema is None, \"Cannot register state source to already instantiated EMA.\"\n for name in args:\n self._ema_sources[name] = getattr(self, name)\n\n def wrap_with_fsdp(self, model: torch.nn.Module, *args, **kwargs):\n model = fsdp.wrap_with_fsdp(self.cfg.fsdp, model, *args, **kwargs)\n if isinstance(model, fsdp.FSDP):\n self._fsdp_modules.append(model)\n return model\n\n def update_best_state_from_stage(self, stage_name: str = 'valid'):\n \"\"\"Update latest best state based on pending metrics of a given stage. This method relies\n on the `BestStateDictManager.update` method to update the best state_dict with latest weights\n if the registered states happen to match to the best performing setup.\n \"\"\"\n if self.best_metric_name is None:\n # when no best metric is defined, the last state is always the best\n self._new_best_state = True\n self.logger.info(\"Updating best state with current state.\")\n else:\n assert stage_name in self._pending_metrics, f\"Metrics for stage {stage_name} not found.\"\n assert self.best_metric_name in self._pending_metrics[stage_name], \\\n f\"Best metric not found in {stage_name} metrics. Cannot register best state\"\n current_score = self._pending_metrics[stage_name][self.best_metric_name]\n all_best_metric_scores = [\n past_metrics[stage_name][self.best_metric_name]\n for past_metrics in self.history\n ]\n all_best_metric_scores.append(current_score)\n best_score = min(all_best_metric_scores)\n self._new_best_state = current_score == best_score\n if self._new_best_state:\n old_best = min(all_best_metric_scores[:-1] + [float('inf')])\n self.logger.info(\n f\"New best state with {self.best_metric_name}={current_score:.3f} (was {old_best:.3f})\")\n\n if self._new_best_state:\n if self.cfg.fsdp.use:\n # this will give an empty state dict on all ranks but the rank 0\n # which will have a copy in memory of the full model.\n with fsdp.switch_to_full_state_dict(self._fsdp_modules):\n for name in self.best_state.states.keys():\n state_source = self._get_state_source(name)\n self.best_state.update(name, state_source)\n # we save to a different dict.\n self.fsdp_best_state.update(self.best_state.state_dict())\n # We cannot efficiently load fsdp_best_state when using FSDP,\n # so we have do do a second pass, with the local shards.\n for name in self.best_state.states.keys():\n state_source = self._get_state_source(name)\n self.best_state.update(name, state_source)\n\n def _load_new_state_dict(self, state_dict: dict) -> dict:\n old_states = {}\n for name, new_state in state_dict.items():\n state_source = self._get_state_source(name)\n old_states[name] = copy_state(state_source.state_dict())\n state_source.load_state_dict(new_state)\n return old_states\n\n @contextmanager\n def swap_best_state(self):\n self.logger.debug(f\"Swapping to best state for: {', '.join(self.best_state.state_dict().keys())}\")\n old_states = self._load_new_state_dict(self.best_state.state_dict())\n try:\n yield\n finally:\n self.logger.debug(\"Swapping back from best to original state\")\n for name, old_state in old_states.items():\n state_source = self._get_state_source(name)\n state_source.load_state_dict(old_state)\n\n @contextmanager\n def swap_ema_state(self):\n if self.ema is None:\n yield\n else:\n ema_state_dict = self.ema.state_dict()['state']\n self.logger.debug(f\"Swapping to EMA state for: {', '.join(ema_state_dict.keys())}\")\n old_states = self._load_new_state_dict(ema_state_dict)\n try:\n yield\n finally:\n self.logger.debug(\"Swapping back from EMA state to original state\")\n for name, old_state in old_states.items():\n state_source = self._get_state_source(name)\n state_source.load_state_dict(old_state)\n\n @property\n def is_training(self):\n return self.current_stage == 'train'\n\n def log_model_summary(self, model: nn.Module):\n \"\"\"Log model summary, architecture and size of the model.\"\"\"\n self.logger.info(model)\n mb = sum(p.numel() for p in model.parameters()) * 4 / 2 ** 20\n self.logger.info(\"Size: %.1f MB\", mb)\n\n @abstractmethod\n def build_model(self):\n \"\"\"Method to implement to initialize model.\"\"\"\n ...\n\n def initialize_ema(self):\n \"\"\"Initialize exponential moving average with the registered sources.\n EMA object is created if the optim.ema.model.decay value is non-null.\n \"\"\"\n from .builders import get_ema\n self.ema = get_ema(self._ema_sources, self.cfg.optim.ema)\n if self.ema is None:\n self.logger.info('No EMA on the model.')\n else:\n assert self.cfg.optim.ema.updates > 0\n self.logger.info(\n f'Initializing EMA on the model with decay = {self.ema.decay}'\n f' every {self.cfg.optim.ema.updates} updates'\n )\n\n @abstractmethod\n def build_dataloaders(self):\n \"\"\"Method to implement to initialize dataloaders.\"\"\"\n ...\n\n @abstractmethod\n def show(self):\n \"\"\"Method to log any information without running the job.\"\"\"\n ...\n\n @property\n def log_updates(self):\n # convenient access to log updates\n return self._log_updates\n\n def checkpoint_path(self, **kwargs):\n kwargs.setdefault('use_fsdp', self.cfg.fsdp.use)\n return self.folder / checkpoint.checkpoint_name(**kwargs)\n\n def epoch_checkpoint_path(self, epoch: int, **kwargs):\n kwargs.setdefault('use_fsdp', self.cfg.fsdp.use)\n return self.folder / checkpoint.checkpoint_name(str(epoch), **kwargs)\n\n def checkpoint_path_with_name(self, name: str, **kwargs):\n kwargs.setdefault('use_fsdp', self.cfg.fsdp.use)\n return self.folder / checkpoint.checkpoint_name(name=name, **kwargs)\n\n def save_checkpoints(self):\n \"\"\"Save checkpoint, optionally keeping a copy for a given epoch.\"\"\"\n is_sharded = self.cfg.fsdp.use\n if not flashy.distrib.is_rank_zero() and not is_sharded:\n return\n self.logger.info(\"Model hash: %s\", model_hash(self.model))\n state = self.state_dict()\n epoch = self.epoch - 1 # pushing metrics will increase the epoch in Flashy, so we do -1 here\n\n # save minimal state_dict as new checkpoint every X epoch\n if self.cfg.checkpoint.save_every:\n if epoch % self.cfg.checkpoint.save_every == 0:\n minimal_state = state\n if self.cfg.checkpoint.keep_every_states is not None and len(self.cfg.checkpoint.keep_every_states) > 0:\n minimal_state = {\n name: source for name, source in state.items()\n if name in self.cfg.checkpoint.keep_every_states\n }\n epoch_checkpoint_path = self.epoch_checkpoint_path(epoch)\n checkpoint.save_checkpoint(minimal_state, epoch_checkpoint_path, is_sharded)\n\n # save checkpoint as latest checkpoint\n if self.cfg.checkpoint.save_last:\n last_checkpoint_path = self.checkpoint_path()\n checkpoint.save_checkpoint(state, last_checkpoint_path, is_sharded)\n\n # flush any stale checkpoint to reduce disk footprint\n checkpoint.flush_stale_checkpoints(self.checkpoint_path())\n\n def load_from_pretrained(self, name: str) -> dict:\n raise NotImplementedError(\"Solver does not provide a way to load pretrained models.\")\n\n def load_checkpoints(self, load_best: bool = False, ignore_state_keys: tp.List[str] = []) -> tp.Optional[dict]:\n \"\"\"Load last checkpoint or the one specified in continue_from.\n\n Args:\n load_best (bool): Whether to load from best state dict or not.\n Best state dict is always used when not loading the current xp.\n ignore_state_keys (list of str): List of sources to ignore when loading the state, e.g. `optimizer`.\n Returns:\n state (dict, optional): The loaded state dictionary.\n \"\"\"\n # load checkpoints from xp folder or cfg.continue_from\n is_sharded = self.cfg.fsdp.use\n load_from_path: tp.Optional[Path] = None\n checkpoint_source: tp.Optional[checkpoint.CheckpointSource] = None\n\n if load_best:\n self.logger.info(\"Trying to load state_dict from best state.\")\n\n state: tp.Optional[dict] = None\n rank0_checkpoint_path = self.checkpoint_path(use_fsdp=False)\n current_checkpoint_path = self.checkpoint_path()\n _pretrained_prefix = '//pretrained/'\n continue_pretrained = (self.cfg.continue_from or '').startswith(_pretrained_prefix)\n if rank0_checkpoint_path.exists():\n self.logger.info(f\"Loading existing checkpoint: {current_checkpoint_path}\")\n load_from_path = current_checkpoint_path\n checkpoint.check_sharded_checkpoint(current_checkpoint_path, rank0_checkpoint_path)\n checkpoint_source = checkpoint.CheckpointSource.CURRENT_XP\n elif self.cfg.continue_from and not continue_pretrained:\n self.logger.info(f\"Continuing from provided checkpoint: {self.cfg.continue_from}\")\n # we're always continuing from consolidated checkpoints: self.cfg.use_fsdp and not continue_best\n load_from_path = checkpoint.resolve_checkpoint_path(self.cfg.continue_from, use_fsdp=False)\n if load_from_path is None:\n self.logger.error('Could not resolve the continue_from checkpoint %s', self.cfg.continue_from)\n raise RuntimeError(f'Could not resolve continue_from checkpoint {self.cfg.continue_from}')\n checkpoint_source = checkpoint.CheckpointSource.OTHER\n\n if load_from_path is not None:\n state = checkpoint.load_checkpoint(load_from_path, is_sharded)\n elif continue_pretrained:\n self.logger.info(\"Loading a pretrained model. Ignoring 'load_best' and 'ignore_state_keys' params.\")\n state = self.load_from_pretrained(self.cfg.continue_from[len(_pretrained_prefix):])\n checkpoint_source = checkpoint.CheckpointSource.PRETRAINED\n load_best = True\n\n # checkpoints are not from the current xp, we only retrieve the best state\n if checkpoint_source is not None and checkpoint_source != checkpoint.CheckpointSource.CURRENT_XP:\n assert state is not None\n self.logger.info(\"Checkpoint source is not the current xp: Load state_dict from best state.\")\n load_best = True\n state = {key: state[key] for key in self._continue_best_source_keys if key in state}\n # loaded checkpoints are FSDP checkpoints: we're reading the best state\n # from FSDP and we drop the regular best_state\n if 'fsdp_best_state' in state and state['fsdp_best_state']:\n state.pop('best_state', None)\n self.logger.info(\"... Loaded checkpoint has FSDP best state\")\n # FSDP is enabled in the solver, if the loaded checkpoints do not have FSDP support\n # then we're initializing FSDP best state with the regular best state\n elif self.cfg.fsdp.use:\n if 'fsdp_best_state' not in state or not state['fsdp_best_state']:\n # we swap non-FSDP checkpoints best_state to FSDP-compatible best state\n state['fsdp_best_state'] = state.pop('best_state')\n self.logger.info(\"... Loaded checkpoint does not have FSDP best state. Use regular best state\")\n\n if state is not None:\n if load_best:\n self.logger.info(\"Ignoring keys when loading best %r\", ignore_state_keys)\n for key in set(ignore_state_keys):\n if key in state:\n state.pop(key)\n has_best_state = 'best_state' in state or 'fsdp_best_state' in state\n assert has_best_state, (\"Trying to load best state but neither 'best_state'\",\n \" or 'fsdp_best_state' found in checkpoints.\")\n self.load_state_dict(state)\n\n # for FSDP, let's make extra sure nothing bad happened with out of sync\n # checkpoints across workers.\n epoch = float(self.epoch)\n avg_epoch = flashy.distrib.average_metrics({'epoch': epoch})['epoch']\n if avg_epoch != epoch:\n raise RuntimeError(\n f\"Inconsistent loading of checkpoints happened, our epoch is {epoch} \"\n f\"but average of epochs is {avg_epoch}, at least one gpu must have a \"\n \"different epoch number.\")\n\n # on load_best, properly reinitialize state_dict, best states and ema\n # otherwise we load from the current xp and don't alter anything\n if load_best:\n self.logger.info(\"Loading state_dict from best state.\")\n if not self.cfg.fsdp.use and self.fsdp_best_state:\n # loading from an FSDP checkpoint but with FSDP deactivated\n self.logger.info(\"... Loading from FSDP best state dict.\")\n self.best_state.load_state_dict(self.fsdp_best_state)\n\n # if load_best, we permanently override the regular state_dict with the best state\n if self.cfg.fsdp.use:\n self.logger.info(\"FSDP is used, loading from FSDP best state.\")\n with fsdp.switch_to_full_state_dict(self._fsdp_modules):\n # this might be really fragile but okay for now.\n self.load_state_dict(self.fsdp_best_state)\n else:\n # we permanently swap the stateful objects to their best state\n self._load_new_state_dict(self.best_state.state_dict())\n\n # the EMA modules should also be instantiated with best state.\n # the easiest way to do so is to reinitialize a new EMA with best state loaded.\n if self.ema is not None:\n self.logger.info(\"Re-initializing EMA from best state\")\n self.initialize_ema()\n\n if self.cfg.fsdp.use:\n self.logger.info(\"Re-initializing best state after using FSDP best state.\")\n for name in self.best_state.states.keys():\n state_source = self._get_state_source(name)\n self.best_state.update(name, state_source)\n\n return state\n\n def restore(self, load_best: bool = False, replay_metrics: bool = False,\n ignore_state_keys: tp.List[str] = []) -> bool:\n \"\"\"Restore the status of a solver for a given xp.\n\n Args:\n load_best (bool): if `True`, load the best state from the checkpoint.\n replay_metrics (bool): if `True`, logs all the metrics from past epochs.\n ignore_state_keys (list of str): list of sources to ignore when loading the state, e.g. `optimizer`.\n \"\"\"\n self.logger.info(\"Restoring weights and history.\")\n restored_checkpoints = self.load_checkpoints(load_best, ignore_state_keys)\n\n self.logger.info(\"Model hash: %s\", model_hash(self.model))\n\n if replay_metrics and len(self.history) > 0:\n self.logger.info(\"Replaying past metrics...\")\n for epoch, stages in enumerate(self.history):\n for stage_name, metrics in stages.items():\n # We manually log the metrics summary to the result logger\n # as we don't want to add them to the pending metrics\n self.result_logger._log_summary(stage_name, metrics, step=epoch + 1, step_name='epoch',\n formatter=self.get_formatter(stage_name))\n return restored_checkpoints is not None\n\n def commit(self, save_checkpoints: bool = True):\n \"\"\"Commit metrics to dora and save checkpoints at the end of an epoch.\"\"\"\n # we override commit to introduce more complex checkpoint saving behaviors\n self.history.append(self._pending_metrics) # This will increase self.epoch\n if save_checkpoints:\n self.save_checkpoints()\n self._start_epoch()\n if flashy.distrib.is_rank_zero():\n self.xp.link.update_history(self.history)\n\n def run_epoch(self):\n \"\"\"Run a single epoch with all stages.\n\n Metrics for a given stage are stored in _pending_metrics and committed by the solver afterwards.\n Children solvers can extend this method with custom behavior, e.g.:\n\n def run_epoch(self):\n ... # custom code\n super().run_epoch()\n ... # custom code\n \"\"\"\n self.run_stage('train', self.train)\n with torch.no_grad():\n with self.swap_ema_state():\n self.run_stage('valid', self.valid)\n # the best state is updated with EMA states if available\n self.update_best_state_from_stage('valid')\n with self.swap_best_state():\n if self.should_run_stage('evaluate'):\n self.run_stage('evaluate', self.evaluate)\n if self.should_run_stage('generate'):\n self.run_stage('generate', with_rank_rng()(self.generate))\n\n def run(self):\n \"\"\"Training loop.\"\"\"\n assert len(self.state_dict()) > 0\n self.restore(replay_metrics=True) # load checkpoint and replay history\n self.log_hyperparams(dict_from_config(self.cfg))\n for epoch in range(self.epoch, self.cfg.optim.epochs + 1):\n if self.should_stop_training():\n return\n self.run_epoch()\n # Commit will send the metrics to Dora and save checkpoints by default.\n self.commit()\n\n def should_stop_training(self) -> bool:\n \"\"\"Check whether we should stop training or not.\"\"\"\n return self.epoch > self.cfg.optim.epochs\n\n def should_run_stage(self, stage_name) -> bool:\n \"\"\"Check whether we want to run the specified stages.\"\"\"\n stage_every = self.cfg[stage_name].get('every', None)\n is_last_epoch = self.epoch == self.cfg.optim.epochs\n is_epoch_every = (stage_every and self.epoch % stage_every == 0)\n return is_last_epoch or is_epoch_every\n\n @abstractmethod\n def run_step(self, idx: int, batch: tp.Any, metrics: dict):\n \"\"\"Perform one training or valid step on a given batch.\"\"\"\n ...\n\n def common_train_valid(self, dataset_split: str, **kwargs: tp.Any):\n \"\"\"Common logic for train and valid stages.\"\"\"\n self.model.train(self.is_training)\n\n loader = self.dataloaders[dataset_split]\n # get a different order for distributed training, otherwise this will get ignored\n if flashy.distrib.world_size() > 1 \\\n and isinstance(loader.sampler, torch.utils.data.distributed.DistributedSampler):\n loader.sampler.set_epoch(self.epoch)\n updates_per_epoch = self.train_updates_per_epoch if self.is_training else len(loader)\n if self.cfg.benchmark_no_load:\n self.logger.warning(\"Fake loading for benchmarking: re-using first batch\")\n batch = next(iter(loader))\n loader = [batch] * updates_per_epoch # type: ignore\n lp = self.log_progress(self.current_stage, loader, total=updates_per_epoch, updates=self.log_updates)\n average = flashy.averager() # epoch wise average\n instant_average = flashy.averager() # average between two logging\n metrics: dict = {}\n\n with self.profiler, self.deadlock_detect: # profiler will only run for the first 20 updates.\n for idx, batch in enumerate(lp):\n self.deadlock_detect.update('batch')\n if idx >= updates_per_epoch:\n break\n metrics = {}\n metrics = self.run_step(idx, batch, metrics)\n self.deadlock_detect.update('step')\n # run EMA step\n if self.ema is not None and self.is_training and (idx + 1) % self.cfg.optim.ema.updates == 0:\n self.logger.debug(\"EMA model step\")\n self.ema.step()\n self.deadlock_detect.update('ema')\n self.profiler.step()\n instant_metrics = instant_average(metrics)\n if lp.update(**instant_metrics):\n instant_average = flashy.averager() # reset averager between two logging\n metrics = average(metrics) # epoch wise average\n self.deadlock_detect.update('end_batch')\n\n metrics = flashy.distrib.average_metrics(metrics, updates_per_epoch)\n return metrics\n\n def train(self):\n \"\"\"Train stage.\"\"\"\n return self.common_train_valid('train')\n\n def valid(self):\n \"\"\"Valid stage.\"\"\"\n return self.common_train_valid('valid')\n\n @abstractmethod\n def evaluate(self):\n \"\"\"Evaluate stage.\"\"\"\n ...\n\n @abstractmethod\n def generate(self):\n \"\"\"Generate stage.\"\"\"\n ...\n\n def run_one_stage(self, stage_name: str):\n \"\"\"Run only the specified stage.\n This method is useful to only generate samples from a trained experiment\n or rerun the validation or evaluation stages.\n \"\"\"\n fn = {\n 'generate': with_rank_rng()(self.generate),\n 'evaluate': self.evaluate,\n 'valid': self.valid,\n }\n if stage_name not in fn:\n raise ValueError(f'Trying to run stage {stage_name} is not supported.')\n assert len(self.state_dict()) > 0\n self._start_epoch()\n with torch.no_grad(), self.swap_best_state():\n self.run_stage(stage_name, fn[stage_name])\n if not self.cfg.execute_inplace:\n self.commit(save_checkpoints=False)\n\n @staticmethod\n def get_eval_solver_from_sig(sig: str, dtype: tp.Optional[str] = None,\n device: tp.Optional[str] = None, autocast: bool = True,\n batch_size: tp.Optional[int] = None,\n override_cfg: tp.Optional[tp.Union[dict, omegaconf.DictConfig]] = None,\n **kwargs):\n \"\"\"Mostly a convenience function around audiocraft.train.get_solver_from_sig,\n populating all the proper param, deactivating EMA, FSDP, loading the best state,\n basically all you need to get a solver ready to \"play\" with in single GPU mode\n and with minimal memory overhead.\n\n Args:\n sig (str): signature to load.\n dtype (str or None): potential dtype, as a string, i.e. 'float16'.\n device (str or None): potential device, as a string, i.e. 'cuda'.\n override_cfg (dict or omegaconf.DictConfig or None): potential device, as a string, i.e. 'cuda'.\n \"\"\"\n from audiocraft import train\n our_override_cfg: tp.Dict[str, tp.Any] = {'optim': {'ema': {'use': False}}}\n our_override_cfg['autocast'] = autocast\n if dtype is not None:\n our_override_cfg['dtype'] = dtype\n if device is not None:\n our_override_cfg['device'] = device\n if batch_size is not None:\n our_override_cfg['dataset'] = {'batch_size': batch_size}\n if override_cfg is None:\n override_cfg = {}\n override_cfg = omegaconf.OmegaConf.merge(\n omegaconf.DictConfig(override_cfg), omegaconf.DictConfig(our_override_cfg)) # type: ignore\n solver = train.get_solver_from_sig(\n sig, override_cfg=override_cfg,\n load_best=True, disable_fsdp=True,\n ignore_state_keys=['optimizer', 'ema'], **kwargs)\n solver.model.eval()\n return solver\n" }, { "path": "audiocraft/solvers/builders.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nAll the functions to build the relevant solvers and used objects\nfrom the Hydra config.\n\"\"\"\n\nfrom enum import Enum\nimport logging\nimport typing as tp\n\nimport dora\nimport flashy\nimport omegaconf\nimport torch\nfrom torch import nn\nfrom torch.optim import Optimizer\n\n# LRScheduler was renamed in some torch versions\ntry:\n from torch.optim.lr_scheduler import LRScheduler # type: ignore\nexcept ImportError:\n from torch.optim.lr_scheduler import _LRScheduler as LRScheduler\n\nfrom .base import StandardSolver\nfrom .. import adversarial, data, losses, metrics, optim\nfrom ..utils.utils import dict_from_config, get_loader\n\n\nlogger = logging.getLogger(__name__)\n\n\nclass DatasetType(Enum):\n AUDIO = \"audio\"\n MUSIC = \"music\"\n SOUND = \"sound\"\n JASCO = \"jasco\"\n\n\ndef get_solver(cfg: omegaconf.DictConfig) -> StandardSolver:\n \"\"\"Instantiate solver from config.\"\"\"\n from .audiogen import AudioGenSolver\n from .compression import CompressionSolver\n from .musicgen import MusicGenSolver\n from .diffusion import DiffusionSolver\n from .magnet import MagnetSolver, AudioMagnetSolver\n from .watermark import WatermarkSolver\n from .jasco import JascoSolver\n klass = {\n 'compression': CompressionSolver,\n 'musicgen': MusicGenSolver,\n 'audiogen': AudioGenSolver,\n 'magnet': MagnetSolver,\n 'audio_magnet': AudioMagnetSolver,\n 'lm': MusicGenSolver, # backward compatibility\n 'diffusion': DiffusionSolver,\n 'sound_lm': AudioGenSolver, # backward compatibility\n 'watermarking': WatermarkSolver,\n 'jasco': JascoSolver,\n }[cfg.solver]\n return klass(cfg) # type: ignore\n\n\ndef get_optim_parameter_groups(model: nn.Module):\n \"\"\"Create parameter groups for the model using the appropriate method\n if defined for each modules, to create the different groups.\n\n Args:\n model (nn.Module): torch model\n Returns:\n List of parameter groups\n \"\"\"\n seen_params: tp.Set[nn.parameter.Parameter] = set()\n other_params = []\n groups = []\n for name, module in model.named_modules():\n if hasattr(module, 'make_optim_group'):\n group = module.make_optim_group()\n params = set(group['params'])\n assert params.isdisjoint(seen_params)\n seen_params |= set(params)\n groups.append(group)\n for param in model.parameters():\n if param not in seen_params:\n other_params.append(param)\n groups.insert(0, {'params': other_params})\n parameters = groups\n return parameters\n\n\ndef get_optimizer(params: tp.Union[nn.Module, tp.Iterable[torch.Tensor]], cfg: omegaconf.DictConfig) -> Optimizer:\n \"\"\"Build torch optimizer from config and set of parameters.\n Supported optimizers: Adam, AdamW\n\n Args:\n params (nn.Module or iterable of torch.Tensor): Parameters to optimize.\n cfg (DictConfig): Optimization-related configuration.\n Returns:\n torch.optim.Optimizer.\n \"\"\"\n if 'optimizer' not in cfg:\n if getattr(cfg, 'optim', None) is not None:\n raise KeyError(\"Optimizer not found in config. Try instantiating optimizer from cfg.optim?\")\n else:\n raise KeyError(\"Optimizer not found in config.\")\n\n parameters = get_optim_parameter_groups(params) if isinstance(params, nn.Module) else params\n optimizer: torch.optim.Optimizer\n if cfg.optimizer == 'adam':\n optimizer = torch.optim.Adam(parameters, lr=cfg.lr, **cfg.adam)\n elif cfg.optimizer == 'adamw':\n optimizer = torch.optim.AdamW(parameters, lr=cfg.lr, **cfg.adam)\n elif cfg.optimizer == 'dadam':\n optimizer = optim.DAdaptAdam(parameters, lr=cfg.lr, **cfg.adam)\n else:\n raise ValueError(f\"Unsupported Optimizer: {cfg.optimizer}\")\n return optimizer\n\n\ndef get_lr_scheduler(optimizer: torch.optim.Optimizer,\n cfg: omegaconf.DictConfig,\n total_updates: int) -> tp.Optional[LRScheduler]:\n \"\"\"Build torch learning rate scheduler from config and associated optimizer.\n Supported learning rate schedulers: ExponentialLRScheduler, PlateauLRScheduler\n\n Args:\n optimizer (torch.optim.Optimizer): Optimizer.\n cfg (DictConfig): Schedule-related configuration.\n total_updates (int): Total number of updates.\n Returns:\n torch.optim.Optimizer.\n \"\"\"\n if 'lr_scheduler' not in cfg:\n raise KeyError(\"LR Scheduler not found in config\")\n\n lr_sched: tp.Optional[LRScheduler] = None\n if cfg.lr_scheduler == 'step':\n lr_sched = torch.optim.lr_scheduler.StepLR(optimizer, **cfg.step)\n elif cfg.lr_scheduler == 'exponential':\n lr_sched = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=cfg.exponential)\n elif cfg.lr_scheduler == 'cosine':\n kwargs = dict_from_config(cfg.cosine)\n warmup_steps = kwargs.pop('warmup')\n lr_sched = optim.CosineLRScheduler(\n optimizer, warmup_steps=warmup_steps, total_steps=total_updates, **kwargs)\n elif cfg.lr_scheduler == 'polynomial_decay':\n kwargs = dict_from_config(cfg.polynomial_decay)\n warmup_steps = kwargs.pop('warmup')\n lr_sched = optim.PolynomialDecayLRScheduler(\n optimizer, warmup_steps=warmup_steps, total_steps=total_updates, **kwargs)\n elif cfg.lr_scheduler == 'inverse_sqrt':\n kwargs = dict_from_config(cfg.inverse_sqrt)\n warmup_steps = kwargs.pop('warmup')\n lr_sched = optim.InverseSquareRootLRScheduler(optimizer, warmup_steps=warmup_steps, **kwargs)\n elif cfg.lr_scheduler == 'linear_warmup':\n kwargs = dict_from_config(cfg.linear_warmup)\n warmup_steps = kwargs.pop('warmup')\n lr_sched = optim.LinearWarmupLRScheduler(optimizer, warmup_steps=warmup_steps, **kwargs)\n elif cfg.lr_scheduler is not None:\n raise ValueError(f\"Unsupported LR Scheduler: {cfg.lr_scheduler}\")\n return lr_sched\n\n\ndef get_ema(module_dict: nn.ModuleDict, cfg: omegaconf.DictConfig) -> tp.Optional[optim.ModuleDictEMA]:\n \"\"\"Initialize Exponential Moving Average.\n\n Args:\n module_dict (nn.ModuleDict): ModuleDict for which to compute the EMA.\n cfg (omegaconf.DictConfig): Optim EMA configuration.\n Returns:\n optim.ModuleDictEMA: EMA version of the ModuleDict.\n \"\"\"\n kw: tp.Dict[str, tp.Any] = dict(cfg)\n use = kw.pop('use', False)\n decay = kw.pop('decay', None)\n device = kw.pop('device', None)\n if not use:\n return None\n if len(module_dict) == 0:\n raise ValueError(\"Trying to build EMA but an empty module_dict source is provided!\")\n ema_module = optim.ModuleDictEMA(module_dict, decay=decay, device=device)\n return ema_module\n\n\ndef get_loss(loss_name: str, cfg: omegaconf.DictConfig):\n \"\"\"Instantiate loss from configuration.\"\"\"\n klass = {\n 'l1': torch.nn.L1Loss,\n 'l2': torch.nn.MSELoss,\n 'mel': losses.MelSpectrogramL1Loss,\n 'mrstft': losses.MRSTFTLoss,\n 'msspec': losses.MultiScaleMelSpectrogramLoss,\n 'sisnr': losses.SISNR,\n 'wm_detection': losses.WMDetectionLoss,\n 'wm_mb': losses.WMMbLoss,\n 'tf_loudnessratio': losses.TFLoudnessRatio\n }[loss_name]\n kwargs = dict(getattr(cfg, loss_name))\n return klass(**kwargs)\n\n\ndef get_balancer(loss_weights: tp.Dict[str, float], cfg: omegaconf.DictConfig) -> losses.Balancer:\n \"\"\"Instantiate loss balancer from configuration for the provided weights.\"\"\"\n kwargs: tp.Dict[str, tp.Any] = dict_from_config(cfg)\n return losses.Balancer(loss_weights, **kwargs)\n\n\ndef get_adversary(name: str, cfg: omegaconf.DictConfig) -> nn.Module:\n \"\"\"Initialize adversary from config.\"\"\"\n klass = {\n 'msd': adversarial.MultiScaleDiscriminator,\n 'mpd': adversarial.MultiPeriodDiscriminator,\n 'msstftd': adversarial.MultiScaleSTFTDiscriminator,\n }[name]\n adv_cfg: tp.Dict[str, tp.Any] = dict(getattr(cfg, name))\n return klass(**adv_cfg)\n\n\ndef get_adversarial_losses(cfg) -> nn.ModuleDict:\n \"\"\"Initialize dict of adversarial losses from config.\"\"\"\n device = cfg.device\n adv_cfg = getattr(cfg, 'adversarial')\n adversaries = adv_cfg.get('adversaries', [])\n adv_loss_name = adv_cfg['adv_loss']\n feat_loss_name = adv_cfg.get('feat_loss')\n normalize = adv_cfg.get('normalize', True)\n feat_loss: tp.Optional[adversarial.FeatureMatchingLoss] = None\n if feat_loss_name:\n assert feat_loss_name in ['l1', 'l2'], f\"Feature loss only support L1 or L2 but {feat_loss_name} found.\"\n loss = get_loss(feat_loss_name, cfg)\n feat_loss = adversarial.FeatureMatchingLoss(loss, normalize)\n loss = adversarial.get_adv_criterion(adv_loss_name)\n loss_real = adversarial.get_real_criterion(adv_loss_name)\n loss_fake = adversarial.get_fake_criterion(adv_loss_name)\n adv_losses = nn.ModuleDict()\n for adv_name in adversaries:\n adversary = get_adversary(adv_name, cfg).to(device)\n optimizer = get_optimizer(adversary.parameters(), cfg.optim)\n adv_loss = adversarial.AdversarialLoss(\n adversary,\n optimizer,\n loss=loss,\n loss_real=loss_real,\n loss_fake=loss_fake,\n loss_feat=feat_loss,\n normalize=normalize\n )\n adv_losses[adv_name] = adv_loss\n return adv_losses\n\n\ndef get_visqol(cfg: omegaconf.DictConfig) -> metrics.ViSQOL:\n \"\"\"Instantiate ViSQOL metric from config.\"\"\"\n kwargs = dict_from_config(cfg)\n return metrics.ViSQOL(**kwargs)\n\n\ndef get_fad(cfg: omegaconf.DictConfig) -> metrics.FrechetAudioDistanceMetric:\n \"\"\"Instantiate Frechet Audio Distance metric from config.\"\"\"\n kwargs = dict_from_config(cfg.tf)\n xp = dora.get_xp()\n kwargs['log_folder'] = xp.folder\n return metrics.FrechetAudioDistanceMetric(**kwargs)\n\n\ndef get_kldiv(cfg: omegaconf.DictConfig) -> metrics.KLDivergenceMetric:\n \"\"\"Instantiate KL-Divergence metric from config.\"\"\"\n kld_metrics = {\n 'passt': metrics.PasstKLDivergenceMetric,\n }\n klass = kld_metrics[cfg.model]\n kwargs = dict_from_config(cfg.get(cfg.model))\n return klass(**kwargs)\n\n\ndef get_text_consistency(cfg: omegaconf.DictConfig) -> metrics.TextConsistencyMetric:\n \"\"\"Instantiate Text Consistency metric from config.\"\"\"\n text_consistency_metrics = {\n 'clap': metrics.CLAPTextConsistencyMetric\n }\n klass = text_consistency_metrics[cfg.model]\n kwargs = dict_from_config(cfg.get(cfg.model))\n return klass(**kwargs)\n\n\ndef get_chroma_cosine_similarity(cfg: omegaconf.DictConfig) -> metrics.ChromaCosineSimilarityMetric:\n \"\"\"Instantiate Chroma Cosine Similarity metric from config.\"\"\"\n assert cfg.model == 'chroma_base', \"Only support 'chroma_base' method for chroma cosine similarity metric\"\n kwargs = dict_from_config(cfg.get(cfg.model))\n return metrics.ChromaCosineSimilarityMetric(**kwargs)\n\n\ndef get_audio_datasets(cfg: omegaconf.DictConfig,\n dataset_type: DatasetType = DatasetType.AUDIO) -> tp.Dict[str, torch.utils.data.DataLoader]:\n \"\"\"Build AudioDataset from configuration.\n\n Args:\n cfg (omegaconf.DictConfig): Configuration.\n dataset_type: The type of dataset to create.\n Returns:\n dict[str, torch.utils.data.DataLoader]: Map of dataloader for each data split.\n \"\"\"\n dataloaders: dict = {}\n\n sample_rate = cfg.sample_rate\n channels = cfg.channels\n seed = cfg.seed\n max_sample_rate = cfg.datasource.max_sample_rate\n max_channels = cfg.datasource.max_channels\n\n assert cfg.dataset is not None, \"Could not find dataset definition in config\"\n\n dataset_cfg = dict_from_config(cfg.dataset)\n splits_cfg: dict = {}\n splits_cfg['train'] = dataset_cfg.pop('train')\n splits_cfg['valid'] = dataset_cfg.pop('valid')\n splits_cfg['evaluate'] = dataset_cfg.pop('evaluate')\n splits_cfg['generate'] = dataset_cfg.pop('generate')\n execute_only_stage = cfg.get('execute_only', None)\n\n for split, path in cfg.datasource.items():\n if not isinstance(path, str):\n continue # skipping this as not a path\n if execute_only_stage is not None and split != execute_only_stage:\n continue\n logger.info(f\"Loading audio data split {split}: {str(path)}\")\n assert (\n cfg.sample_rate <= max_sample_rate\n ), f\"Expecting a max sample rate of {max_sample_rate} for datasource but {sample_rate} found.\"\n assert (\n cfg.channels <= max_channels\n ), f\"Expecting a max number of channels of {max_channels} for datasource but {channels} found.\"\n\n split_cfg = splits_cfg[split]\n split_kwargs = {k: v for k, v in split_cfg.items()}\n kwargs = {**dataset_cfg, **split_kwargs} # split kwargs overrides default dataset_cfg\n kwargs['sample_rate'] = sample_rate\n kwargs['channels'] = channels\n\n if kwargs.get('permutation_on_files') and cfg.optim.updates_per_epoch:\n kwargs['num_samples'] = (\n flashy.distrib.world_size() * cfg.dataset.batch_size * cfg.optim.updates_per_epoch)\n\n num_samples = kwargs['num_samples']\n shuffle = kwargs['shuffle']\n\n return_info = kwargs.pop('return_info')\n batch_size = kwargs.pop('batch_size', None)\n num_workers = kwargs.pop('num_workers')\n\n if dataset_type == DatasetType.MUSIC:\n dataset = data.music_dataset.MusicDataset.from_meta(path, **kwargs)\n elif dataset_type == DatasetType.SOUND:\n dataset = data.sound_dataset.SoundDataset.from_meta(path, **kwargs)\n elif dataset_type == DatasetType.AUDIO:\n dataset = data.info_audio_dataset.InfoAudioDataset.from_meta(path, return_info=return_info, **kwargs)\n elif dataset_type == DatasetType.JASCO:\n dataset = data.jasco_dataset.JascoDataset.from_meta(path, return_info=return_info, **kwargs)\n else:\n raise ValueError(f\"Dataset type is unsupported: {dataset_type}\")\n\n loader = get_loader(\n dataset,\n num_samples,\n batch_size=batch_size,\n num_workers=num_workers,\n seed=seed,\n collate_fn=dataset.collater if return_info else None,\n shuffle=shuffle,\n )\n dataloaders[split] = loader\n\n return dataloaders\n" }, { "path": "audiocraft/solvers/compression.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nimport multiprocessing\nfrom pathlib import Path\nimport typing as tp\n\nimport flashy\nimport omegaconf\nimport torch\nfrom torch import nn\n\nfrom . import base, builders\nfrom .. import models, quantization\nfrom ..utils import checkpoint\nfrom ..utils.samples.manager import SampleManager\nfrom ..utils.utils import get_pool_executor\n\n\nlogger = logging.getLogger(__name__)\n\n\nclass CompressionSolver(base.StandardSolver):\n \"\"\"Solver for compression task.\n\n The compression task combines a set of perceptual and objective losses\n to train an EncodecModel (composed of an encoder-decoder and a quantizer)\n to perform high fidelity audio reconstruction.\n \"\"\"\n def __init__(self, cfg: omegaconf.DictConfig):\n super().__init__(cfg)\n self.rng: torch.Generator # set at each epoch\n self.adv_losses = builders.get_adversarial_losses(self.cfg)\n self.aux_losses = nn.ModuleDict()\n self.info_losses = nn.ModuleDict()\n assert not cfg.fsdp.use, \"FSDP not supported by CompressionSolver.\"\n loss_weights = dict()\n for loss_name, weight in self.cfg.losses.items():\n if loss_name in ['adv', 'feat']:\n for adv_name, _ in self.adv_losses.items():\n loss_weights[f'{loss_name}_{adv_name}'] = weight\n elif weight > 0:\n self.aux_losses[loss_name] = builders.get_loss(loss_name, self.cfg)\n loss_weights[loss_name] = weight\n else:\n self.info_losses[loss_name] = builders.get_loss(loss_name, self.cfg)\n self.balancer = builders.get_balancer(loss_weights, self.cfg.balancer)\n self.register_stateful('adv_losses')\n\n @property\n def best_metric_name(self) -> tp.Optional[str]:\n # best model is the last for the compression model\n return None\n\n def build_model(self):\n \"\"\"Instantiate model and optimizer.\"\"\"\n # Model and optimizer\n self.model = models.builders.get_compression_model(self.cfg).to(self.device)\n self.optimizer = builders.get_optimizer(self.model.parameters(), self.cfg.optim)\n self.register_stateful('model', 'optimizer')\n self.register_best_state('model')\n self.register_ema('model')\n\n def build_dataloaders(self):\n \"\"\"Instantiate audio dataloaders for each stage.\"\"\"\n self.dataloaders = builders.get_audio_datasets(self.cfg)\n\n def show(self):\n \"\"\"Show the compression model and employed adversarial loss.\"\"\"\n self.logger.info(f\"Compression model with {self.model.quantizer.total_codebooks} codebooks:\")\n self.log_model_summary(self.model)\n self.logger.info(\"Adversarial loss:\")\n self.log_model_summary(self.adv_losses)\n self.logger.info(\"Auxiliary losses:\")\n self.logger.info(self.aux_losses)\n self.logger.info(\"Info losses:\")\n self.logger.info(self.info_losses)\n\n def run_step(self, idx: int, batch: torch.Tensor, metrics: dict):\n \"\"\"Perform one training or valid step on a given batch.\"\"\"\n x = batch.to(self.device)\n y = x.clone()\n\n qres = self.model(x)\n assert isinstance(qres, quantization.QuantizedResult)\n y_pred = qres.x\n # Log bandwidth in kb/s\n metrics['bandwidth'] = qres.bandwidth.mean()\n\n if self.is_training:\n d_losses: dict = {}\n if len(self.adv_losses) > 0 and torch.rand(1, generator=self.rng).item() <= 1 / self.cfg.adversarial.every:\n for adv_name, adversary in self.adv_losses.items():\n disc_loss = adversary.train_adv(y_pred, y)\n d_losses[f'd_{adv_name}'] = disc_loss\n metrics['d_loss'] = torch.sum(torch.stack(list(d_losses.values())))\n metrics.update(d_losses)\n\n balanced_losses: dict = {}\n other_losses: dict = {}\n\n # penalty from quantization\n if qres.penalty is not None and qres.penalty.requires_grad:\n other_losses['penalty'] = qres.penalty # penalty term from the quantizer\n\n # adversarial losses\n for adv_name, adversary in self.adv_losses.items():\n adv_loss, feat_loss = adversary(y_pred, y)\n balanced_losses[f'adv_{adv_name}'] = adv_loss\n balanced_losses[f'feat_{adv_name}'] = feat_loss\n\n # auxiliary losses\n for loss_name, criterion in self.aux_losses.items():\n loss = criterion(y_pred, y)\n balanced_losses[loss_name] = loss\n\n # weighted losses\n metrics.update(balanced_losses)\n metrics.update(other_losses)\n metrics.update(qres.metrics)\n\n if self.is_training:\n # backprop losses that are not handled by balancer\n other_loss = torch.tensor(0., device=self.device)\n if 'penalty' in other_losses:\n other_loss += other_losses['penalty']\n if other_loss.requires_grad:\n other_loss.backward(retain_graph=True)\n ratio1 = sum(p.grad.data.norm(p=2).pow(2)\n for p in self.model.parameters() if p.grad is not None)\n assert isinstance(ratio1, torch.Tensor)\n metrics['ratio1'] = ratio1.sqrt()\n\n # balancer losses backward, returns effective training loss\n # with effective weights at the current batch.\n metrics['g_loss'] = self.balancer.backward(balanced_losses, y_pred)\n # add metrics corresponding to weight ratios\n metrics.update(self.balancer.metrics)\n ratio2 = sum(p.grad.data.norm(p=2).pow(2)\n for p in self.model.parameters() if p.grad is not None)\n assert isinstance(ratio2, torch.Tensor)\n metrics['ratio2'] = ratio2.sqrt()\n\n # optim\n flashy.distrib.sync_model(self.model)\n if self.cfg.optim.max_norm:\n torch.nn.utils.clip_grad_norm_(\n self.model.parameters(), self.cfg.optim.max_norm\n )\n self.optimizer.step()\n self.optimizer.zero_grad()\n\n # informative losses only\n info_losses: dict = {}\n with torch.no_grad():\n for loss_name, criterion in self.info_losses.items():\n loss = criterion(y_pred, y)\n info_losses[loss_name] = loss\n\n metrics.update(info_losses)\n\n # aggregated GAN losses: this is useful to report adv and feat across different adversarial loss setups\n adv_losses = [loss for loss_name, loss in metrics.items() if loss_name.startswith('adv')]\n if len(adv_losses) > 0:\n metrics['adv'] = torch.sum(torch.stack(adv_losses))\n feat_losses = [loss for loss_name, loss in metrics.items() if loss_name.startswith('feat')]\n if len(feat_losses) > 0:\n metrics['feat'] = torch.sum(torch.stack(feat_losses))\n\n return metrics\n\n def run_epoch(self):\n # reset random seed at the beginning of the epoch\n self.rng = torch.Generator()\n self.rng.manual_seed(1234 + self.epoch)\n # run epoch\n super().run_epoch()\n\n def evaluate(self):\n \"\"\"Evaluate stage. Runs audio reconstruction evaluation.\"\"\"\n self.model.eval()\n evaluate_stage_name = str(self.current_stage)\n\n loader = self.dataloaders['evaluate']\n updates = len(loader)\n lp = self.log_progress(f'{evaluate_stage_name} inference', loader, total=updates, updates=self.log_updates)\n average = flashy.averager()\n\n pendings = []\n ctx = multiprocessing.get_context('spawn')\n with get_pool_executor(self.cfg.evaluate.num_workers, mp_context=ctx) as pool:\n for idx, batch in enumerate(lp):\n x = batch.to(self.device)\n with torch.no_grad():\n qres = self.model(x)\n\n y_pred = qres.x.cpu()\n y = batch.cpu() # should already be on CPU but just in case\n pendings.append(pool.submit(evaluate_audio_reconstruction, y_pred, y, self.cfg))\n\n metrics_lp = self.log_progress(f'{evaluate_stage_name} metrics', pendings, updates=self.log_updates)\n for pending in metrics_lp:\n metrics = pending.result()\n metrics = average(metrics)\n\n metrics = flashy.distrib.average_metrics(metrics, len(loader))\n return metrics\n\n def generate(self):\n \"\"\"Generate stage.\"\"\"\n self.model.eval()\n sample_manager = SampleManager(self.xp, map_reference_to_sample_id=True)\n generate_stage_name = str(self.current_stage)\n\n loader = self.dataloaders['generate']\n updates = len(loader)\n lp = self.log_progress(generate_stage_name, loader, total=updates, updates=self.log_updates)\n\n for batch in lp:\n reference, _ = batch\n reference = reference.to(self.device)\n with torch.no_grad():\n qres = self.model(reference)\n assert isinstance(qres, quantization.QuantizedResult)\n\n reference = reference.cpu()\n estimate = qres.x.cpu()\n sample_manager.add_samples(estimate, self.epoch, ground_truth_wavs=reference)\n\n flashy.distrib.barrier()\n\n def load_from_pretrained(self, name: str) -> dict:\n model = models.CompressionModel.get_pretrained(name)\n if isinstance(model, models.DAC):\n raise RuntimeError(\"Cannot fine tune a DAC model.\")\n elif isinstance(model, models.HFEncodecCompressionModel):\n self.logger.warning('Trying to automatically convert a HuggingFace model '\n 'to AudioCraft, this might fail!')\n state = model.model.state_dict()\n new_state = {}\n for k, v in state.items():\n if k.startswith('decoder.layers') and '.conv.' in k and '.block.' not in k:\n # We need to determine if this a convtr or a regular conv.\n layer = int(k.split('.')[2])\n if isinstance(model.model.decoder.layers[layer].conv, torch.nn.ConvTranspose1d):\n\n k = k.replace('.conv.', '.convtr.')\n k = k.replace('encoder.layers.', 'encoder.model.')\n k = k.replace('decoder.layers.', 'decoder.model.')\n k = k.replace('conv.', 'conv.conv.')\n k = k.replace('convtr.', 'convtr.convtr.')\n k = k.replace('quantizer.layers.', 'quantizer.vq.layers.')\n k = k.replace('.codebook.', '._codebook.')\n new_state[k] = v\n state = new_state\n elif isinstance(model, models.EncodecModel):\n state = model.state_dict()\n else:\n raise RuntimeError(f\"Cannot fine tune model type {type(model)}.\")\n return {\n 'best_state': {'model': state}\n }\n\n @staticmethod\n def model_from_checkpoint(checkpoint_path: tp.Union[Path, str],\n device: tp.Union[torch.device, str] = 'cpu') -> models.CompressionModel:\n \"\"\"Instantiate a CompressionModel from a given checkpoint path or dora sig.\n This method is a convenient endpoint to load a CompressionModel to use in other solvers.\n\n Args:\n checkpoint_path (Path or str): Path to checkpoint or dora sig from where the checkpoint is resolved.\n This also supports pre-trained models by using a path of the form //pretrained/NAME.\n See `model_from_pretrained` for a list of supported pretrained models.\n use_ema (bool): Use EMA variant of the model instead of the actual model.\n device (torch.device or str): Device on which the model is loaded.\n \"\"\"\n checkpoint_path = str(checkpoint_path)\n if checkpoint_path.startswith('//pretrained/'):\n name = checkpoint_path.split('/', 3)[-1]\n return models.CompressionModel.get_pretrained(name, device)\n logger = logging.getLogger(__name__)\n logger.info(f\"Loading compression model from checkpoint: {checkpoint_path}\")\n _checkpoint_path = checkpoint.resolve_checkpoint_path(checkpoint_path, use_fsdp=False)\n assert _checkpoint_path is not None, f\"Could not resolve compression model checkpoint path: {checkpoint_path}\"\n state = checkpoint.load_checkpoint(_checkpoint_path)\n assert state is not None and 'xp.cfg' in state, f\"Could not load compression model from ckpt: {checkpoint_path}\"\n cfg = state['xp.cfg']\n cfg.device = device\n compression_model = models.builders.get_compression_model(cfg).to(device)\n assert compression_model.sample_rate == cfg.sample_rate, \"Compression model sample rate should match\"\n\n assert 'best_state' in state and state['best_state'] != {}\n assert 'exported' not in state, \"When loading an exported checkpoint, use the //pretrained/ prefix.\"\n compression_model.load_state_dict(state['best_state']['model'])\n compression_model.eval()\n logger.info(\"Compression model loaded!\")\n return compression_model\n\n @staticmethod\n def wrapped_model_from_checkpoint(cfg: omegaconf.DictConfig,\n checkpoint_path: tp.Union[Path, str],\n device: tp.Union[torch.device, str] = 'cpu') -> models.CompressionModel:\n \"\"\"Instantiate a wrapped CompressionModel from a given checkpoint path or dora sig.\n\n Args:\n cfg (omegaconf.DictConfig): Configuration to read from for wrapped mode.\n checkpoint_path (Path or str): Path to checkpoint or dora sig from where the checkpoint is resolved.\n use_ema (bool): Use EMA variant of the model instead of the actual model.\n device (torch.device or str): Device on which the model is loaded.\n \"\"\"\n compression_model = CompressionSolver.model_from_checkpoint(checkpoint_path, device)\n compression_model = models.builders.get_wrapped_compression_model(compression_model, cfg)\n return compression_model\n\n\ndef evaluate_audio_reconstruction(y_pred: torch.Tensor, y: torch.Tensor, cfg: omegaconf.DictConfig) -> dict:\n \"\"\"Audio reconstruction evaluation method that can be conveniently pickled.\"\"\"\n metrics = {}\n if cfg.evaluate.metrics.visqol:\n visqol = builders.get_visqol(cfg.metrics.visqol)\n metrics['visqol'] = visqol(y_pred, y, cfg.sample_rate)\n sisnr = builders.get_loss('sisnr', cfg)\n metrics['sisnr'] = sisnr(y_pred, y)\n return metrics\n" }, { "path": "audiocraft/solvers/diffusion.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport typing as tp\n\nimport flashy\nimport julius\nimport omegaconf\nimport torch\nimport torch.nn.functional as F\n\nfrom . import builders\nfrom . import base\nfrom .. import models\nfrom ..modules.diffusion_schedule import NoiseSchedule\nfrom ..metrics import RelativeVolumeMel\nfrom ..models.builders import get_processor\nfrom ..utils.samples.manager import SampleManager\nfrom ..solvers.compression import CompressionSolver\n\n\nclass PerStageMetrics:\n \"\"\"Handle prompting the metrics per stage.\n It outputs the metrics per range of diffusion states.\n e.g. avg loss when t in [250, 500]\n \"\"\"\n def __init__(self, num_steps: int, num_stages: int = 4):\n self.num_steps = num_steps\n self.num_stages = num_stages\n\n def __call__(self, losses: dict, step: tp.Union[int, torch.Tensor]):\n if type(step) is int:\n stage = int((step / self.num_steps) * self.num_stages)\n return {f\"{name}_{stage}\": loss for name, loss in losses.items()}\n elif type(step) is torch.Tensor:\n stage_tensor = ((step / self.num_steps) * self.num_stages).long()\n out: tp.Dict[str, float] = {}\n for stage_idx in range(self.num_stages):\n mask = (stage_tensor == stage_idx)\n N = mask.sum()\n stage_out = {}\n if N > 0: # pass if no elements in the stage\n for name, loss in losses.items():\n stage_loss = (mask * loss).sum() / N\n stage_out[f\"{name}_{stage_idx}\"] = stage_loss\n out = {**out, **stage_out}\n return out\n\n\nclass DataProcess:\n \"\"\"Apply filtering or resampling.\n\n Args:\n initial_sr (int): Initial sample rate.\n target_sr (int): Target sample rate.\n use_resampling: Whether to use resampling or not.\n use_filter (bool):\n n_bands (int): Number of bands to consider.\n idx_band (int):\n device (torch.device or str):\n cutoffs ():\n boost (bool):\n \"\"\"\n def __init__(self, initial_sr: int = 24000, target_sr: int = 16000, use_resampling: bool = False,\n use_filter: bool = False, n_bands: int = 4,\n idx_band: int = 0, device: torch.device = torch.device('cpu'), cutoffs=None, boost=False):\n \"\"\"Apply filtering or resampling\n Args:\n initial_sr (int): sample rate of the dataset\n target_sr (int): sample rate after resampling\n use_resampling (bool): whether or not performs resampling\n use_filter (bool): when True filter the data to keep only one frequency band\n n_bands (int): Number of bands used\n cuts (none or list): The cutoff frequencies of the band filtering\n if None then we use mel scale bands.\n idx_band (int): index of the frequency band. 0 are lows ... (n_bands - 1) highs\n boost (bool): make the data scale match our music dataset.\n \"\"\"\n assert idx_band < n_bands\n self.idx_band = idx_band\n if use_filter:\n if cutoffs is not None:\n self.filter = julius.SplitBands(sample_rate=initial_sr, cutoffs=cutoffs).to(device)\n else:\n self.filter = julius.SplitBands(sample_rate=initial_sr, n_bands=n_bands).to(device)\n self.use_filter = use_filter\n self.use_resampling = use_resampling\n self.target_sr = target_sr\n self.initial_sr = initial_sr\n self.boost = boost\n\n def process_data(self, x, metric=False):\n if x is None:\n return None\n if self.boost:\n x /= torch.clamp(x.std(dim=(1, 2), keepdim=True), min=1e-4)\n x * 0.22\n if self.use_filter and not metric:\n x = self.filter(x)[self.idx_band]\n if self.use_resampling:\n x = julius.resample_frac(x, old_sr=self.initial_sr, new_sr=self.target_sr)\n return x\n\n def inverse_process(self, x):\n \"\"\"Upsampling only.\"\"\"\n if self.use_resampling:\n x = julius.resample_frac(x, old_sr=self.target_sr, new_sr=self.target_sr)\n return x\n\n\nclass DiffusionSolver(base.StandardSolver):\n \"\"\"Solver for compression task.\n\n The diffusion task allows for MultiBand diffusion model training.\n\n Args:\n cfg (DictConfig): Configuration.\n \"\"\"\n def __init__(self, cfg: omegaconf.DictConfig):\n super().__init__(cfg)\n self.cfg = cfg\n self.device = cfg.device\n self.sample_rate: int = self.cfg.sample_rate\n self.codec_model = CompressionSolver.model_from_checkpoint(\n cfg.compression_model_checkpoint, device=self.device)\n\n self.codec_model.set_num_codebooks(cfg.n_q)\n assert self.codec_model.sample_rate == self.cfg.sample_rate, (\n f\"Codec model sample rate is {self.codec_model.sample_rate} but \"\n f\"Solver sample rate is {self.cfg.sample_rate}.\"\n )\n assert self.codec_model.sample_rate == self.sample_rate, \\\n f\"Sample rate of solver {self.sample_rate} and codec {self.codec_model.sample_rate} \" \\\n \"don't match.\"\n\n self.sample_processor = get_processor(cfg.processor, sample_rate=self.sample_rate)\n self.register_stateful('sample_processor')\n self.sample_processor.to(self.device)\n\n self.schedule = NoiseSchedule(\n **cfg.schedule, device=self.device, sample_processor=self.sample_processor)\n\n self.eval_metric: tp.Optional[torch.nn.Module] = None\n\n self.rvm = RelativeVolumeMel()\n self.data_processor = DataProcess(initial_sr=self.sample_rate, target_sr=cfg.resampling.target_sr,\n use_resampling=cfg.resampling.use, cutoffs=cfg.filter.cutoffs,\n use_filter=cfg.filter.use, n_bands=cfg.filter.n_bands,\n idx_band=cfg.filter.idx_band, device=self.device)\n\n @property\n def best_metric_name(self) -> tp.Optional[str]:\n if self._current_stage == \"evaluate\":\n return 'rvm'\n else:\n return 'loss'\n\n @torch.no_grad()\n def get_condition(self, wav: torch.Tensor) -> torch.Tensor:\n codes, scale = self.codec_model.encode(wav)\n assert scale is None, \"Scaled compression models not supported.\"\n emb = self.codec_model.decode_latent(codes)\n return emb\n\n def build_model(self):\n \"\"\"Build model and optimizer as well as optional Exponential Moving Average of the model.\n \"\"\"\n # Model and optimizer\n self.model = models.builders.get_diffusion_model(self.cfg).to(self.device)\n self.optimizer = builders.get_optimizer(self.model.parameters(), self.cfg.optim)\n self.register_stateful('model', 'optimizer')\n self.register_best_state('model')\n self.register_ema('model')\n\n def build_dataloaders(self):\n \"\"\"Build audio dataloaders for each stage.\"\"\"\n self.dataloaders = builders.get_audio_datasets(self.cfg)\n\n def show(self):\n # TODO\n raise NotImplementedError()\n\n def run_step(self, idx: int, batch: torch.Tensor, metrics: dict):\n \"\"\"Perform one training or valid step on a given batch.\"\"\"\n x = batch.to(self.device)\n loss_fun = F.mse_loss if self.cfg.loss.kind == 'mse' else F.l1_loss\n\n condition = self.get_condition(x) # [bs, 128, T/hop, n_emb]\n sample = self.data_processor.process_data(x)\n\n input_, target, step = self.schedule.get_training_item(sample,\n tensor_step=self.cfg.schedule.variable_step_batch)\n out = self.model(input_, step, condition=condition).sample\n\n base_loss = loss_fun(out, target, reduction='none').mean(dim=(1, 2))\n reference_loss = loss_fun(input_, target, reduction='none').mean(dim=(1, 2))\n loss = base_loss / reference_loss ** self.cfg.loss.norm_power\n\n if self.is_training:\n loss.mean().backward()\n flashy.distrib.sync_model(self.model)\n self.optimizer.step()\n self.optimizer.zero_grad()\n metrics = {\n 'loss': loss.mean(), 'normed_loss': (base_loss / reference_loss).mean(),\n }\n metrics.update(self.per_stage({'loss': loss, 'normed_loss': base_loss / reference_loss}, step))\n metrics.update({\n 'std_in': input_.std(), 'std_out': out.std()})\n return metrics\n\n def run_epoch(self):\n # reset random seed at the beginning of the epoch\n self.rng = torch.Generator()\n self.rng.manual_seed(1234 + self.epoch)\n self.per_stage = PerStageMetrics(self.schedule.num_steps, self.cfg.metrics.num_stage)\n # run epoch\n super().run_epoch()\n\n def evaluate(self):\n \"\"\"Evaluate stage.\n Runs audio reconstruction evaluation.\n \"\"\"\n self.model.eval()\n evaluate_stage_name = f'{self.current_stage}'\n loader = self.dataloaders['evaluate']\n updates = len(loader)\n lp = self.log_progress(f'{evaluate_stage_name} estimate', loader, total=updates, updates=self.log_updates)\n\n metrics = {}\n n = 1\n for idx, batch in enumerate(lp):\n x = batch.to(self.device)\n with torch.no_grad():\n y_pred = self.regenerate(x)\n\n y_pred = y_pred.cpu()\n y = batch.cpu() # should already be on CPU but just in case\n rvm = self.rvm(y_pred, y)\n lp.update(**rvm)\n if len(metrics) == 0:\n metrics = rvm\n else:\n for key in rvm.keys():\n metrics[key] = (metrics[key] * n + rvm[key]) / (n + 1)\n metrics = flashy.distrib.average_metrics(metrics)\n return metrics\n\n @torch.no_grad()\n def regenerate(self, wav: torch.Tensor, step_list: tp.Optional[list] = None):\n \"\"\"Regenerate the given waveform.\"\"\"\n condition = self.get_condition(wav)\n initial = self.schedule.get_initial_noise(self.data_processor.process_data(wav)) # sampling rate changes.\n result = self.schedule.generate_subsampled(self.model, initial=initial, condition=condition,\n step_list=step_list)\n result = self.data_processor.inverse_process(result)\n return result\n\n def generate(self):\n \"\"\"Generate stage.\"\"\"\n sample_manager = SampleManager(self.xp)\n self.model.eval()\n generate_stage_name = f'{self.current_stage}'\n\n loader = self.dataloaders['generate']\n updates = len(loader)\n lp = self.log_progress(generate_stage_name, loader, total=updates, updates=self.log_updates)\n\n for batch in lp:\n reference, _ = batch\n reference = reference.to(self.device)\n estimate = self.regenerate(reference)\n reference = reference.cpu()\n estimate = estimate.cpu()\n sample_manager.add_samples(estimate, self.epoch, ground_truth_wavs=reference)\n flashy.distrib.barrier()\n" }, { "path": "audiocraft/solvers/jasco.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom omegaconf import DictConfig\nfrom . import builders, musicgen\nfrom .compression import CompressionSolver\nfrom .. import models\nfrom ..modules.conditioners import JascoCondConst, SegmentWithAttributes\nimport torch\nimport typing as tp\nimport flashy\nimport time\nimport math\n\n\nclass JascoSolver(musicgen.MusicGenSolver):\n \"\"\"Solver for JASCO - Joint Audio and Symbolic Conditioning for Temporally Controlled Text-to-Music Generation\n https://arxiv.org/abs/2406.10970.\n \"\"\"\n DATASET_TYPE: builders.DatasetType = builders.DatasetType.JASCO\n\n def __init__(self, cfg: DictConfig):\n super().__init__(cfg)\n\n # initialize generation parameters by config\n self.generation_params = {\n 'cfg_coef_all': self.cfg.generate.lm.cfg_coef_all,\n 'cfg_coef_txt': self.cfg.generate.lm.cfg_coef_txt\n }\n\n self.latent_mean = cfg.compression_model_latent_mean\n self.latent_std = cfg.compression_model_latent_std\n self.mse = torch.nn.MSELoss(reduction='none')\n self._best_metric_name = 'loss'\n\n def build_model(self) -> None:\n \"\"\"Instantiate model and optimization.\"\"\"\n assert self.cfg.efficient_attention_backend == \"xformers\", \"JASCO v1 models support only xformers backend.\"\n\n self.compression_model = CompressionSolver.wrapped_model_from_checkpoint(\n self.cfg, self.cfg.compression_model_checkpoint, device=self.device)\n assert self.compression_model.sample_rate == self.cfg.sample_rate, (\n f\"Compression model sample rate is {self.compression_model.sample_rate} but \"\n f\"Solver sample rate is {self.cfg.sample_rate}.\"\n )\n # instantiate JASCO model\n self.model: models.FlowMatchingModel = models.builders.get_jasco_model(self.cfg,\n self.compression_model).to(self.device)\n # initialize optimization\n self.initialize_optimization()\n\n def _get_latents(self, audio):\n with torch.no_grad():\n latents = self.compression_model.model.encoder(audio)\n return latents.permute(0, 2, 1) # [B, D, T] -> [B, T, D]\n\n def _prepare_latents_and_attributes(\n self, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]],\n ) -> tp.Tuple[dict, torch.Tensor, torch.Tensor]:\n \"\"\"Prepare input batchs for language model training.\n\n Args:\n batch (tuple[torch.Tensor, list[SegmentWithAttributes]]): Input batch with audio tensor of shape [B, C, T]\n and corresponding metadata as SegmentWithAttributes (with B items).\n Returns:\n Condition tensors (dict[str, any]): Preprocessed condition attributes.\n Tokens (torch.Tensor): Audio tokens from compression model, of shape [B, K, T_s],\n with B the batch size, K the number of codebooks, T_s the token timesteps.\n Padding mask (torch.Tensor): Mask with valid positions in the tokens tensor, of shape [B, K, T_s].\n \"\"\"\n audio, infos = batch\n audio = audio.to(self.device)\n assert audio.size(0) == len(infos), (\n f\"Mismatch between number of items in audio batch ({audio.size(0)})\",\n f\" and in metadata ({len(infos)})\"\n )\n\n latents = self._get_latents(audio)\n\n # prepare attributes\n if JascoCondConst.CRD.value in self.cfg.conditioners:\n null_chord_idx = self.cfg.conditioners.chords.chords_emb.card\n else:\n null_chord_idx = -1\n attributes = [info.to_condition_attributes() for info in infos]\n if self.model.cfg_dropout is not None:\n attributes = self.model.cfg_dropout(samples=attributes,\n cond_types=[\"wav\", \"text\", \"symbolic\"],\n null_chord_idx=null_chord_idx)\n attributes = self.model.att_dropout(attributes)\n tokenized = self.model.condition_provider.tokenize(attributes)\n\n with self.autocast:\n condition_tensors = self.model.condition_provider(tokenized)\n\n # create a padding mask to hold valid vs invalid positions\n padding_mask = torch.ones_like(latents, dtype=torch.bool, device=latents.device)\n\n return condition_tensors, latents, padding_mask\n\n def _normalized_latents(self, latents: torch.Tensor) -> torch.Tensor:\n \"\"\"Normalize latents.\"\"\"\n return (latents - self.latent_mean) / self.latent_std\n\n def _unnormalized_latents(self, latents: torch.Tensor) -> torch.Tensor:\n \"\"\"Unnormalize latents.\"\"\"\n return (latents * self.latent_std) + self.latent_mean\n\n def _z(self, z_0: torch.Tensor, z_1: torch.Tensor, t: torch.Tensor, sigma_min: float = 1e-5) -> torch.Tensor:\n \"\"\"Interpolate data and prior.\"\"\"\n return (1 - (1 - sigma_min) * t) * z_0 + t * z_1\n\n def _vector_field(self, z_0: torch.Tensor, z_1: torch.Tensor, sigma_min: float = 1e-5) -> torch.Tensor:\n \"\"\"Compute the GT vector field.\n sigma_min is a small value to avoid numerical instabilities.\"\"\"\n return z_1 - (1 - sigma_min) * z_0\n\n def _compute_loss(self, t: torch.Tensor, v_theta: torch.Tensor, v: torch.Tensor) -> torch.Tensor:\n \"\"\"Compute the loss.\"\"\"\n loss_func = self.cfg.get('loss_func', 'increasing')\n if loss_func == 'uniform':\n scales = 1\n elif loss_func == 'increasing':\n scales = 1 + t # type: ignore\n elif loss_func == 'decreasing':\n scales = 2 - t # type: ignore\n else:\n raise ValueError('unsupported loss_func was passed in config')\n return (scales * self.mse(v_theta, v)).mean()\n\n def run_step(self, idx: int, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]], metrics: dict) -> dict:\n \"\"\"Perform one training or valid step on a given batch.\"\"\"\n\n condition_tensors, latents, padding_mask = self._prepare_latents_and_attributes(batch)\n\n self.deadlock_detect.update('tokens_and_conditions')\n\n B, T, D = latents.shape\n device = self.device\n\n # normalize latents\n z_1 = self._normalized_latents(latents)\n\n # sample the N(0,1) prior\n z_0 = torch.randn(B, T, D, device=device)\n\n # random time parameter, between 0 to 1\n t = torch.rand((B, 1, 1), device=device)\n\n # interpolate data and prior\n z = self._z(z_0, z_1, t)\n\n # compute the GT vector field\n v = self._vector_field(z_0, z_1)\n\n with self.autocast:\n v_theta = self.model(latents=z,\n t=t,\n conditions=[],\n condition_tensors=condition_tensors)\n\n loss = self._compute_loss(t, v_theta, v)\n unscaled_loss = loss.clone()\n\n self.deadlock_detect.update('loss')\n\n if self.is_training:\n metrics['lr'] = self.optimizer.param_groups[0]['lr']\n if self.scaler is not None:\n loss = self.scaler.scale(loss)\n self.deadlock_detect.update('scale')\n if self.cfg.fsdp.use:\n loss.backward()\n flashy.distrib.average_tensors(self.model.buffers())\n elif self.cfg.optim.eager_sync:\n with flashy.distrib.eager_sync_model(self.model):\n loss.backward()\n else:\n # this should always be slower but can be useful\n # for weird use cases like multiple backwards.\n loss.backward()\n flashy.distrib.sync_model(self.model)\n self.deadlock_detect.update('backward')\n\n if self.scaler is not None:\n self.scaler.unscale_(self.optimizer)\n if self.cfg.optim.max_norm:\n if self.cfg.fsdp.use:\n metrics['grad_norm'] = self.model.clip_grad_norm_(self.cfg.optim.max_norm) # type: ignore\n else:\n metrics['grad_norm'] = torch.nn.utils.clip_grad_norm_(\n self.model.parameters(), self.cfg.optim.max_norm\n )\n if self.scaler is None:\n self.optimizer.step()\n else:\n self.scaler.step(self.optimizer)\n self.scaler.update()\n if self.lr_scheduler:\n self.lr_scheduler.step()\n self.optimizer.zero_grad()\n self.deadlock_detect.update('optim')\n if self.scaler is not None:\n scale = self.scaler.get_scale()\n metrics['grad_scale'] = scale\n if not loss.isfinite().all():\n raise RuntimeError(\"Model probably diverged.\")\n\n metrics['loss'] = unscaled_loss\n\n return metrics\n\n def _decode_latents(self, latents):\n return self.compression_model.model.decoder(latents.permute(0, 2, 1))\n\n @torch.no_grad()\n def run_generate_step(self, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]],\n gen_duration: float, prompt_duration: tp.Optional[float] = None,\n remove_text_conditioning: bool = False,\n **generation_params) -> dict:\n \"\"\"Run generate step on a batch of optional audio tensor and corresponding attributes.\n\n Args:\n batch (tuple[torch.Tensor, list[SegmentWithAttributes]]):\n use_prompt (bool): Whether to do audio continuation generation with prompt from audio batch.\n gen_duration (float): Target audio duration for the generation.\n prompt_duration (float, optional): Duration for the audio prompt to use for continuation.\n remove_text_conditioning (bool, optional): Whether to remove the prompt from the generated audio.\n generation_params: Additional generation parameters.\n Returns:\n gen_outputs (dict): Generation outputs, consisting in audio, audio tokens from both the generation\n and the prompt along with additional information.\n \"\"\"\n bench_start = time.time()\n audio, meta = batch\n assert audio.size(0) == len(meta), (\n f\"Mismatch between number of items in audio batch ({audio.size(0)})\",\n f\" and in metadata ({len(meta)})\"\n )\n # prepare attributes\n attributes = [x.to_condition_attributes() for x in meta]\n\n # prepare audio prompt\n if prompt_duration is None:\n prompt_audio = None\n else:\n assert prompt_duration < gen_duration, \"Prompt duration must be lower than target generation duration\"\n prompt_audio_frames = int(prompt_duration * self.compression_model.sample_rate)\n prompt_audio = audio[..., :prompt_audio_frames]\n\n # get audio tokens from compression model\n if prompt_audio is None or prompt_audio.nelement() == 0:\n num_samples = len(attributes)\n prompt_tokens = None\n else:\n num_samples = None\n prompt_audio = prompt_audio.to(self.device)\n prompt_tokens, scale = self.compression_model.encode(prompt_audio)\n assert scale is None, \"Compression model in MusicGen should not require rescaling.\"\n\n # generate by sampling from the LM\n with self.autocast:\n total_gen_len = math.ceil(gen_duration * self.compression_model.frame_rate)\n gen_latents = self.model.generate(\n prompt_tokens, attributes, max_gen_len=total_gen_len,\n num_samples=num_samples, **self.generation_params)\n\n # generate audio from latents\n assert gen_latents.dim() == 3 # [B, T, D]\n\n # unnormalize latents\n gen_latents = self._unnormalized_latents(gen_latents)\n gen_audio = self._decode_latents(gen_latents)\n\n bench_end = time.time()\n gen_outputs = {\n 'rtf': (bench_end - bench_start) / gen_duration,\n 'ref_audio': audio,\n 'gen_audio': gen_audio,\n 'gen_tokens': gen_latents,\n 'prompt_audio': prompt_audio,\n 'prompt_tokens': prompt_tokens,\n }\n return gen_outputs\n" }, { "path": "audiocraft/solvers/magnet.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom omegaconf import DictConfig\nfrom . import builders, musicgen\nfrom einops import rearrange\nfrom torch.nn import functional as F\nfrom ..modules.conditioners import SegmentWithAttributes\n\nimport torch\nimport numpy as np\nimport random\nimport typing as tp\nimport math\nimport flashy\n\n\nclass MagnetSolver(musicgen.MusicGenSolver):\n \"\"\"Solver for MAGNeT - Masked Audio Generation using\n a single Non-autoregressive Transformer https://arxiv.org/abs/2401.04577.\n \"\"\"\n def __init__(self, cfg: DictConfig):\n super().__init__(cfg)\n\n # initialize generation parameters by config\n self.generation_params = {\n 'use_sampling': self.cfg.generate.lm.use_sampling,\n 'temp': self.cfg.generate.lm.temp,\n 'top_k': self.cfg.generate.lm.top_k,\n 'top_p': self.cfg.generate.lm.top_p,\n 'max_cfg_coef': self.cfg.generate.lm.max_cfg_coef,\n 'min_cfg_coef': self.cfg.generate.lm.min_cfg_coef,\n 'decoding_steps': list(self.cfg.generate.lm.decoding_steps),\n 'anneal_temp': self.cfg.generate.lm.anneal_temp,\n 'span_scoring': self.cfg.generate.lm.span_scoring,\n 'span_arrangement': self.cfg.generate.lm.span_arrangement\n }\n\n sequence_len = int(cfg.dataset.segment_duration * self.compression_model.frame_rate)\n self.mean_maskrate_to_u = torch.tensor(self._calc_mean_maskrate_to_u_LUT(sequence_len), device=self.device)\n self.ce_per_codebook = [torch.log(torch.tensor(self.compression_model.cardinality, device=self.device))\n for _ in range(cfg.transformer_lm.n_q)]\n\n def build_model(self) -> None:\n self.cfg.transformer_lm.segment_duration = self.cfg.dataset.segment_duration\n self.cfg.transformer_lm.span_len = self.cfg.masking.span_len\n assert self.cfg.efficient_attention_backend == \"xformers\", \"MAGNeT v1 models support only xformers backend.\"\n super().build_model()\n\n def _calc_mean_maskrate_to_u_LUT(self, T: int):\n \"\"\" Create a Look Up Table (LUT) transforming a discrete masking percentage m in 0,1,...,100 to u,\n the number of overlapping spans of length L to place s.t. the masking rate is approximately m/float(100).\n It first creates the inverse transformation, of the masking rate as function of u,\n using the expression choose(T - L, u) / choose(T, u), where L is the atomic span length used\n during masking. See https://arxiv.org/abs/2401.04577,\n appendix C, for the mean mask rate derivation.\n\n We leverage the fact that:\n choose(T - L, u) / choose(T, u) = Prod_{j = 0}^{u - 1}((T - L - j)/(T - j))\n in the provided implementation, in order to avoid overflow.\n Args:\n T (float): Sequence length.\n Returns:\n (List) A LUT transforming m in 0,1,...,100 to u,\n s.t. the masking rate of the span-L mask is approximately m/float(100).\n \"\"\"\n\n L = self.cfg.masking.span_len\n\n u2mean = [0.0] # mean mask rate is 0.0 for u = 0\n v = (T - L) / float(T)\n for u in range(1, T):\n u2mean.append(1 - v)\n v *= (T - L - u) / (T - u) # Overflow-safe implementation of choose(T - L, u) / choose(T, u).\n\n mean2u = []\n for maskperc in range(101):\n maskrate = maskperc / float(100)\n u = int(np.searchsorted(u2mean, maskrate))\n mean2u.append(u)\n\n return mean2u\n\n def _non_spans_mask(self, mask_probs: torch.Tensor, B: int, T: int, device: torch.device) -> torch.Tensor:\n \"\"\" Construct a boolean mask of shape [B, T, 1], with masking rates defined by mask_probs.\n The masked tokens are singletons, placed uniformly at random.\n Args:\n mask_probs (torch.Tensor): The desired masking rate per sample, of shape [B,]\n B (int): Batch size.\n T (int): Sequence length.\n device (torch.device): device of the output tensor\n Returns:\n (torch.Tensor): A mask of shape [B, T]\n \"\"\"\n num_token_masked = (T * mask_probs).round().clamp(min=1)\n batch_randperm = torch.rand((B, T), device=device).argsort(dim=-1)\n return batch_randperm < rearrange(num_token_masked, 'b -> b 1')\n\n def _spans_mask(self, mask_probs: torch.Tensor, B: int, T: int, device: torch.device) -> torch.Tensor:\n \"\"\" Construct a spans mask with masking rates defined by mask_probs,\n where the atomic span length ( > 1 ) is defined by cfg.masking.span_len.\n Args:\n mask_probs (torch.Tensor): The desired masking rate per sample, of shape [B,]\n B (int): Batch size.\n T (int): Sequence length.\n device (torch.device): device of the output tensor\n Returns:\n (torch.Tensor): A spans mask of shape [B, T]\n \"\"\"\n rounded_probs = torch.round(100 * mask_probs).long()\n k = self.mean_maskrate_to_u[rounded_probs].clamp(min=1) # k is the number of span starts\n\n # sample random span starts\n batch_randperm = torch.rand((B, T), device=device).argsort(dim=-1)\n mask = batch_randperm < rearrange(k, 'b -> b 1')\n B, T = mask.shape\n shifted_mask = mask.clone()\n for _ in range(self.cfg.masking.span_len - 1):\n shifted_mask = torch.concat((torch.full((B, 1), False, device=device), shifted_mask[:, :-1]), dim=1)\n mask = torch.logical_or(mask, shifted_mask)\n\n return mask\n\n def _get_mask(self, mask_probs: torch.Tensor, B: int, T: int, device: torch.device) -> torch.Tensor:\n \"\"\" Construct a boolean mask with masking rates defined by mask_probs, and atomic\n span length defined by cfg.masking.span_len.\n Args:\n mask_probs (torch.Tensor): The desired masking rate per sample, of shape [B,]\n B (int): Batch size.\n T (int): Sequence length.\n device (torch.device): device of the output tensor\n Returns:\n (torch.Tensor): A boolean tensor of shape [B, T]\n \"\"\"\n if self.cfg.masking.span_len <= 1:\n return self._non_spans_mask(mask_probs, B, T, device)\n\n return self._spans_mask(mask_probs, B, T, device)\n\n def _compute_cross_entropy_magnet(self, logits: torch.Tensor,\n targets: torch.Tensor, mask: torch.Tensor, stage: torch.Tensor) -> torch.Tensor:\n \"\"\" Compute cross entropy between multi-codebook targets and model's logits.\n The cross entropy is computed only on a specific codebook, defined by the stage argument.\n Valid timesteps for each codebook are pulled from the mask, where invalid\n timesteps are set to 0.\n\n Args:\n logits (torch.Tensor): Model's logits of shape [B, K, T, card].\n targets (torch.Tensor): Target codes, of shape [B, K, T].\n mask (torch.Tensor): Mask for valid target codes, of shape [B, K, T].\n stage (torch.Tensor): The codebook (idx) that is being optimized, as a scalar tensor.\n Returns:\n ce (torch.Tensor): Cross entropy of the codebook that is being optimized.\n \"\"\"\n assert logits.shape[:-1] == targets.shape\n assert mask.shape == targets.shape\n ce = torch.zeros([], device=targets.device)\n logits_k = logits[:, stage, ...].contiguous().view(-1, logits.size(-1)) # [B x T, card]\n targets_k = targets[:, stage, ...].contiguous().view(-1) # [B x T]\n mask_k = mask[:, stage, ...].contiguous().view(-1) # [B x T]\n\n IGNORE_IDX = -1\n targets_k[~mask_k] = IGNORE_IDX\n q_ce = F.cross_entropy(logits_k, targets_k, ignore_index=IGNORE_IDX)\n\n ce += q_ce\n return ce\n\n def run_step(self, idx: int, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]], metrics: dict) -> dict:\n \"\"\"Perform one training or valid step on a given batch.\"\"\"\n check_synchronization_points = idx == 1 and self.device == 'cuda'\n\n condition_tensors, audio_tokens, padding_mask = self._prepare_tokens_and_attributes(\n batch, check_synchronization_points)\n\n self.deadlock_detect.update('tokens_and_conditions')\n\n if check_synchronization_points:\n torch.cuda.set_sync_debug_mode('warn')\n\n B, K, T = audio_tokens.shape\n device = self.device\n\n # Choose the stage (codebook idx) for update, uniformly at random.\n stage_ = random.randint(0, K - 1)\n stage = torch.full((1, ), stage_, device=device)\n\n # masking\n rand_time = torch.zeros((B,), device=device).float().uniform_(0, 1)\n rand_mask_probs = torch.cos(rand_time * math.pi * 0.5)\n\n # stage mask\n stage_mask = self._get_mask(rand_mask_probs, B, T, device) # [B, T]\n stage_mask = stage_mask.unsqueeze(1) # [B, 1, T]\n\n # Keep all preceding codebooks.\n mask = torch.full((B, K, T), False, device=device)\n mask[:, stage, :] = stage_mask\n\n # Mask all codebooks larger than stage_\n mask_id = self.model.special_token_id\n mask[:, (stage_+1):, :] = torch.full((B, K - stage_ - 1, T), True, device=device)\n input_tokens = torch.where(mask, mask_id, audio_tokens)\n\n # Take loss only on the chosen stage, and only on the masked tokens.\n loss_mask = torch.full((B, K, T), False, device=device)\n loss_mask[:, stage, :] = stage_mask\n\n with self.autocast:\n model_output = self.model.compute_predictions(input_tokens, [], condition_tensors, stage=stage_)\n logits = model_output.logits\n loss_mask &= padding_mask\n ce = self._compute_cross_entropy_magnet(logits, audio_tokens, loss_mask, stage)\n loss = ce\n self.deadlock_detect.update('loss')\n\n if check_synchronization_points:\n torch.cuda.set_sync_debug_mode('default')\n\n if self.is_training:\n metrics['lr'] = self.optimizer.param_groups[0]['lr']\n if self.scaler is not None:\n loss = self.scaler.scale(loss)\n self.deadlock_detect.update('scale')\n if self.cfg.fsdp.use:\n loss.backward()\n flashy.distrib.average_tensors(self.model.buffers())\n elif self.cfg.optim.eager_sync:\n with flashy.distrib.eager_sync_model(self.model):\n loss.backward()\n else:\n # this should always be slower but can be useful\n # for weird use cases like multiple backwards.\n loss.backward()\n flashy.distrib.sync_model(self.model)\n self.deadlock_detect.update('backward')\n\n if self.scaler is not None:\n self.scaler.unscale_(self.optimizer)\n if self.cfg.optim.max_norm:\n if self.cfg.fsdp.use:\n metrics['grad_norm'] = self.model.clip_grad_norm_(self.cfg.optim.max_norm) # type: ignore\n else:\n metrics['grad_norm'] = torch.nn.utils.clip_grad_norm_(\n self.model.parameters(), self.cfg.optim.max_norm\n )\n if self.scaler is None:\n self.optimizer.step()\n else:\n self.scaler.step(self.optimizer)\n self.scaler.update()\n if self.lr_scheduler:\n self.lr_scheduler.step()\n self.optimizer.zero_grad()\n self.deadlock_detect.update('optim')\n if self.scaler is not None:\n scale = self.scaler.get_scale()\n metrics['grad_scale'] = scale\n if not loss.isfinite().all():\n raise RuntimeError(\"Model probably diverged.\")\n\n metrics['ce'] = ce\n metrics['ppl'] = torch.exp(ce)\n\n return metrics\n\n\nclass AudioMagnetSolver(MagnetSolver):\n \"\"\"Solver for audio-MAGNeT. A MAGNeT model for sound generation.\n\n More information can be found in the MAGNeT model card.\n \"\"\"\n DATASET_TYPE: builders.DatasetType = builders.DatasetType.SOUND\n" }, { "path": "audiocraft/solvers/musicgen.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom pathlib import Path\nimport time\nimport typing as tp\nimport warnings\n\nimport flashy\nimport math\nimport omegaconf\nimport torch\nfrom torch.nn import functional as F\n\nfrom . import base, builders\nfrom .compression import CompressionSolver\nfrom .. import metrics as eval_metrics\nfrom .. import models\nfrom ..data.audio_dataset import AudioDataset\nfrom ..data.music_dataset import MusicDataset, MusicInfo, AudioInfo\nfrom ..data.audio_utils import normalize_audio\nfrom ..modules.conditioners import JointEmbedCondition, SegmentWithAttributes, WavCondition, \\\n StyleConditioner, _drop_description_condition\nfrom ..utils.cache import CachedBatchWriter, CachedBatchLoader\nfrom ..utils.samples.manager import SampleManager\nfrom ..utils.utils import get_dataset_from_loader, is_jsonable, warn_once, model_hash\n\n\nclass MusicGenSolver(base.StandardSolver):\n \"\"\"Solver for MusicGen training task.\n\n Used in: https://arxiv.org/abs/2306.05284\n \"\"\"\n DATASET_TYPE: builders.DatasetType = builders.DatasetType.MUSIC\n\n def __init__(self, cfg: omegaconf.DictConfig):\n super().__init__(cfg)\n # easier access to sampling parameters\n self.generation_params = {\n 'use_sampling': self.cfg.generate.lm.use_sampling,\n 'temp': self.cfg.generate.lm.temp,\n 'top_k': self.cfg.generate.lm.top_k,\n 'top_p': self.cfg.generate.lm.top_p,\n }\n self._best_metric_name: tp.Optional[str] = 'ce'\n\n self._cached_batch_writer = None\n self._cached_batch_loader = None\n if cfg.cache.path:\n if cfg.cache.write:\n self._cached_batch_writer = CachedBatchWriter(Path(cfg.cache.path))\n if self.cfg.cache.write_num_shards:\n self.logger.warning(\"Multiple shard cache, best_metric_name will be set to None.\")\n self._best_metric_name = None\n else:\n self._cached_batch_loader = CachedBatchLoader(\n Path(cfg.cache.path), cfg.dataset.batch_size, cfg.dataset.num_workers,\n min_length=self.cfg.optim.updates_per_epoch or 1)\n self.dataloaders['original_train'] = self.dataloaders['train']\n self.dataloaders['train'] = self._cached_batch_loader # type: ignore\n\n @staticmethod\n def get_eval_solver_from_sig(sig: str, dtype: tp.Optional[str] = None,\n device: tp.Optional[str] = None, autocast: bool = True,\n batch_size: tp.Optional[int] = None,\n override_cfg: tp.Optional[tp.Union[dict, omegaconf.DictConfig]] = None,\n **kwargs):\n \"\"\"Mostly a convenience function around magma.train.get_solver_from_sig,\n populating all the proper param, deactivating EMA, FSDP, loading the best state,\n basically all you need to get a solver ready to \"play\" with in single GPU mode\n and with minimal memory overhead.\n\n Args:\n sig (str): signature to load.\n dtype (str or None): potential dtype, as a string, i.e. 'float16'.\n device (str or None): potential device, as a string, i.e. 'cuda'.\n override_cfg (dict or omegaconf.DictConfig or None): potential device, as a string, i.e. 'cuda'.\n \"\"\"\n from audiocraft import train\n our_override_cfg: tp.Dict[str, tp.Any] = {'optim': {'ema': {'use': False}}}\n our_override_cfg['autocast'] = autocast\n if dtype is not None:\n our_override_cfg['dtype'] = dtype\n if device is not None:\n our_override_cfg['device'] = device\n if batch_size is not None:\n our_override_cfg['dataset'] = {'batch_size': batch_size}\n if override_cfg is None:\n override_cfg = {}\n override_cfg = omegaconf.OmegaConf.merge(\n omegaconf.DictConfig(override_cfg), omegaconf.DictConfig(our_override_cfg)) # type: ignore\n solver = train.get_solver_from_sig(\n sig, override_cfg=override_cfg,\n load_best=True, disable_fsdp=True,\n ignore_state_keys=['optimizer', 'ema'], **kwargs)\n solver.model.eval()\n return solver\n\n def get_formatter(self, stage_name: str) -> flashy.Formatter:\n return flashy.Formatter({\n 'lr': '.2E',\n 'ce': '.3f',\n 'ppl': '.3f',\n 'grad_norm': '.3E',\n }, exclude_keys=['ce_q*', 'ppl_q*'])\n\n @property\n def best_metric_name(self) -> tp.Optional[str]:\n return self._best_metric_name\n\n def initialize_optimization(self) -> None:\n if self.cfg.fsdp.use:\n assert not self.cfg.autocast, \"Cannot use autocast with fsdp\"\n self.model = self.wrap_with_fsdp(self.model)\n self.register_ema('model')\n # initialize optimization\n self.optimizer = builders.get_optimizer(builders.get_optim_parameter_groups(self.model), self.cfg.optim)\n self.lr_scheduler = builders.get_lr_scheduler(self.optimizer, self.cfg.schedule, self.total_updates)\n self.register_stateful('model', 'optimizer', 'lr_scheduler')\n self.register_best_state('model')\n self.autocast_dtype = {\n 'float16': torch.float16, 'bfloat16': torch.bfloat16\n }[self.cfg.autocast_dtype]\n self.scaler: tp.Optional[torch.cuda.amp.GradScaler] = None\n if self.cfg.fsdp.use:\n need_scaler = self.cfg.fsdp.param_dtype == 'float16'\n else:\n need_scaler = self.cfg.autocast and self.autocast_dtype is torch.float16\n if need_scaler:\n if self.cfg.fsdp.use:\n from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler\n self.scaler = ShardedGradScaler() # type: ignore\n else:\n self.scaler = torch.cuda.amp.GradScaler()\n self.register_stateful('scaler')\n\n def build_model(self) -> None:\n \"\"\"Instantiate models and optimizer.\"\"\"\n # we can potentially not use all quantizers with which the EnCodec model was trained\n # (e.g. we trained the model with quantizers dropout)\n self.compression_model = CompressionSolver.wrapped_model_from_checkpoint(\n self.cfg, self.cfg.compression_model_checkpoint, device=self.device)\n assert self.compression_model.sample_rate == self.cfg.sample_rate, (\n f\"Compression model sample rate is {self.compression_model.sample_rate} but \"\n f\"Solver sample rate is {self.cfg.sample_rate}.\"\n )\n # ensure we have matching configuration between LM and compression model\n assert self.cfg.transformer_lm.card == self.compression_model.cardinality, (\n \"Cardinalities of the LM and compression model don't match: \",\n f\"LM cardinality is {self.cfg.transformer_lm.card} vs \",\n f\"compression model cardinality is {self.compression_model.cardinality}\"\n )\n assert self.cfg.transformer_lm.n_q == self.compression_model.num_codebooks, (\n \"Numbers of codebooks of the LM and compression models don't match: \",\n f\"LM number of codebooks is {self.cfg.transformer_lm.n_q} vs \",\n f\"compression model numer of codebooks is {self.compression_model.num_codebooks}\"\n )\n self.logger.info(\"Compression model has %d codebooks with %d cardinality, and a framerate of %d\",\n self.compression_model.num_codebooks, self.compression_model.cardinality,\n self.compression_model.frame_rate)\n # instantiate LM model\n self.model: tp.Union[models.LMModel, models.FlowMatchingModel] = models.builders.get_lm_model(\n self.cfg).to(self.device)\n\n # initialize optimization\n self.initialize_optimization()\n\n def build_dataloaders(self) -> None:\n \"\"\"Instantiate audio dataloaders for each stage.\"\"\"\n self.dataloaders = builders.get_audio_datasets(self.cfg, dataset_type=self.DATASET_TYPE)\n\n def show(self) -> None:\n \"\"\"Show the compression model and LM model.\"\"\"\n self.logger.info(\"Compression model:\")\n self.log_model_summary(self.compression_model)\n self.logger.info(\"LM model:\")\n self.log_model_summary(self.model)\n\n def load_state_dict(self, state: dict) -> None:\n if 'condition_provider' in state:\n model_state = state['model']\n condition_provider_state = state.pop('condition_provider')\n prefix = 'condition_provider.'\n for key, value in condition_provider_state.items():\n key = prefix + key\n assert key not in model_state\n model_state[key] = value\n if 'compression_model' in state:\n # We used to store the `compression_model` state in the checkpoint, however\n # this is in general not needed, as the compression model should always be readable\n # from the original `cfg.compression_model_checkpoint` location.\n compression_model_state = state.pop('compression_model')\n before_hash = model_hash(self.compression_model)\n self.compression_model.load_state_dict(compression_model_state)\n after_hash = model_hash(self.compression_model)\n if before_hash != after_hash:\n raise RuntimeError(\n \"The compression model state inside the checkpoint is different\"\n \" from the one obtained from compression_model_checkpoint...\"\n \"We do not support altering the compression model inside the LM \"\n \"checkpoint as parts of the code, in particular for running eval post-training \"\n \"will use the compression_model_checkpoint as the source of truth.\")\n\n super().load_state_dict(state)\n\n def load_from_pretrained(self, name: str):\n # TODO: support native HF versions of MusicGen.\n lm_pkg = models.loaders.load_lm_model_ckpt(name)\n state: dict = {\n 'best_state': {\n 'model': lm_pkg['best_state'],\n },\n }\n return state\n\n def _compute_cross_entropy(\n self, logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor\n ) -> tp.Tuple[torch.Tensor, tp.List[torch.Tensor]]:\n \"\"\"Compute cross entropy between multi-codebook targets and model's logits.\n The cross entropy is computed per codebook to provide codebook-level cross entropy.\n Valid timesteps for each of the codebook are pulled from the mask, where invalid\n timesteps are set to 0.\n\n Args:\n logits (torch.Tensor): Model's logits of shape [B, K, T, card].\n targets (torch.Tensor): Target codes, of shape [B, K, T].\n mask (torch.Tensor): Mask for valid target codes, of shape [B, K, T].\n Returns:\n ce (torch.Tensor): Cross entropy averaged over the codebooks\n ce_per_codebook (list of torch.Tensor): Cross entropy per codebook (detached).\n \"\"\"\n B, K, T = targets.shape\n assert logits.shape[:-1] == targets.shape\n assert mask.shape == targets.shape\n ce = torch.zeros([], device=targets.device)\n ce_per_codebook: tp.List[torch.Tensor] = []\n for k in range(K):\n logits_k = logits[:, k, ...].contiguous().view(-1, logits.size(-1)) # [B x T, card]\n targets_k = targets[:, k, ...].contiguous().view(-1) # [B x T]\n mask_k = mask[:, k, ...].contiguous().view(-1) # [B x T]\n ce_targets = targets_k[mask_k]\n ce_logits = logits_k[mask_k]\n q_ce = F.cross_entropy(ce_logits, ce_targets)\n ce += q_ce\n ce_per_codebook.append(q_ce.detach())\n # average cross entropy across codebooks\n ce = ce / K\n return ce, ce_per_codebook\n\n def _get_audio_tokens(self, audio: torch.Tensor):\n with torch.no_grad():\n audio_tokens, scale = self.compression_model.encode(audio)\n assert scale is None, \"Scaled compression model not supported with LM.\"\n return audio_tokens\n\n def _prepare_tokens_and_attributes(\n self, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]],\n check_synchronization_points: bool = False\n ) -> tp.Tuple[dict, torch.Tensor, torch.Tensor]:\n \"\"\"Prepare input batchs for language model training.\n\n Args:\n batch (tuple[torch.Tensor, list[SegmentWithAttributes]]): Input batch with audio tensor of shape [B, C, T]\n and corresponding metadata as SegmentWithAttributes (with B items).\n check_synchronization_points (bool): Whether to check for synchronization points slowing down training.\n Returns:\n Condition tensors (dict[str, any]): Preprocessed condition attributes.\n Tokens (torch.Tensor): Audio tokens from compression model, of shape [B, K, T_s],\n with B the batch size, K the number of codebooks, T_s the token timesteps.\n Padding mask (torch.Tensor): Mask with valid positions in the tokens tensor, of shape [B, K, T_s].\n \"\"\"\n if self.model.training:\n warnings.warn(\n \"Up to version 1.0.1, the _prepare_tokens_and_attributes was evaluated with `torch.no_grad()`. \"\n \"This is inconsistent with how model were trained in the MusicGen paper. We removed the \"\n \"`torch.no_grad()` in version 1.1.0. Small changes to the final performance are expected. \"\n \"Really sorry about that.\")\n if self._cached_batch_loader is None or self.current_stage != \"train\":\n audio, infos = batch\n audio = audio.to(self.device)\n audio_tokens = None\n assert audio.size(0) == len(infos), (\n f\"Mismatch between number of items in audio batch ({audio.size(0)})\",\n f\" and in metadata ({len(infos)})\"\n )\n else:\n audio = None\n # In that case the batch will be a tuple coming from the _cached_batch_writer bit below.\n infos, = batch # type: ignore\n assert all([isinstance(info, AudioInfo) for info in infos])\n assert all([info.audio_tokens is not None for info in infos]) # type: ignore\n audio_tokens = torch.stack([info.audio_tokens for info in infos]).to(self.device) # type: ignore\n audio_tokens = audio_tokens.long()\n for info in infos:\n if isinstance(info, MusicInfo):\n # Careful here, if you want to use this condition_wav (e.b. chroma conditioning),\n # then you must be using the chroma cache! otherwise the code will try\n # to use this segment and fail (by that I mean you will see NaN everywhere).\n info.self_wav = WavCondition(\n torch.full([1, info.channels, info.total_frames], float('NaN')),\n length=torch.tensor([info.n_frames]),\n sample_rate=[info.sample_rate],\n path=[info.meta.path],\n seek_time=[info.seek_time])\n dataset = get_dataset_from_loader(self.dataloaders['original_train'])\n assert isinstance(dataset, MusicDataset), type(dataset)\n if dataset.paraphraser is not None and info.description is not None:\n # Hackingly reapplying paraphraser when using cache.\n info.description = dataset.paraphraser.sample_paraphrase(\n info.meta.path, info.description)\n # prepare attributes\n attributes = [info.to_condition_attributes() for info in infos]\n attributes = self.model.cfg_dropout(attributes)\n attributes = self.model.att_dropout(attributes)\n tokenized = self.model.condition_provider.tokenize(attributes)\n\n # Now we should be synchronization free.\n if self.device == \"cuda\" and check_synchronization_points:\n torch.cuda.set_sync_debug_mode(\"warn\")\n\n if audio_tokens is None:\n audio_tokens = self._get_audio_tokens(audio)\n\n with self.autocast:\n condition_tensors = self.model.condition_provider(tokenized)\n\n # create a padding mask to hold valid vs invalid positions\n padding_mask = torch.ones_like(audio_tokens, dtype=torch.bool, device=audio_tokens.device)\n # replace encodec tokens from padded audio with special_token_id\n if self.cfg.tokens.padding_with_special_token:\n audio_tokens = audio_tokens.clone()\n padding_mask = padding_mask.clone()\n token_sample_rate = self.compression_model.frame_rate\n B, K, T_s = audio_tokens.shape\n for i in range(B):\n n_samples = infos[i].n_frames\n audio_sample_rate = infos[i].sample_rate\n # take the last token generated from actual audio frames (non-padded audio)\n valid_tokens = math.floor(float(n_samples) / audio_sample_rate * token_sample_rate)\n audio_tokens[i, :, valid_tokens:] = self.model.special_token_id\n padding_mask[i, :, valid_tokens:] = 0\n\n if self.device == \"cuda\" and check_synchronization_points:\n torch.cuda.set_sync_debug_mode(\"default\")\n\n if self._cached_batch_writer is not None and self.current_stage == 'train':\n assert self._cached_batch_loader is None\n assert audio_tokens is not None\n for info, one_audio_tokens in zip(infos, audio_tokens):\n assert isinstance(info, AudioInfo)\n if isinstance(info, MusicInfo):\n assert not info.joint_embed, \"joint_embed and cache not supported yet.\"\n info.self_wav = None\n assert one_audio_tokens.max() < 2**15, one_audio_tokens.max().item()\n info.audio_tokens = one_audio_tokens.short().cpu()\n self._cached_batch_writer.save(infos)\n\n return condition_tensors, audio_tokens, padding_mask\n\n def run_step(self, idx: int, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]], metrics: dict) -> dict:\n \"\"\"Perform one training or valid step on a given batch.\"\"\"\n check_synchronization_points = idx == 1 and self.device == 'cuda'\n\n condition_tensors, audio_tokens, padding_mask = self._prepare_tokens_and_attributes(\n batch, check_synchronization_points)\n\n self.deadlock_detect.update('tokens_and_conditions')\n\n if check_synchronization_points:\n torch.cuda.set_sync_debug_mode('warn')\n\n with self.autocast:\n style_mask = None\n if hasattr(self.model.condition_provider.conditioners, 'self_wav'):\n if isinstance(self.model.condition_provider.conditioners.self_wav, StyleConditioner):\n style_mask = self.model.condition_provider.conditioners.self_wav.mask\n\n model_output = self.model.compute_predictions(audio_tokens, [], condition_tensors) # type: ignore\n logits = model_output.logits\n if style_mask is not None:\n mask = padding_mask & model_output.mask & style_mask\n else:\n mask = padding_mask & model_output.mask\n ce, ce_per_codebook = self._compute_cross_entropy(logits, audio_tokens, mask)\n loss = ce\n self.deadlock_detect.update('loss')\n\n if check_synchronization_points:\n torch.cuda.set_sync_debug_mode('default')\n\n if self.is_training:\n metrics['lr'] = self.optimizer.param_groups[0]['lr']\n if self.scaler is not None:\n loss = self.scaler.scale(loss)\n self.deadlock_detect.update('scale')\n if self.cfg.fsdp.use:\n loss.backward()\n flashy.distrib.average_tensors(self.model.buffers())\n elif self.cfg.optim.eager_sync:\n with flashy.distrib.eager_sync_model(self.model):\n loss.backward()\n else:\n # this should always be slower but can be useful\n # for weird use cases like multiple backwards.\n loss.backward()\n flashy.distrib.sync_model(self.model)\n self.deadlock_detect.update('backward')\n\n if self.scaler is not None:\n self.scaler.unscale_(self.optimizer)\n if self.cfg.optim.max_norm:\n if self.cfg.fsdp.use:\n metrics['grad_norm'] = self.model.clip_grad_norm_(self.cfg.optim.max_norm) # type: ignore\n else:\n metrics['grad_norm'] = torch.nn.utils.clip_grad_norm_(\n self.model.parameters(), self.cfg.optim.max_norm\n )\n if self.scaler is None:\n self.optimizer.step()\n else:\n self.scaler.step(self.optimizer)\n self.scaler.update()\n if self.lr_scheduler:\n self.lr_scheduler.step()\n self.optimizer.zero_grad()\n self.deadlock_detect.update('optim')\n if self.scaler is not None:\n scale = self.scaler.get_scale()\n metrics['grad_scale'] = scale\n if not loss.isfinite().all():\n raise RuntimeError(\"Model probably diverged.\")\n\n metrics['ce'] = ce\n metrics['ppl'] = torch.exp(ce)\n for k, ce_q in enumerate(ce_per_codebook):\n metrics[f'ce_q{k + 1}'] = ce_q\n metrics[f'ppl_q{k + 1}'] = torch.exp(ce_q)\n\n return metrics\n\n @torch.no_grad()\n def run_generate_step(self, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]],\n gen_duration: float, prompt_duration: tp.Optional[float] = None,\n remove_text_conditioning: bool = False,\n ) -> dict:\n \"\"\"Run generate step on a batch of optional audio tensor and corresponding attributes.\n\n Args:\n batch (tuple[torch.Tensor, list[SegmentWithAttributes]]):\n use_prompt (bool): Whether to do audio continuation generation with prompt from audio batch.\n gen_duration (float): Target audio duration for the generation.\n prompt_duration (float, optional): Duration for the audio prompt to use for continuation.\n Returns:\n gen_outputs (dict): Generation outputs, consisting in audio, audio tokens from both the generation\n and the prompt along with additional information.\n \"\"\"\n bench_start = time.time()\n audio, meta = batch\n assert audio.size(0) == len(meta), (\n f\"Mismatch between number of items in audio batch ({audio.size(0)})\",\n f\" and in metadata ({len(meta)})\"\n )\n # prepare attributes\n attributes = [x.to_condition_attributes() for x in meta]\n if remove_text_conditioning:\n attributes = _drop_description_condition(attributes)\n\n # prepare audio prompt\n if prompt_duration is None:\n prompt_audio = None\n else:\n assert prompt_duration < gen_duration, \"Prompt duration must be lower than target generation duration\"\n prompt_audio_frames = int(prompt_duration * self.compression_model.sample_rate)\n prompt_audio = audio[..., :prompt_audio_frames]\n\n # get audio tokens from compression model\n if prompt_audio is None or prompt_audio.nelement() == 0:\n num_samples = len(attributes)\n prompt_tokens = None\n else:\n num_samples = None\n prompt_audio = prompt_audio.to(self.device)\n prompt_tokens, scale = self.compression_model.encode(prompt_audio)\n assert scale is None, \"Compression model in MusicGen should not require rescaling.\"\n\n # generate by sampling from the LM\n with self.autocast:\n total_gen_len = math.ceil(gen_duration * self.compression_model.frame_rate)\n gen_tokens = self.model.generate(\n prompt_tokens, attributes, max_gen_len=total_gen_len,\n num_samples=num_samples, **self.generation_params)\n\n # generate audio from tokens\n assert gen_tokens.dim() == 3\n gen_audio = self.compression_model.decode(gen_tokens, None)\n\n bench_end = time.time()\n gen_outputs = {\n 'rtf': (bench_end - bench_start) / gen_duration,\n 'ref_audio': audio,\n 'gen_audio': gen_audio,\n 'gen_tokens': gen_tokens,\n 'prompt_audio': prompt_audio,\n 'prompt_tokens': prompt_tokens,\n }\n return gen_outputs\n\n def generate_audio(self) -> dict:\n \"\"\"Audio generation stage.\"\"\"\n generate_stage_name = f'{self.current_stage}'\n sample_manager = SampleManager(self.xp)\n self.logger.info(f\"Generating samples in {sample_manager.base_folder}\")\n loader = self.dataloaders['generate']\n updates = len(loader)\n lp = self.log_progress(generate_stage_name, loader, total=updates, updates=self.log_updates)\n\n dataset = get_dataset_from_loader(loader)\n dataset_duration = dataset.segment_duration\n assert dataset_duration is not None\n assert isinstance(dataset, AudioDataset)\n target_duration = self.cfg.generate.lm.gen_duration\n prompt_duration = self.cfg.generate.lm.prompt_duration\n if target_duration is None:\n target_duration = dataset_duration\n if prompt_duration is None:\n prompt_duration = dataset_duration / 4\n assert prompt_duration < dataset_duration, (\n f\"Specified prompt duration ({prompt_duration}s) is longer\",\n f\" than reference audio duration ({dataset_duration}s)\"\n )\n\n def get_hydrated_conditions(meta: tp.List[SegmentWithAttributes]):\n hydrated_conditions = []\n for sample in [x.to_condition_attributes() for x in meta]:\n cond_dict = {}\n for cond_type in sample.__annotations__.keys():\n for cond_key, cond_val in getattr(sample, cond_type).items():\n if cond_key not in self.model.condition_provider.conditioners.keys():\n continue\n if is_jsonable(cond_val):\n cond_dict[cond_key] = cond_val\n elif isinstance(cond_val, WavCondition):\n cond_dict[cond_key] = cond_val.path\n elif isinstance(cond_val, JointEmbedCondition):\n cond_dict[cond_key] = cond_val.text # only support text at inference for now\n else:\n # if we reached this point, it is not clear how to log the condition\n # so we just log the type.\n cond_dict[cond_key] = str(type(cond_val))\n continue\n hydrated_conditions.append(cond_dict)\n return hydrated_conditions\n\n metrics: dict = {}\n average = flashy.averager()\n for batch in lp:\n audio, meta = batch\n # metadata for sample manager\n hydrated_conditions = get_hydrated_conditions(meta)\n sample_generation_params = {\n **{f'classifier_free_guidance_{k}': v for k, v in self.cfg.classifier_free_guidance.items()},\n **self.generation_params\n }\n if self.cfg.generate.lm.unprompted_samples:\n if self.cfg.generate.lm.gen_gt_samples:\n # get the ground truth instead of generation\n self.logger.warn(\n \"Use ground truth instead of audio generation as generate.lm.gen_gt_samples=true\")\n gen_unprompted_audio = audio\n rtf = 1.\n else:\n gen_unprompted_outputs = self.run_generate_step(\n batch, gen_duration=target_duration, prompt_duration=None)\n gen_unprompted_audio = gen_unprompted_outputs['gen_audio'].cpu()\n rtf = gen_unprompted_outputs['rtf']\n sample_manager.add_samples(\n gen_unprompted_audio, self.epoch, hydrated_conditions,\n ground_truth_wavs=audio, generation_args=sample_generation_params)\n\n if self.cfg.generate.lm.prompted_samples:\n gen_outputs = self.run_generate_step(\n batch, gen_duration=target_duration, prompt_duration=prompt_duration)\n gen_audio = gen_outputs['gen_audio'].cpu()\n prompt_audio = gen_outputs['prompt_audio'].cpu()\n sample_manager.add_samples(\n gen_audio, self.epoch, hydrated_conditions,\n prompt_wavs=prompt_audio, ground_truth_wavs=audio,\n generation_args=sample_generation_params)\n if self.cfg.generate.lm.no_text_conditioning:\n gen_outputs = self.run_generate_step(\n batch, gen_duration=target_duration, prompt_duration=None,\n remove_text_conditioning=self.cfg.generate.lm.no_text_conditioning)\n gen_audio = gen_outputs['gen_audio'].cpu()\n rtf = gen_outputs['rtf']\n # Here, the prompt is the original audio provided for the style conditioning\n prompt_audio = gen_outputs['ref_audio'].cpu()\n sample_manager.add_samples(\n gen_audio, self.epoch, hydrated_conditions,\n prompt_wavs=prompt_audio, ground_truth_wavs=audio,\n generation_args=sample_generation_params)\n\n metrics['rtf'] = rtf\n metrics = average(metrics)\n\n flashy.distrib.barrier()\n return metrics\n\n def generate(self) -> dict:\n \"\"\"Generate stage.\"\"\"\n self.model.eval()\n with torch.no_grad():\n return self.generate_audio()\n\n def run_epoch(self):\n if self.cfg.cache.write:\n if ((self.epoch - 1) % self.cfg.cache.write_num_shards) != self.cfg.cache.write_shard:\n return\n super().run_epoch()\n\n def train(self):\n \"\"\"Train stage.\n \"\"\"\n if self._cached_batch_writer is not None:\n self._cached_batch_writer.start_epoch(self.epoch)\n if self._cached_batch_loader is None:\n dataset = get_dataset_from_loader(self.dataloaders['train'])\n assert isinstance(dataset, AudioDataset)\n dataset.current_epoch = self.epoch\n else:\n self._cached_batch_loader.start_epoch(self.epoch)\n return super().train()\n\n def evaluate_audio_generation(self) -> dict:\n \"\"\"Evaluate audio generation with off-the-shelf metrics.\"\"\"\n evaluate_stage_name = f'{self.current_stage}_generation'\n # instantiate evaluation metrics, if at least one metric is defined, run audio generation evaluation\n fad: tp.Optional[eval_metrics.FrechetAudioDistanceMetric] = None\n kldiv: tp.Optional[eval_metrics.KLDivergenceMetric] = None\n text_consistency: tp.Optional[eval_metrics.TextConsistencyMetric] = None\n chroma_cosine: tp.Optional[eval_metrics.ChromaCosineSimilarityMetric] = None\n should_run_eval = False\n eval_chroma_wavs: tp.Optional[torch.Tensor] = None\n if self.cfg.evaluate.metrics.fad:\n fad = builders.get_fad(self.cfg.metrics.fad).to(self.device)\n should_run_eval = True\n if self.cfg.evaluate.metrics.kld:\n kldiv = builders.get_kldiv(self.cfg.metrics.kld).to(self.device)\n should_run_eval = True\n if self.cfg.evaluate.metrics.text_consistency:\n text_consistency = builders.get_text_consistency(self.cfg.metrics.text_consistency).to(self.device)\n should_run_eval = True\n if self.cfg.evaluate.metrics.chroma_cosine:\n chroma_cosine = builders.get_chroma_cosine_similarity(self.cfg.metrics.chroma_cosine).to(self.device)\n # if we have predefind wavs for chroma we should purge them for computing the cosine metric\n has_predefined_eval_chromas = 'self_wav' in self.model.condition_provider.conditioners and \\\n self.model.condition_provider.conditioners['self_wav'].has_eval_wavs()\n if has_predefined_eval_chromas:\n warn_once(self.logger, \"Attempting to run cosine eval for config with pre-defined eval chromas! \"\n 'Resetting eval chromas to None for evaluation.')\n eval_chroma_wavs = self.model.condition_provider.conditioners.self_wav.eval_wavs # type: ignore\n self.model.condition_provider.conditioners.self_wav.reset_eval_wavs(None) # type: ignore\n should_run_eval = True\n\n def get_compressed_audio(audio: torch.Tensor) -> torch.Tensor:\n audio_tokens, scale = self.compression_model.encode(audio.to(self.device))\n compressed_audio = self.compression_model.decode(audio_tokens, scale)\n return compressed_audio[..., :audio.shape[-1]]\n\n metrics: dict = {}\n if should_run_eval:\n loader = self.dataloaders['evaluate']\n updates = len(loader)\n lp = self.log_progress(f'{evaluate_stage_name} inference', loader, total=updates, updates=self.log_updates)\n average = flashy.averager()\n dataset = get_dataset_from_loader(loader)\n assert isinstance(dataset, AudioDataset)\n self.logger.info(f\"Computing evaluation metrics on {len(dataset)} samples\")\n\n for idx, batch in enumerate(lp):\n audio, meta = batch\n assert all([self.cfg.sample_rate == m.sample_rate for m in meta])\n\n target_duration = audio.shape[-1] / self.cfg.sample_rate\n if self.cfg.evaluate.fixed_generation_duration:\n target_duration = self.cfg.evaluate.fixed_generation_duration\n\n gen_outputs = self.run_generate_step(\n batch, gen_duration=target_duration,\n remove_text_conditioning=self.cfg.evaluate.get('remove_text_conditioning', False)\n )\n y_pred = gen_outputs['gen_audio'].detach()\n y_pred = y_pred[..., :audio.shape[-1]]\n\n normalize_kwargs = dict(self.cfg.generate.audio)\n normalize_kwargs.pop('format', None)\n y_pred = torch.stack([normalize_audio(w, **normalize_kwargs) for w in y_pred], dim=0).cpu()\n y = audio.cpu() # should already be on CPU but just in case\n sizes = torch.tensor([m.n_frames for m in meta]) # actual sizes without padding\n sample_rates = torch.tensor([m.sample_rate for m in meta]) # sample rates for audio samples\n audio_stems = [Path(m.meta.path).stem + f\"_{m.seek_time}\" for m in meta]\n\n if fad is not None:\n if self.cfg.metrics.fad.use_gt:\n y_pred = get_compressed_audio(y).cpu()\n fad.update(y_pred, y, sizes, sample_rates, audio_stems)\n if kldiv is not None:\n if self.cfg.metrics.kld.use_gt:\n y_pred = get_compressed_audio(y).cpu()\n kldiv.update(y_pred, y, sizes, sample_rates)\n if text_consistency is not None:\n texts = [m.description for m in meta]\n if self.cfg.metrics.text_consistency.use_gt:\n y_pred = y\n text_consistency.update(y_pred, texts, sizes, sample_rates)\n if chroma_cosine is not None:\n if self.cfg.metrics.chroma_cosine.use_gt:\n y_pred = get_compressed_audio(y).cpu()\n chroma_cosine.update(y_pred, y, sizes, sample_rates)\n # restore chroma conditioner's eval chroma wavs\n if eval_chroma_wavs is not None:\n self.model.condition_provider.conditioners['self_wav'].reset_eval_wavs(eval_chroma_wavs)\n\n flashy.distrib.barrier()\n if fad is not None:\n metrics['fad'] = fad.compute()\n if kldiv is not None:\n kld_metrics = kldiv.compute()\n metrics.update(kld_metrics)\n if text_consistency is not None:\n metrics['text_consistency'] = text_consistency.compute()\n if chroma_cosine is not None:\n metrics['chroma_cosine'] = chroma_cosine.compute()\n metrics = average(metrics)\n metrics = flashy.distrib.average_metrics(metrics, len(loader))\n\n return metrics\n\n def evaluate(self) -> dict:\n \"\"\"Evaluate stage.\"\"\"\n self.model.eval()\n with torch.no_grad():\n metrics: dict = {}\n if self.cfg.evaluate.metrics.base:\n metrics.update(self.common_train_valid('evaluate'))\n gen_metrics = self.evaluate_audio_generation()\n return {**metrics, **gen_metrics}\n" }, { "path": "audiocraft/solvers/watermark.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nimport typing as tp\nfrom functools import partial\nimport os\nfrom pathlib import Path\n\nimport flashy\nfrom omegaconf import DictConfig\nimport multiprocessing\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nfrom . import base, builders\nfrom ..models.builders import get_watermark_model\nfrom ..modules.watermark import pad, mix\n\nfrom ..metrics.miou import calculate_miou\nfrom ..metrics.pesq import PesqMetric\n\nfrom ..utils import checkpoint\nfrom ..utils.audio_effects import (\n compress_with_encodec,\n get_audio_effects,\n select_audio_effects,\n)\nfrom ..utils.samples.manager import SampleManager\nfrom ..data.audio import save_spectrograms\nfrom ..utils.utils import get_pool_executor\n\nfrom torchmetrics.audio.snr import ScaleInvariantSignalNoiseRatio\nfrom torchmetrics.audio.stoi import ShortTimeObjectiveIntelligibility\n\n\nif tp.TYPE_CHECKING:\n from ..models.watermark import WMModel\n\n\ndef get_encodec_audio_effect(encodec_cfg: DictConfig, sr: int) -> tp.Dict:\n \"\"\"\n Construct encodec-based compression data agumentation. This method is\n is put here instead of in `audiocraft.utils.audio_effects` because\n it depends on the package `audiocraft.solvers`, which is one layer\n higher than `audiocraft.utils`, so we avoid the circle dependency\n from any solvers using `audiocraft.utils.audio_effects` to do the\n augmentation\n \"\"\"\n from ..solvers.compression import CompressionSolver\n\n codec_model = CompressionSolver.model_from_checkpoint(encodec_cfg.ckpt)\n codec_model.train()\n return {\n f\"encodec_nq={n_q}\": partial(\n compress_with_encodec,\n model=codec_model,\n n_q=n_q,\n sample_rate=sr,\n )\n for n_q in encodec_cfg.n_qs\n }\n\n\ndef random_message(nbits: int, batch_size: int) -> torch.Tensor:\n \"\"\"Return random message as 0/1 tensor.\"\"\"\n if nbits == 0:\n return torch.tensor([])\n return torch.randint(0, 2, (batch_size, nbits))\n\n\nclass WatermarkSolver(base.StandardSolver):\n \"\"\"Solver for different watermarking models\"\"\"\n\n def __init__(self, cfg: DictConfig):\n super().__init__(cfg)\n self.rng: torch.Generator # set at each epoch\n self.model: WMModel\n if hasattr(cfg, \"fsdp\"):\n assert not getattr(\n cfg.fsdp, \"use\", False\n ), \"FSDP not supported by WatermarkSolver.\"\n self._init_losses()\n self._init_augmentations()\n self.balancer = builders.get_balancer(self.loss_weights, self.cfg.balancer)\n self.path_specs = os.path.join(self.folder, \"spectrograms\")\n os.makedirs(self.path_specs, exist_ok=True)\n\n def _init_losses(self):\n assert hasattr(self.cfg, \"losses\") and isinstance(\n self.cfg.losses, (DictConfig, tp.Mapping)\n ), \"WatermarkSolver must declare training losses in the config\"\n\n self.adv_losses = builders.get_adversarial_losses(self.cfg) # noqa\n self.register_stateful(\"adv_losses\")\n\n self.aux_losses = nn.ModuleDict() # noqa\n self.info_losses = nn.ModuleDict() # noqa\n self.wm_losses = nn.ModuleDict() # noqa\n loss_weights = {}\n for loss_name, weight in self.cfg.losses.items():\n\n # explicitly skip this loss calculation by setting a -1 as weight\n # if weight == 0 it will be calculated but kept as info\n if weight == -1:\n continue\n\n if loss_name in [\"adv\", \"feat\"]:\n for adv_name, _ in self.adv_losses.items():\n loss_weights[f\"{loss_name}_{adv_name}\"] = weight\n elif weight > 0:\n if loss_name[:3] == \"wm_\":\n self.wm_losses[loss_name] = builders.get_loss(\n loss_name, self.cfg\n ).to(self.device)\n loss_weights[loss_name] = weight\n else:\n self.aux_losses[loss_name] = builders.get_loss(\n loss_name, self.cfg\n ).to(self.device)\n loss_weights[loss_name] = weight\n else:\n self.info_losses[loss_name] = builders.get_loss(loss_name, self.cfg).to(\n self.device\n )\n\n self.loss_weights = loss_weights # noqa\n\n def _init_augmentations(self):\n if not hasattr(self.cfg, \"aug_weights\") or not hasattr(\n self.cfg, \"audio_effects\"\n ):\n return\n\n aug_weights = {}\n cfg_audio_effects = dict(self.cfg.audio_effects)\n\n # Handle `encodec` augmentation separately as this requires loading a\n # CompressionSolver checkpoint\n encodec_cfg = cfg_audio_effects.pop(\"encodec\", None)\n if encodec_cfg:\n encodec_effects = get_encodec_audio_effect(\n encodec_cfg, self.cfg.sample_rate\n )\n for aug_name in encodec_effects.keys():\n aug_weights[aug_name] = getattr(self.cfg.aug_weights, \"encodec\", -1)\n else:\n encodec_effects = {}\n\n other_effects = get_audio_effects(self.cfg) # noqa\n for name in other_effects.keys():\n aug_weights[name] = self.cfg.aug_weights.get(name, -1)\n\n self.aug_weights = aug_weights # noqa\n self.augmentations = {**encodec_effects, **other_effects} # noqa\n\n @property\n def best_metric_name(self) -> tp.Optional[str]:\n # best model is the last for the watermark model for now\n return None\n\n def build_model(self):\n \"\"\"Instantiate model and optimizer.\"\"\"\n # Model and optimizer\n self.model = get_watermark_model(self.cfg)\n # Need two optimizers ?\n self.optimizer = builders.get_optimizer(self.model.parameters(), self.cfg.optim)\n self.register_stateful(\"model\", \"optimizer\")\n self.register_best_state(\"model\")\n self.register_ema(\"model\")\n\n def build_dataloaders(self):\n \"\"\"Instantiate audio dataloaders for each stage.\"\"\"\n self.dataloaders = builders.get_audio_datasets(self.cfg)\n\n def show(self):\n \"\"\"Show the Watermark model and employed adversarial loss.\"\"\"\n self.log_model_summary(self.model)\n self.logger.info(\"Sould print losses here:\")\n\n def crop(\n self, signal: torch.Tensor, watermark: torch.Tensor\n ) -> tp.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"\n Applies a transformation to modify the watermarked signal to train localization.\n It can be one of the following:\n - zero padding: add zeros at the begining and the end of the signal\n - crop: crop the watermark apply a watermark only on some parts of the signal\n - shuffle: replace some part of the audio with other non watermarked parts\n from the batch\n In every cases the function returns a mask that contains indicates the parts that are or\n not watermarked\n\n Args:\n watermark (torch.Tensor): The watermark to apply on the signal.\n signal (torch.Tensor): clean signal\n Returns:\n watermark (torch.Tensor): modified watermark\n signal (torch.Tensor): modified signal\n mask (torch.Tensor): mask indicating which portion is still watermarked\n \"\"\"\n assert (\n self.cfg.crop.prob + self.cfg.crop.shuffle_prob + self.cfg.crop.pad_prob\n <= 1\n ), f\"The sum of the probabilities {self.cfg.crop.prob=} {self.cfg.crop.shuffle_prob=} \\\n {self.cfg.crop.pad_prob=} should be less than 1\"\n mask = torch.ones_like(watermark)\n p = torch.rand(1)\n if p < self.cfg.crop.pad_prob: # Pad with some probability\n start = int(torch.rand(1) * 0.33 * watermark.size(-1))\n finish = int((0.66 + torch.rand(1) * 0.33) * watermark.size(-1))\n mask[:, :, :start] = 0\n mask[:, :, finish:] = 0\n if torch.rand(1) > 0.5:\n mask = 1 - mask\n signal *= mask # pad signal\n\n elif (\n p < self.cfg.crop.prob + self.cfg.crop.pad_prob + self.cfg.crop.shuffle_prob\n ):\n # Define a mask, then crop or shuffle\n mask_size = round(watermark.shape[-1] * self.cfg.crop.size)\n n_windows = int(\n torch.randint(1, self.cfg.crop.max_n_windows + 1, (1,)).item()\n )\n window_size = int(mask_size / n_windows)\n for _ in range(n_windows): # Create multiple windows in the mask\n mask_start = torch.randint(0, watermark.shape[-1] - window_size, (1,))\n mask[:, :, mask_start: mask_start + window_size] = (\n 0 # Apply window to mask\n )\n # inverse the mask half the time\n if torch.rand(1) > 0.5:\n mask = 1 - mask\n\n if p < self.cfg.crop.pad_prob + self.cfg.crop.shuffle_prob: # shuffle\n # shuffle\n signal_cloned = signal.clone().detach() # detach to be sure\n shuffle_idx = torch.randint(0, signal.size(0), (signal.size(0),))\n signal = signal * mask + signal_cloned[shuffle_idx] * (\n 1 - mask\n ) # shuffle signal where not wm\n\n watermark *= mask # Apply mask to the watermark\n return signal, watermark, mask\n\n def run_step(self, idx: int, batch: torch.Tensor, metrics: dict):\n \"\"\"Perform one training or valid step on a given batch.\"\"\"\n x = batch.to(self.device)\n y = x.clone()\n nbits = getattr(self.model, \"nbits\")\n message = random_message(nbits, y.shape[0]).to(self.device)\n watermark = self.model.get_watermark(x, message=message)\n y, watermark, mask = self.crop(y, watermark)\n\n y_wm = y + watermark\n\n if (\n self.cfg.losses.adv != 0 or self.cfg.losses.feat != 0\n ) and self.is_training: # train quality adv\n d_losses: dict = {}\n if (\n len(self.adv_losses) > 0\n and torch.rand(1, generator=self.rng).item()\n <= 1 / self.cfg.adversarial.every\n ):\n for adv_name, adversary in self.adv_losses.items():\n disc_loss = adversary.train_adv(y_wm, y)\n d_losses[f\"d_{adv_name}\"] = disc_loss\n metrics[\"d_loss\"] = torch.sum(torch.stack(list(d_losses.values())))\n metrics.update(d_losses)\n\n balanced_losses: dict = {}\n other_losses: dict = {}\n\n # adversarial losses\n if self.cfg.losses.adv != 0 or self.cfg.losses.feat != 0:\n for adv_name, adversary in self.adv_losses.items():\n adv_loss, feat_loss = adversary(y_wm, y)\n balanced_losses[f\"adv_{adv_name}\"] = adv_loss\n balanced_losses[f\"feat_{adv_name}\"] = feat_loss\n\n # auxiliary losses on quality/similarity\n for loss_name, criterion in self.aux_losses.items():\n loss = criterion(y_wm, y)\n balanced_losses[loss_name] = loss\n\n # apply augmentations\n mode = \"all\" if self.cfg.select_aug_mode == \"all\" else \"weighted\"\n selected_augs = select_audio_effects(\n self.augmentations,\n self.aug_weights,\n mode=mode,\n max_length=self.cfg.n_max_aug,\n )\n N_augs = len(selected_augs)\n for (\n augmentation_name,\n augmentation_method,\n ) in selected_augs.items():\n # concatenate to use the augmentation function only once\n y_y_wm = torch.cat([y, y_wm], dim=0)\n aug_cat, mask_aug = augmentation_method(y_y_wm, mask=mask)\n aug_y = aug_cat[: y.size(0)]\n aug_y_wm = aug_cat[y.size(0):]\n positive = self.model.detect_watermark(aug_y_wm)\n negative = self.model.detect_watermark(aug_y)\n for loss_name, criterion in self.wm_losses.items():\n loss = criterion(positive, negative, mask_aug, message)\n other_losses[f\"{loss_name}_{augmentation_name}\"] = loss\n\n # weighted losses\n metrics.update(balanced_losses)\n metrics.update(other_losses)\n if self.is_training: # something is weird about the loss balancer not\n other_loss = torch.tensor(0.0, device=self.device)\n for name, o_loss in other_losses.items():\n if \"wm_detection\" in name:\n # here we include the detection losses for augmentation\n other_loss += (self.loss_weights[\"wm_detection\"] / N_augs) * o_loss\n elif \"wm_mb\" in name:\n other_loss += (self.loss_weights[\"wm_mb\"] / N_augs) * o_loss\n else:\n other_loss += self.loss_weights[name] * o_loss\n if other_loss.requires_grad:\n other_loss.backward(retain_graph=True)\n ratio1 = sum(\n p.grad.data.norm(p=2).pow(2)\n for p in self.model.parameters()\n if p.grad is not None\n )\n assert isinstance(ratio1, torch.Tensor)\n metrics[\"ratio1\"] = ratio1.sqrt()\n\n # balancer losses backward, returns effective training loss\n # with effective weights at the current batch.\n metrics[\"g_loss\"] = self.balancer.backward(balanced_losses, y_wm)\n # add metrics corresponding to weight ratios\n metrics.update(self.balancer.metrics)\n ratio2 = sum(\n p.grad.data.norm(p=2).pow(2)\n for p in self.model.parameters()\n if p.grad is not None\n )\n assert isinstance(ratio2, torch.Tensor)\n metrics[\"ratio2\"] = ratio2.sqrt()\n\n # optim\n flashy.distrib.sync_model(self.model)\n if self.cfg.optim.max_norm:\n torch.nn.utils.clip_grad_norm_(\n self.model.parameters(), self.cfg.optim.max_norm\n )\n\n self.optimizer.step()\n self.optimizer.zero_grad()\n\n # informative losses only\n info_losses: dict = {}\n with torch.no_grad():\n for loss_name, criterion in self.info_losses.items():\n loss = criterion(y_wm, y)\n info_losses[loss_name] = loss\n # pesq\n metrics[\"pesq\"] = tensor_pesq(y_wm, y, sr=self.cfg.sample_rate)\n # max allocated memory\n metrics[\"max_mem\"] = torch.cuda.max_memory_allocated() / 1e9\n\n metrics.update(info_losses)\n if self.cfg.losses.adv != 0 or self.cfg.losses.feat != 0:\n # aggregated GAN losses: this is useful to report adv and feat across different adversarial loss setups\n adv_losses = [\n loss\n for loss_name, loss in metrics.items()\n if loss_name.startswith(\"adv\")\n ]\n if len(adv_losses) > 0:\n metrics[\"adv\"] = torch.sum(torch.stack(adv_losses))\n feat_losses = [\n loss\n for loss_name, loss in metrics.items()\n if loss_name.startswith(\"feat\")\n ]\n if len(feat_losses) > 0:\n metrics[\"feat\"] = torch.sum(torch.stack(feat_losses))\n\n return metrics\n\n def run_epoch(self):\n # reset random seed at the beginning of the epoch\n self.rng = torch.Generator()\n self.rng.manual_seed(1234 + self.epoch)\n # run epoch\n super().run_epoch()\n\n def evaluate(self) -> dict:\n \"\"\"Evaluate stage. Runs audio reconstruction evaluation.\"\"\"\n self.model.eval()\n evaluate_stage_name = str(self.current_stage)\n\n loader = self.dataloaders[\"evaluate\"]\n updates = len(loader)\n lp = self.log_progress(\n f\"{evaluate_stage_name} inference\",\n loader,\n total=updates,\n updates=self.log_updates,\n )\n average = flashy.averager()\n\n pendings = []\n ctx = multiprocessing.get_context(\"spawn\")\n with get_pool_executor(self.cfg.evaluate.num_workers, mp_context=ctx) as pool:\n for batch in lp:\n x = batch.to(self.device)\n with torch.no_grad():\n message = random_message(self.model.nbits, x.shape[0])\n watermark = self.model.get_watermark(x, message)\n x_wm = x + watermark\n y_pred = x_wm.cpu()\n y = batch.cpu() # should already be on CPU but just in case\n pendings.append(\n pool.submit(\n evaluate_audio_watermark,\n y_pred,\n y,\n self.cfg,\n )\n )\n # evaluate augmentations\n # evaluation is run on all the augmentations\n for (\n augmentation_name,\n augmentation_method,\n ) in self.augmentations.items():\n # if (\n # \"mp3\" in augmentation_name\n # and idx >= 8\n # and self.cfg.evaluate.every <= 2\n # ):\n # # When evaluating often do not compute mp3 on the full eval dset to make things faster\n # continue\n with torch.no_grad():\n aug_positive = self.model.detect_watermark(\n augmentation_method(x_wm)\n )\n aug_negative = self.model.detect_watermark(\n augmentation_method(x)\n )\n\n pendings.append(\n pool.submit(\n evaluate_augmentations,\n aug_positive.cpu(),\n aug_negative.cpu(),\n augmentation_name,\n message.cpu(),\n )\n )\n # end eval of augmentations\n\n # evaluate localization cropping\n for window_size in np.linspace(0.1, 0.9, 9):\n\n mixed, true_predictions = mix(x, x_wm, window_size=window_size)\n model_predictions = self.model.detect_watermark(mixed)\n pendings.append(\n pool.submit(\n evaluate_localizations,\n model_predictions.cpu(),\n true_predictions.cpu(),\n f\"crop_{window_size:0.1f}\",\n )\n )\n mixed, true_predictions = mix(\n x, x_wm, window_size=window_size, shuffle=True\n )\n model_predictions = self.model.detect_watermark(mixed)\n pendings.append(\n pool.submit(\n evaluate_localizations,\n model_predictions.cpu(),\n true_predictions.cpu(),\n f\"shuffle_{window_size:0.1f}\",\n )\n )\n # evaluate localization padding\n mixed, true_predictions = pad(x_wm)\n model_predictions = self.model.detect_watermark(mixed)\n pendings.append(\n pool.submit(\n evaluate_localizations,\n model_predictions.cpu(),\n true_predictions.cpu(),\n \"padding\",\n )\n )\n mixed, true_predictions = pad(x_wm, central=True)\n model_predictions = self.model.detect_watermark(mixed)\n pendings.append(\n pool.submit(\n evaluate_localizations,\n model_predictions.cpu(),\n true_predictions.cpu(),\n \"central_padding\",\n )\n )\n # end of evaluate localization\n\n metrics_lp = self.log_progress(\n f\"{evaluate_stage_name} metrics\", pendings, updates=self.log_updates\n )\n for pending in metrics_lp:\n metrics = pending.result()\n metrics = average(metrics)\n\n metrics = flashy.distrib.average_metrics(metrics, len(loader))\n if self.cfg.select_aug_mode == \"use_eval_acc\":\n # Adjust augmentation weights based on evaluation loss.\n # Higher accuracy results in lower probability of selecting this augmentation.\n for name in self.augmentations.keys():\n if (\n self.aug_weights[name] != -1\n ): # keep weight to -1 for unwanted augmentations\n # set to 0.05 to ensure that an augmentation is never completely removed during a full epoch.\n self.aug_weights[name] = max(1 - metrics[f\"aug_{name}_acc\"], 0.05)\n return metrics\n\n def generate(self):\n \"\"\"Generate stage.\"\"\"\n self.model.eval()\n sample_manager = SampleManager(self.xp, map_reference_to_sample_id=True)\n generate_stage_name = str(self.current_stage)\n\n loader = self.dataloaders[\"generate\"]\n updates = len(loader)\n lp = self.log_progress(\n generate_stage_name, loader, total=updates, updates=self.log_updates\n )\n path_dir = os.path.join(self.path_specs, f\"epoch={self.epoch}\")\n os.makedirs(path_dir, exist_ok=True)\n first_batch = True\n for batch in lp:\n reference, _ = batch\n reference = reference.to(self.device)\n with torch.no_grad():\n message = random_message(self.model.nbits, reference.shape[0])\n watermark = self.model.get_watermark(reference, message)\n x_wm = reference + watermark\n\n reference = reference.cpu()\n sample_manager.add_samples(\n x_wm.cpu(), self.epoch, ground_truth_wavs=reference\n )\n if first_batch and flashy.distrib.is_rank_zero():\n for i in range(reference.size(0)):\n ys = [\n reference.cpu()[i].squeeze(0).numpy(),\n x_wm.cpu()[i].squeeze(0).numpy(),\n watermark.cpu()[i].squeeze(0).numpy(),\n ]\n path = os.path.join(path_dir, f\"spec_{i}.pdf\")\n save_spectrograms(\n ys,\n names=[\"Ground Truth\", \"Audio Watermarked\", \"Watermark\"],\n sr=self.cfg.sample_rate,\n path=path,\n )\n first_batch = False\n flashy.distrib.barrier()\n\n def load_from_pretrained(self, name: str) -> dict:\n raise ValueError(\"No pretrained model\")\n\n @staticmethod\n def model_from_checkpoint(\n checkpoint_path: tp.Union[Path, str],\n device: tp.Union[torch.device, str] = \"cpu\",\n ) -> \"WMModel\":\n \"\"\"Instantiate a WatermarkModel from a given checkpoint path or dora sig.\n\n Args:\n checkpoint_path (Path or str): Path to checkpoint or dora sig from where the checkpoint is resolved.\n device (torch.device or str): Device on which the model is loaded.\n \"\"\"\n checkpoint_path = str(checkpoint_path)\n logger = logging.getLogger(__name__)\n logger.info(f\"Loading WatermarkModel from checkpoint: {checkpoint_path}\")\n _checkpoint_path = checkpoint.resolve_checkpoint_path(\n checkpoint_path, use_fsdp=False\n )\n assert (\n _checkpoint_path is not None\n ), f\"Could not resolve WatermarkModel checkpoint path: {checkpoint_path}\"\n state = checkpoint.load_checkpoint(_checkpoint_path)\n assert (\n state is not None and \"xp.cfg\" in state\n ), f\"Could not load WatermarkModel from ckpt: {checkpoint_path}\"\n cfg = state[\"xp.cfg\"]\n cfg.device = device\n watermarking_model = get_watermark_model(cfg).to(device)\n\n assert \"best_state\" in state and state[\"best_state\"] != {}\n assert (\n \"exported\" not in state\n ), \"When loading an exported checkpoint, use the //pretrained/ prefix.\"\n watermarking_model.load_state_dict(state[\"best_state\"][\"model\"])\n watermarking_model.eval()\n logger.info(\"Watermarking model loaded!\")\n return watermarking_model\n\n\ndef evaluate_localizations(predictions, true_predictions, name):\n metrics = {}\n # predictions are output of the detector shape [bsz, 2, frames]\n # true_predictions is output of the mix method shape [bsz, 2, frames]\n metrics[f\"localization_acc_{name}\"] = (\n ((predictions[:, 1, :] > 0.5) == true_predictions[:, 1, :])\n .float()\n .mean()\n .item()\n )\n metrics[f\"localization_miou_{name}\"] = calculate_miou(\n predictions[:, 1, :], true_predictions[:, 1, :]\n )\n return metrics\n\n\ndef evaluate_augmentations(\n positive: torch.Tensor,\n negative: torch.Tensor,\n augmentation_name: str,\n message: torch.Tensor,\n) -> dict:\n \"\"\"calculating evaluation metrics but take name of the augmentation\n method that has been done before getting positive and negative results\"\"\"\n metrics = {}\n metrics[f\"aug_{augmentation_name}_acc\"] = compute_accuracy(positive, negative)\n metrics[f\"aug_{augmentation_name}_fpr\"] = compute_FPR(negative)\n metrics[f\"aug_{augmentation_name}_fnr\"] = compute_FNR(positive)\n if message.shape[0] != 0:\n metrics[f\"aug_{augmentation_name}_bit_acc\"] = compute_bit_acc(positive, message)\n\n # add one metric which is average overall score of all augmentations\n metrics[\"all_aug_acc\"] = compute_accuracy(positive, negative)\n\n return metrics\n\n\ndef evaluate_audio_watermark(\n y_pred: torch.Tensor,\n y: torch.Tensor,\n cfg: DictConfig,\n) -> dict:\n \"\"\"Audio reconstruction evaluation method that can be conveniently pickled.\"\"\"\n metrics = {}\n if cfg.evaluate.metrics.visqol:\n visqol = builders.get_visqol(cfg.metrics.visqol)\n metrics[\"visqol\"] = visqol(y_pred, y, cfg.sample_rate)\n sisnr = ScaleInvariantSignalNoiseRatio().to(y.device)\n stoi = ShortTimeObjectiveIntelligibility(fs=cfg.sample_rate)\n metrics[\"sisnr\"] = sisnr(y_pred, y)\n metrics[\"stoi\"] = stoi(y_pred, y)\n metrics[\"pesq\"] = tensor_pesq(y_pred, y, sr=cfg.sample_rate)\n return metrics\n\n\ndef tensor_pesq(y_pred: torch.Tensor, y: torch.Tensor, sr: int):\n # pesq returns error if no speech is detected, so we catch it\n return PesqMetric(sr)(y_pred, y).item()\n\n\ndef compute_accuracy(positive, negative):\n N = (positive[:, 1, :].mean(dim=1) > 0.5).sum() + (\n negative[:, 0, :].mean(dim=1) > 0.5\n ).sum()\n acc = N / (2 * positive.size(0))\n return acc\n\n\ndef compute_FPR(negative):\n N = (negative[:, 1, :].mean(dim=1) > 0.5).sum()\n fpr = N / (negative.size(0))\n return fpr\n\n\ndef compute_FNR(positive):\n N = (positive[:, 0, :].mean(dim=1) > 0.5).sum()\n fpr = N / (positive.size(0))\n return fpr\n\n\ndef _bit_acc(decoded, original):\n bit_acc = (decoded == original).float().mean()\n return bit_acc\n\n\ndef compute_bit_acc(positive, original, mask=None):\n \"\"\"Compute bit accuracy.\n Args:\n positive: detector outputs [bsz, 2+nbits, time_steps]\n original: original message (0 or 1) [bsz, nbits]\n mask: mask of the watermark [bsz, 1, time_steps]\n \"\"\"\n decoded = positive[:, 2:, :] # b 2+nbits t -> b nbits t\n if mask is not None:\n # cut last dim of positive to keep only where mask is 1\n new_shape = [*decoded.shape[:-1], -1] # b nbits t -> b nbits -1\n decoded = torch.masked_select(decoded, mask == 1).reshape(new_shape)\n # average decision over time, then threshold\n decoded = decoded.mean(dim=-1) > 0 # b nbits\n return _bit_acc(decoded, original)\n" }, { "path": "audiocraft/train.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nEntry point for dora to launch solvers for running training loops.\nSee more info on how to use dora: https://github.com/facebookresearch/dora\n\"\"\"\n\nimport logging\nimport multiprocessing\nimport os\nfrom pathlib import Path\nimport sys\nimport typing as tp\n\nfrom dora import git_save, hydra_main, XP\nimport flashy\nimport hydra\nimport omegaconf\n\nfrom .environment import AudioCraftEnvironment\nfrom .utils.cluster import get_slurm_parameters\n\nlogger = logging.getLogger(__name__)\n\n\ndef resolve_config_dset_paths(cfg):\n \"\"\"Enable Dora to load manifest from git clone repository.\"\"\"\n # manifest files for the different splits\n for key, value in cfg.datasource.items():\n if isinstance(value, str):\n cfg.datasource[key] = git_save.to_absolute_path(value)\n\n\ndef get_solver(cfg):\n from . import solvers\n # Convert batch size to batch size for each GPU\n assert cfg.dataset.batch_size % flashy.distrib.world_size() == 0\n cfg.dataset.batch_size //= flashy.distrib.world_size()\n for split in ['train', 'valid', 'evaluate', 'generate']:\n if hasattr(cfg.dataset, split) and hasattr(cfg.dataset[split], 'batch_size'):\n assert cfg.dataset[split].batch_size % flashy.distrib.world_size() == 0\n cfg.dataset[split].batch_size //= flashy.distrib.world_size()\n resolve_config_dset_paths(cfg)\n solver = solvers.get_solver(cfg)\n return solver\n\n\ndef get_solver_from_xp(xp: XP, override_cfg: tp.Optional[tp.Union[dict, omegaconf.DictConfig]] = None,\n restore: bool = True, load_best: bool = True,\n ignore_state_keys: tp.List[str] = [], disable_fsdp: bool = True):\n \"\"\"Given a XP, return the Solver object.\n\n Args:\n xp (XP): Dora experiment for which to retrieve the solver.\n override_cfg (dict or None): If not None, should be a dict used to\n override some values in the config of `xp`. This will not impact\n the XP signature or folder. The format is different\n than the one used in Dora grids, nested keys should actually be nested dicts,\n not flattened, e.g. `{'optim': {'batch_size': 32}}`.\n restore (bool): If `True` (the default), restore state from the last checkpoint.\n load_best (bool): If `True` (the default), load the best state from the checkpoint.\n ignore_state_keys (list[str]): List of sources to ignore when loading the state, e.g. `optimizer`.\n disable_fsdp (bool): if True, disables FSDP entirely. This will\n also automatically skip loading the EMA. For solver specific\n state sources, like the optimizer, you might want to\n use along `ignore_state_keys=['optimizer']`. Must be used with `load_best=True`.\n \"\"\"\n logger.info(f\"Loading solver from XP {xp.sig}. \"\n f\"Overrides used: {xp.argv}\")\n cfg = xp.cfg\n if override_cfg is not None:\n cfg = omegaconf.OmegaConf.merge(cfg, omegaconf.DictConfig(override_cfg))\n if disable_fsdp and cfg.fsdp.use:\n cfg.fsdp.use = False\n assert load_best is True\n # ignoring some keys that were FSDP sharded like model, ema, and best_state.\n # fsdp_best_state will be used in that case. When using a specific solver,\n # one is responsible for adding the relevant keys, e.g. 'optimizer'.\n # We could make something to automatically register those inside the solver, but that\n # seem overkill at this point.\n ignore_state_keys = ignore_state_keys + ['model', 'ema', 'best_state']\n\n try:\n with xp.enter():\n solver = get_solver(cfg)\n if restore:\n solver.restore(load_best=load_best, ignore_state_keys=ignore_state_keys)\n return solver\n finally:\n hydra.core.global_hydra.GlobalHydra.instance().clear()\n\n\ndef get_solver_from_sig(sig: str, *args, **kwargs):\n \"\"\"Return Solver object from Dora signature, i.e. to play with it from a notebook.\n See `get_solver_from_xp` for more information.\n \"\"\"\n xp = main.get_xp_from_sig(sig)\n return get_solver_from_xp(xp, *args, **kwargs)\n\n\ndef init_seed_and_system(cfg):\n import numpy as np\n import torch\n import random\n from audiocraft.modules.transformer import set_efficient_attention_backend\n\n multiprocessing.set_start_method(cfg.mp_start_method)\n logger.debug('Setting mp start method to %s', cfg.mp_start_method)\n random.seed(cfg.seed)\n np.random.seed(cfg.seed)\n # torch also initialize cuda seed if available\n torch.manual_seed(cfg.seed)\n torch.set_num_threads(cfg.num_threads)\n os.environ['MKL_NUM_THREADS'] = str(cfg.num_threads)\n os.environ['OMP_NUM_THREADS'] = str(cfg.num_threads)\n logger.debug('Setting num threads to %d', cfg.num_threads)\n set_efficient_attention_backend(cfg.efficient_attention_backend)\n logger.debug('Setting efficient attention backend to %s', cfg.efficient_attention_backend)\n if 'SLURM_JOB_ID' in os.environ:\n tmpdir = Path('/scratch/slurm_tmpdir/' + os.environ['SLURM_JOB_ID'])\n if tmpdir.exists():\n logger.info(\"Changing tmpdir to %s\", tmpdir)\n os.environ['TMPDIR'] = str(tmpdir)\n\n\n@hydra_main(config_path='../config', config_name='config', version_base='1.1')\ndef main(cfg):\n init_seed_and_system(cfg)\n\n # Setup logging both to XP specific folder, and to stderr.\n log_name = '%s.log.{rank}' % cfg.execute_only if cfg.execute_only else 'solver.log.{rank}'\n flashy.setup_logging(level=str(cfg.logging.level).upper(), log_name=log_name)\n # Initialize distributed training, no need to specify anything when using Dora.\n flashy.distrib.init()\n solver = get_solver(cfg)\n if cfg.show:\n solver.show()\n return\n\n if cfg.execute_only:\n assert cfg.execute_inplace or cfg.continue_from is not None, \\\n \"Please explicitly specify the checkpoint to continue from with continue_from= \" + \\\n \"when running with execute_only or set execute_inplace to True.\"\n solver.restore(replay_metrics=False) # load checkpoint\n solver.run_one_stage(cfg.execute_only)\n return\n\n return solver.run()\n\n\nmain.dora.dir = AudioCraftEnvironment.get_dora_dir()\nmain._base_cfg.slurm = get_slurm_parameters(main._base_cfg.slurm)\n\nif main.dora.shared is not None and not os.access(main.dora.shared, os.R_OK):\n print(\"No read permission on dora.shared folder, ignoring it.\", file=sys.stderr)\n main.dora.shared = None\n\nif __name__ == '__main__':\n main()\n" }, { "path": "audiocraft/utils/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Utilities.\"\"\"\n" }, { "path": "audiocraft/utils/audio_effects.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\nimport inspect\nimport random\nimport typing as tp\nfrom functools import partial\n\nimport julius\nimport omegaconf\nimport torch\nfrom julius import fft_conv1d, resample_frac\n\nfrom ..data.audio_utils import get_aac, get_mp3\n\nif tp.TYPE_CHECKING:\n from ..models.encodec import CompressionModel\n\n\ndef select_audio_effects(\n audio_effects: tp.Dict,\n weights: tp.Optional[tp.Dict] = None,\n mode: str = \"all\",\n max_length: tp.Optional[int] = None,\n):\n \"\"\"Samples a subset of audio effects methods from the `AudioEffects` class.\n\n This function allows you to select a subset of audio effects\n based on the chosen selection mode and optional weights.\n\n Args:\n audio_effects (dict): A dictionary of available audio augmentations, usually\n obtained from the output of the 'get_audio_effects' function.\n weights (dict): A dictionary mapping augmentation names to their corresponding\n probabilities of being selected. This argument is used when 'mode' is set\n to \"weighted.\" If 'weights' is None, all augmentations have equal\n probability of being selected.\n mode (str): The selection mode, which can be one of the following:\n - \"all\": Select all available augmentations.\n - \"weighted\": Select augmentations based on their probabilities in the\n 'weights' dictionary.\n max_length (int): The maximum number of augmentations to select. If 'max_length'\n is None, no limit is applied.\n\n Returns:\n dict: A subset of the 'audio_effects' dictionary containing the selected audio\n augmentations.\n\n Note:\n - In \"all\" mode, all available augmentations are selected.\n - In \"weighted\" mode, augmentations are selected with a probability\n proportional to their weights specified in the 'weights' dictionary.\n - If 'max_length' is set, the function limits the number of selected\n augmentations.\n - If no augmentations are selected or 'audio_effects' is empty, the function\n defaults to including an \"identity\" augmentation.\n - The \"identity\" augmentation means that no audio effect is applied.\n \"\"\"\n if mode == \"all\": # original code\n out = audio_effects\n elif mode == \"weighted\":\n # Probability proportionnal to weights\n assert weights is not None\n out = {\n name: value\n for name, value in audio_effects.items()\n if random.random() < weights.get(name, 1.0)\n }\n else:\n raise ValueError(f\"Unknown mode {mode}\")\n if max_length is not None:\n # Help having a deterministic limit of the gpu memory usage\n random_keys = random.sample(list(out.keys()), max_length)\n out = {key: out[key] for key in random_keys}\n if len(out) == 0: # Check not to return empty dict\n out = {\"identity\": AudioEffects.identity}\n return out\n\n\ndef get_audio_effects(cfg: omegaconf.DictConfig):\n \"\"\"Automatically pull the list all effects available in this class based on the parameters from the cfg\n\n Returns:\n dict: A dict of names and pointers to all methods in this class.\n \"\"\"\n assert hasattr(cfg, \"audio_effects\")\n cfg_audio_effects = dict(cfg[\"audio_effects\"])\n return {\n name: partial(value, **cfg_audio_effects.get(name, {}))\n for name, value in inspect.getmembers(AudioEffects)\n if inspect.isfunction(value)\n }\n\n\ndef audio_effect_return(\n tensor: torch.Tensor, mask: tp.Optional[torch.Tensor]\n) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Return the mask if it was in the input otherwise only the output tensor\"\"\"\n if mask is None:\n return tensor\n else:\n return tensor, mask\n\n\ndef generate_pink_noise(length: int) -> torch.Tensor:\n \"\"\"Generate pink noise using Voss-McCartney algorithm with PyTorch.\"\"\"\n num_rows = 16\n array = torch.randn(num_rows, length // num_rows + 1)\n reshaped_array = torch.cumsum(array, dim=1)\n reshaped_array = reshaped_array.reshape(-1)\n reshaped_array = reshaped_array[:length]\n # Normalize\n pink_noise = reshaped_array / torch.max(torch.abs(reshaped_array))\n return pink_noise\n\n\ndef compress_with_encodec(\n tensor: torch.Tensor,\n n_q: int,\n model: \"CompressionModel\",\n sample_rate: int,\n mask: tp.Optional[torch.Tensor] = None,\n) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Special augmentation function that compresses and decompresses wav tensor\n using a compression model with the n_q codebooks\n \"\"\"\n\n model.to(tensor.device)\n model.set_num_codebooks(n_q)\n codes, scale = model.encode(\n julius.resample_frac(tensor, old_sr=sample_rate, new_sr=model.sample_rate)\n )\n compressed = model.decode(codes=codes, scale=scale)\n return audio_effect_return(\n tensor=julius.resample_frac(\n compressed, old_sr=model.sample_rate, new_sr=sample_rate\n ),\n mask=mask,\n )\n\n\ndef apply_compression_skip_grad(tensor: torch.Tensor, compression_fn, **kwargs):\n \"\"\"Applies a specified compression function to the audio tensor.\n Whire carrying over the grads to the output tensor with skip through estimator\n this is a straight through estimator to make mp3/aac compression differentiable\n see more: Yin et al. 2019 https://arxiv.org/pdf/1903.05662.pdf\n\n Args:\n tensor (torch.Tensor): The input audio tensor.\n compression_fn (function): The compression function to apply.\n **kwargs: Additional keyword arguments for the compression function.\n\n Returns:\n torch.Tensor: The output tensor after applying compression and straight through estimator.\n \"\"\"\n compressed = compression_fn(tensor.detach(), **kwargs)\n\n # Trim compressed output if needed\n compressed = compressed[:, :, : tensor.size(-1)]\n\n # Straight through estimator for differentiable compression\n out = tensor + (compressed - tensor).detach()\n\n # Check that gradients are not broken\n if out.requires_grad:\n assert (\n out.grad_fn\n ), \"The computation graph might be broken due to compression augmentation.\"\n\n return out\n\n\nclass AudioEffects:\n @staticmethod\n def speed(\n tensor: torch.Tensor,\n speed_range: tuple = (0.5, 1.5),\n sample_rate: int = 16000,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Function to change the speed of a batch of audio data.\n The output will have a different length !\n\n Args:\n audio_batch (torch.Tensor): The batch of audio data in torch tensor format.\n speed (float): The speed to change the audio to.\n\n Returns:\n torch.Tensor: The batch of audio data with the speed changed.\n \"\"\"\n speed = torch.FloatTensor(1).uniform_(*speed_range)\n new_sr = int(sample_rate * 1 / speed)\n resampled_tensor = julius.resample.resample_frac(tensor, sample_rate, new_sr)\n if mask is None:\n return resampled_tensor\n else:\n return resampled_tensor, torch.nn.functional.interpolate(\n mask, size=resampled_tensor.size(-1), mode=\"nearest-exact\"\n )\n\n @staticmethod\n def updownresample(\n tensor: torch.Tensor,\n sample_rate: int = 16000,\n intermediate_freq: int = 32000,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n\n orig_shape = tensor.shape\n # upsample\n tensor = resample_frac(tensor, sample_rate, intermediate_freq)\n # downsample\n tensor = resample_frac(tensor, intermediate_freq, sample_rate)\n\n assert tensor.shape == orig_shape\n return audio_effect_return(tensor=tensor, mask=mask)\n\n @staticmethod\n def echo(\n tensor: torch.Tensor,\n volume_range: tuple = (0.1, 0.5),\n duration_range: tuple = (0.1, 0.5),\n sample_rate: int = 16000,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Attenuating the audio volume by a factor of 0.4, delaying it by 100ms,\n and then overlaying it with the original.\n\n Args:\n tensor: 3D Tensor representing the audio signal [bsz, channels, frames]\n volumne range: volume range of the echo signal\n duration range: duration range of the echo signal\n sample_rate: Sample rate of the audio signal.\n Returns:\n Audio signal with reverb.\n \"\"\"\n\n # Create a simple impulse response\n # Duration of the impulse response in seconds\n duration = torch.FloatTensor(1).uniform_(*duration_range)\n volume = torch.FloatTensor(1).uniform_(*volume_range)\n\n n_samples = int(sample_rate * duration)\n impulse_response = torch.zeros(n_samples).type(tensor.type()).to(tensor.device)\n\n # Define a few reflections with decreasing amplitude\n impulse_response[0] = 1.0 # Direct sound\n\n impulse_response[\n int(sample_rate * duration) - 1\n ] = volume # First reflection after 100ms\n\n # Add batch and channel dimensions to the impulse response\n impulse_response = impulse_response.unsqueeze(0).unsqueeze(0)\n\n # Convolve the audio signal with the impulse response\n reverbed_signal = fft_conv1d(tensor, impulse_response)\n\n # Normalize to the original amplitude range for stability\n reverbed_signal = (\n reverbed_signal\n / torch.max(torch.abs(reverbed_signal))\n * torch.max(torch.abs(tensor))\n )\n\n # Ensure tensor size is not changed\n tmp = torch.zeros_like(tensor)\n tmp[..., : reverbed_signal.shape[-1]] = reverbed_signal\n reverbed_signal = tmp\n\n return audio_effect_return(tensor=reverbed_signal, mask=mask)\n\n @staticmethod\n def random_noise(\n waveform: torch.Tensor,\n noise_std: float = 0.001,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Add Gaussian noise to the waveform.\"\"\"\n noise = torch.randn_like(waveform) * noise_std\n noisy_waveform = waveform + noise\n return audio_effect_return(tensor=noisy_waveform, mask=mask)\n\n @staticmethod\n def pink_noise(\n waveform: torch.Tensor,\n noise_std: float = 0.01,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Add pink background noise to the waveform.\"\"\"\n noise = generate_pink_noise(waveform.shape[-1]) * noise_std\n noise = noise.to(waveform.device)\n # Assuming waveform is of shape (bsz, channels, length)\n noisy_waveform = waveform + noise.unsqueeze(0).unsqueeze(0).to(waveform.device)\n return audio_effect_return(tensor=noisy_waveform, mask=mask)\n\n @staticmethod\n def lowpass_filter(\n waveform: torch.Tensor,\n cutoff_freq: float = 5000,\n sample_rate: int = 16000,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Filter the lowpass frequency from the waveform\"\"\"\n return audio_effect_return(\n tensor=julius.lowpass_filter(waveform, cutoff=cutoff_freq / sample_rate),\n mask=mask,\n )\n\n @staticmethod\n def highpass_filter(\n waveform: torch.Tensor,\n cutoff_freq: float = 500,\n sample_rate: int = 16000,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Filter the highpass frequency from the waveform\"\"\"\n return audio_effect_return(\n tensor=julius.highpass_filter(waveform, cutoff=cutoff_freq / sample_rate),\n mask=mask,\n )\n\n @staticmethod\n def bandpass_filter(\n waveform: torch.Tensor,\n cutoff_freq_low: float = 300,\n cutoff_freq_high: float = 8000,\n sample_rate: int = 16000,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Apply a bandpass filter to the waveform by cascading\n a high-pass filter followed by a low-pass filter.\n\n Args:\n waveform (torch.Tensor): Input audio waveform.\n low_cutoff (float): Lower cutoff frequency.\n high_cutoff (float): Higher cutoff frequency.\n sample_rate (int): The sample rate of the waveform.\n\n Returns:\n torch.Tensor: Filtered audio waveform.\n \"\"\"\n\n return audio_effect_return(\n tensor=julius.bandpass_filter(\n waveform,\n cutoff_low=cutoff_freq_low / sample_rate,\n cutoff_high=cutoff_freq_high / sample_rate,\n ),\n mask=mask,\n )\n\n @staticmethod\n def smooth(\n tensor: torch.Tensor,\n window_size_range: tuple = (2, 10),\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Smooths the input tensor (audio signal) using a moving average filter with the\n given window size.\n\n Args:\n tensor (torch.Tensor): Input audio tensor. Assumes tensor shape is (batch_size,\n channels, time).\n window_size (int): Size of the moving average window.\n mask: Masks for the input wave\n\n Returns:\n torch.Tensor: Smoothed audio tensor.\n \"\"\"\n\n window_size = int(torch.FloatTensor(1).uniform_(*window_size_range))\n # Create a uniform smoothing kernel\n kernel = torch.ones(1, 1, window_size).type(tensor.type()) / window_size\n kernel = kernel.to(tensor.device)\n\n smoothed = fft_conv1d(tensor, kernel)\n # Ensure tensor size is not changed\n tmp = torch.zeros_like(tensor)\n tmp[..., : smoothed.shape[-1]] = smoothed\n smoothed = tmp\n\n return audio_effect_return(tensor=smoothed, mask=mask)\n\n @staticmethod\n def boost_audio(\n tensor: torch.Tensor,\n amount: float = 20,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Filter the lowpass frequency from the waveform\"\"\"\n return audio_effect_return(tensor=tensor * (1 + amount / 100), mask=mask)\n\n @staticmethod\n def duck_audio(\n tensor: torch.Tensor,\n amount: float = 20,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Mask input wav with some ducked signnals\"\"\"\n return audio_effect_return(tensor=tensor * (1 - amount / 100), mask=mask)\n\n @staticmethod\n def identity(\n tensor: torch.Tensor, mask: tp.Optional[torch.Tensor] = None\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n return audio_effect_return(tensor=tensor, mask=mask)\n\n @staticmethod\n def mp3_compression(\n tensor: torch.Tensor,\n sample_rate: int = 16000,\n bitrate: str = \"128k\",\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"\n Compress audio using MP3 algorithm\n Args:\n tensor (torch.Tensor): The input audio tensor.\n sample_rate (int): The sample rate of the audio.\n bitrate (str): The bitrate for MP3 compression.\n\n Returns:\n torch.Tensor: The output tensor after applying MP3 compression.\n \"\"\"\n out = apply_compression_skip_grad(\n tensor, get_mp3, sr=sample_rate, bitrate=bitrate\n )\n return audio_effect_return(tensor=out, mask=mask)\n\n @staticmethod\n def aac_compression(\n tensor: torch.Tensor,\n sample_rate: int = 16000,\n bitrate: str = \"128k\",\n lowpass_freq: tp.Optional[int] = None,\n mask: tp.Optional[torch.Tensor] = None,\n ) -> tp.Union[tp.Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:\n \"\"\"Applies AAC compression to an audio tensor.\n\n Args:\n tensor (torch.Tensor): The input audio tensor.\n sample_rate (int): The sample rate of the audio.\n bitrate (str): The bitrate for AAC compression.\n lowpass_freq (Optional[int]): The frequency for a low-pass filter.\n\n Returns:\n torch.Tensor: The output tensor after applying AAC compression.\n \"\"\"\n out = apply_compression_skip_grad(\n tensor, get_aac, sr=sample_rate, bitrate=bitrate, lowpass_freq=lowpass_freq\n )\n return audio_effect_return(tensor=out, mask=mask)\n" }, { "path": "audiocraft/utils/autocast.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\n\n\nclass TorchAutocast:\n \"\"\"TorchAutocast utility class.\n Allows you to enable and disable autocast. This is specially useful\n when dealing with different architectures and clusters with different\n levels of support.\n\n Args:\n enabled (bool): Whether to enable torch.autocast or not.\n args: Additional args for torch.autocast.\n kwargs: Additional kwargs for torch.autocast\n \"\"\"\n def __init__(self, enabled: bool, *args, **kwargs):\n self.autocast = torch.autocast(*args, **kwargs) if enabled else None\n\n def __enter__(self):\n if self.autocast is None:\n return\n try:\n self.autocast.__enter__()\n except RuntimeError:\n device = self.autocast.device\n dtype = self.autocast.fast_dtype\n raise RuntimeError(\n f\"There was an error autocasting with dtype={dtype} device={device}\\n\"\n \"If you are on the FAIR Cluster, you might need to use autocast_dtype=float16\"\n )\n\n def __exit__(self, *args, **kwargs):\n if self.autocast is None:\n return\n self.autocast.__exit__(*args, **kwargs)\n" }, { "path": "audiocraft/utils/best_state.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom collections import defaultdict\nimport logging\nimport typing as tp\n\nimport flashy\nimport torch\n\nfrom ..optim import ModuleDictEMA\nfrom .utils import copy_state\n\n\nlogger = logging.getLogger(__name__)\n\n\nclass BestStateDictManager(flashy.state.StateDictSource):\n \"\"\"BestStateDictManager maintains a copy of best state_dict() for registered sources.\n\n BestStateDictManager has two main attributes:\n states (dict): State dict of the registered StateDictSource.\n param_ids (dict): Dict of parameter ids for registered states from ModuleDictEMA and other sources.\n\n When registering new sources, the BestStateDictManager will ensure two conflicting sources between\n ModuleDictEMA and original modules are not both registered as it would otherwise create ambiguity about\n what to consider for best state.\n\n Args:\n device (torch.device or str): Device on which we keep the copy.\n dtype (torch.dtype): Data type for the state parameters.\n \"\"\"\n def __init__(self, device: tp.Union[torch.device, str] = 'cpu',\n dtype: tp.Optional[torch.dtype] = None):\n self.device = device\n self.states: dict = {}\n self.param_ids: dict = defaultdict(dict)\n self.dtype = dtype\n\n def _get_parameter_ids(self, state_dict):\n return {id(p): name for name, p in state_dict.items() if isinstance(p, torch.Tensor)}\n\n def _validate_no_parameter_ids_overlap(self, name: str, param_ids: dict):\n for registered_name, registered_param_ids in self.param_ids.items():\n if registered_name != name:\n overlap = set.intersection(registered_param_ids.keys(), param_ids.keys())\n assert len(overlap) == 0, f\"Found {len(overlap)} / {len(param_ids.keys())} overlapping parameters\"\n f\" in {name} and already registered {registered_name}: {' '.join(overlap)}\"\n\n def update(self, name: str, source: flashy.state.StateDictSource):\n if name not in self.states:\n raise ValueError(f\"{name} missing from registered states.\")\n self.states[name] = copy_state(source.state_dict(), device=self.device, dtype=self.dtype)\n\n def register(self, name: str, source: flashy.state.StateDictSource):\n if name in self.states:\n raise ValueError(f\"{name} already present in states.\")\n # Registering parameter ids for EMA and non-EMA states allows us to check that\n # there is no overlap that would create ambiguity about how to handle the best state\n param_ids = self._get_parameter_ids(source.state_dict())\n if isinstance(source, ModuleDictEMA):\n logger.debug(f\"Registering to best state: ModuleDictEMA '{name}' with {len(param_ids)} params\")\n self._validate_no_parameter_ids_overlap(name, param_ids)\n self.param_ids[name] = param_ids\n else:\n logger.debug(f\"Registering to best state: StateDictSource '{name}' with {len(param_ids)} params\")\n self._validate_no_parameter_ids_overlap('base', param_ids)\n self.param_ids['base'].update(param_ids)\n # Register state\n self.states[name] = copy_state(source.state_dict(), device=self.device, dtype=self.dtype)\n\n def state_dict(self) -> flashy.state.StateDict:\n return self.states\n\n def load_state_dict(self, state: flashy.state.StateDict):\n for name, sub_state in state.items():\n for k, v in sub_state.items():\n self.states[name][k].copy_(v)\n" }, { "path": "audiocraft/utils/cache.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom concurrent.futures import ThreadPoolExecutor\nfrom collections import deque\nfrom functools import partial\nfrom hashlib import sha1\nimport logging\nfrom pathlib import Path\nimport sys\nimport typing as tp\nimport zipfile\n\nimport flashy\nimport torch\n\n\nlogger = logging.getLogger(__name__)\n\n\ndef get_full_embed(full_embed: torch.Tensor, x: tp.Any, idx: int, device: tp.Union[str, torch.device]) -> torch.Tensor:\n \"\"\"Utility function for the EmbeddingCache, returning the full embedding without any chunking.\n This method can be used in case there is no need in extracting a chunk of the full embedding\n read from the cache.\n\n Args:\n full_embed (torch.Tensor): The full embedding.\n x (any): Batch object from which the full embedding is derived.\n idx (torch.Tensor): Index of object to consider in the batch object.\n Returns:\n full_embed (torch.Tensor): The full embedding\n \"\"\"\n return full_embed.to(device)\n\n\nclass EmbeddingCache:\n \"\"\"Cache around embeddings computation for faster execution.\n The EmbeddingCache is storing pre-computed embeddings on disk and provides a simple API\n to retrieve the pre-computed embeddings on full inputs and extract only a given chunk\n using a user-provided function. When the cache is warm (all embeddings are pre-computed),\n the EmbeddingCache allows for faster training as it removes the need of computing the embeddings.\n Additionally, it provides in-memory cache around the loaded embeddings to limit IO footprint\n and synchronization points in the forward calls.\n\n Args:\n cache_path (Path): Path to folder where all pre-computed embeddings are saved on disk.\n device (str or torch.device): Device on which the embedding is returned.\n compute_embed_fn (callable[[Path, any, int], torch.Tensor], optional): Function to compute\n the embedding from a given object and path. This user provided function can compute the\n embedding from the provided object or using the provided path as entry point. The last parameter\n specify the index corresponding to the current embedding in the object that can represent batch metadata.\n extract_embed_fn (callable[[torch.Tensor, any, int], torch.Tensor], optional): Function to extract\n the desired embedding chunk from the full embedding loaded from the cache. The last parameter\n specify the index corresponding to the current embedding in the object that can represent batch metadata.\n If not specified, will return the full embedding unmodified.\n \"\"\"\n def __init__(self, cache_path: tp.Union[str, Path], device: tp.Union[str, torch.device],\n compute_embed_fn: tp.Callable[[Path, tp.Any, int], torch.Tensor],\n extract_embed_fn: tp.Optional[tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor]] = None):\n self.cache_path = Path(cache_path)\n self.device = device\n self._compute_embed_fn = compute_embed_fn\n self._extract_embed_fn: tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor]\n if extract_embed_fn is not None:\n self._extract_embed_fn = extract_embed_fn\n else:\n self._extract_embed_fn = partial(get_full_embed, device=device)\n if self.cache_path is not None:\n self.cache_path.mkdir(exist_ok=True, parents=True)\n logger.info(f\"Cache instantiated at: {self.cache_path}\")\n self.pool = ThreadPoolExecutor(8)\n self.pool.__enter__()\n self._current_batch_cache: dict = {}\n self._memory_cache: dict = {}\n\n def _get_cache_path(self, path: tp.Union[Path, str]):\n \"\"\"Get cache path for the given file path.\"\"\"\n sig = sha1(str(path).encode()).hexdigest()\n return self.cache_path / sig\n\n @staticmethod\n def _get_full_embed_from_cache(cache: Path):\n \"\"\"Loads full pre-computed embedding from the cache.\"\"\"\n try:\n embed = torch.load(cache, 'cpu')\n except Exception as exc:\n logger.error(\"Error loading %s: %r\", cache, exc)\n embed = None\n return embed\n\n def get_embed_from_cache(self, paths: tp.List[Path], x: tp.Any) -> torch.Tensor:\n \"\"\"Get embedding from cache, computing and storing it to cache if not already cached.\n The EmbeddingCache first tries to load the embedding from the in-memory cache\n containing the pre-computed chunks populated through `populate_embed_cache`.\n If not found, the full embedding is computed and stored on disk to be later accessed\n to populate the in-memory cache, and the desired embedding chunk is extracted and returned.\n\n Args:\n paths (list[Path or str]): List of paths from where the embeddings can be loaded.\n x (any): Object from which the embedding is extracted.\n \"\"\"\n embeds = []\n for idx, path in enumerate(paths):\n cache = self._get_cache_path(path)\n if cache in self._current_batch_cache:\n embed = self._current_batch_cache[cache]\n else:\n full_embed = self._compute_embed_fn(path, x, idx)\n try:\n with flashy.utils.write_and_rename(cache, pid=True) as f:\n torch.save(full_embed.cpu(), f)\n except Exception as exc:\n logger.error('Error saving embed %s (%s): %r', cache, full_embed.shape, exc)\n else:\n logger.info('New embed cache saved: %s (%s)', cache, full_embed.shape)\n embed = self._extract_embed_fn(full_embed, x, idx)\n embeds.append(embed)\n embed = torch.stack(embeds, dim=0)\n return embed\n\n def populate_embed_cache(self, paths: tp.List[Path], x: tp.Any) -> None:\n \"\"\"Populate in-memory caches for embeddings reading from the embeddings stored on disk.\n The in-memory caches consist in a cache for the full embedding and another cache for the\n final embedding chunk. Such caches are used to limit the IO access when computing the actual embeddings\n and reduce the IO footprint and synchronization points during forward passes.\n\n Args:\n paths (list[Path]): List of paths from where the embeddings can be loaded.\n x (any): Object from which the embedding is extracted.\n \"\"\"\n self._current_batch_cache.clear()\n if self.cache_path is not None:\n futures: list = []\n for path in paths:\n assert path is not None, \"Path is required for computation from cache\"\n cache = self._get_cache_path(path)\n if cache in self._memory_cache or not cache.exists():\n futures.append(None)\n else:\n futures.append(self.pool.submit(EmbeddingCache._get_full_embed_from_cache, cache))\n for idx, (path, future) in enumerate(zip(paths, futures)):\n assert path is not None\n cache = self._get_cache_path(path)\n full_embed = None\n if future is None:\n if cache in self._memory_cache:\n full_embed = self._memory_cache[cache]\n else:\n full_embed = future.result()\n if full_embed is not None:\n self._memory_cache[cache] = full_embed\n full_embed = full_embed.to(self.device)\n if full_embed is not None:\n embed = self._extract_embed_fn(full_embed, x, idx)\n self._current_batch_cache[cache] = embed\n\n\nclass CachedBatchWriter:\n \"\"\"Write pre computed caches for mini batches. This can\n make loading a lot more efficient depending on your filesystem.\n\n Args:\n cache_folder (Path): folder in which the cached minibatches\n will be stored.\n\n Inside cache folder, the structure is the following:\n `epoch_number / update_number.zip`\n And the zip file contains one entry per batch item.\n\n It is possible to use the cache with a batch size smaller than\n created with but obviously not larger. Make sure to call the\n `start_epoch(epoch)` method for indicating changes of epochs.\n\n See the grid `audiocraft/grids/musicgen/musicgen_warmup_cache.py`\n for an example of how to warmup the cache.\n \"\"\"\n def __init__(self, cache_folder: Path):\n self.cache_folder = cache_folder\n self._current_epoch: tp.Optional[int] = None\n self._current_index = 0\n\n def start_epoch(self, epoch: int):\n \"\"\"Call at the beginning of each epoch.\n \"\"\"\n self._current_epoch = epoch\n self._current_index = 0\n self._zip_path.parent.mkdir(exist_ok=True, parents=True)\n\n @staticmethod\n def _get_zip_path(cache_folder: Path, epoch: int, index: int):\n return cache_folder / f\"{epoch:05d}\" / f\"{index:06d}.zip\"\n\n @property\n def _zip_path(self):\n assert self._current_epoch is not None\n return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, self._current_index)\n\n def save(self, *content):\n \"\"\"Save one mini batch. This function is distributed-aware\n and will automatically merge all the items from the different\n workers.\n \"\"\"\n all_contents = []\n for rank in range(flashy.distrib.world_size()):\n their_content = flashy.distrib.broadcast_object(content, src=rank)\n all_contents.append(their_content)\n\n if flashy.distrib.is_rank_zero():\n idx = 0\n with flashy.utils.write_and_rename(self._zip_path) as tmp:\n with zipfile.ZipFile(tmp, 'w') as zf:\n for content in all_contents:\n for vals in zip(*content):\n with zf.open(f'{idx}', 'w') as f: # type: ignore\n torch.save(vals, f)\n idx += 1\n flashy.distrib.barrier()\n self._current_index += 1\n\n\nclass CachedBatchLoader:\n \"\"\"Loader for cached mini-batches dumped with `CachedBatchWriter`.\n\n Args:\n cache_folder (Path): folder in which the cached minibatches are stored.\n batch_size (int): batch size (per GPU) expected.\n num_workers (int): number of workers to use for loading.\n min_length (int): minimum expected length for each epoch. If some\n mini-batches are missing, and error is raised.\n\n This is iterable just like a regular DataLoader.\n \"\"\"\n\n def __init__(self, cache_folder: Path, batch_size: int,\n num_workers: int = 10, min_length: int = 1):\n self.cache_folder = cache_folder\n self.batch_size = batch_size\n self.num_workers = num_workers\n self.min_length = min_length\n self._current_epoch: tp.Optional[int] = None\n self.sampler = None # for compatibility with the regular DataLoader\n\n def __len__(self):\n path = CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch or 0, 0).parent\n return len([p for p in path.iterdir() if p.suffix == \".zip\"])\n\n def start_epoch(self, epoch: int):\n \"\"\"Call at the beginning of each epoch.\n \"\"\"\n self._current_epoch = epoch\n\n def _zip_path(self, index: int):\n assert self._current_epoch is not None\n return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, index)\n\n def _load_one(self, index: int):\n zip_path = self._zip_path(index)\n if not zip_path.exists():\n if index < self.min_length:\n raise RuntimeError(f\"Cache should have at least {self.min_length} batches, but {index} doesn't exist\")\n\n return None\n mode = \"rb\" if sys.version_info >= (3, 9) else \"r\"\n try:\n with zipfile.ZipFile(zip_path, 'r') as zf:\n rank = flashy.distrib.rank()\n world_size = flashy.distrib.world_size()\n root = zipfile.Path(zf)\n items = list(root.iterdir())\n total_batch_size = self.batch_size * world_size\n if len(items) < total_batch_size:\n raise RuntimeError(\n f\"The cache can handle a max batch size of {len(items)}, \"\n f\"but {total_batch_size} is needed.\")\n start = rank * self.batch_size\n items = items[start: start + self.batch_size]\n assert len(items) == self.batch_size\n entries = []\n entries = [torch.load(item.open(mode), 'cpu') for item in items] # type: ignore\n transposed = zip(*entries)\n out = []\n for part in transposed:\n assert len(part) > 0\n if isinstance(part[0], torch.Tensor):\n out.append(torch.stack(part))\n else:\n assert isinstance(part, torch.Tensor)\n out.append(part)\n return out\n except Exception:\n logger.error(\"Error when reading zip path %s\", zip_path)\n raise\n\n def __iter__(self):\n \"\"\"This will yields tuples, exactly as provided to the\n `CachedBatchWriter.save` method.\n \"\"\"\n pool = ThreadPoolExecutor(self.num_workers)\n next_index = 0\n queue = deque()\n\n def _get_next():\n nonlocal next_index\n r = queue.popleft().result()\n if r is None:\n return None\n else:\n queue.append(pool.submit(self._load_one, next_index))\n next_index += 1\n return r\n\n with pool:\n # fill the buffer of fetching jobs.\n for _ in range(2 * self.num_workers):\n queue.append(pool.submit(self._load_one, next_index))\n next_index += 1\n while True:\n batch = _get_next()\n if batch is None:\n return\n yield batch\n" }, { "path": "audiocraft/utils/checkpoint.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom enum import Enum\nimport logging\nfrom pathlib import Path\nimport re\nimport typing as tp\n\nimport flashy\nimport torch\n\nfrom ..environment import AudioCraftEnvironment\n\n\nlogger = logging.getLogger(__name__)\n\n\nclass CheckpointSource(Enum):\n CURRENT_XP = \"current_xp\"\n PRETRAINED = \"pretrained\"\n OTHER = \"other\"\n\n\ndef checkpoint_name(name: tp.Optional[str] = None, rank: tp.Optional[int] = None, use_fsdp: bool = False) -> str:\n \"\"\"Checkpoint name formatted for all use in AudioCraft codebase and has the following format:\n `checkpoint_.th(.)`. By convention, name is expected to be empty for last checkpoint,\n 'best' for the best checkpoint or the epoch number.\n\n Args:\n name (str, optional): Name suffix for the checkpoint file stem.\n rank (optional, int): Rank for distributed processing, retrieved with flashy if not provided.\n use_fsdp (bool): Whether the calling solver relies on FSDP.\n Returns:\n str: The checkpoint name.\n \"\"\"\n suffix = ''\n if rank is None:\n rank = flashy.distrib.rank()\n if rank > 0 and use_fsdp:\n suffix = '.' + str(rank)\n name_part = ''\n if name is not None:\n name_part = f'_{name}'\n return f'checkpoint{name_part}.th{suffix}'\n\n\ndef is_sharded_checkpoint(path: Path) -> bool:\n \"\"\"Whether the checkpoint at the given path corresponds to a sharded checkpoint across rank.\"\"\"\n return re.search(r'\\.th\\.\\d+$', path.name) is not None\n\n\ndef resolve_checkpoint_path(sig_or_path: tp.Union[Path, str], name: tp.Optional[str] = None,\n use_fsdp: bool = False) -> tp.Optional[Path]:\n \"\"\"Resolve a given checkpoint path for a provided dora sig or path.\n\n Args:\n sig_or_path (Path or str): Checkpoint path or dora signature.\n name (str, optional): Name suffix for the checkpoint file stem.\n rank (optional, int): Rank for distributed processing, retrieved with flashy if not provided.\n use_fsdp (bool): Whether the calling solver relies on FSDP.\n Returns:\n Path, optional: Resolved checkpoint path, if it exists.\n \"\"\"\n from audiocraft import train\n xps_root = train.main.dora.dir / 'xps'\n sig_or_path = str(sig_or_path)\n if sig_or_path.startswith('//sig/'):\n sig = sig_or_path[len('//sig/'):]\n path = xps_root / sig\n else:\n path = Path(sig_or_path)\n path = AudioCraftEnvironment.resolve_reference_path(path)\n\n if path.is_dir():\n path = path / checkpoint_name(name, use_fsdp=use_fsdp)\n\n if path.exists():\n return path\n else:\n return None\n\n\ndef load_checkpoint(checkpoint_path: Path, is_sharded: bool = False) -> tp.Any:\n \"\"\"Load state from checkpoints at the specified checkpoint path.\"\"\"\n if is_sharded:\n rank0_checkpoint_path = checkpoint_path.parent / checkpoint_name(use_fsdp=False)\n if rank0_checkpoint_path.exists():\n check_sharded_checkpoint(checkpoint_path, rank0_checkpoint_path)\n state = torch.load(checkpoint_path, 'cpu')\n logger.info(\"Checkpoint loaded from %s\", checkpoint_path)\n return state\n\n\ndef save_checkpoint(state: tp.Any, checkpoint_path: Path, is_sharded: bool = False) -> None:\n \"\"\"Save state to disk to the specified checkpoint_path.\"\"\"\n _safe_save_checkpoint(state, checkpoint_path, is_sharded)\n logger.info(\"Checkpoint saved to %s\", checkpoint_path)\n\n\ndef flush_stale_checkpoints(checkpoint_path: Path, keep_last: tp.Optional[int] = None) -> None:\n \"\"\"Flush checkpoints to only keep last N checkpoints.\"\"\"\n if keep_last is None or keep_last <= 0:\n return\n checkpoint_dir = checkpoint_path.parent\n suffix = ''\n if flashy.distrib.rank() > 0:\n suffix = f'.{flashy.distrib.rank()}'\n checkpoint_files_with_epoch = []\n for path in Path(checkpoint_dir).glob(f'checkpoint_*.th{suffix}'):\n epoch_part = path.name.split('.', 1)[0].split('_', 1)[1]\n if epoch_part.isdigit():\n checkpoint_files_with_epoch.append((path, int(epoch_part)))\n checkpoint_files = [path for path, _ in list(sorted(checkpoint_files_with_epoch, key=lambda t: t[1]))]\n total_to_flush = max(0, len(checkpoint_files) - keep_last)\n files_to_flush = checkpoint_files[:total_to_flush]\n for path in files_to_flush:\n logger.debug(\"Removing checkpoint: %s\", str(path))\n path.unlink(missing_ok=True)\n\n\ndef check_sharded_checkpoint(checkpoint_path: Path, rank0_checkpoint_path: Path) -> None:\n \"\"\"Check sharded checkpoint state, ensuring the checkpoints are not corrupted.\"\"\"\n # Finish the work of a previous run that got interrupted while dumping.\n old_path = Path(str(checkpoint_path) + '.old')\n if old_path.exists():\n raise RuntimeError(\n f\"Old checkpoint {old_path} from previous version of this code exist, cannot safely proceed.\")\n token = Path(str(rank0_checkpoint_path) + '.tmp.done')\n tmp_path = Path(str(checkpoint_path) + '.tmp')\n if token.exists():\n if tmp_path.exists():\n tmp_path.rename(checkpoint_path)\n flashy.distrib.barrier()\n if flashy.distrib.is_rank_zero() and token.exists():\n token.unlink()\n\n\ndef _safe_save_checkpoint(state: tp.Any, checkpoint_path: Path, is_sharded: bool = False) -> None:\n \"\"\"Save checkpoints in a safe manner even with when sharded checkpoints across nodes.\"\"\"\n def _barrier_if_sharded():\n if is_sharded:\n flashy.distrib.barrier()\n\n if flashy.distrib.is_rank_zero():\n token = Path(str(checkpoint_path) + '.tmp.done')\n if token.exists():\n token.unlink()\n _barrier_if_sharded()\n with flashy.utils.write_and_rename(checkpoint_path) as f:\n torch.save(state, f)\n _barrier_if_sharded()\n if flashy.distrib.is_rank_zero():\n token.touch()\n _barrier_if_sharded()\n _barrier_if_sharded()\n if flashy.distrib.rank() == 0:\n token.unlink()\n" }, { "path": "audiocraft/utils/cluster.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nUtility functions for SLURM configuration and cluster settings.\n\"\"\"\n\nfrom enum import Enum\nimport os\nimport socket\nimport typing as tp\n\nimport omegaconf\n\n\nclass ClusterType(Enum):\n AWS = \"aws\"\n FAIR = \"fair\"\n RSC = \"rsc\"\n LOCAL_DARWIN = \"darwin\"\n DEFAULT = \"default\" # used for any other cluster.\n\n\ndef _guess_cluster_type() -> ClusterType:\n uname = os.uname()\n fqdn = socket.getfqdn()\n if uname.sysname == \"Linux\" and (uname.release.endswith(\"-aws\") or \".ec2\" in fqdn):\n return ClusterType.AWS\n\n if fqdn.endswith(\".fair\"):\n return ClusterType.FAIR\n\n if fqdn.endswith(\".facebook.com\"):\n return ClusterType.RSC\n\n if uname.sysname == \"Darwin\":\n return ClusterType.LOCAL_DARWIN\n\n return ClusterType.DEFAULT\n\n\ndef get_cluster_type(\n cluster_type: tp.Optional[ClusterType] = None,\n) -> tp.Optional[ClusterType]:\n if cluster_type is None:\n return _guess_cluster_type()\n\n return cluster_type\n\n\ndef get_slurm_parameters(\n cfg: omegaconf.DictConfig, cluster_type: tp.Optional[ClusterType] = None\n) -> omegaconf.DictConfig:\n \"\"\"Update SLURM parameters in configuration based on cluster type.\n If the cluster type is not specify, it infers it automatically.\n \"\"\"\n from ..environment import AudioCraftEnvironment\n cluster_type = get_cluster_type(cluster_type)\n # apply cluster-specific adjustments\n if cluster_type == ClusterType.AWS:\n cfg[\"mem_per_gpu\"] = None\n cfg[\"constraint\"] = None\n cfg[\"setup\"] = []\n elif cluster_type == ClusterType.RSC:\n cfg[\"mem_per_gpu\"] = None\n cfg[\"setup\"] = []\n cfg[\"constraint\"] = None\n cfg[\"partition\"] = \"learn\"\n slurm_exclude = AudioCraftEnvironment.get_slurm_exclude()\n if slurm_exclude is not None:\n cfg[\"exclude\"] = slurm_exclude\n return cfg\n" }, { "path": "audiocraft/utils/deadlock.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nimport os\nfrom queue import Queue, Empty\nimport signal\nimport sys\nimport threading\nimport traceback\n\nlogger = logging.getLogger(__name__)\n\n\nclass DeadlockDetect:\n def __init__(self, use: bool = False, timeout: float = 120.):\n self.use = use\n self.timeout = timeout\n self._queue: Queue = Queue()\n\n def update(self, stage: str):\n if self.use:\n self._queue.put(stage)\n\n def __enter__(self):\n if self.use:\n self._thread = threading.Thread(target=self._detector_thread)\n self._thread.start()\n\n def __exit__(self, exc_type, exc_val, exc_tb):\n if self.use:\n self._queue.put(None)\n self._thread.join()\n\n def _detector_thread(self):\n logger.debug(\"Deadlock detector started\")\n last_stage = \"init\"\n while True:\n try:\n stage = self._queue.get(timeout=self.timeout)\n except Empty:\n break\n if stage is None:\n logger.debug(\"Exiting deadlock detector thread\")\n return\n else:\n last_stage = stage\n logger.error(\"Deadlock detector timed out, last stage was %s\", last_stage)\n for th in threading.enumerate():\n print(th, file=sys.stderr)\n traceback.print_stack(sys._current_frames()[th.ident])\n print(file=sys.stderr)\n sys.stdout.flush()\n sys.stderr.flush()\n os.kill(os.getpid(), signal.SIGKILL)\n" }, { "path": "audiocraft/utils/export.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nUtility to export a training checkpoint to a lightweight release checkpoint.\n\"\"\"\n\nfrom pathlib import Path\nimport typing as tp\n\nfrom omegaconf import OmegaConf\nimport torch\n\nfrom audiocraft import __version__\n\n\ndef export_encodec(checkpoint_path: tp.Union[Path, str], out_file: tp.Union[Path, str]):\n \"\"\"Export only the best state from the given EnCodec checkpoint. This\n should be used if you trained your own EnCodec model.\n \"\"\"\n pkg = torch.load(checkpoint_path, 'cpu')\n new_pkg = {\n 'best_state': pkg['best_state']['model'],\n 'xp.cfg': OmegaConf.to_yaml(pkg['xp.cfg']),\n 'version': __version__,\n 'exported': True,\n }\n Path(out_file).parent.mkdir(exist_ok=True, parents=True)\n torch.save(new_pkg, out_file)\n return out_file\n\n\ndef export_pretrained_compression_model(pretrained_encodec: str, out_file: tp.Union[Path, str]):\n \"\"\"Export a compression model (potentially EnCodec) from a pretrained model.\n This is required for packaging the audio tokenizer along a MusicGen or AudioGen model.\n Do not include the //pretrained/ prefix. For instance if you trained a model\n with `facebook/encodec_32khz`, just put that as a name. Same for `dac_44khz`.\n\n In that case, this will not actually include a copy of the model, simply the reference\n to the model used.\n \"\"\"\n if Path(pretrained_encodec).exists():\n pkg = torch.load(pretrained_encodec)\n assert 'best_state' in pkg\n assert 'xp.cfg' in pkg\n assert 'version' in pkg\n assert 'exported' in pkg\n else:\n pkg = {\n 'pretrained': pretrained_encodec,\n 'exported': True,\n 'version': __version__,\n }\n Path(out_file).parent.mkdir(exist_ok=True, parents=True)\n torch.save(pkg, out_file)\n\n\ndef export_lm(checkpoint_path: tp.Union[Path, str], out_file: tp.Union[Path, str]):\n \"\"\"Export only the best state from the given MusicGen or AudioGen checkpoint.\n \"\"\"\n pkg = torch.load(checkpoint_path, 'cpu')\n if pkg['fsdp_best_state']:\n best_state = pkg['fsdp_best_state']['model']\n else:\n assert pkg['best_state']\n best_state = pkg['best_state']['model']\n new_pkg = {\n 'best_state': best_state,\n 'xp.cfg': OmegaConf.to_yaml(pkg['xp.cfg']),\n 'version': __version__,\n 'exported': True,\n }\n\n Path(out_file).parent.mkdir(exist_ok=True, parents=True)\n torch.save(new_pkg, out_file)\n return out_file\n" }, { "path": "audiocraft/utils/export_legacy.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nLegacy functions used at the time of the first release, kept for referencd.\n\"\"\"\n\nfrom pathlib import Path\nimport typing as tp\n\nfrom omegaconf import OmegaConf, DictConfig\nimport torch\n\nfrom audiocraft import __version__\n\n\ndef _clean_lm_cfg(cfg: DictConfig):\n OmegaConf.set_struct(cfg, False)\n # This used to be set automatically in the LM solver, need a more robust solution\n # for the future.\n cfg['transformer_lm']['card'] = 2048\n n_q = 4\n stereo_cfg = getattr(cfg, 'interleave_stereo_codebooks', None)\n if stereo_cfg is not None and stereo_cfg.use:\n if 'downsample' in stereo_cfg:\n del stereo_cfg['downsample']\n n_q = 8\n cfg['transformer_lm']['n_q'] = n_q\n # Experimental params no longer supported.\n bad_params = ['spectral_norm_attn_iters', 'spectral_norm_ff_iters',\n 'residual_balancer_attn', 'residual_balancer_ff', 'layer_drop']\n for name in bad_params:\n del cfg['transformer_lm'][name]\n OmegaConf.set_struct(cfg, True)\n return cfg\n\n\ndef export_encodec(checkpoint_path: tp.Union[Path, str], out_file: tp.Union[Path, str]):\n pkg = torch.load(checkpoint_path, 'cpu')\n new_pkg = {\n 'best_state': pkg['ema']['state']['model'],\n 'xp.cfg': OmegaConf.to_yaml(pkg['xp.cfg']),\n # The following params were NOT exported for the first release of MusicGen.\n 'version': __version__,\n 'exported': True,\n }\n Path(out_file).parent.mkdir(exist_ok=True, parents=True)\n torch.save(new_pkg, out_file)\n return out_file\n\n\ndef export_lm(checkpoint_path: tp.Union[Path, str], out_file: tp.Union[Path, str]):\n pkg = torch.load(checkpoint_path, 'cpu')\n if pkg['fsdp_best_state']:\n best_state = pkg['fsdp_best_state']['model']\n else:\n best_state = pkg['best_state']['model']\n new_pkg = {\n 'best_state': best_state,\n 'xp.cfg': OmegaConf.to_yaml(_clean_lm_cfg(pkg['xp.cfg'])),\n # The following params were NOT exported for the first release of MusicGen.\n 'version': __version__,\n 'exported': True,\n }\n Path(out_file).parent.mkdir(exist_ok=True, parents=True)\n torch.save(new_pkg, out_file)\n return out_file\n" }, { "path": "audiocraft/utils/notebook.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\ntry:\n import IPython.display as ipd # type: ignore\nexcept ImportError:\n # Note in a notebook...\n pass\n\n\nimport torch\n\n\ndef display_audio(samples: torch.Tensor, sample_rate: int):\n \"\"\"Renders an audio player for the given audio samples.\n\n Args:\n samples (torch.Tensor): a Tensor of decoded audio samples\n with shapes [B, C, T] or [C, T]\n sample_rate (int): sample rate audio should be displayed with.\n \"\"\"\n assert samples.dim() == 2 or samples.dim() == 3\n\n samples = samples.detach().cpu()\n if samples.dim() == 2:\n samples = samples[None, ...]\n\n for audio in samples:\n ipd.display(ipd.Audio(audio, rate=sample_rate))\n" }, { "path": "audiocraft/utils/profiler.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nimport typing as tp\n\nimport dora\nimport torch\n\n\nlogger = logging.getLogger(__name__)\n\n\nclass Profiler:\n \"\"\"Context manager wrapper for xformers profiler.\n \"\"\"\n def __init__(self, module: torch.nn.Module, enabled: bool = False):\n self.profiler: tp.Optional[tp.Any] = None\n if enabled:\n from xformers.profiler import profile\n output_dir = dora.get_xp().folder / 'profiler_data'\n logger.info(\"Profiling activated, results with be saved to %s\", output_dir)\n self.profiler = profile(output_dir=output_dir, module=module)\n\n def step(self):\n if self.profiler is not None:\n self.profiler.step() # type: ignore\n\n def __enter__(self):\n if self.profiler is not None:\n return self.profiler.__enter__() # type: ignore\n\n def __exit__(self, exc_type, exc_value, exc_tb):\n if self.profiler is not None:\n return self.profiler.__exit__(exc_type, exc_value, exc_tb) # type: ignore\n" }, { "path": "audiocraft/utils/samples/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "audiocraft/utils/samples/manager.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"\nAPI that can manage the storage and retrieval of generated samples produced by experiments.\n\nIt offers the following benefits:\n* Samples are stored in a consistent way across epoch\n* Metadata about the samples can be stored and retrieved\n* Can retrieve audio\n* Identifiers are reliable and deterministic for prompted and conditioned samples\n* Can request the samples for multiple XPs, grouped by sample identifier\n* For no-input samples (not prompt and no conditions), samples across XPs are matched\n by sorting their identifiers\n\"\"\"\n\nfrom concurrent.futures import ThreadPoolExecutor\nfrom dataclasses import asdict, dataclass\nfrom functools import lru_cache\nimport hashlib\nimport json\nimport logging\nfrom pathlib import Path\nimport re\nimport typing as tp\nimport unicodedata\nimport uuid\n\nimport dora\nimport torch\n\nfrom ...data.audio import audio_read, audio_write\n\n\nlogger = logging.getLogger(__name__)\n\n\n@dataclass\nclass ReferenceSample:\n id: str\n path: str\n duration: float\n\n\n@dataclass\nclass Sample:\n id: str\n path: str\n epoch: int\n duration: float\n conditioning: tp.Optional[tp.Dict[str, tp.Any]]\n prompt: tp.Optional[ReferenceSample]\n reference: tp.Optional[ReferenceSample]\n generation_args: tp.Optional[tp.Dict[str, tp.Any]]\n\n def __hash__(self):\n return hash(self.id)\n\n def audio(self) -> tp.Tuple[torch.Tensor, int]:\n return audio_read(self.path)\n\n def audio_prompt(self) -> tp.Optional[tp.Tuple[torch.Tensor, int]]:\n return audio_read(self.prompt.path) if self.prompt is not None else None\n\n def audio_reference(self) -> tp.Optional[tp.Tuple[torch.Tensor, int]]:\n return audio_read(self.reference.path) if self.reference is not None else None\n\n\nclass SampleManager:\n \"\"\"Audio samples IO handling within a given dora xp.\n\n The sample manager handles the dumping and loading logic for generated and\n references samples across epochs for a given xp, providing a simple API to\n store, retrieve and compare audio samples.\n\n Args:\n xp (dora.XP): Dora experiment object. The XP contains information on the XP folder\n where all outputs are stored and the configuration of the experiment,\n which is useful to retrieve audio-related parameters.\n map_reference_to_sample_id (bool): Whether to use the sample_id for all reference samples\n instead of generating a dedicated hash id. This is useful to allow easier comparison\n with ground truth sample from the files directly without having to read the JSON metadata\n to do the mapping (at the cost of potentially dumping duplicate prompts/references\n depending on the task).\n \"\"\"\n def __init__(self, xp: dora.XP, map_reference_to_sample_id: bool = False):\n self.xp = xp\n self.base_folder: Path = xp.folder / xp.cfg.generate.path\n self.reference_folder = self.base_folder / 'reference'\n self.map_reference_to_sample_id = map_reference_to_sample_id\n self.samples: tp.List[Sample] = []\n self._load_samples()\n\n @property\n def latest_epoch(self):\n \"\"\"Latest epoch across all samples.\"\"\"\n return max(self.samples, key=lambda x: x.epoch).epoch if self.samples else 0\n\n def _load_samples(self):\n \"\"\"Scan the sample folder and load existing samples.\"\"\"\n jsons = self.base_folder.glob('**/*.json')\n with ThreadPoolExecutor(6) as pool:\n self.samples = list(pool.map(self._load_sample, jsons))\n\n @staticmethod\n @lru_cache(2**26)\n def _load_sample(json_file: Path) -> Sample:\n with open(json_file, 'r') as f:\n data: tp.Dict[str, tp.Any] = json.load(f)\n # fetch prompt data\n prompt_data = data.get('prompt')\n prompt = ReferenceSample(id=prompt_data['id'], path=prompt_data['path'],\n duration=prompt_data['duration']) if prompt_data else None\n # fetch reference data\n reference_data = data.get('reference')\n reference = ReferenceSample(id=reference_data['id'], path=reference_data['path'],\n duration=reference_data['duration']) if reference_data else None\n # build sample object\n return Sample(id=data['id'], path=data['path'], epoch=data['epoch'], duration=data['duration'],\n prompt=prompt, conditioning=data.get('conditioning'), reference=reference,\n generation_args=data.get('generation_args'))\n\n def _init_hash(self):\n return hashlib.sha1()\n\n def _get_tensor_id(self, tensor: torch.Tensor) -> str:\n hash_id = self._init_hash()\n hash_id.update(tensor.numpy().data)\n return hash_id.hexdigest()\n\n def _get_sample_id(self, index: int, prompt_wav: tp.Optional[torch.Tensor],\n conditions: tp.Optional[tp.Dict[str, str]]) -> str:\n \"\"\"Computes an id for a sample given its input data.\n This id is deterministic if prompt and/or conditions are provided by using a sha1 hash on the input.\n Otherwise, a random id of the form \"noinput_{uuid4().hex}\" is returned.\n\n Args:\n index (int): Batch index, Helpful to differentiate samples from the same batch.\n prompt_wav (torch.Tensor): Prompt used during generation.\n conditions (dict[str, str]): Conditioning used during generation.\n \"\"\"\n # For totally unconditioned generations we will just use a random UUID.\n # The function get_samples_for_xps will do a simple ordered match with a custom key.\n if prompt_wav is None and not conditions:\n return f\"noinput_{uuid.uuid4().hex}\"\n\n # Human readable portion\n hr_label = \"\"\n # Create a deterministic id using hashing\n hash_id = self._init_hash()\n hash_id.update(f\"{index}\".encode())\n if prompt_wav is not None:\n hash_id.update(prompt_wav.numpy().data)\n hr_label += \"_prompted\"\n else:\n hr_label += \"_unprompted\"\n if conditions:\n encoded_json = json.dumps(conditions, sort_keys=True).encode()\n hash_id.update(encoded_json)\n cond_str = \"-\".join([f\"{key}={slugify(value)}\"\n for key, value in sorted(conditions.items())])\n cond_str = cond_str[:100] # some raw text might be too long to be a valid filename\n cond_str = cond_str if len(cond_str) > 0 else \"unconditioned\"\n hr_label += f\"_{cond_str}\"\n else:\n hr_label += \"_unconditioned\"\n\n return hash_id.hexdigest() + hr_label\n\n def _store_audio(self, wav: torch.Tensor, stem_path: Path, overwrite: bool = False) -> Path:\n \"\"\"Stores the audio with the given stem path using the XP's configuration.\n\n Args:\n wav (torch.Tensor): Audio to store.\n stem_path (Path): Path in sample output directory with file stem to use.\n overwrite (bool): When False (default), skips storing an existing audio file.\n Returns:\n Path: The path at which the audio is stored.\n \"\"\"\n existing_paths = [\n path for path in stem_path.parent.glob(stem_path.stem + '.*')\n if path.suffix != '.json'\n ]\n exists = len(existing_paths) > 0\n if exists and overwrite:\n logger.warning(f\"Overwriting existing audio file with stem path {stem_path}\")\n elif exists:\n return existing_paths[0]\n\n audio_path = audio_write(stem_path, wav, **self.xp.cfg.generate.audio)\n return audio_path\n\n def add_sample(self, sample_wav: torch.Tensor, epoch: int, index: int = 0,\n conditions: tp.Optional[tp.Dict[str, str]] = None, prompt_wav: tp.Optional[torch.Tensor] = None,\n ground_truth_wav: tp.Optional[torch.Tensor] = None,\n generation_args: tp.Optional[tp.Dict[str, tp.Any]] = None) -> Sample:\n \"\"\"Adds a single sample.\n The sample is stored in the XP's sample output directory, under a corresponding epoch folder.\n Each sample is assigned an id which is computed using the input data. In addition to the\n sample itself, a json file containing associated metadata is stored next to it.\n\n Args:\n sample_wav (torch.Tensor): sample audio to store. Tensor of shape [channels, shape].\n epoch (int): current training epoch.\n index (int): helpful to differentiate samples from the same batch.\n conditions (dict[str, str], optional): conditioning used during generation.\n prompt_wav (torch.Tensor, optional): prompt used during generation. Tensor of shape [channels, shape].\n ground_truth_wav (torch.Tensor, optional): reference audio where prompt was extracted from.\n Tensor of shape [channels, shape].\n generation_args (dict[str, any], optional): dictionary of other arguments used during generation.\n Returns:\n Sample: The saved sample.\n \"\"\"\n sample_id = self._get_sample_id(index, prompt_wav, conditions)\n reuse_id = self.map_reference_to_sample_id\n prompt, ground_truth = None, None\n if prompt_wav is not None:\n prompt_id = sample_id if reuse_id else self._get_tensor_id(prompt_wav.sum(0, keepdim=True))\n prompt_duration = prompt_wav.shape[-1] / self.xp.cfg.sample_rate\n prompt_path = self._store_audio(prompt_wav, self.base_folder / str(epoch) / 'prompt' / prompt_id)\n prompt = ReferenceSample(prompt_id, str(prompt_path), prompt_duration)\n if ground_truth_wav is not None:\n ground_truth_id = sample_id if reuse_id else self._get_tensor_id(ground_truth_wav.sum(0, keepdim=True))\n ground_truth_duration = ground_truth_wav.shape[-1] / self.xp.cfg.sample_rate\n ground_truth_path = self._store_audio(ground_truth_wav, self.base_folder / 'reference' / ground_truth_id)\n ground_truth = ReferenceSample(ground_truth_id, str(ground_truth_path), ground_truth_duration)\n sample_path = self._store_audio(sample_wav, self.base_folder / str(epoch) / sample_id, overwrite=True)\n duration = sample_wav.shape[-1] / self.xp.cfg.sample_rate\n sample = Sample(sample_id, str(sample_path), epoch, duration, conditions, prompt, ground_truth, generation_args)\n self.samples.append(sample)\n with open(sample_path.with_suffix('.json'), 'w') as f:\n json.dump(asdict(sample), f, indent=2)\n return sample\n\n def add_samples(self, samples_wavs: torch.Tensor, epoch: int,\n conditioning: tp.Optional[tp.List[tp.Dict[str, tp.Any]]] = None,\n prompt_wavs: tp.Optional[torch.Tensor] = None,\n ground_truth_wavs: tp.Optional[torch.Tensor] = None,\n generation_args: tp.Optional[tp.Dict[str, tp.Any]] = None) -> tp.List[Sample]:\n \"\"\"Adds a batch of samples.\n The samples are stored in the XP's sample output directory, under a corresponding\n epoch folder. Each sample is assigned an id which is computed using the input data and their batch index.\n In addition to the sample itself, a json file containing associated metadata is stored next to it.\n\n Args:\n sample_wavs (torch.Tensor): Batch of audio wavs to store. Tensor of shape [batch_size, channels, shape].\n epoch (int): Current training epoch.\n conditioning (list of dict[str, str], optional): List of conditions used during generation,\n one per sample in the batch.\n prompt_wavs (torch.Tensor, optional): Prompts used during generation. Tensor of shape\n [batch_size, channels, shape].\n ground_truth_wav (torch.Tensor, optional): Reference audio where prompts were extracted from.\n Tensor of shape [batch_size, channels, shape].\n generation_args (dict[str, Any], optional): Dictionary of other arguments used during generation.\n Returns:\n samples (list of Sample): The saved audio samples with prompts, ground truth and metadata.\n \"\"\"\n samples = []\n for idx, wav in enumerate(samples_wavs):\n prompt_wav = prompt_wavs[idx] if prompt_wavs is not None else None\n gt_wav = ground_truth_wavs[idx] if ground_truth_wavs is not None else None\n conditions = conditioning[idx] if conditioning is not None else None\n samples.append(self.add_sample(wav, epoch, idx, conditions, prompt_wav, gt_wav, generation_args))\n return samples\n\n def get_samples(self, epoch: int = -1, max_epoch: int = -1, exclude_prompted: bool = False,\n exclude_unprompted: bool = False, exclude_conditioned: bool = False,\n exclude_unconditioned: bool = False) -> tp.Set[Sample]:\n \"\"\"Returns a set of samples for this XP. Optionally, you can filter which samples to obtain.\n Please note that existing samples are loaded during the manager's initialization, and added samples through this\n manager are also tracked. Any other external changes are not tracked automatically, so creating a new manager\n is the only way detect them.\n\n Args:\n epoch (int): If provided, only return samples corresponding to this epoch.\n max_epoch (int): If provided, only return samples corresponding to the latest epoch that is <= max_epoch.\n exclude_prompted (bool): If True, does not include samples that used a prompt.\n exclude_unprompted (bool): If True, does not include samples that did not use a prompt.\n exclude_conditioned (bool): If True, excludes samples that used conditioning.\n exclude_unconditioned (bool): If True, excludes samples that did not use conditioning.\n Returns:\n Samples (set of Sample): The retrieved samples matching the provided filters.\n \"\"\"\n if max_epoch >= 0:\n samples_epoch = max(sample.epoch for sample in self.samples if sample.epoch <= max_epoch)\n else:\n samples_epoch = self.latest_epoch if epoch < 0 else epoch\n samples = {\n sample\n for sample in self.samples\n if (\n (sample.epoch == samples_epoch) and\n (not exclude_prompted or sample.prompt is None) and\n (not exclude_unprompted or sample.prompt is not None) and\n (not exclude_conditioned or not sample.conditioning) and\n (not exclude_unconditioned or sample.conditioning)\n )\n }\n return samples\n\n\ndef slugify(value: tp.Any, allow_unicode: bool = False):\n \"\"\"Process string for safer file naming.\n\n Taken from https://github.com/django/django/blob/master/django/utils/text.py\n\n Convert to ASCII if 'allow_unicode' is False. Convert spaces or repeated\n dashes to single dashes. Remove characters that aren't alphanumerics,\n underscores, or hyphens. Convert to lowercase. Also strip leading and\n trailing whitespace, dashes, and underscores.\n \"\"\"\n value = str(value)\n if allow_unicode:\n value = unicodedata.normalize(\"NFKC\", value)\n else:\n value = (\n unicodedata.normalize(\"NFKD\", value)\n .encode(\"ascii\", \"ignore\")\n .decode(\"ascii\")\n )\n value = re.sub(r\"[^\\w\\s-]\", \"\", value.lower())\n return re.sub(r\"[-\\s]+\", \"-\", value).strip(\"-_\")\n\n\ndef _match_stable_samples(samples_per_xp: tp.List[tp.Set[Sample]]) -> tp.Dict[str, tp.List[Sample]]:\n # Create a dictionary of stable id -> sample per XP\n stable_samples_per_xp = [{\n sample.id: sample for sample in samples\n if sample.prompt is not None or sample.conditioning\n } for samples in samples_per_xp]\n # Set of all stable ids\n stable_ids = {id for samples in stable_samples_per_xp for id in samples.keys()}\n # Dictionary of stable id -> list of samples. If an XP does not have it, assign None\n stable_samples = {id: [xp.get(id) for xp in stable_samples_per_xp] for id in stable_ids}\n # Filter out ids that contain None values (we only want matched samples after all)\n # cast is necessary to avoid mypy linter errors.\n return {id: tp.cast(tp.List[Sample], samples) for id, samples in stable_samples.items() if None not in samples}\n\n\ndef _match_unstable_samples(samples_per_xp: tp.List[tp.Set[Sample]]) -> tp.Dict[str, tp.List[Sample]]:\n # For unstable ids, we use a sorted list since we'll match them in order\n unstable_samples_per_xp = [[\n sample for sample in sorted(samples, key=lambda x: x.id)\n if sample.prompt is None and not sample.conditioning\n ] for samples in samples_per_xp]\n # Trim samples per xp so all samples can have a match\n min_len = min([len(samples) for samples in unstable_samples_per_xp])\n unstable_samples_per_xp = [samples[:min_len] for samples in unstable_samples_per_xp]\n # Dictionary of index -> list of matched samples\n return {\n f'noinput_{i}': [samples[i] for samples in unstable_samples_per_xp] for i in range(min_len)\n }\n\n\ndef get_samples_for_xps(xps: tp.List[dora.XP], **kwargs) -> tp.Dict[str, tp.List[Sample]]:\n \"\"\"Gets a dictionary of matched samples across the given XPs.\n Each dictionary entry maps a sample id to a list of samples for that id. The number of samples per id\n will always match the number of XPs provided and will correspond to each XP in the same order given.\n In other words, only samples that can be match across all provided XPs will be returned\n in order to satisfy this rule.\n\n There are two types of ids that can be returned: stable and unstable.\n * Stable IDs are deterministic ids that were computed by the SampleManager given a sample's inputs\n (prompts/conditioning). This is why we can match them across XPs.\n * Unstable IDs are of the form \"noinput_{idx}\" and are generated on-the-fly, in order to map samples\n that used non-deterministic, random ids. This is the case for samples that did not use prompts or\n conditioning for their generation. This function will sort these samples by their id and match them\n by their index.\n\n Args:\n xps: a list of XPs to match samples from.\n start_epoch (int): If provided, only return samples corresponding to this epoch or newer.\n end_epoch (int): If provided, only return samples corresponding to this epoch or older.\n exclude_prompted (bool): If True, does not include samples that used a prompt.\n exclude_unprompted (bool): If True, does not include samples that did not use a prompt.\n exclude_conditioned (bool): If True, excludes samples that used conditioning.\n exclude_unconditioned (bool): If True, excludes samples that did not use conditioning.\n \"\"\"\n managers = [SampleManager(xp) for xp in xps]\n samples_per_xp = [manager.get_samples(**kwargs) for manager in managers]\n stable_samples = _match_stable_samples(samples_per_xp)\n unstable_samples = _match_unstable_samples(samples_per_xp)\n return dict(stable_samples, **unstable_samples)\n" }, { "path": "audiocraft/utils/utils.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom concurrent.futures import ProcessPoolExecutor\nfrom contextlib import contextmanager\nfrom functools import wraps, lru_cache\nimport hashlib\nimport json\nimport logging\nfrom pathlib import Path\nimport typing as tp\nimport flashy\nimport flashy.distrib\nimport omegaconf\nimport torch\nfrom torch.nn.utils.rnn import pad_sequence\n\n\nlogger = logging.getLogger(__name__)\n\n\ndef model_hash(model: torch.nn.Module) -> str:\n \"\"\"Return a model hash. This should allow us to track regressions in model init\n from the logs of past experiments.\n \"\"\"\n hasher = hashlib.sha1()\n for p in model.parameters():\n hasher.update(p.data.cpu().numpy().tobytes())\n return hasher.hexdigest()\n\n\ndef dict_from_config(cfg: omegaconf.DictConfig) -> dict:\n \"\"\"Convenience function to map an omegaconf configuration to a dictionary.\n\n Args:\n cfg (omegaconf.DictConfig): Original configuration to map to dict.\n Returns:\n dict: Config as dictionary object.\n \"\"\"\n dct = omegaconf.OmegaConf.to_container(cfg, resolve=True)\n assert isinstance(dct, dict)\n return dct\n\n\ndef random_subset(dataset, max_samples: int, seed: int = 42) -> torch.utils.data.Subset:\n if max_samples >= len(dataset):\n return dataset\n\n generator = torch.Generator().manual_seed(seed)\n perm = torch.randperm(len(dataset), generator=generator)\n return torch.utils.data.Subset(dataset, perm[:max_samples].tolist())\n\n\ndef get_loader(dataset, num_samples: tp.Optional[int], batch_size: int,\n num_workers: int, seed: int, **kwargs) -> torch.utils.data.DataLoader:\n \"\"\"Convenience function to load dataset into a dataloader with optional subset sampling.\n\n Args:\n dataset: Dataset to load.\n num_samples (Optional[int]): Number of samples to limit subset size.\n batch_size (int): Batch size.\n num_workers (int): Number of workers for data loading.\n seed (int): Random seed.\n \"\"\"\n if num_samples is not None:\n dataset = random_subset(dataset, num_samples, seed)\n\n dataloader = flashy.distrib.loader(\n dataset,\n batch_size=batch_size,\n num_workers=num_workers,\n **kwargs\n )\n return dataloader\n\n\ndef get_dataset_from_loader(dataloader):\n dataset = dataloader.dataset\n if isinstance(dataset, torch.utils.data.Subset):\n return dataset.dataset\n else:\n return dataset\n\n\ndef multinomial(input: torch.Tensor, num_samples: int, replacement=False, *, generator=None):\n \"\"\"torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension.\n\n Args:\n input (torch.Tensor): The input tensor containing probabilities.\n num_samples (int): Number of samples to draw.\n replacement (bool): Whether to draw with replacement or not.\n Keywords args:\n generator (torch.Generator): A pseudorandom number generator for sampling.\n Returns:\n torch.Tensor: Last dimension contains num_samples indices\n sampled from the multinomial probability distribution\n located in the last dimension of tensor input.\n \"\"\"\n input_ = input.reshape(-1, input.shape[-1])\n output_ = torch.multinomial(input_, num_samples=num_samples, replacement=replacement, generator=generator)\n output = output_.reshape(*list(input.shape[:-1]), -1)\n return output\n\n\ndef sample_top_k(probs: torch.Tensor, k: int) -> torch.Tensor:\n \"\"\"Sample next token from top K values along the last dimension of the input probs tensor.\n\n Args:\n probs (torch.Tensor): Input probabilities with token candidates on the last dimension.\n k (int): The k in “top-k”.\n Returns:\n torch.Tensor: Sampled tokens.\n \"\"\"\n top_k_value, _ = torch.topk(probs, k, dim=-1)\n min_value_top_k = top_k_value[..., [-1]]\n probs *= (probs >= min_value_top_k).float()\n probs.div_(probs.sum(dim=-1, keepdim=True))\n next_token = multinomial(probs, num_samples=1)\n return next_token\n\n\ndef sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:\n \"\"\"Sample next token from top P probabilities along the last dimension of the input probs tensor.\n\n Args:\n probs (torch.Tensor): Input probabilities with token candidates on the last dimension.\n p (int): The p in “top-p”.\n Returns:\n torch.Tensor: Sampled tokens.\n \"\"\"\n probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)\n probs_sum = torch.cumsum(probs_sort, dim=-1)\n mask = probs_sum - probs_sort > p\n probs_sort *= (~mask).float()\n probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))\n next_token = multinomial(probs_sort, num_samples=1)\n next_token = torch.gather(probs_idx, -1, next_token)\n return next_token\n\n\nclass DummyPoolExecutor:\n \"\"\"Dummy pool executor to use when we actually have only 1 worker.\n (e.g. instead of ProcessPoolExecutor).\n \"\"\"\n class DummyResult:\n def __init__(self, func, *args, **kwargs):\n self.func = func\n self.args = args\n self.kwargs = kwargs\n\n def result(self):\n return self.func(*self.args, **self.kwargs)\n\n def __init__(self, workers, mp_context=None):\n pass\n\n def submit(self, func, *args, **kwargs):\n return DummyPoolExecutor.DummyResult(func, *args, **kwargs)\n\n def __enter__(self):\n return self\n\n def __exit__(self, exc_type, exc_value, exc_tb):\n return\n\n\ndef get_pool_executor(num_workers: int, mp_context=None):\n return ProcessPoolExecutor(num_workers, mp_context) if num_workers > 1 else DummyPoolExecutor(1)\n\n\ndef length_to_mask(lengths: torch.Tensor, max_len: tp.Optional[int] = None) -> torch.Tensor:\n \"\"\"Utility function to convert a tensor of sequence lengths to a mask (useful when working on padded sequences).\n For example: [3, 5] => [[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]]\n\n Args:\n lengths (torch.Tensor): tensor with lengths\n max_len (int): can set the max length manually. Defaults to None.\n Returns:\n torch.Tensor: mask with 0s where there is pad tokens else 1s\n \"\"\"\n assert len(lengths.shape) == 1, \"Length shape should be 1 dimensional.\"\n final_length = lengths.max().item() if not max_len else max_len\n final_length = max(final_length, 1) # if all seqs are of len zero we don't want a zero-size tensor\n return torch.arange(final_length, device=lengths.device)[None, :] < lengths[:, None]\n\n\ndef hash_trick(word: str, vocab_size: int) -> int:\n \"\"\"Hash trick to pair each word with an index\n\n Args:\n word (str): word we wish to convert to an index\n vocab_size (int): size of the vocabulary\n Returns:\n int: index of the word in the embedding LUT\n \"\"\"\n hash = int(hashlib.sha256(word.encode(\"utf-8\")).hexdigest(), 16)\n return hash % vocab_size\n\n\ndef with_rank_rng(base_seed: int = 1234):\n \"\"\"Decorator for a function so that the function will use a Random Number Generator\n whose state depend on the GPU rank. The original RNG state is restored upon returning.\n\n Args:\n base_seed (int): Random seed.\n \"\"\"\n def _decorator(fun: tp.Callable):\n @wraps(fun)\n def _decorated(*args, **kwargs):\n state = torch.get_rng_state()\n seed = base_seed ^ flashy.distrib.rank()\n torch.manual_seed(seed)\n logger.debug('Rank dependent seed set to %d', seed)\n try:\n return fun(*args, **kwargs)\n finally:\n torch.set_rng_state(state)\n logger.debug('RNG state restored.')\n return _decorated\n return _decorator\n\n\ndef collate(tensors: tp.List[torch.Tensor], dim: int = 0) -> tp.Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Get a list of tensors and collate them to a single tensor. according to the following logic:\n - `dim` specifies the time dimension which will be stacked and padded.\n - The output will contain 1 new dimension (dimension index 0) which will be the size of\n of the original list.\n\n Args:\n tensors (tp.List[torch.Tensor]): List of tensors to collate.\n dim (int): Dimension which will be stacked and padded.\n Returns:\n tp.Tuple[torch.Tensor, torch.Tensor]:\n torch.Tensor: Stacked and padded tensor. The output will contain 1 new dimension\n (dimension index 0) which will be the size of the original list.\n torch.Tensor: Tensor containing length of original tensor sizes (without padding).\n \"\"\"\n tensors = [x.transpose(0, dim) for x in tensors]\n lens = torch.LongTensor([len(x) for x in tensors])\n padded_tensors = pad_sequence(tensors)\n padded_tensors = padded_tensors.transpose(0, 1)\n padded_tensors = padded_tensors.transpose(1, dim + 1)\n return padded_tensors, lens\n\n\n# TODO: Move to flashy?\ndef copy_state(state: tp.Any, device: tp.Union[torch.device, str] = 'cpu',\n dtype: tp.Optional[torch.dtype] = None) -> tp.Any:\n if isinstance(state, torch.Tensor):\n if dtype is None or not state.is_floating_point():\n dtype = state.dtype\n return state.detach().to(device=device, dtype=dtype, copy=True)\n elif isinstance(state, dict):\n return {k: copy_state(v, device, dtype) for k, v in state.items()}\n elif isinstance(state, list):\n return [copy_state(v, device, dtype) for v in state]\n\n\n# TODO: Move to flashy?\n@contextmanager\ndef swap_state(model, state, **kwargs):\n old_state = copy_state(model.state_dict())\n model.load_state_dict(state, **kwargs)\n try:\n yield\n finally:\n model.load_state_dict(old_state)\n\n\n@lru_cache(None)\ndef warn_once(logger, msg):\n \"\"\"Warn about a given message only once.\"\"\"\n logger.warning(msg)\n\n\ndef is_jsonable(x: tp.Any):\n \"\"\"Check if an object can be serialized into a json:\"\"\"\n try:\n json.dumps(x)\n return True\n except (TypeError, OverflowError):\n return False\n\n\ndef load_clap_state_dict(clap_model, path: tp.Union[str, Path]):\n \"\"\"Wrapper around state dict loading of CLAP model\n addressing compatibility issues between CLAP and AudioCraft\n HuggingFace transformer version.\n See: https://github.com/LAION-AI/CLAP/issues/118\n \"\"\"\n from clap_module.factory import load_state_dict # type: ignore\n pkg = load_state_dict(path)\n pkg.pop('text_branch.embeddings.position_ids', None)\n clap_model.model.load_state_dict(pkg)\n\n\ndef construct_frame_chords(\n min_timestamp: int,\n chord_changes: tp.List[tp.Tuple[float, str]],\n mapping_dict: tp.Dict,\n prev_chord: str,\n frame_rate: float,\n segment_duration: float,\n ) -> tp.List[str]:\n \"\"\" Translate symbolic chords [(start_time, tuples),...] into a frame-level int sequence\"\"\"\n\n frames = [\n frame / frame_rate\n for frame in range(\n min_timestamp, int(min_timestamp + segment_duration * frame_rate)\n )\n ]\n\n frame_chords = []\n current_chord = prev_chord\n\n for frame in frames:\n while chord_changes and frame >= chord_changes[0][0]:\n current_chord = chord_changes.pop(0)[1]\n current_chord = 'N' if current_chord in {None, ''} else current_chord\n frame_chords.append(mapping_dict[current_chord])\n\n return frame_chords\n" }, { "path": "config/augmentations/default.yaml", "content": "# @package __global__\n\naudio_effects:\n speed:\n sample_rate: ${sample_rate}\n speed_range: [0.8, 1.2]\n updownresample:\n sample_rate: ${sample_rate}\n intermediate_freq: 32000\n echo:\n sample_rate: ${sample_rate}\n volume_range: [0.1, 0.5]\n duration_range: [0.1, 0.5]\n random_noise:\n noise_std: 0.001\n pink_noise:\n noise_std: 0.01\n lowpass_filter:\n sample_rate: ${sample_rate}\n cutoff_freq: 5000\n highpass_filter:\n cutoff_freq: 500\n sample_rate: ${sample_rate}\n bandpass_filter:\n cutoff_freq_low: 300\n cutoff_freq_high: 8000\n sample_rate: ${sample_rate}\n smooth:\n window_size_range: [2, 10]\n boost_audio:\n amount: 20\n duck_audio:\n amount: 20\n mp3_compression:\n sample_rate: ${sample_rate}\n bitrate: 128k # should be a string e.g. \"8k\", \"32k\".. cf ffmpeg to see available bitrates\n aac_compression:\n sample_rate: ${sample_rate}\n bitrate: 128k # should be a string e.g. \"8k\", \"32k\".. cf ffmpeg to see available bitrates\n lowpass_freq: null # don't apply low pass freq to ffmpeg aac compression\n encodec:\n ckpt: \"//pretrained/facebook/encodec_24khz\"\n n_qs: [4, 8, 16]\n\nselect_aug_mode:\n \"use_eval\" # other are 'all' and 'use_eval_acc', used to sample augmentations, `fixed` uses the prob from aug_weights, `all` uses all agmentations every step\n # `use_eval_acc` changes the weights based on the accuracies at evaluation time\n\naug_weights:\n speed: 0.1\n updownresample: 0.1\n echo: 0.1\n pink_noise: 0.1\n lowpass_filter: 0.1\n highpass_filter: 0.1\n bandpass_filter: 0.1\n smooth: 0.1\n boost_audio: 0.1\n duck_audio: 0.1\n mp3_compression: 0.1 # eval only never use in training even if eval_acc low\n aac_compression: 0.1 # eval only never use in training even if eval_acc low\n encodec: 0.1\n identity: 1 # no augmentation\n\nn_max_aug: null" }, { "path": "config/conditioner/chords2music.yaml", "content": "# @package __global__\n\nclassifier_free_guidance:\n training_dropout: 0.3 # dropout of all conditions\n inference_coef: 3.0\n\nattribute_dropout:\n symbolic:\n chords: 0.3 # independent dropout of chords\n\nfuser:\n cross_attention_pos_emb: false\n cross_attention_pos_emb_scale: 1\n sum: []\n prepend: []\n cross: [description]\n ignore: [chords]\n input_interpolate: []\n\nconditioners:\n description:\n model: t5\n t5:\n name: t5-base\n finetune: false\n word_dropout: 0.3\n normalize_text: false\n chords:\n model: chords_emb\n chords_emb:\n card: 194 # Chordino\n out_dim: 16\n\ndataset:\n train:\n merge_text_p: 0.25\n drop_desc_p: 0.5\n drop_other_p: 0.5\n" }, { "path": "config/conditioner/chroma2music.yaml", "content": "# @package __global__\n\nclassifier_free_guidance:\n training_dropout: 0.2\n inference_coef: 3.0\n\nattribute_dropout:\n args:\n active_on_eval: false\n text: {}\n wav:\n self_wav: 0.5\n\nfuser:\n cross_attention_pos_emb: false\n cross_attention_pos_emb_scale: 1\n sum: []\n prepend: [self_wav, description]\n cross: []\n input_interpolate: []\n\nconditioners:\n self_wav:\n model: chroma_stem\n chroma_stem:\n sample_rate: ${sample_rate}\n n_chroma: 12\n radix2_exp: 14\n argmax: true\n match_len_on_eval: false\n eval_wavs: null\n n_eval_wavs: 100\n cache_path: null\n description:\n model: t5\n t5:\n name: t5-base\n finetune: false\n word_dropout: 0.2\n normalize_text: false\n\ndataset:\n train:\n merge_text_p: 0.25\n drop_desc_p: 0.5\n drop_other_p: 0.5\n" }, { "path": "config/conditioner/clapemb2music.yaml", "content": "# @package __global__\n\nclassifier_free_guidance:\n training_dropout: 0.3\n inference_coef: 3.0\n\nattribute_dropout:\n text: {}\n wav: {}\n\nfuser:\n cross_attention_pos_emb: false\n cross_attention_pos_emb_scale: 1\n sum: []\n prepend: []\n cross: [description]\n input_interpolate: []\n\nconditioners:\n description:\n model: clap\n clap:\n checkpoint: //reference/clap/music_audioset_epoch_15_esc_90.14.pt\n model_arch: 'HTSAT-base'\n enable_fusion: false\n sample_rate: 48000\n max_audio_length: 10\n audio_stride: 1\n dim: 512\n attribute: description\n normalize: true\n quantize: true # use RVQ quantization\n n_q: 12\n bins: 1024\n kmeans_iters: 50\n text_p: 0. # probability of using text embed at train time\n cache_path: null\n\ndataset:\n joint_embed_attributes: [description]\n train:\n merge_text_p: 0.25\n drop_desc_p: 0.5\n drop_other_p: 0.5\n" }, { "path": "config/conditioner/drums2music.yaml", "content": "# @package __global__\n\nclassifier_free_guidance:\n training_dropout: 0.3 # dropout of all conditions\n inference_coef: 3.0\n\nattribute_dropout:\n text: {}\n wav:\n self_wav: 0.3 # independent dropout of drums\n\nfuser:\n cross_attention_pos_emb: false\n cross_attention_pos_emb_scale: 1\n sum: []\n prepend: []\n cross: [description]\n ignore: [self_wav]\n input_interpolate: []\n\nconditioners:\n self_wav:\n model: drum_latents\n drum_latents:\n sample_rate: ${sample_rate}\n out_dim: 2\n blurring_factor: 3\n cache_path: null\n\n description:\n model: t5\n t5:\n name: t5-base\n finetune: false\n word_dropout: 0.3\n normalize_text: false\n\ndataset:\n train:\n merge_text_p: 0.25\n drop_desc_p: 0.5\n drop_other_p: 0.5\n" }, { "path": "config/conditioner/jasco_chords_drums.yaml", "content": "# @package __global__\n\nclassifier_free_guidance:\n training_dropout: 0.3 # dropout of all conditions\n inference_coef: 3.0\n\nattribute_dropout:\n text: {}\n symbolic:\n chords: 0.3 # independent dropout of chords\n wav:\n self_wav: 0.3 # independent dropout of drums\n\nfuser:\n cross_attention_pos_emb: false\n cross_attention_pos_emb_scale: 1\n sum: []\n prepend: []\n cross: [description]\n ignore: [chords, self_wav]\n input_interpolate: []\n\nconditioners:\n self_wav:\n model: drum_latents\n drum_latents:\n sample_rate: ${sample_rate}\n out_dim: 2\n blurring_factor: 3\n cache_path: null\n\n description:\n model: t5\n t5:\n name: t5-base\n finetune: false\n word_dropout: 0.3\n normalize_text: false\n chords:\n model: chords_emb\n chords_emb:\n card: 194 # Chordino\n out_dim: 16\n\ndataset:\n train:\n merge_text_p: 0.25\n drop_desc_p: 0.5\n drop_other_p: 0.5\n\n" }, { "path": "config/conditioner/jasco_chords_drums_melody.yaml", "content": "# @package __global__\n\nclassifier_free_guidance:\n training_dropout: 0.2 # dropout of all conditions\n inference_coef: 3.0\n\nattribute_dropout:\n text:\n description: 0.0\n symbolic:\n chords: 0.5 # independent dropout of chords\n melody: 0.5 # independent dropout of melody\n wav:\n self_wav: 0.5 # independent dropout of drums\n\n\nfuser:\n cross_attention_pos_emb: false\n cross_attention_pos_emb_scale: 1\n sum: []\n prepend: []\n cross: [description]\n ignore: [chords, self_wav, melody]\n input_interpolate: []\n\nconditioners:\n self_wav:\n model: drum_latents\n drum_latents:\n sample_rate: ${sample_rate}\n out_dim: 2\n blurring_factor: 3\n cache_path: ???\n read_only_cache: true\n\n description:\n model: t5\n t5:\n name: t5-base\n finetune: false\n word_dropout: 0.3\n normalize_text: false\n\n chords:\n model: chords_emb\n chords_emb:\n card: 194 # Chordino\n out_dim: 16\n\n melody:\n model: melody\n melody:\n card: 53 # Preprocessed salience dim\n out_dim: 16\n\ndataset:\n train:\n merge_text_p: 0.25\n drop_desc_p: 0.5\n drop_other_p: 0.5\n" }, { "path": "config/conditioner/none.yaml", "content": "# @package __global__\n\n# No conditioning\n\nclassifier_free_guidance:\n training_dropout: 0\n inference_coef: 1\n\nattribute_dropout:\n text: {}\n wav: {}\n\nfuser:\n sum: []\n prepend: []\n cross: []\n input_interpolate: []\n\nconditioners: null\n" }, { "path": "config/conditioner/style2music.yaml", "content": "# @package __global__\n\nclassifier_free_guidance:\n training_dropout: 0.1\n inference_coef: 3.0\n\nattribute_dropout:\n args:\n active_on_eval: false\n text: \n description: 0.4\n wav:\n self_wav: 0.4\n\nfuser:\n cross_attention_pos_emb: false\n cross_attention_pos_emb_scale: 1\n sum: []\n prepend: [self_wav, description]\n cross: []\n input_interpolate: []\n\nconditioners:\n self_wav:\n model: style\n style:\n model_name: mert\n transformer_scale: default\n sample_rate: ${sample_rate}\n encodec_checkpoint: '//pretrained/facebook/encodec_32khz'\n encodec_n_q: 3\n length: 3.0\n ds_factor: 15 # Since MERT is 75Hz, 75/15 results into 5Hz representations\n n_q_out: 6\n eval_q: 3\n q_dropout: true\n bins: 1024\n varying_lengths: [1.5, 4.5]\n batch_norm: true\n compute_mask: true\n num_codebooks_lm: ${transformer_lm.n_q}\n ds_rate_compression: 640\n use_middle_of_segment: false\n rvq_threshold_ema_dead_code: 0.1\n\n description:\n model: t5\n t5:\n name: t5-base\n finetune: false\n word_dropout: 0.2\n normalize_text: false\n\ndataset:\n train:\n merge_text_p: 0.25\n drop_desc_p: 0.5\n drop_other_p: 0.5\n shuffle: true\n" }, { "path": "config/conditioner/text2music.yaml", "content": "# @package __global__\n\nclassifier_free_guidance:\n training_dropout: 0.3\n inference_coef: 3.0\n\nattribute_dropout: {}\n\nfuser:\n cross_attention_pos_emb: false\n cross_attention_pos_emb_scale: 1\n sum: []\n prepend: []\n cross: [description]\n input_interpolate: []\n\nconditioners:\n description:\n model: t5\n t5:\n name: t5-base\n finetune: false\n word_dropout: 0.3\n normalize_text: false\n\ndataset:\n train:\n merge_text_p: 0.25\n drop_desc_p: 0.5\n drop_other_p: 0.5\n" }, { "path": "config/conditioner/text2sound.yaml", "content": "# @package __global__\n\nclassifier_free_guidance:\n training_dropout: 0.1\n inference_coef: 3.0\n\nattribute_dropout: {}\n\nfuser:\n cross_attention_pos_emb: false\n cross_attention_pos_emb_scale: 1\n sum: []\n prepend: []\n cross: [description]\n input_interpolate: []\n\nconditioners:\n description:\n model: t5\n t5:\n name: t5-large\n finetune: false\n word_dropout: 0.\n normalize_text: false\n" }, { "path": "config/config.yaml", "content": "# WARNING: This is the base configuration file shared across ALL solvers in AudioCraft\n# Please don't update this file directly. Instead use distinct configuration files\n# to override the below configuration.\ndefaults:\n - _self_\n - dset: default\n - solver: default\n\ndevice: cuda\ndtype: float32\nautocast: false\nautocast_dtype: bfloat16\nseed: 2036\nshow: false # just show the model and its size and exit\ncontinue_from: # continue from a given sig or path\nexecute_only: # can be set to generate/evaluate/valid to run that stage\nexecute_inplace: false # don't enforce continue_from to be set\n # to enable inplace execution of the stage. This assume\n # that you know what you are doing and execute stage\n # preserving the original xp sig.\nbenchmark_no_load: false # if set to true, will repeat the same batch instead of loading them\n\nefficient_attention_backend: torch # can be torch or xformers.\nnum_threads: 1 # called with torch.set_num_thread.\nmp_start_method: forkserver # multiprocessing method (spawn, fork or fork_server).\n\n\nlabel: # use this if you want twice the same exp, with a name.\n\n# logging parameters\nlogging:\n level: INFO\n log_updates: 10\n log_tensorboard: false\n log_wandb: false\ntensorboard:\n with_media_logging: false\n name: # optional name for the experiment\n sub_dir: # optional sub directory to store tensorboard data\nwandb:\n with_media_logging: true\n project: # project name\n name: # optional name for the experiment\n group: # optional group\n\n# SLURM launcher configuration.\nslurm:\n gpus: 4 # convenience parameter, number of GPUs to use.\n mem_per_gpu: 40 # in GB, total mem is automatically scaled with `gpus`.\n time: 3600\n constraint:\n partition:\n comment:\n setup: []\n exclude: ''\n\n# dora parameters\ndora:\n # Output folder for all artifacts of an experiment.\n dir: /checkpoint/${oc.env:USER}/experiments/audiocraft/outputs\n # The following entries will be ignored by dora when computing the unique XP signature.\n # Note that slurm.* and dora.* are automatically ignored.\n exclude: [\n 'device', 'wandb.*', 'tensorboard.*', 'logging.*',\n 'dataset.num_workers', 'eval.num_workers', 'special.*',\n 'metrics.visqol.bin', 'metrics.fad.bin',\n 'execute_only', 'execute_best', 'generate.every',\n 'optim.eager_sync', 'profiler.*', 'deadlock.*',\n 'efficient_attention_backend', 'num_threads', 'mp_start_method',\n ]\n use_rendezvous: false\n # for grids, always run from a clean repo, allowing reliable runs and storing\n # the exact commit. Your repo must be absolutely pristine clean.\n # Local `dora run` are not impacted for easier debugging.\n git_save: true\n" }, { "path": "config/dset/audio/audiocaps_16khz.yaml", "content": "# @package __global__\n\n# AudioCaps dataset\ndatasource:\n max_sample_rate: 16000\n max_channels: 1\n\n train: null # only evaluation set\n valid: null # only evaluation set\n evaluate: egs/audiocaps/audiocaps_16khz\n generate: egs/audiocaps/audiocaps_16khz # identical to evaluate\n" }, { "path": "config/dset/audio/default.yaml", "content": "# @package __global__\n\ndatasource:\n max_sample_rate: ???\n max_channels: ???\n\n train: ???\n valid: ???\n evaluate: ???\n generate: null\n" }, { "path": "config/dset/audio/example.yaml", "content": "# @package __global__\n\ndatasource:\n max_sample_rate: 44100\n max_channels: 2\n\n train: egs/example\n valid: egs/example\n evaluate: egs/example\n generate: egs/example\n" }, { "path": "config/dset/audio/musiccaps_32khz.yaml", "content": "# @package __global__\n\n# total samples obtained from MusicCaps = 5469\n# (out of 5521 due to AudioSet corrupted samples)\ndatasource:\n max_sample_rate: 32000\n max_channels: 2\n\n train: null # only evaluation set\n valid: null # only evaluation set\n evaluate: egs/musiccaps/musiccaps_32khz\n generate: egs/musiccaps/musiccaps_32khz # identical to evaluate\n" }, { "path": "config/dset/default.yaml", "content": "# @package __global__\n\n# WARNING: This is a base configuration file shared across ALL solvers in AudioCraft\n# Please don't update this file directly. Instead use distinct configuration files\n# to override the below configuration.\ndatasource:\n train: ???\n valid: ???\n evaluate: ???\n generate: ???\n" }, { "path": "config/dset/internal/music_10k_32khz.yaml", "content": "# @package __global__\n\n# high quality music dataset with no artist overlap between splits\ndatasource:\n max_sample_rate: 32000\n max_channels: 1\n\n train: egs/music/music_10k_32khz/train\n valid: egs/music/music_10k_32khz/valid\n evaluate: egs/music/music_10k_32khz/test\n generate: egs/music/music_10k_32khz/test # identical to evaluate\n" }, { "path": "config/dset/internal/music_400k_32khz.yaml", "content": "# @package __global__\n\ndatasource:\n max_sample_rate: 32000\n max_channels: 1\n\n train: egs/music/music_400k_32khz/train\n valid: egs/music/music_400k_32khz/valid\n evaluate: egs/music/music_400k_32khz/test\n generate: egs/music/music_400k_32khz/test # identical to evaluate\n" }, { "path": "config/dset/internal/sounds_16khz.yaml", "content": "# @package __global__\n\n# environmental sounds dataset compiling all datasets\n# with applied filters on tags\ndatasource:\n max_sample_rate: 16000\n max_channels: 1\n\n train: egs/sound/sounds_16khz/train\n valid: egs/sound/sounds_16khz/valid\n evaluate: egs/sound/sounds_16khz/test\n generate: egs/sound/sounds_16khz/test # identical to evaluate\n" }, { "path": "config/model/encodec/default.yaml", "content": "# @package __global__\n\ncompression_model: encodec\n\nencodec:\n autoencoder: seanet\n quantizer: rvq\n sample_rate: ${sample_rate}\n channels: ${channels}\n causal: false\n renormalize: false\n\nseanet:\n dimension: 128\n channels: ${channels}\n causal: ${encodec.causal}\n n_filters: 32\n n_residual_layers: 1\n ratios: [8, 5, 4, 2]\n activation: ELU\n activation_params: {\"alpha\": 1.}\n norm: weight_norm\n norm_params: {}\n kernel_size: 7\n residual_kernel_size: 3\n last_kernel_size: 7\n dilation_base: 2\n pad_mode: constant\n true_skip: true\n compress: 2\n lstm: 2\n disable_norm_outer_blocks: 0\n # Specific encoder or decoder params.\n # You can also override any param for the encoder or decoder only\n # by using Hydra `+param=` syntax, i.e.`\n # `+seanet.decoder.n_filters=64`.\n decoder:\n trim_right_ratio: 1.0\n final_activation: null\n final_activation_params: null\n encoder: {}\n\nrvq:\n n_q: 8\n q_dropout: false\n bins: 1024\n decay: 0.99\n kmeans_init: true\n kmeans_iters: 50\n threshold_ema_dead_code: 2\n orthogonal_reg_weight: 0.0\n orthogonal_reg_active_codes_only: false\n\nno_quant: {}\n" }, { "path": "config/model/encodec/encodec_base_causal.yaml", "content": "# @package __global__\n\ndefaults:\n - encodec/default\n\nencodec:\n causal: true\n\nrvq:\n n_q: 32\n q_dropout: true\n" }, { "path": "config/model/encodec/encodec_large_nq4_s320.yaml", "content": "# @package __global__\n\ndefaults:\n - encodec/default\n\nseanet:\n # default ratios are [8, 5, 4, 2]\n n_filters: 64\n\nrvq:\n bins: 2048\n n_q: 4\n q_dropout: false\n" }, { "path": "config/model/encodec/encodec_large_nq4_s640.yaml", "content": "# @package __global__\n\ndefaults:\n - encodec/default\n\nseanet:\n ratios: [8, 5, 4, 4]\n n_filters: 64\n\nrvq:\n bins: 2048\n n_q: 4\n q_dropout: false\n" }, { "path": "config/model/lm/audiogen_lm.yaml", "content": "# @package __global__\n\ndefaults:\n - lm/default\n - override /conditioner: text2sound\n - override /model/lm/model_scale: small # prefer this group to set model scale instead of transformer_lm keys directly\n\nlm_model: transformer_lm\n\ncodebooks_pattern:\n modeling: delay\n delay:\n delays: [0, 1, 2, 3]\n flatten_first: 0\n empty_initial: 0\n unroll:\n flattening: [0, 1, 2, 3]\n delays: [0, 0, 0, 0]\n music_lm:\n group_by: 2\n coarse_first:\n delays: [0, 0, 0]\n\ntransformer_lm:\n n_q: 4\n card: 2048\n memory_efficient: true\n bias_proj: false\n bias_ff: false\n bias_attn: false\n norm_first: true\n layer_scale: null\n weight_init: gaussian\n depthwise_init: current\n zero_bias_init: true\n attention_as_float32: false\n" }, { "path": "config/model/lm/default.yaml", "content": "# @package __global__\ndefaults:\n - _self_\n - /model/lm/model_scale: base # prefer this group to set model scale instead of transformer_lm keys directly\n\nlm_model: transformer_lm\n\ncodebooks_pattern:\n modeling: parallel\n\ntransformer_lm:\n dim: 512\n num_heads: 8\n num_layers: 8\n hidden_scale: 4\n n_q: 8 # number of streams to model\n card: 1024\n dropout: 0.\n emb_lr: null\n activation: gelu\n norm_first: false # use pre-norm instead of post-norm\n bias_ff: true # use bias for the feedforward\n bias_attn: true # use bias for the attention\n bias_proj: true # use bias for the output projections\n past_context: null\n causal: true\n custom: false # use custom MHA implementation\n memory_efficient: false # use flash attention\n attention_as_float32: false # use float32 for the attention part,\n # recommended at the moment when memory_efficient is True.\n layer_scale: null\n positional_embedding: sin # positional embedding strategy (sin, rope, or sin_rope).\n xpos: false # apply xpos decay (rope only).\n checkpointing: none # layer checkpointing method, can be none, torch, xformers_default.\n # torch is the slowest but uses the least memory,\n # xformers_default is somewhere in between.\n weight_init: null # weight initialization (null, gaussian or uniform)\n depthwise_init: null # perform depthwise initialization (null, current, global)\n zero_bias_init: false # initialize bias to zero if bias in linears and\n # if a weight_init method is used.\n norm: layer_norm # normalization method to use in transformer.\n cross_attention: false\n qk_layer_norm: false\n qk_layer_norm_cross: false\n attention_dropout: null\n kv_repeat: 1\n two_step_cfg: false # whether to do true 2 steps CFG, potentially resolving some padding issues or not...\n" }, { "path": "config/model/lm/model_scale/base.yaml", "content": "# @package __global__\n\n# overrides nothing because default is already transformer base (~ 60M params)\n" }, { "path": "config/model/lm/model_scale/large.yaml", "content": "# @package _global_\n\n# gpt2 inspired, even bigger (~3.3B params)\ntransformer_lm:\n dim: 2048\n num_heads: 32\n num_layers: 48\n" }, { "path": "config/model/lm/model_scale/medium.yaml", "content": "# @package _global_\n\n# gpt2 like (~1.5B params)\ntransformer_lm:\n dim: 1536\n num_heads: 24\n num_layers: 48\n" }, { "path": "config/model/lm/model_scale/small.yaml", "content": "# @package _global_\n\n# 300M Param.\n\ntransformer_lm:\n dim: 1024\n num_heads: 16\n num_layers: 24\n" }, { "path": "config/model/lm/model_scale/xsmall.yaml", "content": "# @package _global_\n# just used for debugging or when we just want to populate the cache\n# and do not care about training.\n\ntransformer_lm:\n dim: 64\n num_heads: 2\n num_layers: 2\n" }, { "path": "config/model/lm/musicgen_lm.yaml", "content": "# @package __global__\n\ndefaults:\n - lm/default\n - override /conditioner: text2music\n - override /model/lm/model_scale: small # prefer this group to set model scale instead of transformer_lm keys directly\n\nlm_model: transformer_lm\n\ncodebooks_pattern:\n modeling: delay\n delay:\n delays: [0, 1, 2, 3]\n flatten_first: 0\n empty_initial: 0\n unroll:\n flattening: [0, 1, 2, 3]\n delays: [0, 0, 0, 0]\n music_lm:\n group_by: 2\n coarse_first:\n delays: [0, 0, 0]\n\ntransformer_lm:\n n_q: 4\n card: 2048\n memory_efficient: true\n bias_proj: false\n bias_ff: false\n bias_attn: false\n norm_first: true\n layer_scale: null\n weight_init: gaussian\n depthwise_init: current\n zero_bias_init: true\n attention_as_float32: false\n" }, { "path": "config/model/none.yaml", "content": "# @package __global__\n\n# This file exist so that model is recognized as a config group\n# by Hydra, and Dora. A bit weird we might need a better fix someday.\n" }, { "path": "config/model/score/basic.yaml", "content": "# @package _global_\n\ndiffusion_unet:\n hidden: 48\n depth: 4\n res_blocks: 1\n norm_groups: 4\n kernel: 8\n stride: 4\n growth: 4\n max_channels: 10_000\n dropout: 0.\n emb_all_layers: true\n bilstm: false\n codec_dim: null\n transformer: false\n cross_attention: false" }, { "path": "config/model/watermark/default.yaml", "content": "# @package __global__\n\naudioseal:\n autoencoder: seanet\n sample_rate: 16000\n channels: 1\n nbits: 16\n\nseanet:\n dimension: 128\n channels: 1\n causal: false\n n_filters: 32\n n_residual_layers: 1\n ratios: [8, 5, 4, 2]\n activation: ELU\n activation_params: { \"alpha\": 1. }\n norm: weight_norm\n norm_params: {}\n kernel_size: 7\n residual_kernel_size: 3\n last_kernel_size: 7\n dilation_base: 2\n pad_mode: constant\n true_skip: true\n compress: 2\n lstm: 2\n disable_norm_outer_blocks: 0\n # Specific encoder or decoder params.\n # You can also override any param for the encoder or decoder only\n # by using Hydra `+param=` syntax, i.e.`\n # `+seanet.decoder.n_filters=64`.\n decoder:\n trim_right_ratio: 1.0\n final_activation: null\n final_activation_params: null\n encoder: {}\n\ndetector: {\n \"output_dim\": 32, # output channels of detector upsampling\n} " }, { "path": "config/solver/audiogen/audiogen_base_16khz.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the base AudioGen model (text-to-sound)\n# on monophonic audio sampled at 16 kHz\n# using a similar EnCodec+LM setup to MusicGen\ndefaults:\n - audiogen/default\n - /model: lm/audiogen_lm\n - override /dset: audio/default\n - _self_\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 16khz\n# with a total stride of 320 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //reference/bd44a852/checkpoint.th\n\nchannels: 1\nsample_rate: 16000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n batch_size: 128 # matching AudioGen paper setup (256 * mix_p=0.5 = 128)\n num_workers: 10\n segment_duration: 10\n min_segment_ratio: 1.0\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n external_metadata_source: null\n # sample mixing augmentation at train time\n train:\n batch_size: 256 # matching AudioGen paper setup\n aug_p: 0.5 # perform audio mixing 50% of the time\n mix_p: 0.5 # proportion of batch items mixed together\n # important: note that this will reduce the\n # actual batch size used at train time\n # which will be equal to mix_p * batch_size\n mix_snr_low: -5\n mix_snr_high: 5\n mix_min_overlap: 0.5\n\ngenerate:\n lm:\n use_sampling: true\n top_k: 250\n top_p: 0.0\n\noptim:\n epochs: 100\n optimizer: adamw\n lr: 5e-4\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: inverse_sqrt\n inverse_sqrt:\n warmup: 3000\n warmup_init_lr: 0.0\n" }, { "path": "config/solver/audiogen/debug.yaml", "content": "# @package __global__\n\n# This is a minimal debugging configuration\n# for MusicGen training solver\ndefaults:\n - audiogen/default\n - /model: lm/audiogen_lm\n - override /model/lm/model_scale: xsmall\n - override /dset: audio/example\n - _self_\n\nautocast: false\ncompression_model_checkpoint: null\ntransformer_lm:\n n_q: 4\n card: 400\n\nconditioners:\n description:\n model: t5\n t5:\n name: t5-small\n\ncodebooks_pattern:\n modeling: parallel\n\nchannels: 1\nsample_rate: 16000\n\ndeadlock:\n use: false # deadlock detection\n\ndataset:\n batch_size: 4\n segment_duration: 5\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n\ngenerate:\n audio:\n strategy: peak\n lm:\n use_sampling: false\n top_k: 0\n top_p: 0.0\n\ncheckpoint:\n save_every: 0\n keep_last: 0\n\noptim:\n epochs: 2\n updates_per_epoch: 10\n optimizer: adamw\n lr: 1e-4\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: null\n" }, { "path": "config/solver/audiogen/default.yaml", "content": "# @package __global__\n\ndefaults:\n - /solver/musicgen/default\n - _self_\n - /solver/audiogen/evaluation: none\n - override /dset: audio/default\n\n# See config/solver/musicgen/default.yaml for a list of possible values.\n# We only keep the most important here.\n\nautocast: true\nautocast_dtype: float16\n\nsolver: audiogen\nsample_rate: ???\nchannels: ???\ncompression_model_checkpoint: ???\n\ntokens:\n padding_with_special_token: false\n\ndataset:\n batch_size: 128\n segment_duration: 10\n min_segment_ratio: 1.0 # lower values such as 0.5 result in generations with a lot of silence.\n\noptim:\n epochs: 100\n updates_per_epoch: 2000\n lr: 1e-4\n optimizer: adamw\n max_norm: 1.0\n adam:\n betas: [0.9, 0.95]\n weight_decay: 0.1\n eps: 1e-8\n\nschedule:\n lr_scheduler: null\n" }, { "path": "config/solver/audiogen/evaluation/none.yaml", "content": "# @package __global__\n\ndataset:\n evaluate:\n num_samples: 10000\n" }, { "path": "config/solver/audiogen/evaluation/objective_eval.yaml", "content": "# @package __global__\n\n# Setup for execute only on audiocaps for audio generation\n# evaluation with objective metrics\n# execute_only=evaluate\n\ndataset:\n max_audio_duration: null\n # ensure the proper values are broadcasted here for evaluate\n evaluate:\n min_audio_duration: 1. # some metrics requires a minimum audio length\n max_audio_duration: null # all samples from audiocaps should be ~10s\n num_samples: null\n segment_duration: null\n generate:\n min_audio_duration: 1.\n max_audio_duration: null\n num_samples: 500\n\nevaluate:\n metrics:\n fad: true\n kld: true\n text_consistency: true\n\nmetrics:\n kld:\n passt:\n pretrained_length: 10 # similarly to reported results in AudioGen paper\n" }, { "path": "config/solver/compression/debug.yaml", "content": "# @package __global__\n\ndefaults:\n - compression/default\n - /model: encodec/encodec_base_causal\n - override /dset: audio/example\n - _self_\n\nchannels: 1\nsample_rate: 16000\n\n# debug config uses just L1\nlosses:\n adv: 0.\n feat: 0.\n l1: 1.\n mel: 0.\n msspec: 0.\n# no balancer\nbalancer:\n balance_grads: false\n ema_decay: 1.\n total_norm: 1.\n per_batch_item: false\n# no adversaries\nadversarial:\n adversaries: []\n adv_loss: hinge\n feat_loss: l1\n\n# faster model for local dev\nseanet:\n dimension: 16\n n_filters: 4\n\n# very small dataset\ndataset:\n batch_size: 8\n num_workers: 10\n num_samples: 100\n segment_duration: 1\n evaluate:\n batch_size: 32\n generate:\n batch_size: 1\n num_samples: 5\n segment_duration: 10\n\n# limited training\nevaluate:\n every: 5\ngenerate:\n every: 5\noptim:\n epochs: 50\n" }, { "path": "config/solver/compression/default.yaml", "content": "# @package __global__\n\ndefaults:\n - ../default\n - override /dset: audio/default\n - _self_\n\nsolver: compression\nsample_rate: ???\nchannels: ???\n\n# loss balancing\nlosses:\n adv: 4.\n feat: 4.\n l1: 0.1\n mel: 0.\n msspec: 2.\n sisnr: 0.\nbalancer:\n balance_grads: true\n ema_decay: 0.999\n per_batch_item: true\n total_norm: 1.\n\nadversarial:\n every: 1\n adversaries: [msstftd]\n adv_loss: hinge\n feat_loss: l1\n\n# losses hyperparameters\nl1: {}\nl2: {}\nmrstft:\n factor_sc: .5\n factor_mag: .5\n normalized: false\nmel:\n sample_rate: ${sample_rate}\n n_fft: 1024\n hop_length: 256\n win_length: 1024\n n_mels: 64\n f_min: 64\n f_max: null\n normalized: false\n floor_level: 1e-5\nsisnr:\n sample_rate: ${sample_rate}\n segment: 5.\nmsspec:\n sample_rate: ${sample_rate}\n range_start: 6\n range_end: 11\n n_mels: 64\n f_min: 64\n f_max: null\n normalized: true\n alphas: false\n floor_level: 1e-5\n\n# metrics\nmetrics:\n visqol:\n mode: audio\n bin: null # path to visqol install\n model: tcdaudio14_aacvopus_coresv_svrnsim_n.68_g.01_c1.model # visqol v3\n\n# adversaries hyperparameters\nmsstftd:\n in_channels: 1\n out_channels: 1\n filters: 32\n norm: weight_norm\n n_ffts: [1024, 2048, 512, 256, 128]\n hop_lengths: [256, 512, 128, 64, 32]\n win_lengths: [1024, 2048, 512, 256, 128]\n activation: LeakyReLU\n activation_params: {negative_slope: 0.3}\nmsd:\n in_channels: 1\n out_channels: 1\n scale_norms: [spectral_norm, weight_norm, weight_norm]\n kernel_sizes: [5, 3]\n filters: 16\n max_filters: 1024\n downsample_scales: [4, 4, 4, 4]\n inner_kernel_sizes: null\n groups: [4, 4, 4, 4]\n strides: null\n paddings: null\n activation: LeakyReLU\n activation_params: {negative_slope: 0.3}\nmpd:\n in_channels: 1\n out_channels: 1\n periods: [2, 3, 5, 7, 11]\n n_layers: 5\n kernel_size: 5\n stride: 3\n filters: 8\n filter_scales: 4\n max_filters: 1024\n activation: LeakyReLU\n activation_params: {negative_slope: 0.3}\n norm: weight_norm\n\n# data hyperparameters\ndataset:\n batch_size: 64\n num_workers: 10\n segment_duration: 1\n train:\n num_samples: 500000\n valid:\n num_samples: 10000\n evaluate:\n batch_size: 32\n num_samples: 10000\n generate:\n batch_size: 32\n num_samples: 50\n segment_duration: 10\n\n# solver hyperparameters\nevaluate:\n every: 25\n num_workers: 5\n metrics:\n visqol: false\n sisnr: true\ngenerate:\n every: 25\n num_workers: 5\n audio:\n sample_rate: ${sample_rate}\n\n# checkpointing schedule\ncheckpoint:\n save_last: true\n save_every: 25\n keep_last: 10\n keep_every_states: null\n\n# optimization hyperparameters\noptim:\n epochs: 200\n updates_per_epoch: 2000\n lr: 3e-4\n max_norm: 0.\n optimizer: adam\n adam:\n betas: [0.5, 0.9]\n weight_decay: 0.\n ema:\n use: true # whether to use EMA or not\n updates: 1 # update at every step\n device: ${device} # device for EMA, can be put on GPU if more frequent updates\n decay: 0.99 # EMA decay value, if null, no EMA is used\n" }, { "path": "config/solver/compression/encodec_audiogen_16khz.yaml", "content": "# @package __global__\n\ndefaults:\n - compression/default\n - /model: encodec/encodec_large_nq4_s320\n - override /dset: audio/default\n - _self_\n\nchannels: 1\nsample_rate: 16000\n" }, { "path": "config/solver/compression/encodec_base_24khz.yaml", "content": "# @package __global__\n\ndefaults:\n - compression/default\n - /model: encodec/encodec_base_causal\n - override /dset: audio/default\n - _self_\n\nchannels: 1\nsample_rate: 24000\n" }, { "path": "config/solver/compression/encodec_musicgen_32khz.yaml", "content": "# @package __global__\n\ndefaults:\n - compression/default\n - /model: encodec/encodec_large_nq4_s640\n - override /dset: audio/default\n - _self_\n\nchannels: 1\nsample_rate: 32000\n" }, { "path": "config/solver/default.yaml", "content": "# @package __global__\n\n# WARNING: This is a base configuration file shared across ALL solvers in AudioCraft\n# Please don't update this file directly. Instead use distinct configuration files\n# to override the below configuration.\nsolver: ???\n\nfsdp:\n use: false # should we use FSDP.\n param_dtype: float16 # equivalent to autocast_dtype for FSDP.\n reduce_dtype: float32 # gradient averaging dtype, float32 will give max stability.\n buffer_dtype: float32 # dtype used for buffers, we don't have much buffers, so let's leave it.\n sharding_strategy: shard_grad_op # can be shard_grad_op or full_shard.\n # full_shard will use less memory but slower ??\n per_block: true # If True, uses nested FSDP.\n\nprofiler:\n enabled: false\n\ndeadlock:\n use: false\n timeout: 600\n\ndataset:\n batch_size: ???\n num_workers: 10\n segment_duration: null\n num_samples: null\n return_info: false\n shuffle: false\n sample_on_duration: true\n sample_on_weight: true\n min_segment_ratio: 0.5\n train:\n num_samples: null\n shuffle: true\n shuffle_seed: 0 # if you want to sample the data differently.\n permutation_on_files: false\n valid:\n num_samples: null\n evaluate:\n num_samples: null\n generate:\n num_samples: null\n return_info: true\n\ncheckpoint:\n save_last: true\n save_every: null\n keep_last: null\n keep_every_states: null\n\ngenerate:\n every: null\n path: 'samples'\n audio:\n format: 'mp3'\n strategy: 'clip'\n sample_rate: null\n lm:\n use_sampling: false\n temp: 1.0\n top_k: 0\n top_p: 0.0\nevaluate:\n every: null\n num_workers: 5\n truncate_audio: null\n fixed_generation_duration: null # in secs\n metrics:\n base: true # run default evaluation (e.g. like train/valid stage)\n\noptim:\n epochs: ???\n updates_per_epoch: null\n lr: ???\n optimizer: ???\n adam:\n betas: [0.9, 0.999]\n weight_decay: 0.\n ema:\n use: false # whether to use EMA or not\n updates: ${optim.updates_per_epoch} # frequency of updates of the EMA\n device: cpu # device for EMA, can be put on GPU if more frequent updates\n decay: 0.99 # EMA decay value, if null, no EMA is used\n\nschedule:\n lr_scheduler: null\n step:\n step_size: null\n gamma: null\n exponential:\n lr_decay: null\n cosine:\n warmup: null\n lr_min_ratio: 0.0\n cycle_length: 1.0\n polynomial_decay:\n warmup: null\n zero_lr_warmup_steps: 0\n end_lr: 0.0\n power: 1\n inverse_sqrt:\n warmup: null\n warmup_init_lr: 0.0\n linear_warmup:\n warmup: null\n warmup_init_lr: 0.0\n" }, { "path": "config/solver/diffusion/debug.yaml", "content": "# @package __global__\n\ndefaults:\n - /solver/default\n - /model: score/basic\n - override /dset: audio/default\n - _self_\n\nsolver: diffusion\n\nsample_rate: 16000\nchannels: 1\ncompression_model_checkpoint: //sig/5091833e\nn_q: 2 # number of codebooks to keep\n\ndataset:\n batch_size: 8\n num_workers: 10\n segment_duration: 1\n train:\n num_samples: 100\n valid:\n num_samples: 100\n evaluate:\n batch_size: 8\n num_samples: 10\n generate:\n batch_size: 8\n num_samples: 10\n segment_duration: 10\n\nloss:\n kind: mse\n norm_power: 0.\n\nvalid:\n every: 1\n\nevaluate:\n every: 5\n num_workers: 5\n metrics:\n visqol: false\n sisnr: false\n rvm: true\n\ngenerate:\n every: 5\n num_workers: 5\n audio:\n sample_rate: ${sample_rate}\n\ncheckpoint:\n save_last: true\n save_every: 25\n keep_last: 10\n keep_every_states: null\n\n\noptim:\n epochs: 50\n updates_per_epoch: 2000\n lr: 2e-4\n max_norm: 0\n optimizer: adam\n adam:\n betas: [0.9, 0.999]\n weight_decay: 0.\n ema:\n use: true # whether to use EMA or not\n updates: 1 # update at every step\n device: ${device} # device for EMA, can be put on GPU if more frequent updates\n decay: 0.99 # EMA decay value, if null, no EMA is used\n\nprocessor:\n name: multi_band_processor\n use: false\n n_bands: 8\n num_samples: 10_000\n power_std: 1.\n\nresampling:\n use: false\n target_sr: 16000\n\nfilter:\n use: false\n n_bands: 4\n idx_band: 0\n cutoffs: null\n\nschedule:\n repartition: \"power\"\n variable_step_batch: true\n beta_t0: 1.0e-5\n beta_t1: 2.9e-2\n beta_exp: 7.5\n num_steps: 1000\n variance: 'beta'\n clip: 5.\n rescale: 1.\n n_bands: null\n noise_scale: 1.0\n\nmetrics:\n num_stage: 4\n" }, { "path": "config/solver/diffusion/default.yaml", "content": "# @package __global__\n\ndefaults:\n - /solver/default\n - /model: score/basic\n - override /dset: audio/default\n - _self_\n\nsolver: diffusion\n\nsample_rate: ???\nchannels: ???\ncompression_model_checkpoint: ???\nn_q: ??? # number of codebooks to keep\n\n\ndataset:\n batch_size: 128\n num_workers: 10\n segment_duration: 1\n train:\n num_samples: 500000\n valid:\n num_samples: 10000\n evaluate:\n batch_size: 16\n num_samples: 10000\n generate:\n batch_size: 32\n num_samples: 50\n segment_duration: 10\n audio:\n sample_rate: ${sample_rate}\n\nloss:\n kind: mse\n norm_power: 0.\n\nvalid:\n every: 1\n\nevaluate:\n every: 20\n num_workers: 5\n metrics:\n visqol: false\n sisnr: false\n rvm: true\n\ngenerate:\n every: 25\n num_workers: 5\n\ncheckpoint:\n save_last: true\n save_every: 25\n keep_last: 10\n keep_every_states: null\n\n\noptim:\n epochs: 20000\n updates_per_epoch: 2000\n lr: 2e-4\n max_norm: 0\n optimizer: adam\n adam:\n betas: [0.9, 0.999]\n weight_decay: 0.\n ema:\n use: true # whether to use EMA or not\n updates: 1 # update at every step\n device: ${device} # device for EMA, can be put on GPU if more frequent updates\n decay: 0.99 # EMA decay value, if null, no EMA is used\n\nprocessor:\n name: multi_band_processor\n use: false\n n_bands: 8\n num_samples: 10_000\n power_std: 1.\n\nresampling:\n use: false\n target_sr: 16000\n\nfilter:\n use: false\n n_bands: 4\n idx_band: 0\n cutoffs: null\n\nschedule:\n repartition: \"power\"\n variable_step_batch: true\n beta_t0: 1.0e-5\n beta_t1: 2.9e-2\n beta_exp: 7.5\n num_steps: 1000\n variance: 'beta'\n clip: 5.\n rescale: 1.\n n_bands: null\n noise_scale: 1.0\n\nmetrics:\n num_stage: 4\n" }, { "path": "config/solver/diffusion/encodec_24khz.yaml", "content": "# @package __global__\n\ndefaults:\n - diffusion/default\n - _self_\n\n\nsample_rate: 24000\nchannels: 1\ncompression_model_checkpoint: //pretrained/facebook/encodec_24khz\nn_q: 4 # num quantizers, 3kbps\n" }, { "path": "config/solver/jasco/chords.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the base MusicGen model (text-to-music)\n# on monophonic audio sampled at 32 kHz\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /dset: audio/example\n - override /conditioner: chords2music\n - _self_\n\nlm_model: flow_matching\nsolver: jasco\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 32khz\n# with a total stride of 640 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //pretrained/facebook/encodec_32khz\n# precomputed mean std accross training data sub-partition\ncompression_model_latent_std: 4.0102\ncompression_model_latent_mean: -0.0074\ncompression_model_framerate: 50\ncompression_model_latent_dim: 128 \n\nefficient_attention_backend: xformers # restricted attention implementation supports only xformers at the moment\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n segment_duration: 10\n batch_size: 320 # 32 GPUs\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n chords_card: ${conditioners.chords.chords_emb.card}\n compression_model_framerate: ${compression_model_framerate}\n\noptim:\n epochs: 500\n optimizer: adamw\n lr: 1e-4\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: cosine\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n\ntransformer_lm:\n causal: false\n skip_connections: false\n flow_dim: ${compression_model_latent_dim}\n chords_dim: ${conditioners.chords.chords_emb.out_dim}\n\ngenerate:\n lm: \n max_prompt_len: null\n max_gen_len: null\n remove_prompts: false\n cfg_coef_all: 3.0\n cfg_coef_txt: 1.0\n prompted_samples: false\n samples:\n prompted: false\n unprompted: true\n" }, { "path": "config/solver/jasco/chords_drums.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the base MusicGen model (text-to-music)\n# on monophonic audio sampled at 32 kHz\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /conditioner: jasco_chords_drums\n - override /dset: audio/default\n - _self_\n\nlm_model: flow_matching\nsolver: jasco\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 32khz\n# with a total stride of 640 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //pretrained/facebook/encodec_32khz\n# precomputed mean std accross training data sub-partition\ncompression_model_latent_std: 4.0102\ncompression_model_latent_mean: -0.0074\ncompression_model_framerate: 50\ncompression_model_latent_dim: 128 \n\nefficient_attention_backend: xformers # restricted attention implementation supports only xformers at the moment\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n segment_duration: 10\n batch_size: 336\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n compression_model_framerate: ${compression_model_framerate}\n\noptim:\n epochs: 500\n optimizer: adamw\n lr: 1e-4\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: cosine\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n\ntransformer_lm:\n causal: false\n skip_connections: false\n flow_dim: ${compression_model_latent_dim}\n chords_dim: ${conditioners.chords.chords_emb.out_dim}\n drums_dim: ${conditioners.self_wav.drum_latents.out_dim}\n\ngenerate:\n lm: \n max_prompt_len: null\n max_gen_len: null\n remove_prompts: false\n cfg_coef_all: 5.0\n cfg_coef_txt: 0.0\n prompted_samples: false\n samples:\n prompted: false\n unprompted: true\n\nconditioners:\n self_wav:\n drum_latents:\n compression_model_latent_dim: ${compression_model_latent_dim}\n compression_model_framerate: ${compression_model_framerate}\n segment_duration: ${dataset.segment_duration}\n" }, { "path": "config/solver/jasco/chords_drums_melody.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the base MusicGen model (text-to-music)\n# on monophonic audio sampled at 32 kHz\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /dset: audio/default\n - override /conditioner: jasco_chords_drums_melody\n - _self_\n\nlm_model: flow_matching\nsolver: jasco\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 32khz\n# with a total stride of 640 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //pretrained/facebook/encodec_32khz\n# precomputed mean std accross training data sub-partition\ncompression_model_latent_std: 4.0102\ncompression_model_latent_mean: -0.0074\ncompression_model_framerate: 50\ncompression_model_latent_dim: 128 \n\nefficient_attention_backend: xformers # restricted attention implementation supports only xformers at the moment\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n segment_duration: 10\n batch_size: 336\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n compression_model_framerate: ${compression_model_framerate}\n melody_kwargs:\n chroma_root: ??? # path to parsed chroma files\n segment_duration: ${dataset.segment_duration}\n melody_fr: 86\n latent_fr: ${compression_model_framerate}\n melody_salience_dim: 53\n override_cache: false\n do_argmax: true\n\noptim:\n epochs: 500\n optimizer: adamw\n lr: 1e-4\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: cosine\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n\ntransformer_lm:\n causal: false\n skip_connections: false\n flow_dim: ${compression_model_latent_dim}\n chords_dim: ${conditioners.chords.chords_emb.out_dim}\n drums_dim: ${conditioners.self_wav.drum_latents.out_dim}\n melody_dim: ${conditioners.melody.melody.out_dim}\n\ngenerate:\n lm: \n max_prompt_len: null\n max_gen_len: null\n remove_prompts: false\n cfg_coef_all: 3.0\n cfg_coef_txt: 1.0\n prompted_samples: false\n samples:\n prompted: false\n unprompted: true\n\nconditioners:\n self_wav:\n drum_latents:\n compression_model_latent_dim: ${compression_model_latent_dim}\n compression_model_framerate: ${compression_model_framerate}\n segment_duration: ${dataset.segment_duration}\n" }, { "path": "config/solver/jasco/drums.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the base MusicGen model (text-to-music)\n# on monophonic audio sampled at 32 kHz\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /dset: audio/default\n - override /conditioner: drums2music\n - _self_\n\nlm_model: flow_matching\nsolver: jasco\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 32khz\n# with a total stride of 640 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //pretrained/facebook/encodec_32khz\n# precomputed mean std accross training data sub-partition\ncompression_model_latent_std: 4.0102\ncompression_model_latent_mean: -0.0074\ncompression_model_framerate: 50\ncompression_model_latent_dim: 128 \n\nefficient_attention_backend: xformers # restricted attention implementation supports only xformers at the moment\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n segment_duration: 10\n batch_size: 336\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n compression_model_framerate: ${compression_model_framerate}\n\noptim:\n epochs: 500\n optimizer: adamw\n lr: 1e-4\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: cosine\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n\ntransformer_lm:\n causal: false\n skip_connections: false\n flow_dim: ${compression_model_latent_dim}\n drums_dim: ${conditioners.self_wav.drum_latents.out_dim}\n\ngenerate:\n lm: \n max_prompt_len: null\n max_gen_len: null\n remove_prompts: false\n cfg_coef_all: 3.0\n cfg_coef_txt: 1.0\n prompted_samples: false\n samples:\n prompted: false\n unprompted: true\n\nconditioners:\n self_wav:\n drum_latents:\n compression_model_latent_dim: ${compression_model_latent_dim}\n compression_model_framerate: ${compression_model_framerate}\n segment_duration: ${dataset.segment_duration}\n" }, { "path": "config/solver/jasco/jasco_32khz_base.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the base MusicGen model (text-to-music)\n# on monophonic audio sampled at 32 kHz\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /dset: audio/default\n - _self_\n\nlm_model: flow_matching\nsolver: jasco\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 32khz\n# with a total stride of 640 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //pretrained/facebook/encodec_32khz\n# precomputed mean std accross training data sub-partition\ncompression_model_latent_std: 4.0102\ncompression_model_latent_mean: -0.0074\ncompression_model_framerate: 50\ncompression_model_latent_dim: 128 \n\nefficient_attention_backend: xformers # restricted attention implementation supports only xformers at the moment\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n segment_duration: 10\n batch_size: 320 # 32 GPUs\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n\noptim:\n epochs: 500\n optimizer: adamw\n lr: 1e-4\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: cosine\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n\ntransformer_lm:\n causal: false\n skip_connections: false\n flow_dim: ${compression_model_latent_dim}\n\ngenerate:\n lm: \n max_prompt_len: null\n max_gen_len: null\n remove_prompts: false\n cfg_coef_all: 3.0\n cfg_coef_txt: 0.0\n\n prompted_samples: false\n samples:\n prompted: false\n unprompted: true\n" }, { "path": "config/solver/magnet/audio_magnet_16khz.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the base audio-MAGNeT model (text-to-sound)\n# on monophonic audio sampled at 16 kHz\n# using a similar EnCodec+LM setup to MAGNeT\ndefaults:\n - audiogen/default\n - /model: lm/audiogen_lm\n - override /dset: audio/default\n - _self_\n\nlm_model: transformer_lm_magnet\nsolver: audio_magnet\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 16khz\n# with a total stride of 320 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //reference/bd44a852/checkpoint.th\n\nchannels: 1\nsample_rate: 16000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n batch_size: 128 # matching AudioGen paper setup (256 * mix_p=0.5 = 128)\n num_workers: 10\n segment_duration: 10\n min_segment_ratio: 1.0\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n external_metadata_source: null\n # sample mixing augmentation at train time\n train:\n batch_size: 256 # matching AudioGen paper setup\n aug_p: 0.5 # perform audio mixing 50% of the time\n mix_p: 0.5 # proportion of batch items mixed together\n # important: note that this will reduce the\n # actual batch size used at train time\n # which will be equal to mix_p * batch_size\n mix_snr_low: -5\n mix_snr_high: 5\n mix_min_overlap: 0.5\n\noptim:\n epochs: 100\n optimizer: adamw\n lr: 5e-4\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: inverse_sqrt\n inverse_sqrt:\n warmup: 3000\n warmup_init_lr: 0.0\n\ncodebooks_pattern:\n modeling: parallel\n parallel:\n empty_initial: -1\n \ntransformer_lm:\n card: 2048\n causal: false\n subcodes_context: 5\n compression_model_framerate: 50 # NOTE: Must match the actual frame rate of the used compression model \n segment_duration: 0\n span_len: -1\n\nmasking:\n span_len: 3\n\ngenerate:\n lm: \n max_prompt_len: null\n max_gen_len: null\n remove_prompts: false\n use_sampling: true\n temp: 3.5\n top_k: 0\n top_p: 0.8\n max_cfg_coef: 20.0\n min_cfg_coef: 1.0\n decoding_steps: [20, 10, 10, 10]\n anneal_temp: true\n span_scoring: 'max'\n span_arrangement: 'nonoverlap'\n prompted_samples: false\n samples:\n prompted: false\n unprompted: true\n\n" }, { "path": "config/solver/magnet/magnet_32khz.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the base MusicGen model (text-to-music)\n# on monophonic audio sampled at 32 kHz\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /dset: audio/default\n - _self_\n\nlm_model: transformer_lm_magnet\nsolver: magnet\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 32khz\n# with a total stride of 640 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //pretrained/facebook/encodec_32khz\n\nefficient_attention_backend: xformers # restricted attention implementation supports only xformers at the moment\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n batch_size: 192 # 32 GPUs\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n\noptim:\n epochs: 500\n optimizer: dadam\n lr: 1\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: cosine\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n\ncodebooks_pattern:\n modeling: parallel\n parallel:\n empty_initial: -1\n \ntransformer_lm:\n card: 2048\n causal: false\n subcodes_context: 5\n compression_model_framerate: 50 # NOTE: Must match the actual frame rate of the used compression model \n segment_duration: 0\n span_len: -1\n\nmasking:\n span_len: 3\n\ngenerate:\n lm: \n max_prompt_len: null\n max_gen_len: null\n remove_prompts: false\n use_sampling: true\n temp: 3.0\n top_k: 0\n top_p: 0.9\n max_cfg_coef: 10.0\n min_cfg_coef: 1.0\n decoding_steps: [60, 10, 10, 10]\n anneal_temp: true\n span_scoring: 'max'\n span_arrangement: 'nonoverlap'\n prompted_samples: false\n samples:\n prompted: false\n unprompted: true\n" }, { "path": "config/solver/musicgen/debug.yaml", "content": "# @package __global__\n\n# This is a minimal debugging configuration\n# for MusicGen training solver\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /model/lm/model_scale: xsmall\n - override /dset: audio/example\n - _self_\n\nautocast: false\ncompression_model_checkpoint: //pretrained/debug_compression_model\ntransformer_lm:\n n_q: 4\n card: 400\n\nconditioners:\n description:\n model: t5\n t5:\n name: t5-small\n\ncodebooks_pattern:\n modeling: parallel\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: false # deadlock detection\n\ndataset:\n batch_size: 4\n segment_duration: 5\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n\ngenerate:\n audio:\n strategy: peak\n lm:\n use_sampling: false\n top_k: 0\n top_p: 0.0\n\ncheckpoint:\n save_every: 0\n keep_last: 0\n\noptim:\n epochs: 2\n updates_per_epoch: 10\n optimizer: adamw\n lr: 1e-4\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: null\n" }, { "path": "config/solver/musicgen/default.yaml", "content": "# @package __global__\n\ndefaults:\n - /solver/default\n - /conditioner: none\n - _self_\n - /solver/musicgen/evaluation: none\n - override /dset: audio/default\n\nautocast: true\nautocast_dtype: float16\n\nsolver: musicgen\nsample_rate: ???\nchannels: ???\ncompression_model_checkpoint: ???\n# The following will set the num codebooks on the underlying\n# model, this might be different from the actual value for n_q\n# given to the transformer, when the model output is postprocessed, for instance\n# for stereo channels. If not provided, default value for the compression model\n# will be used.\ncompression_model_n_q: null\n\ntokens:\n padding_with_special_token: false\n\ninterleave_stereo_codebooks:\n use: false\n per_timestep: false\n\ncache:\n path:\n write: false\n write_shard: 0\n write_num_shards: 1\n\n\ndataset:\n batch_size: 128\n num_workers: 10\n segment_duration: 30\n min_segment_ratio: 0.8 # lower values such as 0.5 result in generations with a lot of silence.\n return_info: true\n train:\n num_samples: 1000000 # need a randomly large number here for AudioDataset\n valid:\n num_samples: 10000\n generate:\n num_samples: 50\n\nmetrics:\n fad:\n use_gt: false\n model: tf\n tf:\n bin: null # path to local frechet_audio_distance code\n model_path: //reference/fad/vggish_model.ckpt\n kld:\n use_gt: false\n model: passt\n passt:\n pretrained_length: 20\n text_consistency:\n use_gt: false\n model: clap\n clap:\n model_path: //reference/clap/music_audioset_epoch_15_esc_90.14.pt\n model_arch: 'HTSAT-base'\n enable_fusion: false\n chroma_cosine:\n use_gt: false\n model: chroma_base\n chroma_base:\n sample_rate: ${sample_rate}\n n_chroma: 12\n radix2_exp: 14\n argmax: true\n\ngenerate:\n every: 25\n num_workers: 5\n path: samples\n audio:\n format: wav\n strategy: loudness\n sample_rate: ${sample_rate}\n loudness_headroom_db: 14\n lm:\n prompted_samples: true\n unprompted_samples: true\n no_text_conditioning: false\n gen_gt_samples: false\n prompt_duration: null # if not set, will use dataset.generate.segment_duration / 4\n gen_duration: null # if not set, will use dataset.generate.segment_duration\n remove_prompts: false\n # generation params\n use_sampling: false\n temp: 1.0\n top_k: 0\n top_p: 0.0\n\nevaluate:\n every: 25\n num_workers: 5\n metrics:\n base: false\n fad: false\n kld: false\n text_consistency: false\n chroma_cosine: false\n\ncheckpoint:\n save_last: true\n save_every: 50\n keep_last: 10\n keep_every_states: null\n\noptim:\n epochs: 200\n updates_per_epoch: 2000\n lr: 1e-4\n optimizer: adamw\n max_norm: 1.0\n eager_sync: true\n adam:\n betas: [0.9, 0.95]\n weight_decay: 0.1\n eps: 1e-8\n\nschedule:\n lr_scheduler: null\n" }, { "path": "config/solver/musicgen/evaluation/none.yaml", "content": "# @package __global__\n\ndataset:\n evaluate:\n num_samples: 10000\n" }, { "path": "config/solver/musicgen/evaluation/objective_eval.yaml", "content": "# @package __global__\n\n# Setup for execute only on musiccaps for audio generation\n# evaluation with objective metrics\n# execute_only=evaluate\n\ndataset:\n max_audio_duration: null\n # ensure the proper values are broadcasted here for evaluate\n evaluate:\n min_audio_duration: 1. # some metrics requires a minimum audio length\n max_audio_duration: null # all samples from musiccaps should be < 20s\n num_samples: null\n segment_duration: null\n generate:\n min_audio_duration: 1.\n max_audio_duration: null\n num_samples: 500\n\nevaluate:\n metrics:\n fad: true\n kld: true\n text_consistency: true\n" }, { "path": "config/solver/musicgen/musicgen_base_32khz.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the base MusicGen model (text-to-music)\n# on monophonic audio sampled at 32 kHz\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /dset: audio/default\n - _self_\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 32khz\n# with a total stride of 640 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //pretrained/facebook/encodec_32khz\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n batch_size: 192 # 32 GPUs\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n\ngenerate:\n lm:\n use_sampling: true\n top_k: 250\n top_p: 0.0\n\noptim:\n epochs: 500\n optimizer: dadam\n lr: 1\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: cosine\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n" }, { "path": "config/solver/musicgen/musicgen_melody_32khz.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for the melody MusicGen model (text+chroma to music)\n# on monophonic audio sampled at 32 kHz\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /conditioner: chroma2music\n - override /dset: audio/default\n - _self_\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 32khz\n# with a total stride of 640 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //pretrained/facebook/encodec_32khz\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n batch_size: 192 # 32 GPUs\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n\ngenerate:\n lm:\n use_sampling: true\n top_k: 250\n top_p: 0.0\n\noptim:\n epochs: 500\n optimizer: dadam\n lr: 1\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: cosine\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n" }, { "path": "config/solver/musicgen/musicgen_style_32khz.yaml", "content": "# @package __global__\n\n# This is the training loop solver\n# for MusicGen-Style model (text-and-style-to-music)\n# on monophonic audio sampled at 32 kHz\ndefaults:\n - musicgen/default\n - /model: lm/musicgen_lm\n - override /conditioner: style2music\n - override /dset: audio/default\n - _self_\n\nautocast: true\nautocast_dtype: float16\n\n# EnCodec large trained on mono-channel music audio sampled at 32khz\n# with a total stride of 640 leading to 50 frames/s.\n# rvq.n_q=4, rvq.bins=2048, no quantization dropout\n# (transformer_lm card and n_q must be compatible)\ncompression_model_checkpoint: //pretrained/facebook/encodec_32khz\n\nchannels: 1\nsample_rate: 32000\n\ndeadlock:\n use: true # deadlock detection\n\ndataset:\n batch_size: 192 # 32 GPUs\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n\ngenerate:\n lm:\n use_sampling: true\n top_k: 250\n top_p: 0.0\n cfg_coef: 3.0\n cfg_coef_beta:\n\noptim:\n epochs: 500\n optimizer: dadam\n lr: 1\n ema:\n use: true\n updates: 10\n device: cuda\n\nlogging:\n log_tensorboard: true\n\nschedule:\n lr_scheduler: cosine\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n" }, { "path": "config/solver/watermark/debug.yaml", "content": "# @package __global__\n\ndefaults:\n - /solver/default\n - /augmentations/default\n - /model: watermark/default\n - override /dset: audio/example\n - _self_\n\nsolver: watermarking # standard name to load the solver using builders\nsample_rate: 48000\nchannels: 1\n\n# all the defaults form compression\nlosses:\n adv: 4.\n feat: 4.\n l1: 0.1\n mel: 0.0\n msspec: 2.0\n sisnr: 0.0\n wm_detection: 1.0 # loss for first 2 bits cannot be 0 \n wm_mb: 1.0 # loss for the rest of the bits (wm message)\n tf_loudnessratio: 10.0\n\nbalancer:\n balance_grads: true\n ema_decay: 0.999\n per_batch_item: true\n total_norm: 1.\n\ncrop:\n prob: 0.4\n shuffle_prob: 0.2\n pad_prob: 0.2 # shuffle_prob + pad_prob + prob <= 1\n size: 0.5\n max_n_windows: 5\n\nadversarial:\n every: 1\n adversaries: [msstftd]\n adv_loss: hinge\n feat_loss: l1\n\ntf_loudnessratio:\n sample_rate: ${sample_rate}\n segment: 0.5\n overlap: 0.5\n n_bands: 16\n temperature: 1.0\n\n# watermarking: audioseal\n\n# losses hyperparameters\nl1: {}\nl2: {}\n\nwm_detection:\n p_weight: 1\n n_weight: 1\n\nwm_mb:\n loss_type: bce # loss between decoded and original\n temperature: 0.1 # decoded is divided by temperature before loss computation\n\nspec_range:\n n_fft: 2048\n min_frequency: 300.0\n max_frequency: 15000.0\n sample_rate: ${sample_rate}\nspec_entropy_range:\n n_fft: 2048\n min_frequency: 300.0\n max_frequency: 15000.0\n sample_rate: ${sample_rate}\nmrstft:\n factor_sc: .5\n factor_mag: .5\n normalized: false\nmel:\n sample_rate: ${sample_rate}\n n_fft: 1024\n hop_length: 256\n win_length: 1024\n n_mels: 64\n f_min: 64\n f_max: null\n normalized: false\n floor_level: 1e-5\nsisnr:\n sample_rate: ${sample_rate}\n segment: 5.\nmsspec:\n sample_rate: ${sample_rate}\n range_start: 6\n range_end: 11\n n_mels: 64\n f_min: 64\n f_max: null\n normalized: true\n alphas: false\n floor_level: 1e-5\n\n# metrics\nmetrics:\n visqol:\n mode: audio\n bin: null # path to visqol install\n model: tcdaudio14_aacvopus_coresv_svrnsim_n.68_g.01_c1.model # visqol v3\n\n# adversaries hyperparameters\nmsstftd:\n in_channels: 1\n out_channels: 1\n filters: 32\n norm: weight_norm\n n_ffts: [1024, 2048, 512, 256, 128]\n hop_lengths: [256, 512, 128, 64, 32]\n win_lengths: [1024, 2048, 512, 256, 128]\n activation: LeakyReLU\n activation_params: { negative_slope: 0.3 }\nmsd:\n in_channels: 1\n out_channels: 1\n scale_norms: [spectral_norm, weight_norm, weight_norm]\n kernel_sizes: [5, 3]\n filters: 16\n max_filters: 1024\n downsample_scales: [4, 4, 4, 4]\n inner_kernel_sizes: null\n groups: [4, 4, 4, 4]\n strides: null\n paddings: null\n activation: LeakyReLU\n activation_params: { negative_slope: 0.3 }\nmpd:\n in_channels: 1\n out_channels: 1\n periods: [2, 3, 5, 7, 11]\n n_layers: 5\n kernel_size: 5\n stride: 3\n filters: 8\n filter_scales: 4\n max_filters: 1024\n activation: LeakyReLU\n activation_params: { negative_slope: 0.3 }\n norm: weight_norm\n\n# data hyperparameters\ndataset:\n batch_size: 16\n num_workers: 10\n segment_duration: 1\n sample_on_weight: false # Uniform sampling all the way\n sample_on_duration: false # Uniform sampling all the way\n\n generate:\n batch_size: 16\n num_samples: 50\n segment_duration: 30\n\n# solver hyperparameters\nevaluate:\n every: 10\n num_workers: 5\n metrics:\n visqol: false\n sisnr: true\ngenerate:\n every: 10\n num_workers: 5\n audio:\n sample_rate: ${sample_rate}\n\n# checkpointing schedule\ncheckpoint:\n save_last: true\n save_every: 25\n keep_last: 10\n keep_every_states: null\n\n\n\n# optimization hyperparameters\noptim:\n epochs: 2\n updates_per_epoch: 10\n lr: 5e-5\n max_norm: 3.0\n optimizer: adam\n adam:\n betas: [0.5, 0.9]\n weight_decay: 0.\n ema:\n use: true # whether to use EMA or not\n updates: 1 # update at every step\n device: ${device} # device for EMA, can be put on GPU if more frequent updates\n decay: 0.99 # EMA decay value, if null, no EMA is used\n\n \nschedule:\n lr_scheduler: \"cosine\"\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n" }, { "path": "config/solver/watermark/default.yaml", "content": "# @package __global__\n\ndefaults:\n - /solver/default\n - /augmentations/default\n - override /dset: audio/example\n - _self_\n\nsolver: watermarking # standard name to load the solver using builders\nsample_rate: ???\nchannels: ???\n\n# all the defaults form compression\nlosses:\n adv: 4.\n feat: 4.\n l1: 0.1\n mel: 0.0\n msspec: 2.0\n sisnr: 0.0\n wm_detection: 1.0 # loss for first 2 bits cannot be 0 \n wm_mb: 1.0 # loss for the rest of the bits (wm message)\n tf_loudnessratio: 10.0\n\nbalancer:\n balance_grads: true\n ema_decay: 0.999\n per_batch_item: true\n total_norm: 1.\n\ncrop:\n prob: 0.4\n shuffle_prob: 0.2\n pad_prob: 0.2 # shuffle_prob + pad_prob + prob <= 1\n size: 0.5\n max_n_windows: 5\n\nadversarial:\n every: 1\n adversaries: [msstftd]\n adv_loss: hinge\n feat_loss: l1\n\ntf_loudnessratio:\n sample_rate: ${sample_rate}\n segment: 0.5\n overlap: 0.5\n n_bands: 16\n temperature: 1.0\n\n# watermarking: audioseal\n\n# losses hyperparameters\nl1: {}\nl2: {}\n\nwm_detection:\n p_weight: 1\n n_weight: 1\n\nwm_mb:\n loss_type: bce # loss between decoded and original\n temperature: 0.1 # decoded is divided by temperature before loss computation\n\nspec_range:\n n_fft: 2048\n min_frequency: 300.0\n max_frequency: 15000.0\n sample_rate: ${sample_rate}\nspec_entropy_range:\n n_fft: 2048\n min_frequency: 300.0\n max_frequency: 15000.0\n sample_rate: ${sample_rate}\nmrstft:\n factor_sc: .5\n factor_mag: .5\n normalized: false\nmel:\n sample_rate: ${sample_rate}\n n_fft: 1024\n hop_length: 256\n win_length: 1024\n n_mels: 64\n f_min: 64\n f_max: null\n normalized: false\n floor_level: 1e-5\nsisnr:\n sample_rate: ${sample_rate}\n segment: 5.\nmsspec:\n sample_rate: ${sample_rate}\n range_start: 6\n range_end: 11\n n_mels: 64\n f_min: 64\n f_max: null\n normalized: true\n alphas: false\n floor_level: 1e-5\n\n# metrics\nmetrics:\n visqol:\n mode: audio\n bin: null # path to visqol install\n model: tcdaudio14_aacvopus_coresv_svrnsim_n.68_g.01_c1.model # visqol v3\n\n# adversaries hyperparameters\nmsstftd:\n in_channels: 1\n out_channels: 1\n filters: 32\n norm: weight_norm\n n_ffts: [1024, 2048, 512, 256, 128]\n hop_lengths: [256, 512, 128, 64, 32]\n win_lengths: [1024, 2048, 512, 256, 128]\n activation: LeakyReLU\n activation_params: { negative_slope: 0.3 }\nmsd:\n in_channels: 1\n out_channels: 1\n scale_norms: [spectral_norm, weight_norm, weight_norm]\n kernel_sizes: [5, 3]\n filters: 16\n max_filters: 1024\n downsample_scales: [4, 4, 4, 4]\n inner_kernel_sizes: null\n groups: [4, 4, 4, 4]\n strides: null\n paddings: null\n activation: LeakyReLU\n activation_params: { negative_slope: 0.3 }\nmpd:\n in_channels: 1\n out_channels: 1\n periods: [2, 3, 5, 7, 11]\n n_layers: 5\n kernel_size: 5\n stride: 3\n filters: 8\n filter_scales: 4\n max_filters: 1024\n activation: LeakyReLU\n activation_params: { negative_slope: 0.3 }\n norm: weight_norm\n\n# data hyperparameters\ndataset:\n batch_size: 16\n num_workers: 10\n segment_duration: 1\n train:\n num_samples: 500000\n valid:\n num_samples: 10000\n evaluate:\n batch_size: 16\n num_samples: 10000\n segment_duration: 10\n\n generate:\n batch_size: 16\n num_samples: 50\n segment_duration: 30\n\n# solver hyperparameters\nevaluate:\n every: 10\n num_workers: 5\n metrics:\n visqol: false\n sisnr: true\ngenerate:\n every: 10\n num_workers: 5\n audio:\n sample_rate: ${sample_rate}\n\n# checkpointing schedule\ncheckpoint:\n save_last: true\n save_every: 25\n keep_last: 10\n keep_every_states: null\n\n\n\n# optimization hyperparameters\noptim:\n epochs: 300\n updates_per_epoch: 2000\n lr: 5e-5\n max_norm: 3.0\n optimizer: adam\n adam:\n betas: [0.5, 0.9]\n weight_decay: 0.\n ema:\n use: true # whether to use EMA or not\n updates: 1 # update at every step\n device: ${device} # device for EMA, can be put on GPU if more frequent updates\n decay: 0.99 # EMA decay value, if null, no EMA is used\n\n \nschedule:\n lr_scheduler: \"cosine\"\n cosine:\n warmup: 4000\n lr_min_ratio: 0.0\n cycle_length: 1.0\n" }, { "path": "config/solver/watermark/robustness.yaml", "content": "# @package __global__\ndefaults:\n - watermark/default\n - /augmentations/default\n - /model: watermark/default\n - _self_\n\nsample_rate: 16000\nchannels: 1\n\nbalancer:\n balance_grads: true\n ema_decay: 0.999\n per_batch_item: true\n total_norm: 1.\n" }, { "path": "config/teams/default.yaml", "content": "default:\n dora_dir: /tmp/audiocraft_${oc.env:USER}\n partitions:\n global: debug\n team: debug\n reference_dir: /tmp\ndarwin: # if we detect we are on a Mac, then most likely we are doing unit testing etc.\n dora_dir: /tmp/audiocraft_${oc.env:USER}\n partitions:\n global: debug\n team: debug\n reference_dir: /tmp\n" }, { "path": "config/teams/labs.yaml", "content": "aws:\n dora_dir: /fsx-audio-craft-llm/${oc.env:USER}/experiments/audiocraft/outputs\n partitions:\n global: learnlab\n team: learnlab\n reference_dir: /fsx-audio-craft-llm/shared/audiocraft/reference\n dataset_mappers:\n \"^/checkpoint/[a-z]+\": \"/fsx-audio-craft-llm\"\nfair:\n dora_dir: /checkpoint/${oc.env:USER}/experiments/audiocraft/outputs\n partitions:\n global: learnlab\n team: learnlab\n reference_dir: /large_experiments/audiocraft/reference\n dataset_mappers:\n \"^/datasets01/datasets01\": \"/datasets01\"\ndarwin:\n dora_dir: /tmp/audiocraft_${oc.env:USER}\n partitions:\n global: debug\n team: debug\n reference_dir: /tmp\nrsc:\n dora_dir: /checkpoint/audiocraft/${oc.env:USER}/experiments/audiocraft/outputs\n partitions:\n global: learn\n team: learn\n reference_dir: /checkpoint/audiocraft/shared/reference" }, { "path": "dataset/example/electro_1.json", "content": "{\"key\": \"\", \"artist\": \"Voyager I\", \"sample_rate\": 48000, \"file_extension\": \"mp3\", \"description\": \"A cool song from Voyager.\", \"keywords\": \"bright, pulsing, cool\", \"duration\": 15.0, \"bpm\": \"\", \"genre\": \"electronic\", \"title\": \"Enracinement\", \"name\": \"electro_1\", \"instrument\": \"Mix\", \"moods\": [\"uplifting\", \"motivational\"]}\n" }, { "path": "dataset/example/electro_2.json", "content": "{\"key\": \"\", \"artist\": \"Voyager I\", \"sample_rate\": 44100, \"file_extension\": \"mp3\", \"description\": \"This is an electronic song sending positive vibes.\", \"keywords\": \"\", \"duration\": 20.0, \"bpm\": \"\", \"genre\": \"electronic\", \"title\": \"Untitled song\", \"name\": \"electro_2\", \"instrument\": \"Mix\", \"moods\": []}\n" }, { "path": "demos/audiogen_demo.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# AudioGen\\n\",\n \"Welcome to AudioGen's demo jupyter notebook. Here you will find a series of self-contained examples of how to use AudioGen in different settings.\\n\",\n \"\\n\",\n \"First, we start by initializing AudioGen. For now, we provide only a medium sized model for AudioGen: `facebook/audiogen-medium` - 1.5B transformer decoder. \\n\",\n \"\\n\",\n \"**Important note:** This variant is different from the original AudioGen model presented at [\\\"AudioGen: Textually-guided audio generation\\\"](https://arxiv.org/abs/2209.15352) as the model architecture is similar to MusicGen with a smaller frame rate and multiple streams of tokens, allowing to reduce generation time.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.models import AudioGen\\n\",\n \"\\n\",\n \"model = AudioGen.get_pretrained('facebook/audiogen-medium')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, let us configure the generation parameters. Specifically, you can control the following:\\n\",\n \"* `use_sampling` (bool, optional): use sampling if True, else do argmax decoding. Defaults to True.\\n\",\n \"* `top_k` (int, optional): top_k used for sampling. Defaults to 250.\\n\",\n \"* `top_p` (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.0.\\n\",\n \"* `temperature` (float, optional): softmax temperature parameter. Defaults to 1.0.\\n\",\n \"* `duration` (float, optional): duration of the generated waveform. Defaults to 10.0.\\n\",\n \"* `cfg_coef` (float, optional): coefficient used for classifier free guidance. Defaults to 3.0.\\n\",\n \"\\n\",\n \"When left unchanged, AudioGen will revert to its default parameters.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model.set_generation_params(\\n\",\n \" use_sampling=True,\\n\",\n \" top_k=250,\\n\",\n \" duration=5\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, we can go ahead and start generating sound using one of the following modes:\\n\",\n \"* Audio continuation using `model.generate_continuation`\\n\",\n \"* Text-conditional samples using `model.generate`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Audio Continuation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import math\\n\",\n \"import torchaudio\\n\",\n \"import torch\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"def get_bip_bip(bip_duration=0.125, frequency=440,\\n\",\n \" duration=0.5, sample_rate=16000, device=\\\"cuda\\\"):\\n\",\n \" \\\"\\\"\\\"Generates a series of bip bip at the given frequency.\\\"\\\"\\\"\\n\",\n \" t = torch.arange(\\n\",\n \" int(duration * sample_rate), device=\\\"cuda\\\", dtype=torch.float) / sample_rate\\n\",\n \" wav = torch.cos(2 * math.pi * frequency * t)[None]\\n\",\n \" tp = (t % (2 * bip_duration)) / (2 * bip_duration)\\n\",\n \" envelope = (tp >= 0.5).float()\\n\",\n \" return wav * envelope\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Here we use a synthetic signal to prompt the generated audio.\\n\",\n \"res = model.generate_continuation(\\n\",\n \" get_bip_bip(0.125).expand(2, -1, -1), \\n\",\n \" 16000, ['Whistling with wind blowing', \\n\",\n \" 'Typing on a typewriter'], \\n\",\n \" progress=True)\\n\",\n \"display_audio(res, 16000)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# You can also use any audio from a file. Make sure to trim the file if it is too long!\\n\",\n \"prompt_waveform, prompt_sr = torchaudio.load(\\\"../assets/sirens_and_a_humming_engine_approach_and_pass.mp3\\\")\\n\",\n \"prompt_duration = 2\\n\",\n \"prompt_waveform = prompt_waveform[..., :int(prompt_duration * prompt_sr)]\\n\",\n \"output = model.generate_continuation(prompt_waveform, prompt_sample_rate=prompt_sr, progress=True)\\n\",\n \"display_audio(output, sample_rate=16000)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Text-conditional Generation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"output = model.generate(\\n\",\n \" descriptions=[\\n\",\n \" 'Subway train blowing its horn',\\n\",\n \" 'A cat meowing',\\n\",\n \" ],\\n\",\n \" progress=True\\n\",\n \")\\n\",\n \"display_audio(output, sample_rate=16000)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": []\n }\n ],\n \"metadata\": {\n \"kernelspec\": {\n \"display_name\": \"Python 3 (ipykernel)\",\n \"language\": \"python\",\n \"name\": \"python3\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.9.7\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "demos/jasco_app.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under thmage license found in the\n# LICENSE file in the root directory of this source tree.\nimport argparse\nfrom concurrent.futures import ProcessPoolExecutor\nimport logging\nimport os\nfrom pathlib import Path\nimport subprocess as sp\nimport sys\nfrom tempfile import NamedTemporaryFile\nimport time\nimport typing as tp\nimport torch\nimport gradio as gr # type: ignore\nfrom audiocraft.data.audio_utils import f32_pcm, normalize_audio\nfrom audiocraft.data.audio import audio_write\nfrom audiocraft.models import JASCO\n# flake8: noqa\n\nMODEL = None # Last used model\nSPACE_ID = os.environ.get('SPACE_ID', '')\nMAX_BATCH_SIZE = 12\nINTERRUPTING = False\nMBD = None\n# We have to wrap subprocess call to clean a bit the log when using gr.make_waveform\n_old_call = sp.call\n\n\ndef _call_nostderr(*args, **kwargs):\n # Avoid ffmpeg vomiting on the logs.\n kwargs['stderr'] = sp.DEVNULL\n kwargs['stdout'] = sp.DEVNULL\n _old_call(*args, **kwargs)\n\n\nsp.call = _call_nostderr\n# Preallocating the pool of processes.\npool = ProcessPoolExecutor(4)\npool.__enter__()\n\n\ndef interrupt():\n global INTERRUPTING\n INTERRUPTING = True\n\n\nclass FileCleaner:\n def __init__(self, file_lifetime: float = 3600):\n self.file_lifetime = file_lifetime\n self.files = [] # type: ignore\n\n def add(self, path: tp.Union[str, Path]):\n self._cleanup()\n self.files.append((time.time(), Path(path)))\n\n def _cleanup(self):\n now = time.time()\n for time_added, path in list(self.files):\n if now - time_added > self.file_lifetime:\n if path.exists():\n path.unlink()\n self.files.pop(0)\n else:\n break\n\n\nfile_cleaner = FileCleaner()\n\n\ndef chords_string_to_list(chords: str):\n if chords == '':\n return []\n\n # clean white spaces or [ ] chars\n chords = chords.replace('[', '')\n chords = chords.replace(']', '')\n chords = chords.replace(' ', '')\n chrd_times = [x.split(',') for x in chords[1:-1].split('),(')]\n return [(x[0], float(x[1])) for x in chrd_times]\n\n\ndef load_model(version='facebook/jasco-chords-drums-400M'):\n global MODEL\n print(\"Loading model\", version)\n if MODEL is None or MODEL.name != version:\n MODEL = None # in case loading would crash\n MODEL = JASCO.get_pretrained(version)\n\n\ndef _do_predictions(texts, chords, melody_matrix, drum_prompt, progress=False, gradio_progress=None, **gen_kwargs):\n MODEL.set_generation_params(**gen_kwargs)\n be = time.time()\n\n # preprocess chords: str to list of tuples\n chords = chords_string_to_list(chords)\n\n if melody_matrix is not None:\n melody_matrix = torch.load(melody_matrix.name, weights_only=True)\n if len(melody_matrix.shape) != 2:\n raise gr.Error(f\"Melody matrix should be a torch tensor of shape [n_melody_bins, T]; got: {melody_matrix.shape}\")\n if melody_matrix.shape[0] > melody_matrix.shape[1]:\n melody_matrix = melody_matrix.permute(1, 0)\n\n # preprocess drums\n if drum_prompt is None:\n preprocessed_drums_wav = None\n drums_sr = 32000\n else:\n # gradio loads audio in int PCM 16-bit, we need to convert it to float32\n drums_sr, drums = drum_prompt[0], f32_pcm(torch.from_numpy(drum_prompt[1])).t()\n if drums.dim() == 1:\n drums = drums[None]\n\n drums = normalize_audio(drums, strategy=\"loudness\", loudness_headroom_db=16, sample_rate=drums_sr)\n preprocessed_drums_wav = drums\n try:\n outputs = MODEL.generate_music(descriptions=texts, chords=chords,\n drums_wav=preprocessed_drums_wav,\n melody_salience_matrix=melody_matrix,\n drums_sample_rate=drums_sr, progress=progress)\n except RuntimeError as e:\n raise gr.Error(\"Error while generating \" + e.args[0])\n outputs = outputs.detach().cpu().float()\n out_wavs = []\n for output in outputs:\n with NamedTemporaryFile(\"wb\", suffix=\".wav\", delete=False) as file:\n audio_write(\n file.name, output, MODEL.sample_rate, strategy=\"loudness\",\n loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)\n out_wavs.append(file.name)\n file_cleaner.add(file.name)\n print(\"batch finished\", len(texts), time.time() - be)\n print(\"Tempfiles currently stored: \", len(file_cleaner.files))\n return out_wavs\n\n\ndef predict_full(model,\n text, chords_sym, melody_file,\n drums_file, drums_mic, drum_input_src,\n cfg_coef_all, cfg_coef_txt,\n ode_rtol, ode_atol,\n ode_solver, ode_steps,\n progress=gr.Progress()):\n global INTERRUPTING\n INTERRUPTING = False\n progress(0, desc=\"Loading model...\")\n load_model(model)\n\n max_generated = 0\n\n def _progress(generated, to_generate):\n nonlocal max_generated\n max_generated = max(generated, max_generated)\n progress((min(max_generated, to_generate), to_generate))\n if INTERRUPTING:\n raise gr.Error(\"Interrupted.\")\n MODEL.set_custom_progress_callback(_progress)\n\n drums = drums_mic if drum_input_src == \"mic\" else drums_file\n wavs = _do_predictions(\n texts=[text] * 2, # we generate two audio outputs for each input prompt\n chords=chords_sym,\n drum_prompt=drums,\n melody_matrix=melody_file,\n progress=True,\n gradio_progress=progress,\n cfg_coef_all=cfg_coef_all,\n cfg_coef_txt=cfg_coef_txt,\n ode_rtol=ode_rtol,\n ode_atol=ode_atol,\n euler=ode_solver == 'euler',\n euler_steps=ode_steps)\n\n return wavs\n\n\ndef ui_full(launch_kwargs):\n with gr.Blocks() as interface:\n gr.Markdown(\n \"\"\"\n # JASCO\n This is your private demo for [JASCO](https://github.com/facebookresearch/audiocraft),\n A text-to-music model, with temporal control over melodies, chords or beats.\n\n presented at: [\"Joint Audio and Symbolic Conditioning for Temporally Controlled Text-to-Music Generation\"]\n (https://arxiv.org/abs/2406.10970)\n \"\"\"\n )\n # Submit | generated\n with gr.Row():\n with gr.Column():\n submit = gr.Button(\"Submit\")\n # Adapted from https://github.com/rkfg/audiocraft/blob/long/app.py, MIT license.\n _ = gr.Button(\"Interrupt\").click(fn=interrupt, queue=False)\n\n with gr.Column():\n audio_output_0 = gr.Audio(label=\"Generated Audio\", type='filepath')\n audio_output_1 = gr.Audio(label=\"Generated Audio\", type='filepath')\n\n # TEXT | models\n with gr.Row():\n with gr.Column():\n text = gr.Text(label=\"Input Text\",\n value=\"Strings, woodwind, orchestral, symphony.\",\n interactive=True)\n with gr.Column():\n model = gr.Radio([\n 'facebook/jasco-chords-drums-400M', 'facebook/jasco-chords-drums-1B',\n 'facebook/jasco-chords-drums-melody-400M', 'facebook/jasco-chords-drums-melody-1B',\n ],\n label=\"Model\", value='facebook/jasco-chords-drums-melody-400M', interactive=True)\n\n # CHORDS\n gr.Markdown(\"Chords conditions\")\n with gr.Row():\n chords_sym = gr.Text(label=\"Chord Progression\",\n value=\"(C, 0.0), (D, 2.0), (F, 4.0), (Ab, 6.0), (Bb, 7.0), (C, 8.0)\",\n interactive=True)\n\n # DRUMS\n gr.Markdown(\"Drums conditions\")\n with gr.Row():\n drum_input_src = gr.Radio([\"file\", \"mic\"], value=\"file\",\n label=\"Condition on drums (optional) File or Mic\")\n drums_file = gr.Audio(sources=[\"upload\"], type=\"numpy\", label=\"File\",\n interactive=True, elem_id=\"drums-input\")\n\n drums_mic = gr.Audio(sources=[\"microphone\"], type=\"numpy\", label=\"Mic\",\n interactive=True, elem_id=\"drums-mic-input\")\n\n # MELODY\n gr.Markdown(\"Melody conditions\")\n with gr.Row():\n melody_file = gr.File(label=\"Melody File\", interactive=True, elem_id=\"melody-file-input\")\n\n # CFG params\n gr.Markdown(\"Classifier-Free Guidance (CFG) Coefficients:\")\n with gr.Row():\n cfg_coef_all = gr.Number(label=\"ALL\", value=1.25, step=0.25, interactive=True)\n cfg_coef_txt = gr.Number(label=\"TEXT\", value=2.5, step=0.25, interactive=True)\n ode_tol = gr.Number(label=\"ODE solver tolerance (defines error approx stop threshold for dynammic solver)\",\n value=1e-4, step=1e-5, interactive=True)\n ode_solver = gr.Radio([\n 'euler', 'dopri5'\n ],\n label=\"ODE Solver\", value='euler', interactive=True)\n ode_steps = gr.Number(label=\"Steps (for euler solver)\", value=10, step=1, interactive=True)\n\n submit.click(fn=predict_full,\n inputs=[model,\n text, chords_sym, melody_file,\n drums_file, drums_mic, drum_input_src,\n cfg_coef_all, cfg_coef_txt, ode_tol, ode_tol, ode_solver, ode_steps],\n outputs=[audio_output_0, audio_output_1])\n gr.Examples(\n fn=predict_full,\n examples=[\n [\n \"80s pop with groovy synth bass and electric piano\",\n \"(N, 0.0), (C, 0.32), (Dm7, 3.456), (Am, 4.608), (F, 8.32), (C, 9.216)\",\n \"./assets/salience_2.th\",\n \"./assets/salience_2.wav\",\n ],\n [\n \"Strings, woodwind, orchestral, symphony.\", # text\n \"(C, 0.0), (D, 2.0), (F, 4.0), (Ab, 6.0), (Bb, 7.0), (C, 8.0)\", # chords\n None, # melody\n None, # drums\n ],\n [\n \"distortion guitars, heavy rock, catchy beat\",\n \"\",\n None,\n \"./assets/sep_drums_1.mp3\",\n ],\n [\n \"hip hop beat with a catchy melody and a groovy bass line\",\n \"\",\n None,\n \"./assets/CJ_Beatbox_Loop_05_90.wav\",\n ],\n [\n \"hip hop beat with a catchy melody and a groovy bass line\",\n \"(C, 0.0), (D, 2.0), (F, 4.0), (Ab, 6.0), (Bb, 7.0), (C, 8.0)\",\n None,\n \"./assets/CJ_Beatbox_Loop_05_90.wav\",\n ],\n\n ],\n inputs=[text, chords_sym, melody_file, drums_file],\n outputs=[audio_output_0, audio_output_1]\n )\n gr.Markdown(\n \"\"\"\n ### More details\n\n \"JASCO\" model will generate a 10 seconds of music based on textual descriptions together with\n temporal controls such as chords and drum tracks.\n These models were trained with descriptions from a stock music catalog. Descriptions that will work best\n should include some level of details on the instruments present, along with some intended use case\n (e.g. adding \"perfect for a commercial\" can somehow help).\n\n We present 4 model variants:\n 1. facebook/jasco-chords-drums-400M - 10s music generation conditioned on text, chords and drums,400M parameters.\n 2. facebook/jasco-chords-drums-1B - 10s music generation conditioned on text, chords and drums, 1B parameters.\n 3. facebook/jasco-chords-drums-melody-400M - 10s music generation conditioned on text, chords, drums and melody,400M parameters.\n 4. facebook/jasco-chords-drums-melody-1B - 10s music generation conditioned on text, chords, drums and melody, 1B parameters.\n\n See https://github.com/facebookresearch/audiocraft/blob/main/docs/JASCO.md\n for more details.\n \"\"\"\n )\n\n interface.queue().launch(**launch_kwargs)\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\n '--listen',\n type=str,\n default='0.0.0.0' if 'SPACE_ID' in os.environ else '127.0.0.1',\n help='IP to listen on for connections to Gradio',\n )\n parser.add_argument(\n '--username', type=str, default='', help='Username for authentication'\n )\n parser.add_argument(\n '--password', type=str, default='', help='Password for authentication'\n )\n parser.add_argument(\n '--server_port',\n type=int,\n default=0,\n help='Port to run the server listener on',\n )\n parser.add_argument(\n '--inbrowser', action='store_true', help='Open in browser'\n )\n parser.add_argument(\n '--share', action='store_true', help='Share the gradio UI'\n )\n\n args = parser.parse_args()\n\n launch_kwargs = {}\n launch_kwargs['server_name'] = args.listen\n\n if args.username and args.password:\n launch_kwargs['auth'] = (args.username, args.password)\n if args.server_port:\n launch_kwargs['server_port'] = args.server_port\n if args.inbrowser:\n launch_kwargs['inbrowser'] = args.inbrowser\n if args.share:\n launch_kwargs['share'] = args.share\n\n logging.basicConfig(level=logging.INFO, stream=sys.stderr)\n\n # Show the interface\n ui_full(launch_kwargs)\n" }, { "path": "demos/jasco_demo.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# JASCO\\n\",\n \"Welcome to JASCO's demo jupyter notebook. \\n\",\n \"Here you will find a self-contained example of how to use JASCO for temporally controlled music generation.\\n\",\n \"\\n\",\n \"You can choose a model from the following selection:\\n\",\n \"1. facebook/jasco-chords-drums-400M - 10s music generation conditioned on text, chords and drums, 400M parameters\\n\",\n \"2. facebook/jasco-chords-drums-1B - 10s music generation conditioned on text, chords and drums, 1B parameters\\n\",\n \"3. facebook/jasco-chords-drums-melody-400M - 10s music generation conditioned on text, chords, drums and melody, 400M parameters\\n\",\n \"4. facebook/jasco-chords-drums-melody-1B - 10s music generation conditioned on text, chords, drums and melody, 1B parameters\\n\",\n \"\\n\",\n \"First, we start by initializing the JASCO model:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import os \\n\",\n \"from audiocraft.models import JASCO\\n\",\n \"\\n\",\n \"model = JASCO.get_pretrained('facebook/jasco-chords-drums-melody-400M', chords_mapping_path='../assets/chord_to_index_mapping.pkl')\\n\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, let us configure the generation parameters. Specifically, you can control the following:\\n\",\n \"* `cfg_coef_all` (float, optional): Coefficient used for classifier free guidance - fully conditional term. \\n\",\n \" Defaults to 5.0.\\n\",\n \"* `cfg_coef_txt` (float, optional): Coefficient used for classifier free guidance - additional text conditional term. \\n\",\n \" Defaults to 0.0.\\n\",\n \"\\n\",\n \"When left unchanged, JASCO will revert to its default parameters.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model.set_generation_params(\\n\",\n \" cfg_coef_all=0.0,\\n\",\n \" cfg_coef_txt=5.0\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, we can go ahead and start generating music given textual prompts.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Text-conditional Generation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"# set textual prompt\\n\",\n \"text = \\\"Funky groove with electric piano playing blue chords rhythmically\\\"\\n\",\n \"\\n\",\n \"# run the model\\n\",\n \"print(\\\"Generating...\\\") \\n\",\n \"output = model.generate(descriptions=[text], progress=True)\\n\",\n \"\\n\",\n \"# display the result\\n\",\n \"print(f\\\"Text: {text}\\\\n\\\")\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Now we can start adding temporal controls! We begin with conditioning on chord progressions:\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Chords-conditional Generation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model.set_generation_params(\\n\",\n \" cfg_coef_all=1.5,\\n\",\n \" cfg_coef_txt=3.0\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"# set textual prompt\\n\",\n \"text = \\\"Strings, woodwind, orchestral, symphony.\\\"\\n\",\n \"\\n\",\n \"# define chord progression\\n\",\n \"chords = [('C', 0.0), ('D', 2.0), ('F', 4.0), ('Ab', 6.0), ('Bb', 7.0), ('C', 8.0)]\\n\",\n \"\\n\",\n \"# run the model\\n\",\n \"print(\\\"Generating...\\\")\\n\",\n \"output = model.generate_music(descriptions=[text], chords=chords, progress=True)\\n\",\n \"\\n\",\n \"# display the result\\n\",\n \"print(f'Text: {text}')\\n\",\n \"print(f'Chord progression: {chords}')\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, we can condition the generation on drum tracks:\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Drums-conditional Generation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import torchaudio\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"\\n\",\n \"# load drum prompt\\n\",\n \"drums_waveform, sr = torchaudio.load(\\\"../assets/sep_drums_1.mp3\\\")\\n\",\n \"\\n\",\n \"# set textual prompt \\n\",\n \"text = \\\"distortion guitars, heavy rock, catchy beat\\\"\\n\",\n \"\\n\",\n \"# run the model\\n\",\n \"print(\\\"Generating...\\\")\\n\",\n \"output = model.generate_music(\\n\",\n \" descriptions=[text],\\n\",\n \" drums_wav=drums_waveform,\\n\",\n \" drums_sample_rate=sr,\\n\",\n \" progress=True\\n\",\n \")\\n\",\n \"\\n\",\n \"# display the result\\n\",\n \"print('drum prompt:')\\n\",\n \"display_audio(drums_waveform, sample_rate=sr)\\n\",\n \"print(f'Text: {text}')\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We can also combine multiple temporal controls! Let's move on to generating with both chords and drums conditioning:\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Drums + Chords conditioning\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import torchaudio\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"\\n\",\n \"# load drum prompt\\n\",\n \"drums_waveform, sr = torchaudio.load(\\\"../assets/sep_drums_1.mp3\\\")\\n\",\n \"\\n\",\n \"# set textual prompt \\n\",\n \"text = \\\"string quartet, orchestral, dramatic\\\"\\n\",\n \"\\n\",\n \"# define chord progression\\n\",\n \"chords = [('C', 0.0), ('D', 2.0), ('F', 4.0), ('Ab', 6.0), ('Bb', 7.0), ('C', 8.0)]\\n\",\n \"\\n\",\n \"# run the model\\n\",\n \"print(\\\"Generating...\\\")\\n\",\n \"output = model.generate_music(\\n\",\n \" descriptions=[text],\\n\",\n \" drums_wav=drums_waveform,\\n\",\n \" drums_sample_rate=sr,\\n\",\n \" chords=chords,\\n\",\n \" progress=True\\n\",\n \")\\n\",\n \"\\n\",\n \"# display the result\\n\",\n \"print('drum prompt:')\\n\",\n \"display_audio(drums_waveform, sample_rate=sr)\\n\",\n \"print(f'Chord progression: {chords}')\\n\",\n \"print(f'Text: {text}')\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Melody + Drums + Chords conditioning - inference example\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"%matplotlib inline\\n\",\n \"from demucs import pretrained\\n\",\n \"from demucs.apply import apply_model\\n\",\n \"from demucs.audio import convert_audio\\n\",\n \"import torch\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"import matplotlib.pyplot as plt\\n\",\n \"\\n\",\n \"# --------------------------\\n\",\n \"# First, choose file to load\\n\",\n \"# --------------------------\\n\",\n \"fnames = ['salience_1', 'salience_2']\\n\",\n \"chords = [\\n\",\n \" [('N', 0.0), ('Eb7', 1.088000000), ('C#', 4.352000000), ('D', 4.864000000), ('Dm7', 6.720000000), ('G7', 8.256000000), ('Am7b5/G', 9.152000000)], # for salience 1\\n\",\n \" [('N', 0.0), ('C', 0.320000000), ('Dm7', 3.456000000), ('Am', 4.608000000), ('F', 8.320000000), ('C', 9.216000000)] # for salience 2\\n\",\n \"]\\n\",\n \"file_idx = 0 # either 0 or 1\\n\",\n \"\\n\",\n \"\\n\",\n \"# ------------------------------------\\n\",\n \"# display audio, melody map and chords\\n\",\n \"# ------------------------------------\\n\",\n \"def plot_chromagram(tensor):\\n\",\n \" # Check if tensor is a PyTorch tensor\\n\",\n \" if not torch.is_tensor(tensor):\\n\",\n \" raise ValueError('Input should be a PyTorch tensor')\\n\",\n \" tensor = tensor.numpy().T # C, T\\n\",\n \" plt.figure(figsize=(20, 20))\\n\",\n \" plt.imshow(tensor, cmap='binary', interpolation='nearest', origin='lower')\\n\",\n \" plt.show()\\n\",\n \"\\n\",\n \"# load salience and display the corresponding wav\\n\",\n \"melody_prompt_wav, melody_prompt_sr = torchaudio.load(f\\\"../assets/{fnames[file_idx]}.wav\\\")\\n\",\n \"print(\\\"Source melody:\\\")\\n\",\n \"display_audio(melody_prompt_wav, sample_rate=melody_prompt_sr)\\n\",\n \"melody = torch.load(f\\\"../assets/{fnames[file_idx]}.th\\\", weights_only=True)\\n\",\n \"plot_chromagram(melody)\\n\",\n \"print(\\\"Chords:\\\")\\n\",\n \"print(chords[file_idx])\\n\",\n \"\\n\",\n \"# --------------------------------------------------\\n\",\n \"# use demucs to seperate the drums stem from src mix\\n\",\n \"# --------------------------------------------------\\n\",\n \"def _get_drums_stem(wav: torch.Tensor, sample_rate: int) -> torch.Tensor:\\n\",\n \" \\\"\\\"\\\"Get parts of the wav that holds the drums, extracting the main stems from the wav.\\\"\\\"\\\"\\n\",\n \" demucs_model = pretrained.get_model('htdemucs').to('cuda')\\n\",\n \" wav = convert_audio(\\n\",\n \" wav, sample_rate, demucs_model.samplerate, demucs_model.audio_channels) # type: ignore\\n\",\n \" stems = apply_model(demucs_model, wav.cuda().unsqueeze(0), device='cuda').squeeze(0)\\n\",\n \" drum_stem = stems[demucs_model.sources.index('drums')] # extract relevant stems for drums conditioning\\n\",\n \" return convert_audio(drum_stem.cpu(), demucs_model.samplerate, sample_rate, 1) # type: ignore\\n\",\n \"drums_wav = _get_drums_stem(melody_prompt_wav, melody_prompt_sr)\\n\",\n \"print(\\\"Separated drums:\\\")\\n\",\n \"display_audio(drums_wav, sample_rate=melody_prompt_sr)\\n\",\n \"\\n\",\n \"# ----------------------------------\\n\",\n \"# Generate using the loaded controls\\n\",\n \"# ----------------------------------\\n\",\n \"# these are free-form texts written randomly\\n\",\n \"texts = [\\n\",\n \" '90s rock with heavy drums and hammond',\\n\",\n \" '80s pop with groovy synth bass and drum machine',\\n\",\n \" 'folk song with leading accordion',\\n\",\n \"]\\n\",\n \"\\n\",\n \"print(\\\"Generating...\\\")\\n\",\n \"# replacing dynammic solver with simple euler solver\\n\",\n \"model.set_generation_params(cfg_coef_all=1.5, cfg_coef_txt=2.5, euler=True, euler_steps=50) # manually set with euler solver\\n\",\n \"output = model.generate_music(\\n\",\n \" descriptions=texts,\\n\",\n \" chords=chords[file_idx],\\n\",\n \" drums_wav=drums_wav,\\n\",\n \" drums_sample_rate=melody_prompt_sr,\\n\",\n \" melody_salience_matrix=melody.permute(1, 0),\\n\",\n \" progress=True\\n\",\n \")\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n }\n ],\n \"metadata\": {\n \"kernelspec\": {\n \"display_name\": \"jasco_dev\",\n \"language\": \"python\",\n \"name\": \"python3\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.9.19\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "demos/magnet_app.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under thmage license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport argparse\nfrom concurrent.futures import ProcessPoolExecutor\nimport logging\nimport os\nfrom pathlib import Path\nimport subprocess as sp\nimport sys\nfrom tempfile import NamedTemporaryFile\nimport time\nimport typing as tp\nimport warnings\n\nimport gradio as gr\n\nfrom audiocraft.data.audio import audio_write\nfrom audiocraft.models import MAGNeT\n\n\nMODEL = None # Last used model\nSPACE_ID = os.environ.get('SPACE_ID', '')\nMAX_BATCH_SIZE = 12\nN_REPEATS = 2\nINTERRUPTING = False\nMBD = None\n# We have to wrap subprocess call to clean a bit the log when using gr.make_waveform\n_old_call = sp.call\n\nPROD_STRIDE_1 = \"prod-stride1 (new!)\"\n\n\ndef _call_nostderr(*args, **kwargs):\n # Avoid ffmpeg vomiting on the logs.\n kwargs['stderr'] = sp.DEVNULL\n kwargs['stdout'] = sp.DEVNULL\n _old_call(*args, **kwargs)\n\n\nsp.call = _call_nostderr\n# Preallocating the pool of processes.\npool = ProcessPoolExecutor(4)\npool.__enter__()\n\n\ndef interrupt():\n global INTERRUPTING\n INTERRUPTING = True\n\n\nclass FileCleaner:\n def __init__(self, file_lifetime: float = 3600):\n self.file_lifetime = file_lifetime\n self.files = []\n\n def add(self, path: tp.Union[str, Path]):\n self._cleanup()\n self.files.append((time.time(), Path(path)))\n\n def _cleanup(self):\n now = time.time()\n for time_added, path in list(self.files):\n if now - time_added > self.file_lifetime:\n if path.exists():\n path.unlink()\n self.files.pop(0)\n else:\n break\n\n\nfile_cleaner = FileCleaner()\n\n\ndef make_waveform(*args, **kwargs):\n # Further remove some warnings.\n be = time.time()\n with warnings.catch_warnings():\n warnings.simplefilter('ignore')\n out = gr.make_waveform(*args, **kwargs)\n print(\"Make a video took\", time.time() - be)\n return out\n\n\ndef load_model(version='facebook/magnet-small-10secs'):\n global MODEL\n print(\"Loading model\", version)\n if MODEL is None or MODEL.name != version:\n MODEL = None # in case loading would crash\n MODEL = MAGNeT.get_pretrained(version)\n\n\ndef _do_predictions(texts, progress=False, gradio_progress=None, **gen_kwargs):\n MODEL.set_generation_params(**gen_kwargs)\n print(\"new batch\", len(texts), texts)\n be = time.time()\n\n try:\n outputs = MODEL.generate(texts, progress=progress, return_tokens=False)\n except RuntimeError as e:\n raise gr.Error(\"Error while generating \" + e.args[0])\n outputs = outputs.detach().cpu().float()\n pending_videos = []\n out_wavs = []\n for i, output in enumerate(outputs):\n with NamedTemporaryFile(\"wb\", suffix=\".wav\", delete=False) as file:\n audio_write(\n file.name, output, MODEL.sample_rate, strategy=\"loudness\",\n loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)\n if i == 0:\n pending_videos.append(pool.submit(make_waveform, file.name))\n out_wavs.append(file.name)\n file_cleaner.add(file.name)\n out_videos = [pending_video.result() for pending_video in pending_videos]\n for video in out_videos:\n file_cleaner.add(video)\n print(\"batch finished\", len(texts), time.time() - be)\n print(\"Tempfiles currently stored: \", len(file_cleaner.files))\n return out_videos, out_wavs\n\n\ndef predict_batched(texts, melodies):\n max_text_length = 512\n texts = [text[:max_text_length] for text in texts]\n load_model('facebook/magnet-small-10secs')\n res = _do_predictions(texts, melodies)\n return res\n\n\ndef predict_full(model, model_path, text, temperature, topp,\n max_cfg_coef, min_cfg_coef, \n decoding_steps1, decoding_steps2, decoding_steps3, decoding_steps4, \n span_score,\n progress=gr.Progress()):\n global INTERRUPTING\n INTERRUPTING = False\n progress(0, desc=\"Loading model...\")\n model_path = model_path.strip()\n if model_path:\n if not Path(model_path).exists():\n raise gr.Error(f\"Model path {model_path} doesn't exist.\")\n if not Path(model_path).is_dir():\n raise gr.Error(f\"Model path {model_path} must be a folder containing \"\n \"state_dict.bin and compression_state_dict_.bin.\")\n model = model_path\n if temperature < 0:\n raise gr.Error(\"Temperature must be >= 0.\")\n\n load_model(model)\n\n max_generated = 0\n\n def _progress(generated, to_generate):\n nonlocal max_generated\n max_generated = max(generated, max_generated)\n progress((min(max_generated, to_generate), to_generate))\n if INTERRUPTING:\n raise gr.Error(\"Interrupted.\")\n MODEL.set_custom_progress_callback(_progress)\n \n videos, wavs = _do_predictions(\n [text] * N_REPEATS, progress=True,\n temperature=temperature, top_p=topp,\n max_cfg_coef=max_cfg_coef, min_cfg_coef=min_cfg_coef, \n decoding_steps=[decoding_steps1, decoding_steps2, decoding_steps3, decoding_steps4],\n span_arrangement='stride1' if (span_score == PROD_STRIDE_1) else 'nonoverlap',\n gradio_progress=progress)\n\n outputs_ = [videos[0]] + [wav for wav in wavs]\n return tuple(outputs_)\n\ndef ui_full(launch_kwargs):\n with gr.Blocks() as interface:\n gr.Markdown(\n \"\"\"\n # MAGNeT\n This is your private demo for [MAGNeT](https://github.com/facebookresearch/audiocraft),\n A fast text-to-music model, consists of a single, non-autoregressive transformer.\n presented at: [\"Masked Audio Generation using a Single Non-Autoregressive Transformer\"] (https://huggingface.co/papers/2401.04577)\n \"\"\"\n )\n with gr.Row():\n with gr.Column():\n with gr.Row():\n text = gr.Text(label=\"Input Text\", value=\"80s electronic track with melodic synthesizers, catchy beat and groovy bass\", interactive=True)\n with gr.Row():\n submit = gr.Button(\"Submit\")\n # Adapted from https://github.com/rkfg/audiocraft/blob/long/app.py, MIT license.\n _ = gr.Button(\"Interrupt\").click(fn=interrupt, queue=False)\n with gr.Row():\n model = gr.Radio(['facebook/magnet-small-10secs', 'facebook/magnet-medium-10secs',\n 'facebook/magnet-small-30secs', 'facebook/magnet-medium-30secs',\n 'facebook/audio-magnet-small', 'facebook/audio-magnet-medium'],\n label=\"Model\", value='facebook/magnet-small-10secs', interactive=True)\n model_path = gr.Text(label=\"Model Path (custom models)\") \n with gr.Row():\n span_score = gr.Radio([\"max-nonoverlap\", PROD_STRIDE_1],\n label=\"Span Scoring\", value=PROD_STRIDE_1, interactive=True) \n with gr.Row():\n decoding_steps1 = gr.Number(label=\"Decoding Steps (stage 1)\", value=20, interactive=True)\n decoding_steps2 = gr.Number(label=\"Decoding Steps (stage 2)\", value=10, interactive=True)\n decoding_steps3 = gr.Number(label=\"Decoding Steps (stage 3)\", value=10, interactive=True)\n decoding_steps4 = gr.Number(label=\"Decoding Steps (stage 4)\", value=10, interactive=True)\n with gr.Row():\n temperature = gr.Number(label=\"Temperature\", value=3.0, step=0.25, minimum=0, interactive=True)\n topp = gr.Number(label=\"Top-p\", value=0.9, step=0.1, minimum=0, maximum=1, interactive=True)\n max_cfg_coef = gr.Number(label=\"Max CFG coefficient\", value=10.0, minimum=0, interactive=True)\n min_cfg_coef = gr.Number(label=\"Min CFG coefficient\", value=1.0, minimum=0, interactive=True) \n with gr.Column():\n output = gr.Video(label=\"Generated Audio - variation 1\")\n audio_outputs = [gr.Audio(label=f\"Generated Audio - variation {i+1}\", type='filepath') for i in range(N_REPEATS)]\n submit.click(fn=predict_full, \n inputs=[model, model_path, text, \n temperature, topp,\n max_cfg_coef, min_cfg_coef,\n decoding_steps1, decoding_steps2, decoding_steps3, decoding_steps4,\n span_score],\n outputs=[output] + [o for o in audio_outputs])\n gr.Examples(\n fn=predict_full,\n examples=[\n [\n \"80s electronic track with melodic synthesizers, catchy beat and groovy bass\",\n 'facebook/magnet-small-10secs',\n 20, 3.0, 0.9, 10.0,\n ],\n [\n \"80s electronic track with melodic synthesizers, catchy beat and groovy bass. 170 bpm\",\n 'facebook/magnet-small-10secs',\n 20, 3.0, 0.9, 10.0,\n ],\n [\n \"Earthy tones, environmentally conscious, ukulele-infused, harmonic, breezy, easygoing, organic instrumentation, gentle grooves\",\n 'facebook/magnet-medium-10secs',\n 20, 3.0, 0.9, 10.0,\n ],\n [ \"Funky groove with electric piano playing blue chords rhythmically\",\n 'facebook/magnet-medium-10secs',\n 20, 3.0, 0.9, 10.0,\n ],\n [\n \"Rock with saturated guitars, a heavy bass line and crazy drum break and fills.\",\n 'facebook/magnet-small-30secs',\n 60, 3.0, 0.9, 10.0,\n ],\n [ \"A grand orchestral arrangement with thunderous percussion, epic brass fanfares, and soaring strings, creating a cinematic atmosphere fit for a heroic battle\",\n 'facebook/magnet-medium-30secs',\n 60, 3.0, 0.9, 10.0,\n ],\n [ \"Seagulls squawking as ocean waves crash while wind blows heavily into a microphone.\",\n 'facebook/audio-magnet-small', \n 20, 3.5, 0.8, 20.0,\n ],\n [ \"A toilet flushing as music is playing and a man is singing in the distance.\",\n 'facebook/audio-magnet-medium', \n 20, 3.5, 0.8, 20.0,\n ],\n ],\n\n inputs=[text, model, decoding_steps1, temperature, topp, max_cfg_coef],\n outputs=[output]\n )\n\n gr.Markdown(\n \"\"\"\n ### More details\n \n #### Music Generation\n \"magnet\" models will generate a short music extract based on the textual description you provided.\n These models can generate either 10 seconds or 30 seconds of music.\n These models were trained with descriptions from a stock music catalog. Descriptions that will work best\n should include some level of details on the instruments present, along with some intended use case\n (e.g. adding \"perfect for a commercial\" can somehow help).\n\n We present 4 model variants:\n 1. facebook/magnet-small-10secs - a 300M non-autoregressive transformer capable of generating 10-second music conditioned\n on text.\n 2. facebook/magnet-medium-10secs - 1.5B parameters, 10 seconds audio.\n 3. facebook/magnet-small-30secs - 300M parameters, 30 seconds audio.\n 4. facebook/magnet-medium-30secs - 1.5B parameters, 30 seconds audio.\n \n #### Sound-Effect Generation\n \"audio-magnet\" models will generate a 10-second sound effect based on the description you provide. \n\n These models were trained on the following data sources: a subset of AudioSet (Gemmeke et al., 2017), \n [BBC sound effects](https://sound-effects.bbcrewind.co.uk/), AudioCaps (Kim et al., 2019), \n Clotho v2 (Drossos et al., 2020), VGG-Sound (Chen et al., 2020), FSD50K (Fonseca et al., 2021), \n [Free To Use Sounds](https://www.freetousesounds.com/all-in-one-bundle/), [Sonniss Game Effects](https://sonniss.com/gameaudiogdc), \n [WeSoundEffects](https://wesoundeffects.com/we-sound-effects-bundle-2020/), \n [Paramount Motion - Odeon Cinematic Sound Effects](https://www.paramountmotion.com/odeon-sound-effects).\n\n We present 2 model variants:\n 1. facebook/audio-magnet-small - 10 second sound effect generation, 300M parameters.\n 2. facebook/audio-magnet-medium - 10 second sound effect generation, 1.5B parameters.\n\n See [github.com/facebookresearch/audiocraft](https://github.com/facebookresearch/audiocraft/blob/main/docs/MAGNET.md)\n for more details.\n \"\"\"\n )\n\n interface.queue().launch(**launch_kwargs)\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\n '--listen',\n type=str,\n default='0.0.0.0' if 'SPACE_ID' in os.environ else '127.0.0.1',\n help='IP to listen on for connections to Gradio',\n )\n parser.add_argument(\n '--username', type=str, default='', help='Username for authentication'\n )\n parser.add_argument(\n '--password', type=str, default='', help='Password for authentication'\n )\n parser.add_argument(\n '--server_port',\n type=int,\n default=0,\n help='Port to run the server listener on',\n )\n parser.add_argument(\n '--inbrowser', action='store_true', help='Open in browser'\n )\n parser.add_argument(\n '--share', action='store_true', help='Share the gradio UI'\n )\n\n args = parser.parse_args()\n\n launch_kwargs = {}\n launch_kwargs['server_name'] = args.listen\n\n if args.username and args.password:\n launch_kwargs['auth'] = (args.username, args.password)\n if args.server_port:\n launch_kwargs['server_port'] = args.server_port\n if args.inbrowser:\n launch_kwargs['inbrowser'] = args.inbrowser\n if args.share:\n launch_kwargs['share'] = args.share\n\n logging.basicConfig(level=logging.INFO, stream=sys.stderr)\n\n # Show the interface\n ui_full(launch_kwargs)\n" }, { "path": "demos/magnet_demo.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# MAGNeT\\n\",\n \"Welcome to MAGNeT's demo jupyter notebook. \\n\",\n \"Here you will find a self-contained example of how to use MAGNeT for music/sound-effect generation.\\n\",\n \"\\n\",\n \"First, we start by initializing MAGNeT for music generation, you can choose a model from the following selection:\\n\",\n \"1. facebook/magnet-small-10secs - a 300M non-autoregressive transformer capable of generating 10-second music conditioned on text.\\n\",\n \"2. facebook/magnet-medium-10secs - 1.5B parameters, 10 seconds music samples.\\n\",\n \"3. facebook/magnet-small-30secs - 300M parameters, 30 seconds music samples.\\n\",\n \"4. facebook/magnet-medium-30secs - 1.5B parameters, 30 seconds music samples.\\n\",\n \"\\n\",\n \"We will use the `facebook/magnet-small-10secs` variant for the purpose of this demonstration.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.models import MAGNeT\\n\",\n \"\\n\",\n \"model = MAGNeT.get_pretrained('facebook/magnet-small-10secs')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, let us configure the generation parameters. Specifically, you can control the following:\\n\",\n \"* `use_sampling` (bool, optional): use sampling if True, else do argmax decoding. Defaults to True.\\n\",\n \"* `top_k` (int, optional): top_k used for sampling. Defaults to 0.\\n\",\n \"* `top_p` (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.9.\\n\",\n \"* `temperature` (float, optional): Initial softmax temperature parameter. Defaults to 3.0.\\n\",\n \"* `max_clsfg_coef` (float, optional): Initial coefficient used for classifier free guidance. Defaults to 10.0.\\n\",\n \"* `min_clsfg_coef` (float, optional): Final coefficient used for classifier free guidance. Defaults to 1.0.\\n\",\n \"* `decoding_steps` (list of n_q ints, optional): The number of iterative decoding steps, for each of the n_q RVQ codebooks.\\n\",\n \"* `span_arrangement` (str, optional): Use either non-overlapping spans ('nonoverlap') or overlapping spans ('stride1') \\n\",\n \" in the masking scheme. \\n\",\n \"\\n\",\n \"When left unchanged, MAGNeT will revert to its default parameters.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model.set_generation_params(\\n\",\n \" use_sampling=True,\\n\",\n \" top_k=0,\\n\",\n \" top_p=0.9,\\n\",\n \" temperature=3.0,\\n\",\n \" max_cfg_coef=10.0,\\n\",\n \" min_cfg_coef=1.0,\\n\",\n \" decoding_steps=[int(20 * model.lm.cfg.dataset.segment_duration // 10), 10, 10, 10],\\n\",\n \" span_arrangement='stride1'\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, we can go ahead and start generating music given textual prompts.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Text-conditional Generation - Music\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"###### Text-to-music prompts - examples ######\\n\",\n \"text = \\\"80s electronic track with melodic synthesizers, catchy beat and groovy bass\\\"\\n\",\n \"# text = \\\"80s electronic track with melodic synthesizers, catchy beat and groovy bass. 170 bpm\\\"\\n\",\n \"# text = \\\"Earthy tones, environmentally conscious, ukulele-infused, harmonic, breezy, easygoing, organic instrumentation, gentle grooves\\\"\\n\",\n \"# text = \\\"Funky groove with electric piano playing blue chords rhythmically\\\"\\n\",\n \"# text = \\\"Rock with saturated guitars, a heavy bass line and crazy drum break and fills.\\\"\\n\",\n \"# text = \\\"A grand orchestral arrangement with thunderous percussion, epic brass fanfares, and soaring strings, creating a cinematic atmosphere fit for a heroic battle\\\"\\n\",\n \" \\n\",\n \"N_VARIATIONS = 3\\n\",\n \"descriptions = [text for _ in range(N_VARIATIONS)]\\n\",\n \"\\n\",\n \"print(f\\\"text prompt: {text}\\\\n\\\")\\n\",\n \"output = model.generate(descriptions=descriptions, progress=True, return_tokens=True)\\n\",\n \"display_audio(output[0], sample_rate=model.compression_model.sample_rate)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Text-conditional Generation - Sound Effects\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Besides music, MAGNeT models can generate sound effects given textual prompts. \\n\",\n \"First, let's load an Audio-MAGNeT model, out of the following collection: \\n\",\n \"1. facebook/audio-magnet-small - a 300M non-autoregressive transformer capable of generating 10 second sound effects conditioned on text.\\n\",\n \"2. facebook/audio-magnet-medium - 10 second sound effect generation, 1.5B parameters.\\n\",\n \"\\n\",\n \"We will use the `facebook/audio-magnet-small` variant for the purpose of this demonstration.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.models import MAGNeT\\n\",\n \"\\n\",\n \"model = MAGNeT.get_pretrained('facebook/audio-magnet-small')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The recommended parameters for sound generation are a bit different than the defaults in MAGNeT, let's initialize it: \"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model.set_generation_params(\\n\",\n \" use_sampling=True,\\n\",\n \" top_k=0,\\n\",\n \" top_p=0.8,\\n\",\n \" temperature=3.5,\\n\",\n \" max_cfg_coef=20.0,\\n\",\n \" min_cfg_coef=1.0,\\n\",\n \" decoding_steps=[int(20 * model.lm.cfg.dataset.segment_duration // 10), 10, 10, 10],\\n\",\n \" span_arrangement='stride1'\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, we can go ahead and start generating sounds given textual prompts.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \" \\n\",\n \"###### Text-to-audio prompts - examples ######\\n\",\n \"text = \\\"Seagulls squawking as ocean waves crash while wind blows heavily into a microphone.\\\"\\n\",\n \"# text = \\\"A toilet flushing as music is playing and a man is singing in the distance.\\\"\\n\",\n \"\\n\",\n \"N_VARIATIONS = 3\\n\",\n \"descriptions = [text for _ in range(N_VARIATIONS)]\\n\",\n \"\\n\",\n \"print(f\\\"text prompt: {text}\\\\n\\\")\\n\",\n \"output = model.generate(descriptions=descriptions, progress=True, return_tokens=True)\\n\",\n \"display_audio(output[0], sample_rate=model.compression_model.sample_rate)\"\n ]\n }\n ],\n \"metadata\": {\n \"kernelspec\": {\n \"display_name\": \"Python 3 (ipykernel)\",\n \"language\": \"python\",\n \"name\": \"python3\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.10.11\"\n },\n \"vscode\": {\n \"interpreter\": {\n \"hash\": \"b02c911f9b3627d505ea4a19966a915ef21f28afb50dbf6b2115072d27c69103\"\n }\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "demos/musicgen_app.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# Updated to account for UI changes from https://github.com/rkfg/audiocraft/blob/long/app.py\n# also released under the MIT license.\n\nimport argparse\nfrom concurrent.futures import ProcessPoolExecutor\nimport logging\nimport os\nfrom pathlib import Path\nimport subprocess as sp\nimport sys\nfrom tempfile import NamedTemporaryFile\nimport time\nimport typing as tp\nimport warnings\n\nfrom einops import rearrange\nimport torch\nimport gradio as gr\n\nfrom audiocraft.data.audio_utils import convert_audio\nfrom audiocraft.data.audio import audio_write\nfrom audiocraft.models.encodec import InterleaveStereoCompressionModel\nfrom audiocraft.models import MusicGen, MultiBandDiffusion\n\n\nMODEL = None # Last used model\nSPACE_ID = os.environ.get('SPACE_ID', '')\nIS_BATCHED = \"facebook/MusicGen\" in SPACE_ID or 'musicgen-internal/musicgen_dev' in SPACE_ID\nprint(IS_BATCHED)\nMAX_BATCH_SIZE = 12\nBATCHED_DURATION = 15\nINTERRUPTING = False\nMBD = None\n# We have to wrap subprocess call to clean a bit the log when using gr.make_waveform\n_old_call = sp.call\n\n\ndef _call_nostderr(*args, **kwargs):\n # Avoid ffmpeg vomiting on the logs.\n kwargs['stderr'] = sp.DEVNULL\n kwargs['stdout'] = sp.DEVNULL\n _old_call(*args, **kwargs)\n\n\nsp.call = _call_nostderr\n# Preallocating the pool of processes.\npool = ProcessPoolExecutor(4)\npool.__enter__()\n\n\ndef interrupt():\n global INTERRUPTING\n INTERRUPTING = True\n\n\nclass FileCleaner:\n def __init__(self, file_lifetime: float = 3600):\n self.file_lifetime = file_lifetime\n self.files = []\n\n def add(self, path: tp.Union[str, Path]):\n self._cleanup()\n self.files.append((time.time(), Path(path)))\n\n def _cleanup(self):\n now = time.time()\n for time_added, path in list(self.files):\n if now - time_added > self.file_lifetime:\n if path.exists():\n path.unlink()\n self.files.pop(0)\n else:\n break\n \nfile_cleaner = FileCleaner()\n\n\ndef make_waveform(*args, **kwargs):\n # Further remove some warnings.\n be = time.time()\n with warnings.catch_warnings():\n warnings.simplefilter('ignore')\n out = gr.make_waveform(*args, **kwargs)\n print(\"Make a video took\", time.time() - be)\n return out\n\n\ndef load_model(version='facebook/musicgen-melody'):\n global MODEL\n print(\"Loading model\", version)\n if MODEL is None or MODEL.name != version:\n # Clear PyTorch CUDA cache and delete model\n del MODEL\n torch.cuda.empty_cache()\n MODEL = None # in case loading would crash\n MODEL = MusicGen.get_pretrained(version)\n\n\ndef load_diffusion():\n global MBD\n if MBD is None:\n print(\"loading MBD\")\n MBD = MultiBandDiffusion.get_mbd_musicgen()\n\n\ndef _do_predictions(texts, melodies, duration, progress=False, gradio_progress=None, **gen_kwargs):\n MODEL.set_generation_params(duration=duration, **gen_kwargs)\n print(\"new batch\", len(texts), texts, [None if m is None else (m[0], m[1].shape) for m in melodies])\n be = time.time()\n processed_melodies = []\n target_sr = 32000\n target_ac = 1\n for melody in melodies:\n if melody is None:\n processed_melodies.append(None)\n else:\n sr, melody = melody[0], torch.from_numpy(melody[1]).to(MODEL.device).float().t()\n if melody.dim() == 1:\n melody = melody[None]\n melody = melody[..., :int(sr * duration)]\n melody = convert_audio(melody, sr, target_sr, target_ac)\n processed_melodies.append(melody)\n\n try:\n if any(m is not None for m in processed_melodies):\n outputs = MODEL.generate_with_chroma(\n descriptions=texts,\n melody_wavs=processed_melodies,\n melody_sample_rate=target_sr,\n progress=progress,\n return_tokens=USE_DIFFUSION\n )\n else:\n outputs = MODEL.generate(texts, progress=progress, return_tokens=USE_DIFFUSION)\n except RuntimeError as e:\n raise gr.Error(\"Error while generating \" + e.args[0])\n if USE_DIFFUSION:\n if gradio_progress is not None:\n gradio_progress(1, desc='Running MultiBandDiffusion...')\n tokens = outputs[1]\n if isinstance(MODEL.compression_model, InterleaveStereoCompressionModel):\n left, right = MODEL.compression_model.get_left_right_codes(tokens)\n tokens = torch.cat([left, right])\n outputs_diffusion = MBD.tokens_to_wav(tokens)\n if isinstance(MODEL.compression_model, InterleaveStereoCompressionModel):\n assert outputs_diffusion.shape[1] == 1 # output is mono\n outputs_diffusion = rearrange(outputs_diffusion, '(s b) c t -> b (s c) t', s=2)\n outputs = torch.cat([outputs[0], outputs_diffusion], dim=0)\n outputs = outputs.detach().cpu().float()\n pending_videos = []\n out_wavs = []\n for output in outputs:\n with NamedTemporaryFile(\"wb\", suffix=\".wav\", delete=False) as file:\n audio_write(\n file.name, output, MODEL.sample_rate, strategy=\"loudness\",\n loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)\n pending_videos.append(pool.submit(make_waveform, file.name))\n out_wavs.append(file.name)\n file_cleaner.add(file.name)\n out_videos = [pending_video.result() for pending_video in pending_videos]\n for video in out_videos:\n file_cleaner.add(video)\n print(\"batch finished\", len(texts), time.time() - be)\n print(\"Tempfiles currently stored: \", len(file_cleaner.files))\n return out_videos, out_wavs\n\n\ndef predict_batched(texts, melodies):\n max_text_length = 512\n texts = [text[:max_text_length] for text in texts]\n load_model('facebook/musicgen-stereo-melody')\n res = _do_predictions(texts, melodies, BATCHED_DURATION)\n return res\n\n\ndef predict_full(model, model_path, decoder, text, melody, duration, topk, topp, temperature, cfg_coef, progress=gr.Progress()):\n global INTERRUPTING\n global USE_DIFFUSION\n INTERRUPTING = False\n progress(0, desc=\"Loading model...\")\n model_path = model_path.strip()\n if model_path:\n if not Path(model_path).exists():\n raise gr.Error(f\"Model path {model_path} doesn't exist.\")\n if not Path(model_path).is_dir():\n raise gr.Error(f\"Model path {model_path} must be a folder containing \"\n \"state_dict.bin and compression_state_dict_.bin.\")\n model = model_path\n if temperature < 0:\n raise gr.Error(\"Temperature must be >= 0.\")\n if topk < 0:\n raise gr.Error(\"Topk must be non-negative.\")\n if topp < 0:\n raise gr.Error(\"Topp must be non-negative.\")\n\n topk = int(topk)\n if decoder == \"MultiBand_Diffusion\":\n USE_DIFFUSION = True\n progress(0, desc=\"Loading diffusion model...\")\n load_diffusion()\n else:\n USE_DIFFUSION = False\n load_model(model)\n\n max_generated = 0\n\n def _progress(generated, to_generate):\n nonlocal max_generated\n max_generated = max(generated, max_generated)\n progress((min(max_generated, to_generate), to_generate))\n if INTERRUPTING:\n raise gr.Error(\"Interrupted.\")\n MODEL.set_custom_progress_callback(_progress)\n\n videos, wavs = _do_predictions(\n [text], [melody], duration, progress=True,\n top_k=topk, top_p=topp, temperature=temperature, cfg_coef=cfg_coef,\n gradio_progress=progress)\n if USE_DIFFUSION:\n return videos[0], wavs[0], videos[1], wavs[1]\n return videos[0], wavs[0], None, None\n\n\ndef toggle_audio_src(choice):\n if choice == \"mic\":\n return gr.update(source=\"microphone\", value=None, label=\"Microphone\")\n else:\n return gr.update(source=\"upload\", value=None, label=\"File\")\n\n\ndef toggle_diffusion(choice):\n if choice == \"MultiBand_Diffusion\":\n return [gr.update(visible=True)] * 2\n else:\n return [gr.update(visible=False)] * 2\n\n\ndef ui_full(launch_kwargs):\n with gr.Blocks() as interface:\n gr.Markdown(\n \"\"\"\n # MusicGen\n This is your private demo for [MusicGen](https://github.com/facebookresearch/audiocraft),\n a simple and controllable model for music generation\n presented at: [\"Simple and Controllable Music Generation\"](https://huggingface.co/papers/2306.05284)\n \"\"\"\n )\n with gr.Row():\n with gr.Column():\n with gr.Row():\n text = gr.Text(label=\"Input Text\", interactive=True)\n with gr.Column():\n radio = gr.Radio([\"file\", \"mic\"], value=\"file\",\n label=\"Condition on a melody (optional) File or Mic\")\n melody = gr.Audio(sources=[\"upload\"], type=\"numpy\", label=\"File\",\n interactive=True, elem_id=\"melody-input\")\n with gr.Row():\n submit = gr.Button(\"Submit\")\n # Adapted from https://github.com/rkfg/audiocraft/blob/long/app.py, MIT license.\n _ = gr.Button(\"Interrupt\").click(fn=interrupt, queue=False)\n with gr.Row():\n model = gr.Radio([\"facebook/musicgen-melody\", \"facebook/musicgen-medium\", \"facebook/musicgen-small\",\n \"facebook/musicgen-large\", \"facebook/musicgen-melody-large\",\n \"facebook/musicgen-stereo-small\", \"facebook/musicgen-stereo-medium\",\n \"facebook/musicgen-stereo-melody\", \"facebook/musicgen-stereo-large\",\n \"facebook/musicgen-stereo-melody-large\"],\n label=\"Model\", value=\"facebook/musicgen-stereo-melody\", interactive=True)\n model_path = gr.Text(label=\"Model Path (custom models)\")\n with gr.Row():\n decoder = gr.Radio([\"Default\", \"MultiBand_Diffusion\"],\n label=\"Decoder\", value=\"Default\", interactive=True)\n with gr.Row():\n duration = gr.Slider(minimum=1, maximum=120, value=10, label=\"Duration\", interactive=True)\n with gr.Row():\n topk = gr.Number(label=\"Top-k\", value=250, interactive=True)\n topp = gr.Number(label=\"Top-p\", value=0, interactive=True)\n temperature = gr.Number(label=\"Temperature\", value=1.0, interactive=True)\n cfg_coef = gr.Number(label=\"Classifier Free Guidance\", value=3.0, interactive=True)\n with gr.Column():\n output = gr.Video(label=\"Generated Music\")\n audio_output = gr.Audio(label=\"Generated Music (wav)\", type='filepath')\n diffusion_output = gr.Video(label=\"MultiBand Diffusion Decoder\")\n audio_diffusion = gr.Audio(label=\"MultiBand Diffusion Decoder (wav)\", type='filepath')\n submit.click(toggle_diffusion, decoder, [diffusion_output, audio_diffusion], queue=False,\n show_progress=False).then(predict_full, inputs=[model, model_path, decoder, text, melody, duration, topk, topp,\n temperature, cfg_coef],\n outputs=[output, audio_output, diffusion_output, audio_diffusion])\n radio.change(toggle_audio_src, radio, [melody], queue=False, show_progress=False)\n\n gr.Examples(\n fn=predict_full,\n examples=[\n [\n \"An 80s driving pop song with heavy drums and synth pads in the background\",\n \"./assets/bach.mp3\",\n \"facebook/musicgen-stereo-melody\",\n \"Default\"\n ],\n [\n \"A cheerful country song with acoustic guitars\",\n \"./assets/bolero_ravel.mp3\",\n \"facebook/musicgen-stereo-melody\",\n \"Default\"\n ],\n [\n \"90s rock song with electric guitar and heavy drums\",\n None,\n \"facebook/musicgen-stereo-medium\",\n \"Default\"\n ],\n [\n \"a light and cheerly EDM track, with syncopated drums, aery pads, and strong emotions\",\n \"./assets/bach.mp3\",\n \"facebook/musicgen-stereo-melody\",\n \"Default\"\n ],\n [\n \"lofi slow bpm electro chill with organic samples\",\n None,\n \"facebook/musicgen-stereo-medium\",\n \"Default\"\n ],\n [\n \"Punk rock with loud drum and power guitar\",\n None,\n \"facebook/musicgen-stereo-medium\",\n \"MultiBand_Diffusion\"\n ],\n ],\n inputs=[text, melody, model, decoder],\n outputs=[output]\n )\n gr.Markdown(\n \"\"\"\n ### More details\n\n The model will generate a short music extract based on the description you provided.\n The model can generate up to 30 seconds of audio in one pass.\n\n The model was trained with description from a stock music catalog, descriptions that will work best\n should include some level of details on the instruments present, along with some intended use case\n (e.g. adding \"perfect for a commercial\" can somehow help).\n\n Using one of the `melody` model (e.g. `musicgen-melody-*`), you can optionally provide a reference audio\n from which a broad melody will be extracted.\n The model will then try to follow both the description and melody provided.\n For best results, the melody should be 30 seconds long (I know, the samples we provide are not...)\n\n It is now possible to extend the generation by feeding back the end of the previous chunk of audio.\n This can take a long time, and the model might lose consistency. The model might also\n decide at arbitrary positions that the song ends.\n\n **WARNING:** Choosing long durations will take a long time to generate (2min might take ~10min).\n An overlap of 12 seconds is kept with the previously generated chunk, and 18 \"new\" seconds\n are generated each time.\n\n We present 10 model variations:\n 1. facebook/musicgen-melody -- a music generation model capable of generating music condition\n on text and melody inputs. **Note**, you can also use text only.\n 2. facebook/musicgen-small -- a 300M transformer decoder conditioned on text only.\n 3. facebook/musicgen-medium -- a 1.5B transformer decoder conditioned on text only.\n 4. facebook/musicgen-large -- a 3.3B transformer decoder conditioned on text only.\n 5. facebook/musicgen-melody-large -- a 3.3B transformer decoder conditioned on and melody.\n 6. facebook/musicgen-stereo-*: same as the previous models but fine tuned to output stereo audio.\n\n We also present two way of decoding the audio tokens\n 1. Use the default GAN based compression model. It can suffer from artifacts especially\n for crashes, snares etc.\n 2. Use [MultiBand Diffusion](https://arxiv.org/abs/2308.02560). Should improve the audio quality,\n at an extra computational cost. When this is selected, we provide both the GAN based decoded\n audio, and the one obtained with MBD.\n\n See [github.com/facebookresearch/audiocraft](https://github.com/facebookresearch/audiocraft/blob/main/docs/MUSICGEN.md)\n for more details.\n \"\"\"\n )\n\n interface.queue().launch(**launch_kwargs)\n\n\ndef ui_batched(launch_kwargs):\n with gr.Blocks() as demo:\n gr.Markdown(\n \"\"\"\n # MusicGen\n\n This is the demo for [MusicGen](https://github.com/facebookresearch/audiocraft/blob/main/docs/MUSICGEN.md),\n a simple and controllable model for music generation\n presented at: [\"Simple and Controllable Music Generation\"](https://huggingface.co/papers/2306.05284).\n
\n \n \"Duplicate\n for longer sequences, more control and no queue.

\n \"\"\"\n )\n with gr.Row():\n with gr.Column():\n with gr.Row():\n text = gr.Text(label=\"Describe your music\", lines=2, interactive=True)\n with gr.Column():\n radio = gr.Radio([\"file\", \"mic\"], value=\"file\",\n label=\"Condition on a melody (optional) File or Mic\")\n melody = gr.Audio(source=\"upload\", type=\"numpy\", label=\"File\",\n interactive=True, elem_id=\"melody-input\")\n with gr.Row():\n submit = gr.Button(\"Generate\")\n with gr.Column():\n output = gr.Video(label=\"Generated Music\")\n audio_output = gr.Audio(label=\"Generated Music (wav)\", type='filepath')\n submit.click(predict_batched, inputs=[text, melody],\n outputs=[output, audio_output], batch=True, max_batch_size=MAX_BATCH_SIZE)\n radio.change(toggle_audio_src, radio, [melody], queue=False, show_progress=False)\n gr.Examples(\n fn=predict_batched,\n examples=[\n [\n \"An 80s driving pop song with heavy drums and synth pads in the background\",\n \"./assets/bach.mp3\",\n ],\n [\n \"A cheerful country song with acoustic guitars\",\n \"./assets/bolero_ravel.mp3\",\n ],\n [\n \"90s rock song with electric guitar and heavy drums\",\n None,\n ],\n [\n \"a light and cheerly EDM track, with syncopated drums, aery pads, and strong emotions bpm: 130\",\n \"./assets/bach.mp3\",\n ],\n [\n \"lofi slow bpm electro chill with organic samples\",\n None,\n ],\n ],\n inputs=[text, melody],\n outputs=[output]\n )\n gr.Markdown(\"\"\"\n ### More details\n\n The model will generate 15 seconds of audio based on the description you provided.\n The model was trained with description from a stock music catalog, descriptions that will work best\n should include some level of details on the instruments present, along with some intended use case\n (e.g. adding \"perfect for a commercial\" can somehow help).\n\n You can optionally provide a reference audio from which a broad melody will be extracted.\n The model will then try to follow both the description and melody provided.\n For best results, the melody should be 30 seconds long (I know, the samples we provide are not...)\n\n You can access more control (longer generation, more models etc.) by clicking\n the \n \"Duplicate\n (you will then need a paid GPU from HuggingFace).\n If you have a GPU, you can run the gradio demo locally (click the link to our repo below for more info).\n Finally, you can get a GPU for free from Google\n and run the demo in [a Google Colab.](https://ai.honu.io/red/musicgen-colab).\n\n See [github.com/facebookresearch/audiocraft](https://github.com/facebookresearch/audiocraft/blob/main/docs/MUSICGEN.md)\n for more details. All samples are generated with the `stereo-melody` model.\n \"\"\")\n\n demo.queue(max_size=8 * 4).launch(**launch_kwargs)\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\n '--listen',\n type=str,\n default='0.0.0.0' if 'SPACE_ID' in os.environ else '127.0.0.1',\n help='IP to listen on for connections to Gradio',\n )\n parser.add_argument(\n '--username', type=str, default='', help='Username for authentication'\n )\n parser.add_argument(\n '--password', type=str, default='', help='Password for authentication'\n )\n parser.add_argument(\n '--server_port',\n type=int,\n default=0,\n help='Port to run the server listener on',\n )\n parser.add_argument(\n '--inbrowser', action='store_true', help='Open in browser'\n )\n parser.add_argument(\n '--share', action='store_true', help='Share the gradio UI'\n )\n\n args = parser.parse_args()\n\n launch_kwargs = {}\n launch_kwargs['server_name'] = args.listen\n\n if args.username and args.password:\n launch_kwargs['auth'] = (args.username, args.password)\n if args.server_port:\n launch_kwargs['server_port'] = args.server_port\n if args.inbrowser:\n launch_kwargs['inbrowser'] = args.inbrowser\n if args.share:\n launch_kwargs['share'] = args.share\n\n logging.basicConfig(level=logging.INFO, stream=sys.stderr)\n\n # Show the interface\n if IS_BATCHED:\n global USE_DIFFUSION\n USE_DIFFUSION = False\n ui_batched(launch_kwargs)\n else:\n ui_full(launch_kwargs)\n" }, { "path": "demos/musicgen_demo.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# MusicGen\\n\",\n \"Welcome to MusicGen's demo jupyter notebook. Here you will find a series of self-contained examples of how to use MusicGen in different settings.\\n\",\n \"\\n\",\n \"First, we start by initializing MusicGen, you can choose a model from the following selection:\\n\",\n \"1. `facebook/musicgen-small` - 300M transformer decoder.\\n\",\n \"2. `facebook/musicgen-medium` - 1.5B transformer decoder.\\n\",\n \"3. `facebook/musicgen-melody` - 1.5B transformer decoder also supporting melody conditioning.\\n\",\n \"4. `facebook/musicgen-large` - 3.3B transformer decoder.\\n\",\n \"\\n\",\n \"We will use the `facebook/musicgen-small` variant for the purpose of this demonstration.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.models import MusicGen\\n\",\n \"from audiocraft.models import MultiBandDiffusion\\n\",\n \"\\n\",\n \"USE_DIFFUSION_DECODER = False\\n\",\n \"# Using small model, better results would be obtained with `medium` or `large`.\\n\",\n \"model = MusicGen.get_pretrained('facebook/musicgen-small')\\n\",\n \"if USE_DIFFUSION_DECODER:\\n\",\n \" mbd = MultiBandDiffusion.get_mbd_musicgen()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, let us configure the generation parameters. Specifically, you can control the following:\\n\",\n \"* `use_sampling` (bool, optional): use sampling if True, else do argmax decoding. Defaults to True.\\n\",\n \"* `top_k` (int, optional): top_k used for sampling. Defaults to 250.\\n\",\n \"* `top_p` (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.0.\\n\",\n \"* `temperature` (float, optional): softmax temperature parameter. Defaults to 1.0.\\n\",\n \"* `duration` (float, optional): duration of the generated waveform. Defaults to 30.0.\\n\",\n \"* `cfg_coef` (float, optional): coefficient used for classifier free guidance. Defaults to 3.0.\\n\",\n \"\\n\",\n \"When left unchanged, MusicGen will revert to its default parameters.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model.set_generation_params(\\n\",\n \" use_sampling=True,\\n\",\n \" top_k=250,\\n\",\n \" duration=30\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, we can go ahead and start generating music using one of the following modes:\\n\",\n \"* Unconditional samples using `model.generate_unconditional`\\n\",\n \"* Music continuation using `model.generate_continuation`\\n\",\n \"* Text-conditional samples using `model.generate`\\n\",\n \"* Melody-conditional samples using `model.generate_with_chroma`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Music Continuation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import math\\n\",\n \"import torchaudio\\n\",\n \"import torch\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"def get_bip_bip(bip_duration=0.125, frequency=440,\\n\",\n \" duration=0.5, sample_rate=32000, device=\\\"cuda\\\"):\\n\",\n \" \\\"\\\"\\\"Generates a series of bip bip at the given frequency.\\\"\\\"\\\"\\n\",\n \" t = torch.arange(\\n\",\n \" int(duration * sample_rate), device=\\\"cuda\\\", dtype=torch.float) / sample_rate\\n\",\n \" wav = torch.cos(2 * math.pi * 440 * t)[None]\\n\",\n \" tp = (t % (2 * bip_duration)) / (2 * bip_duration)\\n\",\n \" envelope = (tp >= 0.5).float()\\n\",\n \" return wav * envelope\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Here we use a synthetic signal to prompt both the tonality and the BPM\\n\",\n \"# of the generated audio.\\n\",\n \"res = model.generate_continuation(\\n\",\n \" get_bip_bip(0.125).expand(2, -1, -1), \\n\",\n \" 32000, ['Jazz jazz and only jazz', \\n\",\n \" 'Heartful EDM with beautiful synths and chords'], \\n\",\n \" progress=True)\\n\",\n \"display_audio(res, 32000)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# You can also use any audio from a file. Make sure to trim the file if it is too long!\\n\",\n \"prompt_waveform, prompt_sr = torchaudio.load(\\\"../assets/bach.mp3\\\")\\n\",\n \"prompt_duration = 2\\n\",\n \"prompt_waveform = prompt_waveform[..., :int(prompt_duration * prompt_sr)]\\n\",\n \"output = model.generate_continuation(prompt_waveform, prompt_sample_rate=prompt_sr, progress=True, return_tokens=True)\\n\",\n \"display_audio(output[0], sample_rate=32000)\\n\",\n \"if USE_DIFFUSION_DECODER:\\n\",\n \" out_diffusion = mbd.tokens_to_wav(output[1])\\n\",\n \" display_audio(out_diffusion, sample_rate=32000)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Text-conditional Generation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"output = model.generate(\\n\",\n \" descriptions=[\\n\",\n \" #'80s pop track with bassy drums and synth',\\n\",\n \" #'90s rock song with loud guitars and heavy drums',\\n\",\n \" #'Progressive rock drum and bass solo',\\n\",\n \" #'Punk Rock song with loud drum and power guitar',\\n\",\n \" #'Bluesy guitar instrumental with soulful licks and a driving rhythm section',\\n\",\n \" #'Jazz Funk song with slap bass and powerful saxophone',\\n\",\n \" 'drum and bass beat with intense percussions'\\n\",\n \" ],\\n\",\n \" progress=True, return_tokens=True\\n\",\n \")\\n\",\n \"display_audio(output[0], sample_rate=32000)\\n\",\n \"if USE_DIFFUSION_DECODER:\\n\",\n \" out_diffusion = mbd.tokens_to_wav(output[1])\\n\",\n \" display_audio(out_diffusion, sample_rate=32000)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Melody-conditional Generation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import torchaudio\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"model = MusicGen.get_pretrained('facebook/musicgen-melody')\\n\",\n \"model.set_generation_params(duration=8)\\n\",\n \"\\n\",\n \"melody_waveform, sr = torchaudio.load(\\\"../assets/bach.mp3\\\")\\n\",\n \"melody_waveform = melody_waveform.unsqueeze(0).repeat(2, 1, 1)\\n\",\n \"output = model.generate_with_chroma(\\n\",\n \" descriptions=[\\n\",\n \" '80s pop track with bassy drums and synth',\\n\",\n \" '90s rock song with loud guitars and heavy drums',\\n\",\n \" ],\\n\",\n \" melody_wavs=melody_waveform,\\n\",\n \" melody_sample_rate=sr,\\n\",\n \" progress=True, return_tokens=True\\n\",\n \")\\n\",\n \"display_audio(output[0], sample_rate=32000)\\n\",\n \"if USE_DIFFUSION_DECODER:\\n\",\n \" out_diffusion = mbd.tokens_to_wav(output[1])\\n\",\n \" display_audio(out_diffusion, sample_rate=32000)\"\n ]\n }\n ],\n \"metadata\": {\n \"kernelspec\": {\n \"display_name\": \"Python 3 (ipykernel)\",\n \"language\": \"python\",\n \"name\": \"python3\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.9.16\"\n },\n \"vscode\": {\n \"interpreter\": {\n \"hash\": \"b02c911f9b3627d505ea4a19966a915ef21f28afb50dbf6b2115072d27c69103\"\n }\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "demos/musicgen_style_app.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# Updated to account for UI changes from https://github.com/rkfg/audiocraft/blob/long/app.py\n# also released under the MIT license.\n\nimport argparse\nfrom concurrent.futures import ProcessPoolExecutor\nimport logging\nimport os\nfrom pathlib import Path\nimport subprocess as sp\nimport sys\nfrom tempfile import NamedTemporaryFile\nimport time\nimport typing as tp\nimport warnings\n\nfrom einops import rearrange\nimport torch\nimport gradio as gr\n\nfrom audiocraft.data.audio_utils import convert_audio\nfrom audiocraft.data.audio import audio_write\nfrom audiocraft.models import MusicGen, MultiBandDiffusion\n\n\nMODEL = None # Last used model\nSPACE_ID = os.environ.get('SPACE_ID', '')\nINTERRUPTING = False\nMBD = None\n# We have to wrap subprocess call to clean a bit the log when using gr.make_waveform\n_old_call = sp.call\n\n\ndef _call_nostderr(*args, **kwargs):\n # Avoid ffmpeg vomiting on the logs.\n kwargs['stderr'] = sp.DEVNULL\n kwargs['stdout'] = sp.DEVNULL\n _old_call(*args, **kwargs)\n\n\nsp.call = _call_nostderr\n# Preallocating the pool of processes.\npool = ProcessPoolExecutor(4)\npool.__enter__()\n\n\ndef interrupt():\n global INTERRUPTING\n INTERRUPTING = True\n\n\nclass FileCleaner:\n def __init__(self, file_lifetime: float = 3600):\n self.file_lifetime = file_lifetime\n self.files = []\n\n def add(self, path: tp.Union[str, Path]):\n self._cleanup()\n self.files.append((time.time(), Path(path)))\n\n def _cleanup(self):\n now = time.time()\n for time_added, path in list(self.files):\n if now - time_added > self.file_lifetime:\n if path.exists():\n path.unlink()\n self.files.pop(0)\n else:\n break\n \nfile_cleaner = FileCleaner()\n\n\ndef make_waveform(*args, **kwargs):\n # Further remove some warnings.\n be = time.time()\n with warnings.catch_warnings():\n warnings.simplefilter('ignore')\n out = gr.make_waveform(*args, **kwargs)\n print(\"Make a video took\", time.time() - be)\n return out\n\n\ndef load_model(version='facebook/musicgen-style'):\n global MODEL\n print(\"Loading model\", version)\n if MODEL is None or MODEL.name != version:\n # Clear PyTorch CUDA cache and delete model\n del MODEL\n torch.cuda.empty_cache()\n MODEL = None # in case loading would crash\n MODEL = MusicGen.get_pretrained(version)\n\n\ndef load_diffusion():\n global MBD\n if MBD is None:\n print(\"loading MBD\")\n MBD = MultiBandDiffusion.get_mbd_musicgen()\n\n\ndef _do_predictions(texts, melodies, duration, top_k, top_p, temperature, cfg_coef, cfg_coef_beta, eval_q, excerpt_length, progress=False, gradio_progress=None):\n MODEL.set_generation_params(duration=duration, top_k=top_k, top_p=top_p, temperature=temperature, cfg_coef=cfg_coef, cfg_coef_beta=cfg_coef_beta)\n MODEL.set_style_conditioner_params(eval_q=eval_q, excerpt_length=excerpt_length)\n print(\"new batch\", len(texts), texts, [None if m is None else (m[0], m[1].shape) for m in melodies])\n be = time.time()\n processed_melodies = []\n target_sr = 32000\n target_ac = 1\n for melody in melodies:\n if melody is None:\n processed_melodies.append(None)\n else:\n sr, melody = melody[0], torch.from_numpy(melody[1]).to(MODEL.device).float().t()\n if melody.dim() == 1:\n melody = melody[None]\n melody = melody[..., :int(sr * duration)]\n melody = convert_audio(melody, sr, target_sr, target_ac)\n processed_melodies.append(melody)\n\n try:\n if any(m is not None for m in processed_melodies):\n outputs = MODEL.generate_with_chroma(\n descriptions=texts,\n melody_wavs=processed_melodies,\n melody_sample_rate=target_sr,\n progress=progress,\n return_tokens=USE_DIFFUSION\n )\n else:\n outputs = MODEL.generate(texts, progress=progress, return_tokens=USE_DIFFUSION)\n except RuntimeError as e:\n raise gr.Error(\"Error while generating \" + e.args[0])\n if USE_DIFFUSION:\n if gradio_progress is not None:\n gradio_progress(1, desc='Running MultiBandDiffusion...')\n tokens = outputs[1]\n outputs_diffusion = MBD.tokens_to_wav(tokens)\n outputs = torch.cat([outputs[0], outputs_diffusion], dim=0)\n outputs = outputs.detach().cpu().float()\n pending_videos = []\n out_wavs = []\n for output in outputs:\n with NamedTemporaryFile(\"wb\", suffix=\".wav\", delete=False) as file:\n audio_write(\n file.name, output, MODEL.sample_rate, strategy=\"loudness\",\n loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)\n pending_videos.append(pool.submit(make_waveform, file.name))\n out_wavs.append(file.name)\n file_cleaner.add(file.name)\n out_videos = [pending_video.result() for pending_video in pending_videos]\n for video in out_videos:\n file_cleaner.add(video)\n print(\"batch finished\", len(texts), time.time() - be)\n print(\"Tempfiles currently stored: \", len(file_cleaner.files))\n return out_videos, out_wavs\n\n\ndef predict_full(model, model_path, decoder, text, melody, duration, topk, topp, temperature, cfg_coef, double_cfg, cfg_coef_beta, eval_q, excerpt_length, progress=gr.Progress()):\n global INTERRUPTING\n global USE_DIFFUSION\n INTERRUPTING = False\n progress(0, desc=\"Loading model...\")\n model_path = model_path.strip()\n if model_path:\n if not Path(model_path).exists():\n raise gr.Error(f\"Model path {model_path} doesn't exist.\")\n if not Path(model_path).is_dir():\n raise gr.Error(f\"Model path {model_path} must be a folder containing \"\n \"state_dict.bin and compression_state_dict_.bin.\")\n model = model_path\n if temperature < 0:\n raise gr.Error(\"Temperature must be >= 0.\")\n if topk < 0:\n raise gr.Error(\"Topk must be non-negative.\")\n if topp < 0:\n raise gr.Error(\"Topp must be non-negative.\")\n if eval_q < 1 or eval_q > 6:\n raise gr.Error(\"eval_q must be an integer between 1 and 6 included.\")\n if excerpt_length > 4.5:\n raise gr.Error(\"excerpt_length must be <= 4.5 seconds\")\n\n topk = int(topk)\n eval_q = int(eval_q)\n if decoder == \"MultiBand_Diffusion\":\n USE_DIFFUSION = True\n progress(0, desc=\"Loading diffusion model...\")\n load_diffusion()\n else:\n USE_DIFFUSION = False\n load_model(model)\n\n if double_cfg != \"Yes\":\n cfg_coef_beta = None\n max_generated = 0\n\n def _progress(generated, to_generate):\n nonlocal max_generated\n max_generated = max(generated, max_generated)\n progress((min(max_generated, to_generate), to_generate))\n if INTERRUPTING:\n raise gr.Error(\"Interrupted.\")\n MODEL.set_custom_progress_callback(_progress)\n\n videos, wavs = _do_predictions(\n [text], [melody], duration, progress=True,\n top_k=topk, top_p=topp, temperature=temperature, cfg_coef=cfg_coef,\n cfg_coef_beta=cfg_coef_beta, eval_q=eval_q, excerpt_length=excerpt_length,\n gradio_progress=progress)\n if USE_DIFFUSION:\n return videos[0], wavs[0], videos[1], wavs[1]\n return videos[0], wavs[0], None, None\n\n\ndef toggle_audio_src(choice):\n if choice == \"mic\":\n return gr.update(source=\"microphone\", value=None, label=\"Microphone\")\n else:\n return gr.update(source=\"upload\", value=None, label=\"File\")\n\n\ndef toggle_diffusion(choice):\n if choice == \"MultiBand_Diffusion\":\n return [gr.update(visible=True)] * 2\n else:\n return [gr.update(visible=False)] * 2\n\n\ndef ui_full(launch_kwargs):\n with gr.Blocks() as interface:\n gr.Markdown(\n \"\"\"\n # MusicGen-Style\n This is your private demo for [MusicGen-Style](https://github.com/facebookresearch/audiocraft),\n a simple and controllable model for music generation\n presented at: [\"Audio Conditioning for Music Generation via Discrete Bottleneck Features\"](https://arxiv.org/abs/2407.12563)\n \"\"\"\n )\n with gr.Row():\n with gr.Column():\n with gr.Row():\n text = gr.Text(label=\"Input Text\", interactive=True)\n with gr.Column():\n radio = gr.Radio([\"file\", \"mic\"], value=\"file\",\n label=\"Condition on a melody (optional) File or Mic\")\n melody = gr.Audio(sources=[\"upload\"], type=\"numpy\", label=\"File\",\n interactive=True, elem_id=\"melody-input\")\n with gr.Row():\n submit = gr.Button(\"Submit\")\n # Adapted from https://github.com/rkfg/audiocraft/blob/long/app.py, MIT license.\n _ = gr.Button(\"Interrupt\").click(fn=interrupt, queue=False)\n with gr.Row():\n model = gr.Radio([\"facebook/musicgen-style\"],\n label=\"Model\", value=\"facebook/musicgen-style\", interactive=True)\n model_path = gr.Text(label=\"Model Path (custom models)\")\n with gr.Row():\n decoder = gr.Radio([\"Default\", \"MultiBand_Diffusion\"],\n label=\"Decoder\", value=\"Default\", interactive=True)\n with gr.Row():\n duration = gr.Slider(minimum=1, maximum=120, value=10, label=\"Duration\", interactive=True)\n eval_q = gr.Slider(minimum=1, maximum=6, value=3, step=1, label=\"Number of RVQ in the style conditioner\", interactive=True)\n with gr.Row():\n topk = gr.Number(label=\"Top-k\", value=250, interactive=True)\n topp = gr.Number(label=\"Top-p\", value=0, interactive=True)\n temperature = gr.Number(label=\"Temperature\", value=1.0, interactive=True)\n cfg_coef = gr.Number(label=\"CFG alpha\", value=3.0, interactive=True)\n double_cfg = gr.Radio([\"Yes\", \"No\"], \n label=\"Use Double Classifier Free Guidance (if No, CFG beta is useless). Only use it if you have input text and a melody file.\", value=\"Yes\", interactive=True)\n cfg_coef_beta = gr.Number(label=\"CFG beta (double CFG)\", value=5.0, interactive=True)\n excerpt_length = gr.Number(label=\"length used of the conditioning (has to be <= 4.5 seconds)\", value=3.0, interactive=True)\n with gr.Column():\n output = gr.Video(label=\"Generated Music\")\n audio_output = gr.Audio(label=\"Generated Music (wav)\", type='filepath')\n diffusion_output = gr.Video(label=\"MultiBand Diffusion Decoder\")\n audio_diffusion = gr.Audio(label=\"MultiBand Diffusion Decoder (wav)\", type='filepath')\n submit.click(toggle_diffusion, decoder, [diffusion_output, audio_diffusion], queue=False,\n show_progress=False).then(predict_full, inputs=[model, model_path, decoder, text, melody, duration, topk, topp,\n temperature, cfg_coef, double_cfg, cfg_coef_beta, eval_q, excerpt_length],\n outputs=[output, audio_output, diffusion_output, audio_diffusion])\n radio.change(toggle_audio_src, radio, [melody], queue=False, show_progress=False)\n\n gr.Examples(\n fn=predict_full,\n examples=[\n [\n \"80s New Wave with synthesizer\",\n \"./assets/electronic.mp3\",\n \"facebook/musicgen-style\",\n \"Default\"\n ],\n ],\n inputs=[text, melody, model, decoder],\n outputs=[output]\n )\n gr.Markdown(\n \"\"\"\n ### More details\n\n The model can generate a short music extract based on 3 different input setups:\n 1) A textual description. In that case we recommend to use simple (not double!) classifier free guidance with the CFG coef = 3.\n\n 2) A audio excerpt that it use for style conditioning. The audio shouldn't be longer that 4.5 seconds. If so, \n a random subsequence will be subsample with the length being chosen by the user. We recommend this length to be between 1.5 and 4.5 seconds. \n We recommend simple CFG with the coef = 3.\n\n 3) Both a textual description and an audio input. In that case the user should use double CFG with alpha=3 and beta=4. Then, if the model \n adheres too much to the text description, the user should lower beta. If the model adheres too much to the style, the user can augment beta. \n The model can generate up to 30 seconds of audio in one pass.\n\n The model was trained with description from a stock music catalog, descriptions that will work best\n should include some level of details on the instruments present, along with some intended use case\n (e.g. adding \"perfect for a commercial\" can somehow help).\n\n We also present two way of decoding the audio tokens\n 1. Use the default GAN based compression model. It can suffer from artifacts especially\n for crashes, snares etc.\n 2. Use [MultiBand Diffusion](https://arxiv.org/abs/2308.02560). Should improve the audio quality,\n at an extra computational cost. When this is selected, we provide both the GAN based decoded\n audio, and the one obtained with MBD.\n\n See [github.com/facebookresearch/audiocraft](https://github.com/facebookresearch/audiocraft/blob/main/docs/MUSICGEN_STYLE.md)\n for more details.\n \"\"\"\n )\n\n interface.queue().launch(**launch_kwargs)\n\n\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\n '--listen',\n type=str,\n default='0.0.0.0' if 'SPACE_ID' in os.environ else '127.0.0.1',\n help='IP to listen on for connections to Gradio',\n )\n parser.add_argument(\n '--username', type=str, default='', help='Username for authentication'\n )\n parser.add_argument(\n '--password', type=str, default='', help='Password for authentication'\n )\n parser.add_argument(\n '--server_port',\n type=int,\n default=0,\n help='Port to run the server listener on',\n )\n parser.add_argument(\n '--inbrowser', action='store_true', help='Open in browser'\n )\n parser.add_argument(\n '--share', action='store_true', help='Share the gradio UI'\n )\n\n args = parser.parse_args()\n\n launch_kwargs = {}\n launch_kwargs['server_name'] = args.listen\n\n if args.username and args.password:\n launch_kwargs['auth'] = (args.username, args.password)\n if args.server_port:\n launch_kwargs['server_port'] = args.server_port\n if args.inbrowser:\n launch_kwargs['inbrowser'] = args.inbrowser\n if args.share:\n launch_kwargs['share'] = args.share\n\n logging.basicConfig(level=logging.INFO, stream=sys.stderr)\n\n # Show the interface\n ui_full(launch_kwargs)" }, { "path": "demos/musicgen_style_demo.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# MusicGen-Style\\n\",\n \"Welcome to MusicGen-Style's demo jupyter notebook. Here you will find a series of self-contained examples of how to use MusicGen-Style in different settings.\\n\",\n \"\\n\",\n \"First, we start by initializing MusicGen-Style.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.models import MusicGen\\n\",\n \"from audiocraft.models import MultiBandDiffusion\\n\",\n \"\\n\",\n \"USE_DIFFUSION_DECODER = False\\n\",\n \"\\n\",\n \"model = MusicGen.get_pretrained('facebook/musicgen-style')\\n\",\n \"if USE_DIFFUSION_DECODER:\\n\",\n \" mbd = MultiBandDiffusion.get_mbd_musicgen()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, let us configure the generation parameters. Specifically, you can control the following:\\n\",\n \"* `use_sampling` (bool, optional): use sampling if True, else do argmax decoding. Defaults to True.\\n\",\n \"* `top_k` (int, optional): top_k used for sampling. Defaults to 250.\\n\",\n \"* `top_p` (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.0.\\n\",\n \"* `temperature` (float, optional): softmax temperature parameter. Defaults to 1.0.\\n\",\n \"* `duration` (float, optional): duration of the generated waveform. Defaults to 30.0.\\n\",\n \"* `cfg_coef` (float, optional): coefficient used for classifier free guidance. Defaults to 3.0.\\n\",\n \"* `cfg_coef_beta` (float, optional): If not None, we use double CFG. cfg_coef_beta is the parameter that pushes the text. Defaults to None, user should start at 5.\\n\",\n \" If the generated music adheres to much to the text, the user should reduce this parameter. If the music adheres too much to the style conditioning, \\n\",\n \" the user should increase it\\n\",\n \"\\n\",\n \"When left unchanged, MusicGen will revert to its default parameters.\\n\",\n \"\\n\",\n \"These are the conditioner parameters for the style conditioner:\\n\",\n \"* `eval_q` (int): integer between 1 and 6 included that tells how many quantizers are used in the RVQ bottleneck\\n\",\n \" of the style conditioner. The higher eval_q is, the more style information passes through the model.\\n\",\n \"* `excerpt_length` (float): float between 1.5 and 4.5 that indicates which length is taken from the audio \\n\",\n \" conditioning to extract style. \\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model.set_generation_params(\\n\",\n \" use_sampling=True,\\n\",\n \" top_k=250,\\n\",\n \" duration=30\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The model can perform text-to-music, style-to-music and text-and-style-to-music.\\n\",\n \"* Text-to-music can be done using `model.generate`, or `model.generate_with_chroma` with the wav condition being None. \\n\",\n \"* Style-to-music and Text-and-Style-to-music can be done using `model.generate_with_chroma`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Text-to-Music\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"model.set_generation_params(\\n\",\n \" duration=8, # generate 8 seconds, can go up to 30\\n\",\n \" use_sampling=True, \\n\",\n \" top_k=250,\\n\",\n \" cfg_coef=3., # Classifier Free Guidance coefficient \\n\",\n \" cfg_coef_beta=None, # double CFG is only useful for text-and-style conditioning\\n\",\n \")\\n\",\n \"\\n\",\n \"output = model.generate(\\n\",\n \" descriptions=[\\n\",\n \" '80s pop track with bassy drums and synth',\\n\",\n \" '90s rock song with loud guitars and heavy drums',\\n\",\n \" 'Progressive rock drum and bass solo',\\n\",\n \" 'Punk Rock song with loud drum and power guitar',\\n\",\n \" 'Bluesy guitar instrumental with soulful licks and a driving rhythm section',\\n\",\n \" 'Jazz Funk song with slap bass and powerful saxophone',\\n\",\n \" 'drum and bass beat with intense percussions'\\n\",\n \" ],\\n\",\n \" progress=True, return_tokens=True\\n\",\n \")\\n\",\n \"display_audio(output[0], sample_rate=32000)\\n\",\n \"if USE_DIFFUSION_DECODER:\\n\",\n \" out_diffusion = mbd.tokens_to_wav(output[1])\\n\",\n \" display_audio(out_diffusion, sample_rate=32000)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Style-to-Music\\n\",\n \"For Style-to-Music, we don't need double CFG. \"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import torchaudio\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"model.set_generation_params(\\n\",\n \" duration=8, # generate 8 seconds, can go up to 30\\n\",\n \" use_sampling=True, \\n\",\n \" top_k=250,\\n\",\n \" cfg_coef=3., # Classifier Free Guidance coefficient \\n\",\n \" cfg_coef_beta=None, # double CFG is only useful for text-and-style conditioning\\n\",\n \")\\n\",\n \"\\n\",\n \"model.set_style_conditioner_params(\\n\",\n \" eval_q=1, # integer between 1 and 6\\n\",\n \" # eval_q is the level of quantization that passes\\n\",\n \" # through the conditioner. When low, the models adheres less to the \\n\",\n \" # audio conditioning\\n\",\n \" excerpt_length=3., # the length in seconds that is taken by the model in the provided excerpt\\n\",\n \" )\\n\",\n \"\\n\",\n \"melody_waveform, sr = torchaudio.load(\\\"../assets/electronic.mp3\\\")\\n\",\n \"melody_waveform = melody_waveform.unsqueeze(0).repeat(2, 1, 1)\\n\",\n \"output = model.generate_with_chroma(\\n\",\n \" descriptions=[None, None], \\n\",\n \" melody_wavs=melody_waveform,\\n\",\n \" melody_sample_rate=sr,\\n\",\n \" progress=True, return_tokens=True\\n\",\n \")\\n\",\n \"display_audio(output[0], sample_rate=32000)\\n\",\n \"if USE_DIFFUSION_DECODER:\\n\",\n \" out_diffusion = mbd.tokens_to_wav(output[1])\\n\",\n \" display_audio(out_diffusion, sample_rate=32000)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Text-and-Style-to-Music\\n\",\n \"For Text-and-Style-to-Music, if we use simple classifier free guidance, the models tends to ignore the text conditioning. We then, introduce double classifier free guidance \\n\",\n \"$$l_{\\\\text{double CFG}} = l_{\\\\emptyset} + \\\\alpha [l_{style} + \\\\beta(l_{text, style} - l_{style}) - l_{\\\\emptyset}]$$\\n\",\n \"\\n\",\n \"For $\\\\beta=1$ we retrieve classic CFG but if $\\\\beta > 1$ we boost the text condition\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import torchaudio\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"model.set_generation_params(\\n\",\n \" duration=8, # generate 8 seconds, can go up to 30\\n\",\n \" use_sampling=True, \\n\",\n \" top_k=250,\\n\",\n \" cfg_coef=3., # Classifier Free Guidance coefficient \\n\",\n \" cfg_coef_beta=5., # double CFG is necessary for text-and-style conditioning\\n\",\n \" # Beta in the double CFG formula. between 1 and 9. When set to 1 \\n\",\n \" # it is equivalent to normal CFG. \\n\",\n \")\\n\",\n \"\\n\",\n \"model.set_style_conditioner_params(\\n\",\n \" eval_q=1, # integer between 1 and 6\\n\",\n \" # eval_q is the level of quantization that passes\\n\",\n \" # through the conditioner. When low, the models adheres less to the \\n\",\n \" # audio conditioning\\n\",\n \" excerpt_length=3., # the length in seconds that is taken by the model in the provided excerpt\\n\",\n \" )\\n\",\n \"\\n\",\n \"melody_waveform, sr = torchaudio.load(\\\"../assets/electronic.mp3\\\")\\n\",\n \"melody_waveform = melody_waveform.unsqueeze(0).repeat(3, 1, 1)\\n\",\n \"\\n\",\n \"descriptions = [\\\"8-bit old video game music\\\", \\\"Chill lofi remix\\\", \\\"80s New wave with synthesizer\\\"]\\n\",\n \"\\n\",\n \"output = model.generate_with_chroma(\\n\",\n \" descriptions=descriptions,\\n\",\n \" melody_wavs=melody_waveform,\\n\",\n \" melody_sample_rate=sr,\\n\",\n \" progress=True, return_tokens=True\\n\",\n \")\\n\",\n \"display_audio(output[0], sample_rate=32000)\\n\",\n \"if USE_DIFFUSION_DECODER:\\n\",\n \" out_diffusion = mbd.tokens_to_wav(output[1])\\n\",\n \" display_audio(out_diffusion, sample_rate=32000)\"\n ]\n }\n ],\n \"metadata\": {\n \"kernelspec\": {\n \"display_name\": \"Python 3 (ipykernel)\",\n \"language\": \"python\",\n \"name\": \"python3\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.9.16\"\n },\n \"vscode\": {\n \"interpreter\": {\n \"hash\": \"b02c911f9b3627d505ea4a19966a915ef21f28afb50dbf6b2115072d27c69103\"\n }\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "docs/AUDIOGEN.md", "content": "# AudioGen: Textually-guided audio generation\n\nAudioCraft provides the code and a model re-implementing AudioGen, a [textually-guided audio generation][audiogen_arxiv]\nmodel that performs text-to-sound generation.\n\nThe provided AudioGen reimplementation follows the LM model architecture introduced in [MusicGen][musicgen_arxiv]\nand is a single stage auto-regressive Transformer model trained over a 16kHz\nEnCodec tokenizer with 4 codebooks sampled at 50 Hz.\nThis model variant reaches similar audio quality than the original implementation introduced in the AudioGen publication\nwhile providing faster generation speed given the smaller frame rate.\n\n**Important note:** The provided models are NOT the original models used to report numbers in the\n[AudioGen publication][audiogen_arxiv]. Refer to the model card to learn more about architectural changes.\n\nListen to samples from the **original AudioGen implementation** in our [sample page][audiogen_samples].\n\n\n## Model Card\n\nSee [the model card](../model_cards/AUDIOGEN_MODEL_CARD.md).\n\n\n## Installation\n\nPlease follow the AudioCraft installation instructions from the [README](../README.md).\n\nAudioCraft requires a GPU with at least 16 GB of memory for running inference with the medium-sized models (~1.5B parameters).\n\n## API and usage\n\nWe provide a simple API and 1 pre-trained models for AudioGen:\n\n`facebook/audiogen-medium`: 1.5B model, text to sound - [🤗 Hub](https://huggingface.co/facebook/audiogen-medium)\n\nYou can play with AudioGen by running the jupyter notebook at [`demos/audiogen_demo.ipynb`](../demos/audiogen_demo.ipynb) locally (if you have a GPU).\n\nSee after a quick example for using the API.\n\n```python\nimport torchaudio\nfrom audiocraft.models import AudioGen\nfrom audiocraft.data.audio import audio_write\n\nmodel = AudioGen.get_pretrained('facebook/audiogen-medium')\nmodel.set_generation_params(duration=5) # generate 5 seconds.\ndescriptions = ['dog barking', 'sirene of an emergency vehicle', 'footsteps in a corridor']\nwav = model.generate(descriptions) # generates 3 samples.\n\nfor idx, one_wav in enumerate(wav):\n # Will save under {idx}.wav, with loudness normalization at -14 db LUFS.\n audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\n## Training\n\nThe [AudioGenSolver](../audiocraft/solvers/audiogen.py) implements the AudioGen's training pipeline\nused to develop the released model. Note that this may not fully reproduce the results presented in the paper.\nSimilarly to MusicGen, it defines an autoregressive language modeling task over multiple streams of\ndiscrete tokens extracted from a pre-trained EnCodec model (see [EnCodec documentation](./ENCODEC.md)\nfor more details on how to train such model) with dataset-specific changes for environmental sound\nprocessing.\n\nNote that **we do NOT provide any of the datasets** used for training AudioGen.\n\n### Example configurations and grids\n\nWe provide configurations to reproduce the released models and our research.\nAudioGen solvers configuration are available in [config/solver/audiogen](../config/solver/audiogen).\nThe base training configuration used for the released models is the following:\n[`solver=audiogen/audiogen_base_16khz`](../config/solver/audiogen/audiogen_base_16khz.yaml)\n\nPlease find some example grids to train AudioGen at\n[audiocraft/grids/audiogen](../audiocraft/grids/audiogen/).\n\n```shell\n# text-to-sound\ndora grid audiogen.audiogen_base_16khz\n```\n\n### Sound dataset and metadata\n\nAudioGen's underlying dataset is an AudioDataset augmented with description metadata.\nThe AudioGen dataset implementation expects the metadata to be available as `.json` files\nat the same location as the audio files or through specified external folder.\nLearn more in the [datasets section](./DATASETS.md).\n\n### Evaluation stage\n\nBy default, evaluation stage is also computing the cross-entropy and the perplexity over the\nevaluation dataset. Indeed the objective metrics used for evaluation can be costly to run\nor require some extra dependencies. Please refer to the [metrics documentation](./METRICS.md)\nfor more details on the requirements for each metric.\n\nWe provide an off-the-shelf configuration to enable running the objective metrics\nfor audio generation in\n[config/solver/audiogen/evaluation/objective_eval](../config/solver/audiogen/evaluation/objective_eval.yaml).\n\nOne can then activate evaluation the following way:\n```shell\n# using the configuration\ndora run solver=audiogen/debug solver/audiogen/evaluation=objective_eval\n# specifying each of the fields, e.g. to activate KL computation\ndora run solver=audiogen/debug evaluate.metrics.kld=true\n```\n\nSee [an example evaluation grid](../audiocraft/grids/audiogen/audiogen_pretrained_16khz_eval.py).\n\n### Generation stage\n\nThe generation stage allows to generate samples conditionally and/or unconditionally and to perform\naudio continuation (from a prompt). We currently support greedy sampling (argmax), sampling\nfrom softmax with a given temperature, top-K and top-P (nucleus) sampling. The number of samples\ngenerated and the batch size used are controlled by the `dataset.generate` configuration\nwhile the other generation parameters are defined in `generate.lm`.\n\n```shell\n# control sampling parameters\ndora run solver=audiogen/debug generate.lm.gen_duration=5 generate.lm.use_sampling=true generate.lm.top_k=15\n```\n\n## More information\n\nRefer to [MusicGen's instructions](./MUSICGEN.md).\n\n### Learn more\n\nLearn more about AudioCraft training pipelines in the [dedicated section](./TRAINING.md).\n\n\n## Citation\n\nAudioGen\n```\n@article{kreuk2022audiogen,\n title={Audiogen: Textually guided audio generation},\n author={Kreuk, Felix and Synnaeve, Gabriel and Polyak, Adam and Singer, Uriel and D{\\'e}fossez, Alexandre and Copet, Jade and Parikh, Devi and Taigman, Yaniv and Adi, Yossi},\n journal={arXiv preprint arXiv:2209.15352},\n year={2022}\n}\n```\n\nMusicGen\n```\n@article{copet2023simple,\n title={Simple and Controllable Music Generation},\n author={Jade Copet and Felix Kreuk and Itai Gat and Tal Remez and David Kant and Gabriel Synnaeve and Yossi Adi and Alexandre Défossez},\n year={2023},\n journal={arXiv preprint arXiv:2306.05284},\n}\n```\n\n## License\n\nSee license information in the [model card](../model_cards/AUDIOGEN_MODEL_CARD.md).\n\n[audiogen_arxiv]: https://arxiv.org/abs/2209.15352\n[musicgen_arxiv]: https://arxiv.org/abs/2306.05284\n[audiogen_samples]: https://felixkreuk.github.io/audiogen/\n" }, { "path": "docs/CONDITIONING.md", "content": "# AudioCraft conditioning modules\n\nAudioCraft provides a\n[modular implementation of conditioning modules](../audiocraft/modules/conditioners.py)\nthat can be used with the language model to condition the generation.\nThe codebase was developed in order to easily extend the set of modules\ncurrently supported to easily develop new ways of controlling the generation.\n\n\n## Conditioning methods\n\nFor now, we support 3 main types of conditioning within AudioCraft:\n* Text-based conditioning methods\n* Waveform-based conditioning methods\n* Joint embedding conditioning methods for text and audio projected in a shared latent space.\n\nThe Language Model relies on 2 core components that handle processing information:\n* The `ConditionProvider` class, that maps metadata to processed conditions, leveraging\nall the defined conditioners for the given task.\n* The `ConditionFuser` class, that takes preprocessed conditions and properly fuse the\nconditioning embedding to the language model inputs following a given fusing strategy.\n\nDifferent conditioners (for text, waveform, joint embeddings...) are provided as torch\nmodules in AudioCraft and are used internally in the language model to process the\nconditioning signals and feed them to the language model.\n\n\n## Core concepts\n\n### Conditioners\n\nThe `BaseConditioner` torch module is the base implementation for all conditioners in AudioCraft.\n\nEach conditioner is expected to implement 2 methods:\n* The `tokenize` method that is used as a preprocessing method that contains all processing\nthat can lead to synchronization points (e.g. BPE tokenization with transfer to the GPU).\nThe output of the tokenize method will then be used to feed the forward method.\n* The `forward` method that takes the output of the tokenize method and contains the core computation\nto obtain the conditioning embedding along with a mask indicating valid indices (e.g. padding tokens).\n\n### ConditionProvider\n\nThe ConditionProvider prepares and provides conditions given a dictionary of conditioners.\n\nConditioners are specified as a dictionary of attributes and the corresponding conditioner\nproviding the processing logic for the given attribute.\n\nSimilarly to the conditioners, the condition provider works in two steps to avoid synchronization points:\n* A `tokenize` method that takes a list of conditioning attributes for the batch,\nand runs all tokenize steps for the set of conditioners.\n* A `forward` method that takes the output of the tokenize step and runs all the forward steps\nfor the set of conditioners.\n\nThe list of conditioning attributes is passed as a list of `ConditioningAttributes`\nthat is presented just below.\n\n### ConditionFuser\n\nOnce all conditioning signals have been extracted and processed by the `ConditionProvider`\nas dense embeddings, they remain to be passed to the language model along with the original\nlanguage model inputs.\n\nThe `ConditionFuser` handles specifically the logic to combine the different conditions\nto the actual model input, supporting different strategies to combine them.\n\nOne can therefore define different strategies to combine or fuse the condition to the input, in particular:\n* Prepending the conditioning signal to the input with the `prepend` strategy,\n* Summing the conditioning signal to the input with the `sum` strategy,\n* Combining the conditioning relying on a cross-attention mechanism with the `cross` strategy,\n* Using input interpolation with the `input_interpolate` strategy.\n\n### SegmentWithAttributes and ConditioningAttributes: From metadata to conditions\n\nThe `ConditioningAttributes` dataclass is the base class for metadata\ncontaining all attributes used for conditioning the language model.\n\nIt currently supports the following types of attributes:\n* Text conditioning attributes: Dictionary of textual attributes used for text-conditioning.\n* Wav conditioning attributes: Dictionary of waveform attributes used for waveform-based\nconditioning such as the chroma conditioning.\n* JointEmbed conditioning attributes: Dictionary of text and waveform attributes\nthat are expected to be represented in a shared latent space.\n\nThese different types of attributes are the attributes that are processed\nby the different conditioners.\n\n`ConditioningAttributes` are extracted from metadata loaded along the audio in the datasets,\nprovided that the metadata used by the dataset implements the `SegmentWithAttributes` abstraction.\n\nAll metadata-enabled datasets to use for conditioning in AudioCraft inherits\nthe [`audiocraft.data.info_dataset.InfoAudioDataset`](../audiocraft/data/info_audio_dataset.py) class\nand the corresponding metadata inherits and implements the `SegmentWithAttributes` abstraction.\nRefer to the [`audiocraft.data.music_dataset.MusicAudioDataset`](../audiocraft/data/music_dataset.py)\nclass as an example.\n\n\n## Available conditioners\n\n### Text conditioners\n\nAll text conditioners are expected to inherit from the `TextConditioner` class.\n\nAudioCraft currently provides two text conditioners:\n* The `LUTConditioner` that relies on look-up-table of embeddings learned at train time,\nand relying on either no tokenizer or a spacy tokenizer. This conditioner is particularly\nuseful for simple experiments and categorical labels.\n* The `T5Conditioner` that relies on a\n[pre-trained T5 model](https://huggingface.co/docs/transformers/model_doc/t5)\nfrozen or fine-tuned at train time to extract the text embeddings.\n\n### Waveform conditioners\n\nAll waveform conditioners are expected to inherit from the `WaveformConditioner` class and\nconsist of a conditioning method that takes a waveform as input. The waveform conditioner\nmust implement the logic to extract the embedding from the waveform and define the downsampling\nfactor from the waveform to the resulting embedding.\n\nThe `ChromaStemConditioner` conditioner is a waveform conditioner for the chroma features\nconditioning used by MusicGen. It takes a given waveform, extracts relevant stems for melody\n(namely all non drums and bass stems) using a\n[pre-trained Demucs model](https://github.com/facebookresearch/demucs)\nand then extracts the chromagram bins from the remaining mix of stems.\n\n### Joint embeddings conditioners\n\nWe finally provide support for conditioning based on joint text and audio embeddings through\nthe `JointEmbeddingConditioner` class and the `CLAPEmbeddingConditioner` that implements such\na conditioning method relying on a [pretrained CLAP model](https://github.com/LAION-AI/CLAP).\n\n## Classifier Free Guidance\n\nWe provide a Classifier Free Guidance implementation in AudioCraft. With the classifier free\nguidance dropout, all attributes are dropped with the same probability.\n\n## Attribute Dropout\n\nWe further provide an attribute dropout strategy. Unlike the classifier free guidance dropout,\nthe attribute dropout drops given attributes with a defined probability, allowing the model\nnot to expect all conditioning signals to be provided at once.\n\n## Faster computation of conditions\n\nConditioners that require some heavy computation on the waveform can be cached, in particular\nthe `ChromaStemConditioner` or `CLAPEmbeddingConditioner`. You just need to provide the\n`cache_path` parameter to them. We recommend running dummy jobs for filling up the cache quickly.\nAn example is provided in the [musicgen.musicgen_melody_32khz grid](../audiocraft/grids/musicgen/musicgen_melody_32khz.py)." }, { "path": "docs/DATASETS.md", "content": "# AudioCraft datasets\n\nOur dataset manifest files consist in 1-json-per-line files, potentially gzipped,\nas `data.jsons` or `data.jsons.gz` files. This JSON contains the path to the audio\nfile and associated metadata. The manifest files are then provided in the configuration,\nas `datasource` sub-configuration. A datasource contains the pointers to the paths of\nthe manifest files for each AudioCraft stage (or split) along with additional information\n(eg. maximum sample rate to use against this dataset). All the datasources are under the\n`dset` group config, with a dedicated configuration file for each dataset.\n\n## Getting started\n\n### Example\n\nSee the provided example in the directory that provides a manifest to use the example dataset\nprovided under the [dataset folder](../dataset/example).\n\nThe manifest files are stored in the [egs folder](../egs/example).\n\n```shell\negs/\n example/data.json.gz\n```\n\nA datasource is defined in the configuration folder, in the dset group config for this dataset\nat [config/dset/audio/example](../config/dset/audio/example.yaml):\n\n```shell\n# @package __global__\n\ndatasource:\n max_sample_rate: 44100\n max_channels: 2\n\n train: egs/example\n valid: egs/example\n evaluate: egs/example\n generate: egs/example\n```\n\nFor proper dataset, one should create manifest for each of the splits and specify the correct path\nto the given manifest in the datasource for each split.\n\nThen, using a dataset through the configuration can be done pointing to the\ncorresponding dataset configuration:\n```shell\ndset= # should match the yaml file name\n\n# for example\ndset=audio/example\n```\n\n### Creating manifest files\n\nAssuming you want to create manifest files to load with AudioCraft's AudioDataset, you can use\nthe following command to create new manifest files from a given folder containing audio files:\n\n```shell\npython -m audiocraft.data.audio_dataset egs/my_dataset/my_dataset_split/data.jsonl.gz\n\n# For example to generate the manifest for dset=audio/example\n# note: we don't use any split and we don't compress the jsonl file for this dummy example\npython -m audiocraft.data.audio_dataset dataset/example egs/example/data.jsonl\n\n# More info with: python -m audiocraft.data.audio_dataset --help\n```\n\n## Additional information\n\n### MusicDataset and metadata\n\nThe MusicDataset is an AudioDataset with additional metadata. The MusicDataset expects\nthe additional metadata to be stored in a JSON file that has the same path as the corresponding\naudio file, but with a `.json` extension.\n\n### SoundDataset and metadata\n\nThe SoundDataset is an AudioDataset with descriptions metadata. Similarly to the MusicDataset,\nthe SoundDataset expects the additional metadata to be stored in a JSON file that has the same\npath as the corresponding audio file, but with a `.json` extension. Additionally, the SoundDataset\nsupports an additional parameter pointing to an extra folder `external_metadata_source` containing\nall the JSON metadata files given they have the same filename as the audio file.\n" }, { "path": "docs/ENCODEC.md", "content": "# EnCodec: High Fidelity Neural Audio Compression\n\nAudioCraft provides the training code for EnCodec, a state-of-the-art deep learning\nbased audio codec supporting both mono and stereo audio, presented in the\n[High Fidelity Neural Audio Compression][arxiv] paper.\nCheck out our [sample page][encodec_samples].\n\n## Original EnCodec models\n\nThe EnCodec models presented in High Fidelity Neural Audio Compression can be accessed\nand used with the [EnCodec repository](https://github.com/facebookresearch/encodec).\n\n**Note**: We do not guarantee compatibility between the AudioCraft and EnCodec codebases\nand released checkpoints at this stage.\n\n\n## Installation\n\nPlease follow the AudioCraft installation instructions from the [README](../README.md).\n\n\n## Training\n\nThe [CompressionSolver](../audiocraft/solvers/compression.py) implements the audio reconstruction\ntask to train an EnCodec model. Specifically, it trains an encoder-decoder with a quantization\nbottleneck - a SEANet encoder-decoder with Residual Vector Quantization bottleneck for EnCodec -\nusing a combination of objective and perceptual losses in the forms of discriminators.\n\nThe default configuration matches a causal EnCodec training at a single bandwidth.\n\n### Example configuration and grids\n\nWe provide sample configuration and grids for training EnCodec models.\n\nThe compression configuration are defined in\n[config/solver/compression](../config/solver/compression).\n\nThe example grids are available at\n[audiocraft/grids/compression](../audiocraft/grids/compression).\n\n```shell\n# base causal encodec on monophonic audio sampled at 24 khz\ndora grid compression.encodec_base_24khz\n# encodec model used for MusicGen on monophonic audio sampled at 32 khz\ndora grid compression.encodec_musicgen_32khz\n```\n\n### Training and validation stages\n\nThe model is trained using a combination of objective and perceptual losses.\nMore specifically, EnCodec is trained with the MS-STFT discriminator along with\nobjective losses through the use of a loss balancer to effectively weight\nthe different losses, in an intuitive manner.\n\n### Evaluation stage\n\nEvaluation metrics for audio generation:\n* SI-SNR: Scale-Invariant Signal-to-Noise Ratio.\n* ViSQOL: Virtual Speech Quality Objective Listener.\n\nNote: Path to the ViSQOL binary (compiled with bazel) needs to be provided in\norder to run the ViSQOL metric on the reference and degraded signals.\nThe metric is disabled by default.\nPlease refer to the [metrics documentation](../METRICS.md) to learn more.\n\n### Generation stage\n\nThe generation stage consists in generating the reconstructed audio from samples\nwith the current model. The number of samples generated and the batch size used are\ncontrolled by the `dataset.generate` configuration. The output path and audio formats\nare defined in the generate stage configuration.\n\n```shell\n# generate samples every 5 epoch\ndora run solver=compression/encodec_base_24khz generate.every=5\n# run with a different dset\ndora run solver=compression/encodec_base_24khz generate.path=\n# limit the number of samples or use a different batch size\ndora grid solver=compression/encodec_base_24khz dataset.generate.num_samples=10 dataset.generate.batch_size=4\n```\n\n### Playing with the model\n\nOnce you have a model trained, it is possible to get the entire solver, or just\nthe trained model with the following functions:\n\n```python\nfrom audiocraft.solvers import CompressionSolver\n\n# If you trained a custom model with signature SIG.\nmodel = CompressionSolver.model_from_checkpoint('//sig/SIG')\n# If you want to get one of the pretrained models with the `//pretrained/` prefix.\nmodel = CompressionSolver.model_from_checkpoint('//pretrained/facebook/encodec_32khz')\n# Or load from a custom checkpoint path\nmodel = CompressionSolver.model_from_checkpoint('/my_checkpoints/foo/bar/checkpoint.th')\n\n\n# If you only want to use a pretrained model, you can also directly get it\n# from the CompressionModel base model class.\nfrom audiocraft.models import CompressionModel\n\n# Here do not put the `//pretrained/` prefix!\nmodel = CompressionModel.get_pretrained('facebook/encodec_32khz')\nmodel = CompressionModel.get_pretrained('dac_44khz')\n\n# Finally, you can also retrieve the full Solver object, with its dataloader etc.\nfrom audiocraft import train\nfrom pathlib import Path\nimport logging\nimport os\nimport sys\n\n# Uncomment the following line if you want some detailed logs when loading a Solver.\n# logging.basicConfig(stream=sys.stderr, level=logging.INFO)\n\n# You must always run the following function from the root directory.\nos.chdir(Path(train.__file__).parent.parent)\n\n\n# You can also get the full solver (only for your own experiments).\n# You can provide some overrides to the parameters to make things more convenient.\nsolver = train.get_solver_from_sig('SIG', {'device': 'cpu', 'dataset': {'batch_size': 8}})\nsolver.model\nsolver.dataloaders\n```\n\n### Importing / Exporting models\n\nAt the moment we do not have a definitive workflow for exporting EnCodec models, for\ninstance to Hugging Face (HF). We are working on supporting automatic conversion between\nAudioCraft and Hugging Face implementations.\n\nWe still have some support for fine-tuning an EnCodec model coming from HF in AudioCraft,\nusing for instance `continue_from=//pretrained/facebook/encodec_32k`.\n\nAn AudioCraft checkpoint can be exported in a more compact format (excluding the optimizer etc.)\nusing `audiocraft.utils.export.export_encodec`. For instance, you could run\n\n```python\nfrom audiocraft.utils import export\nfrom audiocraft import train\nxp = train.main.get_xp_from_sig('SIG')\nexport.export_encodec(\n xp.folder / 'checkpoint.th',\n '/checkpoints/my_audio_lm/compression_state_dict.bin')\n\n\nfrom audiocraft.models import CompressionModel\nmodel = CompressionModel.get_pretrained('/checkpoints/my_audio_lm/compression_state_dict.bin')\n\nfrom audiocraft.solvers import CompressionSolver\n# The two are strictly equivalent, but this function supports also loading from non-already exported models.\nmodel = CompressionSolver.model_from_checkpoint('//pretrained//checkpoints/my_audio_lm/compression_state_dict.bin')\n```\n\nWe will see then how to use this model as a tokenizer for MusicGen/AudioGen in the\n[MusicGen documentation](./MUSICGEN.md).\n\n### Learn more\n\nLearn more about AudioCraft training pipelines in the [dedicated section](./TRAINING.md).\n\n\n## Citation\n```\n@article{defossez2022highfi,\n title={High Fidelity Neural Audio Compression},\n author={Défossez, Alexandre and Copet, Jade and Synnaeve, Gabriel and Adi, Yossi},\n journal={arXiv preprint arXiv:2210.13438},\n year={2022}\n}\n```\n\n\n## License\n\nSee license information in the [README](../README.md).\n\n[arxiv]: https://arxiv.org/abs/2210.13438\n[encodec_samples]: https://ai.honu.io/papers/encodec/samples.html\n" }, { "path": "docs/JASCO.md", "content": "# JASCO: Joint Audio And Symbolic Conditioning for Temporally Controlled Text-To-Music Generation\n\nAudioCraft provides the code and models for JASCO, [Joint Audio And Symbolic Conditioning for Temporally Controlled Text-To-Music Generation][arxiv].\n\nWe present JASCO, a temporally controlled text-to-music generation model utilizing both symbolic and audio-based conditions.\nJASCO can generate high-quality music samples conditioned on global text descriptions along with fine-grained local controls.\nJASCO is based on the Flow Matching modeling paradigm together with a novel conditioning method, allowing for music generation controlled both locally (e.g., chords) and globally (text description).\n\nCheck out our [sample page][sample_page] or test the available demo!\n\nWe use ~16K hours of licensed music to train JASCO. \n\n\n## Model Card\n\nSee [the model card](../model_cards/JASCO_MODEL_CARD.md).\n\n\n## Installation\n\nFirst, Please follow the AudioCraft installation instructions from the [README](../README.md).\n\nThen, download and install chord_extractor from [source](http://www.isophonics.net/nnls-chroma)\n\nSee further required installation under **Data Preprocessing** section\n\n## Usage\n\nWe currently offer two ways to interact with JASCO:\n1. You can use the gradio demo locally by running [`python -m demos.jasco_app`](../demos/jasco_app.py), you can add `--share` to deploy a sharable space mounted on your device.\n2. You can play with JASCO by running the jupyter notebook at [`demos/jasco_demo.ipynb`](../demos/jasco_demo.ipynb) locally.\n\n## API\n\nWe provide a simple API and pre-trained models:\n- `facebook/jasco-chords-drums-400M`: 400M model, text to music with chords and drums support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-400M)\n- `facebook/jasco-chords-drums-1B`: 1B model, text to music with chords and drums support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-1B)\n- `facebook/jasco-chords-drums-melody-400M`: 400M model, text to music with chords, drums and melody support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-melody-400M)\n- `facebook/jasco-chords-drums-melody-1B`: 1B model, text to music with chords, drums and melody support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-melody-1B)\n\n\nSee after a quick example for using the API.\n\n```python\nfrom audiocraft.models import JASCO\n\nmodel = JASCO.get_pretrained('facebook/jasco-chords-drums-400M', chords_mapping_path='../assets/chord_to_index_mapping.pkl')\n\nmodel.set_generation_params(\n cfg_coef_all=5.0,\n cfg_coef_txt=0.0\n)\n\n# set textual prompt\ntext = \"Strings, woodwind, orchestral, symphony.\"\n\n# define chord progression\nchords = [('C', 0.0), ('D', 2.0), ('F', 4.0), ('Ab', 6.0), ('Bb', 7.0), ('C', 8.0)]\n\n# run inference\noutput = model.generate_music(descriptions=[text], chords=chords, progress=True)\n\naudio_write('output', output.cpu().squeeze(0), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\nFor more examples check out `demos/jasco_demo.ipynb`\n\n## 🤗 Transformers Usage\n\nComing soon...\n\n## Data Preprocessing\nIn order to to use the JascoDataset with chords / melody conditioning, please follow the instructions below:\n\n\n### Chords conditioning\nTo extract chords from your desired data follow the following steps:\n\n1. Prepare a `*.jsonl` containing list of absolute file paths in your dataset, should simply be absolute paths seperated by newlines.\n2. Download and install chord_extractor from [source](http://www.isophonics.net/nnls-chroma)\n3. For training purposes run: `python scripts/chords/extract_chords.py --src_jsonl_file= --target_output_dir=`\n
\nand then run: `python scripts/chords/build_chord_map.py --chords_folder= --output_directory=`\n\n4. For evaluation of our released models run: `python scripts/chords/extract_chords.py --src_jsonl_file= --target_output_dir= --path_to_pre_defined_map=`\n
\nand then run: `python scripts/chords/build_chord_map.py --chords_folder= --output_directory= --path_to_pre_defined_map=`\n\n\nNOTE: current scripts assume that all audio files are of `.wav` format, some changes may be required if your data consists of other formats.\n\nNOTE: predefined chord mapping file is available in `assets` directory.\n\n### Melody conditioning\n\nThis section relies on [Deepsalience repo](https://github.com/rabitt/ismir2017-deepsalience) with slight custom scripts written.\n\n#### Clone repo and create virtual environment\n1. `git clone git@github.com:lonzi/ismir2017-deepsalience.git forked_deepsalience_repo`\n2. `cd forked_deepsalience_repo`\n3. `conda create --name deep_salience python=3.7`\n4. `conda activate deep_salience`\n5. `pip install -r requirements.txt`\n\n\n#### Salience map dumps (of entire directory, using slurm job)\n\n##### From src dir\n\n1. create job array: `python predict/create_predict_saliency_cmds.py --src_dir= --out_dir= --n_shards= --multithread`\n2. run job array: `sbatch predict_saliency.sh`\n\n##### From track list\n\n1. create job array: `python predict/create_predict_saliency_cmds.py --tracks_list=tracks_train.txt --out_dir= --n_shards=2 --multithread --sbatch_script_name=predict_saliency_train.sh --saliency_threshold=`\n2. run job array: `sbatch predict_saliency_train.sh`\n\ntracks_train.txt: a list of track paths to process seperated by new lines\n\n\n## Training\n\nThe [JascoSolver](../audiocraft/solvers/jasco.py) implements JASCO's training pipeline.\nconditional flow matching objective over the continuous extracted latents from a pre-trained EnCodec model (see [EnCodec documentation](./ENCODEC.md)\nfor more details on how to train such model).\n\nNote that **we do NOT provide any of the datasets** used for training JASCO.\nWe provide a dummy dataset containing just a few examples for illustrative purposes.\n\nPlease read first the [TRAINING documentation](./TRAINING.md), in particular the Environment Setup section.\n\n\n### Fine tuning existing models\n\nYou can initialize your model to one of the pretrained models by using the `continue_from` argument, in particular\n\n```bash\n# Using pretrained JASCO model.\ndora run solver=jasco/chords_drums model/lm/model_scale=small continue_from=//pretrained/facebook/jasco-chords-drums-400M conditioner=jasco_chords_drums\n\n# Using another model you already trained with a Dora signature SIG.\ndora run solver=jasco/chords_drums model/lm/model_scale=small continue_from=//sig/SIG conditioner=jasco_chords_drums\n\n# Or providing manually a path\ndora run solver=jasco/chords_drums model/lm/model_scale=small conditioner=jasco_chords_drums continue_from=/checkpoints/my_other_xp/checkpoint.th\n```\n\n**Warning:** You are responsible for selecting the other parameters accordingly, in a way that make it compatible\n with the model you are fine tuning. Configuration is NOT automatically inherited from the model you continue from. In particular make sure to select the proper `conditioner` and `model/lm/model_scale`.\n\n**Warning:** We currently do not support fine tuning a model with slightly different layers. If you decide\n to change some parts, like the conditioning or some other parts of the model, you are responsible for manually crafting a checkpoint file from which we can safely run `load_state_dict`.\n If you decide to do so, make sure your checkpoint is saved with `torch.save` and contains a dict\n `{'best_state': {'model': model_state_dict_here}}`. Directly give the path to `continue_from` without a `//pretrained/` prefix.\n\n\n### Evaluation & Generation stage\n\nSee [MusicGen](./MUSICGEN.md)\n\n### Playing with the model\n\nOnce you have launched some experiments, you can easily get access\nto the Solver with the latest trained model using the following snippet.\n\n```python\nfrom audiocraft.solvers.jasco import JascoSolver\n\nsolver = JascoSolver.get_eval_solver_from_sig('SIG', device='cpu', batch_size=8)\nsolver.model\nsolver.dataloaders\n```\n\n### Importing / Exporting models\n\nWe do not support currently loading a model from the Hugging Face implementation or exporting to it.\nIf you want to export your model in a way that is compatible with `audiocraft.models.JASCO`\nAPI, you can run:\n\n```python\nfrom audiocraft.utils import export\nfrom audiocraft import train\nxp = train.main.get_xp_from_sig('SIG_OF_LM')\nexport.export_lm(xp.folder / 'checkpoint.th', '/checkpoints/my_audio_lm/state_dict.bin')\n# You also need to bundle the EnCodec model you used !!\n## Case 1) you trained your own\nxp_encodec = train.main.get_xp_from_sig('SIG_OF_ENCODEC')\nexport.export_encodec(xp_encodec.folder / 'checkpoint.th', '/checkpoints/my_audio_lm/compression_state_dict.bin')\n## Case 2) you used a pretrained model. Give the name you used without the //pretrained/ prefix.\n## This will actually not dump the actual model, simply a pointer to the right model to download.\nexport.export_pretrained_compression_model('facebook/encodec_32khz', '/checkpoints/my_audio_lm/compression_state_dict.bin')\n```\n\nNow you can load your custom model with:\n```python\nimport audiocraft.models\njasco = audiocraft.models.JASCO.get_pretrained('/checkpoints/my_audio_lm/')\n```\n\n\n### Learn more\n\nLearn more about AudioCraft training pipelines in the [dedicated section](./TRAINING.md).\n\n\n## Citation\n```\n@misc{tal2024joint,\n title={Joint Audio and Symbolic Conditioning for Temporally Controlled Text-to-Music Generation}, \n author={Or Tal and Alon Ziv and Itai Gat and Felix Kreuk and Yossi Adi},\n year={2024},\n eprint={2406.10970},\n archivePrefix={arXiv},\n primaryClass={cs.SD}\n}\n```\n\n## License\n\nSee license information in the [model card](../model_cards/JASCO_MODEL_CARD.md).\n\n[arxiv]: https://arxiv.org/pdf/2406.10970\n[sample_page]: https://pages.cs.huji.ac.il/adiyoss-lab/JASCO/\n" }, { "path": "docs/MAGNET.md", "content": "# MAGNeT: Masked Audio Generation using a Single Non-Autoregressive Transformer\n\nAudioCraft provides the code and models for MAGNeT, [Masked Audio Generation using a Single Non-Autoregressive Transformer][arxiv].\n\nMAGNeT is a text-to-music and text-to-sound model capable of generating high-quality audio samples conditioned on text descriptions.\nIt is a masked generative non-autoregressive Transformer trained over a 32kHz EnCodec tokenizer with 4 codebooks sampled at 50 Hz. \nUnlike prior work on masked generative audio Transformers, such as [SoundStorm](https://arxiv.org/abs/2305.09636) and [VampNet](https://arxiv.org/abs/2307.04686), \nMAGNeT doesn't require semantic token conditioning, model cascading or audio prompting, and employs a full text-to-audio using a single non-autoregressive Transformer.\n\nCheck out our [sample page][magnet_samples] or test the available demo!\n\nWe use 16K hours of licensed music to train MAGNeT. Specifically, we rely on an internal dataset\nof 10K high-quality music tracks, and on the ShutterStock and Pond5 music data.\n\n\n## Model Card\n\nSee [the model card](../model_cards/MAGNET_MODEL_CARD.md).\n\n\n## Installation\n\nPlease follow the AudioCraft installation instructions from the [README](../README.md).\n\nAudioCraft requires a GPU with at least 16 GB of memory for running inference with the medium-sized models (~1.5B parameters).\n\n## Usage\n\nWe currently offer two ways to interact with MAGNeT:\n1. You can use the gradio demo locally by running [`python -m demos.magnet_app --share`](../demos/magnet_app.py).\n2. You can play with MAGNeT by running the jupyter notebook at [`demos/magnet_demo.ipynb`](../demos/magnet_demo.ipynb) locally (if you have a GPU).\n\n## API\n\nWe provide a simple API and 6 pre-trained models. The pre trained models are:\n- `facebook/magnet-small-10secs`: 300M model, text to music, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/magnet-small-10secs)\n- `facebook/magnet-medium-10secs`: 1.5B model, text to music, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/magnet-medium-10secs)\n- `facebook/magnet-small-30secs`: 300M model, text to music, generates 30-second samples - [🤗 Hub](https://huggingface.co/facebook/magnet-small-30secs)\n- `facebook/magnet-medium-30secs`: 1.5B model, text to music, generates 30-second samples - [🤗 Hub](https://huggingface.co/facebook/magnet-medium-30secs)\n- `facebook/audio-magnet-small`: 300M model, text to sound-effect - [🤗 Hub](https://huggingface.co/facebook/audio-magnet-small)\n- `facebook/audio-magnet-medium`: 1.5B model, text to sound-effect - [🤗 Hub](https://huggingface.co/facebook/audio-magnet-medium)\n\nIn order to use MAGNeT locally **you must have a GPU**. We recommend 16GB of memory, especially for \nthe medium size models. \n\nSee after a quick example for using the API.\n\n```python\nimport torchaudio\nfrom audiocraft.models import MAGNeT\nfrom audiocraft.data.audio import audio_write\n\nmodel = MAGNeT.get_pretrained('facebook/magnet-small-10secs')\ndescriptions = ['disco beat', 'energetic EDM', 'funky groove']\nwav = model.generate(descriptions) # generates 3 samples.\n\nfor idx, one_wav in enumerate(wav):\n # Will save under {idx}.wav, with loudness normalization at -14 db LUFS.\n audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\n## 🤗 Transformers Usage\n\nComing soon...\n\n## Training\n\nThe [MagnetSolver](../audiocraft/solvers/magnet.py) implements MAGNeT's training pipeline.\nIt defines a masked generation task over multiple streams of discrete tokens\nextracted from a pre-trained EnCodec model (see [EnCodec documentation](./ENCODEC.md)\nfor more details on how to train such model).\n\nNote that **we do NOT provide any of the datasets** used for training MAGNeT.\nWe provide a dummy dataset containing just a few examples for illustrative purposes.\n\nPlease read first the [TRAINING documentation](./TRAINING.md), in particular the Environment Setup section.\n\n\n### Example configurations and grids\n\nWe provide configurations to reproduce the released models and our research.\nMAGNeT solvers configuration are available in [config/solver/magnet](../config/solver/magnet),\nin particular:\n* MAGNeT model for text-to-music:\n[`solver=magnet/magnet_32khz`](../config/solver/magnet/magnet_32khz.yaml)\n* MAGNeT model for text-to-sound:\n[`solver=magnet/audio_magnet_16khz`](../config/solver/magnet/audio_magnet_16khz.yaml)\n\nWe provide 3 different scales, e.g. `model/lm/model_scale=small` (300M), or `medium` (1.5B), and `large` (3.3B).\n\nPlease find some example grids to train MAGNeT at\n[audiocraft/grids/magnet](../audiocraft/grids/magnet/).\n\n```shell\n# text-to-music\ndora grid magnet.magnet_32khz --dry_run --init\n\n# text-to-sound\ndora grid magnet.audio_magnet_16khz --dry_run --init\n\n# Remove the `--dry_run --init` flags to actually schedule the jobs once everything is setup.\n```\n\n### dataset and metadata\nLearn more in the [datasets section](./DATASETS.md).\n\n#### Music Models\nMAGNeT's underlying dataset is an AudioDataset augmented with music-specific metadata.\nThe MAGNeT dataset implementation expects the metadata to be available as `.json` files\nat the same location as the audio files. \n\n#### Sound Models\nAudio-MAGNeT's underlying dataset is an AudioDataset augmented with description metadata.\nThe Audio-MAGNeT dataset implementation expects the metadata to be available as `.json` files\nat the same location as the audio files or through specified external folder.\n\n### Audio tokenizers\n\nSee [MusicGen](./MUSICGEN.md)\n\n### Fine tuning existing models\n\nYou can initialize your model to one of the pretrained models by using the `continue_from` argument, in particular\n\n```bash\n# Using pretrained MAGNeT model.\ndora run solver=magnet/magnet_32khz model/lm/model_scale=medium continue_from=//pretrained/facebook/magnet-medium-10secs conditioner=text2music\n\n# Using another model you already trained with a Dora signature SIG.\ndora run solver=magnet/magnet_32khz model/lm/model_scale=medium continue_from=//sig/SIG conditioner=text2music\n\n# Or providing manually a path\ndora run solver=magnet/magnet_32khz model/lm/model_scale=medium continue_from=/checkpoints/my_other_xp/checkpoint.th\n```\n\n**Warning:** You are responsible for selecting the other parameters accordingly, in a way that make it compatible\n with the model you are fine tuning. Configuration is NOT automatically inherited from the model you continue from. In particular make sure to select the proper `conditioner` and `model/lm/model_scale`.\n\n**Warning:** We currently do not support fine tuning a model with slightly different layers. If you decide\n to change some parts, like the conditioning or some other parts of the model, you are responsible for manually crafting a checkpoint file from which we can safely run `load_state_dict`.\n If you decide to do so, make sure your checkpoint is saved with `torch.save` and contains a dict\n `{'best_state': {'model': model_state_dict_here}}`. Directly give the path to `continue_from` without a `//pretrained/` prefix.\n\n### Evaluation stage\nFor the 6 pretrained MAGNeT models, objective metrics could be reproduced using the following grids:\n\n```shell\n# text-to-music\nREGEN=1 dora grid magnet.magnet_pretrained_32khz_eval --dry_run --init\n\n# text-to-sound\nREGEN=1 dora grid magnet.audio_magnet_pretrained_16khz_eval --dry_run --init\n\n# Remove the `--dry_run --init` flags to actually schedule the jobs once everything is setup.\n```\n\nSee [MusicGen](./MUSICGEN.md) for more details. \n\n### Generation stage\n\nSee [MusicGen](./MUSICGEN.md)\n\n### Playing with the model\n\nOnce you have launched some experiments, you can easily get access\nto the Solver with the latest trained model using the following snippet.\n\n```python\nfrom audiocraft.solvers.magnet import MagnetSolver\n\nsolver = MagnetSolver.get_eval_solver_from_sig('SIG', device='cpu', batch_size=8)\nsolver.model\nsolver.dataloaders\n```\n\n### Importing / Exporting models\n\nWe do not support currently loading a model from the Hugging Face implementation or exporting to it.\nIf you want to export your model in a way that is compatible with `audiocraft.models.MAGNeT`\nAPI, you can run:\n\n```python\nfrom audiocraft.utils import export\nfrom audiocraft import train\nxp = train.main.get_xp_from_sig('SIG_OF_LM')\nexport.export_lm(xp.folder / 'checkpoint.th', '/checkpoints/my_audio_lm/state_dict.bin')\n# You also need to bundle the EnCodec model you used !!\n## Case 1) you trained your own\nxp_encodec = train.main.get_xp_from_sig('SIG_OF_ENCODEC')\nexport.export_encodec(xp_encodec.folder / 'checkpoint.th', '/checkpoints/my_audio_lm/compression_state_dict.bin')\n## Case 2) you used a pretrained model. Give the name you used without the //pretrained/ prefix.\n## This will actually not dump the actual model, simply a pointer to the right model to download.\nexport.export_pretrained_compression_model('facebook/encodec_32khz', '/checkpoints/my_audio_lm/compression_state_dict.bin')\n```\n\nNow you can load your custom model with:\n```python\nimport audiocraft.models\nmagnet = audiocraft.models.MAGNeT.get_pretrained('/checkpoints/my_audio_lm/')\n```\n\n\n### Learn more\n\nLearn more about AudioCraft training pipelines in the [dedicated section](./TRAINING.md).\n\n## FAQ\n\n#### What are top-k, top-p, temperature and classifier-free guidance?\n\nCheck out [@FurkanGozukara tutorial](https://github.com/FurkanGozukara/Stable-Diffusion/blob/main/Tutorials/AI-Music-Generation-Audiocraft-Tutorial.md#more-info-about-top-k-top-p-temperature-and-classifier-free-guidance-from-chatgpt).\n\n#### Should I use FSDP or autocast ?\n\nThe two are mutually exclusive (because FSDP does autocast on its own).\nYou can use autocast up to 1.5B (medium), if you have enough RAM on your GPU.\nFSDP makes everything more complex but will free up some memory for the actual\nactivations by sharding the optimizer state.\n\n## Citation\n```\n@misc{ziv2024masked,\n title={Masked Audio Generation using a Single Non-Autoregressive Transformer}, \n author={Alon Ziv and Itai Gat and Gael Le Lan and Tal Remez and Felix Kreuk and Alexandre Défossez and Jade Copet and Gabriel Synnaeve and Yossi Adi},\n year={2024},\n eprint={2401.04577},\n archivePrefix={arXiv},\n primaryClass={cs.SD}\n}\n```\n\n## License\n\nSee license information in the [model card](../model_cards/MAGNET_MODEL_CARD.md).\n\n[arxiv]: https://arxiv.org/abs/2401.04577\n[magnet_samples]: https://pages.cs.huji.ac.il/adiyoss-lab/MAGNeT/\n" }, { "path": "docs/MBD.md", "content": "# MultiBand Diffusion\n\nAudioCraft provides the code and models for MultiBand Diffusion, [From Discrete Tokens to High Fidelity Audio using MultiBand Diffusion][arxiv].\nMultiBand diffusion is a collection of 4 models that can decode tokens from\nEnCodec tokenizer into waveform audio. You can listen to some examples on the sample page.\n\n\n \"Open\n\n
\n\n\n## Installation\n\nPlease follow the AudioCraft installation instructions from the [README](../README.md).\n\n\n## Usage\n\nWe offer a number of way to use MultiBand Diffusion:\n1. The MusicGen demo includes a toggle to try diffusion decoder. You can use the demo locally by running [`python -m demos.musicgen_app --share`](../demos/musicgen_app.py), or through the [MusicGen Colab](https://colab.research.google.com/drive/1JlTOjB-G0A2Hz3h8PK63vLZk4xdCI5QB?usp=sharing).\n2. You can play with MusicGen by running the jupyter notebook at [`demos/musicgen_demo.ipynb`](../demos/musicgen_demo.ipynb) locally (if you have a GPU).\n\n## API\n\nWe provide a simple API and pre-trained models for MusicGen and for EnCodec at 24 khz for 3 bitrates (1.5 kbps, 3 kbps and 6 kbps).\n\nSee after a quick example for using MultiBandDiffusion with the MusicGen API:\n\n```python\nimport torchaudio\nfrom audiocraft.models import MusicGen, MultiBandDiffusion\nfrom audiocraft.data.audio import audio_write\n\nmodel = MusicGen.get_pretrained('facebook/musicgen-melody')\nmbd = MultiBandDiffusion.get_mbd_musicgen()\nmodel.set_generation_params(duration=8) # generate 8 seconds.\nwav, tokens = model.generate_unconditional(4, return_tokens=True) # generates 4 unconditional audio samples and keep the tokens for MBD generation\ndescriptions = ['happy rock', 'energetic EDM', 'sad jazz']\nwav_diffusion = mbd.tokens_to_wav(tokens)\nwav, tokens = model.generate(descriptions, return_tokens=True) # generates 3 samples and keep the tokens.\nwav_diffusion = mbd.tokens_to_wav(tokens)\nmelody, sr = torchaudio.load('./assets/bach.mp3')\n# Generates using the melody from the given audio and the provided descriptions, returns audio and audio tokens.\nwav, tokens = model.generate_with_chroma(descriptions, melody[None].expand(3, -1, -1), sr, return_tokens=True)\nwav_diffusion = mbd.tokens_to_wav(tokens)\n\nfor idx, one_wav in enumerate(wav):\n # Will save under {idx}.wav and {idx}_diffusion.wav, with loudness normalization at -14 db LUFS for comparing the methods.\n audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n audio_write(f'{idx}_diffusion', wav_diffusion[idx].cpu(), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\nFor the compression task (and to compare with [EnCodec](https://github.com/facebookresearch/encodec)):\n\n```python\nimport torch\nfrom audiocraft.models import MultiBandDiffusion\nfrom encodec import EncodecModel\nfrom audiocraft.data.audio import audio_read, audio_write\n\nbandwidth = 3.0 # 1.5, 3.0, 6.0\nmbd = MultiBandDiffusion.get_mbd_24khz(bw=bandwidth)\nencodec = EncodecModel.encodec_model_24khz()\n\nsomepath = ''\nwav, sr = audio_read(somepath)\nwith torch.no_grad():\n compressed_encodec = encodec(wav)\n compressed_diffusion = mbd.regenerate(wav, sample_rate=sr)\n\naudio_write('sample_encodec', compressed_encodec.squeeze(0).cpu(), mbd.sample_rate, strategy=\"loudness\", loudness_compressor=True)\naudio_write('sample_diffusion', compressed_diffusion.squeeze(0).cpu(), mbd.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\n\n## Training\n\nThe [DiffusionSolver](../audiocraft/solvers/diffusion.py) implements our diffusion training pipeline.\nIt generates waveform audio conditioned on the embeddings extracted from a pre-trained EnCodec model\n(see [EnCodec documentation](./ENCODEC.md) for more details on how to train such model).\n\nNote that **we do NOT provide any of the datasets** used for training our diffusion models.\nWe provide a dummy dataset containing just a few examples for illustrative purposes.\n\n### Example configurations and grids\n\nOne can train diffusion models as described in the paper by using this [dora grid](../audiocraft/grids/diffusion/4_bands_base_32khz.py).\n```shell\n# 4 bands MBD trainning\ndora grid diffusion.4_bands_base_32khz\n```\n\n### Learn more\n\nLearn more about AudioCraft training pipelines in the [dedicated section](./TRAINING.md).\n\n\n## Citation\n\n```\n@article{sanroman2023fromdi,\n title={From Discrete Tokens to High-Fidelity Audio Using Multi-Band Diffusion},\n author={San Roman, Robin and Adi, Yossi and Deleforge, Antoine and Serizel, Romain and Synnaeve, Gabriel and Défossez, Alexandre},\n journal={arXiv preprint arXiv:},\n year={2023}\n}\n```\n\n\n## License\n\nSee license information in the [README](../README.md).\n\n\n[arxiv]: https://arxiv.org/abs/2308.02560\n[mbd_samples]: https://ai.honu.io/papers/mbd/\n" }, { "path": "docs/METRICS.md", "content": "# AudioCraft objective metrics\n\nIn addition to training losses, AudioCraft provides a set of objective metrics\nfor audio synthesis and audio generation. As these metrics may require\nextra dependencies and can be costly to train, they are often disabled by default.\nThis section provides guidance for setting up and using these metrics in\nthe AudioCraft training pipelines.\n\n## Available metrics\n\n### Audio synthesis quality metrics\n\n#### SI-SNR\n\nWe provide an implementation of the Scale-Invariant Signal-to-Noise Ratio in PyTorch.\nNo specific requirement is needed for this metric. Please activate the metric at the\nevaluation stage with the appropriate flag:\n\n**Warning:** We report the opposite of the SI-SNR, e.g. multiplied by -1. This is due to internal \n details where the SI-SNR score can also be used as a training loss function, where lower\n values should indicate better reconstruction. Negative values are such expected and a good sign! Those should be again multiplied by `-1` before publication :)\n\n```shell\ndora run <...> evaluate.metrics.sisnr=true\n```\n\n#### ViSQOL\n\nWe provide a Python wrapper around the ViSQOL [official implementation](https://github.com/google/visqol)\nto conveniently run ViSQOL within the training pipelines.\n\nOne must specify the path to the ViSQOL installation through the configuration in order\nto enable ViSQOL computations in AudioCraft:\n\n```shell\n# the first parameter is used to activate visqol computation while the second specify\n# the path to visqol's library to be used by our python wrapper\ndora run <...> evaluate.metrics.visqol=true metrics.visqol.bin=\n```\n\nSee an example grid: [Compression with ViSQOL](../audiocraft/grids/compression/encodec_musicgen_32khz.py)\n\nTo learn more about ViSQOL and how to build ViSQOL binary using bazel, please refer to the\ninstructions available in the [open source repository](https://github.com/google/visqol).\n\n### Audio generation metrics\n\n#### Frechet Audio Distance\n\nSimilarly to ViSQOL, we use a Python wrapper around the Frechet Audio Distance\n[official implementation](https://github.com/google-research/google-research/tree/master/frechet_audio_distance)\nin TensorFlow.\n\nNote that we had to make several changes to the actual code in order to make it work.\nPlease refer to the [FrechetAudioDistanceMetric](../audiocraft/metrics/fad.py) class documentation\nfor more details. We do not plan to provide further support in obtaining a working setup for the\nFrechet Audio Distance at this stage.\n\n```shell\n# the first parameter is used to activate FAD metric computation while the second specify\n# the path to FAD library to be used by our python wrapper\ndora run <...> evaluate.metrics.fad=true metrics.fad.bin=\n```\n\nSee an example grid: [Evaluation with FAD](../audiocraft/grids/musicgen/musicgen_pretrained_32khz_eval.py)\n\n#### Kullback-Leibler Divergence\n\nWe provide a PyTorch implementation of the Kullback-Leibler Divergence computed over the probabilities\nof the labels obtained by a state-of-the-art audio classifier. We provide our implementation of the KLD\nusing the [PaSST classifier](https://github.com/kkoutini/PaSST).\n\nIn order to use the KLD metric over PaSST, you must install the PaSST library as an extra dependency:\n```shell\npip install 'git+https://github.com/kkoutini/passt_hear21@0.0.19#egg=hear21passt'\n```\n\nThen similarly, you can use the metric activating the corresponding flag:\n\n```shell\n# one could extend the kld metric with additional audio classifier models that can then be picked through the configuration\ndora run <...> evaluate.metrics.kld=true metrics.kld.model=passt\n```\n\n#### Text consistency\n\nWe provide a text-consistency metric, similarly to the MuLan Cycle Consistency from\n[MusicLM](https://arxiv.org/pdf/2301.11325.pdf) or the CLAP score used in\n[Make-An-Audio](https://arxiv.org/pdf/2301.12661v1.pdf).\nMore specifically, we provide a PyTorch implementation of a Text consistency metric\nrelying on a pre-trained [Contrastive Language-Audio Pretraining (CLAP)](https://github.com/LAION-AI/CLAP).\n\nPlease install the CLAP library as an extra dependency prior to using the metric:\n```shell\npip install laion_clap\n```\n\nThen similarly, you can use the metric activating the corresponding flag:\n\n```shell\n# one could extend the text consistency metric with additional audio classifier models that can then be picked through the configuration\ndora run ... evaluate.metrics.text_consistency=true metrics.text_consistency.model=clap\n```\n\nNote that the text consistency metric based on CLAP will require the CLAP checkpoint to be\nprovided in the configuration.\n\n#### Chroma cosine similarity\n\nFinally, as introduced in MusicGen, we provide a Chroma Cosine Similarity metric in PyTorch.\nNo specific requirement is needed for this metric. Please activate the metric at the\nevaluation stage with the appropriate flag:\n\n```shell\ndora run ... evaluate.metrics.chroma_cosine=true\n```\n\n#### Comparing against reconstructed audio\n\nFor all the above audio generation metrics, we offer the option to compute the metric on the reconstructed audio\nfed in EnCodec instead of the generated sample using the flag `.use_gt=true`.\n\n## Example usage\n\nYou will find example of configuration for the different metrics introduced above in:\n* The [musicgen's default solver](../config/solver/musicgen/default.yaml) for all audio generation metrics\n* The [compression's default solver](../config/solver/compression/default.yaml) for all audio synthesis metrics\n\nSimilarly, we provide different examples in our grids:\n* [Evaluation with ViSQOL](../audiocraft/grids/compression/encodec_musicgen_32khz.py)\n* [Evaluation with FAD and others](../audiocraft/grids/musicgen/musicgen_pretrained_32khz_eval.py)\n" }, { "path": "docs/MUSICGEN.md", "content": "# MusicGen: Simple and Controllable Music Generation\n\nAudioCraft provides the code and models for MusicGen, [a simple and controllable model for music generation][arxiv].\nMusicGen is a single stage auto-regressive Transformer model trained over a 32kHz\nEnCodec tokenizer with 4 codebooks sampled at 50 Hz.\nUnlike existing methods like [MusicLM](https://arxiv.org/abs/2301.11325), MusicGen doesn't require\na self-supervised semantic representation, and it generates all 4 codebooks in one pass. By introducing\na small delay between the codebooks, we show we can predict them in parallel, thus having only 50 auto-regressive\nsteps per second of audio.\nCheck out our [sample page][musicgen_samples] or test the available demo!\n\n\n \"Open\n\n\n \"Open\n\n
\n\nWe use 20K hours of licensed music to train MusicGen. Specifically, we rely on an internal dataset\nof 10K high-quality music tracks, and on the ShutterStock and Pond5 music data.\n\n\n## Model Card\n\nSee [the model card](../model_cards/MUSICGEN_MODEL_CARD.md).\n\n\n## Installation\n\nPlease follow the AudioCraft installation instructions from the [README](../README.md).\n\nAudioCraft requires a GPU with at least 16 GB of memory for running inference with the medium-sized models (~1.5B parameters).\n\n## Usage\n\nWe offer a number of way to interact with MusicGen:\n1. A demo is also available on the [`facebook/MusicGen` Hugging Face Space](https://huggingface.co/spaces/facebook/MusicGen)\n(huge thanks to all the HF team for their support).\n2. You can run the extended demo on a Colab:\n[colab notebook](https://ai.honu.io/red/musicgen-colab)\n3. You can use the gradio demo locally by running [`python -m demos.musicgen_app --share`](../demos/musicgen_app.py).\n4. You can play with MusicGen by running the jupyter notebook at [`demos/musicgen_demo.ipynb`](../demos/musicgen_demo.ipynb) locally (if you have a GPU).\n5. Finally, checkout [@camenduru Colab page](https://github.com/camenduru/MusicGen-colab)\nwhich is regularly updated with contributions from @camenduru and the community.\n\n\n## API\n\nWe provide a simple API and 10 pre-trained models. The pre trained models are:\n- `facebook/musicgen-small`: 300M model, text to music only - [🤗 Hub](https://huggingface.co/facebook/musicgen-small)\n- `facebook/musicgen-medium`: 1.5B model, text to music only - [🤗 Hub](https://huggingface.co/facebook/musicgen-medium)\n- `facebook/musicgen-melody`: 1.5B model, text to music and text+melody to music - [🤗 Hub](https://huggingface.co/facebook/musicgen-melody)\n- `facebook/musicgen-large`: 3.3B model, text to music only - [🤗 Hub](https://huggingface.co/facebook/musicgen-large)\n- `facebook/musicgen-melody-large`: 3.3B model, text to music and text+melody to music - [🤗 Hub](https://huggingface.co/facebook/musicgen-melody-large)\n- `facebook/musicgen-stereo-*`: All the previous models fine tuned for stereo generation -\n [small](https://huggingface.co/facebook/musicgen-stereo-small),\n [medium](https://huggingface.co/facebook/musicgen-stereo-medium),\n [large](https://huggingface.co/facebook/musicgen-stereo-large),\n [melody](https://huggingface.co/facebook/musicgen-stereo-melody),\n [melody large](https://huggingface.co/facebook/musicgen-stereo-melody-large).\n\nWe observe the best trade-off between quality and compute with the `facebook/musicgen-medium` or `facebook/musicgen-melody` model.\nIn order to use MusicGen locally **you must have a GPU**. We recommend 16GB of memory, but smaller\nGPUs will be able to generate short sequences, or longer sequences with the `facebook/musicgen-small` model.\n\nSee after a quick example for using the API.\n\n```python\nimport torchaudio\nfrom audiocraft.models import MusicGen\nfrom audiocraft.data.audio import audio_write\n\nmodel = MusicGen.get_pretrained('facebook/musicgen-melody')\nmodel.set_generation_params(duration=8) # generate 8 seconds.\nwav = model.generate_unconditional(4) # generates 4 unconditional audio samples\ndescriptions = ['happy rock', 'energetic EDM', 'sad jazz']\nwav = model.generate(descriptions) # generates 3 samples.\n\nmelody, sr = torchaudio.load('./assets/bach.mp3')\n# generates using the melody from the given audio and the provided descriptions.\nwav = model.generate_with_chroma(descriptions, melody[None].expand(3, -1, -1), sr)\n\nfor idx, one_wav in enumerate(wav):\n # Will save under {idx}.wav, with loudness normalization at -14 db LUFS.\n audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\n## 🤗 Transformers Usage\n\nMusicGen is available in the 🤗 Transformers library from version 4.31.0 onwards, requiring minimal dependencies\nand additional packages. Steps to get started:\n\n1. First install the 🤗 [Transformers library](https://github.com/huggingface/transformers) from main:\n\n```shell\npip install git+https://github.com/huggingface/transformers.git\n```\n\n2. Run the following Python code to generate text-conditional audio samples:\n\n```py\nfrom transformers import AutoProcessor, MusicgenForConditionalGeneration\n\n\nprocessor = AutoProcessor.from_pretrained(\"facebook/musicgen-small\")\nmodel = MusicgenForConditionalGeneration.from_pretrained(\"facebook/musicgen-small\")\n\ninputs = processor(\n text=[\"80s pop track with bassy drums and synth\", \"90s rock song with loud guitars and heavy drums\"],\n padding=True,\n return_tensors=\"pt\",\n)\n\naudio_values = model.generate(**inputs, max_new_tokens=256)\n```\n\n3. Listen to the audio samples either in an ipynb notebook:\n\n```py\nfrom IPython.display import Audio\n\nsampling_rate = model.config.audio_encoder.sampling_rate\nAudio(audio_values[0].numpy(), rate=sampling_rate)\n```\n\nOr save them as a `.wav` file using a third-party library, e.g. `scipy`:\n\n```py\nimport scipy\n\nsampling_rate = model.config.audio_encoder.sampling_rate\nscipy.io.wavfile.write(\"musicgen_out.wav\", rate=sampling_rate, data=audio_values[0, 0].numpy())\n```\n\nFor more details on using the MusicGen model for inference using the 🤗 Transformers library, refer to the\n[MusicGen docs](https://huggingface.co/docs/transformers/main/en/model_doc/musicgen) or the hands-on\n[Google Colab](https://colab.research.google.com/github/sanchit-gandhi/notebooks/blob/main/MusicGen.ipynb).\n\n\n## Training\n\nThe [MusicGenSolver](../audiocraft/solvers/musicgen.py) implements MusicGen's training pipeline.\nIt defines an autoregressive language modeling task over multiple streams of discrete tokens\nextracted from a pre-trained EnCodec model (see [EnCodec documentation](./ENCODEC.md)\nfor more details on how to train such model).\n\nNote that **we do NOT provide any of the datasets** used for training MusicGen.\nWe provide a dummy dataset containing just a few examples for illustrative purposes.\n\nPlease read first the [TRAINING documentation](./TRAINING.md), in particular the Environment Setup section.\n\n\n**Warning:** As of version 1.1.0, a few breaking changes were introduced. Check the [CHANGELOG.md](../CHANGELOG.md)\nfile for more information. You might need to retrain some of your models.\n\n### Example configurations and grids\n\nWe provide configurations to reproduce the released models and our research.\nMusicGen solvers configuration are available in [config/solver/musicgen](../config/solver/musicgen),\nin particular:\n* MusicGen base model for text-to-music:\n[`solver=musicgen/musicgen_base_32khz`](../config/solver/musicgen/musicgen_base_32khz.yaml)\n* MusicGen model with chromagram-conditioning support:\n[`solver=musicgen/musicgen_melody_32khz`](../config/solver/musicgen/musicgen_melody_32khz.yaml)\n\nWe provide 3 different scales, e.g. `model/lm/model_scale=small` (300M), or `medium` (1.5B), and `large` (3.3B).\n\nPlease find some example grids to train MusicGen at\n[audiocraft/grids/musicgen](../audiocraft/grids/musicgen/).\n\n```shell\n# text-to-music\ndora grid musicgen.musicgen_base_32khz --dry_run --init\n# melody-guided music generation\ndora grid musicgen.musicgen_melody_base_32khz --dry_run --init\n# Remove the `--dry_run --init` flags to actually schedule the jobs once everything is setup.\n```\n\n### Music dataset and metadata\n\nMusicGen's underlying dataset is an AudioDataset augmented with music-specific metadata.\nThe MusicGen dataset implementation expects the metadata to be available as `.json` files\nat the same location as the audio files. Learn more in the [datasets section](./DATASETS.md).\n\n\n### Audio tokenizers\n\nWe support a number of audio tokenizers: either pretrained EnCodec models, [DAC](https://github.com/descriptinc/descript-audio-codec), or your own models.\nThe tokenizer is controlled with the setting `compression_model_checkpoint`.\nFor instance,\n\n```bash\n# Using the 32kHz EnCodec trained on music\ndora run solver=musicgen/debug \\\n compression_model_checkpoint=//pretrained/facebook/encodec_32khz \\\n transformer_lm.n_q=4 transformer_lm.card=2048\n\n# Using DAC\ndora run solver=musicgen/debug \\\n compression_model_checkpoint=//pretrained/dac_44khz \\\n transformer_lm.n_q=9 transformer_lm.card=1024 \\\n 'codebooks_pattern.delay.delays=[0,1,2,3,4,5,6,7,8]'\n\n# Using your own model after export (see ENCODEC.md)\ndora run solver=musicgen/debug \\\n compression_model_checkpoint=//pretrained//checkpoints/my_audio_lm/compression_state_dict.bin \\\n transformer_lm.n_q=... transformer_lm.card=...\n\n# Using your own model from its training checkpoint.\ndora run solver=musicgen/debug \\\n compression_model_checkpoint=//sig/SIG \\ # where SIG is the Dora signature of the EnCodec XP.\n transformer_lm.n_q=... transformer_lm.card=...\n```\n\n**Warning:** you are responsible for setting the proper value for `transformer_lm.n_q` and `transformer_lm.card` (cardinality of the codebooks). You also have to update the codebook_pattern to match `n_q` as shown in the example for using DAC. .\n\n\n### Training stereo models\n\nUse the option `interleave_stereo_codebooks.use` set to `True` to activate stereo training along with `channels=2`. Left and right channels will be\nencoded separately by the compression model, then their codebook will be interleaved, e.g. order of codebook is\n`[1_L, 1_R, 2_L, 2_R, ...]`. You will also need to update the delays for the codebook patterns to match the number of codebooks, and the `n_q` value passed to the transformer LM:\n```\ndora run solver=musicgen/debug \\\n compression_model_checkpoint=//pretrained/facebook/encodec_32khz \\\n channels=2 interleave_stereo_codebooks.use=True \\\n transformer_lm.n_q=8 transformer_lm.card=2048 \\\n codebooks_pattern.delay.delays='[0, 0, 1, 1, 2, 2, 3, 3]'\n```\n\n### Fine tuning existing models\n\nYou can initialize your model to one of the pretrained models by using the `continue_from` argument, in particular\n\n```bash\n# Using pretrained MusicGen model.\ndora run solver=musicgen/musicgen_base_32khz model/lm/model_scale=medium continue_from=//pretrained/facebook/musicgen-medium conditioner=text2music\n\n# Using another model you already trained with a Dora signature SIG.\ndora run solver=musicgen/musicgen_base_32khz model/lm/model_scale=medium continue_from=//sig/SIG conditioner=text2music\n\n# Or providing manually a path\ndora run solver=musicgen/musicgen_base_32khz model/lm/model_scale=medium continue_from=/checkpoints/my_other_xp/checkpoint.th\n```\n\n**Warning:** You are responsible for selecting the other parameters accordingly, in a way that make it compatible\n with the model you are fine tuning. Configuration is NOT automatically inherited from the model you continue from. In particular make sure to select the proper `conditioner` and `model/lm/model_scale`.\n\n**Warning:** We currently do not support fine tuning a model with slightly different layers. If you decide\n to change some parts, like the conditioning or some other parts of the model, you are responsible for manually crafting a checkpoint file from which we can safely run `load_state_dict`.\n If you decide to do so, make sure your checkpoint is saved with `torch.save` and contains a dict\n `{'best_state': {'model': model_state_dict_here}}`. Directly give the path to `continue_from` without a `//pretrained/` prefix.\n\n\n#### Fine tuning mono model to stereo\n\nYou will not be able to `continue_from` a mono model with stereo training, as the shape of the embeddings and output linears\nwould not match. You can use the following snippet to prepare a proper finetuning checkpoint.\n\n```python\nfrom pathlib import Path\nimport torch\n\n# Download the pretrained model, e.g. from\n# https://huggingface.co/facebook/musicgen-melody/blob/main/state_dict.bin\n\nmodel_name = 'musicgen-melody'\nroot = Path.home() / 'checkpoints'\n# You are responsible for downloading the following checkpoint in the proper location\ninput_state_dict_path = root / model_name / 'state_dict.bin'\nstate = torch.load(input_state_dict_path, 'cpu')\nbs = state['best_state']\n# there is a slight different in format between training checkpoints and exported public checkpoints.\n# If you want to use your own mono models from one of your training checkpont, following the instructions\n# for exporting a model explained later on this page.\nassert 'model' not in bs, 'The following code is for using an exported pretrained model'\nnbs = dict(bs)\nfor k in range(8):\n # We will just copy mono embeddings and linears twice, once for left and right channels.\n nbs[f'linears.{k}.weight'] = bs[f'linears.{k//2}.weight']\n nbs[f'emb.{k}.weight'] = bs[f'emb.{k//2}.weight']\ntorch.save({'best_state': {'model': nbs}}, root / f'stereo_finetune_{model_name}.th')\n```\n\nNow, you can use `$HOME/checkpoints/stereo_finetune_musicgen-melody.th` as a `continue_from` target (without a `//pretrained` prefix!).\n\n### Caching of EnCodec tokens\n\nIt is possible to precompute the EnCodec tokens and other metadata.\nAn example of generating and using this cache provided in the [musicgen.musicgen_base_cached_32khz grid](../audiocraft/grids/musicgen/musicgen_base_cached_32khz.py).\n\n### Evaluation stage\n\nBy default, evaluation stage is also computing the cross-entropy and the perplexity over the\nevaluation dataset. Indeed the objective metrics used for evaluation can be costly to run\nor require some extra dependencies. Please refer to the [metrics documentation](./METRICS.md)\nfor more details on the requirements for each metric.\n\nWe provide an off-the-shelf configuration to enable running the objective metrics\nfor audio generation in\n[config/solver/musicgen/evaluation/objective_eval](../config/solver/musicgen/evaluation/objective_eval.yaml).\n\nOne can then activate evaluation the following way:\n```shell\n# using the configuration\ndora run solver=musicgen/debug solver/musicgen/evaluation=objective_eval\n# specifying each of the fields, e.g. to activate KL computation\ndora run solver=musicgen/debug evaluate.metrics.kld=true\n```\n\nSee [an example evaluation grid](../audiocraft/grids/musicgen/musicgen_pretrained_32khz_eval.py).\n\n### Generation stage\n\nThe generation stage allows to generate samples conditionally and/or unconditionally and to perform\naudio continuation (from a prompt). We currently support greedy sampling (argmax), sampling\nfrom softmax with a given temperature, top-K and top-P (nucleus) sampling. The number of samples\ngenerated and the batch size used are controlled by the `dataset.generate` configuration\nwhile the other generation parameters are defined in `generate.lm`.\n\n```shell\n# control sampling parameters\ndora run solver=musicgen/debug generate.lm.gen_duration=10 generate.lm.use_sampling=true generate.lm.top_k=15\n```\n\n#### Listening to samples\n\nNote that generation happens automatically every 25 epochs. You can easily access and\ncompare samples between models (as long as they are trained) on the same dataset using the\nMOS tool. For that first `pip install Flask gunicorn`. Then\n```\ngunicorn -w 4 -b 127.0.0.1:8895 -t 120 'scripts.mos:app' --access-logfile -\n```\nAnd access the tool at [https://127.0.0.1:8895](https://127.0.0.1:8895).\n\n### Playing with the model\n\nOnce you have launched some experiments, you can easily get access\nto the Solver with the latest trained model using the following snippet.\n\n```python\nfrom audiocraft.solvers.musicgen import MusicGenSolver\n\nsolver = MusicGenSolver.get_eval_solver_from_sig('SIG', device='cpu', batch_size=8)\nsolver.model\nsolver.dataloaders\n```\n\n### Importing / Exporting models\n\nWe do not support currently loading a model from the Hugging Face implementation or exporting to it.\nIf you want to export your model in a way that is compatible with `audiocraft.models.MusicGen`\nAPI, you can run:\n\n```python\nfrom audiocraft.utils import export\nfrom audiocraft import train\nxp = train.main.get_xp_from_sig('SIG_OF_LM')\nexport.export_lm(xp.folder / 'checkpoint.th', '/checkpoints/my_audio_lm/state_dict.bin')\n# You also need to bundle the EnCodec model you used !!\n## Case 1) you trained your own\nxp_encodec = train.main.get_xp_from_sig('SIG_OF_ENCODEC')\nexport.export_encodec(xp_encodec.folder / 'checkpoint.th', '/checkpoints/my_audio_lm/compression_state_dict.bin')\n## Case 2) you used a pretrained model. Give the name you used without the //pretrained/ prefix.\n## This will actually not dump the actual model, simply a pointer to the right model to download.\nexport.export_pretrained_compression_model('facebook/encodec_32khz', '/checkpoints/my_audio_lm/compression_state_dict.bin')\n```\n\nNow you can load your custom model with:\n```python\nimport audiocraft.models\nmusicgen = audiocraft.models.MusicGen.get_pretrained('/checkpoints/my_audio_lm/')\n```\n\n\n### Learn more\n\nLearn more about AudioCraft training pipelines in the [dedicated section](./TRAINING.md).\n\n## FAQ\n\n#### I need help on Windows\n\n@FurkanGozukara made a complete tutorial for [AudioCraft/MusicGen on Windows](https://youtu.be/v-YpvPkhdO4)\n\n#### I need help for running the demo on Colab\n\nCheck [@camenduru tutorial on YouTube](https://www.youtube.com/watch?v=EGfxuTy9Eeo).\n\n#### What are top-k, top-p, temperature and classifier-free guidance?\n\nCheck out [@FurkanGozukara tutorial](https://github.com/FurkanGozukara/Stable-Diffusion/blob/main/Tutorials/AI-Music-Generation-Audiocraft-Tutorial.md#more-info-about-top-k-top-p-temperature-and-classifier-free-guidance-from-chatgpt).\n\n#### Should I use FSDP or autocast ?\n\nThe two are mutually exclusive (because FSDP does autocast on its own).\nYou can use autocast up to 1.5B (medium), if you have enough RAM on your GPU.\nFSDP makes everything more complex but will free up some memory for the actual\nactivations by sharding the optimizer state.\n\n## Citation\n```\n@inproceedings{copet2023simple,\n title={Simple and Controllable Music Generation},\n author={Jade Copet and Felix Kreuk and Itai Gat and Tal Remez and David Kant and Gabriel Synnaeve and Yossi Adi and Alexandre Défossez},\n booktitle={Thirty-seventh Conference on Neural Information Processing Systems},\n year={2023},\n}\n```\n\n\n## License\n\nSee license information in the [model card](../model_cards/MUSICGEN_MODEL_CARD.md).\n\n\n[arxiv]: https://arxiv.org/abs/2306.05284\n[musicgen_samples]: https://ai.honu.io/papers/musicgen/\n" }, { "path": "docs/MUSICGEN_STYLE.md", "content": "# MusicGen-Style: Audio Conditioning for Music Generation via Discrete Bottleneck Features\n\nAudioCraft provides the code and models for MusicGen-Style, [Audio Conditioning for Music Generation via Discrete Bottleneck Features][arxiv].\n\nMusicGen-Style is a text-and-audio-to-music model that can be conditioned on textual and audio data (thanks to a style conditioner). \nThe style conditioner takes as input a music excerpt of a few seconds (between 1.5 and 4.5) extracts some features that are used by the model to generate music in the same style. \nThis style conditioning can be mixed with textual description. \n\nCheck out our [sample page][musicgen_style_samples] or test the available demo!\n\nWe use 16K hours of licensed music to train MusicGen-Style. Specifically, we rely on an internal dataset\nof 10K high-quality music tracks, and on the ShutterStock and Pond5 music data.\n\n\n## Model Card\n\nSee [the model card](../model_cards/MUSICGEN_STYLE_MODEL_CARD.md).\n\n\n## Installation\n\nPlease follow the AudioCraft installation instructions from the [README](../README.md).\n\nMusicGen-Stem requires a GPU with at least 16 GB of memory for running inference with the medium-sized models (~1.5B parameters).\n\n## Usage\n\n1. You can play with MusicGen-Style by running the jupyter notebook at [`demos/musicgen_style_demo.ipynb`](../demos/musicgen_style_demo.ipynb) locally (if you have a GPU).\n2. You can use the gradio demo locally by running python -m demos.musicgen_style_app --share.\n3. You can play with MusicGen by running the jupyter notebook at demos/musicgen_style_demo.ipynb locally (if you have a GPU).\n\n## API\n\nWe provide a simple API 1 pre-trained model with MERT used as a feature extractor for the style conditioner:\n- `facebook/musicgen-style`: medium (1.5B) MusicGen model, text and style to music, generates 30-second samples - [🤗 Hub](https://huggingface.co/facebook/musicgen-style)\n\nIn order to use MusicGen-Style locally **you must have a GPU**. We recommend 16GB of memory. \n\nSee after a quick example for using the API.\n\nTo perform text-to-music:\n```python\nimport torchaudio\nfrom audiocraft.models import MusicGen\nfrom audiocraft.data.audio import audio_write\n\nmodel = MusicGen.get_pretrained('facebook/musicgen-style')\n\n\nmodel.set_generation_params(\n duration=8, # generate 8 seconds, can go up to 30\n use_sampling=True, \n top_k=250,\n cfg_coef=3., # Classifier Free Guidance coefficient \n cfg_coef_beta=None, # double CFG is only useful for text-and-style conditioning\n) \n\ndescriptions = ['disco beat', 'energetic EDM', 'funky groove']\nwav = model.generate(descriptions) # generates 3 samples.\n\nfor idx, one_wav in enumerate(wav):\n # Will save under {idx}.wav, with loudness normalization at -14 db LUFS.\n audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\nTo perform style-to-music:\n```python\nimport torchaudio\nfrom audiocraft.models import MusicGen\nfrom audiocraft.data.audio import audio_write\n\nmodel = MusicGen.get_pretrained('facebook/musicgen-style')\n\n\nmodel.set_generation_params(\n duration=8, # generate 8 seconds, can go up to 30\n use_sampling=True, \n top_k=250,\n cfg_coef=3., # Classifier Free Guidance coefficient \n cfg_coef_beta=None, # double CFG is only useful for text-and-style conditioning\n)\n\nmodel.set_style_conditioner_params(\n eval_q=1, # integer between 1 and 6\n # eval_q is the level of quantization that passes\n # through the conditioner. When low, the models adheres less to the \n # audio conditioning\n excerpt_length=3., # the length in seconds that is taken by the model in the provided excerpt\n )\n\nmelody, sr = torchaudio.load('./assets/electronic.mp3')\n\n\nwav = model.generate_with_chroma(descriptions=[None, None, None], \n melody[None].expand(3, -1, -1), sr) # generates 3 samples.\n\nfor idx, one_wav in enumerate(wav):\n # Will save under {idx}.wav, with loudness normalization at -14 db LUFS.\n audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\nTo perform style-and-text-to-music:\n```python\nimport torchaudio\nfrom audiocraft.models import MusicGen\nfrom audiocraft.data.audio import audio_write\n\nmodel = MusicGen.get_pretrained('facebook/musicgen-style')\n\n\nmodel.set_generation_params(\n duration=8, # generate 8 seconds, can go up to 30\n use_sampling=True, \n top_k=250,\n cfg_coef=3., # Classifier Free Guidance coefficient \n cfg_coef_beta=5., # double CFG is necessary for text-and-style conditioning\n # Beta in the double CFG formula. between 1 and 9. When set to 1 it is equivalent to normal CFG. \n # When we increase this parameter, the text condition is pushed. See the bottom of https://musicgenstyle.github.io/ \n # to better understand the effects of the double CFG coefficients. \n)\n\nmodel.set_style_conditioner_params(\n eval_q=1, # integer between 1 and 6\n # eval_q is the level of quantization that passes\n # through the conditioner. When low, the models adheres less to the \n # audio conditioning\n excerpt_length=3., # the length in seconds that is taken by the model in the provided excerpt, can be \n # between 1.5 and 4.5 seconds but it has to be shortest to the length of the provided conditioning\n )\n\nmelody, sr = torchaudio.load('./assets/electronic.mp3')\n\ndescriptions = [\"8-bit old video game music\", \"Chill lofi remix\", \"80s New wave with synthesizer\"]\nwav = model.generate_with_chroma(descriptions=[\"8-bit old video game music\"], \n melody[None].expand(3, -1, -1), sr) # generates 3 samples.\n\nfor idx, one_wav in enumerate(wav):\n # Will save under {idx}.wav, with loudness normalization at -14 db LUFS.\n audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\n\n## Training\nTo train MusicGen-Style, we use the [MusicGenSolver](../audiocraft/solvers/musicgen.py).\n\nNote that **we do NOT provide any of the datasets** used for training MusicGen-Style.\nWe provide a dummy dataset containing just a few examples for illustrative purposes.\n\nPlease read first the [TRAINING documentation](./TRAINING.md), in particular the Environment Setup section.\n\n\n### Example configurations and grids\n\nWe provide the configuration to reproduce the training of MusicGen-Style in [config/solver/musicgen/musicgen_style_32khz.yaml](../config/solver/musicgen/musicgen_style_32khz.yaml),\n\nIn particular, the conditioner configuration is provided in [/config/conditioner/style2music.yaml](../config/conditioner/style2music.yaml).\n\nThe grid to train the model is \n[audiocraft/grids/musicgen/musicgen_style_32khz.py](../audiocraft/grids/musicgen/musicgen_style_32khz.py).\n\n```shell\n# text-and-style-to-music\ndora grid musicgen.musicgen_style_32khz --dry_run --init\n\n# Remove the `--dry_run --init` flags to actually schedule the jobs once everything is setup.\n```\n\n### dataset and metadata\nLearn more in the [datasets section](./DATASETS.md).\n\n### Audio tokenizers\n\nSee [MusicGen](./MUSICGEN.md)\n\n### Fine tuning existing models\n\nYou can initialize your model to one of the pretrained models by using the `continue_from` argument, in particular\n\n```bash\n# Using pretrained MusicGen-Style model.\ndora run solver=musicgen/musicgen_style_32khz model/lm/model_scale=medium continue_from=//pretrained/facebook/musicgen-style conditioner=style2music\n\n# Using another model you already trained with a Dora signature SIG.\ndora run solver=musicgen/musicgen_style_32khz model/lm/model_scale=medium continue_from=//sig/SIG conditioner=style2music\n\n# Or providing manually a path\ndora run solver=musicgen/musicgen_style_32khz model/lm/model_scale=medium continue_from=/checkpoints/my_other_xp/checkpoint.th\n```\n\n**Warning:** You are responsible for selecting the other parameters accordingly, in a way that make it compatible\n with the model you are fine tuning. Configuration is NOT automatically inherited from the model you continue from. In particular make sure to select the proper `conditioner` and `model/lm/model_scale`.\n\n**Warning:** We currently do not support fine tuning a model with slightly different layers. If you decide\n to change some parts, like the conditioning or some other parts of the model, you are responsible for manually crafting a checkpoint file from which we can safely run `load_state_dict`.\n If you decide to do so, make sure your checkpoint is saved with `torch.save` and contains a dict\n `{'best_state': {'model': model_state_dict_here}}`. Directly give the path to `continue_from` without a `//pretrained/` prefix.\n\n\n[arxiv]: https://arxiv.org/abs/2407.12563\n[musicgen_samples]: https://musicgenstyle.github.io/\n" }, { "path": "docs/TRAINING.md", "content": "# AudioCraft training pipelines\n\nAudioCraft training pipelines are built on top of PyTorch as our core deep learning library\nand [Flashy](https://github.com/facebookresearch/flashy) as our training pipeline design library,\nand [Dora](https://github.com/facebookresearch/dora) as our experiment manager.\nAudioCraft training pipelines are designed to be research and experiment-friendly.\n\n\n## Environment setup\n\nFor the base installation, follow the instructions from the [README.md](../README.md).\nBelow are some additional instructions for setting up the environment to train new models.\n\n### Team and cluster configuration\n\nIn order to support multiple teams and clusters, AudioCraft uses an environment configuration.\nThe team configuration allows to specify cluster-specific configurations (e.g. SLURM configuration),\nor convenient mapping of paths between the supported environments.\n\nEach team can have a yaml file under the [configuration folder](../config). To select a team set the\n`AUDIOCRAFT_TEAM` environment variable to a valid team name (e.g. `labs` or `default`):\n```shell\nconda env config vars set AUDIOCRAFT_TEAM=default\n```\n\nAlternatively, you can add it to your `.bashrc`:\n```shell\nexport AUDIOCRAFT_TEAM=default\n```\n\nIf not defined, the environment will default to the `default` team.\n\nThe cluster is automatically detected, but it is also possible to override it by setting\nthe `AUDIOCRAFT_CLUSTER` environment variable.\n\nBased on this team and cluster, the environment is then configured with:\n* The dora experiment outputs directory.\n* The available slurm partitions: categorized by global and team.\n* A shared reference directory: In order to facilitate sharing research models while remaining\nagnostic to the used compute cluster, we created the `//reference` symbol that can be used in\nYAML config to point to a defined reference folder containing shared checkpoints\n(e.g. baselines, models for evaluation...).\n\n**Important:** The default output dir for trained models and checkpoints is under `/tmp/`. This is suitable\nonly for quick testing. If you are doing anything serious you MUST edit the file `default.yaml` and\nproperly set the `dora_dir` entries.\n\n#### Overriding environment configurations\n\nYou can set the following environment variables to bypass the team's environment configuration:\n* `AUDIOCRAFT_CONFIG`: absolute path to a team config yaml file.\n* `AUDIOCRAFT_DORA_DIR`: absolute path to a custom dora directory.\n* `AUDIOCRAFT_REFERENCE_DIR`: absolute path to the shared reference directory.\n\n## Training pipelines\n\nEach task supported in AudioCraft has its own training pipeline and dedicated solver.\nLearn more about solvers and key designs around AudioCraft training pipeline below.\nPlease refer to the documentation of each task and model for specific information on a given task.\n\n\n### Solvers\n\nThe core training component in AudioCraft is the solver. A solver holds the definition\nof how to solve a given task: It implements the training pipeline logic, combining the datasets,\nmodel, optimization criterion and components and the full training loop. We refer the reader\nto [Flashy](https://github.com/facebookresearch/flashy) for core principles around solvers.\n\nAudioCraft proposes an initial solver, the `StandardSolver` that is used as the base implementation\nfor downstream solvers. This standard solver provides a nice base management of logging,\ncheckpoints loading/saving, xp restoration, etc. on top of the base Flashy implementation.\nIn AudioCraft, we made the assumption that all tasks are following the same set of stages:\ntrain, valid, evaluate and generation, each relying on a dedicated dataset.\n\nEach solver is responsible for defining the task to solve and the associated stages\nof the training loop in order to leave the full ownership of the training pipeline\nto the researchers. This includes loading the datasets, building the model and\noptimisation components, registering them and defining the execution of each stage.\nTo create a new solver for a given task, one should extend the StandardSolver\nand define each stage of the training loop. One can further customise its own solver\nstarting from scratch instead of inheriting from the standard solver.\n\n```python\nfrom . import base\nfrom .. import optim\n\n\nclass MyNewSolver(base.StandardSolver):\n\n def __init__(self, cfg: omegaconf.DictConfig):\n super().__init__(cfg)\n # one can add custom attributes to the solver\n self.criterion = torch.nn.L1Loss()\n\n def best_metric(self):\n # here optionally specify which metric to use to keep track of best state\n return 'loss'\n\n def build_model(self):\n # here you can instantiate your models and optimization related objects\n # this method will be called by the StandardSolver init method\n self.model = ...\n # the self.cfg attribute contains the raw configuration\n self.optimizer = optim.build_optimizer(self.model.parameters(), self.cfg.optim)\n # don't forget to register the states you'd like to include in your checkpoints!\n self.register_stateful('model', 'optimizer')\n # keep the model best state based on the best value achieved at validation for the given best_metric\n self.register_best('model')\n # if you want to add EMA around the model\n self.register_ema('model')\n\n def build_dataloaders(self):\n # here you can instantiate your dataloaders\n # this method will be called by the StandardSolver init method\n self.dataloaders = ...\n\n ...\n\n # For both train and valid stages, the StandardSolver relies on\n # a share common_train_valid implementation that is in charge of\n # accessing the appropriate loader, iterate over the data up to\n # the specified number of updates_per_epoch, run the ``run_step``\n # function that you need to implement to specify the behavior\n # and finally update the EMA and collect the metrics properly.\n @abstractmethod\n def run_step(self, idx: int, batch: tp.Any, metrics: dict):\n \"\"\"Perform one training or valid step on a given batch.\n \"\"\"\n ... # provide your implementation of the solver over a batch\n\n def train(self):\n \"\"\"Train stage.\n \"\"\"\n return self.common_train_valid('train')\n\n def valid(self):\n \"\"\"Valid stage.\n \"\"\"\n return self.common_train_valid('valid')\n\n @abstractmethod\n def evaluate(self):\n \"\"\"Evaluate stage.\n \"\"\"\n ... # provide your implementation here!\n\n @abstractmethod\n def generate(self):\n \"\"\"Generate stage.\n \"\"\"\n ... # provide your implementation here!\n```\n\n### About Epochs\n\nAudioCraft Solvers uses the concept of Epoch. One epoch doesn't necessarily mean one pass over the entire\ndataset, but instead represent the smallest amount of computation that we want to work with before checkpointing.\nTypically, we find that having an Epoch time around 30min is ideal both in terms of safety (checkpointing often enough)\nand getting updates often enough. One Epoch is at least a `train` stage that lasts for `optim.updates_per_epoch` (2000 by default),\nand a `valid` stage. You can control how long the valid stage takes with `dataset.valid.num_samples`.\nOther stages (`evaluate`, `generate`) will only happen every X epochs, as given by `evaluate.every` and `generate.every`).\n\n\n### Models\n\nIn AudioCraft, a model is a container object that wraps one or more torch modules together\nwith potential processing logic to use in a solver. For example, a model would wrap an encoder module,\na quantisation bottleneck module, a decoder and some tensor processing logic. Each of the previous components\ncan be considered as a small « model unit » on its own but the container model is a practical component\nto manipulate and train a set of modules together.\n\n### Datasets\n\nSee the [dedicated documentation on datasets](./DATASETS.md).\n\n### Metrics\n\nSee the [dedicated documentation on metrics](./METRICS.md).\n\n### Conditioners\n\nAudioCraft language models can be conditioned in various ways and the codebase offers a modular implementation\nof different conditioners that can be potentially combined together.\nLearn more in the [dedicated documentation on conditioning](./CONDITIONING.md).\n\n### Configuration\n\nAudioCraft's configuration is defined in yaml files and the framework relies on\n[hydra](https://hydra.cc/docs/intro/) and [omegaconf](https://omegaconf.readthedocs.io/) to parse\nand manipulate the configuration through Dora.\n\n##### :warning: Important considerations around configurations\n\nOur configuration management relies on Hydra and the concept of group configs to structure\nand compose configurations. Updating the root default configuration files will then have\nan impact on all solvers and tasks.\n**One should never change the default configuration files. Instead they should use Hydra config groups in order to store custom configuration.**\nOnce this configuration is created and used for running experiments, you should not edit it anymore.\n\nNote that as we are using Dora as our experiment manager, all our experiment tracking is based on\nsignatures computed from delta between configurations.\n**One must therefore ensure backward compatibility of the configuration at all time.**\nSee [Dora's README](https://github.com/facebookresearch/dora) and the\n[section below introduction Dora](#running-experiments-with-dora).\n\n##### Configuration structure\n\nThe configuration is organized in config groups:\n* `conditioner`: default values for conditioning modules.\n* `dset`: contains all data source related information (paths to manifest files\nand metadata for a given dataset).\n* `model`: contains configuration for each model defined in AudioCraft and configurations\nfor different variants of models.\n* `solver`: contains the default configuration for each solver as well as configuration\nfor each solver task, combining all the above components.\n* `teams`: contains the cluster configuration per teams. See environment setup for more details.\n\nThe `config.yaml` file is the main configuration that composes the above groups\nand contains default configuration for AudioCraft.\n\n##### Solver's core configuration structure\n\nThe core configuration structure shared across solver is available in `solvers/default.yaml`.\n\n##### Other configuration modules\n\nAudioCraft configuration contains the different setups we used for our research and publications.\n\n## Running experiments with Dora\n\n### Launching jobs\n\nTry launching jobs for different tasks locally with dora run:\n\n```shell\n# run compression task with lightweight encodec\ndora run solver=compression/debug\n```\n\nMost of the time, the jobs are launched through dora grids, for example:\n\n```shell\n# run compression task through debug grid\ndora grid compression.debug\n```\n\nLearn more about running experiments with Dora below.\n\n### A small introduction to Dora\n\n[Dora](https://github.com/facebookresearch/dora) is the experiment manager tool used in AudioCraft.\nCheck out the README to learn how Dora works. Here is a quick summary of what to know:\n* An XP is a unique set of hyper-parameters with a given signature. The signature is a hash\nof those hyper-parameters. We always refer to an XP with its signature, e.g. 9357e12e. We will see\nafter that one can retrieve the hyper-params and re-rerun it in a single command.\n* In fact, the hash is defined as a delta between the base config and the one obtained\nwith the config overrides you passed from the command line. This means you must never change\nthe `conf/**.yaml` files directly, except for editing things like paths. Changing the default values\nin the config files means the XP signature won't reflect that change, and wrong checkpoints might be reused.\nI know, this is annoying, but the reason is that otherwise, any change to the config file would mean\nthat all XPs ran so far would see their signature change.\n\n#### Dora commands\n\n```shell\ndora info -f 81de367c # this will show the hyper-parameter used by a specific XP.\n # Be careful some overrides might present twice, and the right most one\n # will give you the right value for it.\n\ndora run -d -f 81de367c # run an XP with the hyper-parameters from XP 81de367c.\n # `-d` is for distributed, it will use all available GPUs.\n\ndora run -d -f 81de367c dataset.batch_size=32 # start from the config of XP 81de367c but change some hyper-params.\n # This will give you a new XP with a new signature (e.g. 3fe9c332).\n\ndora info -f SIG -t # will tail the log (if the XP has scheduled).\n# if you need to access the logs of the process for rank > 0, in particular because a crash didn't happen in the main\n# process, then use `dora info -f SIG` to get the main log name (finished into something like `/5037674_0_0_log.out`)\n# and worker K can be accessed as `/5037674_0_{K}_log.out`.\n# This is only for scheduled jobs, for local distributed runs with `-d`, then you should go into the XP folder,\n# and look for `worker_{K}.log` logs.\n```\n\nAn XP runs from a specific folder based on its signature, under the\n`//experiments/audiocraft/outputs/` folder.\nYou can safely interrupt a training and resume it, it will reuse any existing checkpoint,\nas it will reuse the same folder. If you made some change to the code and need to ignore\na previous checkpoint you can use `dora run --clear [RUN ARGS]`.\n\nIf you have a Slurm cluster, you can also use the dora grid command, e.g.\n\n```shell\n# Run a dummy grid located at `audiocraft/grids/my_grid_folder/my_grid_name.py`\ndora grid my_grid_folder.my_grid_name\n# The following will simply display the grid and also initialize the Dora experiments database.\n# You can then simply refer to a config using its signature (e.g. as `dora run -f SIG`).\ndora grid my_grid_folder.my_grid_name --dry_run --init\n```\n\nPlease refer to the [Dora documentation](https://github.com/facebookresearch/dora) for more information.\n\n\n#### Clearing up past experiments\n\n```shell\n# This will cancel all the XPs and delete their folder and checkpoints.\n# It will then reschedule them starting from scratch.\ndora grid my_grid_folder.my_grid_name --clear\n# The following will delete the folder and checkpoint for a single XP,\n# and then run it afresh.\ndora run [-f BASE_SIG] [ARGS] --clear\n```\n" }, { "path": "docs/WATERMARKING.md", "content": "# AudioSeal: Proactive Localized Watermarking\n\nAudioCraft provides the training code and models for AudioSeal, a method for speech localized watermarking [Proactive Detection of Voice Cloning with Localized Watermarking][arxiv], with state-of-the-art robustness and detector speed. It jointly trains a generator that embeds a watermark in the audio, and a detector that detects the watermarked fragments in longer audios, even in the presence of editing.\n\n## Installation and setup\n\nMake sure to install audiocraft version `1.4.0a1` or later, and with the `[wm]` extra (see [README](../README.md)).\nAlternatively, you can just install audioseal yourself. To install AudioSeal, follow [Installation](https://github.com/facebookresearch/audioseal) guidelines in the AudioSeal repo.\n\n_NOTE_: Since we use AAC augmentation in our training loop, you need to install ffmpeg, or it will not work (See Section \"Installation\" in [README](../README.md)).\n\nMake sure you follow [steps for basic training setup](TRAINING.md) before starting.\n\n## API\nCheck the [Github repository](https://github.com/facebookresearch/audioseal) for more details.\n\n## Training\n\nThe [WatermarkSolver](../audiocraft/solvers/watermark.py) implements the AudioSeal's training pipeline. It joins the generator and detector that wrap\n`audioseal.AudioSealWM` and `audioseal.AudioSealDetector` respectively. For the training recipe, see [config/solver/watermark/robustness.yaml](../config/solver/watermark/robustness.yaml).\n\nFor illustration, we use the three example audios in `datasets`, with datasourc definition in [dset/audio/example.yaml](../config/dset/audio/example.yaml) (Please read [DATASET](./DATASETS.md) to understand AudioCraft's dataset structure.)\n\nTo run the Watermarking training pipeline locally:\n\n```bash\ndora run solver=watermark/robustness dset=audio/example\n```\n\nyou can override model / experiment parameters here directly like:\n\n```bash\ndora run solver=watermark/robustness dset=audio/example sample_rate=24000\n```\n\nIf you want to run in debug mode:\n\n```bash\npython3 -m pdb -c c -m dora run solver=watermark/robustness dset=audio/example\n```\n" }, { "path": "egs/example/data.jsonl", "content": "{\"path\": \"dataset/example/electro_1.mp3\", \"duration\": 15.024, \"sample_rate\": 48000, \"amplitude\": null, \"weight\": null, \"info_path\": null}\n{\"path\": \"dataset/example/electro_2.mp3\", \"duration\": 20.035918367346937, \"sample_rate\": 44100, \"amplitude\": null, \"weight\": null, \"info_path\": null}\n" }, { "path": "jasco_demo.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# JASCO\\n\",\n \"Welcome to JASCO's demo jupyter notebook. \\n\",\n \"Here you will find a self-contained example of how to use JASCO for temporally controlled music generation.\\n\",\n \"\\n\",\n \"You can choose a model from the following selection:\\n\",\n \"1. facebook/jasco-chords-drums-400M - 10s music generation conditioned on text, chords and drums, 400M parameters\\n\",\n \"2. facebook/jasco-chords-drums-1B - 10s music generation conditioned on text, chords and drums, 1B parameters\\n\",\n \"\\n\",\n \"\\n\",\n \"First, we start by initializing the JASCO model:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"/private/home/ortal1/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\\n\",\n \" @torch.library.impl_abstract(\\\"xformers_flash::flash_fwd\\\")\\n\",\n \"/private/home/ortal1/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\\n\",\n \" @torch.library.impl_abstract(\\\"xformers_flash::flash_bwd\\\")\\n\",\n \"/private/home/ortal1/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\\n\",\n \" from .autonotebook import tqdm as notebook_tqdm\\n\",\n \"/checkpoint/ortal1/Projects/jasco_release/audiocraft/models/loaders.py:71: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\\n\",\n \" return torch.load(file, map_location=device)\\n\",\n \"/private/home/ortal1/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/transformers/models/encodec/modeling_encodec.py:124: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\\n\",\n \" self.register_buffer(\\\"padding_total\\\", torch.tensor(kernel_size - stride, dtype=torch.int64), persistent=False)\\n\"\n ]\n }\n ],\n \"source\": [\n \"import os \\n\",\n \"from audiocraft.models import JASCO\\n\",\n \"\\n\",\n \"chords_mapping_path = os.path.abspath('./assets/chord_to_index_mapping.pkl')\\n\",\n \"model = JASCO.get_pretrained('facebook/jasco-chords-drums-1B', chords_mapping_path='./assets/chord_to_index_mapping.pkl')\\n\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, let us configure the generation parameters. Specifically, you can control the following:\\n\",\n \"* `cfg_coef_all` (float, optional): Coefficient used for classifier free guidance - fully conditional term. \\n\",\n \" Defaults to 5.0.\\n\",\n \"* `cfg_coef_txt` (float, optional): Coefficient used for classifier free guidance - additional text conditional term. \\n\",\n \" Defaults to 0.0.\\n\",\n \"\\n\",\n \"When left unchanged, JASCO will revert to its default parameters.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model.set_generation_params(\\n\",\n \" cfg_coef_all=0.0,\\n\",\n \" cfg_coef_txt=5.0\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, we can go ahead and start generating music given textual prompts.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Text-conditional Generation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"# set textual prompt\\n\",\n \"text = \\\"Funky groove with electric piano playing blue chords rhythmically\\\"\\n\",\n \"\\n\",\n \"# run the model\\n\",\n \"print(\\\"Generating...\\\") \\n\",\n \"output = model.generate(descriptions=[text], progress=True)\\n\",\n \"\\n\",\n \"# display the result\\n\",\n \"print(f\\\"Text: {text}\\\\n\\\")\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Now we can start adding temporal controls! We begin with conditioning on chord progressions:\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Chords-conditional Generation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"model.set_generation_params(\\n\",\n \" cfg_coef_all=1.5,\\n\",\n \" cfg_coef_txt=2.5\\n\",\n \")\\n\",\n \"\\n\",\n \"# set textual prompt\\n\",\n \"text = \\\"Strings, woodwind, orchestral, symphony.\\\"\\n\",\n \"\\n\",\n \"# define chord progression\\n\",\n \"chords = [('C', 0.0), ('D', 2.0), ('F', 4.0), ('Ab', 6.0), ('Bb', 7.0), ('C', 8.0)]\\n\",\n \"\\n\",\n \"# run the model\\n\",\n \"print(\\\"Generating...\\\")\\n\",\n \"output = model.generate_music(descriptions=[text], chords=chords, progress=True)\\n\",\n \"\\n\",\n \"# display the result\\n\",\n \"print(f'Text: {text}')\\n\",\n \"print(f'Chord progression: {chords}')\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, we can condition the generation on drum tracks:\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Drums-conditional Generation\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import torchaudio\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"\\n\",\n \"# load drum prompt\\n\",\n \"drums_waveform, sr = torchaudio.load(\\\"./assets/sep_drums_1.mp3\\\")\\n\",\n \"\\n\",\n \"# set textual prompt \\n\",\n \"text = \\\"distortion guitars, heavy rock, catchy beat\\\"\\n\",\n \"\\n\",\n \"# run the model\\n\",\n \"print(\\\"Generating...\\\")\\n\",\n \"output = model.generate_music(\\n\",\n \" descriptions=[text],\\n\",\n \" drums_wav=drums_waveform,\\n\",\n \" drums_sample_rate=sr,\\n\",\n \" progress=True\\n\",\n \")\\n\",\n \"\\n\",\n \"# display the result\\n\",\n \"print('drum prompt:')\\n\",\n \"display_audio(drums_waveform, sample_rate=sr)\\n\",\n \"print(f'Text: {text}')\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We can also combine multiple temporal controls! Let's move on to generating with both chords and drums conditioning:\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Drums + Chords conditioning\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import torchaudio\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"\\n\",\n \"\\n\",\n \"# load drum prompt\\n\",\n \"drums_waveform, sr = torchaudio.load(\\\"./assets/sep_drums_1.mp3\\\")\\n\",\n \"\\n\",\n \"# set textual prompt \\n\",\n \"text = \\\"string quartet, orchestral, dramatic\\\"\\n\",\n \"\\n\",\n \"# define chord progression\\n\",\n \"chords = [('C', 0.0), ('D', 2.0), ('F', 4.0), ('Ab', 6.0), ('Bb', 7.0), ('C', 8.0)]\\n\",\n \"\\n\",\n \"# run the model\\n\",\n \"print(\\\"Generating...\\\")\\n\",\n \"output = model.generate_music(\\n\",\n \" descriptions=[text],\\n\",\n \" drums_wav=drums_waveform,\\n\",\n \" drums_sample_rate=sr,\\n\",\n \" chords=chords,\\n\",\n \" progress=True\\n\",\n \")\\n\",\n \"\\n\",\n \"# display the result\\n\",\n \"print('drum prompt:')\\n\",\n \"display_audio(drums_waveform, sample_rate=sr)\\n\",\n \"print(f'Chord progression: {chords}')\\n\",\n \"print(f'Text: {text}')\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Melody + Drums + Chords conditioning - inference example\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Source melody:\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"\\n\",\n \" \\n\",\n \" \"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAABj0AAADQCAYAAABcDaP2AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAYDElEQVR4nO3df2jc9f0H8NfFtKm2udTUNVlpioXJanGt2GobHPuhmdkoY8UOHJStlqIwErFmbKOwVTYGEQfqCv5iP3R/WCsddGKZSqkuMoxaUwrVzbKB0GCXRJEmNdC0Np/vH6P37fVnkia5yzuPBxz0Pvfu3evzufe97uzT9+eTy7IsCwAAAAAAgCmuotQFAAAAAAAAjAehBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkITKUhdwtuHh4Thy5EhUV1dHLpcrdTkAAAAAAEAJZVkWx44diwULFkRFxcXXcpRd6HHkyJFoaGgodRkAAAAAAEAZ6e7ujoULF150TNmFHtXV1RHxv+Lz+XyJqwEAAAAAAEppYGAgGhoaCvnBxZRd6HH6lFb5fF7oAQAAAAAARESM6JIYLmQOAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAk4bJCj4ceeihyuVxs3ry5sO348ePR0tIS8+bNizlz5sS6deuit7f3cusEAAAAAAC4qDGHHvv27Yunn346li1bVrT9gQceiJdeeil27twZHR0dceTIkbjzzjsvu1AAAAAAAICLGVPo8dlnn8X69evj97//fVx99dWF7f39/fHHP/4xHnnkkbjttttixYoV8cwzz8Sbb74Zb7311rgVDQAAAAAAcLYxhR4tLS2xZs2aaGpqKtre1dUVJ0+eLNq+ZMmSWLRoUXR2dl5epQAAAAAAABdROdq/sGPHjti/f3/s27fvnMd6enpi5syZMXfu3KLtdXV10dPTc97nGxoaiqGhocL9gYGB0ZYEAAAAAAAwupUe3d3dcf/998dzzz0Xs2bNGpcC2tvbo6ampnBraGgYl+cFAAAAAACml1GFHl1dXdHX1xc33XRTVFZWRmVlZXR0dMS2bduisrIy6urq4sSJE3H06NGiv9fb2xv19fXnfc4tW7ZEf39/4dbd3T3mnQEAAAAAAKavUZ3e6vbbb4+DBw8Wbdu4cWMsWbIkfv7zn0dDQ0PMmDEj9u7dG+vWrYuIiEOHDsXhw4ejsbHxvM9ZVVUVVVVVYywfAAAAAADgf0YVelRXV8cNN9xQtG327Nkxb968wvZNmzZFW1tb1NbWRj6fj/vuuy8aGxtj9erV41c1AAAAAADAWUZ9IfNLefTRR6OioiLWrVsXQ0ND0dzcHE888cR4vwwAAAAAAECRXJZlWamLONPAwEDU1NREf39/5PP5UpcDAAAAAACU0Ghyg1FdyBwAAAAAAKBcCT0AAAAAAIAkCD0AAAAAAIAkCD0AAAAAAIAkVJa6ACZXLpeb8NfIsuySr32hMUy+iZoTI5kH4/WcAACjNRm/iy+H3z2lNdnzw39DAQClUsrfxRP1+8ZKDwAAAAAAIAlCDwAAAAAAIAlObzXNlHJJtOXY5Wmy3xfzAAAoB36TcDHlMj/KpQ4AIF0p/t6w0gMAAAAAAEiC0AMAAAAAAEiC0AMAAAAAAEiC0AMAAAAAAEiC0AMAAAAAAEhCZakLSFkulxvV+CzLJqgSSmm082CipDi/JvrYpnjMUlcun7eRML+AlJWyH+uv5aUcv5un0xy5nOM/nY4TAIynifr9c+Z383i+Rorf+VZ6AAAAAAAASRB6AAAAAAAASXB6qwmU4tIgRs88mDiOLWczJwDKg37MaeZCaTn+ADD5JuP713f8xVnpAQAAAAAAJEHoAQAAAAAAJMHprc5yoSvfWzKUvgu99yNljqRnqs2Jy633fCZqHyai1gifQ4CRmqg+fJp+PPX4buZs4zknJmMeTKXfwgCUp/H6LvH9UXpWegAAAAAAAEkQegAAAAAAAElwequzWH40fXnvOdtUmxNTqd6pVCtAivRhzmZOcLapNiemWr0AlB/fJemw0gMAAAAAAEiC0AMAAAAAAEiC0AMAAAAAAEjClLqmRy6XG7fnco62NI12jpgHlMJ49rIzTcR8nkq1AqRIH2YsxjJvzAnKSbnM4YnqwRPlzGMwktrL/XN/sX0o99qB8THZfVhvSYeVHgAAAAAAQBKEHgAAAAAAQBKm1OmtLDHiUswRpoKpNE+nUq0AKdKHGQvzhqmuXOZwudQxFlO59tNS2Afg8ugDjJWVHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBIqS10AAAAAAMB0kMvlCn/OsqyElUC6rPQAAAAAAACSIPQAAAAAAACSIPQAAAAAAACSMKrQo729PW6++eaorq6O+fPnx9q1a+PQoUNFY44fPx4tLS0xb968mDNnTqxbty56e3vHtWgAAAAAgKkmy7LCDZgYowo9Ojo6oqWlJd56663Ys2dPnDx5Mu64444YHBwsjHnggQfipZdeip07d0ZHR0ccOXIk7rzzznEvHAAAAAAA4Ey57DJixY8//jjmz58fHR0d8bWvfS36+/vjC1/4Qmzfvj2+//3vR0TEBx98ENdff310dnbG6tWrL/mcAwMDUVNTE/39/ZHP58daGgAAAAAAkIDR5AaXdU2P/v7+iIiora2NiIiurq44efJkNDU1FcYsWbIkFi1aFJ2dnZfzUgAAAAAAABdVOda/ODw8HJs3b45bb701brjhhoiI6OnpiZkzZ8bcuXOLxtbV1UVPT895n2doaCiGhoYK9wcGBsZaEgAAAAAAMI2NeaVHS0tLvPfee7Fjx47LKqC9vT1qamoKt4aGhst6PgAAAAAAYHoaU+jR2toau3fvjtdffz0WLlxY2F5fXx8nTpyIo0ePFo3v7e2N+vr68z7Xli1bor+/v3Dr7u4eS0kAAAAAAMA0N6rQI8uyaG1tjV27dsVrr70WixcvLnp8xYoVMWPGjNi7d29h26FDh+Lw4cPR2Nh43uesqqqKfD5fdAMAAAAAABitUV3To6WlJbZv3x4vvvhiVFdXF67TUVNTE1deeWXU1NTEpk2boq2tLWprayOfz8d9990XjY2NsXr16gnZAQAAAAAAgIiIXJZl2YgH53Ln3f7MM8/E3XffHRERx48fj5/85Cfx/PPPx9DQUDQ3N8cTTzxxwdNbnW1gYCBqamqiv7/fqg8AAAAAAJjmRpMbjCr0mAxCDwAAAAAA4LTR5AZjupA5AAAAAABAuRnVNT0mU01NzTnbSrko5UKn9rqYyax3pPWV2cKeMbnYvo52/858rhSOTTm50PvkOMPkGcl3w3h+Ji+np47le3aijKT2yT62AOfjtyxQjkb7u24q9y99GKA8WekBAAAAAAAkQegBAAAAAAAkoWxPb1VuFzIv92WK5V7feBrPfZ1Ox22yObZQepP9Obyc15tqPWOq1QukSS8CytF06k3TaV8BphIrPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCRUlroAAAAARiaXyxX+nGXZJbczNZz5/p3pQu/xSJkLcK6J6pcj+RxTWiPpo1Pt/Upxn2A8WOkBAAAAAAAkQegBAAAAAAAkwemtAAAApogLnaLCqSumtpG8f95jGB8T9VnyGS1/Kb5HKe4TjAcrPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCRUlroAAADSlMvlRjQuy7IJrgQAAIDpwkoPAAAAAAAgCUIPAAAAAAAgCU5vBQDAhHDaKgAAACablR4AAAAAAEAShB4AAAAAAEASnN6KcZfL5c67PfVTXJT7fpd7fSlyzBmLC82bsZiouWZuAwAAAJdrJP8GMpZ/a7DSAwAAAAAASILQAwAAAAAASILTWzHupuvpTcp9v8u9vhQ55ozFVJg3U6FGAAAAoLxN1L8vWOkBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkobLUBQAAAMBkyeVylxyTZdkkVDJ9jOSYn8nxh/Mb7WdpMvi8AuXISg8AAAAAACAJQg8AAAAAACAJExZ6PP7443HttdfGrFmzYtWqVfHOO+9M1EsBAADAiGRZdskb42skx9zxh0sb7WdpMm4A5WhCQo8XXngh2tra4sEHH4z9+/fH8uXLo7m5Ofr6+ibi5QAAAAAAACYm9HjkkUfinnvuiY0bN8bSpUvjqaeeiquuuir+9Kc/TcTLAQAAAAAAjH/oceLEiejq6oqmpqb/f5GKimhqaorOzs7xfjkAAAAAAICIiKgc7yf85JNP4tSpU1FXV1e0va6uLj744INzxg8NDcXQ0FDh/sDAwHiXBAAAAAAATAPjHnqMVnt7e/zqV786Z7vwAwAAAAAAOJ0XZFl2ybHjHnpcc801ccUVV0Rvb2/R9t7e3qivrz9n/JYtW6Ktra1w/6OPPoqlS5dGQ0PDeJcGAAAAAABMUceOHYuampqLjhn30GPmzJmxYsWK2Lt3b6xduzYiIoaHh2Pv3r3R2tp6zviqqqqoqqoq3J8zZ050d3dHlmWxaNGi6O7ujnw+P95lAky6gYGBaGho0NeAZOhrQGr0NSA1+hqQiizL4tixY7FgwYJLjp2Q01u1tbXFhg0bYuXKlXHLLbfEY489FoODg7Fx48ZL/t2KiopYuHBhYblKPp/XlIGk6GtAavQ1IDX6GpAafQ1IwaVWeJw2IaHHXXfdFR9//HFs3bo1enp64sYbb4xXXnnlnIubAwAAAAAAjJcJu5B5a2vreU9nBQAAAAAAMBEqSl3AhVRVVcWDDz5YdL0PgKlMXwNSo68BqdHXgNToa8B0lMuyLCt1EQAAAAAAAJerbFd6AAAAAAAAjIbQAwAAAAAASILQAwAAAAAASILQAwAAAAAASEJZhh6PP/54XHvttTFr1qxYtWpVvPPOO6UuCeC83njjjfjud78bCxYsiFwuF3/961+LHs+yLLZu3Rpf/OIX48orr4ympqb497//XTTm008/jfXr10c+n4+5c+fGpk2b4rPPPpvEvQD4f+3t7XHzzTdHdXV1zJ8/P9auXRuHDh0qGnP8+PFoaWmJefPmxZw5c2LdunXR29tbNObw4cOxZs2auOqqq2L+/Pnx05/+ND7//PPJ3BWAiIh48sknY9myZZHP5yOfz0djY2O8/PLLhcf1NGAqe+ihhyKXy8XmzZsL2/Q1YLoru9DjhRdeiLa2tnjwwQdj//79sXz58mhubo6+vr5SlwZwjsHBwVi+fHk8/vjj53384Ycfjm3btsVTTz0Vb7/9dsyePTuam5vj+PHjhTHr16+P999/P/bs2RO7d++ON954I+69997J2gWAIh0dHdHS0hJvvfVW7NmzJ06ePBl33HFHDA4OFsY88MAD8dJLL8XOnTujo6Mjjhw5EnfeeWfh8VOnTsWaNWvixIkT8eabb8af//znePbZZ2Pr1q2l2CVgmlu4cGE89NBD0dXVFe+++27cdttt8b3vfS/ef//9iNDTgKlr37598fTTT8eyZcuKtutrwLSXlZlbbrkla2lpKdw/depUtmDBgqy9vb2EVQFcWkRku3btKtwfHh7O6uvrs9/+9reFbUePHs2qqqqy559/PsuyLPvnP/+ZRUS2b9++wpiXX345y+Vy2UcffTRptQNcSF9fXxYRWUdHR5Zl/+tjM2bMyHbu3FkY869//SuLiKyzszPLsiz729/+llVUVGQ9PT2FMU8++WSWz+ezoaGhyd0BgPO4+uqrsz/84Q96GjBlHTt2LLvuuuuyPXv2ZF//+tez+++/P8syv9UAsizLymqlx4kTJ6KrqyuampoK2yoqKqKpqSk6OztLWBnA6H344YfR09NT1NNqampi1apVhZ7W2dkZc+fOjZUrVxbGNDU1RUVFRbz99tuTXjPA2fr7+yMiora2NiIiurq64uTJk0W9bcmSJbFo0aKi3vaVr3wl6urqCmOam5tjYGCg8H9WA5TCqVOnYseOHTE4OBiNjY16GjBltbS0xJo1a4r6V4TfagAREZWlLuBMn3zySZw6daqo6UZE1NXVxQcffFCiqgDGpqenJyLivD3t9GM9PT0xf/78oscrKyujtra2MAagVIaHh2Pz5s1x6623xg033BAR/+tbM2fOjLlz5xaNPbu3na/3nX4MYLIdPHgwGhsb4/jx4zFnzpzYtWtXLF26NA4cOKCnAVPOjh07Yv/+/bFv375zHvNbDaDMQg8AAMpHS0tLvPfee/GPf/yj1KUAXJYvf/nLceDAgejv74+//OUvsWHDhujo6Ch1WQCj1t3dHffff3/s2bMnZs2aVepyAMpSWZ3e6pprrokrrrgient7i7b39vZGfX19iaoCGJvTfetiPa2+vj76+vqKHv/888/j008/1feAkmptbY3du3fH66+/HgsXLixsr6+vjxMnTsTRo0eLxp/d287X+04/BjDZZs6cGV/60pdixYoV0d7eHsuXL4/f/e53ehow5XR1dUVfX1/cdNNNUVlZGZWVldHR0RHbtm2LysrKqKur09eAaa+sQo+ZM2fGihUrYu/evYVtw8PDsXfv3mhsbCxhZQCjt3jx4qivry/qaQMDA/H2228XelpjY2McPXo0urq6CmNee+21GB4ejlWrVk16zQBZlkVra2vs2rUrXnvttVi8eHHR4ytWrIgZM2YU9bZDhw7F4cOHi3rbwYMHi0LdPXv2RD6fj6VLl07OjgBcxPDwcAwNDelpwJRz++23x8GDB+PAgQOF28qVK2P9+vWFP+trwHRXdqe3amtriw0bNsTKlSvjlltuicceeywGBwdj48aNpS4N4ByfffZZ/Oc//ync//DDD+PAgQNRW1sbixYtis2bN8dvfvObuO6662Lx4sXxy1/+MhYsWBBr166NiIjrr78+vv3tb8c999wTTz31VJw8eTJaW1vjBz/4QSxYsKBEewVMZy0tLbF9+/Z48cUXo7q6unBe55qamrjyyiujpqYmNm3aFG1tbVFbWxv5fD7uu+++aGxsjNWrV0dExB133BFLly6NH/7wh/Hwww9HT09P/OIXv4iWlpaoqqoq5e4B09CWLVviO9/5TixatCiOHTsW27dvj7///e/x6quv6mnAlFNdXV241tpps2fPjnnz5hW262vAdFd2ocddd90VH3/8cWzdujV6enrixhtvjFdeeeWcCywBlIN33303vvnNbxbut7W1RUTEhg0b4tlnn42f/exnMTg4GPfee28cPXo0vvrVr8Yrr7xSdO7V5557LlpbW+P222+PioqKWLduXWzbtm3S9wUgIuLJJ5+MiIhvfOMbRdufeeaZuPvuuyMi4tFHHy30q6GhoWhubo4nnniiMPaKK66I3bt3x49//ONobGyM2bNnx4YNG+LXv/71ZO0GQEFfX1/86Ec/iv/+979RU1MTy5Yti1dffTW+9a1vRYSeBqRHXwOmu1yWZVmpiwAAAAAAALhcZXVNDwAAAAAAgLESegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEn4P8oFZFHDCXiRAAAAAElFTkSuQmCC\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Chords:\\n\",\n \"[('N', 0.0), ('C', 0.32), ('Dm7', 3.456), ('Am', 4.608), ('F', 8.32), ('C', 9.216)]\\n\"\n ]\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Separated drums:\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"\\n\",\n \" \\n\",\n \" \"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Generating...\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"\\n\",\n \" \\n\",\n \" \"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"\\n\",\n \" \\n\",\n \" \"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"\\n\",\n \" \\n\",\n \" \"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"%matplotlib inline\\n\",\n \"import torchaudio \\n\",\n \"from audiocraft.models import JASCO\\n\",\n \"from demucs import pretrained\\n\",\n \"from demucs.apply import apply_model\\n\",\n \"from demucs.audio import convert_audio\\n\",\n \"import torch\\n\",\n \"from audiocraft.utils.notebook import display_audio\\n\",\n \"import matplotlib.pyplot as plt\\n\",\n \"\\n\",\n \"# --------------------------\\n\",\n \"# First, choose file to load\\n\",\n \"# --------------------------\\n\",\n \"fnames = ['salience_1', 'salience_2']\\n\",\n \"chords = [\\n\",\n \" [('N', 0.0), ('Eb7', 1.088000000), ('C#', 4.352000000), ('D', 4.864000000), ('Dm7', 6.720000000), ('G7', 8.256000000), ('Am7b5/G', 9.152000000)], # for salience 1\\n\",\n \" [('N', 0.0), ('C', 0.320000000), ('Dm7', 3.456000000), ('Am', 4.608000000), ('F', 8.320000000), ('C', 9.216000000)] # for salience 2\\n\",\n \"]\\n\",\n \"file_idx = 1 # either 0 or 1\\n\",\n \"\\n\",\n \"\\n\",\n \"# ------------------------------------\\n\",\n \"# display audio, melody map and chords\\n\",\n \"# ------------------------------------\\n\",\n \"def plot_chromagram(tensor):\\n\",\n \" # Check if tensor is a PyTorch tensor\\n\",\n \" if not torch.is_tensor(tensor):\\n\",\n \" raise ValueError('Input should be a PyTorch tensor')\\n\",\n \" tensor = tensor.numpy().T # C, T\\n\",\n \" plt.figure(figsize=(20, 20))\\n\",\n \" plt.imshow(tensor, cmap='binary', interpolation='nearest', origin='lower')\\n\",\n \" plt.show()\\n\",\n \"\\n\",\n \"# load salience and display the corresponding wav\\n\",\n \"melody_prompt_wav, melody_prompt_sr = torchaudio.load(f\\\"./assets/{fnames[file_idx]}.wav\\\")\\n\",\n \"print(\\\"Source melody:\\\")\\n\",\n \"display_audio(melody_prompt_wav, sample_rate=melody_prompt_sr)\\n\",\n \"melody = torch.load(f\\\"./assets/{fnames[file_idx]}.th\\\", weights_only=True)\\n\",\n \"plot_chromagram(melody)\\n\",\n \"print(\\\"Chords:\\\")\\n\",\n \"print(chords[file_idx])\\n\",\n \"\\n\",\n \"# --------------------------------------------------\\n\",\n \"# use demucs to seperate the drums stem from src mix\\n\",\n \"# --------------------------------------------------\\n\",\n \"def _get_drums_stem(wav: torch.Tensor, sample_rate: int) -> torch.Tensor:\\n\",\n \" \\\"\\\"\\\"Get parts of the wav that holds the drums, extracting the main stems from the wav.\\\"\\\"\\\"\\n\",\n \" demucs_model = pretrained.get_model('htdemucs').to('cuda')\\n\",\n \" wav = convert_audio(\\n\",\n \" wav, sample_rate, demucs_model.samplerate, demucs_model.audio_channels) # type: ignore\\n\",\n \" stems = apply_model(demucs_model, wav.cuda().unsqueeze(0), device='cuda').squeeze(0)\\n\",\n \" drum_stem = stems[demucs_model.sources.index('drums')] # extract relevant stems for drums conditioning\\n\",\n \" return convert_audio(drum_stem.cpu(), demucs_model.samplerate, sample_rate, 1) # type: ignore\\n\",\n \"drums_wav = _get_drums_stem(melody_prompt_wav, melody_prompt_sr)\\n\",\n \"print(\\\"Separated drums:\\\")\\n\",\n \"display_audio(drums_wav, sample_rate=melody_prompt_sr)\\n\",\n \"\\n\",\n \"# ----------------------------------\\n\",\n \"# Generate using the loaded controls\\n\",\n \"# ----------------------------------\\n\",\n \"# these are free-form texts written randomly\\n\",\n \"texts = [\\n\",\n \" '90s rock with heavy drums and hammond',\\n\",\n \" '80s pop with groovy synth bass and drum machine',\\n\",\n \" 'folk song with leading accordion',\\n\",\n \"]\\n\",\n \"\\n\",\n \"print(\\\"Generating...\\\")\\n\",\n \"# replacing dynammic solver with simple euler solver\\n\",\n \"model.set_generation_params(cfg_coef_all=1.5, cfg_coef_txt=2.5, euler=True, euler_steps=50) # manually set with euler solver\\n\",\n \"output = model.generate_music(\\n\",\n \" descriptions=texts,\\n\",\n \" chords=chords[file_idx],\\n\",\n \" drums_wav=drums_wav,\\n\",\n \" drums_sample_rate=melody_prompt_sr,\\n\",\n \" melody_salience_matrix=melody.permute(1, 0),\\n\",\n \" progress=True\\n\",\n \")\\n\",\n \"display_audio(output, sample_rate=model.compression_model.sample_rate)\"\n ]\n }\n ],\n \"metadata\": {\n \"kernelspec\": {\n \"display_name\": \"jasco_dev\",\n \"language\": \"python\",\n \"name\": \"python3\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.9.19\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "model_cards/AUDIOGEN_MODEL_CARD.md", "content": "# AudioGen Model Card\n\n## Model details\n**Organization developing the model:** The FAIR team of Meta AI.\n\n**Model date:** This version of AudioGen was trained between July 2023 and August 2023.\n\n**Model version:** This is version 2 of the model, not to be confused with the original AudioGen model published in [\"AudioGen: Textually Guided Audio Generation\"][audiogen].\nIn this version (v2), AudioGen was trained on the same data, but with some other differences:\n1. This model was trained on 10 seconds (vs. 5 seconds in v1).\n2. The discrete representation used under the hood is extracted using a retrained EnCodec model on the environmental sound data, following the EnCodec setup detailed in the [\"Simple and Controllable Music Generation\" paper][musicgen].\n3. No audio mixing augmentations.\n\n**Model type:** AudioGen consists of an EnCodec model for audio tokenization, and an auto-regressive language model based on the transformer architecture for audio modeling. The released model has 1.5B parameters.\n\n**Paper or resource for more information:** More information can be found in the paper [AudioGen: Textually Guided Audio Generation](https://arxiv.org/abs/2209.15352).\n\n**Citation details:** See [AudioGen paper][audiogen]\n\n**License:** Code is released under MIT, model weights are released under CC-BY-NC 4.0.\n\n**Where to send questions or comments about the model:** Questions and comments about AudioGen can be sent via the [GitHub repository](https://github.com/facebookresearch/audiocraft) of the project, or by opening an issue.\n\n## Intended use\n**Primary intended use:** The primary use of AudioGen is research on AI-based audio generation, including:\n- Research efforts, such as probing and better understanding the limitations of generative models to further improve the state of science\n- Generation of sound guided by text to understand current abilities of generative AI models by machine learning amateurs\n\n**Primary intended users:** The primary intended users of the model are researchers in audio, machine learning and artificial intelligence, as well as amateur seeking to better understand those models.\n\n**Out-of-scope use cases** The model should not be used on downstream applications without further risk evaluation and mitigation. The model should not be used to intentionally create or disseminate audio pieces that create hostile or alienating environments for people. This includes generating audio that people would foreseeably find disturbing, distressing, or offensive; or content that propagates historical or current stereotypes.\n\n## Metrics\n\n**Models performance measures:** We used the following objective measure to evaluate the model on a standard audio benchmark:\n- Frechet Audio Distance computed on features extracted from a pre-trained audio classifier (VGGish)\n- Kullback-Leibler Divergence on label distributions extracted from a pre-trained audio classifier (PaSST)\n\nAdditionally, we run qualitative studies with human participants, evaluating the performance of the model with the following axes:\n- Overall quality of the audio samples;\n- Text relevance to the provided text input;\n\nMore details on performance measures and human studies can be found in the paper.\n\n**Decision thresholds:** Not applicable.\n\n## Evaluation datasets\n\nThe model was evaluated on the [AudioCaps benchmark](https://audiocaps.github.io/).\n\n## Training datasets\n\nThe model was trained on the following data sources: a subset of AudioSet (Gemmeke et al., 2017), [BBC sound effects](https://sound-effects.bbcrewind.co.uk/), AudioCaps (Kim et al., 2019), Clotho v2 (Drossos et al., 2020), VGG-Sound (Chen et al., 2020), FSD50K (Fonseca et al., 2021), [Free To Use Sounds](https://www.freetousesounds.com/all-in-one-bundle/), [Sonniss Game Effects](https://sonniss.com/gameaudiogdc), [WeSoundEffects](https://wesoundeffects.com/we-sound-effects-bundle-2020/), [Paramount Motion - Odeon Cinematic Sound Effects](https://www.paramountmotion.com/odeon-sound-effects).\n\n## Evaluation results\n\nBelow are the objective metrics obtained with the released model on AudioCaps (consisting of 10-second long samples). Note that the model differs from the original AudioGen model introduced in the paper, hence the difference in the metrics.\n\n| Model | Frechet Audio Distance | KLD | Text consistency |\n|---|---|---|---|\n| facebook/audiogen-medium | 1.77 | 1.58 | 0.30 |\n\nMore information can be found in the paper [AudioGen: Textually Guided Audio Generation][audiogen], in the Experiments section.\n\n## Limitations and biases\n\n**Limitations:**\n- The model is not able to generate realistic vocals.\n- The model has been trained with English descriptions and will not perform as well in other languages.\n- It is sometimes difficult to assess what types of text descriptions provide the best generations. Prompt engineering may be required to obtain satisfying results.\n\n**Biases:** The datasets used for training may be lacking of diversity and are not representative of all possible sound events. The generated samples from the model will reflect the biases from the training data.\n\n**Risks and harms:** Biases and limitations of the model may lead to generation of samples that may be considered as biased, inappropriate or offensive. We believe that providing the code to reproduce the research and train new models will allow to broaden the application to new and more representative data.\n\n**Use cases:** Users must be aware of the biases, limitations and risks of the model. AudioGen is a model developed for artificial intelligence research on audio generation. As such, it should not be used for downstream applications without further investigation and mitigation of risks.\n\n[musicgen]: https://arxiv.org/abs/2306.05284\n[audiogen]: https://arxiv.org/abs/2209.15352\n" }, { "path": "model_cards/JASCO_MODEL_CARD.md", "content": "## Model details\n\n**Organization developing the model:** The FAIR team of Meta AI.\n\n**Model date:** JASCO was trained in November 2024.\n\n**Model version:** This is the version 1 of the model.\n\n**Model type:** JASCO consists of an EnCodec model for audio tokenization, and a flow-matching model based on the transformer architecture for music modeling.\nThe model comes in different sizes: 400M and 1B; and currently have a two variant: text-to-music + {chords, drums} controls and text-to-music + {chords, drums, melody} controls.\nJASCO is trained with condition dropout and could be used for inference with dropped conditions.\n \n**Paper or resources for more information:** More information can be found in the paper [Joint Audio And Symbolic Conditioning for Temporally Controlled Text-To-Music Generation][arxiv].\n\n**Citation details:** \n\nCode was implemented by Or Tal and Alon Ziv.\n\n```\n@misc{tal2024joint,\n title={Joint Audio and Symbolic Conditioning for Temporally Controlled Text-to-Music Generation}, \n author={Or Tal and Alon Ziv and Itai Gat and Felix Kreuk and Yossi Adi},\n year={2024},\n eprint={2406.10970},\n archivePrefix={arXiv},\n primaryClass={cs.SD}\n}\n```\n\n**License:** Code is released under MIT, model weights are released under CC-BY-NC 4.0.\n\n**Where to send questions or comments about the model:** Questions and comments about JASCO can be sent via the [GitHub repository](https://github.com/facebookresearch/audiocraft) of the project, or by opening an issue.\n\n## Intended use\n**Primary intended use:** The primary use of JASCO is research on AI-based music generation, including:\n\n- Research efforts, such as probing and better understanding the limitations of generative models to further improve the state of science\n- Generation of music guided by text and (opt) local controls, to understand current abilities of generative AI models by machine learning amateurs\n\n**Primary intended users:** The primary intended users of the model are researchers in audio, machine learning and artificial intelligence, as well as amateur seeking to better understand those models.\n\n**Out-of-scope use cases:** The model should not be used on downstream applications without further risk evaluation and mitigation. The model should not be used to intentionally create or disseminate music pieces that create hostile or alienating environments for people. This includes generating music that people would foreseeably find disturbing, distressing, or offensive; or content that propagates historical or current stereotypes.\n\n## Metrics\n\n**Models performance measures:** We used the following objective measure to evaluate the model on a standard music benchmark:\n\n- Frechet Audio Distance computed on features extracted from a pre-trained audio classifier (VGGish).\n- CLAP Score between audio embedding and text embedding extracted from a pre-trained CLAP model.\n- Melody cosine similarity - pairwise comparison of chromagram extracted from refrence and generated waveforms.\n- Onset F1 - pairwise comparison of onsets extracted from refrence and generated waveforms.\n- Chords Intersection over union (IOU) - pairwise comparison of symbolic chords extracted from refrence and generated waveforms.\n\nAdditionally, we run qualitative studies with human participants, evaluating the performance of the model with the following axes:\n\n- Overall quality of the music samples;\n- Text relevance to the provided text input;\n- Melody match w.r.t reference signal;\n- Drum beat match w.r.t reference signal;\n\nMore details on performance measures and human studies can be found in the [paper][arxiv].\n\n**Decision thresholds:** Not applicable.\n\n## Evaluation datasets\n\nThe model was evaluated on the [MusicCaps benchmark](https://www.kaggle.com/datasets/googleai/musiccaps) and on an in-domain held-out evaluation set, with no artist overlap with the training set.\n\n## Training datasets\n\nThe model was trained on licensed data using the following sources: the [Meta Music Initiative Sound Collection](https://www.fb.com/sound), [Shutterstock music collection](https://www.shutterstock.com/music) and the [Pond5 music collection](https://www.pond5.com/). See the paper for more details about the training set and corresponding preprocessing.\n\n## Evaluation results\n\nBelow are the objective metrics obtained on MusicCaps with the released model. \n\nText-to-music with temporal controls\n\n| Model | Frechet Audio Distance | Text Consistency | Chord IOU | Onset F1 | Melody Cosine Similarity |\n|---|---|---|---|---|---|\n| facebook/jasco-chords-drums-400M | 5.866 | 0.284 | 0.588 | 0.328 | 0.096 |\n| facebook/jasco-chords-drums-1B | 5.587 | 0.291 | 0.589 | 0.331 | 0.097 |\n| facebook/jasco-chords-drums-melody-400M | 4.730 | 0.317 | 0.689 | 0.379 | 0.423 |\n| facebook/jasco-chords-drums-melody-1B | 5.098 | 0.313 | 0.690 | 0.378 | 0.427 |\n\nNote: reccommanded CFG coefficient ratio stands at 1:2 - 'all':'text', results for chords-drums-melody were sampled with all: 1.75, text: 3.5 \n\nText-to-music w.o temporal controls (dropped)\n\n\n| Model | Frechet Audio Distance | Text Consistency | Chord IOU | Onset F1 | Melody Cosine Similarity |\n|---|---|---|---|---|---|\n| facebook/jasco-chords-drums-400M | 5.648 | 0.272 | 0.070 | 0.204 | 0.093 |\n| facebook/jasco-chords-drums-1B | 5.602 | 0.281 | 0.071 | 0.214 | 0.093 |\n| facebook/jasco-chords-drums-melody-400M | 5.816 | 0.293 | 0.091 | 0.203 | 0.098 |\n| facebook/jasco-chords-drums-melody-1B | 5.470 | 0.297 | 0.097 | 0.208 | 0.097 |\n\n## Limitations and biases\n\n**Data:** The data sources used to train the model are created by music professionals and covered by legal agreements with the right holders. The model is trained on ~16k hours of data, we believe that scaling the model on larger datasets can further improve the performance of the model.\n\n**Mitigations:** \nPre-trained models were used to obtain pseudo symbolic supervision. Refer to **Data Preprocessing** section in [Jasco's docs](../docs/JASCO.md)\n\n**Limitations:**\n\n- The model is not able to generate realistic vocals.\n- The model has been trained with English descriptions and will not perform as well in other languages.\n- The model does not perform equally well for all music styles and cultures.\n- It is sometimes difficult to assess what types of text descriptions provide the best generations. Prompt engineering and experimentation with classifier free guidance coefficients may be required to obtain satisfying results.\n- Model could be sensitive to CFG coefficients as melody introduces a strong bias that would require higher text coefficient during generation, some hyper-parameter search could be necessary to obtain desired results.\n\n**Biases:** The source of data is potentially lacking diversity and all music cultures are not equally represented in the dataset. The model may not perform equally well on the wide variety of music genres that exists. The generated samples from the model will reflect the biases from the training data. Further work on this model should include methods for balanced and just representations of cultures, for example, by scaling the training data to be both diverse and inclusive.\n\n**Risks and harms:** Biases and limitations of the model may lead to generation of samples that may be considered as biased, inappropriate or offensive. We believe that providing the code to reproduce the research and train new models will allow to broaden the application to new and more representative data.\n\n**Use cases:** Users must be aware of the biases, limitations and risks of the model. JASCO is a model developed for artificial intelligence research on controllable music generation. As such, it should not be used for downstream applications without further investigation and mitigation of risks.\n\n## API\n\nWe provide a simple API and pre-trained models:\n- `facebook/jasco-chords-drums-400M`: 400M model, text to music with chords and drums support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-400M)\n- `facebook/jasco-chords-drums-1B`: 1B model, text to music with chords and drums support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-1B)\n- `facebook/jasco-chords-drums-melody-400M`: 400M model, text to music with chords, drums and melody support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-melody-400M)\n- `facebook/jasco-chords-drums-melody-1B`: 1B model, text to music with chords, drums and melody support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-melody-1B)\n\n\nSee after a quick example for using the API.\n\n```python\nfrom audiocraft.models import JASCO\n\nmodel = JASCO.get_pretrained('facebook/jasco-chords-drums-400M', chords_mapping_path='../assets/chord_to_index_mapping.pkl')\n\nmodel.set_generation_params(\n cfg_coef_all=1.5,\n cfg_coef_txt=0.5\n)\n\n# set textual prompt\ntext = \"Strings, woodwind, orchestral, symphony.\"\n\n# define chord progression\nchords = [('C', 0.0), ('D', 2.0), ('F', 4.0), ('Ab', 6.0), ('Bb', 7.0), ('C', 8.0)]\n\n# run inference\noutput = model.generate_music(descriptions=[text], chords=chords, progress=True)\n\naudio_write('output', output.cpu().squeeze(0), model.sample_rate, strategy=\"loudness\", loudness_compressor=True)\n```\n\n[arxiv]: https://arxiv.org/pdf/2406.10970" }, { "path": "model_cards/MAGNET_MODEL_CARD.md", "content": "# MAGNeT Model Card\n\n## Model details\n\n**Organization developing the model:** The FAIR team of Meta AI.\n\n**Model date:** MAGNeT was trained between November 2023 and January 2024.\n\n**Model version:** This is the version 1 of the model.\n\n**Model type:** MAGNeT consists of an EnCodec model for audio tokenization, and a non-autoregressive model based on the transformer architecture for music modeling. The model comes in different sizes: 300M and 1.5B; and two variants: a model trained for text-to-music generation, and a model trained for text-to-sound generation.\n \n**Paper or resources for more information:** More information can be found in the paper [Masked Audio Generation using a Single Non-Autoregressive Transformer][arxiv].\n\n**Citation details:** See [our paper][arxiv]\n\n**License:** Code is released under MIT, model weights are released under CC-BY-NC 4.0.\n\n**Where to send questions or comments about the model:** Questions and comments about MAGNeT can be sent via the [GitHub repository](https://github.com/facebookresearch/audiocraft) of the project, or by opening an issue.\n\n## Intended use\n**Primary intended use:** The primary use of MAGNeT is research on AI-based music generation, including:\n\n- Research efforts, such as probing and better understanding the limitations of generative models to further improve the state of science\n- Generation of music guided by text to understand current abilities of generative AI models by machine learning amateurs\n\n**Primary intended users:** The primary intended users of the model are researchers in audio, machine learning and artificial intelligence, as well as amateur seeking to better understand those models.\n\n**Out-of-scope use cases:** The model should not be used on downstream applications without further risk evaluation and mitigation. The model should not be used to intentionally create or disseminate music pieces that create hostile or alienating environments for people. This includes generating music that people would foreseeably find disturbing, distressing, or offensive; or content that propagates historical or current stereotypes.\n\n## Metrics\n\n**Models performance measures:** We used the following objective measure to evaluate the model on a standard music benchmark:\n\n- Frechet Audio Distance computed on features extracted from a pre-trained audio classifier (VGGish)\n- Kullback-Leibler Divergence on label distributions extracted from a pre-trained audio classifier (PaSST)\n- CLAP Score between audio embedding and text embedding extracted from a pre-trained CLAP model\n\nAdditionally, we run qualitative studies with human participants, evaluating the performance of the model with the following axes:\n\n- Overall quality of the music samples;\n- Text relevance to the provided text input;\n\nMore details on performance measures and human studies can be found in the paper.\n\n**Decision thresholds:** Not applicable.\n\n## Evaluation datasets\n\nThe model was evaluated on the [MusicCaps benchmark](https://www.kaggle.com/datasets/googleai/musiccaps) and on an in-domain held-out evaluation set, with no artist overlap with the training set.\n\n## Training datasets\n\nThe model was trained on licensed data using the following sources: the [Meta Music Initiative Sound Collection](https://www.fb.com/sound), [Shutterstock music collection](https://www.shutterstock.com/music) and the [Pond5 music collection](https://www.pond5.com/). See the paper for more details about the training set and corresponding preprocessing.\n\n## Evaluation results\n\nBelow are the objective metrics obtained on MusicCaps with the released model. Note that for the publicly released models, we used the state-of-the-art music source separation method, namely the open source [Hybrid Transformer for Music Source Separation](https://github.com/facebookresearch/demucs) (HT-Demucs), in order to keep only instrumental tracks. This explains the difference in objective metrics with the models used in the paper.\n\n| Model | Frechet Audio Distance | KLD | Text Consistency |\n|---|---|---|---|\n| **facebook/magnet-small-10secs** | 4.22 | 1.11 | 0.28 |\n| facebook/magnet-medium-10secs | 4.61 | 1.14 | 0.28 |\n| facebook/magnet-small-30secs | 4.35 | 1.17 | 0.28 |\n| facebook/magnet-medium-30secs | 4.63 | 1.20 | 0.28 |\n\nMore information can be found in the paper [Masked Audio Generation using a Single Non-Autoregressive Transformer][arxiv], in the Results section.\n\n## Limitations and biases\n\n**Data:** The data sources used to train the model are created by music professionals and covered by legal agreements with the right holders. The model is trained on 16K hours of data, we believe that scaling the model on larger datasets can further improve the performance of the model.\n\n**Mitigations:** Tracks that include vocals have been removed from the data source using corresponding tags, and using a state-of-the-art music source separation method, namely using the open source [Hybrid Transformer for Music Source Separation](https://github.com/facebookresearch/demucs) (HT-Demucs).\n\n**Limitations:**\n\n- The model is not able to generate realistic vocals.\n- The model has been trained with English descriptions and will not perform as well in other languages.\n- The model does not perform equally well for all music styles and cultures.\n- The model sometimes generates end of songs, collapsing to silence.\n- It is sometimes difficult to assess what types of text descriptions provide the best generations. Prompt engineering may be required to obtain satisfying results.\n\n**Biases:** The source of data is potentially lacking diversity and all music cultures are not equally represented in the dataset. The model may not perform equally well on the wide variety of music genres that exists. The generated samples from the model will reflect the biases from the training data. Further work on this model should include methods for balanced and just representations of cultures, for example, by scaling the training data to be both diverse and inclusive.\n\n**Risks and harms:** Biases and limitations of the model may lead to generation of samples that may be considered as biased, inappropriate or offensive. We believe that providing the code to reproduce the research and train new models will allow to broaden the application to new and more representative data.\n\n**Use cases:** Users must be aware of the biases, limitations and risks of the model. MAGNeT is a model developed for artificial intelligence research on controllable music generation. As such, it should not be used for downstream applications without further investigation and mitigation of risks.\n\n[arxiv]: https://arxiv.org/abs/2401.04577\n\n## Audio-MAGNeT - Sound-effect generation models\n\n### Training datasets\n\nThe audio-MAGNeT models were trained on the following data sources: a subset of AudioSet (Gemmeke et al., 2017), [BBC sound effects](https://sound-effects.bbcrewind.co.uk/), AudioCaps (Kim et al., 2019), Clotho v2 (Drossos et al., 2020), VGG-Sound (Chen et al., 2020), FSD50K (Fonseca et al., 2021), [Free To Use Sounds](https://www.freetousesounds.com/all-in-one-bundle/), [Sonniss Game Effects](https://sonniss.com/gameaudiogdc), [WeSoundEffects](https://wesoundeffects.com/we-sound-effects-bundle-2020/), [Paramount Motion - Odeon Cinematic Sound Effects](https://www.paramountmotion.com/odeon-sound-effects).\n\n\n### Evaluation datasets\n\nThe audio-magnet models (sound effect generation) were evaluated on the [AudioCaps benchmark](https://audiocaps.github.io/).\n\n### Evaluation results\n\nBelow are the objective metrics obtained with the released audio-magnet models on AudioCaps (consisting of 10-second long samples). \n\n| Model | Frechet Audio Distance | KLD |\n|---|---|---|\n| facebook/audio-magnet-small | 3.21 | 1.42 |\n| facebook/audio-magnet-medium | 2.32 | 1.64 |\n" }, { "path": "model_cards/MUSICGEN_MODEL_CARD.md", "content": "# MusicGen Model Card\n\n## Model details\n\n**Organization developing the model:** The FAIR team of Meta AI.\n\n**Model date:** MusicGen was trained between April 2023 and May 2023.\n\n**Model version:** This is the version 1 of the model.\n\n**Model type:** MusicGen consists of an EnCodec model for audio tokenization, an auto-regressive language model based on the transformer architecture for music modeling. The model comes in different sizes: 300M, 1.5B and 3.3B parameters ; and two variants: a model trained for text-to-music generation task and a model trained for melody-guided music generation.\n\n**Paper or resources for more information:** More information can be found in the paper [Simple and Controllable Music Generation][arxiv].\n\n**Citation details:** See [our paper][arxiv]\n\n**License:** Code is released under MIT, model weights are released under CC-BY-NC 4.0.\n\n**Where to send questions or comments about the model:** Questions and comments about MusicGen can be sent via the [GitHub repository](https://github.com/facebookresearch/audiocraft) of the project, or by opening an issue.\n\n## Intended use\n**Primary intended use:** The primary use of MusicGen is research on AI-based music generation, including:\n\n- Research efforts, such as probing and better understanding the limitations of generative models to further improve the state of science\n- Generation of music guided by text or melody to understand current abilities of generative AI models by machine learning amateurs\n\n**Primary intended users:** The primary intended users of the model are researchers in audio, machine learning and artificial intelligence, as well as amateur seeking to better understand those models.\n\n**Out-of-scope use cases:** The model should not be used on downstream applications without further risk evaluation and mitigation. The model should not be used to intentionally create or disseminate music pieces that create hostile or alienating environments for people. This includes generating music that people would foreseeably find disturbing, distressing, or offensive; or content that propagates historical or current stereotypes.\n\n## Metrics\n\n**Models performance measures:** We used the following objective measure to evaluate the model on a standard music benchmark:\n\n- Frechet Audio Distance computed on features extracted from a pre-trained audio classifier (VGGish)\n- Kullback-Leibler Divergence on label distributions extracted from a pre-trained audio classifier (PaSST)\n- CLAP Score between audio embedding and text embedding extracted from a pre-trained CLAP model\n\nAdditionally, we run qualitative studies with human participants, evaluating the performance of the model with the following axes:\n\n- Overall quality of the music samples;\n- Text relevance to the provided text input;\n- Adherence to the melody for melody-guided music generation.\n\nMore details on performance measures and human studies can be found in the paper.\n\n**Decision thresholds:** Not applicable.\n\n## Evaluation datasets\n\nThe model was evaluated on the [MusicCaps benchmark](https://www.kaggle.com/datasets/googleai/musiccaps) and on an in-domain held-out evaluation set, with no artist overlap with the training set.\n\n## Training datasets\n\nThe model was trained on licensed data using the following sources: the [Meta Music Initiative Sound Collection](https://www.fb.com/sound), [Shutterstock music collection](https://www.shutterstock.com/music) and the [Pond5 music collection](https://www.pond5.com/). See the paper for more details about the training set and corresponding preprocessing.\n\n## Evaluation results\n\nBelow are the objective metrics obtained on MusicCaps with the released model. Note that for the publicly released models, we had all the datasets go through a state-of-the-art music source separation method, namely using the open source [Hybrid Transformer for Music Source Separation](https://github.com/facebookresearch/demucs) (HT-Demucs), in order to keep only the instrumental part. This explains the difference in objective metrics with the models used in the paper.\n\n| Model | Frechet Audio Distance | KLD | Text Consistency | Chroma Cosine Similarity |\n|---|---|---|---|---|\n| facebook/musicgen-small | 4.88 | 1.42 | 0.27 | - |\n| facebook/musicgen-medium | 5.14 | 1.38 | 0.28 | - |\n| facebook/musicgen-large | 5.48 | 1.37 | 0.28 | - |\n| facebook/musicgen-melody | 4.93 | 1.41 | 0.27 | 0.44 |\n\nMore information can be found in the paper [Simple and Controllable Music Generation][arxiv], in the Results section.\n\n## Limitations and biases\n\n**Data:** The data sources used to train the model are created by music professionals and covered by legal agreements with the right holders. The model is trained on 20K hours of data, we believe that scaling the model on larger datasets can further improve the performance of the model.\n\n**Mitigations:** Vocals have been removed from the data source using corresponding tags, and then using a state-of-the-art music source separation method, namely using the open source [Hybrid Transformer for Music Source Separation](https://github.com/facebookresearch/demucs) (HT-Demucs).\n\n**Limitations:**\n\n- The model is not able to generate realistic vocals.\n- The model has been trained with English descriptions and will not perform as well in other languages.\n- The model does not perform equally well for all music styles and cultures.\n- The model sometimes generates end of songs, collapsing to silence.\n- It is sometimes difficult to assess what types of text descriptions provide the best generations. Prompt engineering may be required to obtain satisfying results.\n\n**Biases:** The source of data is potentially lacking diversity and all music cultures are not equally represented in the dataset. The model may not perform equally well on the wide variety of music genres that exists. The generated samples from the model will reflect the biases from the training data. Further work on this model should include methods for balanced and just representations of cultures, for example, by scaling the training data to be both diverse and inclusive.\n\n**Risks and harms:** Biases and limitations of the model may lead to generation of samples that may be considered as biased, inappropriate or offensive. We believe that providing the code to reproduce the research and train new models will allow to broaden the application to new and more representative data.\n\n**Use cases:** Users must be aware of the biases, limitations and risks of the model. MusicGen is a model developed for artificial intelligence research on controllable music generation. As such, it should not be used for downstream applications without further investigation and mitigation of risks.\n\n## Update: stereo models and large melody.\n\nWe further release a set of stereophonic capable models. Those were fine tuned for 200k updates starting\nfrom the mono models. The training data is otherwise identical and capabilities and limitations are shared with the base modes. The stereo models work by getting 2 streams of tokens from the EnCodec model, and interleaving those using\nthe delay pattern. We also release a mono large model with melody conditioning capabilities. The list of new models\nis as follow:\n\n- facebook/musicgen-stereo-small\n- facebook/musicgen-stereo-medium\n- facebook/musicgen-stereo-large\n- facebook/musicgen-stereo-melody\n- facebook/musicgen-melody-large\n- facebook/musicgen-stereo-melody-large\n\n\n[arxiv]: https://arxiv.org/abs/2306.05284\n" }, { "path": "model_cards/MUSICGEN_STYLE_MODEL_CARD.md", "content": "# MusicGen Model Card\n\n## Model details\n\n**Organization developing the model:** The FAIR team of Meta AI.\n\n**Model date:** MusicGen-Style was trained between November 2023 and February 2024.\n\n**Model version:** This is the version 1 of the model.\n\n**Model type:** MusicGen-Style consists of an EnCodec model for audio tokenization, a 1.5B parameters auto-regressive language model based on the transformer architecture for music modeling conditioned by a text conditioner as well as a style conditioner.\n\n**Paper or resources for more information:** More information can be found in the paper [Audio Conditioning for Music Generation via Discrete Bottleneck Features][arxiv].\n\n**Citation details:** See [our paper][arxiv]\n\n**License:** Code is released under MIT, model weights are released under CC-BY-NC 4.0.\n\n**Where to send questions or comments about the model:** Questions and comments about MusicGen-Style can be sent via the [GitHub repository](https://github.com/facebookresearch/audiocraft) of the project, or by opening an issue.\n\n## Intended use\n**Primary intended use:** The primary use of MusicGen-Style is research on AI-based music generation, including:\n\n- Research efforts, such as probing and better understanding the limitations of generative models to further improve the state of science\n- Generation of music guided by text or style to understand current abilities of generative AI models by machine learning amateurs\n\n**Primary intended users:** The primary intended users of the model are researchers in audio, machine learning and artificial intelligence, as well as amateur seeking to better understand those models.\n\n**Out-of-scope use cases:** The model should not be used on downstream applications without further risk evaluation and mitigation. The model should not be used to intentionally create or disseminate music pieces that create hostile or alienating environments for people. This includes generating music that people would foreseeably find disturbing, distressing, or offensive; or content that propagates historical or current stereotypes.\n\n## Training datasets\n\nThe model was trained on licensed data using the following sources: the [Meta Music Initiative Sound Collection](https://www.fb.com/sound), [Shutterstock music collection](https://www.shutterstock.com/music) and the [Pond5 music collection](https://www.pond5.com/). See the paper for more details about the training set and corresponding preprocessing.\n\n\n## Limitations and biases\n\n**Data:** The data sources used to train the model are created by music professionals and covered by legal agreements with the right holders. The model is trained on 20K hours of data, we believe that scaling the model on larger datasets can further improve the performance of the model.\n\n**Mitigations:** Vocals have been removed from the data source using corresponding tags, and then using a state-of-the-art music source separation method, namely using the open source [Hybrid Transformer for Music Source Separation](https://github.com/facebookresearch/demucs) (HT-Demucs).\n\n**Limitations:**\n\n- The model is not able to generate realistic vocals.\n- The model has been trained with English descriptions and will not perform as well in other languages.\n- The model does not perform equally well for all music styles and cultures.\n- The model sometimes generates end of songs, collapsing to silence.\n- It is sometimes difficult to assess what types of text descriptions provide the best generations. Prompt engineering may be required to obtain satisfying results.\n\n**Biases:** The source of data is potentially lacking diversity and all music cultures are not equally represented in the dataset. The model may not perform equally well on the wide variety of music genres that exists. The generated samples from the model will reflect the biases from the training data. Further work on this model should include methods for balanced and just representations of cultures, for example, by scaling the training data to be both diverse and inclusive.\n\n**Risks and harms:** Biases and limitations of the model may lead to generation of samples that may be considered as biased, inappropriate or offensive. We believe that providing the code to reproduce the research and train new models will allow to broaden the application to new and more representative data.\n\n**Use cases:** Users must be aware of the biases, limitations and risks of the model. MusicGen-Style is a model developed for artificial intelligence research on controllable music generation. As such, it should not be used for downstream applications without further investigation and mitigation of risks.\n\n\n\n[arxiv]: https://arxiv.org/abs/2407.12563\n" }, { "path": "mypy.ini", "content": "[mypy]\n\n[mypy-treetable,torchaudio.*,soundfile,einops.*,av.*,tqdm.*,num2words.*,spacy,xformers.*,scipy,flashy.*,torchmetrics.*,hydra,pesq,demucs.*,huggingface_hub,transformers,dac.*]\nignore_missing_imports = True\n" }, { "path": "requirements.txt", "content": "# please make sure you have already a pytorch install that is cuda enabled!\nav==11.0.0\neinops\nflashy>=0.0.1\nhydra-core>=1.1\nhydra_colorlog\njulius\nnum2words\nnumpy<2.0.0\nsentencepiece\nspacy==3.7.6\ntorch==2.1.0\ntorchaudio>=2.0.0,<2.1.2\nhuggingface_hub\ntqdm\ntransformers>=4.31.0 # need Encodec there.\nxformers<0.0.23\ndemucs\nlibrosa\nsoundfile\ngradio\ntorchmetrics\nencodec\nprotobuf\ntorchvision==0.16.0\ntorchtext==0.16.0\npesq\npystoi\ntorchdiffeq\n" }, { "path": "scripts/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "scripts/chords/build_chord_maps.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\nimport os\nimport pickle\nfrom tqdm import tqdm\nimport argparse\n\n\ndef parse_args():\n parser = argparse.ArgumentParser()\n parser.add_argument('--chords_folder', type=str, required=True,\n help='path to directory containing parsed chords files')\n parser.add_argument('--output_directory', type=str, required=False,\n help='path to output directory to generate code maps to, \\\n if not given - chords_folder would be used', default='')\n parser.add_argument('--path_to_pre_defined_map', type=str, required=False,\n help='for evaluation purpose, use pre-defined chord-to-index map', default='')\n args = parser.parse_args()\n return args\n\n\ndef get_chord_dict(chord_folder: str):\n chord_dict = {}\n distinct_chords = set()\n\n chord_to_index = {} # Mapping between chord and index\n index_counter = 0\n\n for filename in tqdm(os.listdir(chord_folder)):\n if filename.endswith(\".chords\"):\n idx = filename.split(\".\")[0]\n\n with open(os.path.join(chord_folder, filename), \"rb\") as file:\n chord_data = pickle.load(file)\n\n for chord, _ in chord_data:\n distinct_chords.add(chord)\n if chord not in chord_to_index:\n chord_to_index[chord] = index_counter\n index_counter += 1\n\n chord_dict[idx] = chord_data\n chord_to_index[\"UNK\"] = index_counter\n return chord_dict, distinct_chords, chord_to_index\n\n\ndef get_predefined_chord_to_index_map(path_to_chords_to_index_map: str):\n def inner(chord_folder: str):\n chords_to_index = pickle.load(open(path_to_chords_to_index_map, \"rb\"))\n distinct_chords = set(chords_to_index.keys())\n chord_dict = {}\n for filename in tqdm(os.listdir(chord_folder), desc=f'iterating: {chord_folder}'):\n if filename.endswith(\".chords\"):\n idx = filename.split(\".\")[0]\n\n with open(os.path.join(chord_folder, filename), \"rb\") as file:\n chord_data = pickle.load(file)\n\n chord_dict[idx] = chord_data\n return chord_dict, distinct_chords, chords_to_index\n return inner\n\n\nif __name__ == \"__main__\":\n '''This script processes and maps chord data from a directory of parsed chords files,\n generating two output files: a combined chord dictionary and a chord-to-index mapping.'''\n args = parse_args()\n chord_folder = args.chords_folder\n output_dir = args.output_directory\n if output_dir == '':\n output_dir = chord_folder\n func = get_chord_dict\n if args.path_to_pre_defined_map != \"\":\n func = get_predefined_chord_to_index_map(args.path_to_pre_defined_map)\n\n chord_dict, distinct_chords, chord_to_index = func(chord_folder)\n\n # Save the combined chord dictionary as a pickle file\n combined_filename = os.path.join(output_dir, \"combined_chord_dict.pkl\")\n with open(combined_filename, \"wb\") as file:\n pickle.dump(chord_dict, file)\n\n # Save the chord-to-index mapping as a pickle file\n mapping_filename = os.path.join(output_dir, \"chord_to_index_mapping.pkl\")\n with open(mapping_filename, \"wb\") as file:\n pickle.dump(chord_to_index, file)\n\n print(\"Number of distinct chords:\", len(distinct_chords))\n print(\"Chord dictionary:\", chord_to_index)\n" }, { "path": "scripts/chords/extract_chords.py", "content": "# Env - chords_extraction on devfair\n\nimport pickle\nimport argparse\nfrom chord_extractor.extractors import Chordino # type: ignore\nfrom chord_extractor import clear_conversion_cache, LabelledChordSequence # type: ignore\nimport os\nfrom tqdm import tqdm\n\n\ndef parse_args():\n parser = argparse.ArgumentParser()\n parser.add_argument('--src_jsonl_file', type=str, required=True,\n help='abs path to .jsonl file containing list of absolute file paths seperated by new line')\n parser.add_argument('--target_output_dir', type=str, required=True,\n help='target directory to save parsed chord files to, individual files will be saved inside')\n parser.add_argument(\"--override\", action=\"store_true\")\n args = parser.parse_args()\n return args\n\n\ndef save_to_db_cb(tgt_dir: str):\n # Every time one of the files has had chords extracted, receive the chords here\n # along with the name of the original file and then run some logic here, e.g. to\n # save the latest data to DB\n def inner(results: LabelledChordSequence):\n path = results.id.split(\".wav\")\n\n sequence = [(item.chord, item.timestamp) for item in results.sequence]\n\n if len(path) != 2:\n print(\"Something\")\n print(path)\n else:\n file_idx = path[0].split(\"/\")[-1]\n with open(f\"{tgt_dir}/{file_idx}.chords\", \"wb\") as f:\n # dump the object to the file\n pickle.dump(sequence, f)\n return inner\n\n\nif __name__ == \"__main__\":\n '''This script extracts chord data from a list of audio files using the Chordino extractor,\n and saves the extracted chords to individual files in a target directory.'''\n print(\"parsed args\")\n args = parse_args()\n files_to_extract_from = list()\n with open(args.src_jsonl_file, \"r\") as json_file:\n for line in tqdm(json_file.readlines()):\n # fpath = json.loads(line.replace(\"\\n\", \"\"))['path']\n fpath = line.replace(\"\\n\", \"\")\n if not args.override:\n fname = fpath.split(\"/\")[-1].replace(\".wav\", \".chords\")\n if os.path.exists(f\"{args.target_output_dir}/{fname}\"):\n continue\n files_to_extract_from.append(line.replace(\"\\n\", \"\"))\n\n print(f\"num files to parse: {len(files_to_extract_from)}\")\n\n chordino = Chordino()\n\n # Optionally clear cache of file conversions (e.g. wav files that have been converted from midi)\n clear_conversion_cache()\n\n # Run bulk extraction\n res = chordino.extract_many(\n files_to_extract_from,\n callback=save_to_db_cb(args.target_output_dir),\n num_extractors=80,\n num_preprocessors=80,\n max_files_in_cache=400,\n stop_on_error=False,\n )\n" }, { "path": "scripts/chords/job_array_example.sh", "content": "#!/bin/zsh\n#SBATCH --job-name=my_job_array\n#SBATCH --array=0-N # adjust the range of indices as needed\n#SBATCH --output=logs/%A_%a.out # output file name format, this assumes there exists a /logs directory\n#SBATCH --error=logs/%A_%a.err # error file name format, this assumes there exists a /logs directory\n#SBATCH --time=01:00:00 # adjust the time limit as needed\n#SBATCH --nodes=1 # adjust the number of nodes as needed\n#SBATCH --ntasks-per-node=1 # adjust the number of tasks per node as needed\n#SBATCH --cpus-per-task=8 # adjust the number of CPUs per task as needed\n#SBATCH --mem-per-cpu=16G # adjust the memory per CPU as needed\n\n# Load any necessary modules or dependencies\nconda activate your_env\n\n# run extraction of chords in job array\npython scripts/chords/extract_chords.py --src_jsonl_file /path/to/parsed/filepaths_${SLURM_ARRAY_TASK_ID}.jsonl --target_output_dir /target/directory/to/save/chords/to --path_to_pre_defined_map /path/to/predefined/chord_to_index_mapping.pkl\n\n" }, { "path": "scripts/mos.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\n\"\"\"\nTo run this script, from the root of the repo. Make sure to have Flask installed\n\n FLASK_DEBUG=1 FLASK_APP=scripts.mos flask run -p 4567\n # or if you have gunicorn\n gunicorn -w 4 -b 127.0.0.1:8895 -t 120 'scripts.mos:app' --access-logfile -\n\n\"\"\"\nfrom collections import defaultdict\nfrom functools import wraps\nfrom hashlib import sha1\nimport json\nimport math\nfrom pathlib import Path\nimport random\nimport typing as tp\n\nfrom flask import Flask, redirect, render_template, request, session, url_for\n\nfrom audiocraft import train\nfrom audiocraft.utils.samples.manager import get_samples_for_xps\n\n\nSAMPLES_PER_PAGE = 8\nMAX_RATING = 5\nstorage = Path(train.main.dora.dir / 'mos_storage')\nstorage.mkdir(exist_ok=True)\nsurveys = storage / 'surveys'\nsurveys.mkdir(exist_ok=True)\nmagma_root = Path(train.__file__).parent.parent\napp = Flask('mos', static_folder=str(magma_root / 'scripts/static'),\n template_folder=str(magma_root / 'scripts/templates'))\napp.secret_key = b'audiocraft makes the best songs'\n\n\ndef normalize_path(path: Path):\n \"\"\"Just to make path a bit nicer, make them relative to the Dora root dir.\n \"\"\"\n path = path.resolve()\n dora_dir = train.main.dora.dir.resolve() / 'xps'\n return path.relative_to(dora_dir)\n\n\ndef get_full_path(normalized_path: Path):\n \"\"\"Revert `normalize_path`.\n \"\"\"\n return train.main.dora.dir.resolve() / 'xps' / normalized_path\n\n\ndef get_signature(xps: tp.List[str]):\n \"\"\"Return a signature for a list of XP signatures.\n \"\"\"\n return sha1(json.dumps(xps).encode()).hexdigest()[:10]\n\n\ndef ensure_logged(func):\n \"\"\"Ensure user is logged in.\n \"\"\"\n @wraps(func)\n def _wrapped(*args, **kwargs):\n user = session.get('user')\n if user is None:\n return redirect(url_for('login', redirect_to=request.url))\n return func(*args, **kwargs)\n return _wrapped\n\n\n@app.route('/login', methods=['GET', 'POST'])\ndef login():\n \"\"\"Login user if not already, then redirect.\n \"\"\"\n user = session.get('user')\n if user is None:\n error = None\n if request.method == 'POST':\n user = request.form['user']\n if not user:\n error = 'User cannot be empty'\n if user is None or error:\n return render_template('login.html', error=error)\n assert user\n session['user'] = user\n redirect_to = request.args.get('redirect_to')\n if redirect_to is None:\n redirect_to = url_for('index')\n return redirect(redirect_to)\n\n\n@app.route('/', methods=['GET', 'POST'])\n@ensure_logged\ndef index():\n \"\"\"Offer to create a new study.\n \"\"\"\n errors = []\n if request.method == 'POST':\n xps_or_grids = [part.strip() for part in request.form['xps'].split()]\n xps = set()\n for xp_or_grid in xps_or_grids:\n xp_path = train.main.dora.dir / 'xps' / xp_or_grid\n if xp_path.exists():\n xps.add(xp_or_grid)\n continue\n grid_path = train.main.dora.dir / 'grids' / xp_or_grid\n if grid_path.exists():\n for child in grid_path.iterdir():\n if child.is_symlink():\n xps.add(child.name)\n continue\n errors.append(f'{xp_or_grid} is neither an XP nor a grid!')\n assert xps or errors\n blind = 'true' if request.form.get('blind') == 'on' else 'false'\n xps = list(xps)\n if not errors:\n signature = get_signature(xps)\n manifest = {\n 'xps': xps,\n }\n survey_path = surveys / signature\n survey_path.mkdir(exist_ok=True)\n with open(survey_path / 'manifest.json', 'w') as f:\n json.dump(manifest, f, indent=2)\n return redirect(url_for('survey', blind=blind, signature=signature))\n return render_template('index.html', errors=errors)\n\n\n@app.route('/survey/', methods=['GET', 'POST'])\n@ensure_logged\ndef survey(signature):\n success = request.args.get('success', False)\n seed = int(request.args.get('seed', 4321))\n blind = request.args.get('blind', 'false') in ['true', 'on', 'True']\n exclude_prompted = request.args.get('exclude_prompted', 'false') in ['true', 'on', 'True']\n exclude_unprompted = request.args.get('exclude_unprompted', 'false') in ['true', 'on', 'True']\n max_epoch = int(request.args.get('max_epoch', '-1'))\n survey_path = surveys / signature\n assert survey_path.exists(), survey_path\n\n user = session['user']\n result_folder = survey_path / 'results'\n result_folder.mkdir(exist_ok=True)\n result_file = result_folder / f'{user}_{seed}.json'\n\n with open(survey_path / 'manifest.json') as f:\n manifest = json.load(f)\n\n xps = [train.main.get_xp_from_sig(xp) for xp in manifest['xps']]\n names, ref_name = train.main.get_names(xps)\n\n samples_kwargs = {\n 'exclude_prompted': exclude_prompted,\n 'exclude_unprompted': exclude_unprompted,\n 'max_epoch': max_epoch,\n }\n matched_samples = get_samples_for_xps(xps, epoch=-1, **samples_kwargs) # fetch latest epoch\n models_by_id = {\n id: [{\n 'xp': xps[idx],\n 'xp_name': names[idx],\n 'model_id': f'{xps[idx].sig}-{sample.id}',\n 'sample': sample,\n 'is_prompted': sample.prompt is not None,\n 'errors': [],\n } for idx, sample in enumerate(samples)]\n for id, samples in matched_samples.items()\n }\n experiments = [\n {'xp': xp, 'name': names[idx], 'epoch': list(matched_samples.values())[0][idx].epoch}\n for idx, xp in enumerate(xps)\n ]\n\n keys = list(matched_samples.keys())\n keys.sort()\n rng = random.Random(seed)\n rng.shuffle(keys)\n model_ids = keys[:SAMPLES_PER_PAGE]\n\n if blind:\n for key in model_ids:\n rng.shuffle(models_by_id[key])\n\n ok = True\n if request.method == 'POST':\n all_samples_results = []\n for id in model_ids:\n models = models_by_id[id]\n result = {\n 'id': id,\n 'is_prompted': models[0]['is_prompted'],\n 'models': {}\n }\n all_samples_results.append(result)\n for model in models:\n rating = request.form[model['model_id']]\n if rating:\n rating = int(rating)\n assert rating <= MAX_RATING and rating >= 1\n result['models'][model['xp'].sig] = rating\n model['rating'] = rating\n else:\n ok = False\n model['errors'].append('Please rate this model.')\n if ok:\n result = {\n 'results': all_samples_results,\n 'seed': seed,\n 'user': user,\n 'blind': blind,\n 'exclude_prompted': exclude_prompted,\n 'exclude_unprompted': exclude_unprompted,\n }\n print(result)\n with open(result_file, 'w') as f:\n json.dump(result, f)\n seed = seed + 1\n return redirect(url_for(\n 'survey', signature=signature, blind=blind, seed=seed,\n exclude_prompted=exclude_prompted, exclude_unprompted=exclude_unprompted,\n max_epoch=max_epoch, success=True))\n\n ratings = list(range(1, MAX_RATING + 1))\n return render_template(\n 'survey.html', ratings=ratings, blind=blind, seed=seed, signature=signature, success=success,\n exclude_prompted=exclude_prompted, exclude_unprompted=exclude_unprompted, max_epoch=max_epoch,\n experiments=experiments, models_by_id=models_by_id, model_ids=model_ids, errors=[],\n ref_name=ref_name, already_filled=result_file.exists())\n\n\n@app.route('/audio/')\ndef audio(path: str):\n full_path = Path('/') / path\n assert full_path.suffix in [\".mp3\", \".wav\"]\n return full_path.read_bytes(), {'Content-Type': 'audio/mpeg'}\n\n\ndef mean(x):\n return sum(x) / len(x)\n\n\ndef std(x):\n m = mean(x)\n return math.sqrt(sum((i - m)**2 for i in x) / len(x))\n\n\n@app.route('/results/')\n@ensure_logged\ndef results(signature):\n\n survey_path = surveys / signature\n assert survey_path.exists(), survey_path\n result_folder = survey_path / 'results'\n result_folder.mkdir(exist_ok=True)\n\n # ratings per model, then per user.\n ratings_per_model = defaultdict(list)\n users = []\n for result_file in result_folder.iterdir():\n if result_file.suffix != '.json':\n continue\n with open(result_file) as f:\n results = json.load(f)\n users.append(results['user'])\n for result in results['results']:\n for sig, rating in result['models'].items():\n ratings_per_model[sig].append(rating)\n\n fmt = '{:.2f}'\n models = []\n for model in sorted(ratings_per_model.keys()):\n ratings = ratings_per_model[model]\n\n models.append({\n 'sig': model,\n 'samples': len(ratings),\n 'mean_rating': fmt.format(mean(ratings)),\n # the value 1.96 was probably chosen to achieve some\n # confidence interval assuming gaussianity.\n 'std_rating': fmt.format(1.96 * std(ratings) / len(ratings)**0.5),\n })\n return render_template('results.html', signature=signature, models=models, users=users)\n" }, { "path": "scripts/resample_dataset.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Resampling script.\n\"\"\"\nimport argparse\nfrom pathlib import Path\nimport shutil\nimport typing as tp\n\nimport submitit\nimport tqdm\n\nfrom audiocraft.data.audio import audio_read, audio_write\nfrom audiocraft.data.audio_dataset import load_audio_meta, find_audio_files\nfrom audiocraft.data.audio_utils import convert_audio\nfrom audiocraft.environment import AudioCraftEnvironment\n\n\ndef read_txt_files(path: tp.Union[str, Path]):\n with open(args.files_path) as f:\n lines = [line.rstrip() for line in f]\n print(f\"Read {len(lines)} in .txt\")\n lines = [line for line in lines if Path(line).suffix not in ['.json', '.txt', '.csv']]\n print(f\"Filtered and keep {len(lines)} from .txt\")\n return lines\n\n\ndef read_egs_files(path: tp.Union[str, Path]):\n path = Path(path)\n if path.is_dir():\n if (path / 'data.jsonl').exists():\n path = path / 'data.jsonl'\n elif (path / 'data.jsonl.gz').exists():\n path = path / 'data.jsonl.gz'\n else:\n raise ValueError(\"Don't know where to read metadata from in the dir. \"\n \"Expecting either a data.jsonl or data.jsonl.gz file but none found.\")\n meta = load_audio_meta(path)\n return [m.path for m in meta]\n\n\ndef process_dataset(args, n_shards: int, node_index: int, task_index: tp.Optional[int] = None):\n if task_index is None:\n env = submitit.JobEnvironment()\n task_index = env.global_rank\n shard_index = node_index * args.tasks_per_node + task_index\n\n if args.files_path is None:\n lines = [m.path for m in find_audio_files(args.root_path, resolve=False, progress=True, workers=8)]\n else:\n files_path = Path(args.files_path)\n if files_path.suffix == '.txt':\n print(f\"Reading file list from .txt file: {args.files_path}\")\n lines = read_txt_files(args.files_path)\n else:\n print(f\"Reading file list from egs: {args.files_path}\")\n lines = read_egs_files(args.files_path)\n\n total_files = len(lines)\n print(\n f\"Total of {total_files} processed with {n_shards} shards. \" +\n f\"Current idx = {shard_index} -> {total_files // n_shards} files to process\"\n )\n for idx, line in tqdm.tqdm(enumerate(lines)):\n\n # skip if not part of this shard\n if idx % n_shards != shard_index:\n continue\n\n path = str(AudioCraftEnvironment.apply_dataset_mappers(line))\n root_path = str(args.root_path)\n if not root_path.endswith('/'):\n root_path += '/'\n assert path.startswith(str(root_path)), \\\n f\"Mismatch between path and provided root: {path} VS {root_path}\"\n\n try:\n metadata_path = Path(path).with_suffix('.json')\n out_path = args.out_path / path[len(root_path):]\n out_metadata_path = out_path.with_suffix('.json')\n out_done_token = out_path.with_suffix('.done')\n\n # don't reprocess existing files\n if out_done_token.exists():\n continue\n\n print(idx, out_path, path)\n mix, sr = audio_read(path)\n mix_channels = args.channels if args.channels is not None and args.channels > 0 else mix.size(0)\n # enforce simple stereo\n out_channels = mix_channels\n if out_channels > 2:\n print(f\"Mix has more than two channels: {out_channels}, enforcing 2 channels\")\n out_channels = 2\n out_sr = args.sample_rate if args.sample_rate is not None else sr\n out_wav = convert_audio(mix, sr, out_sr, out_channels)\n audio_write(out_path.with_suffix(''), out_wav, sample_rate=out_sr,\n format=args.format, normalize=False, strategy='clip')\n if metadata_path.exists():\n shutil.copy(metadata_path, out_metadata_path)\n else:\n print(f\"No metadata found at {str(metadata_path)}\")\n out_done_token.touch()\n except Exception as e:\n print(f\"Error processing file line: {line}, {e}\")\n\n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser(description=\"Resample dataset with SLURM.\")\n parser.add_argument(\n \"--log_root\",\n type=Path,\n default=Path.home() / 'tmp' / 'resample_logs',\n )\n parser.add_argument(\n \"--files_path\",\n type=Path,\n help=\"List of files to process, either .txt (one file per line) or a jsonl[.gz].\",\n )\n parser.add_argument(\n \"--root_path\",\n type=Path,\n required=True,\n help=\"When rewriting paths, this will be the prefix to remove.\",\n )\n parser.add_argument(\n \"--out_path\",\n type=Path,\n required=True,\n help=\"When rewriting paths, `root_path` will be replaced by this.\",\n )\n parser.add_argument(\"--xp_name\", type=str, default=\"shutterstock\")\n parser.add_argument(\n \"--nodes\",\n type=int,\n default=4,\n )\n parser.add_argument(\n \"--tasks_per_node\",\n type=int,\n default=20,\n )\n parser.add_argument(\n \"--cpus_per_task\",\n type=int,\n default=4,\n )\n parser.add_argument(\n \"--memory_gb\",\n type=int,\n help=\"Memory in GB.\"\n )\n parser.add_argument(\n \"--format\",\n type=str,\n default=\"wav\",\n )\n parser.add_argument(\n \"--sample_rate\",\n type=int,\n default=32000,\n )\n parser.add_argument(\n \"--channels\",\n type=int,\n )\n parser.add_argument(\n \"--partition\",\n default='learnfair',\n )\n parser.add_argument(\"--qos\")\n parser.add_argument(\"--account\")\n parser.add_argument(\"--timeout\", type=int, default=4320)\n parser.add_argument('--debug', action='store_true', help='debug mode (local run)')\n args = parser.parse_args()\n n_shards = args.tasks_per_node * args.nodes\n if args.files_path is None:\n print(\"Warning: --files_path not provided, not recommended when processing more than 10k files.\")\n if args.debug:\n print(\"Debugging mode\")\n process_dataset(args, n_shards=n_shards, node_index=0, task_index=0)\n else:\n\n log_folder = Path(args.log_root) / args.xp_name / '%j'\n print(f\"Logging to: {log_folder}\")\n log_folder.parent.mkdir(parents=True, exist_ok=True)\n executor = submitit.AutoExecutor(folder=str(log_folder))\n if args.qos:\n executor.update_parameters(slurm_partition=args.partition, slurm_qos=args.qos, slurm_account=args.account)\n else:\n executor.update_parameters(slurm_partition=args.partition)\n executor.update_parameters(\n slurm_job_name=args.xp_name, timeout_min=args.timeout,\n cpus_per_task=args.cpus_per_task, tasks_per_node=args.tasks_per_node, nodes=1)\n if args.memory_gb:\n executor.update_parameters(mem=f'{args.memory_gb}GB')\n jobs = []\n with executor.batch():\n for node_index in range(args.nodes):\n job = executor.submit(process_dataset, args, n_shards=n_shards, node_index=node_index)\n jobs.append(job)\n for job in jobs:\n print(f\"Waiting on job {job.job_id}\")\n job.results()\n" }, { "path": "scripts/static/style.css", "content": "body {\n background-color: #fbfbfb;\n margin: 0;\n}\n\nselect, input {\n font-size: 1em;\n max-width: 100%;\n}\n\n.xp_name {\n font-family: monospace;\n}\n\n.simple_form {\n background-color: #dddddd;\n padding: 1em;\n margin: 0.5em;\n}\n\ntextarea {\n margin-top: 0.5em;\n margin-bottom: 0.5em;\n}\n\n.rating {\n background-color: grey;\n padding-top: 5px;\n padding-bottom: 5px;\n padding-left: 8px;\n padding-right: 8px;\n margin-right: 2px;\n cursor:pointer;\n}\n\n.rating_selected {\n background-color: purple;\n}\n\n.content {\n font-family: sans-serif;\n background-color: #f6f6f6;\n padding: 40px;\n margin: 0 auto;\n max-width: 1000px;\n}\n\n.track label {\n padding-top: 10px;\n padding-bottom: 10px;\n}\n.track {\n padding: 15px;\n margin: 5px;\n background-color: #c8c8c8;\n}\n\n.submit-big {\n width:400px;\n height:30px;\n font-size: 20px;\n}\n\n.error {\n color: red;\n}\n\n.ratings {\n margin-left: 10px;\n}\n\n.important {\n font-weight: bold;\n}\n\n.survey {\n margin-bottom: 100px;\n}\n\n.success {\n color: #25901b;\n font-weight: bold;\n}\n.warning {\n color: #8a1f19;\n font-weight: bold;\n}\n.track>section {\n display: flex;\n align-items: center;\n}\n\n.prompt {\n display: flex;\n align-items: center;\n}\n\n.track>section>div {\n padding-left: 10px;\n}\n\naudio {\n max-width: 280px;\n max-height: 40px;\n margin-left: 10px;\n margin-right: 10px;\n}\n\n.special {\n font-weight: bold;\n color: #2c2c2c;\n}\n\n" }, { "path": "scripts/templates/base.html", "content": "\n\n\n {% block head %}\n \n \n AudioCraft — MOS\n {% endblock %}\n\n\n
\n\t

AudioCraft — MOS

\n {% block content %}{% endblock %}\n
\n\n\n" }, { "path": "scripts/templates/index.html", "content": "{% extends \"base.html\" %}\n{% block content %}\n\n

\n Welcome {{session['user']}} to the internal MOS assistant for AudioCraft.\n You can create custom surveys between your models, that you can\n evaluate yourself, or with the help of your teammates, by simply\n sharing a link!\n

\n\n{% for error in errors %}\n

{{error}}

\n{% endfor %}\n
\n
\n
\n \n
\n
\n \n
\n \n\n\n{% endblock %}\n" }, { "path": "scripts/templates/login.html", "content": "{% extends \"base.html\" %}\n{% block content %}\n\n

\n You must identify yourself first! We use a highly secured protocol\n where you just decide your username, and that's it. No password, no encryption,\n just pure trust.\n

\n\n{% if error %}\n

{{error}}

\n{% endif %}\n\n \n \n\n\n{% endblock %}\n" }, { "path": "scripts/templates/results.html", "content": "{% extends \"base.html\" %}\n{% block content %}\n\n

Results for survey #{{signature}}

\n

Checkout the survey page for details on the models.

\n

The following users voted:\n\t{% for user in users %}\n\t\t{{user}}\n\t{% endfor %}\n\n{% for model in models %}\n

{{model['sig']}} ({{model['samples']}} samples)

\n

Ratings: {{model['mean_rating']}} ± {{model['std_rating']}}

\n\n{% endfor %}\n\n{% endblock %}\n" }, { "path": "scripts/templates/survey.html", "content": "{% extends \"base.html\" %}\n{% block content %}\n

Survey #{{signature}}

\n{% if success %}\n

Your ratings have been saved!\nYou have been moved to the next random seed, if you want\nto keep rating more samples.

\n{% endif %}\n{% if already_filled %}\n

You already rated those samples in the past,\n filling this form will override your previous ratings.\n

\n{% endif %}\n

Welcome {{session['user']}} to the survey #{{signature}}.\nGo to the result page to check the results. Go to the home page to start a new survey.\n

\n\n{% for error in errors %}\n

{{error}}

\n{% endfor %}\n\n{% if not blind %}\n

Base config is: {{ref_name}}

\n

The following experiments are compared:

\n
    \n {% for experiment in experiments %}\n
  • {{experiment.xp.sig}} ({{experiment.epoch}} epochs): {{experiment.name}}
    • \n {% endfor %}\n
\n{% else %}\n

This is a blind experiment, the order of all XPs is shuffled with every sample.

\n{% endif %}\n

The current random seed is {{seed}}. You can change it with the following form, and also update blind/non blind.\n

\n\n \n \n \n \n \n \n
\n\n

Samples

\n
\n
\n{% for id in model_ids %}\n
\n

{{id}}

\n {% for model in models_by_id[id] %}\n {% if loop.index == 1 and model.is_prompted %}\n
\n

Prompt is

\n \n

Ground truth is

\n \n
\n {% endif %}\n {% for err in model['errors'] %}\n

{{err}}

\n {% endfor %}\n
\n {% if not blind %}\n

{{model.xp.sig}}:

\n {% endif %}\n \n

Rating:

\n
\n {% for rating in ratings %}\n {{rating}}\n {% endfor %}\n \n
\n

\n
\n {% endfor %}\n
\n
\n{% endfor %}\n \n\n
\n\n{% endblock %}\n" }, { "path": "setup.cfg", "content": "[pep8]\nmax-line-length = 120\n\n[flake8]\nmax-line-length = 120\n\n[coverage:report]\ninclude = audiocraft/*\nomit =\n audiocraft/environment.py\n audiocraft/solvers/*\n audiocraft/utils/*\n audiocraft/*/loaders.py\n audiocraft/*/builders.py\n" }, { "path": "setup.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom pathlib import Path\n\nfrom setuptools import setup, find_packages\n\n\nNAME = 'audiocraft'\nDESCRIPTION = 'Audio generation research library for PyTorch'\n\nURL = 'https://github.com/facebookresearch/audiocraft'\nAUTHOR = 'FAIR Speech & Audio'\nEMAIL = 'defossez@meta.com, jadecopet@meta.com'\nREQUIRES_PYTHON = '>=3.8.0'\n\nfor line in open('audiocraft/__init__.py'):\n line = line.strip()\n if '__version__' in line:\n context = {}\n exec(line, context)\n VERSION = context['__version__']\n\nHERE = Path(__file__).parent\n\ntry:\n with open(HERE / \"README.md\", encoding='utf-8') as f:\n long_description = '\\n' + f.read()\nexcept FileNotFoundError:\n long_description = DESCRIPTION\n\nREQUIRED = [i.strip() for i in open(HERE / 'requirements.txt') if not i.startswith('#')]\n\nsetup(\n name=NAME,\n version=VERSION,\n description=DESCRIPTION,\n author_email=EMAIL,\n long_description=long_description,\n long_description_content_type='text/markdown',\n author=AUTHOR,\n url=URL,\n python_requires=REQUIRES_PYTHON,\n install_requires=REQUIRED,\n extras_require={\n 'dev': ['coverage', 'flake8', 'mypy', 'pdoc3', 'pytest'],\n 'wm': ['audioseal'],\n },\n packages=[p for p in find_packages() if p.startswith('audiocraft')],\n package_data={'audiocraft': ['py.typed']},\n include_package_data=True,\n license='MIT License',\n classifiers=[\n # Trove classifiers\n # Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers\n 'License :: OSI Approved :: MIT License',\n 'Topic :: Multimedia :: Sound/Audio',\n 'Topic :: Scientific/Engineering :: Artificial Intelligence',\n ],\n)\n" }, { "path": "tests/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "tests/adversarial/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "tests/adversarial/test_discriminators.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport random\n\nimport torch\n\nfrom audiocraft.adversarial.discriminators import (\n MultiPeriodDiscriminator,\n MultiScaleDiscriminator,\n MultiScaleSTFTDiscriminator\n)\n\n\nclass TestMultiPeriodDiscriminator:\n\n def test_mpd_discriminator(self):\n N, C, T = 2, 2, random.randrange(1, 100_000)\n t0 = torch.randn(N, C, T)\n periods = [1, 2, 3]\n mpd = MultiPeriodDiscriminator(periods=periods, in_channels=C)\n logits, fmaps = mpd(t0)\n\n assert len(logits) == len(periods)\n assert len(fmaps) == len(periods)\n assert all([logit.shape[0] == N and len(logit.shape) == 4 for logit in logits])\n assert all([feature.shape[0] == N for fmap in fmaps for feature in fmap])\n\n\nclass TestMultiScaleDiscriminator:\n\n def test_msd_discriminator(self):\n N, C, T = 2, 2, random.randrange(1, 100_000)\n t0 = torch.randn(N, C, T)\n\n scale_norms = ['weight_norm', 'weight_norm']\n msd = MultiScaleDiscriminator(scale_norms=scale_norms, in_channels=C)\n logits, fmaps = msd(t0)\n\n assert len(logits) == len(scale_norms)\n assert len(fmaps) == len(scale_norms)\n assert all([logit.shape[0] == N and len(logit.shape) == 3 for logit in logits])\n assert all([feature.shape[0] == N for fmap in fmaps for feature in fmap])\n\n\nclass TestMultiScaleStftDiscriminator:\n\n def test_msstftd_discriminator(self):\n N, C, T = 2, 2, random.randrange(1, 100_000)\n t0 = torch.randn(N, C, T)\n\n n_filters = 4\n n_ffts = [128, 256, 64]\n hop_lengths = [32, 64, 16]\n win_lengths = [128, 256, 64]\n\n msstftd = MultiScaleSTFTDiscriminator(filters=n_filters, n_ffts=n_ffts, hop_lengths=hop_lengths,\n win_lengths=win_lengths, in_channels=C)\n logits, fmaps = msstftd(t0)\n\n assert len(logits) == len(n_ffts)\n assert len(fmaps) == len(n_ffts)\n assert all([logit.shape[0] == N and len(logit.shape) == 4 for logit in logits])\n assert all([feature.shape[0] == N for fmap in fmaps for feature in fmap])\n" }, { "path": "tests/adversarial/test_losses.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport pytest\nimport random\n\nimport torch\n\nfrom audiocraft.adversarial import (\n AdversarialLoss,\n get_adv_criterion,\n get_real_criterion,\n get_fake_criterion,\n FeatureMatchingLoss,\n MultiScaleDiscriminator,\n)\n\n\nclass TestAdversarialLoss:\n\n def test_adversarial_single_multidiscriminator(self):\n adv = MultiScaleDiscriminator()\n optimizer = torch.optim.Adam(\n adv.parameters(),\n lr=1e-4,\n )\n loss, loss_real, loss_fake = get_adv_criterion('mse'), get_real_criterion('mse'), get_fake_criterion('mse')\n adv_loss = AdversarialLoss(adv, optimizer, loss, loss_real, loss_fake)\n\n B, C, T = 4, 1, random.randint(1000, 5000)\n real = torch.randn(B, C, T)\n fake = torch.randn(B, C, T)\n\n disc_loss = adv_loss.train_adv(fake, real)\n assert isinstance(disc_loss, torch.Tensor) and isinstance(disc_loss.item(), float)\n\n loss, loss_feat = adv_loss(fake, real)\n assert isinstance(loss, torch.Tensor) and isinstance(loss.item(), float)\n # we did not specify feature loss\n assert loss_feat.item() == 0.\n\n def test_adversarial_feat_loss(self):\n adv = MultiScaleDiscriminator()\n optimizer = torch.optim.Adam(\n adv.parameters(),\n lr=1e-4,\n )\n loss, loss_real, loss_fake = get_adv_criterion('mse'), get_real_criterion('mse'), get_fake_criterion('mse')\n feat_loss = FeatureMatchingLoss()\n adv_loss = AdversarialLoss(adv, optimizer, loss, loss_real, loss_fake, feat_loss)\n\n B, C, T = 4, 1, random.randint(1000, 5000)\n real = torch.randn(B, C, T)\n fake = torch.randn(B, C, T)\n\n loss, loss_feat = adv_loss(fake, real)\n\n assert isinstance(loss, torch.Tensor) and isinstance(loss.item(), float)\n assert isinstance(loss_feat, torch.Tensor) and isinstance(loss.item(), float)\n\n\nclass TestGeneratorAdversarialLoss:\n\n def test_hinge_generator_adv_loss(self):\n adv_loss = get_adv_criterion(loss_type='hinge')\n\n t0 = torch.randn(1, 2, 0)\n t1 = torch.FloatTensor([1.0, 2.0, 3.0])\n\n assert adv_loss(t0).item() == 0.0\n assert adv_loss(t1).item() == -2.0\n\n def test_mse_generator_adv_loss(self):\n adv_loss = get_adv_criterion(loss_type='mse')\n\n t0 = torch.randn(1, 2, 0)\n t1 = torch.FloatTensor([1.0, 1.0, 1.0])\n t2 = torch.FloatTensor([2.0, 5.0, 5.0])\n\n assert adv_loss(t0).item() == 0.0\n assert adv_loss(t1).item() == 0.0\n assert adv_loss(t2).item() == 11.0\n\n\nclass TestDiscriminatorAdversarialLoss:\n\n def _disc_loss(self, loss_type: str, fake: torch.Tensor, real: torch.Tensor):\n disc_loss_real = get_real_criterion(loss_type)\n disc_loss_fake = get_fake_criterion(loss_type)\n\n loss = disc_loss_fake(fake) + disc_loss_real(real)\n return loss\n\n def test_hinge_discriminator_adv_loss(self):\n loss_type = 'hinge'\n t0 = torch.FloatTensor([0.0, 0.0, 0.0])\n t1 = torch.FloatTensor([1.0, 2.0, 3.0])\n\n assert self._disc_loss(loss_type, t0, t0).item() == 2.0\n assert self._disc_loss(loss_type, t1, t1).item() == 3.0\n\n def test_mse_discriminator_adv_loss(self):\n loss_type = 'mse'\n\n t0 = torch.FloatTensor([0.0, 0.0, 0.0])\n t1 = torch.FloatTensor([1.0, 1.0, 1.0])\n\n assert self._disc_loss(loss_type, t0, t0).item() == 1.0\n assert self._disc_loss(loss_type, t1, t0).item() == 2.0\n\n\nclass TestFeatureMatchingLoss:\n\n def test_features_matching_loss_base(self):\n ft_matching_loss = FeatureMatchingLoss()\n length = random.randrange(1, 100_000)\n t1 = torch.randn(1, 2, length)\n\n loss = ft_matching_loss([t1], [t1])\n assert isinstance(loss, torch.Tensor)\n assert loss.item() == 0.0\n\n def test_features_matching_loss_raises_exception(self):\n ft_matching_loss = FeatureMatchingLoss()\n length = random.randrange(1, 100_000)\n t1 = torch.randn(1, 2, length)\n t2 = torch.randn(1, 2, length + 1)\n\n with pytest.raises(AssertionError):\n ft_matching_loss([], [])\n\n with pytest.raises(AssertionError):\n ft_matching_loss([t1], [t1, t1])\n\n with pytest.raises(AssertionError):\n ft_matching_loss([t1], [t2])\n\n def test_features_matching_loss_output(self):\n loss_nonorm = FeatureMatchingLoss(normalize=False)\n loss_layer_normed = FeatureMatchingLoss(normalize=True)\n\n length = random.randrange(1, 100_000)\n t1 = torch.randn(1, 2, length)\n t2 = torch.randn(1, 2, length)\n\n assert loss_nonorm([t1, t2], [t1, t2]).item() == 0.0\n assert loss_layer_normed([t1, t2], [t1, t2]).item() == 0.0\n\n t3 = torch.FloatTensor([1.0, 2.0, 3.0])\n t4 = torch.FloatTensor([2.0, 10.0, 3.0])\n\n assert loss_nonorm([t3], [t4]).item() == 3.0\n assert loss_nonorm([t3, t3], [t4, t4]).item() == 6.0\n\n assert loss_layer_normed([t3], [t4]).item() == 3.0\n assert loss_layer_normed([t3, t3], [t4, t4]).item() == 3.0\n" }, { "path": "tests/common_utils/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# flake8: noqa\nfrom .temp_utils import TempDirMixin\nfrom .wav_utils import get_batch_white_noise, get_white_noise, save_wav\n" }, { "path": "tests/common_utils/temp_utils.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport os\nimport tempfile\n\n\nclass TempDirMixin:\n \"\"\"Mixin to provide easy access to temp dir.\n \"\"\"\n\n temp_dir_ = None\n\n @classmethod\n def get_base_temp_dir(cls):\n # If AUDIOCRAFT_TEST_DIR is set, use it instead of temporary directory.\n # this is handy for debugging.\n key = \"AUDIOCRAFT_TEST_DIR\"\n if key in os.environ:\n return os.environ[key]\n if cls.temp_dir_ is None:\n cls.temp_dir_ = tempfile.TemporaryDirectory()\n return cls.temp_dir_.name\n\n @classmethod\n def tearDownClass(cls):\n if cls.temp_dir_ is not None:\n try:\n cls.temp_dir_.cleanup()\n cls.temp_dir_ = None\n except PermissionError:\n # On Windows there is a know issue with `shutil.rmtree`,\n # which fails intermittently.\n # https://github.com/python/cpython/issues/74168\n # Following the above thread, we ignore it.\n pass\n super().tearDownClass()\n\n @property\n def id(self):\n return self.__class__.__name__\n\n def get_temp_path(self, *paths):\n temp_dir = os.path.join(self.get_base_temp_dir(), self.id)\n path = os.path.join(temp_dir, *paths)\n os.makedirs(os.path.dirname(path), exist_ok=True)\n return path\n\n def get_temp_dir(self, *paths):\n temp_dir = os.path.join(self.get_base_temp_dir(), self.id)\n path = os.path.join(temp_dir, *paths)\n os.makedirs(path, exist_ok=True)\n return path\n" }, { "path": "tests/common_utils/wav_utils.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom pathlib import Path\n\nimport torch\n\nfrom audiocraft.data.audio import audio_write\n\n\ndef get_white_noise(chs: int = 1, num_frames: int = 1):\n wav = torch.randn(chs, num_frames)\n return wav\n\n\ndef get_batch_white_noise(bs: int = 1, chs: int = 1, num_frames: int = 1):\n wav = torch.randn(bs, chs, num_frames)\n return wav\n\n\ndef save_wav(path: str, wav: torch.Tensor, sample_rate: int):\n assert wav.dim() == 2, wav.shape\n fp = Path(path)\n assert fp.suffix in ['.mp3', '.ogg', '.wav', '.flac'], fp\n audio_write(fp.parent / fp.stem, wav, sample_rate, fp.suffix[1:],\n normalize=False, strategy='clip', peak_clip_headroom_db=0)\n" }, { "path": "tests/data/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "tests/data/test_audio.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom itertools import product\nimport random\n\nimport numpy as np\nimport torch\nimport torchaudio\n\nfrom audiocraft.data.audio import audio_info, audio_read, audio_write, _av_read\n\nfrom ..common_utils import TempDirMixin, get_white_noise, save_wav\n\n\nclass TestInfo(TempDirMixin):\n\n def test_info_mp3(self):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 1.\n for sample_rate, ch in product(sample_rates, channels):\n wav = get_white_noise(ch, int(sample_rate * duration))\n path = self.get_temp_path('sample_wav.mp3')\n save_wav(path, wav, sample_rate)\n info = audio_info(path)\n assert info.sample_rate == sample_rate\n assert info.channels == ch\n # we cannot trust torchaudio for num_frames, so we don't check\n\n def _test_info_format(self, ext: str):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 1.\n for sample_rate, ch in product(sample_rates, channels):\n n_frames = int(sample_rate * duration)\n wav = get_white_noise(ch, n_frames)\n path = self.get_temp_path(f'sample_wav{ext}')\n save_wav(path, wav, sample_rate)\n info = audio_info(path)\n assert info.sample_rate == sample_rate\n assert info.channels == ch\n assert np.isclose(info.duration, duration, atol=1e-5)\n\n def test_info_wav(self):\n self._test_info_format('.wav')\n\n def test_info_flac(self):\n self._test_info_format('.flac')\n\n def test_info_ogg(self):\n self._test_info_format('.ogg')\n\n def test_info_m4a(self):\n # TODO: generate m4a file programmatically\n # self._test_info_format('.m4a')\n pass\n\n\nclass TestRead(TempDirMixin):\n\n def test_read_full_wav(self):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 1.\n for sample_rate, ch in product(sample_rates, channels):\n n_frames = int(sample_rate * duration)\n wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99)\n path = self.get_temp_path('sample_wav.wav')\n save_wav(path, wav, sample_rate)\n read_wav, read_sr = audio_read(path)\n assert read_sr == sample_rate\n assert read_wav.shape[0] == wav.shape[0]\n assert read_wav.shape[1] == wav.shape[1]\n assert torch.allclose(read_wav, wav, rtol=1e-03, atol=1e-04)\n\n def test_read_partial_wav(self):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 1.\n read_duration = torch.rand(1).item()\n for sample_rate, ch in product(sample_rates, channels):\n n_frames = int(sample_rate * duration)\n read_frames = int(sample_rate * read_duration)\n wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99)\n path = self.get_temp_path('sample_wav.wav')\n save_wav(path, wav, sample_rate)\n read_wav, read_sr = audio_read(path, 0, read_duration)\n assert read_sr == sample_rate\n assert read_wav.shape[0] == wav.shape[0]\n assert read_wav.shape[1] == read_frames\n assert torch.allclose(read_wav[..., 0:read_frames], wav[..., 0:read_frames], rtol=1e-03, atol=1e-04)\n\n def test_read_seek_time_wav(self):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 1.\n read_duration = 1.\n for sample_rate, ch in product(sample_rates, channels):\n n_frames = int(sample_rate * duration)\n wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99)\n path = self.get_temp_path('sample_wav.wav')\n save_wav(path, wav, sample_rate)\n seek_time = torch.rand(1).item()\n read_wav, read_sr = audio_read(path, seek_time, read_duration)\n seek_frames = int(sample_rate * seek_time)\n expected_frames = n_frames - seek_frames\n assert read_sr == sample_rate\n assert read_wav.shape[0] == wav.shape[0]\n assert read_wav.shape[1] == expected_frames\n assert torch.allclose(read_wav, wav[..., seek_frames:], rtol=1e-03, atol=1e-04)\n\n def test_read_seek_time_wav_padded(self):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 1.\n read_duration = 1.\n for sample_rate, ch in product(sample_rates, channels):\n n_frames = int(sample_rate * duration)\n read_frames = int(sample_rate * read_duration)\n wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99)\n path = self.get_temp_path('sample_wav.wav')\n save_wav(path, wav, sample_rate)\n seek_time = torch.rand(1).item()\n seek_frames = int(sample_rate * seek_time)\n expected_frames = n_frames - seek_frames\n read_wav, read_sr = audio_read(path, seek_time, read_duration, pad=True)\n expected_pad_wav = torch.zeros(wav.shape[0], read_frames - expected_frames)\n assert read_sr == sample_rate\n assert read_wav.shape[0] == wav.shape[0]\n assert read_wav.shape[1] == read_frames\n assert torch.allclose(read_wav[..., :expected_frames], wav[..., seek_frames:], rtol=1e-03, atol=1e-04)\n assert torch.allclose(read_wav[..., expected_frames:], expected_pad_wav)\n\n\nclass TestAvRead(TempDirMixin):\n\n def test_avread_seek_base(self):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 2.\n for sample_rate, ch in product(sample_rates, channels):\n n_frames = int(sample_rate * duration)\n wav = get_white_noise(ch, n_frames)\n path = self.get_temp_path(f'reference_a_{sample_rate}_{ch}.wav')\n save_wav(path, wav, sample_rate)\n for _ in range(100):\n # seek will always load a full duration segment in the file\n seek_time = random.uniform(0.0, 1.0)\n seek_duration = random.uniform(0.001, 1.0)\n read_wav, read_sr = _av_read(path, seek_time, seek_duration)\n assert read_sr == sample_rate\n assert read_wav.shape[0] == wav.shape[0]\n assert read_wav.shape[-1] == int(seek_duration * sample_rate)\n\n def test_avread_seek_partial(self):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 1.\n for sample_rate, ch in product(sample_rates, channels):\n n_frames = int(sample_rate * duration)\n wav = get_white_noise(ch, n_frames)\n path = self.get_temp_path(f'reference_b_{sample_rate}_{ch}.wav')\n save_wav(path, wav, sample_rate)\n for _ in range(100):\n # seek will always load a partial segment\n seek_time = random.uniform(0.5, 1.)\n seek_duration = 1.\n expected_num_frames = n_frames - int(seek_time * sample_rate)\n read_wav, read_sr = _av_read(path, seek_time, seek_duration)\n assert read_sr == sample_rate\n assert read_wav.shape[0] == wav.shape[0]\n assert read_wav.shape[-1] == expected_num_frames\n\n def test_avread_seek_outofbound(self):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 1.\n for sample_rate, ch in product(sample_rates, channels):\n n_frames = int(sample_rate * duration)\n wav = get_white_noise(ch, n_frames)\n path = self.get_temp_path(f'reference_c_{sample_rate}_{ch}.wav')\n save_wav(path, wav, sample_rate)\n seek_time = 1.5\n read_wav, read_sr = _av_read(path, seek_time, 1.)\n assert read_sr == sample_rate\n assert read_wav.shape[0] == wav.shape[0]\n assert read_wav.shape[-1] == 0\n\n def test_avread_seek_edge(self):\n sample_rates = [8000, 16_000]\n # some of these values will have\n # int(((frames - 1) / sample_rate) * sample_rate) != (frames - 1)\n n_frames = [1000, 1001, 1002]\n channels = [1, 2]\n for sample_rate, ch, frames in product(sample_rates, channels, n_frames):\n duration = frames / sample_rate\n wav = get_white_noise(ch, frames)\n path = self.get_temp_path(f'reference_d_{sample_rate}_{ch}.wav')\n save_wav(path, wav, sample_rate)\n seek_time = (frames - 1) / sample_rate\n seek_frames = int(seek_time * sample_rate)\n read_wav, read_sr = _av_read(path, seek_time, duration)\n assert read_sr == sample_rate\n assert read_wav.shape[0] == wav.shape[0]\n assert read_wav.shape[-1] == (frames - seek_frames)\n\n\nclass TestAudioWrite(TempDirMixin):\n\n def test_audio_write_wav(self):\n torch.manual_seed(1234)\n sample_rates = [8000, 16_000]\n n_frames = [1000, 1001, 1002]\n channels = [1, 2]\n strategies = [\"peak\", \"clip\", \"rms\"]\n formats = [\"wav\", \"mp3\"]\n for sample_rate, ch, frames in product(sample_rates, channels, n_frames):\n for format_, strategy in product(formats, strategies):\n wav = get_white_noise(ch, frames)\n path = self.get_temp_path(f'pred_{sample_rate}_{ch}')\n audio_write(path, wav, sample_rate, format_, strategy=strategy)\n read_wav, read_sr = torchaudio.load(f'{path}.{format_}')\n if format_ == \"wav\":\n assert read_wav.shape == wav.shape\n\n if format_ == \"wav\" and strategy in [\"peak\", \"rms\"]:\n rescaled_read_wav = read_wav / read_wav.abs().max() * wav.abs().max()\n # for a Gaussian, the typical max scale will be less than ~5x the std.\n # The error when writing to disk will ~ 1/2**15, and when rescaling, 5x that.\n # For RMS target, rescaling leaves more headroom by default, leading\n # to a 20x rescaling typically\n atol = (5 if strategy == \"peak\" else 20) / 2**15\n delta = (rescaled_read_wav - wav).abs().max()\n assert torch.allclose(wav, rescaled_read_wav, rtol=0, atol=atol), (delta, atol)\n formats = [\"wav\"] # faster unit tests\n" }, { "path": "tests/data/test_audio_dataset.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom functools import partial\nfrom itertools import product\nimport json\nimport math\nimport os\nimport random\nimport typing as tp\n\nimport pytest\nimport torch\nfrom torch.utils.data import DataLoader\n\nfrom audiocraft.data.audio_dataset import (\n AudioDataset,\n AudioMeta,\n _get_audio_meta,\n load_audio_meta,\n save_audio_meta\n)\nfrom audiocraft.data.zip import PathInZip\n\nfrom ..common_utils import TempDirMixin, get_white_noise, save_wav\n\n\nclass TestAudioMeta(TempDirMixin):\n\n def test_get_audio_meta(self):\n sample_rates = [8000, 16_000]\n channels = [1, 2]\n duration = 1.\n for sample_rate, ch in product(sample_rates, channels):\n n_frames = int(duration * sample_rate)\n wav = get_white_noise(ch, n_frames)\n path = self.get_temp_path('sample.wav')\n save_wav(path, wav, sample_rate)\n m = _get_audio_meta(path, minimal=True)\n assert m.path == path, 'path does not match'\n assert m.sample_rate == sample_rate, 'sample rate does not match'\n assert m.duration == duration, 'duration does not match'\n assert m.amplitude is None\n assert m.info_path is None\n\n def test_save_audio_meta(self):\n audio_meta = [\n AudioMeta(\"mypath1\", 1., 16_000, None, None, PathInZip('/foo/bar.zip:/relative/file1.json')),\n AudioMeta(\"mypath2\", 2., 16_000, None, None, PathInZip('/foo/bar.zip:/relative/file2.json'))\n ]\n empty_audio_meta = []\n for idx, meta in enumerate([audio_meta, empty_audio_meta]):\n path = self.get_temp_path(f'data_{idx}_save.jsonl')\n save_audio_meta(path, meta)\n with open(path, 'r') as f:\n lines = f.readlines()\n read_meta = [AudioMeta.from_dict(json.loads(line)) for line in lines]\n assert len(read_meta) == len(meta)\n for m, read_m in zip(meta, read_meta):\n assert m == read_m\n\n def test_load_audio_meta(self):\n try:\n import dora\n except ImportError:\n dora = None # type: ignore\n\n audio_meta = [\n AudioMeta(\"mypath1\", 1., 16_000, None, None, PathInZip('/foo/bar.zip:/relative/file1.json')),\n AudioMeta(\"mypath2\", 2., 16_000, None, None, PathInZip('/foo/bar.zip:/relative/file2.json'))\n ]\n empty_meta = []\n for idx, meta in enumerate([audio_meta, empty_meta]):\n path = self.get_temp_path(f'data_{idx}_load.jsonl')\n with open(path, 'w') as f:\n for m in meta:\n json_str = json.dumps(m.to_dict()) + '\\n'\n f.write(json_str)\n read_meta = load_audio_meta(path)\n assert len(read_meta) == len(meta)\n for m, read_m in zip(meta, read_meta):\n if dora:\n m.path = dora.git_save.to_absolute_path(m.path)\n assert m == read_m, f'original={m}, read={read_m}'\n\n\nclass TestAudioDataset(TempDirMixin):\n\n def _create_audio_files(self,\n root_name: str,\n num_examples: int,\n durations: tp.Union[float, tp.Tuple[float, float]] = (0.1, 1.),\n sample_rate: int = 16_000,\n channels: int = 1):\n root_dir = self.get_temp_dir(root_name)\n for i in range(num_examples):\n if isinstance(durations, float):\n duration = durations\n elif isinstance(durations, tuple) and len(durations) == 1:\n duration = durations[0]\n elif isinstance(durations, tuple) and len(durations) == 2:\n duration = random.uniform(durations[0], durations[1])\n else:\n assert False\n n_frames = int(duration * sample_rate)\n wav = get_white_noise(channels, n_frames)\n path = os.path.join(root_dir, f'example_{i}.wav')\n save_wav(path, wav, sample_rate)\n return root_dir\n\n def _create_audio_dataset(self,\n root_name: str,\n total_num_examples: int,\n durations: tp.Union[float, tp.Tuple[float, float]] = (0.1, 1.),\n sample_rate: int = 16_000,\n channels: int = 1,\n segment_duration: tp.Optional[float] = None,\n num_examples: int = 10,\n shuffle: bool = True,\n return_info: bool = False):\n root_dir = self._create_audio_files(root_name, total_num_examples, durations, sample_rate, channels)\n dataset = AudioDataset.from_path(root_dir,\n minimal_meta=True,\n segment_duration=segment_duration,\n num_samples=num_examples,\n sample_rate=sample_rate,\n channels=channels,\n shuffle=shuffle,\n return_info=return_info)\n return dataset\n\n def test_dataset_full(self):\n total_examples = 10\n min_duration, max_duration = 1., 4.\n sample_rate = 16_000\n channels = 1\n dataset = self._create_audio_dataset(\n 'dset', total_examples, durations=(min_duration, max_duration),\n sample_rate=sample_rate, channels=channels, segment_duration=None)\n assert len(dataset) == total_examples\n assert dataset.sample_rate == sample_rate\n assert dataset.channels == channels\n for idx in range(len(dataset)):\n sample = dataset[idx]\n assert sample.shape[0] == channels\n assert sample.shape[1] <= int(max_duration * sample_rate)\n assert sample.shape[1] >= int(min_duration * sample_rate)\n\n def test_dataset_segment(self):\n total_examples = 10\n num_samples = 20\n min_duration, max_duration = 1., 4.\n segment_duration = 1.\n sample_rate = 16_000\n channels = 1\n dataset = self._create_audio_dataset(\n 'dset', total_examples, durations=(min_duration, max_duration), sample_rate=sample_rate,\n channels=channels, segment_duration=segment_duration, num_examples=num_samples)\n assert len(dataset) == num_samples\n assert dataset.sample_rate == sample_rate\n assert dataset.channels == channels\n for idx in range(len(dataset)):\n sample = dataset[idx]\n assert sample.shape[0] == channels\n assert sample.shape[1] == int(segment_duration * sample_rate)\n\n def test_dataset_equal_audio_and_segment_durations(self):\n total_examples = 1\n num_samples = 2\n audio_duration = 1.\n segment_duration = 1.\n sample_rate = 16_000\n channels = 1\n dataset = self._create_audio_dataset(\n 'dset', total_examples, durations=audio_duration, sample_rate=sample_rate,\n channels=channels, segment_duration=segment_duration, num_examples=num_samples)\n assert len(dataset) == num_samples\n assert dataset.sample_rate == sample_rate\n assert dataset.channels == channels\n for idx in range(len(dataset)):\n sample = dataset[idx]\n assert sample.shape[0] == channels\n assert sample.shape[1] == int(segment_duration * sample_rate)\n # the random seek_time adds variability on audio read\n sample_1 = dataset[0]\n sample_2 = dataset[1]\n assert not torch.allclose(sample_1, sample_2)\n\n def test_dataset_samples(self):\n total_examples = 1\n num_samples = 2\n audio_duration = 1.\n segment_duration = 1.\n sample_rate = 16_000\n channels = 1\n\n create_dataset = partial(\n self._create_audio_dataset,\n 'dset', total_examples, durations=audio_duration, sample_rate=sample_rate,\n channels=channels, segment_duration=segment_duration, num_examples=num_samples,\n )\n\n dataset = create_dataset(shuffle=True)\n # when shuffle = True, we have different inputs for the same index across epoch\n sample_1 = dataset[0]\n sample_2 = dataset[0]\n assert not torch.allclose(sample_1, sample_2)\n\n dataset_noshuffle = create_dataset(shuffle=False)\n # when shuffle = False, we have same inputs for the same index across epoch\n sample_1 = dataset_noshuffle[0]\n sample_2 = dataset_noshuffle[0]\n assert torch.allclose(sample_1, sample_2)\n\n def test_dataset_return_info(self):\n total_examples = 10\n num_samples = 20\n min_duration, max_duration = 1., 4.\n segment_duration = 1.\n sample_rate = 16_000\n channels = 1\n dataset = self._create_audio_dataset(\n 'dset', total_examples, durations=(min_duration, max_duration), sample_rate=sample_rate,\n channels=channels, segment_duration=segment_duration, num_examples=num_samples, return_info=True)\n assert len(dataset) == num_samples\n assert dataset.sample_rate == sample_rate\n assert dataset.channels == channels\n for idx in range(len(dataset)):\n sample, segment_info = dataset[idx]\n assert sample.shape[0] == channels\n assert sample.shape[1] == int(segment_duration * sample_rate)\n assert segment_info.sample_rate == sample_rate\n assert segment_info.total_frames == int(segment_duration * sample_rate)\n assert segment_info.n_frames <= int(segment_duration * sample_rate)\n assert segment_info.seek_time >= 0\n\n def test_dataset_return_info_no_segment_duration(self):\n total_examples = 10\n num_samples = 20\n min_duration, max_duration = 1., 4.\n segment_duration = None\n sample_rate = 16_000\n channels = 1\n dataset = self._create_audio_dataset(\n 'dset', total_examples, durations=(min_duration, max_duration), sample_rate=sample_rate,\n channels=channels, segment_duration=segment_duration, num_examples=num_samples, return_info=True)\n assert len(dataset) == total_examples\n assert dataset.sample_rate == sample_rate\n assert dataset.channels == channels\n for idx in range(len(dataset)):\n sample, segment_info = dataset[idx]\n assert sample.shape[0] == channels\n assert sample.shape[1] == segment_info.total_frames\n assert segment_info.sample_rate == sample_rate\n assert segment_info.n_frames <= segment_info.total_frames\n\n def test_dataset_collate_fn(self):\n total_examples = 10\n num_samples = 20\n min_duration, max_duration = 1., 4.\n segment_duration = 1.\n sample_rate = 16_000\n channels = 1\n dataset = self._create_audio_dataset(\n 'dset', total_examples, durations=(min_duration, max_duration), sample_rate=sample_rate,\n channels=channels, segment_duration=segment_duration, num_examples=num_samples, return_info=False)\n batch_size = 4\n dataloader = DataLoader(\n dataset,\n batch_size=batch_size,\n num_workers=0\n )\n for idx, batch in enumerate(dataloader):\n assert batch.shape[0] == batch_size\n\n @pytest.mark.parametrize(\"segment_duration\", [1.0, None])\n def test_dataset_with_meta_collate_fn(self, segment_duration):\n total_examples = 10\n num_samples = 20\n min_duration, max_duration = 1., 4.\n segment_duration = 1.\n sample_rate = 16_000\n channels = 1\n dataset = self._create_audio_dataset(\n 'dset', total_examples, durations=(min_duration, max_duration), sample_rate=sample_rate,\n channels=channels, segment_duration=segment_duration, num_examples=num_samples, return_info=True)\n batch_size = 4\n dataloader = DataLoader(\n dataset,\n batch_size=batch_size,\n collate_fn=dataset.collater,\n num_workers=0\n )\n for idx, batch in enumerate(dataloader):\n wav, infos = batch\n assert wav.shape[0] == batch_size\n assert len(infos) == batch_size\n\n @pytest.mark.parametrize(\"segment_duration,sample_on_weight,sample_on_duration,a_hist,b_hist,c_hist\", [\n [1, True, True, 0.5, 0.5, 0.0],\n [1, False, True, 0.25, 0.5, 0.25],\n [1, True, False, 0.666, 0.333, 0.0],\n [1, False, False, 0.333, 0.333, 0.333],\n [None, False, False, 0.333, 0.333, 0.333]])\n def test_sample_with_weight(self, segment_duration, sample_on_weight, sample_on_duration, a_hist, b_hist, c_hist):\n random.seed(1234)\n rng = torch.Generator()\n rng.manual_seed(1234)\n\n def _get_histogram(dataset, repetitions=20_000):\n counts = {file_meta.path: 0. for file_meta in meta}\n for _ in range(repetitions):\n file_meta = dataset.sample_file(0, rng)\n counts[file_meta.path] += 1\n return {name: count / repetitions for name, count in counts.items()}\n\n meta = [\n AudioMeta(path='a', duration=5, sample_rate=1, weight=2),\n AudioMeta(path='b', duration=10, sample_rate=1, weight=None),\n AudioMeta(path='c', duration=5, sample_rate=1, weight=0),\n ]\n dataset = AudioDataset(\n meta, segment_duration=segment_duration, sample_on_weight=sample_on_weight,\n sample_on_duration=sample_on_duration)\n hist = _get_histogram(dataset)\n assert math.isclose(hist['a'], a_hist, abs_tol=0.01)\n assert math.isclose(hist['b'], b_hist, abs_tol=0.01)\n assert math.isclose(hist['c'], c_hist, abs_tol=0.01)\n\n def test_meta_duration_filter_all(self):\n meta = [\n AudioMeta(path='a', duration=5, sample_rate=1, weight=2),\n AudioMeta(path='b', duration=10, sample_rate=1, weight=None),\n AudioMeta(path='c', duration=5, sample_rate=1, weight=0),\n ]\n try:\n AudioDataset(meta, segment_duration=11, min_segment_ratio=1)\n assert False\n except AssertionError:\n assert True\n\n def test_meta_duration_filter_long(self):\n meta = [\n AudioMeta(path='a', duration=5, sample_rate=1, weight=2),\n AudioMeta(path='b', duration=10, sample_rate=1, weight=None),\n AudioMeta(path='c', duration=5, sample_rate=1, weight=0),\n ]\n dataset = AudioDataset(meta, segment_duration=None, min_segment_ratio=1, max_audio_duration=7)\n assert len(dataset) == 2\n" }, { "path": "tests/data/test_audio_utils.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport julius\nimport torch\nimport pytest\n\nfrom audiocraft.data.audio_utils import (\n _clip_wav,\n convert_audio_channels,\n convert_audio,\n f32_pcm,\n i16_pcm,\n normalize_audio\n)\nfrom ..common_utils import get_batch_white_noise\n\n\nclass TestConvertAudioChannels:\n\n def test_convert_audio_channels_downmix(self):\n b, c, t = 2, 3, 100\n audio = get_batch_white_noise(b, c, t)\n mixed = convert_audio_channels(audio, channels=2)\n assert list(mixed.shape) == [b, 2, t]\n\n def test_convert_audio_channels_nochange(self):\n b, c, t = 2, 3, 100\n audio = get_batch_white_noise(b, c, t)\n mixed = convert_audio_channels(audio, channels=c)\n assert list(mixed.shape) == list(audio.shape)\n\n def test_convert_audio_channels_upmix(self):\n b, c, t = 2, 1, 100\n audio = get_batch_white_noise(b, c, t)\n mixed = convert_audio_channels(audio, channels=3)\n assert list(mixed.shape) == [b, 3, t]\n\n def test_convert_audio_channels_upmix_error(self):\n b, c, t = 2, 2, 100\n audio = get_batch_white_noise(b, c, t)\n with pytest.raises(ValueError):\n convert_audio_channels(audio, channels=3)\n\n\nclass TestConvertAudio:\n\n def test_convert_audio_channels_downmix(self):\n b, c, dur = 2, 3, 4.\n sr = 128\n audio = get_batch_white_noise(b, c, int(sr * dur))\n out = convert_audio(audio, from_rate=sr, to_rate=sr, to_channels=2)\n assert list(out.shape) == [audio.shape[0], 2, audio.shape[-1]]\n\n def test_convert_audio_channels_upmix(self):\n b, c, dur = 2, 1, 4.\n sr = 128\n audio = get_batch_white_noise(b, c, int(sr * dur))\n out = convert_audio(audio, from_rate=sr, to_rate=sr, to_channels=3)\n assert list(out.shape) == [audio.shape[0], 3, audio.shape[-1]]\n\n def test_convert_audio_upsample(self):\n b, c, dur = 2, 1, 4.\n sr = 2\n new_sr = 3\n audio = get_batch_white_noise(b, c, int(sr * dur))\n out = convert_audio(audio, from_rate=sr, to_rate=new_sr, to_channels=c)\n out_j = julius.resample.resample_frac(audio, old_sr=sr, new_sr=new_sr)\n assert torch.allclose(out, out_j)\n\n def test_convert_audio_resample(self):\n b, c, dur = 2, 1, 4.\n sr = 3\n new_sr = 2\n audio = get_batch_white_noise(b, c, int(sr * dur))\n out = convert_audio(audio, from_rate=sr, to_rate=new_sr, to_channels=c)\n out_j = julius.resample.resample_frac(audio, old_sr=sr, new_sr=new_sr)\n assert torch.allclose(out, out_j)\n\n def test_convert_pcm(self):\n b, c, dur = 2, 1, 4.\n sr = 3\n i16_audio = torch.randint(-2**15, 2**15, (b, c, int(sr * dur)), dtype=torch.int16)\n f32_audio = f32_pcm(i16_audio)\n another_i16_audio = i16_pcm(f32_audio)\n assert torch.allclose(i16_audio, another_i16_audio)\n\n\nclass TestNormalizeAudio:\n\n def test_clip_wav(self):\n b, c, dur = 2, 1, 4.\n sr = 3\n audio = 10.0 * get_batch_white_noise(b, c, int(sr * dur))\n _clip_wav(audio)\n assert audio.abs().max() <= 1\n\n def test_normalize_audio_clip(self):\n b, c, dur = 2, 1, 4.\n sr = 3\n audio = 10.0 * get_batch_white_noise(b, c, int(sr * dur))\n norm_audio = normalize_audio(audio, strategy='clip')\n assert norm_audio.abs().max() <= 1\n\n def test_normalize_audio_rms(self):\n b, c, dur = 2, 1, 4.\n sr = 3\n audio = 10.0 * get_batch_white_noise(b, c, int(sr * dur))\n norm_audio = normalize_audio(audio, strategy='rms')\n assert norm_audio.abs().max() <= 1\n\n def test_normalize_audio_peak(self):\n b, c, dur = 2, 1, 4.\n sr = 3\n audio = 10.0 * get_batch_white_noise(b, c, int(sr * dur))\n norm_audio = normalize_audio(audio, strategy='peak')\n assert norm_audio.abs().max() <= 1\n" }, { "path": "tests/losses/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "tests/losses/test_losses.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport random\n\nimport torch\n\nfrom audiocraft.losses import (\n MelSpectrogramL1Loss,\n MultiScaleMelSpectrogramLoss,\n MRSTFTLoss,\n SISNR,\n STFTLoss,\n)\nfrom audiocraft.losses.loudnessloss import TFLoudnessRatio\nfrom audiocraft.losses.wmloss import WMMbLoss\nfrom tests.common_utils.wav_utils import get_white_noise\n\n\ndef test_mel_l1_loss():\n N, C, T = 2, 2, random.randrange(1000, 100_000)\n t1 = torch.randn(N, C, T)\n t2 = torch.randn(N, C, T)\n\n mel_l1 = MelSpectrogramL1Loss(sample_rate=22_050)\n loss = mel_l1(t1, t2)\n loss_same = mel_l1(t1, t1)\n\n assert isinstance(loss, torch.Tensor)\n assert isinstance(loss_same, torch.Tensor)\n assert loss_same.item() == 0.0\n\n\ndef test_msspec_loss():\n N, C, T = 2, 2, random.randrange(1000, 100_000)\n t1 = torch.randn(N, C, T)\n t2 = torch.randn(N, C, T)\n\n msspec = MultiScaleMelSpectrogramLoss(sample_rate=22_050)\n loss = msspec(t1, t2)\n loss_same = msspec(t1, t1)\n\n assert isinstance(loss, torch.Tensor)\n assert isinstance(loss_same, torch.Tensor)\n assert loss_same.item() == 0.0\n\n\ndef test_mrstft_loss():\n N, C, T = 2, 2, random.randrange(1000, 100_000)\n t1 = torch.randn(N, C, T)\n t2 = torch.randn(N, C, T)\n\n mrstft = MRSTFTLoss()\n loss = mrstft(t1, t2)\n\n assert isinstance(loss, torch.Tensor)\n\n\ndef test_sisnr_loss():\n N, C, T = 2, 2, random.randrange(1000, 100_000)\n t1 = torch.randn(N, C, T)\n t2 = torch.randn(N, C, T)\n\n sisnr = SISNR()\n loss = sisnr(t1, t2)\n\n assert isinstance(loss, torch.Tensor)\n\n\ndef test_stft_loss():\n N, C, T = 2, 2, random.randrange(1000, 100_000)\n t1 = torch.randn(N, C, T)\n t2 = torch.randn(N, C, T)\n\n mrstft = STFTLoss()\n loss = mrstft(t1, t2)\n\n assert isinstance(loss, torch.Tensor)\n\n\ndef test_wm_loss():\n N, nbits, T = 2, 16, random.randrange(1000, 100_000)\n positive = torch.randn(N, 2 + nbits, T)\n t2 = torch.randn(N, 1, T)\n message = torch.randn(N, nbits)\n\n wmloss = WMMbLoss(0.3, \"mse\")\n loss = wmloss(positive, None, t2, message)\n\n assert isinstance(loss, torch.Tensor)\n\n\ndef test_loudness_loss():\n sr = 16_000\n duration = 1.0\n wav = get_white_noise(1, int(sr * duration)).unsqueeze(0)\n tflrloss = TFLoudnessRatio(sample_rate=sr, n_bands=1)\n\n loss = tflrloss(wav, wav)\n assert isinstance(loss, torch.Tensor)\n" }, { "path": "tests/metrics/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "tests/metrics/test_pesq.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport julius\nimport pesq\nimport torch\nfrom audiocraft.metrics.pesq import PesqMetric\nfrom ..common_utils import TempDirMixin, get_batch_white_noise\n\n\ndef tensor_pesq(y_pred: torch.Tensor, y: torch.Tensor, sr: int):\n # pesq returns error if no speech is detected, so we catch it\n if sr != 16000:\n y_pred = julius.resample_frac(y_pred, sr, 16000)\n y = julius.resample_frac(y, sr, 16000)\n P, n = 0, 0\n for ii in range(y_pred.size(0)):\n try: # torchmetrics crashes when there is one error in the batch so doing it manually..\n P += pesq.pesq(16000, y[ii, 0].cpu().numpy(), y_pred[ii, 0].cpu().numpy())\n n += 1\n except pesq.NoUtterancesError: # this error can append when the sample don't contain speech\n pass\n p = P / n if n != 0 else 0.0\n return p\n\n\nclass TestPesq(TempDirMixin):\n\n def test(self):\n sample_rate = 16_000\n duration = 20\n channel = 1\n bs = 10\n wavs = get_batch_white_noise(bs, channel, int(sample_rate * duration))\n\n pesq_metric = PesqMetric(sample_rate=sample_rate)\n pesq1 = pesq_metric(wavs, wavs)\n print(f\"Pesq between 2 identical white noises: {pesq1}\")\n assert pesq1 > 1\n\n pesq2 = tensor_pesq(wavs, wavs, 16000)\n assert torch.allclose(pesq1, torch.tensor(pesq2))\n" }, { "path": "tests/models/test_audiogen.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport pytest\nimport torch\n\nfrom audiocraft.models import AudioGen\n\n\nclass TestAudioGenModel:\n def get_audiogen(self):\n ag = AudioGen.get_pretrained(name='debug', device='cpu')\n ag.set_generation_params(duration=2.0, extend_stride=2.)\n return ag\n\n def test_base(self):\n ag = self.get_audiogen()\n assert ag.frame_rate == 25\n assert ag.sample_rate == 16000\n assert ag.audio_channels == 1\n\n def test_generate_continuation(self):\n ag = self.get_audiogen()\n prompt = torch.randn(3, 1, 16000)\n wav = ag.generate_continuation(prompt, 16000)\n assert list(wav.shape) == [3, 1, 32000]\n\n prompt = torch.randn(2, 1, 16000)\n wav = ag.generate_continuation(\n prompt, 16000, ['youpi', 'lapin dort'])\n assert list(wav.shape) == [2, 1, 32000]\n\n prompt = torch.randn(2, 1, 16000)\n with pytest.raises(AssertionError):\n wav = ag.generate_continuation(\n prompt, 16000, ['youpi', 'lapin dort', 'one too many'])\n\n def test_generate(self):\n ag = self.get_audiogen()\n wav = ag.generate(\n ['youpi', 'lapin dort'])\n assert list(wav.shape) == [2, 1, 32000]\n\n def test_generate_long(self):\n ag = self.get_audiogen()\n ag.max_duration = 3.\n ag.set_generation_params(duration=4., extend_stride=2.)\n wav = ag.generate(\n ['youpi', 'lapin dort'])\n assert list(wav.shape) == [2, 1, 16000 * 4]\n" }, { "path": "tests/models/test_encodec_model.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport random\n\nimport numpy as np\nimport torch\n\nfrom audiocraft.models import EncodecModel\nfrom audiocraft.modules import SEANetEncoder, SEANetDecoder\nfrom audiocraft.quantization import DummyQuantizer\n\n\nclass TestEncodecModel:\n\n def _create_encodec_model(self,\n sample_rate: int,\n channels: int,\n dim: int = 5,\n n_filters: int = 3,\n n_residual_layers: int = 1,\n ratios: list = [5, 4, 3, 2],\n **kwargs):\n frame_rate = np.prod(ratios)\n encoder = SEANetEncoder(channels=channels, dimension=dim, n_filters=n_filters,\n n_residual_layers=n_residual_layers, ratios=ratios)\n decoder = SEANetDecoder(channels=channels, dimension=dim, n_filters=n_filters,\n n_residual_layers=n_residual_layers, ratios=ratios)\n quantizer = DummyQuantizer()\n model = EncodecModel(encoder, decoder, quantizer, frame_rate=frame_rate,\n sample_rate=sample_rate, channels=channels, **kwargs)\n return model\n\n def test_model(self):\n random.seed(1234)\n sample_rate = 24_000\n channels = 1\n model = self._create_encodec_model(sample_rate, channels)\n for _ in range(10):\n length = random.randrange(1, 10_000)\n x = torch.randn(2, channels, length)\n res = model(x)\n assert res.x.shape == x.shape\n\n def test_model_renorm(self):\n random.seed(1234)\n sample_rate = 24_000\n channels = 1\n model_nonorm = self._create_encodec_model(sample_rate, channels, renormalize=False)\n model_renorm = self._create_encodec_model(sample_rate, channels, renormalize=True)\n\n for _ in range(10):\n length = random.randrange(1, 10_000)\n x = torch.randn(2, channels, length)\n codes, scales = model_nonorm.encode(x)\n codes, scales = model_renorm.encode(x)\n assert scales is not None\n" }, { "path": "tests/models/test_multibanddiffusion.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport random\n\nimport numpy as np\nimport torch\nfrom audiocraft.models.multibanddiffusion import MultiBandDiffusion, DiffusionProcess\nfrom audiocraft.models import EncodecModel, DiffusionUnet\nfrom audiocraft.modules import SEANetEncoder, SEANetDecoder\nfrom audiocraft.modules.diffusion_schedule import NoiseSchedule\nfrom audiocraft.quantization import DummyQuantizer\n\n\nclass TestMBD:\n\n def _create_mbd(self,\n sample_rate: int,\n channels: int,\n n_filters: int = 3,\n n_residual_layers: int = 1,\n ratios: list = [5, 4, 3, 2],\n num_steps: int = 1000,\n codec_dim: int = 128,\n **kwargs):\n frame_rate = np.prod(ratios)\n encoder = SEANetEncoder(channels=channels, dimension=codec_dim, n_filters=n_filters,\n n_residual_layers=n_residual_layers, ratios=ratios)\n decoder = SEANetDecoder(channels=channels, dimension=codec_dim, n_filters=n_filters,\n n_residual_layers=n_residual_layers, ratios=ratios)\n quantizer = DummyQuantizer()\n compression_model = EncodecModel(encoder, decoder, quantizer, frame_rate=frame_rate,\n sample_rate=sample_rate, channels=channels, **kwargs)\n diffusion_model = DiffusionUnet(chin=channels, num_steps=num_steps, codec_dim=codec_dim)\n schedule = NoiseSchedule(device='cpu', num_steps=num_steps)\n DP = DiffusionProcess(model=diffusion_model, noise_schedule=schedule)\n mbd = MultiBandDiffusion(DPs=[DP], codec_model=compression_model)\n return mbd\n\n def test_model(self):\n random.seed(1234)\n sample_rate = 24_000\n channels = 1\n codec_dim = 128\n mbd = self._create_mbd(sample_rate=sample_rate, channels=channels, codec_dim=codec_dim)\n for _ in range(10):\n length = random.randrange(1, 10_000)\n x = torch.randn(2, channels, length)\n res = mbd.regenerate(x, sample_rate)\n assert res.shape == x.shape\n" }, { "path": "tests/models/test_musicgen.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport pytest\nimport torch\n\nfrom audiocraft.models import MusicGen\n\n\nclass TestMusicGenModel:\n def get_musicgen(self):\n mg = MusicGen.get_pretrained(name='debug', device='cpu')\n mg.set_generation_params(duration=2.0, extend_stride=2.)\n return mg\n\n def test_base(self):\n mg = self.get_musicgen()\n assert mg.frame_rate == 25\n assert mg.sample_rate == 32000\n assert mg.audio_channels == 1\n\n def test_generate_unconditional(self):\n mg = self.get_musicgen()\n wav = mg.generate_unconditional(3)\n assert list(wav.shape) == [3, 1, 64000]\n\n def test_generate_continuation(self):\n mg = self.get_musicgen()\n prompt = torch.randn(3, 1, 32000)\n wav = mg.generate_continuation(prompt, 32000)\n assert list(wav.shape) == [3, 1, 64000]\n\n prompt = torch.randn(2, 1, 32000)\n wav = mg.generate_continuation(\n prompt, 32000, ['youpi', 'lapin dort'])\n assert list(wav.shape) == [2, 1, 64000]\n\n prompt = torch.randn(2, 1, 32000)\n with pytest.raises(AssertionError):\n wav = mg.generate_continuation(\n prompt, 32000, ['youpi', 'lapin dort', 'one too many'])\n\n def test_generate(self):\n mg = self.get_musicgen()\n wav = mg.generate(\n ['youpi', 'lapin dort'])\n assert list(wav.shape) == [2, 1, 64000]\n\n def test_generate_long(self):\n mg = self.get_musicgen()\n mg.max_duration = 3.\n mg.set_generation_params(duration=4., extend_stride=2.)\n wav = mg.generate(\n ['youpi', 'lapin dort'])\n assert list(wav.shape) == [2, 1, 32000 * 4]\n\n def test_generate_two_step_cfg(self):\n mg = self.get_musicgen()\n mg.set_generation_params(duration=2.0, extend_stride=2., two_step_cfg=True)\n wav = mg.generate(\n ['youpi', 'lapin dort'])\n assert list(wav.shape) == [2, 1, 64000]\n" }, { "path": "tests/models/test_watermark.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\nimport torch\nfrom audiocraft.models.watermark import AudioSeal\nfrom tests.common_utils.wav_utils import get_white_noise\n\n\nclass TestWatermarkModel:\n\n def test_base(self):\n sr = 16_000\n duration = 1.0\n wav = get_white_noise(1, int(sr * duration)).unsqueeze(0)\n wm = AudioSeal.get_pretrained(name=\"base\")\n\n secret_message = torch.randint(0, 2, (1, 16), dtype=torch.int32)\n watermarked_wav = wm(wav, message=secret_message, sample_rate=sr, alpha=0.8)\n result = wm.detect_watermark(watermarked_wav)\n\n detected = (\n torch.count_nonzero(torch.gt(result[:, 1, :], 0.5)) / result.shape[-1]\n )\n detect_prob = detected.cpu().item() # type: ignore\n\n assert detect_prob >= 0.0\n" }, { "path": "tests/modules/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "tests/modules/test_activations.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\nfrom torch import nn\n\nfrom audiocraft.modules.activations import CustomGLU\n\n\nclass TestActivations:\n def test_custom_glu_calculation(self):\n\n activation = CustomGLU(nn.Identity())\n\n initial_shape = (4, 8, 8)\n\n part_a = torch.ones(initial_shape) * 2\n part_b = torch.ones(initial_shape) * -1\n input = torch.cat((part_a, part_b), dim=-1)\n\n output = activation(input)\n\n # ensure all dimensions match initial shape\n assert output.shape == initial_shape\n # ensure the gating was calculated correctly a * f(b)\n assert torch.all(output == -2).item()\n" }, { "path": "tests/modules/test_codebooks_patterns.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport pytest\nimport torch\n\nfrom audiocraft.modules.codebooks_patterns import (\n DelayedPatternProvider,\n ParallelPatternProvider,\n Pattern,\n UnrolledPatternProvider,\n)\n\n\nclass TestParallelPatternProvider:\n\n @pytest.mark.parametrize(\"n_q\", [1, 4, 32])\n @pytest.mark.parametrize(\"timesteps\", [0, 1, 16, 100])\n def test_get_pattern(self, n_q: int, timesteps: int):\n provider = ParallelPatternProvider(n_q)\n pattern = provider.get_pattern(timesteps)\n # + 1 to account for 1st step\n assert len(pattern.layout) == timesteps + 1\n\n @pytest.mark.parametrize(\"n_q\", [1, 4, 32])\n @pytest.mark.parametrize(\"timesteps\", [8, 16, 100])\n def test_pattern_content(self, n_q: int, timesteps: int):\n provider = ParallelPatternProvider(n_q)\n pattern = provider.get_pattern(timesteps)\n for s, v in enumerate(pattern.layout):\n for i, code in enumerate(v):\n assert i == code.q\n assert code.t == s - 1 # account for the 1st empty step\n\n @pytest.mark.parametrize(\"n_q\", [1, 4, 32])\n @pytest.mark.parametrize(\"timesteps\", [8, 16, 100])\n def test_pattern_max_delay(self, n_q: int, timesteps: int):\n provider = ParallelPatternProvider(n_q)\n pattern = provider.get_pattern(timesteps)\n assert pattern.max_delay == 0\n assert len(pattern.valid_layout) == len(pattern.layout) - pattern.max_delay\n\n\nclass TestDelayedPatternProvider:\n\n @pytest.mark.parametrize(\"n_q\", [1, 4, 32])\n @pytest.mark.parametrize(\"timesteps\", [0, 1, 16, 100])\n def test_get_pattern(self, n_q: int, timesteps: int):\n delays = [\n list(range(n_q)),\n [0] + [1] * (n_q - 1),\n [0] + [4] * (n_q - 1),\n ]\n for delay in delays:\n provider = DelayedPatternProvider(n_q, delay)\n pattern = provider.get_pattern(timesteps)\n # + 1 to account for 1st step\n assert len(pattern.layout) == timesteps + max(delay) + 1\n\n @pytest.mark.parametrize(\"n_q\", [1, 4, 32])\n @pytest.mark.parametrize(\"timesteps\", [8, 16, 100])\n def test_pattern_content(self, n_q: int, timesteps: int):\n provider = DelayedPatternProvider(n_q)\n pattern = provider.get_pattern(timesteps)\n for s, v in enumerate(pattern.layout):\n for i, code in enumerate(v):\n assert i == code.q\n assert code.t == max(0, s - code.q - 1)\n\n @pytest.mark.parametrize(\"timesteps\", [8, 16, 100])\n @pytest.mark.parametrize(\"delay\", [[0, 1, 2, 3], [0, 1, 1, 1], [0, 3, 3, 3], [0, 3]])\n def test_pattern_max_delay(self, timesteps: int, delay: list):\n provider = DelayedPatternProvider(len(delay), delay)\n pattern = provider.get_pattern(timesteps)\n assert pattern.max_delay == max(delay)\n assert len(pattern.valid_layout) == len(pattern.layout) - pattern.max_delay\n\n\nclass TestUnrolledPatternProvider:\n\n @pytest.mark.parametrize(\"timesteps\", [0, 1, 16])\n @pytest.mark.parametrize(\"flattening\", [[0, 1, 2], [0, 1, 1]])\n @pytest.mark.parametrize(\"delays\", [[0, 0, 0], [0, 5, 5]])\n def test_get_pattern(self, timesteps: int, flattening: list, delays: list):\n n_q = len(flattening)\n max_delay = max(delays)\n provider = UnrolledPatternProvider(n_q, flattening, delays)\n pattern = provider.get_pattern(timesteps)\n assert len(pattern.layout) == provider.num_virtual_steps(timesteps) + max_delay\n\n @pytest.mark.parametrize(\"timesteps\", [0, 1, 16])\n @pytest.mark.parametrize(\"flattening\", [[0, 1, 2], [0, 1, 1]])\n @pytest.mark.parametrize(\"delays\", [[0, 0, 0], [0, 5, 5]])\n def test_pattern_max_delay(self, timesteps: int, flattening: list, delays: list):\n n_q = len(flattening)\n max_delay = max(delays)\n provider = UnrolledPatternProvider(n_q, flattening, delays)\n pattern = provider.get_pattern(timesteps)\n assert pattern.max_delay == max_delay\n\n\nclass TestPattern:\n\n def ref_build_pattern_sequence(self, z: torch.Tensor, pattern: Pattern, special_token: int):\n \"\"\"Reference method to build the sequence from the pattern without using fancy scatter.\"\"\"\n bs, n_q, T = z.shape\n z = z.cpu().numpy()\n assert n_q == pattern.n_q\n assert T <= pattern.timesteps\n inp = torch.full((bs, n_q, len(pattern.layout)), special_token, dtype=torch.long).numpy()\n inp[:] = special_token\n for s, v in enumerate(pattern.layout):\n for (t, q) in v:\n if t < T:\n inp[:, q, s] = z[:, q, t]\n return torch.from_numpy(inp)\n\n def ref_revert_pattern_sequence(self, z: torch.Tensor, pattern: Pattern, special_token: int):\n \"\"\"Reference method to revert the sequence from the pattern without using fancy scatter.\"\"\"\n z = z.cpu().numpy()\n bs, n_q, S = z.shape\n assert pattern.n_q == n_q\n inp = torch.full((bs, pattern.n_q, pattern.timesteps), special_token, dtype=torch.long).numpy()\n inp[:] = special_token\n for s, v in enumerate(pattern.layout):\n for (t, q) in v:\n if t < pattern.timesteps:\n inp[:, q, t] = z[:, q, s]\n return torch.from_numpy(inp)\n\n def ref_revert_pattern_logits(self, z: torch.Tensor, pattern: Pattern, special_token: float):\n \"\"\"Reference method to revert the logits from the pattern without using fancy scatter.\"\"\"\n z = z.cpu().numpy()\n bs, card, n_q, S = z.shape\n assert pattern.n_q == n_q\n ref_layout = pattern.layout\n inp = torch.full((bs, card, pattern.n_q, pattern.timesteps), special_token, dtype=torch.float).numpy()\n inp[:] = special_token\n for s, v in enumerate(ref_layout[1:]):\n if s < S:\n for (t, q) in v:\n if t < pattern.timesteps:\n inp[:, :, q, t] = z[:, :, q, s]\n return torch.from_numpy(inp)\n\n def _get_pattern_providers(self, n_q: int):\n pattern_provider_1 = ParallelPatternProvider(n_q)\n pattern_provider_2 = DelayedPatternProvider(n_q, list(range(n_q)))\n pattern_provider_3 = DelayedPatternProvider(n_q, [0] + [1] * (n_q - 1))\n pattern_provider_4 = UnrolledPatternProvider(\n n_q, flattening=list(range(n_q)), delays=[0] * n_q\n )\n pattern_provider_5 = UnrolledPatternProvider(\n n_q, flattening=[0] + [1] * (n_q - 1), delays=[0] * n_q\n )\n pattern_provider_6 = UnrolledPatternProvider(\n n_q, flattening=[0] + [1] * (n_q - 1), delays=[0] + [5] * (n_q - 1)\n )\n return [\n pattern_provider_1,\n pattern_provider_2,\n pattern_provider_3,\n pattern_provider_4,\n pattern_provider_5,\n pattern_provider_6,\n ]\n\n @pytest.mark.parametrize(\"n_q\", [1, 4, 32])\n @pytest.mark.parametrize(\"timesteps\", [16, 72])\n def test_build_pattern_sequence(self, n_q: int, timesteps: int):\n bs = 2\n card = 256\n special_token = card\n\n pattern_providers = self._get_pattern_providers(n_q)\n for pattern_provider in pattern_providers:\n pattern = pattern_provider.get_pattern(timesteps)\n # we can correctly build the sequence from the pattern\n z = torch.randint(0, card, (bs, n_q, timesteps))\n ref_res = self.ref_build_pattern_sequence(z, pattern, special_token)\n res, indexes, mask = pattern.build_pattern_sequence(z, special_token)\n assert (res == ref_res).float().mean() == 1.0\n\n # expected assertion fails on the number of timesteps\n invalid_timesteps = [timesteps + 1]\n if pattern.num_sequence_steps != pattern.timesteps:\n invalid_timesteps.append(pattern.num_sequence_steps)\n for i_timesteps in invalid_timesteps:\n z2 = torch.randint(0, card, (bs, n_q, i_timesteps))\n with pytest.raises(AssertionError):\n pattern.build_pattern_sequence(z2, special_token)\n\n # expected assertion fails on the number of codebooks\n invalid_qs = [0, n_q - 1, n_q + 1]\n for i_q in invalid_qs:\n z3 = torch.randint(0, card, (bs, i_q, timesteps))\n with pytest.raises(AssertionError):\n pattern.build_pattern_sequence(z3, special_token)\n\n @pytest.mark.parametrize(\"n_q\", [1, 4, 32])\n @pytest.mark.parametrize(\"timesteps\", [16, 72])\n def test_revert_pattern_sequence(self, n_q: int, timesteps: int):\n bs = 2\n card = 256\n special_token = card\n\n pattern_providers = self._get_pattern_providers(n_q)\n for pattern_provider in pattern_providers:\n pattern = pattern_provider.get_pattern(timesteps)\n # this works assuming previous tests are successful\n z = torch.randint(0, card, (bs, n_q, timesteps))\n s = self.ref_build_pattern_sequence(z, pattern, special_token)\n ref_out = self.ref_revert_pattern_sequence(s, pattern, special_token)\n # ensure our reference script retrieve the original sequence\n assert z.shape == ref_out.shape\n assert (z == ref_out).float().mean() == 1.0\n # now we can test the scatter version\n out, indexes, mask = pattern.revert_pattern_sequence(s, special_token)\n assert out.shape == ref_out.shape\n assert (out == ref_out).float().mean() == 1.0\n\n @pytest.mark.parametrize(\"n_q\", [1, 4, 32])\n @pytest.mark.parametrize(\"timesteps\", [16, 72])\n @pytest.mark.parametrize(\"card\", [1, 2, 256, 1024])\n def test_revert_pattern_logits(self, n_q: int, timesteps: int, card: int):\n bs = 2\n special_token = card\n logits_special_token = float('nan')\n\n pattern_providers = self._get_pattern_providers(n_q)\n for pattern_provider in pattern_providers:\n pattern = pattern_provider.get_pattern(timesteps)\n # this works assuming previous tests are successful\n z = torch.randint(0, card, (bs, n_q, timesteps))\n s = self.ref_build_pattern_sequence(z, pattern, special_token)\n logits = torch.randn((bs, card, n_q, s.shape[-1]))\n ref_out = self.ref_revert_pattern_logits(logits, pattern, logits_special_token)\n # ensure our reference script retrieve the original sequence\n assert ref_out.shape == torch.Size([bs, card, n_q, timesteps])\n # now we can test the scatter version\n out, indexes, mask = pattern.revert_pattern_logits(logits, logits_special_token)\n assert out.shape == ref_out.shape\n assert (out == ref_out).float().mean() == 1.0\n" }, { "path": "tests/modules/test_conv.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom itertools import product\nimport math\nimport random\n\nimport pytest\nimport torch\nfrom torch import nn\n\nfrom audiocraft.modules import (\n NormConv1d,\n NormConvTranspose1d,\n StreamableConv1d,\n StreamableConvTranspose1d,\n pad1d,\n unpad1d,\n)\n\n\ndef test_get_extra_padding_for_conv1d():\n # TODO: Implement me!\n pass\n\n\ndef test_pad1d_zeros():\n x = torch.randn(1, 1, 20)\n\n xp1 = pad1d(x, (0, 5), mode='constant', value=0.)\n assert xp1.shape[-1] == 25\n xp2 = pad1d(x, (5, 5), mode='constant', value=0.)\n assert xp2.shape[-1] == 30\n xp3 = pad1d(x, (0, 0), mode='constant', value=0.)\n assert xp3.shape[-1] == 20\n xp4 = pad1d(x, (10, 30), mode='constant', value=0.)\n assert xp4.shape[-1] == 60\n\n with pytest.raises(AssertionError):\n pad1d(x, (-1, 0), mode='constant', value=0.)\n\n with pytest.raises(AssertionError):\n pad1d(x, (0, -1), mode='constant', value=0.)\n\n with pytest.raises(AssertionError):\n pad1d(x, (-1, -1), mode='constant', value=0.)\n\n\ndef test_pad1d_reflect():\n x = torch.randn(1, 1, 20)\n\n xp1 = pad1d(x, (0, 5), mode='reflect', value=0.)\n assert xp1.shape[-1] == 25\n xp2 = pad1d(x, (5, 5), mode='reflect', value=0.)\n assert xp2.shape[-1] == 30\n xp3 = pad1d(x, (0, 0), mode='reflect', value=0.)\n assert xp3.shape[-1] == 20\n xp4 = pad1d(x, (10, 30), mode='reflect', value=0.)\n assert xp4.shape[-1] == 60\n\n with pytest.raises(AssertionError):\n pad1d(x, (-1, 0), mode='reflect', value=0.)\n\n with pytest.raises(AssertionError):\n pad1d(x, (0, -1), mode='reflect', value=0.)\n\n with pytest.raises(AssertionError):\n pad1d(x, (-1, -1), mode='reflect', value=0.)\n\n\ndef test_unpad1d():\n x = torch.randn(1, 1, 20)\n\n u1 = unpad1d(x, (5, 5))\n assert u1.shape[-1] == 10\n u2 = unpad1d(x, (0, 5))\n assert u2.shape[-1] == 15\n u3 = unpad1d(x, (5, 0))\n assert u3.shape[-1] == 15\n u4 = unpad1d(x, (0, 0))\n assert u4.shape[-1] == x.shape[-1]\n\n with pytest.raises(AssertionError):\n unpad1d(x, (-1, 0))\n\n with pytest.raises(AssertionError):\n unpad1d(x, (0, -1))\n\n with pytest.raises(AssertionError):\n unpad1d(x, (-1, -1))\n\n\nclass TestNormConv1d:\n\n def test_norm_conv1d_modules(self):\n N, C, T = 2, 2, random.randrange(1, 100_000)\n t0 = torch.randn(N, C, T)\n\n C_out, kernel_size, stride = 1, 4, 1\n expected_out_length = int((T - kernel_size) / stride + 1)\n wn_conv = NormConv1d(C, 1, kernel_size=4, norm='weight_norm')\n gn_conv = NormConv1d(C, 1, kernel_size=4, norm='time_group_norm')\n nn_conv = NormConv1d(C, 1, kernel_size=4, norm='none')\n\n assert isinstance(wn_conv.norm, nn.Identity)\n assert isinstance(wn_conv.conv, nn.Conv1d)\n\n assert isinstance(gn_conv.norm, nn.GroupNorm)\n assert isinstance(gn_conv.conv, nn.Conv1d)\n\n assert isinstance(nn_conv.norm, nn.Identity)\n assert isinstance(nn_conv.conv, nn.Conv1d)\n\n for conv_layer in [wn_conv, gn_conv, nn_conv]:\n out = conv_layer(t0)\n assert isinstance(out, torch.Tensor)\n assert list(out.shape) == [N, C_out, expected_out_length]\n\n\nclass TestNormConvTranspose1d:\n\n def test_normalizations(self):\n N, C, T = 2, 2, random.randrange(1, 100_000)\n t0 = torch.randn(N, C, T)\n\n C_out, kernel_size, stride = 1, 4, 1\n expected_out_length = (T - 1) * stride + (kernel_size - 1) + 1\n\n wn_convtr = NormConvTranspose1d(C, C_out, kernel_size=kernel_size, stride=stride, norm='weight_norm')\n gn_convtr = NormConvTranspose1d(C, C_out, kernel_size=kernel_size, stride=stride, norm='time_group_norm')\n nn_convtr = NormConvTranspose1d(C, C_out, kernel_size=kernel_size, stride=stride, norm='none')\n\n assert isinstance(wn_convtr.norm, nn.Identity)\n assert isinstance(wn_convtr.convtr, nn.ConvTranspose1d)\n\n assert isinstance(gn_convtr.norm, nn.GroupNorm)\n assert isinstance(gn_convtr.convtr, nn.ConvTranspose1d)\n\n assert isinstance(nn_convtr.norm, nn.Identity)\n assert isinstance(nn_convtr.convtr, nn.ConvTranspose1d)\n\n for convtr_layer in [wn_convtr, gn_convtr, nn_convtr]:\n out = convtr_layer(t0)\n assert isinstance(out, torch.Tensor)\n assert list(out.shape) == [N, C_out, expected_out_length]\n\n\nclass TestStreamableConv1d:\n\n def get_streamable_conv1d_output_length(self, length, kernel_size, stride, dilation):\n # StreamableConv1d internally pads to make sure that the last window is full\n padding_total = (kernel_size - 1) * dilation - (stride - 1)\n n_frames = (length - kernel_size + padding_total) / stride + 1\n ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)\n return ideal_length // stride\n\n def test_streamable_conv1d(self):\n N, C, T = 2, 2, random.randrange(1, 100_000)\n t0 = torch.randn(N, C, T)\n C_out = 1\n\n # conv params are [(kernel_size, stride, dilation)]\n conv_params = [(4, 1, 1), (4, 2, 1), (3, 1, 3), (10, 5, 1), (3, 2, 3)]\n for causal, (kernel_size, stride, dilation) in product([False, True], conv_params):\n expected_out_length = self.get_streamable_conv1d_output_length(T, kernel_size, stride, dilation)\n sconv = StreamableConv1d(C, C_out, kernel_size=kernel_size, stride=stride, dilation=dilation, causal=causal)\n out = sconv(t0)\n assert isinstance(out, torch.Tensor)\n print(list(out.shape), [N, C_out, expected_out_length])\n assert list(out.shape) == [N, C_out, expected_out_length]\n\n\nclass TestStreamableConvTranspose1d:\n\n def get_streamable_convtr1d_output_length(self, length, kernel_size, stride):\n padding_total = (kernel_size - stride)\n return (length - 1) * stride - padding_total + (kernel_size - 1) + 1\n\n def test_streamable_convtr1d(self):\n N, C, T = 2, 2, random.randrange(1, 100_000)\n t0 = torch.randn(N, C, T)\n\n C_out = 1\n\n with pytest.raises(AssertionError):\n StreamableConvTranspose1d(C, C_out, kernel_size=4, causal=False, trim_right_ratio=0.5)\n StreamableConvTranspose1d(C, C_out, kernel_size=4, causal=True, trim_right_ratio=-1.)\n StreamableConvTranspose1d(C, C_out, kernel_size=4, causal=True, trim_right_ratio=2)\n\n # causal params are [(causal, trim_right)]\n causal_params = [(False, 1.0), (True, 1.0), (True, 0.5), (True, 0.0)]\n # conv params are [(kernel_size, stride)]\n conv_params = [(4, 1), (4, 2), (3, 1), (10, 5)]\n for ((causal, trim_right_ratio), (kernel_size, stride)) in product(causal_params, conv_params):\n expected_out_length = self.get_streamable_convtr1d_output_length(T, kernel_size, stride)\n sconvtr = StreamableConvTranspose1d(C, C_out, kernel_size=kernel_size, stride=stride,\n causal=causal, trim_right_ratio=trim_right_ratio)\n out = sconvtr(t0)\n assert isinstance(out, torch.Tensor)\n assert list(out.shape) == [N, C_out, expected_out_length]\n" }, { "path": "tests/modules/test_lstm.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport random\nimport torch\n\nfrom audiocraft.modules.lstm import StreamableLSTM\n\n\nclass TestStreamableLSTM:\n\n def test_lstm(self):\n B, C, T = 4, 2, random.randint(1, 100)\n\n lstm = StreamableLSTM(C, 3, skip=False)\n x = torch.randn(B, C, T)\n y = lstm(x)\n\n print(y.shape)\n assert y.shape == torch.Size([B, C, T])\n\n def test_lstm_skip(self):\n B, C, T = 4, 2, random.randint(1, 100)\n\n lstm = StreamableLSTM(C, 3, skip=True)\n x = torch.randn(B, C, T)\n y = lstm(x)\n\n assert y.shape == torch.Size([B, C, T])\n" }, { "path": "tests/modules/test_rope.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\n\nfrom audiocraft.modules.rope import RotaryEmbedding\nfrom audiocraft.modules.transformer import StreamingTransformer, set_efficient_attention_backend\n\n\ndef test_rope():\n set_efficient_attention_backend('torch')\n B, T, H, C = 8, 75, 16, 128\n\n rope = RotaryEmbedding(dim=C)\n xq = torch.rand((B, T, H, C))\n xk = torch.rand((B, T, H, C))\n xq_out, xk_out = rope.rotate_qk(xq, xk, start=7)\n\n assert list(xq_out.shape) == [B, T, H, C]\n assert list(xk_out.shape) == [B, T, H, C]\n\n\ndef test_rope_io_dtypes():\n set_efficient_attention_backend('torch')\n B, T, H, C = 8, 75, 16, 128\n\n rope_32 = RotaryEmbedding(dim=C, dtype=torch.float32)\n rope_64 = RotaryEmbedding(dim=C, dtype=torch.float64)\n\n # Test bfloat16 inputs w/ both 32 and 64 precision rope.\n xq_16 = torch.rand((B, T, H, C)).to(torch.bfloat16)\n xk_16 = torch.rand((B, T, H, C)).to(torch.bfloat16)\n xq_out, xk_out = rope_32.rotate_qk(xq_16, xk_16)\n assert xq_out.dtype == torch.bfloat16\n xq_out, xk_out = rope_64.rotate_qk(xq_16, xk_16)\n assert xq_out.dtype == torch.bfloat16\n\n # Test float32 inputs w/ both 32 and 64 precision rope.\n xq_32 = torch.rand((B, T, H, C)).to(torch.float32)\n xk_32 = torch.rand((B, T, H, C)).to(torch.float32)\n xq_out, xk_out = rope_32.rotate_qk(xq_32, xk_32)\n assert xq_out.dtype == torch.float32\n xq_out, xk_out = rope_64.rotate_qk(xq_32, xk_32)\n assert xq_out.dtype == torch.float32\n\n\ndef test_transformer_with_rope():\n set_efficient_attention_backend('torch')\n torch.manual_seed(1234)\n for pos in ['rope', 'sin_rope']:\n tr = StreamingTransformer(\n 16, 4, 2, custom=True, dropout=0., layer_scale=0.1,\n positional_embedding=pos)\n tr.eval()\n steps = 12\n x = torch.randn(3, steps, 16)\n\n out = tr(x)\n assert list(out.shape) == list(x.shape)\n\n\n@torch.no_grad()\ndef test_rope_streaming():\n set_efficient_attention_backend('torch')\n torch.manual_seed(1234)\n tr = StreamingTransformer(\n 16, 4, 2, causal=True, dropout=0.,\n custom=True, positional_embedding='rope')\n tr.eval()\n steps = 12\n x = torch.randn(3, steps, 16)\n\n ref = tr(x)\n\n with tr.streaming():\n outs = []\n frame_sizes = [1] * steps\n\n for frame_size in frame_sizes:\n frame = x[:, :frame_size]\n x = x[:, frame_size:]\n outs.append(tr(frame))\n\n out = torch.cat(outs, dim=1)\n assert list(out.shape) == [3, steps, 16]\n delta = torch.norm(out - ref) / torch.norm(out)\n assert delta < 1e-6, delta\n\n\n@torch.no_grad()\ndef test_rope_streaming_past_context():\n set_efficient_attention_backend('torch')\n torch.manual_seed(1234)\n\n for context in [None, 10]:\n tr = StreamingTransformer(\n 16, 4, 1 if context else 2,\n causal=True, past_context=context, custom=True,\n dropout=0., positional_embedding='rope')\n tr.eval()\n\n steps = 20\n x = torch.randn(3, steps, 16)\n ref = tr(x)\n\n with tr.streaming():\n outs = []\n frame_sizes = [1] * steps\n\n for frame_size in frame_sizes:\n frame = x[:, :frame_size]\n x = x[:, frame_size:]\n outs.append(tr(frame))\n\n out = torch.cat(outs, dim=1)\n assert list(out.shape) == [3, steps, 16]\n delta = torch.norm(out - ref) / torch.norm(out)\n assert delta < 1e-6, delta\n\n\ndef test_rope_memory_efficient():\n set_efficient_attention_backend('torch')\n torch.manual_seed(1234)\n tr = StreamingTransformer(\n 16, 4, 2, custom=True, dropout=0., layer_scale=0.1,\n positional_embedding='rope')\n tr_mem_efficient = StreamingTransformer(\n 16, 4, 2, dropout=0., memory_efficient=True, layer_scale=0.1,\n positional_embedding='rope')\n tr_mem_efficient.load_state_dict(tr.state_dict())\n tr.eval()\n steps = 12\n x = torch.randn(3, steps, 16)\n\n with torch.no_grad():\n y = tr(x)\n y2 = tr_mem_efficient(x)\n # Check at float precision b/c this is the rope default.\n assert torch.allclose(y, y2, atol=1e-7), (y - y2).norm()\n\n\ndef test_rope_with_xpos():\n set_efficient_attention_backend('torch')\n B, T, H, C = 8, 75, 16, 128\n\n rope = RotaryEmbedding(dim=C, xpos=True)\n xq = torch.rand((B, T, H, C))\n xk = torch.rand((B, T, H, C))\n xq_out, xk_out = rope.rotate_qk(xq, xk, start=7)\n\n assert list(xq_out.shape) == [B, T, H, C]\n assert list(xk_out.shape) == [B, T, H, C]\n\n\ndef test_positional_scale():\n set_efficient_attention_backend('torch')\n B, T, H, C = 8, 75, 16, 128\n\n rope = RotaryEmbedding(dim=C, xpos=True, scale=0.0)\n xq = torch.rand((B, T, H, C))\n xk = torch.rand((B, T, H, C))\n xq_out, xk_out = rope.rotate_qk(xq, xk, start=7)\n\n assert torch.allclose(xq, xq_out)\n assert torch.allclose(xk, xk_out)\n" }, { "path": "tests/modules/test_seanet.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom itertools import product\n\nimport pytest\nimport torch\n\nfrom audiocraft.modules.seanet import SEANetEncoder, SEANetDecoder, SEANetResnetBlock\nfrom audiocraft.modules import StreamableConv1d, StreamableConvTranspose1d\n\n\nclass TestSEANetModel:\n\n def test_base(self):\n encoder = SEANetEncoder()\n decoder = SEANetDecoder()\n\n x = torch.randn(1, 1, 24000)\n z = encoder(x)\n assert list(z.shape) == [1, 128, 75], z.shape\n y = decoder(z)\n assert y.shape == x.shape, (x.shape, y.shape)\n\n def test_causal(self):\n encoder = SEANetEncoder(causal=True)\n decoder = SEANetDecoder(causal=True)\n x = torch.randn(1, 1, 24000)\n\n z = encoder(x)\n assert list(z.shape) == [1, 128, 75], z.shape\n y = decoder(z)\n assert y.shape == x.shape, (x.shape, y.shape)\n\n def test_conv_skip_connection(self):\n encoder = SEANetEncoder(true_skip=False)\n decoder = SEANetDecoder(true_skip=False)\n\n x = torch.randn(1, 1, 24000)\n z = encoder(x)\n assert list(z.shape) == [1, 128, 75], z.shape\n y = decoder(z)\n assert y.shape == x.shape, (x.shape, y.shape)\n\n def test_seanet_encoder_decoder_final_act(self):\n encoder = SEANetEncoder(true_skip=False)\n decoder = SEANetDecoder(true_skip=False, final_activation='Tanh')\n\n x = torch.randn(1, 1, 24000)\n z = encoder(x)\n assert list(z.shape) == [1, 128, 75], z.shape\n y = decoder(z)\n assert y.shape == x.shape, (x.shape, y.shape)\n\n def _check_encoder_blocks_norm(self, encoder: SEANetEncoder, n_disable_blocks: int, norm: str):\n n_blocks = 0\n for layer in encoder.model:\n if isinstance(layer, StreamableConv1d):\n n_blocks += 1\n assert layer.conv.norm_type == 'none' if n_blocks <= n_disable_blocks else norm\n elif isinstance(layer, SEANetResnetBlock):\n for resnet_layer in layer.block:\n if isinstance(resnet_layer, StreamableConv1d):\n # here we add + 1 to n_blocks as we increment n_blocks just after the block\n assert resnet_layer.conv.norm_type == 'none' if (n_blocks + 1) <= n_disable_blocks else norm\n\n def test_encoder_disable_norm(self):\n n_residuals = [0, 1, 3]\n disable_blocks = [0, 1, 2, 3, 4, 5, 6]\n norms = ['weight_norm', 'none']\n for n_res, disable_blocks, norm in product(n_residuals, disable_blocks, norms):\n encoder = SEANetEncoder(n_residual_layers=n_res, norm=norm,\n disable_norm_outer_blocks=disable_blocks)\n self._check_encoder_blocks_norm(encoder, disable_blocks, norm)\n\n def _check_decoder_blocks_norm(self, decoder: SEANetDecoder, n_disable_blocks: int, norm: str):\n n_blocks = 0\n for layer in decoder.model:\n if isinstance(layer, StreamableConv1d):\n n_blocks += 1\n assert layer.conv.norm_type == 'none' if (decoder.n_blocks - n_blocks) < n_disable_blocks else norm\n elif isinstance(layer, StreamableConvTranspose1d):\n n_blocks += 1\n assert layer.convtr.norm_type == 'none' if (decoder.n_blocks - n_blocks) < n_disable_blocks else norm\n elif isinstance(layer, SEANetResnetBlock):\n for resnet_layer in layer.block:\n if isinstance(resnet_layer, StreamableConv1d):\n assert resnet_layer.conv.norm_type == 'none' \\\n if (decoder.n_blocks - n_blocks) < n_disable_blocks else norm\n\n def test_decoder_disable_norm(self):\n n_residuals = [0, 1, 3]\n disable_blocks = [0, 1, 2, 3, 4, 5, 6]\n norms = ['weight_norm', 'none']\n for n_res, disable_blocks, norm in product(n_residuals, disable_blocks, norms):\n decoder = SEANetDecoder(n_residual_layers=n_res, norm=norm,\n disable_norm_outer_blocks=disable_blocks)\n self._check_decoder_blocks_norm(decoder, disable_blocks, norm)\n\n def test_disable_norm_raises_exception(self):\n # Invalid disable_norm_outer_blocks values raise exceptions\n with pytest.raises(AssertionError):\n SEANetEncoder(disable_norm_outer_blocks=-1)\n\n with pytest.raises(AssertionError):\n SEANetEncoder(ratios=[1, 1, 2, 2], disable_norm_outer_blocks=7)\n\n with pytest.raises(AssertionError):\n SEANetDecoder(disable_norm_outer_blocks=-1)\n\n with pytest.raises(AssertionError):\n SEANetDecoder(ratios=[1, 1, 2, 2], disable_norm_outer_blocks=7)\n" }, { "path": "tests/modules/test_transformer.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom itertools import product\n\nimport pytest\nimport torch\n\nfrom audiocraft.modules.transformer import (\n StreamingMultiheadAttention, StreamingTransformer, set_efficient_attention_backend)\n\n\ndef test_transformer_causal_streaming():\n torch.manual_seed(1234)\n\n for context, custom in product([None, 10], [False, True]):\n # Test that causality and receptive fields are properly handled.\n # looking at the gradients\n tr = StreamingTransformer(\n 16, 4, 1 if context else 2,\n causal=True, past_context=context, custom=custom,\n dropout=0.)\n steps = 20\n for k in [0, 10, 15, 19]:\n x = torch.randn(4, steps, 16, requires_grad=True)\n y = tr(x)\n y[:, k].abs().sum().backward()\n if k + 1 < steps:\n assert torch.allclose(x.grad[:, k + 1:], torch.tensor(0.)), x.grad[:, k + 1:].norm()\n assert not torch.allclose(x.grad[:, :k + 1], torch.tensor(0.)), x.grad[:, :k + 1].norm()\n if context is not None and k > context:\n limit = k - context - 1\n assert torch.allclose(x.grad[:, :limit],\n torch.tensor(0.)), x.grad[:, :limit].norm()\n\n # Now check that streaming gives the same result at batch eval.\n x = torch.randn(4, steps, 16)\n y = tr(x)\n ys = []\n with tr.streaming():\n for k in range(steps):\n chunk = x[:, k:k + 1, :]\n ys.append(tr(chunk))\n y_stream = torch.cat(ys, dim=1)\n delta = torch.norm(y_stream - y) / torch.norm(y)\n assert delta < 1e-6, delta\n\n\ndef test_transformer_vs_pytorch():\n torch.manual_seed(1234)\n # Check that in the non causal setting, we get the same result as\n # PyTorch Transformer encoder.\n for custom in [False, True]:\n tr = StreamingTransformer(\n 16, 4, 2,\n causal=False, custom=custom, dropout=0., positional_scale=0.)\n layer = torch.nn.TransformerEncoderLayer(16, 4, dropout=0., batch_first=True)\n tr_ref = torch.nn.TransformerEncoder(layer, 2)\n tr.load_state_dict(tr_ref.state_dict())\n\n x = torch.randn(4, 20, 16)\n y = tr(x)\n y2 = tr_ref(x)\n delta = torch.norm(y2 - y) / torch.norm(y)\n assert delta < 1e-6, delta\n\n\ndef test_streaming_api():\n tr = StreamingTransformer(16, 4, 2, causal=True, dropout=0.)\n tr.eval()\n steps = 12\n x = torch.randn(1, steps, 16)\n\n with torch.no_grad():\n with tr.streaming():\n _ = tr(x[:, :1])\n state = {k: v.clone() for k, v in tr.get_streaming_state().items()}\n y = tr(x[:, 1:2])\n tr.set_streaming_state(state)\n y2 = tr(x[:, 1:2])\n assert torch.allclose(y, y2), (y - y2).norm()\n assert tr.flush() is None\n\n\ndef test_memory_efficient():\n for backend in ['torch']:\n torch.manual_seed(1234)\n set_efficient_attention_backend(backend)\n\n tr = StreamingTransformer(\n 16, 4, 2, custom=True, dropout=0., layer_scale=0.1)\n tr_mem_efficient = StreamingTransformer(\n 16, 4, 2, dropout=0., memory_efficient=True, layer_scale=0.1)\n tr_mem_efficient.load_state_dict(tr.state_dict())\n tr.eval()\n steps = 12\n x = torch.randn(3, steps, 16)\n\n with torch.no_grad():\n y = tr(x)\n y2 = tr_mem_efficient(x)\n assert torch.allclose(y, y2), ((y - y2).norm(), backend)\n\n\ndef test_attention_as_float32():\n torch.manual_seed(1234)\n cases = [\n {'custom': True},\n {'custom': False},\n ]\n for case in cases:\n tr = StreamingTransformer(16, 4, 2, dropout=0., dtype=torch.bfloat16, **case)\n tr_float32 = StreamingTransformer(\n 16, 4, 2, dropout=0., attention_as_float32=True, dtype=torch.bfloat16, **case)\n if not case['custom']:\n # we are not using autocast here because it doesn't really\n # work as expected on CPU, so we have to manually cast the weights of the MHA.\n for layer in tr_float32.layers:\n layer.self_attn.mha.to(torch.float32)\n tr_float32.load_state_dict(tr.state_dict())\n steps = 12\n x = torch.randn(3, steps, 16, dtype=torch.bfloat16)\n\n with torch.no_grad():\n y = tr(x)\n y2 = tr_float32(x)\n assert not torch.allclose(y, y2), (y - y2).norm()\n\n\n@torch.no_grad()\ndef test_streaming_memory_efficient():\n for backend in ['torch']:\n torch.manual_seed(1234)\n set_efficient_attention_backend(backend)\n tr = StreamingTransformer(16, 4, 2, causal=True, dropout=0., custom=True)\n tr_mem_efficient = StreamingTransformer(\n 16, 4, 2, dropout=0., memory_efficient=True, causal=True)\n tr.load_state_dict(tr_mem_efficient.state_dict())\n tr.eval()\n tr_mem_efficient.eval()\n steps = 12\n x = torch.randn(3, steps, 16)\n\n ref = tr(x)\n\n with tr_mem_efficient.streaming():\n outs = []\n # frame_sizes = [2] + [1] * (steps - 2)\n frame_sizes = [1] * steps\n\n for frame_size in frame_sizes:\n frame = x[:, :frame_size]\n x = x[:, frame_size:]\n outs.append(tr_mem_efficient(frame))\n\n out = torch.cat(outs, dim=1)\n delta = torch.norm(out - ref) / torch.norm(out)\n assert delta < 1e-6, delta\n\n\ndef test_cross_attention():\n torch.manual_seed(1234)\n for norm_first in [True, False]:\n m = StreamingTransformer(\n 16, 4, 2, cross_attention=False, norm_first=norm_first, dropout=0., custom=True)\n m_cross = StreamingTransformer(\n 16, 4, 2, cross_attention=True, norm_first=norm_first, dropout=0., custom=True)\n m_cross.load_state_dict(m.state_dict(), strict=False)\n x = torch.randn(2, 5, 16)\n cross_x = torch.randn(2, 3, 16)\n y_ref = m(x)\n y_cross_zero = m_cross(x, cross_attention_src=0 * cross_x)\n # With norm_first, the two should be exactly the same,\n # but with norm_first=False, we get 2 normalization in a row\n # and the epsilon value leads to a tiny change.\n atol = 0. if norm_first else 1e-6\n print((y_ref - y_cross_zero).norm() / y_ref.norm())\n assert torch.allclose(y_ref, y_cross_zero, atol=atol)\n\n # We now expect a difference even with a generous atol of 1e-2.\n y_cross = m_cross(x, cross_attention_src=cross_x)\n assert not torch.allclose(y_cross, y_cross_zero, atol=1e-2)\n\n with pytest.raises(AssertionError):\n _ = m_cross(x)\n _ = m(x, cross_attention_src=cross_x)\n\n\ndef test_cross_attention_compat():\n torch.manual_seed(1234)\n num_heads = 2\n dim = num_heads * 64\n with pytest.raises(AssertionError):\n StreamingMultiheadAttention(dim, num_heads, causal=True, cross_attention=True)\n\n cross_attn = StreamingMultiheadAttention(\n dim, num_heads, dropout=0, cross_attention=True, custom=True)\n ref_attn = torch.nn.MultiheadAttention(dim, num_heads, dropout=0, batch_first=True)\n\n # We can load the regular attention state dict\n # so we have compat when loading old checkpoints.\n cross_attn.load_state_dict(ref_attn.state_dict())\n\n queries = torch.randn(3, 7, dim)\n keys = torch.randn(3, 9, dim)\n values = torch.randn(3, 9, dim)\n\n y = cross_attn(queries, keys, values)[0]\n y_ref = ref_attn(queries, keys, values)[0]\n assert torch.allclose(y, y_ref, atol=1e-7), (y - y_ref).norm() / y_ref.norm()\n\n # Now let's check that streaming is working properly.\n with cross_attn.streaming():\n ys = []\n for step in range(queries.shape[1]):\n ys.append(cross_attn(queries[:, step: step + 1], keys, values)[0])\n y_streaming = torch.cat(ys, dim=1)\n assert torch.allclose(y_streaming, y, atol=1e-7)\n\n\ndef test_repeat_kv():\n torch.manual_seed(1234)\n num_heads = 8\n kv_repeat = 4\n dim = num_heads * 64\n with pytest.raises(AssertionError):\n mha = StreamingMultiheadAttention(\n dim, num_heads, causal=True, kv_repeat=kv_repeat, cross_attention=True)\n mha = StreamingMultiheadAttention(\n dim, num_heads, causal=True, kv_repeat=kv_repeat)\n mha = StreamingMultiheadAttention(\n dim, num_heads, causal=True, kv_repeat=kv_repeat, custom=True)\n x = torch.randn(4, 18, dim)\n y = mha(x, x, x)[0]\n assert x.shape == y.shape\n\n\ndef test_qk_layer_norm():\n torch.manual_seed(1234)\n tr = StreamingTransformer(\n 16, 4, 2, custom=True, dropout=0., qk_layer_norm=True, bias_attn=False)\n steps = 12\n x = torch.randn(3, steps, 16)\n y = tr(x)\n\n tr = StreamingTransformer(\n 16, 4, 2, custom=True, dropout=0., qk_layer_norm=True, cross_attention=True)\n z = torch.randn(3, 21, 16)\n y = tr(x, cross_attention_src=z)\n assert y.shape == x.shape\n" }, { "path": "tests/quantization/test_vq.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\n\nfrom audiocraft.quantization.vq import ResidualVectorQuantizer\n\n\nclass TestResidualVectorQuantizer:\n\n def test_rvq(self):\n x = torch.randn(1, 16, 2048, requires_grad=True)\n vq = ResidualVectorQuantizer(n_q=8, dimension=16, bins=8)\n res = vq(x, 1.)\n assert res.x.shape == torch.Size([1, 16, 2048])\n res.x.sum().backward()\n assert torch.allclose(x.grad.data, torch.ones(1))\n" }, { "path": "tests/utils/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "tests/utils/test_audio_effects.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport pytest\nfrom omegaconf import OmegaConf\n\nfrom audiocraft.utils.audio_effects import AudioEffects, get_audio_effects, select_audio_effects\n\nfrom ..common_utils import get_batch_white_noise\n\n\nclass TestAudioEffect:\n SR = 16_000\n\n @pytest.fixture(autouse=True)\n def audio_effects(self):\n cfg = {\n \"audio_effects\": {\n \"speed\": {\n \"sample_rate\": self.SR,\n \"speed_range\": [0.8, 1.2]\n },\n \"updownresample\": {\n \"sample_rate\": self.SR,\n \"intermediate_freq\": 32_000,\n },\n \"echo\": {\n \"sample_rate\": self.SR,\n \"volume_range\": [0.1, 0.5],\n },\n \"random_noise\": {\n \"noise_std\": 0.001,\n },\n \"pink_noise\": {\n \"noise_std\": 0.01,\n },\n \"lowpass_filter\": {\n \"sample_rate\": self.SR,\n \"cutoff_freq\": 5_000,\n },\n \"highpass_filter\": {\n \"sample_rate\": self.SR,\n \"cutoff_freq\": 500,\n },\n \"bandpass_filter\": {\n \"sample_rate\": self.SR,\n \"cutoff_freq_low\": 300,\n \"cutoff_freq_high\": 8_000,\n },\n \"smooth\": {\n \"window_size_range\": [2, 10],\n },\n \"boost_audio\": {\n \"amount\": 20,\n },\n \"duck_audio\": {\n \"amount\": 20,\n },\n \"mp3_compression\": {\n \"sample_rate\": self.SR,\n \"bitrate\": \"128k\",\n },\n \"aac_compression\": {\n \"sample_rate\": self.SR,\n \"bitrate\": \"128k\",\n \"lowpass_freq\": None,\n }\n }\n }\n weights = {\n \"speed\": 2.0,\n \"updownresample\": 0.4,\n \"echo\": 1.0,\n \"random_noise\": 3.0,\n \"pink_noise\": 0.5,\n \"lowpass_filter\": 4.0,\n \"highpass_filter\": 5.0,\n \"bandpass_filter\": 6.0,\n \"smooth\": 1.0,\n }\n return get_audio_effects(OmegaConf.structured(cfg)), weights\n\n def test_select_empty_effects(self):\n effects = select_audio_effects({})\n assert \"identity\" in effects and effects[\"identity\"] == AudioEffects.identity\n\n def test_select_wrong_strategy(self):\n with pytest.raises(ValueError):\n _ = select_audio_effects(\n audio_effects={},\n mode=\"some invalid mode\"\n )\n\n def test_selection(self, audio_effects):\n effect_cfg, weights = audio_effects\n effects = select_audio_effects(\n audio_effects=effect_cfg,\n weights=weights,\n mode=\"weighted\"\n )\n b, c, t = 2, 4, 32000\n audio = get_batch_white_noise(b, c, t)\n for effect_name, effect_func in effects.items():\n modified_audio = effect_func(audio)\n # It is quite hard to unit test the content of the modified_audio though\n if effect_name == \"speed\": # Speeding up audio should return in more frames\n assert modified_audio.size()[-1] > audio.size()[-1]\n else:\n assert modified_audio.size() == audio.size(), f\"Wrong dimension in {effect_name}\"\n" } ]