import gradio as gr
from transformers import pipeline

# Load the model from Hugging Face
model_pipeline = pipeline("text-generation", model="Mar2Ding/songcomposer_sft")

# Define the function for text generation
def generate_song(prompt):
    result = model_pipeline(prompt, max_length=50)
    return result[0]['generated_text']

# Set up the Gradio interface
iface = gr.Interface(
    fn=generate_song,                     # The function to call
    inputs="text",                         # The input type (text box for user input)
    outputs="text",                        # The output type (generated text)
    live=True                               # Whether to generate live as the user types
)

# Launch the Gradio interface
iface.launch()