from dataclasses import dataclass from typing import Optional, Tuple, Dict, List import math import torch import torch.nn as nn import torch.nn.functional as F try: from .architecture import ( RMSNorm, precompute_rope_freqs, apply_rope, TransformerBlock, ) except (ImportError, ValueError): from architecture import ( RMSNorm, precompute_rope_freqs, apply_rope, TransformerBlock, ) @dataclass class TTSConfig: # Sözlük ve Codebook Parametreleri text_vocab_size: int = 151936 # Qwen2.5 Tokenizer Vocab Size (Breeze-TTS standardı) audio_vocab_size: int = 2048 # Kyutai Mimi her codebook için 2048 token num_codebooks: int = 8 # Kyutai Mimi 8 codebook # Model Boyutları (~35M - 45M parametre) d_model: int = 512 num_heads: int = 8 num_kv_heads: int = 4 # GQA num_layers: int = 8 # Main AR Backbone num_depth_layers: int = 4 # Depth Decoder d_ff: int = 1536 # ~3x d_model (SwiGLU) max_seq_len: int = 2048 rope_theta: float = 10000.0 dropout_rate: float = 0.0 activation: str = "swiglu" dtype: str = "float32" use_qk_norm: bool = False # Qwen2.5'te QK-Norm yoktur; tam uyum için False # Özel Token ID'leri pad_token_id: int = 0 bos_token_id: int = 1 eos_token_id: int = 2 unk_token_id: int = 3 instruct_token_id: int = 4 text_token_id: int = 5 ref_audio_token_id: int = 6 audio_start_token_id: int = 7 audio_end_token_id: int = 8 def __init__(self, **kwargs): valid = {f.name for f in self.__dataclass_fields__.values()} for k, v in kwargs.items(): if k in valid: setattr(self, k, v) @property def torch_dtype(self) -> torch.dtype: if self.dtype == "bfloat16": return torch.bfloat16 elif self.dtype == "float16": return torch.float16 return torch.float32 @classmethod def from_qwen(cls, model_id: str = "Qwen/Qwen2.5-0.5B", **kwargs) -> "TTSConfig": """ Qwen2.5 model konfigürasyonunu doğrudan okuyarak uyumlu TTSConfig oluşturur. """ from transformers import AutoConfig try: qwen_cfg = AutoConfig.from_pretrained(model_id, local_files_only=True) except Exception: qwen_cfg = AutoConfig.from_pretrained(model_id) params = { "text_vocab_size": qwen_cfg.vocab_size, "d_model": qwen_cfg.hidden_size, "num_heads": qwen_cfg.num_attention_heads, "num_kv_heads": qwen_cfg.num_key_value_heads, "num_layers": qwen_cfg.num_hidden_layers, "d_ff": qwen_cfg.intermediate_size, "max_seq_len": getattr(qwen_cfg, "max_position_embeddings", 2048), "rope_theta": getattr(qwen_cfg, "rope_theta", 10000.0), } params.update(kwargs) return cls(**params) class MultiCodebookEmbedding(nn.Module): """ 8 codebook'luk ses tensörünü her codebook için ayrı embedding tablosundan geçirip toplayarak tek bir d_model vektörüne indirger (VALL-E / AudioCraft standardı): e_frame(t) = sum_{k=0}^{K-1} E_k(codes[k, t]) """ def __init__(self, num_codebooks: int, audio_vocab_size: int, d_model: int, dtype: torch.dtype = torch.float32): super().__init__() self.num_codebooks = num_codebooks self.embeddings = nn.ModuleList([ nn.Embedding(audio_vocab_size, d_model, dtype=dtype) for _ in range(num_codebooks) ]) def forward(self, audio_codes: torch.Tensor) -> torch.Tensor: # audio_codes shape: (B, num_codebooks, T) B, K, T = audio_codes.shape out = torch.zeros(B, T, self.embeddings[0].embedding_dim, device=audio_codes.device, dtype=self.embeddings[0].weight.dtype) for k in range(min(K, self.num_codebooks)): out = out + self.embeddings[k](audio_codes[:, k, :]) return out class MainAudioTransformer(nn.Module): """ Stage 1: Metin ve geçmiş ses tokenlarını alarak sıradaki ses karesinin Codebook 0 (Semantik) tokenını tahmin eden Autoregressive Decoder Transformer. """ def __init__(self, config: TTSConfig): super().__init__() self.config = config self.d_model = config.d_model self.layers = nn.ModuleList([ TransformerBlock( d_model=config.d_model, num_heads=config.num_heads, num_kv_heads=config.num_kv_heads, d_ff=config.d_ff, dropout_rate=config.dropout_rate, dtype=config.torch_dtype, ) for _ in range(config.num_layers) ]) self.final_norm = RMSNorm(config.d_model, dtype=config.torch_dtype) self.lm_head_cb0 = nn.Linear(config.d_model, config.audio_vocab_size, bias=False, dtype=config.torch_dtype) def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, position_ids: Optional[torch.Tensor] = None, kv_caches: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None, use_cache: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[List[Tuple[torch.Tensor, torch.Tensor]]]]: new_kv_caches = [] if use_cache else None for i, layer in enumerate(self.layers): layer_cache = kv_caches[i] if kv_caches is not None else None x, new_cache = layer( x, mask=mask, rope=rope, position_ids=position_ids, kv_cache=layer_cache, use_cache=use_cache, ) if use_cache: new_kv_caches.append(new_cache) hidden_states = self.final_norm(x) logits_cb0 = self.lm_head_cb0(hidden_states) return logits_cb0, hidden_states, new_kv_caches class DepthDecoder(nn.Module): """ Stage 2: Main Transformer'dan çıkan ses karesi gizli durumunu (hidden state) ve CB0 tokenını alıp geriye kalan Codebook 1..7 (Akustik detaylar) tokenlarını tahmin eder. """ def __init__(self, config: TTSConfig): super().__init__() self.config = config self.num_codebooks = config.num_codebooks self.d_model = config.d_model # Codebook 0..6 embeddingleri (bir önceki codebook'ları girdi olarak beslemek için) self.cb_embeddings = nn.ModuleList([ nn.Embedding(config.audio_vocab_size, config.d_model, dtype=config.torch_dtype) for _ in range(config.num_codebooks - 1) ]) # Derinlik Transformer Blokları self.layers = nn.ModuleList([ TransformerBlock( d_model=config.d_model, num_heads=config.num_heads, num_kv_heads=config.num_kv_heads, d_ff=config.d_ff, dropout_rate=config.dropout_rate, dtype=config.torch_dtype, ) for _ in range(config.num_depth_layers) ]) self.final_norm = RMSNorm(config.d_model, dtype=config.torch_dtype) # Codebook 1..7 için ayrı tahmin kafaları (Heads) self.heads = nn.ModuleList([ nn.Linear(config.d_model, config.audio_vocab_size, bias=False, dtype=config.torch_dtype) for _ in range(config.num_codebooks - 1) ]) def forward( self, audio_hidden_states: torch.Tensor, audio_codes: Optional[torch.Tensor] = None, rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: """ Giriş: audio_hidden_states: (B, T_audio, d_model) audio_codes: (B, num_codebooks, T_audio) Çıkış: depth_logits: (B, num_codebooks - 1, T_audio, audio_vocab_size) """ B, T_audio, D = audio_hidden_states.shape device = audio_hidden_states.device dtype = audio_hidden_states.dtype # Her frame için codebook 1..7 tahminlerini yap logits_list = [] accumulated_emb = audio_hidden_states # Eğer tam audio_codes verilmişse (Eğitim - Teacher Forcing) if audio_codes is not None and audio_codes.shape[1] >= self.num_codebooks: for k in range(self.num_codebooks - 1): prev_code = audio_codes[:, k, :] accumulated_emb = accumulated_emb + self.cb_embeddings[k](prev_code) x = accumulated_emb for layer in self.layers: x, _ = layer(x, rope=rope) x = self.final_norm(x) head_logits = self.heads[k](x) # (B, T_audio, vocab_size) logits_list.append(head_logits) else: # Çıkarım (Inference): CB0'dan başlayarak CB1..CB7'yi zincirleme (autoregressive) tahmin et curr_code = audio_codes[:, 0, :] if (audio_codes is not None and audio_codes.shape[1] > 0) else None for k in range(self.num_codebooks - 1): if curr_code is not None: accumulated_emb = accumulated_emb + self.cb_embeddings[k](curr_code) x = accumulated_emb for layer in self.layers: x, _ = layer(x, rope=rope) x = self.final_norm(x) head_logits = self.heads[k](x) # (B, T_audio, vocab_size) logits_list.append(head_logits) curr_code = torch.argmax(head_logits, dim=-1) # Bir sonraki codebook için girdi # (B, K-1, T_audio, vocab_size) depth_logits = torch.stack(logits_list, dim=1) return depth_logits class TTSModel(nn.Module): """ Modern Modüler TTS Modeli: - Standart TTS (Text -> Speech) - Voice Design (Instruction Prompting) - Voice Clone (In-Context Reference Audio) - Hibrit Mod """ def __init__(self, config: TTSConfig): super().__init__() self.config = config # 1. Embedding Katmanları self.text_embedding = nn.Embedding( config.text_vocab_size, config.d_model, padding_idx=config.pad_token_id, dtype=config.torch_dtype ) # Sadece CB0 için embedding (Main AR üretimde kullanılır) self.cb0_embedding = nn.Embedding( config.audio_vocab_size, config.d_model, dtype=config.torch_dtype ) # Referans ses ve tam ses kareleri için 8 codebook'lu embedding self.multi_cb_embedding = MultiCodebookEmbedding( config.num_codebooks, config.audio_vocab_size, config.d_model, dtype=config.torch_dtype ) # 2. Stage 1: Main Autoregressive Backbone self.main_backbone = MainAudioTransformer(config) # 3. Stage 2: Depth Decoder self.depth_decoder = DepthDecoder(config) # 4. RoPE frekans tamponları cos, sin = precompute_rope_freqs( head_dim=config.d_model // config.num_heads, seq_len=config.max_seq_len, theta=config.rope_theta, device="cpu", ) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward( self, text_ids: torch.Tensor, target_audio_codes: torch.Tensor, ref_audio_codes: Optional[torch.Tensor] = None, depth_loss_weight: float = 1.0, ) -> Dict[str, torch.Tensor]: """ Eğitim İleri Beslemesi: text_ids: (B, T_text) target_audio_codes: (B, num_codebooks, T_audio) ref_audio_codes: (B, num_codebooks, T_ref) [Opsiyonel - Voice Clone] """ B, T_text = text_ids.shape _, K, T_audio = target_audio_codes.shape device = text_ids.device # 1. Girdileri Embedding Uzayına Taşı text_emb = self.text_embedding(text_ids) # (B, T_text, d_model) prefix_emb = text_emb if ref_audio_codes is not None: # Voice Clone: Referans sesi 8 codebook toplamı olarak göm ref_emb = self.multi_cb_embedding(ref_audio_codes) # (B, T_ref, d_model) prefix_emb = torch.cat([ref_emb, text_emb], dim=1) T_prefix = prefix_emb.shape[1] # Hedef sesin CB0 tokenları (girdi olarak t-1 anı verilir) # İlk ses karesi için ses başlangıç belirteci (BOS yerine 0 tokenı veya son prefix adımı kullanılır) target_cb0_in = target_audio_codes[:, 0, :-1] # (B, T_audio - 1) audio_emb_in = self.cb0_embedding(target_cb0_in) # (B, T_audio - 1, d_model) # 2. Tam diziyi uç uca ekle: [Prefix, Audio_in] full_seq = torch.cat([prefix_emb, audio_emb_in], dim=1) # (B, T_total, d_model) T_total = full_seq.shape[1] # 3. Causal Maske Hazırla # Ses tokenları nedensel olmalı, prefix tokenları birbirini görebilir mask = torch.ones(T_total, T_total, device=device, dtype=torch.bool).tril() # Prefix içi çift yönlü görüşe izin ver mask[:T_prefix, :T_prefix] = True mask = mask.unsqueeze(0).unsqueeze(1) # (1, 1, T_total, T_total) # 4. RoPE rope = (self.rope_cos[:T_total].to(device), self.rope_sin[:T_total].to(device)) # 5. Stage 1: Main Backbone Forward logits_cb0, hidden_states, _ = self.main_backbone(full_seq, mask=mask, rope=rope) # Ses kısmının logitlerini ve gizli durumlarını ayıkla: # Prefix'in son token'ı (<|audio|>, index T_prefix - 1) -> Hedef Frame 0'ı tahmin eder # target_cb0_in'in her karesi t -> Hedef Frame t+1'i tahmin eder # Böylece toplam T_audio adet tahmin ve gizli durum elde edilir! audio_logits_cb0 = logits_cb0[:, T_prefix - 1 :, :] # (B, T_audio, vocab_size) audio_hiddens = hidden_states[:, T_prefix - 1 :, :] # (B, T_audio, d_model) target_cb0_labels = target_audio_codes[:, 0, :] # (B, T_audio) # 6. Loss Stage 1: Codebook 0 Cross Entropy loss_cb0 = F.cross_entropy( audio_logits_cb0.reshape(-1, self.config.audio_vocab_size), target_cb0_labels.reshape(-1), reduction="mean", ) # 7. Stage 2: Depth Decoder Forward (Tüm T_audio kareleri için) depth_rope = (self.rope_cos[:audio_hiddens.shape[1]].to(device), self.rope_sin[:audio_hiddens.shape[1]].to(device)) depth_logits = self.depth_decoder( audio_hiddens, audio_codes=target_audio_codes, rope=depth_rope, ) # (B, K-1, T_audio, vocab_size) # 8. Loss Stage 2: Codebooks 1..7 Cross Entropy target_depth_labels = target_audio_codes[:, 1:, :] # (B, K-1, T_audio) loss_depth = F.cross_entropy( depth_logits.reshape(-1, self.config.audio_vocab_size), target_depth_labels.reshape(-1), reduction="mean", ) total_loss = loss_cb0 + depth_loss_weight * loss_depth return { "total_loss": total_loss, "loss_cb0": loss_cb0, "loss_depth": loss_depth, "logits_cb0": audio_logits_cb0, "logits_depth": depth_logits, } @torch.inference_mode() def generate( self, text_ids: torch.Tensor, ref_audio_codes: Optional[torch.Tensor] = None, max_new_tokens: int = 150, temperature: float = 0.6, top_k: int = 30, top_p: float = 0.95, depth_temperature: float = 0.6, ) -> torch.Tensor: """ Autoregressive Ses Üretimi: text_ids: (1, T_text) ref_audio_codes: (1, num_codebooks, T_ref) [Opsiyonel] Çıktı: generated_codes: (1, num_codebooks, T_audio) """ self.eval() device = text_ids.device # 1. Prefix Embedding text_emb = self.text_embedding(text_ids) if ref_audio_codes is not None: ref_emb = self.multi_cb_embedding(ref_audio_codes) prefix_emb = torch.cat([ref_emb, text_emb], dim=1) else: prefix_emb = text_emb # 2. Autoregressive Üretim (Codebook 0) curr_emb = prefix_emb kv_caches = None cb0_generated: List[int] = [] audio_hiddens_list: List[torch.Tensor] = [] total_len = prefix_emb.shape[1] for step in range(max_new_tokens): if (step + 1) % 20 == 0 or (step + 1) == max_new_tokens: print(f" - Ses karesi üretiliyor: {step + 1}/{max_new_tokens} ({(step + 1) * 100 // max_new_tokens}%)", flush=True) seq_len = curr_emb.shape[1] pos_ids = torch.arange(total_len - seq_len, total_len, device=device).unsqueeze(0) rope = (self.rope_cos[:total_len].to(device), self.rope_sin[:total_len].to(device)) logits_cb0, hidden, kv_caches = self.main_backbone( curr_emb, mask=None, rope=rope, position_ids=pos_ids, kv_caches=kv_caches, use_cache=True, ) # Son adımın logit'ini al last_logits = logits_cb0[:, -1, :] / max(temperature, 1e-5) last_hidden = hidden[:, -1:, :] # (1, 1, d_model) audio_hiddens_list.append(last_hidden) # Top-K / Top-P Sampling if top_k > 0: indices_to_remove = last_logits < torch.topk(last_logits, top_k)[0][..., -1, None] last_logits[indices_to_remove] = -float("Inf") probs = F.softmax(last_logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1).item() cb0_generated.append(next_token) # Bir sonraki adım için embedding curr_emb = self.cb0_embedding(torch.tensor([[next_token]], device=device)) total_len += 1 # 3. Stage 2: Depth Decoder ile Codebook 1..7'yi Tamamla audio_hiddens_tensor = torch.cat(audio_hiddens_list, dim=1) # (1, T_gen, d_model) cb0_tensor = torch.tensor([cb0_generated], device=device).unsqueeze(1) # (1, 1, T_gen) T_gen = audio_hiddens_tensor.shape[1] depth_rope = (self.rope_cos[:T_gen].to(device), self.rope_sin[:T_gen].to(device)) # Depth inference: CB0 tensörünü besleyerek CB1..CB7'yi üret depth_logits = self.depth_decoder(audio_hiddens_tensor, audio_codes=cb0_tensor, rope=depth_rope) # depth_logits shape: (1, K-1, T_gen, vocab_size) if depth_temperature > 0.0: # ⚡ Yumuşak Akustik Sampling: Robotik metalik sesi kırıp doğal tını ve rezonansı kazandırır scaled_depth_logits = depth_logits / max(depth_temperature, 1e-5) if top_k > 0: k_val = min(top_k, scaled_depth_logits.shape[-1]) indices_to_remove = scaled_depth_logits < torch.topk(scaled_depth_logits, k_val)[0][..., -1, None] scaled_depth_logits[indices_to_remove] = -float("Inf") depth_probs = F.softmax(scaled_depth_logits, dim=-1) B, K_minus_1, T_g, V = depth_probs.shape flat_probs = depth_probs.view(-1, V) flat_tokens = torch.multinomial(flat_probs, num_samples=1) depth_tokens = flat_tokens.view(B, K_minus_1, T_g) else: depth_tokens = torch.argmax(depth_logits, dim=-1) # Tüm 8 codebook'u birleştir: [CB0, CB1..7] full_codes = torch.cat([cb0_tensor, depth_tokens], dim=1) # (1, 8, T_gen) return full_codes def load_qwen_backbone(self, qwen_model_id: str = "Qwen/Qwen2.5-0.5B"): """ Qwen2.5 pretrained ağırlıklarını text_embedding ve main_backbone katmanlarına aktarır. """ from transformers import AutoModelForCausalLM print(f"[TTSModel] Qwen pretrained ağırlıkları yükleniyor ({qwen_model_id})...") qwen = AutoModelForCausalLM.from_pretrained(qwen_model_id, torch_dtype=self.config.torch_dtype) # 1. Text Embedding aktarımı qwen_embed = qwen.model.embed_tokens.weight.data min_vocab = min(self.text_embedding.weight.shape[0], qwen_embed.shape[0]) self.text_embedding.weight.data[:min_vocab].copy_(qwen_embed[:min_vocab].to(self.text_embedding.weight.device)) print(f" - Text Embedding yüklendi ({min_vocab} token)") # 2. Transformer katmanları aktarımı num_layers_to_load = min(len(self.main_backbone.layers), len(qwen.model.layers)) for i in range(num_layers_to_load): q_layer = qwen.model.layers[i] m_layer = self.main_backbone.layers[i] # Attention weights & biases m_layer.self_attn.q_proj.weight.data.copy_(q_layer.self_attn.q_proj.weight.data) if m_layer.self_attn.q_proj.bias is not None and q_layer.self_attn.q_proj.bias is not None: m_layer.self_attn.q_proj.bias.data.copy_(q_layer.self_attn.q_proj.bias.data) m_layer.self_attn.k_proj.weight.data.copy_(q_layer.self_attn.k_proj.weight.data) if m_layer.self_attn.k_proj.bias is not None and q_layer.self_attn.k_proj.bias is not None: m_layer.self_attn.k_proj.bias.data.copy_(q_layer.self_attn.k_proj.bias.data) m_layer.self_attn.v_proj.weight.data.copy_(q_layer.self_attn.v_proj.weight.data) if m_layer.self_attn.v_proj.bias is not None and q_layer.self_attn.v_proj.bias is not None: m_layer.self_attn.v_proj.bias.data.copy_(q_layer.self_attn.v_proj.bias.data) m_layer.self_attn.out_proj.weight.data.copy_(q_layer.self_attn.o_proj.weight.data) # FFN weights m_layer.ffn.gate_proj.weight.data.copy_(q_layer.mlp.gate_proj.weight.data) m_layer.ffn.up_proj.weight.data.copy_(q_layer.mlp.up_proj.weight.data) m_layer.ffn.down_proj.weight.data.copy_(q_layer.mlp.down_proj.weight.data) # Normalizasyonlar (Standart RMSNorm - weight birebir kopyalanır) m_layer.norm1.weight.data.copy_(q_layer.input_layernorm.weight.data) m_layer.norm2.weight.data.copy_(q_layer.post_attention_layernorm.weight.data) print(f" - {num_layers_to_load} adet Transformer katmanı başarıyla yüklendi!") # 3. Final Norm self.main_backbone.final_norm.weight.data.copy_(qwen.model.norm.weight.data) print(" - Final RMSNorm yüklendi!") print("[TTSModel] Qwen2.5 omurga ağırlıkları başarıyla entegre edildi!")