runabot/Christmas/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 Christmas.Database import database
from Christmas.UI.Embed import ChristmasEmbed, Aiart_Embed
from Christmas.UI.Modal import Aiart
from Christmas.Tagging import Tagging
from Christmas.Cogs.Event import model
class CAiart(Cog):
    def __init__(self, bot):
        self.bot = bot
        
    
    ART = SlashCommandGroup("그림", "그림 기능")
    
    @ART.command(name="생성", description="크돌이가 그림을 그려줍니다.")
    @cooldown(1, 10, BucketType.user)
    async def _생성(self, ctx, ress: 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),
                     highquality: 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=ChristmasEmbed(title="❌ 에러!", description="서버가 가입되어있지 않아요! 서버를 가입해주세요!", color=Color.red()),ephemeral=True)
        nsfw = False
        if highquality == "사용하기": afterprocess = True
        else: afterprocess = False
        if ctx.channel.is_nsfw():
            nsfw = True
        else:
            nsfw = False
        if shows == "보여주기": shows = False
        else: shows = True
        if ress == "1:1": resoultion = [512,512]
        elif ress == "2:3": resoultion = [512,768]
        elif ress == "7:4": resoultion = [896,512]
        elif ress == "3:4": resoultion = [600,800]
        await database.put_use_art(ctx.guild.id)
        if afterprocess == True:
            result = await database.use_guild_art_point(ctx.guild.id)
            if result is False: return await ctx.respond(embed=ChristmasEmbed(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=afterprocess, show=shows)
        await ctx.send_modal(modal)
        
    @ART.command(name="분석", description="그림을 분석합니다.")
    @cooldown(1, 10, BucketType.user)
    @guild_only()
    async def _분석(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=ChristmasEmbed(title="❌ 에러!", description="서버가 가입되어있지 않아요! 서버를 가입해주세요!", color=Color.red()),ephemeral=True)
        if not file.content_type.startswith("image/"): return await ctx.respond(embed=ChristmasEmbed(title="❌ 에러!", description="그림 파일만 업로드해주세요!", color=Color.red()),ephemeral=True)
        buffer = await file.read()
        taging = Tagging(model=model)
        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 _포인트(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=ChristmasEmbed(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=ChristmasEmbed(title="🎨 아트포인트", description=f"현재 서버의 아트포인트는 {point}개에요!\n 이서버에서는 그림 기능을 {use}번 사용했어요! ", color=Color.green()), ephemeral=False)
        

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