├── MANIFEST.in ├── nonebot_plugin_colab_novelai ├── permissionManager │ ├── __init__.py │ ├── cd.py │ └── nsfw.py ├── js │ ├── keepPageActive.js │ ├── objKeySort.js │ └── undefineWebDriver.js ├── utils │ ├── __init__.py │ ├── webdriver.py │ └── distributed.py ├── saveto │ ├── __init__.py │ ├── local.py │ └── webdav.py ├── access │ ├── cpolar.py │ ├── bce.py │ ├── naifu.py │ └── colab.py ├── config.py ├── __init__.py ├── argparsers.py ├── __meta__.py └── _main.py ├── .idea ├── misc.xml ├── vcs.xml ├── .gitignore ├── inspectionProfiles │ ├── profiles_settings.xml │ └── Project_Default.xml ├── nonebot-plugin-colab-novelai.iml └── modules.xml ├── LICENSE ├── setup.py ├── .gitignore └── README.md /MANIFEST.in: -------------------------------------------------------------------------------- 1 | graft nonebot_plugin_colab_novelai/js 2 | graft nonebot_plugin_colab_novelai/access -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/permissionManager/__init__.py: -------------------------------------------------------------------------------- 1 | from .cd import CooldownManager 2 | from .nsfw import NotSafeForWorkManager 3 | -------------------------------------------------------------------------------- /.idea/misc.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /.idea/vcs.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /.idea/.gitignore: -------------------------------------------------------------------------------- 1 | # Default ignored files 2 | /shelf/ 3 | /workspace.xml 4 | # Editor-based HTTP Client requests 5 | /httpRequests/ 6 | # Datasource local storage ignored files 7 | /dataSources/ 8 | /dataSources.local.xml 9 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/profiles_settings.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 6 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/js/keepPageActive.js: -------------------------------------------------------------------------------- 1 | function ConnectButton(){ 2 | console.log("Connect pushed"); 3 | document.querySelector("#top-toolbar > colab-connect-button").shadowRoot.querySelector("#connect").click() 4 | } 5 | setInterval(ConnectButton,60000); -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/js/objKeySort.js: -------------------------------------------------------------------------------- 1 | function objKeySort(obj) { 2 | let newkey = Object.keys(obj).sort(); 3 | let resStr = ''; 4 | for (let i = 0; i < newkey.length; i++) { 5 | let str = obj[newkey[i]]; 6 | console.log(i, newkey[i], str); 7 | resStr += str; 8 | } 9 | } -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/js/undefineWebDriver.js: -------------------------------------------------------------------------------- 1 | Object.defineProperties(navigator,{ webdriver:{ get: () => undefined } }) 2 | window.navigator.chrome = { runtime: {}, }; 3 | Object.defineProperty(navigator, 'languages', { get: () => ['en-US', 'en'] }); 4 | Object.defineProperty(navigator, 'plugins', { get: () => [1, 2, 3, 4, 5,6], }); -------------------------------------------------------------------------------- /.idea/nonebot-plugin-colab-novelai.iml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .webdriver import ( 2 | chrome_driver, 3 | force_refresh_webpage, wait_and_click_element, 4 | ) 5 | from .distributed import ( 6 | PLUGIN_DIR, NSFW_TAGS, 7 | T_UserID, T_AuthorizedUserID, T_GroupID, T_AuthorizedGroupID, 8 | get_mac_address, convert_audio2wav, 9 | fetch_image_in_message, inject_image_to_state, preprocess_painting_parameters 10 | ) 11 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/saveto/__init__.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional 2 | 3 | from .local import save_img_to_local 4 | from .webdav import save_img_to_webdav 5 | 6 | 7 | async def save_content(images: List[bytes], prompts: str, uc: str, baseimage: Optional[bytes] = None) -> None: 8 | await save_img_to_local(images, prompts, uc, baseimage=baseimage) 9 | await save_img_to_webdav(images, prompts, uc, baseimage=baseimage) 10 | -------------------------------------------------------------------------------- /.idea/modules.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/Project_Default.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 20 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 ThetaPilla. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/access/cpolar.py: -------------------------------------------------------------------------------- 1 | import re 2 | 3 | from httpx import AsyncClient 4 | 5 | from ..config import plugin_config 6 | 7 | 8 | CPOLAR_USER_INFO = { 9 | "login": plugin_config.cpolar_username, 10 | "password": plugin_config.cpolar_password 11 | } 12 | 13 | 14 | async def get_cpolar_authtoken() -> str: 15 | async with AsyncClient() as client: 16 | dashboard_resp = await client.post( 17 | url="https://dashboard.cpolar.com/login", 18 | data=CPOLAR_USER_INFO, 19 | follow_redirects=True, 20 | timeout=None 21 | ) 22 | 23 | try: 24 | cpolar_authtoken = re.findall(r"authtoken\s.+<", dashboard_resp.text)[0][10:-1] 25 | return cpolar_authtoken 26 | except IndexError: 27 | raise ValueError("cpolar帐密填写有误！") 28 | 29 | 30 | async def get_cpolar_url() -> str: 31 | async with AsyncClient() as client: 32 | await client.post( 33 | url="https://dashboard.cpolar.com/login", 34 | data=CPOLAR_USER_INFO, 35 | timeout=None 36 | ) 37 | dashboard_resp = await client.get("https://dashboard.cpolar.com/status", timeout=None) 38 | 39 | try: 40 | cpolar_url = re.findall(r">https://.+\.cpolar\..+", dashboard_resp.text)[0][1:-4] 41 | return cpolar_url 42 | except IndexError: 43 | raise RuntimeError("cpolar不存在正在运行的tunnel！") 44 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/config.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, Dict, List, Union 2 | 3 | from nonebot import get_driver 4 | from pydantic import BaseModel, Extra 5 | 6 | 7 | class Config(BaseModel, extra=Extra.ignore): 8 | headless_webdriver: bool = True 9 | colab_proxy: Optional[str] = None 10 | google_accounts: Dict[str, str] = {} 11 | cpolar_username: str = None 12 | cpolar_password: str = None 13 | bce_apikey: str = None 14 | bce_secretkey: str = None 15 | naifu_max: int = 1 16 | naifu_cd: int = 0 17 | nai_save2local_path: Optional[str] = None 18 | nai_save2webdav_info: Dict[str, Optional[str]] = { 19 | "url": None, 20 | "username": None, "password": None, 21 | "path": None 22 | } 23 | nai_nsfw_tags: Optional[Union[List[str], str]] = None 24 | 25 | 26 | plugin_config = Config.parse_obj(get_driver().config.dict()) 27 | 28 | # Startup autocheck 29 | assert plugin_config.google_accounts, "至少需要填写一个Google账号！" 30 | assert( 31 | plugin_config.cpolar_username is not None and plugin_config.cpolar_password is not None 32 | ), "需要同时填写cpolar的账号与密码！" 33 | assert( 34 | plugin_config.bce_apikey is not None and plugin_config.bce_secretkey is not None 35 | ), "需要同时填写百度智能云的apiKey与SecretKey！" 36 | assert( 37 | all(plugin_config.nai_save2webdav_info.values()) or 38 | all(i is None for i in plugin_config.nai_save2webdav_info.values()) 39 | ), "请检查WebDAV的Config是否完整填写！" 40 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | 4 | with open("README.md", "r", encoding="utf-8") as fh: 5 | long_description = fh.read() 6 | 7 | setuptools.setup( 8 | name="nonebot_plugin_colab_novelai", 9 | version="0.2.2", 10 | author="T_EtherLeaF", 11 | author_email="thetapilla@gmail.com", 12 | keywords=["pip", "nonebot2", "nonebot", "nonebot_plugin", "NovelAI", "Colaboratory", "QQ-bot", "chatbot"], 13 | description="""NoneBot2 由Colab驱动的AI作画插件""", 14 | long_description=long_description, 15 | long_description_content_type="text/markdown", 16 | url="https://github.com/EtherLeaF/nonebot-plugin-colab-novelai", 17 | packages=setuptools.find_packages(), 18 | classifiers=[ 19 | "Programming Language :: Python :: 3", 20 | "License :: OSI Approved :: MIT License", 21 | "Operating System :: OS Independent", 22 | ], 23 | include_package_data=True, 24 | platforms="any", 25 | install_requires=[ 26 | 'nonebot2>=2.0.0b4', 27 | 'nonebot-adapter-onebot>=2.1.5', 28 | 'nonebot-plugin-apscheduler>=0.1.4', 29 | 'httpx>=0.23.0', 30 | 'asyncer>=0.0.2', 31 | 'webdav4>=0.9.8', 32 | 'selenium>=4.6.0', 33 | 'selenium-stealth>=1.0.6', 34 | 'webdriver-manager>=3.8.4', 35 | 'av>=10.0.0', 36 | 'pyyaml>=6.0', 37 | 'packaging>=21.3', 38 | 'pillow>=9.3.0' 39 | ], 40 | python_requires=">=3.8" 41 | ) 42 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/saveto/local.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | from pathlib import Path 4 | from typing import List, Optional 5 | 6 | import anyio 7 | from nonebot.log import logger 8 | 9 | from ..config import plugin_config 10 | 11 | 12 | nai_save2local_path = plugin_config.nai_save2local_path 13 | 14 | if nai_save2local_path is not None: 15 | nai_save2local_path = Path(nai_save2local_path).absolute() 16 | os.makedirs(nai_save2local_path, exist_ok=True) 17 | logger.info(f"NovelAI返回的数据将保存到{nai_save2local_path}！") 18 | 19 | 20 | async def save_img_to_local(images: List[bytes], prompts: str, uc: str, baseimage: Optional[bytes] = None) -> None: 21 | if nai_save2local_path is None: 22 | return 23 | 24 | localtime = time.asctime(time.localtime()) 25 | folder_path = nai_save2local_path/localtime.replace(' ', '_').replace(':', '-') 26 | os.makedirs(folder_path, exist_ok=True) 27 | 28 | for i, image in enumerate(images): 29 | async with await anyio.open_file(folder_path/f"{i}.png", "wb") as f: 30 | await f.write(image) 31 | 32 | async with await anyio.open_file(folder_path/"prompts.txt", "w") as f: 33 | await f.write("[Prompts]\n" + prompts + "\n[Undesired Content]\n" + uc) 34 | 35 | if baseimage is not None: 36 | async with await anyio.open_file(folder_path/"original.png", "wb") as f: 37 | await f.write(baseimage) 38 | 39 | logger.success("图片已保存至本地！") 40 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/saveto/webdav.py: -------------------------------------------------------------------------------- 1 | import time 2 | from io import BytesIO 3 | from typing import Dict, List, Optional 4 | 5 | import asyncio 6 | from asyncer import asyncify 7 | from nonebot.log import logger 8 | 9 | from webdav4.client import Client as WebDavClient, HTTPError 10 | 11 | from ..config import plugin_config 12 | 13 | 14 | webdav_config: Dict[str, Optional[str]] = plugin_config.nai_save2webdav_info 15 | 16 | 17 | @asyncify 18 | def upload_image(client: WebDavClient, img: bytes, path: str) -> None: 19 | try: 20 | client.upload_fileobj(BytesIO(img), to_path=path, overwrite=True) 21 | logger.info(f"WebDAV: 图片已保存至{path}！") 22 | except HTTPError as e: 23 | logger.warning(f"图片保存失败：{e}") 24 | 25 | 26 | async def save_img_to_webdav(images: List[bytes], prompts: str, uc: str, baseimage: Optional[bytes] = None) -> None: 27 | if None in webdav_config.values(): 28 | return 29 | 30 | client = WebDavClient( 31 | webdav_config["url"], 32 | auth=(webdav_config["username"], webdav_config["password"]), 33 | timeout=None 34 | ) 35 | 36 | localtime = time.asctime(time.localtime()) 37 | folder_path = f"{webdav_config['path'].strip('/')}/{localtime}".replace(' ', '_').replace(':', '-') 38 | client.mkdir(folder_path) 39 | 40 | client.upload_fileobj( 41 | BytesIO(bytes("[Prompts]\n" + prompts + "\n[Undesired Content]\n" + uc, 'utf-8')), 42 | to_path=f"{folder_path}/prompts.txt", 43 | overwrite=True 44 | ) 45 | img_upload_tasks = [ 46 | asyncio.create_task(upload_image(client, img=image, path=f"{folder_path}/{i}.png")) 47 | for i, image in enumerate(images) 48 | ] 49 | await asyncio.gather(*img_upload_tasks) 50 | 51 | if baseimage is not None: 52 | await upload_image(client, img=baseimage, path=f"{folder_path}/original.png") 53 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/access/bce.py: -------------------------------------------------------------------------------- 1 | import requests 2 | from io import BytesIO 3 | 4 | from ..config import plugin_config 5 | from ..utils import convert_audio2wav, get_mac_address 6 | 7 | API_KEY = plugin_config.bce_apikey 8 | SECRET_KEY = plugin_config.bce_secretkey 9 | API_URL = "http://vop.baidu.com/server_api" 10 | TOKEN_URL = "https://aip.baidubce.com/oauth/2.0/token" 11 | 12 | 13 | def fetch_token(api_key: str, secret_key: str) -> str: 14 | params = { 15 | "grant_type": "client_credentials", 16 | "client_id": api_key, 17 | "client_secret": secret_key 18 | } 19 | result = requests.post(TOKEN_URL, data=params).json() 20 | 21 | try: 22 | err_des = result["error_description"] 23 | if err_des == "unknown client id": 24 | raise RuntimeError("请检查百度API Key！") 25 | elif err_des == "Client authentication failed": 26 | raise RuntimeError("请检查百度Secret Key！") 27 | # as expected 28 | except KeyError: 29 | return result["access_token"] 30 | 31 | 32 | def recognize_audio(url: str) -> str: 33 | audio_content = requests.get(url).content 34 | 35 | buf_in = BytesIO(audio_content) 36 | buf_out = BytesIO() 37 | convert_audio2wav(buf_in, buf_out) 38 | 39 | audio_content = buf_out.getvalue() 40 | bce_token = fetch_token(API_KEY, SECRET_KEY) 41 | mac_address = get_mac_address() 42 | 43 | header = { 44 | "Content-Type": "audio/wav;rate=16000" 45 | } 46 | params = { 47 | "cuid": mac_address, 48 | "token": bce_token, 49 | "dev_pid": 1737, 50 | } 51 | result = requests.post(API_URL, params=params, data=audio_content, headers=header).json() 52 | if result["err_no"] != 0: 53 | raise RuntimeError( 54 | f"百度API返回错误码：{result['err_no']}，" 55 | f"请参考文档：https://ai.baidu.com/ai-doc/SPEECH/pkgw0bw1p" 56 | ) 57 | 58 | answer = ' '.join(result["result"]) 59 | if answer == '': 60 | raise ValueError("语音未识别出结果！") 61 | return answer 62 | 63 | 64 | # Startup autocheck 65 | fetch_token(API_KEY, SECRET_KEY) 66 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/__init__.py: -------------------------------------------------------------------------------- 1 | # @Env: Python3.10 2 | # -*- coding: utf-8 -*- 3 | # @Author : T_EtherLeaF 4 | # @Email : thetapilla@gmail.com 5 | # @Software: PyCharm 6 | 7 | from argparse import Namespace 8 | 9 | from nonebot import require, get_driver 10 | from nonebot.permission import SUPERUSER 11 | from nonebot.plugin import on_shell_command 12 | from nonebot.params import T_State, Depends, Arg, ShellCommandArgs 13 | from nonebot.adapters.onebot.v11 import MessageEvent 14 | from nonebot_plugin_apscheduler import scheduler 15 | 16 | from .utils import chrome_driver, inject_image_to_state 17 | from .argparsers import naifu_draw_parser, naifu_perm_parser 18 | from ._main import handle_recaptcha, access_colab_with_accounts, naifu_img2img 19 | from .__meta__ import __plugin_meta__ 20 | 21 | 22 | nb_driver = get_driver() 23 | require("nonebot_plugin_apscheduler") 24 | scheduler.add_job(handle_recaptcha, "interval", id="checkReCaptcha", seconds=10) 25 | scheduler.add_job(access_colab_with_accounts, "interval", id="runColab", seconds=30) 26 | 27 | naifu_draw = on_shell_command("naifu", priority=10, parser=naifu_draw_parser) 28 | naifu_permission = on_shell_command("naifu", priority=10, parser=naifu_perm_parser, permission=SUPERUSER) 29 | 30 | 31 | @naifu_draw.handle() 32 | async def _naifu_draw(event: MessageEvent, state: T_State, args: Namespace = ShellCommandArgs()) -> None: 33 | await args.draw(matcher=naifu_draw, event=event, state=state, args=args) 34 | 35 | 36 | @naifu_draw.got('baseimage', prompt="请发送基准图片！", parameterless=[Depends(inject_image_to_state)]) 37 | async def _complete_img2img( 38 | event: MessageEvent, state: T_State, args: Namespace = ShellCommandArgs(), 39 | image: bytes = Arg('baseimage') 40 | ) -> None: 41 | await naifu_img2img(img=image, matcher=naifu_draw, event=event, state=state, args=args) 42 | 43 | 44 | @naifu_permission.handle() 45 | async def _operate_permission(event: MessageEvent, args: Namespace = ShellCommandArgs()) -> None: 46 | await args.operate(matcher=naifu_permission, event=event, user_id=args.uid, group_id=args.gid) 47 | 48 | 49 | @nb_driver.on_shutdown 50 | def _quit_webdriver() -> None: 51 | chrome_driver.quit() 52 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/permissionManager/cd.py: -------------------------------------------------------------------------------- 1 | import time 2 | from typing import Any, Tuple, List, Dict, Set, Type 3 | 4 | import yaml 5 | 6 | from nonebot import get_driver 7 | from nonebot.matcher import Matcher 8 | 9 | from ..config import plugin_config 10 | from ..utils import T_UserID, T_AuthorizedUserID 11 | 12 | 13 | class CooldownManager(object): 14 | @staticmethod 15 | def _load_yml() -> Tuple[Dict[T_UserID, float], Set[T_AuthorizedUserID]]: 16 | with open("./data/colab-novelai/cd.yml", "a+") as f: 17 | f.seek(0) 18 | cd_data = yaml.load(f, Loader=yaml.CFullLoader) 19 | 20 | if cd_data is None: 21 | initial_data = ({}, get_driver().config.superusers) 22 | yaml.dump(initial_data, f) 23 | return initial_data 24 | 25 | return cd_data 26 | 27 | @staticmethod 28 | def _save_yml(data: Any) -> None: 29 | with open("./data/colab-novelai/cd.yml", 'w') as f: 30 | yaml.dump(data, f) 31 | 32 | @classmethod 33 | async def list_authorized_users(cls, matcher: Type[Matcher], **kwargs: Any) -> None: 34 | authorized_users = cls._load_yml()[1] 35 | await matcher.send("当前有以下白名单用户：\n{}".format('\n'.join(authorized_users))) 36 | 37 | @classmethod 38 | async def add_authorized_user( 39 | cls, 40 | matcher: Type[Matcher], 41 | user_id: List[T_UserID], 42 | **kwargs: Any 43 | ) -> None: 44 | cd_data, authorized_users = cls._load_yml() 45 | authorized_users.update(user_id) 46 | cls._save_yml((cd_data, authorized_users)) 47 | 48 | await matcher.send(f"成功将以下用户添加白名单：{', '.join(set(user_id))}") 49 | 50 | @classmethod 51 | async def remove_authorized_user( 52 | cls, 53 | matcher: Type[Matcher], 54 | user_id: List[T_AuthorizedUserID], 55 | **kwargs: Any 56 | ) -> None: 57 | user_id = set(user_id) 58 | cd_data, authorized_users = cls._load_yml() 59 | await matcher.send(f"成功将以下用户移出白名单：{', '.join(authorized_users & user_id)}") 60 | 61 | authorized_users -= user_id 62 | cls._save_yml((cd_data, authorized_users)) 63 | 64 | @classmethod 65 | def record_cd(cls, user_id: T_UserID, num: int) -> None: 66 | cd_data, authorized_users = cls._load_yml() 67 | if user_id in authorized_users: 68 | return 69 | 70 | cd_data[user_id] = time.time() + plugin_config.naifu_cd * num 71 | cls._save_yml((cd_data, authorized_users)) 72 | 73 | @classmethod 74 | def get_user_cd(cls, user_id: T_UserID) -> float: 75 | cd_data, authorized_user = cls._load_yml() 76 | if user_id in authorized_user: 77 | return 0 78 | 79 | try: 80 | return cd_data[user_id] - time.time() 81 | except KeyError: 82 | return 0 83 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/utils/webdriver.py: -------------------------------------------------------------------------------- 1 | import time 2 | import asyncio 3 | from typing import Any 4 | 5 | from selenium import webdriver 6 | from selenium.webdriver.chrome.service import Service 7 | from selenium.webdriver.support.ui import WebDriverWait 8 | from selenium.common.exceptions import NoAlertPresentException 9 | from selenium.webdriver.support import expected_conditions as ec 10 | from selenium.webdriver.chrome.webdriver import WebDriver as ChromeWebDriver 11 | from selenium_stealth import stealth 12 | from webdriver_manager.chrome import ChromeDriverManager 13 | 14 | from .distributed import PLUGIN_DIR 15 | from ..config import plugin_config 16 | from ..access.cpolar import get_cpolar_authtoken 17 | 18 | 19 | options = webdriver.ChromeOptions() 20 | 21 | if plugin_config.headless_webdriver: 22 | options.add_argument('--headless') 23 | options.add_argument("--no-sandbox") 24 | options.add_argument("window-size=1920,1080") 25 | options.add_argument('--disable-dev-shm-usage') 26 | options.add_argument("--disable-setuid-sandbox") 27 | options.add_argument("--remote-debugging-port=9222") 28 | else: 29 | options.add_argument("--start-maximized") 30 | if (proxy := plugin_config.colab_proxy) is not None: 31 | options.add_argument(f'--proxy-server={proxy}') 32 | 33 | options.add_argument("--disable-blink-features") 34 | options.add_argument("--disable-blink-features=AutomationControlled") 35 | options.add_argument('--incognito') 36 | options.add_argument("--disable-extensions") 37 | options.add_argument("--disable-infobars") 38 | options.add_argument("--no-default-browser-check") 39 | options.add_experimental_option("excludeSwitches", ["enable-automation"]) 40 | options.add_experimental_option("useAutomationExtension", False) 41 | 42 | driver_path = ChromeDriverManager(path="./data/colab-novelai/").install() 43 | chrome_driver = webdriver.Chrome(service=Service(driver_path), options=options) 44 | 45 | for filename in ("undefineWebDriver.js", "objKeySort.js"): 46 | with open(PLUGIN_DIR / "js" / filename, 'r', encoding='utf-8') as js: 47 | chrome_driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", { 48 | "source": js.read() 49 | }) 50 | 51 | stealth( 52 | chrome_driver, 53 | languages=["en-US", "en"], 54 | vendor="Google Inc.", 55 | platform="Win32", 56 | webgl_vendor="Intel Inc.", 57 | renderer="Intel Iris OpenGL Engine", 58 | fix_hairline=True, 59 | ) 60 | 61 | 62 | def force_refresh_webpage(driver: ChromeWebDriver, url: str) -> None: 63 | driver.get(url) 64 | try: 65 | driver.switch_to.alert.accept() 66 | except NoAlertPresentException: 67 | pass 68 | 69 | 70 | def wait_and_click_element(driver: ChromeWebDriver, by: str, value: str) -> Any: 71 | element = WebDriverWait(driver, 5).until( 72 | lambda t_driver: t_driver.find_element(by, value) 73 | ) 74 | WebDriverWait(driver, 3).until( 75 | ec.element_to_be_clickable((by, value)) 76 | ) 77 | element.click() 78 | 79 | time.sleep(0.1) 80 | return element 81 | 82 | 83 | # Startup autocheck 84 | try: 85 | asyncio.run(get_cpolar_authtoken()) 86 | chrome_driver.quit() 87 | # already in event loop 88 | except RuntimeError: 89 | pass 90 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/permissionManager/nsfw.py: -------------------------------------------------------------------------------- 1 | from typing import Type, Tuple, List, Set, Any, Optional 2 | 3 | import yaml 4 | 5 | from nonebot import get_driver 6 | from nonebot.matcher import Matcher 7 | from nonebot.adapters.onebot.v11 import MessageEvent, GroupMessageEvent 8 | 9 | from ..utils import T_UserID, T_AuthorizedUserID, T_GroupID, T_AuthorizedGroupID 10 | 11 | 12 | class NotSafeForWorkManager(object): 13 | @staticmethod 14 | def _load_yml() -> Tuple[Set[T_AuthorizedUserID], Set[T_AuthorizedGroupID]]: 15 | with open("./data/colab-novelai/nsfw.yml", "a+") as f: 16 | f.seek(0) 17 | nsfw_conf = yaml.load(f, Loader=yaml.CFullLoader) 18 | 19 | if nsfw_conf is None: 20 | initial_data = (get_driver().config.superusers, set()) 21 | yaml.dump(initial_data, f) 22 | return initial_data 23 | 24 | return nsfw_conf 25 | 26 | @staticmethod 27 | def _save_yml(data: Any) -> None: 28 | with open("./data/colab-novelai/nsfw.yml", 'w') as f: 29 | yaml.dump(data, f) 30 | 31 | @classmethod 32 | async def list_authorized_users(cls, matcher: Type[Matcher], **kwargs: Any) -> None: 33 | authorized_users, authorized_groups = cls._load_yml() 34 | await matcher.send( 35 | "当前允许以下用户使用nsfw tag:\n{}".format('\n'.join(authorized_users)) + 36 | "\n当前允许以下群组使用nsfw tag:\n{}".format('\n'.join(authorized_groups)) 37 | ) 38 | 39 | @classmethod 40 | async def add_authorized_user( 41 | cls, 42 | matcher: Type[Matcher], 43 | event: MessageEvent, 44 | user_id: List[T_UserID], group_id: List[T_GroupID] 45 | ) -> None: 46 | authorized_users, authorized_groups = cls._load_yml() 47 | 48 | if user_id: 49 | authorized_users.update(user_id) 50 | await matcher.send(f"已允许以下用户使用nsfw tag:{', '.join(set(user_id))}") 51 | if group_id: 52 | authorized_groups.update(group_id) 53 | await matcher.send(f"已允许以下群组使用nsfw tag:{', '.join(set(group_id))}") 54 | 55 | if not user_id and not group_id and isinstance(event, GroupMessageEvent): 56 | group_id = [str(event.group_id)] 57 | authorized_groups.update(group_id) 58 | await matcher.send(f"已允许本群使用nsfw tag!") 59 | 60 | cls._save_yml((authorized_users, authorized_groups)) 61 | 62 | @classmethod 63 | async def remove_authorized_user( 64 | cls, 65 | matcher: Type[Matcher], 66 | event: MessageEvent, 67 | user_id: List[T_AuthorizedUserID], group_id: List[T_AuthorizedGroupID] 68 | ) -> None: 69 | user_id = set(user_id) 70 | group_id = set(group_id) 71 | authorized_users, authorized_groups = cls._load_yml() 72 | 73 | if user_id: 74 | await matcher.send(f"已禁止以下用户使用nsfw tag:{', '.join(authorized_users & user_id)}") 75 | authorized_users -= user_id 76 | if group_id: 77 | await matcher.send(f"已禁止以下群组使用nsfw tag:{', '.join(authorized_groups & group_id)}") 78 | authorized_groups -= group_id 79 | 80 | if not user_id and not group_id and isinstance(event, GroupMessageEvent): 81 | group_id = {str(event.group_id)} 82 | await matcher.send(f"已禁止本群使用nsfw tag!") 83 | authorized_groups -= group_id 84 | 85 | cls._save_yml((authorized_users, authorized_groups)) 86 | 87 | @classmethod 88 | def check_nsfw_available(cls, user_id: T_UserID, group_id: Optional[T_GroupID]) -> bool: 89 | authorized_users, authorized_groups = cls._load_yml() 90 | 91 | if user_id in authorized_users and group_id is None: 92 | return True 93 | if user_id in authorized_users and group_id in authorized_groups: 94 | return True 95 | return False 96 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/access/naifu.py: -------------------------------------------------------------------------------- 1 | import re 2 | import base64 3 | from io import BytesIO 4 | from typing import List, Any 5 | 6 | from PIL import Image 7 | from httpx import AsyncClient 8 | 9 | from .cpolar import get_cpolar_url 10 | 11 | SIZE = [i*64 for i in range(1, 17)] 12 | 13 | 14 | # textual prompts -> picture 15 | async def txt2img( 16 | prompt: str, 17 | uc: str, 18 | seed: int, 19 | width: int = 512, height: int = 768, 20 | n_samples: int = 1, 21 | sampler: str = "k_euler_ancestral", 22 | steps: int = 28, 23 | scale: int = 12, 24 | uc_preset: int = 0, 25 | **kwargs: Any 26 | ) -> List[bytes]: 27 | try: 28 | cpolar_url = await get_cpolar_url() 29 | api_url = cpolar_url + "/generate-stream" 30 | headers = { 31 | "content-type": "application/json", 32 | "accept-encoding": "gzip, deflate, br", 33 | "authorization": "Bearer" 34 | } 35 | data = { 36 | "prompt": prompt, 37 | "width": width, "height": height, 38 | "n_samples": n_samples, 39 | "sampler": sampler, 40 | "steps": steps, 41 | "scale": scale, 42 | "seed": seed, 43 | "uc": uc, "ucPreset": uc_preset 44 | } 45 | 46 | async with AsyncClient() as client: 47 | resp = await client.get(cpolar_url, timeout=None) 48 | if resp.status_code in [404, 502]: 49 | raise RuntimeError("APP暂时还未就绪！") 50 | 51 | result = (await client.post(api_url, json=data, headers=headers, timeout=600)).text 52 | result = re.findall(r"data:\S+", result) 53 | images = [base64.b64decode(i[5:]) for i in result] 54 | 55 | return images 56 | 57 | # catch all unexpected events 58 | except RuntimeError as exc: 59 | raise RuntimeError("暂时没有资源可供作图哦，可以稍后再来！") from exc 60 | 61 | 62 | # original image + textual prompts -> picture 63 | async def img2img( 64 | prompt: str, 65 | uc: str, 66 | image: bytes, 67 | seed: int, 68 | n_samples: int = 1, 69 | sampler: str = "k_euler_ancestral", 70 | steps: int = 28, 71 | scale: int = 12, 72 | strength: float = 0.7, 73 | noise: float = 0.2, 74 | uc_preset: int = 0, 75 | **kwargs: Any 76 | ) -> List[bytes]: 77 | # image preprocessing 78 | image = Image.open(BytesIO(image)) 79 | if (bound := max(image.size)) > 1024: 80 | target_size = tuple(map(lambda x: x/bound*1024, image.size)) 81 | else: 82 | target_size = image.size 83 | width, height = map( 84 | lambda s: min(SIZE, key=lambda x: abs(x - s)), 85 | target_size 86 | ) 87 | image = image.resize((width, height), resample=Image.ANTIALIAS) 88 | 89 | image_buf = BytesIO() 90 | image.save(image_buf, format='PNG') 91 | image_b64 = base64.b64encode(image_buf.getvalue()).decode() 92 | 93 | # fetch output 94 | try: 95 | cpolar_url = await get_cpolar_url() 96 | api_url = cpolar_url + "/generate-stream" 97 | headers = { 98 | "content-type": "application/json", 99 | "accept-encoding": "gzip, deflate, br", 100 | "authorization": "Bearer" 101 | } 102 | data = { 103 | "prompt": prompt, 104 | "image": image_b64, 105 | "width": width, "height": height, 106 | "n_samples": n_samples, 107 | "sampler": sampler, 108 | "steps": steps, 109 | "scale": scale, 110 | "strength": strength, 111 | "noise": noise, 112 | "seed": seed, 113 | "uc": uc, "ucPreset": uc_preset 114 | } 115 | 116 | async with AsyncClient() as client: 117 | resp = await client.get(cpolar_url, timeout=None) 118 | if resp.status_code in [404, 502]: 119 | raise RuntimeError("APP暂时还未就绪！") 120 | 121 | result = (await client.post(api_url, json=data, headers=headers, timeout=600)).text 122 | result = re.findall(r"data:\S+", result) 123 | images = [base64.b64decode(i[5:]) for i in result] 124 | 125 | return images 126 | 127 | except RuntimeError as exc: 128 | raise RuntimeError("暂时没有资源可供作图哦，可以稍后再来！") from exc 129 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/argparsers.py: -------------------------------------------------------------------------------- 1 | from nonebot.rule import ArgumentParser 2 | 3 | from ._main import naifu_txt2img, naifu_img2img 4 | from .permissionManager import CooldownManager, NotSafeForWorkManager 5 | 6 | 7 | '''handle user commands''' 8 | naifu_draw_parser = ArgumentParser() 9 | naifu_draw_subparsers = naifu_draw_parser.add_subparsers() 10 | 11 | # text to image 12 | naifu_txt2img_parser = naifu_draw_subparsers.add_parser("draw") 13 | naifu_txt2img_parser.add_argument('prompt', type=str, nargs='+') 14 | naifu_txt2img_parser.add_argument('-i', '--undesired-content', type=str, nargs='+', default=[]) 15 | naifu_txt2img_parser.add_argument( 16 | '-a', '--sampling', type=str, default="k_euler_ancestral", 17 | choices=[ 18 | "k_euler_ancestral", "k_euler", "k_lms", # Recommended 19 | "plms", "ddim" # Other 20 | ] 21 | ) 22 | naifu_txt2img_parser.add_argument('-t', '--steps', type=int, default=28) 23 | naifu_txt2img_parser.add_argument('-c', '--scale', type=float, default=12) 24 | naifu_txt2img_parser.add_argument( 25 | '-s', '--size', type=str, default="512x768", 26 | choices=[ 27 | "384x640", "512x768", "512x1024", # Portrait 28 | "640x384", "768x512", "1024x512", # Landscape 29 | "512x512", "640x640", "1024x1024" # Square 30 | ] 31 | ) 32 | naifu_txt2img_parser.add_argument('-n', '--num', type=int, default=1) 33 | naifu_txt2img_parser.add_argument('-r', '--seed', type=int, default=-1) 34 | naifu_txt2img_parser.set_defaults(draw=naifu_txt2img) 35 | 36 | # image to image 37 | naifu_img2img_parser = naifu_draw_subparsers.add_parser("imgdraw") 38 | naifu_img2img_parser.add_argument('prompt', type=str, nargs='+') 39 | naifu_img2img_parser.add_argument('-i', '--undesired-content', type=str, nargs='+', default=[]) 40 | naifu_img2img_parser.add_argument( 41 | '-a', '--sampling', type=str, default="k_euler_ancestral", 42 | choices=[ 43 | "k_euler_ancestral", "k_euler", "k_lms", # Recommended 44 | "plms", "ddim" # Other 45 | ] 46 | ) 47 | naifu_img2img_parser.add_argument('-t', '--steps', type=int, default=50) 48 | naifu_img2img_parser.add_argument('-c', '--scale', type=float, default=12) 49 | naifu_img2img_parser.add_argument('-n', '--num', type=int, default=1) 50 | naifu_img2img_parser.add_argument('-r', '--seed', type=int, default=-1) 51 | naifu_img2img_parser.add_argument('-e', '--strength', type=float, default=0.7) 52 | naifu_img2img_parser.add_argument('-o', '--noise', type=float, default=0.2) 53 | naifu_img2img_parser.set_defaults(draw=naifu_img2img) 54 | 55 | 56 | '''permission management''' 57 | naifu_perm_parser = ArgumentParser() 58 | naifu_perm_subparsers = naifu_perm_parser.add_subparsers() 59 | 60 | # superuser (without cd) 61 | naifu_su_parser = naifu_perm_subparsers.add_parser("su") 62 | naifu_su_subparsers = naifu_su_parser.add_subparsers() 63 | 64 | naifu_ls_su_parser = naifu_su_subparsers.add_parser("ls") 65 | naifu_ls_su_parser.set_defaults(operate=CooldownManager.list_authorized_users, uid=[], gid=[]) 66 | 67 | naifu_add_su_parser = naifu_su_subparsers.add_parser("add") 68 | naifu_add_su_parser.add_argument('uid', type=str, nargs='+') 69 | naifu_add_su_parser.set_defaults(operate=CooldownManager.add_authorized_user, gid=[]) 70 | 71 | naifu_remove_su_parser = naifu_su_subparsers.add_parser("rm") 72 | naifu_remove_su_parser.add_argument('uid', type=str, nargs='+') 73 | naifu_remove_su_parser.set_defaults(operate=CooldownManager.remove_authorized_user, gid=[]) 74 | 75 | # nsfw mode 76 | naifu_nsfw_parser = naifu_perm_subparsers.add_parser("nsfw") 77 | naifu_nsfw_subparsers = naifu_nsfw_parser.add_subparsers() 78 | 79 | naifu_ls_nsfw_parser = naifu_nsfw_subparsers.add_parser("ls") 80 | naifu_ls_nsfw_parser.set_defaults(operate=NotSafeForWorkManager.list_authorized_users, uid=[], gid=[]) 81 | 82 | naifu_add_nsfw_parser = naifu_nsfw_subparsers.add_parser("add") 83 | naifu_add_nsfw_parser.add_argument('-u', '--uid', type=str, nargs='+', default=[]) 84 | naifu_add_nsfw_parser.add_argument('-g', '--gid', type=str, nargs='+', default=[]) 85 | naifu_add_nsfw_parser.set_defaults(operate=NotSafeForWorkManager.add_authorized_user) 86 | 87 | naifu_remove_nsfw_parser = naifu_nsfw_subparsers.add_parser("rm") 88 | naifu_remove_nsfw_parser.add_argument('-u', '--uid', type=str, nargs='+', default=[]) 89 | naifu_remove_nsfw_parser.add_argument('-g', '--gid', type=str, nargs='+', default=[]) 90 | naifu_remove_nsfw_parser.set_defaults(operate=NotSafeForWorkManager.remove_authorized_user) 91 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | ### VisualStudioCode template 2 | .vscode/* 3 | !.vscode/settings.json 4 | !.vscode/tasks.json 5 | !.vscode/launch.json 6 | !.vscode/extensions.json 7 | *.code-workspace 8 | 9 | # Local History for Visual Studio Code 10 | .history/ 11 | 12 | ### JetBrains template 13 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider 14 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 15 | 16 | # User-specific stuff 17 | .idea/**/workspace.xml 18 | .idea/**/tasks.xml 19 | .idea/**/usage.statistics.xml 20 | .idea/**/dictionaries 21 | .idea/**/shelf 22 | 23 | # Generated files 24 | .idea/**/contentModel.xml 25 | 26 | # Sensitive or high-churn files 27 | .idea/**/dataSources/ 28 | .idea/**/dataSources.ids 29 | .idea/**/dataSources.local.xml 30 | .idea/**/sqlDataSources.xml 31 | .idea/**/dynamic.xml 32 | .idea/**/uiDesigner.xml 33 | .idea/**/dbnavigator.xml 34 | 35 | # Gradle 36 | .idea/**/gradle.xml 37 | .idea/**/libraries 38 | 39 | # Gradle and Maven with auto-import 40 | # When using Gradle or Maven with auto-import, you should exclude module files, 41 | # since they will be recreated, and may cause churn. Uncomment if using 42 | # auto-import. 43 | # .idea/artifacts 44 | # .idea/compiler.xml 45 | # .idea/jarRepositories.xml 46 | # .idea/modules.xml 47 | # .idea/*.iml 48 | # .idea/modules 49 | # *.iml 50 | # *.ipr 51 | 52 | # CMake 53 | cmake-build-*/ 54 | 55 | # Mongo Explorer plugin 56 | .idea/**/mongoSettings.xml 57 | 58 | # File-based project format 59 | *.iws 60 | 61 | # IntelliJ 62 | out/ 63 | 64 | # mpeltonen/sbt-idea plugin 65 | .idea_modules/ 66 | 67 | # JIRA plugin 68 | atlassian-ide-plugin.xml 69 | 70 | # Cursive Clojure plugin 71 | .idea/replstate.xml 72 | 73 | # Crashlytics plugin (for Android Studio and IntelliJ) 74 | com_crashlytics_export_strings.xml 75 | crashlytics.properties 76 | crashlytics-build.properties 77 | fabric.properties 78 | 79 | # Editor-based Rest Client 80 | .idea/httpRequests 81 | 82 | # Android studio 3.1+ serialized cache file 83 | .idea/caches/build_file_checksums.ser 84 | 85 | ### Python template 86 | # Byte-compiled / optimized / DLL files 87 | __pycache__/ 88 | *.py[cod] 89 | *$py.class 90 | 91 | # C extensions 92 | *.so 93 | 94 | # Distribution / packaging 95 | .Python 96 | build/ 97 | develop-eggs/ 98 | dist/ 99 | downloads/ 100 | eggs/ 101 | .eggs/ 102 | lib/ 103 | lib64/ 104 | parts/ 105 | sdist/ 106 | var/ 107 | wheels/ 108 | share/python-wheels/ 109 | *.egg-info/ 110 | .installed.cfg 111 | *.egg 112 | MANIFEST 113 | 114 | # PyInstaller 115 | # Usually these files are written by a python script from a template 116 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 117 | *.manifest 118 | *.spec 119 | 120 | # Installer logs 121 | pip-log.txt 122 | pip-delete-this-directory.txt 123 | 124 | # Unit test / coverage reports 125 | htmlcov/ 126 | .tox/ 127 | .nox/ 128 | .coverage 129 | .coverage.* 130 | .cache 131 | nosetests.xml 132 | coverage.xml 133 | *.cover 134 | *.py,cover 135 | .hypothesis/ 136 | .pytest_cache/ 137 | cover/ 138 | 139 | # Translations 140 | *.mo 141 | *.pot 142 | 143 | # Django stuff: 144 | *.log 145 | local_settings.py 146 | db.sqlite3 147 | db.sqlite3-journal 148 | 149 | # Flask stuff: 150 | instance/ 151 | .webassets-cache 152 | 153 | # Scrapy stuff: 154 | .scrapy 155 | 156 | # Sphinx documentation 157 | docs/_build/ 158 | 159 | # PyBuilder 160 | .pybuilder/ 161 | target/ 162 | 163 | # Jupyter Notebook 164 | .ipynb_checkpoints 165 | 166 | # IPython 167 | profile_default/ 168 | ipython_config.py 169 | 170 | # pyenv 171 | # For a library or package, you might want to ignore these files since the code is 172 | # intended to run in multiple environments; otherwise, check them in: 173 | # .python-version 174 | 175 | # pipenv 176 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 177 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 178 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 179 | # install all needed dependencies. 180 | #Pipfile.lock 181 | 182 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 183 | __pypackages__/ 184 | 185 | # Celery stuff 186 | celerybeat-schedule 187 | celerybeat.pid 188 | 189 | # SageMath parsed files 190 | *.sage.py 191 | 192 | # Environments 193 | .env 194 | .venv 195 | env/ 196 | venv/ 197 | ENV/ 198 | env.bak/ 199 | venv.bak/ 200 | 201 | # Spyder project settings 202 | .spyderproject 203 | .spyproject 204 | 205 | # Rope project settings 206 | .ropeproject 207 | 208 | # mkdocs documentation 209 | /site 210 | 211 | # mypy 212 | .mypy_cache/ 213 | .dmypy.json 214 | dmypy.json 215 | 216 | # Pyre type checker 217 | .pyre/ 218 | 219 | # pytype static type analyzer 220 | .pytype/ 221 | 222 | # Cython debug symbols 223 | cython_debug/ 224 | 225 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/utils/distributed.py: -------------------------------------------------------------------------------- 1 | import re 2 | import os 3 | import uuid 4 | import random 5 | from pathlib import Path 6 | from argparse import Namespace 7 | from typing import Tuple, Dict, Any, Type, TypeVar, Optional, Union 8 | 9 | import av 10 | 11 | from httpx import AsyncClient 12 | from nonebot.matcher import Matcher 13 | from nonebot.params import T_State, Arg 14 | from nonebot.adapters.onebot.v11 import Message 15 | 16 | from ..config import plugin_config 17 | 18 | 19 | T_UserID = TypeVar("T_UserID", str, int) 20 | T_AuthorizedUserID = TypeVar("T_AuthorizedUserID", str, int) 21 | T_GroupID = TypeVar("T_GroupID", str, int) 22 | T_AuthorizedGroupID = TypeVar("T_AuthorizedGroupID", str, int) 23 | 24 | os.makedirs("./data/colab-novelai", exist_ok=True) 25 | PLUGIN_DIR = Path(__file__).absolute().parent.parent 26 | 27 | if plugin_config.nai_nsfw_tags is None: 28 | nai_nsfw_tags = set() 29 | 30 | elif isinstance(plugin_config.nai_nsfw_tags, list): 31 | nai_nsfw_tags = set(plugin_config.nai_nsfw_tags) 32 | 33 | elif isinstance(plugin_config.nai_nsfw_tags, str): 34 | with open(plugin_config.nai_nsfw_tags, 'r', encoding='utf-8') as f: 35 | content = f.read().replace('，', ',') 36 | nai_nsfw_tags = set(content.split(',')) 37 | 38 | NSFW_TAGS = nai_nsfw_tags | { 39 | 'nsfw', 'r18', 'nude', 'dick', 'cock', 'penis', 'pussy', 'cum', 'condom', 'nipple', 'penis', 'sex', 'vaginal', 40 | 'straddling', 'doggystyle', 'doggy style', 'doggy-style', 'missionary', 'lick', 'bukkake', 'armpit', 'breasts out', 41 | 'pov', 'rape', 'anal', 'double penetration', 'bdsm', 'milking', 'vibrator', 'ball gag', 'not safe for work' 42 | 'ejaculation', 'piercing', 'bukakke' 43 | } 44 | 45 | 46 | def get_mac_address() -> str: 47 | address = hex(uuid.getnode())[2:] 48 | return '-'.join(address[i:i+2] for i in range(0, len(address), 2)) 49 | 50 | 51 | def convert_audio2wav(fp_in: Any, fp_out: Any, sample_rate: int = 16000) -> None: 52 | with av.open(fp_in) as buf_in: 53 | in_stream = buf_in.streams.audio[0] 54 | 55 | with av.open(fp_out, 'w', 'wav') as buf_out: 56 | out_stream = buf_out.add_stream( 57 | "pcm_s16le", 58 | rate=sample_rate, 59 | layout="mono" 60 | ) 61 | for frame in buf_in.decode(in_stream): 62 | for packet in out_stream.encode(frame): 63 | buf_out.mux(packet) 64 | 65 | 66 | async def fetch_image_in_message(message: Message) -> Optional[bytes]: 67 | try: 68 | img_url = message["image"][0].data["url"] 69 | except IndexError: 70 | return None 71 | 72 | async with AsyncClient() as client: 73 | resp = await client.get(url=img_url, timeout=30) 74 | img = await resp.aread() 75 | return img 76 | 77 | 78 | async def inject_image_to_state( 79 | matcher: Matcher, state: T_State, 80 | image_arg: Union[bytes, Message] = Arg('baseimage') 81 | ) -> None: 82 | if isinstance(image_arg, bytes): 83 | return 84 | 85 | if (image := await fetch_image_in_message(image_arg)) is None: 86 | await matcher.finish("格式错误，请重新触发指令..") 87 | state["baseimage"] = image 88 | 89 | 90 | def nsfw_tag_filter(prompt: str, uc: str, on: bool = True) -> Tuple[str, str]: 91 | if on: 92 | uc += ','.join(NSFW_TAGS) 93 | 94 | prompts = set(prompt.split(',')) 95 | for p in prompts.copy(): 96 | for nsfw_tag in NSFW_TAGS: 97 | if re.match(nsfw_tag, p, flags=re.I) is not None: 98 | prompts.discard(p) 99 | if not prompts: 100 | raise ValueError("色鬼，好好检查一下你都放了些什么tag进来！") 101 | prompt = ','.join(prompts) 102 | 103 | return prompt, uc 104 | 105 | 106 | async def preprocess_painting_parameters(matcher: Type[Matcher], args: Namespace, on_nsfw: bool) -> Dict[str, Any]: 107 | try: 108 | assert 1 <= args.num <= plugin_config.naifu_max, "图片数量超过上限！" 109 | assert -1 <= args.seed <= 2**32 - 1, "设置的种子需要在-1与2^32-1之间！" 110 | assert 1 <= args.steps <= 50, "设置的steps需要在1到50之间！" 111 | assert 1.1 <= args.scale <= 100, "设置的scale需要在1.1到100之间！" 112 | except AssertionError as e: 113 | await matcher.finish(str(e), at_sender=True) 114 | 115 | prompts, uc = nsfw_tag_filter( 116 | prompt="masterpiece,best quality," + ' '.join(args.prompt).replace('，', ','), 117 | uc="lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, " 118 | "worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry," 119 | + ' '.join(args.undesired_content).replace('，', ',').strip(',') + ',', 120 | on=on_nsfw 121 | ) 122 | return { 123 | "prompt": prompts, 124 | "uc": uc, 125 | "n_samples": args.num, 126 | "seed": random.randint(0, 2**32 - 1) if args.seed == -1 else args.seed, 127 | "sampler": args.sampling, 128 | "steps": args.steps, 129 | "scale": round(args.scale, 1) 130 | } 131 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/access/colab.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from asyncer import asyncify 4 | from nonebot.log import logger 5 | 6 | from selenium.webdriver.common.by import By 7 | from selenium.webdriver.common.keys import Keys 8 | from selenium.webdriver.support.ui import WebDriverWait 9 | from selenium.webdriver.support import expected_conditions as ec 10 | from selenium.common.exceptions import TimeoutException, JavascriptException 11 | 12 | from ..utils import chrome_driver as driver, force_refresh_webpage, wait_and_click_element, PLUGIN_DIR 13 | 14 | NOTEBOOK_URL = "https://colab.research.google.com/drive/1oAMaO-0_SxFSr8OEC1jJVkC84pGFzw-I?usp=sharing" 15 | 16 | 17 | def login_google_acc(gmail: str, password: str) -> None: 18 | try: 19 | # No account logged in yet 20 | try: 21 | # click "Sign in" 22 | login = WebDriverWait(driver, 5).until( 23 | lambda t_driver: t_driver.find_element(By.XPATH, '//*[@id="gb"]/div/div/a') 24 | ) 25 | driver.get(login.get_attribute('href')) 26 | 27 | # Already logged in 28 | except TimeoutException: 29 | # logout current account 30 | logout = WebDriverWait(driver, 5).until( 31 | lambda t_driver: t_driver.find_element( 32 | By.XPATH, '//*[@id="gb"]/div/div[1]/div[2]/div/a' 33 | ) 34 | ) 35 | driver.get(logout.get_attribute('href')) 36 | driver.find_element(By.XPATH, '//*[@id="signout"]').click() 37 | 38 | # click "Sign in" 39 | login = WebDriverWait(driver, 5).until( 40 | lambda t_driver: t_driver.find_element(By.XPATH, '//*[@id="gb"]/div/div/a') 41 | ) 42 | driver.get(login.get_attribute('href')) 43 | 44 | # if prompt, choose "Use another account" when login 45 | try: 46 | wait_and_click_element( 47 | driver, 48 | by=By.XPATH, 49 | value='//*[@id="view_container"]/div/div/div[2]/div/div[1]/div/form/span/section/div/div/div/div/ul' 50 | '/li[@class="JDAKTe eARute W7Aapd zpCp3 SmR8" and not(@jsname="fKeql")]' 51 | ) 52 | except TimeoutException: 53 | pass 54 | 55 | # input gmail and password 56 | gmail_input = WebDriverWait(driver, 5).until(ec.element_to_be_clickable( 57 | (By.XPATH, '//*[@id="identifierId"]') 58 | )) 59 | driver.execute_script("arguments[0].click();", gmail_input) 60 | time.sleep(0.5) 61 | gmail_input.send_keys(gmail, Keys.ENTER) 62 | 63 | pwd_input = WebDriverWait(driver, 5).until(ec.element_to_be_clickable( 64 | (By.XPATH, '//*[@id="password"]/div[1]/div/div[1]/input') 65 | )) 66 | driver.execute_script("arguments[0].click();", pwd_input) 67 | time.sleep(0.5) 68 | pwd_input.send_keys(password, Keys.ENTER) 69 | 70 | # check if the password is incorrect 71 | try: 72 | WebDriverWait(driver, 3).until( 73 | lambda t_driver: t_driver.find_element( 74 | By.XPATH, '//*[@id="yDmH0d"]/c-wiz/div/div[2]/div/div[1]/div/form/span/div[1]/div[2]/div[1]' 75 | ) 76 | ) 77 | raise RuntimeError(f"Google账号{gmail}的密码填写有误！") 78 | except TimeoutException: 79 | logger.success(f"成功登入Google账号：{gmail}！") 80 | 81 | except TimeoutException: 82 | raise RuntimeError(f"登陆Google账号{gmail}发生超时，请检查网络和账密！") 83 | 84 | # In case of Google asking you to complete your account info 85 | try: 86 | # Wait for "not now" button occurs 87 | wait_and_click_element( 88 | driver, 89 | by=By.XPATH, value='//*[@id="yDmH0d"]/c-wiz/div/div/div/div[2]/div[4]/div[1]/button' 90 | ) 91 | 92 | # If that doesn't happen 93 | except TimeoutException: 94 | pass 95 | 96 | 97 | @asyncify 98 | def run_colab(gmail: str, password: str, cpolar_authtoken: str) -> None: 99 | force_refresh_webpage(driver, NOTEBOOK_URL) 100 | 101 | login_google_acc(gmail, password) 102 | 103 | # input cpolar authtoken 104 | time.sleep(3) 105 | try: 106 | authtoken_box = driver.execute_script( 107 | 'return document.querySelector("#cell-54WF-Om0X6tf > div.main-content > div.codecell-input-output > ' 108 | 'div.inputarea.horizontal.layout.both > colab-form > div > colab-form-input > div.layout.horizontal.grow > ' 109 | 'paper-input").shadowRoot.querySelector("#input-1 > input")' 110 | ) 111 | authtoken_box.clear() 112 | authtoken_box.send_keys(cpolar_authtoken) 113 | except JavascriptException: 114 | # failed to fill input box 115 | # mostly, this happens when Google is asking you to do extra verification i.e. phone number 116 | # Colab page won't be loaded normally, then result in this error. 117 | raise RuntimeError( 118 | f"Google账密验证成功，但Colab页面没有被成功加载。可能是因为Google正在要求账号进行额外验证或账号不再可用！" 119 | f"当前账号：{gmail}" 120 | ) 121 | 122 | # run all cells 123 | driver.find_element(By.XPATH, '/html/body').send_keys(Keys.CONTROL + Keys.F9) 124 | 125 | # If Google asks you to confirm running this notebook 126 | try: 127 | wait_and_click_element( 128 | driver, 129 | by=By.XPATH, value='/html/body/colab-dialog/paper-dialog/div[2]/paper-button[2]' 130 | ) 131 | except TimeoutException: 132 | pass 133 | 134 | # keep webpage active 135 | with open(PLUGIN_DIR / "js" / "keepPageActive.js", 'r', encoding='utf-8') as js: 136 | driver.execute_script(js.read()) 137 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |
2 | NoneBotPluginLogo 3 |
4 |

NoneBotPluginText

5 |
6 | 7 |
8 | 9 | # Nonebot-Plugin-Colab-NovelAI 10 | 11 | _✨ 基于框架 [NoneBot2](https://v2.nonebot.dev/) 的AI绘图插件 ✨_ 12 | 13 |

14 | license 15 | Python 16 | NoneBot 17 | 18 | pypi 19 | 20 | 21 | pypi download 22 | 23 |

24 | 25 |
26 | 27 | ## 功能 28 | 29 | - 提供prompt让AI进行绘图 30 | - 可选将图片保存至本地或WebDAV 31 | - 权限管理: 绘图冷却时间与是否允许使用NSFW tag 32 | 33 | ## 安装 34 | 35 | - 使用 nb-cli 36 | 37 | ``` 38 | nb plugin install nonebot_plugin_colab_novelai 39 | ``` 40 | 41 | - 使用 pip 42 | 43 | ``` 44 | pip install nonebot_plugin_colab_novelai 45 | ``` 46 | 47 | ## 获取插件帮助与拓展功能 48 | 49 | - 可选择接入 [nonebot-plugin-PicMenu](https://github.com/hamo-reid/nonebot_plugin_PicMenu) 以便用户获取插件相关信息与用法 50 | - 可选择接入 [nonebot-plugin-manager](https://github.com/nonepkg/nonebot-plugin-manager) 管理插件黑名单 51 | - 可选择接入 [nonebot-plugin-savor](https://github.com/A-kirami/nonebot-plugin-savor) 通过图片反推tag 52 | 53 | ## Requirements 54 | 55 | - 一台能正常访问外网的服务器 (Colab在中国大陆无法访问） 56 | 57 | - 确保服务器已正确安装了Chrome浏览器 58 | 59 | - 注册一堆Google新帐号(建议六个以上)，建议绑定手机号以免登录时出现麻烦，千万不要开启多余的安全设置。 60 | 61 | - 前往[百度智能云](https://ai.baidu.com/tech/speech)申请免费语音识别服务，注册APP并获取相关密钥 62 | - 用于绕过Colab ReCaptcha 63 | 64 | - 前往[cpolar](https://www.cpolar.com/)注册免费账号 65 | - 用于Colab的内网穿透 66 | 67 | ## .env | .env.dev | .env.prod 配置项 68 | 69 | ```ini 70 | headless_webdriver: bool = True # 是否使用无头模式启动浏览器 71 | colab_proxy: Optional[str] = None # 如有需要可填写代理地址 72 | google_accounts: Dict[str, str] = {} # Required, 填写要使用的谷歌账密 {"account": "password", ...} 73 | cpolar_username: str = None # Required, 填写cpolar账号邮箱 74 | cpolar_password: str = None # Required, 填写cpolar账号的密码 75 | bce_apikey: str = None # Required, 填写百度智能云的API Key 76 | bce_secretkey: str = None # Required, 填写百度智能云的Secret Key 77 | naifu_max: int = 1 # 一次作图生成的最大图片数量 78 | naifu_cd: int = 0 # 每个用户每生成一张图片的冷却时间 79 | nai_save2local_path: Optional[str] = None # 将图片保存至本地的存储目录, 不填写则不保存 80 | nai_save2webdav_info: Dict[str, Optional[str]] = { 81 | "url": None, 82 | "username": None, "password": None, # 将图片保存至WebDAV需要的相关配置，不填写则不保存 83 | "path": None 84 | } 85 | nai_nsfw_tags: Optional[List[str] | str] = None # 自定义可能会生成NSFW图片的tag, 填写一个列表或者一个文件路径 86 | # 列表: ["tag1", "tag2", "tag3", ...] 87 | # 若使用文件存储, 需要将tag以逗号分隔，无需引号。 88 | ``` 89 | 90 | ### 配置项额外说明 91 | 92 | - 如果你正在使用没有图形界面的Linux服务器，请不要更改```headless_webdriver``` 93 | 94 | - 插件会尝试禁止未授权的用户绘画NSFW图片，通过屏蔽特定tag来实现。预设的一些tag集合位于[/utils/distributed.py](https://github.com/EtherLeaF/nonebot-plugin-colab-novelai/blob/main/nonebot_plugin_colab_novelai/utils/distributed.py)，如果有其他好的预设想法，欢迎pr。 95 | - 屏蔽的tag集合为```.env```配置项与预设项的并集，匹配时不区分大小写。 96 | 97 | - 如需使用代理，支持填写```http://```or```https://```or```socks5://```+```ip:port``` 98 | 99 | ## 如何使用？ 100 | 101 | 触发指令: ```naifu [] []``` 102 | 103 | - Command: ```draw``` 104 | - CommandPermission: ```Anyone``` 105 | - 用于告诉AI开始作图 106 | 107 | - 用法: ```naifu draw ... [-i --undesired-content ...] [-a --sampling ] [-t --steps ] [-c --scale ] [-n --num ] [-s --size ] [-r --seed ]``` 108 | - ```PROMPT``` 必选参数，指定作画的关键词，以逗号分隔，必须为英语 109 | - ```-i``` 可选参数，指定作画中想避免出现的内容，以逗号分隔，必须为英语 110 | - ```-a``` 可选参数，指定采样器，支持以下几种，默认为```k_euler_ancestral```： 111 | - ```k_euler_ancestral, k_euler, k_lms``` 112 | - ```plms, ddim``` 113 | - ```-t``` 可选参数，指定优化图像的迭代次数，取值范围```1~50```，默认值为```28``` 114 | - ```-c``` 可选参数，值越大越接近描述意图，值越小细节越少自由度越大，取值范围```1.1~100```，默认值为```12``` 115 | - ```-s``` 可选参数，指定图片生成大小，支持以下几种，默认为```512x768```： 116 | - ```384x640, 512x768, 512x1024 # Portrait``` 117 | - ```640x384, 768x512, 1024x512 # Landscape``` 118 | - ```512x512, 640x640, 1024x1024 # Square``` 119 | - ```-n``` 可选参数，指定图片生成数量，最大值参考```.env```配置项，默认值为```1``` 120 | - ```-r``` 可选参数，指定图片生成种子，取值范围```0 ~ 2³²-1```，默认值为```-1```即随机 121 |
122 | 123 | - Command: ```imgdraw``` 124 | - CommandPermission: ```Anyone``` 125 | - 提供基准图片作图 126 | 127 | - 用法: ```naifu imgdraw ... [-i --undesired-content ...] [-a --sampling ] [-t --steps ] [-c --scale ] [-n --num ] [-r --seed ] [-e strength ] [-o noise ]``` 128 | - ```PROMPT``` 必选参数，指定作画的关键词，以逗号分隔，必须为英语 129 | - ```IMAGE``` 必选参数，指定作画基准图片 130 | - ```-i``` 可选参数，指定作画中想避免出现的内容，以逗号分隔，必须为英语 131 | - ```-a``` 可选参数，指定采样器，支持以下几种，默认为```k_euler_ancestral```： 132 | - ```k_euler_ancestral, k_euler, k_lms``` 133 | - ```plms, ddim``` 134 | - ```-t``` 可选参数，指定优化图像的迭代次数，取值范围```1~50```，默认值为```50``` 135 | - ```-c``` 可选参数，值越大越接近描述意图，值越小细节越少自由度越大，取值范围```1.1~100```，默认值为```12``` 136 | - ```-n``` 可选参数，指定图片生成数量，最大值参考```.env```配置项，默认值为```1``` 137 | - ```-r``` 可选参数，指定图片生成种子，取值范围```0 ~ 2³²-1```，默认值为```-1```即随机 138 | - ```-e``` 可选参数，值越低越接近原始图像，取值范围```0~0.99```，默认值为```0.7``` 139 | - ```-o``` 可选参数，值增加会增加细节，一般应低于参数``````，取值范围```0~0.99```，默认值为```0.2``` 140 |
141 | 142 | - Command: ```su``` 143 | - CommandPermission: ```Superuser``` 144 | - 用于管理插件白名单用户 (白名单用户无绘图cd，在```.env```中```naifu_cd```值为非零时生效) 145 | 146 | - Subcommand: ```ls``` 147 | - 列出当前所有白名单用户 148 | - 用法: ```naifu su ls``` 149 |
150 | 151 | - Subcommand: ```add``` 152 | - 添加白名单用户 153 | - 用法: ```naifu su add ...``` 154 | - 必须指定用户QQ号，可填写多个并以空格分隔 155 |
156 | 157 | - Subcommand: ```rm``` 158 | - 移除白名单用户 159 | - 用法: ```naifu su rm ...``` 160 | - 必须指定用户QQ号，可填写多个并以空格分隔 161 |
162 | 163 | - Command: ```nsfw``` 164 | - CommandPermission: ```Superuser``` 165 | - 管理允许绘制NSFW内容的用户与群组 166 | - 注意: 群聊中只有当用户和群聊均有权限时才能绘制NSFW内容！ 167 | 168 | - Subcommand: ```ls``` 169 | - 列出当前所有允许NSFW内容的用户与群组 170 | - 用法: ```naifu nsfw ls``` 171 |
172 | 173 | - Subcommand: ```add``` 174 | - 添加允许NSFW内容的用户或群组 175 | - 用法: ```naifu nsfw add [-u --uid ...] [-g --gid ...]``` 176 | - ```-u``` 可选参数，为用户QQ号，可填写多个并以空格分隔 177 | - ```-g``` 可选参数，为群号，可填写多个并以空格分隔 178 | - 当两个可选参数均未填写时，默认添加当前所处群聊的群号。 179 |
180 | 181 | - Subcommand: ```rm``` 182 | - 移除允许NSFW内容的用户或群组 183 | - 用法: ```naifu nsfw rm [-u --uid ...] [-g --gid ...]``` 184 | - ```-u``` 可选参数，为用户QQ号，可填写多个并以空格分隔 185 | - ```-g``` 可选参数，为群号，可填写多个并以空格分隔 186 | - 当两个可选参数均未填写时，默认移除当前所处群聊的群号。 187 | 188 | 在权限配置文件第一次加载时，会自动添加```.env```的```SUPERUSERS```为插件白名单用户以及分配NSFW权限。 189 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/__meta__.py: -------------------------------------------------------------------------------- 1 | from nonebot.plugin import PluginMetadata 2 | from nonebot.adapters.onebot.v11 import Adapter as OneBotV11Adapter 3 | 4 | from .config import Config as __PluginConfigModel__, plugin_config 5 | 6 | 7 | __plugin_name__ = "Colab-NovelAI" 8 | __plugin_version__ = "0.2.2" 9 | __plugin_author__ = "T_EtherLeaF " 10 | 11 | __plugin_adapters__ = [OneBotV11Adapter] 12 | __plugin_des__ = "由Colab提供算力进行AI作画！" 13 | __plugin_usage__ = ( 14 | "让AI帮你画出好康的图片吧！\n" 15 | "触发指令：naifu [] []\n" 16 | " \n" 17 | "查看具体用法请发送指令：菜单 NovelAI <功能序号>\n" 18 | "子命令集如下：" 19 | ) 20 | 21 | 22 | __plugin_meta__ = PluginMetadata( 23 | name="NovelAI", 24 | description=__plugin_des__, 25 | usage=__plugin_usage__, 26 | config=__PluginConfigModel__, 27 | extra={ 28 | 'unique_name': __plugin_name__, 29 | 'author': __plugin_author__, 30 | 'version': __plugin_version__, 31 | 'adapters': __plugin_adapters__, 32 | 'menu_data': [ 33 | { 34 | 'func': '作图', 35 | 'trigger_method': '子命令：draw', 36 | 'trigger_condition': 'Anyone', 37 | 'brief_des': '告诉AI开始作图', 38 | 'detail_des': f'naifu draw ... [-i --undesired-content ...] ' 39 | f'[-a --sampling ] [-t --steps ] [-c --scale ] [-n --num ] ' 40 | f'[-s --size ] [-r --seed ]\n' 41 | f' \n' 42 | f'PROMPT: 必选参数，指定作画的关键词，以逗号分隔，必须为英语\n' 43 | f' \n' 44 | f'-i: 可选参数，指定作画中想避免出现的内容，以逗号分隔，必须为英语\n' 45 | f' \n' 46 | f'-a: 可选参数，指定采样器，支持以下几种，默认为k_euler_ancestral：\n' 47 | f' k_euler_ancestral, k_euler, k_lms\n' 48 | f' plms, ddim\n' 49 | f' \n' 50 | f'-t: 可选参数，指定优化图像的迭代次数，取值范围1~50，默认值为28\n' 51 | f' \n' 52 | f'-c: 可选参数，值越大越接近描述意图，值越小细节越少自由度越大，取值范围1.1~100，默认值为12\n' 53 | f' \n' 54 | f'-s: 可选参数，指定图片生成大小，支持以下几种，默认为512x768：\n' 55 | f' 384x640, 512x768, 512x1024 # Portrait\n' 56 | f' 640x384, 768x512, 1024x512 # Landscape\n' 57 | f' 512x512, 640x640, 1024x1024 # Square\n' 58 | f' \n' 59 | f'-n: 可选参数，指定图片生成数量，最大为{plugin_config.naifu_max}，默认值为1\n' 60 | f' \n' 61 | f'-r: 可选参数，指定图片生成种子，取值范围0 ~ 2^32-1，默认值为-1即随机' 62 | }, 63 | { 64 | 'func': '以图作图', 65 | 'trigger_method': '子命令：imgdraw', 66 | 'trigger_condition': 'Anyone', 67 | 'brief_des': '输入一张基准图片开始作图', 68 | 'detail_des': f'naifu imgdraw ... [-i --undesired-content ...] ' 69 | f'[-a --sampling ] [-t --steps ] [-c --scale ] [-n --num ] ' 70 | f'[-r --seed ] [-e strength ] [-o noise ]\n' 71 | f' \n' 72 | f'PROMPT: 必选参数，指定作画的关键词，以逗号分隔，必须为英语\n' 73 | f' \n' 74 | f'IMAGE: 必选参数，指定作画基准图片\n' 75 | f' \n' 76 | f'-i: 可选参数，指定作画中想避免出现的内容，以逗号分隔，必须为英语\n' 77 | f' \n' 78 | f'-a: 可选参数，指定采样器，支持以下几种，默认为k_euler_ancestral：\n' 79 | f' k_euler_ancestral, k_euler, k_lms\n' 80 | f' plms, ddim\n' 81 | f' \n' 82 | f'-t: 可选参数，指定优化图像的迭代次数，取值范围1~50，默认值为50\n' 83 | f' \n' 84 | f'-c: 可选参数，值越大越接近描述意图，值越小细节越少自由度越大，取值范围1.1~100，默认值为12\n' 85 | f' \n' 86 | f'-n: 可选参数，指定图片生成数量，最大为{plugin_config.naifu_max}，默认值为1\n' 87 | f' \n' 88 | f'-r: 可选参数，指定图片生成种子，取值范围0 ~ 2^32-1，默认值为-1即随机\n' 89 | f' \n' 90 | f'-e: 可选参数，值越低越接近原始图像，取值范围0~0.99，默认值为0.7\n' 91 | f' \n' 92 | f'-o: 可选参数，值增加会增加细节，一般应低于参数，取值范围0~0.99，默认值为0.2' 93 | }, 94 | { 95 | 'func': '白名单管理', 96 | 'trigger_method': '子命令：su', 97 | 'trigger_condition': 'Superuser', 98 | 'brief_des': '插件白名单用户组管理(无cd)', 99 | 'detail_des': 'naifu su []\n' 100 | ' \n' 101 | ' \n' 102 | '# Subcommand 1:\n' 103 | 'naifu su ls\n' 104 | '列出当前所有白名单用户\n' 105 | ' \n' 106 | '# Subcommand 2:\n' 107 | 'naifu su add ...\n' 108 | '添加白名单用户\n' 109 | '必须指定用户QQ号，可填写多个并以空格分隔\n' 110 | ' \n' 111 | '# Subcommand 3:\n' 112 | 'naifu su rm ...\n' 113 | '移除白名单用户\n' 114 | '必须指定用户QQ号，可填写多个并以空格分隔' 115 | }, 116 | { 117 | 'func': 'NSFW权限管理', 118 | 'trigger_method': '子命令：nsfw', 119 | 'trigger_condition': 'Superuser', 120 | 'brief_des': '管理允许绘制NSFW内容的用户与群组', 121 | 'detail_des': 'naifu nsfw []\n' 122 | ' \n' 123 | ' \n' 124 | '# Subcommand 1:\n' 125 | ' \n' 126 | 'naifu nsfw ls\n' 127 | ' \n' 128 | '列出当前所有允许NSFW内容的用户与群组\n' 129 | ' \n' 130 | '# Subcommand 2:\n' 131 | ' \n' 132 | 'naifu nsfw add [-u --uid ...] [-g --gid ...]\n' 133 | ' \n' 134 | '添加允许NSFW内容的用户或群组\n' 135 | ' \n' 136 | '-u: 可选参数，为用户QQ号，可填写多个并以空格分隔\n' 137 | ' \n' 138 | '-g: 可选参数，为群号，可填写多个并以空格分隔\n' 139 | ' \n' 140 | '当两个可选参数均未填写时，默认添加当前所处群聊的群号。\n' 141 | ' \n' 142 | '# Subcommand 3:\n' 143 | ' \n' 144 | 'naifu nsfw rm [-u --uid ...] [-g --gid ...]\n' 145 | ' \n' 146 | '移除允许NSFW内容的用户或群组\n' 147 | ' \n' 148 | '-u: 可选参数，为用户QQ号，可填写多个并以空格分隔\n' 149 | ' \n' 150 | '-g: 可选参数，为群号，可填写多个并以空格分隔\n' 151 | ' \n' 152 | '当两个可选参数均未填写时，默认移除当前所处群聊的群号。' 153 | } 154 | ] 155 | } 156 | ) 157 | -------------------------------------------------------------------------------- /nonebot_plugin_colab_novelai/_main.py: -------------------------------------------------------------------------------- 1 | import time 2 | from functools import reduce 3 | from argparse import Namespace 4 | from typing import Type, Any, Optional 5 | 6 | import asyncio 7 | from asyncer import asyncify 8 | from nonebot.log import logger 9 | from nonebot.params import T_State 10 | from nonebot.matcher import Matcher 11 | from nonebot.adapters.onebot.v11 import MessageEvent, GroupMessageEvent, MessageSegment 12 | 13 | from selenium.webdriver.common.by import By 14 | from selenium.webdriver.common.keys import Keys 15 | from selenium.webdriver.support.ui import WebDriverWait 16 | from selenium.webdriver.support import expected_conditions as ec 17 | from selenium.common.exceptions import NoSuchElementException, TimeoutException 18 | 19 | from .saveto import save_content 20 | from .access.bce import recognize_audio 21 | from .access.colab import NOTEBOOK_URL, run_colab 22 | from .access.cpolar import get_cpolar_authtoken, get_cpolar_url 23 | from .access.naifu import txt2img, img2img 24 | from .config import plugin_config 25 | from .utils import ( 26 | chrome_driver as driver, 27 | fetch_image_in_message, 28 | force_refresh_webpage, wait_and_click_element, 29 | preprocess_painting_parameters 30 | ) 31 | from .permissionManager import CooldownManager, NotSafeForWorkManager 32 | 33 | 34 | # ———————————————————— scheduled jobs ———————————————————— # 35 | @asyncify 36 | def handle_recaptcha() -> None: 37 | # listen to recaptcha iframe 38 | try: 39 | driver.switch_to.frame(driver.find_element(By.XPATH, '/html/body/colab-recaptcha-dialog/div/div/iframe')) 40 | except NoSuchElementException: 41 | return 42 | 43 | # click recaptcha checkbox 44 | checkbox = WebDriverWait(driver, 3).until(ec.element_to_be_clickable( 45 | (By.CSS_SELECTOR, "div.recaptcha-checkbox-checkmark") 46 | )) 47 | driver.execute_script("arguments[0].click();", checkbox) 48 | driver.switch_to.default_content() 49 | 50 | # detect if reCaptcha iframe still exists 51 | try: 52 | time.sleep(2) 53 | driver.switch_to.frame(driver.find_element( 54 | By.XPATH, '//iframe[contains(@src,"https://www.google.com/recaptcha/api2/bframe")]' 55 | )) 56 | except NoSuchElementException: 57 | logger.success("Colab ReCaptcha passed!") 58 | return 59 | # switch to audio task 60 | try: 61 | wait_and_click_element( 62 | driver, 63 | by=By.CSS_SELECTOR, value='button#recaptcha-audio-button' 64 | ) 65 | except TimeoutException: 66 | pass 67 | 68 | # analyze audio 69 | while True: 70 | # get audio link 71 | try: 72 | recap_audio_link = WebDriverWait(driver, 3).until( 73 | lambda t_driver: t_driver.find_element( 74 | By.XPATH, '//*[@id="rc-audio"]/div[7]/a' 75 | ) 76 | ).get_attribute("href") 77 | # if blocked by Google 78 | except TimeoutException: 79 | force_refresh_webpage(driver, NOTEBOOK_URL) 80 | logger.error("获取Colab ReCaptcha语音时被拦截！") 81 | return 82 | 83 | answer = recognize_audio(url=recap_audio_link) 84 | # input answer 85 | answer_box = wait_and_click_element( 86 | driver, 87 | by=By.XPATH, value='//*[@id="audio-response"]' 88 | ) 89 | answer_box.send_keys(answer, Keys.ENTER) 90 | 91 | try: 92 | # if answer is wrong 93 | WebDriverWait(driver, 3).until(ec.visibility_of_element_located( 94 | (By.XPATH, '//div[@class="rc-audiochallenge-error-message"]') 95 | )) 96 | # refresh audio 97 | wait_and_click_element( 98 | driver, 99 | by=By.XPATH, value='//*[@id="recaptcha-reload-button"]' 100 | ) 101 | logger.warning("Colab ReCaptcha音频验证答案错误！") 102 | # answer is correct 103 | except TimeoutException: 104 | logger.success("Colab ReCaptcha passed!") 105 | break 106 | 107 | # detect warning dialog: exceeding GPU usage limit 108 | driver.switch_to.default_content() 109 | try: 110 | WebDriverWait(driver, 3).until( 111 | lambda t_driver: t_driver.find_element( 112 | By.XPATH, '/html/body/colab-dialog' 113 | ) 114 | ) 115 | 116 | logger.error("当前账号已达到GPU用量上限！") 117 | force_refresh_webpage(driver, NOTEBOOK_URL) 118 | # no such warning 119 | except TimeoutException: 120 | pass 121 | 122 | 123 | async def access_colab_with_accounts() -> None: 124 | cpolar_authtoken = await get_cpolar_authtoken() 125 | 126 | # iter Google accounts 127 | for gmail, password in plugin_config.google_accounts.items(): 128 | try: 129 | await run_colab(gmail, password, cpolar_authtoken) 130 | except RuntimeError as e: 131 | logger.error(e) 132 | logger.info("尝试切换Google账号中...") 133 | continue 134 | 135 | # wait 10min until application startup complete 136 | start = time.time() 137 | while time.time() - start < 600: 138 | try: 139 | await get_cpolar_url() 140 | break 141 | except RuntimeError: 142 | logger.info(f"等待APP启动中... ({round(time.time() - start, 1)}s/600s) 当前账号：{gmail}") 143 | await asyncio.sleep(5) 144 | if time.time() - start >= 600: 145 | logger.error(f"Colab未成功启动，可能是因为ReCaptcha验证失败或已达到用量上限！当前账号：{gmail}") 146 | logger.info("尝试切换Google账号中...") 147 | continue 148 | 149 | logger.success("成功连接至APP！") 150 | while True: 151 | await asyncio.sleep(30) 152 | 153 | try: 154 | await get_cpolar_url() 155 | logger.info(f"当前Colab账号在线中：{gmail}") 156 | except RuntimeError: 157 | logger.warning(f"当前Colab账号已掉线：{gmail}") 158 | logger.info("尝试切换Google账号中...") 159 | break 160 | 161 | 162 | # ———————————————————— user interactions ———————————————————— # 163 | async def naifu_txt2img(matcher: Type[Matcher], event: MessageEvent, args: Namespace, **kwargs: Any) -> None: 164 | # check user cd 165 | user_id = event.get_user_id() 166 | remaining_cd = CooldownManager.get_user_cd(user_id) 167 | if remaining_cd > 0: 168 | await matcher.finish(f"你的cd还有{round(remaining_cd)}秒哦，可以稍后再来！", at_sender=True) 169 | # record user cd 170 | CooldownManager.record_cd(user_id, num=args.num) 171 | 172 | # check nsfw tag availability 173 | if isinstance(event, GroupMessageEvent): 174 | group_id = event.group_id 175 | else: 176 | group_id = None 177 | nsfw_available = NotSafeForWorkManager.check_nsfw_available(user_id, group_id) 178 | 179 | # preprocess parameters for painting 180 | try: 181 | params = await preprocess_painting_parameters(matcher=matcher, args=args, on_nsfw=nsfw_available) 182 | except ValueError as e: 183 | await matcher.finish(str(e), at_sender=True) 184 | width, height = map(int, args.size.split('x')) 185 | 186 | # draw the images 187 | try: 188 | await matcher.send("少女作画中...", at_sender=True) 189 | images = await txt2img(**params, width=width, height=height) 190 | image_segment = reduce( 191 | lambda img1, img2: img1+img2, 192 | [MessageSegment.image(image) for image in images] 193 | ) 194 | await matcher.send(image_segment, at_sender=True) 195 | 196 | # save the images 197 | await save_content(images, params["prompt"], params["uc"]) 198 | await matcher.finish() 199 | 200 | # if any exception occurs 201 | except (ValueError, RuntimeError) as e: 202 | # reset user cd 203 | CooldownManager.record_cd(user_id, num=0) 204 | await matcher.finish(str(e), at_sender=True) 205 | 206 | 207 | async def naifu_img2img( 208 | matcher: Type[Matcher], event: MessageEvent, state: T_State, args: Namespace, 209 | img: Optional[bytes] = None 210 | ) -> None: 211 | user_id = event.get_user_id() 212 | 213 | # baseimage already fetched 214 | if img is not None: 215 | # record user cd 216 | CooldownManager.record_cd(user_id, num=args.num) 217 | 218 | # draw the images 219 | try: 220 | await matcher.send("少女作画中...", at_sender=True) 221 | images = await img2img(**state, image=img) 222 | image_segment = reduce( 223 | lambda img1, img2: img1 + img2, 224 | [MessageSegment.image(image) for image in images] 225 | ) 226 | await matcher.send(image_segment, at_sender=True) 227 | 228 | # save the images 229 | await save_content(images, state["prompt"], state["uc"], baseimage=img) 230 | await matcher.finish() 231 | 232 | except (ValueError, RuntimeError) as e: 233 | CooldownManager.record_cd(user_id, num=0) 234 | await matcher.finish(str(e), at_sender=True) 235 | 236 | # baseimage isn't fetched yet 237 | 238 | # check user cd 239 | remaining_cd = CooldownManager.get_user_cd(user_id) 240 | if remaining_cd > 0: 241 | await matcher.finish(f"你的cd还有{round(remaining_cd)}秒哦，可以稍后再来！", at_sender=True) 242 | 243 | # check if nsfw tags are allowed in current chat 244 | if isinstance(event, GroupMessageEvent): 245 | group_id = event.group_id 246 | else: 247 | group_id = None 248 | nsfw_available = NotSafeForWorkManager.check_nsfw_available(user_id, group_id) 249 | 250 | # preprocess parameters for painting 251 | try: 252 | params = await preprocess_painting_parameters(matcher=matcher, args=args, on_nsfw=nsfw_available) 253 | assert 0 <= args.strength <= 0.99, "设置的strength需要在0到0.99之间！" 254 | assert 0 <= args.noise <= 0.99, "设置的noise需要在0到0.99之间！" 255 | except (ValueError, AssertionError) as e: 256 | await matcher.finish(str(e), at_sender=True) 257 | 258 | # inject parameters 259 | state.update( 260 | params, 261 | strength=round(args.strength, 2), 262 | noise=round(args.noise, 2) 263 | ) 264 | 265 | if event.reply: 266 | if (image := await fetch_image_in_message(event.reply.message)) is not None: 267 | state["baseimage"] = image 268 | if (image := await fetch_image_in_message(event.message)) is not None: 269 | state["baseimage"] = image 270 | 271 | # wait for "got" procedure and call this function again, with image injected 272 | # draw the images 273 | --------------------------------------------------------------------------------