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

web/containers/ModelLabel/NotEnoughRamLabel.tsx → 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,13 +66,21 @@ export default function useGetSystemResources() {
  )
  }
  setTotalNvidiaVram(totalNvidiaVram)
+ setAvailableVram(
+ gpus.reduce(
+ (total: number, gpu: { memoryFree: string }) =>
+ total + Number(gpu.memoryFree),
+ 0
+ )
+ )
  }, [
  setUsedRam,
  setTotalRam,
  setRamUtilitized,
  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