- 共享引擎(app/src/main):GestureEngine 6 手势识别、AudioEngine 三层音频 (ExoPlayer BGM 交叉淡入 / 对白 ducking / SoundPool 低延迟)、Game/GameEngine 状态机接口、MainActivity 全屏封面 + TalkBack - 《声命线》storyline:末世互动广播剧,3 幕 8 抉择点 4 结局 - 《声东击西》reflex:听声辨位反应,12 波战斗 + BGM 三档 BPM 变速 - 《暗室》escape:3x3 网格密室逃脱,钥匙→铁柜→保险丝→电箱→大门谜题链 - 文档:setup(从零搭建)/ architecture(架构)/ engine-api(引擎 API)/ asset-pipeline(资产流水线)/ games/*(三款游戏设计文档) - 资产管线:Qwen3-TTS 对白 + 程序合成音效/BGM + Mureka BGM 的 gen 脚本与台账, 全部 ogg 资产入库可复现 - env.d 环境脚本:android/tts 入库;sfx/mureka 含密钥不入库,仅提交 .example.sh 模板
511 lines
19 KiB
Python
511 lines
19 KiB
Python
import numpy as np
|
|
import json
|
|
import os
|
|
import dill
|
|
import random
|
|
import time
|
|
import torch
|
|
import torchaudio
|
|
|
|
from os import path
|
|
from torchaudio import transforms as T
|
|
from typing import Optional, Callable, List
|
|
|
|
from .utils import Stereo, Mono, PhaseFlipper, PadCrop_Normalized_T, VolumeNorm, strip_trailing_silence
|
|
|
|
AUDIO_KEYS = ("flac", "wav", "mp3", "m4a", "ogg", "opus")
|
|
|
|
# fast_scandir implementation by Scott Hawley originally in https://github.com/zqevans/audio-diffusion/blob/main/dataset/dataset.py
|
|
|
|
def fast_scandir(
|
|
dir:str, # top-level directory at which to begin scanning
|
|
ext:list, # list of allowed file extensions,
|
|
#max_size = 1 * 1000 * 1000 * 1000 # Only files < 1 GB
|
|
):
|
|
"very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243"
|
|
subfolders, files = [], []
|
|
ext = ['.'+x if x[0]!='.' else x for x in ext] # add starting period to extensions if needed
|
|
try: # hope to avoid 'permission denied' by this try
|
|
for f in os.scandir(dir):
|
|
try: # 'hope to avoid too many levels of symbolic links' error
|
|
if f.is_dir():
|
|
subfolders.append(f.path)
|
|
elif f.is_file():
|
|
file_ext = os.path.splitext(f.name)[1].lower()
|
|
is_hidden = os.path.basename(f.path).startswith(".")
|
|
|
|
if file_ext in ext and not is_hidden:
|
|
files.append(f.path)
|
|
except:
|
|
pass
|
|
except:
|
|
pass
|
|
|
|
for dir in list(subfolders):
|
|
sf, f = fast_scandir(dir, ext)
|
|
subfolders.extend(sf)
|
|
files.extend(f)
|
|
return subfolders, files
|
|
|
|
def keyword_scandir(
|
|
dir: str, # top-level directory at which to begin scanning
|
|
ext: list, # list of allowed file extensions
|
|
keywords: list, # list of keywords to search for in the file name
|
|
):
|
|
"very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243"
|
|
subfolders, files = [], []
|
|
# make keywords case insensitive
|
|
keywords = [keyword.lower() for keyword in keywords]
|
|
# add starting period to extensions if needed
|
|
ext = ['.'+x if x[0] != '.' else x for x in ext]
|
|
banned_words = ["paxheader", "__macosx"]
|
|
try: # hope to avoid 'permission denied' by this try
|
|
for f in os.scandir(dir):
|
|
try: # 'hope to avoid too many levels of symbolic links' error
|
|
if f.is_dir():
|
|
subfolders.append(f.path)
|
|
elif f.is_file():
|
|
is_hidden = f.name.split("/")[-1][0] == '.'
|
|
has_ext = os.path.splitext(f.name)[1].lower() in ext
|
|
name_lower = f.name.lower()
|
|
has_keyword = any(
|
|
[keyword in name_lower for keyword in keywords])
|
|
has_banned = any(
|
|
[banned_word in name_lower for banned_word in banned_words])
|
|
if has_ext and has_keyword and not has_banned and not is_hidden and not os.path.basename(f.path).startswith("._"):
|
|
files.append(f.path)
|
|
except:
|
|
pass
|
|
except:
|
|
pass
|
|
|
|
for dir in list(subfolders):
|
|
sf, f = keyword_scandir(dir, ext, keywords)
|
|
subfolders.extend(sf)
|
|
files.extend(f)
|
|
return subfolders, files
|
|
|
|
def get_audio_filenames(
|
|
paths: list, # directories in which to search
|
|
keywords=None,
|
|
exts=['.wav', '.mp3', '.flac', '.ogg', '.aif', '.opus'],
|
|
filelist_path=None
|
|
):
|
|
"recursively get a list of audio filenames"
|
|
filenames = []
|
|
if type(paths) is str:
|
|
paths = [paths]
|
|
for path in paths: # get a list of relevant filenames
|
|
|
|
if filelist_path is None:
|
|
# Check for filelist.txt at the root of the directory
|
|
filelist_path = os.path.join(path, "filelist.txt")
|
|
|
|
if os.path.exists(filelist_path):
|
|
with open(filelist_path, "r") as f:
|
|
files = f.readlines()
|
|
files = [os.path.join(path, file.strip()) for file in files]
|
|
filenames.extend(files)
|
|
continue
|
|
|
|
if keywords is not None:
|
|
subfolders, files = keyword_scandir(path, exts, keywords)
|
|
else:
|
|
subfolders, files = fast_scandir(path, exts)
|
|
filenames.extend(files)
|
|
return filenames
|
|
|
|
def get_latent_filenames(
|
|
paths, # directories in which to search
|
|
extension='npy',
|
|
filelist_path=None
|
|
):
|
|
"recursively get a list of pre-encoded filenames"
|
|
filenames = []
|
|
if type(paths) is str:
|
|
paths = [paths]
|
|
for path in paths: # get a list of relevant filenames
|
|
|
|
if filelist_path is None:
|
|
# Check for filelist.txt at the root of the directory
|
|
filelist_path = os.path.join(path, "filelist.txt")
|
|
|
|
if os.path.exists(filelist_path):
|
|
with open(filelist_path, "r") as f:
|
|
files = f.readlines()
|
|
files = [os.path.join(path, file.strip()) for file in files]
|
|
filenames.extend(files)
|
|
continue
|
|
|
|
_, files = fast_scandir(path, [extension])
|
|
filenames.extend(files)
|
|
|
|
# Filter out silence.npy (used for silence latent padding, not a data sample)
|
|
filenames = [f for f in filenames if os.path.basename(f) != "silence.npy"]
|
|
|
|
# Add metadata paths
|
|
filenames = [(filename, filename.replace(f".{extension}", ".json")) for filename in filenames]
|
|
|
|
return filenames
|
|
|
|
class LocalDatasetConfig:
|
|
def __init__(
|
|
self,
|
|
id: str,
|
|
path: str,
|
|
keywords: Optional[List[str]]=None,
|
|
custom_metadata_fn: Optional[Callable[[str], str]] = None,
|
|
filelist_path = None,
|
|
weight: float = 1.0,
|
|
):
|
|
self.id = id
|
|
self.path = path
|
|
self.custom_metadata_fn = custom_metadata_fn
|
|
self.keywords = keywords
|
|
self.filelist_path = filelist_path
|
|
self.weight = weight
|
|
|
|
class LatentDatasetConfig(LocalDatasetConfig):
|
|
def __init__(
|
|
self,
|
|
latent_extension: str = "npy",
|
|
filelist_path = None,
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
super().__init__(*args, **kwargs)
|
|
self.latent_extension = latent_extension
|
|
self.filelist_path = filelist_path
|
|
# weight is inherited from LocalDatasetConfig via **kwargs
|
|
|
|
class SampleDataset(torch.utils.data.Dataset):
|
|
def __init__(
|
|
self,
|
|
configs,
|
|
sample_size=65536,
|
|
sample_rate=48000,
|
|
random_crop=True,
|
|
force_channels="stereo",
|
|
volume_norm=False,
|
|
volume_norm_param=(-16, 2),
|
|
strip_silence=False,
|
|
pad=True,
|
|
):
|
|
super().__init__()
|
|
self.filenames = []
|
|
self.sample_weights = []
|
|
|
|
self.augs = torch.nn.Sequential(
|
|
PhaseFlipper(),
|
|
#nn.Identity()
|
|
)
|
|
|
|
|
|
self.root_paths = []
|
|
|
|
self.pad_crop = PadCrop_Normalized_T(sample_size, sample_rate, randomize=random_crop, pad=pad)
|
|
self.strip_silence = strip_silence
|
|
|
|
self.force_channels = force_channels
|
|
|
|
self.encoding = torch.nn.Sequential(
|
|
Stereo() if self.force_channels == "stereo" else torch.nn.Identity(),
|
|
Mono() if self.force_channels == "mono" else torch.nn.Identity()
|
|
)
|
|
|
|
self.sr = sample_rate
|
|
|
|
self.volume_norm = VolumeNorm(volume_norm_param, self.sr) if volume_norm else torch.nn.Identity()
|
|
|
|
self.custom_metadata_fns = {}
|
|
|
|
for config in configs:
|
|
self.root_paths.append(config.path)
|
|
new_files = get_audio_filenames(config.path, config.keywords, filelist_path=config.filelist_path)
|
|
self.filenames.extend(new_files)
|
|
self.sample_weights.extend([config.weight] * len(new_files))
|
|
if config.custom_metadata_fn is not None:
|
|
self.custom_metadata_fns[config.path] = dill.dumps(config.custom_metadata_fn)
|
|
|
|
print(f'Found {len(self.filenames)} files')
|
|
|
|
def load_file(self, filename):
|
|
ext = filename.split(".")[-1]
|
|
|
|
audio, in_sr = torchaudio.load(filename, format=ext)
|
|
|
|
if in_sr != self.sr:
|
|
resample_tf = T.Resample(in_sr, self.sr)
|
|
audio = resample_tf(audio)
|
|
|
|
return audio
|
|
|
|
def __len__(self):
|
|
return len(self.filenames)
|
|
|
|
def __getitem__(self, idx):
|
|
audio_filename = self.filenames[idx]
|
|
try:
|
|
start_time = time.time()
|
|
audio = self.load_file(audio_filename)
|
|
|
|
audio = self.volume_norm(audio)
|
|
|
|
if self.strip_silence:
|
|
audio = strip_trailing_silence(audio, self.sr)
|
|
|
|
audio, t_start, t_end, seconds_start, seconds_total, padding_mask = self.pad_crop(audio)
|
|
|
|
# Check for silence
|
|
if is_silence(audio):
|
|
return self[random.randrange(len(self))]
|
|
|
|
# Run augmentations on this sample (including random crop)
|
|
if self.augs is not None:
|
|
audio = self.augs(audio)
|
|
|
|
audio = audio.clamp(-1, 1)
|
|
|
|
# Encode the file to assist in prediction
|
|
if self.encoding is not None:
|
|
audio = self.encoding(audio)
|
|
|
|
info = {}
|
|
|
|
info["path"] = audio_filename
|
|
|
|
for root_path in self.root_paths:
|
|
if root_path in audio_filename:
|
|
info["relpath"] = path.relpath(audio_filename, root_path)
|
|
|
|
info["timestamps"] = (t_start, t_end)
|
|
info["seconds_start"] = seconds_start
|
|
info["seconds_total"] = seconds_total
|
|
info["padding_mask"] = [padding_mask]
|
|
info["sample_rate"] = self.sr
|
|
|
|
end_time = time.time()
|
|
|
|
info["load_time"] = end_time - start_time
|
|
|
|
for custom_md_path in self.custom_metadata_fns.keys():
|
|
if custom_md_path in audio_filename:
|
|
custom_metadata_fn = dill.loads(self.custom_metadata_fns[custom_md_path])
|
|
custom_metadata = custom_metadata_fn(info, audio)
|
|
info.update(custom_metadata)
|
|
|
|
if "__reject__" in info and info["__reject__"]:
|
|
return self[random.randrange(len(self))]
|
|
|
|
# Provide audio inputs as their own dictionary to be merged into info, each audio element will be normalized in the same way as the main audio
|
|
if "__audio__" in info:
|
|
for audio_key, audio_value in info["__audio__"].items():
|
|
# Process the audio_value tensor, which should be a torch tensor
|
|
audio_value, _, _, _, _, _ = self.pad_crop(audio_value)
|
|
audio_value = audio_value.clamp(-1, 1)
|
|
if self.encoding is not None:
|
|
audio_value = self.encoding(audio_value)
|
|
info[audio_key] = audio_value
|
|
|
|
del info["__audio__"]
|
|
|
|
return (audio, info)
|
|
except Exception as e:
|
|
print(f'Couldn\'t load file {audio_filename}: {e}')
|
|
return self[random.randrange(len(self))]
|
|
|
|
|
|
class PreEncodedDataset(torch.utils.data.Dataset):
|
|
def __init__(
|
|
self,
|
|
configs: List[LatentDatasetConfig],
|
|
latent_crop_length=None,
|
|
min_length_sec=None,
|
|
max_length_sec=None,
|
|
random_crop=False,
|
|
tokenizers: Optional[dict] = None,
|
|
):
|
|
super().__init__()
|
|
self.filenames = []
|
|
self.sample_weights = []
|
|
|
|
self.custom_metadata_fns = {}
|
|
|
|
self.silence_latents = {}
|
|
|
|
for config in configs:
|
|
new_files = get_latent_filenames(config.path, config.latent_extension, config.filelist_path)
|
|
self.filenames.extend(new_files)
|
|
self.sample_weights.extend([config.weight] * len(new_files))
|
|
if config.custom_metadata_fn is not None:
|
|
self.custom_metadata_fns[config.path] = dill.dumps(config.custom_metadata_fn)
|
|
|
|
# Load silence latent if available (for variable-length padding)
|
|
paths = config.path if isinstance(config.path, list) else [config.path]
|
|
for path in paths:
|
|
silence_path = os.path.join(path, "silence.npy")
|
|
if os.path.exists(silence_path):
|
|
self.silence_latents[path] = np.load(silence_path).squeeze(0) # [C, N]
|
|
print(f'Loaded silence latent from {silence_path}')
|
|
|
|
self.latent_crop_length = latent_crop_length
|
|
self.random_crop = random_crop
|
|
|
|
self.min_length_sec = min_length_sec
|
|
self.max_length_sec = max_length_sec
|
|
|
|
# tokenizers: dict mapping metadata key -> (tokenizer, max_length)
|
|
# If provided, text fields will be pre-tokenized in DataLoader workers
|
|
self.tokenizers = tokenizers
|
|
|
|
print(f'Found {len(self.filenames)} files')
|
|
|
|
def __len__(self):
|
|
return len(self.filenames)
|
|
|
|
def _get_silence_for_file(self, latent_filename):
|
|
"""Return the silence latent for the dataset that contains this file, or None."""
|
|
for path, silence in self.silence_latents.items():
|
|
if path in latent_filename:
|
|
return silence
|
|
return None
|
|
|
|
def __getitem__(self, idx):
|
|
latent_filename, md_filename = self.filenames[idx]
|
|
try:
|
|
latents = torch.from_numpy(np.load(latent_filename)) # [C, N]
|
|
|
|
with open(md_filename, "r") as f:
|
|
try:
|
|
info = json.load(f)
|
|
except:
|
|
raise Exception(f"Couldn't load metadata file {md_filename}")
|
|
|
|
info["latent_filename"] = latent_filename
|
|
|
|
if self.latent_crop_length is not None:
|
|
stored_length = latents.shape[1]
|
|
|
|
if stored_length > self.latent_crop_length:
|
|
# Crop to latent_crop_length (existing logic)
|
|
# Get the last index from the padding mask, the index of the last 1 in the sequence
|
|
last_ix = len(info["padding_mask"]) - 1 - info["padding_mask"][::-1].index(1)
|
|
|
|
if self.random_crop and last_ix > self.latent_crop_length:
|
|
start = random.randint(0, last_ix - self.latent_crop_length)
|
|
else:
|
|
start = 0
|
|
|
|
latents = latents[:, start:start+self.latent_crop_length]
|
|
info["padding_mask"] = info["padding_mask"][start:start+self.latent_crop_length]
|
|
info["latent_crop_start"] = start
|
|
|
|
elif stored_length < self.latent_crop_length:
|
|
# Pad with silence latent to reach latent_crop_length
|
|
pad_needed = self.latent_crop_length - stored_length
|
|
silence = self._get_silence_for_file(latent_filename)
|
|
|
|
if silence is not None:
|
|
# Slice or tile silence latent to cover pad_needed frames
|
|
if silence.shape[1] >= pad_needed:
|
|
silence_pad = silence[:, :pad_needed]
|
|
else:
|
|
silence_pad = np.tile(silence, (1, (pad_needed // silence.shape[1]) + 1))[:, :pad_needed]
|
|
latents = torch.cat([latents, torch.from_numpy(silence_pad)], dim=1)
|
|
else:
|
|
# No silence latent available — zero-pad as fallback
|
|
latents = torch.nn.functional.pad(latents, (0, pad_needed))
|
|
|
|
# Build padding_mask: valid frames from stored mask, zeros for padding
|
|
info["padding_mask"] = info["padding_mask"][:stored_length] + [0] * pad_needed
|
|
info["latent_crop_start"] = 0
|
|
|
|
else:
|
|
# Exact match
|
|
info["latent_crop_start"] = 0
|
|
|
|
info["latent_crop_length"] = self.latent_crop_length
|
|
|
|
info["padding_mask"] = [torch.tensor(info["padding_mask"])]
|
|
|
|
seconds_total = info["seconds_total"]
|
|
|
|
if self.min_length_sec is not None and seconds_total < self.min_length_sec:
|
|
return self[random.randrange(len(self))]
|
|
|
|
if self.max_length_sec is not None and seconds_total > self.max_length_sec:
|
|
return self[random.randrange(len(self))]
|
|
|
|
for custom_md_path in self.custom_metadata_fns.keys():
|
|
if custom_md_path in latent_filename:
|
|
custom_metadata_fn = dill.loads(self.custom_metadata_fns[custom_md_path])
|
|
custom_metadata = custom_metadata_fn(info, latents)
|
|
info.update(custom_metadata)
|
|
|
|
if "__reject__" in info and info["__reject__"]:
|
|
return self[random.randrange(len(self))]
|
|
|
|
if "__replace__" in info and info["__replace__"] is not None:
|
|
# Replace the latents with the new latents if the custom metadata function returns a new set of latents
|
|
latents = info["__replace__"]
|
|
|
|
info["audio"] = latents
|
|
|
|
# Pre-tokenize text fields in DataLoader workers to avoid
|
|
# CPU contention with the main training thread
|
|
if self.tokenizers is not None:
|
|
for key, (tokenizer, max_length) in self.tokenizers.items():
|
|
if key in info and isinstance(info[key], str):
|
|
# Save raw text before replacing with tokens (needed by CLAP and other text-based losses)
|
|
info[f"{key}_text"] = info[key]
|
|
encoded = tokenizer(
|
|
info[key],
|
|
truncation=True,
|
|
max_length=max_length,
|
|
padding="max_length",
|
|
return_tensors="pt",
|
|
)
|
|
info[key] = {
|
|
"input_ids": encoded["input_ids"].squeeze(0),
|
|
"attention_mask": encoded["attention_mask"].squeeze(0),
|
|
}
|
|
|
|
return (latents, info)
|
|
except Exception as e:
|
|
print(f'Couldn\'t load file {latent_filename}: {e}')
|
|
return self[random.randrange(len(self))]
|
|
|
|
# get_dbmax and is_silence copied from https://github.com/drscotthawley/aeiou/blob/main/aeiou/core.py under Apache 2.0 License
|
|
# License can be found in LICENSES/LICENSE_AEIOU.txt
|
|
def get_dbmax(
|
|
audio, # torch tensor of (multichannel) audio
|
|
):
|
|
"finds the loudest value in the entire clip and puts that into dB (full scale)"
|
|
return 20*torch.log10(torch.flatten(audio.abs()).max()).cpu().numpy()
|
|
|
|
def is_silence(
|
|
audio, # torch tensor of (multichannel) audio
|
|
thresh=-60, # threshold in dB below which we declare to be silence
|
|
):
|
|
"checks if entire clip is 'silence' below some dB threshold"
|
|
dBmax = get_dbmax(audio)
|
|
return dBmax < thresh
|
|
|
|
|
|
def collation_fn(samples):
|
|
batched = list(zip(*samples))
|
|
result = []
|
|
for b in batched:
|
|
if isinstance(b[0], (int, float)):
|
|
b = np.array(b)
|
|
elif isinstance(b[0], torch.Tensor):
|
|
b = torch.stack(b)
|
|
elif isinstance(b[0], np.ndarray):
|
|
b = np.array(b)
|
|
else:
|
|
b = b
|
|
result.append(b)
|
|
return result
|
|
|
|
|