初始化项目，由ModelHub XC社区提供模型
Model: RL-gang/vocence-orpheus-amir Source: Original Platform
This commit is contained in:
ModelHub XC
232
miner.py
Normal file
232
miner.py
Normal file
@@ -0,0 +1,232 @@
|
||||
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"<Miner root={self.root.name} engine=orpheus-amir>"
|
||||
|
||||
@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 {}
|
||||
Reference in New Issue
Block a user