Allow py 3.14, higher versions for a few packages

moshi/README.md

@@ -28,6 +28,10 @@ pip install -U -e "git+https://git@github.com/kyutai-labs/moshi#egg=moshi&subdir
 While we hope that the present codebase will work on Windows, we do not provide official support for it.
 At the moment, we do not support quantization for the PyTorch version, so you will need a GPU with a significant amount of memory (24GB).
 
+### About using Python 3.14
+
+If you want to use Python 3.14, you'll have to disable `torch.compile()` with the environment variable
+`NO_TORCH_COMPILE=1` because as the time of writing, `torch.compile()` does not support Python 3.14.
 
 ## Usage
 

moshi/moshi/client.py

@@ -21,6 +21,7 @@ def __init__(
         self,
         printer: AnyPrinter,
         websocket: aiohttp.ClientWebSocketResponse,
+        event_loop: asyncio.AbstractEventLoop,
         sample_rate: float = 24000,
         channels: int = 1,
         frame_size: int = 1920,
@@ -48,35 +49,26 @@ def __init__(
         self._opus_writer = sphn.OpusStreamWriter(sample_rate)
         self._opus_reader = sphn.OpusStreamReader(sample_rate)
         self._output_queue = queue.Queue()
-
-    async def _queue_loop(self) -> None:
-        while True:
-            if self._done:
-                return
-            await asyncio.sleep(0.001)
-            msg = self._opus_writer.read_bytes()
-            if len(msg) > 0:
-                try:
-                    await self.websocket.send_bytes(b"\x01" + msg)
-                except Exception as e:
-                    print(e)
-                    self._lost_connection()
-                    return
-
-    async def _decoder_loop(self) -> None:
-        all_pcm_data = None
-        while True:
-            if self._done:
+        self._all_pcm_data = None
+        self._event_loop = event_loop
+
+    async def _ws_send(self, msg) -> None:
+        if len(msg) > 0:
+            try:
+                await self.websocket.send_bytes(b"\x01" + msg)
+            except Exception as e:
+                print(e)
+                self._lost_connection()
                 return
-            await asyncio.sleep(0.001)
-            pcm = self._opus_reader.read_pcm()
-            if all_pcm_data is None:
-                all_pcm_data = pcm
-            else:
-                all_pcm_data = np.concatenate((all_pcm_data, pcm))
-            while all_pcm_data.shape[-1] >= self.frame_size:
-                self._output_queue.put(all_pcm_data[: self.frame_size])
-                all_pcm_data = np.array(all_pcm_data[self.frame_size :])
+
+    async def _decode(self, pcm) -> None:
+        if self._all_pcm_data is None:
+            self._all_pcm_data = pcm
+        else:
+            self._all_pcm_data = np.concatenate((self._all_pcm_data, pcm))
+        while self._all_pcm_data.shape[-1] >= self.frame_size:
+            self._output_queue.put(self._all_pcm_data[: self.frame_size])
+            self._all_pcm_data = np.array(self._all_pcm_data[self.frame_size :])
 
     async def _recv_loop(self) -> None:
         try:
@@ -102,7 +94,8 @@ async def _recv_loop(self) -> None:
                 kind = message[0]
                 if kind == 1:  # audio
                     payload = message[1:]
-                    self._opus_reader.append_bytes(payload)
+                    pcm = self._opus_reader.append_bytes(payload)
+                    await self._decode(pcm)
                     self.printer.print_pending()
                 elif kind == 2:  # text
                     payload = message[1:]
@@ -121,7 +114,9 @@ def _lost_connection(self) -> None:
 
     def _on_audio_input(self, in_data, frames, time_, status) -> None:
         assert in_data.shape == (self.frame_size, self.channels), in_data.shape
-        self._opus_writer.append_pcm(in_data[:, 0])
+        opus_bytes = self._opus_writer.append_pcm(in_data[:, 0])
+        if len(opus_bytes) > 0:
+            asyncio.run_coroutine_threadsafe(self._ws_send(opus_bytes), self._event_loop)
 
     def _on_audio_output(self, out_data, frames, time_, status) -> None:
         assert out_data.shape == (self.frame_size, self.channels), out_data.shape
@@ -136,9 +131,7 @@ def _on_audio_output(self, out_data, frames, time_, status) -> None:
 
     async def run(self) -> None:
         with self._in_stream, self._out_stream:
-            await asyncio.gather(
-                self._recv_loop(), self._decoder_loop(), self._queue_loop()
-            )
+            await self._recv_loop()
 
 
 async def run(printer: AnyPrinter, args):
@@ -167,7 +160,7 @@ async def run(printer: AnyPrinter, args):
         async with session.ws_connect(uri) as ws:
             printer.log("info", "connected!")
             printer.print_header()
-            connection = Connection(printer, ws)
+            connection = Connection(printer, ws, event_loop=asyncio.get_event_loop())
             await connection.run()
 
 

moshi/moshi/client_gradio.py

@@ -52,8 +52,7 @@ def receive(self, frame: tuple[int, NDArray]) -> None:
             self.ws = websockets.sync.client.connect(self.ws_url)
         _, array = frame
         array = array.squeeze().astype(np.float32) / 32768.0
-        self.stream_writer.append_pcm(array)
-        bytes = b"\x01" + self.stream_writer.read_bytes()
+        bytes = b"\x01" + self.stream_writer.append_pcm(array)
         self.ws.send(bytes)
 
     def generator(
@@ -65,8 +64,7 @@ def generator(
             kind = message[0]
             if kind == 1:
                 payload = message[1:]
-                self.stream_reader.append_bytes(payload)
-                pcm = self.stream_reader.read_pcm()
+                pcm = self.stream_reader.append_bytes(payload)
                 if self.all_output_data is None:
                     self.all_output_data = pcm
                 else:

moshi/moshi/server.py

@@ -71,104 +71,100 @@ def warmup(self):
 
         torch.cuda.synchronize()
 
-    async def handle_chat(self, request):
-        ws = web.WebSocketResponse()
-        await ws.prepare(request)
+    async def decode_and_send(
+        self,
+        tokens: torch.Tensor,
+        ws: web.WebSocketResponse,
+        opus_writer: sphn.OpusStreamWriter
+    ):
+        assert tokens.shape[1] == self.lm_gen.lm_model.dep_q + 1
+        main_pcm = self.mimi.decode(tokens[:, 1:])
+        main_pcm = main_pcm.cpu()
+        opus_bytes = opus_writer.append_pcm(main_pcm[0, 0].numpy())
+        if len(opus_bytes) > 0:
+            await ws.send_bytes(b"\x01" + opus_bytes)
+        text_token = tokens[0, 0, 0].item()
+        if text_token not in (0, 3):
+            _text = self.text_tokenizer.id_to_piece(text_token)  # type: ignore
+            _text = _text.replace("▁", " ")
+            msg = b"\x02" + bytes(_text, encoding="utf8")
+            log("info", f"text token '{_text}'")
+            await ws.send_bytes(msg)
 
-        async def recv_loop():
-            nonlocal close
-            try:
-                async for message in ws:
-                    if message.type == aiohttp.WSMsgType.ERROR:
-                        log("error", f"{ws.exception()}")
-                        break
-                    elif message.type == aiohttp.WSMsgType.CLOSED:
-                        break
-                    elif message.type != aiohttp.WSMsgType.BINARY:
-                        log("error", f"unexpected message type {message.type}")
-                        continue
-                    message = message.data
-                    if not isinstance(message, bytes):
-                        log("error", f"unsupported message type {type(message)}")
-                        continue
-                    if len(message) == 0:
-                        log("warning", "empty message")
+    async def recv_loop(
+        self,
+        ws: web.WebSocketResponse,
+        opus_reader: sphn.OpusStreamReader,
+        opus_writer: sphn.OpusStreamWriter
+    ):
+        all_pcm_data = None
+        skip_frames = 1
+        try:
+            async for message in ws:
+                if message.type == aiohttp.WSMsgType.ERROR:
+                    log("error", f"{ws.exception()}")
+                    break
+                elif message.type == aiohttp.WSMsgType.CLOSED:
+                    break
+                elif message.type != aiohttp.WSMsgType.BINARY:
+                    log("error", f"unexpected message type {message.type}")
+                    continue
+                message = message.data
+                if not isinstance(message, bytes):
+                    log("error", f"unsupported message type {type(message)}")
+                    continue
+                if len(message) == 0:
+                    log("warning", "empty message")
+                    continue
+                kind = message[0]
+                if kind == 1:  # audio
+                    payload = message[1:]
+                    pcm = opus_reader.append_bytes(payload)
+                    if pcm.shape[-1] == 0:
                         continue
-                    kind = message[0]
-                    if kind == 1:  # audio
-                        payload = message[1:]
-                        opus_reader.append_bytes(payload)
+                    if all_pcm_data is None:
+                        all_pcm_data = pcm
                     else:
-                        log("warning", f"unknown message kind {kind}")
-            finally:
-                close = True
-                log("info", "connection closed")
-
-        async def opus_loop():
-            all_pcm_data = None
-            skip_frames = 1
-
-            while True:
-                if close:
-                    return
-                await asyncio.sleep(0.001)
-                pcm = opus_reader.read_pcm()
-                if pcm.shape[-1] == 0:
-                    continue
-                if all_pcm_data is None:
-                    all_pcm_data = pcm
+                        all_pcm_data = np.concatenate((all_pcm_data, pcm))
+                    while all_pcm_data.shape[-1] >= self.frame_size:
+                        be = time.time()
+                        chunk = all_pcm_data[: self.frame_size]
+                        all_pcm_data = all_pcm_data[self.frame_size:]
+                        chunk = torch.from_numpy(chunk)
+                        chunk = chunk.to(device=self.device)[None, None]
+                        codes = self.mimi.encode(chunk)
+                        if skip_frames:
+                            # The first input audio frame is ignored, as from the point of
+                            # view of the model it is in the past. We still `mimi.encode` for simplicity,
+                            # however as the first encoded frame has a specific structure (due to the left padding),
+                            # we reset the streaming state of the encoder to reapply the padding on the next call.
+                            self.mimi.reset_streaming()
+                            skip_frames -= 1
+                        for c in range(codes.shape[-1]):
+                            tokens = self.lm_gen.step(codes[:, :, c: c + 1])
+                            if tokens is None:
+                                continue
+                            await self.decode_and_send(tokens, ws, opus_writer)
+                        log("info", f"frame handled in {1000 * (time.time() - be):.1f}ms")
                 else:
-                    all_pcm_data = np.concatenate((all_pcm_data, pcm))
-                while all_pcm_data.shape[-1] >= self.frame_size:
-                    be = time.time()
-                    chunk = all_pcm_data[: self.frame_size]
-                    all_pcm_data = all_pcm_data[self.frame_size:]
-                    chunk = torch.from_numpy(chunk)
-                    chunk = chunk.to(device=self.device)[None, None]
-                    codes = self.mimi.encode(chunk)
-                    if skip_frames:
-                        # The first input audio frame is ignored, as from the point of
-                        # view of the model it is in the past. We still `mimi.encode` for simplicity,
-                        # however as the first encoded frame has a specific structure (due to the left padding),
-                        # we reset the streaming state of the encoder to reapply the padding on the next call.
-                        self.mimi.reset_streaming()
-                        skip_frames -= 1
-                    for c in range(codes.shape[-1]):
-                        tokens = self.lm_gen.step(codes[:, :, c: c + 1])
-                        if tokens is None:
-                            continue
-                        assert tokens.shape[1] == self.lm_gen.lm_model.dep_q + 1
-                        main_pcm = self.mimi.decode(tokens[:, 1:])
-                        main_pcm = main_pcm.cpu()
-                        opus_writer.append_pcm(main_pcm[0, 0].numpy())
-                        text_token = tokens[0, 0, 0].item()
-                        if text_token not in (0, 3):
-                            _text = self.text_tokenizer.id_to_piece(text_token)  # type: ignore
-                            _text = _text.replace("▁", " ")
-                            msg = b"\x02" + bytes(_text, encoding="utf8")
-                            log("info", f"text token '{_text}'")
-                            await ws.send_bytes(msg)
-                    log("info", f"frame handled in {1000 * (time.time() - be):.1f}ms")
+                    log("warning", f"unknown message kind {kind}")
+        finally:
+            log("info", "connection closed")
 
-        async def send_loop():
-            while True:
-                if close:
-                    return
-                await asyncio.sleep(0.001)
-                msg = opus_writer.read_bytes()
-                if len(msg) > 0:
-                    await ws.send_bytes(b"\x01" + msg)
+    async def handle_chat(self, request):
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
 
         log("info", "accepted connection")
-        close = False
+
         async with self.lock:
             opus_writer = sphn.OpusStreamWriter(self.mimi.sample_rate)
             opus_reader = sphn.OpusStreamReader(self.mimi.sample_rate)
             self.mimi.reset_streaming()
             self.lm_gen.reset_streaming()
             # Send the handshake.
             await ws.send_bytes(b"\x00")
-            await asyncio.gather(opus_loop(), recv_loop(), send_loop())
+            await self.recv_loop(ws, opus_reader, opus_writer)
         log("info", "done with connection")
         return ws
 

moshi/moshi/utils/quantize.py

@@ -23,7 +23,7 @@ def __init__(self, linear: nn.Linear):
 
     def forward(self, x):
         import bitsandbytes as bnb  # type: ignore
-        state = bnb.MatmulLtState()
+        state = bnb.MatmulLtState()  # pyright: ignore
         state.CB = self.weight  # type: ignore
         assert isinstance(state.CB, torch.Tensor)
         state.SCB = self.weight_scb  # type: ignore
@@ -35,7 +35,7 @@ def forward(self, x):
                 "the model once initialized.")
         assert state.SCB.dtype == torch.float, state.SCB.dtype
         state.has_fp16_weights = False
-        y = bnb.matmul(x.half(), state.CB, state=state)
+        y = bnb.matmul(x.half(), state.CB, state=state)  # pyright: ignore
         assert isinstance(y, torch.Tensor)
         return y
 

moshi/pyproject.toml

@@ -1,19 +1,18 @@
 [project]
 name = "moshi"
-requires-python = ">= 3.10"
+requires-python = ">= 3.10,<3.15"
 description = "Moshi is moshi"
 dependencies = [
     "numpy >= 1.26, < 2.3",
-    "safetensors >= 0.4.0, < 0.7",
-    "huggingface-hub >= 0.24, < 0.34",
-    "bitsandbytes >= 0.45, < 0.46; sys_platform == 'linux'",
+    "safetensors >= 0.4.0, < 0.8.0",
+    "huggingface-hub >= 0.24, < 1.0.0",
+    "bitsandbytes >= 0.45, < 0.50.0; sys_platform == 'linux'",
     "einops >= 0.7, < 0.9",
-    "sentencepiece == 0.2",
+    "sentencepiece >= 0.2.0, < 0.3",
     "sounddevice == 0.5",
-    "sphn >= 0.1.4, < 0.2.0",
-    "torch >= 2.2.0, < 2.8",
+    "sphn >= 0.2.0, < 0.3.0",
+    "torch >= 2.2.0, < 2.10",
     "aiohttp>=3.10.5, <3.12",
-    "pytest >= 8.3.3",
     "tqdm >= 4.48",
 ]
 authors = [{name="Laurent Mazaré", email="laurent@kyutai.org"}]