From e64fa22de63b92a34b56fc208a45645c8356e3aa Mon Sep 17 00:00:00 2001
From: niten <niten@fudo.org>
Date: Fri, 24 May 2024 15:55:54 -0700
Subject: [PATCH] Add threads & workers options

---
 immich-ml-container.nix | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/immich-ml-container.nix b/immich-ml-container.nix
index fc8db86..bc067e0 100644
--- a/immich-ml-container.nix
+++ b/immich-ml-container.nix
@@ -29,6 +29,18 @@ in {
       description = "Path on which to store service state.";
     };
 
+    workers = mkOption {
+      type = int;
+      description = "Number of workers to run.";
+      default = 3;
+    };
+
+    threads = mkOption {
+      type = int;
+      description = "Number of request threads to run.";
+      default = 12;
+    };
+
     immich-version = mkOption { type = str; };
 
     debug = mkEnableOption "Enable debugging logs.";
@@ -48,7 +60,11 @@ in {
         ports =
           [ "${toString cfg.port}:3003" "${toString cfg.metrics-port}:9090" ];
         extraOptions = [ "--gpus=all" ];
-        environment.IMMICH_METRICS = "true";
+        environment = {
+          IMMICH_METRICS = "true";
+          MACHINE_LEARNING_WORKERS = cfg.workers;
+          MACHINE_LEARNING_REQUEST_THREADS = cfg.threads;
+        };
       };
     };