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

class VectorQuantizer(nn.Module):
    """
    Vector Quantization layer for VQ-VAE.
    """
    def __init__(self, num_embeddings, embedding_dim, beta=0.25):
        super(VectorQuantizer, self).__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.beta = beta

        self.embeddings = nn.Embedding(num_embeddings, embedding_dim)
        self.embeddings.weight.data.uniform_(-1 / num_embeddings, 1 / num_embeddings)

    def forward(self, z):
        # Flatten input
        z_flattened = z.view(-1, self.embedding_dim)

        # Compute distances to embeddings
        distances = torch.sum(z_flattened ** 2, dim=1, keepdim=True) - \
                    2 * torch.matmul(z_flattened, self.embeddings.weight.T) + \
                    torch.sum(self.embeddings.weight ** 2, dim=1)

        # Get closest embeddings
        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
        quantized = F.embedding(encoding_indices, self.embeddings.weight).view_as(z)

        # Compute loss
        commitment_loss = F.mse_loss(quantized.detach(), z)
        codebook_loss = F.mse_loss(quantized, z.detach())
        loss = codebook_loss + self.beta * commitment_loss

        quantized = z + (quantized - z).detach()  # Straight-through estimator

        return quantized, loss, encoding_indices.view(*z.shape[:-1])

class Encoder(nn.Module):
    """
    Encoder with convolutional layers for SVD input.
    """
    def __init__(self, input_channels=3, embedding_dim=64):
        super(Encoder, self).__init__()
        self.conv1 = nn.Conv2d(input_channels, 32, kernel_size=4, stride=2, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1)
        self.conv3 = nn.Conv2d(64, embedding_dim, kernel_size=4, stride=2, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.conv3(x)  # Output embedding
        return x

class Decoder(nn.Module):
    """
    Decoder with transposed convolutional layers for SVD reconstruction.
    """
    def __init__(self, embedding_dim=64, output_channels=3):
        super(Decoder, self).__init__()
        self.conv1 = nn.ConvTranspose2d(embedding_dim, 64, kernel_size=4, stride=2, padding=1)
        self.conv2 = nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1)
        self.conv3 = nn.ConvTranspose2d(32, output_channels, kernel_size=4, stride=2, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = torch.sigmoid(self.conv3(x))  # Normalize to [0, 1]
        return x

class VQVAE(nn.Module):
    """
    VQ-VAE model for SVD input.
    """
    def __init__(self, input_channels=3, embedding_dim=64, num_embeddings=512):
        super(VQVAE, self).__init__()
        self.encoder = Encoder(input_channels, embedding_dim)
        self.quantizer = VectorQuantizer(num_embeddings, embedding_dim)
        self.decoder = Decoder(embedding_dim, input_channels)

    def forward(self, x):
        z = self.encoder(x)
        z = z.permute(0, 2, 3, 1)  # Rearrange dimensions for quantization
        quantized, vq_loss, _ = self.quantizer(z)
        quantized = quantized.permute(0, 3, 1, 2)  # Restore dimensions for decoder
        recon = self.decoder(quantized)
        return recon, vq_loss

# Example usage:
if __name__ == "__main__":
    input_channels = 3  # For S, D, V
    embedding_dim = 64
    num_embeddings = 512
    batch_size = 16
    height, width = 128, 256  # Example dimensions for input

    model = VQVAE(input_channels, embedding_dim, num_embeddings)
    input_data = torch.rand(batch_size, input_channels, height, width)
    recon, vq_loss = model(input_data)

    print(f"Reconstructed shape: {recon.shape}")
    print(f"VQ Loss: {vq_loss.item()}")


Вот тебе пример архитектуры для SVD. Доработай ее до рабочего состояния, чтобы в коде уже было обучение и реализуй все необходимые и предложенные тобой полезные методики в коде. Реализуй предобработку midi в svd из папки data. Реализуй чекпоинты у модели и загрузку из последнего сохраненного чекпоинта и добавь параметр регулирующий кол-во чекпоинтов для сохранения. Добавь генерацию svd с дальнейшим преобразованием и сохранением в midi каждые n шагов для контроля за обучением. Реализуй переменный размер L на входе в модель, чтобы он мог принимать абсолютно любые значения