Skip to content

Commit

Permalink
Fix startup error
Browse files Browse the repository at this point in the history
  • Loading branch information
konieshadow committed Dec 29, 2023
1 parent ae34eae commit 2ea7b3e
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 9 deletions.
10 changes: 7 additions & 3 deletions fooocusapi/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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():
Expand Down
17 changes: 13 additions & 4 deletions fooocusapi/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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")
Expand All @@ -415,6 +423,7 @@ class AllModelNamesResponse(BaseModel):
model_config = ConfigDict(
protected_namespaces=('protect_me_', 'also_protect_')
)


class StopResponse(BaseModel):
msg: str
4 changes: 2 additions & 2 deletions fooocusapi/task_queue.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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'
Expand Down

0 comments on commit 2ea7b3e

Please sign in to comment.