feat(eta): Improve duration UX

- Render result without waiting for last poll interval to complete
- Calculate ETA based on past completions
- Update ETA during generation based on place in queue

app.py

@@ -1,4 +1,5 @@
-import time
+from time import time
+from statistics import mean
 from flask import Flask, jsonify, render_template, request
 
 from modules.details import load_lists, rand_details
@@ -6,6 +7,8 @@ from modules.inference import generate_image
 
 app = Flask(__name__)
 
+TEMPLATES_AUTO_RELOAD = True
+
 
 @app.route('/')
 def index():
@@ -19,10 +22,13 @@ tasks = {}
 def create_task():
     prompt = request.args.get('prompt') or "покемон"
 
-    task_id = f"{str(time.time())}_{prompt}"
+    created_at = time()
+
+    task_id = f"{str(created_at)}_{prompt}"
 
     tasks[task_id] = {
         "task_id": task_id,
+        "created_at": created_at,
         "prompt": prompt,
         "status": "pending",
         "poll_count": 0,
@@ -37,7 +43,9 @@ def queue_task():
 
     tasks[task_id]["value"] = generate_image(tasks[task_id]["prompt"])
 
-    tasks[task_id]["status"] = "complete"
+    tasks[task_id]["status"] = "completed"
+
+    tasks[task_id]["completed_at"] = time()
 
     return jsonify(tasks[task_id])
 
@@ -46,6 +54,27 @@ def queue_task():
 def poll_task():
     task_id = request.args.get('task_id')
 
+    pending_tasks = []
+    completed_durations = []
+
+    for task in tasks.values():
+        if task["status"] == "pending":
+            pending_tasks.append(task["task_id"])
+        elif task["status"] == "completed":
+            completed_durations.append(task["completed_at"] - task["created_at"])
+
+    try:
+        place_in_queue = pending_tasks.index(task_id) + 1
+    except:
+        place_in_queue = 0
+
+    if (len(completed_durations)):
+        eta = sum(completed_durations) / len(completed_durations) * (place_in_queue or 1)
+    else:
+        eta = 40 * (place_in_queue or 1)
+
+    tasks[task_id]["place_in_queue"] = place_in_queue
+    tasks[task_id]["eta"] = round(eta, 1)
     tasks[task_id]["poll_count"] += 1
 
     return jsonify(tasks[task_id])

modules/inference.py

@@ -13,14 +13,11 @@ fp16 = torch.cuda.is_available()
 
 file_dir = "./models"
 file_name = "pytorch_model.bin"
-config_file_url = hf_hub_url(
-    repo_id="minimaxir/ai-generated-pokemon-rudalle", filename=file_name)
+config_file_url = hf_hub_url(repo_id="minimaxir/ai-generated-pokemon-rudalle", filename=file_name)
 cached_download(config_file_url, cache_dir=file_dir, force_filename=file_name)
 
-model = get_rudalle_model('Malevich', pretrained=False,
-                          fp16=fp16, device=device)
-model.load_state_dict(torch.load(
-    f"{file_dir}/{file_name}", map_location=f"{'cuda:0' if torch.cuda.is_available() else 'cpu'}"))
+model = get_rudalle_model('Malevich', pretrained=False, fp16=fp16, device=device)
+model.load_state_dict(torch.load(f"{file_dir}/{file_name}", map_location=f"{'cuda:0' if torch.cuda.is_available() else 'cpu'}"))
 
 vae = get_vae().to(device)
 tokenizer = get_tokenizer()
@@ -50,8 +47,7 @@ def generate_image(prompt):
     if prompt.lower() in ['grass', 'fire', 'water', 'lightning', 'fighting', 'psychic', 'colorless', 'darkness', 'metal', 'dragon', 'fairy']:
         prompt = english_to_russian(prompt)
 
-    result, _ = generate_images(
-        prompt, tokenizer, model, vae, top_k=2048, images_num=1, top_p=0.995)
+    result, _ = generate_images(prompt, tokenizer, model, vae, top_k=2048, images_num=1, top_p=0.995)
 
     buffer = BytesIO()
     result[0].save(buffer, format="PNG")

static/index.js

@@ -129,8 +129,9 @@ const createTask = async (prompt) => {
   return task;
 };
 
-const queueTask = (task_id) => {
-  fetch(`${getBasePath()}task/queue?task_id=${task_id}`);
+const queueTask = async (task_id) => {
+  const queueResponse = await fetch(`${getBasePath()}task/queue?task_id=${task_id}`);
+  return queueResponse.json();
 };
 
 const pollTask = async (task) => {
@@ -140,18 +141,16 @@ const pollTask = async (task) => {
 };
 
 const longPollTask = async (task, interval = 10_000, max) => {
-  if (task.status === 'complete' || (max && task.poll_count > max)) {
-    return task;
-  }
+  const etaDisplay = document.querySelector('.eta');
 
-  const taskResponse = await fetch(`${getBasePath()}task/poll?task_id=${task.task_id}`);
+  task = await pollTask(task);
 
-  task = await taskResponse.json();
-
-  if (task.status === 'complete' || task.poll_count > max) {
+  if (task.status === 'completed' || (max && task.poll_count > max)) {
     return task;
   }
 
+  etaDisplay.textContent = Math.round(task.eta);
+
   await new Promise((resolve) => setTimeout(resolve, interval));
 
   return await longPollTask(task, interval, max);
@@ -162,26 +161,34 @@ const longPollTask = async (task, interval = 10_000, max) => {
 const generateButton = document.querySelector('button.generate');
 
 const durationTimer = () => {
+  const elapsedDisplay = document.querySelector('.elapsed');
   let duration = 0.0;
 
-  return (secondsElement) => {
+  return () => {
     const startTime = performance.now();
 
     const incrementSeconds = setInterval(() => {
       duration += 0.1;
-      secondsElement.textContent = duration.toFixed(1);
+      elapsedDisplay.textContent = duration.toFixed(1);
     }, 100);
 
-    const updateDuration = () => (duration = Number(((performance.now() - startTime) / 1_000).toFixed(1)));
+    const updateDuration = (task) => {
+      if (task?.status == 'completed') {
+        duration = task.completed_at - task.created_at;
+        return;
+      }
+
+      duration = Number(((performance.now() - startTime) / 1_000).toFixed(1));
+    };
 
     window.addEventListener('focus', updateDuration);
 
     return {
-      cleanup: () => {
-        updateDuration();
+      cleanup: (completedTask) => {
+        updateDuration(completedTask);
         clearInterval(incrementSeconds);
         window.removeEventListener('focus', updateDuration);
-        secondsElement.textContent = duration.toFixed(1);
+        elapsedDisplay.textContent = duration.toFixed(1);
       },
     };
   };
@@ -238,7 +245,7 @@ generateButton.addEventListener('click', async () => {
   }
 
   const renderSection = document.querySelector('section.render');
-  const durationSeconds = document.querySelector('.duration > .seconds');
+  const durationDisplay = document.querySelector('.duration');
   const initialiseCardRotation = cardRotationInitiator(renderSection);
 
   try {
@@ -246,14 +253,15 @@ generateButton.addEventListener('click', async () => {
 
     const details = await generateDetails();
     const task = await createTask(details.energy_type);
-    queueTask(task.task_id);
 
-    const timer = durationTimer();
-    const cleanupTimer = timer(durationSeconds).cleanup;
+    const timer = durationTimer(durationDisplay);
+    const timerCleanup = timer().cleanup;
+
+    const longPromises = [queueTask(task.task_id), longPollTask(task)];
+    const completedTask = await Promise.any(longPromises);
 
-    const completedTask = await longPollTask(task);
     generating = false;
-    cleanupTimer();
+    timerCleanup(completedTask);
 
     renderSection.innerHTML = cardHTML(details);
     const picture = document.querySelector('img.picture');

templates/index.html

@@ -15,7 +15,7 @@
   </head>
   <body>
     <h1>This Pokémon Does Not Exist</h1>
-    <div class="duration"><span class="seconds">0.0</span>s (ETA: 490s)</div>
+    <div class="duration"><span class="elapsed">0.0</span>s (ETA: <span class="eta">40</span>s)</div>
     <button class="generate">Generate Pokémon Card with AI</button>
     <section class="render"></section>
   </body>