Class LatentQuantizer

model.py:18–57

@si_module
class LatentQuantizer(nn.Module):
    class Config:
        compressor_config: Optional[FSQ.Config] = None

        dim: Optional[int] = None
        ff_dim: Optional[int] = None
        input_dim: int = None

        from_pretrained: Optional[Tuple[str, str]] = None

    def __init__(self, c: Config):
        super().__init__()

        if exists(c.from_pretrained):
            checkpoint = load_ckpt(*c.from_pretrained)
        else:
            assert exists(c.compressor_config), f'hmm {c}'

        self.compressor = c.compressor_config()
        self.ffnn = FFNN(c.dim, c.ff_dim)
        self.input = nn.Linear(c.input_dim, c.dim) if exists(c.input_dim) else nn.Identity()

        if exists(c.from_pretrained):
            self.load_state_dict(checkpoint)

    @T.no_grad()
    def forward(self, x, return_latent=False, known_latent=None):  
        """
        x: (B, S, D)
        """
        if exists(known_latent): 
            return self.compressor.indices_to_codes(known_latent)

        x = self.input(x)
        x = self.ffnn(x)
        x, tokens = self.compressor(x)
        
        if return_latent:
            return x, tokens
        return x


@si_module