Skip to content

Commit 2b6a08c

Browse files
guohuiyuanguohuiyuan
committed
🎨 优化ai画图代码
1 parent 8945e97 commit 2b6a08c

File tree

3 files changed

+50
-50
lines changed

3 files changed

+50
-50
lines changed

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -418,7 +418,7 @@ print("run[CQ:image,file="+j["img"]+"]")
418418

419419
`import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/aiimage"`
420420

421-
- [x] 设置AI画图密钥
421+
- [x] 设置AI画图密钥xxx
422422
- [x] 设置AI画图接口地址https://api.siliconflow.cn/v1/images/generations
423423
- [x] 设置AI画图模型名Kwai-Kolors/Kolors
424424
- [x] 查看AI画图配置

‎plugin/aiimage/config.go‎

Lines changed: 12 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -3,60 +3,40 @@ package aiimage
33

44
import (
55
"fmt"
6-
"os"
76
"strings"
87
"sync"
9-
"time"
108

119
sql "github.com/FloatTech/sqlite"
1210
)
1311

14-
// Storage 管理画图配置存储
15-
type Storage struct {
12+
// storage 管理画图配置存储
13+
type storage struct {
1614
sync.RWMutex
1715
db sql.Sqlite
1816
}
1917

20-
var (
21-
sdb = &Storage{
22-
db: sql.New("data/aiimage/config.db"),
23-
}
24-
)
25-
26-
func init() {
27-
if err := os.MkdirAll("data/aiimage", 0755); err != nil {
28-
panic(err)
29-
}
30-
if err := sdb.db.Open(time.Hour * 24); err != nil {
31-
panic(err)
32-
}
33-
if err := sdb.db.Create("config", &ImageConfig{}); err != nil {
34-
panic(err)
35-
}
36-
}
37-
38-
// ImageConfig 存储AI画图配置信息
39-
type ImageConfig struct {
18+
// imageConfig 存储AI画图配置信息
19+
type imageConfig struct {
4020
ID int64 `db:"id"` // 主键ID
4121
APIKey string `db:"apiKey"` // API密钥
4222
APIURL string `db:"apiUrl"` // API地址
4323
ModelName string `db:"modelName"` // 画图模型名称
4424
}
4525

46-
// GetConfig 获取当前配置
47-
func GetConfig() ImageConfig {
26+
// getConfig 获取当前配置
27+
func (sdb *storage) getConfig() imageConfig {
4828
sdb.RLock()
4929
defer sdb.RUnlock()
50-
cfg := ImageConfig{}
30+
cfg := imageConfig{}
5131
_ = sdb.db.Find("config", &cfg, "WHERE id = 1")
5232
return cfg
5333
}
5434

55-
// SetConfig 设置AI画图配置
56-
func SetConfig(apiKey, apiURL, modelName string) error {
35+
// setConfig 设置AI画图配置
36+
func (sdb *storage) setConfig(apiKey, apiURL, modelName string) error {
5737
sdb.Lock()
5838
defer sdb.Unlock()
59-
return sdb.db.Insert("config", &ImageConfig{
39+
return sdb.db.Insert("config", &imageConfig{
6040
ID: 1,
6141
APIKey: apiKey,
6242
APIURL: apiURL,
@@ -65,8 +45,8 @@ func SetConfig(apiKey, apiURL, modelName string) error {
6545
}
6646

6747
// PrintConfig 返回格式化后的配置信息
68-
func PrintConfig() string {
69-
cfg := GetConfig()
48+
func (sdb *storage) PrintConfig() string {
49+
cfg := sdb.getConfig()
7050
var builder strings.Builder
7151
builder.WriteString("当前AI画图配置:\n")
7252
builder.WriteString(fmt.Sprintf("• 密钥: %s\n", cfg.APIKey))

‎plugin/aiimage/main.go‎

Lines changed: 37 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,11 @@ import (
66
"encoding/json"
77
"net/http"
88
"strings"
9+
"time"
910

11+
fcext "github.com/FloatTech/floatbox/ctxext"
1012
"github.com/FloatTech/floatbox/web"
13+
sql "github.com/FloatTech/sqlite"
1114
"github.com/tidwall/gjson"
1215
zero "github.com/wdvxdr1123/ZeroBot"
1316
"github.com/wdvxdr1123/ZeroBot/message"
@@ -18,60 +21,78 @@ import (
1821
)
1922

2023
func init() {
24+
var sdb = &storage{}
25+
2126
en := control.AutoRegister(&ctrl.Options[*zero.Ctx]{
2227
DisableOnDefault: false,
2328
Extra: control.ExtraFromString("aiimage"),
2429
Brief: "AI画图",
25-
Help: "- 设置AI画图密钥\n" +
30+
Help: "- 设置AI画图密钥xxx\n" +
2631
"- 设置AI画图接口地址https://api.siliconflow.cn/v1/images/generations\n" +
2732
"- 设置AI画图模型名Kwai-Kolors/Kolors\n" +
2833
"- 查看AI画图配置\n" +
2934
"- AI画图 [描述]",
3035
PrivateDataFolder: "aiimage",
3136
})
3237

33-
en.OnPrefix("设置AI画图密钥", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
38+
getdb := fcext.DoOnceOnSuccess(func(ctx *zero.Ctx) bool {
39+
sdb.db = sql.New(en.DataFolder() + "aiimage.db")
40+
err := sdb.db.Open(time.Hour)
41+
if err == nil {
42+
// 创建配置表
43+
err = sdb.db.Create("config", &imageConfig{})
44+
if err != nil {
45+
ctx.SendChain(message.Text("[ERROR]:", err))
46+
return false
47+
}
48+
return true
49+
}
50+
ctx.SendChain(message.Text("[ERROR]:", err))
51+
return false
52+
})
53+
54+
en.OnPrefix("设置AI画图密钥", getdb, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
3455
Handle(func(ctx *zero.Ctx) {
3556
apiKey := strings.TrimSpace(ctx.State["args"].(string))
36-
cfg := GetConfig()
37-
err := SetConfig(apiKey, cfg.APIURL, cfg.ModelName)
57+
cfg := sdb.getConfig()
58+
err := sdb.setConfig(apiKey, cfg.APIURL, cfg.ModelName)
3859
if err != nil {
3960
ctx.SendChain(message.Text("ERROR: 设置API密钥失败: ", err))
4061
return
4162
}
4263
ctx.SendChain(message.Text("成功设置API密钥"))
4364
})
4465

45-
en.OnPrefix("设置AI画图接口地址", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
66+
en.OnPrefix("设置AI画图接口地址", getdb, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
4667
Handle(func(ctx *zero.Ctx) {
4768
apiURL := strings.TrimSpace(ctx.State["args"].(string))
48-
cfg := GetConfig()
49-
err := SetConfig(cfg.APIKey, apiURL, cfg.ModelName)
69+
cfg := sdb.getConfig()
70+
err := sdb.setConfig(cfg.APIKey, apiURL, cfg.ModelName)
5071
if err != nil {
5172
ctx.SendChain(message.Text("ERROR: 设置API地址失败: ", err))
5273
return
5374
}
5475
ctx.SendChain(message.Text("成功设置API地址"))
5576
})
5677

57-
en.OnPrefix("设置AI画图模型名", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
78+
en.OnPrefix("设置AI画图模型名", getdb, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
5879
Handle(func(ctx *zero.Ctx) {
5980
modelName := strings.TrimSpace(ctx.State["args"].(string))
60-
cfg := GetConfig()
61-
err := SetConfig(cfg.APIKey, cfg.APIURL, modelName)
81+
cfg := sdb.getConfig()
82+
err := sdb.setConfig(cfg.APIKey, cfg.APIURL, modelName)
6283
if err != nil {
6384
ctx.SendChain(message.Text("ERROR: 设置模型失败: ", err))
6485
return
6586
}
6687
ctx.SendChain(message.Text("成功设置模型: ", modelName))
6788
})
6889

69-
en.OnFullMatch("查看AI画图配置", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
90+
en.OnFullMatch("查看AI画图配置", getdb, zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
7091
Handle(func(ctx *zero.Ctx) {
71-
ctx.SendChain(message.Text(PrintConfig()))
92+
ctx.SendChain(message.Text(sdb.PrintConfig()))
7293
})
7394

74-
en.OnPrefix("AI画图").SetBlock(true).
95+
en.OnPrefix("AI画图", getdb).SetBlock(true).
7596
Handle(func(ctx *zero.Ctx) {
7697
ctx.SendChain(message.Text("少女思考中..."))
7798
prompt := strings.TrimSpace(ctx.State["args"].(string))
@@ -80,22 +101,21 @@ func init() {
80101
return
81102
}
82103

83-
cfg := GetConfig()
104+
cfg := sdb.getConfig()
84105
if cfg.APIKey == "" || cfg.APIURL == "" || cfg.ModelName == "" {
85106
ctx.SendChain(message.Text("请先配置API密钥、地址和模型"))
86107
return
87108
}
88109

89110
// 准备请求数据
90-
reqData := map[string]interface{}{
111+
reqBytes, _ := json.Marshal(map[string]interface{}{
91112
"model": cfg.ModelName,
92113
"prompt": prompt,
93114
"image_size": "1024x1024",
94115
"batch_size": 4,
95116
"num_inference_steps": 20,
96117
"guidance_scale": 7.5,
97-
}
98-
reqBytes, _ := json.Marshal(reqData)
118+
})
99119

100120
// 发送API请求
101121
data, err := web.RequestDataWithHeaders(

0 commit comments

Comments
 (0)