fix: 3911 - inconsistent between download progress and model hub

This commit is contained in:
Louis 2024-11-06 09:20:54 +07:00
parent 2c8c76afa6
commit 964269dc46
No known key found for this signature in database
GPG Key ID: 44FA9F4D33C37DE2
7 changed files with 41 additions and 17 deletions

View File

@ -50,11 +50,6 @@ export class Downloader implements Processor {
const initialDownloadState: DownloadState = {
modelId,
fileName,
time: {
elapsed: 0,
remaining: 0,
},
speed: 0,
percent: 0,
size: {
total: 0,

View File

@ -6,8 +6,8 @@ export type FileStat = {
export type DownloadState = {
modelId: string // TODO: change to download id
fileName: string
time: DownloadTime
speed: number
time?: DownloadTime
speed?: number
percent: number
size: DownloadSize

View File

@ -2,15 +2,19 @@ import { Fragment } from 'react'
import { Progress, Modal, Button } from '@janhq/joi'
import { useAtomValue } from 'jotai'
import { useAtomValue, useSetAtom } from 'jotai'
import useDownloadModel from '@/hooks/useDownloadModel'
import { modelDownloadStateAtom } from '@/hooks/useDownloadState'
import {
modelDownloadStateAtom,
removeDownloadStateAtom,
} from '@/hooks/useDownloadState'
import { formatDownloadPercentage } from '@/utils/converter'
export default function DownloadingState() {
const downloadStates = useAtomValue(modelDownloadStateAtom)
const removeDownloadState = useSetAtom(removeDownloadStateAtom)
const { abortModelDownload } = useDownloadModel()
const totalCurrentProgress = Object.values(downloadStates)
@ -73,6 +77,7 @@ export default function DownloadingState() {
theme="destructive"
onClick={() => {
if (item?.modelId) {
removeDownloadState(item?.modelId)
abortModelDownload(item?.modelId)
}
}}

View File

@ -8,12 +8,13 @@ import { useAtomValue, useSetAtom } from 'jotai'
import useDownloadModel from '@/hooks/useDownloadModel'
import { modelDownloadStateAtom } from '@/hooks/useDownloadState'
import {
modelDownloadStateAtom,
removeDownloadStateAtom,
} from '@/hooks/useDownloadState'
import { formatDownloadPercentage } from '@/utils/converter'
import { removeDownloadingModelAtom } from '@/helpers/atoms/Model.atom'
type Props = {
model: Model
isFromList?: boolean
@ -21,16 +22,16 @@ type Props = {
const ModalCancelDownload = ({ model, isFromList }: Props) => {
const { abortModelDownload } = useDownloadModel()
const removeModelDownload = useSetAtom(removeDownloadingModelAtom)
const removeDownloadState = useSetAtom(removeDownloadStateAtom)
const allDownloadStates = useAtomValue(modelDownloadStateAtom)
const downloadState = allDownloadStates[model.id]
const cancelText = `Cancel ${formatDownloadPercentage(downloadState?.percent ?? 0)}`
const onAbortDownloadClick = useCallback(() => {
removeModelDownload(model.id)
removeDownloadState(model.id)
abortModelDownload(downloadState?.modelId ?? model.id)
}, [downloadState, abortModelDownload, removeModelDownload, model])
}, [downloadState, abortModelDownload, removeDownloadState, model])
return (
<Modal

View File

@ -6,6 +6,8 @@ import { useSetAtom } from 'jotai'
import { toaster } from '@/containers/Toast'
import { setDownloadStateAtom } from './useDownloadState'
import { extensionManager } from '@/extension/ExtensionManager'
import {
@ -16,10 +18,21 @@ import {
export default function useDownloadModel() {
const removeDownloadingModel = useSetAtom(removeDownloadingModelAtom)
const addDownloadingModel = useSetAtom(addDownloadingModelAtom)
const setDownloadStates = useSetAtom(setDownloadStateAtom)
const downloadModel = useCallback(
async (model: string, id?: string, name?: string) => {
addDownloadingModel(id ?? model)
setDownloadStates({
modelId: id ?? model,
downloadState: 'downloading',
fileName: id ?? model,
size: {
total: 0,
transferred: 0,
},
percent: 0,
})
downloadLocalModel(model, id, name).catch((error) => {
if (error.message) {
toaster({
@ -32,7 +45,7 @@ export default function useDownloadModel() {
removeDownloadingModel(model)
})
},
[removeDownloadingModel, addDownloadingModel]
[removeDownloadingModel, addDownloadingModel, setDownloadStates]
)
const abortModelDownload = useCallback(async (model: string) => {

View File

@ -10,8 +10,18 @@ import {
} from '@/helpers/atoms/Model.atom'
// download states
export const modelDownloadStateAtom = atom<Record<string, DownloadState>>({})
/**
* Remove a download state for a particular model.
*/
export const removeDownloadStateAtom = atom(null, (get, set, id: string) => {
const currentState = { ...get(modelDownloadStateAtom) }
delete currentState[id]
set(modelDownloadStateAtom, currentState)
set(removeDownloadingModelAtom, id)
})
/**
* Used to set the download state for a particular model.
*/

View File

@ -1,6 +1,6 @@
import { memo, useState } from 'react'
import { InferenceEngine, Model } from '@janhq/core'
import { Model } from '@janhq/core'
import { Badge, Button, Tooltip, useClickOutside } from '@janhq/joi'
import { useAtom } from 'jotai'
import {