初始提交：3 款纯音频游戏 + 共享引擎 + 完整文档与资产管线
- 共享引擎（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 模板
This commit is contained in:
SoundGame Dev
@@ -0,0 +1,510 @@
|
||||
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
|
||||
|
||||
|
||||
@@ -0,0 +1,257 @@
|
||||
import math
|
||||
import random
|
||||
import torch
|
||||
|
||||
from torch import nn
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from torchaudio import transforms as T
|
||||
|
||||
class PadCrop(nn.Module):
|
||||
def __init__(self, n_samples, randomize=True):
|
||||
super().__init__()
|
||||
self.n_samples = n_samples
|
||||
self.randomize = randomize
|
||||
|
||||
def __call__(self, signal):
|
||||
n, s = signal.shape
|
||||
start = 0 if (not self.randomize) else torch.randint(0, max(0, s - self.n_samples) + 1, []).item()
|
||||
end = start + self.n_samples
|
||||
output = signal.new_zeros([n, self.n_samples])
|
||||
output[:, :min(s, self.n_samples)] = signal[:, start:end]
|
||||
return output
|
||||
|
||||
class PadCrop_Normalized_T(nn.Module):
|
||||
|
||||
def __init__(self, n_samples: int, sample_rate: int, randomize: bool = True, pad: bool = True):
|
||||
|
||||
super().__init__()
|
||||
|
||||
self.n_samples = n_samples
|
||||
self.sample_rate = sample_rate
|
||||
self.randomize = randomize
|
||||
self.pad = pad
|
||||
|
||||
def __call__(self, source: torch.Tensor) -> Tuple[torch.Tensor, float, float, int, int, torch.Tensor]:
|
||||
|
||||
n_channels, n_samples = source.shape
|
||||
|
||||
# Calculate bounds and offset
|
||||
upper_bound = max(0, n_samples - self.n_samples)
|
||||
offset = 0
|
||||
if self.randomize and n_samples > self.n_samples:
|
||||
offset = random.randint(0, upper_bound)
|
||||
|
||||
# Calculate normalized times
|
||||
norm_denom = upper_bound + self.n_samples
|
||||
t_start = offset / norm_denom
|
||||
t_end = (offset + self.n_samples) / norm_denom
|
||||
|
||||
# Calculate timing info
|
||||
seconds_start = math.floor(offset / self.sample_rate)
|
||||
seconds_total = math.ceil(n_samples / self.sample_rate)
|
||||
|
||||
# Optimize for different cases
|
||||
if n_samples >= self.n_samples:
|
||||
# No padding needed - use view (zero-copy)
|
||||
chunk = source[:, offset:offset + self.n_samples]
|
||||
# Create full mask efficiently
|
||||
padding_mask = torch.ones(self.n_samples, dtype=source.dtype, device=source.device)
|
||||
elif not self.pad:
|
||||
# No padding mode - return audio at natural length
|
||||
chunk = source
|
||||
padding_mask = torch.ones(n_samples, dtype=source.dtype, device=source.device)
|
||||
else:
|
||||
# Padding needed - create chunk and fill in-place
|
||||
chunk = torch.zeros(n_channels, self.n_samples, dtype=source.dtype, device=source.device)
|
||||
chunk[:, :n_samples] = source # Use in-place assignment
|
||||
|
||||
# Create padding mask in-place
|
||||
padding_mask = torch.zeros(self.n_samples, dtype=source.dtype, device=source.device)
|
||||
padding_mask[:n_samples] = 1
|
||||
|
||||
return (
|
||||
chunk,
|
||||
t_start,
|
||||
t_end,
|
||||
seconds_start,
|
||||
seconds_total,
|
||||
padding_mask
|
||||
)
|
||||
|
||||
def strip_trailing_silence(audio, sample_rate, threshold_db=-60, min_silence_duration=0.1):
|
||||
"""Strip silence from the end of an audio tensor.
|
||||
|
||||
Args:
|
||||
audio: tensor [channels, samples]
|
||||
sample_rate: audio sample rate
|
||||
threshold_db: dB threshold below which audio is considered silent
|
||||
min_silence_duration: minimum trailing silence duration in seconds to strip
|
||||
Returns:
|
||||
Truncated audio tensor [channels, trimmed_samples], or original if no significant trailing silence
|
||||
"""
|
||||
n_samples = audio.shape[-1]
|
||||
hop_length = max(1, int(sample_rate * 0.01)) # 10ms frames
|
||||
min_silence_samples = int(sample_rate * min_silence_duration)
|
||||
n_frames = n_samples // hop_length
|
||||
|
||||
if n_frames == 0:
|
||||
return audio
|
||||
|
||||
# Work in float32 for precision
|
||||
audio_f = audio.float()
|
||||
|
||||
# Reshape into frames and compute max absolute amplitude per frame across channels
|
||||
trimmed = audio_f[:, :n_frames * hop_length]
|
||||
frames = trimmed.reshape(audio_f.shape[0], n_frames, hop_length)
|
||||
frame_peak = frames.abs().amax(dim=(0, 2)) # [n_frames] - max across channels and samples
|
||||
frame_db = 20 * torch.log10(frame_peak + 1e-10)
|
||||
|
||||
# Find last frame above threshold
|
||||
above_thresh = (frame_db > threshold_db).nonzero(as_tuple=True)[0]
|
||||
|
||||
if len(above_thresh) == 0:
|
||||
# Entire audio is silent
|
||||
return audio[:, :0]
|
||||
|
||||
last_active_frame = above_thresh[-1].item()
|
||||
content_end = min((last_active_frame + 1) * hop_length, n_samples)
|
||||
|
||||
# Only strip if trailing silence is long enough
|
||||
if (n_samples - content_end) < min_silence_samples:
|
||||
return audio
|
||||
|
||||
return audio[:, :content_end]
|
||||
|
||||
|
||||
class PhaseFlipper(nn.Module):
|
||||
"Randomly invert the phase of a signal"
|
||||
def __init__(self, p=0.5):
|
||||
super().__init__()
|
||||
self.p = p
|
||||
def __call__(self, signal):
|
||||
return -signal if (random.random() < self.p) else signal
|
||||
|
||||
class Mono(nn.Module):
|
||||
def __call__(self, signal):
|
||||
return torch.mean(signal, dim=0, keepdims=True) if len(signal.shape) > 1 else signal
|
||||
|
||||
class Stereo(nn.Module):
|
||||
def __call__(self, signal):
|
||||
signal_shape = signal.shape
|
||||
# Check if it's mono
|
||||
if len(signal_shape) == 1: # s -> 2, s
|
||||
signal = signal.unsqueeze(0).repeat(2, 1)
|
||||
elif len(signal_shape) == 2:
|
||||
if signal_shape[0] == 1: #1, s -> 2, s
|
||||
signal = signal.repeat(2, 1)
|
||||
elif signal_shape[0] > 2: #?, s -> 2,s
|
||||
signal = signal[:2, :]
|
||||
|
||||
return signal
|
||||
|
||||
class VolumeNorm(nn.Module):
|
||||
"Volume normalization and augmentation of a signal [LUFS standard]"
|
||||
def __init__(self, params=[-16, 2], sample_rate=16000, energy_threshold=1e-6):
|
||||
super().__init__()
|
||||
self.loudness = T.Loudness(sample_rate)
|
||||
self.value = params[0]
|
||||
self.gain_range = [-params[1], params[1]]
|
||||
self.energy_threshold = energy_threshold
|
||||
|
||||
def __call__(self, signal):
|
||||
"""
|
||||
signal: torch.Tensor [channels, time]
|
||||
"""
|
||||
# avoid do normalisation for silence
|
||||
energy = torch.mean(signal**2)
|
||||
if energy < self.energy_threshold:
|
||||
return signal
|
||||
|
||||
input_loudness = self.loudness(signal)
|
||||
# Generate a random target loudness within the specified range
|
||||
target_loudness = self.value + (torch.rand(1).item() * (self.gain_range[1] - self.gain_range[0]) + self.gain_range[0])
|
||||
delta_loudness = target_loudness - input_loudness
|
||||
gain = torch.pow(10.0, delta_loudness / 20.0)
|
||||
output = gain * signal
|
||||
|
||||
# Check for potentially clipped samples
|
||||
if torch.max(torch.abs(output)) >= 1.0:
|
||||
output = self.declip(output)
|
||||
|
||||
return output
|
||||
|
||||
def declip(self, signal):
|
||||
"""
|
||||
Declip the signal by scaling down if any samples are clipped
|
||||
"""
|
||||
max_val = torch.max(torch.abs(signal))
|
||||
if max_val > 1.0:
|
||||
signal = signal / max_val
|
||||
signal *= 0.95
|
||||
return signal
|
||||
|
||||
|
||||
def create_padding_mask_from_lengths(
|
||||
valid_lengths: torch.Tensor,
|
||||
total_seq_len: int,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Create a boolean padding mask from per-batch valid sequence lengths.
|
||||
|
||||
Args:
|
||||
valid_lengths: Tensor of shape (batch_size,) with valid length per sample
|
||||
total_seq_len: Total sequence length of the latent
|
||||
|
||||
Returns:
|
||||
Boolean tensor of shape (batch_size, total_seq_len) where True = valid, False = padding
|
||||
"""
|
||||
device = valid_lengths.device
|
||||
positions = torch.arange(total_seq_len, device=device).unsqueeze(0) # (1, T)
|
||||
padding_mask = positions < valid_lengths.unsqueeze(1) # (B, T)
|
||||
return padding_mask
|
||||
|
||||
|
||||
def compute_effective_seq_len_from_conditioning(
|
||||
conditioning: list,
|
||||
sample_rate: int,
|
||||
downsampling_ratio: int = 1,
|
||||
device: str = "cuda"
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Compute effective sequence lengths from seconds_total in conditioning dicts.
|
||||
|
||||
Args:
|
||||
conditioning: List of conditioning dicts, one per batch element
|
||||
sample_rate: Audio sample rate
|
||||
downsampling_ratio: Pretransform downsampling ratio (1 if no pretransform)
|
||||
device: Device to place the tensor on
|
||||
|
||||
Returns:
|
||||
Tensor of shape (batch_size,) with effective sequence lengths in latent space,
|
||||
or None if seconds_total is not present in conditioning
|
||||
"""
|
||||
if conditioning is None:
|
||||
return None
|
||||
|
||||
# Check if seconds_total is present in any conditioning dict
|
||||
if not any("seconds_total" in c for c in conditioning):
|
||||
return None
|
||||
|
||||
effective_lengths = []
|
||||
for cond_dict in conditioning:
|
||||
if "seconds_total" in cond_dict:
|
||||
seconds = cond_dict["seconds_total"]
|
||||
# Convert seconds to latent sequence length
|
||||
audio_samples = int(seconds * sample_rate)
|
||||
latent_length = math.ceil(audio_samples / downsampling_ratio)
|
||||
effective_lengths.append(latent_length)
|
||||
else:
|
||||
# If seconds_total not present for this item, use None as marker
|
||||
effective_lengths.append(None)
|
||||
|
||||
# If any item is missing seconds_total, return None to fall back to full length
|
||||
if any(l is None for l in effective_lengths):
|
||||
return None
|
||||
|
||||
return torch.tensor(effective_lengths, dtype=torch.float32, device=device)
|
||||
Reference in New Issue
Block a user