Compare commits

..

No commits in common. "ae9a6f19ba6b604f5ff4e330fa5e388446e4f758" and "87e63cb770a9924145495ddd088e5f451bf6bced" have entirely different histories.

View file

@ -7,7 +7,6 @@ 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};
@ -37,7 +36,7 @@ async fn main() {
// build our application with a single route
let app = Router::new()
.route("/", post(upload_file))
.layer(axum::extract::DefaultBodyLimit::max(512 * 1024 * 1024));
.layer(axum::extract::DefaultBodyLimit::max(300 * 1024 * 1024));
let upstream_runner = get_upstream_runner();
@ -118,7 +117,7 @@ async fn send_image_to_dd(
.await?;
let body = resp.text().await?;
log::info!("sd body: {}", &body);
log::info!("body: {}", &body);
let json_response: WrappedResponse = serde_json::from_str(&body)?;
log::debug!("called!");
@ -135,8 +134,6 @@ 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))
@ -145,7 +142,7 @@ async fn send_image_to_dd(
.await?;
let body = resp.text().await?;
log::info!("wd14 body: {}", &body);
log::info!("body: {}", &body);
let json_response: WD14Response = serde_json::from_str(&body)?;
// turn WD14Response into WrappedResponse
@ -210,41 +207,6 @@ 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>,
@ -273,8 +235,7 @@ 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(".webm");
|| file_name.ends_with(".gif");
if is_video {
let mut final_tag_set = HashSet::new();
@ -286,27 +247,21 @@ async fn upload_file(
let info = ffprobe::ffprobe(temp_file.path())?;
let stream = info.streams.get(0).unwrap();
log::debug!("stream = {:?}", stream);
log::debug!("format = {:?}", info.format);
let total_frame_count = stream.nb_frames.clone().unwrap().parse::<u32>()?;
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..=1000 => 30,
1001..=1200 => 40,
1201.. => 60,
_ => 63,
301.. => 30,
_ => 33,
} as f64;
let wanted_frame_skip = wanted_frame_skip_seconds * frame_rate;