Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| # | |
| from transformers import Wav2Vec2FeatureExtractor | |
| from transformers import AutoModel | |
| import torch | |
| from torch import nn | |
| import torchaudio | |
| import torchaudio.transforms as T | |
| import logging | |
| import json | |
| import os | |
| import re | |
| import pandas as pd | |
| import importlib | |
| modeling_MERT = importlib.import_module("MERT-v1-95M.modeling_MERT") | |
| from Prediction_Head.MTGGenre_head import MLPProberBase | |
| # input cr: https://huggingface.co/spaces/thealphhamerc/audio-to-text/blob/main/app.py | |
| logger = logging.getLogger("MERT-v1-95M-app") | |
| logger.setLevel(logging.INFO) | |
| ch = logging.StreamHandler() | |
| ch.setLevel(logging.INFO) | |
| formatter = logging.Formatter( | |
| "%(asctime)s;%(levelname)s;%(message)s", "%Y-%m-%d %H:%M:%S") | |
| ch.setFormatter(formatter) | |
| logger.addHandler(ch) | |
| inputs = [ | |
| gr.components.Audio(type="filepath", label="Add music audio file"), | |
| gr.inputs.Audio(source="microphone", type="filepath"), | |
| ] | |
| live_inputs = [ | |
| gr.Audio(source="microphone",streaming=True, type="filepath"), | |
| ] | |
| title = "One Model for All Music Understanding Tasks" | |
| description = "An example of using the [MERT-v1-95M](https://huggingface.co/m-a-p/MERT-v1-95M) model as backbone to conduct multiple music understanding tasks with the universal representation. \n Due the hardware limitation of the machine hosting this demo (2 CPU and 16GB RAM) only the first 4 seconds of audio are used!" | |
| # article = "The tasks include EMO, GS, MTGInstrument, MTGGenre, MTGTop50, MTGMood, NSynthI, NSynthP, VocalSetS, VocalSetT. \n\n More models can be referred at the [map organization page](https://huggingface.co/m-a-p)." | |
| with open('./README.md', 'r') as f: | |
| # skip the header | |
| header_count = 0 | |
| for line in f: | |
| if '---' in line: | |
| header_count += 1 | |
| if header_count >= 2: | |
| break | |
| # read the rest conent | |
| article = f.read() | |
| audio_examples = [ | |
| # ["input/example-1.wav"], | |
| # ["input/example-2.wav"], | |
| ] | |
| df_init = pd.DataFrame(columns=['Task', 'Top 1', 'Top 2', 'Top 3', 'Top 4', 'Top 5']) | |
| transcription_df = gr.DataFrame(value=df_init, label="Output Dataframe", row_count=( | |
| 0, "dynamic"), max_rows=30, wrap=True, overflow_row_behaviour='paginate') | |
| # outputs = [gr.components.Textbox()] | |
| outputs = transcription_df | |
| df_init_live = pd.DataFrame(columns=['Task', 'Top 1', 'Top 2', 'Top 3', 'Top 4', 'Top 5']) | |
| transcription_df_live = gr.DataFrame(value=df_init_live, label="Output Dataframe", row_count=( | |
| 0, "dynamic"), max_rows=30, wrap=True, overflow_row_behaviour='paginate') | |
| outputs_live = transcription_df_live | |
| # Load the model and the corresponding preprocessor config | |
| # model = AutoModel.from_pretrained("m-a-p/MERT-v0-public", trust_remote_code=True) | |
| # processor = Wav2Vec2FeatureExtractor.from_pretrained("m-a-p/MERT-v0-public",trust_remote_code=True) | |
| model = modeling_MERT.MERTModel.from_pretrained("./MERT-v1-95M") | |
| processor = Wav2Vec2FeatureExtractor.from_pretrained("./MERT-v1-95M") | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| MERT_BEST_LAYER_IDX = { | |
| 'EMO': 5, | |
| 'GS': 8, | |
| 'GTZAN': 7, | |
| 'MTGGenre': 7, | |
| 'MTGInstrument': 'all', | |
| 'MTGMood': 6, | |
| 'MTGTop50': 6, | |
| 'MTT': 'all', | |
| 'NSynthI': 6, | |
| 'NSynthP': 1, | |
| 'VocalSetS': 2, | |
| 'VocalSetT': 9, | |
| } | |
| MERT_BEST_LAYER_IDX = { | |
| 'EMO': 5, | |
| 'GS': 8, | |
| 'GTZAN': 7, | |
| 'MTGGenre': 7, | |
| 'MTGInstrument': 'all', | |
| 'MTGMood': 6, | |
| 'MTGTop50': 6, | |
| 'MTT': 'all', | |
| 'NSynthI': 6, | |
| 'NSynthP': 1, | |
| 'VocalSetS': 2, | |
| 'VocalSetT': 9, | |
| } | |
| CLASSIFIERS = { | |
| } | |
| ID2CLASS = { | |
| } | |
| TASKS = ['GS', 'MTGInstrument', 'MTGGenre', 'MTGTop50', 'MTGMood', 'NSynthI', 'NSynthP', 'VocalSetS', 'VocalSetT','EMO',] | |
| Regression_TASKS = ['EMO'] | |
| head_dir = './Prediction_Head/best-layer-MERT-v1-95M' | |
| for task in TASKS: | |
| print('loading', task) | |
| with open(os.path.join(head_dir,f'{task}.id2class.json'), 'r') as f: | |
| ID2CLASS[task]=json.load(f) | |
| num_class = len(ID2CLASS[task].keys()) | |
| CLASSIFIERS[task] = MLPProberBase(d=768, layer=MERT_BEST_LAYER_IDX[task], num_outputs=num_class) | |
| CLASSIFIERS[task].load_state_dict(torch.load(f'{head_dir}/{task}.ckpt')['state_dict']) | |
| CLASSIFIERS[task].to(device) | |
| model.to(device) | |
| def model_infernce(inputs): | |
| waveform, sample_rate = torchaudio.load(inputs) | |
| resample_rate = processor.sampling_rate | |
| # make sure the sample_rate aligned | |
| if resample_rate != sample_rate: | |
| # print(f'setting rate from {sample_rate} to {resample_rate}') | |
| resampler = T.Resample(sample_rate, resample_rate) | |
| waveform = resampler(waveform) | |
| #waveform = waveform.view(-1,) # make it (n_sample, ) | |
| waveform = waveform[0][0:4*resample_rate] # cut to 4s samples | |
| model_inputs = processor(waveform, sampling_rate=resample_rate, return_tensors="pt") | |
| model_inputs.to(device) | |
| with torch.no_grad(): | |
| model_outputs = model(**model_inputs, output_hidden_states=True) | |
| # take a look at the output shape, there are 13 layers of representation | |
| # each layer performs differently in different downstream tasks, you should choose empirically | |
| all_layer_hidden_states = torch.stack(model_outputs.hidden_states).squeeze()[1:,:,:].unsqueeze(0) | |
| print(all_layer_hidden_states.shape) # [13 layer, Time steps, 768 feature_dim] | |
| all_layer_hidden_states = all_layer_hidden_states.mean(dim=2) | |
| task_output_texts = "" | |
| df = pd.DataFrame(columns=['Task', 'Top 1', 'Top 2', 'Top 3', 'Top 4', 'Top 5']) | |
| df_objects = [] | |
| for task in TASKS: | |
| num_class = len(ID2CLASS[task].keys()) | |
| if MERT_BEST_LAYER_IDX[task] == 'all': | |
| logits = CLASSIFIERS[task](all_layer_hidden_states) # [1, 87] | |
| else: | |
| logits = CLASSIFIERS[task](all_layer_hidden_states[:, MERT_BEST_LAYER_IDX[task]]) | |
| # print(f'task {task} logits:', logits.shape, 'num class:', num_class) | |
| sorted_idx = torch.argsort(logits, dim = -1, descending=True)[0] # batch =1 | |
| sorted_prob,_ = torch.sort(nn.functional.softmax(logits[0], dim=-1), dim=-1, descending=True) | |
| # print(sorted_prob) | |
| # print(sorted_prob.shape) | |
| top_n_show = 5 if num_class >= 5 else num_class | |
| # task_output_texts = task_output_texts + f"TASK {task} output:\n" + "\n".join([str(ID2CLASS[task][str(sorted_idx[idx].item())])+f', probability: {sorted_prob[idx].item():.2%}' for idx in range(top_n_show)]) + '\n' | |
| # task_output_texts = task_output_texts + '----------------------\n' | |
| row_elements = [task] | |
| for idx in range(top_n_show): | |
| print(ID2CLASS[task]) | |
| # print('id', str(sorted_idx[idx].item())) | |
| output_class_name = str(ID2CLASS[task][str(sorted_idx[idx].item())]) | |
| output_class_name = re.sub(r'^\w+---', '', output_class_name) | |
| output_class_name = re.sub(r'^\w+\/\w+---', '', output_class_name) | |
| # print('output name', output_class_name) | |
| output_prob = f' {sorted_prob[idx].item():.2%}' | |
| row_elements.append(output_class_name+output_prob) | |
| # fill empty elment | |
| for _ in range(5+1 - len(row_elements)): | |
| row_elements.append(' ') | |
| df_objects.append(row_elements) | |
| df = pd.DataFrame(df_objects, columns=['Task', 'Top 1', 'Top 2', 'Top 3', 'Top 4', 'Top 5']) | |
| return df | |
| def convert_audio(inputs, microphone): | |
| if (microphone is not None): | |
| inputs = microphone | |
| df = model_infernce(inputs) | |
| return df | |
| def live_convert_audio(microphone): | |
| if (microphone is not None): | |
| inputs = microphone | |
| df = model_infernce(inputs) | |
| return df | |
| audio_chunked = gr.Interface( | |
| fn=convert_audio, | |
| inputs=inputs, | |
| outputs=outputs, | |
| allow_flagging="never", | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=audio_examples, | |
| ) | |
| live_audio_chunked = gr.Interface( | |
| fn=live_convert_audio, | |
| inputs=live_inputs, | |
| outputs=outputs_live, | |
| allow_flagging="never", | |
| title=title, | |
| description=description, | |
| article=article, | |
| # examples=audio_examples, | |
| live=True, | |
| ) | |
| demo = gr.Blocks() | |
| with demo: | |
| gr.TabbedInterface( | |
| [ | |
| audio_chunked, | |
| live_audio_chunked, | |
| ], | |
| [ | |
| "Audio File or Recording", | |
| "Live Streaming Music" | |
| ] | |
| ) | |
| # demo.queue(concurrency_count=1, max_size=5) | |
| demo.launch(show_api=False) |