theluyuan/RWKV-Runner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

webgpu support

Browse Source
This commit is contained in:
josc146 2023-08-16 23:07:58 +08:00
parent 74f1a1c033
commit ef53951a16
18 changed files with 67109 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitattributes vendored
View File

@ -2,6 +2,7 @@ backend-python/rwkv_pip/** linguist-vendored
backend-python/wkv_cuda_utils/** linguist-vendored
backend-python/get-pip.py linguist-vendored
backend-python/convert_model.py linguist-vendored
backend-python/convert_safetensors.py linguist-vendored
backend-python/utils/midi.py linguist-vendored
build/** linguist-vendored
finetune/lora/** linguist-vendored

2
.gitignore vendored
View File

@ -5,6 +5,8 @@ __pycache__
.idea
.vs
*.pth
*.st
*.safetensors
*.bin
/config.json
/cache.json

18
backend-golang/rwkv.go
View File

@ -26,6 +26,13 @@ func (a *App) StartServer(python string, port int, host string, rwkvBeta bool) (
return Cmd(args...)
}
func (a *App) StartWebGPUServer(port int, host string) (string, error) {
args := []string{"./backend-rust/webgpu_server"}
args = append(args, "-a", "0", "-t", "backend-rust/assets/rwkv_vocab_v20230424.json",
"--port", strconv.Itoa(port), "--ip", host)
return Cmd(args...)
}
func (a *App) ConvertModel(python string, modelPath string, strategy string, outPath string) (string, error) {
var err error
if python == "" {
@ -37,6 +44,17 @@ func (a *App) ConvertModel(python string, modelPath string, strategy string, out
return Cmd(python, "./backend-python/convert_model.py", "--in", modelPath, "--out", outPath, "--strategy", strategy)
}
func (a *App) ConvertSafetensors(python string, modelPath string, outPath string) (string, error) {
var err error
if python == "" {
python, err = GetPython()
}
if err != nil {
return "", err
}
return Cmd(python, "./backend-python/convert_safetensors.py", "--input", modelPath, "--output", outPath)
}
func (a *App) ConvertData(python string, input string, outputPrefix string, vocab string) (string, error) {
var err error
if python == "" {

53
backend-python/convert_safetensors.py vendored Normal file
View File

@ -0,0 +1,53 @@
import json
import os
import sys
import copy
import torch
from safetensors.torch import load_file, save_file
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--input", type=str, help="Path to input pth model")
parser.add_argument(
"--output",
type=str,
default="./converted.st",
help="Path to output safetensors model",
)
args = parser.parse_args()
def convert_file(
pt_filename: str,
sf_filename: str,
):
loaded = torch.load(pt_filename, map_location="cpu")
if "state_dict" in loaded:
loaded = loaded["state_dict"]
loaded = {k: v.clone().half() for k, v in loaded.items()}
for k, v in loaded.items():
print(f"{k}\t{v.shape}\t{v.dtype}")
# For tensors to be contiguous
loaded = {k: v.contiguous() for k, v in loaded.items()}
dirname = os.path.dirname(sf_filename)
os.makedirs(dirname, exist_ok=True)
save_file(loaded, sf_filename, metadata={"format": "pt"})
reloaded = load_file(sf_filename)
for k in loaded:
pt_tensor = loaded[k]
sf_tensor = reloaded[k]
if not torch.equal(pt_tensor, sf_tensor):
raise RuntimeError(f"The output tensors do not match for key {k}")
if __name__ == "__main__":
try:
convert_file(args.input, args.output)
print(f"Saved to {args.output}")
except Exception as e:
with open("error.txt", "w") as f:
f.write(str(e))

1
backend-python/dep_check.py
View File

@ -1,3 +1,4 @@
import safetensors
import midi2audio
import mido
import lm_dataformat

BIN
backend-python/requirements.txt
View File

Binary file not shown.

BIN
backend-python/requirements_versions.txt
View File

Binary file not shown.

BIN
backend-python/requirements_without_cyac.txt
View File

Binary file not shown.

66861
backend-rust/assets/rwkv_vocab_v20230424.json Normal file
View File

File diff suppressed because it is too large Load Diff

5
frontend/src/_locales/ja/main.json
View File

@ -241,5 +241,8 @@
"Auto Play At The End": "最後に自動再生",
"No File to save": "保存するファイルがありません",
"File Saved": "ファイルが保存されました",
"Failed to load local sound font, please check if the files exist - assets/sound-font": "ローカルサウンドフォントの読み込みに失敗しました、ファイルが存在するか確認してください - assets/sound-font"
"Failed to load local sound font, please check if the files exist - assets/sound-font": "ローカルサウンドフォントの読み込みに失敗しました、ファイルが存在するか確認してください - assets/sound-font",
"Please convert model to safe tensors format first": "モデルを安全なテンソル形式に変換してください",
"Convert To Safe Tensors Format": "安全なテンソル形式に変換",
"Please change Strategy to WebGPU to use safetensors format": "StrategyをWebGPUに変更して、安全なテンソル形式を使用してください"
}

5
frontend/src/_locales/zh-hans/main.json
View File

@ -241,5 +241,8 @@
"Auto Play At The End": "结束时自动播放",
"No File to save": "无文件可保存",
"File Saved": "文件已保存",
"Failed to load local sound font, please check if the files exist - assets/sound-font": "加载本地音色资源失败，请检查文件是否存在 - assets/sound-font"
"Failed to load local sound font, please check if the files exist - assets/sound-font": "加载本地音色资源失败，请检查文件是否存在 - assets/sound-font",
"Please convert model to safe tensors format first": "请先将模型转换为Safetensors格式",
"Convert To Safe Tensors Format": "转换为Safetensors格式",
"Please change Strategy to WebGPU to use safetensors format": "请将Strategy改为WebGPU以使用safetensors格式"
}

71
frontend/src/components/RunButton.tsx
View File

@ -1,6 +1,12 @@
import React, { FC, MouseEventHandler, ReactElement } from 'react';
import commonStore, { ModelStatus } from '../stores/commonStore';
import { AddToDownloadList, CopyFile, FileExists, StartServer } from '../../wailsjs/go/backend_golang/App';
import {
AddToDownloadList,
CopyFile,
FileExists,
StartServer,
StartWebGPUServer
} from '../../wailsjs/go/backend_golang/App';
import { Button } from '@fluentui/react-components';
import { observer } from 'mobx-react-lite';
import { exit, getStatus, readRoot, switchModel, updateConfig } from '../apis';
@ -39,6 +45,7 @@ export const RunButton: FC<{ onClickRun?: MouseEventHandler, iconMode?: boolean
commonStore.setStatus({ status: ModelStatus.Starting });
const modelConfig = commonStore.getCurrentModelConfig();
const webgpu = modelConfig.modelParameters.device === 'WebGPU';
let modelName = '';
let modelPath = '';
if (modelConfig && modelConfig.modelParameters) {
@ -50,9 +57,32 @@ export const RunButton: FC<{ onClickRun?: MouseEventHandler, iconMode?: boolean
return;
}
const ok = await checkDependencies(navigate);
if (!ok)
return;
if (webgpu) {
if (!['.st', '.safetensors'].some(ext => modelPath.endsWith(ext))) {
const stModelPath = modelPath.replace(/\.pth$/, '.st');
if (await FileExists(stModelPath)) {
modelPath = stModelPath;
} else {
toast(t('Please convert model to safe tensors format first'), { type: 'error' });
commonStore.setStatus({ status: ModelStatus.Offline });
return;
}
}
}
if (!webgpu) {
if (['.st', '.safetensors'].some(ext => modelPath.endsWith(ext))) {
toast(t('Please change Strategy to WebGPU to use safetensors format'), { type: 'error' });
commonStore.setStatus({ status: ModelStatus.Offline });
return;
}
}
if (!webgpu) {
const ok = await checkDependencies(navigate);
if (!ok)
return;
}
const currentModelSource = commonStore.modelSourceList.find(item => item.name === modelName);
@ -85,7 +115,12 @@ export const RunButton: FC<{ onClickRun?: MouseEventHandler, iconMode?: boolean
await exit(1000).catch(() => {
});
StartServer(commonStore.settings.customPythonPath, port, commonStore.settings.host !== '127.0.0.1' ? '0.0.0.0' : '127.0.0.1',
const startServer = webgpu ?
(_: string, port: number, host: string) => StartWebGPUServer(port, host)
: StartServer;
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'
).catch((e) => {
const errMsg = e.message || e;
@ -104,19 +139,23 @@ export const RunButton: FC<{ onClickRun?: MouseEventHandler, iconMode?: boolean
if (r.ok && !loading) {
loading = true;
clearInterval(intervalId);
await getStatus().then(status => {
if (status)
commonStore.setStatus(status);
});
if (!webgpu) {
await getStatus().then(status => {
if (status)
commonStore.setStatus(status);
});
}
commonStore.setStatus({ status: ModelStatus.Loading });
toast(t('Loading Model'), { type: 'info' });
updateConfig({
max_tokens: modelConfig.apiParameters.maxResponseToken,
temperature: modelConfig.apiParameters.temperature,
top_p: modelConfig.apiParameters.topP,
presence_penalty: modelConfig.apiParameters.presencePenalty,
frequency_penalty: modelConfig.apiParameters.frequencyPenalty
});
if (!webgpu) {
updateConfig({
max_tokens: modelConfig.apiParameters.maxResponseToken,
temperature: modelConfig.apiParameters.temperature,
top_p: modelConfig.apiParameters.topP,
presence_penalty: modelConfig.apiParameters.presencePenalty,
frequency_penalty: modelConfig.apiParameters.frequencyPenalty
});
}
const strategy = getStrategy(modelConfig);
let customCudaFile = '';

125
frontend/src/pages/Configs.tsx
View File

@ -13,13 +13,14 @@ import { Page } from '../components/Page';
import { useNavigate } from 'react-router';
import { RunButton } from '../components/RunButton';
import { updateConfig } from '../apis';
import { ConvertModel, FileExists, GetPyError } from '../../wailsjs/go/backend_golang/App';
import { getStrategy } from '../utils';
import { ConvertModel, ConvertSafetensors, FileExists, GetPyError } from '../../wailsjs/go/backend_golang/App';
import { checkDependencies, getStrategy } from '../utils';
import { useTranslation } from 'react-i18next';
import { WindowShow } from '../../wailsjs/runtime/runtime';
import strategyImg from '../assets/images/strategy.jpg';
import strategyZhImg from '../assets/images/strategy_zh.jpg';
import { ResetConfigsButton } from '../components/ResetConfigsButton';
import { useMediaQuery } from 'usehooks-ts';
export type ApiParameters = {
apiPort: number
@ -56,6 +57,7 @@ export const Configs: FC = observer(() => {
const [selectedIndex, setSelectedIndex] = React.useState(commonStore.currentModelConfigIndex);
const [selectedConfig, setSelectedConfig] = React.useState(commonStore.modelConfigs[selectedIndex]);
const [displayStrategyImg, setDisplayStrategyImg] = React.useState(false);
const mq = useMediaQuery('(min-width: 640px)');
const navigate = useNavigate();
const port = selectedConfig.apiParameters.apiPort;
@ -128,7 +130,8 @@ export const Configs: FC = observer(() => {
setSelectedIndex(0);
setSelectedConfig(commonStore.modelConfigs[0]);
}} />
<ToolTipButton desc={t('Save Config')} icon={<Save20Regular />} onClick={onClickSave} />
<ToolTipButton desc={mq ? '' : t('Save Config')} icon={<Save20Regular />} text={mq ? t('Save Config') : ''}
onClick={onClickSave} />
</div>
<div className="flex items-center gap-4">
<Label>{t('Config Name')}</Label>
@ -237,40 +240,84 @@ export const Configs: FC = observer(() => {
}} />
</div>
} />
<ToolTipButton text={t('Convert')}
desc={t('Convert model with these configs. Using a converted model will greatly improve the loading speed, but model parameters of the converted model cannot be modified.')}
onClick={async () => {
if (commonStore.platform === 'darwin') {
toast(t('MacOS is not yet supported for performing this operation, please do it manually.'), { type: 'info' });
return;
} else if (commonStore.platform === 'linux') {
toast(t('Linux is not yet supported for performing this operation, please do it manually.'), { type: 'info' });
return;
}
const modelPath = `${commonStore.settings.customModelsPath}/${selectedConfig.modelParameters.modelName}`;
if (await FileExists(modelPath)) {
const strategy = getStrategy(selectedConfig);
const newModelPath = modelPath + '-' + strategy.replace(/[:> *+]/g, '-');
toast(t('Start Converting'), { autoClose: 1000, type: 'info' });
ConvertModel(commonStore.settings.customPythonPath, modelPath, strategy, newModelPath).then(async () => {
if (!await FileExists(newModelPath + '.pth')) {
toast(t('Convert Failed') + ' - ' + await GetPyError(), { type: 'error' });
} else {
toast(`${t('Convert Success')} - ${newModelPath}`, { type: 'success' });
{
selectedConfig.modelParameters.device !== 'WebGPU' ?
<ToolTipButton text={t('Convert')}
desc={t('Convert model with these configs. Using a converted model will greatly improve the loading speed, but model parameters of the converted model cannot be modified.')}
onClick={async () => {
if (commonStore.platform === 'darwin') {
toast(t('MacOS is not yet supported for performing this operation, please do it manually.') + ' (backend-python/convert_model.py)', { type: 'info' });
return;
} else if (commonStore.platform === 'linux') {
toast(t('Linux is not yet supported for performing this operation, please do it manually.') + ' (backend-python/convert_model.py)', { type: 'info' });
return;
}
}).catch(e => {
const errMsg = e.message || e;
if (errMsg.includes('path contains space'))
toast(`${t('Convert Failed')} - ${t('File Path Cannot Contain Space')}`, { type: 'error' });
else
toast(`${t('Convert Failed')} - ${e.message || e}`, { type: 'error' });
});
setTimeout(WindowShow, 1000);
} else {
toast(`${t('Model Not Found')} - ${modelPath}`, { type: 'error' });
}
}} />
const ok = await checkDependencies(navigate);
if (!ok)
return;
const modelPath = `${commonStore.settings.customModelsPath}/${selectedConfig.modelParameters.modelName}`;
if (await FileExists(modelPath)) {
const strategy = getStrategy(selectedConfig);
const newModelPath = modelPath + '-' + strategy.replace(/[:> *+]/g, '-');
toast(t('Start Converting'), { autoClose: 1000, type: 'info' });
ConvertModel(commonStore.settings.customPythonPath, modelPath, strategy, newModelPath).then(async () => {
if (!await FileExists(newModelPath + '.pth')) {
toast(t('Convert Failed') + ' - ' + await GetPyError(), { type: 'error' });
} else {
toast(`${t('Convert Success')} - ${newModelPath}`, { type: 'success' });
}
}).catch(e => {
const errMsg = e.message || e;
if (errMsg.includes('path contains space'))
toast(`${t('Convert Failed')} - ${t('File Path Cannot Contain Space')}`, { type: 'error' });
else
toast(`${t('Convert Failed')} - ${e.message || e}`, { type: 'error' });
});
setTimeout(WindowShow, 1000);
} else {
toast(`${t('Model Not Found')} - ${modelPath}`, { type: 'error' });
}
}} /> :
<ToolTipButton text={t('Convert To Safe Tensors Format')}
desc=""
onClick={async () => {
if (commonStore.platform === 'darwin') {
toast(t('MacOS is not yet supported for performing this operation, please do it manually.') + ' (backend-python/convert_safetensors.py)', { type: 'info' });
return;
} else if (commonStore.platform === 'linux') {
toast(t('Linux is not yet supported for performing this operation, please do it manually.') + ' (backend-python/convert_safetensors.py)', { type: 'info' });
return;
}
const ok = await checkDependencies(navigate);
if (!ok)
return;
const modelPath = `${commonStore.settings.customModelsPath}/${selectedConfig.modelParameters.modelName}`;
if (await FileExists(modelPath)) {
toast(t('Start Converting'), { autoClose: 1000, type: 'info' });
const newModelPath = modelPath.replace(/\.pth$/, '.st');
ConvertSafetensors(commonStore.settings.customPythonPath, modelPath, newModelPath).then(async () => {
if (!await FileExists(newModelPath)) {
toast(t('Convert Failed') + ' - ' + await GetPyError(), { type: 'error' });
} else {
toast(`${t('Convert Success')} - ${newModelPath}`, { type: 'success' });
}
}).catch(e => {
const errMsg = e.message || e;
if (errMsg.includes('path contains space'))
toast(`${t('Convert Failed')} - ${t('File Path Cannot Contain Space')}`, { type: 'error' });
else
toast(`${t('Convert Failed')} - ${e.message || e}`, { type: 'error' });
});
setTimeout(WindowShow, 1000);
} else {
toast(`${t('Model Not Found')} - ${modelPath}`, { type: 'error' });
}
}} />
}
<Labeled label={t('Strategy')} content={
<Dropdown style={{ minWidth: 0 }} className="grow" value={t(selectedConfig.modelParameters.device)!}
selectedOptions={[selectedConfig.modelParameters.device]}
@ -285,7 +332,7 @@ export const Configs: FC = observer(() => {
{commonStore.platform === 'darwin' && <Option value="MPS">MPS</Option>}
<Option value="CUDA">CUDA</Option>
<Option value="CUDA-Beta">{t('CUDA (Beta, Faster)')!}</Option>
<Option value="WebGPU" disabled>WebGPU</Option>
<Option value="WebGPU">WebGPU</Option>
<Option value="Custom">{t('Custom')!}</Option>
</Dropdown>
} />
@ -305,7 +352,7 @@ export const Configs: FC = observer(() => {
}}>
<Option>fp16</Option>
<Option>int8</Option>
<Option>fp32</Option>
{selectedConfig.modelParameters.device !== 'WebGPU' && <Option>fp32</Option>}
</Dropdown>
} />
}
@ -353,7 +400,7 @@ export const Configs: FC = observer(() => {
}
{selectedConfig.modelParameters.device === 'Custom' && <div />}
{
selectedConfig.modelParameters.device !== 'CPU' && selectedConfig.modelParameters.device !== 'MPS' &&
(selectedConfig.modelParameters.device.includes('CUDA') || selectedConfig.modelParameters.device === 'Custom') &&
<Labeled label={t('Use Custom CUDA kernel to Accelerate')}
desc={t('Enabling this option can greatly improve inference speed and save some VRAM, but there may be compatibility issues. If it fails to start, please turn off this option.')}
content={

9
frontend/src/utils/index.tsx
View File

@ -57,6 +57,8 @@ export async function refreshBuiltInModels(readCache: boolean = false) {
return cache;
}
const modelSuffix = ['.pth', '.st', '.safetensors'];
export async function refreshLocalModels(cache: {
models: ModelSourceItem[]
}, filter: boolean = true, initUnfinishedModels: boolean = false) {
@ -65,7 +67,7 @@ export async function refreshLocalModels(cache: {
await ListDirFiles(commonStore.settings.customModelsPath).then((data) => {
cache.models.push(...data.flatMap(d => {
if (!d.isDir && d.name.endsWith('.pth'))
if (!d.isDir && modelSuffix.some((ext => d.name.endsWith(ext))))
return [{
name: d.name,
size: d.size,
@ -146,7 +148,7 @@ export async function refreshRemoteModels(cache: { models: ModelSourceItem[] })
.catch(() => {
});
cache.models = cache.models.filter((model, index, self) => {
return model.name.endsWith('.pth')
return modelSuffix.some((ext => model.name.endsWith(ext)))
&& index === self.findIndex(
m => m.name === model.name || (m.SHA256 && m.SHA256 === model.SHA256 && m.size === model.size));
});
@ -176,6 +178,9 @@ export const getStrategy = (modelConfig: ModelConfig | undefined = undefined) =>
strategy += 'cpu ';
strategy += params.precision === 'int8' ? 'fp32i8' : 'fp32';
break;
case 'WebGPU':
strategy += params.precision === 'int8' ? 'fp16i8' : 'fp16';
break;
case 'CUDA':
case 'CUDA-Beta':
if (avoidOverflow)

4
frontend/wailsjs/go/backend_golang/App.d.ts generated vendored
View File

@ -10,6 +10,8 @@ export function ConvertData(arg1:string,arg2:string,arg3:string,arg4:string):Pro
export function ConvertModel(arg1:string,arg2:string,arg3:string,arg4:string):Promise<string>;
export function ConvertSafetensors(arg1:string,arg2:string,arg3:string):Promise<string>;
export function CopyFile(arg1:string,arg2:string):Promise<void>;
export function DeleteFile(arg1:string):Promise<void>;
@ -48,6 +50,8 @@ export function SaveJson(arg1:string,arg2:any):Promise<void>;
export function StartServer(arg1:string,arg2:number,arg3:string,arg4:boolean):Promise<string>;
export function StartWebGPUServer(arg1:number,arg2:string):Promise<string>;
export function UpdateApp(arg1:string):Promise<boolean>;
export function WslCommand(arg1:string):Promise<void>;

8
frontend/wailsjs/go/backend_golang/App.js generated
View File

@ -18,6 +18,10 @@ export function ConvertModel(arg1, arg2, arg3, arg4) {
return window['go']['backend_golang']['App']['ConvertModel'](arg1, arg2, arg3, arg4);
}
export function ConvertSafetensors(arg1, arg2, arg3) {
return window['go']['backend_golang']['App']['ConvertSafetensors'](arg1, arg2, arg3);
}
export function CopyFile(arg1, arg2) {
return window['go']['backend_golang']['App']['CopyFile'](arg1, arg2);
}
@ -94,6 +98,10 @@ export function StartServer(arg1, arg2, arg3, arg4) {
return window['go']['backend_golang']['App']['StartServer'](arg1, arg2, arg3, arg4);
}
export function StartWebGPUServer(arg1, arg2) {
return window['go']['backend_golang']['App']['StartWebGPUServer'](arg1, arg2);
}
export function UpdateApp(arg1) {
return window['go']['backend_golang']['App']['UpdateApp'](arg1);
}

4
main.go
View File

@ -49,6 +49,9 @@ var cyacInfo embed.FS
//go:embed backend-python
var py embed.FS
//go:embed backend-rust
var webgpu embed.FS
//go:embed finetune
var finetune embed.FS
@ -63,6 +66,7 @@ func main() {
backend.CopyEmbed(cyac)
backend.CopyEmbed(cyacInfo)
backend.CopyEmbed(py)
backend.CopyEmbed(webgpu)
backend.CopyEmbed(finetune)
backend.CopyEmbed(midi)
backend.CopyEmbed(midiAssets)

1
vendor.yml
View File

@ -2,6 +2,7 @@
- ^backend-python/wkv_cuda_utils/
- ^backend-python/get-pip\.py
- ^backend-python/convert_model\.py
- ^backend-python/convert_safetensors\.py
- ^backend-python/utils/midi\.py
- ^build/
- ^finetune/lora/