model.py

from typing import Optional, Tuple

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

from ioblocks import GaussianMixtureIOLayer, FSQ

from transformer import Stack, ShapeRotator, Block as PerfBlock, GPTOutput, CACHE_FILL_VALUE, FFNN, Norm
from tokenizer import make_tokenizer


from utils import si_module, exists, isnt, tqdm0, print0, default, print0_colored
from utils import load_ckpt    


@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
class TransformerVAE(nn.Module):
    class Config:
        io_config: Optional[GaussianMixtureIOLayer.Config] = None
        stack_config: Optional[Stack.Config] = None
        quantizer_config: Optional[LatentQuantizer.Config] = None

        plex_layer: int = None
        plex_roll: int = 1
        split: bool = True

        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.io_config) and exists(c.stack_config) and exists(c.quantizer_config)), f'hmm {c}'

        self.io = c.io_config()
        self.stack = c.stack_config()

        self.plex_layer = c.stack_config.layers//2
        self.plex_roll = c.plex_roll
        self.plex_dim = c.quantizer_config.dim

        assert self.plex_dim is not None and c.stack_config.dim is not None, f'One of the following are None: self.plex_dim: {self.plex_dim}, c.stack_config.dim: {c.stack_config.dim}'
        self.plex_projection = nn.Linear(self.plex_dim, c.stack_config.dim)
        self.out_norm = Norm(c.stack_config.dim)

        if c.split:
            self.io2 = c.io_config()
            self.plex_projection2 = nn.Linear(self.plex_dim, c.stack_config.dim)

            self.io2.fc_loc = None
            self.io2.fc_scale = None
            self.io2.fc_weight = None

        kv_heads = c.stack_config.kv_heads or c.stack_config.n_head
        head_dim = c.stack_config.dim // c.stack_config.n_head
        self.cache_num_layers = c.stack_config.layers + ((c.stack_config.layers - self.plex_layer) if c.split else 0)
        cache_shape = [self.cache_num_layers, c.stack_config.seq_len, 2, kv_heads, head_dim]
        self.cache_shape = cache_shape
        self.cache = [None] * self.cache_num_layers
        
        if exists(c.from_pretrained):
            result = self.load_state_dict(checkpoint, strict=False)
            print0_colored(result, 'yellow')

        self.quantizer = c.quantizer_config().eval()
        self.quantizer.requires_grad = False

    @T.no_grad()
    def quantize(self, x):
        if self.c.split:
            x1, x2 = x.chunk(2, dim=-1)
            with T.autocast(device_type='cuda', dtype=T.bfloat16):
                quantized1 = self.quantizer(x1)
                quantized2 = self.quantizer(x2)
            return quantized1, quantized2
        else:
            with T.autocast(device_type='cuda', dtype=T.bfloat16):
                return self.quantizer(x)
    
    @T.no_grad()
    def untokenize(self, token_data):
        return self.quantizer(None, known_latent=token_data)

    def init_cache(self, bsize, device, dtype, length:int=None):
        cache_shape = self.cache_shape.copy()
        cache_shape[1] = length or cache_shape[1]
        self.cache = T.full((bsize, *cache_shape), CACHE_FILL_VALUE, device=device, dtype=dtype).transpose(0, 1)

    def deinit_cache(self):
        self.cache = [None] * self.cache_num_layers
    
    @T.no_grad()
    def forward(self, data, next_tokens: Optional[Tuple[T.Tensor, T.Tensor]] = None, temps: Optional[Tuple[float, Tuple[float, float]]] = None):
        if self.c.split:
            x1, x2 = data.chunk(2, dim=-1)
            x = self.io.input(x1) + self.io2.input(x2)
        else:
            x = self.io.input(data)

        cache_idx = 0
        for l, layer in enumerate(self.stack.layers):
            if l == self.plex_layer:
                if self.c.split:
                    plex1, plex2 = self.quantize(data)
                    plex1 = T.roll(plex1, -self.c.plex_roll, dims=1)
                    plex2 = T.roll(plex2, -self.c.plex_roll, dims=1)
                    if exists(next_tokens):
                        plex1[:, -1:] = self.untokenize(next_tokens[0])
                        plex2[:, -1:] = self.untokenize(next_tokens[1])
                    x1 = x + self.plex_projection(plex1)
                    x2 = x + self.plex_projection2(plex2)
                else:
                    plex = self.quantize(data)
                    plex = T.roll(plex, -self.c.plex_roll, dims=1)
                    if exists(next_tokens):
                        plex[:, -1:] = self.untokenize(next_tokens)
                    x = x + self.plex_projection(plex)

            if l < self.plex_layer:
                x = layer(x, kv=self.cache[l])
            else:
                if self.c.split:
                    x1 = layer(x1, kv=self.cache[self.plex_layer + cache_idx])
                    cache_idx += 1
                    x2 = layer(x2, kv=self.cache[self.plex_layer + cache_idx])
                    cache_idx += 1
                else:
                    x = layer(x, kv=self.cache[l])

        with T.autocast(device_type='cuda', dtype=T.bfloat16):
            if self.c.split:
                x1, x2 = self.out_norm(x1), self.out_norm(x2)
                out1, out2 = self.io.output(x1), self.io.output(x2)
            else:
                x = self.out_norm(x)
                out = self.io.output(x)

        if isnt(temps):
            if self.c.split:
                return out1, out2
            else:
                return out
        else:
            if self.c.split:
                next_data1 = self.io.temp_sample(out1, temps)[:, -1:, :]
                next_data2 = self.io2.temp_sample(out2, temps)[:, -1:, :]
                next_data = T.cat([next_data1, next_data2], dim=-1)
                return next_data
            else:
                next_data = self.io.temp_sample(out, temps)[:, -1:, :]
                return next_data

@si_module
class HertzDevModel(nn.Module):
    class Config:
        dim: int
        vocab_size: int
        stack_config: Optional[Stack.Config] = None
        latent_size: int = 32

        split: bool = True

        quantizer_config: Optional[LatentQuantizer.Config] = None
        resynthesizer_config: Optional[TransformerVAE.Config] = 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.stack_config)), f'hmm {c}'

        self.input = nn.Linear(c.latent_size, c.dim)
        if self.c.split:
            self.input2 = nn.Linear(c.latent_size, c.dim)

        self.shape_rotator = ShapeRotator(c.stack_config.dim//c.stack_config.n_head, c.stack_config.seq_len, theta=c.stack_config.theta)

        self.layers = nn.ModuleList([
            PerfBlock(
                dim=c.stack_config.dim,
                layer_id=l,
                n_head=c.stack_config.n_head,
                kv_heads=c.stack_config.kv_heads,
                ff_dim=c.stack_config.ff_dim,
                eps=c.stack_config.eps,
                shape_rotator=self.shape_rotator,
            ) for l in range(c.stack_config.layers)
        ])

        self.output = GPTOutput(c.dim, c.vocab_size)
        if self.c.split:
            self.output2 = GPTOutput(c.dim, c.vocab_size)

        self.cache = [None] * c.stack_config.layers
        self.kv_heads = c.stack_config.kv_heads or c.stack_config.n_head
        self.head_dim = c.stack_config.dim // c.stack_config.n_head

        if exists(c.from_pretrained):
            result = self.load_state_dict(checkpoint, strict=False)
            print0_colored(result, 'yellow')

        self.resynthesizer = c.resynthesizer_config().eval()
        self.resynthesizer.requires_grad = False

        self.audio_tokenizer = make_tokenizer(device='cpu')
        self.audio_cache = None
        self.audio_latent_cache = None
        self.use_audio_cache = False

    @T.no_grad()
    def tokenize(self, audio_data):
        orig_audio_shape = audio_data.shape
        if exists(self.audio_cache):
            audio_data = T.cat([self.audio_cache, audio_data], dim=-1)
            self.audio_cache = audio_data[..., -(6*16_000):]
        elif self.use_audio_cache:
            self.audio_cache = audio_data[..., -(6*16_000):]

        if audio_data.shape[1] == 2:
            enc_ch1 = self.audio_tokenizer.latent_from_data(audio_data[:, 0:1])
            enc_ch2 = self.audio_tokenizer.latent_from_data(audio_data[:, 1:2])
            return T.cat([enc_ch1, enc_ch2], dim=-1)[:, -(orig_audio_shape[-1]//2000):]
        else:
            return self.audio_tokenizer.latent_from_data(audio_data)[:, -(orig_audio_shape[-1]//2000):]

    @T.no_grad()
    def untokenize(self, token_data):
        if exists(self.audio_latent_cache):
            token_data = T.cat([self.audio_latent_cache, token_data], dim=1)
            self.audio_latent_cache = token_data[:, -(6*8):]
        elif self.use_audio_cache:
            self.audio_latent_cache = token_data[:, -(6*8):]
        
        if token_data.shape[-1] == 2*self.c.latent_size:
            dec_ch1 = self.audio_tokenizer.data_from_latent(token_data[:, :self.c.latent_size])
            dec_ch2 = self.audio_tokenizer.data_from_latent(token_data[:, self.c.latent_size:])
            return T.cat([dec_ch1, dec_ch2], dim=1)[..., -(token_data.shape[1]*2000):]
        else:
            return self.audio_tokenizer.data_from_latent(token_data)[..., -(token_data.shape[1]*2000):]

    def init_cache(self, bsize, device, dtype, length:int=None):
        cache_shape = [self.c.stack_config.layers, length or self.c.stack_config.seq_len, 2, self.kv_heads, self.head_dim]
        self.cache = T.full((bsize, *cache_shape), CACHE_FILL_VALUE, device=device, dtype=dtype).transpose(0, 1)
        self.resynthesizer.init_cache(bsize, device, dtype, length)
        self.use_audio_cache = True

    def deinit_cache(self):
        self.cache = [None] * len(self.layers)
        self.resynthesizer.deinit_cache()
        self.audio_cache = None
        self.audio_latent_cache = None
        self.use_audio_cache = False
    
    @T.no_grad()
    def forward(self, data):
        if self.c.split:    
            x1, x2 = data.chunk(2, dim=-1)
            x = self.input(x1) + self.input2(x2)
        else:
            x = self.input(data)

        for l, layer in enumerate(self.layers):
            x = layer(x, kv=self.cache[l])

        if self.c.split:
            return self.output(x), self.output2(x)
        else:
            return self.output(x)
        
    @T.no_grad()
    def next_audio_from_audio(self, audio_data: T.Tensor, temps=(0.8, (0.5, 0.1))):
        latents_in = self.tokenize(audio_data)
        next_latents = self.next_latent(latents_in, temps)
        next_model_latent = next_latents[..., self.c.latent_size:]
        audio_decoded = self.untokenize(next_model_latent)[..., -2000:]
        return audio_decoded

        
    @T.no_grad()
    def next_latent(self, model_input: T.Tensor, temps=(0.8, (0.5, 0.1))):

        if self.c.split:
            logits1, logits2 = self.forward(model_input)
            next_logits1 = logits1[:, -1]
            next_logits2 = logits2[:, -1]
            next_token1 = F.softmax(next_logits1 / temps[0], dim=-1).multinomial(1)
            next_token2 = F.softmax(next_logits2 / temps[0], dim=-1).multinomial(1)

            next_input = self.resynthesizer(model_input, next_tokens=(next_token1, next_token2), temps=temps[1])
        else:
            logits = self.forward(model_input)
            next_logits = logits[:, -1]
            next_token = F.softmax(next_logits / temps[0], dim=-1).multinomial(1)

            next_input = self.resynthesizer(model_input, next_tokens=next_token, temps=temps[1])

        return next_input


    @T.no_grad()
    def completion(self, data: T.Tensor, temps=(0.8, (0.5, 0.1)), gen_len=None, use_cache=True) -> T.Tensor:
        """
        only accepts latent-space data.
        """
        if use_cache:
            self.init_cache(data.shape[0], data.device, T.bfloat16)

        next_input = generated = data

        target_len = min(data.shape[1] + default(gen_len, data.shape[1]), self.c.stack_config.seq_len)
        
        for _ in tqdm0(range(data.shape[1], target_len)):
            model_input = next_input if use_cache else generated

            next_input = self.next_latent(model_input, temps)
    
            generated = T.cat([generated, next_input], dim=1)

        if use_cache:
            self.deinit_cache()
        return generated
    


def get_hertz_dev_config(is_split=True, use_pure_audio_ablation=False):
    if is_split:
        checkpoints = [('inference_care_50000', 'e4ff4fe5c7e9f066410d2a5673b7a935'), ('inference_scion_54000', 'cb8bc484423922747b277ebc2933af5d')]
    elif not use_pure_audio_ablation:
        checkpoints = [('inference_whip_72000', '5e7cee7316900737d55fc5d44cc7a8f7'), ('inference_caraway_112000', 'fcb8368ef8ebf7712f3e31e6856da580')]
    else:
        checkpoints = [('inference_whip_72000', '5e7cee7316900737d55fc5d44cc7a8f7'), ('inference_syrup_110000', '353c48f553f1706824c11f3bb6a049e9')]

    quantizer_config=LatentQuantizer.Config(
        from_pretrained=('inference_volcano_3', 'd42bf674022c5f84b051d5d7794f6169'),
        compressor_config=FSQ.Config(
            levels=[8,8,8,8,8],
            dim=2048,
            num_codebooks=1,
            keep_num_codebooks_dim=None,
            scale=None,
            allowed_dtypes=['float32', 'float64', 'bfloat16'],
            channel_first=False,
            projection_has_bias=True,
            return_indices=True,
            force_quantization_f32=True,
            use_rms=False
        ),
        dim=2048,
        ff_dim=8192,
        input_dim=32
    )

    resynthesizer_config=TransformerVAE.Config(
        io_config=GaussianMixtureIOLayer.Config(
            latent_dim=32,
            dim=4096,
            num_components=8,
        ),
        stack_config=Stack.Config(
            layers=8,
            dim=4096,
            seq_len=8192,
            n_head=16,
            ff_dim=11008,
            kv_heads=16,
            eps=1e-5,
            theta=10_000
        ),
        quantizer_config=quantizer_config,
        plex_layer=None,
        plex_roll=1,
        split=is_split,
        from_pretrained=checkpoints[0],
    )

    return HertzDevModel.Config(
        dim=4096,
        vocab_size=32_768,
        stack_config=Stack.Config(
            layers=32,
            dim=4096,
            seq_len=2048,
            n_head=32,
            ff_dim=None,
            kv_heads=None,
            eps=1e-5,
            theta=10_000,
        ),
        quantizer_config=quantizer_config,
        resynthesizer_config=resynthesizer_config,
        split=is_split,
        from_pretrained=checkpoints[1],
    )