Upload 58 files

4d1131c verified 5 days ago

13.3 kB

"""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]

	"""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]