Upload 5 files

best_model.pt

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

clap_tag_embeddings.npy

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

mulan_tag_embeddings.npy

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

music_tagger_gui.py

@@ -0,0 +1,554 @@
+"""
+Gradio Music Tagging GUI
+
+A simple web interface for tagging music using:
+1. Zero-Shot: CLAP + MuLan models with pre-computed tag embeddings
+2. MTG-Jamendo: Trained MERT classifier
+
+Usage:
+    python music_tagger_gui.py
+
+Requirements:
+    pip install -r requirements_gui.txt
+"""
+
+import os
+import torch
+import numpy as np
+import librosa
+import torch.nn as nn
+import torch.nn.functional as F
+import gradio as gr
+from pathlib import Path
+
+
+# ============================================================
+# Configuration
+# ============================================================
+class Config:
+    """Configuration for both inference methods."""
+
+    # Device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Zero-Shot paths (relative to script directory)
+    clap_embeddings_path = "clap_tag_embeddings.npy"
+    mulan_embeddings_path = "mulan_tag_embeddings.npy"
+    tag_names_path = "musiccaps_tag_names.txt"
+
+    # MTG-Jamendo paths
+    mtg_checkpoint_path = "best_model.pt"
+
+    # MERT settings (for MTG-Jamendo)
+    mert_model_name = "m-a-p/MERT-v1-95M"
+    mert_layer = 11
+    mert_feature_dim = 768
+    mert_sample_rate = 24000
+    max_duration = 30  # seconds
+
+
+config = Config()
+
+# Get script directory for relative paths
+SCRIPT_DIR = Path(__file__).parent.resolve()
+
+
+# ============================================================
+# Model Classes
+# ============================================================
+class GenreClassifier(nn.Module):
+    """Classifier for MTG-Jamendo trained model."""
+
+    def __init__(self, input_dim=768, num_classes=50, hidden_dim=512, dropout=0.3):
+        super().__init__()
+        self.classifier = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim // 2, num_classes)
+        )
+
+    def forward(self, x):
+        return self.classifier(x)
+
+
+# ============================================================
+# Global Model Cache (lazy loading)
+# ============================================================
+class ModelCache:
+    """Lazy-loaded model cache to avoid loading models until needed."""
+
+    def __init__(self):
+        self._zero_shot_models = None
+        self._zero_shot_embeddings = None
+        self._zero_shot_tag_names = None
+        self._mtg_model = None
+        self._mtg_tag_names = None
+        self._mert_model = None
+        self._mert_processor = None
+
+    def get_zero_shot_models(self):
+        """Load and cache Zero-Shot models (CLAP + MuLan)."""
+        if self._zero_shot_models is None:
+            print("Loading Zero-Shot models (this may take a moment)...")
+
+            # Import here to avoid loading if not needed
+            try:
+                from muq import MuQMuLan
+                from transformers import ClapModel, ClapProcessor
+            except ImportError as e:
+                raise ImportError(
+                    f"Missing dependencies for Zero-Shot tagging: {e}\n"
+                    "Install with: pip install muq laion-clap transformers"
+                )
+
+            # Load MuLan
+            print("  Loading MuQ-MuLan...")
+            mulan_model = MuQMuLan.from_pretrained("OpenMuQ/MuQ-MuLan-large")
+            mulan_model = mulan_model.to(config.device).eval()
+
+            # Load CLAP
+            print("  Loading CLAP...")
+            clap_model = ClapModel.from_pretrained("laion/larger_clap_music_and_speech")
+            clap_model = clap_model.to(config.device).eval()
+            clap_processor = ClapProcessor.from_pretrained("laion/larger_clap_music_and_speech")
+
+            self._zero_shot_models = (mulan_model, clap_model, clap_processor)
+            print("  Zero-Shot models loaded!")
+
+        return self._zero_shot_models
+
+    def get_zero_shot_embeddings(self):
+        """Load and cache pre-computed tag embeddings."""
+        if self._zero_shot_embeddings is None:
+            clap_path = SCRIPT_DIR / config.clap_embeddings_path
+            mulan_path = SCRIPT_DIR / config.mulan_embeddings_path
+
+            if not clap_path.exists() or not mulan_path.exists():
+                raise FileNotFoundError(
+                    f"Pre-computed embeddings not found!\n"
+                    f"Expected:\n"
+                    f"  - {clap_path}\n"
+                    f"  - {mulan_path}\n\n"
+                    f"Run create_embeddings.py first to generate these files."
+                )
+
+            print("Loading pre-computed embeddings...")
+            clap_embeddings = np.load(str(clap_path))
+            mulan_embeddings = np.load(str(mulan_path))
+
+            self._zero_shot_embeddings = (clap_embeddings, mulan_embeddings)
+            print(f"  Loaded embeddings: CLAP {clap_embeddings.shape}, MuLan {mulan_embeddings.shape}")
+
+        return self._zero_shot_embeddings
+
+    def get_zero_shot_tag_names(self):
+        """Load and cache tag names for Zero-Shot tagging."""
+        if self._zero_shot_tag_names is None:
+            tag_path = SCRIPT_DIR / config.tag_names_path
+
+            if not tag_path.exists():
+                raise FileNotFoundError(
+                    f"Tag names file not found: {tag_path}\n"
+                    "Run create_embeddings.py first to generate this file."
+                )
+
+            with open(tag_path, 'r', encoding='utf-8') as f:
+                self._zero_shot_tag_names = [line.strip() for line in f if line.strip()]
+
+            print(f"  Loaded {len(self._zero_shot_tag_names)} tag names")
+
+        return self._zero_shot_tag_names
+
+    def get_mtg_model(self):
+        """Load and cache MTG-Jamendo trained model."""
+        if self._mtg_model is None:
+            checkpoint_path = SCRIPT_DIR / config.mtg_checkpoint_path
+
+            if not checkpoint_path.exists():
+                raise FileNotFoundError(
+                    f"MTG-Jamendo checkpoint not found: {checkpoint_path}\n"
+                    "Train the model first using mtg_jamendo_training.py"
+                )
+
+            print("Loading MTG-Jamendo model...")
+            checkpoint = torch.load(str(checkpoint_path), map_location=config.device, weights_only=False)
+
+            # Get tag names (handle both old and new keys)
+            self._mtg_tag_names = checkpoint.get('tag_names', checkpoint.get('genre_names', []))
+
+            # Initialize and load model
+            model = GenreClassifier(
+                input_dim=config.mert_feature_dim,
+                num_classes=len(self._mtg_tag_names)
+            )
+            model.load_state_dict(checkpoint['model_state_dict'])
+            model = model.to(config.device).eval()
+
+            self._mtg_model = model
+            print(f"  Loaded model with {len(self._mtg_tag_names)} tags")
+            print(f"  Validation accuracy: {checkpoint.get('val_acc', 0)*100:.2f}%")
+
+        return self._mtg_model, self._mtg_tag_names
+
+    def get_mert_model(self):
+        """Load and cache MERT model for feature extraction."""
+        if self._mert_model is None:
+            try:
+                from transformers import Wav2Vec2FeatureExtractor, AutoModel
+            except ImportError as e:
+                raise ImportError(
+                    f"Missing transformers library: {e}\n"
+                    "Install with: pip install transformers"
+                )
+
+            print("Loading MERT model...")
+            self._mert_processor = Wav2Vec2FeatureExtractor.from_pretrained(config.mert_model_name)
+            self._mert_model = AutoModel.from_pretrained(config.mert_model_name, trust_remote_code=True)
+            self._mert_model = self._mert_model.to(config.device).eval()
+            print("  MERT model loaded!")
+
+        return self._mert_model, self._mert_processor
+
+
+# Global model cache
+model_cache = ModelCache()
+
+
+# ============================================================
+# Inference Functions
+# ============================================================
+@torch.no_grad()
+def tag_audio_zero_shot(audio_path: str, top_k: int = 20, normalization: str = "individual"):
+    """
+    Tag audio using Zero-Shot CLAP + MuLan approach.
+
+    Args:
+        audio_path: Path to audio file
+        top_k: Number of top tags to return
+        normalization: "mulan_only", "global", or "individual"
+
+    Returns:
+        List of (tag_name, confidence) tuples
+    """
+    # Load models and embeddings
+    mulan_model, clap_model, clap_processor = model_cache.get_zero_shot_models()
+    clap_embeddings, mulan_embeddings = model_cache.get_zero_shot_embeddings()
+    tag_names = model_cache.get_zero_shot_tag_names()
+
+    # Ensure embedding/tag count matches
+    min_len = min(len(tag_names), len(clap_embeddings), len(mulan_embeddings))
+    if len(tag_names) != min_len:
+        tag_names = tag_names[:min_len]
+        clap_embeddings = clap_embeddings[:min_len]
+        mulan_embeddings = mulan_embeddings[:min_len]
+
+    # Embed audio with MuLan (24kHz)
+    wav_mulan, _ = librosa.load(audio_path, sr=24000, mono=True)
+    wavs = torch.tensor(wav_mulan, dtype=torch.float32).unsqueeze(0).to(config.device)
+    mulan_audio_embed = mulan_model(wavs=wavs)
+
+    # Embed audio with CLAP (48kHz)
+    wav_clap, _ = librosa.load(audio_path, sr=48000, mono=True)
+    inputs = clap_processor(audio=wav_clap, sampling_rate=48000, return_tensors="pt").to(config.device)
+    clap_audio_embed = clap_model.get_audio_features(**inputs)
+
+    # Convert embeddings to tensors
+    mulan_text_e = torch.tensor(mulan_embeddings, dtype=torch.float32).to(config.device)
+    clap_text_e = torch.tensor(clap_embeddings, dtype=torch.float32).to(config.device)
+
+    # Calculate similarities
+    mulan_sims = F.cosine_similarity(mulan_audio_embed, mulan_text_e, dim=1)
+    clap_sims = F.cosine_similarity(clap_audio_embed, clap_text_e, dim=1)
+
+    # Apply normalization strategy
+    if normalization == "mulan_only":
+        combined = mulan_sims
+    elif normalization == "global":
+        all_sims = torch.cat([mulan_sims, clap_sims])
+        g_min, g_max = all_sims.min(), all_sims.max()
+        mulan_norm = (mulan_sims - g_min) / (g_max - g_min + 1e-8)
+        clap_norm = (clap_sims - g_min) / (g_max - g_min + 1e-8)
+        combined = 0.5 * mulan_norm + 0.5 * clap_norm
+    else:  # individual
+        mulan_norm = (mulan_sims - mulan_sims.min()) / (mulan_sims.max() - mulan_sims.min() + 1e-8)
+        clap_norm = (clap_sims - clap_sims.min()) / (clap_sims.max() - clap_sims.min() + 1e-8)
+        combined = 0.5 * mulan_norm + 0.5 * clap_norm
+
+    # Get top predictions
+    top_scores, top_idx = torch.topk(combined, k=min(top_k, len(tag_names)))
+
+    predictions = []
+    for i, idx in enumerate(top_idx):
+        tag = tag_names[idx.item()]
+        score = top_scores[i].item()
+        predictions.append((tag, score))
+
+    return predictions
+
+
+@torch.no_grad()
+def tag_audio_mtg_jamendo(audio_path: str, top_k: int = 10):
+    """
+    Tag audio using trained MTG-Jamendo MERT classifier.
+
+    Args:
+        audio_path: Path to audio file
+        top_k: Number of top tags to return
+
+    Returns:
+        List of (tag_name, confidence) tuples
+    """
+    # Load models
+    model, tag_names = model_cache.get_mtg_model()
+    mert_model, processor = model_cache.get_mert_model()
+
+    # Load and preprocess audio
+    wav, sr = librosa.load(audio_path, sr=config.mert_sample_rate, mono=True)
+
+    # Limit to max duration
+    max_samples = config.mert_sample_rate * config.max_duration
+    if len(wav) > max_samples:
+        wav = wav[:max_samples]
+
+    # Extract MERT features
+    inputs = processor(wav, sampling_rate=config.mert_sample_rate, return_tensors="pt")
+    inputs = {k: v.to(config.device) for k, v in inputs.items()}
+
+    outputs = mert_model(**inputs, output_hidden_states=True)
+    embeddings = outputs.hidden_states[config.mert_layer]
+    features = embeddings.mean(dim=1).squeeze(0)
+
+    # Predict
+    logits = model(features.unsqueeze(0))
+    probs = torch.softmax(logits, dim=1).squeeze(0).cpu().numpy()
+
+    # Get top predictions
+    top_indices = np.argsort(probs)[::-1][:top_k]
+    predictions = [(tag_names[i], float(probs[i])) for i in top_indices]
+
+    return predictions
+
+
+# ============================================================
+# Gradio Interface Functions
+# ============================================================
+def format_predictions(predictions: list, method: str) -> str:
+    """Format predictions as a readable string."""
+    if not predictions:
+        return "No predictions available."
+
+    lines = [f"## {method} Results\n"]
+
+    for i, (tag, score) in enumerate(predictions, 1):
+        # Create a visual bar
+        bar_length = int(score * 30)
+        bar = "█" * bar_length + "░" * (30 - bar_length)
+        lines.append(f"{i:2d}. **{tag}** — {score*100:.1f}% `{bar}`")
+
+    return "\n".join(lines)
+
+
+def analyze_audio(audio_file, method: str, normalization: str, top_k: int):
+    """
+    Main analysis function called by Gradio interface.
+
+    Args:
+        audio_file: Uploaded audio file path
+        method: "Zero-Shot (CLAP + MuLan)" or "MTG-Jamendo (MERT)"
+        normalization: Normalization strategy for Zero-Shot
+        top_k: Number of top tags to show
+
+    Returns:
+        Formatted prediction results
+    """
+    if audio_file is None:
+        return "Please upload an audio file."
+
+    try:
+        if method == "Zero-Shot (CLAP + MuLan)":
+            norm_map = {
+                "Individual (recommended)": "individual",
+                "Global": "global",
+                "MuLan Only": "mulan_only"
+            }
+            predictions = tag_audio_zero_shot(
+                audio_file,
+                top_k=int(top_k),
+                normalization=norm_map.get(normalization, "individual")
+            )
+            return format_predictions(predictions, "Zero-Shot Tagging")
+
+        elif method == "MTG-Jamendo (MERT)":
+            predictions = tag_audio_mtg_jamendo(audio_file, top_k=int(top_k))
+            return format_predictions(predictions, "MTG-Jamendo Tagging")
+
+        else:
+            return f"Unknown method: {method}"
+
+    except FileNotFoundError as e:
+        return f"**Error: Missing Required Files**\n\n{str(e)}"
+    except ImportError as e:
+        return f"**Error: Missing Dependencies**\n\n{str(e)}"
+    except Exception as e:
+        return f"**Error during analysis:**\n\n{str(e)}"
+
+
+def check_available_methods():
+    """Check which methods are available based on existing files."""
+    available = []
+    messages = []
+
+    # Check Zero-Shot files
+    clap_exists = (SCRIPT_DIR / config.clap_embeddings_path).exists()
+    mulan_exists = (SCRIPT_DIR / config.mulan_embeddings_path).exists()
+    tags_exists = (SCRIPT_DIR / config.tag_names_path).exists()
+
+    if clap_exists and mulan_exists and tags_exists:
+        available.append("Zero-Shot (CLAP + MuLan)")
+    else:
+        missing = []
+        if not clap_exists:
+            missing.append("clap_tag_embeddings.npy")
+        if not mulan_exists:
+            missing.append("mulan_tag_embeddings.npy")
+        if not tags_exists:
+            missing.append("musiccaps_tag_names.txt")
+        messages.append(f"Zero-Shot: Missing {', '.join(missing)}")
+
+    # Check MTG-Jamendo checkpoint
+    if (SCRIPT_DIR / config.mtg_checkpoint_path).exists():
+        available.append("MTG-Jamendo (MERT)")
+    else:
+        messages.append(f"MTG-Jamendo: Missing {config.mtg_checkpoint_path}")
+
+    return available, messages
+
+
+def create_interface():
+    """Create and configure the Gradio interface."""
+    available_methods, status_messages = check_available_methods()
+
+    # Build status message
+    if available_methods:
+        status = f"**Available methods:** {', '.join(available_methods)}"
+    else:
+        status = "**Warning:** No tagging methods available!"
+
+    if status_messages:
+        status += "\n\n**Missing files:**\n" + "\n".join(f"- {m}" for m in status_messages)
+
+    # Default method
+    default_method = available_methods[0] if available_methods else "Zero-Shot (CLAP + MuLan)"
+
+    with gr.Blocks(title="Music Tagger", theme=gr.themes.Soft()) as interface:
+        gr.Markdown("""
+        # Music Auto-Tagger
+
+        Upload a song to analyze it with AI-powered music tagging models.
+
+        **Two methods available:**
+        - **Zero-Shot (CLAP + MuLan):** Uses ~1,300 tags from MusicCaps without training
+        - **MTG-Jamendo (MERT):** Uses a trained classifier for genre/instrument/mood tags
+        """)
+
+        gr.Markdown(status)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                audio_input = gr.Audio(
+                    label="Upload Audio File",
+                    type="filepath",
+                    sources=["upload", "microphone"]
+                )
+
+                method_dropdown = gr.Dropdown(
+                    choices=["Zero-Shot (CLAP + MuLan)", "MTG-Jamendo (MERT)"],
+                    value=default_method,
+                    label="Tagging Method",
+                    info="Choose which model to use for tagging"
+                )
+
+                normalization_dropdown = gr.Dropdown(
+                    choices=["Individual (recommended)", "Global", "MuLan Only"],
+                    value="Individual (recommended)",
+                    label="Normalization (Zero-Shot only)",
+                    info="How to combine CLAP and MuLan scores",
+                    visible=True
+                )
+
+                top_k_slider = gr.Slider(
+                    minimum=5,
+                    maximum=50,
+                    value=15,
+                    step=5,
+                    label="Number of Tags",
+                    info="How many top tags to show"
+                )
+
+                analyze_btn = gr.Button("Analyze", variant="primary", size="lg")
+
+            with gr.Column(scale=2):
+                output_text = gr.Markdown(
+                    label="Results",
+                    value="Upload an audio file and click 'Analyze' to see predictions."
+                )
+
+        # Show/hide normalization based on method
+        def update_normalization_visibility(method):
+            return gr.update(visible=(method == "Zero-Shot (CLAP + MuLan)"))
+
+        method_dropdown.change(
+            fn=update_normalization_visibility,
+            inputs=[method_dropdown],
+            outputs=[normalization_dropdown]
+        )
+
+        # Analyze button
+        analyze_btn.click(
+            fn=analyze_audio,
+            inputs=[audio_input, method_dropdown, normalization_dropdown, top_k_slider],
+            outputs=[output_text]
+        )
+
+        gr.Markdown("""
+        ---
+        **Tips:**
+        - Supported formats: WAV, MP3, FLAC, OGG, and more
+        - Audio is automatically resampled to the required sample rate
+        - First analysis may take longer as models are loaded into memory
+        """)
+
+    return interface
+
+
+# ============================================================
+# Main Entry Point
+# ============================================================
+if __name__ == "__main__":
+    print("=" * 60)
+    print("Music Tagger GUI")
+    print("=" * 60)
+    print(f"Device: {config.device}")
+    print(f"Script directory: {SCRIPT_DIR}")
+
+    # Check available methods
+    available, messages = check_available_methods()
+    print(f"\nAvailable methods: {available if available else 'None'}")
+    if messages:
+        print("Status messages:")
+        for msg in messages:
+            print(f"  - {msg}")
+
+    # Launch interface
+    print("\nLaunching Gradio interface...")
+    interface = create_interface()
+    interface.launch(
+        share=False,  # Set to True to create a public link
+        server_name="0.0.0.0",  # Allow external connections
+        server_port=7860,
+        show_error=True
+    )