main(backup).py

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)