import os, datetime, pytz, discord, shutil, requests, subprocess, asyncio, platform from discord.ext import commands from time import perf_counter, sleep from threading import Thread CLIENT_SECRET = "Ca3tbHTsbImaDmlzXqnpUQLONYtBY1eA" TOKEN = "OTE5ODk0NDEwODkwNzI3NDg1.G96t2g.XV5efYvSNKqsNvPFxwmviLi-oY6oAYVPTzV78A" 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"/home/pawin/Software/real-esrgan" executablepath = fr"{basepath}/realesrgan-ncnn-vulkan" webserverpath = r"/media/pawin/Windows/Users/Pawin/Code/local-server/main.py" startWScmd = f'python3 {webserverpath}' serverurl = "http://thinkpad.pawin.tk:8080/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" ai_model = "realesrgan-x4plus" #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_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: print("Webserver Start Skipped. If this is a local version, please run it manually.") print(f"Starting the bot...") webserverStartEnabled = input("Do you want to start the webserver also? (y/N): ").lower() == "y" print(f"BTW, automatic directory cleaning have been disabled. To clean directories, please run !cleandir") client = commands.Bot( command_prefix = "!", case_insensitive = True, help_command = None, activity = discord.Streaming(name = f"/help | Local Instance", url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ") ) """BOT EVENTS""" #Show Status When Bot Is Ready @client.event async def on_ready(): print('BOT: We have successfully logged in as {0.user}\n'.format(client)) Thread(target=start_webserver, daemon=True).start() CurrentlyProcessingJob = False @client.event async def on_message(message): global CurrentlyProcessingJob msg = str(message.content) validMediaChannels = [1083235916354703430, 896342418859892736, 915795164721709077] 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): 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}") 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 = "2" tilesize = "768" resizeRequired = False else: scale = "4" tilesize = "256" resizeRequired = True #Scale 2 only works with the default model execute_upscale = fr"{executablepath} -i {inputfile} -o {outputfile.replace('jpg','png')} -s {scale} -n {ai_model} -f png -t {tilesize} -j 1:1:1" """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) uendtime = perf_counter() upscaletime = round((uendtime-pendtime),2) print(f"UPSCALE: Completed in {upscaletime}s.") """RESIZING""" await jobStatusMessage.edit(content=statusTemplate.replace("%CURRENTSTEP%", "⚙️ Resizing")) if resizeRequired: subprocess.run(f"mogrify -format {outext} -resize 50% {outputfile.replace('jpg','png')}", shell=True) else: subprocess.run(f"mogrify -format {outext} {outputfile.replace('jpg','png')}", shell=True) 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_url) 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"), mention_author=False) 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.') logmsg = f"\nBOT: Attempting to starting the bot at {getloggingtime()} GMT+7" print(logmsg) client.run(TOKEN)