from __future__ import annotations import threading from functools import cached_property from pathlib import Path from types import SimpleNamespace from typing import Any, Dict, List, Tuple import numpy as np SOH_ID = 128259 EOH_ID = 128260 CODE_START_TOKEN_ID = 128257 CODE_END_TOKEN_ID = 128258 TEXT_EOT_ID = 128009 CODE_TOKEN_OFFSET = 128266 SNAC_MIN_ID = 128266 SNAC_MAX_ID = 156937 SNAC_TOKENS_PER_FRAME = 7 SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz" VOCENCE_EMOTION_HINTS = { "happy": "cheerfully", "sad": "sadly", "angry": "firmly", "excited": "excitedly", "calm": "calmly", "neutral": "", } def _parse_vocence_instruction(instruction: str) -> Dict[str, str]: parts: Dict[str, str] = {} for segment in instruction.split("|"): segment = segment.strip() if ":" in segment: key, val = segment.split(":", 1) parts[key.strip().lower()] = val.strip().lower() return parts def _build_orpheus_prompt(voice: str, instruction: str, text: str) -> str: traits = _parse_vocence_instruction(instruction) hints: List[str] = [] emotion = traits.get("emotion", "") if emotion in VOCENCE_EMOTION_HINTS and VOCENCE_EMOTION_HINTS[emotion]: hints.append(VOCENCE_EMOTION_HINTS[emotion]) speed = traits.get("speed", "") if speed in ("slow", "very_slow"): hints.append("slowly") elif speed in ("fast", "very_fast"): hints.append("quickly") if instruction.strip() and not traits: hints.append(instruction.strip()[:120]) body = text.strip() or "Hello." if hints: body = f"{' '.join(hints)} {body}".strip() return f"{voice}: {body}" def _redistribute_codes(code_list: List[int]) -> List[Any]: import torch layer_1: List[int] = [] layer_2: List[int] = [] layer_3: List[int] = [] frames = len(code_list) // 7 for i in range(frames): layer_1.append(code_list[7 * i]) layer_2.append(code_list[7 * i + 1] - 4096) layer_3.append(code_list[7 * i + 2] - (2 * 4096)) layer_3.append(code_list[7 * i + 3] - (3 * 4096)) layer_2.append(code_list[7 * i + 4] - (4 * 4096)) layer_3.append(code_list[7 * i + 5] - (5 * 4096)) layer_3.append(code_list[7 * i + 6] - (6 * 4096)) return [ torch.tensor(layer_1).unsqueeze(0), torch.tensor(layer_2).unsqueeze(0), torch.tensor(layer_3).unsqueeze(0), ] class Miner: REPO_SENTINEL = "config.json" SETTINGS_FILE = "vocence_config.yaml" WARMUP_TIMEOUT = 600.0 def __init__(self, path_hf_repo: Path) -> None: self.root = Path(path_hf_repo).resolve() if not (self.root / self.REPO_SENTINEL).is_file(): raise FileNotFoundError(f"{self.REPO_SENTINEL} not present in {self.root}") _ = self.settings _ = self.model def __repr__(self) -> str: return f"" @cached_property def settings(self) -> SimpleNamespace: raw = self._load_yaml(self.root / self.SETTINGS_FILE) rt = raw.get("runtime") or {} gen = raw.get("generation") or {} lim = raw.get("limits") or {} return SimpleNamespace( voice=str(rt.get("voice", "Amir")), sample_rate=int(gen.get("sample_rate", 24000)), max_instruction_chars=int(lim.get("max_instruction_chars", 600)), max_text_chars=int(lim.get("max_text_chars", 2000)), max_new_tokens=int(gen.get("max_new_tokens", 1200)), temperature=float(gen.get("temperature", 0.6)), top_p=float(gen.get("top_p", 0.95)), repetition_penalty=float(gen.get("repetition_penalty", 1.1)), prefer_cuda=str(rt.get("device_preference", "cuda")).lower() == "cuda", prefer_bf16=str(rt.get("dtype", "bfloat16")).lower() == "bfloat16", ) @cached_property def model(self) -> SimpleNamespace: return self._instantiate_engine() def warmup(self) -> None: outcome: dict[str, Any] = {"done": False, "err": None} def _trial() -> None: try: self.generate_wav(instruction="Neutral voice.", text="Warming up.") outcome["done"] = True except Exception as exc: outcome["err"] = repr(exc) worker = threading.Thread(target=_trial, daemon=True) worker.start() worker.join(timeout=self.WARMUP_TIMEOUT) if not outcome["done"]: raise RuntimeError( f"warmup did not complete within {self.WARMUP_TIMEOUT}s: " f"{outcome['err'] or 'no completion signal'}" ) def generate_wav(self, instruction: str, text: str) -> Tuple[np.ndarray, int]: import torch s = self.settings instruction = instruction[: s.max_instruction_chars] text = text[: s.max_text_chars] prompt = _build_orpheus_prompt(s.voice, instruction, text) input_ids = self.model.tokenizer(prompt, return_tensors="pt").input_ids start_token = torch.tensor([[SOH_ID]], dtype=torch.int64) end_tokens = torch.tensor([[TEXT_EOT_ID, EOH_ID]], dtype=torch.int64) modified = torch.cat([start_token, input_ids, end_tokens], dim=1).to(self.model.device) with torch.inference_mode(): generated_ids = self.model.llm.generate( modified, max_new_tokens=s.max_new_tokens, do_sample=True, temperature=s.temperature, top_p=s.top_p, repetition_penalty=s.repetition_penalty, num_return_sequences=1, eos_token_id=CODE_END_TOKEN_ID, use_cache=True, ) row = generated_ids[0] token_indices = (row == CODE_START_TOKEN_ID).nonzero(as_tuple=True) if len(token_indices[0]) > 0: last_idx = token_indices[0][-1].item() cropped = row[last_idx + 1 :] else: cropped = row masked = cropped[cropped != CODE_END_TOKEN_ID] row_length = masked.size(0) new_length = (row_length // SNAC_TOKENS_PER_FRAME) * SNAC_TOKENS_PER_FRAME trimmed = masked[:new_length] if trimmed.size(0) < SNAC_TOKENS_PER_FRAME: raise ValueError("orpheus-amir produced insufficient SNAC tokens") code_list = [(int(t) - CODE_TOKEN_OFFSET) for t in trimmed.tolist()] codes = _redistribute_codes(code_list) codes = [c.to(self.model.device) for c in codes] audio_hat = self.model.snac.decode(codes) wave = audio_hat.detach().squeeze().cpu().numpy().astype(np.float32) return wave, s.sample_rate def _instantiate_engine(self) -> SimpleNamespace: import torch from snac import SNAC from transformers import AutoModelForCausalLM, AutoTokenizer s = self.settings cuda_ready = torch.cuda.is_available() device = "cuda:0" if (s.prefer_cuda and cuda_ready) else "cpu" dtype = torch.bfloat16 if (s.prefer_bf16 and cuda_ready) else torch.float32 model_name = str(self.root) print("[Miner] Loading Orpheus tokenizer...") tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) print(f"[Miner] Loading Orpheus model ({dtype})...") llm = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=dtype, device_map=device, trust_remote_code=True, ) print("[Miner] Loading SNAC 24kHz decoder...") snac_dir = self.root / "snac" if not snac_dir.is_dir(): raise FileNotFoundError(f"snac/ not present in {self.root} (required for no-egress deploy)") model_name = str(snac_dir) snac = SNAC.from_pretrained(model_name) snac.to(device) print(f"[Miner] orpheus-amir ready :: device={device} dtype={dtype}") return SimpleNamespace(llm=llm, tokenizer=tokenizer, snac=snac, device=device) @staticmethod def _load_yaml(path: Path) -> dict[str, Any]: if not path.is_file(): return {} from yaml import safe_load with path.open("r", encoding="utf-8") as fh: return safe_load(fh) or {}