2024-03-17 00:42:36 +00:00
import re
import os
import random
import base64
from discord import InputTextStyle , Interaction , Member
from discord . ui import Modal , InputText
from dotenv import load_dotenv , find_dotenv
#from RUNA.Module import check_curse
from RUNA . UI . Embed import Mail_Embed , Aiart_Embed
from RUNA . UI . Buttons import Mail_Confirm_Button
from RUNA . Module import process_prompt , post_gpu_server , base64_to_image
from RUNA . Tagging import Tagging
from RUNA . Cogs . Event import model
class Send_Mail_Modal ( Modal ) :
def __init__ ( self , reciveuser : Member , editmode : bool , * args , * * kwargs ) :
self . reciveuser = reciveuser
self . editmode = editmode
super ( ) . __init__ ( timeout = None , * args , * * kwargs )
self . add_item ( InputText ( label = " 제목 " , placeholder = " 제목을 입력해주세요. " , style = InputTextStyle . short , required = True , custom_id = " mail_title " ) )
self . add_item ( InputText ( label = " 내용 " , placeholder = " 내용을 입력해주세요. " , style = InputTextStyle . long , required = True , custom_id = " mail_content " ) )
async def callback ( self , interaction : Interaction ) :
if self . editmode == False :
await interaction . response . send_message ( embed = Mail_Embed . mail_confirm ( self . children [ 0 ] . value , self . children [ 1 ] . value , self . reciveuser ) , view = Mail_Confirm_Button ( recive_user = self . reciveuser , title = self . children [ 0 ] . value , description = self . children [ 1 ] . value , editmode = False ) , ephemeral = True )
else :
await interaction . response . send_message ( embed = Mail_Embed . mail_confirm_edit ( self . children [ 0 ] . value , self . children [ 1 ] . value , self . reciveuser ) , view = Mail_Confirm_Button ( recive_user = self . reciveuser , title = self . children [ 0 ] . value , description = self . children [ 1 ] . value , editmode = True ) , ephemeral = True )
class Aiart ( Modal ) :
def __init__ ( self , show : bool , allownsfw : bool , res : list , style1 : float , style2 : float , afterprocess : bool , * args , * * kwargs ) :
self . show = show
self . allownsfw = allownsfw
self . res = res
self . style1 = style1
self . style2 = style2
self . afterprocess = afterprocess
super ( ) . __init__ ( timeout = None , * args , * * kwargs )
self . add_item ( InputText ( label = " 그람에 포함시키고 싶은 태그를 입력해주세요. " , placeholder = " 한글도 입력 가능해요! 태그는 문장형도 알아듣지만 ,를 통해 구별해요 " , style = InputTextStyle . long , required = True , custom_id = " aiart_tag " ) )
self . add_item ( InputText ( label = " 그림에서 제거하고 싶은 태그를 입력해주세요. " , placeholder = " 한글도 입력 가능해요! 태그는 문장형도 알아듣지만 ,를 통해 구별해요 " , style = InputTextStyle . long , required = False , custom_id = " aiart_remove_tag " ) )
async def callback ( self , interaction : Interaction ) :
prompt = self . children [ 0 ] . value
remove = self . children [ 1 ] . value
#255자까지 자름
original_prompt = prompt
original_remove = remove
if len ( prompt ) > 255 :
original_prompt = original_prompt [ : 255 ]
if len ( remove ) > 255 :
original_remove = original_remove [ : 255 ]
await interaction . response . send_message ( embed = Aiart_Embed . generating ( ) , ephemeral = self . show )
avatars = None
if re . search ( r " \ { [0-9] {18} \ } " , prompt ) :
user_id = re . search ( r " \ { [0-9] {18} \ } " , prompt ) . group ( ) [ 1 : - 1 ]
user = await interaction . guild . fetch_member ( int ( user_id ) )
if user == None :
avatars = None
else :
avatars = user . avatar
#prompt: str, remove: str, res: list, isnsfw: bool, style1: float, style2: float, afterprocess: float
payload = await process_prompt ( prompt , remove , self . res , self . allownsfw , self . style1 , self . style2 , self . afterprocess , avatars )
result = await post_gpu_server ( f " { random . choices ( list ( os . getenv ( " AI_GATEWAT " ) ) ) } /sdapi/v1/txt2img " , payload )
if result [ " status " ] != True :
return await interaction . edit_original_response ( embed = Aiart_Embed . failed_generate ( ) )
if self . allownsfw is False :
tag = Tagging ( )
result2 = await tag . predict ( base64 . b64decode ( result [ " data " ] [ " images " ] [ 0 ] ) )
if result2 is None :
return await interaction . edit_original_response ( embed = Aiart_Embed . failed_generate ( ) )
rating = max ( result2 [ 2 ] , key = result2 [ 2 ] . get )
if rating == " explicit " or rating == " questionable " :
return await interaction . edit_original_response ( embed = Aiart_Embed . gen_nsfw ( ) )
try :
attachment = await base64_to_image ( result [ " data " ] [ " images " ] [ 0 ] )
except Exception as e :
return await interaction . edit_original_response ( embed = Aiart_Embed . failed_generate ( ) )
await interaction . edit_original_response ( embed = Aiart_Embed . sucess_generate ( original_prompt , original_remove ) , file = attachment )
2023-11-26 15:27:32 +00:00