converted WAN nodes to use new client; polishing

This commit is contained in:
bigcat88 2025-10-17 19:51:45 +03:00
parent 50f6a5e10d
commit b34bc7987b
No known key found for this signature in database
GPG Key ID: 1F0BF0EC3CF22721
10 changed files with 340 additions and 333 deletions

View File

@ -1,32 +1,31 @@
import logging import logging
from typing import Optional from typing import Optional
import torch import torch
from typing_extensions import override from typing_extensions import override
from comfy_api.input import VideoInput
from comfy_api.latest import IO, ComfyExtension
from comfy_api_nodes.apis import ( from comfy_api_nodes.apis import (
MoonvalleyTextToVideoRequest, MoonvalleyPromptResponse,
MoonvalleyTextToVideoInferenceParams, MoonvalleyTextToVideoInferenceParams,
MoonvalleyTextToVideoRequest,
MoonvalleyVideoToVideoInferenceParams, MoonvalleyVideoToVideoInferenceParams,
MoonvalleyVideoToVideoRequest, MoonvalleyVideoToVideoRequest,
MoonvalleyPromptResponse,
) )
from comfy_api_nodes.util import ( 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_container_format_is_mp4,
validate_image_dimensions, validate_image_dimensions,
download_url_to_video_output,
upload_video_to_comfyapi,
upload_images_to_comfyapi,
sync_op,
ApiEndpoint,
poll_op,
validate_string, validate_string,
trim_video,
) )
from comfy_api.input import VideoInput
from comfy_api.latest import ComfyExtension, IO
API_UPLOADS_ENDPOINT = "/proxy/moonvalley/uploads" API_UPLOADS_ENDPOINT = "/proxy/moonvalley/uploads"
API_PROMPTS_ENDPOINT = "/proxy/moonvalley/prompts" API_PROMPTS_ENDPOINT = "/proxy/moonvalley/prompts"
API_VIDEO2VIDEO_ENDPOINT = "/proxy/moonvalley/prompts/video-to-video" 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: if (width, height) not in supported_resolutions:
supported_list = ", ".join( supported_list = ", ".join([f"{w}x{h}" for w, h in sorted(supported_resolutions)])
[f"{w}x{h}" for w, h in sorted(supported_resolutions)] raise ValueError(f"Resolution {width}x{height} not supported. Supported: {supported_list}")
)
raise ValueError(
f"Resolution {width}x{height} not supported. Supported: {supported_list}"
)
def _validate_and_trim_duration(video: VideoInput) -> VideoInput: 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}"), ApiEndpoint(path=f"{API_PROMPTS_ENDPOINT}/{task_id}"),
response_model=MoonvalleyPromptResponse, response_model=MoonvalleyPromptResponse,
status_extractor=lambda r: (r.status if r and r.status else None), status_extractor=lambda r: (r.status if r and r.status else None),
poll_interval=16.0,
max_poll_attempts=240,
) )

View File

@ -21,7 +21,6 @@ from comfy_api_nodes.apis import (
RunwayImageToVideoRequest, RunwayImageToVideoRequest,
RunwayImageToVideoResponse, RunwayImageToVideoResponse,
RunwayTaskStatusResponse as TaskStatusResponse, RunwayTaskStatusResponse as TaskStatusResponse,
RunwayTaskStatusEnum as TaskStatus,
RunwayModelEnum as Model, RunwayModelEnum as Model,
RunwayDurationEnum as Duration, RunwayDurationEnum as Duration,
RunwayAspectRatioEnum as AspectRatio, RunwayAspectRatioEnum as AspectRatio,
@ -33,7 +32,18 @@ from comfy_api_nodes.apis import (
ReferenceImage, ReferenceImage,
RunwayTextToImageAspectRatioEnum, 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.input_impl import VideoFromFile
from comfy_api.latest import ComfyExtension, IO 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( async def get_response(
cls: type[IO.ComfyNode], cls: type[IO.ComfyNode], task_id: str, estimated_duration: Optional[int] = None
task_id: str, estimated_duration: Optional[int] = None
) -> TaskStatusResponse: ) -> TaskStatusResponse:
"""Poll the task status until it is finished then get the response.""" """Poll the task status until it is finished then get the response."""
return await poll_op( return await poll_op(
cls, cls,
ApiEndpoint(path=f"{PATH_GET_TASK_STATUS}/{task_id}"), 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, response_model=TaskStatusResponse,
status_extractor=lambda r: r.status.value, status_extractor=lambda r: r.status.value,
estimated_duration=estimated_duration, estimated_duration=estimated_duration,
@ -217,11 +219,7 @@ class RunwayImageToVideoNodeGen3a(IO.ComfyNode):
duration=Duration(duration), duration=Duration(duration),
ratio=AspectRatio(ratio), ratio=AspectRatio(ratio),
promptImage=RunwayPromptImageObject( promptImage=RunwayPromptImageObject(
root=[ root=[RunwayPromptImageDetailedObject(uri=str(download_urls[0]), position="first")]
RunwayPromptImageDetailedObject(
uri=str(download_urls[0]), position="first"
)
]
), ),
), ),
) )
@ -311,11 +309,7 @@ class RunwayImageToVideoNodeGen4(IO.ComfyNode):
duration=Duration(duration), duration=Duration(duration),
ratio=AspectRatio(ratio), ratio=AspectRatio(ratio),
promptImage=RunwayPromptImageObject( promptImage=RunwayPromptImageObject(
root=[ root=[RunwayPromptImageDetailedObject(uri=str(download_urls[0]), position="first")]
RunwayPromptImageDetailedObject(
uri=str(download_urls[0]), position="first"
)
]
), ),
), ),
estimated_duration=AVERAGE_DURATION_FLF_SECONDS, estimated_duration=AVERAGE_DURATION_FLF_SECONDS,
@ -420,12 +414,8 @@ class RunwayFirstLastFrameNode(IO.ComfyNode):
ratio=AspectRatio(ratio), ratio=AspectRatio(ratio),
promptImage=RunwayPromptImageObject( promptImage=RunwayPromptImageObject(
root=[ root=[
RunwayPromptImageDetailedObject( RunwayPromptImageDetailedObject(uri=str(download_urls[0]), position="first"),
uri=str(download_urls[0]), position="first" RunwayPromptImageDetailedObject(uri=str(download_urls[1]), position="last"),
),
RunwayPromptImageDetailedObject(
uri=str(download_urls[1]), position="last"
),
] ]
), ),
), ),
@ -527,5 +517,6 @@ class RunwayExtension(ComfyExtension):
RunwayTextToImageNode, RunwayTextToImageNode,
] ]
async def comfy_entrypoint() -> RunwayExtension: async def comfy_entrypoint() -> RunwayExtension:
return RunwayExtension() return RunwayExtension()

View File

@ -1,25 +1,24 @@
import logging import logging
from enum import Enum from enum import Enum
from typing import Optional, Literal, TypeVar from typing import Literal, Optional, TypeVar
from typing_extensions import override
import torch import torch
from pydantic import BaseModel, Field 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 ( 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, ApiEndpoint,
sync_op, download_url_to_video_output,
get_number_of_images,
poll_op, 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_TEXT_TO_VIDEO = "/proxy/vidu/text2video"
VIDU_IMAGE_TO_VIDEO = "/proxy/vidu/img2video" VIDU_IMAGE_TO_VIDEO = "/proxy/vidu/img2video"
VIDU_REFERENCE_VIDEO = "/proxy/vidu/reference2video" VIDU_REFERENCE_VIDEO = "/proxy/vidu/reference2video"
@ -28,8 +27,9 @@ VIDU_GET_GENERATION_STATUS = "/proxy/vidu/tasks/%s/creations"
R = TypeVar("R") R = TypeVar("R")
class VideoModelName(str, Enum): class VideoModelName(str, Enum):
vidu_q1 = 'viduq1' vidu_q1 = "viduq1"
class AspectRatio(str, Enum): class AspectRatio(str, Enum):
@ -102,7 +102,7 @@ async def execute_task(
) -> R: ) -> R:
response = await sync_op( response = await sync_op(
cls, cls,
endpoint=ApiEndpoint(path=vidu_endpoint,method="POST"), endpoint=ApiEndpoint(path=vidu_endpoint, method="POST"),
response_model=TaskCreationResponse, response_model=TaskCreationResponse,
data=payload, data=payload,
) )
@ -560,5 +560,6 @@ class ViduExtension(ComfyExtension):
ViduStartEndToVideoNode, ViduStartEndToVideoNode,
] ]
async def comfy_entrypoint() -> ViduExtension: async def comfy_entrypoint() -> ViduExtension:
return ViduExtension() return ViduExtension()

View File

@ -1,20 +1,22 @@
import re import re
from typing import Optional, Type, Union from typing import Optional
from typing_extensions import override
import torch import torch
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from comfy_api.latest import ComfyExtension, Input, IO from typing_extensions import override
from comfy_api_nodes.apis.client import (
from comfy_api.latest import IO, ComfyExtension, Input
from comfy_api_nodes.util import (
ApiEndpoint, ApiEndpoint,
HttpMethod, audio_to_base64_string,
SynchronousOperation, download_url_to_image_tensor,
PollingOperation, download_url_to_video_output,
EmptyRequest, get_number_of_images,
R, poll_op,
T, 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): class Text2ImageInputField(BaseModel):
@ -140,53 +142,7 @@ class VideoTaskStatusResponse(BaseModel):
request_id: str = Field(...) request_id: str = Field(...)
RES_IN_PARENS = re.compile(r'\((\d+)\s*[x×]\s*(\d+)\)') 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()
class WanTextToImageApi(IO.ComfyNode): class WanTextToImageApi(IO.ComfyNode):
@ -253,7 +209,7 @@ class WanTextToImageApi(IO.ComfyNode):
IO.Boolean.Input( IO.Boolean.Input(
"watermark", "watermark",
default=True, 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, optional=True,
), ),
], ],
@ -280,7 +236,11 @@ class WanTextToImageApi(IO.ComfyNode):
prompt_extend: bool = True, prompt_extend: bool = True,
watermark: bool = True, watermark: bool = True,
): ):
payload = Text2ImageTaskCreationRequest( 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, model=model,
input=Text2ImageInputField(prompt=prompt, negative_prompt=negative_prompt), input=Text2ImageInputField(prompt=prompt, negative_prompt=negative_prompt),
parameters=Txt2ImageParametersField( parameters=Txt2ImageParametersField(
@ -289,17 +249,15 @@ class WanTextToImageApi(IO.ComfyNode):
prompt_extend=prompt_extend, prompt_extend=prompt_extend,
watermark=watermark, watermark=watermark,
), ),
),
) )
response = await process_task( if not initial_response.output:
{ raise Exception(f"Unknown error occurred: {initial_response.code} - {initial_response.message}")
"auth_token": cls.hidden.auth_token_comfy_org, response = await poll_op(
"comfy_api_key": cls.hidden.api_key_comfy_org, cls,
}, ApiEndpoint(path=f"/proxy/wan/api/v1/tasks/{initial_response.output.task_id}"),
"/proxy/wan/api/v1/services/aigc/text2image/image-synthesis",
request_model=Text2ImageTaskCreationRequest,
response_model=ImageTaskStatusResponse, response_model=ImageTaskStatusResponse,
payload=payload, status_extractor=lambda x: x.output.task_status,
node_id=cls.hidden.unique_id,
estimated_duration=9, estimated_duration=9,
poll_interval=3, poll_interval=3,
) )
@ -370,7 +328,7 @@ class WanImageToImageApi(IO.ComfyNode):
IO.Boolean.Input( IO.Boolean.Input(
"watermark", "watermark",
default=True, 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, optional=True,
), ),
], ],
@ -402,8 +360,12 @@ class WanImageToImageApi(IO.ComfyNode):
raise ValueError(f"Expected 1 or 2 input images, got {n_images}.") raise ValueError(f"Expected 1 or 2 input images, got {n_images}.")
images = [] images = []
for i in image: for i in image:
images.append("data:image/png;base64," + tensor_to_base64_string(i, total_pixels=4096*4096)) images.append("data:image/png;base64," + tensor_to_base64_string(i, total_pixels=4096 * 4096))
payload = Image2ImageTaskCreationRequest( 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, model=model,
input=Image2ImageInputField(prompt=prompt, negative_prompt=negative_prompt, images=images), input=Image2ImageInputField(prompt=prompt, negative_prompt=negative_prompt, images=images),
parameters=Image2ImageParametersField( parameters=Image2ImageParametersField(
@ -411,19 +373,17 @@ class WanImageToImageApi(IO.ComfyNode):
seed=seed, seed=seed,
watermark=watermark, watermark=watermark,
), ),
),
) )
response = await process_task( if not initial_response.output:
{ raise Exception(f"Unknown error occurred: {initial_response.code} - {initial_response.message}")
"auth_token": cls.hidden.auth_token_comfy_org, response = await poll_op(
"comfy_api_key": cls.hidden.api_key_comfy_org, cls,
}, ApiEndpoint(path=f"/proxy/wan/api/v1/tasks/{initial_response.output.task_id}"),
"/proxy/wan/api/v1/services/aigc/image2image/image-synthesis",
request_model=Image2ImageTaskCreationRequest,
response_model=ImageTaskStatusResponse, response_model=ImageTaskStatusResponse,
payload=payload, status_extractor=lambda x: x.output.task_status,
node_id=cls.hidden.unique_id,
estimated_duration=42, estimated_duration=42,
poll_interval=3, poll_interval=4,
) )
return IO.NodeOutput(await download_url_to_image_tensor(str(response.output.results[0].url))) 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( IO.Boolean.Input(
"watermark", "watermark",
default=True, 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, optional=True,
), ),
], ],
@ -551,7 +511,12 @@ class WanTextToVideoApi(IO.ComfyNode):
if audio is not None: if audio is not None:
validate_audio_duration(audio, 3.0, 29.0) validate_audio_duration(audio, 3.0, 29.0)
audio_url = "data:audio/mp3;base64," + audio_to_base64_string(audio, "mp3", "libmp3lame") audio_url = "data:audio/mp3;base64," + audio_to_base64_string(audio, "mp3", "libmp3lame")
payload = Text2VideoTaskCreationRequest(
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, model=model,
input=Text2VideoInputField(prompt=prompt, negative_prompt=negative_prompt, audio_url=audio_url), input=Text2VideoInputField(prompt=prompt, negative_prompt=negative_prompt, audio_url=audio_url),
parameters=Text2VideoParametersField( parameters=Text2VideoParametersField(
@ -562,17 +527,15 @@ class WanTextToVideoApi(IO.ComfyNode):
prompt_extend=prompt_extend, prompt_extend=prompt_extend,
watermark=watermark, watermark=watermark,
), ),
),
) )
response = await process_task( if not initial_response.output:
{ raise Exception(f"Unknown error occurred: {initial_response.code} - {initial_response.message}")
"auth_token": cls.hidden.auth_token_comfy_org, response = await poll_op(
"comfy_api_key": cls.hidden.api_key_comfy_org, cls,
}, ApiEndpoint(path=f"/proxy/wan/api/v1/tasks/{initial_response.output.task_id}"),
"/proxy/wan/api/v1/services/aigc/video-generation/video-synthesis",
request_model=Text2VideoTaskCreationRequest,
response_model=VideoTaskStatusResponse, response_model=VideoTaskStatusResponse,
payload=payload, status_extractor=lambda x: x.output.task_status,
node_id=cls.hidden.unique_id,
estimated_duration=120 * int(duration / 5), estimated_duration=120 * int(duration / 5),
poll_interval=6, poll_interval=6,
) )
@ -661,7 +624,7 @@ class WanImageToVideoApi(IO.ComfyNode):
IO.Boolean.Input( IO.Boolean.Input(
"watermark", "watermark",
default=True, 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, optional=True,
), ),
], ],
@ -693,12 +656,16 @@ class WanImageToVideoApi(IO.ComfyNode):
): ):
if get_number_of_images(image) != 1: if get_number_of_images(image) != 1:
raise ValueError("Exactly one input image is required.") 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 audio_url = None
if audio is not None: if audio is not None:
validate_audio_duration(audio, 3.0, 29.0) validate_audio_duration(audio, 3.0, 29.0)
audio_url = "data:audio/mp3;base64," + audio_to_base64_string(audio, "mp3", "libmp3lame") audio_url = "data:audio/mp3;base64," + audio_to_base64_string(audio, "mp3", "libmp3lame")
payload = Image2VideoTaskCreationRequest( 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, model=model,
input=Image2VideoInputField( input=Image2VideoInputField(
prompt=prompt, negative_prompt=negative_prompt, img_url=image_url, audio_url=audio_url prompt=prompt, negative_prompt=negative_prompt, img_url=image_url, audio_url=audio_url
@ -711,17 +678,15 @@ class WanImageToVideoApi(IO.ComfyNode):
prompt_extend=prompt_extend, prompt_extend=prompt_extend,
watermark=watermark, watermark=watermark,
), ),
),
) )
response = await process_task( if not initial_response.output:
{ raise Exception(f"Unknown error occurred: {initial_response.code} - {initial_response.message}")
"auth_token": cls.hidden.auth_token_comfy_org, response = await poll_op(
"comfy_api_key": cls.hidden.api_key_comfy_org, cls,
}, ApiEndpoint(path=f"/proxy/wan/api/v1/tasks/{initial_response.output.task_id}"),
"/proxy/wan/api/v1/services/aigc/video-generation/video-synthesis",
request_model=Image2VideoTaskCreationRequest,
response_model=VideoTaskStatusResponse, response_model=VideoTaskStatusResponse,
payload=payload, status_extractor=lambda x: x.output.task_status,
node_id=cls.hidden.unique_id,
estimated_duration=120 * int(duration / 5), estimated_duration=120 * int(duration / 5),
poll_interval=6, poll_interval=6,
) )

View File

@ -35,7 +35,7 @@ def default_base_url() -> str:
async def sleep_with_interrupt( async def sleep_with_interrupt(
seconds: float, seconds: float,
node_cls: type[IO.ComfyNode], node_cls: Optional[type[IO.ComfyNode]],
label: Optional[str] = None, label: Optional[str] = None,
start_ts: Optional[float] = None, start_ts: Optional[float] = None,
estimated_total: Optional[int] = None, estimated_total: Optional[int] = None,

View File

@ -8,7 +8,7 @@ import uuid
from dataclasses import dataclass from dataclasses import dataclass
from enum import Enum from enum import Enum
from io import BytesIO 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 from urllib.parse import urljoin, urlparse
import aiohttp import aiohttp
@ -79,7 +79,7 @@ class _PollUIState:
_RETRY_STATUS = {408, 429, 500, 502, 503, 504} _RETRY_STATUS = {408, 429, 500, 502, 503, 504}
COMPLETED_STATUSES = ["succeeded", "succeed", "success", "completed"] COMPLETED_STATUSES = ["succeeded", "succeed", "success", "completed"]
FAILED_STATUSES = ["cancelled", "failed", "error"] FAILED_STATUSES = ["cancelled", "canceled", "failed", "error"]
QUEUED_STATUSES = ["created", "queued", "queueing", "submitted"] QUEUED_STATUSES = ["created", "queued", "queueing", "submitted"]
@ -97,7 +97,7 @@ async def sync_op(
retry_delay: float = 1.0, retry_delay: float = 1.0,
retry_backoff: float = 2.0, retry_backoff: float = 2.0,
wait_label: str = "Waiting for server", wait_label: str = "Waiting for server",
estimated_total: Optional[int] = None, estimated_duration: Optional[int] = None,
final_label_on_success: Optional[str] = "Completed", final_label_on_success: Optional[str] = "Completed",
progress_origin_ts: Optional[float] = None, progress_origin_ts: Optional[float] = None,
monitor_progress: bool = True, monitor_progress: bool = True,
@ -114,7 +114,7 @@ async def sync_op(
retry_delay=retry_delay, retry_delay=retry_delay,
retry_backoff=retry_backoff, retry_backoff=retry_backoff,
wait_label=wait_label, wait_label=wait_label,
estimated_total=estimated_total, estimated_duration=estimated_duration,
as_binary=False, as_binary=False,
final_label_on_success=final_label_on_success, final_label_on_success=final_label_on_success,
progress_origin_ts=progress_origin_ts, progress_origin_ts=progress_origin_ts,
@ -130,7 +130,7 @@ async def poll_op(
poll_endpoint: ApiEndpoint, poll_endpoint: ApiEndpoint,
*, *,
response_model: Type[M], 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, progress_extractor: Optional[Callable[[M], Optional[int]]] = None,
price_extractor: Optional[Callable[[M], Optional[float]]] = None, price_extractor: Optional[Callable[[M], Optional[float]]] = None,
completed_statuses: Optional[list[Union[str, int]]] = None, completed_statuses: Optional[list[Union[str, int]]] = None,
@ -183,7 +183,7 @@ async def sync_op_raw(
retry_delay: float = 1.0, retry_delay: float = 1.0,
retry_backoff: float = 2.0, retry_backoff: float = 2.0,
wait_label: str = "Waiting for server", wait_label: str = "Waiting for server",
estimated_total: Optional[int] = None, estimated_duration: Optional[int] = None,
as_binary: bool = False, as_binary: bool = False,
final_label_on_success: Optional[str] = "Completed", final_label_on_success: Optional[str] = "Completed",
progress_origin_ts: Optional[float] = None, progress_origin_ts: Optional[float] = None,
@ -212,7 +212,7 @@ async def sync_op_raw(
retry_backoff=retry_backoff, retry_backoff=retry_backoff,
wait_label=wait_label, wait_label=wait_label,
monitor_progress=monitor_progress, monitor_progress=monitor_progress,
estimated_total=estimated_total, estimated_total=estimated_duration,
final_label_on_success=final_label_on_success, final_label_on_success=final_label_on_success,
progress_origin_ts=progress_origin_ts, progress_origin_ts=progress_origin_ts,
) )
@ -223,7 +223,7 @@ async def poll_op_raw(
cls: type[IO.ComfyNode], cls: type[IO.ComfyNode],
poll_endpoint: ApiEndpoint, 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, progress_extractor: Optional[Callable[[dict[str, Any]], Optional[int]]] = None,
price_extractor: Optional[Callable[[dict[str, Any]], Optional[float]]] = None, price_extractor: Optional[Callable[[dict[str, Any]], Optional[float]]] = None,
completed_statuses: Optional[list[Union[str, int]]] = 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. Returns the final JSON response from the poll endpoint.
""" """
completed_states = COMPLETED_STATUSES if completed_statuses is None else completed_statuses completed_states = _normalize_statuses(COMPLETED_STATUSES if completed_statuses is None else completed_statuses)
failed_states = FAILED_STATUSES if failed_statuses is None else failed_statuses failed_states = _normalize_statuses(FAILED_STATUSES if failed_statuses is None else failed_statuses)
queued_states = QUEUED_STATUSES if queued_statuses is None else queued_statuses queued_states = _normalize_statuses(QUEUED_STATUSES if queued_statuses is None else queued_statuses)
started = time.monotonic() started = time.monotonic()
consumed_attempts = 0 # counts only non-queued polls consumed_attempts = 0 # counts only non-queued polls
@ -271,7 +271,7 @@ async def poll_op_raw(
) )
_display_time_progress( _display_time_progress(
cls, cls,
label=state.status_label, status=state.status_label,
elapsed_seconds=int(now - state.started), elapsed_seconds=int(now - state.started),
estimated_total=state.estimated_duration, estimated_total=state.estimated_duration,
price=state.price, price=state.price,
@ -294,7 +294,7 @@ async def poll_op_raw(
retry_delay=retry_delay_per_poll, retry_delay=retry_delay_per_poll,
retry_backoff=retry_backoff_per_poll, retry_backoff=retry_backoff_per_poll,
wait_label="Checking", wait_label="Checking",
estimated_total=None, estimated_duration=None,
as_binary=False, as_binary=False,
final_label_on_success=None, final_label_on_success=None,
monitor_progress=False, monitor_progress=False,
@ -310,7 +310,7 @@ async def poll_op_raw(
timeout=cancel_timeout, timeout=cancel_timeout,
max_retries=0, max_retries=0,
wait_label="Cancelling task", wait_label="Cancelling task",
estimated_total=None, estimated_duration=None,
as_binary=False, as_binary=False,
final_label_on_success=None, final_label_on_success=None,
monitor_progress=False, monitor_progress=False,
@ -318,7 +318,7 @@ async def poll_op_raw(
raise raise
try: try:
status = status_extractor(resp_json) status = _normalize_status_value(status_extractor(resp_json))
except Exception as e: except Exception as e:
logging.error("Status extraction failed: %s", e) logging.error("Status extraction failed: %s", e)
status = None status = None
@ -360,7 +360,7 @@ async def poll_op_raw(
_display_time_progress( _display_time_progress(
cls, cls,
label=status if status else "Completed", status=status if status else "Completed",
elapsed_seconds=int(now_ts - started), elapsed_seconds=int(now_ts - started),
estimated_total=estimated_duration, estimated_total=estimated_duration,
price=state.price, price=state.price,
@ -385,7 +385,7 @@ async def poll_op_raw(
timeout=cancel_timeout, timeout=cancel_timeout,
max_retries=0, max_retries=0,
wait_label="Cancelling task", wait_label="Cancelling task",
estimated_total=None, estimated_duration=None,
as_binary=False, as_binary=False,
final_label_on_success=None, final_label_on_success=None,
monitor_progress=False, monitor_progress=False,
@ -414,12 +414,12 @@ def _display_text(
node_cls: type[IO.ComfyNode], node_cls: type[IO.ComfyNode],
text: Optional[str], text: Optional[str],
*, *,
status: Optional[str] = None, status: Optional[Union[str, int]] = None,
price: Optional[float] = None, price: Optional[float] = None,
) -> None: ) -> None:
display_lines: list[str] = [] display_lines: list[str] = []
if status: 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: if price is not None:
display_lines.append(f"Price: ${float(price):,.4f}") display_lines.append(f"Price: ${float(price):,.4f}")
if text is not None: if text is not None:
@ -430,7 +430,7 @@ def _display_text(
def _display_time_progress( def _display_time_progress(
node_cls: type[IO.ComfyNode], node_cls: type[IO.ComfyNode],
label: str, status: Optional[Union[str, int]],
elapsed_seconds: int, elapsed_seconds: int,
estimated_total: Optional[int] = None, 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)" time_line = f"Time elapsed: {int(elapsed_seconds)}s (~{remaining}s remaining)"
else: else:
time_line = f"Time elapsed: {int(elapsed_seconds)}s" 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]: async def _diagnose_connectivity() -> dict[str, bool]:
@ -640,7 +640,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
file_value.seek(0) file_value.seek(0)
form.add_field(field_name, file_value, filename=filename, content_type=content_type) 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": elif cfg.content_type == "application/x-www-form-urlencoded" and method != "GET":
payload_headers["Content-Type"] = "application/x-www-form-urlencoded" payload_headers["Content-Type"] = "application/x-www-form-urlencoded"
payload_kw["data"] = cfg.data or {} payload_kw["data"] = cfg.data or {}
@ -660,7 +660,6 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
except Exception as _log_e: except Exception as _log_e:
logging.debug("[DEBUG] request logging failed: %s", _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_coro = sess.request(method, url, params=params, **payload_kw)
req_task = asyncio.create_task(req_coro) req_task = asyncio.create_task(req_coro)
@ -684,7 +683,6 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
body = await resp.json() body = await resp.json()
except (ContentTypeError, json.JSONDecodeError): except (ContentTypeError, json.JSONDecodeError):
body = await resp.text() body = await resp.text()
# Retryable?
if resp.status in _RETRY_STATUS and attempt <= cfg.max_retries: if resp.status in _RETRY_STATUS and attempt <= cfg.max_retries:
logging.warning( logging.warning(
"HTTP %s %s -> %s. Retrying in %.2fs (retry %d of %d).", "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) logging.debug("[DEBUG] response logging failed: %s", _log_e)
raise Exception(msg) raise Exception(msg)
# Success
if expect_binary: if expect_binary:
# Read stream in chunks so that cancellation is fast when user interrupts
buff = bytearray() buff = bytearray()
last_tick = time.monotonic() last_tick = time.monotonic()
async for chunk in resp.content.iter_chunked(64 * 1024): 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") logging.debug("Polling was interrupted by user")
raise raise
except (ClientError, asyncio.TimeoutError, socket.gaierror) as e: except (ClientError, asyncio.TimeoutError, socket.gaierror) as e:
# Retry transient connection issues
if attempt <= cfg.max_retries: if attempt <= cfg.max_retries:
logging.warning( logging.warning(
"Connection error calling %s %s. Retrying in %.2fs (%d/%d): %s", "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: if operation_succeeded and cfg.monitor_progress and cfg.final_label_on_success:
_display_time_progress( _display_time_progress(
cfg.node_cls, cfg.node_cls,
label=cfg.final_label_on_success, status=cfg.final_label_on_success,
elapsed_seconds=( elapsed_seconds=(
final_elapsed_seconds final_elapsed_seconds
if final_elapsed_seconds is not None if final_elapsed_seconds is not None
@ -926,3 +921,20 @@ def _wrap_model_extractor(
raise raise
return _wrapped 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

View File

@ -11,9 +11,7 @@ import torch
from PIL import Image from PIL import Image
from comfy.utils import common_upscale from comfy.utils import common_upscale
from comfy_api.input import VideoInput from comfy_api.latest import Input, InputImpl
from comfy_api.input.basic_types import AudioInput
from comfy_api.latest import InputImpl
from ._helpers import mimetype_to_extension from ._helpers import mimetype_to_extension
@ -165,7 +163,7 @@ def tensor_to_data_uri(
return f"data:{mime_type};base64,{base64_string}" 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.""" """Converts an audio input to a base64 string."""
sample_rate: int = audio["sample_rate"] sample_rate: int = audio["sample_rate"]
waveform: torch.Tensor = audio["waveform"] waveform: torch.Tensor = audio["waveform"]
@ -232,7 +230,7 @@ def audio_tensor_to_contiguous_ndarray(waveform: torch.Tensor) -> np.ndarray:
return audio_data_np 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() waveform = audio["waveform"].cpu()
output_buffer = BytesIO() output_buffer = BytesIO()
@ -255,7 +253,7 @@ def audio_input_to_mp3(audio: AudioInput) -> BytesIO:
return output_buffer 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, Returns a new VideoInput object trimmed from the beginning to the specified duration,
using av to avoid loading entire video into memory. using av to avoid loading entire video into memory.

View File

@ -1,7 +1,5 @@
import asyncio import asyncio
import contextlib import contextlib
import logging
import time
import uuid import uuid
from io import BytesIO from io import BytesIO
from pathlib import Path from pathlib import Path
@ -13,10 +11,15 @@ import torch
from aiohttp.client_exceptions import ClientError, ContentTypeError from aiohttp.client_exceptions import ClientError, ContentTypeError
from comfy_api.input_impl import VideoFromFile 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 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 .client import _diagnose_connectivity
from .common_exceptions import ApiServerError, LocalNetworkError, ProcessingInterrupted from .common_exceptions import ApiServerError, LocalNetworkError, ProcessingInterrupted
from .conversions import bytesio_to_image_tensor from .conversions import bytesio_to_image_tensor
@ -32,7 +35,7 @@ async def download_url_to_bytesio(
max_retries: int = 3, max_retries: int = 3,
retry_delay: float = 1.0, retry_delay: float = 1.0,
retry_backoff: float = 2.0, retry_backoff: float = 2.0,
cls: type[ComfyIO.ComfyNode] = None, cls: type[COMFY_IO.ComfyNode] = None,
) -> None: ) -> None:
"""Stream-download a URL to `dest`. """Stream-download a URL to `dest`.
@ -52,7 +55,8 @@ async def download_url_to_bytesio(
attempt = 0 attempt = 0
delay = retry_delay delay = retry_delay
headers = {} headers: dict[str, str] = {}
if url.startswith("/proxy/"): if url.startswith("/proxy/"):
if cls is None: if cls is None:
raise ValueError("For relative 'cloud' paths, the `cls` parameter is required.") raise ValueError("For relative 'cloud' paths, the `cls` parameter is required.")
@ -66,12 +70,44 @@ async def download_url_to_bytesio(
is_path_sink = isinstance(dest, (str, Path)) is_path_sink = isinstance(dest, (str, Path))
fhandle = None 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: try:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
request_logger.log_request_response(operation_id=op_id, request_method="GET", request_url=url) request_logger.log_request_response(operation_id=op_id, request_method="GET", request_url=url)
async with aiohttp.ClientSession(timeout=timeout_cfg) as session: session = aiohttp.ClientSession(timeout=timeout_cfg)
async with session.get(url, headers=headers) as resp: stop_evt = asyncio.Event()
async def _monitor():
try:
while not stop_evt.is_set():
if is_processing_interrupted():
return
await asyncio.sleep(1.0)
except asyncio.CancelledError:
return
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: if resp.status >= 400:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
try: try:
@ -90,7 +126,7 @@ async def download_url_to_bytesio(
) )
if resp.status in _RETRY_STATUS and attempt <= max_retries: if resp.status in _RETRY_STATUS and attempt <= max_retries:
await _sleep_with_cancel(delay) await sleep_with_interrupt(delay, cls, None, None, None)
delay *= retry_backoff delay *= retry_backoff
continue continue
raise Exception(f"Failed to download (HTTP {resp.status}).") raise Exception(f"Failed to download (HTTP {resp.status}).")
@ -105,12 +141,25 @@ async def download_url_to_bytesio(
sink = dest # BytesIO or file-like sink = dest # BytesIO or file-like
written = 0 written = 0
async for chunk in resp.content.iter_chunked(1024 * 1024): while True:
sink.write(chunk) try:
written += len(chunk) 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(): if is_processing_interrupted():
raise ProcessingInterrupted("Task cancelled") raise ProcessingInterrupted("Task cancelled")
if not chunk:
if resp.content.at_eof():
break
continue
sink.write(chunk)
written += len(chunk)
if isinstance(dest, BytesIO): if isinstance(dest, BytesIO):
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
dest.seek(0) dest.seek(0)
@ -125,10 +174,8 @@ async def download_url_to_bytesio(
response_content=f"[streamed {written} bytes to dest]", response_content=f"[streamed {written} bytes to dest]",
) )
return return
except asyncio.CancelledError:
except ProcessingInterrupted: raise ProcessingInterrupted("Task cancelled") from None
logging.debug("Download was interrupted by user")
raise
except (ClientError, asyncio.TimeoutError) as e: except (ClientError, asyncio.TimeoutError) as e:
if attempt <= max_retries: if attempt <= max_retries:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
@ -138,7 +185,7 @@ async def download_url_to_bytesio(
request_url=url, request_url=url,
error_message=f"{type(e).__name__}: {str(e)} (will retry)", 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 delay *= retry_backoff
continue continue
@ -149,8 +196,21 @@ async def download_url_to_bytesio(
) from e ) from e
raise ApiServerError("The remote service appears unreachable at this time.") from e raise ApiServerError("The remote service appears unreachable at this time.") from e
finally: finally:
if stop_evt is not None:
stop_evt.set()
if monitor_task:
monitor_task.cancel()
with contextlib.suppress(Exception): 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: if fhandle:
with contextlib.suppress(Exception):
fhandle.flush() fhandle.flush()
fhandle.close() fhandle.close()
@ -159,7 +219,7 @@ async def download_url_to_image_tensor(
url: str, url: str,
*, *,
timeout: float = None, timeout: float = None,
cls: type[ComfyIO.ComfyNode] = None, cls: type[COMFY_IO.ComfyNode] = None,
) -> torch.Tensor: ) -> torch.Tensor:
"""Downloads an image from a URL and returns a [B, H, W, C] tensor.""" """Downloads an image from a URL and returns a [B, H, W, C] tensor."""
result = BytesIO() result = BytesIO()
@ -171,7 +231,7 @@ async def download_url_to_video_output(
video_url: str, video_url: str,
*, *,
timeout: float = None, timeout: float = None,
cls: type[ComfyIO.ComfyNode] = None, cls: type[COMFY_IO.ComfyNode] = None,
) -> VideoFromFile: ) -> VideoFromFile:
"""Downloads a video from a URL and returns a `VIDEO` output.""" """Downloads a video from a URL and returns a `VIDEO` output."""
result = BytesIO() result = BytesIO()
@ -186,15 +246,3 @@ def _generate_operation_id(method: str, url: str, attempt: int) -> str:
except Exception: except Exception:
slug = "download" slug = "download"
return f"{method}_{slug}_try{attempt}_{uuid.uuid4().hex[:8]}" 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))

View File

@ -11,9 +11,7 @@ import aiohttp
import torch import torch
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from comfy_api.input.basic_types import AudioInput from comfy_api.latest import IO, Input
from comfy_api.input.video_types import VideoInput
from comfy_api.latest import IO
from comfy_api.util import VideoCodec, VideoContainer from comfy_api.util import VideoCodec, VideoContainer
from comfy_api_nodes.apis import request_logger from comfy_api_nodes.apis import request_logger
@ -72,7 +70,7 @@ async def upload_images_to_comfyapi(
async def upload_audio_to_comfyapi( async def upload_audio_to_comfyapi(
cls: type[IO.ComfyNode], cls: type[IO.ComfyNode],
audio: AudioInput, audio: Input.Audio,
*, *,
container_format: str = "mp4", container_format: str = "mp4",
codec_name: str = "aac", codec_name: str = "aac",
@ -92,7 +90,7 @@ async def upload_audio_to_comfyapi(
async def upload_video_to_comfyapi( async def upload_video_to_comfyapi(
cls: type[IO.ComfyNode], cls: type[IO.ComfyNode],
video: VideoInput, video: Input.Video,
*, *,
container: VideoContainer = VideoContainer.MP4, container: VideoContainer = VideoContainer.MP4,
codec: VideoCodec = VideoCodec.H264, codec: VideoCodec = VideoCodec.H264,
@ -104,8 +102,8 @@ async def upload_video_to_comfyapi(
""" """
if max_duration is not None: if max_duration is not None:
try: try:
actual_duration = video.duration_seconds actual_duration = video.get_duration()
if actual_duration is not None and actual_duration > max_duration: if actual_duration > max_duration:
raise ValueError( raise ValueError(
f"Video duration ({actual_duration:.2f}s) exceeds the maximum allowed ({max_duration}s)." 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() req_task.cancel()
raise ProcessingInterrupted("Upload cancelled") raise ProcessingInterrupted("Upload cancelled")
try:
resp = await req_task resp = await req_task
except asyncio.CancelledError:
raise ProcessingInterrupted("Upload cancelled") from None
async with resp: async with resp:
if resp.status >= 400: if resp.status >= 400:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
@ -286,6 +288,8 @@ async def upload_file(
except Exception as e: except Exception as e:
logging.debug("[DEBUG] upload response logging failed: %s", e) logging.debug("[DEBUG] upload response logging failed: %s", e)
return return
except asyncio.CancelledError:
raise ProcessingInterrupted("Task cancelled") from None
except (aiohttp.ClientError, asyncio.TimeoutError) as e: except (aiohttp.ClientError, asyncio.TimeoutError) as e:
if attempt <= max_retries: if attempt <= max_retries:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):

View File

@ -69,12 +69,3 @@ messages_control.disable = [
"no-else-return", "no-else-return",
"unused-variable", "unused-variable",
] ]
[tool.black]
line-length = 120
preview = true
[tool.isort]
profile = "black"