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

# Load the model and tokenizer from Hugging Face
model_name = "openai/jukebox-1b-lyrics"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, ignore_mismatched_sizes=True)

# Define a function for Gradio that uses the model to generate music
def generate_music(prompt, genre):
    # Tokenize input with the specified genre
    inputs = tokenizer(prompt, genres=genre, return_tensors="pt")
    outputs = model(**inputs)
    return "Generated instrumental music or lyrics"  # Replace with actual music processing

# Set up Gradio interface
interface = gr.Interface(
    fn=generate_music,
    inputs=[
        gr.inputs.Textbox(label="Prompt"),
        gr.inputs.Textbox(label="Genre")  # Adding genre as input
    ],
    outputs="text",
    title="Jukebox Music Generator",
    description="Enter a prompt and genre to generate instrumental music using the Jukebox 1B Lyrics model."
)

# Launch the interface with share=True for a public link
interface.launch(share=True)