import React, { createContext, useContext, useEffect, useState } from 'react' | |
import { ModelInfo } from '../types' | |
interface ModelContextType { | |
progress: number | |
status: string | |
setProgress: (progress: number) => void | |
setStatus: (status: string) => void | |
modelInfo: ModelInfo | |
setModelInfo: (model: ModelInfo) => void | |
} | |
const ModelContext = createContext<ModelContextType | undefined>(undefined) | |
export function ModelProvider({ children }: { children: React.ReactNode }) { | |
const [progress, setProgress] = useState<number>(0) | |
const [status, setStatus] = useState<string>('idle') | |
const [modelInfo, setModelInfo] = useState<ModelInfo>({} as ModelInfo) | |
// set progress to 0 when model is changed | |
useEffect(() => { | |
setProgress(0) | |
}, [modelInfo.name]) | |
return ( | |
<ModelContext.Provider | |
value={{ | |
progress, | |
setProgress, | |
status, | |
setStatus, | |
modelInfo, | |
setModelInfo | |
}} | |
> | |
{children} | |
</ModelContext.Provider> | |
) | |
} | |
export function useModel() { | |
const context = useContext(ModelContext) | |
if (context === undefined) { | |
throw new Error('useModel must be used within a ModelProvider') | |
} | |
return context | |
} | |