-
Notifications
You must be signed in to change notification settings - Fork 1
/
bot.py
47 lines (36 loc) · 1.33 KB
/
bot.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
import os
from io import BytesIO
import discord
import requests
from PIL import Image
from captioning import Pipeline
client = discord.Client()
pipe = Pipeline()
def prepare_attachment(im, filename, resize=False):
if resize:
im = im.resize((256, 256)) # Super lazy, consider rescaling
temp = BytesIO()
temp.name = filename
im.save(temp, format='PNG')
temp.seek(0)
return discord.File(temp, filename=temp.name)
@client.event
async def on_message(message):
await client.wait_until_ready()
if message.author.id == client.user.id:
return
if message.attachments:
attachment = message.attachments[0]
im = Image.open(requests.get(attachment.url, stream=True).raw)
caption = pipe(im)
try:
await message.channel.send(caption, file=prepare_attachment(im, attachment.filename))
await message.delete()
except discord.errors.HTTPException:
await message.channel.send(caption, file=prepare_attachment(im, attachment.filename, resize=True))
await message.delete()
except Exception as e:
await message.channel.send(f'Error: {e}')
else:
await message.channel.send("No attachments to be found...Can't caption dat! Try sending me an image 😉")
client.run(os.environ.get("DISCORD_TOKEN", None))