- 共享引擎(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 模板
442 lines
16 KiB
Python
442 lines
16 KiB
Python
"""Pre-encode audio files into SAME latents for LoRA finetuning (MLX, torch-free).
|
|
|
|
Offline audio → latent CLI mirroring underfit's ``dataset_processing/pre_encode.py``
|
|
conventions, ported to the standalone MLX runtime: single-process, soundfile-based
|
|
loading (no torch / torchaudio), and ``models.defs.audio_encoding.encode_audio``
|
|
doing the padding / chunking / masking.
|
|
|
|
Every audio file under ``--audio-dir`` (recursive) becomes a ``.npy`` + ``.json``
|
|
pair at the same relative path under ``--output-dir``:
|
|
|
|
<output-dir>/<relpath-with-npy-suffix> latents [D, T] float32
|
|
<output-dir>/<relpath-with-json-suffix> metadata (underfit field parity)
|
|
|
|
Conventions matched to underfit:
|
|
* Whole files, no clip splitting; cap at ``--max-duration`` seconds (default
|
|
600) with the sample cap aligned DOWN to a multiple of 4096 (ds_ratio).
|
|
* 44.1 kHz stereo — mono is duplicated to stereo, >2 channels keep the first
|
|
two, other sample rates are linearly resampled. NO loudness/peak norm.
|
|
* Encoder runs fp32; ``chunked=True`` for clips longer than 30 s.
|
|
* Metadata: ``path, relpath, src_relpath, seconds_total`` (rounded 3dp),
|
|
``seconds_start: 0, audio_samples, latent_shape``, latent-resolution 0/1
|
|
``padding_mask``, merged with tags extracted from sidecars.
|
|
* Skip files whose outputs already exist unless ``--overwrite``.
|
|
* ``<output-dir>/details.json`` records the global encode params.
|
|
|
|
Latent scaling note: the MLX SAME encoders output DiT-space (softnorm) latents
|
|
directly — the equivalent of torch's unscaled-latents-then-divide-by-scale
|
|
convention with scale 1.0 — so latents are saved as-is and the trainer uses
|
|
them directly (no pretransform-scale division at train time).
|
|
|
|
Tag extraction (ported from underfit's ``extract_tags``): JSON sidecar
|
|
(``<stem>.json`` next to the audio, or ``<parent>/json/<stem>.json``) keeps all
|
|
non-empty string/number values; else a ``.txt`` sidecar (same dir or
|
|
``<parent>/txt/``) whose whole stripped content becomes the ``prompt`` tag.
|
|
Embedded ID3/Vorbis tags are a DELIBERATE OMISSION — they need an extra
|
|
dependency (``audio_metadata``), and sidecars are the underfit-recommended
|
|
path anyway.
|
|
|
|
Usage:
|
|
python pre_encode_mlx.py --audio-dir ~/clips --output-dir ~/latents \
|
|
--codec same-s [--max-duration 600] [--overwrite]
|
|
|
|
``--codec same-s`` pairs with the sm-music / sm-sfx DiTs, ``same-l`` with medium.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
REPO = Path(__file__).resolve().parent.parent # optimized/mlx (scripts/ is one level down)
|
|
if str(REPO) not in sys.path:
|
|
sys.path.insert(0, str(REPO))
|
|
if str(REPO / "scripts") not in sys.path:
|
|
sys.path.insert(0, str(REPO / "scripts"))
|
|
|
|
import mlx.core as mx # noqa: E402
|
|
|
|
from models.defs.audio_encoding import ( # noqa: E402
|
|
SAME_ENCODER_PAD_MODULO,
|
|
SAMPLES_PER_LATENT,
|
|
encode_audio,
|
|
)
|
|
|
|
SAMPLE_RATE = 44100
|
|
DS_RATIO = SAMPLES_PER_LATENT # 4096 samples per latent
|
|
CHUNKED_THRESHOLD_SECONDS = 30.0
|
|
|
|
AUDIO_EXTS = {".wav", ".mp3", ".flac", ".ogg", ".aif", ".aiff", ".m4a"}
|
|
_MIN_AUDIO_SIZE = 4096 # skip files smaller than this (macOS resource forks, corrupt)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Audio discovery
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def find_audio_files(root) -> list[Path]:
|
|
"""Recursively find audio files under ``root``, sorted by path.
|
|
|
|
Skips macOS resource forks (``._*``) and files too small to be real audio.
|
|
"""
|
|
files = []
|
|
for dirpath, _, filenames in os.walk(root):
|
|
for fn in filenames:
|
|
if fn.startswith("._"):
|
|
continue
|
|
fp = Path(dirpath) / fn
|
|
if fp.suffix.lower() in AUDIO_EXTS:
|
|
try:
|
|
if fp.stat().st_size < _MIN_AUDIO_SIZE:
|
|
continue
|
|
except OSError:
|
|
continue
|
|
files.append(fp)
|
|
files.sort()
|
|
return files
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Audio loading (torch-free)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_audio(path) -> np.ndarray:
|
|
"""Load an audio file → ``(2, T)`` float32 at 44.1 kHz.
|
|
|
|
soundfile decodes; mono is duplicated to stereo, >2 channels keep the
|
|
first two; non-44.1 kHz input is linearly resampled at
|
|
align_corners=False positions (same pattern as sa3_gradio.read_audio_any).
|
|
No loudness/peak normalization.
|
|
"""
|
|
import soundfile as sf
|
|
|
|
data, sr = sf.read(str(path), dtype="float32", always_2d=True) # (T, ch)
|
|
y = data.T
|
|
y = np.stack([y[0], y[0]]) if y.shape[0] == 1 else y[:2]
|
|
if sr != SAMPLE_RATE:
|
|
in_len = y.shape[-1]
|
|
new_len = int(round(in_len * SAMPLE_RATE / sr))
|
|
scale = in_len / new_len
|
|
pos = np.clip((np.arange(new_len) + 0.5) * scale - 0.5, 0, in_len - 1)
|
|
y = np.stack([np.interp(pos, np.arange(in_len), ch) for ch in y])
|
|
return np.ascontiguousarray(y, dtype=np.float32)
|
|
|
|
|
|
def max_samples_for_duration(max_duration: float) -> int:
|
|
"""Sample cap for ``max_duration`` seconds, aligned DOWN to ds_ratio (4096)."""
|
|
return (int(max_duration * SAMPLE_RATE) // DS_RATIO) * DS_RATIO
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Sidecar tag extraction (ported from underfit pre_encode.extract_tags;
|
|
# embedded ID3/Vorbis tags deliberately omitted — see module docstring)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def extract_tags(filepath) -> dict:
|
|
"""Extract tag fields for an audio file from sidecar files.
|
|
|
|
Priority:
|
|
1. JSON sidecar (``<stem>.json`` same dir, or ``<parent>/json/<stem>.json``)
|
|
— all non-empty string/number values kept (custom keys preserved).
|
|
2. Plain ``.txt`` sidecar (same dir or ``<parent>/txt/``) — whole stripped
|
|
content becomes the ``prompt`` tag.
|
|
Returns ``{}`` when neither sidecar exists (no embedded-tag fallback).
|
|
"""
|
|
fp = Path(filepath)
|
|
stem = fp.stem
|
|
|
|
# 1. JSON sidecar — same dir, then sibling json/ dir
|
|
sidecar_candidates = [
|
|
fp.with_suffix(".json"),
|
|
fp.parent.parent / "json" / (stem + ".json"),
|
|
]
|
|
for sidecar in sidecar_candidates:
|
|
if sidecar.exists():
|
|
try:
|
|
with open(sidecar) as f:
|
|
sc = json.load(f)
|
|
tags_out = {
|
|
k: str(v)
|
|
for k, v in sc.items()
|
|
if v and isinstance(v, (str, int, float))
|
|
}
|
|
if tags_out:
|
|
return tags_out
|
|
except Exception:
|
|
pass
|
|
|
|
# 2. Plain .txt sidecar (SA3 convention) — same dir, then sibling txt/ dir.
|
|
txt_candidates = [
|
|
fp.with_suffix(".txt"),
|
|
fp.parent.parent / "txt" / (stem + ".txt"),
|
|
]
|
|
for txt in txt_candidates:
|
|
if txt.exists():
|
|
try:
|
|
content = txt.read_text(encoding="utf-8", errors="replace").strip()
|
|
if content:
|
|
return {"prompt": content}
|
|
except Exception:
|
|
pass
|
|
|
|
return {}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Core encode
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def encode_file(
|
|
encoder,
|
|
audio_np: np.ndarray,
|
|
*,
|
|
pad_modulo: int,
|
|
actual_samples: int | None = None,
|
|
):
|
|
"""Encode one ``(2, T)`` float32 waveform → (latents, padding_mask, seconds).
|
|
|
|
Returns:
|
|
latents_np: ``[D, T_latent]`` float32 (softnorm/DiT space, saved as-is)
|
|
padding_mask: latent-resolution 0/1 int list, length ``T_latent`` —
|
|
zeros over the alignment-padding tail
|
|
seconds_total: ``round(actual_samples / 44100, 3)``
|
|
|
|
``chunked=True`` is used automatically for clips longer than 30 s.
|
|
"""
|
|
audio_np = np.asarray(audio_np)
|
|
if audio_np.ndim != 2 or audio_np.shape[0] != 2:
|
|
raise ValueError(f"Expected audio shaped (2, T), got {audio_np.shape}.")
|
|
if actual_samples is None:
|
|
actual_samples = int(audio_np.shape[-1])
|
|
duration = actual_samples / SAMPLE_RATE
|
|
encoded = encode_audio(
|
|
encoder,
|
|
audio_np[None, ...],
|
|
valid_sample_lengths=[actual_samples],
|
|
pad_modulo=pad_modulo,
|
|
chunked=duration > CHUNKED_THRESHOLD_SECONDS,
|
|
)
|
|
latents_np = np.asarray(encoded.latents.astype(mx.float32))[0] # [D, T]
|
|
padding_mask = [int(v) for v in np.asarray(encoded.padding_mask)[0]]
|
|
seconds_total = round(actual_samples / SAMPLE_RATE, 3)
|
|
return latents_np, padding_mask, seconds_total
|
|
|
|
|
|
def build_metadata(
|
|
fpath,
|
|
rel,
|
|
out_rel,
|
|
*,
|
|
actual_samples: int,
|
|
latent_shape,
|
|
padding_mask,
|
|
seconds_total,
|
|
) -> dict:
|
|
"""Underfit-parity metadata dict for one encoded file (tags merged last)."""
|
|
meta = {
|
|
"path": str(fpath),
|
|
"relpath": str(out_rel),
|
|
"src_relpath": str(rel),
|
|
"seconds_total": seconds_total,
|
|
"seconds_start": 0,
|
|
"audio_samples": int(actual_samples),
|
|
"latent_shape": [int(s) for s in latent_shape],
|
|
"padding_mask": padding_mask,
|
|
}
|
|
meta.update(extract_tags(fpath))
|
|
return meta
|
|
|
|
|
|
def run(
|
|
audio_dir,
|
|
output_dir,
|
|
codec: str,
|
|
encoder,
|
|
*,
|
|
max_duration: float = 600.0,
|
|
overwrite: bool = False,
|
|
exclude: set | None = None,
|
|
) -> dict:
|
|
"""Encode every audio file under ``audio_dir`` into ``output_dir``.
|
|
|
|
Returns stats: ``{"encoded", "skipped", "errors", "count"}``.
|
|
"""
|
|
if codec not in SAME_ENCODER_PAD_MODULO:
|
|
raise ValueError(f"Unknown codec {codec!r}; expected same-s or same-l.")
|
|
pad_modulo = SAME_ENCODER_PAD_MODULO[codec]
|
|
|
|
audio_dir = Path(audio_dir).expanduser().resolve()
|
|
output_dir = Path(output_dir).expanduser().resolve()
|
|
files = find_audio_files(audio_dir)
|
|
if exclude:
|
|
files = [f for f in files if str(f.relative_to(audio_dir)) not in exclude]
|
|
max_samples = max_samples_for_duration(max_duration)
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
details = {
|
|
"codec": codec,
|
|
"sample_rate": SAMPLE_RATE,
|
|
"max_duration": max_duration,
|
|
"max_samples": max_samples,
|
|
"pad_modulo": pad_modulo,
|
|
"audio_dir": str(audio_dir),
|
|
"count": len(files),
|
|
}
|
|
with open(output_dir / "details.json", "w") as f:
|
|
json.dump(details, f, indent=2)
|
|
|
|
total = len(files)
|
|
encoded = skipped = errors = 0
|
|
t0 = time.time()
|
|
for index, fpath in enumerate(files, 1):
|
|
rel = fpath.relative_to(audio_dir)
|
|
out_rel = rel.with_suffix(".npy")
|
|
npy_path = output_dir / out_rel
|
|
json_path = npy_path.with_suffix(".json")
|
|
tag = f"[{index}/{total}]"
|
|
if npy_path.exists() and json_path.exists() and not overwrite:
|
|
skipped += 1
|
|
print(f" {tag} SKIP {out_rel} (exists)")
|
|
continue
|
|
try:
|
|
audio = load_audio(fpath)
|
|
actual_samples = int(audio.shape[-1])
|
|
if actual_samples > max_samples:
|
|
audio = audio[:, :max_samples]
|
|
actual_samples = max_samples
|
|
|
|
latents_np, padding_mask, seconds_total = encode_file(
|
|
encoder, audio, pad_modulo=pad_modulo, actual_samples=actual_samples
|
|
)
|
|
|
|
npy_path.parent.mkdir(parents=True, exist_ok=True)
|
|
np.save(str(npy_path), latents_np)
|
|
meta = build_metadata(
|
|
fpath,
|
|
rel,
|
|
out_rel,
|
|
actual_samples=actual_samples,
|
|
latent_shape=latents_np.shape,
|
|
padding_mask=padding_mask,
|
|
seconds_total=seconds_total,
|
|
)
|
|
with open(json_path, "w") as f:
|
|
json.dump(meta, f)
|
|
|
|
encoded += 1
|
|
shape_str = "x".join(str(s) for s in latents_np.shape)
|
|
print(f" {tag} {out_rel} {seconds_total:.1f}s -> [{shape_str}]")
|
|
except Exception as e:
|
|
errors += 1
|
|
print(f" {tag} ERROR {rel}: {e}")
|
|
|
|
elapsed = time.time() - t0
|
|
mins, secs = divmod(int(elapsed), 60)
|
|
print(
|
|
f"\nDone in {mins}m{secs:02d}s: {encoded} encoded, {skipped} skipped, "
|
|
f"{errors} errors (of {total})"
|
|
)
|
|
print(f"Output: {output_dir}")
|
|
return {"encoded": encoded, "skipped": skipped, "errors": errors, "count": total}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_codec_encoder(codec: str):
|
|
"""Load the real SAME encoder for ``codec`` in fp32. Returns (encoder, pad_modulo)."""
|
|
from sa3_mlx import load_encoder # lazy: pulls weight-download machinery
|
|
|
|
return load_encoder(codec, mx.float32)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Pre-encode audio into SAME latents for MLX LoRA finetuning",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
parser.add_argument(
|
|
"--audio-dir", "-i", required=True,
|
|
help="Directory of audio files (searched recursively)",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir", "-o", required=True,
|
|
help="Output root; each audio file becomes <output-dir>/<relpath>.npy + .json",
|
|
)
|
|
parser.add_argument(
|
|
"--codec", required=True, choices=sorted(SAME_ENCODER_PAD_MODULO),
|
|
help="Autoencoder: same-s pairs with sm-music/sm-sfx, same-l with medium",
|
|
)
|
|
parser.add_argument(
|
|
"--max-duration", type=float, default=600.0,
|
|
help="Crop files longer than N seconds (default: 600; aligned down to 4096 samples)",
|
|
)
|
|
parser.add_argument(
|
|
"--overwrite", action="store_true",
|
|
help="Re-encode files whose outputs already exist",
|
|
)
|
|
parser.add_argument(
|
|
"--exclude-file", type=str, default=None,
|
|
help="Text file of relpaths (one per line, relative to --audio-dir) to skip",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
audio_dir = Path(args.audio_dir).expanduser().resolve()
|
|
if not audio_dir.is_dir():
|
|
print(f"Error: not a directory: {audio_dir}")
|
|
sys.exit(1)
|
|
files = find_audio_files(audio_dir)
|
|
if not files:
|
|
print(f"No audio files found in {audio_dir}")
|
|
sys.exit(1)
|
|
|
|
exclude = set()
|
|
if args.exclude_file and Path(args.exclude_file).is_file():
|
|
exclude = {ln.strip() for ln in Path(args.exclude_file).read_text().splitlines()
|
|
if ln.strip()}
|
|
n_before = len(files)
|
|
if exclude:
|
|
files = [f for f in files if str(f.relative_to(audio_dir)) not in exclude]
|
|
if not files:
|
|
print(f"All {n_before} file(s) excluded by {args.exclude_file}")
|
|
sys.exit(1)
|
|
|
|
max_samples = max_samples_for_duration(args.max_duration)
|
|
import mlx.core as mx
|
|
print(f"Engine: MLX · device: {mx.default_device()}")
|
|
print(f"Codec: {args.codec} (pad_modulo={SAME_ENCODER_PAD_MODULO[args.codec]}, fp32)")
|
|
print(f"Input: {audio_dir} ({len(files)} audio files"
|
|
+ (f", {n_before - len(files)} excluded)" if exclude else ")"))
|
|
print(f"Output: {Path(args.output_dir).expanduser().resolve()}")
|
|
print(f"Max: {max_samples / SAMPLE_RATE:.1f}s ({max_samples:,} samples)\n")
|
|
|
|
print("Loading encoder...")
|
|
encoder, pad_modulo = load_codec_encoder(args.codec)
|
|
assert pad_modulo == SAME_ENCODER_PAD_MODULO[args.codec]
|
|
|
|
run(
|
|
audio_dir,
|
|
args.output_dir,
|
|
args.codec,
|
|
encoder,
|
|
max_duration=args.max_duration,
|
|
overwrite=args.overwrite,
|
|
exclude=exclude,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|