File size: 5,354 Bytes
f48ab3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | from __future__ import annotations
import argparse
import csv
import json
import os
import time
from multiprocessing import Pool
import numpy as np
import torch
from diffusers.pipelines.deprecated.audio_diffusion.mel import Mel
from transformers import AutoTokenizer, ClapModel
X_RES, Y_RES = 384, 256
HOP, NFFT, SR, TOPDB = 1024, 2048, 22050, 80
MAX_TOKENS = 32
def read_captions(csv_path):
rows = []
with open(csv_path, newline="", encoding="utf-8", errors="replace") as f:
reader = csv.DictReader(f)
for row in reader:
fn = row["file_name"]
caps = [row[f"caption_{i}"] for i in range(1, 6) if row.get(f"caption_{i}")]
rows.append((fn, caps))
return rows
_mel = None
def _get_mel():
global _mel
if _mel is None:
_mel = Mel(x_res=X_RES, y_res=Y_RES, sample_rate=SR, n_fft=NFFT, hop_length=HOP, top_db=TOPDB)
return _mel
def to_mel_array(path):
try:
mel = _get_mel()
mel.load_audio(audio_file=path)
img = mel.audio_slice_to_image(0)
return path, np.array(img, dtype=np.uint8)
except Exception:
return path, None
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--audio-dir", required=True)
ap.add_argument("--captions-csv", required=True)
ap.add_argument("--split", required=True)
ap.add_argument("--out", default="/root/data")
ap.add_argument("--clap", default="laion/clap-htsat-unfused")
ap.add_argument("--workers", type=int, default=16)
ap.add_argument("--batch", type=int, default=128)
ap.add_argument("--device", default="cuda")
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
rows = read_captions(args.captions_csv)
print(f"[build] {args.split}: {len(rows)} clips listed in captions csv", flush=True)
paths = [os.path.join(args.audio_dir, fn) for fn, _ in rows]
missing = [p for p in paths if not os.path.exists(p)]
if missing:
print(f"[build] WARNING {len(missing)} missing audio files, e.g. {missing[:3]}", flush=True)
t0 = time.time()
mel_by_path, failed = {}, 0
with Pool(args.workers) as pool:
for i, (p, arr) in enumerate(pool.imap(to_mel_array, paths, chunksize=8)):
if arr is None:
failed += 1
else:
mel_by_path[p] = arr
if (i + 1) % 500 == 0:
print(f"[build] mel {i+1}/{len(paths)} failed={failed} {time.time()-t0:.0f}s", flush=True)
print(f"[build] mel done: {len(mel_by_path)} ok, {failed} failed, {time.time()-t0:.0f}s", flush=True)
clip_paths = list(mel_by_path.keys())
clip_index = {p: i for i, p in enumerate(clip_paths)}
mel_arr = np.stack([mel_by_path[p] for p in clip_paths])
print(f"[build] mel_arr {mel_arr.shape} {mel_arr.nbytes/2**20:.0f} MiB", flush=True)
pair_clip_idx, pair_captions = [], []
for fn, caps in rows:
p = os.path.join(args.audio_dir, fn)
if p not in clip_index:
continue
for c in caps:
pair_clip_idx.append(clip_index[p])
pair_captions.append(c)
print(f"[build] {len(pair_captions)} (clip, caption) pairs", flush=True)
dev = args.device
tok = AutoTokenizer.from_pretrained(args.clap)
clap = ClapModel.from_pretrained(args.clap).to(dev).eval()
pool_dim = clap.config.projection_dim
seq_dim = clap.config.text_config.hidden_size
@torch.no_grad()
def encode(strings):
enc = tok(strings, padding="max_length", truncation=True, max_length=MAX_TOKENS,
return_tensors="pt").to(dev)
out = clap.text_model(**enc)
seq = out.last_hidden_state.float()
pooled = clap.text_projection(out.pooler_output).float()
return seq.cpu().numpy().astype(np.float16), pooled.cpu().numpy().astype(np.float16)
seq_chunks, pool_chunks = [], []
t0 = time.time()
for i in range(0, len(pair_captions), args.batch):
chunk = pair_captions[i:i + args.batch]
s, p = encode(chunk)
seq_chunks.append(s)
pool_chunks.append(p)
if (i // args.batch) % 20 == 0:
print(f"[build] clap {i}/{len(pair_captions)} {time.time()-t0:.0f}s", flush=True)
text_seq = np.concatenate(seq_chunks)
text_pool = np.concatenate(pool_chunks)
print(f"[build] text_seq {text_seq.shape} text_pool {text_pool.shape}", flush=True)
np.save(f"{args.out}/{args.split}_mel.npy", mel_arr)
np.save(f"{args.out}/{args.split}_text_seq.npy", text_seq)
np.save(f"{args.out}/{args.split}_text_pool.npy", text_pool)
json.dump({"pair_clip_idx": pair_clip_idx, "captions": pair_captions,
"clip_files": [os.path.basename(p) for p in clip_paths],
"n_clips": len(clip_paths), "seq_dim": seq_dim, "pool_dim": pool_dim,
"x_res": X_RES, "y_res": Y_RES, "hop_length": HOP, "n_fft": NFFT,
"sample_rate": SR, "top_db": TOPDB, "max_tokens": MAX_TOKENS},
open(f"{args.out}/{args.split}_meta.json", "w"))
if args.split == "development":
ns, npz = encode([""])
np.save(f"{args.out}/null_seq.npy", ns[0])
np.save(f"{args.out}/null_pool.npy", npz[0])
print("[build] wrote null embeddings", flush=True)
print("BUILDDONE", flush=True)
if __name__ == "__main__":
main()
|