Update app.py

app.py

@@ -133,81 +133,174 @@ class InferenceModel(object):
         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'
+	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):
+		"""解析用于训练模型的 gin 文件。"""
+		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):
+		"""在解析训练 gin 配置后加载 T5X `Model`。"""
+		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):
+		"""从检查点中恢复训练状态，重置 self._predict_fn()。"""
+		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):
+		"""从预处理的数据集批次中预测 tokens。"""
+		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):
+		"""从音频样本推断出音符序列。
+
+		参数：
+			audio：16kHz 的单个音频样本的 1 维 numpy 数组。
+		返回：
+			转录音频的音符序列。
+		"""
+		print(f"[{current_time()}] 运行：从音频样本中推断音符序列")
+		ds = self.audio_to_dataset(audio)
+		ds = self.preprocess(ds)
+
+		model_ds = self.model.FEATURE_CONVERTER_CLS(pack=False)(
+				ds, task_feature_lengths=self.sequence_length)
+		model_ds = model_ds.batch(self.batch_size)
+
+		inferences = (tokens for batch in model_ds.as_numpy_iterator()
+									for tokens in self.predict_tokens(batch))
+
+		predictions = []
+		for example, tokens in zip(ds.as_numpy_iterator(), inferences):
+			predictions.append(self.postprocess(tokens, example))
+
+		result = metrics_utils.event_predictions_to_ns(
+				predictions, codec=self.codec, encoding_spec=self.encoding_spec)
+		return result['est_ns']
+
+	def audio_to_dataset(self, audio):
+		"""从输入音频创建一个包含频谱图的 TF Dataset。"""
+		print(f"[{current_time()}] 运行：从音频创建包含频谱图的 TF Dataset")
+		frames, frame_times = self._audio_to_frames(audio)
+		return tf.data.Dataset.from_tensors({
+				'inputs': frames,
+				'input_times': frame_times,
+		})
+
+	def _audio_to_frames(self, audio):
+		"""从音频计算频谱图帧。"""
+		print(f"[{current_time()}] 运行：从音频计算频谱图帧")
+		frame_size = self.spectrogram_config.hop_width
+		padding = [0, frame_size - len(audio) % frame_size]
+		audio = np.pad(audio, padding, mode='constant')
+		frames = spectrograms.split_audio(audio, self.spectrogram_config)
+		num_frames = len(audio) // frame_size
+		times = np.arange(num_frames) / self.spectrogram_config.frames_per_second
+		return frames, times
+
+	def preprocess(self, ds):
+		pp_chain = [
+				functools.partial(
+						t5.data.preprocessors.split_tokens_to_inputs_length,
+						sequence_length=self.sequence_length,
+						output_features=self.output_features,
+						feature_key='inputs',
+						additional_feature_keys=['input_times']),
+				# 在训练期间进行缓存。
+				preprocessors.add_dummy_targets,
+				functools.partial(
+						preprocessors.compute_spectrograms,
+						spectrogram_config=self.spectrogram_config)
+		]
+		for pp in pp_chain:
+			ds = pp(ds)
+		return ds
+
+	def postprocess(self, tokens, example):
+		tokens = self._trim_eos(tokens)
+		start_time = example['input_times'][0]
+		# 向下取整到最接近的符号化时间步。
+		start_time -= start_time % (1 / self.codec.steps_per_second)
+		return {
+				'est_tokens': tokens,
+				'start_time': start_time,
+				# 内部 MT3 代码期望原始输入，这里不使用。
+				'raw_inputs': []
+		}
+
+	@staticmethod
+	def _trim_eos(tokens):
+		tokens = np.array(tokens, np.int32)
+		if vocabularies.DECODED_EOS_ID in tokens:
+			tokens = tokens[:np.argmax(tokens == vocabularies.DECODED_EOS_ID)]
+		return tokens
+
+
+inference_model = InferenceModel('/home/user/app/checkpoints/mt3/', 'mt3')
+
+
+def inference(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 演示。要使用它，只需上传音频文件，或点击示例以查看效果。更多信息请参阅下面的链接。"
@@ -224,4 +317,4 @@ gr.Interface(
     description=description,
     article=article,
     examples=examples
-).launch(server_port=7861)
+).launch(server_port=7861)