Initial commit - Orpheus ComfyUI mode via LMStudio stream

node/decoder.py

@@ -0,0 +1,72 @@
+import torch
+import numpy as np
+from snac import SNAC
+
+# One global SNAC, placed on GPU if available.
+_SNAC = None
+_DEVICE = None
+
+def _init_snac():
+    global _SNAC, _DEVICE
+    if _SNAC is not None:
+        return
+    if torch.cuda.is_available():
+        _DEVICE = torch.device("cuda")
+    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        _DEVICE = torch.device("cpu")  # MPS path not stable for SNAC; keep CPU.
+    else:
+        _DEVICE = torch.device("cpu")
+
+    _SNAC = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to(_DEVICE)
+
+def get_snac_device():
+    _init_snac()
+    return str(_DEVICE)
+
+def convert_to_audio(multiframe, count):
+    """
+    multiframe: list[int] flat, length multiple of 7 (codebook layout)
+    Returns: bytes of int16 PCM at 24kHz, or None if invalid.
+    """
+    _init_snac()
+
+    if len(multiframe) < 7:
+        return None
+
+    # Lay out codebooks: c0 (1), c1 (2), c2 (4) per 7-tuple frame
+    num_frames = len(multiframe) // 7
+    frame = multiframe[: num_frames * 7]
+
+    c0 = torch.tensor([], device=_DEVICE, dtype=torch.int32)
+    c1 = torch.tensor([], device=_DEVICE, dtype=torch.int32)
+    c2 = torch.tensor([], device=_DEVICE, dtype=torch.int32)
+
+    for j in range(num_frames):
+        i = 7 * j
+        # c0: [i]
+        c0 = torch.cat([c0, torch.tensor([frame[i]], device=_DEVICE, dtype=torch.int32)])
+        # c1: [i+1, i+4]
+        c1 = torch.cat([c1,
+                        torch.tensor([frame[i + 1]], device=_DEVICE, dtype=torch.int32),
+                        torch.tensor([frame[i + 4]], device=_DEVICE, dtype=torch.int32)])
+        # c2: [i+2, i+3, i+5, i+6]
+        c2 = torch.cat([c2,
+                        torch.tensor([frame[i + 2]], device=_DEVICE, dtype=torch.int32),
+                        torch.tensor([frame[i + 3]], device=_DEVICE, dtype=torch.int32),
+                        torch.tensor([frame[i + 5]], device=_DEVICE, dtype=torch.int32),
+                        torch.tensor([frame[i + 6]], device=_DEVICE, dtype=torch.int32)])
+
+    # Bounds check
+    for cb in (c0, c1, c2):
+        if torch.any(cb < 0) or torch.any(cb > 4096):
+            return None
+
+    codes = [c0.unsqueeze(0), c1.unsqueeze(0), c2.unsqueeze(0)]
+
+    with torch.inference_mode():
+        audio_hat = _SNAC.decode(codes)  # (B, 1, n)
+    # Empirically use middle slice for stable chunking (as in upstream)
+    audio_slice = audio_hat[:, :, 2048:4096]
+    audio_np = audio_slice.detach().cpu().numpy()
+    pcm16 = (audio_np * 32767.0).astype(np.int16).tobytes()
+    return pcm16