csm-mlx

An implementation of the CSM(Conversation Speech Model) for Apple Silicon using MLX.

We also have a CLI!

And fine-tuning too!

Installation

Note that this is not a CLI installation. Please refer to the CLI section for more information.

Recommendation: Give uv a try. It's truly magical.

uv add git+https://github.com/senstella/csm-mlx --upgrade

Or, you can install it via pip:

pip install git+https://github.com/senstella/csm-mlx --upgrade

Make sure to use Python < 3.13. It tends to have sentencepiece compiler error.

Usage

Basic generation

from mlx_lm.sample_utils import make_sampler
from huggingface_hub import hf_hub_download
from csm_mlx import CSM, csm_1b, generate

import audiofile
import numpy as np

# Initialize the model
csm = CSM(csm_1b())  # csm_1b() is a configuration for the CSM model.
weight = hf_hub_download(repo_id="senstella/csm-1b-mlx", filename="ckpt.safetensors")
csm.load_weights(weight)

# Generate audio from text
audio = generate(
    csm,
    text="Hello from Sesame.",
    speaker=0,
    context=[],
    max_audio_length_ms=10_000,
    sampler=make_sampler(temp=0.8, top_k=50), # Put mlx_lm's sampler here! Supports: temp, top_p, min_p, min_tokens_to_keep, top_k.
    # Additionally, you can provide `stream` argument to specify what device to use for generation.
    # https://ml-explore.github.io/mlx/build/html/usage/using_streams.html
)

audiofile.write("./audio.wav", np.asarray(audio), 24000)

Providing Context

from csm_mlx import CSM, csm_1b, generate, Segment
import mlx.core as mx

# Initialize the model
csm = CSM(csm_1b())
weight = hf_hub_download(repo_id="senstella/csm-1b-mlx", filename="ckpt.safetensors")
csm.load_weights(weight)

# Create previous conversation segments
context = [
    Segment(
        speaker=0,
        text="How are you doing today?",
        audio=mx.array(...)  # Previous audio for this segment
    ),
    Segment(
        speaker=1,
        text="I'm doing great, thank you!",
        audio_path="~/somewhere_in_the_universe/stuff.wav"  # Or you can specify the audio path too!
    )
]

# Generate a response in the conversation
audio = generate(
    csm,
    text="That's wonderful to hear!",
    speaker=0,
    context=context,
    max_audio_length_ms=5_000
    # If you don't provide any sampler, greedy sampling will be used.
)

Speed up via quantization

from mlx import nn
from mlx_lm.sample_utils import make_sampler, make_logits_processors
from huggingface_hub import hf_hub_download
from csm_mlx import CSM, csm_1b, generate

import audiofile
import numpy as np

# Initialize the model
csm = CSM(csm_1b())  # csm_1b() is a configuration for the CSM model.
weight = hf_hub_download(repo_id="senstella/csm-1b-mlx", filename="ckpt.safetensors")
csm.load_weights(weight)

# Quantize
# If you want, specify the `group_size`, `bits`.
# https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.quantize.html
nn.quantize(csm)

# Generate audio from text
audio = generate(
    csm,
    text="Hello from Sesame.",
    speaker=0,
    context=[],
    max_audio_length_ms=10_000,
    sampler=make_sampler(temp=0.8, top_k=50), # Put mlx_lm's sampler here! Supports: temp, top_p, min_p, min_tokens_to_keep, top_k.
    # If the model doesn't end the sequence properly, you can try to adjust the sampler parameters
    # or set eos_token logit bias via logits_processors=make_logits_processors(logit_bias={0: 4})
    # Additionally, you can provide `stream` argument to specify what device to use for generation.
    # https://ml-explore.github.io/mlx/build/html/usage/using_streams.html
)

audiofile.write("./audio.wav", np.asarray(audio), 24000)

Streaming generation

from mlx import nn
from mlx_lm.sample_utils import make_sampler
from csm_mlx import CSM, csm_1b
from csm_mlx.generation import stream_generate

# Initialize model
csm = CSM(csm_1b())
csm.load_weights("ckpt.safetensors", strict=True)

# nn.quantize(csm) — Speed up nearly real-time on M2 Air, but loses quality.

# Stream generate audio chunks
for chunk in stream_generate(
    csm,
    text="This is an example of streaming audio generation.",
    speaker=0,
    context=[],
    max_audio_length_ms=5000,
    sampler=make_sampler(temp=0.8, top_k=50),
):
    # Process each chunk as it's generated

Loading audio for a segment

If you want to load an audio for a segment, you need to resample it to 24000.

import mlx.core as mx
import audiofile
import audresample

def read_audio(audio_path, sampling_rate=24000) -> mx.array:
    signal, original_sampling_rate = audiofile.read(audio_path, always_2d=True)
    signal = audresample.resample(signal, original_sampling_rate, sampling_rate)

    signal = mx.array(signal)

    if signal.shape[0] >= 1:
        signal = signal.mean(axis=0)
    else:
        signal = signal.squeeze(0)

    return signal  # (audio_length, )

CLI

Installation

If you're having sentencepiece error on uv, please try adding --python 3.12 flag on the uv tool install command.

# Recommendation: uv tools - works best!
uv tool install "git+https://github.com/senstella/csm-mlx[cli]" --upgrade

# Or with pipx
pipx install "git+https://github.com/senstella/csm-mlx[cli]" --upgrade

Usage

Finetuning CLI usage is here.

Basic Text-to-Speech

csm-mlx generate "Hello from Sesame." -o output.wav

With Options

csm-mlx generate "Hello from Sesame." \
  --output output.wav \
  --model 1b \
  --speaker 0 \
  --temperature 0.8 \
  --min-p 0.05 \
  --max-audio-length 10000

Feeding Contexts

You can provide conversation context to make the generated speech more natural — or clone a voice with it.

You must provide audio & text & speaker in the pair.

csm-mlx generate "Nice to meet you too!" \
  --output response.wav \
  --input-audios previous.wav \
  --input-texts "Hello, nice to meet you." \
  --input-speakers 1

Portable with uv

uv run --with 'git+https://github.com/senstella/csm-mlx[cli]' --python 3.12 python -m csm_mlx "Hello from Sesame." -o output.wav

CLI Reference

csm-mlx generate [TEXT] [OPTIONS]

Arguments

TEXT: The text to convert to speech

Options

-o, --output PATH: Output audio file path [required]
-m, --model [1b]: Model size (default: 1b)
-s, --speaker INT: Speaker ID (default: 0)
-l, --max-audio-length INT: Maximum audio length in milliseconds (default: 10000 — 10 seconds)
-t, --temperature, --temp FLOAT: Sampling temperature (default: 0.8)
-p, --top-p FLOAT: Top-p sampling parameter
-m, --min-p FLOAT: Minimum probability for sampling (default: 0.05)
-k, --top-k INT: Top-k sampling parameter
-kt, --min-tokens-to-keep INT: Minimum tokens to keep during sampling (default: 1)
-is, --input-speakers LIST: List of speaker IDs for context
-ia, --input-audios LIST: List of audio files for context
-it, --input-texts LIST: List of text transcripts for context

Todo

Fix up RoPE
Implement watermarking
Add streaming generation
Optimize performance further for real-time inference

Acknowledgments

Thanks to Sesame for original PyTorch implementation and weights!
Thanks to Moshi project for creating the mimi codec and mimi_mlx implementation.
Thanks to torchtune project for providing LLaMA attention implementation.
Thanks to MLX project for providing the framework that made this implementation possible.
Thanks to typer for powering the CLI interface.
Thanks to audiofile and audresample for audio processing.

License

Apache 2.0

License

senstella/csm-mlx

ErrorLooks like something went wrong!

About

Resources

License

Uh oh!

Stars

Watchers

Forks

Releases

Packages 0

Uh oh!

Contributors 5

Languages

Error
Looks like something went wrong!

Packages