mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-09 00:57:04 +08:00
Remove polling operations.
This commit is contained in:
Robin Huang
parent
945ff14be1
commit
280997b383
@ -99,21 +99,15 @@ from typing import (
|
|||||||
Any,
|
Any,
|
||||||
TypeVar,
|
TypeVar,
|
||||||
Generic,
|
Generic,
|
||||||
Callable,
|
|
||||||
)
|
)
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
import time
|
|
||||||
import json
|
import json
|
||||||
import requests
|
import requests
|
||||||
from urllib.parse import urljoin
|
from urllib.parse import urljoin
|
||||||
|
|
||||||
# Import models from your generated stubs
|
|
||||||
|
|
||||||
T = TypeVar("T", bound=BaseModel)
|
T = TypeVar("T", bound=BaseModel)
|
||||||
R = TypeVar("R", bound=BaseModel)
|
R = TypeVar("R", bound=BaseModel)
|
||||||
P = TypeVar("P", bound=BaseModel) # For poll response
|
|
||||||
|
|
||||||
|
|
||||||
class EmptyRequest(BaseModel):
|
class EmptyRequest(BaseModel):
|
||||||
"""Base class for empty request bodies.
|
"""Base class for empty request bodies.
|
||||||
@ -337,125 +331,3 @@ class SynchronousOperation(Generic[T, R]):
|
|||||||
self.response = self.endpoint.response_model.model_validate(resp)
|
self.response = self.endpoint.response_model.model_validate(resp)
|
||||||
logging.debug(f"[DEBUG] Parsed Response: {self.response}")
|
logging.debug(f"[DEBUG] Parsed Response: {self.response}")
|
||||||
return self.response
|
return self.response
|
||||||
|
|
||||||
|
|
||||||
class TaskStatus(str, Enum):
|
|
||||||
"""Enum for task status values"""
|
|
||||||
|
|
||||||
COMPLETED = "completed"
|
|
||||||
FAILED = "failed"
|
|
||||||
PENDING = "pending"
|
|
||||||
|
|
||||||
|
|
||||||
class PollingOperation(Generic[T, R]):
|
|
||||||
"""
|
|
||||||
Represents an asynchronous API operation that requires polling for completion.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
poll_endpoint: ApiEndpoint[EmptyRequest, R],
|
|
||||||
completed_statuses: list,
|
|
||||||
failed_statuses: list,
|
|
||||||
status_extractor: Callable[[R], str],
|
|
||||||
request: Optional[T] = None,
|
|
||||||
api_base: str = "https://stagingapi.comfy.org",
|
|
||||||
auth_token: Optional[str] = None,
|
|
||||||
poll_interval: float = 1.0,
|
|
||||||
):
|
|
||||||
self.poll_endpoint = poll_endpoint
|
|
||||||
self.request = request
|
|
||||||
self.api_base = api_base
|
|
||||||
self.auth_token = auth_token
|
|
||||||
self.poll_interval = poll_interval
|
|
||||||
|
|
||||||
# Polling configuration
|
|
||||||
self.status_extractor = status_extractor or (
|
|
||||||
lambda x: getattr(x, "status", None)
|
|
||||||
)
|
|
||||||
self.completed_statuses = completed_statuses
|
|
||||||
self.failed_statuses = failed_statuses
|
|
||||||
|
|
||||||
# For storing response data
|
|
||||||
self.final_response = None
|
|
||||||
self.error = None
|
|
||||||
|
|
||||||
def execute(self, client: Optional[ApiClient] = None) -> R:
|
|
||||||
"""Execute the polling operation using the provided client. If failed, raise an exception."""
|
|
||||||
try:
|
|
||||||
if client is None:
|
|
||||||
client = ApiClient(
|
|
||||||
base_url=self.api_base,
|
|
||||||
api_key=self.auth_token,
|
|
||||||
)
|
|
||||||
return self._poll_until_complete(client)
|
|
||||||
except Exception as e:
|
|
||||||
raise Exception(f"Error during polling: {str(e)}")
|
|
||||||
|
|
||||||
def _check_task_status(self, response: R) -> TaskStatus:
|
|
||||||
"""Check task status using the status extractor function"""
|
|
||||||
try:
|
|
||||||
status = self.status_extractor(response)
|
|
||||||
if status in self.completed_statuses:
|
|
||||||
return TaskStatus.COMPLETED
|
|
||||||
elif status in self.failed_statuses:
|
|
||||||
return TaskStatus.FAILED
|
|
||||||
return TaskStatus.PENDING
|
|
||||||
except Exception as e:
|
|
||||||
logging.debug(f"Error extracting status: {e}")
|
|
||||||
return TaskStatus.PENDING
|
|
||||||
|
|
||||||
def _poll_until_complete(self, client: ApiClient) -> R:
|
|
||||||
"""Poll until the task is complete"""
|
|
||||||
poll_count = 0
|
|
||||||
while True:
|
|
||||||
try:
|
|
||||||
poll_count += 1
|
|
||||||
logging.debug(f"[DEBUG] Polling attempt #{poll_count}")
|
|
||||||
|
|
||||||
request_dict = (
|
|
||||||
self.request.model_dump(exclude_none=True)
|
|
||||||
if self.request is not None
|
|
||||||
else None
|
|
||||||
)
|
|
||||||
|
|
||||||
if poll_count == 1:
|
|
||||||
logging.debug(
|
|
||||||
f"[DEBUG] Poll Request: {self.poll_endpoint.method.value} {self.poll_endpoint.path}"
|
|
||||||
)
|
|
||||||
logging.debug(
|
|
||||||
f"[DEBUG] Poll Request Data: {json.dumps(request_dict, indent=2) if request_dict else 'None'}"
|
|
||||||
)
|
|
||||||
|
|
||||||
# Query task status
|
|
||||||
resp = client.request(
|
|
||||||
method=self.poll_endpoint.method.value,
|
|
||||||
path=self.poll_endpoint.path,
|
|
||||||
params=self.poll_endpoint.query_params,
|
|
||||||
json=request_dict,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Parse response
|
|
||||||
response_obj = self.poll_endpoint.response_model.model_validate(resp)
|
|
||||||
|
|
||||||
# Check if task is complete
|
|
||||||
status = self._check_task_status(response_obj)
|
|
||||||
logging.debug(f"[DEBUG] Task Status: {status}")
|
|
||||||
|
|
||||||
if status == TaskStatus.COMPLETED:
|
|
||||||
logging.debug("[DEBUG] Task completed successfully")
|
|
||||||
self.final_response = response_obj
|
|
||||||
return self.final_response
|
|
||||||
elif status == TaskStatus.FAILED:
|
|
||||||
logging.debug(f"[DEBUG] Task failed: {json.dumps(resp)}")
|
|
||||||
raise Exception(f"Task failed: {json.dumps(resp)}")
|
|
||||||
else:
|
|
||||||
logging.debug("[DEBUG] Task still pending, continuing to poll...")
|
|
||||||
|
|
||||||
# Wait before polling again
|
|
||||||
logging.debug(f"[DEBUG] Waiting {self.poll_interval} seconds before next poll")
|
|
||||||
time.sleep(self.poll_interval)
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logging.debug(f"[DEBUG] Polling error: {str(e)}")
|
|
||||||
raise Exception(f"Error while polling: {str(e)}")
|
|
||||||
|
|||||||
@ -1,14 +1,7 @@
|
|||||||
import os
|
|
||||||
import requests
|
|
||||||
from inspect import cleandoc
|
from inspect import cleandoc
|
||||||
from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, IO
|
from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, IO
|
||||||
from comfy_api_nodes.apis.client import ApiEndpoint, SynchronousOperation, HttpMethod, PollingOperation, EmptyRequest
|
from comfy_api_nodes.apis.client import ApiEndpoint, SynchronousOperation, HttpMethod
|
||||||
from comfy_api_nodes.apis.stubs import IdeogramGenerateRequest, IdeogramGenerateResponse, ImageRequest, MinimaxVideoGenerationRequest, MinimaxVideoGenerationResponse, MinimaxFileRetrieveResponse, MinimaxTaskResultResponse, Model
|
from comfy_api_nodes.apis.stubs import IdeogramGenerateRequest, IdeogramGenerateResponse, ImageRequest
|
||||||
import folder_paths
|
|
||||||
import logging
|
|
||||||
from comfy.comfy_types.node_typing import FileLocator
|
|
||||||
import json
|
|
||||||
import av
|
|
||||||
|
|
||||||
class IdeogramTextToImage(ComfyNodeABC):
|
class IdeogramTextToImage(ComfyNodeABC):
|
||||||
"""
|
"""
|
||||||
@ -164,185 +157,13 @@ class IdeogramTextToImage(ComfyNodeABC):
|
|||||||
# return ""
|
# return ""
|
||||||
|
|
||||||
|
|
||||||
class MinimaxVideoNode:
|
|
||||||
"""
|
|
||||||
Generates videos synchronously based on a prompt, and optional parameters using Minimax's API.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self):
|
|
||||||
self.output_dir = folder_paths.get_output_directory()
|
|
||||||
self.type = "output"
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(s):
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"prompt_text": (
|
|
||||||
"STRING",
|
|
||||||
{
|
|
||||||
"multiline": True,
|
|
||||||
"default": "",
|
|
||||||
"tooltip": "Text prompt to guide the video generation",
|
|
||||||
},
|
|
||||||
),
|
|
||||||
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
|
||||||
"model": (
|
|
||||||
[
|
|
||||||
"T2V-01",
|
|
||||||
"I2V-01-Director",
|
|
||||||
"S2V-01",
|
|
||||||
"I2V-01",
|
|
||||||
"I2V-01-live",
|
|
||||||
"T2V-01",
|
|
||||||
],
|
|
||||||
{
|
|
||||||
"default": "T2V-01",
|
|
||||||
"tooltip": "Model to use for video generation",
|
|
||||||
},
|
|
||||||
),
|
|
||||||
},
|
|
||||||
"optional": {
|
|
||||||
"seed": (
|
|
||||||
IO.INT,
|
|
||||||
{
|
|
||||||
"default": 0,
|
|
||||||
"min": 0,
|
|
||||||
"max": 0xFFFFFFFFFFFFFFFF,
|
|
||||||
"control_after_generate": True,
|
|
||||||
"tooltip": "The random seed used for creating the noise.",
|
|
||||||
},
|
|
||||||
),
|
|
||||||
},
|
|
||||||
"hidden": {
|
|
||||||
"prompt": "PROMPT",
|
|
||||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
||||||
"auth_token": "AUTH_TOKEN_COMFY_ORG",
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("VIDEO",)
|
|
||||||
DESCRIPTION = "Generates videos from prompts using Minimax's API"
|
|
||||||
FUNCTION = "generate_video"
|
|
||||||
CATEGORY = "video"
|
|
||||||
API_NODE = True
|
|
||||||
OUTPUT_NODE = True
|
|
||||||
|
|
||||||
def generate_video(
|
|
||||||
self,
|
|
||||||
prompt_text,
|
|
||||||
filename_prefix,
|
|
||||||
seed=0,
|
|
||||||
model="T2V-01",
|
|
||||||
prompt=None,
|
|
||||||
extra_pnginfo=None,
|
|
||||||
auth_token=None,
|
|
||||||
):
|
|
||||||
video_generate_operation = SynchronousOperation(
|
|
||||||
endpoint=ApiEndpoint(
|
|
||||||
path="/proxy/minimax/video_generation",
|
|
||||||
method=HttpMethod.POST,
|
|
||||||
request_model=MinimaxVideoGenerationRequest,
|
|
||||||
response_model=MinimaxVideoGenerationResponse,
|
|
||||||
),
|
|
||||||
request=MinimaxVideoGenerationRequest(
|
|
||||||
model=Model(model),
|
|
||||||
prompt=prompt_text,
|
|
||||||
callback_url=None,
|
|
||||||
first_frame_image=None,
|
|
||||||
subject_reference=None,
|
|
||||||
prompt_optimizer=None,
|
|
||||||
),
|
|
||||||
auth_token=auth_token,
|
|
||||||
)
|
|
||||||
response = video_generate_operation.execute()
|
|
||||||
|
|
||||||
task_id = response.task_id
|
|
||||||
|
|
||||||
video_generate_operation = PollingOperation(
|
|
||||||
poll_endpoint=ApiEndpoint(
|
|
||||||
path="/proxy/minimax/query/video_generation",
|
|
||||||
method=HttpMethod.GET,
|
|
||||||
request_model=EmptyRequest,
|
|
||||||
response_model=MinimaxTaskResultResponse,
|
|
||||||
query_params={"task_id": task_id},
|
|
||||||
),
|
|
||||||
completed_statuses=["Success"],
|
|
||||||
failed_statuses=["Fail"],
|
|
||||||
status_extractor=lambda x: x.status.value,
|
|
||||||
auth_token=auth_token,
|
|
||||||
)
|
|
||||||
task_result = video_generate_operation.execute()
|
|
||||||
|
|
||||||
file_id = task_result.file_id
|
|
||||||
if file_id is None:
|
|
||||||
raise Exception("Request was not successful. Missing file ID.")
|
|
||||||
file_retrieve_operation = SynchronousOperation(
|
|
||||||
endpoint=ApiEndpoint(
|
|
||||||
path="/proxy/minimax/files/retrieve",
|
|
||||||
method=HttpMethod.GET,
|
|
||||||
request_model=EmptyRequest,
|
|
||||||
response_model=MinimaxFileRetrieveResponse,
|
|
||||||
query_params={"file_id": int(file_id)},
|
|
||||||
),
|
|
||||||
request=EmptyRequest(),
|
|
||||||
auth_token=auth_token,
|
|
||||||
)
|
|
||||||
file_result = file_retrieve_operation.execute()
|
|
||||||
|
|
||||||
file_url = file_result.file.download_url
|
|
||||||
if file_url is None:
|
|
||||||
raise Exception(f"No video was found in the response. Full response: {file_result.model_dump()}")
|
|
||||||
logging.info(f"Generated video URL: {file_url}")
|
|
||||||
|
|
||||||
# Construct the save path
|
|
||||||
full_output_folder, filename, counter, subfolder, filename_prefix = (
|
|
||||||
folder_paths.get_save_image_path(filename_prefix, self.output_dir)
|
|
||||||
)
|
|
||||||
file_basename = f"{filename}_{counter:05}_.mp4"
|
|
||||||
save_path = os.path.join(full_output_folder, file_basename)
|
|
||||||
|
|
||||||
# Download the video data
|
|
||||||
video_response = requests.get(file_url)
|
|
||||||
video_data = video_response.content
|
|
||||||
|
|
||||||
# Save the video data to a file
|
|
||||||
with open(save_path, "wb") as video_file:
|
|
||||||
video_file.write(video_data)
|
|
||||||
|
|
||||||
# Add workflow metadata to the video container
|
|
||||||
if prompt is not None or extra_pnginfo is not None:
|
|
||||||
try:
|
|
||||||
container = av.open(save_path, mode="r+")
|
|
||||||
if prompt is not None:
|
|
||||||
container.metadata["prompt"] = json.dumps(prompt)
|
|
||||||
if extra_pnginfo is not None:
|
|
||||||
for x in extra_pnginfo:
|
|
||||||
container.metadata[x] = json.dumps(extra_pnginfo[x])
|
|
||||||
container.close()
|
|
||||||
except Exception as e:
|
|
||||||
logging.warning(f"Failed to add metadata to video: {e}")
|
|
||||||
|
|
||||||
# Create a FileLocator for the frontend to use for the preview
|
|
||||||
results: list[FileLocator] = [
|
|
||||||
{
|
|
||||||
"filename": file_basename,
|
|
||||||
"subfolder": subfolder,
|
|
||||||
"type": self.type,
|
|
||||||
}
|
|
||||||
]
|
|
||||||
|
|
||||||
return {"ui": {"images": results, "animated": (True,)}}
|
|
||||||
|
|
||||||
|
|
||||||
# A dictionary that contains all nodes you want to export with their names
|
# A dictionary that contains all nodes you want to export with their names
|
||||||
# NOTE: names should be globally unique
|
# NOTE: names should be globally unique
|
||||||
NODE_CLASS_MAPPINGS = {
|
NODE_CLASS_MAPPINGS = {
|
||||||
"IdeogramTextToImage": IdeogramTextToImage,
|
"IdeogramTextToImage": IdeogramTextToImage,
|
||||||
"MinimaxVideoNode": MinimaxVideoNode
|
|
||||||
}
|
}
|
||||||
|
|
||||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
"IdeogramTextToImage": "Ideogram Text to Image",
|
"IdeogramTextToImage": "Ideogram Text to Image",
|
||||||
"MinimaxVideoNode": "Minimax Video Generator"
|
|
||||||
}
|
}
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user