import gradio as gr
import librosa
import pydub
import profanity_check
import openai

def clean_song(file_path):
    # Load the audio file and isolate the acapella.
    audio_file = librosa.load(file_path)
    acapella = librosa.effects.trim(audio_file, top_db=60)

    # Transcribe the acapella with the OpenAI Whisper model.
    transcript = openai.engine("text-davinci-002").generate(
        text="What is the acapella of this song?",
        prompt="Listen to this audio file: " + acapella.to_wav().hex(),
        temperature=0.7,
        max_tokens=200,
    )

    # Find the timestamps of the profane words in the transcript.
    profane_words = profanity_check.get_profanity(transcript)
    timestamps = [
        (m.start(), m.end()) for m in profanity_check.match_all(transcript)
    ]

    # Mute the profane words in the audio file.
    audio = pydub.AudioSegment.from_wav(file_path)
    for start, end in timestamps:
        audio[start:end].set_volume(0)

    # Save the clean audio file.
    audio.export("clean_song.wav", format="wav")

    return "Clean audio file saved as clean_song.wav"

gr.Interface(clean_song, inputs="file", outputs="text").launch()