init
This commit is contained in:
HfCloud
commit
9a0f4e4a3e
|
|
@ -0,0 +1,5 @@
|
||||||
|
storage.json
|
||||||
|
*.png
|
||||||
|
.vscode
|
||||||
|
myconf.json
|
||||||
|
.DS_Store
|
||||||
|
|
@ -0,0 +1,139 @@
|
||||||
|
from distutils.command.config import config
|
||||||
|
import flask
|
||||||
|
import requests
|
||||||
|
import json
|
||||||
|
import openai
|
||||||
|
import re
|
||||||
|
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
|
||||||
|
import torch
|
||||||
|
import argparse
|
||||||
|
|
||||||
|
ps = argparse.ArgumentParser()
|
||||||
|
ps.add_argument("--config", default="config.json", help="Configuration file")
|
||||||
|
args = ps.parse_args()
|
||||||
|
|
||||||
|
with open(args.config) as f:
|
||||||
|
config_json = json.load(f)
|
||||||
|
|
||||||
|
class GlobalData:
|
||||||
|
# OPENAI_ORGID = config_json[""]
|
||||||
|
OPENAI_APIKEY = config_json["OpenAI-API-Key"]
|
||||||
|
OPENAI_MODEL = config_json["GPT-Model"]
|
||||||
|
OPENAI_MODEL_TEMPERATURE = 0.66
|
||||||
|
OPENAI_MODEL_MAXTOKENS = 2048
|
||||||
|
|
||||||
|
context_for_users = {}
|
||||||
|
context_for_groups = {}
|
||||||
|
|
||||||
|
GENERATE_PICTURE_ARG_PAT = re.compile("(\(|（)([0-9]+)[ \n\t]+([0-9]+)[ \n\t]+([0-9]+)(\)|）)")
|
||||||
|
GENERATE_PICTURE_NEG_PROMPT_DELIMETER = re.compile("\n+")
|
||||||
|
GENERATE_PICTURE_MAX_ITS = 200 #最大迭代次数
|
||||||
|
|
||||||
|
app = flask.Flask(__name__)
|
||||||
|
sd_pipe = StableDiffusionPipeline.from_pretrained(config_json["Diffusion-Model"], torch_dtype=torch.float32)
|
||||||
|
sd_pipe.scheduler = DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config)
|
||||||
|
if torch.backends.mps.is_available():
|
||||||
|
sd_pipe = sd_pipe.to("mps")
|
||||||
|
elif torch.cuda.is_available():
|
||||||
|
sd_pipe = sd_pipe.to("cuda")
|
||||||
|
|
||||||
|
def send_text_to_user(user_id : str, text : str):
|
||||||
|
requests.post(url="http://localhost:11110/send_text",
|
||||||
|
data=json.dumps({"user_id" : user_id, "text" : text, "in_group" : False}))
|
||||||
|
|
||||||
|
def call_gpt(prompt : str):
|
||||||
|
try:
|
||||||
|
res = openai.Completion.create(
|
||||||
|
model=GlobalData.OPENAI_MODEL,
|
||||||
|
prompt=prompt,
|
||||||
|
max_tokens=GlobalData.OPENAI_MODEL_MAXTOKENS,
|
||||||
|
temperature=GlobalData.OPENAI_MODEL_TEMPERATURE)
|
||||||
|
if len(res["choices"]) > 0:
|
||||||
|
return res["choices"][0]["text"].strip()
|
||||||
|
else:
|
||||||
|
return ""
|
||||||
|
except:
|
||||||
|
return "上下文长度超出模型限制，请对我说\“重置上下文\"，然后再试一次"
|
||||||
|
|
||||||
|
@app.route("/chat_clear", methods=['POST'])
|
||||||
|
def app_chat_clear():
|
||||||
|
data = json.loads(flask.globals.request.get_data())
|
||||||
|
GlobalData.context_for_users[data["user_id"]] = ""
|
||||||
|
print(f"Cleared context for {data['user_id']}")
|
||||||
|
return ""
|
||||||
|
|
||||||
|
@app.route("/chat", methods=['POST'])
|
||||||
|
def app_chat():
|
||||||
|
data = json.loads(flask.globals.request.get_data())
|
||||||
|
#print(data)
|
||||||
|
prompt = GlobalData.context_for_users.get(data["user_id"], "")
|
||||||
|
|
||||||
|
if not data["text"][-1] in ['?', '？', '.', '。', ',', '，', '!', '！']:
|
||||||
|
data["text"] += "。"
|
||||||
|
|
||||||
|
prompt += "\n" + data["text"]
|
||||||
|
|
||||||
|
if len(prompt) > 4000:
|
||||||
|
prompt = prompt[:4000]
|
||||||
|
|
||||||
|
resp = call_gpt(prompt=prompt)
|
||||||
|
GlobalData.context_for_users[data["user_id"]] = (prompt + resp)
|
||||||
|
|
||||||
|
print(f"Prompt = {prompt}\nResponse = {resp}")
|
||||||
|
|
||||||
|
return json.dumps({"user_id" : data["user_id"], "text" : resp, "in_group" : False})
|
||||||
|
|
||||||
|
@app.route("/draw", methods=['POST'])
|
||||||
|
def app_draw():
|
||||||
|
data = json.loads(flask.globals.request.get_data())
|
||||||
|
|
||||||
|
prompt = data["prompt"]
|
||||||
|
|
||||||
|
i = 0
|
||||||
|
for i in range(len(prompt)):
|
||||||
|
if prompt[i] == ':' or prompt[i] == '：':
|
||||||
|
break
|
||||||
|
if i == len(prompt):
|
||||||
|
return json.dumps({"user_name" : data["user_name"], "filename" : "", "error" : True, "error_msg" : "格式不对，正确的格式是：生成图片：Prompt 或者 生成图片(宽 高 迭代次数)：Prompt"})
|
||||||
|
|
||||||
|
|
||||||
|
match_args = re.match(GlobalData.GENERATE_PICTURE_ARG_PAT, prompt[:i])
|
||||||
|
if not match_args is None:
|
||||||
|
W = int(match_args.group(2))
|
||||||
|
H = int(match_args.group(3))
|
||||||
|
ITS = int(match_args.group(4))
|
||||||
|
|
||||||
|
if W > 2000 or H > 2000:
|
||||||
|
return json.dumps({"user_name" : data["user_name"], "filename" : "", "error" : True, "error_msg" : "你要求的图片太大了，我不干了～"})
|
||||||
|
|
||||||
|
if ITS > GlobalData.GENERATE_PICTURE_MAX_ITS:
|
||||||
|
return json.dumps({"user_name" : data["user_name"], "filename" : "", "error" : True, "error_msg" : f"迭代次数太多了，不要超过{GlobalData.GENERATE_PICTURE_MAX_ITS}次"})
|
||||||
|
else:
|
||||||
|
if len(prompt[:i].strip()) != 0:
|
||||||
|
return json.dumps({"user_name" : data["user_name"], "filename" : "", "error" : True, "error_msg" : "格式不对，正确的格式是：生成图片：Prompt 或者 生成图片(宽 高 迭代次数)：Prompt"})
|
||||||
|
else:
|
||||||
|
W = 768
|
||||||
|
H = 768
|
||||||
|
ITS = 50
|
||||||
|
|
||||||
|
prompt = prompt[(i+1):].strip()
|
||||||
|
|
||||||
|
prompts = re.split(GlobalData.GENERATE_PICTURE_NEG_PROMPT_DELIMETER, prompt)
|
||||||
|
prompt = prompts[0]
|
||||||
|
|
||||||
|
neg_prompt = None
|
||||||
|
if len(prompts) > 1:
|
||||||
|
neg_prompt = prompts[1]
|
||||||
|
|
||||||
|
print(f"Generating picture with prompt = {prompt} , negative prompt = {neg_prompt}")
|
||||||
|
|
||||||
|
image = sd_pipe(prompt=prompt, negative_prompt=neg_prompt, width=W, height=H, num_inference_steps=ITS).images[0]
|
||||||
|
|
||||||
|
image.save("latest.png")
|
||||||
|
return json.dumps({"user_name" : data["user_name"], "filename" : "latest.png", "error" : False, "error_msg" : ""})
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
#openai.organization = GlobalData.OPENAI_ORGID
|
||||||
|
openai.api_key = GlobalData.OPENAI_APIKEY
|
||||||
|
|
||||||
|
app.run(host="0.0.0.0", port=11111)
|
||||||
|
|
@ -0,0 +1,5 @@
|
||||||
|
{
|
||||||
|
"OpenAI-API-Key" : "",
|
||||||
|
"GPT-Model" : "text-davinci-003",
|
||||||
|
"Diffusion-Model" : "stabilityai/stable-diffusion-2-1"
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,5 @@
|
||||||
|
module SXYWechatBot
|
||||||
|
|
||||||
|
go 1.19
|
||||||
|
|
||||||
|
require github.com/eatmoreapple/openwechat v1.3.8
|
||||||
|
|
@ -0,0 +1,2 @@
|
||||||
|
github.com/eatmoreapple/openwechat v1.3.8 h1:dGQy/UeuSb7eVaCgJMoxM0fBG5DhboKfYk3g+FLZOWE=
|
||||||
|
github.com/eatmoreapple/openwechat v1.3.8/go.mod h1:61HOzTyvLobGdgWhL68jfGNwTJEv0mhQ1miCXQrvWU8=
|
||||||
|
|
@ -0,0 +1,2 @@
|
||||||
|
go run wechat_client.go & python bot.py --config $1
|
||||||
|
|
||||||
|
|
@ -0,0 +1,187 @@
|
||||||
|
package main
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"net/http"
|
||||||
|
"io/ioutil"
|
||||||
|
"time"
|
||||||
|
"os"
|
||||||
|
"bytes"
|
||||||
|
"encoding/json"
|
||||||
|
"github.com/eatmoreapple/openwechat"
|
||||||
|
)
|
||||||
|
|
||||||
|
func Use(vals ...interface{}) {
|
||||||
|
for _, val := range vals {
|
||||||
|
_ = val
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type SendTextRequest struct {
|
||||||
|
InGroup bool `json:"in_group"` //本来想用于区分在群聊和非群聊时的上下文记忆规则，但是最终没有实现...
|
||||||
|
UserID string `json:"user_id"`
|
||||||
|
Text string `json:"text"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SendImageRequest struct {
|
||||||
|
UserName string `json:"user_name"`
|
||||||
|
FileName string `json:"filename"`
|
||||||
|
HasError bool `json:"error"`
|
||||||
|
ErrorMessage string `json:"error_msg"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type GenerateImageRequest struct {
|
||||||
|
UserName string `json:"user_name"`
|
||||||
|
Prompt string `json:"prompt"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func HttpPost(url string, data interface{}, timelim int) []byte {
|
||||||
|
// 超时时间
|
||||||
|
timeout, _ := time.ParseDuration(fmt.Sprintf("%ss", timelim))
|
||||||
|
|
||||||
|
client := &http.Client{Timeout: timeout}
|
||||||
|
jsonStr, _ := json.Marshal(data)
|
||||||
|
resp, err := client.Post(url, "application/json", bytes.NewBuffer(jsonStr))
|
||||||
|
if err != nil {
|
||||||
|
return []byte("")
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
result, _ := ioutil.ReadAll(resp.Body)
|
||||||
|
return result
|
||||||
|
|
||||||
|
// ———————————————
|
||||||
|
// 版权声明：本文为CSDN博主「gaoluhua」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
|
||||||
|
// 原文链接：https://blog.csdn.net/gaoluhua/article/details/124855716
|
||||||
|
}
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
|
||||||
|
bot := openwechat.DefaultBot(openwechat.Desktop) // 桌面模式，上面登录不上的可以尝试切换这种模式
|
||||||
|
reloadStorage := openwechat.NewJsonFileHotReloadStorage("storage.json")
|
||||||
|
defer reloadStorage.Close()
|
||||||
|
|
||||||
|
err := bot.PushLogin(reloadStorage, openwechat.NewRetryLoginOption())
|
||||||
|
if err != nil {
|
||||||
|
fmt.Println(err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取登陆的用户
|
||||||
|
self, err := bot.GetCurrentUser()
|
||||||
|
if err != nil {
|
||||||
|
fmt.Println(err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
Use(self)
|
||||||
|
|
||||||
|
// 注册消息处理函数
|
||||||
|
bot.MessageHandler = func(msg *openwechat.Message) {
|
||||||
|
if msg.IsTickledMe() {
|
||||||
|
msg.ReplyText("别拍了，机器人是会被拍坏掉的。")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if !msg.IsText() {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
content := msg.Content
|
||||||
|
if msg.IsSendByGroup() && !msg.IsAt() {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if msg.IsSendByGroup() && msg.IsAt() {
|
||||||
|
atheader := fmt.Sprintf("@%s", self.NickName)
|
||||||
|
//fmt.Println(atheader)
|
||||||
|
if strings.HasPrefix(content, atheader) {
|
||||||
|
content = strings.TrimLeft(content[len(atheader):], " \t\n")
|
||||||
|
for {
|
||||||
|
if strings.HasPrefix(content, " ") {
|
||||||
|
content = content[1:]
|
||||||
|
} else {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
//fmt.Println(content)
|
||||||
|
|
||||||
|
if strings.HasPrefix(content, "生成图片") {
|
||||||
|
// 调用Stable Diffusion
|
||||||
|
// msg.ReplyText("这个功能还没有实现，可以先期待一下~")
|
||||||
|
sender, _ := msg.Sender()
|
||||||
|
|
||||||
|
content = strings.TrimLeft(content[len("生成图片"):], " \t\n")
|
||||||
|
|
||||||
|
resp_raw := HttpPost("http://localhost:11111/draw", GenerateImageRequest{UserName : sender.ID(), Prompt : content}, 120)
|
||||||
|
if len(resp_raw) == 0 {
|
||||||
|
msg.ReplyText("生成图片出错啦QwQ，或许可以再试一次")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
resp := SendImageRequest{}
|
||||||
|
json.Unmarshal(resp_raw, &resp)
|
||||||
|
//fmt.Println(resp.FileName)
|
||||||
|
if resp.HasError {
|
||||||
|
msg.ReplyText( fmt.Sprintf("生成图片出错啦QwQ，错误信息是：%s", resp.ErrorMessage) )
|
||||||
|
} else {
|
||||||
|
img, _ := os.Open(resp.FileName)
|
||||||
|
defer img.Close()
|
||||||
|
msg.ReplyImage(img)
|
||||||
|
}
|
||||||
|
|
||||||
|
} else {
|
||||||
|
// 调用ChatGPT
|
||||||
|
|
||||||
|
sender, _ := msg.Sender()
|
||||||
|
//var group openwechat.Group{} = nil
|
||||||
|
var group *openwechat.Group = nil
|
||||||
|
|
||||||
|
if msg.IsSendByGroup() {
|
||||||
|
group = &openwechat.Group{User : sender}
|
||||||
|
}
|
||||||
|
|
||||||
|
if content == "重置上下文" {
|
||||||
|
if !msg.IsSendByGroup() {
|
||||||
|
HttpPost("http://localhost:11111/chat_clear", SendTextRequest{InGroup : msg.IsSendByGroup(), UserID : sender.ID(), Text : ""}, 60)
|
||||||
|
} else {
|
||||||
|
HttpPost("http://localhost:11111/chat_clear", SendTextRequest{InGroup : msg.IsSendByGroup(), UserID : group.ID(), Text : ""}, 60)
|
||||||
|
}
|
||||||
|
msg.ReplyText("OK，我忘掉了之前的上下文。")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
resp := SendTextRequest{}
|
||||||
|
resp_raw := []byte("")
|
||||||
|
|
||||||
|
if !msg.IsSendByGroup() {
|
||||||
|
resp_raw = HttpPost("http://localhost:11111/chat", SendTextRequest{InGroup : false, UserID : sender.ID(), Text : msg.Content}, 60)
|
||||||
|
} else {
|
||||||
|
resp_raw = HttpPost("http://localhost:11111/chat", SendTextRequest{InGroup : false, UserID : group.ID(), Text : msg.Content}, 60)
|
||||||
|
}
|
||||||
|
if len(resp_raw) == 0 {
|
||||||
|
msg.ReplyText("运算超时了QAQ，或许可以再试一次。")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
json.Unmarshal(resp_raw, &resp)
|
||||||
|
|
||||||
|
if len(resp.Text) == 0 {
|
||||||
|
msg.ReplyText("GPT对此没有什么想说的，换个话题吧。")
|
||||||
|
} else {
|
||||||
|
if msg.IsSendByGroup() {
|
||||||
|
sender_in_group, _ := msg.SenderInGroup()
|
||||||
|
nickname := sender_in_group.NickName
|
||||||
|
msg.ReplyText(fmt.Sprintf("@%s\n%s\n-------------------\n%s", nickname, content, resp.Text))
|
||||||
|
} else {
|
||||||
|
msg.ReplyText(resp.Text)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
bot.Block()
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue