from torch.utils.data import Dataset from beartype.typing import Sequence, Callable, Optional, Dict, Tuple, List, Union from beartype import beartype from beartype.door import is_bearable import random import pandas as pd import os from torchaudio.functional import resample import torch import typing as tp from pathlib import Path import torchaudio as ta import torch.nn.functional as F import numpy as np import json import yaml import torchaudio import math import re from loguru import logger import ffmpeg class Read_and_PadCrop_Normalized_T(torch.nn.Module): def __init__(self, n_samples: int, sample_rate: int, randomize: bool = True): super().__init__() self.n_samples = n_samples self.sample_rate = sample_rate self.randomize = randomize def __call__(self, filename: str, duration: float, cur_sample_rate: int) -> Tuple[torch.Tensor, float, float, int, int]: if self.n_samples < 0: #means not clip chunk, _ = torchaudio.load(filename, frame_offset=0, num_frames=-1) t_start = 0. t_end = 1.0 offset = 0 else: if(duration<(float(self.n_samples)/self.sample_rate+1)): # print(duration,(float(self.n_samples)/self.sample_rate+1)) chunk, _ = torchaudio.load(filename, frame_offset=0, num_frames=-1) t_start = 0. t_end = min(1.0, float(self.n_samples) / float(self.sample_rate) / duration) offset = 0 # print('c1:',chunk.shape) else: offset = np.random.randint(0,int(duration*cur_sample_rate)-int(float(self.n_samples)/self.sample_rate*cur_sample_rate)) t_start = offset / float(cur_sample_rate) / duration t_end = t_start + float(self.n_samples) / float(self.sample_rate) / duration chunk, _ = torchaudio.load(filename, frame_offset=offset, num_frames=int(float(self.n_samples)/self.sample_rate*cur_sample_rate)) # print('offset:',offset) # print('c0:',chunk.shape) # Pad with silence if necessary. if(chunk.shape[0]>1): chunk = chunk[torch.randint(chunk.shape[0], size=(1,)),:].float() else: chunk = chunk[[0],:].float() if(cur_sample_rate!=self.sample_rate): # print('a:',cur_sample_rate,chunk.shape) chunk = torchaudio.functional.resample(chunk, cur_sample_rate, self.sample_rate) # print('b:',self.sample_rate,chunk.shape) if self.n_samples > 0: if chunk.shape[-1] < self.n_samples: chunk = torch.cat([chunk, torch.zeros((1, self.n_samples - chunk.shape[-1],))],-1) else: chunk = chunk[:,0:self.n_samples] seconds_start = math.floor(offset / cur_sample_rate) seconds_total = math.floor(duration) return ( chunk, t_start, t_end, seconds_start, seconds_total ) class Read_and_PadCrop_Normalized_T_Avoid_Watermark(torch.nn.Module): def __init__(self, n_samples: int, sample_rate: int, randomize: bool = True, w_start = 0, w_interval = 11.3): super().__init__() self.n_samples = n_samples self.sample_rate = sample_rate self.randomize = randomize self.w_start = w_start self.w_interval = w_interval def __call__(self, filename: str, duration: float, cur_sample_rate: int) -> Tuple[torch.Tensor, float, float, int, int]: if self.n_samples < 0: #means not clip chunk, _ = torchaudio.load(filename, frame_offset=0, num_frames=-1) t_start = 0. t_end = 1.0 offset = 0 else: if(duration<(float(self.n_samples)/self.sample_rate+1)): # print(duration,(float(self.n_samples)/self.sample_rate+1)) chunk, _ = torchaudio.load(filename, frame_offset=0, num_frames=-1) t_start = 0. t_end = min(1.0, float(self.n_samples) / float(self.sample_rate) / duration) offset = 0 # print('c1:',chunk.shape) else: n_offset_option = (duration - self.w_start) // self.w_interval if n_offset_option <= 1: offset = 0 else: offset = int((random.randint(0,n_offset_option-1) * self.w_interval + self.w_start) * cur_sample_rate) # offset = np.random.randint(0,int(duration*cur_sample_rate)-int(float(self.n_samples)/self.sample_rate*cur_sample_rate)) t_start = offset / float(cur_sample_rate) / duration t_end = t_start + float(self.n_samples) / float(self.sample_rate) / duration chunk, _ = torchaudio.load(filename, frame_offset=offset, num_frames=int(float(self.n_samples)/self.sample_rate*cur_sample_rate)) # print('offset:',offset) # print('c0:',chunk.shape) # Pad with silence if necessary. if(chunk.shape[0]>1): chunk = chunk[torch.randint(chunk.shape[0], size=(1,)),:].float() else: chunk = chunk[[0],:].float() if(cur_sample_rate!=self.sample_rate): # print('a:',cur_sample_rate,chunk.shape) chunk = torchaudio.functional.resample(chunk, cur_sample_rate, self.sample_rate) # print('b:',self.sample_rate,chunk.shape) if self.n_samples > 0: if chunk.shape[-1] < self.n_samples: chunk = torch.cat([chunk, torch.zeros((1, self.n_samples - chunk.shape[-1],))],-1) else: chunk = chunk[:,0:self.n_samples] seconds_start = math.floor(offset / cur_sample_rate) seconds_total = math.floor(duration) return ( chunk, t_start, t_end, seconds_start, seconds_total ) USE_DUMMY_AUDIO = False #当测试代码时，可以将其置为True，这样就不会读取实际数据，而是用生成的静默音频代替 if USE_DUMMY_AUDIO: logger.warning("USE_DUMMY_AUDIO flag is True, don't use it when train or test!") class SafeAudioReader: """ This class is an adaptor to Read_and_PadCrop_Normalized_T, make it safe to read audio data. """ def __init__(self, duration: float, # 返回音频长度 sample_rate: int, # 返回音频的采样率，如与实际音频采样率不同，会作resample randomize: bool = True, use_avoid_watermark_policy = False, ): self.n_samples = int(sample_rate * duration) self.reader = ( Read_and_PadCrop_Normalized_T_Avoid_Watermark if use_avoid_watermark_policy \ else Read_and_PadCrop_Normalized_T )(n_samples=self.n_samples, sample_rate=sample_rate, randomize=randomize) #NOTE:这个是核心的函数，所有数据集读取音频都是调用的这个函数！ def __call__(self, filepath: os.PathLike, # 音频路径 origin_sample_rate: Optional[int] = None, # 从json文件中读取的实际采样率，如果不给定，则会从文件头中读取 origin_duration: float = None, # 从json文件中读取的实际时长，如果不给定，则会从文件头中读取 ) -> torch.Tensor: if USE_DUMMY_AUDIO: wav = torch.zeros(self.n_samples, dtype=torch.float32) return wav try: if origin_sample_rate is None or origin_duration is None: # audio_info = torchaudio.info(filepath) # origin_sample_rate = audio_info.sample_rate # origin_duration = audio_info.num_frames / origin_sample_rate info = ffmpeg.probe(filepath) origin_duration = float(info['format']['duration']) origin_sample_rate = int(info['streams'][0]['sample_rate']) wav, *ignored = self.reader(filepath, origin_duration, origin_sample_rate) wav = wav.squeeze_(0) except Exception as e: logger.error(f"Error reading {filepath}: {e}") wav = torch.zeros(self.n_samples, dtype=torch.float32) return wav class PromptTemplate: def __init__(self, template_text: str, tag_map: Dict[str, str], lang:str ='en'): self.template_text = template_text self.tag_map = tag_map self.lang = lang @property def tags(self): return tuple(self.tag_map.keys()) def apply(self, **kwargs): for tag in list(kwargs.keys()): if kwargs[tag] == '': kwargs.pop(tag) for tag in self.tags: if tag in kwargs: kwargs[tag] = self.tag_map[tag].format(**{tag: kwargs[tag]}).strip('[]') else: kwargs[tag] = '' prompt = self.template_text.format(**kwargs) return self.beautify(prompt) def beautify(self, text): if self.lang == 'en': return self._beautify_en(text) elif self.lang == 'zh': return self._beautify_zh(text) else: raise ValueError(f'Unknown language {self.lang}') @staticmethod def _beautify_en(text): # no continuous commas without content between them text = re.sub(r'[,\s]*,[,\s]*', r', ', text) # no continuous whitespace text = re.sub(r'\s+', ' ', text) # the comma is NOT followed by whitespace, and should be followed by ONE whitespace text = re.sub(r'\s+,', r',', text) text = re.sub(r',\s+', r', ', text) # no whitespace before the full stop text = re.sub(r'\s+\.', r'.', text) # strip whitespace, comma, and replace ',.' text = text.strip(' ,') text = text.replace(',.', '.') return text @staticmethod def _beautify_zh(text): # no continuous commas without content between them text = re.sub(r'[，、\s]*，[，、\s]*', r'，', text) text = re.sub(r'[，、\s]*、[，、\s]*', r'、', text) # assume there should be NO whitespace in Chinese text = re.sub(r'\s+', r'', text) # strip whitespace, comma, and replace '，。' text = text.strip('， 、') text = text.replace('，。', '。') return text def __repr__(self): return f'PromptTemplate({self.template_text!r}, {self.tag_map!r})' __str__ = __repr__ def parse_prompt_template(prompt_template_text, lang='en'): span_pattern = re.compile(r'\[.*?{.+?}.*?\]', re.DOTALL) tag_pattern = re.compile(r'{.+?}', re.DOTALL) template_text = prompt_template_text.strip() span_texts = span_pattern.findall(prompt_template_text) tag_map = {} for span_text in span_texts: tag = tag_pattern.findall(span_text)[0].strip('{}') tag_map[tag] = span_text template_text = template_text.replace(span_text, '{'+tag+'}') return PromptTemplate(template_text=template_text, tag_map=tag_map, lang=lang) def load_prompt_templates(path, num = 5, lang='en') -> List[PromptTemplate]: with open(path, 'r') as f: lines = f.readlines() cnt = 0 pts = [] for line in lines: pt = parse_prompt_template(line, lang=lang) cnt += 1 if len(pt.tags) < num: logger.error(f'Not enough tags on {path} in line {cnt}: {pt.tags}') pts.append(pt) return pts def get_base_dir_file(key: os.PathLike): base = os.path.basename(key) dirname = os.path.basename(os.path.dirname(key)) return os.path.join(dirname, base) def read_jsonlike(path: os.PathLike): #json or jsonl if str(path).endswith(".json"): with open(path, 'r', encoding='utf8') as f: data = json.load(f) return data elif str(path).endswith(".jsonl"): with open(path, 'r', encoding='utf8') as f: data = [json.loads(line) for line in f.readlines()] return data else: raise ValueError("Unknown file format") dist_prob_map = { 1: (1.0,), 2: (0.5, 0.5), 3: (0.3, 0.4, 0.3), 4: (0.2, 0.3, 0.3, 0.2), 5: (0.2, 0.2, 0.3, 0.2, 0.1), 6: (0.1, 0.15, 0.2, 0.2, 0.2, 0.15), 7: (0.05, 0.1, 0.1, 0.2, 0.25, 0.2, 0.1), 8: (0.03, 0.05, 0.1, 0.15, 0.25, 0.2, 0.1, 0.12), 9: (0.02, 0.1, 0.1, 0.1, 0.15, 0.2, 0.15, 0.1, 0.08), 10: (0.01, 0.1, 0.1, 0.15, 0.2, 0.15, 0.1, 0.05, 0.05, 0.09) } ''' #更加偏向短文本的方案 dist_prob_map = { 1: (1.0,), 2: (0.7, 0.3), 3: (0.7, 0.2, 0.1), 4: (0.6, 0.2, 0.1, 0.1), 5: (0.6, 0.2, 0.1, 0.05, 0.05), 6: (0.6, 0.15, 0.1, 0.05, 0.05, 0.05), 7: (0.05, 0.1, 0.1, 0.2, 0.25, 0.2, 0.1), 8: (0.03, 0.05, 0.1, 0.15, 0.25, 0.2, 0.1, 0.12), 9: (0.02, 0.1, 0.1, 0.1, 0.15, 0.2, 0.15, 0.1, 0.08), 10: (0.01, 0.1, 0.1, 0.15, 0.2, 0.15, 0.1, 0.05, 0.05, 0.09) } ''' #全部都用的方案 # dist_prob_map = { # 1: (1.0,), # 2: (0, 1.0), # 3: (0, 0, 1.0), # 4: (0, 0, 0, 1.0), # 5: (0, 0, 0, 0, 1.0), # 6: (0, 0, 0, 0, 0, 1.0), # 7: (0, 0, 0, 0, 0, 0, 1.0), # 8: (0, 0, 0, 0, 0, 0, 0, 1.0), # 9: (0, 0, 0, 0, 0, 0, 0, 0, 1.0), # 10: (0, 0, 0, 0, 0, 0, 0, 0, 0, 1.0) # } dist_prob_map_low = { 1: (1.0,), 2: (0.8, 0.2), 3: (0.8, 0.1, 0.1), 4: (0.7, 0.1, 0.1, 0.1), 5: (0.7, 0.1, 0.1, 0.05, 0.05), 6: (0.7, 0.1, 0.05, 0.05, 0.05, 0.05), } _bpm_range_rights = ( (40, '20-40'), (60, '40-60'), (66, '60-66'), (76, '66-76'), (108, '76-108'), (120, '108-120'), (168, '120-168'), (176, '168-176'), (200, '176-200') ) _bpm_desc_map = { '20-40': ("glacial pace", "extremely slow tempo", "crawl-like speed", "snail's pace", "almost motionless rhythm", "Larghissimo"), '40-60': ("broad and slow", "spacious tempo", "unhurried pace", "calm rhythm", "relaxed speed", "Largo"), '60-66': ("gentle tempo", "leisurely pace", "easy-going rhythm", "unrushed speed", "smooth and slow", 'Larghetto'), '66-76': ("slow and steady", "deliberate tempo", "unhurried pace", "relaxed rhythm", "easy speed", 'Adagio'), '76-108': ("walking pace", "moderate tempo", "steady rhythm", "balanced speed", "easy-flowing tempo", "Andante"), '108-120': ("medium pace", "comfortable tempo", "even rhythm", "measured speed", "controlled tempo", 'Moderato'), '120-168': ("quick and lively", "brisk pace", "energetic tempo", "upbeat rhythm", "spirited speed", 'Allegro'), '168-176': ("lively and fast", "bright tempo", "sprightly pace", "vibrant rhythm", "animated speed", 'Vivace'), '176-200': ("very fast tempo", "rapid pace", "high-speed rhythm", "hurried speed", "accelerated tempo", 'Presto'), '>200': ("extremely fast", "breakneck speed", "blazing tempo", "lightning-fast rhythm", "supercharged pace", 'Prestissimo') } _bpm_desc_map_zh = { '20-40': ("极度缓慢", "极慢的节奏", "悠长的旋律", "迟缓的节奏", "几乎静止的节奏", "甚缓"), '40-60': ("宽广而缓慢", "宽敞的节奏", "从容不迫的速度", "平静的节奏", "轻松的速度", "广板"), '60-66': ("柔和的节奏", "悠闲的速度", "轻松的节奏", "不慌不忙的速度", "平滑而缓慢", '小广板'), '66-76': ("缓慢而稳定", "沉稳的旋律", "从容不迫的速度", "轻松的节奏", "轻松的速度", '慢板'), '76-108': ("步行速度", "适中的节奏", "稳定的节奏", "平衡的速度", "流畅的节奏", "行板"), '108-120': ("中等速度", "舒适的节奏", "均匀的节奏", "有节制的速度", "稳定的氛围", '中板'), '120-168': ("快速而生动", "轻快的速度", "充满活力的节奏", "欢快的节奏", "富有精神的速度", '快板'), '168-176': ("生动而快速", "明快的节奏", "活泼的速度", "充满活力的节奏", "生气勃勃的速度", '活泼的'), '176-200': ("非常快的节奏", "快速的速度", "高速的节奏", "匆忙的速度", "加速的节奏", '急板'), '>200': ("极快的速度", "极速旋律", "炽热的节奏", "闪电般的节奏", "疾驰的速度", '最急板') } def get_bpm_range(bpm): bpm = int(bpm) for right, tag in _bpm_range_rights: if bpm <= right: return tag return '>200' def gen_bpm_descript(bpm, lang='en'): bpm_range = get_bpm_range(bpm) if lang == 'en': return random.choice(_bpm_desc_map[bpm_range]) elif lang == 'zh': return random.choice(_bpm_desc_map_zh[bpm_range]) else: raise ValueError(f"Unknown language {lang}") def read_translate(translate: Union[Dict[str, os.PathLike], os.PathLike, None]): if translate is None: return None if isinstance(translate, str): return read_jsonlike(translate) return {k: read_jsonlike(path) for k, path in translate.items()} def gen_plain_prompt(key_list, sep=', '): if len(key_list) == 0: return 'none' key_list = [k.strip() for k in key_list] if len(key_list) > 10: random.shuffle(key_list) key_list = key_list[:10] probs = dist_prob_map[len(key_list)] num_tags = random.choices(range(1, len(key_list)+1), probs, k=1)[0] random.shuffle(key_list) tags = key_list[:num_tags] tags_str = sep.join(tags) return tags_str class MagnaTagATuneDataset(Dataset): def __init__(self): pass def tags_to_desc(tag_list, sep=',') -> str: if not isinstance(tag_list, Sequence): return str(tag_list) if isinstance(tag_list, str): return tag_list if len(tag_list) <= 0: return '' elif len(tag_list) <= 5: probs = dist_prob_map[len(tag_list)] tags_num = random.choices(range(1, len(tag_list)+1), probs)[0] random.shuffle(tag_list) tag_list = tag_list[:tags_num] return sep.join(tag_list) else: probs = dist_prob_map[5] tags_num = random.choices(range(1, 6), probs)[0] random.shuffle(tag_list) tag_list = tag_list[:tags_num] return sep.join(tag_list) def get_sr_and_duration_info(item): return item.get('sample_rate', None), item.get('duration', None) class MtgJamendoDatasetFromJson(Dataset): def __init__(self, data_dir:str, json_path:str, duration:float=10, sr:int = 0, lang = 'en', plain_rate = 0, return_audio = True, return_path = False, prompt_template_path: os.PathLike = None, tag_types = [], translate:Optional[Dict[str, os.PathLike]] = None, use_literal_none = True, ): self.audio_reader = SafeAudioReader(duration, sr) self.data_dir = data_dir self._load_metadata_json(json_path) self.sr = sr self.duration = duration self.plain_rate = plain_rate self.return_audio = return_audio self.return_path = return_path self.use_literal_none = use_literal_none self.lang = lang self.use_dynamic_prompt = prompt_template_path is not None and plain_rate < 1.0 if self.use_dynamic_prompt: self.prompt_templates = load_prompt_templates(prompt_template_path, num = len(tag_types)) self.tag_types = tag_types self.translate = read_translate(translate) #这些tag被认为是弱语义的，会避免产生仅包含这些tag的文本提示 WEAK_TAG_LIST = ["title", "artist"] def _load_metadata_json(self, json_path): with open(json_path) as fp: self.data = json.load(fp) def convert_key_to_path(self, key): return os.path.join(self.data_dir, get_base_dir_file(key)) def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] path = self.convert_key_to_path(item['key']) description = self.generate_description(item) if self.return_audio: sr, duration = get_sr_and_duration_info(item) audio = self.audio_reader(path, sr, duration) else: audio = None if self.return_path: return audio, description, path return audio, description def tags_to_desc(self, tag_list, tag_type) -> str: if self.lang == 'en': return tags_to_desc(tag_list) elif self.lang == 'zh': translator = self.translate[tag_type] translated_tag_list = [translator[tag] for tag in tag_list if tag in translator ] return tags_to_desc(translated_tag_list, sep='、') def generate_description(self, item): if random.random() > self.plain_rate: # dynamically generate prompt from given prompt template prompt_template = random.choice(self.prompt_templates) description = self.generate_description_dynamic(item, prompt_template) else: # use plain prompt, i.e. tags sequence separated by comma description = self.generate_description_plain(item) return description def generate_description_dynamic(self, data, prompt_template: PromptTemplate): exists_tag = [key for key in data if (key in self.tag_types) and (data[key] is not None) and (len(data[key]) > 0)] exists_weak_tag = list(filter(lambda t: t in self.WEAK_TAG_LIST, exists_tag)) exists_strong_tag = list(filter(lambda t: t not in self.WEAK_TAG_LIST, exists_tag)) if len(exists_strong_tag) > 0: probs = dist_prob_map[len(exists_strong_tag)] tags_num = random.choices(range(1, len(exists_strong_tag)+1), probs)[0] random.shuffle(exists_strong_tag) tags = exists_strong_tag[:tags_num] weak_probs = dist_prob_map_low[len(exists_weak_tag) + 1] weak_tags_num = random.choices(range(0, len(exists_weak_tag) + 1), weak_probs)[0] random.shuffle(exists_weak_tag) weak_tags = exists_weak_tag[:weak_tags_num] tags += weak_tags tags_args = {tag: self.tags_to_desc(data[tag], tag) for tag in tags} prompt = prompt_template.apply(**tags_args) else: # no strong tags, use all weak tags instead tags_args = {tag: self.tags_to_desc(data[tag], tag) for tag in exists_weak_tag} prompt = prompt_template.apply(**tags_args) if self.use_literal_none and len(tags_args) == 0: return 'none' return prompt def generate_description_plain(self, item): keywords = [] for tag_t in self.tag_types: this_key = item[tag_t] if this_key is None: continue if isinstance(this_key, str): this_key = [this_key] if self.lang != 'en': this_key = [self.get_translation(tag_t, k) for k in this_key] keywords += this_key return gen_plain_prompt(keywords, sep=self.keysep) def get_translation(self, tag_t, k): k = k.strip() if k in self.translate[tag_t]: return self.translate[tag_t][k] else: return k @property def keysep(self): if self.lang == 'zh': return '，' if random.random() > 0.5 else '、' elif self.lang == 'en': return ', ' class AudioStockDataset(Dataset): def __init__(self, metadata_path:str, duration:float=10, sr:int = 0, plain_rate = 0, return_path = False, return_audio = True, prompt_template_path: os.PathLike = None, tag_types = [], lang = 'en', translate:Optional[Dict[str, os.PathLike]] = None, use_literal_none = True, ): self.audio_reader = SafeAudioReader(duration, sr) self._load_metadata(metadata_path) self.sr = sr self.duration = duration self.plain_rate = plain_rate self.return_path = return_path self.return_audio = return_audio self.use_literal_none = use_literal_none self.use_dynamic_prompt = prompt_template_path is not None and plain_rate < 1.0 if self.use_dynamic_prompt: self.prompt_templates = load_prompt_templates(prompt_template_path, num = len(tag_types), lang = lang) self.tag_types = tag_types self.lang = lang self.translate = read_translate(translate) def _load_metadata(self, metadata_path): with open(metadata_path) as fp: lines = fp.readlines() self.data = [] for line in lines: item = json.loads(line) self.data.append(item) self.is_info_recorded = bool('Tags' in self.data[0]) def __len__(self): return len(self.data) def __getitem__(self, idx): path:str = self.data[idx]["path"] json_path = path[:path.rfind('.')] + ".json" if self.is_info_recorded: item = self.data[idx] else: try: with open(json_path) as fp: item:dict = json.load(fp) except Exception as e: print(f"Error loading json file {json_path} :\n{e}") item = {} description = self.generate_description(item) if self.return_audio: sr, duration = get_sr_and_duration_info(item) audio = self.audio_reader(path, sr, duration) else: audio = None if self.return_path: return audio, description, path return audio, description def generate_description(self, item): if random.random() > self.plain_rate: # dynamically generate prompt from given prompt template prompt_template = random.choice(self.prompt_templates) description = self.generate_description_dynamic(item, prompt_template) else: # use plain prompt, i.e. tags sequence separated by comma description = self.generate_description_plain(item) return description def generate_description_dynamic(self, data, prompt_template: PromptTemplate): exists_tag = [key for key in data if (key in self.tag_types) and (data[key] is not None) and (len(data[key]) > 0)] if len(exists_tag) > 0: probs = dist_prob_map[len(exists_tag)] tags_num = random.choices(range(1, len(exists_tag)+1), probs)[0] random.shuffle(exists_tag) tags = exists_tag[:tags_num] tags_args = {tag: self.tags_to_desc(data[tag], tag) for tag in tags} tags_args = self.handle_BPM_tag(tags_args) prompt = prompt_template.apply(**tags_args) else: return 'none' if self.use_literal_none and len(tags_args) == 0: return 'none' return prompt def get_translation(self, tag_t, k): k = k.strip() if k in self.translate[tag_t]: return self.translate[tag_t][k] else: return k def generate_description_plain(self, item): keywords = [] for tag_t in self.tag_types: if tag_t == 'BPMDescript': bpm = item['BPM'] if bpm is None or bpm.strip() == '' or bpm.strip() == '0': continue this_key = gen_bpm_descript(bpm.strip(), lang=self.lang) elif tag_t == 'BPM': bpm = item['BPM'] if bpm is None or bpm.strip() == '' or bpm.strip() == '0': continue this_key = f"{bpm.strip()} bpm" else: this_key = item[tag_t] if this_key is None: continue if isinstance(this_key, str): this_key = [this_key] if self.lang != 'en': this_key = [self.get_translation(tag_t, k) for k in this_key] if this_key is None: continue if isinstance(this_key, str): this_key = [this_key] keywords += this_key return gen_plain_prompt(keywords, sep=self.keysep) @property def keysep(self): if self.lang == 'zh': return '，' if random.random() > 0.5 else '、' elif self.lang == 'en': return ', ' def tags_to_desc(self, tag_list, tag_type) -> str: if self.lang == 'en': return tags_to_desc(tag_list) elif self.lang == 'zh': if tag_type == 'BPM': return tags_to_desc(tag_list, sep='、') translator = self.translate[tag_type] translated_tag_list = [translator[tag] for tag in tag_list if tag in translator ] return tags_to_desc(translated_tag_list, sep='、') def handle_BPM_tag(self, tags_args): if "BPM" in tags_args and 'BPMDescript' in self.tag_types: bpm = tags_args["BPM"] del tags_args["BPM"] tag_types_used = random.choice((('BPM',), ('BPMDescript',), ('BPM', 'BPMDescript'))) for tag_type in tag_types_used: tags_args[tag_type] = bpm if tag_type == 'BPM' else gen_bpm_descript(bpm, lang=self.lang) return tags_args def mp3_path_to_id(mp3_path): return int( mp3_path[mp3_path.rindex('/') + 1 : mp3_path.rindex('.')] ) class TmeDataset(Dataset): def __init__(self, data_index:str, music_info:str = None, duration:float = 10, sr:int = 0, plain_rate = 0, return_path = False, return_audio = True, return_ID = False, prompt_format_path: os.PathLike = None, tag_types = ['*'], lang = 'zh', translate: Optional[os.PathLike] = None, prompt_dir: os.PathLike = None, #使用GPT生成的预有的prompt ): if plain_rate > 0: print("Tme Dataset do not support plain rate > 0, use plain_rate = 0 instead.") plain_rate = 0 self.audio_reader = SafeAudioReader(duration, sr) self.sr = sr self.duration = duration self.plain_rate = plain_rate self.return_path = return_path self.return_audio = return_audio self.return_ID = return_ID self.lang = lang self.use_ready_prompt = prompt_dir is not None data_index = read_jsonlike(data_index) self.data_index_dict = {mp3_path_to_id(d['path']) : d for d in data_index} self.data_ids = list(self.data_index_dict.keys()) if not self.use_ready_prompt: #读取音乐的信息文件 music_info = read_jsonlike(music_info) if 'music' in music_info: music_info = music_info['music'] self.music_info_dict = {d["歌曲ID"]:d for d in music_info} self.data_index_dict = {k:v for k,v in self.data_index_dict.items() if k in self.music_info_dict} self.data_ids = list(self.data_index_dict.keys()) with open(prompt_format_path) as fp: self.prompt_formats = yaml.load(fp, Loader=yaml.FullLoader) #加载tag types，并分成一般的tag_types和关键的key_tag_types if '*' in tag_types: self.tag_types = ['歌曲名', 'bpm', '专辑名', '歌手名', '作曲', 'tag'] else: self.tag_types = tag_types self.key_tag_types = [] if 'tag' in self.tag_types: self.tag_types.remove('tag') self.key_tag_types = list(self.prompt_formats['tag'].keys()) #加载translate翻译 if translate is not None: self.translator = read_jsonlike(translate) else: data_ids_set = set(self.data_ids) self.prompts_dict = {} for fname in os.listdir(prompt_dir): items = read_jsonlike(os.path.join(prompt_dir, fname)) for item in items: if item['ID'] not in data_ids_set or not self.is_valid_prompt_text(item['Text']): continue if item['ID'] not in self.prompts_dict: self.prompts_dict[item['ID']] = [] self.prompts_dict[item['ID']].append(item['Text']) self.data_index_dict = {k:v for k,v in self.data_index_dict.items() if k in self.prompts_dict} self.data_ids = list(self.data_index_dict.keys()) def tags_to_desc(self, tag_list) -> str: if is_bearable(tag_list, int): return str(tag_list) if self.lang == 'zh': return tags_to_desc(tag_list, sep=self.sep) else: translated_tag_list = [self.translator[tag] for tag in tag_list if tag in self.translator ] return tags_to_desc(translated_tag_list, sep=self.sep) def gen_desc_of_tag(self, formats, tags): fmt = random.choice(formats) return fmt.format(self.tags_to_desc(tags)) @staticmethod def check_valid(value): if isinstance(value, int) or isinstance(value, float): return value > 0 if (value is not None) and (not isinstance(value, Sequence) or len(value) > 0): return True return False @staticmethod def remove_repeat(data): #若专辑名和歌曲名相同，则只使用后者 album_name = data.get('专辑名', None) if album_name is not None and album_name == data.get('歌曲名', None): del data['专辑名'] return data @property def comma(self): if self.lang == 'zh': return '，' elif self.lang == 'en': return ', ' @property def sep(self): if self.lang == 'zh': return '、' elif self.lang == 'en': return ', ' def generate_description(self, item): if random.random() > self.plain_rate: # dynamically generate prompt from given prompt template description = self.generate_description_dynamic(item) else: # use plain prompt, i.e. tags sequence separated by comma description = self.generate_description_plain(item) return description def generate_description_dynamic(self, data): data = self.remove_repeat(data) weak_tags = [key for key in data if (key in self.tag_types and self.check_valid(data[key]))] #弱语义的tag，这些tag的出现比例会放低 key_tags = [key for key in data['tag'] if (key in self.key_tag_types and self.check_valid(data['tag'][key]))] #关键的tag，这些tag必须出现至少一个 prompts = [] if len(weak_tags) > 0: probs = dist_prob_map_low[len(weak_tags)] if len(key_tags) > 0: tags_num = random.choices(range(0, len(weak_tags)), probs)[0] else: tags_num = random.choices(range(1, len(weak_tags) + 1), probs)[0] random.shuffle(weak_tags) tags = weak_tags[:tags_num] for tag_type in tags: tag_desc = self.gen_desc_of_tag(self.prompt_formats[tag_type], int(data[tag_type]) if tag_type == 'bpm' else data[tag_type]) prompts.append(tag_desc) if len(key_tags) > 0: probs = dist_prob_map[len(key_tags)] tags_num = random.choices(range(1, len(key_tags) + 1), probs)[0] random.shuffle(key_tags) tags = key_tags[:tags_num] for tag_type in tags: tag_desc = self.gen_desc_of_tag(self.prompt_formats['tag'][tag_type], data['tag'][tag_type]) prompts.append(tag_desc) random.shuffle(prompts) return self.comma.join(prompts) def generate_description_plain(self, item): keywords = item['tag'] if self.lang != 'en': keywords = [self.translator[k.strip()] for k in keywords] return gen_plain_prompt(keywords, sep=self.keysep) @property def keysep(self): if self.lang == 'zh': return '，' if random.random() > 0.5 else '、' elif self.lang == 'en': return ', ' def is_valid_prompt_text(self, text): for bad in ('抱歉','sorry', 'Sorry'): if bad in text: return False return True def get_ready_prompt(self, path): sid = mp3_path_to_id(path) return random.choice(self.prompts_dict[sid]) def __len__(self): return len(self.data_ids) def __getitem__(self, idx): data_id = self.data_ids[idx] item = self.data_index_dict[data_id] path = item['path'] if not self.use_ready_prompt: info = self.music_info_dict[data_id] description = self.generate_description(info) else: description = self.get_ready_prompt(path) if self.return_audio: sr, duration = get_sr_and_duration_info(item) audio = self.audio_reader(path, sr, duration) else: audio = None if self.return_path: if self.return_ID: return audio, description, path, info['歌曲ID'] return audio, description, path if self.return_ID: return audio, description, info['歌曲ID'] return audio, description class Pond5Dataset(Dataset): MAX_PROMPT_LEN = 200 def __init__(self, metadata_path:str, index_path:str, duration:float=10, sr:int = 0, plain_rate = 0, return_path = False, return_audio = True, lang = 'en', translate:Optional[Dict[str, os.PathLike]] = None, use_literal_none = True, use_avoid_watermark_policy = None, ): if use_avoid_watermark_policy is None: raise ValueError("`use_avoid_watermark_policy` is an important param, you need to explicitly specify it with bool type") self.use_avoid_watermark_policy = use_avoid_watermark_policy self.audio_reader = SafeAudioReader(duration, sr, use_avoid_watermark_policy=use_avoid_watermark_policy) self._load_metadata(metadata_path, index_path) self.sr = sr self.duration = duration self.plain_rate = plain_rate self.return_path = return_path self.return_audio = return_audio self.use_literal_none = use_literal_none self.lang = lang self.translate = read_translate(translate) def _load_metadata(self, metadata_path, index_path): data_index = read_jsonlike(index_path) data_ids = set([item['id'] for item in data_index]) with open(metadata_path) as fp: lines = fp.readlines() append_ids = set() self.data = [] for line in lines: item = json.loads(line) if item['id'] in data_ids and item['id'] not in append_ids: self.data.append(item) append_ids.add(item['id']) def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] path:str = item["path"] description = self.generate_description(item) if self.return_audio: sr, duration = get_sr_and_duration_info(item) audio = self.audio_reader(path, sr, duration) else: audio = None if self.return_path: return audio, description, path return audio, description @property def keysep(self): if self.lang == 'zh': return '，' if random.random() > 0.5 else '、' elif self.lang == 'en': return ', ' def generate_description(self, item): if random.random() > self.plain_rate: # dynamically generate prompt from given prompt template description = self.generate_description_dynamic(item) else: # use plain prompt, i.e. tags sequence separated by comma description = self.generate_description_plain(item) return description def get_translation(self, k): k = k.strip() if k in self.translate: return self.translate[k] else: return k def generate_description_plain(self, item): keywords = item['keywords'] if self.lang != 'en': keywords = [self.get_translation(k) for k in keywords] return gen_plain_prompt(keywords, sep=self.keysep) def generate_description_dynamic(self,item): desc = item.get('desc', 'none') if desc is None: desc = 'none' desc = desc.strip() if len(desc) > self.MAX_PROMPT_LEN: shorter_desc = desc[:self.MAX_PROMPT_LEN] # find last stop stop_idx = shorter_desc.rfind('.') if stop_idx == -1: stop_idx = shorter_desc.rfind('!') if stop_idx == -1: stop_idx = shorter_desc.rfind(',') if stop_idx == -1: stop_idx = self.MAX_PROMPT_LEN - 1 desc = desc[:stop_idx+1] return desc class SoundDataset(Dataset): def __init__(self, metadata_index: str, duration:float = 10, min_non_silent_duration:float = 3, sr:int = 0, return_path = False, return_audio = True, ): self.data = read_jsonlike(metadata_index) self.sr = sr self.reader = SafeAudioReader(duration, sr) self.duration = duration self.min_non_silent_duration = min_non_silent_duration self.return_audio = return_audio self.return_path = return_path def __getitem__(self, index): item = self.data[index] if self.return_audio: origin_duration = item['duration'] if origin_duration < self.min_non_silent_duration: audio = self.read_and_repeat_and_pad(item) else: audio = self.reader(item['path'], item['sample_rate'], origin_duration) else: audio = None desc = item['caption'] if self.return_path: return audio, desc, item['path'] else: return audio, desc def __len__(self): return len(self.data) def read_and_repeat_and_pad(self, item): path = item['path'] try: # read clip, sr = torchaudio.load(path) if len(clip.shape) > 1: clip = torch.mean(clip, dim=0, keepdim=True) clip = resample(clip, sr, self.sr) #repeat n_repeats = math.ceil(self.min_non_silent_duration/item['duration']) clip = torch.repeat_interleave(clip, n_repeats, dim=0).reshape(-1) #pad n_samples = int(self.duration * self.sr) if clip.shape[0] >= n_samples: audio = clip[:n_samples] else: audio = torch.zeros(int(self.duration * self.sr), dtype=clip.dtype) start_pos = np.random.randint(0, max(0,(n_samples - clip.shape[0]))) audio[start_pos:start_pos+clip.shape[0]] = clip return audio except Exception as e: logger.error(f"Error reading {path}: {e}") wav = torch.zeros(int(self.duration * self.sr), dtype=torch.float32) return wav class CombinedDataset(Dataset): @beartype def __init__(self, datasets: Sequence[Dataset], ratios: Sequence[int]): self.datasets = datasets self.datasets_index = [] for i,dataset in enumerate(datasets): if dataset is None: continue for dup in range(ratios[i]): for j in range(len(dataset)): self.datasets_index.append((i,j)) def __len__(self): return len(self.datasets_index) def __getitem__(self, idx): index = self.datasets_index[idx] i,j = index return self.datasets[i][j] class CombinedDataset_random(Dataset): @beartype def __init__(self, num_examples:int, datasets: Sequence[Dataset], ratios: Sequence[int]): self.datasets = datasets self.datasets_index = [] for i,dataset in enumerate(datasets): if dataset is None: continue for dup in range(ratios[i]): for j in range(len(dataset)): self.datasets_index.append((i,j)) if num_examples > 0: self.random_choose = True self.dataset_len = num_examples else: self.random_choose = False self.dataset_len = len(self.datasets_index) def __len__(self): return self.dataset_len def __getitem__(self, idx): first_try = True try_cnt = 0 while True: try: if(self.random_choose or not first_try): index2 = [] index2.append(np.random.randint(0,len(self.datasets))) index2.append(np.random.randint(0,len(self.datasets[index2[-1]]))) else: index2 = self.datasets_index[idx] first_try = False out = list(self.datasets[index2[0]][index2[1]]) return out except: print("Error loadding ", index2) try_cnt += 1 if(try_cnt>10): raise ValueError() class SoundMixedDataset(Dataset): @staticmethod def music_desc(desc): return f'Music:<{desc}>' @staticmethod def sound_desc(desc): return f'Effect:<{desc}>' def __init__(self, music_dataset: Dataset, sound_dataset: Dataset, mixed_ratios: Tuple[float, float, float] = (0.3, 0.3, 0.4) # 只有音乐：只有音效：音乐音效混合 的比例 ) -> None: self.music_dataset = music_dataset self.sound_dataset = sound_dataset music_r, sound_r, mix_r = [r/sum(mixed_ratios) for r in mixed_ratios] #化为0-1间的比例 #三个概率区间的左端点 self.music_anchor = 0 self.sound_anchor = music_r self.mix_anchor = music_r + sound_r def __len__(self): return len(self.music_dataset) def get_random_sound_data(self): idx = random.randint(0, len(self.sound_dataset)-1) return self.sound_dataset[idx] def __getitem__(self, idx): p = random.random() if p >= self.mix_anchor: music, m_desc = self.music_dataset[idx] sound, s_desc = self.get_random_sound_data() audio = music + sound if(audio.abs().max()>1.0): music = music / audio.abs().max() * 0.95 audio = audio / audio.abs().max() * 0.95 desc = self.music_desc(m_desc) + self.sound_desc(s_desc) return audio[None,:], music[None,:], desc elif p >= self.sound_anchor: audio, desc = self.get_random_sound_data() return audio[None,:], torch.zeros_like(audio[None,:]), self.sound_desc(desc) else: audio, desc = self.music_dataset[idx] return audio[None,:], audio[None,:], self.music_desc(desc) class DecoTagDataset(Dataset): '''这个类把普通的datatset包装成适用于标签解耦学习的dataset''' TAG_TYPES = ('genre', 'mood', 'insrument') def __init__(self, dataset_class: type, tag_map: Dict[str, str], *args, **kwargs): self.datasets = [] for i, tag_t in enumerate(self.TAG_TYPES): kwargs['tag_types'] = [tag_map[tag_t]] kwargs['return_audio'] = (i == 0) #只有第0个需要返回音频和文本，其余只需要返回文本 self.datasets.append(dataset_class(*args, **kwargs)) def __len__(self): return len(self.datasets[0]) def __getitem__(self, idx): audio, text = self.datasets[0][idx] texts = (text, self.datasets[1][idx][1], self.datasets[2][idx][1]) return audio, texts class DecoTagWrapper: '''这是一个包装器，便于选择是否使用标签解耦学习''' def __init__(self, dataset_class: Dataset, deco_tag_types: List[str] = list(), switch_on: bool = False): self.dataset_class = dataset_class self.tag_map = dict(zip(DecoTagDataset.TAG_TYPES, deco_tag_types)) self.switch_on = switch_on def __call__(self, *args, **kwargs): if self.switch_on: return DecoTagDataset(self.dataset_class, self.tag_map, *args, **kwargs) else: return self.dataset_class(*args, **kwargs)