File size: 2,424 Bytes
657c0f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import torch
import torch.nn as nn
from transformers import EncodecModel, EncodecConfig
from model.attention_weighted_pooling import AttentionWeightedPooling
from model.projection import ProjectionHead
from model.chordbeat_encoder import ChordBeatEncoder
class MusicConRec(nn.Module):
def __init__(self, codebook_size=1024, feature_dim=128, proj_dim=128):
super().__init__()
# === AUDIO SIDE ===
self.encodec = EncodecModel.from_pretrained("facebook/encodec_24khz")
# Do not freeze EncodecModel parameters — allow fine-tuning
for param in self.encodec.parameters():
param.requires_grad = True
self.code_embedding = nn.Embedding(codebook_size, feature_dim)
self.audio_pool = AttentionWeightedPooling(feature_dim)
self.audio_proj = ProjectionHead(feature_dim, out_dim=proj_dim)
# === CHORD SIDE ===
self.chord_encoder = ChordBeatEncoder(
input_dim=13,
d_model=feature_dim
)
def forward(self, audio, chord):
"""
audio: (B, 1, T)
chord: (B, T_chord, 13)
"""
# =========================
# ENCODE
# =========================
encoder_outputs = self.encodec.encode(audio)
audio_codes = encoder_outputs['audio_codes'].long()
audio_scales = encoder_outputs['audio_scales']
codes = audio_codes.squeeze(0).permute(0, 2, 1)
codes = self.code_embedding(codes) # (B, T, Q, D)
codes = codes.mean(dim=2) # (B, T, D)
# =========================
# POOL + PROJECT
# =========================
h_audio = self.audio_pool(codes) # (B, D)
z_audio = self.audio_proj(h_audio) # (B, proj_dim)
# =========================
# RECONSTRUCTION
# =========================
x_recon = self.encodec.decode(audio_codes, audio_scales)['audio_values']
x_recon = torch.tanh(x_recon).clamp(-1.0, 1.0)
# =========================
# CHORD BRANCH
# =========================
z_chord, h_chord = self.chord_encoder(chord)
return {
"x_recon": x_recon,
"z_audio": z_audio,
"z_chord": z_chord,
"h_audio": h_audio,
"h_chord": h_chord
} |