462 lines
18 KiB
Python
462 lines
18 KiB
Python
from torch.utils.data import Dataset
|
||
from beartype.typing import Sequence, Callable, Optional, Dict, List
|
||
from beartype.door import is_bearable
|
||
import random
|
||
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 soundfile
|
||
import numpy as np
|
||
import json
|
||
import yaml
|
||
import random
|
||
import librosa
|
||
from loguru import logger
|
||
import re
|
||
|
||
|
||
def _av_read(filepath, seek_time=0, duration=None):
|
||
if duration is not None:
|
||
sr = librosa.get_samplerate(filepath)
|
||
offset = seek_time
|
||
num_samples = int(duration * sr)
|
||
wav, _ = librosa.load(filepath, sr=sr, offset=offset, duration=duration)
|
||
else:
|
||
wav, sr = librosa.load(filepath, sr=None, offset=seek_time)
|
||
|
||
return wav, sr
|
||
|
||
def audio_read(filepath: tp.Union[str, Path], seek_time: float = 0.,
|
||
duration: float = -1., pad: bool = True) -> tp.Tuple[torch.Tensor, int]:
|
||
"""Read audio by picking the most appropriate backend tool based on the audio format.
|
||
|
||
Args:
|
||
filepath (str or Path): Path to audio file to read.
|
||
seek_time (float): Time at which to start reading in the file.
|
||
duration (float): Duration to read from the file. If set to -1, the whole file is read.
|
||
pad (bool): Pad output audio if not reaching expected duration.
|
||
Returns:
|
||
tuple of torch.Tensor, int: Tuple containing audio data and sample rate.
|
||
"""
|
||
fp = Path(filepath)
|
||
if fp.suffix in ['.flac', '.ogg']: # TODO: check if we can safely use av_read for .ogg
|
||
# There is some bug with ffmpeg and reading flac
|
||
info = soundfile.info(filepath)
|
||
frames = -1 if duration <= 0 else int(duration * info.samplerate)
|
||
frame_offset = int(seek_time * info.samplerate)
|
||
wav, sr = soundfile.read(filepath, start=frame_offset, frames=frames, dtype=np.float32)
|
||
assert info.samplerate == sr, f"Mismatch of sample rates {info.samplerate} {sr}"
|
||
wav = torch.from_numpy(wav).t().contiguous()
|
||
if len(wav.shape) == 1:
|
||
wav = torch.unsqueeze(wav, 0)
|
||
elif (
|
||
fp.suffix in ['.wav', '.mp3'] and fp.suffix[1:] in ta.utils.sox_utils.list_read_formats()
|
||
and duration <= 0 and seek_time == 0
|
||
):
|
||
# Torchaudio is faster if we load an entire file at once.
|
||
wav, sr = librosa.load(fp, sr=None, mono=True)
|
||
else:
|
||
wav, sr = _av_read(filepath, seek_time, duration)
|
||
if pad and duration > 0:
|
||
expected_frames = int(duration * sr)
|
||
wav = F.pad(torch.tensor(wav), (0, expected_frames - wav.shape[-1]))
|
||
if not isinstance(wav, torch.Tensor):
|
||
wav = torch.tensor(wav)
|
||
return wav, sr
|
||
|
||
def random_seek_read(filepath, duration):
|
||
if duration > 0:
|
||
total_duration = librosa.get_duration(path=filepath)
|
||
acceptable_start = max(0, total_duration - duration)
|
||
wav, sr = audio_read(filepath, random.uniform(0, acceptable_start), duration, pad=True)
|
||
else:
|
||
wav, sr = audio_read(filepath, 0, -1, pad=False)
|
||
return wav, sr
|
||
|
||
def safe_random_seek_read(filepath, duration, sample_rate):
|
||
try:
|
||
wav, sr = random_seek_read(filepath, duration)
|
||
if sr != sample_rate:
|
||
wav = resample(wav, sr, sample_rate)
|
||
sr = sample_rate
|
||
except Exception as e:
|
||
logger.error(f"Error reading {filepath}: {e}")
|
||
sr = sample_rate
|
||
wav = torch.zeros(sr * max(duration, 0), dtype=torch.float32)
|
||
return wav, sr
|
||
|
||
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_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: Optional[Dict[str, os.PathLike]]):
|
||
if translate is None:
|
||
return None
|
||
return {k: read_jsonlike(path) for k, path in translate.items()}
|
||
|
||
|
||
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)
|
||
|
||
|
||
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
|
||
|
||
|
||
class AudioStockDataset(Dataset):
|
||
def __init__(self,
|
||
num_examples:int,
|
||
metadata_path:str,
|
||
duration:float=60,
|
||
sr:int = 0,
|
||
return_path = False,
|
||
return_audio = True,
|
||
prompt_template_path: os.PathLike = None,
|
||
tag_types = [],
|
||
lang = 'en',
|
||
translate:Optional[Dict[str, os.PathLike]] = None
|
||
):
|
||
self.duration = duration
|
||
self.MAX_DURATION = 360
|
||
self._load_metadata(metadata_path)
|
||
if num_examples > 0:
|
||
self.random_choose = True
|
||
self.dataset_len = num_examples
|
||
else:
|
||
self.random_choose = False
|
||
self.dataset_len = len(self.data)
|
||
self.sr = sr
|
||
self.return_path = return_path
|
||
self.return_audio = return_audio
|
||
|
||
self.use_dynamic_prompt = prompt_template_path is not None
|
||
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):
|
||
total_len = 0; valid_len = 0
|
||
with open(metadata_path) as fp:
|
||
lines = fp.readlines()
|
||
self.data = []
|
||
for line in lines:
|
||
item = json.loads(line)
|
||
total_len += 1
|
||
if(item['duration']>self.duration and item['duration']<self.MAX_DURATION):
|
||
valid_len += 1
|
||
self.data.append(item)
|
||
print("Filter data from {} to {}".format(total_len, valid_len))
|
||
self.is_info_recorded = bool('Tags' in self.data[0])
|
||
|
||
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 = np.random.randint(0,len(self.data))
|
||
else:
|
||
index2 = idx
|
||
first_try = False
|
||
return self.getitem_main(index2)
|
||
except:
|
||
print("Error loadding ", self.data[idx]["path"])
|
||
try_cnt += 1
|
||
if(try_cnt>10):
|
||
raise ValueError()
|
||
|
||
def getitem_main(self, idx):
|
||
path:str = self.data[idx]["path"]
|
||
json_path = path[:path.rfind('.')] + ".json"
|
||
if self.is_info_recorded:
|
||
item = self.data[idx]
|
||
else:
|
||
with open(json_path) as fp:
|
||
item:dict = json.load(fp)
|
||
description = self.generate_description(item)
|
||
if self.return_audio:
|
||
audio, sr = safe_random_seek_read(path, duration=self.duration, sample_rate=self.sr)
|
||
else:
|
||
audio = None
|
||
if self.return_path:
|
||
return audio, description, path
|
||
return audio, description
|
||
|
||
|
||
|
||
def generate_description(self, item):
|
||
if self.use_dynamic_prompt:
|
||
# dynamically generate prompt from given prompt template
|
||
prompt_template = random.choice(self.prompt_templates)
|
||
description = self.generate_description_dynamic(item, prompt_template)
|
||
else:
|
||
# use ordinary static prompt instead
|
||
description = self.generate_description_ordinary(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:
|
||
# no strong tags, use all weak tags instead
|
||
prompt = prompt_template.apply()
|
||
|
||
return prompt
|
||
|
||
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 generate_description_ordinary(self, data, thresh = 0.3):
|
||
if self.lang != 'en':
|
||
raise ValueError(f'Language {self.lang} is not supported for ordinary description generation')
|
||
description = f'a piece of music by {data["Artist"]}'
|
||
|
||
# Add genre if available
|
||
if data["Genre"] and random.random() > thresh:
|
||
genres = ', '.join(data["Genre"])
|
||
description += f', belonging to the {genres} genres'
|
||
|
||
# Add moods if available
|
||
if data["Tags"] and random.random() > thresh:
|
||
tags = ', '.join(data["Tags"])
|
||
description += f'. This track contains the tags:{tags}'
|
||
|
||
# Add moods if available
|
||
if data["Mood"] and random.random() > thresh:
|
||
moods = ', '.join(data["Mood"])
|
||
description += f'. This track conveys a {moods} mood.'
|
||
|
||
# Add instruments if available
|
||
if data["Instrument"] and random.random() > thresh:
|
||
instruments = ', '.join(data["Instrument"])
|
||
description += f'. and primarily features the following instruments: {instruments}'
|
||
|
||
# Add a period to end the description
|
||
description += '.'
|
||
|
||
return description
|
||
|