Spaces:
Running
Running
| """Automatic audio annotation for bulk dataset creation. | |
| Two tiers: | |
| - basic: librosa-only DSP (tempo, key). No downloads. CPU. ~instant per file. | |
| - rich: basic + LAION-CLAP zero-shot tagging (genre, mood, instrument). | |
| Lazy-loaded; downloads ~2.35 GB checkpoint on first use. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import logging | |
| import os | |
| import threading | |
| from pathlib import Path | |
| from typing import Any, Callable, Dict, List, Optional | |
| logger = logging.getLogger(__name__) | |
| AUDIO_EXTENSIONS = (".wav", ".mp3", ".flac", ".m4a", ".ogg", ".aac") | |
| CLAP_CKPT_FILENAME = "music_audioset_epoch_15_esc_90.14.pt" | |
| CLAP_REPO = "lukewys/laion_clap" | |
| # Text-side dependencies laion_clap pulls from HF on construction. | |
| # We stage these into models/pretrained/clap/hub/ so the rich tier is | |
| # fully offline after a single download and nothing leaks to ~/.cache. | |
| # | |
| # These are laion_clap's own hardcoded ids (see clap_module/model.py), and two | |
| # of them are legacy aliases: `roberta-base` now redirects to | |
| # `FacebookAI/roberta-base`, `bert-base-uncased` to | |
| # `google-bert/bert-base-uncased`. They MUST stay spelled the legacy way — the | |
| # cache dir is named after the id, and `from_pretrained('roberta-base')` only | |
| # finds a pre-staged snapshot under `models--roberta-base`. Downloading them | |
| # under the canonical name would leave laion_clap fetching them again at | |
| # construction, defeating the offline staging. | |
| CLAP_TEXT_DEPS = ("roberta-base", "bert-base-uncased", "facebook/bart-base") | |
| # Whole-clip CLAP embedding windowing. laion_clap (enable_fusion=False) only | |
| # looks at max_len=480000 samples (10 s @ 48 kHz) and picks a RANDOM offset | |
| # for longer clips — non-deterministic tags and a single arbitrary window. | |
| # We instead embed non-overlapping 10 s windows and mean-pool, which is | |
| # deterministic (rand_trunc is a no-op on exact-length windows) and hears the | |
| # whole clip. CLAP_MAX_WINDOWS caps compute (12 ≈ 120 s/clip); a trailing | |
| # window shorter than CLAP_MIN_TAIL is dropped so near-silent tails don't | |
| # skew the pooled embedding. | |
| CLAP_SR = 48000 | |
| CLAP_WIN = 480000 | |
| CLAP_MAX_WINDOWS = 12 | |
| CLAP_MIN_TAIL = CLAP_WIN // 2 | |
| # Confidence gates: a tag is only emitted when the top-1 cosine similarity | |
| # clears the per-group floor, and (genre/mood) beats the runner-up by | |
| # CLAP_MARGIN — near-ties mean "model can't tell", not a coin-flip. Missing | |
| # keys are intentional: apply_template() drops the segment and | |
| # _compose_prompt() guards each field, same as a missing BPM/key. | |
| # Starting points — calibrate against a hand-labeled set if tags get sparse. | |
| # Measured on the bundled demo clips: real-music top-1 sims run ~0.13 (mood) | |
| # to ~0.22 (genre) with top1-top2 margins of 0.012-0.034, while an SFX-style | |
| # impulse clip sits at 0.04-0.09 across all groups. A margin of 0.02 would | |
| # have dropped genre/mood from real music, hence 0.01. Caveat: stationary | |
| # noise legitimately scores ~0.18 as "experimental", too close to real music | |
| # to gate on tau alone. | |
| CLAP_TAU = {"genre": 0.10, "mood": 0.10, "instruments": 0.12} | |
| CLAP_MARGIN = 0.01 | |
| KEY_NAMES_SHARP = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"] | |
| KEY_NAMES_FLAT = ["C", "Db", "D", "Eb", "E", "F", "Gb", "G", "Ab", "A", "Bb", "B"] | |
| # Krumhansl-Schmuckler key profiles. | |
| KRUMHANSL_MAJOR = [6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88] | |
| KRUMHANSL_MINOR = [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17] | |
| def _iter_audio_files(folder: Path) -> List[Path]: | |
| results: List[Path] = [] | |
| for root, _, files in os.walk(folder): | |
| for name in files: | |
| if name.startswith("."): | |
| continue | |
| if name.lower().endswith(AUDIO_EXTENSIONS): | |
| results.append(Path(root) / name) | |
| results.sort() | |
| return results | |
| def _estimate_tempo(y, sr) -> Optional[int]: | |
| import librosa | |
| import numpy as np | |
| try: | |
| tempo, _ = librosa.beat.beat_track(y=y, sr=sr) | |
| # librosa 0.10+ returns tempo as np.ndarray (shape (1,) typically). | |
| # numpy 2.x removed implicit float() conversion of N-d arrays — | |
| # `float(np.array([120.]))` now raises TypeError instead of returning | |
| # 120.0 like numpy 1.x did. Normalize via .flat[0] which handles | |
| # scalar, 0-d, 1-d, and N-d uniformly. | |
| arr = np.atleast_1d(np.asarray(tempo)) | |
| if arr.size == 0: | |
| return None | |
| bpm = float(arr.flat[0]) | |
| if bpm <= 0: | |
| return None | |
| return int(round(bpm)) | |
| except Exception as exc: | |
| logger.debug("tempo estimation failed: %s", exc) | |
| return None | |
| def _estimate_brightness(y, sr) -> Optional[str]: | |
| import librosa | |
| try: | |
| centroid = float(librosa.feature.spectral_centroid(y=y, sr=sr).mean()) | |
| except Exception as exc: | |
| logger.debug("centroid estimation failed: %s", exc) | |
| return None | |
| if centroid <= 0: | |
| return None | |
| if centroid < 1500: | |
| return "dark" | |
| if centroid > 3500: | |
| return "bright" | |
| return None | |
| def _estimate_character(y, sr) -> Optional[str]: | |
| import librosa | |
| import numpy as np | |
| try: | |
| harm, perc = librosa.effects.hpss(y) | |
| eh = float(np.mean(harm ** 2)) | |
| ep = float(np.mean(perc ** 2)) | |
| except Exception as exc: | |
| logger.debug("HPSS failed: %s", exc) | |
| return None | |
| total = eh + ep | |
| if total <= 0: | |
| return None | |
| perc_ratio = ep / total | |
| if perc_ratio > 0.65: | |
| return "percussion-driven" | |
| if perc_ratio < 0.20: | |
| return "melodic" | |
| return None | |
| def _estimate_key(y, sr) -> Optional[str]: | |
| import librosa | |
| import numpy as np | |
| try: | |
| chroma = librosa.feature.chroma_cqt(y=y, sr=sr) | |
| chroma_mean = chroma.mean(axis=1) | |
| if chroma_mean.sum() <= 0: | |
| return None | |
| chroma_mean = chroma_mean / chroma_mean.sum() | |
| major = np.asarray(KRUMHANSL_MAJOR) | |
| minor = np.asarray(KRUMHANSL_MINOR) | |
| best_score = -1.0 | |
| best_key = None | |
| for i in range(12): | |
| maj_score = float(np.corrcoef(chroma_mean, np.roll(major, i))[0, 1]) | |
| min_score = float(np.corrcoef(chroma_mean, np.roll(minor, i))[0, 1]) | |
| if maj_score > best_score: | |
| best_score = maj_score | |
| best_key = f"{KEY_NAMES_SHARP[i]} major" | |
| if min_score > best_score: | |
| best_score = min_score | |
| best_key = f"{KEY_NAMES_SHARP[i]} minor" | |
| return best_key | |
| except Exception as exc: | |
| logger.debug("key estimation failed: %s", exc) | |
| return None | |
| def _compose_prompt(parts: Dict[str, Any]) -> str: | |
| genre = parts.get("genre") | |
| mood = parts.get("mood") | |
| instruments = parts.get("instruments") or [] | |
| bpm = parts.get("bpm") | |
| key = parts.get("key") | |
| brightness = parts.get("brightness") | |
| character = parts.get("character") | |
| head_bits: List[str] = [] | |
| if mood: | |
| head_bits.append(str(mood)) | |
| if genre: | |
| head_bits.append(f"{genre} track") | |
| elif head_bits: | |
| head_bits[-1] = f"{head_bits[-1]} track" | |
| opening = " ".join(head_bits) | |
| descriptors = [d for d in (brightness, character) if d] | |
| fragments: List[str] = [] | |
| if opening: | |
| fragments.append(opening) | |
| if descriptors: | |
| fragments.append(", ".join(descriptors)) | |
| if bpm: | |
| fragments.append(f"{bpm} BPM") | |
| if key: | |
| fragments.append(f"in {key}") | |
| if instruments: | |
| fragments.append("with " + ", ".join(instruments)) | |
| out = ", ".join(fragments) | |
| return out[:1].upper() + out[1:] if out else "" | |
| # Serializes checkpoint loads that need relaxed unpickling. The CLAP loader | |
| # below holds this while it reads the (trusted) .pt; any other model-load site | |
| # that has to allow weights_only=False should acquire the same lock so the | |
| # windows never overlap. It is module-level (process-wide) on purpose. | |
| CHECKPOINT_LOAD_LOCK = threading.Lock() | |
| def _clap_load_state_dict_trusted(checkpoint_path: str, map_location: str = "cpu", | |
| skip_params: bool = True): | |
| """Vendored copy of laion_clap.clap_module.factory.load_state_dict. | |
| We replicate it here so weights_only=False is passed to torch.load for THIS | |
| call only. The previous approach swapped torch.load process-wide while CLAP | |
| loaded, which silently leaked weights_only=False into any concurrent | |
| torch.load on another thread (e.g. a generation/LoRA model load) — an unsafe | |
| unpickling foot-gun. Loading the state dict directly keeps the relaxation | |
| local to this one trusted checkpoint and touches no global state. | |
| """ | |
| import torch | |
| checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False) | |
| if isinstance(checkpoint, dict) and "state_dict" in checkpoint: | |
| state_dict = checkpoint["state_dict"] | |
| else: | |
| state_dict = checkpoint | |
| if skip_params and state_dict: | |
| if next(iter(state_dict.items()))[0].startswith("module"): | |
| state_dict = {k[7:]: v for k, v in state_dict.items()} | |
| # Newer transformers dropped the roberta position_ids buffer; remove it | |
| # so the strict-ish load doesn't choke. pop() is a no-op if absent. | |
| try: | |
| from packaging import version | |
| import transformers | |
| if version.parse(transformers.__version__) >= version.parse("4.31.0"): | |
| state_dict.pop("text_branch.embeddings.position_ids", None) | |
| except Exception: | |
| state_dict.pop("text_branch.embeddings.position_ids", None) | |
| return state_dict | |
| class _ClapTagger: | |
| """Lazy holder for a LAION-CLAP model used for zero-shot tagging.""" | |
| def __init__(self, ckpt_path: Path): | |
| self.ckpt_path = ckpt_path | |
| self._model = None | |
| self._lock = threading.Lock() | |
| self._label_embeds: Dict[str, Any] = {} | |
| def ensure_loaded(self): | |
| if self._model is not None: | |
| return | |
| with self._lock: | |
| if self._model is not None: | |
| return | |
| if not self.ckpt_path.exists(): | |
| raise FileNotFoundError( | |
| f"CLAP checkpoint not found at {self.ckpt_path}. " | |
| "Download it first via /api/bulk-annotate/download-clap." | |
| ) | |
| logging.getLogger("transformers").setLevel(logging.ERROR) | |
| # Point HF resolution at our project-local cache and disable the | |
| # HEAD-revalidation traffic. After download_clap_checkpoint() has | |
| # staged the text deps under <pretrained>/clap/hub/, CLAP_Module | |
| # loads them offline with zero HF hub requests. | |
| # | |
| # Two reasons env vars alone aren't enough: | |
| # 1. huggingface_hub.constants.HF_HUB_OFFLINE is captured at | |
| # module-import time (constants.py:185). model_manager.py | |
| # imports huggingface_hub at app startup, so the constant is | |
| # already False by the time we set the env var here. | |
| # transformers.utils.hub.is_offline_mode reads that same | |
| # constant — patching the attribute makes both libraries see | |
| # offline mode. | |
| # 2. laion_clap/training/data.py:44-46 runs three from_pretrained | |
| # calls at MODULE LEVEL — those fire the first time we do | |
| # `import laion_clap` and predate any patch we do after the | |
| # import. So we patch BEFORE the import, not after. | |
| hub_dir = self.ckpt_path.parent / "hub" | |
| env_keys = ("HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE", "TRANSFORMERS_CACHE", | |
| "HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") | |
| prev_env = {k: os.environ.get(k) for k in env_keys} | |
| os.environ["HF_HUB_CACHE"] = str(hub_dir) | |
| os.environ["HUGGINGFACE_HUB_CACHE"] = str(hub_dir) | |
| os.environ["TRANSFORMERS_CACHE"] = str(hub_dir) | |
| os.environ["HF_HUB_OFFLINE"] = "1" | |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" | |
| import huggingface_hub.constants as _hhc | |
| prev_offline_attr = _hhc.HF_HUB_OFFLINE | |
| prev_cache_attr = _hhc.HF_HUB_CACHE | |
| _hhc.HF_HUB_OFFLINE = True | |
| # HF_HUB_CACHE is captured at import time too (same reason as | |
| # HF_HUB_OFFLINE), so the env vars above don't actually redirect it. | |
| # Patch the constant directly — otherwise transformers falls back to | |
| # the default ~/.cache/huggingface, which may hold only the tokenizer | |
| # (not config.json / weights) and then fails in offline mode. | |
| _hhc.HF_HUB_CACHE = str(hub_dir) | |
| try: | |
| import laion_clap # noqa: E402 — must follow the offline patch | |
| import torch | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = laion_clap.CLAP_Module(enable_fusion=False, amodel="HTSAT-base", device=device) | |
| finally: | |
| _hhc.HF_HUB_OFFLINE = prev_offline_attr | |
| _hhc.HF_HUB_CACHE = prev_cache_attr | |
| for k, v in prev_env.items(): | |
| if v is None: | |
| os.environ.pop(k, None) | |
| else: | |
| os.environ[k] = v | |
| # torch >= 2.6 flipped torch.load(weights_only=True) and newer | |
| # transformers dropped the roberta position_ids buffer, so | |
| # laion_clap's own load_ckpt errors twice: unpickling, then strict | |
| # state_dict mismatch. We load the (trusted) checkpoint via a | |
| # vendored helper that passes weights_only=False for this call only, | |
| # under a shared lock — no process-wide torch.load patch, so a | |
| # concurrent generation load can never inherit unsafe unpickling. | |
| with CHECKPOINT_LOAD_LOCK: | |
| state = _clap_load_state_dict_trusted(str(self.ckpt_path), skip_params=True) | |
| missing, unexpected = model.model.load_state_dict(state, strict=False) | |
| if unexpected: | |
| logger.debug("CLAP unexpected keys ignored: %s", unexpected[:5]) | |
| if missing: | |
| logger.debug("CLAP missing keys: %s", missing[:5]) | |
| self._model = model | |
| self._device = device | |
| logger.info("CLAP loaded on %s from %s", device, self.ckpt_path) | |
| def _embed_labels(self, group: str, prompts: List[str]): | |
| import torch | |
| key = f"{group}:{'|'.join(prompts)}" | |
| if key in self._label_embeds: | |
| return self._label_embeds[key] | |
| with torch.no_grad(): | |
| embed = self._model.get_text_embedding(prompts, use_tensor=True) | |
| embed = embed / embed.norm(dim=-1, keepdim=True).clamp_min(1e-8) | |
| self._label_embeds[key] = embed | |
| return embed | |
| def _embed_audio(self, audio_path: Path): | |
| """Deterministic whole-clip embedding: mean-pool over 10 s windows. | |
| Each window is exactly CLAP_WIN samples, so laion_clap's rand_trunc | |
| path never fires (overflow=0) and repeated runs produce identical | |
| embeddings. If laion_clap is ever run with enable_fusion=True this | |
| windowing must be revisited. | |
| """ | |
| import librosa | |
| import numpy as np | |
| import torch | |
| # Decode only what the windowing below can use (CLAP_MAX_WINDOWS x | |
| # 10 s). Without the cap a 2-hour file decodes fully into RAM | |
| # (~1.3 GB of float32 at 48 kHz) just to embed its first 2 minutes. | |
| y, _ = librosa.load(str(audio_path), sr=CLAP_SR, mono=True, | |
| duration=CLAP_MAX_WINDOWS * 10) | |
| if y.size == 0: | |
| raise ValueError("empty audio") | |
| # Non-overlapping 10 s windows. Pad a short *first/only* window; | |
| # drop a short trailing window so near-silent tails don't skew the mean. | |
| windows = [] | |
| for start in range(0, y.size, CLAP_WIN): | |
| chunk = y[start:start + CLAP_WIN] | |
| if chunk.size < CLAP_WIN: | |
| if windows and chunk.size < CLAP_MIN_TAIL: | |
| break # negligible tail — ignore it | |
| chunk = np.pad(chunk, (0, CLAP_WIN - chunk.size)) | |
| windows.append(chunk) | |
| if len(windows) >= CLAP_MAX_WINDOWS: | |
| break | |
| batch = torch.from_numpy(np.stack(windows)).float() | |
| with torch.no_grad(): | |
| embs = self._model.get_audio_embedding_from_data(x=batch, use_tensor=True) | |
| # Mean-pool across windows, then L2-normalize once. | |
| emb = embs.mean(dim=0, keepdim=True) | |
| return emb / emb.norm(dim=-1, keepdim=True).clamp_min(1e-8) | |
| def tag(self, audio_path: Path, label_sets: Dict[str, List[str]], top_k_instruments: int = 2) -> Dict[str, Any]: | |
| self.ensure_loaded() | |
| import torch | |
| audio_embed = self._embed_audio(audio_path) | |
| out: Dict[str, Any] = {} | |
| for group in ("genre", "mood"): | |
| labels = label_sets.get(group) or [] | |
| if not labels: | |
| continue | |
| prompts = [f"a {lab} music track" if group == "genre" else f"a {lab} sounding music track" for lab in labels] | |
| text_embed = self._embed_labels(group, prompts) | |
| sims = (audio_embed @ text_embed.T).squeeze(0) | |
| if sims.numel() >= 2: | |
| top2 = torch.topk(sims, k=2) | |
| top_val, top_idx = float(top2.values[0]), int(top2.indices[0]) | |
| margin = top_val - float(top2.values[1]) | |
| else: | |
| top_val, top_idx, margin = float(sims[0]), 0, float("inf") | |
| logger.debug( | |
| "CLAP %s for %s: top=%r sim=%.4f margin=%.4f", | |
| group, audio_path.name, labels[top_idx], top_val, margin, | |
| ) | |
| if top_val >= CLAP_TAU[group] and margin >= CLAP_MARGIN: | |
| out[group] = labels[top_idx] | |
| # else: omit — the prompt template drops the segment | |
| instruments = label_sets.get("instruments") or [] | |
| if instruments: | |
| prompts = [f"music featuring {lab}" for lab in instruments] | |
| text_embed = self._embed_labels("instruments", prompts) | |
| sims = (audio_embed @ text_embed.T).squeeze(0) | |
| k = min(top_k_instruments, len(instruments)) | |
| top = torch.topk(sims, k=k) | |
| logger.debug( | |
| "CLAP instruments for %s: %s", | |
| audio_path.name, | |
| [(instruments[i], round(v, 4)) for v, i in zip(top.values.tolist(), top.indices.tolist())], | |
| ) | |
| kept = [instruments[i] for v, i in zip(top.values.tolist(), top.indices.tolist()) | |
| if v >= CLAP_TAU["instruments"]] | |
| if kept: # only set the key if at least one passes | |
| out["instruments"] = kept | |
| return out | |
| _clap_tagger_singleton: Optional[_ClapTagger] = None | |
| _clap_tagger_lock = threading.Lock() | |
| def get_clap_tagger(clap_ckpt_path: Path) -> _ClapTagger: | |
| global _clap_tagger_singleton | |
| with _clap_tagger_lock: | |
| if _clap_tagger_singleton is None or _clap_tagger_singleton.ckpt_path != clap_ckpt_path: | |
| _clap_tagger_singleton = _ClapTagger(clap_ckpt_path) | |
| return _clap_tagger_singleton | |
| def clap_checkpoint_path(models_pretrained_dir: Path) -> Path: | |
| return models_pretrained_dir / "clap" / CLAP_CKPT_FILENAME | |
| def clap_hub_dir(models_pretrained_dir: Path) -> Path: | |
| """HF cache for laion_clap's text-side deps. Sibling of the .pt.""" | |
| return models_pretrained_dir / "clap" / "hub" | |
| def clap_checkpoint_available(models_pretrained_dir: Path) -> bool: | |
| return clap_checkpoint_path(models_pretrained_dir).exists() | |
| def _text_dep_snapshot_present(hub_dir: Path, repo_id: str) -> bool: | |
| safe = "models--" + repo_id.replace("/", "--") | |
| snap_root = hub_dir / safe / "snapshots" | |
| if not snap_root.exists(): | |
| return False | |
| return any(snap_root.iterdir()) | |
| def download_clap_checkpoint( | |
| models_pretrained_dir: Path, | |
| progress_cb: Optional[Callable[[str], None]] = None, | |
| phase_cb: Optional[Callable[[int, int, str], None]] = None, | |
| ) -> Path: | |
| """Download the CLAP audio .pt plus laion_clap's text-side HF snapshots. | |
| Four sequential phases — emit a phase update (current, total, label) at the | |
| start of each so a multi-phase progress UI can show real context. Skips | |
| phases whose artifacts are already on disk. | |
| `progress_cb` (str-only) is kept for the bulk-annotate API. | |
| `phase_cb` (current, total, label) is the structured channel. | |
| """ | |
| target = clap_checkpoint_path(models_pretrained_dir) | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| hub_dir = clap_hub_dir(models_pretrained_dir) | |
| hub_dir.mkdir(parents=True, exist_ok=True) | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| import os | |
| from utils.hf_transfer import xet_disabled | |
| total_phases = 1 + len(CLAP_TEXT_DEPS) | |
| def _emit(phase_index: int, label: str) -> None: | |
| if phase_cb: | |
| phase_cb(phase_index, total_phases, label) | |
| if progress_cb: | |
| progress_cb(f"[{phase_index}/{total_phases}] {label}") | |
| if not target.exists(): | |
| _emit(1, "CLAP audio model (~2.35 GB)") | |
| # Use custom CLAP from fragmenta-models on HF Spaces | |
| use_custom_repo = os.getenv('FRAGMENTA_USE_CUSTOM_MODELS', '').lower() == 'true' | |
| if use_custom_repo: | |
| repo_id = "MazCodes/fragmenta-models" | |
| else: | |
| repo_id = CLAP_REPO | |
| downloaded = hf_hub_download( | |
| repo_id=repo_id, | |
| filename=CLAP_CKPT_FILENAME, | |
| local_dir=str(target.parent), | |
| ) | |
| downloaded_path = Path(downloaded) | |
| if downloaded_path != target: | |
| try: | |
| downloaded_path.replace(target) | |
| except OSError: | |
| import shutil | |
| shutil.copy2(downloaded_path, target) | |
| # laion_clap's CLAP_Module(...) constructor instantiates a Roberta text | |
| # branch plus bert/bart tokenizers at import time. Pre-stage them into | |
| # our own cache so the rich tier is fully offline after this step. | |
| # safetensors only — pytorch_model.bin is a redundant copy. | |
| # Forced onto the classic HTTP transfer: hf_xet asks for a read token at | |
| # /api/models/<repo>/xet-read-token/<sha> without following the 307 that | |
| # maps a legacy alias to its canonical id, so `roberta-base` 404s there and | |
| # takes the whole CLAP install down with it. The regular transfer follows | |
| # the redirect. These are small files, so Xet buys nothing here anyway. | |
| with xet_disabled(): | |
| for i, repo_id in enumerate(CLAP_TEXT_DEPS, start=2): | |
| if _text_dep_snapshot_present(hub_dir, repo_id): | |
| continue | |
| _emit(i, f"Text encoder: {repo_id}") | |
| snapshot_download( | |
| repo_id=repo_id, | |
| cache_dir=str(hub_dir), | |
| allow_patterns=[ | |
| "config.json", | |
| "tokenizer*", | |
| "vocab*", | |
| "merges.txt", | |
| "special_tokens_map.json", | |
| "model.safetensors", | |
| ], | |
| ) | |
| return target | |
| def load_label_sets(label_sets_path: Optional[Path]) -> Dict[str, List[str]]: | |
| if not label_sets_path or not label_sets_path.exists(): | |
| return {"genre": [], "mood": [], "instruments": []} | |
| with open(label_sets_path, "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| return { | |
| "genre": list(data.get("genre") or []), | |
| "mood": list(data.get("mood") or []), | |
| "instruments": list(data.get("instruments") or []), | |
| } | |
| def annotate_file( | |
| audio_path: Path, | |
| tier: str, | |
| clap_tagger: Optional[_ClapTagger], | |
| label_sets: Dict[str, List[str]], | |
| sr: int = 22050, | |
| max_seconds: float = 60.0, | |
| prompt_template: Optional[str] = None, | |
| ) -> Dict[str, Any]: | |
| import librosa | |
| import warnings | |
| parts: Dict[str, Any] = {} | |
| try: | |
| y, loaded_sr = librosa.load(str(audio_path), sr=sr, mono=True, duration=max_seconds) | |
| except Exception as exc: | |
| logger.warning("librosa failed to load %s: %s", audio_path.name, exc) | |
| return { | |
| "file_name": audio_path.name, | |
| "prompt": "", | |
| "path": str(audio_path), | |
| "error": f"load failed: {exc}", | |
| } | |
| # Silent / harmonically flat clips trip librosa's "Trying to estimate | |
| # tuning from empty frequency set" warning during chroma extraction. | |
| # The warning is benign — the analysis returns sensible defaults — but | |
| # it spams stderr on every silent file, so we mute it here. | |
| with warnings.catch_warnings(): | |
| warnings.filterwarnings( | |
| "ignore", | |
| message="Trying to estimate tuning from empty frequency set", | |
| ) | |
| parts["bpm"] = _estimate_tempo(y, loaded_sr) | |
| parts["key"] = _estimate_key(y, loaded_sr) | |
| parts["brightness"] = _estimate_brightness(y, loaded_sr) | |
| parts["character"] = _estimate_character(y, loaded_sr) | |
| if tier == "rich" and clap_tagger is not None: | |
| try: | |
| tags = clap_tagger.tag(audio_path, label_sets) | |
| parts.update(tags) | |
| except Exception as exc: | |
| logger.warning("CLAP tagging failed for %s: %s", audio_path.name, exc) | |
| # Template-driven prompt assembly. Falls back to the legacy descriptive | |
| # prose if no template is supplied (call sites that haven't been | |
| # threaded with project metadata yet). | |
| if prompt_template is not None and prompt_template.strip(): | |
| from app.backend.data.projects import apply_template | |
| prompt = apply_template(prompt_template, parts) | |
| else: | |
| prompt = _compose_prompt(parts) | |
| return { | |
| "file_name": audio_path.name, | |
| "prompt": prompt, | |
| "path": str(audio_path), | |
| "attributes": parts, | |
| } | |
| def annotate_folder( | |
| folder: Path, | |
| tier: str, | |
| label_sets: Dict[str, List[str]], | |
| clap_ckpt_path: Optional[Path] = None, | |
| progress_cb: Optional[Callable[[int, int, str], None]] = None, | |
| ) -> List[Dict[str, Any]]: | |
| folder = Path(folder) | |
| if not folder.exists() or not folder.is_dir(): | |
| raise ValueError(f"Folder not found: {folder}") | |
| files = _iter_audio_files(folder) | |
| if not files: | |
| raise ValueError(f"No audio files found in {folder}") | |
| clap_tagger: Optional[_ClapTagger] = None | |
| if tier == "rich": | |
| if not clap_ckpt_path or not Path(clap_ckpt_path).exists(): | |
| raise FileNotFoundError( | |
| "Rich tier requires the CLAP checkpoint; download it first." | |
| ) | |
| clap_tagger = get_clap_tagger(Path(clap_ckpt_path)) | |
| clap_tagger.ensure_loaded() | |
| results: List[Dict[str, Any]] = [] | |
| total = len(files) | |
| for i, audio_path in enumerate(files, start=1): | |
| if progress_cb: | |
| progress_cb(i, total, audio_path.name) | |
| entry = annotate_file(audio_path, tier, clap_tagger, label_sets) | |
| results.append(entry) | |
| return results | |
| def unload_clap(): | |
| """Free CLAP weights from VRAM. Call before training starts.""" | |
| global _clap_tagger_singleton | |
| with _clap_tagger_lock: | |
| if _clap_tagger_singleton is not None: | |
| _clap_tagger_singleton._model = None | |
| _clap_tagger_singleton._label_embeds = {} | |
| _clap_tagger_singleton = None | |
| try: | |
| import torch | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| except Exception: | |
| pass | |