chapter_09_example_03.py

"""
This example shows a basic Drums RNN generation with a
looping synthesizer playback, using a MIDI hub to send the sequence
to an external device.

VERSION: Magenta 1.1.7
"""
import argparse
import os
import time
from decimal import Decimal

import magenta.music as mm
import mido
import tensorflow as tf
from magenta.common import concurrency
from magenta.interfaces.midi.midi_hub import MidiHub
from magenta.interfaces.midi.midi_interaction import adjust_sequence_times
from magenta.models.drums_rnn import drums_rnn_sequence_generator
from magenta.music import constants
from magenta.protobuf import generator_pb2
from magenta.protobuf import music_pb2
from visual_midi import Plotter

parser = argparse.ArgumentParser()
parser.add_argument("--midi_port", type=str, default="FLUID Synth")
args = parser.parse_args()


def generate(unused_argv):
  # Downloads the bundle from the magenta website
  mm.notebook_utils.download_bundle("drum_kit_rnn.mag", "bundles")
  bundle = mm.sequence_generator_bundle.read_bundle_file(
    os.path.join("bundles", "drum_kit_rnn.mag"))

  # Initialize the generator "drum_kit"
  generator_map = drums_rnn_sequence_generator.get_generator_map()
  generator = generator_map["drum_kit"](checkpoint=None, bundle=bundle)
  generator.initialize()

  # Define constants
  qpm = 120
  num_bars = 3
  seconds_per_step = 60.0 / qpm / generator.steps_per_quarter
  num_steps_per_bar = constants.DEFAULT_STEPS_PER_BAR
  seconds_per_bar = num_steps_per_bar * seconds_per_step

  # Use a priming sequence
  primer_sequence = mm.midi_io.midi_file_to_note_sequence(
    os.path.join("primers", "Jazz_Drum_Basic_1_bar.mid"))
  primer_start_time = 0
  primer_end_time = primer_start_time + seconds_per_bar

  # Calculates the generation start and end time
  generation_start_time = primer_end_time
  generation_end_time = generation_start_time + (seconds_per_bar * num_bars)
  generator_options = generator_pb2.GeneratorOptions()
  generator_options.args['temperature'].float_value = 1.1
  generator_options.generate_sections.add(
    start_time=generation_start_time,
    end_time=generation_end_time)

  # Generates on primer sequence
  sequence = generator.generate(primer_sequence, generator_options)

  # Outputs the plot
  os.makedirs("output", exist_ok=True)
  plot_file = os.path.join("output", "out.html")
  pretty_midi = mm.midi_io.note_sequence_to_pretty_midi(sequence)
  plotter = Plotter()
  plotter.show(pretty_midi, plot_file)
  print(f"Generated plot file: {os.path.abspath(plot_file)}")

  # We find the proper input port for the software synth
  # (which is the output port for Magenta)
  output_ports = [name for name in mido.get_output_names()
                  if args.midi_port in name]
  if not output_ports:
    raise Exception(f"Cannot find proper output ports in: "
                    f"{mido.get_output_names()}")
  print(f"Playing generated MIDI in output port names: {output_ports}")

  # Start a new MIDI hub on that port (output only)
  midi_hub = MidiHub(input_midi_ports=[],
                     output_midi_ports=output_ports,
                     texture_type=None)

  # Start on a empty sequence, allowing the update of the
  # sequence for later.
  empty_sequence = music_pb2.NoteSequence()
  player = midi_hub.start_playback(empty_sequence,
                                   allow_updates=True)
  player._channel = 9

  # We want a period in seconds of 4 bars (which is the loop
  # length). Using 240 / qpm, we have a period of 1 bar, or
  # 2 seconds at 120 qpm. We multiply that by 4 bars.
  # (using the Decimal class for more accuracy)
  period = Decimal(240) / qpm
  period = period * (num_bars + 1)
  sleeper = concurrency.Sleeper()
  while True:
    try:
      # We get the next tick time by using the period
      # to find the absolute tick number (since epoch),
      # and multiplying by the period length. This is
      # used to sleep until that time.
      # We also find the current tick time for the player
      # to update.
      # (using the Decimal class for more accuracy)
      now = Decimal(time.time())
      tick_number = int(now // period)
      tick_number_next = tick_number + 1
      tick_time = tick_number * period
      tick_time_next = tick_number_next * period

      print(f"now {now} tick_time {tick_time} tick_time_next {tick_time_next}")

      # Update the player time to the current tick time
      sequence_adjusted = music_pb2.NoteSequence()
      sequence_adjusted.CopyFrom(sequence)
      sequence_adjusted = adjust_sequence_times(sequence_adjusted,
                                                float(tick_time))
      player.update_sequence(sequence_adjusted,
                             start_time=float(tick_time))

      # Sleep until the next tick time
      sleeper.sleep_until(float(tick_time_next))
    except KeyboardInterrupt:
      print(f"Stopping")
      return 0


if __name__ == "__main__":
  tf.app.run(generate)