models/flowsep/diffusers/pipelines/spectrogram_diffusion/midi_utils.py

# Copyright 2022 The Music Spectrogram Diffusion Authors.
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import dataclasses
import math
import os
from typing import Any, Callable, List, Mapping, MutableMapping, Optional, Sequence, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F

from ...utils import is_note_seq_available
from .pipeline_spectrogram_diffusion import TARGET_FEATURE_LENGTH


if is_note_seq_available():
    import note_seq
else:
    raise ImportError("Please install note-seq via `pip install note-seq`")


INPUT_FEATURE_LENGTH = 2048

SAMPLE_RATE = 16000
HOP_SIZE = 320
FRAME_RATE = int(SAMPLE_RATE // HOP_SIZE)

DEFAULT_STEPS_PER_SECOND = 100
DEFAULT_MAX_SHIFT_SECONDS = 10
DEFAULT_NUM_VELOCITY_BINS = 1

SLAKH_CLASS_PROGRAMS = {
    "Acoustic Piano": 0,
    "Electric Piano": 4,
    "Chromatic Percussion": 8,
    "Organ": 16,
    "Acoustic Guitar": 24,
    "Clean Electric Guitar": 26,
    "Distorted Electric Guitar": 29,
    "Acoustic Bass": 32,
    "Electric Bass": 33,
    "Violin": 40,
    "Viola": 41,
    "Cello": 42,
    "Contrabass": 43,
    "Orchestral Harp": 46,
    "Timpani": 47,
    "String Ensemble": 48,
    "Synth Strings": 50,
    "Choir and Voice": 52,
    "Orchestral Hit": 55,
    "Trumpet": 56,
    "Trombone": 57,
    "Tuba": 58,
    "French Horn": 60,
    "Brass Section": 61,
    "Soprano/Alto Sax": 64,
    "Tenor Sax": 66,
    "Baritone Sax": 67,
    "Oboe": 68,
    "English Horn": 69,
    "Bassoon": 70,
    "Clarinet": 71,
    "Pipe": 73,
    "Synth Lead": 80,
    "Synth Pad": 88,
}


@dataclasses.dataclass
class NoteRepresentationConfig:
    """Configuration note representations."""

    onsets_only: bool
    include_ties: bool


@dataclasses.dataclass
class NoteEventData:
    pitch: int
    velocity: Optional[int] = None
    program: Optional[int] = None
    is_drum: Optional[bool] = None
    instrument: Optional[int] = None


@dataclasses.dataclass
class NoteEncodingState:
    """Encoding state for note transcription, keeping track of active pitches."""

    # velocity bin for active pitches and programs
    active_pitches: MutableMapping[Tuple[int, int], int] = dataclasses.field(default_factory=dict)


@dataclasses.dataclass
class EventRange:
    type: str
    min_value: int
    max_value: int


@dataclasses.dataclass
class Event:
    type: str
    value: int


class Tokenizer:
    def __init__(self, regular_ids: int):
        # The special tokens: 0=PAD, 1=EOS, and 2=UNK
        self._num_special_tokens = 3
        self._num_regular_tokens = regular_ids

    def encode(self, token_ids):
        encoded = []
        for token_id in token_ids:
            if not 0 <= token_id < self._num_regular_tokens:
                raise ValueError(
                    f"token_id {token_id} does not fall within valid range of [0, {self._num_regular_tokens})"
                )
            encoded.append(token_id + self._num_special_tokens)

        # Add EOS token
        encoded.append(1)

        # Pad to till INPUT_FEATURE_LENGTH
        encoded = encoded + [0] * (INPUT_FEATURE_LENGTH - len(encoded))

        return encoded


class Codec:
    """Encode and decode events.

    Useful for declaring what certain ranges of a vocabulary should be used for. This is intended to be used from
    Python before encoding or after decoding with GenericTokenVocabulary. This class is more lightweight and does not
    include things like EOS or UNK token handling.

    To ensure that 'shift' events are always the first block of the vocab and start at 0, that event type is required
    and specified separately.
    """

    def __init__(self, max_shift_steps: int, steps_per_second: float, event_ranges: List[EventRange]):
        """Define Codec.

        Args:
          max_shift_steps: Maximum number of shift steps that can be encoded.
          steps_per_second: Shift steps will be interpreted as having a duration of
              1 / steps_per_second.
          event_ranges: Other supported event types and their ranges.
        """
        self.steps_per_second = steps_per_second
        self._shift_range = EventRange(type="shift", min_value=0, max_value=max_shift_steps)
        self._event_ranges = [self._shift_range] + event_ranges
        # Ensure all event types have unique names.
        assert len(self._event_ranges) == len({er.type for er in self._event_ranges})

    @property
    def num_classes(self) -> int:
        return sum(er.max_value - er.min_value + 1 for er in self._event_ranges)

    # The next couple methods are simplified special case methods just for shift
    # events that are intended to be used from within autograph functions.

    def is_shift_event_index(self, index: int) -> bool:
        return (self._shift_range.min_value <= index) and (index <= self._shift_range.max_value)

    @property
    def max_shift_steps(self) -> int:
        return self._shift_range.max_value

    def encode_event(self, event: Event) -> int:
        """Encode an event to an index."""
        offset = 0
        for er in self._event_ranges:
            if event.type == er.type:
                if not er.min_value <= event.value <= er.max_value:
                    raise ValueError(
                        f"Event value {event.value} is not within valid range "
                        f"[{er.min_value}, {er.max_value}] for type {event.type}"
                    )
                return offset + event.value - er.min_value
            offset += er.max_value - er.min_value + 1

        raise ValueError(f"Unknown event type: {event.type}")

    def event_type_range(self, event_type: str) -> Tuple[int, int]:
        """Return [min_id, max_id] for an event type."""
        offset = 0
        for er in self._event_ranges:
            if event_type == er.type:
                return offset, offset + (er.max_value - er.min_value)
            offset += er.max_value - er.min_value + 1

        raise ValueError(f"Unknown event type: {event_type}")

    def decode_event_index(self, index: int) -> Event:
        """Decode an event index to an Event."""
        offset = 0
        for er in self._event_ranges:
            if offset <= index <= offset + er.max_value - er.min_value:
                return Event(type=er.type, value=er.min_value + index - offset)
            offset += er.max_value - er.min_value + 1

        raise ValueError(f"Unknown event index: {index}")


@dataclasses.dataclass
class ProgramGranularity:
    # both tokens_map_fn and program_map_fn should be idempotent
    tokens_map_fn: Callable[[Sequence[int], Codec], Sequence[int]]
    program_map_fn: Callable[[int], int]


def drop_programs(tokens, codec: Codec):
    """Drops program change events from a token sequence."""
    min_program_id, max_program_id = codec.event_type_range("program")
    return tokens[(tokens < min_program_id) | (tokens > max_program_id)]


def programs_to_midi_classes(tokens, codec):
    """Modifies program events to be the first program in the MIDI class."""
    min_program_id, max_program_id = codec.event_type_range("program")
    is_program = (tokens >= min_program_id) & (tokens <= max_program_id)
    return np.where(is_program, min_program_id + 8 * ((tokens - min_program_id) // 8), tokens)


PROGRAM_GRANULARITIES = {
    # "flat" granularity; drop program change tokens and set NoteSequence
    # programs to zero
    "flat": ProgramGranularity(tokens_map_fn=drop_programs, program_map_fn=lambda program: 0),
    # map each program to the first program in its MIDI class
    "midi_class": ProgramGranularity(
        tokens_map_fn=programs_to_midi_classes, program_map_fn=lambda program: 8 * (program // 8)
    ),
    # leave programs as is
    "full": ProgramGranularity(tokens_map_fn=lambda tokens, codec: tokens, program_map_fn=lambda program: program),
}


def frame(signal, frame_length, frame_step, pad_end=False, pad_value=0, axis=-1):
    """
    equivalent of tf.signal.frame
    """
    signal_length = signal.shape[axis]
    if pad_end:
        frames_overlap = frame_length - frame_step
        rest_samples = np.abs(signal_length - frames_overlap) % np.abs(frame_length - frames_overlap)
        pad_size = int(frame_length - rest_samples)

        if pad_size != 0:
            pad_axis = [0] * signal.ndim
            pad_axis[axis] = pad_size
            signal = F.pad(signal, pad_axis, "constant", pad_value)
    frames = signal.unfold(axis, frame_length, frame_step)
    return frames


def program_to_slakh_program(program):
    # this is done very hackily, probably should use a custom mapping
    for slakh_program in sorted(SLAKH_CLASS_PROGRAMS.values(), reverse=True):
        if program >= slakh_program:
            return slakh_program


def audio_to_frames(
    samples,
    hop_size: int,
    frame_rate: int,
) -> Tuple[Sequence[Sequence[int]], torch.Tensor]:
    """Convert audio samples to non-overlapping frames and frame times."""
    frame_size = hop_size
    samples = np.pad(samples, [0, frame_size - len(samples) % frame_size], mode="constant")

    # Split audio into frames.
    frames = frame(
        torch.Tensor(samples).unsqueeze(0),
        frame_length=frame_size,
        frame_step=frame_size,
        pad_end=False,  # TODO check why its off by 1 here when True
    )

    num_frames = len(samples) // frame_size

    times = np.arange(num_frames) / frame_rate
    return frames, times


def note_sequence_to_onsets_and_offsets_and_programs(
    ns: note_seq.NoteSequence,
) -> Tuple[Sequence[float], Sequence[NoteEventData]]:
    """Extract onset & offset times and pitches & programs from a NoteSequence.

    The onset & offset times will not necessarily be in sorted order.

    Args:
      ns: NoteSequence from which to extract onsets and offsets.

    Returns:
      times: A list of note onset and offset times. values: A list of NoteEventData objects where velocity is zero for
      note
          offsets.
    """
    # Sort by program and pitch and put offsets before onsets as a tiebreaker for
    # subsequent stable sort.
    notes = sorted(ns.notes, key=lambda note: (note.is_drum, note.program, note.pitch))
    times = [note.end_time for note in notes if not note.is_drum] + [note.start_time for note in notes]
    values = [
        NoteEventData(pitch=note.pitch, velocity=0, program=note.program, is_drum=False)
        for note in notes
        if not note.is_drum
    ] + [
        NoteEventData(pitch=note.pitch, velocity=note.velocity, program=note.program, is_drum=note.is_drum)
        for note in notes
    ]
    return times, values


def num_velocity_bins_from_codec(codec: Codec):
    """Get number of velocity bins from event codec."""
    lo, hi = codec.event_type_range("velocity")
    return hi - lo


# segment an array into segments of length n
def segment(a, n):
    return [a[i : i + n] for i in range(0, len(a), n)]


def velocity_to_bin(velocity, num_velocity_bins):
    if velocity == 0:
        return 0
    else:
        return math.ceil(num_velocity_bins * velocity / note_seq.MAX_MIDI_VELOCITY)


def note_event_data_to_events(
    state: Optional[NoteEncodingState],
    value: NoteEventData,
    codec: Codec,
) -> Sequence[Event]:
    """Convert note event data to a sequence of events."""
    if value.velocity is None:
        # onsets only, no program or velocity
        return [Event("pitch", value.pitch)]
    else:
        num_velocity_bins = num_velocity_bins_from_codec(codec)
        velocity_bin = velocity_to_bin(value.velocity, num_velocity_bins)
        if value.program is None:
            # onsets + offsets + velocities only, no programs
            if state is not None:
                state.active_pitches[(value.pitch, 0)] = velocity_bin
            return [Event("velocity", velocity_bin), Event("pitch", value.pitch)]
        else:
            if value.is_drum:
                # drum events use a separate vocabulary
                return [Event("velocity", velocity_bin), Event("drum", value.pitch)]
            else:
                # program + velocity + pitch
                if state is not None:
                    state.active_pitches[(value.pitch, value.program)] = velocity_bin
                return [
                    Event("program", value.program),
                    Event("velocity", velocity_bin),
                    Event("pitch", value.pitch),
                ]


def note_encoding_state_to_events(state: NoteEncodingState) -> Sequence[Event]:
    """Output program and pitch events for active notes plus a final tie event."""
    events = []
    for pitch, program in sorted(state.active_pitches.keys(), key=lambda k: k[::-1]):
        if state.active_pitches[(pitch, program)]:
            events += [Event("program", program), Event("pitch", pitch)]
    events.append(Event("tie", 0))
    return events


def encode_and_index_events(
    state, event_times, event_values, codec, frame_times, encode_event_fn, encoding_state_to_events_fn=None
):
    """Encode a sequence of timed events and index to audio frame times.

    Encodes time shifts as repeated single step shifts for later run length encoding.

    Optionally, also encodes a sequence of "state events", keeping track of the current encoding state at each audio
    frame. This can be used e.g. to prepend events representing the current state to a targets segment.

    Args:
      state: Initial event encoding state.
      event_times: Sequence of event times.
      event_values: Sequence of event values.
      encode_event_fn: Function that transforms event value into a sequence of one
          or more Event objects.
      codec: An Codec object that maps Event objects to indices.
      frame_times: Time for every audio frame.
      encoding_state_to_events_fn: Function that transforms encoding state into a
          sequence of one or more Event objects.

    Returns:
      events: Encoded events and shifts. event_start_indices: Corresponding start event index for every audio frame.
          Note: one event can correspond to multiple audio indices due to sampling rate differences. This makes
          splitting sequences tricky because the same event can appear at the end of one sequence and the beginning of
          another.
      event_end_indices: Corresponding end event index for every audio frame. Used
          to ensure when slicing that one chunk ends where the next begins. Should always be true that
          event_end_indices[i] = event_start_indices[i + 1].
      state_events: Encoded "state" events representing the encoding state before
          each event.
      state_event_indices: Corresponding state event index for every audio frame.
    """
    indices = np.argsort(event_times, kind="stable")
    event_steps = [round(event_times[i] * codec.steps_per_second) for i in indices]
    event_values = [event_values[i] for i in indices]

    events = []
    state_events = []
    event_start_indices = []
    state_event_indices = []

    cur_step = 0
    cur_event_idx = 0
    cur_state_event_idx = 0

    def fill_event_start_indices_to_cur_step():
        while (
            len(event_start_indices) < len(frame_times)
            and frame_times[len(event_start_indices)] < cur_step / codec.steps_per_second
        ):
            event_start_indices.append(cur_event_idx)
            state_event_indices.append(cur_state_event_idx)

    for event_step, event_value in zip(event_steps, event_values):
        while event_step > cur_step:
            events.append(codec.encode_event(Event(type="shift", value=1)))
            cur_step += 1
            fill_event_start_indices_to_cur_step()
            cur_event_idx = len(events)
            cur_state_event_idx = len(state_events)
        if encoding_state_to_events_fn:
            # Dump state to state events *before* processing the next event, because
            # we want to capture the state prior to the occurrence of the event.
            for e in encoding_state_to_events_fn(state):
                state_events.append(codec.encode_event(e))

        for e in encode_event_fn(state, event_value, codec):
            events.append(codec.encode_event(e))

    # After the last event, continue filling out the event_start_indices array.
    # The inequality is not strict because if our current step lines up exactly
    # with (the start of) an audio frame, we need to add an additional shift event
    # to "cover" that frame.
    while cur_step / codec.steps_per_second <= frame_times[-1]:
        events.append(codec.encode_event(Event(type="shift", value=1)))
        cur_step += 1
        fill_event_start_indices_to_cur_step()
        cur_event_idx = len(events)

    # Now fill in event_end_indices. We need this extra array to make sure that
    # when we slice events, each slice ends exactly where the subsequent slice
    # begins.
    event_end_indices = event_start_indices[1:] + [len(events)]

    events = np.array(events).astype(np.int32)
    state_events = np.array(state_events).astype(np.int32)
    event_start_indices = segment(np.array(event_start_indices).astype(np.int32), TARGET_FEATURE_LENGTH)
    event_end_indices = segment(np.array(event_end_indices).astype(np.int32), TARGET_FEATURE_LENGTH)
    state_event_indices = segment(np.array(state_event_indices).astype(np.int32), TARGET_FEATURE_LENGTH)

    outputs = []
    for start_indices, end_indices, event_indices in zip(event_start_indices, event_end_indices, state_event_indices):
        outputs.append(
            {
                "inputs": events,
                "event_start_indices": start_indices,
                "event_end_indices": end_indices,
                "state_events": state_events,
                "state_event_indices": event_indices,
            }
        )

    return outputs


def extract_sequence_with_indices(features, state_events_end_token=None, feature_key="inputs"):
    """Extract target sequence corresponding to audio token segment."""
    features = features.copy()
    start_idx = features["event_start_indices"][0]
    end_idx = features["event_end_indices"][-1]

    features[feature_key] = features[feature_key][start_idx:end_idx]

    if state_events_end_token is not None:
        # Extract the state events corresponding to the audio start token, and
        # prepend them to the targets array.
        state_event_start_idx = features["state_event_indices"][0]
        state_event_end_idx = state_event_start_idx + 1
        while features["state_events"][state_event_end_idx - 1] != state_events_end_token:
            state_event_end_idx += 1
        features[feature_key] = np.concatenate(
            [
                features["state_events"][state_event_start_idx:state_event_end_idx],
                features[feature_key],
            ],
            axis=0,
        )

    return features


def map_midi_programs(
    feature, codec: Codec, granularity_type: str = "full", feature_key: str = "inputs"
) -> Mapping[str, Any]:
    """Apply MIDI program map to token sequences."""
    granularity = PROGRAM_GRANULARITIES[granularity_type]

    feature[feature_key] = granularity.tokens_map_fn(feature[feature_key], codec)
    return feature


def run_length_encode_shifts_fn(
    features,
    codec: Codec,
    feature_key: str = "inputs",
    state_change_event_types: Sequence[str] = (),
) -> Callable[[Mapping[str, Any]], Mapping[str, Any]]:
    """Return a function that run-length encodes shifts for a given codec.

    Args:
      codec: The Codec to use for shift events.
      feature_key: The feature key for which to run-length encode shifts.
      state_change_event_types: A list of event types that represent state
          changes; tokens corresponding to these event types will be interpreted as state changes and redundant ones
          will be removed.

    Returns:
      A preprocessing function that run-length encodes single-step shifts.
    """
    state_change_event_ranges = [codec.event_type_range(event_type) for event_type in state_change_event_types]

    def run_length_encode_shifts(features: MutableMapping[str, Any]) -> Mapping[str, Any]:
        """Combine leading/interior shifts, trim trailing shifts.

        Args:
          features: Dict of features to process.

        Returns:
          A dict of features.
        """
        events = features[feature_key]

        shift_steps = 0
        total_shift_steps = 0
        output = np.array([], dtype=np.int32)

        current_state = np.zeros(len(state_change_event_ranges), dtype=np.int32)

        for event in events:
            if codec.is_shift_event_index(event):
                shift_steps += 1
                total_shift_steps += 1

            else:
                # If this event is a state change and has the same value as the current
                # state, we can skip it entirely.
                is_redundant = False
                for i, (min_index, max_index) in enumerate(state_change_event_ranges):
                    if (min_index <= event) and (event <= max_index):
                        if current_state[i] == event:
                            is_redundant = True
                        current_state[i] = event
                if is_redundant:
                    continue

                # Once we've reached a non-shift event, RLE all previous shift events
                # before outputting the non-shift event.
                if shift_steps > 0:
                    shift_steps = total_shift_steps
                    while shift_steps > 0:
                        output_steps = np.minimum(codec.max_shift_steps, shift_steps)
                        output = np.concatenate([output, [output_steps]], axis=0)
                        shift_steps -= output_steps
                output = np.concatenate([output, [event]], axis=0)

        features[feature_key] = output
        return features

    return run_length_encode_shifts(features)


def note_representation_processor_chain(features, codec: Codec, note_representation_config: NoteRepresentationConfig):
    tie_token = codec.encode_event(Event("tie", 0))
    state_events_end_token = tie_token if note_representation_config.include_ties else None

    features = extract_sequence_with_indices(
        features, state_events_end_token=state_events_end_token, feature_key="inputs"
    )

    features = map_midi_programs(features, codec)

    features = run_length_encode_shifts_fn(features, codec, state_change_event_types=["velocity", "program"])

    return features


class MidiProcessor:
    def __init__(self):
        self.codec = Codec(
            max_shift_steps=DEFAULT_MAX_SHIFT_SECONDS * DEFAULT_STEPS_PER_SECOND,
            steps_per_second=DEFAULT_STEPS_PER_SECOND,
            event_ranges=[
                EventRange("pitch", note_seq.MIN_MIDI_PITCH, note_seq.MAX_MIDI_PITCH),
                EventRange("velocity", 0, DEFAULT_NUM_VELOCITY_BINS),
                EventRange("tie", 0, 0),
                EventRange("program", note_seq.MIN_MIDI_PROGRAM, note_seq.MAX_MIDI_PROGRAM),
                EventRange("drum", note_seq.MIN_MIDI_PITCH, note_seq.MAX_MIDI_PITCH),
            ],
        )
        self.tokenizer = Tokenizer(self.codec.num_classes)
        self.note_representation_config = NoteRepresentationConfig(onsets_only=False, include_ties=True)

    def __call__(self, midi: Union[bytes, os.PathLike, str]):
        if not isinstance(midi, bytes):
            with open(midi, "rb") as f:
                midi = f.read()

        ns = note_seq.midi_to_note_sequence(midi)
        ns_sus = note_seq.apply_sustain_control_changes(ns)

        for note in ns_sus.notes:
            if not note.is_drum:
                note.program = program_to_slakh_program(note.program)

        samples = np.zeros(int(ns_sus.total_time * SAMPLE_RATE))

        _, frame_times = audio_to_frames(samples, HOP_SIZE, FRAME_RATE)
        times, values = note_sequence_to_onsets_and_offsets_and_programs(ns_sus)

        events = encode_and_index_events(
            state=NoteEncodingState(),
            event_times=times,
            event_values=values,
            frame_times=frame_times,
            codec=self.codec,
            encode_event_fn=note_event_data_to_events,
            encoding_state_to_events_fn=note_encoding_state_to_events,
        )

        events = [
            note_representation_processor_chain(event, self.codec, self.note_representation_config) for event in events
        ]
        input_tokens = [self.tokenizer.encode(event["inputs"]) for event in events]

        return input_tokens