diff --git a/nodes/base.py b/nodes/base.py index b964e33..1e144a5 100644 --- a/nodes/base.py +++ b/nodes/base.py @@ -107,8 +107,8 @@ class BaseFlux: if Status(status) == Status.READY: print(f"[BFL] Task {task_id} ready after {elapsed:.1f}s — downloading image") return self.process_result(result, output_format=output_format) - elif Status(status) == Status.PENDING: - print(f"[BFL] Attempt {attempt}/{max_attempts}: pending — retrying in 5s") + elif Status(status) in [Status.PENDING, Status.REASONING, Status.GENERATING]: + print(f"[BFL] Attempt {attempt}/{max_attempts}: {status} — retrying in 5s") attempt += 1 if attempt <= max_attempts: time.sleep(5) diff --git a/nodes/status.py b/nodes/status.py index ba4efe5..826828d 100644 --- a/nodes/status.py +++ b/nodes/status.py @@ -1,9 +1,12 @@ from enum import Enum + class Status(Enum): TASK_NOT_FOUND = "Task not found" PENDING = "Pending" + REASONING = "Reasoning" + GENERATING = "Generating" REQUEST_MODERATED = "Request Moderated" CONTENT_MODERATED = "Content Moderated" READY = "Ready" - ERROR = "Error" \ No newline at end of file + ERROR = "Error"