Skip to content

Commit

Permalink
🐝 Change Imagine command name to pollinate and fix regeneration bugs
Browse files Browse the repository at this point in the history
  • Loading branch information
@Zingzy
Zingzy committed Mar 26, 2024
1 parent c3776ca commit 1c13e3c
Showing 1 changed file with 162 additions and 158 deletions.
320 changes: 162 additions & 158 deletions cogs/imagine_cog.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,207 +6,211 @@
from utils import *
from constants import *

class ImagineButtonView(discord.ui.View):
def __init__(self, link: str = None):
super().__init__(timeout=None)
self.link = link

class Imagine(commands.Cog):
def __init__(self, bot):
self.bot = bot
if link is not None:
self.add_item(discord.ui.Button(label="Link", url=self.link))

async def cog_load(self):
await self.bot.wait_until_ready()
self.bot.add_view(self.ImagineButtonView())
@discord.ui.button(style=discord.ButtonStyle.secondary, custom_id="regenerate-button", emoji="<:redo:1187101382101180456>")
async def regenerate(self, interaction: discord.Interaction, button: discord.ui.Button):
message_id = interaction.message.id
await interaction.response.send_message(embed=discord.Embed(title="Regenerating Your Image", description="Please wait while we generate your image", color=discord.Color.blurple()), ephemeral=True)

class ImagineButtonView(discord.ui.View):
def __init__(self, link: str = None):
super().__init__(timeout=None)
self.link = link
message_data = get_prompt_data(message_id)

if link is not None:
self.add_item(discord.ui.Button(label="Link", url=self.link))
if not message_data:
await interaction.followup.send(embed=discord.Embed(title="Error", description="Message not found", color=discord.Color.red()), ephemeral=True)
return

@discord.ui.button(style=discord.ButtonStyle.secondary, custom_id="regenerate-button", emoji="<:redo:1187101382101180456>")
async def regenerate(self, interaction: discord.Interaction, button: discord.ui.Button):
message_id = interaction.message.id
await interaction.response.send_message(embed=discord.Embed(title="Regenerating Your Image", description="Please wait while we generate your image", color=discord.Color.blurple()), ephemeral=True)
start = datetime.datetime.now()

message_data = get_prompt_data(message_id)
prompt = message_data["prompt"]
width = message_data["width"]
height = message_data["height"]
model = message_data["model"]
negative = message_data["negative"]
cached = message_data["cached"]
nologo = message_data["nologo"]
enhance = message_data["enhance"]

if not message_data:
await interaction.followup.send(embed=discord.Embed(title="Error", description="Message not found", color=discord.Color.red()), ephemeral=True)
return
try:
dic, image = await generate_image(prompt, width, height, model, negative, cached, nologo, enhance)
except Exception as e:
print(e)
await interaction.followup.send(embed=discord.Embed(title="Error", description=f"Error generating image : {e}", color=discord.Color.red()), ephemeral=True)
return

prompt = message_data["prompt"]
width = message_data["width"]
height = message_data["height"]
model = message_data["model"]
negative = message_data["negative"]
cached = message_data["cached"]
nologo = message_data["nologo"]
enhance = message_data["enhance"]
image_file = discord.File(image, filename="image.png")

try:
dic, image = await generate_image(prompt, width, height, model, negative, cached, nologo, enhance)
except Exception as e:
print(e)
await interaction.followup.send(embed=discord.Embed(title="Error", description=f"Error generating image : {e}", color=discord.Color.red()), ephemeral=True)
return
time_taken = datetime.datetime.now() - start

image_file = discord.File(image, filename="image.png")
context = f"## {prompt} - {interaction.user.mention}\n### Model - `{model}` | Time Taken - `{round(time_taken.total_seconds(), 2)} s`\n### Width - `{width} px` | Height - `{height} px`\n### Enchance - `{enhance}`"

response = await interaction.channel.send(f"## {prompt} - {interaction.user.mention}", file=image_file, view=self)
response = await interaction.channel.send(context, file=image_file, view=ImagineButtonView(link=dic["bookmark_url"]))

dic["_id"] = response.id
dic["channel_id"] = interaction.channel.id
dic["user_id"] = interaction.user.id
dic["guild_id"] = interaction.guild.id
dic["author"] = interaction.user.id
dic["bookmarks"] = []
dic["likes"] = []
dic["_id"] = response.id
dic["channel_id"] = interaction.channel.id
dic["user_id"] = interaction.user.id
dic["guild_id"] = interaction.guild.id
dic["author"] = interaction.user.id
dic["bookmarks"] = []
dic["likes"] = []

user_data = get_user_data(interaction.user.id)
if user_data is None:
user_data = {
"_id": interaction.user.id,
"bookmarks": [],
"likes": [],
"prompts": [],
"last_prompt": None,
}
save_user_data(interaction.user.id, user_data)

user_data["prompts"].append(response.id)

update_user_data(interaction.user.id, user_data)
save_prompt_data(message_id, dic)

@discord.ui.button(label="0", style=discord.ButtonStyle.secondary, custom_id="like-button", emoji="<:like:1187101385230143580>")
async def like(self, interaction: discord.Interaction, button: discord.ui.Button):
try:
id = interaction.message.id
message_data = get_prompt_data(id)
likes = message_data["likes"]

user_data = get_user_data(interaction.user.id)
if user_data is None:
user_data = {
"_id": interaction.user.id,
"bookmarks": [],
"likes": [],
"prompts": [],
"bookmarks": {},
"likes": {},
"prompts": {},
"last_prompt": None,
}
save_user_data(interaction.user.id, user_data)

user_data["prompts"].append(response.id)
if interaction.user.id in likes:
likes.remove(interaction.user.id)
update_prompt_data(id, {"likes": likes})
button.label = f"{len(likes)}"
await interaction.response.edit_message(view=self)

update_user_data(interaction.user.id, user_data)
save_prompt_data(message_id, dic)
user_data["likes"].remove(id)
update_user_data(interaction.user.id, user_data)

@discord.ui.button(label="0", style=discord.ButtonStyle.secondary, custom_id="like-button", emoji="<:like:1187101385230143580>")
async def like(self, interaction: discord.Interaction, button: discord.ui.Button):
try:
id = interaction.message.id
message_data = get_prompt_data(id)
likes = message_data["likes"]
return
else:
likes.append(interaction.user.id)
update_prompt_data(id, {"likes": likes})
button.label = f"{len(likes)}"
await interaction.response.edit_message(view=self)

user_data["likes"].append(id)
update_user_data(interaction.user.id, user_data)

return
except Exception as e:
print(e)
interaction.response.send_message(embed=discord.Embed(title="Error Liking the Image", description=f"{e}", color=discord.Color.red()), ephemeral=True)


@discord.ui.button(label = "0", style=discord.ButtonStyle.secondary, custom_id="bookmark-button", emoji="<:save:1187101389822902344>")
async def bookmark(self, interaction: discord.Interaction, button: discord.ui.Button):
try:
id = interaction.message.id
message_data = get_prompt_data(id)
bookmarks = message_data["bookmarks"]

if interaction.user.id in bookmarks:
await interaction.response.send_message(embed=discord.Embed(title="Error", description="You have already bookmarked this image", color=discord.Color.red()), ephemeral=True)
else:
bookmarks.append(interaction.user.id)
update_prompt_data(id, {"bookmarks": bookmarks})
button.label = f"{len(bookmarks)}"
await interaction.response.edit_message(view=self)

embed = discord.Embed(title=f"Prompt : {message_data['prompt']}", description=f"url : {message_data['bookmark_url']}", color=discord.Color.og_blurple())
embed.set_image(url=message_data["bookmark_url"])

await interaction.user.send(embed=embed)

user_data = get_user_data(interaction.user.id)
if user_data is None:
user_data = {
"_id": interaction.user.id,
"bookmarks": {},
"likes": {},
"prompts": {},
"bookmarks": [],
"likes": [],
"prompts": [],
"last_prompt": None,
}
save_user_data(interaction.user.id, user_data)

if interaction.user.id in likes:
likes.remove(interaction.user.id)
update_prompt_data(id, {"likes": likes})
button.label = f"{len(likes)}"
await interaction.response.edit_message(view=self)
user_data["bookmarks"].append(id)
update_user_data(interaction.user.id, user_data)

user_data["likes"].remove(id)
update_user_data(interaction.user.id, user_data)
return

return
else:
likes.append(interaction.user.id)
update_prompt_data(id, {"likes": likes})
button.label = f"{len(likes)}"
await interaction.response.edit_message(view=self)
except Exception as e:
print(e)
await interaction.response.send_message(embed=discord.Embed(title="Error Bookmarking the Image", description=f"{e}", color=discord.Color.red()), ephemeral=True)

user_data["likes"].append(id)
update_user_data(interaction.user.id, user_data)
@discord.ui.button(style=discord.ButtonStyle.red, custom_id="delete-button", emoji="<:delete:1187102382312652800>")
async def delete(self, interaction: discord.Interaction, button: discord.ui.Button):
try:
data = get_prompt_data(interaction.message.id)
author_id = data["author"]
likes = data["likes"]
bookmarks = data["bookmarks"]
try:
int(author_id)
except:
pass

return
except Exception as e:
print(e)
interaction.response.send_message(embed=discord.Embed(title="Error Liking the Image", description=f"{e}", color=discord.Color.red()), ephemeral=True)
if interaction.user.id != author_id:
await interaction.response.send_message(embed=discord.Embed(title="Error", description="You can only delete your own images", color=discord.Color.red()), ephemeral=True)
return

delete_prompt_data(interaction.message.id)
await interaction.message.delete()

@discord.ui.button(label = "0", style=discord.ButtonStyle.secondary, custom_id="bookmark-button", emoji="<:save:1187101389822902344>")
async def bookmark(self, interaction: discord.Interaction, button: discord.ui.Button):
try:
id = interaction.message.id
message_data = get_prompt_data(id)
bookmarks = message_data["bookmarks"]

if interaction.user.id in bookmarks:
await interaction.response.send_message(embed=discord.Embed(title="Error", description="You have already bookmarked this image", color=discord.Color.red()), ephemeral=True)
else:
bookmarks.append(interaction.user.id)
update_prompt_data(id, {"bookmarks": bookmarks})
button.label = f"{len(bookmarks)}"
await interaction.response.edit_message(view=self)

embed = discord.Embed(title=f"Prompt : {message_data['prompt']}", description=f"url : {message_data['bookmark_url']}", color=discord.Color.og_blurple())
embed.set_image(url=message_data["bookmark_url"])

await interaction.user.send(embed=embed)

user_data = get_user_data(interaction.user.id)
if user_data is None:
user_data = {
"_id": interaction.user.id,
"bookmarks": [],
"likes": [],
"prompts": [],
"last_prompt": None,
}
save_user_data(interaction.user.id, user_data)

user_data["bookmarks"].append(id)
update_user_data(interaction.user.id, user_data)

return

except Exception as e:
print(e)
await interaction.response.send_message(embed=discord.Embed(title="Error Bookmarking the Image", description=f"{e}", color=discord.Color.red()), ephemeral=True)

@discord.ui.button(style=discord.ButtonStyle.red, custom_id="delete-button", emoji="<:delete:1187102382312652800>")
async def delete(self, interaction: discord.Interaction, button: discord.ui.Button):
try:
data = get_prompt_data(interaction.message.id)
author_id = data["author"]
likes = data["likes"]
bookmarks = data["bookmarks"]
user_data = get_user_data(interaction.user.id)
if user_data is None:
return

user_data["prompts"].remove(interaction.message.id)
update_user_data(interaction.user.id, user_data)

for i in likes:
try:
int(author_id)
user_data = get_user_data(i)
user_data["likes"].remove(interaction.message.id)
update_user_data(i, user_data)
except:
pass

if interaction.user.id != author_id:
await interaction.response.send_message(embed=discord.Embed(title="Error", description="You can only delete your own images", color=discord.Color.red()), ephemeral=True)
return

delete_prompt_data(interaction.message.id)
await interaction.message.delete()

user_data = get_user_data(interaction.user.id)
if user_data is None:
return

user_data["prompts"].remove(interaction.message.id)
update_user_data(interaction.user.id, user_data)

for i in likes:
try:
user_data = get_user_data(i)
user_data["likes"].remove(interaction.message.id)
update_user_data(i, user_data)
except:
pass
for i in bookmarks:
try:
user_data = get_user_data(i)
user_data["bookmarks"].remove(interaction.message.id)
update_user_data(i, user_data)
except:
pass

for i in bookmarks:
try:
user_data = get_user_data(i)
user_data["bookmarks"].remove(interaction.message.id)
update_user_data(i, user_data)
except:
pass
except Exception as e:
print(e)
await interaction.response.send_message(embed=discord.Embed(title="Error Deleting the Image", description=f"{e}", color=discord.Color.red()), ephemeral=True)

except Exception as e:
print(e)
await interaction.response.send_message(embed=discord.Embed(title="Error Deleting the Image", description=f"{e}", color=discord.Color.red()), ephemeral=True)
class Imagine(commands.Cog):
def __init__(self, bot):
self.bot = bot

async def cog_load(self):
await self.bot.wait_until_ready()
self.bot.add_view(ImagineButtonView())

async def model_autocomplete(self,
interaction: discord.Interaction,
Expand All @@ -219,7 +223,7 @@ async def model_autocomplete(self,
]


@app_commands.command(name="imagine", description="Imagine a prompt")
@app_commands.command(name="pollinate", description="Generate AI Images")
@app_commands.autocomplete(model=model_autocomplete)
@app_commands.guild_only()
@app_commands.checks.cooldown(1, 15)
Expand Down Expand Up @@ -250,7 +254,7 @@ async def imagine_command(self, interaction, prompt:str, model: str = "Swizz8",
if i.lower() in NSFW_WORDS:
image_file.filename = f"SPOILER_{image_file.filename}"

view = self.ImagineButtonView(link=dic["bookmark_url"])
view = ImagineButtonView(link=dic["bookmark_url"])

time_taken = datetime.datetime.now() - start

Expand Down

0 comments on commit 1c13e3c

Please sign in to comment.