mirror of https://gitee.com/godoos/godoos.git
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
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
|
|
}
|
|
|