| # # Windowed (local bidirectional) streaming decoder for MioCodec's wave | |
| # # decoder path — an alternative to chunk_streamer.py's causal-only design. | |
| # # | |
| # # WHY THIS EXISTS / WHAT PROBLEM IT SOLVES: | |
| # # chunk_streamer.py (causal wave_prenet/wave_decoder + causal ResNet convs) | |
| # # only ever sees PAST tokens for each chunk. That's genuinely streamable | |
| # # with zero lookahead, but chunk boundaries can still show minor | |
| # # discontinuities from context-window proportional-splitting error, and | |
| # # any manual overlap/crossfade approach (as in the author's separate | |
| # # no-finetuning streaming script) still needs hand-tuned blending to hide | |
| # # the seam. | |
| # # | |
| # # This module instead gives the model REAL future context (a fixed, small | |
| # # lookahead window of already-available tokens — since the TTS LLM has, at | |
| # # decode time, already generated those tokens; they are not being | |
| # # predicted, just not yet "released" as audio) for every chunk except the | |
| # # very last one. With genuine future context, the model's own computation | |
| # # smooths the chunk boundary the same way it would if you'd decoded the | |
| # # whole utterance in one shot — no manual crossfade/overlap-add stitching | |
| # # needed on the audio side at all. | |
| # # | |
| # # ARCHITECTURAL CHANGE FROM chunk_streamer.py: | |
| # # - wave_prenet / wave_decoder: causal=False (plain bidirectional | |
| # # attention) INSIDE each local window. This is not "attend to | |
| # # everything ever generated" — it's bounded to | |
| # # left_context + chunk + right_lookahead tokens, so latency and | |
| # # compute stay bounded. Only the mask changes; weights/shapes don't. | |
| # # - wave_prior_net / wave_post_net: the ORIGINAL symmetric-padded | |
| # # ResNetBlock/ResNetStack (from istft_head.py), NOT CausalResNetStack. | |
| # # Symmetric padding is fine here because the whole window (including | |
| # # lookahead) is known upfront when the window is processed — there's | |
| # # no "future doesn't exist yet" problem within one already-buffered | |
| # # window. | |
| # # - istft_head: still the streaming ISTFTHead (streaming_istft_head.py), | |
| # # used only for its offline-equivalent forward() path per window (NOT | |
| # # stream_step — each window is decoded as one shot since we're not | |
| # # doing frame-by-frame emission anymore, we're doing window-by-window | |
| # # with exact trimming instead of an incremental accumulator). | |
| # # | |
| # # EDGE CASES (explicitly handled, not hardcoded around): | |
| # # - FIRST chunk: no left context exists yet (nothing generated before | |
| # # it). left_context_tokens is naturally 0 here; this is correct | |
| # # behavior, not a bug to patch around. | |
| # # - LAST chunk (LLM finished generating): no right lookahead exists | |
| # # (nothing more will ever be generated). right_lookahead_tokens is 0 | |
| # # here for the same reason. | |
| # # - Both context lengths are therefore VARIABLE per call, driven by | |
| # # actual token availability, never assumed fixed. | |
| # # | |
| # # EXACT (not approximate) frame trimming: chunk_streamer.py used a | |
| # # proportional estimate (context_frames = round(stft_length * fraction)) | |
| # # to decide which output frames to discard. That approximation was a | |
| # # known, flagged source of boundary error. This module instead computes | |
| # # the exact frame count for the left/right portions using the SAME | |
| # # _calculate_target_stft_length function the model itself uses, applied | |
| # # separately to left_context_tokens and right_lookahead_tokens durations, | |
| # # so the split lines up with how the model actually maps tokens to frames. | |
| # from dataclasses import dataclass | |
| # import torch | |
| # import torch.nn.functional as F | |
| # MIN_CHUNK_SIZE = 2 # user-facing floor per this task: "variable, any number > 1" | |
| # DEFAULT_CHUNK_SIZE = 20 | |
| # DEFAULT_CONTEXT_TOKENS = 10 # "around 10" per spec, applies to both left and right | |
| # @dataclass | |
| # class WindowedStreamConfig: | |
| # chunk_size: int = DEFAULT_CHUNK_SIZE # tokens released as audio per step (variable, user-chosen) | |
| # left_context_tokens: int = DEFAULT_CONTEXT_TOKENS # past tokens given as extra context (not released as audio) | |
| # right_lookahead_tokens: int = DEFAULT_CONTEXT_TOKENS # future tokens given as extra context (not released as audio) | |
| # token_rate_hz: float = 25.0 | |
| # def __post_init__(self): | |
| # if self.chunk_size <= 1: | |
| # raise ValueError( | |
| # f"chunk_size={self.chunk_size} must be > 1 (variable at inference time, " | |
| # f"but always more than a single token)." | |
| # ) | |
| # if self.left_context_tokens < 0 or self.right_lookahead_tokens < 0: | |
| # raise ValueError("context/lookahead token counts must be >= 0.") | |
| # def set_transformer_causal(transformer, causal: bool): | |
| # """ | |
| # Flips Attention.causal on every layer of a Transformer post-hoc (same | |
| # technique as build_model_with_causal_patches in finetuning_streaming_lora.py, | |
| # just reusable/importable here). Does not change any weight shapes — | |
| # only the attention mask applied at forward time. | |
| # """ | |
| # for layer in transformer.layers: | |
| # layer.attention.causal = causal | |
| # class MioCodecWindowedStreamer: | |
| # """ | |
| # Stateful per-session windowed streaming decoder. One instance per | |
| # active voice-bot call / TTS session. | |
| # Requires a model where wave_prenet / wave_decoder / wave_prior_net / | |
| # wave_post_net are the ORIGINAL (non-causal, symmetric-padded) | |
| # versions -- NOT the causal-patched ones from chunk_streamer.py's | |
| # companion build. Use build_windowed_model() below, or manually ensure | |
| # set_transformer_causal(model.wave_prenet, False) and | |
| # set_transformer_causal(model.wave_decoder, False) have been called and | |
| # wave_prior_net/wave_post_net are plain ResNetStack (not | |
| # CausalResNetStack). | |
| # Usage (token-by-token arrival from an autoregressive TTS LLM): | |
| # streamer = MioCodecWindowedStreamer( | |
| # model, global_embedding, | |
| # config=WindowedStreamConfig(chunk_size=20, left_context_tokens=10, right_lookahead_tokens=10) | |
| # ) | |
| # for new_token_embedding in llm_token_stream: | |
| # audio_chunk = streamer.push_token(new_token_embedding) | |
| # if audio_chunk is not None: | |
| # play_audio(audio_chunk) # no crossfade needed before playing this | |
| # final_audio = streamer.finish() # releases the last chunk(s), with | |
| # # right_lookahead_tokens=0 since nothing more exists | |
| # if final_audio is not None: | |
| # play_audio(final_audio) | |
| # Chunk size is passed once at construction, matching the "variable at | |
| # inference time, user picks it, must be > 1" requirement — instantiate | |
| # a new streamer (or read chunk_size from config) per session/request | |
| # with whatever value the caller wants. | |
| # """ | |
| # def __init__(self, model, global_embedding: torch.Tensor, config: WindowedStreamConfig | None = None): | |
| # self.model = model | |
| # self.config = config or WindowedStreamConfig() | |
| # self.device = global_embedding.device | |
| # self.global_embedding = global_embedding.unsqueeze(0) # (1, dim) | |
| # self._all_tokens: list[torch.Tensor] = [] # every token embedding seen so far, in order (kept for left-context slicing) | |
| # self._released_count = 0 # how many tokens' worth of audio has already been emitted | |
| # self._finished = False | |
| # def push_token(self, token_embedding: torch.Tensor) -> torch.Tensor | None: | |
| # return self.push_tokens(token_embedding.unsqueeze(0)) | |
| # def push_tokens(self, token_embeddings: torch.Tensor) -> torch.Tensor | None: | |
| # """ | |
| # token_embeddings: (k, dim) — newly-arrived tokens, in order. | |
| # Returns (samples,) if a full chunk_size window could be released | |
| # (with right_lookahead available), else None (still buffering, | |
| # waiting for enough lookahead tokens to arrive). | |
| # """ | |
| # if self._finished: | |
| # raise RuntimeError("This streamer has already been finished(); create a new one for a new session.") | |
| # self._all_tokens.append(token_embeddings) | |
| # all_tokens_flat = torch.cat(self._all_tokens, dim=0) | |
| # total_available = all_tokens_flat.shape[0] | |
| # pending = total_available - self._released_count | |
| # needed = self.config.chunk_size + self.config.right_lookahead_tokens | |
| # if pending < needed: | |
| # return None # not enough tokens yet to release a full chunk + see its lookahead | |
| # chunk_start = self._released_count | |
| # chunk_end = chunk_start + self.config.chunk_size | |
| # # Take exactly one chunk_size window per call; extra pending tokens | |
| # # beyond chunk_end (+lookahead) simply stay buffered in | |
| # # self._all_tokens for the next call. | |
| # audio = self._decode_window(all_tokens_flat, chunk_start, chunk_end, is_final_chunk=False) | |
| # self._released_count = chunk_end | |
| # return audio | |
| # def _decode_window( | |
| # self, all_tokens_flat: torch.Tensor, chunk_start: int, chunk_end: int, is_final_chunk: bool | |
| # ) -> torch.Tensor: | |
| # total_available = all_tokens_flat.shape[0] | |
| # left_start = max(0, chunk_start - self.config.left_context_tokens) | |
| # actual_left = chunk_start - left_start # variable: 0 for the first chunk, else up to left_context_tokens | |
| # if is_final_chunk: | |
| # right_end = total_available | |
| # else: | |
| # right_end = min(total_available, chunk_end + self.config.right_lookahead_tokens) | |
| # actual_right = right_end - chunk_end # variable: 0 for the final chunk, else up to right_lookahead_tokens | |
| # window_tokens = all_tokens_flat[left_start:right_end] # (left+chunk+right, dim) | |
| # content_embedding = window_tokens.unsqueeze(0).to(self.device) # (1, W, dim) | |
| # # Exact (not proportional-estimate) frame accounting: compute the | |
| # # target STFT length for left / chunk / right SEPARATELY using the | |
| # # model's own token-rate -> sample-rate -> frame-length formula, | |
| # # then sum, so the boundary between "context frames" and "audio to | |
| # # keep" is derived the same way the model derives frame counts | |
| # # everywhere else, not approximated from a token-count ratio. | |
| # def tokens_to_frames(n_tokens: int) -> int: | |
| # if n_tokens <= 0: | |
| # return 0 | |
| # audio_len = int(round(n_tokens * self.model.config.sample_rate / self.config.token_rate_hz)) | |
| # return self.model._calculate_target_stft_length(audio_len) | |
| # left_frames = tokens_to_frames(actual_left) | |
| # chunk_frames = tokens_to_frames(chunk_end - chunk_start) | |
| # right_frames = tokens_to_frames(actual_right) | |
| # window_frames = left_frames + chunk_frames + right_frames | |
| # # --- inline forward_wave, non-causal (bidirectional-in-window) --- | |
| # local_latent = self.model.wave_prenet(content_embedding) | |
| # if self.model.wave_conv_upsample is not None: | |
| # local_latent = self.model.wave_conv_upsample(local_latent.transpose(1, 2)).transpose(1, 2) | |
| # local_latent = F.interpolate( | |
| # local_latent.transpose(1, 2), size=window_frames, mode=self.model.config.wave_interpolation_mode | |
| # ).transpose(1, 2) | |
| # local_latent = self.model.wave_prior_net(local_latent.transpose(1, 2)).transpose(1, 2) | |
| # local_latent = self.model.wave_decoder(local_latent, condition=self.global_embedding.unsqueeze(1)) | |
| # local_latent = self.model.wave_post_net(local_latent.transpose(1, 2)).transpose(1, 2) | |
| # if self.model.wave_upsampler is not None: | |
| # local_latent = self.model.wave_upsampler(local_latent.transpose(1, 2)) | |
| # # Exact frame slice for the chunk-only portion. | |
| # chunk_only_frames = local_latent[:, left_frames : left_frames + chunk_frames, :] | |
| # # One-shot ISTFT over exactly the chunk's own frames. Uses the | |
| # # offline-equivalent forward() (full window in, full audio out for | |
| # # THIS SLICE), not the incremental stream_step accumulator — each | |
| # # windowed chunk is processed independently here, which is exact | |
| # # by construction (no accumulator state to drift), unlike | |
| # # chunk_streamer.py's approach. | |
| # audio_chunk = self.model.istft_head(chunk_only_frames).squeeze(0) | |
| # return audio_chunk | |
| # def finish(self) -> torch.Tensor | None: | |
| # """ | |
| # Call once when the LLM has finished emitting tokens. Releases all | |
| # remaining un-released tokens as the final chunk, with | |
| # right_lookahead_tokens=0 (nothing more will ever exist) — this is | |
| # the correct, not-a-bug edge case for the last chunk. | |
| # """ | |
| # if self._finished: | |
| # return None | |
| # self._finished = True | |
| # if not self._all_tokens: | |
| # return None | |
| # all_tokens_flat = torch.cat(self._all_tokens, dim=0) | |
| # total_available = all_tokens_flat.shape[0] | |
| # if self._released_count >= total_available: | |
| # return None # everything was already released by push_tokens calls | |
| # audio = self._decode_window( | |
| # all_tokens_flat, self._released_count, total_available, is_final_chunk=True | |
| # ) | |
| # self._released_count = total_available | |
| # return audio | |
| # def build_windowed_model(repo_id: str = "Aratako/MioCodec-25Hz-24kHz", device: str = "cuda"): | |
| # """ | |
| # Loads the pretrained model and configures it for WINDOWED (local | |
| # bidirectional) streaming: wave_prenet / wave_decoder set to | |
| # causal=False (plain attention, unchanged from pretrained default — | |
| # MioCodec's original decoder was never causal to begin with, so this | |
| # is actually a no-op vs. the stock model) and wave_prior_net / | |
| # wave_post_net left as the original ResNetStack (also a no-op vs. | |
| # stock). The only thing this function actually needs to do beyond | |
| # loading is swap istft_head for the streaming-capable ISTFTHead | |
| # (streaming_istft_head.py) so MioCodecWindowedStreamer can call it. | |
| # In other words: for the windowed-streaming approach, wave_prenet / | |
| # wave_decoder / wave_prior_net / wave_post_net do NOT need patching at | |
| # all — they're used exactly as pretrained. Only istft_head needs the | |
| # streaming-compatible version (and even then, only its offline- | |
| # equivalent forward() path is used per window, so this is also a | |
| # behavior-preserving swap, not an architecture change). | |
| # """ | |
| # from miocodec.model import MioCodecModel | |
| # from streaming_istft_head import ISTFTHead as StreamingISTFTHead | |
| # model = MioCodecModel.from_pretrained(repo_id).to(device) | |
| # cfg = model.config | |
| # old_head = model.istft_head | |
| # new_head = StreamingISTFTHead( | |
| # dim=cfg.wave_decoder_dim, n_fft=cfg.n_fft, hop_length=cfg.hop_length, padding=cfg.istft_padding | |
| # ).to(device) # BUG FIX: newly constructed modules default to CPU even | |
| # # when assigned onto a model already moved to `device` -- | |
| # # nn.Module.to() only moves parameters that exist AT CALL | |
| # # TIME, it does not retroactively affect submodules | |
| # # attached afterward. Without this, model.istft_head.out | |
| # # (a fresh nn.Linear) stayed on CPU while every other | |
| # # module was on CUDA, causing | |
| # # "Expected all tensors to be on the same device" only | |
| # # when running on GPU (silently fine on CPU, since | |
| # # everything was CPU there anyway -- which is exactly | |
| # # why this didn't surface until a real GPU run). | |
| # new_head.out.load_state_dict(old_head.out.state_dict()) | |
| # model.istft_head = new_head | |
| # # Explicit, not implicit: confirm causal=False on both transformers. | |
| # # This matches the model's pretrained state already, but setting it | |
| # # explicitly documents the requirement and protects against a future | |
| # # change to MioCodecModelConfig defaults silently breaking this. | |
| # set_transformer_causal(model.wave_prenet, False) | |
| # set_transformer_causal(model.wave_decoder, False) | |
| # return model | |
| # Windowed (local bidirectional) streaming decoder for MioCodec's wave | |
| # decoder path — an alternative to chunk_streamer.py's causal-only design. | |
| # | |
| # WHY THIS EXISTS / WHAT PROBLEM IT SOLVES: | |
| # chunk_streamer.py (causal wave_prenet/wave_decoder + causal ResNet convs) | |
| # only ever sees PAST tokens for each chunk. That's genuinely streamable | |
| # with zero lookahead, but chunk boundaries can still show minor | |
| # discontinuities from context-window proportional-splitting error, and | |
| # any manual overlap/crossfade approach (as in the author's separate | |
| # no-finetuning streaming script) still needs hand-tuned blending to hide | |
| # the seam. | |
| # | |
| # This module instead gives the model REAL future context (a fixed, small | |
| # lookahead window of already-available tokens — since the TTS LLM has, at | |
| # decode time, already generated those tokens; they are not being | |
| # predicted, just not yet "released" as audio) for every chunk except the | |
| # very last one. With genuine future context, the model's own computation | |
| # smooths the chunk boundary the same way it would if you'd decoded the | |
| # whole utterance in one shot — no manual crossfade/overlap-add stitching | |
| # needed on the audio side at all. | |
| # | |
| # ARCHITECTURAL CHANGE FROM chunk_streamer.py: | |
| # - wave_prenet / wave_decoder: causal=False (plain bidirectional | |
| # attention) INSIDE each local window. This is not "attend to | |
| # everything ever generated" — it's bounded to | |
| # left_context + chunk + right_lookahead tokens, so latency and | |
| # compute stay bounded. Only the mask changes; weights/shapes don't. | |
| # - wave_prior_net / wave_post_net: the ORIGINAL symmetric-padded | |
| # ResNetBlock/ResNetStack (from istft_head.py), NOT CausalResNetStack. | |
| # Symmetric padding is fine here because the whole window (including | |
| # lookahead) is known upfront when the window is processed — there's | |
| # no "future doesn't exist yet" problem within one already-buffered | |
| # window. | |
| # - istft_head: still the streaming ISTFTHead (streaming_istft_head.py), | |
| # used only for its offline-equivalent forward() path per window (NOT | |
| # stream_step — each window is decoded as one shot since we're not | |
| # doing frame-by-frame emission anymore, we're doing window-by-window | |
| # with exact trimming instead of an incremental accumulator). | |
| # | |
| # EDGE CASES (explicitly handled, not hardcoded around): | |
| # - FIRST chunk: no left context exists yet (nothing generated before | |
| # it). left_context_tokens is naturally 0 here; this is correct | |
| # behavior, not a bug to patch around. | |
| # - LAST chunk (LLM finished generating): no right lookahead exists | |
| # (nothing more will ever be generated). right_lookahead_tokens is 0 | |
| # here for the same reason. | |
| # - Both context lengths are therefore VARIABLE per call, driven by | |
| # actual token availability, never assumed fixed. | |
| # | |
| # EXACT (not approximate) frame trimming: chunk_streamer.py used a | |
| # proportional estimate (context_frames = round(stft_length * fraction)) | |
| # to decide which output frames to discard. That approximation was a | |
| # known, flagged source of boundary error. This module instead computes | |
| # the exact frame count for the left/right portions using the SAME | |
| # _calculate_target_stft_length function the model itself uses, applied | |
| # separately to left_context_tokens and right_lookahead_tokens durations, | |
| # so the split lines up with how the model actually maps tokens to frames. | |
| from dataclasses import dataclass | |
| import torch | |
| import torch.nn.functional as F | |
| MIN_CHUNK_SIZE = 2 # user-facing floor per this task: "variable, any number > 1" | |
| DEFAULT_CHUNK_SIZE = 20 | |
| DEFAULT_CONTEXT_TOKENS = 10 # "around 10" per spec, applies to both left and right | |
| class WindowedStreamConfig: | |
| chunk_size: int = DEFAULT_CHUNK_SIZE # tokens released as audio per step (variable, user-chosen) | |
| left_context_tokens: int = DEFAULT_CONTEXT_TOKENS # past tokens given as extra context (not released as audio) | |
| right_lookahead_tokens: int = DEFAULT_CONTEXT_TOKENS # future tokens given as extra context (not released as audio) | |
| token_rate_hz: float = 25.0 | |
| def __post_init__(self): | |
| if self.chunk_size <= 1: | |
| raise ValueError( | |
| f"chunk_size={self.chunk_size} must be > 1 (variable at inference time, " | |
| f"but always more than a single token)." | |
| ) | |
| if self.left_context_tokens < 0 or self.right_lookahead_tokens < 0: | |
| raise ValueError("context/lookahead token counts must be >= 0.") | |
| def set_transformer_causal(transformer, causal: bool): | |
| """ | |
| Flips Attention.causal on every layer of a Transformer post-hoc (same | |
| technique as build_model_with_causal_patches in finetuning_streaming_lora.py, | |
| just reusable/importable here). Does not change any weight shapes — | |
| only the attention mask applied at forward time. | |
| """ | |
| for layer in transformer.layers: | |
| layer.attention.causal = causal | |
| class MioCodecWindowedStreamer: | |
| """ | |
| Stateful per-session windowed streaming decoder. One instance per | |
| active voice-bot call / TTS session. | |
| Requires a model where wave_prenet / wave_decoder / wave_prior_net / | |
| wave_post_net are the ORIGINAL (non-causal, symmetric-padded) | |
| versions -- NOT the causal-patched ones from chunk_streamer.py's | |
| companion build. Use build_windowed_model() below, or manually ensure | |
| set_transformer_causal(model.wave_prenet, False) and | |
| set_transformer_causal(model.wave_decoder, False) have been called and | |
| wave_prior_net/wave_post_net are plain ResNetStack (not | |
| CausalResNetStack). | |
| Usage (token-by-token arrival from an autoregressive TTS LLM): | |
| streamer = MioCodecWindowedStreamer( | |
| model, global_embedding, | |
| config=WindowedStreamConfig(chunk_size=20, left_context_tokens=10, right_lookahead_tokens=10) | |
| ) | |
| for new_token_embedding in llm_token_stream: | |
| audio_chunk = streamer.push_token(new_token_embedding) | |
| if audio_chunk is not None: | |
| play_audio(audio_chunk) # no crossfade needed before playing this | |
| final_audio = streamer.finish() # releases the last chunk(s), with | |
| # right_lookahead_tokens=0 since nothing more exists | |
| if final_audio is not None: | |
| play_audio(final_audio) | |
| Chunk size is passed once at construction, matching the "variable at | |
| inference time, user picks it, must be > 1" requirement — instantiate | |
| a new streamer (or read chunk_size from config) per session/request | |
| with whatever value the caller wants. | |
| """ | |
| def __init__(self, model, global_embedding: torch.Tensor, config: WindowedStreamConfig | None = None): | |
| self.model = model | |
| self.config = config or WindowedStreamConfig() | |
| self.device = global_embedding.device | |
| self.global_embedding = global_embedding.unsqueeze(0) # (1, dim) | |
| self._all_tokens: list[torch.Tensor] = [] # every token embedding seen so far, in order (kept for left-context slicing) | |
| self._released_count = 0 # how many tokens' worth of audio has already been emitted | |
| self._finished = False | |
| def is_finished(self) -> bool: | |
| """True once finish() has been called on this streamer. Public, | |
| stable name for callers (e.g. app.py) that need to check session | |
| state without reaching into a private attribute.""" | |
| return self._finished | |
| def push_token(self, token_embedding: torch.Tensor) -> torch.Tensor | None: | |
| return self.push_tokens(token_embedding.unsqueeze(0)) | |
| def push_tokens(self, token_embeddings: torch.Tensor) -> torch.Tensor | None: | |
| """ | |
| token_embeddings: (k, dim) — newly-arrived tokens, in order. | |
| Returns (samples,) if a full chunk_size window could be released | |
| (with right_lookahead available), else None (still buffering, | |
| waiting for enough lookahead tokens to arrive). | |
| Releases EXACTLY chunk_size tokens per call, even if more than | |
| chunk_size + right_lookahead_tokens are already available. This is | |
| the right choice when you want steady, predictable chunk sizes | |
| (e.g. for a fixed playback cadence). If your caller can end up | |
| with a large backlog (e.g. the LLM generated a burst of tokens | |
| faster than you're consuming them) and you'd rather release | |
| everything available in one go instead of releasing it as several | |
| smaller chunk_size-sized pieces, use push_all_available() instead | |
| — fewer, larger chunks means fewer join points and each chunk | |
| gets more of its own tokens as real (not context-only) input. | |
| """ | |
| if self._finished: | |
| raise RuntimeError("This streamer has already been finished(); create a new one for a new session.") | |
| self._all_tokens.append(token_embeddings) | |
| all_tokens_flat = torch.cat(self._all_tokens, dim=0) | |
| total_available = all_tokens_flat.shape[0] | |
| pending = total_available - self._released_count | |
| needed = self.config.chunk_size + self.config.right_lookahead_tokens | |
| if pending < needed: | |
| return None # not enough tokens yet to release a full chunk + see its lookahead | |
| chunk_start = self._released_count | |
| chunk_end = chunk_start + self.config.chunk_size | |
| # Take exactly one chunk_size window per call; extra pending tokens | |
| # beyond chunk_end (+lookahead) simply stay buffered in | |
| # self._all_tokens for the next call. | |
| audio = self._decode_window(all_tokens_flat, chunk_start, chunk_end, is_final_chunk=False) | |
| self._released_count = chunk_end | |
| return audio | |
| def push_all_available(self, token_embeddings: torch.Tensor, generation_finished: bool) -> torch.Tensor | None: | |
| """ | |
| Like push_tokens(), but releases AS MANY tokens as are currently | |
| buffered in one shot, instead of always exactly chunk_size. | |
| Use this when the caller wants dynamic-sized chunks: if the LLM | |
| produced a burst of tokens (e.g. it's generating faster than | |
| real-time and several chunk_size-worths built up while you were | |
| waiting on something else), releasing them as ONE larger chunk | |
| instead of several chunk_size-sized ones means fewer join points | |
| and better in-chunk context for the model (each token gets to | |
| "see" more of its real neighbors, not just the reserved lookahead | |
| window, since a single call's real chunk tokens now span all of | |
| them). | |
| Behavior: | |
| - If generation_finished is True: release EVERYTHING pending, | |
| with zero right-lookahead (nothing more will ever exist) — | |
| same edge case finish() handles, made available mid-stream so | |
| callers don't need a separate code path for "the LLM just | |
| finished on this exact call." | |
| - If generation_finished is False: release everything pending | |
| EXCEPT the last right_lookahead_tokens tokens, which are held | |
| back so the released portion still gets real future context | |
| (not zero-lookahead) rather than being downgraded to a | |
| causal-only decode. If pending <= right_lookahead_tokens, | |
| nothing is released yet (returns None) — there's nothing to | |
| safely release without sacrificing all lookahead. | |
| Args: | |
| token_embeddings: (k, dim) newly-arrived tokens, in order. | |
| Pass an empty tensor (shape (0, dim)) if you're calling | |
| this purely to flush already-buffered tokens without new | |
| arrivals. | |
| generation_finished: whether the LLM has stopped producing | |
| further tokens for this session. | |
| Returns: | |
| (samples,) audio if anything was released, else None. | |
| """ | |
| if self._finished: | |
| raise RuntimeError("This streamer has already been finished(); create a new one for a new session.") | |
| if token_embeddings.shape[0] > 0: | |
| self._all_tokens.append(token_embeddings) | |
| if not self._all_tokens: | |
| return None | |
| all_tokens_flat = torch.cat(self._all_tokens, dim=0) | |
| total_available = all_tokens_flat.shape[0] | |
| pending = total_available - self._released_count | |
| if pending <= 0: | |
| return None | |
| chunk_start = self._released_count | |
| if generation_finished: | |
| chunk_end = total_available # release everything, zero lookahead (handled by _decode_window's is_final_chunk=True) | |
| audio = self._decode_window(all_tokens_flat, chunk_start, chunk_end, is_final_chunk=True) | |
| self._released_count = chunk_end | |
| return audio | |
| # Not finished: hold back right_lookahead_tokens so the released | |
| # portion still gets real future context. | |
| releasable = pending - self.config.right_lookahead_tokens | |
| if releasable <= 0: | |
| return None # not enough buffered beyond the reserved lookahead to release anything yet | |
| chunk_end = chunk_start + releasable | |
| audio = self._decode_window(all_tokens_flat, chunk_start, chunk_end, is_final_chunk=False) | |
| self._released_count = chunk_end | |
| return audio | |
| def _decode_window( | |
| self, all_tokens_flat: torch.Tensor, chunk_start: int, chunk_end: int, is_final_chunk: bool | |
| ) -> torch.Tensor: | |
| total_available = all_tokens_flat.shape[0] | |
| left_start = max(0, chunk_start - self.config.left_context_tokens) | |
| actual_left = chunk_start - left_start # variable: 0 for the first chunk, else up to left_context_tokens | |
| if is_final_chunk: | |
| right_end = total_available | |
| else: | |
| right_end = min(total_available, chunk_end + self.config.right_lookahead_tokens) | |
| actual_right = right_end - chunk_end # variable: 0 for the final chunk, else up to right_lookahead_tokens | |
| window_tokens = all_tokens_flat[left_start:right_end] # (left+chunk+right, dim) | |
| content_embedding = window_tokens.unsqueeze(0).to(self.device) # (1, W, dim) | |
| # Exact (not proportional-estimate) frame accounting: compute the | |
| # target STFT length for left / chunk / right SEPARATELY using the | |
| # model's own token-rate -> sample-rate -> frame-length formula, | |
| # then sum, so the boundary between "context frames" and "audio to | |
| # keep" is derived the same way the model derives frame counts | |
| # everywhere else, not approximated from a token-count ratio. | |
| def tokens_to_frames(n_tokens: int) -> int: | |
| if n_tokens <= 0: | |
| return 0 | |
| audio_len = int(round(n_tokens * self.model.config.sample_rate / self.config.token_rate_hz)) | |
| return self.model._calculate_target_stft_length(audio_len) | |
| left_frames = tokens_to_frames(actual_left) | |
| chunk_frames = tokens_to_frames(chunk_end - chunk_start) | |
| right_frames = tokens_to_frames(actual_right) | |
| window_frames = left_frames + chunk_frames + right_frames | |
| # --- inline forward_wave, non-causal (bidirectional-in-window) --- | |
| local_latent = self.model.wave_prenet(content_embedding) | |
| if self.model.wave_conv_upsample is not None: | |
| local_latent = self.model.wave_conv_upsample(local_latent.transpose(1, 2)).transpose(1, 2) | |
| local_latent = F.interpolate( | |
| local_latent.transpose(1, 2), size=window_frames, mode=self.model.config.wave_interpolation_mode | |
| ).transpose(1, 2) | |
| local_latent = self.model.wave_prior_net(local_latent.transpose(1, 2)).transpose(1, 2) | |
| local_latent = self.model.wave_decoder(local_latent, condition=self.global_embedding.unsqueeze(1)) | |
| local_latent = self.model.wave_post_net(local_latent.transpose(1, 2)).transpose(1, 2) | |
| if self.model.wave_upsampler is not None: | |
| local_latent = self.model.wave_upsampler(local_latent.transpose(1, 2)) | |
| # Exact frame slice for the chunk-only portion. | |
| chunk_only_frames = local_latent[:, left_frames : left_frames + chunk_frames, :] | |
| # One-shot ISTFT over exactly the chunk's own frames. Uses the | |
| # offline-equivalent forward() (full window in, full audio out for | |
| # THIS SLICE), not the incremental stream_step accumulator — each | |
| # windowed chunk is processed independently here, which is exact | |
| # by construction (no accumulator state to drift), unlike | |
| # chunk_streamer.py's approach. | |
| audio_chunk = self.model.istft_head(chunk_only_frames).squeeze(0) | |
| return audio_chunk | |
| def finish(self) -> torch.Tensor | None: | |
| """ | |
| Call once when the LLM has finished emitting tokens. Releases all | |
| remaining un-released tokens as the final chunk, with | |
| right_lookahead_tokens=0 (nothing more will ever exist) — this is | |
| the correct, not-a-bug edge case for the last chunk. | |
| """ | |
| if self._finished: | |
| return None | |
| self._finished = True | |
| if not self._all_tokens: | |
| return None | |
| all_tokens_flat = torch.cat(self._all_tokens, dim=0) | |
| total_available = all_tokens_flat.shape[0] | |
| if self._released_count >= total_available: | |
| return None # everything was already released by push_tokens calls | |
| audio = self._decode_window( | |
| all_tokens_flat, self._released_count, total_available, is_final_chunk=True | |
| ) | |
| self._released_count = total_available | |
| return audio | |
| def build_windowed_model(repo_id: str = "Aratako/MioCodec-25Hz-24kHz", device: str = "cuda"): | |
| """ | |
| Loads the pretrained model and configures it for WINDOWED (local | |
| bidirectional) streaming: wave_prenet / wave_decoder set to | |
| causal=False (plain attention, unchanged from pretrained default — | |
| MioCodec's original decoder was never causal to begin with, so this | |
| is actually a no-op vs. the stock model) and wave_prior_net / | |
| wave_post_net left as the original ResNetStack (also a no-op vs. | |
| stock). The only thing this function actually needs to do beyond | |
| loading is swap istft_head for the streaming-capable ISTFTHead | |
| (streaming_istft_head.py) so MioCodecWindowedStreamer can call it. | |
| In other words: for the windowed-streaming approach, wave_prenet / | |
| wave_decoder / wave_prior_net / wave_post_net do NOT need patching at | |
| all — they're used exactly as pretrained. Only istft_head needs the | |
| streaming-compatible version (and even then, only its offline- | |
| equivalent forward() path is used per window, so this is also a | |
| behavior-preserving swap, not an architecture change). | |
| """ | |
| from miocodec.model import MioCodecModel | |
| from streaming_istft_head import ISTFTHead as StreamingISTFTHead | |
| model = MioCodecModel.from_pretrained(repo_id).to(device) | |
| cfg = model.config | |
| old_head = model.istft_head | |
| new_head = StreamingISTFTHead( | |
| dim=cfg.wave_decoder_dim, n_fft=cfg.n_fft, hop_length=cfg.hop_length, padding=cfg.istft_padding | |
| ).to(device) # BUG FIX: newly constructed modules default to CPU even | |
| # when assigned onto a model already moved to `device` -- | |
| # nn.Module.to() only moves parameters that exist AT CALL | |
| # TIME, it does not retroactively affect submodules | |
| # attached afterward. Without this, model.istft_head.out | |
| # (a fresh nn.Linear) stayed on CPU while every other | |
| # module was on CUDA, causing | |
| # "Expected all tensors to be on the same device" only | |
| # when running on GPU (silently fine on CPU, since | |
| # everything was CPU there anyway -- which is exactly | |
| # why this didn't surface until a real GPU run). | |
| new_head.out.load_state_dict(old_head.out.state_dict()) | |
| model.istft_head = new_head | |
| # Explicit, not implicit: confirm causal=False on both transformers. | |
| # This matches the model's pretrained state already, but setting it | |
| # explicitly documents the requirement and protects against a future | |
| # change to MioCodecModelConfig defaults silently breaking this. | |
| set_transformer_causal(model.wave_prenet, False) | |
| set_transformer_causal(model.wave_decoder, False) | |
| return model |