From 9704525fdf811195d5b2c71fbc5ea9202b7365e1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torjus=20H=C3=A5kestad?= <torjus@usit.uio.no>
Date: Tue, 28 May 2024 19:44:52 +0200
Subject: [PATCH] Add cuda for obs-backgroundremoval

---
 home/programs/obs-studio/default.nix | 40 ++++++++++++++++------------
 1 file changed, 23 insertions(+), 17 deletions(-)

diff --git a/home/programs/obs-studio/default.nix b/home/programs/obs-studio/default.nix
index 0cc0c3e..447a6b4 100644
--- a/home/programs/obs-studio/default.nix
+++ b/home/programs/obs-studio/default.nix
@@ -1,25 +1,31 @@
-{ pkgs, ... }:
-# let
-#   obs-backgroundremoval_1_1_10 = (pkgs.obs-studio-plugins.obs-backgroundremoval.overrideAttrs {
-#     version = "1.1.10";
-#     src = pkgs.fetchFromGitHub {
-#       owner = "occ-ai";
-#       repo = "obs-backgroundremoval";
-#       rev = "1.1.10";
-#       hash = "sha256-6OTtmjXS3fiorEPUSjjqrS8u9fVJnnZZoD722W3JQO8=";
-#     };
-#   }).override {
-#     onnxruntime = pkgs.unstable.onnxruntime;
-#   };
-# in
+{ pkgs, lib, osConfig, ... }:
+let
+  withCuda = if (osConfig.system.name == "gunter") then true else false;
+  onnxruntime-gpu = (pkgs.onnxruntime.override {
+    cudaSupport = withCuda;
+  });
+  obs-backgrounremoval-gpu = ((pkgs.obs-studio-plugins.obs-backgroundremoval.override {
+    onnxruntime = onnxruntime-gpu;
+  }).overrideAttrs (old: {
+    version = "1.1.14-beta";
+    src = pkgs.fetchFromGitHub {
+      owner = "occ-ai";
+      repo = "obs-backgroundremoval";
+      rev = "94be8c35fe077be93a6f5ef347a802295a36dddd";
+      hash = "sha256-qnxDNeTWQYiRMqT6jNp8GC8ef6aaAAY+OXAak54dVc8=";
+    };
+    cmakeFlags =
+      if withCuda then
+        (lib.lists.remove "-DDISABLE_ONNXRUNTIME_GPU=ON" old.cmakeFlags)
+      else old.cmakeFlags;
+  }));
+in
 {
   programs.obs-studio = {
     enable = true;
     plugins = with pkgs.obs-studio-plugins; [
-      obs-backgroundremoval
       obs-pipewire-audio-capture
       obs-shaderfilter
-    ];
-    #plugins = [ obs-backgroundremoval_1_1_10 ];
+    ] ++ [ obs-backgrounremoval-gpu ];
   };
 }