| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
|
|
| import numpy as np |
| import soundfile as sf |
| import torch |
| from PIL import Image |
| from diffusers.pipelines.deprecated.audio_diffusion.mel import Mel |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer, ClapModel |
|
|
| from audio_dit import AudioDiT |
|
|
| MAX_TOKENS = 32 |
|
|
| @torch.no_grad() |
| def generate(model, seq, pool, null_seq, null_pool, steps, cfg, dev): |
| B = seq.shape[0] |
| x = torch.randn(B, 1, model.y_res, model.x_res, device=dev) |
| ns = null_seq.expand(B, -1, -1) |
| npool = null_pool.expand(B, -1) |
| dt = 1.0 / steps |
| for i in range(steps): |
| t = torch.full((B,), i * dt, device=dev) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| vc = model(x, t, seq, pool) |
| vu = model(x, t, ns, npool) |
| x = x + (vu + cfg * (vc - vu)).float() * dt |
| return x |
|
|
| def image_from_tensor(row): |
| arr = ((row.clamp(-1, 1).float().cpu().numpy() + 1) * 127.5 + 0.5).astype(np.uint8) |
| return Image.fromarray(arr) |
|
|
| def load_model(weights, config, dev): |
| cfgd = json.load(open(config)) |
| dit_cfg = cfgd["dit"] |
| model = AudioDiT(x_res=dit_cfg["x_res"], y_res=dit_cfg["y_res"], |
| text_seq_dim=dit_cfg["text_seq_dim"], |
| text_pool_dim=dit_cfg["text_pool_dim"]).to(dev).eval() |
| model.load_state_dict(load_file(weights)) |
| mel = Mel(x_res=dit_cfg["x_res"], y_res=dit_cfg["y_res"], sample_rate=cfgd["mel"]["sample_rate"], |
| n_fft=cfgd["mel"]["n_fft"], hop_length=cfgd["mel"]["hop_length"], top_db=cfgd["mel"]["top_db"]) |
| return model, mel, cfgd |
|
|
| def load_clap(name, dev): |
| tok = AutoTokenizer.from_pretrained(name) |
| clap = ClapModel.from_pretrained(name).to(dev).eval() |
|
|
| @torch.no_grad() |
| def enc(strings): |
| t = tok(strings, padding="max_length", max_length=MAX_TOKENS, truncation=True, |
| return_tensors="pt").to(dev) |
| out = clap.text_model(**t) |
| return out.last_hidden_state.float(), clap.text_projection(out.pooler_output).float() |
|
|
| return enc |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("prompts", nargs="+") |
| ap.add_argument("--out-dir", default="samples") |
| ap.add_argument("--cfg", type=float, default=4.0) |
| ap.add_argument("--steps", type=int, default=50) |
| ap.add_argument("--device", default="cuda") |
| ap.add_argument("--weights", default="/root/runs/audio_v1/model.safetensors") |
| ap.add_argument("--config", default="/root/runs/audio_v1/config.json") |
| ap.add_argument("--clap", default="laion/clap-htsat-unfused") |
| args = ap.parse_args() |
| dev = args.device |
| os.makedirs(args.out_dir, exist_ok=True) |
|
|
| model, mel, _ = load_model(args.weights, args.config, dev) |
| enc = load_clap(args.clap, dev) |
|
|
| seq, pool = enc(args.prompts) |
| null_seq, null_pool = enc([""]) |
| x = generate(model, seq, pool, null_seq, null_pool, args.steps, args.cfg, dev) |
|
|
| for prompt, row in zip(args.prompts, x[:, 0]): |
| img = image_from_tensor(row) |
| audio = mel.image_to_audio(img) |
| name = "_".join(prompt.lower().split())[:40] |
| sf.write(os.path.join(args.out_dir, f"{name}.wav"), audio, mel.get_sample_rate()) |
| print(f" wrote {name}.wav ({len(audio)/mel.get_sample_rate():.1f}s) {prompt}", flush=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
