runabot/Christmas/UI/Modal.py

import re
import base64
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
from Christmas.Tagging import Tagging
from Christmas.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.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]
        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())
        if self.allownsfw is False:
            tag = Tagging(model=model)
            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)