800 lines
34 KiB
Python
800 lines
34 KiB
Python
# wavespeed_image_api.py
|
|
import asyncio
|
|
import logging
|
|
import httpx
|
|
from typing import Any, Dict, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
async def send_wavespeed_image_edit_request(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: str,
|
|
image_base64: str,
|
|
guidance_scale: Optional[float] = None,
|
|
seed: int = -1,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Sends a request to the WaveSpeedAI Image Edit/Enhancement APIs and polls for the result.
|
|
Handles models like SeedEdit V3 and Portrait.
|
|
"""
|
|
if not api_key:
|
|
raise ValueError("WaveSpeed API key is required.")
|
|
if not image_base64:
|
|
raise ValueError("Input image is required for this model.")
|
|
|
|
# Always request base64 output for direct use in ComfyUI
|
|
payload = {
|
|
"prompt": prompt,
|
|
"image": f"data:image/png;base64,{image_base64}",
|
|
"enable_base64_output": True,
|
|
}
|
|
if seed != -1:
|
|
payload["seed"] = seed
|
|
|
|
# Conditionally add guidance_scale for models that support it
|
|
if guidance_scale is not None:
|
|
payload["guidance_scale"] = guidance_scale
|
|
|
|
submit_url = f"https://api.wavespeed.ai/api/v3/{model}"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
request_id = None
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
try:
|
|
logger.info(f"Submitting image edit task to WaveSpeedAI model: {model}")
|
|
response = await client.post(submit_url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
if result.get("data") and result["data"].get("id"):
|
|
request_id = result["data"]["id"]
|
|
logger.info(f"Task submitted successfully. Request ID: {request_id}")
|
|
else:
|
|
logger.error(f"WaveSpeedAI API Error: Unexpected submission response format. {result}")
|
|
return {"error": "Unexpected submission response format.", "details": result}
|
|
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error(f"WaveSpeedAI API Error on submission: {e.response.status_code} - {e.response.text}")
|
|
return {"error": f"HTTP {e.response.status_code}", "details": e.response.text}
|
|
except Exception as e:
|
|
logger.error(f"Error submitting to WaveSpeedAI: {e}", exc_info=True)
|
|
return {"error": str(e)}
|
|
|
|
if not request_id:
|
|
return {"error": "Failed to get a request ID from WaveSpeedAI."}
|
|
|
|
# --- Polling for result ---
|
|
result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
|
|
polling_headers = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
for attempt in range(60): # Poll for up to 60 seconds (60 * 1s)
|
|
try:
|
|
await asyncio.sleep(1) # 1-second polling interval
|
|
logger.debug(f"Polling for result... (Attempt {attempt + 1})")
|
|
|
|
poll_response = await client.get(result_url, headers=polling_headers)
|
|
poll_response.raise_for_status()
|
|
|
|
poll_result = poll_response.json().get("data", {})
|
|
status = poll_result.get("status")
|
|
|
|
if status == "completed":
|
|
logger.info("WaveSpeedAI task completed.")
|
|
outputs = poll_result.get("outputs", [])
|
|
if not outputs:
|
|
return {"error": "Task completed but no outputs found."}
|
|
|
|
# Expecting a list with one base64 string
|
|
b64_data = outputs[0]
|
|
# The API returns the full data URI, strip the prefix
|
|
if "base64," in b64_data:
|
|
b64_json = b64_data.split("base64,")[1]
|
|
else:
|
|
b64_json = b64_data
|
|
|
|
return {"data": [{"b64_json": b64_json}]}
|
|
|
|
elif status == "failed":
|
|
error_msg = poll_result.get("error", "Unknown error.")
|
|
logger.error(f"WaveSpeedAI task failed: {error_msg}")
|
|
return {"error": f"Task failed: {error_msg}"}
|
|
else:
|
|
logger.info(f"Task status: {status}. Continuing to poll.")
|
|
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error(f"WaveSpeedAI API Error during polling: {e.response.status_code} - {e.response.text}")
|
|
return {"error": f"HTTP {e.response.status_code} while polling", "details": e.response.text}
|
|
except Exception as e:
|
|
logger.error(f"Error polling WaveSpeedAI: {e}", exc_info=True)
|
|
return {"error": f"Polling failed: {str(e)}"}
|
|
|
|
return {"error": "Polling timed out after 60 seconds."}
|
|
|
|
|
|
async def send_wavespeed_flux_request(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: str,
|
|
image_base64: Optional[str] = None, # For single image models
|
|
images_base64: Optional[list[str]] = None, # For multi-image models
|
|
size: Optional[str] = None,
|
|
num_inference_steps: int = 28,
|
|
guidance_scale: float = 2.5,
|
|
num_images: int = 1,
|
|
seed: int = -1,
|
|
enable_safety_checker: bool = True,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Sends a request to the WaveSpeedAI FLUX Kontext Dev API and polls for the result.
|
|
"""
|
|
if not api_key:
|
|
raise ValueError("WaveSpeed API key is required.")
|
|
|
|
payload = {
|
|
"prompt": prompt,
|
|
"num_inference_steps": num_inference_steps,
|
|
"guidance_scale": guidance_scale,
|
|
"num_images": num_images,
|
|
"enable_safety_checker": enable_safety_checker,
|
|
"enable_base64_output": True,
|
|
}
|
|
if seed != -1:
|
|
payload["seed"] = seed
|
|
if size:
|
|
payload["size"] = size
|
|
if images_base64:
|
|
payload["images"] = [f"data:image/png;base64,{b64}" for b64 in images_base64]
|
|
elif image_base64:
|
|
payload["image"] = f"data:image/png;base64,{image_base64}"
|
|
|
|
submit_url = f"https://api.wavespeed.ai/api/v3/{model}"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
request_id = None
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
try:
|
|
logger.info(f"Submitting task to WaveSpeedAI Flux model: {model}")
|
|
response = await client.post(submit_url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
if result.get("data") and result["data"].get("id"):
|
|
request_id = result["data"]["id"]
|
|
logger.info(f"Task submitted successfully. Request ID: {request_id}")
|
|
else:
|
|
logger.error(f"WaveSpeedAI API Error: Unexpected submission response. {result}")
|
|
return {"error": "Unexpected submission response format.", "details": result}
|
|
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error(f"WaveSpeedAI API Error on submission: {e.response.status_code} - {e.response.text}")
|
|
return {"error": f"HTTP {e.response.status_code}", "details": e.response.text}
|
|
except Exception as e:
|
|
logger.error(f"Error submitting to WaveSpeedAI: {e}", exc_info=True)
|
|
return {"error": str(e)}
|
|
|
|
if not request_id:
|
|
return {"error": "Failed to get a request ID from WaveSpeedAI."}
|
|
|
|
# --- Polling for result ---
|
|
result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
|
|
polling_headers = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
for attempt in range(120): # Poll for up to 120 seconds
|
|
try:
|
|
await asyncio.sleep(1)
|
|
logger.debug(f"Polling for result... (Attempt {attempt + 1})")
|
|
|
|
poll_response = await client.get(result_url, headers=polling_headers)
|
|
poll_response.raise_for_status()
|
|
|
|
poll_result = poll_response.json().get("data", {})
|
|
status = poll_result.get("status")
|
|
|
|
if status == "completed":
|
|
logger.info("WaveSpeedAI Flux task completed.")
|
|
outputs = poll_result.get("outputs", [])
|
|
if not outputs:
|
|
return {"error": "Task completed but no outputs found."}
|
|
|
|
# Process multiple potential outputs
|
|
b64_jsons = []
|
|
for out in outputs:
|
|
if "base64," in out:
|
|
b64_jsons.append(out.split("base64,")[1])
|
|
else:
|
|
b64_jsons.append(out)
|
|
|
|
return {"data": [{"b64_json": b64} for b64 in b64_jsons]}
|
|
|
|
elif status == "failed":
|
|
error_msg = poll_result.get("error", "Unknown error.")
|
|
logger.error(f"WaveSpeedAI Flux task failed: {error_msg}")
|
|
return {"error": f"Task failed: {error_msg}"}
|
|
else:
|
|
logger.info(f"Task status: {status}. Continuing to poll.")
|
|
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error(f"WaveSpeedAI API Error during polling: {e.response.status_code} - {e.response.text}")
|
|
return {"error": f"HTTP {e.response.status_code} while polling", "details": e.response.text}
|
|
except Exception as e:
|
|
logger.error(f"Error polling WaveSpeedAI: {e}", exc_info=True)
|
|
return {"error": f"Polling failed: {str(e)}"}
|
|
|
|
return {"error": "Polling timed out after 120 seconds."}
|
|
|
|
async def send_wavespeed_seedream_request(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: str,
|
|
size: Optional[str] = "2048*2048",
|
|
seed: Optional[int] = -1,
|
|
images_base64: Optional[list[str]] = None,
|
|
max_images: Optional[int] = 1,
|
|
enable_sync_mode: bool = False,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Generic function for Bytedance Seedream V4 models on WaveSpeed.
|
|
Handles task submission and polling for results.
|
|
"""
|
|
if not api_key:
|
|
raise ValueError("WaveSpeed API key is required.")
|
|
|
|
# Correct the model path for the API endpoint
|
|
api_model_path = model
|
|
if model.startswith("bytedance/seedream-v4-"):
|
|
suffix = model.replace("bytedance/seedream-v4-", "")
|
|
api_model_path = f"bytedance/seedream-v4/{suffix}"
|
|
|
|
submit_url = f"https://api.wavespeed.ai/api/v3/{api_model_path}"
|
|
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
|
|
payload = {
|
|
"enable_sync_mode": enable_sync_mode,
|
|
"enable_base64_output": True,
|
|
}
|
|
|
|
if "edit" in model:
|
|
if not images_base64:
|
|
return {"error": "Edit mode requires at least one image."}
|
|
payload["images"] = [f"data:image/png;base64,{b64}" for b64 in images_base64]
|
|
payload["prompt"] = prompt
|
|
else:
|
|
payload["prompt"] = prompt
|
|
|
|
if size:
|
|
payload["size"] = size
|
|
|
|
if seed is not None and seed != -1:
|
|
payload["seed"] = seed
|
|
|
|
if "sequential" in model:
|
|
payload["max_images"] = max_images
|
|
|
|
logger.info(f"WaveSpeed Seedream Request: URL={submit_url}, Model={model}")
|
|
|
|
request_id = None
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
try:
|
|
response = await client.post(submit_url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
if result.get("data") and result["data"].get("id"):
|
|
request_id = result["data"]["id"]
|
|
logger.info(f"Task submitted successfully. Request ID: {request_id}")
|
|
else:
|
|
return {"error": "Unexpected submission response format.", "details": result}
|
|
except httpx.HTTPStatusError as e:
|
|
return {"error": f"HTTP {e.response.status_code}", "details": e.response.text}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
if not request_id:
|
|
return {"error": "Failed to get a request ID from WaveSpeedAI."}
|
|
|
|
result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
|
|
polling_headers = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
for attempt in range(120):
|
|
try:
|
|
await asyncio.sleep(1)
|
|
poll_response = await client.get(result_url, headers=polling_headers)
|
|
poll_response.raise_for_status()
|
|
|
|
poll_result = poll_response.json().get("data", {})
|
|
status = poll_result.get("status")
|
|
|
|
if status == "completed":
|
|
logger.info("WaveSpeedAI Seedream task completed.")
|
|
outputs = poll_result.get("outputs", [])
|
|
if not outputs:
|
|
return {"error": "Task completed but no outputs found."}
|
|
|
|
b64_jsons = []
|
|
for out in outputs:
|
|
if "base64," in out:
|
|
b64_jsons.append(out.split("base64,")[1])
|
|
else:
|
|
b64_jsons.append(out)
|
|
|
|
return {"data": [{"b64_json": b64} for b64 in b64_jsons]}
|
|
|
|
elif status == "failed":
|
|
error_msg = poll_result.get("error", "Unknown error.")
|
|
return {"error": f"Task failed: {error_msg}"}
|
|
except httpx.HTTPStatusError as e:
|
|
return {"error": f"HTTP {e.response.status_code} while polling", "details": e.response.text}
|
|
except Exception as e:
|
|
return {"error": f"Polling failed: {str(e)}"}
|
|
|
|
return {"error": "Polling timed out after 120 seconds."}
|
|
|
|
|
|
async def send_wavespeed_hunyuan_request(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: str,
|
|
size: Optional[str] = None,
|
|
seed: int = -1,
|
|
output_format: Optional[str] = None,
|
|
enable_sync_mode: bool = False,
|
|
) -> Dict[str, Any]:
|
|
"""Submit and poll a WaveSpeed Hunyuan Image 3 task."""
|
|
if not api_key:
|
|
raise ValueError("WaveSpeed API key is required.")
|
|
if not prompt:
|
|
raise ValueError("Prompt is required for Hunyuan Image generation.")
|
|
|
|
def _normalize_size(value: Optional[str]) -> Optional[str]:
|
|
if not value:
|
|
return None
|
|
normalized = value.replace('x', '*').replace('X', '*')
|
|
return normalized
|
|
|
|
payload: Dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"enable_base64_output": True,
|
|
"enable_sync_mode": enable_sync_mode,
|
|
}
|
|
normalized_size = _normalize_size(size)
|
|
if normalized_size:
|
|
payload["size"] = normalized_size
|
|
if seed is not None and seed != -1:
|
|
payload["seed"] = seed
|
|
if output_format:
|
|
payload["output_format"] = output_format
|
|
|
|
submit_url = f"https://api.wavespeed.ai/api/v3/{model}"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
|
|
request_id = None
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
try:
|
|
response = await client.post(submit_url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
data = result.get("data", {})
|
|
request_id = data.get("id")
|
|
if not request_id:
|
|
return {"error": "Unexpected submission response format.", "details": result}
|
|
logger.info("WaveSpeed Hunyuan task submitted. Request ID: %s", request_id)
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error("WaveSpeed Hunyuan submission failed: %s - %s", e.response.status_code, e.response.text)
|
|
return {"error": f"HTTP {e.response.status_code}", "details": e.response.text}
|
|
except Exception as exc:
|
|
logger.error("Failed to submit WaveSpeed Hunyuan request", exc_info=True)
|
|
return {"error": str(exc)}
|
|
|
|
result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
|
|
polling_headers = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
for attempt in range(120):
|
|
try:
|
|
await asyncio.sleep(1)
|
|
poll_response = await client.get(result_url, headers=polling_headers)
|
|
poll_response.raise_for_status()
|
|
poll_result = poll_response.json().get("data", {})
|
|
status = poll_result.get("status")
|
|
|
|
if status == "completed":
|
|
outputs = poll_result.get("outputs", [])
|
|
if not outputs:
|
|
return {"error": "Task completed but no outputs found."}
|
|
b64_values = []
|
|
for output in outputs:
|
|
if "base64," in output:
|
|
b64_values.append(output.split("base64,")[1])
|
|
else:
|
|
b64_values.append(output)
|
|
return {"data": [{"b64_json": value} for value in b64_values]}
|
|
if status == "failed":
|
|
error_msg = poll_result.get("error", "Unknown error.")
|
|
return {"error": f"Task failed: {error_msg}"}
|
|
except httpx.HTTPStatusError as e:
|
|
return {"error": f"HTTP {e.response.status_code} while polling", "details": e.response.text}
|
|
except Exception as exc:
|
|
logger.error("Error polling WaveSpeed Hunyuan request", exc_info=True)
|
|
return {"error": f"Polling failed: {str(exc)}"}
|
|
|
|
return {"error": "Polling timed out after 120 seconds."}
|
|
|
|
|
|
async def send_wavespeed_qwen_edit_plus_request(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: str,
|
|
images_base64: list[str],
|
|
size: Optional[str] = None,
|
|
seed: int = -1,
|
|
output_format: Optional[str] = None,
|
|
enable_sync_mode: bool = False,
|
|
) -> Dict[str, Any]:
|
|
"""Submit and poll a WaveSpeed Qwen Image Edit Plus task."""
|
|
if not api_key:
|
|
raise ValueError("WaveSpeed API key is required.")
|
|
if not images_base64:
|
|
raise ValueError("At least one reference image is required for Qwen Image Edit Plus.")
|
|
|
|
def _normalize_size(value: Optional[str]) -> Optional[str]:
|
|
if not value:
|
|
return None
|
|
return value.replace('x', '*').replace('X', '*')
|
|
|
|
payload: Dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"images": [f"data:image/png;base64,{b64}" for b64 in images_base64],
|
|
"enable_base64_output": True,
|
|
"enable_sync_mode": enable_sync_mode,
|
|
}
|
|
normalized_size = _normalize_size(size)
|
|
if normalized_size:
|
|
payload["size"] = normalized_size
|
|
if seed is not None and seed != -1:
|
|
payload["seed"] = seed
|
|
if output_format:
|
|
payload["output_format"] = output_format
|
|
|
|
submit_url = f"https://api.wavespeed.ai/api/v3/{model}"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
|
|
request_id = None
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
try:
|
|
response = await client.post(submit_url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
data = result.get("data", {})
|
|
request_id = data.get("id")
|
|
if not request_id:
|
|
return {"error": "Unexpected submission response format.", "details": result}
|
|
logger.info("WaveSpeed Qwen Edit Plus task submitted. Request ID: %s", request_id)
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error("WaveSpeed Qwen submission failed: %s - %s", e.response.status_code, e.response.text)
|
|
return {"error": f"HTTP {e.response.status_code}", "details": e.response.text}
|
|
except Exception as exc:
|
|
logger.error("Failed to submit WaveSpeed Qwen request", exc_info=True)
|
|
return {"error": str(exc)}
|
|
|
|
result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
|
|
polling_headers = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
for attempt in range(120):
|
|
try:
|
|
await asyncio.sleep(1)
|
|
poll_response = await client.get(result_url, headers=polling_headers)
|
|
poll_response.raise_for_status()
|
|
poll_result = poll_response.json().get("data", {})
|
|
status = poll_result.get("status")
|
|
|
|
if status == "completed":
|
|
outputs = poll_result.get("outputs", [])
|
|
if not outputs:
|
|
return {"error": "Task completed but no outputs found."}
|
|
b64_values = []
|
|
for output in outputs:
|
|
if "base64," in output:
|
|
b64_values.append(output.split("base64,")[1])
|
|
else:
|
|
b64_values.append(output)
|
|
return {"data": [{"b64_json": value} for value in b64_values]}
|
|
if status == "failed":
|
|
error_msg = poll_result.get("error", "Unknown error.")
|
|
return {"error": f"Task failed: {error_msg}"}
|
|
except httpx.HTTPStatusError as e:
|
|
return {"error": f"HTTP {e.response.status_code} while polling", "details": e.response.text}
|
|
except Exception as exc:
|
|
logger.error("Error polling WaveSpeed Qwen request", exc_info=True)
|
|
return {"error": f"Polling failed: {str(exc)}"}
|
|
|
|
return {"error": "Polling timed out after 120 seconds."}
|
|
|
|
|
|
async def send_wavespeed_imagen4_request(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: str,
|
|
aspect_ratio: Optional[str] = None,
|
|
resolution: Optional[str] = None,
|
|
num_images: int = 1,
|
|
seed: int = -1,
|
|
negative_prompt: Optional[str] = None,
|
|
enable_sync_mode: bool = False,
|
|
) -> Dict[str, Any]:
|
|
"""Submit and poll a WaveSpeed Imagen 4 generation task."""
|
|
|
|
if not api_key:
|
|
raise ValueError("WaveSpeed API key is required.")
|
|
if not prompt:
|
|
raise ValueError("Prompt is required for Imagen 4 generation.")
|
|
|
|
payload: Dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"num_images": max(1, min(4, num_images)),
|
|
"enable_base64_output": True,
|
|
"enable_sync_mode": enable_sync_mode,
|
|
}
|
|
|
|
if aspect_ratio and aspect_ratio.lower() != "auto":
|
|
payload["aspect_ratio"] = aspect_ratio
|
|
if resolution and resolution.lower() != "auto":
|
|
payload["resolution"] = resolution.lower()
|
|
if negative_prompt:
|
|
payload["negative_prompt"] = negative_prompt
|
|
if seed is not None and seed != -1:
|
|
payload["seed"] = seed
|
|
|
|
submit_url = f"https://api.wavespeed.ai/api/v3/{model}"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
|
|
request_id = None
|
|
async with httpx.AsyncClient(timeout=90.0) as client:
|
|
try:
|
|
response = await client.post(submit_url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
data = result.get("data", {})
|
|
request_id = data.get("id")
|
|
if not request_id:
|
|
return {"error": "Unexpected submission response format.", "details": result}
|
|
logger.info("WaveSpeed Imagen4 task submitted. Request ID: %s", request_id)
|
|
except httpx.HTTPStatusError as exc:
|
|
logger.error("WaveSpeed Imagen4 submission failed: %s - %s", exc.response.status_code, exc.response.text)
|
|
return {"error": f"HTTP {exc.response.status_code}", "details": exc.response.text}
|
|
except Exception as exc:
|
|
logger.error("Failed to submit WaveSpeed Imagen4 request", exc_info=True)
|
|
return {"error": str(exc)}
|
|
|
|
result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
|
|
polling_headers = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
for attempt in range(150):
|
|
try:
|
|
await asyncio.sleep(1)
|
|
poll_response = await client.get(result_url, headers=polling_headers)
|
|
poll_response.raise_for_status()
|
|
poll_result = poll_response.json().get("data", {})
|
|
status = poll_result.get("status")
|
|
|
|
if status == "completed":
|
|
outputs = poll_result.get("outputs", [])
|
|
if not outputs:
|
|
return {"error": "Task completed but no outputs found."}
|
|
images = []
|
|
for output in outputs:
|
|
if isinstance(output, str) and "base64," in output:
|
|
images.append({"b64_json": output.split("base64,")[1]})
|
|
elif isinstance(output, str) and output.strip():
|
|
images.append({"url": output})
|
|
if images:
|
|
return {"data": images}
|
|
return {"error": "Completed task but no usable outputs returned."}
|
|
|
|
if status == "failed":
|
|
error_msg = poll_result.get("error", "Unknown error.")
|
|
return {"error": f"Task failed: {error_msg}"}
|
|
except httpx.HTTPStatusError as exc:
|
|
return {"error": f"HTTP {exc.response.status_code} while polling", "details": exc.response.text}
|
|
except Exception as exc:
|
|
logger.error("Error polling WaveSpeed Imagen4 request", exc_info=True)
|
|
return {"error": f"Polling failed: {str(exc)}"}
|
|
|
|
return {"error": "Polling timed out after 150 seconds."}
|
|
|
|
|
|
async def send_wavespeed_dreamina_request(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: str,
|
|
size: Optional[str] = None,
|
|
seed: int = -1,
|
|
prompt_expansion: bool = True,
|
|
enable_sync_mode: bool = False,
|
|
) -> Dict[str, Any]:
|
|
"""Submit and poll a WaveSpeed Dreamina V3.1 generation task."""
|
|
|
|
if not api_key:
|
|
raise ValueError("WaveSpeed API key is required.")
|
|
if not prompt:
|
|
raise ValueError("Prompt is required for Dreamina generation.")
|
|
|
|
def _normalize_size(value: Optional[str]) -> Optional[str]:
|
|
if not value:
|
|
return None
|
|
return value.replace("x", "*").replace("X", "*")
|
|
|
|
payload: Dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"enable_prompt_expansion": bool(prompt_expansion),
|
|
"enable_base64_output": True,
|
|
"enable_sync_mode": enable_sync_mode,
|
|
}
|
|
normalized_size = _normalize_size(size)
|
|
if normalized_size:
|
|
payload["size"] = normalized_size
|
|
if seed is not None and seed != -1:
|
|
payload["seed"] = seed
|
|
|
|
submit_url = f"https://api.wavespeed.ai/api/v3/{model}"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
|
|
request_id = None
|
|
async with httpx.AsyncClient(timeout=90.0) as client:
|
|
try:
|
|
response = await client.post(submit_url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
data = result.get("data", {})
|
|
request_id = data.get("id")
|
|
if not request_id:
|
|
return {"error": "Unexpected submission response format.", "details": result}
|
|
logger.info("WaveSpeed Dreamina task submitted. Request ID: %s", request_id)
|
|
except httpx.HTTPStatusError as exc:
|
|
logger.error("WaveSpeed Dreamina submission failed: %s - %s", exc.response.status_code, exc.response.text)
|
|
return {"error": f"HTTP {exc.response.status_code}", "details": exc.response.text}
|
|
except Exception as exc:
|
|
logger.error("Failed to submit WaveSpeed Dreamina request", exc_info=True)
|
|
return {"error": str(exc)}
|
|
|
|
result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
|
|
polling_headers = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
for attempt in range(150):
|
|
try:
|
|
await asyncio.sleep(1)
|
|
poll_response = await client.get(result_url, headers=polling_headers)
|
|
poll_response.raise_for_status()
|
|
poll_result = poll_response.json().get("data", {})
|
|
status = poll_result.get("status")
|
|
|
|
if status == "completed":
|
|
outputs = poll_result.get("outputs", [])
|
|
if not outputs:
|
|
return {"error": "Task completed but no outputs found."}
|
|
images = []
|
|
for output in outputs:
|
|
if isinstance(output, str) and "base64," in output:
|
|
images.append({"b64_json": output.split("base64,")[1]})
|
|
elif isinstance(output, str) and output.strip():
|
|
images.append({"url": output})
|
|
if images:
|
|
return {"data": images}
|
|
return {"error": "Completed task but no usable outputs returned."}
|
|
|
|
if status == "failed":
|
|
error_msg = poll_result.get("error", "Unknown error.")
|
|
return {"error": f"Task failed: {error_msg}"}
|
|
except httpx.HTTPStatusError as exc:
|
|
return {"error": f"HTTP {exc.response.status_code} while polling", "details": exc.response.text}
|
|
except Exception as exc:
|
|
logger.error("Error polling WaveSpeed Dreamina request", exc_info=True)
|
|
return {"error": f"Polling failed: {str(exc)}"}
|
|
|
|
return {"error": "Polling timed out after 150 seconds."}
|
|
|
|
|
|
async def send_wavespeed_nano_banana_edit_request(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: Optional[str],
|
|
image_base64: str,
|
|
output_format: str = "png",
|
|
enable_sync_mode: bool = False,
|
|
) -> Dict[str, Any]:
|
|
"""Submit and poll a WaveSpeed Nano Banana edit task."""
|
|
|
|
if not api_key:
|
|
raise ValueError("WaveSpeed API key is required.")
|
|
if not image_base64:
|
|
raise ValueError("An input image is required for Nano Banana edit requests.")
|
|
|
|
payload: Dict[str, Any] = {
|
|
"image": f"data:image/png;base64,{image_base64}",
|
|
"output_format": (output_format or "png").lower(),
|
|
"enable_base64_output": True,
|
|
"enable_sync_mode": enable_sync_mode,
|
|
}
|
|
if prompt:
|
|
payload["prompt"] = prompt
|
|
|
|
submit_url = f"https://api.wavespeed.ai/api/v3/{model}"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
|
|
request_id = None
|
|
async with httpx.AsyncClient(timeout=90.0) as client:
|
|
try:
|
|
response = await client.post(submit_url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
data = result.get("data", {})
|
|
request_id = data.get("id")
|
|
if not request_id:
|
|
return {"error": "Unexpected submission response format.", "details": result}
|
|
logger.info("WaveSpeed Nano Banana task submitted. Request ID: %s", request_id)
|
|
except httpx.HTTPStatusError as exc:
|
|
logger.error("WaveSpeed Nano Banana submission failed: %s - %s", exc.response.status_code, exc.response.text)
|
|
return {"error": f"HTTP {exc.response.status_code}", "details": exc.response.text}
|
|
except Exception as exc:
|
|
logger.error("Failed to submit WaveSpeed Nano Banana request", exc_info=True)
|
|
return {"error": str(exc)}
|
|
|
|
result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
|
|
polling_headers = {"Authorization": f"Bearer {api_key}"}
|
|
|
|
for attempt in range(150):
|
|
try:
|
|
await asyncio.sleep(1)
|
|
poll_response = await client.get(result_url, headers=polling_headers)
|
|
poll_response.raise_for_status()
|
|
poll_result = poll_response.json().get("data", {})
|
|
status = poll_result.get("status")
|
|
|
|
if status == "completed":
|
|
outputs = poll_result.get("outputs", [])
|
|
if not outputs:
|
|
return {"error": "Task completed but no outputs found."}
|
|
images = []
|
|
for output in outputs:
|
|
if isinstance(output, str) and "base64," in output:
|
|
images.append({"b64_json": output.split("base64,")[1]})
|
|
elif isinstance(output, str) and output.strip():
|
|
images.append({"url": output})
|
|
if images:
|
|
return {"data": images}
|
|
return {"error": "Completed task but no usable outputs returned."}
|
|
|
|
if status == "failed":
|
|
error_msg = poll_result.get("error", "Unknown error.")
|
|
return {"error": f"Task failed: {error_msg}"}
|
|
except httpx.HTTPStatusError as exc:
|
|
return {"error": f"HTTP {exc.response.status_code} while polling", "details": exc.response.text}
|
|
except Exception as exc:
|
|
logger.error("Error polling WaveSpeed Nano Banana request", exc_info=True)
|
|
return {"error": f"Polling failed: {str(exc)}"}
|
|
|
|
return {"error": "Polling timed out after 150 seconds."}
|