import spaces import gradio as gr import torch from pathlib import Path from audiotools import AudioSignal from tria.model.tria import TRIA from tria.pipelines.tokenizer import Tokenizer from tria.features import rhythm_features from functools import partial from pyharp.core import ModelCard, build_endpoint from pyharp.media.audio import load_audio, save_audio from pyharp.labels import LabelList # Global Config DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") N_OUTPUTS = 3 # Model Zoo MODEL_ZOO = { "small_musdb_moises_2b": { "checkpoint": "pretrained/tria/small_musdb_moises_2b/80000/model.pt", "model_cfg": { "codebook_size": 1024, "n_codebooks": 9, "n_channels": 512, "n_feats": 2, "n_heads": 8, "n_layers": 12, "mult": 4, "p_dropout": 0.0, "bias": True, "max_len": 1000, "pos_enc": "rope", "qk_norm": True, "use_sdpa": True, "interp": "nearest", "share_emb": True, }, "tokenizer_cfg": {"name": "dac"}, "feature_cfg": { "sample_rate": 16_000, "n_bands": 2, "n_mels": 40, "window_length": 384, "hop_length": 192, "quantization_levels": 5, "slow_ma_ms": 200, "post_smooth_ms": 100, "legacy_normalize": False, "clamp_max": 50.0, "normalize_quantile": 0.98, }, "infer_cfg": { "top_p": 0.95, "top_k": None, "temp": 1.0, "mask_temp": 10.5, "iterations": [8, 8, 8, 8, 4, 4, 4, 4, 4], "guidance_scale": 2.0, "causal_bias": 1.0, }, "max_duration": 6.0, }, } # Loaded model cache LOADED = dict(name=None, model=None, tokenizer=None, feature_fn=None, infer_cfg=None, sample_rate=None, max_duration=None) # Model loading def load_model_by_name(name: str): """Load a TRIA model by name (cached).""" if LOADED["name"] == name and LOADED["model"] is not None: return LOADED["model"] cfg = MODEL_ZOO[name] model = TRIA(**cfg["model_cfg"]) sd = torch.load(cfg["checkpoint"], map_location="cpu") model.load_state_dict(sd, strict=True) model.to(DEVICE).eval() tokenizer = Tokenizer(**cfg["tokenizer_cfg"]).to(DEVICE) feat_fn = partial(rhythm_features, **cfg.get("feature_cfg", {})) LOADED.update( dict( name=name, model=model, tokenizer=tokenizer, feature_fn=feat_fn, infer_cfg=cfg["infer_cfg"], sample_rate=tokenizer.sample_rate, max_duration=cfg["max_duration"], ) ) return model # Inference logic @spaces.GPU @torch.inference_mode() def generate_audio(model_name, timbre_path, rhythm_path, cfg_scale, top_p, mask_temperature, seed): model = load_model_by_name(model_name) tokenizer = LOADED["tokenizer"] feat_fn = LOADED["feature_fn"] sample_rate = LOADED["sample_rate"] infer_cfg = LOADED["infer_cfg"] timbre_sig = load_audio(timbre_path).resample(sample_rate) rhythm_sig = load_audio(rhythm_path).resample(sample_rate) timbre_sig.ensure_max_of_audio() rhythm_sig.ensure_max_of_audio() prefix_dur = int(LOADED["max_duration"] / 3) timbre_tokens = tokenizer.encode(timbre_sig) rhythm_tokens = tokenizer.encode(rhythm_sig) tokens = torch.cat([timbre_tokens.tokens, rhythm_tokens.tokens], dim=-1) n_batch, n_codebooks, n_frames = tokens.shape prefix_frames = timbre_tokens.tokens.shape[-1] feats = feat_fn(rhythm_sig) feats = torch.nn.functional.interpolate(feats, n_frames - prefix_frames, mode=model.interp) full_feats = torch.zeros(n_batch, feats.shape[1], n_frames, device=DEVICE) full_feats[..., prefix_frames:] = feats prefix_mask = torch.arange(n_frames, device=DEVICE)[None, :].repeat(n_batch, 1) < prefix_frames buffer_mask = prefix_mask[:, None, :].repeat(1, n_codebooks, 1) feats_mask = ~prefix_mask outputs = [] for i in range(N_OUTPUTS): torch.manual_seed(seed + i) gen = model.inference( tokens.clone().to(DEVICE), full_feats.to(DEVICE), buffer_mask.clone().to(DEVICE), feats_mask.to(DEVICE), top_p=float(top_p), mask_temp=float(mask_temperature), iterations=infer_cfg["iterations"], guidance_scale=float(cfg_scale), )[..., prefix_frames:] rhythm_tokens.tokens = gen out_sig = tokenizer.decode(rhythm_tokens) out_sig.ensure_max_of_audio() output_path = f"tria_out_{i+1}.wav" save_audio(out_sig, output_path) path_i = output_path outputs.append(str(path_i)) return tuple(outputs) # PyHARP Metadata model_card = ModelCard( name="TRIA: The Rhythm In Anything", description=( "Transform your rhythmic ideas into full drum performances. TRIA takes two short audio prompts: \n " "Timbre Prompt: an example recording for the desired sound (e.g. drum sound) \n " "Rhythm Prompt: the sound gesture expressing the desired pattern (e.g. tapping or beatboxing) \n" "It generates 3 drum arrangements that match your groove and chosen timbre. " ), author="Patrick O'Reilly, Julia Barnett, Hugo Flores García, Annie Chu, Nathan Pruyne, Prem Seetharaman, Bryan Pardo", tags=["tria", "rhythm-generation", "pyharp"], ) # Gradio and PyHARP Endpoint with gr.Blocks(title="TRIA") as demo: timbre_in = gr.Audio(type="filepath", label="Timbre Prompt").harp_required(True) rhythm_in = gr.Audio(type="filepath", label="Rhythm Prompt").harp_required(True) model_names = list(MODEL_ZOO.keys()) model_dropdown = gr.Dropdown(choices=model_names, value=model_names[0], label="Model") with gr.Row(): cfg_scale = gr.Slider(0.0, 10.0, value=2.0, step=0.1, label="CFG Scale", info=( "Controls how strongly the model follows your prompts/conditions.\n" "Low values: Model is more creative/random, High values: Model adheres strictly to prompts, less variation.")) top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.01, label="Top P", info=( "Nucleus Sampling: Limits the pool of possible next tokens to the most likely ones.\n" "(Default behav.) only sample from tokens that make up the top 95% of probability")) mask_temperature = gr.Slider(0.0, 20.0, value=10.5, step=0.1, label="Mask Temperature", info=( "In masked models, controls the randomness of which tokens get unmasked first.\n" "Low temperature: More deterministic, predictable unmasking order, High temperature: More random unmasking, can add variety")) seed = gr.Slider(0, 1000, value=0, step=1, label="Random Seed", info=( "Same seed and same inputs: Produces identical outputs (reproducible).\n" "Set a specific seed to get consistent results, or use fifferent seeds to try variations")) out1 = gr.Audio(type="filepath", label="Generated #1") out2 = gr.Audio(type="filepath", label="Generated #2") out3 = gr.Audio(type="filepath", label="Generated #3") app = build_endpoint( model_card=model_card, input_components=[model_dropdown, timbre_in, rhythm_in, cfg_scale, top_p, mask_temperature, seed], output_components=[out1, out2, out3], process_fn=generate_audio, ) demo.queue().launch(share=True, show_error=True)