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),
                     afterprocess: Option(float, name="후처리정도", description="그림의 후처리 정도를 정합니다. 0.99에 가까울수록 후처리랑 강하게 합니다", required=False, min_value=0,max_value=0.99, default=0.5)  
                ):
        if not await database.get_guild(ctx.guild.id): return await ctx.respond(embed=ChristmasEmbed(title="❌ 에러!", description="서버가 가입되어있지 않아요! 서버를 가입해주세요!", color=Color.red()),ephemeral=True)
        nsfw = False
        afterprocess = 1 - afterprocess
        if ctx.channel.is_nsfw(): nsfw = True
        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]
        # 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=True)
        if 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=True, file=File(fp=io.BytesIO(buffer), filename="image.png"))
        
def setup(bot):
    bot.add_cog(CAiart(bot))