From 2c1a6673f3af002b2cb0d66340e0b0b3f1439d2b Mon Sep 17 00:00:00 2001
From: Nick Mowen <nickmowen213@gmail.com>
Date: Mon, 6 Feb 2023 14:24:27 -0700
Subject: [PATCH] Add ability to GPU device to be automatically detected when
 multiple exist

---
 frigate/ffmpeg_presets.py | 32 +++++++++++++++++++++++++++++---
 frigate/util.py           |  6 +++---
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/frigate/ffmpeg_presets.py b/frigate/ffmpeg_presets.py
index a9bd3059e..93d9e9443 100644
--- a/frigate/ffmpeg_presets.py
+++ b/frigate/ffmpeg_presets.py
@@ -6,8 +6,34 @@ from typing import Any
 
 from frigate.version import VERSION
 from frigate.const import BTBN_PATH
+from frigate.util import vainfo_hwaccel
 
 TIMEOUT_PARAM = "-timeout" if os.path.exists(BTBN_PATH) else "-stimeout"
+GPU_DEVICE_PARAM = None
+
+
+def get_gpu_device() -> str:
+    """Gets the appropriate Intel/AMD GPU device."""
+    if GPU_DEVICE_PARAM:
+        return GPU_DEVICE_PARAM
+
+    devices = filter(lambda d: d.startswith("render"), os.listdir("/dev/dri"))
+
+    if len(devices) < 2:
+        GPU_DEVICE_PARAM = "renderD128"
+        return GPU_DEVICE_PARAM
+    else:
+        for device in devices:
+            check = vainfo_hwaccel(device_name=device)
+
+            if check.returncode == 0:
+                GPU_DEVICE_PARAM = device
+                return GPU_DEVICE_PARAM
+
+    raise ValueError(
+        "Hardware acceleration was requested but no suitable GPU was found."
+    )
+
 
 _user_agent_args = [
     "-user_agent",
@@ -23,7 +49,7 @@ PRESETS_HW_ACCEL_DECODE = {
         "-hwaccel",
         "vaapi",
         "-hwaccel_device",
-        "/dev/dri/renderD128",
+        get_gpu_device(),
         "-hwaccel_output_format",
         "vaapi",
     ],
@@ -31,7 +57,7 @@ PRESETS_HW_ACCEL_DECODE = {
         "-hwaccel",
         "qsv",
         "-qsv_device",
-        "/dev/dri/renderD128",
+        get_gpu_device(),
         "-hwaccel_output_format",
         "qsv",
         "-c:v",
@@ -43,7 +69,7 @@ PRESETS_HW_ACCEL_DECODE = {
         "-hwaccel",
         "qsv",
         "-qsv_device",
-        "/dev/dri/renderD128",
+        get_gpu_device(),
         "-hwaccel_output_format",
         "qsv",
         "-c:v",
diff --git a/frigate/util.py b/frigate/util.py
index 02433b358..3a0e0258e 100755
--- a/frigate/util.py
+++ b/frigate/util.py
@@ -14,7 +14,7 @@ from abc import ABC, abstractmethod
 from collections import Counter
 from collections.abc import Mapping
 from multiprocessing import shared_memory
-from typing import Any, AnyStr, Tuple
+from typing import Any, AnyStr, Optional, Tuple
 
 import cv2
 import numpy as np
@@ -976,9 +976,9 @@ def ffprobe_stream(path: str) -> sp.CompletedProcess:
     return sp.run(ffprobe_cmd, capture_output=True)
 
 
-def vainfo_hwaccel() -> sp.CompletedProcess:
+def vainfo_hwaccel(device_name: Optional[str]) -> sp.CompletedProcess:
     """Run vainfo."""
-    ffprobe_cmd = ["vainfo"]
+    ffprobe_cmd = ["vainfo"] if not device_name else ["vainfo", "--display", "drm", "--device", f"/dev/dri/{device_name}"]
     return sp.run(ffprobe_cmd, capture_output=True)