267 lines
9.1 KiB
Python
267 lines
9.1 KiB
Python
import asyncio
|
||
import base64
|
||
import random
|
||
import time
|
||
from io import BytesIO
|
||
|
||
import aiohttp
|
||
from nonebot import get_driver
|
||
from nonebot.log import logger
|
||
from PIL import Image
|
||
|
||
from ..config import config
|
||
from ..utils import png2jpg
|
||
from ..utils.data import shapemap
|
||
|
||
|
||
class AIDRAW_BASE:
|
||
max_resolution: int = 16
|
||
sampler: str
|
||
|
||
def __init__(
|
||
self,
|
||
user_id: str,
|
||
group_id: str,
|
||
tags: str = "",
|
||
ntags: str = "",
|
||
seed: int = None,
|
||
scale: int = None,
|
||
steps: int = None,
|
||
batch: int = None,
|
||
strength: float = None,
|
||
noise: float = None,
|
||
shape: str = "p",
|
||
model: str = None,
|
||
sampler: str = "",
|
||
hires: bool = False,
|
||
**kwargs,
|
||
):
|
||
"""
|
||
AI绘画的核心部分,将与服务器通信的过程包装起来,并方便扩展服务器类型
|
||
|
||
:user_id: 用户id,必须
|
||
:group_id: 群聊id,如果是私聊则应置为0,必须
|
||
:tags: 图像的标签
|
||
:ntags: 图像的反面标签
|
||
:seed: 生成的种子，不指定的情况下随机生成
|
||
:scale: 标签的参考度，值越高越贴近于标签,但可能产生过度锐化。范围为0-30,默认11
|
||
:steps: 训练步数。范围为1-50,默认28.以图生图时强制50
|
||
:batch: 同时生成数量
|
||
:strength: 以图生图时使用,变化的强度。范围为0-1,默认0.7
|
||
:noise: 以图生图时使用,变化的噪音,数值越大细节越多,但可能产生伪影,不建议超过strength。范围0-1,默认0.2
|
||
:shape: 图像的形状，支持"p""s""l"三种，同时支持以"x"分割宽高的指定分辨率。
|
||
该值会被设置限制，并不会严格遵守输入
|
||
类初始化后,该参数会被拆分为:width:和:height:
|
||
:model: 指定的模型，模型名称在配置文件中手动设置。不指定模型则按照负载均衡自动选择
|
||
|
||
AIDRAW还包含了以下几种内置的参数
|
||
:status: 记录了AIDRAW的状态,默认为0等待中(处理中)
|
||
非0的值为运行完成后的状态值,200和201为正常运行,其余值为产生错误
|
||
:result: 当正常运行完成后,该参数为一个包含了生成图片bytes信息的数组
|
||
:maxresolution: 一般不用管，用于限制不同服务器的最大分辨率
|
||
如果你的SD经过了魔改支持更大分辨率可以修改该值并重新设置宽高
|
||
:cost: 记录了本次生成需要花费多少点数，自动计算
|
||
:signal: asyncio.Event类,可以作为信号使用。仅占位，需要自行实现相关方法
|
||
"""
|
||
self.status: int = 0
|
||
self.result: list = []
|
||
self.signal: asyncio.Event = None
|
||
self.model = model
|
||
self.time = time.strftime("%Y-%m-%d %H:%M:%S")
|
||
self.user_id: str = user_id
|
||
self.tags: str = tags
|
||
self.seed: list[int] = [seed or random.randint(0, 4294967295)]
|
||
self.group_id: str = group_id
|
||
self.scale: int = int(scale or 11)
|
||
self.strength: float = strength or 0.7
|
||
self.batch: int = batch or 1
|
||
self.steps: int = steps or 28
|
||
self.noise: float = noise or 0.2
|
||
self.ntags: str = ntags
|
||
self.img2img: bool = False
|
||
self.image: str = None
|
||
self.width, self.height = self.extract_shape(shape)
|
||
self.sampler = sampler
|
||
self.hires = hires
|
||
# 数值合法检查
|
||
if self.steps <= 0 or self.steps > (50 if not config.novelai_paid else 28):
|
||
self.steps = 28
|
||
if self.strength < 0 or self.strength > 1:
|
||
self.strength = 0.7
|
||
if self.noise < 0 or self.noise > 1:
|
||
self.noise = 0.2
|
||
if self.scale <= 0 or self.scale > 30:
|
||
self.scale = 11
|
||
# 多图时随机填充剩余seed
|
||
for i in range(self.batch - 1):
|
||
self.seed.append(random.randint(0, 4294967295))
|
||
# 计算cost
|
||
self.update_cost()
|
||
|
||
def extract_shape(self, shape: str):
|
||
"""
|
||
将shape拆分为width和height
|
||
"""
|
||
if shape:
|
||
if "x" in shape:
|
||
width, height, *_ = shape.split("x")
|
||
if width.isdigit() and height.isdigit():
|
||
return self.shape_set(int(width), int(height))
|
||
else:
|
||
return shapemap.get(shape)
|
||
else:
|
||
return shapemap.get(shape)
|
||
else:
|
||
return (512, 768)
|
||
|
||
def update_cost(self):
|
||
"""
|
||
更新cost
|
||
"""
|
||
if config.novelai_paid == 1:
|
||
anlas = 0
|
||
if (self.width * self.height > 409600) or self.image or self.batch > 1:
|
||
anlas = round(
|
||
self.width
|
||
* self.height
|
||
* self.strength
|
||
* self.batch
|
||
* self.steps
|
||
/ 2293750
|
||
)
|
||
if anlas < 2:
|
||
anlas = 2
|
||
if self.user_id in get_driver().config.superusers:
|
||
self.cost = 0
|
||
else:
|
||
self.cost = anlas
|
||
elif config.novelai_paid == 2:
|
||
anlas = round(
|
||
self.width
|
||
* self.height
|
||
* self.strength
|
||
* self.batch
|
||
* self.steps
|
||
/ 2293750
|
||
)
|
||
if anlas < 2:
|
||
anlas = 2
|
||
if self.user_id in get_driver().config.superusers:
|
||
self.cost = 0
|
||
else:
|
||
self.cost = anlas
|
||
else:
|
||
self.cost = 0
|
||
|
||
def add_image(self, image: bytes):
|
||
"""
|
||
向类中添加图片，将其转化为以图生图模式
|
||
也可用于修改类中已经存在的图片
|
||
"""
|
||
# 根据图片重写长宽
|
||
tmpfile = BytesIO(image)
|
||
image_ = Image.open(tmpfile)
|
||
width, height = image_.size
|
||
self.width, self.height = self.shape_set(width, height)
|
||
self.image = str(base64.b64encode(image), "utf-8")
|
||
self.steps = 50
|
||
self.img2img = True
|
||
self.update_cost()
|
||
|
||
def shape_set(self, width: int, height: int):
|
||
"""
|
||
设置宽高
|
||
"""
|
||
limit = 1024 if config.paid else 640
|
||
if width * height > pow(min(config.novelai_size, limit), 2):
|
||
if width <= height:
|
||
ratio = height / width
|
||
width: float = config.novelai_size / pow(ratio, 0.5)
|
||
height: float = width * ratio
|
||
else:
|
||
ratio = width / height
|
||
height: float = config.novelai_size / pow(ratio, 0.5)
|
||
width: float = height * ratio
|
||
base = round(max(width, height) / 64)
|
||
if base > self.max_resolution:
|
||
base = self.max_resolution
|
||
if width <= height:
|
||
return (round(width / height * base) * 64, 64 * base)
|
||
else:
|
||
return (64 * base, round(height / width * base) * 64)
|
||
|
||
async def post_(self, header: dict, post_api: str, json: dict):
|
||
"""
|
||
向服务器发送请求的核心函数，不要直接调用，请使用post函数
|
||
:header: 请求头
|
||
:post_api: 请求地址
|
||
:json: 请求体
|
||
"""
|
||
# 请求交互
|
||
async with aiohttp.ClientSession(headers=header) as session:
|
||
# 向服务器发送请求
|
||
async with session.post(post_api, json=json) as resp:
|
||
if resp.status not in [200, 201]:
|
||
logger.error(await resp.text())
|
||
raise RuntimeError(f"与服务器沟通时发生{resp.status}错误")
|
||
img = await self.fromresp(resp)
|
||
logger.debug(f"获取到返回图片，正在处理")
|
||
|
||
# 将图片转化为jpg
|
||
if config.novelai_save == 1:
|
||
image_new = await png2jpg(img)
|
||
else:
|
||
image_new = base64.b64decode(img)
|
||
self.result.append(image_new)
|
||
return image_new
|
||
|
||
async def fromresp(self, resp):
|
||
"""
|
||
处理请求的返回内容，不要直接调用，请使用post函数
|
||
"""
|
||
img: str = await resp.text()
|
||
return img.split("data:")[1]
|
||
|
||
def run(self):
|
||
"""
|
||
运行核心函数，发送请求并处理
|
||
"""
|
||
pass
|
||
|
||
def keys(self):
|
||
return (
|
||
"seed",
|
||
"scale",
|
||
"strength",
|
||
"noise",
|
||
"sampler",
|
||
"model",
|
||
"steps",
|
||
"width",
|
||
"height",
|
||
"img2img",
|
||
)
|
||
|
||
def __getitem__(self, item):
|
||
return getattr(self, item)
|
||
|
||
def format(self):
|
||
dict_self = dict(self)
|
||
list = []
|
||
str = ""
|
||
for i, v in dict_self.items():
|
||
str += f"{i}={v}\n"
|
||
list.append(str)
|
||
list.append(f"tags={self.tags}\n")
|
||
list.append(f"ntags={self.ntags}")
|
||
return list
|
||
|
||
def __repr__(self):
|
||
return (
|
||
f"time={self.time}\nuser_id={self.user_id}\ngroup_id={self.group_id}\ncost={self.cost}\nbatch={self.batch}\n"
|
||
+ "".join(self.format())
|
||
)
|
||
|
||
def __str__(self):
|
||
return self.__repr__().replace("\n", ";")
|