Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,51 @@
+Third-Party License Attribution for Audio Processing Module
+===========================================================
+
+This directory contains code derived from multiple open-source projects. 
+The following sections detail the licenses and attributions for third-party code.
+
+## XCodec Repository
+The code in this directory is derived from:
+https://github.com/zhenye234/xcodec
+
+## Individual File Attributions
+
+### Quantization Module (quantization/)
+- Several files contain code derived from Meta Platforms, Inc. and the vector-quantize-pytorch repository
+- Individual files contain their own license headers where applicable
+- The vector-quantize-pytorch portions are licensed under the MIT License
+
+## License Terms
+
+### MIT License (for applicable portions)
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+## Attribution Requirements
+When using this code, please ensure proper attribution to:
+1. The original xcodec repository: https://github.com/zhenye234/xcodec
+2. Any other repositories mentioned in individual file headers
+3. This derivative work and its modifications
+
+## Disclaimer
+This directory contains modified versions of the original code. Please refer to
+the original repositories for the canonical implementations and their specific
+license terms.
+
+For any questions about licensing or attribution, please check the individual
+file headers and the original source repositories. 

boson_codeit.py

@@ -0,0 +1,651 @@
+# #!/usr/bin/env python3
+# """
+# Audio Processing Script for Boson Codes
+# Processes audio files in parallel using Higgs Audio Tokenizer
+# and saves encoded representations as .pt files.
+# """
+
+# import os
+# import sys
+# import json
+# import torch
+# import librosa
+# import numpy as np
+# import warnings
+# import argparse
+# from pathlib import Path
+# from multiprocessing import Pool
+# from tqdm import tqdm
+
+# from datasets import load_from_disk
+# from higgs_audio_tokenizer import HiggsAudioTokenizer
+
+# # Suppress PyTorch FutureWarnings
+# warnings.filterwarnings("ignore", category=FutureWarning)
+
+# # Global configuration
+# DEFAULT_OUTPUT_DIR = "/home/ubuntu/boson_codes"
+# DEFAULT_NUM_CORES = 48
+# DEFAULT_SAMPLE_RATE = 44100
+# DEFAULT_DATASET_PATH = "/home/ubuntu/ttsar/Layla/src_bpe_2/data"
+
+# # Model paths
+# CONFIG_PATH = "/home/ubuntu/.cache/huggingface/hub/models--bosonai--higgs-audio-v2-tokenizer/snapshots/9d4988fbd4ad07b4cac3a5fa462741a41810dbec/config.json"
+# MODEL_PATH = "/home/ubuntu/.cache/huggingface/hub/models--bosonai--higgs-audio-v2-tokenizer/snapshots/9d4988fbd4ad07b4cac3a5fa462741a41810dbec/model.pth"
+
+# # Global model variable (initialized in each worker)
+# model = None
+
+
+# def init_worker():
+#     """Initialize model once per worker process."""
+#     global model
+#     device = 'cpu'
+    
+#     # Load config
+#     with open(CONFIG_PATH, 'r') as f:
+#         config = json.load(f)
+    
+#     # Initialize model
+#     model = HiggsAudioTokenizer(
+#         **config,
+#         device=device,
+#     )
+    
+#     # Load weights
+#     parameter_dict = torch.load(MODEL_PATH, map_location=device)
+#     _ = model.load_state_dict(parameter_dict, strict=False)
+#     model = model.to(device)
+#     _ = model.eval()
+    
+#     print(f"Model loaded in worker {os.getpid()}")
+
+
+# def process_audio_file(args):
+#     """Process a single audio file using pre-loaded model."""
+#     filename, output_dir, sample_rate = args
+    
+#     try:
+#         # Output filename - same name, just change extension to .pt
+#         base_name = Path(filename).stem
+#         output_path = os.path.join(output_dir, f"{base_name}.pt")
+        
+#         # Skip if exists (double-check in case of race conditions)
+#         if os.path.exists(output_path):
+#             return ("skipped", filename)
+        
+#         # Load and process audio
+#         wav, sr = librosa.load(filename, sr=sample_rate)
+#         wav = torch.from_numpy(wav).unsqueeze(0).float().to('cpu')
+        
+#         # Encode using the pre-loaded model
+#         with torch.no_grad():
+#             encoded = model._xcodec_encode(wav.unsqueeze(0))
+        
+#         # Save codes only
+#         torch.save(encoded.audio_codes, output_path)
+        
+#         return ("success", filename)
+        
+#     except Exception as e:
+#         return ("error", filename, str(e))
+
+
+# def load_dataset(dataset_path):
+#     """Load and prepare the dataset."""
+#     print(f"Loading dataset from: {dataset_path}")
+#     ds = load_from_disk(dataset_path)
+#     print(f"Dataset info: {ds}")
+    
+#     # Remove unnecessary columns
+#     columns_to_remove = ['spk', 'duration', 'codes', 'input_ids', 'attention_mask']
+#     existing_columns = [col for col in columns_to_remove if col in ds.column_names]
+#     if existing_columns:
+#         ds = ds.remove_columns(existing_columns)
+#         print(f"Removed columns: {existing_columns}")
+    
+#     # Convert to pandas DataFrame
+#     df = ds.to_pandas()
+#     print(f"Loaded {len(df)} files from dataset")
+#     return df
+
+
+# def main(args):
+#     """Main processing function."""
+#     # Change to audio processing directory
+#     os.chdir("/home/ubuntu/ttsar/boson_audio_codec/audio_processing")
+#     print(f"Working directory: {os.getcwd()}")
+    
+#     # Create output directory
+#     os.makedirs(args.output_dir, exist_ok=True)
+#     print(f"Output directory: {args.output_dir}")
+    
+#     # Check if model files exist
+#     if not os.path.exists(CONFIG_PATH):
+#         print(f"Error: Config file not found at {CONFIG_PATH}")
+#         sys.exit(1)
+#     if not os.path.exists(MODEL_PATH):
+#         print(f"Error: Model file not found at {MODEL_PATH}")
+#         sys.exit(1)
+    
+#     # Load dataset
+#     df = load_dataset(args.dataset_path)
+    
+#     # Get filenames from dataframe
+#     all_filenames = df['filename'].tolist()
+    
+#     # Pre-filter to exclude already processed files
+#     filenames_to_process = []
+#     already_processed = []
+    
+#     print(f"\nChecking for already processed files...")
+#     for filename in all_filenames:
+#         base_name = Path(filename).stem
+#         output_path = os.path.join(args.output_dir, f"{base_name}.pt")
+#         if os.path.exists(output_path):
+#             already_processed.append(filename)
+#         else:
+#             filenames_to_process.append(filename)
+    
+#     print(f"\nTotal files: {len(all_filenames)}")
+#     print(f"Already processed: {len(already_processed)}")
+#     print(f"To process: {len(filenames_to_process)}")
+    
+#     if len(filenames_to_process) == 0:
+#         print("\nAll files have already been processed!")
+#         return
+    
+#     print(f"\nProcessing {len(filenames_to_process)} files using {args.num_cores} cores...")
+#     print(f"Sample rate: {args.sample_rate} Hz")
+    
+#     # Prepare arguments for multiprocessing
+#     process_args = [(filename, args.output_dir, args.sample_rate) 
+#                     for filename in filenames_to_process]
+    
+#     # Process in parallel with model reuse
+#     with Pool(processes=args.num_cores, initializer=init_worker) as pool:
+#         results = list(tqdm(
+#             pool.imap(process_audio_file, process_args, chunksize=args.chunksize),
+#             total=len(filenames_to_process),
+#             desc="Processing audio files"
+#         ))
+    
+#     # Count results
+#     processed = sum(1 for r in results if r[0] == "success")
+#     skipped = sum(1 for r in results if r[0] == "skipped")
+#     errors = sum(1 for r in results if r[0] == "error")
+    
+#     print(f"\nProcessing complete!")
+#     print(f"  Successfully processed: {processed}")
+#     print(f"  Previously processed: {len(already_processed)}")
+#     print(f"  Skipped (race condition): {skipped}")
+#     print(f"  Errors: {errors}")
+    
+#     # Show errors if any
+#     if errors > 0:
+#         print("\nErrors encountered:")
+#         error_log_path = os.path.join(args.output_dir, "processing_errors.log")
+#         with open(error_log_path, 'w') as f:
+#             for r in results:
+#                 if r[0] == "error":
+#                     error_msg = f"{r[1]}: {r[2]}"
+#                     print(f"  {error_msg}")
+#                     f.write(error_msg + "\n")
+#         print(f"\nError log saved to: {error_log_path}")
+    
+#     # Show summary of all processed files
+#     total_processed_files = len(list(Path(args.output_dir).glob("*.pt")))
+#     print(f"\nTotal .pt files in {args.output_dir}: {total_processed_files}")
+
+
+# if __name__ == "__main__":
+#     parser = argparse.ArgumentParser(
+#         description="Process audio files using Higgs Audio Tokenizer and save as .pt files"
+#     )
+    
+#     parser.add_argument(
+#         "--dataset-path", 
+#         type=str, 
+#         default=DEFAULT_DATASET_PATH,
+#         help=f"Path to the dataset (default: {DEFAULT_DATASET_PATH})"
+#     )
+    
+#     parser.add_argument(
+#         "--output-dir", 
+#         type=str, 
+#         default=DEFAULT_OUTPUT_DIR,
+#         help=f"Output directory for .pt files (default: {DEFAULT_OUTPUT_DIR})"
+#     )
+    
+#     parser.add_argument(
+#         "--num-cores", 
+#         type=int, 
+#         default=DEFAULT_NUM_CORES,
+#         help=f"Number of CPU cores to use (default: {DEFAULT_NUM_CORES})"
+#     )
+    
+#     parser.add_argument(
+#         "--sample-rate", 
+#         type=int, 
+#         default=DEFAULT_SAMPLE_RATE,
+#         help=f"Sample rate for audio processing (default: {DEFAULT_SAMPLE_RATE})"
+#     )
+    
+#     parser.add_argument(
+#         "--chunksize", 
+#         type=int, 
+#         default=1,
+#         help="Chunksize for multiprocessing pool (default: 1)"
+#     )
+    
+#     args = parser.parse_args()
+    
+#     # Run main processing
+#     try:
+#         main(args)
+#     except KeyboardInterrupt:
+#         print("\n\nProcessing interrupted by user")
+#         sys.exit(1)
+#     except Exception as e:
+#         print(f"\n\nError: {e}")
+#         sys.exit(1)
+
+#!/usr/bin/env python3
+"""
+GPU Batch Processing Script for Boson Codes with Dataset Loading
+"""
+
+import os
+import sys
+import json
+import torch
+import torch.nn.functional as F
+import librosa
+import numpy as np
+from pathlib import Path
+from tqdm import tqdm
+import warnings
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+
+# from boson_multimodal.audio_processing.higgs_audio_tokenizer import load_higgs_audio_tokenizer
+# model = load_higgs_audio_tokenizer("bosonai/higgs-audio-v2-tokenizer")
+import librosa
+import torch
+import torch.nn.functional as F
+import numpy as np
+import json
+import torch
+
+from higgs_audio_tokenizer import HiggsAudioTokenizer
+# model = load_higgs_audio_tokenizer("bosonai/higgs-audio-v2-tokenizer")
+
+import torch
+import torch.nn as nn
+import warnings
+
+# Suppress warnings
+warnings.filterwarnings('ignore')
+
+def remove_weight_norms_from_model(model):
+    for module in model.modules():
+        try:
+            remove_weight_norm(module)
+        except:
+            continue
+    return model
+
+
+class EncodedResult:
+    def __init__(self, audio_codes):
+        self.audio_codes = audio_codes
+
+def encode_batch(model, x_batch):
+    """
+    Encodes a batch of audio tensors using the HiggsAudioTokenizer model.
+    Args:
+        model: The loaded HiggsAudioTokenizer model.
+        x_batch: A tensor of shape [B, 1, T]
+    """
+    # Acoustic and Semantic Feature Extraction
+    e_semantic_input = model.get_regress_target(x_batch).detach()
+    e_semantic = model.encoder_semantic(e_semantic_input.transpose(1, 2))
+    e_acoustic = model.encoder(x_batch)
+
+    # This block contains the fix for batch processing
+    if e_acoustic.shape[2] != e_semantic.shape[2]:
+        pad_size = 160 * model.semantic_downsample_factor
+        
+        # 1. Remove channel dim, preserving batch dim -> [B, T]
+        x_slice = x_batch[:, 0, :]
+        
+        # 2. Pad the tensor
+        x_padded = F.pad(x_slice, (pad_size, pad_size))
+        
+        # 3. Re-add channel dim before passing to encoder -> [B, 1, T_padded]
+        e_acoustic = model.encoder(x_padded.unsqueeze(1))
+
+    # Ensure dimensions match before concatenating
+    min_len = min(e_acoustic.shape[2], e_semantic.shape[2])
+    e_acoustic = e_acoustic[:, :, :min_len]
+    e_semantic = e_semantic[:, :, :min_len]
+
+    # Remainder of the original encoding logic
+    e = torch.cat([e_acoustic, e_semantic], dim=1)
+    e = model.fc_prior(e.transpose(1, 2))
+
+    if model.quantizer_type == "RVQ":
+        e = e.transpose(1, 2)
+        _, codes, _, _ = model.quantizer(e, model.frame_rate, None)
+        codes = codes.permute(1, 0, 2)
+    else: # RFSQ
+        quantized, codes = model.quantizer(e)
+        codes = codes.permute(0, 2, 1)
+
+    return EncodedResult(audio_codes=codes)
+
+
+def fix_all_inference_issues(model):
+    """
+    Comprehensive fix for all potential inference issues
+    """
+    device = next(model.parameters()).device
+    
+    # 1. Force everything to eval mode
+    model.eval()
+    with torch.no_grad():
+        for module in model.modules():
+            if isinstance(module, nn.Module):
+                module.eval()
+                if hasattr(module, 'training'):
+                    module.training = False
+    
+    # 2. Fix semantic model specifically
+    if hasattr(model, 'semantic_model'):
+        print("Fixing semantic model...")
+        
+        # Move to correct device
+        model.semantic_model = model.semantic_model.to(device)
+        model.semantic_model.eval()
+        
+        # Disable ALL gradient checkpointing
+        def disable_gradient_checkpointing(module):
+            if hasattr(module, 'gradient_checkpointing'):
+                module.gradient_checkpointing = False
+            if hasattr(module, 'gradient_checkpointing_disable'):
+                try:
+                    module.gradient_checkpointing_disable()
+                except:
+                    pass
+            for child in module.children():
+                disable_gradient_checkpointing(child)
+        
+        disable_gradient_checkpointing(model.semantic_model)
+        
+        # For HuBERT specifically
+        if hasattr(model.semantic_model, 'encoder'):
+            model.semantic_model.encoder.gradient_checkpointing = False
+            if hasattr(model.semantic_model.encoder, 'layers'):
+                for layer in model.semantic_model.encoder.layers:
+                    if hasattr(layer, 'gradient_checkpointing'):
+                        layer.gradient_checkpointing = False
+    
+    # 3. Set all dropout to eval mode
+    def set_dropout_eval(module):
+        if isinstance(module, nn.Dropout):
+            module.eval()
+            module.training = False
+        for child in module.children():
+            set_dropout_eval(child)
+    
+    set_dropout_eval(model)
+    
+    # 4. Clear any cached computations
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None
+    
+    return model
+
+def inference_pipeline(checkpoint_path, config_path, device='cuda'):
+    """
+    Complete pipeline for inference with your trained model
+    """
+    # Load config
+    print("Loading config...")
+    with open(config_path, 'r') as f:
+        config = json.load(f)
+    
+    # Create model
+    print("Creating model...")
+    model = HiggsAudioTokenizer(
+        n_filters=config['n_filters'],
+        D=config['D'],
+        target_bandwidths=config['target_bandwidths'],
+        ratios=config['ratios'],
+        sample_rate=config['sample_rate'],
+        bins=config['bins'],
+        n_q=config['n_q'],
+        codebook_dim=config.get('codebook_dim', None),
+        semantic_techer=config['semantic_techer'],
+        device=device
+    ).to(device)
+    
+    # Load checkpoint
+    print("Loading checkpoint...")
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    
+    if 'model_state_dict' in checkpoint:
+        state_dict = checkpoint['model_state_dict']
+    else:
+        state_dict = checkpoint
+    
+    # Remove 'module.' prefix if present (from DDP)
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        if k.startswith('module.'):
+            new_state_dict[k[7:]] = v
+        else:
+            new_state_dict[k] = v
+    
+    model.load_state_dict(new_state_dict, strict=False)
+    
+    # Fix all inference issues
+    print("Fixing inference issues...")
+    model = fix_all_inference_issues(model)
+
+    
+    return model
+
+
+
+# # Add paths
+# sys.path.insert(0, "/home/ubuntu/AP-BWE")
+
+# Suppress warnings
+warnings.filterwarnings("ignore")
+
+# Configuration
+OUTPUT_DIR = "/home/ubuntu/data_boson_44.1khz"
+BATCH_SIZE = 32
+SAMPLE_RATE = 44100
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+DATASET_PATH = "/home/ubuntu/ttsar/Layla/src_bpe_2/Qanary_data"
+
+# # Model paths
+# CONFIG_PATH = "/home/ubuntu/.cache/huggingface/hub/models--bosonai--higgs-audio-v2-tokenizer/snapshots/9d4988fbd4ad07b4cac3a5fa462741a41810dbec/config.json"
+# MODEL_PATH = "/home/ubuntu/.cache/huggingface/hub/models--bosonai--higgs-audio-v2-tokenizer/snapshots/9d4988fbd4ad07b4cac3a5fa462741a41810dbec/model.pth"
+
+# --- Setup ---
+print(f"Using device: {DEVICE}")
+
+# Change to working directory
+os.chdir("/home/ubuntu/ttsar/boson_audio_codec/audio_processing")
+
+# Load dataset
+from datasets import load_from_disk
+
+
+print(f"Loading dataset from: {DATASET_PATH}")
+ds = load_from_disk(DATASET_PATH)
+print(f"Dataset info: {ds}")
+
+# Remove unnecessary columns
+columns_to_remove = ['spk', 'duration', 'codes', 'input_ids', 'attention_mask']
+existing_columns = [col for col in columns_to_remove if col in ds.column_names]
+if existing_columns:
+    ds = ds.remove_columns(existing_columns)
+
+df = ds.to_pandas()
+print(f"Loaded {len(df)} files from dataset")
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+print(f"Output directory '{OUTPUT_DIR}' is ready.")
+
+# --- Filter already processed ---
+print("Checking for already processed files...")
+
+def get_output_path(audio_path):
+    base_name = Path(audio_path).stem
+    return os.path.join(OUTPUT_DIR, f"{base_name}.pt")
+
+# Filter
+original_count = len(df)
+df['output_exists'] = df['filename'].apply(lambda x: os.path.exists(get_output_path(x)))
+df_filtered = df[~df['output_exists']].copy()
+skipped_count = original_count - len(df_filtered)
+
+print(f"Found {skipped_count} already processed files. Skipping them.")
+print(f"Processing {len(df_filtered)} remaining files.")
+
+if len(df_filtered) == 0:
+    print("All files have already been processed!")
+    exit()
+
+# --- Load Model ---
+print("Loading Higgs Audio Tokenizer model...")
+
+from transformers import HubertModel
+from higgs_audio_tokenizer import HiggsAudioTokenizer
+
+# Load config
+# with open(CONFIG_PATH, 'r') as f:
+#     config = json.load(f)
+
+# # Initialize model
+# model = HiggsAudioTokenizer(
+#     **config,
+#     device=DEVICE,
+# )
+
+# Load weights
+# parameter_dict = torch.load(MODEL_PATH, map_location=DEVICE)
+# _ = model.load_state_dict(parameter_dict, strict=False)
+# model = model.to(DEVICE)
+# _ = model.eval()
+
+
+checkpoint_path = '/home/ubuntu/ttsar/boson_audio_codec/audio_processing/outputs_CQT/checkpoints/step_99000.pth'
+config_path = '/home/ubuntu/ttsar/boson_audio_codec/audio_processing/config copy.json'
+device = 'cuda'
+model = inference_pipeline(checkpoint_path, config_path, device)
+_ = model.eval()
+
+model = remove_weight_norms_from_model(model)
+
+print(f"Model loaded on {DEVICE}")
+
+# Get hop length
+hop_length = model.hop_length
+print(f"Encoder hop length: {hop_length}")
+
+# --- Batch Processing ---
+print(f"\nStarting batch processing with batch size {BATCH_SIZE}...")
+
+# Process in batches
+filenames = df_filtered['filename'].tolist()
+total_processed = 0
+total_errors = 0
+
+with torch.no_grad():
+    for batch_start in tqdm(range(0, len(filenames), BATCH_SIZE), desc="Processing batches"):
+        batch_end = min(batch_start + BATCH_SIZE, len(filenames))
+        batch_filenames = filenames[batch_start:batch_end]
+        
+        batch_audio = []
+        batch_lengths = []
+        batch_outputs = []
+        
+        # Load batch
+        for filename in batch_filenames:
+            output_path = get_output_path(filename)
+            
+            # Skip if exists (race condition check)
+            if os.path.exists(output_path):
+                continue
+            
+            try:
+                # Load audio
+                wav, _ = librosa.load(filename, sr=SAMPLE_RATE)
+                wav_tensor = torch.from_numpy(wav).float()
+                
+                batch_audio.append(wav_tensor)
+                batch_lengths.append(len(wav))
+                batch_outputs.append(output_path)
+                
+            except Exception as e:
+                print(f"\nError loading {filename}: {e}")
+                total_errors += 1
+                continue
+        
+        if not batch_audio:
+            continue
+        
+        # Pad batch to same length
+        max_len = max(len(x) for x in batch_audio)
+        padded_batch = []
+        
+        for audio in batch_audio:
+            pad_len = max_len - len(audio)
+            if pad_len > 0:
+                audio = F.pad(audio, (0, pad_len), mode='constant', value=0)
+            # Don't add extra dimensions here, just collect the padded audio
+            padded_batch.append(audio)
+        
+        # Convert list to tensor and add channel dimension
+        # Stack along batch dimension to get [B, T]
+        batch_tensor = torch.stack(padded_batch, dim=0)  # [B, T]
+        # Add channel dimension
+        batch_tensor = batch_tensor.unsqueeze(1)  # [B, 1, T]
+        batch_tensor = batch_tensor.to(DEVICE)
+        
+        # Encode batch
+        try:
+            encoded = encode_batch(model, batch_tensor)
+            codes = encoded.audio_codes  # [B, n_codebooks, T_compressed]
+            
+            # Save each item
+            for idx, (output_path, orig_len) in enumerate(zip(batch_outputs, batch_lengths)):
+                # Calculate true code length
+                true_code_len = int(np.ceil(orig_len / hop_length))
+                
+                # Extract non-padded codes
+                item_codes = codes[idx, :, :true_code_len].cpu()
+                
+                # Save
+                torch.save(item_codes, output_path)
+                total_processed += 1
+                
+        except Exception as e:
+            print(f"\nError encoding batch: {e}")
+            total_errors += len(batch_outputs)
+
+print("\n" + "="*50)
+print("PROCESSING COMPLETE!")
+print("="*50)
+print(f"Successfully processed: {total_processed} files")
+print(f"Previously processed: {skipped_count} files")
+print(f"Errors encountered: {total_errors} files")
+print(f"Output directory: {OUTPUT_DIR}")
+
+# Final count
+final_count = len(list(Path(OUTPUT_DIR).glob("*.pt")))
+print(f"Total .pt files in output: {final_count}")

descriptaudiocodec/dac/model/base.py

@@ -0,0 +1,286 @@
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+import torch
+import tqdm
+from audiotools import AudioSignal
+from torch import nn
+
+SUPPORTED_VERSIONS = ["1.0.0"]
+
+
+@dataclass
+class DACFile:
+    codes: torch.Tensor
+
+    # Metadata
+    chunk_length: int
+    original_length: int
+    input_db: float
+    channels: int
+    sample_rate: int
+    padding: bool
+    dac_version: str
+
+    def save(self, path):
+        artifacts = {
+            "codes": self.codes.numpy().astype(np.uint16),
+            "metadata": {
+                "input_db": self.input_db.numpy().astype(np.float32),
+                "original_length": self.original_length,
+                "sample_rate": self.sample_rate,
+                "chunk_length": self.chunk_length,
+                "channels": self.channels,
+                "padding": self.padding,
+                "dac_version": SUPPORTED_VERSIONS[-1],
+            },
+        }
+        path = Path(path).with_suffix(".dac")
+        with open(path, "wb") as f:
+            np.save(f, artifacts)
+        return path
+
+    @classmethod
+    def load(cls, path):
+        artifacts = np.load(path, allow_pickle=True)[()]
+        codes = torch.from_numpy(artifacts["codes"].astype(int))
+        if artifacts["metadata"].get("dac_version", None) not in SUPPORTED_VERSIONS:
+            raise RuntimeError(f"Given file {path} can't be loaded with this version of descript-audio-codec.")
+        return cls(codes=codes, **artifacts["metadata"])
+
+
+class CodecMixin:
+    @property
+    def padding(self):
+        if not hasattr(self, "_padding"):
+            self._padding = True
+        return self._padding
+
+    @padding.setter
+    def padding(self, value):
+        assert isinstance(value, bool)
+
+        layers = [l for l in self.modules() if isinstance(l, (nn.Conv1d, nn.ConvTranspose1d))]
+
+        for layer in layers:
+            if value:
+                if hasattr(layer, "original_padding"):
+                    layer.padding = layer.original_padding
+            else:
+                layer.original_padding = layer.padding
+                layer.padding = tuple(0 for _ in range(len(layer.padding)))
+
+        self._padding = value
+
+    def get_delay(self):
+        # Any number works here, delay is invariant to input length
+        l_out = self.get_output_length(0)
+        L = l_out
+
+        layers = []
+        for layer in self.modules():
+            if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
+                layers.append(layer)
+
+        for layer in reversed(layers):
+            d = layer.dilation[0]
+            k = layer.kernel_size[0]
+            s = layer.stride[0]
+
+            if isinstance(layer, nn.ConvTranspose1d):
+                L = ((L - d * (k - 1) - 1) / s) + 1
+            elif isinstance(layer, nn.Conv1d):
+                L = (L - 1) * s + d * (k - 1) + 1
+
+            L = math.ceil(L)
+
+        l_in = L
+
+        return (l_in - l_out) // 2
+
+    def get_output_length(self, input_length):
+        L = input_length
+        # Calculate output length
+        for layer in self.modules():
+            if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
+                d = layer.dilation[0]
+                k = layer.kernel_size[0]
+                s = layer.stride[0]
+
+                if isinstance(layer, nn.Conv1d):
+                    L = ((L - d * (k - 1) - 1) / s) + 1
+                elif isinstance(layer, nn.ConvTranspose1d):
+                    L = (L - 1) * s + d * (k - 1) + 1
+
+                L = math.floor(L)
+        return L
+
+    @torch.no_grad()
+    def compress(
+        self,
+        audio_path_or_signal: Union[str, Path, AudioSignal],
+        win_duration: float = 1.0,
+        verbose: bool = False,
+        normalize_db: float = -16,
+        n_quantizers: int = None,
+    ) -> DACFile:
+        """Processes an audio signal from a file or AudioSignal object into
+        discrete codes. This function processes the signal in short windows,
+        using constant GPU memory.
+
+        Parameters
+        ----------
+        audio_path_or_signal : Union[str, Path, AudioSignal]
+            audio signal to reconstruct
+        win_duration : float, optional
+            window duration in seconds, by default 5.0
+        verbose : bool, optional
+            by default False
+        normalize_db : float, optional
+            normalize db, by default -16
+
+        Returns
+        -------
+        DACFile
+            Object containing compressed codes and metadata
+            required for decompression
+        """
+        audio_signal = audio_path_or_signal
+        if isinstance(audio_signal, (str, Path)):
+            audio_signal = AudioSignal.load_from_file_with_ffmpeg(str(audio_signal))
+
+        self.eval()
+        original_padding = self.padding
+        original_device = audio_signal.device
+
+        audio_signal = audio_signal.clone()
+        original_sr = audio_signal.sample_rate
+
+        resample_fn = audio_signal.resample
+        loudness_fn = audio_signal.loudness
+
+        # If audio is > 10 minutes long, use the ffmpeg versions
+        if audio_signal.signal_duration >= 10 * 60 * 60:
+            resample_fn = audio_signal.ffmpeg_resample
+            loudness_fn = audio_signal.ffmpeg_loudness
+
+        original_length = audio_signal.signal_length
+        resample_fn(self.sample_rate)
+        input_db = loudness_fn()
+
+        if normalize_db is not None:
+            audio_signal.normalize(normalize_db)
+        audio_signal.ensure_max_of_audio()
+
+        nb, nac, nt = audio_signal.audio_data.shape
+        audio_signal.audio_data = audio_signal.audio_data.reshape(nb * nac, 1, nt)
+        win_duration = audio_signal.signal_duration if win_duration is None else win_duration
+
+        if audio_signal.signal_duration <= win_duration:
+            # Unchunked compression (used if signal length < win duration)
+            self.padding = True
+            n_samples = nt
+            hop = nt
+        else:
+            # Chunked inference
+            self.padding = False
+            # Zero-pad signal on either side by the delay
+            audio_signal.zero_pad(self.delay, self.delay)
+            n_samples = int(win_duration * self.sample_rate)
+            # Round n_samples to nearest hop length multiple
+            n_samples = int(math.ceil(n_samples / self.hop_length) * self.hop_length)
+            hop = self.get_output_length(n_samples)
+
+        codes = []
+        range_fn = range if not verbose else tqdm.trange
+
+        for i in range_fn(0, nt, hop):
+            x = audio_signal[..., i : i + n_samples]
+            x = x.zero_pad(0, max(0, n_samples - x.shape[-1]))
+
+            audio_data = x.audio_data.to(self.device)
+            audio_data = self.preprocess(audio_data, self.sample_rate)
+            _, c, _, _, _ = self.encode(audio_data, n_quantizers)
+            codes.append(c.to(original_device))
+            chunk_length = c.shape[-1]
+
+        codes = torch.cat(codes, dim=-1)
+
+        dac_file = DACFile(
+            codes=codes,
+            chunk_length=chunk_length,
+            original_length=original_length,
+            input_db=input_db,
+            channels=nac,
+            sample_rate=original_sr,
+            padding=self.padding,
+            dac_version=SUPPORTED_VERSIONS[-1],
+        )
+
+        if n_quantizers is not None:
+            codes = codes[:, :n_quantizers, :]
+
+        self.padding = original_padding
+        return dac_file
+
+    @torch.no_grad()
+    def decompress(
+        self,
+        obj: Union[str, Path, DACFile],
+        verbose: bool = False,
+    ) -> AudioSignal:
+        """Reconstruct audio from a given .dac file
+
+        Parameters
+        ----------
+        obj : Union[str, Path, DACFile]
+            .dac file location or corresponding DACFile object.
+        verbose : bool, optional
+            Prints progress if True, by default False
+
+        Returns
+        -------
+        AudioSignal
+            Object with the reconstructed audio
+        """
+        self.eval()
+        if isinstance(obj, (str, Path)):
+            obj = DACFile.load(obj)
+
+        original_padding = self.padding
+        self.padding = obj.padding
+
+        range_fn = range if not verbose else tqdm.trange
+        codes = obj.codes
+        original_device = codes.device
+        chunk_length = obj.chunk_length
+        recons = []
+
+        for i in range_fn(0, codes.shape[-1], chunk_length):
+            c = codes[..., i : i + chunk_length].to(self.device)
+            z = self.quantizer.from_codes(c)[0]
+            r = self.decode(z)
+            recons.append(r.to(original_device))
+
+        recons = torch.cat(recons, dim=-1)
+        recons = AudioSignal(recons, self.sample_rate)
+
+        resample_fn = recons.resample
+        loudness_fn = recons.loudness
+
+        # If audio is > 10 minutes long, use the ffmpeg versions
+        if recons.signal_duration >= 10 * 60 * 60:
+            resample_fn = recons.ffmpeg_resample
+            loudness_fn = recons.ffmpeg_loudness
+
+        recons.normalize(obj.input_db)
+        resample_fn(obj.sample_rate)
+        recons = recons[..., : obj.original_length]
+        loudness_fn()
+        recons.audio_data = recons.audio_data.reshape(-1, obj.channels, obj.original_length)
+
+        self.padding = original_padding
+        return recons

descriptaudiocodec/dac/model/dac.py

@@ -0,0 +1,365 @@
+import math
+from typing import List
+from typing import Union
+
+import numpy as np
+import torch
+from audiotools import AudioSignal
+from audiotools.ml import BaseModel
+from torch import nn
+
+from .base import CodecMixin
+from dac.nn.layers import Snake1d
+from dac.nn.layers import WNConv1d
+from dac.nn.layers import WNConvTranspose1d
+from dac.nn.quantize import ResidualVectorQuantize
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int = 16, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            ResidualUnit(dim // 2, dilation=1),
+            ResidualUnit(dim // 2, dilation=3),
+            ResidualUnit(dim // 2, dilation=9),
+            Snake1d(dim // 2),
+            WNConv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 64,
+        strides: list = [2, 4, 8, 8],
+        d_latent: int = 256,
+    ):
+        super().__init__()
+        # Create first convolution
+        self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
+
+        # Create EncoderBlocks that double channels as they downsample by `stride`
+        for stride in strides:
+            d_model *= 2
+            self.block += [EncoderBlock(d_model, stride=stride)]
+
+        # Create last convolution
+        self.block += [
+            Snake1d(d_model),
+            WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
+        ]
+
+        # Wrap black into nn.Sequential
+        self.block = nn.Sequential(*self.block)
+        self.enc_dim = d_model
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1, out_pad=0):
+        super().__init__()
+        self.block = nn.Sequential(
+            Snake1d(input_dim),
+            WNConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,  # out_pad,
+            ),
+            ResidualUnit(output_dim, dilation=1),
+            ResidualUnit(output_dim, dilation=3),
+            ResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        input_channel,
+        channels,
+        rates,
+        d_out: int = 1,
+    ):
+        super().__init__()
+
+        # Add first conv layer
+        layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3)]
+
+        # Add upsampling + MRF blocks
+        for i, stride in enumerate(rates):
+            input_dim = channels // 2**i
+            output_dim = channels // 2 ** (i + 1)
+            if i == 1:
+                out_pad = 1
+            else:
+                out_pad = 0
+            layers += [DecoderBlock(input_dim, output_dim, stride, out_pad)]
+
+        # Add final conv layer
+        layers += [
+            Snake1d(output_dim),
+            WNConv1d(output_dim, d_out, kernel_size=7, padding=3),
+            # nn.Tanh(),
+        ]
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class DAC(BaseModel, CodecMixin):
+    def __init__(
+        self,
+        encoder_dim: int = 64,
+        encoder_rates: List[int] = [2, 4, 8, 8],
+        latent_dim: int = None,
+        decoder_dim: int = 1536,
+        decoder_rates: List[int] = [8, 8, 4, 2],
+        n_codebooks: int = 9,
+        codebook_size: int = 1024,
+        codebook_dim: Union[int, list] = 8,
+        quantizer_dropout: bool = False,
+        sample_rate: int = 44100,
+    ):
+        super().__init__()
+
+        self.encoder_dim = encoder_dim
+        self.encoder_rates = encoder_rates
+        self.decoder_dim = decoder_dim
+        self.decoder_rates = decoder_rates
+        self.sample_rate = sample_rate
+
+        if latent_dim is None:
+            latent_dim = encoder_dim * (2 ** len(encoder_rates))
+
+        self.latent_dim = latent_dim
+
+        self.hop_length = np.prod(encoder_rates)
+        self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
+
+        self.n_codebooks = n_codebooks
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.quantizer = ResidualVectorQuantize(
+            input_dim=latent_dim,
+            n_codebooks=n_codebooks,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            quantizer_dropout=quantizer_dropout,
+        )
+
+        self.decoder = Decoder(
+            latent_dim,
+            decoder_dim,
+            decoder_rates,
+        )
+        self.sample_rate = sample_rate
+        self.apply(init_weights)
+
+        self.delay = self.get_delay()
+
+    def preprocess(self, audio_data, sample_rate):
+        if sample_rate is None:
+            sample_rate = self.sample_rate
+        assert sample_rate == self.sample_rate
+
+        length = audio_data.shape[-1]
+        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
+        audio_data = nn.functional.pad(audio_data, (0, right_pad))
+
+        return audio_data
+
+    def encode(
+        self,
+        audio_data: torch.Tensor,
+        n_quantizers: int = None,
+    ):
+        """Encode given audio data and return quantized latent codes
+
+        Parameters
+        ----------
+        audio_data : Tensor[B x 1 x T]
+            Audio data to encode
+        n_quantizers : int, optional
+            Number of quantizers to use, by default None
+            If None, all quantizers are used.
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "z" : Tensor[B x D x T]
+                Quantized continuous representation of input
+            "codes" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+            "latents" : Tensor[B x N*D x T]
+                Projected latents (continuous representation of input before quantization)
+            "vq/commitment_loss" : Tensor[1]
+                Commitment loss to train encoder to predict vectors closer to codebook
+                entries
+            "vq/codebook_loss" : Tensor[1]
+                Codebook loss to update the codebook
+            "length" : int
+                Number of samples in input audio
+        """
+        z = self.encoder(audio_data)
+        z, codes, latents, commitment_loss, codebook_loss = self.quantizer(z, n_quantizers)
+        return z, codes, latents, commitment_loss, codebook_loss
+
+    def decode(self, z: torch.Tensor):
+        """Decode given latent codes and return audio data
+
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+            Quantized continuous representation of input
+        length : int, optional
+            Number of samples in output audio, by default None
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "audio" : Tensor[B x 1 x length]
+                Decoded audio data.
+        """
+        return self.decoder(z)
+
+    def forward(
+        self,
+        audio_data: torch.Tensor,
+        sample_rate: int = None,
+        n_quantizers: int = None,
+    ):
+        """Model forward pass
+
+        Parameters
+        ----------
+        audio_data : Tensor[B x 1 x T]
+            Audio data to encode
+        sample_rate : int, optional
+            Sample rate of audio data in Hz, by default None
+            If None, defaults to `self.sample_rate`
+        n_quantizers : int, optional
+            Number of quantizers to use, by default None.
+            If None, all quantizers are used.
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "z" : Tensor[B x D x T]
+                Quantized continuous representation of input
+            "codes" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+            "latents" : Tensor[B x N*D x T]
+                Projected latents (continuous representation of input before quantization)
+            "vq/commitment_loss" : Tensor[1]
+                Commitment loss to train encoder to predict vectors closer to codebook
+                entries
+            "vq/codebook_loss" : Tensor[1]
+                Codebook loss to update the codebook
+            "length" : int
+                Number of samples in input audio
+            "audio" : Tensor[B x 1 x length]
+                Decoded audio data.
+        """
+        length = audio_data.shape[-1]
+        audio_data = self.preprocess(audio_data, sample_rate)
+        z, codes, latents, commitment_loss, codebook_loss = self.encode(audio_data, n_quantizers)
+
+        x = self.decode(z)
+        return {
+            "audio": x[..., :length],
+            "z": z,
+            "codes": codes,
+            "latents": latents,
+            "vq/commitment_loss": commitment_loss,
+            "vq/codebook_loss": codebook_loss,
+        }
+
+
+if __name__ == "__main__":
+    import numpy as np
+    from functools import partial
+
+    model = DAC().to("cpu")
+
+    for n, m in model.named_modules():
+        o = m.extra_repr()
+        p = sum([np.prod(p.size()) for p in m.parameters()])
+        fn = lambda o, p: o + f" {p / 1e6:<.3f}M params."
+        setattr(m, "extra_repr", partial(fn, o=o, p=p))
+    print(model)
+    print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()]))
+
+    length = 88200 * 2
+    x = torch.randn(1, 1, length).to(model.device)
+    x.requires_grad_(True)
+    x.retain_grad()
+
+    # Make a forward pass
+    out = model(x)["audio"]
+    print("Input shape:", x.shape)
+    print("Output shape:", out.shape)
+
+    # Create gradient variable
+    grad = torch.zeros_like(out)
+    grad[:, :, grad.shape[-1] // 2] = 1
+
+    # Make a backward pass
+    out.backward(grad)
+
+    # Check non-zero values
+    gradmap = x.grad.squeeze(0)
+    gradmap = (gradmap != 0).sum(0)  # sum across features
+    rf = (gradmap != 0).sum()
+
+    print(f"Receptive field: {rf.item()}")
+
+    x = AudioSignal(torch.randn(1, 1, 44100 * 60), 44100)
+    model.decompress(model.compress(x, verbose=True), verbose=True)

descriptaudiocodec/dac/nn/layers.py

@@ -0,0 +1,33 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)

descriptaudiocodec/dac/nn/quantize.py

@@ -0,0 +1,251 @@
+from typing import Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+from dac.nn.layers import WNConv1d
+
+
+class VectorQuantize(nn.Module):
+    """
+    Implementation of VQ similar to Karpathy's repo:
+    https://github.com/karpathy/deep-vector-quantization
+    Additionally uses following tricks from Improved VQGAN
+    (https://arxiv.org/pdf/2110.04627.pdf):
+        1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
+            for improved codebook usage
+        2. l2-normalized codes: Converts euclidean distance to cosine similarity which
+            improves training stability
+    """
+
+    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
+        self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    def forward(self, z):
+        """Quantized the input tensor using a fixed codebook and returns
+        the corresponding codebook vectors
+
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+
+        Returns
+        -------
+        Tensor[B x D x T]
+            Quantized continuous representation of input
+        Tensor[1]
+            Commitment loss to train encoder to predict vectors closer to codebook
+            entries
+        Tensor[1]
+            Codebook loss to update the codebook
+        Tensor[B x T]
+            Codebook indices (quantized discrete representation of input)
+        Tensor[B x D x T]
+            Projected latents (continuous representation of input before quantization)
+        """
+
+        # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
+        z_e = self.in_proj(z)  # z_e : (B x D x T)
+        z_q, indices = self.decode_latents(z_e)
+
+        commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
+        codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
+
+        z_q = z_e + (z_q - z_e).detach()  # noop in forward pass, straight-through gradient estimator in backward pass
+
+        z_q = self.out_proj(z_q)
+
+        return z_q, commitment_loss, codebook_loss, indices, z_e
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight  # codebook: (N x D)
+
+        # L2 normalize encodings and codebook (ViT-VQGAN)
+        encodings = F.normalize(encodings)
+        codebook = F.normalize(codebook)
+
+        # Compute euclidean distance with codebook
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+        return z_q, indices
+
+
+class ResidualVectorQuantize(nn.Module):
+    """
+    Introduced in SoundStream: An end2end neural audio codec
+    https://arxiv.org/abs/2107.03312
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 512,
+        n_codebooks: int = 9,
+        codebook_size: int = 1024,
+        codebook_dim: Union[int, list] = 8,
+        quantizer_dropout: float = 0.0,
+    ):
+        super().__init__()
+        if isinstance(codebook_dim, int):
+            codebook_dim = [codebook_dim for _ in range(n_codebooks)]
+
+        self.n_codebooks = n_codebooks
+        self.codebook_dim = codebook_dim
+        self.codebook_size = codebook_size
+
+        self.quantizers = nn.ModuleList(
+            [VectorQuantize(input_dim, codebook_size, codebook_dim[i]) for i in range(n_codebooks)]
+        )
+        self.quantizer_dropout = quantizer_dropout
+
+    def forward(self, z, n_quantizers: int = None):
+        """Quantized the input tensor using a fixed set of `n` codebooks and returns
+        the corresponding codebook vectors
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+        n_quantizers : int, optional
+            No. of quantizers to use
+            (n_quantizers < self.n_codebooks ex: for quantizer dropout)
+            Note: if `self.quantizer_dropout` is True, this argument is ignored
+                when in training mode, and a random number of quantizers is used.
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+
+            "z" : Tensor[B x D x T]
+                Quantized continuous representation of input
+            "codes" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+            "latents" : Tensor[B x N*D x T]
+                Projected latents (continuous representation of input before quantization)
+            "vq/commitment_loss" : Tensor[1]
+                Commitment loss to train encoder to predict vectors closer to codebook
+                entries
+            "vq/codebook_loss" : Tensor[1]
+                Codebook loss to update the codebook
+        """
+        z_q = 0
+        residual = z
+        commitment_loss = 0
+        codebook_loss = 0
+
+        codebook_indices = []
+        latents = []
+
+        if n_quantizers is None:
+            n_quantizers = self.n_codebooks
+        if self.training:
+            n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1
+            dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))
+            n_dropout = int(z.shape[0] * self.quantizer_dropout)
+            n_quantizers[:n_dropout] = dropout[:n_dropout]
+            n_quantizers = n_quantizers.to(z.device)
+
+        for i, quantizer in enumerate(self.quantizers):
+            if self.training is False and i >= n_quantizers:
+                break
+
+            z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(residual)
+
+            # Create mask to apply quantizer dropout
+            mask = torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
+            z_q = z_q + z_q_i * mask[:, None, None]
+            residual = residual - z_q_i
+
+            # Sum losses
+            commitment_loss += (commitment_loss_i * mask).mean()
+            codebook_loss += (codebook_loss_i * mask).mean()
+
+            codebook_indices.append(indices_i)
+            latents.append(z_e_i)
+
+        codes = torch.stack(codebook_indices, dim=1)
+        latents = torch.cat(latents, dim=1)
+
+        return z_q, codes, latents, commitment_loss, codebook_loss
+
+    def from_codes(self, codes: torch.Tensor):
+        """Given the quantized codes, reconstruct the continuous representation
+        Parameters
+        ----------
+        codes : Tensor[B x N x T]
+            Quantized discrete representation of input
+        Returns
+        -------
+        Tensor[B x D x T]
+            Quantized continuous representation of input
+        """
+        z_q = 0.0
+        z_p = []
+        n_codebooks = codes.shape[1]
+        for i in range(n_codebooks):
+            z_p_i = self.quantizers[i].decode_code(codes[:, i, :])
+            z_p.append(z_p_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+        return z_q, torch.cat(z_p, dim=1), codes
+
+    def from_latents(self, latents: torch.Tensor):
+        """Given the unquantized latents, reconstruct the
+        continuous representation after quantization.
+
+        Parameters
+        ----------
+        latents : Tensor[B x N x T]
+            Continuous representation of input after projection
+
+        Returns
+        -------
+        Tensor[B x D x T]
+            Quantized representation of full-projected space
+        Tensor[B x D x T]
+            Quantized representation of latent space
+        """
+        z_q = 0
+        z_p = []
+        codes = []
+        dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])
+
+        n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[0]
+        for i in range(n_codebooks):
+            j, k = dims[i], dims[i + 1]
+            z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
+            z_p.append(z_p_i)
+            codes.append(codes_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+
+        return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)
+
+
+if __name__ == "__main__":
+    rvq = ResidualVectorQuantize(quantizer_dropout=True)
+    x = torch.randn(16, 512, 80)
+    y = rvq(x)
+    print(y["latents"].shape)

discriminator.py

@@ -0,0 +1,596 @@
+# import torch
+# import torch.nn as nn
+# import torch.nn.functional as F
+# from audiotools import AudioSignal
+# from audiotools import ml
+# from audiotools import STFTParams
+# from einops import rearrange
+# from torch.nn.utils import weight_norm
+
+
+# def WNConv1d(*args, **kwargs):
+#     act = kwargs.pop("act", True)
+#     conv = weight_norm(nn.Conv1d(*args, **kwargs))
+#     if not act:
+#         return conv
+#     return nn.Sequential(conv, nn.LeakyReLU(0.1))
+
+
+# def WNConv2d(*args, **kwargs):
+#     act = kwargs.pop("act", True)
+#     conv = weight_norm(nn.Conv2d(*args, **kwargs))
+#     if not act:
+#         return conv
+#     return nn.Sequential(conv, nn.LeakyReLU(0.1))
+
+
+# class MPD(nn.Module):
+#     def __init__(self, period):
+#         super().__init__()
+#         self.period = period
+#         self.convs = nn.ModuleList(
+#             [
+#                 WNConv2d(1, 32, (5, 1), (3, 1), padding=(2, 0)),
+#                 WNConv2d(32, 128, (5, 1), (3, 1), padding=(2, 0)),
+#                 WNConv2d(128, 512, (5, 1), (3, 1), padding=(2, 0)),
+#                 WNConv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0)),
+#                 WNConv2d(1024, 1024, (5, 1), 1, padding=(2, 0)),
+#             ]
+#         )
+#         self.conv_post = WNConv2d(
+#             1024, 1, kernel_size=(3, 1), padding=(1, 0), act=False
+#         )
+
+#     def pad_to_period(self, x):
+#         t = x.shape[-1]
+#         x = F.pad(x, (0, self.period - t % self.period), mode="reflect")
+#         return x
+
+#     def forward(self, x):
+#         fmap = []
+
+#         x = self.pad_to_period(x)
+#         x = rearrange(x, "b c (l p) -> b c l p", p=self.period)
+
+#         for layer in self.convs:
+#             x = layer(x)
+#             fmap.append(x)
+
+#         x = self.conv_post(x)
+#         fmap.append(x)
+
+#         return fmap
+
+
+# class MSD(nn.Module):
+#     def __init__(self, rate: int = 1, sample_rate: int = 44100):
+#         super().__init__()
+#         self.convs = nn.ModuleList(
+#             [
+#                 WNConv1d(1, 16, 15, 1, padding=7),
+#                 WNConv1d(16, 64, 41, 4, groups=4, padding=20),
+#                 WNConv1d(64, 256, 41, 4, groups=16, padding=20),
+#                 WNConv1d(256, 1024, 41, 4, groups=64, padding=20),
+#                 WNConv1d(1024, 1024, 41, 4, groups=256, padding=20),
+#                 WNConv1d(1024, 1024, 5, 1, padding=2),
+#             ]
+#         )
+#         self.conv_post = WNConv1d(1024, 1, 3, 1, padding=1, act=False)
+#         self.sample_rate = sample_rate
+#         self.rate = rate
+
+#     def forward(self, x):
+#         x = AudioSignal(x, self.sample_rate)
+#         x.resample(self.sample_rate // self.rate)
+#         x = x.audio_data
+
+#         fmap = []
+
+#         for l in self.convs:
+#             x = l(x)
+#             fmap.append(x)
+#         x = self.conv_post(x)
+#         fmap.append(x)
+
+#         return fmap
+
+
+# BANDS = [(0.0, 0.1), (0.1, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)]
+
+
+# class MRD(nn.Module):
+#     def __init__(
+#         self,
+#         window_length: int,
+#         hop_factor: float = 0.25,
+#         sample_rate: int = 44100,
+#         bands: list = BANDS,
+#     ):
+#         """Complex multi-band spectrogram discriminator.
+#         Parameters
+#         ----------
+#         window_length : int
+#             Window length of STFT.
+#         hop_factor : float, optional
+#             Hop factor of the STFT, defaults to ``0.25 * window_length``.
+#         sample_rate : int, optional
+#             Sampling rate of audio in Hz, by default 44100
+#         bands : list, optional
+#             Bands to run discriminator over.
+#         """
+#         super().__init__()
+
+#         self.window_length = window_length
+#         self.hop_factor = hop_factor
+#         self.sample_rate = sample_rate
+#         self.stft_params = STFTParams(
+#             window_length=window_length,
+#             hop_length=int(window_length * hop_factor),
+#             match_stride=True,
+#         )
+
+#         n_fft = window_length // 2 + 1
+#         bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
+#         self.bands = bands
+
+#         ch = 32
+#         convs = lambda: nn.ModuleList(
+#             [
+#                 WNConv2d(2, ch, (3, 9), (1, 1), padding=(1, 4)),
+#                 WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
+#                 WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
+#                 WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
+#                 WNConv2d(ch, ch, (3, 3), (1, 1), padding=(1, 1)),
+#             ]
+#         )
+#         self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
+#         self.conv_post = WNConv2d(ch, 1, (3, 3), (1, 1), padding=(1, 1), act=False)
+
+#     def spectrogram(self, x):
+#         x = AudioSignal(x, self.sample_rate, stft_params=self.stft_params)
+#         x = torch.view_as_real(x.stft())
+#         x = rearrange(x, "b 1 f t c -> (b 1) c t f")
+#         # Split into bands
+#         x_bands = [x[..., b[0] : b[1]] for b in self.bands]
+#         return x_bands
+
+#     def forward(self, x):
+#         x_bands = self.spectrogram(x)
+#         fmap = []
+
+#         x = []
+#         for band, stack in zip(x_bands, self.band_convs):
+#             for layer in stack:
+#                 band = layer(band)
+#                 fmap.append(band)
+#             x.append(band)
+
+#         x = torch.cat(x, dim=-1)
+#         x = self.conv_post(x)
+#         fmap.append(x)
+
+#         return fmap
+
+
+# class Discriminator(ml.BaseModel):
+#     def __init__(
+#         self,
+#         rates: list = [],
+#         periods: list = [2, 3, 5, 7, 11],
+#         fft_sizes: list = [2048, 1024, 512],
+#         sample_rate: int = 44100,
+#         bands: list = BANDS,
+#     ):
+#         """Discriminator that combines multiple discriminators.
+
+#         Parameters
+#         ----------
+#         rates : list, optional
+#             sampling rates (in Hz) to run MSD at, by default []
+#             If empty, MSD is not used.
+#         periods : list, optional
+#             periods (of samples) to run MPD at, by default [2, 3, 5, 7, 11]
+#         fft_sizes : list, optional
+#             Window sizes of the FFT to run MRD at, by default [2048, 1024, 512]
+#         sample_rate : int, optional
+#             Sampling rate of audio in Hz, by default 44100
+#         bands : list, optional
+#             Bands to run MRD at, by default `BANDS`
+#         """
+#         super().__init__()
+#         discs = []
+#         discs += [MPD(p) for p in periods]
+#         discs += [MSD(r, sample_rate=sample_rate) for r in rates]
+#         discs += [MRD(f, sample_rate=sample_rate, bands=bands) for f in fft_sizes]
+#         self.discriminators = nn.ModuleList(discs)
+
+#     def preprocess(self, y):
+#         # Remove DC offset
+#         y = y - y.mean(dim=-1, keepdims=True)
+#         # Peak normalize the volume of input audio
+#         y = 0.8 * y / (y.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
+#         return y
+
+#     def forward(self, x):
+#         x = self.preprocess(x)
+#         fmaps = [d(x) for d in self.discriminators]
+#         return fmaps
+
+
+# if __name__ == "__main__":
+#     disc = Discriminator()
+#     x = torch.zeros(1, 1, 44100)
+#     results = disc(x)
+#     for i, result in enumerate(results):
+#         print(f"disc{i}")
+#         for i, r in enumerate(result):
+#             print(r.shape, r.mean(), r.min(), r.max())
+#         print()
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from audiotools import AudioSignal, STFTParams
+from audiotools import ml
+from einops import rearrange
+from torch.nn.utils import weight_norm
+import torchaudio
+import nnAudio.features as features
+from munch import Munch
+
+
+BANDS = [(0.0, 0.1), (0.1, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)]
+
+
+def WNConv1d(*args, **kwargs):
+    act = kwargs.pop("act", True)
+    conv = weight_norm(nn.Conv1d(*args, **kwargs))
+    if not act:
+        return conv
+    return nn.Sequential(conv, nn.LeakyReLU(0.1))
+
+
+def WNConv2d(*args, **kwargs):
+    act = kwargs.pop("act", True)
+    conv = weight_norm(nn.Conv2d(*args, **kwargs))
+    if not act:
+        return conv
+    return nn.Sequential(conv, nn.LeakyReLU(0.1))
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def get_2d_padding(kernel_size, dilation=(1, 1)):
+    return (int((kernel_size[0] * dilation[0] - dilation[0]) / 2),
+            int((kernel_size[1] * dilation[1] - dilation[1]) / 2))
+
+
+class NormConv2d(nn.Module):
+    """Conv2d with normalization"""
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, 
+                 padding=0, dilation=1, groups=1, bias=True, norm="weight_norm"):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, 
+                             stride, padding, dilation, groups, bias)
+        if norm == "weight_norm":
+            self.conv = weight_norm(self.conv)
+    
+    def forward(self, x):
+        return self.conv(x)
+
+
+class MPD(nn.Module):
+    def __init__(self, period):
+        super().__init__()
+        self.period = period
+        self.convs = nn.ModuleList([
+            WNConv2d(1, 32, (5, 1), (3, 1), padding=(2, 0)),
+            WNConv2d(32, 128, (5, 1), (3, 1), padding=(2, 0)),
+            WNConv2d(128, 512, (5, 1), (3, 1), padding=(2, 0)),
+            WNConv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0)),
+            WNConv2d(1024, 1024, (5, 1), 1, padding=(2, 0)),
+        ])
+        self.conv_post = WNConv2d(1024, 1, kernel_size=(3, 1), padding=(1, 0), act=False)
+
+    def pad_to_period(self, x):
+        t = x.shape[-1]
+        x = F.pad(x, (0, self.period - t % self.period), mode="reflect")
+        return x
+
+    def forward(self, x):
+        fmap = []
+        x = self.pad_to_period(x)
+        x = rearrange(x, "b c (l p) -> b c l p", p=self.period)
+
+        for layer in self.convs:
+            x = layer(x)
+            fmap.append(x)
+
+        x = self.conv_post(x)
+        fmap.append(x)
+        return fmap
+
+
+class MSD(nn.Module):
+    def __init__(self, rate: int = 1, sample_rate: int = 44100):
+        super().__init__()
+        self.convs = nn.ModuleList([
+            WNConv1d(1, 16, 15, 1, padding=7),
+            WNConv1d(16, 64, 41, 4, groups=4, padding=20),
+            WNConv1d(64, 256, 41, 4, groups=16, padding=20),
+            WNConv1d(256, 1024, 41, 4, groups=64, padding=20),
+            WNConv1d(1024, 1024, 41, 4, groups=256, padding=20),
+            WNConv1d(1024, 1024, 5, 1, padding=2),
+        ])
+        self.conv_post = WNConv1d(1024, 1, 3, 1, padding=1, act=False)
+        self.sample_rate = sample_rate
+        self.rate = rate
+
+    def forward(self, x):
+        x = AudioSignal(x, self.sample_rate)
+        x.resample(self.sample_rate // self.rate)
+        x = x.audio_data
+
+        fmap = []
+        for l in self.convs:
+            x = l(x)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return fmap
+
+
+class DiscriminatorCQT(nn.Module):
+    def __init__(self, cfg, hop_length, n_octaves, bins_per_octave):
+        super().__init__()
+        self.cfg = cfg
+        self.filters = cfg.filters
+        self.max_filters = cfg.max_filters
+        self.filters_scale = cfg.filters_scale
+        self.kernel_size = (3, 9)
+        self.dilations = cfg.dilations
+        self.stride = (1, 2)
+        self.in_channels = cfg.in_channels
+        self.out_channels = cfg.out_channels
+        self.fs = cfg.sampling_rate
+        self.hop_length = hop_length
+        self.n_octaves = n_octaves
+        self.bins_per_octave = bins_per_octave
+
+        self.cqt_transform = features.cqt.CQT2010v2(
+            sr=self.fs * 2,
+            hop_length=self.hop_length,
+            n_bins=self.bins_per_octave * self.n_octaves,
+            bins_per_octave=self.bins_per_octave,
+            output_format="Complex",
+            pad_mode="constant",
+        )
+
+        self.conv_pres = nn.ModuleList()
+        for i in range(self.n_octaves):
+            self.conv_pres.append(
+                NormConv2d(
+                    self.in_channels * 2,  # Real + Imaginary
+                    self.in_channels * 2,
+                    kernel_size=self.kernel_size,
+                    padding=get_2d_padding(self.kernel_size),
+                    norm="weight_norm",
+                )
+            )
+
+        self.convs = nn.ModuleList()
+        self.convs.append(
+            NormConv2d(
+                self.in_channels * 2,
+                self.filters,
+                kernel_size=self.kernel_size,
+                padding=get_2d_padding(self.kernel_size),
+            )
+        )
+
+        in_chs = min(self.filters_scale * self.filters, self.max_filters)
+        for i, dilation in enumerate(self.dilations):
+            out_chs = min((self.filters_scale ** (i + 1)) * self.filters, self.max_filters)
+            self.convs.append(
+                NormConv2d(
+                    in_chs,
+                    out_chs,
+                    kernel_size=self.kernel_size,
+                    stride=self.stride,
+                    dilation=(dilation, 1),
+                    padding=get_2d_padding(self.kernel_size, (dilation, 1)),
+                    norm="weight_norm",
+                )
+            )
+            in_chs = out_chs
+            
+        out_chs = min(
+            (self.filters_scale ** (len(self.dilations) + 1)) * self.filters,
+            self.max_filters,
+        )
+        self.convs.append(
+            NormConv2d(
+                in_chs,
+                out_chs,
+                kernel_size=(self.kernel_size[0], self.kernel_size[0]),
+                padding=get_2d_padding((self.kernel_size[0], self.kernel_size[0])),
+                norm="weight_norm",
+            )
+        )
+
+        self.conv_post = NormConv2d(
+            out_chs,
+            self.out_channels,
+            kernel_size=(self.kernel_size[0], self.kernel_size[0]),
+            padding=get_2d_padding((self.kernel_size[0], self.kernel_size[0])),
+            norm="weight_norm",
+        )
+
+        self.activation = torch.nn.LeakyReLU(negative_slope=0.1)
+        self.resample = torchaudio.transforms.Resample(
+            orig_freq=self.fs, new_freq=self.fs * 2
+        )
+
+    def forward(self, x):
+        fmap = []
+        x = self.resample(x)
+        z = self.cqt_transform(x)
+        
+
+        z_amplitude = z[:, :, :, 0].unsqueeze(1)
+        z_phase = z[:, :, :, 1].unsqueeze(1)
+        z = torch.cat([z_amplitude, z_phase], dim=1)
+        z = rearrange(z, "b c w t -> b c t w")
+
+        latent_z = []
+        for i in range(self.n_octaves):
+            octave_band = z[:, :, :, i * self.bins_per_octave : (i + 1) * self.bins_per_octave]
+            processed_band = self.conv_pres[i](octave_band)
+            latent_z.append(processed_band)
+        latent_z = torch.cat(latent_z, dim=-1)
+
+        for i, l in enumerate(self.convs):
+            latent_z = l(latent_z)
+            latent_z = self.activation(latent_z)
+            fmap.append(latent_z)
+
+        latent_z = self.conv_post(latent_z)
+        fmap.append(latent_z)
+        
+        return fmap
+
+
+class MultiScaleSubbandCQT(nn.Module):
+    """CQT discriminator at multiple scales"""
+    def __init__(self, sample_rate=44100):
+        super().__init__()
+        cfg = Munch({
+            "hop_lengths": [1024, 512, 512],
+            "sampling_rate": sample_rate,
+            "filters": 32,
+            "max_filters": 1024,
+            "filters_scale": 1,
+            "dilations": [1, 2, 4],
+            "in_channels": 1,
+            "out_channels": 1,
+            "n_octaves": [10, 10, 10],
+            "bins_per_octaves": [24, 36, 48],
+        })
+        self.cfg = cfg
+        self.discriminators = nn.ModuleList([
+            DiscriminatorCQT(
+                cfg,
+                hop_length=cfg.hop_lengths[i],
+                n_octaves=cfg.n_octaves[i],
+                bins_per_octave=cfg.bins_per_octaves[i],
+            )
+            for i in range(len(cfg.hop_lengths))
+        ])
+
+    def forward(self, x):
+        fmap = []
+        for disc in self.discriminators:
+            fmap.extend(disc(x))
+        return fmap
+
+
+BANDS = [(0.0, 0.1), (0.1, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)]
+
+class MRD(nn.Module):
+    def __init__(self, window_length: int, hop_factor: float = 0.25, 
+                 sample_rate: int = 44100, bands: list = BANDS):
+        """Multi-resolution spectrogram discriminator."""
+        super().__init__()
+        self.window_length = window_length
+        self.hop_factor = hop_factor
+        self.sample_rate = sample_rate
+        self.stft_params = STFTParams(
+            window_length=window_length,
+            hop_length=int(window_length * hop_factor),
+            match_stride=True,
+        )
+
+        n_fft = window_length // 2 + 1
+        bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
+        self.bands = bands
+
+        ch = 32
+        convs = lambda: nn.ModuleList([
+            WNConv2d(2, ch, (3, 9), (1, 1), padding=(1, 4)),
+            WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
+            WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
+            WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
+            WNConv2d(ch, ch, (3, 3), (1, 1), padding=(1, 1)),
+        ])
+        self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
+        self.conv_post = WNConv2d(ch, 1, (3, 3), (1, 1), padding=(1, 1), act=False)
+
+    def spectrogram(self, x):
+        x = AudioSignal(x, self.sample_rate, stft_params=self.stft_params)
+        x = torch.view_as_real(x.stft())
+        x = rearrange(x, "b 1 f t c -> (b 1) c t f")
+        x_bands = [x[..., b[0] : b[1]] for b in self.bands]
+        return x_bands
+
+    def forward(self, x):
+        x_bands = self.spectrogram(x)
+        fmap = []
+
+        x = []
+        for band, stack in zip(x_bands, self.band_convs):
+            for layer in stack:
+                band = layer(band)
+                fmap.append(band)
+            x.append(band)
+
+        x = torch.cat(x, dim=-1)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return fmap
+
+
+class Discriminator(ml.BaseModel):
+    def __init__(
+        self,
+        rates: list = [],
+        periods: list = [2, 3, 5, 7, 11],
+        fft_sizes: list = [2048, 1024, 512],
+        sample_rate: int = 44100,
+    ):
+        """Discriminator combining MPD, MSD, MRD and CQT.
+        
+        Parameters
+        ----------
+        rates : list, optional
+            Sampling rates for MSD, by default []
+        periods : list, optional
+            Periods for MPD, by default [2, 3, 5, 7, 11]
+        fft_sizes : list, optional
+            FFT sizes for MRD, by default [2048, 1024, 512]
+        sample_rate : int, optional
+            Sampling rate of audio in Hz, by default 44100
+        """
+        super().__init__()
+        discs = []
+        # Time-domain discriminators
+        discs += [MPD(p) for p in periods]
+        discs += [MSD(r, sample_rate=sample_rate) for r in rates]
+        
+        # Frequency-domain discriminators (both STFT and CQT)
+        discs += [MRD(f, sample_rate=sample_rate) for f in fft_sizes]
+        discs += [MultiScaleSubbandCQT(sample_rate=sample_rate)]
+            
+        self.discriminators = nn.ModuleList(discs)
+
+    def preprocess(self, y):
+        # Remove DC offset
+        y = y - y.mean(dim=-1, keepdims=True)
+        # Peak normalize
+        y = 0.8 * y / (y.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
+        return y
+
+    def forward(self, x):
+        x = self.preprocess(x)
+        fmaps = [d(x) for d in self.discriminators]
+        return fmaps

higgs_audio_tokenizer.py

@@ -0,0 +1,373 @@
+# Based on code from: https://github.com/zhenye234/xcodec
+# Licensed under MIT License
+# Modifications by BosonAI
+
+import math
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Union, Sequence
+import numpy as np
+from transformers import AutoModel
+import torchaudio
+import json
+import librosa
+from huggingface_hub import snapshot_download
+
+from vector_quantize_pytorch import ResidualFSQ
+from descriptaudiocodec.dac.model import dac as dac2
+from quantization.vq import ResidualVectorQuantizer
+from semantic_module import Encoder, Decoder
+
+from transformers import HubertModel
+
+
+# At the top of higgs_audio_tokenizer.py, after the imports
+
+def WNConv1d(*args, **kwargs):
+    """Applies weight normalization to a 1D Convolutional layer."""
+    return nn.utils.weight_norm(nn.Conv1d(*args, **kwargs))
+
+def WNLinear(*args, **kwargs):
+    """Applies weight normalization to a Linear layer."""
+    return nn.utils.weight_norm(nn.Linear(*args, **kwargs))
+
+def init_weights(m):
+    """
+    Applies Xavier (Glorot) uniform initialization to Conv and Linear layers.
+    This is a robust, "classic" initialization scheme.
+    """
+    if isinstance(m, (nn.Conv1d, nn.Conv2d)):
+        # Truncated normal initialization for convolutional layers
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.Linear):
+        # Also apply to linear layers for consistency
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.Embedding):
+        # Initialize the codebook gently as well
+        nn.init.trunc_normal_(m.weight, std=0.02)
+
+
+class EncodedResult:
+    def __init__(self, audio_codes):
+        self.audio_codes = audio_codes
+
+class HiggsAudioFeatureExtractor(nn.Module):
+    def __init__(self, sampling_rate=16000):
+        super().__init__()
+        self.sampling_rate = sampling_rate
+
+    def forward(self, raw_audio, sampling_rate=16000, return_tensors="pt"):
+        audio_signal = torch.tensor(raw_audio)
+        audio_signal = audio_signal.unsqueeze(0)
+        if len(audio_signal.shape) < 3:
+            audio_signal = audio_signal.unsqueeze(0)
+        return {"input_values": audio_signal}
+
+
+class HiggsAudioTokenizer(nn.Module):
+    def __init__(
+        self,
+        n_filters: int = 32,
+        D: int = 128,
+        target_bandwidths: Sequence[Union[int, float]] = [1, 1.5, 2, 4, 6],
+        ratios: Sequence[int] = [8, 5, 4, 2],  #  downsampling by 320
+        sample_rate: int = 16000,
+        bins: int = 1024,
+        n_q: int = 8,
+        codebook_dim: int = None,
+        normalize: bool = False,
+        causal: bool = False,
+        semantic_techer: str = "hubert_base_general",
+        last_layer_semantic: bool = True,
+        merge_mode: str = "concat",
+        downsample_mode: str = "step_down",
+        semantic_mode: str = "classic",
+        vq_scale: int = 1,
+        semantic_sample_rate: int = None,
+        device: str = "cuda",
+    ):
+        super().__init__()
+        self.hop_length = np.prod(ratios)
+        self.semantic_techer = semantic_techer
+
+        self.frame_rate = math.ceil(sample_rate / np.prod(ratios))  # 50 Hz
+
+        self.target_bandwidths = target_bandwidths
+        self.n_q = n_q
+        self.sample_rate = sample_rate
+        self.encoder = dac2.Encoder(64, ratios, D)
+
+        self.decoder_2 = dac2.Decoder(D, 1024, ratios)
+        self.last_layer_semantic = last_layer_semantic
+        self.device = device
+        if semantic_techer == "hubert_base":
+            self.semantic_model = AutoModel.from_pretrained("facebook/hubert-base-ls960")
+            self.semantic_sample_rate = 16000
+            self.semantic_dim = 768
+            self.encoder_semantic_dim = 768
+
+        elif semantic_techer == "wavlm_base_plus":
+            self.semantic_model = AutoModel.from_pretrained("microsoft/wavlm-base-plus")
+            self.semantic_sample_rate = 16000
+            self.semantic_dim = 768
+            self.encoder_semantic_dim = 768
+            
+        elif semantic_techer == "mHubert_base":
+            self.semantic_model = AutoModel.from_pretrained("utter-project/mHuBERT-147")
+            self.semantic_sample_rate = 16000
+            self.semantic_dim = 768
+            self.encoder_semantic_dim = 768
+
+        elif semantic_techer == "hubert_base_general":
+            self.semantic_model = HubertModel.from_pretrained("/home/ubuntu/.cache/huggingface/hub/models--bosonai--hubert_base/snapshots/b4b85f1652c16ad63fdc818221b215b79ff55934", trust_remote_code=False)
+            self.semantic_sample_rate = 16000
+            self.semantic_dim = 768
+            self.encoder_semantic_dim = 768
+
+        # Overwrite semantic model sr to ensure semantic_downsample_factor is an integer
+        if semantic_sample_rate is not None:
+            self.semantic_sample_rate = semantic_sample_rate
+
+        self.semantic_model.eval()
+
+        # make the semantic model parameters do not need gradient
+        for param in self.semantic_model.parameters():
+            param.requires_grad = False
+
+        self.semantic_downsample_factor = int(self.hop_length / (self.sample_rate / self.semantic_sample_rate) / 320)
+
+        self.quantizer_dim = int((D + self.encoder_semantic_dim) // vq_scale)
+        self.encoder_semantic = Encoder(input_channels=self.semantic_dim, encode_channels=self.encoder_semantic_dim)
+        self.decoder_semantic = Decoder(
+            code_dim=self.encoder_semantic_dim, output_channels=self.semantic_dim, decode_channels=self.semantic_dim
+        )
+
+        # out_D=D+768
+        if isinstance(bins, int):  # RVQ
+            self.quantizer = ResidualVectorQuantizer(
+                dimension=self.quantizer_dim, codebook_dim=codebook_dim, n_q=n_q, bins=bins
+            )
+            self.quantizer_type = "RVQ"
+        else:  # RFSQ
+            self.quantizer = ResidualFSQ(dim=self.quantizer_dim, levels=bins, num_quantizers=n_q)
+            self.quantizer_type = "RFSQ"
+
+        # self.fc_prior = nn.Linear(D + self.encoder_semantic_dim, self.quantizer_dim)
+        # self.fc_post1 = nn.Linear(self.quantizer_dim, self.encoder_semantic_dim)
+        # self.fc_post2 = nn.Linear(self.quantizer_dim, D)
+
+
+        self.fc_prior = WNLinear(D + self.encoder_semantic_dim, self.quantizer_dim)
+        self.fc_post1 = WNLinear(self.quantizer_dim, self.encoder_semantic_dim)
+        self.fc_post2 = WNLinear(self.quantizer_dim, D)
+        
+        
+        self.downsample_mode = downsample_mode
+        if downsample_mode == "avg":
+            self.semantic_pooling = nn.AvgPool1d(
+                kernel_size=self.semantic_downsample_factor, stride=self.semantic_downsample_factor
+            )
+
+        self.audio_tokenizer_feature_extractor = HiggsAudioFeatureExtractor(sampling_rate=self.sample_rate)
+
+        self.apply(init_weights)
+
+    @property
+    def tps(self):
+        return self.frame_rate
+
+    @property
+    def sampling_rate(self):
+        return self.sample_rate
+
+    @property
+    def num_codebooks(self):
+        return self.n_q
+
+    @property
+    def codebook_size(self):
+        return self.quantizer_dim
+
+    def get_last_layer(self):
+        return self.decoder.layers[-1].weight
+
+    def calculate_rec_loss(self, rec, target):
+        target = target / target.norm(dim=-1, keepdim=True)
+        rec = rec / rec.norm(dim=-1, keepdim=True)
+        rec_loss = (1 - (target * rec).sum(-1)).mean()
+
+        return rec_loss
+
+    @torch.no_grad()
+    def get_regress_target(self, x):
+        x = torchaudio.functional.resample(x, self.sample_rate, self.semantic_sample_rate)
+
+        if (
+            self.semantic_techer == "hubert_base"
+            or self.semantic_techer == "hubert_base_general"
+            or self.semantic_techer == "wavlm_base_plus"
+        ):
+            x = x[:, 0, :]
+            x = F.pad(x, (160, 160))
+            target = self.semantic_model(x, output_hidden_states=True).hidden_states
+            target = torch.stack(target, dim=1)  # .transpose(-1, -2)#.flatten(start_dim=1, end_dim=2)
+
+            # average for all layers
+            target = target.mean(1)
+            # target = target[9]
+            # if self.hop_length > 320:
+            #     target = self.semantic_pooling(target.transpose(1, 2)).transpose(1, 2)
+
+        elif self.semantic_techer == "w2v_bert2":
+            target = self.semantic_model(x)
+
+        elif self.semantic_techer.startswith("whisper"):
+            if self.last_layer_semantic:
+                target = self.semantic_model(x, avg_layers=False)
+            else:
+                target = self.semantic_model(x, avg_layers=True)
+
+        elif self.semantic_techer.startswith("mert_music"):
+            if self.last_layer_semantic:
+                target = self.semantic_model(x, avg_layers=False)
+            else:
+                target = self.semantic_model(x, avg_layers=True)
+
+        elif self.semantic_techer.startswith("qwen_audio_omni"):
+            target = self.semantic_model(x)
+
+        if self.downsample_mode == "step_down":
+            if self.semantic_downsample_factor > 1:
+                target = target[:, :: self.semantic_downsample_factor, :]
+
+        elif self.downsample_mode == "avg":
+            target = self.semantic_pooling(target.transpose(1, 2)).transpose(1, 2)
+        return target
+
+    def forward(self, x: torch.Tensor, bw: int):
+        e_semantic_input = self.get_regress_target(x).detach()
+
+        e_semantic = self.encoder_semantic(e_semantic_input.transpose(1, 2))
+        e_acoustic = self.encoder(x)
+
+        e = torch.cat([e_acoustic, e_semantic], dim=1)
+
+        e = self.fc_prior(e.transpose(1, 2))
+
+        if self.quantizer_type == "RVQ":
+            e = e.transpose(1, 2)
+            quantized, codes, bandwidth, commit_loss = self.quantizer(e, self.frame_rate, bw)
+            quantized = quantized.transpose(1, 2)
+        else:
+            quantized, codes = self.quantizer(e)
+            commit_loss = torch.tensor(0.0)
+
+        quantized_semantic = self.fc_post1(quantized).transpose(1, 2)
+        quantized_acoustic = self.fc_post2(quantized).transpose(1, 2)
+
+        o = self.decoder_2(quantized_acoustic)
+
+        o_semantic = self.decoder_semantic(quantized_semantic)
+        semantic_recon_loss = F.mse_loss(e_semantic_input.transpose(1, 2).detach(), o_semantic)
+
+        return o, commit_loss, semantic_recon_loss, None
+
+    def encode(self, audio_path_or_wv, sr=None, loudness_normalize=False, loudness_threshold=-23.0):
+        if isinstance(audio_path_or_wv, str):
+            wv, sr = librosa.load(audio_path_or_wv, mono=True, sr=None)
+        else:
+            wv = audio_path_or_wv
+            assert sr is not None
+        if loudness_normalize:
+            import pyloudnorm as pyln
+
+            meter = pyln.Meter(sr)
+            l = meter.integrated_loudness(wv)
+            wv = pyln.normalize.loudness(wv, l, loudness_threshold)
+        if sr != self.sampling_rate:
+            wv = librosa.resample(wv, orig_sr=sr, target_sr=self.sampling_rate)
+        if self.audio_tokenizer_feature_extractor is not None:
+            inputs = self.audio_tokenizer_feature_extractor(
+                raw_audio=wv, sampling_rate=self.audio_tokenizer_feature_extractor.sampling_rate, return_tensors="pt"
+            )
+            input_values = inputs["input_values"].to(self.device)
+        else:
+            input_values = torch.from_numpy(wv).float().unsqueeze(0)
+        with torch.no_grad():
+            encoder_outputs = self._xcodec_encode(input_values)
+            vq_code = encoder_outputs.audio_codes[0]
+        return vq_code
+    
+
+
+    def _xcodec_encode(self, x: torch.Tensor, target_bw: Optional[int] = None) -> torch.Tensor:
+        bw = target_bw
+
+        e_semantic_input = self.get_regress_target(x).detach()
+
+        e_semantic = self.encoder_semantic(e_semantic_input.transpose(1, 2))
+        e_acoustic = self.encoder(x)
+
+        if e_acoustic.shape[2] != e_semantic.shape[2]:
+            pad_size = 160 * self.semantic_downsample_factor
+            e_acoustic = self.encoder(F.pad(x[:, 0, :], (pad_size, pad_size)).unsqueeze(0))
+
+        if e_acoustic.shape[2] != e_semantic.shape[2]:
+            if e_acoustic.shape[2] > e_semantic.shape[2]:
+                e_acoustic = e_acoustic[:, :, : e_semantic.shape[2]]
+            else:
+                e_semantic = e_semantic[:, :, : e_acoustic.shape[2]]
+
+        e = torch.cat([e_acoustic, e_semantic], dim=1)
+
+        e = self.fc_prior(e.transpose(1, 2))
+
+        if self.quantizer_type == "RVQ":
+            e = e.transpose(1, 2)
+            quantized, codes, bandwidth, commit_loss = self.quantizer(e, self.frame_rate, bw)
+            codes = codes.permute(1, 0, 2)
+        else:
+            quantized, codes = self.quantizer(e)
+            codes = codes.permute(0, 2, 1)
+
+        # return codes
+        return EncodedResult(codes)
+
+    def decode(self, vq_code: torch.Tensor) -> torch.Tensor:
+        if self.quantizer_type == "RVQ":
+            vq_code = vq_code.permute(1, 0, 2)
+            quantized = self.quantizer.decode(vq_code)
+            quantized = quantized.transpose(1, 2)
+        else:
+            vq_code = vq_code.permute(0, 2, 1)
+            quantized = self.quantizer.get_output_from_indices(vq_code)
+        quantized_acoustic = self.fc_post2(quantized).transpose(1, 2)
+
+        o = self.decoder_2(quantized_acoustic)
+        return o.cpu().numpy()
+
+
+def load_higgs_audio_tokenizer(tokenizer_name_or_path, device="cuda"):
+    is_local = os.path.exists(tokenizer_name_or_path)
+    if not is_local:
+        tokenizer_path = snapshot_download(tokenizer_name_or_path)
+    else:
+        tokenizer_path = tokenizer_name_or_path
+    config_path = os.path.join(tokenizer_path, "config.json")
+    model_path = os.path.join(tokenizer_path, "model.pth")
+    config = json.load(open(config_path))
+    model = HiggsAudioTokenizer(
+        **config,
+        device=device,
+    )
+    parameter_dict = torch.load(model_path, map_location=device, weights_only=False)
+    model.load_state_dict(parameter_dict, strict=False)
+    model.to(device)
+    model.eval()
+    return model

loss.py

@@ -0,0 +1,368 @@
+import torch
+import torch.nn.functional as F
+from audiotools import AudioSignal
+from audiotools import STFTParams
+from torch import nn
+import typing
+from typing import List
+
+class L1Loss(nn.L1Loss):
+    """L1 Loss between AudioSignals. Defaults
+    to comparing ``audio_data``, but any
+    attribute of an AudioSignal can be used.
+
+    Parameters
+    ----------
+    attribute : str, optional
+        Attribute of signal to compare, defaults to ``audio_data``.
+    weight : float, optional
+        Weight of this loss, defaults to 1.0.
+
+    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/distance.py
+    """
+
+    def __init__(self, attribute: str = "audio_data", weight: float = 1.0, **kwargs):
+        self.attribute = attribute
+        self.weight = weight
+        super().__init__(**kwargs)
+
+    def forward(self, x: AudioSignal, y: AudioSignal):
+        """
+        Parameters
+        ----------
+        x : AudioSignal
+            Estimate AudioSignal
+        y : AudioSignal
+            Reference AudioSignal
+
+        Returns
+        -------
+        torch.Tensor
+            L1 loss between AudioSignal attributes.
+        """
+        if isinstance(x, AudioSignal):
+            x = getattr(x, self.attribute)
+            y = getattr(y, self.attribute)
+        return super().forward(x, y)
+
+
+class SISDRLoss(nn.Module):
+    """
+    Computes the Scale-Invariant Source-to-Distortion Ratio between a batch
+    of estimated and reference audio signals or aligned features.
+
+    Parameters
+    ----------
+    scaling : int, optional
+        Whether to use scale-invariant (True) or
+        signal-to-noise ratio (False), by default True
+    reduction : str, optional
+        How to reduce across the batch (either 'mean',
+        'sum', or none).], by default ' mean'
+    zero_mean : int, optional
+        Zero mean the references and estimates before
+        computing the loss, by default True
+    clip_min : int, optional
+        The minimum possible loss value. Helps network
+        to not focus on making already good examples better, by default None
+    weight : float, optional
+        Weight of this loss, defaults to 1.0.
+
+    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/distance.py
+    """
+
+    def __init__(
+        self,
+        scaling: int = True,
+        reduction: str = "mean",
+        zero_mean: int = True,
+        clip_min: int = None,
+        weight: float = 1.0,
+    ):
+        self.scaling = scaling
+        self.reduction = reduction
+        self.zero_mean = zero_mean
+        self.clip_min = clip_min
+        self.weight = weight
+        super().__init__()
+
+    def forward(self, x: AudioSignal, y: AudioSignal):
+        eps = 1e-8
+        # nb, nc, nt
+        if isinstance(x, AudioSignal):
+            references = x.audio_data
+            estimates = y.audio_data
+        else:
+            references = x
+            estimates = y
+
+        nb = references.shape[0]
+        references = references.reshape(nb, 1, -1).permute(0, 2, 1)
+        estimates = estimates.reshape(nb, 1, -1).permute(0, 2, 1)
+
+        # samples now on axis 1
+        if self.zero_mean:
+            mean_reference = references.mean(dim=1, keepdim=True)
+            mean_estimate = estimates.mean(dim=1, keepdim=True)
+        else:
+            mean_reference = 0
+            mean_estimate = 0
+
+        _references = references - mean_reference
+        _estimates = estimates - mean_estimate
+
+        references_projection = (_references**2).sum(dim=-2) + eps
+        references_on_estimates = (_estimates * _references).sum(dim=-2) + eps
+
+        scale = (
+            (references_on_estimates / references_projection).unsqueeze(1)
+            if self.scaling
+            else 1
+        )
+
+        e_true = scale * _references
+        e_res = _estimates - e_true
+
+        signal = (e_true**2).sum(dim=1)
+        noise = (e_res**2).sum(dim=1)
+        sdr = -10 * torch.log10(signal / noise + eps)
+
+        if self.clip_min is not None:
+            sdr = torch.clamp(sdr, min=self.clip_min)
+
+        if self.reduction == "mean":
+            sdr = sdr.mean()
+        elif self.reduction == "sum":
+            sdr = sdr.sum()
+        return sdr
+
+
+class MultiScaleSTFTLoss(nn.Module):
+    """Computes the multi-scale STFT loss from [1].
+
+    Parameters
+    ----------
+    window_lengths : List[int], optional
+        Length of each window of each STFT, by default [2048, 512]
+    loss_fn : typing.Callable, optional
+        How to compare each loss, by default nn.L1Loss()
+    clamp_eps : float, optional
+        Clamp on the log magnitude, below, by default 1e-5
+    mag_weight : float, optional
+        Weight of raw magnitude portion of loss, by default 1.0
+    log_weight : float, optional
+        Weight of log magnitude portion of loss, by default 1.0
+    pow : float, optional
+        Power to raise magnitude to before taking log, by default 2.0
+    weight : float, optional
+        Weight of this loss, by default 1.0
+    match_stride : bool, optional
+        Whether to match the stride of convolutional layers, by default False
+
+    References
+    ----------
+
+    1.  Engel, Jesse, Chenjie Gu, and Adam Roberts.
+        "DDSP: Differentiable Digital Signal Processing."
+        International Conference on Learning Representations. 2019.
+
+    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
+    """
+
+    def __init__(
+        self,
+        window_lengths: List[int] = [2048, 512],
+        loss_fn: typing.Callable = nn.L1Loss(),
+        clamp_eps: float = 1e-5,
+        mag_weight: float = 1.0,
+        log_weight: float = 1.0,
+        pow: float = 2.0,
+        weight: float = 1.0,
+        match_stride: bool = False,
+        window_type: str = None,
+    ):
+        super().__init__()
+        self.stft_params = [
+            STFTParams(
+                window_length=w,
+                hop_length=w // 4,
+                match_stride=match_stride,
+                window_type=window_type,
+            )
+            for w in window_lengths
+        ]
+        self.loss_fn = loss_fn
+        self.log_weight = log_weight
+        self.mag_weight = mag_weight
+        self.clamp_eps = clamp_eps
+        self.weight = weight
+        self.pow = pow
+
+    def forward(self, x: AudioSignal, y: AudioSignal):
+        """Computes multi-scale STFT between an estimate and a reference
+        signal.
+
+        Parameters
+        ----------
+        x : AudioSignal
+            Estimate signal
+        y : AudioSignal
+            Reference signal
+
+        Returns
+        -------
+        torch.Tensor
+            Multi-scale STFT loss.
+        """
+        loss = 0.0
+        for s in self.stft_params:
+            x.stft(s.window_length, s.hop_length, s.window_type)
+            y.stft(s.window_length, s.hop_length, s.window_type)
+            loss += self.log_weight * self.loss_fn(
+                x.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
+                y.magnitude.clamp(self.clamp_eps).pow(self.pow).log10(),
+            )
+            loss += self.mag_weight * self.loss_fn(x.magnitude, y.magnitude)
+        return loss
+
+
+class MelSpectrogramLoss(nn.Module):
+    """Compute distance between mel spectrograms. Can be used
+    in a multi-scale way.
+
+    Parameters
+    ----------
+    n_mels : List[int]
+        Number of mels per STFT, by default [150, 80],
+    window_lengths : List[int], optional
+        Length of each window of each STFT, by default [2048, 512]
+    loss_fn : typing.Callable, optional
+        How to compare each loss, by default nn.L1Loss()
+    clamp_eps : float, optional
+        Clamp on the log magnitude, below, by default 1e-5
+    mag_weight : float, optional
+        Weight of raw magnitude portion of loss, by default 1.0
+    log_weight : float, optional
+        Weight of log magnitude portion of loss, by default 1.0
+    pow : float, optional
+        Power to raise magnitude to before taking log, by default 2.0
+    weight : float, optional
+        Weight of this loss, by default 1.0
+    match_stride : bool, optional
+        Whether to match the stride of convolutional layers, by default False
+
+    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
+    """
+
+    def __init__(
+        self,
+        n_mels: List[int] = [150, 80],
+        window_lengths: List[int] = [2048, 512],
+        loss_fn: typing.Callable = nn.L1Loss(),
+        clamp_eps: float = 1e-5,
+        mag_weight: float = 1.0,
+        log_weight: float = 1.0,
+        pow: float = 2.0,
+        weight: float = 1.0,
+        match_stride: bool = False,
+        mel_fmin: List[float] = [0.0, 0.0],
+        mel_fmax: List[float] = [None, None],
+        window_type: str = None,
+    ):
+        super().__init__()
+        self.stft_params = [
+            STFTParams(
+                window_length=w,
+                hop_length=w // 4,
+                match_stride=match_stride,
+                window_type=window_type,
+            )
+            for w in window_lengths
+        ]
+        self.n_mels = n_mels
+        self.loss_fn = loss_fn
+        self.clamp_eps = clamp_eps
+        self.log_weight = log_weight
+        self.mag_weight = mag_weight
+        self.weight = weight
+        self.mel_fmin = mel_fmin
+        self.mel_fmax = mel_fmax
+        self.pow = pow
+
+    def forward(self, x: AudioSignal, y: AudioSignal):
+        """Computes mel loss between an estimate and a reference
+        signal.
+
+        Parameters
+        ----------
+        x : AudioSignal
+            Estimate signal
+        y : AudioSignal
+            Reference signal
+
+        Returns
+        -------
+        torch.Tensor
+            Mel loss.
+        """
+        loss = 0.0
+        for n_mels, fmin, fmax, s in zip(
+            self.n_mels, self.mel_fmin, self.mel_fmax, self.stft_params
+        ):
+            kwargs = {
+                "window_length": s.window_length,
+                "hop_length": s.hop_length,
+                "window_type": s.window_type,
+            }
+            x_mels = x.mel_spectrogram(n_mels, mel_fmin=fmin, mel_fmax=fmax, **kwargs)
+            y_mels = y.mel_spectrogram(n_mels, mel_fmin=fmin, mel_fmax=fmax, **kwargs)
+
+            loss += self.log_weight * self.loss_fn(
+                x_mels.clamp(self.clamp_eps).pow(self.pow).log10(),
+                y_mels.clamp(self.clamp_eps).pow(self.pow).log10(),
+            )
+            loss += self.mag_weight * self.loss_fn(x_mels, y_mels)
+        return loss
+
+
+class GANLoss(nn.Module):
+    """
+    Computes a discriminator loss, given a discriminator on
+    generated waveforms/spectrograms compared to ground truth
+    waveforms/spectrograms. Computes the loss for both the
+    discriminator and the generator in separate functions.
+    """
+
+    def __init__(self, discriminator):
+        super().__init__()
+        self.discriminator = discriminator
+
+    def forward(self, fake, real):
+        d_fake = self.discriminator(fake.audio_data)
+        d_real = self.discriminator(real.audio_data)
+        return d_fake, d_real
+
+    def discriminator_loss(self, fake, real):
+        d_fake, d_real = self.forward(fake.clone().detach(), real)
+
+        loss_d = 0
+        for x_fake, x_real in zip(d_fake, d_real):
+            loss_d += torch.mean(x_fake[-1] ** 2)
+            loss_d += torch.mean((1 - x_real[-1]) ** 2)
+        return loss_d
+
+    def generator_loss(self, fake, real):
+        d_fake, d_real = self.forward(fake, real)
+
+        loss_g = 0
+        for x_fake in d_fake:
+            loss_g += torch.mean((1 - x_fake[-1]) ** 2)
+
+        loss_feature = 0
+
+        for i in range(len(d_fake)):
+            for j in range(len(d_fake[i]) - 1):
+                loss_feature += F.l1_loss(d_fake[i][j], d_real[i][j].detach())
+        return loss_g, loss_feature
+    
+    

outputs/logs/250801-104649/events.out.tfevents.1754045209.192-222-50-191.575849.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:adf08230fc10f50d2a7fcb4c2deaf1b2cbb45116b6d44f1cce5eb365471f516b
+size 657

outputs/logs/250801-104824/events.out.tfevents.1754045304.192-222-50-191.577752.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85e63a92198fbc1d513003ee532294feaedf87505a51e15ec2fc520900ec53c4
+size 657

outputs/logs/250801-104944/events.out.tfevents.1754045384.192-222-50-191.579650.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4e7f045c676534be11f5ad39b553611caa786d3e1559cbf100ca84133d54a77
+size 88

outputs/logs/250801-105034/events.out.tfevents.1754045434.192-222-50-191.581483.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5c97f5d3b5b6cc60223a60dceb2b06f6a847e682935f8d45086892f176bd98d
+size 657

outputs/logs/250801-105133/events.out.tfevents.1754045493.192-222-50-191.583409.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7bae6adc53141720bb44a2fda69e2f97bce51b98365458ec53ce59c66bf255f
+size 5751664

outputs/logs/250801-134657/events.out.tfevents.1754056017.192-222-50-191.688744.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c021361932098e4ba601f59134b8ea6c0192a31d103eb6374070e394bb740060
+size 61388

outputs/logs/250801-135301/events.out.tfevents.1754056381.192-222-50-191.693590.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61a297a283372869b74886f9fa2a17d11bd77fe284275aae505fbaa28957cda8
+size 88

outputs/logs/250801-135344/events.out.tfevents.1754056424.192-222-50-191.695388.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0954cc521ca9b594ace0fb485ccf8556621e017c08da9d525e4071c94489c3ec
+size 657

outputs/logs/250801-135510/events.out.tfevents.1754056510.192-222-50-191.697490.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1dcf0296a6058bce4cfc81563b9ffaac98d0e171d3f69e15008f179aeb4e5d8e
+size 3419391

outputs/logs/250801-202235/events.out.tfevents.1754079755.192-222-50-191.6026.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:163634220bcd180cfe60f9d73024faeeeaed6c2849ea5d44909e6666e4a4ac54
+size 88

outputs/logs/250801-202320/events.out.tfevents.1754079800.192-222-50-191.6708.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7088dc97ef965f0fb3bfff3b1f28e1b63644f666b6464a8524235f5bdd68a24
+size 7234042

outputs/logs/250802-065733/events.out.tfevents.1754117853.192-222-50-191.86944.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b13f240f199d42345f13428575a813d039fdd4e8899a915b7b55741b504ea46
+size 208834

outputs/logs/250802-072035/events.out.tfevents.1754119235.192-222-50-191.100690.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93fc6699e79f38e10ada3f7431617725b9b35daee4cda1978fb73800f7304113
+size 3373

outputs_24/logs/250730-112649/events.out.tfevents.1753874809.192-222-50-191.3556345.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34843a6a8d0849900de8943c4bbf4dcf77424e605eb1e4a721b8cbf894f1fc73
+size 88

outputs_24/logs/250730-112910/events.out.tfevents.1753874950.192-222-50-191.3557426.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d9cca204af08430087351d78de76e17d16ad4ed8e2e31540b0857406c77bb86
+size 1990

outputs_24/logs/250730-113135/events.out.tfevents.1753875095.192-222-50-191.3558918.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0f7180018cd52196492be018bc160ce1222c437e27d4019bd87d5207e9f1fe5
+size 63081

outputs_24/logs/250730-114727/events.out.tfevents.1753876047.192-222-50-191.3567432.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:395a91280aa73bfaf31de37cb66f76ca9916e25d59fa84568e2124ad19162e7a
+size 3808

outputs_24/logs/250730-115006/events.out.tfevents.1753876206.192-222-50-191.3569242.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b47cceec7b3b651af1fa77129099bd8f0fcebfb60ba0950b1a0efeba2ea28a2f
+size 6744012

outputs_24/logs/250730-151325/events.out.tfevents.1753888405.192-222-50-191.3660307.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69269e593e8fb146793020448eedf00cda64a27f3b3a3cebf47dbc5c7df5e8dc
+size 5976

outputs_24/logs/250730-152054/events.out.tfevents.1753888854.192-222-50-191.3663830.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f87195745b3babee5ab3afbc4b46d85decbf2c39f31376891b3d001f1e1db52
+size 88

outputs_24/logs/250730-152132/events.out.tfevents.1753888892.192-222-50-191.3664702.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7e10c22d4a06fb02d74592b2ac393196affce38724e0c2fa3117dfec6931acf
+size 88

outputs_24/logs/250730-152218/events.out.tfevents.1753888938.192-222-50-191.3665630.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c9e1081361a51f466892f52660c7126dd4e281ad1e1a86b01ecf914a03318ed
+size 657

outputs_24/logs/250730-152329/events.out.tfevents.1753889009.192-222-50-191.3666743.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbbfa9f88f5e7790426f8bee52f171526f4d7fb1ab5efb51e25c77581a3afee2
+size 657

outputs_24/logs/250730-152554/events.out.tfevents.1753889154.192-222-50-191.3668339.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7bd5fbc186ff41408a06c8d8ecb91ea3429ab1971c864edb0411ae2391af8965
+size 88

outputs_24/logs/250730-152702/events.out.tfevents.1753889222.192-222-50-191.3669391.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df5419dd06bca60309a60b28ee070c8fc816845de941a0b311ebfbb6777dccac
+size 88

outputs_24/logs/250730-152902/events.out.tfevents.1753889342.192-222-50-191.3671654.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f2377fd533270911cdd47226a5b11222bb0d037b3a47eaac9c66b3ffc605d03
+size 1526378

outputs_24/logs/250730-161025/events.out.tfevents.1753891825.192-222-50-191.3698156.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f3427b4010f32544ebc7c050f6ed85ce380c15bddcc039e3fffa295b4a2b4813
+size 1528786

outputs_24/logs/250730-165034/events.out.tfevents.1753894234.192-222-50-191.3717308.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:242a5126eb7e28c25d6ccf4377c745c970065110502df9dbf74e3392eff098c8
+size 4794

outputs_24/logs/250730-165327/events.out.tfevents.1753894407.192-222-50-191.3719515.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f02d9f09b42a51e0879c5de3e059703520e0aa2c0a1fc329e2878ecd0deb1c23
+size 657