cache extracted subtitles & deepl translations HEAD main

10 files changed, 184 insertions, 53 deletions

diff --git a/src/app.rs b/src/app.rs
index 49efd49..35a501e 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -1,4 +1,5 @@
 use adw::prelude::*;
+use cached::{DiskCache, IOCached};
 use relm4::{WorkerController, prelude::*};
 
 use crate::{
@@ -13,13 +14,15 @@ use crate::{
     subtitle_view::{SubtitleView, SubtitleViewMsg, SubtitleViewOutput},
     subtitles::{
         MetadataCollection, SUBTITLE_TRACKS, StreamIndex, SubtitleCue, SubtitleTrack,
-        TrackMetadata,
-        extraction::{SubtitleExtractor, SubtitleExtractorMsg, SubtitleExtractorOutput},
+        SubtitleTrackCollection, TrackMetadata,
+        extraction::{
+            ExtractionArgs, SubtitleExtractor, SubtitleExtractorMsg, SubtitleExtractorOutput,
+        },
         state::SubtitleState,
     },
     transcript::{Transcript, TranscriptMsg, TranscriptOutput},
-    translation::{DeeplTranslator, deepl::DeeplTranslatorMsg},
-    util::Tracker,
+    translation::{DeeplTranslator, TRANSLATIONS, deepl::DeeplTranslatorMsg},
+    util::{self, Tracker},
 };
 
 pub struct App {
@@ -34,6 +37,9 @@ pub struct App {
     open_url_dialog: Controller<OpenDialog>,
     subtitle_selection_dialog: Option<Controller<SubtitleSelectionDialog>>,
 
+    subtitle_extraction_args: Option<ExtractionArgs>,
+    subtitle_cache: DiskCache<ExtractionArgs, SubtitleTrackCollection>,
+
     primary_subtitle_state: SubtitleState,
     secondary_subtitle_state: SubtitleState,
 
@@ -153,6 +159,8 @@ impl SimpleComponent for App {
             },
         );
 
+        let subtitle_cache = util::make_cache("subtitles");
+
         let model = Self {
             root: root.clone(),
             player,
@@ -165,6 +173,9 @@ impl SimpleComponent for App {
             open_url_dialog,
             subtitle_selection_dialog: None,
 
+            subtitle_extraction_args: None,
+            subtitle_cache,
+
             primary_subtitle_state: SubtitleState::default(),
             secondary_subtitle_state: SubtitleState::default(),
 
@@ -193,6 +204,14 @@ impl SimpleComponent for App {
             }
             AppMsg::SubtitleExtractionComplete => {
                 log::info!("Subtitle extraction complete");
+                if let Some(ref args) = self.subtitle_extraction_args {
+                    if let Err(e) = self
+                        .subtitle_cache
+                        .cache_set(args.clone(), SUBTITLE_TRACKS.read().clone())
+                    {
+                        log::error!("error caching extracted subtitles: {}", e);
+                    }
+                }
             }
             AppMsg::ApplySubtitleSettings(settings) => {
                 self.primary_subtitle_state
@@ -249,6 +268,8 @@ impl SimpleComponent for App {
                 mut metadata,
                 whisper_stream_index,
             } => {
+                self.reset();
+
                 if let Some(ix) = whisper_stream_index {
                     let audio_metadata = metadata.audio.get(&ix).unwrap();
                     let subs_metadata = TrackMetadata {
@@ -263,17 +284,33 @@ impl SimpleComponent for App {
                     metadata.subtitles.insert(ix, subs_metadata);
                 }
 
-                self.player
-                    .sender()
-                    .send(PlayerMsg::SetUrl(url.clone()))
-                    .unwrap();
-                self.extractor
-                    .sender()
-                    .send(SubtitleExtractorMsg::ExtractFromUrl {
-                        url,
-                        whisper_stream_index,
-                    })
-                    .unwrap();
+                let extraction_args = ExtractionArgs {
+                    url: url.clone(),
+                    whisper_stream_index,
+                };
+
+                match self.subtitle_cache.cache_get(&extraction_args) {
+                    Ok(Some(track_collection)) => {
+                        log::debug!("subtitle cache hit");
+                        *(SUBTITLE_TRACKS.write()) = track_collection;
+                    }
+                    Ok(None) => {
+                        log::debug!("subtitle cache miss");
+                        self.extractor
+                            .sender()
+                            .send(SubtitleExtractorMsg::Extract(extraction_args.clone()))
+                            .unwrap();
+                    }
+                    Err(e) => {
+                        log::error!("error querying subtitle cache: {}", e);
+                        self.extractor
+                            .sender()
+                            .send(SubtitleExtractorMsg::Extract(extraction_args.clone()))
+                            .unwrap();
+                    }
+                }
+
+                self.subtitle_extraction_args = Some(extraction_args);
 
                 let subtitle_selection_dialog = SubtitleSelectionDialog::builder()
                     .launch((self.root.clone().into(), metadata))
@@ -283,12 +320,28 @@ impl SimpleComponent for App {
                         }
                     });
                 self.subtitle_selection_dialog = Some(subtitle_selection_dialog);
+
+                self.player.sender().send(PlayerMsg::SetUrl(url)).unwrap();
             }
         }
     }
 }
 
 impl App {
+    fn reset(&mut self) {
+        SUBTITLE_TRACKS.write().clear();
+        TRANSLATIONS.write().clear();
+
+        self.subtitle_selection_dialog = None;
+        self.subtitle_extraction_args = None;
+        self.primary_subtitle_state = SubtitleState::default();
+        self.secondary_subtitle_state = SubtitleState::default();
+        self.autopaused = false;
+        self.hovering_primary_cue = false;
+
+        // TODO also clear transcript?
+    }
+
     fn update_subtitle_states(&mut self, position: gst::ClockTime) {
         self.update_primary_subtitle_state(position);
         self.update_secondary_subtitle_state(position);
@@ -356,10 +409,10 @@ impl App {
 fn update_subtitle_state(state: &mut SubtitleState, position: gst::ClockTime) {
     if let Some(stream_ix) = state.stream_ix {
         let lock = SUBTITLE_TRACKS.read();
-        let track = lock.get(&stream_ix).unwrap();
-
-        update_last_time_ix(&track.start_times, &mut state.last_started_cue_ix, position);
-        update_last_time_ix(&track.end_times, &mut state.last_ended_cue_ix, position);
+        if let Some(track) = lock.get(&stream_ix) {
+            update_last_time_ix(&track.start_times, &mut state.last_started_cue_ix, position);
+            update_last_time_ix(&track.end_times, &mut state.last_ended_cue_ix, position);
+        }
     }
 }
 
diff --git a/src/open_dialog.rs b/src/open_dialog.rs
index 3b822be..b84ff3b 100644
--- a/src/open_dialog.rs
+++ b/src/open_dialog.rs
@@ -68,6 +68,7 @@ impl Component for OpenDialog {
                 set_child = &adw::NavigationView {
                     add = &adw::NavigationPage {
                         set_title: "Open File or Stream",
+                        set_tag: Some("file_selection"),
 
                         #[wrap(Some)]
                         set_child = &adw::ToolbarView {
@@ -304,6 +305,9 @@ impl OpenDialog {
         self.url.get_mut().clear();
         self.do_whisper_extraction = false;
         self.whisper_stream_index = None;
+        if let Some(ref nav) = self.navigation_view {
+            nav.pop_to_tag("file_selection");
+        }
     }
 
     fn fetch_metadata(&mut self, sender: ComponentSender<Self>) {
diff --git a/src/player.rs b/src/player.rs
index d533f48..0804fed 100644
--- a/src/player.rs
+++ b/src/player.rs
@@ -205,9 +205,12 @@ impl SimpleComponent for Player {
                             }
                         }
                         PlayMessage::Buffering(_) => {
+                            // println!("buffering")
                             // TODO
                         }
-                        _ => {}
+                        msg => {
+                            // println!("msg: {:?}", msg);
+                        }
                     }
 
                     glib::ControlFlow::Continue
diff --git a/src/subtitles/extraction/embedded.rs b/src/subtitles/extraction/embedded.rs
index 920f52b..39698cf 100644
--- a/src/subtitles/extraction/embedded.rs
+++ b/src/subtitles/extraction/embedded.rs
@@ -5,8 +5,7 @@ use anyhow::Context;
 use crate::{subtitles::SubtitleCue, subtitles::extraction::*};
 
 pub fn extract_embedded_subtitles(
-    // stream index to use when storing extracted subtitles, this index already
-    // has to be in TRACKS when this function is called!
+    // stream index to use when storing extracted subtitles
     stream_ix: StreamIndex,
     context: ffmpeg::codec::Context,
     time_base: ffmpeg::Rational,
diff --git a/src/subtitles/extraction/mod.rs b/src/subtitles/extraction/mod.rs
index 5070fdb..6495b62 100644
--- a/src/subtitles/extraction/mod.rs
+++ b/src/subtitles/extraction/mod.rs
@@ -3,22 +3,30 @@ mod embedded;
 /// Synthesis of subtitles from audio using whisper.cpp
 mod whisper;
 
-use std::{collections::BTreeMap, sync::mpsc, thread};
+use std::{collections::BTreeMap, fmt::Display, sync::mpsc, thread};
 
 use ffmpeg::Rational;
 use relm4::{ComponentSender, Worker};
 
-use crate::subtitles::{SUBTITLE_TRACKS, StreamIndex, SubtitleCue};
+use crate::subtitles::{StreamIndex, SubtitleCue};
 
 pub struct SubtitleExtractor {}
 
+#[derive(Debug, Clone)]
+pub struct ExtractionArgs {
+    pub url: String,
+    pub whisper_stream_index: Option<usize>,
+}
+
+impl Display for ExtractionArgs {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{} {:?}", self.url, self.whisper_stream_index)
+    }
+}
+
 #[derive(Debug)]
 pub enum SubtitleExtractorMsg {
-    ExtractFromUrl {
-        url: String,
-        // the index of the audio stream on which to run a whisper transcription
-        whisper_stream_index: Option<usize>,
-    },
+    Extract(ExtractionArgs),
 }
 
 #[derive(Debug)]
@@ -38,10 +46,10 @@ impl Worker for SubtitleExtractor {
 
     fn update(&mut self, msg: SubtitleExtractorMsg, sender: ComponentSender<Self>) {
         match msg {
-            SubtitleExtractorMsg::ExtractFromUrl {
+            SubtitleExtractorMsg::Extract(ExtractionArgs {
                 url,
                 whisper_stream_index: whisper_audio_stream_ix,
-            } => {
+            }) => {
                 self.handle_extract_from_url(url, whisper_audio_stream_ix, sender);
             }
         }
@@ -55,12 +63,8 @@ impl SubtitleExtractor {
         whisper_audio_stream_ix: Option<usize>,
         sender: ComponentSender<Self>,
     ) {
-        // Clear existing tracks
-        SUBTITLE_TRACKS.write().clear();
-
         match self.extract_subtitles(&url, whisper_audio_stream_ix, sender.clone()) {
             Ok(_) => {
-                log::info!("Subtitle extraction completed successfully");
                 sender
                     .output(SubtitleExtractorOutput::ExtractionComplete)
                     .unwrap();
@@ -125,7 +129,8 @@ impl SubtitleExtractor {
         }
 
         // wait for extraction to complete
-        for (_, (_, join_handle)) in subtitle_extractors {
+        for (packet_tx, join_handle) in subtitle_extractors.into_values() {
+            drop(packet_tx);
             join_handle
                 .join()
                 .unwrap()
diff --git a/src/subtitles/extraction/whisper.rs b/src/subtitles/extraction/whisper.rs
index bd6fba7..be4346a 100644
--- a/src/subtitles/extraction/whisper.rs
+++ b/src/subtitles/extraction/whisper.rs
@@ -21,8 +21,7 @@ struct WhisperCue {
 }
 
 pub fn generate_whisper_subtitles(
-    // stream index to use when storing generated subtitles, this index
-    // already has to be in TRACKS when this function is called!
+    // stream index to use when storing generated subtitles
     stream_ix: StreamIndex,
     context: ffmpeg::codec::Context,
     time_base: ffmpeg::Rational,
diff --git a/src/subtitles/mod.rs b/src/subtitles/mod.rs
index acb73dc..de747f1 100644
--- a/src/subtitles/mod.rs
+++ b/src/subtitles/mod.rs
@@ -4,6 +4,7 @@ pub mod state;
 use std::collections::BTreeMap;
 
 use relm4::SharedState;
+use serde::{Deserialize, Serialize};
 
 pub type StreamIndex = usize;
 
@@ -26,7 +27,7 @@ pub struct SubtitleCue {
     pub end_time: gst::ClockTime,
 }
 
-#[derive(Default, Debug, Clone)]
+#[derive(Default, Debug, Clone, Serialize, Deserialize)]
 pub struct SubtitleTrack {
     // SoA of cue text, start timestamp, end timestamp
     pub texts: Vec<String>,
@@ -34,7 +35,9 @@ pub struct SubtitleTrack {
     pub end_times: Vec<gst::ClockTime>,
 }
 
-pub static SUBTITLE_TRACKS: SharedState<BTreeMap<StreamIndex, SubtitleTrack>> = SharedState::new();
+pub type SubtitleTrackCollection = BTreeMap<StreamIndex, SubtitleTrack>;
+
+pub static SUBTITLE_TRACKS: SharedState<SubtitleTrackCollection> = SharedState::new();
 
 impl TrackMetadata {
     pub fn from_ffmpeg_stream(stream: &ffmpeg::Stream) -> Self {
diff --git a/src/translation/deepl.rs b/src/translation/deepl.rs
index f2e84d7..11510d5 100644
--- a/src/translation/deepl.rs
+++ b/src/translation/deepl.rs
@@ -1,6 +1,9 @@
-use std::{collections::BTreeMap, time::Duration};
+use std::{collections::BTreeMap, fmt::Display, time::Duration};
 
+use crate::util;
+use cached::proc_macro::io_cached;
 use deepl::DeepLApi;
+use deepl::ModelType;
 use relm4::prelude::*;
 
 use crate::{
@@ -66,30 +69,28 @@ impl AsyncComponent for DeeplTranslator {
 impl DeeplTranslator {
     async fn do_translate(&mut self, sender: AsyncComponentSender<Self>) {
         if let Some(stream_ix) = self.stream_ix {
-            let deepl = DeepLApi::with(&Settings::default().deepl_api_key()).new();
-
             let next_cue_to_translate = self.next_cues_to_translate.entry(stream_ix).or_insert(0);
 
-            if let Some(cue) = {
+            while let Some(cue) = {
                 SUBTITLE_TRACKS
                     .read()
                     .get(&stream_ix)
-                    .unwrap()
-                    .texts
-                    .get(*next_cue_to_translate)
+                    .and_then(|stream| stream.texts.get(*next_cue_to_translate))
                     .cloned()
             } {
-                match deepl
-                    .translate_text(cue, deepl::Lang::EN)
-                    .model_type(deepl::ModelType::PreferQualityOptimized)
-                    .await
+                match translate_text(TranslateRequest {
+                    input: cue.clone(),
+                    source_lang: None,            // TODO
+                    target_lang: deepl::Lang::EN, // TODO
+                })
+                .await
                 {
-                    Ok(mut resp) => {
+                    Ok(translated) => {
                         TRANSLATIONS
                             .write()
                             .entry(stream_ix)
                             .or_insert(Vec::new())
-                            .push(resp.translations.pop().unwrap().text);
+                            .push(translated);
 
                         *next_cue_to_translate = *next_cue_to_translate + 1;
                     }
@@ -100,7 +101,45 @@ impl DeeplTranslator {
             }
         }
 
+        // check every second for new cues to be translated
         relm4::tokio::time::sleep(Duration::from_secs(1)).await;
         sender.input(DeeplTranslatorMsg::DoTranslate);
     }
 }
+
+#[derive(Clone)]
+struct TranslateRequest {
+    input: String,
+    source_lang: Option<deepl::Lang>,
+    target_lang: deepl::Lang,
+}
+
+impl Display for TranslateRequest {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{:?}-{}:{}",
+            self.source_lang, self.target_lang, self.input
+        )
+    }
+}
+
+#[io_cached(
+    disk = true,
+    create = r#"{ util::make_cache("deepl_translations") }"#,
+    map_error = r#"|err| err"#
+)]
+async fn translate_text(req: TranslateRequest) -> anyhow::Result<String> {
+    let deepl = DeepLApi::with(&Settings::default().deepl_api_key()).new();
+    let mut requester = deepl.translate_text(req.input, req.target_lang);
+    requester.model_type(ModelType::PreferQualityOptimized);
+    if let Some(source_lang) = req.source_lang {
+        requester.source_lang(source_lang);
+    }
+    let translated = requester.await?.translations.pop().unwrap().text;
+
+    // try to respect deepl's rate-limit
+    relm4::tokio::time::sleep(Duration::from_millis(500)).await;
+
+    Ok(translated)
+}
diff --git a/src/util/cache.rs b/src/util/cache.rs
new file mode 100644
index 0000000..6e4672f
--- /dev/null
+++ b/src/util/cache.rs
@@ -0,0 +1,24 @@
+use std::env;
+
+use cached::DiskCache;
+use directories::BaseDirs;
+use serde::{Serialize, de::DeserializeOwned};
+
+pub fn make_cache<K, V>(name: &str) -> DiskCache<K, V>
+where
+    K: ToString,
+    V: Serialize + DeserializeOwned,
+{
+    let dir = match BaseDirs::new() {
+        Some(base_dirs) => base_dirs.cache_dir().join("lleap"),
+        None => env::current_dir()
+            .expect("unable to determine current directory")
+            .join("lleap_cache"),
+    };
+
+    DiskCache::new(name)
+        .set_disk_directory(dir)
+        .set_sync_to_disk_on_cache_change(true)
+        .build()
+        .expect("unable to open disk cache")
+}
diff --git a/src/util/mod.rs b/src/util/mod.rs
index 4d19eff..2098a35 100644
--- a/src/util/mod.rs
+++ b/src/util/mod.rs
@@ -1,3 +1,5 @@
+mod cache;
 mod tracker;
 
+pub use cache::make_cache;
 pub use tracker::Tracker;