diff --git a/README.md b/README.md index e2a6318..89936b8 100644 --- a/README.md +++ b/README.md @@ -8,12 +8,20 @@ This repository contains training, inference, and evaluation code for the paper ## Installation -Installation requires Python 3.11+. To install the package and all dependencies with pip: +Installation requires Python 3.11+. + +To install the package with `pip`: ```bash git clone https://github.com/EleutherAI/aria -cd aria -pip install -e ".[all]" +cd aria && pip install -e ".[all]" +``` + +To install the package with `uv`: + +```bash +git clone https://github.com/EleutherAI/aria +cd aria && uv sync --extra all ``` ## Quickstart @@ -79,7 +87,7 @@ Our embedding model was trained to capture composition-level and performance-lev ## Real-time demo -In `demo/` we provide an MLX (Apple Silicon) implementation of the real-time interactive piano-continuation demo showcased in our release blog post. In order to use the demo, you must download the demo-specific model checkpoint which enhances the model to additionally control the sustain pedal ([direct-download](https://huggingface.co/loubb/aria-medium-base/resolve/main/model-demo.safetensors?download=true)). +In `demo/` we provide an MLX (Apple Silicon) implementation of the real-time interactive piano-continuation demo showcased in our release blog post. In order to use the demo, you must download the demo-specific model checkpoint which enhances the model to additionally control the sustain pedal ([direct-download](https://huggingface.co/loubb/aria-medium-base/resolve/main/model-demo.safetensors?download=true)). If using `uv`, install the demo dependencies with `uv sync --extra demo`. For our demonstration, we used an acoustic Yamaha Disklavier piano with simultaneous MIDI input and output ports connected via a standard MIDI interface. We disabled the built-in Disklavier playback mode, instead manually calibrating key-velocity latency to enhance responsiveness. You may recreate this in your own environment with our acoustic calibration settings, using the following script: @@ -93,8 +101,8 @@ python ./demo/demo_mlx.py \ --hardware ./demo/hardware/c4dm-disklavier.json \ --midi_control_signal 67 \ --midi_reset_control_signal 66 \ - --temp 0.85 \ - --min_p 0.05 + --temp 0.9 \ + --min_p 0.035 ``` A MIDI input device is not strictly required to play around with the demo: By using the `--midi_path` and `--midi_through` arguments you can mock real-time input by playing from a MIDI file. All that is required are MIDI drivers (e.g., CoreMIDI) and a virtual software instrument (e.g., Fluidsynth, Pianoteq) to render the output. In this mode, you can initiate the model takeover by pressing the enter key. @@ -107,8 +115,8 @@ python ./demo/demo_mlx.py \ --midi_path ${MIDI_PATH} \ --midi_through \ --midi_out \ - --temp 0.85 \ - --min_p 0.05 + --temp 0.9 \ + --min_p 0.035 ``` ❗**NOTE**: Responsiveness of the real-time demo is dependent on your system configuration, specifically GPU memory bandwidth. diff --git a/aria/inference/model_mlx.py b/aria/inference/model_mlx.py index 51ad32e..0a743dc 100644 --- a/aria/inference/model_mlx.py +++ b/aria/inference/model_mlx.py @@ -178,13 +178,13 @@ def fill_condition_kv(self, emb: mx.array): assert self.model_config.emb_size is not None input_pos = mx.array([0], dtype=mx.int32) - mask = self.causal_mask[None, None, input_pos] + mask = self.causal_mask[None, None, input_pos, :1] offset = 0 x = mx.expand_dims(emb, axis=1) for layer in self.encode_layers: - x = layer(x, input_pos, offset, mask) + x = layer(x, input_pos, 0, offset, mask) def __call__( self, diff --git a/aria/training/train.py b/aria/training/train.py index 6521272..5882c67 100644 --- a/aria/training/train.py +++ b/aria/training/train.py @@ -94,9 +94,9 @@ def get_tokenizer_name( train_config = TrainingDataset.get_config_from_path(train_data_paths[0]) val_config = TrainingDataset.get_config_from_path(val_data_path) - assert ( - train_config["tokenizer_name"] == val_config["tokenizer_name"] - ), "Dataset tokenizers don't match" + assert train_config["tokenizer_name"] == val_config["tokenizer_name"], ( + "Dataset tokenizers don't match" + ) return train_config["tokenizer_name"] @@ -127,9 +127,9 @@ def setup_project_dir(project_dir: str | None): elif project_dir: # Run checks on project directory if os.path.isdir(project_dir): - assert ( - len(os.listdir(project_dir)) == 0 - ), "Provided project directory is not empty" + assert len(os.listdir(project_dir)) == 0, ( + "Provided project directory is not empty" + ) project_dir_abs = os.path.abspath(project_dir) elif os.path.isfile(project_dir): raise FileExistsError( @@ -226,9 +226,9 @@ def get_dataloaders( if init_epoch: train_dataset.init_epoch(idx=init_epoch) - assert ( - len(val_dataset.epoch_files_by_dir[0]) == 1 - ), "val-data directory should only contain one epoch" + assert len(val_dataset.epoch_files_by_dir[0]) == 1, ( + "val-data directory should only contain one epoch" + ) if apply_aug: train_dataset.set_transform(tokenizer.export_data_aug()) @@ -427,9 +427,9 @@ def val_loop(dataloader, _epoch: int): return avg_val_loss if steps_per_checkpoint: - assert ( - steps_per_checkpoint > 1 - ), "Invalid checkpoint mode value (too small)" + assert steps_per_checkpoint > 1, ( + "Invalid checkpoint mode value (too small)" + ) TRAILING_LOSS_STEPS = 200 PAD_ID = train_dataloader.dataset.tokenizer.pad_id @@ -516,9 +516,9 @@ def resume_train( assert torch.cuda.is_available() is True, "CUDA not available" assert os.path.isdir(checkpoint_dir), f"No dir at {checkpoint_dir}" for train_data_path in train_data_paths: - assert os.path.isdir( - train_data_path - ), f"No dir found at {train_data_path}" + assert os.path.isdir(train_data_path), ( + f"No dir found at {train_data_path}" + ) assert os.path.isdir(val_data_path), f"No dir found at {val_data_path}" tokenizer_name = get_tokenizer_name(train_data_paths, val_data_path) @@ -649,9 +649,9 @@ def train( assert batch_size > 0, "Invalid batch size" assert torch.cuda.is_available() is True, "CUDA not available" for train_data_path in train_data_paths: - assert os.path.isdir( - train_data_path - ), f"No dir found at {train_data_path}" + assert os.path.isdir(train_data_path), ( + f"No dir found at {train_data_path}" + ) assert os.path.isdir(val_data_path), f"No dir found at {val_data_path}" tokenizer_name = get_tokenizer_name(train_data_paths, val_data_path) diff --git a/demo/config.json b/demo/config.json deleted file mode 100644 index c648810..0000000 --- a/demo/config.json +++ /dev/null @@ -1,51 +0,0 @@ -{ - "tokenizer": { - "abs": { - "ignore_instruments": { - "piano": false, - "chromatic": true, - "organ": false, - "guitar": false, - "bass": false, - "strings": false, - "ensemble": false, - "brass": false, - "reed": false, - "pipe": false, - "synth_lead": false, - "synth_pad": true, - "synth_effect": true, - "ethnic": true, - "percussive": true, - "sfx": true - }, - "instrument_programs": { - "piano": 0, - "chromatic": 13, - "organ": 16, - "guitar": 24, - "bass": 32, - "strings": 40, - "ensemble": 48, - "brass": 56, - "reed": 64, - "pipe": 73, - "synth_lead": 80, - "synth_pad": 88, - "synth_effect": 96, - "ethnic": 104, - "percussive": 112, - "sfx": 120 - }, - "drum_velocity": 60, - "velocity_quantization_step": 10, - "abs_time_step_ms": 5000, - "max_dur_ms": 5000, - "time_step_ms": 10, - "include_pedal": true, - "composer_names": ["bach", "beethoven", "mozart", "chopin", "rachmaninoff", "liszt", "debussy", "schubert", "brahms", "ravel", "satie", "scarlatti"], - "form_names": ["sonata", "prelude", "nocturne", "étude", "waltz", "mazurka", "impromptu", "fugue"], - "genre_names": ["jazz", "classical"] - } - } -} diff --git a/demo/demo_mlx.py b/demo/demo_mlx.py index 37f4f2f..e20e50c 100644 --- a/demo/demo_mlx.py +++ b/demo/demo_mlx.py @@ -36,8 +36,8 @@ BEAM_WIDTH: int = 3 TIME_TOK_WEIGHTING: int = -5 -FIRST_ONSET_BUFFER_MS: int = -200 -MAX_STREAM_DELAY_MS: int = 100 +FIRST_ONSET_BUFFER_MS: int = -150 +MAX_STREAM_DELAY_MS: int = 50 MIN_NOTE_DELTA_MS: int = 0 MIN_PEDAL_DELTA_MS: int = 0 @@ -47,7 +47,11 @@ VELOCITY_OUTPUT_LATENCY_MS: dict[int, int] = {v: 0 for v in range(0, 127, 10)} -config_path = pathlib.Path(__file__).parent.resolve().joinpath("config.json") +TOKENIZER_CONFIG_PATH = ( + pathlib.Path(__file__) + .parent.resolve() + .joinpath("demo-tokenizer-config.json") +) file_handler = logging.FileHandler("./demo.log", mode="w") file_handler.setLevel(logging.DEBUG) @@ -210,7 +214,9 @@ def prefill( logits = model( idxs=idxs, input_pos=input_pos + EMBEDDING_OFFSET, - max_kv_pos=math.ceil(input_pos[-1].item() / KV_CHUNK_SIZE) + max_kv_pos=math.ceil( + (input_pos[-1].item() + EMBEDDING_OFFSET) / KV_CHUNK_SIZE + ) * KV_CHUNK_SIZE, offset=input_pos[0] + EMBEDDING_OFFSET, ) @@ -228,7 +234,9 @@ def decode_one( logits = model( idxs=idxs, input_pos=input_pos + EMBEDDING_OFFSET, - max_kv_pos=math.ceil(input_pos[-1].item() / KV_CHUNK_SIZE) + max_kv_pos=math.ceil( + (input_pos[-1].item() + EMBEDDING_OFFSET) / KV_CHUNK_SIZE + ) * KV_CHUNK_SIZE, offset=input_pos[0] + EMBEDDING_OFFSET, )[:, -1] @@ -361,7 +369,7 @@ def warmup_model(model: TransformerLM): def load_model(checkpoint_path: str): logger = get_logger() - tokenizer = AbsTokenizer(config_path=config_path) + tokenizer = AbsTokenizer(config_path=TOKENIZER_CONFIG_PATH) model_config = ModelConfig(**load_model_config("medium-emb")) model_config.set_vocab_size(tokenizer.vocab_size) @@ -377,7 +385,11 @@ def load_model(checkpoint_path: str): assert ( tokenizer.vocab_size == weights["model.tok_embeddings.weight"].shape[0] - ), "Embedding shape mismatch. Ensure that you are loading the demo-specific checkpoint." + ), ( + "Embedding shape mismatch. Ensure that you are loading the " + f"demo-specific checkpoint. tokenizer={tokenizer.vocab_size}, " + f"weights={weights['model.tok_embeddings.weight'].shape[0]}" + ) model.load_weights(list(weights.items()), strict=False) model.eval() @@ -460,7 +472,7 @@ def recalc_dur_tokens_chunked( window_pos = mx.arange(idx - 1, end_idx - 1, dtype=mx.int32) logger.info( - f"Recalculating chunked durations for positions: {idx-1} - {end_idx-2}" + f"Recalculating chunked durations for positions: {idx - 1} - {end_idx - 2}" ) logits = prefill(model, idxs=window_ids, input_pos=window_pos) @@ -531,7 +543,7 @@ def decode_first_tokens( input_pos=mx.array([idx - 1], dtype=mx.int32), ) - logger.info(f"Inserted time_tok at position {idx-1}") + logger.info(f"Inserted time_tok at position {idx - 1}") num_time_toks_to_add -= 1 enc_seq[:, idx - 1] = time_tok_id idx += 1 @@ -588,7 +600,7 @@ def decode_first_tokens( input_pos=mx.array([idx - 1], dtype=mx.int32), ) logger.debug( - f"Sampled logits for positions {idx} by inserting {tok} at position {idx-1}" + f"Sampled logits for positions {idx} by inserting {tok} at position {idx - 1}" ) next_log_probs = nn.log_softmax(next_logits, axis=-1) @@ -637,7 +649,7 @@ def decode_first_tokens( ) logger.info( - f"Updated KV-Cache by re-inserting {best_tok_1} at position {idx-1}" + f"Updated KV-Cache by re-inserting {best_tok_1} at position {idx - 1}" ) logger.debug(f"Internal KV-state: {tokenizer.decode(model.get_kv_ctx())}") @@ -682,7 +694,7 @@ def decode_tokens( ) logger.debug( - f"Sampled logits for positions {idx} by inserting {prev_tok} at position {idx-1}" + f"Sampled logits for positions {idx} by inserting {prev_tok} at position {idx - 1}" ) logits[:, tokenizer.tok_to_id[tokenizer.ped_off_tok]] += 3 # Manual adj @@ -703,7 +715,7 @@ def decode_tokens( enc_seq[:, idx] = next_token_ids next_token = tokenizer.id_to_tok[next_token_ids[0].item()] logger.debug( - f"({(time.time() - decode_one_start_time_s)*1000:.2f}ms) {idx}: {next_token}" + f"({(time.time() - decode_one_start_time_s) * 1000:.2f}ms) {idx}: {next_token}" ) if next_token in {tokenizer.ped_on_tok, tokenizer.ped_off_tok}: @@ -756,7 +768,7 @@ def generate_tokens( logger.debug(f"Priming sequence {priming_seq}") logger.info(f"Priming sequence length: {priming_seq_len}") - logger.info(f"Prefilling up to (and including) position: {start_idx-1}") + logger.info(f"Prefilling up to (and including) position: {start_idx - 1}") prefill_start_s = time.time() chunked_prefill( @@ -770,7 +782,9 @@ def generate_tokens( logger.info( f"Prefill took {(time.time() - prefill_start_s) * 1000:.2f} milliseconds" ) - logger.info(f"Starting duration recalculation from position: {start_idx-1}") + logger.info( + f"Starting duration recalculation from position: {start_idx - 1}" + ) recalculate_dur_start_s = time.time() enc_seq, priming_seq, next_token_logits = recalc_dur_tokens_chunked( @@ -1160,7 +1174,7 @@ def stream_midi( and msg["pitch"] != "pedal" and active_pitch_uuid.get(msg["pitch"]) == msg["uuid"] ] - remaining_off_msgs.sort(key=lambda m: (m["epoch_time_ms"])) + remaining_off_msgs.sort(key=lambda m: m["epoch_time_ms"]) for msg in remaining_off_msgs: mido_msg = _create_mido_message( @@ -1213,7 +1227,12 @@ def stream_msgs( logger.info("Removing final pedal_off from tokenized sequence") priming_seq = priming_seq[:-2] + priming_seq.append(tokenizer.delimiter_tok) + if is_ending is True: + raise NotImplementedError( + "I've removed this functionality so this should never trigger" + ) priming_seq.append(tokenizer.dim_tok) generated_tokens_queue = queue.Queue() @@ -1421,7 +1440,7 @@ def chunked_prefill( break logger.info( - f"KV stored up to idx={max(0, len(prev_context)- 1)} (curr_context_len={len(curr_context)})" + f"KV stored up to idx={max(0, len(prev_context) - 1)} (curr_context_len={len(curr_context)})" ) return prev_context @@ -1433,7 +1452,7 @@ def continuous_prefill( received_messages_queue: queue.Queue, prev_context: list[int], ): - tokenizer = AbsTokenizer(config_path=config_path) + tokenizer = AbsTokenizer(config_path=TOKENIZER_CONFIG_PATH) logger = get_logger("PREFILL") msg_cnt = 0 seen_sentinel = False @@ -1486,6 +1505,7 @@ def capture_and_update_kv( midi_performance_queue: queue.Queue, midi_capture_channel: int, first_msg_epoch_time_ms: int | None = None, + idle_timeout_ms: int | None = None, ): received_messages_queue = queue.Queue() results_queue = queue.Queue() @@ -1500,6 +1520,7 @@ def capture_and_update_kv( "first_msg_epoch_time_ms": first_msg_epoch_time_ms, "results_queue": results_queue, "wait_for_close": wait_for_close, + "idle_timeout_ms": idle_timeout_ms, }, ) capture_midi_thread.start() @@ -1525,6 +1546,7 @@ def capture_midi_input( results_queue: queue.Queue, first_msg_epoch_time_ms: int | None = None, wait_for_close: bool = False, + idle_timeout_ms: int | None = None, ): logger = get_logger("CAPTURE") first_on_msg_epoch_ms = None @@ -1532,6 +1554,7 @@ def capture_midi_input( pedal_down = False pitches_held_down = set() pitches_sustained_by_pedal = set() + waiting_since_ms = None while not midi_performance_queue.empty(): try: @@ -1542,6 +1565,10 @@ def capture_midi_input( logger.info("Listening for input") logger.info("Commencing generation upon keypress or control signal") + if idle_timeout_ms is not None: + waiting_since_ms = get_epoch_time_ms() + logger.info(f"Idle timeout enabled: {idle_timeout_ms}ms") + while True: epoch_time_ms = get_epoch_time_ms() active_notes = pitches_held_down.union(pitches_sustained_by_pedal) @@ -1554,6 +1581,18 @@ def capture_midi_input( try: msg = midi_performance_queue.get(block=True, timeout=0.01) except queue.Empty: + if ( + idle_timeout_ms is not None + and waiting_since_ms is not None + and first_on_msg_epoch_ms is None + and (get_epoch_time_ms() - waiting_since_ms) > idle_timeout_ms + ): + logger.info( + f"Idle timeout reached ({idle_timeout_ms}ms), resetting" + ) + reset_sentinel.set() + control_sentinel.set() + break continue if msg.is_meta or msg.type == "program_change": @@ -1805,7 +1844,7 @@ def run( back_and_forth: bool, ): logger = get_logger() - tokenizer = AbsTokenizer(config_path=config_path) + tokenizer = AbsTokenizer(config_path=TOKENIZER_CONFIG_PATH) control_sentinel = threading.Event() currently_generating_sentinel = threading.Event() @@ -1911,6 +1950,7 @@ def run( midi_performance_queue=midi_performance_queue, midi_capture_channel=curr_midi_channel, first_msg_epoch_time_ms=first_on_msg_epoch_ms, + idle_timeout_ms=3000, ) keypress_thread.join() @@ -1969,12 +2009,18 @@ def main(args): logger = get_logger() model = load_model(checkpoint_path=args.checkpoint) model = warmup_model(model=model) + global EMBEDDING_OFFSET if args.embedding_checkpoint and args.embedding_midi_path: insert_embedding( model=model, embedding_model_checkpoint_path=args.embedding_checkpoint, embedding_midi_path=args.embedding_midi_path, ) + else: + model.fill_condition_kv( + mx.zeros((1, model.model_config.emb_size), dtype=DTYPE) + ) + EMBEDDING_OFFSET = 1 assert (args.midi_path and os.path.isfile(args.midi_path)) or args.midi_in @@ -2022,7 +2068,7 @@ def playback(midi_path: str, midi_out: str, save_path: str | None = None): close_notes(midi_out) starting_epoch_time_ms = get_epoch_time_ms() - tokenizer = AbsTokenizer(config_path=config_path) + tokenizer = AbsTokenizer(config_path=TOKENIZER_CONFIG_PATH) tokens_queue = queue.Queue() midi_messages_queue = queue.Queue() stream_midi_results_queue = queue.Queue() diff --git a/pyproject.toml b/pyproject.toml index 199cfb8..7b6e56d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,18 +12,18 @@ requires-python = ">=3.11" dependencies = [ "ariautils @ git+https://github.com/EleutherAI/aria-utils.git", "torch>=2.3", - "mlx", + "mlx<=0.26", "safetensors", "jsonlines", "tqdm", ] [project.optional-dependencies] -dev = ["black"] +dev = ["ruff", "ty"] train = ["accelerate"] eval = ["transformers", "torchaudio", "mido"] -demo = ["python-rtmidi"] -all = ["black", "accelerate", "transformers", "torchaudio", "mido", "python-rtmidi"] +demo = ["python-rtmidi", "numpy"] +all = ["black", "accelerate", "transformers", "torchaudio", "mido", "python-rtmidi", "numpy"] [tool.black] line-length = 80