| 23 | ) |
| 24 | |
| 25 | func init() { // 插件主体 |
| 26 | engine := control.AutoRegister(&ctrl.Options[*zero.Ctx]{ |
| 27 | DisableOnDefault: false, |
| 28 | Brief: "MagicPrompt-Stable-Diffusion吟唱提示", |
| 29 | Help: "- 吟唱提示 xxx", |
| 30 | PrivateDataFolder: "magicprompt", |
| 31 | }) |
| 32 | |
| 33 | // 开启 |
| 34 | engine.OnPrefixGroup([]string{`吟唱提示`, "吟唱补全"}).SetBlock(true). |
| 35 | Handle(func(ctx *zero.Ctx) { |
| 36 | _ctx, _cancel := context.WithTimeout(context.Background(), hf.TimeoutMax*time.Second) |
| 37 | defer _cancel() |
| 38 | ctx.SendChain(message.Text("少女祈祷中...")) |
| 39 | |
| 40 | magicpromptURL := fmt.Sprintf(hf.WssJoinPath, magicpromptRepo) |
| 41 | args := ctx.State["args"].(string) |
| 42 | c, _, err := websocket.DefaultDialer.Dial(magicpromptURL, nil) |
| 43 | if err != nil { |
| 44 | ctx.SendChain(message.Text("ERROR: ", err)) |
| 45 | return |
| 46 | } |
| 47 | defer c.Close() |
| 48 | |
| 49 | r := hf.PushRequest{ |
| 50 | FnIndex: 0, |
| 51 | Data: []interface{}{args}, |
| 52 | } |
| 53 | b, err := json.Marshal(r) |
| 54 | if err != nil { |
| 55 | ctx.SendChain(message.Text("ERROR: ", err)) |
| 56 | return |
| 57 | } |
| 58 | |
| 59 | err = c.WriteMessage(websocket.TextMessage, b) |
| 60 | if err != nil { |
| 61 | ctx.SendChain(message.Text("ERROR: ", err)) |
| 62 | return |
| 63 | } |
| 64 | t := time.NewTicker(time.Second * 1) |
| 65 | defer t.Stop() |
| 66 | for { |
| 67 | select { |
| 68 | case <-t.C: |
| 69 | _, data, err := c.ReadMessage() |
| 70 | if err != nil { |
| 71 | ctx.SendChain(message.Text("ERROR: ", err)) |
| 72 | return |
| 73 | } |
| 74 | j := gjson.ParseBytes(data) |
| 75 | if j.Get("msg").String() == hf.WssCompleteStatus { |
| 76 | m := message.Message{} |
| 77 | for _, v := range strings.Split(j.Get("output.data.0").String(), "\n\n") { |
| 78 | m = append(m, ctxext.FakeSenderForwardNode(ctx, message.Text(v))) |
| 79 | } |
| 80 | if id := ctx.Send(m).ID(); id == 0 { |
| 81 | ctx.SendChain(message.Text("ERROR: 可能被风控或下载图片用时过长，请耐心等待")) |
| 82 | } |