SFlowTTS / Modules /ar_style_predictor.py

Upload full repo excluding dump_40, dump_100, precomputed_tokens, precomputed_data

fb0011a verified 8 months ago

16.5 kB

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

	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 2num_levels to 3num_levels-1
	idx_pr_flat = flat_indices[:, 2num_levels:3num_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: [Bbeam_width] -> [Bbeam_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