diff --git a/comfy_api_nodes/apis/client.py b/comfy_api_nodes/apis/client.py index 9e0676a62..6240a0241 100644 --- a/comfy_api_nodes/apis/client.py +++ b/comfy_api_nodes/apis/client.py @@ -99,21 +99,15 @@ 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. @@ -337,125 +331,3 @@ 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 a2b840c27..24bd2f3ff 100644 --- a/comfy_api_nodes/nodes_api.py +++ b/comfy_api_nodes/nodes_api.py @@ -1,14 +1,7 @@ -import os -import requests from inspect import cleandoc 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.stubs import IdeogramGenerateRequest, IdeogramGenerateResponse, ImageRequest, MinimaxVideoGenerationRequest, MinimaxVideoGenerationResponse, MinimaxFileRetrieveResponse, MinimaxTaskResultResponse, Model -import folder_paths -import logging -from comfy.comfy_types.node_typing import FileLocator -import json -import av +from comfy_api_nodes.apis.client import ApiEndpoint, SynchronousOperation, HttpMethod +from comfy_api_nodes.apis.stubs import IdeogramGenerateRequest, IdeogramGenerateResponse, ImageRequest class IdeogramTextToImage(ComfyNodeABC): """ @@ -164,185 +157,13 @@ class IdeogramTextToImage(ComfyNodeABC): # 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 # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "IdeogramTextToImage": IdeogramTextToImage, - "MinimaxVideoNode": MinimaxVideoNode } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "IdeogramTextToImage": "Ideogram Text to Image", - "MinimaxVideoNode": "Minimax Video Generator" }