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)