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] 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);