diff --git a/reference.xml b/reference.xml deleted file mode 100644 index f971ec4e..00000000 --- a/reference.xml +++ /dev/null @@ -1,17481 +0,0 @@ -This file is a merged representation of the entire codebase, combined into a single document by Repomix. - - -This section contains a summary of this file. - - -This file contains a packed representation of the entire repository's contents. -It is designed to be easily consumable by AI systems for analysis, code review, -or other automated processes. - - - -The content is organized as follows: -1. This summary section -2. Repository information -3. Directory structure -4. Repository files (if enabled) -5. Multiple file entries, each consisting of: - - File path as an attribute - - Full contents of the file - - - -- This file should be treated as read-only. Any changes should be made to the - original repository files, not this packed version. -- When processing this file, use the file path to distinguish - between different files in the repository. -- Be aware that this file may contain sensitive information. Handle it with - the same level of security as you would the original repository. - - - -- Some files may have been excluded based on .gitignore rules and Repomix's configuration -- Binary files are not included in this packed representation. Please refer to the Repository Structure section for a complete list of file paths, including binary files -- Files matching patterns in .gitignore are excluded -- Files matching default ignore patterns are excluded -- Files are sorted by Git change count (files with more changes are at the bottom) - - - - - -.github/ - workflows/ - rust-ci.yml -mimi-pyo3/ - py_src/ - rustymimi/ - __init__.py - __init__.pyi - src/ - lib.rs - Cargo.toml - pyproject.toml - stub.py -moshi-backend/ - src/ - audio.rs - benchmark.rs - build.rs - main.rs - standalone.rs - stream_both.rs - utils.rs - build.rs - Cargo.toml - config-q8.json - config.json -moshi-cli/ - src/ - audio_io.rs - gen.rs - main.rs - multistream.rs - Cargo.toml -moshi-core/ - src/ - asr.rs - batched_transformer.rs - conditioner.rs - conv.rs - kv_cache.rs - lib.rs - lm_generate_multistream.rs - lm_generate.rs - lm.rs - mimi.rs - nn.rs - quantization.rs - seanet.rs - streaming.rs - transformer.rs - tts_streaming.rs - tts.rs - wav.rs - Cargo.toml -moshi-server/ - src/ - asr.rs - batched_asr.rs - lib.rs - lm.rs - main.rs - metrics.rs - mimi.rs - protocol.rs - py_module_post.rs - py_module.rs - tts.rs - utils.rs - build.rs - Cargo.toml - pyproject.toml - tts.py - voice.py -Cargo.toml -LICENSE -protocol.md -README.md -rustfmt.toml -s2st-1b.toml - - - -This section contains the contents of the repository's files. - - -on: [push, pull_request] - -name: Continuous integration - -jobs: - check: - name: Check - defaults: - run: - working-directory: ./rust - runs-on: ${{ matrix.os }} - strategy: - matrix: - os: [ubuntu-latest, windows-latest, macOS-latest] - rust: [stable, nightly] - steps: - - uses: actions/checkout@v2 - - uses: actions-rs/toolchain@v1 - with: - profile: minimal - toolchain: ${{ matrix.rust }} - override: true - - uses: actions-rs/cargo@v1 - with: - command: check - - test: - name: Test Suite - defaults: - run: - working-directory: ./rust - runs-on: ${{ matrix.os }} - strategy: - matrix: - os: [ubuntu-latest, windows-latest, macOS-latest] - rust: [stable, nightly] - steps: - - uses: actions/checkout@v2 - - uses: actions-rs/toolchain@v1 - with: - profile: minimal - toolchain: ${{ matrix.rust }} - override: true - - uses: actions-rs/cargo@v1 - with: - command: test - - fmt: - name: Rustfmt - defaults: - run: - working-directory: ./rust - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v2 - - uses: actions-rs/toolchain@v1 - with: - profile: minimal - toolchain: stable - override: true - - run: rustup component add rustfmt - - uses: actions-rs/cargo@v1 - with: - command: fmt - args: --all -- --check - - clippy: - name: Clippy - defaults: - run: - working-directory: ./rust - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v2 - - uses: actions-rs/toolchain@v1 - with: - profile: minimal - toolchain: stable - override: true - - run: rustup component add clippy - - uses: actions-rs/cargo@v1 - with: - command: clippy - args: -- -D warnings - - - -from .rustymimi import * - - - -# Generated content DO NOT EDIT -from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Sequence -from os import PathLike - -@staticmethod -def write_wav(filename, data, sample_rate): - """ - Writes an audio file using the wav format based on pcm data from a numpy array. - - This only supports a single channel at the moment so the input array data is expected to have a - single dimension. - """ - pass - -class StreamTokenizer: - def __init__(path, *, dtype="f32", max_seq_len=None): - pass - - def decode(self, codes): - """ """ - pass - - def encode(self, pcm_data): - """ """ - pass - - def get_decoded(self): - """ """ - pass - - def get_encoded(self): - """ """ - pass - -class Tokenizer: - def __init__(path, *, dtype="f32", max_seq_len=None): - pass - - def decode(self, codes): - """ """ - pass - - def decode_step(self, codes): - """ """ - pass - - def encode(self, pcm_data): - """ """ - pass - - def encode_step(self, pcm_data): - """ """ - pass - - def reset(self): - """ """ - pass - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use pyo3::prelude::*; - -use ::moshi as mm; -use mm::{candle, candle_nn, conv, mimi, seanet, transformer}; -use std::sync::{mpsc, Mutex}; - -trait PyRes { - #[allow(unused)] - fn w(self) -> PyResult; - fn w_f>(self, p: P) -> PyResult; -} - -impl> PyRes for Result { - fn w(self) -> PyResult { - self.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.into().to_string())) - } - fn w_f>(self, p: P) -> PyResult { - self.map_err(|e| { - let e = e.into().to_string(); - let msg = format!("{:?}: {e}", p.as_ref()); - pyo3::exceptions::PyValueError::new_err(msg) - }) - } -} - -#[macro_export] -macro_rules! py_bail { - ($msg:literal $(,)?) => { - return Err(pyo3::exceptions::PyValueError::new_err(format!($msg))) - }; - ($err:expr $(,)?) => { - return Err(pyo3::exceptions::PyValueError::new_err(format!($err))) - }; - ($fmt:expr, $($arg:tt)*) => { - return Err(pyo3::exceptions::PyValueError::new_err(format!($fmt, $($arg)*))) - }; -} - -fn mimi_cfg(num_codebooks: usize, max_seq_len: Option) -> mimi::Config { - let seanet_cfg = seanet::Config { - dimension: 512, - channels: 1, - causal: true, - n_filters: 64, - n_residual_layers: 1, - activation: candle_nn::Activation::Elu(1.), - compress: 2, - dilation_base: 2, - disable_norm_outer_blocks: 0, - final_activation: None, - kernel_size: 7, - residual_kernel_size: 3, - last_kernel_size: 3, - lstm: 0, - norm: conv::Norm::WeightNorm, - pad_mode: conv::PadMode::Constant, - ratios: vec![8, 6, 5, 4], - true_skip: true, - }; - let transformer_cfg = transformer::Config { - d_model: seanet_cfg.dimension, - num_heads: 8, - num_layers: 8, - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: Some(0.01), - context: 250, - conv_kernel_size: 5, - use_conv_bias: true, - use_conv_block: false, - max_period: 10000, - positional_embedding: transformer::PositionalEmbedding::Rope, - gating: None, - norm: mm::NormType::LayerNorm, - - dim_feedforward: 2048, - kv_repeat: 1, - conv_layout: true, // see builders.py - cross_attention: None, - shared_cross_attn: true, - max_seq_len: max_seq_len.unwrap_or(8192), // the transformer works at 25hz so this is ~5 mins. - }; - mimi::Config { - channels: 1, - sample_rate: 24_000., - frame_rate: 12.5, - renormalize: true, - resample_method: mimi::ResampleMethod::Conv, - seanet: seanet_cfg, - transformer: transformer_cfg, - quantizer_n_q: num_codebooks, - quantizer_bins: 2048, - quantizer_dim: 256, - } -} - -#[pyclass] -struct Tokenizer { - mimi: mimi::Mimi, - device: candle::Device, - dtype: candle::DType, -} - -#[pymethods] -impl Tokenizer { - #[pyo3(signature = (path, *, num_codebooks=8, dtype="f32", max_seq_len=None))] - #[new] - fn new( - path: std::path::PathBuf, - num_codebooks: usize, - dtype: &str, - max_seq_len: Option, - ) -> PyResult { - let device = candle::Device::Cpu; - let dtype = match dtype { - "f32" => candle::DType::F32, - "f16" => candle::DType::F16, - "bf16" => candle::DType::BF16, - dtype => py_bail!("unsupported dtype '{dtype}'"), - }; - let vb = - unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[path], dtype, &device).w()? }; - let cfg = mimi_cfg(num_codebooks, max_seq_len); - let mimi = mimi::Mimi::new(cfg, vb).w()?; - Ok(Self { mimi, device, dtype }) - } - - fn encode(&mut self, pcm_data: numpy::PyReadonlyArray3) -> PyResult { - let py = pcm_data.py(); - let pcm_data = pcm_data.as_array(); - let pcm_shape = pcm_data.shape().to_vec(); - let pcm_data = match pcm_data.to_slice() { - None => py_bail!("input data is not contiguous"), - Some(data) => data, - }; - let codes = py - .allow_threads(|| { - let pcm_data = candle::Tensor::from_slice(pcm_data, pcm_shape, &self.device)? - .to_dtype(self.dtype)?; - let codes = self.mimi.encode(&pcm_data)?; - codes.to_vec3::() - }) - .w()?; - let codes = numpy::PyArray3::from_vec3(py, &codes)?; - Ok(codes.into_any().unbind()) - } - - fn encode_step(&mut self, pcm_data: numpy::PyReadonlyArray3) -> PyResult { - let py = pcm_data.py(); - let pcm_data = pcm_data.as_array(); - let pcm_shape = pcm_data.shape().to_vec(); - let pcm_data = match pcm_data.to_slice() { - None => py_bail!("input data is not contiguous"), - Some(data) => data, - }; - let codes = py - .allow_threads(|| { - let pcm_data = candle::Tensor::from_slice(pcm_data, pcm_shape, &self.device)? - .to_dtype(self.dtype)?; - let codes = self.mimi.encode_step(&pcm_data.into(), &().into())?; - match codes.as_option() { - Some(codes) => Ok::<_, candle::Error>(Some(codes.to_vec3::()?)), - None => Ok(None), - } - }) - .w()?; - match codes { - Some(codes) => { - let codes = numpy::PyArray3::from_vec3(py, &codes)?; - Ok(codes.into_any().unbind()) - } - None => Ok(py.None()), - } - } - - fn decode(&mut self, codes: numpy::PyReadonlyArray3, py: Python) -> PyResult { - let codes = codes.as_array(); - let codes_shape = codes.shape().to_vec(); - let codes = match codes.to_slice() { - None => py_bail!("input data is not contiguous"), - Some(data) => data, - }; - let pcm = py - .allow_threads(|| { - let codes = candle::Tensor::from_slice(codes, codes_shape, &self.device)?; - let pcm = self.mimi.decode(&codes)?.to_dtype(candle::DType::F32)?; - pcm.to_vec3::() - }) - .w()?; - let pcm = numpy::PyArray3::from_vec3(py, &pcm)?; - Ok(pcm.into_any().unbind()) - } - - fn decode_step( - &mut self, - codes: numpy::PyReadonlyArray3, - py: Python, - ) -> PyResult { - let codes = codes.as_array(); - let codes_shape = codes.shape().to_vec(); - let codes = match codes.to_slice() { - None => py_bail!("input data is not contiguous"), - Some(data) => data, - }; - let pcm = py - .allow_threads(|| { - let codes = candle::Tensor::from_slice(codes, codes_shape, &self.device)?; - let pcm = self.mimi.decode_step(&codes.into(), &().into())?; - match pcm.as_option() { - Some(pcm) => { - let pcm = pcm.to_dtype(candle::DType::F32)?; - Ok::<_, candle::Error>(Some(pcm.to_vec3::()?)) - } - None => Ok(None), - } - }) - .w()?; - match pcm { - Some(pcm) => { - let pcm = numpy::PyArray3::from_vec3(py, &pcm)?; - Ok(pcm.into_any().unbind()) - } - None => Ok(py.None()), - } - } - - fn reset(&mut self) { - self.mimi.reset_state() - } -} - -#[pyclass] -struct StreamTokenizer { - #[allow(unused)] - dtype: candle::DType, - encoder_rx: Mutex>>>, - encoder_tx: mpsc::Sender>, - decoder_rx: Mutex>>, - decoder_tx: mpsc::Sender>>, -} - -#[pymethods] -impl StreamTokenizer { - #[pyo3(signature = (path, *, num_codebooks=8, dtype="f32", max_seq_len=None))] - #[new] - fn new( - path: std::path::PathBuf, - num_codebooks: usize, - dtype: &str, - max_seq_len: Option, - ) -> PyResult { - let device = candle::Device::Cpu; - let dtype = match dtype { - "f32" => candle::DType::F32, - "f16" => candle::DType::F16, - "bf16" => candle::DType::BF16, - dtype => py_bail!("unsupported dtype '{dtype}'"), - }; - let vb = - unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[path], dtype, &device).w()? }; - let cfg = mimi_cfg(num_codebooks, max_seq_len); - let mut e_mimi = mimi::Mimi::new(cfg, vb).w()?; - let mut d_mimi = e_mimi.clone(); - let (encoder_tx, e_rx) = mpsc::channel::>(); - let (decoder_tx, d_rx) = mpsc::channel::>>(); - let (d_tx, decoder_rx) = mpsc::channel::>(); - let (e_tx, encoder_rx) = mpsc::channel::>>(); - std::thread::spawn(move || { - while let Ok(pcm_data) = e_rx.recv() { - // Can't wait for try blocks to be a thing - if let Err(err) = (|| { - let l = pcm_data.len(); - let pcm_data = - candle::Tensor::from_vec(pcm_data, (1, 1, l), &candle::Device::Cpu)? - .to_dtype(dtype)?; - let codes = e_mimi.encode_step(&pcm_data.into(), &().into())?; - if let Some(codes) = codes.as_option() { - let mut codes = codes.to_vec3::()?; - e_tx.send(codes.remove(0))?; - } - Ok::<_, anyhow::Error>(()) - })() { - eprintln!("error in encoder thread {err:?}") - } - } - }); - std::thread::spawn(move || { - while let Ok(codes) = d_rx.recv() { - if let Err(err) = (|| { - let codes = candle::Tensor::new(codes, &candle::Device::Cpu)?.unsqueeze(2)?; - let pcm_data = d_mimi.decode_step(&codes.into(), &().into())?; - if let Some(pcm_data) = pcm_data.as_option() { - let mut pcm_data = pcm_data.to_vec3::()?; - d_tx.send(pcm_data.remove(0).remove(0))?; - } - Ok::<_, anyhow::Error>(()) - })() { - eprintln!("error in decoder thread {err:?}") - } - } - }); - Ok(Self { - dtype, - encoder_rx: Mutex::new(encoder_rx), - encoder_tx, - decoder_rx: Mutex::new(decoder_rx), - decoder_tx, - }) - } - - fn encode(&mut self, pcm_data: numpy::PyReadonlyArray1) -> PyResult<()> { - self.encoder_tx.send(pcm_data.as_array().to_vec()).w()?; - Ok(()) - } - - fn decode(&mut self, codes: numpy::PyReadonlyArray2) -> PyResult<()> { - let codes = codes.as_array(); - let dims = codes.shape(); - let codes = match codes.to_slice() { - None => py_bail!("input data is not contiguous"), - Some(data) => data.to_vec(), - }; - let codes = codes.chunks_exact(dims[1]).map(|v| v.to_vec()).collect::>(); - self.decoder_tx.send(codes).w()?; - Ok(()) - } - - fn get_encoded(&mut self, py: Python) -> PyResult { - match self.encoder_rx.lock().unwrap().try_recv() { - Ok(codes) => { - let codes = numpy::PyArray2::from_vec2(py, &codes)?; - Ok(codes.into_any().unbind()) - } - Err(mpsc::TryRecvError::Disconnected) => { - py_bail!("worker thread disconnected") - } - Err(mpsc::TryRecvError::Empty) => Ok(py.None()), - } - } - - fn get_decoded(&mut self, py: Python) -> PyResult { - match self.decoder_rx.lock().unwrap().try_recv() { - Ok(pcm) => { - let pcm = numpy::PyArray1::from_vec(py, pcm); - Ok(pcm.into_any().unbind()) - } - Err(mpsc::TryRecvError::Disconnected) => { - py_bail!("worker thread disconnected") - } - Err(mpsc::TryRecvError::Empty) => Ok(py.None()), - } - } -} - -/// Writes an audio file using the wav format based on pcm data from a numpy array. -/// -/// This only supports a single channel at the moment so the input array data is expected to have a -/// single dimension. -#[pyfunction] -#[pyo3(signature = (filename, data, sample_rate))] -fn write_wav( - filename: std::path::PathBuf, - data: numpy::PyReadonlyArray1, - sample_rate: u32, -) -> PyResult<()> { - let w = std::fs::File::create(&filename).w_f(&filename)?; - let mut w = std::io::BufWriter::new(w); - let data = data.as_array().to_vec(); - mm::wav::write_pcm_as_wav(&mut w, &data, sample_rate).w_f(&filename)?; - Ok(()) -} - -#[pymodule] -fn rustymimi(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_function(wrap_pyfunction!(write_wav, m)?)?; - Ok(()) -} - - - -[package] -name = "mimi-pyo3" -version.workspace = true -edition.workspace = true -description.workspace = true -repository.workspace = true -keywords.workspace = true -categories.workspace = true -license.workspace = true - -[lib] -name = "rustymimi" -crate-type = ["cdylib"] - -[dependencies] -anyhow = { workspace = true } -numpy = { workspace = true } -pyo3 = { workspace = true } -moshi = { workspace = true } - - - -[build-system] -requires = ["maturin>=1.4,<2.0"] -build-backend = "maturin" - -[project] -name = "rustymimi" -requires-python = ">=3.8" -classifiers = [ - "Programming Language :: Rust", - "Programming Language :: Python :: Implementation :: CPython", - "Programming Language :: Python :: Implementation :: PyPy", -] -dynamic = ["version"] - -[tool.maturin] -python-source = "py_src" -module-name = "rustymimi.rustymimi" -bindings = 'pyo3' -features = ["pyo3/extension-module"] - - - -# See: https://raw.githubusercontent.com/huggingface/tokenizers/main/bindings/python/stub.py -import argparse -import inspect -import os -from typing import Optional -import black -from pathlib import Path -import re - - -INDENT = " " * 4 -GENERATED_COMMENT = "# Generated content DO NOT EDIT\n" -TYPING = """from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Sequence -from os import PathLike -""" -RETURN_TYPE_MARKER = "&RETURNS&: " -FORWARD_REF_PATTERN = re.compile(r"ForwardRef\('([^']+)'\)") - - -def do_indent(text: Optional[str], indent: str): - if text is None: - return "" - return text.replace("\n", f"\n{indent}") - - -def function(obj, indent: str, text_signature: str = None): - if text_signature is None: - text_signature = obj.__text_signature__ - - text_signature = text_signature.replace("$self", "self").lstrip().rstrip() - doc_string = obj.__doc__ - if doc_string is None: - doc_string = "" - - # Check if we have a return type annotation in the docstring - return_type = None - doc_lines = doc_string.split("\n") - if doc_lines[-1].lstrip().startswith(RETURN_TYPE_MARKER): - # Extract the return type and remove it from the docstring - return_type = doc_lines[-1].lstrip()[len(RETURN_TYPE_MARKER) :].strip() - doc_string = "\n".join(doc_lines[:-1]) - - string = "" - if return_type: - string += f"{indent}def {obj.__name__}{text_signature} -> {return_type}:\n" - else: - string += f"{indent}def {obj.__name__}{text_signature}:\n" - indent += INDENT - string += f'{indent}"""\n' - string += f"{indent}{do_indent(doc_string, indent)}\n" - string += f'{indent}"""\n' - string += f"{indent}pass\n" - string += "\n" - string += "\n" - return string - - -def member_sort(member): - if inspect.isclass(member): - value = 10 + len(inspect.getmro(member)) - else: - value = 1 - return value - - -def fn_predicate(obj): - value = inspect.ismethoddescriptor(obj) or inspect.isbuiltin(obj) - if value: - return obj.__text_signature__ and not obj.__name__.startswith("_") - if inspect.isgetsetdescriptor(obj): - return not obj.__name__.startswith("_") - return False - - -def get_module_members(module): - members = [ - member - for name, member in inspect.getmembers(module) - if not name.startswith("_") and not inspect.ismodule(member) - ] - members.sort(key=member_sort) - return members - - -def pyi_file(obj, indent=""): - string = "" - if inspect.ismodule(obj): - string += GENERATED_COMMENT - string += TYPING - members = get_module_members(obj) - for member in members: - string += pyi_file(member, indent) - - elif inspect.isclass(obj): - indent += INDENT - mro = inspect.getmro(obj) - if len(mro) > 2: - inherit = f"({mro[1].__name__})" - else: - inherit = "" - string += f"class {obj.__name__}{inherit}:\n" - - body = "" - if obj.__doc__: - body += f'{indent}"""\n{indent}{do_indent(obj.__doc__, indent)}\n{indent}"""\n' - - fns = inspect.getmembers(obj, fn_predicate) - - # Init - if obj.__text_signature__: - body += f"{indent}def __init__{obj.__text_signature__}:\n" - body += f"{indent+INDENT}pass\n" - body += "\n" - - - for name, fn in fns: - body += pyi_file(fn, indent=indent) - - if not body: - body += f"{indent}pass\n" - - string += body - string += "\n\n" - - elif inspect.isbuiltin(obj): - string += f"{indent}@staticmethod\n" - string += function(obj, indent) - - elif inspect.ismethoddescriptor(obj): - string += function(obj, indent) - - elif inspect.isgetsetdescriptor(obj): - # TODO it would be interesting to add the setter maybe ? - string += f"{indent}@property\n" - string += function(obj, indent, text_signature="(self)") - - elif obj.__class__.__name__ == "DType": - string += f"class {str(obj).lower()}(DType):\n" - string += f"{indent+INDENT}pass\n" - else: - raise Exception(f"Object {obj} is not supported") - return string - - -def py_file(module, origin): - members = get_module_members(module) - - string = GENERATED_COMMENT - string += f"from .. import {origin}\n" - string += "\n" - for member in members: - if hasattr(member, "__name__"): - name = member.__name__ - else: - name = str(member) - string += f"{name} = {origin}.{name}\n" - return string - - -def do_black(content, is_pyi): - mode = black.Mode( - target_versions={black.TargetVersion.PY35}, - line_length=119, - is_pyi=is_pyi, - string_normalization=True, - ) - try: - return black.format_file_contents(content, fast=True, mode=mode) - except black.NothingChanged: - return content - - -def write(module, directory, origin, check=False): - submodules = [(name, member) for name, member in inspect.getmembers(module) if inspect.ismodule(member)] - - filename = os.path.join(directory, "__init__.pyi") - pyi_content = pyi_file(module) - pyi_content = do_black(pyi_content, is_pyi=True) - os.makedirs(directory, exist_ok=True) - if check: - with open(filename, "r") as f: - data = f.read() - assert data == pyi_content, f"The content of {filename} seems outdated, please run `python stub.py`" - else: - with open(filename, "w") as f: - f.write(pyi_content) - - filename = os.path.join(directory, "__init__.py") - py_content = py_file(module, origin) - py_content = do_black(py_content, is_pyi=False) - os.makedirs(directory, exist_ok=True) - - is_auto = False - if not os.path.exists(filename): - is_auto = True - else: - with open(filename, "r") as f: - line = f.readline() - if line == GENERATED_COMMENT: - is_auto = True - - if is_auto: - if check: - with open(filename, "r") as f: - data = f.read() - assert data == py_content, f"The content of {filename} seems outdated, please run `python stub.py`" - else: - with open(filename, "w") as f: - f.write(py_content) - - for name, submodule in submodules: - write(submodule, os.path.join(directory, name), f"{name}", check=check) - - -def extract_additional_types(module): - additional_types = {} - for name, member in inspect.getmembers(module): - if inspect.isclass(member): - if hasattr(member, "__name__"): - name = member.__name__ - else: - name = str(member) - if name not in additional_types: - additional_types[name] = member - return additional_types - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--check", action="store_true") - - args = parser.parse_args() - - cwd = Path.cwd() - directory = "py_src/rustymimi/" - - import rustymimi - write(rustymimi.rustymimi, directory, "rustymimi", check=args.check) - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -#![allow(unused)] -use std::io::prelude::*; - -pub trait Sample { - fn to_i16(&self) -> i16; -} - -impl Sample for f32 { - fn to_i16(&self) -> i16 { - (self.clamp(-1.0, 1.0) * 32767.0) as i16 - } -} - -impl Sample for f64 { - fn to_i16(&self) -> i16 { - (self.clamp(-1.0, 1.0) * 32767.0) as i16 - } -} - -impl Sample for i16 { - fn to_i16(&self) -> i16 { - *self - } -} - -pub fn write_pcm_as_wav( - w: &mut W, - samples: &[S], - sample_rate: u32, -) -> std::io::Result<()> { - let len = 12u32; // header - let len = len + 24u32; // fmt - let len = len + samples.len() as u32 * 2 + 8; // data - let n_channels = 1u16; - let bytes_per_second = sample_rate * 2 * n_channels as u32; - w.write_all(b"RIFF")?; - w.write_all(&(len - 8).to_le_bytes())?; // total length minus 8 bytes - w.write_all(b"WAVE")?; - - // Format block - w.write_all(b"fmt ")?; - w.write_all(&16u32.to_le_bytes())?; // block len minus 8 bytes - w.write_all(&1u16.to_le_bytes())?; // PCM - w.write_all(&n_channels.to_le_bytes())?; // one channel - w.write_all(&sample_rate.to_le_bytes())?; - w.write_all(&bytes_per_second.to_le_bytes())?; - w.write_all(&2u16.to_le_bytes())?; // 2 bytes of data per sample - w.write_all(&16u16.to_le_bytes())?; // bits per sample - - // Data block - w.write_all(b"data")?; - w.write_all(&(samples.len() as u32 * 2).to_le_bytes())?; - for sample in samples.iter() { - w.write_all(&sample.to_i16().to_le_bytes())? - } - Ok(()) -} - -fn conv(samples: &mut Vec, data: std::borrow::Cow>) -where - T: symphonia::core::sample::Sample, - f32: symphonia::core::conv::FromSample, -{ - use symphonia::core::audio::Signal; - use symphonia::core::conv::FromSample; - samples.extend(data.chan(0).iter().map(|v| f32::from_sample(*v))) -} - -pub(crate) fn pcm_decode>(path: P) -> anyhow::Result<(Vec, u32)> { - use symphonia::core::audio::{AudioBufferRef, Signal}; - - let src = std::fs::File::open(path)?; - let mss = symphonia::core::io::MediaSourceStream::new(Box::new(src), Default::default()); - let hint = symphonia::core::probe::Hint::new(); - let meta_opts: symphonia::core::meta::MetadataOptions = Default::default(); - let fmt_opts: symphonia::core::formats::FormatOptions = Default::default(); - let probed = symphonia::default::get_probe().format(&hint, mss, &fmt_opts, &meta_opts)?; - let mut format = probed.format; - let track = format - .tracks() - .iter() - .find(|t| t.codec_params.codec != symphonia::core::codecs::CODEC_TYPE_NULL) - .expect("no supported audio tracks"); - let mut decoder = symphonia::default::get_codecs() - .make(&track.codec_params, &Default::default()) - .expect("unsupported codec"); - let track_id = track.id; - let sample_rate = track.codec_params.sample_rate.unwrap_or(0); - let mut pcm_data = Vec::new(); - while let Ok(packet) = format.next_packet() { - while !format.metadata().is_latest() { - format.metadata().pop(); - } - if packet.track_id() != track_id { - continue; - } - match decoder.decode(&packet)? { - AudioBufferRef::F32(buf) => pcm_data.extend(buf.chan(0)), - AudioBufferRef::U8(data) => conv(&mut pcm_data, data), - AudioBufferRef::U16(data) => conv(&mut pcm_data, data), - AudioBufferRef::U24(data) => conv(&mut pcm_data, data), - AudioBufferRef::U32(data) => conv(&mut pcm_data, data), - AudioBufferRef::S8(data) => conv(&mut pcm_data, data), - AudioBufferRef::S16(data) => conv(&mut pcm_data, data), - AudioBufferRef::S24(data) => conv(&mut pcm_data, data), - AudioBufferRef::S32(data) => conv(&mut pcm_data, data), - AudioBufferRef::F64(data) => conv(&mut pcm_data, data), - } - } - Ok((pcm_data, sample_rate)) -} - -pub(crate) fn resample(pcm_in: &[f32], sr_in: usize, sr_out: usize) -> anyhow::Result> { - use rubato::Resampler; - - let mut pcm_out = - Vec::with_capacity((pcm_in.len() as f64 * sr_out as f64 / sr_in as f64) as usize + 1024); - - let mut resampler = rubato::FftFixedInOut::::new(sr_in, sr_out, 1024, 1)?; - let mut output_buffer = resampler.output_buffer_allocate(true); - let mut pos_in = 0; - while pos_in + resampler.input_frames_next() < pcm_in.len() { - let (in_len, out_len) = - resampler.process_into_buffer(&[&pcm_in[pos_in..]], &mut output_buffer, None)?; - pos_in += in_len; - pcm_out.extend_from_slice(&output_buffer[0][..out_len]); - } - - if pos_in < pcm_in.len() { - let (_in_len, out_len) = resampler.process_partial_into_buffer( - Some(&[&pcm_in[pos_in..]]), - &mut output_buffer, - None, - )?; - pcm_out.extend_from_slice(&output_buffer[0][..out_len]); - } - - Ok(pcm_out) -} - -pub(crate) fn write_opus_header(w: &mut W) -> std::io::Result<()> { - use byteorder::WriteBytesExt; - - // https://wiki.xiph.org/OggOpus#ID_Header - w.write_all(b"OpusHead")?; - w.write_u8(1)?; // version - w.write_u8(1)?; // channel count - w.write_u16::(3840)?; // pre-skip - w.write_u32::(48000)?; // sample-rate in Hz - w.write_i16::(0)?; // output gain Q7.8 in dB - w.write_u8(0)?; // channel map - Ok(()) -} - -pub(crate) fn write_opus_tags(w: &mut W) -> std::io::Result<()> { - use byteorder::WriteBytesExt; - - // https://wiki.xiph.org/OggOpus#Comment_Header - let vendor = "KyutaiMoshi"; - w.write_all(b"OpusTags")?; - w.write_u32::(vendor.len() as u32)?; // vendor string length - w.write_all(vendor.as_bytes())?; // vendor string, UTF8 encoded - w.write_u32::(0u32)?; // number of tags - Ok(()) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::stream_both::{AppStateInner, Config, SessionConfigReq, StreamOut, StreamingModel}; -use anyhow::Result; -use std::sync::mpsc; - -#[derive(serde::Serialize)] -#[serde(tag = "type")] -enum Event { - InputPcm { pcm_len: usize, time: f64 }, - Step { step: usize, time: f64 }, - StepPostSampling { step: usize, time: f64 }, - SendPcm { pcm_len: usize, time: f64 }, -} - -#[derive(serde::Serialize)] -struct StatsTracker { - events: Vec, -} - -fn system_time() -> f64 { - std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs_f64() -} - -impl StatsTracker { - fn new() -> Self { - Self { events: vec![] } - } - - fn on_update(&mut self, out: StreamOut) { - match out { - StreamOut::Pcm { pcm } => { - self.events.push(Event::SendPcm { time: system_time(), pcm_len: pcm.len() }); - } - StreamOut::MetaData { metadata } => { - tracing::info!(?metadata, "send-metadata"); - } - StreamOut::Text { text } => { - tracing::info!(text, "send-text"); - } - StreamOut::InputPcm { pcm_len } => { - self.events.push(Event::InputPcm { time: system_time(), pcm_len }); - } - StreamOut::StepStart { step } => { - self.events.push(Event::Step { time: system_time(), step }); - } - StreamOut::StepPostSampling { step } => { - self.events.push(Event::StepPostSampling { time: system_time(), step }); - } - StreamOut::Ready => {} - } - } -} - -pub async fn run(args: &crate::BenchmarkArgs, config: &Config) -> Result<()> { - tracing::info!( - avx = ?candle::utils::with_avx(), - neon = ?candle::utils::with_neon(), - simd128 = ?candle::utils::with_simd128(), - f16c = ?candle::utils::with_f16c(), - ?config, - "cpu" - ); - tracing::info!(?config, "starting benchmark"); - let session_config = SessionConfigReq { - text_temperature: None, - text_topk: None, - audio_temperature: None, - audio_topk: None, - max_steps: Some(args.steps), - audio_seed: Some(299792458), - text_seed: Some(299792458), - email: None, - pad_mult: None, - repetition_penalty_context: None, - repetition_penalty: None, - }; - if args.mimi_only { - let device = crate::standalone::device(args.cpu)?; - let mimi_device = if config.use_cpu_for_mimi { &candle::Device::Cpu } else { &device }; - let mut mimi_model = moshi::mimi::load( - &config.mimi_model_file, - Some(config.mimi_num_codebooks), - mimi_device, - )?; - let config = mimi_model.config(); - let frame_length = (config.sample_rate / config.frame_rate).ceil() as usize; - for _step in 0..args.steps { - let fake_pcm = - candle::Tensor::zeros((1, 1, frame_length), candle::DType::F32, mimi_device)?; - let codes = mimi_model.encode_step(&fake_pcm.into(), &().into())?; - let ys = mimi_model.decode_step(&codes, &().into())?; - if ys.as_option().is_none() { - anyhow::bail!("Expected mimi to output some stuff, but nothing came out."); - } - device.synchronize()?; - } - } else { - let standalone_args = crate::StandaloneArgs { cpu: args.cpu }; - let state = std::sync::Arc::new(AppStateInner::new(&standalone_args, config)?); - for _i in 0..args.reps { - let sm = StreamingModel::new(&state, session_config.clone()); - let (in_pcm_tx, in_pcm_rx) = mpsc::channel(); - let (stream_out_tx, mut stream_out_rx) = tokio::sync::mpsc::unbounded_channel(); - let w = tokio::task::spawn_blocking(move || sm.run(in_pcm_rx, stream_out_tx, None)); - - let task = tokio::spawn({ - let stat_file = args.stat_file.clone(); - async move { - let mut st = StatsTracker::new(); - while let Some(out) = stream_out_rx.recv().await { - st.on_update(out) - } - tracing::info!("stream-out receiver closed"); - if let Some(stat_file) = stat_file { - use std::io::Write; - - let json_string = serde_json::to_string_pretty(&st).unwrap(); - let mut stat_file = std::fs::File::create(stat_file).unwrap(); - stat_file.write_all(json_string.as_bytes()).unwrap() - } - } - }); - let zeros = vec![0f32; 48000 / 25]; - let start_time = tokio::time::Instant::now(); - for step in 0..args.steps + 20 { - let target_time = - start_time + tokio::time::Duration::from_millis(80).mul_f64(step as f64); - tokio::time::sleep_until(target_time).await; - in_pcm_tx.send(zeros.to_vec())?; - } - task.await?; - w.await??; - } - } - Ok(()) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::Result; -use vergen::EmitBuilder; - -pub fn main() -> Result<()> { - // NOTE: This will output everything, and requires all features enabled. - // NOTE: See the EmitBuilder documentation for configuration options. - EmitBuilder::builder().all_build().all_cargo().all_git().all_rustc().all_sysinfo().emit()?; - Ok(()) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::Result; -use clap::Parser; -use std::str::FromStr; - -mod audio; -mod benchmark; -mod standalone; -mod stream_both; -mod utils; - -#[derive(Parser, Debug)] -#[clap(name = "server", about = "moshi web server")] -struct Args { - #[clap(short = 'l', long = "log", default_value = "info")] - log_level: String, - - #[clap(long)] - config: String, - - #[clap(long)] - silent: bool, - - #[command(subcommand)] - command: Command, -} - -#[derive(Parser, Debug)] -struct StandaloneArgs { - #[clap(long)] - cpu: bool, -} - -#[derive(Clone, Parser, Debug)] -pub struct BenchmarkArgs { - #[clap(long)] - cpu: bool, - - #[clap(short = 'n', long, default_value_t = 200)] - steps: usize, - - #[clap(short = 'n', long, default_value_t = 1)] - reps: usize, - - #[clap(short = 's', long)] - stat_file: Option, - - #[clap(long)] - chrome_tracing: bool, - - #[clap(long)] - asr: bool, - - #[clap(long)] - mimi_only: bool, -} - -#[derive(Debug, clap::Subcommand)] -enum Command { - Standalone(StandaloneArgs), - Benchmark(BenchmarkArgs), -} - -/// A TLS acceptor that sets `TCP_NODELAY` on accepted streams. -#[derive(Clone, Debug)] -pub struct NoDelayAcceptor; - -impl axum_server::accept::Accept for NoDelayAcceptor { - type Stream = tokio::net::TcpStream; - type Service = S; - type Future = - futures_util::future::BoxFuture<'static, std::io::Result<(Self::Stream, Self::Service)>>; - - fn accept(&self, stream: tokio::net::TcpStream, service: S) -> Self::Future { - Box::pin(async move { - // Disable Nagle's algorithm. - stream.set_nodelay(true)?; - Ok::<_, std::io::Error>((stream, service)) - }) - } -} - -fn tracing_init( - log_dir: &str, - instance_name: &str, - log_level: &str, - silent: bool, -) -> Result { - use tracing_subscriber::prelude::*; - - let build_info = utils::BuildInfo::new(); - let file_appender = tracing_appender::rolling::daily(log_dir, format!("log.{instance_name}")); - let (non_blocking, guard) = tracing_appender::non_blocking(file_appender); - let filter = tracing_subscriber::filter::LevelFilter::from_str(log_level)?; - let mut layers = vec![tracing_subscriber::fmt::layer() - .with_writer(non_blocking) - .with_filter(filter) - .boxed()]; - if !silent { - layers.push(Box::new( - tracing_subscriber::fmt::layer().with_writer(std::io::stdout).with_filter(filter), - )) - }; - tracing_subscriber::registry().with(layers).init(); - tracing::info!(?build_info); - Ok(guard) -} - -#[tokio::main(flavor = "multi_thread")] -async fn main() -> Result<()> { - let args = Args::parse(); - match args.command { - Command::Standalone(standalone_args) => { - let mut config = standalone::Config::load(&args.config)?; - let _guard = tracing_init( - &config.stream.log_dir, - &config.stream.instance_name, - &args.log_level, - args.silent, - )?; - tracing::info!("starting process with pid {}", std::process::id()); - - if config.stream.requires_model_download() { - standalone::download_from_hub(&mut config.stream).await?; - } - if !std::path::PathBuf::from(&config.static_dir).exists() { - use hf_hub::api::tokio::Api; - let api = Api::new()?; - let repo = api.model("kyutai/moshi-artifacts".to_string()); - let dist_tgz = repo.get("dist.tgz").await?; - if let Some(parent) = dist_tgz.parent() { - let dist = parent.join("dist"); - if !dist.exists() { - let output = std::process::Command::new("tar") - .arg("-xzf") - .arg(&dist_tgz) - .arg("-C") - .arg(parent) - .output()?; - if !output.status.success() { - anyhow::bail!( - "error extract {dist_tgz:?}: {}", - String::from_utf8_lossy(&output.stderr) - ); - } - } - config.static_dir = dist.to_string_lossy().to_string() - } - } - standalone::run(&standalone_args, &config).await?; - } - Command::Benchmark(standalone_args) => { - let config = stream_both::Config::load(&args.config)?; - let _guard = if standalone_args.chrome_tracing { - use tracing_chrome::ChromeLayerBuilder; - use tracing_subscriber::prelude::*; - let (chrome_layer, guard) = ChromeLayerBuilder::new().build(); - tracing_subscriber::registry().with(chrome_layer).init(); - let b: Box = Box::new(guard); - b - } else { - let guard = tracing_init( - &config.log_dir, - &config.instance_name, - &args.log_level, - args.silent, - )?; - let b: Box = Box::new(guard); - b - }; - benchmark::run(&standalone_args, &config).await?; - } - } - Ok(()) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::{Context, Result}; -use axum::extract::ws; -use std::sync::Arc; -use std::{path::Path, str::FromStr}; - -use crate::{stream_both, StandaloneArgs}; - -#[derive(serde::Deserialize, Debug, Clone)] -pub struct Config { - cert_dir: String, - pub static_dir: String, - addr: String, - port: u16, - - #[serde(flatten)] - pub stream: stream_both::Config, -} - -impl Config { - pub fn load>(p: P) -> Result { - let config = std::fs::read_to_string(p)?; - let mut config: Self = serde_json::from_str(&config)?; - config.static_dir = crate::utils::replace_env_vars(&config.static_dir); - config.cert_dir = crate::utils::replace_env_vars(&config.cert_dir); - config.stream.log_dir = crate::utils::replace_env_vars(&config.stream.log_dir); - config.stream.text_tokenizer_file = - crate::utils::replace_env_vars(&config.stream.text_tokenizer_file); - config.stream.mimi_model_file = - crate::utils::replace_env_vars(&config.stream.mimi_model_file); - config.stream.lm_model_file = crate::utils::replace_env_vars(&config.stream.lm_model_file); - Ok(config) - } - - pub fn cert_file(&self, name: &str) -> std::path::PathBuf { - let cert_dir = std::path::PathBuf::from(&self.cert_dir); - cert_dir.join(name) - } -} - -pub(crate) fn device(cpu: bool) -> Result { - use candle::Device; - if cpu { - Ok(Device::Cpu) - } else if candle::utils::cuda_is_available() { - Ok(Device::new_cuda(0)?) - } else if candle::utils::metal_is_available() { - Ok(Device::new_metal(0)?) - } else { - Ok(Device::Cpu) - } -} - -impl stream_both::AppStateInner { - pub fn new(args: &StandaloneArgs, config: &stream_both::Config) -> Result { - let device = device(args.cpu)?; - let dtype = if device.is_cuda() { candle::DType::BF16 } else { candle::DType::F32 }; - let lm_model = moshi::lm::load_streaming(&config.lm_model_file, dtype, &device)?; - let mimi_device = if config.use_cpu_for_mimi { &candle::Device::Cpu } else { &device }; - let mimi_model = moshi::mimi::load( - &config.mimi_model_file, - Some(config.mimi_num_codebooks), - mimi_device, - )?; - let text_tokenizer = - sentencepiece::SentencePieceProcessor::open(&config.text_tokenizer_file)?; - // Warm-up code. - { - tracing::info!(?dtype, ?device, "warming up the model"); - let mut lm_model = lm_model.clone(); - let (_v, ys) = - lm_model.forward(None, vec![None; config.mimi_num_codebooks], &().into())?; - let mut lp = candle_transformers::generation::LogitsProcessor::new(123, None, None); - let _ = lm_model.depformer_sample(&ys, None, &[], &mut lp)?; - let mut mimi_model = mimi_model.clone(); - let config = mimi_model.config(); - let frame_length = (config.sample_rate / config.frame_rate).ceil() as usize; - let fake_pcm = - candle::Tensor::zeros((1, 1, frame_length), candle::DType::F32, mimi_device)?; - let codes = mimi_model.encode_step(&fake_pcm.into(), &().into())?; - let ys = mimi_model.decode_step(&codes, &().into())?; - if ys.as_option().is_none() { - anyhow::bail!("Expected Mimi to output some stuff, but nothing came out."); - } - device.synchronize()?; - tracing::info!("model is ready to roll!"); - } - Ok(Self { lm_model, mimi_model, device, config: config.clone(), text_tokenizer }) - } -} - -async fn handle_socket(socket: ws::WebSocket, sm: stream_both::StreamingModel) { - if let Err(err) = stream_both::handle_socket(socket, sm, None).await { - tracing::error!(err = err.to_string(), "handle_socket") - } -} - -pub async fn stream_handler( - ws: ws::WebSocketUpgrade, - axum::extract::ConnectInfo(addr): axum::extract::ConnectInfo, - state: axum::extract::State, - req: axum::extract::Query, -) -> impl axum::response::IntoResponse { - tracing::info!(?addr, "received connection"); - let sm = stream_both::StreamingModel::new(&state.0, req.0); - ws.on_upgrade(move |v| handle_socket(v, sm)) -} - -pub async fn download_from_hub(config: &mut stream_both::Config) -> Result<()> { - use hf_hub::api::tokio::Api; - let api = Api::new()?; - let repo = api.model(config.hf_repo.clone()); - let extract_filename = |path: &str| -> Result { - Path::new(path) - .file_name() - .and_then(|f| f.to_str()) - .map(String::from) - .ok_or_else(|| anyhow::anyhow!("'{path}' has no file name")) - }; - for file_path in - [&mut config.lm_model_file, &mut config.mimi_model_file, &mut config.text_tokenizer_file] - .iter_mut() - { - let filename = extract_filename(file_path) - .with_context(|| format!("Failed to extract filename for '{file_path}'"))?; - let downloaded_path = repo - .get(&filename) - .await - .with_context(|| format!("Failed to download '{file_path}' file"))?; - **file_path = downloaded_path - .into_os_string() - .into_string() - .map_err(|_| anyhow::anyhow!("'{file_path}' path is not a valid string"))?; - } - Ok(()) -} - -pub async fn run(args: &StandaloneArgs, config: &Config) -> Result<()> { - let cert_pem = config.cert_file("cert.pem"); - let key_pem = config.cert_file("key.pem"); - if !cert_pem.exists() || !key_pem.exists() { - let rcgen::CertifiedKey { cert, key_pair } = - rcgen::generate_simple_self_signed(vec!["localhost".to_string()])?; - std::fs::write(&cert_pem, cert.pem())?; - std::fs::write(&key_pem, key_pair.serialize_pem())?; - } - - let tls_config = - axum_server::tls_rustls::RustlsConfig::from_pem_file(cert_pem, key_pem).await?; - let sock_addr = std::net::SocketAddr::from(( - std::net::IpAddr::from_str(config.addr.as_str()) - .unwrap_or(std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)), - config.port, - )); - let state = Arc::new(stream_both::AppStateInner::new(args, &config.stream)?); - tracing::info!("serving static dir {}", config.static_dir); - let app = axum::Router::new() - .route("/api/chat", axum::routing::get(stream_handler)) - .fallback_service( - tower_http::services::ServeDir::new(&config.static_dir) - .append_index_html_on_directories(true), - ) - .layer(tower::ServiceBuilder::new().layer(tower_http::trace::TraceLayer::new_for_http())) - .with_state(state); - tracing::info!("standalone worker listening on https://{}", sock_addr); - axum_server::bind_rustls(sock_addr, tls_config) - .serve(app.into_make_service_with_connect_info::()) - .await?; - Ok(()) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::Result; -use axum::extract::ws; -use futures_util::{ - stream::{SplitSink, SplitStream, StreamExt}, - SinkExt, -}; -use std::sync::Arc; - -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct Config { - pub instance_name: String, - #[serde(default)] - pub hf_repo: String, - pub lm_model_file: String, - pub log_dir: String, - pub text_tokenizer_file: String, - pub mimi_model_file: String, - pub mimi_num_codebooks: usize, - pub lm_config: Option, - #[serde(default = "default_false")] - pub use_cpu_for_mimi: bool, - pub asr_delay_in_tokens: Option, -} - -fn default_false() -> bool { - false -} - -impl Config { - pub fn load>(p: P) -> Result { - let config = std::fs::read_to_string(p)?; - let mut config: Self = serde_json::from_str(&config)?; - config.log_dir = crate::utils::replace_env_vars(&config.log_dir); - config.text_tokenizer_file = crate::utils::replace_env_vars(&config.text_tokenizer_file); - config.mimi_model_file = crate::utils::replace_env_vars(&config.mimi_model_file); - config.lm_model_file = crate::utils::replace_env_vars(&config.lm_model_file); - Ok(config) - } - - /// Check if all modelling files are available on machine. - pub fn requires_model_download(&self) -> bool { - [&self.lm_model_file, &self.mimi_model_file, &self.text_tokenizer_file] - .iter() - .any(|file| !std::path::Path::new(file).exists()) - } -} - -pub type AppState = Arc; -pub struct AppStateInner { - pub lm_model: moshi::lm::LmModel, - pub mimi_model: moshi::mimi::Mimi, - pub text_tokenizer: sentencepiece::SentencePieceProcessor, - pub device: candle::Device, - pub config: Config, -} - -impl AppStateInner { - fn text( - &self, - prev_text_token: u32, - text_token: u32, - config: &moshi::lm_generate_multistream::Config, - ) -> Option { - if text_token != config.text_start_token - && text_token != config.text_pad_token - && text_token != config.text_eop_token - { - if prev_text_token == config.text_start_token { - self.text_tokenizer.decode_piece_ids(&[text_token]).ok() - } else { - let prev_ids = self.text_tokenizer.decode_piece_ids(&[prev_text_token]).ok(); - let ids = self.text_tokenizer.decode_piece_ids(&[prev_text_token, text_token]).ok(); - prev_ids.and_then(|prev_ids| { - ids.map(|ids| { - if ids.len() > prev_ids.len() { - ids[prev_ids.len()..].to_string() - } else { - String::new() - } - }) - }) - } - } else { - None - } - } -} - -#[derive(serde::Deserialize, Debug, Clone)] -pub struct SessionConfigReq { - pub text_temperature: Option, - pub text_topk: Option, - pub audio_temperature: Option, - pub audio_topk: Option, - pub max_steps: Option, - pub audio_seed: Option, - pub text_seed: Option, - pub email: Option, - pub pad_mult: Option, - pub repetition_penalty_context: Option, - pub repetition_penalty: Option, -} - -#[derive(serde::Serialize, Debug, Clone)] -pub struct SessionConfig { - pub text_temperature: f64, - pub text_topk: usize, - pub audio_temperature: f64, - pub audio_topk: usize, - pub max_steps: usize, - pub audio_seed: u64, - pub text_seed: u64, - pub pad_mult: Option, - pub repetition_penalty: Option<(usize, f32)>, - pub email: Option, - pub user_feedback: Option, -} - -#[derive(serde::Serialize, Debug, Clone)] -struct SessionSummary<'a> { - #[serde(flatten)] - session_config: &'a SessionConfig, - last_step_idx: usize, - transcript: String, - addr: Option, - lm_model_file: &'a str, - mimi_model_file: &'a str, - #[serde(flatten)] - lm_config: &'a Option, -} - -impl SessionConfigReq { - fn into_session_config(self) -> SessionConfig { - use rand::Rng; - - let repetition_penalty = self.repetition_penalty_context.zip(self.repetition_penalty); - SessionConfig { - text_temperature: self.text_temperature.unwrap_or(0.8), - text_topk: self.text_topk.unwrap_or(250), - text_seed: self.text_seed.unwrap_or_else(|| rand::thread_rng().gen()), - audio_temperature: self.audio_temperature.unwrap_or(0.8), - audio_topk: self.audio_topk.unwrap_or(250), - audio_seed: self.audio_seed.unwrap_or_else(|| rand::thread_rng().gen()), - email: self.email, - user_feedback: None, - max_steps: self.max_steps.unwrap_or(4500).min(4500), - pad_mult: self.pad_mult, - repetition_penalty, - } - } -} - -#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)] -pub struct MetaData { - text_temperature: f64, - text_topk: usize, - audio_temperature: f64, - audio_topk: usize, - pad_mult: f32, - repetition_penalty_context: usize, - repetition_penalty: f32, - lm_model_file: String, - mimi_model_file: String, - build_info: crate::utils::BuildInfo, - instance_name: String, -} - -#[derive(Debug, Clone)] -pub enum StreamOut { - Ready, - InputPcm { pcm_len: usize }, - MetaData { metadata: Box }, - StepStart { step: usize }, - StepPostSampling { step: usize }, - Text { text: String }, - Pcm { pcm: Vec }, -} - -// This must be an allowed value among 120, 240, 480, 960, 1920, and 2880. -// Using a different value would result in a BadArg "invalid argument" error when calling encode. -// https://opus-codec.org/docs/opus_api-1.2/group__opus__encoder.html#ga4ae9905859cd241ef4bb5c59cd5e5309 -const OPUS_ENCODER_FRAME_SIZE: usize = 960; - -#[derive(Debug, Clone, Copy)] -pub enum MsgType { - Handshake, - Audio, - Text, - Control, - Metadata, - Error, - Ping, -} - -impl MsgType { - pub fn from_u8(v: u8) -> Result { - let s = match v { - 0 => MsgType::Handshake, - 1 => MsgType::Audio, - 2 => MsgType::Text, - 3 => MsgType::Control, - 4 => MsgType::Metadata, - 5 => MsgType::Error, - 6 => MsgType::Ping, - _ => anyhow::bail!("unexpected msg type {v}"), - }; - Ok(s) - } - - pub fn to_u8(self) -> u8 { - match self { - MsgType::Handshake => 0, - MsgType::Audio => 1, - MsgType::Text => 2, - MsgType::Control => 3, - MsgType::Metadata => 4, - MsgType::Error => 5, - MsgType::Ping => 6, - } - } -} - -pub struct MsgSender { - pw: ogg::PacketWriter<'static, Vec>, - encoder: opus::Encoder, - out_pcm: std::collections::VecDeque, - out_pcm_buf: Vec, - total_data: usize, - sender: SplitSink, -} - -impl MsgSender { - fn new(sender: SplitSink) -> Result { - let encoder = opus::Encoder::new(24000, opus::Channels::Mono, opus::Application::Voip)?; - // Not sure what the appropriate buffer size would be here. - let out_pcm_buf = vec![0u8; 50_000]; - let out_pcm = std::collections::VecDeque::with_capacity(2 * OPUS_ENCODER_FRAME_SIZE); - - let all_data = Vec::new(); - let mut pw = ogg::PacketWriter::new(all_data); - let mut head = Vec::new(); - crate::audio::write_opus_header(&mut head)?; - pw.write_packet(head, 42, ogg::PacketWriteEndInfo::EndPage, 0)?; - let mut tags = Vec::new(); - crate::audio::write_opus_tags(&mut tags)?; - pw.write_packet(tags, 42, ogg::PacketWriteEndInfo::EndPage, 0)?; - Ok(Self { pw, encoder, out_pcm, out_pcm_buf, total_data: 0, sender }) - } - - async fn send_text(&mut self, text: String) -> Result<()> { - let msg: Vec = [&[MsgType::Text.to_u8()], text.as_bytes()].concat(); - let msg = ws::Message::Binary(msg.into()); - self.sender.send(msg).await?; - Ok(()) - } - - async fn send_ready(&mut self) -> Result<()> { - // The payload is made of two fields. - // 1. Protocol version (`u32`) - always 0 for now. - // 2. Model version (`u32`). - let msg: Vec = [&[MsgType::Handshake.to_u8()], [0u8; 8].as_slice()].concat(); - let msg = ws::Message::Binary(msg.into()); - self.sender.send(msg).await?; - Ok(()) - } - - async fn send_metadata(&mut self, md: Box) -> Result<()> { - let bytes = serde_json::to_vec(&md)?; - let msg: Vec = [&[MsgType::Metadata.to_u8()], bytes.as_slice()].concat(); - let msg = ws::Message::Binary(msg.into()); - self.sender.send(msg).await?; - Ok(()) - } - - async fn send_pcm(&mut self, pcm: Vec) -> Result<()> { - self.out_pcm.extend(pcm.iter()); - self.total_data += pcm.len(); - let nchunks = self.out_pcm.len() / OPUS_ENCODER_FRAME_SIZE; - for _chunk_id in 0..nchunks { - let mut chunk = Vec::with_capacity(OPUS_ENCODER_FRAME_SIZE); - for _i in 0..OPUS_ENCODER_FRAME_SIZE { - let v = match self.out_pcm.pop_front() { - None => anyhow::bail!("unexpected err popping from pcms"), - Some(v) => v, - }; - chunk.push(v) - } - let size = self.encoder.encode_float(&chunk, &mut self.out_pcm_buf)?; - if size > 0 { - let msg = self.out_pcm_buf[..size].to_vec(); - self.pw.write_packet( - msg, - 42, - ogg::PacketWriteEndInfo::EndPage, - self.total_data as u64, - )? - } else { - tracing::error!("OPUS SIZE 0") - } - let data = self.pw.inner_mut(); - if !data.is_empty() { - let msg: Vec = [&[MsgType::Audio.to_u8()], data.as_slice()].concat(); - let msg = ws::Message::Binary(msg.into()); - self.sender.send(msg).await?; - self.sender.flush().await?; - data.clear(); - } else { - tracing::error!("OGG SIZE 0") - } - } - Ok(()) - } -} - -pub struct StreamingModel { - state: AppState, - device: candle::Device, - config: moshi::lm_generate_multistream::Config, - session_config: SessionConfig, -} - -impl StreamingModel { - fn run_with_state_asr( - &self, - state: &mut moshi::lm_generate_multistream::State, - receiver: std::sync::mpsc::Receiver>, - sender: tokio::sync::mpsc::UnboundedSender, - asr_delay_in_tokens: usize, - ) -> Result<()> { - use candle::IndexOp; - - let app_state = &self.state; - - let mut mimi = app_state.mimi_model.clone(); - let config = state.config().clone(); - - mimi.reset_state(); - tracing::info!("processing loop"); - let mut prev_text_token = config.text_start_token; - let mimi_device = - if self.state.config.use_cpu_for_mimi { &candle::Device::Cpu } else { &self.device }; - mimi_device.synchronize()?; - sender.send(StreamOut::Ready)?; - while let Ok(in_pcm) = receiver.recv() { - if in_pcm.is_empty() { - continue; - } - let pcm_len = in_pcm.len(); - sender.send(StreamOut::InputPcm { pcm_len })?; - let pcms = candle::Tensor::from_vec(in_pcm, (1, 1, pcm_len), mimi_device)?; - let audio_tokens = mimi.encode_step(&pcms.into(), &().into())?; - let audio_tokens = match audio_tokens.as_option() { - None => continue, - Some(audio_tokens) => audio_tokens, - }; - let (_one, _codebooks, steps) = audio_tokens.dims3()?; - - for step in 0..steps { - let codes = audio_tokens.i((0, .., step))?.to_vec1::()?; - // For the ASR, we don't provide text tokens during the initial steps except the - // initial one. - if state.step_idx() > 0 && state.step_idx() < asr_delay_in_tokens { - prev_text_token = state.step_(None, &codes, None, None, None)?; - } else { - sender.send(StreamOut::StepStart { step })?; - let text_token = state.step(prev_text_token, &codes, None, None)?; - sender.send(StreamOut::StepPostSampling { step })?; - if let Some(text) = app_state.text(prev_text_token, text_token, &config) { - sender.send(StreamOut::Text { text })?; - } - prev_text_token = text_token; - } - } - } - tracing::info!("finished the processing loop"); - Ok(()) - } - - fn run_with_state( - &self, - state: &mut moshi::lm_generate_multistream::State, - receiver: std::sync::mpsc::Receiver>, - sender: tokio::sync::mpsc::UnboundedSender, - ) -> Result<()> { - use candle::IndexOp; - - let app_state = &self.state; - - let mut mimi = app_state.mimi_model.clone(); - let config = state.config().clone(); - - mimi.reset_state(); - tracing::info!("processing loop"); - let mut prev_text_token = config.text_start_token; - let mut tensor_tokens = vec![]; - let mimi_device = - if self.state.config.use_cpu_for_mimi { &candle::Device::Cpu } else { &self.device }; - mimi_device.synchronize()?; - sender.send(StreamOut::Ready)?; - while let Ok(in_pcm) = receiver.recv() { - if in_pcm.is_empty() { - continue; - } - let pcm_len = in_pcm.len(); - sender.send(StreamOut::InputPcm { pcm_len })?; - let pcms = candle::Tensor::from_vec(in_pcm, (1, 1, pcm_len), mimi_device)?; - let audio_tokens = mimi.encode_step(&pcms.into(), &().into())?; - let audio_tokens = match audio_tokens.as_option() { - None => continue, - Some(audio_tokens) => audio_tokens, - }; - let (_one, _codebooks, steps) = audio_tokens.dims3()?; - - for step in 0..steps { - let codes = audio_tokens.i((0, .., step))?.to_vec1::()?; - sender.send(StreamOut::StepStart { step })?; - let text_token = state.step(prev_text_token, &codes, None, None)?; - sender.send(StreamOut::StepPostSampling { step })?; - if let Some(audio_tokens) = state.last_audio_tokens() { - let audio_tokens = { - let cb = app_state.config.mimi_num_codebooks; - candle::Tensor::from_slice(&audio_tokens[..cb], (1, cb, 1), mimi_device)? - }; - tensor_tokens.push(audio_tokens.clone()); - let pcm = mimi.decode_step(&audio_tokens.into(), &().into())?; - if let Some(pcm) = pcm.as_option() { - let pcm = pcm.i((0, 0))?.to_vec1::()?; - sender.send(StreamOut::Pcm { pcm })?; - } - } - if let Some(text) = app_state.text(prev_text_token, text_token, &config) { - sender.send(StreamOut::Text { text })?; - } - prev_text_token = text_token; - } - } - tracing::info!("finished the processing loop"); - Ok(()) - } - - fn run_with_state_mt( - &self, - state: &mut moshi::lm_generate_multistream::State, - receiver: std::sync::mpsc::Receiver>, - sender: tokio::sync::mpsc::UnboundedSender, - ) -> Result<()> { - use candle::IndexOp; - - let app_state = &self.state; - - let mut mimi = app_state.mimi_model.clone(); - let config = state.config().clone(); - - mimi.reset_state(); - tracing::info!("processing loop"); - let mut prev_text_token = config.text_start_token; - let mut tensor_tokens = vec![]; - let (tx_i, rx_i) = std::sync::mpsc::channel::<(Vec, usize)>(); - let (tx_o, rx_o) = std::sync::mpsc::channel::>(); - let sender = Arc::new(sender); - let status = std::thread::scope(|s| { - s.spawn({ - let mut mimi = mimi.clone(); - let sender = sender.clone(); - move || { - 'outer: while let Ok(in_pcm) = receiver.recv() { - if in_pcm.is_empty() { - continue; - } - let pcm_len = in_pcm.len(); - sender.send(StreamOut::InputPcm { pcm_len })?; - let pcms = candle::Tensor::from_vec( - in_pcm, - (1, 1, pcm_len), - &candle::Device::Cpu, - )?; - let audio_tokens = mimi.encode_step(&pcms.into(), &().into())?; - let audio_tokens = match audio_tokens.as_option() { - None => continue, - Some(audio_tokens) => audio_tokens, - }; - let (_one, _codebooks, steps) = audio_tokens.dims3()?; - for step in 0..steps { - let codes = audio_tokens.i((0, .., step))?.to_vec1::()?; - if tx_i.send((codes, step)).is_err() { - break 'outer; - } - } - } - Ok::<_, anyhow::Error>(()) - } - }); - s.spawn({ - let cb = app_state.config.mimi_num_codebooks; - let sender = sender.clone(); - move || { - while let Ok(audio_tokens) = rx_o.recv() { - let audio_tokens = { - candle::Tensor::from_slice( - &audio_tokens[..cb], - (1, cb, 1), - &candle::Device::Cpu, - )? - }; - tensor_tokens.push(audio_tokens.clone()); - let pcm = mimi.decode_step(&audio_tokens.into(), &().into())?; - if let Some(pcm) = pcm.as_option() { - let pcm = pcm.i((0, 0))?.to_vec1::()?; - sender.send(StreamOut::Pcm { pcm })?; - } - } - Ok::<_, anyhow::Error>(()) - } - }); - sender.send(StreamOut::Ready)?; - while let Ok((codes, step)) = rx_i.recv() { - tracing::info!("received codes"); - sender.send(StreamOut::StepStart { step })?; - let text_token = state.step(prev_text_token, &codes, None, None); - sender.send(StreamOut::StepPostSampling { step })?; - tracing::info!(?text_token, "codes"); - if text_token.is_err() { - drop(rx_i); - drop(tx_o); - break; - } - let text_token = text_token?; - if let Some(audio_tokens) = state.last_audio_tokens() { - tx_o.send(audio_tokens)? - } - if let Some(text) = app_state.text(prev_text_token, text_token, &config) { - sender.send(StreamOut::Text { text })?; - } - prev_text_token = text_token; - } - Ok::<_, anyhow::Error>(()) - }); - match status { - Ok(()) => tracing::info!("finished the processing loop"), - Err(err) => tracing::error!(?err, "processing loop"), - }; - Ok(()) - } - - pub fn new(state: &AppState, session_config: SessionConfigReq) -> Self { - let config = match state.config.lm_config.as_ref() { - None => moshi::lm_generate_multistream::Config::v0_1(), - Some(config) => config.clone(), - }; - let session_config = session_config.into_session_config(); - Self { state: state.clone(), device: state.device.clone(), config, session_config } - } - - pub fn run( - &self, - receiver: std::sync::mpsc::Receiver>, - sender: tokio::sync::mpsc::UnboundedSender, - addr: Option, - ) -> Result<()> { - let app_state = &self.state; - let (repetition_penalty_context, repetition_penalty) = - self.session_config.repetition_penalty.unwrap_or((32, 1.)); - let metadata = MetaData { - text_temperature: self.session_config.text_temperature, - text_topk: self.session_config.text_topk, - audio_temperature: self.session_config.audio_temperature, - audio_topk: self.session_config.audio_topk, - pad_mult: self.session_config.pad_mult.unwrap_or(0.), - repetition_penalty, - repetition_penalty_context, - lm_model_file: self.state.config.lm_model_file.to_string(), - mimi_model_file: self.state.config.mimi_model_file.to_string(), - build_info: crate::utils::BuildInfo::new(), - instance_name: self.state.config.instance_name.to_string(), - }; - sender.send(StreamOut::MetaData { metadata: Box::new(metadata) })?; - let lm_model = app_state.lm_model.clone(); - let audio_lp = candle_transformers::generation::LogitsProcessor::from_sampling( - self.session_config.audio_seed, - candle_transformers::generation::Sampling::TopK { - k: self.session_config.audio_topk, - temperature: self.session_config.audio_temperature, - }, - ); - let text_lp = candle_transformers::generation::LogitsProcessor::from_sampling( - self.session_config.text_seed, - candle_transformers::generation::Sampling::TopK { - k: self.session_config.text_topk, - temperature: self.session_config.text_temperature, - }, - ); - let mut state = moshi::lm_generate_multistream::State::new( - lm_model, - self.session_config.max_steps, - audio_lp, - text_lp, - self.session_config.pad_mult, - self.session_config.repetition_penalty, - None, - self.config.clone(), - ); - - // We want to log the output even if the run function returns an error. - let run_result = if self.state.config.use_cpu_for_mimi { - self.run_with_state_mt(&mut state, receiver, sender) - } else if let Some(asr_delay_in_tokens) = self.state.config.asr_delay_in_tokens { - self.run_with_state_asr(&mut state, receiver, sender, asr_delay_in_tokens) - } else { - self.run_with_state(&mut state, receiver, sender) - }; - { - let text_tokens = state.text_tokens(false); - let transcript = { - let text_tokens = text_tokens - .iter() - .filter_map(|v| { - let v = *v; - if v != moshi::lm_generate_multistream::UNGENERATED - && v != self.config.text_pad_token - && v != self.config.text_eop_token - && v != self.config.text_start_token - { - Some(v) - } else { - None - } - }) - .collect::>(); - self.state - .text_tokenizer - .decode_piece_ids(&text_tokens) - .unwrap_or_else(|_| String::new()) - }; - let audio_tokens = state.audio_tokens(false); - let audio_tokens = audio_tokens - .iter() - .map(|v| { - v.iter() - .map(|v| { - if *v == moshi::lm_generate_multistream::UNGENERATED { - -1 - } else { - *v as i64 - } - }) - .collect::>() - }) - .collect::>(); - let text_tokens = candle::Tensor::new(text_tokens, &candle::Device::Cpu)? - .to_dtype(candle::DType::I64)?; - let audio_tokens = candle::Tensor::new(audio_tokens, &candle::Device::Cpu)?; - let since_epoch = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?; - let (secs, us) = (since_epoch.as_secs(), since_epoch.subsec_micros()); - let log_dir = &app_state.config.log_dir; - let base_path = format!("{log_dir}/{}-{secs}-{us}", app_state.config.instance_name); - let json_filename = format!("{base_path}.json"); - let json_content = serde_json::to_string_pretty(&SessionSummary { - session_config: &self.session_config, - last_step_idx: state.step_idx(), - transcript, - addr, - mimi_model_file: &self.state.config.mimi_model_file, - lm_model_file: &self.state.config.lm_model_file, - lm_config: &self.state.config.lm_config, - })?; - std::fs::write(json_filename, json_content)?; - let st_filename = format!("{base_path}.safetensors"); - let st_content = - std::collections::HashMap::from([("text", text_tokens), ("audio", audio_tokens)]); - candle::safetensors::save(&st_content, st_filename)?; - } - run_result - } -} - -type Handle = tokio::task::JoinHandle>; - -fn spawn_recv_loops( - mut receiver: SplitStream, - sender: std::sync::mpsc::Sender>, -) -> Result<(Handle, Handle)> { - use tokio::io::AsyncWriteExt; - - let (mut tx, rx) = tokio::io::duplex(100_000); - let mut pr = ogg::reading::async_api::PacketReader::new(rx); - let mut decoder = opus::Decoder::new(24000, opus::Channels::Mono)?; - let handle1 = tokio::spawn({ - async move { - loop { - match receiver.next().await { - None => { - // The close logic is that if this loop exits, then tx gets dropped so pr - // gets closed and the second thread gets dropped resulting in sender - // getting dropped. - break; - } - Some(v) => { - let v = v?.into_data(); - if v.is_empty() { - continue; - } - let msg_type = MsgType::from_u8(v[0])?; - match msg_type { - MsgType::Metadata => {} - MsgType::Handshake => {} - MsgType::Control => {} - MsgType::Text => {} - MsgType::Error => {} - MsgType::Ping => {} - MsgType::Audio => tx.write_all(&v[1..]).await?, - } - } - } - } - tracing::info!("socket closed"); - Ok::<_, anyhow::Error>(()) - } - }); - let handle2 = tokio::spawn(async move { - // TODO: dynamic sizing? - let mut pcm_buf = vec![0f32; 24_000 * 10]; - let mut size_in_buf = 0; - loop { - match pr.next().await { - None => { - break; - } - Some(packet) => { - let packet = packet?; - if packet.data.starts_with(b"OpusHead") || packet.data.starts_with(b"OpusTags") - { - continue; - } - let read_size = decoder.decode_float( - &packet.data, - &mut pcm_buf[size_in_buf..], - /* Forward Error Correction */ false, - )?; - size_in_buf += read_size; - // flush the data every half timestep - if size_in_buf >= 24_000 / 25 { - if sender.send(pcm_buf[..size_in_buf].to_vec()).is_err() { - break; - } - size_in_buf = 0; - } - } - } - } - tracing::info!("decoder closed"); - Ok::<_, anyhow::Error>(()) - }); - Ok((handle1, handle2)) -} - -async fn sender_loop( - mut stream_out_rx: tokio::sync::mpsc::UnboundedReceiver, - mut sender: MsgSender, -) -> Result<()> { - // It is important for the recv here to be an async enabled one. Otherwise this could lead - // to some weird deadlocks. - while let Some(v) = stream_out_rx.recv().await { - match v { - StreamOut::Pcm { pcm } => sender.send_pcm(pcm).await?, - StreamOut::Ready => sender.send_ready().await?, - StreamOut::MetaData { metadata } => sender.send_metadata(metadata).await?, - StreamOut::Text { text } => sender.send_text(text).await?, - StreamOut::InputPcm { .. } - | StreamOut::StepStart { .. } - | StreamOut::StepPostSampling { .. } => {} - } - } - Ok::<_, anyhow::Error>(()) -} - -pub async fn handle_socket( - socket: ws::WebSocket, - sm: StreamingModel, - addr: Option, -) -> Result<()> { - tracing::info!("accepted websocket connection"); - let (sender, receiver) = socket.split(); - let sender = MsgSender::new(sender)?; - - tracing::info!("starting streaming"); - - let (in_pcm_tx, in_pcm_rx) = std::sync::mpsc::channel(); - let (stream_out_tx, stream_out_rx) = tokio::sync::mpsc::unbounded_channel(); - let (loop1, loop2) = spawn_recv_loops(receiver, in_pcm_tx)?; - std::thread::spawn(move || { - if let Err(err) = sm.run(in_pcm_rx, stream_out_tx, addr) { - tracing::error!("{err}") - } - }); - let sender_loop = tokio::spawn(async move { - match sender_loop(stream_out_rx, sender).await { - Ok(()) => tracing::info!("sender closed"), - Err(err) => { - // Using the Display trait rather than the Debug one so as not to include the backtrace. - let err = format!("{err}"); - tracing::info!(err, "sender err") - } - } - }); - - let sleep = tokio::time::sleep(std::time::Duration::from_secs(360)); - tokio::pin!(sleep); - // select should ensure that all the threads get aborted on timeout. - tokio::select! { - _ = &mut sleep => { - tracing::error!("reached timeout"); - } - r = loop1 => { - tracing::error!(?r, "loop1 ended") - } - r = loop2 => { - tracing::error!(?r, "loop2 ended") - } - r = sender_loop => { - tracing::error!(?r, "sender loop ended") - } - } - Ok(()) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -#[derive(Debug, PartialEq, Clone, serde::Deserialize, serde::Serialize)] -pub struct BuildInfo { - build_timestamp: String, - build_date: String, - git_branch: String, - git_timestamp: String, - git_date: String, - git_hash: String, - git_describe: String, - rustc_host_triple: String, - rustc_version: String, - cargo_target_triple: String, -} - -impl BuildInfo { - pub fn new() -> BuildInfo { - BuildInfo { - build_timestamp: String::from(env!("VERGEN_BUILD_TIMESTAMP")), - build_date: String::from(env!("VERGEN_BUILD_DATE")), - git_branch: String::from(env!("VERGEN_GIT_BRANCH")), - git_timestamp: String::from(env!("VERGEN_GIT_COMMIT_TIMESTAMP")), - git_date: String::from(env!("VERGEN_GIT_COMMIT_DATE")), - git_hash: String::from(env!("VERGEN_GIT_SHA")), - git_describe: String::from(env!("VERGEN_GIT_DESCRIBE")), - rustc_host_triple: String::from(env!("VERGEN_RUSTC_HOST_TRIPLE")), - rustc_version: String::from(env!("VERGEN_RUSTC_SEMVER")), - cargo_target_triple: String::from(env!("VERGEN_CARGO_TARGET_TRIPLE")), - } - } -} - -pub struct WrapJson(pub anyhow::Result); - -impl axum::response::IntoResponse for WrapJson { - fn into_response(self) -> axum::response::Response { - match self.0 { - Ok(v) => axum::Json(v).into_response(), - Err(err) => { - tracing::error!(?err, "returning internal server error 500"); - (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("{err}")).into_response() - } - } - } -} - -pub fn replace_env_vars(input: &str) -> String { - let re = regex::Regex::new(r"\$([A-Za-z_][A-Za-z0-9_]*)").unwrap(); - re.replace_all(input, |caps: ®ex::Captures| { - let var_name = &caps[1]; - std::env::var(var_name).unwrap_or_else(|_| "".to_string()) - }) - .to_string() -} - -pub struct WrapBincode(pub anyhow::Result); - -impl axum::response::IntoResponse for WrapBincode { - fn into_response(self) -> axum::response::Response { - match self.0.and_then(|v| Ok(bincode::serialize(&v)?)) { - Ok(v) => (axum::http::StatusCode::OK, v).into_response(), - Err(err) => { - tracing::error!(?err, "returning internal server error 500"); - (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("{err}")).into_response() - } - } - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::Result; -use vergen::EmitBuilder; - -pub fn main() -> Result<()> { - // NOTE: This will output everything, and requires all features enabled. - // NOTE: See the EmitBuilder documentation for configuration options. - EmitBuilder::builder().all_build().all_cargo().all_git().all_rustc().all_sysinfo().emit()?; - Ok(()) -} - - - -[package] -name = "moshi-backend" -version.workspace = true -edition.workspace = true -description.workspace = true -repository.workspace = true -keywords.workspace = true -categories.workspace = true -license.workspace = true - -[dependencies] -anyhow = { workspace = true } -axum = { workspace = true } -axum-server = { workspace = true } -base64ct = { workspace = true } -bincode = { workspace = true } -byteorder = { workspace = true } -candle = { workspace = true } -candle-nn = { workspace = true } -candle-transformers = { workspace = true } -clap = { workspace = true } -env_logger = { workspace = true } -futures-util = { workspace = true } -hf-hub = { workspace = true } -rcgen = { workspace = true } -http = { workspace = true } -lazy_static = { workspace = true } -log = { workspace = true } -moshi = { workspace = true } -ogg = { workspace = true } -opus = { workspace = true } -rand = { workspace = true } -rand_chacha = { workspace = true } -regex = { workspace = true } -rubato = { workspace = true } -sentencepiece = { workspace = true } -serde = { workspace = true } -serde_json = { workspace = true } -sha3 = { workspace = true } -symphonia = { workspace = true } -tokenizers = { workspace = true } -tokio = { workspace = true } -tokio-rustls = { workspace = true } -tower = { workspace = true } -tower-http = { workspace = true } -tracing = { workspace = true } -tracing-appender = { workspace = true } -tracing-chrome = { workspace = true } -tracing-subscriber = { workspace = true } - -[build-dependencies] -anyhow = { workspace = true } -vergen = { workspace = true } - -[features] -default = [] -cuda = ["moshi/cuda", "candle/cuda", "candle-nn/cuda", "candle-transformers/cuda"] -metal = ["moshi/metal", "candle/metal", "candle-nn/metal", "candle-transformers/metal"] - -[profile.release] -debug = true - -[profile.release-no-debug] -inherits = "release" -debug = false - - - -{ - "instance_name": "foo", - "hf_repo": "kyutai/moshiko-candle-q8", - "lm_model_file": "$HOME/tmp/moshiko_rs_301e30bf@120/model.q8.gguf", - "text_tokenizer_file": "$HOME/tmp/tokenizer_spm_32k_3.model", - "log_dir": "$HOME/tmp/moshi-logs", - "mimi_model_file": "$HOME/tmp/tokenizer-e351c8d8-checkpoint125.safetensors", - "mimi_num_codebooks": 8, - "static_dir": "../client/dist", - "addr": "0.0.0.0", - "port": 8998, - "cert_dir": "." -} - - - -{ - "instance_name": "foo", - "hf_repo": "kyutai/moshiko-candle-bf16", - "lm_model_file": "$HOME/tmp/moshiko_rs_301e30bf@120/model.safetensors", - "text_tokenizer_file": "$HOME/tmp/tokenizer_spm_32k_3.model", - "log_dir": "$HOME/tmp/moshi-logs", - "mimi_model_file": "$HOME/tmp/tokenizer-e351c8d8-checkpoint125.safetensors", - "mimi_num_codebooks": 8, - "static_dir": "../client/dist", - "addr": "0.0.0.0", - "port": 8998, - "cert_dir": "." -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -#![allow(unused)] -use anyhow::{Context, Result}; -use std::collections::VecDeque; -use std::sync::{Arc, Mutex}; - -pub const SAMPLE_RATE: usize = 24_000; - -pub(crate) struct AudioOutputData_ { - resampled_data: std::collections::VecDeque, - resampler: rubato::FastFixedIn, - output_buffer: Vec, - input_buffer: Vec, - input_len: usize, - // The number of (resampled) samples that have been seen so far. - total_samples: usize, - // Some subtitle index together with the index at which it should get printed. - subs: VecDeque<(usize, String)>, - mean_squares: f32, -} - -impl AudioOutputData_ { - pub(crate) fn new(input_sample_rate: usize, output_sample_rate: usize) -> Result { - use rubato::Resampler; - - let resampled_data = std::collections::VecDeque::with_capacity(output_sample_rate * 10); - let resample_ratio = output_sample_rate as f64 / input_sample_rate as f64; - let resampler = rubato::FastFixedIn::new( - resample_ratio, - f64::max(resample_ratio, 1.0), - rubato::PolynomialDegree::Septic, - 1024, - 1, - )?; - let input_buffer = resampler.input_buffer_allocate(true).remove(0); - let output_buffer = resampler.output_buffer_allocate(true).remove(0); - Ok(Self { - resampled_data, - resampler, - input_buffer, - output_buffer, - input_len: 0, - total_samples: 0, - subs: VecDeque::new(), - mean_squares: 0., - }) - } - - pub(crate) fn total_samples(&self) -> usize { - self.total_samples - } - - pub(crate) fn samples_in_buffer(&self) -> usize { - self.resampled_data.len() - } - - pub fn reset(&mut self) { - use rubato::Resampler; - self.output_buffer.fill(0.); - self.input_buffer.fill(0.); - self.total_samples = 0; - self.resampler.reset(); - self.resampled_data.clear(); - self.subs.clear(); - self.mean_squares = 0.; - } - - pub(crate) fn take_all(&mut self) -> Vec { - let mut data = Vec::with_capacity(self.resampled_data.len()); - while let Some(elem) = self.resampled_data.pop_back() { - data.push(elem); - } - data - } - - pub(crate) fn db10(&self) -> f32 { - 10. + (self.mean_squares + 1e-10).log10() - } - - pub(crate) fn clear(&mut self) { - self.resampled_data.clear(); - self.subs.clear(); - } - - pub(crate) fn is_empty(&self) -> bool { - self.resampled_data.is_empty() - } - - // Assumes that the input buffer is large enough. - fn push_input_buffer(&mut self, samples: &[f32]) { - self.input_buffer[self.input_len..self.input_len + samples.len()].copy_from_slice(samples); - self.input_len += samples.len(); - self.total_samples += samples.len(); - } - - pub(crate) fn push_sub(&mut self, sub: String) { - self.subs.push_back((self.total_samples, sub)) - } - - pub(crate) fn push_samples(&mut self, samples: &[f32]) -> Result<()> { - use rubato::Resampler; - - let mut pos_in = 0; - loop { - let rem = self.input_buffer.len() - self.input_len; - let pos_end = usize::min(pos_in + rem, samples.len()); - self.push_input_buffer(&samples[pos_in..pos_end]); - pos_in = pos_end; - if self.input_len < self.input_buffer.len() { - break; - } - let (_, out_len) = self.resampler.process_into_buffer( - &[&self.input_buffer], - &mut [&mut self.output_buffer], - None, - )?; - for &elem in self.output_buffer[..out_len].iter() { - self.resampled_data.push_front(elem) - } - self.input_len = 0; - } - Ok(()) - } -} - -type AudioOutputData = Arc>; - -pub(crate) fn setup_output_stream(real_time: bool) -> Result<(cpal::Stream, AudioOutputData)> { - setup_output_stream_map(real_time, |s| { - use std::io::Write; - print!("{s}"); - let _ = std::io::stdout().flush(); - }) -} - -pub(crate) fn setup_output_stream_map( - real_time: bool, - mut f: F, -) -> Result<(cpal::Stream, AudioOutputData)> { - use cpal::traits::{DeviceTrait, HostTrait, StreamTrait}; - - println!("Setup audio output stream!"); - let host = cpal::default_host(); - let device = host.default_output_device().context("no output device available")?; - let mut supported_configs_range = device.supported_output_configs()?; - let config_range = match supported_configs_range.find(|c| c.channels() == 1) { - // On macOS, it's commonly the case that there are only stereo outputs. - None => device.supported_output_configs()?.next().context("no audio output available")?, - Some(config_range) => config_range, - }; - let sample_rate = cpal::SampleRate(SAMPLE_RATE as u32) - .clamp(config_range.min_sample_rate(), config_range.max_sample_rate()); - let config: cpal::StreamConfig = config_range.with_sample_rate(sample_rate).into(); - let channels = config.channels as usize; - println!( - "cpal device: {} {} {config:?}", - device.name().unwrap_or_else(|_| "unk".to_string()), - config.sample_rate.0 - ); - let audio_data = - Arc::new(Mutex::new(AudioOutputData_::new(SAMPLE_RATE, config.sample_rate.0 as usize)?)); - let ad = audio_data.clone(); - let mut total_samples = 0; - let stream = device.build_output_stream( - &config, - move |data: &mut [f32], _: &cpal::OutputCallbackInfo| { - data.fill(0.); - let mut ad = ad.lock().unwrap(); - let mut last_elem = 0f32; - loop { - let should_pop = match ad.subs.front() { - None => false, - Some((i, _)) => *i < total_samples, - }; - if !should_pop { - break; - } - if let Some((_, s)) = ad.subs.pop_front() { - f(s) - } - } - for (idx, elem) in data.iter_mut().enumerate() { - if idx % channels == 0 { - match ad.resampled_data.pop_back() { - None => break, - Some(v) => { - last_elem = v; - total_samples += 1; - *elem = v - } - } - } else { - *elem = last_elem - } - } - if real_time && ad.resampled_data.len() > SAMPLE_RATE / 2 { - let pcm_data = ad.resampled_data.drain(..).collect::>(); - if let Ok(pcm_data) = resample(&pcm_data, SAMPLE_RATE, SAMPLE_RATE * 2 / 3) { - for v in pcm_data.into_iter().rev() { - ad.resampled_data.push_back(v) - } - } - } - }, - move |err| eprintln!("cpal error: {err}"), - None, // None=blocking, Some(Duration)=timeout - )?; - stream.play()?; - Ok((stream, audio_data)) -} - -pub(crate) fn setup_input_stream() -> Result<(cpal::Stream, AudioOutputData)> { - use cpal::traits::{DeviceTrait, HostTrait, StreamTrait}; - - println!("Setup audio input stream!"); - let host = cpal::default_host(); - let device = host.default_input_device().context("no input device available")?; - let mut supported_configs_range = device.supported_input_configs()?; - let config_range = - supported_configs_range.find(|c| c.channels() == 1).context("no audio input available")?; - let sample_rate = cpal::SampleRate(SAMPLE_RATE as u32) - .clamp(config_range.min_sample_rate(), config_range.max_sample_rate()); - let config: cpal::StreamConfig = config_range.with_sample_rate(sample_rate).into(); - println!( - "cpal device: {} {} {config:?}", - device.name().unwrap_or_else(|_| "unk".to_string()), - config.sample_rate.0 - ); - let audio_data = - Arc::new(Mutex::new(AudioOutputData_::new(config.sample_rate.0 as usize, SAMPLE_RATE)?)); - let ad = audio_data.clone(); - let sample_rate = config.sample_rate.0 as f32; - let stream = device.build_input_stream( - &config, - move |data: &[f32], _: &cpal::InputCallbackInfo| { - let mut ad = ad.lock().unwrap(); - if !data.is_empty() { - let l = data.len() as f32; - let mean = data.iter().sum::() / l; - let mean_squares = data.iter().map(|v| (v - mean) * (v - mean)).sum::() / l; - let decay = (-l / sample_rate * 10.).exp2(); - ad.mean_squares = decay * ad.mean_squares + (1. - decay) * mean_squares; - } - if let Err(err) = ad.push_samples(data) { - eprintln!("error processing audio input {err:?}") - } - }, - move |err| eprintln!("cpal error: {err}"), - None, // None=blocking, Some(Duration)=timeout - )?; - stream.play()?; - Ok((stream, audio_data)) -} - -fn conv(samples: &mut Vec, data: std::borrow::Cow>) -where - T: symphonia::core::sample::Sample, - f32: symphonia::core::conv::FromSample, -{ - use symphonia::core::audio::Signal; - use symphonia::core::conv::FromSample; - samples.extend(data.chan(0).iter().map(|v| f32::from_sample(*v))) -} - -pub(crate) fn pcm_decode>(path: P) -> Result<(Vec, u32)> { - use symphonia::core::audio::{AudioBufferRef, Signal}; - - let src = std::fs::File::open(path)?; - let mss = symphonia::core::io::MediaSourceStream::new(Box::new(src), Default::default()); - let hint = symphonia::core::probe::Hint::new(); - let meta_opts: symphonia::core::meta::MetadataOptions = Default::default(); - let fmt_opts: symphonia::core::formats::FormatOptions = Default::default(); - let probed = symphonia::default::get_probe().format(&hint, mss, &fmt_opts, &meta_opts)?; - let mut format = probed.format; - let track = format - .tracks() - .iter() - .find(|t| t.codec_params.codec != symphonia::core::codecs::CODEC_TYPE_NULL) - .expect("no supported audio tracks"); - let mut decoder = symphonia::default::get_codecs() - .make(&track.codec_params, &Default::default()) - .expect("unsupported codec"); - let track_id = track.id; - let sample_rate = track.codec_params.sample_rate.unwrap_or(0); - let mut pcm_data = Vec::new(); - while let Ok(packet) = format.next_packet() { - while !format.metadata().is_latest() { - format.metadata().pop(); - } - if packet.track_id() != track_id { - continue; - } - match decoder.decode(&packet)? { - AudioBufferRef::F32(buf) => pcm_data.extend(buf.chan(0)), - AudioBufferRef::U8(data) => conv(&mut pcm_data, data), - AudioBufferRef::U16(data) => conv(&mut pcm_data, data), - AudioBufferRef::U24(data) => conv(&mut pcm_data, data), - AudioBufferRef::U32(data) => conv(&mut pcm_data, data), - AudioBufferRef::S8(data) => conv(&mut pcm_data, data), - AudioBufferRef::S16(data) => conv(&mut pcm_data, data), - AudioBufferRef::S24(data) => conv(&mut pcm_data, data), - AudioBufferRef::S32(data) => conv(&mut pcm_data, data), - AudioBufferRef::F64(data) => conv(&mut pcm_data, data), - } - } - Ok((pcm_data, sample_rate)) -} - -pub(crate) fn resample(pcm_in: &[f32], sr_in: usize, sr_out: usize) -> Result> { - use rubato::Resampler; - - let mut pcm_out = - Vec::with_capacity((pcm_in.len() as f64 * sr_out as f64 / sr_in as f64) as usize + 1024); - - let mut resampler = rubato::FftFixedInOut::::new(sr_in, sr_out, 1024, 1)?; - let mut output_buffer = resampler.output_buffer_allocate(true); - let mut pos_in = 0; - while pos_in + resampler.input_frames_next() < pcm_in.len() { - let (in_len, out_len) = - resampler.process_into_buffer(&[&pcm_in[pos_in..]], &mut output_buffer, None)?; - pos_in += in_len; - pcm_out.extend_from_slice(&output_buffer[0][..out_len]); - } - - if pos_in < pcm_in.len() { - let (_in_len, out_len) = resampler.process_partial_into_buffer( - Some(&[&pcm_in[pos_in..]]), - &mut output_buffer, - None, - )?; - pcm_out.extend_from_slice(&output_buffer[0][..out_len]); - } - - Ok(pcm_out) -} - - - -use anyhow::Result; -use candle::{Device, IndexOp, Tensor}; - -pub struct Args { - pub lm_model_file: String, - pub lm_config_file: String, - pub mimi_model_file: String, - pub audio_input_file: String, - pub text_tokenizer: String, - pub audio_output_file: String, - pub seed: u64, - pub cfg_alpha: Option, -} - -pub fn run(args: &Args, dev: &Device) -> Result<()> { - let dtype = dev.bf16_default_to_f32(); - tracing::info!(?dtype, ?dev); - - tracing::info!("loading the audio input"); - let (in_pcm, in_pcm_len) = { - let (mut pcm, sample_rate) = crate::audio_io::pcm_decode(&args.audio_input_file)?; - pcm.extend_from_slice(&vec![0.0; 12000]); - let pcm = if sample_rate != 24_000 { - crate::audio_io::resample(&pcm, sample_rate as usize, 24_000)? - } else { - pcm - }; - let pcm_len = pcm.len(); - let pcm = Tensor::from_vec(pcm, (1, 1, pcm_len), dev)?; - (pcm, pcm_len) - }; - tracing::info!(in_pcm_len, "loaded the audio input"); - - tracing::info!("loading the config"); - let lm_config = std::fs::read_to_string(&args.lm_config_file)?; - let lm_config: moshi::lm::Config = toml::from_str(&lm_config)?; - tracing::info!("loading the audio tokenizer"); - let mut mimi = moshi::mimi::load(&args.mimi_model_file, Some(8), dev)?; - tracing::info!("loading the lm"); - let lm_model = moshi::lm::load_lm_model(lm_config.clone(), &args.lm_model_file, dtype, dev)?; - tracing::info!("loading the text tokenizer"); - let text_tokenizer = sentencepiece::SentencePieceProcessor::open(&args.text_tokenizer)?; - tracing::info!("done loading models"); - - let audio_lp = candle_transformers::generation::LogitsProcessor::from_sampling( - args.seed, - candle_transformers::generation::Sampling::TopK { k: 250, temperature: 0.8 }, - ); - let text_lp = candle_transformers::generation::LogitsProcessor::from_sampling( - args.seed, - candle_transformers::generation::Sampling::TopK { k: 250, temperature: 0.8 }, - ); - let generated_audio_codebooks = lm_config.depformer.as_ref().map_or(8, |v| v.num_slices); - - let conditions = match lm_model.condition_provider() { - None => None, - Some(cp) => { - let conditions = if args.cfg_alpha.is_some() { - use moshi::conditioner::Condition::AddToInput; - let AddToInput(c1) = cp.condition_lut("description", "very_good")?; - let AddToInput(c2) = cp.condition_lut("description", "very_bad")?; - AddToInput(Tensor::cat(&[c1, c2], 0)?) - } else { - cp.condition_lut("description", "very_good")? - }; - tracing::info!(?conditions, "generated conditions"); - Some(conditions) - } - }; - let max_steps = 2500; - let cfg_alpha = if args.cfg_alpha == Some(1.) { None } else { args.cfg_alpha }; - let mut state = { - let config = moshi::lm_generate_multistream::Config { - acoustic_delay: 2, - audio_vocab_size: lm_config.audio_vocab_size, - generated_audio_codebooks, - input_audio_codebooks: lm_config.audio_codebooks - generated_audio_codebooks, - text_start_token: lm_config.text_out_vocab_size as u32, - text_eop_token: 0, - text_pad_token: 3, - }; - moshi::lm_generate_multistream::State::new( - lm_model, - max_steps + 20, - audio_lp, - text_lp, - None, - None, - cfg_alpha, - config, - ) - }; - - let mut prev_text_token = state.config().text_start_token; - let mut out_pcms = vec![]; - let mut text_tokens = vec![]; - let mut nsteps = 0; - tracing::info!("starting the inference loop"); - let start_time = std::time::Instant::now(); - for start_index in 0..(in_pcm_len / 1920).min(max_steps) { - nsteps += 1; - let in_pcm = in_pcm.i((.., .., start_index * 1920..(start_index + 1) * 1920))?; - let codes = mimi.encode_step(&in_pcm.into(), &().into())?; - if let Some(codes) = codes.as_option() { - let (_b, _codebooks, steps) = codes.dims3()?; - for step in 0..steps { - let codes = codes.i((.., .., step..step + 1))?; - let codes = codes.i((0, .., 0))?.to_vec1::()?; - prev_text_token = - state.step_(Some(prev_text_token), &codes, None, None, conditions.as_ref())?; - if prev_text_token != 0 && prev_text_token != 3 { - text_tokens.push(prev_text_token) - } - if let Some(audio_tokens) = state.last_audio_tokens() { - let audio_tokens = - Tensor::new(&audio_tokens[..generated_audio_codebooks], dev)? - .reshape((1, 1, ()))? - .t()?; - let out_pcm = mimi.decode_step(&audio_tokens.into(), &().into())?; - if let Some(out_pcm) = out_pcm.as_option() { - out_pcms.push(out_pcm.clone()); - } - } - } - } - } - let dt = start_time.elapsed().as_secs_f32(); - tracing::info!( - "generated {nsteps} steps in {dt:.2}s, {:.0}ms/token", - dt * 1000. / (nsteps as f32) - ); - let str = text_tokenizer.decode_piece_ids(&text_tokens)?; - tracing::info!(str, "generated text"); - let out_pcms = Tensor::cat(&out_pcms, 2)?; - tracing::info!(shape = ?out_pcms.shape(), "generated audio"); - let out_pcms = out_pcms.i((0, 0))?.to_vec1::()?; - let mut out_wav = std::fs::File::create(&args.audio_output_file)?; - moshi::wav::write_pcm_as_wav(&mut out_wav, &out_pcms, 24_000)?; - tracing::info!(audio = args.audio_output_file, "generated audio"); - Ok(()) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::Result; -use clap::Parser; - -mod audio_io; -mod gen; -mod multistream; - -use candle::Device; - -#[derive(Debug, Parser)] -struct Args { - #[command(subcommand)] - command: Command, - - /// Enable tracing (generates a trace-timestamp.json file). - #[arg(long)] - tracing: bool, -} - -#[derive(Debug, clap::Subcommand)] -enum Command { - Client { - #[arg(long)] - host: String, - - #[arg(long, default_value_t = 8998)] - port: usize, - }, - Tui { - #[arg(long)] - host: String, - - #[arg(long, default_value_t = 8998)] - port: usize, - }, - Gen { - #[arg(long)] - lm_model_file: String, - - #[arg(long)] - mimi_model_file: String, - - #[arg(long)] - lm_config_file: String, - - #[arg(long)] - text_tokenizer: String, - - #[arg(long)] - audio_input_file: String, - - #[arg(long)] - audio_output_file: String, - - #[arg(long, default_value_t = 299_792_458)] - seed: u64, - - #[arg(long)] - cfg_alpha: Option, - - /// Run on cpu - #[arg(long)] - cpu: bool, - }, -} - -pub fn device(cpu: bool) -> Result { - if cpu { - Ok(Device::Cpu) - } else if candle::utils::cuda_is_available() { - Ok(Device::new_cuda(0)?) - } else if candle::utils::metal_is_available() { - Ok(Device::new_metal(0)?) - } else { - Ok(Device::Cpu) - } -} - -#[tokio::main(flavor = "multi_thread", worker_threads = 10)] -async fn main() -> Result<()> { - use tracing_chrome::ChromeLayerBuilder; - use tracing_subscriber::prelude::*; - - let args = Args::parse(); - let _guard = if args.tracing { - let (chrome_layer, guard) = ChromeLayerBuilder::new().build(); - tracing_subscriber::registry().with(chrome_layer).init(); - Some(guard) - } else { - None - }; - match args.command { - Command::Client { host, port } => { - tracing_subscriber::fmt::init(); - multistream::client::run(host, port).await? - } - Command::Tui { host, port } => { - tracing_subscriber::fmt::init(); - multistream::client_tui::run(host, port).await? - } - Command::Gen { - seed, - text_tokenizer, - lm_model_file, - lm_config_file, - mimi_model_file, - audio_input_file, - audio_output_file, - cfg_alpha, - cpu, - } => { - let dev = device(cpu)?; - tracing_subscriber::fmt::init(); - let args = gen::Args { - lm_model_file, - mimi_model_file, - text_tokenizer, - lm_config_file, - audio_input_file, - audio_output_file, - seed, - cfg_alpha, - }; - gen::run(&args, &dev)? - } - } - Ok(()) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -pub mod client { - use anyhow::Result; - use futures_util::{ - stream::{SplitSink, StreamExt}, - SinkExt, - }; - use std::io::Write; - use tokio::io::AsyncWriteExt; - use tokio_tungstenite::tungstenite::protocol::Message; - - type WebSocket = tokio_tungstenite::WebSocketStream< - tokio_tungstenite::MaybeTlsStream, - >; - - const OPUS_ENCODER_FRAME_SIZE: usize = 960; - - pub struct MsgSender { - pw: ogg::PacketWriter<'static, Vec>, - encoder: opus::Encoder, - out_pcm: std::collections::VecDeque, - out_pcm_buf: Vec, - total_data: usize, - sender: SplitSink, - } - - pub(crate) fn write_opus_header(w: &mut W) -> std::io::Result<()> { - use byteorder::WriteBytesExt; - - // https://wiki.xiph.org/OggOpus#ID_Header - w.write_all(b"OpusHead")?; - w.write_u8(1)?; // version - w.write_u8(1)?; // channel count - w.write_u16::(3840)?; // pre-skip - w.write_u32::(48000)?; // sample-rate in Hz - w.write_i16::(0)?; // output gain Q7.8 in dB - w.write_u8(0)?; // channel map - Ok(()) - } - - pub(crate) fn write_opus_tags(w: &mut W) -> std::io::Result<()> { - use byteorder::WriteBytesExt; - - // https://wiki.xiph.org/OggOpus#Comment_Header - let vendor = "KyutaiMoshi"; - w.write_all(b"OpusTags")?; - w.write_u32::(vendor.len() as u32)?; // vendor string length - w.write_all(vendor.as_bytes())?; // vendor string, UTF8 encoded - w.write_u32::(0u32)?; // number of tags - Ok(()) - } - - impl MsgSender { - pub fn new(sender: SplitSink) -> Result { - let encoder = opus::Encoder::new(24000, opus::Channels::Mono, opus::Application::Voip)?; - // Not sure what the appropriate buffer size would be here. - let out_pcm_buf = vec![0u8; 50_000]; - let out_pcm = std::collections::VecDeque::with_capacity(2 * OPUS_ENCODER_FRAME_SIZE); - - let all_data = Vec::new(); - let mut pw = ogg::PacketWriter::new(all_data); - let mut head = Vec::new(); - write_opus_header(&mut head)?; - pw.write_packet(head, 42, ogg::PacketWriteEndInfo::EndPage, 0)?; - let mut tags = Vec::new(); - write_opus_tags(&mut tags)?; - pw.write_packet(tags, 42, ogg::PacketWriteEndInfo::EndPage, 0)?; - Ok(Self { pw, encoder, out_pcm, out_pcm_buf, total_data: 0, sender }) - } - - pub async fn send_control(&mut self, control: u8) -> Result<()> { - let msg = Message::Binary(vec![3u8, control]); - self.sender.send(msg).await?; - Ok(()) - } - - pub async fn send_pcm(&mut self, pcm: &[f32]) -> Result<()> { - self.out_pcm.extend(pcm.iter()); - self.total_data += pcm.len(); - let nchunks = self.out_pcm.len() / OPUS_ENCODER_FRAME_SIZE; - for _chunk_id in 0..nchunks { - let mut chunk = Vec::with_capacity(OPUS_ENCODER_FRAME_SIZE); - for _i in 0..OPUS_ENCODER_FRAME_SIZE { - let v = match self.out_pcm.pop_front() { - None => anyhow::bail!("unexpected err popping from pcms"), - Some(v) => v, - }; - chunk.push(v) - } - let size = self.encoder.encode_float(&chunk, &mut self.out_pcm_buf)?; - if size > 0 { - let msg = self.out_pcm_buf[..size].to_vec(); - self.pw.write_packet( - msg, - 42, - ogg::PacketWriteEndInfo::EndPage, - self.total_data as u64, - )? - } - let data = self.pw.inner_mut(); - if !data.is_empty() { - let msg: Vec = [&[1u8], data.as_slice()].concat(); - let msg = Message::Binary(msg); - self.sender.send(msg).await?; - data.clear(); - } - } - Ok(()) - } - } - - pub async fn run(host: String, port: usize) -> Result<()> { - let uri = format!("wss://{host}:{port}/api/chat"); - tracing::info!("connecting to {uri}"); - let (_stream, ad) = crate::audio_io::setup_output_stream(true)?; - let (_in_stream, input_audio) = crate::audio_io::setup_input_stream()?; - let connector = - native_tls::TlsConnector::builder().danger_accept_invalid_certs(true).build()?; - let (stream, response) = tokio_tungstenite::connect_async_tls_with_config( - uri, - None, - false, - Some(tokio_tungstenite::Connector::NativeTls(connector)), - ) - .await?; - tracing::info!("connected, got {response:?}"); - let (sender, mut receiver) = stream.split(); - let mut sender = MsgSender::new(sender)?; - let (mut tx, rx) = tokio::io::duplex(100_000); - tokio::spawn(async move { - let mut decoder = opus::Decoder::new(24000, opus::Channels::Mono)?; - let mut pr = ogg::reading::async_api::PacketReader::new(rx); - let mut pcm_buf = vec![0f32; 24_000 * 120]; - let mut all_pcms = vec![]; - let mut total_size = 0; - tracing::info!("waiting for audio data"); - while let Some(packet) = pr.next().await { - let packet = packet?; - if packet.data.starts_with(b"OpusHead") || packet.data.starts_with(b"OpusTags") { - continue; - } - let size = decoder.decode_float( - &packet.data, - &mut pcm_buf, - /* Forward Error Correction */ false, - )?; - if size > 0 { - tracing::info!(total_size, size, "received audio"); - let pcm = &pcm_buf[..size]; - total_size += size; - all_pcms.push(pcm.to_vec()); - let mut ad = ad.lock().unwrap(); - ad.push_samples(pcm)?; - } - } - let all_pcms = all_pcms.concat(); - tracing::info!(len = all_pcms.len(), "saving pcms with shape"); - let mut w = std::fs::File::create("received.wav")?; - moshi::wav::write_pcm_as_wav(&mut w, &all_pcms, 24000)?; - Ok::<(), anyhow::Error>(()) - }); - tokio::spawn(async move { - loop { - let input = input_audio.lock().unwrap().take_all(); - if sender.send_pcm(&input).await.is_err() { - break; - }; - tokio::time::sleep(std::time::Duration::from_millis(20)).await - } - }); - while let Some(received) = receiver.next().await { - match received? { - Message::Close(_) => break, - Message::Text(text) => { - tracing::error!("unexpected text message {text}"); - continue; - } - Message::Frame(_) | Message::Ping(_) | Message::Pong(_) => continue, - Message::Binary(bin) => { - if bin.is_empty() { - continue; - } - match bin[0] { - // Handshake - 0 => {} - // Audio - 1 => { - tx.write_all(&bin[1..]).await?; - } - 2 => { - let txt = String::from_utf8_lossy(&bin[1..]); - print!("{txt}"); - std::io::stdout().flush()?; - } - 3 => { - tracing::error!("unsupported control message") - } - 4 => { - tracing::error!("unsupported metadata message") - } - mt => { - tracing::error!("unexpected message type {mt}"); - continue; - } - } - } - }; - } - println!("\n"); - Ok(()) - } -} - -pub mod client_tui { - use super::client::MsgSender; - use anyhow::Result; - use futures_util::stream::StreamExt; - use ratatui::{prelude::*, widgets::*}; - use std::sync::{Arc, Mutex}; - use tokio::io::AsyncWriteExt; - use tokio::sync::mpsc; - use tokio_tungstenite::tungstenite::protocol::Message; - - fn initialize_panic_handler() { - let original_hook = std::panic::take_hook(); - std::panic::set_hook(Box::new(move |panic_info| { - shutdown().unwrap(); - original_hook(panic_info); - })); - } - - fn startup() -> Result<()> { - crossterm::terminal::enable_raw_mode()?; - crossterm::execute!(std::io::stderr(), crossterm::terminal::EnterAlternateScreen)?; - Ok(()) - } - - fn shutdown() -> Result<()> { - crossterm::execute!(std::io::stderr(), crossterm::terminal::LeaveAlternateScreen)?; - crossterm::terminal::disable_raw_mode()?; - Ok(()) - } - - struct Stats { - recv_messages: usize, - recv_text_messages: usize, - recv_audio_messages: usize, - sent_audio_messages: usize, - } - - impl Stats { - fn new() -> Self { - Self { - recv_messages: 0, - recv_text_messages: 0, - recv_audio_messages: 0, - sent_audio_messages: 0, - } - } - } - - struct App { - action_tx: mpsc::UnboundedSender, - should_quit: bool, - ticker: i64, - state: Arc>, - stats: Arc>, - tui_log_state: tui_logger::TuiWidgetState, - input_audio: Arc>, - output_audio: Arc>, - subs: Arc>>, - current_db10: u64, - sender: Arc>, - } - - impl App { - fn current_db10(&mut self) -> u64 { - let db10 = self.input_audio.lock().unwrap().db10(); - if self.current_db10 as f32 + 1.3 < db10 || db10 < self.current_db10 as f32 - 0.3 { - self.current_db10 = db10 as u64 - } - self.current_db10 - } - } - - fn ui(f: &mut Frame, app: &mut App) { - let area = f.size(); - let instructions = block::Title::from(Line::from(vec![ - " Quit ".into(), - " ".yellow().bold(), - " Restart ".into(), - " ".yellow().bold(), - ])); - let state = *app.state.lock().unwrap(); - let block = Block::default() - .title("MoshiMoshi") - .title_alignment(Alignment::Center) - .title(instructions.alignment(Alignment::Center).position(block::Position::Bottom)) - .borders(Borders::ALL) - .border_type(BorderType::Rounded); - let chunks = Layout::default() - .direction(Direction::Vertical) - .constraints([Constraint::Length(8), Constraint::Min(0)]) - .split(block.inner(area)); - - let (stats1, stats2) = { - let input_audio = app.input_audio.lock().unwrap(); - let output_audio = app.output_audio.lock().unwrap(); - let stats = app.stats.lock().unwrap(); - let stats1 = format!( - "msgs: {}\naudio msgs: {}\ntext msgs: {}\nplay len: {} ({:.1}s)\nplay buf: {} ({:.1}s)\n", - stats.recv_messages, stats.recv_audio_messages, stats.recv_text_messages, - output_audio.total_samples(), - output_audio.total_samples() as f32 / 24000., - output_audio.samples_in_buffer(), - output_audio.samples_in_buffer() as f32 / 24000., - ); - let stats2 = format!( - "audio msgs: {}\nrecd len: {} ({:.1}s)\nrecd buf: {} ({:.1}s)", - stats.sent_audio_messages, - input_audio.total_samples(), - input_audio.total_samples() as f32 / 24000., - input_audio.samples_in_buffer(), - input_audio.samples_in_buffer() as f32 / 24000., - ); - (stats1, stats2) - }; - let header_chunks = Layout::default() - .direction(Direction::Horizontal) - .constraints([Constraint::Min(0), Constraint::Length(30), Constraint::Length(30)]) - .split(chunks[0]); - let (header_bg_color, header_fg_color) = match state { - State::Running => { - if app.ticker / 4 % 2 == 0 { - (Color::Black, Color::Red) - } else { - (Color::Black, Color::Green) - } - } - State::Quit => (Color::Red, Color::White), - }; - let state = match state { - State::Running => "\nRUNNING...", - State::Quit => "\nEXITING...", - }; - let header_in_block = Block::default() - .title("state") - .title_alignment(Alignment::Center) - .border_style(Style::default().bg(Color::Black).fg(Color::White)) - .borders(Borders::ALL); - let header_sub_chunks = Layout::default() - .direction(Direction::Vertical) - .constraints([Constraint::Min(0), Constraint::Length(1)]) - .split(header_in_block.inner(header_chunks[0])); - - let header = Paragraph::new(state.bold()) - .style(Style::default().bg(header_bg_color).fg(header_fg_color)) - .alignment(Alignment::Center); - let stats1 = Paragraph::new(stats1) - .block( - Block::default() - .title("received") - .title_alignment(Alignment::Center) - .borders(Borders::ALL), - ) - .style(Style::default().bg(Color::Black).fg(Color::White)) - .alignment(Alignment::Left); - let stats2 = Paragraph::new(stats2) - .block( - Block::default() - .title("sent") - .title_alignment(Alignment::Center) - .borders(Borders::ALL), - ) - .style(Style::default().bg(Color::Black).fg(Color::White)) - .alignment(Alignment::Left); - let bar = BarChart::default() - .bar_width(1) - .direction(Direction::Horizontal) - .bar_style(Style::new().red().on_black()) - .value_style(Style::new().white().bold()) - .label_style(Style::new().white().bold()) - .data(&[("mic", app.current_db10())]) - .max(10); - f.render_widget(block, area); - f.render_widget(header_in_block, header_chunks[0]); - f.render_widget(header, header_sub_chunks[0]); - f.render_widget(bar, header_sub_chunks[1]); - f.render_widget(stats1, header_chunks[1]); - f.render_widget(stats2, header_chunks[2]); - let chunks = Layout::default() - .constraints([Constraint::Percentage(70), Constraint::Percentage(30)]) - .split(chunks[1]); - let subs: String = { - let subs = app.subs.lock().unwrap(); - subs.join("") - }; - let subs = Paragraph::new(subs) - .style(Style::default().bg(Color::Black).fg(Color::White)) - .alignment(Alignment::Left) - .wrap(Wrap { trim: true }); - f.render_widget(subs, chunks[0]); - let footer = tui_logger::TuiLoggerSmartWidget::default() - .style_error(Style::default().fg(Color::Red)) - .style_debug(Style::default().fg(Color::Green)) - .style_warn(Style::default().fg(Color::Yellow)) - .style_trace(Style::default().fg(Color::Magenta)) - .style_info(Style::default().fg(Color::Cyan)) - .output_separator(':') - .output_timestamp(Some("%H:%M:%S".to_string())) - .output_level(Some(tui_logger::TuiLoggerLevelOutput::Abbreviated)) - .output_target(true) - .output_file(true) - .output_line(true) - .state(&app.tui_log_state); - f.render_widget(footer, chunks[1]); - } - - #[derive(Debug, Copy, Clone, PartialEq, Eq)] - enum State { - Running, - Quit, - } - - #[derive(PartialEq)] - enum Action { - None, - Enter, - Space, - Quit, - } - - async fn update(app: &mut App, msg: Action) -> Result<()> { - match msg { - Action::Quit => { - log::info!("exiting"); - *app.state.lock().unwrap() = State::Quit; - app.should_quit = true - } - Action::None => {} - Action::Enter => app.sender.lock().await.send_control(0).await?, - Action::Space => app.sender.lock().await.send_control(1).await?, - }; - Ok(()) - } - - fn handle_event(tx: mpsc::UnboundedSender) -> tokio::task::JoinHandle<()> { - let tick_rate = std::time::Duration::from_millis(250); - tokio::spawn(async move { - loop { - let action = if crossterm::event::poll(tick_rate).unwrap() { - if let crossterm::event::Event::Key(key) = crossterm::event::read().unwrap() { - if key.kind == crossterm::event::KeyEventKind::Press { - if key.modifiers.contains(crossterm::event::KeyModifiers::CONTROL) { - match key.code { - crossterm::event::KeyCode::Char('c' | 'C') => Action::Quit, - _ => Action::None, - } - } else { - match key.code { - crossterm::event::KeyCode::Char('q' | 'Q') => Action::Quit, - crossterm::event::KeyCode::Enter => Action::Enter, - crossterm::event::KeyCode::Char(' ') => Action::Space, - _ => Action::None, - } - } - } else { - Action::None - } - } else { - Action::None - } - } else { - Action::None - }; - if tx.send(action).is_err() { - break; - } - } - }) - } - - pub async fn run(host: String, port: usize) -> Result<()> { - let uri = format!("wss://{host}:{port}/api/chat"); - tracing::info!("connecting to {uri}"); - let subs = Arc::new(Mutex::new(vec![])); - let (_out_stream, output_audio) = crate::audio_io::setup_output_stream(true)?; - let (_in_stream, input_audio) = crate::audio_io::setup_input_stream()?; - let connector = - native_tls::TlsConnector::builder().danger_accept_invalid_certs(true).build()?; - let (stream, response) = tokio_tungstenite::connect_async_tls_with_config( - uri, - None, - false, - Some(tokio_tungstenite::Connector::NativeTls(connector)), - ) - .await?; - tracing::info!("connected, got {response:?}"); - - initialize_panic_handler(); - startup()?; - let mut t = Terminal::new(CrosstermBackend::new(std::io::stderr()))?; - - let (sender, mut receiver) = stream.split(); - let sender = Arc::new(tokio::sync::Mutex::new(MsgSender::new(sender)?)); - let (mut tx, rx) = tokio::io::duplex(100_000); - - let (action_tx, mut action_rx) = mpsc::unbounded_channel(); - let state = Arc::new(Mutex::new(State::Running)); - let stats = Arc::new(Mutex::new(Stats::new())); - let mut app = App { - should_quit: false, - action_tx, - ticker: 0, - state: state.clone(), - input_audio: input_audio.clone(), - output_audio: output_audio.clone(), - tui_log_state: tui_logger::TuiWidgetState::new(), - subs: subs.clone(), - stats: stats.clone(), - current_db10: 0, - sender: sender.clone(), - }; - handle_event(app.action_tx.clone()); - - tokio::spawn({ - let output_audio = output_audio.clone(); - async move { - let mut decoder = opus::Decoder::new(24000, opus::Channels::Mono)?; - let mut pr = ogg::reading::async_api::PacketReader::new(rx); - let mut pcm_buf = vec![0f32; 24_000 * 120]; - let mut all_pcms = vec![]; - tracing::info!("waiting for audio data"); - while let Some(packet) = pr.next().await { - let packet = packet?; - if packet.data.starts_with(b"OpusHead") || packet.data.starts_with(b"OpusTags") - { - continue; - } - let size = decoder.decode_float( - &packet.data, - &mut pcm_buf, - /* Forward Error Correction */ false, - )?; - if size > 0 { - let pcm = &pcm_buf[..size]; - all_pcms.push(pcm.to_vec()); - // TODO: if the buffer is already containing more than x secs of audio, we - // should probably trim it. - output_audio.lock().unwrap().push_samples(pcm)? - } - } - let all_pcms = all_pcms.concat(); - tracing::info!(len = all_pcms.len(), "saving pcms with shape"); - let mut w = std::fs::File::create("received.wav")?; - moshi::wav::write_pcm_as_wav(&mut w, &all_pcms, 24000)?; - Ok::<(), anyhow::Error>(()) - } - }); - tokio::spawn(async move { - loop { - let input = input_audio.lock().unwrap().take_all(); - if sender.lock().await.send_pcm(&input).await.is_err() { - break; - }; - tokio::time::sleep(std::time::Duration::from_millis(20)).await - } - }); - tokio::spawn(async move { - while let Some(received) = receiver.next().await { - match received? { - Message::Close(_) => break, - Message::Text(text) => { - tracing::error!("unexpected text message {text}"); - continue; - } - Message::Frame(_) | Message::Ping(_) | Message::Pong(_) => continue, - Message::Binary(bin) => { - if bin.is_empty() { - continue; - } - match bin[0] { - // Handshake - 0 => {} - // Audio - 1 => { - { - let mut stats = stats.lock().unwrap(); - stats.recv_messages += 1; - stats.recv_audio_messages += 1; - } - tx.write_all(&bin[1..]).await?; - } - 2 => { - { - let mut stats = stats.lock().unwrap(); - stats.recv_messages += 1; - stats.recv_text_messages += 1; - } - let text = String::from_utf8_lossy(&bin[1..]).to_string(); - subs.lock().unwrap().push(text) - } - 3 => { - tracing::error!("unsupported control message") - } - 4 => { - tracing::error!("unsupported metadata message") - } - mt => { - tracing::error!("unexpected message type {mt}"); - } - } - } - }; - } - Ok::<_, anyhow::Error>(()) - }); - - loop { - t.draw(|f| { - ui(f, &mut app); - })?; - if let Some(action) = action_rx.recv().await { - update(&mut app, action).await?; - } - if app.should_quit { - break; - } - app.ticker += 1; - } - - shutdown()?; - Ok(()) - } -} - - - -[package] -name = "moshi-cli" -version.workspace = true -edition.workspace = true -description.workspace = true -repository.workspace = true -keywords.workspace = true -categories.workspace = true -license.workspace = true - -[dependencies] -anyhow = { workspace = true } -byteorder = { workspace = true } -candle = { workspace = true } -candle-nn = { workspace = true } -candle-transformers = { workspace = true } -clap = { workspace = true } -color-eyre = { workspace = true } -cpal = { workspace = true } -crossterm = { workspace = true } -env_logger = { workspace = true } -futures = { workspace = true } -futures-util = { workspace = true } -log = { workspace = true } -moshi = { workspace = true } -native-tls = { workspace = true } -ogg = { workspace = true } -opus = { workspace = true } -rand = { workspace = true } -ratatui = { workspace = true } -rubato = { workspace = true } -rustls = { workspace = true } -sentencepiece = { workspace = true } -serde_json = { workspace = true } -symphonia = { workspace = true } -tokio = { workspace = true } -tokio-tungstenite = { workspace = true } -toml = { workspace = true } -tracing = { workspace = true } -tracing-chrome = { workspace = true } -tracing-subscriber = { workspace = true } -tui-logger = { workspace = true } - -[features] -default = [] -cuda = ["moshi/cuda", "candle/cuda", "candle-nn/cuda", "candle-transformers/cuda"] -metal = ["moshi/metal", "candle/metal", "candle-nn/metal", "candle-transformers/metal"] - -[profile.release] -debug = true - -[profile.release-no-debug] -inherits = "release" -debug = false - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. -use crate::lm::LmModel; -use crate::mimi::Mimi; -use candle::{IndexOp, Result, Tensor}; - -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub enum AsrMsg { - Step { step_idx: usize, prs: Vec> }, - Word { tokens: Vec, start_time: f64, batch_idx: usize }, - EndWord { stop_time: f64, batch_idx: usize }, -} - -#[derive(Debug, Clone)] -pub struct ItemState { - step_idx: usize, - text_token: u32, - word_tokens: Vec, - unended_word: bool, - last_stop_time: f64, - audio_pad_token: u32, - next_codebooks: Vec, -} - -impl ItemState { - fn reset(&mut self) { - self.step_idx = 0; - self.text_token = 0; - self.word_tokens.clear(); - self.unended_word = false; - self.last_stop_time = 0.; - self.next_codebooks.fill(self.audio_pad_token); - } - - pub fn text_token(&self) -> u32 { - self.text_token - } - - pub fn is_first_step(&self) -> bool { - self.step_idx == 0 - } - - pub fn next_token(&mut self, codebook_idx: usize, token: u32) -> u32 { - let v = self.next_codebooks[codebook_idx]; - self.next_codebooks[codebook_idx] = token; - if self.is_first_step() { - self.audio_pad_token - } else { - v - } - } -} - -pub struct State { - asr_delay_in_tokens: usize, - model_step_idx: usize, - temperature: f64, - lm: LmModel, - audio_tokenizer: Mimi, - device: candle::Device, - batch: Vec, -} - -impl State { - pub fn new( - batch_size: usize, - asr_delay_in_tokens: usize, - temperature: f64, - audio_tokenizer: Mimi, - lm: LmModel, - ) -> Result { - let text_token = lm.text_start_token(); - let device = lm.device().clone(); - let item_state = ItemState { - text_token, - word_tokens: vec![], - unended_word: false, - step_idx: 0, - last_stop_time: 0., - audio_pad_token: lm.audio_pad_token(), - next_codebooks: vec![lm.audio_pad_token(); lm.in_audio_codebooks()], - }; - let mut s = Self { - asr_delay_in_tokens, - lm, - model_step_idx: 0, - audio_tokenizer, - temperature, - device, - batch: vec![item_state; batch_size], - }; - s.reset()?; - Ok(s) - } - - pub fn model_step_idx(&self) -> usize { - self.model_step_idx - } - - pub fn device(&self) -> &candle::Device { - &self.device - } - - pub fn batch_size(&self) -> usize { - self.batch.len() - } - - pub fn asr_delay_in_tokens(&self) -> usize { - self.asr_delay_in_tokens - } - - pub fn reset(&mut self) -> Result<()> { - self.lm.reset_state(); - self.audio_tokenizer.reset_state(); - self.batch.iter_mut().for_each(|s| s.reset()); - Ok(()) - } - - pub fn step_pcm( - &mut self, - pcm: Tensor, - conditions: Option<&crate::conditioner::Condition>, - mask: &crate::StreamMask, - f: F, - ) -> Result> - where - F: Fn(&[ItemState], &Tensor, &[Tensor]), - { - let audio_tokens = self.audio_tokenizer.encode_step(&pcm.into(), mask)?; - if let Some(audio_tokens) = audio_tokens.as_option() { - self.step_tokens(audio_tokens, conditions, mask, f) - } else { - Ok(vec![]) - } - } - - fn text_tokens(&self) -> Result { - let batch_size = self.batch_size(); - let text_start_token = self.lm.text_start_token(); - // We used to have literal 0s for the first asr_delay_in_tokens - 1 steps - // This is not the case anymore. - let dev = self.lm.device(); - let text_tokens = self - .batch - .iter() - .map(|s| if s.is_first_step() { text_start_token } else { s.text_token() }) - .collect::>(); - Tensor::from_vec(text_tokens, (batch_size, 1), dev) - } - - pub fn step_tokens( - &mut self, - audio_tokens: &Tensor, - conditions: Option<&crate::conditioner::Condition>, - mask: &crate::StreamMask, - f: F, - ) -> Result> - where - F: Fn(&[ItemState], &Tensor, &[Tensor]), - { - let (batch_size, codebooks, steps) = audio_tokens.dims3()?; - if batch_size != self.batch_size() { - candle::bail!("batch size mismatch: {batch_size} != {}", self.batch_size()); - } - let mut words = vec![]; - for step in 0..steps { - let audio_tokens = audio_tokens.narrow(2, step, 1)?; - let audio_tokens = audio_tokens.reshape((batch_size, codebooks))?.to_vec2::()?; - let audio_tokens = (0..codebooks) - .map(|codebook_idx| { - let audio_tokens = audio_tokens - .iter() - .zip(self.batch.iter_mut()) - .enumerate() - .map(|(batch_idx, (audio_token, item))| { - if !mask.is_active(batch_idx) { - 0 - } else { - item.next_token(codebook_idx, audio_token[codebook_idx]) - } - }) - .collect(); - let audio_tokens = - Tensor::from_vec(audio_tokens, (batch_size, 1), self.device())?; - Ok(audio_tokens) - }) - .collect::>>()?; - let text = self.text_tokens()?; - f(self.batch.as_slice(), &text, &audio_tokens); - let audio_tokens = audio_tokens.into_iter().map(Some).collect::>(); - let (text_logits, transformer_out) = - self.lm.forward_cond(Some(text), audio_tokens, conditions, mask)?; - self.model_step_idx += 1; - let extra_heads = self.lm.extra_heads(&transformer_out)?; - let mut prs = vec![]; - for extra_head in extra_heads.iter() { - // Only retrieve the first element for each extra-head. - let prs_ = - candle_nn::ops::softmax_last_dim(&extra_head.to_dtype(candle::DType::F32)?)? - .i((.., 0, 0))? - .to_vec1::()?; - prs.push(prs_); - } - if !prs.is_empty() { - words.push(AsrMsg::Step { step_idx: self.model_step_idx(), prs }); - } - - let text_tokens = if self.temperature <= 0.0 { - text_logits.i((.., 0))?.argmax(candle::D::Minus1)? - } else { - candle_nn::sampling::gumbel_softmax( - &text_logits.i((.., 0))?.to_dtype(candle::DType::F32)?, - self.temperature, - candle::D::Minus1, - )? - }; - let text_tokens = text_tokens.to_vec1::()?; - for (batch_idx, (text_token, item)) in - text_tokens.into_iter().zip(self.batch.iter_mut()).enumerate() - { - if !mask.is_active(batch_idx) { - continue; - } - item.text_token = text_token; - item.step_idx += 1; - if item.step_idx >= self.asr_delay_in_tokens { - if text_token == 3 || text_token == 0 { - if !item.word_tokens.is_empty() { - let mut tokens = vec![]; - std::mem::swap(&mut item.word_tokens, &mut tokens); - words.push(AsrMsg::Word { - tokens, - start_time: item.last_stop_time, - batch_idx, - }); - item.unended_word = true; - } - } else { - item.word_tokens.push(item.text_token) - } - if item.text_token == 0 { - let stop_time = (item.step_idx - self.asr_delay_in_tokens) as f64 / 12.5; - if item.unended_word { - item.unended_word = false; - words.push(AsrMsg::EndWord { stop_time, batch_idx }); - } - item.last_stop_time = stop_time; - } - } - } - } - Ok(words) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize) -> Result<()> { - if batch_idx >= self.batch_size() { - candle::bail!("batch index out of range: {batch_idx} >= {}", self.batch_size()); - } - self.batch[batch_idx].reset(); - self.lm.reset_batch_idx(batch_idx, self.batch_size())?; - self.audio_tokenizer.reset_batch_idx(batch_idx, self.batch_size())?; - Ok(()) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. -use crate::nn::{ - linear, linear_from, matmul_dtype, MaybeQuantizedLinear, MaybeQuantizedVarBuilder, -}; -use crate::streaming::{StreamMask, StreamTensor, StreamingModule}; -use candle::{IndexOp, Module, Result, Tensor}; - -use crate::kv_cache::{ - IndicesAndMask, ScatteredCacheBuilder as KvCacheBuilder, ScatteredKvCache as KvCache, -}; - -use crate::transformer::{ - CaSrc, Config, LayerScale, PositionalEmbedding, Rope, RotaryEmbedding, - StreamingMultiheadCrossAttention, -}; - -#[derive(Debug, Clone)] -pub struct StreamingMultiheadAttention { - // Self-attention with KV Cache - in_proj: MaybeQuantizedLinear, - out_proj: MaybeQuantizedLinear, - kv_repeat: usize, - num_heads: usize, - context: usize, - kv_cache: KvCache, - span: tracing::Span, -} - -impl StreamingMultiheadAttention { - pub fn new( - cfg: &Config, - builder: &KvCacheBuilder, - vb: MaybeQuantizedVarBuilder, - ) -> Result { - let embed_dim = cfg.d_model; - let head_dim = embed_dim / cfg.num_heads; - let num_kv = cfg.num_heads / cfg.kv_repeat; - let out_dim = embed_dim + 2 * num_kv * (embed_dim / cfg.num_heads); - let in_proj_weight = vb.get((out_dim, embed_dim), "in_proj_weight")?; - let in_proj_bias = - if cfg.bias_attn { Some(vb.get_unquantized(out_dim, "in_proj_bias")?) } else { None }; - let in_proj = linear_from(in_proj_weight, in_proj_bias)?; - let out_proj = linear(embed_dim, embed_dim, cfg.bias_attn, vb.pp("out_proj"))?; - Ok(Self { - in_proj, - out_proj, - kv_repeat: cfg.kv_repeat, - num_heads: cfg.num_heads, - context: cfg.context, - kv_cache: builder.make_cache(num_kv, head_dim)?, - span: tracing::span!(tracing::Level::TRACE, "mha"), - }) - } - - pub fn is_quantized(&self) -> bool { - match self.in_proj { - MaybeQuantizedLinear::Quantized(_) => true, - MaybeQuantizedLinear::Real(_) => false, - } - } - - pub fn forward( - &mut self, - xs: &Tensor, - rope: Option<&Rope>, - iam: &IndicesAndMask, - ) -> Result { - let _enter = self.span.enter(); - if self.kv_repeat != 1 { - candle::bail!("only kv-repeat = 1 is supported") - } - let (b, t, hd) = xs.dims3()?; - let head_dim = hd / self.num_heads; - // time_dim = 1, layout: b,t,h,d - let qkv = xs.apply(&self.in_proj)?.reshape((b, t, 3, self.num_heads, head_dim))?; - let original_dtype = qkv.dtype(); - let qkv = if self.is_quantized() { qkv.to_dtype(matmul_dtype(xs.device()))? } else { qkv }; - let q = qkv.i((.., .., 0))?; - let k = qkv.i((.., .., 1))?; - let v = qkv.i((.., .., 2))?; - // qk_layer_norm = None - // kv_repeat = 1, otherwise we would need repeat_kv - let mut q = q.transpose(1, 2)?.contiguous()?; // b,h,t,d - let mut k = k.transpose(1, 2)?.contiguous()?; // b,h,k,d - let v = v.transpose(1, 2)?.contiguous()?; // b,h,k,d - if let Some(rope) = rope.as_ref() { - q = rope.apply_rotary_emb(&q)?; - k = rope.apply_rotary_emb(&k)?; - } - - let (k, v) = { self.kv_cache.append(&k.contiguous()?, &v.contiguous()?, iam)? }; - // The KV cache keeps all the data at the moment, we want to trim - // down the part that comes from the cache to at most context to - // be coherent with the mask shape we provide. - let k_len = k.dim(2)?; - let k_target_len = t + usize::min(self.context, k_len - t); - let (k, v) = if k_target_len < k_len { - let k = k.narrow(2, k_len - k_target_len, k_target_len)?; - let v = v.narrow(2, k_len - k_target_len, k_target_len)?; - (k, v) - } else { - (k.clone(), v.clone()) - }; - - let xs = { - let pre_ws = q.matmul(&k.t()?)?; // b,h,t,k - let pre_ws = (pre_ws * (head_dim as f64).powf(-0.5))?; - let pre_ws = pre_ws.broadcast_add(iam.mask())?; - let ws = candle_nn::ops::softmax_last_dim(&pre_ws)?; // b,h,t,k - ws.matmul(&v)? // b,h,t,d - }; - - let xs = xs - .transpose(1, 2)? // b,t,h,d - .reshape((b, t, hd))? - .to_dtype(original_dtype)? - .apply(&self.out_proj)?; - Ok(xs) - } - - pub fn set_kv_cache(&mut self, kv_cache: KvCache) { - self.kv_cache = kv_cache - } -} - -#[derive(Debug, Clone)] -pub enum Mlp { - //Feed Forward layers - NoGating { - linear1: MaybeQuantizedLinear, - linear2: MaybeQuantizedLinear, - }, - Gating { - linear_in: MaybeQuantizedLinear, - linear_out: MaybeQuantizedLinear, - activation: candle_nn::Activation, - }, -} - -impl Mlp { - pub fn new(cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - let d_model = cfg.d_model; - match cfg.gating { - None => { - let linear1 = linear(d_model, cfg.dim_feedforward, cfg.bias_ff, vb.pp("linear1"))?; - let linear2 = linear(cfg.dim_feedforward, d_model, cfg.bias_ff, vb.pp("linear2"))?; - Ok(Self::NoGating { linear1, linear2 }) - } - Some(activation) => { - let vb = vb.pp("gating"); - let hidden = if cfg.dim_feedforward == 4 * d_model { - 11 * d_model / 4 - } else { - 2 * cfg.dim_feedforward / 3 - }; - let linear_in = linear(d_model, 2 * hidden, cfg.bias_ff, vb.pp("linear_in"))?; - let linear_out = linear(hidden, d_model, cfg.bias_ff, vb.pp("linear_out"))?; - Ok(Self::Gating { linear_in, linear_out, activation }) - } - } - } -} - -impl Module for Mlp { - fn forward(&self, xs: &Tensor) -> Result { - match self { - Self::NoGating { linear1, linear2 } => xs.apply(linear1)?.gelu_erf()?.apply(linear2), - Self::Gating { linear_in, linear_out, activation } => { - let xs = xs.apply(linear_in)?; - let (b, t, _) = xs.dims3()?; - let xs = xs.reshape((b, t, 2, ()))?; - let xs = (xs.i((.., .., 0))?.apply(activation)? * xs.i((.., .., 1))?)?; - xs.apply(linear_out) - } - } - } -} - -#[derive(Debug, Clone)] -pub struct RmsNorm { - pub(crate) alpha: Tensor, - pub(crate) eps: f32, -} - -impl RmsNorm { - pub fn new(d_model: usize, eps: f32, vb: MaybeQuantizedVarBuilder) -> Result { - let alpha = vb.get_unquantized((1, 1, d_model), "alpha")?.reshape(d_model)?; - Ok(Self { alpha, eps }) - } -} - -impl Module for RmsNorm { - fn forward(&self, xs: &Tensor) -> Result { - candle_nn::ops::rms_norm(xs, &self.alpha, self.eps) - } -} - -#[derive(Debug, Clone)] -pub struct LayerNorm { - inner: candle_nn::LayerNorm, -} - -impl LayerNorm { - pub fn new(d_model: usize, eps: f32, vb: MaybeQuantizedVarBuilder) -> Result { - let bias = vb.get_unquantized(d_model, "bias")?; - let alpha = if vb.contains_key("alpha") { - vb.get_unquantized((1, 1, d_model), "alpha")?.reshape(d_model)? - } else { - vb.get_unquantized(d_model, "weight")?.reshape(d_model)? - }; - let inner = candle_nn::LayerNorm::new(alpha, bias, eps as f64); - Ok(Self { inner }) - } -} - -impl Module for LayerNorm { - fn forward(&self, xs: &Tensor) -> Result { - self.inner.forward(xs) - } -} - -#[derive(Debug, Clone)] -pub enum Norm { - LayerNorm(LayerNorm), - RmsNorm(RmsNorm), -} - -impl Norm { - pub fn new(d_model: usize, cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - let norm = Self::new_shortcut(d_model, cfg.norm, vb)?; - Ok(norm) - } - - pub fn new_shortcut( - d_model: usize, - typ: crate::NormType, - vb: MaybeQuantizedVarBuilder, - ) -> Result { - let norm = match typ { - crate::NormType::LayerNorm => { - let norm = LayerNorm::new(d_model, 1e-5, vb)?; - Self::LayerNorm(norm) - } - crate::NormType::RmsNorm => { - let norm = RmsNorm::new(d_model, 1e-8, vb)?; - Self::RmsNorm(norm) - } - }; - Ok(norm) - } -} - -impl Module for Norm { - fn forward(&self, xs: &Tensor) -> Result { - match self { - Self::LayerNorm(m) => m.forward(xs), - Self::RmsNorm(m) => m.forward(xs), - } - } -} - -#[derive(Debug, Clone)] -pub struct StreamingTransformerLayer { - self_attn: StreamingMultiheadAttention, - mlp: Mlp, - norm1: Norm, - norm2: Norm, - layer_scale_1: Option, - layer_scale_2: Option, - cross_attn: Option<(Norm, StreamingMultiheadCrossAttention)>, - norm_first: bool, - span: tracing::Span, -} - -impl StreamingTransformerLayer { - pub fn new( - cfg: &Config, - builder: &KvCacheBuilder, - vb: MaybeQuantizedVarBuilder, - shared_ca_vb: Option, - ) -> Result { - if cfg.use_conv_block { - candle::bail!("conv-block is not supported") - } - let d_model = cfg.d_model; - let mlp = Mlp::new(cfg, vb.clone())?; - let norm1 = Norm::new(d_model, cfg, vb.pp("norm1"))?; - let norm2 = Norm::new(d_model, cfg, vb.pp("norm2"))?; - let layer_scale_1 = match cfg.layer_scale { - None => None, - Some(ls) => { - let ls = LayerScale::new(d_model, ls, vb.pp("layer_scale_1"))?; - Some(ls) - } - }; - let layer_scale_2 = match cfg.layer_scale { - None => None, - Some(ls) => { - let ls = LayerScale::new(d_model, ls, vb.pp("layer_scale_2"))?; - Some(ls) - } - }; - let self_attn = StreamingMultiheadAttention::new(cfg, builder, vb.pp("self_attn"))?; - let cross_attn = match cfg.cross_attention.map(|v| v.1) { - Some(norm_type) => { - let norm_cross = Norm::new_shortcut(d_model, norm_type, vb.pp("norm_cross"))?; - let cross_attn = match shared_ca_vb { - None => { - StreamingMultiheadCrossAttention::new(cfg, vb.pp("cross_attention"), None)? - } - Some(shared_vb) => StreamingMultiheadCrossAttention::new( - cfg, - shared_vb.pp("cross_attention"), - Some(vb.pp("cross_attention.gate")), - )?, - }; - Some((norm_cross, cross_attn)) - } - None => None, - }; - Ok(Self { - self_attn, - mlp, - norm1, - norm2, - layer_scale_1, - layer_scale_2, - cross_attn, - norm_first: cfg.norm_first, - span: tracing::span!(tracing::Level::TRACE, "transformer-layer"), - }) - } - - pub fn forward( - &mut self, - xs: &Tensor, - rope: Option<&Rope>, - ca_src: Option<&CaSrc>, - iam: &IndicesAndMask, - ) -> Result { - let _enter = self.span.enter(); - if !self.norm_first { - candle::bail!("only norm_first = true is supported") - } - let norm1 = xs.apply(&self.norm1)?; - let xs = (xs - + self.self_attn.forward(&norm1, rope, iam)?.apply(&self.layer_scale_1.as_ref())?)?; - - let xs = match (self.cross_attn.as_mut(), ca_src) { - (Some((norm_cross, cross_attn)), Some(ca_src)) => { - let residual = &xs; - let xs = xs.apply(norm_cross)?; - (residual + cross_attn.forward(&xs, ca_src, None)?)? - } - _ => xs, - }; - - let xs = - (&xs + xs.apply(&self.norm2)?.apply(&self.mlp)?.apply(&self.layer_scale_2.as_ref()))?; - Ok(xs) - } - - pub fn set_kv_cache(&mut self, kv_cache: KvCache) { - self.self_attn.set_kv_cache(kv_cache); - } -} - -#[derive(Debug, Clone)] -pub struct StreamingTransformer { - // Main transformer - layers: Vec, - positional_embedding: PositionalEmbedding, - causal: bool, - builder: KvCacheBuilder, - rope: Option, -} - -impl StreamingTransformer { - pub fn new(batch_size: usize, cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - let vb_l = vb.pp("layers"); - let rope = match cfg.positional_embedding { - PositionalEmbedding::Rope => { - let rope = RotaryEmbedding::new( - cfg.d_model / cfg.num_heads, - cfg.max_period as f32, - vb.device(), - )?; - Some(rope) - } - PositionalEmbedding::None | PositionalEmbedding::Sin => None, - }; - let mut layers = Vec::with_capacity(cfg.num_layers); - let builder = KvCacheBuilder::new(batch_size, cfg.context, vb.dtype(), vb.device())?; - for layer_idx in 0..cfg.num_layers { - // Also send weights of first layer as only it contains the KQV proj weights - // for shared cross-attention layers - let shared_vb = if cfg.shared_cross_attn { Some(vb_l.pp(0)) } else { None }; - let layer = - StreamingTransformerLayer::new(cfg, &builder, vb_l.pp(layer_idx), shared_vb)?; - layers.push(layer) - } - Ok(Self { - layers, - positional_embedding: cfg.positional_embedding, - causal: cfg.causal, - builder, - rope, - }) - } - - pub fn forward(&mut self, xs: &Tensor, m: &StreamMask) -> Result { - self.forward_ca(xs, None, m) - } - - pub fn batch_size(&self) -> usize { - self.builder.batch_size() - } - - fn positions(&self) -> &[usize] { - self.builder.positions() - } - - pub fn forward_ca( - &mut self, - xs: &Tensor, - ca_src: Option<&CaSrc>, - m: &StreamMask, - ) -> Result { - let (b, t, _c) = xs.dims3()?; - if b != self.batch_size() { - candle::bail!("unexpected batch size {b} != {}", self.batch_size()) - } - if !self.causal { - candle::bail!("only causal mode is supported") - } - let iam = match m.cpu() { - None => candle::bail!("batched-transformer expects a mask"), - Some(m) => self.builder.indices_and_mask(t, m)?, - }; - let rope = match self.rope { - Some(ref rope) => { - let pos = self - .positions() - .iter() - .map(|&v| (0..t).map(|i| (v + i) as u32).collect::>()) - .collect::>(); - let pos = Tensor::new(pos, xs.device())?; - Some(rope.rope(&pos)?) - } - None => None, - }; - let mut xs = match self.positional_embedding { - PositionalEmbedding::Rope | PositionalEmbedding::None => xs.clone(), - PositionalEmbedding::Sin => candle::bail!("sin positional embedding is not supported"), - }; - for layer in self.layers.iter_mut() { - xs = layer.forward(&xs, rope.as_ref(), ca_src, &iam)? - } - Ok(xs) - } - - pub fn maybe_precompute_ca_kv(&self, ca_src: Option) -> Result> { - let ca_src = match ca_src { - None => None, - Some(CaSrc::KeysValues(_)) => ca_src, - Some(tokens) => { - if self.layers.is_empty() { - Some(tokens) - } else { - match &self.layers[0].cross_attn { - None => Some(tokens), - Some((_, ca_module)) => { - let (k, v) = ca_module.compute_kv(&tokens)?; - Some(CaSrc::KeysValues((k, v))) - } - } - } - } - }; - Ok(ca_src) - } - - pub fn copy_state(&mut self, from: &Self) -> Result<()> { - if self.layers.len() != from.layers.len() { - candle::bail!("cannot copy kv-caches as the transformers have different depths") - } - self.layers - .iter_mut() - .zip(from.layers.iter()) - .for_each(|(v, w)| v.set_kv_cache(w.self_attn.kv_cache.clone())); - Ok(()) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize) -> Result<()> { - if batch_idx >= self.batch_size() { - candle::bail!("batch_idx {batch_idx} is out of bounds for last_reset_pos") - } - self.builder.reset_batch_index(batch_idx); - Ok(()) - } -} - -impl StreamingModule for StreamingTransformer { - fn reset_state(&mut self) { - self.builder.reset(); - } - - fn step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - match xs.as_option() { - None => Ok(StreamTensor::empty()), - Some(xs) => Ok(StreamTensor::from_tensor(self.forward(xs, m)?)), - } - } -} - -#[derive(Debug, Clone)] -pub struct ProjectedTransformer { - // Projected transformer with unquantized projection - transformer: StreamingTransformer, - input_proj: Option, - output_projs: Vec>, - conv_layout: bool, - span: tracing::Span, -} - -impl ProjectedTransformer { - pub fn new( - input_dim: usize, - output_dims: &[usize], - batch_size: usize, - cfg: &Config, - vb: MaybeQuantizedVarBuilder, - ) -> Result { - let transformer = StreamingTransformer::new(batch_size, cfg, vb.pp("transformer"))?; - let input_proj = if input_dim == cfg.d_model { - None - } else { - let l = linear(input_dim, cfg.d_model, false, vb.pp("input_proj"))?; - Some(l) - }; - let mut output_projs = Vec::with_capacity(output_dims.len()); - let vb_o = vb.pp("output_projs"); - for (i, &output_dim) in output_dims.iter().enumerate() { - let output_proj = if output_dim == cfg.d_model { - None - } else { - let l = linear(cfg.d_model, output_dim, false, vb_o.pp(i))?; - Some(l) - }; - output_projs.push(output_proj) - } - Ok(Self { - transformer, - input_proj, - output_projs, - conv_layout: cfg.conv_layout, - span: tracing::span!(tracing::Level::TRACE, "proj-transformer"), - }) - } - - pub fn forward(&mut self, xs: &Tensor, m: &StreamMask) -> Result> { - let _enter = self.span.enter(); - let xs = if self.conv_layout { xs.transpose(1, 2)? } else { xs.clone() }; - let xs = xs.apply(&self.input_proj.as_ref())?; - let xs = self.transformer.forward(&xs, m)?; - let mut ys = Vec::with_capacity(self.output_projs.len()); - for output_proj in self.output_projs.iter() { - let ys_ = xs.apply(&output_proj.as_ref())?; - let ys_ = if self.conv_layout { ys_.transpose(1, 2)? } else { ys_ }; - ys.push(ys_) - } - Ok(ys) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize) -> Result<()> { - self.transformer.reset_batch_idx(batch_idx) - } -} - -impl StreamingModule for ProjectedTransformer { - fn reset_state(&mut self) { - self.transformer.reset_state() - } - - fn step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - let xs = xs.apply(&|x: &Tensor| { - if self.conv_layout { - x.transpose(1, 2) - } else { - Ok(x.clone()) - } - })?; - let xs = xs.apply(&self.input_proj.as_ref())?; - let xs = self.transformer.step(&xs, m)?; - let ys = xs.apply(&self.output_projs[0].as_ref())?; - ys.apply(&|y: &Tensor| { - if self.conv_layout { - y.transpose(1, 2) - } else { - Ok(y.clone()) - } - }) - } -} - - - -use crate::nn::{ - linear, MaybeQuantizedEmbedding as Embedding, MaybeQuantizedLinear as Linear, - MaybeQuantizedVarBuilder as VarBuilder, -}; -use candle::{DType, Result, Tensor}; -use std::collections::HashMap; - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct LutConfig { - pub n_bins: usize, - pub dim: usize, - pub possible_values: Vec, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct ContinuousAttributeConfig { - pub dim: usize, - pub scale_factor: f32, - pub max_period: f32, -} - -#[derive(Debug, Clone, serde::Deserialize)] -#[serde(tag = "type")] -pub enum ConditionerConfig { - Lut(LutConfig), - ContinuousAttribute(ContinuousAttributeConfig), -} - -pub type Config = HashMap; - -#[derive(Debug, Clone)] -pub struct LutConditioner { - embed: Embedding, - output_proj: Linear, - #[allow(unused)] - learnt_padding: Tensor, - possible_values: HashMap, -} - -impl LutConditioner { - pub fn new(output_dim: usize, cfg: &LutConfig, vb: VarBuilder) -> Result { - let embed = Embedding::new(cfg.n_bins + 1, cfg.dim, vb.pp("embed"))?; - let output_proj = linear(cfg.dim, output_dim, false, vb.pp("output_proj"))?; - let learnt_padding = vb.get_as_tensor((1, 1, output_dim), "learnt_padding")?; - let possible_values: HashMap = - cfg.possible_values.iter().enumerate().map(|(i, v)| (v.to_string(), i)).collect(); - Ok(Self { embed, output_proj, learnt_padding, possible_values }) - } - - pub fn condition(&self, value: &str) -> Result { - let idx = match self.possible_values.get(value) { - None => candle::bail!("unknown value for lut conditioner '{value}'"), - Some(idx) => *idx, - }; - let cond = Tensor::from_vec(vec![idx as u32], (1, 1), self.embed.embeddings().device())? - .apply(&self.embed)? - .apply(&self.output_proj)?; - Ok(Condition::AddToInput(cond)) - } -} - -#[derive(Debug, Clone)] -pub struct ContinuousAttributeConditioner { - scale_factor: f32, - max_period: f32, - dim: usize, - output_proj: Linear, - #[allow(unused)] - learnt_padding: Tensor, - device: candle::Device, -} - -impl ContinuousAttributeConditioner { - pub fn new(output_dim: usize, cfg: &ContinuousAttributeConfig, vb: VarBuilder) -> Result { - let output_proj = linear(cfg.dim, output_dim, false, vb.pp("output_proj"))?; - let learnt_padding = vb.get_as_tensor((1, 1, output_dim), "learnt_padding")?; - Ok(Self { - scale_factor: cfg.scale_factor, - max_period: cfg.max_period, - dim: cfg.dim, - output_proj, - learnt_padding, - device: vb.device().clone(), - }) - } - - // `positions` should have shape (b, t, 1), the output will be (b, t, dim) - pub fn create_sin_embeddings(&self, positions: &Tensor, dtype: DType) -> Result { - let dev = positions.device(); - let half_dim = self.dim / 2; - let positions = positions.to_dtype(dtype)?; - let adim: Vec<_> = (0..half_dim) - .map(|i| 1f32 / self.max_period.powf(i as f32 / (half_dim - 1) as f32)) - .collect(); - let adim = Tensor::from_vec(adim, (1, 1, ()), dev)?; - let freqs = positions.broadcast_mul(&adim)?; - let pos_emb = Tensor::cat(&[freqs.cos()?, freqs.sin()?], candle::D::Minus1)?; - Ok(pos_emb) - } - - // TODO(laurent): should we support different values per batch element? - pub fn condition(&self, value: f32) -> Result { - let value = value * self.scale_factor; - let positions = Tensor::full(value, (1, 1, 1), &self.device)?; - let cond = self - .create_sin_embeddings(&positions, DType::F32)? - .to_dtype(self.output_proj.dtype())? - .apply(&self.output_proj)?; - Ok(Condition::AddToInput(cond)) - } -} - -#[derive(Debug, Clone)] -pub enum Conditioner { - Lut(LutConditioner), - ContinuousAttribute(ContinuousAttributeConditioner), -} - -#[derive(Debug, Clone)] -pub struct ConditionProvider { - conditioners: HashMap, -} - -#[derive(Debug, Clone)] -pub enum Condition { - AddToInput(Tensor), -} - -impl ConditionProvider { - pub fn new(output_dim: usize, cfg: &Config, vb: VarBuilder) -> Result { - let vb = vb.pp("conditioners"); - let mut conditioners = HashMap::new(); - for (conditioner_name, conditioner_cfg) in cfg.iter() { - let vb = vb.pp(conditioner_name); - let conditioner = match conditioner_cfg { - ConditionerConfig::Lut(cfg) => { - Conditioner::Lut(LutConditioner::new(output_dim, cfg, vb)?) - } - ConditionerConfig::ContinuousAttribute(cfg) => Conditioner::ContinuousAttribute( - ContinuousAttributeConditioner::new(output_dim, cfg, vb)?, - ), - }; - conditioners.insert(conditioner_name.to_string(), conditioner); - } - Ok(Self { conditioners }) - } - - pub fn condition_lut(&self, name: &str, value: &str) -> Result { - let lut = match self.conditioners.get(name) { - None => candle::bail!("unknown conditioner {name}"), - Some(Conditioner::Lut(l)) => l, - Some(_) => candle::bail!("cannot use conditioner with a str value {name}"), - }; - let cond = lut.condition(value)?; - Ok(cond) - } - - pub fn condition_cont(&self, name: &str, value: f32) -> Result { - let c = match self.conditioners.get(name) { - None => candle::bail!("unknown conditioner {name}"), - Some(Conditioner::ContinuousAttribute(c)) => c, - Some(_) => candle::bail!("cannot use conditioner with a str value {name}"), - }; - let cond = c.condition(value)?; - Ok(cond) - } - - pub fn learnt_padding(&self, name: &str) -> Result { - let c = match self.conditioners.get(name) { - None => candle::bail!("unknown conditioner {name}"), - Some(Conditioner::ContinuousAttribute(c)) => c.learnt_padding.clone(), - Some(Conditioner::Lut(c)) => c.learnt_padding.clone(), - }; - Ok(Condition::AddToInput(c)) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::streaming::{StreamMask, StreamTensor, StreamingModule}; -use candle::{IndexOp, Module, Result, Tensor, D}; -use candle_nn::{Conv1d, VarBuilder}; - -#[allow(clippy::enum_variant_names)] -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub enum Norm { - WeightNorm, - SpectralNorm, - TimeGroupNorm, -} - -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub enum PadMode { - Constant, - Reflect, - Replicate, -} - -// Applies weight norm for inference by recomputing the weight tensor. This -// does not apply to training. -// https://pytorch.org/docs/stable/generated/torch.nn.utils.weight_norm.html -fn conv1d_weight_norm( - in_c: usize, - out_c: usize, - kernel_size: usize, - bias: bool, - config: candle_nn::Conv1dConfig, - vb: VarBuilder, -) -> Result { - let weight = if vb.contains_tensor("weight") { - vb.get((out_c, in_c, kernel_size), "weight")? - } else { - let weight_g = vb.get((out_c, 1, 1), "weight_g")?; - let weight_v = vb.get((out_c, in_c, kernel_size), "weight_v")?; - let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; - weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)? - }; - let bias = if bias { Some(vb.get(out_c, "bias")?) } else { None }; - Ok(Conv1d::new(weight, bias, config)) -} - -#[derive(Debug, Clone)] -pub struct NormConv1d { - conv: Conv1d, - norm: Option, - span: tracing::Span, -} - -impl NormConv1d { - #[allow(clippy::too_many_arguments)] - pub fn new( - in_c: usize, - out_c: usize, - k_size: usize, - causal: bool, - norm: Option, - bias: bool, - cfg: candle_nn::Conv1dConfig, - vb: VarBuilder, - ) -> Result { - let conv = match norm { - None | Some(Norm::TimeGroupNorm) => { - if bias { - candle_nn::conv1d(in_c, out_c, k_size, cfg, vb.pp("conv"))? - } else { - candle_nn::conv1d_no_bias(in_c, out_c, k_size, cfg, vb.pp("conv"))? - } - } - Some(Norm::WeightNorm) => { - conv1d_weight_norm(in_c, out_c, k_size, bias, cfg, vb.pp("conv"))? - } - Some(Norm::SpectralNorm) => candle::bail!("SpectralNorm is not supported yet."), - }; - let norm = match norm { - None | Some(Norm::WeightNorm) | Some(Norm::SpectralNorm) => None, - Some(Norm::TimeGroupNorm) => { - if causal { - candle::bail!("GroupNorm doesn't support causal evaluation.") - } - let norm = candle_nn::group_norm(1, out_c, 1e-5, vb.pp("norm"))?; - Some(norm) - } - }; - Ok(Self { conv, norm, span: tracing::span!(tracing::Level::TRACE, "norm-conv1d") }) - } -} - -impl Module for NormConv1d { - fn forward(&self, xs: &Tensor) -> Result { - let _enter = self.span.enter(); - let xs = xs.apply(&self.conv)?; - match self.norm.as_ref() { - None => Ok(xs), - Some(norm) => xs.apply(norm), - } - } -} - -#[derive(Debug, Clone)] -pub struct NormConvTranspose1d { - ws: Tensor, - bs: Option, - k_size: usize, - stride: usize, - groups: usize, - norm: Option, - span: tracing::Span, -} - -impl NormConvTranspose1d { - #[allow(clippy::too_many_arguments)] - pub fn new( - in_c: usize, - out_c: usize, - k_size: usize, - causal: bool, - norm: Option, - bias: bool, - stride: usize, - groups: usize, - vb: VarBuilder, - ) -> Result { - let vb = vb.pp("convtr"); - let bs = if bias { Some(vb.get(out_c, "bias")?) } else { None }; - let ws = match norm { - None | Some(Norm::TimeGroupNorm) => vb.get((in_c, out_c / groups, k_size), "weight")?, - Some(Norm::WeightNorm) => { - if vb.contains_tensor("weight") { - vb.get((in_c, out_c, k_size), "weight")? - } else { - let weight_g = vb.get((in_c, 1, 1), "weight_g")?; - let weight_v = vb.get((in_c, out_c, k_size), "weight_v")?; - let norm_v = weight_v.sqr()?.sum_keepdim((1, 2))?.sqrt()?; - weight_v.broadcast_mul(&weight_g)?.broadcast_div(&norm_v)? - } - } - Some(Norm::SpectralNorm) => candle::bail!("SpectralNorm is not supported yet."), - }; - let (ws, groups) = if groups == out_c && in_c == out_c { - let eye = Tensor::eye(out_c, ws.dtype(), ws.device())?; - let ws = ws.repeat((1, out_c, 1))?.mul(&eye.unsqueeze(2)?.repeat((1, 1, k_size))?)?; - (ws, 1) - } else { - (ws, groups) - }; - let norm = match norm { - None | Some(Norm::WeightNorm) | Some(Norm::SpectralNorm) => None, - Some(Norm::TimeGroupNorm) => { - if causal { - candle::bail!("GroupNorm doesn't support causal evaluation.") - } - let norm = candle_nn::group_norm(1, out_c, 1e-5, vb.pp("norm"))?; - Some(norm) - } - }; - Ok(Self { - ws, - bs, - k_size, - stride, - groups, - norm, - span: tracing::span!(tracing::Level::TRACE, "norm-conv-tr1d"), - }) - } -} - -impl Module for NormConvTranspose1d { - fn forward(&self, xs: &Tensor) -> Result { - let _enter = self.span.enter(); - // conv-transpose1d seems to be broken on metal after enough iterations. Causing - // the following error: - // _status < MTLCommandBufferStatusCommitted > - // -[IOGPUMetalCommandBuffer setCurrentCommandEncoder:] - // This is now fixed in candle. - let xs = Tensor::conv_transpose1d(xs, &self.ws, 0, 0, self.stride, 1, self.groups)?; - let xs = match &self.bs { - None => xs, - Some(bias) => { - let b = bias.dims1()?; - let bias = bias.reshape((1, b, 1))?; - xs.broadcast_add(&bias)? - } - }; - match self.norm.as_ref() { - None => Ok(xs), - Some(norm) => xs.apply(norm), - } - } -} - -fn get_extra_padding_for_conv1d( - xs: &Tensor, - k_size: usize, - stride: usize, - padding_total: usize, -) -> Result { - let len = xs.dim(D::Minus1)?; - let n_frames = (len + padding_total).saturating_sub(k_size) as f64 / stride as f64 + 1.0; - let ideal_len = - ((n_frames.ceil() as usize - 1) * stride + k_size).saturating_sub(padding_total); - Ok(ideal_len.saturating_sub(len)) -} - -fn pad1d(xs: &Tensor, pad_l: usize, pad_r: usize, mode: PadMode) -> Result { - match mode { - PadMode::Constant => xs.pad_with_zeros(D::Minus1, pad_l, pad_r), - PadMode::Reflect => candle::bail!("pad-mode 'reflect' is not supported"), - PadMode::Replicate => xs.pad_with_same(D::Minus1, pad_l, pad_r), - } -} - -fn unpad1d(xs: &Tensor, unpad_l: usize, unpad_r: usize) -> Result { - let len = xs.dim(D::Minus1)?; - if len < unpad_l + unpad_r { - candle::bail!("unpad1d: tensor len {len} is too low, {unpad_l} + {unpad_r}") - } - xs.narrow(D::Minus1, unpad_l, len - (unpad_l + unpad_r)) -} - -#[derive(Debug, Clone)] -pub struct StreamableConv1d { - conv: NormConv1d, - causal: bool, - pad_mode: PadMode, - state_prev_xs: StreamTensor, - left_pad_applied: bool, - kernel_size: usize, - span: tracing::Span, -} - -impl StreamableConv1d { - #[allow(clippy::too_many_arguments)] - pub fn new( - in_c: usize, - out_c: usize, - k_size: usize, - stride: usize, - dilation: usize, - groups: usize, - bias: bool, - causal: bool, - norm: Option, - pad_mode: PadMode, - vb: VarBuilder, - ) -> Result { - let cfg = candle_nn::Conv1dConfig { - padding: 0, - stride, - dilation, - groups, - cudnn_fwd_algo: Some(candle::conv::CudnnFwdAlgo::ImplicitGemm), - }; - let conv = NormConv1d::new(in_c, out_c, k_size, causal, norm, bias, cfg, vb.pp("conv"))?; - if k_size < stride { - candle::bail!("kernel-size {k_size} is smaller than stride {stride}") - } - Ok(Self { - conv, - causal, - pad_mode, - state_prev_xs: StreamTensor::empty(), - left_pad_applied: false, - kernel_size: k_size, - span: tracing::span!(tracing::Level::TRACE, "streamable-conv1d"), - }) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, _batch_size: usize) -> Result<()> { - if let Some(v) = self.state_prev_xs.as_option() { - let v = v.contiguous()?; - v.i(batch_idx..(1 + batch_idx))?.zero_set()?; - self.state_prev_xs = StreamTensor::from_tensor(v); - } - Ok(()) - } -} - -impl Module for StreamableConv1d { - fn forward(&self, xs: &Tensor) -> Result { - let _enter = self.span.enter(); - let (_b, _t, _c) = xs.dims3()?; - let k_size = self.conv.conv.weight().dim(D::Minus1)?; - let conv_cfg = self.conv.conv.config(); - // Effective kernel size with dilations. - let k_size = (k_size - 1) * conv_cfg.dilation + 1; - let padding_total = k_size - conv_cfg.stride; - let extra_padding = - get_extra_padding_for_conv1d(xs, k_size, conv_cfg.stride, padding_total)?; - let xs = if self.causal { - pad1d(xs, padding_total, extra_padding, self.pad_mode)? - } else { - let padding_right = padding_total / 2; - let padding_left = padding_total - padding_right; - pad1d(xs, padding_left, padding_right + extra_padding, self.pad_mode)? - }; - xs.apply(&self.conv) - } -} - -impl StreamingModule for StreamableConv1d { - fn reset_state(&mut self) { - self.state_prev_xs.reset(); - self.left_pad_applied = false; - } - - fn step(&mut self, xs: &StreamTensor, mask: &StreamMask) -> Result { - let _enter = self.span.enter(); - let xs = match xs.as_option() { - None => return Ok(().into()), - Some(xs) => xs.clone(), - }; - let xs = if self.left_pad_applied { - xs - } else { - self.left_pad_applied = true; - let k_size = self.conv.conv.weight().dim(D::Minus1)?; - let conv_cfg = self.conv.conv.config(); - let k_size = (k_size - 1) * conv_cfg.dilation + 1; - let padding_total = k_size - conv_cfg.stride; - pad1d(&xs, padding_total, 0, self.pad_mode)? - }; - let cfg = self.conv.conv.config(); - let stride = cfg.stride; - let dilation = cfg.dilation; - let kernel = (self.kernel_size - 1) * dilation + 1; - let xs = StreamTensor::cat2(&self.state_prev_xs, &xs.into(), D::Minus1)?; - let seq_len = xs.seq_len(D::Minus1)?; - let num_frames = (seq_len + stride).saturating_sub(kernel) / stride; - let (state_prev_xs, ys) = if num_frames > 0 { - let offset = num_frames * stride; - let state_prev_xs = xs.narrow(D::Minus1, offset, seq_len - offset)?; - let in_l = (num_frames - 1) * stride + kernel; - let xs = xs.narrow(D::Minus1, 0, in_l)?; - // We apply the underlying convtr directly rather than through forward so as - // not to apply any padding here. - let ys = xs.apply(&self.conv.conv)?; - (state_prev_xs, ys) - } else { - (xs, StreamTensor::empty()) - }; - let state_prev_xs = match mask.as_option() { - None => state_prev_xs, - Some(mask) => match (state_prev_xs.as_option(), self.state_prev_xs.as_option()) { - (None, None) => state_prev_xs, - (Some(state_prev_xs), None) => { - let z = state_prev_xs.zeros_like()?; - let mask = mask.reshape(((), 1, 1))?.broadcast_as(state_prev_xs.shape())?; - mask.where_cond(state_prev_xs, &z)?.into() - } - (None, Some(_)) => { - candle::bail!("streaming conv1d should only be used with constant steps") - } - (Some(prev_xs), Some(prev_prev_xs)) => { - if prev_xs.shape() != prev_prev_xs.shape() { - candle::bail!("streaming conv1d should only be used with constant steps {prev_xs:?} {prev_prev_xs:?}") - } - let mask = mask.reshape(((), 1, 1))?.broadcast_as(prev_xs.shape())?; - mask.where_cond(prev_xs, prev_prev_xs)?.into() - } - }, - }; - self.state_prev_xs = state_prev_xs; - Ok(ys) - } -} - -#[derive(Debug, Clone)] -pub struct StreamableConvTranspose1d { - convtr: NormConvTranspose1d, - causal: bool, - state_prev_ys: StreamTensor, - kernel_size: usize, - span: tracing::Span, -} - -impl StreamableConvTranspose1d { - #[allow(clippy::too_many_arguments)] - pub fn new( - in_c: usize, - out_c: usize, - k_size: usize, - stride: usize, - groups: usize, - bias: bool, - causal: bool, - norm: Option, - vb: VarBuilder, - ) -> Result { - let convtr = NormConvTranspose1d::new( - in_c, - out_c, - k_size, - causal, - norm, - bias, - stride, - groups, - vb.pp("convtr"), - )?; - Ok(Self { - convtr, - causal, - kernel_size: k_size, - state_prev_ys: StreamTensor::empty(), - span: tracing::span!(tracing::Level::TRACE, "streamable-conv-tr1d"), - }) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, _batch_size: usize) -> Result<()> { - if let Some(v) = self.state_prev_ys.as_option() { - let v = v.contiguous()?; - v.i(batch_idx..(1 + batch_idx))?.zero_set()?; - self.state_prev_ys = v.into(); - } - Ok(()) - } -} - -impl Module for StreamableConvTranspose1d { - fn forward(&self, xs: &Tensor) -> Result { - let _enter = self.span.enter(); - let k_size = self.convtr.k_size; - let stride = self.convtr.stride; - let padding_total = k_size.saturating_sub(stride); - let xs = xs.apply(&self.convtr)?; - if self.causal { - // This corresponds to trim_right_ratio = 1. - unpad1d(&xs, 0, padding_total) - } else { - let padding_right = padding_total / 2; - let padding_left = padding_total - padding_right; - unpad1d(&xs, padding_left, padding_right) - } - } -} - -impl StreamingModule for StreamableConvTranspose1d { - fn reset_state(&mut self) { - self.state_prev_ys.reset() - } - - fn step(&mut self, xs: &StreamTensor, mask: &StreamMask) -> Result { - let _enter = self.span.enter(); - let xs = match xs.as_option() { - Some(xs) => xs, - None => return Ok(StreamTensor::empty()), - }; - let stride = self.convtr.stride; - // We apply the underlying convtr directly rather than through forward so as - // not to apply any padding here. - let ys = self.convtr.forward(xs)?; - let ot = ys.dim(D::Minus1)?; - let ys = match self.state_prev_ys.as_option() { - None => ys, - Some(prev_ys) => { - let pt = prev_ys.dim(D::Minus1)?; - // Remove the bias as it will be applied multiple times. - let prev_ys = match &self.convtr.bs { - None => prev_ys.clone(), - Some(bias) => { - let bias = bias.reshape((1, (), 1))?; - prev_ys.broadcast_sub(&bias)? - } - }; - let ys1 = (ys.narrow(D::Minus1, 0, pt)? + prev_ys)?; - let ys2 = ys.narrow(D::Minus1, pt, ot - pt)?; - Tensor::cat(&[ys1, ys2], D::Minus1)? - } - }; - let invalid_steps = self.kernel_size - stride; - let (ys, prev_ys) = StreamTensor::from(ys).split(D::Minus1, ot - invalid_steps)?; - let prev_ys = match mask.as_option() { - None => prev_ys, - Some(mask) => match (prev_ys.as_option(), self.state_prev_ys.as_option()) { - (None, None) => prev_ys, - (Some(prev_ys), None) => { - let z = prev_ys.zeros_like()?; - let mask = mask.reshape(((), 1, 1))?.broadcast_as(prev_ys.shape())?; - mask.where_cond(prev_ys, &z)?.into() - } - (None, Some(_)) => { - candle::bail!("streaming conv-tr1d should only be used with constant steps") - } - (Some(prev_ys), Some(prev_prev_ys)) => { - if prev_ys.shape() != prev_prev_ys.shape() { - candle::bail!("streaming conv-tr1d should only be used with constant steps {prev_ys:?} {prev_prev_ys:?}") - } - let mask = mask.reshape(((), 1, 1))?.broadcast_as(prev_ys.shape())?; - mask.where_cond(prev_ys, prev_prev_ys)?.into() - } - }, - }; - self.state_prev_ys = prev_ys; - Ok(ys) - } -} - -#[derive(Debug, Clone)] -pub struct ConvDownsample1d { - conv: StreamableConv1d, -} - -impl ConvDownsample1d { - pub fn new( - stride: usize, - dim: usize, - causal: bool, - learnt: bool, - vb: VarBuilder, - ) -> Result { - if !learnt { - candle::bail!("only learnt=true is supported") - } - let conv = StreamableConv1d::new( - /* in_c */ dim, - /* out_c */ dim, - /* k_size_c */ 2 * stride, - /* stride */ stride, - /* dilation */ 1, - /* groups */ 1, // channel_wise = false - /* bias */ false, - /* causal */ causal, - /* norm */ None, - /* pad_mode */ PadMode::Replicate, - vb.pp("conv"), - )?; - Ok(Self { conv }) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - self.conv.reset_batch_idx(batch_idx, batch_size) - } -} - -impl Module for ConvDownsample1d { - fn forward(&self, xs: &Tensor) -> Result { - xs.apply(&self.conv) - } -} - -impl StreamingModule for ConvDownsample1d { - fn reset_state(&mut self) { - self.conv.reset_state() - } - - fn step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - self.conv.step(xs, m) - } -} - -#[derive(Debug, Clone)] -pub struct ConvTrUpsample1d { - convtr: StreamableConvTranspose1d, -} - -impl ConvTrUpsample1d { - pub fn new( - stride: usize, - dim: usize, - causal: bool, - learnt: bool, - vb: VarBuilder, - ) -> Result { - if !learnt { - candle::bail!("only learnt=true is supported") - } - let convtr = StreamableConvTranspose1d::new( - dim, - dim, - /* k_size */ 2 * stride, - /* stride */ stride, - /* groups */ dim, - /* bias */ false, - /* causal */ causal, - /* norm */ None, - vb.pp("convtr"), - )?; - Ok(Self { convtr }) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - self.convtr.reset_batch_idx(batch_idx, batch_size) - } -} - -impl Module for ConvTrUpsample1d { - fn forward(&self, xs: &Tensor) -> Result { - xs.apply(&self.convtr) - } -} - -impl StreamingModule for ConvTrUpsample1d { - fn reset_state(&mut self) { - self.convtr.reset_state() - } - - fn step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - self.convtr.step(xs, m) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use candle::IndexOp; - - fn run_conv1d( - k_size: usize, - stride: usize, - dilation: usize, - step_size: usize, - len: usize, - bias: bool, - ) -> Result<()> { - // TODO: We should ensure for the seed to be constant when running these tests. - let dev = &candle::Device::Cpu; - let vm = candle_nn::VarMap::new(); - let vb = VarBuilder::from_varmap(&vm, candle::DType::F32, dev); - let conv1d = StreamableConv1d::new( - /* in_c */ 2, - /* out_c */ 3, - /* k_size */ k_size, - /* stride */ stride, - /* dilation */ dilation, - /* groups */ 1, - /* bias */ bias, - /* causal */ true, - /* norm */ None, - /* pad_mode */ PadMode::Constant, - vb, - )?; - let xs = Tensor::randn(0f32, 1., (1, 2, step_size * len), dev)?; - let ys = conv1d.forward(&xs)?; - let mut conv1d = conv1d; - let mut ys_steps = vec![]; - for idx in 0..len { - let xs = xs.i((.., .., step_size * idx..step_size * (idx + 1)))?; - let ys = conv1d.step(&xs.into(), &().into())?; - if let Some(ys) = ys.as_option() { - ys_steps.push(ys.clone()) - } - } - let ys_steps = Tensor::cat(&ys_steps, D::Minus1)?; - let diff = (&ys - &ys_steps)?.abs()?.flatten_all()?.max(0)?.to_vec0::()?; - if diff > 1e-5 { - println!("{xs}"); - println!("{ys}"); - println!("{ys_steps}"); - candle::bail!("larger diff than expected {diff}") - } - Ok(()) - } - - fn run_conv_tr1d( - k_size: usize, - stride: usize, - step_size: usize, - len: usize, - bias: bool, - ) -> Result<()> { - // TODO: We should ensure for the seed to be constant when running these tests. - let dev = &candle::Device::Cpu; - let vm = candle_nn::VarMap::new(); - let vb = VarBuilder::from_varmap(&vm, candle::DType::F32, dev); - let conv1d = StreamableConvTranspose1d::new( - /* in_c */ 2, /* out_c */ 3, /* k_size */ k_size, - /* stride */ stride, /* groups */ 1, /* bias */ bias, - /* causal */ true, /* norm */ None, vb, - )?; - let xs = Tensor::randn(0f32, 1., (1, 2, step_size * len), dev)?; - let ys = conv1d.forward(&xs)?; - let mut conv1d = conv1d; - let mut ys_steps = vec![]; - for idx in 0..len { - let xs = xs.i((.., .., step_size * idx..step_size * (idx + 1)))?; - let ys = conv1d.step(&xs.into(), &().into())?; - if let Some(ys) = ys.as_option() { - ys_steps.push(ys.clone()) - } - } - let ys_steps = Tensor::cat(&ys_steps, D::Minus1)?; - let diff = (&ys - &ys_steps)?.abs()?.flatten_all()?.max(0)?.to_vec0::()?; - if diff > 1e-5 { - println!("{xs}"); - println!("{ys}"); - println!("{ys_steps}"); - candle::bail!("larger diff than expected {diff}") - } - Ok(()) - } - - #[test] - fn conv1d() -> Result<()> { - for step_size in [1, 2, 3] { - for bias in [false, true] { - run_conv1d(1, 1, 1, step_size, 5, bias)?; - run_conv1d(2, 1, 1, step_size, 5, bias)?; - run_conv1d(2, 2, 1, step_size, 6, bias)?; - run_conv1d(3, 2, 1, step_size, 8, bias)?; - run_conv1d(3, 2, 2, step_size, 8, bias)?; - } - } - Ok(()) - } - - #[test] - fn conv_tr1d() -> Result<()> { - for step_size in [1, 2, 3] { - for bias in [false, true] { - run_conv_tr1d(1, 1, step_size, 5, bias)?; - run_conv_tr1d(2, 1, step_size, 5, bias)?; - run_conv_tr1d(3, 1, step_size, 5, bias)?; - run_conv_tr1d(3, 2, step_size, 5, bias)?; - } - } - Ok(()) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use candle::{DType, Device, Result, Tensor}; -use candle_nn::kv_cache::RotatingKvCache; - -#[derive(Debug, Clone)] -pub struct IndicesAndMask { - indices: Tensor, - mask: Tensor, -} - -impl IndicesAndMask { - pub fn mask(&self) -> &Tensor { - &self.mask - } -} - -#[derive(Debug, Clone)] -pub struct ScatteredKvCache { - k: Tensor, - v: Tensor, - context: usize, -} - -impl ScatteredKvCache { - pub fn append( - &mut self, - k: &Tensor, - v: &Tensor, - iam: &IndicesAndMask, - ) -> Result<(Tensor, Tensor)> { - if self.context <= k.dim(2)? { - return Ok((k.clone(), v.clone())); - } - let indices = iam.indices.unsqueeze(2)?.unsqueeze(1)?; - let indices = indices.broadcast_as(k.shape())?.contiguous()?; - self.k.scatter_set(&indices, k, 2)?; - self.v.scatter_set(&indices, v, 2)?; - Ok((self.k.clone(), self.v.clone())) - } - - pub fn k(&self) -> &Tensor { - &self.k - } - - pub fn v(&self) -> &Tensor { - &self.v - } -} - -#[derive(Debug, Clone)] -pub struct ScatteredCacheBuilder { - context: usize, - // The current position in the stream, this can be larger than context. - positions: Vec, - // The index where the next element will be stored. - indices: Vec, - dtype: DType, - device: Device, -} - -impl ScatteredCacheBuilder { - pub fn new(batch_size: usize, context: usize, dtype: DType, device: &Device) -> Result { - let positions = vec![0; batch_size]; - let indices = vec![0; batch_size]; - Ok(Self { positions, indices, context, dtype, device: device.clone() }) - } - - pub fn make_cache(&self, num_heads: usize, head_dim: usize) -> Result { - let batch_size = self.batch_size(); - let shape = (batch_size, num_heads, self.context, head_dim); - let k = Tensor::zeros(shape, self.dtype, self.device())?; - let v = Tensor::zeros(shape, self.dtype, self.device())?; - Ok(ScatteredKvCache { k, v, context: self.context }) - } - - pub fn positions(&self) -> &[usize] { - &self.positions - } - - pub fn reset(&mut self) { - self.positions.fill(0); - self.indices.fill(0); - } - - pub fn batch_size(&self) -> usize { - self.positions.len() - } - - pub fn reset_batch_index(&mut self, batch_index: usize) { - self.positions[batch_index] = 0; - self.indices[batch_index] = 0; - } - - #[allow(clippy::needless_range_loop)] - pub fn indices_and_mask( - &mut self, - seq_len: usize, - batch_mask: &[bool], - ) -> Result { - // mask shape is (b, h, t, k) - let context = self.context; - if self.context <= seq_len { - return self.indices_and_mask_abs(seq_len, batch_mask); - } - let mut attention_masks = Vec::with_capacity(self.batch_size()); - let mut cache_indices = Vec::with_capacity(self.batch_size()); - for (batch_i, &batch_mask) in batch_mask.iter().enumerate() { - if !batch_mask { - let masks: Vec> = vec![vec![0.0; context]; seq_len]; - let indices = vec![self.indices[batch_i] as u32; seq_len]; - attention_masks.push(masks); - cache_indices.push(indices); - } else { - let start_index = self.indices[batch_i]; - let start_pos = self.positions[batch_i]; - let mut masks: Vec> = Vec::with_capacity(seq_len); - let mut indices = Vec::with_capacity(seq_len); - let mut all_pos = vec![usize::MAX; context]; - if start_pos < context { - for i in 0..start_pos { - all_pos[i] = i; - } - } else { - let offset = start_pos - start_index; - for i in 0..context { - all_pos[i] = - if i < start_index { i + offset } else { i + offset - context }; - } - } - for seq_i in 0..seq_len { - let index = self.indices[batch_i]; - all_pos[index] = seq_i + start_pos; - indices.push(index as u32); - self.indices[batch_i] += 1; - self.positions[batch_i] += 1; - if self.indices[batch_i] >= self.context { - self.indices[batch_i] = 0; - } - } - - for seq_i in 0..seq_len { - let my_pos = seq_i + start_pos; - let mask = all_pos - .iter() - .map(|&pos| if pos <= my_pos { 0.0 } else { f32::NEG_INFINITY }) - .collect::>(); - masks.push(mask); - } - - attention_masks.push(masks); - cache_indices.push(indices); - } - } - // Flattening the attention mask then using Tensor::from_vec rather using Tensor::new ends - // up being almost 10x faster with candle 0.9.0. The slowness seems to be on the CPU - // copies, to be further investigated. - let attention_masks = - attention_masks.into_iter().flat_map(|m| m.into_iter().flatten()).collect::>(); - let mask = Tensor::from_vec(attention_masks, ((), 1, seq_len, context), self.device())? - .to_dtype(self.dtype)?; - let indices = Tensor::new(cache_indices, self.device())?; - Ok(IndicesAndMask { indices, mask }) - } - - pub fn device(&self) -> &Device { - &self.device - } - - #[allow(clippy::needless_range_loop)] - fn indices_and_mask_abs( - &mut self, - seq_len: usize, - batch_mask: &[bool], - ) -> Result { - let mask = self.get_mask_abs(seq_len, seq_len)?; - let mut cache_indices = Vec::with_capacity(self.batch_size()); - for (batch_i, &batch_mask) in batch_mask.iter().enumerate() { - if !batch_mask { - let indices = vec![self.indices[batch_i] as u32; seq_len]; - cache_indices.push(indices); - } else { - let mut indices = Vec::with_capacity(seq_len); - for _ in 0..seq_len { - let index = self.indices[batch_i]; - indices.push(index as u32); - self.indices[batch_i] += 1; - self.positions[batch_i] += 1; - if self.indices[batch_i] >= self.context { - self.indices[batch_i] = 0; - } - } - cache_indices.push(indices); - } - } - let indices = Tensor::new(cache_indices, self.device())?; - Ok(IndicesAndMask { indices, mask }) - } - - fn get_mask_abs(&self, size1: usize, size2: usize) -> Result { - let context = self.context; - let mask: Vec<_> = (0..size1) - .flat_map(|i| { - (0..size2).map(move |j| { - if size1 + j > size2 + i || size1 + j + context < size2 + i { - f32::NEG_INFINITY - } else { - 0.0 - } - }) - }) - .collect(); - Tensor::from_slice(&mask, (size1, size2), self.device()) - } -} - -#[derive(Debug, Clone)] -pub enum KvCache { - Rotating(RotatingKvCache), -} - -impl KvCache { - pub fn new(dim: usize, max_seq_len: usize) -> Self { - let cache = RotatingKvCache::new(dim, max_seq_len); - Self::Rotating(cache) - } - - pub fn current_seq_len(&self) -> usize { - match self { - KvCache::Rotating(cache) => cache.current_seq_len(), - } - } - - pub fn reset(&mut self) { - match self { - KvCache::Rotating(cache) => cache.reset(), - } - } - - pub fn append(&mut self, key: &Tensor, value: &Tensor) -> Result<(Tensor, Tensor)> { - match self { - KvCache::Rotating(cache) => cache.append(key, value), - } - } - - pub fn positions(&self, seq_len: usize) -> Vec { - match self { - KvCache::Rotating(cache) => cache.positions(seq_len), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use candle::IndexOp; - - #[test] - fn test_scattered_kv_cache() -> Result<()> { - let device = Device::Cpu; - let mut cache = ScatteredCacheBuilder::new(2, 5, DType::F32, &device)?; - let inf = f32::INFINITY; - - let iam = cache.indices_and_mask(1, &[true, false])?; - let mask = iam.mask.i((.., 0))?.to_vec3::()?; - assert_eq!(iam.indices.to_vec2::()?, [[0], [0]]); - assert_eq!(mask, [[[0.0, -inf, -inf, -inf, -inf]], [[0.0, 0.0, 0.0, 0.0, 0.0]]]); - - let iam = cache.indices_and_mask(1, &[true, false])?; - let mask = iam.mask.i((.., 0))?.to_vec3::()?; - assert_eq!(iam.indices.to_vec2::()?, [[1], [0]]); - assert_eq!(mask, [[[0.0, 0.0, -inf, -inf, -inf]], [[0.0, 0.0, 0.0, 0.0, 0.0]]]); - - let iam = cache.indices_and_mask(3, &[false, true])?; - let mask = iam.mask.i((.., 0))?.to_vec3::()?; - assert_eq!(iam.indices.to_vec2::()?, [[2, 2, 2], [0, 1, 2]]); - assert_eq!( - mask, - [ - [[0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]], - [ - [0.0, -inf, -inf, -inf, -inf], - [0.0, 0.0, -inf, -inf, -inf], - [0.0, 0.0, 0.0, -inf, -inf] - ] - ] - ); - - let iam = cache.indices_and_mask(3, &[true, true])?; - let mask = iam.mask.i((.., 0))?.to_vec3::()?; - assert_eq!(iam.indices.to_vec2::()?, [[2, 3, 4], [3, 4, 0]]); - assert_eq!( - mask, - [ - [ - [0.0, 0.0, 0.0, -inf, -inf], - [0.0, 0.0, 0.0, 0.0, -inf], - [0.0, 0.0, 0.0, 0.0, 0.0] - ], - [ - [-inf, 0.0, 0.0, 0.0, -inf], - [-inf, 0.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.0, 0.0, 0.0] - ] - ] - ); - - let iam = cache.indices_and_mask(1, &[true, false])?; - let mask = iam.mask.i((.., 0))?.to_vec3::()?; - assert_eq!(iam.indices.to_vec2::()?, [[0], [1]]); - assert_eq!(mask, [[[0.0, 0.0, 0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0, 0.0, 0.0]]]); - - let iam = cache.indices_and_mask(2, &[true, false])?; - let mask = iam.mask.i((.., 0))?.to_vec3::()?; - assert_eq!(iam.indices.to_vec2::()?, [[1, 2], [1, 1]]); - assert_eq!( - mask, - [ - [[0.0, 0.0, -inf, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]], - [[0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]] - ] - ); - - Ok(()) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -pub use candle; -pub use candle_nn; - -pub mod asr; -pub mod batched_transformer; -pub mod conditioner; -pub mod conv; -pub mod kv_cache; -pub mod lm; -pub mod lm_generate; -pub mod lm_generate_multistream; -pub mod mimi; -pub mod nn; -pub mod quantization; -pub mod seanet; -pub mod streaming; -pub mod transformer; -pub mod tts; -pub mod tts_streaming; -pub mod wav; - -#[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)] -pub enum NormType { - RmsNorm, - LayerNorm, -} - -pub use streaming::{StreamMask, StreamTensor, StreamingModule}; - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use candle::{IndexOp, Tensor}; -use candle_transformers::generation::LogitsProcessor; - -use crate::transformer::CaSrc; - -pub const UNGENERATED: u32 = u32::MAX; - -#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)] -pub struct Config { - pub generated_audio_codebooks: usize, - pub input_audio_codebooks: usize, - pub audio_vocab_size: usize, - pub acoustic_delay: usize, - pub text_pad_token: u32, - pub text_eop_token: u32, - pub text_start_token: u32, -} - -impl Config { - pub fn v0_1() -> Self { - Self { - generated_audio_codebooks: 8, - input_audio_codebooks: 8, - audio_vocab_size: 2049, - acoustic_delay: 2, - text_eop_token: 0, - text_pad_token: 3, - text_start_token: 32000, - } - } - - pub fn v0_1_two_ways() -> Self { - Self { - generated_audio_codebooks: 16, - input_audio_codebooks: 0, - audio_vocab_size: 2049, - acoustic_delay: 2, - text_eop_token: 0, - text_pad_token: 3, - text_start_token: 32000, - } - } - - pub fn v0_1_one_way() -> Self { - Self { - generated_audio_codebooks: 8, - input_audio_codebooks: 0, - audio_vocab_size: 2049, - acoustic_delay: 2, - text_eop_token: 0, - text_pad_token: 3, - text_start_token: 32000, - } - } - - pub fn audio_pad_token(&self) -> u32 { - self.audio_vocab_size as u32 - 1 - } - - pub fn total_audio_codebooks(&self) -> usize { - self.generated_audio_codebooks + self.input_audio_codebooks - } -} - -pub struct State { - model: crate::lm::LmModel, - audio_tokens: Vec>, - text_tokens: Vec, - audio_lp: LogitsProcessor, - text_lp: LogitsProcessor, - step_idx: usize, - pad_mult: Option, - // For repetition penalty, we provide the context len (in text tokens) and the penalty. - repetition_penalty: Option<(usize, f32)>, - forced_audio_tokens: crate::lm::ForcedAudioTokens, - user_rating: u32, - cfg_alpha: Option, - config: Config, -} - -impl State { - #[allow(clippy::too_many_arguments)] - pub fn new( - model: crate::lm::LmModel, - max_step_idx: usize, - audio_lp: LogitsProcessor, - text_lp: LogitsProcessor, - pad_mult: Option, - repetition_penalty: Option<(usize, f32)>, - cfg_alpha: Option, - config: Config, - ) -> Self { - let audio_tokens: Vec> = vec![ - vec![UNGENERATED; config.total_audio_codebooks()]; - max_step_idx + config.acoustic_delay - ]; - let text_tokens = vec![UNGENERATED; max_step_idx + config.acoustic_delay]; - let forced_audio_tokens = crate::lm::ForcedAudioTokens::new( - config.acoustic_delay, - config.audio_pad_token(), - &[8, 8], - ); - Self { - model, - audio_tokens, - text_tokens, - audio_lp, - text_lp, - step_idx: 0, - pad_mult, - repetition_penalty, - forced_audio_tokens, - user_rating: 0, // 0 indicates no ratings have been submitted from the front - cfg_alpha, - config, - } - } - - pub fn step_idx(&self) -> usize { - self.step_idx - } - - fn audio_pad_token(&self) -> u32 { - self.config.audio_pad_token() - } - - pub fn config(&self) -> &Config { - &self.config - } - - pub fn user_rating(&self) -> u32 { - self.user_rating - } - pub fn set_user_rating(&mut self, grade: u32) { - self.user_rating = grade - } - - fn apply_repetition_penalty(&self, logits: Tensor) -> candle::Result { - let logits = match self.repetition_penalty { - None => logits, - Some((_, 1.)) => logits, - Some((context_size, penalty)) => { - let device = logits.device(); - let mut logits = logits.to_dtype(candle::DType::F32)?.to_vec1::()?; - let mut already_seen = std::collections::HashSet::new(); - let mut non_pad_tokens = 0; - for &token_id in self.text_tokens(false).iter().rev() { - if token_id == self.config.text_pad_token - || token_id == self.config.text_eop_token - || token_id == self.config.text_start_token - { - continue; - } - // Look at the last [context_size] tokens at most, count all tokens there even - // if we already saw them. - if non_pad_tokens >= context_size { - break; - } - non_pad_tokens += 1; - - if already_seen.contains(&token_id) { - continue; - } - - already_seen.insert(token_id); - if let Some(logit) = logits.get_mut(token_id as usize) { - if *logit >= 0. { - *logit /= penalty - } else { - *logit *= penalty - } - } - } - let logits_len = logits.len(); - Tensor::from_vec(logits, logits_len, device)? - } - }; - Ok(logits) - } - - // The acoustic tokens are written with a delay, so this can create "gaps" of UNGENERATED - // tokens in the case where we call `step_audio_prompt` *after* `step`. - pub fn step_( - &mut self, - text_token: Option, - input_audio_tokens: &[u32], - force_text_token: Option, - ca_src: Option<&CaSrc>, - conditions: Option<&crate::conditioner::Condition>, - ) -> candle::Result { - let mut codes = Vec::with_capacity(self.config.total_audio_codebooks()); - let dev = self.model.device(); - for (c_idx, &t) in input_audio_tokens.iter().enumerate() { - self.audio_tokens[self.step_idx][c_idx + self.config.generated_audio_codebooks] = t - } - let batch_size = if self.cfg_alpha.is_some() { 2 } else { 1 }; - for codebook in 0..self.config.total_audio_codebooks() { - let t = if codebook == 0 || codebook == self.config.generated_audio_codebooks { - if self.step_idx == 0 { - self.audio_pad_token() - } else { - self.audio_tokens[self.step_idx - 1][codebook] - } - } else if self.step_idx <= self.config.acoustic_delay { - self.audio_pad_token() - } else { - self.audio_tokens[self.step_idx - self.config.acoustic_delay - 1][codebook] - }; - if t == UNGENERATED { - candle::bail!("internal error, ungenerated {} {codebook}", self.step_idx) - } - let t = Tensor::from_vec(vec![t; batch_size], (batch_size, 1), dev)?; - codes.push(Some(t)) - } - let text_token = match text_token { - Some(text_token) => { - Some(Tensor::from_vec(vec![text_token; batch_size], (batch_size, 1), dev)?) - } - None => None, - }; - let (text_logits, ys) = match ca_src.as_ref() { - None => { - let (logits, ys) = - self.model.forward_cond(text_token, codes, conditions, &().into())?; - let logits = match self.cfg_alpha { - None => logits.i((0, 0))?, - Some(a) => match logits.dim(0)? { - 2 => ((logits.i((0, 0))? * a)? - (logits.i((1, 0))? * (a - 1.))?)?, - b_size => candle::bail!("unexpected batch size {b_size}"), - }, - }; - (logits, ys) - } - Some(ca_src) => { - if self.cfg_alpha.is_some() { - candle::bail!("cfg is not supported with cross attention") - } - let (logits, ys) = - self.model.forward_ca(text_token, codes, ca_src, None, &().into())?; - (logits.i((0, 0))?, ys) - } - }; - let text_logits = self.apply_repetition_penalty(text_logits)?; - let text_token = match force_text_token { - Some(tt) => tt, - None => self.text_lp.sample_f(&text_logits, |prs| { - if let Some(pad_mult) = self.pad_mult.as_ref() { - prs[self.config.text_pad_token as usize] *= f32::exp(*pad_mult); - } - })?, - }; - self.text_tokens[self.step_idx] = text_token; - let last_audio_tokens = match self.cfg_alpha { - None => self.model.depformer_sample( - &ys, - Some(text_token), - self.forced_audio_tokens.forced_tokens(self.step_idx), - &mut self.audio_lp, - )?, - Some(cfg_alpha) => self.model.depformer_sample_cfg( - &ys, - cfg_alpha, - Some(text_token), - self.forced_audio_tokens.forced_tokens(self.step_idx), - &mut self.audio_lp, - )?, - }; - let audio_pad_token = self.audio_pad_token(); - for c_idx in 0..self.config.generated_audio_codebooks { - let delay = if c_idx == 0 || c_idx == self.config.generated_audio_codebooks { - 0 - } else { - self.config.acoustic_delay - }; - let pos = &mut self.audio_tokens[self.step_idx.saturating_sub(delay)][c_idx]; - // Overwrite existing positions even if there are non-UNGENERATED values. This - // actually happens for the first few slices because of the saturating_sub. - *pos = last_audio_tokens.as_ref().map_or(audio_pad_token, |l| l[c_idx]); - } - self.step_idx += 1; - if self.step_idx >= self.audio_tokens.len() { - candle::bail!("max step-idx reached") - } - Ok(text_token) - } - - pub fn step_without_ca_src( - &mut self, - text_token: u32, - input_audio_tokens: &[u32], - force_text_token: Option, - ) -> candle::Result { - self.step_(Some(text_token), input_audio_tokens, force_text_token, None, None) - } - - pub fn step( - &mut self, - text_token: u32, - input_audio_tokens: &[u32], - force_text_token: Option, - ca_src: Option<&CaSrc>, - ) -> candle::Result { - self.step_(Some(text_token), input_audio_tokens, force_text_token, ca_src, None) - } - - /// If include_all is set, all the time steps are returned. Otherwise only the timesteps that - /// have been generated are handled. - pub fn audio_tokens(&self, include_all: bool) -> &[Vec] { - if include_all { - &self.audio_tokens - } else { - let max_idx = usize::min(self.step_idx, self.audio_tokens.len()); - &self.audio_tokens[..max_idx] - } - } - - pub fn text_tokens(&self, include_all: bool) -> &[u32] { - if include_all { - &self.text_tokens - } else { - let max_idx = usize::min(self.step_idx, self.text_tokens.len()); - &self.text_tokens[..max_idx] - } - } - - pub fn last_audio_tokens(&self) -> Option> { - if self.step_idx <= self.config.acoustic_delay { - None - } else { - // step_idx is in advance by 1 + there is a 2 token delay on audio tokens. - let audio_tokens = &self.audio_tokens[self.step_idx - self.config.acoustic_delay - 1]; - if audio_tokens.iter().any(|v| *v as usize >= self.config.audio_vocab_size - 1) { - None - } else { - Some(audio_tokens.clone()) - } - } - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -// The state struct in this module handles generation for a LM model: -// - Apply the audio delays. -// - Allow for teacher forcing of the audio/text tokens. -// - Support "literal-zeros" tokens for both text and audio. -// - Make no assumptions on the number of streams. -// - TODO: Handle batch size > 1 -// - TODO: Support CFG. -// - TODO: Use CPU based tensors for storing the tokens? - -use candle::{IndexOp, Result, Tensor}; -use candle_transformers::generation::LogitsProcessor; - -#[derive(Copy, Clone, Debug, PartialEq, Eq)] -pub enum Token { - Set(u32), - Ungenerated, - LiteralZero, -} - -#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)] -pub struct Config { - pub audio_delays: Vec, - pub audio_vocab_size: usize, - pub text_pad_token: u32, - pub text_eop_token: u32, - pub text_start_token: u32, -} - -impl Config { - pub fn audio_pad_token(&self) -> u32 { - self.audio_vocab_size as u32 - 1 - } - - pub fn audio_codebooks(&self) -> usize { - self.audio_delays.len() - } - - pub fn max_audio_delay(&self) -> usize { - self.audio_delays.iter().max().cloned().unwrap_or(0) - } -} - -pub struct State { - model: crate::lm::LmModel, - audio_tokens: Vec>, - text_tokens: Vec, - audio_lp: LogitsProcessor, - text_lp: LogitsProcessor, - step_idx: usize, - config: Config, -} - -impl State { - pub fn new( - model: crate::lm::LmModel, - max_step_idx: usize, - audio_lp: LogitsProcessor, - text_lp: LogitsProcessor, - config: Config, - ) -> Self { - // TODO(laurent): handle a batch dimension. - let total_len = max_step_idx + config.max_audio_delay(); - let audio_tokens = vec![vec![Token::Ungenerated; config.audio_codebooks()]; total_len]; - let text_tokens = vec![Token::Ungenerated; total_len]; - Self { model, audio_tokens, text_tokens, audio_lp, text_lp, step_idx: 0, config } - } - - pub fn step_idx(&self) -> usize { - self.step_idx - } - - pub fn audio_pad_token(&self) -> u32 { - self.config.audio_pad_token() - } - - pub fn config(&self) -> &Config { - &self.config - } - - pub fn set_audio_tokens(&mut self, audio_tokens: &[Option]) -> Result<()> { - for (s, at) in self.audio_tokens[self.step_idx].iter_mut().zip(audio_tokens.iter()) { - if let Some(at) = at { - *s = *at - } - } - Ok(()) - } - - pub fn step(&mut self, conditions: Option<&crate::conditioner::Condition>) -> Result<()> { - let dev = self.model.device(); - - let mut forced_audio_tokens = Vec::with_capacity(self.config.audio_codebooks()); - for (codebook, &delay) in self.config.audio_delays.iter().enumerate() { - let forced_token = if self.step_idx < delay { - Some(self.audio_pad_token()) - } else { - match self.audio_tokens[self.step_idx - delay][codebook] { - Token::Ungenerated | Token::LiteralZero => None, - Token::Set(v) => Some(v), - } - }; - forced_audio_tokens.push(forced_token); - } - - let mut codes = Vec::with_capacity(self.config.audio_codebooks()); - for (codebook, &delay) in self.config.audio_delays.iter().enumerate() { - let t = if self.step_idx <= delay { - Some(self.audio_pad_token()) - } else { - match self.audio_tokens[self.step_idx - delay - 1][codebook] { - Token::LiteralZero => None, - Token::Set(v) => Some(v), - Token::Ungenerated => { - candle::bail!("internal error, ungenerated {} {codebook}", self.step_idx) - } - } - }; - let t = match t { - None => None, - Some(t) => Some(Tensor::from_vec(vec![t; 1], (1, 1), dev)?), - }; - codes.push(t) - } - let text_token = if self.step_idx == 0 { - Some(self.config.text_start_token) - } else { - match self.text_tokens[self.step_idx - 1] { - Token::LiteralZero => None, - Token::Set(t) => Some(t), - Token::Ungenerated => { - candle::bail!("internal error, ungenerated {} text", self.step_idx) - } - } - }; - let text_token = match text_token { - None => None, - Some(t) => Some(Tensor::from_vec(vec![t; 1], (1, 1), dev)?), - }; - let (text_logits, ys) = - self.model.forward_cond(text_token, codes, conditions, &().into())?; - let text_token = match self.text_tokens[self.step_idx] { - Token::Ungenerated => { - let t = self.text_lp.sample(&text_logits.i((0, 0))?)?; - self.text_tokens[self.step_idx] = Token::Set(t); - Some(t) - } - Token::Set(t) => Some(t), - Token::LiteralZero => None, - }; - let audio_tokens = self.model.depformer_sample( - &ys, - text_token, - &forced_audio_tokens, - &mut self.audio_lp, - )?; - if let Some(audio_tokens) = audio_tokens { - for (codebook, audio_token) in audio_tokens.into_iter().enumerate() { - let delay = self.config.audio_delays[codebook]; - if self.step_idx < delay { - continue; - } - let pos = &mut self.audio_tokens[self.step_idx - delay][codebook]; - if *pos == Token::Ungenerated { - *pos = Token::Set(audio_token) - } - } - } - self.step_idx += 1; - if self.step_idx >= self.audio_tokens.len() { - candle::bail!("max step-idx reached") - } - Ok(()) - } - - pub fn last_text_token(&self) -> Result> { - if self.step_idx == 0 { - Ok(None) - } else { - match self.text_tokens[self.step_idx - 1] { - Token::Set(t) => Ok(Some(t)), - Token::LiteralZero => Ok(None), - Token::Ungenerated => { - candle::bail!("internal error, ungenerated step {}, text", self.step_idx) - } - } - } - } - - pub fn last_audio_tokens(&self) -> Result>> { - let max_audio_delay = self.config.max_audio_delay(); - if self.step_idx <= max_audio_delay { - Ok(None) - } else { - let mut audio_tokens = vec![]; - for (cb, audio_token) in - self.audio_tokens[self.step_idx - max_audio_delay - 1].iter().enumerate() - { - match audio_token { - Token::LiteralZero => return Ok(None), - Token::Set(s) => audio_tokens.push(*s), - Token::Ungenerated => { - candle::bail!("internal error, ungenerated step {}, cb {cb}", self.step_idx) - } - } - } - Ok(Some(audio_tokens)) - } - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::nn::{linear, MaybeQuantizedEmbedding, MaybeQuantizedLinear, MaybeQuantizedVarBuilder}; -use crate::{ - batched_transformer, - transformer::{self, CaSrc}, - NormType, StreamMask, -}; -use candle::{DType, Device, IndexOp, Module, Result, Tensor}; - -thread_local! { - pub static VERBOSE: bool = { - match std::env::var("MIMI_VERBOSE") { - Ok(s) => { - !s.is_empty() && s != "0" - }, - Err(_) => false, - } - } -} -#[derive(Debug, Clone, serde::Deserialize)] -pub struct DepFormerConfig { - pub transformer: transformer::Config, - pub num_slices: usize, - pub low_rank_embeddings: Option, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct ExtraHeadsConfig { - pub num_heads: usize, - pub dim: usize, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct Config { - pub transformer: transformer::Config, - pub depformer: Option, - pub text_in_vocab_size: usize, - pub text_out_vocab_size: usize, - pub audio_vocab_size: usize, - pub audio_codebooks: usize, - pub conditioners: Option, - pub extra_heads: Option, -} - -impl Config { - fn depformer_cfg(num_slices: usize) -> DepFormerConfig { - let depformer_cfg = transformer::Config { - d_model: 1024, - num_heads: 16, - num_layers: 6, - dim_feedforward: 1024 * 4, // dim * hidden_scale - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: None, - context: num_slices, - max_period: 10000, - use_conv_block: false, - use_conv_bias: true, - cross_attention: None, - gating: Some(candle_nn::Activation::Silu), - norm: NormType::RmsNorm, - positional_embedding: transformer::PositionalEmbedding::None, - conv_layout: false, - conv_kernel_size: 3, - kv_repeat: 1, - max_seq_len: 4096, - shared_cross_attn: false, - }; - DepFormerConfig { num_slices, transformer: depformer_cfg, low_rank_embeddings: None } - } - - // /lustre/scwpod02/client/kyutai/alex/mimi_exp/xps/af78657c/outputs/hyperparams.json - // Update 2024-03-19: Sin embeddings -> None, RmsNorm fix, scale factor 4.125 - // Update 2024-05-02: split text_vocab_size into text_in_vocab_size and text_out_vocab_size. - // embeddings. - pub fn v0_1() -> Self { - let lm_cfg = transformer::Config { - d_model: 4096, - num_heads: 32, - num_layers: 32, - dim_feedforward: 4096 * 4, // dim * hidden_scale - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: None, - context: 3000, - max_period: 10000, - use_conv_block: false, - use_conv_bias: true, - cross_attention: None, - gating: Some(candle_nn::Activation::Silu), - norm: NormType::RmsNorm, - positional_embedding: transformer::PositionalEmbedding::Rope, - conv_layout: false, - conv_kernel_size: 3, - kv_repeat: 1, - max_seq_len: 4096, - shared_cross_attn: false, - }; - Self { - transformer: lm_cfg, - depformer: Some(Self::depformer_cfg(8)), - audio_vocab_size: 2049, - text_in_vocab_size: 32001, - text_out_vocab_size: 32000, - audio_codebooks: 8, - conditioners: Default::default(), - extra_heads: None, - } - } - - pub fn v0_1_vision() -> Self { - let lm_cfg = transformer::Config { - d_model: 4096, - num_heads: 32, - num_layers: 32, - dim_feedforward: 4096 * 4, // dim * hidden_scale - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: None, - context: 3000, - max_period: 10000, - use_conv_block: false, - use_conv_bias: true, - cross_attention: Some(( - transformer::CrossAttentionGating::ConditionalGatedSigmoid, - NormType::RmsNorm, - None, - )), - gating: Some(candle_nn::Activation::Silu), - norm: NormType::RmsNorm, - positional_embedding: transformer::PositionalEmbedding::Rope, - conv_layout: false, - conv_kernel_size: 3, - kv_repeat: 1, - max_seq_len: 4096, - shared_cross_attn: true, - }; - Self { - transformer: lm_cfg, - depformer: Some(Self::depformer_cfg(8)), - audio_vocab_size: 2049, - text_in_vocab_size: 32001, - text_out_vocab_size: 32000, - audio_codebooks: 8, - conditioners: Default::default(), - extra_heads: None, - } - } - - pub fn v0_1_vision_streaming(num_slices: usize) -> Self { - let mut s = Self::v0_1_vision(); - s.audio_codebooks = 16; - if let Some(depformer) = s.depformer.as_mut() { - depformer.num_slices = num_slices; - depformer.transformer.context = num_slices; - } - s - } - - pub fn v0_1_streaming(num_slices: usize) -> Self { - let mut s = Self::v0_1(); - s.audio_codebooks = 16; - if let Some(depformer) = s.depformer.as_mut() { - depformer.num_slices = num_slices; - depformer.transformer.context = num_slices; - } - s - } - - pub fn v0_1_asr() -> Self { - let mut s = Self::v0_1(); - s.audio_codebooks = 8; - if let Some(depformer) = s.depformer.as_mut() { - depformer.num_slices = 0; - depformer.transformer.context = 0; - } - s - } - - // /lustre/scwpod02/client/kyutai/neilz/mimi_exp/xps/6bbe4692/outputs/hyperparams.json - pub fn tts_v0_1() -> Self { - let lm_cfg = transformer::Config { - d_model: 2048, - num_heads: 32, - num_layers: 48, - dim_feedforward: 4096 * 2, // dim * hidden_scale - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: None, - context: 4096, - max_period: 10000, - use_conv_block: false, - use_conv_bias: true, - cross_attention: Some(( - transformer::CrossAttentionGating::Normal, - NormType::LayerNorm, - None, - )), - gating: None, - norm: NormType::LayerNorm, - positional_embedding: transformer::PositionalEmbedding::Rope, - conv_layout: false, - conv_kernel_size: 3, - kv_repeat: 1, - max_seq_len: 4096, - shared_cross_attn: false, - }; - Self { - transformer: lm_cfg, - depformer: Some(Self::depformer_cfg(16)), - audio_vocab_size: 2050, - text_in_vocab_size: 32001, - text_out_vocab_size: 32001, - audio_codebooks: 16, - conditioners: Default::default(), - extra_heads: None, - } - } - - // /lustre/scwpod02/client/kyutai-interns/tomlab/mimi_exp/xps/c879d080/.hydra/config.yaml - // /lustre/scwpod02/client/kyutai-interns/tomlab/mimi_exp/xps/41e5e07d/.hydra/config.yaml - pub fn s2s_v0_1() -> Self { - let lm_cfg = transformer::Config { - d_model: 2048, - num_heads: 16, - num_layers: 16, - dim_feedforward: 4096 * 2, // dim * hidden_scale - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: None, - context: 3000, - max_period: 10000, - use_conv_block: false, - use_conv_bias: true, - cross_attention: None, - gating: Some(candle_nn::Activation::Silu), - norm: NormType::RmsNorm, - positional_embedding: transformer::PositionalEmbedding::Rope, - conv_layout: false, - conv_kernel_size: 3, - kv_repeat: 1, - max_seq_len: 4096, - shared_cross_attn: false, - }; - Self { - transformer: lm_cfg, - depformer: Some(Self::depformer_cfg(16)), - audio_vocab_size: 2049, - text_in_vocab_size: 48001, - text_out_vocab_size: 48000, - audio_codebooks: 16, - conditioners: Default::default(), - extra_heads: None, - } - } - - pub fn s2s_v0_1_streaming(num_slices: usize) -> Self { - let mut s = Self::s2s_v0_1(); - s.audio_codebooks = 16; - if let Some(depformer) = s.depformer.as_mut() { - depformer.num_slices = num_slices; - depformer.transformer.context = num_slices; - } - s - } - - // /lustre/scwpod02/client/kyutai/neilz/mimi_exp/xps/33e476c7/.hydra/config.yaml - pub fn asr_v0_1_1b() -> Self { - let lm_cfg = transformer::Config { - d_model: 2048, - num_heads: 16, - num_layers: 16, - dim_feedforward: 2048 * 4, - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: None, - context: 750, - max_period: 100_000, - use_conv_block: false, - use_conv_bias: true, - cross_attention: None, - gating: Some(candle_nn::Activation::Silu), - norm: NormType::RmsNorm, - positional_embedding: transformer::PositionalEmbedding::Rope, - conv_layout: false, - conv_kernel_size: 3, - kv_repeat: 1, - max_seq_len: 4096, - shared_cross_attn: false, - }; - Self { - transformer: lm_cfg, - depformer: None, - audio_vocab_size: 2049, - text_in_vocab_size: 48001, - text_out_vocab_size: 48000, - audio_codebooks: 8, - conditioners: Default::default(), - extra_heads: None, - } - } - - pub fn asr_300m_202501() -> Self { - let lm_cfg = transformer::Config { - d_model: 1024, - num_heads: 8, - num_layers: 16, - dim_feedforward: 1024 * 4, - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: None, - context: 750, - max_period: 100_000, - use_conv_block: false, - use_conv_bias: true, - cross_attention: None, - gating: Some(candle_nn::Activation::Silu), - norm: NormType::RmsNorm, - positional_embedding: transformer::PositionalEmbedding::Rope, - conv_layout: false, - conv_kernel_size: 3, - kv_repeat: 1, - max_seq_len: 4096, - shared_cross_attn: false, - }; - Self { - transformer: lm_cfg, - depformer: None, - audio_vocab_size: 2049, - text_in_vocab_size: 48001, - text_out_vocab_size: 48000, - audio_codebooks: 32, - conditioners: Default::default(), - extra_heads: None, - } - } - - // /lustre/scwpod02/client/kyutai/alex/mimi_exp/xps/d50593ae/.hydra/config.yaml - pub fn tts_202501() -> Self { - let lm_cfg = transformer::Config { - d_model: 2048, - num_heads: 32, - num_layers: 48, - dim_feedforward: 2048 * 4, // dim * hidden_scale - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: None, - context: 500, - max_period: 10000, - use_conv_block: false, - use_conv_bias: true, - cross_attention: Some(( - transformer::CrossAttentionGating::Normal, - NormType::LayerNorm, - None, - )), - gating: Some(candle_nn::Activation::Silu), - norm: NormType::RmsNorm, - positional_embedding: transformer::PositionalEmbedding::Rope, - conv_layout: false, - conv_kernel_size: 3, - kv_repeat: 1, - max_seq_len: 4096, - shared_cross_attn: false, - }; - Self { - transformer: lm_cfg, - depformer: Some(Self::depformer_cfg(32)), - audio_vocab_size: 2049, - text_in_vocab_size: 8001, - text_out_vocab_size: 8000, - audio_codebooks: 32, - conditioners: Default::default(), - extra_heads: None, - } - } - - // /lustre/scwpod02/client/kyutai-interns/tomlab/mimi_exp/xps/1d426dfd/.hydra/config.yaml - pub fn s2s_2b_16rvq_202501() -> Self { - let lm_cfg = transformer::Config { - d_model: 2560, - num_heads: 20, - num_layers: 24, - dim_feedforward: 2560 * 4, // dim * hidden_scale - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: None, - context: 3000, - max_period: 100000, - use_conv_block: false, - use_conv_bias: true, - cross_attention: None, - gating: Some(candle_nn::Activation::Silu), - norm: NormType::RmsNorm, - positional_embedding: transformer::PositionalEmbedding::Rope, - conv_layout: false, - conv_kernel_size: 3, - kv_repeat: 1, - max_seq_len: 4096, - shared_cross_attn: false, - }; - Self { - transformer: lm_cfg, - depformer: Some(Self::depformer_cfg(16)), - audio_vocab_size: 2049, - text_in_vocab_size: 48001, - text_out_vocab_size: 48000, - audio_codebooks: 32, - conditioners: Default::default(), - extra_heads: None, - } - } -} - -#[derive(Debug, Clone)] -struct LowRankEmbeddings { - embeddings: MaybeQuantizedEmbedding, - low_rank: Option, -} - -impl LowRankEmbeddings { - fn new( - in_vocab_size: usize, - dim: usize, - low_rank_dim: Option, - vb: MaybeQuantizedVarBuilder, - ) -> Result { - let (low_rank, embeddings) = match low_rank_dim { - None => { - let embeddings = MaybeQuantizedEmbedding::new(in_vocab_size, dim, vb)?; - (None, embeddings) - } - Some(low_rank_dim) => { - let low_rank = linear(low_rank_dim, dim, false, vb.pp("low_rank"))?; - let embeddings = MaybeQuantizedEmbedding::new(in_vocab_size, low_rank_dim, vb)?; - (Some(low_rank), embeddings) - } - }; - Ok(Self { embeddings, low_rank }) - } -} - -impl Module for LowRankEmbeddings { - fn forward(&self, xs: &Tensor) -> Result { - let embs = xs.apply(&self.embeddings)?; - match self.low_rank.as_ref() { - None => Ok(embs), - Some(lr) => embs.apply(lr), - } - } -} - -#[derive(Debug, Clone)] -struct DepFormerSlice { - // There is no need for a streaming+batching mode here as the depformer does not have - // "persistent" caches. - transformer: transformer::StreamingTransformer, - // Note that the embedding for the first slice does not have the same dimension as the - // embedding for the other slices as it takes a text token as input rather than an audio token. - emb: LowRankEmbeddings, - linear_in: MaybeQuantizedLinear, // depformer_in.{idx} - linear_out: MaybeQuantizedLinear, // linears.{idx} -} - -impl DepFormerSlice { - fn new( - in_vocab_size: usize, - out_vocab_size: usize, - main_transformer_dim: usize, - cfg: &DepFormerConfig, - vb: MaybeQuantizedVarBuilder, - ) -> Result { - let dim = cfg.transformer.d_model; - let transformer = - transformer::StreamingTransformer::new(&cfg.transformer, vb.pp("transformer"))?; - let emb = - LowRankEmbeddings::new(in_vocab_size, dim, cfg.low_rank_embeddings, vb.pp("emb"))?; - let linear_in = linear(main_transformer_dim, dim, false, vb.pp("linear_in"))?; - let linear_out = linear(dim, out_vocab_size, false, vb.pp("linear_out"))?; - Ok(Self { transformer, emb, linear_in, linear_out }) - } -} - -#[derive(Debug, Clone)] -pub struct DepFormer { - slices: Vec, -} - -impl DepFormer { - pub fn new( - text_vocab_size: usize, - audio_vocab_size: usize, - main_transformer_dim: usize, - cfg: &DepFormerConfig, - vb: MaybeQuantizedVarBuilder, - ) -> Result { - let mut slices = Vec::with_capacity(cfg.num_slices); - for slice_idx in 0..cfg.num_slices { - let in_vs = if slice_idx == 0 { text_vocab_size } else { audio_vocab_size }; - // The depformer cannot predict the audio padding token. - let slice = DepFormerSlice::new( - in_vs, - audio_vocab_size - 1, // The depformer cannot emit an audio padding token. - main_transformer_dim, - cfg, - vb.pp(slice_idx), - )?; - slices.push(slice) - } - Ok(Self { slices }) - } - - /// Run a transformer sampling step, getting a token id per codebook. - /// - `xs` is the previous layer hidden state. - pub fn sample( - &mut self, - xs: &Tensor, - text_token: Option, - forced_audio_tokens: &[Option], - lp: &mut candle_transformers::generation::LogitsProcessor, - ) -> Result> { - use crate::streaming::StreamingModule; - let dev = xs.device(); - let mut tokens = Vec::with_capacity(self.slices.len()); - let mut last_token = text_token; - for slice_idx in 0..self.slices.len() { - if slice_idx == 0 { - self.slices[slice_idx].transformer.reset_state(); - } else { - let (lhs, rhs) = self.slices.split_at_mut(slice_idx); - rhs[0].transformer.copy_state(&lhs[slice_idx - 1].transformer)? - } - let slice = &mut self.slices[slice_idx]; - let xs = slice.linear_in.forward(xs)?; - let xs = match last_token { - Some(last_token) => { - let token_id = Tensor::from_vec(vec![last_token], (1, 1), dev)?; - let token_emb = slice.emb.forward(&token_id)?; - xs.broadcast_add(&token_emb)? - } - None => xs, - }; - let xs = slice.transformer.forward(&xs)?; - let logits = xs.apply(&slice.linear_out)?; - let logits = match logits.dim(0)? { - 1 => logits.i((0, 0))?, - b_size => candle::bail!("unexpected batch size {b_size}"), - }; - let token = lp.sample(&logits)?; - if VERBOSE.with(|v| *v) { - println!("sampled {token} logits {slice_idx}:\n{logits}"); - } - tokens.push(token); - let token_for_next_layer = - forced_audio_tokens.get(slice_idx).copied().flatten().unwrap_or(token); - last_token = Some(token_for_next_layer); - } - Ok(tokens) - } - - // Sampling with classifier free guidance. - pub fn sample_cfg( - &mut self, - xs: &Tensor, - cfg_alpha: f64, - text_token: Option, - forced_audio_tokens: &[Option], - lp: &mut candle_transformers::generation::LogitsProcessor, - ) -> Result> { - use crate::streaming::StreamingModule; - let dev = xs.device(); - let mut tokens = Vec::with_capacity(self.slices.len()); - let mut last_token = text_token; - for slice_idx in 0..self.slices.len() { - if slice_idx == 0 { - self.slices[slice_idx].transformer.reset_state(); - } else { - let (lhs, rhs) = self.slices.split_at_mut(slice_idx); - rhs[0].transformer.copy_state(&lhs[slice_idx - 1].transformer)? - } - let slice = &mut self.slices[slice_idx]; - let xs = slice.linear_in.forward(xs)?; - let xs = match last_token { - Some(last_token) => { - let token_id = Tensor::from_vec(vec![last_token], (1, 1), dev)?; - let token_emb = slice.emb.forward(&token_id)?; - xs.broadcast_add(&token_emb)? - } - None => xs, - }; - let xs = slice.transformer.forward(&xs)?; - let logits = xs.apply(&slice.linear_out)?; - let logits = match logits.dim(0)? { - 2 => ((logits.i((0, 0))? * cfg_alpha)? - (logits.i((1, 0))? * (cfg_alpha - 1.))?)?, - b_size => candle::bail!("unexpected batch size {b_size}"), - }; - let token = lp.sample(&logits)?; - if VERBOSE.with(|v| *v) { - println!("sampled {token} logits {slice_idx}:\n{logits}"); - } - tokens.push(token); - let token_for_next_layer = - forced_audio_tokens.get(slice_idx).copied().flatten().unwrap_or(token); - last_token = Some(token_for_next_layer); - } - Ok(tokens) - } -} - -#[derive(Debug, Clone)] -enum StreamingTransformer { - Normal(transformer::StreamingTransformer), - Batched(batched_transformer::StreamingTransformer), -} - -impl crate::StreamingModule for StreamingTransformer { - fn reset_state(&mut self) { - match self { - StreamingTransformer::Normal(t) => t.reset_state(), - StreamingTransformer::Batched(t) => t.reset_state(), - } - } - - fn step( - &mut self, - xs: &crate::StreamTensor, - mask: &crate::StreamMask, - ) -> Result { - match self { - StreamingTransformer::Normal(t) => t.step(xs, mask), - StreamingTransformer::Batched(t) => t.step(xs, mask), - } - } -} - -impl StreamingTransformer { - fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - match self { - StreamingTransformer::Normal(t) => t.reset_batch_idx(batch_idx, batch_size), - StreamingTransformer::Batched(t) => t.reset_batch_idx(batch_idx), - } - } - - fn maybe_precompute_ca_kv(&self, ca_src: Option) -> Result> { - match self { - StreamingTransformer::Normal(t) => t.maybe_precompute_ca_kv(ca_src), - StreamingTransformer::Batched(t) => t.maybe_precompute_ca_kv(ca_src), - } - } - - fn forward(&mut self, xs: &Tensor, m: &StreamMask) -> Result { - match self { - StreamingTransformer::Normal(t) => t.forward(xs), - StreamingTransformer::Batched(t) => t.forward(xs, m), - } - } - - fn forward_ca( - &mut self, - xs: &Tensor, - ca_src: Option<&CaSrc>, - m: &StreamMask, - ) -> Result { - match self { - StreamingTransformer::Normal(t) => t.forward_ca(xs, ca_src), - StreamingTransformer::Batched(t) => t.forward_ca(xs, ca_src, m), - } - } -} - -#[derive(Debug, Clone)] -pub struct LmModel { - transformer: StreamingTransformer, - text_emb: MaybeQuantizedEmbedding, - audio_embs: Vec, - text_linear: MaybeQuantizedLinear, - out_norm: transformer::Norm, - depformer: Option, - audio_vocab_size: usize, - text_in_vocab_size: usize, - condition_provider: Option, - extra_heads: Vec, - dtype: DType, -} - -impl LmModel { - pub fn new(cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - Self::new_(None, cfg, vb) - } - - pub fn batched(batch_size: usize, cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - Self::new_(Some(batch_size), cfg, vb) - } - - pub fn new_( - batch_size: Option, - cfg: &Config, - vb: MaybeQuantizedVarBuilder, - ) -> Result { - let d_model = cfg.transformer.d_model; - let depformer = match &cfg.depformer { - None => None, - Some(depformer_cfg) => { - let depformer = DepFormer::new( - cfg.text_in_vocab_size, - cfg.audio_vocab_size, - d_model, - depformer_cfg, - vb.pp("depformer"), - )?; - Some(depformer) - } - }; - let text_emb = - MaybeQuantizedEmbedding::new(cfg.text_in_vocab_size, d_model, vb.pp("text_emb"))?; - let out_norm = transformer::Norm::new(d_model, &cfg.transformer, vb.pp("out_norm"))?; - let text_linear = linear(d_model, cfg.text_out_vocab_size, false, vb.pp("text_linear"))?; - let transformer = match batch_size { - None => { - let transformer = - transformer::StreamingTransformer::new(&cfg.transformer, vb.pp("transformer"))?; - StreamingTransformer::Normal(transformer) - } - Some(batch_size) => { - let transformer = batched_transformer::StreamingTransformer::new( - batch_size, - &cfg.transformer, - vb.pp("transformer"), - )?; - StreamingTransformer::Batched(transformer) - } - }; - let vb_e = vb.pp("emb"); - let mut audio_embs = Vec::with_capacity(cfg.audio_codebooks); - for i in 0..cfg.audio_codebooks { - let emb = MaybeQuantizedEmbedding::new(cfg.audio_vocab_size, d_model, vb_e.pp(i))?; - audio_embs.push(emb) - } - let dtype = vb.dtype(); - let condition_provider = match cfg.conditioners.as_ref() { - None => None, - Some(cfg) => { - let conditioners = crate::conditioner::ConditionProvider::new( - d_model, - cfg, - vb.pp("condition_provider"), - )?; - Some(conditioners) - } - }; - let mut extra_heads = vec![]; - if let Some(ExtraHeadsConfig { num_heads, dim }) = cfg.extra_heads { - for i in 0..num_heads { - let extra_head = linear(d_model, dim, false, vb.pp("extra_heads").pp(i))?; - extra_heads.push(extra_head) - } - } - Ok(Self { - transformer, - text_emb, - text_linear, - audio_embs, - out_norm, - depformer, - text_in_vocab_size: cfg.text_in_vocab_size, - audio_vocab_size: cfg.audio_vocab_size, - condition_provider, - extra_heads, - dtype, - }) - } - - pub fn condition_provider(&self) -> Option<&crate::conditioner::ConditionProvider> { - self.condition_provider.as_ref() - } - - pub fn reset_state(&mut self) { - use crate::streaming::StreamingModule; - self.transformer.reset_state() - } - - pub fn in_audio_codebooks(&self) -> usize { - self.audio_embs.len() - } - - pub fn audio_pad_token(&self) -> u32 { - self.audio_vocab_size as u32 - 1 - } - - pub fn text_start_token(&self) -> u32 { - self.text_in_vocab_size as u32 - 1 - } - - pub fn generated_audio_codebooks(&self) -> usize { - self.depformer.as_ref().map_or(0, |v| v.slices.len()) - } - - pub fn is_quantized(&self) -> bool { - match self.text_linear { - MaybeQuantizedLinear::Quantized(_) => true, - MaybeQuantizedLinear::Real(_) => false, - } - } - - pub fn device(&self) -> &Device { - self.text_emb.embeddings().device() - } - - pub fn dtype(&self) -> DType { - self.text_emb.embeddings().dtype() - } - - pub fn forward( - &mut self, - text_ids: Option, - audio_ids: Vec>, - mask: &StreamMask, - ) -> candle::Result<(Tensor, Tensor)> { - self.forward_cond(text_ids, audio_ids, None, mask) - } - - pub fn extra_heads(&self, vs: &Tensor) -> Result> { - let mut extra_heads = Vec::with_capacity(self.extra_heads.len()); - for extra_head in self.extra_heads.iter() { - let extra_head = vs.apply(extra_head)?; - extra_heads.push(extra_head) - } - Ok(extra_heads) - } - - pub fn forward_cond( - &mut self, - text_ids: Option, - audio_ids: Vec>, - conditions: Option<&crate::conditioner::Condition>, - mask: &StreamMask, - ) -> candle::Result<(Tensor, Tensor)> { - if VERBOSE.with(|v| *v) { - print!("text_ids "); - if let Some(text_ids) = text_ids.as_ref() { - let text_ids = text_ids.flatten_all()?.to_vec1::()?; - println!("{text_ids:?}"); - } else { - println!("none") - } - print!("audio_ids "); - for audio_id in audio_ids.iter() { - if let Some(audio_id) = audio_id { - let audio_id = audio_id.flatten_all()?.to_vec1::()?; - print!(" {audio_id:?}"); - } else { - print!(" none") - } - } - println!(); - } - let mut emb = match text_ids.as_ref() { - Some(text_ids) => text_ids.apply(&self.text_emb)?, - None => { - let device = self.text_emb.embeddings().device(); - Tensor::zeros((1, 1, self.text_emb.hidden_size()?), self.dtype, device)? - } - }; - - for (audio_emb, audio_ids) in self.audio_embs.iter().zip(audio_ids.iter()) { - if let Some(audio_ids) = audio_ids { - let e = audio_ids.apply(audio_emb)?; - emb = (emb + e)? - } - } - if let Some(conditions) = conditions { - match conditions { - crate::conditioner::Condition::AddToInput(v) => emb = emb.broadcast_add(v)?, - } - } - let ys = self.transformer.forward(&emb, mask)?; - let ys = ys.apply(&self.out_norm)?; - let logits = ys.apply(&self.text_linear)?; - if VERBOSE.with(|v| *v) { - println!("logits:\n{logits}"); - } - Ok((logits, ys)) - } - - pub fn maybe_precompute_ca_kv(&self, ca_src: Option) -> Result> { - let ca_src = match ca_src { - None => None, - z => self.transformer.maybe_precompute_ca_kv(z)?, - }; - Ok(ca_src) - } - - pub fn forward_ca( - &mut self, - text_ids: Option, - audio_ids: Vec>, - ca_src: &CaSrc, - conditions: Option<&crate::conditioner::Condition>, - mask: &StreamMask, - ) -> candle::Result<(Tensor, Tensor)> { - if VERBOSE.with(|v| *v) { - print!("text_ids "); - if let Some(text_ids) = text_ids.as_ref() { - let text_ids = text_ids.flatten_all()?.to_vec1::()?; - println!("{text_ids:?}"); - } else { - println!("none") - } - print!("audio_ids "); - for audio_id in audio_ids.iter() { - if let Some(audio_id) = audio_id { - let audio_id = audio_id.flatten_all()?.to_vec1::()?; - print!(" {audio_id:?}"); - } else { - print!(" none") - } - } - println!(); - } - let b_size = match ca_src { - CaSrc::KeysValues((cak, _)) => cak.dim(0)?, - CaSrc::Tokens(catoks) => catoks.dim(0)?, - }; - let mut emb = match text_ids { - Some(text_ids) => text_ids.apply(&self.text_emb)?, - None => { - let device = self.text_emb.embeddings().device(); - Tensor::zeros((b_size, 1, self.text_emb.hidden_size()?), self.dtype, device)? - } - }; - for (audio_emb, audio_ids) in self.audio_embs.iter().zip(audio_ids.iter()) { - if let Some(audio_ids) = audio_ids { - let e = audio_ids.apply(audio_emb)?; - emb = emb.broadcast_add(&e)? - } - } - if let Some(conditions) = conditions { - match conditions { - crate::conditioner::Condition::AddToInput(v) => emb = emb.broadcast_add(v)?, - } - } - let ys = self.transformer.forward_ca(&emb, Some(ca_src), mask)?; - let ys = ys.apply(&self.out_norm)?; - let logits = ys.apply(&self.text_linear)?; - Ok((logits, ys)) - } - - pub fn depformer_sample( - &mut self, - xs: &Tensor, - text_token: Option, - forced_audio_tokens: &[Option], - lp: &mut candle_transformers::generation::LogitsProcessor, - ) -> Result>> { - let sample = match self.depformer.as_mut() { - None => None, - Some(m) => { - let sample = m.sample(xs, text_token, forced_audio_tokens, lp)?; - Some(sample) - } - }; - Ok(sample) - } - - pub fn depformer_sample_cfg( - &mut self, - xs: &Tensor, - cfg_alpha: f64, - text_token: Option, - forced_audio_tokens: &[Option], - lp: &mut candle_transformers::generation::LogitsProcessor, - ) -> Result>> { - let sample = match self.depformer.as_mut() { - None => None, - Some(m) => { - let sample = m.sample_cfg(xs, cfg_alpha, text_token, forced_audio_tokens, lp)?; - Some(sample) - } - }; - Ok(sample) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - self.transformer.reset_batch_idx(batch_idx, batch_size) - } -} - -pub fn load_lm_model>( - cfg: Config, - model_file: P, - dtype: DType, - dev: &Device, -) -> Result { - let quantized = model_file.as_ref().extension().is_some_and(|v| v == "gguf"); - let vb = if quantized { - MaybeQuantizedVarBuilder::Quantized( - candle_transformers::quantized_var_builder::VarBuilder::from_gguf(model_file, dev)?, - ) - } else { - unsafe { - MaybeQuantizedVarBuilder::Real(candle_nn::VarBuilder::from_mmaped_safetensors( - &[model_file], - dtype, - dev, - )?) - } - }; - let model = LmModel::new(&cfg, vb)?; - Ok(model) -} - -pub fn load>( - model_file: P, - dtype: DType, - dev: &Device, -) -> Result { - let cfg = Config::v0_1(); - load_lm_model(cfg, model_file, dtype, dev) -} - -pub fn load_streaming>( - model_file: P, - dtype: DType, - dev: &Device, -) -> Result { - let cfg = Config::v0_1_streaming(8); - load_lm_model(cfg, model_file, dtype, dev) -} - -pub fn load_streaming_both_ways>( - model_file: P, - dtype: DType, - dev: &Device, -) -> Result { - let cfg = Config::v0_1_streaming(16); - load_lm_model(cfg, model_file, dtype, dev) -} - -pub fn load_vision>( - model_file: P, - override_cross_attention_gating: Option, - override_cross_attention_in_dim: Option, - dtype: DType, - dev: &Device, -) -> Result { - // load_vision allows for overriding some hyperparams of the lm from the main config file - let mut cfg = Config::v0_1_vision_streaming(8); - cfg.transformer.cross_attention = override_cross_attention_gating - .map(|v| (v, cfg.transformer.norm, override_cross_attention_in_dim)); - load_lm_model(cfg, model_file, dtype, dev) -} - -pub fn load_s2s>( - model_file: P, - dtype: DType, - dev: &Device, -) -> Result { - let cfg = Config::s2s_2b_16rvq_202501(); - load_lm_model(cfg, model_file, dtype, dev) -} - -pub fn load_asr>( - model_file: P, - dtype: DType, - dev: &Device, -) -> Result { - let cfg = Config::asr_v0_1_1b(); - load_lm_model(cfg, model_file, dtype, dev) -} - -pub struct ForcedAudioTokens { - acoustic_delay: usize, - // Tokens that are teacher forced before the acoustic delay. - pre_delay_tokens: Vec>, -} - -impl ForcedAudioTokens { - pub fn new(acoustic_delay: usize, audio_pad_token: u32, stream_codebooks: &[usize]) -> Self { - let mut pre_delay_tokens = vec![]; - for codebooks in stream_codebooks.iter() { - for c in 0..*codebooks { - let token = if c == 0 { None } else { Some(audio_pad_token) }; - pre_delay_tokens.push(token); - } - } - Self { acoustic_delay, pre_delay_tokens } - } - - pub fn forced_tokens(&self, step_idx: usize) -> &[Option] { - if step_idx < self.acoustic_delay { - &self.pre_delay_tokens - } else { - &[] - } - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::streaming::{StreamMask, StreamTensor, StreamingModule}; -use crate::{conv, quantization, seanet, transformer}; -use candle::{DType, Device, Module, Result, Tensor}; -use candle_nn::VarBuilder; - -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub enum ResampleMethod { - Conv, - Interpolate, -} - -#[derive(Debug, Clone)] -pub struct Config { - pub channels: usize, - pub sample_rate: f64, - pub frame_rate: f64, - pub renormalize: bool, - pub resample_method: ResampleMethod, - pub seanet: seanet::Config, - pub transformer: transformer::Config, - pub quantizer_n_q: usize, - pub quantizer_bins: usize, - pub quantizer_dim: usize, -} - -impl Config { - // /lustre/scwpod02/client/kyutai/alex/mimi_exp/xps/b7d2bd5a/.hydra/config.yaml - pub fn v0_1(num_codebooks: Option) -> Self { - let seanet_cfg = seanet::Config { - dimension: 512, - channels: 1, - causal: true, - n_filters: 64, - n_residual_layers: 1, - activation: candle_nn::Activation::Elu(1.), - compress: 2, - dilation_base: 2, - disable_norm_outer_blocks: 0, - final_activation: None, - kernel_size: 7, - residual_kernel_size: 3, - last_kernel_size: 3, - lstm: 0, - norm: conv::Norm::WeightNorm, - pad_mode: conv::PadMode::Constant, - ratios: vec![8, 6, 5, 4], - true_skip: true, - }; - let transformer_cfg = transformer::Config { - d_model: seanet_cfg.dimension, - num_heads: 8, - num_layers: 8, - causal: true, - norm_first: true, - bias_ff: false, - bias_attn: false, - layer_scale: Some(0.01), - context: 250, - conv_kernel_size: 5, - use_conv_bias: true, - use_conv_block: false, - cross_attention: None, - max_period: 10000, - gating: None, - norm: crate::NormType::LayerNorm, - positional_embedding: transformer::PositionalEmbedding::Rope, - - dim_feedforward: 2048, - kv_repeat: 1, - conv_layout: true, // see builders.py - max_seq_len: 8192, // the transformer works at 25hz so this is ~5 mins. - shared_cross_attn: false, - }; - Config { - channels: 1, - sample_rate: 24_000., - frame_rate: 12.5, - renormalize: true, - resample_method: ResampleMethod::Conv, - seanet: seanet_cfg, - transformer: transformer_cfg, - quantizer_n_q: num_codebooks.unwrap_or(16), - quantizer_bins: 2048, - quantizer_dim: 256, - } - } -} - -#[derive(Debug, Clone)] -pub struct Mimi { - encoder: seanet::SeaNetEncoder, - decoder: seanet::SeaNetDecoder, - encoder_transformer: transformer::Transformer, - decoder_transformer: transformer::Transformer, - downsample: conv::ConvDownsample1d, - upsample: conv::ConvTrUpsample1d, - quantizer: quantization::SplitResidualVectorQuantizer, - config: Config, -} - -impl Mimi { - pub fn new(cfg: Config, vb: VarBuilder) -> Result { - Self::new_(None, cfg, vb) - } - - pub fn batched(batch_size: usize, cfg: Config, vb: VarBuilder) -> Result { - Self::new_(Some(batch_size), cfg, vb) - } - - fn new_(batch_size: Option, cfg: Config, vb: VarBuilder) -> Result { - let dim = cfg.seanet.dimension; - let encoder = seanet::SeaNetEncoder::new(&cfg.seanet, vb.pp("encoder"))?; - let decoder = seanet::SeaNetDecoder::new(&cfg.seanet, vb.pp("decoder"))?; - let encoder_transformer = transformer::Transformer::new( - batch_size, - dim, - &cfg.transformer, - vb.pp("encoder_transformer"), - )?; - let decoder_transformer = transformer::Transformer::new( - batch_size, - dim, - &cfg.transformer, - vb.pp("decoder_transformer"), - )?; - let quantizer = quantization::SplitResidualVectorQuantizer::new( - /* dim */ cfg.quantizer_dim, - /* input_dim */ Some(dim), - /* output_dim */ Some(dim), - /* n_q */ cfg.quantizer_n_q, - /* bins */ cfg.quantizer_bins, - vb.pp("quantizer"), - )?; - let encoder_frame_rate = - cfg.sample_rate / cfg.seanet.ratios.iter().product::() as f64; - - let downsample_stride = (encoder_frame_rate / cfg.frame_rate) as usize; - // `upsample` and `downsample` only apply if frame_rate is different from encoder_frame_rate. - let downsample = conv::ConvDownsample1d::new( - /* stride */ downsample_stride, - /* dim */ dim, - /* causal */ true, - /* learnt */ true, - vb.pp("downsample"), - )?; - let upsample = conv::ConvTrUpsample1d::new( - /* stride */ downsample_stride, - /* dim */ dim, - /* causal */ true, - /* learnt */ true, - vb.pp("upsample"), - )?; - - Ok(Self { - encoder, - decoder, - encoder_transformer, - decoder_transformer, - quantizer, - downsample, - upsample, - config: cfg, - }) - } - - pub fn config(&self) -> &Config { - &self.config - } - - pub fn encode_pre_quantize(&mut self, xs: &Tensor) -> Result { - let xs = self.encoder.forward(xs)?; - self.encoder_transformer.reset_state(); - let xs = self.encoder_transformer.forward(&xs)?; - let xs = &xs[0]; - xs.apply(&self.downsample) - } - - pub fn encode(&mut self, xs: &Tensor) -> Result { - let xs = self.encoder.forward(xs)?; - self.encoder_transformer.reset_state(); - let xs = self.encoder_transformer.forward(&xs)?; - let xs = &xs[0]; - let xs = xs.apply(&self.downsample)?; - let codes = self.quantizer.encode(&xs)?; - Ok(codes) - } - - pub fn encode_step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - let xs = self.encoder.step(xs, m)?; - let xs = self.encoder_transformer.step(&xs, m)?; - let xs = self.downsample.step(&xs, m)?; - match xs.as_option() { - None => Ok(().into()), - Some(xs) => { - let codes = self.quantizer.encode(xs)?; - Ok(codes.into()) - } - } - } - - pub fn decode(&mut self, codes: &Tensor) -> Result { - let emb = self.quantizer.decode(codes)?; - let emb = emb.apply(&self.upsample)?; - self.decoder_transformer.reset_state(); - let outs = self.decoder_transformer.forward(&emb)?; - let out = &outs[0]; - self.decoder.forward(out) - } - - pub fn decode_step(&mut self, codes: &StreamTensor, m: &StreamMask) -> Result { - let emb = match codes.as_option() { - Some(codes) => StreamTensor::from_tensor(self.quantizer.decode(codes)?), - None => StreamTensor::empty(), - }; - let emb = self.upsample.step(&emb, m)?; - let out = self.decoder_transformer.step(&emb, m)?; - self.decoder.step(&out, m) - } - - pub fn reset_state(&mut self) { - self.encoder.reset_state(); - self.encoder_transformer.reset_state(); - self.decoder.reset_state(); - self.decoder_transformer.reset_state(); - self.upsample.reset_state(); - self.downsample.reset_state(); - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - self.encoder_transformer.reset_batch_idx(batch_idx, batch_size)?; - self.encoder_transformer.reset_batch_idx(batch_idx, batch_size)?; - self.encoder.reset_batch_idx(batch_idx, batch_size)?; - self.decoder.reset_batch_idx(batch_idx, batch_size)?; - self.upsample.reset_batch_idx(batch_idx, batch_size)?; - self.downsample.reset_batch_idx(batch_idx, batch_size)?; - Ok(()) - } -} - -pub fn load(model_file: &str, num_codebooks: Option, dev: &Device) -> Result { - let vb = - unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[model_file], DType::F32, dev)? }; - let cfg = Config::v0_1(num_codebooks); - let mimi = Mimi::new(cfg, vb)?; - Ok(mimi) -} - -pub fn load_b( - batch_size: Option, - model_file: &str, - num_codebooks: Option, - dev: &Device, -) -> Result { - let vb = - unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[model_file], DType::F32, dev)? }; - let cfg = Config::v0_1(num_codebooks); - let mimi = Mimi::new_(batch_size, cfg, vb)?; - Ok(mimi) -} - - - -use candle::quantized::QTensor; -use candle::{DType, Device, Module, Result, Shape, Tensor}; -use candle_transformers::quantized_nn as candle_qnn; -use candle_transformers::quantized_var_builder::VarBuilder as QuantizedVarBuilder; - -use std::sync::Arc; - -#[derive(Clone)] -pub enum MaybeQuantizedWeight { - // Enum types around real and quantized model weights - Real(Tensor), - Quantized(Arc), -} - -impl MaybeQuantizedWeight { - fn to_tensor(&self, dev: &Device) -> Result { - match self { - Self::Real(t) => Ok(t.clone()), - Self::Quantized(t) => t.dequantize(dev), - } - } -} - -pub fn matmul_dtype(device: &candle::Device) -> DType { - // Dtype used for intermediate matmul in attention during quantized execution - if device.is_cuda() { - DType::BF16 - } else { - DType::F32 - } -} - -#[derive(Clone)] -pub enum MaybeQuantizedVarBuilder<'a> { - // Enum types around real and quantized var builders - Real(candle_nn::VarBuilder<'a>), - Quantized(QuantizedVarBuilder), -} - -impl MaybeQuantizedVarBuilder<'_> { - pub fn pp(&self, s: S) -> Self { - match self { - Self::Real(weights) => MaybeQuantizedVarBuilder::Real(weights.pp(s)), - Self::Quantized(weights) => MaybeQuantizedVarBuilder::Quantized(weights.pp(s)), - } - } - - pub fn get>(&self, s: S, path: &str) -> Result { - let w = match self { - Self::Real(weights) => MaybeQuantizedWeight::Real(weights.get(s, path)?), - Self::Quantized(weights) => MaybeQuantizedWeight::Quantized(weights.get(s, path)?), - }; - Ok(w) - } - - pub fn get_as_tensor>(&self, s: S, path: &str) -> Result { - let w = match self { - Self::Real(weights) => MaybeQuantizedWeight::Real(weights.get(s, path)?), - Self::Quantized(weights) => MaybeQuantizedWeight::Quantized(weights.get(s, path)?), - }; - w.to_tensor(self.device()) - } - - pub fn get_unquantized>(&self, s: S, path: &str) -> Result { - match self { - Self::Real(weights) => weights.get(s, path), - Self::Quantized(weights) => weights.get(s, path)?.dequantize(weights.device()), - } - } - - pub fn contains_key(&self, name: &str) -> bool { - match self { - Self::Real(weights) => weights.contains_tensor(name), - Self::Quantized(weights) => weights.contains_key(name), - } - } - - pub fn device(&self) -> &Device { - match self { - Self::Real(weights) => weights.device(), - Self::Quantized(weights) => weights.device(), - } - } - - pub fn dtype(&self) -> DType { - match self { - Self::Real(weights) => weights.dtype(), - Self::Quantized(_) => DType::F32, - } - } -} - -#[derive(Debug, Clone)] -pub enum MaybeQuantizedLinear { - Real(candle_nn::Linear), - Quantized(candle_qnn::Linear), -} - -impl Module for MaybeQuantizedLinear { - fn forward(&self, xs: &Tensor) -> Result { - match self { - Self::Real(module) => module.forward(xs), - Self::Quantized(module) => module.forward(xs), - } - } -} - -impl MaybeQuantizedLinear { - pub fn dtype(&self) -> DType { - match self { - Self::Real(l) => l.weight().dtype(), - Self::Quantized(_) => DType::F32, - } - } -} - -#[derive(Debug, Clone)] -pub enum MaybeQuantizedEmbedding { - Real(candle_nn::Embedding), - Quantized(candle_qnn::Embedding), -} - -impl MaybeQuantizedEmbedding { - pub fn new(in_vocab_size: usize, dim: usize, vb: MaybeQuantizedVarBuilder) -> Result { - let emb = match vb { - MaybeQuantizedVarBuilder::Real(weights) => { - MaybeQuantizedEmbedding::Real(candle_nn::embedding(in_vocab_size, dim, weights)?) - } - MaybeQuantizedVarBuilder::Quantized(weights) => MaybeQuantizedEmbedding::Quantized( - candle_transformers::quantized_nn::Embedding::new(in_vocab_size, dim, weights)?, - ), - }; - Ok(emb) - } - - pub fn embeddings(&self) -> &Tensor { - match self { - MaybeQuantizedEmbedding::Real(weights) => weights.embeddings(), - MaybeQuantizedEmbedding::Quantized(weights) => weights.embeddings(), - } - } - - pub fn hidden_size(&self) -> Result { - let size = match self { - MaybeQuantizedEmbedding::Real(weights) => weights.hidden_size(), - MaybeQuantizedEmbedding::Quantized(weights) => weights.embeddings().dim(1)?, - }; - Ok(size) - } - - pub fn dtype(&self) -> DType { - match self { - Self::Real(l) => l.embeddings().dtype(), - Self::Quantized(_) => DType::F32, - } - } -} - -impl Module for MaybeQuantizedEmbedding { - fn forward(&self, xs: &Tensor) -> Result { - match self { - Self::Real(module) => module.forward(xs), - Self::Quantized(module) => module.forward(xs), - } - } -} - -pub fn linear( - in_d: usize, - out_d: usize, - bias: bool, - vb: MaybeQuantizedVarBuilder, -) -> Result { - let output_linear = match vb { - MaybeQuantizedVarBuilder::Real(weights) => { - if bias { - MaybeQuantizedLinear::Real(candle_nn::linear(in_d, out_d, weights)?) - } else { - MaybeQuantizedLinear::Real(candle_nn::linear_no_bias(in_d, out_d, weights)?) - } - } - MaybeQuantizedVarBuilder::Quantized(weights) => { - MaybeQuantizedLinear::Quantized(candle_qnn::linear_b(in_d, out_d, bias, weights)?) - } - }; - Ok(output_linear) -} - -pub fn linear_from( - weight: MaybeQuantizedWeight, - bias: Option, -) -> Result { - let layer = match weight { - MaybeQuantizedWeight::Real(w) => { - MaybeQuantizedLinear::Real(candle_nn::Linear::new(w, bias)) - } - MaybeQuantizedWeight::Quantized(w) => { - MaybeQuantizedLinear::Quantized(candle_qnn::Linear::from_arc(w, bias)?) - } - }; - Ok(layer) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use candle::{IndexOp, Layout, Result, Shape, Tensor, D}; -use candle_nn::{linear, Linear, VarBuilder}; - -struct CodebookEncode; - -impl candle::CustomOp2 for CodebookEncode { - fn name(&self) -> &'static str { - "cb" - } - - fn cpu_fwd( - &self, - lhs_storage: &candle::CpuStorage, - lhs_layout: &Layout, - rhs_storage: &candle::CpuStorage, - rhs_layout: &Layout, - ) -> Result<(candle::CpuStorage, Shape)> { - use rayon::prelude::*; - - let (lhs_dim1, lhs_dim2) = lhs_layout.shape().dims2()?; - let (rhs_dim1, rhs_dim2) = rhs_layout.shape().dims2()?; - if lhs_dim2 != rhs_dim2 { - candle::bail!("CodebookEncode, mismatch on last dim, {lhs_layout:?} {rhs_layout:?}"); - } - if lhs_dim2 == 0 { - candle::bail!("CodebookEncode, empty last dim {lhs_layout:?}") - } - let lhs = match lhs_layout.contiguous_offsets() { - None => candle::bail!("CodebookEncode, lhs has to be contiguous, got {lhs_layout:?}"), - Some((o1, o2)) => { - let slice = lhs_storage.as_slice::()?; - &slice[o1..o2] - } - }; - let rhs = match rhs_layout.contiguous_offsets() { - None => candle::bail!("CodebookEncode, rhs has to be contiguous, got {rhs_layout:?}"), - Some((o1, o2)) => { - let slice = rhs_storage.as_slice::()?; - &slice[o1..o2] - } - }; - let dst = (0..lhs_dim1) - .into_par_iter() - .map(|idx1| { - let mut where_min = 0; - let mut min_dist = f32::INFINITY; - let lhs = &lhs[idx1 * lhs_dim2..(idx1 + 1) * lhs_dim2]; - for idx2 in 0..rhs_dim1 { - let rhs = &rhs[idx2 * rhs_dim2..(idx2 + 1) * rhs_dim2]; - let mut dist = 0f32; - for (a, b) in lhs.iter().zip(rhs.iter()) { - dist += (a - b) * (a - b) - } - if dist < min_dist { - min_dist = dist; - where_min = idx2; - } - } - where_min as u32 - }) - .collect(); - let storage = candle::WithDType::to_cpu_storage_owned(dst); - Ok((storage, (lhs_dim1,).into())) - } -} - -#[allow(unused)] -#[derive(Debug, Clone)] -pub struct EuclideanCodebook { - initialized: Tensor, - cluster_usage: Tensor, - embedding_sum: Tensor, - embedding: Tensor, - c2: Tensor, - epsilon: f64, - dim: usize, - span_encode: tracing::Span, - span_decode: tracing::Span, -} - -impl EuclideanCodebook { - pub fn new(dim: usize, codebook_size: usize, vb: VarBuilder) -> Result { - let epsilon = 1e-5; - let initialized = vb.get(1, "_initialized")?; - let cluster_usage = vb.get(codebook_size, "cluster_usage")?; - let embedding_sum = vb.get((codebook_size, dim), "embedding_sum")?; - let embedding = { - let cluster_usage = cluster_usage.maximum(epsilon)?.unsqueeze(1)?; - embedding_sum.broadcast_div(&cluster_usage)? - }; - let c2 = ((&embedding * &embedding)?.sum(D::Minus1)? / 2.0)?; - Ok(Self { - initialized, - cluster_usage, - embedding_sum, - embedding, - c2, - epsilon, - dim, - span_encode: tracing::span!(tracing::Level::TRACE, "euclidean-encode"), - span_decode: tracing::span!(tracing::Level::TRACE, "euclidean-encode"), - }) - } - - pub fn encode_very_slow(&self, xs: &Tensor) -> Result { - let _enter = self.span_encode.enter(); - let mut target_shape = xs.dims().to_vec(); - target_shape.pop(); - let xs = xs.flatten_to(D::Minus2)?; - let _ = xs.dims2()?; - // TODO: avoid repeating this. - let cluster_usage = self.cluster_usage.maximum(self.epsilon)?.unsqueeze(1)?; - let embedding = self.embedding_sum.broadcast_div(&cluster_usage)?; - // Manual cdist implementation. - let diff = xs.unsqueeze(1)?.broadcast_sub(&embedding.unsqueeze(0)?)?; - let dists = diff.sqr()?.sum(D::Minus1)?; - let codes = dists.argmin(D::Minus1)?; - codes.reshape(target_shape) - } - - pub fn encode_slow(&self, xs: &Tensor) -> Result { - let _enter = self.span_encode.enter(); - let mut target_shape = xs.dims().to_vec(); - target_shape.pop(); - let xs = xs.flatten_to(D::Minus2)?; - let _ = xs.dims2()?; - let dot_prod = xs.matmul(&self.embedding.t()?)?; - let codes = self.c2.broadcast_sub(&dot_prod)?.argmin(D::Minus1)?; - codes.reshape(target_shape) - } - - pub fn encode(&self, xs: &Tensor) -> Result { - let _enter = self.span_encode.enter(); - let mut target_shape = xs.dims().to_vec(); - target_shape.pop(); - let xs = xs.flatten_to(D::Minus2)?; - let _ = xs.dims2()?; - let codes = Tensor::apply_op2(&xs, &self.embedding, CodebookEncode)?; - codes.reshape(target_shape) - } - - pub fn decode(&self, indexes: &Tensor) -> Result { - let _enter = self.span_decode.enter(); - // let ys = candle_nn::Embedding::new(self.embedding.clone(), self.dim).forward(xs)?; - let mut final_dims = indexes.dims().to_vec(); - final_dims.push(self.dim); - let indexes = indexes.flatten_all()?; - let values = self.embedding.index_select(&indexes, 0)?; - let values = values.reshape(final_dims)?; - Ok(values) - } -} - -#[allow(unused)] -#[derive(Debug, Clone)] -pub struct VectorQuantization { - project_in: Option, - project_out: Option, - codebook: EuclideanCodebook, -} - -impl VectorQuantization { - pub fn new( - dim: usize, - codebook_size: usize, - codebook_dim: Option, - vb: VarBuilder, - ) -> Result { - let codebook_dim = codebook_dim.unwrap_or(dim); - let (project_in, project_out) = if codebook_dim == dim { - (None, None) - } else { - let p_in = linear(dim, codebook_dim, vb.pp("project_in"))?; - let p_out = linear(codebook_dim, dim, vb.pp("project_out"))?; - (Some(p_in), Some(p_out)) - }; - let codebook = EuclideanCodebook::new(codebook_dim, codebook_size, vb.pp("_codebook"))?; - Ok(Self { project_in, project_out, codebook }) - } - - pub fn encode(&self, xs: &Tensor) -> Result { - let xs = xs.t()?.apply(&self.project_in.as_ref())?; - self.codebook.encode_slow(&xs) - } - - pub fn decode(&self, codes: &Tensor) -> Result { - let quantized = self.codebook.decode(codes)?; - let quantized = match &self.project_out { - None => quantized, - Some(p) => quantized.apply(p)?, - }; - quantized.t() - } -} - -#[derive(Debug, Clone)] -pub struct ResidualVectorQuantization { - layers: Vec, -} - -impl ResidualVectorQuantization { - pub fn new( - n_q: usize, - dim: usize, - codebook_size: usize, - codebook_dim: Option, - vb: VarBuilder, - ) -> Result { - let vb = vb.pp("layers"); - let mut layers = Vec::with_capacity(n_q); - for i in 0..n_q { - let layer = VectorQuantization::new(dim, codebook_size, codebook_dim, vb.pp(i))?; - layers.push(layer) - } - Ok(Self { layers }) - } - - pub fn encode(&self, xs: &Tensor) -> Result { - let mut codes = Vec::with_capacity(self.layers.len()); - let mut residual = xs.clone(); - for layer in self.layers.iter() { - let indices = layer.encode(&residual)?; - let quantized = layer.decode(&indices)?; - residual = (residual - quantized)?; - codes.push(indices) - } - Tensor::stack(&codes, 0) - } - - pub fn decode(&self, xs: &Tensor) -> Result { - if self.layers.is_empty() { - candle::bail!("empty layers in ResidualVectorQuantization") - } - if self.layers.len() != xs.dim(0)? { - candle::bail!( - "mismatch between the number of layers {} and the code shape {:?}", - self.layers.len(), - xs.shape() - ) - } - let mut quantized = self.layers[0].decode(&xs.i(0)?)?; - for (i, layer) in self.layers.iter().enumerate().skip(1) { - let xs = xs.i(i)?; - quantized = (quantized + layer.decode(&xs))? - } - Ok(quantized) - } -} - -#[allow(unused)] -#[derive(Debug, Clone)] -pub struct ResidualVectorQuantizer { - vq: ResidualVectorQuantization, - input_proj: Option, - output_proj: Option, -} - -impl ResidualVectorQuantizer { - pub fn new( - dim: usize, - input_dim: Option, - output_dim: Option, - n_q: usize, - bins: usize, - force_projection: bool, - vb: VarBuilder, - ) -> Result { - let input_dim = input_dim.unwrap_or(dim); - let output_dim = output_dim.unwrap_or(dim); - - let input_proj = if input_dim == dim && !force_projection { - None - } else { - let c = candle_nn::conv1d_no_bias( - input_dim, - dim, - 1, - Default::default(), - vb.pp("input_proj"), - )?; - Some(c) - }; - let output_proj = if output_dim == dim && !force_projection { - None - } else { - let c = candle_nn::conv1d_no_bias( - dim, - output_dim, - 1, - Default::default(), - vb.pp("output_proj"), - )?; - Some(c) - }; - - let vq = ResidualVectorQuantization::new( - n_q, - dim, - /* codebook_size */ bins, - /* codebook_dim */ None, - vb.pp("vq"), - )?; - Ok(Self { vq, input_proj, output_proj }) - } - - pub fn encode(&self, xs: &Tensor) -> Result { - let codes = self.vq.encode(&xs.apply(&self.input_proj.as_ref())?)?; - codes.transpose(0, 1) - } - - pub fn decode(&self, codes: &Tensor) -> Result { - // codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T]. - let codes = codes.transpose(0, 1)?; - let quantized = self.vq.decode(&codes)?; - match &self.output_proj { - None => Ok(quantized), - Some(p) => quantized.apply(p), - } - } -} - -// we do not use any codebook_offset at the moment. When reconstructing the codes, we could just -// concatenate the indexes. -#[derive(Debug, Clone)] -pub struct SplitResidualVectorQuantizer { - rvq_first: ResidualVectorQuantizer, - rvq_rest: ResidualVectorQuantizer, - n_q: usize, - span_encode: tracing::Span, - span_decode: tracing::Span, -} - -impl SplitResidualVectorQuantizer { - pub fn new( - dim: usize, - input_dim: Option, - output_dim: Option, - n_q: usize, - bins: usize, - vb: VarBuilder, - ) -> Result { - let rvq_first = ResidualVectorQuantizer::new( - dim, - input_dim, - output_dim, - 1, - bins, - true, - vb.pp("rvq_first"), - )?; - let rvq_rest = ResidualVectorQuantizer::new( - dim, - input_dim, - output_dim, - n_q - 1, - bins, - true, - vb.pp("rvq_rest"), - )?; - let span_encode = tracing::span!(tracing::Level::TRACE, "split-rvq-encode"); - let span_decode = tracing::span!(tracing::Level::TRACE, "split-rvq-decode"); - Ok(Self { rvq_first, rvq_rest, n_q, span_encode, span_decode }) - } - - pub fn encode(&self, xs: &Tensor) -> Result { - let _enter = self.span_encode.enter(); - let codes = self.rvq_first.encode(xs)?; - if self.n_q > 1 { - // We encode xs again here rather than the residual. The decomposition is not - // hierarchical but rather having semantic tokens for rvq_first and the acoustic tokens - // for rvq_rest. - let rest_codes = self.rvq_rest.encode(xs)?; - Tensor::cat(&[codes, rest_codes], 1) - } else { - Ok(codes) - } - } - - pub fn decode(&self, codes: &Tensor) -> Result { - // codes is [B, K, T], with T frames, K nb of codebooks. - let _enter = self.span_decode.enter(); - let quantized = self.rvq_first.decode(&codes.i((.., ..1))?)?; - let quantized = if self.n_q > 1 { - (quantized + self.rvq_rest.decode(&codes.i((.., 1..))?))? - } else { - quantized - }; - Ok(quantized) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::streaming::{self, StreamMask, StreamTensor, StreamingModule}; -use candle::{Module, Result, Tensor}; -use candle_nn::VarBuilder; - -use crate::conv::{StreamableConv1d, StreamableConvTranspose1d}; - -#[derive(Debug, Clone)] -pub struct Config { - pub dimension: usize, - pub channels: usize, - pub causal: bool, - pub n_filters: usize, - pub n_residual_layers: usize, - pub ratios: Vec, - pub activation: candle_nn::Activation, - pub norm: crate::conv::Norm, - pub kernel_size: usize, - pub residual_kernel_size: usize, - pub last_kernel_size: usize, - pub dilation_base: usize, - pub pad_mode: crate::conv::PadMode, - pub true_skip: bool, - pub compress: usize, - pub lstm: usize, - pub disable_norm_outer_blocks: usize, - pub final_activation: Option, -} - -#[derive(Debug, Clone)] -pub struct SeaNetResnetBlock { - block: Vec, - shortcut: Option, - activation: candle_nn::Activation, - skip_op: streaming::StreamingBinOp, - span: tracing::Span, -} - -impl SeaNetResnetBlock { - #[allow(clippy::too_many_arguments)] - pub fn new( - dim: usize, - k_sizes_and_dilations: &[(usize, usize)], - activation: candle_nn::Activation, - norm: Option, - causal: bool, - pad_mode: crate::conv::PadMode, - compress: usize, - true_skip: bool, - vb: VarBuilder, - ) -> Result { - let mut block = Vec::with_capacity(k_sizes_and_dilations.len()); - let hidden = dim / compress; - let vb_b = vb.pp("block"); - for (i, (k_size, dilation)) in k_sizes_and_dilations.iter().enumerate() { - let in_c = if i == 0 { dim } else { hidden }; - let out_c = if i == k_sizes_and_dilations.len() - 1 { dim } else { hidden }; - let c = StreamableConv1d::new( - in_c, - out_c, - /* k_size */ *k_size, - /* stride */ 1, - /* dilation */ *dilation, - /* groups */ 1, - /* bias */ true, - /* causal */ causal, - /* norm */ norm, - /* pad_mode */ pad_mode, - vb_b.pp(2 * i + 1), - )?; - block.push(c) - } - let shortcut = if true_skip { - None - } else { - let c = StreamableConv1d::new( - dim, - dim, - /* k_size */ 1, - /* stride */ 1, - /* dilation */ 1, - /* groups */ 1, - /* bias */ true, - /* causal */ causal, - /* norm */ norm, - /* pad_mode */ pad_mode, - vb.pp("shortcut"), - )?; - Some(c) - }; - Ok(Self { - block, - shortcut, - activation, - skip_op: streaming::StreamingBinOp::new(streaming::BinOp::Add, candle::D::Minus1), - span: tracing::span!(tracing::Level::TRACE, "sea-resnet"), - }) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - for b in self.block.iter_mut() { - b.reset_batch_idx(batch_idx, batch_size)?; - } - if let Some(shortcut) = self.shortcut.as_mut() { - shortcut.reset_batch_idx(batch_idx, batch_size)?; - } - self.skip_op.reset_batch_idx(batch_idx, batch_size)?; - Ok(()) - } -} - -impl Module for SeaNetResnetBlock { - fn forward(&self, xs: &Tensor) -> Result { - let _enter = self.span.enter(); - let mut ys = xs.clone(); - for block in self.block.iter() { - ys = ys.apply(&self.activation)?.apply(block)?; - } - match self.shortcut.as_ref() { - None => ys + xs, - Some(shortcut) => ys + xs.apply(shortcut), - } - } -} - -impl StreamingModule for SeaNetResnetBlock { - fn reset_state(&mut self) { - self.skip_op.reset_state(); - for block in self.block.iter_mut() { - block.reset_state() - } - if let Some(shortcut) = self.shortcut.as_mut() { - shortcut.reset_state() - } - } - - fn step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - let _enter = self.span.enter(); - let mut ys = xs.clone(); - for block in self.block.iter_mut() { - ys = block.step(&ys.apply(&self.activation)?, m)?; - } - match self.shortcut.as_mut() { - None => self.skip_op.step(&ys, xs, m), - Some(shortcut) => self.skip_op.step(&ys, &shortcut.step(xs, m)?, m), - } - } -} - -#[derive(Debug, Clone)] -struct EncoderLayer { - residuals: Vec, - downsample: StreamableConv1d, -} - -#[derive(Debug, Clone)] -pub struct SeaNetEncoder { - init_conv1d: StreamableConv1d, - activation: candle_nn::Activation, - layers: Vec, - final_conv1d: StreamableConv1d, - span: tracing::Span, -} - -impl SeaNetEncoder { - pub fn new(cfg: &Config, vb: VarBuilder) -> Result { - if cfg.lstm > 0 { - candle::bail!("seanet lstm is not supported") - } - let n_blocks = 2 + cfg.ratios.len(); - let mut mult = 1usize; - let init_norm = if cfg.disable_norm_outer_blocks >= 1 { None } else { Some(cfg.norm) }; - let mut layer_idx = 0; - let vb = vb.pp("model"); - let init_conv1d = StreamableConv1d::new( - cfg.channels, - mult * cfg.n_filters, - cfg.kernel_size, - /* stride */ 1, - /* dilation */ 1, - /* groups */ 1, - /* bias */ true, - /* causal */ cfg.causal, - /* norm */ init_norm, - /* pad_mode */ cfg.pad_mode, - vb.pp(layer_idx), - )?; - layer_idx += 1; - let mut layers = Vec::with_capacity(cfg.ratios.len()); - - for (i, &ratio) in cfg.ratios.iter().rev().enumerate() { - let norm = if cfg.disable_norm_outer_blocks >= i + 2 { None } else { Some(cfg.norm) }; - let mut residuals = Vec::with_capacity(cfg.n_residual_layers); - for j in 0..cfg.n_residual_layers { - let resnet_block = SeaNetResnetBlock::new( - mult * cfg.n_filters, - &[(cfg.residual_kernel_size, cfg.dilation_base.pow(j as u32)), (1, 1)], - cfg.activation, - norm, - cfg.causal, - cfg.pad_mode, - cfg.compress, - cfg.true_skip, - vb.pp(layer_idx), - )?; - residuals.push(resnet_block); - layer_idx += 1; - } - let downsample = StreamableConv1d::new( - mult * cfg.n_filters, - mult * cfg.n_filters * 2, - /* k_size */ ratio * 2, - /* stride */ ratio, - /* dilation */ 1, - /* groups */ 1, - /* bias */ true, - /* causal */ true, - /* norm */ norm, - /* pad_mode */ cfg.pad_mode, - vb.pp(layer_idx + 1), - )?; - layer_idx += 2; - let layer = EncoderLayer { downsample, residuals }; - layers.push(layer); - mult *= 2 - } - - let final_norm = - if cfg.disable_norm_outer_blocks >= n_blocks { None } else { Some(cfg.norm) }; - let final_conv1d = StreamableConv1d::new( - mult * cfg.n_filters, - cfg.dimension, - cfg.last_kernel_size, - /* stride */ 1, - /* dilation */ 1, - /* groups */ 1, - /* bias */ true, - /* causal */ cfg.causal, - /* norm */ final_norm, - /* pad_mode */ cfg.pad_mode, - vb.pp(layer_idx + 1), - )?; - Ok(Self { - init_conv1d, - activation: cfg.activation, - layers, - final_conv1d, - span: tracing::span!(tracing::Level::TRACE, "sea-encoder"), - }) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - self.init_conv1d.reset_batch_idx(batch_idx, batch_size)?; - self.final_conv1d.reset_batch_idx(batch_idx, batch_size)?; - for layer in self.layers.iter_mut() { - layer.downsample.reset_batch_idx(batch_idx, batch_size)?; - for l in layer.residuals.iter_mut() { - l.reset_batch_idx(batch_idx, batch_size)?; - } - } - Ok(()) - } -} - -impl Module for SeaNetEncoder { - fn forward(&self, xs: &Tensor) -> Result { - let _enter = self.span.enter(); - let mut xs = xs.apply(&self.init_conv1d)?; - for layer in self.layers.iter() { - for residual in layer.residuals.iter() { - xs = xs.apply(residual)? - } - xs = xs.apply(&self.activation)?.apply(&layer.downsample)?; - } - xs.apply(&self.activation)?.apply(&self.final_conv1d) - } -} - -impl StreamingModule for SeaNetEncoder { - fn reset_state(&mut self) { - self.init_conv1d.reset_state(); - self.layers.iter_mut().for_each(|v| { - v.residuals.iter_mut().for_each(|v| v.reset_state()); - v.downsample.reset_state() - }); - self.final_conv1d.reset_state(); - } - - fn step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - let _enter = self.span.enter(); - let mut xs = self.init_conv1d.step(xs, m)?; - for layer in self.layers.iter_mut() { - for residual in layer.residuals.iter_mut() { - xs = residual.step(&xs, m)?; - } - xs = layer.downsample.step(&xs.apply(&self.activation)?, m)?; - } - self.final_conv1d.step(&xs.apply(&self.activation)?, m) - } -} - -#[derive(Debug, Clone)] -struct DecoderLayer { - upsample: StreamableConvTranspose1d, - residuals: Vec, -} - -#[derive(Debug, Clone)] -pub struct SeaNetDecoder { - init_conv1d: StreamableConv1d, - activation: candle_nn::Activation, - layers: Vec, - final_conv1d: StreamableConv1d, - final_activation: Option, - span: tracing::Span, -} - -impl SeaNetDecoder { - pub fn new(cfg: &Config, vb: VarBuilder) -> Result { - if cfg.lstm > 0 { - candle::bail!("seanet lstm is not supported") - } - let n_blocks = 2 + cfg.ratios.len(); - let mut mult = 1 << cfg.ratios.len(); - let init_norm = - if cfg.disable_norm_outer_blocks == n_blocks { None } else { Some(cfg.norm) }; - let mut layer_idx = 0; - let vb = vb.pp("model"); - let init_conv1d = StreamableConv1d::new( - cfg.dimension, - mult * cfg.n_filters, - cfg.kernel_size, - /* stride */ 1, - /* dilation */ 1, - /* groups */ 1, - /* bias */ true, - /* causal */ cfg.causal, - /* norm */ init_norm, - /* pad_mode */ cfg.pad_mode, - vb.pp(layer_idx), - )?; - layer_idx += 1; - let mut layers = Vec::with_capacity(cfg.ratios.len()); - for (i, &ratio) in cfg.ratios.iter().enumerate() { - let norm = if cfg.disable_norm_outer_blocks + i + 1 >= n_blocks { - None - } else { - Some(cfg.norm) - }; - let upsample = StreamableConvTranspose1d::new( - mult * cfg.n_filters, - mult * cfg.n_filters / 2, - /* k_size */ ratio * 2, - /* stride */ ratio, - /* groups */ 1, - /* bias */ true, - /* causal */ true, - /* norm */ norm, - vb.pp(layer_idx + 1), - )?; - layer_idx += 2; - - let mut residuals = Vec::with_capacity(cfg.n_residual_layers); - for j in 0..cfg.n_residual_layers { - let resnet_block = SeaNetResnetBlock::new( - mult * cfg.n_filters / 2, - &[(cfg.residual_kernel_size, cfg.dilation_base.pow(j as u32)), (1, 1)], - cfg.activation, - norm, - cfg.causal, - cfg.pad_mode, - cfg.compress, - cfg.true_skip, - vb.pp(layer_idx), - )?; - residuals.push(resnet_block); - layer_idx += 1; - } - let layer = DecoderLayer { upsample, residuals }; - layers.push(layer); - mult /= 2 - } - let final_norm = if cfg.disable_norm_outer_blocks >= 1 { None } else { Some(cfg.norm) }; - let final_conv1d = StreamableConv1d::new( - cfg.n_filters, - cfg.channels, - cfg.last_kernel_size, - /* stride */ 1, - /* dilation */ 1, - /* groups */ 1, - /* bias */ true, - /* causal */ cfg.causal, - /* norm */ final_norm, - /* pad_mode */ cfg.pad_mode, - vb.pp(layer_idx + 1), - )?; - Ok(Self { - init_conv1d, - activation: cfg.activation, - layers, - final_conv1d, - final_activation: cfg.final_activation, - span: tracing::span!(tracing::Level::TRACE, "sea-decoder"), - }) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - self.init_conv1d.reset_batch_idx(batch_idx, batch_size)?; - self.final_conv1d.reset_batch_idx(batch_idx, batch_size)?; - for layer in self.layers.iter_mut() { - layer.upsample.reset_batch_idx(batch_idx, batch_size)?; - for l in layer.residuals.iter_mut() { - l.reset_batch_idx(batch_idx, batch_size)?; - } - } - Ok(()) - } -} - -impl Module for SeaNetDecoder { - fn forward(&self, xs: &Tensor) -> Result { - let _enter = self.span.enter(); - let mut xs = xs.apply(&self.init_conv1d)?; - for layer in self.layers.iter() { - xs = xs.apply(&self.activation)?.apply(&layer.upsample)?; - for residual in layer.residuals.iter() { - xs = xs.apply(residual)? - } - } - let xs = xs.apply(&self.activation)?.apply(&self.final_conv1d)?; - let xs = match self.final_activation.as_ref() { - None => xs, - Some(act) => xs.apply(act)?, - }; - Ok(xs) - } -} - -impl StreamingModule for SeaNetDecoder { - fn reset_state(&mut self) { - self.init_conv1d.reset_state(); - self.layers.iter_mut().for_each(|v| { - v.residuals.iter_mut().for_each(|v| v.reset_state()); - v.upsample.reset_state() - }); - self.final_conv1d.reset_state(); - } - - fn step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - let _enter = self.span.enter(); - let mut xs = self.init_conv1d.step(xs, m)?; - for layer in self.layers.iter_mut() { - xs = layer.upsample.step(&xs.apply(&self.activation)?, m)?; - for residual in layer.residuals.iter_mut() { - xs = residual.step(&xs, m)?; - } - } - let xs = self.final_conv1d.step(&xs.apply(&self.activation)?, m)?; - let xs = match self.final_activation.as_ref() { - None => xs, - Some(act) => xs.apply(act)?, - }; - Ok(xs) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use candle::{Device, IndexOp, Result, Tensor}; - -pub trait Dim: candle::shape::Dim + Copy {} -impl Dim for T {} - -#[derive(Clone)] -pub struct StreamTensor(Option); - -#[derive(Debug, Clone)] -struct MaskInner { - cpu: Vec, - mask: Tensor, -} - -#[derive(Clone)] -pub struct StreamMask(Option); - -impl std::fmt::Debug for StreamMask { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match &self.0 { - Some(t) => write!(f, "{:?}", t.mask.shape()), - None => write!(f, "Empty"), - } - } -} - -impl std::convert::From<()> for StreamMask { - fn from(_value: ()) -> Self { - Self(None) - } -} - -impl StreamMask { - pub fn empty() -> Self { - Self(None) - } - - pub fn new(cpu: Vec, device: &Device) -> Result { - let mask = cpu.iter().map(|&v| u8::from(v)).collect::>(); - let mask = Tensor::new(mask, device)?; - Ok(Self(Some(MaskInner { cpu, mask }))) - } - - pub fn is_active(&self, batch_idx: usize) -> bool { - self.cpu().is_none_or(|v| v[batch_idx]) - } - - pub fn is_empty(&self) -> bool { - self.0.is_none() - } - - pub fn shape(&self) -> Option<&candle::Shape> { - self.0.as_ref().map(|t| t.mask.shape()) - } - - pub fn as_option(&self) -> Option<&Tensor> { - self.0.as_ref().map(|v| &v.mask) - } - - pub fn cpu(&self) -> Option<&[bool]> { - self.0.as_ref().map(|v| v.cpu.as_slice()) - } -} - -impl std::fmt::Debug for StreamTensor { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match &self.0 { - Some(t) => write!(f, "{:?}", t.shape()), - None => write!(f, "Empty"), - } - } -} - -impl std::convert::From> for StreamTensor { - fn from(value: Option) -> Self { - Self(value) - } -} - -impl std::convert::From for StreamTensor { - fn from(value: Tensor) -> Self { - Self(Some(value)) - } -} - -impl std::convert::From<()> for StreamTensor { - fn from(_value: ()) -> Self { - Self(None) - } -} - -impl StreamTensor { - pub fn empty() -> Self { - Self(None) - } - - pub fn is_empty(&self) -> bool { - self.0.is_none() - } - - pub fn from_tensor(tensor: Tensor) -> Self { - Self(Some(tensor)) - } - - pub fn shape(&self) -> Option<&candle::Shape> { - self.0.as_ref().map(|t| t.shape()) - } - - pub fn cat2(&self, rhs: &Self, dim: D) -> Result { - let xs = match (&self.0, &rhs.0) { - (Some(lhs), Some(rhs)) => { - let xs = Tensor::cat(&[lhs, rhs], dim)?; - Some(xs) - } - (Some(xs), None) | (None, Some(xs)) => Some(xs.clone()), - (None, None) => None, - }; - Ok(Self(xs)) - } - - pub fn seq_len(&self, dim: D) -> Result { - match &self.0 { - None => Ok(0), - Some(v) => v.dim(dim), - } - } - - pub fn reset(&mut self) { - self.0 = None - } - - pub fn narrow(&self, dim: D, offset: usize, len: usize) -> Result { - let t = match &self.0 { - None => None, - Some(t) => { - let seq_len = t.dim(dim)?; - if seq_len <= offset { - None - } else { - let t = t.narrow(dim, offset, usize::min(len, seq_len - offset))?; - Some(t) - } - } - }; - Ok(Self(t)) - } - - /// Splits the Streaming Tensor on the time axis `dim` with the first `lhs_len` elements - /// returned in the first output and the remaining in the second output. - pub fn split(&self, dim: D, lhs_len: usize) -> Result<(Self, Self)> { - match &self.0 { - None => Ok((Self::empty(), Self::empty())), - Some(t) => { - let seq_len = t.dim(dim)?; - let lhs_len = usize::min(seq_len, lhs_len); - if lhs_len == 0 { - Ok((Self::empty(), t.clone().into())) - } else { - let lhs = Self::from_tensor(t.narrow(dim, 0, lhs_len)?); - let rhs_len = seq_len - lhs_len; - let rhs = if rhs_len == 0 { - Self::empty() - } else { - Self::from_tensor(t.narrow(dim, lhs_len, rhs_len)?) - }; - Ok((lhs, rhs)) - } - } - } - } - - pub fn as_option(&self) -> Option<&Tensor> { - self.0.as_ref() - } - - pub fn apply(&self, m: &M) -> Result { - match &self.0 { - None => Ok(Self::empty()), - Some(t) => Ok(Self::from_tensor(t.apply(m)?)), - } - } -} - -pub trait StreamingModule { - // TODO: Should we also have a flush method? - fn step(&mut self, xs: &StreamTensor, mask: &StreamMask) -> Result; - fn reset_state(&mut self); -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub enum BinOp { - Add, - Mul, - Sub, - Div, -} - -#[derive(Debug, Clone)] -pub struct StreamingBinOp { - prev_lhs: StreamTensor, - prev_rhs: StreamTensor, - pub op: BinOp, - pub dim: candle::D, -} - -impl StreamingBinOp { - pub fn new(op: BinOp, dim: candle::D) -> Self { - Self { prev_lhs: StreamTensor::empty(), prev_rhs: StreamTensor::empty(), op, dim } - } - - pub fn reset_state(&mut self) { - self.prev_lhs.reset(); - self.prev_rhs.reset(); - } - - pub fn forward(&self, lhs: &Tensor, rhs: &Tensor) -> Result { - match self.op { - BinOp::Add => Tensor::add(lhs, rhs), - BinOp::Mul => Tensor::mul(lhs, rhs), - BinOp::Sub => Tensor::sub(lhs, rhs), - BinOp::Div => Tensor::div(lhs, rhs), - } - } - - pub fn step( - &mut self, - lhs: &StreamTensor, - rhs: &StreamTensor, - mask: &StreamMask, - ) -> Result { - let lhs = StreamTensor::cat2(&self.prev_lhs, lhs, self.dim)?; - let rhs = StreamTensor::cat2(&self.prev_rhs, rhs, self.dim)?; - let lhs_len = lhs.seq_len(self.dim)?; - let rhs_len = rhs.seq_len(self.dim)?; - let common_len = usize::min(lhs_len, rhs_len); - let (lhs, prev_lhs) = lhs.split(self.dim, common_len)?; - let (rhs, prev_rhs) = rhs.split(self.dim, common_len)?; - let ys = match (&lhs.0, &rhs.0) { - (Some(lhs), Some(rhs)) => { - let ys = self.forward(lhs, rhs)?; - StreamTensor::from_tensor(ys) - } - (None, None) => StreamTensor::empty(), - (lhs, rhs) => candle::bail!("INTERNAL ERROR inconsistent lhs and rhs {lhs:?} {rhs:?}"), - }; - if !mask.is_empty() && (!prev_lhs.is_empty() || !prev_rhs.is_empty()) { - candle::bail!( - "cannot use a stream mask with a streaming bin op {prev_lhs:?} {prev_rhs:?} {lhs:?} {rhs:?}" - ); - } - self.prev_lhs = prev_lhs; - self.prev_rhs = prev_rhs; - Ok(ys) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, _batch_size: usize) -> Result<()> { - if let Some(v) = self.prev_lhs.as_option() { - let v = v.contiguous()?; - v.i(batch_idx..(1 + batch_idx))?.zero_set()?; - self.prev_lhs = StreamTensor::from_tensor(v); - } - if let Some(v) = self.prev_rhs.as_option() { - let v = v.contiguous()?; - v.i(batch_idx..(1 + batch_idx))?.zero_set()?; - self.prev_rhs = StreamTensor::from_tensor(v); - } - Ok(()) - } -} - -/// Simple wrapper that doesn't do any buffering. -pub struct Map(T); - -impl StreamingModule for Map { - fn reset_state(&mut self) {} - - fn step(&mut self, xs: &StreamTensor, _: &StreamMask) -> Result { - xs.apply(&self.0) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -// Implements various modules for transformers with support for both quantized and unquantized forwards -// Main differences between quantized and unquantized execution: -// 1. For quantized models' attention `matmul_dtype`` converts intermediate activations to BF16 for -// more efficient matmuls -// 2. Quantized tensors cannot be easily split (regarding cross attention and QKV proj weights) -// 3. Linear and Quantized linear layers are two different types -use crate::nn::{ - linear, linear_from, matmul_dtype, MaybeQuantizedLinear, MaybeQuantizedVarBuilder, -}; -use crate::streaming::{StreamMask, StreamTensor, StreamingModule}; -use candle::{DType, Device, IndexOp, Module, Result, Tensor, D}; - -use crate::kv_cache::KvCache; -use candle::Context; - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct Config { - pub d_model: usize, - pub num_heads: usize, - pub num_layers: usize, - pub causal: bool, - pub norm_first: bool, - pub bias_ff: bool, - pub bias_attn: bool, - pub layer_scale: Option, - pub positional_embedding: PositionalEmbedding, - pub use_conv_block: bool, - pub cross_attention: Option<(CrossAttentionGating, crate::NormType, Option)>, - pub conv_kernel_size: usize, - pub use_conv_bias: bool, - pub gating: Option, - pub norm: crate::NormType, - pub context: usize, - pub max_period: usize, - pub max_seq_len: usize, - - pub kv_repeat: usize, - pub dim_feedforward: usize, - pub conv_layout: bool, - - #[serde(default)] - pub shared_cross_attn: bool, -} - -#[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)] -pub enum PositionalEmbedding { - Rope, - Sin, - None, -} - -#[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -pub enum CrossAttentionGating { - // Configure Type of gating used at the output of vision cross-attention layers - Normal, - ConstantGatedTanh, - ConstantGatedSigmoid, - ConditionalGatedTanh, - ConditionalGatedSigmoid, - ConditionalGatedSigmoidLearnableBias, - ConditionalGatedTanhLearnableBias, -} - -#[derive(Debug, Clone)] -pub enum CaSrc { - // Input to cross-attention to handle cases where the - // cross-attention source can be shared across timesteps and/or layers - // either a single tensor (has yet to be projected) - // or pre-computed K,V projections; - Tokens(Tensor), - KeysValues((Tensor, Tensor)), -} - -#[derive(Debug, Clone)] -pub struct LayerScale { - scale: Tensor, -} - -impl LayerScale { - pub fn new(d_model: usize, _init: f64, vb: MaybeQuantizedVarBuilder) -> Result { - let scale = vb.get_unquantized(d_model, "scale")?; - Ok(Self { scale }) - } -} - -impl Module for LayerScale { - fn forward(&self, xs: &Tensor) -> Result { - xs.broadcast_mul(&self.scale) - } -} - -#[derive(Debug, Clone)] -pub enum XaGate { - // Optional gating at the output of a cross-attention layer - // Normal: No gating | Identity - Normal, - // ConstantGated: Multiply by a scalar - ConstantGated { - alpha: Tensor, - }, - // ConditionalGated: Pass the input x through a small MLP; - // The output yields a vector of scales (one for each channel) - // that x is then multiplied by - ConditionalGated { - in_proj: MaybeQuantizedLinear, - out_proj: MaybeQuantizedLinear, - activation: candle_nn::init::NonLinearity, - learnable_bias: bool, - }, -} - -impl XaGate { - pub fn new(cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - let gating_cfg = - cfg.cross_attention.map(|v| v.0).context("no cross-attention specified")?; - match gating_cfg { - // no gating - CrossAttentionGating::Normal => Ok(Self::Normal), - // constant (per-layer parameter) with tanh activation - CrossAttentionGating::ConstantGatedTanh => { - let alpha = vb.get_unquantized((1, 1, 1), "alpha")?.tanh()?; - Ok(Self::ConstantGated { alpha }) - } - // constant (per-layer parameter) with sigmoid activation - CrossAttentionGating::ConstantGatedSigmoid => { - let alpha = - candle_nn::ops::sigmoid(&(vb.get_unquantized((1, 1, 1), "alpha")? - 4.0)?)?; - Ok(Self::ConstantGated { alpha }) - } - // input conditional (small MLP) with tanh or sigmoid act - CrossAttentionGating::ConditionalGatedTanh - | CrossAttentionGating::ConditionalGatedSigmoid - | CrossAttentionGating::ConditionalGatedSigmoidLearnableBias - | CrossAttentionGating::ConditionalGatedTanhLearnableBias => { - let dim = cfg.d_model; - let hidden_dims = (0.125 * dim as f32).floor() as usize; - let learnable_bias = matches!( - gating_cfg, - CrossAttentionGating::ConditionalGatedSigmoidLearnableBias - | CrossAttentionGating::ConditionalGatedTanhLearnableBias - ); - let in_proj = linear(dim, hidden_dims, false, vb.pp("alpha.0"))?; - let out_proj = linear(hidden_dims, dim, learnable_bias, vb.pp("alpha.2"))?; - let activation = match gating_cfg { - CrossAttentionGating::ConditionalGatedTanh - | CrossAttentionGating::ConditionalGatedTanhLearnableBias => { - candle_nn::init::NonLinearity::Tanh - } - CrossAttentionGating::ConditionalGatedSigmoid - | CrossAttentionGating::ConditionalGatedSigmoidLearnableBias => { - candle_nn::init::NonLinearity::Sigmoid - } - _ => candle::bail!("Invalid cross-attention config specified."), - }; - Ok(Self::ConditionalGated { in_proj, out_proj, activation, learnable_bias }) - } - } - } -} - -impl Module for XaGate { - fn forward(&self, xs: &Tensor) -> Result { - match self { - Self::Normal => Ok(xs.clone()), - Self::ConstantGated { alpha } => xs.broadcast_mul(alpha), - Self::ConditionalGated { in_proj, out_proj, activation, learnable_bias } => { - let alpha = xs.apply(in_proj)?.relu()?.apply(out_proj)?; - let alpha = match (activation, learnable_bias) { - (candle_nn::init::NonLinearity::Tanh, _) => alpha.tanh(), - (candle_nn::init::NonLinearity::Sigmoid, true) => { - candle_nn::ops::sigmoid(&alpha) - } - (candle_nn::init::NonLinearity::Sigmoid, false) => { - candle_nn::ops::sigmoid(&(alpha - 4.0)?) - } - _ => candle::bail!("Invalid non-linearity specified in cross-attention gating"), - }; - xs * alpha? - } - } - } -} - -#[derive(Debug, Clone)] -pub struct StreamingMultiheadCrossAttention { - //Cross-attention modules. Q and KV projections are separate - // because x (speech tokens) and ca_src (cross-attention source) can have - // different dimensions - in_proj_q: MaybeQuantizedLinear, - in_proj_kv: MaybeQuantizedLinear, - out_proj: MaybeQuantizedLinear, - kv_repeat: usize, - num_heads: usize, - gate: XaGate, - span: tracing::Span, -} - -impl StreamingMultiheadCrossAttention { - pub fn new( - cfg: &Config, - vb: MaybeQuantizedVarBuilder, - gate_vb: Option, - ) -> Result { - let embed_dim = cfg.d_model; - let num_kv = cfg.num_heads / cfg.kv_repeat; - let out_kv_dim = num_kv * (embed_dim / cfg.num_heads); - let out_dim = embed_dim + 2 * out_kv_dim; - // Case 1 (legacy): A single in_proj; i.e., both x and ca_src *must* have - // the same number of dims this is only possible for non-quantized tensors though - // as we will need to split Q/KV weights down the line even when they have the same - // shape since they take different inputs - let (in_proj_q, in_proj_kv) = if vb.contains_key("in_proj_weight") { - match &vb { - MaybeQuantizedVarBuilder::Quantized(_) => candle::bail!("Quantized cross-attention layers require a separate in_proj_weight_q and in_proj_weight_kv"), - MaybeQuantizedVarBuilder::Real(weights) => { - let in_proj_weight = weights.get((out_dim, embed_dim), "in_proj_weight")?; - let in_proj_weight_q = in_proj_weight.narrow(0, 0, embed_dim)?; - let in_proj_weight_kv = in_proj_weight.narrow(0, embed_dim, 2 * out_kv_dim)?; - let (in_proj_bias_q, in_proj_bias_kv) = if cfg.bias_attn { - let b = weights.get(out_dim, "in_proj_bias")?; - let in_proj_bias_q = b.narrow(0, 0, embed_dim)?; - let in_proj_bias_kv = b.narrow(0, embed_dim, 2 * out_kv_dim)?; - (Some(in_proj_bias_q), Some(in_proj_bias_kv)) - } else { - (None, None) - }; - (MaybeQuantizedLinear::Real(candle_nn::Linear::new(in_proj_weight_q, in_proj_bias_q)), - MaybeQuantizedLinear::Real(candle_nn::Linear::new(in_proj_weight_kv, in_proj_bias_kv))) - - } - } - } else { - // Case 2: Separate projections for query (x) and kv (ca_src) - let kv_in_dim = match cfg.cross_attention.map(|v| v.2) { - None => candle::bail!("cfg.cross_attention is None in cross_attention module"), - Some(d) => match d { - None | Some(0) => embed_dim, - Some(dd) => dd, - }, - }; - let in_proj_weight_q = vb.get((embed_dim, embed_dim), "in_proj_weight_q")?; - let in_proj_weight_kv = vb.get((2 * out_kv_dim, kv_in_dim), "in_proj_weight_kv")?; - - // Biases are always unquantized - let (in_proj_bias_q, in_proj_bias_kv) = if cfg.bias_attn { - ( - Some(vb.get_unquantized(embed_dim, "in_proj_bias_q")?), - Some(vb.get_unquantized(2 * out_kv_dim, "in_proj_bias_kv")?), - ) - } else { - (None, None) - }; - - // Finally, we can build the actual linear layers - let in_proj_q = linear_from(in_proj_weight_q, in_proj_bias_q)?; - let in_proj_kv = linear_from(in_proj_weight_kv, in_proj_bias_kv)?; - (in_proj_q, in_proj_kv) - }; - - let out_proj = linear(embed_dim, embed_dim, cfg.bias_attn, vb.pp("out_proj"))?; - let gate = match gate_vb { - None => XaGate::new(cfg, vb.pp("gate"))?, - Some(layer_gate_vb) => XaGate::new(cfg, layer_gate_vb)?, - }; - Ok(Self { - in_proj_q, - in_proj_kv, - out_proj, - kv_repeat: cfg.kv_repeat, - num_heads: cfg.num_heads, - gate, - span: tracing::span!(tracing::Level::TRACE, "mhca"), - }) - } - - pub fn is_quantized(&self) -> bool { - match self.in_proj_q { - MaybeQuantizedLinear::Quantized(_) => true, - MaybeQuantizedLinear::Real(_) => false, - } - } - - pub fn compute_kv(&self, ca_src: &CaSrc) -> Result<(Tensor, Tensor)> { - // this is used twice: - // in the standard forward pass of the cross-attention - // for vision models, after loading an image we can precompute its KV projections - // as the image is constant across multiple timesteps - match ca_src { - CaSrc::KeysValues(cakv) => Ok(cakv.clone()), - CaSrc::Tokens(xs) => { - let kv = xs.apply(&self.in_proj_kv)?; - let (ca_b, ca_t, ca_dim) = kv.dims3()?; - let head_dim = ca_dim / (2 * self.num_heads); - let kv = kv.reshape((ca_b, ca_t, 2, (), head_dim))?; - // convert to correct float point type for quantized models - let kv = - if self.is_quantized() { kv.to_dtype(matmul_dtype(xs.device()))? } else { kv }; - let k = kv.i((.., .., 0))?; - let v = kv.i((.., .., 1))?; - let k = k.transpose(1, 2)?.contiguous()?; // b,h,k,d - let v = v.transpose(1, 2)?.contiguous()?; // b,h,k,d - Ok((k, v)) - } - } - } - - pub fn forward(&self, xs: &Tensor, ca_src: &CaSrc, mask: Option<&Tensor>) -> Result { - let _enter = self.span.enter(); - if self.kv_repeat != 1 { - candle::bail!("only kv-repeat = 1 is supported") - } - let (b, t, hd) = xs.dims3()?; - let head_dim = hd / self.num_heads; - // time_dim = 1, layout: b,t,h,d - let q = xs.apply(&self.in_proj_q)?; - let original_dtype = q.dtype(); - let q = q.reshape((b, t, self.num_heads, head_dim))?; - let q = if self.is_quantized() { q.to_dtype(matmul_dtype(xs.device()))? } else { q }; - let (k, v) = self.compute_kv(ca_src)?; - // qk_layer_norm = None - // kv_repeat = 1, otherwise we would need repeat_kv - let q = q.transpose(1, 2)?.contiguous()?; // b,h,t,d - - let pre_ws = q.matmul(&k.t()?)?; // b,h,t,k - let pre_ws = (pre_ws * (head_dim as f64).powf(-0.5))?; - - let pre_ws = match mask { - None => pre_ws, - Some(mask) => pre_ws.broadcast_add(mask)?, - }; - - let ws = candle_nn::ops::softmax_last_dim(&pre_ws)?; // b,h,t,k - let xs = ws.matmul(&v)?; // b,h,t,d - let xs = xs - .transpose(1, 2)? // b,t,h,d - .reshape((b, t, hd))? - .to_dtype(original_dtype)? - .apply(&self.out_proj)? - .apply(&self.gate)?; - Ok(xs) - } -} - -#[derive(Debug, Clone)] -pub struct Rope { - sin: Tensor, - cos: Tensor, -} - -impl Rope { - pub fn apply_rotary_emb(&self, qk: &Tensor) -> Result { - let qk_dtype = qk.dtype(); - candle_nn::rotary_emb::rope_i(&qk.to_dtype(DType::F32)?, &self.cos, &self.sin)? - .to_dtype(qk_dtype) - } -} - -#[derive(Debug, Clone)] -pub struct RotaryEmbedding { - inv_freq: Tensor, -} - -impl RotaryEmbedding { - pub fn new(dim: usize, theta: f32, dev: &Device) -> Result { - let inv_freq: Vec<_> = - (0..dim).step_by(2).map(|i| 1f32 / theta.powf(i as f32 / dim as f32)).collect(); - let inv_freq_len = inv_freq.len(); - let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; - Ok(Self { inv_freq }) - } - - pub fn rope(&self, pos: &Tensor) -> Result { - let t = pos.to_dtype(DType::F32)?; - let freqs = match *t.dims() { - [d] => t.reshape((d, 1))?.matmul(&self.inv_freq)?, - [b, d] => t.reshape((b * d, 1))?.matmul(&self.inv_freq)?.reshape((b, d, ()))?, - _ => candle::bail!("Invalid shape for rotary embedding {pos:?}"), - }; - Ok(Rope { sin: freqs.sin()?, cos: freqs.cos()? }) - } -} - -#[cfg(feature = "flash-attn")] -fn flash_attn( - q: &Tensor, - k: &Tensor, - v: &Tensor, - softmax_scale: f32, - causal: bool, -) -> Result { - candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) -} - -#[cfg(not(feature = "flash-attn"))] -fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { - unimplemented!("compile with '--features flash-attn'") -} - -#[derive(Debug, Clone)] -pub struct StreamingMultiheadAttention { - // Self-attention with KV Cache - in_proj: MaybeQuantizedLinear, - out_proj: MaybeQuantizedLinear, - kv_repeat: usize, - num_heads: usize, - context: usize, - kv_cache: KvCache, - use_flash_attn: bool, - span: tracing::Span, -} - -impl StreamingMultiheadAttention { - pub fn new(cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - let embed_dim = cfg.d_model; - let num_kv = cfg.num_heads / cfg.kv_repeat; - let out_dim = embed_dim + 2 * num_kv * (embed_dim / cfg.num_heads); - let in_proj_weight = vb.get((out_dim, embed_dim), "in_proj_weight")?; - let in_proj_bias = - if cfg.bias_attn { Some(vb.get_unquantized(out_dim, "in_proj_bias")?) } else { None }; - let in_proj = linear_from(in_proj_weight, in_proj_bias)?; - let out_proj = linear(embed_dim, embed_dim, cfg.bias_attn, vb.pp("out_proj"))?; - Ok(Self { - in_proj, - out_proj, - kv_repeat: cfg.kv_repeat, - num_heads: cfg.num_heads, - context: cfg.context, - kv_cache: KvCache::new(2, cfg.context), - use_flash_attn: false, - span: tracing::span!(tracing::Level::TRACE, "mha"), - }) - } - - pub fn is_quantized(&self) -> bool { - match self.in_proj { - MaybeQuantizedLinear::Quantized(_) => true, - MaybeQuantizedLinear::Real(_) => false, - } - } - - pub fn forward( - &mut self, - xs: &Tensor, - rope: Option<&Rope>, - mask: Option<&Tensor>, - ) -> Result { - let _enter = self.span.enter(); - if self.kv_repeat != 1 { - candle::bail!("only kv-repeat = 1 is supported") - } - let (b, t, hd) = xs.dims3()?; - let head_dim = hd / self.num_heads; - // time_dim = 1, layout: b,t,h,d - let qkv = xs.apply(&self.in_proj)?.reshape((b, t, 3, self.num_heads, head_dim))?; - let original_dtype = qkv.dtype(); - let qkv = if self.is_quantized() { qkv.to_dtype(matmul_dtype(xs.device()))? } else { qkv }; - let q = qkv.i((.., .., 0))?; - let k = qkv.i((.., .., 1))?; - let v = qkv.i((.., .., 2))?; - // qk_layer_norm = None - // kv_repeat = 1, otherwise we would need repeat_kv - let mut q = q.transpose(1, 2)?.contiguous()?; // b,h,t,d - let mut k = k.transpose(1, 2)?.contiguous()?; // b,h,k,d - let v = v.transpose(1, 2)?.contiguous()?; // b,h,k,d - if let Some(rope) = rope.as_ref() { - q = rope.apply_rotary_emb(&q)?; - k = rope.apply_rotary_emb(&k)?; - } - - let (k, v) = { self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)? }; - // The KV cache keeps all the data at the moment, we want to trim - // down the part that comes from the cache to at most context to - // be coherent with the mask shape we provide. - let k_len = k.dim(2)?; - let k_target_len = t + usize::min(self.context, k_len - t); - let (k, v) = if k_target_len < k_len { - let k = k.narrow(2, k_len - k_target_len, k_target_len)?; - let v = v.narrow(2, k_len - k_target_len, k_target_len)?; - (k, v) - } else { - (k.clone(), v.clone()) - }; - - let xs = if q.dtype() == DType::BF16 && self.use_flash_attn { - let q = q.transpose(1, 2)?; - let k = k.transpose(1, 2)?; - let v = v.transpose(1, 2)?; - let softmax_scale = 1f32 / (head_dim as f32).sqrt(); - flash_attn(&q, &k, &v, softmax_scale, mask.is_some())?.transpose(1, 2)? - } else { - let pre_ws = q.matmul(&k.t()?)?; // b,h,t,k - let pre_ws = (pre_ws * (head_dim as f64).powf(-0.5))?; - - let pre_ws = match mask { - None => pre_ws, - Some(mask) => pre_ws.broadcast_add(mask)?, - }; - - let ws = candle_nn::ops::softmax_last_dim(&pre_ws)?; // b,h,t,k - ws.matmul(&v)? // b,h,t,d - }; - - let xs = xs - .transpose(1, 2)? // b,t,h,d - .reshape((b, t, hd))? - .to_dtype(original_dtype)? - .apply(&self.out_proj)?; - Ok(xs) - } - - pub fn reset_kv_cache(&mut self) { - self.kv_cache.reset() - } - - pub fn set_kv_cache(&mut self, kv_cache: KvCache) { - self.kv_cache = kv_cache - } -} - -#[derive(Debug, Clone)] -pub enum Mlp { - //Feed Forward layers - NoGating { - linear1: MaybeQuantizedLinear, - linear2: MaybeQuantizedLinear, - }, - Gating { - linear_in: MaybeQuantizedLinear, - linear_out: MaybeQuantizedLinear, - activation: candle_nn::Activation, - }, -} - -impl Mlp { - pub fn new(cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - let d_model = cfg.d_model; - match cfg.gating { - None => { - let linear1 = linear(d_model, cfg.dim_feedforward, cfg.bias_ff, vb.pp("linear1"))?; - let linear2 = linear(cfg.dim_feedforward, d_model, cfg.bias_ff, vb.pp("linear2"))?; - Ok(Self::NoGating { linear1, linear2 }) - } - Some(activation) => { - let vb = vb.pp("gating"); - let hidden = if cfg.dim_feedforward == 4 * d_model { - 11 * d_model / 4 - } else { - 2 * cfg.dim_feedforward / 3 - }; - let linear_in = linear(d_model, 2 * hidden, cfg.bias_ff, vb.pp("linear_in"))?; - let linear_out = linear(hidden, d_model, cfg.bias_ff, vb.pp("linear_out"))?; - Ok(Self::Gating { linear_in, linear_out, activation }) - } - } - } -} - -impl Module for Mlp { - fn forward(&self, xs: &Tensor) -> Result { - match self { - Self::NoGating { linear1, linear2 } => xs.apply(linear1)?.gelu_erf()?.apply(linear2), - Self::Gating { linear_in, linear_out, activation } => { - let xs = xs.apply(linear_in)?; - let (b, t, _) = xs.dims3()?; - let xs = xs.reshape((b, t, 2, ()))?; - let xs = (xs.i((.., .., 0))?.apply(activation)? * xs.i((.., .., 1))?)?; - xs.apply(linear_out) - } - } - } -} - -#[derive(Debug, Clone)] -pub struct RmsNorm { - pub(crate) alpha: Tensor, - pub(crate) eps: f32, -} - -impl RmsNorm { - pub fn new(d_model: usize, eps: f32, vb: MaybeQuantizedVarBuilder) -> Result { - let alpha = vb.get_unquantized((1, 1, d_model), "alpha")?.reshape(d_model)?; - Ok(Self { alpha, eps }) - } -} - -impl Module for RmsNorm { - fn forward(&self, xs: &Tensor) -> Result { - candle_nn::ops::rms_norm(xs, &self.alpha, self.eps) - } -} - -#[derive(Debug, Clone)] -pub struct LayerNorm { - inner: candle_nn::LayerNorm, -} - -impl LayerNorm { - pub fn new(d_model: usize, eps: f32, vb: MaybeQuantizedVarBuilder) -> Result { - let bias = vb.get_unquantized(d_model, "bias")?; - let alpha = if vb.contains_key("alpha") { - vb.get_unquantized((1, 1, d_model), "alpha")?.reshape(d_model)? - } else { - vb.get_unquantized(d_model, "weight")?.reshape(d_model)? - }; - let inner = candle_nn::LayerNorm::new(alpha, bias, eps as f64); - Ok(Self { inner }) - } -} - -impl Module for LayerNorm { - fn forward(&self, xs: &Tensor) -> Result { - self.inner.forward(xs) - } -} - -#[derive(Debug, Clone)] -pub enum Norm { - LayerNorm(LayerNorm), - RmsNorm(RmsNorm), -} - -impl Norm { - pub fn new(d_model: usize, cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - let norm = Self::new_shortcut(d_model, cfg.norm, vb)?; - Ok(norm) - } - - pub fn new_shortcut( - d_model: usize, - typ: crate::NormType, - vb: MaybeQuantizedVarBuilder, - ) -> Result { - let norm = match typ { - crate::NormType::LayerNorm => { - let norm = LayerNorm::new(d_model, 1e-5, vb)?; - Self::LayerNorm(norm) - } - crate::NormType::RmsNorm => { - let norm = RmsNorm::new(d_model, 1e-8, vb)?; - Self::RmsNorm(norm) - } - }; - Ok(norm) - } -} - -impl Module for Norm { - fn forward(&self, xs: &Tensor) -> Result { - match self { - Self::LayerNorm(m) => m.forward(xs), - Self::RmsNorm(m) => m.forward(xs), - } - } -} - -#[derive(Debug, Clone)] -pub struct StreamingTransformerLayer { - self_attn: StreamingMultiheadAttention, - mlp: Mlp, - norm1: Norm, - norm2: Norm, - layer_scale_1: Option, - layer_scale_2: Option, - cross_attn: Option<(Norm, StreamingMultiheadCrossAttention)>, - norm_first: bool, - span: tracing::Span, -} - -impl StreamingTransformerLayer { - pub fn new( - cfg: &Config, - vb: MaybeQuantizedVarBuilder, - shared_ca_vb: Option, - ) -> Result { - if cfg.use_conv_block { - candle::bail!("conv-block is not supported") - } - let d_model = cfg.d_model; - let mlp = Mlp::new(cfg, vb.clone())?; - let norm1 = Norm::new(d_model, cfg, vb.pp("norm1"))?; - let norm2 = Norm::new(d_model, cfg, vb.pp("norm2"))?; - let layer_scale_1 = match cfg.layer_scale { - None => None, - Some(ls) => { - let ls = LayerScale::new(d_model, ls, vb.pp("layer_scale_1"))?; - Some(ls) - } - }; - let layer_scale_2 = match cfg.layer_scale { - None => None, - Some(ls) => { - let ls = LayerScale::new(d_model, ls, vb.pp("layer_scale_2"))?; - Some(ls) - } - }; - let self_attn = StreamingMultiheadAttention::new(cfg, vb.pp("self_attn"))?; - let cross_attn = match cfg.cross_attention.map(|v| v.1) { - Some(norm_type) => { - let norm_cross = Norm::new_shortcut(d_model, norm_type, vb.pp("norm_cross"))?; - let cross_attn = match shared_ca_vb { - None => { - StreamingMultiheadCrossAttention::new(cfg, vb.pp("cross_attention"), None)? - } - Some(shared_vb) => StreamingMultiheadCrossAttention::new( - cfg, - shared_vb.pp("cross_attention"), - Some(vb.pp("cross_attention.gate")), - )?, - }; - Some((norm_cross, cross_attn)) - } - None => None, - }; - Ok(Self { - self_attn, - mlp, - norm1, - norm2, - layer_scale_1, - layer_scale_2, - cross_attn, - norm_first: cfg.norm_first, - span: tracing::span!(tracing::Level::TRACE, "transformer-layer"), - }) - } - - pub fn forward( - &mut self, - xs: &Tensor, - rope: Option<&Rope>, - ca_src: Option<&CaSrc>, - mask: Option<&Tensor>, - ) -> Result { - let _enter = self.span.enter(); - if !self.norm_first { - candle::bail!("only norm_first = true is supported") - } - let norm1 = xs.apply(&self.norm1)?; - let xs = (xs - + self.self_attn.forward(&norm1, rope, mask)?.apply(&self.layer_scale_1.as_ref())?)?; - - let xs = match (self.cross_attn.as_mut(), ca_src) { - (Some((norm_cross, cross_attn)), Some(ca_src)) => { - let residual = &xs; - let xs = xs.apply(norm_cross)?; - (residual + cross_attn.forward(&xs, ca_src, None)?)? - } - _ => xs, - }; - - let xs = - (&xs + xs.apply(&self.norm2)?.apply(&self.mlp)?.apply(&self.layer_scale_2.as_ref()))?; - Ok(xs) - } - - pub fn reset_kv_cache(&mut self) { - self.self_attn.reset_kv_cache(); - } - - pub fn set_kv_cache(&mut self, kv_cache: KvCache) { - self.self_attn.set_kv_cache(kv_cache); - } -} - -#[derive(Debug, Clone)] -pub struct StreamingTransformer { - // Main transformer - layers: Vec, - positional_embedding: PositionalEmbedding, - max_period: usize, - causal: bool, - num_heads: usize, - context: usize, - last_reset_pos: Vec, - rope: Option, -} - -impl StreamingTransformer { - pub fn new(cfg: &Config, vb: MaybeQuantizedVarBuilder) -> Result { - let vb_l = vb.pp("layers"); - let rope = match cfg.positional_embedding { - PositionalEmbedding::Rope => { - let rope = RotaryEmbedding::new( - cfg.d_model / cfg.num_heads, - cfg.max_period as f32, - vb.device(), - )?; - Some(rope) - } - PositionalEmbedding::None | PositionalEmbedding::Sin => None, - }; - let mut layers = Vec::with_capacity(cfg.num_layers); - for layer_idx in 0..cfg.num_layers { - // Also send weights of first layer as only it contains the KQV proj weights - // for shared cross-attention layers - let shared_vb = if cfg.shared_cross_attn { Some(vb_l.pp(0)) } else { None }; - let layer = StreamingTransformerLayer::new(cfg, vb_l.pp(layer_idx), shared_vb)?; - layers.push(layer) - } - Ok(Self { - layers, - positional_embedding: cfg.positional_embedding, - max_period: cfg.max_period, - causal: cfg.causal, - num_heads: cfg.num_heads, - context: cfg.context, - last_reset_pos: vec![], - rope, - }) - } - - pub fn forward(&mut self, xs: &Tensor) -> Result { - self.forward_ca(xs, None) - } - - fn current_seq_len(&self) -> usize { - self.layers[0].self_attn.kv_cache.current_seq_len() - } - - pub fn forward_ca(&mut self, xs: &Tensor, ca_src: Option<&CaSrc>) -> Result { - let (b, t, c) = xs.dims3()?; - if !self.causal { - candle::bail!("only causal mode is supported") - } - if self.last_reset_pos.is_empty() { - self.last_reset_pos.resize(b, 0); - } - let current_seq_len = self.current_seq_len(); - // We will extract at most "context" from the kv_cache. - // Note that the mask still discards the values that are before context as this can happen - // when t > context. - let mask = { - // mask shape should be b, h, t, k - // self.layers[0].self_attn.kv_cache.attn_mask(t, xs.device())?; - // let mask = mask.broadcast_left((b, self.num_heads))?; - let ks = self.layers[0].self_attn.kv_cache.positions(t); - let min_ks = ks.iter().min().context("no positions, is t == 0?")?; - if t == 1 && self.last_reset_pos.iter().all(|v| v <= min_ks) { - // No need for a mask here. - None - } else { - let mut mask = Vec::with_capacity(b * self.num_heads * t * ks.len()); - for &last_reset_pos in self.last_reset_pos.iter() { - for t_pos in 0..t { - let t_pos = t_pos + current_seq_len; - for &k_pos in ks.iter() { - let m = if last_reset_pos <= k_pos - && k_pos <= t_pos - && t_pos <= k_pos + self.context - { - 0f32 - } else { - f32::NEG_INFINITY - }; - mask.push(m); - } - } - } - let mask = Tensor::from_vec(mask, (b, 1, t, ks.len()), xs.device())? - .to_dtype(xs.dtype())? - .expand((b, self.num_heads, t, ks.len()))?; - Some(mask) - } - }; - // pos is used for the rotary embeddings, as these are relative embeddings there is no need - // to adjust them for the actual position using last_reset_pos. - let pos = - Tensor::arange(current_seq_len as u32, (current_seq_len + t) as u32, xs.device())?; - let rope = match self.rope { - Some(ref rope) => Some(rope.rope(&pos)?), - None => None, - }; - let mut xs = match self.positional_embedding { - PositionalEmbedding::Rope | PositionalEmbedding::None => xs.clone(), - PositionalEmbedding::Sin => { - let dev = xs.device(); - let theta = self.max_period as f32; - let half_dim = c / 2; - let positions = pos.unsqueeze(1)?.to_dtype(DType::F32)?; - let inv_freq: Vec<_> = (0..half_dim) - .map(|i| 1f32 / theta.powf(i as f32 / (half_dim - 1) as f32)) - .collect(); - let inv_freq_len = inv_freq.len(); - let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; - let freqs = positions.broadcast_mul(&inv_freq)?; - let pos_emb = Tensor::cat(&[freqs.cos()?, freqs.sin()?], D::Minus1)?; - xs.broadcast_add(&pos_emb)? - } - }; - for layer in self.layers.iter_mut() { - xs = layer.forward(&xs, rope.as_ref(), ca_src, mask.as_ref())? - } - Ok(xs) - } - - pub fn maybe_precompute_ca_kv(&self, ca_src: Option) -> Result> { - let ca_src = match ca_src { - None => None, - Some(CaSrc::KeysValues(_)) => ca_src, - Some(tokens) => { - if self.layers.is_empty() { - Some(tokens) - } else { - match &self.layers[0].cross_attn { - None => Some(tokens), - Some((_, ca_module)) => { - let (k, v) = ca_module.compute_kv(&tokens)?; - Some(CaSrc::KeysValues((k, v))) - } - } - } - } - }; - Ok(ca_src) - } - - pub fn copy_state(&mut self, from: &Self) -> Result<()> { - if self.layers.len() != from.layers.len() { - candle::bail!("cannot copy kv-caches as the transformers have different depths") - } - self.last_reset_pos = from.last_reset_pos.clone(); - self.layers - .iter_mut() - .zip(from.layers.iter()) - .for_each(|(v, w)| v.set_kv_cache(w.self_attn.kv_cache.clone())); - Ok(()) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - if self.last_reset_pos.is_empty() { - self.last_reset_pos.resize(batch_size, 0); - } - if batch_idx >= self.last_reset_pos.len() { - candle::bail!("batch_idx {} is out of bounds for last_reset_pos", batch_idx) - } - self.last_reset_pos[batch_idx] = self.current_seq_len(); - Ok(()) - } -} - -impl StreamingModule for StreamingTransformer { - fn reset_state(&mut self) { - self.last_reset_pos.clear(); - self.layers.iter_mut().for_each(|v| v.reset_kv_cache()) - } - - fn step(&mut self, xs: &StreamTensor, _: &StreamMask) -> Result { - // TODO: Use the StreamMask - match xs.as_option() { - None => Ok(StreamTensor::empty()), - Some(xs) => Ok(StreamTensor::from_tensor(self.forward(xs)?)), - } - } -} - -#[derive(Debug, Clone)] -pub struct ProjectedTransformer { - // Projected transformer with unquantized projection - transformer: StreamingTransformer, - input_proj: Option, - output_projs: Vec>, - conv_layout: bool, - span: tracing::Span, -} - -impl ProjectedTransformer { - pub fn new( - input_dim: usize, - output_dims: &[usize], - cfg: &Config, - vb: MaybeQuantizedVarBuilder, - ) -> Result { - let transformer = StreamingTransformer::new(cfg, vb.pp("transformer"))?; - let input_proj = if input_dim == cfg.d_model { - None - } else { - let l = linear(input_dim, cfg.d_model, false, vb.pp("input_proj"))?; - Some(l) - }; - let mut output_projs = Vec::with_capacity(output_dims.len()); - let vb_o = vb.pp("output_projs"); - for (i, &output_dim) in output_dims.iter().enumerate() { - let output_proj = if output_dim == cfg.d_model { - None - } else { - let l = linear(cfg.d_model, output_dim, false, vb_o.pp(i))?; - Some(l) - }; - output_projs.push(output_proj) - } - Ok(Self { - transformer, - input_proj, - output_projs, - conv_layout: cfg.conv_layout, - span: tracing::span!(tracing::Level::TRACE, "proj-transformer"), - }) - } - - pub fn forward(&mut self, xs: &Tensor) -> Result> { - let _enter = self.span.enter(); - let xs = if self.conv_layout { xs.transpose(1, 2)? } else { xs.clone() }; - let xs = xs.apply(&self.input_proj.as_ref())?; - let xs = self.transformer.forward(&xs)?; - let mut ys = Vec::with_capacity(self.output_projs.len()); - for output_proj in self.output_projs.iter() { - let ys_ = xs.apply(&output_proj.as_ref())?; - let ys_ = if self.conv_layout { ys_.transpose(1, 2)? } else { ys_ }; - ys.push(ys_) - } - Ok(ys) - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - self.transformer.reset_batch_idx(batch_idx, batch_size) - } -} - -impl StreamingModule for ProjectedTransformer { - fn reset_state(&mut self) { - self.transformer.reset_state() - } - - fn step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - let xs = xs.apply(&|x: &Tensor| { - if self.conv_layout { - x.transpose(1, 2) - } else { - Ok(x.clone()) - } - })?; - let xs = xs.apply(&self.input_proj.as_ref())?; - let xs = self.transformer.step(&xs, m)?; - let ys = xs.apply(&self.output_projs[0].as_ref())?; - ys.apply(&|y: &Tensor| { - if self.conv_layout { - y.transpose(1, 2) - } else { - Ok(y.clone()) - } - }) - } -} - -#[derive(Debug, Clone)] -pub enum Transformer { - Standard(ProjectedTransformer), - Batched(crate::batched_transformer::ProjectedTransformer), -} - -impl StreamingModule for Transformer { - fn reset_state(&mut self) { - match self { - Transformer::Standard(t) => t.reset_state(), - Transformer::Batched(t) => t.reset_state(), - } - } - - fn step(&mut self, xs: &StreamTensor, m: &StreamMask) -> Result { - match self { - Transformer::Standard(t) => t.step(xs, m), - Transformer::Batched(t) => t.step(xs, m), - } - } -} - -impl Transformer { - pub fn new( - batch_size: Option, - dim: usize, - cfg: &Config, - vb: candle_nn::VarBuilder, - ) -> Result { - let transformer = match batch_size { - Some(batch_size) => { - let transformer = crate::batched_transformer::ProjectedTransformer::new( - dim, - &[dim], - batch_size, - cfg, - MaybeQuantizedVarBuilder::Real(vb), - )?; - Transformer::Batched(transformer) - } - None => { - let transformer = ProjectedTransformer::new( - dim, - &[dim], - cfg, - MaybeQuantizedVarBuilder::Real(vb), - )?; - Transformer::Standard(transformer) - } - }; - Ok(transformer) - } - - pub fn forward(&mut self, xs: &Tensor) -> Result> { - match self { - Transformer::Standard(t) => t.forward(xs), - Transformer::Batched(t) => t.forward(xs, &().into()), - } - } - - pub fn reset_batch_idx(&mut self, batch_idx: usize, batch_size: usize) -> Result<()> { - match self { - Transformer::Standard(t) => t.reset_batch_idx(batch_idx, batch_size), - Transformer::Batched(t) => t.reset_batch_idx(batch_idx), - } - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use candle::{IndexOp, Result, Tensor}; -use candle_transformers::generation::LogitsProcessor; - -use crate::transformer::CaSrc; - -pub const UNGENERATED: u32 = u32::MAX; - -#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)] -pub struct Config { - pub acoustic_delay: usize, - pub text_pad_token: u32, - pub text_bos_token: u32, - pub text_eos_token: u32, - pub text_eop_token: u32, - pub text_start_token: u32, - pub text_audio_delay_in_tokens: usize, - pub max_consecutive_pads: usize, - pub extra_steps: usize, - pub speaker_cond_duration_s: f64, - pub speaker_cond_dim: usize, - pub speaker_cond_n_speakers: usize, -} - -impl Config { - pub fn v202501() -> Self { - Self { - acoustic_delay: 2, - text_eop_token: 0, - text_bos_token: 1, - text_eos_token: 2, - text_pad_token: 3, - text_start_token: 8000, - text_audio_delay_in_tokens: 25, // aka interleaver_delay = 2s - max_consecutive_pads: 10, - extra_steps: 5, - speaker_cond_duration_s: 10., - speaker_cond_dim: 2048, - speaker_cond_n_speakers: 5, - } - } -} - -pub struct State { - model: crate::lm::LmModel, - ca_src: Option, - audio_tokens: Vec>, - text_tokens: Vec, - consecutive_pads: usize, - audio_lp: LogitsProcessor, - text_lp: LogitsProcessor, - step_idx: usize, - forced_audio_tokens: crate::lm::ForcedAudioTokens, - cfg_alpha: Option, - config: Config, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum AllowedTokens { - Text(u32), - Pad, - PadOrEpad, -} - -impl State { - pub fn new( - model: crate::lm::LmModel, - ca_src: Option, - max_step_idx: usize, - audio_lp: LogitsProcessor, - text_lp: LogitsProcessor, - cfg_alpha: Option, - config: Config, - ) -> Self { - let audio_tokens: Vec> = vec![ - vec![UNGENERATED; model.generated_audio_codebooks()]; - max_step_idx + config.acoustic_delay - ]; - let text_tokens = vec![UNGENERATED; max_step_idx + config.acoustic_delay]; - let forced_audio_tokens = crate::lm::ForcedAudioTokens::new( - config.acoustic_delay, - model.audio_pad_token(), - &[model.generated_audio_codebooks()], - ); - Self { - model, - ca_src, - audio_tokens, - text_tokens, - consecutive_pads: 0, - audio_lp, - text_lp, - step_idx: 0, - forced_audio_tokens, - cfg_alpha, - config, - } - } - - pub fn step_idx(&self) -> usize { - self.step_idx - } - - fn audio_pad_token(&self) -> u32 { - self.model.audio_pad_token() - } - - pub fn config(&self) -> &Config { - &self.config - } - - // The acoustic tokens are written with a delay, so this can create "gaps" of UNGENERATED - // tokens in the case where we call `step_audio_prompt` *after* `step`. - pub fn step( - &mut self, - prev_text_token: u32, - allowed_tokens: AllowedTokens, - conditions: Option<&crate::conditioner::Condition>, - ) -> Result { - let mut codes = Vec::with_capacity(self.model.generated_audio_codebooks()); - let dev = self.model.device(); - let batch_size = if self.cfg_alpha.is_some() { 2 } else { 1 }; - for codebook in 0..self.model.generated_audio_codebooks() { - let t = if codebook == 0 { - if self.step_idx == 0 { - Some(self.audio_pad_token()) - } else if self.step_idx <= self.config.text_audio_delay_in_tokens { - // The delayed pattern for TTS is a bit special, the audio-pad tokens are used - // in the same way as usual, i.e. for the first slice and until the acoustic - // delay for semantic tokens. - // However for the first couple seconds (set by `text_audio_delay_in_tokens`), - // the tokens that are *not* audio-pad are replaced by "literal zeros". - None - } else { - Some(self.audio_tokens[self.step_idx - 1][codebook]) - } - } else if self.step_idx <= self.config.acoustic_delay { - Some(self.audio_pad_token()) - } else if self.step_idx - <= self.config.text_audio_delay_in_tokens + self.config.acoustic_delay - { - // The same comment as above applies here. - None - } else { - Some(self.audio_tokens[self.step_idx - self.config.acoustic_delay - 1][codebook]) - }; - if t == Some(UNGENERATED) { - candle::bail!("internal error, ungenerated {}", self.step_idx) - } - let t = match t { - Some(t) => Some(Tensor::from_vec(vec![t; batch_size], (batch_size, 1), dev)?), - None => None, - }; - codes.push(t) - } - let prev_text_token = - Some(Tensor::from_vec(vec![prev_text_token; batch_size], (batch_size, 1), dev)?); - let (text_logits, ys) = match self.ca_src.as_ref() { - None => self.model.forward_cond(prev_text_token, codes, conditions, &().into())?, - Some(ca_src) => { - self.model.forward_ca(prev_text_token, codes, ca_src, conditions, &().into())? - } - }; - let text_logits = match self.cfg_alpha { - None => text_logits.i((0, 0))?, - Some(a) => match text_logits.dim(0)? { - 2 => ((text_logits.i((0, 0))? * a)? - (text_logits.i((1, 0))? * (a - 1.))?)?, - b_size => candle::bail!("unexpected batch size {b_size}"), - }, - }; - // When in tts mode, there are only two possible outcomes corresponding to tokens 0 and 3. - // 0 -> EOP or the next text token, this is ambiguous, a list of consecutive 0s correspond to - // word + EOP + word + EOP ... - // 3 -> pad. - // This will change when the simplerleaver lands. - let text_token = match allowed_tokens { - AllowedTokens::Text(v) => v, - AllowedTokens::Pad => self.config.text_pad_token, - AllowedTokens::PadOrEpad => { - if self.consecutive_pads > self.config.max_consecutive_pads { - self.config.text_eop_token - } else { - let text_token = self.text_lp.sample(&text_logits)?; - if text_token == self.config.text_pad_token { - self.config.text_pad_token - } else { - self.config.text_eop_token - } - } - } - }; - if text_token == self.config.text_pad_token { - self.consecutive_pads += 1 - } else { - self.consecutive_pads = 0 - } - self.text_tokens[self.step_idx] = text_token; - let last_audio_tokens = if self.step_idx < self.config.text_audio_delay_in_tokens { - None - } else { - match self.cfg_alpha { - None => self.model.depformer_sample( - &ys, - Some(text_token), - self.forced_audio_tokens.forced_tokens(self.step_idx), - &mut self.audio_lp, - )?, - Some(cfg_alpha) => self.model.depformer_sample_cfg( - &ys, - cfg_alpha, - Some(text_token), - self.forced_audio_tokens.forced_tokens(self.step_idx), - &mut self.audio_lp, - )?, - } - }; - let audio_pad_token = self.audio_pad_token(); - for c_idx in 0..self.model.generated_audio_codebooks() { - let delay = if c_idx == 0 { 0 } else { self.config.acoustic_delay }; - let pos = &mut self.audio_tokens[self.step_idx.saturating_sub(delay)][c_idx]; - match last_audio_tokens.as_ref() { - Some(lat) => { - if *pos == UNGENERATED { - *pos = lat[c_idx] - } - } - None => { - if *pos == UNGENERATED { - *pos = audio_pad_token - } - } - } - } - self.step_idx += 1; - if self.step_idx >= self.audio_tokens.len() { - candle::bail!("max step-idx reached") - } - Ok(text_token) - } - - pub fn overwrite_last_text_token(&mut self, text_token: u32) -> Result<()> { - if self.step_idx == 0 { - candle::bail!("cannot overwrite first token") - } - if text_token == UNGENERATED { - candle::bail!("cannot overwrite with UNGENERATED") - } - self.text_tokens[self.step_idx - 1] = text_token; - Ok(()) - } - - /// If include_all is set, all the time steps are returned. Otherwise only the timesteps that - /// have been generated are handled. - pub fn audio_tokens(&self, include_all: bool) -> &[Vec] { - if include_all { - &self.audio_tokens - } else { - let max_idx = usize::min(self.step_idx, self.audio_tokens.len()); - &self.audio_tokens[..max_idx] - } - } - - pub fn text_tokens(&self, include_all: bool) -> &[u32] { - if include_all { - &self.text_tokens - } else { - let max_idx = usize::min(self.step_idx, self.text_tokens.len()); - &self.text_tokens[..max_idx] - } - } - - pub fn last_audio_tokens(&self) -> Option> { - if self.step_idx <= self.config.acoustic_delay { - None - } else { - // step_idx is in advance by 1 + there is a 2 token delay on audio tokens. - let audio_tokens = &self.audio_tokens[self.step_idx - self.config.acoustic_delay - 1]; - if audio_tokens.iter().any(|v| *v >= self.audio_pad_token()) { - None - } else { - Some(audio_tokens.clone()) - } - } - } - - pub fn audio_codebooks(&self) -> usize { - self.model.generated_audio_codebooks() - } - - pub fn device(&self) -> &candle::Device { - self.model.device() - } - - pub fn dtype(&self) -> candle::DType { - self.model.dtype() - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum Speaker { - Main, - Other, -} - -pub fn tokenize_prompt( - text: &[String], - text_bos_token: u32, - text_eos_token: u32, - encode: impl Fn(&str) -> std::result::Result, E>, -) -> std::result::Result, Speaker)>, E> { - let mut prompt = vec![]; - for (turn_idx, turn) in text.iter().enumerate() { - let (speaker, turn_token) = if turn_idx % 2 == 0 { - (Speaker::Main, text_bos_token) - } else { - (Speaker::Other, text_eos_token) - }; - for (word_idx, word) in turn.split(' ').enumerate() { - let mut word = encode(word)?.into_iter().collect::>(); - if word_idx == 0 && speaker == Speaker::Main { - word.insert(0, turn_token) - } - if !word.is_empty() { - prompt.push((word, speaker)) - } - } - } - Ok(prompt) -} - -#[derive(Debug, Clone)] -pub struct SpeakerEncoder { - mimi: crate::mimi::Mimi, - learnt_padding: Tensor, - proj: candle_nn::Linear, - n_speakers: usize, - cond_dim: usize, - device: candle::Device, - dtype: candle::DType, -} - -impl SpeakerEncoder { - pub fn new( - mimi: crate::mimi::Mimi, - speaker_cond_dim: usize, - speaker_cond_n_speakers: usize, - dtype: candle::DType, - vb: candle_nn::VarBuilder, - ) -> Result { - let learnt_padding = vb.get( - (1, 1, speaker_cond_dim), - "condition_provider.conditioners.speaker_wavs.learnt_padding", - )?; - let mimi_dim = mimi.config().seanet.dimension; - let proj = candle_nn::linear_no_bias( - mimi_dim, - speaker_cond_dim, - vb.pp("condition_provider.conditioners.speaker_wavs.output_proj"), - )?; - Ok(Self { - mimi, - learnt_padding, - proj, - n_speakers: speaker_cond_n_speakers, - cond_dim: speaker_cond_dim, - device: vb.device().clone(), - dtype, - }) - } - - pub fn device(&self) -> &candle::Device { - &self.device - } - - pub fn sample_rate(&self) -> f64 { - self.mimi.config().sample_rate - } - - pub fn encode(&self, speakers: &[Tensor]) -> Result { - if speakers.is_empty() { - candle::bail!("empty speakers in encode") - } - let mut pcms = vec![]; - for pcm in speakers.iter().take(self.n_speakers) { - let stdev = pcm.broadcast_sub(&pcm.mean_all()?)?.sqr()?.mean_all()?.sqrt()?; - let pcm = (pcm * 0.08)?.broadcast_div(&stdev)?; - pcms.push(pcm) - } - let n_speakers = pcms.len(); - let pcm = Tensor::cat(&pcms, 0)?; - let mut mimi = self.mimi.clone(); - mimi.reset_state(); - let embeddings = mimi.encode_pre_quantize(&pcm)?.t()?.apply(&self.proj)?; - let embeddings = if n_speakers < self.n_speakers { - let lp = - embeddings.narrow(0, 0, 1)?.zeros_like()?.broadcast_add(&self.learnt_padding)?; - let mut embs = vec![embeddings]; - embs.resize(self.n_speakers - n_speakers + 1, lp); - Tensor::cat(&embs, 0)? - } else { - embeddings - }; - let embeddings = embeddings.flatten(0, 1)?.unsqueeze(0)?; - let embeddings = crate::tts::add_sin_embeddings(&embeddings)?; - embeddings.to_dtype(self.dtype) - } - - pub fn empty(&self) -> Result { - let embeddings = - self.learnt_padding.broadcast_as((1, self.n_speakers * 125, self.cond_dim))?; - let embeddings = crate::tts::add_sin_embeddings(&embeddings)?; - embeddings.to_dtype(self.dtype) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::transformer::CaSrc; -use candle::{Context, DType, Result, Tensor, D}; -use candle_nn::{linear_no_bias, Linear, VarBuilder}; -use candle_transformers::models::t5; - -pub struct Config { - pub t5: t5::Config, - pub lm: crate::lm::Config, - pub mimi: crate::mimi::Config, - pub max_duration_s: f64, - pub speaker_cond_duration_s: f64, - pub max_speakers: usize, -} - -impl Config { - pub fn v0_1(t5: t5::Config) -> Self { - let lm = crate::lm::Config::tts_v0_1(); - let mimi = crate::mimi::Config::v0_1(None); - Self { t5, lm, mimi, max_duration_s: 60., speaker_cond_duration_s: 4., max_speakers: 5 } - } - - pub fn v0_2(t5: t5::Config) -> Self { - let lm = crate::lm::Config::tts_v0_1(); - let mimi = crate::mimi::Config::v0_1(None); - Self { t5, lm, mimi, max_duration_s: 60., speaker_cond_duration_s: 10., max_speakers: 2 } - } -} - -#[derive(Clone)] -pub struct Model { - t5: t5::T5EncoderModel, - pub lm: crate::lm::LmModel, - speaker_cond: Option<(crate::mimi::Mimi, Linear)>, - t5_proj: Linear, - pub sample_rate: f64, - frame_rate: f64, - audio_vocab_size: u32, - audio_codebooks: usize, - pub max_duration_s: f64, - max_speakers: usize, - end_of_gen: Option, -} - -impl Model { - pub fn new( - cfg: &Config, - vb_t5: VarBuilder, - vb_lm: VarBuilder, - vb_speaker_cond: Option, - ) -> Result { - let t5 = t5::T5EncoderModel::load(vb_t5, &cfg.t5)?; - let speaker_cond = match vb_speaker_cond { - None => None, - Some(vb) => { - let mimi = crate::mimi::Mimi::new(cfg.mimi.clone(), vb)?; - let proj = linear_no_bias( - cfg.mimi.seanet.dimension, - cfg.lm.transformer.d_model, - vb_lm.pp("condition_provider.conditioners.speaker_wavs.output_proj"), - )?; - Some((mimi, proj)) - } - }; - let t5_proj = { - let name = if speaker_cond.is_some() { - "condition_provider.conditioners.diarized_transcript_in_segment.output_proj" - } else { - "condition_provider.conditioners.transcript_in_segment.output_proj" - }; - linear_no_bias(cfg.t5.d_model, cfg.lm.transformer.d_model, vb_lm.pp(name))? - }; - let lm = - crate::lm::LmModel::new(&cfg.lm, crate::nn::MaybeQuantizedVarBuilder::Real(vb_lm))?; - Ok(Self { - t5, - lm, - speaker_cond, - t5_proj, - sample_rate: cfg.mimi.sample_rate, - frame_rate: cfg.mimi.frame_rate, - audio_vocab_size: cfg.lm.audio_vocab_size as u32, - audio_codebooks: cfg.lm.audio_codebooks, - max_duration_s: cfg.max_duration_s, - max_speakers: cfg.max_speakers, - end_of_gen: None, - }) - } -} - -pub fn add_sin_embeddings(xs: &Tensor) -> Result { - let target_dtype = xs.dtype(); - let (_b_size, seq_len, dim) = xs.dims3()?; - let dev = xs.device(); - let half_dim = dim / 2; - let positions = - Tensor::arange(0u32, seq_len as u32, dev)?.unsqueeze(1)?.to_dtype(DType::F32)?; - let inv_freq: Vec<_> = - (0..half_dim).map(|i| 1f32 / 10000f32.powf(i as f32 / (half_dim - 1) as f32)).collect(); - let inv_freq_len = inv_freq.len(); - let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?; - let freqs = positions.broadcast_mul(&inv_freq)?; - let pos_emb = Tensor::cat(&[freqs.cos()?, freqs.sin()?], D::Minus1)?; - let xs = xs.to_dtype(DType::F32)?.broadcast_add(&pos_emb)?; - xs.to_dtype(target_dtype) -} - -impl Model { - pub fn conditions( - &mut self, - token_ids: &Tensor, - speaker_pcm: Option<&Tensor>, - ) -> Result { - let t5_condition = - self.t5.forward(token_ids)?.to_dtype(candle::DType::BF16)?.apply(&self.t5_proj)?; - let conditions = match speaker_pcm { - None => t5_condition, - Some(speaker_pcm) => { - let sc = match self.speaker_cond.as_mut() { - None => candle::bail!("speaker_pcm specified without a speaker-cond model"), - Some((mimi, proj)) => mimi - .encode_pre_quantize(speaker_pcm)? - .t()? - .to_dtype(candle::DType::BF16)? - .apply(proj)?, - }; - let z = sc.zeros_like()?; - let mut c1 = vec![&t5_condition, &sc]; - let mut c2 = vec![&t5_condition, &z]; - for _i in 0..self.max_speakers - 1 { - c1.push(&z); - c2.push(&z); - } - let c1 = Tensor::cat(&c1, 1)?; - let c2 = Tensor::cat(&c2, 1)?; - let xs = Tensor::cat(&[&c1, &c2], 0)?; - add_sin_embeddings(&xs)? - } - }; - Ok(conditions) - } - - pub fn sample(&mut self, conditions: &Tensor, cfg_alpha: f64) -> Result>> { - let lp = candle_transformers::generation::LogitsProcessor::from_sampling( - 299792458, - candle_transformers::generation::Sampling::TopK { k: 100, temperature: 0.8 }, - ); - self.sample_lp(conditions, cfg_alpha, lp) - } - - pub fn sample_lp( - &mut self, - conditions: &Tensor, - cfg_alpha: f64, - mut lp: candle_transformers::generation::LogitsProcessor, - ) -> Result>> { - let max_steps = (self.max_duration_s * self.frame_rate) as usize + 1; - let audio_codebooks = self.audio_codebooks; - let audio_vocab_size = self.audio_vocab_size; - let mut audio_tokens: Vec> = vec![vec![u32::MAX; audio_codebooks]; max_steps + 2]; - let forced_audio_tokens = crate::lm::ForcedAudioTokens::new( - /* acoustic_delay= */ 2, - self.lm.audio_pad_token(), - &[audio_codebooks], - ); - let quantizer_bins = audio_vocab_size - 2; // 2048 - for step_idx in 0..(max_steps + 2) { - let mut codes = Vec::with_capacity(audio_codebooks); - for codebook in 0..audio_codebooks { - let t = if codebook == 0 { - if step_idx == 0 { - audio_vocab_size - 1 - } else { - audio_tokens[step_idx - 1][0] - } - } else if step_idx <= 2 { - audio_vocab_size - 1 - } else { - audio_tokens[step_idx - 3][codebook] - }; - let t = Tensor::new(&[t], conditions.device())?.unsqueeze(0)?; - codes.push(Some(t)) - } - let (_text_logits, ys) = self.lm.forward_ca( - None, - codes, - &CaSrc::Tokens(conditions.clone()), - None, - &().into(), - )?; - let last_audio_tokens = if self.speaker_cond.is_some() { - self.lm.depformer_sample_cfg( - &ys, - cfg_alpha, - None, - forced_audio_tokens.forced_tokens(step_idx), - &mut lp, - )? - } else { - self.lm.depformer_sample( - &ys, - None, - forced_audio_tokens.forced_tokens(step_idx), - &mut lp, - )? - }; - let last_audio_tokens = last_audio_tokens.context("no depformer")?; - for (c_idx, token) in last_audio_tokens.into_iter().enumerate() { - if step_idx > 0 && token >= quantizer_bins && self.end_of_gen.is_none() { - // Continue generating for two steps to get the final acoustic tokens. - self.end_of_gen = Some(step_idx + 2) - } - let delay = if c_idx == 0 { 0 } else { 2 }; - audio_tokens[step_idx.saturating_sub(delay)][c_idx] = token - } - if Some(step_idx) == self.end_of_gen { - break; - } - } - Ok(audio_tokens) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use std::io::prelude::*; - -pub trait Sample { - fn to_i16(&self) -> i16; -} - -impl Sample for f32 { - fn to_i16(&self) -> i16 { - (self.clamp(-1.0, 1.0) * 32767.0) as i16 - } -} - -impl Sample for f64 { - fn to_i16(&self) -> i16 { - (self.clamp(-1.0, 1.0) * 32767.0) as i16 - } -} - -impl Sample for i16 { - fn to_i16(&self) -> i16 { - *self - } -} - -pub fn write_pcm_as_wav( - w: &mut W, - samples: &[S], - sample_rate: u32, -) -> std::io::Result<()> { - let len = 12u32; // header - let len = len + 24u32; // fmt - let len = len + samples.len() as u32 * 2 + 8; // data - let n_channels = 1u16; - let bytes_per_second = sample_rate * 2 * n_channels as u32; - w.write_all(b"RIFF")?; - w.write_all(&(len - 8).to_le_bytes())?; // total length minus 8 bytes - w.write_all(b"WAVE")?; - - // Format block - w.write_all(b"fmt ")?; - w.write_all(&16u32.to_le_bytes())?; // block len minus 8 bytes - w.write_all(&1u16.to_le_bytes())?; // PCM - w.write_all(&n_channels.to_le_bytes())?; // one channel - w.write_all(&sample_rate.to_le_bytes())?; - w.write_all(&bytes_per_second.to_le_bytes())?; - w.write_all(&2u16.to_le_bytes())?; // 2 bytes of data per sample - w.write_all(&16u16.to_le_bytes())?; // bits per sample - - // Data block - w.write_all(b"data")?; - w.write_all(&(samples.len() as u32 * 2).to_le_bytes())?; - for sample in samples.iter() { - w.write_all(&sample.to_i16().to_le_bytes())? - } - Ok(()) -} - - - -[package] -name = "moshi" -version.workspace = true -edition.workspace = true -description.workspace = true -repository.workspace = true -keywords.workspace = true -categories.workspace = true -license.workspace = true -readme = "../README.md" - -[dependencies] -candle = { workspace = true } -candle-nn = { workspace = true } -candle-transformers = { workspace = true } -candle-flash-attn = { workspace = true, optional = true } - -rayon = { workspace = true } -serde = { workspace = true } -tracing = { workspace = true } - -[features] -default = [] -cuda = ["candle/cuda", "candle-nn/cuda"] -metal = ["candle/metal", "candle-nn/metal"] -flash-attn = ["cuda", "dep:candle-flash-attn"] - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::AsrStreamingQuery as Query; -use anyhow::{Context, Result}; -use axum::extract::ws; -use candle::{DType, Device, Tensor}; -use candle_nn::VarBuilder; -use std::collections::VecDeque; -use tokio::task; -use tokio::time::{timeout, Duration}; - -const FRAME_SIZE: usize = 1920; - -#[derive(serde::Serialize, serde::Deserialize)] -#[serde(tag = "type")] -pub enum InMsg { - Init, - Marker { id: i64 }, - Audio { pcm: Vec }, - OggOpus { data: Vec }, -} - -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -#[serde(tag = "type")] -pub enum OutMsg { - Word { text: String, start_time: f64 }, - EndWord { stop_time: f64 }, - Marker { id: i64 }, - Step { step_idx: usize, prs: Vec, buffered_pcm: usize }, - Error { message: String }, - Ready, -} - -#[derive(Debug)] -pub struct Asr { - asr_delay_in_tokens: usize, - temperature: f64, - lm: moshi::lm::LmModel, - audio_tokenizer: moshi::mimi::Mimi, - text_tokenizer: std::sync::Arc, - instance_name: String, - log_dir: std::path::PathBuf, - conditions: Option, -} - -impl Asr { - pub fn new(asr: &crate::AsrConfig, config: &crate::Config, dev: &Device) -> Result { - let dtype = dev.bf16_default_to_f32(); - let vb_lm = - unsafe { VarBuilder::from_mmaped_safetensors(&[&asr.lm_model_file], dtype, dev)? }; - let lm = - moshi::lm::LmModel::new(&asr.model, moshi::nn::MaybeQuantizedVarBuilder::Real(vb_lm))?; - let conditions = match lm.condition_provider() { - None => None, - Some(cp) => { - let delay = - asr.conditioning_delay.context("missing conditioning_delay in config")?; - let conditions = cp.condition_cont("delay", -delay)?; - tracing::info!(?conditions, "generated conditions"); - Some(conditions) - } - }; - let audio_tokenizer = { - let vb = unsafe { - candle_nn::VarBuilder::from_mmaped_safetensors( - &[&asr.audio_tokenizer_file], - DType::F32, - dev, - )? - }; - let mut cfg = moshi::mimi::Config::v0_1(Some(asr.model.audio_codebooks)); - // The mimi transformer runs at 25Hz. - cfg.transformer.max_seq_len = asr.model.transformer.max_seq_len * 2; - moshi::mimi::Mimi::new(cfg, vb)? - }; - let text_tokenizer = sentencepiece::SentencePieceProcessor::open(&asr.text_tokenizer_file) - .with_context(|| asr.text_tokenizer_file.clone())?; - Ok(Self { - asr_delay_in_tokens: asr.asr_delay_in_tokens, - lm, - temperature: asr.temperature.unwrap_or(0.0), - audio_tokenizer, - text_tokenizer: text_tokenizer.into(), - log_dir: config.log_dir.clone().into(), - instance_name: config.instance_name.clone(), - conditions, - }) - } - - pub fn warmup(&self) -> Result<()> { - let lm = self.lm.clone(); - let audio_tokenizer = self.audio_tokenizer.clone(); - let mut state = moshi::asr::State::new( - 1, - self.asr_delay_in_tokens, - self.temperature, - audio_tokenizer, - lm, - )?; - let dev = state.device().clone(); - let pcm = vec![0f32; FRAME_SIZE * state.batch_size()]; - for _ in 0..2 { - let pcm = Tensor::new(pcm.as_slice(), &dev)?.reshape((state.batch_size(), 1, ()))?; - let _asr_msgs = - state.step_pcm(pcm, self.conditions.as_ref(), &().into(), |_, _, _| ())?; - } - Ok(()) - } - - pub async fn handle_socket(&self, socket: ws::WebSocket, query: Query) -> Result<()> { - use futures_util::{SinkExt, StreamExt}; - use serde::Serialize; - - let (mut sender, mut receiver) = socket.split(); - let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); - let (log_tx, log_rx) = std::sync::mpsc::channel(); - let lm = self.lm.clone(); - let audio_tokenizer = self.audio_tokenizer.clone(); - let mut state = moshi::asr::State::new( - 1, - self.asr_delay_in_tokens, - self.temperature, - audio_tokenizer, - lm, - )?; - let text_tokenizer = self.text_tokenizer.clone(); - - let asr_delay_in_tokens = self.asr_delay_in_tokens; - let conditions = self.conditions.clone(); - let mut ogg_opus_decoder = kaudio::ogg_opus::Decoder::new(24000, 1920)?; - let recv_loop = task::spawn(async move { - let dev = state.device().clone(); - // Store the markers in a double ended queue - let mut markers = VecDeque::new(); - while let Some(msg) = receiver.next().await { - let msg = match msg? { - ws::Message::Binary(x) => x, - // ping messages are automatically answered by tokio-tungstenite as long as - // the connection is read from. - ws::Message::Ping(_) | ws::Message::Pong(_) | ws::Message::Text(_) => continue, - ws::Message::Close(_) => break, - }; - let msg: InMsg = rmp_serde::from_slice(&msg)?; - let pcm = match msg { - // Init is only used in batched mode. - InMsg::Init => None, - InMsg::Marker { id } => { - tracing::info!("received marker {id}"); - let step_idx = state.model_step_idx(); - markers.push_back((step_idx, id)); - None - } - InMsg::OggOpus { data } => ogg_opus_decoder.decode(&data)?.map(|v| v.to_vec()), - InMsg::Audio { pcm } => Some(pcm), - }; - if let Some(pcm) = pcm { - tracing::info!("received audio {}", pcm.len()); - let pcm = Tensor::new(pcm.as_slice(), &dev)? - .reshape((1, 1, ()))? - .broadcast_as((state.batch_size(), 1, pcm.len()))?; - let asr_msgs = state.step_pcm( - pcm, - conditions.as_ref(), - &().into(), - |_, text_tokens, audio_tokens| { - let res = || { - let text_tokens = text_tokens.to_device(&Device::Cpu)?; - let audio_tokens: Vec = audio_tokens - .iter() - .map(|t| t.to_device(&Device::Cpu)) - .collect::>>()?; - let audio_tokens = Tensor::stack(&audio_tokens, 1)?; - log_tx.send((text_tokens, audio_tokens))?; - Ok::<_, anyhow::Error>(()) - }; - if let Err(err) = res() { - tracing::error!(?err, "failed to send log"); - } - }, - )?; - for asr_msg in asr_msgs.into_iter() { - let msg = match asr_msg { - moshi::asr::AsrMsg::Word { tokens, start_time, .. } => OutMsg::Word { - text: text_tokenizer.decode_piece_ids(&tokens)?, - start_time, - }, - moshi::asr::AsrMsg::Step { step_idx, prs } => { - let prs = prs.iter().map(|p| p[0]).collect::>(); - OutMsg::Step { step_idx, prs, buffered_pcm: 0 } - } - moshi::asr::AsrMsg::EndWord { stop_time, .. } => { - OutMsg::EndWord { stop_time } - } - }; - tx.send(msg)? - } - while let Some((step_idx, id)) = markers.front() { - if *step_idx + asr_delay_in_tokens <= state.model_step_idx() { - tx.send(OutMsg::Marker { id: *id })?; - markers.pop_front(); - } else { - break; - } - } - } - } - Ok::<(), anyhow::Error>(()) - }); - let send_loop = task::spawn(async move { - loop { - // The recv method is cancel-safe so can be wrapped in a timeout. - let msg = timeout(Duration::from_secs(10), rx.recv()).await; - let msg = match msg { - Ok(None) => break, - Err(_) => ws::Message::Ping(vec![].into()), - Ok(Some(msg)) => { - let mut buf = vec![]; - msg.serialize( - &mut rmp_serde::Serializer::new(&mut buf) - .with_human_readable() - .with_struct_map(), - )?; - ws::Message::Binary(buf.into()) - } - }; - sender.send(msg).await?; - } - tracing::info!("send loop exited"); - Ok::<(), anyhow::Error>(()) - }); - let sleep = tokio::time::sleep(std::time::Duration::from_secs(360)); - tokio::pin!(sleep); - // select should ensure that all the threads get aborted on timeout. - // TODO(laurent): this actually doesn't work as expected, and the background threads don't - // appear to be cancelled properly (at least the websocket connection remains open. - // laurent: Actually I guess this is because we wait for at least one of these to finish - // before exiting this task. - tokio::select! { - _ = &mut sleep => { - tracing::error!("reached timeout"); - } - res = recv_loop => { - match res { - Err(err) => tracing::error!(?err, "recv loop ended"), - Ok(Err(err)) => tracing::error!(?err, "recv loop err"), - Ok(Ok(())) => tracing::info!("recv loop ended"), - } - } - res = send_loop => { - match res { - Err(err) => tracing::error!(?err, "send loop ended"), - Ok(Err(err)) => tracing::error!(?err, "send loop err"), - Ok(Ok(())) => tracing::info!("send loop ended"), - } - } - } - let (text_tokens, audio_tokens): (Vec<_>, Vec<_>) = log_rx.try_iter().unzip(); - let text_tokens = Tensor::cat(&text_tokens, candle::D::Minus1)?; - let audio_tokens = Tensor::cat(&audio_tokens, candle::D::Minus1)?; - self.save_logs(&query, audio_tokens, text_tokens)?; - tracing::info!("exiting handle-socket"); - Ok(()) - } - - fn save_logs(&self, query: &Query, audio_tokens: Tensor, text_tokens: Tensor) -> Result<()> { - let since_epoch = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?; - let (secs, us) = (since_epoch.as_secs(), since_epoch.subsec_micros()); - let base_path = self.log_dir.join(format!("{}-asr-{secs}-{us}", self.instance_name)); - let json_filename = base_path.with_extension("json"); - let json_content = serde_json::to_string_pretty(query)?; - std::fs::write(json_filename, json_content)?; - let st_filename = base_path.with_extension("safetensors"); - let audio_tokens = audio_tokens.to_device(&Device::Cpu)?.to_dtype(DType::I64)?; - let st_content = - std::collections::HashMap::from([("text", text_tokens), ("audio", audio_tokens)]); - candle::safetensors::save(&st_content, st_filename)?; - Ok(()) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::asr::{InMsg, OutMsg}; -use crate::metrics::asr as metrics; -use crate::AsrStreamingQuery as Query; -use anyhow::{Context, Result}; -use axum::extract::ws; -use candle::{DType, Device, Tensor}; -use candle_nn::VarBuilder; -use std::collections::{BinaryHeap, VecDeque}; -use std::sync::{Arc, Mutex}; -use tokio::task; -use tokio::time::{timeout, Duration}; - -const FRAME_SIZE: usize = 1920; -const SEND_PING_EVERY: Duration = Duration::from_secs(10); -const POST_RETRY_DELAY: Duration = Duration::from_millis(100); -const POST_MAX_RETRIES: usize = 1000; - -#[derive(Debug, PartialEq, Eq, Clone)] -struct Marker { - channel_id: ChannelId, - batch_idx: usize, - step_idx: usize, - marker_id: i64, -} - -impl std::cmp::PartialOrd for Marker { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl std::cmp::Ord for Marker { - // We use reverse ordering as this will be embedded in a max heap. - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.step_idx.cmp(&other.step_idx).reverse() - } -} - -type InSend = std::sync::mpsc::Sender; -type InRecv = std::sync::mpsc::Receiver; -type OutSend = tokio::sync::mpsc::UnboundedSender; -type OutRecv = tokio::sync::mpsc::UnboundedReceiver; - -/// Unique identifier. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -pub struct ChannelId(usize); - -impl ChannelId { - fn new() -> Self { - // https://users.rust-lang.org/t/idiomatic-rust-way-to-generate-unique-id/33805 - use std::sync::atomic; - static COUNTER: atomic::AtomicUsize = atomic::AtomicUsize::new(1); - Self(COUNTER.fetch_add(1, atomic::Ordering::Relaxed)) - } -} - -struct Channel { - id: ChannelId, - in_rx: InRecv, - out_tx: OutSend, - data: VecDeque, - decoder: kaudio::ogg_opus::Decoder, - steps: usize, -} - -impl Channel { - fn new(in_rx: InRecv, out_tx: OutSend) -> Result { - metrics::OPEN_CHANNELS.inc(); - let decoder = kaudio::ogg_opus::Decoder::new(24000, FRAME_SIZE)?; - Ok(Self { id: ChannelId::new(), in_rx, out_tx, data: VecDeque::new(), decoder, steps: 0 }) - } - - fn extend_data(&mut self, mut pcm: Vec) -> Option> { - if self.data.is_empty() && pcm.len() >= FRAME_SIZE { - self.data.extend(&pcm[FRAME_SIZE..]); - pcm.truncate(FRAME_SIZE); - Some(pcm) - } else { - self.data.extend(&pcm); - if self.data.len() >= FRAME_SIZE { - Some(self.data.drain(..FRAME_SIZE).collect()) - } else { - None - } - } - } - - fn send(&self, msg: OutMsg, ref_channel_id: Option) -> Result<()> { - // If the channel id has changed compared to the reference. Return Ok(()) - // so as not to disconnect the new user. - if Some(self.id) != ref_channel_id { - return Ok(()); - } - self.out_tx.send(msg)?; - Ok(()) - } -} - -impl Drop for Channel { - fn drop(&mut self) { - metrics::OPEN_CHANNELS.dec(); - metrics::CONNECTION_NUM_STEPS.observe(self.steps as f64); - } -} - -struct Logger { - base_path: std::path::PathBuf, - log_tx: std::sync::mpsc::Sender<(Tensor, Tensor)>, - log_rx: std::sync::mpsc::Receiver<(Tensor, Tensor)>, - log_frequency_s: f64, -} - -impl Logger { - fn new>( - instance_name: &str, - log_dir: P, - log_frequency_s: f64, - ) -> Result { - let since_epoch = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?; - let (secs, us) = (since_epoch.as_secs(), since_epoch.subsec_micros()); - let base_path = log_dir.as_ref().join(format!("{instance_name}-asr-{secs}-{us}")); - let (log_tx, log_rx) = std::sync::mpsc::channel::<(Tensor, Tensor)>(); - Ok(Self { base_path, log_tx, log_rx, log_frequency_s }) - } - - fn log_loop(self) { - tracing::info!(?self.base_path, "starting log loop"); - task::spawn_blocking(move || { - let mut cnt = 0usize; - loop { - std::thread::sleep(std::time::Duration::from_secs_f64(self.log_frequency_s)); - let tokens: Vec<_> = self.log_rx.try_iter().collect(); - if tokens.is_empty() { - tracing::info!("no tokens to log"); - continue; - } - let st_filename = self.base_path.with_extension(format!("{cnt}.safetensors")); - tracing::info!(?st_filename, "writing logs"); - let (text_tokens, audio_tokens): (Vec<_>, Vec<_>) = tokens.into_iter().unzip(); - let write = || { - let text_tokens = Tensor::cat(&text_tokens, candle::D::Minus1)?; - let audio_tokens = Tensor::cat(&audio_tokens, candle::D::Minus1)?; - let st_content = std::collections::HashMap::from([ - ("text", text_tokens), - ("audio", audio_tokens), - ]); - candle::safetensors::save(&st_content, st_filename)?; - Ok::<_, anyhow::Error>(()) - }; - if let Err(err) = write() { - tracing::error!(?err, "failed to write safetensors"); - } - cnt += 1; - } - }); - } -} - -struct BatchedAsrInner { - channels: Channels, - asr_delay_in_tokens: usize, - temperature: f64, - lm: moshi::lm::LmModel, - audio_tokenizer: moshi::mimi::Mimi, - text_tokenizer: std::sync::Arc, -} - -fn warmup( - state: &mut moshi::asr::State, - conditions: Option<&moshi::conditioner::Condition>, -) -> Result<()> { - let dev = state.device().clone(); - let pcm = vec![0f32; FRAME_SIZE * state.batch_size()]; - let pcm = Tensor::from_vec(pcm, (state.batch_size(), 1, FRAME_SIZE), &dev)?; - let mask = moshi::StreamMask::new(vec![true; state.batch_size()], &dev)?; - for _ in 0..2 { - let _asr_msgs = state.step_pcm(pcm.clone(), conditions, &mask, |_, _, _| ())?; - } - dev.synchronize()?; - Ok(()) -} - -impl BatchedAsrInner { - fn start_model_loop( - self, - conditioning_delay: Option, - conditioning_learnt_padding: bool, - batch_size: usize, - logger: Option<&Logger>, - ) -> Result<()> { - let conditions = match self.lm.condition_provider() { - None => None, - Some(cp) => match (conditioning_delay, conditioning_learnt_padding) { - (Some(delay), false) => { - let conditions = cp.condition_cont("delay", -delay)?; - tracing::info!(?conditions, "generated conditions"); - Some(conditions) - } - (None, true) => { - let conditions = cp.learnt_padding("delay")?; - tracing::info!(?conditions, "generated conditions"); - Some(conditions) - } - (Some(_), true) => anyhow::bail!( - "conditioning_delay/conditioning_learnt_padding cannot be both set" - ), - (None, false) => { - anyhow::bail!("conditioning_delay/conditioning_learnt_padding is required") - } - }, - }; - let mut state = moshi::asr::State::new( - batch_size, - self.asr_delay_in_tokens, - self.temperature, - self.audio_tokenizer.clone(), - self.lm.clone(), - )?; - let log_tx = logger.map(|v| v.log_tx.clone()); - let dev = state.device().clone(); - let model_loop: task::JoinHandle> = task::spawn_blocking(move || { - tracing::info!("warming-up the asr"); - warmup(&mut state, conditions.as_ref())?; - tracing::info!("starting asr loop {batch_size}"); - // Store the markers in a double ended queue - let mut markers = BinaryHeap::new(); - // This loop runs in real-time. - let mut step_idx = 0; - loop { - let (batch_pcm, mask, ref_channel_ids) = - self.pre_process(&mut state, step_idx, &mut markers); - let with_data = mask.iter().filter(|v| **v).count(); - if with_data > 0 { - let mask = moshi::StreamMask::new(mask, &dev)?; - let pcm = - Tensor::new(batch_pcm.as_slice(), &dev)?.reshape((batch_size, 1, ()))?; - let start_time = std::time::Instant::now(); - let asr_msgs = state.step_pcm( - pcm, - conditions.as_ref(), - &mask, - |_, text_tokens, audio_tokens| { - let res = || { - if let Some(log_tx) = log_tx.as_ref() { - let text_tokens = text_tokens.to_device(&Device::Cpu)?; - let audio_tokens: Vec = audio_tokens - .iter() - .map(|t| t.to_device(&Device::Cpu)) - .collect::>>()?; - let audio_tokens = Tensor::stack(&audio_tokens, 1)?; - if let Err(err) = log_tx.send((text_tokens, audio_tokens)) { - tracing::error!(?err, "failed to send log"); - }; - } - Ok::<_, anyhow::Error>(()) - }; - if let Err(err) = res() { - tracing::error!(?err, "failed to send log"); - } - }, - )?; - let elapsed = start_time.elapsed().as_secs_f64(); - metrics::MODEL_STEP_DURATION.observe(elapsed); - tracing::info!(step_idx, with_data, "{:.2}ms", elapsed * 1000.); - step_idx += 1; - self.post_process(asr_msgs, step_idx, &mut markers, &mask, &ref_channel_ids)?; - } else { - std::thread::sleep(std::time::Duration::from_millis(2)); - } - } - }); - task::spawn(async { - match model_loop.await { - Err(err) => tracing::error!(?err, "model loop join err"), - Ok(Err(err)) => tracing::error!(?err, "model loop err"), - Ok(Ok(())) => tracing::info!("model loop exited"), - } - }); - Ok(()) - } - - fn pre_process( - &self, - state: &mut moshi::asr::State, - step_idx: usize, - markers: &mut BinaryHeap, - ) -> (Vec, Vec, Vec>) { - use rayon::prelude::*; - enum Todo { - Reset(usize), - Marker(Marker), - } - - let mut mask = vec![false; state.batch_size()]; - let mut channels = self.channels.lock().unwrap(); - let mut batch_pcm = vec![0f32; FRAME_SIZE * channels.len()]; - let channel_ids = channels.iter().map(|c| c.as_ref().map(|c| c.id)).collect::>(); - let todo = batch_pcm - .par_chunks_mut(FRAME_SIZE) - .zip(channels.par_iter_mut()) - .zip(mask.par_iter_mut()) - .enumerate() - .flat_map(|(bid, ((out_pcm, channel), mask))| -> Option { - let c = channel.as_mut()?; - if c.out_tx.is_closed() { - *channel = None; - None - } else { - use std::sync::mpsc::TryRecvError; - match c.in_rx.try_recv() { - Ok(InMsg::Init) => { - if c.out_tx.send(OutMsg::Ready).is_err() { - *channel = None; - } - Some(Todo::Reset(bid)) - } - Ok(InMsg::Marker { id }) => { - tracing::info!(bid, id, "received marker"); - // The marker only gets sent back once all the current data has been - // processed and the asr delay has passed. - let current_data = c.data.len() / FRAME_SIZE; - let step_idx = step_idx + state.asr_delay_in_tokens() + current_data; - let marker = Marker { - channel_id: c.id, - batch_idx: bid, - step_idx, - marker_id: id, - }; - Some(Todo::Marker(marker)) - } - Ok(InMsg::OggOpus { data }) => { - match c.decoder.decode(&data) { - Err(err) => tracing::error!(?err, "oggopus not supported"), - Ok(None) => {} - Ok(Some(pcm)) => { - out_pcm.copy_from_slice(pcm); - c.steps += 1; - *mask = true; - } - } - None - } - Ok(InMsg::Audio { pcm }) => { - if let Some(bpcm) = c.extend_data(pcm) { - out_pcm.copy_from_slice(&bpcm); - c.steps += 1; - *mask = true; - } - None - } - Err(TryRecvError::Empty) => { - // Even if we haven't received new data, we process the existing one. - if let Some(bpcm) = c.extend_data(vec![]) { - out_pcm.copy_from_slice(&bpcm); - c.steps += 1; - *mask = true; - } - None - } - Err(TryRecvError::Disconnected) => { - *channel = None; - None - } - } - } - }) - .collect::>(); - todo.into_iter().for_each(|t| match t { - Todo::Reset(bid) => { - if let Err(err) = state.reset_batch_idx(bid) { - tracing::error!(?err, bid, "failed to reset batch"); - } - } - Todo::Marker(m) => markers.push(m), - }); - (batch_pcm, mask, channel_ids) - } - - fn post_process( - &self, - asr_msgs: Vec, - step_idx: usize, - markers: &mut BinaryHeap, - mask: &moshi::StreamMask, - ref_channel_ids: &[Option], - ) -> Result<()> { - let mut channels = self.channels.lock().unwrap(); - for asr_msg in asr_msgs.into_iter() { - match asr_msg { - moshi::asr::AsrMsg::Word { tokens, start_time, batch_idx } => { - let msg = OutMsg::Word { - text: self.text_tokenizer.decode_piece_ids(&tokens)?, - start_time, - }; - if let Some(c) = channels[batch_idx].as_ref() { - if c.send(msg, ref_channel_ids[batch_idx]).is_err() { - channels[batch_idx] = None; - } - } - } - moshi::asr::AsrMsg::EndWord { stop_time, batch_idx } => { - let msg = OutMsg::EndWord { stop_time }; - if let Some(c) = channels[batch_idx].as_ref() { - if c.send(msg, ref_channel_ids[batch_idx]).is_err() { - channels[batch_idx] = None; - } - } - } - moshi::asr::AsrMsg::Step { step_idx, prs } => { - for (batch_idx, c) in channels.iter_mut().enumerate() { - if !mask.is_active(batch_idx) { - continue; - } - if let Some(ch) = c.as_mut() { - let prs = prs.iter().map(|p| p[batch_idx]).collect(); - let msg = OutMsg::Step { step_idx, prs, buffered_pcm: ch.data.len() }; - if ch.send(msg, ref_channel_ids[batch_idx]).is_err() { - *c = None; - } - } - } - } - } - } - while let Some(m) = markers.peek() { - if m.step_idx <= step_idx { - if let Some(c) = channels[m.batch_idx].as_ref() { - if c.send(OutMsg::Marker { id: m.marker_id }, Some(m.channel_id)).is_err() { - channels[m.batch_idx] = None; - } - } - markers.pop(); - } else { - break; - } - } - Ok(()) - } -} - -type Channels = Arc>>>; - -pub struct BatchedAsr { - channels: Channels, - config: crate::AsrConfig, - batch_size: usize, -} - -impl BatchedAsr { - pub fn new( - batch_size: usize, - asr: &crate::AsrConfig, - config: &crate::Config, - dev: &Device, - ) -> Result { - let dtype = dev.bf16_default_to_f32(); - let vb_lm = - unsafe { VarBuilder::from_mmaped_safetensors(&[&asr.lm_model_file], dtype, dev)? }; - let lm = moshi::lm::LmModel::batched( - batch_size, - &asr.model, - moshi::nn::MaybeQuantizedVarBuilder::Real(vb_lm), - )?; - let audio_tokenizer = { - let vb = unsafe { - candle_nn::VarBuilder::from_mmaped_safetensors( - &[&asr.audio_tokenizer_file], - DType::F32, - dev, - )? - }; - let mut cfg = moshi::mimi::Config::v0_1(Some(asr.model.audio_codebooks)); - // The mimi transformer runs at 25Hz. - cfg.transformer.max_seq_len = asr.model.transformer.max_seq_len * 2; - moshi::mimi::Mimi::batched(batch_size, cfg, vb)? - }; - let text_tokenizer = sentencepiece::SentencePieceProcessor::open(&asr.text_tokenizer_file) - .with_context(|| asr.text_tokenizer_file.clone())?; - let channels = (0..batch_size).map(|_| None).collect::>(); - let channels = Arc::new(Mutex::new(channels)); - let asr_delay_in_tokens = - asr.conditioning_delay.map_or(asr.asr_delay_in_tokens, |v| (v * 12.5) as usize + 1); - let batched_asr = BatchedAsrInner { - asr_delay_in_tokens, - temperature: asr.temperature.unwrap_or(0.0), - lm, - audio_tokenizer, - text_tokenizer: text_tokenizer.into(), - channels: channels.clone(), - }; - let logger = match asr.log_frequency_s { - Some(s) => Some(Logger::new(&config.instance_name, &config.log_dir, s)?), - None => None, - }; - batched_asr.start_model_loop( - asr.conditioning_delay, - asr.conditioning_learnt_padding, - batch_size, - logger.as_ref(), - )?; - if let Some(logger) = logger { - logger.log_loop() - } - Ok(Self { channels, config: asr.clone(), batch_size }) - } - - fn channels(&self) -> Result> { - let mut channels = self.channels.lock().unwrap(); - // Linear scan to find an available channel. This is fairly inefficient, instead we should - // probably have a queue of available slots. - for (batch_idx, channel) in channels.iter_mut().enumerate() { - if channel.is_none() { - let (in_tx, in_rx) = std::sync::mpsc::channel::(); - let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::(); - let c = Channel::new(in_rx, out_tx)?; - *channel = Some(c); - return Ok(Some((batch_idx, in_tx, out_rx))); - } - } - Ok(None) - } - - pub async fn handle_query(&self, query: axum::body::Bytes) -> Result> { - tracing::info!("batched-asr post query"); - let (batch_idx, in_tx, mut out_rx) = { - let mut num_tries = 0; - loop { - match self.channels() { - Ok(Some(x)) => break x, - Ok(None) => { - num_tries += 1; - if num_tries > POST_MAX_RETRIES { - tracing::error!("no free channels after 1000 tries"); - anyhow::bail!("no free channels"); - } - tokio::time::sleep(POST_RETRY_DELAY).await; - } - Err(err) => { - tracing::error!(?err, "no free channels"); - Err(err)? - } - } - } - }; - tracing::info!(batch_idx, "batched-asr channel"); - in_tx.send(InMsg::Init)?; - let (pcm, sample_rate) = crate::utils::pcm_decode(query)?; - let pcm = if sample_rate == 24000 { - pcm - } else { - kaudio::resample(&pcm, sample_rate as usize, 24000)? - }; - in_tx.send(InMsg::Audio { pcm })?; - in_tx.send(InMsg::Marker { id: 0 })?; - in_tx.send(InMsg::Audio { pcm: vec![0f32; 240000] })?; - let mut msgs = vec![]; - while let Some(msg) = out_rx.recv().await { - match msg { - OutMsg::Marker { .. } => break, - OutMsg::Error { .. } | OutMsg::Word { .. } | OutMsg::EndWord { .. } => { - msgs.push(msg) - } - OutMsg::Ready | OutMsg::Step { .. } => {} - } - } - Ok(msgs) - } - - pub async fn handle_socket(&self, socket: ws::WebSocket, query: Query) -> Result<()> { - use futures_util::{SinkExt, StreamExt}; - use serde::Serialize; - - tracing::info!(?query, "batched-asr ws query"); - metrics::CONNECT.inc(); - - let (mut sender, receiver) = socket.split(); - let (batch_idx, in_tx, mut out_rx) = match self.channels()? { - Some(v) => v, - None => { - tracing::error!("no free channels"); - let mut msg = vec![]; - OutMsg::Error { message: "no free channels".into() }.serialize( - &mut rmp_serde::Serializer::new(&mut msg) - .with_human_readable() - .with_struct_map(), - )?; - sender.send(ws::Message::binary(msg)).await?; - sender.close().await?; - anyhow::bail!("no free channels") - } - }; - tracing::info!(batch_idx, "batched-asr channel"); - in_tx.send(InMsg::Init)?; - - let recv_loop = task::spawn(async move { - let mut receiver = receiver; - // There are two timeouts here: - // - The short timeout handles the case where the client does not answer the regular pings. - // - The long timeout handles the case where the client does not send valid data for a - // long time. - let mut last_message_received = std::time::Instant::now(); - let short_timeout_duration = SEND_PING_EVERY * 2; - let long_timeout_duration = std::time::Duration::from_secs(120); - loop { - use ws::Message; - let msg = match timeout(short_timeout_duration, receiver.next()).await { - Ok(Some(msg)) => msg, - Ok(None) => break, - Err(_) => { - tracing::info!(?batch_idx, "recv loop short timeout"); - break; - } - }; - if last_message_received.elapsed() > long_timeout_duration { - tracing::info!(?batch_idx, "recv loop long timeout"); - break; - } - let msg = match msg? { - Message::Binary(x) => x, - // ping messages are automatically answered by tokio-tungstenite as long as - // the connection is read from. - Message::Ping(_) | Message::Pong(_) | Message::Text(_) => continue, - Message::Close(_) => break, - }; - last_message_received = std::time::Instant::now(); - let msg: InMsg = rmp_serde::from_slice(&msg)?; - in_tx.send(msg)?; - } - Ok::<_, anyhow::Error>(()) - }); - let send_loop = task::spawn(async move { - let mut sender = sender; - loop { - // The recv method is cancel-safe so can be wrapped in a timeout. - let msg = timeout(SEND_PING_EVERY, out_rx.recv()).await; - let msg = match msg { - Ok(None) => break, - Err(_) => ws::Message::Ping(vec![].into()), - Ok(Some(msg)) => { - let mut buf = vec![]; - msg.serialize( - &mut rmp_serde::Serializer::new(&mut buf) - .with_human_readable() - .with_struct_map(), - )?; - ws::Message::Binary(buf.into()) - } - }; - sender.send(msg).await?; - } - Ok::<(), anyhow::Error>(()) - }); - - // Keep track of the outputs of the different threads. - task::spawn(async { - match send_loop.await { - Err(err) => tracing::error!(?err, "send loop join err"), - Ok(Err(err)) => tracing::error!(?err, "send loop err"), - Ok(Ok(())) => tracing::info!("send loop exited"), - } - }); - task::spawn(async { - match recv_loop.await { - Err(err) => tracing::error!(?err, "recv loop join err"), - Ok(Err(err)) => tracing::error!(?err, "recv loop err"), - Ok(Ok(())) => tracing::info!("recv loop exited"), - } - }); - - Ok(()) - } - - pub fn config(&self) -> &crate::AsrConfig { - &self.config - } - - pub fn total_slots(&self) -> usize { - self.batch_size - } - - pub fn used_slots(&self) -> usize { - self.channels.lock().unwrap().iter().filter(|v| v.is_some()).count() - } -} - - - -pub mod protocol; - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. -use crate::protocol::MsgType; -use anyhow::{Context, Result}; -use axum::extract::ws; -use candle::{Device, IndexOp, Tensor}; -use candle_nn::VarBuilder; -use candle_transformers::generation::LogitsProcessor; - -use kaudio::ogg_opus; - -struct TextDecoder { - gen_config: moshi::lm_generate_multistream::Config, - text_tokenizer: std::sync::Arc, -} - -impl TextDecoder { - fn text(&self, prev_text_token: u32, text_token: u32) -> Option { - let config = &self.gen_config; - if text_token != config.text_start_token - && text_token != config.text_pad_token - && text_token != config.text_eop_token - { - if prev_text_token == config.text_start_token { - self.text_tokenizer.decode_piece_ids(&[text_token]).ok() - } else { - let prev_ids = self.text_tokenizer.decode_piece_ids(&[prev_text_token]).ok(); - let ids = self.text_tokenizer.decode_piece_ids(&[prev_text_token, text_token]).ok(); - prev_ids.and_then(|prev_ids| { - ids.map(|ids| { - if ids.len() > prev_ids.len() { - ids[prev_ids.len()..].to_string() - } else { - String::new() - } - }) - }) - } - } else { - None - } - } -} - -pub struct Lm { - dev: Device, - gen_config: moshi::lm_generate_multistream::Config, - lm: moshi::lm::LmModel, - audio_tokenizer: moshi::mimi::Mimi, - text_tokenizer: std::sync::Arc, - instance_name: String, - log_dir: std::path::PathBuf, -} - -enum WsEvent { - Text(String), - Pcm(Vec), -} - -enum LogEvent { - TextToken(u32), - AudioTokens(Vec), -} - -impl Lm { - pub fn new(lm: &crate::LmConfig, config: &crate::Config, dev: &Device) -> Result { - let dtype = dev.bf16_default_to_f32(); - let model_config = &lm.model; - let gen_config = lm.gen.clone(); - let audio_tokenizer = moshi::mimi::load(&lm.audio_tokenizer_file, Some(8), dev)?; - let text_tokenizer = sentencepiece::SentencePieceProcessor::open(&lm.text_tokenizer_file) - .with_context(|| lm.text_tokenizer_file.clone())?; - let vb_lm = - unsafe { VarBuilder::from_mmaped_safetensors(&[&lm.lm_model_file], dtype, dev)? }; - let lm = moshi::lm::LmModel::new( - model_config, - moshi::nn::MaybeQuantizedVarBuilder::Real(vb_lm), - )?; - Ok(Self { - audio_tokenizer, - lm, - gen_config, - dev: dev.clone(), - log_dir: config.log_dir.clone().into(), - instance_name: config.instance_name.clone(), - text_tokenizer: text_tokenizer.into(), - }) - } - - pub async fn handle_socket(&self, socket: ws::WebSocket) -> Result<()> { - use futures_util::StreamExt; - - tracing::info!("connected"); - let (opus_in_tx, mut opus_in_rx) = tokio::sync::mpsc::unbounded_channel(); - let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel(); - let (event_tx, event_rx) = std::sync::mpsc::channel(); - let (mut ws_sender, mut ws_receiver) = socket.split(); - let ws_recv_handle = tokio::spawn(async move { - while let Some(msg) = ws_receiver.next().await { - let (msg_type, payload) = match msg? { - ws::Message::Binary(b) => { - if b.is_empty() { - continue; - } - let msg_type = MsgType::from_u8(b[0])?; - let payload = b[1..].to_vec(); - (msg_type, payload) - } - _ => continue, - }; - match msg_type { - MsgType::Audio => { - opus_in_tx.send(payload)?; - } - t => { - tracing::warn!("unexpected msg type {t:?}"); - continue; - } - } - } - Ok::<_, anyhow::Error>(()) - }); - let dev = self.dev.clone(); - let mut audio_tokenizer = self.audio_tokenizer.clone(); - audio_tokenizer.reset_state(); - let text_lp = LogitsProcessor::from_sampling( - 299792458, - candle_transformers::generation::Sampling::TopK { k: 25, temperature: 0.8 }, - ); - let audio_lp = LogitsProcessor::from_sampling( - 299792458, - candle_transformers::generation::Sampling::TopK { k: 250, temperature: 0.8 }, - ); - let conditions = match self.lm.condition_provider() { - None => None, - Some(cp) => { - let conditions = cp.condition_lut("description", "very_good")?; - tracing::info!(?conditions, "generated conditions"); - Some(conditions) - } - }; - - let mut state = moshi::lm_generate_multistream::State::new( - self.lm.clone(), - /* max_steps = */ 4096, - audio_lp, - text_lp, - None, - None, - None, - self.gen_config.clone(), - ); - let text_decoder = TextDecoder { - gen_config: self.gen_config.clone(), - text_tokenizer: self.text_tokenizer.clone(), - }; - let mut decoder = ogg_opus::Decoder::new(24000, 1920)?; - let pcm_recv_handle = tokio::spawn(async move { - let mut prev_text_token = state.config().text_start_token; - tracing::info!("starting pcm recv loop"); - while let Some(opus) = opus_in_rx.recv().await { - if let Some(pcm) = decoder.decode(&opus)? { - let pcm = Tensor::new(pcm, &dev)?.reshape((1, 1, ()))?; - let audio_tokens = audio_tokenizer.encode_step(&pcm.into(), &().into())?; - let audio_tokens = match audio_tokens.as_option() { - None => continue, - Some(audio_tokens) => audio_tokens, - }; - let (_one, _codebooks, steps) = audio_tokens.dims3()?; - - for step in 0..steps { - let codes = audio_tokens.i((0, .., step))?.to_vec1::()?; - let text_token = state.step_( - Some(prev_text_token), - &codes, - None, - None, - conditions.as_ref(), - )?; - - if let Some(text) = text_decoder.text(prev_text_token, text_token) { - out_tx.send(WsEvent::Text(text))? - } - event_tx.send(LogEvent::TextToken(text_token))?; - tracing::info!(text_token, "sampled text token"); - if let Some(audio_tokens) = state.last_audio_tokens() { - let audio_tokens_t = { - let cb = state.config().generated_audio_codebooks; - Tensor::from_slice(&audio_tokens[..cb], (1, cb, 1), &dev)? - }; - event_tx.send(LogEvent::AudioTokens(audio_tokens))?; - let pcm = - audio_tokenizer.decode_step(&audio_tokens_t.into(), &().into())?; - if let Some(pcm) = pcm.as_option() { - let pcm = pcm.i((0, 0))?.to_vec1::()?; - out_tx.send(WsEvent::Pcm(pcm))?; - } - } - prev_text_token = text_token - } - } - } - Ok::<_, anyhow::Error>(()) - }); - let send_handle = tokio::spawn(async move { - use futures_util::SinkExt; - - let mut encoder = ogg_opus::Encoder::new(24000)?; - let mut handshake = vec![MsgType::Handshake.to_u8()]; - handshake.resize(9, 0u8); - if let Err(err) = ws_sender.send(ws::Message::binary(handshake)).await { - tracing::error!("error sending header {err:?}"); - return Ok(()); - } - { - let msg: Vec = [&[MsgType::Audio.to_u8()], encoder.header_data()].concat(); - let msg = ws::Message::Binary(msg.into()); - ws_sender.send(msg).await?; - } - while let Some(evt) = out_rx.recv().await { - let msg: Vec = match evt { - WsEvent::Pcm(pcm) => { - let ogg = encoder.encode_page(&pcm)?; - [&[MsgType::Audio.to_u8()], ogg.as_slice()].concat() - } - WsEvent::Text(text) => [&[MsgType::Text.to_u8()], text.as_bytes()].concat(), - }; - let msg = ws::Message::Binary(msg.into()); - ws_sender.send(msg).await? - } - Ok::<_, anyhow::Error>(()) - }); - let sleep = tokio::time::sleep(std::time::Duration::from_secs(360)); - tokio::pin!(sleep); - // select should ensure that all the threads get aborted on timeout. - tokio::select! { - _ = &mut sleep => { - tracing::error!("reached timeout"); - } - r = pcm_recv_handle => { - tracing::error!(?r, "pcm recv loop ended") - } - r = ws_recv_handle => { - tracing::error!(?r, "ws recv loop ended") - } - r = send_handle => { - tracing::error!(?r, "ws send loop ended") - } - }; - let events: Vec<_> = event_rx.try_iter().collect(); - self.save_logs((), events)?; - Ok(()) - } - - fn save_logs(&self, query: (), events: Vec) -> Result<()> { - let cpu = &Device::Cpu; - let since_epoch = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?; - let (secs, us) = (since_epoch.as_secs(), since_epoch.subsec_micros()); - let base_path = self.log_dir.join(format!("{}-lm-{secs}-{us}", self.instance_name)); - let json_filename = base_path.with_extension("json"); - let json_content = serde_json::to_string_pretty(&query)?; - std::fs::write(json_filename, json_content)?; - let st_filename = base_path.with_extension("safetensors"); - let text_tokens: Vec = events - .iter() - .filter_map(|v| match v { - LogEvent::TextToken(v) => Some(*v as i64), - LogEvent::AudioTokens(_) => None, - }) - .collect(); - let text_len = text_tokens.len(); - let text_tokens = - Tensor::from_vec(text_tokens, text_len, cpu)?.to_dtype(candle::DType::I64)?; - let audio_tokens: Vec<_> = events - .iter() - .filter_map(|v| match v { - LogEvent::TextToken(_) => None, - LogEvent::AudioTokens(a) => { - let a = a.iter().map(|v| *v as i64).collect::>(); - Some(Tensor::from_slice(&a, (1, a.len(), 1), cpu)) - } - }) - .collect::>>()?; - let audio_tokens = Tensor::cat(&audio_tokens, 2)?; - let st_content = - std::collections::HashMap::from([("text", text_tokens), ("audio", audio_tokens)]); - candle::safetensors::save(&st_content, st_filename)?; - Ok(()) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::Result; -use axum::{http::StatusCode, response::IntoResponse, response::Response}; -use candle::Device; -use std::str::FromStr; -use std::sync::Arc; - -mod asr; -mod batched_asr; -mod lm; -mod metrics; -mod mimi; -mod protocol; -mod py_module; -mod py_module_post; -mod tts; -mod utils; - -const ID_HEADER: &str = "kyutai-api-key"; -const ROOM_ID_HEADER: &str = "room_id"; - -pub const TTS_PY: &[u8] = include_bytes!("../tts.py"); -pub const VOICE_PY: &[u8] = include_bytes!("../voice.py"); -pub const UV_LOCK: &[u8] = include_bytes!("../uv.lock"); - -#[derive(clap::Parser, Debug)] -struct WorkerArgs { - #[clap(short = 'l', long = "log", default_value = "info")] - log_level: String, - - #[clap(short = 'a', long = "addr", default_value = "0.0.0.0")] - addr: String, - - #[clap(short = 'p', long = "port", default_value = "8080")] - port: u16, - - #[clap(long)] - cpu: bool, - - #[clap(long)] - config: String, - - #[clap(long)] - silent: bool, -} - -#[derive(Debug, clap::Subcommand)] -enum Command { - Validate { configs: Vec }, - Configs { which: String }, - Worker(WorkerArgs), -} - -#[derive(clap::Parser, Debug)] -#[clap(name = "server", about = "Kyutai moshi server")] -struct Args { - #[command(subcommand)] - command: Command, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct TtsConfig { - pub lm_model_file: String, - pub text_tokenizer_file: String, - pub speaker_tokenizer_file: String, - pub audio_tokenizer_file: String, - pub voices: std::collections::HashMap, - pub voice_dir: String, - pub model: moshi::lm::Config, - pub generation: moshi::tts_streaming::Config, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct AsrConfig { - pub lm_model_file: String, - pub text_tokenizer_file: String, - pub audio_tokenizer_file: String, - pub model: moshi::lm::Config, - pub asr_delay_in_tokens: usize, - #[serde(default)] - pub log_frequency_s: Option, - #[serde(default)] - pub conditioning_delay: Option, - // The default for bools in rust is false. - #[serde(default)] - pub conditioning_learnt_padding: bool, - #[serde(default)] - pub temperature: Option, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct MimiConfig { - pub audio_tokenizer_file: String, - pub auth_recv: bool, - pub rooms: Vec, - pub default_room: Option, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct LmConfig { - pub lm_model_file: String, - pub text_tokenizer_file: String, - pub audio_tokenizer_file: String, - pub model: moshi::lm::Config, - pub gen: moshi::lm_generate_multistream::Config, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct PyConfig { - #[serde(default)] - pub script: Option, - pub batch_size: usize, - pub text_tokenizer_file: String, - pub text_bos_token: u32, - #[serde(default)] - pub py: Option, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct PyPostConfig { - #[serde(default)] - pub script: Option, - #[serde(default)] - pub py: Option, -} - -#[derive(Debug, Clone, serde::Deserialize)] -#[serde(tag = "type")] -pub enum ModuleConfig { - Tts { - path: String, - #[serde(flatten)] - config: TtsConfig, - }, - Asr { - path: String, - #[serde(flatten)] - config: AsrConfig, - }, - BatchedAsr { - path: String, - #[serde(flatten)] - config: AsrConfig, - batch_size: usize, - }, - Mimi { - send_path: String, - recv_path: String, - #[serde(flatten)] - config: MimiConfig, - }, - Lm { - path: String, - #[serde(flatten)] - config: LmConfig, - }, - Py { - path: String, - #[serde(flatten)] - config: PyConfig, - }, - PyPost { - path: String, - #[serde(flatten)] - config: PyPostConfig, - }, -} - -#[derive(Debug, Clone, serde::Deserialize)] -pub struct Config { - pub static_dir: String, - pub log_dir: String, - pub instance_name: String, - #[serde(default)] - pub modules: std::collections::HashMap, - pub authorized_ids: std::collections::HashSet, -} - -impl Config { - pub fn load>(p: P) -> Result { - use utils::resolve_or_download as rod; - let config = std::fs::read_to_string(p)?; - let mut config: Self = toml::from_str(&config)?; - for (_, c) in config.modules.iter_mut() { - match c { - ModuleConfig::Mimi { send_path: _, recv_path: _, config: c } => { - c.audio_tokenizer_file = rod(&c.audio_tokenizer_file)?; - } - ModuleConfig::Tts { path: _, config: c } => { - c.lm_model_file = rod(&c.lm_model_file)?; - c.text_tokenizer_file = rod(&c.text_tokenizer_file)?; - c.speaker_tokenizer_file = rod(&c.speaker_tokenizer_file)?; - c.audio_tokenizer_file = rod(&c.audio_tokenizer_file)?; - for (_, v) in c.voices.iter_mut() { - *v = rod(v)? - } - c.voice_dir = rod(&c.voice_dir)?; - } - ModuleConfig::BatchedAsr { path: _, config: c, batch_size: _ } => { - c.lm_model_file = rod(&c.lm_model_file)?; - c.text_tokenizer_file = rod(&c.text_tokenizer_file)?; - c.audio_tokenizer_file = rod(&c.audio_tokenizer_file)?; - } - ModuleConfig::Asr { path: _, config: c } => { - c.lm_model_file = rod(&c.lm_model_file)?; - c.text_tokenizer_file = rod(&c.text_tokenizer_file)?; - c.audio_tokenizer_file = rod(&c.audio_tokenizer_file)?; - } - ModuleConfig::Lm { path: _, config: c } => { - c.audio_tokenizer_file = rod(&c.audio_tokenizer_file)?; - c.text_tokenizer_file = rod(&c.text_tokenizer_file)?; - c.lm_model_file = rod(&c.lm_model_file)?; - } - ModuleConfig::Py { path: _, config: c } => { - if let Some(script) = &mut c.script { - *script = rod(script)?; - } - c.text_tokenizer_file = rod(&c.text_tokenizer_file)?; - if let Some(t) = c.py.as_mut() { - crate::utils::resolve_or_download_toml(t)?; - } - } - ModuleConfig::PyPost { path: _, config: c } => { - if let Some(script) = &mut c.script { - *script = rod(script)?; - } - if let Some(t) = c.py.as_mut() { - crate::utils::resolve_or_download_toml(t)?; - } - } - } - } - config.static_dir = rod(&config.static_dir)?; - config.log_dir = rod(&config.log_dir)?; - config.instance_name = rod(&config.instance_name)?; - Ok(config) - } -} - -fn device(cpu: bool) -> Result { - if cpu { - Ok(Device::Cpu) - } else if candle::utils::cuda_is_available() { - Ok(Device::new_cuda(0)?) - } else if candle::utils::metal_is_available() { - Ok(Device::new_metal(0)?) - } else { - Ok(Device::Cpu) - } -} - -#[allow(unused)] -enum Module { - Tts { path: String, m: Arc }, - Asr { path: String, m: Arc }, - BatchedAsr { path: String, m: Arc }, - Mimi { send_path: String, recv_path: String, m: Arc }, - Lm { path: String, m: Arc }, - Py { path: String, m: Arc }, - PyPost { path: String, m: Arc }, -} - -struct SharedStateInner { - config: Config, -} - -type SharedState = Arc; - -fn lm_router(s: Arc, path: &str) -> axum::Router<()> { - async fn lm_websocket( - socket: axum::extract::ws::WebSocket, - state: Arc, - _addr: Option, - ) { - if let Err(err) = state.handle_socket(socket).await { - tracing::error!(?err, "lm") - } - } - - async fn lm_streaming( - ws: axum::extract::ws::WebSocketUpgrade, - headers: axum::http::HeaderMap, - state: axum::extract::State>, - ) -> utils::AxumResult { - let addr = headers.get("X-Real-IP").and_then(|v| v.to_str().ok().map(|v| v.to_string())); - tracing::info!(addr, "handling lm-streaming query"); - let state = state.0.clone(); - let upg = ws.write_buffer_size(0).on_upgrade(move |v| lm_websocket(v, state, addr)); - Ok(upg) - } - - axum::Router::new().route(path, axum::routing::get(lm_streaming)).with_state(s) -} - -impl Module { - fn new(module_cfg: &ModuleConfig, full_cfg: &Config, dev: &Device) -> Result { - let m = match module_cfg { - ModuleConfig::Lm { path, config } => { - let m = lm::Lm::new(config, full_cfg, dev)?; - let m = Arc::new(m); - Self::Lm { m, path: path.to_string() } - } - ModuleConfig::Asr { path, config } => { - let m = asr::Asr::new(config, full_cfg, dev)?; - let m = Arc::new(m); - tracing::info!("warming up the asr"); - m.warmup()?; - tracing::info!("done warming up the asr, ready to roll!"); - Self::Asr { m, path: path.to_string() } - } - ModuleConfig::BatchedAsr { path, config, batch_size } => { - let m = batched_asr::BatchedAsr::new(*batch_size, config, full_cfg, dev)?; - let m = Arc::new(m); - Self::BatchedAsr { m, path: path.to_string() } - } - ModuleConfig::Tts { path, config } => { - let voice = config.voices.keys().next(); - let m = tts::Model::new(config, full_cfg, dev)?; - let m = Arc::new(m); - if let Some(voice) = voice { - tracing::info!(voice, "warming up the tts"); - m.run(&TtsQuery { - text: vec!["hello".to_string()], - seed: 42, - temperature: 0.8, - top_k: 250, - voice: Some(voice.clone()), - voices: None, - max_seq_len: None, - return_timestamps: None, - cfg_alpha: None, - })?; - tracing::info!("done warming up the tts, ready to roll!"); - } - Self::Tts { m, path: path.to_string() } - } - ModuleConfig::Mimi { send_path, recv_path, config } => { - let m = mimi::Mimi::new(config, full_cfg, dev)?; - let m = Arc::new(m); - Self::Mimi { m, send_path: send_path.to_string(), recv_path: recv_path.to_string() } - } - ModuleConfig::Py { path, config } => { - let m = py_module::M::new(config.clone())?; - let m = Arc::new(m); - Self::Py { m, path: path.to_string() } - } - ModuleConfig::PyPost { path, config } => { - let m = py_module_post::M::new(config.clone())?; - let m = Arc::new(m); - Self::PyPost { m, path: path.to_string() } - } - }; - Ok(m) - } - - fn router(&self, shared_state: &SharedState) -> Result> { - let router = match self { - Self::Lm { path, m } => lm_router(m.clone(), path), - Self::Asr { path, m } => asr_router(m.clone(), path, shared_state), - Self::BatchedAsr { path, m } => batched_asr_router(m.clone(), path, shared_state), - Self::Tts { path, m } => tts_router(m.clone(), path, shared_state), - Self::Mimi { send_path, recv_path, m } => { - mimi_router(m.clone(), send_path, recv_path, shared_state) - } - Self::Py { path, m } => py_router(m.clone(), path, shared_state), - Self::PyPost { path, m } => py_router_post(m.clone(), path, shared_state), - }; - Ok(router) - } -} - -struct AppStateInner { - modules: Vec, -} - -type AppState = Arc; - -impl AppStateInner { - fn new(args: &WorkerArgs, config: Config) -> Result { - let device = device(args.cpu)?; - - // The following does not have a significant impact as soon as batch sizes are - // large enough so we don't activate it for now. - // #[cfg(feature = "cuda")] - // if let candle::Device::Cuda(d) = &device { - // unsafe { - // d.disable_event_tracking(); - // } - // }; - - let mut modules = Vec::with_capacity(config.modules.len()); - for (_, module_cfg) in config.modules.iter() { - let m = Module::new(module_cfg, &config, &device)?; - modules.push(m) - } - Ok(Self { modules }) - } -} - -fn tracing_init( - log_dir: &str, - instance_name: &str, - log_level: &str, - silent: bool, -) -> Result { - use tracing_subscriber::prelude::*; - - let build_info = utils::BuildInfo::new(); - let file_appender = tracing_appender::rolling::daily(log_dir, format!("log.{instance_name}")); - let (non_blocking, guard) = tracing_appender::non_blocking(file_appender); - let filter = tracing_subscriber::filter::LevelFilter::from_str(log_level)?; - let mut layers = vec![tracing_subscriber::fmt::layer() - .event_format(tracing_subscriber::fmt::format().with_file(true).with_line_number(true)) - .with_writer(non_blocking) - .with_filter(filter) - .boxed()]; - if !silent { - layers.push(Box::new( - tracing_subscriber::fmt::layer() - .event_format( - tracing_subscriber::fmt::format().with_file(true).with_line_number(true), - ) - .with_writer(std::io::stdout) - .with_filter(filter), - )) - }; - tracing_subscriber::registry().with(layers).init(); - tracing::info!(?build_info); - Ok(guard) -} - -async fn metrics( - axum::extract::ConnectInfo(_addr): axum::extract::ConnectInfo, - _state: axum::extract::State, - _req: axum::extract::Query<()>, -) -> impl IntoResponse { - use prometheus::Encoder; - - let encoder = prometheus::TextEncoder::new(); - let metric_families = prometheus::gather(); - let mut buffer = vec![]; - if let Err(err) = encoder.encode(&metric_families, &mut buffer) { - return (axum::http::StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response(); - }; - axum::response::Response::builder() - .status(200) - .header(axum::http::header::CONTENT_TYPE, encoder.format_type()) - .body(axum::body::Body::from(buffer)) - .unwrap() -} - -#[tokio::main(flavor = "multi_thread")] -async fn main() { - // When an error bubbles up in the tokio main function, the whole program does not - // seem to crash if some background tasks are still running. - // This can lead to errors such as "port already in use" not being reported so we - // exit the process explicitely here. - if let Err(err) = main_().await { - eprintln!("Error: {err}"); - std::process::exit(1); - } -} - -async fn main_() -> Result<()> { - let args = ::parse(); - match args.command { - Command::Configs { which } => match which.as_str() { - "tts.py" => { - println!("{}", String::from_utf8_lossy(TTS_PY)) - } - "voice.py" => { - println!("{}", String::from_utf8_lossy(VOICE_PY)) - } - "uv.lock" => { - println!("{}", String::from_utf8_lossy(UV_LOCK)) - } - _ => { - eprintln!("Unknown config: {which}"); - std::process::exit(1); - } - }, - Command::Validate { configs } => { - tracing_subscriber::fmt().init(); - for config in configs.iter() { - let _ = Config::load(config)?; - tracing::info!(?config, "loaded succesfully") - } - } - Command::Worker(args) => { - use axum::routing::get; - - let config = Config::load(&args.config)?; - if std::env::var("RUST_LOG").is_err() { - std::env::set_var("RUST_LOG", format!("{},hyper=info,mio=info", args.log_level)) - } - let _guard = - tracing_init(&config.log_dir, &config.instance_name, &args.log_level, args.silent)?; - let num_workers = tokio::runtime::Handle::current().metrics().num_workers(); - tracing::info!(num_workers, "starting worker"); - - let static_dir = utils::resolve_or_download(&config.static_dir)?; - let shared_state = Arc::new(SharedStateInner { config: config.clone() }); - let state = Arc::new(AppStateInner::new(&args, config)?); - let mut app = axum::Router::new() - .route("/api/build_info", get(build_info)) - .route("/api/modules_info", get(modules_info)) - .route("/metrics", axum::routing::get(metrics)) - .fallback_service( - tower_http::services::ServeDir::new(&static_dir) - .append_index_html_on_directories(true), - ) - .layer( - tower::ServiceBuilder::new() - .layer(tower_http::trace::TraceLayer::new_for_http()), - ) - .with_state(state.clone()); - for module in state.modules.iter() { - app = app.merge(module.router(&shared_state)?) - } - - let sock_addr = std::net::SocketAddr::from(( - std::net::IpAddr::from_str(args.addr.as_str()) - .unwrap_or(std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)), - args.port, - )); - tracing::info!("listening on http://{}", sock_addr); - let listener = tokio::net::TcpListener::bind(sock_addr).await?; - axum::serve( - listener, - app.into_make_service_with_connect_info::(), - ) - .await?; - } - } - Ok(()) -} - -#[derive(serde::Deserialize, serde::Serialize, Debug, Clone, Copy, PartialEq, Eq)] -enum StreamingOutput { - Pcm, - PcmMessagePack, - OggOpus, - OggOpusMessagePack, -} -fn default_seed() -> u64 { - 42 -} -fn default_temperature() -> f64 { - 0.8 -} -fn default_top_k() -> usize { - 250 -} -fn default_format() -> StreamingOutput { - StreamingOutput::OggOpus -} - -#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)] -struct TtsStreamingQuery { - #[serde(default = "default_seed")] - seed: u64, - #[serde(default = "default_temperature")] - temperature: f64, - #[serde(default = "default_top_k")] - top_k: usize, - #[serde(default = "default_format")] - format: StreamingOutput, - voice: Option, - voices: Option>, - max_seq_len: Option, - cfg_alpha: Option, - auth_id: Option, -} - -#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)] -struct TtsQuery { - text: Vec, - seed: u64, - temperature: f64, - top_k: usize, - voice: Option, - voices: Option>, - max_seq_len: Option, - return_timestamps: Option, - cfg_alpha: Option, -} - -#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)] -struct TtsResponse { - wav: String, - transcript: Vec, -} - -fn tts_router(s: Arc, path: &str, ss: &SharedState) -> axum::Router<()> { - use base64::Engine; - - async fn tts_websocket( - socket: axum::extract::ws::WebSocket, - state: Arc, - query: TtsStreamingQuery, - _addr: Option, - ) { - if let Err(err) = state.handle_socket(socket, query).await { - tracing::error!(?err, "tts") - } - } - - async fn t( - state: axum::extract::State<(Arc, SharedState)>, - headers: axum::http::HeaderMap, - req: axum::Json, - ) -> utils::AxumResult { - tracing::info!("handling tts query {req:?}"); - let valid_id = headers - .get(ID_HEADER) - .and_then(|v| v.to_str().ok()) - .is_some_and(|id| state.0 .1.config.authorized_ids.contains(id)); - if !valid_id { - return Ok(StatusCode::UNAUTHORIZED.into_response()); - } - let (wav, transcript) = { - let _guard = state.0 .0.mutex.lock().await; - state.0 .0.run(&req)? - }; - tracing::info!("ok {}", wav.len()); - if req.return_timestamps.unwrap_or(false) { - let data = - TtsResponse { wav: base64::prelude::BASE64_STANDARD.encode(wav), transcript }; - Ok(( - StatusCode::OK, - [(axum::http::header::CONTENT_TYPE, "application/json")], - axum::Json(data), - ) - .into_response()) - } else { - Ok((StatusCode::OK, [(axum::http::header::CONTENT_TYPE, "audio/wav")], wav) - .into_response()) - } - } - - async fn streaming_t( - ws: axum::extract::ws::WebSocketUpgrade, - headers: axum::http::HeaderMap, - state: axum::extract::State<(Arc, SharedState)>, - req: axum::extract::Query, - ) -> utils::AxumResult { - tracing::info!("handling tts streaming query {req:?}"); - let addr = headers.get("X-Real-IP").and_then(|v| v.to_str().ok().map(|v| v.to_string())); - // It's tricky to set the headers of a websocket in javascript so we pass the token via the - // query too. - let auth_id = match headers.get(ID_HEADER) { - Some(v) => v.to_str().ok(), - None => req.auth_id.as_deref(), - }; - let valid_id = auth_id.is_some_and(|id| state.1.config.authorized_ids.contains(id)); - if !valid_id { - return Ok(StatusCode::UNAUTHORIZED.into_response()); - } - let tts_query = req.0.clone(); - let tts = state.0 .0.clone(); - let upg = - ws.write_buffer_size(0).on_upgrade(move |v| tts_websocket(v, tts, tts_query, addr)); - Ok(upg) - } - - axum::Router::new() - .route(path, axum::routing::post(t)) - .route(&format!("{path}_streaming"), axum::routing::get(streaming_t)) - .with_state((s, ss.clone())) -} - -async fn build_info( - axum::extract::ConnectInfo(_addr): axum::extract::ConnectInfo, - _state: axum::extract::State, - _req: axum::extract::Query<()>, -) -> impl IntoResponse { - let build_info = utils::BuildInfo::new(); - utils::WrapJson(Ok(build_info)).into_response() -} - -async fn modules_info( - axum::extract::ConnectInfo(_addr): axum::extract::ConnectInfo, - state: axum::extract::State, - _req: axum::extract::Query<()>, -) -> impl IntoResponse { - let modules: Vec<_> = state - .modules - .iter() - .filter_map(|m| match m { - Module::BatchedAsr { path, m } => { - let config = m.config(); - let mut info = std::collections::HashMap::new(); - info.insert("type", "batched_asr".to_string()); - info.insert("path", path.to_string()); - info.insert("lm", config.lm_model_file.clone()); - info.insert("audio_tokenizer", config.audio_tokenizer_file.clone()); - info.insert("used_slots", m.used_slots().to_string()); - info.insert("total_slots", m.total_slots().to_string()); - Some(info) - } - Module::Py { path, m } => { - let config = m.config(); - let mut info = std::collections::HashMap::new(); - info.insert("type", "py".to_string()); - info.insert("path", path.to_string()); - info.insert("script", config.script.as_ref().map_or("tts.py", |v| v).to_string()); - info.insert("used_slots", m.used_slots().to_string()); - info.insert("total_slots", m.total_slots().to_string()); - Some(info) - } - _ => None, - }) - .collect(); - utils::WrapJson(Ok(modules)).into_response() -} - -#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)] -struct AsrStreamingQuery { - auth_id: Option, -} - -#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)] -struct PyStreamingQuery { - auth_id: Option, - #[serde(default = "default_format")] - format: StreamingOutput, - #[serde(default)] - voice: Option, -} - -fn asr_router(s: Arc, path: &str, ss: &SharedState) -> axum::Router<()> { - async fn asr_websocket( - socket: axum::extract::ws::WebSocket, - state: Arc, - query: AsrStreamingQuery, - _addr: Option, - ) { - if let Err(err) = state.handle_socket(socket, query).await { - tracing::error!(?err, "asr") - } - } - - async fn t( - ws: axum::extract::ws::WebSocketUpgrade, - headers: axum::http::HeaderMap, - state: axum::extract::State<(Arc, SharedState)>, - req: axum::extract::Query, - ) -> utils::AxumResult { - let addr = headers.get("X-Real-IP").and_then(|v| v.to_str().ok().map(|v| v.to_string())); - tracing::info!(addr, "handling asr-streaming query"); - // It's tricky to set the headers of a websocket in javascript so we pass the token via the - // query too. - let auth_id = match headers.get(ID_HEADER) { - Some(v) => v.to_str().ok(), - None => req.auth_id.as_deref(), - }; - let valid_id = auth_id.is_some_and(|id| state.1.config.authorized_ids.contains(id)); - if !valid_id { - return Ok(StatusCode::UNAUTHORIZED.into_response()); - } - let asr_query = req.0.clone(); - let asr = state.0 .0.clone(); - let upg = - ws.write_buffer_size(0).on_upgrade(move |v| asr_websocket(v, asr, asr_query, addr)); - Ok(upg) - } - axum::Router::new().route(path, axum::routing::get(t)).with_state((s, ss.clone())) -} - -fn batched_asr_router( - s: Arc, - path: &str, - ss: &SharedState, -) -> axum::Router<()> { - async fn asr_websocket( - socket: axum::extract::ws::WebSocket, - state: Arc, - query: AsrStreamingQuery, - _addr: Option, - ) { - if let Err(err) = state.handle_socket(socket, query).await { - tracing::error!(?err, "asr") - } - } - - // TODO: add a batch mode. - async fn t( - state: axum::extract::State<(Arc, SharedState)>, - headers: axum::http::HeaderMap, - req: axum::body::Bytes, - ) -> utils::AxumResult { - tracing::info!(len = req.len(), "handling asr post query"); - let valid_id = headers - .get(ID_HEADER) - .and_then(|v| v.to_str().ok()) - .is_some_and(|id| state.0 .1.config.authorized_ids.contains(id)); - if !valid_id { - return Ok(StatusCode::UNAUTHORIZED.into_response()); - } - let transcript = state.0 .0.handle_query(req).await?; - Ok(( - StatusCode::OK, - [(axum::http::header::CONTENT_TYPE, "application/json")], - axum::Json(transcript), - ) - .into_response()) - } - - async fn streaming_t( - ws: axum::extract::ws::WebSocketUpgrade, - headers: axum::http::HeaderMap, - state: axum::extract::State<(Arc, SharedState)>, - req: axum::extract::Query, - ) -> utils::AxumResult { - let addr = headers.get("X-Real-IP").and_then(|v| v.to_str().ok().map(|v| v.to_string())); - tracing::info!(addr, "handling batched asr-streaming query"); - // It's tricky to set the headers of a websocket in javascript so we pass the token via the - // query too. - let auth_id = match headers.get(ID_HEADER) { - Some(v) => v.to_str().ok(), - None => req.auth_id.as_deref(), - }; - let valid_id = auth_id.is_some_and(|id| state.1.config.authorized_ids.contains(id)); - if !valid_id { - return Ok(StatusCode::UNAUTHORIZED.into_response()); - } - let asr_query = req.0.clone(); - let asr = state.0 .0.clone(); - let upg = - ws.write_buffer_size(0).on_upgrade(move |v| asr_websocket(v, asr, asr_query, addr)); - Ok(upg) - } - axum::Router::new() - .route(path, axum::routing::post(t)) - .route(path, axum::routing::get(streaming_t)) - .with_state((s, ss.clone())) -} - -fn py_router_post(s: Arc, path: &str, ss: &SharedState) -> axum::Router<()> { - async fn t( - state: axum::extract::State<(Arc, SharedState)>, - _headers: axum::http::HeaderMap, - req: axum::body::Bytes, - ) -> utils::AxumResult { - tracing::info!("handling py-post query"); - match state.0 .0.run_one(req).await { - Ok(data) => Ok((StatusCode::OK, data).into_response()), - Err(err) => { - tracing::error!(?err, "py-post"); - Ok(StatusCode::INTERNAL_SERVER_ERROR.into_response()) - } - } - } - - axum::Router::new() - .route(path, axum::routing::post(t)) - .with_state((s, ss.clone())) - .layer(axum::extract::DefaultBodyLimit::disable()) - .layer(tower_http::limit::RequestBodyLimitLayer::new(16 * 1024 * 1024)) -} - -fn py_router(s: Arc, path: &str, ss: &SharedState) -> axum::Router<()> { - async fn py_websocket( - socket: axum::extract::ws::WebSocket, - state: Arc, - query: PyStreamingQuery, - _addr: Option, - ) { - if let Err(err) = state.handle_socket(socket, query).await { - tracing::error!(?err, "py") - } - } - - // TODO: add a batch mode. - async fn t( - state: axum::extract::State<(Arc, SharedState)>, - headers: axum::http::HeaderMap, - req: axum::Json, - ) -> utils::AxumResult { - tracing::info!("handling py streaming post query {req:?}"); - let valid_id = headers - .get(ID_HEADER) - .and_then(|v| v.to_str().ok()) - .is_some_and(|id| state.0 .1.config.authorized_ids.contains(id)); - if !valid_id { - return Ok(StatusCode::UNAUTHORIZED.into_response()); - } - let wav = state.0 .0.handle_query(&req).await?; - tracing::info!("ok {}", wav.len()); - Ok((StatusCode::OK, [(axum::http::header::CONTENT_TYPE, "audio/wav")], wav).into_response()) - } - - async fn streaming_t( - ws: axum::extract::ws::WebSocketUpgrade, - headers: axum::http::HeaderMap, - state: axum::extract::State<(Arc, SharedState)>, - req: axum::extract::Query, - ) -> utils::AxumResult { - let addr = headers.get("X-Real-IP").and_then(|v| v.to_str().ok().map(|v| v.to_string())); - tracing::info!(addr, "handling py streaming query"); - // It's tricky to set the headers of a websocket in javascript so we pass the token via the - // query too. - let auth_id = match headers.get(ID_HEADER) { - Some(v) => v.to_str().ok(), - None => req.auth_id.as_deref(), - }; - let valid_id = auth_id.is_some_and(|id| state.1.config.authorized_ids.contains(id)); - if !valid_id { - return Ok(StatusCode::UNAUTHORIZED.into_response()); - } - let py_query = req.0.clone(); - let py = state.0 .0.clone(); - let upg = ws.write_buffer_size(0).on_upgrade(move |v| py_websocket(v, py, py_query, addr)); - Ok(upg) - } - axum::Router::new() - .route(path, axum::routing::post(t)) - .route(path, axum::routing::get(streaming_t)) - .with_state((s, ss.clone())) -} - -#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)] -struct MimiStreamingQuery { - auth_id: Option, - room_id: Option, -} - -fn mimi_router( - s: Arc, - send_path: &str, - recv_path: &str, - ss: &SharedState, -) -> axum::Router<()> { - async fn mimi_recv_websocket( - socket: axum::extract::ws::WebSocket, - state: Arc, - room_id: Option, - _addr: Option, - ) { - if let Err(err) = state.recv_socket(socket, room_id).await { - tracing::error!(?err, "mimi") - } - } - - async fn recv( - ws: axum::extract::ws::WebSocketUpgrade, - headers: axum::http::HeaderMap, - state: axum::extract::State<(Arc, SharedState)>, - req: axum::extract::Query, - ) -> utils::AxumResult { - let addr = headers.get("X-Real-IP").and_then(|v| v.to_str().ok().map(|v| v.to_string())); - tracing::info!(addr, "handling mimi-streaming query"); - // It's tricky to set the headers of a websocket in javascript so we pass the token via the - // query too. - if state.0 .0.auth_recv() { - let auth_id = match headers.get(ID_HEADER) { - Some(v) => v.to_str().ok(), - None => req.auth_id.as_deref(), - }; - let valid_id = auth_id.is_some_and(|id| state.0 .1.config.authorized_ids.contains(id)); - if !valid_id { - return Ok(StatusCode::UNAUTHORIZED.into_response()); - } - } - let room_id = match headers.get(ROOM_ID_HEADER) { - Some(v) => v.to_str().ok().map(|v| v.to_string()), - None => req.room_id.clone(), - }; - let state = state.0 .0.clone(); - let upg = ws - .write_buffer_size(0) - .on_upgrade(move |v| mimi_recv_websocket(v, state, room_id, addr)); - Ok(upg) - } - - async fn mimi_send_websocket( - socket: axum::extract::ws::WebSocket, - state: Arc, - room_id: String, - _addr: Option, - ) { - if let Err(err) = state.send_socket(socket, room_id).await { - tracing::error!(?err, "mimi") - } - } - - async fn send( - ws: axum::extract::ws::WebSocketUpgrade, - headers: axum::http::HeaderMap, - state: axum::extract::State<(Arc, SharedState)>, - req: axum::extract::Query, - ) -> utils::AxumResult { - let addr = headers.get("X-Real-IP").and_then(|v| v.to_str().ok().map(|v| v.to_string())); - tracing::info!(addr, "handling mimi-streaming send query"); - let auth_id = match headers.get(ID_HEADER) { - Some(v) => v.to_str().ok(), - None => req.auth_id.as_deref(), - }; - let valid_id = auth_id.is_some_and(|id| state.0 .1.config.authorized_ids.contains(id)); - if !valid_id { - return Ok(StatusCode::UNAUTHORIZED.into_response()); - } - let room_id = match headers.get(ROOM_ID_HEADER) { - Some(v) => v.to_str().ok().map(|v| v.to_string()), - None => req.room_id.clone(), - }; - let room_id = match room_id { - None => Err(anyhow::format_err!("no room_id"))?, - Some(room_id) => room_id, - }; - let state = state.0 .0; - let upg = ws - .write_buffer_size(0) - .on_upgrade(move |v| mimi_send_websocket(v, state, room_id, addr)); - Ok(upg) - } - axum::Router::new() - .route(send_path, axum::routing::get(send)) - .route(recv_path, axum::routing::get(recv)) - .with_state((s, ss.clone())) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use lazy_static::lazy_static; -use prometheus::{ - histogram_opts, labels, opts, register_counter, register_gauge, register_histogram, -}; -use prometheus::{Counter, Gauge, Histogram}; - -pub mod asr { - use super::*; - lazy_static! { - pub static ref CONNECT: Counter = register_counter!(opts!( - "asr_connect", - "Number of connections to the asr.", - labels! {"handler" => "all",} - )) - .unwrap(); - pub static ref MODEL_STEP_DURATION: Histogram = register_histogram!(histogram_opts!( - "asr_model_step_duration", - "ASR model step duration distribution.", - vec![20e-3, 30e-3, 40e-3, 50e-3, 60e-3, 70e-3, 80e-3], - )) - .unwrap(); - pub static ref CONNECTION_NUM_STEPS: Histogram = register_histogram!(histogram_opts!( - "asr_connection_num_steps", - "ASR model, distribution of number of steps for a connection.", - vec![2., 25., 125., 250., 500., 750., 1125., 1500., 2250., 3000., 4500.], - )) - .unwrap(); - pub static ref OPEN_CHANNELS: Gauge = register_gauge!(opts!( - "asr_open_channels", - "Number of open channels (users currently connected).", - labels! {"handler" => "all",} - )) - .unwrap(); - } -} - -pub mod py { - use super::*; - lazy_static! { - pub static ref CONNECT: Counter = register_counter!(opts!( - "py_connect", - "Number of connections to the py-module.", - labels! {"handler" => "all",} - )) - .unwrap(); - pub static ref TOTAL_STEPS: Counter = register_counter!(opts!( - "py_total_steps", - "Total number of times the python callback was called.", - labels! {"handler" => "all",} - )) - .unwrap(); - pub static ref ACTIVE_STEPS: Counter = register_counter!(opts!( - "py_active_steps", - "Number of times the python callback was called with some active users.", - labels! {"handler" => "all",} - )) - .unwrap(); - pub static ref MISSING_WORDS_STEPS: Counter = register_counter!(opts!( - "py_missing_words_steps", - "Number of times the user failed to send words fast enough.", - labels! {"handler" => "all",} - )) - .unwrap(); - pub static ref COULD_HAVE_RUN_STEPS: Counter = register_counter!(opts!( - "py_could_have_run_steps", - "Number of times we ran the callback with enough words for a user.", - labels! {"handler" => "all",} - )) - .unwrap(); - pub static ref MODEL_STEP_DURATION: Histogram = register_histogram!(histogram_opts!( - "py_model_step_duration", - "py module step duration distribution.", - vec![10e-3, 15e-3, 20e-3, 30e-3, 40e-3, 50e-3, 80e-3], - )) - .unwrap(); - pub static ref CONNECTION_NUM_STEPS: Histogram = register_histogram!(histogram_opts!( - "py_model_connection_num_steps", - "py module number of steps with data being generated.", - vec![2., 25., 62.5, 125., 250., 500., 750.], - )) - .unwrap(); - pub static ref OPEN_CHANNELS: Gauge = register_gauge!(opts!( - "py_open_channels", - "Number of open channels (users currently connected).", - labels! {"handler" => "all",} - )) - .unwrap(); - } -} - -pub mod py_post { - use super::*; - lazy_static! { - pub static ref CONNECT: Counter = register_counter!(opts!( - "py_post_connect", - "Number of connections to the py_post module.", - labels! {"handler" => "all",} - )) - .unwrap(); - pub static ref MODEL_DURATION: Histogram = register_histogram!(histogram_opts!( - "py_post_model_duration", - "py-post model duration distribution.", - vec![20e-3, 30e-3, 40e-3, 50e-3, 60e-3, 70e-3, 80e-3], - )) - .unwrap(); - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::protocol::MsgType; -use anyhow::Result; -use axum::extract::ws; -use candle::{Device, IndexOp, Tensor}; -use std::sync::Arc; - -use kaudio::ogg_opus; - -struct Sender { - tx: tokio::sync::broadcast::Sender, - encoder: kaudio::ogg_opus::Encoder, -} - -impl Sender { - fn send_raw(&mut self, data: &[u8]) -> Result<()> { - let msg = ws::Message::Binary(data.to_vec().into()); - let _ = self.tx.send(msg); - Ok(()) - } - - fn send_pcm(&mut self, pcm: &[f32]) -> Result<()> { - let data = self.encoder.encode_page(pcm)?; - let msg: Vec = [&[MsgType::Audio.to_u8()], data.as_slice()].concat(); - let msg = ws::Message::Binary(msg.into()); - // We do not fail on send errors as these mean that there is no subscribers though - // new subscribers may show up later. - let _ = self.tx.send(msg); - Ok(()) - } - - fn send_ping(&mut self) { - let msg = ws::Message::Binary(vec![MsgType::Ping.to_u8()].into()); - let _ = self.tx.send(msg); - } -} - -struct Room { - sender: Arc>, - header_message: ws::Message, - rx: tokio::sync::broadcast::Receiver, -} - -impl Room { - fn new() -> Result { - let (tx, rx) = tokio::sync::broadcast::channel(10); - let encoder = ogg_opus::Encoder::new(24_000)?; - let header_message: Vec = [&[MsgType::Audio.to_u8()], encoder.header_data()].concat(); - let header_message = ws::Message::Binary(header_message.into()); - let sender = Sender { tx, encoder }; - let sender = Arc::new(tokio::sync::Mutex::new(sender)); - tokio::spawn({ - let sender = sender.clone(); - async move { - loop { - tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; - let mut sender = sender.lock().await; - sender.send_ping(); - } - } - }); - Ok(Self { sender, header_message, rx }) - } -} - -pub struct Mimi { - audio_tokenizer: moshi::mimi::Mimi, - device: Device, - #[allow(unused)] - instance_name: String, - auth_recv: bool, - #[allow(unused)] - log_dir: std::path::PathBuf, - rooms: std::collections::HashMap, - default_room: Option, -} - -impl Mimi { - pub fn new(mimi: &crate::MimiConfig, config: &crate::Config, dev: &Device) -> Result { - let audio_tokenizer = moshi::mimi::load(&mimi.audio_tokenizer_file, Some(8), dev)?; - let mut rooms = std::collections::HashMap::new(); - for room in mimi.rooms.iter() { - rooms.insert(room.to_string(), Room::new()?); - } - - Ok(Self { - audio_tokenizer, - device: dev.clone(), - log_dir: config.log_dir.clone().into(), - instance_name: config.instance_name.clone(), - auth_recv: mimi.auth_recv, - default_room: mimi.default_room.clone(), - rooms, - }) - } - - pub fn auth_recv(&self) -> bool { - self.auth_recv - } - - pub async fn recv_socket(&self, socket: ws::WebSocket, room_id: Option) -> Result<()> { - use futures_util::{SinkExt, StreamExt}; - - let room_id = match (room_id, self.default_room.as_ref()) { - (Some(r), _) => r, - (None, Some(d)) => d.to_string(), - (None, None) => anyhow::bail!("no room_id provided"), - }; - let room = match self.rooms.get(&room_id) { - None => anyhow::bail!("unknown room"), - Some(room) => room, - }; - - // Re-subscribe early to have more chances to have a message immediately available. - let mut rx = room.rx.resubscribe(); - let (mut ws_sender, mut ws_receiver) = socket.split(); - let recv_loop = async move { while ws_receiver.next().await.is_some() {} }; - let mut handshake = vec![MsgType::Handshake.to_u8()]; - handshake.resize(9, 0u8); - if let Err(err) = ws_sender.send(ws::Message::binary(handshake)).await { - tracing::error!("error sending header {err:?}"); - return Ok(()); - } - if let Err(err) = ws_sender.send(room.header_message.clone()).await { - tracing::error!("error sending header {err:?}"); - return Ok(()); - } - let send_loop = async move { - loop { - let msg = match rx.recv().await { - Ok(msg) => msg, - Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue, - Err(tokio::sync::broadcast::error::RecvError::Closed) => { - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - continue; - } - }; - if let Err(err) = ws_sender.send(msg).await { - tracing::error!("exiting recv loop, error in send: {err:?}"); - break; - } - } - }; - tokio::select! { - _ = send_loop => tracing::info!("recv_socket: send loop exited"), - _ = recv_loop => tracing::info!("recv_socket: receiver disconnected"), - } - Ok(()) - } - - pub async fn send_socket(&self, socket: ws::WebSocket, room_id: String) -> Result<()> { - use futures_util::StreamExt; - - tracing::info!("connected to sender for {room_id}"); - let room = match self.rooms.get(&room_id) { - None => anyhow::bail!("unknown room"), - Some(room) => room, - }; - let mut sender = match room.sender.try_lock() { - Ok(s) => s, - Err(_) => anyhow::bail!("already a producer"), - }; - let (_ws_sender, mut ws_receiver) = socket.split(); - let mut audio_tokenizer = self.audio_tokenizer.clone(); - - let mut pcm_all = vec![]; - while let Some(msg) = ws_receiver.next().await { - let msg = match msg? { - ws::Message::Binary(b) => b.to_vec(), - _ => continue, - }; - if msg.is_empty() { - continue; - } - match MsgType::from_u8(msg[0]) { - Ok(MsgType::Text) => { - // Forward directly the text messages. - sender.send_raw(&msg)?; - } - Ok(MsgType::Codes) => { - let codes: Vec = msg[1..] - .chunks_exact(4) - .map(|chunk| u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) - .collect(); - let ncodes = codes.len(); - let codes = Tensor::from_vec(codes, (1, ncodes, 1), &self.device)?; - let pcm = audio_tokenizer.decode_step(&codes.into(), &().into())?; - if let Some(pcm) = pcm.as_option() { - let pcm = pcm.i((0, 0))?.to_vec1::()?; - for v in pcm.into_iter() { - pcm_all.push(v); - if pcm_all.len() == 1920 { - sender.send_pcm(&pcm_all)?; - pcm_all.clear(); - } - } - // Sleep to avoid starving the scheduler. - tokio::time::sleep(tokio::time::Duration::from_millis(1)).await; - } - } - t => { - tracing::warn!("unexpected msg type {t:?}"); - continue; - } - } - } - tracing::info!("send_socket: exiting send loop"); - - Ok(()) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::Result; - -#[derive(Debug, Clone, Copy)] -pub enum MsgType { - Handshake, - Audio, - Text, - Control, - Metadata, - Error, - Ping, - ColoredText, - Image, - Codes, -} - -impl MsgType { - pub fn from_u8(v: u8) -> Result { - let s = match v { - 0 => MsgType::Handshake, - 1 => MsgType::Audio, - 2 => MsgType::Text, - 3 => MsgType::Control, - 4 => MsgType::Metadata, - 5 => MsgType::Error, - 6 => MsgType::Ping, - 7 => MsgType::ColoredText, - 8 => MsgType::Image, - 9 => MsgType::Codes, - _ => anyhow::bail!("unexpected msg type {v}"), - }; - Ok(s) - } - - pub fn to_u8(self) -> u8 { - match self { - MsgType::Handshake => 0, - MsgType::Audio => 1, - MsgType::Text => 2, - MsgType::Control => 3, - MsgType::Metadata => 4, - MsgType::Error => 5, - MsgType::Ping => 6, - MsgType::ColoredText => 7, - MsgType::Image => 8, - MsgType::Codes => 9, - } - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::metrics::py_post as metrics; -use crate::py_module::{toml_to_py, VerbosePyErr}; -use anyhow::{Context, Result}; -use numpy::{PyArrayMethods, PyUntypedArrayMethods}; -use pyo3::prelude::*; -use pyo3_ffi::c_str; -use tokio::task; - -type Out = (Vec, Vec); -struct ModelQuery { - pcm: Vec, - out_tx: tokio::sync::oneshot::Sender, -} - -pub struct Inner { - app: PyObject, - in_rx: std::sync::mpsc::Receiver, -} - -#[derive(serde::Serialize, serde::Deserialize)] -#[serde(tag = "type")] -pub enum OutMsg { - Voice { embeddings: Vec, shape: Vec }, -} - -impl Inner { - fn start_model_loop(self) -> Result<()> { - // use numpy::{PyArrayMethods, ToPyArray}; - let model_loop: task::JoinHandle> = task::spawn_blocking(move || { - while let Ok(req) = self.in_rx.recv() { - if let Err(err) = self.handle_query(req) { - tracing::error!(?err, "failed to handle query"); - } - } - Ok(()) - }); - task::spawn(async { - match model_loop.await { - Err(err) => tracing::error!(?err, "model loop join err"), - Ok(Err(err)) => tracing::error!(?err, "model loop err"), - Ok(Ok(())) => tracing::info!("model loop exited"), - } - }); - Ok(()) - } - - fn handle_query(&self, req: ModelQuery) -> Result<()> { - let start_time = std::time::Instant::now(); - let emb = Python::with_gil(|py| -> Result<_> { - let pcm = numpy::PyArray1::from_vec(py, req.pcm); - let emb = self.app.call_method1(py, "run_one", (pcm,)).map_err(VerbosePyErr::from)?; - let emb = match emb.downcast_bound::>(py) { - Ok(emb) => emb, - Err(_) => { - anyhow::bail!("failed to downcast to PyArrayDyn") - } - }; - let shape = emb.shape().to_vec(); - tracing::info!(?shape, "generated embeddings"); - Ok((emb.to_vec()?, shape)) - })?; - let elapsed = start_time.elapsed().as_secs_f64(); - metrics::MODEL_DURATION.observe(elapsed); - if let Err(err) = req.out_tx.send(emb) { - anyhow::bail!("failed to send response: {err:?}"); - } - Ok(()) - } -} - -pub struct M { - in_tx: std::sync::mpsc::Sender, -} - -impl M { - pub fn new(config: crate::PyPostConfig) -> Result { - crate::py_module::init()?; - let (script, script_name) = match &config.script { - None => { - let script_name = std::ffi::CString::new("voice.py")?; - let script = std::ffi::CString::new(crate::VOICE_PY)?; - (script, script_name) - } - Some(script) => { - let script_name = std::ffi::CString::new(script.as_bytes())?; - let script = - std::fs::read_to_string(script).with_context(|| format!("{script:?}"))?; - let script = std::ffi::CString::new(script)?; - (script, script_name) - } - }; - let app = Python::with_gil(|py| -> Result<_> { - let py_config = pyo3::types::PyDict::new(py); - if let Some(cfg) = config.py.as_ref() { - for (key, value) in cfg.iter() { - py_config.set_item(key, toml_to_py(py, value)?)?; - } - } - let app = - PyModule::from_code(py, script.as_c_str(), script_name.as_c_str(), c_str!("foo")) - .map_err(VerbosePyErr::from)? - .getattr("init")? - .call1((py_config,)) - .map_err(VerbosePyErr::from)?; - Ok(app.unbind()) - })?; - let (in_tx, in_rx) = std::sync::mpsc::channel(); - let inner = Inner { app, in_rx }; - inner.start_model_loop()?; - Ok(Self { in_tx }) - } - - pub async fn run_one(&self, data: axum::body::Bytes) -> Result { - use serde::Serialize; - - metrics::CONNECT.inc(); - let (out_tx, out_rx) = tokio::sync::oneshot::channel::(); - let pcm = task::spawn_blocking(move || -> Result> { - let (pcm, sample_rate) = crate::utils::pcm_decode(data)?; - let mut pcm = if sample_rate == 24000 { - pcm - } else { - kaudio::resample(&pcm, sample_rate as usize, 24000)? - }; - pcm.resize(240000, 0.0); - Ok(pcm) - }); - let pcm = pcm.await??; - let query = ModelQuery { pcm, out_tx }; - self.in_tx.send(query)?; - let (embeddings, shape) = out_rx.await?; - let msg = OutMsg::Voice { embeddings, shape }; - let mut bytes = vec![]; - msg.serialize( - &mut rmp_serde::Serializer::new(&mut bytes).with_human_readable().with_struct_map(), - )?; - - Ok(bytes.into()) - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use crate::metrics::py as metrics; -use crate::PyStreamingQuery as Query; -use crate::StreamingOutput; -use anyhow::{Context, Result}; -use axum::extract::ws; -use numpy::PyArrayMethods; -use pyo3::prelude::*; -use pyo3_ffi::c_str; -use std::sync::{Arc, Mutex}; -use tokio::task; -use tokio::time::{timeout, Duration}; - -const FRAME_SIZE: usize = 1920; -const MASK_HAS_PCM: u8 = 1 << 0; -const MASK_IS_EOS: u8 = 1 << 1; -const MASK_WORD_FINISHED: u8 = 1 << 2; -const MASK_AR_STEP: u8 = 1 << 3; -const MASK_MISSING_WORDS: u8 = 1 << 4; - -const SEND_PING_EVERY: Duration = Duration::from_secs(10); -const POST_RETRY_DELAY: Duration = Duration::from_millis(100); -const POST_MAX_RETRIES: usize = 1000; - -#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)] -pub struct TtsQuery { - text: String, - voice: String, -} - -pub struct VerbosePyErr { - err: PyErr, -} - -impl From for VerbosePyErr { - fn from(err: PyErr) -> Self { - Self { err } - } -} - -fn get_traceback(py: Python<'_>, err: &PyErr) -> Result { - let traceback_mod = PyModule::import(py, "traceback")?; - let func = traceback_mod.getattr("format_exception")?; - let traceback_obj = func.call1((err.get_type(py), err.value(py), err.traceback(py)))?; - let lines = traceback_obj.extract::>()?; - Ok(lines.join("")) -} - -impl std::error::Error for VerbosePyErr {} - -impl std::fmt::Display for VerbosePyErr { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - Python::with_gil(|py| { - let traceback = match get_traceback(py, &self.err) { - Err(_) => "no traceback".to_string(), - Ok(traceback) => traceback, - }; - write!(f, "{}\n{}", self.err, traceback) - }) - } -} - -impl std::fmt::Debug for VerbosePyErr { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - ::fmt(self, f) - } -} - -#[derive(serde::Serialize, serde::Deserialize)] -#[serde(tag = "type")] -pub enum InMsg { - Text { text: String }, - Voice { embeddings: Vec, shape: Vec }, - Eos, -} - -#[derive(Debug, Clone)] -pub enum Msg { - Text(String, Vec), - Voice { embeddings: Vec, shape: Vec }, - Eos, -} - -/// Unique identifier. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -pub struct ChannelId(usize); - -impl ChannelId { - fn new() -> Self { - // https://users.rust-lang.org/t/idiomatic-rust-way-to-generate-unique-id/33805 - use std::sync::atomic; - static COUNTER: atomic::AtomicUsize = atomic::AtomicUsize::new(1); - Self(COUNTER.fetch_add(1, atomic::Ordering::Relaxed)) - } -} - -type InSend = std::sync::mpsc::Sender; -type InRecv = std::sync::mpsc::Receiver; -type OutSend = tokio::sync::mpsc::UnboundedSender>; -type OutRecv = tokio::sync::mpsc::UnboundedReceiver>; - -struct Channel { - id: ChannelId, - in_rx: InRecv, - out_tx: OutSend, - encoder: crate::tts::Encoder, - voice: Option, - sent_init: bool, - words: std::collections::VecDeque, - steps: usize, - prev_word_steps: usize, -} - -impl Channel { - fn new( - in_rx: InRecv, - out_tx: OutSend, - encoder: crate::tts::Encoder, - voice: Option, - ) -> Self { - metrics::OPEN_CHANNELS.inc(); - let words = std::collections::VecDeque::new(); - Self { - id: ChannelId::new(), - in_rx, - out_tx, - encoder, - words, - voice: voice.map(Voice::File), - sent_init: false, - steps: 0, - prev_word_steps: 0, - } - } -} - -impl Drop for Channel { - fn drop(&mut self) { - metrics::CONNECTION_NUM_STEPS.observe(self.steps as f64); - metrics::OPEN_CHANNELS.dec(); - } -} - -pub fn init() -> PyResult<()> { - pyo3::prepare_freethreaded_python(); - Python::with_gil(|py| -> PyResult<()> { - let signal = py.import("signal")?; - // Set SIGINT to have the default action rather than triggering a Python exception - signal.getattr("signal")?.call1((signal.getattr("SIGINT")?, signal.getattr("SIG_DFL")?))?; - Ok(()) - })?; - Ok(()) -} - -type Channels = Arc>>>; - -struct Inner { - channels: Channels, - app: PyObject, -} - -enum Voice { - File(String), - Embeddings { embeddings: Vec, shape: Vec }, -} - -impl<'py> IntoPyObject<'py> for Voice { - type Target = PyAny; - type Output = Bound<'py, PyAny>; - type Error = PyErr; - fn into_pyobject(self, py: Python<'py>) -> std::result::Result { - use numpy::ToPyArray; - let go = |s| -> PyResult<_> { - let any = match s { - Voice::File(v) => v.into_pyobject(py)?.into_any(), - Voice::Embeddings { embeddings, shape } => match *shape.as_slice() { - [dim1] => embeddings.to_pyarray(py).reshape((dim1,))?.into_any(), - [dim1, dim2] => embeddings.to_pyarray(py).reshape((dim1, dim2))?.into_any(), - [d1, d2, d3] => embeddings.to_pyarray(py).reshape((d1, d2, d3))?.into_any(), - _ => return Ok(py.None().into_bound(py)), - }, - }; - Ok(any) - }; - // We convert errors to None, this should result in using the default voice rather than - // crashing the whole process. - match go(self) { - Ok(any) => Ok(any), - Err(_) => Ok(py.None().into_bound(py)), - } - } -} - -// The arguments passed to the python step function, for now this is: -// (batch_idx, tokens, voice) -// tokens can include a -1 to indicate a new user, and a -2 to indicate -// end of stream. -type PyInput = (usize, Vec, Option); - -impl Inner { - fn pre_process(&self, _step_idx: usize) -> Result<(Vec, Vec>)> { - let mut channels = self.channels.lock().unwrap(); - let mut in_data = vec![]; - let mut channel_ids = Vec::with_capacity(channels.len()); - for (batch_idx, channel) in channels.iter_mut().enumerate() { - channel_ids.push(channel.as_ref().map(|c| c.id)); - if let Some(c) = channel.as_mut() { - if c.out_tx.is_closed() { - *channel = None; - } else { - use std::sync::mpsc::TryRecvError; - match c.in_rx.try_recv() { - Ok(Msg::Text(word, tokens)) => { - c.words.push_back(word); - let mut t = Vec::with_capacity(tokens.len() + 1); - if !c.sent_init { - t.push(-1); - c.sent_init = true; - } - for &v in tokens.iter() { - t.push(v as i32); - } - in_data.push((batch_idx, t, c.voice.take())); - } - Ok(Msg::Voice { embeddings, shape }) => { - c.voice = Some(Voice::Embeddings { embeddings, shape }); - } - Ok(Msg::Eos) => { - if c.sent_init { - in_data.push((batch_idx, vec![-2], None)); - } else { - *channel = None - } - } - Err(TryRecvError::Empty) => {} - Err(TryRecvError::Disconnected) => *channel = None, - } - }; - } - } - Ok((in_data, channel_ids)) - } - - fn start_model_loop(self, batch_size: usize) -> Result<()> { - use numpy::{PyArrayMethods, ToPyArray}; - use rayon::prelude::*; - use std::ops::DerefMut; - - let model_loop: task::JoinHandle> = task::spawn_blocking(move || { - // Maybe the model loop could just always hold the gil? - tracing::info!("starting-up the py model loop"); - let pcm_data = numpy::ndarray::Array2::::zeros([batch_size, FRAME_SIZE]); - let pcm_data = Python::with_gil(|py| pcm_data.to_pyarray(py).unbind()); - let mask = numpy::ndarray::Array1::::zeros([batch_size]); - let mask = Python::with_gil(|py| mask.to_pyarray(py).unbind()); - let tokens = numpy::ndarray::Array2::::zeros([batch_size, 33]); - let tokens = Python::with_gil(|py| tokens.to_pyarray(py).unbind()); - - for step_idx in 0.. { - // We store the channel ids here to check that they have not changed when sending - // the data back to the user. - let (in_data, channel_ids) = self.pre_process(step_idx)?; - let start_time = std::time::Instant::now(); - Python::with_gil(|py| -> Result<()> { - self.app - .call_method1(py, "step", (in_data, &pcm_data, &mask, &tokens)) - .map_err(VerbosePyErr::from)?; - let elapsed = start_time.elapsed().as_secs_f64(); - let pcm = pcm_data.bind(py).readonly(); - let mask = mask.bind(py).readonly(); - let tokens = tokens.bind(py).readonly(); - let pcm = pcm.as_slice().context("pcm is not contiguous")?; - let mask = mask.as_slice().context("mask is not contiguous")?; - let _tokens = tokens.as_slice().context("tokens is not contiguous")?; - - // Only store the sample is something was actually done. - if mask.iter().any(|&x| (x & MASK_AR_STEP) > 0) { - metrics::MODEL_STEP_DURATION.observe(elapsed); - metrics::ACTIVE_STEPS.inc(); - } - metrics::TOTAL_STEPS.inc(); - - let mut channels = self.channels.lock().unwrap(); - let c = channels.deref_mut(); - - c.par_iter_mut().enumerate().for_each(|(batch_idx, channel)| { - if let Some(c) = channel.as_mut() { - let mask = mask[batch_idx]; - // The channel has changed so skip the update. - if Some(c.id) != channel_ids[batch_idx] { - return; - } - if (mask & MASK_AR_STEP) > 0 { - c.steps += 1; - } - if c.sent_init { - if (mask & MASK_MISSING_WORDS) > 0 { - metrics::MISSING_WORDS_STEPS.inc(); - } else { - metrics::COULD_HAVE_RUN_STEPS.inc(); - } - } - if (mask & MASK_WORD_FINISHED) > 0 { - if let Some(text) = c.words.pop_front() { - let wwts = crate::tts::WordWithTimestamps { - text, - start_s: c.prev_word_steps as f64 / 12.5, - stop_s: c.steps as f64 / 12.5, - }; - c.prev_word_steps = c.steps; - match c.encoder.encode_word(wwts) { - Ok(Some(msg)) => { - let _ = c.out_tx.send(msg).is_err(); - } - Ok(None) => {} - Err(err) => { - tracing::error!(?err, ?batch_idx, "encoder word error") - } - } - } - } - if (mask & MASK_HAS_PCM) > 0 { - let pcm = pcm[batch_idx * FRAME_SIZE..(batch_idx + 1) * FRAME_SIZE] - .to_vec(); - match c.encoder.encode(pcm) { - Ok(msg) => { - if c.out_tx.send(msg).is_err() { - *channel = None; - } - } - Err(err) => { - tracing::error!(?err, ?batch_idx, "encoder error") - } - } - } - // The TTS has finished generating so we close the channel, this should - // drop out_tx and result in the websock closing. - if (mask & MASK_IS_EOS) > 0 { - tracing::info!(?batch_idx, "tts finished"); - *channel = None; - } - } - }); - Ok(()) - })?; - } - Ok(()) - }); - task::spawn(async { - match model_loop.await { - Err(err) => tracing::error!(?err, "model loop join err"), - Ok(Err(err)) => tracing::error!(?err, "model loop err"), - Ok(Ok(())) => tracing::info!("model loop exited"), - } - }); - Ok(()) - } -} - -#[derive(Clone)] -pub struct M { - channels: Channels, - text_tokenizer: std::sync::Arc, - config: crate::PyConfig, -} - -pub(crate) fn toml_to_py<'a>(py: Python<'a>, value: &toml::Value) -> Result> { - let value = match value { - toml::Value::Float(v) => v.into_pyobject(py)?.into_any(), - toml::Value::Integer(v) => v.into_pyobject(py)?.into_any(), - toml::Value::String(v) => v.into_pyobject(py)?.into_any(), - toml::Value::Boolean(v) => v.into_pyobject(py)?.to_owned().into_any(), - toml::Value::Table(table) => { - let v = pyo3::types::PyDict::new(py); - for (key, value) in table.iter() { - v.set_item(key, toml_to_py(py, value)?)?; - } - v.into_any() - } - toml::Value::Array(vs) => { - let v = pyo3::types::PyList::empty(py); - for value in vs.iter() { - v.append(toml_to_py(py, value)?)?; - } - v.into_any() - } - toml::Value::Datetime(_) => { - anyhow::bail!("unsupported value type DateTime") - } - }; - Ok(value) -} - -fn text_pre_process(text: &str) -> String { - text.replace('’', "'").replace('–', "").replace(':', " ").replace(['(', ')'], "") -} - -impl M { - pub fn new(config: crate::PyConfig) -> Result { - init()?; - let text_tokenizer = - sentencepiece::SentencePieceProcessor::open(&config.text_tokenizer_file) - .with_context(|| config.text_tokenizer_file.clone())?; - let batch_size = config.batch_size; - let (script, script_name) = match &config.script { - None => { - let script_name = std::ffi::CString::new("tts.py")?; - let script = std::ffi::CString::new(crate::TTS_PY)?; - (script, script_name) - } - Some(script) => { - let script_name = std::ffi::CString::new(script.as_bytes())?; - let script = - std::fs::read_to_string(script).with_context(|| format!("{script:?}"))?; - let script = std::ffi::CString::new(script)?; - (script, script_name) - } - }; - let app = Python::with_gil(|py| -> Result<_> { - let py_config = pyo3::types::PyDict::new(py); - if let Some(cfg) = config.py.as_ref() { - for (key, value) in cfg.iter() { - py_config.set_item(key, toml_to_py(py, value)?)?; - } - } - let app = - PyModule::from_code(py, script.as_c_str(), script_name.as_c_str(), c_str!("foo")) - .map_err(VerbosePyErr::from)? - .getattr("init")? - .call1((batch_size.into_pyobject(py)?, py_config)) - .map_err(VerbosePyErr::from)?; - Ok(app.unbind()) - })?; - let channels = (0..batch_size).map(|_| None).collect::>(); - let channels = Arc::new(Mutex::new(channels)); - let text_tokenizer = Arc::new(text_tokenizer); - let inner = Inner { app, channels: channels.clone() }; - inner.start_model_loop(batch_size)?; - Ok(Self { config, channels, text_tokenizer }) - } - - // Returns None if no channel is available at the moment. - fn channels( - &self, - format: StreamingOutput, - voice: Option, - ) -> Result> { - let mut channels = self.channels.lock().unwrap(); - // Linear scan to find an available channel. This is fairly inefficient, instead we should - // probably have a queue of available slots. - for (batch_idx, channel) in channels.iter_mut().enumerate() { - if channel.is_none() { - let (in_tx, in_rx) = std::sync::mpsc::channel::(); - let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::>(); - let mut encoder = crate::tts::Encoder::new(format)?; - if let Some(msg) = encoder.encode_msg(crate::tts::OutMsg::Ready)? { - out_tx.send(msg)? - } - if let Some(header) = encoder.header()? { - out_tx.send(header)? - } - let c = Channel::new(in_rx, out_tx, encoder, voice.clone()); - *channel = Some(c); - return Ok(Some((batch_idx, in_tx, out_rx))); - } - } - Ok(None) - } - - // TODO: Add a proper batch variant that would enqueue the task so that it can be processed - // when there is a free channel. - pub async fn handle_query(&self, query: &TtsQuery) -> Result> { - tracing::info!("py handle-query"); - metrics::CONNECT.inc(); - let text_tokenizer = self.text_tokenizer.clone(); - let text_bos_token = self.config().text_bos_token; - let (batch_idx, in_tx, mut out_rx) = { - let mut num_tries = 0; - loop { - match self.channels(StreamingOutput::Pcm, Some(query.voice.clone())) { - Ok(Some(x)) => break x, - Ok(None) => { - num_tries += 1; - if num_tries > POST_MAX_RETRIES { - tracing::error!("no free channels after 1000 tries"); - anyhow::bail!("no free channels"); - } - tokio::time::sleep(POST_RETRY_DELAY).await; - } - Err(err) => { - tracing::error!(?err, "no free channels"); - Err(err)? - } - } - } - }; - tracing::info!(batch_idx, "batched-py channel"); - let mut inserted_bos = false; - let text = text_pre_process(&query.text); - for word in text.split_whitespace() { - if word.is_empty() { - continue; - } - let mut word_tokens: Vec<_> = - text_tokenizer.encode(word)?.into_iter().map(|v| v.id).collect(); - if !inserted_bos { - inserted_bos = true; - word_tokens.insert(0, text_bos_token) - } - in_tx.send(Msg::Text(word.to_string(), word_tokens))?; - } - in_tx.send(Msg::Eos)?; - let mut pcm = vec![]; - tracing::info!(batch_idx, "starting the receive loop"); - while let Some(data) = out_rx.recv().await { - pcm.push(data) - } - let pcm = pcm.into_iter().flatten().collect::>(); - - let pcm = { - use byteorder::ByteOrder; - let mut buf = vec![0f32; pcm.len() / std::mem::size_of::()]; - byteorder::LittleEndian::read_f32_into(&pcm, &mut buf); - buf - }; - - let mut wav = vec![]; - moshi::wav::write_pcm_as_wav(&mut wav, &pcm, 24_000)?; - Ok(wav) - } - - pub async fn handle_socket(&self, socket: ws::WebSocket, query: Query) -> Result<()> { - use futures_util::{SinkExt, StreamExt}; - - tracing::info!(?query, "py query"); - metrics::CONNECT.inc(); - - let (mut sender, receiver) = socket.split(); - let (bidx, in_tx, mut out_rx) = match self.channels(query.format, query.voice.clone())? { - Some(x) => x, - None => { - tracing::error!("no free channels"); - let mut encoder = crate::tts::Encoder::new(query.format)?; - let msg = crate::tts::OutMsg::Error { message: "no free channels".into() }; - if let Some(msg) = encoder.encode_msg(msg)? { - sender.send(ws::Message::binary(msg)).await?; - sender.close().await?; - } - anyhow::bail!("no free channels") - } - }; - tracing::info!(?bidx, "batched-py channel"); - let text_tokenizer = self.text_tokenizer.clone(); - let text_bos_token = self.config().text_bos_token; - - let recv_loop = task::spawn(async move { - let timeout_duration = SEND_PING_EVERY * 3; - let mut receiver = receiver; - let mut inserted_bos = false; - let mut send_text = |msg: &str| -> Result<()> { - let msg = text_pre_process(msg); - for word in msg.split_whitespace() { - if word.is_empty() { - continue; - } - let mut word_tokens: Vec<_> = - text_tokenizer.encode(word)?.into_iter().map(|v| v.id).collect(); - if !inserted_bos { - inserted_bos = true; - word_tokens.insert(0, text_bos_token) - } - in_tx.send(Msg::Text(word.to_string(), word_tokens))?; - } - Ok(()) - }; - loop { - use ws::Message; - let msg = match timeout(timeout_duration, receiver.next()).await { - Ok(Some(msg)) => msg, - Ok(None) => break, - Err(_) => { - tracing::info!(?bidx, "recv loop short timeout"); - break; - } - }; - match msg? { - Message::Text(text) => send_text(&text)?, - Message::Binary(msg) => { - if msg.as_ref() == b"\0" { - tracing::info!(?bidx, "received end of stream"); - in_tx.send(Msg::Eos)? - } else { - let msg: InMsg = rmp_serde::from_slice(&msg)?; - match msg { - InMsg::Eos => in_tx.send(Msg::Eos)?, - InMsg::Text { text } => send_text(&text)?, - InMsg::Voice { embeddings, shape } => { - in_tx.send(Msg::Voice { embeddings, shape })? - } - } - } - } - // ping messages are automatically answered by tokio-tungstenite as long as - // the connection is read from. - Message::Ping(_) | Message::Pong(_) => {} - Message::Close(_) => break, - }; - } - Ok::<_, anyhow::Error>(()) - }); - let send_loop = task::spawn(async move { - let mut sender = sender; - let mut last_ping_sent = std::time::Instant::now(); - loop { - // The recv method is cancel-safe so can be wrapped in a timeout. - let msg = timeout(SEND_PING_EVERY, out_rx.recv()).await; - let now = std::time::Instant::now(); - if now.duration_since(last_ping_sent) > SEND_PING_EVERY { - last_ping_sent = now; - sender.send(ws::Message::Ping(vec![].into())).await?; - } - if let Ok(msg) = msg { - match msg { - None => break, - Some(msg) => { - let msg = ws::Message::binary(msg); - sender.send(msg).await?; - } - } - }; - } - sender.close().await?; - drop(sender); - Ok::<(), anyhow::Error>(()) - }); - - // Keep track of the outputs of the different threads. - task::spawn(async { - match send_loop.await { - Err(err) => tracing::error!(?err, "send loop join err"), - Ok(Err(err)) => tracing::error!(?err, "send loop err"), - Ok(Ok(())) => tracing::info!("send loop exited"), - } - }); - task::spawn(async { - match recv_loop.await { - Err(err) => tracing::error!(?err, "recv loop join err"), - Ok(Err(err)) => tracing::error!(?err, "recv loop err"), - Ok(Ok(())) => tracing::info!("recv loop exited"), - } - }); - - Ok(()) - } - - pub fn config(&self) -> &crate::PyConfig { - &self.config - } - - pub fn total_slots(&self) -> usize { - self.config.batch_size - } - - pub fn used_slots(&self) -> usize { - self.channels.lock().unwrap().iter().filter(|v| v.is_some()).count() - } -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::{Context, Result}; -use axum::extract::ws; -use candle::{DType, Device, IndexOp, Tensor}; -use candle_nn::VarBuilder; -use moshi::tts_streaming::Speaker; - -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct WordWithTimestamps { - pub text: String, - pub start_s: f64, - pub stop_s: f64, -} - -pub struct Model { - lm: moshi::lm::LmModel, - audio_tokenizer: moshi::mimi::Mimi, - text_tokenizer: std::sync::Arc, - speaker_encoder: moshi::tts_streaming::SpeakerEncoder, - ca_srcs: std::collections::HashMap, - tts_config: moshi::tts_streaming::Config, - instance_name: String, - voice_dir: std::path::PathBuf, - log_dir: std::path::PathBuf, - // Dummy way to ensure that only a single inference can happen. - pub(crate) mutex: tokio::sync::Mutex<()>, -} - -pub enum Encoder { - OggOpus(kaudio::ogg_opus::Encoder), - OggOpusMessagePack(kaudio::ogg_opus::Encoder), - Pcm, - PcmMessagePack, -} - -enum LogMessage { - Text(String), - Slice(u32, Tensor), -} - -#[derive(serde::Serialize)] -struct QueryWithTexts<'a, Q: serde::Serialize> { - #[serde(flatten)] - query: &'a Q, - texts: Vec, -} - -#[derive(Clone)] -struct LogSender(std::sync::mpsc::Sender); -struct Logger(std::sync::mpsc::Receiver); - -fn logger() -> (LogSender, Logger) { - let (log_tx, log_rx) = std::sync::mpsc::channel(); - (LogSender(log_tx), Logger(log_rx)) -} - -impl LogSender { - fn send(&self, msg: LogMessage) { - let _err = self.0.send(msg); - } - - fn send_text(&self, text: String) { - self.send(LogMessage::Text(text)); - } - - fn send_slice(&self, idx: u32, slice: Tensor) { - self.send(LogMessage::Slice(idx, slice)); - } -} - -impl Logger { - fn save, T: serde::Serialize>( - self, - query: &T, - log_dir: P, - instance_name: &str, - ) -> Result<()> { - // Use log_rx.iter() to wait on the process loop being done. - - let mut text_tokens = vec![]; - let mut audio_tokens = vec![]; - let mut texts = vec![]; - for elem in self.0.into_iter() { - match elem { - LogMessage::Text(text) => { - texts.push(text); - } - LogMessage::Slice(idx, slice) => { - audio_tokens.push(slice); - text_tokens.push(idx); - } - } - } - let text_tokens = text_tokens.into_iter().map(|v| (v, Speaker::Main)).collect::>(); - let audio_tokens = Tensor::cat(&audio_tokens, candle::D::Minus1)?; - let since_epoch = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?; - let (secs, us) = (since_epoch.as_secs(), since_epoch.subsec_micros()); - let base_path = log_dir.as_ref().join(format!("{instance_name}-tts-{secs}-{us}")); - let json_filename = base_path.with_extension("json"); - let query = QueryWithTexts { query, texts }; - let json_content = serde_json::to_string_pretty(&query)?; - std::fs::write(json_filename, json_content)?; - let st_filename = base_path.with_extension("safetensors"); - let text_tokens: Vec<_> = text_tokens.iter().map(|v| v.0 as i64).collect(); - let text_len = text_tokens.len(); - let text_tokens = candle::Tensor::from_vec(text_tokens, text_len, &candle::Device::Cpu)? - .to_dtype(DType::I64)?; - let audio_tokens = audio_tokens.to_device(&Device::Cpu)?.to_dtype(DType::I64)?; - let st_content = - std::collections::HashMap::from([("text", text_tokens), ("audio", audio_tokens)]); - candle::safetensors::save(&st_content, st_filename)?; - Ok(()) - } -} - -#[derive(serde::Serialize, serde::Deserialize)] -#[serde(tag = "type")] -pub enum OutMsg { - Text { text: String, start_s: f64, stop_s: f64 }, - Audio { pcm: Vec }, - OggOpus { data: Vec }, - Error { message: String }, - Ready, -} - -impl Encoder { - pub fn new(format: crate::StreamingOutput) -> Result { - match format { - crate::StreamingOutput::OggOpus => Self::ogg_opus(24000), - crate::StreamingOutput::OggOpusMessagePack => Self::ogg_opus_message_pack(24000), - crate::StreamingOutput::Pcm => Ok(Self::pcm()), - crate::StreamingOutput::PcmMessagePack => Ok(Self::pcm_message_pack()), - } - } - - fn ogg_opus(sample_rate: usize) -> Result { - Ok(Self::OggOpus(kaudio::ogg_opus::Encoder::new(sample_rate)?)) - } - - fn ogg_opus_message_pack(sample_rate: usize) -> Result { - Ok(Self::OggOpusMessagePack(kaudio::ogg_opus::Encoder::new(sample_rate)?)) - } - - fn pcm_message_pack() -> Self { - Self::PcmMessagePack - } - - fn pcm() -> Self { - Self::Pcm - } - - pub fn header(&self) -> Result>> { - let header = match self { - Self::OggOpus(oo) => Some(oo.header_data().to_vec()), - Self::OggOpusMessagePack(oo) => { - use serde::Serialize; - let msg = OutMsg::OggOpus { data: oo.header_data().to_vec() }; - let mut buf = vec![]; - msg.serialize( - &mut rmp_serde::Serializer::new(&mut buf) - .with_human_readable() - .with_struct_map(), - )?; - Some(buf) - } - Self::Pcm => None, - Self::PcmMessagePack => None, - }; - Ok(header) - } - - pub fn encode_word(&self, wwts: WordWithTimestamps) -> Result>> { - if wwts.text.is_empty() { - return Ok(None); - } - let buf = match self { - Self::Pcm | Self::OggOpus(_) => None, - Self::OggOpusMessagePack(_) | Self::PcmMessagePack => { - use serde::Serialize; - let mut buf = vec![]; - OutMsg::Text { text: wwts.text, start_s: wwts.start_s, stop_s: wwts.stop_s } - .serialize( - &mut rmp_serde::Serializer::new(&mut buf) - .with_human_readable() - .with_struct_map(), - )?; - Some(buf) - } - }; - Ok(buf) - } - - pub fn encode(&mut self, pcm: Vec) -> Result> { - use serde::Serialize; - let buf = match self { - Self::OggOpus(oo) => oo.encode_page(&pcm)?, - Self::OggOpusMessagePack(oo) => { - let data = oo.encode_page(&pcm)?; - let mut buf = vec![]; - OutMsg::OggOpus { data }.serialize( - &mut rmp_serde::Serializer::new(&mut buf) - .with_human_readable() - .with_struct_map(), - )?; - buf - } - Self::PcmMessagePack => { - let mut buf = vec![]; - OutMsg::Audio { pcm }.serialize( - &mut rmp_serde::Serializer::new(&mut buf) - .with_human_readable() - .with_struct_map(), - )?; - buf - } - Self::Pcm => { - use byteorder::ByteOrder; - let mut buf = vec![0u8; std::mem::size_of_val(pcm.as_slice())]; - byteorder::LittleEndian::write_f32_into(&pcm, &mut buf); - buf - } - }; - Ok(buf) - } - - pub fn encode_msg(&mut self, msg: OutMsg) -> Result>> { - use serde::Serialize; - let buf = match self { - Self::OggOpus(_) | Self::Pcm => None, - Self::OggOpusMessagePack(_) | Self::PcmMessagePack => { - let mut buf = vec![]; - msg.serialize( - &mut rmp_serde::Serializer::new(&mut buf) - .with_human_readable() - .with_struct_map(), - )?; - Some(buf) - } - }; - Ok(buf) - } -} - -impl Model { - pub fn new(tts: &crate::TtsConfig, config: &crate::Config, dev: &Device) -> Result { - let dtype = dev.bf16_default_to_f32(); - let model_config = &tts.model; - let audio_codebooks = model_config.audio_codebooks; - let audio_tokenizer = - moshi::mimi::load(&tts.audio_tokenizer_file, Some(audio_codebooks), dev)?; - let speaker_tokenizer = if tts.speaker_tokenizer_file == tts.audio_tokenizer_file { - audio_tokenizer.clone() - } else if tts.speaker_tokenizer_file.is_empty() { - let vb_lm = unsafe { - VarBuilder::from_mmaped_safetensors(&[&tts.lm_model_file], DType::F32, dev)? - }; - let cfg = moshi::mimi::Config::v0_1(None); - moshi::mimi::Mimi::new( - cfg, - vb_lm.pp("condition_provider.conditioners.speaker_wavs.compression_model"), - )? - } else { - moshi::mimi::load(&tts.speaker_tokenizer_file, None, dev)? - }; - let vb_lm = - unsafe { VarBuilder::from_mmaped_safetensors(&[&tts.lm_model_file], dtype, dev)? }; - let speaker_encoder = moshi::tts_streaming::SpeakerEncoder::new( - speaker_tokenizer, - tts.generation.speaker_cond_dim, - tts.generation.speaker_cond_n_speakers, - dtype, - vb_lm.to_dtype(DType::F32), - )?; - let text_tokenizer = sentencepiece::SentencePieceProcessor::open(&tts.text_tokenizer_file) - .with_context(|| tts.text_tokenizer_file.clone())?; - let mut ca_srcs = std::collections::HashMap::new(); - for (name, path) in tts.voices.iter() { - let ca_src = match candle::safetensors::load(path, dev)?.get("ca_src") { - Some(ca_src) => ca_src.clone(), - None => anyhow::bail!("missing ca_src tensor in {path}"), - }; - let ca_src = ca_src.narrow(0, 0, 1)?.to_dtype(dtype)?; - ca_srcs.insert(name.to_string(), ca_src); - } - let lm = moshi::lm::LmModel::new( - model_config, - moshi::nn::MaybeQuantizedVarBuilder::Real(vb_lm), - )?; - Ok(Self { - lm, - audio_tokenizer, - text_tokenizer: std::sync::Arc::new(text_tokenizer), - speaker_encoder, - ca_srcs, - tts_config: tts.generation.clone(), - instance_name: config.instance_name.to_string(), - log_dir: config.log_dir.clone().into(), - voice_dir: tts.voice_dir.clone().into(), - mutex: tokio::sync::Mutex::new(()), - }) - } - - pub async fn handle_socket( - &self, - socket: ws::WebSocket, - query: crate::TtsStreamingQuery, - ) -> Result<()> { - use futures_util::{SinkExt, StreamExt}; - - let _guard = self.mutex.lock().await; - let config = &self.tts_config; - let (log_tx, log_rx) = logger(); - let log_tx2 = log_tx.clone(); - let sampling = if query.temperature <= 0. || query.top_k <= 1 { - candle_transformers::generation::Sampling::ArgMax - } else { - candle_transformers::generation::Sampling::TopK { - k: query.top_k, - temperature: query.temperature, - } - }; - - let text_lp = candle_transformers::generation::LogitsProcessor::from_sampling( - query.seed, - sampling.clone(), - ); - let audio_lp = - candle_transformers::generation::LogitsProcessor::from_sampling(query.seed, sampling); - let conditions = match self.lm.condition_provider() { - None => None, - Some(cp) => { - let conditions = cp.condition_lut("control", "also_good")?; - tracing::info!(?conditions, "generated conditions"); - Some(conditions) - } - }; - - let mut last_text_token = config.text_start_token; - let ca_src = self.voice_ca_src(query.voice.as_ref(), query.voices.as_ref())?; - ca_src.device().synchronize()?; - let ca_src = if query.cfg_alpha.is_some() { - let lp = self.speaker_encoder.empty()?; - Tensor::cat(&[ca_src, lp], 0)? - } else { - ca_src - }; - let max_seq_len = query.max_seq_len.unwrap_or(2048); - let mut state = moshi::tts_streaming::State::new( - self.lm.clone(), - Some(moshi::transformer::CaSrc::Tokens(ca_src)), - max_seq_len, - audio_lp, - text_lp, - query.cfg_alpha, - config.clone(), - ); - let text_tokenizer = self.text_tokenizer.clone(); - - let (mut sender, mut receiver) = socket.split(); - let (in_tx, in_rx) = std::sync::mpsc::channel(); - let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel(); - let text_bos_token = state.config().text_bos_token; - let recv_loop = tokio::task::spawn(async move { - let mut inserted_bos = false; - while let Some(msg) = receiver.next().await { - let msg = match msg? { - ws::Message::Text(x) => x, - ws::Message::Binary(x) => { - // End of stream, we do not exit the loop so as not to close - // the connection. - if x.as_ref() == b"\0" { - log::info!("received end of stream"); - in_tx.send(None)?; - } - continue; - } - // ping messages are automatically answered by tokio-tungstenite as long as - // the connection is read from. - ws::Message::Ping(_) | ws::Message::Pong(_) => continue, - ws::Message::Close(_) => break, - }; - - let msg: String = msg.to_string(); - for word in msg.split(' ') { - if word.is_empty() { - continue; - } - let mut word_tokens: Vec<_> = - text_tokenizer.encode(word)?.into_iter().map(|v| v.id).collect(); - if !inserted_bos { - inserted_bos = true; - word_tokens.insert(0, text_bos_token) - } - log_tx2.send_text(word.to_string()); - in_tx.send(Some(word_tokens))?; - } - } - tracing::info!("recv loop exited - connection closed"); - Ok::<(), anyhow::Error>(()) - }); - let mut audio_tokenizer = self.audio_tokenizer.clone(); - audio_tokenizer.reset_state(); - let text_tokenizer = self.text_tokenizer.clone(); - let format = query.format; - let process_loop = tokio::task::spawn_blocking(move || { - let err = (|| { - tracing::info!("starting the inference loop"); - let text_audio_delay_in_tokens = state.config().text_audio_delay_in_tokens; - let acoustic_delay = state.config().acoustic_delay; - let text_eop_token = state.config().text_eop_token; - let text_pad_token = state.config().text_pad_token; - let extra_steps = state.config().extra_steps; - - let mut token_idx = 0; - let mut step_past_last_token = 0; - // Start with an empty list to trigger the first bos. - let mut word_tokens = Some(vec![]); - - let mut encoder = Encoder::new(format)?; - if let Some(header) = encoder.header()? { - out_tx.send(header)? - } - let mut last_epad_index = 0usize; - for step_idx in 0..max_seq_len { - let allowed_tokens = match word_tokens.as_ref() { - None => { - step_past_last_token += 1; - if step_past_last_token > extra_steps + text_audio_delay_in_tokens { - break; - } - moshi::tts_streaming::AllowedTokens::Pad - } - Some(word_tokens) => match word_tokens.get(token_idx) { - None => moshi::tts_streaming::AllowedTokens::PadOrEpad, - Some(id) => moshi::tts_streaming::AllowedTokens::Text(*id), - }, - }; - last_text_token = - state.step(last_text_token, allowed_tokens, conditions.as_ref())?; - if last_text_token == text_eop_token { - if let Some(vs) = word_tokens { - if let Ok(text) = text_tokenizer.decode_piece_ids(&vs) { - let start_s = last_epad_index as f64 / 12.5; - let stop_s = step_idx as f64 / 12.5; - let wwts = WordWithTimestamps { text, start_s, stop_s }; - if let Some(oo) = encoder.encode_word(wwts)? { - out_tx.send(oo)?; - } - } - } - last_epad_index = step_idx; - word_tokens = in_rx.recv()?; - if word_tokens.is_none() { - // We teacher force a pad instead of tho eop for the last word. - state.overwrite_last_text_token(text_pad_token)?; - } - token_idx = 0; - } else if last_text_token != text_pad_token { - token_idx += 1; - } - if let Some(audio_tokens) = state.last_audio_tokens() { - let cb = audio_tokens.len(); - let audio_tokens = - candle::Tensor::from_vec(audio_tokens, (1, cb, 1), state.device())?; - if step_idx >= text_audio_delay_in_tokens + acoustic_delay { - let pcm = audio_tokenizer - .decode_step(&audio_tokens.clone().into(), &().into())?; - if let Some(pcm) = pcm.as_option() { - let pcm = pcm.flatten_all()?.to_vec1::()?; - let oo = encoder.encode(pcm)?; - out_tx.send(oo)?; - } - } - log_tx.send_slice(last_text_token, audio_tokens) - } else { - let cb = state.audio_codebooks(); - let audio_tokens = - candle::Tensor::zeros((1, cb, 1), DType::U32, state.device())?; - log_tx.send_slice(last_text_token, audio_tokens) - } - } - std::thread::sleep(std::time::Duration::from_secs(1)); - Ok::<(), anyhow::Error>(()) - })(); - match err { - Err(err) => tracing::error!(?err, "process loop exited"), - Ok(()) => tracing::info!("process loop exited"), - } - }); - let send_loop = tokio::task::spawn(async move { - use tokio::time::{timeout, Duration}; - loop { - // The recv method is cancel-safe so can be wrapped in a timeout. - let msg = timeout(Duration::from_secs(10), out_rx.recv()).await; - let msg = match msg { - Ok(Some(msg)) => ws::Message::binary(msg), - Ok(None) => break, - Err(_) => ws::Message::Ping(vec![].into()), - }; - sender.send(msg).await?; - } - tracing::info!("send loop exited - connection closed"); - sender.close().await?; - tracing::info!("send loop exited - connection really closed"); - drop(sender); - Ok::<(), anyhow::Error>(()) - }); - // select should ensure that all the threads get aborted on timeout. - // TODO(laurent): this actually doesn't work as expected, and the background threads don't - // appear to be cancelled properly (at least the websocket connection remains open. - let sleep = tokio::time::sleep(std::time::Duration::from_secs(360)); - tokio::pin!(sleep); - tokio::select! { - _ = sleep => { - tracing::error!("reached timeout"); - } - res = recv_loop => { - match res { - Err(err) => tracing::error!(?err, "recv loop ended"), - Ok(Err(err)) => tracing::error!(?err, "recv loop err"), - Ok(Ok(())) => tracing::info!("recv loop ended"), - } - } - p = process_loop => { - match p { - Err(err) => tracing::error!(?err, "process loop ended"), - Ok(()) => tracing::info!("process loop ended"), - } - } - res = send_loop => { - match res { - Err(err) => tracing::error!(?err, "send loop ended"), - Ok(Err(err)) => tracing::error!(?err, "send loop err"), - Ok(Ok(())) => tracing::info!("send loop ended"), - } - } - } - tracing::info!("exiting handle-socket"); - if let Err(err) = log_rx.save(&query, &self.log_dir, &self.instance_name) { - tracing::error!(?err, "cannot save logs") - }; - Ok(()) - } - - pub fn voice_ca_src( - &self, - voice: Option<&String>, - voices: Option<&Vec>, - ) -> Result { - match (voice, voices) { - (None, None) => anyhow::bail!("either voice or voices has to be set"), - (Some(_), Some(_)) => { - anyhow::bail!("voice and voices should not be set at the same time") - } - (Some(voice), None) => match self.ca_srcs.get(voice) { - None => { - let voice_dir = std::fs::canonicalize(&self.voice_dir)?; - let mut pcms = vec![]; - let (voice, speaker_cond_start_s) = match voice.split_once('+') { - None => (voice.as_str(), 0.0), - Some((v, delay)) => { - let delay = match delay.parse::() { - Ok(delay) => delay, - Err(_) => anyhow::bail!( - "unexpected format for delay in {voice}: '{delay}'" - ), - }; - (v, delay) - } - }; - let path = std::fs::canonicalize(voice_dir.join(voice))?; - if !path.starts_with(&voice_dir) { - tracing::error!(?voice_dir, ?path, "unable to access voice file"); - anyhow::bail!("unknown voice file '{voice}'") - } - let pcm = speaker_pcm( - self.speaker_encoder.sample_rate(), - speaker_cond_start_s, - self.tts_config.speaker_cond_duration_s, - path, - self.lm.device(), - )?; - pcms.push(pcm.clone()); - pcms.push(pcm); - Ok(self.speaker_encoder.encode(&pcms)?) - } - Some(v) => Ok(v.clone()), - }, - (None, Some(voices)) => { - let voice_dir = std::fs::canonicalize(&self.voice_dir)?; - let mut pcms = vec![]; - for voice in voices.iter() { - let (voice, speaker_cond_start_s) = match voice.split_once('+') { - None => (voice.as_str(), 0.0), - Some((v, delay)) => { - let delay = match delay.parse::() { - Ok(delay) => delay, - Err(_) => anyhow::bail!( - "unexpected format for delay in {voice}: '{delay}'" - ), - }; - (v, delay) - } - }; - let path = std::fs::canonicalize(voice_dir.join(voice))?; - if !path.starts_with(&voice_dir) { - tracing::error!(?voice_dir, ?path, "unable to access voice file"); - anyhow::bail!("unknown voice file '{voice}'") - } - let pcm = speaker_pcm( - self.speaker_encoder.sample_rate(), - speaker_cond_start_s, - self.tts_config.speaker_cond_duration_s, - path, - self.lm.device(), - )?; - pcms.push(pcm) - } - Ok(self.speaker_encoder.encode(&pcms)?) - } - } - } - - pub fn run(&self, query: &crate::TtsQuery) -> Result<(Vec, Vec)> { - let config = &self.tts_config; - let text_audio_delay_in_tokens = config.text_audio_delay_in_tokens; - let text_bos_token = config.text_bos_token; - let text_eos_token = config.text_eos_token; - let text_eop_token = config.text_eop_token; - let text_pad_token = config.text_pad_token; - let mut prompt = moshi::tts_streaming::tokenize_prompt( - &query.text, - text_bos_token, - text_eos_token, - |s| self.text_tokenizer.encode(s).map(|v| v.into_iter().map(|v| v.id).collect()), - )?; - // Insert an empty word to start with and trigger the first bos. - prompt.insert(0, (vec![], Speaker::Other)); - tracing::info!(?prompt, "starting tts"); - let mut transcript = vec![]; - let (log_tx, log_rx) = logger(); - let all_audio_tokens = { - let start_time = std::time::Instant::now(); - let sampling = if query.temperature <= 0. || query.top_k <= 1 { - candle_transformers::generation::Sampling::ArgMax - } else { - candle_transformers::generation::Sampling::TopK { - k: query.top_k, - temperature: query.temperature, - } - }; - - let text_lp = candle_transformers::generation::LogitsProcessor::from_sampling( - query.seed, - sampling.clone(), - ); - let audio_lp = candle_transformers::generation::LogitsProcessor::from_sampling( - query.seed, sampling, - ); - let conditions = match self.lm.condition_provider() { - None => None, - Some(cp) => { - let conditions = cp.condition_lut("control", "also_good")?; - tracing::info!(?conditions, "generated conditions"); - Some(conditions) - } - }; - - let mut last_text_token = config.text_start_token; - let ca_src = self.voice_ca_src(query.voice.as_ref(), query.voices.as_ref())?; - let ca_src = if query.cfg_alpha.is_some() { - let lp = self.speaker_encoder.empty()?; - Tensor::cat(&[ca_src, lp], 0)? - } else { - ca_src - }; - let max_seq_len = query.max_seq_len.unwrap_or(2048); - let config = config.clone(); - let mut state = moshi::tts_streaming::State::new( - self.lm.clone(), - Some(moshi::transformer::CaSrc::Tokens(ca_src)), - max_seq_len, - audio_lp, - text_lp, - query.cfg_alpha, - config.clone(), - ); - let mut all_audio_tokens = vec![]; - tracing::info!("starting the inference loop"); - let mut word_idx = 0; - let mut token_idx = 0; - let mut step_past_last_token = 0; - let mut last_epad_index = 0usize; - for step_idx in 0..max_seq_len { - let word_tokens = prompt.get(word_idx); - let allowed_tokens = match word_tokens.as_ref() { - None => { - step_past_last_token += 1; - if step_past_last_token > 5 + text_audio_delay_in_tokens { - break; - } - moshi::tts_streaming::AllowedTokens::Pad - } - Some(word_tokens) => match word_tokens.0.get(token_idx) { - None => moshi::tts_streaming::AllowedTokens::PadOrEpad, - Some(id) => moshi::tts_streaming::AllowedTokens::Text(*id), - }, - }; - last_text_token = - state.step(last_text_token, allowed_tokens, conditions.as_ref())?; - if last_text_token == text_eop_token { - if let Some(vs) = word_tokens { - if let Ok(text) = self.text_tokenizer.decode_piece_ids(&vs.0) { - let start_s = last_epad_index as f64 / 12.5; - let stop_s = step_idx as f64 / 12.5; - transcript.push(WordWithTimestamps { text, start_s, stop_s }) - } - } - last_epad_index = step_idx; - word_idx += 1; - token_idx = 0; - } else if last_text_token != text_pad_token { - token_idx += 1; - } - if let Some(audio_tokens) = state.last_audio_tokens() { - let cb = audio_tokens.len(); - let audio_tokens = - candle::Tensor::from_vec(audio_tokens, (1, cb, 1), state.device())?; - if step_idx >= text_audio_delay_in_tokens { - all_audio_tokens.push(audio_tokens.clone()) - } - log_tx.send_slice(last_text_token, audio_tokens) - } else { - let cb = state.audio_codebooks(); - let audio_tokens = - candle::Tensor::zeros((1, cb, 1), DType::U32, state.device())?; - log_tx.send_slice(last_text_token, audio_tokens) - } - } - let dt = start_time.elapsed().as_secs_f64(); - let total = all_audio_tokens.len(); - tracing::info!( - "processed {total} total steps in {dt:.2}s, {:.2} steps/s", - total as f64 / dt - ); - Tensor::cat(&all_audio_tokens, candle::D::Minus1)? - }; - let (_one, _codebooks, total_steps) = all_audio_tokens.dims3()?; - let mut all_pcm_chunks = vec![]; - let chunk_by = 25; - let mut mimi = self.audio_tokenizer.clone(); - for start_step in (0..total_steps).step_by(chunk_by) { - let chunk_steps = usize::min(chunk_by, total_steps - start_step); - let pcm = mimi.decode_step( - &all_audio_tokens.narrow(2, start_step, chunk_steps)?.into(), - &().into(), - )?; - if let Some(pcm) = pcm.as_option() { - all_pcm_chunks.push(pcm.clone()) - } - } - // Close the log stream so that log_rx.save does not block. - std::mem::drop(log_tx); - if let Err(err) = log_rx.save(&query, &self.log_dir, &self.instance_name) { - tracing::error!(?err, "cannot save logs") - }; - - let pcm = Tensor::cat(&all_pcm_chunks, 2)?; - let pcm = pcm.i((0, 0))?.to_vec1::()?; - let mut wav = vec![]; - moshi::wav::write_pcm_as_wav(&mut wav, &pcm, 24_000)?; - Ok((wav, transcript)) - } -} - -pub fn speaker_pcm>( - mimi_sample_rate: f64, - speaker_cond_start_s: f64, - speaker_cond_duration_s: f64, - speaker: P, - dev: &Device, -) -> Result { - let (pcm, sample_rate) = kaudio::pcm_decode(speaker)?; - let pcm = if sample_rate != mimi_sample_rate as u32 { - kaudio::resample(&pcm, sample_rate as usize, mimi_sample_rate as usize)? - } else { - pcm - }; - let start_pos = (speaker_cond_start_s * mimi_sample_rate) as usize; - let sample_len = (speaker_cond_duration_s * mimi_sample_rate) as usize; - let pcm = &pcm[start_pos..start_pos + sample_len]; - let pcm = Tensor::new(pcm, dev)?.reshape((1, 1, ()))?; - Ok(pcm) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::Result; - -#[derive(Debug, PartialEq, Clone, serde::Deserialize, serde::Serialize)] -pub struct BuildInfo { - build_timestamp: String, - build_date: String, - git_branch: String, - git_timestamp: String, - git_date: String, - git_hash: String, - git_describe: String, - rustc_host_triple: String, - rustc_version: String, - cargo_target_triple: String, -} - -impl BuildInfo { - pub fn new() -> BuildInfo { - BuildInfo { - build_timestamp: String::from(env!("VERGEN_BUILD_TIMESTAMP")), - build_date: String::from(env!("VERGEN_BUILD_DATE")), - git_branch: String::from(env!("VERGEN_GIT_BRANCH")), - git_timestamp: String::from(env!("VERGEN_GIT_COMMIT_TIMESTAMP")), - git_date: String::from(env!("VERGEN_GIT_COMMIT_DATE")), - git_hash: String::from(env!("VERGEN_GIT_SHA")), - git_describe: String::from(env!("VERGEN_GIT_DESCRIBE")), - rustc_host_triple: String::from(env!("VERGEN_RUSTC_HOST_TRIPLE")), - rustc_version: String::from(env!("VERGEN_RUSTC_SEMVER")), - cargo_target_triple: String::from(env!("VERGEN_CARGO_TARGET_TRIPLE")), - } - } -} - -pub fn replace_env_vars(input: &str) -> String { - let re = regex::Regex::new(r"\$([A-Za-z_][A-Za-z0-9_]*)").unwrap(); - re.replace_all(input, |caps: ®ex::Captures| { - let var_name = &caps[1]; - std::env::var(var_name).unwrap_or_else(|_| "".to_string()) - }) - .to_string() -} - -pub fn resolve_or_download(input: &str) -> Result { - let path = match input.strip_prefix("hf://") { - None => replace_env_vars(input), - Some(path) => { - let s: Vec<&str> = path.split('/').collect(); - if s.len() < 3 { - anyhow::bail!("unexpected format for hf path {input}") - } - let repo = format!("{}/{}", s[0], s[1]); - let file = s[2..].join("/"); - let api = hf_hub::api::sync::Api::new()?.model(repo); - api.get(&file)?.to_string_lossy().to_string() - } - }; - Ok(path) -} - -fn walk_toml(t: &mut toml::Value, f: &impl Fn(&mut String) -> Result<()>) -> Result<()> { - match t { - toml::Value::Table(t) => { - for (_, t) in t.iter_mut() { - walk_toml(t, f)?; - } - } - toml::Value::Array(a) => { - for t in a.iter_mut() { - walk_toml(t, f)? - } - } - toml::Value::String(s) => f(s)?, - toml::Value::Integer(_) - | toml::Value::Float(_) - | toml::Value::Boolean(_) - | toml::Value::Datetime(_) => {} - } - Ok(()) -} - -pub fn resolve_or_download_toml(t: &mut toml::Table) -> Result<()> { - for (_, t) in t.iter_mut() { - walk_toml(t, &|s: &mut String| -> Result<()> { - *s = resolve_or_download(s)?; - Ok(()) - })?; - } - Ok(()) -} - -pub struct WrapJson(pub Result); - -impl axum::response::IntoResponse for WrapJson { - fn into_response(self) -> axum::response::Response { - match self.0 { - Ok(v) => axum::Json(v).into_response(), - Err(err) => { - tracing::error!(?err, "returning internal server error 500"); - (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("{err}")).into_response() - } - } - } -} - -pub struct AxumError(anyhow::Error); - -impl axum::response::IntoResponse for AxumError { - fn into_response(self) -> axum::response::Response { - let err = self.0; - tracing::error!(?err); - (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("{err:?}")).into_response() - } -} - -impl> From for AxumError { - fn from(value: E) -> Self { - Self(value.into()) - } -} - -pub type AxumResult = std::result::Result; - -fn conv(samples: &mut Vec, data: std::borrow::Cow>) -where - T: symphonia::core::sample::Sample, - f32: symphonia::core::conv::FromSample, -{ - use symphonia::core::audio::Signal; - use symphonia::core::conv::FromSample; - samples.extend(data.chan(0).iter().map(|v| f32::from_sample(*v))) -} - -pub fn pcm_decode(bytes: axum::body::Bytes) -> anyhow::Result<(Vec, u32)> { - use symphonia::core::audio::{AudioBufferRef, Signal}; - - let source = std::io::Cursor::new(bytes); - let mss = symphonia::core::io::MediaSourceStream::new(Box::new(source), Default::default()); - let hint = symphonia::core::probe::Hint::new(); - let meta_opts: symphonia::core::meta::MetadataOptions = Default::default(); - let fmt_opts: symphonia::core::formats::FormatOptions = Default::default(); - let probed = symphonia::default::get_probe().format(&hint, mss, &fmt_opts, &meta_opts)?; - let mut format = probed.format; - let track = format - .tracks() - .iter() - .find(|t| t.codec_params.codec != symphonia::core::codecs::CODEC_TYPE_NULL) - .expect("no supported audio tracks"); - let mut decoder = symphonia::default::get_codecs() - .make(&track.codec_params, &Default::default()) - .expect("unsupported codec"); - let track_id = track.id; - let sample_rate = track.codec_params.sample_rate.unwrap_or(0); - let mut pcm_data = Vec::new(); - while let Ok(packet) = format.next_packet() { - while !format.metadata().is_latest() { - format.metadata().pop(); - } - if packet.track_id() != track_id { - continue; - } - match decoder.decode(&packet)? { - AudioBufferRef::F32(buf) => pcm_data.extend(buf.chan(0)), - AudioBufferRef::U8(data) => conv(&mut pcm_data, data), - AudioBufferRef::U16(data) => conv(&mut pcm_data, data), - AudioBufferRef::U24(data) => conv(&mut pcm_data, data), - AudioBufferRef::U32(data) => conv(&mut pcm_data, data), - AudioBufferRef::S8(data) => conv(&mut pcm_data, data), - AudioBufferRef::S16(data) => conv(&mut pcm_data, data), - AudioBufferRef::S24(data) => conv(&mut pcm_data, data), - AudioBufferRef::S32(data) => conv(&mut pcm_data, data), - AudioBufferRef::F64(data) => conv(&mut pcm_data, data), - } - } - Ok((pcm_data, sample_rate)) -} - - - -// Copyright (c) Kyutai, all rights reserved. -// This source code is licensed under the license found in the -// LICENSE file in the root directory of this source tree. - -use anyhow::Result; -use vergen::EmitBuilder; - -pub fn main() -> Result<()> { - // NOTE: This will output everything, and requires all features enabled. - // NOTE: See the EmitBuilder documentation for configuration options. - EmitBuilder::builder().all_build().all_cargo().all_git().all_rustc().all_sysinfo().emit()?; - Ok(()) -} - - - -[package] -name = "moshi-server" -version.workspace = true -edition.workspace = true -description.workspace = true -repository.workspace = true -keywords.workspace = true -categories.workspace = true -license.workspace = true - -[dependencies] -anyhow = { workspace = true } -axum = { workspace = true } -base64 = { workspace = true } -bincode = { workspace = true } -byteorder = { workspace = true } -candle = { workspace = true } -candle-nn = { workspace = true } -candle-transformers = { workspace = true } -clap = { workspace = true } -futures-util = { workspace = true } -hf-hub = { workspace = true } -kaudio = { workspace = true } -lazy_static = { workspace = true } -log = { workspace = true } -moshi = { workspace = true } -numpy = { workspace = true } -ogg = { workspace = true } -opus = { workspace = true } -prometheus = { workspace = true } -pyo3 = { workspace = true } -pyo3-ffi = { workspace = true } -rayon = { workspace = true } -regex = { workspace = true } -rmp-serde = { workspace = true } -rubato = { workspace = true } -sentencepiece = { workspace = true } -serde = { workspace = true } -serde_json = { workspace = true } -symphonia = { workspace = true } -tokio = { workspace = true } -toml = { workspace = true } -tower = { workspace = true } -tower-http = { workspace = true } -tracing = { workspace = true } -tracing-appender = { workspace = true } -tracing-subscriber = { workspace = true } - -[build-dependencies] -anyhow = { workspace = true } -vergen = { workspace = true } - -[features] -default = [] -cuda = ["moshi/cuda", "candle/cuda", "candle-nn/cuda", "candle-transformers/cuda"] -metal = ["moshi/metal", "candle/metal", "candle-nn/metal", "candle-transformers/metal"] - - - -[project] -name = "tts-python" -version = "0.1.0" -readme = "README.md" -requires-python = "==3.12.8" -dependencies = [ - "moshi==0.2.7", - "setuptools", - "xformers", - "pydantic", - "julius", - "torchaudio", -] - - - -# Copyright (c) Kyutai, all rights reserved. - -import argparse -from dataclasses import dataclass, field -from enum import Enum -import huggingface_hub -from pathlib import Path -import random -import time -import typing as tp - -import numpy as np -from safetensors.torch import load_file -import torch - -from moshi.conditioners import ConditionAttributes, dropout_all_conditions, TensorCondition -from moshi.models import loaders -from moshi.models.lm import _LMGenState, LMGen -from moshi.models.tts import TTSModel, Entry, State, StateMachine, DEFAULT_DSM_TTS_REPO -from moshi.modules.transformer import StreamingMultiheadAttention -from pydantic import BaseModel - - -class MaskFlags(Enum): - # Output PCM is ready - HAS_PCM = 1 - # Generation is done, no need to step again. - IS_EOS = 2 - # One word was consumed in the text stream. - WORD_FINISHED = 4 - # One AR step was performed. - AR_STEP = 8 - # AR step was skipped because the client is not sending words fast enough. - MISSING_WORDS = 16 - - -def flags_out_from_mask_(flags_out: np.ndarray, mask: torch.Tensor, value: int): - flags_out[mask.numpy()] |= value - - -def split_at_specific_separator(text: str, separator: str, index_of_separator: int) -> tuple[str, str]: - """ kyutai/tts-voices/unmute-prod-website/*.safetensors - becomes - ('kyutai/tts-voices', 'unmute-prod-website/*.safetensors) - with index_of_separator=1. - """ - if text.count(separator) <= index_of_separator: - raise ValueError(f"Separator '{separator}' not found {index_of_separator + 1} times in `{text}`.") - parts = text.split(separator, index_of_separator + 1) - return separator.join(parts[:-1]), parts[-1] - - -class Config(BaseModel): - log_folder: Path = Path.home() / 'tmp/tts-service' - hf_repo: str = DEFAULT_DSM_TTS_REPO - mimi_weight: Path | None = None - moshi_weight: Path | None = None - config_path: Path | None = None - tokenizer: Path | None = None - device: str = 'cuda' - - n_q: int = 24 - # This can have multiple formats: - # - A path to a folder with voices, e.g. `models/tts` - # - A huggingface snapshot, e.g. `hf-snapshot://kyutai/tts-voices` - # - A huggingface snapshot with a pattern, - # e.g. `hf-snapshot://kyutai/tts-voices/unmute-prod-website/*.safetensors` - voice_folder: str = str(Path.home() / 'models/tts-voices') - default_voice: str = "barack_demo.wav" - - temp: float = 0.6 - cfg_coef: float = 2. - - max_padding: int = 8 - initial_padding: int = 2 - final_padding: int = 4 - padding_between: int = 1 - - interleaved_text_only: int = 2 - debug: bool = False - - -def init(batch_size: int, config_override: dict) -> 'TTSService': - config = Config(**config_override) - config.log_folder.mkdir(parents=True, exist_ok=True) - - print("retrieving checkpoint") - checkpoint_info = loaders.CheckpointInfo.from_hf_repo( - config.hf_repo, moshi_weights=config.moshi_weight, mimi_weights=config.mimi_weight, - config_path=config.config_path, tokenizer=config.tokenizer) - - cfg_condition = None - tts_model = TTSModel.from_checkpoint_info( - checkpoint_info, n_q=config.n_q, temp=config.temp, cfg_coef=config.cfg_coef, - max_padding=config.max_padding, initial_padding=config.initial_padding, final_padding=config.final_padding, - device=config.device) - if tts_model.valid_cfg_conditionings: - # Model was trained with CFG distillation. - cfg_condition = tts_model.cfg_coef - tts_model.cfg_coef = 1. - cfg_is_no_text = False - else: - cfg_is_no_text = True - - voice_suffix = tts_model.voice_suffix - print(f"loading voices from {config.voice_folder}, with suffix {voice_suffix}.") - all_attributes = {} - voice_folder = config.voice_folder - if voice_folder.startswith("hf-snapshot://"): - voice_folder = voice_folder.removeprefix("hf-snapshot://") - # We detect if there is a pattern in the voice folder. - if voice_folder.count("/") > 1: - voice_folder, pattern = split_at_specific_separator(voice_folder, '/', 1) - else: - pattern = None - print(f"retrieving voices from {voice_folder}") - voice_folder = huggingface_hub.snapshot_download(voice_folder, allow_patterns=pattern) - voice_folder = Path(voice_folder) - - for file in voice_folder.glob(f'**/*{voice_suffix}'): - relative = file.relative_to(voice_folder) - name = str(relative.with_name(relative.name.removesuffix(voice_suffix))) - try: - attributes = tts_model.make_condition_attributes([file, file], cfg_coef=cfg_condition) - except Exception: - print(f"[WARNING] failed to load voice {name}") - else: - all_attributes[name] = attributes - - if not all_attributes: - raise RuntimeError( - "No voices found, please check your voice folder. " - f"Searched for files matching {voice_folder}/**/*{voice_suffix}" - ) - - if config.default_voice not in all_attributes: - raise RuntimeError( - f"Default voice {config.default_voice}, please check your voice folder. " - f"Expected {voice_folder}/{config.default_voice}{voice_suffix} to exist" - ) - - service = TTSService( - batch_size=batch_size, default_attribute_name=config.default_voice, - all_attributes=all_attributes, - tts_model=tts_model, - cfg_condition=cfg_condition, - cfg_is_no_text=cfg_is_no_text, - padding_between=config.padding_between, - debug=config.debug, - interleaved_text_only=config.interleaved_text_only) - - return service - - -@dataclass -class ClientState: - is_complete: bool = False - state: State | None = None - offset: int = 0 - - def reset(self, state_machine: StateMachine) -> None: - self.is_complete = False - self.offset = 0 - self.state = state_machine.new_state([]) - - -@dataclass -class TTSService: - batch_size: int - default_attribute_name: str - all_attributes: dict[str, ConditionAttributes] - - tts_model: TTSModel - - cfg_is_no_text: bool = True - cfg_condition: float | None = None - padding_between: int = 1 - n_q: int = 32 - debug: bool = False - interleaved_text_only: int = 0 - - flags_out: np.ndarray | None = None - clients: list[ClientState] = field(default_factory=list) - cross_attention_cache: dict[str, torch.Tensor] = field(default_factory=dict) - cross_attentions: list[StreamingMultiheadAttention] = field(default_factory=list) - - def __post_init__(self): - lm = self.tts_model.lm - tts_model = self.tts_model - mimi = self.tts_model.mimi - machine = self.tts_model.machine - - self.device = lm.device - self.dtype = lm.dtype - self.remaining_text_only = self.interleaved_text_only - - for _ in range(self.batch_size): - client = ClientState() - self.clients.append(client) - - print("Filling cross attention cache.") - for name, attributes in self.all_attributes.items(): - self.cross_attention_cache[name] = self._get_cross_attention_source([attributes]) - - assert lm.condition_provider is not None - - cas = [self.all_attributes[self.default_attribute_name]] * self.batch_size - if self.tts_model.cfg_coef != 1.0: - nulled = make_null(cas) - cas = cas + nulled - prepared = lm.condition_provider.prepare(cas) - condition_tensors = lm.condition_provider(prepared) - - for module in lm.modules(): - if isinstance(module, StreamingMultiheadAttention) and module.cross_attention: - self.cross_attentions.append(module) - - self.lm_gen = LMGen( - lm, temp=tts_model.temp, temp_text=tts_model.temp, cfg_coef=tts_model.cfg_coef, - condition_tensors=condition_tensors, on_text_hook=self._on_text_hook, - on_audio_hook=self._on_audio_hook, cfg_is_no_text=self.cfg_is_no_text, - support_out_of_sync=True) - self.lm_gen.streaming_forever(self.batch_size) - mimi.streaming_forever(self.batch_size) - - missing = lm.n_q - lm.dep_q - self.input_tokens = torch.full( - (self.batch_size, missing, 1), machine.token_ids.zero, - dtype=torch.long, device=self.device) - self.no_depformer_tokens = torch.full( - (self.batch_size, lm.dep_q, 1), machine.token_ids.zero, - dtype=torch.long, device=self.device) - self.last_actives: list[bool] = [False] * self.batch_size - print("warming up.") - for _ in range(3): - mimi.set_exec_mask(torch.ones(self.batch_size, dtype=torch.bool)) - self.lm_gen.set_exec_mask(torch.ones(self.batch_size, dtype=torch.bool)) - frame = self.lm_gen.step(self.input_tokens) - assert frame is not None - mimi.decode(frame[:, 1:].clamp(min=0)) - print("ready to roll.") - - def _get_cross_attention_source(self, all_attributes: list[ConditionAttributes]) -> torch.Tensor: - lm = self.tts_model.lm - assert lm.condition_provider is not None - assert lm.fuser is not None - prepared = lm.condition_provider.prepare(all_attributes) - condition_tensors = lm.condition_provider(prepared) - cross = lm.fuser.get_cross(condition_tensors) - assert cross is not None - return cross.to(device=self.device, dtype=self.dtype) - - @property - def _lm_gen_state(self) -> _LMGenState: - assert self.lm_gen._streaming_state is not None - return self.lm_gen._streaming_state - - def _on_audio_hook(self, audio_tokens: torch.Tensor) -> None: - delays = self.lm_gen.delays_cuda[1: 1 + self.tts_model.lm.dep_q] - mask = self._lm_gen_state.offsets[:, None] < delays + self.tts_model.delay_steps - audio_tokens.masked_fill_(mask, self.tts_model.machine.token_ids.zero) - - def _on_text_hook(self, text_tokens) -> None: - tokens = text_tokens.tolist() - out_tokens = [] - for b, (token, client) in enumerate(zip(tokens, self.clients)): - if not self.last_actives[b]: - out_tokens.append(token) - continue - assert client.state is not None - out_token, consumed_new_word = self.tts_model.machine.process(client.offset, client.state, token) - - if self.flags_out is not None and consumed_new_word: - self.flags_out[b] |= MaskFlags.WORD_FINISHED.value - out_tokens.append(out_token) - text_tokens[:] = torch.tensor(out_tokens, dtype=torch.long, device=text_tokens.device) - - def _print(self, *args, **kwargs): - if self.debug: - print(*args, **kwargs) - - @torch.no_grad() - def step(self, updates: list[tuple[int, list[int], np.ndarray | str | None]], pcm_out: np.ndarray, - flags_out: np.ndarray, code_out: np.ndarray) -> None: - mimi = self.tts_model.mimi - machine = self.tts_model.machine - delay_steps = self.tts_model.delay_steps - - self.flags_out = flags_out - flags_out[:] = 0 - - reset_mask = torch.zeros(self.batch_size, dtype=torch.bool) - # List of pre computed cross attention values. - new_cross_sources: list[torch.Tensor] = [] - new_cross_indexes: list[int] = [] - # List of new dynamic conditioning that we need to compute. - new_voice_indexes: list[int] = [] - new_voice_sources: list[torch.Tensor] = [] - for b, new_entry, voice in updates: - client = self.clients[b] - if not new_entry: - self._print(f"[{b}] NO TOKENS REALLY LAURENT.") - if new_entry[0] == -1: - client.reset(machine) - reset_mask[b] = True - new_entry = new_entry[1:] - if isinstance(voice, np.ndarray): - new_voice_indexes.append(b) - new_voice_sources.append(torch.from_numpy(voice)) - else: - cross_source = self.cross_attention_cache.get(voice or '', None) - if cross_source is None: - cross_source = self.cross_attention_cache[self.default_attribute_name] - new_cross_sources.append(cross_source) - new_cross_indexes.append(b) - self._print(f"[{b}] Reset, voice is {voice}.") - if client.state is None: - self._print(f"[{b}] Trying to push {new_entry}, but not assigned.") - elif new_entry == [-2]: - self._print(f"[{b}] Done.") - client.is_complete = True - else: - self._print(f"[{b}] Pushing {new_entry}.") - padding = 0 - if self.padding_between > 0: - padding = max(0, self.padding_between + len(new_entry) - 1) - client.state.entries.append(Entry(new_entry, '', padding=padding)) - - actives = [] - mimi_actives = [] - in_text_onlys = [] - for b, client in enumerate(self.clients): - if client.state is None: - # client is not currently assigned. - active = False - elif client.is_complete: - # We got all the words from the client and are wrapping up. - active = True - elif client.state.forced_padding > 0: - # We are sure we won't try to consume a word at this point. - active = True - elif len(client.state.entries) > self.tts_model.machine.second_stream_ahead: - # We have some words ready to be consumed. - active = True - else: - flags_out[b] |= MaskFlags.MISSING_WORDS.value - active = False - actives.append(active) - - real_offset = client.offset - self.lm_gen.max_delay - - mimi_active = active and (real_offset >= delay_steps) - mimi_actives.append(mimi_active) - - in_text_only = active and (client.offset < delay_steps) - in_text_onlys.append(in_text_only) - - in_text_only_mask = torch.tensor(in_text_onlys, dtype=torch.bool) - run_in_text_only = self.remaining_text_only > 0 and in_text_only_mask.any() - - if run_in_text_only: - self.remaining_text_only -= 1 - mimi_exec_mask = torch.zeros(self.batch_size, dtype=torch.bool) - exec_mask = in_text_only_mask - actives = in_text_onlys - else: - self.remaining_text_only = self.interleaved_text_only - exec_mask = torch.tensor(actives, dtype=torch.bool) - mimi_exec_mask = torch.tensor(mimi_actives, dtype=torch.bool) - del mimi_actives - self.last_actives = actives - - flags_out_from_mask_(flags_out, exec_mask, MaskFlags.AR_STEP.value) - flags_out_from_mask_(flags_out, mimi_exec_mask, MaskFlags.HAS_PCM.value) - - # We check on exec_mask whether we actually need to run anything, before we move it to CUDA. - # However, we still need to perform the reset and update of cross attention for models - # with a text lookahead stream. - skip_exec = not exec_mask.any() - - exec_mask = exec_mask.to(self.device) - mimi_exec_mask = mimi_exec_mask.to(self.device) - need_reset = reset_mask.any() - reset_mask = reset_mask.to(self.device) - - if new_voice_sources: - all_attributes = [make_condition_attributes([voice_source], cfg_condition=self.cfg_condition) - for voice_source in new_voice_sources] - new_cross_sources += self._get_cross_attention_source(all_attributes).split(1) - new_cross_indexes += new_voice_indexes - if new_cross_sources: - cross_source = torch.cat(new_cross_sources) - cross_indexes = torch.tensor(new_cross_indexes, dtype=torch.long, device=self.device) - for attention in self.cross_attentions: - k, v = attention._compute_cross_attention(cross_source, cross_source) - state = attention._streaming_state - assert state is not None - assert state.k_cross is not None - assert state.v_cross is not None - state.k_cross.index_copy_(0, cross_indexes, k) - state.v_cross.index_copy_(0, cross_indexes, v) - - if need_reset: - self.lm_gen.reset_streaming(reset_mask=reset_mask) - mimi.reset_streaming(reset_mask=reset_mask) - - if skip_exec: - return - - self.lm_gen.set_exec_mask(exec_mask) - mimi.set_exec_mask(mimi_exec_mask) - - depformer_replace_tokens = self.no_depformer_tokens if run_in_text_only else None - frame = self.lm_gen.step(self.input_tokens, depformer_replace_tokens=depformer_replace_tokens) - assert frame is not None - audio_frame = frame[:, 1:] - audio_frame.clamp_(min=0) - - if run_in_text_only: - pcm = None - else: - pcm = mimi.decode(audio_frame) - pcm.clamp_(-0.99, 0.99) - - for b, client in enumerate(self.clients): - if actives[b]: - assert client.state is not None - client.offset += 1 - self._print(f"[{b}] Offset {client.offset: 3d}, pendings={len(client.state.entries): 3d}.") - if client.is_complete and client.state.end_step is not None: - # We were waiting for the end of the generation. - real_end = ( - client.state.end_step + delay_steps + self.tts_model.final_padding + self.lm_gen.max_delay) - if client.offset >= real_end: - self._print(f"[{b}] Done.") - client.reset(machine) - flags_out[b] |= MaskFlags.IS_EOS.value - if pcm is not None: - pcm_out[:] = pcm[:, 0].cpu().numpy() - code_out[:, :frame.shape[1]] = frame[:, :, 0].int().cpu().numpy() - code_out[:, frame.shape[1]:] = 0 - self.flags_out = None - - -class Profiler: - """Context manager wrapper for xformers profiler. - """ - def __init__(self, enabled: bool = False): - self.profiler: tp.Optional[tp.Any] = None - if enabled: - from xformers.profiler import profile - from xformers.profiler.api import PyTorchProfiler - output_dir = './profiler_data' - schedule = ( - (PyTorchProfiler, 6, 12), - ) - self.profiler = profile(output_dir=output_dir, schedule=schedule) - - def step(self): - if self.profiler is not None: - self.profiler.step() # type: ignore - - def __enter__(self): - if self.profiler is not None: - return self.profiler.__enter__() # type: ignore - - def __exit__(self, exc_type, exc_value, exc_tb): - if self.profiler is not None: - return self.profiler.__exit__(exc_type, exc_value, exc_tb) # type: ignore - - -def make_condition_attributes(voices: list[Path | torch.Tensor], - max_speakers: int = 5, - cfg_condition: float | None = None) -> ConditionAttributes: - assert voices - voice_tensor = None - mask = None - for idx in range(5): - if idx < len(voices): - voice = voices[idx] - if isinstance(voice, Path): - emb = load_file(voice, device='cuda')['speaker_wavs'] - else: - emb = voice - assert emb.dim() == 3 - if voice_tensor is None: - voice_tensor = torch.zeros(1, max_speakers, emb.shape[2], emb.shape[1], device='cuda') - if mask is None: - mask = torch.zeros(1, max_speakers, emb.shape[2], dtype=torch.bool, device='cuda') - voice_tensor[:, idx, :, :] = emb.transpose(1, 2) - mask[:, idx, :] = True - assert voice_tensor is not None - assert mask is not None - voice_tensor = voice_tensor.view(1, -1, voice_tensor.shape[-1]) - mask = mask.view(1, -1) - tensors = { - 'speaker_wavs': TensorCondition(voice_tensor, mask) - } - text: dict[str, str | None] = { - 'control': 'ok', - } - if cfg_condition is None: - text['cfg'] = None - else: - text['cfg'] = format(cfg_condition, '.1f') - return ConditionAttributes(text=dict(text), tensor=tensors) - - -def make_null(all_attributes: tp.Sequence[ConditionAttributes]) -> list[ConditionAttributes]: - return dropout_all_conditions(all_attributes) - - -if __name__ == '__main__': - rng = random.Random(1234) - parser = argparse.ArgumentParser() - parser.add_argument('-p', '--profile', action='store_true') - parser.add_argument('-b', '--batch_size', default=32, type=int) - args = parser.parse_args() - bs = args.batch_size - service = init(batch_size=bs, config_override={}) - print("Service initialized") - pcm_out = np.zeros((bs, 1920)) - flags_out = np.zeros(bs, dtype=np.int32) - code_out = np.zeros((bs, 33), dtype=np.int32) - service.step([(0, [-1], '')], pcm_out=pcm_out, flags_out=flags_out, code_out=code_out) - profiler = Profiler(enabled=args.profile) - with profiler: - for _ in range(100): - inp = [] - if rng.random() < 0.1: - word = [13, 34] - inp.append((0, word, None)) - be = time.time() - service.step(inp, pcm_out=pcm_out, flags_out=flags_out, code_out=code_out) - el = time.time() - be - print(f"FR {el * 1000:.1f}ms") - profiler.step() - - - -# Demucs code is coming from -# Copyright (c) Facebook, Inc. and its affiliates. -# All rights reserved. -# - - -from pathlib import Path -from moshi.models import loaders -import numpy as np -from pydantic import BaseModel - - -import math -import time - -import julius -import torch -from torch import nn -from torch.nn import functional as F -import torchaudio.transforms - - -def normalize_loudness( - wav: torch.Tensor, - sample_rate: int, - # The headroom is set more conservatively than in tts_make_voice.py, which is what - # we use for pre-made recordings. - loudness_headroom_db: float = 18, - energy_floor: float = 2e-3 -): - """Normalize an input signal to a user loudness in dB LKFS. - Audio loudness is defined according to the ITU-R BS.1770-4 recommendation. - - Args: - wav (torch.Tensor): Input multichannel audio data. - sample_rate (int): Sample rate. - loudness_headroom_db (float): Target loudness of the output in dB LUFS. - energy_floor (float): anything below that RMS level will not be rescaled. - Returns: - torch.Tensor: Loudness normalized output data. - """ - wav = wav - wav.mean(dim=-1, keepdim=True) - energy = wav.std() - if energy < energy_floor: - # Feeding audio lower than that will fail. - return wav - transform = torchaudio.transforms.Loudness(sample_rate) - try: - input_loudness_db = transform(wav).item() - except RuntimeError: - # audio is too short. - return wav - # calculate the gain needed to scale to the desired loudness level - delta_loudness = -loudness_headroom_db - input_loudness_db - gain = 10.0 ** (delta_loudness / 20.0) - output = gain * wav - assert output.isfinite().all(), (input_loudness_db, wav.pow(2).mean().sqrt()) - return output - - -def sinc(t: torch.Tensor) -> torch.Tensor: - """sinc. - - :param t: the input tensor - """ - return torch.where(t == 0, torch.ones(1, device=t.device, dtype=t.dtype), torch.sin(t) / t) - - -def kernel_upsample2(zeros=56, device=None): - """kernel_upsample2. - - """ - win = torch.hann_window(4 * zeros + 1, periodic=False, device=device) - winodd = win[1::2] - t = torch.linspace(-zeros + 0.5, zeros - 0.5, 2 * zeros, device=device) - t *= math.pi - kernel = (sinc(t) * winodd).view(1, 1, -1) - return kernel - - -def upsample2(x, zeros=56): - """ - Upsampling the input by 2 using sinc interpolation. - Smith, Julius, and Phil Gossett. "A flexible sampling-rate conversion method." - ICASSP'84. IEEE International Conference on Acoustics, Speech, and Signal Processing. - Vol. 9. IEEE, 1984. - """ - *other, time = x.shape - kernel = kernel_upsample2(zeros, x.device).to(x) - out = F.conv1d(x.view(-1, 1, time), kernel, padding=zeros)[..., 1:].view(*other, time) - y = torch.stack([x, out], dim=-1) - return y.view(*other, -1) - - -def kernel_downsample2(zeros=56, device=None): - """kernel_downsample2. - - """ - win = torch.hann_window(4 * zeros + 1, periodic=False, device=device) - winodd = win[1::2] - t = torch.linspace(-zeros + 0.5, zeros - 0.5, 2 * zeros, device=device) - t.mul_(math.pi) - kernel = (sinc(t) * winodd).view(1, 1, -1) - return kernel - - -def downsample2(x, zeros=56): - """ - Downsampling the input by 2 using sinc interpolation. - Smith, Julius, and Phil Gossett. "A flexible sampling-rate conversion method." - ICASSP'84. IEEE International Conference on Acoustics, Speech, and Signal Processing. - Vol. 9. IEEE, 1984. - """ - if x.shape[-1] % 2 != 0: - x = F.pad(x, (0, 1)) - xeven = x[..., ::2] - xodd = x[..., 1::2] - *other, time = xodd.shape - kernel = kernel_downsample2(zeros, x.device).to(x) - out = xeven + F.conv1d(xodd.view(-1, 1, time), kernel, padding=zeros)[..., :-1].view( - *other, time) - return out.view(*other, -1).mul(0.5) - - -class BLSTM(nn.Module): - def __init__(self, dim, layers=2, bi=True): - super().__init__() - klass = nn.LSTM - self.lstm = klass(bidirectional=bi, num_layers=layers, hidden_size=dim, input_size=dim) - self.linear = None - if bi: - self.linear = nn.Linear(2 * dim, dim) - - def forward(self, x, hidden=None): - x, hidden = self.lstm(x, hidden) - if self.linear: - x = self.linear(x) - return x, hidden - - -def rescale_conv(conv, reference): - std = conv.weight.std().detach() - scale = (std / reference)**0.5 - conv.weight.data /= scale - if conv.bias is not None: - conv.bias.data /= scale - - -def rescale_module(module, reference): - for sub in module.modules(): - if isinstance(sub, (nn.Conv1d, nn.ConvTranspose1d)): - rescale_conv(sub, reference) - - -class Demucs(nn.Module): - """ - Demucs speech enhancement model. - Args: - - chin (int): number of input channels. - - chout (int): number of output channels. - - hidden (int): number of initial hidden channels. - - depth (int): number of layers. - - kernel_size (int): kernel size for each layer. - - stride (int): stride for each layer. - - causal (bool): if false, uses BiLSTM instead of LSTM. - - resample (int): amount of resampling to apply to the input/output. - Can be one of 1, 2 or 4. - - growth (float): number of channels is multiplied by this for every layer. - - max_hidden (int): maximum number of channels. Can be useful to - control the size/speed of the model. - - normalize (bool): if true, normalize the input. - - glu (bool): if true uses GLU instead of ReLU in 1x1 convolutions. - - rescale (float): controls custom weight initialization. - See https://arxiv.org/abs/1911.13254. - - floor (float): stability flooring when normalizing. - - sample_rate (float): sample_rate used for training the model. - - """ - def __init__(self, - chin=1, - chout=1, - hidden=48, - depth=5, - kernel_size=8, - stride=4, - causal=True, - resample=4, - growth=2, - max_hidden=10_000, - normalize=True, - glu=True, - rescale=0.1, - floor=1e-3, - sample_rate=16_000): - - super().__init__() - if resample not in [1, 2, 4]: - raise ValueError("Resample should be 1, 2 or 4.") - - self.chin = chin - self.chout = chout - self.hidden = hidden - self.depth = depth - self.kernel_size = kernel_size - self.stride = stride - self.causal = causal - self.floor = floor - self.resample = resample - self.normalize = normalize - self.sample_rate = sample_rate - - self.encoder = nn.ModuleList() - self.decoder = nn.ModuleList() - activation = nn.GLU(1) if glu else nn.ReLU() - ch_scale = 2 if glu else 1 - - for index in range(depth): - encode = [] - encode += [ - nn.Conv1d(chin, hidden, kernel_size, stride), - nn.ReLU(), - nn.Conv1d(hidden, hidden * ch_scale, 1), activation, - ] - self.encoder.append(nn.Sequential(*encode)) - - decode = [] - decode += [ - nn.Conv1d(hidden, ch_scale * hidden, 1), activation, - nn.ConvTranspose1d(hidden, chout, kernel_size, stride), - ] - if index > 0: - decode.append(nn.ReLU()) - self.decoder.insert(0, nn.Sequential(*decode)) - chout = hidden - chin = hidden - hidden = min(int(growth * hidden), max_hidden) - - self.lstm = BLSTM(chin, bi=not causal) - if rescale: - rescale_module(self, reference=rescale) - - def valid_length(self, length): - """ - Return the nearest valid length to use with the model so that - there is no time steps left over in a convolutions, e.g. for all - layers, size of the input - kernel_size % stride = 0. - - If the mixture has a valid length, the estimated sources - will have exactly the same length. - """ - length = math.ceil(length * self.resample) - for idx in range(self.depth): - length = math.ceil((length - self.kernel_size) / self.stride) + 1 - length = max(length, 1) - for idx in range(self.depth): - length = (length - 1) * self.stride + self.kernel_size - length = int(math.ceil(length / self.resample)) - return int(length) - - @property - def total_stride(self): - return self.stride ** self.depth // self.resample - - def forward(self, mix): - if mix.dim() == 2: - mix = mix.unsqueeze(1) - - if self.normalize: - mono = mix.mean(dim=1, keepdim=True) - std = mono.std(dim=-1, keepdim=True) - mix = mix / (self.floor + std) - else: - std = 1 - length = mix.shape[-1] - x = mix - x = F.pad(x, (0, self.valid_length(length) - length)) - if self.resample == 2: - x = upsample2(x) - elif self.resample == 4: - x = upsample2(x) - x = upsample2(x) - skips = [] - for encode in self.encoder: - x = encode(x) - skips.append(x) - x = x.permute(2, 0, 1) - x, _ = self.lstm(x) - x = x.permute(1, 2, 0) - for decode in self.decoder: - skip = skips.pop(-1) - x = x + skip[..., :x.shape[-1]] - x = decode(x) - if self.resample == 2: - x = downsample2(x) - elif self.resample == 4: - x = downsample2(x) - x = downsample2(x) - - x = x[..., :length] - return std * x - - -def fast_conv(conv, x): - """ - Faster convolution evaluation if either kernel size is 1 - or length of sequence is 1. - """ - batch, chin, length = x.shape - chout, chin, kernel = conv.weight.shape - assert batch == 1 - if kernel == 1: - x = x.view(chin, length) - out = torch.addmm( - conv.bias.view(-1, 1), - conv.weight.view(chout, chin), x) - elif length == kernel: - x = x.view(chin * kernel, 1) - out = torch.addmm( - conv.bias.view(-1, 1), - conv.weight.view(chout, chin * kernel), x) - else: - out = conv(x) - return out.view(batch, chout, -1) - - -class DemucsStreamer: - """ - Streaming implementation for Demucs. It supports being fed with any amount - of audio at a time. You will get back as much audio as possible at that - point. - - Args: - - demucs (Demucs): Demucs model. - - dry (float): amount of dry (e.g. input) signal to keep. 0 is maximum - noise removal, 1 just returns the input signal. Small values > 0 - allows to limit distortions. - - num_frames (int): number of frames to process at once. Higher values - will increase overall latency but improve the real time factor. - - resample_lookahead (int): extra lookahead used for the resampling. - - resample_buffer (int): size of the buffer of previous inputs/outputs - kept for resampling. - """ - def __init__(self, demucs, - dry=0, - num_frames=1, - resample_lookahead=64, - resample_buffer=256, - mean_decay_duration: float = 10.): - device = next(iter(demucs.parameters())).device - self.demucs = demucs - self.lstm_state = None - self.conv_state = None - self.dry = dry - self.resample_lookahead = resample_lookahead - resample_buffer = min(demucs.total_stride, resample_buffer) - self.resample_buffer = resample_buffer - self.frame_length = demucs.valid_length(1) + demucs.total_stride * (num_frames - 1) - self.total_length = self.frame_length + self.resample_lookahead - self.stride = demucs.total_stride * num_frames - self.resample_in = torch.zeros(demucs.chin, resample_buffer, device=device) - self.resample_out = torch.zeros(demucs.chin, resample_buffer, device=device) - - self.frames = 0 - self.total_time = 0 - self.mean_variance = 0. - self.mean_total = 0. - mean_receptive_field_in_samples = mean_decay_duration * demucs.sample_rate - mean_receptive_field_in_frames = mean_receptive_field_in_samples / demucs.total_stride - self.mean_decay = 1 - 1 / mean_receptive_field_in_frames - - self.pending = torch.zeros(demucs.chin, 0, device=device) - - bias = demucs.decoder[0][2].bias - weight = demucs.decoder[0][2].weight - chin, chout, kernel = weight.shape - self._bias = bias.view(-1, 1).repeat(1, kernel).view(-1, 1) - self._weight = weight.permute(1, 2, 0).contiguous() - - @property - def variance(self) -> float: - return self.mean_variance / self.mean_total - - def reset_time_per_frame(self): - self.total_time = 0 - self.frames = 0 - - @property - def time_per_frame(self): - return self.total_time / self.frames - - def flush(self): - """ - Flush remaining audio by padding it with zero and initialize the previous - status. Call this when you have no more input and want to get back the last - chunk of audio. - """ - self.lstm_state = None - self.conv_state = None - pending_length = self.pending.shape[1] - padding = torch.zeros(self.demucs.chin, self.total_length, device=self.pending.device) - out = self.feed(padding) - return out[:, :pending_length] - - def feed(self, wav): - """ - Apply the model to mix using true real time evaluation. - Normalization is done online as is the resampling. - """ - begin = time.time() - demucs = self.demucs - resample_buffer = self.resample_buffer - stride = self.stride - resample = demucs.resample - - if wav.dim() != 2: - raise ValueError("input wav should be two dimensional.") - chin, _ = wav.shape - if chin != demucs.chin: - raise ValueError(f"Expected {demucs.chin} channels, got {chin}") - - self.pending = torch.cat([self.pending, wav], dim=1) - outs = [] - while self.pending.shape[1] >= self.total_length: - self.frames += 1 - frame = self.pending[:, :self.total_length] - dry_signal = frame[:, :stride] - if demucs.normalize: - mono = frame.mean(0) - variance = (mono**2).mean() - self.mean_variance = self.mean_variance * self.mean_decay + (1 - self.mean_decay) * variance - self.mean_total = self.mean_total * self.mean_decay + (1 - self.mean_decay) - frame = frame / (demucs.floor + torch.sqrt(self.variance)) - padded_frame = torch.cat([self.resample_in, frame], dim=-1) - self.resample_in[:] = frame[:, stride - resample_buffer:stride] - frame = padded_frame - - if resample == 4: - frame = upsample2(upsample2(frame)) - elif resample == 2: - frame = upsample2(frame) - frame = frame[:, resample * resample_buffer:] # remove pre sampling buffer - frame = frame[:, :resample * self.frame_length] # remove extra samples after window - - out, extra = self._separate_frame(frame) - padded_out = torch.cat([self.resample_out, out, extra], 1) - self.resample_out[:] = out[:, -resample_buffer:] - if resample == 4: - out = downsample2(downsample2(padded_out)) - elif resample == 2: - out = downsample2(padded_out) - else: - out = padded_out - - out = out[:, resample_buffer // resample:] - out = out[:, :stride] - - if demucs.normalize: - out *= torch.sqrt(self.variance) - out = self.dry * dry_signal + (1 - self.dry) * out - outs.append(out) - self.pending = self.pending[:, stride:] - - self.total_time += time.time() - begin - if outs: - out = torch.cat(outs, 1) - else: - out = torch.zeros(chin, 0, device=wav.device) - return out - - def _separate_frame(self, frame): - demucs = self.demucs - skips = [] - next_state = [] - first = self.conv_state is None - stride = self.stride * demucs.resample - x = frame[None] - for idx, encode in enumerate(demucs.encoder): - stride //= demucs.stride - length = x.shape[2] - if idx == demucs.depth - 1: - # This is sligthly faster for the last conv - x = fast_conv(encode[0], x) - x = encode[1](x) - x = fast_conv(encode[2], x) - x = encode[3](x) - else: - if not first: - prev = self.conv_state.pop(0) - prev = prev[..., stride:] - tgt = (length - demucs.kernel_size) // demucs.stride + 1 - missing = tgt - prev.shape[-1] - offset = length - demucs.kernel_size - demucs.stride * (missing - 1) - x = x[..., offset:] - x = encode[1](encode[0](x)) - x = fast_conv(encode[2], x) - x = encode[3](x) - if not first: - x = torch.cat([prev, x], -1) - next_state.append(x) - skips.append(x) - - x = x.permute(2, 0, 1) - x, self.lstm_state = demucs.lstm(x, self.lstm_state) - x = x.permute(1, 2, 0) - # In the following, x contains only correct samples, i.e. the one - # for which each time position is covered by two window of the upper layer. - # extra contains extra samples to the right, and is used only as a - # better padding for the online resampling. - extra = None - for idx, decode in enumerate(demucs.decoder): - skip = skips.pop(-1) - x += skip[..., :x.shape[-1]] - x = fast_conv(decode[0], x) - x = decode[1](x) - - if extra is not None: - skip = skip[..., x.shape[-1]:] - extra += skip[..., :extra.shape[-1]] - extra = decode[2](decode[1](decode[0](extra))) - x = decode[2](x) - next_state.append(x[..., -demucs.stride:] - decode[2].bias.view(-1, 1)) - if extra is None: - extra = x[..., -demucs.stride:] - else: - extra[..., :demucs.stride] += next_state[-1] - x = x[..., :-demucs.stride] - - if not first: - prev = self.conv_state.pop(0) - x[..., :demucs.stride] += prev - if idx != demucs.depth - 1: - x = decode[3](x) - extra = decode[3](extra) - self.conv_state = next_state - return x[0], extra[0] - - -def get_demucs(): - model = Demucs(hidden=64) - url = "https://dl.fbaipublicfiles.com/adiyoss/denoiser/dns64-a7761ff99a7d5bb6.th" - state_dict = torch.hub.load_state_dict_from_url(url, map_location='cpu') - model.load_state_dict(state_dict) - return model - - -class Config(BaseModel): - log_folder: Path = Path.home() / 'tmp/tts-service' - hf_repo: str = loaders.DEFAULT_REPO - mimi_weight: Path = Path.home() / 'models/moshi/moshi_e9d43d50@500/e9d43d50_500_mimi_voice.safetensors' - config_path: Path | None = None - device: str = "cpu" - dry_fraction: float = 0.02 - num_cpu_threads: int = 8 - - -class Processor: - def __init__(self, config_override: dict): - print(config_override) - config = Config(**config_override) - torch.set_num_threads(config.num_cpu_threads) - checkpoint_info = loaders.CheckpointInfo.from_hf_repo( - config.hf_repo, mimi_weights=config.mimi_weight, config_path=config.config_path, - ) - self.dry_fraction = config.dry_fraction - loaders._quantizer_kwargs["n_q"] = 16 - print("loading mimi") - self._mimi = checkpoint_info.get_mimi(device=config.device) - print("mimi loaded") - self._length = 24000 * 10 - self._demucs = get_demucs() - self._lowpass = julius.lowpass.LowPassFilter(8 / 24) - self._downsample = julius.resample.ResampleFrac(24, 16) - self._upsample = julius.resample.ResampleFrac(16, 24) - - @torch.no_grad() - def run_one(self, pcm: np.ndarray): - print(pcm.shape) - wav = torch.from_numpy(pcm[None, None, :self._length]).float() - assert wav.shape[-1] == self._length - - low = self._lowpass(wav) - high = wav - low - low = self._downsample(low, full=True) - - denoised = self._demucs(low) - denoised = (1 - self.dry_fraction) * denoised + self.dry_fraction * low - denoised = self._upsample(denoised, output_length=wav.shape[-1]) - denoised = denoised + high - - denoised = normalize_loudness(denoised, 24000) - latent = self._mimi.encode_to_latent(denoised, quantize=False) - latent = latent.cpu().numpy() - print(latent.shape) - return latent - - -def init(config: dict): - processor = Processor(config) - return processor - - - -[workspace] -members = [ - "mimi-pyo3", - "moshi-backend", - "moshi-cli", - "moshi-core", - "moshi-server", -] -resolver = "2" - -[workspace.package] -version = "0.6.3" -edition = "2021" -license = "MIT/Apache-2.0" -description = "moshi, a real-time voice AI" -repository = "https://github.com/kyutai-labs/moshi" -keywords = ["machine-learning", "audio"] -categories = ["science"] - -[workspace.dependencies] -anyhow = "1" -axum = { version = "0.8.1", features = ["ws"] } -axum-server = { version = "0.6", features = ["tls-rustls"] } -base64 = "0.21.7" -base64ct = { version = "1.6.0", features = ["alloc"] } -bincode = "1.3.3" -byteorder = "1.5.0" -candle = { version = "0.9.1", package = "candle-core" } -candle-flash-attn = "0.9.1" -candle-nn = "0.9.1" -candle-transformers = "0.9.1" -clap = { version = "4.4.12", features = ["derive"] } -color-eyre = "0.6.2" -cpal = "0.15.3" -crossterm = { version = "0.27.0", features = ["event-stream"] } -env_logger = "0.10.1" -futures = "0.3.28" -futures-util = "0.3.30" -hf-hub = { version = "0.4.3", features = ["tokio"] } -http = "1.1.0" -kaudio = "0.2.1" -lazy_static = "1.5.0" -log = "0.4.20" -moshi = { path = "./moshi-core", version = "0.6.3" } -native-tls = "0.2.11" -numpy = "0.23.0" -ogg = { version = "0.9.1", features = ["async"] } -opus = "0.3.0" -prometheus = "0.13.4" -pyo3 = "0.23.0" -pyo3-ffi = "0.23.0" -rand = { version = "0.8.5", features = ["getrandom"] } -rand_chacha = "0.3.1" -ratatui = "0.27.0" -rayon = "1.8.1" -rcgen = "0.13.1" -regex = "1.10.3" -rmp-serde = "1.3.0" -rubato = "0.15.0" -rustls = "0.23.5" -sentencepiece = "0.11.2" -serde = { version = "1.0", features = ["derive"] } -serde_json = "1.0.115" -sha3 = "0.10.8" -symphonia = { version = "0.5.3", features = ["all"] } -tokenizers = "0.15.2" -tokio = { version = "1.35.1", features = ["full"] } -tokio-rustls = "0.24.1" -tokio-tungstenite = { version = "0.21.0", features = ["rustls", "native-tls"] } -toml = "0.8.19" -tower = "0.4.13" -tower-http = { version = "0.5", features = ["full"] } -tracing = "0.1.40" -tracing-appender = "0.2.3" -tracing-chrome = "0.7.2" -tracing-subscriber = "0.3.18" -tui-logger = "0.11.2" -vergen = { version = "8.3.1", features = ["build", "cargo", "git", "gitcl", "rustc", "si"] } - -[profile.release] -debug = true - -[profile.release-no-debug] -inherits = "release" -debug = false - - - -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. - - - -# Protocol - -The connection takes place using a websocket. This handles the message lengths -for us. The binary protocol for messages is as follows. The protocol uses little -endian encoding. - -Each message starts by a single byte indicating the message type `MT`. -The format for the rest of the message, aka the payload, depends on `MT`. - -``` -- Handshake MT=0. The payload is made of two fields. - 1. Protocol version (`u32`) - always 0 for now. - 2. Model version (`u32`). -- Audio MT=1. The payload is made of a single field. - - Binary data for the ogg frames containing opus encoded audio (24kHz, mono). -- Text MT=2. The payload is made of a single field. - - UTF8 encoded string. -- Control MT=3. The payload is made of a single field. This is not used in full - streaming mode. - - One byte B describing the control itself. - - Start B=0. - - EndTurn B=1. - - Pause B=2. - - Restart B=3. -- MetaData MT=4. The payload is made of a single field. - - UTF8 encoded string with json data. -- Error MT=5. The payload is made of a single field. - - UTF8 encoded string containing the error description. -- Ping MT=6. No payload, this message type is currently unused. -``` -Messages with an unknow message types should be discarded. - - - -# moshi - rust - -[![Latest version](https://img.shields.io/crates/v/moshi.svg)](https://crates.io/crates/moshi) -[![Documentation](https://docs.rs/moshi/badge.svg)](https://docs.rs/moshi) -![License](https://img.shields.io/crates/l/moshi.svg) - -See the [top-level README.md](../README.md) for more information. - -This provides the Rust backend (both Mimi and Moshi) and client implementation. -The Mimi implementation is available through Python bindings, through the `rustymimi` package. - -## Requirements - -You will need a recent version of the [Rust toolchain](https://rustup.rs/). -To compile GPU support, you will also need the [CUDA](https://developer.nvidia.com/cuda-toolkit) properly installed for your GPU, in particular with `nvcc`. - - -## Rust based Mimi with Python bindings - -First, a standalone rust based implementation of Mimi is provided, along with Python bindings. -This is the one used by `moshi_mlx`. It is automatically installed with `moshi_mlx`, but you -can install it separately as -```bash -# Install from pip: -pip install rustymimi -# Alternatively, if you want to compile the package run from the root of the repo. -maturin dev -r -m rust/mimi-pyo3/Cargo.toml -``` - -## Rust server - -If you don't have ssl certificates yet, generate a `key.pem` and `cert.pem` file -using the following command. -```bash -openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -days 365 -nodes -subj "/CN=localhost" -``` - -In order to run the rust inference server, use the following command from within -the this directory: - -```bash -cargo run --features cuda --bin moshi-backend -r -- --config moshi-backend/config.json standalone -``` - -When using macOS, you can replace `--features cuda` with `--features metal`. - -Alternatively you can use `config-q8.json` rather than `config.json` to use the -quantified q8 model. You can select a different pretrained model, e.g. Moshika, -by changing the `"hf_repo"` key in either file. - -Once the server has printed 'standalone worker listening', you can use the web -UI. By default the rust version uses https so it will be at -[localhost:8998](https://localhost:8998). - -You will get some warnings about the site being unsafe. When using chrome you -can bypass it by selecting "Details" or "Advanced", then "Visit this unsafe -site" or "Proceed to localhost (unsafe)". - -## Rust client - -We recommend using the web UI as it provides some echo cancellation that helps -the overall model quality. Alternatively we provide some command line interfaces -for the rust and python versions, the protocol is the same as with the web UI so -there is nothing to change on the server side. - -### Rust Command Line - -From within the `rust` directory, run the following: -```bash -cargo run --bin moshi-cli -r -- tui --host localhost -``` - -## License - -The present code is provided under the Apache license. - - - -use_small_heuristics = "Max" -edition = "2021" - - - -text_in_vocab_size = 48001 -text_out_vocab_size = 48000 -audio_vocab_size = 2049 -audio_codebooks = 16 - -[transformer] -d_model = 2048 -num_heads = 16 -num_layers = 16 -dim_feedforward = 8192 -causal = true -norm_first = true -bias_ff = false -bias_attn = false -context = 3000 -max_period = 100000 -use_conv_block = false -use_conv_bias = true -gating = "silu" -norm = "RmsNorm" -positional_embedding = "Rope" -conv_layout = false -conv_kernel_size = 3 -kv_repeat = 1 -max_seq_len = 4096 - -[depformer] -num_slices = 8 - -[depformer.transformer] -d_model = 1024 -num_heads = 16 -num_layers = 6 -dim_feedforward = 4096 -causal = true -norm_first = true -bias_ff = false -bias_attn = false -context = 32 -max_period = 10000 -use_conv_block = false -use_conv_bias = true -gating = "silu" -norm = "RmsNorm" -positional_embedding = "None" -conv_layout = false -conv_kernel_size = 3 -kv_repeat = 1 -max_seq_len = 4096 - -[conditioners.description] -type = "Lut" -n_bins = 31 -dim = 16 -possible_values = ["very_bad", "bad", "neutral", "good", "very_good"] - - - diff --git a/src/lib.rs b/src/lib.rs index 79898813..d91a6dd1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -96,13 +96,13 @@ pub struct WordTimestamp { } /// WebSocket interface for real-time transcription streaming. -/// +/// /// ## Connection /// Connect to `ws://localhost:/` where `` is specified via the `--ws` option. -/// +/// /// ## Message Format /// All messages are JSON objects with a `type` field indicating the message type: -/// +/// /// ### Word Message /// Sent for each transcribed word as it's recognized: /// ```json @@ -113,7 +113,7 @@ pub struct WordTimestamp { /// "end_time": 1.45 // null for real-time words without end time yet /// } /// ``` -/// +/// /// ### Pause Message /// Sent when voice activity detection detects a pause (requires --vad flag): /// ```json @@ -122,7 +122,7 @@ pub struct WordTimestamp { /// "timestamp": 1234567890.123 /// } /// ``` -/// +/// /// ### Final Message /// Sent at the end of transcription with complete results: /// ```json @@ -135,7 +135,7 @@ pub struct WordTimestamp { /// ] /// } /// ``` -/// +/// /// ### Restart Command /// Send from client to restart transcription after timeout or final message: /// ```json @@ -143,12 +143,12 @@ pub struct WordTimestamp { /// "type": "restart" /// } /// ``` -/// +/// /// ## Usage Example /// ```bash /// # Start transcription with WebSocket on port 8080 /// ears --live --ws 8080 -/// +/// /// # With timestamps and VAD /// ears --live --ws 8080 --timestamps --vad /// ``` @@ -162,9 +162,7 @@ pub enum WebSocketMessage { end_time: Option, }, #[serde(rename = "pause")] - Pause { - timestamp: f64, - }, + Pause { timestamp: f64 }, #[serde(rename = "final")] Final { text: String, @@ -243,7 +241,9 @@ impl Model { for chunk in pcm.chunks(1920) { let tensor = Tensor::new(chunk, &self.dev)?.reshape((1, 1, chunk.len()))?; - let _ = self.state.step_pcm(tensor, None, &().into(), |_, _, _| ())?; + let _ = self + .state + .step_pcm(tensor, None, &().into(), |_, _, _| ())?; } Ok(()) } @@ -412,24 +412,23 @@ impl Model { ws_port: u16, ) -> Result { use futures::{SinkExt, StreamExt}; + use std::io::Write; use std::sync::Arc; use tokio::sync::{broadcast, mpsc}; use tokio_tungstenite::{accept_async, tungstenite::Message}; - use std::io::Write; - // Create broadcast channel for WebSocket messages + // WebSocket broadcast channel let (ws_tx, _ws_rx) = broadcast::channel(100); let ws_tx = Arc::new(ws_tx); - // Create channel for restart commands + // Channel used to request a restart of the transcription session let (restart_tx, mut restart_rx) = mpsc::unbounded_channel(); let restart_tx = Arc::new(restart_tx); - // Start WebSocket server + // Spawn WebSocket server let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{}", ws_port)).await?; let ws_tx_clone = ws_tx.clone(); let restart_tx_clone = restart_tx.clone(); - tokio::spawn(async move { while let Ok((stream, _)) = listener.accept().await { let ws_tx = ws_tx_clone.clone(); @@ -446,17 +445,15 @@ impl Model { let (mut ws_sender, mut ws_receiver) = ws_stream.split(); let mut ws_rx = ws_tx.subscribe(); - // Handle incoming WebSocket messages let receive_task = tokio::spawn(async move { while let Some(msg) = ws_receiver.next().await { match msg { Ok(Message::Close(_)) => break, Ok(Message::Text(text)) => { - if let Ok(command) = serde_json::from_str::(&text) { - match command { - WebSocketCommand::Restart => { - let _ = restart_tx.send(()); - } + if let Ok(cmd) = serde_json::from_str::(&text) + { + if let WebSocketCommand::Restart = cmd { + let _ = restart_tx.send(()); } } } @@ -464,12 +461,11 @@ impl Model { eprintln!("WebSocket receive error: {}", e); break; } - _ => {} // Ignore other message types + _ => {} } } }); - // Forward broadcast messages to WebSocket let send_task = tokio::spawn(async move { while let Ok(ws_msg) = ws_rx.recv().await { let json_msg = serde_json::to_string(&ws_msg).unwrap_or_default(); @@ -479,186 +475,166 @@ impl Model { } }); - tokio::select! { - _ = receive_task => {}, - _ = send_task => {}, - } + let _ = tokio::join!(receive_task, send_task); }); } }); - // Process audio synchronously with WebSocket streaming + // Bridge blocking audio receiver to async channel + let (pcm_tx, mut pcm_rx) = mpsc::unbounded_channel(); + std::thread::spawn(move || { + while let Ok(chunk) = audio_rx.recv() { + if pcm_tx.send(chunk).is_err() { + break; + } + } + }); + let mut all_audio = Vec::new(); let mut overall_words = Vec::new(); let mut overall_text = String::new(); - + loop { let mut words = Vec::new(); let mut current_text = String::new(); let mut last_word: Option<(String, f64)> = None; let mut printed_eot = false; let mut last_voice_activity: Option = None; - let mut transcription_active = true; + let mut restart = false; eprintln!("Starting transcription session..."); - while transcription_active { + loop { tokio::select! { - // Handle restart commands _ = restart_rx.recv() => { eprintln!("Received restart command"); + restart = true; break; } - - // Handle audio chunks with timeout - _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => { - match audio_rx.try_recv() { - Ok(pcm_chunk) => { - if save_audio.is_some() { - all_audio.extend_from_slice(&pcm_chunk); - } + Some(pcm_chunk) = pcm_rx.recv() => { + if save_audio.is_some() { + all_audio.extend_from_slice(&pcm_chunk); + } - let mut has_voice_activity = false; - - for pcm in pcm_chunk.chunks(1920) { - let pcm_tensor = Tensor::new(pcm, &self.dev)?.reshape((1, 1, ()))?; - let asr_msgs = self - .state - .step_pcm(pcm_tensor, None, &().into(), |_, _, _| ())?; - - for asr_msg in asr_msgs.iter() { - match asr_msg { - moshi::asr::AsrMsg::Step { prs, .. } => { - if self.vad && prs[2][0] > 0.5 && !printed_eot { - printed_eot = true; - let pause_msg = WebSocketMessage::Pause { - timestamp: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs_f64(), - }; - let _ = ws_tx.send(pause_msg); - - if !self.timestamps { - print!(" "); - std::io::stdout().flush().ok(); - } - } + let mut has_voice_activity = false; + + for pcm in pcm_chunk.chunks(1920) { + let pcm_tensor = Tensor::new(pcm, &self.dev)?.reshape((1, 1, ()))?; + let asr_msgs = self.state.step_pcm(pcm_tensor, None, &().into(), |_, _, _| ())?; + + for asr_msg in asr_msgs.iter() { + match asr_msg { + moshi::asr::AsrMsg::Step { prs, .. } => { + if self.vad && prs[2][0] > 0.5 && !printed_eot { + printed_eot = true; + let pause_msg = WebSocketMessage::Pause { + timestamp: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64(), + }; + let _ = ws_tx.send(pause_msg); + + if !self.timestamps { + print!(" "); + std::io::stdout().flush().ok(); } - moshi::asr::AsrMsg::EndWord { stop_time, .. } => { - printed_eot = false; - has_voice_activity = true; - if self.timestamps { - if let Some((word, start_time)) = last_word.take() { - println!("[{start_time:5.2}-{stop_time:5.2}] {word}"); - - let word_timestamp = WordTimestamp { - word: word.clone(), - start_time, - end_time: Some(*stop_time), - }; - words.push(word_timestamp.clone()); - - let ws_msg = WebSocketMessage::Word { - word: word_timestamp.word, - start_time: word_timestamp.start_time, - end_time: word_timestamp.end_time, - }; - let _ = ws_tx.send(ws_msg); - } - } + } + } + moshi::asr::AsrMsg::EndWord { stop_time, .. } => { + printed_eot = false; + has_voice_activity = true; + if self.timestamps { + if let Some((word, start_time)) = last_word.take() { + println!("[{start_time:5.2}-{stop_time:5.2}] {word}"); + let word_ts = WordTimestamp { + word: word.clone(), + start_time, + end_time: Some(*stop_time), + }; + words.push(word_ts.clone()); + let ws_msg = WebSocketMessage::Word { + word: word_ts.word, + start_time: word_ts.start_time, + end_time: word_ts.end_time, + }; + let _ = ws_tx.send(ws_msg); } - moshi::asr::AsrMsg::Word { - tokens, start_time, .. - } => { - printed_eot = false; - has_voice_activity = true; - let word = self - .text_tokenizer - .decode_piece_ids(tokens) - .unwrap_or_else(|_| String::new()); - - current_text.push(' '); - current_text.push_str(&word); - - if !self.timestamps { - print!(" {}", word); - std::io::stdout().flush().ok(); - - // Send word without end time for real-time streaming - let ws_msg = WebSocketMessage::Word { - word: word.clone(), - start_time: *start_time, - end_time: None, - }; - let _ = ws_tx.send(ws_msg); - } else { - if let Some((prev_word, prev_start_time)) = last_word.take() { - println!( - "[{prev_start_time:5.2}-{start_time:5.2}] {prev_word}" - ); - - let word_timestamp = WordTimestamp { - word: prev_word.clone(), - start_time: prev_start_time, - end_time: Some(*start_time), - }; - words.push(word_timestamp.clone()); - - let ws_msg = WebSocketMessage::Word { - word: word_timestamp.word, - start_time: word_timestamp.start_time, - end_time: word_timestamp.end_time, - }; - let _ = ws_tx.send(ws_msg); - } - last_word = Some((word, *start_time)); - } + } + } + moshi::asr::AsrMsg::Word { tokens, start_time, .. } => { + printed_eot = false; + has_voice_activity = true; + let word = self.text_tokenizer + .decode_piece_ids(tokens) + .unwrap_or_else(|_| String::new()); + + current_text.push(' '); + current_text.push_str(&word); + + if !self.timestamps { + print!(" {}", word); + std::io::stdout().flush().ok(); + + let ws_msg = WebSocketMessage::Word { + word: word.clone(), + start_time: *start_time, + end_time: None, + }; + let _ = ws_tx.send(ws_msg); + } else { + if let Some((prev_word, prev_start_time)) = last_word.take() { + println!("[{prev_start_time:5.2}-{start_time:5.2}] {prev_word}"); + let word_ts = WordTimestamp { + word: prev_word.clone(), + start_time: prev_start_time, + end_time: Some(*start_time), + }; + words.push(word_ts.clone()); + let ws_msg = WebSocketMessage::Word { + word: word_ts.word, + start_time: word_ts.start_time, + end_time: word_ts.end_time, + }; + let _ = ws_tx.send(ws_msg); } + last_word = Some((word, *start_time)); } } } + } + } - // Update voice activity timestamp if we detected voice - if has_voice_activity { - last_voice_activity = Some(std::time::Instant::now()); - } + if has_voice_activity { + last_voice_activity = Some(std::time::Instant::now()); + } - // Check for timeout - if let Some(timeout_secs) = self.vad_timeout { - if let Some(last_activity) = last_voice_activity { - if last_activity.elapsed() > std::time::Duration::from_secs_f64(timeout_secs) { - eprintln!("Voice activity timeout reached"); - transcription_active = false; - } - } + if let Some(timeout_secs) = self.vad_timeout { + if let Some(last_activity) = last_voice_activity { + if last_activity.elapsed() > std::time::Duration::from_secs_f64(timeout_secs) { + eprintln!("Voice activity timeout reached"); + break; } } - Err(_) => { - // No audio data available, continue waiting - } } } } } - // Handle final word for this session if let Some((word, start_time)) = last_word.take() { if self.timestamps { println!("[{start_time:5.2}- ] {word}"); } - - let word_timestamp = WordTimestamp { + let word_ts = WordTimestamp { word: word.clone(), start_time, end_time: None, }; - words.push(word_timestamp.clone()); - + words.push(word_ts.clone()); let ws_msg = WebSocketMessage::Word { - word: word_timestamp.word, - start_time: word_timestamp.start_time, - end_time: word_timestamp.end_time, + word: word_ts.word, + start_time: word_ts.start_time, + end_time: word_ts.end_time, }; let _ = ws_tx.send(ws_msg); } @@ -667,7 +643,6 @@ impl Model { println!(); } - // Add session words to overall collection overall_words.extend(words.clone()); if !current_text.is_empty() { if !overall_text.is_empty() { @@ -676,7 +651,6 @@ impl Model { overall_text.push_str(current_text.trim()); } - // Send final result for this session let session_result = TranscriptionResult { text: current_text.trim().to_string(), words: words.clone(), @@ -688,9 +662,7 @@ impl Model { }; let _ = ws_tx.send(final_msg); - // Check if we should continue or exit - if restart_rx.try_recv().is_err() { - // No restart command pending, exit the loop + if !restart { break; } } @@ -704,7 +676,6 @@ impl Model { words: overall_words, }) } - fn transcribe_pcm(&mut self, mut pcm: Vec) -> Result { if self.config.stt_config.audio_silence_prefix_seconds > 0.0 { let silence_len = @@ -851,10 +822,13 @@ pub mod audio { Ok(()) } - pub fn start_audio_capture(audio_tx: Sender>, device_index: Option) -> Result<()> { + pub fn start_audio_capture( + audio_tx: Sender>, + device_index: Option, + ) -> Result<()> { let max_retries = 3; let mut retry_count = 0; - + loop { match start_audio_capture_internal(audio_tx.clone(), device_index) { Ok(()) => break, @@ -863,7 +837,10 @@ pub mod audio { if retry_count >= max_retries { return Err(e); } - eprintln!("Audio capture failed (attempt {}/{}): {}", retry_count, max_retries, e); + eprintln!( + "Audio capture failed (attempt {}/{}): {}", + retry_count, max_retries, e + ); eprintln!("Retrying audio capture in 1 second..."); thread::sleep(std::time::Duration::from_secs(1)); } @@ -872,7 +849,10 @@ pub mod audio { Ok(()) } - fn start_audio_capture_internal(audio_tx: Sender>, device_index: Option) -> Result<()> { + fn start_audio_capture_internal( + audio_tx: Sender>, + device_index: Option, + ) -> Result<()> { let host = cpal::default_host(); let device = if let Some(index) = device_index { host.input_devices()? @@ -947,4 +927,3 @@ pub mod audio { } } } -