You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

167 lines
3.6 KiB

// Copyright (c) seasonjs. All rights reserved.
// Licensed under the MIT License. See License.txt in the project root for license information.
//go:build windows && amd64
package sd
import (
_ "embed"
"encoding/json"
"log"
"os/exec"
"strings"
"syscall"
"golang.org/x/sys/cpu"
)
//go:embed deps/windows/sd-abi_avx2.dll
var libStableDiffusionAvx2 []byte
//go:embed deps/windows/sd-abi_avx.dll
var libStableDiffusionAvx []byte
//go:embed deps/windows/sd-abi_avx512.dll
var libStableDiffusionAvx512 []byte
//go:embed deps/windows/sd-abi_cuda12.dll
var libStableDiffusionCuda12 []byte
//go:embed deps/windows/sd-abi_rocm5.5.dll
var libStableDiffusionRocm5 []byte
var libName = "stable-diffusion-*.dll"
func getDl(gpu bool) []byte {
if gpu {
info, err := NewGPU()
if err != nil {
log.Println(err)
}
cuda := info.Cuda()
rocm := info.ROCm()
if cuda.Available() {
log.Print("get gpu info: ", cuda.Name)
log.Println("Use GPU CUDA instead.")
return libStableDiffusionCuda12
}
if rocm.Available() {
log.Print("get gpu info: ", cuda.Name)
log.Println("Use GPU ROCm instead.")
return libStableDiffusionRocm5
}
log.Println("GPU not support, use CPU instead.")
}
if cpu.X86.HasAVX512 {
log.Println("Use CPU AVX512 instead.")
return libStableDiffusionAvx512
}
if cpu.X86.HasAVX2 {
log.Println("Use CPU AVX2 instead.")
return libStableDiffusionAvx2
}
if cpu.X86.HasAVX {
log.Println("Use CPU AVX instead.")
return libStableDiffusionAvx
}
panic("Automatic loading of dynamic library failed, please use `NewRwkvModel` method load manually. ")
}
type Driver struct {
Name string `json:"Name"`
AdapterCompatibility string `json:"AdapterCompatibility"`
AdapterRAM int64 `json:"AdapterRAM"`
}
func (d *Driver) Available() bool {
return d.Name != "" && d.AdapterCompatibility != "" && d.AdapterRAM != 0
}
// GPU 类用于管理显卡信息
type GPU struct {
drivers []Driver
cuda []Driver
rocm []Driver
}
func NewGPU() (*GPU, error) {
cmd := exec.Command("powershell", `
$graphicsCards = Get-WmiObject Win32_VideoController
$graphicsArray = @()
$graphicsEmpty = @{
'Name' = ''
'AdapterCompatibility' = ''
'AdapterRAM' = 0
}
$graphicsArray += $graphicsEmpty
foreach ($card in $graphicsCards) {
$graphicsInfo = @{
'Name' = $card.Caption
'AdapterCompatibility' = $card.VideoProcessor
'AdapterRAM' = $card.AdapterRAM
}
$graphicsArray += $graphicsInfo
}
$graphicsArray | ConvertTo-Json
`)
cmd.SysProcAttr = &syscall.SysProcAttr{
HideWindow: true,
//https://learn.microsoft.com/en-us/windows/win32/procthread/process-creation-flags
CreationFlags: 0x08000000,
}
output, err := cmd.CombinedOutput()
if err != nil {
return nil, err
}
var drivers []Driver
err = json.Unmarshal(output, &drivers)
if err != nil {
return nil, err
}
cudaSupport := make([]Driver, 0)
rocmSupport := make([]Driver, 0)
for _, driver := range drivers {
if strings.Contains(strings.ToUpper(driver.Name), "NVIDIA") {
cudaSupport = append(cudaSupport, driver)
} else if strings.Contains(strings.ToUpper(driver.Name), "AMD") {
rocmSupport = append(rocmSupport, driver)
}
}
return &GPU{
drivers: drivers,
cuda: cudaSupport,
rocm: rocmSupport,
}, nil
}
func (g *GPU) Cuda() *Driver {
if len(g.cuda) > 0 {
return &g.cuda[0]
}
return &Driver{}
}
func (g *GPU) ROCm() *Driver {
if len(g.rocm) > 0 {
return &g.rocm[0]
}
return &Driver{}
}
func (g *GPU) Info() []Driver {
return g.drivers
}