From a3ea827aa6b157d80bd855877fca41abb401a18f Mon Sep 17 00:00:00 2001
From: DK <35783662+gottyduke@users.noreply.github.com>
Date: Fri, 21 Jun 2024 19:18:37 -0500
Subject: [PATCH] feat: poll for all jobs in wait mode

---
 civitai/__init__.py | 49 ++++++++++++++++++++++++++-------------------
 1 file changed, 28 insertions(+), 21 deletions(-)

diff --git a/civitai/__init__.py b/civitai/__init__.py
index 8eff7e9..7e8acd9 100644
--- a/civitai/__init__.py
+++ b/civitai/__init__.py
@@ -80,38 +80,45 @@ async def async_create():
         async def _poll_for_job_completion(self, token, interval=1, timeout=300):
             start_time = time.time()
             last_response = None
+            completed_jobs = {}
             while time.time() - start_time < timeout:
-                response = await get_v1consumerjobs(token, api_config_override=self.civitai)
+                response = await get_v1consumerjobs(
+                    token, api_config_override=self.civitai
+                )
                 if response and response.jobs:
                     last_response = response
                     logging.info(f"Job status: {last_response}")
-                    completed_jobs = []
                     for job in response.jobs:
-                        if job.result and job.result.get("blobUrl"):
-                            completed_jobs.append({
-                                "jobId": job.jobId,
-                                "cost": job.cost,
-                                "result": job.result,
-                                "scheduled": job.scheduled,
-                            })
-                    if completed_jobs:
-                        return {
-                            "token": response.token,
-                            "jobs": completed_jobs
-                        }
+                        if job.jobId not in completed_jobs:
+                            if job.result and job.result.get("blobUrl"):
+                                completed_jobs[job.jobId] = {
+                                    "jobId": job.jobId,
+                                    "cost": job.cost,
+                                    "result": job.result,
+                                    "scheduled": job.scheduled,
+                                }
+                    if len(completed_jobs) >= len(response.jobs):
+                        break
                 await asyncio.sleep(interval)
 
+            if completed_jobs:
+                return {"token": response.token, "jobs": list(completed_jobs.values())}
+
             if last_response:
                 logging.warning(
-                    "Job {token} did not complete within 5 minutes.")
+                    f"Job {token} did not complete within {timeout // 60} minutes."
+                )
                 return {
                     "token": last_response.token,
-                    "jobs": [{
-                        "jobId": job.jobId,
-                        "cost": job.cost,
-                        "result": job.result,
-                        "scheduled": job.scheduled,
-                    } for job in last_response.jobs]
+                    "jobs": [
+                        {
+                            "jobId": job.jobId,
+                            "cost": job.cost,
+                            "result": job.result,
+                            "scheduled": job.scheduled,
+                        }
+                        for job in last_response.jobs
+                    ],
                 }
 
     class Jobs: