|
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], |
|
) |
