Uploaded model

model/attention_weighted_pooling.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class AttentionWeightedPooling(nn.Module):
+    def __init__(self, in_dim, hidden_dim=128):
+        super().__init__()
+
+        # equivalent to conv blocks in paper → here MLP over time
+        self.attn = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        """
+        x: (B, T, C)
+        """
+        # compute attention weights
+        weights = self.attn(x)  # (B, T, 1)
+
+        # apply weights
+        weighted = x * weights  # (B, T, C)
+
+        # weighted average pooling
+        pooled = weighted.sum(dim=1) / (weights.sum(dim=1) + 1e-8)
+
+        return pooled  # (B, C)

model/chordbeat_encoder.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=500):
+        super().__init__()
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1)
+
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model)
+        )
+
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        self.pe = pe.unsqueeze(0)  # (1, T, D)
+
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1)].to(x.device)
+    
+
+class ChordBeatEncoder(nn.Module):
+    def __init__(self, input_dim=13, d_model=128, nhead=4, num_layers=3):
+        super().__init__()
+
+        # project input → model dim
+        self.input_proj = nn.Linear(input_dim, d_model)
+
+        self.pos_enc = PositionalEncoding(d_model)
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            batch_first=True
+        )
+
+        self.transformer = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=num_layers
+        )
+
+        # pooling (same idea as your audio side)
+        self.pool = nn.Sequential(
+            nn.Linear(d_model, 1),
+            nn.Softmax(dim=1)
+        )
+
+        self.output_proj = nn.Linear(d_model, d_model)
+
+    def forward(self, x):
+        """
+        x: (B, T, 13)
+        """
+
+        x = self.input_proj(x)      # (B, T, D)
+        x = self.pos_enc(x)
+
+        x = self.transformer(x)     # (B, T, D)
+
+        # attention pooling
+        weights = self.pool(x)      # (B, T, 1)
+        h = (x * weights).sum(dim=1)
+
+        z = self.output_proj(h)
+
+        return z, h

model/model.py

@@ -0,0 +1,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
+        }

model/projection.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ProjectionHead(nn.Module):
+    def __init__(self, in_dim, hidden_dim=512, out_dim=128):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.BatchNorm1d(hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, out_dim)
+        )
+
+    def forward(self, x):
+        return self.net(x)