fix(gpu-detection): default to CPU if there is less than 4GB of GPU available

Signed-off-by: Ettore Di Giacinto <[email protected]>

Author: mudler

pkg/system/capabilities.go

@@ -5,12 +5,15 @@ import (
 	"runtime"
 	"strings"
 
+	"github.com/jaypipes/ghw/pkg/gpu"
 	"github.com/mudler/LocalAI/pkg/xsysinfo"
 	"github.com/rs/zerolog/log"
 )
 
 type SystemState struct {
 	GPUVendor string
+	gpus      []*gpu.GraphicsCard
+	vram      uint64
 }
 
 const (
@@ -91,24 +94,39 @@ func (s *SystemState) getSystemCapabilities() string {
 	}
 
 	log.Info().Str("Capability", s.GPUVendor).Msgf("Capability automatically detected, set %s to override", capabilityEnv)
+	// If vram is less than 4GB, let's default to CPU but warn the user that they can override that via env
+	if s.vram <= 4*1024*1024*1024 {
+		log.Warn().Msgf("VRAM is less than 4GB, defaulting to CPU. Set %s to override", capabilityEnv)
+		return defaultCapability
+	}
+
 	return s.GPUVendor
 }
 
 func GetSystemState() (*SystemState, error) {
-	gpuVendor, _ := detectGPUVendor()
+	gpus, err := xsysinfo.GPUs()
+	if err != nil {
+		return nil, err
+	}
+
+	log.Debug().Any("gpus", gpus).Msg("GPUs")
+	gpuVendor, _ := detectGPUVendor(gpus)
 	log.Debug().Str("gpuVendor", gpuVendor).Msg("GPU vendor")
 
+	vram, err := xsysinfo.TotalAvailableVRAM()
+	if err != nil {
+		return nil, err
+	}
+	log.Debug().Any("vram", vram).Msg("Total available VRAM")
+
 	return &SystemState{
 		GPUVendor: gpuVendor,
+		gpus:      gpus,
+		vram:      vram,
 	}, nil
 }
 
-func detectGPUVendor() (string, error) {
-	gpus, err := xsysinfo.GPUs()
-	if err != nil {
-		return "", err
-	}
-
+func detectGPUVendor(gpus []*gpu.GraphicsCard) (string, error) {
 	for _, gpu := range gpus {
 		if gpu.DeviceInfo != nil {
 			if gpu.DeviceInfo.Vendor != nil {