Cargo.toml

@@ -9,7 +9,6 @@ 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"

src/main.rs

@@ -4,32 +4,13 @@ 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::{HashMap, HashSet};
+use std::collections::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() {
@@ -37,12 +18,9 @@ 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(100 * 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
@@ -95,31 +73,22 @@ async fn test_handler() -> Result<(), AppError> {
     Ok(())
 }
 
-#[derive(Serialize, Deserialize)]
-struct WD14Response {
-    caption: HashMap<String, f32>,
-}
-
 async fn send_image_to_dd(
     file_contents: Vec<u8>,
     file_name: String,
     file_mime_type: &str,
     options: &Options,
 ) -> anyhow::Result<WrappedResponse> {
-    let part = reqwest::multipart::Part::bytes(file_contents.clone())
+    let part = reqwest::multipart::Part::bytes(file_contents)
         .file_name(file_name)
         .mime_str(file_mime_type)
         .unwrap();
     let form = reqwest::multipart::Form::new().part("file", part);
 
-    let upstream_runner = get_upstream_runner();
-
-    match upstream_runner {
-        Runner::DeepDanbooru(url) => {
     log::debug!("calling dd");
 
     let resp = reqwest::Client::new()
-                .post(url)
+        .post("http://localhost:4443")
         .multipart(form)
         .header("authorization", "Bearer 123")
         .query(&[("threshold", options.threshold.clone())])
@@ -132,60 +101,24 @@ async fn send_image_to_dd(
 
     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::<String>::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"))
-        // fast seeking with -ss instead of select filter
+        .input(ffmpeg_cli::File::new(input_path))
-        .input(
+        .output(
-            ffmpeg_cli::File::new(input_path)
+            ffmpeg_cli::File::new(output_path)
-                .option(Parameter::KeyValue("ss", &timeline_index_param)),
+                .option(Parameter::KeyValue("vf", &frame_index_param))
-        )
+                .option(Parameter::KeyValue("vframes", "1")),
-        .output(ffmpeg_cli::File::new(output_path).option(Parameter::KeyValue("vframes", "1")));
+        );
 
     log::debug!("running {:?}", builder);
     let ffmpeg = builder.run().await.unwrap();
@@ -205,7 +138,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");
@@ -244,9 +177,7 @@ 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/")
+        let is_video = file_type.starts_with("video/") || file_name.ends_with(".mp4");
-            || file_name.ends_with(".mp4")
-            || file_name.ends_with(".gif");
         if is_video {
             let mut final_tag_set = HashSet::new();
 
@@ -262,41 +193,36 @@ async fn upload_file(
             let frame_rate_str = stream.r_frame_rate.clone();
 
             let parts = frame_rate_str.split("/").into_iter().collect::<Vec<_>>();
-            let frame_rate: f64 =
+            let frame_rate =
-                parts.get(0).unwrap().parse::<f64>()? / parts.get(1).unwrap().parse::<f64>()?;
+                parts.get(0).unwrap().parse::<u32>()? / parts.get(1).unwrap().parse::<u32>()?;
 
-            let total_length_in_seconds = total_frame_count as f64 / frame_rate;
+            let total_length_in_seconds = total_frame_count / frame_rate;
-            let wanted_frame_skip_seconds = match total_length_in_seconds as usize {
+            let wanted_frame_skip_seconds = match total_length_in_seconds {
                 0..=10 => 2,
                 11..=60 => 10,
                 61..=120 => 15,
-                121..=300 => 20,
+                121.. => 20,
-                301.. => 30,
+            };
-                _ => 33,
+            let wanted_frame_skip = (wanted_frame_skip_seconds * frame_rate).try_into().unwrap();
-            } 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!("frame path: '{}'", &temporary_frame_path);
+            log::info!("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 as usize) {
+            for frame_number in (0..total_frame_count).step_by(wanted_frame_skip) {
                 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?