flask_server.py

from matplotlib.figure import Figure
from sklearn.preprocessing import LabelEncoder
from sample import (sample_model, make_conditioning_tensors,
                    ConditioningMap, make_conditioning_map)

import soundfile
from typing import Union, Tuple, Mapping, Optional, Dict, List
import tempfile
import os
import functools
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import librosa.display
from zipfile import ZipFile
from distutils.util import strtobool
import smart_open

import torch
import torchaudio

import flask
from flask import request, abort
from flask_cors import CORS
from werkzeug.utils import secure_filename

import io
import base64

import logging
from logging import handlers as logging_handlers

import gc

upload_directory = './uploads'
os.makedirs(upload_directory, exist_ok=True)

torchaudio.set_audio_backend('sox_io')

# use matplotlib without an X server
# on desktop, this prevents matplotlib windows from popping around
mpl.use('Agg')

ALLOWED_EXTENSIONS = {'wav'}

# INITIALIZATION
wav_response_headers = {"Content-Type": "audio/wav"
                        }

# get colormap
ncolors = 256
color_array = plt.get_cmap('Greys')(range(ncolors))
# change alpha values
color_array[:, -1] = np.linspace(0.0, 1.0, ncolors)
# create a colormap object
map_object = LinearSegmentedColormap.from_list(name='Greys_alpha',
                                               colors=color_array)
# register this new colormap with matplotlib
plt.register_cmap(cmap=map_object)


def make_label_encoders(
        label_encoder_classes_per_modality: Mapping[str, List[any]]):
    label_encoders: Mapping[str, LabelEncoder] = {}
    for modality, classes in label_encoder_classes_per_modality.items():
        le = LabelEncoder().fit(classes)
        label_encoders[modality] = le
    return label_encoders


@torch.inference_mode()
def create_app():
    app = flask.Flask(__name__, static_folder='uploads')
    CORS(app)
    app.config['UPLOAD_EXTENSIONS'] = (torchaudio.utils.sox_utils
                                       .list_read_formats()) + ['mp3']

    # upload_directory = expand_path(pathlib.Path(tempfile.gettempdir())
    #                                / 'vqvae_uploads/')
    upload_directory = 'uploads/'
    app.config.from_pyfile('./flask_default_config.cfg')
    app.config.from_pyfile('./flask_config.cfg', silent=True)
    app.config['UPLOAD_FOLDER'] = str('./' + upload_directory)

    gunicorn_logger = logging.getLogger('gunicorn.debug')
    app.logger.handlers = gunicorn_logger.handlers
    app.logger.setLevel(gunicorn_logger.level)

    WARMUP_MODELS = strtobool(
        os.environ.get("FLASK_WARMUP_MODELS",
                       default=str(app.config["WARMUP_MODELS"])))

    FS_HZ = app.config["FS_HZ"]
    HOP_LENGTH = app.config["HOP_LENGTH"]
    MAX_SOUND_DURATION_S = app.config["MAX_SOUND_DURATION_S"]
    DEVICE = os.environ.get("FLASK_PYTORCH_DEVICE",
                            default=app.config["PYTORCH_DEVICE"])

    SPECTROGRAMS_UPSAMPLING_FACTOR = app.config["SPECTROGRAMS_UPSAMPLING_FACTOR"]
    DEFAULT_COLORMAP = app.config["DEFAULT_COLORMAP"]
    USE_LOCAL_CONDITIONING = app.config["USE_LOCAL_CONDITIONING"]
    SAMPLING_DATASET_JSON_PATH = app.config.get("SAMPLING_DATASET_JSON_PATH")

    with smart_open.open(app.config['VQVAE_PATH'], 'rb') as vqvae_f:
        print("Load VQ-VAE")
        vqvae = torch.jit.load(vqvae_f, map_location=DEVICE)
        vqvae.eval().to(DEVICE)

    with smart_open.open(app.config['TRANSFORMER_TOP_PATH'], 'rb'
                         ) as transformer_top_f:
        print("Load top-layer Transformer")
        transformer_top = torch.jit.load(
            transformer_top_f,
            map_location=DEVICE)
        transformer_top_f.seek(0)
        inference_transformer_top = load_jit_model(
            transformer_top_f,
            map_location=DEVICE)
        transformer_top.eval().to(DEVICE)

    with smart_open.open(app.config['TRANSFORMER_BOTTOM_PATH'], 'rb'
                         ) as transformer_bottom_f:
        print("Load bottom-layer Transformer")
        transformer_bottom = torch.jit.load(
            transformer_bottom_f,
            map_location=DEVICE)
        transformer_bottom_f.seek(0)
        inference_transformer_bottom = load_jit_model(
            transformer_bottom_f,
            map_location=DEVICE)
        transformer_bottom.eval().to(DEVICE)

    print("Retrieving label encoders")
    label_encoders_per_modality = make_label_encoders(
        transformer_top.class_conditioning_label_encoder_data_per_modality
    )

    codes_dataloader = None
    if SAMPLING_DATASET_JSON_PATH:
        # TODO(@tbazin, 2022/01/10)
        print("Loading dataset for initial sounds sampling")
        # classes_for_conditioning = ['pitch', 'instrument_family_str']
        raise NotImplementedError

    partial_sample_model = functools.partial(
        sample_model,
        device=DEVICE,
        top_k_sampling_k=app.config['SAMPLING_TOP_K'],
        top_p_sampling_p=app.config['SAMPLING_TOP_P'],
        use_predictive_sampling=app.config['USE_PREDICTIVE_SAMPLING'],
        use_parallel_model=False,
        progressbar_decorator=lambda x: x
    )

    def warmup_models():
        """Call the inference models with blank input for TorchScript warmup
        """
        top_code, bottom_code, *_ = generate()
        audio = vqvae.decode_code(top_code.to(DEVICE),
                                  bottom_code.to(DEVICE))
        vqvae.forward(audio)

    def get_codemaps_from_database(
            duration_top: int,
            attribute_constraints: Dict[str, Union[any, List[any]]] = {},
            ):
        if 'pitch_class' or 'octave' in attribute_constraints.keys():
            assert (label_encoders_per_modality is not None
                    and 'pitch' in label_encoders_per_modality)

        def decode_attributes(encoded_attributes):
            decoded_attributes = {
                key: label_encoders_per_modality[key].inverse_transform(
                    [value.item()])[0]
                for key, value in encoded_attributes.items()}
            if 'pitch_class' in attribute_constraints.keys():
                decoded_attributes['pitch_class'] = (
                    decoded_attributes['pitch'] % 12)
            if 'octave' in attribute_constraints.keys():
                decoded_attributes['pitch_class'] = (
                    decoded_attributes['pitch'] // 12)
            return decoded_attributes

        def check_attributes(
                attributes: Dict[str, Union[any, List[any]]]):
            return all([
                attributes[key] == constraint
                for key, constraint in attribute_constraints.items()
            ])

        found_valid_sample = False
        while not found_valid_sample:
            # TODO(theis, 2021/06/29): replace this brute-force search with
            # usign a proper database search with filtering etc.
            top_code, bottom_code, encoded_sample_attributes = next(iter(
                codes_dataloader))
            sample_attributes = decode_attributes(encoded_sample_attributes)
            found_valid_sample = check_attributes(sample_attributes)
        return (resize_codemaps_repeat_last(top_code, bottom_code, duration_top),
                sample_attributes)

    def get_duration_sox_n(audio_file_path: str) -> float:
        """Retrieve duration of a signal without loading it

        This uses the global sampling frequency of the loaded models
        """
        audiometadata = torchaudio.info(audio_file_path)
        num_frames = audiometadata.num_frames
        if num_frames == 0:
            num_frames = torchaudio.load(audio_file_path)[0].size(-1)
        original_fs_hz = audiometadata.sample_rate
        duration_n = num_frames
        # TODO(theis): probably not exact value
        duration_n_resampled = round(duration_n * (app.config['FS_HZ']
                                                   / original_fs_hz))
        return duration_n_resampled

    def get_duration_sox_s(audio_file_path: str) -> float:
        """Retrieve duration of a signal without loading it

        This uses the global sampling frequency of the loaded models
        """
        duration_n = get_duration_sox_n(audio_file_path)
        return duration_n / FS_HZ

    def write_audio_to_bytes(audio: torch.Tensor, sample_id: str = ''
                             ) -> Tuple[bytes, str]:
        """Generate and send WAV file
        """
        audio_format = 'wav'
        audio_np = audio.cpu().numpy()
        # normalize audio
        audio_np = audio_np / audio_np.max()
        with tempfile.NamedTemporaryFile() as f:
            soundfile.write(f.name,
                            audio_np,
                            samplerate=app.config['FS_HZ'],
                            format=audio_format)
            return filelike_to_b64(f), audio_format

    def get_vqvae_top_resolution_n() -> int:
        """Return the duration in samples of one column of the VQVAE's top layer"""
        dummy_codes_top = torch.zeros(transformer_top.shape,
                                      dtype=torch.long).to(DEVICE).unsqueeze(0)
        dummy_codes_bottom = torch.zeros(transformer_bottom.shape,
                                         dtype=torch.long).to(DEVICE).unsqueeze(0)
        decoded_audio = vqvae.spectrograms_helper.to_audio(
            vqvae.decode_code(dummy_codes_top, dummy_codes_bottom))
        _, duration_top = transformer_top.shape
        return decoded_audio.shape[-1] // duration_top

    def adapt_duration(audio_file_path: str) -> float:
        """Adapt duration of a file for loading

        Accounts for both the max duration and the VQ-VAE's resolution
        """
        duration_n = get_duration_sox_n(audio_file_path)
        # trim to max duration
        duration_n = min(MAX_SOUND_DURATION_S * FS_HZ, duration_n)
        # round-up to the resolution of the VQVAE
        vqvae_top_resolution_n = get_vqvae_top_resolution_n()
        duration_n = vqvae_top_resolution_n * (max(
            transformer_top.shape[1],
            round(duration_n / vqvae_top_resolution_n)))
        return duration_n

    def parse_conditioning(request) -> Tuple[torch.LongTensor,
                                             torch.LongTensor,
                                             Mapping[str, ConditioningMap],
                                             Mapping[str, ConditioningMap],
                                             ]:
        json_data = request.get_json(force=True)

        if 'top_conditioning' not in json_data.keys():
            return None, None

        conditioning_top = json_data['top_conditioning']
        conditioning_bottom = json_data['bottom_conditioning']

        class_conditioning_top_map = make_conditioning_map(
            conditioning_top,
            label_encoders_per_modality)
        class_conditioning_bottom_map = make_conditioning_map(
            conditioning_bottom,
            label_encoders_per_modality)

        return (class_conditioning_top_map, class_conditioning_bottom_map,
                conditioning_top, conditioning_bottom)

    @torch.inference_mode()
    def make_spectrogram_image(spectrogram: torch.Tensor,
                               filename: str = 'spectrogram',
                               upsampling_factor: int = 1,
                               format: str = 'png',
                               colormap: Optional[str] = None
                               ) -> Tuple[bytes, str]:
        """Generate and save a png image for the provided spectrogram.

        Assumes melscale frequency axis.

        Arguments:
            spectrogram (torch.Tensor): the mel-scale spectrogram to draw
        Returns:
            output_path (str): the path where the image was written
        """
        colormap = request.args.get('spectrogram_image_colormap', type=str,
                                    default=DEFAULT_COLORMAP)

        if upsampling_factor > 1:
            spectrogram = (
                torch.nn.functional.interpolate(
                    spectrogram.unsqueeze(0).unsqueeze(1),
                    mode='bilinear',
                    scale_factor=upsampling_factor)).squeeze(0).squeeze(0)
        spectrogram_np = spectrogram.cpu().numpy()

        fig, ax = full_frame(width=12, height=8)
        librosa.display.specshow(
            spectrogram_np,
            #  y_axis='mel',
            ax=ax,
            sr=FS_HZ * upsampling_factor,
            cmap=colormap if colormap is not None else DEFAULT_COLORMAP,
            hop_length=HOP_LENGTH
        )

        # write image to bytes
        with tempfile.NamedTemporaryFile() as f:
            fig.savefig(f.name, format=format,
                        dpi=20,
                        pad_inches=0, bbox_inches=0,
                        transparent=True)
            fig.clear()
            plt.close('all')
            gc.collect()

            return filelike_to_b64(f), format

    def filelike_to_b64(f: io.FileIO) -> bytes:
        """Convert bytes to base64 encoding"""
        f.seek(0)
        return base64.b64encode(f.read())

    @app.route('/ping', methods=['GET'])
    def ping():
        return "pong"

    @torch.inference_mode()
    def generate(temperature: Optional[float] = 1.0,
                 pitch: Optional[int] = None,
                 instrument_family_str: Optional[str] = None):
        class_conditioning = {
            'top': {},
            'bottom': {}
        }
        if pitch is not None:
            for layer in class_conditioning.keys():
                class_conditioning[layer]['pitch'] = pitch
        if instrument_family_str is not None:
            for layer in class_conditioning.keys():
                class_conditioning[layer]['instrument_family_str'] = (
                    instrument_family_str)

        class_conditioning_tensors_top = make_conditioning_tensors(
            class_conditioning['top'],
            label_encoders_per_modality)
        class_conditioning_tensors_bottom = make_conditioning_tensors(
            class_conditioning['bottom'],
            label_encoders_per_modality)

        batch_size = 1
        top_code = partial_sample_model(
            model=transformer_top,
            inference_model=inference_transformer_top,
            batch_size=batch_size,
            codemap_size=transformer_top.shape,
            temperature=temperature,
            class_conditioning=class_conditioning_tensors_top
        )
        bottom_code = partial_sample_model(
            model=transformer_bottom,
            inference_model=inference_transformer_bottom,
            condition=top_code,
            batch_size=batch_size,
            codemap_size=transformer_bottom.shape,
            temperature=temperature,
            class_conditioning=class_conditioning_tensors_bottom,
        )

        class_conditioning_top_map = {
            modality: make_matrix(transformer_top.shape,
                                  value)
            for modality, value in class_conditioning['top'].items()
        }
        class_conditioning_bottom_map = {
            modality: make_matrix(transformer_bottom.shape,
                                  value)
            for modality, value in class_conditioning['bottom'].items()
        }

        return (top_code, bottom_code,
                class_conditioning_top_map, class_conditioning_bottom_map)

    @app.route('/generate', methods=['GET'])
    def http_generate():
        """
        Return a new, generated sound
        Usage:
            [GET/POST] /generate?[pitch=XXX]&[instrument_family_str=XXX]&[temperature=XXX]
            - Request: empty payload, a new sound is synthesized from scratch
            - Response: a new, generated sound
        """
        temperature = request.args.get('temperature', 1, type=float)

        pitch = request.args.get('pitch', type=int)
        instrument_family_str = (
            request.args.get('instrument_family_str', type=str))
        generation = generate(temperature, pitch, instrument_family_str)

        response = make_response(*generation)
        return response

    @app.route('/sample-from-dataset', methods=['GET', 'POST'])
    @torch.inference_mode()
    def sample_from_dataset():
        abort(404)

        duration_top = request.args.get('duration_top', type=int)

        # retrieve and check sampling constraints
        constraint_pitch = (
            request.args.get('pitch', type=int))
        assert (
            constraint_pitch is None
            or constraint_pitch in label_encoders_per_modality['pitch'].classes_
            )

        constraint_pitch_class = request.args.get('pitch_class', type=int)
        if (constraint_pitch_class is not None
            and (constraint_pitch_class < 0
                 or constraint_pitch_class > 12)):
            constraint_pitch_class = None

        constraint_octave = request.args.get('octave', type=int)
        if (constraint_octave is not None and constraint_octave < 0):
            constraint_octave = None

        constraint_instrument_family_str = request.args.get(
            'instrument_family_str', type=str)
        assert (
            constraint_instrument_family_str is None
            or constraint_instrument_family_str in (
                label_encoders_per_modality['instrument_family_str'].classes_)
            )

        attribute_constraints = {}
        # TODO(theis, 2021_04_20): simplify this
        if constraint_pitch is not None:
            attribute_constraints['pitch'] = constraint_pitch
        if constraint_pitch_class is not None:
            attribute_constraints['pitch_class'] = constraint_pitch_class
        if constraint_octave is not None:
            attribute_constraints['octave'] = constraint_octave
        if constraint_instrument_family_str is not None:
            attribute_constraints['instrument_family_str'] = (
                constraint_instrument_family_str)
        (top_code, bottom_code), sampled_attributes = get_codemaps_from_database(
            duration_top, attribute_constraints)

        class_conditioning_top = class_conditioning_bottom = {
            'pitch': int(sampled_attributes['pitch']),
            'instrument_family_str': str(
                sampled_attributes['instrument_family_str'])
        }

        class_conditioning_top_map = {
            modality: make_matrix(top_code.shape,
                                  value)
            for modality, value in class_conditioning_top.items()
        }
        class_conditioning_bottom_map = {
            modality: make_matrix(bottom_code.shape,
                                  value)
            for modality, value in class_conditioning_bottom.items()
        }

        response = make_response(top_code, bottom_code,
                                 class_conditioning_top_map,
                                 class_conditioning_bottom_map)
        return response

    @app.route('/sample-from-dataset-audio', methods=['GET'])
    @torch.inference_mode()
    def sample_from_dataset_audio():
        abort(404)

        duration_top = request.args.get('duration_top', type=int)

        # retrieve and check sampling constraints
        constraint_pitch = request.args.get('pitch', type=int)
        assert (
            constraint_pitch is None
            or constraint_pitch in label_encoders_per_modality['pitch'].classes_
            )

        constraint_pitch_class = request.args.get('pitch_class', type=int)
        if (constraint_pitch_class is not None
            and (constraint_pitch_class < 0
                 or constraint_pitch_class > 12)):
            constraint_pitch_class = None

        constraint_octave = request.args.get('octave', type=int)
        if (constraint_octave is not None and constraint_octave < 0):
            constraint_octave = None

        constraint_instrument_family_str = request.args.get(
            'instrument_family_str', type=str, default=None)
        assert (
            constraint_instrument_family_str is None
            or constraint_instrument_family_str in (
                label_encoders_per_modality['instrument_family_str'].classes_)
            )

        attribute_constraints = {}
        # TODO(theis, 2021_04_20): simplify this
        if constraint_pitch is not None:
            attribute_constraints['pitch'] = constraint_pitch
        if constraint_pitch_class is not None:
            attribute_constraints['pitch_class'] = constraint_pitch_class
        if constraint_octave is not None:
            attribute_constraints['octave'] = constraint_octave
        if constraint_instrument_family_str is not None:
            attribute_constraints['instrument_family_str'] = (
                constraint_instrument_family_str)
        (top_code, bottom_code), sampled_attributes = get_codemaps_from_database(
            duration_top, attribute_constraints)
        with torch.inference_mode():
            logmelspectrogram_and_IF = vqvae.decode_code(top_code.to(DEVICE),
                                                        bottom_code.to(DEVICE))

        # convert to audio and write to file
        audio = vqvae.spectrograms_helper.to_audio(logmelspectrogram_and_IF)[0]
        audio_path = write_audio_to_bytes(audio)

        return flask.send_file(audio_path, mimetype="audio/wav",
                               max_age=-1  # disable cache
                               )

    @app.route('/test-generate', methods=['GET', 'POST'])
    @torch.inference_mode()
    def test_generate():
        pitch = request.args.get('pitch', type=int)
        instrument_family_str = request.args.get('instrument_family_str',
                                                 type=str)

        class_conditioning_top = class_conditioning_bottom = {
            'pitch': pitch,
            'instrument_family_str': instrument_family_str
        }

        top_code = torch.randint(size=transformer_top.shape, low=0,
                                 high=vqvae.n_embed_t).unsqueeze(0)
        bottom_code = torch.randint(size=transformer_bottom.shape, low=0,
                                    high=vqvae.n_embed_b).unsqueeze(0)

        class_conditioning_top_map = {
            modality: make_matrix(transformer_top.shape,
                                  value)
            for modality, value in class_conditioning_top.items()
        }
        class_conditioning_bottom_map = {
            modality: make_matrix(transformer_bottom.shape,
                                  value)
            for modality, value in class_conditioning_bottom.items()
        }

        response = make_response(top_code, bottom_code,
                                 class_conditioning_top_map,
                                 class_conditioning_bottom_map)
        return response

    @app.route('/analyze-audio', methods=['POST'])
    @torch.inference_mode()
    def audio_to_codes():
        if 'audio' not in request.files.keys():
            app.logger.error('No `audio` key in payload (expects a MultiPart-Form)')
            abort(400, 'Please provide an audio file through a MultiPart-Form with a file at key `audio`')
        uploaded_audio_file = request.files['audio']
        uploaded_audio_filename = secure_filename(uploaded_audio_file.filename)
        if uploaded_audio_filename != '':
            file_ext = os.path.splitext(uploaded_audio_filename)[1][1:]
            if file_ext not in app.config['UPLOAD_EXTENSIONS']:
                app.logger.error(f"File extension `.{file_ext}` is not supported")
                abort(400, f"File extension `.{file_ext}` is not supported")

        pitch = request.args.get('pitch', type=int)
        instrument_family_str = request.args.get('instrument_family_str',
                                                 type=str)

        class_conditioning_top = class_conditioning_bottom = {
            'pitch': pitch,
            'instrument_family_str': instrument_family_str
        }

        # TODO(@tbazin, 2022/02/28): avoid using temporary file, directly pass
        # around bytes object?
        with tempfile.NamedTemporaryFile(
                'w+b',
                suffix=secure_filename(uploaded_audio_filename)
                ) as f:
            temp_filename = f.name
            uploaded_audio_file.save(temp_filename)
            duration_n = adapt_duration(temp_filename)
            spec_and_IF = vqvae.spectrograms_helper.from_wavfile(
                temp_filename, duration_n=duration_n).to(DEVICE)

        _, _, _, top_code, bottom_code, *_ = vqvae.encode(spec_and_IF)

        class_conditioning_top_map = {
            modality: make_matrix(transformer_top.shape,
                                  value)
            for modality, value in class_conditioning_top.items()
        }
        class_conditioning_bottom_map = {
            modality: make_matrix(transformer_bottom.shape,
                                  value)
            for modality, value in class_conditioning_bottom.items()
        }

        response = make_response(top_code, bottom_code,
                                 class_conditioning_top_map,
                                 class_conditioning_bottom_map)
        return response

    @app.route('/timerange-change', methods=['POST'])
    @torch.inference_mode()
    def timerange_change():
        """
        Perform local re-generation on a sheet and return the updated sheet
        Usage:
            POST /timerange-change?TODO(theis)
            - Request:
            - Response:
        """
        layer = request.args.get('layer', type=str)
        temperature = request.args.get('temperature', default=1, type=float)
        start_index_top = request.args.get('start_index_top', type=int)
        uniform_sampling = bool(strtobool(
            request.args.get('uniform_sampling', type=str,
                             default="False")))

        # try to retrieve local conditioning map in the request's JSON payload
        (class_conditioning_top_map, class_conditioning_bottom_map,
         input_conditioning_top, input_conditioning_bottom) = (
            parse_conditioning(request)
        )
        global_instrument_family_str = request.args.get(
            'instrument_family_str', type=str)
        global_pitch = request.args.get('pitch', type=int)
        global_class_conditioning = {
            'pitch': global_pitch,
            'instrument_family_str': global_instrument_family_str
        }
        if (not USE_LOCAL_CONDITIONING
                or not transformer_bottom.local_class_conditioning):
            class_conditioning_bottom = global_class_conditioning.copy()
            class_conditioning_tensors_bottom = make_conditioning_tensors(
                class_conditioning_bottom,
                label_encoders_per_modality)
            class_conditioning_bottom_map = None
        else:
            class_conditioning_bottom = class_conditioning_tensors_bottom = None

        top_code, bottom_code = parse_codes(request)

        # extract frame to operate on
        end_index_top = start_index_top + transformer_top.shape[1]
        top_code_frame = top_code[..., start_index_top:end_index_top]

        upsampling_ratio_time = (transformer_bottom.shape[1]
                                 // transformer_top.shape[1])
        start_index_bottom = upsampling_ratio_time * start_index_top
        end_index_bottom = start_index_bottom + transformer_bottom.shape[1]
        bottom_code_frame = bottom_code[..., start_index_bottom:end_index_bottom]
        generation_mask_batched = parse_mask(request).to(DEVICE)

        time_indexes_top = make_time_indexes(start_index_top,
                                             top_code.shape[-1],
                                             transformer_top.shape[-1])
        time_indexes_bottom = make_time_indexes(start_index_bottom,
                                                bottom_code.shape[-1],
                                                transformer_bottom.shape[-1])

        if layer == 'bottom':
            if not uniform_sampling:
                bottom_code_resampled_frame = partial_sample_model(
                    model=transformer_bottom,
                    condition=top_code_frame,
                    batch_size=1,
                    codemap_size=transformer_bottom.shape,
                    temperature=temperature,
                    class_conditioning=class_conditioning_tensors_bottom,
                    local_class_conditioning_map=class_conditioning_bottom_map,
                    initial_code=bottom_code_frame,
                    mask=generation_mask_batched,
                    time_indexes_source=time_indexes_top,
                    time_indexes_target=time_indexes_bottom,
                )
            else:
                bottom_code_resampled_frame = bottom_code_frame.masked_scatter(
                    generation_mask_batched,
                    torch.randint_like(bottom_code_frame,
                                       high=transformer_bottom.n_class_target)
                )

            bottom_code_resampled = bottom_code
            bottom_code_resampled[..., start_index_bottom:end_index_bottom] = (
                bottom_code_resampled_frame)

            # create JSON response
            response = make_response(top_code, bottom_code_resampled,
                                     input_conditioning_top,
                                     input_conditioning_bottom)
        elif layer == 'top':
            if (not USE_LOCAL_CONDITIONING
                    or not transformer_top.local_class_conditioning):
                # try to retrieve conditioning from http arguments
                class_conditioning_top = global_class_conditioning.copy()
                class_conditioning_tensors_top = make_conditioning_tensors(
                    class_conditioning_top,
                    label_encoders_per_modality)
                class_conditioning_top_map = None
            else:
                class_conditioning_top = class_conditioning_tensors_top = None

            if not uniform_sampling:
                if transformer_top.self_conditional_model:
                    condition = top_code_frame
                else:
                    condition = None
                top_code_resampled_frame = partial_sample_model(
                    model=transformer_top,
                    condition=condition,
                    device=DEVICE,
                    batch_size=1,
                    codemap_size=transformer_top.shape,
                    temperature=temperature,
                    class_conditioning=class_conditioning_tensors_top,
                    local_class_conditioning_map=class_conditioning_top_map,
                    initial_code=top_code_frame,
                    mask=generation_mask_batched,
                    time_indexes_source=time_indexes_top,
                    time_indexes_target=time_indexes_top,
                )
            else:
                top_code_resampled_frame = top_code_frame.masked_scatter(
                    generation_mask_batched,
                    torch.randint_like(top_code_frame,
                                       high=transformer_top.n_class_target)
                )

            top_code_resampled = top_code
            top_code_resampled[..., start_index_top:end_index_top] = (
                top_code_resampled_frame)

            upsampling_ratio_frequency = (transformer_bottom.shape[0]
                                          // transformer_top.shape[0])
            generation_mask_bottom_batched = (
                generation_mask_batched
                .repeat_interleave(upsampling_ratio_frequency, -2)
                .repeat_interleave(upsampling_ratio_time, -1)
            )
            bottom_code_resampled_frame = partial_sample_model(
                model=transformer_bottom,
                condition=top_code_resampled_frame,
                device=DEVICE,
                batch_size=1,
                codemap_size=transformer_bottom.shape,
                temperature=temperature,
                class_conditioning=class_conditioning_tensors_bottom,
                local_class_conditioning_map=class_conditioning_bottom_map,
                initial_code=bottom_code_frame,
                mask=generation_mask_bottom_batched,
                time_indexes_source=time_indexes_top,
                time_indexes_target=time_indexes_bottom,
            )

            # update conditioning map
            bottom_mask = generation_mask_bottom_batched[0]
            new_conditioning_map_bottom = {
                modality: masked_fill(modality_conditioning,
                                      bottom_mask,
                                      class_conditioning_bottom[modality])
                for modality, modality_conditioning
                in input_conditioning_bottom.items()
            }

            bottom_code_resampled = bottom_code
            bottom_code_resampled[..., start_index_bottom:end_index_bottom] = (
                bottom_code_resampled_frame)

            # create JSON response
            response = make_response(top_code_resampled, bottom_code_resampled,
                                     input_conditioning_top,
                                     new_conditioning_map_bottom)

        return response

    @app.route('/erase', methods=['POST'])
    @torch.inference_mode()
    def erase():
        amplitude = request.args.get('eraser_amplitude', type=float)
        start_index_top = request.args.get('start_index_top', type=int)

        top_code_batched, bottom_code_batched = parse_codes(request)
        generation_mask = parse_mask(request).to(DEVICE)[0]

        logmelspectrogram, IF = vqvae.decode_code(top_code_batched,
                                                  bottom_code_batched)[0]

        top_code = top_code_batched[0]
        upsampling_f = logmelspectrogram.shape[0] // top_code.shape[0]
        upsampling_t = logmelspectrogram.shape[1] // top_code.shape[1]

        upsampled_mask = (generation_mask.float().flip(0)
                          .repeat_interleave(upsampling_f, 0)
                          .repeat_interleave(upsampling_t, 1)
                          ).flip(0)
        amplitude_mask = 200 * amplitude * upsampled_mask

        # zero-pad the amplitude mask
        padding_before = torch.zeros(logmelspectrogram.shape[0],
                                     upsampling_t * start_index_top)
        padding_after = torch.zeros(logmelspectrogram.shape[0],
                                    max(0,
                                        upsampling_t * (
                                            top_code.shape[1] - (
                                                start_index_top + generation_mask.shape[1]))))
        amplitude_mask = torch.cat([
            padding_before.to(DEVICE),
            amplitude_mask,
            padding_after.to(DEVICE)], dim=1)

        masked_logmelspectrogram_and_IF = torch.cat(
            [(logmelspectrogram - amplitude_mask).unsqueeze(0),
             IF.unsqueeze(0)],
            dim=0
        ).unsqueeze(0)

        _, _, _, new_top_code, new_bottom_code, *_ = vqvae.encode(
            masked_logmelspectrogram_and_IF)

        (_, _, input_conditioning_top, input_conditioning_bottom) = (
            parse_conditioning(request))
        return make_response(new_top_code, new_bottom_code,
                             input_conditioning_top,
                             input_conditioning_bottom)

    @torch.inference_mode()
    def parse_codes(request, device: torch.device = DEVICE
                    ) -> Tuple[torch.LongTensor, torch.LongTensor]:
        json_data = request.get_json(force=True)

        top_code_array = json_data['top_code']
        bottom_code_array = json_data['bottom_code']

        top_code = torch.LongTensor(top_code_array
                                    ).unsqueeze(0).to(device)
        bottom_code = torch.LongTensor(bottom_code_array
                                       ).unsqueeze(0).to(device)

        return top_code, bottom_code

    @app.route('/get-audio', methods=['POST'])
    def get_audio():
        top_code, bottom_code = parse_codes(request, DEVICE)
        audio_b64, audio_format = codes_to_audio_bytes(top_code, bottom_code)
        return flask.send_file(io.BytesIO(audio_b64),
                               mimetype=f"audio/{audio_format}",
                               cache_timeout=-1  # disable cache
                               )

    @app.route('/get-spectrogram-image', methods=['POST'])
    @torch.inference_mode()
    def get_spectrogram_image():
        colormap = request.args.get('spectrogram_image_colormap', None,
                                    type=str)
        top_code, bottom_code = parse_codes(request, DEVICE)
        spectrogram_image_b64, image_format = (
            codes_to_spectrogram_image_bytes(top_code, bottom_code,
                                             colormap=colormap)
            )
        return flask.send_file(io.BytesIO(spectrogram_image_b64),
                               mimetype=f"image/{image_format}",
                               cache_timeout=-1  # disable cache
                               )

    @app.route('/top-conditioned-sample', methods=['POST'])
    @torch.inference_mode()
    def top_conditioned_sample():
        """Sample from the bottom prior given the incoming top codemap"""
        BYPASS = False

        top_code, bottom_code = parse_codes(request)
        global_instrument_family_str = request.args.get(
            'instrument_family_str', type=str)
        min_pitch = request.args.get('min_pitch', type=int)
        max_pitch = request.args.get('max_pitch', type=int)

        if not BYPASS:
            temperature = request.args.get('temperature', type=float,
                                           default=1.)
            top_p = request.args.get('top_p', type=float, default=0.)
            top_k = request.args.get('top_k', type=int, default=0)

            class_conditioning_tensors_bottom = make_conditioning_tensors(
                {'pitch': (min_pitch, max_pitch),
                 'instrument_family_str': global_instrument_family_str},
                label_encoders_per_modality
            )

            # repeat the top codemap for all bottom samples
            num_samples = max_pitch - min_pitch
            top_code = top_code.expand(num_samples, -1, -1)

            bottom_code = sample_model(
                transformer_bottom, DEVICE, num_samples,
                transformer_bottom.shape,
                temperature, condition=top_code,
                class_conditioning=class_conditioning_tensors_bottom,
                top_p_sampling_p=top_p,
                top_k_sampling_k=top_k
            )
        else:
            import time
            num_samples = 1
            top_code = top_code.expand(num_samples, -1, -1)
            bottom_code = bottom_code.expand(num_samples, -1, -1)
            time.sleep(2)

        logmelspectrogram_and_IF = vqvae.decode_code(top_code,
                                                     bottom_code)

        zip_path = upload_directory + 'samples.zip'
        with ZipFile(zip_path, 'w') as zf:
            for pitch, sample in zip(range(min_pitch, max_pitch),
                                     logmelspectrogram_and_IF):
                # convert to audio and write to file
                audio = vqvae.spectrograms_helper.to_audio(sample.unsqueeze(0))[0]
                audio_path = write_audio_to_bytes(
                    audio,
                    f'-{global_instrument_family_str}-{pitch}')
                zf.write(audio_path,
                         arcname=f'{global_instrument_family_str}-{pitch}.wav')

        return flask.send_file(zip_path, mimetype="application/zip",
                               cache_timeout=-1  # disable cache
                               )

    @torch.inference_mode()
    def codes_to_audio(top_code: torch.LongTensor,
                       bottom_code: torch.LongTensor) -> torch.Tensor:
        logmelspectrogram_and_IF = vqvae.decode_code(
            top_code, bottom_code)
        audio = vqvae.spectrograms_helper.to_audio(logmelspectrogram_and_IF)[0]
        return audio

    def codes_to_audio_bytes(top_code: torch.LongTensor,
                             bottom_code: torch.LongTensor) -> Tuple[bytes, str]:
        audio = codes_to_audio(top_code, bottom_code)
        audio_b64, audio_format = write_audio_to_bytes(audio)
        return audio_b64, audio_format

    @torch.inference_mode()
    def codes_to_spectrogram_image_bytes(
            top_code: torch.torch.LongTensor,
            bottom_code: torch.torch.LongTensor,
            colormap: Optional[str] = None
            ) -> Tuple[bytes, str]:
        logmelspectrogram_and_IF = vqvae.decode_code(top_code,
                                                     bottom_code)

        # generate spectrogram PNG image
        spectrogram = logmelspectrogram_and_IF[0, 0]
        image_format = 'png'
        spectrogram_image_b64, image_format = make_spectrogram_image(
            spectrogram,
            upsampling_factor=SPECTROGRAMS_UPSAMPLING_FACTOR,
            format=image_format,
            colormap=colormap)
        return spectrogram_image_b64, image_format

    @torch.inference_mode()
    def make_response(top_code: torch.Tensor,
                      bottom_code: torch.Tensor,
                      class_conditioning_top_map: Mapping[str, ConditioningMap],
                      class_conditioning_bottom_map: Mapping[str, ConditioningMap],
                      send_files: bool = False):
        top_code = top_code.to(DEVICE)
        bottom_code = bottom_code.to(DEVICE)

        logmelspectrogram_and_IF = vqvae.decode_code(top_code, bottom_code)

        audio = vqvae.spectrograms_helper.to_audio(logmelspectrogram_and_IF)[0]
        audio_b64, audio_format = write_audio_to_bytes(audio)

        spectrogram = logmelspectrogram_and_IF[0, 0]
        image_format = 'png'
        spectrogram_image_b64, image_format = make_spectrogram_image(
            spectrogram,
            upsampling_factor=SPECTROGRAMS_UPSAMPLING_FACTOR,
            format=image_format)

        return flask.jsonify({
            'top_code': top_code[0].int().cpu().numpy().tolist(),
            'bottom_code': bottom_code[0].int().cpu().numpy().tolist(),

            'top_conditioning': class_conditioning_top_map,
            'bottom_conditioning': class_conditioning_bottom_map,

            'audio': {
                'data': audio_b64.decode('ascii'),
                'content-type': f'audio/{audio_format}'
                },
            'image': {
                'data': spectrogram_image_b64.decode('ascii'),
                'content-type': f'image/{image_format}'
                }})

    if WARMUP_MODELS:
        print("Warming up inference models")
        warmup_models()

    print("Application ready")
    return app


def full_frame(width=None, height=None):
    """Initialize a full-frame matplotlib figure and axes

    Taken from a GitHub Gist by Kile McDonald:
    https://gist.github.com/kylemcdonald/bedcc053db0e7843ef95c531957cb90f
    """
    import matplotlib as mpl
    mpl.rcParams['savefig.pad_inches'] = 0
    figsize = None if width is None else (width, height)
    fig = Figure(figsize=figsize)
    ax = fig.add_axes((0, 0, 1, 1), frameon=False)
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
    ax.autoscale(tight=True)
    return fig, ax


def load_jit_model(path,
                   preserved_attrs: Optional[List[str]] = [],
                   map_location=None):
    return torch.jit.optimize_for_inference(
        torch.jit.freeze(
            torch.jit.load(path, map_location).eval(),
            preserved_attrs=preserved_attrs
        )
    )


def make_matrix(shape: Tuple[int, int],
                value: Union[str, int]
                ) -> ConditioningMap:
    return [[value] * shape[1]] * shape[0]


def masked_fill(array, mask, value):
    return [[value if mask_value else previous_value
             for previous_value, mask_value in zip(array_row, mask_row)]
            for array_row, mask_row in zip(array, mask)]


def resize_codemaps_repeat_last(
        top_code: torch.Tensor, bottom_code: torch.Tensor,
        duration_top: int) -> Tuple[torch.Tensor, torch.Tensor]:
    upsampling_ratio_time = bottom_code.shape[-1] // top_code.shape[-1]
    duration_bottom = upsampling_ratio_time * duration_top

    def resize_codemap(codemap: torch.Tensor, duration: int) -> torch.Tensor:
        codemap = codemap[..., :duration]
        if codemap.shape[-1] < duration:
            codemap = torch.cat([codemap] + [codemap[..., -1:]] * (
                duration - codemap.shape[-1]),
                                dim=-1)
        return codemap

    return tuple(resize_codemap(codemap, duration)
                 for (codemap, duration) in zip((top_code, bottom_code),
                                                (duration_top, duration_bottom)))


def parse_mask(request) -> torch.BoolTensor:
    json_data = request.get_json(force=True)

    generation_mask_array = json_data['mask']
    generation_mask = torch.BoolTensor(generation_mask_array
                                       ).unsqueeze(0)

    return generation_mask


def make_time_indexes(start_index: int, codemap_duration: int,
                      transformer_duration: int) -> List[int]:
    time_indexes_full = [0]  # attack
    num_steps_to_repeat = transformer_duration - 2
    steps_repetitions = (codemap_duration - 2) // num_steps_to_repeat
    for i in range(num_steps_to_repeat - 1):
        time_indexes_full += [i+1] * steps_repetitions
    time_indexes_full += [num_steps_to_repeat] * (
        (codemap_duration - 2) - (len(time_indexes_full)-1))
    time_indexes_full += [transformer_duration-1]

    return time_indexes_full[start_index:
                             start_index+transformer_duration]