| """Load the actual Small export and reuse the existing channel-wise CFG loop.""" |
| import json |
| import sys |
| import threading |
| import time |
| from contextlib import nullcontext |
|
|
| import numpy as np |
| import torch |
| from huggingface_hub import snapshot_download |
|
|
| import settings |
|
|
| sys.path.append(str(settings.PARENT)) |
| from tts_engine import TTSEngine |
|
|
|
|
| def zero_temperature_sampler(sampler): |
| """The upstream sampler requires positive T; T=0 means genuine argmax.""" |
| def sample(*args, **kwargs): |
| if float(kwargs.get("temperature", 1)) == 0: |
| kwargs["do_sample"] = False |
| return sampler(*args, **kwargs) |
| return sample |
|
|
|
|
| class SmallEngine(TTSEngine): |
| def __init__(self): |
| started = time.monotonic() |
| self.device = settings.DEVICE |
| self.lock = threading.Lock() |
| self.refs = {} |
| self.last_codes = [] |
| self.last_performance = None |
|
|
| root = snapshot_download(settings.REPO, revision=settings.REVISION, |
| allow_patterns=["code/*.py", "assets/**", f"checkpoints/{settings.STAGE}/*", "README.md", "LICENSE"]) |
| from pathlib import Path |
| root = Path(root) |
| sys.path.insert(0, str(root / "code")) |
| import moss_small |
| from large_talker import build_fresh |
| from packing import ScorePacker |
| moss_small.SFT3 = str(root / "assets/sft3") |
| moss_small.QWEN = str(root / "assets/qwen3") |
| moss_small.load_qwen_backbone = lambda model, log=print: None |
| schema = json.loads((root / "assets/score_schema.json").read_text()) |
| print(f"[small] loading {settings.STAGE} ({settings.REVISION[:8]})", flush=True) |
| self.model, cfg = build_fresh(schema, log=lambda _: None) |
| state = torch.load(root / "checkpoints" / settings.STAGE / "model_bf16.pt", map_location="cpu", weights_only=True) |
| self.model.load_state_dict(state, strict=True) |
| self.model.tie_weights() |
| self.model = self.model.to(self.device, dtype=torch.bfloat16).eval() |
| self.model._sample_next_token = zero_temperature_sampler(self.model._sample_next_token) |
| del state |
| _, _, Processor = moss_small.export_classes() |
| self.proc = Processor.from_pretrained(moss_small.SFT3, |
| codec_path="OpenMOSS-Team/MOSS-Audio-Tokenizer-v2", |
| codec_weight_dtype="bf16" if settings.FAST_INFERENCE else "fp32", codec_compute_dtype="bf16") |
| self.proc.audio_tokenizer = self.proc.audio_tokenizer.to(self.device).eval() |
| self.packer = ScorePacker(self.proc, cfg, schema) |
| self.sr = int(self.proc.model_config.sampling_rate) |
| self.parameters = sum(p.numel() for p in self.model.parameters()) |
| from performance_core import FastGenerator |
| self.fast = FastGenerator(self) |
| if settings.FAST_INFERENCE and settings.COMPILE_GLOBAL: |
| |
| |
| |
| from torch import _dynamo |
| _dynamo.config.cache_size_limit = 64 |
| self.model.transformer.forward = torch.compile( |
| self.model.transformer.forward, dynamic=True, mode="default") |
| print(f"[small] ready: {self.parameters:,} parameters, {self.sr} Hz, {time.monotonic()-started:.1f}s", flush=True) |
|
|
| @torch.inference_mode() |
| def warm_search(self): |
| """Compile continuation/cache-fork paths before advertising ready.""" |
| from beam_search import ChunkSearch |
| from prompting import build |
| class WarmScorer: |
| device = "none (startup warm-up)" |
| last_timing = {"asr_seconds":0., "quality_seconds":0., "total_seconds":0.} |
| def score(self, waves, text, final): |
| return [{"reward":float(len(w)), "prefix_wer":0.} for w in waves] |
| spec = build({"text":"Ready when you are.","caption":"Warm conversation."},"caption","en",seconds=6.6) |
| path = settings.HERE / "performance/benchmark-reference.wav" |
| for n in [4,8,12,16]: |
| ChunkSearch(self,WarmScorer()).generate(spec,n=n,beam_width=1,chunk_frames=10, |
| max_seconds=6.6,seed=991,reference_path=path if path.is_file() else None) |
| self.last_codes=[] |
| self.last_performance=None |
|
|
| @staticmethod |
| def _sample_continue_or_end(mdl, lh, p, stop_bias=None): |
| |
| return mdl._sample_next_assistant_text_token(local_hidden_states=lh, |
| do_sample=p["text_temperature"] > 0, temperature=p["text_temperature"], top_k=p["text_top_k"], top_p=p["text_top_p"]) |
|
|
| @torch.inference_mode() |
| def reference(self, path): |
| key = str(path) |
| if key not in self.refs: |
| codes = self.proc.encode_audios_from_path(path, n_vq=12)[0].cpu() |
| if not 0 < len(codes) <= 37: |
| raise ValueError("Referenzaudio muss höchstens 37 Codec-Frames lang sein.") |
| self.refs[key] = codes |
| return self.refs[key] |
|
|
| def forget_reference(self, path): |
| self.refs.pop(str(path), None) |
|
|
| def pack(self, prompt, spec, reference, n): |
| record = {"prompt": prompt, "text": spec["text"], "frames": spec["frames"], "lang": spec["lang"]} |
| row = self.packer.pack_mode(record, [], "reference" if reference is not None else "instruction", reference, generation=True) |
| batch = self.packer.collate([row] * n) |
| return batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device) |
|
|
| def pack_many(self, specs, reference, neutral=False): |
| from performance_core import combine_branches |
| |
| |
| rows = [self.pack(s["unconditional"] if neutral else s["prompt"], s, reference, 1) for s in specs] |
| return combine_branches(rows, self.model.config.pad_token_id, self.model.config.audio_pad_token_id) |
|
|
| @torch.inference_mode() |
| def generate_many(self, specs, guidance=1.5, seed=777, reference_path=None, |
| temperature=1.0, progress=None, streaming=False, batch_size=None, chunk_frames=10, |
| max_frames=None): |
| from performance_core import StreamingCodec |
| waves, self.last_codes, batches = [], [], [] |
| with self.lock, torch.cuda.device(self.device), torch.autocast("cuda", dtype=torch.bfloat16): |
| reference = self.reference(reference_path) if reference_path else None |
| batch_size = batch_size or settings.BATCH_SIZE |
| for offset in range(0, len(specs), batch_size): |
| group = specs[offset:offset + batch_size] |
| count = len(group) |
| branches = [self.pack_many(group, reference)] |
| if guidance != 1 and any(s["prompt"] != s["unconditional"] for s in group): |
| branches.append(self.pack_many(group, reference, neutral=True)) |
| torch.manual_seed(seed + offset) |
| torch.cuda.manual_seed_all(seed + offset) |
| p = {"max_new_tokens": min(settings.MAX_FRAMES + 60, max(s["frames"] for s in group) + 60), |
| "audio_temperature": temperature, "audio_top_p": .95, "audio_top_k": 50, |
| "audio_repetition_penalty": 1., "text_temperature": 0. if temperature == 0 else 1., |
| "text_top_p": 1., "text_top_k": 50} |
| if max_frames is not None: |
| p["max_new_tokens"] = min(p["max_new_tokens"], max_frames) |
| torch.cuda.synchronize(self.device) |
| start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] |
| capture_before = self.fast.capture_seconds |
| wall_start = time.perf_counter() |
| start.record() |
| decoder = self.proc.audio_tokenizer |
| context = StreamingCodec(decoder, count, self.model.config.audio_pad_token_id) if streaming else nullcontext() |
| outputs, frames, previous = [], [], 0 |
| first = None |
| with context as incremental: |
| iterator = (self.fast.stream_frames(branches, p, chunk_frames=chunk_frames if streaming else p["max_new_tokens"]+1, |
| guidance=guidance, graph=True) if settings.FAST_INFERENCE else |
| self._stream_frames_cfg(branches, p, chunk_frames=chunk_frames if streaming else p["max_new_tokens"]+1, |
| guidance=guidance)) |
| for frames in iterator: |
| if streaming: |
| outputs.append(incremental.decode(torch.stack(frames[previous:], dim=1))) |
| previous = len(frames) |
| if first is None: |
| first = torch.cuda.Event(enable_timing=True) |
| first.record() |
| if not frames: |
| raise RuntimeError("Das Modell hat keinen Audio-Frame erzeugt.") |
| codes = torch.stack(frames, dim=1) |
| valid = codes[:, :, 0].ne(self.model.config.audio_pad_token_id) |
| if not streaming: |
| result = decoder.decode(codes.masked_fill(~valid.unsqueeze(-1), 0).permute(2, 0, 1).contiguous(), |
| padding_mask=valid, num_quantizers=12, return_dict=True, chunk_duration=8) |
| outputs.append((result.audio.float().mean(1), result.audio_lengths)) |
| end.record() |
| end.synchronize() |
| seconds = start.elapsed_time(end) / 1000 |
| durations = valid.sum(-1).cpu().tolist() |
| batches.append({"batch_size": count, "device_seconds": seconds, |
| "audio_seconds": [f / 12.5 for f in durations], |
| "model_wall_seconds": time.perf_counter() - wall_start, |
| "first_audio_seconds": start.elapsed_time(first) / 1000 if first else seconds, |
| "graph_capture_seconds": self.fast.capture_seconds - capture_before}) |
| for i, c in enumerate(codes): |
| c = c[valid[i]] |
| if len(c) < 2: |
| raise RuntimeError("Ein Take enthält weniger als zwei Audio-Frames.") |
| self.last_codes.append(c.cpu()) |
| pieces = [pcm[i, :int(lengths[i])].cpu() for pcm, lengths in outputs if int(lengths[i]) > 0] |
| wave = torch.cat(pieces).numpy().reshape(-1) |
| if not np.isfinite(wave).all(): |
| raise RuntimeError("Nicht-endliche Audiowerte im Take.") |
| waves.append(wave) |
| if progress: |
| progress(len(waves), len(specs)) |
| total_device = sum(b["device_seconds"] for b in batches) |
| total_audio = sum(sum(b["audio_seconds"]) for b in batches) |
| self.last_performance = {"device_seconds": total_device, "raw_audio_seconds": total_audio, |
| "aggregate_rtf": total_device / total_audio, "audio_seconds_per_second": total_audio / total_device, |
| "batches": batches, "streaming": streaming, "chunk_frames": chunk_frames if streaming else None, |
| "backend": "fused CFG + local CUDA graphs + compiled global" if settings.FAST_INFERENCE and settings.COMPILE_GLOBAL |
| else "fused CFG + local CUDA graphs" if settings.FAST_INFERENCE else "legacy", |
| "includes": "TTS prefill, autoregressive decode and PCM codec; GPU event span including launch gaps", |
| "excludes": "reference encoding, prompt packing, Luna, alignment, scoring, CPU audio copy, MP3, network", |
| "graph_capture_seconds": sum(b["graph_capture_seconds"] for b in batches)} |
| return waves |
|
|
| @torch.inference_mode() |
| def warm_up(self): |
| from prompting import build |
| spec = build({"text": "Ready when you are.", "caption": "Warm, natural conversation."}, "caption", "en", seconds=2.) |
| path = settings.HERE / "performance/benchmark-reference.wav" |
| references = [None] |
| with self.lock, torch.cuda.device(self.device), torch.autocast("cuda", dtype=torch.bfloat16): |
| if path.is_file(): |
| references.append(self.reference(path)) |
| for n in [1, settings.BATCH_SIZE]: |
| for reference in references: |
| branches = [self.pack(spec["prompt"], spec, reference, n), |
| self.pack(spec["unconditional"], spec, reference, n)] |
| params = {"max_new_tokens":20,"audio_temperature":1.,"audio_top_p":.95,"audio_top_k":50, |
| "audio_repetition_penalty":1.,"text_temperature":1.,"text_top_p":1.,"text_top_k":50} |
| frames=[] |
| for frames in self.fast.stream_frames(branches,params,chunk_frames=21,guidance=1.5,min_frames=20,graph=True): |
| pass |
| codes=torch.stack(frames,dim=1).permute(2,0,1).contiguous() |
| self.proc.audio_tokenizer.decode(codes,return_dict=True,chunk_duration=8) |
| torch.cuda.synchronize(self.device) |
| self.last_codes = [] |
| self.last_performance = None |
|
|
| @torch.inference_mode() |
| def generate(self, spec, n=1, guidance=1.5, seed=777, reference_path=None, temperature=1.0, progress=None, streaming=False, chunk_frames=10, |
| max_frames=None): |
| batch_size, retries = min(n, settings.BATCH_SIZE), [] |
| while True: |
| try: |
| waves = self.generate_many([spec] * n, guidance, seed, reference_path, temperature, progress, |
| streaming=streaming, batch_size=batch_size, chunk_frames=chunk_frames, |
| max_frames=max_frames) |
| self.last_performance["oom_batch_retries"] = retries |
| return waves |
| except torch.cuda.OutOfMemoryError: |
| |
| |
| |
| self.fast.clear_graphs() |
| torch.cuda.empty_cache() |
| if batch_size == 1: |
| raise |
| retries.append(batch_size) |
| batch_size = max(1, batch_size // 2) |
|
|
| @torch.inference_mode() |
| def generate_legacy(self, spec, n=1, guidance=1.5, seed=777, reference_path=None, temperature=1.0, progress=None): |
| waves, self.last_codes = [], [] |
| with self.lock, torch.autocast("cuda", dtype=torch.bfloat16): |
| reference = self.reference(reference_path) if reference_path else None |
| for offset in range(0, n, 4): |
| count = min(4, n - offset) |
| branches = [self.pack(spec["prompt"], spec, reference, count)] |
| if guidance != 1.0 and spec["prompt"] != spec["unconditional"]: |
| branches.append(self.pack(spec["unconditional"], spec, reference, count)) |
| torch.manual_seed(seed + offset) |
| torch.cuda.manual_seed_all(seed + offset) |
| p = {"max_new_tokens": min(settings.MAX_FRAMES + 60, spec["frames"] + 60), |
| "audio_temperature": temperature, "audio_top_p": 0.95, |
| "audio_top_k": 50, "audio_repetition_penalty": 1.0, |
| "text_temperature": 0.0 if temperature == 0 else 1.0, "text_top_p": 1.0, "text_top_k": 50} |
| frames = [] |
| for frames in self._stream_frames_cfg(branches, p, chunk_frames=p["max_new_tokens"] + 1, guidance=guidance): |
| pass |
| if not frames: |
| raise RuntimeError("Das Modell hat keinen Audio-Frame erzeugt.") |
| codes = torch.stack(frames, dim=1) |
| for c in codes: |
| c = c[c[:, 0] != self.model.config.audio_pad_token_id] |
| if len(c) < 2: |
| raise RuntimeError("Ein Take enthält weniger als zwei Audio-Frames.") |
| self.last_codes.append(c.cpu()) |
| decoded = self.proc.decode_audio_codes([c], return_stereo=False)[0] |
| wave = decoded.float().cpu().numpy().reshape(-1) |
| if not np.isfinite(wave).all(): |
| raise RuntimeError("Nicht-endliche Audiowerte im Take.") |
| waves.append(wave) |
| if progress: |
| progress(len(waves), n) |
| return waves |
|
|
|
|
| class Scorers: |
| """Use the same proven Judge/Parakeet/VoiceCLAP ranker, without its LoRAs.""" |
| def __init__(self): |
| from asr_engine import ParakeetASR |
| from bestofn import Judge |
| from retrieval import Retriever |
| self.asr = ParakeetASR(device=settings.SCORE_DEVICE) |
| self.retriever = Retriever(device=settings.SCORE_DEVICE) |
| self.retriever._load() |
| self.judge = Judge(device=settings.SCORE_DEVICE, asr=self.asr) |
| if not self.judge.ok or not self.retriever.ok: |
| raise RuntimeError("Die Best-of-N-Scorer konnten nicht geladen werden.") |
| self.judge.attach_clap(self.retriever) |
| if not self.judge._load_heads(): |
| raise RuntimeError("Die Burst-Detektorköpfe fehlen.") |
| |
| |
| |
| self.judge.gen.embedder = self.retriever.model |
| self.judge.bl.embedder = self.retriever.model |
|
|
| def score(self, waves, spec): |
| from bestofn import rank |
| scores = self.judge.score(waves, 48000, spec["text"], general=spec["general"], |
| script=spec["tagged"], tagged=spec["tagged"]) |
| from metrics import word_errors |
| for score in scores: |
| score["legacy_word_distance"] = score["wer"] |
| score.update(word_errors(spec["text"], score["hyp"])) |
| rank(scores) |
| return scores |
|
|
