78 lines
2.1 KiB
Python
78 lines
2.1 KiB
Python
# Copyright (c) 2024 Amphion.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from einops import rearrange
|
|
from torch.nn.utils import weight_norm
|
|
|
|
|
|
def WNConv1d(*args, **kwargs):
|
|
return weight_norm(nn.Conv1d(*args, **kwargs))
|
|
|
|
|
|
def WNConvTranspose1d(*args, **kwargs):
|
|
return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
|
|
|
|
|
|
class LookupFreeQuantize(nn.Module):
|
|
def __init__(
|
|
self,
|
|
input_dim,
|
|
codebook_size,
|
|
codebook_dim,
|
|
):
|
|
super().__init__()
|
|
self.input_dim = input_dim
|
|
self.codebook_size = codebook_size
|
|
self.codebook_dim = codebook_dim
|
|
|
|
assert 2**codebook_dim == codebook_size
|
|
|
|
if self.input_dim != self.codebook_dim:
|
|
self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
|
|
self.out_project = WNConv1d(
|
|
self.codebook_dim, self.input_dim, kernel_size=1
|
|
)
|
|
|
|
else:
|
|
self.in_project = nn.Identity()
|
|
self.out_project = nn.Identity()
|
|
|
|
def forward(self, z):
|
|
z_e = self.in_project(z)
|
|
z_e = F.sigmoid(z_e)
|
|
|
|
z_q = z_e + (torch.round(z_e) - z_e).detach()
|
|
|
|
z_q = self.out_project(z_q)
|
|
|
|
commit_loss = torch.zeros(z.shape[0], device=z.device)
|
|
codebook_loss = torch.zeros(z.shape[0], device=z.device)
|
|
|
|
bits = (
|
|
2
|
|
** torch.arange(self.codebook_dim, device=z.device)
|
|
.unsqueeze(0)
|
|
.unsqueeze(-1)
|
|
.long()
|
|
) # (1, d, 1)
|
|
indices = (torch.round(z_e.clone().detach()).long() * bits).sum(1).long()
|
|
|
|
return z_q, commit_loss, codebook_loss, indices, z_e
|
|
|
|
def vq2emb(self, vq, out_proj=True):
|
|
emb = torch.zeros(
|
|
vq.shape[0], self.codebook_dim, vq.shape[-1], device=vq.device
|
|
) # (B, d, T)
|
|
for i in range(self.codebook_dim):
|
|
emb[:, i, :] = (vq % 2).float()
|
|
vq = vq // 2
|
|
if out_proj:
|
|
emb = self.out_project(emb)
|
|
return emb
|