From f29a3194b4068c587d574fe10b670d3a2b5e6f35 Mon Sep 17 00:00:00 2001
From: Christian Byrne <cbyrne@comfy.org>
Date: Mon, 28 Apr 2025 15:17:57 -0700
Subject: [PATCH] Manually add BFL polling status response schema (#15)

---
 comfy_api_nodes/apis/BFLPolling.py | 29 +++++++++++++++++++++++++++++
 comfy_api_nodes/nodes_api.py       |  2 +-
 2 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 comfy_api_nodes/apis/BFLPolling.py

diff --git a/comfy_api_nodes/apis/BFLPolling.py b/comfy_api_nodes/apis/BFLPolling.py
new file mode 100644
index 000000000..e40bf3344
--- /dev/null
+++ b/comfy_api_nodes/apis/BFLPolling.py
@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field, confloat
+
+
+class BFLStatus(str, Enum):
+    task_not_found = "Task not found"
+    pending = "Pending"
+    request_moderated = "Request Moderated"
+    content_moderated = "Content Moderated"
+    ready = "Ready"
+    error = "Error"
+
+
+class BFLFluxProStatusResponse(BaseModel):
+    id: str = Field(..., description="The unique identifier for the generation task.")
+    status: BFLStatus = Field(..., description="The status of the task.")
+    result: Optional[Dict[str, Any]] = Field(
+        None, description="The result of the task (null if not completed)."
+    )
+    progress: confloat(ge=0.0, le=1.0) = Field(
+        ..., description="The progress of the task (0.0 to 1.0)."
+    )
+    details: Optional[Dict[str, Any]] = Field(
+        None, description="Additional details about the task (null if not available)."
+    )
diff --git a/comfy_api_nodes/nodes_api.py b/comfy_api_nodes/nodes_api.py
index 605f84960..03100aab5 100644
--- a/comfy_api_nodes/nodes_api.py
+++ b/comfy_api_nodes/nodes_api.py
@@ -16,10 +16,10 @@ from comfy_api_nodes.apis import (
     IdeogramGenerateResponse,
     BFLFluxProGenerateRequest,
     BFLFluxProGenerateResponse,
-    BFLStatus,
     ImageRequest,
     Model
 )
+from comfy_api_nodes.apis.BFLPolling import BFLStatus
 from comfy_api_nodes.apis.client import ApiEndpoint, HttpMethod, SynchronousOperation, PollingOperation, EmptyRequest
 
 import numpy as np