diff options
| -rw-r--r-- | comfier.py | 486 | ||||
| -rw-r--r-- | tests/conftest.py | 13 | ||||
| -rw-r--r-- | tests/test_generate.py | 38 |
3 files changed, 332 insertions, 205 deletions
@@ -55,23 +55,6 @@ class Tools: def __init__(self): self.valves = self.Valves() - async def comfier_generate_seed( - self, bits: Annotated[int, Field(ge=8, le=63)] = 63 - ) -> str: - """ - Generate a new random seed for image generation. - - Calling this is optional: comfier_generate_image picks a random seed itself and reports it in its result. - Use this tool only when a seed is needed before generating, for example to create several images with distinct, known seeds. - Returns the seed as a decimal string, which can be passed as the seed of comfier_generate_image. - - :param bits: Maximum bit-width of the seed, between 8 and 63. Larger values allow larger seeds. Defaults to 63. - :return: The generated seed as a decimal string. - """ - - clamped = max(8, min(int(bits), 63)) - return f"{random.getrandbits(clamped)}" - async def comfier_generate_image( self, positive_prompt: str, @@ -199,6 +182,23 @@ class Tools: "images": image_urls, } + async def comfier_generate_seed( + self, bits: Annotated[int, Field(ge=8, le=63)] = 63 + ) -> str: + """ + Generate a new random seed for image generation. + + Calling this is optional: comfier_generate_image picks a random seed itself and reports it in its result. + Use this tool only when a seed is needed before generating, for example to create several images with distinct, known seeds. + Returns the seed as a decimal string, which can be passed as the seed of comfier_generate_image. + + :param bits: Maximum bit-width of the seed, between 8 and 63. Larger values allow larger seeds. Defaults to 63. + :return: The generated seed as a decimal string. + """ + + clamped = max(8, min(int(bits), 63)) + return f"{random.getrandbits(clamped)}" + def _apply_workflow_parameters( self, workflow: dict[str, Any], @@ -269,51 +269,110 @@ class Tools: except Exception: pass - async def _verify_prompt_alive( + async def _collect_outputs( self, + ws: aiohttp.ClientWebSocketResponse, session: aiohttp.ClientSession, http_base: str, + workflow: dict[str, Any], prompt_id: str, - ) -> bool: + event_emitter: Optional[EventEmitter], + ) -> list[dict[str, Any]]: """ - Verify with ComfyUI that the backend is reachable and still tracking prompt_id - in either /queue (running/pending) or /history. + Follow the ComfyUI WebSocket until the prompt finishes, emitting progress. + + Returns the image descriptors reported by the workflow's output nodes. + Raises ConnectionError if the backend disappears, RuntimeError if the + job is interrupted or fails. """ - timeout = aiohttp.ClientTimeout(total=3.0) - try: - async with session.get(f"{http_base}/queue", timeout=timeout) as q_resp: - q_resp.raise_for_status() - q_data = await q_resp.json() + output_images: list[dict[str, Any]] = [] + last_emit_time = 0.0 + idle_ticks = 0 - running = q_data.get("queue_running") or [] - pending = q_data.get("queue_pending") or [] - in_queue = any( - isinstance(item, list) and len(item) > 1 and item[1] == prompt_id - for item in (running + pending) - ) - if in_queue: - return True + while True: + try: + msg = await asyncio.wait_for(ws.receive(), timeout=1.0) + except asyncio.TimeoutError: + idle_ticks += 1 + if idle_ticks >= 5: + idle_ticks = 0 + await self._ensure_prompt_alive(session, http_base, prompt_id) + continue - async with session.get( - f"{http_base}/history/{prompt_id}", timeout=timeout - ) as h_resp: - h_resp.raise_for_status() - h_data = await h_resp.json() or {} - if prompt_id not in h_data: - return False + idle_ticks = 0 - status_info = (h_data[prompt_id] or {}).get("status") or {} - if status_info.get("status_str") == "error": - raise RuntimeError( - "The job was terminated with an error or interruption on the backend." - ) - return True + if msg.type != aiohttp.WSMsgType.TEXT: + self._raise_if_ws_closed(ws, msg.type) + continue - except aiohttp.ClientError as err: - raise ConnectionError( - f"Lost connection to ComfyUI backend ({err})." - ) from err + payload = json.loads(msg.data) or {} + msg_type = payload.get("type") + data = payload.get("data") or {} + + if data.get("prompt_id") and data.get("prompt_id") != prompt_id: + continue + + if msg_type == "executing": + if data.get("node") is None: + return output_images + await self._emit_node_status(event_emitter, workflow, str(data["node"])) + + elif msg_type == "progress": + last_emit_time = await self._emit_progress( + event_emitter, data, last_emit_time + ) + + elif msg_type == "executed": + images = (data.get("output") or {}).get("images") or [] + if isinstance(images, list): + output_images.extend(img for img in images if isinstance(img, dict)) + + elif msg_type == "execution_interrupted": + raise RuntimeError("Generation was interrupted on the backend.") + + elif msg_type == "execution_error": + error_message = data.get("exception_message", "Unknown ComfyUI error") + raise RuntimeError(f"ComfyUI execution error: {error_message}") + + async def _emit_node_status( + self, + event_emitter: Optional[EventEmitter], + workflow: dict[str, Any], + node_id: str, + ) -> None: + """Announce the workflow node ComfyUI has started executing.""" + + node_info = workflow.get(node_id) or {} + node_title = ( + (node_info.get("_meta") or {}).get("title") + or node_info.get("class_type") + or f"Node {node_id}" + ) + await self._emit_status(event_emitter, f"Running: {node_title}...", done=False) + + async def _emit_progress( + self, + event_emitter: Optional[EventEmitter], + data: dict[str, Any], + last_emit_time: float, + ) -> float: + """Emit a throttled sampling progress update. Returns the new last-emit time.""" + + val = data.get("value", 0) + maximum = data.get("max", 1) + now = time.monotonic() + + if (now - last_emit_time) < 0.25 and val != maximum: + return last_emit_time + + percentage = int((val / maximum) * 100) if maximum else 0 + await self._emit_status( + event_emitter, + f"Sampling: step {val}/{maximum} ({percentage}%)", + done=False, + ) + return now async def _emit_status( self, @@ -337,6 +396,16 @@ class Tools: } ) + async def _ensure_prompt_alive( + self, session: aiohttp.ClientSession, http_base: str, prompt_id: str + ) -> None: + """Raise ConnectionError if ComfyUI no longer tracks the prompt.""" + + if not await self._verify_prompt_alive(session, http_base, prompt_id): + raise ConnectionError( + f"ComfyUI backend lost prompt {prompt_id} (backend may have restarted or crashed)." + ) + async def _execute( self, workflow: dict[str, Any], @@ -352,11 +421,8 @@ class Tools: headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} client_id = str(uuid.uuid4()) - output_images: list[dict[str, Any]] = [] stored_urls: list[str] = [] - last_emit_time = 0.0 prompt_id: Optional[str] = None - idle_ticks = 0 try: async with aiohttp.ClientSession(headers=headers) as session: @@ -367,113 +433,15 @@ class Tools: await self._emit_status( event_emitter, "Submitting workflow to ComfyUI...", done=False ) - - async with session.post( - f"{http_base}/prompt", - json={"prompt": workflow, "client_id": client_id}, - ) as response: - if response.status >= 400: - error_body = await response.text() - raise RuntimeError( - f"ComfyUI rejected workflow (HTTP {response.status}): {error_body}" - ) - queued = await response.json() - prompt_id = queued["prompt_id"] - + prompt_id = await self._submit_workflow( + session, http_base, workflow, client_id + ) await self._emit_status( event_emitter, "Waiting for ComfyUI queue...", done=False ) - - while True: - try: - msg = await asyncio.wait_for(ws.receive(), timeout=1.0) - idle_ticks = 0 - except asyncio.TimeoutError: - idle_ticks += 1 - if idle_ticks >= 5 and prompt_id: - idle_ticks = 0 - is_alive = await self._verify_prompt_alive( - session, http_base, prompt_id - ) - if not is_alive: - raise ConnectionError( - f"ComfyUI backend lost prompt {prompt_id} (backend may have restarted or crashed)." - ) from None - continue - - if msg.type != aiohttp.WSMsgType.TEXT: - if msg.type in ( - aiohttp.WSMsgType.CLOSE, - aiohttp.WSMsgType.CLOSING, - aiohttp.WSMsgType.CLOSED, - aiohttp.WSMsgType.ERROR, - ): - ws_err = ws.exception() - detail = ( - f": {ws_err}" - if ws_err - else f" (code {ws.close_code})" - ) - raise ConnectionError( - f"ComfyUI WebSocket disconnected unexpectedly before completion{detail}." - ) - continue - - payload = json.loads(msg.data) or {} - msg_type = payload.get("type") - data = payload.get("data") or {} - - if data.get("prompt_id") and data.get("prompt_id") != prompt_id: - continue - - if msg_type == "executing": - node_id = data.get("node") - if node_id is None: - break - node_info = workflow.get(str(node_id)) or {} - node_meta = node_info.get("_meta") or {} - node_title = ( - node_meta.get("title") - or node_info.get("class_type") - or f"Node {node_id}" - ) - await self._emit_status( - event_emitter, f"Running: {node_title}...", done=False - ) - - elif msg_type == "progress": - val = data.get("value", 0) - maximum = data.get("max", 1) - percentage = int((val / maximum) * 100) if maximum else 0 - now = time.monotonic() - if (now - last_emit_time) >= 0.25 or val == maximum: - last_emit_time = now - await self._emit_status( - event_emitter, - f"Sampling: step {val}/{maximum} ({percentage}%)", - done=False, - ) - - elif msg_type == "executed": - output_data = data.get("output") or {} - images = output_data.get("images") or [] - if isinstance(images, list): - output_images.extend( - img for img in images if isinstance(img, dict) - ) - - elif msg_type == "execution_interrupted": - raise RuntimeError( - "Generation was interrupted on the backend." - ) - - elif msg_type == "execution_error": - error_message = data.get( - "exception_message", "Unknown ComfyUI error" - ) - raise RuntimeError( - f"ComfyUI execution error: {error_message}" - ) + output_images = await self._collect_outputs( + ws, session, http_base, workflow, prompt_id, event_emitter + ) if not output_images: raise RuntimeError( @@ -485,53 +453,11 @@ class Tools: ) for img in output_images: - filename = img.get("filename") - if not filename: - continue - - params = { - "filename": filename, - "subfolder": img.get("subfolder", ""), - "type": img.get("type", "output"), - } - - async with session.get( - f"{http_base}/view", params=params - ) as img_response: - img_response.raise_for_status() - image_bytes = await img_response.read() - content_type = ( - img_response.headers.get("Content-Type") - or mimetypes.guess_type(filename)[0] - or "image/png" - ) - - uploaded: Any = upload_image( - request=request, - image_data=image_bytes, - content_type=content_type, - metadata=metadata, - user=user, + url = await self._store_image( + session, http_base, img, request, user, metadata ) - - if inspect.isawaitable(uploaded): - uploaded = await uploaded - - if isinstance(uploaded, tuple): - uploaded = uploaded[-1] - - if isinstance(uploaded, str): - url = uploaded - elif isinstance(uploaded, dict) and "id" in uploaded: - url = f"/api/v1/files/{uploaded['id']}/content" - elif hasattr(uploaded, "id"): - url = f"/api/v1/files/{uploaded.id}/content" # pyright: ignore[reportAttributeAccessIssue] - else: - raise RuntimeError( - f"Unexpected return type from upload_image: {type(uploaded)!r}" - ) - - stored_urls.append(url) + if url: + stored_urls.append(url) except asyncio.CancelledError: await asyncio.shield(self._cancel(http_base, headers, prompt_id)) @@ -635,6 +561,25 @@ class Tools: return parsed & ((1 << 63) - 1) + def _raise_if_ws_closed( + self, ws: aiohttp.ClientWebSocketResponse, msg_type: aiohttp.WSMsgType + ) -> None: + """Raise ConnectionError if a non-text WebSocket message signals a closed connection.""" + + if msg_type not in ( + aiohttp.WSMsgType.CLOSE, + aiohttp.WSMsgType.CLOSING, + aiohttp.WSMsgType.CLOSED, + aiohttp.WSMsgType.ERROR, + ): + return + + ws_err = ws.exception() + detail = f": {ws_err}" if ws_err else f" (code {ws.close_code})" + raise ConnectionError( + f"ComfyUI WebSocket disconnected unexpectedly before completion{detail}." + ) + async def _resolve_dimensions(self, size: Optional[str]) -> tuple[int, int]: """Parse 'WIDTHxHEIGHT' from the size parameter or fall back to global IMAGE_SIZE.""" @@ -719,3 +664,136 @@ class Tools: raw_steps = await Config.get("image_generation.steps", 25) return int(raw_steps) + + async def _store_image( + self, + session: aiohttp.ClientSession, + http_base: str, + img: dict[str, Any], + request: Request, + user: Any, + metadata: dict[str, Any], + ) -> Optional[str]: + """ + Download one generated image from ComfyUI and store it in Open WebUI. + + Returns the stored image's URL, or None if the descriptor has no filename. + """ + + filename = img.get("filename") + if not filename: + return None + + params = { + "filename": filename, + "subfolder": img.get("subfolder", ""), + "type": img.get("type", "output"), + } + + async with session.get(f"{http_base}/view", params=params) as img_response: + img_response.raise_for_status() + image_bytes = await img_response.read() + content_type = ( + img_response.headers.get("Content-Type") + or mimetypes.guess_type(filename)[0] + or "image/png" + ) + + uploaded: Any = upload_image( + request=request, + image_data=image_bytes, + content_type=content_type, + metadata=metadata, + user=user, + ) + + if inspect.isawaitable(uploaded): + uploaded = await uploaded + + return self._upload_result_to_url(uploaded) + + async def _submit_workflow( + self, + session: aiohttp.ClientSession, + http_base: str, + workflow: dict[str, Any], + client_id: str, + ) -> str: + """Queue the workflow on ComfyUI and return the new prompt ID.""" + + async with session.post( + f"{http_base}/prompt", + json={"prompt": workflow, "client_id": client_id}, + ) as response: + if response.status >= 400: + error_body = await response.text() + raise RuntimeError( + f"ComfyUI rejected workflow (HTTP {response.status}): {error_body}" + ) + queued = await response.json() + return queued["prompt_id"] + + def _upload_result_to_url(self, uploaded: Any) -> str: + """Turn the return value of Open WebUI's upload_image into a file content URL.""" + + if isinstance(uploaded, tuple): + uploaded = uploaded[-1] + + if isinstance(uploaded, str): + return uploaded + + if isinstance(uploaded, dict) and "id" in uploaded: + return f"/api/v1/files/{uploaded['id']}/content" + + if hasattr(uploaded, "id"): + return f"/api/v1/files/{uploaded.id}/content" # pyright: ignore[reportAttributeAccessIssue] + + raise RuntimeError( + f"Unexpected return type from upload_image: {type(uploaded)!r}" + ) + + async def _verify_prompt_alive( + self, + session: aiohttp.ClientSession, + http_base: str, + prompt_id: str, + ) -> bool: + """ + Verify with ComfyUI that the backend is reachable and still tracking prompt_id + in either /queue (running/pending) or /history. + """ + + timeout = aiohttp.ClientTimeout(total=3.0) + try: + async with session.get(f"{http_base}/queue", timeout=timeout) as q_resp: + q_resp.raise_for_status() + q_data = await q_resp.json() + + running = q_data.get("queue_running") or [] + pending = q_data.get("queue_pending") or [] + in_queue = any( + isinstance(item, list) and len(item) > 1 and item[1] == prompt_id + for item in (running + pending) + ) + if in_queue: + return True + + async with session.get( + f"{http_base}/history/{prompt_id}", timeout=timeout + ) as h_resp: + h_resp.raise_for_status() + h_data = await h_resp.json() or {} + if prompt_id not in h_data: + return False + + status_info = (h_data[prompt_id] or {}).get("status") or {} + if status_info.get("status_str") == "error": + raise RuntimeError( + "The job was terminated with an error or interruption on the backend." + ) + return True + + except aiohttp.ClientError as err: + raise ConnectionError( + f"Lost connection to ComfyUI backend ({err})." + ) from err diff --git a/tests/conftest.py b/tests/conftest.py index 7370405..59e3349 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -30,6 +30,8 @@ class FakeComfy: self.interrupts: list[Any] = [] self.drop_ws = False self.hang = False + self.queue_has_prompt = True + self._stalled: tuple[web.WebSocketResponse, str] | None = None self.url = "" self._ws: web.WebSocketResponse | None = None self._ws_ready = asyncio.Event() @@ -75,6 +77,7 @@ class FakeComfy: async def _play(self, ws: web.WebSocketResponse, pid: str) -> None: if self.hang: + self._stalled = (ws, pid) return if self.drop_ws: await ws.close() @@ -82,10 +85,18 @@ class FakeComfy: for msg in self.script or self.success_script(pid): await ws.send_str(json.dumps(msg)) + async def release(self) -> None: + """Finish a stalled (`hang`) job: send the scripted messages now.""" + assert self._stalled is not None + ws, pid = self._stalled + self.hang = False + await self._play(ws, pid) + async def queue(self, request: web.Request) -> web.Response: if request.method == "POST": return web.json_response({}) - return web.json_response({"queue_running": [[0, "pid-1"]], "queue_pending": []}) + running = [[0, "pid-1"]] if self.queue_has_prompt else [] + return web.json_response({"queue_running": running, "queue_pending": []}) async def interrupt(self, request: web.Request) -> web.Response: self.interrupts.append(await request.json()) diff --git a/tests/test_generate.py b/tests/test_generate.py index 37dda32..8eaf44b 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -91,3 +91,41 @@ async def test_cancel_interrupts_backend( with pytest.raises(asyncio.CancelledError): await task assert comfy.interrupts == [{"prompt_id": "pid-1"}] + + +@pytest.fixture +def fast_idle(monkeypatch: pytest.MonkeyPatch) -> None: + """Shrink the 1s WebSocket idle timeout so stall detection takes ms, not seconds.""" + import comfier + + real_wait_for = asyncio.wait_for + + async def wait_for(aw: Any, timeout: float | None = None) -> Any: + return await real_wait_for(aw, 0.01 if timeout == 1.0 else timeout) + + monkeypatch.setattr(comfier.asyncio, "wait_for", wait_for) + + +async def test_stalled_backend_lost_prompt( + make_tools: Callable[[], Any], comfy: FakeComfy, fast_idle: None +) -> None: + comfy.hang = True + comfy.queue_has_prompt = False # not in /queue, and /history is empty + tools, _ = make_tools() + with pytest.raises(ConnectionError, match="lost prompt pid-1"): + await asyncio.wait_for(_run(tools), timeout=5) + + +async def test_stalled_backend_still_queued_keeps_waiting( + make_tools: Callable[[], Any], comfy: FakeComfy, fast_idle: None +) -> None: + comfy.hang = True # silent, but /queue still lists the prompt + tools, _ = make_tools() + task = asyncio.create_task(_run(tools)) + while not comfy.submitted: + await asyncio.sleep(0.01) + await asyncio.sleep(0.3) # many idle ticks and liveness checks pass + assert not task.done() + await comfy.release() + result, _ = await asyncio.wait_for(task, timeout=5) + assert result["status"] == "success" |