Delete utils/vqvae2.py with huggingface_hub

utils/vqvae2.py

@@ -1,570 +0,0 @@
-"""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",
-]