85 lines
4.9 KiB
Python
85 lines
4.9 KiB
Python
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)
|
|
|