import gradio as gr
from transformers import AutoModelWithLMHead, AutoTokenizer
import torch

# Load the model and tokenizer
model_name = "Gertie01/MusicLM"
model = AutoModelWithLMHead.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define the function to generate music
def generate_music(text):
    input_ids = tokenizer.encode(text, return_tensors="pt")
    output = model.generate(input_ids, max_length=100)
    generated_music = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_music

# Create the Gradio interface
input_text = gr.inputs.Textbox(lines=5, label="Input Text")
output_text = gr.outputs.Textbox(label="Generated Music")

title = "Music Generation with GPT-3.5"
description = "Enter a text prompt and the model will generate music based on it."

gr.Interface(fn=generate_music, inputs=input_text, outputs=output_text, title=title, description=description).launch()