Upload 5 files

README.md

@@ -44,6 +44,15 @@ The first generation will be slower due to torch.compile, then speed will increa
 
 The model was trained on vocals but not lyrics. Vocals will not have recognizable words.
 
+## LoRA Training
+
+- Prepare folder of .mp3 files
+- Run python train_lora_encode_latents.py --audio-dir=/path/to/your/mp3s --output-dir=latents to save the latents
+- Run python train_lora.py --latents_dir=latents to train the LoRA. You may need to adjust learning rate, steps or batch size depending on your dataset etc.
+- Run python merge_lora.py --lora-checkpoint=lora_step1000.safetensors --output-checkpoint=merged.safetensors to merge the LoRA checkpoint into the base model for inference
+- Run python gradio_app.py --checkpoint=merged.safetensors to run the merged checkpoint for inference
+- Test inference with tag "soundtrack"; Lora training uses this tag. Additional tags may work.
+
 ## Credits
 
 This project builds upon the following open-source projects:

gradio_app.py

@@ -3,6 +3,7 @@ from pathlib import Path
 from typing import List, Tuple
 import uuid
 import json
+import argparse
 import gradio as gr
 import torch
 import torchaudio
@@ -83,7 +84,7 @@ rf_sampler: RF | None = None
 device: torch.device | None = None
 _available_tags: List[str] | None = None
 
-def load_resources() -> List[str]:
+def load_resources(checkpoint_path) -> List[str]:
 
     torch.set_float32_matmul_precision('high')
 
@@ -107,7 +108,6 @@ def load_resources() -> List[str]:
         max_tags=8,
     ).to(device)
 
-    checkpoint_path = "checkpoints/checkpoint_461260.safetensors"
     print(f"Loading checkpoint: {checkpoint_path}")
 
     state_dict = load_file(checkpoint_path, device=str(device))
@@ -139,7 +139,6 @@ def generate_audio(
     sample_steps: int,
 ) -> Tuple[Tuple[int, object], str]:
 
-    load_resources()
     assert model is not None and vae is not None and rf_sampler is not None and device is not None
 
     if not tags:
@@ -178,8 +177,8 @@ def generate_audio(
 
     return (sr, audio_numpy), str(output_path)
 
-def build_interface() -> gr.Blocks:
-    available_tags = load_resources()
+def build_interface(checkpoint_path) -> gr.Blocks:
+    available_tags = load_resources(checkpoint_path)
 
     # Define preset tag combinations
     presets = [
@@ -259,7 +258,16 @@ def build_interface() -> gr.Blocks:
 
     return demo
 
-demo = build_interface()
-
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="LocalSong Gradio Interface")
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default="checkpoints/checkpoint_461260.safetensors",
+        help="Path to the model checkpoint"
+    )
+    args = parser.parse_args()
+
+    demo = build_interface(args.checkpoint)
+
     demo.launch()

merge_lora.py

@@ -0,0 +1,196 @@
+import torch
+import torch.nn as nn
+import argparse
+from safetensors.torch import load_file, save_file
+from model import LocalSongModel
+from pathlib import Path
+
+class LoRALinear(nn.Module):
+    def __init__(self, original_linear: nn.Linear, rank: int = 8, alpha: float = 16.0):
+        super().__init__()
+        self.original_linear = original_linear
+        self.rank = rank
+        self.alpha = alpha
+        self.scaling = alpha / rank
+
+        self.lora_A = nn.Parameter(torch.zeros(original_linear.in_features, rank))
+        self.lora_B = nn.Parameter(torch.zeros(rank, original_linear.out_features))
+
+        nn.init.kaiming_uniform_(self.lora_A, a=5**0.5)
+        nn.init.zeros_(self.lora_B)
+
+        self.original_linear.weight.requires_grad = False
+        if self.original_linear.bias is not None:
+            self.original_linear.bias.requires_grad = False
+
+    def forward(self, x):
+        result = self.original_linear(x)
+        lora_out = (x @ self.lora_A @ self.lora_B) * self.scaling
+        return result + lora_out
+
+def inject_lora(model, rank=8, alpha=16.0, target_modules=['qkv', 'proj', 'w1', 'w2', 'w3', 'q_proj', 'kv_proj'], device=None):
+    if device is None:
+        device = next(model.parameters()).device
+
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear):
+            if any(target in name for target in target_modules):
+                *parent_path, attr_name = name.split('.')
+                parent = model
+                for p in parent_path:
+                    parent = getattr(parent, p)
+
+                lora_layer = LoRALinear(module, rank=rank, alpha=alpha)
+                lora_layer.lora_A.data = lora_layer.lora_A.data.to(device)
+                lora_layer.lora_B.data = lora_layer.lora_B.data.to(device)
+                setattr(parent, attr_name, lora_layer)
+
+    return model
+
+def load_lora_weights(model, lora_path, device):
+    print(f"Loading LoRA from {lora_path}")
+    lora_state_dict = load_file(lora_path, device=str(device))
+
+    loaded_count = 0
+    for name, module in model.named_modules():
+        if isinstance(module, LoRALinear):
+            lora_a_key = f"{name}.lora_A"
+            lora_b_key = f"{name}.lora_B"
+            if lora_a_key in lora_state_dict and lora_b_key in lora_state_dict:
+                module.lora_A.data = lora_state_dict[lora_a_key].to(device)
+                module.lora_B.data = lora_state_dict[lora_b_key].to(device)
+                loaded_count += 2
+
+    print(f"Loaded {loaded_count} LoRA parameters")
+
+def merge_lora_into_model(model):
+    """
+    Merge LoRA weights into the base model weights.
+    For each LoRALinear layer: W_merged = W_original + (lora_A @ lora_B) * scaling
+    """
+    print("\nMerging LoRA weights into base model...")
+    merged_count = 0
+
+    for name, module in model.named_modules():
+        if isinstance(module, LoRALinear):
+            lora_delta = (module.lora_A @ module.lora_B) * module.scaling
+
+            with torch.no_grad():
+                module.original_linear.weight.data += lora_delta.T 
+
+            merged_count += 1
+
+    print(f"Merged {merged_count} LoRA layers into base weights")
+
+def extract_base_weights(model):
+    """
+    Extract the merged weights from LoRALinear modules back into a regular state dict.
+    """
+    print("\nExtracting merged weights...")
+    new_state_dict = {}
+
+    for name, module in model.named_modules():
+        if isinstance(module, LoRALinear):
+            original_name_weight = f"{name}.weight"
+            original_name_bias = f"{name}.bias"
+
+            new_state_dict[original_name_weight] = module.original_linear.weight.data
+            if module.original_linear.bias is not None:
+                new_state_dict[original_name_bias] = module.original_linear.bias.data
+
+    # Copy over all non-LoRA parameters
+    for name, param in model.named_parameters():
+        if 'lora_A' not in name and 'lora_B' not in name and 'original_linear' not in name:
+            new_state_dict[name] = param.data
+
+    print(f"Extracted {len(new_state_dict)} parameters")
+    return new_state_dict
+
+def main():
+    parser = argparse.ArgumentParser(description="Merge LoRA weights into a base model checkpoint")
+    parser.add_argument(
+        "--base-checkpoint",
+        type=str,
+        default="checkpoints/checkpoint_461260.safetensors",
+        help="Path to the base model checkpoint"
+    )
+    parser.add_argument(
+        "--lora-checkpoint",
+        type=str,
+        default="lora.safetensors",
+        help="Path to the LoRA checkpoint"
+    )
+    parser.add_argument(
+        "--output-checkpoint",
+        type=str,
+        default="checkpoints/checkpoint_461260_merged_lora.safetensors",
+        help="Path to save the merged checkpoint"
+    )
+    args = parser.parse_args()
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+
+    # Configuration
+    base_checkpoint = args.base_checkpoint
+    lora_checkpoint = args.lora_checkpoint
+    output_checkpoint = args.output_checkpoint
+
+    lora_rank = 16
+    lora_alpha = 16.0
+
+    print(f"\nBase checkpoint: {base_checkpoint}")
+    print(f"LoRA checkpoint: {lora_checkpoint}")
+    print(f"Output checkpoint: {output_checkpoint}")
+    print(f"LoRA rank: {lora_rank}, alpha: {lora_alpha}")
+
+    # Load base model
+    print("\nLoading base model...")
+    model = LocalSongModel(
+        in_channels=8,
+        num_groups=16,
+        hidden_size=1024,
+        decoder_hidden_size=2048,
+        num_blocks=36,
+        patch_size=(16, 1),
+        num_classes=2304,
+        max_tags=8,
+    ).to(device)
+
+    state_dict = load_file(base_checkpoint, device=str(device))
+    model.load_state_dict(state_dict, strict=True)
+    print("Base model loaded")
+
+    print("\nInjecting LoRA layers...")
+    model = inject_lora(model, rank=lora_rank, alpha=lora_alpha, device=device)
+
+    load_lora_weights(model, lora_checkpoint, device)
+
+    merge_lora_into_model(model)
+
+    merged_state_dict = extract_base_weights(model)
+
+    print(f"\nSaving merged checkpoint to {output_checkpoint}...")
+    save_file(merged_state_dict, output_checkpoint)
+    print("✓ Merged checkpoint saved successfully!")
+
+    print("\nVerifying merged checkpoint...")
+    test_model = LocalSongModel(
+        in_channels=8,
+        num_groups=16,
+        hidden_size=1024,
+        decoder_hidden_size=2048,
+        num_blocks=36,
+        patch_size=(16, 1),
+        num_classes=2304,
+        max_tags=8,
+    ).to(device)
+
+    merged_loaded = load_file(output_checkpoint, device=str(device))
+    test_model.load_state_dict(merged_loaded, strict=True)
+    print("✓ Merged checkpoint verified successfully!")
+
+    print(f"\nDone! You can now use '{output_checkpoint}' as a standalone checkpoint without needing LoRA.")
+
+if __name__ == '__main__':
+    main()

train_lora.py

@@ -0,0 +1,273 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from pathlib import Path
+import argparse
+from tqdm import tqdm
+from safetensors.torch import save_file, load_file
+from collections import deque
+from model import LocalSongModel
+
+HARDCODED_TAGS = [1908]
+torch.set_float32_matmul_precision('high')
+
+class LoRALinear(nn.Module):
+    def __init__(self, original_linear: nn.Linear, rank: int = 8, alpha: float = 16.0):
+        super().__init__()
+        self.original_linear = original_linear
+        self.rank = rank
+        self.alpha = alpha
+        self.scaling = alpha / rank
+
+        self.lora_A = nn.Parameter(torch.zeros(original_linear.in_features, rank))
+        self.lora_B = nn.Parameter(torch.zeros(rank, original_linear.out_features))
+
+        nn.init.kaiming_uniform_(self.lora_A, a=5**0.5)
+        nn.init.zeros_(self.lora_B)
+
+        self.original_linear.weight.requires_grad = False
+        if self.original_linear.bias is not None:
+            self.original_linear.bias.requires_grad = False
+
+    def forward(self, x):
+        result = self.original_linear(x)
+        lora_out = (x @ self.lora_A @ self.lora_B) * self.scaling
+        return result + lora_out
+
+def inject_lora(model: LocalSongModel, rank: int = 8, alpha: float = 16.0, target_modules=['qkv', 'proj', 'w1', 'w2', 'w3', 'q_proj', 'kv_proj'], device=None):
+    """Inject LoRA layers into the model."""
+
+    lora_modules = []
+
+    if device is None:
+        device = next(model.parameters()).device
+
+    for name, module in model.named_modules():
+
+        if isinstance(module, nn.Linear):
+
+            if any(target in name for target in target_modules):
+
+                *parent_path, attr_name = name.split('.')
+                parent = model
+                for p in parent_path:
+                    parent = getattr(parent, p)
+
+                lora_layer = LoRALinear(module, rank=rank, alpha=alpha)
+
+                lora_layer.lora_A.data = lora_layer.lora_A.data.to(device)
+                lora_layer.lora_B.data = lora_layer.lora_B.data.to(device)
+                setattr(parent, attr_name, lora_layer)
+                lora_modules.append(name)
+
+    print(f"Injected LoRA into {len(lora_modules)} layers:")
+    for name in lora_modules[:5]:
+        print(f"  - {name}")
+    if len(lora_modules) > 5:
+        print(f"  ... and {len(lora_modules) - 5} more")
+
+    return model
+
+def get_lora_parameters(model):
+    """Extract only LoRA parameters for optimization."""
+    lora_params = []
+    for module in model.modules():
+        if isinstance(module, LoRALinear):
+            lora_params.extend([module.lora_A, module.lora_B])
+    return lora_params
+
+def save_lora_weights(model, output_path):
+    """Save LoRA weights to a safetensors file."""
+    lora_state_dict = {}
+
+    for name, module in model.named_modules():
+        if isinstance(module, LoRALinear):
+            lora_state_dict[f"{name}.lora_A"] = module.lora_A
+            lora_state_dict[f"{name}.lora_B"] = module.lora_B
+
+    save_file(lora_state_dict, output_path)
+    print(f"Saved {len(lora_state_dict)} LoRA parameters to {output_path}")
+
+class LatentDataset(Dataset):
+    """Dataset for pre-encoded latents."""
+
+    def __init__(self, latents_dir: str):
+        self.latents_dir = Path(latents_dir)
+
+        self.latent_files = sorted(list(self.latents_dir.glob("*.pt")))
+
+        if len(self.latent_files) == 0:
+            raise ValueError(f"No .pt files found in {latents_dir}")
+
+        print(f"Found {len(self.latent_files)} latent files")
+
+    def __len__(self):
+        return len(self.latent_files)
+
+    def __getitem__(self, idx):
+        latent = torch.load(self.latent_files[idx])
+
+        if latent.ndim == 3:
+            latent = latent.unsqueeze(0)
+
+        return latent
+
+class RectifiedFlow:
+    """Simplified rectified flow matching."""
+
+    def __init__(self, model):
+        self.model = model
+
+    def forward(self, x, cond):
+        """Compute flow matching loss."""
+        b = x.size(0)
+
+        nt = torch.randn((b,), device=x.device)
+        t = torch.sigmoid(nt)
+
+        texp = t.view([b, *([1] * len(x.shape[1:]))])
+        z1 = torch.randn_like(x)
+        zt = (1 - texp) * x + texp * z1
+
+        vtheta = self.model(zt, t, cond)
+
+        target = z1 - x
+        loss = ((vtheta - target) ** 2).mean()
+
+        return loss
+
+def collate_fn(batch, subsection_length=1024):
+    """Custom collate function to sample random subsections."""
+    sampled_latents = []
+
+    for latent in batch:
+        if latent.ndim == 3:
+            latent = latent.unsqueeze(0)
+
+        _, _, _, width = latent.shape
+
+        if width < subsection_length:
+            # Pad if too short
+            pad_amount = subsection_length - width
+            latent = torch.nn.functional.pad(latent, (0, pad_amount), mode='constant', value=0)
+        else:
+            # Randomly sample subsection
+            max_start = width - subsection_length
+            start_idx = torch.randint(0, max_start + 1, (1,)).item()
+            latent = latent[:, :, :, start_idx:start_idx + subsection_length]
+
+        sampled_latents.append(latent.squeeze(0))
+
+    batch_latents = torch.stack(sampled_latents)
+
+    batch_tags = [HARDCODED_TAGS] * len(batch)
+
+    return batch_latents, batch_tags
+
+def main():
+    parser = argparse.ArgumentParser(description='LoRA training for LocalSong model with embedding training')
+
+    parser.add_argument('--latents_dir', type=str, required=True,
+                        help='Directory containing VAE-encoded latents (.pt files)')
+    
+    parser.add_argument('--checkpoint', type=str, default='checkpoints/checkpoint_461260.safetensors',
+                        help='Path to base model checkpoint')
+    parser.add_argument('--lora_rank', type=int, default=16,
+                        help='LoRA rank')
+    parser.add_argument('--lora_alpha', type=float, default=16,
+                        help='LoRA alpha (scaling factor)')
+    parser.add_argument('--batch_size', type=int, default=16,
+                        help='Batch size')
+    parser.add_argument('--lr', type=float, default=2e-4,
+                        help='Learning rate')
+    parser.add_argument('--steps', type=int, default=1500,
+                        help='Number of training steps')
+    parser.add_argument('--subsection_length', type=int, default=512,
+                        help='Latent subsection length')
+    parser.add_argument('--output', type=str, default='lora.safetensors',
+                        help='Output path for LoRA weights')
+    parser.add_argument('--save_every', type=int, default=500,
+                        help='Save checkpoint every N steps')
+
+    args = parser.parse_args()
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+
+    print(f"Using hardcoded tags: {HARDCODED_TAGS}")
+
+    print(f"Loading base model from {args.checkpoint}")
+    model = LocalSongModel(
+        in_channels=8,
+        num_groups=16,
+        hidden_size=1024,
+        decoder_hidden_size=2048,
+        num_blocks=36,
+        patch_size=(16, 1),
+        num_classes=2304,
+        max_tags=8,
+    )
+
+    print(f"Loading checkpoint from {args.checkpoint}")
+    state_dict = load_file(args.checkpoint)
+    model.load_state_dict(state_dict, strict=True)
+    print("Base model loaded")
+
+    model = model.to(device)
+    model = inject_lora(model, rank=args.lora_rank, alpha=args.lora_alpha, device=device)
+
+    model.train()
+
+    lora_params = get_lora_parameters(model)
+    optimizer = optim.Adam(lora_params, lr=args.lr)
+    print(f"Training {len(lora_params)} LoRA parameters")
+
+    dataset = LatentDataset(args.latents_dir)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=0,
+        collate_fn=lambda batch: collate_fn(batch, args.subsection_length)
+    )
+
+    rf = RectifiedFlow(model)
+
+    print("\nStarting training...")
+    step = 0
+    pbar = tqdm(total=args.steps, desc="Training")
+
+    loss_history = deque(maxlen=50)
+
+    while step < args.steps:
+        for batch_latents, batch_tags in dataloader:
+            batch_latents = batch_latents.to(device)
+
+            optimizer.zero_grad()
+            loss = rf.forward(batch_latents, batch_tags)
+
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(lora_params, 1.0)
+            optimizer.step()
+
+            # Track loss and compute average
+            loss_history.append(loss.item())
+            avg_loss = sum(loss_history) / len(loss_history)
+
+            pbar.set_postfix({'loss': f'{avg_loss:.4f}'})
+            pbar.update(1)
+            step += 1
+
+            if step % args.save_every == 0:
+                save_path = args.output.replace('.safetensors', f'_step{step}.safetensors')
+                save_lora_weights(model, save_path)
+
+            if step >= args.steps:
+                break
+
+    save_lora_weights(model, args.output)
+    print(f"\nTraining complete! LoRA weights saved to {args.output}")
+
+if __name__ == '__main__':
+    main()

train_lora_encode_latents.py

@@ -0,0 +1,103 @@
+import torch
+import torchaudio
+from pathlib import Path
+import argparse
+from tqdm import tqdm
+from acestep.music_dcae.music_dcae_pipeline import MusicDCAE
+
+class AudioVAE:
+    def __init__(self, device: torch.device):
+        self.model = MusicDCAE().to(device)
+        self.model.eval()
+        self.device = device
+        self.latent_mean = torch.tensor(
+            [0.1207, -0.0186, -0.0947, -0.3779, 0.5956, 0.3422, 0.1796, -0.0526],
+            device=device,
+        ).view(1, -1, 1, 1)
+        self.latent_std = torch.tensor(
+            [0.4638, 0.3154, 0.6244, 1.5078, 0.4696, 0.4633, 0.5614, 0.2707],
+            device=device,
+        ).view(1, -1, 1, 1)
+
+    def encode(self, audio):
+
+        with torch.no_grad():
+            audio_lengths = torch.tensor([audio.shape[2]] * audio.shape[0]).to(self.device)
+            latents, _ = self.model.encode(audio, audio_lengths, sr=48000)
+            latents = (latents - self.latent_mean) / self.latent_std
+        return latents
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        with torch.no_grad():
+            latents = latents * self.latent_std + self.latent_mean
+            _, audio_list = self.model.decode(latents, sr=48000)
+            audio_batch = torch.stack(audio_list).to(self.device)
+        return audio_batch
+
+def load_audio(audio_path, target_sr=48000):
+    """Load and preprocess audio file."""
+    audio, sr = torchaudio.load(audio_path)
+
+    if audio.shape[0] == 1:
+        audio = audio.repeat(2, 1)
+    elif audio.shape[0] > 2:
+        audio = audio[:2]
+
+    if sr != target_sr:
+        resampler = torchaudio.transforms.Resample(sr, target_sr)
+        audio = resampler(audio)
+
+    return audio
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Encode audio files to VAE latents')
+
+    parser.add_argument('--audio-dir', type=str, required=True,
+                        help='Directory containing audio files')
+    parser.add_argument('--output-dir', type=str, default="latents",
+                        help='Directory to save encoded latents')
+
+    args = parser.parse_args()
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    audio_dir = Path(args.audio_dir)
+    audio_extensions = ['*.mp3', '*.wav', '*.flac', '*.ogg', '*.m4a']
+    audio_files = []
+    for ext in audio_extensions:
+        audio_files.extend(list(audio_dir.glob(ext)))
+    audio_files = sorted(audio_files)
+
+    if len(audio_files) == 0:
+        raise ValueError(f"No audio files found in {args.audio_dir}")
+
+    print(f"Found {len(audio_files)} audio files")
+
+    vae = AudioVAE(device)
+    print("VAE loaded")
+
+    # Encode each audio file
+    print("\nEncoding audio files...")
+    for audio_path in tqdm(audio_files, desc="Encoding"):
+        try:
+            audio = load_audio(audio_path)
+            audio = audio.unsqueeze(0).to(device)
+            latents = vae.encode(audio)
+            latents = latents.squeeze(0)
+
+            output_path = output_dir / f"{audio_path.stem}.pt"
+            torch.save(latents.cpu(), output_path)
+
+        except Exception as e:
+            print(f"\nError encoding {audio_path.name}: {e}")
+            continue
+
+    print(f"\nEncoding complete! Saved {len(list(output_dir.glob('*.pt')))} latent files to {output_dir}")
+
+if __name__ == '__main__':
+    main()