Skip to content

feat: add spark tts BiCodec and spark tts support batch stream - #130

Merged
weedge merged 12 commits into
mainfrom
feat/tts
Mar 17, 2025
Merged

weedge merged 12 commits into
mainfrom
feat/tts

Conversation

@weedge

@weedge weedge commented Mar 16, 2025

Copy link
Copy Markdown
Collaborator

feat:

  • add spark tts BiCodec and unit test
CODEC_TAG=codec_bitokenizer CODEC_MODEL_DIR=./models/SparkAudio/Spark-TTS-0.5B \
    python -m unittest test.modules.codec.test.TestCodec.test_encode_decode
# ---- Inference Overview of Voice Cloning ----
python -m unittest test.modules.speech.tts.test_spark.TestSparkTTS.test_get_voices
REF_AUDIO_PATH=./test/audio_files/asr_example_zh.wav \
    REF_TEXT="欢迎大家来体验达摩院推出的语音识别模型" \
    python -m unittest test.modules.speech.tts.test_spark.TestSparkTTS.test_set_voice

REF_AUDIO_PATH=./test/audio_files/asr_example_zh.wav \
    REF_TEXT="欢迎大家来体验达摩院推出的语音识别模型" \
    TTS_TEXT="万物之始,大道至简,衍化至繁。君不见黄河之水天上来,奔流到海不复回。君不见高堂明镜悲白发,朝如青丝暮成雪。人生得意须尽欢,莫使金樽空对月。天生我材必有用,千金散尽还复来。" \
    python -m unittest test.modules.speech.tts.test_spark.TestSparkTTS.test_synthesize

REF_AUDIO_PATH=./test/audio_files/asr_example_zh.wav \
    REF_TEXT="欢迎大家来体验达摩院推出的语音识别模型" \
    TTS_TEXT="万物之始,大道至简,衍化至繁。" \
    python -m unittest test.modules.speech.tts.test_spark.TestSparkTTS.test_synthesize

REF_AUDIO_PATH=./test/audio_files/asr_example_zh.wav \
    REF_TEXT="欢迎大家来体验达摩院推出的语音识别模型" \
    TTS_TEXT="万物之始,大道至简,衍化至繁。人生得意须尽欢,莫使金樽空对月。天生我材必有用,千金散尽还复来。" \
    python -m unittest test.modules.speech.tts.test_spark.TestSparkTTS.test_synthesize_speak

# ---- Inference Overview of Controlled Generation with gender ----

python -m unittest test.modules.speech.tts.test_spark.TestSparkTTS.test_synthesize
python -m unittest test.modules.speech.tts.test_spark.TestSparkTTS.test_synthesize_speak

TTS_TEXT="万物之始,大道至简,衍化至繁。君不见黄河之水天上来,奔流到海不复回。君不见高堂明镜悲白发,朝如青丝暮成雪。人生得意须尽欢,莫使金樽空对月。天生我材必有用,千金散尽还复来。" \
    python -m unittest test.modules.speech.tts.test_spark.TestSparkTTS.test_synthesize

TTS_TEXT="万物之始,大道至简,衍化至繁。人生得意须尽欢,莫使金樽空对月。天生我材必有用,千金散尽还复来。" \
    python -m unittest test.modules.speech.tts.test_spark.TestSparkTTS.test_synthesize_speak
  • add spark tts grpc test
TTS_TAG=tts_spark IS_SAVE=1 IS_RELOAD=1 \
    TTS_MODEL_DIR=./models/SparkAudio/Spark-TTS-0.5B \
    TTS_LM_MODEL_PATH=./models/SparkAudio/Spark-TTS-0.5B/LLM \
    python -m src.cmd.grpc.speaker.client

cpu/gpu test is ok
colab(T4): https://github.com/weedge/doraemon-nb/blob/main/achatbot_spark_tts.ipynb


Inference Overview of Voice Cloning
Inference Overview of Controlled Generation

Note

  • Inference Overview of Voice Cloning: global_token_ids used by ref audio(audio mel-spec encode token) and generate audio(token decode waveform).
  • Inference Overview of Controlled Generation (gender) : global_token_ids used by generate audio(token decode waveform).

Model struct: (KISS LLM)

image

Spark-TTS 训练过程还未公开,论文中的介绍:

Spark-TTS 语言模型使用整个 VoxBox 训练集进行训练。如果数据集缺少预定义的训练/测试分割,使用整个处理后的数据集进行训练。训练采用 AdamW 优化器,带有 β 1 = 0.9 和 β 2 = 0.96 。该模型在 3 个时期内进行训练,使用 768 个样本的批处理大小。

BiCodec :
image

BiCodec 训练过程还未公开,论文中的介绍:(后续拿 DAC: ⭐️ 2023.6 High-Fidelity Audio Compression with Improved RVQGAN | paper code 的训练框架来训练一下,数据集是公开的,应该可以复现,训练参数规模不大,使用单卡,或者单机多卡 DDP 训练下(pytorch中DDP包括了DP))

BiCodec 在 LibriSpeech 数据集的完整训练集上进行训练,包含 960 小时的英语语音数据。此外,还包含了 Emilia-CN 和 Emilia-EN 的 1,000 小时语音数据,使总训练数据达到约 3,000 小时。所有音频样本都重新采样为 16 kHz。全局标记长度设置为 32。 为了进行优化,使用 AdamW 优化器,其移动平均系数为 β 1 = 0.8 和 β 2 = 0.9 。该模型使用具有 614.4 秒语音的批处理大小,在大约 800k 个训练步骤内收敛。

论文中提到的内部SparkVox训练框架结构,未公开源码,关注数据处理和训练部分
image

数据部分见论文中的附录

Wav2Vec :

image

Speaker Encoder(ECAPA-TDNN):

image

ConvNet: 这块可以对比 f5-tts 的结构 #101

  • ⭐️ 2022. A ConvNet for the 2020s | paper code (对ResNet-200进行改进,按照ViT(encoder)变体 Swin Transformer设计,替换MSA(multiheaded self-attention)-> 7x7 conv2d; MLP(linear->1x1 kernel conv2d,active function ReLU->GeLU), BN->LN, 纯 ConvNet 模型结构, 性能和Swin transformer相当, 但是模型结构轻巧,推理更快; 但是ConvNeXt 可能更适合某些任务,比如: 图像分类、对象检测、实例和语义分割任务; 而 Transformers 对于其他任务可能更灵活,泛化能力强,当用于需要离散、稀疏或结构化输出的任务时,Transformer 可能会更加灵活。所以架构选择应该满足手头任务的需求,同时力求简单。)

CrossAttention: learnable queries (nn.Parameter)

image

在图像生成模型中 Stable Diffusion model (DiT) 也使用了CrossAttention

PS: 算法部分大部分是子神经网络的拼装组合,但是怎么个拼接法,而且可以扩展,这个还是需要多看论文,搞懂子模块为啥需要这么设计和拼装。这个大部分靠论文取胜了,读得多,根据要解决的问题,灵感自然而然会来(特别是多模态的论文一些模块会使用基础模块,需要弄懂设计所要解决的问题和解决方法,而且图像和语音模态的基础模块是可以相互借鉴的)

Signed-off-by: weedge <weege007@gmail.com>

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hello @weedge, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

This pull request introduces a new BiCodec for SparkTTS, leveraging Wav2Vec2 for feature extraction. The BiCodec model is based on the architecture described in the Spark-TTS paper, utilizing semantic and global tokens for speech synthesis. The implementation includes the BiCodec model itself, along with the Wav2Vec2 feature extractor, and necessary utilities for loading configurations and extracting audio features. The changes also update the README with instructions for setting up the model repository and running benchmarks, and fix a typo in the tritonserver.py file.

Highlights

  • BiCodec Implementation: Adds the core BiCodec model for text-to-speech synthesis, based on the SparkTTS architecture.
  • Wav2Vec2 Integration: Integrates Wav2Vec2 feature extractor for semantic token extraction, enhancing the BiCodec model's capabilities.
  • Configuration and Utilities: Includes utilities for loading model configurations and extracting audio features, streamlining the setup and usage of the BiCodec model.
  • Documentation Updates: Updates the README with instructions for setting up the model repository and running benchmarks, providing guidance for users.

Changelog

  • deploy/modal/src/llm/trtllm/tts_spark/README.md
    • Updates the README to include instructions for filling parameters in the model configuration.
    • Adds a 'todo' section for future benchmarking on different GPU architectures and using Perf Analyzer.
    • Adds links to pb_stub documentation for InferenceRequest and InferenceResponse.
  • deploy/modal/src/llm/trtllm/tts_spark/tritonserver.py
    • Changes the curl command to check for server readiness instead of server liveness.
  • src/modules/codec/audio/bicodec.py
    • Adds a new BiCodecTokenizer class implementing the ICodec interface.
    • Includes Wav2Vec2 feature extraction for semantic token generation.
    • Implements encode_code and decode_code methods for tokenizing and detokenizing audio waveforms.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in issue comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.


Did you know?

The field of speech synthesis has roots dating back to the late 18th century with early attempts to create speech-generating machines, but modern TTS systems rely heavily on machine learning and deep learning techniques.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

The pull request introduces a new BiCodec for Spark TTS, integrating Wav2Vec2 for feature extraction. The code includes necessary dependencies and implements encoding and decoding functionalities. Overall, the implementation seems well-structured, but there are a few areas that could benefit from further attention.

Summary of Findings

  • Error Handling: The ModuleNotFoundError exception handling could be improved by providing more specific guidance to the user, such as listing the exact packages to install.
  • Code Clarity: Consider adding comments to explain the purpose of specific calculations or operations, especially in the get_ref_clip and extract_wav2vec2_features methods.
  • Logging: Add more logging to indicate the start and end of key processes, such as model loading and feature extraction, to aid in debugging and monitoring.

Merge Readiness

The pull request introduces a new feature and seems generally well-structured. However, addressing the identified issues related to error handling, code clarity, and logging would improve the overall quality and maintainability of the code. I recommend addressing these points before merging. I am unable to directly approve this pull request, and other reviewers should also review this code before merging.

Comment thread deploy/modal/src/llm/trtllm/tts_spark/README.md Outdated
Comment thread src/modules/codec/audio/bicodec.py Outdated
Comment thread src/modules/codec/audio/bicodec.py
Comment thread src/modules/codec/audio/bicodec.py Outdated
weedge and others added 7 commits March 16, 2025 22:59
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: weedge <weege007@gmail.com>
Signed-off-by: weedge <weege007@gmail.com>
Signed-off-by: weedge <weege007@gmail.com>
Signed-off-by: weedge <weege007@gmail.com>
@weedge weedge changed the title feat: add spark tts BiCodec feat: add spark tts BiCodec and spark tts support batch stream Mar 17, 2025
weedge added 4 commits March 18, 2025 00:13
Signed-off-by: weedge <weege007@gmail.com>
Signed-off-by: weedge <weege007@gmail.com>
Signed-off-by: weedge <weege007@gmail.com>
Signed-off-by: weedge <weege007@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant