import gradio as gr
from transformers import AutoProcessor, AutoModelForTextToWaveform
# Load the processor and model
processor = AutoProcessor.from_pretrained("suno/bark-small")
model = AutoModelForTextToWaveform.from_pretrained("suno/bark-small")

# Define a function to process input and generate waveform
def generate_waveform(text):
    inputs = processor(text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        output = model(**inputs)
    return output

# Input text
input_text = 'Prolongation'

# Generate waveform
waveform_output = generate_waveform(input_text)

# Display waveform
gr.Interface(
    generate_waveform, 
    inputs=gr.inputs.Textbox(default=input_text, label="Enter your text here"),
    outputs=gr.outputs.Audio(type="audio", label="Generated Audio"),
    title="Text to Waveform Generation",
    description="Enter text and generate corresponding waveform"
).launch()