"""BlueMagpie-TTS hybrid model. Architecture = VoxCPM2 with its Text-Semantic LM (MiniCPM4) swapped for Barbet: text tokens ──Barbet.embed_tokens──────────────┐ ├─ interleave ─> Barbet backbone ─> H_b audio latents ─LocEnc─┬─ enc_to_tslm_proj(H_b)─┘ │ └─ enc_to_lm_proj(H_v)──────────┐ tslm_adapter (H_b -> H_v) │ │ │ FSQ (audio positions) │ │ enc_outputs (H_v) fusion_concat_proj(cat) ┴─────────────┤ │ │ RALM (MiniCPM4 8L) lm_to_dit_proj │ res_to_dit_proj │ └───────── concat ──────────┘ │ mu (2 prefix tokens) LocDiT (CFM) │ latent patch AudioVAE Division of labour (matching VoxCPM2's design intent): - Barbet (TSLM): what to say, prosody planning, pacing, emphasis, control text. - RALM + LocDiT: fine-grained acoustic detail, kept verbatim from VoxCPM2 so pretrained weights load unchanged. - ``tslm_adapter`` is the explicit bridge between the two incompatible hidden spaces; ``enc_to_tslm_proj`` is its input-side counterpart feeding LocEnc features into Barbet. Everything outside the TSLM block mirrors ``voxcpm/model/voxcpm2.py`` (Apache-2.0, Copyright 2026 OpenBMB) so that VoxCPM2 checkpoints remain loadable. """ from __future__ import annotations import os import sys import warnings from typing import Generator, List, Optional, Tuple, Union import librosa import torch import torch.nn as nn from einops import rearrange from tqdm import tqdm from bluemagpie._vendor.voxcpm.model.utils import get_dtype, next_and_close, pick_runtime_dtype, resolve_runtime_device from bluemagpie._vendor.voxcpm.modules.audiovae import AudioVAEV2 from bluemagpie._vendor.voxcpm.modules.layers import ScalarQuantizationLayer from bluemagpie._vendor.voxcpm.modules.locdit import UnifiedCFM, VoxCPMLocDiTV2 from bluemagpie._vendor.voxcpm.modules.locenc import VoxCPMLocEnc from bluemagpie._vendor.voxcpm.modules.minicpm4 import MiniCPMModel from .adapter import ProjectionAdapter from .conditioning import SpeakerProjector from .config import BlueMagpieConfig from .tslm import BarbetTSLM class BlueMagpieModel(nn.Module): def __init__( self, config: BlueMagpieConfig, tokenizer=None, audio_vae: AudioVAEV2 = None, device: str | None = None, ): super().__init__() self.config = config self.feat_dim = config.feat_dim self.patch_size = config.patch_size self.device = resolve_runtime_device(device, config.device) self.config.device = self.device resolved_dtype = pick_runtime_dtype(self.device, self.config.dtype) if resolved_dtype != self.config.dtype: print( f"[bluemagpie] adjusted dtype {self.config.dtype} -> {resolved_dtype} for device {self.device}", file=sys.stderr, ) self.config.dtype = resolved_dtype vox_lm = config.vox_lm_config h_vox = vox_lm.hidden_size # ---------------- Text-Semantic LM: Barbet ---------------- # barbet_cfg, token_ids = config.resolve_barbet_config() self.barbet_config = barbet_cfg self.base_lm = BarbetTSLM(barbet_cfg) h_barbet = barbet_cfg.hidden_size self.text_tokenizer = tokenizer self.audio_start_token = token_ids["audio_start"] self.audio_end_token = token_ids["audio_end"] self.ref_audio_start_token = token_ids["ref_audio_start"] self.ref_audio_end_token = token_ids["ref_audio_end"] self.spk_token = token_ids["spk"] # ---------------- Residual Acoustic LM (VoxCPM2) ---------------- # residual_lm_config = vox_lm.model_copy(deep=True) residual_lm_config.num_hidden_layers = config.residual_lm_num_layers residual_lm_config.vocab_size = 0 residual_lm_config.no_rope = config.residual_lm_no_rope self.residual_lm = MiniCPMModel(residual_lm_config) self.residual_lm.setup_cache(1, config.max_length, self.device, get_dtype(self.config.dtype)) # ---------------- Local Encoder (VoxCPM2) ---------------- # encoder_config = vox_lm.model_copy(deep=True) encoder_config.hidden_size = config.encoder_config.hidden_dim encoder_config.intermediate_size = config.encoder_config.ffn_dim encoder_config.num_attention_heads = config.encoder_config.num_heads encoder_config.num_hidden_layers = config.encoder_config.num_layers encoder_config.kv_channels = config.encoder_config.kv_channels encoder_config.vocab_size = 0 self.feat_encoder = VoxCPMLocEnc(encoder_config, input_dim=config.feat_dim) # ---------------- Local DiT (VoxCPM2) ---------------- # decoder_config = vox_lm.model_copy(deep=True) decoder_config.hidden_size = config.dit_config.hidden_dim decoder_config.intermediate_size = config.dit_config.ffn_dim decoder_config.num_attention_heads = config.dit_config.num_heads decoder_config.num_hidden_layers = config.dit_config.num_layers decoder_config.kv_channels = config.dit_config.kv_channels decoder_config.vocab_size = 0 self.feat_decoder = UnifiedCFM( in_channels=config.feat_dim, cfm_params=config.dit_config.cfm_config, estimator=VoxCPMLocDiTV2(decoder_config, in_channels=config.feat_dim), mean_mode=config.dit_config.dit_mean_mode, ) # ---------------- VoxCPM2 projections (semantic space H_v) ---------------- # self.fsq_layer = ScalarQuantizationLayer( h_vox, h_vox, config.scalar_quantization_latent_dim, config.scalar_quantization_scale ) self.enc_to_lm_proj = nn.Linear(config.encoder_config.hidden_dim, h_vox) self.lm_to_dit_proj = nn.Linear(h_vox, config.dit_config.hidden_dim) self.res_to_dit_proj = nn.Linear(h_vox, config.dit_config.hidden_dim) self.fusion_concat_proj = nn.Linear(h_vox * 2, h_vox) # Stop Predictor (semantic space) self.stop_proj = nn.Linear(h_vox, h_vox) self.stop_actn = nn.SiLU() self.stop_head = nn.Linear(h_vox, 2, bias=False) self.stop_loss = nn.CrossEntropyLoss(reduction="none") # ---------------- Barbet <-> VoxCPM2 bridges (new, trained from scratch) -------- # self.enc_to_tslm_proj = nn.Linear(config.encoder_config.hidden_dim, h_barbet) self.tslm_adapter = ProjectionAdapter(h_barbet, h_vox, config.adapter_config) # ---------------- Speaker conditioning (v1: centroid -> [spk] slot) ------------ # self.speaker_embed_dim = config.speaker_embed_dim self.speaker_projector = SpeakerProjector(config.speaker_embed_dim, h_barbet) # ---------------- Audio VAE ---------------- # self.audio_vae = audio_vae if audio_vae is not None: self.chunk_size = audio_vae.chunk_size self._decode_chunk_size = getattr(audio_vae, "decode_chunk_size", audio_vae.chunk_size) self._encode_sample_rate = audio_vae.sample_rate self.sample_rate = getattr(audio_vae, "out_sample_rate", audio_vae.sample_rate) def _dtype(self): return get_dtype(self.config.dtype) def _runtime_dtype(self) -> torch.dtype: """Return the dtype the non-AudioVAE network is currently using.""" return next(self.parameters()).dtype def _runtime_device(self) -> torch.device: """Return the device the module is actually placed on. Distributed wrappers may move parameters after construction. Keep the cached device string in sync so Triton-backed kernels launch on the same CUDA device as their inputs. """ device = next(self.parameters()).device if device.type == "cuda" and device.index is not None and torch.cuda.current_device() != device.index: torch.cuda.set_device(device) device_str = str(device) if device_str != str(self.device): self.device = device_str self.config.device = device_str return device def _inject_speaker(self, combined_embed, speaker_centroids, spk_mask): """Add the projected speaker centroid at the [spk] position. ``combined_embed``: [B, T, H_b]. ``speaker_centroids``: [B, D] (the per-speaker ECAPA centroid). ``spk_mask``: [B, T], 1 only at the [spk] slot of samples whose speaker is present (not dropped). Where ``spk_mask`` is 0 the slot keeps its [spk] token embedding, which acts as the learned null speaker. """ if speaker_centroids is None or spk_mask is None: return combined_embed dtype = combined_embed.dtype spk_vec = self.speaker_projector(speaker_centroids.to(combined_embed.device, dtype=dtype)) # [B, H_b] spk_mask = spk_mask.to(combined_embed.device, dtype=dtype) return combined_embed + spk_mask.unsqueeze(-1) * spk_vec.unsqueeze(1) def _tokenize(self, text: str) -> List[int]: if self.text_tokenizer is None: raise ValueError("No tokenizer attached to BlueMagpieModel") if hasattr(self.text_tokenizer, "encode"): return self.text_tokenizer.encode(text, add_special_tokens=False) return self.text_tokenizer(text) # ------------------------------------------------------------------ # # Training forward # ------------------------------------------------------------------ # def forward( self, text_tokens: torch.Tensor, text_mask: torch.Tensor, audio_feats: torch.Tensor, audio_mask: torch.Tensor, loss_mask: torch.Tensor, position_ids: torch.Tensor, labels: torch.Tensor, *, speaker_centroids: Optional[torch.Tensor] = None, spk_mask: Optional[torch.Tensor] = None, progress: float = 0.0, sample_generate: bool = False, sample_generate_timesteps: int = 10, ): del position_ids # not used (parity with VoxCPM2) device = self._runtime_device() text_tokens = text_tokens.to(device, dtype=torch.long) text_mask = text_mask.to(device, dtype=self._dtype()) audio_feats = audio_feats.to(device, dtype=self._dtype()) audio_mask = audio_mask.to(device, dtype=self._dtype()) loss_mask = loss_mask.to(device, dtype=self._dtype()) labels = labels.to(device, dtype=torch.long) B, T, P, D = audio_feats.shape feat_locenc = self.feat_encoder(audio_feats) feat_embed_tslm = self.enc_to_tslm_proj(feat_locenc) # Barbet input space feat_embed_lm = self.enc_to_lm_proj(feat_locenc) # VoxCPM2 semantic space (RALM fusion) text_embed = self.base_lm.embed_tokens(text_tokens) combined_embed = text_mask.unsqueeze(-1) * text_embed + audio_mask.unsqueeze(-1) * feat_embed_tslm combined_embed = self._inject_speaker(combined_embed, speaker_centroids, spk_mask) barbet_hidden = self.base_lm(inputs_embeds=combined_embed) tslm_hidden = self.tslm_adapter(barbet_hidden).to(self._dtype()) enc_outputs = self.fsq_layer(tslm_hidden) * audio_mask.unsqueeze(-1) + tslm_hidden * text_mask.unsqueeze(-1) lm_hidden = torch.cat((torch.zeros_like(enc_outputs[:, 0:1, :]), enc_outputs[:, :-1, :]), dim=1) residual_inputs = self.fusion_concat_proj( torch.cat((enc_outputs, audio_mask.unsqueeze(-1) * feat_embed_lm), dim=-1) ) residual_outputs, _ = self.residual_lm(inputs_embeds=residual_inputs, is_causal=True) residual_outputs = residual_outputs.to(self._dtype()) residual_hidden = torch.cat( (torch.zeros_like(residual_outputs[:, 0:1, :]), residual_outputs[:, :-1, :]), dim=1, ) dit_hidden = torch.cat((self.lm_to_dit_proj(lm_hidden), self.res_to_dit_proj(residual_hidden)), dim=-1) dit_hidden = rearrange(dit_hidden, "b t c -> (b t) c") target_dtype = self._dtype() feat_gt = rearrange(audio_feats.to(target_dtype), "b t p d -> (b t) p d") feat_cond = torch.cat( (torch.zeros_like(audio_feats[:, 0:1, ...]), audio_feats[:, :-1, ...]), dim=1, ) feat_cond = rearrange(feat_cond.to(target_dtype), "b t p d -> (b t) p d") loss_seq_mask = loss_mask.unsqueeze(-1).repeat(1, 1, self.patch_size) loss_seq_mask = rearrange(loss_seq_mask, "b t p -> (b t) p 1").to(target_dtype) diff_loss = self.feat_decoder.compute_loss( feat_gt.transpose(1, 2).contiguous(), dit_hidden, cond=feat_cond.transpose(1, 2).contiguous(), tgt_mask=loss_seq_mask.transpose(1, 2).contiguous(), progress=progress, ) stop_logits = self.stop_head(self.stop_actn(self.stop_proj(lm_hidden))) stop_losses = self.stop_loss(stop_logits.transpose(1, 2), labels) denom = torch.clamp(loss_mask.sum(), min=1.0) stop_loss = (stop_losses * loss_mask).sum() / denom feat_pred = None if sample_generate: feat_cond_for_sample = feat_cond.transpose(1, 2).contiguous() feat_pred_seq = self.feat_decoder( mu=dit_hidden, patch_size=self.patch_size, cond=feat_cond_for_sample, n_timesteps=max(int(sample_generate_timesteps), 1), ) feat_pred = rearrange(feat_pred_seq.transpose(1, 2), "(b t) p d -> b d (t p)", b=B, p=self.patch_size) feat_gt_tensor = rearrange(feat_gt, "(b t) p d -> b d (t p)", b=B, p=self.patch_size) return { "loss/diff": diff_loss, "loss/stop": stop_loss, "feat_gt": feat_gt_tensor, "feat_pred": feat_pred, "stop_logits": stop_logits, # Pre-FSQ adapted hidden states, for hidden-space distillation # against the original VoxCPM2 TSLM (see README, stage 0). "tslm_hidden": tslm_hidden, } # ------------------------------------------------------------------ # # Input assembly (the four VoxCPM2 prompting modes) # ------------------------------------------------------------------ # def _encode_wav( self, wav_path: str, padding_mode: str = "right", ) -> torch.Tensor: audio, _ = librosa.load(wav_path, sr=self._encode_sample_rate, mono=True) audio = torch.from_numpy(audio).unsqueeze(0) patch_len = self.patch_size * self.chunk_size if audio.size(1) % patch_len != 0: padding_size = patch_len - audio.size(1) % patch_len pad = (padding_size, 0) if padding_mode == "left" else (0, padding_size) audio = torch.nn.functional.pad(audio, pad) device = self._runtime_device() feat = self.audio_vae.encode(audio.to(device), self._encode_sample_rate).cpu() return feat.view(self.audio_vae.latent_dim, -1, self.patch_size).permute(1, 2, 0) def _make_ref_prefix(self, ref_feat: torch.Tensor, device: torch.device): ref_len = ref_feat.size(0) z1 = torch.zeros((1, self.patch_size, self.audio_vae.latent_dim), dtype=torch.float32, device=device) tokens = torch.cat( [ torch.tensor([self.ref_audio_start_token], dtype=torch.int32, device=device), torch.zeros(ref_len, dtype=torch.int32, device=device), torch.tensor([self.ref_audio_end_token], dtype=torch.int32, device=device), ] ) feats = torch.cat([z1, ref_feat, z1], dim=0) t_mask = torch.cat( [ torch.tensor([1], dtype=torch.int32), torch.zeros(ref_len, dtype=torch.int32), torch.tensor([1], dtype=torch.int32), ] ).to(device) a_mask = torch.cat( [ torch.tensor([0], dtype=torch.int32), torch.ones(ref_len, dtype=torch.int32), torch.tensor([0], dtype=torch.int32), ] ).to(device) return tokens, feats, t_mask, a_mask def _build_inputs( self, text: str, ref_feat: Optional[torch.Tensor] = None, prompt_feat: Optional[torch.Tensor] = None, speaker_slot: str = "none", ): """Assemble (text_token, audio_feat, text_mask, audio_mask, spk_mask). Layout: [spk?] [ref prefix?] [text + audio_start] [prompt audio?] ``speaker_slot`` is "none", "null", or "centroid". The null slot keeps the learned [spk] token embedding, matching speaker-dropout training. The centroid slot is filled by the projected speaker vector. """ if speaker_slot not in {"none", "null", "centroid"}: raise ValueError(f"speaker_slot must be one of none/null/centroid, got {speaker_slot!r}") text_token = torch.LongTensor(self._tokenize(text)) text_token = torch.cat( [ text_token, torch.tensor([self.audio_start_token], dtype=torch.long, device=text_token.device), ], dim=-1, ) text_length = text_token.shape[0] device = text_token.device text_pad_feat = torch.zeros( (text_length, self.patch_size, self.audio_vae.latent_dim), dtype=torch.float32, device=device, ) tokens = [text_token] feats = [text_pad_feat] t_masks = [torch.ones(text_length, dtype=torch.int32, device=device)] a_masks = [torch.zeros(text_length, dtype=torch.int32, device=device)] s_masks = [torch.zeros(text_length, dtype=torch.int32, device=device)] if ref_feat is not None: ref_tokens, ref_feats, ref_t_mask, ref_a_mask = self._make_ref_prefix(ref_feat, device) tokens.insert(0, ref_tokens.long()) feats.insert(0, ref_feats) t_masks.insert(0, ref_t_mask) a_masks.insert(0, ref_a_mask) s_masks.insert(0, torch.zeros(ref_tokens.shape[0], dtype=torch.int32, device=device)) if speaker_slot != "none": one_feat = torch.zeros((1, self.patch_size, self.audio_vae.latent_dim), dtype=torch.float32, device=device) tokens.insert(0, torch.tensor([self.spk_token], dtype=torch.long, device=device)) feats.insert(0, one_feat) t_masks.insert( 0, torch.ones(1, dtype=torch.int32, device=device) if speaker_slot == "null" else torch.zeros(1, dtype=torch.int32, device=device), ) a_masks.insert(0, torch.zeros(1, dtype=torch.int32, device=device)) s_masks.insert( 0, torch.ones(1, dtype=torch.int32, device=device) if speaker_slot == "centroid" else torch.zeros(1, dtype=torch.int32, device=device), ) if prompt_feat is not None: prompt_len = prompt_feat.size(0) tokens.append(torch.zeros(prompt_len, dtype=torch.long, device=device)) feats.append(prompt_feat) t_masks.append(torch.zeros(prompt_len, dtype=torch.int32, device=device)) a_masks.append(torch.ones(prompt_len, dtype=torch.int32, device=device)) s_masks.append(torch.zeros(prompt_len, dtype=torch.int32, device=device)) return ( torch.cat(tokens), torch.cat(feats, dim=0), torch.cat(t_masks), torch.cat(a_masks), torch.cat(s_masks), ) # ------------------------------------------------------------------ # # Generation # ------------------------------------------------------------------ # def generate(self, *args, **kwargs) -> torch.Tensor: return next_and_close(self._generate(*args, streaming=False, **kwargs)) def generate_streaming(self, *args, **kwargs) -> Generator[torch.Tensor, None, None]: return self._generate(*args, streaming=True, **kwargs) @torch.inference_mode() def _generate( self, target_text: str, prompt_text: str = "", prompt_wav_path: str = "", reference_wav_path: str = "", speaker_centroid: Optional[torch.Tensor] = None, min_len: int = 2, max_len: int = 2000, inference_timesteps: int = 10, cfg_value: float = 2.0, retry_badcase: bool = False, retry_badcase_max_times: int = 3, retry_badcase_ratio_threshold: float = 6.0, use_null_speaker: bool = True, streaming: bool = False, streaming_prefix_len: int = 4, ) -> Generator[torch.Tensor, None, None]: if retry_badcase and streaming: warnings.warn("Retry on bad cases is not supported in streaming mode, setting retry_badcase=False.") retry_badcase = False device = self._runtime_device() ref_feat = self._encode_wav(reference_wav_path, padding_mode="right") if reference_wav_path else None prompt_feat = self._encode_wav(prompt_wav_path, padding_mode="left") if prompt_wav_path else None text = (prompt_text + target_text) if prompt_feat is not None else target_text speaker_centroids = None if speaker_centroid is not None: speaker_centroids = speaker_centroid.reshape(1, -1).to(device, dtype=self._runtime_dtype()) speaker_slot = "centroid" if speaker_centroids is not None else ("null" if use_null_speaker else "none") text_token, audio_feat, text_mask, audio_mask, spk_mask = self._build_inputs( text, ref_feat, prompt_feat, speaker_slot=speaker_slot ) text_token = text_token.unsqueeze(0).to(device) text_mask = text_mask.unsqueeze(0).to(device, dtype=self._runtime_dtype()) audio_feat = audio_feat.unsqueeze(0).to(device, dtype=self._runtime_dtype()) audio_mask = audio_mask.unsqueeze(0).to(device, dtype=self._runtime_dtype()) spk_mask = spk_mask.unsqueeze(0).to(device, dtype=self._runtime_dtype()) target_text_length = len(self._tokenize(target_text)) retry_badcase_times = 0 while retry_badcase_times < retry_badcase_max_times: inference_result = self._inference( text_token, text_mask, audio_feat, audio_mask, min_len=min_len, max_len=min(int(target_text_length * retry_badcase_ratio_threshold + 10), max_len), inference_timesteps=inference_timesteps, cfg_value=cfg_value, speaker_centroids=speaker_centroids, spk_mask=spk_mask, streaming=streaming, streaming_prefix_len=streaming_prefix_len, ) if streaming: with self.audio_vae.streaming_decode() as vae_dec: for latent_pred, _, _ctx in inference_result: decode_audio = vae_dec.decode_chunk(latent_pred.to(torch.float32)) yield decode_audio.squeeze(1).cpu() break else: latent_pred, pred_audio_feat, context_len = next_and_close(inference_result) if retry_badcase and pred_audio_feat.shape[0] >= target_text_length * retry_badcase_ratio_threshold: print( f" Badcase detected, audio_text_ratio={pred_audio_feat.shape[0] / target_text_length}, retrying...", file=sys.stderr, ) retry_badcase_times += 1 continue break if not streaming: decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)) decode_patch_len = self.patch_size * self._decode_chunk_size if context_len > 0: decode_audio = decode_audio[..., decode_patch_len * context_len :].squeeze(1).cpu() else: decode_audio = decode_audio.squeeze(1).cpu() yield decode_audio # ------------------------------------------------------------------ # # Core AR loop # ------------------------------------------------------------------ # @torch.inference_mode() def _inference( self, text: torch.Tensor, text_mask: torch.Tensor, feat: torch.Tensor, feat_mask: torch.Tensor, min_len: int = 2, max_len: int = 2000, inference_timesteps: int = 10, cfg_value: float = 2.0, speaker_centroids: Optional[torch.Tensor] = None, spk_mask: Optional[torch.Tensor] = None, streaming: bool = False, streaming_prefix_len: int = 4, ) -> Generator[Tuple[torch.Tensor, Union[torch.Tensor, List[torch.Tensor]], int], None, None]: device = self._runtime_device() dtype = self._runtime_dtype() text = text.to(device) text_mask = text_mask.to(device, dtype=dtype) feat = feat.to(device, dtype=dtype) feat_mask = feat_mask.to(device, dtype=dtype) if speaker_centroids is not None: speaker_centroids = speaker_centroids.to(device, dtype=dtype) if spk_mask is not None: spk_mask = spk_mask.to(device, dtype=dtype) B, T, P, D = feat.shape feat_locenc = self.feat_encoder(feat) # [b, t, h_enc] feat_embed_tslm = self.enc_to_tslm_proj(feat_locenc) feat_embed_lm = self.enc_to_lm_proj(feat_locenc) text_embed = self.base_lm.embed_tokens(text) combined_embed = text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed_tslm combined_embed = self._inject_speaker(combined_embed, speaker_centroids, spk_mask) prefix_feat_cond = feat[:, -1, ...] # b, p, d curr_embed = None # Streaming context patches (continuation modes only) has_continuation_audio = feat_mask[0, -1].item() == 1 context_len = 0 if has_continuation_audio: audio_indices = feat_mask.squeeze(0).nonzero(as_tuple=True)[0] context_len = min(streaming_prefix_len - 1, len(audio_indices)) last_audio_indices = audio_indices[-context_len:] pred_feat_seq = list(feat[:, last_audio_indices, :, :].split(1, dim=1)) else: pred_feat_seq = [] # --- TSLM prefill (Barbet stepwise cache) --- # barbet_hidden_seq, tslm_state = self.base_lm.prefill(combined_embed) tslm_hidden_seq = self.tslm_adapter(barbet_hidden_seq) enc_outputs = ( self.fsq_layer(tslm_hidden_seq) * feat_mask.unsqueeze(-1) + tslm_hidden_seq * text_mask.unsqueeze(-1) ) lm_hidden = enc_outputs[:, -1, :] # --- RALM prefill (MiniCPM static KV cache) --- # residual_enc_inputs = self.fusion_concat_proj( torch.cat((enc_outputs, feat_mask.unsqueeze(-1) * feat_embed_lm), dim=-1) ) residual_enc_outputs, residual_kv_cache_tuple = self.residual_lm( inputs_embeds=residual_enc_inputs, is_causal=True, ) self.residual_lm.kv_cache.fill_caches(residual_kv_cache_tuple) residual_hidden = residual_enc_outputs[:, -1, :] for i in tqdm(range(max_len)): dit_hidden_1 = self.lm_to_dit_proj(lm_hidden) # [b, h_dit] dit_hidden_2 = self.res_to_dit_proj(residual_hidden) # [b, h_dit] dit_hidden = torch.cat((dit_hidden_1, dit_hidden_2), dim=-1) pred_feat = self.feat_decoder( mu=dit_hidden, patch_size=self.patch_size, cond=prefix_feat_cond.transpose(1, 2).contiguous(), n_timesteps=inference_timesteps, cfg_value=cfg_value, ).transpose(1, 2) # [b, p, d] curr_locenc = self.feat_encoder(pred_feat.unsqueeze(1)) # b, 1, h_enc curr_embed_tslm = self.enc_to_tslm_proj(curr_locenc) curr_embed_lm = self.enc_to_lm_proj(curr_locenc) curr_embed = curr_embed_tslm # naming parity with VoxCPM2 pred_feat_seq.append(pred_feat.unsqueeze(1)) # b, 1, p, d prefix_feat_cond = pred_feat if streaming: feat_pred = rearrange(pred_feat.unsqueeze(1), "b t p d -> b d (t p)", b=B, p=self.patch_size) yield feat_pred, pred_feat_seq, context_len if len(pred_feat_seq) > streaming_prefix_len: pred_feat_seq = pred_feat_seq[-streaming_prefix_len:] stop_flag = self.stop_head(self.stop_actn(self.stop_proj(lm_hidden))).argmax(dim=-1)[0].cpu().item() if i > min_len and stop_flag == 1: break barbet_step_hidden = self.base_lm.forward_step(curr_embed[:, 0, :], tslm_state) lm_hidden = self.fsq_layer(self.tslm_adapter(barbet_step_hidden)) curr_residual_input = self.fusion_concat_proj(torch.cat((lm_hidden, curr_embed_lm[:, 0, :]), dim=-1)) residual_hidden = self.residual_lm.forward_step( curr_residual_input, torch.tensor([self.residual_lm.kv_cache.step()], device=curr_residual_input.device), ).clone() if not streaming: pred_feat_seq = torch.cat(pred_feat_seq, dim=1) # b, t, p, d feat_pred = rearrange(pred_feat_seq, "b t p d -> b d (t p)", b=B, p=self.patch_size) generated_feat = pred_feat_seq[:, context_len:, :, :].squeeze(0).cpu() yield feat_pred, generated_feat, context_len def inference(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]: feat_pred, generated_feat, _ = next_and_close(self._inference(*args, streaming=False, **kwargs)) return feat_pred, generated_feat # ------------------------------------------------------------------ # # Checkpoint I/O (BlueMagpie's own format) # ------------------------------------------------------------------ # def save_pretrained(self, path: str): os.makedirs(path, exist_ok=True) with open(os.path.join(path, "config.json"), "w", encoding="utf-8") as f: f.write(self.config.model_dump_json(indent=2)) state = {k: v for k, v in self.state_dict().items() if not k.startswith("audio_vae.")} torch.save(state, os.path.join(path, "pytorch_model.bin")) if self.audio_vae is not None: torch.save(self.audio_vae.state_dict(), os.path.join(path, "audiovae.pth")) if self.text_tokenizer is not None and hasattr(self.text_tokenizer, "save_pretrained"): self.text_tokenizer.save_pretrained(path) @classmethod def from_local( cls, path: str, tokenizer=None, training: bool = False, device: str | None = None, ) -> "BlueMagpieModel": from bluemagpie._vendor.voxcpm.modules.audiovae import AudioVAEV2 with open(os.path.join(path, "config.json"), "r", encoding="utf-8") as f: config = BlueMagpieConfig.model_validate_json(f.read()) if tokenizer is None: try: from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(path) except Exception: tokenizer = None audio_vae = AudioVAEV2(config=config.audio_vae_config) if config.audio_vae_config else AudioVAEV2() vae_path = os.path.join(path, "audiovae.pth") if os.path.exists(vae_path): vae_state = torch.load(vae_path, map_location="cpu", weights_only=True) audio_vae.load_state_dict(vae_state.get("state_dict", vae_state)) model = cls(config, tokenizer, audio_vae, device=device) # mmap keeps the checkpoint memory-mapped instead of materializing a # second full copy in RAM (load peak ~halved; matters on Colab-class hosts). state = torch.load( os.path.join(path, "pytorch_model.bin"), map_location="cpu", weights_only=True, mmap=True ) missing, unexpected = model.load_state_dict(state, strict=False) missing = [k for k in missing if not k.startswith("audio_vae.")] if missing or unexpected: print(f"[bluemagpie] missing keys: {missing[:8]}... unexpected: {unexpected[:8]}...", file=sys.stderr) if not training: model = model.to(get_dtype(model.config.dtype)) model.audio_vae = model.audio_vae.to(torch.float32) return model.to(model.device).eval() if not training else model.to(model.device)