diff options
Diffstat (limited to '')
| -rw-r--r-- | .config/mpv/scripts/whisper.lua | 543 |
1 files changed, 543 insertions, 0 deletions
diff --git a/.config/mpv/scripts/whisper.lua b/.config/mpv/scripts/whisper.lua new file mode 100644 index 0000000..424edf8 --- /dev/null +++ b/.config/mpv/scripts/whisper.lua @@ -0,0 +1,543 @@ +--[[ + * whisper.lua + * + * AUTHORS: dyphire,robertgzr + * License: MIT +]] + +local msg = require('mp.msg') +local utils = require('mp.utils') +local options = require('mp.options') + +---- Script Options ---- +local o = { + ffmpeg_path = 'ffmpeg', + model = '~/src/edl-toolbox/whisper-ggml-small.bin', + language = 'auto', + queue = '3', + use_gpu = 'true', + gpu_device = '0', + -- Specify output path, supports absolute and relative paths + -- Special value: "source" saves the subtitle file to the directory + -- where the video file is located + output_path = 'source', + -- Specify how many subtitles are generated before updating + -- to avoid frequent flickering of subtitles + update_interval = 20, + -- Segment duration in seconds + segment_duration = 20, +} + +options.read_options(o, _, function() end) +------------------------ + +o.ffmpeg_path = mp.command_native({ 'expand-path', o.ffmpeg_path }) +o.output_path = mp.command_native({ 'expand-path', o.output_path }) + +local pid = mp.get_property_native('pid') +local temp_path = os.getenv('TEMP') or '/tmp/' + +local subtitles_file +local subtitle_count = 1 +local append_subtitle_count = 1 +local subtitles_written = false +local whisper_running = false +local state = {} +local time_ranges = {} + +local is_windows = package.config:sub(1, 1) == '\\' + +local function is_protocol(path) + return type(path) == 'string' + and (path:find('^%a[%w.+-]-://') ~= nil or path:find('^%a[%w.+-]-:%?') ~= nil) +end + +local function file_exists(path) + if path then + local meta = utils.file_info(path) + return meta and meta.is_file + end + return false +end + +local function is_writable(path) + local file = io.open(path, 'w') + if file then + file:close() + os.remove(path) + return true + end + return false +end + +local function check_and_remove_empty_file(file_path) + if file_exists(file_path) then + local file = io.open(file_path, 'r') + if file then + local content = file:read('*all') + file:close() + if content == '' then + os.remove(file_path) + end + end + end +end + +local function normalize(path) + if normalize_path ~= nil then + if normalize_path then + path = mp.command_native({ 'normalize-path', path }) + else + local directory = mp.get_property('working-directory', '') + path = utils.join_path(directory, path:gsub('^%.[\\/]', '')) + if is_windows then + path = path:gsub('\\', '/') + end + end + return path + end + + normalize_path = false + + local commands = mp.get_property_native('command-list', {}) + for _, command in ipairs(commands) do + if command.name == 'normalize-path' then + normalize_path = true + break + end + end + return normalize(path) +end + +local function format_time(time_str) + local h, m, s, ms = nil, nil, nil, nil + if time_str:match('^%d+:%d+:%d+[%.:]%d+$') then + h, m, s, ms = time_str:match('(%d+):(%d+):(%d+)[%.:](%d+)') + elseif time_str:match('^%d+:%d+[%.:]%d+$') then + h = 0 + m, s, ms = time_str:match('(%d+):(%d+)[%.:](%d+)') + else + return time_str + end + + return string.format('%02d:%02d:%02d,%03d', h, m, s, ms) +end + +local function timestamp_to_seconds(timestamp) + local h, m, s, ms = timestamp:match('(%d+):(%d+):(%d+),(%d+)') + return tonumber(h) * 3600 + tonumber(m) * 60 + tonumber(s) + tonumber(ms) / 1000 +end + +local function seconds_to_timestamp(seconds) + local h = math.floor(seconds / 3600) + local m = math.floor(seconds / 60) % 60 + local s = math.floor(seconds % 60) + local ms = math.floor((seconds - math.floor(seconds)) * 1000) + return string.format('%02d:%02d:%02d,%03d', h, m, s, ms) +end + +local function check_sub(sub_file) + local tracks = mp.get_property_native('track-list') + local _, sub_title = utils.split_path(sub_file) + for _, track in ipairs(tracks) do + local external_filename = track['external-filename'] + local track_title = track['title'] + if external_filename then + _, track_title = utils.split_path(external_filename) + end + + if track['type'] == 'sub' and track_title == sub_title then + return true, track['id'] + end + end + return false, nil +end + +local function append_sub(sub_file, auto) + local sub, id = check_sub(sub_file) + if not sub then + if auto then + mp.commandv('sub-add', sub_file, 'auto') + else + mp.commandv('sub-add', sub_file) + end + else + mp.commandv('sub-reload', id) + end +end + +local function shift_subtitle_timestamps(temp_srt, srt_file, subtitle_count, start_time) + local temp_file = io.open(temp_srt, 'r') + local main_file = io.open(srt_file, 'a') + + if not temp_file or not main_file then + msg.error('Failed to open temporary or main SRT file.') + return subtitle_count + end + + local subtitle_number = subtitle_count + if subtitle_number == 1 then + mp.osd_message('AI subtitles are loaded and updated in real time', 5) + msg.info('AI subtitles are loaded and updated in real time') + end + for line in temp_file:lines() do + if line:match('%d+:%d+:%d+,%d+%D+%d+:%d+:%d+,%d+') then + local start_ts, end_ts = line:match('(%d+:%d+:%d+,%d+)%D+(%d+:%d+:%d+,%d+)') + if start_ts and end_ts then + local start_seconds = timestamp_to_seconds(start_ts) + start_time + local end_seconds = timestamp_to_seconds(end_ts) + start_time + main_file:write(subtitle_number .. '\n') + main_file:write( + seconds_to_timestamp(start_seconds) + .. ' --> ' + .. seconds_to_timestamp(end_seconds) + .. '\n' + ) + subtitle_number = subtitle_number + 1 + end + elseif line ~= '' and not tonumber(line) then + main_file:write(line .. '\n') + end + end + + temp_file:close() + main_file:close() + + os.remove(temp_srt) + + return subtitle_number +end +------------------------ + +local function process_audio_segment(video_path, temp_srt_path, start_time_str, segment_duration) + local whisper_filter = { + 'model=' .. normalize(o.model), + 'language=' .. o.language, + 'use_gpu=' .. o.use_gpu, + 'gpu_device=' .. o.gpu_device, + 'queue=' .. o.queue, + 'format=srt', + 'destination=' .. temp_srt_path, + } + local args = { + o.ffmpeg_path, + '-hide_banner', + '-nostdin', + '-y', + '-loglevel', + 'quiet', + '-i', + video_path, + '-ss', + start_time_str, + '-t', + utils.to_string(segment_duration), + '-map', + string.format('a:%s?', mp.get_property_number('current-tracks/audio/id', 0) - 1), + '-vn', + '-sn', + '-af', + 'whisper=' .. table.concat(whisper_filter, ':'), + '-f', + 'null', + '-', + } + local res = mp.command_native({ + name = 'subprocess', + capture_stdout = true, + capture_stderr = true, + args = args, + }) + + if res and res.status ~= 0 then + msg.error('Processing failed for: ' .. video_path .. '\n' .. res.stdout .. ' -- ' .. res.stderr) + return false + end + + return true +end + +local function process_video_incrementally(video_path, srt_file, segment_duration) + local start_time = 0 + local segment_index = 1 + local file_duration = mp.get_property_number('duration') + + msg.info('hi') + + while true do + if start_time >= file_duration then + break + end + if start_time + segment_duration > file_duration then + segment_duration = file_duration - start_time + end + + local temp_srt_file = utils.join_path(temp_path, 'whisper-' .. pid .. '.srt') + if file_exists(temp_srt_file) then + os.remove(temp_srt_file) + end + + local start_time_str = string.format( + '%02d:%02d:%02d', + math.floor(start_time / 3600), + math.floor(start_time / 60) % 60, + start_time % 60 + ) + msg.verbose( + string.format('Processing segment: %d, Start Time: %s', segment_index, start_time_str) + ) + local success = + process_audio_segment(video_path, temp_srt_file, start_time_str, segment_duration) + if not success or not file_exists(temp_srt_file) then + msg.verbose('Segment processing completed or failed.') + break + end + + if file_exists(temp_srt_file) then + subtitle_count = + shift_subtitle_timestamps(temp_srt_file, srt_file, subtitle_count, start_time) + end + + if file_exists(srt_file) then + append_sub(srt_file) + end + + start_time = start_time + segment_duration + segment_index = segment_index + 1 + end +end + +local function whisper_segment() + local path = mp.get_property('path') + local fname = mp.get_property('filename/no-ext') + if not path or is_protocol(path) then + return + end + if path then + path = normalize(path) + dir = utils.split_path(path) + end + + if o.output_path ~= 'source' then + subtitles_file = utils.join_path(o.output_path, fname .. '.srt') + else + subtitles_file = utils.join_path(dir, fname .. '.srt') + end + + if file_exists(subtitles_file) then + msg.info('Subtitles file already exists: ' .. subtitles_file) + return + end + + if not is_writable(subtitles_file) then + subtitles_file = utils.join_path(temp_path, fname .. '.srt') + end + + mp.osd_message('Subtitle generation in progress', 9) + msg.info('Subtitle generation in progress') + msg.verbose('Subtitle file => ' .. subtitles_file) + + whisper_running = true + process_video_incrementally(path, subtitles_file, o.segment_duration) + whisper_running = false + + if file_exists(subtitles_file) then + mp.osd_message('Subtitles successfully generated', 5) + msg.info('Subtitles successfully generated') + append_sub(subtitles_file) + end +end +------------------------ + +local function adjust_time_range(strat_time, end_time) + for _, range in ipairs(time_ranges) do + if not (end_time <= range.start or strat_time >= range.finish) then + if strat_time >= range.start and end_time <= range.finish then + return nil, nil + end + if strat_time < range.finish and end_time > range.start then + if strat_time < range.finish then + strat_time = range.finish + end + if end_time > range.start then + end_time = range.start + end + end + end + end + return strat_time, end_time +end + +local function get_time_range(strat_time, end_time) + strat_time, end_time = adjust_time_range(strat_time, end_time) + if strat_time and strat_time < end_time then + return true, strat_time, end_time + else + return false + end +end + +local function whisper_cache(current_pos, subtitle_count) + local temp_video_file = utils.join_path(temp_path, 'whisper-' .. pid .. '.mkv') + local srt_file = utils.join_path(temp_path, 'whisper.srt') + local file_duration = mp.get_property_number('duration') + local cache_state = mp.get_property_native('demuxer-cache-state') + local cache_ranges = cache_state and cache_state['seekable-ranges'] or {} + local cache_start = cache_ranges[1] and cache_ranges[1]['start'] or current_pos + local cache_end = cache_ranges[1] and cache_ranges[1]['end'] or current_pos + + if current_pos < cache_start or cache_start < state.pos then + current_pos = cache_start + state.pos = current_pos + end + + if current_pos >= file_duration then + if file_exists(srt_file) then + append_sub(srt_file) + end + return + end + + local valid_range, strat_time, end_time = get_time_range(current_pos, cache_end) + if strat_time and end_time then + current_pos = strat_time + cache_end = end_time + end + + if not valid_range or cache_end <= current_pos then + mp.add_timeout(1, function() + whisper_cache(current_pos, subtitle_count) + end) + return + end + + if subtitle_count == 0 then + mp.osd_message('Subtitle generation in progress (from cache)', 9) + msg.info('Subtitle generation in progress (from cache)') + local files_to_remove = { + temp_srt_file1 = utils.join_path(temp_path, 'whisper.srt'), + temp_srt_file2 = utils.join_path(temp_path, 'whisper-' .. pid .. '.srt'), + } + + for _, file in pairs(files_to_remove) do + if file_exists(file) then + os.remove(file) + end + end + end + + whisper_running = true + mp.commandv('dump-cache', math.ceil(current_pos), math.floor(cache_end), temp_video_file) + local temp_srt = srt_file .. '-1.srt' + local success = + process_audio_segment(temp_video_file, temp_srt, current_pos, (cache_end - current_pos)) + if not success then + msg.verbose('Segment processing completed or failed.') + end + whisper_running = false + + if file_exists(temp_srt) then + subtitle_number = shift_subtitle_timestamps(temp_srt, srt_file, subtitle_number, current_pos) + end + + if file_exists(srt_file) then + subtitle_count = subtitle_count + 1 + append_sub(srt_file) + end + + table.insert(time_ranges, { start = current_pos, finish = cache_end }) + table.sort(time_ranges, function(a, b) + return a.start < b.start + end) + + current_pos = cache_end + + -- Callback + mp.add_timeout(1, function() + whisper_cache(current_pos, subtitle_count) + end) +end +------------------------ + +local function whisper_main() + if whisper_running then + return + end + local path = mp.get_property_native('path') + local cache = mp.get_property_native('cache') + local cache_state = mp.get_property_native('demuxer-cache-state') + local cache_ranges = cache_state and cache_state['seekable-ranges'] or {} + if path and is_protocol(path) or cache == 'auto' and #cache_ranges > 0 then + time_ranges = {} + subtitle_count = 0 + local current_pos = mp.get_property_native('time-pos') + local cache_start = cache_ranges[1]['start'] + state.pos = cache_start or current_pos + whisper_cache(cache_start, subtitle_count) + return + end + whisper_segment() +end + +-- mp.register_event('log-message', function(e) +-- if e.prefix ~= mp.get_script_name() then +-- return +-- end +-- +-- local file = subtitles_file and io.open(subtitles_file, 'a') +-- if file and e.text and e.text ~= '' then +-- local text_pattern = '%[([%d+:]?%d+:%d+%.%d+)%D+([%d+:]?%d+:%d+%.%d+)%]%s*(.*)' +-- local start_time_srt, end_time_srt, subtitle_text = e.text:match(text_pattern) +-- if start_time_srt and end_time_srt and subtitle_text then +-- local start_time = format_time(start_time_srt) +-- local end_time = format_time(end_time_srt) +-- +-- file:write(subtitle_count .. '\n') +-- file:write(start_time .. ' --> ' .. end_time .. '\n') +-- file:write(subtitle_text .. '\n') +-- file:close() +-- +-- subtitle_count = subtitle_count + 1 +-- subtitles_written = true +-- end +-- if subtitle_count % o.update_interval == 1 and subtitles_written then +-- if append_subtitle_count == 1 then +-- mp.osd_message('Subtitles are loaded and updated in real time', 5) +-- msg.info('Subtitles are loaded and updated in real time') +-- end +-- append_sub(subtitles_file) +-- subtitles_written = false +-- append_subtitle_count = append_subtitle_count + 1 +-- end +-- end +-- end) + +mp.add_hook('on_unload', 50, function() + start_index = 1 + in_progress_batches = 0 + time_ranges = nil + progress_cache = nil + collectgarbage() + time_ranges = {} + progress_cache = {} + + temp_path = os.getenv('TEMP') or '/tmp/' + local path = mp.get_property('path') + local dir = utils.split_path(path) + local filename = mp.get_property('filename/no-ext') + local files_to_remove = { + temp_video_file = utils.join_path(temp_path, 'whisper-' .. pid .. '.mkv'), + temp_srt_file1 = utils.join_path(temp_path, 'whisper.srt'), + temp_srt_file2 = utils.join_path(temp_path, 'whisper-' .. pid .. '.srt'), + } + + for _, file in pairs(files_to_remove) do + if file_exists(file) then + os.remove(file) + end + end + + check_and_remove_empty_file(subtitles_file) +end) + +mp.register_script_message('whisper/sub-whisper', whisper_main) |