128 lines
3.5 KiB
TypeScript
128 lines
3.5 KiB
TypeScript
import { invoke } from '@tauri-apps/api/core';
|
|
import { listen } from '@tauri-apps/api/event';
|
|
import { BaseExtension, events } from '@janhq/core';
|
|
|
|
export enum Settings {
|
|
hfToken = 'hf-token',
|
|
}
|
|
|
|
type DownloadEvent = {
|
|
task_id: string
|
|
total_size: number
|
|
downloaded_size: number
|
|
download_type: string
|
|
event_type: string
|
|
}
|
|
|
|
export default class DownloadManager extends BaseExtension {
|
|
hf_token?: string
|
|
|
|
async onLoad() {
|
|
this.registerSettings(SETTINGS)
|
|
this.hf_token = await this.getSetting<string>(Settings.hfToken, undefined)
|
|
}
|
|
|
|
async onUnload() { }
|
|
|
|
async downloadFile(url: string, path: string, taskId: string) {
|
|
// relay tauri events to Jan events
|
|
const unlisten = await listen<DownloadEvent>('download', (event) => {
|
|
let payload = event.payload
|
|
let eventName = {
|
|
Updated: 'onFileDownloadUpdate',
|
|
Error: 'onFileDownloadError',
|
|
Success: 'onFileDownloadSuccess',
|
|
Stopped: 'onFileDownloadStopped',
|
|
Started: 'onFileDownloadStarted',
|
|
}[payload.event_type]
|
|
|
|
// remove this once event system is back in web-app
|
|
console.log(taskId, payload.event_type, payload.downloaded_size / payload.total_size)
|
|
|
|
events.emit(eventName, {
|
|
modelId: taskId,
|
|
percent: payload.downloaded_size / payload.total_size,
|
|
size: {
|
|
transferred: payload.downloaded_size,
|
|
total: payload.total_size,
|
|
},
|
|
downloadType: payload.download_type,
|
|
})
|
|
})
|
|
|
|
try {
|
|
await invoke<void>(
|
|
"download_file",
|
|
{ url, path, taskId, headers: this._getHeaders() },
|
|
)
|
|
} catch (error) {
|
|
console.error("Error downloading file:", error)
|
|
events.emit('onFileDownloadError', {
|
|
modelId: url,
|
|
downloadType: 'Model',
|
|
})
|
|
throw error
|
|
} finally {
|
|
unlisten()
|
|
}
|
|
}
|
|
|
|
async downloadHfRepo(modelId: string, saveDir: string, taskId: string, branch?: string) {
|
|
// relay tauri events to Jan events
|
|
const unlisten = await listen<DownloadEvent>('download', (event) => {
|
|
let payload = event.payload
|
|
let eventName = {
|
|
Updated: 'onFileDownloadUpdate',
|
|
Error: 'onFileDownloadError',
|
|
Success: 'onFileDownloadSuccess',
|
|
Stopped: 'onFileDownloadStopped',
|
|
Started: 'onFileDownloadStarted',
|
|
}[payload.event_type]
|
|
|
|
// remove this once event system is back in web-app
|
|
console.log(taskId, payload.event_type, payload.downloaded_size / payload.total_size)
|
|
|
|
events.emit(eventName, {
|
|
modelId: taskId,
|
|
percent: payload.downloaded_size / payload.total_size,
|
|
size: {
|
|
transferred: payload.downloaded_size,
|
|
total: payload.total_size,
|
|
},
|
|
downloadType: payload.download_type,
|
|
})
|
|
})
|
|
|
|
try {
|
|
await invoke<void>(
|
|
"download_hf_repo",
|
|
{ modelId, saveDir, taskId, branch, headers: this._getHeaders() },
|
|
)
|
|
} catch (error) {
|
|
console.error("Error downloading file:", error)
|
|
events.emit('onFileDownloadError', {
|
|
modelId: modelId,
|
|
downloadType: 'Model',
|
|
})
|
|
throw error
|
|
} finally {
|
|
unlisten()
|
|
}
|
|
}
|
|
|
|
async cancelDownload(taskId: string) {
|
|
try {
|
|
await invoke<void>("cancel_download_task", { taskId })
|
|
} catch (error) {
|
|
console.error("Error cancelling download:", error)
|
|
throw error
|
|
}
|
|
}
|
|
|
|
_getHeaders() {
|
|
return {
|
|
...(this.hf_token && { Authorization: `Bearer ${this.hf_token}` })
|
|
}
|
|
}
|
|
}
|