Memory Usage in Kokoro

[chigkim]

I'm trying to incorporate MLX-Audio into AIChat that generates a conversation between two LLMs as speech.
However, memory consumption just goes up and up and up indefinitely as I generate more. I don't have this problem with kokoro-onnx.
I'm not sure if it's my code or MLX-audio, but it doesn't seem to release previous generation from memory after I append the samples to an array?
Here's the minimal code that I can use to produce.

from mlx_audio.tts.utils import load_model
import numpy as np
import codecs
import random
import json
import soundfile as sf

def random_pause(min_duration=0.5, max_duration=1.0, sample_rate=24000):
    silence_duration = random.uniform(min_duration, max_duration)
    silence = np.zeros(int(silence_duration * sample_rate))
    return silence

def generate_audio(text, voice, l, r, speed=1.0, lang="a", sample_rate=24000):
    res = tts.generate(
        text=text,
        voice=voice,
        speed=speed,
        lang_code=lang,
        temperature=0.7,
        verbose=False,
    )
    samples = list(res)[0].audio
    pause = random_pause(sample_rate=sample_rate)
    samples = np.concatenate([samples, pause])
    samples = np.column_stack((samples*l, samples*r))
    return samples

tts = load_model("mlx-community/Kokoro-82M-bf16")
a_voice= "af_sky"
b_voice = "af_heart"

chat = json.load(codecs.open("chat.json", "r", "utf-8"))
wavs = []
for i in range(0, len(chat), 2):
    content = chat[i]["content"]
    print(a_voice, content)
    samples = generate_audio(content, a_voice, 0.8, 1.0)
    wavs.append(samples)
    content = chat[i+1]["content"]
    print(b_voice, content)
    samples = generate_audio(content, b_voice, 1.0, 0.8)
    wavs.append(samples)

wav = np.concat(wavs)
sf.write(f"podcast.wav", wav, 24000)

I'm also attaching the chat history that you can use to reproduce.

chat.json

Thanks!

[lucasnewman]

If you do del res before returning from generate_audio() does it fix it? You could also try a direct np.array(samples) cast to for evaluation and make sure you aren't holding a reference to the underlying mx.array.

In theory the function scope would control your references to the MLX graph, but there could be a lingering reference somewhere since you're persisting all the downstream audio segments.

[Blaizzy]

In addition to what Lucas said, you can also do mx.metal.clear_cache() at the end of each generation

[Blaizzy]

Yes, he is accumulating results without releasing resources.

In regard to evaluation, on v0.0.2 we added eval on Kokoro.

[chigkim]

I changed to:

audio = list(res)[0].audio
samples = np.array(audio)
del res
del audio

However, it still going up crazy.
Also I tried mlx.metal.clear_cache(), but it says mlx doesn't have attribute metal.
I do collect samples returned from generate_audio into an list and assemble all the generated sample into one wav file at the end.
However, collecting them doesn't seem to have a problem in kokoro-onnx.

[Blaizzy]

https://ml-explore.github.io/mlx/build/html/search.html?q=clear+cache#

[Blaizzy]

Can you share the text you are generating?

And also try to use the generate.py and pass --join_audio flag and see if the same happens

[chigkim]

Yes, If you just use the chat history file in json that I attached in the first post, you should be able to duplicate.
Thanks!

[chigkim]

Ah it's mlx.core.metal.clear_casche().
It kind of works. Memory consumption goes back down after a generation. However, each generation takes more memory. For example:
generation 1: 16gb-22gb
generation 2: 16gb-26gb
generation 3: 16gb-30gb
generation 4: 16gb-34gb
...
I just made up the numbers to illustrate the trend. They're not the actual numbers I saw.
Thanks for your help!

[Blaizzy]

Awesome!

We will take a look and address it in future releases

[chigkim]

I think this is something you might want to address soon. I looked into it further, and the memory consumption differs significantly compared to kokoro-onnx.
When I use mlx-audio, memory usage increases rather quickly. Even though the peak memory usage reported by model.generate() shows around 4.5GB per generation, Activity Monitor shows total used memory jumping from 10GB to 50GB. Wired memory also climbs to about 40GB only after several generations.
This is with using mlx.core.metal.clear_cache()—without clearing the cache, my computer becomes unusable very quickly after a few generations.
In contrast, running the same script with kokoro-onnx, total memory stays around 16GB and increases very slowly as it accumulates the generated audio. It increases by exactly how much I expect the total audio would be. Wired memory stays steady at around 5.5-7GB depending on the generation.
I'm using the largest kokoro-onnx model (about 325MB), so I believe the precision should be comparable.
Thanks again!

[chigkim]

FYI, here's a screen recording.
https://github.com/user-attachments/assets/0d14a67e-76d0-49dc-a1f2-4bbfe2ac9de2
Hope that helps.
Thanks.

[Blaizzy]

Sure, will look into it.