Feature aipaint (#450)

README.md

@@ -346,6 +346,24 @@ print("run[CQ:image,file="+j["img"]+"]")
 
   - [x] 设置默认限速为每 m [分钟 | 秒] n 次触发
 
+</details>
+<details>
+  <summary>ai绘图</summary>
+
+  `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
+
 </details>
 <details>
   <summary>AIWife</summary>

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加解密

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: 可能被风控或下载图片用时过长，请耐心等待"))
+	}
+}

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
+}

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
+}