Spaces:
Running on Zero
Running on Zero
| from typing import List, Sequence | |
| from torch import Tensor | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from dataclasses import dataclass | |
| from utils.audio import load_audio, make_variable_frequency_sinewave | |
| from utils.plotting import plot_waveforms | |
| from conditioning.embedded_condition import EmbeddedCondition | |
| from conditioning.embedder import Embedder, LinearProjectionEmbedder | |
| import config as cfg | |
| class Beat: | |
| """ | |
| beats / downbeats: | |
| Tensor of Int or Long type, containing indices of samples where | |
| beats / downbeats occur | |
| seq_len: int representing the total length of the audio in samples | |
| """ | |
| beats: Tensor | |
| downbeats: Tensor | |
| seq_len: int | |
| class DirectSinusoidalBeatEmbedder(Embedder): | |
| def __init__(self, embedding_dim: int): | |
| super().__init__(input_dim=2, embedding_dim=embedding_dim) | |
| self.encodec_framerate: int = 50 | |
| self.sample_rate: int = 32_000 | |
| def forward(self, | |
| x: Sequence[Beat], | |
| duplicate_for_cfg: bool = False) -> EmbeddedCondition: | |
| """ | |
| Returns: A tensor of shape [B, S, H], where B is batch, H is | |
| embedding dim, and S is the length of the sequence | |
| """ | |
| assert len(set([t.seq_len for t in x])) == 1 | |
| emb_list: List[Tensor] = [None] * len(x) # type: ignore | |
| ratio = self.encodec_framerate / self.sample_rate | |
| for i, track in enumerate(x): | |
| # convert indices from explicit to latent space | |
| beats = torch.floor(track.beats.float() * ratio).long() | |
| downbeats = torch.floor(track.downbeats.float() * ratio).long() | |
| seq_len = round(track.seq_len * ratio) | |
| # make a variable-frequency sinewave lined with the beats | |
| beat_emb = make_variable_frequency_sinewave(seq_len, beats) | |
| downbeat_emb = make_variable_frequency_sinewave(seq_len, downbeats) | |
| emb_list[i] = torch.stack((beat_emb, downbeat_emb), dim=-1) | |
| embeds: Tensor = torch.stack(emb_list) | |
| # create an embedding that's just the beat sinewave in every dimension | |
| embeds = torch.cat( | |
| (embeds[..., :1].repeat(1, 1, self.embedding_dim // 2), | |
| embeds[..., 1:].repeat(1, 1, self.embedding_dim // 2)), | |
| dim=-1) | |
| mask = torch.ones(embeds.shape[:-1], | |
| device=embeds.device, | |
| dtype=torch.bool) | |
| # concatenate an empty condition | |
| if duplicate_for_cfg: | |
| embeds = torch.cat((embeds, torch.zeros_like(embeds)), dim=0) | |
| mask = torch.cat((mask, torch.zeros_like(mask)), dim=0) | |
| return EmbeddedCondition(embeds, mask) | |
| def null_condition(self, batch_size: int) -> EmbeddedCondition: | |
| return EmbeddedCondition( | |
| torch.zeros(batch_size, | |
| 1, | |
| self.embedding_dim, | |
| dtype=torch.float32, | |
| device=self.output_proj.weight.device), | |
| torch.zeros(batch_size, | |
| 1, | |
| dtype=torch.bool, | |
| device=self.output_proj.weight.device)) | |
| class SinusoidalBeatEmbedder(LinearProjectionEmbedder): | |
| def __init__(self, embedding_dim: int): | |
| super().__init__(input_dim=2, embedding_dim=embedding_dim) | |
| self.encodec_framerate: int = 50 | |
| self.sample_rate: int = 32_000 | |
| def forward(self, | |
| x: Sequence[Beat], | |
| duplicate_for_cfg: bool = False) -> EmbeddedCondition: | |
| """ | |
| Returns: A tensor of shape [B, S, H], where B is batch, H is | |
| embedding dim, and S is the length of the sequence | |
| """ | |
| assert len(set([t.seq_len for t in x])) == 1 | |
| emb_list: List[Tensor] = [None] * len(x) # type: ignore | |
| for i, track in enumerate(x): | |
| beats = (track.beats / self.sample_rate * | |
| self.encodec_framerate).round().int() | |
| downbeats = (track.downbeats / self.sample_rate * | |
| self.encodec_framerate).round().int() | |
| seq_len = round(track.seq_len / self.sample_rate * | |
| self.encodec_framerate) | |
| beat_emb = make_variable_frequency_sinewave(seq_len, beats) | |
| downbeat_emb = make_variable_frequency_sinewave(seq_len, downbeats) | |
| emb_list[i] = torch.stack((beat_emb, downbeat_emb), dim=-1) | |
| embeds: Tensor = torch.stack(emb_list) | |
| embeds = self.output_proj(embeds.to(self.output_proj.weight)) | |
| mask = torch.ones(embeds.shape[:-1], | |
| device=embeds.device, | |
| dtype=torch.bool) | |
| if duplicate_for_cfg: | |
| embeds = torch.cat((embeds, torch.zeros_like(embeds)), dim=0) | |
| mask = torch.cat((mask, torch.zeros_like(mask)), dim=0) | |
| return EmbeddedCondition(embeds, mask) | |
| def null_condition(self, batch_size: int) -> EmbeddedCondition: | |
| return EmbeddedCondition( | |
| torch.zeros(batch_size, | |
| 1, | |
| self.embedding_dim, | |
| dtype=torch.float32, | |
| device=self.output_proj.weight.device), | |
| torch.zeros(batch_size, | |
| 1, | |
| dtype=torch.bool, | |
| device=self.output_proj.weight.device)) | |
| class SinusoidalBeatEmbedderMLP(Embedder): | |
| def __init__(self, embedding_dim: int): | |
| super().__init__(input_dim=2, embedding_dim=embedding_dim) | |
| # self.embedding_dim: int = embedding_dim | |
| self.encodec_framerate: int = 50 | |
| self.sample_rate: int = 32_000 | |
| self.output_proj: nn.Sequential = nn.Sequential( | |
| nn.Linear(self.input_dim, 512), nn.ReLU(), | |
| nn.Linear(512, self.embedding_dim)) | |
| def forward(self, | |
| x: Sequence[Beat], | |
| duplicate_for_cfg: bool = False) -> EmbeddedCondition: | |
| """ | |
| Returns: A tensor of shape [B, S, H], where B is batch, H is | |
| embedding dim, and S is the length of the sequence | |
| """ | |
| assert len(set([t.seq_len for t in x])) == 1 | |
| emb_list: List[Tensor] = [None] * len(x) # type: ignore | |
| for i, track in enumerate(x): | |
| beats = (track.beats / self.sample_rate * | |
| self.encodec_framerate).round().int() | |
| downbeats = (track.downbeats / self.sample_rate * | |
| self.encodec_framerate).round().int() | |
| seq_len = round(track.seq_len / self.sample_rate * | |
| self.encodec_framerate) | |
| beat_emb = make_variable_frequency_sinewave(seq_len, beats) | |
| downbeat_emb = make_variable_frequency_sinewave(seq_len, downbeats) | |
| emb_list[i] = torch.stack((beat_emb, downbeat_emb), dim=-1) | |
| embeds: Tensor = torch.stack(emb_list) | |
| embeds = self.output_proj(embeds.to(self.output_proj[-1].weight)) | |
| mask = torch.ones(embeds.shape[:-1], | |
| device=embeds.device, | |
| dtype=torch.bool) | |
| if duplicate_for_cfg: | |
| embeds = torch.cat((embeds, torch.zeros_like(embeds)), dim=0) | |
| mask = torch.cat((mask, torch.zeros_like(mask)), dim=0) | |
| return EmbeddedCondition(embeds, mask) | |
| def null_condition(self, batch_size: int) -> EmbeddedCondition: | |
| return EmbeddedCondition( | |
| torch.zeros(batch_size, | |
| 1, | |
| self.embedding_dim, | |
| dtype=torch.float32, | |
| device=self.output_proj[-1].weight.device), | |
| torch.zeros(batch_size, | |
| 1, | |
| dtype=torch.bool, | |
| device=self.output_proj[-1].weight.device)) | |
| if __name__ == "__main__": | |
| song_path = cfg.moises_path( | |
| ) / "0d528a19-cb0f-4421-b250-444f9343e51c/mixed.wav" | |
| song = load_audio(song_path) | |
| n_samples = song.shape[-1] | |
| loaded = np.load( | |
| "/home/tkol/dev/datasets/moisesdb/moisesdb_v0.1/0d528a19-cb0f-4421-b250-444f9343e51c/beatthis.npz" | |
| ) | |
| beats = torch.tensor([int(b * 32_000) for b in loaded["beats"]], | |
| dtype=torch.int32) | |
| downbeats = torch.tensor([int(b * 32_000) for b in loaded["downbeats"]], | |
| dtype=torch.int32) | |
| assert n_samples > beats.max() | |
| b = Beat(beats=beats, downbeats=downbeats, seq_len=n_samples) | |
| embedder = SinusoidalBeatEmbedder(1024) | |
| emb = embedder([b]) | |
| directembedder = DirectSinusoidalBeatEmbedder(1024) | |
| directemb = directembedder([b]) | |
| # from matplotlib import pyplot as plt | |
| # plt.style.use(cfg.ROOT / "tkol.mplstyle") | |
| # plt.figure(figsize=(100, 10)) | |
| # plt.scatter(x=torch.arange(emb.shape[-1]) / 50 * 32_000, y=emb) | |
| # plt.show() | |
| # plot_waveforms( | |
| # song, | |
| # emb, | |
| # # savepath=cfg.ROOT / "plots/beatontrack.png", | |
| # figsize=(50, 3), | |
| # dpi=100, | |
| # ) | |