src/main.rs

@@ -31,20 +31,13 @@ fn get_upstream_runner() -> Runner {
     }
 }
 
-fn get_sd_model() -> String {
-    match std::env::var("SD_MODEL") {
-        Ok(value) => value,
-        Err(_) => "wd14-vit-v2-git".to_string(),
-    }
-}
-
 #[tokio::main]
 async fn main() {
     pretty_env_logger::init();
     // build our application with a single route
     let app = Router::new()
         .route("/", post(upload_file))
-        .layer(axum::extract::DefaultBodyLimit::max(8 * 1024 * 1024 * 1024));
+        .layer(axum::extract::DefaultBodyLimit::max(5 * 1024 * 1024 * 1024));
 
     let upstream_runner = get_upstream_runner();
 
@@ -95,13 +88,7 @@ where
 
 #[derive(Serialize, Deserialize)]
 struct WD14Response {
-    caption: WD14ResponseTagHolder,
-}
-
-#[derive(Serialize, Deserialize)]
-struct WD14ResponseTagHolder {
-    tag: HashMap<String, f32>,
-    rating: HashMap<String, f32>,
+    caption: HashMap<String, f32>,
 }
 
 async fn send_image_to_dd(
@@ -143,9 +130,7 @@ async fn send_image_to_dd(
             let mut map: HashMap<&str, &str> = HashMap::new();
             let file_base64 = general_purpose::STANDARD.encode(file_contents.clone());
 
-            let sd_model = get_sd_model();
-
-            map.insert("model", &sd_model);
+            map.insert("model", "wd14-vit-v2-git");
             map.insert("threshold", options.threshold.as_ref());
             map.insert("image", &file_base64);
 
@@ -165,10 +150,7 @@ async fn send_image_to_dd(
 
             // turn WD14Response into WrappedResponse
             let mut tags = Vec::<String>::new();
-            for ele in json_response.caption.tag {
-                tags.push(ele.0.clone());
-            }
-            for ele in json_response.caption.rating {
+            for ele in json_response.caption {
                 tags.push(ele.0.clone());
             }
 
@@ -339,7 +321,6 @@ async fn upload_file(
         let is_video = file_type.starts_with("video/")
             || file_name.ends_with(".mp4")
             || file_name.ends_with(".gif")
-            || file_name.ends_with(".mkv")
             || file_name.ends_with(".webm");
         if is_video {
             let mut final_tag_set = HashSet::new();