mirror of
https://github.com/FloatTech/ZeroBot-Plugin.git
synced 2025-12-19 13:59:39 +08:00
92 lines
2.4 KiB
Go
92 lines
2.4 KiB
Go
// 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
|
|
}
|
|
}
|
|
})
|
|
}
|