diff --git a/Christmas/Cogs/Commands_Aiart.py b/Christmas/Cogs/Commands_Aiart.py index f1df62d..b403a48 100644 --- a/Christmas/Cogs/Commands_Aiart.py +++ b/Christmas/Cogs/Commands_Aiart.py @@ -19,18 +19,26 @@ class CAiart(Cog): 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) + highquality: Option(str, name="고퀄리티모드", description="그림을 뽑을때 고퀄리티 모드를 사용할지 정합니다. 고퀄리티 모드는 크돌이의 추천 이후 얻는 아트포인트로 사용할수 있습니다.", required=False, choices=["사용하기", "사용하지 않기"], default="사용하지 않기") ): if ctx.guild is not None and 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 + if highquality == "사용하기": afterprocess = True + else: afterprocess = False + if ctx.channel.is_nsfw(): + nsfw = True + else: + nsfw = False if shows == "보여주기": shows = False else: shows = True 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] + await database.put_use_art(ctx.guild.id) + if afterprocess == True: + result = await database.use_guild_art_point(ctx.guild.id) + if result is False: return await ctx.respond(embed=ChristmasEmbed(title="❌ 에러!", description="아트포인트가 부족해요! 아트포인트를 충전해주세요! 아트포인트는 ``/추천``명령어를 통해 충전 가능해요!", color=Color.red()),ephemeral=True) # 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) @@ -58,6 +66,17 @@ class CAiart(Cog): 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 await ctx.respond(embed=Aiart_Embed.evalate(sorted_tags, rating), ephemeral=False, file=File(fp=io.BytesIO(buffer), filename="image.png")) + + @ART.command(name="포인트", description="아트포인트를 확인합니다.") + @cooldown(1, 10, BucketType.user) + @guild_only() + async def _포인트(self, ctx: Context): + await ctx.defer(ephemeral=False) + if ctx.guild is not None and not await database.get_guild(ctx.guild.id): return await ctx.respond(embed=ChristmasEmbed(title="❌ 에러!", description="서버가 가입되어있지 않아요! 서버를 가입해주세요!", color=Color.red()),ephemeral=True) + point = await database.get_guild_art_point(ctx.guild.id) + use = await database.get_use_art(ctx.guild.id) + await ctx.respond(embed=ChristmasEmbed(title="🎨 아트포인트", description=f"현재 서버의 아트포인트는 {point}개에요!\n 이서버에서는 그림 기능을 {use}번 사용했어요! ", color=Color.green()), ephemeral=False) + def setup(bot): bot.add_cog(CAiart(bot)) \ No newline at end of file diff --git a/Christmas/Database/__init__.py b/Christmas/Database/__init__.py index c76de2a..4c00c21 100644 --- a/Christmas/Database/__init__.py +++ b/Christmas/Database/__init__.py @@ -91,5 +91,44 @@ class database: async def insert_arcade_point(guild_id: int, user_id: int): conn = await MongoDBClient().connect() - await conn.arcadepre.update_one({"_id": guild_id}, {"$inc": {"vote": 1}}) - await conn.arcadepreuser.update_one({"_id": user_id}, {"$inc": {"vote": 1}}) \ No newline at end of file + if await conn.arcadepre.find_one({"_id": guild_id}) == None: + await conn.arcadepre.insert_one({"_id": guild_id, "user": [{"_id": user_id, "point": 1}], "vote": 0}) + else: + await conn.arcadepre.update_one({"_id": guild_id}, {"$inc": {"user.$[elem].point": 1}}, array_filters=[{"elem._id": user_id}]) + + async def insert_guild_art_point(guild_id: int, amount: int): + conn = await MongoDBClient().connect() + if await conn.guild.find_one({"_id": guild_id}) == None: + await conn.guild.insert_one({"_id": guild_id, "art_point": 1}) + else: + await conn.guild.update_one({"_id": guild_id}, {"$inc": {"art_point": amount}}) + + async def use_guild_art_point(guild_id: int): + conn = await MongoDBClient().connect() + if await conn.guild.find_one({"_id": guild_id, "art_point": {"$gt": 0}}) == None: + return False + await conn.guild.update_one({"_id": guild_id}, {"$inc": {"art_point": -1}}) + return True + + async def get_guild_art_point(guild_id: int): + conn = await MongoDBClient().connect() + if await conn.guild.find_one({"_id": guild_id, "art_point": {"$exists": True}}) == None: + return 0 + else: + result = await conn.guild.find_one({"_id": guild_id}) + return result["art_point"] + + async def get_use_art(guild_id: int): + conn = await MongoDBClient().connect() + if await conn.guild.find_one({"_id": guild_id, "art": {"$exists": True}}) == None: + return 0 + else: + result = await conn.guild.find_one({"_id": guild_id}) + return result["art"] + + async def put_use_art(guild_id: int): + conn = await MongoDBClient().connect() + if await conn.guild.find_one({"_id": guild_id, "art": {"$exists": True}}) == None: + await conn.guild.update_one({"_id": guild_id}, {"$set": {"art": 1}}) + else: + await conn.guild.update_one({"_id": guild_id}, {"$inc": {"art": 1}}) \ No newline at end of file diff --git a/Christmas/Module.py b/Christmas/Module.py index cc94462..27bba4a 100644 --- a/Christmas/Module.py +++ b/Christmas/Module.py @@ -6,12 +6,46 @@ import re import random import aiohttp import asyncio +import sentry_sdk from aiogoogletrans import Translator from typing import Any, Dict from discord import File, Asset from Christmas.Tagging import Tagging from Christmas.Cogs.Event import model translator = Translator() +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", + "bikini", + "swimsuit", + "give birth", + "slave" +] + weight = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.7, 0.8, 0.9] cfg = [6, 7, 8, 9, 10, 11] #def check_curse(text: str): @@ -21,7 +55,7 @@ def is_korean(string): match = pattern.search(string) return bool(match) -async def process_prompt(prompt: str, remove: str, res: list, isnsfw: bool, style1: float, style2: float, afterprocess: float, avatar): +async def process_prompt(prompt: str, remove: str, res: list, isnsfw: bool, style1: float, style2: float, afterprocess: bool, avatar): tags = None if not avatar == None: # gif인 경우 png로 변환 @@ -44,6 +78,10 @@ async def process_prompt(prompt: str, remove: str, res: list, isnsfw: bool, styl if is_korean(remove): remove = await translator.translate(remove, dest="en") remove = remove.text + if isnsfw == False: + for i in BLOCKTAG: + if i in prompt: + prompt = prompt.replace(i, "") default_negative = """(KHFB, AuroraNegative, easynegative, negative_hand-neg, verybadimagenegative_v1.3:0.8), (Worst Quality, Low Quality:1.4), border, skimpy, grayscale, multiple_girls, 3d, realistic, string, multiple hands, chinese, thick abs, chubby abs, lowres, bad anatomy, asymmetric wings, elf, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, large areolae, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, Multiple legs, multiple feet, tie, (necktie:1.5), several hands, three feet, four legs, three legs,animal ears""" if style1 != 0: prompt = prompt + f"" @@ -59,9 +97,14 @@ async def process_prompt(prompt: str, remove: str, res: list, isnsfw: bool, styl prompt = prompt + f"" if tags != None: prompt = prompt + "," + tags - payloads = { + payloads = {} + if afterprocess == True: + payloads = { "prompt": prompt, "negative_prompt": negative_prompt, + "enable_hr": True, + "hr_scale": 1.5, + "hr_upscaler": "R-ESRGAN 4x+ Anime6B", "seed": random.randint(0, 1000000000), "steps": 25, "cfg_scale": random.choice(cfg), @@ -69,7 +112,7 @@ async def process_prompt(prompt: str, remove: str, res: list, isnsfw: bool, styl "height": res[1], "sampler_index": "DPM++ 2M Karras", "refiner_checkpoint": "smooREFINERV2R10_half", - "refiner_switch_at": afterprocess, + "refiner_switch_at": 0.45, "alwayson_scripts": { "ADetailer": { 'args': [ @@ -78,6 +121,26 @@ async def process_prompt(prompt: str, remove: str, res: list, isnsfw: bool, styl 'ad_inpaint_only_masked': True }] }}} + else: + payloads = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "seed": random.randint(0, 1000000000), + "steps": 25, + "cfg_scale": random.choice(cfg), + "width": res[0], + "height": res[1], + "sampler_index": "DPM++ 2M Karras", + "refiner_checkpoint": "smooREFINERV2R10_half", + "refiner_switch_at": 0.45, + "alwayson_scripts": { + "ADetailer": { + 'args': [ + { + 'ad_model': 'face_yolov8n.pt', + 'ad_inpaint_only_masked': True + }] + }}} return payloads @@ -120,7 +183,7 @@ async def get_gpuserver_status(url) -> Dict: oomcount = result["cuda"]["events"]["oom"] return {"status": "online", "system_memory_usage": bytes_to_gb(memstatus), "cuda_memory_usage": bytes_to_gb(cudamemstatus), "oom_count": oomcount} except Exception as e: - print(e) + sentry_sdk.capture_exception(e) return {"status": "offline"} def bytes_to_gb(bytes: int) -> float: diff --git a/Christmas/UI/Buttons.py b/Christmas/UI/Buttons.py index 6d4ab13..de88dec 100644 --- a/Christmas/UI/Buttons.py +++ b/Christmas/UI/Buttons.py @@ -67,7 +67,8 @@ class Recommanded(View): async def finished(self, button, interaction: Interaction): await interaction.response.defer() data = await KoreanbotsRequester(self.config.KOREANBOT_TOKEN).get_bot_vote(int(interaction.user.id), interaction.guild.me.id) - if data.voted == True: + if data["data"]["voted"] == True: + await database.insert_guild_art_point(interaction.guild.id, 5) await database.insert_arcade_point(interaction.guild.id, interaction.user.id) await interaction.edit_original_response(embed=Arcade_Embed.recommend_sucess(), view=None) else: diff --git a/Christmas/UI/Embed.py b/Christmas/UI/Embed.py index f5235ab..dada436 100644 --- a/Christmas/UI/Embed.py +++ b/Christmas/UI/Embed.py @@ -67,7 +67,7 @@ class Default_Embed: @staticmethod def cooldown(sec): - return "이 명령어는 " + str(sec) + "초 뒤에 다시 사용할 수 있어요!" + return "이 명령어는 " + f"``{round(sec, 2)}``" + "초 뒤에 다시 사용할 수 있어요!" @staticmethod def help(): @@ -98,7 +98,7 @@ class Default_Embed: gpuserver: [system_memory_usage,cuda_memory_usage,oom_count] """ embed = Embed(title="**봇 정보**", description="크돌이의 정보에요!") - embed.add_field(name="**봇 개요**", value=f"봇 ID: {bot.user.id}\n봇 버전: 1.2.0\n가동시간: {str(uptime)}\n개발자: RunaLab\n**tmddn3070**\nFlag_Fan\nPma\nGunu\ni3l2\nNon", inline=True) + embed.add_field(name="**봇 개요**", value=f"봇 ID: {bot.user.id}\n봇 버전: 1.3.0\n가동시간: {str(uptime)}\n개발자: RunaLab,PvPConnect", inline=True) orin = psutil.virtual_memory().used orin = orin / 1024 / 1024 / 1024 if gpuserver == None or gpuserver["status"] == "offline": @@ -234,6 +234,13 @@ class Aiart_Embed: embed.set_footer() return embed + @staticmethod + def gen_nsfw(): + embed = ChristmasEmbed(title="❌ 그림 생성 실패!", description="그림 생성에 실패했어요!") + embed.add_field(name="안내", value="NSFW 채널이 아닌경우 NSFW 그림을 생성을 할수 없어요!", inline=False) + embed.set_footer() + return embed + class Music_Embed: @staticmethod diff --git a/Christmas/UI/Modal.py b/Christmas/UI/Modal.py index 44cc532..7dcbaa6 100644 --- a/Christmas/UI/Modal.py +++ b/Christmas/UI/Modal.py @@ -1,4 +1,5 @@ import re +import base64 from discord import InputTextStyle, Interaction, Member from discord.ui import Modal, InputText @@ -8,35 +9,9 @@ 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 -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): @@ -54,8 +29,9 @@ class Send_Mail_Modal(Modal): 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): + 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 @@ -65,8 +41,8 @@ class Aiart(Modal): 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")) + 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 @@ -78,10 +54,6 @@ class Aiart(Modal): 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): @@ -98,11 +70,17 @@ class Aiart(Modal): 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: + 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()) - await interaction.edit_original_response(embed=Aiart_Embed.sucess_generate(original_prompt,original_remove), file=attachment) + 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) \ No newline at end of file