mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-01-31 06:19:50 +08:00
89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
# encoding:utf-8
|
||
|
||
import json
|
||
import os
|
||
from bridge.context import ContextType
|
||
from bridge.reply import Reply, ReplyType
|
||
import plugins
|
||
from plugins import *
|
||
from common.log import logger
|
||
import webuiapi
|
||
import io
|
||
|
||
|
||
@plugins.register(name="sdwebui", desc="利用stable-diffusion webui来画图", version="2.0", author="lanvent")
|
||
class SDWebUI(Plugin):
|
||
def __init__(self):
|
||
super().__init__()
|
||
curdir = os.path.dirname(__file__)
|
||
config_path = os.path.join(curdir, "config.json")
|
||
try:
|
||
with open(config_path, "r", encoding="utf-8") as f:
|
||
config = json.load(f)
|
||
self.rules = config["rules"]
|
||
defaults = config["defaults"]
|
||
self.default_params = defaults["params"]
|
||
self.default_options = defaults["options"]
|
||
self.start_args = config["start"]
|
||
self.api = webuiapi.WebUIApi(**self.start_args)
|
||
self.handlers[Event.ON_HANDLE_CONTEXT] = self.on_handle_context
|
||
logger.info("[SD] inited")
|
||
except FileNotFoundError:
|
||
logger.error(f"[SD] init failed, {config_path} not found")
|
||
except Exception as e:
|
||
logger.error("[SD] init failed, exception: %s" % e)
|
||
|
||
def on_handle_context(self, e_context: EventContext):
|
||
|
||
if e_context['context'].type != ContextType.IMAGE_CREATE:
|
||
return
|
||
|
||
logger.debug("[SD] on_handle_context. content: %s" %e_context['context'].content)
|
||
|
||
logger.info("[SD] image_query={}".format(e_context['context'].content))
|
||
reply = Reply()
|
||
try:
|
||
content = e_context['context'].content[:]
|
||
# 解析用户输入 如"横版 高清 二次元:cat"
|
||
keywords, prompt = content.split(":", 1)
|
||
keywords = keywords.split()
|
||
|
||
rule_params = {}
|
||
rule_options = {}
|
||
for keyword in keywords:
|
||
matched = False
|
||
for rule in self.rules:
|
||
if keyword in rule["keywords"]:
|
||
for key in rule["params"]:
|
||
rule_params[key] = rule["params"][key]
|
||
if "options" in rule:
|
||
for key in rule["options"]:
|
||
rule_options[key] = rule["options"][key]
|
||
matched = True
|
||
break # 一个关键词只匹配一个规则
|
||
if not matched:
|
||
logger.warning("[SD] keyword not matched: %s" % keyword)
|
||
|
||
params = {**self.default_params, **rule_params}
|
||
options = {**self.default_options, **rule_options}
|
||
params["prompt"] = params.get("prompt", "")+f", {prompt}"
|
||
if len(options) > 0:
|
||
logger.info("[SD] cover rule_options={}".format(rule_options))
|
||
self.api.set_options(options)
|
||
logger.info("[SD] params={}".format(params))
|
||
result = self.api.txt2img(
|
||
**params
|
||
)
|
||
reply.type = ReplyType.IMAGE
|
||
b_img = io.BytesIO()
|
||
result.image.save(b_img, format="PNG")
|
||
reply.content = b_img
|
||
e_context.action = EventAction.BREAK_PASS # 事件结束后，不跳过处理context的默认逻辑
|
||
except Exception as e:
|
||
reply.type = ReplyType.ERROR
|
||
reply.content = "[SD] "+str(e)
|
||
logger.error("[SD] exception: %s" % e)
|
||
e_context.action = EventAction.CONTINUE # 事件继续，交付给下个插件或默认逻辑
|
||
finally:
|
||
e_context['reply'] = reply
|