"""Ported from FashionFlora/ogg-tts (Apache-2.0), ogg_tts/infer_omnifish.py: block diffusion with KV caches. Kept as it is: prefix_ids, voice_context, _Stream, blocks_cached, speech_frames. Left out: the CLI, streaming and checkpoint loading (nodes.py does those). `of` is the ported omnifish module.\n""" import math import torch import torch.nn.functional as F from . import model as of def prefix_ids(tok, text, lang, before=(), after=(), speaker=None, quality=None): """before / after: up to 2 sentences of text context each (not spoken); speaker: token name ("pl_3"), or the language bag ("pl_x", only with a prompt), or "dialogue" (text from of.dialogue_text, no prompt), or None""" C = of.NUM_CODEBOOKS instruct = "dialogue" if speaker == "dialogue" else f"speaker:{speaker}" if speaker else "" instruct = f"{instruct} quality:{quality}".strip() if quality else instruct # recording id, e.g. q35 instruct = instruct or "None" style = tok(f"<|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>", return_tensors="pt").input_ids txt = tok(of.context_text(text, list(before), list(after)), return_tensors="pt").input_ids return torch.cat([style, txt], 1).repeat(C, 1) def voice_context(tok, voices, text, lang, before=(), after=(), quality=None): """dialogue with defined voices: voices [(name, transcript, codes [10, T]), ...], text from of.dialogue_text([(name, line), ...]) -> context (ids [10, L], audio mask [L]) laid out as in training: [style][S1: transcript][S1 audio][S2: transcript][S2 audio]...[scene text]""" C = of.NUM_CODEBOOKS instruct = f"dialogue quality:{quality}" if quality else "dialogue" parts = [(tok(f"<|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>", return_tensors="pt").input_ids.repeat(C, 1), False)] for name, transcript, codes in voices: parts += [(tok(f"<|quad_start|>{name}<|quad_end|>{transcript.strip()}", return_tensors="pt").input_ids.repeat(C, 1), False), (codes.cpu().long(), True)] parts.append((tok(of.context_text(text, list(before), list(after)), return_tensors="pt").input_ids.repeat(C, 1), False)) return (torch.cat([x for x, _ in parts], 1), torch.cat([torch.full((x.shape[1],), f, dtype=torch.bool) for x, f in parts])) class _Stream: """KV-cached sequences, one batch row per context ((ids [10, L], audio mask [L]): the conditional one and the empty CFG one), left-padded to a common length: the context (text + prompt) and the finished clean blocks are computed once; the current block is a query of `block` frames against the cache. Every allowed key of a block query is in the cache or in the block itself (clean/noisy block b sees context, clean blocks < b and itself), so a row only masks its padding. One forward per unmasking iteration serves both CFG branches.""" def __init__(self, model, contexts, dev): from transformers import DynamicCache self.model, self.dev, self.cache, self.done = model, dev, DynamicCache(), 0 C, B = of.NUM_CODEBOOKS, len(contexts) self.ctx = [c.shape[1] for c, _ in contexts] T = max(self.ctx) self.pad = [T - c for c in self.ctx] if T == 0: return ids = torch.zeros(B, C, T, dtype=torch.long, device=dev) am = torch.zeros(B, T, dtype=torch.bool, device=dev) mask = torch.zeros(B, 1, T, T, dtype=torch.bool, device=dev) pos = torch.zeros(B, T, dtype=torch.long, device=dev) for r, (c, a) in enumerate(contexts): k = self.pad[r] ids[r, :, k:] = c.to(dev) am[r, k:] = a.to(dev) mask[r, 0, k:, k:] = True # context sees context mask[r, 0, :k, :k] = torch.eye(k, dtype=torch.bool, device=dev) # padding rows: themselves only pos[r, k:] = torch.arange(self.ctx[r], device=dev) self.ctx_h = self._run(ids, am, mask, pos) # context states (the MTP reads the prompt frames' ones) def _run(self, ids, am, mask, pos): with torch.autocast(self.dev.type, torch.bfloat16): return self.model.llm(inputs_embeds=self.model.embed(ids, am), attention_mask=mask, position_ids=pos, past_key_values=self.cache, use_cache=True, return_dict=True)[0] def _block(self, semantic, full=None): B, L, past = len(self.ctx), semantic.shape[0], self.cache.get_seq_length() ids = torch.full((B, of.NUM_CODEBOOKS, L), of.MASK_ID, dtype=torch.long, device=self.dev) ids[:, 0] = semantic if full is not None: # a decided block with its residual codes (MASK where END / PAD) ids[:] = full am = torch.ones(B, L, dtype=torch.bool, device=self.dev) mask = torch.ones(B, 1, L, past + L, dtype=torch.bool, device=self.dev) pos = torch.empty(B, L, dtype=torch.long, device=self.dev) for r in range(B): mask[r, 0, :, :self.pad[r]] = False pos[r] = self.ctx[r] + self.done + torch.arange(L, device=self.dev) return ids, am, mask, pos def logits(self, cur): """semantic logits [rows, frames, vocab] of the noisy block `cur`; the cache is left as it was""" h = self._run(*self._block(cur)) self.cache.crop(-cur.shape[0]) # drop the noisy block again head = self.model.semantic_head return head(h.to(head.weight.dtype)).float() def peek(self, cur): """row 0's hidden states of the decided block with semantic ids only (the MTP input); cache unchanged""" h = self._run(*self._block(cur))[0] self.cache.crop(-cur.shape[0]) return h def commit(self, full): """append the decided block with all 10 codes [10, block] as context of the later blocks""" self._run(*self._block(full[0], full)) self.done += full.shape[1] @torch.no_grad() def blocks_cached(model, tok, prompt_codes, prompt_text, text, lang, block=16, steps=6, cfg=1.5, temperature=0.8, max_frames=None, seed=0, before=(), after=(), speaker=None, quality=None, mtp_temperature=0.0, mtp_top_k=0, mtp_rounds=4, full_context=True, voices=None, min_frames=0, semantic_left=(), mtp_prompt=True, target_prefix=None): """block diffusion with KV caches (one batch row for the conditional and one for the CFG branch). Per block: unmask its semantic ids in `steps` iterations, run the MTP on it (semantic-only states), append it with all 10 codes (full_context, for models trained so) or its semantic ids only (default) to the cache as context of the later blocks. Yields (semantic ids of every frame so far, codes [10, frames]) after each block; the last block holds END (or a PAD), or the safety cap is reached.""" dev = next(model.parameters()).device g = torch.Generator(device=dev).manual_seed(seed) if mtp_temperature > 0: torch.manual_seed(seed) if voices: # a dialogue with defined voices (name, transcript, codes): prompt_codes / prompt_text unused frames, chars = sum(v[2].shape[1] for v in voices), sum(len(v[1]) for v in voices) context = voice_context(tok, voices, text, lang, before, after, quality) else: if prompt_codes is None: prompt_codes, prompt_text = torch.zeros(of.NUM_CODEBOOKS, 0, dtype=torch.long), "" frames, chars = prompt_codes.shape[1], len(prompt_text) if semantic_left: instruct = f"speaker:{speaker}" if speaker else "" instruct = f"{instruct} quality:{quality}".strip() if quality else instruct instruct = instruct or "None" style = tok(f"<|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>", return_tensors="pt").input_ids.repeat(of.NUM_CODEBOOKS, 1) parts, audio_parts = [style], [torch.zeros(style.shape[1], dtype=torch.bool)] # as in training: the left lines as text only, in one <|object_ref|> span (their sound comes with the # prompt); entries are texts or (text, codes) pairs, codes unused; as many lines as the model was trained with left_text = " ".join(x if isinstance(x, str) else x[0] for x in semantic_left) left_ids = tok(f"<|object_ref_start|>{left_text}<|object_ref_end|>", return_tensors="pt").input_ids.repeat(of.NUM_CODEBOOKS, 1) parts.append(left_ids) audio_parts.append(torch.zeros(left_ids.shape[1], dtype=torch.bool)) # the target text starts with the prompt's transcript (training: "prompt text target text"); no text-only # `before` (training has none next to the left context) target_ids = tok(of.context_text(f"{prompt_text} {text}".strip(), (), list(after)), return_tensors="pt").input_ids.repeat(of.NUM_CODEBOOKS, 1) parts.append(target_ids) audio_parts.append(torch.zeros(target_ids.shape[1], dtype=torch.bool)) if prompt_codes.shape[1]: parts.append(prompt_codes.cpu().long()) audio_parts.append(torch.ones(prompt_codes.shape[1], dtype=torch.bool)) context = (torch.cat(parts, 1), torch.cat(audio_parts)) else: full_text = f"{prompt_text} {text}".strip() prefix = prefix_ids(tok, full_text, lang, before, after, speaker, quality) context = (torch.cat([prefix, prompt_codes.cpu().long()], 1), torch.cat([torch.zeros(prefix.shape[1], dtype=torch.bool), torch.ones(prompt_codes.shape[1], dtype=torch.bool)])) if max_frames is None: rate = frames / chars if chars else 1.6 max_frames = int(2 * rate * len(text)) + 2 * block none = (torch.zeros(of.NUM_CODEBOOKS, 0, dtype=torch.long), torch.zeros(0, dtype=torch.bool)) st = _Stream(model, [context] + ([none] if cfg > 0 else []), dev) mtp_dtype = next(model.mtp.parameters()).dtype # the MTP sees the prompt (all 10 codes) as decided frames before the target, as in training P = prompt_codes.shape[1] if (prompt_codes is not None and not voices and mtp_prompt) else 0 S = model.mtp.new_state(P + max_frames + block, model.llm.config.hidden_size, dev, mtp_dtype) if P: S["codes"][:, :P] = prompt_codes.to(dev).long() S["states"][0][0, :P] = st.ctx_h[0, -P:].to(mtp_dtype) with torch.autocast(dev.type, torch.bfloat16): model.mtp.prime(S, P) clean = torch.zeros(0, dtype=torch.long, device=dev) if target_prefix is not None: # continuation: earlier audio [10, n·block] of the target as decided blocks tp = target_prefix.to(dev).long() for lo in range(0, tp.shape[1] - tp.shape[1] % block, block): full = tp[:, lo:lo + block] h = st.peek(full[0]) S["codes"][:, P + lo:P + lo + block], S["states"][0][0, P + lo:P + lo + block] = full, h.to(mtp_dtype) st.commit(full) clean = torch.cat([clean, full[0]]) while clean.shape[0] < max_frames: cur = torch.full((block,), of.MASK_ID, dtype=torch.long, device=dev) for s in range(steps): lg = st.logits(cur) logits = lg[0] + cfg * (lg[0] - lg[1]) if cfg > 0 else lg[0] logits[:, of.MASK_ID] = -float("inf") if min_frames: local = torch.arange(block, device=dev) + clean.shape[0] logits[local < min_frames, of.END_ID:] = -float("inf") masked = cur == of.MASK_ID probs = F.softmax(logits / max(temperature, 1e-5), -1) sample = torch.multinomial(probs, 1, generator=g)[:, 0] if temperature > 0 else logits.argmax(-1) conf = probs.gather(1, sample[:, None])[:, 0] conf = torch.where(masked, conf, torch.full_like(conf, -1.0)) keep_masked = int(math.floor(block * math.cos(math.pi / 2 * (s + 1) / steps))) n_fix = int(masked.sum()) - keep_masked if n_fix > 0: idx = conf.topk(n_fix).indices cur[idx] = sample[idx] lo = clean.shape[0] end = (cur >= of.END_ID).nonzero() nb = int(end[0]) if len(end) else block # speech frames of this block full = torch.full((of.NUM_CODEBOOKS, block), of.MASK_ID, dtype=torch.long, device=dev) full[0] = cur if nb: h = st.peek(cur) S["codes"][0, P + lo:P + lo + nb], S["states"][0][0, P + lo:P + lo + nb] = cur[:nb], h[:nb].to(mtp_dtype) with torch.autocast(dev.type, torch.bfloat16): model.mtp.decode_block(S, P + lo, P + lo + nb, block, mtp_rounds, mtp_temperature, mtp_top_k, prefix=P) full[1:, :nb] = S["codes"][1:, P + lo:P + lo + nb] clean = torch.cat([clean, cur]) yield clean, S["codes"][:, P:P + lo + nb] if len(end): return st.commit(full if full_context else full[:1].expand(of.NUM_CODEBOOKS, -1).clone().masked_fill( torch.arange(of.NUM_CODEBOOKS, device=dev)[:, None] > 0, of.MASK_ID)) # semantic-only context def speech_frames(clean): end = (clean >= of.END_ID).nonzero() return int(end[0]) if len(end) else clean.shape[0]