Ivor StrelarandClaude Opus 5.5
17964f3017
A GPT Image run failed with "Connection reset by peer" while polling an already-submitted task, losing an image that WaveSpeed went on to bill. One transient network error was enough to abort the node. Polling, image uploads and result downloads are safe to repeat, so they now retry connection errors and timeouts up to five times (2, 4, 8, 15, 30 s apart) before failing with a message that the task may still finish on WaveSpeed. Submitting a task is deliberately not retried: an unclear failure there could otherwise create a second billed task. Applies to every node on the shared v3 helpers (GPT Image, Nano Banana, Seedream, Seedance, Tripo, Rodin, Meshy). Verified with a stubbed API that drops connections. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
275 lines
11 KiB
Python
275 lines
11 KiB
Python
# ABOUTME: Shared async helpers for nodes that call WaveSpeed's v3 prediction API directly
|
|
# ABOUTME: Lets those nodes run in parallel in ComfyUI: every blocking call runs off the event loop
|
|
|
|
import asyncio
|
|
import io
|
|
import os
|
|
import time
|
|
from urllib.parse import urlparse
|
|
|
|
import numpy as np
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from .utils import tensor2images
|
|
|
|
try:
|
|
from comfy_api.latest import Types
|
|
|
|
File3D = Types.File3D
|
|
except ImportError: # older ComfyUI without 3D file types
|
|
File3D = None
|
|
|
|
RESULT_PATH = "/api/v3/predictions/{id}/result"
|
|
TERMINAL_FAILURES = {"failed", "cancelled", "timeout", "deleted"}
|
|
|
|
# The partner image nodes downscale references to this many pixels before upload.
|
|
MAX_UPLOAD_PIXELS = 2048 * 2048
|
|
|
|
# ComfyUI runs an `async def` node as a background task and moves on to other
|
|
# ready nodes while it waits, which is how partner nodes overlap. That only
|
|
# helps if the coroutine actually yields: WaveSpeedClient uses blocking
|
|
# `requests`, so every call below goes through asyncio.to_thread.
|
|
|
|
|
|
def check_interrupted():
|
|
"""Raise if the user pressed Cancel in ComfyUI; a no-op outside ComfyUI."""
|
|
try:
|
|
import comfy.model_management as mm
|
|
except ImportError:
|
|
return
|
|
mm.throw_exception_if_processing_interrupted()
|
|
|
|
|
|
# Network failures worth retrying: a dropped or reset connection, or a request
|
|
# that timed out. Only used for requests that are safe to repeat (polling,
|
|
# uploads, downloads), never for submitting a task, which WaveSpeed bills.
|
|
TRANSIENT_ERRORS = (requests.exceptions.ConnectionError, requests.exceptions.Timeout,
|
|
requests.exceptions.ChunkedEncodingError)
|
|
RETRY_DELAYS = (2, 4, 8, 15, 30)
|
|
|
|
|
|
async def with_retries(fn, *args, label="WaveSpeed"):
|
|
"""Run a blocking, repeatable request in a worker thread, retrying network
|
|
hiccups with growing pauses before giving up."""
|
|
for attempt, delay in enumerate(RETRY_DELAYS + (None,), start=1):
|
|
try:
|
|
return await asyncio.to_thread(fn, *args)
|
|
except TRANSIENT_ERRORS as e:
|
|
if delay is None:
|
|
raise RuntimeError(
|
|
f"{label}: lost the connection to WaveSpeed {attempt} times in a row ({e}). "
|
|
f"Check your internet connection; any task already submitted may still finish "
|
|
f"on WaveSpeed's side."
|
|
) from e
|
|
print(f"[SF {label}] Network error ({type(e).__name__}); retrying in {delay}s "
|
|
f"({attempt}/{len(RETRY_DELAYS)}).")
|
|
check_interrupted()
|
|
await asyncio.sleep(delay)
|
|
|
|
|
|
async def call(client, fn, *args, retry=False):
|
|
"""Run a blocking API call in a worker thread, turning a 401 into an error
|
|
that says which key was used. `retry` is only for calls that are safe to repeat."""
|
|
try:
|
|
if retry:
|
|
return await with_retries(fn, *args)
|
|
return await asyncio.to_thread(fn, *args)
|
|
except Exception as e:
|
|
if "401" in str(e) or "Unauthorized" in str(e):
|
|
source = client.get("source", "the SF WaveSpeed Client node")
|
|
raise RuntimeError(
|
|
f"WaveSpeed rejected the API key, which came from {source}. "
|
|
f"Check it on wavespeed.ai under API Keys, or clear the api_key field "
|
|
f"on SF WaveSpeed Client to fall back to a key saved elsewhere."
|
|
) from e
|
|
raise
|
|
|
|
|
|
async def wait_for_result(api, task_id, timeout_seconds, label, interval=2):
|
|
"""Poll a prediction until it completes, yielding between polls. Returns the result."""
|
|
loop = asyncio.get_running_loop()
|
|
started = loop.time()
|
|
while True:
|
|
check_interrupted()
|
|
result = await with_retries(api.get, RESULT_PATH.format(id=task_id), label=label)
|
|
status = result.get("status")
|
|
if status == "completed":
|
|
return result
|
|
if status in TERMINAL_FAILURES:
|
|
raise RuntimeError(f"{label} task {status}: {result.get('error') or result}")
|
|
if loop.time() - started > timeout_seconds:
|
|
raise RuntimeError(f"{label} task {task_id} still '{status}' after {timeout_seconds}s.")
|
|
await asyncio.sleep(interval)
|
|
|
|
|
|
async def submit_and_collect(api, client, endpoint, payload, count, timeout_seconds, label):
|
|
"""Submit `count` identical predictions at once and wait for all of them.
|
|
Returns every output URL, in submission order."""
|
|
|
|
async def one():
|
|
task = await call(client, api.post, endpoint, payload)
|
|
if not task.get("id"):
|
|
raise RuntimeError(f"WaveSpeed did not return a task id: {task}")
|
|
result = await wait_for_result(api, task["id"], timeout_seconds, label)
|
|
outputs = result.get("outputs") or []
|
|
return outputs if isinstance(outputs, list) else [outputs]
|
|
|
|
batches = await asyncio.gather(*(one() for _ in range(count)))
|
|
return [url for batch in batches for url in batch]
|
|
|
|
|
|
def flatten_images(tensors):
|
|
"""Every frame of every IMAGE tensor, as PIL images."""
|
|
frames = []
|
|
for tensor in tensors:
|
|
if tensor is None:
|
|
continue
|
|
batch = tensor if tensor.dim() == 4 else tensor.unsqueeze(0)
|
|
frames.extend(tensor2images(batch))
|
|
return frames
|
|
|
|
|
|
def downscale(image, max_pixels=MAX_UPLOAD_PIXELS):
|
|
w, h = image.size
|
|
if w * h <= max_pixels:
|
|
return image
|
|
scale = (max_pixels / (w * h)) ** 0.5
|
|
return image.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS)
|
|
|
|
|
|
async def upload_images(api, client, images):
|
|
"""Upload references concurrently, downscaled first. Returns their URLs in order."""
|
|
return list(await asyncio.gather(*(call(client, api.upload_file, downscale(img), retry=True)
|
|
for img in images)))
|
|
|
|
|
|
async def fetch_images(urls, timeout=120):
|
|
"""Download output images concurrently."""
|
|
|
|
def get(url):
|
|
response = requests.get(url, timeout=timeout)
|
|
response.raise_for_status()
|
|
return Image.open(io.BytesIO(response.content))
|
|
|
|
return list(await asyncio.gather(*(with_retries(get, url, label="download") for url in urls)))
|
|
|
|
|
|
def _has_alpha(image):
|
|
return image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info)
|
|
|
|
|
|
def to_image_batch(frames):
|
|
"""One IMAGE tensor from PIL frames. Keeps alpha when any frame is transparent;
|
|
resizes stragglers to the first frame, since a batch must share one size."""
|
|
if not frames:
|
|
raise RuntimeError("WaveSpeed returned no images.")
|
|
mode = "RGBA" if any(_has_alpha(f) for f in frames) else "RGB"
|
|
width, height = frames[0].size
|
|
tensors = []
|
|
for frame in frames:
|
|
if frame.size != (width, height):
|
|
frame = frame.resize((width, height), Image.LANCZOS)
|
|
tensors.append(torch.from_numpy(np.asarray(frame.convert(mode), dtype=np.float32) / 255.0)[None,])
|
|
return torch.cat(tensors, dim=0)
|
|
|
|
|
|
# --- 3D models -----------------------------------------------------------------
|
|
|
|
|
|
def pick_model_url(outputs, prefer=None):
|
|
"""The model file URL out of a result's `outputs`, which may be a list or a string.
|
|
`prefer` (e.g. ".fbx") wins when a result carries several model files."""
|
|
candidates = outputs if isinstance(outputs, list) else [outputs]
|
|
urls = [u for u in candidates if isinstance(u, str) and u.startswith("http")]
|
|
for url in urls:
|
|
if prefer and urlparse(url).path.lower().endswith(prefer):
|
|
return url
|
|
for url in urls:
|
|
if urlparse(url).path.lower().endswith((".glb", ".fbx")):
|
|
return url
|
|
if urls:
|
|
return urls[0]
|
|
raise RuntimeError(f"No model URL in the result: {outputs!r}")
|
|
|
|
|
|
def output_dir():
|
|
"""ComfyUI's output/3d folder, created if missing."""
|
|
try:
|
|
import folder_paths
|
|
|
|
base = folder_paths.get_output_directory()
|
|
except ImportError:
|
|
base = os.getcwd()
|
|
path = os.path.join(base, "3d")
|
|
os.makedirs(path, exist_ok=True)
|
|
return path
|
|
|
|
|
|
async def submit_and_wait(api, client, endpoint, payload, timeout_seconds, label, interval=3):
|
|
"""Submit one prediction and wait for it. Returns (task_id, result)."""
|
|
task = await call(client, api.post, endpoint, payload)
|
|
task_id = task.get("id")
|
|
if not task_id:
|
|
raise RuntimeError(f"WaveSpeed did not return a task id: {task}")
|
|
print(f"[SF {label}] Submitted {task_id}.")
|
|
return task_id, await wait_for_result(api, task_id, timeout_seconds, label, interval=interval)
|
|
|
|
|
|
async def save_3d_model(url, name, default_ext=".glb"):
|
|
"""Download a model into output/3d as `name` plus the URL's extension.
|
|
Returns (glb, fbx, path); the File3D goes in whichever slot matches the file."""
|
|
ext = os.path.splitext(urlparse(url).path)[1].lower() or default_ext
|
|
path = os.path.join(output_dir(), f"{name}{ext}")
|
|
|
|
def download():
|
|
response = requests.get(url, timeout=300)
|
|
response.raise_for_status()
|
|
with open(path, "wb") as f:
|
|
f.write(response.content)
|
|
|
|
await with_retries(download, label="download")
|
|
|
|
glb = fbx = None
|
|
if File3D is not None:
|
|
model = File3D(path)
|
|
glb = model if ext == ".glb" else None
|
|
fbx = model if ext == ".fbx" else None
|
|
return glb, fbx, path
|
|
|
|
|
|
async def generate_3d_model(api, client, image, endpoint, build_payload, quad, file_prefix,
|
|
timeout_seconds, label):
|
|
"""Image-to-3D from upload to File3D: uploads the first frame, submits the
|
|
job, waits for it, saves the model to output/3d and returns
|
|
(glb, fbx, path, url). `build_payload` receives the uploaded image URL.
|
|
|
|
Tripo returns FBX when quad topology is requested and GLB otherwise, so only
|
|
the output matching the returned file is populated, as on ComfyUI's own
|
|
Tripo nodes.
|
|
"""
|
|
image_url = await call(client, api.upload_file, tensor2images(image)[0], retry=True)
|
|
started = time.time()
|
|
task_id, result = await submit_and_wait(api, client, endpoint, build_payload(image_url), timeout_seconds, label)
|
|
|
|
url = pick_model_url(result.get("outputs"))
|
|
glb, fbx, path = await save_3d_model(url, f"{file_prefix}_{task_id}", ".fbx" if quad else ".glb")
|
|
print(f"[SF {label}] Done in {time.time() - started:.0f}s -> {path}")
|
|
return (glb, fbx, path, url)
|
|
|
|
|
|
# --- Video ---------------------------------------------------------------------
|
|
|
|
|
|
async def fetch_video(url, timeout=600):
|
|
"""Download an output video into memory as a ComfyUI VIDEO."""
|
|
from comfy_api.latest import InputImpl
|
|
|
|
def get():
|
|
response = requests.get(url, timeout=timeout)
|
|
response.raise_for_status()
|
|
return io.BytesIO(response.content)
|
|
|
|
return InputImpl.VideoFromFile(await with_retries(get, label="download"))
|