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/apis/__init__.py b/comfy_api_nodes/apis/__init__.py index ee8a70c3a..4ca4d5ab9 100644 --- a/comfy_api_nodes/apis/__init__.py +++ b/comfy_api_nodes/apis/__init__.py @@ -42,15 +42,6 @@ class BFLFluxProGenerateResponse(BaseModel): polling_url: str = Field(..., description='URL to poll for the generation result.') -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 ErrorResponse(BaseModel): error: str message: str diff --git a/comfy_api_nodes/nodes_api.py b/comfy_api_nodes/nodes_api.py index f7e6a234a..9ab1d72aa 100644 --- a/comfy_api_nodes/nodes_api.py +++ b/comfy_api_nodes/nodes_api.py @@ -20,10 +20,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