|
from dataclasses import dataclass |
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
from einops import rearrange |
|
from vector_quantize_pytorch import GroupedResidualFSQ |
|
|
|
from .firefly import ConvNeXtBlock, FishConvNet, FishTransConvNet |
|
|
|
|
|
@dataclass |
|
class FSQResult: |
|
z: torch.Tensor |
|
codes: torch.Tensor |
|
latents: torch.Tensor |
|
|
|
|
|
class DownsampleFiniteScalarQuantize(nn.Module): |
|
def __init__( |
|
self, |
|
input_dim: int = 512, |
|
n_codebooks: int = 9, |
|
n_groups: int = 1, |
|
levels: tuple[int] = (8, 5, 5, 5), |
|
downsample_factor: tuple[int] = (2, 2), |
|
downsample_dims: tuple[int] | None = None, |
|
): |
|
super().__init__() |
|
|
|
if downsample_dims is None: |
|
downsample_dims = [input_dim for _ in range(len(downsample_factor))] |
|
|
|
all_dims = (input_dim,) + tuple(downsample_dims) |
|
|
|
self.residual_fsq = GroupedResidualFSQ( |
|
dim=all_dims[-1], |
|
levels=levels, |
|
num_quantizers=n_codebooks, |
|
groups=n_groups, |
|
) |
|
|
|
self.downsample_factor = downsample_factor |
|
self.downsample_dims = downsample_dims |
|
|
|
self.downsample = nn.Sequential( |
|
*[ |
|
nn.Sequential( |
|
FishConvNet( |
|
all_dims[idx], |
|
all_dims[idx + 1], |
|
kernel_size=factor, |
|
stride=factor, |
|
), |
|
ConvNeXtBlock(dim=all_dims[idx + 1]), |
|
) |
|
for idx, factor in enumerate(downsample_factor) |
|
] |
|
) |
|
|
|
self.upsample = nn.Sequential( |
|
*[ |
|
nn.Sequential( |
|
FishTransConvNet( |
|
all_dims[idx + 1], |
|
all_dims[idx], |
|
kernel_size=factor, |
|
stride=factor, |
|
), |
|
ConvNeXtBlock(dim=all_dims[idx]), |
|
) |
|
for idx, factor in reversed(list(enumerate(downsample_factor))) |
|
] |
|
) |
|
|
|
self.apply(self._init_weights) |
|
|
|
def _init_weights(self, m): |
|
if isinstance(m, (nn.Conv1d, nn.Linear)): |
|
nn.init.trunc_normal_(m.weight, std=0.02) |
|
nn.init.constant_(m.bias, 0) |
|
|
|
def forward(self, z) -> FSQResult: |
|
original_shape = z.shape |
|
z = self.downsample(z) |
|
quantized, indices = self.residual_fsq(z.mT) |
|
result = FSQResult( |
|
z=quantized.mT, |
|
codes=indices.mT, |
|
latents=z, |
|
) |
|
result.z = self.upsample(result.z) |
|
|
|
|
|
diff = original_shape[-1] - result.z.shape[-1] |
|
left = diff // 2 |
|
right = diff - left |
|
|
|
if diff > 0: |
|
result.z = F.pad(result.z, (left, right)) |
|
elif diff < 0: |
|
result.z = result.z[..., -left:right] |
|
|
|
return result |
|
|
|
def encode(self, z): |
|
z = self.downsample(z) |
|
_, indices = self.residual_fsq(z.mT) |
|
indices = rearrange(indices, "g b l r -> b (g r) l") |
|
return indices |
|
|
|
def decode(self, indices: torch.Tensor): |
|
indices = rearrange(indices, "b (g r) l -> g b l r", g=self.residual_fsq.groups) |
|
z_q = self.residual_fsq.get_output_from_indices(indices) |
|
z_q = self.upsample(z_q.mT) |
|
return z_q |
|
|
