upscaler-bot / main.py
PawinC's picture
Update main.py
381dd73 verified
raw
history blame
12.4 kB
import os, datetime, pytz, discord, shutil, requests, subprocess, asyncio, platform, json
from discord.ext import commands
from time import perf_counter, sleep
from threading import Thread
from dotenv import load_dotenv
load_dotenv()
import torch
from PIL import Image
import numpy as np
from RealESRGAN import RealESRGAN
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = RealESRGAN(device, scale=4)
model.load_weights('weights/RealESRGAN_x4.pth', download=True)
async def do_upscale(infile, outfile):
image = Image.open(infile).convert('RGB')
sr_image = model.predict(image)
sr_image.save(outfile)
CLIENT_SECRET = os.environ['ClientSecret']
TOKEN = os.environ['TestingToken']
ADMINS = [766145655038410763, 937495666471628831]
ostype = platform.system()
print(f"\nDetected OS: {ostype}")
if ostype == "Windows":
basepath = r"C:\Users\Pawin\Software\real-esrgan"
executablepath = fr"{basepath}\realesrgan-ncnn-vulkan.exe"
webserverpath = r"C:\Users\Pawin\Code\local-server\main.py"
startWScmd = f"python {webserverpath}"
serverurl = "http://thinkpad.pawin.tk/upscale"
elif ostype == "Linux":
basepath = r"./real-esrgan"
executablepath = fr"{basepath}/realesrgan-ncnn-vulkan"
webserverpath = r"/Code/Flask/main.py"
startWScmd = f'python3 {webserverpath}'
serverurl = "https://dev.pawin.tk/upscale"
imagepath = os.path.join(basepath, "discordbot")
inputpath = os.path.join(imagepath, "input")
outputpath = os.path.join(imagepath, "output")
outext = "jpg" #or png
ai_model = "realesr-animevideov3"
targetScale = "2" #2 or 4 (string)
#ai_model = "realesrgan-x4plus-anime"
#Can be realesr-animevideov3(default) | realesrgan-x4plus | realesrgan-x4plus-anime | realesrnet-x4plus
def getloggingtime():
return datetime.datetime.now(pytz.timezone('Asia/Bangkok')).strftime('%d/%m/%Y %H:%M:%S')
def generate_id(path):
return str(len(os.listdir(path))+1).zfill(3)
def file_cleanup():
dirlist = [inputpath, outputpath]
contentlength = len(os.listdir(inputpath))
for directory in dirlist:
shutil.rmtree(directory)
os.mkdir(directory)
print("BOT: Sucessfully cleaned directories")
return contentlength
def start_dummy_webserver():
subprocess.run("mkdir -p hi && echo Hi! > hi/index.html && python -m http.server 7860 -d hi", shell=True)
def start_webserver():
if webserverStartEnabled:
try:
print("Calling Webserver...")
subprocess.run(startWScmd, shell=True)
except:
print("Something went wrong with the webserver.")
print("Webserver process launched.")
else:
start_dummy_webserver()
print("Webserver Start Skipped. Running a dummy one just to make huggingface display it as running.")
def loadMediaChannels():
global mediaChannels
with open("resources/mediaChannels.json", "r") as f:
mediaChannels = json.load(f)
#print("Successfully loaded media channels.")
def backupMediaChannels():
with open("resources/mediaChannels.json", "w") as f:
json.dump(mediaChannels, f, indent=2, ensure_ascii=False)
#print("Successfully backed up media channel.")
def toggleMediaChannel(channelId: int):
global mediaChannels
loadMediaChannels()
if channelId in mediaChannels:
mediaChannels.remove(channelId)
actionDone = "Unmarked"
else:
mediaChannels.append(int(channelId))
actionDone = "Marked"
backupMediaChannels()
resultMsg = f"Successfully {actionDone} this channel as a media channel."
print(resultMsg)
return resultMsg
print(f"Starting the upscaler bot...")
#webserverStartEnabled = input("Do you want to start the webserver also? (y/N): ").lower() == "y"
webserverStartEnabled = False
#print(f"BTW, automatic directory cleaning have been disabled. To clean directories, please run u!cleandir")
loadMediaChannels()
client = commands.Bot(
intents = discord.Intents.all(),
command_prefix = "u!",
case_insensitive = True,
help_command = None,
activity = discord.Streaming(name = f"u!help | Currently Running Upscaler Bot ", url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ")
)
"""BOT EVENTS"""
#Show Status When Bot Is Ready
@client.event
async def on_ready():
print('\nBOT: We have successfully logged in as {0.user}'.format(client))
file_cleanup()
Thread(target=start_webserver, daemon=True).start()
CurrentlyProcessingJob = False
@client.event
async def on_message(message):
global CurrentlyProcessingJob
msg = str(message.content)
validMediaChannels = mediaChannels
channel = message.channel
author = message.author
channelId = int(channel.id)
authorId = int(author.id)
#print(f"Channel id is {channelId}, author id is {authorId}")
imageList = []
for item in message.attachments:
if item.content_type.startswith("image"):
imageList.append(item.url)
if message.author.bot or msg.startswith("p!"):
pass
#elif (authorId in ADMINS) and (channelId in validMediaChannels) and (message.attachments) and (imageList):
elif (authorId) and (channelId in validMediaChannels) and (message.attachments) and (imageList):
print(f'IMAGE: New attachment recieved from {author} in channel {channel}. Processing...')
await message.add_reaction("πŸ“¬")
jobStatusMessage = await message.reply(f"Status: Upscale queued.", mention_author=False)
if CurrentlyProcessingJob:
pendingTemplate = f"__Status__: Upscale %CURRENTSTEP% for job *{message.id}*"
await jobStatusMessage.edit(content=pendingTemplate.replace("%CURRENTSTEP%", "Pending πŸ•°οΈ"))
print(f"IMAGE: There is already a job running. Message id {message.id} has been put on queue")
while CurrentlyProcessingJob:
await asyncio.sleep(1)
#wait untill it finishes
await jobStatusMessage.edit(content=pendingTemplate.replace("%CURRENTSTEP%", "Starting πŸƒ"))
print(f"IMAGE: Starting job {message.id}")
CurrentlyProcessingJob = True
attachmentCount = len(imageList)
multipleImages = attachmentCount > 1
if multipleImages:
batchstart = perf_counter()
taskType = "Batch Upscale"
else:
taskType = "Image Upscaling"
for i in range(attachmentCount):
"""PREPROCESSING"""
fileid = f"{message.id}_{i}"
statusTemplate = f"__Status__: %CURRENTSTEP% image from message *{message.id}* **({i+1}/{attachmentCount})**."
await jobStatusMessage.edit(content=statusTemplate.replace("%CURRENTSTEP%", "πŸ”„ Preprocessing"))
starttime = perf_counter()
url = imageList[i]
extension = url.split(".")[-1]
#fileid = generate_id(f"{imagepath}\input")
inputfile = os.path.join(inputpath, f"{fileid}.{extension}").split("?")[0]
outputpng = os.path.join(outputpath, f"{fileid}.png")
outputfile = os.path.join(outputpath, f"{fileid}.{outext}")
with open (inputfile, "wb") as f:
f.write(requests.get(url).content)
"""UPSCALE AUTO-CONFIG"""
if ai_model == "realesr-animevideov3":
scale = targetScale
tilesize = "768"
resizeRequired = False
else:
scale = "4"
tilesize = "256"
resizeRequired = True if targetScale == "2" else False
#Scale 2 only works with the default model
execute_upscale = fr"{executablepath} -i {inputfile} -o {outputpng} -n {ai_model} -s {scale} -f png -t {tilesize}"
print(execute_upscale)
"""UPSCALING"""
pendtime = perf_counter()
preprocessingtime = round((pendtime-starttime), 2)
print(f"PREPROCESS: Completed in {preprocessingtime}s.\nUPSCALE:")
await jobStatusMessage.edit(content=(statusTemplate.replace("%CURRENTSTEP%", "🧠 Upscaling")+"\nThis may take a while..."))
#os.system(execute_upscale)
#subprocess.run(execute_upscale, shell=True)
# upscaleProcess = await asyncio.create_subprocess_shell(execute_upscale)
# await upscaleProcess.wait()
await do_upscale(inputfile, outputpng)
uendtime = perf_counter()
upscaletime = round((uendtime-pendtime),2)
print(f"UPSCALE: Completed in {upscaletime}s.")
"""RESIZING"""
await jobStatusMessage.edit(content=statusTemplate.replace("%CURRENTSTEP%", "βš™οΈ Resizing"))
#qualityArgs = f" -quality 95 " if outext == "jpg" else ""
qualityArgs = ""
resizeCommand = f"mogrify -format {outext}{qualityArgs} -resize 50% {outputpng}"
convertCommand = f"mogrify -format {outext}{qualityArgs} {outputpng}"
resizeProcess = await asyncio.create_subprocess_shell(resizeCommand if resizeRequired else convertCommand)
await resizeProcess.wait()
rendtime = perf_counter()
resizingtime = round((rendtime-uendtime), 2)
print(f"RESIZE: Completed in {resizingtime}s.")
"""DELIVERING"""
await jobStatusMessage.edit(content=statusTemplate.replace("%CURRENTSTEP%", "βœ‰οΈ Sending"))
try:
file=discord.File(outputfile)
imgurl = f"attachment://{fileid}.{outext}"
prepembed=discord.Embed(title="Sucessfully Upscaled Image")
prepembed.set_image(url=imgurl)
jobResultMessage = await message.channel.send(embed=prepembed, file=file)
sendtime = perf_counter()
localImageUrl = f"{serverurl}/{fileid}.{outext}"
outputImageUrl = jobResultMessage.embeds[0].image.url
#print(outputImageUrl)
#outputImageUrl = jobResultMessage.attachments[0].url
#print(outputImageUrl)
embed = discord.Embed(title="Upscaled Image", url=outputImageUrl, description=f"[Local File]({localImageUrl})")
embed.set_author(name=author, icon_url=message.author.avatar)
embed.set_image(url=imgurl)
sendingtime = round((sendtime-rendtime), 2)
processingstats = f"Preprocessing: {preprocessingtime}s | Resizing: {resizingtime}s | Sending: {sendingtime}s."
embed.set_footer(text=f"Took {upscaletime}s to upscale {fileid}.\nUpscaling done with {ai_model}.\n{processingstats}")
await jobResultMessage.edit(embed=embed)
except discord.errors.HTTPException as e:
baseErrorMessage = f"There was an error sending the output for image id {fileid}."
if '413 Payload Too Large' in str(e):
await message.reply(f"{baseErrorMessage} It was too large for discord to handle.\n```python\n{e}```")
else:
await message.reply(f"{baseErrorMessage}\n```python\n{e}```")
except Exception as e:
await message.reply(f"Encountered an error while processing attachment {fileid}\n```python\n{e}```")
"""CLEANING UP"""
#Already finished the whole job.
await jobStatusMessage.edit(content=f"\n__Status__: βœ… Job Completed for all {attachmentCount} attachments from message {message.id} at {getloggingtime()}".replace("all 1 attachments", "an attachment"))
#await message.delete()
await asyncio.sleep(2)
await jobStatusMessage.delete()
print(f"IMAGE: Job finished for messageID {message.id}\n")
CurrentlyProcessingJob = False
else:
print(f"MESSAGE: {channel}- {author}: {msg}")
await client.process_commands(message)
"""BOT COMMANDS"""
@client.command()
async def test(ctx):
await ctx.send("Hello!")
@client.command()
async def ping(ctx):
await ctx.send(f'The ping is {round(client.latency * 1000)}ms.')
@client.command()
async def cleandir(ctx):
if ctx.author.id in ADMINS:
await ctx.channel.purge(limit=file_cleanup()+1)
await ctx.send(f'Successfully images directories', delete_after= 2)
else:
await ctx.send('You do not have the permissions to perform this action.')
@client.command()
async def clear(ctx, amount=None):
if ctx.author.id in ADMINS:
if amount == "all":
amount = 100
else:
try:
amount = int(amount)
except:
amount = 2
await ctx.channel.purge(limit=amount)
await ctx.send(f"Finished clearing {amount} messages {ctx.message.author.mention}", delete_after=4)
else:
await ctx.send('You do not have the permissions to clear messages using this bot.')
@client.command(aliases = ['enable-upscale'])
async def toggleUpscale(ctx):
result = toggleMediaChannel(int(ctx.channel.id))
#print(result)
await ctx.send(result)
logmsg = f"\nBOT: Attempting to starting the bot at {getloggingtime()} GMT+7"
print(logmsg)
client.run(TOKEN)