From 7516e47bb7bbff4adc615bea3879e313c6386296 Mon Sep 17 00:00:00 2001 From: Luna Date: Sun, 11 Jun 2023 18:05:48 -0300 Subject: [PATCH 1/4] add log call --- src/main.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/main.rs b/src/main.rs index 72baebd..db459ac 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,6 +21,7 @@ async fn main() { .layer(axum::extract::DefaultBodyLimit::max(100 * 1024 * 1024)); // run it with hyper on localhost:3000 + log::info!("running on 0.0.0.0:6679 to localhost:4443"); axum::Server::bind(&"0.0.0.0:6679".parse().unwrap()) .serve(app.into_make_service()) .await From 49fb8f4763d1384e770d95a5541d4428b8cd1239 Mon Sep 17 00:00:00 2001 From: Luna Date: Sun, 11 Jun 2023 22:59:44 -0300 Subject: [PATCH 2/4] add support for stable-diffusion-webui-wd14-tagger extension --- Cargo.toml | 1 + src/main.rs | 99 +++++++++++++++++++++++++++++++++++++++++++---------- 2 files changed, 82 insertions(+), 18 deletions(-) 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 db459ac..54828bc 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,10 +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 localhost:4443"); + 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 @@ -74,34 +95,76 @@ 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( @@ -139,7 +202,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"); From 1875d3fd48bb9598ae4cb2e51700132714aa1c40 Mon Sep 17 00:00:00 2001 From: Luna Date: Sun, 11 Jun 2023 23:00:03 -0300 Subject: [PATCH 3/4] add gif to allowed extensions --- src/main.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/main.rs b/src/main.rs index 54828bc..a9384a3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -241,7 +241,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(); From 61250ffaa2d6a91fcbdc91ecd741a7821b69b87c Mon Sep 17 00:00:00 2001 From: Luna Date: Sun, 11 Jun 2023 23:00:12 -0300 Subject: [PATCH 4/4] use faster seeking in ffmpeg --- src/main.rs | 38 +++++++++++++++++++++++--------------- 1 file changed, 23 insertions(+), 15 deletions(-) diff --git a/src/main.rs b/src/main.rs index a9384a3..64b0ad9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -171,18 +171,21 @@ 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(); @@ -259,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?