Spaces:
Sleeping
Sleeping
import torch | |
import gradio as gr | |
from transformers import AutoTokenizer, AutoModelForTextToWaveform | |
def install_model(namemodel,tokenn,namemodelonxx): | |
model = AutoModelForTextToWaveform.from_pretrained(namemodel,token=tokenn) | |
namemodelonxxx=convert_to_onnx(model,namemodelonxx) | |
return namemodelonxxx | |
def convert_to_onnx(model,namemodelonxx): | |
vocab_size = model.text_encoder.embed_tokens.weight.size(0) | |
example_input = torch.randint(0, vocab_size, (1, 100), dtype=torch.long) | |
torch.onnx.export( | |
model, # The model to be exported | |
example_input, # Example input for the model | |
namemodelonxx, # The filename for the exported ONNX model | |
opset_version=11, # Use an appropriate ONNX opset version | |
input_names=['input'], # Name of the input layer | |
output_names=['output'], # Name of the output layer | |
dynamic_axes={ | |
'input': {0: 'batch_size', 1: 'sequence_length'}, # Dynamic axes for variable-length inputs | |
'output': {0: 'batch_size'} | |
} | |
) | |
return namemodelonxx | |
with gr.Blocks() as demo: | |
with gr.Row(): | |
with gr.Column(): | |
text_n_model=gr.Textbox(label="name model") | |
text_n_token=gr.Textbox(label="token") | |
text_n_onxx=gr.Textbox(label="name model onxx") | |
with gr.Column(): | |
btn=gr.Button("convert") | |
label=gr.Label("return name model onxx") | |
btn.click(install_model,[text_n_model,text_n_token,text_n_onxx],[label]) | |
demo.launch() |