Update train.py

fixing tupple issue

train.py

@@ -1,7 +1,6 @@
 import os
 import torch
 import torchaudio
-import numpy as np
 from torch.utils.data import Dataset, DataLoader
 import torch.nn as nn
 import torch.optim as optim
@@ -9,174 +8,286 @@ from pathlib import Path
 import argparse
 from torch.nn.utils.rnn import pad_sequence
 
-# ---- RVC v2 Architecture (Hubert + Pitch + ContentVec) ----
+
+# ---- Simplified RVC-like Architecture ----
 class HubertEncoder(nn.Module):
-    def __init__(self, input_dim=1024, hidden_dim=768):
+    def __init__(self, input_dim=128, hidden_dim=256):
         super().__init__()
+
         self.conv1 = nn.Conv1d(input_dim, hidden_dim, 3, padding=1)
         self.conv2 = nn.Conv1d(hidden_dim, hidden_dim, 3, padding=1)
-        self.lstm = nn.LSTM(hidden_dim, hidden_dim//2, 2, batch_first=True, bidirectional=True)
+        self.lstm = nn.LSTM(
+            hidden_dim,
+            hidden_dim // 2,
+            num_layers=2,
+            batch_first=True,
+            bidirectional=True,
+        )
         self.proj = nn.Linear(hidden_dim, 256)
-    
+
     def forward(self, x):
-        x = x.transpose(1, 2)  # (B, T, F) -> (B, F, T)
+        # x: (B, T, 128)
+        x = x.transpose(1, 2)  # (B, 128, T)
         x = torch.relu(self.conv1(x))
         x = torch.relu(self.conv2(x))
-        x = x.transpose(1, 2)  # Back to (B, T, F)
+        x = x.transpose(1, 2)  # (B, T, hidden)
         out, _ = self.lstm(x)
-        return self.proj(out)  # 256-dim features
+        return self.proj(out)  # (B, T, 256)
+
 
 class PitchEncoder(nn.Module):
     def __init__(self):
         super().__init__()
+
         self.f0_conv = nn.Sequential(
             nn.Conv1d(1, 64, 3, padding=1),
             nn.ReLU(),
             nn.Conv1d(64, 128, 3, padding=1),
-            nn.ReLU()
+            nn.ReLU(),
         )
+
         self.pitch_proj = nn.Linear(128, 256)
-    
+
     def forward(self, f0):
-        f0 = f0.unsqueeze(1).transpose(1, 2)  # (B, T) -> (B, 1, T)
-        out = self.f0_conv(f0)
-        out = out.mean(-1)  # Global avg pool
-        return self.pitch_proj(out)
+        # f0: (B, T)
+        x = f0.unsqueeze(1)  # (B, 1, T)
+        x = self.f0_conv(x)  # (B, 128, T)
+        x = x.transpose(1, 2)  # (B, T, 128)
+        return self.pitch_proj(x)  # (B, T, 256)
+
 
 class RVCDecoder(nn.Module):
-    def __init__(self, dim=256):
+    def __init__(self, dim=256, mel_dim=128):
         super().__init__()
-        self.content_lstm = nn.LSTM(dim, dim, 2, batch_first=True, bidirectional=True)
-        self.pitch_lstm = nn.LSTM(dim, dim//2, 1, batch_first=True)
-        self.fusion = nn.MultiheadAttention(dim*2, 8)
+
+        self.content_lstm = nn.LSTM(
+            dim,
+            dim,
+            num_layers=2,
+            batch_first=True,
+            bidirectional=True,
+        )
+
+        self.pitch_proj = nn.Linear(dim, dim * 2)
+
+        self.fusion = nn.MultiheadAttention(
+            embed_dim=dim * 2,
+            num_heads=8,
+            batch_first=True,
+        )
+
         self.output_proj = nn.Sequential(
-            nn.Linear(dim*2, dim),
+            nn.Linear(dim * 2, dim),
             nn.ReLU(),
-            nn.Linear(dim, 1024)  # Mel output
+            nn.Linear(dim, mel_dim),
         )
-    
+
     def forward(self, content, pitch):
-        content_out, _ = self.content_lstm(content)
-        pitch_out, _ = self.pitch_lstm(pitch)
-        pitch_out = pitch_out.repeat(1, content_out.size(1), 1)
-        
-        fused, _ = self.fusion(content_out, pitch_out, content_out)
-        return self.output_proj(fused)
+        # content: (B, T, 256)
+        # pitch:   (B, T, 256)
+
+        content_out, _ = self.content_lstm(content)  # (B, T, 512)
+        pitch_out = self.pitch_proj(pitch)           # (B, T, 512)
+
+        fused, _ = self.fusion(
+            query=content_out,
+            key=pitch_out,
+            value=content_out,
+        )
+
+        return self.output_proj(fused)  # (B, T, 128)
+
 
 class RVCv2(nn.Module):
     def __init__(self):
         super().__init__()
-        self.hubert = HubertEncoder()
+        self.hubert = HubertEncoder(input_dim=128)
         self.pitch = PitchEncoder()
-        self.decoder = RVCDecoder()
-    
+        self.decoder = RVCDecoder(dim=256, mel_dim=128)
+
     def forward(self, mel, f0):
+        # mel: (B, T, 128)
+        # f0:  (B, T)
         content = self.hubert(mel)
         pitch_feat = self.pitch(f0)
         return self.decoder(content, pitch_feat)
 
-# ---- Advanced Audio Dataset ----
+
+# ---- Dataset ----
 class RVCv2Dataset(Dataset):
     def __init__(self, dataset_dir, sample_rate=40000, duration=10):
         self.files = list(Path(dataset_dir).glob("*.wav"))
+
+        if len(self.files) == 0:
+            raise ValueError(f"No .wav files found in {dataset_dir}")
+
         self.sample_rate = sample_rate
         self.duration = duration
         self.n_samples = int(sample_rate * duration)
-        
+
+        self.mel_transform = torchaudio.transforms.MelSpectrogram(
+            sample_rate=self.sample_rate,
+            n_mels=128,
+            n_fft=2048,
+            hop_length=512,
+        )
+
     def __len__(self):
         return len(self.files)
-    
+
     def __getitem__(self, idx):
         waveform, sr = torchaudio.load(self.files[idx])
-        
+
+        # Convert stereo/multi-channel to mono
+        if waveform.shape[0] > 1:
+            waveform = waveform.mean(dim=0, keepdim=True)
+
         # Resample
         if sr != self.sample_rate:
             resampler = torchaudio.transforms.Resample(sr, self.sample_rate)
             waveform = resampler(waveform)
-        
-        # Trim/pad
+
+        # Trim/pad audio
         if waveform.shape[1] > self.n_samples:
             waveform = waveform[:, :self.n_samples]
         else:
-            waveform = torch.nn.functional.pad(waveform, (0, self.n_samples - waveform.shape[1]))
-        
-        # Mel spectrogram (target)
-        mel_transform = torchaudio.transforms.MelSpectrogram(
-            sample_rate=self.sample_rate, n_mels=128, n_fft=2048, hop_length=512
-        )
-        mel = mel_transform(waveform).squeeze(0)
+            pad_amount = self.n_samples - waveform.shape[1]
+            waveform = torch.nn.functional.pad(waveform, (0, pad_amount))
+
+        # Mel spectrogram: (1, 128, T) -> (128, T)
+        mel = self.mel_transform(waveform).squeeze(0)
         mel = torch.log(mel + 1e-9)
-        
-        # Dummy F0 (real impl needs crepe/dio)
-        f0 = torch.ones(mel.shape[0]) * 200.0  # Placeholder
-        f0 = torch.tensor(f0).float()
-        
+
+        # Convert to (T, 128)
+        mel = mel.transpose(0, 1)
+
+        # Dummy F0 placeholder, one value per time frame
+        f0 = torch.ones(mel.shape[0], dtype=torch.float32) * 200.0
+
         return mel, f0, waveform
 
+
 def collate_fn(batch):
     mels, f0s, waves = zip(*batch)
-    mels = pad_sequence(mels, batch_first=True, padding_value=0.0)
-    f0s = pad_sequence(f0s.unsqueeze(1), batch_first=True, padding_value=0.0).squeeze(1)
+
+    # mels are list of tensors shaped (T, 128)
+    mels = pad_sequence(
+        mels,
+        batch_first=True,
+        padding_value=0.0,
+    )
+
+    # f0s are list of tensors shaped (T,)
+    f0s = pad_sequence(
+        f0s,
+        batch_first=True,
+        padding_value=0.0,
+    )
+
     return mels, f0s, waves
 
+
 # ---- Training Loop ----
-def train_rvc_v2(model_name, dataset_dir, sample_rate=40000, epochs=200, batch_size=8, lr=2e-4):
+def train_rvc_v2(
+    model_name,
+    dataset_dir,
+    sample_rate=40000,
+    epochs=200,
+    batch_size=8,
+    lr=2e-4,
+):
     print(f"🚀 RVC v2 Training Started: {model_name}")
-    print(f"📂 Dataset: {dataset_dir} ({len(os.listdir(dataset_dir))} files)")
-    
+    print(f"📂 Dataset: {dataset_dir}")
+
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"🛠️ Device: {device}")
-    
-    # Data
+
     dataset = RVCv2Dataset(dataset_dir, sample_rate)
-    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
-    
-    # Model
+
+    print(f"🎧 Files found: {len(dataset)}")
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=collate_fn,
+        num_workers=0,
+    )
+
     model = RVCv2().to(device)
-    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)
-    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+
+    optimizer = optim.AdamW(
+        model.parameters(),
+        lr=lr,
+        weight_decay=1e-5,
+    )
+
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(
+        optimizer,
+        T_max=epochs,
+    )
+
     criterion = nn.MSELoss()
-    
+
     os.makedirs("weights", exist_ok=True)
-    
-    best_loss = float('inf')
+
+    best_loss = float("inf")
+
     for epoch in range(epochs):
         model.train()
-        total_loss = 0
-        
+        total_loss = 0.0
+
         for batch_idx, (mel, f0, _) in enumerate(dataloader):
-            mel, f0 = mel.to(device), f0.to(device)
-            
+            mel = mel.to(device)
+            f0 = f0.to(device)
+
             optimizer.zero_grad()
+
             output = model(mel, f0)
-            loss = criterion(output, mel)  # Reconstruction
+
+            loss = criterion(output, mel)
             loss.backward()
-            
-            # Gradient clipping
+
             torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+
             optimizer.step()
-            
+
             total_loss += loss.item()
-        
+
         scheduler.step()
+
         avg_loss = total_loss / len(dataloader)
-        
+
         if avg_loss < best_loss:
             best_loss = avg_loss
             torch.save(model.state_dict(), f"weights/{model_name}.pth")
-        
+
         if epoch % 10 == 0:
-            print(f"Epoch {epoch}/{epochs} | Loss: {avg_loss:.4f} | LR: {scheduler.get_last_lr()[0]:.2e}")
-    
-    print(f"✅ RVC v2 Training Complete! Best model: weights/{model_name}.pth")
+            print(
+                f"Epoch {epoch}/{epochs} | "
+                f"Loss: {avg_loss:.4f} | "
+                f"Best: {best_loss:.4f} | "
+                f"LR: {scheduler.get_last_lr()[0]:.2e}"
+            )
+
+    print(f"✅ Training Complete! Best model: weights/{model_name}.pth")
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="RVC v2 Training")
-    parser.add_argument("--model_name", required=True, help="Model name (e.g., zeynep_rvc)")
-    parser.add_argument("--dataset", required=True, help="Path to dataset folder")
+
+    parser.add_argument("--model_name", required=True)
+    parser.add_argument("--dataset", required=True)
     parser.add_argument("--sample_rate", type=int, default=40000)
     parser.add_argument("--epochs", type=int, default=200)
     parser.add_argument("--batch_size", type=int, default=8)
-    
+
     args = parser.parse_args()
-    train_rvc_v2(args.model_name, args.dataset, args.sample_rate, args.epochs, args.batch_size)
+
+    train_rvc_v2(
+        model_name=args.model_name,
+        dataset_dir=args.dataset,
+        sample_rate=args.sample_rate,
+        epochs=args.epochs,
+        batch_size=args.batch_size,
+    )