Skip to content

Commit 307f66f

Browse files
guohuiyuanguohuiyuan
committed
✨ 添加ai画图
1 parent c690405 commit 307f66f

File tree

3 files changed

+227
-0
lines changed

3 files changed

+227
-0
lines changed

‎main.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ import (
6767
_ "github.com/FloatTech/ZeroBot-Plugin/custom" // 自定义插件合集
6868
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/ahsai" // ahsai tts
6969
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/aifalse" // 服务器监控
70+
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/aiimage" // AI画图
7071
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/aiwife" // 随机老婆
7172
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/alipayvoice" // 支付宝到账语音
7273
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/animetrace" // AnimeTrace 动画/Galgame识别

‎plugin/aiimage/config.go‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
// Package aiimage 提供AI画图功能配置
2+
package aiimage
3+
4+
import (
5+
"fmt"
6+
"os"
7+
"strings"
8+
"sync"
9+
"time"
10+
11+
sql "github.com/FloatTech/sqlite"
12+
)
13+
14+
// Storage 管理画图配置存储
15+
type Storage struct {
16+
sync.RWMutex
17+
db sql.Sqlite
18+
}
19+
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 {
40+
ID int64 `db:"id"` // 主键ID
41+
APIKey string `db:"apiKey"` // API密钥
42+
APIURL string `db:"apiUrl"` // API地址
43+
ModelName string `db:"modelName"` // 画图模型名称
44+
}
45+
46+
// GetConfig 获取当前配置
47+
func GetConfig() ImageConfig {
48+
sdb.RLock()
49+
defer sdb.RUnlock()
50+
cfg := ImageConfig{}
51+
_ = sdb.db.Find("config", &cfg, "WHERE id = 1")
52+
return cfg
53+
}
54+
55+
// SetConfig 设置AI画图配置
56+
func SetConfig(apiKey, apiURL, modelName string) error {
57+
sdb.Lock()
58+
defer sdb.Unlock()
59+
return sdb.db.Insert("config", &ImageConfig{
60+
ID: 1,
61+
APIKey: apiKey,
62+
APIURL: apiURL,
63+
ModelName: modelName,
64+
})
65+
}
66+
67+
// PrintConfig 返回格式化后的配置信息
68+
func PrintConfig() string {
69+
cfg := GetConfig()
70+
var builder strings.Builder
71+
builder.WriteString("当前AI画图配置:\n")
72+
builder.WriteString(fmt.Sprintf("• 密钥: %s\n", cfg.APIKey))
73+
builder.WriteString(fmt.Sprintf("• 接口地址: %s\n", cfg.APIURL))
74+
builder.WriteString(fmt.Sprintf("• 模型名: %s\n", cfg.ModelName))
75+
return builder.String()
76+
}

‎plugin/aiimage/main.go‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
// Package aiimage AI画图
2+
package aiimage
3+
4+
import (
5+
"bytes"
6+
"encoding/json"
7+
"net/http"
8+
"strings"
9+
10+
"github.com/FloatTech/floatbox/web"
11+
"github.com/tidwall/gjson"
12+
zero "github.com/wdvxdr1123/ZeroBot"
13+
"github.com/wdvxdr1123/ZeroBot/message"
14+
15+
ctrl "github.com/FloatTech/zbpctrl"
16+
"github.com/FloatTech/zbputils/control"
17+
"github.com/FloatTech/zbputils/ctxext"
18+
)
19+
20+
func init() {
21+
en := control.AutoRegister(&ctrl.Options[*zero.Ctx]{
22+
DisableOnDefault: false,
23+
Extra: control.ExtraFromString("aiimage"),
24+
Brief: "AI画图",
25+
Help: "- 设置AI画图密钥\n" +
26+
"- 设置AI画图接口地址https://api.siliconflow.cn/v1/images/generations\n" +
27+
"- 设置AI画图模型名Kwai-Kolors/Kolors\n" +
28+
"- 查看AI画图配置\n" +
29+
"- AI画图 [描述]",
30+
PrivateDataFolder: "aiimage",
31+
})
32+
33+
en.OnPrefix("设置AI画图密钥", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
34+
Handle(func(ctx *zero.Ctx) {
35+
apiKey := strings.TrimSpace(ctx.State["args"].(string))
36+
cfg := GetConfig()
37+
err := SetConfig(apiKey, cfg.APIURL, cfg.ModelName)
38+
if err != nil {
39+
ctx.SendChain(message.Text("ERROR: 设置API密钥失败: ", err))
40+
return
41+
}
42+
ctx.SendChain(message.Text("成功设置API密钥"))
43+
})
44+
45+
en.OnPrefix("设置AI画图接口地址", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
46+
Handle(func(ctx *zero.Ctx) {
47+
apiURL := strings.TrimSpace(ctx.State["args"].(string))
48+
cfg := GetConfig()
49+
err := SetConfig(cfg.APIKey, apiURL, cfg.ModelName)
50+
if err != nil {
51+
ctx.SendChain(message.Text("ERROR: 设置API地址失败: ", err))
52+
return
53+
}
54+
ctx.SendChain(message.Text("成功设置API地址"))
55+
})
56+
57+
en.OnPrefix("设置AI画图模型名", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
58+
Handle(func(ctx *zero.Ctx) {
59+
modelName := strings.TrimSpace(ctx.State["args"].(string))
60+
cfg := GetConfig()
61+
err := SetConfig(cfg.APIKey, cfg.APIURL, modelName)
62+
if err != nil {
63+
ctx.SendChain(message.Text("ERROR: 设置模型失败: ", err))
64+
return
65+
}
66+
ctx.SendChain(message.Text("成功设置模型: ", modelName))
67+
})
68+
69+
en.OnFullMatch("查看AI画图配置", zero.OnlyPrivate, zero.SuperUserPermission).SetBlock(true).
70+
Handle(func(ctx *zero.Ctx) {
71+
ctx.SendChain(message.Text(PrintConfig()))
72+
})
73+
74+
en.OnPrefix("AI画图").SetBlock(true).
75+
Handle(func(ctx *zero.Ctx) {
76+
prompt := strings.TrimSpace(ctx.State["args"].(string))
77+
if prompt == "" {
78+
ctx.SendChain(message.Text("请输入图片描述"))
79+
return
80+
}
81+
82+
cfg := GetConfig()
83+
if cfg.APIKey == "" || cfg.APIURL == "" {
84+
ctx.SendChain(message.Text("请先配置API密钥和地址"))
85+
return
86+
}
87+
88+
// 准备请求数据
89+
reqData := map[string]interface{}{
90+
"model": cfg.ModelName,
91+
"prompt": prompt,
92+
"image_size": "1024x1024",
93+
"batch_size": 4,
94+
"num_inference_steps": 20,
95+
"guidance_scale": 7.5,
96+
}
97+
reqBytes, _ := json.Marshal(reqData)
98+
99+
// 发送API请求
100+
data, err := web.RequestDataWithHeaders(
101+
web.NewDefaultClient(),
102+
cfg.APIURL,
103+
"POST",
104+
func(req *http.Request) error {
105+
req.Header.Set("Authorization", "Bearer "+cfg.APIKey)
106+
req.Header.Set("Content-Type", "application/json")
107+
return nil
108+
},
109+
bytes.NewReader(reqBytes),
110+
)
111+
if err != nil {
112+
ctx.SendChain(message.Text("API请求失败: ", err))
113+
return
114+
}
115+
116+
// 解析API响应
117+
jsonData := gjson.ParseBytes(data)
118+
images := jsonData.Get("images")
119+
if !images.Exists() {
120+
images = jsonData.Get("data")
121+
if !images.Exists() {
122+
ctx.SendChain(message.Text("未获取到图片URL"))
123+
return
124+
}
125+
}
126+
127+
// 发送生成的图片和相关信息
128+
inferenceTime := jsonData.Get("timings.inference").Float()
129+
seed := jsonData.Get("seed").Int()
130+
msg := make(message.Message, 0, 1)
131+
msg = append(msg, ctxext.FakeSenderForwardNode(ctx, message.Text("图片生成成功!\n",
132+
"提示词: ", prompt, "\n",
133+
"模型: ", cfg.ModelName, "\n",
134+
"推理时间: ", inferenceTime, "秒\n",
135+
"种子: ", seed)))
136+
137+
// 添加所有图片
138+
images.ForEach(func(_, value gjson.Result) bool {
139+
url := value.Get("url").String()
140+
if url != "" {
141+
msg = append(msg, ctxext.FakeSenderForwardNode(ctx, message.Image(url)))
142+
}
143+
return true
144+
})
145+
146+
if len(msg) > 0 {
147+
ctx.Send(msg)
148+
}
149+
})
150+
}

0 commit comments

Comments
 (0)