| 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)) |
| |
| |
| 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() |
|
|
