app.py

import spaces
import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from moshi.models import loaders, LMGen
import numpy as np
from tqdm import tqdm


MAX_LENGTH = 24000 * 5  # For example, 30 seconds of audio at 24kHz
mimi_weight = hf_hub_download(loaders.DEFAULT_REPO, loaders.MIMI_NAME)
moshi_weight = hf_hub_download(loaders.DEFAULT_REPO, loaders.MOSHI_NAME)


@spaces.GPU
def compute_codes(wav):
    """wav = torch.randn(1, 1, 24000 * 10)  # should be [B, C=1, T]"""
    mimi = loaders.get_mimi(mimi_weight)
    mimi.set_num_codebooks(8)  # up to 32 for mimi, but limited to 8 for moshi.

    with torch.no_grad():
        # Supports streaming too.
        frame_size = int(mimi.sample_rate / mimi.frame_rate)
        all_codes = []
        with mimi.streaming(batch_size=1):
            for offset in tqdm(range(0, wav.shape[-1], frame_size), desc="computing Codes"):
                frame = wav[:, :, offset: offset + frame_size]
                codes = mimi.encode(frame)
                if codes.shape[-1] == 1:
                    all_codes.append(codes)
                else:
                    print(f"Warning: Empty codes for frame at offset {offset}")
    return all_codes

@spaces.GPU
def generate_reponse(all_codes):
    """wav = torch.randn(1, 1, 24000 * 10)  # should be [B, C=1, T]"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Set up Mimi
    mimi = loaders.get_mimi(mimi_weight, device='cpu')
    mimi.set_num_codebooks(8)  # up to 32 for mimi, but limited to 8 for moshi.
    mimi.to(device)

    # Set up Moshi/LM Gen
    moshi = loaders.get_moshi_lm(moshi_weight, device='cpu')
    moshi.to(device)  # Move to GPU after loading
    lm_gen = LMGen(moshi, temp=0.8, temp_text=0.7)  # this handles sampling params etc.

    out_wav_chunks = []


    # Now we will stream over both Moshi I/O, and decode on the fly with Mimi.
    with torch.no_grad(), lm_gen.streaming(1), mimi.streaming(1):
        for idx, code in tqdm(enumerate(all_codes), desc="generate tokens"):
            # print("CODE: ", code.shape)
            tokens_out = lm_gen.step(code.to(device))
            # tokens_out is [B, 1 + 8, 1], with tokens_out[:, 1] representing the text token.
            if tokens_out is not None:
                wav_chunk = mimi.decode(tokens_out[:, 1:])
                out_wav_chunks.append(wav_chunk)
              
            print(idx, end='\r')



    return torch.cat(out_wav_chunks, dim=-1)

def convert2wav(audio):
    if audio is None:
        return None
    
    sr, data = audio
    
    # Convert to mono if stereo
    if len(data.shape) > 1:
        data = np.mean(data, axis=1)
    
    # Convert to torch tensor
    wav = torch.from_numpy(data).float()
    
    # Reshape to (1, 1, samples)
    wav = wav.unsqueeze(0).unsqueeze(0)
    
    # Resample to 24000 Hz if necessary
    if sr != 24000:
        wav = torch.nn.functional.interpolate(wav, size=24000 * 10, mode='linear', align_corners=False)
    
    # Ensure the tensor has the correct shape (1, 1, 24000 * 10)
    wav = wav[:, :, :24000 * 10]
    
    return wav

def truncate_audio(wav, max_length):
    if wav.shape[2] > max_length:
        return wav[:, :, -max_length:]
    return wav


##########################################################################################################
##########################################################################################################

def process_audio(audio, instream):
    log_out = ""
    outwav = torch.randn(1, 1, 24000 * 2)
    stream = torch.randn(1, 1, 24000 * 2)


    print("Audio recieved")
    if audio is None:
        return gr.update(), (24000, outwav.squeeze().cpu().numpy()), instream, gr.update(visible=True,value=f"Audio is None")
    
    try:
        if instream is None:
            instream = (24000, torch.randn(1, 1, 24000 * 10).squeeze().cpu().numpy())
        print("1. COMBINE AUDIO WITH PREVIOUS CONVERSATION TO STORE")
        stream = (audio[0], np.concatenate((instream[1], audio[1])))

        # Assuming instream[1] and audio[1] are valid inputs for convert2wav
        print("2. CONVERT AUDIO TO WAV")
        wav1 = convert2wav(instream)
        wav2 = convert2wav(audio)

        # Concatenate along the last dimension (time axis)
        print("3. COMBINE AUDIOS TO A SINGLE STREAM")
        combined_wav = torch.cat((wav1, wav2), dim=2)

        # Truncate Audio to a defined length to recude computational efforts
        print("4. TRUNCATE AUDIO LENGTH TO GIVEN DURATION")
        combined_wav = truncate_audio(combined_wav, MAX_LENGTH)
        
        # Preprocessing, convert the audio into the processable codes/tokens
        print("5. COMPUTE CODES")
        mimi_codes = compute_codes(combined_wav)

        # Generation of the Model's reponse
        print("6. GENRATE TOKENS")
        outwav = generate_reponse(mimi_codes)


    except Exception as e:
        return gr.update(value=None), (24000, outwav.squeeze().cpu().numpy()), stream, gr.update(visible=True,value=f"LOG: \n{e}")

    return gr.update(value=None), (24000, outwav.squeeze().cpu().numpy()), stream, gr.update(visible=False)


with gr.Blocks() as demo:
    gr.Markdown("# Moshi Demo")
    gr.Markdown(" ")
    gr.Markdown("-----------")

    gr.Markdown("### Model Description")
    gr.Markdown("""Moshi is a speech-text foundation model that casts spoken dialogue as speech-to-speech generation. Starting from a text language model backbone, Moshi generates speech as tokens from the residual quantizer of a neural audio codec, while modeling separately its own speech and that of the user into parallel streams. This allows for the removal of explicit speaker turns, and the modeling of arbitrary conversational dynamics.
Moshi also predicts time-aligned text tokens as a prefix to audio tokens. This “Inner
Monologue” method significantly improves the linguistic quality of generated speech and provides streaming speech recognition and text-to-speech. As a result, Moshi is the first real-time full-duplex spoken large language model, with a theoretical latency of 160ms, 200ms in practice. 
""")
    gr.Markdown("""
                - **Developed by:**  Kyutai
                - **Model type:** Multimodal speech-text foundation model
                - **Language(s) (NLP):** English
                - **License:** CC-BY""")

    gr.Markdown("### Model Sources ")
    gr.Markdown("""
                - **Repository:** [repo](https://github.com/kyutai-labs/moshi)
                - **Paper:** [paper](http://kyutai.org/Moshi.pdf)
                - **Demo:** [demo](https://moshi.chat/) """)


    gr.Markdown("""
                🚨 
                The Model will produce a lot of silence, because it is actually meant to stream the input and output.
                I will try to create a demo which works with the streaming.""")

    input_audio = gr.Audio(sources="microphone", label="Input Audio")
    output_audio = gr.Audio(label="Processed Audio", streaming=True, autoplay=True)
    stream = gr.State()

    log_out = gr.Textbox("Log", visible=False)

    input_audio.stop_recording(
        fn=process_audio,
        inputs=[input_audio, stream],
        outputs=[input_audio, output_audio, stream, log_out]
    )

    with gr.Row():
        with gr.Accordion("📙 Citation", open=False):
            gr.Textbox(
                value="""@techreport{kyutai2024moshi,
    author = {Alexandre D\'efossez and Laurent Mazar\'e and Manu Orsini and Am\'elie Royer and Patrick P\'erez and Herv\'e J\'egou and Edouard Grave and Neil Zeghidour},
    title = {Moshi: a speech-text foundation model for real-time dialogue},
    institution = {Kyutai},
    year={2024},
    month={September},
    url={http://kyutai.org/Moshi.pdf},
}
""", lines=7,
                label="Copy the BibTeX snippet to cite this source",
                elem_id="citation-button",
                show_copy_button=True,
            )

    
demo.launch(debug=True)