Spaces:
Running
Running
Fix WaveSpeed async job handling - add polling for outputs
Browse filesWaveSpeed returns async responses with urls.get for some models.
Added _poll_for_result method to poll until outputs are ready.
Also added _resolve_video_model helper.
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
src/content_engine/services/cloud_providers/wavespeed_provider.py
CHANGED
|
@@ -151,6 +151,62 @@ class WaveSpeedProvider(CloudProvider):
|
|
| 151 |
return model_name
|
| 152 |
return EDIT_MODEL_MAP["default"]
|
| 153 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 154 |
@staticmethod
|
| 155 |
def _ensure_min_image_size(image_bytes: bytes, min_pixels: int = 3686400) -> bytes:
|
| 156 |
"""Upscale image if total pixel count is below the minimum required by the API.
|
|
@@ -425,6 +481,11 @@ class WaveSpeedProvider(CloudProvider):
|
|
| 425 |
image_url = out[0]
|
| 426 |
elif isinstance(out, str):
|
| 427 |
image_url = out
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 428 |
elif isinstance(output, list) and output:
|
| 429 |
image_url = output[0]
|
| 430 |
elif isinstance(output, str):
|
|
|
|
| 151 |
return model_name
|
| 152 |
return EDIT_MODEL_MAP["default"]
|
| 153 |
|
| 154 |
+
def _resolve_video_model(self, model_name: str | None) -> str:
|
| 155 |
+
"""Resolve a friendly name to a WaveSpeed video model API path."""
|
| 156 |
+
if model_name and model_name in VIDEO_MODEL_MAP:
|
| 157 |
+
return VIDEO_MODEL_MAP[model_name]
|
| 158 |
+
if model_name:
|
| 159 |
+
return model_name
|
| 160 |
+
return VIDEO_MODEL_MAP["default"]
|
| 161 |
+
|
| 162 |
+
async def _poll_for_result(self, poll_url: str, max_attempts: int = 60, interval: float = 2.0) -> str:
|
| 163 |
+
"""Poll the WaveSpeed async job URL until outputs are ready.
|
| 164 |
+
|
| 165 |
+
Returns the first output URL when available.
|
| 166 |
+
"""
|
| 167 |
+
import asyncio
|
| 168 |
+
|
| 169 |
+
for attempt in range(max_attempts):
|
| 170 |
+
try:
|
| 171 |
+
resp = await self._http_client.get(
|
| 172 |
+
poll_url,
|
| 173 |
+
headers={"Authorization": f"Bearer {self._api_key}"},
|
| 174 |
+
)
|
| 175 |
+
resp.raise_for_status()
|
| 176 |
+
result = resp.json()
|
| 177 |
+
|
| 178 |
+
data = result.get("data", result)
|
| 179 |
+
status = data.get("status", "")
|
| 180 |
+
|
| 181 |
+
if status == "failed":
|
| 182 |
+
error_msg = data.get("error", "Unknown error")
|
| 183 |
+
raise RuntimeError(f"WaveSpeed job failed: {error_msg}")
|
| 184 |
+
|
| 185 |
+
outputs = data.get("outputs", [])
|
| 186 |
+
if outputs:
|
| 187 |
+
logger.info("WaveSpeed job completed after %d polls", attempt + 1)
|
| 188 |
+
return outputs[0]
|
| 189 |
+
|
| 190 |
+
# Also check for 'output' field
|
| 191 |
+
if "output" in data:
|
| 192 |
+
out = data["output"]
|
| 193 |
+
if isinstance(out, list) and out:
|
| 194 |
+
return out[0]
|
| 195 |
+
elif isinstance(out, str):
|
| 196 |
+
return out
|
| 197 |
+
|
| 198 |
+
if status == "completed" and not outputs:
|
| 199 |
+
raise RuntimeError(f"WaveSpeed job completed but no outputs: {data}")
|
| 200 |
+
|
| 201 |
+
logger.debug("WaveSpeed job pending (attempt %d/%d)", attempt + 1, max_attempts)
|
| 202 |
+
await asyncio.sleep(interval)
|
| 203 |
+
|
| 204 |
+
except httpx.HTTPStatusError as e:
|
| 205 |
+
logger.warning("Poll request failed: %s", e)
|
| 206 |
+
await asyncio.sleep(interval)
|
| 207 |
+
|
| 208 |
+
raise RuntimeError(f"WaveSpeed job timed out after {max_attempts * interval}s")
|
| 209 |
+
|
| 210 |
@staticmethod
|
| 211 |
def _ensure_min_image_size(image_bytes: bytes, min_pixels: int = 3686400) -> bytes:
|
| 212 |
"""Upscale image if total pixel count is below the minimum required by the API.
|
|
|
|
| 481 |
image_url = out[0]
|
| 482 |
elif isinstance(out, str):
|
| 483 |
image_url = out
|
| 484 |
+
# Async response: outputs empty but urls.get exists - need to poll
|
| 485 |
+
elif not outputs and "urls" in data and "get" in data.get("urls", {}):
|
| 486 |
+
poll_url = data["urls"]["get"]
|
| 487 |
+
logger.info("WaveSpeed returned async job, polling: %s", poll_url[:80])
|
| 488 |
+
image_url = await self._poll_for_result(poll_url)
|
| 489 |
elif isinstance(output, list) and output:
|
| 490 |
image_url = output[0]
|
| 491 |
elif isinstance(output, str):
|