import torch import torch.nn as nn import torch.nn.functional as F from .gemma_wrapper import GemmaWrapper from .fsq import FSQ class ARStylePredictor(nn.Module): def __init__(self, gemma_model_name="google/gemma-3-4b-it", style_dim=128, pitch_dim=64, prosodic_dim=64, fsq_levels=[8, 8, 4], hidden_dim=512, num_layers=4, nhead=8): super().__init__() self.gemma = GemmaWrapper(model_name=gemma_model_name) self.fsq_levels = fsq_levels self.style_dim = style_dim self.pitch_dim = pitch_dim self.prosodic_dim = prosodic_dim # FSQ modules self.fsq_acoustic = FSQ(fsq_levels) self.fsq_pitch = FSQ(fsq_levels) self.fsq_prosodic = FSQ(fsq_levels) # Encoders (Style -> FSQ Space) # These project the high-dim style vector to the low-dim FSQ space (e.g. 3 dims) self.proj_acoustic = nn.Linear(style_dim, len(fsq_levels)) self.proj_pitch = nn.Linear(pitch_dim, len(fsq_levels)) self.proj_prosodic = nn.Linear(prosodic_dim, len(fsq_levels)) # Decoders (FSQ Space -> Style) - For stable reconstruction # We use separate linear layers to map back from FSQ values to style vectors. # This ensures we can learn a stable reconstruction even if the FSQ space is fragile. self.dec_acoustic = nn.Linear(len(fsq_levels), style_dim) self.dec_pitch = nn.Linear(len(fsq_levels), pitch_dim) self.dec_prosodic = nn.Linear(len(fsq_levels), prosodic_dim) # Calculate vocabulary size and offsets self.offsets = [] current_offset = 0 # Order: Acoustic -> Pitch -> Prosodic self.styles = ['acoustic', 'pitch', 'prosodic'] self.style_dims = {'acoustic': style_dim, 'pitch': pitch_dim, 'prosodic': prosodic_dim} # We flatten the levels: Acoustic L1, L2, L3, Pitch L1, L2, L3, Prosodic L1, L2, L3 for style in self.styles: for level in fsq_levels: self.offsets.append(current_offset) current_offset += level self.vocab_size = current_offset self.sos_token = self.vocab_size self.eos_token = self.vocab_size + 1 self.full_vocab_size = self.vocab_size + 2 # AR Transformer self.embedding = nn.Embedding(self.full_vocab_size, hidden_dim) self.pos_embedding = nn.Parameter(torch.randn(1, len(self.offsets) + 1, hidden_dim)) # +1 for SOS # Text projection config = self.gemma.model.config if hasattr(config, 'text_config') and config.text_config is not None: gemma_hidden_size = config.text_config.hidden_size else: gemma_hidden_size = config.hidden_size self.text_proj = nn.Linear(gemma_hidden_size, hidden_dim) decoder_layer = nn.TransformerDecoderLayer(d_model=hidden_dim, nhead=nhead, dim_feedforward=2048, batch_first=True, norm_first=True) self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers) # Prediction Head self.head = nn.Linear(hidden_dim, self.full_vocab_size) def get_fsq_indices(self, acoustic, pitch, prosodic): # Project and quantize to get indices # Acoustic z_a = self.proj_acoustic(acoustic) _, idx_a = self.fsq_acoustic.quantize(z_a) # [B, 3] # Pitch z_p = self.proj_pitch(pitch) _, idx_p = self.fsq_pitch.quantize(z_p) # [B, 3] # Prosodic z_pr = self.proj_prosodic(prosodic) _, idx_pr = self.fsq_prosodic.quantize(z_pr) # [B, 3] # Flatten and add offsets # idx_a: [B, 3]. We need to add offsets[0], offsets[1], offsets[2] indices = [] offset_idx = 0 # Acoustic for i in range(idx_a.shape[1]): indices.append(idx_a[:, i] + self.offsets[offset_idx]) offset_idx += 1 # Pitch for i in range(idx_p.shape[1]): indices.append(idx_p[:, i] + self.offsets[offset_idx]) offset_idx += 1 # Prosodic for i in range(idx_pr.shape[1]): indices.append(idx_pr[:, i] + self.offsets[offset_idx]) offset_idx += 1 return torch.stack(indices, dim=1) # [B, 9] def reconstruct(self, flat_indices): # flat_indices: [B, 9] (or length of total levels) # We assume the indices are already offset-adjusted (or we need to subtract them?) # If they come from the model prediction (argmax), they are in [0, full_vocab_size]. # We need to subtract offsets to get back to [0, L-1]. # Split into styles num_levels = len(self.fsq_levels) # Acoustic: indices 0 to num_levels-1 idx_a_flat = flat_indices[:, 0:num_levels] # Pitch: indices num_levels to 2*num_levels-1 idx_p_flat = flat_indices[:, num_levels:2*num_levels] # Prosodic: indices 2*num_levels to 3*num_levels-1 idx_pr_flat = flat_indices[:, 2*num_levels:3*num_levels] def decode_style(indices_flat, fsq_module, dec_module, start_offset_idx): # indices_flat: [B, num_levels] # Subtract offsets indices = [] for i in range(num_levels): offset = self.offsets[start_offset_idx + i] # Ensure we don't have negative indices if prediction was wrong (shouldn't happen with valid masking) val = indices_flat[:, i] - offset val = torch.clamp(val, 0, self.fsq_levels[i] - 1) indices.append(val) indices = torch.stack(indices, dim=1) # [B, num_levels] # Get FSQ values fsq_values = fsq_module.indices_to_values(indices) # [B, num_levels] # Decode to style vector recon = dec_module(fsq_values) return recon acoustic = decode_style(idx_a_flat, self.fsq_acoustic, self.dec_acoustic, 0) pitch = decode_style(idx_p_flat, self.fsq_pitch, self.dec_pitch, num_levels) prosodic = decode_style(idx_pr_flat, self.fsq_prosodic, self.dec_prosodic, 2*num_levels) return acoustic, pitch, prosodic def forward(self, texts, acoustic=None, pitch=None, prosodic=None, return_loss=True): # texts: list of strings # acoustic, pitch, prosodic: [B, dim] - Ground truth styles # 1. Get Text Embeddings gemma_out, gemma_mask = self.gemma(texts) # [B, seq_len, gemma_dim] memory = self.text_proj(gemma_out) # [B, seq_len, hidden_dim] if return_loss: if acoustic is None or pitch is None or prosodic is None: raise ValueError("Must provide styles for training") # 2. Get Target Indices target_indices = self.get_fsq_indices(acoustic, pitch, prosodic) # [B, 9] # 3. Prepare Input Sequence (SOS + Target) B = target_indices.size(0) sos = torch.full((B, 1), self.sos_token, device=target_indices.device, dtype=torch.long) input_seq = torch.cat([sos, target_indices], dim=1) # [B, 10] # 4. Embeddings tgt_emb = self.embedding(input_seq[:, :-1]) # [B, 9, hidden_dim] # Add positional embeddings # We use fixed positions 0..8 positions = torch.arange(tgt_emb.size(1), device=tgt_emb.device).unsqueeze(0) tgt_emb = tgt_emb + self.pos_embedding[:, :tgt_emb.size(1), :] # 5. Transformer Decoder # Causal Mask sz = tgt_emb.size(1) tgt_mask = torch.triu(torch.ones(sz, sz, device=tgt_emb.device) * float('-inf'), diagonal=1) # Memory Mask (padding mask from Gemma) # gemma_mask is True for valid tokens. Transformer expects True for ignored tokens? # PyTorch Transformer: key_padding_mask: True for values to be ignored. # Gemma mask: 1 for valid, 0 for padding. memory_key_padding_mask = ~gemma_mask output = self.transformer_decoder( tgt_emb, memory, tgt_mask=tgt_mask, memory_key_padding_mask=memory_key_padding_mask ) # 6. Prediction logits = self.head(output) # [B, 9, vocab_size] # 7. Loss # We need to mask logits so that at position i, we only predict valid tokens for that level # But standard CrossEntropy handles this if we just pass the targets. # However, to help convergence, we can mask invalid logits. loss = F.cross_entropy(logits.reshape(-1, self.full_vocab_size), target_indices.reshape(-1)) return loss, logits else: # Inference return self.inference(texts) def inference(self, texts, beam_width=8): # 1. Get Text Embeddings gemma_out, gemma_mask = self.gemma(texts) memory = self.text_proj(gemma_out) memory_key_padding_mask = ~gemma_mask B = len(texts) device = memory.device num_tokens = len(self.offsets) if beam_width == 1: # Start with SOS curr_seq = torch.full((B, 1), self.sos_token, device=device, dtype=torch.long) for i in range(num_tokens): tgt_emb = self.embedding(curr_seq) tgt_emb = tgt_emb + self.pos_embedding[:, :curr_seq.size(1), :] # Causal Mask sz = tgt_emb.size(1) tgt_mask = torch.triu(torch.ones(sz, sz, device=device) * float('-inf'), diagonal=1) output = self.transformer_decoder( tgt_emb, memory, tgt_mask=tgt_mask, memory_key_padding_mask=memory_key_padding_mask ) # Get last token logits last_output = output[:, -1, :] # [B, hidden_dim] logits = self.head(last_output) # [B, vocab_size] # Mask invalid tokens for this step # The valid tokens are [offset, offset + level_size) offset = self.offsets[i] level_size = self.fsq_levels[i % len(self.fsq_levels)] # Assuming repeating pattern 8,8,4 # Create a mask mask = torch.ones_like(logits) * float('-inf') mask[:, offset : offset + level_size] = 0 logits = logits + mask # Greedy decoding (or sample) next_token = torch.argmax(logits, dim=-1).unsqueeze(1) curr_seq = torch.cat([curr_seq, next_token], dim=1) # Remove SOS generated_indices = curr_seq[:, 1:] else: # Beam Search # Expand memory: [B, S, D] -> [B * beam_width, S, D] memory = memory.repeat_interleave(beam_width, dim=0) memory_key_padding_mask = memory_key_padding_mask.repeat_interleave(beam_width, dim=0) # Initial sequences: [B * beam_width, 1] curr_seq = torch.full((B * beam_width, 1), self.sos_token, device=device, dtype=torch.long) # Scores: [B, beam_width] # Initialize first beam with 0, others with -inf to force selection from first beam initially scores = torch.full((B, beam_width), float('-inf'), device=device) scores[:, 0] = 0.0 scores = scores.view(-1) # [B * beam_width] for i in range(num_tokens): tgt_emb = self.embedding(curr_seq) tgt_emb = tgt_emb + self.pos_embedding[:, :curr_seq.size(1), :] sz = tgt_emb.size(1) tgt_mask = torch.triu(torch.ones(sz, sz, device=device) * float('-inf'), diagonal=1) output = self.transformer_decoder( tgt_emb, memory, tgt_mask=tgt_mask, memory_key_padding_mask=memory_key_padding_mask ) last_output = output[:, -1, :] # [B*beam_width, hidden_dim] logits = self.head(last_output) # [B*beam_width, vocab_size] # Mask invalid tokens offset = self.offsets[i] level_size = self.fsq_levels[i % len(self.fsq_levels)] mask = torch.ones_like(logits) * float('-inf') mask[:, offset : offset + level_size] = 0 log_probs = F.log_softmax(logits + mask, dim=-1) # [B*beam_width, vocab_size] # Add previous scores # scores: [B*beam_width] -> [B*beam_width, 1] total_scores = log_probs + scores.unsqueeze(1) # [B*beam_width, vocab_size] # Reshape to select topk per batch item # [B, beam_width, vocab_size] total_scores = total_scores.view(B, beam_width, -1) # Flatten beam and vocab dimensions: [B, beam_width * vocab_size] total_scores = total_scores.view(B, -1) # Select top k top_scores, top_indices = torch.topk(total_scores, beam_width, dim=1) # [B, beam_width] # Update scores scores = top_scores.view(-1) # Recover beam index and token index # top_indices is into [beam_width * vocab_size] beam_indices = top_indices // self.full_vocab_size # [B, beam_width] token_indices = top_indices % self.full_vocab_size # [B, beam_width] # Construct new sequences # We need to select the correct previous sequences based on beam_indices # curr_seq: [B * beam_width, len] # Reshape curr_seq to [B, beam_width, len] curr_seq = curr_seq.view(B, beam_width, -1) # Gather the correct beams # beam_indices: [B, beam_width]. We want to select from dim 1 of curr_seq # beam_indices needs to be expanded to [B, beam_width, len] beam_indices_expanded = beam_indices.unsqueeze(2).expand(-1, -1, curr_seq.size(2)) new_seq = torch.gather(curr_seq, 1, beam_indices_expanded) # [B, beam_width, len] # Append new tokens token_indices = token_indices.unsqueeze(2) # [B, beam_width, 1] new_seq = torch.cat([new_seq, token_indices], dim=2) # [B, beam_width, len+1] # Flatten back curr_seq = new_seq.view(B * beam_width, -1) # Select best beam (index 0 since topk sorts them) curr_seq = curr_seq.view(B, beam_width, -1) best_seq = curr_seq[:, 0, :] # [B, len] generated_indices = best_seq[:, 1:] # Remove SOS # Reconstruct acoustic, pitch, prosodic = self.reconstruct(generated_indices) return acoustic, pitch, prosodic def get_reconstruction_loss(self, acoustic, pitch, prosodic): # Acoustic z_a = self.proj_acoustic(acoustic) q_a, _ = self.fsq_acoustic.quantize(z_a) rec_a = self.dec_acoustic(q_a) loss_a = F.mse_loss(rec_a, acoustic) # Pitch z_p = self.proj_pitch(pitch) q_p, _ = self.fsq_pitch.quantize(z_p) rec_p = self.dec_pitch(q_p) loss_p = F.mse_loss(rec_p, pitch) # Prosodic z_pr = self.proj_prosodic(prosodic) q_pr, _ = self.fsq_prosodic.quantize(z_pr) rec_pr = self.dec_prosodic(q_pr) loss_pr = F.mse_loss(rec_pr, prosodic) return loss_a + loss_p + loss_pr