| import os | |
| import time | |
| import json | |
| import requests | |
| import logging | |
| from pathlib import Path | |
| logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(message)s') | |
| logger = logging.getLogger("wavespeed_client") | |
| class WavespeedClient: | |
| def __init__(self, api_key=None): | |
| self.api_key = api_key or os.environ.get("WAVESPEED_API_KEY") | |
| if not self.api_key: | |
| raise ValueError("WAVESPEED_API_KEY not found in environment.") | |
| self.base_url = "https://api.wavespeed.ai/api/v3" | |
| self.headers = { | |
| "Authorization": f"Bearer {self.api_key}", | |
| "Content-Type": "application/json" | |
| } | |
| def submit_task(self, model_endpoint, payload): | |
| """Submits a generation task to a specific model endpoint.""" | |
| logger.info(f"Submitting task to {model_endpoint}...") | |
| url = f"{self.base_url}/{model_endpoint}" | |
| response = requests.post(url, headers=self.headers, json=payload) | |
| response.raise_for_status() | |
| result = response.json() | |
| task_id = result["data"]["id"] | |
| logger.info(f"Task submitted. ID: {task_id}") | |
| return task_id | |
| def poll_result(self, task_id, interval=10, timeout=600): | |
| """Polls for the result of a task until completion or timeout.""" | |
| logger.info(f"Polling for result of task {task_id}...") | |
| url = f"{self.base_url}/predictions/{task_id}/result" | |
| start_time = time.time() | |
| while time.time() - start_time < timeout: | |
| response = requests.get(url, headers=self.headers) | |
| response.raise_for_status() | |
| result = response.json() | |
| data = result["data"] | |
| status = data.get("status") | |
| if status == "completed": | |
| outputs = data.get("outputs", []) | |
| logger.info(f"Task completed successfully. Outputs: {outputs}") | |
| return outputs | |
| elif status == "failed": | |
| error = data.get("error", "Unknown error") | |
| logger.error(f"Task failed: {error}") | |
| raise Exception(f"Wavespeed task failed: {error}") | |
| time.sleep(interval) | |
| raise TimeoutError(f"Task {task_id} timed out after {timeout} seconds.") | |