From 7f85f98519d5114351cd2c93df76af695fde7a8e Mon Sep 17 00:00:00 2001
From: Bob Bregant <bob@spincats.com>
Date: Sat, 17 Sep 2022 04:44:56 +0000
Subject: [PATCH] More accurate wait time estimates. Still not perfect, but not
 bad.

---
 bot.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/bot.py b/bot.py
index ae8ebdf..8743458 100644
--- a/bot.py
+++ b/bot.py
@@ -27,8 +27,11 @@ class EnlargeButton(discord.ui.Button):
         self.logger.info(f'Queueing request for {interaction.user} to embiggen {self.custom_id}.')
         interaction.client.commissions.put_nowait((t, 'embiggen', interaction.message, self.custom_id, None, None, None, None))
         queuelen = interaction.client.commissions.qsize()
-        waittime = (queuelen * 6) - 4 
-        await interaction.response.send_message(f'{interaction.user} you are number {queuelen} in the Happy Trees processing queue. Estimated wait time is {waittime} minutes.')
+        remain = interaction.client.currtasktime - time.time()
+        if remain < 0:
+            remain = 0
+        waittime = (remain / 60) + (queuelen * 7) - 6.5
+        await interaction.response.send_message(f'{interaction.user} you are number {queuelen} in the Happy Trees processing queue. Estimated wait time is {waittime:0.2f} minutes.')
 
 class EnlargeView(discord.ui.View):
     def __init__(self, outfile=None, outpath=None):
@@ -83,6 +86,7 @@ class HappyTreesBot(commands.Bot):
         self.logger = logging.getLogger('discord.HappyTreesBot')
         self.commissions = asyncio.Queue()
         self.paintings = []
+        self.currtasktime = 0
         self.outpath = os.path.abspath(os.path.expanduser('~/happytreesbot/outputs'))
         self.sdpath = os.path.abspath(os.path.expanduser('~/stable-diffusion'))
         self.ganpath = os.path.abspath(os.path.expanduser('~/real-ESRGAN'))
@@ -229,14 +233,18 @@ class HappyTreesBot(commands.Bot):
         self.logger.info(f'Queueing request for {message.author}. Prompt: "{prompt}"; Samples: {samples}; Seed: {seed}; Steps: {steps}; Strength: {strength}.')
         self.commissions.put_nowait((t, 'paint', message, prompt, samples, seed, steps, strength))
         position=self.commissions.qsize()
-        waittime=position * 6
-        await message.reply(f'{message.author} you are number {position} in the Happy Trees processing queue. Estimated wait time is {waittime} minutes.')
+        remain = self.currtasktime - time.time()
+        if remain < 0:
+            remain = 0
+        waittime = ((remain + 30 + (2 * steps * samples)) / 60) + ((position - 1) * 7)
+        await message.reply(f'{message.author} you are number {position} in the Happy Trees processing queue. Estimated wait time is {waittime:0.2f} minutes.')
 
     async def painting(self):
         while True:
             self.logger.debug(f'Bob Ross is ready to paint!')
             t, command, message, prompt, samples, seed, steps, strength = await self.commissions.get()
             if command == "embiggen":
+                self.currtasktime = time.time() + 30
                 os.makedirs(self.outpath, exist_ok=True)
                 with tempfile.NamedTemporaryFile(suffix=".png", dir=self.outpath) as fp:
                     start = time.perf_counter()
@@ -256,6 +264,7 @@ class HappyTreesBot(commands.Bot):
                 self.commissions.task_done()
             elif command == "paint":
                 self.logger.info(f'Bob Ross is painting "{prompt}" for {message.author}.')
+                self.currtasktime = time.time() + 30 + (2 * steps * samples)
                 # This section copies the filename generation code from 
                 # optimized_txt2img.py from optimizedSD
                 os.makedirs(self.outpath, exist_ok=True)