fix: model recommendation label is not based on VRAM (#2517)

web/containers/ModelLabel/NotEnoughMemoryLabel.tsx

@@ -10,9 +10,9 @@ import {
 } from '@janhq/uikit'
 import { InfoIcon } from 'lucide-react'
 
-const NotEnoughRamLabel: React.FC = () => (
+const NotEnoughMemoryLabel = ({ unit }: { unit: string }) => (
   <Badge className="space-x-1 rounded-md" themes="danger">
-    <span>Not enough RAM</span>
+    <span>Not enough {unit}</span>
     <Tooltip>
       <TooltipTrigger>
         <InfoIcon size={16} />
@@ -31,4 +31,4 @@ const NotEnoughRamLabel: React.FC = () => (
   </Badge>
 )
 
-export default React.memo(NotEnoughRamLabel)
+export default React.memo(NotEnoughMemoryLabel)

web/containers/ModelLabel/index.tsx

@@ -4,13 +4,19 @@ import { useAtomValue } from 'jotai'
 
 import { useActiveModel } from '@/hooks/useActiveModel'
 
-import NotEnoughRamLabel from './NotEnoughRamLabel'
+import { useSettings } from '@/hooks/useSettings'
+
+import NotEnoughMemoryLabel from './NotEnoughMemoryLabel'
 
 import RecommendedLabel from './RecommendedLabel'
 
 import SlowOnYourDeviceLabel from './SlowOnYourDeviceLabel'
 
-import { totalRamAtom, usedRamAtom } from '@/helpers/atoms/SystemBar.atom'
+import {
+  availableVramAtom,
+  totalRamAtom,
+  usedRamAtom,
+} from '@/helpers/atoms/SystemBar.atom'
 
 type Props = {
   size: number
@@ -20,12 +26,21 @@ const ModelLabel: React.FC<Props> = ({ size }) => {
   const { activeModel } = useActiveModel()
   const totalRam = useAtomValue(totalRamAtom)
   const usedRam = useAtomValue(usedRamAtom)
+  const availableVram = useAtomValue(availableVramAtom)
+  const { settings } = useSettings()
 
   const getLabel = (size: number) => {
     const minimumRamModel = size * 1.25
-    const availableRam = totalRam - usedRam + (activeModel?.metadata.size ?? 0)
+    const availableRam =
+      settings?.run_mode === 'gpu'
+        ? availableVram * 1000000 // MB to bytes
+        : totalRam - usedRam + (activeModel?.metadata.size ?? 0)
     if (minimumRamModel > totalRam) {
-      return <NotEnoughRamLabel />
+      return (
+        <NotEnoughMemoryLabel
+          unit={settings?.run_mode === 'gpu' ? 'VRAM' : 'RAM'}
+        />
+      )
     }
     if (minimumRamModel < availableRam) {
       return <RecommendedLabel />

web/helpers/atoms/SystemBar.atom.ts

@@ -9,4 +9,5 @@ export const ramUtilitizedAtom = atom<number>(0)
 export const gpusAtom = atom<Record<string, never>[]>([])
 
 export const nvidiaTotalVramAtom = atom<number>(0)
+export const availableVramAtom = atom<number>(0)
 export const systemMonitorCollapseAtom = atom<boolean>(false)

web/hooks/useGetSystemResources.ts

@@ -12,6 +12,7 @@ import {
   nvidiaTotalVramAtom,
   gpusAtom,
   ramUtilitizedAtom,
+  availableVramAtom,
 } from '@/helpers/atoms/SystemBar.atom'
 
 export default function useGetSystemResources() {
@@ -24,6 +25,7 @@ export default function useGetSystemResources() {
   const setUsedRam = useSetAtom(usedRamAtom)
   const setCpuUsage = useSetAtom(cpuUsageAtom)
   const setTotalNvidiaVram = useSetAtom(nvidiaTotalVramAtom)
+  const setAvailableVram = useSetAtom(availableVramAtom)
   const setRamUtilitized = useSetAtom(ramUtilitizedAtom)
 
   const getSystemResources = useCallback(async () => {
@@ -64,6 +66,13 @@ export default function useGetSystemResources() {
       )
     }
     setTotalNvidiaVram(totalNvidiaVram)
+    setAvailableVram(
+      gpus.reduce(
+        (total: number, gpu: { memoryFree: string }) =>
+          total + Number(gpu.memoryFree),
+        0
+      )
+    )
   }, [
     setUsedRam,
     setTotalRam,
@@ -71,6 +80,7 @@ export default function useGetSystemResources() {
     setCpuUsage,
     setGpus,
     setTotalNvidiaVram,
+    setAvailableVram,
   ])
 
   const watch = () => {

web/hooks/useSettings.ts

@@ -5,13 +5,24 @@ import { atom, useAtom } from 'jotai'
 
 export const isShowNotificationAtom = atom<boolean>(false)
 
+export type AppSettings = {
+  run_mode: 'cpu' | 'gpu' | undefined
+  notify: boolean
+  gpus_in_use: string[]
+  vulkan: boolean
+  gpus: string[]
+}
+
 export const useSettings = () => {
   const [isGPUModeEnabled, setIsGPUModeEnabled] = useState(false) // New state for GPU mode
   const [showNotification, setShowNotification] = useAtom(
     isShowNotificationAtom
   )
+  const [settings, setSettings] = useState<AppSettings>()
 
   useEffect(() => {
+    readSettings().then((settings) => setSettings(settings as AppSettings))
+
     setTimeout(() => validateSettings, 3000)
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [])
@@ -79,5 +90,6 @@ export const useSettings = () => {
     saveSettings,
     setShowNotification,
     validateSettings,
+    settings,
   }
 }

web/screens/ExploreModels/ExploreModelItemHeader/index.tsx

@@ -24,6 +24,8 @@ import { MainViewState } from '@/constants/screens'
 import { useCreateNewThread } from '@/hooks/useCreateNewThread'
 import useDownloadModel from '@/hooks/useDownloadModel'
 
+import { useSettings } from '@/hooks/useSettings'
+
 import { toGibibytes } from '@/utils/converter'
 
 import { mainViewStateAtom } from '@/helpers/atoms/App.atom'
@@ -45,11 +47,11 @@ type Props = {
   open: string
 }
 
-const getLabel = (size: number, ram: number) => {
+const getLabel = (size: number, ram: number, unit: string = 'RAM') => {
   if (size * 1.25 >= ram) {
     return (
       <Badge className="rounded-md" themes="danger">
-        Not enough RAM
+        Not enough {unit}
       </Badge>
     )
   } else {
@@ -67,13 +69,14 @@ const ExploreModelItemHeader: React.FC<Props> = ({ model, onClick, open }) => {
   const downloadedModels = useAtomValue(downloadedModelsAtom)
   const { requestCreateNewThread } = useCreateNewThread()
   const totalRam = useAtomValue(totalRamAtom)
+  const { settings } = useSettings()
 
   const nvidiaTotalVram = useAtomValue(nvidiaTotalVramAtom)
   const setMainViewState = useSetAtom(mainViewStateAtom)
 
   // Default nvidia returns vram in MB, need to convert to bytes to match the unit of totalRamW
   let ram = nvidiaTotalVram * 1024 * 1024
-  if (ram === 0) {
+  if (ram === 0 || settings?.run_mode === 'cpu') {
     ram = totalRam
   }
   const serverEnabled = useAtomValue(serverEnabledAtom)
@@ -158,7 +161,11 @@ const ExploreModelItemHeader: React.FC<Props> = ({ model, onClick, open }) => {
           <span className="mr-4 font-semibold text-muted-foreground">
             {toGibibytes(model.metadata.size)}
           </span>
-          {getLabel(model.metadata.size, ram)}
+          {getLabel(
+            model.metadata.size,
+            ram,
+            settings?.run_mode === 'gpu' ? 'VRAM' : 'RAM'
+          )}
 
           {downloadButton}
           <ChevronDownIcon