diff --git a/comfy_api_nodes/nodes_moonvalley.py b/comfy_api_nodes/nodes_moonvalley.py index 426875d32..7c31d95b3 100644 --- a/comfy_api_nodes/nodes_moonvalley.py +++ b/comfy_api_nodes/nodes_moonvalley.py @@ -1,32 +1,31 @@ import logging from typing import Optional + import torch from typing_extensions import override +from comfy_api.input import VideoInput +from comfy_api.latest import IO, ComfyExtension from comfy_api_nodes.apis import ( - MoonvalleyTextToVideoRequest, + MoonvalleyPromptResponse, MoonvalleyTextToVideoInferenceParams, + MoonvalleyTextToVideoRequest, MoonvalleyVideoToVideoInferenceParams, MoonvalleyVideoToVideoRequest, - MoonvalleyPromptResponse, ) from comfy_api_nodes.util import ( + ApiEndpoint, + download_url_to_video_output, + poll_op, + sync_op, + trim_video, + upload_images_to_comfyapi, + upload_video_to_comfyapi, validate_container_format_is_mp4, validate_image_dimensions, - download_url_to_video_output, - upload_video_to_comfyapi, - upload_images_to_comfyapi, - sync_op, - ApiEndpoint, - poll_op, validate_string, - trim_video, ) -from comfy_api.input import VideoInput -from comfy_api.latest import ComfyExtension, IO - - API_UPLOADS_ENDPOINT = "/proxy/moonvalley/uploads" API_PROMPTS_ENDPOINT = "/proxy/moonvalley/prompts" API_VIDEO2VIDEO_ENDPOINT = "/proxy/moonvalley/prompts/video-to-video" @@ -103,12 +102,8 @@ def _validate_video_dimensions(width: int, height: int) -> None: } if (width, height) not in supported_resolutions: - supported_list = ", ".join( - [f"{w}x{h}" for w, h in sorted(supported_resolutions)] - ) - raise ValueError( - f"Resolution {width}x{height} not supported. Supported: {supported_list}" - ) + supported_list = ", ".join([f"{w}x{h}" for w, h in sorted(supported_resolutions)]) + raise ValueError(f"Resolution {width}x{height} not supported. Supported: {supported_list}") def _validate_and_trim_duration(video: VideoInput) -> VideoInput: @@ -160,6 +155,8 @@ async def get_response(cls: type[IO.ComfyNode], task_id: str) -> MoonvalleyPromp ApiEndpoint(path=f"{API_PROMPTS_ENDPOINT}/{task_id}"), response_model=MoonvalleyPromptResponse, status_extractor=lambda r: (r.status if r and r.status else None), + poll_interval=16.0, + max_poll_attempts=240, ) @@ -269,7 +266,7 @@ class MoonvalleyImg2VideoNode(IO.ComfyNode): # Get MIME type from tensor - assuming PNG format for image tensors mime_type = "image/png" - image_url = (await upload_images_to_comfyapi(cls, image, max_images=1, mime_type=mime_type))[0] + image_url = (await upload_images_to_comfyapi(cls, image, max_images=1, mime_type=mime_type))[0] task_creation_response = await sync_op( cls, endpoint=ApiEndpoint(path=API_IMG2VIDEO_ENDPOINT, method="POST"), diff --git a/comfy_api_nodes/nodes_runway.py b/comfy_api_nodes/nodes_runway.py index aac69167e..0543d1d0e 100644 --- a/comfy_api_nodes/nodes_runway.py +++ b/comfy_api_nodes/nodes_runway.py @@ -21,7 +21,6 @@ from comfy_api_nodes.apis import ( RunwayImageToVideoRequest, RunwayImageToVideoResponse, RunwayTaskStatusResponse as TaskStatusResponse, - RunwayTaskStatusEnum as TaskStatus, RunwayModelEnum as Model, RunwayDurationEnum as Duration, RunwayAspectRatioEnum as AspectRatio, @@ -33,7 +32,18 @@ from comfy_api_nodes.apis import ( ReferenceImage, RunwayTextToImageAspectRatioEnum, ) -from comfy_api_nodes.util import image_tensor_pair_to_batch, validate_string, validate_image_dimensions, validate_image_aspect_ratio, upload_images_to_comfyapi, download_url_to_video_output, download_url_to_image_tensor, ApiEndpoint, sync_op, poll_op +from comfy_api_nodes.util import ( + image_tensor_pair_to_batch, + validate_string, + validate_image_dimensions, + validate_image_aspect_ratio, + upload_images_to_comfyapi, + download_url_to_video_output, + download_url_to_image_tensor, + ApiEndpoint, + sync_op, + poll_op, +) from comfy_api.input_impl import VideoFromFile from comfy_api.latest import ComfyExtension, IO @@ -93,20 +103,12 @@ def get_image_url_from_task_status(response: TaskStatusResponse) -> Union[str, N async def get_response( - cls: type[IO.ComfyNode], - task_id: str, estimated_duration: Optional[int] = None + cls: type[IO.ComfyNode], task_id: str, estimated_duration: Optional[int] = None ) -> TaskStatusResponse: """Poll the task status until it is finished then get the response.""" return await poll_op( cls, ApiEndpoint(path=f"{PATH_GET_TASK_STATUS}/{task_id}"), - completed_statuses=[ - TaskStatus.SUCCEEDED.value, - ], - failed_statuses=[ - TaskStatus.FAILED.value, - TaskStatus.CANCELLED.value, - ], response_model=TaskStatusResponse, status_extractor=lambda r: r.status.value, estimated_duration=estimated_duration, @@ -143,9 +145,9 @@ class RunwayImageToVideoNodeGen3a(IO.ComfyNode): display_name="Runway Image to Video (Gen3a Turbo)", category="api node/video/Runway", description="Generate a video from a single starting frame using Gen3a Turbo model. " - "Before diving in, review these best practices to ensure that " - "your input selections will set your generation up for success: " - "https://help.runwayml.com/hc/en-us/articles/33927968552339-Creating-with-Act-One-on-Gen-3-Alpha-and-Turbo.", + "Before diving in, review these best practices to ensure that " + "your input selections will set your generation up for success: " + "https://help.runwayml.com/hc/en-us/articles/33927968552339-Creating-with-Act-One-on-Gen-3-Alpha-and-Turbo.", inputs=[ IO.String.Input( "prompt", @@ -217,11 +219,7 @@ class RunwayImageToVideoNodeGen3a(IO.ComfyNode): duration=Duration(duration), ratio=AspectRatio(ratio), promptImage=RunwayPromptImageObject( - root=[ - RunwayPromptImageDetailedObject( - uri=str(download_urls[0]), position="first" - ) - ] + root=[RunwayPromptImageDetailedObject(uri=str(download_urls[0]), position="first")] ), ), ) @@ -237,9 +235,9 @@ class RunwayImageToVideoNodeGen4(IO.ComfyNode): display_name="Runway Image to Video (Gen4 Turbo)", category="api node/video/Runway", description="Generate a video from a single starting frame using Gen4 Turbo model. " - "Before diving in, review these best practices to ensure that " - "your input selections will set your generation up for success: " - "https://help.runwayml.com/hc/en-us/articles/37327109429011-Creating-with-Gen-4-Video.", + "Before diving in, review these best practices to ensure that " + "your input selections will set your generation up for success: " + "https://help.runwayml.com/hc/en-us/articles/37327109429011-Creating-with-Gen-4-Video.", inputs=[ IO.String.Input( "prompt", @@ -311,11 +309,7 @@ class RunwayImageToVideoNodeGen4(IO.ComfyNode): duration=Duration(duration), ratio=AspectRatio(ratio), promptImage=RunwayPromptImageObject( - root=[ - RunwayPromptImageDetailedObject( - uri=str(download_urls[0]), position="first" - ) - ] + root=[RunwayPromptImageDetailedObject(uri=str(download_urls[0]), position="first")] ), ), estimated_duration=AVERAGE_DURATION_FLF_SECONDS, @@ -332,12 +326,12 @@ class RunwayFirstLastFrameNode(IO.ComfyNode): display_name="Runway First-Last-Frame to Video", category="api node/video/Runway", description="Upload first and last keyframes, draft a prompt, and generate a video. " - "More complex transitions, such as cases where the Last frame is completely different " - "from the First frame, may benefit from the longer 10s duration. " - "This would give the generation more time to smoothly transition between the two inputs. " - "Before diving in, review these best practices to ensure that your input selections " - "will set your generation up for success: " - "https://help.runwayml.com/hc/en-us/articles/34170748696595-Creating-with-Keyframes-on-Gen-3.", + "More complex transitions, such as cases where the Last frame is completely different " + "from the First frame, may benefit from the longer 10s duration. " + "This would give the generation more time to smoothly transition between the two inputs. " + "Before diving in, review these best practices to ensure that your input selections " + "will set your generation up for success: " + "https://help.runwayml.com/hc/en-us/articles/34170748696595-Creating-with-Keyframes-on-Gen-3.", inputs=[ IO.String.Input( "prompt", @@ -420,12 +414,8 @@ class RunwayFirstLastFrameNode(IO.ComfyNode): ratio=AspectRatio(ratio), promptImage=RunwayPromptImageObject( root=[ - RunwayPromptImageDetailedObject( - uri=str(download_urls[0]), position="first" - ), - RunwayPromptImageDetailedObject( - uri=str(download_urls[1]), position="last" - ), + RunwayPromptImageDetailedObject(uri=str(download_urls[0]), position="first"), + RunwayPromptImageDetailedObject(uri=str(download_urls[1]), position="last"), ] ), ), @@ -443,7 +433,7 @@ class RunwayTextToImageNode(IO.ComfyNode): display_name="Runway Text to Image", category="api node/image/Runway", description="Generate an image from a text prompt using Runway's Gen 4 model. " - "You can also include reference image to guide the generation.", + "You can also include reference image to guide the generation.", inputs=[ IO.String.Input( "prompt", @@ -527,5 +517,6 @@ class RunwayExtension(ComfyExtension): RunwayTextToImageNode, ] + async def comfy_entrypoint() -> RunwayExtension: return RunwayExtension() diff --git a/comfy_api_nodes/nodes_vidu.py b/comfy_api_nodes/nodes_vidu.py index 9c4d30bc3..0e0572f8c 100644 --- a/comfy_api_nodes/nodes_vidu.py +++ b/comfy_api_nodes/nodes_vidu.py @@ -1,25 +1,24 @@ import logging from enum import Enum -from typing import Optional, Literal, TypeVar -from typing_extensions import override +from typing import Literal, Optional, TypeVar import torch from pydantic import BaseModel, Field +from typing_extensions import override -from comfy_api.latest import ComfyExtension, IO +from comfy_api.latest import IO, ComfyExtension from comfy_api_nodes.util import ( - validate_aspect_ratio_closeness, - validate_image_dimensions, - validate_image_aspect_ratio_range, - get_number_of_images, - download_url_to_video_output, - upload_images_to_comfyapi, ApiEndpoint, - sync_op, + download_url_to_video_output, + get_number_of_images, poll_op, + sync_op, + upload_images_to_comfyapi, + validate_aspect_ratio_closeness, + validate_image_aspect_ratio_range, + validate_image_dimensions, ) - VIDU_TEXT_TO_VIDEO = "/proxy/vidu/text2video" VIDU_IMAGE_TO_VIDEO = "/proxy/vidu/img2video" VIDU_REFERENCE_VIDEO = "/proxy/vidu/reference2video" @@ -28,8 +27,9 @@ VIDU_GET_GENERATION_STATUS = "/proxy/vidu/tasks/%s/creations" R = TypeVar("R") + class VideoModelName(str, Enum): - vidu_q1 = 'viduq1' + vidu_q1 = "viduq1" class AspectRatio(str, Enum): @@ -102,7 +102,7 @@ async def execute_task( ) -> R: response = await sync_op( cls, - endpoint=ApiEndpoint(path=vidu_endpoint,method="POST"), + endpoint=ApiEndpoint(path=vidu_endpoint, method="POST"), response_model=TaskCreationResponse, data=payload, ) @@ -560,5 +560,6 @@ class ViduExtension(ComfyExtension): ViduStartEndToVideoNode, ] + async def comfy_entrypoint() -> ViduExtension: return ViduExtension() diff --git a/comfy_api_nodes/nodes_wan.py b/comfy_api_nodes/nodes_wan.py index 61d50746b..2aab3c2ff 100644 --- a/comfy_api_nodes/nodes_wan.py +++ b/comfy_api_nodes/nodes_wan.py @@ -1,20 +1,22 @@ import re -from typing import Optional, Type, Union -from typing_extensions import override +from typing import Optional import torch from pydantic import BaseModel, Field -from comfy_api.latest import ComfyExtension, Input, IO -from comfy_api_nodes.apis.client import ( +from typing_extensions import override + +from comfy_api.latest import IO, ComfyExtension, Input +from comfy_api_nodes.util import ( ApiEndpoint, - HttpMethod, - SynchronousOperation, - PollingOperation, - EmptyRequest, - R, - T, + audio_to_base64_string, + download_url_to_image_tensor, + download_url_to_video_output, + get_number_of_images, + poll_op, + sync_op, + tensor_to_base64_string, + validate_audio_duration, ) -from comfy_api_nodes.util import get_number_of_images, validate_audio_duration, tensor_to_base64_string, audio_to_base64_string, download_url_to_video_output, download_url_to_image_tensor class Text2ImageInputField(BaseModel): @@ -140,53 +142,7 @@ class VideoTaskStatusResponse(BaseModel): request_id: str = Field(...) -RES_IN_PARENS = re.compile(r'\((\d+)\s*[x×]\s*(\d+)\)') - - -async def process_task( - auth_kwargs: dict[str, str], - url: str, - request_model: Type[T], - response_model: Type[R], - payload: Union[ - Text2ImageTaskCreationRequest, - Image2ImageTaskCreationRequest, - Text2VideoTaskCreationRequest, - Image2VideoTaskCreationRequest, - ], - node_id: str, - estimated_duration: int, - poll_interval: int, -) -> Type[R]: - initial_response = await SynchronousOperation( - endpoint=ApiEndpoint( - path=url, - method=HttpMethod.POST, - request_model=request_model, - response_model=TaskCreationResponse, - ), - request=payload, - auth_kwargs=auth_kwargs, - ).execute() - - if not initial_response.output: - raise Exception(f"Unknown error occurred: {initial_response.code} - {initial_response.message}") - - return await PollingOperation( - poll_endpoint=ApiEndpoint( - path=f"/proxy/wan/api/v1/tasks/{initial_response.output.task_id}", - method=HttpMethod.GET, - request_model=EmptyRequest, - response_model=response_model, - ), - completed_statuses=["SUCCEEDED"], - failed_statuses=["FAILED", "CANCELED", "UNKNOWN"], - status_extractor=lambda x: x.output.task_status, - estimated_duration=estimated_duration, - poll_interval=poll_interval, - node_id=node_id, - auth_kwargs=auth_kwargs, - ).execute() +RES_IN_PARENS = re.compile(r"\((\d+)\s*[x×]\s*(\d+)\)") class WanTextToImageApi(IO.ComfyNode): @@ -253,7 +209,7 @@ class WanTextToImageApi(IO.ComfyNode): IO.Boolean.Input( "watermark", default=True, - tooltip="Whether to add an \"AI generated\" watermark to the result.", + tooltip='Whether to add an "AI generated" watermark to the result.', optional=True, ), ], @@ -280,26 +236,28 @@ class WanTextToImageApi(IO.ComfyNode): prompt_extend: bool = True, watermark: bool = True, ): - payload = Text2ImageTaskCreationRequest( - model=model, - input=Text2ImageInputField(prompt=prompt, negative_prompt=negative_prompt), - parameters=Txt2ImageParametersField( - size=f"{width}*{height}", - seed=seed, - prompt_extend=prompt_extend, - watermark=watermark, + initial_response = await sync_op( + cls, + ApiEndpoint(path="/proxy/wan/api/v1/services/aigc/text2image/image-synthesis", method="POST"), + response_model=TaskCreationResponse, + data=Text2ImageTaskCreationRequest( + model=model, + input=Text2ImageInputField(prompt=prompt, negative_prompt=negative_prompt), + parameters=Txt2ImageParametersField( + size=f"{width}*{height}", + seed=seed, + prompt_extend=prompt_extend, + watermark=watermark, + ), ), ) - response = await process_task( - { - "auth_token": cls.hidden.auth_token_comfy_org, - "comfy_api_key": cls.hidden.api_key_comfy_org, - }, - "/proxy/wan/api/v1/services/aigc/text2image/image-synthesis", - request_model=Text2ImageTaskCreationRequest, + if not initial_response.output: + raise Exception(f"Unknown error occurred: {initial_response.code} - {initial_response.message}") + response = await poll_op( + cls, + ApiEndpoint(path=f"/proxy/wan/api/v1/tasks/{initial_response.output.task_id}"), response_model=ImageTaskStatusResponse, - payload=payload, - node_id=cls.hidden.unique_id, + status_extractor=lambda x: x.output.task_status, estimated_duration=9, poll_interval=3, ) @@ -314,7 +272,7 @@ class WanImageToImageApi(IO.ComfyNode): display_name="Wan Image to Image", category="api node/image/Wan", description="Generates an image from one or two input images and a text prompt. " - "The output image is currently fixed at 1.6 MP; its aspect ratio matches the input image(s).", + "The output image is currently fixed at 1.6 MP; its aspect ratio matches the input image(s).", inputs=[ IO.Combo.Input( "model", @@ -370,7 +328,7 @@ class WanImageToImageApi(IO.ComfyNode): IO.Boolean.Input( "watermark", default=True, - tooltip="Whether to add an \"AI generated\" watermark to the result.", + tooltip='Whether to add an "AI generated" watermark to the result.', optional=True, ), ], @@ -402,28 +360,30 @@ class WanImageToImageApi(IO.ComfyNode): raise ValueError(f"Expected 1 or 2 input images, got {n_images}.") images = [] for i in image: - images.append("data:image/png;base64," + tensor_to_base64_string(i, total_pixels=4096*4096)) - payload = Image2ImageTaskCreationRequest( - model=model, - input=Image2ImageInputField(prompt=prompt, negative_prompt=negative_prompt, images=images), - parameters=Image2ImageParametersField( - # size=f"{width}*{height}", - seed=seed, - watermark=watermark, + images.append("data:image/png;base64," + tensor_to_base64_string(i, total_pixels=4096 * 4096)) + initial_response = await sync_op( + cls, + ApiEndpoint(path="/proxy/wan/api/v1/services/aigc/image2image/image-synthesis", method="POST"), + response_model=TaskCreationResponse, + data=Image2ImageTaskCreationRequest( + model=model, + input=Image2ImageInputField(prompt=prompt, negative_prompt=negative_prompt, images=images), + parameters=Image2ImageParametersField( + # size=f"{width}*{height}", + seed=seed, + watermark=watermark, + ), ), ) - response = await process_task( - { - "auth_token": cls.hidden.auth_token_comfy_org, - "comfy_api_key": cls.hidden.api_key_comfy_org, - }, - "/proxy/wan/api/v1/services/aigc/image2image/image-synthesis", - request_model=Image2ImageTaskCreationRequest, + if not initial_response.output: + raise Exception(f"Unknown error occurred: {initial_response.code} - {initial_response.message}") + response = await poll_op( + cls, + ApiEndpoint(path=f"/proxy/wan/api/v1/tasks/{initial_response.output.task_id}"), response_model=ImageTaskStatusResponse, - payload=payload, - node_id=cls.hidden.unique_id, + status_extractor=lambda x: x.output.task_status, estimated_duration=42, - poll_interval=3, + poll_interval=4, ) return IO.NodeOutput(await download_url_to_image_tensor(str(response.output.results[0].url))) @@ -517,7 +477,7 @@ class WanTextToVideoApi(IO.ComfyNode): IO.Boolean.Input( "watermark", default=True, - tooltip="Whether to add an \"AI generated\" watermark to the result.", + tooltip='Whether to add an "AI generated" watermark to the result.', optional=True, ), ], @@ -551,28 +511,31 @@ class WanTextToVideoApi(IO.ComfyNode): if audio is not None: validate_audio_duration(audio, 3.0, 29.0) audio_url = "data:audio/mp3;base64," + audio_to_base64_string(audio, "mp3", "libmp3lame") - payload = Text2VideoTaskCreationRequest( - model=model, - input=Text2VideoInputField(prompt=prompt, negative_prompt=negative_prompt, audio_url=audio_url), - parameters=Text2VideoParametersField( - size=f"{width}*{height}", - duration=duration, - seed=seed, - audio=generate_audio, - prompt_extend=prompt_extend, - watermark=watermark, + + initial_response = await sync_op( + cls, + ApiEndpoint(path="/proxy/wan/api/v1/services/aigc/video-generation/video-synthesis", method="POST"), + response_model=TaskCreationResponse, + data=Text2VideoTaskCreationRequest( + model=model, + input=Text2VideoInputField(prompt=prompt, negative_prompt=negative_prompt, audio_url=audio_url), + parameters=Text2VideoParametersField( + size=f"{width}*{height}", + duration=duration, + seed=seed, + audio=generate_audio, + prompt_extend=prompt_extend, + watermark=watermark, + ), ), ) - response = await process_task( - { - "auth_token": cls.hidden.auth_token_comfy_org, - "comfy_api_key": cls.hidden.api_key_comfy_org, - }, - "/proxy/wan/api/v1/services/aigc/video-generation/video-synthesis", - request_model=Text2VideoTaskCreationRequest, + if not initial_response.output: + raise Exception(f"Unknown error occurred: {initial_response.code} - {initial_response.message}") + response = await poll_op( + cls, + ApiEndpoint(path=f"/proxy/wan/api/v1/tasks/{initial_response.output.task_id}"), response_model=VideoTaskStatusResponse, - payload=payload, - node_id=cls.hidden.unique_id, + status_extractor=lambda x: x.output.task_status, estimated_duration=120 * int(duration / 5), poll_interval=6, ) @@ -661,7 +624,7 @@ class WanImageToVideoApi(IO.ComfyNode): IO.Boolean.Input( "watermark", default=True, - tooltip="Whether to add an \"AI generated\" watermark to the result.", + tooltip='Whether to add an "AI generated" watermark to the result.', optional=True, ), ], @@ -693,35 +656,37 @@ class WanImageToVideoApi(IO.ComfyNode): ): if get_number_of_images(image) != 1: raise ValueError("Exactly one input image is required.") - image_url = "data:image/png;base64," + tensor_to_base64_string(image, total_pixels=2000*2000) + image_url = "data:image/png;base64," + tensor_to_base64_string(image, total_pixels=2000 * 2000) audio_url = None if audio is not None: validate_audio_duration(audio, 3.0, 29.0) audio_url = "data:audio/mp3;base64," + audio_to_base64_string(audio, "mp3", "libmp3lame") - payload = Image2VideoTaskCreationRequest( - model=model, - input=Image2VideoInputField( - prompt=prompt, negative_prompt=negative_prompt, img_url=image_url, audio_url=audio_url - ), - parameters=Image2VideoParametersField( - resolution=resolution, - duration=duration, - seed=seed, - audio=generate_audio, - prompt_extend=prompt_extend, - watermark=watermark, + initial_response = await sync_op( + cls, + ApiEndpoint(path="/proxy/wan/api/v1/services/aigc/video-generation/video-synthesis", method="POST"), + response_model=TaskCreationResponse, + data=Image2VideoTaskCreationRequest( + model=model, + input=Image2VideoInputField( + prompt=prompt, negative_prompt=negative_prompt, img_url=image_url, audio_url=audio_url + ), + parameters=Image2VideoParametersField( + resolution=resolution, + duration=duration, + seed=seed, + audio=generate_audio, + prompt_extend=prompt_extend, + watermark=watermark, + ), ), ) - response = await process_task( - { - "auth_token": cls.hidden.auth_token_comfy_org, - "comfy_api_key": cls.hidden.api_key_comfy_org, - }, - "/proxy/wan/api/v1/services/aigc/video-generation/video-synthesis", - request_model=Image2VideoTaskCreationRequest, + if not initial_response.output: + raise Exception(f"Unknown error occurred: {initial_response.code} - {initial_response.message}") + response = await poll_op( + cls, + ApiEndpoint(path=f"/proxy/wan/api/v1/tasks/{initial_response.output.task_id}"), response_model=VideoTaskStatusResponse, - payload=payload, - node_id=cls.hidden.unique_id, + status_extractor=lambda x: x.output.task_status, estimated_duration=120 * int(duration / 5), poll_interval=6, ) diff --git a/comfy_api_nodes/util/_helpers.py b/comfy_api_nodes/util/_helpers.py index bf6c32d49..328fe5227 100644 --- a/comfy_api_nodes/util/_helpers.py +++ b/comfy_api_nodes/util/_helpers.py @@ -35,7 +35,7 @@ def default_base_url() -> str: async def sleep_with_interrupt( seconds: float, - node_cls: type[IO.ComfyNode], + node_cls: Optional[type[IO.ComfyNode]], label: Optional[str] = None, start_ts: Optional[float] = None, estimated_total: Optional[int] = None, diff --git a/comfy_api_nodes/util/client.py b/comfy_api_nodes/util/client.py index 184e23824..d6986a79c 100644 --- a/comfy_api_nodes/util/client.py +++ b/comfy_api_nodes/util/client.py @@ -8,7 +8,7 @@ import uuid from dataclasses import dataclass from enum import Enum from io import BytesIO -from typing import Any, Callable, Literal, Optional, Type, TypeVar, Union +from typing import Any, Callable, Iterable, Literal, Optional, Type, TypeVar, Union from urllib.parse import urljoin, urlparse import aiohttp @@ -79,7 +79,7 @@ class _PollUIState: _RETRY_STATUS = {408, 429, 500, 502, 503, 504} COMPLETED_STATUSES = ["succeeded", "succeed", "success", "completed"] -FAILED_STATUSES = ["cancelled", "failed", "error"] +FAILED_STATUSES = ["cancelled", "canceled", "failed", "error"] QUEUED_STATUSES = ["created", "queued", "queueing", "submitted"] @@ -97,7 +97,7 @@ async def sync_op( retry_delay: float = 1.0, retry_backoff: float = 2.0, wait_label: str = "Waiting for server", - estimated_total: Optional[int] = None, + estimated_duration: Optional[int] = None, final_label_on_success: Optional[str] = "Completed", progress_origin_ts: Optional[float] = None, monitor_progress: bool = True, @@ -114,7 +114,7 @@ async def sync_op( retry_delay=retry_delay, retry_backoff=retry_backoff, wait_label=wait_label, - estimated_total=estimated_total, + estimated_duration=estimated_duration, as_binary=False, final_label_on_success=final_label_on_success, progress_origin_ts=progress_origin_ts, @@ -130,7 +130,7 @@ async def poll_op( poll_endpoint: ApiEndpoint, *, response_model: Type[M], - status_extractor: Callable[[M], Optional[str]], + status_extractor: Callable[[M], Optional[Union[str, int]]], progress_extractor: Optional[Callable[[M], Optional[int]]] = None, price_extractor: Optional[Callable[[M], Optional[float]]] = None, completed_statuses: Optional[list[Union[str, int]]] = None, @@ -183,7 +183,7 @@ async def sync_op_raw( retry_delay: float = 1.0, retry_backoff: float = 2.0, wait_label: str = "Waiting for server", - estimated_total: Optional[int] = None, + estimated_duration: Optional[int] = None, as_binary: bool = False, final_label_on_success: Optional[str] = "Completed", progress_origin_ts: Optional[float] = None, @@ -212,7 +212,7 @@ async def sync_op_raw( retry_backoff=retry_backoff, wait_label=wait_label, monitor_progress=monitor_progress, - estimated_total=estimated_total, + estimated_total=estimated_duration, final_label_on_success=final_label_on_success, progress_origin_ts=progress_origin_ts, ) @@ -223,7 +223,7 @@ async def poll_op_raw( cls: type[IO.ComfyNode], poll_endpoint: ApiEndpoint, *, - status_extractor: Callable[[dict[str, Any]], Optional[str]], + status_extractor: Callable[[dict[str, Any]], Optional[Union[str, int]]], progress_extractor: Optional[Callable[[dict[str, Any]], Optional[int]]] = None, price_extractor: Optional[Callable[[dict[str, Any]], Optional[float]]] = None, completed_statuses: Optional[list[Union[str, int]]] = None, @@ -247,9 +247,9 @@ async def poll_op_raw( Returns the final JSON response from the poll endpoint. """ - completed_states = COMPLETED_STATUSES if completed_statuses is None else completed_statuses - failed_states = FAILED_STATUSES if failed_statuses is None else failed_statuses - queued_states = QUEUED_STATUSES if queued_statuses is None else queued_statuses + completed_states = _normalize_statuses(COMPLETED_STATUSES if completed_statuses is None else completed_statuses) + failed_states = _normalize_statuses(FAILED_STATUSES if failed_statuses is None else failed_statuses) + queued_states = _normalize_statuses(QUEUED_STATUSES if queued_statuses is None else queued_statuses) started = time.monotonic() consumed_attempts = 0 # counts only non-queued polls @@ -271,7 +271,7 @@ async def poll_op_raw( ) _display_time_progress( cls, - label=state.status_label, + status=state.status_label, elapsed_seconds=int(now - state.started), estimated_total=state.estimated_duration, price=state.price, @@ -294,7 +294,7 @@ async def poll_op_raw( retry_delay=retry_delay_per_poll, retry_backoff=retry_backoff_per_poll, wait_label="Checking", - estimated_total=None, + estimated_duration=None, as_binary=False, final_label_on_success=None, monitor_progress=False, @@ -310,7 +310,7 @@ async def poll_op_raw( timeout=cancel_timeout, max_retries=0, wait_label="Cancelling task", - estimated_total=None, + estimated_duration=None, as_binary=False, final_label_on_success=None, monitor_progress=False, @@ -318,7 +318,7 @@ async def poll_op_raw( raise try: - status = status_extractor(resp_json) + status = _normalize_status_value(status_extractor(resp_json)) except Exception as e: logging.error("Status extraction failed: %s", e) status = None @@ -360,7 +360,7 @@ async def poll_op_raw( _display_time_progress( cls, - label=status if status else "Completed", + status=status if status else "Completed", elapsed_seconds=int(now_ts - started), estimated_total=estimated_duration, price=state.price, @@ -385,7 +385,7 @@ async def poll_op_raw( timeout=cancel_timeout, max_retries=0, wait_label="Cancelling task", - estimated_total=None, + estimated_duration=None, as_binary=False, final_label_on_success=None, monitor_progress=False, @@ -414,12 +414,12 @@ def _display_text( node_cls: type[IO.ComfyNode], text: Optional[str], *, - status: Optional[str] = None, + status: Optional[Union[str, int]] = None, price: Optional[float] = None, ) -> None: display_lines: list[str] = [] if status: - display_lines.append(f"Status: {status.capitalize()}") + display_lines.append(f"Status: {status.capitalize() if isinstance(status, str) else status}") if price is not None: display_lines.append(f"Price: ${float(price):,.4f}") if text is not None: @@ -430,7 +430,7 @@ def _display_text( def _display_time_progress( node_cls: type[IO.ComfyNode], - label: str, + status: Optional[Union[str, int]], elapsed_seconds: int, estimated_total: Optional[int] = None, *, @@ -444,7 +444,7 @@ def _display_time_progress( time_line = f"Time elapsed: {int(elapsed_seconds)}s (~{remaining}s remaining)" else: time_line = f"Time elapsed: {int(elapsed_seconds)}s" - _display_text(node_cls, time_line, status=label, price=price) + _display_text(node_cls, time_line, status=status, price=price) async def _diagnose_connectivity() -> dict[str, bool]: @@ -640,7 +640,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool): with contextlib.suppress(Exception): file_value.seek(0) form.add_field(field_name, file_value, filename=filename, content_type=content_type) - payload_kw["data"] = form # do not send body on GET + payload_kw["data"] = form elif cfg.content_type == "application/x-www-form-urlencoded" and method != "GET": payload_headers["Content-Type"] = "application/x-www-form-urlencoded" payload_kw["data"] = cfg.data or {} @@ -660,7 +660,6 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool): except Exception as _log_e: logging.debug("[DEBUG] request logging failed: %s", _log_e) - # Compose the HTTP request coroutine req_coro = sess.request(method, url, params=params, **payload_kw) req_task = asyncio.create_task(req_coro) @@ -684,7 +683,6 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool): body = await resp.json() except (ContentTypeError, json.JSONDecodeError): body = await resp.text() - # Retryable? if resp.status in _RETRY_STATUS and attempt <= cfg.max_retries: logging.warning( "HTTP %s %s -> %s. Retrying in %.2fs (retry %d of %d).", @@ -733,9 +731,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool): logging.debug("[DEBUG] response logging failed: %s", _log_e) raise Exception(msg) - # Success if expect_binary: - # Read stream in chunks so that cancellation is fast when user interrupts buff = bytearray() last_tick = time.monotonic() async for chunk in resp.content.iter_chunked(64 * 1024): @@ -794,7 +790,6 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool): logging.debug("Polling was interrupted by user") raise except (ClientError, asyncio.TimeoutError, socket.gaierror) as e: - # Retry transient connection issues if attempt <= cfg.max_retries: logging.warning( "Connection error calling %s %s. Retrying in %.2fs (%d/%d): %s", @@ -873,7 +868,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool): if operation_succeeded and cfg.monitor_progress and cfg.final_label_on_success: _display_time_progress( cfg.node_cls, - label=cfg.final_label_on_success, + status=cfg.final_label_on_success, elapsed_seconds=( final_elapsed_seconds if final_elapsed_seconds is not None @@ -926,3 +921,20 @@ def _wrap_model_extractor( raise return _wrapped + + +def _normalize_statuses(values: Optional[Iterable[Union[str, int]]]) -> set[Union[str, int]]: + if not values: + return set() + out: set[Union[str, int]] = set() + for v in values: + nv = _normalize_status_value(v) + if nv is not None: + out.add(nv) + return out + + +def _normalize_status_value(val: Union[str, int, None]) -> Union[str, int, None]: + if isinstance(val, str): + return val.strip().lower() + return val diff --git a/comfy_api_nodes/util/conversions.py b/comfy_api_nodes/util/conversions.py index 3cf81459e..10cd1051b 100644 --- a/comfy_api_nodes/util/conversions.py +++ b/comfy_api_nodes/util/conversions.py @@ -11,9 +11,7 @@ import torch from PIL import Image from comfy.utils import common_upscale -from comfy_api.input import VideoInput -from comfy_api.input.basic_types import AudioInput -from comfy_api.latest import InputImpl +from comfy_api.latest import Input, InputImpl from ._helpers import mimetype_to_extension @@ -165,7 +163,7 @@ def tensor_to_data_uri( return f"data:{mime_type};base64,{base64_string}" -def audio_to_base64_string(audio: AudioInput, container_format: str = "mp4", codec_name: str = "aac") -> str: +def audio_to_base64_string(audio: Input.Audio, container_format: str = "mp4", codec_name: str = "aac") -> str: """Converts an audio input to a base64 string.""" sample_rate: int = audio["sample_rate"] waveform: torch.Tensor = audio["waveform"] @@ -232,7 +230,7 @@ def audio_tensor_to_contiguous_ndarray(waveform: torch.Tensor) -> np.ndarray: return audio_data_np -def audio_input_to_mp3(audio: AudioInput) -> BytesIO: +def audio_input_to_mp3(audio: Input.Audio) -> BytesIO: waveform = audio["waveform"].cpu() output_buffer = BytesIO() @@ -255,7 +253,7 @@ def audio_input_to_mp3(audio: AudioInput) -> BytesIO: return output_buffer -def trim_video(video: VideoInput, duration_sec: float) -> VideoInput: +def trim_video(video: Input.Video, duration_sec: float) -> Input.Video: """ Returns a new VideoInput object trimmed from the beginning to the specified duration, using av to avoid loading entire video into memory. diff --git a/comfy_api_nodes/util/download_helpers.py b/comfy_api_nodes/util/download_helpers.py index bffe9bc2f..d0b65424b 100644 --- a/comfy_api_nodes/util/download_helpers.py +++ b/comfy_api_nodes/util/download_helpers.py @@ -1,7 +1,5 @@ import asyncio import contextlib -import logging -import time import uuid from io import BytesIO from pathlib import Path @@ -13,10 +11,15 @@ import torch from aiohttp.client_exceptions import ClientError, ContentTypeError from comfy_api.input_impl import VideoFromFile -from comfy_api.latest import IO as ComfyIO +from comfy_api.latest import IO as COMFY_IO from comfy_api_nodes.apis import request_logger -from ._helpers import default_base_url, get_auth_header, is_processing_interrupted +from ._helpers import ( + default_base_url, + get_auth_header, + is_processing_interrupted, + sleep_with_interrupt, +) from .client import _diagnose_connectivity from .common_exceptions import ApiServerError, LocalNetworkError, ProcessingInterrupted from .conversions import bytesio_to_image_tensor @@ -32,7 +35,7 @@ async def download_url_to_bytesio( max_retries: int = 3, retry_delay: float = 1.0, retry_backoff: float = 2.0, - cls: type[ComfyIO.ComfyNode] = None, + cls: type[COMFY_IO.ComfyNode] = None, ) -> None: """Stream-download a URL to `dest`. @@ -52,7 +55,8 @@ async def download_url_to_bytesio( attempt = 0 delay = retry_delay - headers = {} + headers: dict[str, str] = {} + if url.startswith("/proxy/"): if cls is None: raise ValueError("For relative 'cloud' paths, the `cls` parameter is required.") @@ -66,69 +70,112 @@ async def download_url_to_bytesio( is_path_sink = isinstance(dest, (str, Path)) fhandle = None + session: Optional[aiohttp.ClientSession] = None + stop_evt: Optional[asyncio.Event] = None + monitor_task: Optional[asyncio.Task] = None + req_task: Optional[asyncio.Task] = None + try: with contextlib.suppress(Exception): request_logger.log_request_response(operation_id=op_id, request_method="GET", request_url=url) - async with aiohttp.ClientSession(timeout=timeout_cfg) as session: - async with session.get(url, headers=headers) as resp: - if resp.status >= 400: - with contextlib.suppress(Exception): - try: - body = await resp.json() - except (ContentTypeError, ValueError): - text = await resp.text() - body = text if len(text) <= 4096 else f"[text {len(text)} bytes]" - request_logger.log_request_response( - operation_id=op_id, - request_method="GET", - request_url=url, - response_status_code=resp.status, - response_headers=dict(resp.headers), - response_content=body, - error_message=f"HTTP {resp.status}", - ) + session = aiohttp.ClientSession(timeout=timeout_cfg) + stop_evt = asyncio.Event() - if resp.status in _RETRY_STATUS and attempt <= max_retries: - await _sleep_with_cancel(delay) - delay *= retry_backoff - continue - raise Exception(f"Failed to download (HTTP {resp.status}).") - - if is_path_sink: - p = Path(str(dest)) - with contextlib.suppress(Exception): - p.parent.mkdir(parents=True, exist_ok=True) - fhandle = open(p, "wb") - sink = fhandle - else: - sink = dest # BytesIO or file-like - - written = 0 - async for chunk in resp.content.iter_chunked(1024 * 1024): - sink.write(chunk) - written += len(chunk) + async def _monitor(): + try: + while not stop_evt.is_set(): if is_processing_interrupted(): - raise ProcessingInterrupted("Task cancelled") + return + await asyncio.sleep(1.0) + except asyncio.CancelledError: + return - if isinstance(dest, BytesIO): - with contextlib.suppress(Exception): - dest.seek(0) + monitor_task = asyncio.create_task(_monitor()) + req_task = asyncio.create_task(session.get(url, headers=headers)) + done, pending = await asyncio.wait({req_task, monitor_task}, return_when=asyncio.FIRST_COMPLETED) + + if monitor_task in done and req_task in pending: + req_task.cancel() + with contextlib.suppress(Exception): + await req_task + raise ProcessingInterrupted("Task cancelled") + + try: + resp = await req_task + except asyncio.CancelledError: + raise ProcessingInterrupted("Task cancelled") from None + + async with resp: + if resp.status >= 400: with contextlib.suppress(Exception): + try: + body = await resp.json() + except (ContentTypeError, ValueError): + text = await resp.text() + body = text if len(text) <= 4096 else f"[text {len(text)} bytes]" request_logger.log_request_response( operation_id=op_id, request_method="GET", request_url=url, response_status_code=resp.status, response_headers=dict(resp.headers), - response_content=f"[streamed {written} bytes to dest]", + response_content=body, + error_message=f"HTTP {resp.status}", ) - return - except ProcessingInterrupted: - logging.debug("Download was interrupted by user") - raise + if resp.status in _RETRY_STATUS and attempt <= max_retries: + await sleep_with_interrupt(delay, cls, None, None, None) + delay *= retry_backoff + continue + raise Exception(f"Failed to download (HTTP {resp.status}).") + + if is_path_sink: + p = Path(str(dest)) + with contextlib.suppress(Exception): + p.parent.mkdir(parents=True, exist_ok=True) + fhandle = open(p, "wb") + sink = fhandle + else: + sink = dest # BytesIO or file-like + + written = 0 + while True: + try: + chunk = await asyncio.wait_for(resp.content.read(1024 * 1024), timeout=1.0) + except asyncio.TimeoutError: + chunk = b"" + except asyncio.CancelledError: + raise ProcessingInterrupted("Task cancelled") from None + + if is_processing_interrupted(): + raise ProcessingInterrupted("Task cancelled") + + if not chunk: + if resp.content.at_eof(): + break + continue + + sink.write(chunk) + written += len(chunk) + + if isinstance(dest, BytesIO): + with contextlib.suppress(Exception): + dest.seek(0) + + with contextlib.suppress(Exception): + request_logger.log_request_response( + operation_id=op_id, + request_method="GET", + request_url=url, + response_status_code=resp.status, + response_headers=dict(resp.headers), + response_content=f"[streamed {written} bytes to dest]", + ) + return + except asyncio.CancelledError: + raise ProcessingInterrupted("Task cancelled") from None except (ClientError, asyncio.TimeoutError) as e: if attempt <= max_retries: with contextlib.suppress(Exception): @@ -138,7 +185,7 @@ async def download_url_to_bytesio( request_url=url, error_message=f"{type(e).__name__}: {str(e)} (will retry)", ) - await _sleep_with_cancel(delay) + await sleep_with_interrupt(delay, cls, None, None, None) delay *= retry_backoff continue @@ -149,8 +196,21 @@ async def download_url_to_bytesio( ) from e raise ApiServerError("The remote service appears unreachable at this time.") from e finally: - with contextlib.suppress(Exception): - if fhandle: + if stop_evt is not None: + stop_evt.set() + if monitor_task: + monitor_task.cancel() + with contextlib.suppress(Exception): + await monitor_task + if req_task and not req_task.done(): + req_task.cancel() + with contextlib.suppress(Exception): + await req_task + if session: + with contextlib.suppress(Exception): + await session.close() + if fhandle: + with contextlib.suppress(Exception): fhandle.flush() fhandle.close() @@ -159,7 +219,7 @@ async def download_url_to_image_tensor( url: str, *, timeout: float = None, - cls: type[ComfyIO.ComfyNode] = None, + cls: type[COMFY_IO.ComfyNode] = None, ) -> torch.Tensor: """Downloads an image from a URL and returns a [B, H, W, C] tensor.""" result = BytesIO() @@ -171,7 +231,7 @@ async def download_url_to_video_output( video_url: str, *, timeout: float = None, - cls: type[ComfyIO.ComfyNode] = None, + cls: type[COMFY_IO.ComfyNode] = None, ) -> VideoFromFile: """Downloads a video from a URL and returns a `VIDEO` output.""" result = BytesIO() @@ -186,15 +246,3 @@ def _generate_operation_id(method: str, url: str, attempt: int) -> str: except Exception: slug = "download" return f"{method}_{slug}_try{attempt}_{uuid.uuid4().hex[:8]}" - - -async def _sleep_with_cancel(seconds: float) -> None: - """Sleep in 1s slices while checking for interruption.""" - end = time.monotonic() + seconds - while True: - if is_processing_interrupted(): - raise ProcessingInterrupted("Task cancelled") - now = time.monotonic() - if now >= end: - return - await asyncio.sleep(min(1.0, end - now)) diff --git a/comfy_api_nodes/util/upload_helpers.py b/comfy_api_nodes/util/upload_helpers.py index fa9d10b8e..a345d451d 100644 --- a/comfy_api_nodes/util/upload_helpers.py +++ b/comfy_api_nodes/util/upload_helpers.py @@ -11,9 +11,7 @@ import aiohttp import torch from pydantic import BaseModel, Field -from comfy_api.input.basic_types import AudioInput -from comfy_api.input.video_types import VideoInput -from comfy_api.latest import IO +from comfy_api.latest import IO, Input from comfy_api.util import VideoCodec, VideoContainer from comfy_api_nodes.apis import request_logger @@ -72,7 +70,7 @@ async def upload_images_to_comfyapi( async def upload_audio_to_comfyapi( cls: type[IO.ComfyNode], - audio: AudioInput, + audio: Input.Audio, *, container_format: str = "mp4", codec_name: str = "aac", @@ -92,7 +90,7 @@ async def upload_audio_to_comfyapi( async def upload_video_to_comfyapi( cls: type[IO.ComfyNode], - video: VideoInput, + video: Input.Video, *, container: VideoContainer = VideoContainer.MP4, codec: VideoCodec = VideoCodec.H264, @@ -104,8 +102,8 @@ async def upload_video_to_comfyapi( """ if max_duration is not None: try: - actual_duration = video.duration_seconds - if actual_duration is not None and actual_duration > max_duration: + actual_duration = video.get_duration() + if actual_duration > max_duration: raise ValueError( f"Video duration ({actual_duration:.2f}s) exceeds the maximum allowed ({max_duration}s)." ) @@ -244,7 +242,11 @@ async def upload_file( req_task.cancel() raise ProcessingInterrupted("Upload cancelled") - resp = await req_task + try: + resp = await req_task + except asyncio.CancelledError: + raise ProcessingInterrupted("Upload cancelled") from None + async with resp: if resp.status >= 400: with contextlib.suppress(Exception): @@ -286,6 +288,8 @@ async def upload_file( except Exception as e: logging.debug("[DEBUG] upload response logging failed: %s", e) return + except asyncio.CancelledError: + raise ProcessingInterrupted("Task cancelled") from None except (aiohttp.ClientError, asyncio.TimeoutError) as e: if attempt <= max_retries: with contextlib.suppress(Exception): diff --git a/pyproject.toml b/pyproject.toml index 29023d517..fcbcb3dd9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,12 +69,3 @@ messages_control.disable = [ "no-else-return", "unused-variable", ] - - -[tool.black] -line-length = 120 -preview = true - - -[tool.isort] -profile = "black"