feat: add ROCm provider to list (#10)

ssube · Jan 19, 2023 · 3bcd7a8 · 3bcd7a8
1 parent 137a5a8
commit 3bcd7a8
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py
@@ -67,7 +67,10 @@
 platform_providers = {
     'amd': 'DmlExecutionProvider',
     'cpu': 'CPUExecutionProvider',
+    'cuda': 'CUDAExecutionProvider',
+    'directml': 'DmlExecutionProvider',
     'nvidia': 'CUDAExecutionProvider',
+    'rocm': 'ROCmExecutionProvider',
 }
 pipeline_schedulers = {
     'ddim': DDIMScheduler,

diff --git a/gui/src/strings.ts b/gui/src/strings.ts
@@ -18,7 +18,9 @@ export const MODEL_LABELS = {
 export const PLATFORM_LABELS: Record<string, string> = {
   amd: 'AMD GPU',
   cpu: 'CPU',
-  nvidia: 'Nvidia GPU',
+  directml: 'DirectML',
+  nvidia: 'CUDA',
+  rocm: 'ROCm',
 };
 
 export const SCHEDULER_LABELS: Record<string, string> = {