more robust logic to fetch video frame count

src/main.rs

@@ -7,6 +7,7 @@ use axum_macros::debug_handler;
 use base64::{engine::general_purpose, Engine as _};
 use core::panic;
 use ffmpeg_cli::Parameter;
+use ffprobe::{Format, Stream};
 use futures_util::{future::ready, StreamExt};
 use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, HashSet};
@@ -36,7 +37,7 @@ async fn main() {
     // build our application with a single route
     let app = Router::new()
         .route("/", post(upload_file))
-        .layer(axum::extract::DefaultBodyLimit::max(300 * 1024 * 1024));
+        .layer(axum::extract::DefaultBodyLimit::max(512 * 1024 * 1024));
 
     let upstream_runner = get_upstream_runner();
 
@@ -117,7 +118,7 @@ async fn send_image_to_dd(
                 .await?;
 
             let body = resp.text().await?;
-            log::info!("body: {}", &body);
+            log::info!("sd body: {}", &body);
             let json_response: WrappedResponse = serde_json::from_str(&body)?;
 
             log::debug!("called!");
@@ -134,6 +135,8 @@ async fn send_image_to_dd(
             map.insert("image", &file_base64);
 
             let serialized_map = serde_json::to_vec(&map).unwrap();
+            let len = serialized_map.len();
+            log::info!("wd14 request length {} bytes", len);
 
             let resp = reqwest::Client::new()
                 .post(format!("{}/tagger/v1/interrogate", url))
@@ -142,7 +145,7 @@ async fn send_image_to_dd(
                 .await?;
 
             let body = resp.text().await?;
-            log::info!("body: {}", &body);
+            log::info!("wd14 body: {}", &body);
             let json_response: WD14Response = serde_json::from_str(&body)?;
 
             // turn WD14Response into WrappedResponse
@@ -207,6 +210,41 @@ async fn fetch_frame_as_image(
     Ok(())
 }
 
+fn fetch_frame_count_full_decode(path: &std::path::Path) -> anyhow::Result<u64> {
+    let config = ffprobe::ConfigBuilder::new().count_frames(true).build();
+    let new_info = ffprobe::ffprobe_config(config, path)?;
+    let res = new_info
+        .streams
+        .get(0)
+        .unwrap()
+        .nb_read_frames
+        .clone()
+        .unwrap()
+        .parse::<u64>()?;
+    Ok(res)
+}
+
+fn calculate_frame_count(
+    path: &std::path::Path,
+    stream: &Stream,
+    format: &Format,
+    frame_rate: f64,
+) -> anyhow::Result<u64> {
+    Ok(if let Some(parseable_data) = stream.nb_frames.clone() {
+        // if we can get it from the stream metadata, use it
+        parseable_data.parse::<u64>()?
+    } else if let Some(parseable_data) = format.try_get_duration() {
+        // this is a std::time::duration
+        // multiply that by frame rate and we get total frame count (approximate)
+        log::warn!("fetching duration from format metadata...");
+        let seconds = parseable_data?.as_secs_f64();
+        (seconds * frame_rate) as u64
+    } else {
+        log::warn!("file didn't provide frame metadata, calculating it ourselves...");
+        fetch_frame_count_full_decode(path)?
+    })
+}
+
 #[debug_handler]
 async fn upload_file(
     options: Query<Options>,
@@ -235,7 +273,8 @@ async fn upload_file(
         let file_name = maybe_file_name.unwrap();
         let is_video = file_type.starts_with("video/")
             || file_name.ends_with(".mp4")
-            || file_name.ends_with(".gif");
+            || file_name.ends_with(".gif")
+            || file_name.ends_with(".webm");
         if is_video {
             let mut final_tag_set = HashSet::new();
 
@@ -247,21 +286,27 @@ async fn upload_file(
             let info = ffprobe::ffprobe(temp_file.path())?;
             let stream = info.streams.get(0).unwrap();
 
-            let total_frame_count = stream.nb_frames.clone().unwrap().parse::<u32>()?;
-            let frame_rate_str = stream.r_frame_rate.clone();
+            log::debug!("stream = {:?}", stream);
+            log::debug!("format = {:?}", info.format);
 
+            let frame_rate_str = stream.r_frame_rate.clone();
             let parts = frame_rate_str.split("/").into_iter().collect::<Vec<_>>();
             let frame_rate: f64 =
                 parts.get(0).unwrap().parse::<f64>()? / parts.get(1).unwrap().parse::<f64>()?;
 
+            let total_frame_count =
+                calculate_frame_count(temp_file.path(), &stream, &info.format, frame_rate)?;
+
             let total_length_in_seconds = total_frame_count as f64 / frame_rate;
             let wanted_frame_skip_seconds = match total_length_in_seconds as usize {
                 0..=10 => 2,
                 11..=60 => 10,
                 61..=120 => 15,
                 121..=300 => 20,
-                301.. => 30,
-                _ => 33,
+                301..=1000 => 30,
+                1001..=1200 => 40,
+                1201.. => 60,
+                _ => 63,
             } as f64;
             let wanted_frame_skip = wanted_frame_skip_seconds * frame_rate;