# # 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 @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 @property 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