code/tt_diffrhythm/wire.py

"""Tensors between the host ACE-Step process and the TT DiT worker.

The two sides run different numpy/torch builds (ACE-Step venv vs the TT image), so tensors cross the socket as raw
little-endian bytes with dtype and shape, inside plain dicts that pickle identically on both sides."""
from __future__ import annotations

import torch

_DTYPES = {"float32": torch.float32}


def encode(t: torch.Tensor) -> dict:
    t = t.detach().to("cpu", torch.float32).contiguous()
    return {"dtype": "float32", "shape": list(t.shape), "data": t.numpy().tobytes()}


def decode(d: dict) -> torch.Tensor:
    return torch.frombuffer(bytearray(d["data"]), dtype=_DTYPES[d["dtype"]]).reshape(d["shape"])