| from __future__ import annotations
|
|
|
| from pathlib import Path
|
|
|
| import numpy as np
|
| import torch
|
| from snac import SNAC
|
| from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
| CODE_START_TOKEN_ID = 128257
|
| CODE_END_TOKEN_ID = 128258
|
| CODE_TOKEN_OFFSET = 128266
|
| SNAC_MIN_ID = 128266
|
| SNAC_MAX_ID = 156937
|
| SNAC_TOKENS_PER_FRAME = 7
|
|
|
| SOH_ID = 128259
|
| EOH_ID = 128260
|
| SOA_ID = 128261
|
| BOS_ID = 128000
|
| TEXT_EOT_ID = 128009
|
|
|
|
|
| def build_prompt(tokenizer, description: str, text: str) -> str:
|
| """Build formatted prompt for Maya1."""
|
| soh_token = tokenizer.decode([SOH_ID])
|
| eoh_token = tokenizer.decode([EOH_ID])
|
| soa_token = tokenizer.decode([SOA_ID])
|
| sos_token = tokenizer.decode([CODE_START_TOKEN_ID])
|
| eot_token = tokenizer.decode([TEXT_EOT_ID])
|
| bos_token = tokenizer.bos_token
|
|
|
| formatted_text = f'<description="{description}"> {text}'
|
|
|
| prompt = (
|
| soh_token + bos_token + formatted_text + eot_token +
|
| eoh_token + soa_token + sos_token
|
| )
|
|
|
| return prompt
|
|
|
|
|
| def extract_snac_codes(token_ids: list) -> list:
|
| """Extract SNAC codes from generated tokens."""
|
| try:
|
| eos_idx = token_ids.index(CODE_END_TOKEN_ID)
|
| except ValueError:
|
| eos_idx = len(token_ids)
|
|
|
| snac_codes = [
|
| token_id for token_id in token_ids[:eos_idx]
|
| if SNAC_MIN_ID <= token_id <= SNAC_MAX_ID
|
| ]
|
|
|
| return snac_codes
|
|
|
|
|
| def unpack_snac_from_7(snac_tokens: list) -> list:
|
| """Unpack 7-token SNAC frames to 3 hierarchical levels."""
|
| if snac_tokens and snac_tokens[-1] == CODE_END_TOKEN_ID:
|
| snac_tokens = snac_tokens[:-1]
|
|
|
| frames = len(snac_tokens) // SNAC_TOKENS_PER_FRAME
|
| snac_tokens = snac_tokens[:frames * SNAC_TOKENS_PER_FRAME]
|
|
|
| if frames == 0:
|
| return [[], [], []]
|
|
|
| l1, l2, l3 = [], [], []
|
|
|
| for i in range(frames):
|
| slots = snac_tokens[i*7:(i+1)*7]
|
| l1.append((slots[0] - CODE_TOKEN_OFFSET) % 4096)
|
| l2.extend([
|
| (slots[1] - CODE_TOKEN_OFFSET) % 4096,
|
| (slots[4] - CODE_TOKEN_OFFSET) % 4096,
|
| ])
|
| l3.extend([
|
| (slots[2] - CODE_TOKEN_OFFSET) % 4096,
|
| (slots[3] - CODE_TOKEN_OFFSET) % 4096,
|
| (slots[5] - CODE_TOKEN_OFFSET) % 4096,
|
| (slots[6] - CODE_TOKEN_OFFSET) % 4096,
|
| ])
|
|
|
| return [l1, l2, l3]
|
|
|
|
|
| def format_description(description: str) -> str:
|
| parts = description.strip().split("|")
|
| data = {}
|
|
|
|
|
| for part in parts:
|
| if ":" in part:
|
| key, value = part.split(":", 1)
|
| data[key.strip()] = value.strip()
|
|
|
|
|
| gender = data.get("gender", "")
|
| age_group = data.get("age_group", "")
|
| accent = data.get("accent", "")
|
| pitch = data.get("pitch", "")
|
| speed = data.get("speed", "")
|
| emotion = data.get("emotion", "")
|
| tone = data.get("tone", "")
|
|
|
|
|
| sentence1 = f"Realistic {gender} voice"
|
|
|
| if age_group == "senior":
|
| sentence1 += " in the 40s age"
|
| elif age_group == "adult":
|
| sentence1 += " in the 30s age"
|
| elif age_group == "young_adult":
|
| sentence1 += " in the 20s age"
|
| else:
|
| sentence1 += " in the 20s age"
|
|
|
| if accent:
|
| if accent.lower() == "us":
|
| accent = "American"
|
| elif accent.lower() == "uk":
|
| accent = "British"
|
| elif accent.lower() == "au":
|
| accent = "Australian"
|
| elif accent.lower() == "in":
|
| accent = "Indian"
|
| elif accent.lower() == "neutral":
|
| accent = "Asian American"
|
| elif accent.lower() == "other":
|
| accent = "American"
|
| sentence1 += f" with {accent.lower()} accent"
|
|
|
| sentence2_parts = []
|
| if pitch:
|
| sentence2_parts.append(f"{pitch.capitalize()} pitch")
|
| if emotion:
|
|
|
| if emotion.lower() == "happy":
|
| emotion = "energetic"
|
| elif emotion.lower() == "angry":
|
| emotion = "sarcastic"
|
| elif emotion.lower() == "calm":
|
| emotion = "neutral"
|
| elif emotion.lower() == "serious":
|
| emotion = "dry"
|
| elif emotion.lower() == "fearful":
|
| emotion = "sad"
|
| sentence2_parts.append(f"{emotion} timbre")
|
| if speed:
|
| if speed.lower() == "normal":
|
| speed = "conversational"
|
| sentence2_parts.append(f"{speed} pacing")
|
| if tone:
|
|
|
| if tone.lower() == "cold":
|
| tone = "harsh"
|
| elif tone.lower() == "friendly":
|
| tone = "warm"
|
| elif tone.lower() == "formal":
|
| tone = "smooth"
|
| elif tone.lower() == "casual":
|
| tone = "gravelly"
|
| elif tone.lower() == "authoritative":
|
| tone = "throaty"
|
| sentence2_parts.append(f"{tone} tone")
|
|
|
| sentence2 = ", ".join(sentence2_parts)
|
|
|
| return sentence1 + ". " + sentence2 + "."
|
|
|
|
|
| class Miner:
|
| """Vocence miner wrapper for Maya + SNAC inference."""
|
|
|
| def __init__(self, path_hf_repo: Path) -> None:
|
| self._repo_path = Path(path_hf_repo).resolve()
|
| self._device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
| self.model = AutoModelForCausalLM.from_pretrained(
|
| str(self._repo_path),
|
| torch_dtype=torch.bfloat16,
|
| device_map="auto",
|
| trust_remote_code=True,
|
| )
|
| self.tokenizer = AutoTokenizer.from_pretrained(
|
| str(self._repo_path),
|
| trust_remote_code=True,
|
| )
|
|
|
| snac_path = self._repo_path / "snac_model"
|
| if snac_path.exists():
|
| self.snac_model = SNAC.from_pretrained(str(snac_path)).eval()
|
| else:
|
| self.snac_model = SNAC.from_pretrained("snac_model").eval()
|
| if torch.cuda.is_available():
|
| self.snac_model = self.snac_model.to("cuda")
|
|
|
| def warmup(self) -> None:
|
| _ = self.generate_wav(
|
| instruction="| gender: male | pitch: mid | speed: normal | age_group: adult | emotion: calm | tone: formal | accent: us",
|
| text="This is a warmup utterance for the voice engine.",
|
| )
|
|
|
| def generate_wav(self, instruction: str, text: str) -> tuple[np.ndarray, int]:
|
| description = format_description(instruction)
|
| prompt = build_prompt(self.tokenizer, description, text)
|
|
|
| inputs = self.tokenizer(prompt, return_tensors="pt")
|
| if torch.cuda.is_available():
|
| inputs = {k: v.to("cuda") for k, v in inputs.items()}
|
|
|
| with torch.inference_mode():
|
| outputs = self.model.generate(
|
| **inputs,
|
| max_new_tokens=2048,
|
| min_new_tokens=28,
|
| temperature=0.4,
|
| top_p=0.9,
|
| repetition_penalty=1.1,
|
| do_sample=True,
|
| eos_token_id=CODE_END_TOKEN_ID,
|
| pad_token_id=self.tokenizer.pad_token_id,
|
| )
|
|
|
| generated_ids = outputs[0, inputs["input_ids"].shape[1] :].tolist()
|
| snac_tokens = extract_snac_codes(generated_ids)
|
| if len(snac_tokens) < SNAC_TOKENS_PER_FRAME:
|
| raise RuntimeError("Not enough SNAC tokens generated for decoding.")
|
|
|
| levels = unpack_snac_from_7(snac_tokens)
|
| codes_tensor = [
|
| torch.tensor(level, dtype=torch.long, device=self._device).unsqueeze(0)
|
| for level in levels
|
| ]
|
|
|
| with torch.inference_mode():
|
| z_q = self.snac_model.quantizer.from_codes(codes_tensor)
|
| audio = self.snac_model.decoder(z_q)[0, 0].cpu().numpy()
|
|
|
| if len(audio) > 2048:
|
| audio = audio[2048:]
|
|
|
| return audio.astype(np.float32), 24000
|
|
|
