staskaer / nonebot_plugin_realesrgan Goto Github PK

from .utils import *
from nonebot.adapters.onebot.v11 import Message, MessageEvent, MessageSegment, Bot, Event
from service.gocqhttp_guild_patch import Message, GuildMessageEvent, MessageSegment, Bot, Event
from nonebot.params import T_State, CommandArg
from nonebot.plugin import on_command
from nonebot.adapters.onebot.v11.helpers import HandleCancellation
from json import loads
from nonebot import get_driver

try:
    api = loads(get_driver().config.json())["realesrgan_api"]
except:
    api = 'https://hf.space/embed/mute23/linxi_resolutionRatio/api/predict/'

real_esrgan = on_command(
    "重建", aliases={"real-esrgan", "超分辨率重建", "图片重建", "real_esrgan", "图像重建"}, priority=30
)


@real_esrgan.handle()
async def real_esrgan_handle_first(
    bot: Bot,
    event: MessageEvent | GuildMessageEvent,
    state: T_State,
    args: Message = CommandArg(),
):
    state['id'] = event.get_user_id()
    for seg in args:
        if seg.type == "text":
            state["mode"] = seg.data["text"].strip()
        if seg.type == "image":
            state['img'] = event.message['image']
            break


@real_esrgan.got("mode", prompt="请提供重建模式(二次元:anime，其他:base)，模式不绝对，可以任意选", parameterless=[HandleCancellation("已取消")])
async def real_esrgan_get_mode(event: MessageEvent | GuildMessageEvent, state: T_State):
    mode = str(state["mode"]).strip()
    if mode not in ["anime", "base"]:
        await real_esrgan.reject('"base" | "anime", 二选一')


@real_esrgan.got("img", prompt="请上传需要超分辨率重建的图片", parameterless=[HandleCancellation("已取消")])
async def real_esrgan_handle_img(event: MessageEvent | GuildMessageEvent    , state: T_State):
    # 先拿到需要转换的图
    for seg in state["img"]:
        if seg.type == "image":
            img = await get_img(seg.data["url"])
            break
    else:
        await real_esrgan.finish(Message(f"[CQ:at,qq={state['id']}]不是图捏"))
    # 下面来处理图片
    try:
        json_data = img_encode_to_json(img, state['mode'])  # 先获取图片并进行编码
        if json_data is None:
            await real_esrgan.finish(Message(f"[CQ:at,qq={state['id']}]服务器无法接收到这张图捏,要不重试试试？"))
        result = await get_result(json_data, api=api)  # 然后进行超分辨率重建
        if result is None:
            await real_esrgan.finish(Message(f"[CQ:at,qq={state['id']}]这张图没能被正确解析，可能网络连接失败或者是由于远程服务器免费额度耗尽，如果是后者建议联系管理员使用自建仓库(免费)"))
        img = img_decode_from_json(result)  # 获取重建后图片并进行解码发送
        await real_esrgan.finish(MessageSegment.image(img))
    except Exception as e:
        ...
        # print('错误类型是', e.__class__.__name__)
        # print('错误明细是', e)

from json import loads
from base64 import b64encode, b64decode
import aiohttp


######### 实现对图片的base64编码和从返回数据中解析出原始图片 #########


def img_encode_to_json(img: bytes, mode="base") -> dict:
    '''
    对二进制格式的图片进行base64编码，并组织成json格式返回，用于request请求
    Args:
        img (bytes): 二进制格式图片
    Returns:
        dict: 返回的json数据
    '''
    base64_data = b64encode(img)
    base64_data = str(base64_data, 'utf-8')
    return {
        'fn_index': 0,
        'data': [
            'data:image/jpeg;base64,{}'.format(base64_data),
            f'{mode}',
        ]
    }


def img_decode_from_json(response_str: str) -> bytes:
    '''
    将返回的json解析处img的base64再解码返回
    Args:
        response_str (str): 字符串格式的json
    Returns:
        bytes: img的二进制格式
    '''
    result = loads(response_str)
    img_base64 = result['data'][0].split("base64,")[1]
    return b64decode(img_base64)


######### 网络请求部分 #########

async def get_img(img_url: str) -> bytes:
    '''
    将收到的图片下载下来，并转换成二进制格式
    Args:
        img_url (str): 图片url地址
    Returns:
        bytes: 二进制格式的图片
    '''
    async with aiohttp.ClientSession() as session:
        async with session.get(img_url) as resp:
            result = await resp.content.read()
    if not result:
        return None
    return result


async def get_result(json_data: dict) -> str:
    '''
    来构造请求并获取返回的重建后的图像
    Args:
        json_data (dict): 对图片编码后的数据
    Returns:
        str: 返回的json格式数据
    '''
    headers = {
        'authority': 'hf.space',
        'accept': '*/*',
        'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8',
        'cache-control': 'no-cache',
        # Already added when you pass json=
        # 'content-type': 'application/json',
        'dnt': '1',
        'origin': 'https://hf.space',
        'pragma': 'no-cache',
        'referer': 'https://hf.space/embed/mute23/linxi_resolutionRatio/+?__theme=light',
        'sec-ch-ua': '"Google Chrome";v="105", "Not)A;Brand";v="8", "Chromium";v="105"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"Windows"',
        'sec-fetch-dest': 'empty',
        'sec-fetch-mode': 'cors',
        'sec-fetch-site': 'same-origin',
        'sec-gpc': '1',
        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/105.0.0.0 Safari/537.36',
    }
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post('https://hf.space/embed/mute23/linxi_resolutionRatio/api/predict/', headers=headers, json=json_data, timeout=120) as resp:
                result = await resp.text()
        if not result:
            return None
        return result
    except:
        return None