runabot/RUNA/Cogs/Commands_Aiart.py

import io
from discord import SlashCommandGroup, Option, Member, Color, File, Attachment
from discord.ext.commands import Cog, BucketType, cooldown, guild_only, Context
from RUNA.Database import database
from RUNA.UI.Embed import RunaEmbed, Aiart_Embed
from RUNA.UI.Modal import Aiart
from RUNA.Tagging import Tagging
from RUNA.Cogs.Event import model


class CogAiart(Cog):
    def __init__(self, bot):
        self.bot = bot

    ART = SlashCommandGroup("그림", "AI 그림 기능")

    @ART.command(name="생성", description="루나가 그림을 그려줍니다.")
    @cooldown(1, 10, BucketType.user)
    async def _generate(self, ctx: Context,
                        resolution: Option(str, name="해상도", description="그림의 해상도를 입력해주세요. (기본값: 512x512)",
                                      required=False, choices=["1:1", "2:3", "7:4", "3:4"], default="1:1"),
                        shows: Option(str, name="보기여부", description="생성된 그림을 나를 제외한 사람한테 보여줄지 말지를 지정합니다.",
                                      choices=["보여주기", "보여주지 말기"], required=False, default="보여주기"),
                        style1: Option(float, name="디테일표현",
                                       description="그림의 디테일정도를 지정합니다. -1에 가까울수록 단순한 그림이 나오고 1에 가까울수록 높은 디테일을 구현 합니다.",
                                       required=False, min_value=-1, max_value=1, default=0),
                        style2: Option(float, name="광선표현",
                                       description="그림의 광선 디테일 정도를 지정합니다. 0.8에 가까울수록 더 부드럽고 좋은 광선표현을 하지만 그림의 디테일이 떨어지거나 캐릭터일 경우 손이 제대로 생성되지 않을수 있습니다.",
                                       required=False, min_value=0, max_value=0.8, default=0),
                        high_quality: Option(str, name="고퀄리티모드",
                                        description="그림을 뽑을때 고퀄리티 모드를 사용할지 정합니다. 고퀄리티 모드는 크돌이의 추천 이후 얻는 아트포인트로 사용할수 있습니다.",
                                        required=False, choices=["사용하기", "사용하지 않기"], default="사용하지 않기")
                        ):
        if ctx.guild is not None and not await database.get_guild(ctx.guild.id): return await ctx.respond(
            embed=RunaEmbed(title="❌ 에러!", description="서버가 가입되어있지 않아요! 서버를 가입해주세요!", color=Color.red()),
            ephemeral=True)
        nsfw = False
        if high_quality == "사용하기":
            after_process = True
        else:
            after_process = False
        if ctx.channel.is_nsfw():
            nsfw = True
        else:
            nsfw = False
        if shows == "보여주기":
            shows = False
        else:
            shows = True
        if resolution == "1:1":
            resoultion = [512, 512]
        elif resolution == "2:3":
            resoultion = [512, 768]
        elif resolution == "7:4":
            resoultion = [896, 512]
        elif resolution == "3:4":
            resoultion = [600, 800]
        await database.put_use_art(ctx.guild.id)
        if after_process == True:
            result = await database.use_guild_art_point(ctx.guild.id)
            if result is False: return await ctx.respond(
                embed=RunaEmbed(title="❌ 에러!", description="아트포인트가 부족해요! 아트포인트를 충전해주세요! 아트포인트는 ``/추천``명령어를 통해 충전 가능해요!",
                                color=Color.red()), ephemeral=True)
        # allownsfw: bool, res: list, style1: float, style2: float, afterprocess: float
        modal = Aiart(title="태그를 입력해주세요", allownsfw=nsfw, res=resoultion, style1=style1, style2=style2,
                      afterprocess=after_process, show=shows)
        await ctx.send_modal(modal)

    @ART.command(name="분석", description="그림을 분석합니다.")
    @cooldown(1, 10, BucketType.user)
    @guild_only()
    async def _evulate(self, ctx: Context,
                       file: Option(Attachment, name="파일", description="분석할 그림을 업로드해주세요.", required=True)):
        await ctx.defer(ephemeral=False)
        if ctx.guild is not None and not await database.get_guild(ctx.guild.id): return await ctx.respond(
            embed=RunaEmbed(title="❌ 에러!", description="서버가 가입되어있지 않아요! 서버를 가입해주세요!", color=Color.red()),
            ephemeral=True)
        if not file.content_type.startswith("image/"): return await ctx.respond(
            embed=RunaEmbed(title="❌ 에러!", description="그림 파일만 업로드해주세요!", color=Color.red()), ephemeral=True)
        buffer = await file.read()
        taging = Tagging()
        tag = await taging.predict(buffer)
        rating = tag[2]
        tags = tag[4]
        hangul = {
            "general": "건전함",
            "sensitive": "매우조금 불건전",
            "questionable": "조금 불건전",
            "explicit": "매우 불건전"
        }
        ratings = max(rating, key=rating.get)
        rating = hangul[ratings]
        sorted_tags = sorted(tags.items(), key=lambda x: x[1], reverse=True)[:8]
        #  UnicodeDecodeError: 'utf-8' codec can't decode byte 0xff in position 0: invalid start byte
        await ctx.respond(embed=Aiart_Embed.evalate(sorted_tags, rating), ephemeral=False,
                          file=File(fp=io.BytesIO(buffer), filename="image.png"))

    @ART.command(name="포인트", description="아트포인트를 확인합니다.")
    @cooldown(1, 10, BucketType.user)
    @guild_only()
    async def _point(self, ctx: Context):
        await ctx.defer(ephemeral=False)
        if ctx.guild is not None and not await database.get_guild(ctx.guild.id): return await ctx.respond(
            embed=RunaEmbed(title="❌ 에러!", description="서버가 가입되어있지 않아요! 서버를 가입해주세요!", color=Color.red()),
            ephemeral=True)
        point = await database.get_guild_art_point(ctx.guild.id)
        use = await database.get_use_art(ctx.guild.id)
        await ctx.respond(
            embed=RunaEmbed(title="🎨 아트포인트", description=f"현재 서버의 아트포인트는 {point}개에요!\n 이서버에서는 그림 기능을 {use}번 사용했어요! ",
                            color=Color.green()), ephemeral=False)


def setup(bot):
    bot.add_cog(CogAiart(bot))