From 27f0bf08be1553a07cd809f13b5058ae962a4284 Mon Sep 17 00:00:00 2001 From: Geoffery Powell Date: Sat, 7 Sep 2024 23:06:22 -0500 Subject: [PATCH] =?UTF-8?q?feat=20=E2=9C=A8:=20Add=20anime=20image=20gener?= =?UTF-8?q?ation=20feature=20and=20related=20models.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 12 +++++----- models/ImageRequest.py | 37 ++++++++++++++++++++++++++----- tests/models/test_ImageRequest.py | 6 ++--- 3 files changed, 40 insertions(+), 15 deletions(-) diff --git a/main.py b/main.py index 5874e5f..902489f 100644 --- a/main.py +++ b/main.py @@ -49,6 +49,7 @@ async def on_message(self, message): print( f'{message.author.name} [{message.author.id}] sent: {message.content} on Channel: {message.channel.id}') + async def updateStatus(): global streamers with open('status.json') as fs: @@ -81,17 +82,17 @@ async def activityType(data): # This should have most of the payload options # ================================== # @tree.command(description="Dream of an Image") -async def dream(interaction: discord.Interaction, prompt: str, negative: str = "", steps: int = 20, - seed: int = -1, cfg_scale: int = 7, width: int = 816, height: int = 1024): +async def dream(interaction: discord.Interaction, prompt: str, negative: str = "", steps: int = 20, + seed: int = -1, cfg_scale: int = 7, width: int = 816, height: int = 1024, is_anime: bool = False): # Dream print(await sendLog(log=f'{interaction.user.name} dreaming of {prompt}', client=client)) - img_request = ImageRequest() + img_request = ImageRequest(is_anime=is_anime) # Acknowledge the interaction try: await interaction.response.defer() except Exception as e: - print(await sendLog(log=e)) + print(await sendLog(log=e)) img_request.set_prompt(prompt) img_request.set_negative_prompt(negative) @@ -101,11 +102,10 @@ async def dream(interaction: discord.Interaction, prompt: str, negative: str = " img_request.set_width(width) img_request.set_height(height) img_request.set_request_type(RequestTypes.TXT2IMG) - + if not interaction.channel.is_nsfw(): img_request.set_not_nsfw() - # Add the command to the queue await sd_request(interaction, img_request, defer=True) diff --git a/models/ImageRequest.py b/models/ImageRequest.py index f57d3f0..5765645 100644 --- a/models/ImageRequest.py +++ b/models/ImageRequest.py @@ -6,16 +6,20 @@ class ImageRequest: def __init__(self, prompt: str = "warning sign", - negative_prompt: str = "fewer digits, extra digits, score_6, score_5, score_4", + negative_prompt: str = "", width: int = 816, height: int = 1024, seed: int = -1, steps: int = 20, - cfg_scale: float = 7): + cfg_scale: float = 7, + is_anime: bool = False): self.good_qualities = "score_9, score_8_up, score_7_up" - self.bad_qualities = "fewer digits, extra digits, score_6, score_5, score_4" + self.bad_qualities = "fewer digits, extra digits, score_6, score_5, score_4, watermark" self.sfw_prompt = "rating_safe" self.sfw_negative = "rating_explicit" + self.is_anime = is_anime + self.anime_prompt = "source_anime" + self.anime_negative = "source_pony, source_furry, source_cartoon" self.set_prompt(prompt) self.set_negative_prompt(negative_prompt) @@ -44,15 +48,18 @@ def __init__(self, prompt: str = "warning sign", self.generation_time = None def set_prompt(self, prompt: str): - prompt = self.easy_positive(prompt) if prompt == "": self.prompt = "warning sign" else: self.prompt = prompt + self.prompt = self.easy_positive(self.prompt) + if self.is_anime: + self.set_anime_prompt() def set_negative_prompt(self, negative_prompt: str): - negative_prompt = self.easy_negative(negative_prompt) - self.negative_prompt = negative_prompt + self.negative_prompt = self.easy_negative(negative_prompt) + if self.is_anime: + self.set_anime_negative() def set_width(self, width: int): self.width = self.dimension_clamp(width) @@ -90,6 +97,9 @@ def set_not_nsfw(self): def set_request_type(self, type: RequestTypes): self.request_type = type + def set_is_anime(self, is_anime: bool): + self.is_anime = is_anime + def set_generation_time(self, time: time): self.generation_time = time @@ -134,11 +144,23 @@ def easy_negative(self, negative_prompt): return f"{self.bad_qualities}, " + negative_prompt return negative_prompt + def set_anime_prompt(self): + if not self.anime_prompt in self.prompt: + self.prompt = f"{self.anime_prompt}, " + self.prompt + + def set_anime_negative(self): + if not self.anime_negative in self.negative_prompt: + self.negative_prompt = f"{self.anime_negative}, " + \ + self.negative_prompt + def get_prompt_without_qualities(self): cleaned_prompt = self.prompt if self.good_qualities in self.prompt: cleaned_prompt = cleaned_prompt.replace( f"{self.good_qualities}, ", "") + if self.anime_prompt in self.prompt: + cleaned_prompt = cleaned_prompt.replace( + f"{self.anime_prompt}, ", "") return cleaned_prompt @@ -147,5 +169,8 @@ def get_negative_without_qualities(self): if self.bad_qualities in self.negative_prompt: cleaned_prompt = cleaned_prompt.replace( f"{self.bad_qualities}, ", "") + if self.anime_negative in self.negative_prompt: + cleaned_prompt = cleaned_prompt.replace( + f"{self.anime_negative}, ", "") return cleaned_prompt diff --git a/tests/models/test_ImageRequest.py b/tests/models/test_ImageRequest.py index 2665f77..2baac54 100644 --- a/tests/models/test_ImageRequest.py +++ b/tests/models/test_ImageRequest.py @@ -15,7 +15,7 @@ def test_init(self): self.assertEqual( image_request.prompt, "score_9, score_8_up, score_7_up, warning sign") self.assertEqual(image_request.negative_prompt, - "fewer digits, extra digits, score_6, score_5, score_4") + "fewer digits, extra digits, score_6, score_5, score_4, watermark, ") self.assertGreater(image_request.seed, 0) self.assertEqual(image_request.enable_hr, False) self.assertEqual(image_request.hr_scale, 2) @@ -92,14 +92,14 @@ def test_easy_positive_adds_string_if_missing(self): self.assertEqual(result, expected_result) def test_easy_negative_does_NOT_add_if_already_there(self): - original_negative_prompt = "original negative prompt with fewer digits, extra digits, score_6, score_5, score_4, 1boy" + original_negative_prompt = f"original negative prompt with {self.image_request.bad_qualities} 1boy" expected_result = original_negative_prompt result = self.image_request.easy_negative(original_negative_prompt) self.assertEqual(result, expected_result) def test_easy_negative_adds_string_if_missing(self): original_negative_prompt = "original negative prompt without negative" - expected_result = "fewer digits, extra digits, score_6, score_5, score_4, original negative prompt without negative" + expected_result = f"{self.image_request.bad_qualities}, original negative prompt without negative" result = self.image_request.easy_negative(original_negative_prompt) self.assertEqual(result, expected_result)