|
| 1 | +import { createWriteStream, existsSync, rmSync } from 'fs'; |
| 2 | +import { CommandRunner, SubCommand, InquirerService } from 'nest-commander'; |
| 3 | +import { resolve } from 'path'; |
| 4 | +import { HttpService } from '@nestjs/axios'; |
| 5 | +import { Presets, SingleBar } from 'cli-progress'; |
| 6 | +import decompress from 'decompress'; |
| 7 | + |
| 8 | +@SubCommand({ name: 'init', aliases: ['setup'] }) |
| 9 | +export class InitCommand extends CommandRunner { |
| 10 | + CORTEX_RELEASES_URL = 'https://api.github.com/repos/janhq/cortex/releases'; |
| 11 | + |
| 12 | + constructor( |
| 13 | + private readonly httpService: HttpService, |
| 14 | + private readonly inquirerService: InquirerService, |
| 15 | + ) { |
| 16 | + super(); |
| 17 | + } |
| 18 | + |
| 19 | + async run(input: string[], options?: any): Promise<void> { |
| 20 | + options = await this.inquirerService.ask('create-init-questions', options); |
| 21 | + const version = input[0] ?? 'latest'; |
| 22 | + |
| 23 | + await this.download(this.parseEngineFileName(options), version); |
| 24 | + } |
| 25 | + |
| 26 | + download = async ( |
| 27 | + engineFileName: string, |
| 28 | + version: string = 'latest', |
| 29 | + ): Promise<any> => { |
| 30 | + const res = await this.httpService |
| 31 | + .get( |
| 32 | + this.CORTEX_RELEASES_URL + `${version === 'latest' ? '/latest' : ''}`, |
| 33 | + { |
| 34 | + headers: { |
| 35 | + 'X-GitHub-Api-Version': '2022-11-28', |
| 36 | + Accept: 'application/vnd.github+json', |
| 37 | + }, |
| 38 | + }, |
| 39 | + ) |
| 40 | + .toPromise(); |
| 41 | + |
| 42 | + if (!res?.data) { |
| 43 | + console.log('Failed to fetch releases'); |
| 44 | + process.exit(1); |
| 45 | + } |
| 46 | + |
| 47 | + let release = res?.data; |
| 48 | + if (Array.isArray(res?.data)) { |
| 49 | + release = Array(res?.data)[0].find( |
| 50 | + (e) => e.name === version.replace('v', ''), |
| 51 | + ); |
| 52 | + } |
| 53 | + const toDownloadAsset = release.assets.find((s: any) => |
| 54 | + s.name.includes(engineFileName), |
| 55 | + ); |
| 56 | + |
| 57 | + if (!toDownloadAsset) { |
| 58 | + console.log(`Could not find engine file ${engineFileName}`); |
| 59 | + process.exit(1); |
| 60 | + } |
| 61 | + |
| 62 | + console.log(`Downloading engine file ${engineFileName}`); |
| 63 | + const engineDir = resolve(this.rootDir(), 'cortex-cpp'); |
| 64 | + if (existsSync(engineDir)) rmSync(engineDir, { recursive: true }); |
| 65 | + |
| 66 | + const download = await this.httpService |
| 67 | + .get(toDownloadAsset.browser_download_url, { |
| 68 | + responseType: 'stream', |
| 69 | + }) |
| 70 | + .toPromise(); |
| 71 | + if (!download) { |
| 72 | + throw new Error('Failed to download model'); |
| 73 | + } |
| 74 | + |
| 75 | + const destination = resolve(this.rootDir(), toDownloadAsset.name); |
| 76 | + |
| 77 | + await new Promise((resolve, reject) => { |
| 78 | + const writer = createWriteStream(destination); |
| 79 | + let receivedBytes = 0; |
| 80 | + const totalBytes = download.headers['content-length']; |
| 81 | + |
| 82 | + writer.on('finish', () => { |
| 83 | + bar.stop(); |
| 84 | + resolve(true); |
| 85 | + }); |
| 86 | + |
| 87 | + writer.on('error', (error) => { |
| 88 | + bar.stop(); |
| 89 | + reject(error); |
| 90 | + }); |
| 91 | + |
| 92 | + const bar = new SingleBar({}, Presets.shades_classic); |
| 93 | + bar.start(100, 0); |
| 94 | + |
| 95 | + download.data.on('data', (chunk: any) => { |
| 96 | + receivedBytes += chunk.length; |
| 97 | + bar.update(Math.floor((receivedBytes / totalBytes) * 100)); |
| 98 | + }); |
| 99 | + |
| 100 | + download.data.pipe(writer); |
| 101 | + }); |
| 102 | + |
| 103 | + try { |
| 104 | + await decompress( |
| 105 | + resolve(this.rootDir(), destination), |
| 106 | + resolve(this.rootDir()), |
| 107 | + ); |
| 108 | + } catch (e) { |
| 109 | + console.log(e); |
| 110 | + process.exit(1); |
| 111 | + } |
| 112 | + process.exit(0); |
| 113 | + }; |
| 114 | + |
| 115 | + parseEngineFileName = (options: { |
| 116 | + runMode?: 'CPU' | 'GPU'; |
| 117 | + gpuType?: 'Nvidia' | 'Others (Vulkan)'; |
| 118 | + instructions?: 'AVX' | 'AVX2' | 'AVX-512' | undefined; |
| 119 | + cudaVersion?: '11' | '12'; |
| 120 | + }) => { |
| 121 | + const platform = |
| 122 | + process.platform === 'win32' |
| 123 | + ? 'windows' |
| 124 | + : process.platform === 'darwin' |
| 125 | + ? 'mac' |
| 126 | + : process.platform; |
| 127 | + const arch = process.arch === 'arm64' ? process.arch : 'amd64'; |
| 128 | + const cudaVersion = |
| 129 | + options.runMode === 'GPU' |
| 130 | + ? options.gpuType === 'Nvidia' |
| 131 | + ? '-cuda-' + (options.cudaVersion === '11' ? '11.7' : '12.2') |
| 132 | + : '-vulkan' |
| 133 | + : ''; |
| 134 | + const instructions = options.instructions ? `-${options.instructions}` : ''; |
| 135 | + const engineName = `${platform}-${arch}${instructions.toLowerCase()}${cudaVersion}`; |
| 136 | + return `${engineName}.tar.gz`; |
| 137 | + }; |
| 138 | + |
| 139 | + rootDir = () => resolve(__dirname, `../../../`); |
| 140 | +} |
0 commit comments