from suno_utils.worker.utils import retry_decorator from suno_utils.worker.settings import s3_client from suno_utils.tasks import ss_vad import tempfile import torch import torchaudio import librosa import os import re import numpy as np from dataclasses import dataclass from typing import Optional from suno_utils.audio import Audio from suno_utils.tasks.ditto_v2 import Ditto from suno_utils.worker.loader import _get_compatible_tokens from suno_utils.tasks.lyrics_alignment.shortest_path_aligner import ( ShortestPathAlignerConfig, ShortestPathAligner, ) from suno_utils.utils.clip import SunoClip from suno_utils.tasks.audio_features.instrument import InstrumentExtractor from suno_utils.tasks.audio_features.vocal import VocalExtractor from suno_utils.tasks.hoot import encode as hoot_encode from suno_utils.tasks.ditto_v2 import encode_overlap as encode, SAMPLE_RATE from suno_utils.utils.clip import SunoClip import numpy as np DITTO_EMBEDDING_DIM = 128 DITTO_SR = 24000 DITTO_LEN_S = 120 SILENCE_CUTOFF = 30 SEP_SAMPLE_RATE = 44100 DITTO_S3_PATH = "s3://suno-data/minz/models/ditto_v2_epoch_57.pt" S3_SAVE_PREFIX = "tasks/feature_eval/cover_persona/" retry_s3_upload = retry_decorator(3, wait_seconds=5)(s3_client.upload_file) TEST_SETS = { "all_genres": "../../task_eval/all_genre_tasks.json", "hard": "../../task_eval/hard_tasks.json", } @dataclass class Prompt: lyrics: str tags: str s3_id: str test_tags: Optional[str] = None test_lyrics: Optional[str] = None neg_tags: Optional[str] = None cover_arr: Optional[np.ndarray] = None artist_arr: Optional[np.ndarray] = None def load_prompt_arr(audio_id): with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file: target_version = "4.0" tokens = _get_compatible_tokens(temp_file.name, audio_id, target_version) return tokens[: (25 * 120)] class FeatureEval: def __init__(self, model_name): self.model_name = model_name self.ditto = Ditto( latent_dim=DITTO_EMBEDDING_DIM, model_path=DITTO_S3_PATH, is_flash=False, is_serving=True, ) def run_feature_eval(self, s3_id, audio, task, time_prefix): if isinstance(audio, Audio): with tempfile.NamedTemporaryFile(suffix=".mp3") as temp_file: audio.write_hq_mp3(temp_file.name) waveform, sr = torchaudio.load(temp_file.name) else: # file path waveform, sr = torchaudio.load(audio) waveform = torch.mean(waveform, dim=0).unsqueeze(0) if "vox" in task: # need to split stems if sr != SEP_SAMPLE_RATE: sr = SEP_SAMPLE_RATE resampler = torchaudio.transforms.Resample(sr, SEP_SAMPLE_RATE) vocals = ss_vad.encode(resampler(waveform)) vocals, _ = librosa.effects.trim(vocals, top_db=SILENCE_CUTOFF) if sr != DITTO_SR: resampler = torchaudio.transforms.Resample(sr, DITTO_SR) waveform = resampler(waveform)[:, : (DITTO_SR * DITTO_LEN_S)] ditto_embed = self.ditto.music_to_latent(waveform, task=task)[0].detach().cpu().numpy() ditto_data = { "embedding": ditto_embed, "sample_rate": DITTO_SR, "embed_dim": DITTO_EMBEDDING_DIM, "silence_cutoff_db": SILENCE_CUTOFF, "task": task, "model": self.model_name, } with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file: np.savez(temp_file.name, **ditto_data) retry_s3_upload( temp_file.name, "suno-data-uploads", os.path.join(os.path.join(S3_SAVE_PREFIX, time_prefix), f"{s3_id}_{task}_ditto.npz"), ) def extract_bracketed_content(text): # This regex matches [section_name] or [section_name:description] with optional whitespace pattern = r"\[\s*([^:\[\]]+)(?:\s*:\s*([^:\[\]]+))?\s*\]" # Find all matches matches = re.findall(pattern, text) # Process results into a more usable format results = [] for match in matches: # Each match is a tuple of (section_name, description) # If no description was found, the second element will be an empty string section_name = match[0].strip() description_part = match[1].strip() if match[1] else None results.append({"section_name": section_name, "description": description_part}) return results class InlineTagEval: def __init__(self, model_name, tokenizer, threshold=0.6, tag_type="instrument"): self.model_name = model_name good_spa_config = ShortestPathAlignerConfig( enable_jumps=True, logit_skip_coef=13.2, char_skip_coef=6.1, spelling_error_coef=2.1, section_skip_coef=27.2, line_skip_coef=34.0, char_epsilon=0.0007, silence_threshold_p=0.9, ) self.spa = ShortestPathAligner.from_sentencepiece(tokenizer._tokenizer, good_spa_config) self.tag_type = tag_type if tag_type == "instrument": self.ie = InstrumentExtractor() else: self.ie = VocalExtractor() print(f"Tag type: {tag_type}") self.threshold = 0.6 def run_feature_eval(self, s3_id, prior_text, group_mappings, time_prefix): clip = SunoClip(s3_id) vae_arr = clip.full_arr() example_audio = clip.audio() output = self.spa.align(prior_text, hoot_encode(example_audio, return_logits=True, batch_size=1)) section_timings = [] last_word_start_s = 0.0 for idx, row in enumerate(output): text = row["word"] headers = extract_bracketed_content(text) if len(headers) > 0: if idx > 0: last_word_start_s = output[idx - 1]["start_s"] for header in headers: section_timings.append( { "section_name": header["section_name"], "start_time_s": row["start_s"] if len(section_timings) > 0 else 0.0, "end_time_s": example_audio.duration_s, "description": header["description"], "last_start_s": last_word_start_s, } ) if len(section_timings) > 1: if row["start_s"] == section_timings[-2]["start_time_s"]: section_timings[-2]["start_time_s"] = section_timings[-2]["last_start_s"] section_timings[-2]["end_time_s"] = last_word_start_s last_word_start_s = row["start_s"] sections = [] for idx, row in enumerate(section_timings): from_token = int(row["start_time_s"] * 25) to_token = int(row["end_time_s"] * 25) name = row["section_name"] description = row["description"] # segment_audio = example_audio.get_segment(from_s=row["start_time_s"], to_s=row["end_time_s"]) data = vae_arr[from_token:to_token] if self.tag_type == "instrument": _, _, group_tags, instrument_tags = self.ie.extract(data, threshold=self.threshold) sections.append( { "name": name, "description": description, "found_groups": group_tags, "found_instruments": instrument_tags, } ) else: logits, gender = self.ie.extract(data, threshold=self.threshold) sections.append( {"name": name, "description": description, "logits": logits, "gender": gender} ) tag_data = { "section_labels": sections, "threshold": self.threshold, "group_mappings": group_mappings, "model": self.model_name, } with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file: np.savez(temp_file.name, **tag_data) retry_s3_upload( temp_file.name, "suno-data-uploads", os.path.join(os.path.join(S3_SAVE_PREFIX, time_prefix), f"{s3_id}_section_tags.npz"), ) class GenreEval: def __init__(self, model_name): self.model_name = model_name self.ditto = Ditto( latent_dim=DITTO_EMBEDDING_DIM, model_path=DITTO_S3_PATH, is_flash=False, is_serving=True, ) def run_feature_eval(self, s3_id, time_prefix): clip = SunoClip(s3_id) audio = clip.audio().convert(sample_rate=SAMPLE_RATE, n_channels=1, byte_width=2) encoding_self = np.mean(encode([audio], task="self_sim")[0], axis=0) encoding_genre = np.mean(encode([audio], task="genre_sim")[0], axis=0) ditto_data = { "embed_self": encoding_self, "embed_genre": encoding_genre, "model": self.model_name, } with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file: np.savez(temp_file.name, **ditto_data) retry_s3_upload( temp_file.name, "suno-data-uploads", os.path.join(os.path.join(S3_SAVE_PREFIX, time_prefix), f"{s3_id}_genre_sim_ditto.npz"), )