main.py

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)