|
|
@@ -4,6 +4,7 @@ import json
|
|
|
import os
|
|
|
from bridge.context import ContextType
|
|
|
from bridge.reply import Reply, ReplyType
|
|
|
+from config import conf
|
|
|
import plugins
|
|
|
from plugins import *
|
|
|
from common.log import logger
|
|
|
@@ -45,40 +46,49 @@ class SDWebUI(Plugin):
|
|
|
try:
|
|
|
content = e_context['context'].content[:]
|
|
|
# 解析用户输入 如"横版 高清 二次元:cat"
|
|
|
- keywords, prompt = content.split(":", 1)
|
|
|
+ if ":" in content:
|
|
|
+ keywords, prompt = content.split(":", 1)
|
|
|
+ else:
|
|
|
+ keywords = content
|
|
|
+ prompt = ""
|
|
|
+
|
|
|
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的默认逻辑
|
|
|
+ if "help" in keywords or "帮助" in keywords:
|
|
|
+ reply.type = ReplyType.INFO
|
|
|
+ reply.content = self.get_help_text()
|
|
|
+ else:
|
|
|
+ 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)
|
|
|
@@ -86,3 +96,19 @@ class SDWebUI(Plugin):
|
|
|
e_context.action = EventAction.CONTINUE # 事件继续,交付给下个插件或默认逻辑
|
|
|
finally:
|
|
|
e_context['reply'] = reply
|
|
|
+
|
|
|
+ def get_help_text(self):
|
|
|
+ if not conf().get('image_create_prefix'):
|
|
|
+ return "画图功能未启用"
|
|
|
+ else:
|
|
|
+ trigger = conf()['image_create_prefix'][0]
|
|
|
+ help_text = f"请使用<{trigger}[关键词1] [关键词2]...:提示语>的格式作画,如\"{trigger}横版 高清:cat\"\n"
|
|
|
+ help_text += "目前可用关键词:\n"
|
|
|
+ for rule in self.rules:
|
|
|
+ keywords = [f"[{keyword}]" for keyword in rule['keywords']]
|
|
|
+ help_text += f"{','.join(keywords)}"
|
|
|
+ if "desc" in rule:
|
|
|
+ help_text += f"-{rule['desc']}\n"
|
|
|
+ else:
|
|
|
+ help_text += "\n"
|
|
|
+ return help_text
|