Merge pull request #190 from konieshadow/dev

Optimize worker error handling
mrhan1993 · Jan 24, 2024 · 9dd0242 · 9dd0242
2 parents 95fc3f8 + cef5150
commit 9dd0242
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/fooocusapi/models.py b/fooocusapi/models.py
@@ -125,7 +125,7 @@ class Text2ImgRequest(BaseModel):
     advanced_params: AdvancedParams | None = AdvancedParams()
     require_base64: bool = Field(default=False, description="Return base64 data of generated image")
     async_process: bool = Field(default=False, description="Set to true will run async and return job info for retrieve generataion result later")
-    webhook_url: str | None = Field(default=None, description="Optional URL for a webhook callback. If provided, the system will send a POST request to this URL upon task completion or failure."
+    webhook_url: str | None = Field(default='', description="Optional URL for a webhook callback. If provided, the system will send a POST request to this URL upon task completion or failure."
                                                               " This allows for asynchronous notification of task status.")
 
 def style_selection_parser(style_selections: str) -> List[str]:

diff --git a/fooocusapi/task_queue.py b/fooocusapi/task_queue.py
@@ -120,11 +120,13 @@ def finish_task(self, job_id: str):
             webhook_url = task.webhook_url or self.webhook_url
 
             data = { "job_id": task.job_id, "job_result": [] }
-            for item in task.task_result:
-                data["job_result"].append({
-                    "url": get_file_serve_url(item.im) if item.im else None,
-                    "seed": item.seed if item.seed else "-1",
-                })
+
+            if isinstance(task.task_result, List):
+                for item in task.task_result:
+                    data["job_result"].append({
+                        "url": get_file_serve_url(item.im) if item.im else None,
+                        "seed": item.seed if item.seed else "-1",
+                    })
 
             # Send webhook
             if task.is_finished and webhook_url:

diff --git a/fooocusapi/worker.py b/fooocusapi/worker.py
@@ -842,6 +842,7 @@ def callback(step, x0, x, total_steps, y):
                 break
             except Exception as e:
                 print('Process error:', e)
+                logging.exception(e)
                 results.append(ImageGenerationResult(
                     im=None, seed=task['task_seed'], finish_reason=GenerationFinishReason.error))
                 async_task.set_result(results, True, str(e))
@@ -859,7 +860,7 @@ def callback(step, x0, x, total_steps, y):
         logging.exception(e)
 
         if not async_task.is_finished:
-            task_queue.finish_task(async_task.job_id)
             async_task.set_result([], True, str(e))
+            task_queue.finish_task(async_task.job_id)
             print(f"[Task Queue] Finish task with error, job_id={async_task.job_id}")
         return []