feat ✨: Add anime image generation feature and related models.

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)
 

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

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)