From 81f55fddf65d53734e38a1f9f896eb13ce5cb160 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Mon, 21 Apr 2025 23:39:14 -0700 Subject: [PATCH] Revert "Remove polling operations." This reverts commit 8415404ce8fbc0262b7de54fc700c5c8854a34fc. --- comfy_api_nodes/apis/client.py | 128 +++++++++++++++++++++++++ comfy_api_nodes/nodes_api.py | 170 +++++++++++++++++++++++++++++++++ 2 files changed, 298 insertions(+) diff --git a/comfy_api_nodes/apis/client.py b/comfy_api_nodes/apis/client.py index 05706710e..faee11230 100644 --- a/comfy_api_nodes/apis/client.py +++ b/comfy_api_nodes/apis/client.py @@ -99,15 +99,21 @@ from typing import ( Any, TypeVar, Generic, + Callable, ) from pydantic import BaseModel from enum import Enum +import time import json import requests from urllib.parse import urljoin +# Import models from your generated stubs + T = TypeVar("T", bound=BaseModel) R = TypeVar("R", bound=BaseModel) +P = TypeVar("P", bound=BaseModel) # For poll response + class EmptyRequest(BaseModel): """Base class for empty request bodies. @@ -379,3 +385,125 @@ class SynchronousOperation(Generic[T, R]): self.response = self.endpoint.response_model.model_validate(resp) logging.debug(f"[DEBUG] Parsed Response: {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)}") diff --git a/comfy_api_nodes/nodes_api.py b/comfy_api_nodes/nodes_api.py index 927a05fd2..dbdbcdf60 100644 --- a/comfy_api_nodes/nodes_api.py +++ b/comfy_api_nodes/nodes_api.py @@ -428,6 +428,176 @@ class OpenAIGPTImage1(ComfyNodeABC): return (img_tensor,) +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 # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = {