diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..d6403048500ec8b54d6c7a2aa52ec1a3204cb9f5 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,32 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +ref_1320.wav filter=lfs diff=lfs merge=lfs -text +ref_2094.wav filter=lfs diff=lfs merge=lfs -text +ref_260.wav filter=lfs diff=lfs merge=lfs -text +ref_3575.wav filter=lfs diff=lfs merge=lfs -text +ref_5639.wav filter=lfs diff=lfs merge=lfs -text +ref_6930.wav filter=lfs diff=lfs merge=lfs -text +ref_7127.wav filter=lfs diff=lfs merge=lfs -text +ref_7729.wav filter=lfs diff=lfs merge=lfs -text +samples/1_librispeech_1320.wav filter=lfs diff=lfs merge=lfs -text +samples/1_librispeech_2094.wav filter=lfs diff=lfs merge=lfs -text +samples/1_librispeech_260.wav filter=lfs diff=lfs merge=lfs -text +samples/1_paradee.wav filter=lfs diff=lfs merge=lfs -text +samples/2_librispeech_1320.wav filter=lfs diff=lfs merge=lfs -text +samples/2_librispeech_2094.wav filter=lfs diff=lfs merge=lfs -text +samples/2_librispeech_260.wav filter=lfs diff=lfs merge=lfs -text +samples/2_paradee.wav filter=lfs diff=lfs merge=lfs -text +samples/3_librispeech_1320.wav filter=lfs diff=lfs merge=lfs -text +samples/3_librispeech_2094.wav filter=lfs diff=lfs merge=lfs -text +samples/3_librispeech_260.wav filter=lfs diff=lfs merge=lfs -text +samples/3_paradee.wav filter=lfs diff=lfs merge=lfs -text +samples/4_librispeech_1320.wav filter=lfs diff=lfs merge=lfs -text +samples/4_librispeech_2094.wav filter=lfs diff=lfs merge=lfs -text +samples/4_librispeech_260.wav filter=lfs diff=lfs merge=lfs -text +samples/4_paradee.wav filter=lfs diff=lfs merge=lfs -text +samples/5_librispeech_1320.wav filter=lfs diff=lfs merge=lfs -text +samples/5_librispeech_2094.wav filter=lfs diff=lfs merge=lfs -text +samples/5_librispeech_260.wav filter=lfs diff=lfs merge=lfs -text +samples/5_paradee.wav filter=lfs diff=lfs merge=lfs -text +samples/vc_1320_to_2094.wav filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d9e634b96bde1f421ec81210ad7ac839e08db509 --- /dev/null +++ b/README.md @@ -0,0 +1,260 @@ +--- +license: apache-2.0 +language: + - en +pipeline_tag: text-to-speech +library_name: pytorch +base_model: + - sahilmahendrakar/Paradee-8M-v1.1 + - ASLP-lab/MeanVC2 +tags: + - text-to-speech + - tts + - zero-shot + - voice-cloning + - voice-conversion + - streaming + - flow-matching + - on-device +--- + +# Silencio + +**Small zero-shot text-to-speech that speaks in any voice.** + +Silencio turns English text into 16 kHz speech in a voice it has never heard before. A few seconds +of reference audio are enough. All 33.6M parameters fit in one 134 MB `model.safetensors`. It renders +audio in 120 ms chunks and runs about 9x faster than real time on one CPU thread of an Apple M4. + +The same weights also do voice conversion: give it speech instead of text and it says the same words in +another voice. + +Silencio is built on two open models. The text front-end comes from +[Paradee-8M v1.1](https://huggingface.co/sahilmahendrakar/Paradee-8M-v1.1) by Sahil Mahendrakar, +itself distilled from [Kokoro-82M](https://huggingface.co/hexgrad/Kokoro-82M) by hexgrad. The +voice renderer and vocoder come from [MeanVC2](https://huggingface.co/ASLP-lab/MeanVC2) by ASLP-lab. +A new 3.4M-parameter content adapter connects them, so text goes straight into the renderer with no +intermediate audio. + +## Model details + +- **Model type:** zero-shot text-to-speech. A text encoder predicts phoneme durations, pitch and + phonetic features. A ConvNeXt adapter turns them into content tokens at 25 Hz. A mean-flow DiT + renders a mel spectrogram from those tokens and a 256-d voice embedding in 3 steps. Vocos then + makes the waveform. +- **Parameters:** 33.6M (text encoder 4.2M, adapter 3.4M, renderer 17.7M, vocoder 8.3M) +- **Weights file:** `model.safetensors`, 134 MB (fp32) +- **Language:** English (American pronunciation) +- **Voices:** zero-shot, from a 5–15 s reference clip. 8 ready-made voices are included. +- **Output:** mono, 16 kHz +- **License:** [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0) + +### How it fits together + +``` +text ─► misaki G2P ─► text encoder ─► content adapter ─► DiT renderer ─► Vocos ─► 16 kHz audio + (phonemes) durations, F0, content tokens mean flow, ▲ + features @40 Hz @25 Hz 3 steps │ + ▲ │ + voice embedding ──┘ (256-d) │ +speech ─► Fast-U2++ streaming ASR encoder ─► content tokens ─┘ (voice conversion) +``` + +The text encoder predicts how long each phoneme lasts and how the pitch moves. The renderer only +decides how it sounds in the chosen voice. So timing and intonation come from the text, and timbre +comes from the voice embedding. + +## Quickstart + +```sh +hf download BlueRoseLabs/Silencio --local-dir Silencio +cd Silencio +pip install -r requirements.txt +pip install --no-deps --ignore-requires-python misaki==0.9.4 # G2P; declares python<3.13 but runs on 3.13 +``` + +### Python + +```python +from silencio import SilencioTTS + +tts = SilencioTTS.from_pretrained("BlueRoseLabs/Silencio") # or the local folder +tts.save("speech.wav", "Hello from Silencio.", voice="librispeech_1320") + +# Streaming: one chunk per sentence, float32 at 16 kHz +for pcm in tts.stream("Speech can start playing right away. The next sentence follows.", voice="librispeech_2094"): + ... +``` + +`generate` returns a numpy array, and `speed` (default 1.0) changes the speaking rate. On first use misaki downloads spaCy's small English model (about 13 MB). + +### Any voice from a short clip + +```python +import soundfile as sf +from silencio.enroll import enroll # pip install s3prl + +voice = enroll(*sf.read("reference.wav")) # 256-d embedding; save it with torch.save and reuse it +tts.save("cloned.wav", "This is my voice, read by Silencio.", voice=voice) +``` + +The voice encoder is WavLM-Large with ECAPA-TDNN, the same one MeanVC2 was trained with. It is large +(about 1.3 GB) and only needed to add new voices, so it is not part of `model.safetensors`. It is +downloaded on first use. A saved embedding needs nothing extra. + +### Voice conversion + +```python +wav, sr = sf.read("source.wav") +sf.write("converted.wav", tts.convert(wav, sr, voice="librispeech_2094"), 16000) +``` + +For speech input, the content tokens come from MeanVC2's streaming Fast-U2++ encoder. It is +downloaded from [ASLP-lab/MeanVC2](https://huggingface.co/ASLP-lab/MeanVC2) on first use and runs on CPU. + +## Voices + +`voices.safetensors` holds 8 voices enrolled from LibriSpeech test-clean readers. A reader's voice +is called `librispeech_`. The reference clip used for each one is included as `ref_.wav`. + +| Voice | Reference clip | +|---|---| +| `librispeech_1320` (default) | [ref_1320.wav](ref_1320.wav) | +| `librispeech_2094` | [ref_2094.wav](ref_2094.wav) | +| `librispeech_260` | [ref_260.wav](ref_260.wav) | +| `librispeech_3575` | [ref_3575.wav](ref_3575.wav) | +| `librispeech_5639` | [ref_5639.wav](ref_5639.wav) | +| `librispeech_6930` | [ref_6930.wav](ref_6930.wav) | +| `librispeech_7127` | [ref_7127.wav](ref_7127.wav) | +| `librispeech_7729` | [ref_7729.wav](ref_7729.wav) | + +## Samples + +Five held-out sentences, read by Silencio in three voices. The last column is +Paradee-8M v1.1, whose text encoder Silencio uses, in its single voice, for comparison. The sentences are +in [`samples/sentences.txt`](samples/sentences.txt). + +| # | `librispeech_1320` | `librispeech_2094` | `librispeech_260` | Paradee-8M v1.1 | +|---|---|---|---|---| +| 1 | [listen](samples/1_librispeech_1320.wav) | [listen](samples/1_librispeech_2094.wav) | [listen](samples/1_librispeech_260.wav) | [listen](samples/1_paradee.wav) | +| 2 | [listen](samples/2_librispeech_1320.wav) | [listen](samples/2_librispeech_2094.wav) | [listen](samples/2_librispeech_260.wav) | [listen](samples/2_paradee.wav) | +| 3 | [listen](samples/3_librispeech_1320.wav) | [listen](samples/3_librispeech_2094.wav) | [listen](samples/3_librispeech_260.wav) | [listen](samples/3_paradee.wav) | +| 4 | [listen](samples/4_librispeech_1320.wav) | [listen](samples/4_librispeech_2094.wav) | [listen](samples/4_librispeech_260.wav) | [listen](samples/4_paradee.wav) | +| 5 | [listen](samples/5_librispeech_1320.wav) | [listen](samples/5_librispeech_2094.wav) | [listen](samples/5_librispeech_260.wav) | [listen](samples/5_paradee.wav) | + +Voice conversion: [source](ref_1320.wav) (LibriSpeech 1320) → [converted to `librispeech_2094`](samples/vc_1320_to_2094.wav). Whisper transcribes both identically, and the speaker similarity to the target voice is 0.83. + +## Results + +All numbers are on Paradee's 200 held-out WikiText sentences, scored the way the Paradee card scores +them. Silencio speaks them in its 8 voices in turn. + +| Model (voice) | Params | Weights | UTMOS | WER | Speaker similarity | Speed, T4 | Speed, 1 CPU thread | +|---|---:|---:|---:|---:|---:|---:|---:| +| **Silencio (8 zero-shot voices)** | **33.6M** | **134 MB** | 4.15 | 7.0% | 0.72 | 3.7x | 8.8x (M4) · 1.6x (Colab) | +| Paradee-8M v1.1 (af_heart, single voice), PyTorch | 7.9M | 8.5 MB (int8) | 4.41 | 6.3% | – | – | – | + +- **UTMOS** is a neural network trained on human ratings. It predicts how natural a clip sounds, + on a scale from 1 to 5 (SpeechMOS `utmos22_strong`). +- **WER** is the share of words that Whisper (base) transcribes wrongly. +- **Speaker similarity** is the cosine similarity between the voice of the output and the voice of the + reference clip, measured with MeanVC2's WavLM/ECAPA encoder. 1.0 is the same embedding. +- **Speed** is how many times faster than real time the model runs, counting text to waveform, in PyTorch + fp32. CPU speed is measured with one thread, on an Apple M4 (one two-sentence paragraph) and on the Colab host (10 held-out sentences). + +Paradee was scored by the same script, in PyTorch, which is why its WER differs a little from the 5.9% on its own card. + +Paradee speaks one studio voice at 24 kHz. Silencio speaks voices taken from audiobook recordings, and +it does so at 16 kHz. Its UTMOS therefore also reflects the recording quality of the reference voices, +not only the model. + +## How it was made + +Only the adapter was trained. Everything else is reused frozen from the two base models. + +1. **Content space.** MeanVC2's renderer takes content tokens from a streaming ASR encoder (256-d at + 25 Hz) and a speaker embedding. Content tokens carry what is said and how, but not who says it. +2. **Targets.** Paradee read 6,000 WikiText sentences at speeds 0.9, 1.0 and 1.1 (about 6 hours). For + each one, Paradee's internal text-side features were kept, and MeanVC2's ASR encoder was run on + Paradee's audio to get the content tokens that audio would produce. +3. **Adapter.** A ConvNeXt network (4 blocks at 40 Hz, resampling to 25 Hz, then 4 more blocks) learns to + map the text features, pitch and loudness to those content tokens. It uses an L1 plus cosine loss. + Training took 15,000 steps at batch 32, about 22 minutes on one T4. +4. **One model.** The text encoder, adapter, renderer (EMA weights) and vocoder were put into a single + `model.safetensors`. Weight norm was folded into the weights, and the Vocos TorchScript was converted + to plain PyTorch. The output of the merged model matches the original chained pipeline to within + 7e-4. + +In speech-to-speech mode the original ASR encoder takes the adapter's place, so voice conversion works +exactly as in MeanVC2. + +## Files + +| File | What it is | +|---|---| +| `model.safetensors` | The whole model: text encoder, content adapter, DiT renderer, Vocos vocoder (134 MB, fp32) | +| `config.json` | Phoneme vocabulary, module sizes, chunking and the default voice | +| `voices.safetensors` | 8 ready-made voice embeddings (256-d each) | +| `ref_*.wav` | The LibriSpeech clips those voices were enrolled from | +| `silencio/` | The inference code. Upstream modules are vendored under their Apache-2.0 licenses, with only import paths changed | +| `requirements.txt` | Runtime dependencies | +| `samples/` | The audio above | +| `eval.json` | The numbers in Results | + +## Limitations + +- English only, with American pronunciation. Numbers, abbreviations and rare words are pronounced only + as well as misaki handles them. +- Output is 16 kHz, which is the rate of MeanVC2's vocoder. That is enough for voice assistants + and audiobooks, but not for full-band audio. +- Prosody comes from Paradee's single training voice, so every voice speaks with similar timing and + intonation. Speaker similarity is lower than in voice conversion, where timing comes from the + source speaker. +- Long text is split into sentences that are spoken one after another, so prosody does not carry + across sentences. + +## Responsible use + +Do not use Silencio to impersonate a real person, to deceive listeners about who is speaking, or +to clone or convert a voice without the speaker's consent. Tell listeners when speech is synthetic where +that matters to them. + +## License + +The weights and code in this repository are released under +[Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0), the license of every model they are built +from: + +- [Paradee-8M v1.1](https://huggingface.co/sahilmahendrakar/Paradee-8M-v1.1) by Sahil Mahendrakar + (Apache 2.0): text encoder. +- [Kokoro-82M](https://huggingface.co/hexgrad/Kokoro-82M) by hexgrad (Apache 2.0): the architecture and + teacher behind Paradee, and the [misaki](https://github.com/hexgrad/misaki) G2P. +- [MeanVC2](https://huggingface.co/ASLP-lab/MeanVC2) by ASLP-lab (Apache 2.0): DiT renderer, Vocos + vocoder, Fast-U2++ encoder and the speaker-encoder code. +- [Vocos](https://github.com/gemelo-ai/vocos) by Gemelo AI (MIT): vocoder architecture. + +The voice encoder used by `enroll` (WavLM-Large from +[UniSpeech](https://github.com/microsoft/UniSpeech), with its ECAPA-TDNN speaker-verification +fine-tune) is not redistributed here. It is downloaded at runtime under its own license. The bundled +voices come from [LibriSpeech](https://www.openslr.org/12) (CC BY 4.0) readers, whose recordings are from +the public-domain LibriVox project. + +## Citation + +```bibtex +@misc{silencio2026, + title = {Silencio: Zero-Shot Text-to-Speech by Connecting a Distilled Text Encoder to a Streaming Voice Renderer}, + author = {{BlueRose Labs}}, + year = {2026}, + url = {https://huggingface.co/BlueRoseLabs/Silencio} +} + +@misc{mahendrakar2026paradee, + title = {Paradee: Distilling Kokoro-82M into an 8M-Parameter Single-Voice Text-to-Speech Model}, + author = {Mahendrakar, Sahil}, + year = {2026}, + url = {https://huggingface.co/sahilmahendrakar/Paradee-8M-v1.0} +} +``` + +MeanVC2: [github.com/ASLP-lab/MeanVC2](https://github.com/ASLP-lab/MeanVC2). diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..688181d2a9de0edce197e7f4cf8c6ccec8f4dd00 --- /dev/null +++ b/config.json @@ -0,0 +1,149 @@ +{ + "model_type": "silencio-tts", + "version": "1.0", + "sample_rate": 16000, + "vocab": { + ";": 1, + ":": 2, + ",": 3, + ".": 4, + "!": 5, + "?": 6, + "—": 9, + "…": 10, + "\"": 11, + "(": 12, + ")": 13, + "“": 14, + "”": 15, + " ": 16, + "̃": 17, + "ʣ": 18, + "ʥ": 19, + "ʦ": 20, + "ʨ": 21, + "ᵝ": 22, + "ꭧ": 23, + "A": 24, + "I": 25, + "O": 31, + "Q": 33, + "S": 35, + "T": 36, + "W": 39, + "Y": 41, + "ᵊ": 42, + "a": 43, + "b": 44, + "c": 45, + "d": 46, + "e": 47, + "f": 48, + "h": 50, + "i": 51, + "j": 52, + "k": 53, + "l": 54, + "m": 55, + "n": 56, + "o": 57, + "p": 58, + "q": 59, + "r": 60, + "s": 61, + "t": 62, + "u": 63, + "v": 64, + "w": 65, + "x": 66, + "y": 67, + "z": 68, + "ɑ": 69, + "ɐ": 70, + "ɒ": 71, + "æ": 72, + "β": 75, + "ɔ": 76, + "ɕ": 77, + "ç": 78, + "ɖ": 80, + "ð": 81, + "ʤ": 82, + "ə": 83, + "ɚ": 85, + "ɛ": 86, + "ɜ": 87, + "ɟ": 90, + "ɡ": 92, + "ɥ": 99, + "ɨ": 101, + "ɪ": 102, + "ʝ": 103, + "ɯ": 110, + "ɰ": 111, + "ŋ": 112, + "ɳ": 113, + "ɲ": 114, + "ɴ": 115, + "ø": 116, + "ɸ": 118, + "θ": 119, + "œ": 120, + "ɹ": 123, + "ɾ": 125, + "ɻ": 126, + "ʁ": 128, + "ɽ": 129, + "ʂ": 130, + "ʃ": 131, + "ʈ": 132, + "ʧ": 133, + "ʊ": 135, + "ʋ": 136, + "ʌ": 138, + "ɣ": 139, + "ɤ": 140, + "χ": 142, + "ʎ": 143, + "ʒ": 147, + "ʔ": 148, + "ˈ": 156, + "ˌ": 157, + "ː": 158, + "ʰ": 162, + "ʲ": 164, + "↓": 169, + "→": 171, + "↗": 172, + "↘": 173, + "ᵻ": 177 + }, + "adapter": { + "dim": 256 + }, + "dit": { + "dim": 512, + "depth": 4, + "heads": 2, + "ff_mult": 2, + "bn_dim": 256, + "conv_layers": 4, + "chunk_size": 12, + "block_size": 4, + "dropout": 0.0, + "qk_norm": "rms_norm" + }, + "vocos": { + "input_channels": 80, + "dim": 320, + "intermediate_dim": 1536, + "num_layers": 8, + "n_fft": 640, + "hop_length": 160, + "padding": "center" + }, + "chunk_size": 12, + "block_size": 4, + "steps": 3, + "default_voice": "librispeech_6930" +} \ No newline at end of file diff --git a/eval.json b/eval.json new file mode 100644 index 0000000000000000000000000000000000000000..6dc99f37b58c4a5264457f9a79f9119a269416fc --- /dev/null +++ b/eval.json @@ -0,0 +1,25 @@ +{ + "n": 200, + "voices": [ + "librispeech_1320", + "librispeech_2094", + "librispeech_260", + "librispeech_3575", + "librispeech_5639", + "librispeech_6930", + "librispeech_7127", + "librispeech_7729" + ], + "params_m": 33.61, + "silencio": { + "wer": 0.06953260242354299, + "utmos": 4.151572382450103, + "spk_sim": 0.7226772201061249, + "rtf_t4": 0.2737586687655075, + "rtf_cpu_1thread": 0.6323270854857574 + }, + "paradee_v1_1": { + "wer": 0.06318522792844777, + "utmos": 4.407845231294632 + } +} \ No newline at end of file diff --git a/model.safetensors b/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..2a6005d926a4d68c6f3ef372b19548bfb5c27ed1 --- /dev/null +++ b/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:19f6c837e93857af3d91415e88c08a89f4535f4e0e56e9b313bc87fa5cc24276 +size 134476848 diff --git a/ref_1320.wav b/ref_1320.wav new file mode 100644 index 0000000000000000000000000000000000000000..b533dbc6d1113fabb2a5689dfaad2071a1904b18 --- /dev/null +++ b/ref_1320.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:27e4cf2998ad8b13279721d2d5e21c6a2f9c232b4882a75f7857f760a6899453 +size 250764 diff --git a/ref_2094.wav b/ref_2094.wav new file mode 100644 index 0000000000000000000000000000000000000000..179eca4e8f00cd9651292b24623508b667d504fb --- /dev/null +++ b/ref_2094.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:72fab69f58aa47a6cb82c13c8fcb23979098329e52c42db69ac3a8a5649c123d +size 722284 diff --git a/ref_260.wav b/ref_260.wav new file mode 100644 index 0000000000000000000000000000000000000000..b9fa2748900d8f3c844ee260bc288aa3dda1a202 --- /dev/null +++ b/ref_260.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:69df92e49522c52d424c18c1a560295523ca46ddc5e43e11829fed38118d4692 +size 470924 diff --git a/ref_3575.wav b/ref_3575.wav new file mode 100644 index 0000000000000000000000000000000000000000..7250849365d02c034d432a35de9d78569187bec1 --- /dev/null +++ b/ref_3575.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6aaf21d7747b079cb4f860823e2a9f4b09809474dd7f32754c9bf91a2313d72b +size 263404 diff --git a/ref_5639.wav b/ref_5639.wav new file mode 100644 index 0000000000000000000000000000000000000000..710c2f9ad818f655598e593bab190225acbc8614 --- /dev/null +++ b/ref_5639.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8069fd8ee2a120f435d682a8e27a9435fa6b80122b5d8843779178d19a3c1c2a +size 504684 diff --git a/ref_6930.wav b/ref_6930.wav new file mode 100644 index 0000000000000000000000000000000000000000..f9d05dda0696484db7996d9d54d1559ec9da0fe7 --- /dev/null +++ b/ref_6930.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7a0c28ad2cca523cd7e7d72e360a3f314eeb905688358dbf6f150c0a87b87f76 +size 455244 diff --git a/ref_7127.wav b/ref_7127.wav new file mode 100644 index 0000000000000000000000000000000000000000..615aa6ec56ae642f8ed1f83ed0c36ad8ffc9c4ee --- /dev/null +++ b/ref_7127.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5b5d5625b3cac60e98bddec49206b053d64d7c0d0d871090801b26f8a2755bce +size 575084 diff --git a/ref_7729.wav b/ref_7729.wav new file mode 100644 index 0000000000000000000000000000000000000000..a44330db89d095d5181b4251b2d41c1f7477a8b5 --- /dev/null +++ b/ref_7729.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f8e9574a4743df5adca62c99b13aeee6cd48082aed075708545a898435aabde3 +size 265644 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..4ddea49753edfda2dcf7eada7a50d33fd6bc9bb5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,21 @@ +# Silencio runtime. misaki (the G2P) declares python<3.13 but runs on 3.13; install it with +# pip install --no-deps --ignore-requires-python misaki==0.9.4 +torch>=2.5 +torchaudio>=2.5 +numpy +huggingface-hub +safetensors +transformers +x-transformers==2.2.11 +einops +librosa +soundfile +# misaki's own dependencies +addict +regex +spacy +num2words +phonemizer-fork +espeakng-loader +# only for enrolling new voices (silencio.enroll) +s3prl diff --git a/samples/1_librispeech_1320.wav b/samples/1_librispeech_1320.wav new file mode 100644 index 0000000000000000000000000000000000000000..8f1dd6ec77f6f148e7a21066e2865c35cdb11b8d --- /dev/null +++ b/samples/1_librispeech_1320.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:86acc5e690d395b78d16d99cc8381af523e8a061036a89f44b7258883ce020da +size 306924 diff --git a/samples/1_librispeech_2094.wav b/samples/1_librispeech_2094.wav new file mode 100644 index 0000000000000000000000000000000000000000..2898f041642cb7284641bc5667963a53bf925fee --- /dev/null +++ b/samples/1_librispeech_2094.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ef3d327c72599c3862c772a75b9143c7976079ba830a3b9bba5835bc665713c7 +size 306924 diff --git a/samples/1_librispeech_260.wav b/samples/1_librispeech_260.wav new file mode 100644 index 0000000000000000000000000000000000000000..34a1d93293619de6f04371a1592e2f174fba7b15 --- /dev/null +++ b/samples/1_librispeech_260.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aeed717d7e2d0f6bd98bf628429151d6687aa668a2cb023eaa4bdec72118b1d1 +size 306924 diff --git a/samples/1_paradee.wav b/samples/1_paradee.wav new file mode 100644 index 0000000000000000000000000000000000000000..38680c60a4e4a02a60383b05dadec2a3eb2822ae --- /dev/null +++ b/samples/1_paradee.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:879b44188d8e1d043b6f85f07d89ebc91047f4b40c1143991d5dc778aa7ea5f9 +size 466844 diff --git a/samples/2_librispeech_1320.wav b/samples/2_librispeech_1320.wav new file mode 100644 index 0000000000000000000000000000000000000000..6f1acc6362a13383be21a683d0ab3fcb000585a8 --- /dev/null +++ b/samples/2_librispeech_1320.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4399e04c86c1e83ad22e258fbdc0161f71a56ebe5ad1786d393fb5f141dd63f2 +size 184044 diff --git a/samples/2_librispeech_2094.wav b/samples/2_librispeech_2094.wav new file mode 100644 index 0000000000000000000000000000000000000000..9f6a1b78fccca6deef90621d71508b5b3a9c1056 --- /dev/null +++ b/samples/2_librispeech_2094.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:769d8d22e3a93a5ed42be25f0ebd7629273aaee1bf837eca73217895a0d7c425 +size 184044 diff --git a/samples/2_librispeech_260.wav b/samples/2_librispeech_260.wav new file mode 100644 index 0000000000000000000000000000000000000000..f9ef8de484a6c57f53d05e312f76ed6520b06ba0 --- /dev/null +++ b/samples/2_librispeech_260.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c275d5246288a9d71bcca0d0d90c87e2722d2f415d2579d76eca3c71006dcca8 +size 184044 diff --git a/samples/2_paradee.wav b/samples/2_paradee.wav new file mode 100644 index 0000000000000000000000000000000000000000..0a2cb64cc7a3e3507edc9f9bca0d520138f88e6e --- /dev/null +++ b/samples/2_paradee.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3e166f7a5e7d332e3fcc3b7977d4ee7bdca18591772432165d69d2f01e9cd98a +size 282044 diff --git a/samples/3_librispeech_1320.wav b/samples/3_librispeech_1320.wav new file mode 100644 index 0000000000000000000000000000000000000000..e62b3a3fab92367840af10e3fb20612f22341901 --- /dev/null +++ b/samples/3_librispeech_1320.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0a76e623642c7d3c1c42ab612335c2c7f4debab2d6e4a1a1e1c98358529877ff +size 235244 diff --git a/samples/3_librispeech_2094.wav b/samples/3_librispeech_2094.wav new file mode 100644 index 0000000000000000000000000000000000000000..8b989379b60d7d6d78d1bb18e762cce140198741 --- /dev/null +++ b/samples/3_librispeech_2094.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cc3841776e11eeb75b0472e59480646888a4300d1ed7a7cb3e394762ca500f2f +size 235244 diff --git a/samples/3_librispeech_260.wav b/samples/3_librispeech_260.wav new file mode 100644 index 0000000000000000000000000000000000000000..affec4d2f6490961c1c66f6631900537eba6ac95 --- /dev/null +++ b/samples/3_librispeech_260.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:16f17e71a6ac27d1ca93731212ebc8490f4c97b5c89a7d92912b68deaca08876 +size 235244 diff --git a/samples/3_paradee.wav b/samples/3_paradee.wav new file mode 100644 index 0000000000000000000000000000000000000000..413f4513815aa26b345929a7021fdab3a84a64c6 --- /dev/null +++ b/samples/3_paradee.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ba7d8300a36b960ed4dd95a36fa185d483e18e7cdae175737c807b92ad347e5a +size 361244 diff --git a/samples/4_librispeech_1320.wav b/samples/4_librispeech_1320.wav new file mode 100644 index 0000000000000000000000000000000000000000..1468a4dd93b7615a64a6f0320a4a5cc7e73146de --- /dev/null +++ b/samples/4_librispeech_1320.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3c32ae2b277b885cd31e5799a77c4a51ee64c5502f9f4296d9382cceb89890a5 +size 168684 diff --git a/samples/4_librispeech_2094.wav b/samples/4_librispeech_2094.wav new file mode 100644 index 0000000000000000000000000000000000000000..78e442d15a260925514b7d82480c8cb45820ab0b --- /dev/null +++ b/samples/4_librispeech_2094.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bff1793e67eb2ed58e3b0a0429d85ed5e8551d711ff53c4f258a623ff1c15fa4 +size 168684 diff --git a/samples/4_librispeech_260.wav b/samples/4_librispeech_260.wav new file mode 100644 index 0000000000000000000000000000000000000000..ce9ccc81d1845cc2a3c6386dea9154ec0d0b0c77 --- /dev/null +++ b/samples/4_librispeech_260.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:041da728500f660fc23cdd7eb2a9f2db217e0066728f0c65d3e12576474fad4f +size 168684 diff --git a/samples/4_paradee.wav b/samples/4_paradee.wav new file mode 100644 index 0000000000000000000000000000000000000000..f255d8e484433d6a1b8ee6d8a01f8c739d51b94a --- /dev/null +++ b/samples/4_paradee.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3a0c5dc7a5797da25eb03af806759c22e8e1e6cdec521ebf7151d2cab4b0aaa7 +size 261644 diff --git a/samples/5_librispeech_1320.wav b/samples/5_librispeech_1320.wav new file mode 100644 index 0000000000000000000000000000000000000000..5ba91ad326fa79a2092ed34aa0cad6999514ba51 --- /dev/null +++ b/samples/5_librispeech_1320.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e5738455475737c27bb67eebae87c6233065e87e3f035249510878e19985b53a +size 235244 diff --git a/samples/5_librispeech_2094.wav b/samples/5_librispeech_2094.wav new file mode 100644 index 0000000000000000000000000000000000000000..aef1049860d08a82063a943fe12a39f26ab30d75 --- /dev/null +++ b/samples/5_librispeech_2094.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:597f15e3edefa292ce3489574f884398236d69f042d73d8165159aad1cdbbd31 +size 235244 diff --git a/samples/5_librispeech_260.wav b/samples/5_librispeech_260.wav new file mode 100644 index 0000000000000000000000000000000000000000..858e0873b163e5f16ca3feed5e9ad198151a1b8a --- /dev/null +++ b/samples/5_librispeech_260.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fffda725e02ce7721b3b41937e57c00782e35909145cb50d034b60db25f31f4f +size 235244 diff --git a/samples/5_paradee.wav b/samples/5_paradee.wav new file mode 100644 index 0000000000000000000000000000000000000000..140138d45a7185f8a0654251cddfe29f8199744e --- /dev/null +++ b/samples/5_paradee.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:63135b92c297c2634d21277aa08e00de99aad84bdd5abcaff64b5b89eb322121 +size 358844 diff --git a/samples/sentences.txt b/samples/sentences.txt new file mode 100644 index 0000000000000000000000000000000000000000..594b2d136dcceff145e9d78c884f177f9e5ddb0a --- /dev/null +++ b/samples/sentences.txt @@ -0,0 +1,5 @@ +1. As of August 2015, there were 169 proposed targets for these goals and 304 proposed indicators to show compliance. +2. White dwarf stars, if they have a near companion, may then become Type Ia supernovae. +3. According to Herodotus, this was because they had decided in council that they could not beat the Allies in a naval battle. +4. The ashes were later taken to Veluvana, where they were enshrined in a silver stupa. +5. Ten priests of the Evangelical Lutheran Church of Finland and 90 (moderate) socialists were executed also. diff --git a/samples/vc_1320_to_2094.wav b/samples/vc_1320_to_2094.wav new file mode 100644 index 0000000000000000000000000000000000000000..afd1183ebe3f794ceb19f1a0d2f7da65112bfea9 --- /dev/null +++ b/samples/vc_1320_to_2094.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:683d8c80da6838c05a20f90591a859af647b6fbf2a4463ee56fe31e1c4e3870d +size 245484 diff --git a/silencio/__init__.py b/silencio/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..58e92da6b4cdf12f9373c6a2f663d62dbb0059cb --- /dev/null +++ b/silencio/__init__.py @@ -0,0 +1,11 @@ +"""Silencio: zero-shot English text-to-speech and voice conversion from one weights file.""" +from .adapter import Adapter, bnf_frames + +__all__ = ["Adapter", "SilencioTTS", "bnf_frames"] + + +def __getattr__(name): # torchaudio, misaki and x-transformers load only when the model is asked for + if name == "SilencioTTS": + from .tts import SilencioTTS + return SilencioTTS + raise AttributeError(name) diff --git a/silencio/adapter.py b/silencio/adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..68d3e379785ef7526fadeb29111b231b6637fac1 --- /dev/null +++ b/silencio/adapter.py @@ -0,0 +1,80 @@ +"""Trainable bridge between Paradee's text side and MeanVC2's content space. + +Paradee's text side produces, per 40 Hz frame, a 512-d phonetic feature (`asr`) plus pitch +and loudness curves at 80 Hz. MeanVC2's DiT consumes 256-d Fast-U2++ bottleneck features +(BNF) at 25 Hz. The adapter maps the first onto the second, so the TTS path skips both +Paradee's vocoder and MeanVC2's ASR encoder. +""" +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + +PARADEE_FPS = 40 # one aligned frame = 600 samples at 24 kHz +BNF_FPS = 25 # one BNF frame = 4 fbank frames of 10 ms +ASR_DIM, BNF_DIM = 512, 256 +PROSODY_CH = 6 # log-F0, voicing and loudness for the two 80 Hz sub-frames + + +def bnf_frames(seconds, window=19, stride=16): + """Number of BNF frames the streaming Fast-U2++ (160 ms preset) yields for `seconds` of 16 kHz audio.""" + samples = int(round(seconds * 16000)) + fbank = 1 + (samples - 400) // 160 if samples >= 400 else 0 + return 0 if fbank < window else ((fbank - window) // stride + 1) * (stride // 4) + + +class ConvNeXt(nn.Module): + def __init__(self, dim, mult=3, kernel=7, dropout=0.0): + super().__init__() + self.dw = nn.Conv1d(dim, dim, kernel, padding=kernel // 2, groups=dim) + self.norm = nn.LayerNorm(dim) + self.pw1, self.pw2 = nn.Linear(dim, dim * mult), nn.Linear(dim * mult, dim) + self.drop = nn.Dropout(dropout) + self.gamma = nn.Parameter(torch.full((dim,), 1e-2)) + + def forward(self, x, mask): # x [B,C,T], mask [B,1,T] + h = self.norm(self.dw(x * mask).transpose(1, 2)) + return x + (self.gamma * self.drop(self.pw2(F.gelu(self.pw1(h))))).transpose(1, 2) * mask + + +class Adapter(nn.Module): + """asr [B,512,F] @40 Hz, F0 [B,2F], N [B,2F] -> BNF [B,T,256] @25 Hz.""" + + def __init__(self, dim=256, pre_blocks=4, post_blocks=4, dropout=0.1): + super().__init__() + self.inp = nn.Conv1d(ASR_DIM + PROSODY_CH, dim, 1) + self.pre = nn.ModuleList(ConvNeXt(dim, dropout=dropout) for _ in range(pre_blocks)) + self.post = nn.ModuleList(ConvNeXt(dim, dropout=dropout) for _ in range(post_blocks)) + self.norm = nn.LayerNorm(dim) + self.out = nn.Linear(dim, BNF_DIM) + # BNF statistics from the training targets: the network predicts normalised BNF + self.register_buffer("bnf_mean", torch.zeros(BNF_DIM)) + self.register_buffer("bnf_std", torch.ones(BNF_DIM)) + + @staticmethod + def prosody(F0, N): + B, F2 = F0.shape + f0 = F0.view(B, F2 // 2, 2).transpose(1, 2) + n = N.view(B, F2 // 2, 2).transpose(1, 2) + voiced = (f0 > 10).float() + logf0 = torch.where(f0 > 10, (torch.log(f0.clamp_min(1.0)) - math.log(200.0)) / 0.3, torch.zeros_like(f0)) + return torch.cat([logf0, voiced, n / 10.0], 1) + + def forward(self, asr, F0, N, out_len, in_len=None): + """out_len: [B] long, target BNF lengths. in_len: [B] long, valid 40 Hz frames (None = all).""" + B, _, Fr = asr.shape + in_len = torch.full((B,), Fr, device=asr.device) if in_len is None else in_len + m_in = (torch.arange(Fr, device=asr.device)[None] < in_len[:, None]).float()[:, None] + x = self.inp(torch.cat([asr, self.prosody(F0, N)], 1)) * m_in + for b in self.pre: + x = b(x, m_in) + T = int(out_len.max()) + # each item is resampled on its own valid span so padding never leaks into the content + x = torch.stack([F.pad(F.interpolate(x[i:i + 1, :, :int(in_len[i])], size=int(out_len[i]), + mode="linear", align_corners=False)[0], (0, T - int(out_len[i]))) + for i in range(B)]) + m_out = (torch.arange(T, device=asr.device)[None] < out_len[:, None]).float()[:, None] + for b in self.post: + x = b(x, m_out) + y = self.out(self.norm(x.transpose(1, 2))) * self.bnf_std + self.bnf_mean + return y * m_out.transpose(1, 2) diff --git a/silencio/asr.py b/silencio/asr.py new file mode 100644 index 0000000000000000000000000000000000000000..c94515665528b4c2a1756b2362dd8ec8a2912669 --- /dev/null +++ b/silencio/asr.py @@ -0,0 +1,41 @@ +"""MeanVC2's streaming Fast-U2++ ASR encoder: speech -> content tokens (BNF) for voice conversion, +and the training targets for the adapter. The TorchScript model is downloaded from ASLP-lab/MeanVC2.""" +import torch +import torchaudio +import torchaudio.compliance.kaldi as kaldi +from huggingface_hub import hf_hub_download + +REPO = "ASLP-lab/MeanVC2" +SR = 16000 +# the 160 ms Fast-U2++ that feeds MeanVC2's 120 ms + 40 ms model (see MeanVC2 infer_e2e.py) +ASR_PRESET = dict(asr="fastu2pp_160ms.pt", window=19, stride=16, offset_init=8, offset_step=4, req_cache=8) + + +def to_16k(wav, sr): + wav = torch.as_tensor(wav, dtype=torch.float32) + return wav if sr == SR else torchaudio.functional.resample(wav, sr, SR) + + +class ContentEncoder: + """Streaming Fast-U2++ ASR encoder (TorchScript): 16 kHz wave -> BNF [T,256] at 25 Hz. + The trace has CPU constants baked in, so it always runs on CPU (as in MeanVC2).""" + + def __init__(self): + self.device = "cpu" + self.model = torch.jit.load(hf_hub_download(REPO, ASR_PRESET["asr"]), map_location="cpu").eval() + + @torch.no_grad() + def __call__(self, wav16): + P = ASR_PRESET + x = torch.as_tensor(wav16, dtype=torch.float32).detach().cpu().reshape(1, -1) * (1 << 15) + fb = kaldi.fbank(x, frame_length=25, frame_shift=10, snip_edges=True, num_mel_bins=80, + energy_floor=0.0, dither=0.0, sample_frequency=SR)[None].to(self.device) + att = torch.zeros(6, 4, P["req_cache"], 128, device=self.device) + cnn = torch.zeros(6, 1, 256, 8, device=self.device) + offset, outs, i = P["offset_init"], [], 0 + req = torch.tensor(P["req_cache"], dtype=torch.int64) + while i + P["window"] <= fb.shape[1]: + out, att, cnn = self.model(fb[:, i:i + P["window"]], torch.tensor(offset, dtype=torch.int64), req, att, cnn) + outs.append(out[0]) + offset += P["offset_step"]; i += P["stride"] + return torch.cat(outs) if outs else torch.zeros(0, 256, device=self.device) diff --git a/silencio/ecapa_tdnn.py b/silencio/ecapa_tdnn.py new file mode 100644 index 0000000000000000000000000000000000000000..32fa16ab59226394e68cbc70fc7ea8c513539cde --- /dev/null +++ b/silencio/ecapa_tdnn.py @@ -0,0 +1,308 @@ +# Vendored from https://github.com/ASLP-lab/MeanVC2 (preprocess/models/ecapa_tdnn.py), Apache-2.0; +# the model code follows https://github.com/lawlict/ECAPA-TDNN (MIT). UpstreamExpert is imported at the top +# so the WavLM checkpoint path can be passed in as config_path. +# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN + +import os +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchaudio.transforms as trans +from s3prl.upstream.wavlm.expert import UpstreamExpert +# from .utils import UpstreamExpert + + +''' Res2Conv1d + BatchNorm1d + ReLU +''' + + +class Res2Conv1dReluBn(nn.Module): + ''' + in_channels == out_channels == channels + ''' + + def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True, scale=4): + super().__init__() + assert channels % scale == 0, "{} % {} != 0".format(channels, scale) + self.scale = scale + self.width = channels // scale + self.nums = scale if scale == 1 else scale - 1 + + self.convs = [] + self.bns = [] + for i in range(self.nums): + self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias)) + self.bns.append(nn.BatchNorm1d(self.width)) + self.convs = nn.ModuleList(self.convs) + self.bns = nn.ModuleList(self.bns) + + def forward(self, x): + out = [] + spx = torch.split(x, self.width, 1) + for i in range(self.nums): + if i == 0: + sp = spx[i] + else: + sp = sp + spx[i] + # Order: conv -> relu -> bn + sp = self.convs[i](sp) + sp = self.bns[i](F.relu(sp)) + out.append(sp) + if self.scale != 1: + out.append(spx[self.nums]) + out = torch.cat(out, dim=1) + + return out + + +''' Conv1d + BatchNorm1d + ReLU +''' + + +class Conv1dReluBn(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True): + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias) + self.bn = nn.BatchNorm1d(out_channels) + + def forward(self, x): + return self.bn(F.relu(self.conv(x))) + + +''' The SE connection of 1D case. +''' + + +class SE_Connect(nn.Module): + def __init__(self, channels, se_bottleneck_dim=128): + super().__init__() + self.linear1 = nn.Linear(channels, se_bottleneck_dim) + self.linear2 = nn.Linear(se_bottleneck_dim, channels) + + def forward(self, x): + out = x.mean(dim=2) + out = F.relu(self.linear1(out)) + out = torch.sigmoid(self.linear2(out)) + out = x * out.unsqueeze(2) + + return out + + +''' SE-Res2Block of the ECAPA-TDNN architecture. +''' + + +# def SE_Res2Block(channels, kernel_size, stride, padding, dilation, scale): +# return nn.Sequential( +# Conv1dReluBn(channels, 512, kernel_size=1, stride=1, padding=0), +# Res2Conv1dReluBn(512, kernel_size, stride, padding, dilation, scale=scale), +# Conv1dReluBn(512, channels, kernel_size=1, stride=1, padding=0), +# SE_Connect(channels) +# ) + + +class SE_Res2Block(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim): + super().__init__() + self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale) + self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0) + self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim) + + self.shortcut = None + if in_channels != out_channels: + self.shortcut = nn.Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + ) + + def forward(self, x): + residual = x + if self.shortcut: + residual = self.shortcut(x) + + x = self.Conv1dReluBn1(x) + x = self.Res2Conv1dReluBn(x) + x = self.Conv1dReluBn2(x) + x = self.SE_Connect(x) + + return x + residual + + +''' Attentive weighted mean and standard deviation pooling. +''' + + +class AttentiveStatsPool(nn.Module): + def __init__(self, in_dim, attention_channels=128, global_context_att=False): + super().__init__() + self.global_context_att = global_context_att + + # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs. + if global_context_att: + self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1) # equals W and b in the paper + else: + self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1) # equals W and b in the paper + self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1) # equals V and k in the paper + + def forward(self, x): + + if self.global_context_att: + context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x) + context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x) + x_in = torch.cat((x, context_mean, context_std), dim=1) + else: + x_in = x + + # DON'T use ReLU here! In experiments, I find ReLU hard to converge. + alpha = torch.tanh(self.linear1(x_in)) + # alpha = F.relu(self.linear1(x_in)) + alpha = torch.softmax(self.linear2(alpha), dim=2) + mean = torch.sum(alpha * x, dim=2) + residuals = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2 + std = torch.sqrt(residuals.clamp(min=1e-9)) + return torch.cat([mean, std], dim=1) + + +class ECAPA_TDNN(nn.Module): + def __init__(self, feat_dim=80, channels=512, emb_dim=192, global_context_att=False, + feat_type='fbank', sr=16000, feature_selection="hidden_states", update_extract=False, config_path=None): + super().__init__() + + self.feat_type = feat_type + self.feature_selection = feature_selection + self.update_extract = update_extract + self.sr = sr + + if feat_type == "fbank" or feat_type == "mfcc": + self.update_extract = False + + win_len = int(sr * 0.025) + hop_len = int(sr * 0.01) + + if feat_type == 'fbank': + self.feature_extract = trans.MelSpectrogram(sample_rate=sr, n_fft=512, win_length=win_len, + hop_length=hop_len, f_min=0.0, f_max=sr // 2, + pad=0, n_mels=feat_dim) + elif feat_type == 'mfcc': + melkwargs = { + 'n_fft': 512, + 'win_length': win_len, + 'hop_length': hop_len, + 'f_min': 0.0, + 'f_max': sr // 2, + 'pad': 0 + } + self.feature_extract = trans.MFCC(sample_rate=sr, n_mfcc=feat_dim, log_mels=False, + melkwargs=melkwargs) + else: + if config_path is None: + # TODO: meanvc2 root directory (two levels above preprocess/models/) + _ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + self.feature_extract = UpstreamExpert(ckpt=os.path.join(_ROOT, 'preprocess/ckpts/wavlm_large.pt')) + else: + self.feature_extract = UpstreamExpert(config_path) + if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(self.feature_extract.model.encoder.layers[23].self_attn, "fp32_attention"): + self.feature_extract.model.encoder.layers[23].self_attn.fp32_attention = False + if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(self.feature_extract.model.encoder.layers[11].self_attn, "fp32_attention"): + self.feature_extract.model.encoder.layers[11].self_attn.fp32_attention = False + + self.feat_num = self.get_feat_num() + self.feature_weight = nn.Parameter(torch.zeros(self.feat_num)) + + if feat_type != 'fbank' and feat_type != 'mfcc': + freeze_list = ['final_proj', 'label_embs_concat', 'mask_emb', 'project_q', 'quantizer'] + for name, param in self.feature_extract.named_parameters(): + for freeze_val in freeze_list: + if freeze_val in name: + param.requires_grad = False + break + + if not self.update_extract: + for param in self.feature_extract.parameters(): + param.requires_grad = False + + self.instance_norm = nn.InstanceNorm1d(feat_dim) + # self.channels = [channels] * 4 + [channels * 3] + self.channels = [channels] * 4 + [1536] + + self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2) + self.layer2 = SE_Res2Block(self.channels[0], self.channels[1], kernel_size=3, stride=1, padding=2, dilation=2, scale=8, se_bottleneck_dim=128) + self.layer3 = SE_Res2Block(self.channels[1], self.channels[2], kernel_size=3, stride=1, padding=3, dilation=3, scale=8, se_bottleneck_dim=128) + self.layer4 = SE_Res2Block(self.channels[2], self.channels[3], kernel_size=3, stride=1, padding=4, dilation=4, scale=8, se_bottleneck_dim=128) + + # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1) + cat_channels = channels * 3 + self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1) + self.pooling = AttentiveStatsPool(self.channels[-1], attention_channels=128, global_context_att=global_context_att) + self.bn = nn.BatchNorm1d(self.channels[-1] * 2) + self.linear = nn.Linear(self.channels[-1] * 2, emb_dim) + + + def get_feat_num(self): + self.feature_extract.eval() + wav = [torch.randn(self.sr).to(next(self.feature_extract.parameters()).device)] + with torch.no_grad(): + features = self.feature_extract(wav) + select_feature = features[self.feature_selection] + if isinstance(select_feature, (list, tuple)): + return len(select_feature) + else: + return 1 + + def get_feat(self, x): + if self.update_extract: + x = self.feature_extract([sample for sample in x]) + else: + with torch.no_grad(): + if self.feat_type == 'fbank' or self.feat_type == 'mfcc': + x = self.feature_extract(x) + 1e-6 # B x feat_dim x time_len + else: + x = self.feature_extract([sample for sample in x]) + + if self.feat_type == 'fbank': + x = x.log() + + if self.feat_type != "fbank" and self.feat_type != "mfcc": + x = x[self.feature_selection] + if isinstance(x, (list, tuple)): + x = torch.stack(x, dim=0) + else: + x = x.unsqueeze(0) + norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + x = (norm_weights * x).sum(dim=0) + x = torch.transpose(x, 1, 2) + 1e-6 + + x = self.instance_norm(x) + return x + + def forward(self, x): + x = self.get_feat(x) + + out1 = self.layer1(x) + out2 = self.layer2(out1) + out3 = self.layer3(out2) + out4 = self.layer4(out3) + + out = torch.cat([out2, out3, out4], dim=1) + out = F.relu(self.conv(out)) + out = self.bn(self.pooling(out)) + out = self.linear(out) + + return out + + +def ECAPA_TDNN_SMALL(feat_dim, emb_dim=256, feat_type='fbank', sr=16000, feature_selection="hidden_states", update_extract=False, config_path=None): + return ECAPA_TDNN(feat_dim=feat_dim, channels=512, emb_dim=emb_dim, + feat_type=feat_type, sr=sr, feature_selection=feature_selection, update_extract=update_extract, config_path=config_path) + +if __name__ == '__main__': + x = torch.zeros(2, 32000) + model = ECAPA_TDNN_SMALL(feat_dim=768, emb_dim=256, feat_type='hubert_base', feature_selection="hidden_states", + update_extract=False) + + out = model(x) + # print(model) + print(out.shape) + diff --git a/silencio/enroll.py b/silencio/enroll.py new file mode 100644 index 0000000000000000000000000000000000000000..bdb6c55d25d9e3d30eb6db5550f375747509775a --- /dev/null +++ b/silencio/enroll.py @@ -0,0 +1,44 @@ +"""Voice enrollment: a reference clip -> the 256-d voice embedding Silencio speaks with. + +The encoder is WavLM-Large + ECAPA-TDNN, the speaker model MeanVC2 was trained against. It is large +(~1.3 GB, CC BY-SA 3.0 weights) and only needed to add new voices, so it is not part of model.safetensors; +its weights are downloaded on first use. Needs `pip install s3prl`. +""" +import numpy as np +import torch +import torchaudio +from huggingface_hub import hf_hub_download + +# s3prl-format WavLM-Large, and the UniSpeech speaker-verification fine-tune (sha256 51f07e3b...) +WAVLM = ("s3prl/converted_ckpts", "wavlm_large.pt") +ECAPA = ("bezzam/wavlm_large_finetune_seed_tts_eval", "wavlm_large_finetune.pth") + + +class VoiceEncoder: + def __init__(self, device=None): + from .ecapa_tdnn import ECAPA_TDNN_SMALL + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + self.model = ECAPA_TDNN_SMALL(feat_dim=1024, feat_type="wavlm_large", config_path=hf_hub_download(*WAVLM)) + self.model.load_state_dict(torch.load(hf_hub_download(*ECAPA), map_location="cpu")["model"], strict=False) + self.model.to(self.device).eval() + + @torch.no_grad() + def __call__(self, wav, sr): + """wav: mono float array (5-15 s of clean speech works best) -> voice [256].""" + wav = torch.as_tensor(np.asarray(wav, dtype=np.float32)) + if wav.ndim > 1: + wav = wav.mean(-1) + if sr != 16000: + wav = torchaudio.functional.resample(wav, sr, 16000) + return self.model(wav[None].to(self.device))[0].float().cpu() + + +_encoder = None + + +def enroll(wav, sr, device=None): + """Voice embedding for `wav` at sample rate `sr`; the encoder is loaded once and reused.""" + global _encoder + if _encoder is None: + _encoder = VoiceEncoder(device) + return _encoder(wav, sr) diff --git a/silencio/meanvc2/__init__.py b/silencio/meanvc2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..12441627d409939d362fd2ad78d0478b09661ece --- /dev/null +++ b/silencio/meanvc2/__init__.py @@ -0,0 +1,2 @@ +"""MeanVC2 DiT renderer (ASLP-lab, Apache-2.0), vendored for inference.""" +from .dit import DiT diff --git a/silencio/meanvc2/dit.py b/silencio/meanvc2/dit.py new file mode 100644 index 0000000000000000000000000000000000000000..02c3e3dc655f653b1906b5d4d1df448e2c5dacbe --- /dev/null +++ b/silencio/meanvc2/dit.py @@ -0,0 +1,334 @@ +# Vendored from https://github.com/ASLP-lab/MeanVC2 (src/infer/dit_kvcache.py), Apache-2.0. +# Kept verbatim apart from import paths so Silencio runs without the original repo. +""" +ein notation: +b - batch +n - sequence +nt - text sequence +nw - raw wave length +d - dimension +""" + +from __future__ import annotations + +import torch +from torch import nn +import torch.nn.functional as F +from einops import rearrange + +from x_transformers.x_transformers import RotaryEmbedding + +# from src.model.prompt_vp import MRTE +from .modules import ( + TimestepEmbedding, + ConvNeXtV2Block, + ConvPositionEmbedding, + AdaLayerNorm_Final, + precompute_freqs_cis, + get_pos_embed_indices, +) + +from .kv_modules import ( + DiTBlock, + ChunkDiTBlock, +) + +class GlobalTimbreMemory(nn.Module): + """Global Timbre Memory (GTM) + Decomposes the global speaker embedding into K reusable timbre prototype slots. + A speaker-specific 2-layer MLP generates speaker-specific KV, fused with a universal prior. + """ + def __init__(self, spk_dim=256, memory_slots=8, hidden_dim=256): + super().__init__() + self.memory_slots = memory_slots + self.hidden_dim = hidden_dim + + # Speaker-specific mapping -- 2-layer MLP for increased expressiveness + self.mlp_k = nn.Sequential( + nn.Linear(spk_dim, spk_dim), + nn.SiLU(), + nn.Linear(spk_dim, memory_slots * hidden_dim), + ) + self.mlp_v = nn.Sequential( + nn.Linear(spk_dim, spk_dim), + nn.SiLU(), + nn.Linear(spk_dim, memory_slots * hidden_dim), + ) + + # Universal speaker-agnostic prototypes (learns common pronunciation patterns), small-variance init + self.k_prior = nn.Parameter(torch.zeros(memory_slots, hidden_dim)) + self.v_prior = nn.Parameter(torch.zeros(memory_slots, hidden_dim)) + nn.init.normal_(self.k_prior, std=0.02) + nn.init.normal_(self.v_prior, std=0.02) + + # LayerNorm after fusion for stable training + self.norm_k = nn.LayerNorm(hidden_dim) + self.norm_v = nn.LayerNorm(hidden_dim) + + def forward(self, spks): + # spks: [B, spk_dim] static global speaker embedding + B = spks.shape[0] + # Generate speaker-specific key-value pairs + k_spk = self.mlp_k(spks).reshape(B, self.memory_slots, self.hidden_dim) + v_spk = self.mlp_v(spks).reshape(B, self.memory_slots, self.hidden_dim) + # Fuse with universal prototypes + LayerNorm + k = self.norm_k(k_spk + torch.tanh(self.k_prior).unsqueeze(0)) # [B, K, D] + v = self.norm_v(v_spk + torch.tanh(self.v_prior).unsqueeze(0)) # [B, K, D] + return k, v + + +class TemporalTimbreEncoder(nn.Module): + """Temporal Timbre Encoder (TVT processing block) + Frame-level content vectors serve as Query in Multi-Head Cross-Attention with GTM. + The resulting time-varying timbre features act as frame-level conditioning, combined + with the global speaker embedding at the input. + """ + def __init__(self, content_dim=256, hidden_dim=256, attn_dim=128, num_heads=4): + super().__init__() + self.num_heads = num_heads + self.head_dim = attn_dim // num_heads + assert attn_dim % num_heads == 0 + + # Multi-Head Cross-Attention projections + self.q_proj = nn.Linear(content_dim, attn_dim) + self.k_proj = nn.Linear(hidden_dim, attn_dim) # Input from GTM hidden_dim + self.v_proj = nn.Linear(hidden_dim, attn_dim) # Independent V projection + self.out_proj = nn.Linear(attn_dim, content_dim) # Output projection back to content_dim + + self.attn_scale = self.head_dim ** -0.5 + self.attn_norm = nn.LayerNorm(content_dim) + + + def forward(self, bn, k_mem, v_mem): + """Args: + bn: [B, T, content_dim] frame-level content features + k_mem: [B, K, hidden_dim] GTM keys + v_mem: [B, K, hidden_dim] GTM values + Returns: + timbre_cond: [B, T, content_dim] time-varying timbre conditioning + """ + B, T, _ = bn.shape + K = k_mem.shape[1] + + # ---- 1. Multi-Head Cross-Attention: content queries × GTM ---- + q = self.q_proj(bn) # [B, T, attn_dim] + k = self.k_proj(k_mem) # [B, K, attn_dim] + v = self.v_proj(v_mem) # [B, K, attn_dim] + + # reshape → [B, num_heads, seq_len, head_dim] + q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(B, K, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(B, K, self.num_heads, self.head_dim).transpose(1, 2) + + attn = torch.matmul(q, k.transpose(-2, -1)) * self.attn_scale # [B, H, T, K] + attn = torch.softmax(attn, dim=-1) + v_t = torch.matmul(attn, v) # [B, H, T, head_dim] + + # merge heads -> output projection + v_t = v_t.transpose(1, 2).contiguous().view(B, T, -1) # [B, T, attn_dim] + timbre_cond = self.attn_norm(self.out_proj(v_t)) # [B, T, content_dim] + + return timbre_cond + +class InputEmbedding(nn.Module): + def __init__(self, mel_dim, cond_dim, out_dim): + super().__init__() + self.proj = nn.Linear(mel_dim + cond_dim * 2, out_dim) + # self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim) + + def forward(self, x: float["b n d"], cond: float["b n d"], spks: float["b n d"], drop_audio_cond=False): # noqa: F722 + # def forward(self, x: float["b n d"], cond: float["b n d"], timbre_cond: float["b n d"], drop_audio_cond=False): # noqa: F722 + if drop_audio_cond: # cfg for cond audio + cond = torch.zeros_like(cond) + spks = torch.zeros_like(spks) + + x = self.proj(torch.cat((x, cond, spks), dim=-1)) + # x = self.conv_pos_embed(x) + x + return x + + + +# Transformer backbone using DiT blocks + + +class DiT(nn.Module): + def __init__( + self, + *, + dim, + depth=8, + heads=8, + dim_head=64, + dropout=0.1, + ff_mult=4, + mel_dim=80, + bn_dim=256, + qk_norm=None, + conv_layers=0, + chunk_size=8, + block_size=4, + pe_attn_head=None, + long_skip_connection=False, + checkpoint_activations=False, + forward_layers=[0], # Layer 0 allowed to look ahead + backward_layers=[0,1,2,3], # Layer 3 allowed to look behind + t_f_num=[1,0,0,0], + t_p_num=[2,2,1,1], + ): + super().__init__() + + self.t_time_embed = TimestepEmbedding(dim) + self.r_time_embed = TimestepEmbedding(dim) + self.input_embed = InputEmbedding(mel_dim, bn_dim, dim) + self.rotary_embed = RotaryEmbedding(dim_head) + + self.dim = dim + self.depth = depth + + # GTM + TVT time-varying timbre module (replaces MRTE) + self.gtm = GlobalTimbreMemory(spk_dim=bn_dim, memory_slots=32, hidden_dim=bn_dim) + self.temporal_timbre = TemporalTimbreEncoder( + content_dim=bn_dim, hidden_dim=bn_dim, + attn_dim=128, num_heads=4, + ) + + forward_layers = set(forward_layers) if forward_layers else set() + backward_layers = set(backward_layers) if backward_layers else set() + + self.transformer_blocks = nn.ModuleList( + [ + ChunkDiTBlock( + dim=dim, + heads=heads, + dim_head=dim_head, + ff_mult=ff_mult, + dropout=dropout, + qk_norm=qk_norm, + chunk_size=chunk_size, + block_size=block_size, + pe_attn_head=pe_attn_head, + t_p=t_p_num[i] if i in backward_layers else 0, # backward + t_f=t_f_num[i] if i in forward_layers else 0, # forward + ) + for i in range(depth) + ] + ) + self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None + + self.norm_out = AdaLayerNorm_Final(dim) # final modulation + self.proj_out = nn.Linear(dim, mel_dim) + + self.checkpoint_activations = checkpoint_activations + + self.initialize_weights() + + def initialize_weights(self): + # Zero-out AdaLN layers in DiT blocks: + for block in self.transformer_blocks: + nn.init.constant_(block.attn_norm.linear.weight, 0) + nn.init.constant_(block.attn_norm.linear.bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.norm_out.linear.weight, 0) + nn.init.constant_(self.norm_out.linear.bias, 0) + nn.init.constant_(self.proj_out.weight, 0) + nn.init.constant_(self.proj_out.bias, 0) + + def ckpt_wrapper(self, module): + # https://github.com/chuanyangjin/fast-DiT/blob/main/models.py + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + + return ckpt_forward + + + def forward( + self, + x: float["b n d"], # nosied input audio # noqa: F722 B, T, 80 + t: float["b"] | float[""], # time step # noqa: F821 F722 + r: float["b"] | float[""], # time step # noqa: F821 F722 + cache: float["b n d"], + cond: float["b n d"], # bn # noqa: F722 B, T, 256 + spks: float["b d"], # spks # noqa: F722 B, 256 + offset=0, + mask: bool["b n"] | None = None, # noqa: F722 + is_inference: bool = False, + is_uncondition: bool = False, + cfg_mask: bool["b"] | None = None, # noqa: F722 + kv_cache=None, + ): + + batch, seq_len = x.shape[0], x.shape[1] + + # ---- timestep embedding ---- + t = self.t_time_embed(t) + r = self.r_time_embed(r) + t = t + r + + # ---- GTM: global speaker embedding -> timbre memory key-value pairs ---- + k_mem, v_mem = self.gtm(spks) # spks: [B, 256] -> k,v: [B, K, 256] + + # ---- TVT: frame-level BN x GTM -> timbre-enhanced timbre_cond ---- + timbre_cond = self.temporal_timbre(cond, k_mem, v_mem) # [B, T, bn_dim] + + # Expand spks_global to frame level, as pure global identity condition (independent of timbre_cond) + spks_expanded = spks.unsqueeze(1).expand(-1, cond.shape[1], -1) # [B, T, spk_dim] + + # ---- CFG masking ---- + if cfg_mask is not None: + cfg_mask_ = rearrange(cfg_mask, "b -> b 1 1") + timbre_cond = torch.where(cfg_mask_, torch.zeros_like(timbre_cond), timbre_cond) + spks_expanded = torch.where(cfg_mask_, torch.zeros_like(spks_expanded), spks_expanded) + + # Dual-path input: timbre_cond (content+timbre) + spks_expanded (pure global identity) + x = self.input_embed(x, timbre_cond, spks_expanded, drop_audio_cond=is_uncondition) + + # train + if not is_inference: + + rope = self.rotary_embed.forward_from_seq_len(seq_len) + # infer + else: + if cache != None: + cache = self.cache_embed(cache) + x = torch.concat((cache, x), dim=1) # [b, cache_len + seq_len, dim] + + # inference does not need to consider mask + cache_len = cache.shape[1] + rope_cache = self.rotary_embed.forward_from_seq_len(cache_len) + rope_x = self.rotary_embed.forward_from_seq_len(offset + seq_len) + rope = (torch.concat((rope_cache[0], rope_x[0][:, -seq_len:, :]), dim=1), rope_cache[1]) + else: + rope = self.rotary_embed.forward_from_seq_len(offset + seq_len) + + + if self.long_skip_connection is not None: + residual = x + + new_kv_cache = [] + # inner_hidden_states = [] + for index_block, block in enumerate(self.transformer_blocks): + if kv_cache is not None: + block_kv_cache = kv_cache[index_block] + else: + block_kv_cache = None + if self.checkpoint_activations: + # https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint + x, new_block_kv_cache = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, t, mask, rope, block_kv_cache, use_reentrant=False) + else: + x, new_block_kv_cache = block(x, t, mask=mask, rope=rope, is_inference=is_inference, kv_cache=block_kv_cache) + new_kv_cache.append(new_block_kv_cache) + if self.long_skip_connection is not None: + x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) + + + # x = x[:, -seq_len:, :] + x = self.norm_out(x, t) + + output = self.proj_out(x) + + return output, new_kv_cache + + + diff --git a/silencio/meanvc2/kv_modules.py b/silencio/meanvc2/kv_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..5c780331db74465bb3e1be758f02d63797bfa356 --- /dev/null +++ b/silencio/meanvc2/kv_modules.py @@ -0,0 +1,1060 @@ +# Vendored from https://github.com/ASLP-lab/MeanVC2 (src/infer/modules_kvcahe.py), Apache-2.0. +# Kept verbatim apart from import paths so Silencio runs without the original repo. +""" +ein notation: +b - batch +n - sequence +nt - text sequence +nw - raw wave length +d - dimension +""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import torch.nn.functional as F +import torchaudio +from librosa.filters import mel as librosa_mel_fn +from torch import nn +from x_transformers.x_transformers import apply_rotary_pos_emb + + +# raw wav to mel spec + + +mel_basis_cache = {} +hann_window_cache = {} + + +def get_bigvgan_mel_spectrogram( + waveform, + n_fft=1024, + n_mel_channels=100, + target_sample_rate=24000, + hop_length=256, + win_length=1024, + fmin=0, + fmax=None, + center=False, +): # Copy from https://github.com/NVIDIA/BigVGAN/tree/main + device = waveform.device + key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}" + + if key not in mel_basis_cache: + mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax) + mel_basis_cache[key] = torch.from_numpy(mel).float().to(device) # TODO: why they need .float()? + hann_window_cache[key] = torch.hann_window(win_length).to(device) + + mel_basis = mel_basis_cache[key] + hann_window = hann_window_cache[key] + + padding = (n_fft - hop_length) // 2 + waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1) + + spec = torch.stft( + waveform, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9) + + mel_spec = torch.matmul(mel_basis, spec) + mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5)) + + return mel_spec + + +def get_vocos_mel_spectrogram( + waveform, + n_fft=1024, + n_mel_channels=100, + target_sample_rate=24000, + hop_length=256, + win_length=1024, +): + mel_stft = torchaudio.transforms.MelSpectrogram( + sample_rate=target_sample_rate, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + n_mels=n_mel_channels, + power=1, + center=True, + normalized=False, + norm=None, + ).to(waveform.device) + if len(waveform.shape) == 3: + waveform = waveform.squeeze(1) # 'b 1 nw -> b nw' + + assert len(waveform.shape) == 2 + + mel = mel_stft(waveform) + mel = mel.clamp(min=1e-5).log() + return mel + + +class MelSpec(nn.Module): + def __init__( + self, + n_fft=1024, + hop_length=256, + win_length=1024, + n_mel_channels=100, + target_sample_rate=24_000, + mel_spec_type="vocos", + ): + super().__init__() + assert mel_spec_type in ["vocos", "bigvgan"], print("We only support two extract mel backend: vocos or bigvgan") + + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length + self.n_mel_channels = n_mel_channels + self.target_sample_rate = target_sample_rate + + if mel_spec_type == "vocos": + self.extractor = get_vocos_mel_spectrogram + elif mel_spec_type == "bigvgan": + self.extractor = get_bigvgan_mel_spectrogram + + self.register_buffer("dummy", torch.tensor(0), persistent=False) + + def forward(self, wav): + if self.dummy.device != wav.device: + self.to(wav.device) + + mel = self.extractor( + waveform=wav, + n_fft=self.n_fft, + n_mel_channels=self.n_mel_channels, + target_sample_rate=self.target_sample_rate, + hop_length=self.hop_length, + win_length=self.win_length, + ) + + return mel + + +# sinusoidal position embedding + + +class SinusPositionEmbedding(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x, scale=1000): + 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 + + +# convolutional position embedding + + +class ConvPositionEmbedding(nn.Module): + def __init__(self, dim, kernel_size=31, groups=16): + super().__init__() + assert kernel_size % 2 != 0 + self.conv1d = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), + nn.Mish(), + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), + nn.Mish(), + ) + + def forward(self, x: float["b n d"], mask: bool["b n"] | None = None): # noqa: F722 + if mask is not None: + mask = mask[..., None] + x = x.masked_fill(~mask, 0.0) + + x = x.permute(0, 2, 1) + x = self.conv1d(x) + out = x.permute(0, 2, 1) + + if mask is not None: + out = out.masked_fill(~mask, 0.0) + + return out + + +# rotary positional embedding related + + +def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0): + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ + # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py + theta *= theta_rescale_factor ** (dim / (dim - 2)) + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + t = torch.arange(end, device=freqs.device) # type: ignore + freqs = torch.outer(t, freqs).float() # type: ignore + freqs_cos = torch.cos(freqs) # real part + freqs_sin = torch.sin(freqs) # imaginary part + return torch.cat([freqs_cos, freqs_sin], dim=-1) + + +def get_pos_embed_indices(start, length, max_pos, scale=1.0): + # length = length if isinstance(length, int) else length.max() + scale = scale * torch.ones_like(start, dtype=torch.float32) # in case scale is a scalar + pos = ( + start.unsqueeze(1) + + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long() + ) + # avoid extra long error. + pos = torch.where(pos < max_pos, pos, max_pos - 1) + return pos + + +# Global Response Normalization layer (Instance Normalization ?) + + +class GRN(nn.Module): + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=1, keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py +# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108 + + +class ConvNeXtV2Block(nn.Module): + def __init__( + self, + dim: int, + intermediate_dim: int, + dilation: int = 1, + ): + super().__init__() + padding = (dilation * (7 - 1)) // 2 + self.dwconv = nn.Conv1d( + dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation + ) # depthwise conv + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.grn = GRN(intermediate_dim) + self.pwconv2 = nn.Linear(intermediate_dim, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = x.transpose(1, 2) # b n d -> b d n + x = self.dwconv(x) + x = x.transpose(1, 2) # b d n -> b n d + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + return residual + x + + +# RMSNorm + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + self.native_rms_norm = float(torch.__version__[:3]) >= 2.4 + + def forward(self, x): + if self.native_rms_norm: + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps) + else: + variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = x * self.weight + + return x + + +# AdaLayerNorm +# return with modulated x for attn input, and params for later mlp modulation + + +class AdaLayerNorm(nn.Module): + def __init__(self, dim): + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 6) + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, emb=None): + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1) + + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +# AdaLayerNorm for final layer +# return only with modulated x for attn input, cuz no more mlp modulation + + +class AdaLayerNorm_Final(nn.Module): + def __init__(self, dim): + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 2) + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, emb): + emb = self.linear(self.silu(emb)) + scale, shift = torch.chunk(emb, 2, dim=1) + + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +# FeedForward + + +class FeedForward(nn.Module): + def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + activation = nn.GELU(approximate=approximate) + project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) + self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) + + def forward(self, x): + return self.ff(x) + + +# Attention with possible joint part +# modified from diffusers/src/diffusers/models/attention_processor.py + + +class Attention(nn.Module): + def __init__( + self, + processor: JointAttnProcessor | AttnProcessor | ChunkAttnProcessor | BlockAttnProcessor, + dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + context_dim: Optional[int] = None, # if not None -> joint attention + context_pre_only: bool = False, + qk_norm: Optional[str] = None, + ): + super().__init__() + + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + self.processor = processor + + self.dim = dim + self.heads = heads + self.inner_dim = dim_head * heads + self.dropout = dropout + + self.context_dim = context_dim + self.context_pre_only = context_pre_only + + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + + if qk_norm is None: + self.q_norm = None + self.k_norm = None + elif qk_norm == "rms_norm": + self.q_norm = RMSNorm(dim_head, eps=1e-6) + self.k_norm = RMSNorm(dim_head, eps=1e-6) + else: + raise ValueError(f"Unimplemented qk_norm: {qk_norm}") + + if self.context_dim is not None: + self.to_q_c = nn.Linear(context_dim, self.inner_dim) + self.to_k_c = nn.Linear(context_dim, self.inner_dim) + self.to_v_c = nn.Linear(context_dim, self.inner_dim) + if qk_norm is None: + self.c_q_norm = None + self.c_k_norm = None + elif qk_norm == "rms_norm": + self.c_q_norm = RMSNorm(dim_head, eps=1e-6) + self.c_k_norm = RMSNorm(dim_head, eps=1e-6) + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, dim)) + self.to_out.append(nn.Dropout(dropout)) + + if self.context_dim is not None and not self.context_pre_only: + self.to_out_c = nn.Linear(self.inner_dim, context_dim) + + def forward( + self, + x: float["b n d"], # noised input x # noqa: F722 + c: float["b n d"] = None, # context c # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding for x + c_rope=None, # rotary position embedding for c + is_inference=False, + kv_cache=None, + ) -> torch.Tensor: + if c is not None: + return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope, is_inference=is_inference, kv_cache=kv_cache) + else: + return self.processor(self, x, mask=mask, rope=rope, is_inference=is_inference, kv_cache=kv_cache) + + +# Attention processor + + +class AttnProcessor: + def __init__( + self, + pe_attn_head: int | None = None, # number of attention head to apply rope, None for all + ): + + self.pe_attn_head = pe_attn_head + + def __call__( + self, + attn: Attention, + x: float["b n d"], # noised input x # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding + ) -> torch.FloatTensor: + batch_size = x.shape[0] + + # `sample` projections + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) + + # attention + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + + # apply rotary position embedding + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + + if self.pe_attn_head is not None: + pn = self.pe_attn_head + query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale) + key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + # mask. e.g. inference got a batch with different target durations, mask out the padding + if mask is not None: + attn_mask = mask + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n' + attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2]) + else: + attn_mask = None + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=True) + x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + x = x.to(query.dtype) + + # linear proj + x = attn.to_out[0](x) + # dropout + x = attn.to_out[1](x) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + + return x + +def scaled_dot_product_attention_only(query, key, value, attn_mask=None, dropout_p=0.0, + is_causal=False, scale=None, enable_gqa=False) -> torch.Tensor: + + L, S = query.size(-2), key.size(-2) + B = query.size(0) + scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale + attn_bias = torch.zeros(B, 1, L, S, dtype=query.dtype, device=query.device) + if is_causal: + assert attn_mask is None + temp_mask = torch.ones(B, 1, L, S, dtype=torch.bool).tril(diagonal=0) + attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) + attn_bias.to(query.dtype) + + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_bias.masked_fill_(attn_mask[:,:,-attn_bias.shape[2]:,:].logical_not(), float("-inf")) + else: + attn_bias = attn_mask + attn_bias + + if enable_gqa: + key = key.repeat_interleave(query.size(-3)//key.size(-3), -3) + value = value.repeat_interleave(query.size(-3)//value.size(-3), -3) + + attn_weight = query @ key.transpose(-2, -1) * scale_factor + attn_weight += attn_bias + attn_weight = torch.softmax(attn_weight, dim=-1) + attn_weight = torch.dropout(attn_weight, dropout_p, train=True) + return attn_weight @ value + +class BlockAttnProcessor: + def __init__(self, chunk_size: int, block_size: int, t_p: int, t_f: int, pe_attn_head: int | None = None,): + """ + Args: + chunk_size (int): Number of tokens per chunk + block_size (int): Number of tokens per block. + t_p (int): Number of past chunks to attend to + t_f (int): Number of leading blocks in the future chunk to attend to + """ + self.pe_attn_head = pe_attn_head + self.chunk_size = chunk_size + self.block_size = block_size + self.t_p = t_p + self.t_f = t_f + + def try_cached_mask(self, seq_len, device): + idx = torch.arange(seq_len, device=device) + ci = idx // self.chunk_size + + qi = ci[:, None] + kj = ci[None, :] + + # Within the same chunk + same_chunk = (qi == kj) + + # Previous t_p chunks + # prev_chunk = (kj == (qi - self.t_p)) + prev_chunk = (kj >= (qi - self.t_p)) & (kj < qi) + + # First t_f blocks of the next chunk + # First compute key(j)'s offset within its own chunk; this offset is bounded by block_size * t_f + offset_in_chunk = (idx % self.chunk_size)[None, :] + next_chunk_first_block = (kj == (qi + 1)) & (offset_in_chunk < self.block_size * self.t_f) + + computed_mask = same_chunk | prev_chunk | next_chunk_first_block + + return computed_mask + + def __call__( + self, + attn: Attention, + x: float["b n d"], # noised input x # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding + is_inference=False, + kv_cache=None, + ) -> torch.FloatTensor: + + # batch_size, seq_len, _ = x.shape + batch_size = x.shape[0] + device = x.device + + # 1. Compute query, key, value projections + query = attn.to_q(x) # Linear layer expands dims [b, n, d * heads] + key = attn.to_k(x) + value = attn.to_v(x) #torch.Size([batch, seq, 1024]) + + ## 3. Reshape query, key, value into multi-head format: [batch, heads, seq_len, head_dim] attention + inner_dim = key.shape[-1] # d * heads + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # [b, heads, n, d] + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + + # kvcache + if kv_cache is None: + key_cache = None + value_cache = None + else: + key_cache, value_cache = kv_cache + + if kv_cache is not None: + key = torch.cat([key_cache, key], dim=2) + if value_cache is not None: + value = torch.cat([value_cache, value], dim=2) + + new_kv_cache = (key[:, :, :-self.block_size, :], value[:, :, :-self.block_size, :]) # Subsequent blocks are future info; do not cache them + + batch_size, _, seq_len, _ = key.shape + + # apply rotary position embedding + # Apply rotary position encoding to q and k + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + + if self.pe_attn_head is not None: + pn = self.pe_attn_head + query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale) + key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + # [b, i, j] = True means token i can attend to token j, + # [b, i, j] = False means token i cannot attend to token j, + computed_mask = self.try_cached_mask(seq_len,device).unsqueeze(0).expand(batch_size, -1, -1) # [b, seq_len, seq_len] + + # 4.5 Expand the final mask to multi-head dimensions; shape becomes [batch, heads, seq_len, seq_len] #torch.Size([2, 16, 636, 636]) + attn_mask = computed_mask.unsqueeze(1).expand(batch_size, 1, seq_len, seq_len) + + # 5. Call PyTorch 2.0 scaled_dot_product_attention + attn_output = scaled_dot_product_attention_only( + query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False + )#attn_mask.to(query.dtype) + # attn_output shape: [batch, heads, seq_len, head_dim] + + # 6. Restore shape; concatenate multi-head back to original dims [batch, seq_len, inner_dim] + attn_output = attn_output.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + attn_output = attn_output.to(query.dtype) + + # 7. Pass through output projection and dropout + attn_output = attn.to_out[0](attn_output) + attn_output = attn.to_out[1](attn_output) + + # Expand mask to [batch, seq_len, 1] and zero out the output accordingly + if mask is not None: + mask = mask.unsqueeze(-1) + attn_output = attn_output.masked_fill(~mask, 0.0) + + return attn_output, new_kv_cache + +class ChunkAttnProcessor: + def __init__( + self, + chunk_size: int, + pe_attn_head=None, # number of attention head to apply rope, None for all + ): + self.chunk_size = chunk_size + self.pe_attn_head = pe_attn_head + + def __call__( + self, + attn: Attention, + x: float["b 2*N*chunk_size d"], # noised input x # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding + is_inference=False, + ) -> torch.FloatTensor: + batch_size, seq_len, _ = x.shape + + # `sample` projections + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) + + # attention + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + + # apply rotary position embedding + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + + if self.pe_attn_head is not None: + pn = self.pe_attn_head + query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale) + key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + indices = torch.arange(seq_len, device=x.device) + chunk_indices = indices // self.chunk_size + N = int(seq_len / 2 / self.chunk_size) + + # attn_mask_1 = chunk_indices.unsqueeze(0) <= chunk_indices.unsqueeze(1) + # attn_mask_2 = (chunk_indices.unsqueeze(0) + N < chunk_indices.unsqueeze(1)) | (chunk_indices.unsqueeze(0) == chunk_indices.unsqueeze(1)) + # attn_mask = attn_mask_1 & attn_mask_2 + + # Generate left/right side identifiers (left side = first N*chunk_size frames) + is_right_side = indices >= (N * self.chunk_size) + + # Left blocks (M_i) can attend to <= current block's left blocks + # left_mask = chunk_indices.unsqueeze(0) <= chunk_indices.unsqueeze(1) + + # Right blocks (M'_i) can only attend to left clean blocks (all M_j, j < i) and itself + # right_mask = ( + # (chunk_indices.unsqueeze(0) < (chunk_indices.unsqueeze(1) - N)) | # Access left clean blocks + # (chunk_indices.unsqueeze(0) == chunk_indices.unsqueeze(1)) # Access itself + # ) + + max_lookback = 5 + num_cache_blocks = N # N + + # 3. Expand dims for broadcasting + ci = chunk_indices.unsqueeze(0) # [L,1], row: chunk the query position belongs to ??? shouldn't it be [1, L]? + cj = chunk_indices.unsqueeze(1) # [1,L], col: chunk the key position belongs to ??? shouldn't it be [L, 1]? + + # 4. Compute relative new block index: for block j, rel_j = cj - N; only rel_j >= 0 is a new block + rel_j = cj - num_cache_blocks # [1,L] + + # 5. Self-attention: token can always attend to itself + mask_self = ci == cj # [L,L] + + mask_cache = ( + (rel_j >= 0) & + (ci < num_cache_blocks) & + (ci < rel_j) & + (ci >= rel_j - max_lookback) + ) + + right_mask = mask_self | mask_cache # [L,L] boolean matrix + + + lookback_k = 5 # Look back at most 5 previous blocks + self = 6 blocks total + block_diff = cj - ci + left_mask = (block_diff >= 0) & (block_diff <= lookback_k) + + # Combine masks + if not is_inference: + # attn_mask = torch.where( + # is_right_side.unsqueeze(1), # Apply right_mask to right-side blocks + # right_mask, + # left_mask, # Apply left_mask to left-side blocks + # ) + attn_mask = right_mask + else: + attn_mask = left_mask + + + if mask is not None: + pad_mask = mask.unsqueeze(1) & mask.unsqueeze(2) | torch.eye(seq_len, device=x.device).unsqueeze(0).bool() + attn_mask = attn_mask.unsqueeze(0).expand(batch_size, -1, -1) & pad_mask + else: + attn_mask = attn_mask.unsqueeze(0).expand(batch_size, -1, -1) + + attn_mask = attn_mask.unsqueeze(1).expand(batch_size, 1, seq_len, seq_len) + + # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + x = scaled_dot_product_attention_only(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=True) + x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + x = x.to(query.dtype) + + # linear proj + x = attn.to_out[0](x) + # dropout + x = attn.to_out[1](x) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + + return x + +# Joint Attention processor for MM-DiT +# modified from diffusers/src/diffusers/models/attention_processor.py + + +class JointAttnProcessor: + def __init__(self): + pass + + def __call__( + self, + attn: Attention, + x: float["b n d"], # noised input x # noqa: F722 + c: float["b nt d"] = None, # context c, here text # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding for x + c_rope=None, # rotary position embedding for c + ) -> torch.FloatTensor: + residual = x + + batch_size = c.shape[0] + + # `sample` projections + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) + + # `context` projections + c_query = attn.to_q_c(c) + c_key = attn.to_k_c(c) + c_value = attn.to_v_c(c) + + # attention + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + c_query = c_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + c_key = c_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + c_value = c_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + if attn.c_q_norm is not None: + c_query = attn.c_q_norm(c_query) + if attn.c_k_norm is not None: + c_key = attn.c_k_norm(c_key) + + # apply rope for context and noised input independently + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + if c_rope is not None: + freqs, xpos_scale = c_rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale) + c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale) + + # joint attention + query = torch.cat([query, c_query], dim=2) + key = torch.cat([key, c_key], dim=2) + value = torch.cat([value, c_value], dim=2) + + # mask. e.g. inference got a batch with different target durations, mask out the padding + if mask is not None: + attn_mask = F.pad(mask, (0, c.shape[1]), value=True) # no mask for c (text) + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n' + attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2]) + else: + attn_mask = None + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + x = x.to(query.dtype) + + # Split the attention outputs. + x, c = ( + x[:, : residual.shape[1]], + x[:, residual.shape[1] :], + ) + + # linear proj + x = attn.to_out[0](x) + # dropout + x = attn.to_out[1](x) + if not attn.context_pre_only: + c = attn.to_out_c(c) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + # c = c.masked_fill(~mask, 0.) # no mask for c (text) + + return x, c + + +# DiT Block + + +class DiTBlock(nn.Module): + def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, qk_norm=None, pe_attn_head=None): + super().__init__() + + self.attn_norm = AdaLayerNorm(dim) + self.attn = Attention( + processor=AttnProcessor(pe_attn_head=pe_attn_head), + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + qk_norm=qk_norm, + ) + + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") + + def forward(self, x, t, mask=None, rope=None): # x: noised input, t: time embedding + # pre-norm & modulation for attention input + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + + # attention + attn_output = self.attn(x=norm, mask=mask, rope=rope) + + # process attention output for input x + x = x + gate_msa.unsqueeze(1) * attn_output + + norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(norm) + x = x + gate_mlp.unsqueeze(1) * ff_output + + return x + + +class ChunkDiTBlock(nn.Module): + def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, qk_norm=None, chunk_size=16, block_size=8, t_p=0, t_f=0, pe_attn_head=None): + super().__init__() + + self.attn_norm = AdaLayerNorm(dim) + self.attn = Attention( + processor=BlockAttnProcessor(chunk_size=chunk_size, block_size=block_size, t_p=t_p, t_f=t_f, pe_attn_head=pe_attn_head), + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + qk_norm=qk_norm, + ) + + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") + + def forward(self, x, t, mask=None, rope=None, is_inference=False, kv_cache=None): # x: noised input, t: time embedding + # pre-norm & modulation for attention input + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + + # attention + attn_output, new_kv_cache = self.attn(x=norm, mask=mask, rope=rope, is_inference=is_inference, kv_cache=kv_cache) + + # process attention output for input x + x = x + gate_msa.unsqueeze(1) * attn_output + + norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(norm) + x = x + gate_mlp.unsqueeze(1) * ff_output + + return x, new_kv_cache + +# MMDiT Block https://arxiv.org/abs/2403.03206 + + +class MMDiTBlock(nn.Module): + r""" + modified from diffusers/src/diffusers/models/attention.py + + notes. + _c: context related. text, cond, etc. (left part in sd3 fig2.b) + _x: noised input related. (right part) + context_pre_only: last layer only do prenorm + modulation cuz no more ffn + """ + + def __init__( + self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_dim=None, context_pre_only=False, qk_norm=None + ): + super().__init__() + if context_dim is None: + context_dim = dim + self.context_pre_only = context_pre_only + + self.attn_norm_c = AdaLayerNorm_Final(context_dim) if context_pre_only else AdaLayerNorm(context_dim) + self.attn_norm_x = AdaLayerNorm(dim) + self.attn = Attention( + processor=JointAttnProcessor(), + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + context_dim=context_dim, + context_pre_only=context_pre_only, + qk_norm=qk_norm, + ) + + if not context_pre_only: + self.ff_norm_c = nn.LayerNorm(context_dim, elementwise_affine=False, eps=1e-6) + self.ff_c = FeedForward(dim=context_dim, mult=ff_mult, dropout=dropout, approximate="tanh") + else: + self.ff_norm_c = None + self.ff_c = None + self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") + + def forward(self, x, c, t, mask=None, rope=None, c_rope=None): # x: noised input, c: context, t: time embedding + # pre-norm & modulation for attention input + if self.context_pre_only: + norm_c = self.attn_norm_c(c, t) + else: + norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t) + norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t) + + # attention + x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope) + + # process attention output for context c + if self.context_pre_only: + c = None + else: # if not last layer + c = c + c_gate_msa.unsqueeze(1) * c_attn_output + + norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + c_ff_output = self.ff_c(norm_c) + c = c + c_gate_mlp.unsqueeze(1) * c_ff_output + + # process attention output for input x + x = x + x_gate_msa.unsqueeze(1) * x_attn_output + + norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None] + x_ff_output = self.ff_x(norm_x) + x = x + x_gate_mlp.unsqueeze(1) * x_ff_output + + return c, x + + +# time step conditioning embedding + + +class TimestepEmbedding(nn.Module): + def __init__(self, dim, freq_embed_dim=256): + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: float["b"]): # noqa: F821 + time_hidden = self.time_embed(timestep) + time_hidden = time_hidden.to(timestep.dtype) + time = self.time_mlp(time_hidden) # b d + return time diff --git a/silencio/meanvc2/modules.py b/silencio/meanvc2/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..708927ed7e04a631b8a9cf03e611b2bc1d8583a7 --- /dev/null +++ b/silencio/meanvc2/modules.py @@ -0,0 +1,1042 @@ +# Vendored from https://github.com/ASLP-lab/MeanVC2 (src/model/modules.py), Apache-2.0. +# Kept verbatim apart from import paths so Silencio runs without the original repo. +""" +ein notation: +b - batch +n - sequence +nt - text sequence +nw - raw wave length +d - dimension +""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import torch.nn.functional as F +import torchaudio +from librosa.filters import mel as librosa_mel_fn +from torch import nn +from x_transformers.x_transformers import apply_rotary_pos_emb + + +# raw wav to mel spec + + +mel_basis_cache = {} +hann_window_cache = {} + + +def get_bigvgan_mel_spectrogram( + waveform, + n_fft=1024, + n_mel_channels=100, + target_sample_rate=24000, + hop_length=256, + win_length=1024, + fmin=0, + fmax=None, + center=False, +): # Copy from https://github.com/NVIDIA/BigVGAN/tree/main + device = waveform.device + key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}" + + if key not in mel_basis_cache: + mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax) + mel_basis_cache[key] = torch.from_numpy(mel).float().to(device) # TODO: why they need .float()? + hann_window_cache[key] = torch.hann_window(win_length).to(device) + + mel_basis = mel_basis_cache[key] + hann_window = hann_window_cache[key] + + padding = (n_fft - hop_length) // 2 + waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1) + + spec = torch.stft( + waveform, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9) + + mel_spec = torch.matmul(mel_basis, spec) + mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5)) + + return mel_spec + + +def get_vocos_mel_spectrogram( + waveform, + n_fft=1024, + n_mel_channels=100, + target_sample_rate=24000, + hop_length=256, + win_length=1024, +): + mel_stft = torchaudio.transforms.MelSpectrogram( + sample_rate=target_sample_rate, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + n_mels=n_mel_channels, + power=1, + center=True, + normalized=False, + norm=None, + ).to(waveform.device) + if len(waveform.shape) == 3: + waveform = waveform.squeeze(1) # 'b 1 nw -> b nw' + + assert len(waveform.shape) == 2 + + mel = mel_stft(waveform) + mel = mel.clamp(min=1e-5).log() + return mel + + +class MelSpec(nn.Module): + def __init__( + self, + n_fft=1024, + hop_length=256, + win_length=1024, + n_mel_channels=100, + target_sample_rate=24_000, + mel_spec_type="vocos", + ): + super().__init__() + assert mel_spec_type in ["vocos", "bigvgan"], print("We only support two extract mel backend: vocos or bigvgan") + + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length + self.n_mel_channels = n_mel_channels + self.target_sample_rate = target_sample_rate + + if mel_spec_type == "vocos": + self.extractor = get_vocos_mel_spectrogram + elif mel_spec_type == "bigvgan": + self.extractor = get_bigvgan_mel_spectrogram + + self.register_buffer("dummy", torch.tensor(0), persistent=False) + + def forward(self, wav): + if self.dummy.device != wav.device: + self.to(wav.device) + + mel = self.extractor( + waveform=wav, + n_fft=self.n_fft, + n_mel_channels=self.n_mel_channels, + target_sample_rate=self.target_sample_rate, + hop_length=self.hop_length, + win_length=self.win_length, + ) + + return mel + + +# sinusoidal position embedding + + +class SinusPositionEmbedding(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x, scale=1000): + 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 + + +# convolutional position embedding + + +class ConvPositionEmbedding(nn.Module): + def __init__(self, dim, kernel_size=31, groups=16): + super().__init__() + assert kernel_size % 2 != 0 + self.conv1d = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), + nn.Mish(), + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), + nn.Mish(), + ) + + def forward(self, x: float["b n d"], mask: bool["b n"] | None = None): # noqa: F722 + if mask is not None: + mask = mask[..., None] + x = x.masked_fill(~mask, 0.0) + + x = x.permute(0, 2, 1) + x = self.conv1d(x) + out = x.permute(0, 2, 1) + + if mask is not None: + out = out.masked_fill(~mask, 0.0) + + return out + + +# rotary positional embedding related + + +def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0): + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ + # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py + theta *= theta_rescale_factor ** (dim / (dim - 2)) + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + t = torch.arange(end, device=freqs.device) # type: ignore + freqs = torch.outer(t, freqs).float() # type: ignore + freqs_cos = torch.cos(freqs) # real part + freqs_sin = torch.sin(freqs) # imaginary part + return torch.cat([freqs_cos, freqs_sin], dim=-1) + + +def get_pos_embed_indices(start, length, max_pos, scale=1.0): + # length = length if isinstance(length, int) else length.max() + scale = scale * torch.ones_like(start, dtype=torch.float32) # in case scale is a scalar + pos = ( + start.unsqueeze(1) + + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long() + ) + # avoid extra long error. + pos = torch.where(pos < max_pos, pos, max_pos - 1) + return pos + + +# Global Response Normalization layer (Instance Normalization ?) + + +class GRN(nn.Module): + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=1, keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py +# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108 + + +class ConvNeXtV2Block(nn.Module): + def __init__( + self, + dim: int, + intermediate_dim: int, + dilation: int = 1, + ): + super().__init__() + padding = (dilation * (7 - 1)) // 2 + self.dwconv = nn.Conv1d( + dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation + ) # depthwise conv + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.grn = GRN(intermediate_dim) + self.pwconv2 = nn.Linear(intermediate_dim, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = x.transpose(1, 2) # b n d -> b d n + x = self.dwconv(x) + x = x.transpose(1, 2) # b d n -> b n d + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + return residual + x + + +# RMSNorm + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + self.native_rms_norm = float(torch.__version__[:3]) >= 2.4 + + def forward(self, x): + if self.native_rms_norm: + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps) + else: + variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = x * self.weight + + return x + + +# AdaLayerNorm +# return with modulated x for attn input, and params for later mlp modulation + + +class AdaLayerNorm(nn.Module): + def __init__(self, dim): + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 6) + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, emb=None): + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1) + + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +# AdaLayerNorm for final layer +# return only with modulated x for attn input, cuz no more mlp modulation + + +class AdaLayerNorm_Final(nn.Module): + def __init__(self, dim): + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 2) + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, emb): + emb = self.linear(self.silu(emb)) + scale, shift = torch.chunk(emb, 2, dim=1) + + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +# FeedForward + + +class FeedForward(nn.Module): + def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + activation = nn.GELU(approximate=approximate) + project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) + self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) + + def forward(self, x): + return self.ff(x) + + +# Attention with possible joint part +# modified from diffusers/src/diffusers/models/attention_processor.py + + +class Attention(nn.Module): + def __init__( + self, + processor: JointAttnProcessor | AttnProcessor | ChunkAttnProcessor | BlockAttnProcessor, + dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + context_dim: Optional[int] = None, # if not None -> joint attention + context_pre_only: bool = False, + qk_norm: Optional[str] = None, + ): + super().__init__() + + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + self.processor = processor + + self.dim = dim + self.heads = heads + self.inner_dim = dim_head * heads + self.dropout = dropout + + self.context_dim = context_dim + self.context_pre_only = context_pre_only + + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + + if qk_norm is None: + self.q_norm = None + self.k_norm = None + elif qk_norm == "rms_norm": + self.q_norm = RMSNorm(dim_head, eps=1e-6) + self.k_norm = RMSNorm(dim_head, eps=1e-6) + else: + raise ValueError(f"Unimplemented qk_norm: {qk_norm}") + + if self.context_dim is not None: + self.to_q_c = nn.Linear(context_dim, self.inner_dim) + self.to_k_c = nn.Linear(context_dim, self.inner_dim) + self.to_v_c = nn.Linear(context_dim, self.inner_dim) + if qk_norm is None: + self.c_q_norm = None + self.c_k_norm = None + elif qk_norm == "rms_norm": + self.c_q_norm = RMSNorm(dim_head, eps=1e-6) + self.c_k_norm = RMSNorm(dim_head, eps=1e-6) + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, dim)) + self.to_out.append(nn.Dropout(dropout)) + + if self.context_dim is not None and not self.context_pre_only: + self.to_out_c = nn.Linear(self.inner_dim, context_dim) + + def forward( + self, + x: float["b n d"], # noised input x # noqa: F722 + c: float["b n d"] = None, # context c # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding for x + c_rope=None, # rotary position embedding for c + is_inference=False, + ) -> torch.Tensor: + if c is not None: + return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope, is_inference=is_inference) + else: + return self.processor(self, x, mask=mask, rope=rope, is_inference=is_inference) + + +# Attention processor + + +class AttnProcessor: + def __init__( + self, + pe_attn_head: int | None = None, # number of attention head to apply rope, None for all + ): + + self.pe_attn_head = pe_attn_head + + def __call__( + self, + attn: Attention, + x: float["b n d"], # noised input x # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding + ) -> torch.FloatTensor: + batch_size = x.shape[0] + + # `sample` projections + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) + + # attention + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + + # apply rotary position embedding + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + + if self.pe_attn_head is not None: + pn = self.pe_attn_head + query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale) + key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + # mask. e.g. inference got a batch with different target durations, mask out the padding + if mask is not None: + attn_mask = mask + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n' + attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2]) + else: + attn_mask = None + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=True) + x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + x = x.to(query.dtype) + + # linear proj + x = attn.to_out[0](x) + # dropout + x = attn.to_out[1](x) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + + return x + +def scaled_dot_product_attention_only(query, key, value, attn_mask=None, dropout_p=0.0, + is_causal=False, scale=None, enable_gqa=False) -> torch.Tensor: + + L, S = query.size(-2), key.size(-2) + B = query.size(0) + scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale + attn_bias = torch.zeros(B, 1, L, S, dtype=query.dtype, device=query.device) + if is_causal: + assert attn_mask is None + temp_mask = torch.ones(B, 1, L, S, dtype=torch.bool).tril(diagonal=0) + attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) + attn_bias.to(query.dtype) + + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf")) + else: + attn_bias = attn_mask + attn_bias + + if enable_gqa: + key = key.repeat_interleave(query.size(-3)//key.size(-3), -3) + value = value.repeat_interleave(query.size(-3)//value.size(-3), -3) + + attn_weight = query @ key.transpose(-2, -1) * scale_factor + attn_weight += attn_bias + attn_weight = torch.softmax(attn_weight, dim=-1) + attn_weight = torch.dropout(attn_weight, dropout_p, train=True) + return attn_weight @ value + +class BlockAttnProcessor: + def __init__(self, chunk_size: int, block_size: int, t_p: int, t_f: int, pe_attn_head: int | None = None,): + """ + Args: + chunk_size (int): Number of tokens per chunk + block_size (int): Number of tokens per block. + t_p (int): Number of past chunks to attend to + t_f (int): Number of leading blocks in the future chunk to attend to + """ + self.pe_attn_head = pe_attn_head + self.chunk_size = chunk_size + self.block_size = block_size + self.t_p = t_p + self.t_f = t_f + + def try_cached_mask(self, seq_len, device): + idx = torch.arange(seq_len, device=device) + ci = idx // self.chunk_size + + qi = ci[:, None] + kj = ci[None, :] + + # Within the same chunk + same_chunk = (qi == kj) + + # Previous t_p chunks + # prev_chunk = (kj == (qi - self.t_p)) + prev_chunk = (kj >= (qi - self.t_p)) & (kj < qi) + + # First t_f blocks of the next chunk + # First compute key(j)'s offset within its own chunk; this offset is bounded by block_size * t_f + offset_in_chunk = (idx % self.chunk_size)[None, :] + next_chunk_first_block = (kj == (qi + 1)) & (offset_in_chunk < self.block_size * self.t_f) + + computed_mask = same_chunk | prev_chunk | next_chunk_first_block + + return computed_mask + + def __call__( + self, + attn: Attention, + x: float["b n d"], # noised input x # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding + is_inference=False, + ) -> torch.FloatTensor: + + batch_size, seq_len, _ = x.shape + device = x.device + + # 1. Compute query, key, value projections + query = attn.to_q(x) # Linear layer expands dims [b, n, d * heads] + key = attn.to_k(x) + value = attn.to_v(x) #torch.Size([batch, seq, 1024]) + + ## 3. Reshape query, key, value into multi-head format: [batch, heads, seq_len, head_dim] attention + inner_dim = key.shape[-1] # d * heads + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # [b, heads, n, d] + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + + + # apply rotary position embedding + # Apply rotary position encoding to q and k + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + + if self.pe_attn_head is not None: + pn = self.pe_attn_head + query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale) + key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + # [b, i, j] = True means token i can attend to token j, + # [b, i, j] = False means token i cannot attend to token j, + computed_mask = self.try_cached_mask(seq_len,device).unsqueeze(0).expand(batch_size, -1, -1) # [b, seq_len, seq_len] + + # 4.5 Expand the final mask to multi-head dimensions; shape becomes [batch, heads, seq_len, seq_len] #torch.Size([2, 16, 636, 636]) + attn_mask = computed_mask.unsqueeze(1).expand(batch_size, 1, seq_len, seq_len) + + # 5. Call PyTorch 2.0 scaled_dot_product_attention + attn_output = scaled_dot_product_attention_only( + query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False + )#attn_mask.to(query.dtype) + # attn_output shape: [batch, heads, seq_len, head_dim] + + # 6. Restore shape; concatenate multi-head back to original dims [batch, seq_len, inner_dim] + attn_output = attn_output.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + attn_output = attn_output.to(query.dtype) + + # 7. Pass through output projection and dropout + attn_output = attn.to_out[0](attn_output) + attn_output = attn.to_out[1](attn_output) + + # Expand mask to [batch, seq_len, 1] and zero out the output accordingly + if mask is not None: + mask = mask.unsqueeze(-1) + attn_output = attn_output.masked_fill(~mask, 0.0) + + return attn_output + +class ChunkAttnProcessor: + def __init__( + self, + chunk_size: int, + pe_attn_head=None, # number of attention head to apply rope, None for all + ): + self.chunk_size = chunk_size + self.pe_attn_head = pe_attn_head + + def __call__( + self, + attn: Attention, + x: float["b 2*N*chunk_size d"], # noised input x # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding + is_inference=False, + ) -> torch.FloatTensor: + batch_size, seq_len, _ = x.shape + + # `sample` projections + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) + + # attention + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + + # apply rotary position embedding + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + + if self.pe_attn_head is not None: + pn = self.pe_attn_head + query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale) + key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + indices = torch.arange(seq_len, device=x.device) + chunk_indices = indices // self.chunk_size + N = int(seq_len / 2 / self.chunk_size) + + # attn_mask_1 = chunk_indices.unsqueeze(0) <= chunk_indices.unsqueeze(1) + # attn_mask_2 = (chunk_indices.unsqueeze(0) + N < chunk_indices.unsqueeze(1)) | (chunk_indices.unsqueeze(0) == chunk_indices.unsqueeze(1)) + # attn_mask = attn_mask_1 & attn_mask_2 + + # Generate left/right side identifiers (left side = first N*chunk_size frames) + is_right_side = indices >= (N * self.chunk_size) + + # Left blocks (M_i) can attend to <= current block's left blocks + # left_mask = chunk_indices.unsqueeze(0) <= chunk_indices.unsqueeze(1) + + # Right blocks (M'_i) can only attend to left clean blocks (all M_j, j < i) and itself + # right_mask = ( + # (chunk_indices.unsqueeze(0) < (chunk_indices.unsqueeze(1) - N)) | # Access left clean blocks + # (chunk_indices.unsqueeze(0) == chunk_indices.unsqueeze(1)) # Access itself + # ) + + max_lookback = 5 + num_cache_blocks = N # N + + # 3. Expand dims for broadcasting + ci = chunk_indices.unsqueeze(0) # [L,1], row: chunk the query position belongs to ??? shouldn't it be [1, L]? + cj = chunk_indices.unsqueeze(1) # [1,L], col: chunk the key position belongs to ??? shouldn't it be [L, 1]? + + # 4. Compute relative new block index: for block j, rel_j = cj - N; only rel_j >= 0 is a new block + rel_j = cj - num_cache_blocks # [1,L] + + # 5. Self-attention: token can always attend to itself + mask_self = ci == cj # [L,L] + + mask_cache = ( + (rel_j >= 0) & + (ci < num_cache_blocks) & + (ci < rel_j) & + (ci >= rel_j - max_lookback) + ) + + right_mask = mask_self | mask_cache # [L,L] boolean matrix + + + lookback_k = 5 # Look back at most 5 previous blocks + self = 6 blocks total + block_diff = cj - ci + left_mask = (block_diff >= 0) & (block_diff <= lookback_k) + + # Combine masks + if not is_inference: + # attn_mask = torch.where( + # is_right_side.unsqueeze(1), # Apply right_mask to right-side blocks + # right_mask, + # left_mask, # Apply left_mask to left-side blocks + # ) + attn_mask = right_mask + else: + attn_mask = left_mask + + + if mask is not None: + pad_mask = mask.unsqueeze(1) & mask.unsqueeze(2) | torch.eye(seq_len, device=x.device).unsqueeze(0).bool() + attn_mask = attn_mask.unsqueeze(0).expand(batch_size, -1, -1) & pad_mask + else: + attn_mask = attn_mask.unsqueeze(0).expand(batch_size, -1, -1) + + attn_mask = attn_mask.unsqueeze(1).expand(batch_size, 1, seq_len, seq_len) + + # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + x = scaled_dot_product_attention_only(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=True) + x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + x = x.to(query.dtype) + + # linear proj + x = attn.to_out[0](x) + # dropout + x = attn.to_out[1](x) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + + return x + +# Joint Attention processor for MM-DiT +# modified from diffusers/src/diffusers/models/attention_processor.py + + +class JointAttnProcessor: + def __init__(self): + pass + + def __call__( + self, + attn: Attention, + x: float["b n d"], # noised input x # noqa: F722 + c: float["b nt d"] = None, # context c, here text # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + rope=None, # rotary position embedding for x + c_rope=None, # rotary position embedding for c + ) -> torch.FloatTensor: + residual = x + + batch_size = c.shape[0] + + # `sample` projections + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) + + # `context` projections + c_query = attn.to_q_c(c) + c_key = attn.to_k_c(c) + c_value = attn.to_v_c(c) + + # attention + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + c_query = c_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + c_key = c_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + c_value = c_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + if attn.c_q_norm is not None: + c_query = attn.c_q_norm(c_query) + if attn.c_k_norm is not None: + c_key = attn.c_k_norm(c_key) + + # apply rope for context and noised input independently + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + if c_rope is not None: + freqs, xpos_scale = c_rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale) + c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale) + + # joint attention + query = torch.cat([query, c_query], dim=2) + key = torch.cat([key, c_key], dim=2) + value = torch.cat([value, c_value], dim=2) + + # mask. e.g. inference got a batch with different target durations, mask out the padding + if mask is not None: + attn_mask = F.pad(mask, (0, c.shape[1]), value=True) # no mask for c (text) + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n' + attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2]) + else: + attn_mask = None + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + x = x.to(query.dtype) + + # Split the attention outputs. + x, c = ( + x[:, : residual.shape[1]], + x[:, residual.shape[1] :], + ) + + # linear proj + x = attn.to_out[0](x) + # dropout + x = attn.to_out[1](x) + if not attn.context_pre_only: + c = attn.to_out_c(c) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + # c = c.masked_fill(~mask, 0.) # no mask for c (text) + + return x, c + + +# DiT Block + + +class DiTBlock(nn.Module): + def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, qk_norm=None, pe_attn_head=None): + super().__init__() + + self.attn_norm = AdaLayerNorm(dim) + self.attn = Attention( + processor=AttnProcessor(pe_attn_head=pe_attn_head), + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + qk_norm=qk_norm, + ) + + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") + + def forward(self, x, t, mask=None, rope=None): # x: noised input, t: time embedding + # pre-norm & modulation for attention input + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + + # attention + attn_output = self.attn(x=norm, mask=mask, rope=rope) + + # process attention output for input x + x = x + gate_msa.unsqueeze(1) * attn_output + + norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(norm) + x = x + gate_mlp.unsqueeze(1) * ff_output + + return x + + +class ChunkDiTBlock(nn.Module): + def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, qk_norm=None, chunk_size=16, block_size=8, t_p=0, t_f=0, pe_attn_head=None): + super().__init__() + + self.attn_norm = AdaLayerNorm(dim) + self.attn = Attention( + processor=BlockAttnProcessor(chunk_size=chunk_size, block_size=block_size, t_p=t_p, t_f=t_f, pe_attn_head=pe_attn_head), + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + qk_norm=qk_norm, + ) + + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") + + def forward(self, x, t, mask=None, rope=None, is_inference=False): # x: noised input, t: time embedding + # pre-norm & modulation for attention input + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + + # attention + attn_output = self.attn(x=norm, mask=mask, rope=rope, is_inference=is_inference) + + # process attention output for input x + x = x + gate_msa.unsqueeze(1) * attn_output + + norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(norm) + x = x + gate_mlp.unsqueeze(1) * ff_output + + return x + +# MMDiT Block https://arxiv.org/abs/2403.03206 + + +class MMDiTBlock(nn.Module): + r""" + modified from diffusers/src/diffusers/models/attention.py + + notes. + _c: context related. text, cond, etc. (left part in sd3 fig2.b) + _x: noised input related. (right part) + context_pre_only: last layer only do prenorm + modulation cuz no more ffn + """ + + def __init__( + self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_dim=None, context_pre_only=False, qk_norm=None + ): + super().__init__() + if context_dim is None: + context_dim = dim + self.context_pre_only = context_pre_only + + self.attn_norm_c = AdaLayerNorm_Final(context_dim) if context_pre_only else AdaLayerNorm(context_dim) + self.attn_norm_x = AdaLayerNorm(dim) + self.attn = Attention( + processor=JointAttnProcessor(), + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + context_dim=context_dim, + context_pre_only=context_pre_only, + qk_norm=qk_norm, + ) + + if not context_pre_only: + self.ff_norm_c = nn.LayerNorm(context_dim, elementwise_affine=False, eps=1e-6) + self.ff_c = FeedForward(dim=context_dim, mult=ff_mult, dropout=dropout, approximate="tanh") + else: + self.ff_norm_c = None + self.ff_c = None + self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") + + def forward(self, x, c, t, mask=None, rope=None, c_rope=None): # x: noised input, c: context, t: time embedding + # pre-norm & modulation for attention input + if self.context_pre_only: + norm_c = self.attn_norm_c(c, t) + else: + norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t) + norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t) + + # attention + x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope) + + # process attention output for context c + if self.context_pre_only: + c = None + else: # if not last layer + c = c + c_gate_msa.unsqueeze(1) * c_attn_output + + norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + c_ff_output = self.ff_c(norm_c) + c = c + c_gate_mlp.unsqueeze(1) * c_ff_output + + # process attention output for input x + x = x + x_gate_msa.unsqueeze(1) * x_attn_output + + norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None] + x_ff_output = self.ff_x(norm_x) + x = x + x_gate_mlp.unsqueeze(1) * x_ff_output + + return c, x + + +# time step conditioning embedding + + +class TimestepEmbedding(nn.Module): + def __init__(self, dim, freq_embed_dim=256): + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: float["b"]): # noqa: F821 + time_hidden = self.time_embed(timestep) + time_hidden = time_hidden.to(timestep.dtype) + time = self.time_mlp(time_hidden) # b d + return time diff --git a/silencio/paradee/__init__.py b/silencio/paradee/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..25f48b2e2d4a07de0fec59efb1d9da4393ba5b65 --- /dev/null +++ b/silencio/paradee/__init__.py @@ -0,0 +1,2 @@ +"""Paradee-8M v1.1 text side (Sahil Mahendrakar, Apache-2.0) on Kokoro modules (hexgrad, Apache-2.0).""" +from .student import TextStudent diff --git a/silencio/paradee/kokoro_modules.py b/silencio/paradee/kokoro_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..329af3bb7363ed737576753e9e9c7f7f64e10384 --- /dev/null +++ b/silencio/paradee/kokoro_modules.py @@ -0,0 +1,267 @@ +# Vendored from https://github.com/hexgrad/kokoro (kokoro/modules.py + AdaIN blocks from kokoro/istftnet.py), Apache-2.0. +# Kept verbatim apart from import paths so Silencio runs without the original repo. +from torch.nn.utils.parametrizations import weight_norm +from transformers import AlbertModel +import math +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +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) + +def get_padding(kernel_size, dilation=1): + return int((kernel_size*dilation - dilation)/2) + + +class AdaIN1d(nn.Module): + def __init__(self, style_dim, num_features): + super().__init__() + # affine should be False, however there's a bug in the old torch.onnx.export (not newer dynamo) that causes the channel dimension to be lost if affine=False. When affine is true, there's additional learnably parameters. This shouldn't really matter setting it to True, since we're in inference mode + self.norm = nn.InstanceNorm1d(num_features, affine=True) + self.fc = nn.Linear(style_dim, num_features*2) + + def forward(self, x, s): + h = self.fc(s) + h = h.view(h.size(0), h.size(1), 1) + gamma, beta = torch.chunk(h, chunks=2, dim=1) + return (1 + gamma) * self.norm(x) + beta + + +class UpSample1d(nn.Module): + def __init__(self, layer_type): + super().__init__() + self.layer_type = layer_type + + def forward(self, x): + if self.layer_type == 'none': + return x + else: + return F.interpolate(x, scale_factor=2, mode='nearest') + + +class AdainResBlk1d(nn.Module): + def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2), upsample='none', dropout_p=0.0): + super().__init__() + self.actv = actv + self.upsample_type = upsample + self.upsample = UpSample1d(upsample) + self.learned_sc = dim_in != dim_out + self._build_weights(dim_in, dim_out, style_dim) + self.dropout = nn.Dropout(dropout_p) + if upsample == 'none': + self.pool = nn.Identity() + else: + self.pool = weight_norm(nn.ConvTranspose1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1)) + + def _build_weights(self, dim_in, dim_out, style_dim): + self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1)) + self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1)) + self.norm1 = AdaIN1d(style_dim, dim_in) + self.norm2 = AdaIN1d(style_dim, dim_out) + if self.learned_sc: + self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False)) + + def _shortcut(self, x): + x = self.upsample(x) + if self.learned_sc: + x = self.conv1x1(x) + return x + + def _residual(self, x, s): + x = self.norm1(x, s) + x = self.actv(x) + x = self.pool(x) + x = self.conv1(self.dropout(x)) + x = self.norm2(x, s) + x = self.actv(x) + x = self.conv2(self.dropout(x)) + return x + + def forward(self, x, s): + out = self._residual(x, s) + out = (out + self._shortcut(x)) * torch.rsqrt(torch.tensor(2)) + return out + + + +# https://github.com/yl4579/StyleTTS2/blob/main/models.py + + +class LinearNorm(nn.Module): + def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'): + super(LinearNorm, self).__init__() + self.linear_layer = nn.Linear(in_dim, out_dim, bias=bias) + nn.init.xavier_uniform_(self.linear_layer.weight, gain=nn.init.calculate_gain(w_init_gain)) + + def forward(self, x): + return self.linear_layer(x) + + +class LayerNorm(nn.Module): + def __init__(self, channels, eps=1e-5): + super().__init__() + self.channels = channels + self.eps = eps + self.gamma = nn.Parameter(torch.ones(channels)) + self.beta = nn.Parameter(torch.zeros(channels)) + + def forward(self, x): + x = x.transpose(1, -1) + x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps) + return x.transpose(1, -1) + + +class TextEncoder(nn.Module): + def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)): + super().__init__() + self.embedding = nn.Embedding(n_symbols, channels) + padding = (kernel_size - 1) // 2 + self.cnn = nn.ModuleList() + for _ in range(depth): + self.cnn.append(nn.Sequential( + weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)), + LayerNorm(channels), + actv, + nn.Dropout(0.2), + )) + self.lstm = nn.LSTM(channels, channels//2, 1, batch_first=True, bidirectional=True) + + def forward(self, x, input_lengths, m): + x = self.embedding(x) # [B, T, emb] + x = x.transpose(1, 2) # [B, emb, T] + m = m.unsqueeze(1) + x.masked_fill_(m, 0.0) + for c in self.cnn: + x = c(x) + x.masked_fill_(m, 0.0) + x = x.transpose(1, 2) # [B, T, chn] + lengths = input_lengths if input_lengths.device == torch.device('cpu') else input_lengths.to('cpu') + x = nn.utils.rnn.pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False) + self.lstm.flatten_parameters() + x, _ = self.lstm(x) + x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True) + x = x.transpose(-1, -2) + x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]], device=x.device) + x_pad[:, :, :x.shape[-1]] = x + x = x_pad + x.masked_fill_(m, 0.0) + return x + + +class AdaLayerNorm(nn.Module): + def __init__(self, style_dim, channels, eps=1e-5): + super().__init__() + self.channels = channels + self.eps = eps + self.fc = nn.Linear(style_dim, channels*2) + + def forward(self, x, s): + x = x.transpose(-1, -2) + x = x.transpose(1, -1) + h = self.fc(s) + h = h.view(h.size(0), h.size(1), 1) + gamma, beta = torch.chunk(h, chunks=2, dim=1) + gamma, beta = gamma.transpose(1, -1), beta.transpose(1, -1) + x = F.layer_norm(x, (self.channels,), eps=self.eps) + x = (1 + gamma) * x + beta + return x.transpose(1, -1).transpose(-1, -2) + + +class ProsodyPredictor(nn.Module): + def __init__(self, style_dim, d_hid, nlayers, max_dur=50, dropout=0.1): + super().__init__() + self.text_encoder = DurationEncoder(sty_dim=style_dim, d_model=d_hid,nlayers=nlayers, dropout=dropout) + self.lstm = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True) + self.duration_proj = LinearNorm(d_hid, max_dur) + self.shared = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True) + self.F0 = nn.ModuleList() + self.F0.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout)) + self.F0.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout)) + self.F0.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout)) + self.N = nn.ModuleList() + self.N.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout)) + self.N.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout)) + self.N.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout)) + self.F0_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0) + self.N_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0) + + def forward(self, texts, style, text_lengths, alignment, m): + d = self.text_encoder(texts, style, text_lengths, m) + m = m.unsqueeze(1) + lengths = text_lengths if text_lengths.device == torch.device('cpu') else text_lengths.to('cpu') + x = nn.utils.rnn.pack_padded_sequence(d, lengths, batch_first=True, enforce_sorted=False) + self.lstm.flatten_parameters() + x, _ = self.lstm(x) + x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True) + x_pad = torch.zeros([x.shape[0], m.shape[-1], x.shape[-1]], device=x.device) + x_pad[:, :x.shape[1], :] = x + x = x_pad + duration = self.duration_proj(nn.functional.dropout(x, 0.5, training=False)) + en = (d.transpose(-1, -2) @ alignment) + return duration.squeeze(-1), en + + def F0Ntrain(self, x, s): + x, _ = self.shared(x.transpose(-1, -2)) + F0 = x.transpose(-1, -2) + for block in self.F0: + F0 = block(F0, s) + F0 = self.F0_proj(F0) + N = x.transpose(-1, -2) + for block in self.N: + N = block(N, s) + N = self.N_proj(N) + return F0.squeeze(1), N.squeeze(1) + + +class DurationEncoder(nn.Module): + def __init__(self, sty_dim, d_model, nlayers, dropout=0.1): + super().__init__() + self.lstms = nn.ModuleList() + for _ in range(nlayers): + self.lstms.append(nn.LSTM(d_model + sty_dim, d_model // 2, num_layers=1, batch_first=True, bidirectional=True)) + self.lstms.append(AdaLayerNorm(sty_dim, d_model)) + self.dropout = dropout + self.d_model = d_model + self.sty_dim = sty_dim + + def forward(self, x, style, text_lengths, m): + masks = m + x = x.permute(2, 0, 1) + s = style.expand(x.shape[0], x.shape[1], -1) + x = torch.cat([x, s], axis=-1) + x.masked_fill_(masks.unsqueeze(-1).transpose(0, 1), 0.0) + x = x.transpose(0, 1) + x = x.transpose(-1, -2) + for block in self.lstms: + if isinstance(block, AdaLayerNorm): + x = block(x.transpose(-1, -2), style).transpose(-1, -2) + x = torch.cat([x, s.permute(1, 2, 0)], axis=1) + x.masked_fill_(masks.unsqueeze(-1).transpose(-1, -2), 0.0) + else: + lengths = text_lengths if text_lengths.device == torch.device('cpu') else text_lengths.to('cpu') + x = x.transpose(-1, -2) + x = nn.utils.rnn.pack_padded_sequence( + x, lengths, batch_first=True, enforce_sorted=False) + block.flatten_parameters() + x, _ = block(x) + x, _ = nn.utils.rnn.pad_packed_sequence( + x, batch_first=True) + x = F.dropout(x, p=self.dropout, training=False) + x = x.transpose(-1, -2) + x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]], device=x.device) + x_pad[:, :, :x.shape[-1]] = x + x = x_pad + + return x.transpose(-1, -2) + + +# https://github.com/yl4579/StyleTTS2/blob/main/Utils/PLBERT/util.py +class CustomAlbert(AlbertModel): + def forward(self, *args, **kwargs): + outputs = super().forward(*args, **kwargs) + return outputs.last_hidden_state diff --git a/silencio/paradee/student.py b/silencio/paradee/student.py new file mode 100644 index 0000000000000000000000000000000000000000..bd47f0a275c2153ccd565abeb7c80618fe0aaf55 --- /dev/null +++ b/silencio/paradee/student.py @@ -0,0 +1,64 @@ +# Vendored from https://github.com/sahilmahendrakar/paradee (training/scripts/student.py), Apache-2.0. +# Kept verbatim apart from import paths so Silencio runs without the original repo. +"""Single-voice text-side student: width-scaled instances of Kokoro's own modules +(CustomAlbert, ProsodyPredictor, TextEncoder) with a learned constant style vector in place +of the voice input, plus a projection to the teacher's 512-channel asr space so the frozen +teacher decoder can consume its output.""" +import torch, torch.nn as nn +from transformers import AlbertConfig +from .kokoro_modules import CustomAlbert, ProsodyPredictor, TextEncoder + +PRESETS = { # name: (albert hidden, albert layers, albert heads, hidden_dim, style_dim, n_layer) + # "+mlp" variants use a 2-layer nonlinear asr projection (the teacher asr features are rank ~430 for 99% variance) + "xs": (128, 4, 4, 128, 16, 2), + "s": (256, 6, 4, 192, 32, 3), + "m": (384, 6, 6, 256, 32, 3), + "l": (512, 8, 8, 384, 64, 3), +} + +class TextStudent(nn.Module): + def __init__(self, preset="s", n_token=178, teacher_dim=512, max_dur=50): + super().__init__() + mlp = preset.endswith("+mlp"); preset = preset.replace("+mlp", "") + ah, al, an, hd, sd, nl = PRESETS[preset] + self.bert = CustomAlbert(AlbertConfig(vocab_size=n_token, hidden_size=ah, num_attention_heads=an, + intermediate_size=ah * 3, num_hidden_layers=al, + max_position_embeddings=512, embedding_size=min(128, ah), dropout=0.1)) + self.bert_encoder = nn.Linear(ah, hd) + self.predictor = ProsodyPredictor(style_dim=sd, d_hid=hd, nlayers=nl, max_dur=max_dur, dropout=0.1) + self.text_encoder = TextEncoder(channels=hd, kernel_size=5, depth=nl, n_symbols=n_token) + self.asr_proj = (nn.Sequential(nn.Linear(hd, teacher_dim), nn.GELU(), nn.Linear(teacher_dim, teacher_dim)) if mlp + else nn.Linear(hd, teacher_dim)) + self.style = nn.Parameter(torch.randn(1, sd) * 0.1) # the one voice + + def forward(self, ids, lengths, aln=None): + """ids [B,T] long, lengths [B]. aln: optional teacher alignment [B,T,F] (teacher forcing). + Returns dur [B,T] (pre-round sigmoid sums), F0 [B,2F], N [B,2F], asr_tok [B,512,T].""" + B, T = ids.shape + mask = torch.arange(T, device=ids.device)[None].expand(B, -1) >= lengths[:, None] + s = self.style.expand(B, -1) + bert = self.bert(ids, attention_mask=(~mask).int()) + d_en = self.bert_encoder(bert).transpose(-1, -2) + d = self.predictor.text_encoder(d_en, s, lengths, mask) + x, _ = self.predictor.lstm(d) + dur = torch.sigmoid(self.predictor.duration_proj(x)).sum(-1) + t_en = self.text_encoder(ids, lengths, mask) + asr_tok = self.asr_proj(t_en.transpose(-1, -2)).transpose(-1, -2) + if aln is None: + return dur, None, None, asr_tok + en = d.transpose(-1, -2) @ aln + F0, N = self.predictor.F0Ntrain(en, s) + return dur, F0, N, asr_tok + +def alignment(pred_dur, F=None): + """pred_dur [T] long -> [T, F] one-hot expansion matrix.""" + idx = torch.repeat_interleave(torch.arange(len(pred_dur), device=pred_dur.device), pred_dur) + F = F or len(idx) + a = torch.zeros(len(pred_dur), F, device=pred_dur.device); a[idx[:F], torch.arange(min(F, len(idx)), device=pred_dur.device)] = 1 + return a + +if __name__ == "__main__": + for p in list(PRESETS) + ["s+mlp"]: + m = TextStudent(p); n = sum(x.numel() for x in m.parameters()) + parts = {k: sum(x.numel() for x in getattr(m, k).parameters()) for k in ["bert", "predictor", "text_encoder", "asr_proj"]} + print(f"{p}: {n/1e6:.2f}M " + " ".join(f"{k} {v/1e6:.2f}M" for k, v in parts.items())) diff --git a/silencio/text.py b/silencio/text.py new file mode 100644 index 0000000000000000000000000000000000000000..a806991a40e69a9aee4352aa7b244b5be606d99b --- /dev/null +++ b/silencio/text.py @@ -0,0 +1,67 @@ +"""Silencio's text encoder: English text -> phonemes -> durations, pitch, loudness and phonetic features. + +This is the text side of Paradee-8M v1.1 (Kokoro's architecture at small width), used as a frozen +encoder. Its output is aligned to 40 Hz frames, which the adapter maps into the renderer's content space. +""" +import torch +import torch.nn as nn + +from .paradee import TextStudent + +MAX_PHONEMES = 510 # 512 positions, two of them the pad tokens at each end + + +def strip_weight_norm(module): + """Fold weight-norm parametrisations into plain weights (how the released weights are stored).""" + from torch.nn.utils import remove_weight_norm + from torch.nn.utils.parametrize import is_parametrized, remove_parametrizations + for m in module.modules(): + try: + remove_weight_norm(m) + except Exception: + pass + if is_parametrized(m, "weight"): + remove_parametrizations(m, "weight", leave_parametrized=True) + return module + + +class TextEncoder(nn.Module): + def __init__(self, vocab): + super().__init__() + self.vocab = vocab + self.ts = strip_weight_norm(TextStudent("s+mlp")) + self._g2p = None + + def phonemize(self, text): + if self._g2p is None: # misaki: Kokoro's own G2P, the spelling Paradee was trained on + from misaki import en, espeak + self._g2p = en.G2P(trf=False, british=False, fallback=espeak.EspeakFallback(british=False), unk="") + return self._g2p(text)[0] + + def ids(self, phonemes): + ids = [self.vocab[c] for c in phonemes if c in self.vocab][:MAX_PHONEMES] + return torch.LongTensor([[0] + ids + [0]]) + + @torch.no_grad() + def forward(self, ids, speed=1.0): + """ids [1,T] -> asr [1,512,F], F0 [1,2F], N [1,2F]; F frames at 40 Hz.""" + t = self.ts; p = t.predictor; s = t.style + d_en = t.bert_encoder(t.bert(ids, attention_mask=torch.ones_like(ids))).transpose(-1, -2) + T = ids.shape[1]; sx = s[:, :, None].expand(-1, -1, T); x = torch.cat([d_en, sx], 1) + for b in p.text_encoder.lstms: + if isinstance(b, nn.LSTM): + x = b(x.transpose(-1, -2))[0].transpose(-1, -2) + else: + x = torch.cat([b(x.transpose(-1, -2), s).transpose(-1, -2), sx], 1) + d = x.transpose(-1, -2) + dur = torch.sigmoid(p.duration_proj(p.lstm(d)[0])).sum(-1)[0] + e = t.text_encoder; h = e.embedding(ids).transpose(1, 2) + for c in e.cnn: + h = c(h) + asr_tok = t.asr_proj(e.lstm(h.transpose(1, 2))[0]).transpose(-1, -2) + n = torch.clamp(torch.round(dur / speed), min=1).long() + end = torch.cumsum(n, 0); start = end - n + f = torch.arange(int(end[-1]), device=ids.device)[None] + aln = ((f >= start[:, None]) & (f < end[:, None])).float()[None] + F0, N = p.F0Ntrain(d.transpose(-1, -2) @ aln, s) + return asr_tok @ aln, F0, N diff --git a/silencio/tts.py b/silencio/tts.py new file mode 100644 index 0000000000000000000000000000000000000000..b93c542d26804a3e1f426d07f8dbe7172a97d9b7 --- /dev/null +++ b/silencio/tts.py @@ -0,0 +1,136 @@ +"""Silencio: zero-shot English text-to-speech from one weights file. + + text -> TextEncoder (durations, pitch, phonetic features @40 Hz) + -> Adapter (content tokens @25 Hz) + -> DiT renderer (mean-flow, 3 steps, streaming chunks, conditioned on a voice) -> 80-bin mel + -> Vocos -> 16 kHz speech + +`convert` swaps the text front-end for MeanVC2's streaming ASR encoder, turning the same renderer into +voice conversion. A voice is a 256-d speaker embedding. Bundled voices ship in voices.safetensors; new ones come from +`silencio.enroll` (needs the WavLM/ECAPA speaker encoder, which is not part of the weights file). +""" +import json +import os +import re +import numpy as np +import torch +import torch.nn as nn + +from .adapter import Adapter, bnf_frames, PARADEE_FPS +from .meanvc2 import DiT +from .text import TextEncoder +from .vocos import Vocos + +SAMPLE_RATE = 16000 +DEFAULT_REPO = "BlueRoseLabs/Silencio" + + +class SilencioTTS(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.text = TextEncoder(config["vocab"]) + self.adapter = Adapter(**config["adapter"]) + self.dit = DiT(**config["dit"]) + self.vocos = Vocos(**config["vocos"]) + self.voices = {} + + # ---- loading ------------------------------------------------------------------------- + @classmethod + def from_pretrained(cls, path_or_repo=DEFAULT_REPO, device=None, revision=None): + from safetensors.torch import load_file + if os.path.isdir(path_or_repo): + get = lambda f: os.path.join(path_or_repo, f) + else: + from huggingface_hub import hf_hub_download + get = lambda f: hf_hub_download(path_or_repo, f, revision=revision) + model = cls(json.load(open(get("config.json")))) + model.load_state_dict(load_file(get("model.safetensors")), strict=True) + try: + model.voices = {k: v.float() for k, v in load_file(get("voices.safetensors")).items()} + except Exception: + model.voices = {} + device = device or ("cuda" if torch.cuda.is_available() else "cpu") + return model.to(device).eval() + + @property + def device(self): + return next(self.parameters()).device + + def voice(self, voice): + """A voice name from voices.safetensors, a .npy/.pt path, or a [256] tensor/array.""" + if voice is None: + voice = self.config.get("default_voice") or next(iter(self.voices)) + if isinstance(voice, str): + if voice in self.voices: + return self.voices[voice] + v = np.load(voice) if voice.endswith(".npy") else torch.load(voice, map_location="cpu") + return torch.as_tensor(v, dtype=torch.float32).reshape(-1) + return torch.as_tensor(voice, dtype=torch.float32).reshape(-1) + + # ---- synthesis ----------------------------------------------------------------------- + @torch.inference_mode() + def content(self, phonemes, speed=1.0): + """Phonemes -> content tokens [T, 256] at 25 Hz.""" + dev = self.device + asr, F0, N = self.text(self.text.ids(phonemes).to(dev), speed) + T = bnf_frames(asr.shape[2] / PARADEE_FPS) + return self.adapter(asr, F0, N, torch.tensor([T], device=dev))[0] + + @torch.inference_mode() + def render(self, content, voice, steps=3, seed=0): + """Content tokens [T, 256] + voice [256] -> 16 kHz waveform. Chunked exactly as MeanVC2 streams.""" + dev, cs, bs = self.device, self.config["chunk_size"], self.config["block_size"] + g = torch.Generator(device=dev).manual_seed(seed) if seed is not None else None + ts = torch.linspace(1.0, 0.0, steps + 1, device=dev) + b = content.to(dev, torch.float32)[None].transpose(1, 2) + b = nn.functional.interpolate(b, size=b.shape[2] * 4, mode="linear", align_corners=True).transpose(1, 2) + spk = self.voice(voice).to(dev).reshape(1, -1) + L = b.shape[1] + noise = torch.randn(1, L, 80, device=dev, generator=g) + outs, offset, kv = [], 0, None + for c in range((L + cs - 1) // cs): + s0 = c * cs; mid = min(s0 + cs, L); end = min(mid + bs, L) + x = noise[:, s0:end] + for i in range(steps): + t, r = ts[i], ts[i + 1] + u, tmp = self.dit(x, t[None], r[None], cache=None, cond=b[:, s0:end], spks=spk, + offset=offset, is_inference=True, kv_cache=kv) + x = x - (t - r) * u + kv = tmp + outs.append(x[:, :mid - s0]); offset += mid - s0 + mel = (torch.cat(outs, 1).transpose(1, 2) + 1) / 2 + return self.vocos.decode(mel).reshape(-1).float().cpu() + + def sentences(self, text): + for s in re.split(r"(?<=[.!?…])\s+|\n+", text.strip()): + if s.strip(): + ps = self.text.phonemize(s) + if ps: + yield ps + + def stream(self, text, voice=None, speed=1.0, steps=3, seed=0): + """Yields float32 numpy chunks at 16 kHz, one sentence at a time.""" + for ps in self.sentences(text): + yield self.render(self.content(ps, speed), voice, steps, seed).numpy() + + def generate(self, text, voice=None, speed=1.0, steps=3, seed=0, pause=0.12): + """Text -> float32 numpy waveform at 16 kHz.""" + gap = np.zeros(int(pause * SAMPLE_RATE), np.float32) + parts = [p for chunk in self.stream(text, voice, speed, steps, seed) for p in (chunk, gap)] + return np.concatenate(parts[:-1]) if parts else np.zeros(0, np.float32) + + __call__ = generate + + def convert(self, wav, sr, voice=None, steps=3, seed=0): + """Voice conversion: speech `wav` at `sr` -> the same words in `voice`, 16 kHz numpy. + Content comes from MeanVC2's streaming Fast-U2++ encoder (downloaded on first use, CPU).""" + from .asr import ContentEncoder, to_16k + if getattr(self, "_asr", None) is None: + self._asr = ContentEncoder() + wav = np.asarray(wav, dtype=np.float32) + return self.render(self._asr(to_16k(wav.mean(-1) if wav.ndim > 1 else wav, sr)), voice, steps, seed).numpy() + + def save(self, path, text, voice=None, **kw): + import soundfile as sf + sf.write(path, self.generate(text, voice, **kw), SAMPLE_RATE) diff --git a/silencio/vocos.py b/silencio/vocos.py new file mode 100644 index 0000000000000000000000000000000000000000..cd8247ff72bdc92894bdac37011f1006d9590ea6 --- /dev/null +++ b/silencio/vocos.py @@ -0,0 +1,81 @@ +"""Vocos vocoder (Siuzdak, MIT) at MeanVC2's 16 kHz setting: 80-bin mel -> waveform. + +The same architecture as the `vocos` package (ConvNeXt backbone + iSTFT head), written out so its +weights can sit in Silencio's single safetensors file instead of MeanVC2's TorchScript archive. +""" +import torch +import torch.nn as nn + + +class ConvNeXtBlock(nn.Module): + def __init__(self, dim, intermediate_dim): + super().__init__() + self.dwconv = nn.Conv1d(dim, dim, 7, padding=3, groups=dim) + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) + self.act = nn.GELU() + self.pwconv2 = nn.Linear(intermediate_dim, dim) + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): # [B, C, T] + h = self.norm(self.dwconv(x).transpose(1, 2)) + return x + (self.gamma * self.pwconv2(self.act(self.pwconv1(h)))).transpose(1, 2) + + +class VocosBackbone(nn.Module): + def __init__(self, input_channels=80, dim=320, intermediate_dim=1536, num_layers=8): + super().__init__() + self.embed = nn.Conv1d(input_channels, dim, 7, padding=3) + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.convnext = nn.ModuleList(ConvNeXtBlock(dim, intermediate_dim) for _ in range(num_layers)) + self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6) + + def forward(self, x): # [B, n_mels, T] -> [B, T, dim] + x = self.norm(self.embed(x).transpose(1, 2)).transpose(1, 2) + for block in self.convnext: + x = block(x) + return self.final_layer_norm(x.transpose(1, 2)) + + +class ISTFT(nn.Module): + """Inverse STFT with Vocos' "center" or "same" padding.""" + + def __init__(self, n_fft, hop_length, win_length, padding="same"): + super().__init__() + self.n_fft, self.hop_length, self.win_length, self.padding = n_fft, hop_length, win_length, padding + self.register_buffer("window", torch.hann_window(win_length)) + + def forward(self, spec): # complex [B, n_fft/2+1, T] + if self.padding == "center": + return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True) + pad = (self.win_length - self.hop_length) // 2 + T = spec.shape[-1] + frames = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") * self.window[None, :, None] + size = (T - 1) * self.hop_length + self.win_length + y = nn.functional.fold(frames, (1, size), (1, self.win_length), stride=(1, self.hop_length))[:, 0, 0, pad:-pad] + env = nn.functional.fold(self.window.square().expand(1, T, -1).transpose(1, 2), (1, size), + (1, self.win_length), stride=(1, self.hop_length)).squeeze()[pad:-pad] + return y / env.clamp_min(1e-11) + + +class ISTFTHead(nn.Module): + def __init__(self, dim=320, n_fft=640, hop_length=160, padding="same"): + super().__init__() + self.out = nn.Linear(dim, n_fft + 2) + self.istft = ISTFT(n_fft, hop_length, n_fft, padding) + + def forward(self, x): # [B, T, dim] -> [B, samples] + mag, p = self.out(x).transpose(1, 2).chunk(2, dim=1) + mag = torch.exp(mag).clip(max=1e2) + return self.istft(mag * (torch.cos(p) + 1j * torch.sin(p))) + + +class Vocos(nn.Module): + def __init__(self, input_channels=80, dim=320, intermediate_dim=1536, num_layers=8, n_fft=640, hop_length=160, + padding="same"): + super().__init__() + self.backbone = VocosBackbone(input_channels, dim, intermediate_dim, num_layers) + self.head = ISTFTHead(dim, n_fft, hop_length, padding) + + def decode(self, mel): # [B, 80, T] -> [B, samples] + return self.head(self.backbone(mel)) diff --git a/voices.safetensors b/voices.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f96c3f2a8870051ab07341b1feeb35199df8f8a0 --- /dev/null +++ b/voices.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:576ca3ab21ab756df8dad72c4a5cd618eb11e1a8a2e7e73dfe3cb733f8ebb348 +size 8808