百度文心AI画图插件 (#442)

README.md

@@ -1083,6 +1083,26 @@ print("run[CQ:image,file="+j["img"]+"]")
 
   - [x] 来份网易云热评
 
+</details>
+<details>
+  <summary>百度文心AI画图</summary>
+
+  `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/wenxinAI"`
+
+  - 基于百度文心的免费AI画图插件(因为是免费的,图片质量你懂的)
+
+  - key申请链接：https://wenxin.baidu.com/moduleApi/key
+
+  - [x] 为[自己/本群/QQ号/群+群号]设置AI画图key [API Key] [Secret Key]
+
+  - 例：“为10086设置AI画图key 123 456”；“为群10010设置AI画图key 789 101”
+
+  - [x] [bot名称]画几张[图片描述]的[图片类型][图片尺寸]
+
+  - 指令示例：
+
+  - 椛椛帮我画几张金凤凰，背景绚烂，高饱和，古风，仙境，高清，4K，古风的油画方图
+
 </details>
 <details>
   <summary>月幕galgame图</summary>

main.go

@@ -121,6 +121,7 @@ import (
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/translation"   // 翻译
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/vtb_quotation" // vtb语录
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/wangyiyun"     // 网易云音乐热评
+	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/wenxinAI"      // 百度文心AI画图
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/word_count"    // 聊天热词
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/wordle"        // 猜单词
 	_ "github.com/FloatTech/ZeroBot-Plugin/plugin/ymgal"         // 月幕galgame

plugin/wenxinAI/ernievilg.go

@@ -0,0 +1,397 @@
+// Package ernie AI画图
+package ernie
+
+import (
+	"errors"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	fcext "github.com/FloatTech/floatbox/ctxext"
+	"github.com/FloatTech/floatbox/process"
+	ctrl "github.com/FloatTech/zbpctrl"
+	"github.com/FloatTech/zbputils/control"
+	"github.com/FloatTech/zbputils/ctxext"
+	zero "github.com/wdvxdr1123/ZeroBot"
+	"github.com/wdvxdr1123/ZeroBot/extension/single"
+	"github.com/wdvxdr1123/ZeroBot/message"
+
+	// 数据库
+	sql "github.com/FloatTech/sqlite"
+	// 百度文心AI画图API
+	wenxin "github.com/FloatTech/AnimeAPI/wenxinAI/ernievilg"
+)
+
+const (
+	serviceName = "AIdraw"
+	serviceErr  = "[" + serviceName + "]ERROR:\n"
+)
+
+type keydb struct {
+	db *sql.Sqlite
+	sync.RWMutex
+}
+
+// db内容
+type apikey struct {
+	ID         int64  // 群号
+	APIKey     string // API Key
+	SecretKey  string // Secret Key
+	Token      string // AccessToken
+	Updatetime int64  // token的有效时间
+	MaxLimit   int    // 总使用次数
+	DayLimit   int    // 当天的使用次数
+	Lasttime   string // 记录使用的时间，用于刷新使用次数
+}
+
+var (
+	groupinfo = &keydb{
+		db: &sql.Sqlite{},
+	}
+	limit = 50
+	dtype = [...]string{
+		"古风", "油画", "水彩画", "卡通画", "二次元", "浮世绘", "蒸汽波艺术", "low poly", "像素风格", "概念艺术", "未来主义", "赛博朋克", "写实风格", "洛丽塔风格", "巴洛克风格", "超现实主义",
+	}
+)
+
+func init() { // 插件主体
+	engine := control.Register(serviceName, &ctrl.Options[*zero.Ctx]{
+		DisableOnDefault: false,
+		Help: "AI画图\n" +
+			"基于百度文心的免费AI画图插件,\n因为是免费的,图片质量你懂的。\n" +
+			"key申请链接：https://wenxin.baidu.com/moduleApi/key\n" +
+			"注意：每个apikey每日上限50次,总上限500次请求。次数超过了请自行更新apikey\n" +
+			"- 为[自己/本群/QQ号/群+群号]设置AI画图key [API Key] [Secret Key]\n" +
+			"例：\n[为10086设置AI画图key 123 456]\n[为群10010设置AI画图key 789 101]\n" +
+			"- [bot名称]画几张[图片描述]的[图片类型][图片尺寸]\n" +
+			"————————————————————\n" +
+			"图片描述指南:\n图片主体，细节词(请用逗号连接)\n官方prompt指南:https://wenxin.baidu.com/wenxin/docs#Ol7ece95m\n" +
+			"————————————————————\n" +
+			"图片类型当前支持：" + strings.Join(dtype[:], "、") +
+			"\n————————————————————\n" +
+			"图片尺寸当前只支持：方图/长图/横图\n" +
+			"————————————————————\n" +
+			"指令示例：\n" +
+			"椛椛帮我画几张金凤凰，背景绚烂，高饱和，古风，仙境，高清，4K，古风的油画方图",
+		PrivateDataFolder: "ernievilg",
+	}).ApplySingle(single.New(
+		single.WithKeyFn(func(ctx *zero.Ctx) int64 { return ctx.Event.GroupID }),
+		single.WithPostFn[int64](func(ctx *zero.Ctx) {
+			ctx.Break()
+			ctx.Send(
+				message.ReplyWithMessage(ctx.Event.MessageID,
+					message.Text(zero.BotConfig.NickName[0], "正在给别人画图，请不要打扰哦"),
+				),
+			)
+		}),
+	))
+	getdb := fcext.DoOnceOnSuccess(func(ctx *zero.Ctx) bool {
+		groupinfo.db.DBPath = engine.DataFolder() + "keydb.db"
+		err := groupinfo.db.Open(time.Hour * 24)
+		if err != nil {
+			ctx.SendChain(message.Text(serviceErr, err))
+			return false
+		}
+		return true
+	})
+	// 画图
+	engine.OnRegex(`画几张(.*[^的$])的(.*[^\s$])(方图|长图|横图)$`, zero.OnlyToMe, getdb).SetBlock(true).
+		Handle(func(ctx *zero.Ctx) {
+			uid := -ctx.Event.UserID
+			gid := ctx.Event.GroupID
+			// 获取个人和群的key
+			userinfo, err1 := groupinfo.checkGroup(uid)
+			info, err2 := groupinfo.checkGroup(gid)
+			switch {
+			// 如果是个人请求且报错
+			case gid == 0 && err1 != nil:
+				ctx.SendChain(message.Text(serviceErr, err1))
+				return
+			// 如果群报错而个人没有,就切换成个人的
+			case err2 != nil && err1 == nil:
+				gid = uid
+				info = userinfo
+			// 如果都报错就以群为优先级
+			case err1 != nil && err2 != nil:
+				ctx.SendChain(message.Text(serviceErr, err2))
+				return
+			}
+			// 判断使用次数
+			check := false
+			switch {
+			// 群和个人都没有次数了
+			case info.DayLimit == 0 && userinfo.DayLimit == 0:
+				ctx.SendChain(message.Text("我已经画了", limit, "张了！我累了！不画不画，就不画！"))
+				return
+			// 个人还有次数的话
+			case info.DayLimit == 0 && userinfo.DayLimit != 0:
+				check = true
+			}
+			switch {
+			// 群和个人都没有总次数了
+			case info.MaxLimit == 0 && userinfo.MaxLimit == 0:
+				ctx.SendChain(message.Text("设置的key使用次数超过了限额，请更换key。"))
+				return
+			// 个人还有总次数的话
+			case info.MaxLimit == 0 && userinfo.MaxLimit != 0:
+				check = true
+			}
+			if check { // 如果只有个人有次数就切换回个人key
+				gid = uid
+				info = userinfo
+			}
+			// 创建任务
+			keyword := ctx.State["regex_matched"].([]string)[1]
+			if len([]rune(keyword)) >= 64 { // 描述不能超过64个字
+				ctx.SendChain(message.Text("要求太多了啦！减少点！"))
+				return
+			}
+			picType := ctx.State["regex_matched"].([]string)[2]
+			chooseSize := ctx.State["regex_matched"].([]string)[3]
+			wtime := 3
+			picSize := "1024*1024"
+			switch chooseSize {
+			case "长图":
+				wtime = 5
+				picSize = "1024*1536"
+			case "横图":
+				wtime = 5
+				picSize = "1536*1024"
+			}
+			taskID, err := wenxin.BuildWork(info.Token, keyword, picType, picSize)
+			if err != nil {
+				ctx.SendChain(message.Text(serviceErr, err))
+				return
+			}
+			if taskID < 1 {
+				ctx.SendChain(message.Text("要求太复杂力！想不出来..."))
+				return
+			}
+			// 开始画图
+			ctx.SendChain(message.Text(zero.BotConfig.NickName[0], "知道了，我可能需要", time.Duration(wtime*10)*time.Second, "左右才能画好哦，请等待..."))
+			i := 0
+			for range time.NewTicker(10 * time.Second).C {
+				// 等待 wtime * 10秒
+				i++
+				if i <= wtime {
+					continue
+				}
+				/*
+					if i > 60{// 十分钟还不出图就放弃
+						ctx.SendChain(message.Text("呜呜呜，要求太复杂力！画不出来..."))
+						return
+					}
+				// 获取结果*/
+				picURL, status, err := wenxin.GetPic(info.Token, taskID)
+				if err != nil {
+					ctx.SendChain(message.Text(serviceErr, err))
+					return
+				}
+				if status == "0" {
+					msg := message.Message{ctxext.FakeSenderForwardNode(ctx, message.Text("我画好了！"))}
+					for _, imginfo := range picURL {
+						msg = append(msg,
+							ctxext.FakeSenderForwardNode(ctx,
+								message.Image(imginfo.Image)))
+					}
+					if id := ctx.Send(msg).ID(); id == 0 {
+						ctx.SendChain(message.Text("ERROR: 可能被风控了"))
+					}
+					break
+				}
+			}
+			err = groupinfo.update(gid)
+			if err != nil {
+				ctx.SendChain(message.Text(serviceErr, err))
+			}
+			process.SleepAbout1sTo2s()
+			ctx.SendChain(message.Text("累死了，今天我最多只能画", info.DayLimit-1, "张图哦"))
+		})
+	engine.OnRegex(`^为(群)?(自己|本群|\d+)设置AI画图key\s(.*[^\s$])\s(.+)$`, getdb).SetBlock(true).
+		Handle(func(ctx *zero.Ctx) {
+			mode := ctx.State["regex_matched"].([]string)[1]
+			user := ctx.State["regex_matched"].([]string)[2]
+			aKey := ctx.State["regex_matched"].([]string)[3]
+			sKey := ctx.State["regex_matched"].([]string)[4]
+			dbID := -ctx.Event.UserID // 默认给自己
+			switch {
+			case mode != "": // 指定群的话
+				gid, err := strconv.ParseInt(user, 10, 64)
+				if err != nil {
+					ctx.SendChain(message.Text(serviceErr, err))
+					return
+				}
+				dbID = gid
+			case user == "本群": // 用于本群
+				gid := ctx.Event.GroupID
+				if gid == 0 {
+					ctx.SendChain(message.Text(serviceErr, "请指定群聊，或者使用指令；\n为群xxx设置AI画图key xxx xxx"))
+					return
+				}
+				dbID = gid
+			case user != "自己": // 给别人开key
+				uid, err := strconv.ParseInt(user, 10, 64)
+				if err != nil {
+					ctx.SendChain(message.Text(serviceErr, err))
+					return
+				}
+				dbID = -uid
+			}
+			err := groupinfo.insert(dbID, aKey, sKey)
+			if err != nil {
+				ctx.SendChain(message.Text(serviceErr, err))
+				return
+			}
+			ctx.SendChain(message.Text("成功!"))
+		})
+}
+
+// 登记group的key
+func (sql *keydb) insert(gid int64, akey, skey string) error {
+	sql.Lock()
+	defer sql.Unlock()
+	// 给db文件创建表格(没有才创建)，表格名称groupinfo，表格结构apikey
+	err := sql.db.Create("groupinfo", &apikey{})
+	if err != nil {
+		return err
+	}
+	// 获取group信息
+	groupinfo := apikey{} // 用于暂存数据
+	err = sql.db.Find("groupinfo", &groupinfo, "where ID is "+strconv.FormatInt(gid, 10))
+	if err != nil {
+		// 如果该group没有注册过
+		err = sql.db.Find("groupinfo", &groupinfo, "where APIKey is '"+akey+"' and SecretKey is '"+skey+"'")
+		if err == nil {
+			// 如果key存在过将当前的数据迁移过去
+			groupinfo.ID = gid
+		} else {
+			groupinfo = apikey{
+				ID:        gid,
+				APIKey:    akey,
+				SecretKey: skey,
+				MaxLimit:  500,
+			}
+		}
+		return sql.db.Insert("groupinfo", &groupinfo)
+	}
+	// 进行更新
+	groupinfo.APIKey = akey
+	groupinfo.SecretKey = skey
+	groupinfo.MaxLimit = 500
+	return sql.db.Insert("groupinfo", &groupinfo)
+}
+
+// 获取group信息
+func (sql *keydb) checkGroup(gid int64) (groupinfo apikey, err error) {
+	sql.Lock()
+	defer sql.Unlock()
+	// 给db文件创建表格(没有才创建)，表格名称groupinfo，表格结构apikey
+	err = sql.db.Create("groupinfo", &apikey{})
+	if err != nil {
+		return
+	}
+	// 先判断该群是否已经设置过key了
+	if ok := sql.db.CanFind("groupinfo", "where ID is "+strconv.FormatInt(gid, 10)); !ok {
+		if gid > 0 {
+			err = errors.New("该群没有设置过apikey，请前往https://wenxin.baidu.com/moduleApi/key获取key值后，发送指令:\n为本群设置AI画图key [API Key] [Secret Key]")
+		} else {
+			err = errors.New("你没有设置过apikey，请前往https://wenxin.baidu.com/moduleApi/key获取key值后，发送指令:\n为自己设置AI画图key [API Key] [Secret Key]")
+		}
+		return
+	}
+	// 获取group信息
+	err = sql.db.Find("groupinfo", &groupinfo, "where ID is "+strconv.FormatInt(gid, 10))
+	if err != nil {
+		return
+	}
+	// 如果隔天使用刷新次数
+	if time.Now().Format("2006/01/02") != groupinfo.Lasttime {
+		groupinfo.DayLimit = limit
+		groupinfo.Lasttime = time.Now().Format("2006/01/02")
+	}
+	if err = sql.db.Insert("groupinfo", &groupinfo); err != nil {
+		return
+	}
+	// 如果token有效期过期
+	if time.Since(time.Unix(groupinfo.Updatetime, 0)).Hours() > 24 {
+		token, err1 := wenxin.GetToken(groupinfo.APIKey, groupinfo.SecretKey)
+		if err1 != nil {
+			err = err1
+			return
+		}
+		groupinfo.Token = token
+		groupinfo.Updatetime = time.Now().Unix()
+		err = sql.db.Insert("groupinfo", &groupinfo)
+		if err == nil {
+			// 更新相同key的他人次数
+			condition := "where not ID is " + strconv.FormatInt(gid, 10) +
+				" and APIKey = '" + groupinfo.APIKey +
+				"' and SecretKey = '" + groupinfo.SecretKey + "'"
+			otherinfo := apikey{}
+			var groups []int64 // 将相同的key的ID暂存
+			// 无视没有找到相同的key的err
+			_ = sql.db.FindFor("groupinfo", &otherinfo, condition, func() error {
+				groups = append(groups, otherinfo.ID)
+				return nil
+			})
+			if len(groups) != 0 { // 如果有相同的key就更新
+				for _, group := range groups {
+					err = sql.db.Find("groupinfo", &otherinfo, "where ID is "+strconv.FormatInt(group, 10))
+					if err == nil {
+						otherinfo.Token = groupinfo.Token
+						otherinfo.Updatetime = groupinfo.Updatetime
+						err = sql.db.Insert("groupinfo", &otherinfo)
+					}
+				}
+			}
+		}
+	}
+	return
+}
+
+// 记录次数(-1)
+func (sql *keydb) update(gid int64) error {
+	sql.Lock()
+	defer sql.Unlock()
+	// 给db文件创建表格(没有才创建)，表格名称groupinfo，表格结构apikey
+	err := sql.db.Create("groupinfo", &apikey{})
+	if err != nil {
+		return err
+	}
+	groupinfo := apikey{} // 用于暂存数据
+	// 获取group信息
+	err = sql.db.Find("groupinfo", &groupinfo, "where ID is "+strconv.FormatInt(gid, 10))
+	if err != nil {
+		return err
+	}
+	groupinfo.MaxLimit--
+	groupinfo.DayLimit--
+	err = sql.db.Insert("groupinfo", &groupinfo)
+	if err != nil {
+		return err
+	}
+	// 更新相同key的他人次数
+	condition := "where not ID is " + strconv.FormatInt(gid, 10) +
+		" and APIKey = '" + groupinfo.APIKey +
+		"' and SecretKey = '" + groupinfo.SecretKey + "'"
+	otherinfo := apikey{}
+	var groups []int64 // 将相同的key的ID暂存
+	// 无视没有找到相同的key的err
+	_ = sql.db.FindFor("groupinfo", &otherinfo, condition, func() error {
+		groups = append(groups, otherinfo.ID)
+		return nil
+	})
+	if len(groups) != 0 { // 如果有相同的key就更新
+		for _, group := range groups {
+			err = sql.db.Find("groupinfo", &otherinfo, "where ID is "+strconv.FormatInt(group, 10))
+			if err == nil {
+				otherinfo.MaxLimit = groupinfo.MaxLimit
+				otherinfo.DayLimit = groupinfo.DayLimit
+				otherinfo.Lasttime = groupinfo.Lasttime
+				err = sql.db.Insert("groupinfo", &otherinfo)
+			}
+		}
+	}
+	return err
+}