#!/usr/bin/env python3 """mio_codec_wrapper.py Wrapper to run MioCodec in two modes: - streaming: connects to a websocket (e.g. vLLM/your LLM) that emits token embeddings (JSON messages) and decodes them to audio in near real-time using the windowed streamer. - offline: loads a JSONL file of embeddings and decodes to a WAV file. Requirements (minimal): torch, numpy, websockets (for streaming), soundfile (for offline write), sounddevice (optional playback). Install quickly without cloning full repo (if upstream exists): pip install git+https://github.com/Aratako/MioCodec.git If you already have local helper files (`streaming_istft_head.py`, `windowed_streamer.py`), the wrapper will import them; otherwise it will try to use `miocodec` package's model and swap in the streaming head. """ from __future__ import annotations import argparse import asyncio import json import sys from pathlib import Path import numpy as np import torch def build_args(): p = argparse.ArgumentParser() p.add_argument("--mode", choices=("stream", "offline"), default="stream") p.add_argument("--uri", default="ws://localhost:8000/stream", help="Websocket URI for streaming mode") p.add_argument("--input", help="JSONL file with embeddings for offline mode") p.add_argument("--out", default="out.wav", help="Output WAV file for offline mode") p.add_argument("--device", default="cpu") p.add_argument("--repo-id", default="Aratako/MioCodec-25Hz-24kHz") p.add_argument("--play", action="store_true", help="Play decoded audio while streaming") p.add_argument("--chunk-size", type=int, default=20) p.add_argument("--left", type=int, default=10) p.add_argument("--right", type=int, default=10) return p.parse_args() def load_model_and_streamer(repo_id: str, device: str, chunk_size: int, left: int, right: int): """Return (model, streamer, sample_rate). Attempts to import local windowed_streamer first.""" try: # Prefer local helper module if present from windowed_streamer import build_windowed_model, MioCodecWindowedStreamer, WindowedStreamConfig model = build_windowed_model(repo_id=repo_id, device=device) cfg = model.config stream_cfg = WindowedStreamConfig(chunk_size=chunk_size, left_context_tokens=left, right_lookahead_tokens=right) # create a default global embedding (zero) sized to model's expected dim gdim = getattr(cfg, "global_embedding_dim", getattr(cfg, "wave_decoder_dim", 512)) global_embedding = torch.zeros(gdim, dtype=torch.float32).to(device) streamer = MioCodecWindowedStreamer(model, global_embedding, config=stream_cfg) return model, streamer, getattr(cfg, "sample_rate", 24000) except Exception: # Fallback: import MioCodecModel from installed package and swap head from miocodec.model import MioCodecModel try: from streaming_istft_head import ISTFTHead as StreamingISTFTHead except Exception: StreamingISTFTHead = None model = MioCodecModel.from_pretrained(repo_id).to(device) cfg = model.config if StreamingISTFTHead is not None: old_head = model.istft_head new_head = StreamingISTFTHead(dim=cfg.wave_decoder_dim, n_fft=cfg.n_fft, hop_length=cfg.hop_length, padding=cfg.istft_padding).to(device) new_head.out.load_state_dict(old_head.out.state_dict()) model.istft_head = new_head # create a simple streamer-like helper using WindowedStreamConfig from windowed_streamer if available try: from windowed_streamer import MioCodecWindowedStreamer, WindowedStreamConfig stream_cfg = WindowedStreamConfig(chunk_size=chunk_size, left_context_tokens=left, right_lookahead_tokens=right) gdim = getattr(cfg, "global_embedding_dim", getattr(cfg, "wave_decoder_dim", 512)) global_embedding = torch.zeros(gdim, dtype=torch.float32).to(device) streamer = MioCodecWindowedStreamer(model, global_embedding, config=stream_cfg) except Exception: # Minimal fallback: use a naive one-shot decoder wrapper for full sequences streamer = None return model, streamer, getattr(cfg, "sample_rate", 24000) async def stream_mode(uri: str, model, streamer, samplerate: int, play: bool, device: str): """Connect to websocket; expect messages with {"embedding": [...]} or {"finished": true}.""" try: import websockets except Exception: print("Install websockets: pip install websockets", file=sys.stderr) return sd = None if play: try: import sounddevice as sd except Exception: sd = None async with websockets.connect(uri) as ws: print("connected", uri) async for raw in ws: try: msg = json.loads(raw) except Exception: continue if msg.get("finished"): if streamer is not None: final = streamer.finish() if final is not None: arr = final.cpu().numpy().astype(np.float32) if sd is not None: sd.play(arr, samplerate) sd.wait() else: print("final chunk ready (not played)") break emb = msg.get("embedding") if emb is None: continue tensor = torch.tensor(emb, dtype=torch.float32).to(device) if streamer is not None: audio = streamer.push_token(tensor) if audio is not None: arr = audio.cpu().numpy().astype(np.float32) if sd is not None: sd.play(arr, samplerate) # do not wait to allow continuous streaming else: sys.stdout.buffer.write(arr.tobytes()) else: # no streamer: buffer and decode once finished # simple approach: accumulate embeddings on model directly if model supports method if not hasattr(stream_mode, "buffer"): stream_mode.buffer = [] stream_mode.buffer.append(tensor.cpu().numpy()) def write_wav(path: Path, arr: np.ndarray, samplerate: int): try: import soundfile as sf sf.write(str(path), arr, samplerate) except Exception: # fallback: write raw float32 path.write_bytes(arr.astype(np.float32).tobytes()) def offline_mode(input_path: Path, out_path: Path, model, streamer, samplerate: int, device: str): """Read JSONL file of {embedding: [...]} lines and decode offline.""" embeddings = [] with input_path.open("r") as f: for line in f: try: m = json.loads(line) except Exception: continue if "embedding" in m: embeddings.append(torch.tensor(m["embedding"], dtype=torch.float32).to(device)) if streamer is not None: # push all and finish for emb in embeddings: streamer.push_token(emb) final = streamer.finish() if final is not None: arr = final.cpu().numpy().astype(np.float32) write_wav(out_path, arr, samplerate) print("Wrote", out_path) return # fallback one-shot: try to concatenate and run model-specific decode path if hasattr(model, "istft_head") and hasattr(model, "wave_prenet"): # build a content embedding and run forward once (best-effort) all_tokens = torch.stack(embeddings, dim=0).unsqueeze(0).to(device) with torch.no_grad(): local_latent = model.wave_prenet(all_tokens) if model.wave_conv_upsample is not None: local_latent = model.wave_conv_upsample(local_latent.transpose(1, 2)).transpose(1, 2) # approximate frames frames = local_latent.shape[1] * 100 # heuristic; prefer streamer for correctness local_latent = torch.nn.functional.interpolate(local_latent.transpose(1, 2), size=frames, mode=model.config.wave_interpolation_mode).transpose(1, 2) local_latent = model.wave_prior_net(local_latent.transpose(1, 2)).transpose(1, 2) local_latent = model.wave_decoder(local_latent, condition=torch.zeros(1, 1, local_latent.shape[-1], device=device)) local_latent = model.wave_post_net(local_latent.transpose(1, 2)).transpose(1, 2) audio = model.istft_head(local_latent).squeeze(0).cpu().numpy().astype(np.float32) write_wav(out_path, audio, samplerate) print("Wrote", out_path) return print("Could not decode offline: no streamer and model lacks decode path.") def main(): args = build_args() device = args.device model, streamer, samplerate = load_model_and_streamer(args.repo_id, device, args.chunk_size, args.left, args.right) if args.mode == "stream": asyncio.run(stream_mode(args.uri, model, streamer, samplerate, args.play, device)) else: if not args.input: print("--input required for offline mode", file=sys.stderr) return offline_mode(Path(args.input), Path(args.out), model, streamer, samplerate, device) if __name__ == "__main__": main()