| from dataclasses import dataclass, asdict
|
| from typing import NamedTuple
|
| import torch
|
| from torch import Tensor, nn
|
| import torch.nn.functional as F
|
| MOODS=()
|
| ROLES=()
|
| INSTRUMENTS=()
|
| FRAME_FEATURES=()
|
| CONTROL_FEATURES=()
|
|
|
| @dataclass(frozen=True)
|
| class A02bConfig:
|
| width: int = 96
|
| layers: int = 4
|
| heads: int = 4
|
| dropout: float = 0.1
|
| frames_per_beat: int = 16
|
| instrument_embedding_dim: int = 8
|
| use_convolution: bool = True
|
| use_relative_bias: bool = True
|
| use_stop_head: bool = False
|
| relative_transitions_beats: tuple[float, ...] = (2 / 16, 3 / 16, 4 / 16, 5 / 16, 6 / 16, 7 / 16, 8 / 16, 3 / 4, 1, 2, 4, 8, 16, 32)
|
|
|
| def __post_init__(self) -> None:
|
| if self.instrument_embedding_dim <= 0:
|
| raise ValueError('instrument_embedding_dim must be positive')
|
| if self.width <= 0 or self.heads <= 0 or self.width % self.heads:
|
| raise ValueError('width must be positive and divisible by heads')
|
| if self.layers <= 0 or not 0 <= self.dropout < 1:
|
| raise ValueError('layers must be positive; dropout must be in [0,1)')
|
| if self.frames_per_beat not in (16, 48):
|
| raise ValueError('A02b supports separately configured rates 16 and 48')
|
| edges = self.relative_transitions_beats
|
| if len(edges) != 14 or not all((x > 0 for x in edges)):
|
| raise ValueError('32 signed buckets require 14 positive transitions')
|
| if not all((a < b for a, b in zip(edges, edges[1:]))):
|
| raise ValueError('relative distance transitions must increase')
|
| if not all((abs(x * self.frames_per_beat - round(x * self.frames_per_beat)) < 1e-08 for x in edges)):
|
| raise ValueError('bucket boundaries must align with the frame grid')
|
|
|
| class A02bOutput(NamedTuple):
|
| normalized_cc1: Tensor
|
| stop_logits: Tensor | None
|
|
|
| def zero_padding(h: Tensor, mask: Tensor) -> Tensor:
|
| return h.masked_fill(~mask[..., None], 0.0)
|
|
|
| class SignedRelativeBias(nn.Module):
|
|
|
| def __init__(self, config: A02bConfig) -> None:
|
| super().__init__()
|
| self.register_buffer('transitions_frames', torch.tensor([round(x * config.frames_per_beat) for x in config.relative_transitions_beats], dtype=torch.long))
|
| self.table = nn.Embedding(32, config.heads)
|
| nn.init.zeros_(self.table.weight)
|
|
|
| def bucket_ids(self, length: int) -> Tensor:
|
| positions = torch.arange(length, device=self.transitions_frames.device)
|
| offset = positions[None, :] - positions[:, None]
|
| magnitude_bucket = torch.ones_like(offset)
|
| for transition in self.transitions_frames.unbind():
|
| magnitude_bucket = magnitude_bucket + (offset.abs() >= transition)
|
| signed_bucket = magnitude_bucket + (offset > 0).long() * 16
|
| return torch.where(offset == 0, 0, signed_bucket)
|
|
|
| def forward(self, length: int) -> Tensor:
|
| return self.table(self.bucket_ids(length)).permute(2, 0, 1)[None]
|
|
|
| class FiLM(nn.Module):
|
|
|
| def __init__(self, width: int) -> None:
|
| super().__init__()
|
| self.modulation = nn.Linear(64, 2 * width)
|
| nn.init.zeros_(self.modulation.weight)
|
| nn.init.zeros_(self.modulation.bias)
|
|
|
| def forward(self, u: Tensor, c: Tensor) -> Tensor:
|
| delta_gamma, beta = self.modulation(c).chunk(2, dim=-1)
|
| return (1 + delta_gamma[:, None]) * u + beta[:, None]
|
|
|
| class BidirectionalAttention(nn.Module):
|
|
|
| def __init__(self, config: A02bConfig) -> None:
|
| super().__init__()
|
| self.heads = config.heads
|
| self.head_dim = config.width // config.heads
|
| self.dropout = config.dropout
|
| self.qkv = nn.Linear(config.width, 3 * config.width)
|
| self.output = nn.Linear(config.width, config.width)
|
|
|
| def forward(self, u: Tensor, attention_bias: Tensor) -> Tensor:
|
| batch, length, width = u.shape
|
| qkv = self.qkv(u).reshape(batch, length, 3, self.heads, self.head_dim)
|
| q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
|
| result = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_bias, dropout_p=self.dropout if self.training else 0.0, is_causal=False)
|
| result = result.transpose(1, 2).reshape(batch, length, width)
|
| return self.output(result)
|
|
|
| class EncoderBlock(nn.Module):
|
|
|
| def __init__(self, config: A02bConfig) -> None:
|
| super().__init__()
|
| d = config.width
|
| self.norm_attention = nn.LayerNorm(d)
|
| self.norm_ff = nn.LayerNorm(d)
|
| self.film_attention = FiLM(d)
|
| self.film_ff = FiLM(d)
|
| self.attention = BidirectionalAttention(config)
|
| self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
|
| self.dropout = nn.Dropout(config.dropout)
|
|
|
| def forward(self, h: Tensor, c: Tensor, mask: Tensor, bias: Tensor) -> Tensor:
|
| u = self.film_attention(self.norm_attention(h), c)
|
| h = zero_padding(h + self.dropout(self.attention(u, bias)), mask)
|
| u = self.film_ff(self.norm_ff(h), c)
|
| return zero_padding(h + self.dropout(self.ff(u)), mask)
|
|
|
| class AireA02b(nn.Module):
|
|
|
| def __init__(self, config: A02bConfig | None=None) -> None:
|
| super().__init__()
|
| self.config = config or A02bConfig()
|
| cfg = self.config
|
| d = cfg.width
|
| self.mood_embedding = nn.Embedding(len(MOODS), 16)
|
| self.role_embedding = nn.Embedding(len(ROLES), 8)
|
| self.instrument_embedding = nn.Embedding(len(INSTRUMENTS), cfg.instrument_embedding_dim)
|
| self.control_mlp = nn.Sequential(nn.Linear(30 + cfg.instrument_embedding_dim, 64), nn.GELU(), nn.Linear(64, 64))
|
| self.lead_embedding = nn.Embedding(129, d)
|
| self.pitch_projection = nn.Linear(128, d, bias=False)
|
| self.onset_projection = nn.Linear(128, d, bias=False)
|
| self.offset_projection = nn.Linear(128, d, bias=False)
|
| self.numeric_projection = nn.Linear(25, d)
|
| self.control_projection = nn.Linear(64, d)
|
| self.convolutions = nn.ModuleList([nn.Conv1d(d, d, kernel_size=5, padding=2) for _ in range(2)] if cfg.use_convolution else [])
|
| self.relative_bias = SignedRelativeBias(cfg) if cfg.use_relative_bias else None
|
| self.blocks = nn.ModuleList([EncoderBlock(cfg) for _ in range(cfg.layers)])
|
| self.final_norm = nn.LayerNorm(d)
|
| self.cc_head = nn.Linear(d, 1)
|
| self.stop_head = nn.Linear(d, 1) if cfg.use_stop_head else None
|
|
|
| def manifest(self) -> dict:
|
| return {'model': 'aire-A02b', 'input_schema': 'A02b-25-instrument-v1', 'config': asdict(self.config), 'moods': list(MOODS), 'roles': list(ROLES), 'frame_features': list(FRAME_FEATURES), 'control_features': list(CONTROL_FEATURES), 'instruments': list(INSTRUMENTS), 'velocity_output': 'preserve-input-only', 'rest_pitch_index': 128, 'relative_bucket_rule': 'signed-piecewise-beats-v1; zero=0; past=1:15; future=17:31'}
|
|
|
| def forward(self, *, lead_pitch: Tensor, pitch_roll: Tensor, onset_roll: Tensor, offset_roll: Tensor, frame_numeric: Tensor, mood_id: Tensor, role_id: Tensor, instrument_id: Tensor, ctrl_numeric: Tensor, mask: Tensor) -> A02bOutput:
|
| if mask.ndim != 2 or mask.dtype != torch.bool:
|
| raise ValueError('mask must be bool [B,T]')
|
| batch, length = mask.shape
|
| inputs = {'lead_pitch': (lead_pitch, (batch, length), torch.long), 'pitch_roll': (pitch_roll, (batch, length, 128), self.cc_head.weight.dtype), 'onset_roll': (onset_roll, (batch, length, 128), self.cc_head.weight.dtype), 'offset_roll': (offset_roll, (batch, length, 128), self.cc_head.weight.dtype), 'frame_numeric': (frame_numeric, (batch, length, 25), self.cc_head.weight.dtype), 'mood_id': (mood_id, (batch,), torch.long), 'role_id': (role_id, (batch,), torch.long), 'instrument_id': (instrument_id, (batch,), torch.long), 'ctrl_numeric': (ctrl_numeric, (batch, 6), self.cc_head.weight.dtype)}
|
| device = self.cc_head.weight.device
|
| if mask.device != device:
|
| raise ValueError('mask and model must be on the same device')
|
| for name, (tensor, shape, dtype) in inputs.items():
|
| if tuple(tensor.shape) != shape or tensor.dtype != dtype or tensor.device != device:
|
| raise ValueError(f'{name} must have shape {shape}, dtype {dtype}, device {device}')
|
| if batch == 0 or length == 0:
|
| empty = self.cc_head.weight.new_zeros((batch, length))
|
| return A02bOutput(empty, empty if self.stop_head is not None else None)
|
| for name, tensor, count in (('mood_id', mood_id, len(MOODS)), ('role_id', role_id, len(ROLES)), ('instrument_id', instrument_id, len(INSTRUMENTS))):
|
| if bool(((tensor < 0) | (tensor >= count)).any()):
|
| raise ValueError(f'{name} is outside the frozen category range')
|
| if not bool(mask.any(dim=1).all()):
|
| raise ValueError('all-padding samples must bypass the model')
|
| if bool((~mask[:, :-1] & mask[:, 1:]).any()):
|
| raise ValueError('mask must consist of a valid prefix and right padding')
|
| lead_pitch = lead_pitch.masked_fill(~mask, 128)
|
| pitch_roll = zero_padding(pitch_roll, mask)
|
| onset_roll = zero_padding(onset_roll, mask)
|
| offset_roll = zero_padding(offset_roll, mask)
|
| numeric = zero_padding(frame_numeric, mask)
|
| velocity_enabled = ctrl_numeric[:, 5:6] != 0
|
| gate = velocity_enabled[:, None, :]
|
| numeric = torch.cat([numeric[..., :3], torch.where(gate, numeric[..., 3:4], 0.0), numeric[..., 4:24], torch.where(gate, numeric[..., 24:25], 0.0)], dim=-1)
|
| c = self.control_mlp(torch.cat([self.mood_embedding(mood_id), self.role_embedding(role_id), self.instrument_embedding(instrument_id), ctrl_numeric], dim=-1))
|
| h = zero_padding(self.lead_embedding(lead_pitch) + self.pitch_projection(pitch_roll) + self.onset_projection(onset_roll) + self.offset_projection(offset_roll) + self.numeric_projection(numeric) + self.control_projection(c)[:, None], mask)
|
| if self.convolutions:
|
| residual = h
|
| for convolution in self.convolutions:
|
| h = F.gelu(convolution(h.transpose(1, 2))).transpose(1, 2)
|
| h = zero_padding(h, mask)
|
| h = zero_padding(residual + h, mask)
|
| bias = h.new_zeros((batch, 1, 1, length)).masked_fill(~mask[:, None, None, :], float('-inf'))
|
| if self.relative_bias is not None:
|
| bias = bias + self.relative_bias(length).to(h.dtype)
|
| for block in self.blocks:
|
| h = block(h, c, mask, bias)
|
| h = self.final_norm(h)
|
| y = torch.sigmoid(self.cc_head(h).squeeze(-1)).masked_fill(~mask, 0.0)
|
| stop = None
|
| if self.stop_head is not None:
|
| stop = self.stop_head(h).squeeze(-1).masked_fill(~mask, 0.0)
|
| return A02bOutput(y, stop)
|
|
|
