app.py

import gradio as gr
import os
import datetime
import pytz
from pathlib import Path

def current_time():
    current = datetime.datetime.now(pytz.timezone('Asia/Shanghai')).strftime("%Y年-%m月-%d日 %H时:%M分:%S秒")
    return current

print(f"[{current_time()}] 开始部署空间...")

print(f"[{current_time()}] 日志：安装 - gsutil")
os.system("pip install gsutil")
print(f"[{current_time()}] 日志：Git - 克隆 Github 的 T5X 训练框架到当前目录")
os.system("git clone --branch=main https://github.com/google-research/t5x")
print(f"[{current_time()}] 日志：文件 - 移动 t5x 到当前目录并重命名为 t5x_tmp 并删除")
os.system("mv t5x t5x_tmp; mv t5x_tmp/* .; rm -r t5x_tmp")
print(f"[{current_time()}] 日志：编辑 - 替换 setup.py 内的文本“jax[tpu]”为“jax”")
os.system("sed -i 's:jax\[tpu\]:jax:' setup.py")
print(f"[{current_time()}] 日志：Python - 使用 pip 安装 当前目录内的 Python 包")
os.system("python3 -m pip install -e .")
print(f"[{current_time()}] 日志：Python - 更新 Python 包管理器 pip")
os.system("python3 -m pip install --upgrade pip")
print(f"[{current_time()}] 日志：安装 - langchain")
os.system("pip install langchain")
print(f"[{current_time()}] 日志：安装 - sentence-transformers")
os.system("pip install sentence-transformers")

print(f"[{current_time()}] 日志：Git - 克隆 Github 的 airio 到当前目录")
os.system("git clone --branch=main https://github.com/google/airio")
print(f"[{current_time()}] 日志：文件 - 移动 airio 到当前目录并重命名为 airio_tmp 并删除")
os.system("mv airio airio_tmp; mv airio_tmp/* .; rm -r airio_tmp")
print(f"[{current_time()}] 日志：Python - 使用 pip 安装 当前目录内的 Python 包")
os.system("python3 -m pip install -e .")

print(f"[{current_time()}] 日志：Git - 克隆 Github 的 MT3 模型到当前目录")
os.system("git clone --branch=main https://github.com/magenta/mt3")
print(f"[{current_time()}] 日志：文件 - 移动 mt3 到当前目录并重命名为 mt3_tmp 并删除")
os.system("mv mt3 mt3_tmp; mv mt3_tmp/* .; rm -r mt3_tmp")
print(f"[{current_time()}] 日志：Python - 使用 pip 从 storage.googleapis.com 安装 jax[cuda11_local] nest-asyncio pyfluidsynth")
os.system("python3 -m pip install jax[cuda11_local] nest-asyncio pyfluidsynth==1.3.0 -e . -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html")
print(f"[{current_time()}] 日志：安装 - 更新 jaxlib")
os.system("pip install --upgrade jaxlib")
print(f"[{current_time()}] 日志：Python - 使用 pip 安装 当前目录内的 Python 包")
os.system("python3 -m pip install -e .")
print(f"[{current_time()}] 日志：安装 - TensorFlow CPU")
os.system("pip install tensorflow_cpu")

print(f"[{current_time()}] 日志：gsutil - 复制 MT3 检查点到当前目录")
os.system("gsutil -q -m cp -r gs://mt3/checkpoints .")

print(f"[{current_time()}] 日志：gsutil - 复制 SoundFont 文件到当前目录")
os.system("gsutil -q -m cp gs://magentadata/soundfonts/SGM-v2.01-Sal-Guit-Bass-V1.3.sf2 .")

print(f"[{current_time()}] 日志：导入 - 必要工具")
import functools
import os
import numpy as np
import tensorflow.compat.v2 as tf

import gin
import jax
import librosa
import note_seq
import seqio
import t5
import t5x

from mt3 import metrics_utils
from mt3 import models
from mt3 import network
from mt3 import note_sequences
from mt3 import preprocessors
from mt3 import spectrograms
from mt3 import vocabularies

import nest_asyncio
nest_asyncio.apply()

SAMPLE_RATE = 16000
SF2_PATH = 'SGM-v2.01-Sal-Guit-Bass-V1.3.sf2'

def upload_audio(audio, sample_rate):
    return note_seq.audio_io.wav_data_to_samples_librosa(
        audio, sample_rate=sample_rate)


print(f"[{current_time()}] 日志：开始包装模型...")
class InferenceModel(object):
    """音乐转录的 T5X 模型包装器。"""

    def __init__(self, checkpoint_path, model_type='mt3'):
        if model_type == 'ismir2021':
            num_velocity_bins = 127
            self.encoding_spec = note_sequences.NoteEncodingSpec
            self.inputs_length = 512
        elif model_type == 'mt3':
            num_velocity_bins = 1
            self.encoding_spec = note_sequences.NoteEncodingWithTiesSpec
            self.inputs_length = 256
        else:
            raise ValueError('unknown model_type: %s' % model_type)

        gin_files = ['/home/user/app/mt3/gin/model.gin',
                    '/home/user/app/mt3/gin/mt3.gin']

        self.batch_size = 8
        self.outputs_length = 1024
        self.sequence_length = {'inputs': self.inputs_length,
                                'targets': self.outputs_length}

        self.partitioner = t5x.partitioning.PjitPartitioner(
                model_parallel_submesh=None, num_partitions=1)

        print(f"[{current_time()}] 日志：构建编解码器")
        self.spectrogram_config = spectrograms.SpectrogramConfig()
        self.codec = vocabularies.build_codec(
                vocab_config=vocabularies.VocabularyConfig(
                num_velocity_bins=num_velocity_bins)
                )
        self.vocabulary = vocabularies.vocabulary_from_codec(self.codec)
        self.output_features = {
                'inputs': seqio.ContinuousFeature(dtype=tf.float32, rank=2),
                'targets': seqio.Feature(vocabulary=self.vocabulary),
        }

        print(f"[{current_time()}] 日志：创建 T5X 模型")
        self._parse_gin(gin_files)
        self.model = self._load_model()

        print(f"[{current_time()}] 日志：恢复模型检查点")
        self.restore_from_checkpoint(checkpoint_path)

    @property
    def input_shapes(self):
        return {
                    'encoder_input_tokens': (self.batch_size, self.inputs_length),
                    'decoder_input_tokens': (self.batch_size, self.outputs_length)
        }

    def _parse_gin(self, gin_files):
        print(f"[{current_time()}] 日志：解析 gin 文件")
        gin_bindings = [
                'from __gin__ import dynamic_registration',
                'from mt3 import vocabularies',
                'VOCAB_CONFIG=@vocabularies.VocabularyConfig()',
                'vocabularies.VocabularyConfig.num_velocity_bins=%NUM_VELOCITY_BINS'
        ]
        with gin.unlock_config():
            gin.parse_config_files_and_bindings(
                    gin_files, gin_bindings, finalize_config=False)

    def _load_model(self):
        print(f"[{current_time()}] 日志：加载 T5X 模型")
        model_config = gin.get_configurable(network.T5Config)()
        module = network.Transformer(config=model_config)
        return models.ContinuousInputsEncoderDecoderModel(
                module=module,
                input_vocabulary=self.output_features['inputs'].vocabulary,
                output_vocabulary=self.output_features['targets'].vocabulary,
                optimizer_def=t5x.adafactor.Adafactor(decay_rate=0.8, step_offset=0),
                input_depth=spectrograms.input_depth(self.spectrogram_config))


    def restore_from_checkpoint(self, checkpoint_path):
        print(f"[{current_time()}] 日志：从检查点恢复训练状态")
        train_state_initializer = t5x.utils.TrainStateInitializer(
            optimizer_def=self.model.optimizer_def,
            init_fn=self.model.get_initial_variables,
            input_shapes=self.input_shapes,
            partitioner=self.partitioner)

        restore_checkpoint_cfg = t5x.utils.RestoreCheckpointConfig(
                path=checkpoint_path, mode='specific', dtype='float32')

        train_state_axes = train_state_initializer.train_state_axes
        self._predict_fn = self._get_predict_fn(train_state_axes)
        self._train_state = train_state_initializer.from_checkpoint_or_scratch(
                [restore_checkpoint_cfg], init_rng=jax.random.PRNGKey(0))

    @functools.lru_cache()
    def _get_predict_fn(self, train_state_axes):
        print(f"[{current_time()}] 日志：生成用于解码的预测函数")
        def partial_predict_fn(params, batch, decode_rng):
            return self.model.predict_batch_with_aux(
                    params, batch, decoder_params={'decode_rng': None})
        return self.partitioner.partition(
                partial_predict_fn,
                in_axis_resources=(
                        train_state_axes.params,
                        t5x.partitioning.PartitionSpec('data',), None),
                out_axis_resources=t5x.partitioning.PartitionSpec('data',)
        )

    def predict_tokens(self, batch, seed=0):
        print(f"[{current_time()}] 运行：从预处理数据集中预测音符序列")
        prediction, _ = self._predict_fn(
    self._train_state.params, batch, jax.random.PRNGKey(seed))
        return self.vocabulary.decode_tf(prediction).numpy()

    def __call__(self, audio):
        filename = os.path.basename(audio)  # 获取输入文件的文件名
        print(f"[{current_time()}] 运行：输入文件: {filename}")
        with open(audio, 'rb') as fd:
            contents = fd.read()
        audio = upload_audio(contents,sample_rate=16000)
        est_ns = inference_model(audio)
        note_seq.sequence_proto_to_midi_file(est_ns, './transcribed.mid')
        return './transcribed.mid'

title = "MT3"
description = "MT3：多任务多音轨音乐转录的 Gradio 演示。要使用它，只需上传音频文件，或点击示例以查看效果。更多信息请参阅下面的链接。"

article = "<p style='text-align: center'>出错了？试试把文件转换为MP3后再上传吧~</p><p style='text-align: center'><a href='https://arxiv.org/abs/2111.03017' target='_blank'>MT3: 多任务多音轨音乐转录</a> | <a href='https://github.com/hmjz100/mt3' target='_blank'>Github 仓库</a></p>"

examples=[['canon.flac'], ['download.wav']]

gr.Interface(
    inference,
    gr.Audio(type="filepath", label="输入"),
    outputs=gr.File(label="输出"),
    title=title,
    description=description,
    article=article,
    examples=examples
).launch(server_port=7861)