WaveSpeedAI Python SDK

Official Python SDK for the WaveSpeedAI inference platform

🌐 Visit wavespeed.ai • 📖 Documentation • 💬 Issues

---

Installation

pip install wavespeed

API Client

Run WaveSpeed AI models with a simple API:

import wavespeed

output = wavespeed.run(
    "wavespeed-ai/z-image/turbo",
    input={"prompt": "A cat holds a sign that says 'Hello World'"},
)

print(output["outputs"][0])  # Output URL

Authentication

Set your API key via environment variable:

export WAVESPEED_API_KEY="your-api-key"

Or pass it directly:

from wavespeed import Client

client = Client(api_key="your-api-key")
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds a sign that says 'Hello World'"})

Options

output = wavespeed.run(
    "wavespeed-ai/z-image/turbo",
    input={"prompt": "A cat holds a sign that says 'Hello World'"},
    timeout=36000.0     # Max wait time in seconds (default: 36000.0)
    poll_interval=1.0,  # Status check interval (default: 1.0)
)

Upload Files

Upload images, videos, or audio files:

import wavespeed

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

Serverless Worker

Build serverless workers for the WaveSpeed platform.

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"}

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

Then use the interactive Swagger UI at http://localhost:8000/ or make requests:

# 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

Option	Description
--test_input JSON	Run locally with JSON test input
--waverless_serve_api	Start FastAPI development server
--waverless_api_host HOST	API server host (default: localhost)
--waverless_api_port PORT	API server port (default: 8000)
--waverless_log_level LEVEL	Log level (DEBUG, INFO, WARN, ERROR)

Environment Variables

API Client

Variable	Description
WAVESPEED_API_KEY	WaveSpeed API key

Serverless Worker

Variable	Description
WAVERLESS_POD_ID	Worker/pod identifier
WAVERLESS_API_KEY	API authentication key
WAVERLESS_WEBHOOK_GET_JOB	Job fetch endpoint
WAVERLESS_WEBHOOK_POST_OUTPUT	Result submission endpoint

License

MIT