path: root/.config/mpv/scripts/whisper.lua

--[[
    * 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)