File size: 4,125 Bytes
a89ad01 | 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 | from __future__ import annotations
import json
import secrets
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
MODEL_DTYPE = (torch.bfloat16 if DEVICE == "cuda" else
torch.float16 if DEVICE == "mps" else torch.float32)
AUTOCAST = MODEL_DTYPE
STEPS = 50
DEFAULT_CFG = 4.0
TIME_SHIFT = 3.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"]
AESTHETIC_TAG_ID = OTHER_TAGS["aesthetic"]
TAGS = METADATA_TAGS | OTHER_TAGS
del TAGS["synthetic"]
del TAGS["aesthetic"]
VISIBLE_MAX_TAGS = 6
SAMPLE_MP3S = {
path.stem: path
for path in sorted((Path(__file__).parent / "samples").glob("*.mp3"))
}
model = load_model(DEVICE, MODEL_DTYPE)
decoder = load_decoder(device=DEVICE, dtype=DTYPE)
def sampling_timesteps(device, dtype):
timesteps = torch.linspace(1.0, 0.0, STEPS + 1, device=device, dtype=dtype)[:-1]
return TIME_SHIFT * timesteps / (1 + (TIME_SHIFT - 1) * timesteps)
@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 len(ids) > VISIBLE_MAX_TAGS:
raise gr.Error(f"Choose at most {VISIBLE_MAX_TAGS} tags.")
if all(name in METADATA_TAGS for name in tag_names):
ids.append(SYNTHETIC_TAG_ID)
else:
ids.append(AESTHETIC_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)
timesteps = sampling_timesteps(x.device, x.dtype)
for i in range(STEPS):
t = timesteps[i]
t_next = (timesteps[i + 1] if i + 1 < STEPS else
torch.tensor(0.0, device=x.device, dtype=x.dtype))
# The model predicts clean latents. Guidance alternates between the
# shallow path-drop branch and a full unconditional pass.
weak = (model.shallow, model)[i % 2]
with torch.autocast(DEVICE, dtype=AUTOCAST, enabled=DEVICE != "cpu"):
cond_x0 = model(x, t, tags).float()
uncond_x0 = weak(x, t, null).float()
cond = (x - cond_x0) / t
uncond = (x - uncond_x0) / t
x = x - (t - t_next) * (uncond + cfg_scale * (cond - uncond))
audio = decoder.decode(x.to(DTYPE))
return SAMPLE_RATE, audio[0].float().clamp(-1, 1).T.cpu().numpy()
def load_sample(tag_names: list[str]):
return SAMPLE_MP3S[tag_names[0]]
with gr.Blocks(title="Audio DiT") as demo:
gr.Markdown(
f"# Audio DiT\nChoose 1-{VISIBLE_MAX_TAGS} tags to generate a song, "
"or click a sample song to load it."
)
tag_box = gr.Dropdown(choices=list(TAGS), value=[], multiselect=True,
max_choices=VISIBLE_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", format="mp3", interactive=False)
gr.Examples(
examples=[[[tag]] for tag in SAMPLE_MP3S],
inputs=tag_box,
outputs=audio_out,
fn=load_sample,
cache_examples=False,
run_on_click=True,
example_labels=list(SAMPLE_MP3S),
label="Sample songs",
)
button.click(generate, inputs=[tag_box, cfg_scale, seed], outputs=audio_out,
concurrency_limit=1)
if __name__ == "__main__":
demo.queue().launch()
|