Upload 27 files

Podcaster Dynamic site using Zonos-v0.1, a leading open-weight text-to-speech model trained on more than 200k hours of varied multilingual speech, delivering expressiveness and quality, for TTS podcast creation.

.gitattributes

@@ -33,3 +33,5 @@ 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
+Zonos-main/assets/ArchitectureDiagram.png filter=lfs diff=lfs merge=lfs -text
+Zonos-main/assets/exampleaudio.mp3 filter=lfs diff=lfs merge=lfs -text

Zonos-main/.gitignore

@@ -0,0 +1,13 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+# Misc.
+.ipynb_checkpoints/

Zonos-main/.python-version

@@ -0,0 +1 @@
+3.12

Zonos-main/CONDITIONING_README.md

@@ -0,0 +1,120 @@
+# Conditioning explanations
+Here we will list out all the conditionings the model accepts as well as a short description and some tips for optimal use. For conditionings with a learned unconditional, they can be set to that to allow the model to infer an appropriate setting.
+### espeak
+- **Type:** `EspeakPhonemeConditioner`
+- **Description:**  
+  Responsible for cleaning, phonemicizing, tokenizing, and embedding the text provided to the model. This is the text pre-processing pipeline. If you would like to change how a word is pronounced or enter raw phonemes you can do that here.
+
+  Supported by transformer and hybrid models.
+---
+### speaker
+- **Type:** `PassthroughConditioner`
+- **Attributes:**
+  - **cond_dim:** `128`
+  - **uncond_type:** `learned`
+  - **projection:** `linear`
+- **Description:**  
+  An embedded representation of the speakers voice. We use [these](https://huggingface.co/Zyphra/Zonos-v0.1-speaker-embedding) speaker embedding models. It can capture a surprising amount of detail from the reference clip and supports arbitrary length input. Try to input clean reference clips containing only speech. It can be valid to concatenate multiple clean samples from the same speaker into one long sample and may lead to better cloning. If the speaker clip is very long, it is advisable to cut out long speech-free background music segments if they exist. If the reference clip is yielding noisy outputs with denoising enabled we recommend doing source separation before cloning.
+
+  Supported by transformer and hybrid models.
+---
+### emotion
+- **Type:** `FourierConditioner`
+- **Attributes:**
+  - **input_dim:** `8`
+  - **uncond_type:** `learned`
+- **Description:**  
+  Encodes emotion in an 8D vector. Included emotions are Happiness, Sadness, Disgust, Fear, Surprise, Anger, Other, Neutral in that order. This vector tends to be entangled with various other conditioning inputs. More notably, it's entangled with text based on the text sentiment (eg. Angry texts will be more effectively conditioned to be angry, but if you try to make it sound sad it will be a lot less effective). It's also entangled with pitch standard deviation since larger values there tend to correlate to more emotional utterances. It's also heavily correlated with VQScore and DNSMOS as these conditionings favor neutral speech. It's also possible to do a form of "negative prompting" by doing CFG where the unconditional branch is set to a highly neutral emotion vector instead of the true unconditional value, doing this will exaggerate the emotions as it pushes the model away from being neutral.
+
+  Supported by transformer and hybrid models.
+---
+### fmax
+- **Type:** `FourierConditioner`
+- **Attributes:**
+  - **min_val:** `0`
+  - **max_val:** `24000`
+  - **uncond_type:** `learned`
+- **Description:**  
+  Specifies the max frequency of the audio. For best results select 22050 or 24000 as these correspond to 44.1 and 48KHz audio respectively. They should not be any different in terms of actual max frequency since the model's sampling rate is 44.1KHz but they represent different slices of data which lead to slightly different voicing. Selecting a lower value generally produces lower-quality results both in terms of acoustics and voicing.
+
+  For voice cloning it is recommended to use 22050.
+
+  Supported by transformer and hybrid models.
+---
+### pitch_std
+- **Type:** `FourierConditioner`
+- **Attributes:**
+  - **min_val:** `0`
+  - **max_val:** `400`
+  - **uncond_type:** `learned`
+- **Description:**  
+  Specifies the standard deviation of the pitch of the output audio. Wider variations of pitch tend to be more correlated with expressive speech. Good values are from 20-45 for normal speech and 60-150 for expressive speech. Higher than that generally tend to be crazier samples.
+
+  Supported by transformer and hybrid models.
+---
+### speaking_rate
+- **Type:** `FourierConditioner`
+- **Attributes:**
+  - **min_val:** `0`
+  - **max_val:** `40`
+  - **uncond_type:** `learned`
+- **Description:**  
+  Specifies the number of phonemes to be read per second. When entering a long text, it is advisable to adjust the speaking rate such that the number of phonemes is readable within the generation length. For example, if your generation length is 10 seconds, and your input is 300 phonemes, you would want either 30 phonemes per second (which is very very fast) or to generate a longer sample. The model's maximum is 30 seconds. Please note that unrealistic speaking rates can be OOD for the model and create undesirable effects, so at the 30-second limit, it can be better to cut the text short and do multiple generations than to feed the model the entire prompt and have an unrealistically low speaking rate.
+  
+  Supported by transformer and hybrid models.
+---
+### language_id
+- **Type:** `IntegerConditioner`
+- **Attributes:**
+  - **min_val:** `-1`
+  - **max_val:** `126`
+  - **uncond_type:** `learned`
+- **Description:**  
+  Indicates which language the output should be in. A mapping for these values can be found in the [conditioning section](https://github.com/Zyphra/Zonos/blob/3807c8e04bd4beaadb9502b3df1ffa4b0350e3f7/zonos/conditioning.py#L308C1-L376C21) of Zonos.
+
+  Supported by transformer and hybrid models.
+---
+### vqscore_8
+- **Type:** `FourierConditioner`
+- **Attributes:**
+  - **input_dim:** `8`
+  - **min_val:** `0.5`
+  - **max_val:** `0.8`
+  - **uncond_type:** `learned`
+- **Description:**  
+  Encodes the desired [VQScore](https://github.com/JasonSWFu/VQscore) value for the output audio. VQScore is an unsupervised speech quality (cleanliness) estimation method that we found has superior generalization and reduced biases compared to supervised methods like DNSMOS. A good value for our model is 0.78 for high-quality speech. The eight dimensions correspond to consecutive 1/8th chunks of the audio. (eg. for an 8-second output, the first dimension represents the quality of the first second only). For inference, we generally set all 8 dimensions to the same value. This has an unfortunately strong correlation with expressiveness, so for expressive speech, we recommend setting it to unconditional.
+
+  Only applicable for the hybrid model.
+---
+### ctc_loss
+- **Type:** `FourierConditioner`
+- **Attributes:**
+  - **min_val:** `-1.0`
+  - **max_val:** `1000`
+  - **uncond_type:** `learned`
+- **Description:**  
+  Encodes loss values from a [CTC](https://en.wikipedia.org/wiki/Connectionist_temporal_classification) (Connectionist Temporal Classification) setup, this indicates how well the training-time transcription matched with the audio according to a CTC model. For inference always use low values (eg. 0.0 or 1.0)
+
+  Only applicable for the hybrid model.
+---
+### dnsmos_ovrl
+- **Type:** `FourierConditioner`
+- **Attributes:**
+  - **min_val:** `1`
+  - **max_val:** `5`
+  - **uncond_type:** `learned`
+- **Description:**  
+  A [MOS](https://arxiv.org/abs/2110.01763) score for the output audio. This is similar to VQScore and tends to have a stronger entanglement with emotions. It additionally has a strong entanglement with languages. Set to 4.0 for very clean and neutral English speech, else we recommend setting it to unconditional.
+
+  Only applicable for the hybrid model.
+---
+### speaker_noised
+- **Type:** `IntegerConditioner`
+- **Attributes:**
+  - **min_val:** `0`
+  - **max_val:** `1`
+  - **uncond_type:** `learned`
+- **Description:**  
+  Indicates if the speaker embedding is noisy or not. If checked this lets the model clean (denoise) the input speaker embedding. When this is set to True, VQScore and DNSMOS will have a lot more power to clean the speaker embedding, so for very noisy input samples we recommend setting this to True and specifying a high VQScore value. If your speaker cloning outputs sound echo-y or do weird things, setting this to True will help.
+
+  Only applicable for the hybrid model.

Zonos-main/Dockerfile

@@ -0,0 +1,11 @@
+FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel
+RUN pip install uv
+
+RUN apt update && \
+    apt install -y espeak-ng && \
+    rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+COPY . ./
+
+RUN uv pip install --system -e . && uv pip install --system -e .[compile]

Zonos-main/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

Zonos-main/README.md

@@ -0,0 +1,159 @@
+# Zonos-v0.1
+
+<div align="center">
+<img src="assets/ZonosHeader.png" 
+     alt="Alt text" 
+     style="width: 500px;
+            height: auto;
+            object-position: center top;">
+</div>
+
+<div align="center">
+  <a href="https://discord.gg/gTW9JwST8q" target="_blank">
+    <img src="https://img.shields.io/badge/Join%20Our%20Discord-7289DA?style=for-the-badge&logo=discord&logoColor=white" alt="Discord">
+  </a>
+</div>
+
+---
+
+Zonos-v0.1 is a leading open-weight text-to-speech model trained on more than 200k hours of varied multilingual speech, delivering expressiveness and quality on par with—or even surpassing—top TTS providers.
+
+Our model enables highly natural speech generation from text prompts when given a speaker embedding or audio prefix, and can accurately perform speech cloning when given a reference clip spanning just a few seconds. The conditioning setup also allows for fine control over speaking rate, pitch variation, audio quality, and emotions such as happiness, fear, sadness, and anger. The model outputs speech natively at 44kHz.
+
+##### For more details and speech samples, check out our blog [here](https://www.zyphra.com/post/beta-release-of-zonos-v0-1)
+
+##### We also have a hosted version available at [playground.zyphra.com/audio](https://playground.zyphra.com/audio)
+
+---
+
+Zonos follows a straightforward architecture: text normalization and phonemization via eSpeak, followed by DAC token prediction through a transformer or hybrid backbone. An overview of the architecture can be seen below.
+
+<div align="center">
+<img src="assets/ArchitectureDiagram.png" 
+     alt="Alt text" 
+     style="width: 1000px;
+            height: auto;
+            object-position: center top;">
+</div>
+
+---
+
+## Usage
+
+### Python
+
+```python
+import torch
+import torchaudio
+from zonos.model import Zonos
+from zonos.conditioning import make_cond_dict
+from zonos.utils import DEFAULT_DEVICE as device
+
+# model = Zonos.from_pretrained("Zyphra/Zonos-v0.1-hybrid", device=device)
+model = Zonos.from_pretrained("Zyphra/Zonos-v0.1-transformer", device=device)
+
+wav, sampling_rate = torchaudio.load("assets/exampleaudio.mp3")
+speaker = model.make_speaker_embedding(wav, sampling_rate)
+
+cond_dict = make_cond_dict(text="Hello, world!", speaker=speaker, language="en-us")
+conditioning = model.prepare_conditioning(cond_dict)
+
+codes = model.generate(conditioning)
+
+wavs = model.autoencoder.decode(codes).cpu()
+torchaudio.save("sample.wav", wavs[0], model.autoencoder.sampling_rate)
+```
+
+### Gradio interface (recommended)
+
+```bash
+uv run gradio_interface.py
+# python gradio_interface.py
+```
+
+This should produce a `sample.wav` file in your project root directory.
+
+_For repeated sampling we highly recommend using the gradio interface instead, as the minimal example needs to load the model every time it is run._
+
+## Features
+
+- Zero-shot TTS with voice cloning: Input desired text and a 10-30s speaker sample to generate high quality TTS output
+- Audio prefix inputs: Add text plus an audio prefix for even richer speaker matching. Audio prefixes can be used to elicit behaviours such as whispering which can otherwise be challenging to replicate when cloning from speaker embeddings
+- Multilingual support: Zonos-v0.1 supports English, Japanese, Chinese, French, and German
+- Audio quality and emotion control: Zonos offers fine-grained control of many aspects of the generated audio. These include speaking rate, pitch, maximum frequency, audio quality, and various emotions such as happiness, anger, sadness, and fear.
+- Fast: our model runs with a real-time factor of ~2x on an RTX 4090 (i.e. generates 2 seconds of audio per 1 second of compute time)
+- Gradio WebUI: Zonos comes packaged with an easy to use gradio interface to generate speech
+- Simple installation and deployment: Zonos can be installed and deployed simply using the docker file packaged with our repository.
+
+## Installation
+
+#### System requirements
+
+- **Operating System:** Linux (preferably Ubuntu 22.04/24.04), macOS
+- **GPU:** 6GB+ VRAM, Hybrid additionally requires a 3000-series or newer Nvidia GPU
+
+Note: Zonos can also run on CPU provided there is enough free RAM. However, this will be a lot slower than running on a dedicated GPU, and likely won't be sufficient for interactive use.
+
+For experimental windows support check out [this fork](https://github.com/sdbds/Zonos-for-windows).
+
+See also [Docker Installation](#docker-installation)
+
+#### System dependencies
+
+Zonos depends on the eSpeak library phonemization. You can install it on Ubuntu with the following command:
+
+```bash
+apt install -y espeak-ng # For Ubuntu
+# brew install espeak-ng # For MacOS
+```
+
+#### Python dependencies
+
+We highly recommend using a recent version of [uv](https://docs.astral.sh/uv/#installation) for installation. If you don't have uv installed, you can install it via pip: `pip install -U uv`.
+
+##### Installing into a new uv virtual environment (recommended)
+
+```bash
+uv sync
+uv sync --extra compile # optional but needed to run the hybrid
+uv pip install -e .
+```
+
+##### Installing into the system/actived environment using uv
+
+```bash
+uv pip install -e .
+uv pip install -e .[compile] # optional but needed to run the hybrid
+```
+
+##### Installing into the system/actived environment using pip
+
+```bash
+pip install -e .
+pip install --no-build-isolation -e .[compile] # optional but needed to run the hybrid
+```
+
+##### Confirm that it's working
+
+For convenience we provide a minimal example to check that the installation works:
+
+```bash
+uv run sample.py
+# python sample.py
+```
+
+## Docker installation
+
+```bash
+git clone https://github.com/Zyphra/Zonos.git
+cd Zonos
+
+# For gradio
+docker compose up
+
+# Or for development you can do
+docker build -t zonos .
+docker run -it --gpus=all --net=host -v /path/to/Zonos:/Zonos -t zonos
+cd /Zonos
+python sample.py # this will generate a sample.wav in /Zonos
+```

Zonos-main/assets/exampleaudio.mp3

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

Zonos-main/docker-compose.yml

@@ -0,0 +1,16 @@
+version: '3.8'
+
+services:
+  zonos:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    container_name: zonos_container
+    runtime: nvidia
+    network_mode: "host"
+    stdin_open: true
+    tty: true
+    command: ["python3", "gradio_interface.py"]
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=0
+      - GRADIO_SHARE=False

Zonos-main/gradio_interface.py

@@ -0,0 +1,419 @@
+import torch
+import torchaudio
+import gradio as gr
+from os import getenv
+
+from zonos.model import Zonos, DEFAULT_BACKBONE_CLS as ZonosBackbone
+from zonos.conditioning import make_cond_dict, supported_language_codes
+from zonos.utils import DEFAULT_DEVICE as device
+
+CURRENT_MODEL_TYPE = None
+CURRENT_MODEL = None
+
+SPEAKER_EMBEDDING = None
+SPEAKER_AUDIO_PATH = None
+
+
+def load_model_if_needed(model_choice: str):
+    global CURRENT_MODEL_TYPE, CURRENT_MODEL
+    if CURRENT_MODEL_TYPE != model_choice:
+        if CURRENT_MODEL is not None:
+            del CURRENT_MODEL
+            torch.cuda.empty_cache()
+        print(f"Loading {model_choice} model...")
+        CURRENT_MODEL = Zonos.from_pretrained(model_choice, device=device)
+        CURRENT_MODEL.requires_grad_(False).eval()
+        CURRENT_MODEL_TYPE = model_choice
+        print(f"{model_choice} model loaded successfully!")
+    return CURRENT_MODEL
+
+
+def update_ui(model_choice):
+    """
+    Dynamically show/hide UI elements based on the model's conditioners.
+    We do NOT display 'language_id' or 'ctc_loss' even if they exist in the model.
+    """
+    model = load_model_if_needed(model_choice)
+    cond_names = [c.name for c in model.prefix_conditioner.conditioners]
+    print("Conditioners in this model:", cond_names)
+
+    text_update = gr.update(visible=("espeak" in cond_names))
+    language_update = gr.update(visible=("espeak" in cond_names))
+    speaker_audio_update = gr.update(visible=("speaker" in cond_names))
+    prefix_audio_update = gr.update(visible=True)
+    emotion1_update = gr.update(visible=("emotion" in cond_names))
+    emotion2_update = gr.update(visible=("emotion" in cond_names))
+    emotion3_update = gr.update(visible=("emotion" in cond_names))
+    emotion4_update = gr.update(visible=("emotion" in cond_names))
+    emotion5_update = gr.update(visible=("emotion" in cond_names))
+    emotion6_update = gr.update(visible=("emotion" in cond_names))
+    emotion7_update = gr.update(visible=("emotion" in cond_names))
+    emotion8_update = gr.update(visible=("emotion" in cond_names))
+    vq_single_slider_update = gr.update(visible=("vqscore_8" in cond_names))
+    fmax_slider_update = gr.update(visible=("fmax" in cond_names))
+    pitch_std_slider_update = gr.update(visible=("pitch_std" in cond_names))
+    speaking_rate_slider_update = gr.update(visible=("speaking_rate" in cond_names))
+    dnsmos_slider_update = gr.update(visible=("dnsmos_ovrl" in cond_names))
+    speaker_noised_checkbox_update = gr.update(visible=("speaker_noised" in cond_names))
+    unconditional_keys_update = gr.update(
+        choices=[name for name in cond_names if name not in ("espeak", "language_id")]
+    )
+
+    return (
+        text_update,
+        language_update,
+        speaker_audio_update,
+        prefix_audio_update,
+        emotion1_update,
+        emotion2_update,
+        emotion3_update,
+        emotion4_update,
+        emotion5_update,
+        emotion6_update,
+        emotion7_update,
+        emotion8_update,
+        vq_single_slider_update,
+        fmax_slider_update,
+        pitch_std_slider_update,
+        speaking_rate_slider_update,
+        dnsmos_slider_update,
+        speaker_noised_checkbox_update,
+        unconditional_keys_update,
+    )
+
+
+def generate_audio(
+    model_choice,
+    text,
+    language,
+    speaker_audio,
+    prefix_audio,
+    e1,
+    e2,
+    e3,
+    e4,
+    e5,
+    e6,
+    e7,
+    e8,
+    vq_single,
+    fmax,
+    pitch_std,
+    speaking_rate,
+    dnsmos_ovrl,
+    speaker_noised,
+    cfg_scale,
+    top_p,
+    top_k,
+    min_p,
+    linear,
+    confidence,
+    quadratic,
+    seed,
+    randomize_seed,
+    unconditional_keys,
+    progress=gr.Progress(),
+):
+    """
+    Generates audio based on the provided UI parameters.
+    We do NOT use language_id or ctc_loss even if the model has them.
+    """
+    selected_model = load_model_if_needed(model_choice)
+
+    speaker_noised_bool = bool(speaker_noised)
+    fmax = float(fmax)
+    pitch_std = float(pitch_std)
+    speaking_rate = float(speaking_rate)
+    dnsmos_ovrl = float(dnsmos_ovrl)
+    cfg_scale = float(cfg_scale)
+    top_p = float(top_p)
+    top_k = int(top_k)
+    min_p = float(min_p)
+    linear = float(linear)
+    confidence = float(confidence)
+    quadratic = float(quadratic)
+    seed = int(seed)
+    max_new_tokens = 86 * 30
+
+    # This is a bit ew, but works for now.
+    global SPEAKER_AUDIO_PATH, SPEAKER_EMBEDDING
+
+    if randomize_seed:
+        seed = torch.randint(0, 2**32 - 1, (1,)).item()
+    torch.manual_seed(seed)
+
+    if speaker_audio is not None and "speaker" not in unconditional_keys:
+        if speaker_audio != SPEAKER_AUDIO_PATH:
+            print("Recomputed speaker embedding")
+            wav, sr = torchaudio.load(speaker_audio)
+            SPEAKER_EMBEDDING = selected_model.make_speaker_embedding(wav, sr)
+            SPEAKER_EMBEDDING = SPEAKER_EMBEDDING.to(device, dtype=torch.bfloat16)
+            SPEAKER_AUDIO_PATH = speaker_audio
+
+    audio_prefix_codes = None
+    if prefix_audio is not None:
+        wav_prefix, sr_prefix = torchaudio.load(prefix_audio)
+        wav_prefix = wav_prefix.mean(0, keepdim=True)
+        wav_prefix = selected_model.autoencoder.preprocess(wav_prefix, sr_prefix)
+        wav_prefix = wav_prefix.to(device, dtype=torch.float32)
+        audio_prefix_codes = selected_model.autoencoder.encode(wav_prefix.unsqueeze(0))
+
+    emotion_tensor = torch.tensor(list(map(float, [e1, e2, e3, e4, e5, e6, e7, e8])), device=device)
+
+    vq_val = float(vq_single)
+    vq_tensor = torch.tensor([vq_val] * 8, device=device).unsqueeze(0)
+
+    cond_dict = make_cond_dict(
+        text=text,
+        language=language,
+        speaker=SPEAKER_EMBEDDING,
+        emotion=emotion_tensor,
+        vqscore_8=vq_tensor,
+        fmax=fmax,
+        pitch_std=pitch_std,
+        speaking_rate=speaking_rate,
+        dnsmos_ovrl=dnsmos_ovrl,
+        speaker_noised=speaker_noised_bool,
+        device=device,
+        unconditional_keys=unconditional_keys,
+    )
+    conditioning = selected_model.prepare_conditioning(cond_dict)
+
+    estimated_generation_duration = 30 * len(text) / 400
+    estimated_total_steps = int(estimated_generation_duration * 86)
+
+    def update_progress(_frame: torch.Tensor, step: int, _total_steps: int) -> bool:
+        progress((step, estimated_total_steps))
+        return True
+
+    codes = selected_model.generate(
+        prefix_conditioning=conditioning,
+        audio_prefix_codes=audio_prefix_codes,
+        max_new_tokens=max_new_tokens,
+        cfg_scale=cfg_scale,
+        batch_size=1,
+        sampling_params=dict(top_p=top_p, top_k=top_k, min_p=min_p, linear=linear, conf=confidence, quad=quadratic),
+        callback=update_progress,
+    )
+
+    wav_out = selected_model.autoencoder.decode(codes).cpu().detach()
+    sr_out = selected_model.autoencoder.sampling_rate
+    if wav_out.dim() == 2 and wav_out.size(0) > 1:
+        wav_out = wav_out[0:1, :]
+    return (sr_out, wav_out.squeeze().numpy()), seed
+
+
+def build_interface():
+    supported_models = []
+    if "transformer" in ZonosBackbone.supported_architectures:
+        supported_models.append("Zyphra/Zonos-v0.1-transformer")
+
+    if "hybrid" in ZonosBackbone.supported_architectures:
+        supported_models.append("Zyphra/Zonos-v0.1-hybrid")
+    else:
+        print(
+            "| The current ZonosBackbone does not support the hybrid architecture, meaning only the transformer model will be available in the model selector.\n"
+            "| This probably means the mamba-ssm library has not been installed."
+        )
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                model_choice = gr.Dropdown(
+                    choices=supported_models,
+                    value=supported_models[0],
+                    label="Zonos Model Type",
+                    info="Select the model variant to use.",
+                )
+                text = gr.Textbox(
+                    label="Text to Synthesize",
+                    value="Zonos uses eSpeak for text to phoneme conversion!",
+                    lines=4,
+                    max_length=500,  # approximately
+                )
+                language = gr.Dropdown(
+                    choices=supported_language_codes,
+                    value="en-us",
+                    label="Language Code",
+                    info="Select a language code.",
+                )
+            prefix_audio = gr.Audio(
+                value="assets/silence_100ms.wav",
+                label="Optional Prefix Audio (continue from this audio)",
+                type="filepath",
+            )
+            with gr.Column():
+                speaker_audio = gr.Audio(
+                    label="Optional Speaker Audio (for cloning)",
+                    type="filepath",
+                )
+                speaker_noised_checkbox = gr.Checkbox(label="Denoise Speaker?", value=False)
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("## Conditioning Parameters")
+                dnsmos_slider = gr.Slider(1.0, 5.0, value=4.0, step=0.1, label="DNSMOS Overall")
+                fmax_slider = gr.Slider(0, 24000, value=24000, step=1, label="Fmax (Hz)")
+                vq_single_slider = gr.Slider(0.5, 0.8, 0.78, 0.01, label="VQ Score")
+                pitch_std_slider = gr.Slider(0.0, 300.0, value=45.0, step=1, label="Pitch Std")
+                speaking_rate_slider = gr.Slider(5.0, 30.0, value=15.0, step=0.5, label="Speaking Rate")
+
+            with gr.Column():
+                gr.Markdown("## Generation Parameters")
+                cfg_scale_slider = gr.Slider(1.0, 5.0, 2.0, 0.1, label="CFG Scale")
+                seed_number = gr.Number(label="Seed", value=420, precision=0)
+                randomize_seed_toggle = gr.Checkbox(label="Randomize Seed (before generation)", value=True)
+
+        with gr.Accordion("Sampling", open=False):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### NovelAi's unified sampler")
+                    linear_slider = gr.Slider(-2.0, 2.0, 0.5, 0.01, label="Linear (set to 0 to disable unified sampling)", info="High values make the output less random.")
+                    #Conf's theoretical range is between -2 * Quad and 0.
+                    confidence_slider = gr.Slider(-2.0, 2.0, 0.40, 0.01, label="Confidence", info="Low values make random outputs more random.")
+                    quadratic_slider = gr.Slider(-2.0, 2.0, 0.00, 0.01, label="Quadratic", info="High values make low probablities much lower.")
+                with gr.Column():
+                    gr.Markdown("### Legacy sampling")
+                    top_p_slider = gr.Slider(0.0, 1.0, 0, 0.01, label="Top P")
+                    min_k_slider = gr.Slider(0.0, 1024, 0, 1, label="Min K")
+                    min_p_slider = gr.Slider(0.0, 1.0, 0, 0.01, label="Min P")
+
+        with gr.Accordion("Advanced Parameters", open=False):
+            gr.Markdown(
+                "### Unconditional Toggles\n"
+                "Checking a box will make the model ignore the corresponding conditioning value and make it unconditional.\n"
+                'Practically this means the given conditioning feature will be unconstrained and "filled in automatically".'
+            )
+            with gr.Row():
+                unconditional_keys = gr.CheckboxGroup(
+                    [
+                        "speaker",
+                        "emotion",
+                        "vqscore_8",
+                        "fmax",
+                        "pitch_std",
+                        "speaking_rate",
+                        "dnsmos_ovrl",
+                        "speaker_noised",
+                    ],
+                    value=["emotion"],
+                    label="Unconditional Keys",
+                )
+
+            gr.Markdown(
+                "### Emotion Sliders\n"
+                "Warning: The way these sliders work is not intuitive and may require some trial and error to get the desired effect.\n"
+                "Certain configurations can cause the model to become unstable. Setting emotion to unconditional may help."
+            )
+            with gr.Row():
+                emotion1 = gr.Slider(0.0, 1.0, 1.0, 0.05, label="Happiness")
+                emotion2 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Sadness")
+                emotion3 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Disgust")
+                emotion4 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Fear")
+            with gr.Row():
+                emotion5 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Surprise")
+                emotion6 = gr.Slider(0.0, 1.0, 0.05, 0.05, label="Anger")
+                emotion7 = gr.Slider(0.0, 1.0, 0.1, 0.05, label="Other")
+                emotion8 = gr.Slider(0.0, 1.0, 0.2, 0.05, label="Neutral")
+
+        with gr.Column():
+            generate_button = gr.Button("Generate Audio")
+            output_audio = gr.Audio(label="Generated Audio", type="numpy", autoplay=True)
+
+        model_choice.change(
+            fn=update_ui,
+            inputs=[model_choice],
+            outputs=[
+                text,
+                language,
+                speaker_audio,
+                prefix_audio,
+                emotion1,
+                emotion2,
+                emotion3,
+                emotion4,
+                emotion5,
+                emotion6,
+                emotion7,
+                emotion8,
+                vq_single_slider,
+                fmax_slider,
+                pitch_std_slider,
+                speaking_rate_slider,
+                dnsmos_slider,
+                speaker_noised_checkbox,
+                unconditional_keys,
+            ],
+        )
+
+        # On page load, trigger the same UI refresh
+        demo.load(
+            fn=update_ui,
+            inputs=[model_choice],
+            outputs=[
+                text,
+                language,
+                speaker_audio,
+                prefix_audio,
+                emotion1,
+                emotion2,
+                emotion3,
+                emotion4,
+                emotion5,
+                emotion6,
+                emotion7,
+                emotion8,
+                vq_single_slider,
+                fmax_slider,
+                pitch_std_slider,
+                speaking_rate_slider,
+                dnsmos_slider,
+                speaker_noised_checkbox,
+                unconditional_keys,
+            ],
+        )
+
+        # Generate audio on button click
+        generate_button.click(
+            fn=generate_audio,
+            inputs=[
+                model_choice,
+                text,
+                language,
+                speaker_audio,
+                prefix_audio,
+                emotion1,
+                emotion2,
+                emotion3,
+                emotion4,
+                emotion5,
+                emotion6,
+                emotion7,
+                emotion8,
+                vq_single_slider,
+                fmax_slider,
+                pitch_std_slider,
+                speaking_rate_slider,
+                dnsmos_slider,
+                speaker_noised_checkbox,
+                cfg_scale_slider,
+                top_p_slider,
+                min_k_slider,
+                min_p_slider,
+                linear_slider,
+                confidence_slider,
+                quadratic_slider,
+                seed_number,
+                randomize_seed_toggle,
+                unconditional_keys,
+            ],
+            outputs=[output_audio, seed_number],
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = build_interface()
+    share = getenv("GRADIO_SHARE", "False").lower() in ("true", "1", "t")
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=share)

Zonos-main/pyproject.toml

@@ -0,0 +1,39 @@
+[project]
+name = "zonos"
+version = "0.1.0"
+description = "Text-to-speech by Zyphra"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=2.5.1",
+    "setuptools",
+    "packaging",
+    "inflect>=7.5.0",
+    "kanjize>=1.5.0",
+    "numpy>=2.2.2",
+    "phonemizer>=3.3.0",
+    "sudachidict-full>=20241021",
+    "sudachipy>=0.6.10",
+    "torchaudio>=2.5.1",
+    "transformers>=4.48.1",
+    "soundfile>=0.13.1",
+    "huggingface-hub>=0.28.1",
+    "gradio>=5.15.0",
+]
+
+# These are technically optional, but mamba-ssm is required to run hybrid models.
+[project.optional-dependencies]
+compile = [
+    "flash-attn>=2.7.3",
+    "mamba-ssm>=2.2.4",
+    "causal-conv1d>=1.5.0.post8",
+]
+
+[tool.setuptools.packages.find]
+include = ["zonos"]
+
+[tool.uv]
+no-build-isolation-package = ["flash-attn", "mamba-ssm", "causal-conv1d"]
+
+[tool.ruff]
+line-length = 120

Zonos-main/sample.py

@@ -0,0 +1,21 @@
+import torch
+import torchaudio
+from zonos.model import Zonos
+from zonos.conditioning import make_cond_dict
+from zonos.utils import DEFAULT_DEVICE as device
+
+# model = Zonos.from_pretrained("Zyphra/Zonos-v0.1-hybrid", device=device)
+model = Zonos.from_pretrained("Zyphra/Zonos-v0.1-transformer", device=device)
+
+wav, sampling_rate = torchaudio.load("assets/exampleaudio.mp3")
+speaker = model.make_speaker_embedding(wav, sampling_rate)
+
+torch.manual_seed(421)
+
+cond_dict = make_cond_dict(text="Hello, world!", speaker=speaker, language="en-us")
+conditioning = model.prepare_conditioning(cond_dict)
+
+codes = model.generate(conditioning)
+
+wavs = model.autoencoder.decode(codes).cpu()
+torchaudio.save("sample.wav", wavs[0], model.autoencoder.sampling_rate)

Zonos-main/zonos/autoencoder.py

@@ -0,0 +1,27 @@
+import math
+
+import torch
+import torchaudio
+from transformers.models.dac import DacModel
+
+
+class DACAutoencoder:
+    def __init__(self):
+        super().__init__()
+        self.dac = DacModel.from_pretrained("descript/dac_44khz")
+        self.dac.eval().requires_grad_(False)
+        self.codebook_size = self.dac.config.codebook_size
+        self.num_codebooks = self.dac.quantizer.n_codebooks
+        self.sampling_rate = self.dac.config.sampling_rate
+
+    def preprocess(self, wav: torch.Tensor, sr: int) -> torch.Tensor:
+        wav = torchaudio.functional.resample(wav, sr, 44_100)
+        right_pad = math.ceil(wav.shape[-1] / 512) * 512 - wav.shape[-1]
+        return torch.nn.functional.pad(wav, (0, right_pad))
+
+    def encode(self, wav: torch.Tensor) -> torch.Tensor:
+        return self.dac.encode(wav).audio_codes
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        with torch.autocast(self.dac.device.type, torch.float16, enabled=self.dac.device.type != "cpu"):
+            return self.dac.decode(audio_codes=codes).audio_values.unsqueeze(1).float()

Zonos-main/zonos/backbone/__init__.py

@@ -0,0 +1,12 @@
+BACKBONES = {}
+
+try:
+    from ._mamba_ssm import MambaSSMZonosBackbone
+
+    BACKBONES["mamba_ssm"] = MambaSSMZonosBackbone
+except ImportError:
+    pass
+
+from ._torch import TorchZonosBackbone
+
+BACKBONES["torch"] = TorchZonosBackbone

Zonos-main/zonos/backbone/_mamba_ssm.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+from mamba_ssm.models.mixer_seq_simple import create_block
+from mamba_ssm.ops.triton.layer_norm import layer_norm_fn
+
+from zonos.config import BackboneConfig, InferenceParams
+
+
+class MambaSSMZonosBackbone(nn.Module):
+    supported_architectures = ["transformer", "hybrid"]
+
+    def __init__(self, config: BackboneConfig):
+        super().__init__()
+        self.config = config
+
+        self.layers = nn.ModuleList(
+            [
+                create_block(
+                    d_model=config.d_model,
+                    d_intermediate=config.d_intermediate
+                    if (i not in config.attn_layer_idx)
+                    else config.attn_mlp_d_intermediate,
+                    ssm_cfg=config.ssm_cfg,
+                    layer_idx=i,
+                    attn_layer_idx=config.attn_layer_idx,
+                    attn_cfg=config.attn_cfg,
+                    norm_epsilon=config.norm_epsilon,
+                    residual_in_fp32=config.residual_in_fp32,
+                    fused_add_norm=True,
+                    rms_norm=config.rms_norm,
+                )
+                for i in range(config.n_layer)
+            ]
+        )
+
+        self.norm_f = nn.LayerNorm(config.d_model, eps=config.norm_epsilon)
+
+    def allocate_inference_cache(self, batch_size: int, max_seqlen: int, dtype: torch.dtype = torch.bfloat16):
+        return {
+            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype)
+            for i, layer in enumerate(self.layers)
+        }
+
+    def forward(self, hidden_states: torch.Tensor, inference_params: InferenceParams | None = None):
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(hidden_states, residual, inference_params)
+
+        return layer_norm_fn(
+            hidden_states,
+            self.norm_f.weight,
+            self.norm_f.bias,
+            residual,
+            eps=self.norm_f.eps,
+            residual_in_fp32=self.config.residual_in_fp32,
+            is_rms_norm=self.config.rms_norm,
+        )

Zonos-main/zonos/backbone/_torch.py

@@ -0,0 +1,152 @@
+# Based on gpt-fast: https://github.com/pytorch-labs/gpt-fast/blob/095b2229ee3a40e379c11f05b94bd6923db63b4b/model.py
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from zonos.config import BackboneConfig, InferenceParams
+
+
+def precompute_freqs_cis(seq_len: int, n_elem: int, base: float = 10000) -> torch.Tensor:
+    freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem))
+    t = torch.arange(seq_len, device=freqs.device)
+    freqs = torch.outer(t, freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+    cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
+    return cache
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+    freqs_cis = freqs_cis.view(-1, xshaped.size(1), 1, xshaped.size(3), 2)
+    x_out2 = torch.stack(
+        [
+            xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
+            xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
+        ],
+        -1,
+    )
+
+    x_out2 = x_out2.flatten(3)
+    return x_out2.type_as(x)
+
+
+def _update_kv_cache(
+    k: torch.Tensor, v: torch.Tensor, inference_params: InferenceParams, layer_idx: int
+) -> torch.Tensor:
+    """k/v: (batch_size, seqlen, nheads, head_dim) or (batch_size, 1, nheads, head_dim)"""
+    assert layer_idx in inference_params.key_value_memory_dict
+    kv_cache, _ = inference_params.key_value_memory_dict[layer_idx]
+    # Adjust key and value for inference
+    batch_start = inference_params.batch_size_offset
+    batch_end = batch_start + k.shape[0]
+    sequence_start = inference_params.seqlen_offset
+    sequence_end = sequence_start + k.shape[1]
+    assert batch_end <= kv_cache.shape[0]
+    assert sequence_end <= kv_cache.shape[1]
+    assert kv_cache is not None
+    kv_cache[batch_start:batch_end, sequence_start:sequence_end, 0, ...] = k
+    kv_cache[batch_start:batch_end, sequence_start:sequence_end, 1, ...] = v
+    return kv_cache[batch_start:batch_end, :sequence_end, ...]
+
+
+class TorchZonosBackbone(nn.Module):
+    supported_architectures = ["transformer"]
+    freqs_cis: torch.Tensor
+
+    def __init__(self, config: BackboneConfig):
+        assert not config.ssm_cfg, "This backbone implementation only supports the Transformer model."
+        super().__init__()
+        self.config = config
+
+        self.layers = nn.ModuleList(TransformerBlock(config, i) for i in range(config.n_layer))
+        self.norm_f = nn.LayerNorm(config.d_model, eps=config.norm_epsilon)
+
+    def allocate_inference_cache(self, batch_size: int, max_seqlen: int, dtype: torch.dtype = torch.bfloat16):
+        # TODO: This function should be pure
+        head_dim = self.config.d_model // self.config.attn_cfg["num_heads"]
+        self.freqs_cis = precompute_freqs_cis(16384, head_dim)
+        return {
+            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype)
+            for i, layer in enumerate(self.layers)
+        }
+
+    def forward(self, hidden_states: torch.Tensor, inference_params: InferenceParams) -> torch.Tensor:
+        input_pos = torch.arange(0, hidden_states.shape[1], device=hidden_states.device)
+        input_pos = input_pos + inference_params.lengths_per_sample.unsqueeze(-1)
+
+        freqs_cis = self.freqs_cis[input_pos].expand(hidden_states.shape[0], -1, -1, -1)
+        for i, layer in enumerate(self.layers):
+            hidden_states = layer(hidden_states, inference_params, freqs_cis)
+        return self.norm_f(hidden_states)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: BackboneConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+
+        self.norm = nn.LayerNorm(config.d_model, eps=config.norm_epsilon)
+        self.mixer = Attention(config, layer_idx)
+        self.norm2 = nn.LayerNorm(config.d_model, eps=config.norm_epsilon)
+        self.mlp = FeedForward(config)
+
+        self.num_heads_kv = config.attn_cfg["num_heads_kv"]
+        self.head_dim = config.d_model // config.attn_cfg["num_heads"]
+
+    def allocate_inference_cache(self, batch_size: int, max_seqlen: int, dtype: torch.dtype = torch.bfloat16):
+        return torch.empty(batch_size, max_seqlen, 2, self.num_heads_kv, self.head_dim, dtype=dtype), None
+
+    def forward(self, x: torch.Tensor, inference_params: InferenceParams, freqs_cis: torch.Tensor) -> torch.Tensor:
+        x = x + self.mixer(self.norm(x), inference_params, freqs_cis)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, config: BackboneConfig, layer_idx: int):
+        super().__init__()
+        self.num_heads = config.attn_cfg["num_heads"]
+        self.num_heads_kv = config.attn_cfg["num_heads_kv"]
+        self.head_dim = config.d_model // self.num_heads
+        self.layer_idx = layer_idx
+
+        total_head_dim = (self.num_heads + 2 * self.num_heads_kv) * self.head_dim
+        self.in_proj = nn.Linear(config.d_model, total_head_dim, bias=False)
+        self.out_proj = nn.Linear(self.num_heads * self.head_dim, config.d_model, bias=False)
+
+    def forward(self, x: torch.Tensor, inference_params: InferenceParams, freqs_cis: torch.Tensor) -> torch.Tensor:
+        batch_size, seqlen, _ = x.shape
+
+        q_size = self.num_heads * self.head_dim
+        kv_size = self.num_heads_kv * self.head_dim
+        q, k, v = self.in_proj(x).split([q_size, kv_size, kv_size], dim=-1)
+
+        q = q.view(batch_size, seqlen, self.num_heads, self.head_dim)
+        k = k.view(batch_size, seqlen, self.num_heads_kv, self.head_dim)
+        v = v.view(batch_size, seqlen, self.num_heads_kv, self.head_dim)
+
+        q = apply_rotary_emb(q, freqs_cis)
+        k = apply_rotary_emb(k, freqs_cis)
+
+        kv = _update_kv_cache(k, v, inference_params, self.layer_idx)
+        k, v = kv.unbind(dim=-3)
+
+        q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=seqlen > 1, enable_gqa=True)
+
+        y = y.transpose(1, 2).contiguous().view(batch_size, seqlen, q_size)
+
+        y = self.out_proj(y)
+        return y
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config: BackboneConfig) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(config.d_model, 2 * config.attn_mlp_d_intermediate, bias=False)
+        self.fc2 = nn.Linear(config.attn_mlp_d_intermediate, config.d_model, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y, gate = self.fc1(x).chunk(2, dim=-1)
+        return self.fc2(y * F.silu(gate))

Zonos-main/zonos/codebook_pattern.py

@@ -0,0 +1,12 @@
+import torch
+import torch.nn.functional as F
+
+
+def apply_delay_pattern(codes: torch.Tensor, mask_token: int):
+    codes = F.pad(codes, (0, codes.shape[1]), value=mask_token)
+    return torch.stack([codes[:, k].roll(k + 1) for k in range(codes.shape[1])], dim=1)
+
+
+def revert_delay_pattern(codes: torch.Tensor):
+    _, n_q, seq_len = codes.shape
+    return torch.stack([codes[:, k, k + 1 : seq_len - n_q + k + 1] for k in range(n_q)], dim=1)

Zonos-main/zonos/conditioning.py

@@ -0,0 +1,405 @@
+from functools import cache
+from typing import Any, Literal, Iterable
+
+import torch
+import torch.nn as nn
+
+from zonos.config import PrefixConditionerConfig
+from zonos.utils import DEFAULT_DEVICE
+
+
+class Conditioner(nn.Module):
+    def __init__(
+        self,
+        output_dim: int,
+        name: str,
+        cond_dim: int | None = None,
+        projection: Literal["none", "linear", "mlp"] = "none",
+        uncond_type: Literal["learned", "none"] = "none",
+        **kwargs,
+    ):
+        super().__init__()
+        self.name = name
+        self.output_dim = output_dim
+        self.cond_dim = cond_dim = cond_dim or output_dim
+
+        if projection == "linear":
+            self.project = nn.Linear(cond_dim, output_dim)
+        elif projection == "mlp":
+            self.project = nn.Sequential(
+                nn.Linear(cond_dim, output_dim),
+                nn.SiLU(),
+                nn.Linear(output_dim, output_dim),
+            )
+        else:
+            self.project = nn.Identity()
+
+        self.uncond_vector = None
+        if uncond_type == "learned":
+            self.uncond_vector = nn.Parameter(torch.zeros(output_dim))
+
+    def apply_cond(self, *inputs: Any) -> torch.Tensor:
+        raise NotImplementedError()
+
+    def forward(self, inputs: tuple[Any, ...] | None) -> torch.Tensor:
+        if inputs is None:
+            assert self.uncond_vector is not None
+            return self.uncond_vector.data.view(1, 1, -1)
+
+        cond = self.apply_cond(*inputs)
+        cond = self.project(cond)
+        return cond
+
+
+# ------- ESPEAK CONTAINMENT ZONE ------------------------------------------------------------------------------------------------------------------------------------------------
+import os
+import sys
+import re
+import unicodedata
+
+import inflect
+import torch
+import torch.nn as nn
+from kanjize import number2kanji
+from phonemizer.backend import EspeakBackend
+from sudachipy import Dictionary, SplitMode
+
+if sys.platform == "darwin":
+    os.environ["PHONEMIZER_ESPEAK_LIBRARY"] = "/opt/homebrew/lib/libespeak-ng.dylib"
+
+# --- Number normalization code from https://github.com/daniilrobnikov/vits2/blob/main/text/normalize_numbers.py ---
+
+_inflect = inflect.engine()
+_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])")
+_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)")
+_pounds_re = re.compile(r"£([0-9\,]*[0-9]+)")
+_dollars_re = re.compile(r"\$([0-9\.\,]*[0-9]+)")
+_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)")
+_number_re = re.compile(r"[0-9]+")
+
+
+def _remove_commas(m: re.Match) -> str:
+    return m.group(1).replace(",", "")
+
+
+def _expand_decimal_point(m: re.Match) -> str:
+    return m.group(1).replace(".", " point ")
+
+
+def _expand_dollars(m: re.Match) -> str:
+    match = m.group(1)
+    parts = match.split(".")
+    if len(parts) > 2:
+        return match + " dollars"  # Unexpected format
+    dollars = int(parts[0]) if parts[0] else 0
+    cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if dollars and cents:
+        dollar_unit = "dollar" if dollars == 1 else "dollars"
+        cent_unit = "cent" if cents == 1 else "cents"
+        return "%s %s, %s %s" % (dollars, dollar_unit, cents, cent_unit)
+    elif dollars:
+        dollar_unit = "dollar" if dollars == 1 else "dollars"
+        return "%s %s" % (dollars, dollar_unit)
+    elif cents:
+        cent_unit = "cent" if cents == 1 else "cents"
+        return "%s %s" % (cents, cent_unit)
+    else:
+        return "zero dollars"
+
+
+def _expand_ordinal(m: re.Match) -> str:
+    return _inflect.number_to_words(m.group(0))
+
+
+def _expand_number(m: re.Match) -> str:
+    num = int(m.group(0))
+    if num > 1000 and num < 3000:
+        if num == 2000:
+            return "two thousand"
+        elif num > 2000 and num < 2010:
+            return "two thousand " + _inflect.number_to_words(num % 100)
+        elif num % 100 == 0:
+            return _inflect.number_to_words(num // 100) + " hundred"
+        else:
+            return _inflect.number_to_words(num, andword="", zero="oh", group=2).replace(", ", " ")
+    else:
+        return _inflect.number_to_words(num, andword="")
+
+
+def normalize_numbers(text: str) -> str:
+    text = re.sub(_comma_number_re, _remove_commas, text)
+    text = re.sub(_pounds_re, r"\1 pounds", text)
+    text = re.sub(_dollars_re, _expand_dollars, text)
+    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
+    text = re.sub(_ordinal_re, _expand_ordinal, text)
+    text = re.sub(_number_re, _expand_number, text)
+    return text
+
+
+# --- Number normalization code end ---
+
+
+PAD_ID, UNK_ID, BOS_ID, EOS_ID = 0, 1, 2, 3
+SPECIAL_TOKEN_IDS = [PAD_ID, UNK_ID, BOS_ID, EOS_ID]
+
+_punctuation = ';:,.!?¡¿—…"«»“”() *~-/\\&'
+_letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+_letters_ipa = (
+    "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌː��ʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+)
+
+symbols = [*_punctuation, *_letters, *_letters_ipa]
+_symbol_to_id = {s: i for i, s in enumerate(symbols, start=len(SPECIAL_TOKEN_IDS))}
+
+
+def _get_symbol_id(s: str) -> int:
+    return _symbol_to_id.get(s, 1)
+
+
+def get_symbol_ids(text: str) -> list[int]:
+    return list(map(_get_symbol_id, text))
+
+
+def tokenize_phonemes(phonemes: list[str]) -> tuple[torch.Tensor, list[int]]:
+    phoneme_ids = [[BOS_ID, *get_symbol_ids(phonemes), EOS_ID] for phonemes in phonemes]
+    lengths = list(map(len, phoneme_ids))
+    longest = max(lengths)
+    phoneme_ids = [[PAD_ID] * (longest - len(ids)) + ids for ids in phoneme_ids]
+    return torch.tensor(phoneme_ids), lengths
+
+
+def normalize_jp_text(text: str, tokenizer=Dictionary(dict="full").create()) -> str:
+    text = unicodedata.normalize("NFKC", text)
+    text = re.sub(r"\d+", lambda m: number2kanji(int(m[0])), text)
+    final_text = " ".join([x.reading_form() for x in tokenizer.tokenize(text, SplitMode.A)])
+    return final_text
+
+
+def clean(texts: list[str], languages: list[str]) -> list[str]:
+    texts_out = []
+    for text, language in zip(texts, languages):
+        if "ja" in language:
+            text = normalize_jp_text(text)
+        else:
+            text = normalize_numbers(text)
+        texts_out.append(text)
+    return texts_out
+
+
+@cache
+def get_backend(language: str) -> "EspeakBackend":
+    import logging
+
+    from phonemizer.backend import EspeakBackend
+
+    logger = logging.getLogger("phonemizer")
+    backend = EspeakBackend(
+        language,
+        preserve_punctuation=True,
+        with_stress=True,
+        punctuation_marks=_punctuation,
+        logger=logger,
+    )
+    logger.setLevel(logging.ERROR)
+    return backend
+
+
+def phonemize(texts: list[str], languages: list[str]) -> list[str]:
+    texts = clean(texts, languages)
+
+    batch_phonemes = []
+    for text, language in zip(texts, languages):
+        backend = get_backend(language)
+        phonemes = backend.phonemize([text], strip=True)
+        batch_phonemes.append(phonemes[0])
+
+    return batch_phonemes
+
+
+class EspeakPhonemeConditioner(Conditioner):
+    def __init__(self, output_dim: int, **kwargs):
+        super().__init__(output_dim, **kwargs)
+        self.phoneme_embedder = nn.Embedding(len(SPECIAL_TOKEN_IDS) + len(symbols), output_dim)
+
+    def apply_cond(self, texts: list[str], languages: list[str]) -> torch.Tensor:
+        """
+        Args:
+            texts: list of texts to convert to phonemes
+            languages: ISO 639-1 -or otherwise eSpeak compatible- language code
+        """
+        device = self.phoneme_embedder.weight.device
+
+        phonemes = phonemize(texts, languages)
+        phoneme_ids, _ = tokenize_phonemes(phonemes)
+        phoneme_embeds = self.phoneme_embedder(phoneme_ids.to(device))
+
+        return phoneme_embeds
+
+
+# ------- ESPEAK CONTAINMENT ZONE ------------------------------------------------------------------------------------------------------------------------------------------------
+
+
+class FourierConditioner(Conditioner):
+    def __init__(
+        self,
+        output_dim: int,
+        input_dim: int = 1,
+        std: float = 1.0,
+        min_val: float = 0.0,
+        max_val: float = 1.0,
+        **kwargs,
+    ):
+        assert output_dim % 2 == 0
+        super().__init__(output_dim, **kwargs)
+        self.register_buffer("weight", torch.randn([output_dim // 2, input_dim]) * std)
+        self.input_dim, self.min_val, self.max_val = input_dim, min_val, max_val
+
+    def apply_cond(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-1] == self.input_dim
+        x = (x - self.min_val) / (self.max_val - self.min_val)  # [batch_size, seq_len, input_dim]
+        f = 2 * torch.pi * x.to(self.weight.dtype) @ self.weight.T  # [batch_size, seq_len, output_dim // 2]
+        return torch.cat([f.cos(), f.sin()], dim=-1)  # [batch_size, seq_len, output_dim]
+
+
+class IntegerConditioner(Conditioner):
+    def __init__(self, output_dim: int, min_val: int = 0, max_val: int = 512, **kwargs):
+        super().__init__(output_dim, **kwargs)
+        self.min_val = min_val
+        self.max_val = max_val
+        self.int_embedder = nn.Embedding(max_val - min_val + 1, output_dim)
+
+    def apply_cond(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-1] == 1
+        return self.int_embedder(x.squeeze(-1) - self.min_val)  # [batch_size, seq_len, output_dim]
+
+
+class PassthroughConditioner(Conditioner):
+    def __init__(self, output_dim: int, **kwargs):
+        super().__init__(output_dim, **kwargs)
+
+    def apply_cond(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-1] == self.cond_dim
+        return x
+
+
+_cond_cls_map = {
+    "PassthroughConditioner": PassthroughConditioner,
+    "EspeakPhonemeConditioner": EspeakPhonemeConditioner,
+    "FourierConditioner": FourierConditioner,
+    "IntegerConditioner": IntegerConditioner,
+}
+
+
+def build_conditioners(conditioners: list[dict], output_dim: int) -> list[Conditioner]:
+    return [_cond_cls_map[config["type"]](output_dim, **config) for config in conditioners]
+
+
+class PrefixConditioner(Conditioner):
+    def __init__(self, config: PrefixConditionerConfig, output_dim: int):
+        super().__init__(output_dim, "prefix", projection=config.projection)
+        self.conditioners = nn.ModuleList(build_conditioners(config.conditioners, output_dim))
+        self.norm = nn.LayerNorm(output_dim)
+        self.required_keys = {c.name for c in self.conditioners if c.uncond_vector is None}
+
+    def forward(self, cond_dict: dict) -> torch.Tensor:
+        if not set(cond_dict).issuperset(self.required_keys):
+            raise ValueError(f"Missing required keys: {self.required_keys - set(cond_dict)}")
+        conds = []
+        for conditioner in self.conditioners:
+            conds.append(conditioner(cond_dict.get(conditioner.name)))
+        max_bsz = max(map(len, conds))
+        assert all(c.shape[0] in (max_bsz, 1) for c in conds)
+        conds = [c.expand(max_bsz, -1, -1) for c in conds]
+        return self.norm(self.project(torch.cat(conds, dim=-2)))
+
+
+supported_language_codes = [
+    'af', 'am', 'an', 'ar', 'as', 'az', 'ba', 'bg', 'bn', 'bpy', 'bs', 'ca', 'cmn',
+    'cs', 'cy', 'da', 'de', 'el', 'en-029', 'en-gb', 'en-gb-scotland', 'en-gb-x-gbclan',
+    'en-gb-x-gbcwmd', 'en-gb-x-rp', 'en-us', 'eo', 'es', 'es-419', 'et', 'eu', 'fa',
+    'fa-latn', 'fi', 'fr-be', 'fr-ch', 'fr-fr', 'ga', 'gd', 'gn', 'grc', 'gu', 'hak',
+    'hi', 'hr', 'ht', 'hu', 'hy', 'hyw', 'ia', 'id', 'is', 'it', 'ja', 'jbo', 'ka',
+    'kk', 'kl', 'kn', 'ko', 'kok', 'ku', 'ky', 'la', 'lfn', 'lt', 'lv', 'mi', 'mk',
+    'ml', 'mr', 'ms', 'mt', 'my', 'nb', 'nci', 'ne', 'nl', 'om', 'or', 'pa', 'pap',
+    'pl', 'pt', 'pt-br', 'py', 'quc', 'ro', 'ru', 'ru-lv', 'sd', 'shn', 'si', 'sk',
+    'sl', 'sq', 'sr', 'sv', 'sw', 'ta', 'te', 'tn', 'tr', 'tt', 'ur', 'uz', 'vi',
+    'vi-vn-x-central', 'vi-vn-x-south', 'yue'
+]  # fmt: off
+
+
+def make_cond_dict(
+    text: str = "It would be nice to have time for testing, indeed.",
+    language: str = "en-us",
+    speaker: torch.Tensor | None = None,
+    
+    # Emotion vector from 0.0 to 1.0
+    #   Is entangled with pitch_std because more emotion => more pitch variation
+    #                     VQScore and DNSMOS because they favor neutral speech
+    #
+    #                       Happiness, Sadness, Disgust, Fear, Surprise, Anger, Other, Neutral
+    emotion: list[float] = [0.3077, 0.0256, 0.0256, 0.0256, 0.0256, 0.0256, 0.2564, 0.3077],
+
+    # Maximum frequency (0 to 24000), should be 22050 or 24000 for 44.1 or 48 kHz audio
+    # For voice cloning use 22050
+    fmax: float = 22050.0,
+    
+    # Standard deviation for pitch (0 to 400), should be 
+    #   20-45 for normal speech, 
+    #   60-150 for expressive speech, 
+    #   higher values => crazier samples
+    pitch_std: float = 20.0,
+
+    # Speaking rate in phonemes per minute (0 to 40). 30 is very fast, 10 is slow.
+    speaking_rate: float = 15.0,
+
+    # Target VoiceQualityScore for the generated speech (0.5 to 0.8).
+    #   A list of values must be provided which represent each 1/8th of the audio.
+    #   You should unset for expressive speech.
+    # According to discord Chat this is only used for the hybrid model
+    vqscore_8: list[float] = [0.78] * 8,
+
+    # CTC target loss
+    # Only used for the hybrid model
+    ctc_loss: float = 0.0,
+    # Only used for the hybrid model
+    dnsmos_ovrl: float = 4.0,
+    # Only used for the hybrid model
+    speaker_noised: bool = False,
+    unconditional_keys: Iterable[str] = {"vqscore_8", "dnsmos_ovrl"},
+    device: torch.device | str = DEFAULT_DEVICE,
+) -> dict:
+    """
+    A helper to build the 'cond_dict' that the model expects.
+    By default, it will generate a random speaker embedding
+    """
+    assert language.lower() in supported_language_codes, "Please pick a supported language"
+
+    language_code_to_id = {lang: i for i, lang in enumerate(supported_language_codes)}
+
+    cond_dict = {
+        "espeak": ([text], [language]),
+        "speaker": speaker,
+        "emotion": emotion,
+        "fmax": fmax,
+        "pitch_std": pitch_std,
+        "speaking_rate": speaking_rate,
+        "language_id": language_code_to_id[language],
+        "vqscore_8": vqscore_8,
+        "ctc_loss": ctc_loss,
+        "dnsmos_ovrl": dnsmos_ovrl,
+        "speaker_noised": int(speaker_noised),
+    }
+
+    for k in unconditional_keys:
+        cond_dict.pop(k, None)
+
+    for k, v in cond_dict.items():
+        if isinstance(v, (float, int, list)):
+            v = torch.tensor(v)
+        if isinstance(v, torch.Tensor):
+            cond_dict[k] = v.view(1, 1, -1).to(device)
+
+        if k == "emotion":
+            cond_dict[k] /= cond_dict[k].sum(dim=-1)
+
+    return cond_dict

Zonos-main/zonos/config.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass, field
+from typing import Literal
+
+import torch
+
+
+# https://github.com/state-spaces/mamba/blob//mamba_ssm/utils/generation.py#L18
+@dataclass
+class InferenceParams:
+    """Inference parameters that are passed to the main model in order
+    to efficienly calculate and store the context during inference."""
+
+    max_seqlen: int
+    max_batch_size: int
+    seqlen_offset: int = 0
+    batch_size_offset: int = 0
+    key_value_memory_dict: dict = field(default_factory=dict)
+    lengths_per_sample: torch.Tensor | None = None
+
+    def reset(self, max_seqlen, max_batch_size):
+        self.max_seqlen = max_seqlen
+        self.max_batch_size = max_batch_size
+        self.seqlen_offset = 0
+        if self.lengths_per_sample is not None:
+            self.lengths_per_sample.zero_()
+
+
+@dataclass
+class BackboneConfig:
+    d_model: int = 1024
+    d_intermediate: int = 0
+    attn_mlp_d_intermediate: int = 0
+    n_layer: int = 16
+    ssm_cfg: dict = field(default_factory=dict)
+    attn_layer_idx: list = field(default_factory=list)
+    attn_cfg: dict = field(default_factory=dict)
+    rms_norm: bool = False
+    residual_in_fp32: bool = False
+    norm_epsilon: float = 1e-5
+
+
+@dataclass
+class PrefixConditionerConfig:
+    conditioners: list[dict]
+    projection: Literal["none", "linear", "mlp"]
+
+
+@dataclass
+class ZonosConfig:
+    backbone: BackboneConfig
+    prefix_conditioner: PrefixConditionerConfig
+    eos_token_id: int = 1024
+    masked_token_id: int = 1025
+    pad_vocab_to_multiple_of: int = 8
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "ZonosConfig":
+        d = d.copy()
+        backbone_config = BackboneConfig(**d.pop("backbone"))
+        prefix_conditioner_config = PrefixConditionerConfig(**d.pop("prefix_conditioner"))
+        config = cls(backbone_config, prefix_conditioner_config, **d)
+        return config

Zonos-main/zonos/model.py

@@ -0,0 +1,315 @@
+import json
+from typing import Callable
+
+import safetensors
+import torch
+import torch.nn as nn
+from huggingface_hub import hf_hub_download
+from tqdm import tqdm
+
+from zonos.autoencoder import DACAutoencoder
+from zonos.backbone import BACKBONES
+from zonos.codebook_pattern import apply_delay_pattern, revert_delay_pattern
+from zonos.conditioning import PrefixConditioner
+from zonos.config import InferenceParams, ZonosConfig
+from zonos.sampling import sample_from_logits
+from zonos.speaker_cloning import SpeakerEmbeddingLDA
+from zonos.utils import DEFAULT_DEVICE, find_multiple, pad_weight_
+
+DEFAULT_BACKBONE_CLS = next(iter(BACKBONES.values()))
+
+
+class Zonos(nn.Module):
+    def __init__(self, config: ZonosConfig, backbone_cls=DEFAULT_BACKBONE_CLS):
+        super().__init__()
+        self.config = config
+        dim = config.backbone.d_model
+        self.eos_token_id = config.eos_token_id
+        self.masked_token_id = config.masked_token_id
+
+        self.autoencoder = DACAutoencoder()
+        self.backbone = backbone_cls(config.backbone)
+        self.prefix_conditioner = PrefixConditioner(config.prefix_conditioner, dim)
+        self.spk_clone_model = None
+
+        # TODO: pad to multiple of at least 8
+        self.embeddings = nn.ModuleList([nn.Embedding(1026, dim) for _ in range(self.autoencoder.num_codebooks)])
+        self.heads = nn.ModuleList([nn.Linear(dim, 1025, bias=False) for _ in range(self.autoencoder.num_codebooks)])
+
+        self._cg_graph = None
+        self._cg_batch_size = None
+        self._cg_input_ids = None
+        self._cg_logits = None
+        self._cg_inference_params = None
+        self._cg_scale = None
+
+        if config.pad_vocab_to_multiple_of:
+            self.register_load_state_dict_post_hook(self._pad_embeddings_and_heads)
+
+    def _pad_embeddings_and_heads(self, *args, **kwargs):
+        for w in [*self.embeddings, *self.heads]:
+            pad_weight_(w, self.config.pad_vocab_to_multiple_of)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @classmethod
+    def from_pretrained(
+        cls, repo_id: str, revision: str | None = None, device: str = DEFAULT_DEVICE, **kwargs
+    ) -> "Zonos":
+        config_path = hf_hub_download(repo_id=repo_id, filename="config.json", revision=revision)
+        model_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", revision=revision)
+        return cls.from_local(config_path, model_path, device, **kwargs)
+
+    @classmethod
+    def from_local(
+        cls, config_path: str, model_path: str, device: str = DEFAULT_DEVICE, backbone: str | None = None
+    ) -> "Zonos":
+        config = ZonosConfig.from_dict(json.load(open(config_path)))
+        if backbone:
+            backbone_cls = BACKBONES[backbone]
+        else:
+            is_transformer = not bool(config.backbone.ssm_cfg)
+            backbone_cls = DEFAULT_BACKBONE_CLS
+            # Preferentially route to pure torch backbone for increased performance and lower latency.
+            if is_transformer and "torch" in BACKBONES:
+                backbone_cls = BACKBONES["torch"]
+
+        model = cls(config, backbone_cls).to(device, torch.bfloat16)
+        model.autoencoder.dac.to(device)
+
+        sd = model.state_dict()
+        with safetensors.safe_open(model_path, framework="pt") as f:
+            for k in f.keys():
+                sd[k] = f.get_tensor(k)
+        model.load_state_dict(sd)
+
+        return model
+
+    def make_speaker_embedding(self, wav: torch.Tensor, sr: int) -> torch.Tensor:
+        """Generate a speaker embedding from an audio clip."""
+        if self.spk_clone_model is None:
+            self.spk_clone_model = SpeakerEmbeddingLDA()
+        _, spk_embedding = self.spk_clone_model(wav.to(self.spk_clone_model.device), sr)
+        return spk_embedding.unsqueeze(0).bfloat16()
+
+    def embed_codes(self, codes: torch.Tensor) -> torch.Tensor:
+        return sum(emb(codes[:, i]) for i, emb in enumerate(self.embeddings))
+
+    def apply_heads(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return torch.stack([head(hidden_states) for head in self.heads], dim=1)
+
+    def _compute_logits(
+        self, hidden_states: torch.Tensor, inference_params: InferenceParams, cfg_scale: float
+    ) -> torch.Tensor:
+        """
+        Pass `hidden_states` into `backbone` and `multi_head`, applying
+        classifier-free guidance if `cfg_scale != 1.0`.
+        """
+        last_hidden_states = self.backbone(hidden_states, inference_params)[:, -1, :].unsqueeze(1)
+        logits = self.apply_heads(last_hidden_states).squeeze(2).float()
+        if cfg_scale != 1.0:
+            cond_logits, uncond_logits = logits.chunk(2)
+            logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale
+        logits[..., 1025:].fill_(-torch.inf)  # ensures padding is ignored
+        return logits
+
+    def _decode_one_token(
+        self,
+        input_ids: torch.Tensor,
+        inference_params: InferenceParams,
+        cfg_scale: float,
+        allow_cudagraphs: bool = True,
+    ) -> torch.Tensor:
+        """
+        Single-step decode. Prepares the hidden states, possibly replicates them
+        for CFG, and then delegates to `_compute_logits`.
+
+        Below we wrap this function with a simple CUDA Graph capturing mechanism,
+        doing 3 warmup steps if needed and then capturing or replaying the graph.
+        We only recapture if the batch size changes.
+        """
+        # TODO: support cfg_scale==1
+        if cfg_scale == 1.0:
+            hidden_states = self.embed_codes(input_ids)
+            return self._compute_logits(hidden_states, inference_params, cfg_scale)
+
+        bsz = input_ids.size(0)
+
+        if not allow_cudagraphs or input_ids.device.type != "cuda":
+            hidden_states_local = self.embed_codes(input_ids)
+            hidden_states_local = hidden_states_local.repeat(2, 1, 1)
+            return self._compute_logits(hidden_states_local, inference_params, cfg_scale)
+
+        need_capture = (self._cg_graph is None) or (self._cg_batch_size != bsz)
+
+        if need_capture:
+            self._cg_graph = None
+
+            self._cg_batch_size = bsz
+            self._cg_inference_params = inference_params
+            self._cg_scale = cfg_scale
+
+            for _ in range(3):
+                hidden_states = self.embed_codes(input_ids)
+                hidden_states = hidden_states.repeat(2, 1, 1)  # because cfg != 1.0
+                logits = self._compute_logits(hidden_states, inference_params, cfg_scale)
+
+            self._cg_input_ids = input_ids.clone()
+            self._cg_logits = torch.empty_like(logits)
+
+            g = torch.cuda.CUDAGraph()
+
+            def capture_region():
+                hidden_states_local = self.embed_codes(self._cg_input_ids)
+                hidden_states_local = hidden_states_local.repeat(2, 1, 1)
+                self._cg_logits = self._compute_logits(hidden_states_local, self._cg_inference_params, self._cg_scale)
+
+            with torch.cuda.graph(g):
+                capture_region()
+
+            self._cg_graph = g
+
+        else:
+            self._cg_input_ids.copy_(input_ids)
+
+        self._cg_graph.replay()
+
+        return self._cg_logits
+
+    def _prefill(
+        self,
+        prefix_hidden_states: torch.Tensor,
+        input_ids: torch.Tensor,
+        inference_params: InferenceParams,
+        cfg_scale: float,
+    ) -> torch.Tensor:
+        """
+        "Prefill" mode: we already have `prefix_hidden_states`, and we want
+        to append new embeddings, then compute the logits.
+        """
+        # Replicate input_ids if CFG is enabled
+        if cfg_scale != 1.0:
+            input_ids = input_ids.expand(prefix_hidden_states.shape[0], -1, -1)
+        hidden_states = torch.cat([prefix_hidden_states, self.embed_codes(input_ids)], dim=1)
+        return self._compute_logits(hidden_states, inference_params, cfg_scale)
+
+    def setup_cache(self, batch_size: int, max_seqlen: int, dtype: torch.dtype = torch.bfloat16) -> InferenceParams:
+        max_seqlen = find_multiple(max_seqlen, 8)
+        key_value_memory_dict = self.backbone.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype)
+        lengths_per_sample = torch.full((batch_size,), 0, dtype=torch.int32)
+        return InferenceParams(max_seqlen, batch_size, 0, 0, key_value_memory_dict, lengths_per_sample)
+
+    def prepare_conditioning(self, cond_dict: dict, uncond_dict: dict | None = None) -> torch.Tensor:
+        if uncond_dict is None:
+            uncond_dict = {k: cond_dict[k] for k in self.prefix_conditioner.required_keys}
+        return torch.cat(
+            [
+                self.prefix_conditioner(cond_dict),
+                self.prefix_conditioner(uncond_dict),
+            ]
+        )
+
+    def can_use_cudagraphs(self) -> bool:
+        # Only the mamba-ssm backbone supports CUDA Graphs at the moment
+        return self.device.type == "cuda" and "_mamba_ssm" in str(self.backbone.__class__)
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        prefix_conditioning: torch.Tensor,  # [bsz, cond_seq_len, d_model]
+        audio_prefix_codes: torch.Tensor | None = None,  # [bsz, 9, prefix_audio_seq_len]
+        max_new_tokens: int = 86 * 30,
+        cfg_scale: float = 2.0,
+        batch_size: int = 1,
+        sampling_params: dict = dict(min_p=0.1),
+        progress_bar: bool = True,
+        disable_torch_compile: bool = False,
+        callback: Callable[[torch.Tensor, int, int], bool] | None = None,
+    ):
+        assert cfg_scale != 1, "TODO: add support for cfg_scale=1"
+        prefix_audio_len = 0 if audio_prefix_codes is None else audio_prefix_codes.shape[2]
+        device = self.device
+
+        # Use CUDA Graphs if supported, and torch.compile otherwise.
+        cg = self.can_use_cudagraphs()
+        decode_one_token = self._decode_one_token
+        decode_one_token = torch.compile(decode_one_token, dynamic=True, disable=cg or disable_torch_compile)
+
+        unknown_token = -1
+        audio_seq_len = prefix_audio_len + max_new_tokens
+        seq_len = prefix_conditioning.shape[1] + audio_seq_len + 9
+
+        with torch.device(device):
+            inference_params = self.setup_cache(batch_size=batch_size * 2, max_seqlen=seq_len)
+            codes = torch.full((batch_size, 9, audio_seq_len), unknown_token)
+
+        if audio_prefix_codes is not None:
+            codes[..., :prefix_audio_len] = audio_prefix_codes
+
+        delayed_codes = apply_delay_pattern(codes, self.masked_token_id)
+
+        delayed_prefix_audio_codes = delayed_codes[..., : prefix_audio_len + 1]
+
+        logits = self._prefill(prefix_conditioning, delayed_prefix_audio_codes, inference_params, cfg_scale)
+        next_token = sample_from_logits(logits, **sampling_params)
+
+        offset = delayed_prefix_audio_codes.shape[2]
+        frame = delayed_codes[..., offset : offset + 1]
+        frame.masked_scatter_(frame == unknown_token, next_token)
+
+        prefix_length = prefix_conditioning.shape[1] + prefix_audio_len + 1
+        inference_params.seqlen_offset += prefix_length
+        inference_params.lengths_per_sample[:] += prefix_length
+
+        logit_bias = torch.zeros_like(logits)
+        logit_bias[:, 1:, self.eos_token_id] = -torch.inf  # only allow codebook 0 to predict EOS
+
+        stopping = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        max_steps = delayed_codes.shape[2] - offset
+        remaining_steps = torch.full((batch_size,), max_steps, device=device)
+        progress = tqdm(total=max_steps, desc="Generating", disable=not progress_bar)
+        cfg_scale = torch.tensor(cfg_scale)
+
+        step = 0
+        while torch.max(remaining_steps) > 0:
+            offset += 1
+            input_ids = delayed_codes[..., offset - 1 : offset]
+            logits = decode_one_token(input_ids, inference_params, cfg_scale, allow_cudagraphs=cg)
+            logits += logit_bias
+
+            next_token = sample_from_logits(logits, generated_tokens=delayed_codes[..., :offset], **sampling_params)
+            eos_in_cb0 = next_token[:, 0] == self.eos_token_id
+
+            remaining_steps[eos_in_cb0[:, 0]] = torch.minimum(remaining_steps[eos_in_cb0[:, 0]], torch.tensor(9))
+            stopping |= eos_in_cb0[:, 0]
+
+            eos_codebook_idx = 9 - remaining_steps
+            eos_codebook_idx = torch.clamp(eos_codebook_idx, max=9 - 1)
+            for i in range(next_token.shape[0]):
+                if stopping[i]:
+                    idx = eos_codebook_idx[i].item()
+                    next_token[i, :idx] = self.masked_token_id
+                    next_token[i, idx] = self.eos_token_id
+
+            frame = delayed_codes[..., offset : offset + 1]
+            frame.masked_scatter_(frame == unknown_token, next_token)
+            inference_params.seqlen_offset += 1
+            inference_params.lengths_per_sample[:] += 1
+
+            remaining_steps -= 1
+
+            progress.update()
+            step += 1
+
+            if callback is not None and not callback(frame, step, max_steps):
+                break
+
+        out_codes = revert_delay_pattern(delayed_codes)
+        out_codes.masked_fill_(out_codes >= 1024, 0)
+        out_codes = out_codes[..., : offset - 9]
+
+        self._cg_graph = None  # reset cuda graph to avoid cache changes
+
+        return out_codes

Zonos-main/zonos/sampling.py

@@ -0,0 +1,182 @@
+import torch
+
+
+def multinomial(input: torch.Tensor, num_samples: int, replacement=False, *, generator=None):
+    """torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension.
+
+    Args:
+        input (torch.Tensor): The input tensor containing probabilities.
+        num_samples (int): Number of samples to draw.
+        replacement (bool): Whether to draw with replacement or not.
+    Keywords args:
+        generator (torch.Generator): A pseudorandom number generator for sampling.
+    Returns:
+        torch.Tensor: Last dimension contains num_samples indices
+            sampled from the multinomial probability distribution
+            located in the last dimension of tensor input.
+    """
+
+    if num_samples == 1:
+        q = torch.empty_like(input).exponential_(1, generator=generator)
+        return torch.argmax(input / q, dim=-1, keepdim=True).to(torch.int64)
+
+    input_ = input.reshape(-1, input.shape[-1])
+    output_ = torch.multinomial(input_, num_samples=num_samples, replacement=replacement, generator=generator)
+    output = output_.reshape(*list(input.shape[:-1]), -1)
+    return output
+
+
+def apply_unified(probs: torch.Tensor, linear: float, conf: float, quad: float) -> torch.Tensor:
+    """Sample next token using unified sampling approach that combines linear scaling, confidence, and quadratic terms.
+    
+    Args:
+        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
+        linear (float): Linear scaling factor applied to log probabilities.
+        conf (float): Confidence factor that scales the entropy term.
+        quad (float): Quadratic penalty factor applied to squared log probabilities.
+    Returns:
+        torch.Tensor: Modified probability distribution after applying unified sampling.
+    """
+    logprobs = torch.log(probs.clamp_min(1e-20))
+    entropy = -torch.sum(probs * logprobs, dim=-1, keepdim=True)
+    raw = logprobs * (linear + entropy * conf) - logprobs**2 * quad
+    return raw.softmax(dim=-1)
+
+def apply_top_k(
+    probs: torch.Tensor,
+    k: int,
+) -> torch.Tensor:
+    """Sample next token from top K values along the last dimension of the input probs tensor.
+
+    Args:
+        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
+        k (int): The k in “top-k”.
+    Returns:
+        torch.Tensor: Sampled tokens.
+    """
+    v, _ = torch.topk(probs, min(k, probs.size(-1)))
+    pivot = v.select(-1, -1).unsqueeze(-1)
+    probs = torch.where(probs < pivot, 0.0, probs)
+    probs.div_(probs.sum(dim=-1, keepdim=True))
+    return probs
+
+
+def apply_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
+    """Sample next token from top P probabilities along the last dimension of the input probs tensor.
+
+    Args:
+        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
+        p (int): The p in “top-p”.
+    Returns:
+        torch.Tensor: Sampled tokens.
+    """
+    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    mask = probs_sum - probs_sort > p
+    probs_sort *= (~mask).float()
+    probs = probs.scatter(-1, probs_idx, probs_sort)
+    probs.div_(probs.sum(dim=-1, keepdim=True))
+    return probs
+
+
+def apply_min_p(probs: torch.Tensor, min_p: float) -> torch.Tensor:
+    """Sample next token using min-p sampling.
+
+    Args:
+        scores (torch.FloatTensor): Input logits with token candidates on the last dimension.
+        min_p (float): Minimum token probability, scaled by the probability of the most likely token.
+                       Must be between 0 and 1. Typical values are in the 0.01-0.2 range.
+    Returns:
+        torch.Tensor: Sampled tokens.
+    """
+    top_probs, _ = probs.max(dim=-1, keepdim=True)
+    tokens_to_remove = probs < (min_p * top_probs)
+    probs = probs.masked_fill(tokens_to_remove, 0.0)
+    probs.div_(probs.sum(dim=-1, keepdim=True))
+    return probs
+
+
+def modify_logit_for_repetition_penalty(
+    logits: torch.Tensor,
+    generated_tokens: torch.Tensor,
+    repetition_penalty: float,
+    repetition_penalty_window: int,
+):
+    """See https://arxiv.org/abs/1909.05858
+    Apply repetition penalty over a sliding window of the last `repetition_penalty_window` tokens.
+    logits: (batch_size, n_codebooks, vocab_size)
+    generated_tokens: (batch_size, n_codebooks, seq_len)
+    """
+    generated_tokens = generated_tokens[..., -repetition_penalty_window:]
+    generated_tokens = generated_tokens.clamp_max(logits.shape[-1] - 1).to(torch.int64)
+    rp = torch.full_like(logits, repetition_penalty)
+    factors = torch.ones_like(logits).scatter_reduce(2, generated_tokens, rp, reduce="prod")
+    return torch.where(logits <= 0, logits * factors, logits / factors)
+
+
+def sample_from_logits(
+    logits: torch.Tensor,
+    temperature: float = 1.0,
+    top_p: float = 0.0,
+    top_k: int = 0,
+    min_p: float = 0.0,
+    linear: float = 0.0,
+    conf: float = 0.0,
+    quad: float = 0.0,
+    generated_tokens: torch.Tensor | None = None,
+    repetition_penalty: float = 3.0,
+    repetition_penalty_window: int = 2,
+) -> torch.Tensor:
+    """Sample next token from logits using either top_k/p/min_p OR using NovelAI's Unified Sampler.
+    
+    Args:
+        logits (torch.Tensor): Input logits with token candidates on the last dimension.
+
+        temperature (float): Randomness of the sampling. Lower temperature results in more deterministic samples.
+            To disable sampling entirely, set it to 0. For NovelAI's Unified Sampler, set it to 1.0
+
+        top_p (float): Only sample from the most probable tokens whose cumulative probability is less than p.
+            This is called nucleus sampling. Must be between 0 and 1. Typical values are in the 0.1-0.9 range.
+
+            Set to 0 to disable.
+
+        top_k (int): Only sample from the top k most probable tokens. Set to 0 to disable.
+        
+        min_p (float): Minimum token probability, scaled by the probability of the most likely token.
+                       Must be between 0 and 1. Typical values are in the 0.01-0.2 range.
+                       If too high, no token might be sampled leading to silence (?)
+
+        linear (float): NovelAI's Unified Sampler -> 0.0 to 1.0, default from gradio 0.5
+
+            Set Linear between 0 and 1 according to how unusual you want tokens to be. 
+            Lower numbers will produce more unusual/creative outputs, 
+            but you will have to reroll or edit more. 
+
+        conf (float): Confidence - Low values make random outputs more random. -> -2.0 * Quad to 2.0, default from gradio 0.4
+
+            As a starting point, set Quad = 1/3 - Linear * 4 / 15, and Conf = -Quad / 2.
+
+        quad (float): Quadratic - High values make low probablities much lower. -> -2.0 to 2.0, default from gradio 0.0
+
+    Returns:
+        torch.Tensor: Sampled tokens.
+    """
+    if repetition_penalty != 1.0 and generated_tokens is not None:
+        logits = modify_logit_for_repetition_penalty(logits, generated_tokens, repetition_penalty, repetition_penalty_window)
+
+    if temperature > 0:
+        probs = torch.softmax(logits / temperature, dim=-1)
+        if linear > 0.0:
+            probs = apply_unified(probs, linear, conf, quad)
+        if top_p > 0:
+            probs = apply_top_p(probs, top_p)
+        if top_k > 0:
+            probs = apply_top_k(probs, top_k)
+        if min_p > 0:
+            probs = apply_min_p(probs, min_p)
+
+        next_token = multinomial(probs, num_samples=1)
+    else:
+        next_token = torch.argmax(logits, dim=-1, keepdim=True)
+
+    return next_token  # [batch_size, num_codebooks, 1]

Zonos-main/zonos/speaker_cloning.py

@@ -0,0 +1,412 @@
+import math
+from functools import cache
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from huggingface_hub import hf_hub_download
+
+from zonos.utils import DEFAULT_DEVICE
+
+
+class logFbankCal(nn.Module):
+    def __init__(
+        self,
+        sample_rate: int = 16_000,
+        n_fft: int = 512,
+        win_length: float = 0.025,
+        hop_length: float = 0.01,
+        n_mels: int = 80,
+    ):
+        super().__init__()
+        self.fbankCal = torchaudio.transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=n_fft,
+            win_length=int(win_length * sample_rate),
+            hop_length=int(hop_length * sample_rate),
+            n_mels=n_mels,
+        )
+
+    def forward(self, x):
+        out = self.fbankCal(x)
+        out = torch.log(out + 1e-6)
+        out = out - out.mean(axis=2).unsqueeze(dim=2)
+        return out
+
+
+class ASP(nn.Module):
+    # Attentive statistics pooling
+    def __init__(self, in_planes, acoustic_dim):
+        super(ASP, self).__init__()
+        outmap_size = int(acoustic_dim / 8)
+        self.out_dim = in_planes * 8 * outmap_size * 2
+
+        self.attention = nn.Sequential(
+            nn.Conv1d(in_planes * 8 * outmap_size, 128, kernel_size=1),
+            nn.ReLU(),
+            nn.BatchNorm1d(128),
+            nn.Conv1d(128, in_planes * 8 * outmap_size, kernel_size=1),
+            nn.Softmax(dim=2),
+        )
+
+    def forward(self, x):
+        x = x.reshape(x.size()[0], -1, x.size()[-1])
+        w = self.attention(x)
+        mu = torch.sum(x * w, dim=2)
+        sg = torch.sqrt((torch.sum((x**2) * w, dim=2) - mu**2).clamp(min=1e-5))
+        x = torch.cat((mu, sg), 1)
+
+        x = x.view(x.size()[0], -1)
+        return x
+
+
+class SimAMBasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, ConvLayer, NormLayer, in_planes, planes, stride=1, block_id=1):
+        super(SimAMBasicBlock, self).__init__()
+        self.conv1 = ConvLayer(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = NormLayer(planes)
+        self.conv2 = ConvLayer(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = NormLayer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.sigmoid = nn.Sigmoid()
+
+        self.downsample = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.downsample = nn.Sequential(
+                ConvLayer(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
+                NormLayer(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out = self.SimAM(out)
+        out += self.downsample(x)
+        out = self.relu(out)
+        return out
+
+    def SimAM(self, X, lambda_p=1e-4):
+        n = X.shape[2] * X.shape[3] - 1
+        d = (X - X.mean(dim=[2, 3], keepdim=True)).pow(2)
+        v = d.sum(dim=[2, 3], keepdim=True) / n
+        E_inv = d / (4 * (v + lambda_p)) + 0.5
+        return X * self.sigmoid(E_inv)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, ConvLayer, NormLayer, in_planes, planes, stride=1, block_id=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = ConvLayer(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = NormLayer(planes)
+        self.conv2 = ConvLayer(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = NormLayer(planes)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.downsample = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.downsample = nn.Sequential(
+                ConvLayer(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
+                NormLayer(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.downsample(x)
+        out = self.relu(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, ConvLayer, NormLayer, in_planes, planes, stride=1, block_id=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, in_planes, block, num_blocks, in_ch=1, feat_dim="2d", **kwargs):
+        super(ResNet, self).__init__()
+        if feat_dim == "1d":
+            self.NormLayer = nn.BatchNorm1d
+            self.ConvLayer = nn.Conv1d
+        elif feat_dim == "2d":
+            self.NormLayer = nn.BatchNorm2d
+            self.ConvLayer = nn.Conv2d
+        elif feat_dim == "3d":
+            self.NormLayer = nn.BatchNorm3d
+            self.ConvLayer = nn.Conv3d
+        else:
+            print("error")
+
+        self.in_planes = in_planes
+
+        self.conv1 = self.ConvLayer(in_ch, in_planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = self.NormLayer(in_planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(block, in_planes, num_blocks[0], stride=1, block_id=1)
+        self.layer2 = self._make_layer(block, in_planes * 2, num_blocks[1], stride=2, block_id=2)
+        self.layer3 = self._make_layer(block, in_planes * 4, num_blocks[2], stride=2, block_id=3)
+        self.layer4 = self._make_layer(block, in_planes * 8, num_blocks[3], stride=2, block_id=4)
+
+    def _make_layer(self, block, planes, num_blocks, stride, block_id=1):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.ConvLayer, self.NormLayer, self.in_planes, planes, stride, block_id))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x
+
+
+def ResNet293(in_planes: int, **kwargs):
+    return ResNet(in_planes, SimAMBasicBlock, [10, 20, 64, 3], **kwargs)
+
+
+class ResNet293_based(nn.Module):
+    def __init__(
+        self,
+        in_planes: int = 64,
+        embd_dim: int = 256,
+        acoustic_dim: int = 80,
+        featCal=None,
+        dropout: float = 0,
+        **kwargs,
+    ):
+        super(ResNet293_based, self).__init__()
+        self.featCal = featCal
+        self.front = ResNet293(in_planes)
+        block_expansion = SimAMBasicBlock.expansion
+        self.pooling = ASP(in_planes * block_expansion, acoustic_dim)
+        self.bottleneck = nn.Linear(self.pooling.out_dim, embd_dim)
+        self.drop = nn.Dropout(dropout) if dropout else None
+
+    def forward(self, x):
+        x = self.featCal(x)
+        x = self.front(x.unsqueeze(dim=1))
+        x = self.pooling(x)
+        if self.drop:
+            x = self.drop(x)
+        x = self.bottleneck(x)
+        return x
+
+
+class SEModule(nn.Module):
+    def __init__(self, channels, bottleneck=128):
+        super(SEModule, self).__init__()
+        self.se = nn.Sequential(
+            nn.AdaptiveAvgPool1d(1),
+            nn.Conv1d(channels, bottleneck, kernel_size=1, padding=0),
+            nn.ReLU(),
+            # nn.BatchNorm1d(bottleneck), # Removed
+            nn.Conv1d(bottleneck, channels, kernel_size=1, padding=0),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, input):
+        x = self.se(input)
+        return input * x
+
+
+class Bottle2neck(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size=None, dilation=None, scale=8):
+        super(Bottle2neck, self).__init__()
+        width = int(math.floor(planes / scale))
+        self.conv1 = nn.Conv1d(inplanes, width * scale, kernel_size=1)
+        self.bn1 = nn.BatchNorm1d(width * scale)
+        self.nums = scale - 1
+        convs = []
+        bns = []
+        num_pad = math.floor(kernel_size / 2) * dilation
+        for i in range(self.nums):
+            convs.append(nn.Conv1d(width, width, kernel_size=kernel_size, dilation=dilation, padding=num_pad))
+            bns.append(nn.BatchNorm1d(width))
+        self.convs = nn.ModuleList(convs)
+        self.bns = nn.ModuleList(bns)
+        self.conv3 = nn.Conv1d(width * scale, planes, kernel_size=1)
+        self.bn3 = nn.BatchNorm1d(planes)
+        self.relu = nn.ReLU()
+        self.width = width
+        self.se = SEModule(planes)
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.relu(out)
+        out = self.bn1(out)
+
+        spx = torch.split(out, self.width, 1)
+        for i in range(self.nums):
+            if i == 0:
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            sp = self.convs[i](sp)
+            sp = self.relu(sp)
+            sp = self.bns[i](sp)
+            if i == 0:
+                out = sp
+            else:
+                out = torch.cat((out, sp), 1)
+        out = torch.cat((out, spx[self.nums]), 1)
+
+        out = self.conv3(out)
+        out = self.relu(out)
+        out = self.bn3(out)
+
+        out = self.se(out)
+        out += residual
+        return out
+
+
+class ECAPA_TDNN(nn.Module):
+    def __init__(self, C, featCal):
+        super(ECAPA_TDNN, self).__init__()
+        self.featCal = featCal
+        self.conv1 = nn.Conv1d(80, C, kernel_size=5, stride=1, padding=2)
+        self.relu = nn.ReLU()
+        self.bn1 = nn.BatchNorm1d(C)
+        self.layer1 = Bottle2neck(C, C, kernel_size=3, dilation=2, scale=8)
+        self.layer2 = Bottle2neck(C, C, kernel_size=3, dilation=3, scale=8)
+        self.layer3 = Bottle2neck(C, C, kernel_size=3, dilation=4, scale=8)
+        # I fixed the shape of the output from MFA layer, that is close to the setting from ECAPA paper.
+        self.layer4 = nn.Conv1d(3 * C, 1536, kernel_size=1)
+        self.attention = nn.Sequential(
+            nn.Conv1d(4608, 256, kernel_size=1),
+            nn.ReLU(),
+            nn.BatchNorm1d(256),
+            nn.Tanh(),  # Added
+            nn.Conv1d(256, 1536, kernel_size=1),
+            nn.Softmax(dim=2),
+        )
+        self.bn5 = nn.BatchNorm1d(3072)
+        self.fc6 = nn.Linear(3072, 192)
+        self.bn6 = nn.BatchNorm1d(192)
+
+    def forward(self, x):
+        x = self.featCal(x)
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.bn1(x)
+
+        x1 = self.layer1(x)
+        x2 = self.layer2(x + x1)
+        x3 = self.layer3(x + x1 + x2)
+
+        x = self.layer4(torch.cat((x1, x2, x3), dim=1))
+        x = self.relu(x)
+
+        t = x.size()[-1]
+
+        global_x = torch.cat(
+            (
+                x,
+                torch.mean(x, dim=2, keepdim=True).repeat(1, 1, t),
+                torch.sqrt(torch.var(x, dim=2, keepdim=True).clamp(min=1e-4)).repeat(1, 1, t),
+            ),
+            dim=1,
+        )
+
+        w = self.attention(global_x)
+
+        mu = torch.sum(x * w, dim=2)
+        sg = torch.sqrt((torch.sum((x**2) * w, dim=2) - mu**2).clamp(min=1e-4))
+
+        x = torch.cat((mu, sg), 1)
+        x = self.bn5(x)
+        x = self.fc6(x)
+        x = self.bn6(x)
+
+        return x
+
+
+class SpeakerEmbedding(nn.Module):
+    def __init__(self, ckpt_path: str = "ResNet293_SimAM_ASP_base.pt", device: str = DEFAULT_DEVICE):
+        super().__init__()
+        self.device = device
+        with torch.device(device):
+            self.model = ResNet293_based()
+            state_dict = torch.load(ckpt_path, weights_only=True, mmap=True, map_location="cpu")
+            self.model.load_state_dict(state_dict)
+            self.model.featCal = logFbankCal()
+
+        self.requires_grad_(False).eval()
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    @cache
+    def _get_resampler(self, orig_sample_rate: int):
+        return torchaudio.transforms.Resample(orig_sample_rate, 16_000).to(self.device)
+
+    def prepare_input(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:
+        assert wav.ndim < 3
+        if wav.ndim == 2:
+            wav = wav.mean(0, keepdim=True)
+        wav = self._get_resampler(sample_rate)(wav)
+        return wav
+
+    def forward(self, wav: torch.Tensor, sample_rate: int):
+        wav = self.prepare_input(wav, sample_rate).to(self.device, self.dtype)
+        return self.model(wav).to(wav.device)
+
+
+class SpeakerEmbeddingLDA(nn.Module):
+    def __init__(self, device: str = DEFAULT_DEVICE):
+        super().__init__()
+        spk_model_path = hf_hub_download(
+            repo_id="Zyphra/Zonos-v0.1-speaker-embedding",
+            filename="ResNet293_SimAM_ASP_base.pt",
+        )
+        lda_spk_model_path = hf_hub_download(
+            repo_id="Zyphra/Zonos-v0.1-speaker-embedding",
+            filename="ResNet293_SimAM_ASP_base_LDA-128.pt",
+        )
+
+        self.device = device
+        with torch.device(device):
+            self.model = SpeakerEmbedding(spk_model_path, device)
+            lda_sd = torch.load(lda_spk_model_path, weights_only=True)
+            out_features, in_features = lda_sd["weight"].shape
+            self.lda = nn.Linear(in_features, out_features, bias=True, dtype=torch.float32)
+            self.lda.load_state_dict(lda_sd)
+
+        self.requires_grad_(False).eval()
+
+    def forward(self, wav: torch.Tensor, sample_rate: int):
+        emb = self.model(wav, sample_rate).to(torch.float32)
+        return emb, self.lda(emb)

Zonos-main/zonos/utils.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def find_multiple(n: int, k: int) -> int:
+    if k == 0 or n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
+def pad_weight_(w: nn.Embedding | nn.Linear, multiple: int):
+    """Pad the weight of an embedding or linear layer to a multiple of `multiple`."""
+    if isinstance(w, nn.Embedding):
+        # Pad input dim
+        if w.weight.shape[1] % multiple == 0:
+            return
+        w.weight.data = F.pad(w.weight.data, (0, 0, 0, w.weight.shape[1] % multiple))
+        w.num_embeddings, w.embedding_dim = w.weight.shape
+    elif isinstance(w, nn.Linear):
+        # Pad output dim
+        if w.weight.shape[0] % multiple == 0:
+            return
+        w.weight.data = F.pad(w.weight.data, (0, 0, 0, w.weight.shape[0] % multiple))
+        w.out_features, w.in_features = w.weight.shape
+    else:
+        raise ValueError(f"Unsupported weight type: {type(w)}")
+
+
+def get_device() -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device(torch.cuda.current_device())
+    # MPS breaks for whatever reason. Uncomment when it's working.
+    # if torch.mps.is_available():
+    #     return torch.device("mps")
+    return torch.device("cpu")
+
+
+DEFAULT_DEVICE = get_device()