From 2ea7b3efc0f9d2897a9adf19a924c9600a8f6fdb Mon Sep 17 00:00:00 2001
From: Konie <guofengzuo@hotmail.com>
Date: Fri, 29 Dec 2023 12:01:58 +0800
Subject: [PATCH] Fix startup error

---
 fooocusapi/api.py        | 10 +++++++---
 fooocusapi/models.py     | 17 +++++++++++++----
 fooocusapi/task_queue.py |  4 ++--
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/fooocusapi/api.py b/fooocusapi/api.py
index 73c42c6..49fc559 100644
--- a/fooocusapi/api.py
+++ b/fooocusapi/api.py
@@ -7,7 +7,7 @@
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
 
-from fooocusapi.models import AllModelNamesResponse, AsyncJobResponse, QueryJobRequest,StopResponse , GeneratedImageResult, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, JobQueueInfo, HistoryResponse, Text2ImgRequest
+from fooocusapi.models import AllModelNamesResponse, AsyncJobResponse, QueryJobRequest,StopResponse , GeneratedImageResult, ImgInpaintOrOutpaintRequest, ImgPromptRequest, ImgUpscaleOrVaryRequest, JobQueueInfo, JobHistoryResponse, Text2ImgRequest
 from fooocusapi.api_utils import generation_output, req_to_params
 import fooocusapi.file_utils as file_utils
 from fooocusapi.parameters import GenerationFinishReason, ImageGenerationResult
@@ -254,10 +254,14 @@ def query_job(req: QueryJobRequest = Depends()):
 def job_queue():
     return JobQueueInfo(running_size=len(task_queue.queue), finished_size=len(task_queue.history), last_job_id=task_queue.last_job_id)
 
-@app.get("/v1/generation/job-history", response_model=HistoryResponse, description="Query historical job data")
+
+@app.get("/v1/generation/job-history", response_model=JobHistoryResponse, description="Query historical job data")
 def get_history():
     # Fetch and return the historical tasks
-    return HistoryResponse(history=task_queue.history, queue=task_queue.queue)
+    hitory = [JobHistoryInfo(job_id=item.job_id, is_finished=item.is_finished) for item in task_queue.history]
+    queue = [JobHistoryInfo(job_id=item.job_id, is_finished=item.is_finished) for item in task_queue.queue]
+    return JobHistoryResponse(history=hitory, queue=queue)
+
 
 @app.post("/v1/generation/stop", response_model=StopResponse, description="Job stoping")
 def stop():
diff --git a/fooocusapi/models.py b/fooocusapi/models.py
index 724c5ac..316da77 100644
--- a/fooocusapi/models.py
+++ b/fooocusapi/models.py
@@ -11,7 +11,7 @@
 from enum import Enum
 
 from fooocusapi.parameters import GenerationFinishReason, defualt_styles, default_base_model_name, default_refiner_model_name, default_refiner_switch, default_loras, default_cfg_scale, default_prompt_negative, default_aspect_ratio, default_sampler, default_scheduler
-from fooocusapi.task_queue import TaskType
+from fooocusapi.task_queue import QueueTask, TaskType
 
 from modules import flags
 
@@ -403,10 +403,18 @@ class JobQueueInfo(BaseModel):
     finished_size: int = Field(description="Finished job cound (after auto clean)")
     last_job_id: str = Field(description="Last submit generation job id")
 
+
+# TODO May need more detail fields, will add later when someone need
+class JobHistoryInfo(BaseModel):
+    job_id: str
+    is_finished: bool = False
+
+
 # Response model for the historical tasks
-class HistoryResponse(BaseModel):
-    queue: List[QueueTask] = []
-    history: List[QueueTask] = []
+class JobHistoryResponse(BaseModel):
+    queue: List[JobHistoryInfo] = []
+    history: List[JobHistoryInfo] = []
+
 
 class AllModelNamesResponse(BaseModel):
     model_filenames: List[str] = Field(description="All available model filenames")
@@ -415,6 +423,7 @@ class AllModelNamesResponse(BaseModel):
     model_config = ConfigDict(
         protected_namespaces=('protect_me_', 'also_protect_')
     )
+
     
 class StopResponse(BaseModel):
     msg: str
diff --git a/fooocusapi/task_queue.py b/fooocusapi/task_queue.py
index 76f4a74..40371c3 100644
--- a/fooocusapi/task_queue.py
+++ b/fooocusapi/task_queue.py
@@ -26,7 +26,7 @@ class QueueTask(object):
     finish_with_error: bool = False
     task_status: str | None = None
     task_step_preview: str | None = None
-    task_result: any = None
+    task_result: List[ImageGenerationResult] = None
     error_message: str | None = None
     webhook_url: str | None = None  # attribute for individual webhook_url
 
@@ -47,7 +47,7 @@ def set_progress(self, progress: int, status: str | None):
     def set_step_preview(self, task_step_preview: str | None):
         self.task_step_preview = task_step_preview
 
-    def set_result(self, task_result: any, finish_with_error: bool, error_message: str | None = None):
+    def set_result(self, task_result: List[ImageGenerationResult], finish_with_error: bool, error_message: str | None = None):
         if not finish_with_error:
             self.finish_progress = 100
             self.task_status = 'Finished'