80 lines
3.2 KiB
Python
80 lines
3.2 KiB
Python
import os
|
|
|
|
import click
|
|
import stable_whisper
|
|
import srt
|
|
import datetime
|
|
|
|
song_file_extensions = set([
|
|
"mp3", # explicitly supported by whisper
|
|
"wav",
|
|
"mp4",
|
|
"mpeg",
|
|
"mpga",
|
|
"m4a",
|
|
"webm",
|
|
"flac", # stable-ts likely uses ffmpeg to convert this for whisper
|
|
])
|
|
|
|
@click.command()
|
|
@click.option('--model', '-m', default='small', help='Which whisper model to use (choices are those of whisper.available_models)')
|
|
@click.option('--language', '-l', required=True)
|
|
@click.argument('directory')
|
|
def main(model, language, directory):
|
|
print(f"Loading {model} model for use with language {language}...")
|
|
model = stable_whisper.load_model(model)
|
|
|
|
for (dirpath, _, filenames) in os.walk(directory):
|
|
if any([ file_is_audio(f) for f in filenames ]):
|
|
# audio files exist here, so let's process this folder
|
|
for f in filenames:
|
|
if get_file_extension(f) in song_file_extensions:
|
|
filename = os.path.join(dirpath, f)
|
|
# get corresponding txt file name
|
|
# (we expect unaligned files to have the same filename as the song, except for the extension)
|
|
unaligned_lyrics_filename = os.path.join(dirpath, os.path.splitext(f)[0] + ".txt")
|
|
# (we also expect existing .lrc files to contain aligned lyrics)
|
|
aligned_lyrics_filename = os.path.join(dirpath, os.path.splitext(f)[0] + ".lrc")
|
|
|
|
# run model to get aligned srt
|
|
if not os.path.exists(unaligned_lyrics_filename):
|
|
print(f"No corresponding unaligned lyric txt exists for {filename}")
|
|
continue
|
|
if os.path.exists(aligned_lyrics_filename):
|
|
print(f"Corresponding aligned lyric txt already exists for {filename}")
|
|
continue
|
|
print(f"Aligning lyrics for {filename}")
|
|
|
|
result: stable_whisper.WhisperResult = model.align(filename, open(unaligned_lyrics_filename).read(), language=language, original_split=True, regroup=False)
|
|
|
|
# turn srt to lrc
|
|
srt = result.to_srt_vtt(filepath=None, word_level=False)
|
|
lrc = srt_to_lrc(srt)
|
|
print(f"Writing aligned lyrics at: {aligned_lyrics_filename}")
|
|
open(aligned_lyrics_filename, "x").write(lrc)
|
|
|
|
def get_file_extension(filename: str) -> str:
|
|
# "asdf.omg.lol" -> [asdf.omg, .lol] -> lol
|
|
return os.path.splitext(filename)[1][1:]
|
|
|
|
def file_is_audio(filename: str) -> bool:
|
|
return get_file_extension(filename) in song_file_extensions
|
|
|
|
def timedelta_to_hhmmssss(td: datetime.timedelta) -> str:
|
|
dt = datetime.datetime(1969, 1, 1) + td
|
|
return dt.strftime('%M:%S.%f')[:-4]
|
|
|
|
def srt_to_lrc(srt_text: str) -> str:
|
|
subs = list(srt.parse(srt_text))
|
|
lines = [f"[{timedelta_to_hhmmssss(s.start)}]{s.content}" for s in subs]
|
|
|
|
# add the end of lyrics marker
|
|
end_time = subs[-1].end
|
|
lines.append(f"[{timedelta_to_hhmmssss(end_time)}]")
|
|
return "\n".join(lines)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|