mio_codec_wrapper.py

#!/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()