diff --git a/Cargo.toml b/Cargo.toml index 74f6016..de1d817 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ edition = "2021" anyhow = "1.0.70" axum = { version = "0.6.12", features = ["tokio", "multipart"] } axum-macros = "0.3.7" +base64 = "0.21.2" ffmpeg-cli = "0.1.0" ffprobe = "0.3.3" filepath = "0.1.2" diff --git a/src/main.rs b/src/main.rs index 72baebd..64b0ad9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,13 +4,32 @@ use axum::response::{IntoResponse, Response}; use axum::routing::post; use axum::{Json, Router}; use axum_macros::debug_handler; +use base64::{engine::general_purpose, Engine as _}; +use core::panic; use ffmpeg_cli::Parameter; use futures_util::{future::ready, StreamExt}; use serde::{Deserialize, Serialize}; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::io::Read; use std::io::Write; use std::process::Stdio; +use std::rc::Rc; + +#[derive(Debug)] +enum Runner { + DeepDanbooru(String), + StableDiffusionWebUI(String), +} + +fn get_upstream_runner() -> Runner { + match std::env::var("DD_ADDRESS") { + Ok(value) => Runner::DeepDanbooru(value), + Err(_) => match std::env::var("SD_ADDRESS") { + Ok(value) => Runner::StableDiffusionWebUI(value), + Err(_) => panic!("shit no addr"), + }, + } +} #[tokio::main] async fn main() { @@ -18,9 +37,12 @@ async fn main() { // build our application with a single route let app = Router::new() .route("/", post(upload_file)) - .layer(axum::extract::DefaultBodyLimit::max(100 * 1024 * 1024)); + .layer(axum::extract::DefaultBodyLimit::max(300 * 1024 * 1024)); + + let upstream_runner = get_upstream_runner(); // run it with hyper on localhost:3000 + log::info!("running on 0.0.0.0:6679 to {:?}", upstream_runner); axum::Server::bind(&"0.0.0.0:6679".parse().unwrap()) .serve(app.into_make_service()) .await @@ -73,52 +95,97 @@ async fn test_handler() -> Result<(), AppError> { Ok(()) } +#[derive(Serialize, Deserialize)] +struct WD14Response { + caption: HashMap, +} + async fn send_image_to_dd( file_contents: Vec, file_name: String, file_mime_type: &str, options: &Options, ) -> anyhow::Result { - let part = reqwest::multipart::Part::bytes(file_contents) + let part = reqwest::multipart::Part::bytes(file_contents.clone()) .file_name(file_name) .mime_str(file_mime_type) .unwrap(); let form = reqwest::multipart::Form::new().part("file", part); - log::debug!("calling dd"); + let upstream_runner = get_upstream_runner(); - let resp = reqwest::Client::new() - .post("http://localhost:4443") - .multipart(form) - .header("authorization", "Bearer 123") - .query(&[("threshold", options.threshold.clone())]) - .send() - .await?; + match upstream_runner { + Runner::DeepDanbooru(url) => { + log::debug!("calling dd"); - let body = resp.text().await?; - log::info!("body: {}", &body); - let json_response: WrappedResponse = serde_json::from_str(&body)?; + let resp = reqwest::Client::new() + .post(url) + .multipart(form) + .header("authorization", "Bearer 123") + .query(&[("threshold", options.threshold.clone())]) + .send() + .await?; - log::debug!("called!"); - Ok(json_response) + let body = resp.text().await?; + log::info!("body: {}", &body); + let json_response: WrappedResponse = serde_json::from_str(&body)?; + + log::debug!("called!"); + Ok(json_response) + } + Runner::StableDiffusionWebUI(url) => { + log::debug!("calling sd"); + + let mut map: HashMap<&str, &str> = HashMap::new(); + let file_base64 = general_purpose::STANDARD.encode(file_contents.clone()); + + map.insert("model", "wd14-vit-v2-git"); + map.insert("threshold", options.threshold.as_ref()); + map.insert("image", &file_base64); + + let serialized_map = serde_json::to_vec(&map).unwrap(); + + let resp = reqwest::Client::new() + .post(format!("{}/tagger/v1/interrogate", url)) + .body(serialized_map) + .send() + .await?; + + let body = resp.text().await?; + log::info!("body: {}", &body); + let json_response: WD14Response = serde_json::from_str(&body)?; + + // turn WD14Response into WrappedResponse + let mut tags = Vec::::new(); + for ele in json_response.caption { + tags.push(ele.0.clone()); + } + + log::debug!("called!"); + Ok(WrappedResponse::Tags(tags)) + } + } } async fn fetch_frame_as_image( input_path: &str, output_path: &str, frame_index: usize, + frame_rate: f64, // X/1sec ) -> anyhow::Result<()> { let frame_index_param = format!("select=eq(n\\,{})", frame_index); + let timeline_index: f64 = frame_index as f64 / frame_rate; + let timeline_index_param = format!("{:.5}", timeline_index); let builder = ffmpeg_cli::FfmpegBuilder::new() .stderr(Stdio::piped()) .option(Parameter::Single("nostdin")) .option(Parameter::Single("y")) - .input(ffmpeg_cli::File::new(input_path)) - .output( - ffmpeg_cli::File::new(output_path) - .option(Parameter::KeyValue("vf", &frame_index_param)) - .option(Parameter::KeyValue("vframes", "1")), - ); + // fast seeking with -ss instead of select filter + .input( + ffmpeg_cli::File::new(input_path) + .option(Parameter::KeyValue("ss", &timeline_index_param)), + ) + .output(ffmpeg_cli::File::new(output_path).option(Parameter::KeyValue("vframes", "1"))); log::debug!("running {:?}", builder); let ffmpeg = builder.run().await.unwrap(); @@ -138,7 +205,7 @@ async fn fetch_frame_as_image( }) .await; - log::debug!("run"); + log::debug!("run!"); let output = ffmpeg.process.wait_with_output().unwrap(); log::debug!("out"); @@ -177,7 +244,9 @@ async fn upload_file( if let Some(file_contents) = maybe_file_contents { let file_type = maybe_file_type.unwrap(); let file_name = maybe_file_name.unwrap(); - let is_video = file_type.starts_with("video/") || file_name.ends_with(".mp4"); + let is_video = file_type.starts_with("video/") + || file_name.ends_with(".mp4") + || file_name.ends_with(".gif"); if is_video { let mut final_tag_set = HashSet::new(); @@ -193,36 +262,41 @@ async fn upload_file( let frame_rate_str = stream.r_frame_rate.clone(); let parts = frame_rate_str.split("/").into_iter().collect::>(); - let frame_rate = - parts.get(0).unwrap().parse::()? / parts.get(1).unwrap().parse::()?; + let frame_rate: f64 = + parts.get(0).unwrap().parse::()? / parts.get(1).unwrap().parse::()?; - let total_length_in_seconds = total_frame_count / frame_rate; - let wanted_frame_skip_seconds = match total_length_in_seconds { + 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.. => 20, - }; - let wanted_frame_skip = (wanted_frame_skip_seconds * frame_rate).try_into().unwrap(); + 121..=300 => 20, + 301.. => 30, + _ => 33, + } as f64; + let wanted_frame_skip = wanted_frame_skip_seconds * frame_rate; let temporary_frame_dir = tempfile::tempdir()?; let temporary_frame_path = format!("{}/frame.png", temporary_frame_dir.path().to_string_lossy()); - log::info!("path: '{}'", &temporary_frame_path); + log::info!("frame path: '{}'", &temporary_frame_path); log::info!("wanted_frame_skip: {}", &wanted_frame_skip_seconds); - for frame_number in (0..total_frame_count).step_by(wanted_frame_skip) { + for frame_number in (0..total_frame_count).step_by(wanted_frame_skip as usize) { log::info!("extracting frame {:?}", frame_number); fetch_frame_as_image( temp_file.path().to_str().unwrap(), &temporary_frame_path, frame_number.try_into().unwrap(), + frame_rate, ) .await?; + log::info!("extracted frame {:?}", frame_number); let mut actual_frame_file = std::fs::File::open(&temporary_frame_path)?; let mut frame_data = vec![]; actual_frame_file.read_to_end(&mut frame_data)?; + log::info!("sending frame {:?}", frame_number); let tags_from_frame = if let WrappedResponse::Tags(tags_from_frame) = send_image_to_dd(frame_data, "amongus.png".to_string(), "image/png", &options) .await?