diff --git a/README.md b/README.md
index 0cf67195f6..61de32326f 100644
--- a/README.md
+++ b/README.md
@@ -346,6 +346,24 @@ print("run[CQ:image,file="+j["img"]+"]")
- [x] 设置默认限速为每 m [分钟 | 秒] n 次触发
+
+
+ ai绘图
+
+ `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/aipaint"`
+
+ - [x] [ ai绘图 | 生成色图 | 生成涩图 | ai画图 ] xxx
+
+ - [x] [ 以图绘图 | 以图生图 | 以图画图 ] xxx [图片]|@xxx|[qq号]
+
+ - [ ] 设置ai绘图配置 [server] [token]
+
+ 例1: 设置ai绘图配置 http://91.216.169.75:5010 abc
+
+ 例2: 设置ai绘图配置 http://91.217.139.190:5010 abc
+
+ 通过 http://91.217.139.190:5010/token 获取token
+
AIWife
diff --git a/main.go b/main.go
index 37c34f07ff..4d9758dd2a 100644
--- a/main.go
+++ b/main.go
@@ -60,6 +60,7 @@ import (
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/ahsai" // ahsai tts
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/ai_false" // 服务器监控
+ _ "github.com/FloatTech/ZeroBot-Plugin/plugin/aipaint" // ai绘图
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/aiwife" // 随机老婆
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/alipayvoice" // 支付宝到账语音
_ "github.com/FloatTech/ZeroBot-Plugin/plugin/b14" // base16384加解密
diff --git a/plugin/aipaint/aipaint.go b/plugin/aipaint/aipaint.go
new file mode 100644
index 0000000000..38d2445dbd
--- /dev/null
+++ b/plugin/aipaint/aipaint.go
@@ -0,0 +1,167 @@
+// Package aipaint ai绘图
+package aipaint
+
+import (
+ "bytes"
+ "encoding/base64"
+ "encoding/json"
+ "fmt"
+ "image"
+ "net/url"
+ "os"
+ "regexp"
+ "strconv"
+ "strings"
+
+ "github.com/FloatTech/floatbox/binary"
+ "github.com/FloatTech/floatbox/file"
+ "github.com/FloatTech/floatbox/img/writer"
+ "github.com/FloatTech/floatbox/web"
+ ctrl "github.com/FloatTech/zbpctrl"
+ "github.com/FloatTech/zbputils/control"
+ "github.com/FloatTech/zbputils/ctxext"
+ zero "github.com/wdvxdr1123/ZeroBot"
+ "github.com/wdvxdr1123/ZeroBot/message"
+)
+
+var (
+ datapath string
+ predictRe = regexp.MustCompile(`{"steps".+?}`)
+ // 参考host http://91.217.139.190:5010 http://91.216.169.75:5010
+ aipaintTxt2ImgURL = "/got_image?token=%v&tags=%v"
+ aipaintImg2ImgURL = "/got_image2image?token=%v&tags=%v"
+ cfg = newServerConfig("data/aipaint/config.json")
+)
+
+type result struct {
+ Steps int `json:"steps"`
+ Sampler string `json:"sampler"`
+ Seed int `json:"seed"`
+ Strength float64 `json:"strength"`
+ Noise float64 `json:"noise"`
+ Scale float64 `json:"scale"`
+ Uc string `json:"uc"`
+}
+
+func (r *result) String() string {
+ return fmt.Sprintf("steps: %v\nsampler: %v\nseed: %v\nstrength: %v\nnoise: %v\nscale: %v\nuc: %v\n", r.Steps, r.Sampler, r.Seed, r.Strength, r.Noise, r.Scale, r.Uc)
+}
+
+func init() { // 插件主体
+ engine := control.Register("aipaint", &ctrl.Options[*zero.Ctx]{
+ DisableOnDefault: false,
+ Help: "ai绘图\n" +
+ "- [ ai绘图 | 生成色图 | 生成涩图 | ai画图 ] xxx\n" +
+ "- [ 以图绘图 | 以图生图 | 以图画图 ] xxx [图片]|@xxx|[qq号]\n" +
+ "- 设置ai绘图配置 [server] [token]\n" +
+ "例1: 设置ai绘图配置 http://91.216.169.75:5010 abc\n" +
+ "例2: 设置ai绘图配置 http://91.217.139.190:5010 abc\n" +
+ "通过 http://91.217.139.190:5010/token 获取token",
+ PrivateDataFolder: "aipaint",
+ })
+ datapath = file.BOTPATH + "/" + engine.DataFolder()
+ engine.OnPrefixGroup([]string{`ai绘图`, `生成色图`, `生成涩图`, `ai画图`}).SetBlock(true).
+ Handle(func(ctx *zero.Ctx) {
+ server, token, err := cfg.load()
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ ctx.SendChain(message.Text("少女祈祷中..."))
+ args := ctx.State["args"].(string)
+ data, err := web.GetData(server + fmt.Sprintf(aipaintTxt2ImgURL, token, url.QueryEscape(strings.TrimSpace(strings.ReplaceAll(args, " ", "%20")))))
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ sendAiImg(ctx, data)
+ })
+ engine.OnRegex(`^(以图绘图|以图生图|以图画图)[\s\S]*?(\[CQ:(image\,file=([0-9a-zA-Z]{32}).*|at.+?(\d{5,11}))\].*|(\d+))$`).SetBlock(true).
+ Handle(func(ctx *zero.Ctx) {
+ server, token, err := cfg.load()
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ c := newContext(ctx.Event.UserID)
+ list := ctx.State["regex_matched"].([]string)
+ err = c.prepareLogos(list[4]+list[5]+list[6], strconv.FormatInt(ctx.Event.UserID, 10))
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ args := strings.TrimSuffix(strings.TrimPrefix(list[0], list[1]), list[2])
+ if args == "" {
+ ctx.SendChain(message.Text("ERROR: 以图绘图必须添加tag"))
+ return
+ }
+ ctx.SendChain(message.Text("少女祈祷中..."))
+ postURL := server + fmt.Sprintf(aipaintImg2ImgURL, token, url.QueryEscape(strings.TrimSpace(strings.ReplaceAll(args, " ", "%20"))))
+
+ f, err := os.Open(c.headimgsdir[0])
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ defer f.Close()
+
+ img, _, err := image.Decode(f)
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ imageShape := ""
+ switch {
+ case img.Bounds().Dx() > img.Bounds().Dy():
+ imageShape = "Landscape"
+ case img.Bounds().Dx() == img.Bounds().Dy():
+ imageShape = "Square"
+ default:
+ imageShape = "Portrait"
+ }
+
+ // 图片转base64
+ base64Bytes, err := writer.ToBase64(img)
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ data, err := web.PostData(postURL+"&shape="+imageShape, "text/plain", bytes.NewReader(base64Bytes))
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ sendAiImg(ctx, data)
+ })
+ engine.OnRegex(`^设置ai绘图配置\s(.*[^\s$])\s(.+)$`, zero.SuperUserPermission).SetBlock(true).
+ Handle(func(ctx *zero.Ctx) {
+ regexMatched := ctx.State["regex_matched"].([]string)
+ err := cfg.save(regexMatched[1], regexMatched[2])
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ ctx.SendChain(message.Text("成功设置server为", regexMatched[1], ", token为", regexMatched[2]))
+ })
+}
+
+func sendAiImg(ctx *zero.Ctx, data []byte) {
+ var loadData string
+ if predictRe.MatchString(binary.BytesToString(data)) {
+ loadData = predictRe.FindStringSubmatch(binary.BytesToString(data))[0]
+ }
+ var r result
+ if loadData != "" {
+ err := json.Unmarshal(binary.StringToBytes(loadData), &r)
+ if err != nil {
+ ctx.SendChain(message.Text("ERROR: ", err))
+ return
+ }
+ }
+ encodeStr := base64.StdEncoding.EncodeToString(data)
+ m := message.Message{ctxext.FakeSenderForwardNode(ctx, message.Image("base64://"+encodeStr))}
+ m = append(m, ctxext.FakeSenderForwardNode(ctx, message.Text(r.String())))
+ if id := ctx.Send(m).ID(); id == 0 {
+ ctx.SendChain(message.Text("ERROR: 可能被风控或下载图片用时过长，请耐心等待"))
+ }
+}
diff --git a/plugin/aipaint/config.go b/plugin/aipaint/config.go
new file mode 100644
index 0000000000..445a0cec2f
--- /dev/null
+++ b/plugin/aipaint/config.go
@@ -0,0 +1,54 @@
+package aipaint
+
+import (
+ "encoding/json"
+ "errors"
+ "os"
+
+ "github.com/FloatTech/floatbox/file"
+)
+
+// 配置结构体
+type serverConfig struct {
+ BaseURL string `json:"base_url"`
+ Token string `json:"token"`
+ file string
+}
+
+func newServerConfig(file string) *serverConfig {
+ return &serverConfig{
+ file: file,
+ }
+}
+
+func (cfg *serverConfig) save(baseURL, token string) (err error) {
+ cfg.BaseURL = baseURL
+ cfg.Token = token
+ reader, err := os.Create(cfg.file)
+ if err != nil {
+ return err
+ }
+ defer reader.Close()
+ return json.NewEncoder(reader).Encode(cfg)
+}
+
+func (cfg *serverConfig) load() (aipaintServer, token string, err error) {
+ if cfg.BaseURL != "" && cfg.Token != "" {
+ aipaintServer = cfg.BaseURL
+ token = cfg.Token
+ return
+ }
+ if file.IsNotExist(cfg.file) {
+ err = errors.New("no server config")
+ return
+ }
+ reader, err := os.Open(cfg.file)
+ if err != nil {
+ return
+ }
+ defer reader.Close()
+ err = json.NewDecoder(reader).Decode(cfg)
+ aipaintServer = cfg.BaseURL
+ token = cfg.Token
+ return
+}
diff --git a/plugin/aipaint/context.go b/plugin/aipaint/context.go
new file mode 100644
index 0000000000..30082e9c6e
--- /dev/null
+++ b/plugin/aipaint/context.go
@@ -0,0 +1,39 @@
+package aipaint
+
+import (
+ "os"
+ "strconv"
+ "strings"
+
+ "github.com/FloatTech/floatbox/file"
+)
+
+type context struct {
+ usrdir string
+ headimgsdir []string
+}
+
+func newContext(user int64) *context {
+ c := new(context)
+ c.usrdir = datapath + "users/" + strconv.FormatInt(user, 10) + `/`
+ _ = os.MkdirAll(c.usrdir, 0755)
+ c.headimgsdir = make([]string, 2)
+ c.headimgsdir[0] = c.usrdir + "0.gif"
+ c.headimgsdir[1] = c.usrdir + "1.gif"
+ return c
+}
+
+func (cc *context) prepareLogos(s ...string) error {
+ for i, v := range s {
+ _, err := strconv.Atoi(v)
+ if err != nil {
+ err = file.DownloadTo("https://gchat.qpic.cn/gchatpic_new//--"+strings.ToUpper(v)+"/0", cc.usrdir+strconv.Itoa(i)+".gif", true)
+ } else {
+ err = file.DownloadTo("http://q4.qlogo.cn/g?b=qq&nk="+v+"&s=640", cc.usrdir+strconv.Itoa(i)+".gif", true)
+ }
+ if err != nil {
+ return err
+ }
+ }
+ return nil
+}