cnn_bigru/utils/vqvae2.py

"""vqvae2.py — VQ-VAE-2 Hierárquico para compressão multimodal.

Implementa Vector Quantized Variational Autoencoder 2 (Roy et al., 2019),
uma versão hierárquica do VQ-VAE que usa dois níveis de quantização:
    - Top level: captura padrões semânticos globais (low-resolution)
    - Bottom level: captura detalhes locais (high-resolution)

==============================================================================
ANÁLISE MATEMÁTICA E LÓGICA
==============================================================================

ARQUITETURA:
    Input x ∈ R^{B×C×H×W}
        |
        v
    [Bottom Encoder] -> z_b ∈ R^{B×d_b×H×W}
        |
        v
    [Top Encoder] -> z_t ∈ R^{B×d_t×H/2×W/2}
        |
        v
    [Top Quantizer] -> z_t_q (codebook lookup)
        |
        v
    [Top Decoder] -> z_t_dec ∈ R^{B×d_b×H×W}
        |
        v
    [Bottom Quantizer with z_t_dec] -> z_b_q
        |
        v
    [Bottom Decoder] -> x_recon

LOSS:
    L = L_recon + beta * L_commit_top + beta * L_commit_bottom
      + gamma * L_codebook_top + gamma * L_codebook_bottom

    onde:
        L_recon = MSE(x, x_recon)               ou L1
        L_commit_k = ||sg(z_k) - e_k||^2        (commit loss)
        L_codebook_k = ||z_k - sg(e_k)||^2      (codebook loss)
        sg = stop-gradient

    Update via EMA (Exponential Moving Average) — alternativo ao gradiente:
        n_k^{t+1} = gamma_ema * n_k^t + (1 - gamma_ema) * sum_k 1[assigned]
        m_k^{t+1} = gamma_ema * m_k^t + (1 - gamma_ema) * sum_k z_k
        e_k = m_k / n_k

CODEBOOK:
    Top codebook: K_t vetores de dimensão d_t (tipicamente 512-1024 vetores)
    Bottom codebook: K_b vetores de dimensão d_b (tipicamente 512-1024 vetores)

==============================================================================
INTEGRAÇÃO COM CNN-BiGRU
==============================================================================

O VQ-VAE-2 é usado para comprimir representações multimodais (imagem, áudio)
em códigos discretos que podem ser processados pelo backbone CNN-BiGRU:

    1. Imagem -> VQ-VAE-2 -> códigos top + bottom (tokens discretos)
    2. Códigos -> embedding lookup -> [B, T, D]
    3. [B, T, D] -> CooperativeCNNBiGRU (como se fosse texto tokenizado)

Isto permite que o modelo trate imagem/áudio da mesma forma que texto,
unificando o pipeline multimodal.

Autor: CNN-BiGRU Project
"""
from __future__ import annotations

import logging
import math
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

logger = logging.getLogger(__name__)


# ============================================================================
# Configuração
# ============================================================================

@dataclass
class VQVAE2Config:
    """Configuração do VQ-VAE-2 hierárquico."""
    # Dimensões de entrada (assumindo features [B, C, H, W])
    in_channels: int = 1
    # Dimensões latentes
    bottom_channels: int = 64    # d_b
    top_channels: int = 32       # d_t
    # Tamanhos dos codebooks
    n_bottom_codes: int = 512    # K_b
    n_top_codes: int = 512       # K_t
    # Beta (commit loss weight)
    commitment_cost: float = 0.25
    # Gamma (codebook loss weight — apenas para grad update, não EMA)
    codebook_cost: float = 0.0
    # Usar EMA para atualizar codebook (default: True)
    use_ema: bool = True
    ema_decay: float = 0.99
    # Usar residual connection (z + delta)
    use_residual: bool = False
    # Número de camadas de down/upsampling
    n_downsample: int = 2
    # Base channels do encoder/decoder
    hidden_channels: int = 64
    # Device
    device: str = "cpu"


# ============================================================================
# Residual Blocks (para encoder/decoder)
# ============================================================================

def _adaptive_groups(channels: int, max_groups: int = 8) -> int:
    """Computa número de grupos para GroupNorm que divide channels."""
    if channels <= 0:
        return 1
    g = min(max_groups, channels)
    while channels % g != 0 and g > 1:
        g -= 1
    return max(1, g)


class ResidualBlock(nn.Module):
    """Bloco residual simples com Conv2d + GroupNorm + SiLU."""

    def __init__(self, channels: int, hidden_channels: Optional[int] = None):
        super().__init__()
        h = hidden_channels or channels
        self.conv1 = nn.Conv2d(channels, h, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(h, channels, kernel_size=3, padding=1)
        # GroupNorm adaptativo por número de canais
        self.norm1 = nn.GroupNorm(_adaptive_groups(channels), channels)
        # norm2 aplica-se à saída de conv1 (h canais)
        self.norm2 = nn.GroupNorm(_adaptive_groups(h), h)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = F.silu(self.norm1(x))
        h = self.conv1(h)
        h = F.silu(self.norm2(h))
        h = self.conv2(h)
        return x + h


# ============================================================================
# Encoder / Decoder (Conv2d-based)
# ============================================================================

class ConvEncoder(nn.Module):
    """Encoder convolucional: reduz resolução espacial por fator 2^n_downsample."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int = 64,
        n_downsample: int = 2,
    ):
        super().__init__()
        layers = []
        c_in = in_channels
        for i in range(n_downsample):
            c_out = hidden_channels * (2 ** i)
            layers.extend([
                nn.Conv2d(c_in, c_out, kernel_size=4, stride=2, padding=1),
                nn.GroupNorm(_adaptive_groups(c_out), c_out),
                nn.SiLU(),
            ])
            c_in = c_out
        # Residual blocks para refinar
        layers.append(ResidualBlock(c_in, out_channels))
        # Projeta para out_channels
        layers.append(nn.Conv2d(c_in, out_channels, kernel_size=1))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class ConvDecoder(nn.Module):
    """Decoder convolucional: aumenta resolução espacial por fator 2^n_upsample."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int = 64,
        n_upsample: int = 2,
    ):
        super().__init__()
        layers = []
        # Primeiro projeta de in_channels para hidden_channels * 2^(n-1)
        c_start = hidden_channels * (2 ** (n_upsample - 1))
        layers.append(nn.Conv2d(in_channels, c_start, kernel_size=1))
        layers.append(ResidualBlock(c_start, c_start))

        c_in = c_start
        for i in range(n_upsample - 1, -1, -1):
            c_out = hidden_channels * (2 ** max(0, i - 1)) if i > 0 else out_channels
            layers.extend([
                nn.ConvTranspose2d(c_in, c_out, kernel_size=4, stride=2, padding=1),
                nn.GroupNorm(_adaptive_groups(c_out), c_out) if i > 0 else nn.Identity(),
                nn.SiLU() if i > 0 else nn.Identity(),
            ])
            c_in = c_out
        # Camada final 1x1 para ajustar canais
        layers.append(nn.Conv2d(c_in, out_channels, kernel_size=1))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


# ============================================================================
# Vector Quantizer (com EMA opcional)
# ============================================================================

class VectorQuantizer(nn.Module):
    """Quantizador por codebook com Straight-Through Estimator (STE).

    Args:
        n_codes: número de códigos no codebook (K)
        dim: dimensão de cada código (d)
        commitment_cost: peso da commitment loss
        use_ema: se True, atualiza codebook via EMA em vez de gradiente
        ema_decay: decay do EMA
    """

    def __init__(
        self,
        n_codes: int,
        dim: int,
        commitment_cost: float = 0.25,
        use_ema: bool = True,
        ema_decay: float = 0.99,
    ):
        super().__init__()
        self.n_codes = n_codes
        self.dim = dim
        self.commitment_cost = commitment_cost
        self.use_ema = use_ema
        self.ema_decay = ema_decay

        # Codebook: [K, d]
        codebook = torch.randn(n_codes, dim) * 0.01
        if use_ema:
            self.register_buffer("codebook", codebook)
            self.register_buffer("ema_count", torch.zeros(n_codes))
            self.register_buffer("ema_weight", codebook.clone())
        else:
            self.codebook = nn.Parameter(codebook)

    def forward(
        self,
        z: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
        """Quantiza z para o codebook mais próximo.

        Args:
            z: [B, d, H, W] ou [B, d, T]

        Returns:
            z_q: quantizado [B, d, H, W]
            indices: [B, H, W] índices do codebook
            loss_dict: dict com commitment loss, codebook loss, etc.
        """
        # Reshape z para [N, d] onde N = B * H * W
        original_shape = z.shape
        if z.dim() == 4:
            B, d, H, W = z.shape
            z_flat = z.permute(0, 2, 3, 1).reshape(-1, d)  # [B*H*W, d]
        elif z.dim() == 3:
            B, d, T = z.shape
            z_flat = z.permute(0, 2, 1).reshape(-1, d)  # [B*T, d]
        else:
            z_flat = z

        # Distância euclidiana ao codebook
        # ||z - e||^2 = ||z||^2 - 2*z·e + ||e||^2
        dist = (
            z_flat.pow(2).sum(dim=-1, keepdim=True)
            - 2 * z_flat @ self.codebook.t()
            + self.codebook.pow(2).sum(dim=-1, keepdim=False).unsqueeze(0)
        )  # [N, K]

        # Encontrar código mais próximo
        indices = dist.argmin(dim=-1)  # [N]
        one_hot = F.one_hot(indices, self.n_codes).float()  # [N, K]

        # Quantizado
        z_q = one_hot @ self.codebook  # [N, d]

        # Commitment loss
        commit_loss = F.mse_loss(z_flat, z_q.detach())
        # Codebook loss (apenas se não usar EMA)
        if not self.use_ema:
            codebook_loss = F.mse_loss(z_q, z_flat.detach())
        else:
            codebook_loss = torch.zeros(1, device=z.device)

        # EMA update
        if self.use_ema and self.training:
            with torch.no_grad():
                # Conta quantas vezes cada código foi usado
                code_count = one_hot.sum(dim=0)  # [K]
                # Soma dos z_flat para cada código
                code_sum = one_hot.t() @ z_flat  # [K, d]

                # EMA
                self.ema_count.mul_(self.ema_decay).add_(
                    code_count, alpha=1 - self.ema_decay
                )
                self.ema_weight.mul_(self.ema_decay).add_(
                    code_sum, alpha=1 - self.ema_decay
                )
                # Atualiza codebook
                self.codebook.copy_(self.ema_weight / self.ema_count.unsqueeze(-1).clamp(min=1e-8))

        # Straight-Through Estimator: z_q = z + (z_q - z).detach()
        z_q_st = z_flat + (z_q - z_flat).detach()

        # Reshape de volta
        if z.dim() == 4:
            z_q_st = z_q_st.view(B, H, W, d).permute(0, 3, 1, 2).contiguous()
            indices = indices.view(B, H, W)
        elif z.dim() == 3:
            z_q_st = z_q_st.view(B, T, d).permute(0, 2, 1).contiguous()
            indices = indices.view(B, T)
        else:
            z_q_st = z_q_st.view(original_shape)

        # Taxa de uso do codebook (para diagnóstico)
        with torch.no_grad():
            used = (one_hot.sum(dim=0) > 0).float().sum()
            usage = used / self.n_codes

        loss_dict = {
            "commit_loss": commit_loss * self.commitment_cost,
            "codebook_loss": codebook_loss,
            "usage": usage.detach(),
        }

        return z_q_st, indices, loss_dict


# ============================================================================
# VQ-VAE-2 Hierárquico
# ============================================================================

class VQVAE2(nn.Module):
    """VQ-VAE-2 Hierárquico.

    Estrutura:
        x -> BottomEnc -> z_b -> TopEnc -> z_t -> TopQ -> z_t_q
                                                          |
                                                          v
                                                       TopDec -> z_t_dec
                                                          |
                                                          v
        z_b + z_t_dec -> BottomQ -> z_b_q -> BottomDec -> x_recon
    """

    def __init__(self, config: VQVAE2Config):
        super().__init__()
        self.config = config

        # Bottom encoder/decoder
        self.bottom_encoder = ConvEncoder(
            in_channels=config.in_channels,
            out_channels=config.bottom_channels,
            hidden_channels=config.hidden_channels,
            n_downsample=config.n_downsample,
        )
        self.bottom_decoder = ConvDecoder(
            in_channels=config.bottom_channels,
            out_channels=config.in_channels,
            hidden_channels=config.hidden_channels,
            n_upsample=config.n_downsample,
        )

        # Top encoder/decoder (operam sobre z_b)
        self.top_encoder = ConvEncoder(
            in_channels=config.bottom_channels,
            out_channels=config.top_channels,
            hidden_channels=config.hidden_channels,
            n_downsample=1,  # downsample adicional por 2
        )
        self.top_decoder = ConvDecoder(
            in_channels=config.top_channels,
            out_channels=config.bottom_channels,
            hidden_channels=config.hidden_channels,
            n_upsample=1,
        )

        # Quantizadores
        self.top_quantizer = VectorQuantizer(
            n_codes=config.n_top_codes,
            dim=config.top_channels,
            commitment_cost=config.commitment_cost,
            use_ema=config.use_ema,
            ema_decay=config.ema_decay,
        )
        self.bottom_quantizer = VectorQuantizer(
            n_codes=config.n_bottom_codes,
            dim=config.bottom_channels,
            commitment_cost=config.commitment_cost,
            use_ema=config.use_ema,
            ema_decay=config.ema_decay,
        )

    def forward(
        self,
        x: torch.Tensor,
    ) -> Dict[str, torch.Tensor]:
        """Forward pass completo.

        Args:
            x: [B, C, H, W] entrada (imagem, espectrograma, etc.)

        Returns:
            dict com:
                x_recon: [B, C, H, W] reconstrução
                top_indices: [B, H_t, W_t] índices top
                bottom_indices: [B, H_b, W_b] índices bottom
                loss: loss total (recon + commit + codebook)
                loss_dict: breakdown
        """
        # 1. Bottom encoder
        z_b = self.bottom_encoder(x)  # [B, d_b, H_b, W_b]

        # 2. Top encoder (sobre z_b)
        z_t = self.top_encoder(z_b)  # [B, d_t, H_t, W_t]

        # 3. Top quantizer
        z_t_q, top_indices, top_loss = self.top_quantizer(z_t)

        # 4. Top decoder (reconstroi z_b aproximado do top)
        z_t_dec = self.top_decoder(z_t_q)  # [B, d_b, H_b, W_b]

        # 5. Bottom quantizer sobre (z_b + z_t_dec) ou (z_b - z_t_dec)
        if self.config.use_residual:
            z_b_residual = z_b - z_t_dec
            z_b_q, bottom_indices, bottom_loss = self.bottom_quantizer(z_b_residual)
            z_b_full = z_b_q + z_t_dec  # reconstrói
        else:
            z_b_q, bottom_indices, bottom_loss = self.bottom_quantizer(z_b)
            z_b_full = z_b_q + z_t_dec  # combina top + bottom

        # 6. Bottom decoder
        x_recon = self.bottom_decoder(z_b_full)  # [B, C, H, W]

        # 7. Loss
        recon_loss = F.mse_loss(x_recon, x)
        total_loss = (
            recon_loss
            + top_loss["commit_loss"]
            + bottom_loss["commit_loss"]
            + top_loss["codebook_loss"]
            + bottom_loss["codebook_loss"]
        )

        return {
            "x_recon": x_recon,
            "top_indices": top_indices,
            "bottom_indices": bottom_indices,
            "loss": total_loss,
            "loss_dict": {
                "recon_loss": recon_loss.detach(),
                "top_commit": top_loss["commit_loss"].detach(),
                "bottom_commit": bottom_loss["commit_loss"].detach(),
                "top_usage": top_loss["usage"],
                "bottom_usage": bottom_loss["usage"],
            },
        }

    def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Apenas encoder (para uso em pipeline multimodal).

        Returns:
            (top_indices, bottom_indices) — códigos discretos
        """
        z_b = self.bottom_encoder(x)
        z_t = self.top_encoder(z_b)
        _, top_idx, _ = self.top_quantizer(z_t)
        z_t_dec = self.top_decoder(self.top_quantizer.codebook[top_idx].permute(
            0, 3, 1, 2
        ) if top_idx.dim() == 3 else self.top_quantizer.codebook[top_idx])
        # Simplificado: apenas retorna top_idx e bottom_idx brutos
        if self.config.use_residual:
            z_b_resid = z_b - z_t_dec
            _, bottom_idx, _ = self.bottom_quantizer(z_b_resid)
        else:
            _, bottom_idx, _ = self.bottom_quantizer(z_b)
        return top_idx, bottom_idx

    def decode_codes(
        self,
        top_indices: torch.Tensor,
        bottom_indices: torch.Tensor,
    ) -> torch.Tensor:
        """Decodifica a partir de códigos discretos."""
        # Top
        z_t_q = self.top_quantizer.codebook[top_indices]  # [B, H_t, W_t, d_t]
        # Adicionar dim de canal se necessário
        if z_t_q.dim() == 4:
            z_t_q = z_t_q.permute(0, 3, 1, 2).contiguous()
        z_t_dec = self.top_decoder(z_t_q)

        # Bottom
        z_b_q = self.bottom_quantizer.codebook[bottom_indices]
        if z_b_q.dim() == 4:
            z_b_q = z_b_q.permute(0, 3, 1, 2).contiguous()
        z_b_full = z_b_q + z_t_dec

        # Decoder
        x_recon = self.bottom_decoder(z_b_full)
        return x_recon


# ============================================================================
# Self-test
# ============================================================================

def _self_test():
    """Teste rápido do VQ-VAE-2."""
    torch.manual_seed(42)
    config = VQVAE2Config(
        in_channels=1,
        bottom_channels=16,
        top_channels=8,
        n_bottom_codes=64,
        n_top_codes=64,
        n_downsample=2,
        hidden_channels=16,
        use_ema=True,
        commitment_cost=0.25,
    )
    vqvae = VQVAE2(config)
    print(f"VQ-VAE-2 params: {sum(p.numel() for p in vqvae.parameters())}")

    # Forward
    x = torch.randn(2, 1, 16, 16)
    out = vqvae(x)
    print(f"Recon: {out['x_recon'].shape}")
    print(f"Top indices: {out['top_indices'].shape}")
    print(f"Bottom indices: {out['bottom_indices'].shape}")
    print(f"Loss: {out['loss'].item():.4f}")
    print(f"Loss dict: {out['loss_dict']}")

    # Test encode/decode
    top_idx, bot_idx = vqvae.encode(x)
    print(f"Encode: top {top_idx.shape}, bot {bot_idx.shape}")


if __name__ == "__main__":
    _self_test()


__all__ = [
    "VQVAE2Config",
    "ResidualBlock",
    "ConvEncoder",
    "ConvDecoder",
    "VectorQuantizer",
    "VQVAE2",
]