-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathapp.py
198 lines (164 loc) · 7.95 KB
/
app.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
import os
import re
import discord
from discord.ext import commands
from ditk import logging
from hbutils.string import plural_word
from hbutils.system import TemporaryDirectory
from maid_assistant.calc import safe_eval
from maid_assistant.explain import tag_explain
from maid_assistant.sites.danbooru import query_danbooru_images, download_danbooru_images
from maid_assistant.sites.gelbooru import query_gelbooru_images, download_gelbooru_images
logging.try_init_root(logging.INFO)
intents = discord.Intents.default()
intents.message_content = True
intents.messages = True
bot = commands.Bot(command_prefix='maid ', intents=intents)
@bot.command(name='calc',
help="Calculate python-based math expression.")
async def calc_command(ctx, *, expression: str):
logging.info(f'Calculate expression {expression!r} ...')
try:
result = safe_eval(expression)
ret_text = f"Result: {result}"
except Exception as e:
ret_text = f'Calculation Error: {e}'
await ctx.message.reply(ret_text)
@bot.command(name='danbooru',
help='Search danbooru images')
async def danbooru_command(ctx, *, tags_text: str):
tags = list(filter(bool, re.split(r'\s+', tags_text)))
if hasattr(ctx.channel, 'is_nsfw'):
is_nsfw = ctx.channel.is_nsfw()
level_name = f'{"NSFW" if ctx.channel.is_nsfw() else "SFW"} '
allowed_ratings = {'g', 's'} if not is_nsfw else {'q', 'e'}
else:
level_name = ''
allowed_ratings = {'g', 's', 'q', 'e'}
reply_message = await ctx.message.reply(
f'Cute maid is searching {level_name}images '
f'with tags {", ".join([f"`{tag}`" for tag in tags])} from danbooru ...')
with TemporaryDirectory() as td:
result = query_danbooru_images(tags, count=10, allowed_ratings=allowed_ratings)
embed = discord.Embed(
title="Danbooru Images",
description=f"This is the search result of tags: {tags!r}.\n"
f"{plural_word(len(result), 'image')} found in total.\n"
f"Powered by [deepghs/danbooru2023-webp-4Mpixel_index](https://huggingface.co/datasets/deepghs/danbooru2023-webp-4Mpixel_index) "
f"and [deepghs/cheesechaser](https://github.com/deepghs/cheesechaser).",
color=0x00ff00
)
files = []
for id_, image in result:
dst_file = os.path.join(td, f'{id_}.webp')
image.save(dst_file, quality=90)
files.append(discord.File(dst_file, filename=os.path.basename(dst_file)))
await reply_message.delete()
await ctx.message.reply(embed=embed, files=files)
@bot.command(name='danbooru_dl',
help='Batch download danbooru images')
async def danbooru_dl_command(ctx, *, tags_text: str):
tags = list(filter(bool, re.split(r'\s+', tags_text)))
reply_message = await ctx.message.reply(
f'Cute maid is downloading and packing images '
f'with tags {", ".join([f"`{tag}`" for tag in tags])} from danbooru ...')
with download_danbooru_images(tags, max_total_size=25 * 1024 ** 2) as (file_count, package_file):
embed = discord.Embed(
title="Danbooru Image Pack",
description=f"This is the image package of tags: {tags!r}.\n"
f"{plural_word(len(file_count), 'image')} inside.\n"
f"Powered by [deepghs/danbooru2023-webp-4Mpixel_index](https://huggingface.co/datasets/deepghs/danbooru2023-webp-4Mpixel_index) "
f"and [deepghs/cheesechaser](https://github.com/deepghs/cheesechaser).",
color=0x00ff00,
)
await reply_message.delete()
await ctx.message.reply(
embed=embed,
files=[
discord.File(package_file, filename=os.path.basename(package_file))
]
)
@bot.command(name='gelbooru',
help='Search gelbooru images')
async def gelbooru_command(ctx, *, tags_text: str):
tags = list(filter(bool, re.split(r'\s+', tags_text)))
if hasattr(ctx.channel, 'is_nsfw'):
is_nsfw = ctx.channel.is_nsfw()
level_name = f'{"NSFW" if ctx.channel.is_nsfw() else "SFW"} '
allowed_ratings = {'general', 'sensitive'} if not is_nsfw else {'questionable', 'explicit'}
else:
level_name = ''
allowed_ratings = {'general', 'sensitive', 'questionable', 'explicit'}
reply_message = await ctx.message.reply(
f'Cute maid is searching {level_name}images '
f'with tags {", ".join([f"`{tag}`" for tag in tags])} from gelbooru ...')
with TemporaryDirectory() as td:
result = query_gelbooru_images(tags, count=10, allowed_ratings=allowed_ratings)
embed = discord.Embed(
title="Gelbooru Images",
description=f"This is the search result of tags: {tags!r}.\n"
f"{plural_word(len(result), 'image')} found in total.\n"
f"Powered by [deepghs/gelbooru-webp-4Mpixel](https://huggingface.co/datasets/deepghs/gelbooru-webp-4Mpixel) "
f"and [deepghs/cheesechaser](https://github.com/deepghs/cheesechaser).",
color=0x0000ff
)
files = []
for id_, image in result:
dst_file = os.path.join(td, f'{id_}.webp')
image.save(dst_file, quality=90)
files.append(discord.File(dst_file, filename=os.path.basename(dst_file)))
await reply_message.delete()
await ctx.message.reply(embed=embed, files=files)
@bot.command(name='gelbooru_dl',
help='Batch download gelbooru images')
async def gelbooru_dl_command(ctx, *, tags_text: str):
tags = list(filter(bool, re.split(r'\s+', tags_text)))
reply_message = await ctx.message.reply(
f'Cute maid is downloading and packing images '
f'with tags {", ".join([f"`{tag}`" for tag in tags])} from gelbooru ...')
with download_gelbooru_images(tags, max_total_size=25 * 1024 ** 2) as (file_count, package_file):
embed = discord.Embed(
title="Danbooru Image Pack",
description=f"This is the image package of tags: {tags!r}.\n"
f"{plural_word(len(file_count), 'image')} inside.\n"
f"Powered by [deepghs/gelbooru-webp-4Mpixel](https://huggingface.co/datasets/deepghs/gelbooru-webp-4Mpixel) "
f"and [deepghs/cheesechaser](https://github.com/deepghs/cheesechaser).",
color=0x0000ff,
)
await reply_message.delete()
await ctx.message.reply(
embed=embed,
files=[
discord.File(package_file, filename=os.path.basename(package_file))
]
)
async def explain_command_raw(ctx, *, tag: str, lang: str):
reply_message = await ctx.message.reply(f'Cute maid is trying to understand '
f'and explain tag `{tag}` in {lang} ...')
try:
reply_text = tag_explain(tag, lang, use_other_names=True)
except Exception as err:
reply_text = f'Explain error - {err!r}'
await reply_message.delete()
await ctx.message.reply(reply_text)
@bot.command(name='explain',
help='Explain tags in english')
async def explain_en_command(ctx, *, tag: str):
await explain_command_raw(ctx, tag=tag, lang='english')
@bot.command(name='explain_cn',
help='Explain tags in chinese')
async def explain_cn_command(ctx, *, tag: str):
await explain_command_raw(ctx, tag=tag, lang='simplified chinese')
@bot.command(name='explain_jp',
help='Explain tags in japanese')
async def explain_jp_command(ctx, *, tag: str):
await explain_command_raw(ctx, tag=tag, lang='japanese')
@bot.command(name='explain_kr',
help='Explain tags in korean')
async def explain_kr_command(ctx, *, tag: str):
await explain_command_raw(ctx, tag=tag, lang='korean')
@bot.event
async def on_ready():
logging.info(f'Bot logged in as {bot.user}')
if __name__ == '__main__':
bot.run(os.environ['DC_BOT_TOKEN'])