2023-12-02 06:03:50 +00:00
import io
from discord import SlashCommandGroup , Option , Member , Color , File , Attachment
from discord . ext . commands import Cog , BucketType , cooldown , guild_only , Context
2023-11-26 15:27:32 +00:00
from Christmas . Database import database
2023-12-02 06:03:50 +00:00
from Christmas . UI . Embed import ChristmasEmbed , Aiart_Embed
2023-11-26 15:27:32 +00:00
from Christmas . UI . Modal import Aiart
2023-12-02 06:03:50 +00:00
from Christmas . Tagging import Tagging
from Christmas . Cogs . Event import model
2023-11-25 17:28:48 +00:00
class CAiart ( Cog ) :
def __init__ ( self , bot ) :
self . bot = bot
2023-11-26 15:27:32 +00:00
ART = SlashCommandGroup ( " 그림 " , " 그림 기능 " )
2023-11-25 17:28:48 +00:00
2023-11-26 15:27:32 +00:00
@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
2023-12-02 06:03:50 +00:00
if shows == " 보여주기 " : shows = False
else : shows = True
2023-11-26 15:27:32 +00:00
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 )
2023-12-02 06:03:50 +00:00
@ART.command ( name = " 분석 " , description = " 그림을 분석합니다. " )
@cooldown ( 1 , 10 , BucketType . user )
@guild_only ( )
async def _분석 ( self , ctx : Context , file : Option ( Attachment , name = " 파일 " , description = " 분석할 그림을 업로드해주세요. " , required = True ) ) :
2023-12-05 13:30:58 +00:00
await ctx . defer ( ephemeral = False )
2023-12-02 06:03:50 +00:00
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
2023-12-05 13:30:58 +00:00
await ctx . respond ( embed = Aiart_Embed . evalate ( sorted_tags , rating ) , ephemeral = False , file = File ( fp = io . BytesIO ( buffer ) , filename = " image.png " ) )
2023-12-02 06:03:50 +00:00
2023-11-25 17:28:48 +00:00
def setup ( bot ) :
bot . add_cog ( CAiart ( bot ) )