From 71abd357a45b3b541007f155a96b7a038aeb41b3 Mon Sep 17 00:00:00 2001
From: josc146 <josStorer@outlook.com>
Date: Tue, 3 Oct 2023 13:50:58 +0800
Subject: [PATCH] update startup

---
 frontend/src/components/RunButton.tsx | 41 +++++++++++++++------------
 frontend/src/utils/index.tsx          | 13 +++++----
 2 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/frontend/src/components/RunButton.tsx b/frontend/src/components/RunButton.tsx
index 2f43faa..9903460 100644
--- a/frontend/src/components/RunButton.tsx
+++ b/frontend/src/components/RunButton.tsx
@@ -11,7 +11,7 @@ import { Button } from '@fluentui/react-components';
 import { observer } from 'mobx-react-lite';
 import { exit, getStatus, readRoot, switchModel, updateConfig } from '../apis';
 import { toast } from 'react-toastify';
-import { checkDependencies, getStrategy, getSupportedCustomCudaFile, toastWithButton } from '../utils';
+import { checkDependencies, getStrategy, toastWithButton } from '../utils';
 import { useTranslation } from 'react-i18next';
 import { ToolTipButton } from './ToolTipButton';
 import { Play16Regular, Stop16Regular } from '@fluentui/react-icons';
@@ -119,9 +119,10 @@ export const RunButton: FC<{ onClickRun?: MouseEventHandler, iconMode?: boolean
       const startServer = webgpu ?
         (_: string, port: number, host: string) => StartWebGPUServer(port, host)
         : StartServer;
+      const isUsingCudaBeta = modelConfig.modelParameters.device === 'CUDA-Beta';
 
       startServer(commonStore.settings.customPythonPath, port, commonStore.settings.host !== '127.0.0.1' ? '0.0.0.0' : '127.0.0.1',
-        modelConfig.modelParameters.device === 'CUDA-Beta'
+        isUsingCudaBeta
       ).catch((e) => {
         const errMsg = e.message || e;
         if (errMsg.includes('path contains space'))
@@ -162,22 +163,26 @@ export const RunButton: FC<{ onClickRun?: MouseEventHandler, iconMode?: boolean
             if ((modelConfig.modelParameters.device.includes('CUDA') || modelConfig.modelParameters.device === 'Custom')
               && modelConfig.modelParameters.useCustomCuda && !strategy.includes('fp32')) {
               if (commonStore.platform === 'windows') {
-                customCudaFile = getSupportedCustomCudaFile();
-                if (customCudaFile) {
-                  FileExists('./py310/Lib/site-packages/rwkv/model.py').then((exist) => {
-                    // defensive measure. As Python has already been launched, will only take effect the next time it runs.
-                    if (!exist) CopyFile('./backend-python/wkv_cuda_utils/wkv_cuda_model.py', './py310/Lib/site-packages/rwkv/model.py');
-                  });
-                  await CopyFile(customCudaFile, './py310/Lib/site-packages/rwkv/wkv_cuda.pyd').catch(() => {
-                    FileExists('./py310/Lib/site-packages/rwkv/wkv_cuda.pyd').then((exist) => {
-                      if (!exist) {
-                        customCudaFile = '';
-                        toast(t('Failed to copy custom cuda file'), { type: 'error' });
-                      }
-                    });
-                  });
-                } else
-                  toast(t('Supported custom cuda file not found'), { type: 'warning' });
+                // this part is currently unused because there's no longer a need to use different kernels for different GPUs, but it might still be needed in the future
+                //
+                // customCudaFile = getSupportedCustomCudaFile(isUsingCudaBeta);
+                // if (customCudaFile) {
+                //   let kernelTargetPath: string;
+                //   if (isUsingCudaBeta)
+                //     kernelTargetPath = './backend-python/rwkv_pip/beta/wkv_cuda.pyd';
+                //   else
+                //     kernelTargetPath = './backend-python/rwkv_pip/wkv_cuda.pyd';
+                //   await CopyFile(customCudaFile, kernelTargetPath).catch(() => {
+                //     FileExists(kernelTargetPath).then((exist) => {
+                //       if (!exist) {
+                //         customCudaFile = '';
+                //         toast(t('Failed to copy custom cuda file'), { type: 'error' });
+                //       }
+                //     });
+                //   });
+                // } else
+                //   toast(t('Supported custom cuda file not found'), { type: 'warning' });
+                customCudaFile = 'any';
               } else {
                 customCudaFile = 'any';
               }
diff --git a/frontend/src/utils/index.tsx b/frontend/src/utils/index.tsx
index 65a7c2e..63dfe20 100644
--- a/frontend/src/utils/index.tsx
+++ b/frontend/src/utils/index.tsx
@@ -1,6 +1,5 @@
 import {
   AddToDownloadList,
-  CopyFile,
   DeleteFile,
   DepCheck,
   InstallPyDep,
@@ -402,8 +401,6 @@ export const checkDependencies = async (navigate: NavigateFunction) => {
       return false;
     }
     commonStore.setDepComplete(true);
-    if (commonStore.platform === 'windows')
-      CopyFile('./backend-python/wkv_cuda_utils/wkv_cuda_model.py', './py310/Lib/site-packages/rwkv/model.py');
   }
   return true;
 };
@@ -428,12 +425,16 @@ export function toastWithButton(text: string, buttonText: string, onClickButton:
   return id;
 }
 
-export function getSupportedCustomCudaFile() {
+export function getSupportedCustomCudaFile(isBeta: boolean) {
   if ([' 10', ' 16', ' 20', ' 30', 'MX', 'Tesla P', 'Quadro P', 'NVIDIA P', 'TITAN X', 'TITAN RTX', 'RTX A',
     'Quadro RTX 4000', 'Quadro RTX 5000', 'Tesla T4', 'NVIDIA A10', 'NVIDIA A40'].some(v => commonStore.status.device_name.includes(v)))
-    return './backend-python/wkv_cuda_utils/wkv_cuda10_30.pyd';
+    return isBeta ?
+      './backend-python/wkv_cuda_utils/beta/wkv_cuda10_30.pyd' :
+      './backend-python/wkv_cuda_utils/wkv_cuda10_30.pyd';
   else if ([' 40', 'RTX 5000 Ada', 'RTX 6000 Ada', 'RTX TITAN Ada', 'NVIDIA L40'].some(v => commonStore.status.device_name.includes(v)))
-    return './backend-python/wkv_cuda_utils/wkv_cuda40.pyd';
+    return isBeta ?
+      './backend-python/wkv_cuda_utils/beta/wkv_cuda40.pyd' :
+      './backend-python/wkv_cuda_utils/wkv_cuda40.pyd';
   else
     return '';
 }
\ No newline at end of file