Update main.py

main.py

@@ -5,6 +5,21 @@ 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]
@@ -210,8 +225,12 @@ async def on_message(message):
       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()
+
+        
+      # upscaleProcess = await asyncio.create_subprocess_shell(execute_upscale)
+      # await upscaleProcess.wait()
+
+      await do_upscale(inputfile, outputpng)
 
       uendtime = perf_counter()