model.py

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!")