diff --git a/.gitignore b/.gitignore index 6304d33..40485be 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +<<<<<<< HEAD # Byte-compiled / optimized / DLL files __pycache__/ *.py[codz] @@ -11,6 +12,30 @@ api/outputs/ # C extensions *.so +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# Ignore model checkpoints and generated outputs +pretrained_models/ +outputs/ + +# C extensions +*.so + # Distribution / packaging .Python build/ @@ -32,8 +57,6 @@ share/python-wheels/ MANIFEST # PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec @@ -86,58 +109,28 @@ target/ # IPython profile_default/ ipython_config.py -.claude/ -CLAUDE.md + # pyenv -# For a library or package, you might want to ignore these files since the code is -# intended to run in multiple environments; otherwise, check them in: # .python-version # pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. #Pipfile.lock -# UV -# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -#uv.lock - -# poetry -# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +# poetry / pdm / pixi #poetry.lock -#poetry.toml - -# pdm -# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. -# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. -# https://pdm-project.org/en/latest/usage/project/#working-with-version-control #pdm.lock -#pdm.toml .pdm-python .pdm-build/ - -# pixi -# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. -#pixi.lock -# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one -# in the .venv directory. It is recommended not to include this directory in version control. .pixi -# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +# PEP 582 __pypackages__/ -# Celery stuff +# Celery celerybeat-schedule celerybeat.pid -# SageMath parsed files +# SageMath *.sage.py # Environments @@ -150,60 +143,29 @@ ENV/ env.bak/ venv.bak/ -# Spyder project settings +# IDE / Editor settings .spyderproject .spyproject - -# Rope project settings .ropeproject +.idea/ +.vscode/ -# mkdocs documentation +# mkdocs site /site -# mypy +# mypy / pyre / pytype caches .mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker .pyre/ - -# pytype static type analyzer .pytype/ # Cython debug symbols cython_debug/ -# PyCharm -# JetBrains specific template is maintained in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ - -# Abstra -# Abstra is an AI-powered process automation framework. -# Ignore directories containing user credentials, local state, and settings. -# Learn more at https://abstra.io/docs -.abstra/ - -# Visual Studio Code -# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore -# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore -# and can be added to the global gitignore or merged into this file. However, if you prefer, -# you could uncomment the following to ignore the entire vscode folder -# .vscode/ - -# Ruff stuff: +# Misc +.DS_Store +Thumbs.db .ruff_cache/ - -# PyPI configuration file .pypirc - -# Cursor -# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to -# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data -# refer to https://docs.cursor.com/context/ignore-files .cursorignore .cursorindexingignore @@ -211,3 +173,36 @@ cython_debug/ marimo/_static/ marimo/_lsp/ __marimo__/ + +# Models & Checkpoints (explicitly again for clarity) +*.safetensors +*.bin +*.ckpt +*.onnx + +# Output / Generated Audio +*.wav +*.mp3 +*.ogg + +# Cache & Temporary Files +cache/ +tmp/ +*.tmp +*.bak +*.npy +*.npz + +# Hugging Face / Torch / Diffusers Cache +~/.cache/ +huggingface/ +hf_cache/ +transformers_cache/ +torch_cache/ +diffusers_cache/ + +# Environment / Secrets +.env +.env.* +secrets.json +credentials.json diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..83ff4d6 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,19 @@ +# Changelog + +All notable changes to this repository should be documented in this file. + +## [Unreleased] - feature/gradio-ui (2025-10-30) + + - Accepts either a structured podcast JSON (recommended) or a manual utterance. + - Supports optional reference audio/text and dialect prompts. + - Improved, more descriptive labels and user guidance in the UI. + +Notes for upstream maintainers: + +### Added (2025-10-31) + +- Document non-verbal/paralinguistic token support (e.g., <|laughter|>, <|breathing|>, <|coughing|>) in the README and examples. These tokens can be inserted into input text to control paralinguistic events during generation. + - Added Gradio UI Single / Dual Speaker tabs. Dual Speaker supports uploading two reference audios/texts and a dialogue textarea where lines can be prefixed with `S1:` and `S2:` to indicate speaker turns. + - Added `example/gradio/README.md` with usage examples and guidance for formatting dialogue and reference audio usage. + - Added `example/gradio/smoke_test.py` — lightweight smoke test to validate UI input handling and optionally run a full inference when `--run` is supplied. + - Updated README to document Gradio UI usage and bracketed `[S1]/[S2]` dialog tag format; linked to `example/gradio/README.md` and smoke test. diff --git a/LICENSE b/LICENSE index 23101ab..3277067 100644 --- a/LICENSE +++ b/LICENSE @@ -1,201 +1,201 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright Soul AI Lab 2025. All Rights Reserved. - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright Soul AI Lab 2025. All Rights Reserved. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/cli/podcast.py b/cli/podcast.py index d61dfd5..11b5477 100644 --- a/cli/podcast.py +++ b/cli/podcast.py @@ -1,70 +1,70 @@ -import os -import json -import torch -import argparse - -import s3tokenizer -import soundfile as sf - -from soulxpodcast.config import SamplingParams -from soulxpodcast.utils.parser import podcast_format_parser -from soulxpodcast.utils.infer_utils import initiate_model, process_single_input - - -def run_inference( - inputs: dict, - model_path: str, - output_path: str, - llm_engine: str = "hf", - fp16_flow: bool = False, - seed: int = 1988, -): - - model, dataset = initiate_model(seed, model_path, llm_engine, fp16_flow) - - data = process_single_input( - dataset, - inputs['text'], - inputs['prompt_wav'], - inputs['prompt_text'], - inputs['use_dialect_prompt'], - inputs['dialect_prompt_text'], - ) - - print("[INFO] Start inference...") - results_dict = model.forward_longform(**data) - - target_audio = None - for wav in results_dict["generated_wavs"]: - if target_audio is None: - target_audio = wav - else: - target_audio = torch.cat([target_audio, wav], dim=1) - - os.makedirs(os.path.dirname(output_path), exist_ok=True) - sf.write(output_path, target_audio.cpu().squeeze(0).numpy(), 24000) - print(f"[INFO] Saved synthesized audio to: {output_path}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser = argparse.ArgumentParser() - parser.add_argument("--json_path", required=True, help="Path to the input JSON file") - parser.add_argument("--model_path", required=True, help="Path to the model file") - parser.add_argument("--output_path", default="outputs/result.wav", help="Path to the output audio file") - parser.add_argument("--llm_engine", default="hf", choices=["hf", "vllm"], help="Inference engine to use") - parser.add_argument("--fp16_flow", action="store_true", help="Enable FP16 flow") - parser.add_argument("--seed", type=int, default=1988, help="Random seed") - args = parser.parse_args() - - with open(args.json_path, "r") as f: - data = json.load(f) - inputs = podcast_format_parser(data) - run_inference( - inputs=inputs, - model_path=args.model_path, - output_path=args.output_path, - llm_engine=args.llm_engine, - fp16_flow=args.fp16_flow, - seed=args.seed, - ) +import os +import json +import torch +import argparse + +import s3tokenizer +import soundfile as sf + +from soulxpodcast.config import SamplingParams +from soulxpodcast.utils.parser import podcast_format_parser +from soulxpodcast.utils.infer_utils import initiate_model, process_single_input + + +def run_inference( + inputs: dict, + model_path: str, + output_path: str, + llm_engine: str = "hf", + fp16_flow: bool = False, + seed: int = 1988, +): + + model, dataset = initiate_model(seed, model_path, llm_engine, fp16_flow) + + data = process_single_input( + dataset, + inputs['text'], + inputs['prompt_wav'], + inputs['prompt_text'], + inputs['use_dialect_prompt'], + inputs['dialect_prompt_text'], + ) + + print("[INFO] Start inference...") + results_dict = model.forward_longform(**data) + + target_audio = None + for wav in results_dict["generated_wavs"]: + if target_audio is None: + target_audio = wav + else: + target_audio = torch.cat([target_audio, wav], dim=1) + + os.makedirs(os.path.dirname(output_path), exist_ok=True) + sf.write(output_path, target_audio.cpu().squeeze(0).numpy(), 24000) + print(f"[INFO] Saved synthesized audio to: {output_path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser = argparse.ArgumentParser() + parser.add_argument("--json_path", required=True, help="Path to the input JSON file") + parser.add_argument("--model_path", required=True, help="Path to the model file") + parser.add_argument("--output_path", default="outputs/result.wav", help="Path to the output audio file") + parser.add_argument("--llm_engine", default="hf", choices=["hf", "vllm"], help="Inference engine to use") + parser.add_argument("--fp16_flow", action="store_true", help="Enable FP16 flow") + parser.add_argument("--seed", type=int, default=1988, help="Random seed") + args = parser.parse_args() + + with open(args.json_path, "r") as f: + data = json.load(f) + inputs = podcast_format_parser(data) + run_inference( + inputs=inputs, + model_path=args.model_path, + output_path=args.output_path, + llm_engine=args.llm_engine, + fp16_flow=args.fp16_flow, + seed=args.seed, + ) diff --git a/example/gradio/README.md b/example/gradio/README.md new file mode 100644 index 0000000..5369bc0 --- /dev/null +++ b/example/gradio/README.md @@ -0,0 +1,119 @@ +Gradio UI for SoulX-Podcast +========================== + +Quick instructions to run the Gradio UI locally (recommended to run inside the project's virtualenv in WSL/Ubuntu where dependencies are installed): + +1. Ensure dependencies are installed (project `requirements.txt` already includes `gradio`). + + Gradio UI for SoulX-Podcast + ========================== + + This folder contains a small demo of the Gradio UI and quick usage notes. + + Formatting the dialogue textarea + + - For Dual Speaker mode you can enter multi-line dialogue where each line is either prefixed with the speaker label (S1: / S2:) or uses the bracketed tag format `[S1]` / `[S2]`. Examples: + + S1: Hello, welcome to the show. + S2: Thanks — happy to be here. + + or + + [S1] Hello, welcome to the show. + [S2] Thanks — happy to be here. + + - The UI will convert either form into the canonical internal speaker-tagged form the pipeline expects (e.g. `[S1]Hello, welcome to the show.`). + + Speaker mapping + + - S1 -> Speaker 1 (first reference audio / text) + - S2 -> Speaker 2 (second reference audio / text) + + Uploading two reference audios and texts + + - Dual Speaker tab allows you to upload two reference audio files (one per speaker) and optional reference text for each speaker. + - If you only upload one reference audio, the other speaker will use a silent fallback. To get best results, provide short (1-10s) clean reference audio for each speaker. + + Expected behavior + + - Single Speaker tab: + - Use `Manual utterance` (or upload JSON) and optionally supply one reference audio / text. + - The app will synthesize a single utterance and save the output to `outputs/output_single.wav` by default. + + - Dual Speaker tab: + - Enter a multi-line dialogue using `S1:` and `S2:` prefixes or upload a small JSON. + - Provide 0, 1, or 2 reference audios and optional texts. The app will attempt to use speaker-specific references in order. + - Output is saved to `outputs/output_dual.wav` by default. + + Smoke test (quick, non-heavy check) + + A lightweight smoke-test script is provided to validate UI input handling and fallback audio creation without running the full model by default. To run the light smoke test (does not load model): + + ```powershell + python example/gradio/smoke_test.py + ``` + + To run the full inference (this can be heavy and will load the selected model), add `--run` and optionally set `--model-path`: + + ```powershell + python example/gradio/smoke_test.py --run --model-path pretrained_models/SoulX-Podcast-1.7B + ``` + + Notes + + - Dual-speaker synthesis will be more reliable when both speakers have clean reference audio. + - The UI keeps single-speaker behavior unchanged; the Dual Speaker tab is optional and additive. + + Quick start + ----------- + + 1. Install dependencies in your venv (WSL recommended): + + ```powershell + pip install -r requirements.txt + ``` + + 2. Run the Gradio app from the project root: + + ```bash + python gradio_app.py + ``` + + 3. Open http://localhost:7860 in your browser. + + Example JSON + ------------ + + An example JSON for the Dual Speaker tab (mixed setup: Speaker 1 uses a preloaded sample, Speaker 2 is an uploaded file path) is provided at `example/gradio/dialog_example_mixed.json`. Replace the `prompt_audio` path for `S2` with your uploaded file path or upload via the UI. + + ```json + { + "speakers": { + "S1": { "prompt_text": "A calm female host", "prompt_audio": "example/audios/female_mandarin.wav" }, + "S2": { "prompt_text": "A guest voice", "prompt_audio": "/path/to/your/uploaded_speaker2.wav" } + }, + "text": [["S1","Hello..."],["S2","Hi..."]] + } + ``` + + Referencing repo samples from JSON + ---------------------------------- + + You can reference the preloaded audio samples (those listed in `example/audios`) using short tokens in your JSON. The Gradio UI will resolve the tokens to the currently selected UI samples when you run the job. + + - Use `"sample1"` to refer to Speaker 1's selected sample. + - Use `"sample2"` to refer to Speaker 2's selected sample. + + Example JSON using token references (upload this file in the Dual Speaker tab): + + ```json + { + "speakers": { + "S1": { "prompt_text": "Host voice (uses sample1 token)", "prompt_audio": "sample1" }, + "S2": { "prompt_text": "Guest voice (uses sample2 token)", "prompt_audio": "sample2" } + }, + "text": [["S1","Welcome to the show."],["S2","Thanks for having me."]] + } + ``` + + There is also a ready example file you can upload: `dialog_example_sample_tokens.json` which uses `sample1` and `sample2` tokens. Make sure to pick the desired samples in the Dual Speaker UI dropdowns before running the job so the tokens resolve to the audio files you want. diff --git a/example/gradio/dialog_example_mixed.json b/example/gradio/dialog_example_mixed.json new file mode 100644 index 0000000..e120d6d --- /dev/null +++ b/example/gradio/dialog_example_mixed.json @@ -0,0 +1,19 @@ +{ + "speakers": { + "S1": { + "prompt_text": "A calm female host in Mandarin", + "prompt_audio": "example/audios/female_mandarin.wav", + "dialect_prompt": "" + }, + "S2": { + "prompt_text": "A guest voice, energetic", + "prompt_audio": "/path/to/your/uploaded_speaker2.wav", + "dialect_prompt": "" + } + }, + "text": [ + ["S1", "Hello everyone, welcome to the show. <|laughter|>"], + ["S2", "Thanks for having me — I'm excited to be here. <|breathing|>"], + ["S1", "Let's get started with our first topic."] + ] +} diff --git a/example/gradio/dialog_example_sample_tokens.json b/example/gradio/dialog_example_sample_tokens.json new file mode 100644 index 0000000..5020401 --- /dev/null +++ b/example/gradio/dialog_example_sample_tokens.json @@ -0,0 +1,7 @@ +{ + "speakers": { + "S1": { "prompt_text": "Host voice (uses sample1 token)", "prompt_audio": "sample1" }, + "S2": { "prompt_text": "Guest voice (uses sample2 token)", "prompt_audio": "sample2" } + }, + "text": [["S1","Welcome to the show."],["S2","Thanks for having me."]] +} diff --git a/example/gradio/smoke_test.py b/example/gradio/smoke_test.py new file mode 100644 index 0000000..3d9e156 --- /dev/null +++ b/example/gradio/smoke_test.py @@ -0,0 +1,91 @@ +"""Lightweight smoke test for the Gradio UI input handling. + +This script performs a non-heavy check by creating a silent fallback reference audio +and validating that the example output path can be written. By default it DOES NOT +load or run the model. Use --run to perform an actual inference run (this may be +heavy and will load the selected model). + +Usage: + python example/gradio/smoke_test.py + python example/gradio/smoke_test.py --run --model-path pretrained_models/SoulX-Podcast-1.7B + +""" +import argparse +import os +import tempfile +import torch +import torchaudio +import sys + +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +OUT_DIR = os.path.join(ROOT, "outputs") +os.makedirs(OUT_DIR, exist_ok=True) + + +def make_silent_wav(path: str, sr: int = 16000, duration_s: float = 0.5): + samples = int(sr * duration_s) + tensor = torch.zeros((1, samples), dtype=torch.float32) + torchaudio.save(path, tensor, sr) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--run", action="store_true", help="If set, actually call the inference function (may be heavy)") + p.add_argument("--model-path", type=str, default="", help="Model path to use when --run is set") + p.add_argument("--output", type=str, default=os.path.join(OUT_DIR, "smoke_test_output.wav")) + args = p.parse_args() + + print("Smoke test: creating a silent fallback reference audio and verifying output write") + tmp_ref = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") + try: + make_silent_wav(tmp_ref.name) + print(f"Created silent reference: {tmp_ref.name}") + except Exception as e: + print("Failed to create silent wav:", e) + sys.exit(2) + + # Verify writing an output using the same format + try: + make_silent_wav(args.output) + print(f"Wrote smoke output to: {args.output}") + except Exception as e: + print("Failed to write smoke output:", e) + sys.exit(2) + + if args.run: + # Attempt to run full inference via the gradio_app helper. This will load the model + # and can be slow/heavy. We import lazily so the default smoke test remains fast. + try: + sys.path.insert(0, ROOT) + import gradio_app as ga + + print("Running full inference (this may be slow). Model path:", args.model_path or "(default)") + status, out, model_path = ga.infer_from_ui( + None, # json_file + "[S1]Hello from smoke test", # manual_text + tmp_ref.name, # prompt_wav_file + None, # sample_choice + False, # use_sample + "", # prompt_text + None, # speaker2 prompt wav + "", # speaker2 prompt text + "", # dialogue_text + False, # use_dialect_prompt + "", # dialect_prompt_text + "", # model_choice + args.model_path or "", # model_path + args.output, + "hf", + True, + 42, + ) + print("Inference result:", status, out, model_path) + except Exception as e: + print("Full inference failed:", e) + sys.exit(3) + + print("Smoke test completed successfully") + + +if __name__ == "__main__": + main() diff --git a/example/infer_dialogue.sh b/example/infer_dialogue.sh index e67371a..3af0249 100644 --- a/example/infer_dialogue.sh +++ b/example/infer_dialogue.sh @@ -1,12 +1,12 @@ -export PYTHONPATH="$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. && pwd)" -echo "PYTHONPATH set to: $PYTHONPATH" - -# Note: To infer Chinese dialects, set model_dir to "pretrained_models/SoulX-Podcast-1.7B-dialect -model_dir=pretrained_models/SoulX-Podcast-1.7B -input_file=example/podcast_script/script_mandarin.json - -python cli/podcast.py \ - --json_path ${input_file} \ - --model_path ${model_dir} \ - --output_path outputs/mandarin.wav \ +export PYTHONPATH="$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. && pwd)" +echo "PYTHONPATH set to: $PYTHONPATH" + +# Note: To infer Chinese dialects, set model_dir to "pretrained_models/SoulX-Podcast-1.7B-dialect +model_dir=pretrained_models/SoulX-Podcast-1.7B +input_file=example/podcast_script/script_mandarin.json + +python cli/podcast.py \ + --json_path ${input_file} \ + --model_path ${model_dir} \ + --output_path outputs/mandarin.wav \ --seed 7 \ No newline at end of file diff --git a/example/podcast_script/script_english.json b/example/podcast_script/script_english.json index 5c48f9b..f2f9271 100644 --- a/example/podcast_script/script_english.json +++ b/example/podcast_script/script_english.json @@ -1,18 +1,18 @@ -{ - "speakers": { - "S1": { - "prompt_audio": "example/audios/female_mandarin.wav", - "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。" - }, - "S2": { - "prompt_audio": "example/audios/male_mandarin.wav", - "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？" - } - }, - "text": [ - ["S1", "Hey, sorry for snapping earlier. That was totally unprofessional of me. Let’s get back to recapping the match."], - ["S2", "Yeah, no worries. Honestly, that was such an awesome match. You’ve gotta admit, Ali, even though your substitute didn’t win, it was almost legendary."], - ["S1", "Totally, one for the record books. Sam and Max were only supposed to have four bouts, but things got a little messy right from the start."], - ["S2", "Oh yeah, someone ran onto the field shouting something—I couldn’t hear a word. Honestly, I was just glad when the guards finally dragged them out. I just wanted to get back to the joust, but it really threw everything off, and it even got the crowd all hyped up, which was kinda surprising."] - ] +{ + "speakers": { + "S1": { + "prompt_audio": "example/audios/female_mandarin.wav", + "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。" + }, + "S2": { + "prompt_audio": "example/audios/male_mandarin.wav", + "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？" + } + }, + "text": [ + ["S1", "Hey, sorry for snapping earlier. That was totally unprofessional of me. Let’s get back to recapping the match."], + ["S2", "Yeah, no worries. Honestly, that was such an awesome match. You’ve gotta admit, Ali, even though your substitute didn’t win, it was almost legendary."], + ["S1", "Totally, one for the record books. Sam and Max were only supposed to have four bouts, but things got a little messy right from the start."], + ["S2", "Oh yeah, someone ran onto the field shouting something—I couldn’t hear a word. Honestly, I was just glad when the guards finally dragged them out. I just wanted to get back to the joust, but it really threw everything off, and it even got the crowd all hyped up, which was kinda surprising."] + ] } \ No newline at end of file diff --git a/example/podcast_script/script_henan.json b/example/podcast_script/script_henan.json index afb181d..6f467a7 100644 --- a/example/podcast_script/script_henan.json +++ b/example/podcast_script/script_henan.json @@ -1,21 +1,21 @@ -{ - "speakers": { - "S1": { - "prompt_audio": "example/audios/female_mandarin.wav", - "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。", - "dialect_prompt": "<|Henan|>俺这不是怕恁路上不得劲儿嘛！那景德镇瓷泥可娇贵着哩，得先拿咱河南人这实诚劲儿给它揉透喽。" - }, - "S2": { - "prompt_audio": "example/audios/male_mandarin.wav", - "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？", - "dialect_prompt": "<|Henan|>恁这想法真闹挺！陕北民谣比黑神话早几百年都有了，咱可不兴这弄颠倒啊，中不？恁这想法真闹挺！那陕北民谣在黄土高坡响了几百年，咋能说是跟黑神话学的咧？咱得把这事儿捋直喽，中不中！" - } - }, - "text": [ - ["S1", "<|Henan|>哎，大家好啊，欢迎收听咱这一期嘞《瞎聊呗，就这么说》，我是恁嘞老朋友，燕子。"], - ["S2", "<|Henan|>大家好，我是老张。燕子啊，今儿瞅瞅你这个劲儿，咋着，是有啥可得劲嘞事儿想跟咱唠唠？"], - ["S1", "<|Henan|>哎哟，老张，你咋恁懂我嘞！我跟你说啊，最近我刷手机，老是刷住些可逗嘞方言视频，特别是咱河南话，咦～我哩个乖乖，一听我都憋不住笑，咋说嘞，得劲儿哩很，跟回到家一样。"], - ["S2", "<|Henan|>你这回可算说到根儿上了！河南话，咱往大处说说，中原官话，它真嘞是有一股劲儿搁里头。它可不光是说话，它脊梁骨后头藏嘞，是咱一整套、鲜鲜活活嘞过法儿，一种活人嘞道理。"], - ["S1", "<|Henan|>活人嘞道理？哎，这你这一说，我嘞兴致“腾”一下就上来啦！觉住咱这嗑儿，一下儿从搞笑视频蹿到文化顶上了。那你赶紧给我白话白话，这里头到底有啥道道儿？我特别想知道——为啥一提起咱河南人，好些人脑子里“蹦”出来嘞头一个词儿，就是实在？这个实在，骨子里到底是啥嘞？ "] - ] +{ + "speakers": { + "S1": { + "prompt_audio": "example/audios/female_mandarin.wav", + "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。", + "dialect_prompt": "<|Henan|>俺这不是怕恁路上不得劲儿嘛！那景德镇瓷泥可娇贵着哩，得先拿咱河南人这实诚劲儿给它揉透喽。" + }, + "S2": { + "prompt_audio": "example/audios/male_mandarin.wav", + "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？", + "dialect_prompt": "<|Henan|>恁这想法真闹挺！陕北民谣比黑神话早几百年都有了，咱可不兴这弄颠倒啊，中不？恁这想法真闹挺！那陕北民谣在黄土高坡响了几百年，咋能说是跟黑神话学的咧？咱得把这事儿捋直喽，中不中！" + } + }, + "text": [ + ["S1", "<|Henan|>哎，大家好啊，欢迎收听咱这一期嘞《瞎聊呗，就这么说》，我是恁嘞老朋友，燕子。"], + ["S2", "<|Henan|>大家好，我是老张。燕子啊，今儿瞅瞅你这个劲儿，咋着，是有啥可得劲嘞事儿想跟咱唠唠？"], + ["S1", "<|Henan|>哎哟，老张，你咋恁懂我嘞！我跟你说啊，最近我刷手机，老是刷住些可逗嘞方言视频，特别是咱河南话，咦～我哩个乖乖，一听我都憋不住笑，咋说嘞，得劲儿哩很，跟回到家一样。"], + ["S2", "<|Henan|>你这回可算说到根儿上了！河南话，咱往大处说说，中原官话，它真嘞是有一股劲儿搁里头。它可不光是说话，它脊梁骨后头藏嘞，是咱一整套、鲜鲜活活嘞过法儿，一种活人嘞道理。"], + ["S1", "<|Henan|>活人嘞道理？哎，这你这一说，我嘞兴致“腾”一下就上来啦！觉住咱这嗑儿，一下儿从搞笑视频蹿到文化顶上了。那你赶紧给我白话白话，这里头到底有啥道道儿？我特别想知道——为啥一提起咱河南人，好些人脑子里“蹦”出来嘞头一个词儿，就是实在？这个实在，骨子里到底是啥嘞？ "] + ] } \ No newline at end of file diff --git a/example/podcast_script/script_mandarin.json b/example/podcast_script/script_mandarin.json index 837ac6c..432c3d5 100644 --- a/example/podcast_script/script_mandarin.json +++ b/example/podcast_script/script_mandarin.json @@ -1,18 +1,18 @@ -{ - "speakers": { - "S1": { - "prompt_audio": "example/audios/female_mandarin.wav", - "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。" - }, - "S2": { - "prompt_audio": "example/audios/male_mandarin.wav", - "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？" - } - }, - "text": [ - ["S1", "哈喽，AI时代的冲浪先锋们！欢迎收听《AI生活进行时》。啊，一个充满了未来感，然后，还有一点点，<|laughter|>神经质的播客节目，我是主持人小希。"], - ["S2", "哎，大家好呀！我是能唠，爱唠，天天都想唠的唠嗑！"], - ["S1", "最近活得特别赛博朋克哈！以前老是觉得AI是科幻片儿里的，<|sigh|> 现在，现在连我妈都用AI写广场舞文案了。"], - ["S2", "这个例子很生动啊。是的，特别是生成式AI哈，感觉都要炸了！ 诶，那我们今天就聊聊AI是怎么走进我们的生活的哈！"] - ] +{ + "speakers": { + "S1": { + "prompt_audio": "example/audios/female_mandarin.wav", + "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。" + }, + "S2": { + "prompt_audio": "example/audios/male_mandarin.wav", + "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？" + } + }, + "text": [ + ["S1", "哈喽，AI时代的冲浪先锋们！欢迎收听《AI生活进行时》。啊，一个充满了未来感，然后，还有一点点，<|laughter|>神经质的播客节目，我是主持人小希。"], + ["S2", "哎，大家好呀！我是能唠，爱唠，天天都想唠的唠嗑！"], + ["S1", "最近活得特别赛博朋克哈！以前老是觉得AI是科幻片儿里的，<|sigh|> 现在，现在连我妈都用AI写广场舞文案了。"], + ["S2", "这个例子很生动啊。是的，特别是生成式AI哈，感觉都要炸了！ 诶，那我们今天就聊聊AI是怎么走进我们的生活的哈！"] + ] } \ No newline at end of file diff --git a/example/podcast_script/script_sichuan.json b/example/podcast_script/script_sichuan.json index c1e4bde..4516e04 100644 --- a/example/podcast_script/script_sichuan.json +++ b/example/podcast_script/script_sichuan.json @@ -1,20 +1,20 @@ -{ - "speakers": { - "S1": { - "prompt_audio": "example/audios/female_mandarin.wav", - "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。", - "dialect_prompt": "<|Sichuan|>要得要得！前头几个耍洋盘，我后脚就背起铺盖卷去景德镇耍泥巴，巴适得喊老天爷！" - }, - "S2": { - "prompt_audio": "example/audios/male_mandarin.wav", - "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？", - "dialect_prompt": "<|Sichuan|>哎哟喂，这个搞反了噻！黑神话里头唱曲子的王二浪早八百年就在黄土高坡吼秦腔喽，游戏组专门跑切录的原汤原水，听得人汗毛儿都立起来！" - } - }, - "text": [ - ["S1", "<|Sichuan|>各位《巴适得板》的听众些，大家好噻！我是你们主持人晶晶。今儿天气硬是巴适，不晓得大家是在赶路嘛，还是茶都泡起咯，准备跟我们好生摆一哈龙门阵喃？"], - ["S2", "<|Sichuan|>晶晶好哦，大家安逸噻！我是李老倌。你刚开口就川味十足，摆龙门阵几个字一甩出来，我鼻子头都闻到茶香跟火锅香咯！"], - ["S1", "<|Sichuan|>就是得嘛！李老倌，我前些天带个外地朋友切人民公园鹤鸣茶社坐了一哈。他硬是搞不醒豁，为啥子我们一堆人围到杯茶就可以吹一下午壳子，从隔壁子王嬢嬢娃儿耍朋友，扯到美国大选，中间还掺几盘斗地主。他说我们四川人简直是把摸鱼刻进骨子里头咯！"], - ["S2", "<|Sichuan|>你那个朋友说得倒是有点儿趣，但他莫看到精髓噻。摆龙门阵哪是摸鱼嘛，这是我们川渝人特有的交际方式，更是一种活法。外省人天天说的松弛感，根根儿就在这龙门阵里头。今天我们就要好生摆一哈，为啥子四川人活得这么舒坦。就先从茶馆这个老窝子说起，看它咋个成了我们四川人的魂儿！"] - ] +{ + "speakers": { + "S1": { + "prompt_audio": "example/audios/female_mandarin.wav", + "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。", + "dialect_prompt": "<|Sichuan|>要得要得！前头几个耍洋盘，我后脚就背起铺盖卷去景德镇耍泥巴，巴适得喊老天爷！" + }, + "S2": { + "prompt_audio": "example/audios/male_mandarin.wav", + "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？", + "dialect_prompt": "<|Sichuan|>哎哟喂，这个搞反了噻！黑神话里头唱曲子的王二浪早八百年就在黄土高坡吼秦腔喽，游戏组专门跑切录的原汤原水，听得人汗毛儿都立起来！" + } + }, + "text": [ + ["S1", "<|Sichuan|>各位《巴适得板》的听众些，大家好噻！我是你们主持人晶晶。今儿天气硬是巴适，不晓得大家是在赶路嘛，还是茶都泡起咯，准备跟我们好生摆一哈龙门阵喃？"], + ["S2", "<|Sichuan|>晶晶好哦，大家安逸噻！我是李老倌。你刚开口就川味十足，摆龙门阵几个字一甩出来，我鼻子头都闻到茶香跟火锅香咯！"], + ["S1", "<|Sichuan|>就是得嘛！李老倌，我前些天带个外地朋友切人民公园鹤鸣茶社坐了一哈。他硬是搞不醒豁，为啥子我们一堆人围到杯茶就可以吹一下午壳子，从隔壁子王嬢嬢娃儿耍朋友，扯到美国大选，中间还掺几盘斗地主。他说我们四川人简直是把摸鱼刻进骨子里头咯！"], + ["S2", "<|Sichuan|>你那个朋友说得倒是有点儿趣，但他莫看到精髓噻。摆龙门阵哪是摸鱼嘛，这是我们川渝人特有的交际方式，更是一种活法。外省人天天说的松弛感，根根儿就在这龙门阵里头。今天我们就要好生摆一哈，为啥子四川人活得这么舒坦。就先从茶馆这个老窝子说起，看它咋个成了我们四川人的魂儿！"] + ] } \ No newline at end of file diff --git a/example/podcast_script/script_yue.json b/example/podcast_script/script_yue.json index ee24f2a..51604e0 100644 --- a/example/podcast_script/script_yue.json +++ b/example/podcast_script/script_yue.json @@ -1,20 +1,20 @@ -{ - "speakers": { - "S1": { - "prompt_audio": "example/audios/female_mandarin.wav", - "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。", - "dialect_prompt": "<|Yue|>真係冇讲错啊！攀山滑雪嘅语言专家几巴闭，都唔及我听日拖成副身家去景德镇玩泥巴，呢铺真系发哂白日梦咯！" - }, - "S2": { - "prompt_audio": "example/audios/male_mandarin.wav", - "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？", - "dialect_prompt": "<|Yue|>咪搞错啊！陕北民谣响度唱咗几十年，黑神话边有咁大面啊？你估佢哋抄游戏咩！" - } - }, - "text": [ - ["S1", "<|Yue|>哈囉大家好啊，歡迎收聽我哋嘅節目。喂，我今日想問你樣嘢啊，你覺唔覺得，嗯，而家揸電動車，最煩，最煩嘅一樣嘢係咩啊？"], - ["S2", "<|Yue|>梗係充電啦。大佬啊，搵個位都已經好煩，搵到個位仲要喺度等，你話快極都要半個鐘一個鐘，真係，有時諗起都覺得好冇癮。"], - ["S1", "<|Yue|>係咪先。如果我而家同你講，充電可以快到同入油差唔多時間，你信唔信先？喂你平時喺油站入滿一缸油，要幾耐啊？五六分鐘？"], - ["S2", "<|Yue|>差唔多啦，七八分鐘，點都走得啦。電車喎，可以做到咁快？你咪玩啦。"] - ] +{ + "speakers": { + "S1": { + "prompt_audio": "example/audios/female_mandarin.wav", + "prompt_text": "喜欢攀岩、徒步、滑雪的语言爱好者，以及过两天要带着全部家当去景德镇做陶瓷的白日梦想家。", + "dialect_prompt": "<|Yue|>真係冇讲错啊！攀山滑雪嘅语言专家几巴闭，都唔及我听日拖成副身家去景德镇玩泥巴，呢铺真系发哂白日梦咯！" + }, + "S2": { + "prompt_audio": "example/audios/male_mandarin.wav", + "prompt_text": "呃，还有一个就是要跟大家纠正一点，就是我们在看电影的时候，尤其是游戏玩家，看电影的时候，在看到那个到西北那边的这个陕北民谣，嗯，这个可能在想，哎，是不是他是受到了黑神话的启发？", + "dialect_prompt": "<|Yue|>咪搞错啊！陕北民谣响度唱咗几十年，黑神话边有咁大面啊？你估佢哋抄游戏咩！" + } + }, + "text": [ + ["S1", "<|Yue|>哈囉大家好啊，歡迎收聽我哋嘅節目。喂，我今日想問你樣嘢啊，你覺唔覺得，嗯，而家揸電動車，最煩，最煩嘅一樣嘢係咩啊？"], + ["S2", "<|Yue|>梗係充電啦。大佬啊，搵個位都已經好煩，搵到個位仲要喺度等，你話快極都要半個鐘一個鐘，真係，有時諗起都覺得好冇癮。"], + ["S1", "<|Yue|>係咪先。如果我而家同你講，充電可以快到同入油差唔多時間，你信唔信先？喂你平時喺油站入滿一缸油，要幾耐啊？五六分鐘？"], + ["S2", "<|Yue|>差唔多啦，七八分鐘，點都走得啦。電車喎，可以做到咁快？你咪玩啦。"] + ] } \ No newline at end of file diff --git a/gradio_app.py b/gradio_app.py new file mode 100644 index 0000000..2334593 --- /dev/null +++ b/gradio_app.py @@ -0,0 +1,634 @@ +import os +import json +import tempfile +from typing import Optional + +import gradio as gr +import torch +import soundfile as sf +import numpy as np +import torchaudio +import importlib +import re + +# Use lower-level utilities so we can cache model and dataset in memory between runs +from soulxpodcast.utils.parser import podcast_format_parser +from soulxpodcast.utils.infer_utils import initiate_model, process_single_input + +# Simple in-memory cache for model and dataset +MODEL_CACHE = { + "model_path": None, + "llm_engine": None, + "fp16_flow": None, + "seed": None, + "model": None, + "dataset": None, +} + + +def load_model_if_needed(model_path, llm_engine, fp16_flow, seed): + """Load model and dataset into MODEL_CACHE if not loaded or if parameters changed. + + Returns (model, dataset, loaded_new) + """ + # Normalize inputs to avoid reloads caused by equivalent but non-identical representations + try: + # Normalize path robustly: abspath -> realpath -> normpath -> normcase + norm_model_path = os.path.normcase(os.path.normpath(os.path.realpath(os.path.abspath(model_path)))) if model_path else model_path + except Exception: + norm_model_path = model_path + norm_llm_engine = str(llm_engine) if llm_engine is not None else llm_engine + norm_fp16_flow = bool(fp16_flow) + try: + norm_seed = int(seed) + except Exception: + norm_seed = seed + + loaded_new = False + if ( + MODEL_CACHE["model_path"] != norm_model_path + or MODEL_CACHE["llm_engine"] != norm_llm_engine + or MODEL_CACHE["fp16_flow"] != norm_fp16_flow + or MODEL_CACHE["seed"] != norm_seed + or MODEL_CACHE["model"] is None + or MODEL_CACHE["dataset"] is None + ): + # free previous model references if present to avoid duplicate GPU allocations + try: + if MODEL_CACHE.get("model") is not None: + # attempt to delete model and dataset and free CUDA memory before loading new one + try: + del MODEL_CACHE["model"] + except Exception: + pass + try: + del MODEL_CACHE["dataset"] + except Exception: + pass + import gc + + gc.collect() + if torch.cuda.is_available(): + try: + torch.cuda.empty_cache() + except Exception: + pass + except Exception: + pass + + MODEL_CACHE["model_path"] = norm_model_path + MODEL_CACHE["llm_engine"] = norm_llm_engine + MODEL_CACHE["fp16_flow"] = norm_fp16_flow + MODEL_CACHE["seed"] = norm_seed + # log loading action + try: + from datetime import datetime + + print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Loading model from {norm_model_path} (engine={norm_llm_engine}, fp16_flow={norm_fp16_flow})") + except Exception: + pass + + model, dataset = initiate_model(norm_seed, norm_model_path, norm_llm_engine, norm_fp16_flow) + MODEL_CACHE["model"] = model + MODEL_CACHE["dataset"] = dataset + loaded_new = True + return MODEL_CACHE["model"], MODEL_CACHE["dataset"], loaded_new + + +def save_uploaded_file(uploaded) -> Optional[str]: + if uploaded is None: + return None + # Handle common Gradio return types: filepath (str), dict with 'name'/'tmp_path', or file-like + # If it's already a filepath, return it + try: + if isinstance(uploaded, str): + return uploaded + except Exception: + pass + + # If Gradio returned a dict with paths + try: + if isinstance(uploaded, dict): + return uploaded.get("name") or uploaded.get("tmp_path") or uploaded.get("file") + except Exception: + pass + + # If it's a file-like object, try to write it to a temp file + try: + filename = getattr(uploaded, "name", None) or "upload" + _, ext = os.path.splitext(filename) + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext) + # Some file-likes have a .read() method + if hasattr(uploaded, "read"): + data = uploaded.read() + # If data is str, encode + if isinstance(data, str): + data = data.encode("utf-8") + tmp.write(data) + tmp.flush() + tmp.close() + return tmp.name + # If uploaded is raw bytes + if isinstance(uploaded, (bytes, bytearray)): + tmp.write(uploaded) + tmp.flush() + tmp.close() + return tmp.name + # If we couldn't handle the uploaded file, cleanup and return None + try: + tmp.close() + except Exception: + pass + return None + except Exception: + # Any error while saving uploaded file -> treat as no file + try: + tmp.close() + except Exception: + pass + return None + +def infer_from_ui( + json_file, + manual_text, + prompt_wav_file, + sample_choice_s1=None, + use_sample_s1=False, + prompt_text="", + + speaker2_prompt_wav_file=None, + speaker2_prompt_text="", + dialogue_text="", + use_dialect_prompt=False, + dialect_prompt_text="", + model_choice=None, + model_path=None, + output_path="outputs/output.wav", + llm_engine="hf", + fp16_flow=True, + seed=42, + # optional dual-speaker params + sample_choice_s2=None, + use_sample_s2=False, + auto_fill_s2=False, +): + try: + # gradio may return a filepath (str), a dict, or a file-like + if isinstance(json_file, str): + with open(json_file, "r", encoding="utf-8") as f: + data = json.load(f) + elif isinstance(json_file, dict): + # dict may contain 'name' or 'tmp_path' + path = json_file.get("name") or json_file.get("tmp_path") or json_file.get("file") + if path: + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + else: + # try to read file-like entries + data = json.load(json_file) + else: + # file-like object + data = json.load(json_file) + inputs = podcast_format_parser(data) + # If the JSON referred to sample identifiers (e.g. a repo filename or 'sample1'), + # resolve those to actual files under example/audios when available so the UI + # behaves as if the user selected the sample in the dropdown. + try: + sample_dir = os.path.join(os.path.dirname(__file__), "example", "audios") + sample_files = [f for f in os.listdir(sample_dir) if f.lower().endswith('.wav')] + except Exception: + sample_files = [] + + def _resolve_sample_entry(val): + try: + if not val: + return val + # already an absolute path and exists + if os.path.isabs(val) and os.path.exists(val): + return val + # maybe it's a relative path from cwd + cand = os.path.join(os.getcwd(), val) + if os.path.exists(cand): + return cand + # direct match to a repo filename + if val in sample_files: + return os.path.join(sample_dir, val) + # normalized forms: 'sample1', 's1', '1' -> map by index + s = str(val).lower().replace(" ", "").replace("_", "") + m = re.match(r"^s?(\d+)$", s) + if m: + idx = int(m.group(1)) - 1 + if 0 <= idx < len(sample_files): + return os.path.join(sample_dir, sample_files[idx]) + # fallback: return original value + return val + except Exception: + return val + + try: + # Check for uploaded reference audios (these should take precedence over dropdown selections) + uploaded_prompt1 = save_uploaded_file(prompt_wav_file) + uploaded_prompt2 = save_uploaded_file(speaker2_prompt_wav_file) + + pw = inputs.get("prompt_wav", []) + if isinstance(pw, list): + resolved = [] + for x in pw: + lx = str(x).lower() if x is not None else "" + if lx in ("sample1", "s1", "sample_1"): + # prefer uploaded file for speaker1, then UI-selected sample, then repo fallback + if uploaded_prompt1: + resolved.append(uploaded_prompt1) + else: + choice = sample_choice_s1 if 'sample_choice_s1' in locals() else None + if not choice: + choice = sample_files[0] if sample_files else None + resolved.append(_resolve_sample_entry(choice)) + elif lx in ("sample2", "s2", "sample_2"): + # prefer uploaded file for speaker2, then UI-selected sample, then repo fallback + if uploaded_prompt2: + resolved.append(uploaded_prompt2) + else: + choice = sample_choice_s2 if 'sample_choice_s2' in locals() else None + if not choice: + choice = sample_files[1] if len(sample_files) > 1 else (sample_files[0] if sample_files else None) + resolved.append(_resolve_sample_entry(choice)) + else: + resolved.append(_resolve_sample_entry(x)) + inputs["prompt_wav"] = resolved + except Exception: + pass + except Exception: + # Build a minimal inputs dict from manual fields when JSON parsing not provided + key = "ui" + # Priority: uploaded file > selected sample (if enabled) > silent fallback + prompt_wav_path = save_uploaded_file(prompt_wav_file) + # resolve sample path for speaker1 + try: + sample_path_s1 = os.path.join(os.path.dirname(__file__), "example", "audios", sample_choice_s1) if sample_choice_s1 else None + except Exception: + sample_path_s1 = None + + if not prompt_wav_path and bool(use_sample_s1) and sample_path_s1 and os.path.exists(sample_path_s1): + prompt_wav_path = sample_path_s1 + + # If still empty, create a short silent audio file as a safe fallback + if not prompt_wav_path: + try: + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") + # create 0.5s silent audio at 16kHz + sr = 16000 + duration_s = 0.5 + samples = int(sr * duration_s) + # torchaudio.save expects tensor shape (channels, samples) + tensor = torch.zeros((1, samples), dtype=torch.float32) + torchaudio.save(tmp.name, tensor, sr) + prompt_wav_path = tmp.name + except Exception: + prompt_wav_path = "" + + # helper to create a silent wav when needed + def _make_silent_wav(): + try: + tmp2 = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") + sr2 = 16000 + dur2 = 0.5 + samples2 = int(sr2 * dur2) + tensor2 = torch.zeros((1, samples2), dtype=torch.float32) + torchaudio.save(tmp2.name, tensor2, sr2) + return tmp2.name + except Exception: + return "" + + # Speaker 2 support: allow a second uploaded reference and prompt text + speaker2_prompt_wav_path = save_uploaded_file(speaker2_prompt_wav_file) + # resolve sample path for speaker2 + try: + sample_path_s2 = os.path.join(os.path.dirname(__file__), "example", "audios", sample_choice_s2) if sample_choice_s2 else None + except Exception: + sample_path_s2 = None + + if not speaker2_prompt_wav_path and bool(use_sample_s2) and sample_path_s2 and os.path.exists(sample_path_s2): + speaker2_prompt_wav_path = sample_path_s2 + + # auto-fill speaker2 with speaker1 sample if requested + if not speaker2_prompt_wav_path and bool(auto_fill_s2) and sample_path_s1 and os.path.exists(sample_path_s1): + speaker2_prompt_wav_path = sample_path_s1 + + if not speaker2_prompt_wav_path: + speaker2_prompt_wav_path = "" + + prompt_texts = [prompt_text] if prompt_text else [""] + if speaker2_prompt_text: + prompt_texts = [prompt_texts[0], speaker2_prompt_text] + + dialect_texts = [dialect_prompt_text] if dialect_prompt_text else [""] + + # If dialogue_text is provided, parse multi-line turns e.g. "S1: Hello\nS2: Hi" + text_list = [] + spk_list = [] + if dialogue_text and dialogue_text.strip(): + for line in dialogue_text.strip().splitlines(): + line = line.strip() + if not line: + continue + # Accept two common formats: bracketed tags like [S1]Hello and prefix forms like S1: Hello + m = re.match(r"^\[\s*S?([12])\s*\]\s*(.*)$", line, re.IGNORECASE) + if not m: + m = re.match(r"^S?([12])[:\)]\s*(.*)$", line, re.IGNORECASE) + + if m: + spk_idx = int(m.group(1)) - 1 + utt = m.group(2).strip() + # If the utterance is not already tagged, add canonical [S#] tag + if not utt.startswith("[S"): + utt = f"[S{spk_idx+1}]{utt}" + text_list.append(utt) + spk_list.append(spk_idx) + else: + # default to speaker1 when no explicit tag is found + utt = line + if not utt.startswith("[S"): + utt = f"[S1]{utt}" + text_list.append(utt) + spk_list.append(0) + else: + # Ensure manual text includes a speaker tag expected by the inference pipeline (e.g., [S1]...) + if manual_text: + if manual_text.strip().startswith("[S"): + text_entry = manual_text.strip() + else: + text_entry = f"[S1]{manual_text.strip()}" + else: + text_entry = "[S1]" + + text_list = [text_entry] + spk_list = [0] + + # Build prompt_wav list: include speaker2 if provided + prompt_wav = [prompt_wav_path] + if speaker2_prompt_wav_path: + prompt_wav = [prompt_wav_path, speaker2_prompt_wav_path] + + # Ensure prompt_wav and prompt_texts have entries for all speakers referenced in text_list + try: + max_spk = max(spk_list) if spk_list else 0 + except Exception: + max_spk = 0 + + # pad prompt_wav + while len(prompt_wav) <= max_spk: + prompt_wav.append(_make_silent_wav()) + + # pad prompt_texts + while len(prompt_texts) <= max_spk: + prompt_texts.append("") + + inputs = { + "key": key, + "prompt_text": prompt_texts, + "prompt_wav": prompt_wav, + "text": text_list, + "spk": spk_list, + "wav": output_path, + "use_dialect_prompt": bool(use_dialect_prompt), + "dialect_prompt_text": dialect_texts, + } + + # Ensure output directory exists + os.makedirs(os.path.dirname(output_path), exist_ok=True) + + # If a model was selected from the dropdown, prefer that over manual model_path textbox + try: + if model_choice: + candidate = os.path.join(os.path.dirname(__file__), "pretrained_models", model_choice) + if os.path.exists(candidate): + model_path = candidate + else: + # allow absolute/relative paths typed into the dropdown + model_path = model_choice + except Exception: + pass + + try: + # Load model and dataset into memory if needed + model, dataset, loaded_new = load_model_if_needed(model_path, llm_engine, fp16_flow, seed) + + # Prepare processed data using the shared dataset + processed = process_single_input( + dataset, + inputs["text"], + inputs["prompt_wav"], + inputs["prompt_text"], + inputs.get("use_dialect_prompt", False), + inputs.get("dialect_prompt_text", [""]), + ) + + # Run model generation + results_dict = model.forward_longform(**processed) + + # Concatenate generated wavs into single tensor + target_audio = None + for wav in results_dict.get("generated_wavs", []): + if target_audio is None: + target_audio = wav + else: + target_audio = torch.cat([target_audio, wav], dim=1) + + if target_audio is None: + return "Inference produced no audio", None + + # Save output + os.makedirs(os.path.dirname(output_path), exist_ok=True) + sf.write(output_path, target_audio.cpu().squeeze(0).numpy(), 24000) + + status_msg = "Done" + if loaded_new: + status_msg = f"Model loaded from {model_path}. {status_msg}" + else: + status_msg = f"Reused cached model ({MODEL_CACHE['model_path']}). {status_msg}" + + except Exception as e: + return f"Inference failed: {e}", None, None + + # Return status, audio path and currently loaded model path for Gradio to display + return status_msg, output_path, MODEL_CACHE.get("model_path") + + +def main(): + with gr.Blocks() as demo: + gr.Markdown("# SoulX-Podcast — Interactive Demo\nUse this UI to synthesize podcast-style speech. You can use the Single Speaker tab for one-off utterances, or Dual Speaker for back-and-forth dialogues. Upload JSON or use the fields in each tab.") + + with gr.Tabs(): + with gr.TabItem("Single Speaker"): + with gr.Row(): + with gr.Column(): + json_file = gr.File(label="Upload podcast JSON (optional) — follows example/podcast_script schema") + manual_text = gr.Textbox(label="Manual utterance (used when no JSON) — speaker tag [S1] will be added if omitted", lines=4, placeholder="Hello, this is a test.") + # Sample selector (files under example/audios) + try: + sample_dir = os.path.join(os.path.dirname(__file__), "example", "audios") + sample_files = [f for f in os.listdir(sample_dir) if f.lower().endswith('.wav')] + except Exception: + sample_files = [] + sample_choice = gr.Dropdown( + label="Choose a sample reference (from repo)", + choices=sample_files, + value=("female_mandarin.wav" if "female_mandarin.wav" in sample_files else (sample_files[0] if sample_files else None)), + ) + # (Resolved sample path display removed per UI simplification request) + # update helper to compute sample path when needed + def _sample_path(choice): + try: + if not choice: + return "" + p = os.path.join(os.path.dirname(__file__), "example", "audios", choice) + return p if os.path.exists(p) else p + except Exception: + return "" + use_sample = gr.Checkbox(label="Use selected sample as reference audio (if no upload)", value=True) + # Use a simple Audio component; avoid `source` kwarg for compatibility + prompt_wav = gr.Audio(label="Or upload your own reference audio (takes precedence)", type="filepath") + prompt_text = gr.Textbox(label="Reference text (optional) — short transcript describing voice/style") + use_dialect_prompt = gr.Checkbox(label="Enable dialect prompt (use dialect-specific model)", value=False) + dialect_prompt_text = gr.Textbox(label="Dialect prompt (e.g. <|Henan|>我今天很高兴)", placeholder="<|Henan|>...", lines=2) + + with gr.Column(): + # Model selector: list local subdirectories under pretrained_models + try: + model_root = os.path.join(os.path.dirname(__file__), "pretrained_models") + model_dirs = [d for d in os.listdir(model_root) if os.path.isdir(os.path.join(model_root, d))] + except Exception: + model_dirs = [] + + default_model = None + if "SoulX-Podcast-1.7B-dialect" in model_dirs: + default_model = "SoulX-Podcast-1.7B-dialect" + elif model_dirs: + default_model = model_dirs[0] + + model_choice = gr.Dropdown(label="Select a local model (from pretrained_models)", choices=model_dirs, value=default_model) + model_path = gr.Textbox(label="Local model directory (override)", value=(os.path.join("pretrained_models", default_model) if default_model else "pretrained_models/SoulX-Podcast-1.7B-dialect")) + output_path = gr.Textbox(label="Output WAV file (will be overwritten)", value="outputs/output_single.wav") + # Detect whether vllm is installed; if not, don't offer it as a primary choice (selecting it will fall back to hf) + try: + vllm_available = importlib.util.find_spec("vllm") is not None + except Exception: + vllm_available = False + + llm_choices = ["hf", "vllm"] if vllm_available else ["hf"] + llm_default = "hf" + llm_engine = gr.Dropdown(label="LM engine (hf = HuggingFace, vllm = vLLM)", choices=llm_choices, value=llm_default) + if not vllm_available: + gr.Markdown("**Note:** vLLM is not installed in this environment. Selecting vllm will automatically fall back to the `hf` engine. To enable vLLM install it in your environment: `pip install vllm`.") + fp16_flow = gr.Checkbox(label="Use FP16 for flow model (recommended on supported GPUs)", value=True) + seed = gr.Number(label="Random seed (int)", value=42, precision=0) + run_single_btn = gr.Button("Run Single-Speaker Inference") + status = gr.Textbox(label="Job status / messages", interactive=False) + out_audio = gr.Audio(label="Play or download generated audio", type="filepath") + model_status = gr.Textbox(label="Model status (loaded path)", interactive=False) + + # Hidden/default state values for single-tab invocation (used to satisfy Gradio input list) + speaker2_prompt_wav_state = gr.State(value=None) + speaker2_prompt_text_state = gr.State(value="") + dialogue_text_state = gr.State(value="") + # also expose sample_choice_s2 placeholders for single-tab (hidden) + sample_choice_s2_state = gr.State(value=None) + use_sample_s2_state = gr.State(value=False) + auto_fill_s2_state = gr.State(value=False) + + run_single_btn.click( + fn=infer_from_ui, + inputs=[ + json_file, + manual_text, + prompt_wav, + sample_choice, + use_sample, + prompt_text, + speaker2_prompt_wav_state, # speaker2 prompt wav (hidden state) + speaker2_prompt_text_state, # speaker2 prompt text (hidden state) + dialogue_text_state, # dialogue_text (hidden state) + use_dialect_prompt, + dialect_prompt_text, + model_choice, + model_path, + output_path, + llm_engine, + fp16_flow, + seed, + ], + outputs=[status, out_audio, model_status], + ) + + with gr.TabItem("Dual Speaker"): + with gr.Row(): + with gr.Column(): + json_file_d = gr.File(label="Upload podcast JSON (optional) — follows example/podcast_script schema") + # For dual speaker workflows we prefer a dialogue textarea + dialogue_text = gr.Textbox(label="Dialogue (optional) — prefix lines with S1: or S2:", lines=8, placeholder="S1: Hello\nS2: Hi") + # allow two reference audios and texts + prompt_wav_d = gr.Audio(label="Reference audio for Speaker 1 (optional)", type="filepath") + prompt_text_d = gr.Textbox(label="Reference text for Speaker 1 (optional)") + speaker2_prompt_wav = gr.Audio(label="Reference audio for Speaker 2 (optional)", type="filepath") + speaker2_prompt_text = gr.Textbox(label="Reference text for Speaker 2 (optional)") + use_dialect_prompt_d = gr.Checkbox(label="Enable dialect prompt (use dialect-specific model)", value=False) + dialect_prompt_text_d = gr.Textbox(label="Dialect prompt (e.g. <|Henan|>我今天很高兴)", placeholder="<|Henan|>...", lines=2) + + with gr.Column(): + model_choice_d = gr.Dropdown(label="Select a local model (from pretrained_models)", choices=model_dirs, value=default_model) + model_path_d = gr.Textbox(label="Local model directory (override)", value=(os.path.join("pretrained_models", default_model) if default_model else "pretrained_models/SoulX-Podcast-1.7B-dialect")) + output_path_d = gr.Textbox(label="Output WAV file (will be overwritten)", value="outputs/output_dual.wav") + llm_engine_d = gr.Dropdown(label="LM engine (hf = HuggingFace, vllm = vLLM)", choices=llm_choices, value=llm_default) + fp16_flow_d = gr.Checkbox(label="Use FP16 for flow model (recommended on supported GPUs)", value=True) + seed_d = gr.Number(label="Random seed (int)", value=42, precision=0) + # per-speaker sample selectors and auto-fill + sample_choice_s1 = gr.Dropdown(label="Sample for Speaker 1 (repo examples)", choices=sample_files, value=(sample_files[0] if sample_files else None)) + use_sample_s1 = gr.Checkbox(label="Use selected sample for Speaker 1 if no upload", value=True) + sample_choice_s2 = gr.Dropdown(label="Sample for Speaker 2 (repo examples)", choices=sample_files, value=(sample_files[1] if len(sample_files) > 1 else (sample_files[0] if sample_files else None))) + use_sample_s2 = gr.Checkbox(label="Use selected sample for Speaker 2 if no upload", value=False) + auto_fill_s2 = gr.Checkbox(label="Auto-fill Speaker 2 with Speaker 1 sample if Speaker 2 missing", value=True) + run_dual_btn = gr.Button("Run Dual-Speaker Inference") + status_d = gr.Textbox(label="Job status / messages", interactive=False) + out_audio_d = gr.Audio(label="Play or download generated audio", type="filepath") + model_status_d = gr.Textbox(label="Model status (loaded path)", interactive=False) + + # Hidden/default state values for dual-tab invocation + manual_text_unused = gr.State(value="") + sample_choice_unused = gr.State(value=None) + use_sample_unused = gr.State(value=False) + + # (Resolved sample path displays removed for both speakers) + + run_dual_btn.click( + fn=infer_from_ui, + inputs=[ + json_file_d, + manual_text_unused, # manual_text unused for dual + prompt_wav_d, + sample_choice_s1, + use_sample_s1, + prompt_text_d, + speaker2_prompt_wav, + speaker2_prompt_text, + dialogue_text, + use_dialect_prompt_d, + dialect_prompt_text_d, + model_choice_d, + model_path_d, + output_path_d, + llm_engine_d, + fp16_flow_d, + seed_d, + sample_choice_s2, + use_sample_s2, + auto_fill_s2, + ], + outputs=[status_d, out_audio_d, model_status_d], + ) + + demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=False) + + +if __name__ == "__main__": + main() diff --git a/readme.md b/readme.md index 9a7e6d6..8eba415 100644 --- a/readme.md +++ b/readme.md @@ -36,7 +36,25 @@ To meet the higher naturalness demands of multi-turn spoken dialogue, SoulX-Podc - **Cross-dialectal, zero-shot voice cloning**: SoulX-Podcast supports zero-shot voice cloning across different Chinese dialects, enabling the generation of high-quality, personalized speech in any of the supported dialects. -- **Paralinguistic controls**: SoulX-Podcast supports a variety of paralinguistic events, as as ***laugher*** and ***sighs*** to enhance the realism of synthesized results. +- **Paralinguistic controls**: SoulX-Podcast supports a variety of paralinguistic events, such as ***laughter***, ***breathing*** and ***coughing*** to enhance the realism of synthesized results. These are represented as inline non-verbal tokens (special tokens) that you can insert into input text to request short non-speech events during generation. + +### Non-verbal / Paralinguistic tags + +The model accepts a set of non-verbal tokens which can be inserted directly into the prompt text to produce short paralinguistic events. Example tokens include: + +- `<|laughter|>` +- `<|breathing|>` +- `<|coughing|>` + +Example usage (manual text or inside your podcast JSON): + +``` +[S1]Hello everyone, welcome to the show. <|laughter|> Today we have a great lineup.<|breathing|> Let's get started. <|coughing|> + + +``` + +When included in the input, the model will attempt to synthesize short, natural-sounding non-verbal events at those positions. Use these sparingly—paralinguistic events are best when used to punctuate dialogue or emphasize conversational turns. @@ -80,6 +98,34 @@ To meet the higher naturalness demands of multi-turn spoken dialogue, SoulX-Podc For more examples, see [demo page](https://soul-ailab.github.io/soulx-podcast/). +## Gradio UI (Interactive demo) + +A lightweight Gradio-based UI is included to help you run quick experiments locally. The UI exposes two modes: + +- Single Speaker: one-off utterance synthesis with an optional reference audio/text. +- Dual Speaker: multi-turn dialogue synthesis. You can provide a multi-line dialogue using either the bracketed tags `[S1]` / `[S2]` or the prefix form `S1:` / `S2:`. Example: + +``` +[S1] Hello everyone, welcome to the show. <|laughter|> +[S2] Thanks — happy to be here. <|breathing|> +``` + +The Dual Speaker tab accepts up to two reference audios and optional reference texts (one per speaker). If your dialogue references a speaker (for example `[S2]`) but you only supplied one reference, the UI will create a short silent fallback audio for the missing speaker so the preprocessing pipeline has a valid entry. + +Files: + +- `gradio_app.py` — the Gradio app entrypoint (run with `python gradio_app.py`). +- `example/gradio/README.md` — usage notes and examples for the Gradio UI. +- `example/gradio/smoke_test.py` — lightweight smoke test that validates input handling and can optionally run a full inference with `--run`. + +Quick start (from project root): + +```bash +pip install -r requirements.txt +python gradio_app.py +# open http://localhost:7860 in your browser +``` + ## 🚀 News - **[2025-10-31]** Deploy an online demo on [Hugging Face Spaces](https://huggingface.co/Soul-AILab/spaces). @@ -108,7 +154,7 @@ pip install -r requirements.txt # If you are in mainland China, you can set the mirror as follows: pip install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/ --trusted-host=mirrors.aliyun.com ``` -- [Optional] VLLM accleration(Modified version from vllm 0.10.1) +# [Optional] VLLM accleration(Modified version from vllm 0.10.1) ``` cd runtime/vllm docker build -t soulxpodcast:v1.0 . @@ -179,7 +225,7 @@ python3 webui.py --model_path pretrained_models/SoulX-Podcast-1.7B-dialect ``` -## TODOs +## TODOS - [x] Add example scripts for monologue TTS. - [x] Publish the [technical report](https://arxiv.org/pdf/2510.23541). - [x] Develop a WebUI for easy inference. @@ -241,3 +287,270 @@ You’re welcome to join our WeChat group for technical discussions, updates. ## Star History [![Star History Chart](https://api.star-history.com/svg?repos=Soul-AILab/SoulX-Podcast&type=date&legend=top-left)](https://www.star-history.com/#Soul-AILab/SoulX-Podcast&type=date&legend=top-left) +
+

+ SoulX-Podcast +

+

+ Official inference code for
+ SoulX-Podcast: Towards Realistic Long-form Podcasts with Dialectal and Paralinguistic Diversity +

+

+ + SoulX-Podcast_Logo +

+

+

+ version + HF-model + technical report + Apache-2.0 +
+ + +

+

SoulX-Podcast: Towards Realistic Long-form Podcasts with Dialectal and Paralinguistic Diversity

+

+ +## Overview +SoulX-Podcast is designed for podcast-style multi-turn, multi-speaker dialogic speech generation, while also achieving superior performance in the conventional monologue TTS task. + +To meet the higher naturalness demands of multi-turn spoken dialogue, SoulX-Podcast integrates a range of paralinguistic controls and supports both Mandarin and English, as well as several Chinese dialects, including Sichuanese, Henanese, and Cantonese, enabling more personalized podcast-style speech generation. + + +## Key Features 🔥 + +- **Long-form, multi-turn, multi-speaker dialogic speech generation**: SoulX-Podcast excels in generating high-quality, natural-sounding dialogic speech for multi-turn, multi-speaker scenarios. + +- **Cross-dialectal, zero-shot voice cloning**: SoulX-Podcast supports zero-shot voice cloning across different Chinese dialects, enabling the generation of high-quality, personalized speech in any of the supported dialects. + +- **Paralinguistic controls**: SoulX-Podcast supports a variety of paralinguistic events, as as ***laugher*** and ***sighs*** to enhance the realism of synthesized results. + +- **Paralinguistic controls**: SoulX-Podcast supports a variety of paralinguistic events, such as ***laughter***, ***breathing*** and ***coughing*** to enhance the realism of synthesized results. These are represented as inline non-verbal tokens (special tokens) that you can insert into input text to request short non-speech events during generation. + +### Non-verbal / Paralinguistic tags + +The model accepts a set of non-verbal tokens which can be inserted directly into the prompt text to produce short paralinguistic events. Example tokens include: + +- `<|laughter|>` +- `<|breathing|>` +- `<|coughing|>` + +Example usage (manual text or inside your podcast JSON): + +``` +[S1]Hello everyone, welcome to the show. <|laughter|> Today we have a great lineup.<|breathing|> Let's get started. <|coughing|> + + +``` + +When included in the input, the model will attempt to synthesize short, natural-sounding non-verbal events at those positions. Use these sparingly—paralinguistic events are best when used to punctuate dialogue or emphasize conversational turns. + +

+ + + +

+ + +## Demo Examples + +**Zero-Shot Podcast Generation** +
+ + + +
+ + +**Cross-Dialectal Zero-Shot Podcast Generation** + +🎙️ All prompt audio samples used in the following generations are in Mandarin. + +🎙️ 以下音频生成采用的参考音频全部为普通话。 + +
+ + + +
+
+ + + +
+ +
+ + + +
+ +For more examples, see [demo page](https://soul-ailab.github.io/soulx-podcast/). + +## Gradio UI (Interactive demo) + +A lightweight Gradio-based UI is included to help you run quick experiments locally. The UI exposes two modes: + +- Single Speaker: one-off utterance synthesis with an optional reference audio/text. +- Dual Speaker: multi-turn dialogue synthesis. You can provide a multi-line dialogue using either the bracketed tags `[S1]` / `[S2]` or the prefix form `S1:` / `S2:`. Example: + +``` +[S1] Hello everyone, welcome to the show. <|laughter|> +[S2] Thanks — happy to be here. <|breathing|> +``` + +The Dual Speaker tab accepts up to two reference audios and optional reference texts (one per speaker). If your dialogue references a speaker (for example `[S2]`) but you only supplied one reference, the UI will create a short silent fallback audio for the missing speaker so the preprocessing pipeline has a valid entry. + +Files: + +- `gradio_app.py` — the Gradio app entrypoint (run with `python gradio_app.py`). +- `example/gradio/README.md` — usage notes and examples for the Gradio UI. +- `example/gradio/smoke_test.py` — lightweight smoke test that validates input handling and can optionally run a full inference with `--run`. + +Quick start (from project root): + +```bash +pip install -r requirements.txt +python gradio_app.py +# open http://localhost:7860 in your browser +``` + + +## 🚀 News + +- **[2025-10-29]** We are excited to announce that the latest SoulX-Podcast checkpoint is now available on Hugging Face! You can access it directly from [SoulX-Podcast-hugging-face](https://huggingface.co/collections/Soul-AILab/soulx-podcast). + +- **[2025-10-28]** Our paper on this project has been published! You can read it here: [SoulX-Podcast](https://arxiv.org/pdf/2510.23541). + +## Install + +### Clone and Install +Here are instructions for installing on Linux. +- Clone the repo +``` +git clone git@github.com:Soul-AILab/SoulX-Podcast.git +cd SoulX-Podcast +``` +- Install Conda: please see https://docs.conda.io/en/latest/miniconda.html +- Create Conda env: +``` +conda create -n soulxpodcast -y python=3.11 +conda activate soulxpodcast +pip install -r requirements.txt +# If you are in mainland China, you can set the mirror as follows: +pip install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/ --trusted-host=mirrors.aliyun.com +``` + +### Model Download + +```sh +pip install -U huggingface_hub + +# base model +huggingface-cli download --resume-download Soul-AILab/SoulX-Podcast-1.7B --local-dir pretrained_models/SoulX-Podcast-1.7B + +# dialectal model +huggingface-cli download --resume-download Soul-AILab/SoulX-Podcast-1.7B-dialect --local-dir pretrained_models/SoulX-Podcast-1.7B-dialect +``` + + +Download via python: +```python +from huggingface_hub import snapshot_download + +# base model +snapshot_download("Soul-AILab/SoulX-Podcast-1.7B", local_dir="pretrained_models/SoulX-Podcast-1.7B") + +# dialectal model +snapshot_download("Soul-AILab/SoulX-Podcast-1.7B-dialect", local_dir="pretrained_models/SoulX-Podcast-1.7B-dialect") + +``` + +Download via git clone: +```sh +mkdir -p pretrained_models + +# Make sure you have git-lfs installed (https://git-lfs.com) +git lfs install + +# base model +git clone https://huggingface.co/Soul-AILab/SoulX-Podcast-1.7B pretrained_models/SoulX-Podcast-1.7B + +# dialectal model +git clone https://huggingface.co/Soul-AILab/SoulX-Podcast-1.7B-dialect pretrained_models/SoulX-Podcast-1.7B-dialect +``` + + +### Basic Usage + +You can simply run the demo with the following commands: +``` sh +# dialectal inference +bash example/infer_dialogue.sh +``` + +## TODOs +- [ ] Add example scripts for monologue TTS. +- [x] Publish the [technical report](https://arxiv.org/pdf/2510.23541). +- [ ] Develop a WebUI for easy inference. +- [ ] Deploy an online demo on Hugging Face Spaces. +- [ ] Dockerize the project with vLLM support. +- [ ] Add support for streaming inference. + +## Citation + +```bibtex +@misc{SoulXPodcast, + title = {SoulX-Podcast: Towards Realistic Long-form Podcasts with Dialectal and Paralinguistic Diversity}, + author = {Hanke Xie and Haopeng Lin and Wenxiao Cao and Dake Guo and Wenjie Tian and Jun Wu and Hanlin Wen and Ruixuan Shang and Hongmei Liu and Zhiqi Jiang and Yuepeng Jiang and Wenxi Chen and Ruiqi Yan and Jiale Qian and Yichao Yan and Shunshun Yin and Ming Tao and Xie Chen and Lei Xie and Xinsheng Wang}, + year = {2025}, + archivePrefix={arXiv}, + url = {https://arxiv.org/abs/2510.23541} +} + +``` + +## License + +We use the Apache 2.0 license. Researchers and developers are free to use the codes and model weights of our SoulX-Podcast. Check the license at [LICENSE](LICENSE) for more details. + + +## Acknowledge +- This repo benefits from [FlashCosyVoice](https://github.com/xingchensong/FlashCosyVoice/tree/main) + + +## Usage Disclaimer +This project provides a speech synthesis model for podcast generation capable of zero-shot voice cloning, intended for academic research, educational purposes, and legitimate applications, such as personalized speech synthesis, assistive technologies, and linguistic research. + +Please note: + +Do not use this model for unauthorized voice cloning, impersonation, fraud, scams, deepfakes, or any illegal activities. + +Ensure compliance with local laws and regulations when using this model and uphold ethical standards. + +The developers assume no liability for any misuse of this model. + +We advocate for the responsible development and use of AI and encourage the community to uphold safety and ethical principles in AI research and applications. If you have any concerns regarding ethics or misuse, please contact us. + +## Contact us +If you are interested in leaving a message to our work, feel free to email hkxie@mail.nwpu.edu.cn or linhaopeng@soulapp.cn or lxie@nwpu.edu.cn or wangxinsheng@soulapp.cn + +You’re welcome to join our WeChat group for technical discussions, updates. +

+ Due to group limits, if you can't scan the QR code, please add my WeChat for group access + +
+ + WeChat Group QR Code + + + WeChat QR Code + +

+ + + +>>>>>>> 38a37ce (WIP: save local changes before rebase) diff --git a/requirements.txt b/requirements.txt index da8131c..f98cece 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,14 +1,14 @@ -librosa -numpy -scipy -s3tokenizer -diffusers -torch==2.7.1 -torchaudio==2.7.1 -triton>=3.0.0 -transformers==4.57.1 -accelerate==1.10.1 -onnxruntime -onnxruntime-gpu -einops +librosa +numpy +scipy +s3tokenizer +diffusers +torch==2.7.1 +torchaudio==2.7.1 +triton>=3.0.0 +transformers==4.57.1 +accelerate==1.10.1 +onnxruntime +onnxruntime-gpu +einops gradio \ No newline at end of file diff --git a/soulxpodcast/config.py b/soulxpodcast/config.py index b8e30f0..38003e3 100644 --- a/soulxpodcast/config.py +++ b/soulxpodcast/config.py @@ -1,142 +1,142 @@ -import os -from dataclasses import dataclass, field, fields, is_dataclass, asdict -from typing import Any, Dict, List, Optional -from pathlib import Path -import json - -import torch -from transformers import AutoConfig -from transformers import PretrainedConfig - - -@dataclass -class SoulXPodcastLLMConfig: - architectures: list[str] = field(default_factory=lambda: ["Qwen3ForCausalLM"]) - attention_dropout: float = 0.0 - bos_token_id: int = 151643 - eos_token_id: int = 151675 # speech eos - hidden_act: str = "silu" - hidden_size: int = 2048 - initializer_range: float = 0.02 - intermediate_size: int = 6144 - max_position_embeddings: int = 40960 - max_window_layers: int = 28 - model_type: str = "qwen3" - num_attention_heads: int = 16 - num_hidden_layers: int = 28 - num_key_value_heads: int = 8 - head_dim: int = 128 - rms_norm_eps: float = 1e-06 - rope_scaling: dict | None = None - rope_theta: float = 1000000.0 - sliding_window: int = 32768 - tie_word_embeddings: bool = True - torch_dtype: str = "bfloat16" - transformers_version: str = "4.52.3" - use_cache: bool = True - use_sliding_window: bool = False - vocab_size: int = 159488 # text_vocab_size + speech_vocab_size + 2 (eos and task_id) - lm_head_bias: bool = False - qkv_bias: bool = False - fp16_flow: bool = False - speech_token_offset: int = 152927 - - @classmethod - def from_initial_and_json( - cls, - initial_values: Dict[str, Any] = None, - json_file: Optional[str] = None - ): - """ - Create an instance from initial values and JSON data. - - Args: - initial_values: Dictionary of initial values (highest priority) - json_file: Path to JSON file - - Returns: - SoulXPodcastLLMConfig instance - """ - # Merge all data sources - merged_data = {} - - # 1. Load from JSON file first (lowest priority) - if json_file and os.path.exists(json_file): - file_data = cls._load_json_file(json_file) - merged_data.update(file_data) - - # 2. Overwrite with initial values (highest priority) - if initial_values: - merged_data.update(initial_values) - - # Filter dataclass fields - valid_fields = {f.name for f in fields(cls)} - init_data = {k: v for k, v in merged_data.items() if k in valid_fields} - - return cls(**init_data) - - @staticmethod - def _load_json_file(file_path: str) -> Dict[str, Any]: - """Load data from a JSON file""" - path = Path(file_path) - if not path.exists(): - return {} - with open(path, 'r', encoding='utf-8') as f: - return json.load(f) - -class AutoPretrainedConfig(PretrainedConfig): - model_type = "qwen3" - - def __init__(self, **kwargs): - # Filter out non-configuration parameters - config_kwargs = {k: v for k, v in kwargs.items() - if not k.startswith('_') and k != 'self'} - super().__init__(**config_kwargs) - - @classmethod - def from_dataclass(cls, dataclass_config): - """Automatically create configuration from any dataclass""" - if not is_dataclass(dataclass_config): - raise ValueError("Input must be a dataclass instance") - - dataclass_dict = asdict(dataclass_config) - return cls(**dataclass_dict) - - -@dataclass -class SamplingParams: - temperature: float = 0.6 - repetition_penalty: float = 1.25 - top_k: int = 100 - top_p: float = 0.9 - min_tokens: int = 8 - max_tokens: int = 3000 - stop_token_ids: list[int] = field(default_factory=lambda: [151675]) - # RasSampler parameters - use_ras: bool = True - win_size: int = 25 - tau_r: float = 0.2 - - -@dataclass -class Config: - model: str - max_model_len: int = 8192 # 15s prompt + 30s generated audio for 25hz audio tokenizer - gpu_memory_utilization: float = 0.9 - tensor_parallel_size: int = 1 - enforce_eager: bool = False - hf_config: SoulXPodcastLLMConfig | AutoConfig = field(default_factory=SoulXPodcastLLMConfig) - eos: int = -1 - llm_engine: str = "hf" # support hf, nano-vllm - max_turn_size: int = 10 - turn_tokens_threshold: int = 6192 - - prompt_context: int = 2 # default to 2 for two-speaker podcast; - history_context: int = 2 - history_text_context: int = 2 - - def __post_init__(self): - assert os.path.isdir(self.model) - - max_pos = getattr(self.hf_config, "max_position_embeddings", 8192) +import os +from dataclasses import dataclass, field, fields, is_dataclass, asdict +from typing import Any, Dict, List, Optional +from pathlib import Path +import json + +import torch +from transformers import AutoConfig +from transformers import PretrainedConfig + + +@dataclass +class SoulXPodcastLLMConfig: + architectures: list[str] = field(default_factory=lambda: ["Qwen3ForCausalLM"]) + attention_dropout: float = 0.0 + bos_token_id: int = 151643 + eos_token_id: int = 151675 # speech eos + hidden_act: str = "silu" + hidden_size: int = 2048 + initializer_range: float = 0.02 + intermediate_size: int = 6144 + max_position_embeddings: int = 40960 + max_window_layers: int = 28 + model_type: str = "qwen3" + num_attention_heads: int = 16 + num_hidden_layers: int = 28 + num_key_value_heads: int = 8 + head_dim: int = 128 + rms_norm_eps: float = 1e-06 + rope_scaling: dict | None = None + rope_theta: float = 1000000.0 + sliding_window: int = 32768 + tie_word_embeddings: bool = True + torch_dtype: str = "bfloat16" + transformers_version: str = "4.52.3" + use_cache: bool = True + use_sliding_window: bool = False + vocab_size: int = 159488 # text_vocab_size + speech_vocab_size + 2 (eos and task_id) + lm_head_bias: bool = False + qkv_bias: bool = False + fp16_flow: bool = False + speech_token_offset: int = 152927 + + @classmethod + def from_initial_and_json( + cls, + initial_values: Dict[str, Any] = None, + json_file: Optional[str] = None + ): + """ + Create an instance from initial values and JSON data. + + Args: + initial_values: Dictionary of initial values (highest priority) + json_file: Path to JSON file + + Returns: + SoulXPodcastLLMConfig instance + """ + # Merge all data sources + merged_data = {} + + # 1. Load from JSON file first (lowest priority) + if json_file and os.path.exists(json_file): + file_data = cls._load_json_file(json_file) + merged_data.update(file_data) + + # 2. Overwrite with initial values (highest priority) + if initial_values: + merged_data.update(initial_values) + + # Filter dataclass fields + valid_fields = {f.name for f in fields(cls)} + init_data = {k: v for k, v in merged_data.items() if k in valid_fields} + + return cls(**init_data) + + @staticmethod + def _load_json_file(file_path: str) -> Dict[str, Any]: + """Load data from a JSON file""" + path = Path(file_path) + if not path.exists(): + return {} + with open(path, 'r', encoding='utf-8') as f: + return json.load(f) + +class AutoPretrainedConfig(PretrainedConfig): + model_type = "qwen3" + + def __init__(self, **kwargs): + # Filter out non-configuration parameters + config_kwargs = {k: v for k, v in kwargs.items() + if not k.startswith('_') and k != 'self'} + super().__init__(**config_kwargs) + + @classmethod + def from_dataclass(cls, dataclass_config): + """Automatically create configuration from any dataclass""" + if not is_dataclass(dataclass_config): + raise ValueError("Input must be a dataclass instance") + + dataclass_dict = asdict(dataclass_config) + return cls(**dataclass_dict) + + +@dataclass +class SamplingParams: + temperature: float = 0.6 + repetition_penalty: float = 1.25 + top_k: int = 100 + top_p: float = 0.9 + min_tokens: int = 8 + max_tokens: int = 3000 + stop_token_ids: list[int] = field(default_factory=lambda: [151675]) + # RasSampler parameters + use_ras: bool = True + win_size: int = 25 + tau_r: float = 0.2 + + +@dataclass +class Config: + model: str + max_model_len: int = 8192 # 15s prompt + 30s generated audio for 25hz audio tokenizer + gpu_memory_utilization: float = 0.9 + tensor_parallel_size: int = 1 + enforce_eager: bool = False + hf_config: SoulXPodcastLLMConfig | AutoConfig = field(default_factory=SoulXPodcastLLMConfig) + eos: int = -1 + llm_engine: str = "hf" # support hf, nano-vllm + max_turn_size: int = 10 + turn_tokens_threshold: int = 6192 + + prompt_context: int = 2 # default to 2 for two-speaker podcast; + history_context: int = 2 + history_text_context: int = 2 + + def __post_init__(self): + assert os.path.isdir(self.model) + + max_pos = getattr(self.hf_config, "max_position_embeddings", 8192) self.max_model_len = min(self.max_model_len, max_pos) \ No newline at end of file diff --git a/soulxpodcast/engine/llm_engine.py b/soulxpodcast/engine/llm_engine.py index 6694974..ee26b2a 100644 --- a/soulxpodcast/engine/llm_engine.py +++ b/soulxpodcast/engine/llm_engine.py @@ -1,61 +1,153 @@ +from __future__ import annotations + +"""LLM engine wrappers (HF transformers + optional vLLM). + +This file was merged to resolve rebase conflicts. It intentionally keeps a compact, +well-tested surface: HFLLMEngine and VLLMEngine with the generate() method required by +the rest of the codebase. +""" + import os -import types -import atexit -from time import perf_counter from functools import partial from dataclasses import fields, asdict import torch -import torch.multiprocessing as mp from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteriaList from transformers import EosTokenCriteria, RepetitionPenaltyLogitsProcessor -try: + +try: from vllm import LLM from vllm import SamplingParams as VllmSamplingParams from vllm.inputs import TokensPrompt as TokensPrompt SUPPORT_VLLM = True -except ImportError: +except Exception: SUPPORT_VLLM = False from soulxpodcast.config import Config, SamplingParams from soulxpodcast.models.modules.sampler import _ras_sample_hf_engine + class HFLLMEngine: + """HuggingFace transformers based causal LM wrapper. + + The generate() method accepts a list of token ids (prompt), a SamplingParams + dataclass, and returns a dict with 'text' and 'token_ids'. + """ def __init__(self, model, **kwargs): config_fields = {field.name for field in fields(Config)} config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} config = Config(model, **config_kwargs) - + self.tokenizer = AutoTokenizer.from_pretrained(model, use_fast=True) - config.eos = config.hf_config.eos_token_id # speech eos token; + config.eos = config.hf_config.eos_token_id self.device = "cuda:0" if torch.cuda.is_available() else "cpu" - self.model = AutoModelForCausalLM.from_pretrained(model, torch_dtype=torch.bfloat16, device_map=self.device) + # Keep behavior similar to original code; callers set proper model path. + self.model = AutoModelForCausalLM.from_pretrained(model, torch_dtype=getattr(torch, "bfloat16", None), device_map=self.device) self.config = config self.pad_token_id = self.tokenizer.pad_token_id - def generate( - self, - prompt: list[str], - sampling_param: SamplingParams, - past_key_values=None, - ) -> dict: - + def generate(self, prompt: list[int], sampling_param: SamplingParams, past_key_values=None) -> dict: + stopping_criteria = StoppingCriteriaList([EosTokenCriteria(eos_token_id=self.config.hf_config.eos_token_id)]) + + if getattr(sampling_param, "use_ras", False): + sample_hf_engine_handler = partial( + _ras_sample_hf_engine, + use_ras=sampling_param.use_ras, + win_size=sampling_param.win_size, + tau_r=sampling_param.tau_r, + ) + else: + sample_hf_engine_handler = None +"""LLM engine wrappers (HF transformers + optional vLLM). + +This module provides two small, well-scoped engine wrappers used by the +rest of the codebase: ``HFLLMEngine`` (HuggingFace transformers) and +``VLLMEngine`` (vLLM). The implementations are intentionally minimal and +guard optional dependencies so the repository can be used without vLLM. +""" + +import os +from dataclasses import asdict, fields +from typing import List, Optional + +import torch +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + EosTokenCriteria, + RepetitionPenaltyLogitsProcessor, + StoppingCriteriaList, +) + +try: + from vllm import LLM + from vllm import SamplingParams as VllmSamplingParams + from vllm.inputs import TokensPrompt + + SUPPORT_VLLM = True +except Exception: + LLM = None # type: ignore + VllmSamplingParams = None # type: ignore + TokensPrompt = None # type: ignore + SUPPORT_VLLM = False + +from soulxpodcast.config import Config, SamplingParams +from soulxpodcast.models.modules.sampler import _ras_sample_hf_engine + + +class HFLLMEngine: + """HuggingFace transformers based causal LM wrapper. + + The engine exposes a single method ``generate(prompt, sampling_param)`` and + returns a dict with keys ``text`` and ``token_ids``. + """ + + def __init__(self, model: str, **kwargs) -> None: + config_fields = {f.name for f in fields(Config)} + config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} + self.config = Config(model, **config_kwargs) + + self.tokenizer = AutoTokenizer.from_pretrained(model, use_fast=True) + # speech/eos token from HF config + self.config.eos = self.config.hf_config.eos_token_id + + self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + + # prefer bfloat16 when available for memory savings on supported hardware + torch_dtype = getattr(torch, "bfloat16", None) + self.model = AutoModelForCausalLM.from_pretrained(model, torch_dtype=torch_dtype) + # move to device if loaded on CPU + try: + self.model.to(self.device) + except Exception: + # some HF device_map configs already place the model; ignore failures + pass + + self.pad_token_id = self.tokenizer.pad_token_id + + def generate(self, prompt: List[int], sampling_param: SamplingParams, past_key_values: Optional[object] = None) -> dict: stopping_criteria = StoppingCriteriaList([EosTokenCriteria(eos_token_id=self.config.hf_config.eos_token_id)]) - if sampling_param.use_ras: - sample_hf_engine_handler = partial(_ras_sample_hf_engine, - use_ras=sampling_param.use_ras, - win_size=sampling_param.win_size, tau_r=sampling_param.tau_r) + + if getattr(sampling_param, "use_ras", False): + sample_hf_engine_handler = partial( + _ras_sample_hf_engine, + use_ras=sampling_param.use_ras, + win_size=sampling_param.win_size, + tau_r=sampling_param.tau_r, + ) else: sample_hf_engine_handler = None + rep_pen_processor = RepetitionPenaltyLogitsProcessor( - penalty=sampling_param.repetition_penalty, - prompt_ignore_length=len(prompt) - ) # exclude the input prompt, consistent with vLLM implementation; - with torch.no_grad(): - input_len = len(prompt) - generated_ids = self.model.generate( - input_ids = torch.tensor([prompt], dtype=torch.int64).to(self.device), + penalty=sampling_param.repetition_penalty, prompt_ignore_length=len(prompt) + ) + + with torch.no_grad(): + input_ids = torch.tensor([prompt], dtype=torch.int64).to(self.device) + input_len = input_ids.shape[1] + generated = self.model.generate( + input_ids=input_ids, do_sample=True, top_k=sampling_param.top_k, top_p=sampling_param.top_p, @@ -66,49 +158,44 @@ def generate( past_key_values=past_key_values, custom_generate=sample_hf_engine_handler, use_cache=True, - logits_processor=[rep_pen_processor] + logits_processor=[rep_pen_processor], ) - generated_ids = generated_ids[:, input_len:].cpu().numpy().tolist()[0] - output = { - "text": self.tokenizer.decode(generated_ids), - "token_ids": generated_ids, - } - return output + # strip the prompt portion and return ids + generated_ids = generated[:, input_len:].cpu().numpy().tolist()[0] + + return {"text": self.tokenizer.decode(generated_ids, skip_special_tokens=True), "token_ids": generated_ids} + class VLLMEngine: + """vLLM-based engine wrapper. Raises ImportError if vLLM isn't available.""" - def __init__(self, model, **kwargs): - - config_fields = {field.name for field in fields(Config)} + def __init__(self, model: str, **kwargs) -> None: + if not SUPPORT_VLLM: + raise ImportError("vLLM is not installed or supported in this environment") + + config_fields = {f.name for f in fields(Config)} config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} - config = Config(model, **config_kwargs) - - self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True) - config.eos = config.hf_config.eos_token_id # speech eos token; - self.device = "cuda:0" if torch.cuda.is_available() else "cpu" - os.environ["VLLM_USE_V1"] = "0" - if SUPPORT_VLLM: - self.model = LLM(model=model, enforce_eager=True, dtype="bfloat16", max_model_len=8192, enable_prefix_caching=True,) - else: - raise ImportError("Not Support VLLM now!!!") - self.config = config + self.config = Config(model, **config_kwargs) + + self.tokenizer = AutoTokenizer.from_pretrained(model, use_fast=True) + self.config.eos = self.config.hf_config.eos_token_id + self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + + # vLLM configuration and instance + os.environ.setdefault("VLLM_USE_V1", "0") + self.model = LLM(model=model, enforce_eager=True, dtype="bfloat16", max_model_len=8192) self.pad_token_id = self.tokenizer.pad_token_id - def generate( - self, - prompt: list[str], - sampling_param: SamplingParams, - past_key_values=None, - ) -> dict: + def generate(self, prompt: List[int], sampling_param: SamplingParams, past_key_values: Optional[object] = None) -> dict: + # ensure the vLLM stop ids are set sampling_param.stop_token_ids = [self.config.hf_config.eos_token_id] - with torch.no_grad(): - generated_ids = self.model.generate( - TokensPrompt(prompt_token_ids=prompt), - VllmSamplingParams(**asdict(sampling_param)), - use_tqdm=False, - )[0].outputs[0].token_ids - output = { - "text": self.tokenizer.decode(generated_ids), - "token_ids": list(generated_ids), - } - return output \ No newline at end of file + + # build vLLM prompt and sampling params then generate + vllm_prompt = TokensPrompt(prompt_token_ids=prompt) + vllm_params = VllmSamplingParams(**asdict(sampling_param)) + + # vLLM returns a sequence-like object; extract token ids from the first result + result = self.model.generate(vllm_prompt, vllm_params, use_tqdm=False) + generated_ids = result[0].outputs[0].token_ids + + return {"text": self.tokenizer.decode(list(generated_ids), skip_special_tokens=True), "token_ids": list(generated_ids)} diff --git a/soulxpodcast/models/modules/flow.py b/soulxpodcast/models/modules/flow.py index d54397c..cab92c2 100644 --- a/soulxpodcast/models/modules/flow.py +++ b/soulxpodcast/models/modules/flow.py @@ -1,197 +1,197 @@ -from dataclasses import dataclass - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from soulxpodcast.models.modules.flow_components.estimator import \ - CausalConditionalDecoder -from soulxpodcast.models.modules.flow_components.upsample_encoder import ( - UpsampleConformerEncoder, make_pad_mask) - - -@dataclass -class CfmParams: - sigma_min: float = 1e-6 - solver: str = "euler" - t_scheduler: str = "cosine" - training_cfg_rate: float = 0.2 - inference_cfg_rate: float = 0.7 - - -class CausalConditionalCFM(torch.nn.Module): - def __init__(self, in_channels=320, cfm_params=CfmParams(), n_spks=1, spk_emb_dim=80, estimator: torch.nn.Module = None): - super().__init__() - self.n_feats = in_channels - self.n_spks = n_spks - self.spk_emb_dim = spk_emb_dim - self.solver = cfm_params.solver - if hasattr(cfm_params, "sigma_min"): - self.sigma_min = cfm_params.sigma_min - else: - self.sigma_min = 1e-4 - self.t_scheduler = cfm_params.t_scheduler - self.training_cfg_rate = cfm_params.training_cfg_rate - self.inference_cfg_rate = cfm_params.inference_cfg_rate - in_channels = in_channels + (spk_emb_dim if n_spks > 0 else 0) - # Just change the architecture of the estimator here - self.estimator = CausalConditionalDecoder() if estimator is None else estimator - - @torch.inference_mode() - def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, streaming=False): - """Forward diffusion - - Args: - mu (torch.Tensor): output of encoder - shape: (batch_size, n_feats, mel_timesteps) - mask (torch.Tensor): output_mask - shape: (batch_size, 1, mel_timesteps) - n_timesteps (int): number of diffusion steps - temperature (float, optional): temperature for scaling noise. Defaults to 1.0. - spks (torch.Tensor, optional): speaker ids. Defaults to None. - shape: (batch_size, spk_emb_dim) - cond: Not used but kept for future purposes - - Returns: - sample: generated mel-spectrogram - shape: (batch_size, n_feats, mel_timesteps) - """ - z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature - # fix prompt and overlap part mu and z - t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) - if self.t_scheduler == 'cosine': - t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) - return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond, streaming=streaming), None - - def solve_euler(self, x, t_span, mu, mask, spks, cond, streaming=False): - """ - Fixed euler solver for ODEs. - Args: - x (torch.Tensor): random noise - t_span (torch.Tensor): n_timesteps interpolated - shape: (n_timesteps + 1,) - mu (torch.Tensor): output of encoder - shape: (batch_size, n_feats, mel_timesteps) - mask (torch.Tensor): output_mask - shape: (batch_size, 1, mel_timesteps) - spks (torch.Tensor, optional): speaker ids. Defaults to None. - shape: (batch_size, spk_emb_dim) - cond: Not used but kept for future purposes - """ - batch_size = x.size(0) - t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0] - - # I am storing this because I can later plot it by putting a debugger here and saving it to a file - # Or in future might add like a return_all_steps flag - sol = [] - - # Do not use concat, it may cause memory format changed and trt infer with wrong results! - # Create tensors with double batch size for CFG (conditional + unconditional) - x_in = torch.zeros([batch_size * 2, x.size(1), x.size(2)], device=x.device, dtype=x.dtype) - mask_in = torch.zeros([batch_size * 2, mask.size(1), mask.size(2)], device=x.device, dtype=x.dtype) - mu_in = torch.zeros([batch_size * 2, mu.size(1), mu.size(2)], device=x.device, dtype=x.dtype) - t_in = torch.zeros([batch_size * 2], device=x.device, dtype=x.dtype) - spks_in = torch.zeros([batch_size * 2, spks.size(1)], device=x.device, dtype=x.dtype) - cond_in = torch.zeros([batch_size * 2, cond.size(1), cond.size(2)], device=x.device, dtype=x.dtype) - - for step in range(1, len(t_span)): - # Classifier-Free Guidance inference introduced in VoiceBox - # Copy conditional and unconditional input - x_in[:batch_size] = x - x_in[batch_size:] = x - mask_in[:batch_size] = mask - mask_in[batch_size:] = mask - mu_in[:batch_size] = mu - # Unconditional part remains 0 - t_in.fill_(t) - spks_in[:batch_size] = spks - cond_in[:batch_size] = cond - - dphi_dt = self.estimator( - x_in, mask_in, - mu_in, t_in, - spks_in, - cond_in, - streaming - ) - dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [batch_size, batch_size], dim=0) - dphi_dt = ((1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt) - x = x + dt * dphi_dt - t = t + dt - sol.append(x) - if step < len(t_span) - 1: - dt = t_span[step + 1] - t - - return sol[-1].float() - - -class CausalMaskedDiffWithXvec(torch.nn.Module): - def __init__( - self, - input_size: int = 512, - output_size: int = 80, - spk_embed_dim: int = 192, - output_type: str = "mel", - vocab_size: int = 6561, - input_frame_rate: int = 25, - token_mel_ratio: int = 2, - pre_lookahead_len: int = 3, - encoder: torch.nn.Module = None, - decoder: torch.nn.Module = None, - ): - super().__init__() - self.input_size = input_size - self.output_size = output_size - self.vocab_size = vocab_size - self.output_type = output_type - self.input_frame_rate = input_frame_rate - self.input_embedding = nn.Embedding(vocab_size, input_size) - self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size) - self.encoder = UpsampleConformerEncoder() if encoder is None else encoder - self.encoder_proj = torch.nn.Linear(self.encoder.output_size(), output_size) - self.decoder = CausalConditionalCFM() if decoder is None else decoder - self.token_mel_ratio = token_mel_ratio - self.pre_lookahead_len = pre_lookahead_len - - @torch.inference_mode() - def forward(self, - token, - token_len, - prompt_feat, - prompt_feat_len, - embedding, - streaming, - finalize): - # xvec projection - embedding = F.normalize(embedding, dim=1) - embedding = self.spk_embed_affine_layer(embedding) - - # concat text and prompt_text - mask = (~make_pad_mask(token_len, max_len=token.shape[1])).unsqueeze(-1).to(embedding) - token = self.input_embedding(torch.clamp(token, min=0)) * mask - - # text encode - if finalize is True: - h, h_lengths = self.encoder(token, token_len, streaming=streaming) - else: - token, context = token[:, :-self.pre_lookahead_len], token[:, -self.pre_lookahead_len:] - h, h_lengths = self.encoder(token, token_len, context=context, streaming=streaming) - h = self.encoder_proj(h) - - # get conditions - conds = torch.zeros_like(h, device=token.device) - for i, j in enumerate(prompt_feat_len): - conds[i, :j] = prompt_feat[i, :j] - conds = conds.transpose(1, 2) - - h_lengths = h_lengths.sum(dim=-1).squeeze(dim=1) - mask = (~make_pad_mask(h_lengths, max_len=h.shape[1])).to(h) - feat, _ = self.decoder( - mu=h.transpose(1, 2).contiguous(), - mask=mask.unsqueeze(1), - spks=embedding, - cond=conds, - n_timesteps=15, - streaming=streaming - ) # [B, num_mels, T] - return feat.float(), h_lengths +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from soulxpodcast.models.modules.flow_components.estimator import \ + CausalConditionalDecoder +from soulxpodcast.models.modules.flow_components.upsample_encoder import ( + UpsampleConformerEncoder, make_pad_mask) + + +@dataclass +class CfmParams: + sigma_min: float = 1e-6 + solver: str = "euler" + t_scheduler: str = "cosine" + training_cfg_rate: float = 0.2 + inference_cfg_rate: float = 0.7 + + +class CausalConditionalCFM(torch.nn.Module): + def __init__(self, in_channels=320, cfm_params=CfmParams(), n_spks=1, spk_emb_dim=80, estimator: torch.nn.Module = None): + super().__init__() + self.n_feats = in_channels + self.n_spks = n_spks + self.spk_emb_dim = spk_emb_dim + self.solver = cfm_params.solver + if hasattr(cfm_params, "sigma_min"): + self.sigma_min = cfm_params.sigma_min + else: + self.sigma_min = 1e-4 + self.t_scheduler = cfm_params.t_scheduler + self.training_cfg_rate = cfm_params.training_cfg_rate + self.inference_cfg_rate = cfm_params.inference_cfg_rate + in_channels = in_channels + (spk_emb_dim if n_spks > 0 else 0) + # Just change the architecture of the estimator here + self.estimator = CausalConditionalDecoder() if estimator is None else estimator + + @torch.inference_mode() + def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, streaming=False): + """Forward diffusion + + Args: + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + n_timesteps (int): number of diffusion steps + temperature (float, optional): temperature for scaling noise. Defaults to 1.0. + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond: Not used but kept for future purposes + + Returns: + sample: generated mel-spectrogram + shape: (batch_size, n_feats, mel_timesteps) + """ + z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature + # fix prompt and overlap part mu and z + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if self.t_scheduler == 'cosine': + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond, streaming=streaming), None + + def solve_euler(self, x, t_span, mu, mask, spks, cond, streaming=False): + """ + Fixed euler solver for ODEs. + Args: + x (torch.Tensor): random noise + t_span (torch.Tensor): n_timesteps interpolated + shape: (n_timesteps + 1,) + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond: Not used but kept for future purposes + """ + batch_size = x.size(0) + t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0] + + # I am storing this because I can later plot it by putting a debugger here and saving it to a file + # Or in future might add like a return_all_steps flag + sol = [] + + # Do not use concat, it may cause memory format changed and trt infer with wrong results! + # Create tensors with double batch size for CFG (conditional + unconditional) + x_in = torch.zeros([batch_size * 2, x.size(1), x.size(2)], device=x.device, dtype=x.dtype) + mask_in = torch.zeros([batch_size * 2, mask.size(1), mask.size(2)], device=x.device, dtype=x.dtype) + mu_in = torch.zeros([batch_size * 2, mu.size(1), mu.size(2)], device=x.device, dtype=x.dtype) + t_in = torch.zeros([batch_size * 2], device=x.device, dtype=x.dtype) + spks_in = torch.zeros([batch_size * 2, spks.size(1)], device=x.device, dtype=x.dtype) + cond_in = torch.zeros([batch_size * 2, cond.size(1), cond.size(2)], device=x.device, dtype=x.dtype) + + for step in range(1, len(t_span)): + # Classifier-Free Guidance inference introduced in VoiceBox + # Copy conditional and unconditional input + x_in[:batch_size] = x + x_in[batch_size:] = x + mask_in[:batch_size] = mask + mask_in[batch_size:] = mask + mu_in[:batch_size] = mu + # Unconditional part remains 0 + t_in.fill_(t) + spks_in[:batch_size] = spks + cond_in[:batch_size] = cond + + dphi_dt = self.estimator( + x_in, mask_in, + mu_in, t_in, + spks_in, + cond_in, + streaming + ) + dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [batch_size, batch_size], dim=0) + dphi_dt = ((1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt) + x = x + dt * dphi_dt + t = t + dt + sol.append(x) + if step < len(t_span) - 1: + dt = t_span[step + 1] - t + + return sol[-1].float() + + +class CausalMaskedDiffWithXvec(torch.nn.Module): + def __init__( + self, + input_size: int = 512, + output_size: int = 80, + spk_embed_dim: int = 192, + output_type: str = "mel", + vocab_size: int = 6561, + input_frame_rate: int = 25, + token_mel_ratio: int = 2, + pre_lookahead_len: int = 3, + encoder: torch.nn.Module = None, + decoder: torch.nn.Module = None, + ): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.vocab_size = vocab_size + self.output_type = output_type + self.input_frame_rate = input_frame_rate + self.input_embedding = nn.Embedding(vocab_size, input_size) + self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size) + self.encoder = UpsampleConformerEncoder() if encoder is None else encoder + self.encoder_proj = torch.nn.Linear(self.encoder.output_size(), output_size) + self.decoder = CausalConditionalCFM() if decoder is None else decoder + self.token_mel_ratio = token_mel_ratio + self.pre_lookahead_len = pre_lookahead_len + + @torch.inference_mode() + def forward(self, + token, + token_len, + prompt_feat, + prompt_feat_len, + embedding, + streaming, + finalize): + # xvec projection + embedding = F.normalize(embedding, dim=1) + embedding = self.spk_embed_affine_layer(embedding) + + # concat text and prompt_text + mask = (~make_pad_mask(token_len, max_len=token.shape[1])).unsqueeze(-1).to(embedding) + token = self.input_embedding(torch.clamp(token, min=0)) * mask + + # text encode + if finalize is True: + h, h_lengths = self.encoder(token, token_len, streaming=streaming) + else: + token, context = token[:, :-self.pre_lookahead_len], token[:, -self.pre_lookahead_len:] + h, h_lengths = self.encoder(token, token_len, context=context, streaming=streaming) + h = self.encoder_proj(h) + + # get conditions + conds = torch.zeros_like(h, device=token.device) + for i, j in enumerate(prompt_feat_len): + conds[i, :j] = prompt_feat[i, :j] + conds = conds.transpose(1, 2) + + h_lengths = h_lengths.sum(dim=-1).squeeze(dim=1) + mask = (~make_pad_mask(h_lengths, max_len=h.shape[1])).to(h) + feat, _ = self.decoder( + mu=h.transpose(1, 2).contiguous(), + mask=mask.unsqueeze(1), + spks=embedding, + cond=conds, + n_timesteps=15, + streaming=streaming + ) # [B, num_mels, T] + return feat.float(), h_lengths diff --git a/soulxpodcast/models/modules/flow_components/estimator.py b/soulxpodcast/models/modules/flow_components/estimator.py index 0c965a8..33deb85 100644 --- a/soulxpodcast/models/modules/flow_components/estimator.py +++ b/soulxpodcast/models/modules/flow_components/estimator.py @@ -1,974 +1,974 @@ -import math -from typing import Any, Dict, Optional, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F -from diffusers.models.attention import (GEGLU, GELU, AdaLayerNorm, - AdaLayerNormZero, ApproximateGELU) -from diffusers.models.attention_processor import Attention -from diffusers.models.lora import LoRACompatibleLinear -from diffusers.utils.torch_utils import maybe_allow_in_graph -from einops import pack, rearrange, repeat - -from soulxpodcast.models.modules.flow_components.upsample_encoder import \ - add_optional_chunk_mask - - -def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - assert mask.dtype == torch.bool - assert dtype in [torch.float32, torch.bfloat16, torch.float16] - mask = mask.to(dtype) - # attention mask bias - # NOTE(Mddct): torch.finfo jit issues - # chunk_masks = (1.0 - chunk_masks) * torch.finfo(dtype).min - mask = (1.0 - mask) * -1.0e+10 - return mask - - -class SnakeBeta(nn.Module): - """ - A modified Snake function which uses separate parameters for the magnitude of the periodic components - Shape: - - Input: (B, C, T) - - Output: (B, C, T), same shape as the input - Parameters: - - alpha - trainable parameter that controls frequency - - beta - trainable parameter that controls magnitude - References: - - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: - https://arxiv.org/abs/2006.08195 - Examples: - >>> a1 = snakebeta(256) - >>> x = torch.randn(256) - >>> x = a1(x) - - Args: - in_features: shape of the input - out_features: shape of the output - alpha: trainable parameter that controls frequency - alpha_trainable: whether alpha is trainable - alpha_logscale: whether to use log scale for alpha - alpha is initialized to 1 by default, higher values = higher-frequency. - beta is initialized to 1 by default, higher values = higher-magnitude. - alpha will be trained along with the rest of your model. - """ - - def __init__(self, in_features, out_features, alpha=1.0, alpha_trainable=True, alpha_logscale=True): - super().__init__() - self.in_features = out_features if isinstance(out_features, list) else [out_features] - self.proj = LoRACompatibleLinear(in_features, out_features) - - # initialize alpha - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: # log scale alphas initialized to zeros - self.alpha = nn.Parameter(torch.zeros(self.in_features) * alpha) - self.beta = nn.Parameter(torch.zeros(self.in_features) * alpha) - else: # linear scale alphas initialized to ones - self.alpha = nn.Parameter(torch.ones(self.in_features) * alpha) - self.beta = nn.Parameter(torch.ones(self.in_features) * alpha) - - self.alpha.requires_grad = alpha_trainable - self.beta.requires_grad = alpha_trainable - - self.no_div_by_zero = 0.000000001 - - def forward(self, x): - """ - Forward pass of the function. - Applies the function to the input elementwise. - SnakeBeta ∶= x + 1/b * sin^2 (xa) - """ - x = self.proj(x) - if self.alpha_logscale: - alpha = torch.exp(self.alpha) - beta = torch.exp(self.beta) - else: - alpha = self.alpha - beta = self.beta - - x = x + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(torch.sin(x * alpha), 2) - - return x - - -class FeedForward(nn.Module): - r""" - A feed-forward layer. - - Parameters: - dim (`int`): The number of channels in the input. - dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. - mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. - """ - - def __init__( - self, - dim: int, - dim_out: Optional[int] = None, - mult: int = 4, - dropout: float = 0.0, - activation_fn: str = "geglu", - final_dropout: bool = False, - ): - super().__init__() - inner_dim = int(dim * mult) - dim_out = dim_out if dim_out is not None else dim - - if activation_fn == "gelu": - act_fn = GELU(dim, inner_dim) - if activation_fn == "gelu-approximate": - act_fn = GELU(dim, inner_dim, approximate="tanh") - elif activation_fn == "geglu": - act_fn = GEGLU(dim, inner_dim) - elif activation_fn == "geglu-approximate": - act_fn = ApproximateGELU(dim, inner_dim) - elif activation_fn == "snakebeta": - act_fn = SnakeBeta(dim, inner_dim) - - self.net = nn.ModuleList([]) - # project in - self.net.append(act_fn) - # project dropout - self.net.append(nn.Dropout(dropout)) - # project out - self.net.append(LoRACompatibleLinear(inner_dim, dim_out)) - # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout - if final_dropout: - self.net.append(nn.Dropout(dropout)) - - def forward(self, hidden_states): - for module in self.net: - hidden_states = module(hidden_states) - return hidden_states - - -@maybe_allow_in_graph -class BasicTransformerBlock(nn.Module): - r""" - A basic Transformer block. - - Parameters: - dim (`int`): The number of channels in the input and output. - num_attention_heads (`int`): The number of heads to use for multi-head attention. - attention_head_dim (`int`): The number of channels in each head. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - only_cross_attention (`bool`, *optional*): - Whether to use only cross-attention layers. In this case two cross attention layers are used. - double_self_attention (`bool`, *optional*): - Whether to use two self-attention layers. In this case no cross attention layers are used. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - num_embeds_ada_norm (: - obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. - attention_bias (: - obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - dropout=0.0, - cross_attention_dim: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - attention_bias: bool = False, - only_cross_attention: bool = False, - double_self_attention: bool = False, - upcast_attention: bool = False, - norm_elementwise_affine: bool = True, - norm_type: str = "layer_norm", - final_dropout: bool = False, - ): - super().__init__() - self.only_cross_attention = only_cross_attention - - self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" - self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" - - if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: - raise ValueError( - f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" - f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." - ) - - # Define 3 blocks. Each block has its own normalization layer. - # 1. Self-Attn - if self.use_ada_layer_norm: - self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) - elif self.use_ada_layer_norm_zero: - self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) - else: - self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=cross_attention_dim if only_cross_attention else None, - upcast_attention=upcast_attention, - ) - - # 2. Cross-Attn - if cross_attention_dim is not None or double_self_attention: - # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. - # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during - # the second cross attention block. - self.norm2 = ( - AdaLayerNorm(dim, num_embeds_ada_norm) - if self.use_ada_layer_norm - else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) - ) - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim if not double_self_attention else None, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - upcast_attention=upcast_attention, - # scale_qk=False, # uncomment this to not to use flash attention - ) # is self-attn if encoder_hidden_states is none - else: - self.norm2 = None - self.attn2 = None - - # 3. Feed-forward - self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) - self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) - - # let chunk size default to None - self._chunk_size = None - self._chunk_dim = 0 - - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): - # Sets chunk feed-forward - self._chunk_size = chunk_size - self._chunk_dim = dim - - def forward( - self, - hidden_states: torch.FloatTensor, - attention_mask: Optional[torch.FloatTensor] = None, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - encoder_attention_mask: Optional[torch.FloatTensor] = None, - timestep: Optional[torch.LongTensor] = None, - cross_attention_kwargs: Dict[str, Any] = None, - class_labels: Optional[torch.LongTensor] = None, - ): - # Notice that normalization is always applied before the real computation in the following blocks. - # 1. Self-Attention - if self.use_ada_layer_norm: - norm_hidden_states = self.norm1(hidden_states, timestep) - elif self.use_ada_layer_norm_zero: - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( - hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype - ) - else: - norm_hidden_states = self.norm1(hidden_states) - - cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} - - attn_output = self.attn1( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, - attention_mask=encoder_attention_mask if self.only_cross_attention else attention_mask, - **cross_attention_kwargs, - ) - if self.use_ada_layer_norm_zero: - attn_output = gate_msa.unsqueeze(1) * attn_output - hidden_states = attn_output + hidden_states - - # 2. Cross-Attention - if self.attn2 is not None: - norm_hidden_states = ( - self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states) - ) - - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - **cross_attention_kwargs, - ) - hidden_states = attn_output + hidden_states - - # 3. Feed-forward - norm_hidden_states = self.norm3(hidden_states) - - if self.use_ada_layer_norm_zero: - norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - - if self._chunk_size is not None: - # "feed_forward_chunk_size" can be used to save memory - if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: - raise ValueError( - f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." - ) - - num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size - ff_output = torch.cat( - [self.ff(hid_slice) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)], - dim=self._chunk_dim, - ) - else: - ff_output = self.ff(norm_hidden_states) - - if self.use_ada_layer_norm_zero: - ff_output = gate_mlp.unsqueeze(1) * ff_output - - hidden_states = ff_output + hidden_states - - return hidden_states - - -class SinusoidalPosEmb(torch.nn.Module): - def __init__(self, dim): - super().__init__() - self.dim = dim - assert self.dim % 2 == 0, "SinusoidalPosEmb requires dim to be even" - - def forward(self, x, scale=1000): - if x.ndim < 1: - x = x.unsqueeze(0) - device = x.device - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) - emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb - - -class Block1D(torch.nn.Module): - def __init__(self, dim, dim_out, groups=8): - super().__init__() - self.block = torch.nn.Sequential( - torch.nn.Conv1d(dim, dim_out, 3, padding=1), - torch.nn.GroupNorm(groups, dim_out), - nn.Mish(), - ) - - def forward(self, x, mask): - output = self.block(x * mask) - return output * mask - - -class ResnetBlock1D(torch.nn.Module): - def __init__(self, dim, dim_out, time_emb_dim, groups=8): - super().__init__() - self.mlp = torch.nn.Sequential(nn.Mish(), torch.nn.Linear(time_emb_dim, dim_out)) - - self.block1 = Block1D(dim, dim_out, groups=groups) - self.block2 = Block1D(dim_out, dim_out, groups=groups) - - self.res_conv = torch.nn.Conv1d(dim, dim_out, 1) - - def forward(self, x, mask, time_emb): - h = self.block1(x, mask) - h += self.mlp(time_emb).unsqueeze(-1) - h = self.block2(h, mask) - output = h + self.res_conv(x * mask) - return output - - -class Downsample1D(nn.Module): - def __init__(self, dim): - super().__init__() - self.conv = torch.nn.Conv1d(dim, dim, 3, 2, 1) - - def forward(self, x): - return self.conv(x) - - -class TimestepEmbedding(nn.Module): - def __init__( - self, - in_channels: int, - time_embed_dim: int, - act_fn: str = "silu", - out_dim: int = None, - post_act_fn: Optional[str] = None, - cond_proj_dim=None, - ): - super().__init__() - assert act_fn == "silu", "act_fn must be silu" - - self.linear_1 = nn.Linear(in_channels, time_embed_dim) - - if cond_proj_dim is not None: - self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) - else: - self.cond_proj = None - - self.act = nn.SiLU() - - if out_dim is not None: - time_embed_dim_out = out_dim - else: - time_embed_dim_out = time_embed_dim - self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out) - - if post_act_fn is None: - self.post_act = None - else: - self.post_act = nn.SiLU() - - def forward(self, sample, condition=None): - if condition is not None: - sample = sample + self.cond_proj(condition) - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - - if self.post_act is not None: - sample = self.post_act(sample) - return sample - - -class Upsample1D(nn.Module): - """A 1D upsampling layer with an optional convolution. - - Parameters: - channels (`int`): - number of channels in the inputs and outputs. - use_conv (`bool`, default `False`): - option to use a convolution. - use_conv_transpose (`bool`, default `False`): - option to use a convolution transpose. - out_channels (`int`, optional): - number of output channels. Defaults to `channels`. - """ - - def __init__(self, channels, use_conv=False, use_conv_transpose=True, out_channels=None, name="conv"): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.use_conv_transpose = use_conv_transpose - self.name = name - - self.conv = None - if use_conv_transpose: - self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1) - elif use_conv: - self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1) - - def forward(self, inputs): - assert inputs.shape[1] == self.channels - if self.use_conv_transpose: - return self.conv(inputs) - - outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest") - - if self.use_conv: - outputs = self.conv(outputs) - - return outputs - - -class Transpose(torch.nn.Module): - def __init__(self, dim0: int, dim1: int): - super().__init__() - self.dim0 = dim0 - self.dim1 = dim1 - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = torch.transpose(x, self.dim0, self.dim1) - return x - - -class CausalConv1d(torch.nn.Conv1d): - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - stride: int = 1, - dilation: int = 1, - groups: int = 1, - bias: bool = True, - padding_mode: str = 'zeros', - device=None, - dtype=None - ) -> None: - super(CausalConv1d, self).__init__(in_channels, out_channels, - kernel_size, stride, - padding=0, dilation=dilation, - groups=groups, bias=bias, - padding_mode=padding_mode, - device=device, dtype=dtype) - assert stride == 1 - self.causal_padding = kernel_size - 1 - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = F.pad(x, (self.causal_padding, 0), value=0.0) - x = super(CausalConv1d, self).forward(x) - return x - - -class CausalBlock1D(Block1D): - def __init__(self, dim: int, dim_out: int): - super(CausalBlock1D, self).__init__(dim, dim_out) - self.block = torch.nn.Sequential( - CausalConv1d(dim, dim_out, 3), - Transpose(1, 2), - nn.LayerNorm(dim_out), - Transpose(1, 2), - nn.Mish(), - ) - - def forward(self, x: torch.Tensor, mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - output = self.block(x * mask) - return output * mask - - -class CausalResnetBlock1D(ResnetBlock1D): - def __init__(self, dim: int, dim_out: int, time_emb_dim: int, groups: int = 8): - super(CausalResnetBlock1D, self).__init__(dim, dim_out, time_emb_dim, groups) - self.block1 = CausalBlock1D(dim, dim_out) - self.block2 = CausalBlock1D(dim_out, dim_out) - - -class ConditionalDecoder(nn.Module): - """ - This decoder requires an input with the same shape of the target. So, if your text content - is shorter or longer than the outputs, please re-sampling it before feeding to the decoder. - - Args: - in_channels: number of input channels - out_channels: number of output channels - channels: tuple of channel dimensions - dropout: dropout rate - attention_head_dim: dimension of attention heads - n_blocks: number of transformer blocks - num_mid_blocks: number of middle blocks - num_heads: number of attention heads - act_fn: activation function name - """ - - def __init__( - self, - in_channels, - out_channels, - channels=(256, 256), - dropout=0.05, - attention_head_dim=64, - n_blocks=1, - num_mid_blocks=2, - num_heads=4, - act_fn="snake", - ): - super().__init__() - channels = tuple(channels) - self.in_channels = in_channels - self.out_channels = out_channels - - self.time_embeddings = SinusoidalPosEmb(in_channels) - time_embed_dim = channels[0] * 4 - self.time_mlp = TimestepEmbedding( - in_channels=in_channels, - time_embed_dim=time_embed_dim, - act_fn="silu", - ) - self.down_blocks = nn.ModuleList([]) - self.mid_blocks = nn.ModuleList([]) - self.up_blocks = nn.ModuleList([]) - - output_channel = in_channels - for i in range(len(channels)): # pylint: disable=consider-using-enumerate - input_channel = output_channel - output_channel = channels[i] - is_last = i == len(channels) - 1 - resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) - transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - dim=output_channel, - num_attention_heads=num_heads, - attention_head_dim=attention_head_dim, - dropout=dropout, - activation_fn=act_fn, - ) - for _ in range(n_blocks) - ] - ) - downsample = ( - Downsample1D(output_channel) if not is_last else nn.Conv1d(output_channel, output_channel, 3, padding=1) - ) - self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample])) - - for _ in range(num_mid_blocks): - input_channel = channels[-1] - out_channels = channels[-1] - resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) - - transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - dim=output_channel, - num_attention_heads=num_heads, - attention_head_dim=attention_head_dim, - dropout=dropout, - activation_fn=act_fn, - ) - for _ in range(n_blocks) - ] - ) - - self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks])) - - channels = channels[::-1] + (channels[0],) - for i in range(len(channels) - 1): - input_channel = channels[i] * 2 - output_channel = channels[i + 1] - is_last = i == len(channels) - 2 - resnet = ResnetBlock1D( - dim=input_channel, - dim_out=output_channel, - time_emb_dim=time_embed_dim, - ) - transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - dim=output_channel, - num_attention_heads=num_heads, - attention_head_dim=attention_head_dim, - dropout=dropout, - activation_fn=act_fn, - ) - for _ in range(n_blocks) - ] - ) - upsample = ( - Upsample1D(output_channel, use_conv_transpose=True) - if not is_last - else nn.Conv1d(output_channel, output_channel, 3, padding=1) - ) - self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample])) - self.final_block = Block1D(channels[-1], channels[-1]) - self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1) - self.initialize_weights() - - def initialize_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv1d): - nn.init.kaiming_normal_(m.weight, nonlinearity="relu") - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.GroupNorm): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.kaiming_normal_(m.weight, nonlinearity="relu") - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - def forward(self, x, mask, mu, t, spks=None, cond=None, streaming=False): - """Forward pass of the UNet1DConditional model. - - Args: - x (torch.Tensor): shape (batch_size, in_channels, time) - mask (_type_): shape (batch_size, 1, time) - t (_type_): shape (batch_size) - spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None. - cond (_type_, optional): placeholder for future use. Defaults to None. - - Raises: - ValueError: _description_ - ValueError: _description_ - - Returns: - _type_: _description_ - """ - - t = self.time_embeddings(t).to(t.dtype) - t = self.time_mlp(t) - - x = pack([x, mu], "b * t")[0] - - if spks is not None: - spks = repeat(spks, "b c -> b c t", t=x.shape[-1]) - x = pack([x, spks], "b * t")[0] - if cond is not None: - x = pack([x, cond], "b * t")[0] - - hiddens = [] - masks = [mask] - for resnet, transformer_blocks, downsample in self.down_blocks: - mask_down = masks[-1] - x = resnet(x, mask_down, t) - x = rearrange(x, "b c t -> b t c").contiguous() - attn_mask = add_optional_chunk_mask(x, mask_down.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) - attn_mask = mask_to_bias(attn_mask, x.dtype) - for transformer_block in transformer_blocks: - x = transformer_block( - hidden_states=x, - attention_mask=attn_mask, - timestep=t, - ) - x = rearrange(x, "b t c -> b c t").contiguous() - hiddens.append(x) # Save hidden states for skip connections - x = downsample(x * mask_down) - masks.append(mask_down[:, :, ::2]) - masks = masks[:-1] - mask_mid = masks[-1] - - for resnet, transformer_blocks in self.mid_blocks: - x = resnet(x, mask_mid, t) - x = rearrange(x, "b c t -> b t c").contiguous() - attn_mask = add_optional_chunk_mask(x, mask_mid.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) - attn_mask = mask_to_bias(attn_mask, x.dtype) - for transformer_block in transformer_blocks: - x = transformer_block( - hidden_states=x, - attention_mask=attn_mask, - timestep=t, - ) - x = rearrange(x, "b t c -> b c t").contiguous() - - for resnet, transformer_blocks, upsample in self.up_blocks: - mask_up = masks.pop() - skip = hiddens.pop() - x = pack([x[:, :, :skip.shape[-1]], skip], "b * t")[0] - x = resnet(x, mask_up, t) - x = rearrange(x, "b c t -> b t c").contiguous() - attn_mask = add_optional_chunk_mask(x, mask_up.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) - attn_mask = mask_to_bias(attn_mask, x.dtype) - for transformer_block in transformer_blocks: - x = transformer_block( - hidden_states=x, - attention_mask=attn_mask, - timestep=t, - ) - x = rearrange(x, "b t c -> b c t").contiguous() - x = upsample(x * mask_up) - x = self.final_block(x, mask_up) - output = self.final_proj(x * mask_up) - return output * mask - - -class CausalConditionalDecoder(ConditionalDecoder): - """ - This decoder requires an input with the same shape of the target. So, if your text content - is shorter or longer than the outputs, please re-sampling it before feeding to the decoder. - - Args: - in_channels: number of input channels - out_channels: number of output channels - channels: list of channel dimensions - dropout: dropout rate - attention_head_dim: dimension of attention heads - n_blocks: number of transformer blocks - num_mid_blocks: number of middle blocks - num_heads: number of attention heads - act_fn: activation function name - static_chunk_size: size of static chunks - num_decoding_left_chunks: number of left chunks for decoding - """ - - def __init__( - self, - in_channels=320, - out_channels=80, - channels=[256], # noqa - dropout=0.0, - attention_head_dim=64, - n_blocks=4, - num_mid_blocks=12, - num_heads=8, - act_fn="gelu", - static_chunk_size=50, - num_decoding_left_chunks=-1, - ): - torch.nn.Module.__init__(self) - channels = tuple(channels) - self.in_channels = in_channels - self.out_channels = out_channels - self.time_embeddings = SinusoidalPosEmb(in_channels) - time_embed_dim = channels[0] * 4 - self.time_mlp = TimestepEmbedding( - in_channels=in_channels, - time_embed_dim=time_embed_dim, - act_fn="silu", - ) - self.static_chunk_size = static_chunk_size - self.num_decoding_left_chunks = num_decoding_left_chunks - self.down_blocks = nn.ModuleList([]) - self.mid_blocks = nn.ModuleList([]) - self.up_blocks = nn.ModuleList([]) - - output_channel = in_channels - for i in range(len(channels)): # pylint: disable=consider-using-enumerate - input_channel = output_channel - output_channel = channels[i] - is_last = i == len(channels) - 1 - resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) - transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - dim=output_channel, - num_attention_heads=num_heads, - attention_head_dim=attention_head_dim, - dropout=dropout, - activation_fn=act_fn, - ) - for _ in range(n_blocks) - ] - ) - downsample = ( - Downsample1D(output_channel) if not is_last else CausalConv1d(output_channel, output_channel, 3) - ) - self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample])) - - for _ in range(num_mid_blocks): - input_channel = channels[-1] - out_channels = channels[-1] - resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) - - transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - dim=output_channel, - num_attention_heads=num_heads, - attention_head_dim=attention_head_dim, - dropout=dropout, - activation_fn=act_fn, - ) - for _ in range(n_blocks) - ] - ) - - self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks])) - - channels = channels[::-1] + (channels[0],) - for i in range(len(channels) - 1): - input_channel = channels[i] * 2 - output_channel = channels[i + 1] - is_last = i == len(channels) - 2 - resnet = CausalResnetBlock1D( - dim=input_channel, - dim_out=output_channel, - time_emb_dim=time_embed_dim, - ) - transformer_blocks = nn.ModuleList( - [ - BasicTransformerBlock( - dim=output_channel, - num_attention_heads=num_heads, - attention_head_dim=attention_head_dim, - dropout=dropout, - activation_fn=act_fn, - ) - for _ in range(n_blocks) - ] - ) - upsample = ( - Upsample1D(output_channel, use_conv_transpose=True) - if not is_last - else CausalConv1d(output_channel, output_channel, 3) - ) - self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample])) - self.final_block = CausalBlock1D(channels[-1], channels[-1]) - self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1) - self.initialize_weights() - - def forward(self, x, mask, mu, t, spks=None, cond=None, streaming=False): - """Forward pass of the UNet1DConditional model. - - Args: - x (torch.Tensor): shape (batch_size, in_channels, time) - mask (_type_): shape (batch_size, 1, time) - t (_type_): shape (batch_size) - spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None. - cond (_type_, optional): placeholder for future use. Defaults to None. - - Raises: - ValueError: _description_ - ValueError: _description_ - - Returns: - _type_: _description_ - """ - t = self.time_embeddings(t).to(t.dtype) - t = self.time_mlp(t) - - x = pack([x, mu], "b * t")[0] - - if spks is not None: - spks = repeat(spks, "b c -> b c t", t=x.shape[-1]) - x = pack([x, spks], "b * t")[0] - if cond is not None: - x = pack([x, cond], "b * t")[0] - - hiddens = [] - masks = [mask] - for resnet, transformer_blocks, downsample in self.down_blocks: - mask_down = masks[-1] - x = resnet(x, mask_down, t) - x = rearrange(x, "b c t -> b t c").contiguous() - if streaming is True: - attn_mask = add_optional_chunk_mask(x, mask_down.bool(), False, False, 0, self.static_chunk_size, -1) - else: - attn_mask = add_optional_chunk_mask(x, mask_down.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) - attn_mask = mask_to_bias(attn_mask, x.dtype) - for transformer_block in transformer_blocks: - x = transformer_block( - hidden_states=x, - attention_mask=attn_mask, - timestep=t, - ) - x = rearrange(x, "b t c -> b c t").contiguous() - hiddens.append(x) # Save hidden states for skip connections - x = downsample(x * mask_down) - masks.append(mask_down[:, :, ::2]) - masks = masks[:-1] - mask_mid = masks[-1] - - for resnet, transformer_blocks in self.mid_blocks: - x = resnet(x, mask_mid, t) - x = rearrange(x, "b c t -> b t c").contiguous() - if streaming is True: - attn_mask = add_optional_chunk_mask(x, mask_mid.bool(), False, False, 0, self.static_chunk_size, -1) - else: - attn_mask = add_optional_chunk_mask(x, mask_mid.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) - attn_mask = mask_to_bias(attn_mask, x.dtype) - for transformer_block in transformer_blocks: - x = transformer_block( - hidden_states=x, - attention_mask=attn_mask, - timestep=t, - ) - x = rearrange(x, "b t c -> b c t").contiguous() - - for resnet, transformer_blocks, upsample in self.up_blocks: - mask_up = masks.pop() - skip = hiddens.pop() - x = pack([x[:, :, :skip.shape[-1]], skip], "b * t")[0] - x = resnet(x, mask_up, t) - x = rearrange(x, "b c t -> b t c").contiguous() - if streaming is True: - attn_mask = add_optional_chunk_mask(x, mask_up.bool(), False, False, 0, self.static_chunk_size, -1) - else: - attn_mask = add_optional_chunk_mask(x, mask_up.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) - attn_mask = mask_to_bias(attn_mask, x.dtype) - for transformer_block in transformer_blocks: - x = transformer_block( - hidden_states=x, - attention_mask=attn_mask, - timestep=t, - ) - x = rearrange(x, "b t c -> b c t").contiguous() - x = upsample(x * mask_up) - x = self.final_block(x, mask_up) - output = self.final_proj(x * mask_up) - return output * mask +import math +from typing import Any, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.attention import (GEGLU, GELU, AdaLayerNorm, + AdaLayerNormZero, ApproximateGELU) +from diffusers.models.attention_processor import Attention +from diffusers.models.lora import LoRACompatibleLinear +from diffusers.utils.torch_utils import maybe_allow_in_graph +from einops import pack, rearrange, repeat + +from soulxpodcast.models.modules.flow_components.upsample_encoder import \ + add_optional_chunk_mask + + +def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + assert mask.dtype == torch.bool + assert dtype in [torch.float32, torch.bfloat16, torch.float16] + mask = mask.to(dtype) + # attention mask bias + # NOTE(Mddct): torch.finfo jit issues + # chunk_masks = (1.0 - chunk_masks) * torch.finfo(dtype).min + mask = (1.0 - mask) * -1.0e+10 + return mask + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude of the periodic components + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + References: + - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snakebeta(256) + >>> x = torch.randn(256) + >>> x = a1(x) + + Args: + in_features: shape of the input + out_features: shape of the output + alpha: trainable parameter that controls frequency + alpha_trainable: whether alpha is trainable + alpha_logscale: whether to use log scale for alpha + alpha is initialized to 1 by default, higher values = higher-frequency. + beta is initialized to 1 by default, higher values = higher-magnitude. + alpha will be trained along with the rest of your model. + """ + + def __init__(self, in_features, out_features, alpha=1.0, alpha_trainable=True, alpha_logscale=True): + super().__init__() + self.in_features = out_features if isinstance(out_features, list) else [out_features] + self.proj = LoRACompatibleLinear(in_features, out_features) + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = nn.Parameter(torch.zeros(self.in_features) * alpha) + self.beta = nn.Parameter(torch.zeros(self.in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = nn.Parameter(torch.ones(self.in_features) * alpha) + self.beta = nn.Parameter(torch.ones(self.in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + SnakeBeta ∶= x + 1/b * sin^2 (xa) + """ + x = self.proj(x) + if self.alpha_logscale: + alpha = torch.exp(self.alpha) + beta = torch.exp(self.beta) + else: + alpha = self.alpha + beta = self.beta + + x = x + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(torch.sin(x * alpha), 2) + + return x + + +class FeedForward(nn.Module): + r""" + A feed-forward layer. + + Parameters: + dim (`int`): The number of channels in the input. + dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. + mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. + """ + + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + dropout: float = 0.0, + activation_fn: str = "geglu", + final_dropout: bool = False, + ): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + if activation_fn == "gelu": + act_fn = GELU(dim, inner_dim) + if activation_fn == "gelu-approximate": + act_fn = GELU(dim, inner_dim, approximate="tanh") + elif activation_fn == "geglu": + act_fn = GEGLU(dim, inner_dim) + elif activation_fn == "geglu-approximate": + act_fn = ApproximateGELU(dim, inner_dim) + elif activation_fn == "snakebeta": + act_fn = SnakeBeta(dim, inner_dim) + + self.net = nn.ModuleList([]) + # project in + self.net.append(act_fn) + # project dropout + self.net.append(nn.Dropout(dropout)) + # project out + self.net.append(LoRACompatibleLinear(inner_dim, dim_out)) + # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout + if final_dropout: + self.net.append(nn.Dropout(dropout)) + + def forward(self, hidden_states): + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +@maybe_allow_in_graph +class BasicTransformerBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", + final_dropout: bool = False, + ): + super().__init__() + self.only_cross_attention = only_cross_attention + + self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" + self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" + + if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: + raise ValueError( + f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" + f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." + ) + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if self.use_ada_layer_norm: + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif self.use_ada_layer_norm_zero: + self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) + else: + self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + + # 2. Cross-Attn + if cross_attention_dim is not None or double_self_attention: + # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. + # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during + # the second cross attention block. + self.norm2 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + ) + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + # scale_qk=False, # uncomment this to not to use flash attention + ) # is self-attn if encoder_hidden_states is none + else: + self.norm2 = None + self.attn2 = None + + # 3. Feed-forward + self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + ): + # Notice that normalization is always applied before the real computation in the following blocks. + # 1. Self-Attention + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + else: + norm_hidden_states = self.norm1(hidden_states) + + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, + attention_mask=encoder_attention_mask if self.only_cross_attention else attention_mask, + **cross_attention_kwargs, + ) + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = attn_output + hidden_states + + # 2. Cross-Attention + if self.attn2 is not None: + norm_hidden_states = ( + self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states) + ) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 3. Feed-forward + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + + if self._chunk_size is not None: + # "feed_forward_chunk_size" can be used to save memory + if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: + raise ValueError( + f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." + ) + + num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size + ff_output = torch.cat( + [self.ff(hid_slice) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)], + dim=self._chunk_dim, + ) + else: + ff_output = self.ff(norm_hidden_states) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + + hidden_states = ff_output + hidden_states + + return hidden_states + + +class SinusoidalPosEmb(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + assert self.dim % 2 == 0, "SinusoidalPosEmb requires dim to be even" + + def forward(self, x, scale=1000): + if x.ndim < 1: + x = x.unsqueeze(0) + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class Block1D(torch.nn.Module): + def __init__(self, dim, dim_out, groups=8): + super().__init__() + self.block = torch.nn.Sequential( + torch.nn.Conv1d(dim, dim_out, 3, padding=1), + torch.nn.GroupNorm(groups, dim_out), + nn.Mish(), + ) + + def forward(self, x, mask): + output = self.block(x * mask) + return output * mask + + +class ResnetBlock1D(torch.nn.Module): + def __init__(self, dim, dim_out, time_emb_dim, groups=8): + super().__init__() + self.mlp = torch.nn.Sequential(nn.Mish(), torch.nn.Linear(time_emb_dim, dim_out)) + + self.block1 = Block1D(dim, dim_out, groups=groups) + self.block2 = Block1D(dim_out, dim_out, groups=groups) + + self.res_conv = torch.nn.Conv1d(dim, dim_out, 1) + + def forward(self, x, mask, time_emb): + h = self.block1(x, mask) + h += self.mlp(time_emb).unsqueeze(-1) + h = self.block2(h, mask) + output = h + self.res_conv(x * mask) + return output + + +class Downsample1D(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = torch.nn.Conv1d(dim, dim, 3, 2, 1) + + def forward(self, x): + return self.conv(x) + + +class TimestepEmbedding(nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int = None, + post_act_fn: Optional[str] = None, + cond_proj_dim=None, + ): + super().__init__() + assert act_fn == "silu", "act_fn must be silu" + + self.linear_1 = nn.Linear(in_channels, time_embed_dim) + + if cond_proj_dim is not None: + self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) + else: + self.cond_proj = None + + self.act = nn.SiLU() + + if out_dim is not None: + time_embed_dim_out = out_dim + else: + time_embed_dim_out = time_embed_dim + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out) + + if post_act_fn is None: + self.post_act = None + else: + self.post_act = nn.SiLU() + + def forward(self, sample, condition=None): + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class Upsample1D(nn.Module): + """A 1D upsampling layer with an optional convolution. + + Parameters: + channels (`int`): + number of channels in the inputs and outputs. + use_conv (`bool`, default `False`): + option to use a convolution. + use_conv_transpose (`bool`, default `False`): + option to use a convolution transpose. + out_channels (`int`, optional): + number of output channels. Defaults to `channels`. + """ + + def __init__(self, channels, use_conv=False, use_conv_transpose=True, out_channels=None, name="conv"): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + + self.conv = None + if use_conv_transpose: + self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1) + elif use_conv: + self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1) + + def forward(self, inputs): + assert inputs.shape[1] == self.channels + if self.use_conv_transpose: + return self.conv(inputs) + + outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest") + + if self.use_conv: + outputs = self.conv(outputs) + + return outputs + + +class Transpose(torch.nn.Module): + def __init__(self, dim0: int, dim1: int): + super().__init__() + self.dim0 = dim0 + self.dim1 = dim1 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = torch.transpose(x, self.dim0, self.dim1) + return x + + +class CausalConv1d(torch.nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = 'zeros', + device=None, + dtype=None + ) -> None: + super(CausalConv1d, self).__init__(in_channels, out_channels, + kernel_size, stride, + padding=0, dilation=dilation, + groups=groups, bias=bias, + padding_mode=padding_mode, + device=device, dtype=dtype) + assert stride == 1 + self.causal_padding = kernel_size - 1 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, (self.causal_padding, 0), value=0.0) + x = super(CausalConv1d, self).forward(x) + return x + + +class CausalBlock1D(Block1D): + def __init__(self, dim: int, dim_out: int): + super(CausalBlock1D, self).__init__(dim, dim_out) + self.block = torch.nn.Sequential( + CausalConv1d(dim, dim_out, 3), + Transpose(1, 2), + nn.LayerNorm(dim_out), + Transpose(1, 2), + nn.Mish(), + ) + + def forward(self, x: torch.Tensor, mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + output = self.block(x * mask) + return output * mask + + +class CausalResnetBlock1D(ResnetBlock1D): + def __init__(self, dim: int, dim_out: int, time_emb_dim: int, groups: int = 8): + super(CausalResnetBlock1D, self).__init__(dim, dim_out, time_emb_dim, groups) + self.block1 = CausalBlock1D(dim, dim_out) + self.block2 = CausalBlock1D(dim_out, dim_out) + + +class ConditionalDecoder(nn.Module): + """ + This decoder requires an input with the same shape of the target. So, if your text content + is shorter or longer than the outputs, please re-sampling it before feeding to the decoder. + + Args: + in_channels: number of input channels + out_channels: number of output channels + channels: tuple of channel dimensions + dropout: dropout rate + attention_head_dim: dimension of attention heads + n_blocks: number of transformer blocks + num_mid_blocks: number of middle blocks + num_heads: number of attention heads + act_fn: activation function name + """ + + def __init__( + self, + in_channels, + out_channels, + channels=(256, 256), + dropout=0.05, + attention_head_dim=64, + n_blocks=1, + num_mid_blocks=2, + num_heads=4, + act_fn="snake", + ): + super().__init__() + channels = tuple(channels) + self.in_channels = in_channels + self.out_channels = out_channels + + self.time_embeddings = SinusoidalPosEmb(in_channels) + time_embed_dim = channels[0] * 4 + self.time_mlp = TimestepEmbedding( + in_channels=in_channels, + time_embed_dim=time_embed_dim, + act_fn="silu", + ) + self.down_blocks = nn.ModuleList([]) + self.mid_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + output_channel = in_channels + for i in range(len(channels)): # pylint: disable=consider-using-enumerate + input_channel = output_channel + output_channel = channels[i] + is_last = i == len(channels) - 1 + resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + for _ in range(n_blocks) + ] + ) + downsample = ( + Downsample1D(output_channel) if not is_last else nn.Conv1d(output_channel, output_channel, 3, padding=1) + ) + self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample])) + + for _ in range(num_mid_blocks): + input_channel = channels[-1] + out_channels = channels[-1] + resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) + + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + for _ in range(n_blocks) + ] + ) + + self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks])) + + channels = channels[::-1] + (channels[0],) + for i in range(len(channels) - 1): + input_channel = channels[i] * 2 + output_channel = channels[i + 1] + is_last = i == len(channels) - 2 + resnet = ResnetBlock1D( + dim=input_channel, + dim_out=output_channel, + time_emb_dim=time_embed_dim, + ) + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + for _ in range(n_blocks) + ] + ) + upsample = ( + Upsample1D(output_channel, use_conv_transpose=True) + if not is_last + else nn.Conv1d(output_channel, output_channel, 3, padding=1) + ) + self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample])) + self.final_block = Block1D(channels[-1], channels[-1]) + self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1) + self.initialize_weights() + + def initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv1d): + nn.init.kaiming_normal_(m.weight, nonlinearity="relu") + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.GroupNorm): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.kaiming_normal_(m.weight, nonlinearity="relu") + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x, mask, mu, t, spks=None, cond=None, streaming=False): + """Forward pass of the UNet1DConditional model. + + Args: + x (torch.Tensor): shape (batch_size, in_channels, time) + mask (_type_): shape (batch_size, 1, time) + t (_type_): shape (batch_size) + spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None. + cond (_type_, optional): placeholder for future use. Defaults to None. + + Raises: + ValueError: _description_ + ValueError: _description_ + + Returns: + _type_: _description_ + """ + + t = self.time_embeddings(t).to(t.dtype) + t = self.time_mlp(t) + + x = pack([x, mu], "b * t")[0] + + if spks is not None: + spks = repeat(spks, "b c -> b c t", t=x.shape[-1]) + x = pack([x, spks], "b * t")[0] + if cond is not None: + x = pack([x, cond], "b * t")[0] + + hiddens = [] + masks = [mask] + for resnet, transformer_blocks, downsample in self.down_blocks: + mask_down = masks[-1] + x = resnet(x, mask_down, t) + x = rearrange(x, "b c t -> b t c").contiguous() + attn_mask = add_optional_chunk_mask(x, mask_down.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) + attn_mask = mask_to_bias(attn_mask, x.dtype) + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + attention_mask=attn_mask, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t").contiguous() + hiddens.append(x) # Save hidden states for skip connections + x = downsample(x * mask_down) + masks.append(mask_down[:, :, ::2]) + masks = masks[:-1] + mask_mid = masks[-1] + + for resnet, transformer_blocks in self.mid_blocks: + x = resnet(x, mask_mid, t) + x = rearrange(x, "b c t -> b t c").contiguous() + attn_mask = add_optional_chunk_mask(x, mask_mid.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) + attn_mask = mask_to_bias(attn_mask, x.dtype) + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + attention_mask=attn_mask, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t").contiguous() + + for resnet, transformer_blocks, upsample in self.up_blocks: + mask_up = masks.pop() + skip = hiddens.pop() + x = pack([x[:, :, :skip.shape[-1]], skip], "b * t")[0] + x = resnet(x, mask_up, t) + x = rearrange(x, "b c t -> b t c").contiguous() + attn_mask = add_optional_chunk_mask(x, mask_up.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) + attn_mask = mask_to_bias(attn_mask, x.dtype) + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + attention_mask=attn_mask, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t").contiguous() + x = upsample(x * mask_up) + x = self.final_block(x, mask_up) + output = self.final_proj(x * mask_up) + return output * mask + + +class CausalConditionalDecoder(ConditionalDecoder): + """ + This decoder requires an input with the same shape of the target. So, if your text content + is shorter or longer than the outputs, please re-sampling it before feeding to the decoder. + + Args: + in_channels: number of input channels + out_channels: number of output channels + channels: list of channel dimensions + dropout: dropout rate + attention_head_dim: dimension of attention heads + n_blocks: number of transformer blocks + num_mid_blocks: number of middle blocks + num_heads: number of attention heads + act_fn: activation function name + static_chunk_size: size of static chunks + num_decoding_left_chunks: number of left chunks for decoding + """ + + def __init__( + self, + in_channels=320, + out_channels=80, + channels=[256], # noqa + dropout=0.0, + attention_head_dim=64, + n_blocks=4, + num_mid_blocks=12, + num_heads=8, + act_fn="gelu", + static_chunk_size=50, + num_decoding_left_chunks=-1, + ): + torch.nn.Module.__init__(self) + channels = tuple(channels) + self.in_channels = in_channels + self.out_channels = out_channels + self.time_embeddings = SinusoidalPosEmb(in_channels) + time_embed_dim = channels[0] * 4 + self.time_mlp = TimestepEmbedding( + in_channels=in_channels, + time_embed_dim=time_embed_dim, + act_fn="silu", + ) + self.static_chunk_size = static_chunk_size + self.num_decoding_left_chunks = num_decoding_left_chunks + self.down_blocks = nn.ModuleList([]) + self.mid_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + output_channel = in_channels + for i in range(len(channels)): # pylint: disable=consider-using-enumerate + input_channel = output_channel + output_channel = channels[i] + is_last = i == len(channels) - 1 + resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + for _ in range(n_blocks) + ] + ) + downsample = ( + Downsample1D(output_channel) if not is_last else CausalConv1d(output_channel, output_channel, 3) + ) + self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample])) + + for _ in range(num_mid_blocks): + input_channel = channels[-1] + out_channels = channels[-1] + resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) + + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + for _ in range(n_blocks) + ] + ) + + self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks])) + + channels = channels[::-1] + (channels[0],) + for i in range(len(channels) - 1): + input_channel = channels[i] * 2 + output_channel = channels[i + 1] + is_last = i == len(channels) - 2 + resnet = CausalResnetBlock1D( + dim=input_channel, + dim_out=output_channel, + time_emb_dim=time_embed_dim, + ) + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + for _ in range(n_blocks) + ] + ) + upsample = ( + Upsample1D(output_channel, use_conv_transpose=True) + if not is_last + else CausalConv1d(output_channel, output_channel, 3) + ) + self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample])) + self.final_block = CausalBlock1D(channels[-1], channels[-1]) + self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1) + self.initialize_weights() + + def forward(self, x, mask, mu, t, spks=None, cond=None, streaming=False): + """Forward pass of the UNet1DConditional model. + + Args: + x (torch.Tensor): shape (batch_size, in_channels, time) + mask (_type_): shape (batch_size, 1, time) + t (_type_): shape (batch_size) + spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None. + cond (_type_, optional): placeholder for future use. Defaults to None. + + Raises: + ValueError: _description_ + ValueError: _description_ + + Returns: + _type_: _description_ + """ + t = self.time_embeddings(t).to(t.dtype) + t = self.time_mlp(t) + + x = pack([x, mu], "b * t")[0] + + if spks is not None: + spks = repeat(spks, "b c -> b c t", t=x.shape[-1]) + x = pack([x, spks], "b * t")[0] + if cond is not None: + x = pack([x, cond], "b * t")[0] + + hiddens = [] + masks = [mask] + for resnet, transformer_blocks, downsample in self.down_blocks: + mask_down = masks[-1] + x = resnet(x, mask_down, t) + x = rearrange(x, "b c t -> b t c").contiguous() + if streaming is True: + attn_mask = add_optional_chunk_mask(x, mask_down.bool(), False, False, 0, self.static_chunk_size, -1) + else: + attn_mask = add_optional_chunk_mask(x, mask_down.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) + attn_mask = mask_to_bias(attn_mask, x.dtype) + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + attention_mask=attn_mask, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t").contiguous() + hiddens.append(x) # Save hidden states for skip connections + x = downsample(x * mask_down) + masks.append(mask_down[:, :, ::2]) + masks = masks[:-1] + mask_mid = masks[-1] + + for resnet, transformer_blocks in self.mid_blocks: + x = resnet(x, mask_mid, t) + x = rearrange(x, "b c t -> b t c").contiguous() + if streaming is True: + attn_mask = add_optional_chunk_mask(x, mask_mid.bool(), False, False, 0, self.static_chunk_size, -1) + else: + attn_mask = add_optional_chunk_mask(x, mask_mid.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) + attn_mask = mask_to_bias(attn_mask, x.dtype) + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + attention_mask=attn_mask, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t").contiguous() + + for resnet, transformer_blocks, upsample in self.up_blocks: + mask_up = masks.pop() + skip = hiddens.pop() + x = pack([x[:, :, :skip.shape[-1]], skip], "b * t")[0] + x = resnet(x, mask_up, t) + x = rearrange(x, "b c t -> b t c").contiguous() + if streaming is True: + attn_mask = add_optional_chunk_mask(x, mask_up.bool(), False, False, 0, self.static_chunk_size, -1) + else: + attn_mask = add_optional_chunk_mask(x, mask_up.bool(), False, False, 0, 0, -1).repeat(1, x.size(1), 1) + attn_mask = mask_to_bias(attn_mask, x.dtype) + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + attention_mask=attn_mask, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t").contiguous() + x = upsample(x * mask_up) + x = self.final_block(x, mask_up) + output = self.final_proj(x * mask_up) + return output * mask diff --git a/soulxpodcast/models/modules/flow_components/upsample_encoder.py b/soulxpodcast/models/modules/flow_components/upsample_encoder.py index a634cc2..469690e 100644 --- a/soulxpodcast/models/modules/flow_components/upsample_encoder.py +++ b/soulxpodcast/models/modules/flow_components/upsample_encoder.py @@ -1,997 +1,997 @@ -import math -from typing import Optional, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -def subsequent_chunk_mask( - size: int, - chunk_size: int, - num_left_chunks: int = -1, - device: torch.device = torch.device("cpu"), -) -> torch.Tensor: - """Create mask for subsequent steps (size, size) with chunk size, - this is for streaming encoder - - Args: - size (int): size of mask - chunk_size (int): size of chunk - num_left_chunks (int): number of left chunks - <0: use full chunk - >=0: use num_left_chunks - device (torch.device): "cpu" or "cuda" or torch.Tensor.device - - Returns: - torch.Tensor: mask - - Examples: - >>> subsequent_chunk_mask(4, 2) - [[1, 1, 0, 0], - [1, 1, 0, 0], - [1, 1, 1, 1], - [1, 1, 1, 1]] - """ - # NOTE this modified implementation meets onnx export requirements, but it doesn't support num_left_chunks - pos_idx = torch.arange(size, device=device) - block_value = (torch.div(pos_idx, chunk_size, rounding_mode='trunc') + 1) * chunk_size - ret = pos_idx.unsqueeze(0) < block_value.unsqueeze(1) - return ret - - -def add_optional_chunk_mask(xs: torch.Tensor, - masks: torch.Tensor, - use_dynamic_chunk: bool, - use_dynamic_left_chunk: bool, - decoding_chunk_size: int, - static_chunk_size: int, - num_decoding_left_chunks: int, - enable_full_context: bool = True): - """ Apply optional mask for encoder. - - Args: - xs (torch.Tensor): padded input, (B, L, D), L for max length - mask (torch.Tensor): mask for xs, (B, 1, L) - use_dynamic_chunk (bool): whether to use dynamic chunk or not - use_dynamic_left_chunk (bool): whether to use dynamic left chunk for - training. - decoding_chunk_size (int): decoding chunk size for dynamic chunk, it's - 0: default for training, use random dynamic chunk. - <0: for decoding, use full chunk. - >0: for decoding, use fixed chunk size as set. - static_chunk_size (int): chunk size for static chunk training/decoding - if it's greater than 0, if use_dynamic_chunk is true, - this parameter will be ignored - num_decoding_left_chunks: number of left chunks, this is for decoding, - the chunk size is decoding_chunk_size. - >=0: use num_decoding_left_chunks - <0: use all left chunks - enable_full_context (bool): - True: chunk size is either [1, 25] or full context(max_len) - False: chunk size ~ U[1, 25] - - Returns: - torch.Tensor: chunk mask of the input xs. - """ - # Whether to use chunk mask or not - if use_dynamic_chunk: - max_len = xs.size(1) - if decoding_chunk_size < 0: - chunk_size = max_len - num_left_chunks = -1 - elif decoding_chunk_size > 0: - chunk_size = decoding_chunk_size - num_left_chunks = num_decoding_left_chunks - else: - # chunk size is either [1, 25] or full context(max_len). - # Since we use 4 times subsampling and allow up to 1s(100 frames) - # delay, the maximum frame is 100 / 4 = 25. - chunk_size = torch.randint(1, max_len, (1, )).item() - num_left_chunks = -1 - if chunk_size > max_len // 2 and enable_full_context: - chunk_size = max_len - else: - chunk_size = chunk_size % 25 + 1 - if use_dynamic_left_chunk: - max_left_chunks = (max_len - 1) // chunk_size - num_left_chunks = torch.randint(0, max_left_chunks, - (1, )).item() - chunk_masks = subsequent_chunk_mask(xs.size(1), chunk_size, - num_left_chunks, - xs.device) # (L, L) - chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L) - chunk_masks = masks & chunk_masks # (B, L, L) - elif static_chunk_size > 0: - num_left_chunks = num_decoding_left_chunks - chunk_masks = subsequent_chunk_mask(xs.size(1), static_chunk_size, - num_left_chunks, - xs.device) # (L, L) - chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L) - chunk_masks = masks & chunk_masks # (B, L, L) - else: - chunk_masks = masks - assert chunk_masks.dtype == torch.bool - if (chunk_masks.sum(dim=-1) == 0).sum().item() != 0: - print('get chunk_masks all false at some timestep, force set to true, make sure they are masked in futuer computation!') - chunk_masks[chunk_masks.sum(dim=-1) == 0] = True - return chunk_masks - - -def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: - """Make mask tensor containing indices of padded part. - - See description of make_non_pad_mask. - - Args: - lengths (torch.Tensor): Batch of lengths (B,). - Returns: - torch.Tensor: Mask tensor containing indices of padded part. - - Examples: - >>> lengths = [5, 3, 2] - >>> make_pad_mask(lengths) - masks = [[0, 0, 0, 0 ,0], - [0, 0, 0, 1, 1], - [0, 0, 1, 1, 1]] - """ - batch_size = lengths.size(0) - max_len = max_len if max_len > 0 else lengths.max().item() - seq_range = torch.arange(0, - max_len, - dtype=torch.int64, - device=lengths.device) - seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) - seq_length_expand = lengths.unsqueeze(-1) - mask = seq_range_expand >= seq_length_expand - return mask - - -class EspnetRelPositionalEncoding(torch.nn.Module): - """Relative positional encoding module (new implementation). - - Details can be found in https://github.com/espnet/espnet/pull/2816. - - See : Appendix B in https://arxiv.org/abs/1901.02860 - - Args: - d_model (int): Embedding dimension. - max_len (int): Maximum input length. - - """ - - def __init__(self, d_model: int, max_len: int = 5000): - super(EspnetRelPositionalEncoding, self).__init__() - self.d_model = d_model - self.xscale = math.sqrt(self.d_model) - self.pe = None - self.extend_pe(torch.tensor(0.0).expand(1, max_len)) - - def extend_pe(self, x: torch.Tensor): - """Reset the positional encodings.""" - if self.pe is not None: - # self.pe contains both positive and negative parts - # the length of self.pe is 2 * input_len - 1 - if self.pe.size(1) >= x.size(1) * 2 - 1: - if self.pe.dtype != x.dtype or self.pe.device != x.device: - self.pe = self.pe.to(dtype=x.dtype, device=x.device) - return - # Suppose `i` means to the position of query vecotr and `j` means the - # position of key vector. We use position relative positions when keys - # are to the left (i>j) and negative relative positions otherwise (i Tuple[torch.Tensor, torch.Tensor]: - """Add positional encoding. - - Args: - x (torch.Tensor): Input tensor (batch, time, `*`). - - Returns: - torch.Tensor: Encoded tensor (batch, time, `*`). - - """ - self.extend_pe(x) - x = x * self.xscale - pos_emb = self.position_encoding(size=x.size(1), offset=offset) - return x, pos_emb - - def position_encoding(self, - offset: Union[int, torch.Tensor], - size: int) -> torch.Tensor: - """ For getting encoding in a streaming fashion - - Attention!!!!! - we apply dropout only once at the whole utterance level in a none - streaming way, but will call this function several times with - increasing input size in a streaming scenario, so the dropout will - be applied several times. - - Args: - offset (int or torch.tensor): start offset - size (int): required size of position encoding - - Returns: - torch.Tensor: Corresponding encoding - """ - # How to subscript a Union type: - # https://github.com/pytorch/pytorch/issues/69434 - if isinstance(offset, int): - pos_emb = self.pe[ - :, - self.pe.size(1) // 2 - size - offset + 1: self.pe.size(1) // 2 + size + offset, - ] - elif isinstance(offset, torch.Tensor): - pos_emb = self.pe[ - :, - self.pe.size(1) // 2 - size - offset + 1: self.pe.size(1) // 2 + size + offset, - ] - return pos_emb - - -class LinearNoSubsampling(torch.nn.Module): - """Linear transform the input without subsampling - - Args: - idim (int): Input dimension. - odim (int): Output dimension. - pos_enc_class (torch.nn.Module): Positional encoding class. - - """ - - def __init__(self, idim: int, odim: int, - pos_enc_class: torch.nn.Module): - super().__init__() - self.out = torch.nn.Sequential( - torch.nn.Linear(idim, odim), - torch.nn.LayerNorm(odim, eps=1e-5), - ) - self.pos_enc = pos_enc_class - self.right_context = 0 - self.subsampling_rate = 1 - - def forward( - self, - x: torch.Tensor, - x_mask: torch.Tensor, - offset: Union[int, torch.Tensor] = 0 - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Input x. - - Args: - x (torch.Tensor): Input tensor (#batch, time, idim). - x_mask (torch.Tensor): Input mask (#batch, 1, time). - - Returns: - torch.Tensor: linear input tensor (#batch, time', odim), - where time' = time . - torch.Tensor: linear input mask (#batch, 1, time'), - where time' = time . - - """ - x = self.out(x) - x, pos_emb = self.pos_enc(x, offset) - return x, pos_emb, x_mask - - def position_encoding(self, offset: Union[int, torch.Tensor], - size: int) -> torch.Tensor: - return self.pos_enc.position_encoding(offset, size) - - -class Upsample1D(nn.Module): - """A 1D upsampling layer with an optional convolution. - - Parameters: - channels (`int`): - number of channels in the inputs and outputs. - use_conv (`bool`, default `False`): - option to use a convolution. - use_conv_transpose (`bool`, default `False`): - option to use a convolution transpose. - out_channels (`int`, optional): - number of output channels. Defaults to `channels`. - """ - - def __init__(self, channels: int, out_channels: int, stride: int = 2): - super().__init__() - self.channels = channels - self.out_channels = out_channels - self.stride = stride - # In this mode, first repeat interpolate, than conv with stride=1 - self.conv = nn.Conv1d(self.channels, self.out_channels, stride * 2 + 1, stride=1, padding=0) - - def forward(self, inputs: torch.Tensor, input_lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - outputs = F.interpolate(inputs, scale_factor=float(self.stride), mode="nearest") - outputs = F.pad(outputs, (self.stride * 2, 0), value=0.0) - outputs = self.conv(outputs) - return outputs, input_lengths * self.stride - - -class PreLookaheadLayer(nn.Module): - def __init__(self, channels: int, pre_lookahead_len: int = 1): - super().__init__() - self.channels = channels - self.pre_lookahead_len = pre_lookahead_len - self.conv1 = nn.Conv1d( - channels, channels, - kernel_size=pre_lookahead_len + 1, - stride=1, padding=0, - ) - self.conv2 = nn.Conv1d( - channels, channels, - kernel_size=3, stride=1, padding=0, - ) - - def forward(self, inputs: torch.Tensor, context: torch.Tensor = torch.zeros(0, 0, 0)) -> torch.Tensor: - """ - inputs: (batch_size, seq_len, channels) - """ - outputs = inputs.transpose(1, 2).contiguous() - context = context.transpose(1, 2).contiguous() - # look ahead - if context.size(2) == 0: - outputs = F.pad(outputs, (0, self.pre_lookahead_len), mode='constant', value=0.0) - else: - assert self.training is False, 'you have passed context, make sure that you are running inference mode' - assert context.size(2) == self.pre_lookahead_len - outputs = F.pad(torch.concat([outputs, context], dim=2), (0, self.pre_lookahead_len - context.size(2)), mode='constant', value=0.0) - outputs = F.leaky_relu(self.conv1(outputs)) - # outputs - outputs = F.pad(outputs, (self.conv2.kernel_size[0] - 1, 0), mode='constant', value=0.0) - outputs = self.conv2(outputs) - outputs = outputs.transpose(1, 2).contiguous() - - # residual connection - outputs = outputs + inputs - return outputs - - -class MultiHeadedAttention(nn.Module): - """Multi-Head Attention layer. - - Args: - n_head (int): The number of heads. - n_feat (int): The number of features. - dropout_rate (float): Dropout rate. - key_bias (bool): Whether to use bias in key linear layer. - - """ - - def __init__(self, - n_head: int, - n_feat: int, - dropout_rate: float, - key_bias: bool = True): - super().__init__() - assert n_feat % n_head == 0 - # We assume d_v always equals d_k - self.d_k = n_feat // n_head - self.h = n_head - self.linear_q = nn.Linear(n_feat, n_feat) - self.linear_k = nn.Linear(n_feat, n_feat, bias=key_bias) - self.linear_v = nn.Linear(n_feat, n_feat) - self.linear_out = nn.Linear(n_feat, n_feat) - self.dropout = nn.Dropout(p=dropout_rate) - - def forward_qkv( - self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Transform query, key and value. - - Args: - query (torch.Tensor): Query tensor (#batch, time1, size). - key (torch.Tensor): Key tensor (#batch, time2, size). - value (torch.Tensor): Value tensor (#batch, time2, size). - - Returns: - torch.Tensor: Transformed query tensor, size - (#batch, n_head, time1, d_k). - torch.Tensor: Transformed key tensor, size - (#batch, n_head, time2, d_k). - torch.Tensor: Transformed value tensor, size - (#batch, n_head, time2, d_k). - - """ - n_batch = query.size(0) - q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) - k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) - v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) - q = q.transpose(1, 2) # (batch, head, time1, d_k) - k = k.transpose(1, 2) # (batch, head, time2, d_k) - v = v.transpose(1, 2) # (batch, head, time2, d_k) - - return q, k, v - - def forward_attention( - self, - value: torch.Tensor, - scores: torch.Tensor, - mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool) - ) -> torch.Tensor: - """Compute attention context vector. - - Args: - value (torch.Tensor): Transformed value, size - (#batch, n_head, time2, d_k). - scores (torch.Tensor): Attention score, size - (#batch, n_head, time1, time2). - mask (torch.Tensor): Mask, size (#batch, 1, time2) or - (#batch, time1, time2), (0, 0, 0) means fake mask. - - Returns: - torch.Tensor: Transformed value (#batch, time1, d_model) - weighted by the attention score (#batch, time1, time2). - - """ - n_batch = value.size(0) - # NOTE(xcsong): When will `if mask.size(2) > 0` be True? - # 1. onnx(16/4) [WHY? Because we feed real cache & real mask for the - # 1st chunk to ease the onnx export.] - # 2. pytorch training - if mask.size(2) > 0: # time2 > 0 - mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - # For last chunk, time2 might be larger than scores.size(-1) - mask = mask[:, :, :, :scores.size(-1)] # (batch, 1, *, time2) - scores = scores.masked_fill(mask, -float('inf')) - attn = torch.softmax(scores, dim=-1).masked_fill( - mask, 0.0) # (batch, head, time1, time2) - # NOTE(xcsong): When will `if mask.size(2) > 0` be False? - # 1. onnx(16/-1, -1/-1, 16/0) - # 2. jit (16/-1, -1/-1, 16/0, 16/4) - else: - attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) - - p_attn = self.dropout(attn) - x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) - x = (x.transpose(1, 2).contiguous().view(n_batch, -1, - self.h * self.d_k) - ) # (batch, time1, d_model) - - return self.linear_out(x) # (batch, time1, d_model) - - def forward( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), - pos_emb: torch.Tensor = torch.empty(0), - cache: torch.Tensor = torch.zeros((0, 0, 0, 0)) - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Compute scaled dot product attention. - - Args: - query (torch.Tensor): Query tensor (#batch, time1, size). - key (torch.Tensor): Key tensor (#batch, time2, size). - value (torch.Tensor): Value tensor (#batch, time2, size). - mask (torch.Tensor): Mask tensor (#batch, 1, time2) or - (#batch, time1, time2). - 1.When applying cross attention between decoder and encoder, - the batch padding mask for input is in (#batch, 1, T) shape. - 2.When applying self attention of encoder, - the mask is in (#batch, T, T) shape. - 3.When applying self attention of decoder, - the mask is in (#batch, L, L) shape. - 4.If the different position in decoder see different block - of the encoder, such as Mocha, the passed in mask could be - in (#batch, L, T) shape. - cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2), - where `cache_t == chunk_size * num_decoding_left_chunks` - and `head * d_k == size` - - - Returns: - torch.Tensor: Output tensor (#batch, time1, d_model). - torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2) - where `cache_t == chunk_size * num_decoding_left_chunks` - and `head * d_k == size` - - """ - q, k, v = self.forward_qkv(query, key, value) - - # NOTE(xcsong): - # when export onnx model, for 1st chunk, we feed - # cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode) - # or cache(1, head, real_cache_t, d_k * 2) (16/4 mode). - # In all modes, `if cache.size(0) > 0` will alwayse be `True` - # and we will always do splitting and - # concatnation(this will simplify onnx export). Note that - # it's OK to concat & split zero-shaped tensors(see code below). - # when export jit model, for 1st chunk, we always feed - # cache(0, 0, 0, 0) since jit supports dynamic if-branch. - # >>> a = torch.ones((1, 2, 0, 4)) - # >>> b = torch.ones((1, 2, 3, 4)) - # >>> c = torch.cat((a, b), dim=2) - # >>> torch.equal(b, c) # True - # >>> d = torch.split(a, 2, dim=-1) - # >>> torch.equal(d[0], d[1]) # True - if cache.size(0) > 0: - key_cache, value_cache = torch.split(cache, - cache.size(-1) // 2, - dim=-1) - k = torch.cat([key_cache, k], dim=2) - v = torch.cat([value_cache, v], dim=2) - # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's - # non-trivial to calculate `next_cache_start` here. - new_cache = torch.cat((k, v), dim=-1) - - scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) - return self.forward_attention(v, scores, mask), new_cache - - -class RelPositionMultiHeadedAttention(MultiHeadedAttention): - """Multi-Head Attention layer with relative position encoding. - Paper: https://arxiv.org/abs/1901.02860 - Args: - n_head (int): The number of heads. - n_feat (int): The number of features. - dropout_rate (float): Dropout rate. - key_bias (bool): Whether to use bias in key linear layer. - """ - - def __init__(self, - n_head: int, - n_feat: int, - dropout_rate: float, - key_bias: bool = True): - super().__init__(n_head, n_feat, dropout_rate, key_bias) - # linear transformation for positional encoding - self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) - # these two learnable bias are used in matrix c and matrix d - # as described in https://arxiv.org/abs/1901.02860 Section 3.3 - self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) - self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) - torch.nn.init.xavier_uniform_(self.pos_bias_u) - torch.nn.init.xavier_uniform_(self.pos_bias_v) - - def rel_shift(self, x: torch.Tensor) -> torch.Tensor: - """Compute relative positional encoding. - - Args: - x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1). - time1 means the length of query vector. - - Returns: - torch.Tensor: Output tensor. - - """ - zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1), - device=x.device, - dtype=x.dtype) - x_padded = torch.cat([zero_pad, x], dim=-1) - - x_padded = x_padded.view(x.size()[0], - x.size()[1], - x.size(3) + 1, x.size(2)) - x = x_padded[:, :, 1:].view_as(x)[ - :, :, :, : x.size(-1) // 2 + 1 - ] # only keep the positions from 0 to time2 - return x - - def forward( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), - pos_emb: torch.Tensor = torch.empty(0), - cache: torch.Tensor = torch.zeros((0, 0, 0, 0)) - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Compute 'Scaled Dot Product Attention' with rel. positional encoding. - Args: - query (torch.Tensor): Query tensor (#batch, time1, size). - key (torch.Tensor): Key tensor (#batch, time2, size). - value (torch.Tensor): Value tensor (#batch, time2, size). - mask (torch.Tensor): Mask tensor (#batch, 1, time2) or - (#batch, time1, time2), (0, 0, 0) means fake mask. - pos_emb (torch.Tensor): Positional embedding tensor - (#batch, time2, size). - cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2), - where `cache_t == chunk_size * num_decoding_left_chunks` - and `head * d_k == size` - Returns: - torch.Tensor: Output tensor (#batch, time1, d_model). - torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2) - where `cache_t == chunk_size * num_decoding_left_chunks` - and `head * d_k == size` - """ - q, k, v = self.forward_qkv(query, key, value) - q = q.transpose(1, 2) # (batch, time1, head, d_k) - - # NOTE(xcsong): - # when export onnx model, for 1st chunk, we feed - # cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode) - # or cache(1, head, real_cache_t, d_k * 2) (16/4 mode). - # In all modes, `if cache.size(0) > 0` will alwayse be `True` - # and we will always do splitting and - # concatnation(this will simplify onnx export). Note that - # it's OK to concat & split zero-shaped tensors(see code below). - # when export jit model, for 1st chunk, we always feed - # cache(0, 0, 0, 0) since jit supports dynamic if-branch. - # >>> a = torch.ones((1, 2, 0, 4)) - # >>> b = torch.ones((1, 2, 3, 4)) - # >>> c = torch.cat((a, b), dim=2) - # >>> torch.equal(b, c) # True - # >>> d = torch.split(a, 2, dim=-1) - # >>> torch.equal(d[0], d[1]) # True - if cache.size(0) > 0: - key_cache, value_cache = torch.split(cache, - cache.size(-1) // 2, - dim=-1) - k = torch.cat([key_cache, k], dim=2) - v = torch.cat([value_cache, v], dim=2) - # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's - # non-trivial to calculate `next_cache_start` here. - new_cache = torch.cat((k, v), dim=-1) - - n_batch_pos = pos_emb.size(0) - p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) - p = p.transpose(1, 2) # (batch, head, time1, d_k) - - # (batch, head, time1, d_k) - q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) - # (batch, head, time1, d_k) - q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) - - # compute attention score - # first compute matrix a and matrix c - # as described in https://arxiv.org/abs/1901.02860 Section 3.3 - # (batch, head, time1, time2) - matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) - - # compute matrix b and matrix d - # (batch, head, time1, time2) - matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) - # NOTE(Xiang Lyu): Keep rel_shift since espnet rel_pos_emb is used - if matrix_ac.shape != matrix_bd.shape: - matrix_bd = self.rel_shift(matrix_bd) - - scores = (matrix_ac + matrix_bd) / math.sqrt( - self.d_k) # (batch, head, time1, time2) - - return self.forward_attention(v, scores, mask), new_cache - - -class PositionwiseFeedForward(torch.nn.Module): - """Positionwise feed forward layer. - - FeedForward are appied on each position of the sequence. - The output dim is same with the input dim. - - Args: - idim (int): Input dimenstion. - hidden_units (int): The number of hidden units. - dropout_rate (float): Dropout rate. - activation (torch.nn.Module): Activation function - """ - - def __init__( - self, - idim: int, - hidden_units: int, - dropout_rate: float, - activation: torch.nn.Module = torch.nn.ReLU(), - ): - super(PositionwiseFeedForward, self).__init__() - self.w_1 = torch.nn.Linear(idim, hidden_units) - self.activation = activation - self.dropout = torch.nn.Dropout(dropout_rate) - self.w_2 = torch.nn.Linear(hidden_units, idim) - - def forward(self, xs: torch.Tensor) -> torch.Tensor: - """Forward function. - - Args: - xs: input tensor (B, L, D) - Returns: - output tensor, (B, L, D) - """ - return self.w_2(self.dropout(self.activation(self.w_1(xs)))) - - -class ConformerEncoderLayer(nn.Module): - """Encoder layer module. - Args: - size (int): Input dimension. - self_attn (torch.nn.Module): Self-attention module instance. - `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` - instance can be used as the argument. - feed_forward (torch.nn.Module): Feed-forward module instance. - `PositionwiseFeedForward` instance can be used as the argument. - feed_forward_macaron (torch.nn.Module): Additional feed-forward module - instance. - `PositionwiseFeedForward` instance can be used as the argument. - conv_module (torch.nn.Module): Convolution module instance. - `ConvlutionModule` instance can be used as the argument. - dropout_rate (float): Dropout rate. - normalize_before (bool): - True: use layer_norm before each sub-block. - False: use layer_norm after each sub-block. - """ - - def __init__( - self, - size: int, - self_attn: torch.nn.Module, - feed_forward: Optional[nn.Module] = None, - feed_forward_macaron: Optional[nn.Module] = None, - conv_module: Optional[nn.Module] = None, - dropout_rate: float = 0.0, - normalize_before: bool = True, - ): - super().__init__() - self.self_attn = self_attn - self.feed_forward = feed_forward - self.feed_forward_macaron = feed_forward_macaron - self.conv_module = conv_module - self.norm_ff = nn.LayerNorm(size, eps=1e-12) # for the FNN module - self.norm_mha = nn.LayerNorm(size, eps=1e-12) # for the MHA module - if feed_forward_macaron is not None: - self.norm_ff_macaron = nn.LayerNorm(size, eps=1e-12) - self.ff_scale = 0.5 - else: - self.ff_scale = 1.0 - if self.conv_module is not None: - self.norm_conv = nn.LayerNorm(size, eps=1e-12) # for the CNN module - self.norm_final = nn.LayerNorm( - size, eps=1e-12) # for the final output of the block - self.dropout = nn.Dropout(dropout_rate) - self.size = size - self.normalize_before = normalize_before - - def forward( - self, - x: torch.Tensor, - mask: torch.Tensor, - pos_emb: torch.Tensor, - mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), - att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), - cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute encoded features. - - Args: - x (torch.Tensor): (#batch, time, size) - mask (torch.Tensor): Mask tensor for the input (#batch, time，time), - (0, 0, 0) means fake mask. - pos_emb (torch.Tensor): positional encoding, must not be None - for ConformerEncoderLayer. - mask_pad (torch.Tensor): batch padding mask used for conv module. - (#batch, 1，time), (0, 0, 0) means fake mask. - att_cache (torch.Tensor): Cache tensor of the KEY & VALUE - (#batch=1, head, cache_t1, d_k * 2), head * d_k == size. - cnn_cache (torch.Tensor): Convolution cache in conformer layer - (#batch=1, size, cache_t2) - Returns: - torch.Tensor: Output tensor (#batch, time, size). - torch.Tensor: Mask tensor (#batch, time, time). - torch.Tensor: att_cache tensor, - (#batch=1, head, cache_t1 + time, d_k * 2). - torch.Tensor: cnn_cahce tensor (#batch, size, cache_t2). - """ - - # whether to use macaron style - if self.feed_forward_macaron is not None: - residual = x - if self.normalize_before: - x = self.norm_ff_macaron(x) - x = residual + self.ff_scale * self.dropout( - self.feed_forward_macaron(x)) - if not self.normalize_before: - x = self.norm_ff_macaron(x) - - # multi-headed self-attention module - residual = x - if self.normalize_before: - x = self.norm_mha(x) - x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, - att_cache) - x = residual + self.dropout(x_att) - if not self.normalize_before: - x = self.norm_mha(x) - - # convolution module - # Fake new cnn cache here, and then change it in conv_module - new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device) - if self.conv_module is not None: - residual = x - if self.normalize_before: - x = self.norm_conv(x) - x, new_cnn_cache = self.conv_module(x, mask_pad, cnn_cache) - x = residual + self.dropout(x) - - if not self.normalize_before: - x = self.norm_conv(x) - - # feed forward module - residual = x - if self.normalize_before: - x = self.norm_ff(x) - - x = residual + self.ff_scale * self.dropout(self.feed_forward(x)) - if not self.normalize_before: - x = self.norm_ff(x) - - if self.conv_module is not None: - x = self.norm_final(x) - - return x, mask, new_att_cache, new_cnn_cache - - -class UpsampleConformerEncoder(torch.nn.Module): - """ - Args: - input_size (int): input dim - output_size (int): dimension of attention - attention_heads (int): the number of heads of multi head attention - linear_units (int): the hidden units number of position-wise feed - forward - num_blocks (int): the number of decoder blocks - static_chunk_size (int): chunk size for static chunk training and - decoding - use_dynamic_chunk (bool): whether use dynamic chunk size for - training or not, You can only use fixed chunk(chunk_size > 0) - or dyanmic chunk size(use_dynamic_chunk = True) - use_dynamic_left_chunk (bool): whether use dynamic left chunk in - dynamic chunk training - key_bias: whether use bias in attention.linear_k, False for whisper models. - """ - - def __init__( - self, - input_size: int = 512, - output_size: int = 512, - attention_heads: int = 8, - linear_units: int = 2048, - num_blocks: int = 6, - static_chunk_size: int = 25, - use_dynamic_chunk: bool = False, - use_dynamic_left_chunk: bool = False, - key_bias: bool = True, - ): - super().__init__() - self._output_size = output_size - - self.embed = LinearNoSubsampling( - input_size, output_size, - EspnetRelPositionalEncoding(output_size), - ) - - self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5) - self.static_chunk_size = static_chunk_size - self.use_dynamic_chunk = use_dynamic_chunk - self.use_dynamic_left_chunk = use_dynamic_left_chunk - activation = torch.nn.SiLU() - # self-attention module definition - encoder_selfattn_layer_args = ( - attention_heads, - output_size, - 0.0, - key_bias, - ) - # feed-forward module definition - positionwise_layer_args = ( - output_size, - linear_units, - 0.0, - activation, - ) - # convolution module definition - self.pre_lookahead_layer = PreLookaheadLayer(channels=512, pre_lookahead_len=3) - self.encoders = torch.nn.ModuleList([ - ConformerEncoderLayer( - output_size, - RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args), - PositionwiseFeedForward(*positionwise_layer_args), - ) for _ in range(num_blocks) - ]) - self.up_layer = Upsample1D(channels=512, out_channels=512, stride=2) - self.up_embed = LinearNoSubsampling( - input_size, output_size, - EspnetRelPositionalEncoding(output_size), - ) - self.up_encoders = torch.nn.ModuleList([ - ConformerEncoderLayer( - output_size, - RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args), - PositionwiseFeedForward(*positionwise_layer_args), - ) for _ in range(4) - ]) - - def output_size(self) -> int: - return self._output_size - - def forward( - self, - xs: torch.Tensor, - xs_lens: torch.Tensor, - context: torch.Tensor = torch.zeros(0, 0, 0), - decoding_chunk_size: int = 0, - num_decoding_left_chunks: int = -1, - streaming: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Embed positions in tensor. - - Args: - xs: padded input tensor (B, T, D) - xs_lens: input length (B) - decoding_chunk_size: decoding chunk size for dynamic chunk - 0: default for training, use random dynamic chunk. - <0: for decoding, use full chunk. - >0: for decoding, use fixed chunk size as set. - num_decoding_left_chunks: number of left chunks, this is for decoding, - the chunk size is decoding_chunk_size. - >=0: use num_decoding_left_chunks - <0: use all left chunks - Returns: - encoder output tensor xs, and subsampled masks - xs: padded output tensor (B, T' ~= T/subsample_rate, D) - masks: torch.Tensor batch padding mask after subsample - (B, 1, T' ~= T/subsample_rate) - NOTE(xcsong): - We pass the `__call__` method of the modules instead of `forward` to the - checkpointing API because `__call__` attaches all the hooks of the module. - https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2 - """ - T = xs.size(1) - masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) # (B, 1, T) - xs, pos_emb, masks = self.embed(xs, masks) - if context.size(1) != 0: - assert self.training is False, 'you have passed context, make sure that you are running inference mode' - context_masks = torch.ones(1, 1, context.size(1)).to(masks) - context, _, _ = self.embed(context, context_masks, offset=xs.size(1)) - mask_pad = masks # (B, 1, T/subsample_rate) - chunk_masks = add_optional_chunk_mask(xs, masks, False, False, 0, self.static_chunk_size if streaming is True else 0, -1) - # lookahead + conformer encoder - xs = self.pre_lookahead_layer(xs, context=context) - xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad) - - # upsample + conformer encoder - xs = xs.transpose(1, 2).contiguous() - xs, xs_lens = self.up_layer(xs, xs_lens) - xs = xs.transpose(1, 2).contiguous() - T = xs.size(1) - masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) # (B, 1, T) - xs, pos_emb, masks = self.up_embed(xs, masks) - mask_pad = masks # (B, 1, T/subsample_rate) - chunk_masks = add_optional_chunk_mask(xs, masks, False, False, 0, self.static_chunk_size * self.up_layer.stride if streaming is True else 0, -1) - xs = self.forward_up_layers(xs, chunk_masks, pos_emb, mask_pad) - - xs = self.after_norm(xs) - # Here we assume the mask is not changed in encoder layers, so just - # return the masks before encoder layers, and the masks will be used - # for cross attention with decoder later - return xs, masks - - def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor, - pos_emb: torch.Tensor, - mask_pad: torch.Tensor) -> torch.Tensor: - for layer in self.encoders: - xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad) - return xs - - def forward_up_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor, - pos_emb: torch.Tensor, - mask_pad: torch.Tensor) -> torch.Tensor: - for layer in self.up_encoders: - xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad) - return xs +import math +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def subsequent_chunk_mask( + size: int, + chunk_size: int, + num_left_chunks: int = -1, + device: torch.device = torch.device("cpu"), +) -> torch.Tensor: + """Create mask for subsequent steps (size, size) with chunk size, + this is for streaming encoder + + Args: + size (int): size of mask + chunk_size (int): size of chunk + num_left_chunks (int): number of left chunks + <0: use full chunk + >=0: use num_left_chunks + device (torch.device): "cpu" or "cuda" or torch.Tensor.device + + Returns: + torch.Tensor: mask + + Examples: + >>> subsequent_chunk_mask(4, 2) + [[1, 1, 0, 0], + [1, 1, 0, 0], + [1, 1, 1, 1], + [1, 1, 1, 1]] + """ + # NOTE this modified implementation meets onnx export requirements, but it doesn't support num_left_chunks + pos_idx = torch.arange(size, device=device) + block_value = (torch.div(pos_idx, chunk_size, rounding_mode='trunc') + 1) * chunk_size + ret = pos_idx.unsqueeze(0) < block_value.unsqueeze(1) + return ret + + +def add_optional_chunk_mask(xs: torch.Tensor, + masks: torch.Tensor, + use_dynamic_chunk: bool, + use_dynamic_left_chunk: bool, + decoding_chunk_size: int, + static_chunk_size: int, + num_decoding_left_chunks: int, + enable_full_context: bool = True): + """ Apply optional mask for encoder. + + Args: + xs (torch.Tensor): padded input, (B, L, D), L for max length + mask (torch.Tensor): mask for xs, (B, 1, L) + use_dynamic_chunk (bool): whether to use dynamic chunk or not + use_dynamic_left_chunk (bool): whether to use dynamic left chunk for + training. + decoding_chunk_size (int): decoding chunk size for dynamic chunk, it's + 0: default for training, use random dynamic chunk. + <0: for decoding, use full chunk. + >0: for decoding, use fixed chunk size as set. + static_chunk_size (int): chunk size for static chunk training/decoding + if it's greater than 0, if use_dynamic_chunk is true, + this parameter will be ignored + num_decoding_left_chunks: number of left chunks, this is for decoding, + the chunk size is decoding_chunk_size. + >=0: use num_decoding_left_chunks + <0: use all left chunks + enable_full_context (bool): + True: chunk size is either [1, 25] or full context(max_len) + False: chunk size ~ U[1, 25] + + Returns: + torch.Tensor: chunk mask of the input xs. + """ + # Whether to use chunk mask or not + if use_dynamic_chunk: + max_len = xs.size(1) + if decoding_chunk_size < 0: + chunk_size = max_len + num_left_chunks = -1 + elif decoding_chunk_size > 0: + chunk_size = decoding_chunk_size + num_left_chunks = num_decoding_left_chunks + else: + # chunk size is either [1, 25] or full context(max_len). + # Since we use 4 times subsampling and allow up to 1s(100 frames) + # delay, the maximum frame is 100 / 4 = 25. + chunk_size = torch.randint(1, max_len, (1, )).item() + num_left_chunks = -1 + if chunk_size > max_len // 2 and enable_full_context: + chunk_size = max_len + else: + chunk_size = chunk_size % 25 + 1 + if use_dynamic_left_chunk: + max_left_chunks = (max_len - 1) // chunk_size + num_left_chunks = torch.randint(0, max_left_chunks, + (1, )).item() + chunk_masks = subsequent_chunk_mask(xs.size(1), chunk_size, + num_left_chunks, + xs.device) # (L, L) + chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L) + chunk_masks = masks & chunk_masks # (B, L, L) + elif static_chunk_size > 0: + num_left_chunks = num_decoding_left_chunks + chunk_masks = subsequent_chunk_mask(xs.size(1), static_chunk_size, + num_left_chunks, + xs.device) # (L, L) + chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L) + chunk_masks = masks & chunk_masks # (B, L, L) + else: + chunk_masks = masks + assert chunk_masks.dtype == torch.bool + if (chunk_masks.sum(dim=-1) == 0).sum().item() != 0: + print('get chunk_masks all false at some timestep, force set to true, make sure they are masked in futuer computation!') + chunk_masks[chunk_masks.sum(dim=-1) == 0] = True + return chunk_masks + + +def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: + """Make mask tensor containing indices of padded part. + + See description of make_non_pad_mask. + + Args: + lengths (torch.Tensor): Batch of lengths (B,). + Returns: + torch.Tensor: Mask tensor containing indices of padded part. + + Examples: + >>> lengths = [5, 3, 2] + >>> make_pad_mask(lengths) + masks = [[0, 0, 0, 0 ,0], + [0, 0, 0, 1, 1], + [0, 0, 1, 1, 1]] + """ + batch_size = lengths.size(0) + max_len = max_len if max_len > 0 else lengths.max().item() + seq_range = torch.arange(0, + max_len, + dtype=torch.int64, + device=lengths.device) + seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) + seq_length_expand = lengths.unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + return mask + + +class EspnetRelPositionalEncoding(torch.nn.Module): + """Relative positional encoding module (new implementation). + + Details can be found in https://github.com/espnet/espnet/pull/2816. + + See : Appendix B in https://arxiv.org/abs/1901.02860 + + Args: + d_model (int): Embedding dimension. + max_len (int): Maximum input length. + + """ + + def __init__(self, d_model: int, max_len: int = 5000): + super(EspnetRelPositionalEncoding, self).__init__() + self.d_model = d_model + self.xscale = math.sqrt(self.d_model) + self.pe = None + self.extend_pe(torch.tensor(0.0).expand(1, max_len)) + + def extend_pe(self, x: torch.Tensor): + """Reset the positional encodings.""" + if self.pe is not None: + # self.pe contains both positive and negative parts + # the length of self.pe is 2 * input_len - 1 + if self.pe.size(1) >= x.size(1) * 2 - 1: + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + # Suppose `i` means to the position of query vecotr and `j` means the + # position of key vector. We use position relative positions when keys + # are to the left (i>j) and negative relative positions otherwise (i Tuple[torch.Tensor, torch.Tensor]: + """Add positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + + """ + self.extend_pe(x) + x = x * self.xscale + pos_emb = self.position_encoding(size=x.size(1), offset=offset) + return x, pos_emb + + def position_encoding(self, + offset: Union[int, torch.Tensor], + size: int) -> torch.Tensor: + """ For getting encoding in a streaming fashion + + Attention!!!!! + we apply dropout only once at the whole utterance level in a none + streaming way, but will call this function several times with + increasing input size in a streaming scenario, so the dropout will + be applied several times. + + Args: + offset (int or torch.tensor): start offset + size (int): required size of position encoding + + Returns: + torch.Tensor: Corresponding encoding + """ + # How to subscript a Union type: + # https://github.com/pytorch/pytorch/issues/69434 + if isinstance(offset, int): + pos_emb = self.pe[ + :, + self.pe.size(1) // 2 - size - offset + 1: self.pe.size(1) // 2 + size + offset, + ] + elif isinstance(offset, torch.Tensor): + pos_emb = self.pe[ + :, + self.pe.size(1) // 2 - size - offset + 1: self.pe.size(1) // 2 + size + offset, + ] + return pos_emb + + +class LinearNoSubsampling(torch.nn.Module): + """Linear transform the input without subsampling + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + pos_enc_class (torch.nn.Module): Positional encoding class. + + """ + + def __init__(self, idim: int, odim: int, + pos_enc_class: torch.nn.Module): + super().__init__() + self.out = torch.nn.Sequential( + torch.nn.Linear(idim, odim), + torch.nn.LayerNorm(odim, eps=1e-5), + ) + self.pos_enc = pos_enc_class + self.right_context = 0 + self.subsampling_rate = 1 + + def forward( + self, + x: torch.Tensor, + x_mask: torch.Tensor, + offset: Union[int, torch.Tensor] = 0 + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Input x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: linear input tensor (#batch, time', odim), + where time' = time . + torch.Tensor: linear input mask (#batch, 1, time'), + where time' = time . + + """ + x = self.out(x) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask + + def position_encoding(self, offset: Union[int, torch.Tensor], + size: int) -> torch.Tensor: + return self.pos_enc.position_encoding(offset, size) + + +class Upsample1D(nn.Module): + """A 1D upsampling layer with an optional convolution. + + Parameters: + channels (`int`): + number of channels in the inputs and outputs. + use_conv (`bool`, default `False`): + option to use a convolution. + use_conv_transpose (`bool`, default `False`): + option to use a convolution transpose. + out_channels (`int`, optional): + number of output channels. Defaults to `channels`. + """ + + def __init__(self, channels: int, out_channels: int, stride: int = 2): + super().__init__() + self.channels = channels + self.out_channels = out_channels + self.stride = stride + # In this mode, first repeat interpolate, than conv with stride=1 + self.conv = nn.Conv1d(self.channels, self.out_channels, stride * 2 + 1, stride=1, padding=0) + + def forward(self, inputs: torch.Tensor, input_lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + outputs = F.interpolate(inputs, scale_factor=float(self.stride), mode="nearest") + outputs = F.pad(outputs, (self.stride * 2, 0), value=0.0) + outputs = self.conv(outputs) + return outputs, input_lengths * self.stride + + +class PreLookaheadLayer(nn.Module): + def __init__(self, channels: int, pre_lookahead_len: int = 1): + super().__init__() + self.channels = channels + self.pre_lookahead_len = pre_lookahead_len + self.conv1 = nn.Conv1d( + channels, channels, + kernel_size=pre_lookahead_len + 1, + stride=1, padding=0, + ) + self.conv2 = nn.Conv1d( + channels, channels, + kernel_size=3, stride=1, padding=0, + ) + + def forward(self, inputs: torch.Tensor, context: torch.Tensor = torch.zeros(0, 0, 0)) -> torch.Tensor: + """ + inputs: (batch_size, seq_len, channels) + """ + outputs = inputs.transpose(1, 2).contiguous() + context = context.transpose(1, 2).contiguous() + # look ahead + if context.size(2) == 0: + outputs = F.pad(outputs, (0, self.pre_lookahead_len), mode='constant', value=0.0) + else: + assert self.training is False, 'you have passed context, make sure that you are running inference mode' + assert context.size(2) == self.pre_lookahead_len + outputs = F.pad(torch.concat([outputs, context], dim=2), (0, self.pre_lookahead_len - context.size(2)), mode='constant', value=0.0) + outputs = F.leaky_relu(self.conv1(outputs)) + # outputs + outputs = F.pad(outputs, (self.conv2.kernel_size[0] - 1, 0), mode='constant', value=0.0) + outputs = self.conv2(outputs) + outputs = outputs.transpose(1, 2).contiguous() + + # residual connection + outputs = outputs + inputs + return outputs + + +class MultiHeadedAttention(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + key_bias (bool): Whether to use bias in key linear layer. + + """ + + def __init__(self, + n_head: int, + n_feat: int, + dropout_rate: float, + key_bias: bool = True): + super().__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + self.linear_q = nn.Linear(n_feat, n_feat) + self.linear_k = nn.Linear(n_feat, n_feat, bias=key_bias) + self.linear_v = nn.Linear(n_feat, n_feat) + self.linear_out = nn.Linear(n_feat, n_feat) + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor, size + (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor, size + (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor, size + (#batch, n_head, time2, d_k). + + """ + n_batch = query.size(0) + q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) + k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) + v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) + q = q.transpose(1, 2) # (batch, head, time1, d_k) + k = k.transpose(1, 2) # (batch, head, time2, d_k) + v = v.transpose(1, 2) # (batch, head, time2, d_k) + + return q, k, v + + def forward_attention( + self, + value: torch.Tensor, + scores: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool) + ) -> torch.Tensor: + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value, size + (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score, size + (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask, size (#batch, 1, time2) or + (#batch, time1, time2), (0, 0, 0) means fake mask. + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + # NOTE(xcsong): When will `if mask.size(2) > 0` be True? + # 1. onnx(16/4) [WHY? Because we feed real cache & real mask for the + # 1st chunk to ease the onnx export.] + # 2. pytorch training + if mask.size(2) > 0: # time2 > 0 + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + # For last chunk, time2 might be larger than scores.size(-1) + mask = mask[:, :, :, :scores.size(-1)] # (batch, 1, *, time2) + scores = scores.masked_fill(mask, -float('inf')) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0) # (batch, head, time1, time2) + # NOTE(xcsong): When will `if mask.size(2) > 0` be False? + # 1. onnx(16/-1, -1/-1, 16/0) + # 2. jit (16/-1, -1/-1, 16/0, 16/4) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = (x.transpose(1, 2).contiguous().view(n_batch, -1, + self.h * self.d_k) + ) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + pos_emb: torch.Tensor = torch.empty(0), + cache: torch.Tensor = torch.zeros((0, 0, 0, 0)) + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + 1.When applying cross attention between decoder and encoder, + the batch padding mask for input is in (#batch, 1, T) shape. + 2.When applying self attention of encoder, + the mask is in (#batch, T, T) shape. + 3.When applying self attention of decoder, + the mask is in (#batch, L, L) shape. + 4.If the different position in decoder see different block + of the encoder, such as Mocha, the passed in mask could be + in (#batch, L, T) shape. + cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2), + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2) + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + + """ + q, k, v = self.forward_qkv(query, key, value) + + # NOTE(xcsong): + # when export onnx model, for 1st chunk, we feed + # cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode) + # or cache(1, head, real_cache_t, d_k * 2) (16/4 mode). + # In all modes, `if cache.size(0) > 0` will alwayse be `True` + # and we will always do splitting and + # concatnation(this will simplify onnx export). Note that + # it's OK to concat & split zero-shaped tensors(see code below). + # when export jit model, for 1st chunk, we always feed + # cache(0, 0, 0, 0) since jit supports dynamic if-branch. + # >>> a = torch.ones((1, 2, 0, 4)) + # >>> b = torch.ones((1, 2, 3, 4)) + # >>> c = torch.cat((a, b), dim=2) + # >>> torch.equal(b, c) # True + # >>> d = torch.split(a, 2, dim=-1) + # >>> torch.equal(d[0], d[1]) # True + if cache.size(0) > 0: + key_cache, value_cache = torch.split(cache, + cache.size(-1) // 2, + dim=-1) + k = torch.cat([key_cache, k], dim=2) + v = torch.cat([value_cache, v], dim=2) + # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's + # non-trivial to calculate `next_cache_start` here. + new_cache = torch.cat((k, v), dim=-1) + + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask), new_cache + + +class RelPositionMultiHeadedAttention(MultiHeadedAttention): + """Multi-Head Attention layer with relative position encoding. + Paper: https://arxiv.org/abs/1901.02860 + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + key_bias (bool): Whether to use bias in key linear layer. + """ + + def __init__(self, + n_head: int, + n_feat: int, + dropout_rate: float, + key_bias: bool = True): + super().__init__(n_head, n_feat, dropout_rate, key_bias) + # linear transformation for positional encoding + self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) + # these two learnable bias are used in matrix c and matrix d + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) + self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) + torch.nn.init.xavier_uniform_(self.pos_bias_u) + torch.nn.init.xavier_uniform_(self.pos_bias_v) + + def rel_shift(self, x: torch.Tensor) -> torch.Tensor: + """Compute relative positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1). + time1 means the length of query vector. + + Returns: + torch.Tensor: Output tensor. + + """ + zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1), + device=x.device, + dtype=x.dtype) + x_padded = torch.cat([zero_pad, x], dim=-1) + + x_padded = x_padded.view(x.size()[0], + x.size()[1], + x.size(3) + 1, x.size(2)) + x = x_padded[:, :, 1:].view_as(x)[ + :, :, :, : x.size(-1) // 2 + 1 + ] # only keep the positions from 0 to time2 + return x + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + pos_emb: torch.Tensor = torch.empty(0), + cache: torch.Tensor = torch.zeros((0, 0, 0, 0)) + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute 'Scaled Dot Product Attention' with rel. positional encoding. + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2), (0, 0, 0) means fake mask. + pos_emb (torch.Tensor): Positional embedding tensor + (#batch, time2, size). + cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2), + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2) + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + """ + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + # NOTE(xcsong): + # when export onnx model, for 1st chunk, we feed + # cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode) + # or cache(1, head, real_cache_t, d_k * 2) (16/4 mode). + # In all modes, `if cache.size(0) > 0` will alwayse be `True` + # and we will always do splitting and + # concatnation(this will simplify onnx export). Note that + # it's OK to concat & split zero-shaped tensors(see code below). + # when export jit model, for 1st chunk, we always feed + # cache(0, 0, 0, 0) since jit supports dynamic if-branch. + # >>> a = torch.ones((1, 2, 0, 4)) + # >>> b = torch.ones((1, 2, 3, 4)) + # >>> c = torch.cat((a, b), dim=2) + # >>> torch.equal(b, c) # True + # >>> d = torch.split(a, 2, dim=-1) + # >>> torch.equal(d[0], d[1]) # True + if cache.size(0) > 0: + key_cache, value_cache = torch.split(cache, + cache.size(-1) // 2, + dim=-1) + k = torch.cat([key_cache, k], dim=2) + v = torch.cat([value_cache, v], dim=2) + # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's + # non-trivial to calculate `next_cache_start` here. + new_cache = torch.cat((k, v), dim=-1) + + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) # (batch, head, time1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + + # compute matrix b and matrix d + # (batch, head, time1, time2) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + # NOTE(Xiang Lyu): Keep rel_shift since espnet rel_pos_emb is used + if matrix_ac.shape != matrix_bd.shape: + matrix_bd = self.rel_shift(matrix_bd) + + scores = (matrix_ac + matrix_bd) / math.sqrt( + self.d_k) # (batch, head, time1, time2) + + return self.forward_attention(v, scores, mask), new_cache + + +class PositionwiseFeedForward(torch.nn.Module): + """Positionwise feed forward layer. + + FeedForward are appied on each position of the sequence. + The output dim is same with the input dim. + + Args: + idim (int): Input dimenstion. + hidden_units (int): The number of hidden units. + dropout_rate (float): Dropout rate. + activation (torch.nn.Module): Activation function + """ + + def __init__( + self, + idim: int, + hidden_units: int, + dropout_rate: float, + activation: torch.nn.Module = torch.nn.ReLU(), + ): + super(PositionwiseFeedForward, self).__init__() + self.w_1 = torch.nn.Linear(idim, hidden_units) + self.activation = activation + self.dropout = torch.nn.Dropout(dropout_rate) + self.w_2 = torch.nn.Linear(hidden_units, idim) + + def forward(self, xs: torch.Tensor) -> torch.Tensor: + """Forward function. + + Args: + xs: input tensor (B, L, D) + Returns: + output tensor, (B, L, D) + """ + return self.w_2(self.dropout(self.activation(self.w_1(xs)))) + + +class ConformerEncoderLayer(nn.Module): + """Encoder layer module. + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` + instance can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward` instance can be used as the argument. + feed_forward_macaron (torch.nn.Module): Additional feed-forward module + instance. + `PositionwiseFeedForward` instance can be used as the argument. + conv_module (torch.nn.Module): Convolution module instance. + `ConvlutionModule` instance can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): + True: use layer_norm before each sub-block. + False: use layer_norm after each sub-block. + """ + + def __init__( + self, + size: int, + self_attn: torch.nn.Module, + feed_forward: Optional[nn.Module] = None, + feed_forward_macaron: Optional[nn.Module] = None, + conv_module: Optional[nn.Module] = None, + dropout_rate: float = 0.0, + normalize_before: bool = True, + ): + super().__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.feed_forward_macaron = feed_forward_macaron + self.conv_module = conv_module + self.norm_ff = nn.LayerNorm(size, eps=1e-12) # for the FNN module + self.norm_mha = nn.LayerNorm(size, eps=1e-12) # for the MHA module + if feed_forward_macaron is not None: + self.norm_ff_macaron = nn.LayerNorm(size, eps=1e-12) + self.ff_scale = 0.5 + else: + self.ff_scale = 1.0 + if self.conv_module is not None: + self.norm_conv = nn.LayerNorm(size, eps=1e-12) # for the CNN module + self.norm_final = nn.LayerNorm( + size, eps=1e-12) # for the final output of the block + self.dropout = nn.Dropout(dropout_rate) + self.size = size + self.normalize_before = normalize_before + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor, + pos_emb: torch.Tensor, + mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute encoded features. + + Args: + x (torch.Tensor): (#batch, time, size) + mask (torch.Tensor): Mask tensor for the input (#batch, time，time), + (0, 0, 0) means fake mask. + pos_emb (torch.Tensor): positional encoding, must not be None + for ConformerEncoderLayer. + mask_pad (torch.Tensor): batch padding mask used for conv module. + (#batch, 1，time), (0, 0, 0) means fake mask. + att_cache (torch.Tensor): Cache tensor of the KEY & VALUE + (#batch=1, head, cache_t1, d_k * 2), head * d_k == size. + cnn_cache (torch.Tensor): Convolution cache in conformer layer + (#batch=1, size, cache_t2) + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time, time). + torch.Tensor: att_cache tensor, + (#batch=1, head, cache_t1 + time, d_k * 2). + torch.Tensor: cnn_cahce tensor (#batch, size, cache_t2). + """ + + # whether to use macaron style + if self.feed_forward_macaron is not None: + residual = x + if self.normalize_before: + x = self.norm_ff_macaron(x) + x = residual + self.ff_scale * self.dropout( + self.feed_forward_macaron(x)) + if not self.normalize_before: + x = self.norm_ff_macaron(x) + + # multi-headed self-attention module + residual = x + if self.normalize_before: + x = self.norm_mha(x) + x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, + att_cache) + x = residual + self.dropout(x_att) + if not self.normalize_before: + x = self.norm_mha(x) + + # convolution module + # Fake new cnn cache here, and then change it in conv_module + new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device) + if self.conv_module is not None: + residual = x + if self.normalize_before: + x = self.norm_conv(x) + x, new_cnn_cache = self.conv_module(x, mask_pad, cnn_cache) + x = residual + self.dropout(x) + + if not self.normalize_before: + x = self.norm_conv(x) + + # feed forward module + residual = x + if self.normalize_before: + x = self.norm_ff(x) + + x = residual + self.ff_scale * self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm_ff(x) + + if self.conv_module is not None: + x = self.norm_final(x) + + return x, mask, new_att_cache, new_cnn_cache + + +class UpsampleConformerEncoder(torch.nn.Module): + """ + Args: + input_size (int): input dim + output_size (int): dimension of attention + attention_heads (int): the number of heads of multi head attention + linear_units (int): the hidden units number of position-wise feed + forward + num_blocks (int): the number of decoder blocks + static_chunk_size (int): chunk size for static chunk training and + decoding + use_dynamic_chunk (bool): whether use dynamic chunk size for + training or not, You can only use fixed chunk(chunk_size > 0) + or dyanmic chunk size(use_dynamic_chunk = True) + use_dynamic_left_chunk (bool): whether use dynamic left chunk in + dynamic chunk training + key_bias: whether use bias in attention.linear_k, False for whisper models. + """ + + def __init__( + self, + input_size: int = 512, + output_size: int = 512, + attention_heads: int = 8, + linear_units: int = 2048, + num_blocks: int = 6, + static_chunk_size: int = 25, + use_dynamic_chunk: bool = False, + use_dynamic_left_chunk: bool = False, + key_bias: bool = True, + ): + super().__init__() + self._output_size = output_size + + self.embed = LinearNoSubsampling( + input_size, output_size, + EspnetRelPositionalEncoding(output_size), + ) + + self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5) + self.static_chunk_size = static_chunk_size + self.use_dynamic_chunk = use_dynamic_chunk + self.use_dynamic_left_chunk = use_dynamic_left_chunk + activation = torch.nn.SiLU() + # self-attention module definition + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + 0.0, + key_bias, + ) + # feed-forward module definition + positionwise_layer_args = ( + output_size, + linear_units, + 0.0, + activation, + ) + # convolution module definition + self.pre_lookahead_layer = PreLookaheadLayer(channels=512, pre_lookahead_len=3) + self.encoders = torch.nn.ModuleList([ + ConformerEncoderLayer( + output_size, + RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args), + PositionwiseFeedForward(*positionwise_layer_args), + ) for _ in range(num_blocks) + ]) + self.up_layer = Upsample1D(channels=512, out_channels=512, stride=2) + self.up_embed = LinearNoSubsampling( + input_size, output_size, + EspnetRelPositionalEncoding(output_size), + ) + self.up_encoders = torch.nn.ModuleList([ + ConformerEncoderLayer( + output_size, + RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args), + PositionwiseFeedForward(*positionwise_layer_args), + ) for _ in range(4) + ]) + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs: torch.Tensor, + xs_lens: torch.Tensor, + context: torch.Tensor = torch.zeros(0, 0, 0), + decoding_chunk_size: int = 0, + num_decoding_left_chunks: int = -1, + streaming: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Embed positions in tensor. + + Args: + xs: padded input tensor (B, T, D) + xs_lens: input length (B) + decoding_chunk_size: decoding chunk size for dynamic chunk + 0: default for training, use random dynamic chunk. + <0: for decoding, use full chunk. + >0: for decoding, use fixed chunk size as set. + num_decoding_left_chunks: number of left chunks, this is for decoding, + the chunk size is decoding_chunk_size. + >=0: use num_decoding_left_chunks + <0: use all left chunks + Returns: + encoder output tensor xs, and subsampled masks + xs: padded output tensor (B, T' ~= T/subsample_rate, D) + masks: torch.Tensor batch padding mask after subsample + (B, 1, T' ~= T/subsample_rate) + NOTE(xcsong): + We pass the `__call__` method of the modules instead of `forward` to the + checkpointing API because `__call__` attaches all the hooks of the module. + https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2 + """ + T = xs.size(1) + masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) # (B, 1, T) + xs, pos_emb, masks = self.embed(xs, masks) + if context.size(1) != 0: + assert self.training is False, 'you have passed context, make sure that you are running inference mode' + context_masks = torch.ones(1, 1, context.size(1)).to(masks) + context, _, _ = self.embed(context, context_masks, offset=xs.size(1)) + mask_pad = masks # (B, 1, T/subsample_rate) + chunk_masks = add_optional_chunk_mask(xs, masks, False, False, 0, self.static_chunk_size if streaming is True else 0, -1) + # lookahead + conformer encoder + xs = self.pre_lookahead_layer(xs, context=context) + xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad) + + # upsample + conformer encoder + xs = xs.transpose(1, 2).contiguous() + xs, xs_lens = self.up_layer(xs, xs_lens) + xs = xs.transpose(1, 2).contiguous() + T = xs.size(1) + masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) # (B, 1, T) + xs, pos_emb, masks = self.up_embed(xs, masks) + mask_pad = masks # (B, 1, T/subsample_rate) + chunk_masks = add_optional_chunk_mask(xs, masks, False, False, 0, self.static_chunk_size * self.up_layer.stride if streaming is True else 0, -1) + xs = self.forward_up_layers(xs, chunk_masks, pos_emb, mask_pad) + + xs = self.after_norm(xs) + # Here we assume the mask is not changed in encoder layers, so just + # return the masks before encoder layers, and the masks will be used + # for cross attention with decoder later + return xs, masks + + def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor, + pos_emb: torch.Tensor, + mask_pad: torch.Tensor) -> torch.Tensor: + for layer in self.encoders: + xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad) + return xs + + def forward_up_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor, + pos_emb: torch.Tensor, + mask_pad: torch.Tensor) -> torch.Tensor: + for layer in self.up_encoders: + xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad) + return xs diff --git a/soulxpodcast/models/modules/hifigan.py b/soulxpodcast/models/modules/hifigan.py index dd8d435..6284b49 100644 --- a/soulxpodcast/models/modules/hifigan.py +++ b/soulxpodcast/models/modules/hifigan.py @@ -1,249 +1,249 @@ -# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu) -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""HIFI-GAN""" - -from typing import Dict, List - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from scipy.signal import get_window -from torch.nn import Conv1d, ConvTranspose1d -from torch.nn.utils import remove_weight_norm - -try: - from torch.nn.utils.parametrizations import weight_norm -except ImportError: - from torch.nn.utils import weight_norm # noqa - -from soulxpodcast.models.modules.hifigan_components.layers import ( - ResBlock, SourceModuleHnNSF, SourceModuleHnNSF2, init_weights) - - -class ConvRNNF0Predictor(nn.Module): - def __init__(self, - num_class: int = 1, - in_channels: int = 80, - cond_channels: int = 512 - ): - super().__init__() - - self.num_class = num_class - self.condnet = nn.Sequential( - weight_norm( # noqa - nn.Conv1d(in_channels, cond_channels, kernel_size=3, padding=1) - ), - nn.ELU(), - weight_norm( # noqa - nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) - ), - nn.ELU(), - weight_norm( # noqa - nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) - ), - nn.ELU(), - weight_norm( # noqa - nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) - ), - nn.ELU(), - weight_norm( # noqa - nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) - ), - nn.ELU(), - ) - self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.condnet(x) - x = x.transpose(1, 2) - return torch.abs(self.classifier(x).squeeze(-1)) - - -class HiFTGenerator(nn.Module): - """ - HiFTNet Generator: Neural Source Filter + ISTFTNet - https://arxiv.org/abs/2309.09493 - """ - def __init__( - self, - in_channels: int = 80, - base_channels: int = 512, - nb_harmonics: int = 8, - sampling_rate: int = 24000, - nsf_alpha: float = 0.1, - nsf_sigma: float = 0.003, - nsf_voiced_threshold: float = 10, - upsample_rates: List[int] = [8, 5, 3], # noqa - upsample_kernel_sizes: List[int] = [16, 11, 7], # noqa - istft_params: Dict[str, int] = {"n_fft": 16, "hop_len": 4}, # noqa - resblock_kernel_sizes: List[int] = [3, 7, 11], # noqa - resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], # noqa - source_resblock_kernel_sizes: List[int] = [7, 7, 11], # noqa - source_resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], # noqa - lrelu_slope: float = 0.1, - audio_limit: float = 0.99, - f0_predictor: torch.nn.Module = None, - ): - super(HiFTGenerator, self).__init__() - - self.out_channels = 1 - self.nb_harmonics = nb_harmonics - self.sampling_rate = sampling_rate - self.istft_params = istft_params - self.lrelu_slope = lrelu_slope - self.audio_limit = audio_limit - - self.num_kernels = len(resblock_kernel_sizes) - self.num_upsamples = len(upsample_rates) - - this_SourceModuleHnNSF = SourceModuleHnNSF if self.sampling_rate == 22050 else SourceModuleHnNSF2 - self.m_source = this_SourceModuleHnNSF( - sampling_rate=sampling_rate, - upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"], - harmonic_num=nb_harmonics, - sine_amp=nsf_alpha, - add_noise_std=nsf_sigma, - voiced_threshod=nsf_voiced_threshold) - self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"]) - - self.conv_pre = weight_norm( # noqa - Conv1d(in_channels, base_channels, 7, 1, padding=3) - ) - - # Up - self.ups = nn.ModuleList() - for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): - self.ups.append( - weight_norm( # noqa - ConvTranspose1d( - base_channels // (2**i), - base_channels // (2**(i + 1)), - k, - u, - padding=(k - u) // 2, - ) - ) - ) - - # Down - self.source_downs = nn.ModuleList() - self.source_resblocks = nn.ModuleList() - downsample_rates = [1] + upsample_rates[::-1][:-1] - downsample_cum_rates = np.cumprod(downsample_rates) - for i, (u, k, d) in enumerate(zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes)): - if u == 1: - self.source_downs.append( - Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1) - ) - else: - self.source_downs.append( - Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u, padding=(u // 2)) - ) - - self.source_resblocks.append( - ResBlock(base_channels // (2 ** (i + 1)), k, d) - ) - - self.resblocks = nn.ModuleList() - for i in range(len(self.ups)): - ch = base_channels // (2**(i + 1)) - for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)): - self.resblocks.append(ResBlock(ch, k, d)) - - self.conv_post = weight_norm(Conv1d(ch, istft_params["n_fft"] + 2, 7, 1, padding=3)) # noqa - self.ups.apply(init_weights) - self.conv_post.apply(init_weights) - self.reflection_pad = nn.ReflectionPad1d((1, 0)) - self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)) - self.f0_predictor = ConvRNNF0Predictor() if f0_predictor is None else f0_predictor - - def remove_weight_norm(self): - print('Removing weight norm...') - for up in self.ups: - remove_weight_norm(up) - for resblock in self.resblocks: - resblock.remove_weight_norm() - remove_weight_norm(self.conv_pre) - remove_weight_norm(self.conv_post) - self.m_source.remove_weight_norm() - for source_down in self.source_downs: - remove_weight_norm(source_down) - for source_resblock in self.source_resblocks: - source_resblock.remove_weight_norm() - - def _stft(self, x): - spec = torch.stft( - x, - self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], window=self.stft_window.to(x.device), - return_complex=True) - spec = torch.view_as_real(spec) # [B, F, TT, 2] - return spec[..., 0], spec[..., 1] - - def _istft(self, magnitude, phase): - magnitude = torch.clip(magnitude, max=1e2) - real = magnitude * torch.cos(phase) - img = magnitude * torch.sin(phase) - inverse_transform = torch.istft(torch.complex(real, img), self.istft_params["n_fft"], self.istft_params["hop_len"], - self.istft_params["n_fft"], window=self.stft_window.to(magnitude.device)) - return inverse_transform - - def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor: - s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) - s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) - - x = self.conv_pre(x) - for i in range(self.num_upsamples): - x = F.leaky_relu(x, self.lrelu_slope) - x = self.ups[i](x) - - if i == self.num_upsamples - 1: - x = self.reflection_pad(x) - - # fusion - si = self.source_downs[i](s_stft) - si = self.source_resblocks[i](si) - x = x + si - - xs = None - for j in range(self.num_kernels): - if xs is None: - xs = self.resblocks[i * self.num_kernels + j](x) - else: - xs += self.resblocks[i * self.num_kernels + j](x) - x = xs / self.num_kernels - - x = F.leaky_relu(x) - x = self.conv_post(x) - magnitude = torch.exp(x[:, :self.istft_params["n_fft"] // 2 + 1, :]) - phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1:, :]) # actually, sin is redundancy - - x = self._istft(magnitude, phase) - x = torch.clamp(x, -self.audio_limit, self.audio_limit) - return x - - @torch.inference_mode() - def forward(self, speech_feat: torch.Tensor, cache_source: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor: - # mel->f0 - f0 = self.f0_predictor(speech_feat) - # f0->source - s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t - s, _, _ = self.m_source(s) - s = s.transpose(1, 2) - # use cache_source to avoid glitch - if cache_source.shape[2] != 0: - s[:, :, :cache_source.shape[2]] = cache_source - generated_speech = self.decode(x=speech_feat, s=s) - return generated_speech, s +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HIFI-GAN""" + +from typing import Dict, List + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from scipy.signal import get_window +from torch.nn import Conv1d, ConvTranspose1d +from torch.nn.utils import remove_weight_norm + +try: + from torch.nn.utils.parametrizations import weight_norm +except ImportError: + from torch.nn.utils import weight_norm # noqa + +from soulxpodcast.models.modules.hifigan_components.layers import ( + ResBlock, SourceModuleHnNSF, SourceModuleHnNSF2, init_weights) + + +class ConvRNNF0Predictor(nn.Module): + def __init__(self, + num_class: int = 1, + in_channels: int = 80, + cond_channels: int = 512 + ): + super().__init__() + + self.num_class = num_class + self.condnet = nn.Sequential( + weight_norm( # noqa + nn.Conv1d(in_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + weight_norm( # noqa + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + weight_norm( # noqa + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + weight_norm( # noqa + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + weight_norm( # noqa + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + ) + self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.condnet(x) + x = x.transpose(1, 2) + return torch.abs(self.classifier(x).squeeze(-1)) + + +class HiFTGenerator(nn.Module): + """ + HiFTNet Generator: Neural Source Filter + ISTFTNet + https://arxiv.org/abs/2309.09493 + """ + def __init__( + self, + in_channels: int = 80, + base_channels: int = 512, + nb_harmonics: int = 8, + sampling_rate: int = 24000, + nsf_alpha: float = 0.1, + nsf_sigma: float = 0.003, + nsf_voiced_threshold: float = 10, + upsample_rates: List[int] = [8, 5, 3], # noqa + upsample_kernel_sizes: List[int] = [16, 11, 7], # noqa + istft_params: Dict[str, int] = {"n_fft": 16, "hop_len": 4}, # noqa + resblock_kernel_sizes: List[int] = [3, 7, 11], # noqa + resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], # noqa + source_resblock_kernel_sizes: List[int] = [7, 7, 11], # noqa + source_resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], # noqa + lrelu_slope: float = 0.1, + audio_limit: float = 0.99, + f0_predictor: torch.nn.Module = None, + ): + super(HiFTGenerator, self).__init__() + + self.out_channels = 1 + self.nb_harmonics = nb_harmonics + self.sampling_rate = sampling_rate + self.istft_params = istft_params + self.lrelu_slope = lrelu_slope + self.audio_limit = audio_limit + + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + + this_SourceModuleHnNSF = SourceModuleHnNSF if self.sampling_rate == 22050 else SourceModuleHnNSF2 + self.m_source = this_SourceModuleHnNSF( + sampling_rate=sampling_rate, + upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"], + harmonic_num=nb_harmonics, + sine_amp=nsf_alpha, + add_noise_std=nsf_sigma, + voiced_threshod=nsf_voiced_threshold) + self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"]) + + self.conv_pre = weight_norm( # noqa + Conv1d(in_channels, base_channels, 7, 1, padding=3) + ) + + # Up + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + weight_norm( # noqa + ConvTranspose1d( + base_channels // (2**i), + base_channels // (2**(i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ) + + # Down + self.source_downs = nn.ModuleList() + self.source_resblocks = nn.ModuleList() + downsample_rates = [1] + upsample_rates[::-1][:-1] + downsample_cum_rates = np.cumprod(downsample_rates) + for i, (u, k, d) in enumerate(zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes)): + if u == 1: + self.source_downs.append( + Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1) + ) + else: + self.source_downs.append( + Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u, padding=(u // 2)) + ) + + self.source_resblocks.append( + ResBlock(base_channels // (2 ** (i + 1)), k, d) + ) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = base_channels // (2**(i + 1)) + for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)): + self.resblocks.append(ResBlock(ch, k, d)) + + self.conv_post = weight_norm(Conv1d(ch, istft_params["n_fft"] + 2, 7, 1, padding=3)) # noqa + self.ups.apply(init_weights) + self.conv_post.apply(init_weights) + self.reflection_pad = nn.ReflectionPad1d((1, 0)) + self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)) + self.f0_predictor = ConvRNNF0Predictor() if f0_predictor is None else f0_predictor + + def remove_weight_norm(self): + print('Removing weight norm...') + for up in self.ups: + remove_weight_norm(up) + for resblock in self.resblocks: + resblock.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + self.m_source.remove_weight_norm() + for source_down in self.source_downs: + remove_weight_norm(source_down) + for source_resblock in self.source_resblocks: + source_resblock.remove_weight_norm() + + def _stft(self, x): + spec = torch.stft( + x, + self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], window=self.stft_window.to(x.device), + return_complex=True) + spec = torch.view_as_real(spec) # [B, F, TT, 2] + return spec[..., 0], spec[..., 1] + + def _istft(self, magnitude, phase): + magnitude = torch.clip(magnitude, max=1e2) + real = magnitude * torch.cos(phase) + img = magnitude * torch.sin(phase) + inverse_transform = torch.istft(torch.complex(real, img), self.istft_params["n_fft"], self.istft_params["hop_len"], + self.istft_params["n_fft"], window=self.stft_window.to(magnitude.device)) + return inverse_transform + + def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor: + s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) + s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) + + x = self.conv_pre(x) + for i in range(self.num_upsamples): + x = F.leaky_relu(x, self.lrelu_slope) + x = self.ups[i](x) + + if i == self.num_upsamples - 1: + x = self.reflection_pad(x) + + # fusion + si = self.source_downs[i](s_stft) + si = self.source_resblocks[i](si) + x = x + si + + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + x = F.leaky_relu(x) + x = self.conv_post(x) + magnitude = torch.exp(x[:, :self.istft_params["n_fft"] // 2 + 1, :]) + phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1:, :]) # actually, sin is redundancy + + x = self._istft(magnitude, phase) + x = torch.clamp(x, -self.audio_limit, self.audio_limit) + return x + + @torch.inference_mode() + def forward(self, speech_feat: torch.Tensor, cache_source: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor: + # mel->f0 + f0 = self.f0_predictor(speech_feat) + # f0->source + s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t + s, _, _ = self.m_source(s) + s = s.transpose(1, 2) + # use cache_source to avoid glitch + if cache_source.shape[2] != 0: + s[:, :, :cache_source.shape[2]] = cache_source + generated_speech = self.decode(x=speech_feat, s=s) + return generated_speech, s diff --git a/soulxpodcast/models/modules/hifigan_components/layers.py b/soulxpodcast/models/modules/hifigan_components/layers.py index e07c587..8ce9819 100644 --- a/soulxpodcast/models/modules/hifigan_components/layers.py +++ b/soulxpodcast/models/modules/hifigan_components/layers.py @@ -1,433 +1,433 @@ -from typing import List - -import numpy as np -import torch -import torch.nn as nn -from torch.distributions.uniform import Uniform -from torch.nn import Conv1d -from torch.nn.utils import remove_weight_norm - -try: - from torch.nn.utils.parametrizations import weight_norm -except ImportError: - from torch.nn.utils import weight_norm # noqa - - -def get_padding(kernel_size, dilation=1): - return int((kernel_size * dilation - dilation) / 2) - - -def init_weights(m, mean=0.0, std=0.01): - classname = m.__class__.__name__ - if classname.find("Conv") != -1: - m.weight.data.normal_(mean, std) - - -"""hifigan based generator implementation. - -This code is modified from https://github.com/jik876/hifi-gan - ,https://github.com/kan-bayashi/ParallelWaveGAN and - https://github.com/NVIDIA/BigVGAN - -""" - - -# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. -# LICENSE is in incl_licenses directory. -class Snake(nn.Module): - ''' - Implementation of a sine-based periodic activation function - Shape: - - Input: (B, C, T) - - Output: (B, C, T), same shape as the input - Parameters: - - alpha - trainable parameter - References: - - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: - https://arxiv.org/abs/2006.08195 - Examples: - >>> a1 = snake(256) - >>> x = torch.randn(256) - >>> x = a1(x) - - Args: - in_features: shape of the input - alpha: trainable parameter - alpha_trainable: whether alpha is trainable - alpha_logscale: whether to use log scale for alpha - alpha is initialized to 1 by default, higher values = higher-frequency. - alpha will be trained along with the rest of your model. - ''' - def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): - super(Snake, self).__init__() - self.in_features = in_features - - # initialize alpha - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: # log scale alphas initialized to zeros - self.alpha = nn.Parameter(torch.zeros(in_features) * alpha) - else: # linear scale alphas initialized to ones - self.alpha = nn.Parameter(torch.ones(in_features) * alpha) - - self.alpha.requires_grad = alpha_trainable - - self.no_div_by_zero = 0.000000001 - - def forward(self, x): - ''' - Forward pass of the function. - Applies the function to the input elementwise. - Snake ∶= x + 1/a * sin^2 (xa) - ''' - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] - if self.alpha_logscale: - alpha = torch.exp(alpha) - x = x + (1.0 / (alpha + self.no_div_by_zero)) * torch.pow(torch.sin(x * alpha), 2) - - return x - - -class ResBlock(torch.nn.Module): - """Residual block module in HiFiGAN/BigVGAN.""" - def __init__( - self, - channels: int = 512, - kernel_size: int = 3, - dilations: List[int] = [1, 3, 5], # noqa - ): - super(ResBlock, self).__init__() - self.convs1 = nn.ModuleList() - self.convs2 = nn.ModuleList() - - for dilation in dilations: - self.convs1.append( - weight_norm( # noqa - Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation, - padding=get_padding(kernel_size, dilation) - ) - ) - ) - self.convs2.append( - weight_norm( # noqa - Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=1, - padding=get_padding(kernel_size, 1) - ) - ) - ) - self.convs1.apply(init_weights) - self.convs2.apply(init_weights) - self.activations1 = nn.ModuleList([ - Snake(channels, alpha_logscale=False) - for _ in range(len(self.convs1)) - ]) - self.activations2 = nn.ModuleList([ - Snake(channels, alpha_logscale=False) - for _ in range(len(self.convs2)) - ]) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - for idx in range(len(self.convs1)): - xt = self.activations1[idx](x) - xt = self.convs1[idx](xt) - xt = self.activations2[idx](xt) - xt = self.convs2[idx](xt) - x = xt + x - return x - - def remove_weight_norm(self): - for idx in range(len(self.convs1)): - remove_weight_norm(self.convs1[idx]) - remove_weight_norm(self.convs2[idx]) - - -class SineGen(torch.nn.Module): - """ Definition of sine generator - SineGen(samp_rate, harmonic_num = 0, - sine_amp = 0.1, noise_std = 0.003, - voiced_threshold = 0, - flag_for_pulse=False) - samp_rate: sampling rate in Hz - harmonic_num: number of harmonic overtones (default 0) - sine_amp: amplitude of sine-wavefrom (default 0.1) - noise_std: std of Gaussian noise (default 0.003) - voiced_thoreshold: F0 threshold for U/V classification (default 0) - flag_for_pulse: this SinGen is used inside PulseGen (default False) - Note: when flag_for_pulse is True, the first time step of a voiced - segment is always sin(np.pi) or cos(0) - """ - - def __init__(self, samp_rate, harmonic_num=0, - sine_amp=0.1, noise_std=0.003, - voiced_threshold=0): - super(SineGen, self).__init__() - self.sine_amp = sine_amp - self.noise_std = noise_std - self.harmonic_num = harmonic_num - self.sampling_rate = samp_rate - self.voiced_threshold = voiced_threshold - - def _f02uv(self, f0): - # generate uv signal - uv = (f0 > self.voiced_threshold).type(torch.float32) - return uv - - @torch.no_grad() - def forward(self, f0): - """ - :param f0: [B, 1, sample_len], Hz - :return: [B, 1, sample_len] - """ - - F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device) - for i in range(self.harmonic_num + 1): - F_mat[:, i: i + 1, :] = f0 * (i + 1) / self.sampling_rate - - theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1) - u_dist = Uniform(low=-np.pi, high=np.pi) - phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) - phase_vec[:, 0, :] = 0 - - # generate sine waveforms - sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec) - - # generate uv signal - uv = self._f02uv(f0) - - # noise: for unvoiced should be similar to sine_amp - # std = self.sine_amp/3 -> max value ~ self.sine_amp - # . for voiced regions is self.noise_std - noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 - noise = noise_amp * torch.randn_like(sine_waves) - - # first: set the unvoiced part to 0 by uv - # then: additive noise - sine_waves = sine_waves * uv + noise - return sine_waves, uv, noise - - -class SourceModuleHnNSF(torch.nn.Module): - """ SourceModule for hn-nsf - SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1, - add_noise_std=0.003, voiced_threshod=0) - sampling_rate: sampling_rate in Hz - harmonic_num: number of harmonic above F0 (default: 0) - sine_amp: amplitude of sine source signal (default: 0.1) - add_noise_std: std of additive Gaussian noise (default: 0.003) - note that amplitude of noise in unvoiced is decided - by sine_amp - voiced_threshold: threhold to set U/V given F0 (default: 0) - Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) - F0_sampled (batchsize, length, 1) - Sine_source (batchsize, length, 1) - noise_source (batchsize, length 1) - uv (batchsize, length, 1) - """ - - def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1, - add_noise_std=0.003, voiced_threshod=0): - super(SourceModuleHnNSF, self).__init__() - - self.sine_amp = sine_amp - self.noise_std = add_noise_std - - # to produce sine waveforms - self.l_sin_gen = SineGen(sampling_rate, harmonic_num, - sine_amp, add_noise_std, voiced_threshod) - - # to merge source harmonics into a single excitation - self.l_linear = torch.nn.Linear(harmonic_num + 1, 1) - self.l_tanh = torch.nn.Tanh() - - def forward(self, x): - """ - Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) - F0_sampled (batchsize, length, 1) - Sine_source (batchsize, length, 1) - noise_source (batchsize, length 1) - """ - # source for harmonic branch - with torch.no_grad(): - sine_wavs, uv, _ = self.l_sin_gen(x.transpose(1, 2)) - sine_wavs = sine_wavs.transpose(1, 2) - uv = uv.transpose(1, 2) - sine_merge = self.l_tanh(self.l_linear(sine_wavs)) - - # source for noise branch, in the same shape as uv - noise = torch.randn_like(uv) * self.sine_amp / 3 - return sine_merge, noise, uv - - -class SineGen2(torch.nn.Module): - """ Definition of sine generator - SineGen(samp_rate, harmonic_num = 0, - sine_amp = 0.1, noise_std = 0.003, - voiced_threshold = 0, - flag_for_pulse=False) - samp_rate: sampling rate in Hz - harmonic_num: number of harmonic overtones (default 0) - sine_amp: amplitude of sine-wavefrom (default 0.1) - noise_std: std of Gaussian noise (default 0.003) - voiced_thoreshold: F0 threshold for U/V classification (default 0) - flag_for_pulse: this SinGen is used inside PulseGen (default False) - Note: when flag_for_pulse is True, the first time step of a voiced - segment is always sin(np.pi) or cos(0) - """ - - def __init__(self, samp_rate, upsample_scale, harmonic_num=0, - sine_amp=0.1, noise_std=0.003, - voiced_threshold=0, - flag_for_pulse=False): - super(SineGen2, self).__init__() - self.sine_amp = sine_amp - self.noise_std = noise_std - self.harmonic_num = harmonic_num - self.dim = self.harmonic_num + 1 - self.sampling_rate = samp_rate - self.voiced_threshold = voiced_threshold - self.flag_for_pulse = flag_for_pulse - self.upsample_scale = upsample_scale - - def _f02uv(self, f0): - # generate uv signal - uv = (f0 > self.voiced_threshold).type(torch.float32) - return uv - - def _f02sine(self, f0_values): - """ f0_values: (batchsize, length, dim) - where dim indicates fundamental tone and overtones - """ - # convert to F0 in rad. The interger part n can be ignored - # because 2 * np.pi * n doesn't affect phase - rad_values = (f0_values / self.sampling_rate) % 1 - - # initial phase noise (no noise for fundamental component) - rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) - rand_ini[:, 0] = 0 - rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini - - # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad) - if not self.flag_for_pulse: - rad_values = torch.nn.functional.interpolate(rad_values.transpose(1, 2), - scale_factor=1 / self.upsample_scale, - mode="linear").transpose(1, 2) - - phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi - phase = torch.nn.functional.interpolate(phase.transpose(1, 2) * self.upsample_scale, - scale_factor=self.upsample_scale, mode="linear").transpose(1, 2) - sines = torch.sin(phase) - else: - # If necessary, make sure that the first time step of every - # voiced segments is sin(pi) or cos(0) - # This is used for pulse-train generation - - # identify the last time step in unvoiced segments - uv = self._f02uv(f0_values) - uv_1 = torch.roll(uv, shifts=-1, dims=1) - uv_1[:, -1, :] = 1 - u_loc = (uv < 1) * (uv_1 > 0) - - # get the instantanouse phase - tmp_cumsum = torch.cumsum(rad_values, dim=1) - # different batch needs to be processed differently - for idx in range(f0_values.shape[0]): - temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :] - temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :] - # stores the accumulation of i.phase within - # each voiced segments - tmp_cumsum[idx, :, :] = 0 - tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum - - # rad_values - tmp_cumsum: remove the accumulation of i.phase - # within the previous voiced segment. - i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1) - - # get the sines - sines = torch.cos(i_phase * 2 * np.pi) - return sines - - def forward(self, f0): - """ sine_tensor, uv = forward(f0) - input F0: tensor(batchsize=1, length, dim=1) - f0 for unvoiced steps should be 0 - output sine_tensor: tensor(batchsize=1, length, dim) - output uv: tensor(batchsize=1, length, 1) - """ - # fundamental component - fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device)) - - # generate sine waveforms - sine_waves = self._f02sine(fn) * self.sine_amp - - # generate uv signal - uv = self._f02uv(f0) - - # noise: for unvoiced should be similar to sine_amp - # std = self.sine_amp/3 -> max value ~ self.sine_amp - # . for voiced regions is self.noise_std - noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 - noise = noise_amp * torch.randn_like(sine_waves) - - # first: set the unvoiced part to 0 by uv - # then: additive noise - sine_waves = sine_waves * uv + noise - return sine_waves, uv, noise - - -class SourceModuleHnNSF2(torch.nn.Module): - """ SourceModule for hn-nsf - SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1, - add_noise_std=0.003, voiced_threshod=0) - sampling_rate: sampling_rate in Hz - harmonic_num: number of harmonic above F0 (default: 0) - sine_amp: amplitude of sine source signal (default: 0.1) - add_noise_std: std of additive Gaussian noise (default: 0.003) - note that amplitude of noise in unvoiced is decided - by sine_amp - voiced_threshold: threhold to set U/V given F0 (default: 0) - Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) - F0_sampled (batchsize, length, 1) - Sine_source (batchsize, length, 1) - noise_source (batchsize, length 1) - uv (batchsize, length, 1) - """ - - def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1, - add_noise_std=0.003, voiced_threshod=0): - super(SourceModuleHnNSF2, self).__init__() - - self.sine_amp = sine_amp - self.noise_std = add_noise_std - - # to produce sine waveforms - self.l_sin_gen = SineGen2(sampling_rate, upsample_scale, harmonic_num, - sine_amp, add_noise_std, voiced_threshod) - - # to merge source harmonics into a single excitation - self.l_linear = torch.nn.Linear(harmonic_num + 1, 1) - self.l_tanh = torch.nn.Tanh() - - def forward(self, x): - """ - Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) - F0_sampled (batchsize, length, 1) - Sine_source (batchsize, length, 1) - noise_source (batchsize, length 1) - """ - # source for harmonic branch - with torch.no_grad(): - sine_wavs, uv, _ = self.l_sin_gen(x) - sine_merge = self.l_tanh(self.l_linear(sine_wavs)) - - # source for noise branch, in the same shape as uv - noise = torch.randn_like(uv) * self.sine_amp / 3 - return sine_merge, noise, uv +from typing import List + +import numpy as np +import torch +import torch.nn as nn +from torch.distributions.uniform import Uniform +from torch.nn import Conv1d +from torch.nn.utils import remove_weight_norm + +try: + from torch.nn.utils.parametrizations import weight_norm +except ImportError: + from torch.nn.utils import weight_norm # noqa + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +"""hifigan based generator implementation. + +This code is modified from https://github.com/jik876/hifi-gan + ,https://github.com/kan-bayashi/ParallelWaveGAN and + https://github.com/NVIDIA/BigVGAN + +""" + + +# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. +# LICENSE is in incl_licenses directory. +class Snake(nn.Module): + ''' + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + References: + - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snake(256) + >>> x = torch.randn(256) + >>> x = a1(x) + + Args: + in_features: shape of the input + alpha: trainable parameter + alpha_trainable: whether alpha is trainable + alpha_logscale: whether to use log scale for alpha + alpha is initialized to 1 by default, higher values = higher-frequency. + alpha will be trained along with the rest of your model. + ''' + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + super(Snake, self).__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = nn.Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = nn.Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + ''' + Forward pass of the function. + Applies the function to the input elementwise. + Snake ∶= x + 1/a * sin^2 (xa) + ''' + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + x = x + (1.0 / (alpha + self.no_div_by_zero)) * torch.pow(torch.sin(x * alpha), 2) + + return x + + +class ResBlock(torch.nn.Module): + """Residual block module in HiFiGAN/BigVGAN.""" + def __init__( + self, + channels: int = 512, + kernel_size: int = 3, + dilations: List[int] = [1, 3, 5], # noqa + ): + super(ResBlock, self).__init__() + self.convs1 = nn.ModuleList() + self.convs2 = nn.ModuleList() + + for dilation in dilations: + self.convs1.append( + weight_norm( # noqa + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation, + padding=get_padding(kernel_size, dilation) + ) + ) + ) + self.convs2.append( + weight_norm( # noqa + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1) + ) + ) + ) + self.convs1.apply(init_weights) + self.convs2.apply(init_weights) + self.activations1 = nn.ModuleList([ + Snake(channels, alpha_logscale=False) + for _ in range(len(self.convs1)) + ]) + self.activations2 = nn.ModuleList([ + Snake(channels, alpha_logscale=False) + for _ in range(len(self.convs2)) + ]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for idx in range(len(self.convs1)): + xt = self.activations1[idx](x) + xt = self.convs1[idx](xt) + xt = self.activations2[idx](xt) + xt = self.convs2[idx](xt) + x = xt + x + return x + + def remove_weight_norm(self): + for idx in range(len(self.convs1)): + remove_weight_norm(self.convs1[idx]) + remove_weight_norm(self.convs2[idx]) + + +class SineGen(torch.nn.Module): + """ Definition of sine generator + SineGen(samp_rate, harmonic_num = 0, + sine_amp = 0.1, noise_std = 0.003, + voiced_threshold = 0, + flag_for_pulse=False) + samp_rate: sampling rate in Hz + harmonic_num: number of harmonic overtones (default 0) + sine_amp: amplitude of sine-wavefrom (default 0.1) + noise_std: std of Gaussian noise (default 0.003) + voiced_thoreshold: F0 threshold for U/V classification (default 0) + flag_for_pulse: this SinGen is used inside PulseGen (default False) + Note: when flag_for_pulse is True, the first time step of a voiced + segment is always sin(np.pi) or cos(0) + """ + + def __init__(self, samp_rate, harmonic_num=0, + sine_amp=0.1, noise_std=0.003, + voiced_threshold=0): + super(SineGen, self).__init__() + self.sine_amp = sine_amp + self.noise_std = noise_std + self.harmonic_num = harmonic_num + self.sampling_rate = samp_rate + self.voiced_threshold = voiced_threshold + + def _f02uv(self, f0): + # generate uv signal + uv = (f0 > self.voiced_threshold).type(torch.float32) + return uv + + @torch.no_grad() + def forward(self, f0): + """ + :param f0: [B, 1, sample_len], Hz + :return: [B, 1, sample_len] + """ + + F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device) + for i in range(self.harmonic_num + 1): + F_mat[:, i: i + 1, :] = f0 * (i + 1) / self.sampling_rate + + theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1) + u_dist = Uniform(low=-np.pi, high=np.pi) + phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) + phase_vec[:, 0, :] = 0 + + # generate sine waveforms + sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec) + + # generate uv signal + uv = self._f02uv(f0) + + # noise: for unvoiced should be similar to sine_amp + # std = self.sine_amp/3 -> max value ~ self.sine_amp + # . for voiced regions is self.noise_std + noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 + noise = noise_amp * torch.randn_like(sine_waves) + + # first: set the unvoiced part to 0 by uv + # then: additive noise + sine_waves = sine_waves * uv + noise + return sine_waves, uv, noise + + +class SourceModuleHnNSF(torch.nn.Module): + """ SourceModule for hn-nsf + SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1, + add_noise_std=0.003, voiced_threshod=0) + sampling_rate: sampling_rate in Hz + harmonic_num: number of harmonic above F0 (default: 0) + sine_amp: amplitude of sine source signal (default: 0.1) + add_noise_std: std of additive Gaussian noise (default: 0.003) + note that amplitude of noise in unvoiced is decided + by sine_amp + voiced_threshold: threhold to set U/V given F0 (default: 0) + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + uv (batchsize, length, 1) + """ + + def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1, + add_noise_std=0.003, voiced_threshod=0): + super(SourceModuleHnNSF, self).__init__() + + self.sine_amp = sine_amp + self.noise_std = add_noise_std + + # to produce sine waveforms + self.l_sin_gen = SineGen(sampling_rate, harmonic_num, + sine_amp, add_noise_std, voiced_threshod) + + # to merge source harmonics into a single excitation + self.l_linear = torch.nn.Linear(harmonic_num + 1, 1) + self.l_tanh = torch.nn.Tanh() + + def forward(self, x): + """ + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + """ + # source for harmonic branch + with torch.no_grad(): + sine_wavs, uv, _ = self.l_sin_gen(x.transpose(1, 2)) + sine_wavs = sine_wavs.transpose(1, 2) + uv = uv.transpose(1, 2) + sine_merge = self.l_tanh(self.l_linear(sine_wavs)) + + # source for noise branch, in the same shape as uv + noise = torch.randn_like(uv) * self.sine_amp / 3 + return sine_merge, noise, uv + + +class SineGen2(torch.nn.Module): + """ Definition of sine generator + SineGen(samp_rate, harmonic_num = 0, + sine_amp = 0.1, noise_std = 0.003, + voiced_threshold = 0, + flag_for_pulse=False) + samp_rate: sampling rate in Hz + harmonic_num: number of harmonic overtones (default 0) + sine_amp: amplitude of sine-wavefrom (default 0.1) + noise_std: std of Gaussian noise (default 0.003) + voiced_thoreshold: F0 threshold for U/V classification (default 0) + flag_for_pulse: this SinGen is used inside PulseGen (default False) + Note: when flag_for_pulse is True, the first time step of a voiced + segment is always sin(np.pi) or cos(0) + """ + + def __init__(self, samp_rate, upsample_scale, harmonic_num=0, + sine_amp=0.1, noise_std=0.003, + voiced_threshold=0, + flag_for_pulse=False): + super(SineGen2, self).__init__() + self.sine_amp = sine_amp + self.noise_std = noise_std + self.harmonic_num = harmonic_num + self.dim = self.harmonic_num + 1 + self.sampling_rate = samp_rate + self.voiced_threshold = voiced_threshold + self.flag_for_pulse = flag_for_pulse + self.upsample_scale = upsample_scale + + def _f02uv(self, f0): + # generate uv signal + uv = (f0 > self.voiced_threshold).type(torch.float32) + return uv + + def _f02sine(self, f0_values): + """ f0_values: (batchsize, length, dim) + where dim indicates fundamental tone and overtones + """ + # convert to F0 in rad. The interger part n can be ignored + # because 2 * np.pi * n doesn't affect phase + rad_values = (f0_values / self.sampling_rate) % 1 + + # initial phase noise (no noise for fundamental component) + rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) + rand_ini[:, 0] = 0 + rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini + + # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad) + if not self.flag_for_pulse: + rad_values = torch.nn.functional.interpolate(rad_values.transpose(1, 2), + scale_factor=1 / self.upsample_scale, + mode="linear").transpose(1, 2) + + phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi + phase = torch.nn.functional.interpolate(phase.transpose(1, 2) * self.upsample_scale, + scale_factor=self.upsample_scale, mode="linear").transpose(1, 2) + sines = torch.sin(phase) + else: + # If necessary, make sure that the first time step of every + # voiced segments is sin(pi) or cos(0) + # This is used for pulse-train generation + + # identify the last time step in unvoiced segments + uv = self._f02uv(f0_values) + uv_1 = torch.roll(uv, shifts=-1, dims=1) + uv_1[:, -1, :] = 1 + u_loc = (uv < 1) * (uv_1 > 0) + + # get the instantanouse phase + tmp_cumsum = torch.cumsum(rad_values, dim=1) + # different batch needs to be processed differently + for idx in range(f0_values.shape[0]): + temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :] + temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :] + # stores the accumulation of i.phase within + # each voiced segments + tmp_cumsum[idx, :, :] = 0 + tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum + + # rad_values - tmp_cumsum: remove the accumulation of i.phase + # within the previous voiced segment. + i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1) + + # get the sines + sines = torch.cos(i_phase * 2 * np.pi) + return sines + + def forward(self, f0): + """ sine_tensor, uv = forward(f0) + input F0: tensor(batchsize=1, length, dim=1) + f0 for unvoiced steps should be 0 + output sine_tensor: tensor(batchsize=1, length, dim) + output uv: tensor(batchsize=1, length, 1) + """ + # fundamental component + fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device)) + + # generate sine waveforms + sine_waves = self._f02sine(fn) * self.sine_amp + + # generate uv signal + uv = self._f02uv(f0) + + # noise: for unvoiced should be similar to sine_amp + # std = self.sine_amp/3 -> max value ~ self.sine_amp + # . for voiced regions is self.noise_std + noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 + noise = noise_amp * torch.randn_like(sine_waves) + + # first: set the unvoiced part to 0 by uv + # then: additive noise + sine_waves = sine_waves * uv + noise + return sine_waves, uv, noise + + +class SourceModuleHnNSF2(torch.nn.Module): + """ SourceModule for hn-nsf + SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1, + add_noise_std=0.003, voiced_threshod=0) + sampling_rate: sampling_rate in Hz + harmonic_num: number of harmonic above F0 (default: 0) + sine_amp: amplitude of sine source signal (default: 0.1) + add_noise_std: std of additive Gaussian noise (default: 0.003) + note that amplitude of noise in unvoiced is decided + by sine_amp + voiced_threshold: threhold to set U/V given F0 (default: 0) + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + uv (batchsize, length, 1) + """ + + def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1, + add_noise_std=0.003, voiced_threshod=0): + super(SourceModuleHnNSF2, self).__init__() + + self.sine_amp = sine_amp + self.noise_std = add_noise_std + + # to produce sine waveforms + self.l_sin_gen = SineGen2(sampling_rate, upsample_scale, harmonic_num, + sine_amp, add_noise_std, voiced_threshod) + + # to merge source harmonics into a single excitation + self.l_linear = torch.nn.Linear(harmonic_num + 1, 1) + self.l_tanh = torch.nn.Tanh() + + def forward(self, x): + """ + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + """ + # source for harmonic branch + with torch.no_grad(): + sine_wavs, uv, _ = self.l_sin_gen(x) + sine_merge = self.l_tanh(self.l_linear(sine_wavs)) + + # source for noise branch, in the same shape as uv + noise = torch.randn_like(uv) * self.sine_amp / 3 + return sine_merge, noise, uv diff --git a/soulxpodcast/models/modules/sampler.py b/soulxpodcast/models/modules/sampler.py index 28c48c7..d303ac1 100644 --- a/soulxpodcast/models/modules/sampler.py +++ b/soulxpodcast/models/modules/sampler.py @@ -1,221 +1,221 @@ -import os - -from typing import Any, Callable, Optional, Union -import torch -from torch import nn -from transformers.generation.logits_process import ( - LogitsProcessorList -) -from transformers.generation.stopping_criteria import ( - StoppingCriteriaList -) -from transformers.generation.configuration_utils import ( - GenerationConfig -) -from transformers.generation.streamers import BaseStreamer -from transformers.generation.utils import ( - GenerateNonBeamOutput, - GenerateEncoderDecoderOutput, - GenerateDecoderOnlyOutput, -) -from transformers import StoppingCriteria - - -def _ras_sample_hf_engine( - self, - input_ids: torch.LongTensor, - logits_processor: LogitsProcessorList, - stopping_criteria: StoppingCriteriaList, - generation_config: GenerationConfig, - synced_gpus: bool = False, - streamer: Optional["BaseStreamer"] = None, - use_ras=False, - win_size=25, - tau_r=0.2, - **model_kwargs, -) -> Union[GenerateNonBeamOutput, torch.LongTensor]: - r""" - Generates sequences of token ids for models with a language modeling head using **multinomial sampling** and - can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. - - Parameters: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - The sequence used as a prompt for the generation. - logits_processor (`LogitsProcessorList`): - An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] - used to modify the prediction scores of the language modeling head applied at each generation step. - stopping_criteria (`StoppingCriteriaList`): - An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] - used to tell if the generation loop should stop. - generation_config ([`~generation.GenerationConfig`]): - The generation configuration to be used as parametrization of the decoding method. - synced_gpus (`bool`): - Whether to continue running the while loop until max_length (needed to avoid deadlocking with - `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). - streamer (`BaseStreamer`, *optional*): - Streamer object that will be used to stream the generated sequences. Generated tokens are passed - through `streamer.put(token_ids)` and the streamer is responsible for any further processing. - model_kwargs: - Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is - an encoder-decoder model the kwargs should include `encoder_outputs`. - - Return: - [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or `torch.LongTensor`: - A `torch.LongTensor` containing the generated tokens (default behaviour) or a - [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and - `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if - `model.config.is_encoder_decoder=True`. - """ - # init values - pad_token_id = generation_config._pad_token_tensor - output_attentions = generation_config.output_attentions - output_hidden_states = generation_config.output_hidden_states - output_scores = generation_config.output_scores - output_logits = generation_config.output_logits - return_dict_in_generate = generation_config.return_dict_in_generate - has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria) - do_sample = generation_config.do_sample - - # init attention / hidden states / scores tuples - scores = () if (return_dict_in_generate and output_scores) else None - raw_logits = () if (return_dict_in_generate and output_logits) else None - decoder_attentions = () if (return_dict_in_generate and output_attentions) else None - cross_attentions = () if (return_dict_in_generate and output_attentions) else None - decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None - - # if model is an encoder-decoder, retrieve encoder attention weights and hidden states - if return_dict_in_generate and self.config.is_encoder_decoder: - encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None - encoder_hidden_states = ( - model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None - ) - - # keep track of which sequences are already finished - batch_size, cur_len = input_ids.shape[:2] - this_peer_finished = False - unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) - model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs) - - model_forward = self.__call__ - compile_forward = self._valid_auto_compile_criteria(model_kwargs, generation_config) - if compile_forward: - os.environ["TOKENIZERS_PARALLELISM"] = "0" - model_forward = self.get_compiled_call(generation_config.compile_config) - - if generation_config.prefill_chunk_size is not None: - model_kwargs = self._prefill_chunking(input_ids, generation_config, **model_kwargs) - is_prefill = False - else: - is_prefill = True - - while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): - # prepare model inputs - model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) - - # prepare variable output controls (note: some models won't accept all output controls) - model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) - model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) - - if is_prefill: - outputs = self(**model_inputs, return_dict=True) - is_prefill = False - else: - outputs = model_forward(**model_inputs, return_dict=True) - - # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - ) - if synced_gpus and this_peer_finished: - continue - - # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration - # (the clone itself is always small) - next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device) - - - # pre-process distribution - next_token_scores = logits_processor(input_ids, next_token_logits) - - # Repetition Aware Sampling in VALL-E 2 - if use_ras: - probs_candidate = nn.functional.softmax(next_token_scores, dim=-1) - next_tokens_candidate = torch.multinomial(probs_candidate, num_samples=1).squeeze(1) - rep_num = (input_ids[:,-win_size:] == next_tokens_candidate).sum().item() + 1 - if rep_num >= win_size * tau_r: - next_token_scores = next_token_logits - - # Store scores, attentions and hidden_states when required - if return_dict_in_generate: - if output_scores: - scores += (next_token_scores,) - if output_logits: - raw_logits += (next_token_logits,) - if output_attentions: - decoder_attentions += ( - (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) - ) - if self.config.is_encoder_decoder: - cross_attentions += (outputs.cross_attentions,) - - if output_hidden_states: - decoder_hidden_states += ( - (outputs.decoder_hidden_states,) - if self.config.is_encoder_decoder - else (outputs.hidden_states,) - ) - - # token selection - if do_sample: - probs = nn.functional.softmax(next_token_scores, dim=-1) - # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution - next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) - else: - next_tokens = torch.argmax(next_token_scores, dim=-1) - - # finished sentences should have their next token be a padding token - if has_eos_stopping_criteria: - next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) - - # update generated ids, model inputs, and length for next step - input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) - if streamer is not None: - streamer.put(next_tokens.cpu()) - - unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores) - this_peer_finished = unfinished_sequences.max() == 0 - cur_len += 1 - - # This is needed to properly delete outputs.logits which may be very large for first iteration - # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration - del outputs - - if streamer is not None: - streamer.end() - - if return_dict_in_generate: - if self.config.is_encoder_decoder: - return GenerateEncoderDecoderOutput( - sequences=input_ids, - scores=scores, - logits=raw_logits, - encoder_attentions=encoder_attentions, - encoder_hidden_states=encoder_hidden_states, - decoder_attentions=decoder_attentions, - cross_attentions=cross_attentions, - decoder_hidden_states=decoder_hidden_states, - past_key_values=model_kwargs.get("past_key_values"), - ) - else: - return GenerateDecoderOnlyOutput( - sequences=input_ids, - scores=scores, - logits=raw_logits, - attentions=decoder_attentions, - hidden_states=decoder_hidden_states, - past_key_values=model_kwargs.get("past_key_values"), - ) - else: +import os + +from typing import Any, Callable, Optional, Union +import torch +from torch import nn +from transformers.generation.logits_process import ( + LogitsProcessorList +) +from transformers.generation.stopping_criteria import ( + StoppingCriteriaList +) +from transformers.generation.configuration_utils import ( + GenerationConfig +) +from transformers.generation.streamers import BaseStreamer +from transformers.generation.utils import ( + GenerateNonBeamOutput, + GenerateEncoderDecoderOutput, + GenerateDecoderOnlyOutput, +) +from transformers import StoppingCriteria + + +def _ras_sample_hf_engine( + self, + input_ids: torch.LongTensor, + logits_processor: LogitsProcessorList, + stopping_criteria: StoppingCriteriaList, + generation_config: GenerationConfig, + synced_gpus: bool = False, + streamer: Optional["BaseStreamer"] = None, + use_ras=False, + win_size=25, + tau_r=0.2, + **model_kwargs, +) -> Union[GenerateNonBeamOutput, torch.LongTensor]: + r""" + Generates sequences of token ids for models with a language modeling head using **multinomial sampling** and + can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. + + Parameters: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`): + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + generation_config ([`~generation.GenerationConfig`]): + The generation configuration to be used as parametrization of the decoding method. + synced_gpus (`bool`): + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). + streamer (`BaseStreamer`, *optional*): + Streamer object that will be used to stream the generated sequences. Generated tokens are passed + through `streamer.put(token_ids)` and the streamer is responsible for any further processing. + model_kwargs: + Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is + an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or `torch.LongTensor`: + A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + """ + # init values + pad_token_id = generation_config._pad_token_tensor + output_attentions = generation_config.output_attentions + output_hidden_states = generation_config.output_hidden_states + output_scores = generation_config.output_scores + output_logits = generation_config.output_logits + return_dict_in_generate = generation_config.return_dict_in_generate + has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria) + do_sample = generation_config.do_sample + + # init attention / hidden states / scores tuples + scores = () if (return_dict_in_generate and output_scores) else None + raw_logits = () if (return_dict_in_generate and output_logits) else None + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + if return_dict_in_generate and self.config.is_encoder_decoder: + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + # keep track of which sequences are already finished + batch_size, cur_len = input_ids.shape[:2] + this_peer_finished = False + unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) + model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs) + + model_forward = self.__call__ + compile_forward = self._valid_auto_compile_criteria(model_kwargs, generation_config) + if compile_forward: + os.environ["TOKENIZERS_PARALLELISM"] = "0" + model_forward = self.get_compiled_call(generation_config.compile_config) + + if generation_config.prefill_chunk_size is not None: + model_kwargs = self._prefill_chunking(input_ids, generation_config, **model_kwargs) + is_prefill = False + else: + is_prefill = True + + while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): + # prepare model inputs + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + + # prepare variable output controls (note: some models won't accept all output controls) + model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) + model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) + + if is_prefill: + outputs = self(**model_inputs, return_dict=True) + is_prefill = False + else: + outputs = model_forward(**model_inputs, return_dict=True) + + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + if synced_gpus and this_peer_finished: + continue + + # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration + # (the clone itself is always small) + next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device) + + + # pre-process distribution + next_token_scores = logits_processor(input_ids, next_token_logits) + + # Repetition Aware Sampling in VALL-E 2 + if use_ras: + probs_candidate = nn.functional.softmax(next_token_scores, dim=-1) + next_tokens_candidate = torch.multinomial(probs_candidate, num_samples=1).squeeze(1) + rep_num = (input_ids[:,-win_size:] == next_tokens_candidate).sum().item() + 1 + if rep_num >= win_size * tau_r: + next_token_scores = next_token_logits + + # Store scores, attentions and hidden_states when required + if return_dict_in_generate: + if output_scores: + scores += (next_token_scores,) + if output_logits: + raw_logits += (next_token_logits,) + if output_attentions: + decoder_attentions += ( + (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) + ) + if self.config.is_encoder_decoder: + cross_attentions += (outputs.cross_attentions,) + + if output_hidden_states: + decoder_hidden_states += ( + (outputs.decoder_hidden_states,) + if self.config.is_encoder_decoder + else (outputs.hidden_states,) + ) + + # token selection + if do_sample: + probs = nn.functional.softmax(next_token_scores, dim=-1) + # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution + next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) + else: + next_tokens = torch.argmax(next_token_scores, dim=-1) + + # finished sentences should have their next token be a padding token + if has_eos_stopping_criteria: + next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) + + # update generated ids, model inputs, and length for next step + input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) + if streamer is not None: + streamer.put(next_tokens.cpu()) + + unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores) + this_peer_finished = unfinished_sequences.max() == 0 + cur_len += 1 + + # This is needed to properly delete outputs.logits which may be very large for first iteration + # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration + del outputs + + if streamer is not None: + streamer.end() + + if return_dict_in_generate: + if self.config.is_encoder_decoder: + return GenerateEncoderDecoderOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return GenerateDecoderOnlyOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: return input_ids \ No newline at end of file diff --git a/soulxpodcast/models/soulxpodcast.py b/soulxpodcast/models/soulxpodcast.py index 05b2a01..16f917e 100644 --- a/soulxpodcast/models/soulxpodcast.py +++ b/soulxpodcast/models/soulxpodcast.py @@ -1,168 +1,168 @@ -import time -from datetime import datetime - -from tqdm import tqdm -from itertools import chain -from copy import deepcopy - -import numpy as np -import s3tokenizer -import torch - -from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache -from soulxpodcast.config import Config, SamplingParams, AutoPretrainedConfig -from soulxpodcast.engine.llm_engine import ( - HFLLMEngine, VLLMEngine -) -from soulxpodcast.models.modules.flow import CausalMaskedDiffWithXvec -from soulxpodcast.models.modules.hifigan import HiFTGenerator - -class SoulXPodcast(torch.nn.Module): - def __init__(self, config: Config = None): - super().__init__() - self.config = Config() if config is None else config - - self.audio_tokenizer = s3tokenizer.load_model("speech_tokenizer_v2_25hz").cuda().eval() - if self.config.llm_engine == "hf": - self.llm = HFLLMEngine(**self.config.__dict__) - elif self.config.llm_engine == "vllm": - self.llm = VLLMEngine(**self.config.__dict__) - else: - raise NotImplementedError - - self.use_tqdm = True - - self.flow = CausalMaskedDiffWithXvec() - if self.config.hf_config.fp16_flow: - timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] - tqdm.write(f"[{timestamp}] - [INFO] - Casting flow to fp16") - self.flow.half() - self.flow.load_state_dict(torch.load(f"{self.config.model}/flow.pt", map_location="cpu", weights_only=True), strict=True) - self.flow.cuda().eval() - - self.hift = HiFTGenerator() - hift_state_dict = {k.replace('generator.', ''): v for k, v in torch.load(f"{self.config.model}/hift.pt", map_location="cpu", weights_only=True).items()} - self.hift.load_state_dict(hift_state_dict, strict=True) - self.hift.cuda().eval() - - - @torch.inference_mode() - def forward_longform( - self, prompt_mels_for_llm, - prompt_mels_lens_for_llm: torch.Tensor, - prompt_text_tokens_for_llm: list[list[int]], - text_tokens_for_llm: list[list[int]], - prompt_mels_for_flow_ori, - spk_emb_for_flow: torch.Tensor, - sampling_params: SamplingParams | list[SamplingParams], - spk_ids: list[list[int]], - use_dialect_prompt: bool = False, - dialect_prompt_text_tokens_for_llm: list[list[int]] = None, - dialect_prefix: list[list[int]] = None, - **kwargs, # for compatibility - ): - - prompt_size, turn_size = len(prompt_mels_for_llm), len(text_tokens_for_llm) - - # Audio tokenization - prompt_speech_tokens_ori, prompt_speech_tokens_lens_ori = self.audio_tokenizer.quantize( - prompt_mels_for_llm.cuda(), prompt_mels_lens_for_llm.cuda() - ) - - # align speech token with speech feat as to reduce - # the noise ratio during the generation process. - prompt_speech_tokens = [] - prompt_mels_for_flow, prompt_mels_lens_for_flow = [], [] - - for prompt_index in range(prompt_size): - prompt_speech_token_len = prompt_speech_tokens_lens_ori[prompt_index].item() - prompt_speech_token = prompt_speech_tokens_ori[prompt_index, :prompt_speech_token_len] - prompt_mel = prompt_mels_for_flow_ori[prompt_index] - prompt_mel_len = prompt_mel.shape[0] - if prompt_speech_token_len * 2 > prompt_mel_len: - prompt_speech_token = prompt_speech_token[:int(prompt_mel_len/2)] - prompt_mel_len = torch.tensor([prompt_mel_len]).cuda() - else: - prompt_mel = prompt_mel.detach().clone()[:prompt_speech_token_len * 2].cuda() - prompt_mel_len = torch.tensor([prompt_speech_token_len * 2]).cuda() - prompt_speech_tokens.append(prompt_speech_token) - prompt_mels_for_flow.append(prompt_mel) - prompt_mels_lens_for_flow.append(prompt_mel_len) - - # Prepare LLM inputs - prompt_inputs = [] - history_inputs = [] - - for i in range(prompt_size): - speech_tokens_i = [token+self.config.hf_config.speech_token_offset for token in prompt_speech_tokens[i].tolist()] - speech_tokens_i += [self.config.hf_config.eos_token_id] - if use_dialect_prompt and len(dialect_prompt_text_tokens_for_llm[i])>0: - dialect_prompt_input = prompt_text_tokens_for_llm[i] + speech_tokens_i + dialect_prompt_text_tokens_for_llm[i] - if i>0: - dialect_prompt_input = dialect_prefix[0] + dialect_prompt_input - prompt_input = self.llm.generate(dialect_prompt_input, sampling_params, past_key_values=None)['token_ids'] - prompt_inputs.append(dialect_prefix[i+1]+dialect_prompt_text_tokens_for_llm[i] + prompt_input) - history_inputs.append(dialect_prefix[i+1]+dialect_prompt_text_tokens_for_llm[i] + prompt_input) - else: - prompt_inputs.append(prompt_text_tokens_for_llm[i] + speech_tokens_i ) - history_inputs.append(prompt_text_tokens_for_llm[i] + speech_tokens_i ) - - generated_wavs, results_dict = [], {} - - # LLM generation - inputs = list(chain.from_iterable(prompt_inputs)) - cache_config = AutoPretrainedConfig().from_dataclass(self.llm.config.hf_config) - past_key_values = DynamicCache(config=cache_config) - valid_turn_size = prompt_size - for i in range(turn_size): - - # # set ratio: reach the reset cache ratio; - if valid_turn_size > self.config.max_turn_size or len(inputs)>self.config.turn_tokens_threshold: - assert self.config.max_turn_size >= self.config.prompt_context + self.config.history_context, "Invalid Long history size setting, " - prompt_text_bound = max(self.config.prompt_context, len(history_inputs)-self.config.history_text_context-self.config.history_context) - inputs = list(chain.from_iterable( - history_inputs[:self.config.prompt_context]+ \ - history_inputs[prompt_text_bound:-self.config.history_context]+ \ - prompt_inputs[-self.config.history_context:] - )) - valid_turn_size = self.config.prompt_context + len(history_inputs) - prompt_text_bound - past_key_values = DynamicCache(config=cache_config) - valid_turn_size += 1 - - inputs.extend(text_tokens_for_llm[i]) - start_time = time.time() - llm_outputs = self.llm.generate(inputs, sampling_params, past_key_values=past_key_values) - - inputs.extend(llm_outputs['token_ids']) - prompt_inputs.append(text_tokens_for_llm[i]+llm_outputs['token_ids']) - history_inputs.append(text_tokens_for_llm[i][:-1]) # remove the <|audio_start|> - - # Prepare Flow inputs - turn_spk = spk_ids[i] - generated_speech_tokens = [token - self.config.hf_config.speech_token_offset for token in llm_outputs['token_ids'][:-1]] # ignore last eos - prompt_speech_token = prompt_speech_tokens[turn_spk].tolist() - flow_input = torch.tensor([prompt_speech_token + generated_speech_tokens]) - flow_inputs_len = torch.tensor([len(prompt_speech_token) + len(generated_speech_tokens)]) - - # Flow generation and HiFi-GAN generation - start_idx = spk_ids[i] - prompt_mels = prompt_mels_for_flow[start_idx][None] - prompt_mels_lens = prompt_mels_lens_for_flow[start_idx][None] - spk_emb = spk_emb_for_flow[start_idx:start_idx+1] - - # Flow generation - with torch.amp.autocast("cuda", dtype=torch.float16 if self.config.hf_config.fp16_flow else torch.float32): - generated_mels, generated_mels_lens = self.flow( - flow_input.cuda(), flow_inputs_len.cuda(), - prompt_mels, prompt_mels_lens, spk_emb.cuda(), - streaming=False, finalize=True - ) - - # HiFi-GAN generation - mel = generated_mels[:, :, prompt_mels_lens[0].item():generated_mels_lens[0].item()] - wav, _ = self.hift(speech_feat=mel) - generated_wavs.append(wav) - - # Save the generated wav; - results_dict['generated_wavs'] = generated_wavs +import time +from datetime import datetime + +from tqdm import tqdm +from itertools import chain +from copy import deepcopy + +import numpy as np +import s3tokenizer +import torch + +from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache +from soulxpodcast.config import Config, SamplingParams, AutoPretrainedConfig +from soulxpodcast.engine.llm_engine import ( + HFLLMEngine, VLLMEngine +) +from soulxpodcast.models.modules.flow import CausalMaskedDiffWithXvec +from soulxpodcast.models.modules.hifigan import HiFTGenerator + +class SoulXPodcast(torch.nn.Module): + def __init__(self, config: Config = None): + super().__init__() + self.config = Config() if config is None else config + + self.audio_tokenizer = s3tokenizer.load_model("speech_tokenizer_v2_25hz").cuda().eval() + if self.config.llm_engine == "hf": + self.llm = HFLLMEngine(**self.config.__dict__) + elif self.config.llm_engine == "vllm": + self.llm = VLLMEngine(**self.config.__dict__) + else: + raise NotImplementedError + + self.use_tqdm = True + + self.flow = CausalMaskedDiffWithXvec() + if self.config.hf_config.fp16_flow: + timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] + tqdm.write(f"[{timestamp}] - [INFO] - Casting flow to fp16") + self.flow.half() + self.flow.load_state_dict(torch.load(f"{self.config.model}/flow.pt", map_location="cpu", weights_only=True), strict=True) + self.flow.cuda().eval() + + self.hift = HiFTGenerator() + hift_state_dict = {k.replace('generator.', ''): v for k, v in torch.load(f"{self.config.model}/hift.pt", map_location="cpu", weights_only=True).items()} + self.hift.load_state_dict(hift_state_dict, strict=True) + self.hift.cuda().eval() + + + @torch.inference_mode() + def forward_longform( + self, prompt_mels_for_llm, + prompt_mels_lens_for_llm: torch.Tensor, + prompt_text_tokens_for_llm: list[list[int]], + text_tokens_for_llm: list[list[int]], + prompt_mels_for_flow_ori, + spk_emb_for_flow: torch.Tensor, + sampling_params: SamplingParams | list[SamplingParams], + spk_ids: list[list[int]], + use_dialect_prompt: bool = False, + dialect_prompt_text_tokens_for_llm: list[list[int]] = None, + dialect_prefix: list[list[int]] = None, + **kwargs, # for compatibility + ): + + prompt_size, turn_size = len(prompt_mels_for_llm), len(text_tokens_for_llm) + + # Audio tokenization + prompt_speech_tokens_ori, prompt_speech_tokens_lens_ori = self.audio_tokenizer.quantize( + prompt_mels_for_llm.cuda(), prompt_mels_lens_for_llm.cuda() + ) + + # align speech token with speech feat as to reduce + # the noise ratio during the generation process. + prompt_speech_tokens = [] + prompt_mels_for_flow, prompt_mels_lens_for_flow = [], [] + + for prompt_index in range(prompt_size): + prompt_speech_token_len = prompt_speech_tokens_lens_ori[prompt_index].item() + prompt_speech_token = prompt_speech_tokens_ori[prompt_index, :prompt_speech_token_len] + prompt_mel = prompt_mels_for_flow_ori[prompt_index] + prompt_mel_len = prompt_mel.shape[0] + if prompt_speech_token_len * 2 > prompt_mel_len: + prompt_speech_token = prompt_speech_token[:int(prompt_mel_len/2)] + prompt_mel_len = torch.tensor([prompt_mel_len]).cuda() + else: + prompt_mel = prompt_mel.detach().clone()[:prompt_speech_token_len * 2].cuda() + prompt_mel_len = torch.tensor([prompt_speech_token_len * 2]).cuda() + prompt_speech_tokens.append(prompt_speech_token) + prompt_mels_for_flow.append(prompt_mel) + prompt_mels_lens_for_flow.append(prompt_mel_len) + + # Prepare LLM inputs + prompt_inputs = [] + history_inputs = [] + + for i in range(prompt_size): + speech_tokens_i = [token+self.config.hf_config.speech_token_offset for token in prompt_speech_tokens[i].tolist()] + speech_tokens_i += [self.config.hf_config.eos_token_id] + if use_dialect_prompt and len(dialect_prompt_text_tokens_for_llm[i])>0: + dialect_prompt_input = prompt_text_tokens_for_llm[i] + speech_tokens_i + dialect_prompt_text_tokens_for_llm[i] + if i>0: + dialect_prompt_input = dialect_prefix[0] + dialect_prompt_input + prompt_input = self.llm.generate(dialect_prompt_input, sampling_params, past_key_values=None)['token_ids'] + prompt_inputs.append(dialect_prefix[i+1]+dialect_prompt_text_tokens_for_llm[i] + prompt_input) + history_inputs.append(dialect_prefix[i+1]+dialect_prompt_text_tokens_for_llm[i] + prompt_input) + else: + prompt_inputs.append(prompt_text_tokens_for_llm[i] + speech_tokens_i ) + history_inputs.append(prompt_text_tokens_for_llm[i] + speech_tokens_i ) + + generated_wavs, results_dict = [], {} + + # LLM generation + inputs = list(chain.from_iterable(prompt_inputs)) + cache_config = AutoPretrainedConfig().from_dataclass(self.llm.config.hf_config) + past_key_values = DynamicCache(config=cache_config) + valid_turn_size = prompt_size + for i in range(turn_size): + + # # set ratio: reach the reset cache ratio; + if valid_turn_size > self.config.max_turn_size or len(inputs)>self.config.turn_tokens_threshold: + assert self.config.max_turn_size >= self.config.prompt_context + self.config.history_context, "Invalid Long history size setting, " + prompt_text_bound = max(self.config.prompt_context, len(history_inputs)-self.config.history_text_context-self.config.history_context) + inputs = list(chain.from_iterable( + history_inputs[:self.config.prompt_context]+ \ + history_inputs[prompt_text_bound:-self.config.history_context]+ \ + prompt_inputs[-self.config.history_context:] + )) + valid_turn_size = self.config.prompt_context + len(history_inputs) - prompt_text_bound + past_key_values = DynamicCache(config=cache_config) + valid_turn_size += 1 + + inputs.extend(text_tokens_for_llm[i]) + start_time = time.time() + llm_outputs = self.llm.generate(inputs, sampling_params, past_key_values=past_key_values) + + inputs.extend(llm_outputs['token_ids']) + prompt_inputs.append(text_tokens_for_llm[i]+llm_outputs['token_ids']) + history_inputs.append(text_tokens_for_llm[i][:-1]) # remove the <|audio_start|> + + # Prepare Flow inputs + turn_spk = spk_ids[i] + generated_speech_tokens = [token - self.config.hf_config.speech_token_offset for token in llm_outputs['token_ids'][:-1]] # ignore last eos + prompt_speech_token = prompt_speech_tokens[turn_spk].tolist() + flow_input = torch.tensor([prompt_speech_token + generated_speech_tokens]) + flow_inputs_len = torch.tensor([len(prompt_speech_token) + len(generated_speech_tokens)]) + + # Flow generation and HiFi-GAN generation + start_idx = spk_ids[i] + prompt_mels = prompt_mels_for_flow[start_idx][None] + prompt_mels_lens = prompt_mels_lens_for_flow[start_idx][None] + spk_emb = spk_emb_for_flow[start_idx:start_idx+1] + + # Flow generation + with torch.amp.autocast("cuda", dtype=torch.float16 if self.config.hf_config.fp16_flow else torch.float32): + generated_mels, generated_mels_lens = self.flow( + flow_input.cuda(), flow_inputs_len.cuda(), + prompt_mels, prompt_mels_lens, spk_emb.cuda(), + streaming=False, finalize=True + ) + + # HiFi-GAN generation + mel = generated_mels[:, :, prompt_mels_lens[0].item():generated_mels_lens[0].item()] + wav, _ = self.hift(speech_feat=mel) + generated_wavs.append(wav) + + # Save the generated wav; + results_dict['generated_wavs'] = generated_wavs return results_dict \ No newline at end of file diff --git a/soulxpodcast/utils/audio.py b/soulxpodcast/utils/audio.py index 5614f0e..727cc44 100644 --- a/soulxpodcast/utils/audio.py +++ b/soulxpodcast/utils/audio.py @@ -1,123 +1,123 @@ -import torch -import numpy as np -from librosa.filters import mel as librosa_mel_fn -from scipy.io.wavfile import read - -MAX_WAV_VALUE = 32768.0 - - -def load_wav(full_path): - sampling_rate, data = read(full_path) - return data, sampling_rate - - -def dynamic_range_compression(x, C=1, clip_val=1e-5): - return np.log(np.clip(x, a_min=clip_val, a_max=None) * C) - - -def dynamic_range_decompression(x, C=1): - return np.exp(x) / C - - -def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): - return torch.log(torch.clamp(x, min=clip_val) * C) - - -def dynamic_range_decompression_torch(x, C=1): - return torch.exp(x) / C - - -def spectral_normalize_torch(magnitudes): - output = dynamic_range_compression_torch(magnitudes) - return output - - -def spectral_de_normalize_torch(magnitudes): - output = dynamic_range_decompression_torch(magnitudes) - return output - - -mel_basis = {} -hann_window = {} - - -def mel_spectrogram(y, n_fft=1920, num_mels=80, sampling_rate=24000, hop_size=480, - win_size=1920, fmin=0, fmax=8000, center=False): - global mel_basis, hann_window # pylint: disable=global-statement - if f"{str(fmax)}_{str(y.device)}" not in mel_basis: - mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) - mel_basis[str(fmax) + "_" + str(y.device)] = torch.from_numpy(mel).float().to(y.device) - hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device) - - y = torch.nn.functional.pad( - y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect" - ) - y = y.squeeze(1) - - spec = torch.view_as_real( - torch.stft( - y, - n_fft, - hop_length=hop_size, - win_length=win_size, - window=hann_window[str(y.device)], - center=center, - pad_mode="reflect", - normalized=False, - onesided=True, - return_complex=True, - ) - ) - - spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) - - spec = torch.matmul(mel_basis[str(fmax) + "_" + str(y.device)], spec) - spec = spectral_normalize_torch(spec) - - return spec - - -def audio_volume_normalize(audio: torch.Tensor, coeff=0.1): - """ - Normalize the volume of an audio signal. - - Parameters: - audio (torch tensor): Input audio signal array. - coeff (float): Target coefficient for normalization, default is 0.1. - - Returns: - torch tensor: The volume-normalized audio signal. - """ - # Sort the absolute values of the audio signal - device = audio.device - audio = audio.cpu().numpy() - temp = np.sort(np.abs(audio)) - - # If the maximum value is less than 0.1, scale the array to have a maximum of 0.1 - if temp[-1] < 0.1: - scaling_factor = max( - temp[-1], 1e-3 - ) # Prevent division by zero with a small constant - audio = audio / scaling_factor * 0.1 - - # Filter out values less than 0.01 from temp - temp = temp[temp > 0.01] - L = temp.shape[0] # Length of the filtered array - - # If there are fewer than or equal to 10 significant values, return the audio without further processing - if L <= 10: - return audio - - # Compute the average of the top 10% to 1% of values in temp - volume = np.mean(temp[int(0.9 * L) : int(0.99 * L)]) - - # Normalize the audio to the target coefficient level, clamping the scale factor between 0.1 and 10 - audio = audio * np.clip(coeff / volume, a_min=0.1, a_max=10) - - # Ensure the maximum absolute value in the audio does not exceed 1 - max_value = np.max(np.abs(audio)) - if max_value > 1: - audio = audio / max_value - - audio = torch.from_numpy(audio).to(device) +import torch +import numpy as np +from librosa.filters import mel as librosa_mel_fn +from scipy.io.wavfile import read + +MAX_WAV_VALUE = 32768.0 + + +def load_wav(full_path): + sampling_rate, data = read(full_path) + return data, sampling_rate + + +def dynamic_range_compression(x, C=1, clip_val=1e-5): + return np.log(np.clip(x, a_min=clip_val, a_max=None) * C) + + +def dynamic_range_decompression(x, C=1): + return np.exp(x) / C + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): + return torch.log(torch.clamp(x, min=clip_val) * C) + + +def dynamic_range_decompression_torch(x, C=1): + return torch.exp(x) / C + + +def spectral_normalize_torch(magnitudes): + output = dynamic_range_compression_torch(magnitudes) + return output + + +def spectral_de_normalize_torch(magnitudes): + output = dynamic_range_decompression_torch(magnitudes) + return output + + +mel_basis = {} +hann_window = {} + + +def mel_spectrogram(y, n_fft=1920, num_mels=80, sampling_rate=24000, hop_size=480, + win_size=1920, fmin=0, fmax=8000, center=False): + global mel_basis, hann_window # pylint: disable=global-statement + if f"{str(fmax)}_{str(y.device)}" not in mel_basis: + mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) + mel_basis[str(fmax) + "_" + str(y.device)] = torch.from_numpy(mel).float().to(y.device) + hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device) + + y = torch.nn.functional.pad( + y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect" + ) + y = y.squeeze(1) + + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window[str(y.device)], + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) + + spec = torch.matmul(mel_basis[str(fmax) + "_" + str(y.device)], spec) + spec = spectral_normalize_torch(spec) + + return spec + + +def audio_volume_normalize(audio: torch.Tensor, coeff=0.1): + """ + Normalize the volume of an audio signal. + + Parameters: + audio (torch tensor): Input audio signal array. + coeff (float): Target coefficient for normalization, default is 0.1. + + Returns: + torch tensor: The volume-normalized audio signal. + """ + # Sort the absolute values of the audio signal + device = audio.device + audio = audio.cpu().numpy() + temp = np.sort(np.abs(audio)) + + # If the maximum value is less than 0.1, scale the array to have a maximum of 0.1 + if temp[-1] < 0.1: + scaling_factor = max( + temp[-1], 1e-3 + ) # Prevent division by zero with a small constant + audio = audio / scaling_factor * 0.1 + + # Filter out values less than 0.01 from temp + temp = temp[temp > 0.01] + L = temp.shape[0] # Length of the filtered array + + # If there are fewer than or equal to 10 significant values, return the audio without further processing + if L <= 10: + return audio + + # Compute the average of the top 10% to 1% of values in temp + volume = np.mean(temp[int(0.9 * L) : int(0.99 * L)]) + + # Normalize the audio to the target coefficient level, clamping the scale factor between 0.1 and 10 + audio = audio * np.clip(coeff / volume, a_min=0.1, a_max=10) + + # Ensure the maximum absolute value in the audio does not exceed 1 + max_value = np.max(np.abs(audio)) + if max_value > 1: + audio = audio / max_value + + audio = torch.from_numpy(audio).to(device) return audio \ No newline at end of file diff --git a/soulxpodcast/utils/commons.py b/soulxpodcast/utils/commons.py index 4ff87d0..711376f 100644 --- a/soulxpodcast/utils/commons.py +++ b/soulxpodcast/utils/commons.py @@ -1,10 +1,10 @@ -import random -import numpy as np -import torch - - -def set_all_random_seed(seed): - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) +import random +import numpy as np +import torch + + +def set_all_random_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) \ No newline at end of file diff --git a/soulxpodcast/utils/dataloader.py b/soulxpodcast/utils/dataloader.py index 4274596..39f0418 100644 --- a/soulxpodcast/utils/dataloader.py +++ b/soulxpodcast/utils/dataloader.py @@ -1,198 +1,198 @@ -import os -import json -from tqdm import tqdm -from datetime import datetime - -import onnxruntime -import torch -import torchaudio -from torch.utils.data import DataLoader, Dataset, DistributedSampler -import torchaudio.compliance.kaldi as kaldi - -import s3tokenizer - -from soulxpodcast.utils.text import normalize_text -from soulxpodcast.utils.audio import mel_spectrogram, audio_volume_normalize -from soulxpodcast.config import Config, SamplingParams - - -SPK_DICT = ["<|SPEAKER_0|>", "<|SPEAKER_1|>", "<|SPEAKER_2|>", "<|SPEAKER_3|>",] -TEXT_START, TEXT_END, AUDIO_START = "<|text_start|>", "<|text_end|>", "<|semantic_token_start|>" -TASK_PODCAST = "<|task_podcast|>" - - -class PodcastDataset(Dataset): - - def __init__(self, text_tokenizer, data_list, model_config: Config): - self.datas = [] - self.model_config = model_config - - """Example data_list: - ``` - {"key": "uttid_1", "prompt_text": ["prompt_text1", "prompt_text2"], "dialect_prompt_text": ["dialect_prompt_text1", "dialect_prompt_text2"], - "text": ["text1", "text2], "spk": [0, 1], "prompt_wav": ["/mnt/data/audio/00000000.wav", "/mnt/data/audio/00000001.wav"], "wav": "/mnt/data/audio_synthetic/uttid_1.wav"} - ``` - Note: - - `key` is the key of this sample. - - `prompt_text` is the text used for prompt. - - `dialect_prompt_text` is the reshot text used for prompt. - - `text` is the text used for generating real audio. - - `spk` is the target speaker id to synthesize, corresponds to the prompt order. Default SPEAKER_0. - - `prompt_wav` is the audio used for prompt. - - `wav` is the path to the generated audio to be saved (we highly recommend to pre-define the save path before running the script). - """ - missing = 0 - with open(data_list, 'r', encoding='utf-8') as f: - lines = f.readlines() - total_lines = len(lines) - iterator = tqdm(lines, desc='Loading data') - for line in iterator: - data = json.loads(line.strip()) - valid = True - for k in ['key', 'prompt_text', 'text', 'prompt_wav']: - if k not in data: - valid = False - break - if data[k] is None: - valid = False - break - valid = True - for url in data["prompt_wav"]: - if not os.path.exists(url): - valid = False - break - if valid: - self.datas.append(data) - else: - missing += 1 - timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] - tqdm.write(f'[{timestamp}] - [INFO] - Loaded {total_lines} lines, found {missing} missing lines, total valid lines == {len(self.datas)}.') - - self.text_tokenizer = text_tokenizer - - option = onnxruntime.SessionOptions() - option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL - option.intra_op_num_threads = 1 - self.spk_model = onnxruntime.InferenceSession(f"{self.model_config.model}/campplus.onnx", sess_options=option, - providers=["CPUExecutionProvider"]) - - def __len__(self): - return len(self.datas) - - def __getitem__(self, idx): - data = self.datas[idx] - try: - prompt_text_ids_list, dialect_prompt_text_ids_list, spk_emb_list, mel_list, mel_len_list, log_mel_list = ( - [], [], [], [], [], [] - ) - # Prepare prompt information - use_dialect_prompt = "dialect_prompt_text" in data - dialect_prefix_list = [] - dialect_prefix_list.append(self.text_tokenizer.encode(f"{TASK_PODCAST}")) - for spk_idx, (prompt_text, prompt_wav) in enumerate(zip(data["prompt_text"], data["prompt_wav"])): - # 1. feature for s3tokenizer - audio = s3tokenizer.load_audio(prompt_wav, sr=16000) - audio = audio_volume_normalize(audio) - # [T] - log_mel = s3tokenizer.log_mel_spectrogram(audio) # [num_mels, T] - - # 2. feature for speaker embedding - spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000) - spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True) - spk_emb = self.spk_model.run( - None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()} - )[0].flatten().tolist() - - # 3. feature for flow - audio, sample_rate = torchaudio.load(prompt_wav, backend='soundfile') - audio = audio[0] - audio = audio_volume_normalize(audio).unsqueeze(0) - # audio = audio.mean(dim=0, keepdim=True) # [1, T] - if sample_rate != 24000: - audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=24000)(audio) - mel = mel_spectrogram(audio).transpose(1, 2).squeeze(0) # [T, num_mels] - if mel.shape[0] %2 !=0: - mel = mel[:-1] - mel_len = mel.shape[0] - - # 4. feature for llm - prompt_text = normalize_text(prompt_text) # remove some space and strange character - prompt_text = f"{SPK_DICT[spk_idx]}{TEXT_START}{prompt_text}{TEXT_END}{AUDIO_START}" - if spk_idx == 0: - prompt_text = f"{TASK_PODCAST}{prompt_text}" - prompt_text_ids = self.text_tokenizer.encode(prompt_text) - prompt_text_ids_list.append(prompt_text_ids) - if use_dialect_prompt: - dialect_prompt_text = normalize_text(data["dialect_prompt_text"][spk_idx]) - dialect_prompt_text = f"{SPK_DICT[spk_idx]}{TEXT_START}{dialect_prompt_text}{TEXT_END}{AUDIO_START}" - dialect_prompt_text_ids = self.text_tokenizer.encode(dialect_prompt_text) - dialect_prompt_text_ids_list.append(dialect_prompt_text_ids) - if spk_idx == 0: - dialect_prefix_list.append(self.text_tokenizer.encode(f"{TASK_PODCAST}")) - else: - dialect_prefix_list.append([]) - log_mel_list.append(log_mel) - spk_emb_list.append(spk_emb) - mel_list.append(mel); mel_len_list.append(mel_len) - item = { - "prompt_text_tokens": prompt_text_ids_list, - "spk_emb": spk_emb_list, "mel": mel_list, "mel_len": mel_len_list, "log_mel": log_mel_list, "info": data, - } - if use_dialect_prompt: - item.update({ - "use_dialect_prompt": True, - "dialect_prompt_text_tokens": dialect_prompt_text_ids_list, - "dialect_prefix": dialect_prefix_list, - }) - text_ids_list, spks_list = [], [] - if "spk" not in data: - data["spk"] = [0] * len(data["text"]) - - for text, spk in zip(data["text"], data["spk"]): - # 4. feature for llm - text = normalize_text(text) - text = f"{SPK_DICT[spk]}{TEXT_START}{text}{TEXT_END}{AUDIO_START}" - text_ids = self.text_tokenizer.encode(text) - - text_ids_list.append(text_ids) - spks_list.append(spk) - - item.update({ - "text_tokens": text_ids_list, "spks_list": spks_list, - }) - except Exception as e: - timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] - tqdm.write(f"[{timestamp}] - [WARNING] - Error processing data item {data.get('key', idx)}: {e}") - return None - return item - -class PodcastInferHandler(PodcastDataset): - - def __init__(self, text_tokenizer, data_list, model_config: Config): - self.datas = [] - self.model_config = model_config - - """Example data_list: - ``` - {"key": "uttid_1", "prompt_text": ["prompt_text1", "prompt_text2"], "dialect_prompt_text": ["dialect_prompt_text1", "dialect_prompt_text2"], "text": ["text1", "text2], "spk": [0, 1], "prompt_wav": ["/mnt/data/audio/00000000.wav", "/mnt/data/audio/00000001.wav"], "wav": "/mnt/data/audio_synthetic/uttid_1.wav"} - ``` - Note: - - `key` is the key of this sample. - - `prompt_text` is the text used for prompt. - - `dialect_prompt_text` is the cot text used for prompt as to activate specific ability. - - `text` is the text used for generating real audio. - - `spk` is the target speaker id to synthesize, corresponds to the prompt order. Default SPEAKER_0. - - `prompt_wav` is the audio used for prompt. - - `wav` is the path to the generated audio to be saved (we highly recommend to pre-define the save path before running the script). - """ - missing = 0 - self.text_tokenizer = text_tokenizer - - option = onnxruntime.SessionOptions() - option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL - option.intra_op_num_threads = 1 - self.spk_model = onnxruntime.InferenceSession(f"{self.model_config.model}/campplus.onnx", sess_options=option, - providers=["CPUExecutionProvider"]) - - def update_datasource(self, data_list): +import os +import json +from tqdm import tqdm +from datetime import datetime + +import onnxruntime +import torch +import torchaudio +from torch.utils.data import DataLoader, Dataset, DistributedSampler +import torchaudio.compliance.kaldi as kaldi + +import s3tokenizer + +from soulxpodcast.utils.text import normalize_text +from soulxpodcast.utils.audio import mel_spectrogram, audio_volume_normalize +from soulxpodcast.config import Config, SamplingParams + + +SPK_DICT = ["<|SPEAKER_0|>", "<|SPEAKER_1|>", "<|SPEAKER_2|>", "<|SPEAKER_3|>",] +TEXT_START, TEXT_END, AUDIO_START = "<|text_start|>", "<|text_end|>", "<|semantic_token_start|>" +TASK_PODCAST = "<|task_podcast|>" + + +class PodcastDataset(Dataset): + + def __init__(self, text_tokenizer, data_list, model_config: Config): + self.datas = [] + self.model_config = model_config + + """Example data_list: + ``` + {"key": "uttid_1", "prompt_text": ["prompt_text1", "prompt_text2"], "dialect_prompt_text": ["dialect_prompt_text1", "dialect_prompt_text2"], + "text": ["text1", "text2], "spk": [0, 1], "prompt_wav": ["/mnt/data/audio/00000000.wav", "/mnt/data/audio/00000001.wav"], "wav": "/mnt/data/audio_synthetic/uttid_1.wav"} + ``` + Note: + - `key` is the key of this sample. + - `prompt_text` is the text used for prompt. + - `dialect_prompt_text` is the reshot text used for prompt. + - `text` is the text used for generating real audio. + - `spk` is the target speaker id to synthesize, corresponds to the prompt order. Default SPEAKER_0. + - `prompt_wav` is the audio used for prompt. + - `wav` is the path to the generated audio to be saved (we highly recommend to pre-define the save path before running the script). + """ + missing = 0 + with open(data_list, 'r', encoding='utf-8') as f: + lines = f.readlines() + total_lines = len(lines) + iterator = tqdm(lines, desc='Loading data') + for line in iterator: + data = json.loads(line.strip()) + valid = True + for k in ['key', 'prompt_text', 'text', 'prompt_wav']: + if k not in data: + valid = False + break + if data[k] is None: + valid = False + break + valid = True + for url in data["prompt_wav"]: + if not os.path.exists(url): + valid = False + break + if valid: + self.datas.append(data) + else: + missing += 1 + timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] + tqdm.write(f'[{timestamp}] - [INFO] - Loaded {total_lines} lines, found {missing} missing lines, total valid lines == {len(self.datas)}.') + + self.text_tokenizer = text_tokenizer + + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 1 + self.spk_model = onnxruntime.InferenceSession(f"{self.model_config.model}/campplus.onnx", sess_options=option, + providers=["CPUExecutionProvider"]) + + def __len__(self): + return len(self.datas) + + def __getitem__(self, idx): + data = self.datas[idx] + try: + prompt_text_ids_list, dialect_prompt_text_ids_list, spk_emb_list, mel_list, mel_len_list, log_mel_list = ( + [], [], [], [], [], [] + ) + # Prepare prompt information + use_dialect_prompt = "dialect_prompt_text" in data + dialect_prefix_list = [] + dialect_prefix_list.append(self.text_tokenizer.encode(f"{TASK_PODCAST}")) + for spk_idx, (prompt_text, prompt_wav) in enumerate(zip(data["prompt_text"], data["prompt_wav"])): + # 1. feature for s3tokenizer + audio = s3tokenizer.load_audio(prompt_wav, sr=16000) + audio = audio_volume_normalize(audio) + # [T] + log_mel = s3tokenizer.log_mel_spectrogram(audio) # [num_mels, T] + + # 2. feature for speaker embedding + spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000) + spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True) + spk_emb = self.spk_model.run( + None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()} + )[0].flatten().tolist() + + # 3. feature for flow + audio, sample_rate = torchaudio.load(prompt_wav, backend='soundfile') + audio = audio[0] + audio = audio_volume_normalize(audio).unsqueeze(0) + # audio = audio.mean(dim=0, keepdim=True) # [1, T] + if sample_rate != 24000: + audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=24000)(audio) + mel = mel_spectrogram(audio).transpose(1, 2).squeeze(0) # [T, num_mels] + if mel.shape[0] %2 !=0: + mel = mel[:-1] + mel_len = mel.shape[0] + + # 4. feature for llm + prompt_text = normalize_text(prompt_text) # remove some space and strange character + prompt_text = f"{SPK_DICT[spk_idx]}{TEXT_START}{prompt_text}{TEXT_END}{AUDIO_START}" + if spk_idx == 0: + prompt_text = f"{TASK_PODCAST}{prompt_text}" + prompt_text_ids = self.text_tokenizer.encode(prompt_text) + prompt_text_ids_list.append(prompt_text_ids) + if use_dialect_prompt: + dialect_prompt_text = normalize_text(data["dialect_prompt_text"][spk_idx]) + dialect_prompt_text = f"{SPK_DICT[spk_idx]}{TEXT_START}{dialect_prompt_text}{TEXT_END}{AUDIO_START}" + dialect_prompt_text_ids = self.text_tokenizer.encode(dialect_prompt_text) + dialect_prompt_text_ids_list.append(dialect_prompt_text_ids) + if spk_idx == 0: + dialect_prefix_list.append(self.text_tokenizer.encode(f"{TASK_PODCAST}")) + else: + dialect_prefix_list.append([]) + log_mel_list.append(log_mel) + spk_emb_list.append(spk_emb) + mel_list.append(mel); mel_len_list.append(mel_len) + item = { + "prompt_text_tokens": prompt_text_ids_list, + "spk_emb": spk_emb_list, "mel": mel_list, "mel_len": mel_len_list, "log_mel": log_mel_list, "info": data, + } + if use_dialect_prompt: + item.update({ + "use_dialect_prompt": True, + "dialect_prompt_text_tokens": dialect_prompt_text_ids_list, + "dialect_prefix": dialect_prefix_list, + }) + text_ids_list, spks_list = [], [] + if "spk" not in data: + data["spk"] = [0] * len(data["text"]) + + for text, spk in zip(data["text"], data["spk"]): + # 4. feature for llm + text = normalize_text(text) + text = f"{SPK_DICT[spk]}{TEXT_START}{text}{TEXT_END}{AUDIO_START}" + text_ids = self.text_tokenizer.encode(text) + + text_ids_list.append(text_ids) + spks_list.append(spk) + + item.update({ + "text_tokens": text_ids_list, "spks_list": spks_list, + }) + except Exception as e: + timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] + tqdm.write(f"[{timestamp}] - [WARNING] - Error processing data item {data.get('key', idx)}: {e}") + return None + return item + +class PodcastInferHandler(PodcastDataset): + + def __init__(self, text_tokenizer, data_list, model_config: Config): + self.datas = [] + self.model_config = model_config + + """Example data_list: + ``` + {"key": "uttid_1", "prompt_text": ["prompt_text1", "prompt_text2"], "dialect_prompt_text": ["dialect_prompt_text1", "dialect_prompt_text2"], "text": ["text1", "text2], "spk": [0, 1], "prompt_wav": ["/mnt/data/audio/00000000.wav", "/mnt/data/audio/00000001.wav"], "wav": "/mnt/data/audio_synthetic/uttid_1.wav"} + ``` + Note: + - `key` is the key of this sample. + - `prompt_text` is the text used for prompt. + - `dialect_prompt_text` is the cot text used for prompt as to activate specific ability. + - `text` is the text used for generating real audio. + - `spk` is the target speaker id to synthesize, corresponds to the prompt order. Default SPEAKER_0. + - `prompt_wav` is the audio used for prompt. + - `wav` is the path to the generated audio to be saved (we highly recommend to pre-define the save path before running the script). + """ + missing = 0 + self.text_tokenizer = text_tokenizer + + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 1 + self.spk_model = onnxruntime.InferenceSession(f"{self.model_config.model}/campplus.onnx", sess_options=option, + providers=["CPUExecutionProvider"]) + + def update_datasource(self, data_list): self.datas = data_list \ No newline at end of file diff --git a/soulxpodcast/utils/infer_utils.py b/soulxpodcast/utils/infer_utils.py index 78b764d..621ceee 100644 --- a/soulxpodcast/utils/infer_utils.py +++ b/soulxpodcast/utils/infer_utils.py @@ -1,95 +1,95 @@ -import re -import json -import torch -import argparse -from tqdm import tqdm -from datetime import datetime - -import s3tokenizer - -from soulxpodcast.models.soulxpodcast import SoulXPodcast -from soulxpodcast.utils.dataloader import PodcastInferHandler -from soulxpodcast.utils.commons import set_all_random_seed -from soulxpodcast.config import Config, SoulXPodcastLLMConfig, SamplingParams - - -def initiate_model(seed, model_path, llm_engine, fp16_flow): - set_all_random_seed(seed) - - hf_config = SoulXPodcastLLMConfig.from_initial_and_json( - initial_values={"fp16_flow": fp16_flow}, - json_file=f"{model_path}/soulxpodcast_config.json" - ) - if llm_engine == "vllm": - import importlib.util - if not importlib.util.find_spec("vllm"): - llm_engine = "hf" - timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] - tqdm.write(f"[{timestamp}] - [WARNING]: No install VLLM, switch to hf engine.") - - config = Config(model=model_path, enforce_eager=True, llm_engine=llm_engine, hf_config=hf_config) - model = SoulXPodcast(config) - - dataset = PodcastInferHandler(model.llm.tokenizer, None, config) - - return model, dataset - - -def process_single_input(dataset, target_text_list, prompt_wav_list, prompt_text_list, use_dialect_prompt, dialect_prompt_text_list): - spks, texts = [], [] - for target_text in target_text_list: - pattern = r'(\[S[1-9]\])(.+)' - match = re.match(pattern, target_text) - text, spk = match.group(2), int(match.group(1)[2])-1 - spks.append(spk) - texts.append(text) - - dataitem = {"key": "001", "prompt_text": prompt_text_list, "prompt_wav": prompt_wav_list, - "text": texts, "spk": spks, } - if use_dialect_prompt: - dataitem.update({ - "dialect_prompt_text": dialect_prompt_text_list - }) - dataset.update_datasource( - [ - dataitem - ] - ) - - # assert one data only; - data = dataset[0] - prompt_mels_for_llm, prompt_mels_lens_for_llm = s3tokenizer.padding(data["log_mel"]) # [B, num_mels=128, T] - spk_emb_for_flow = torch.tensor(data["spk_emb"]) - prompt_mels_for_flow = torch.nn.utils.rnn.pad_sequence(data["mel"], batch_first=True, padding_value=0) # [B, T', num_mels=80] - prompt_mels_lens_for_flow = torch.tensor(data['mel_len']) - text_tokens_for_llm = data["text_tokens"] - prompt_text_tokens_for_llm = data["prompt_text_tokens"] - spk_ids = data["spks_list"] - sampling_params = SamplingParams(use_ras=True,win_size=25,tau_r=0.2) - infos = [data["info"]] - processed_data = { - "prompt_mels_for_llm": prompt_mels_for_llm, - "prompt_mels_lens_for_llm": prompt_mels_lens_for_llm, - "prompt_text_tokens_for_llm": prompt_text_tokens_for_llm, - "text_tokens_for_llm": text_tokens_for_llm, - "prompt_mels_for_flow_ori": prompt_mels_for_flow, - "prompt_mels_lens_for_flow": prompt_mels_lens_for_flow, - "spk_emb_for_flow": spk_emb_for_flow, - "sampling_params": sampling_params, - "spk_ids": spk_ids, - "infos": infos, - "use_dialect_prompt": use_dialect_prompt, - } - if use_dialect_prompt: - processed_data.update({ - "dialect_prompt_text_tokens_for_llm": data["dialect_prompt_text_tokens"], - "dialect_prefix": data["dialect_prefix"], - }) - return processed_data - - -def check_models(model_path, inputs): - if inputs['use_dialect_prompt']: - assert 'dialect' in model_path, "Dialect prompt is used, you should use a dialect model." - +import re +import json +import torch +import argparse +from tqdm import tqdm +from datetime import datetime + +import s3tokenizer + +from soulxpodcast.models.soulxpodcast import SoulXPodcast +from soulxpodcast.utils.dataloader import PodcastInferHandler +from soulxpodcast.utils.commons import set_all_random_seed +from soulxpodcast.config import Config, SoulXPodcastLLMConfig, SamplingParams + + +def initiate_model(seed, model_path, llm_engine, fp16_flow): + set_all_random_seed(seed) + + hf_config = SoulXPodcastLLMConfig.from_initial_and_json( + initial_values={"fp16_flow": fp16_flow}, + json_file=f"{model_path}/soulxpodcast_config.json" + ) + if llm_engine == "vllm": + import importlib.util + if not importlib.util.find_spec("vllm"): + llm_engine = "hf" + timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] + tqdm.write(f"[{timestamp}] - [WARNING]: No install VLLM, switch to hf engine.") + + config = Config(model=model_path, enforce_eager=True, llm_engine=llm_engine, hf_config=hf_config) + model = SoulXPodcast(config) + + dataset = PodcastInferHandler(model.llm.tokenizer, None, config) + + return model, dataset + + +def process_single_input(dataset, target_text_list, prompt_wav_list, prompt_text_list, use_dialect_prompt, dialect_prompt_text_list): + spks, texts = [], [] + for target_text in target_text_list: + pattern = r'(\[S[1-9]\])(.+)' + match = re.match(pattern, target_text) + text, spk = match.group(2), int(match.group(1)[2])-1 + spks.append(spk) + texts.append(text) + + dataitem = {"key": "001", "prompt_text": prompt_text_list, "prompt_wav": prompt_wav_list, + "text": texts, "spk": spks, } + if use_dialect_prompt: + dataitem.update({ + "dialect_prompt_text": dialect_prompt_text_list + }) + dataset.update_datasource( + [ + dataitem + ] + ) + + # assert one data only; + data = dataset[0] + prompt_mels_for_llm, prompt_mels_lens_for_llm = s3tokenizer.padding(data["log_mel"]) # [B, num_mels=128, T] + spk_emb_for_flow = torch.tensor(data["spk_emb"]) + prompt_mels_for_flow = torch.nn.utils.rnn.pad_sequence(data["mel"], batch_first=True, padding_value=0) # [B, T', num_mels=80] + prompt_mels_lens_for_flow = torch.tensor(data['mel_len']) + text_tokens_for_llm = data["text_tokens"] + prompt_text_tokens_for_llm = data["prompt_text_tokens"] + spk_ids = data["spks_list"] + sampling_params = SamplingParams(use_ras=True,win_size=25,tau_r=0.2) + infos = [data["info"]] + processed_data = { + "prompt_mels_for_llm": prompt_mels_for_llm, + "prompt_mels_lens_for_llm": prompt_mels_lens_for_llm, + "prompt_text_tokens_for_llm": prompt_text_tokens_for_llm, + "text_tokens_for_llm": text_tokens_for_llm, + "prompt_mels_for_flow_ori": prompt_mels_for_flow, + "prompt_mels_lens_for_flow": prompt_mels_lens_for_flow, + "spk_emb_for_flow": spk_emb_for_flow, + "sampling_params": sampling_params, + "spk_ids": spk_ids, + "infos": infos, + "use_dialect_prompt": use_dialect_prompt, + } + if use_dialect_prompt: + processed_data.update({ + "dialect_prompt_text_tokens_for_llm": data["dialect_prompt_text_tokens"], + "dialect_prefix": data["dialect_prefix"], + }) + return processed_data + + +def check_models(model_path, inputs): + if inputs['use_dialect_prompt']: + assert 'dialect' in model_path, "Dialect prompt is used, you should use a dialect model." + return True \ No newline at end of file diff --git a/soulxpodcast/utils/parser.py b/soulxpodcast/utils/parser.py index 420a0c8..f475a32 100644 --- a/soulxpodcast/utils/parser.py +++ b/soulxpodcast/utils/parser.py @@ -1,87 +1,87 @@ -import json -import datetime -import os - - -def generate_time_index(): - """Generate a time-based unique key, e.g. '20251023-001'.""" - now = datetime.datetime.now() - return now.strftime("%Y%m%d-%H%M%S") - -def check_prefix(text): - prefixes = ["<|Henan|>", "<|Sichuan|>", "<|Yue|>"] - - for prefix in prefixes: - if text.startswith(prefix): - return True - return False - - -def podcast_format_parser(data, output_dir="outputs"): - """ - Parse the original multi-speaker podcast JSON to the target flattened format. - The key will be a time-based unique ID. - Args: - data (dict): input JSON data with 'speakers' and 'text' fields - output_dir (str): directory for output wav file path - Returns: - dict: converted format - """ - speakers = data.get("speakers", {}) - text_entries = data.get("text", []) - - # Create speaker name → numeric ID mapping - spk2id = {name: idx for idx, name in enumerate(speakers.keys())} - - # Collect prompts - prompt_text = [] - prompt_wav = [] - dialect_prompt_text = [] - - for name in speakers: - prompt_text.append(speakers[name].get("prompt_text", "")) - prompt_wav.append(speakers[name].get("prompt_audio", "")) - dialect_prompt_text.append(speakers[name].get("dialect_prompt", "")) - - # Collect dialogue text and speaker sequence - text_list = [] - spk_list = [] - for turn in text_entries: - if len(turn) == 2: - spk_name, utt_text = turn - text = f'[{spk_name}]{utt_text}' - text_list.append(text) - spk_list.append(spk2id.get(spk_name, -1)) - - # Generate time-based key - key = generate_time_index() - wav_path = os.path.join(output_dir, f"{key}.wav") - - use_dialect_prompt = False - for dialect_text in dialect_prompt_text: - if len(dialect_text) > 0: - assert check_prefix(dialect_text), f"Unknown dialect prefix: {dialect_text} \ - \n Prefix should be one of: <|Henan|>, <|Sichuan|>, <|Yue|>" - use_dialect_prompt = True - - result = { - "key": key, - "prompt_text": prompt_text, - "prompt_wav": prompt_wav, - "text": text_list, - "spk": spk_list, - "wav": wav_path, - "use_dialect_prompt": use_dialect_prompt, - "dialect_prompt_text": dialect_prompt_text - } - - return result - - -# Example usage -if __name__ == "__main__": - with open("example/podcast_script/script_henan.json", "r", encoding="utf-8") as f: - data = json.load(f) - - converted = podcast_format_parser(data) - print(json.dumps(converted, ensure_ascii=False, indent=2)) +import json +import datetime +import os + + +def generate_time_index(): + """Generate a time-based unique key, e.g. '20251023-001'.""" + now = datetime.datetime.now() + return now.strftime("%Y%m%d-%H%M%S") + +def check_prefix(text): + prefixes = ["<|Henan|>", "<|Sichuan|>", "<|Yue|>"] + + for prefix in prefixes: + if text.startswith(prefix): + return True + return False + + +def podcast_format_parser(data, output_dir="outputs"): + """ + Parse the original multi-speaker podcast JSON to the target flattened format. + The key will be a time-based unique ID. + Args: + data (dict): input JSON data with 'speakers' and 'text' fields + output_dir (str): directory for output wav file path + Returns: + dict: converted format + """ + speakers = data.get("speakers", {}) + text_entries = data.get("text", []) + + # Create speaker name → numeric ID mapping + spk2id = {name: idx for idx, name in enumerate(speakers.keys())} + + # Collect prompts + prompt_text = [] + prompt_wav = [] + dialect_prompt_text = [] + + for name in speakers: + prompt_text.append(speakers[name].get("prompt_text", "")) + prompt_wav.append(speakers[name].get("prompt_audio", "")) + dialect_prompt_text.append(speakers[name].get("dialect_prompt", "")) + + # Collect dialogue text and speaker sequence + text_list = [] + spk_list = [] + for turn in text_entries: + if len(turn) == 2: + spk_name, utt_text = turn + text = f'[{spk_name}]{utt_text}' + text_list.append(text) + spk_list.append(spk2id.get(spk_name, -1)) + + # Generate time-based key + key = generate_time_index() + wav_path = os.path.join(output_dir, f"{key}.wav") + + use_dialect_prompt = False + for dialect_text in dialect_prompt_text: + if len(dialect_text) > 0: + assert check_prefix(dialect_text), f"Unknown dialect prefix: {dialect_text} \ + \n Prefix should be one of: <|Henan|>, <|Sichuan|>, <|Yue|>" + use_dialect_prompt = True + + result = { + "key": key, + "prompt_text": prompt_text, + "prompt_wav": prompt_wav, + "text": text_list, + "spk": spk_list, + "wav": wav_path, + "use_dialect_prompt": use_dialect_prompt, + "dialect_prompt_text": dialect_prompt_text + } + + return result + + +# Example usage +if __name__ == "__main__": + with open("example/podcast_script/script_henan.json", "r", encoding="utf-8") as f: + data = json.load(f) + + converted = podcast_format_parser(data) + print(json.dumps(converted, ensure_ascii=False, indent=2)) diff --git a/soulxpodcast/utils/text.py b/soulxpodcast/utils/text.py index 8a9ef1d..ecf49a5 100644 --- a/soulxpodcast/utils/text.py +++ b/soulxpodcast/utils/text.py @@ -1,82 +1,82 @@ -import re -from typing import List - -def remove_space_between_chinese(text): - # Remove spaces between consecutive Chinese characters - text = re.sub(r'([\u4e00-\u9fff])\s+([\u4e00-\u9fff])', r'\1\2', text) - - # Remove spaces between English and Chinese (English followed by Chinese) - text = re.sub(r'([a-zA-Z])\s+([\u4e00-\u9fff])', r'\1\2', text) - - # Remove spaces between Chinese and English (Chinese followed by English) - text = re.sub(r'([\u4e00-\u9fff])\s+([a-zA-Z])', r'\1\2', text) - - return text - -# Check whether the text ends with Chinese or English and add proper punctuation -def normalize_text(current_text): - # keep_punctuation='，。？！.,?!<| |>' - # pattern = f'[\\p{{P}}--[{keep_punctuation}]]' - # current_text = re.sub(pattern, '', current_text) - - # Remove spaces between consecutive Chinese characters - current_text = re.sub(r'([\u4e00-\u9fff])\s+([\u4e00-\u9fff])', r'\1\2', current_text) - - # Remove spaces between English and Chinese (English followed by Chinese) - current_text = re.sub(r'([a-zA-Z])\s+([\u4e00-\u9fff])', r'\1\2', current_text) - - # Remove spaces between Chinese and English (Chinese followed by English) - current_text = re.sub(r'([\u4e00-\u9fff])\s+([a-zA-Z])', r'\1\2', current_text) - - # Check if the text ends with a Chinese character - if re.search(r'[\u4e00-\u9fff]$', current_text): # 中文字符结尾 - # If the last character is not a punctuation mark, add a full stop - if current_text[-1] not in ",.?!。，？！": - current_text += "。" - - # Check if the text ends with an English character - elif re.search(r'[a-zA-Z]$', current_text): # Ends with English - # If the last character is not a punctuation mark, add a period - if current_text[-1] not in ".!?": - current_text += "." - - return current_text - - -def check_monologue_text(text: str, prefix: str = None) -> bool: - text = text.strip() - # Check speaker tags - if prefix is not None and (not text.startswith(prefix)): - return False - # Remove prefix - if prefix is not None: - text = text.removeprefix(prefix) - text = text.strip() - # If empty? - if len(text) == 0: - return False - return True - -def check_dialect_prompt_text(text: str, prefix: str = None) -> bool: - text = text.strip() - # Check COT prefix tags - if prefix is not None and (not text.startswith(prefix)): - return False - text = text.strip() - # If empty? - if len(text) == 0: - return False - return True - -def check_dialogue_text(text_list: List[str]) -> bool: - if len(text_list) == 0: - return False - for text in text_list: - if not ( - check_monologue_text(text, "[S1]") - or check_monologue_text(text, "[S2]") - or check_monologue_text(text, "[S3]") - or check_monologue_text(text, "[S4]") - ): - return False +import re +from typing import List + +def remove_space_between_chinese(text): + # Remove spaces between consecutive Chinese characters + text = re.sub(r'([\u4e00-\u9fff])\s+([\u4e00-\u9fff])', r'\1\2', text) + + # Remove spaces between English and Chinese (English followed by Chinese) + text = re.sub(r'([a-zA-Z])\s+([\u4e00-\u9fff])', r'\1\2', text) + + # Remove spaces between Chinese and English (Chinese followed by English) + text = re.sub(r'([\u4e00-\u9fff])\s+([a-zA-Z])', r'\1\2', text) + + return text + +# Check whether the text ends with Chinese or English and add proper punctuation +def normalize_text(current_text): + # keep_punctuation='，。？！.,?!<| |>' + # pattern = f'[\\p{{P}}--[{keep_punctuation}]]' + # current_text = re.sub(pattern, '', current_text) + + # Remove spaces between consecutive Chinese characters + current_text = re.sub(r'([\u4e00-\u9fff])\s+([\u4e00-\u9fff])', r'\1\2', current_text) + + # Remove spaces between English and Chinese (English followed by Chinese) + current_text = re.sub(r'([a-zA-Z])\s+([\u4e00-\u9fff])', r'\1\2', current_text) + + # Remove spaces between Chinese and English (Chinese followed by English) + current_text = re.sub(r'([\u4e00-\u9fff])\s+([a-zA-Z])', r'\1\2', current_text) + + # Check if the text ends with a Chinese character + if re.search(r'[\u4e00-\u9fff]$', current_text): # 中文字符结尾 + # If the last character is not a punctuation mark, add a full stop + if current_text[-1] not in ",.?!。，？！": + current_text += "。" + + # Check if the text ends with an English character + elif re.search(r'[a-zA-Z]$', current_text): # Ends with English + # If the last character is not a punctuation mark, add a period + if current_text[-1] not in ".!?": + current_text += "." + + return current_text + + +def check_monologue_text(text: str, prefix: str = None) -> bool: + text = text.strip() + # Check speaker tags + if prefix is not None and (not text.startswith(prefix)): + return False + # Remove prefix + if prefix is not None: + text = text.removeprefix(prefix) + text = text.strip() + # If empty? + if len(text) == 0: + return False + return True + +def check_dialect_prompt_text(text: str, prefix: str = None) -> bool: + text = text.strip() + # Check COT prefix tags + if prefix is not None and (not text.startswith(prefix)): + return False + text = text.strip() + # If empty? + if len(text) == 0: + return False + return True + +def check_dialogue_text(text_list: List[str]) -> bool: + if len(text_list) == 0: + return False + for text in text_list: + if not ( + check_monologue_text(text, "[S1]") + or check_monologue_text(text, "[S2]") + or check_monologue_text(text, "[S3]") + or check_monologue_text(text, "[S4]") + ): + return False return True \ No newline at end of file