Text-to-Audio / app.py
MaxGab's picture
Update app.py
ed02bf5 verified
raw
history blame
1.27 kB
import gradio as gr
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
import torch
import soundfile as sf
model_name = "facebook/musicgen-small"
config = AutoConfig.from_pretrained(model_name)
if not hasattr(config, 'dropout'):
config.dropout = 0.1
if not hasattr(config, 'layerdrop'):
config.layerdrop = 0.1
if not hasattr(config, 'max_position_embeddings'):
config.max_position_embeddings = 2048
if not hasattr(config, 'num_attention_heads'):
config.num_attention_heads = 16
if not hasattr(config, 'num_hidden_layers'):
config.num_hidden_layers = 24
if not hasattr(config, 'num_hidden_layers'):
config.num_hidden_layers = 24
if not hasattr(config, 'hidden_size'):
config.hidden_size = 1024
model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
tokenizer = AutoTokenizer.from_pretrained(model_name)
def text_to_audio(prompt):
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
with torch.no_grad():
output = model.generate(input_ids)
audio_data = output[0].cpu().numpy()
audio_file = "generated_audio.wav"
sf.write(audio_file, audio_data, 22050)
return audio_file
gr.Interface(fn=text_to_audio, inputs="text", outputs="audio").launch()