230 lines
9.8 KiB
Python
230 lines
9.8 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."} |