Feature huggingface (#478)

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