Update app.py

app.py

@@ -6,18 +6,15 @@ from pathlib import Path
 os.system("pip install gsutil")
 
 
-os.system("git clone --branch=main https://github.com/google-research/t5x")
-os.system("mv t5x t5x_tmp; mv t5x_tmp/* .; rm -r t5x_tmp")
-os.system("sed -i 's:jax\[tpu\]:jax:' setup.py")
-os.system("python3 -m pip install -e .")
-os.system("python3 -m pip install --upgrade pip")
+
 
 
 
 # install mt3
 os.system("git clone --branch=main https://github.com/magenta/mt3")
 os.system("mv mt3 mt3_tmp; mv mt3_tmp/* .; rm -r mt3_tmp")
-os.system("python3 -m pip install -e .")
+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
+")
 
 # copy checkpoints
 os.system("gsutil -q -m cp -r gs://mt3/checkpoints .")
@@ -35,18 +32,13 @@ import functools
 import os
 
 import numpy as np
-
 import tensorflow.compat.v2 as tf
 
 import functools
 import gin
-import jax.linear_util
-jax.extend.linear_util = jax.linear_util
+import jax
 import librosa
 import note_seq
-
-
-
 import seqio
 import t5
 import t5x
@@ -59,6 +51,7 @@ from mt3 import preprocessors
 from mt3 import spectrograms
 from mt3 import vocabularies
 
+from google.colab import files
 
 import nest_asyncio
 nest_asyncio.apply()
@@ -66,9 +59,12 @@ nest_asyncio.apply()
 SAMPLE_RATE = 16000
 SF2_PATH = 'SGM-v2.01-Sal-Guit-Bass-V1.3.sf2'
 
-def upload_audio(audio, sample_rate):
+def upload_audio(sample_rate):
+  data = list(files.upload().values())
+  if len(data) > 1:
+    print('Multiple files uploaded; using only one.')
   return note_seq.audio_io.wav_data_to_samples_librosa(
-    audio, sample_rate=sample_rate)
+    data[0], sample_rate=sample_rate)
 
 
 
@@ -89,16 +85,16 @@ class InferenceModel(object):
     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']
+    gin_files = ['/content/mt3/gin/model.gin',
+                 f'/content/mt3/gin/{model_type}.gin']
 
     self.batch_size = 8
     self.outputs_length = 1024
-    self.sequence_length = {'inputs': self.inputs_length, 
+    self.sequence_length = {'inputs': self.inputs_length,
                             'targets': self.outputs_length}
 
     self.partitioner = t5x.partitioning.PjitPartitioner(
-        model_parallel_submesh=(1, 1, 1, 1), num_partitions=1)
+        num_partitions=1)
 
     # Build Codecs and Vocabularies.
     self.spectrogram_config = spectrograms.SpectrogramConfig()
@@ -187,9 +183,10 @@ class InferenceModel(object):
 
   def __call__(self, audio):
     """Infer note sequence from audio samples.
-    
+
     Args:
       audio: 1-d numpy array of audio samples (16kHz) for a single example.
+
     Returns:
       A note_sequence of the transcribed audio.
     """