diff --git a/comfy_api_nodes/nodes_comfy_cloud.py b/comfy_api_nodes/nodes_comfy_cloud.py index 37243b339..9a31951fa 100644 --- a/comfy_api_nodes/nodes_comfy_cloud.py +++ b/comfy_api_nodes/nodes_comfy_cloud.py @@ -1,4 +1,7 @@ +import math +import re from typing import ClassVar +from urllib.parse import quote from typing_extensions import override @@ -31,6 +34,11 @@ from comfy_api_nodes.util import ( _GENERATE_ENDPOINT = ApiEndpoint(path="/proxy/comfy-cloud/workflow/generate", method="POST") +def _task_endpoints(task_id: str) -> tuple[ApiEndpoint, ApiEndpoint]: + task_path = f"/proxy/comfy-cloud/workflow/tasks/{quote(task_id, safe='')}" + return ApiEndpoint(path=task_path), ApiEndpoint(path=f"{task_path}/cancel", method="POST") + + class _ComfyCloudWorkflowNode(IO.ComfyNode): workflow: ClassVar[ComfyCloudWorkflow] node_id: ClassVar[str] @@ -69,7 +77,8 @@ class _ComfyCloudWorkflowNode(IO.ComfyNode): @classmethod async def execute(cls, prompt: str, image: Input.Image | None = None) -> IO.NodeOutput: - validate_string(prompt, min_length=1) + prompt = prompt.strip() + validate_string(prompt, min_length=1, max_length=4096) image_url = None if cls.requires_image: @@ -94,17 +103,17 @@ class _ComfyCloudWorkflowNode(IO.ComfyNode): inputs=inputs, ), ) + polling_endpoint, cancel_endpoint = _task_endpoints(task.task_id) result = await poll_op( cls, - ApiEndpoint(path=task.polling_url), + polling_endpoint, response_model=ComfyCloudStatusResponse, status_extractor=lambda response: response.status, progress_extractor=lambda response: response.progress, - cancel_endpoint=ApiEndpoint(path=task.cancel_url, method="POST"), + cancel_endpoint=cancel_endpoint, ) if not result.output_url: - detail = f": {result.error}" if result.error else "" - raise RuntimeError(f"Comfy Cloud task {result.task_id} completed without an output URL{detail}") + raise RuntimeError("Comfy Cloud task completed without an output URL.") if cls.returns_video: output = await download_url_to_video_output(result.output_url, cls=cls) @@ -406,17 +415,17 @@ class ComfyCloudSeedVR2ImageUpscaleNode(_ComfyCloudWorkflowNode): async def _run_video_workflow(cls: type[IO.ComfyNode], workflow: ComfyCloudWorkflow, inputs: ComfyCloudWorkflowInputs) -> IO.NodeOutput: task = await sync_op(cls, _GENERATE_ENDPOINT, response_model=ComfyCloudGenerateResponse, data=ComfyCloudGenerateRequest(workflow=workflow, inputs=inputs)) + polling_endpoint, cancel_endpoint = _task_endpoints(task.task_id) result = await poll_op( cls, - ApiEndpoint(path=task.polling_url), + polling_endpoint, response_model=ComfyCloudStatusResponse, status_extractor=lambda response: response.status, progress_extractor=lambda response: response.progress, - cancel_endpoint=ApiEndpoint(path=task.cancel_url, method="POST"), + cancel_endpoint=cancel_endpoint, ) if not result.output_url: - detail = f": {result.error}" if result.error else "" - raise RuntimeError(f"Comfy Cloud task {result.task_id} completed without an output URL{detail}") + raise RuntimeError("Comfy Cloud task completed without an output URL.") return IO.NodeOutput(await download_url_to_video_output(result.output_url, cls=cls)) @@ -499,8 +508,8 @@ class ComfyCloudLTX23ImageAudioPerformanceNode(IO.ComfyNode): validate_string(prompt, min_length=1, max_length=4096) if get_number_of_images(image) != 1: raise ValueError("Exactly one input image is required.") - audio_duration = audio["waveform"].shape[-1] / audio["sample_rate"] - if duration_seconds > audio_duration: + audio_duration = _audio_duration(audio) + if duration_seconds - min(1 / float(audio["sample_rate"]), 1e-3) > audio_duration: raise ValueError(f"Duration ({duration_seconds:g}s) exceeds input audio duration ({audio_duration:.2f}s).") image_url = await upload_image_to_comfyapi(cls, image) audio_url = await upload_audio_to_comfyapi(cls, audio) @@ -552,17 +561,17 @@ class ComfyCloudSCAIL2CharacterReplacementNode(IO.ComfyNode): return _video_schema( "ComfyCloudSCAIL2CharacterReplacementNode", "SCAIL-2 Character Replacement", - [IO.Image.Input("reference_character"), IO.Video.Input("driving_video", tooltip="Must contain 81–157 decoded frames."), _prompt_input("scene_prompt"), IO.String.Input("driving_subject", default=""), IO.String.Input("reference_subject", default="human"), _video_seed_input(1)], + [IO.Image.Input("reference_character"), IO.Video.Input("driving_video", tooltip="Must contain 81–157 decoded frames."), _prompt_input("scene_prompt"), IO.String.Input("driving_subject", default="human"), IO.String.Input("reference_subject", default="human"), _video_seed_input(1)], ) @classmethod async def execute(cls, reference_character: Input.Image, driving_video: Input.Video, scene_prompt: str, driving_subject: str, reference_subject: str, seed: int) -> IO.NodeOutput: - validate_string(scene_prompt, min_length=1, max_length=4096) - validate_string(driving_subject, min_length=1, max_length=256) - validate_string(reference_subject, min_length=1, max_length=256) + validate_string(scene_prompt, min_length=1, max_length=4096, field_name="scene_prompt") + validate_string(driving_subject, min_length=1, max_length=256, field_name="driving_subject") + validate_string(reference_subject, min_length=1, max_length=256, field_name="reference_subject") if get_number_of_images(reference_character) != 1: raise ValueError("Exactly one reference character image is required.") - validate_video_frame_count(driving_video, min_frame_count=81, max_frame_count=157) + validate_video_frame_count(driving_video, min_frame_count=81, max_frame_count=157, fail_on_error=True) image_url = await upload_image_to_comfyapi(cls, reference_character) video_url = await upload_video_to_comfyapi(cls, driving_video) return await _run_video_workflow(cls, "video.scail-2-character-replacement.v1", ComfyCloudWorkflowInputs(scene_prompt=scene_prompt, driving_subject=driving_subject, reference_subject=reference_subject, reference_character_url=image_url, driving_video_url=video_url, seed=seed)) @@ -587,38 +596,68 @@ def _audio_schema(node_id: str, display_name: str, inputs: list[IO.Input], outpu def _audio_duration(audio: Input.Audio) -> float: - return audio["waveform"].shape[-1] / audio["sample_rate"] + sample_rate = float(audio["sample_rate"]) + if not math.isfinite(sample_rate) or sample_rate <= 0: + raise ValueError("Audio sample rate must be a positive number.") + return audio["waveform"].shape[-1] / sample_rate def _validate_audio_duration(name: str, audio: Input.Audio, minimum: float, maximum: float) -> None: duration = _audio_duration(audio) - if duration < minimum or duration > maximum: + tolerance = min(1 / float(audio["sample_rate"]), 1e-3) + if duration + tolerance < minimum or duration - tolerance > maximum: raise ValueError(f"{name} duration must be between {minimum:g} and {maximum:g} seconds.") +def _normalize_dialogue(script: str) -> str: + utterances: list[tuple[str, list[str]]] = [] + for raw_line in script.splitlines(): + line = raw_line.strip() + if not line: + continue + match = re.fullmatch(r"(?:SPEAKER\s+)?([A-Z])\s*:\s*(.*)", line, flags=re.IGNORECASE) + if match: + speaker, text = match.groups() + if speaker.upper() not in ("A", "B"): + raise ValueError("Dialogue supports only speakers A and B.") + if utterances and not utterances[-1][1]: + raise ValueError("Dialogue utterances cannot be blank.") + utterances.append((speaker.upper(), [text.strip()] if text.strip() else [])) + elif re.match(r"(?:SPEAKER\s+[A-Z]|NARRATOR)\s*:", line, flags=re.IGNORECASE): + raise ValueError("Dialogue supports only speakers A and B.") + elif utterances: + utterances[-1][1].append(line) + else: + raise ValueError("Dialogue must start with speaker A or B.") + if not utterances: + raise ValueError("Dialogue must contain at least one utterance.") + if not utterances[-1][1]: + raise ValueError("Dialogue utterances cannot be blank.") + return "\n".join(f"SPEAKER {speaker}: {' '.join(lines)}" for speaker, lines in utterances) + + async def _audio_asset(cls: type[IO.ComfyNode], name: str, audio: Input.Audio) -> dict[str, ComfyCloudAssetInput]: return {name: ComfyCloudAssetInput(type="AUDIO", url=await upload_audio_to_comfyapi(cls, audio))} async def _run_audio_workflow(cls: type[IO.ComfyNode], workflow: ComfyCloudWorkflow, inputs: ComfyCloudWorkflowInputs, output_names: tuple[str, ...] = ()) -> IO.NodeOutput: task = await sync_op(cls, _GENERATE_ENDPOINT, response_model=ComfyCloudGenerateResponse, data=ComfyCloudGenerateRequest(workflow=workflow, inputs=inputs)) + polling_endpoint, cancel_endpoint = _task_endpoints(task.task_id) result = await poll_op( cls, - ApiEndpoint(path=task.polling_url), + polling_endpoint, response_model=ComfyCloudStatusResponse, status_extractor=lambda response: response.status, progress_extractor=lambda response: response.progress, - cancel_endpoint=ApiEndpoint(path=task.cancel_url, method="POST"), + cancel_endpoint=cancel_endpoint, ) if output_names: if not result.output_urls or any(not result.output_urls.get(name) for name in output_names): - detail = f": {result.error}" if result.error else "" - raise RuntimeError(f"Comfy Cloud task {result.task_id} completed without all named output URLs{detail}") + raise RuntimeError("Comfy Cloud task completed without all named output URLs.") outputs = [await download_url_to_audio_input(result.output_urls[name], cls=cls) for name in output_names] return IO.NodeOutput(*outputs) if not result.output_url: - detail = f": {result.error}" if result.error else "" - raise RuntimeError(f"Comfy Cloud task {result.task_id} completed without an output URL{detail}") + raise RuntimeError("Comfy Cloud task completed without an output URL.") return IO.NodeOutput(await download_url_to_audio_input(result.output_url, cls=cls)) @@ -709,8 +748,8 @@ class ComfyCloudChatterboxDialogueNode(IO.ComfyNode): @classmethod async def execute(cls, script: str, speaker_a_reference: Input.Audio, speaker_b_reference: Input.Audio, exaggeration: float, cfg_weight: float, temperature: float, seed: int) -> IO.NodeOutput: validate_string(script, min_length=1, max_length=10000, field_name="script") - if any(line.strip() and not line.strip().startswith(("SPEAKER A:", "SPEAKER B:", "SPEAKER C:", "SPEAKER D:")) for line in script.splitlines()): - raise ValueError("Every nonblank utterance must start with SPEAKER A: through SPEAKER D:.") + script = _normalize_dialogue(script) + validate_string(script, min_length=1, max_length=10000, field_name="script") _validate_audio_duration("Speaker A reference", speaker_a_reference, 1, 30) _validate_audio_duration("Speaker B reference", speaker_b_reference, 1, 30) assets = { @@ -758,17 +797,17 @@ class ComfyCloudMelBandRoFormerStemSeparationNode(IO.ComfyNode): async def _run_3d_workflow(cls: type[IO.ComfyNode], workflow: ComfyCloudWorkflow, inputs: ComfyCloudWorkflowInputs, file_format: str) -> IO.NodeOutput: task = await sync_op(cls, _GENERATE_ENDPOINT, response_model=ComfyCloudGenerateResponse, data=ComfyCloudGenerateRequest(workflow=workflow, inputs=inputs)) + polling_endpoint, cancel_endpoint = _task_endpoints(task.task_id) result = await poll_op( cls, - ApiEndpoint(path=task.polling_url), + polling_endpoint, response_model=ComfyCloudStatusResponse, status_extractor=lambda response: response.status, progress_extractor=lambda response: response.progress, - cancel_endpoint=ApiEndpoint(path=task.cancel_url, method="POST"), + cancel_endpoint=cancel_endpoint, ) if not result.output_url: - detail = f": {result.error}" if result.error else "" - raise RuntimeError(f"Comfy Cloud task {result.task_id} completed without an output URL{detail}") + raise RuntimeError("Comfy Cloud task completed without an output URL.") return IO.NodeOutput(await download_url_to_file_3d(result.output_url, file_format, cls=cls)) @@ -840,6 +879,8 @@ class ComfyCloudHunyuan3DMultiViewTo3DNode(IO.ComfyNode): @classmethod async def execute(cls, front_image: Input.Image, back_image: Input.Image, seed: int) -> IO.NodeOutput: + if get_number_of_images(front_image) != 1 or get_number_of_images(back_image) != 1: + raise ValueError("Exactly one front image and one back image are required.") assets = await _image_asset(cls, "front_image", front_image, "Uploading front image") assets.update(await _image_asset(cls, "back_image", back_image, "Uploading back image")) return await _run_3d_workflow(cls, "3d.hunyuan3d-multiview-to-3d.v1", ComfyCloudWorkflowInputs(assets=assets, seed=seed), "glb") diff --git a/comfy_api_nodes/util/download_helpers.py b/comfy_api_nodes/util/download_helpers.py index 3ed2c58c9..91d811e80 100644 --- a/comfy_api_nodes/util/download_helpers.py +++ b/comfy_api_nodes/util/download_helpers.py @@ -27,6 +27,7 @@ from .common_exceptions import ApiServerError, LocalNetworkError, ProcessingInte from .conversions import audio_bytes_to_audio_input, bytesio_to_image_tensor _RETRY_STATUS = {408, 429, 500, 502, 503, 504} +_MAX_IN_MEMORY_DOWNLOAD_BYTES = 512 * 1024 * 1024 async def download_url_to_bytesio( @@ -58,6 +59,8 @@ async def download_url_to_bytesio( attempt = 0 delay = retry_delay headers: dict[str, str] = {} + is_path_sink = isinstance(dest, (str, Path)) + can_reset_sink = is_path_sink or (callable(getattr(dest, "seek", None)) and callable(getattr(dest, "truncate", None))) parsed_url = urlparse(url) if not parsed_url.scheme and not parsed_url.netloc: # is URL relative? @@ -68,10 +71,12 @@ async def download_url_to_bytesio( while True: attempt += 1 + if not is_path_sink and can_reset_sink: + dest.seek(0) + dest.truncate(0) op_id = _generate_operation_id("GET", url, attempt) timeout_cfg = aiohttp.ClientTimeout(total=timeout) - is_path_sink = isinstance(dest, (str, Path)) fhandle = None session: aiohttp.ClientSession | None = None stop_evt: asyncio.Event | None = None @@ -129,11 +134,17 @@ async def download_url_to_bytesio( ) if resp.status in _RETRY_STATUS and attempt <= max_retries: + if not can_reset_sink: + raise Exception(f"Failed to download (HTTP {resp.status}); destination cannot be reset for retry.") await sleep_with_interrupt(delay, cls, None, None, None) delay *= retry_backoff continue raise Exception(f"Failed to download (HTTP {resp.status}).") + max_bytes = None if is_path_sink else _MAX_IN_MEMORY_DOWNLOAD_BYTES + if max_bytes is not None and resp.content_length is not None and resp.content_length > max_bytes: + raise ValueError(f"Download exceeds the {max_bytes}-byte in-memory limit.") + if is_path_sink: p = Path(str(dest)) with contextlib.suppress(Exception): @@ -160,10 +171,12 @@ async def download_url_to_bytesio( break continue - sink.write(chunk) written += len(chunk) + if max_bytes is not None and written > max_bytes: + raise ValueError(f"Download exceeds the {max_bytes}-byte in-memory limit.") + sink.write(chunk) - if isinstance(dest, BytesIO): + if not is_path_sink and hasattr(dest, "seek"): with contextlib.suppress(Exception): dest.seek(0) @@ -180,6 +193,8 @@ async def download_url_to_bytesio( raise ProcessingInterrupted("Task cancelled") from None except (ClientError, OSError) as e: if attempt <= max_retries: + if not can_reset_sink: + raise ApiServerError("The download failed and its destination cannot be reset for retry.") from e request_logger.log_request_response( operation_id=op_id, request_method="GET", diff --git a/comfy_api_nodes/util/validation_utils.py b/comfy_api_nodes/util/validation_utils.py index f01edea96..47da53e92 100644 --- a/comfy_api_nodes/util/validation_utils.py +++ b/comfy_api_nodes/util/validation_utils.py @@ -138,11 +138,14 @@ def validate_video_frame_count( video: Input.Video, min_frame_count: int | None = None, max_frame_count: int | None = None, + fail_on_error: bool = False, ): try: frame_count = video.get_frame_count() except Exception as e: logging.error("Error getting frame count of video: %s", e) + if fail_on_error: + raise ValueError("Unable to determine video frame count.") from e return if min_frame_count is not None and min_frame_count > frame_count: diff --git a/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py b/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py index e052075b8..5a69893fb 100644 --- a/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py +++ b/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py @@ -3,6 +3,7 @@ from io import BytesIO from typing import get_args from unittest.mock import AsyncMock, Mock +import aiohttp import pytest import torch @@ -110,6 +111,61 @@ def test_contract_omits_optional_status_fields(): assert status.model_dump(exclude_none=True) == {"task_id": "task-1", "status": "queued"} +@pytest.mark.parametrize( + "node", + [ + nodes_comfy_cloud.ComfyCloudTextToImageNode, + nodes_comfy_cloud.ComfyCloudTextToVideoNode, + nodes_comfy_cloud.ComfyCloudImageToVideoNode, + nodes_comfy_cloud.ComfyCloudImageEditNode, + ], +) +def test_legacy_nodes_reject_oversized_prompts(monkeypatch, node): + sync = AsyncMock() + monkeypatch.setattr(nodes_comfy_cloud, "sync_op", sync) + + with pytest.raises(Exception, match="4096"): + asyncio.run(node.execute("x" * 4097, object())) + sync.assert_not_awaited() + + +def test_legacy_nodes_strip_prompts_before_submission(monkeypatch): + run = AsyncMock(return_value=("output",)) + monkeypatch.setattr(nodes_comfy_cloud.ComfyCloudTextToImageNode, "_run", run) + + asyncio.run(nodes_comfy_cloud.ComfyCloudTextToImageNode.execute(" prompt ")) + + assert run.call_args.args[0].prompt == "prompt" + + +def test_task_routes_ignore_response_urls_and_errors_hide_task_token(monkeypatch): + sync = AsyncMock( + return_value=ComfyCloudGenerateResponse( + task_id="secret/task-token", + status="queued", + polling_url="https://attacker.example/poll", + cancel_url="https://attacker.example/cancel", + ) + ) + poll = AsyncMock( + return_value=ComfyCloudStatusResponse( + task_id="secret/task-token", + status="completed", + error="provider details with secret/task-token", + ) + ) + monkeypatch.setattr(nodes_comfy_cloud, "sync_op", sync) + monkeypatch.setattr(nodes_comfy_cloud, "poll_op", poll) + + with pytest.raises(RuntimeError) as error: + asyncio.run(nodes_comfy_cloud.ComfyCloudTextToVideoNode.execute("A prompt")) + + assert poll.call_args.args[1].path == "/proxy/comfy-cloud/workflow/tasks/secret%2Ftask-token" + assert poll.call_args.kwargs["cancel_endpoint"].path == "/proxy/comfy-cloud/workflow/tasks/secret%2Ftask-token/cancel" + assert "task-token" not in str(error.value) + assert "provider details" not in str(error.value) + + IMAGE_POC_NODES = [ ( nodes_comfy_cloud.ComfyCloudIdeogram4DesignNode, @@ -302,6 +358,157 @@ def test_scail_stages_reference_image_and_driving_video(monkeypatch): video.get_frame_count.assert_called_once() +def test_scail_defaults_and_frame_count_fail_closed_before_upload(monkeypatch): + schema = {input.id: input for input in nodes_comfy_cloud.ComfyCloudSCAIL2CharacterReplacementNode.define_schema().inputs} + image_upload = AsyncMock() + video_upload = AsyncMock() + video = Mock() + video.get_frame_count.side_effect = RuntimeError("decode failed") + monkeypatch.setattr(nodes_comfy_cloud, "upload_image_to_comfyapi", image_upload) + monkeypatch.setattr(nodes_comfy_cloud, "upload_video_to_comfyapi", video_upload) + monkeypatch.setattr(nodes_comfy_cloud, "get_number_of_images", lambda image: 1) + + assert schema["driving_subject"].default == "human" + with pytest.raises(ValueError, match="Unable to determine video frame count"): + asyncio.run(nodes_comfy_cloud.ComfyCloudSCAIL2CharacterReplacementNode.execute(object(), video, "park", "human", "human", 1)) + image_upload.assert_not_awaited() + video_upload.assert_not_awaited() + + +@pytest.mark.parametrize("frame_count", [80, 158]) +def test_scail_rejects_out_of_range_frames_before_upload(monkeypatch, frame_count): + upload = AsyncMock() + video = Mock() + video.get_frame_count.return_value = frame_count + monkeypatch.setattr(nodes_comfy_cloud, "upload_image_to_comfyapi", upload) + monkeypatch.setattr(nodes_comfy_cloud, "upload_video_to_comfyapi", upload) + monkeypatch.setattr(nodes_comfy_cloud, "get_number_of_images", lambda image: 1) + + with pytest.raises(ValueError, match="frame count"): + asyncio.run(nodes_comfy_cloud.ComfyCloudSCAIL2CharacterReplacementNode.execute(object(), video, "park", "human", "human", 1)) + upload.assert_not_awaited() + + +def test_in_memory_download_resets_retry_and_enforces_stream_limit(monkeypatch): + class Content: + def __init__(self, chunks): + self.chunks = iter(chunks) + self.finished = False + + async def read(self, size): + chunk = next(self.chunks) + if isinstance(chunk, Exception): + raise chunk + if not chunk: + self.finished = True + return chunk + + def at_eof(self): + return self.finished + + class Response: + status = 200 + headers = {} + + def __init__(self, chunks, content_length=None): + self.content = Content(chunks) + self.content_length = content_length + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + responses = [Response([b"partial", aiohttp.ClientPayloadError("retry")]), Response([b"final", b""])] + + class Session: + def __init__(self, timeout): + pass + + async def get(self, url, headers): + return responses.pop(0) + + async def close(self): + pass + + monkeypatch.setattr(download_helpers.aiohttp, "ClientSession", Session) + monkeypatch.setattr(download_helpers, "sleep_with_interrupt", AsyncMock()) + destination = BytesIO() + + asyncio.run(download_helpers.download_url_to_bytesio("https://example.com/result", destination)) + assert destination.read() == b"final" + + monkeypatch.setattr(download_helpers, "_MAX_IN_MEMORY_DOWNLOAD_BYTES", 4) + responses.append(Response([b"12345", b""])) + with pytest.raises(ValueError, match="in-memory limit"): + asyncio.run(download_helpers.download_url_to_bytesio("https://example.com/result", BytesIO())) + + responses.append(Response([], content_length=5)) + with pytest.raises(ValueError, match="in-memory limit"): + asyncio.run(download_helpers.download_url_to_bytesio("https://example.com/result", BytesIO())) + + +def test_file_object_download_resets_retry_and_enforces_stream_limit(monkeypatch, tmp_path): + destination = (tmp_path / "result.bin").open("w+b") + destination.write(b"stale") + + class Content: + def __init__(self, chunks): + self.chunks = iter(chunks) + self.finished = False + + async def read(self, size): + chunk = next(self.chunks) + if isinstance(chunk, Exception): + raise chunk + if not chunk: + self.finished = True + return chunk + + def at_eof(self): + return self.finished + + class Response: + status = 200 + headers = {} + content_length = None + + def __init__(self, chunks): + self.content = Content(chunks) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + responses = [Response([b"partial", aiohttp.ClientPayloadError("retry")]), Response([b"done", b""])] + + class Session: + def __init__(self, timeout): + pass + + async def get(self, url, headers): + return responses.pop(0) + + async def close(self): + pass + + monkeypatch.setattr(download_helpers.aiohttp, "ClientSession", Session) + monkeypatch.setattr(download_helpers, "sleep_with_interrupt", AsyncMock()) + monkeypatch.setattr(download_helpers, "_MAX_IN_MEMORY_DOWNLOAD_BYTES", 10) + + asyncio.run(download_helpers.download_url_to_bytesio("https://example.com/result", destination)) + assert destination.read() == b"done" + + monkeypatch.setattr(download_helpers, "_MAX_IN_MEMORY_DOWNLOAD_BYTES", 4) + responses.append(Response([b"12345", b""])) + with pytest.raises(ValueError, match="in-memory limit"): + asyncio.run(download_helpers.download_url_to_bytesio("https://example.com/result", destination)) + destination.close() + + def test_download_cloud_audio_url_to_audio_input(monkeypatch): node = nodes_comfy_cloud.ComfyCloudTextToImageNode downloaded = b"encoded audio" @@ -399,7 +606,7 @@ def test_audio_poc_request_mapping_and_named_result_decoding(monkeypatch): assert request.workflow == "audio.melbandroformer-stem-separation.v1" assert request.inputs.model_dump(exclude_none=True) == {"assets": {"audio": {"type": "AUDIO", "url": "/uploads/song.m4a"}}} assert [call.args[0] for call in download.await_args_list] == ["/vocals.mp3", "/instruments.mp3"] - assert poll.call_args.kwargs["cancel_endpoint"].path == "/tasks/task-audio/cancel" + assert poll.call_args.kwargs["cancel_endpoint"].path == "/proxy/comfy-cloud/workflow/tasks/task-audio/cancel" assert tuple(output) == ("vocals-audio", "instruments-audio") @@ -429,11 +636,50 @@ def test_chatterbox_dialogue_rejects_invalid_speaker_labels(monkeypatch): monkeypatch.setattr(nodes_comfy_cloud, "upload_audio_to_comfyapi", upload) audio = {"waveform": torch.zeros(1, 1, 48000), "sample_rate": 48000} - with pytest.raises(ValueError, match="Every nonblank utterance"): + with pytest.raises(ValueError, match="only speakers A and B"): asyncio.run(nodes_comfy_cloud.ComfyCloudChatterboxDialogueNode.execute("NARRATOR: Hello", audio, audio, 0.5, 0.5, 0.8, 0)) upload.assert_not_awaited() +def test_chatterbox_dialogue_normalizes_labels_and_continuations(monkeypatch): + run = AsyncMock(return_value=("audio-output",)) + monkeypatch.setattr(nodes_comfy_cloud, "_run_audio_workflow", run) + monkeypatch.setattr(nodes_comfy_cloud, "upload_audio_to_comfyapi", AsyncMock(side_effect=["/a", "/b"])) + audio = {"waveform": torch.zeros(1, 1, 48000), "sample_rate": 48000} + + asyncio.run(nodes_comfy_cloud.ComfyCloudChatterboxDialogueNode.execute("a: Hello\ncontinued\n\nSpeaker B: Hi", audio, audio, 0.5, 0.5, 0.8, 0)) + + assert run.call_args.args[2].script == "SPEAKER A: Hello continued\nSPEAKER B: Hi" + + +def test_chatterbox_dialogue_accepts_colons_and_label_only_lines(): + assert nodes_comfy_cloud._normalize_dialogue("SPEAKER A :\nMeet at 10:30\nhttps://example.com\nB: Done") == ( + "SPEAKER A: Meet at 10:30 https://example.com\nSPEAKER B: Done" + ) + + +@pytest.mark.parametrize("script", ["SPEAKER A:", "just a continuation", " "]) +def test_chatterbox_dialogue_rejects_blank_or_unattributed_text(monkeypatch, script): + upload = AsyncMock() + monkeypatch.setattr(nodes_comfy_cloud, "upload_audio_to_comfyapi", upload) + audio = {"waveform": torch.zeros(1, 1, 48000), "sample_rate": 48000} + + with pytest.raises(Exception): + asyncio.run(nodes_comfy_cloud.ComfyCloudChatterboxDialogueNode.execute(script, audio, audio, 0.5, 0.5, 0.8, 0)) + upload.assert_not_awaited() + + +def test_audio_duration_tolerates_one_sample_and_rejects_invalid_sample_rate(): + one_sample_over = {"waveform": torch.zeros(1, 1, 48001), "sample_rate": 48000} + nodes_comfy_cloud._validate_audio_duration("Audio", one_sample_over, 0.5, 1) + + with pytest.raises(ValueError, match="sample rate"): + nodes_comfy_cloud._validate_audio_duration("Audio", {"waveform": torch.zeros(1, 1, 1), "sample_rate": 0}, 0.5, 1) + + with pytest.raises(ValueError, match="between"): + nodes_comfy_cloud._validate_audio_duration("Audio", {"waveform": torch.zeros(1, 1, 31), "sample_rate": 1}, 1, 30) + + @pytest.mark.parametrize(("file_format", "expected_format"), [(".GLB", "glb"), ("SPZ", "spz")]) def test_download_cloud_3d_url_to_file_3d(monkeypatch, file_format, expected_format): node = nodes_comfy_cloud.ComfyCloudTextToImageNode @@ -521,6 +767,28 @@ def test_3d_poc_schema_defaults_and_ranges(): assert (panorama["merge_resolution"].default, panorama["merge_resolution"].min, panorama["merge_resolution"].max) == (1024, 256, 8192) +def test_hunyuan_multiview_validates_both_images_before_upload(monkeypatch): + upload = AsyncMock() + monkeypatch.setattr(nodes_comfy_cloud, "upload_image_to_comfyapi", upload) + monkeypatch.setattr(nodes_comfy_cloud, "get_number_of_images", lambda image: image) + + with pytest.raises(ValueError, match="Exactly one front image and one back image"): + asyncio.run(nodes_comfy_cloud.ComfyCloudHunyuan3DMultiViewTo3DNode.execute(1, 2, 9)) + upload.assert_not_awaited() + + +def test_extension_preserves_all_23_poc_node_registrations(): + legacy_nodes = { + nodes_comfy_cloud.ComfyCloudTextToImageNode, + nodes_comfy_cloud.ComfyCloudTextToVideoNode, + nodes_comfy_cloud.ComfyCloudImageToVideoNode, + nodes_comfy_cloud.ComfyCloudImageEditNode, + } + registered = set(asyncio.run(nodes_comfy_cloud.ComfyCloudExtension().get_node_list())) + + assert len(registered - legacy_nodes) == 23 + + def test_3d_workflow_submission_polling_cancel_and_download(monkeypatch): sync = AsyncMock(return_value=ComfyCloudGenerateResponse(task_id="task-3d", status="queued", polling_url="/tasks/task-3d", cancel_url="/tasks/task-3d/cancel")) poll = AsyncMock(return_value=ComfyCloudStatusResponse(task_id="task-3d", status="completed", output_url="/results/model.spz")) @@ -534,8 +802,8 @@ def test_3d_workflow_submission_polling_cancel_and_download(monkeypatch): request = sync.call_args.kwargs["data"] assert request.workflow == "3d.triposplat-image-to-gaussian-splat.v1" assert request.inputs.model_dump(exclude_none=True) == {"seed": 46} - assert poll.call_args.args[1].path == "/tasks/task-3d" - assert poll.call_args.kwargs["cancel_endpoint"].path == "/tasks/task-3d/cancel" + assert poll.call_args.args[1].path == "/proxy/comfy-cloud/workflow/tasks/task-3d" + assert poll.call_args.kwargs["cancel_endpoint"].path == "/proxy/comfy-cloud/workflow/tasks/task-3d/cancel" assert poll.call_args.kwargs["cancel_endpoint"].method == "POST" download.assert_awaited_once_with("/results/model.spz", "spz", cls=nodes_comfy_cloud.ComfyCloudTripoSplatImageToGaussianSplatNode) assert output[0] == "spz-output"