File size: 3,157 Bytes
9781faf | 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 | from __future__ import annotations
import json
import secrets
from itertools import pairwise
from pathlib import Path
import gradio as gr
import torch
from audio_dit import CHANNELS, FRAMES, MAX_TAGS, load_model
from same_l_decoder import SAMPLE_RATE, load_decoder
DEVICE = ("cuda" if torch.cuda.is_available() else
"mps" if torch.backends.mps.is_available() else "cpu")
DTYPE = torch.float32 if DEVICE == "cpu" else torch.float16
AUTOCAST = torch.bfloat16 if DEVICE == "cuda" else torch.float16
STEPS = 50
DEFAULT_CFG = 3.5
TIME_SHIFT = 2.0
TAG_GROUPS = json.loads((Path(__file__).parent / "tags.json").read_text(encoding="utf-8"))
METADATA_TAGS = TAG_GROUPS["metadata"]
OTHER_TAGS = TAG_GROUPS["other"]
SYNTHETIC_TAG_ID = METADATA_TAGS["synthetic"]
TAGS = METADATA_TAGS | OTHER_TAGS
del TAGS["synthetic"] # Source conditioning is added automatically, not shown in the UI.
TIMESTEPS = [1 - i / STEPS for i in range(STEPS + 1)]
TIMESTEPS = [TIME_SHIFT * t / (1 + (TIME_SHIFT - 1) * t) for t in TIMESTEPS]
model, latent_mean, latent_std = load_model(DEVICE)
decoder = load_decoder(device=DEVICE, dtype=DTYPE)
@torch.no_grad()
def generate(tag_names: list[str], cfg_scale: float, seed: int):
if not tag_names:
raise gr.Error("Pick at least one tag.")
seed = secrets.randbits(63) if seed == -1 else seed
torch.manual_seed(seed)
ids = [TAGS[name] for name in tag_names]
if all(name in METADATA_TAGS for name in tag_names):
if len(ids) == MAX_TAGS:
raise gr.Error(f"Choose at most {MAX_TAGS - 1} metadata tags.")
ids.append(SYNTHETIC_TAG_ID)
tags = torch.tensor([ids + [0] * (MAX_TAGS - len(ids))], device=DEVICE)
null = torch.zeros_like(tags)
x = torch.randn(1, CHANNELS, FRAMES, device=DEVICE)
for i, (t, t_next) in enumerate(pairwise(TIMESTEPS)):
# Guidance alternates between the shallow path-drop branch and a full
# unconditional pass as the weak model to steer away from.
weak = (model.shallow, model)[i % 2]
with torch.autocast(DEVICE, dtype=AUTOCAST, enabled=DEVICE != "cpu"):
cond = model(x, t, tags).float()
uncond = weak(x, t, null).float()
x = x + (t_next - t) * (uncond + cfg_scale * (cond - uncond))
audio = decoder.decode((x * latent_std + latent_mean).to(DTYPE))
return SAMPLE_RATE, audio[0].float().clamp(-1, 1).T.cpu().numpy()
with gr.Blocks(title="Audio DiT") as demo:
gr.Markdown("# Audio DiT\nChoose 1-8 tags.")
tag_box = gr.Dropdown(choices=list(TAGS), value=[], multiselect=True,
max_choices=MAX_TAGS, label="Tags", filterable=True)
cfg_scale = gr.Slider(2.0, 6.0, value=DEFAULT_CFG, step=0.5,
label="CFG scale")
seed = gr.Number(value=-1, precision=0, minimum=-1,
label="Seed (-1 = random)")
button = gr.Button("Generate", variant="primary")
audio_out = gr.Audio(label="Output", interactive=False)
button.click(generate, inputs=[tag_box, cfg_scale, seed], outputs=audio_out,
concurrency_limit=1)
if __name__ == "__main__":
demo.queue().launch()
|