From deb655de3635307e376806de368c13d0bd18510f Mon Sep 17 00:00:00 2001 From: himawari <54976075+guohuiyuan@users.noreply.github.com> Date: Sat, 5 Nov 2022 17:32:31 +0800 Subject: [PATCH] Feature huggingface (#478) --- README.md | 72 ++++++++++----- go.mod | 4 +- go.sum | 4 +- main.go | 3 + plugin/aipaint/aipaint.go | 72 +++++++++++---- plugin/aipaint/config.go | 26 +++--- plugin/magicprompt/magicprompt.go | 91 +++++++++++++++++++ plugin/realcugan/realcugan.go | 105 ++++++++++++++++++++++ plugin/vitsnyaru/vitsnyaru.go | 141 ++++++++++++++++++++++++++++++ 9 files changed, 463 insertions(+), 55 deletions(-) create mode 100644 plugin/magicprompt/magicprompt.go create mode 100644 plugin/realcugan/realcugan.go create mode 100644 plugin/vitsnyaru/vitsnyaru.go diff --git a/README.md b/README.md index 08042ebd..d950b81d 100644 --- a/README.md +++ b/README.md @@ -368,6 +368,10 @@ print("run[CQ:image,file="+j["img"]+"]") - [x] 设置ai绘图配置 [server] [token] + - [x] 设置ai绘图撤回时间90s + + - [x] 查看ai绘图配置 + 例: 设置ai绘图配置 http://91.216.169.75:5010 abc 参考服务器 http://91.217.139.190:5010, http://91.216.169.75:5010, http://185.80.202.180:5010 @@ -860,6 +864,14 @@ print("run[CQ:image,file="+j["img"]+"]") 来份萝莉 ``` + +
+ MagicPrompt-Stable-Diffusion吟唱提示 + + `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/magicprompt"` + + - [x] 吟唱提示[xxxx] +
简易midi音乐制作 @@ -1045,6 +1057,14 @@ print("run[CQ:image,file="+j["img"]+"]") - [x] 重置花名册 +
+
+ Real-CUGAN清晰术 + + `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/realcugan"` + + - [x] 清晰术(双重吟唱|三重吟唱|四重吟唱)(强力术式|中等术式|弱术式|不变式|原式)[图片] +
投胎 @@ -1142,6 +1162,14 @@ print("run[CQ:image,file="+j["img"]+"]") - [x] 解塔罗牌[牌名] - [x] [塔罗|大阿卡纳|小阿卡纳|混合]牌阵[圣三角|时间之流|四要素|五牌阵|吉普赛十字|马蹄|六芒星] +
+
+ 舔狗日记 + + `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/tiangou"` + + - [x] 舔狗日记 +
搜番 @@ -1150,20 +1178,6 @@ print("run[CQ:image,file="+j["img"]+"]") - [x] 搜番 | 搜索番剧[图片] -
-
- 猜单词 - - `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/wordle"` - - - [x] 个人猜单词 - - - [x] 团队猜单词 - - - [x] 团队六阶猜单词 - - - [x] 团队七阶猜单词 -
翻译 @@ -1172,6 +1186,14 @@ print("run[CQ:image,file="+j["img"]+"]") - [x] >TL 你好 +
+
+ vits猫雷 + + `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/vitsnyaru"` + + - [x] 让猫雷说[xxxx] +
vtb语录 @@ -1252,14 +1274,6 @@ print("run[CQ:image,file="+j["img"]+"]") - [x] 更新gal -
-
- 舔狗日记 - - `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/tiangou"` - - - [x] 舔狗日记 -
聊天热词 @@ -1268,6 +1282,20 @@ print("run[CQ:image,file="+j["img"]+"]") - [x] 热词 [群号] [消息数目]|热词 123456 1000 +
+
+ 猜单词 + + `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin/wordle"` + + - [x] 个人猜单词 + + - [x] 团队猜单词 + + - [x] 团队六阶猜单词 + + - [x] 团队七阶猜单词 +
鬼东西 diff --git a/go.mod b/go.mod index 852de31e..2fb5c1df 100644 --- a/go.mod +++ b/go.mod @@ -5,13 +5,14 @@ go 1.19 require ( github.com/Baidu-AIP/golang-sdk v1.1.1 github.com/Coloured-glaze/gg v1.3.4 - github.com/FloatTech/AnimeAPI v1.5.2-0.20221101144606-ec2f4e59d0ef + github.com/FloatTech/AnimeAPI v1.5.2-0.20221105044443-0c9004b2f051 github.com/FloatTech/floatbox v0.0.0-20221029160423-446812ec82d9 github.com/FloatTech/sqlite v0.5.0 github.com/FloatTech/ttl v0.0.0-20220715042055-15612be72f5b github.com/FloatTech/zbpctrl v1.5.2 github.com/FloatTech/zbputils v1.5.1-0.20221101032142-05b2a4825567 github.com/RomiChan/syncx v0.0.0-20220404072119-d7ea0ae15a4c + github.com/RomiChan/websocket v1.4.3-0.20220123145318-307a86b127bc github.com/antchfx/htmlquery v1.2.5 github.com/corona10/goimagehash v1.1.0 github.com/fumiama/ahsai v0.1.0 @@ -38,7 +39,6 @@ require ( ) require ( - github.com/RomiChan/websocket v1.4.3-0.20220123145318-307a86b127bc // indirect github.com/antchfx/xpath v1.2.1 // indirect github.com/disintegration/imaging v1.6.2 // indirect github.com/ericpauley/go-quantize v0.0.0-20200331213906-ae555eb2afa4 // indirect diff --git a/go.sum b/go.sum index c91b427b..77c4b507 100644 --- a/go.sum +++ b/go.sum @@ -4,8 +4,8 @@ github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym github.com/Coloured-glaze/gg v1.3.4 h1:l31zIF/HaVwkzjrj+A56RGQoSKyKuR1IWtIrqXGFStI= github.com/Coloured-glaze/gg v1.3.4/go.mod h1:Ih5NLNNDHOy3RJbB0EPqGTreIzq/H02TGThIagh8HJg= github.com/DATA-DOG/go-sqlmock v1.3.3/go.mod h1:f/Ixk793poVmq4qj/V1dPUg2JEAKC73Q5eFN3EC/SaM= -github.com/FloatTech/AnimeAPI v1.5.2-0.20221101144606-ec2f4e59d0ef h1:IO4ZSaTrRTidAXqUH86dHgzk8vRlw2+YX0Oz+4n/oeg= -github.com/FloatTech/AnimeAPI v1.5.2-0.20221101144606-ec2f4e59d0ef/go.mod h1:7ZXHFEgAOFvfbgIwgyFl8G/ISnUl1KvWBzngoCpQVLg= +github.com/FloatTech/AnimeAPI v1.5.2-0.20221105044443-0c9004b2f051 h1:NEW8HzBNybMXAV0NrDpMF31n/e9BFGvNiGkDPGSkouc= +github.com/FloatTech/AnimeAPI v1.5.2-0.20221105044443-0c9004b2f051/go.mod h1:Z+Q4kIPNo/OX4RWw6WGQOQcaNsbkv/wPmPDw8p4aQIY= github.com/FloatTech/floatbox v0.0.0-20221029160423-446812ec82d9 h1:HYJ7lwaqaOKmbYooPUMWxMhXRTp+JItoyeqa320ARD4= github.com/FloatTech/floatbox v0.0.0-20221029160423-446812ec82d9/go.mod h1:w+ND28mRaJvxUJ6pRXS6i4cLzutpXsWyroutCzBdL78= github.com/FloatTech/sqlite v0.5.0 h1:U7J5Omc534PqmH6csfu+ypCo3DS8L91l5lTsxUu3b/U= diff --git a/main.go b/main.go index 0dae3b39..11bba083 100644 --- a/main.go +++ b/main.go @@ -99,6 +99,7 @@ import ( _ "github.com/FloatTech/ZeroBot-Plugin/plugin/jptingroom" // 日语听力学习材料 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/juejuezi" // 绝绝子生成器 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/lolicon" // lolicon 随机图片 + _ "github.com/FloatTech/ZeroBot-Plugin/plugin/magicprompt" // magicprompt吟唱提示 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/midicreate" // 简易midi音乐制作 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/moegoe" // 日韩 VITS 模型拟声 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/moyu" // 摸鱼 @@ -112,6 +113,7 @@ import ( _ "github.com/FloatTech/ZeroBot-Plugin/plugin/nsfw" // nsfw图片识别 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/omikuji" // 浅草寺求签 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/qqwife" // 一群一天一夫一妻制群老婆 + _ "github.com/FloatTech/ZeroBot-Plugin/plugin/realcugan" // realcugan清晰术 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/reborn" // 投胎 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/runcode" // 在线运行代码 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/saucenao" // 以图搜图 @@ -124,6 +126,7 @@ import ( _ "github.com/FloatTech/ZeroBot-Plugin/plugin/tiangou" // 舔狗日记 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/tracemoe" // 搜番 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/translation" // 翻译 + _ "github.com/FloatTech/ZeroBot-Plugin/plugin/vitsnyaru" // vits猫雷 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/vtb_quotation" // vtb语录 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/wangyiyun" // 网易云音乐热评 _ "github.com/FloatTech/ZeroBot-Plugin/plugin/wenxinAI" // 百度文心AI画图 diff --git a/plugin/aipaint/aipaint.go b/plugin/aipaint/aipaint.go index e1daab2f..dffa615c 100644 --- a/plugin/aipaint/aipaint.go +++ b/plugin/aipaint/aipaint.go @@ -56,6 +56,8 @@ func init() { // 插件主体 "- [ ai高级绘图 | 高级生成色图 | 高级生成涩图 | ai高级画图 ] [prompt]\n" + "- [ 以图绘图 | 以图生图 | 以图画图 ] xxx [图片]|@xxx|[qq号]\n" + "- 设置ai绘图配置 [server] [token]\n" + + "- 设置ai绘图撤回时间90s\n" + + "- 查看ai绘图配置\n" + "例: 设置ai绘图配置 http://91.217.139.190:5010 abc\n" + "参考服务器 http://91.217.139.190:5010, http://91.216.169.75:5010, http://185.80.202.180:5010\n" + "通过 http://91.217.139.190:5010/token 获取token\n" + @@ -67,23 +69,23 @@ func init() { // 插件主体 datapath = file.BOTPATH + "/" + engine.DataFolder() engine.OnPrefixGroup([]string{`ai绘图`, `生成色图`, `生成涩图`, `ai画图`}).SetBlock(true). Handle(func(ctx *zero.Ctx) { - server, token, err := cfg.load() + 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"))))) + data, err := web.GetData(cfg.BaseURL + fmt.Sprintf(aipaintTxt2ImgURL, cfg.Token, url.QueryEscape(strings.TrimSpace(strings.ReplaceAll(args, " ", "%20"))))) if err != nil { ctx.SendChain(message.Text("ERROR: ", err)) return } - sendAiImg(ctx, data) + sendAiImg(ctx, data, cfg.Interval) }) 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() + err := cfg.load() if err != nil { ctx.SendChain(message.Text("ERROR: ", err)) return @@ -101,7 +103,7 @@ func init() { // 插件主体 return } ctx.SendChain(message.Text("少女祈祷中...")) - postURL := server + fmt.Sprintf(aipaintImg2ImgURL, token, url.QueryEscape(strings.TrimSpace(strings.ReplaceAll(args, " ", "%20")))) + postURL := cfg.BaseURL + fmt.Sprintf(aipaintImg2ImgURL, cfg.Token, url.QueryEscape(strings.TrimSpace(strings.ReplaceAll(args, " ", "%20")))) f, err := os.Open(c.headimgsdir[0]) if err != nil { @@ -136,11 +138,11 @@ func init() { // 插件主体 ctx.SendChain(message.Text("ERROR: ", err)) return } - sendAiImg(ctx, data) + sendAiImg(ctx, data, cfg.Interval) }) engine.OnPrefixGroup([]string{`ai高级绘图`, `高级生成色图`, `高级生成涩图`, `ai高级画图`}).SetBlock(true). Handle(func(ctx *zero.Ctx) { - server, token, err := cfg.load() + err := cfg.load() if err != nil { ctx.SendChain(message.Text("ERROR: ", err)) return @@ -156,13 +158,13 @@ func init() { // 插件主体 if len(value) > 1 { if value[0] == "R18" && value[1] == "1" { value[1] = "0" - ctx.SendChain(message.Text("不准涩涩!已将R18设置为0。")) + ctx.SendChain(message.Text("不准涩涩! 已将R18设置为0. ")) } tags[value[0]] = strings.Join(value[1:], ":") } } ctx.SendChain(message.Text("少女祈祷中...")) - apiurl := "/got_image?token=" + token + apiurl := "/got_image?token=" + cfg.Token if _, ok := tags["tags"]; ok { apiurl += "&tags=" + url.QueryEscape(strings.ReplaceAll(strings.TrimSpace(tags["tags"]), " ", "%20")) } @@ -181,26 +183,63 @@ func init() { // 插件主体 if _, ok := tags["seed"]; ok { apiurl += "&seed=" + url.QueryEscape(strings.TrimSpace(tags["seed"])) } - data, err := web.GetData(server + apiurl) + data, err := web.GetData(cfg.BaseURL + apiurl) if err != nil { ctx.SendChain(message.Text("ERROR: ", err)) return } - sendAiImg(ctx, data) + sendAiImg(ctx, data, cfg.Interval) }) 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]) + err := cfg.load() if err != nil { ctx.SendChain(message.Text("ERROR: ", err)) return } - ctx.SendChain(message.Text("成功设置server为", regexMatched[1], ", token为", regexMatched[2])) + err = cfg.update(regexMatched[1], regexMatched[2], cfg.Interval) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + text := fmt.Sprintf("成功设置\nbase_url: %v\ntoken: %v\ninterval: %v\n", cfg.BaseURL, cfg.Token, cfg.Interval) + ctx.SendChain(message.Text(text)) + }) + engine.OnRegex(`^设置ai绘图撤回时间(\d{1,3})s$`, zero.SuperUserPermission).SetBlock(true). + Handle(func(ctx *zero.Ctx) { + regexMatched := ctx.State["regex_matched"].([]string) + interval, err := strconv.Atoi(regexMatched[1]) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + err = cfg.load() + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + err = cfg.update(cfg.BaseURL, cfg.Token, interval) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + text := fmt.Sprintf("成功设置\nbase_url: %v\ntoken: %v\ninterval: %v\n", cfg.BaseURL, cfg.Token, cfg.Interval) + ctx.SendChain(message.Text(text)) + }) + engine.OnFullMatch(`查看ai绘图配置`, zero.SuperUserPermission).SetBlock(true). + Handle(func(ctx *zero.Ctx) { + err := cfg.load() + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + text := fmt.Sprintf("base_url: %v\ntoken: %v\ninterval: %v\n", cfg.BaseURL, cfg.Token, cfg.Interval) + ctx.SendChain(message.Text(text)) }) } -func sendAiImg(ctx *zero.Ctx, data []byte) { +func sendAiImg(ctx *zero.Ctx, data []byte, interval int) { var loadData string if predictRe.MatchString(binary.BytesToString(data)) { loadData = predictRe.FindStringSubmatch(binary.BytesToString(data))[0] @@ -223,11 +262,10 @@ func sendAiImg(ctx *zero.Ctx, data []byte) { m = append(m, ctxext.FakeSenderForwardNode(ctx, message.Text(r.String()))) if mid := ctx.Send(m); mid.ID() == 0 { ctx.SendChain(message.Text("ERROR: 可能被风控或下载图片用时过长，请耐心等待")) - } else { + } else if interval > 0 { go func(i message.MessageID) { - time.Sleep(90 * time.Second) + time.Sleep(time.Duration(interval) * time.Second) ctx.DeleteMessage(i) }(mid) } - } diff --git a/plugin/aipaint/config.go b/plugin/aipaint/config.go index 445a0cec..12f80967 100644 --- a/plugin/aipaint/config.go +++ b/plugin/aipaint/config.go @@ -10,9 +10,10 @@ import ( // 配置结构体 type serverConfig struct { - BaseURL string `json:"base_url"` - Token string `json:"token"` - file string + BaseURL string `json:"base_url"` + Token string `json:"token"` + Interval int `json:"interval"` + file string } func newServerConfig(file string) *serverConfig { @@ -21,9 +22,14 @@ func newServerConfig(file string) *serverConfig { } } -func (cfg *serverConfig) save(baseURL, token string) (err error) { - cfg.BaseURL = baseURL - cfg.Token = token +func (cfg *serverConfig) update(baseURL, token string, interval int) (err error) { + if baseURL != "" { + cfg.BaseURL = baseURL + } + if token != "" { + cfg.Token = token + } + cfg.Interval = interval reader, err := os.Create(cfg.file) if err != nil { return err @@ -32,10 +38,8 @@ func (cfg *serverConfig) save(baseURL, token string) (err error) { 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 +func (cfg *serverConfig) load() (err error) { + if cfg.BaseURL != "" && cfg.Token != "" && cfg.Interval != 0 { return } if file.IsNotExist(cfg.file) { @@ -48,7 +52,5 @@ func (cfg *serverConfig) load() (aipaintServer, token string, err error) { } defer reader.Close() err = json.NewDecoder(reader).Decode(cfg) - aipaintServer = cfg.BaseURL - token = cfg.Token return } diff --git a/plugin/magicprompt/magicprompt.go b/plugin/magicprompt/magicprompt.go new file mode 100644 index 00000000..1f0d4d32 --- /dev/null +++ b/plugin/magicprompt/magicprompt.go @@ -0,0 +1,91 @@ +// Package magicprompt MagicPrompt-Stable-Diffusion吟唱提示 +package magicprompt + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + hf "github.com/FloatTech/AnimeAPI/huggingface" + ctrl "github.com/FloatTech/zbpctrl" + "github.com/FloatTech/zbputils/control" + "github.com/FloatTech/zbputils/ctxext" + "github.com/RomiChan/websocket" + "github.com/tidwall/gjson" + zero "github.com/wdvxdr1123/ZeroBot" + "github.com/wdvxdr1123/ZeroBot/message" +) + +const ( + magicpromptRepo = "Gustavosta/MagicPrompt-Stable-Diffusion" +) + +func init() { // 插件主体 + engine := control.Register("magicprompt", &ctrl.Options[*zero.Ctx]{ + DisableOnDefault: false, + Brief: "MagicPrompt-Stable-Diffusion吟唱提示", + Help: "- 吟唱提示 xxx", + PrivateDataFolder: "magicprompt", + }) + + // 开启 + engine.OnPrefixGroup([]string{`吟唱提示`, "吟唱补全"}).SetBlock(true). + Handle(func(ctx *zero.Ctx) { + _ctx, _cancel := context.WithTimeout(context.Background(), hf.TimeoutMax*time.Second) + defer _cancel() + ctx.SendChain(message.Text("少女祈祷中...")) + + magicpromptURL := fmt.Sprintf(hf.WssJoinPath, magicpromptRepo) + args := ctx.State["args"].(string) + c, _, err := websocket.DefaultDialer.Dial(magicpromptURL, nil) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + defer c.Close() + + r := hf.PushRequest{ + FnIndex: 0, + Data: []interface{}{args}, + } + b, err := json.Marshal(r) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + + err = c.WriteMessage(websocket.TextMessage, b) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + t := time.NewTicker(time.Second * 1) + defer t.Stop() + for { + select { + case <-t.C: + _, data, err := c.ReadMessage() + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + j := gjson.ParseBytes(data) + if j.Get("msg").String() == hf.WssCompleteStatus { + m := message.Message{} + for _, v := range strings.Split(j.Get("output.data.0").String(), "\n\n") { + m = append(m, ctxext.FakeSenderForwardNode(ctx, message.Text(v))) + } + if id := ctx.Send(m).ID(); id == 0 { + ctx.SendChain(message.Text("ERROR: 可能被风控或下载图片用时过长，请耐心等待")) + } + return + } + case <-_ctx.Done(): + ctx.SendChain(message.Text("ERROR: 吟唱提示指令超时")) + return + } + } + }) +} diff --git a/plugin/realcugan/realcugan.go b/plugin/realcugan/realcugan.go new file mode 100644 index 00000000..0c3f65d5 --- /dev/null +++ b/plugin/realcugan/realcugan.go @@ -0,0 +1,105 @@ +// Package realcugan Real-CUGAN清晰术 +package realcugan + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "image" + "strings" + + hf "github.com/FloatTech/AnimeAPI/huggingface" + "github.com/FloatTech/floatbox/web" + ctrl "github.com/FloatTech/zbpctrl" + "github.com/FloatTech/zbputils/control" + "github.com/FloatTech/zbputils/ctxext" + "github.com/tidwall/gjson" + zero "github.com/wdvxdr1123/ZeroBot" + "github.com/wdvxdr1123/ZeroBot/message" +) + +const ( + realcuganRepo = "shichen1231/Real-CUGAN" +) + +func init() { // 插件主体 + engine := control.Register("realcugan", &ctrl.Options[*zero.Ctx]{ + DisableOnDefault: false, + Brief: "Real-CUGAN清晰术", + Help: "- 清晰术(双重吟唱|三重吟唱|四重吟唱)(强力术式|中等术式|弱术式|不变式|原式)[图片]", + PrivateDataFolder: "realcugan", + }) + engine.OnPrefix("清晰术", zero.MustProvidePicture).SetBlock(true). + Handle(func(ctx *zero.Ctx) { + ctx.SendChain(message.Text("少女祈祷中...")) + realcuganURL := fmt.Sprintf(hf.HTTPSPredictPath, realcuganRepo) + for _, url := range ctx.State["image_url"].([]string) { + imgdata, err := web.GetData(url) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + img, _, err := image.Decode(bytes.NewReader(imgdata)) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + // 初始化参数 + var ( + fashu = ctx.Event.Message.ExtractPlainText() + scale = 2 + con = "conservative" + ) + switch { + case strings.Contains(fashu, "双重吟唱"): + scale = 2 + case strings.Contains(fashu, "三重吟唱") && img.Bounds().Dx()*img.Bounds().Dy() < 400000: + scale = 3 + case strings.Contains(fashu, "四重吟唱") && img.Bounds().Dx()*img.Bounds().Dy() < 400000: + scale = 4 + } + switch { + case strings.Contains(fashu, "强力术式"): + con = "denoise3x" + case strings.Contains(fashu, "中等术式"): + con = "no-denoise" + if scale == 2 { + con = "denoise2x" + } + case strings.Contains(fashu, "弱术式"): + con = "no-denoise" + if scale == 2 { + con = "denoise1x" + } + case strings.Contains(fashu, "不变式"): + con = "no-denoise" + case strings.Contains(fashu, "原式"): + con = "conservative" + } + modelname := fmt.Sprintf("up%vx-latest-%v.pth", scale, con) + encodeStr := base64.StdEncoding.EncodeToString(imgdata) + encodeStr = "data:image/jpeg;base64," + encodeStr + pr := hf.PushRequest{ + Data: []interface{}{encodeStr, modelname, 2}, + } + buf := bytes.NewBuffer([]byte{}) + err = json.NewEncoder(buf).Encode(pr) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + data, err := web.PostData(realcuganURL, "application/json", buf) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + imgStr := gjson.ParseBytes(data).Get("data.0").String() + m := message.Message{ctxext.FakeSenderForwardNode(ctx, message.Text(scale, "重唱", con, "分支大清晰术!")), + ctxext.FakeSenderForwardNode(ctx, message.Image("base64://"+strings.TrimPrefix(imgStr, "data:image/png;base64,")))} + if id := ctx.Send(m).ID(); id == 0 { + ctx.SendChain(message.Text("ERROR: 可能被风控或下载图片用时过长，请耐心等待")) + } + } + }) +} diff --git a/plugin/vitsnyaru/vitsnyaru.go b/plugin/vitsnyaru/vitsnyaru.go new file mode 100644 index 00000000..a729e82d --- /dev/null +++ b/plugin/vitsnyaru/vitsnyaru.go @@ -0,0 +1,141 @@ +// Package vitsnyaru vits猫雷 +package vitsnyaru + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + hf "github.com/FloatTech/AnimeAPI/huggingface" + ctrl "github.com/FloatTech/zbpctrl" + "github.com/FloatTech/zbputils/control" + "github.com/tidwall/gjson" + zero "github.com/wdvxdr1123/ZeroBot" + "github.com/wdvxdr1123/ZeroBot/message" +) + +const ( + vitsnyaruRepo = "innnky/vits-nyaru" +) + +func init() { // 插件主体 + engine := control.Register("vitsnyaru", &ctrl.Options[*zero.Ctx]{ + DisableOnDefault: false, + Brief: "vits猫雷", + Help: "- 让猫雷说 xxx", + PrivateDataFolder: "vitsnyaru", + }) + + // 开启 + engine.OnPrefix(`让猫雷说`).SetBlock(true). + Handle(func(ctx *zero.Ctx) { + _ctx, _cancel := context.WithTimeout(context.Background(), hf.TimeoutMax*time.Second) + defer _cancel() + ch := make(chan []byte, 1) + + args := ctx.State["args"].(string) + pushURL := fmt.Sprintf(hf.HTTPSPushPath, vitsnyaruRepo) + statusURL := fmt.Sprintf(hf.HTTPSStatusPath, vitsnyaruRepo) + ctx.SendChain(message.Text("少女祈祷中...")) + var ( + pushReq hf.PushRequest + pushRes hf.PushResponse + statusReq hf.StatusRequest + statusRes hf.StatusResponse + data []byte + ) + + // 获取clean后的文本 + pushReq = hf.PushRequest{ + Action: hf.DefaultAction, + Data: []interface{}{args}, + FnIndex: 1, + } + pushRes, err := hf.Push(pushURL, pushReq) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + statusReq = hf.StatusRequest{ + Hash: pushRes.Hash, + } + + t := time.NewTicker(time.Second * 1) + defer t.Stop() + LOOP: + for { + select { + case <-t.C: + data, err = hf.Status(statusURL, statusReq) + if err != nil { + ch <- data + break LOOP + } + if gjson.ParseBytes(data).Get("status").String() == hf.CompleteStatus { + ch <- data + break LOOP + } + case <-_ctx.Done(): + ch <- data + break LOOP + } + } + + data = <-ch + err = json.Unmarshal(data, &statusRes) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + + // 用clean的文本预测语音 + pushReq = hf.PushRequest{ + Action: hf.DefaultAction, + Data: statusRes.Data.Data, + FnIndex: 2, + } + pushRes, err = hf.Push(pushURL, pushReq) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + statusReq = hf.StatusRequest{ + Hash: pushRes.Hash, + } + + LOOP2: + for { + select { + case <-t.C: + data, err = hf.Status(statusURL, statusReq) + if err != nil { + ch <- data + break LOOP2 + } + if gjson.ParseBytes(data).Get("status").String() == hf.CompleteStatus { + ch <- data + break LOOP2 + } + case <-_ctx.Done(): + ch <- data + break LOOP2 + } + } + + data = <-ch + err = json.Unmarshal(data, &statusRes) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + return + } + + // 发送语音 + if len(statusRes.Data.Data) < 2 { + ctx.SendChain(message.Text("ERROR: 未能获取语音")) + return + } + ctx.SendChain(message.Record("base64://" + strings.TrimPrefix(statusRes.Data.Data[1].(string), "data:audio/wav;base64,"))) + }) +}