Upload app.py

app.py

@@ -2,8 +2,8 @@ import subprocess
 import sys
 
 subprocess.check_call([sys.executable,"-m","pip","install", 'torch'])
-subprocess.check_call([sys.executable,"-m","pip","install",'causal-conv1d'])
-subprocess.check_call([sys.executable, "-m", "pip", "install", 'numpy', 'miditok','mamba-ssm','gradio'])
+# subprocess.check_call([sys.executable,"-m","pip","install",'causal-conv1d'])
+subprocess.check_call([sys.executable, "-m", "pip", "install", 'numpy', 'miditok','transformers','gradio'])
 subprocess.check_call(["apt-get", "install", "timidity", "-y"])
 
 # !pip install pretty_midi midi2audio
@@ -25,31 +25,15 @@ subprocess.check_call(["apt-get", "install", "timidity", "-y"])
 # subprocess.check_call(['export', 'LIBRARY_PATH="/usr/local/cuda/lib64/stubs"'])
 import os
 
-os.environ['LC_ALL'] = "en_US.UTF-8"
-os.environ['LD_LIBRARY_PATH'] = "/usr/lib64-nvidia"
-os.environ['LIBRARY_PATH'] = "/usr/local/cuda/lib64/stubs"
-
+from transformers import MambaConfig, MambaForCausalLM
 
 
 import gradio as gr
 import torch
-from mamba_ssm import Mamba
-from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
-from mamba_ssm.models.config_mamba import MambaConfig
 import numpy as np
 
 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-if torch.cuda.is_available():
-    subprocess.check_call(['ldconfig', '/usr/lib64-nvidia'])
-    # !ldconfig /usr/lib64-nvidia
-
-# !wget "https://huggingface.co/krystv/MIDI_Mamba-159M/resolve/main/MIDI_Mamba-159M_1536VS.pt"
-# !wget "https://huggingface.co/krystv/MIDI_Mamba-159M/resolve/main/tokenizer_1536mix_BPE.json"
-if os.path.isfile("MIDI_Mamba-159M_1536VS.pt") == False:
-    subprocess.check_call(['wget', 'https://huggingface.co/krystv/MIDI_Mamba-159M/resolve/main/MIDI_Mamba-159M_1536VS.pt'])
-
-if os.path.isfile("tokenizer_1536mix_BPE.json") == False:
-    subprocess.check_call(['wget', 'https://huggingface.co/krystv/MIDI_Mamba-159M/resolve/main/tokenizer_1536mix_BPE.json'])
+model = MambaForCausalLM.from_pretrained("krystv/MIDI_Mamba-159M")
 
 
 
@@ -61,13 +45,15 @@ mc.vocab_size = 1536
 from miditok import MIDILike,REMI,TokenizerConfig
 from pathlib import Path
 import torch
-
+if os.path.isfile("tokenizer_1536mix_BPE.json") == False:
+    subprocess.check_call(['wget', 'https://huggingface.co/krystv/MIDI_Mamba-159M/resolve/main/tokenizer_1536mix_BPE.json'])
 tokenizer = REMI(params='tokenizer_1536mix_BPE.json')
 
 
 
-mf = MambaLMHeadModel(config=mc,device=device)
-mf.load_state_dict(torch.load("/content/MIDI_Mamba-159M_1536VS.pt",map_location=device))
+if torch.cuda.is_available():
+    model= model.to(device)
+
 
 
 
@@ -86,19 +72,20 @@ custom_html = f"""
 """
 
 
-@spaces.GPU(duration=120)
 def generate(number,top_k_selector,top_p_selector, temperature_selector):
-    input_ids = torch.tensor([[1,]]).to(device)
-    out = mf.generate(
+    input_ids = torch.tensor([[1,]])
+    if torch.cuda.is_available():
+        input_ids = input_ids.to(device)
+    out = model.generate(
     input_ids=input_ids,
+    do_sample=True,
     max_length=int(number),
     temperature=temperature_selector,
     top_p=top_p_selector,
     top_k=top_k_selector,
 
     eos_token_id=2,)
-    m = tokenizer.decode(np.array(out[0].to('cpu')))
-    np.array(out.to('cpu')).shape
+    m = tokenizer.decode(np.array(out[0].cpu()))
     m.dump_midi('output.mid')
     # !timidity output.mid -Ow -o - | ffmpeg -y -f wav -i - output.mp3
     timidity_cmd = ['timidity', 'output.mid', '-Ow', '-o', 'output.wav']
@@ -123,17 +110,10 @@ def generate_and_save(number,top_k_selector,top_p_selector, temperature_selector
 
 
 
-# iface = gr.Interface(fn=generate_and_save,
-#  inputs=[number_selector,top_k_selector,top_p_selector, temperature_selector,generate_button,custom_html_wid],
-#   outputs=[output_box,download_midi_button],
-#    title="MIDI Mamba-159M",submit_btn=False,
-#    clear_btn=False,
-#    description="MIDI Mamba is a Mamba based model trained on MIDI data collected from open internet to train music model.",
-#    allow_flagging=False,)
 
 with gr.Blocks() as b1:
     gr.Markdown("<h1 style='text-align: center;'>MIDI Mamba-159M <h1/> ")
-    gr.Markdown("<h3 style='text-align: center;'>MIDI Mamba is a Mamba based model trained on MIDI data collected from open internet to train music model. <br> by Hemant Kumar<h3/>")
+    gr.Markdown("<h3 style='text-align: center;'>MIDI Mamba is a Mamba based model trained on MIDI data collected from open internet to train music model. Current accuracy is performance is decreased due to model conversion make sure to checkout the <a href='https://colab.research.google.com/github/HemantKArya/MIDI_Mamba/blob/main/MIDI_Mamba.ipynb'>colab notebook<a/> <br> by Hemant Kumar<h3/>")
     with gr.Row():
         with gr.Column():
             number_selector = gr.Number(label="Select Length of output",value=512)