Silencio v1.0: zero-shot TTS and voice conversion in one model

.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

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_<speaker id>`. The reference clip used for each one is included as `ref_<id>.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).

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"
+}

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
+ }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19f6c837e93857af3d91415e88c08a89f4535f4e0e56e9b313bc87fa5cc24276
+size 134476848

ref_1320.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27e4cf2998ad8b13279721d2d5e21c6a2f9c232b4882a75f7857f760a6899453
+size 250764

ref_2094.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72fab69f58aa47a6cb82c13c8fcb23979098329e52c42db69ac3a8a5649c123d
+size 722284

ref_260.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69df92e49522c52d424c18c1a560295523ca46ddc5e43e11829fed38118d4692
+size 470924

ref_3575.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6aaf21d7747b079cb4f860823e2a9f4b09809474dd7f32754c9bf91a2313d72b
+size 263404

ref_5639.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8069fd8ee2a120f435d682a8e27a9435fa6b80122b5d8843779178d19a3c1c2a
+size 504684

ref_6930.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a0c28ad2cca523cd7e7d72e360a3f314eeb905688358dbf6f150c0a87b87f76
+size 455244

ref_7127.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b5d5625b3cac60e98bddec49206b053d64d7c0d0d871090801b26f8a2755bce
+size 575084

ref_7729.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8e9574a4743df5adca62c99b13aeee6cd48082aed075708545a898435aabde3
+size 265644

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

samples/1_librispeech_1320.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86acc5e690d395b78d16d99cc8381af523e8a061036a89f44b7258883ce020da
+size 306924

samples/1_librispeech_2094.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef3d327c72599c3862c772a75b9143c7976079ba830a3b9bba5835bc665713c7
+size 306924

samples/1_librispeech_260.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aeed717d7e2d0f6bd98bf628429151d6687aa668a2cb023eaa4bdec72118b1d1
+size 306924

samples/1_paradee.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:879b44188d8e1d043b6f85f07d89ebc91047f4b40c1143991d5dc778aa7ea5f9
+size 466844

samples/2_librispeech_1320.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4399e04c86c1e83ad22e258fbdc0161f71a56ebe5ad1786d393fb5f141dd63f2
+size 184044

samples/2_librispeech_2094.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:769d8d22e3a93a5ed42be25f0ebd7629273aaee1bf837eca73217895a0d7c425
+size 184044

samples/2_librispeech_260.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c275d5246288a9d71bcca0d0d90c87e2722d2f415d2579d76eca3c71006dcca8
+size 184044

samples/2_paradee.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e166f7a5e7d332e3fcc3b7977d4ee7bdca18591772432165d69d2f01e9cd98a
+size 282044

samples/3_librispeech_1320.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a76e623642c7d3c1c42ab612335c2c7f4debab2d6e4a1a1e1c98358529877ff
+size 235244

samples/3_librispeech_2094.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc3841776e11eeb75b0472e59480646888a4300d1ed7a7cb3e394762ca500f2f
+size 235244

samples/3_librispeech_260.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16f17e71a6ac27d1ca93731212ebc8490f4c97b5c89a7d92912b68deaca08876
+size 235244

samples/3_paradee.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba7d8300a36b960ed4dd95a36fa185d483e18e7cdae175737c807b92ad347e5a
+size 361244

samples/4_librispeech_1320.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c32ae2b277b885cd31e5799a77c4a51ee64c5502f9f4296d9382cceb89890a5
+size 168684

samples/4_librispeech_2094.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bff1793e67eb2ed58e3b0a0429d85ed5e8551d711ff53c4f258a623ff1c15fa4
+size 168684

samples/4_librispeech_260.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:041da728500f660fc23cdd7eb2a9f2db217e0066728f0c65d3e12576474fad4f
+size 168684

samples/4_paradee.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a0c5dc7a5797da25eb03af806759c22e8e1e6cdec521ebf7151d2cab4b0aaa7
+size 261644

samples/5_librispeech_1320.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5738455475737c27bb67eebae87c6233065e87e3f035249510878e19985b53a
+size 235244

samples/5_librispeech_2094.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:597f15e3edefa292ce3489574f884398236d69f042d73d8165159aad1cdbbd31
+size 235244

samples/5_librispeech_260.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fffda725e02ce7721b3b41937e57c00782e35909145cb50d034b60db25f31f4f
+size 235244

samples/5_paradee.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63135b92c297c2634d21277aa08e00de99aad84bdd5abcaff64b5b89eb322121
+size 358844

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.

samples/vc_1320_to_2094.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:683d8c80da6838c05a20f90591a859af647b6fbf2a4463ee56fe31e1c4e3870d
+size 245484

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)

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)

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)

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)
+

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)

silencio/meanvc2/__init__.py

@@ -0,0 +1,2 @@
+"""MeanVC2 DiT renderer (ASLP-lab, Apache-2.0), vendored for inference."""
+from .dit import DiT

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
+
+
+

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

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

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

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

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)
+    # "<name>+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()))

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

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)