WaveSpeedAI Python SDK

Official Python SDK for the WaveSpeedAI inference platform

🌐 Visit wavespeed.ai • 📖 Documentation • 💬 Issues

Installation

pip install wavespeed

API Client

import wavespeed
output = wavespeed.run(
 "wavespeed-ai/z-image/turbo",
 {"prompt": "Cat"},
)
print(output["outputs"][0]) # Output URL

Authentication

export WAVESPEED_API_KEY="your-api-key"

from wavespeed import Client
client = Client(api_key="your-api-key")
output = client.run("wavespeed-ai/z-image/turbo", {"prompt": "Cat"})

Options

output = wavespeed.run(
 "wavespeed-ai/z-image/turbo",
 {"prompt": "Cat"},
 timeout=36000.0, # Max wait time in seconds (default: 36000.0)
 poll_interval=1.0, # Status check interval (default: 1.0)
 enable_sync_mode=False, # Best-effort sync result attempt (default: False)
)

Sync Mode

output = wavespeed.run(
 "wavespeed-ai/z-image/turbo",
 {"prompt": "Cat"},
 enable_sync_mode=True,
)

Retry Configuration

from wavespeed import Client
client = Client(
 api_key="your-api-key",
 max_retries=0, # Replacement task attempts (default: 0)
 max_connection_retries=5, # Result-query GET retries; POST is never retried
 retry_interval=1.0, # Base delay between retries in seconds (default: 1.0)
)

Upload Files

import wavespeed
url = wavespeed.upload("/path/to/image.png")
print(url)

Serverless Worker

Basic Handler

import wavespeed.serverless as serverless
def handler(job):
 job_input = job["input"]
 result = job_input.get("prompt", "").upper()
 return {"output": result}
serverless.start({"handler": handler})

Async Handler

import wavespeed.serverless as serverless
async def handler(job):
 job_input = job["input"]
 result = await process_async(job_input)
 return {"output": result}
serverless.start({"handler": handler})

Generator Handler (Streaming)

import wavespeed.serverless as serverless
def handler(job):
 for i in range(10):
 yield {"progress": i, "partial": f"chunk-{i}"}
serverless.start({"handler": handler})

Input Validation

from wavespeed.serverless.utils import validate
INPUT_SCHEMA = {
 "prompt": {"type": str, "required": True},
 "max_tokens": {"type": int, "required": False, "default": 100},
 "temperature": {
 "type": float,
 "required": False,
 "default": 0.7,
 "constraints": lambda x: 0 <= x <= 2,
 },
}
def handler(job):
 result = validate(job["input"], INPUT_SCHEMA)
 if "errors" in result:
 return {"error": result["errors"]}
 validated = result["validated_input"]
 # process with validated input...
 return {"output": "done"}

Concurrent Execution

import wavespeed.serverless as serverless
def handler(job):
 return {"output": job["input"]["data"]}
def concurrency_modifier(current_concurrency):
 return 2 # Process 2 jobs concurrently
serverless.start({
 "handler": handler,
 "concurrency_modifier": concurrency_modifier
})

Local Development

Test with JSON Input

# Using CLI argument
python handler.py --test_input '{"input": {"prompt": "hello"}}'
# Using test_input.json file (auto-detected)
echo '{"input": {"prompt": "hello"}}' > test_input.json
python handler.py

Running Tests

# Run all tests
python -m pytest
# Run a single test file
python -m pytest tests/test_api.py
# Run a specific test
python -m pytest tests/test_api.py::TestClient::test_run_success -v

FastAPI Development Server

python handler.py --waverless_serve_api --waverless_api_port 8000

# Synchronous execution
curl -X POST http://localhost:8000/runsync \
 -H "Content-Type: application/json" \
 -d '{"input": {"prompt": "hello"}}'
# Async execution
curl -X POST http://localhost:8000/run \
 -H "Content-Type: application/json" \
 -d '{"input": {"prompt": "hello"}}'

CLI Options

Environment Variables

API Client

Serverless Worker

License