Compare commits

...

4 commits

Author SHA1 Message Date
ae9a6f19ba more robust logic to fetch video frame count 2023-06-12 17:08:28 -03:00
ad268357bd tweak values for larger videos 2023-06-12 16:33:47 -03:00
0f29ba76c2 add webm support 2023-06-12 16:33:44 -03:00
ef64ba589b bump max body size 2023-06-12 16:33:34 -03:00

View file

@ -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;