runabot/Christmas/UI/Modal.py

import re
from discord import InputTextStyle, Interaction, Member
from discord.ui import Modal, InputText

#from Christmas.Module import check_curse

from Christmas.UI.Embed import Mail_Embed, Aiart_Embed
from Christmas.UI.Buttons import Mail_Confirm_Button
from Christmas.Module import process_prompt, post_gpu_server, base64_to_image
from Christmas.config import ChristmasConfig

BLOCKTAG = [
    "nsfw",
    "nude",
    "nipples",
    "nipple",
    "pussy",
    "public hair",
    "gay",
    "lesbian",
    "corpse",
    "no panties",
    "no panty",
    "no bra",
    "bra",
    "panty",
    "panties",
    "underwear",
    "undergarment",
    "underpants",
    "underpant",
    "blowjob",
    "sex",
    "sexy",
    "pennis"
    "realistic",
    "open breasts",
    "breasts"
]

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: float, *args, **kwargs):
        self.show = show
        self.allownsfw = allownsfw
        self.res = res
        self.style1 = style1
        self.style2 = style2
        self.config = ChristmasConfig()
        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]
        if self.allownsfw is False:
            for i in BLOCKTAG:
                if i in prompt:
                    prompt = prompt.replace(i, "")
        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
        config = ChristmasConfig()
        #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)
        config = ChristmasConfig()
        result = await post_gpu_server(f"{config.AI}/sdapi/v1/txt2img", payload)
        if result["status"] != True:
            return await interaction.edit_original_response(embed=Aiart_Embed.failed_generate())
        else:
            try:
                #print(result)
                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)