From 22f79bbdf19595a2d5df9a58c0fe655ea772ea40 Mon Sep 17 00:00:00 2001
From: Luna <git@l4.pm>
Date: Sun, 4 Aug 2024 18:35:21 -0300
Subject: [PATCH 1/4] let model be configurable

---
 src/main.rs | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/main.rs b/src/main.rs
index fa0e5f5..840ce24 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -31,6 +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();
@@ -130,7 +137,9 @@ async fn send_image_to_dd(
             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");
+            let sd_model = get_sd_model();
+
+            map.insert("model", &sd_model);
             map.insert("threshold", options.threshold.as_ref());
             map.insert("image", &file_base64);
 

From ee275f0fa54af1ea5dc3539bf28f3de6417593fb Mon Sep 17 00:00:00 2001
From: Luna <git@l4.pm>
Date: Sun, 4 Aug 2024 18:36:49 -0300
Subject: [PATCH 2/4] bump body limit lol

---
 src/main.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/main.rs b/src/main.rs
index 840ce24..04c05d9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -44,7 +44,7 @@ async fn main() {
     // build our application with a single route
     let app = Router::new()
         .route("/", post(upload_file))
-        .layer(axum::extract::DefaultBodyLimit::max(5 * 1024 * 1024 * 1024));
+        .layer(axum::extract::DefaultBodyLimit::max(8 * 1024 * 1024 * 1024));
 
     let upstream_runner = get_upstream_runner();
 

From 657436f80db4f8a688c91a0f18b45194956babac Mon Sep 17 00:00:00 2001
From: Luna <git@l4.pm>
Date: Sun, 4 Aug 2024 18:36:56 -0300
Subject: [PATCH 3/4] support new api

---
 src/main.rs | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 04c05d9..47f1f72 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -95,7 +95,13 @@ where
 
 #[derive(Serialize, Deserialize)]
 struct WD14Response {
-    caption: HashMap<String, f32>,
+    caption: WD14ResponseTagHolder,
+}
+
+#[derive(Serialize, Deserialize)]
+struct WD14ResponseTagHolder {
+    tag: HashMap<String, f32>,
+    rating: HashMap<String, f32>,
 }
 
 async fn send_image_to_dd(
@@ -159,7 +165,10 @@ async fn send_image_to_dd(
 
             // turn WD14Response into WrappedResponse
             let mut tags = Vec::<String>::new();
-            for ele in json_response.caption {
+            for ele in json_response.caption.tag {
+                tags.push(ele.0.clone());
+            }
+            for ele in json_response.caption.rating {
                 tags.push(ele.0.clone());
             }
 

From 6d917bdf1f3c60b83e082850994978e67a731178 Mon Sep 17 00:00:00 2001
From: Luna <git@l4.pm>
Date: Sun, 4 Aug 2024 18:37:00 -0300
Subject: [PATCH 4/4] allow mkv

---
 src/main.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/main.rs b/src/main.rs
index 47f1f72..2d8649c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -339,6 +339,7 @@ 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();