wavespeed 0.1.0__tar.gz → 0.1.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/python-publish.yml +2 -2
- {wavespeed-0.1.0 → wavespeed-0.1.2}/PKG-INFO +56 -7
- {wavespeed-0.1.0 → wavespeed-0.1.2}/README.md +55 -6
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/_version.py +1 -1
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/api/__init__.py +20 -0
- wavespeed-0.1.2/src/wavespeed/api/client.py +393 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/config.py +9 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/PKG-INFO +56 -7
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_api.py +13 -1
- wavespeed-0.1.0/src/wavespeed/api/client.py +0 -238
- {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/claude-code-review.yml +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/claude.yml +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/pre-commit.yml +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/.github/workflows/python-packages.yml +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/.gitignore +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/.pre-commit-config.yaml +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/CLAUDE.md +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/LICENSE +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/MANIFEST.in +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/VERSIONING.md +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/images/build_and_push.sh +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/Dockerfile +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/build_and_test.sh +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/handler.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/requirements.txt +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/images/test_worker/test_input.json +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/pyproject.toml +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/pytest.ini +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/setup.cfg +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/__init__.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/_config_module.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/__init__.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/__init__.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/fastapi.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/handler.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/heartbeat.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/http.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/job.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/local.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/logger.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/progress.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/scaler.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/modules/state.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/utils/__init__.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/utils/boto3_utils.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/utils/validator.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed/serverless/worker.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/SOURCES.txt +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/dependency_links.txt +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/not-zip-safe +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/requires.txt +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/src/wavespeed.egg-info/top_level.txt +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/__init__.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/conftest.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_config.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/__init__.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/__init__.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_fastapi.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_handler.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_http.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_job.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_local.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_logger.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_progress.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_scaler.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_modules/test_state.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_utils/__init__.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_utils/test_validator.py +0 -0
- {wavespeed-0.1.0 → wavespeed-0.1.2}/tests/test_serverless/test_worker.py +0 -0
|
@@ -39,8 +39,8 @@ jobs:
|
|
|
39
39
|
- name: Install dependencies
|
|
40
40
|
run: |
|
|
41
41
|
python -m pip install --upgrade pip
|
|
42
|
-
pip install build 'setuptools>=64' 'setuptools_scm>=8'
|
|
42
|
+
pip install build wheel 'setuptools>=64' 'setuptools_scm>=8'
|
|
43
43
|
- name: Build package
|
|
44
|
-
run: python -m build
|
|
44
|
+
run: python -m build
|
|
45
45
|
- name: Publish package
|
|
46
46
|
uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: wavespeed
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.2
|
|
4
4
|
Summary: WaveSpeedAI Python Client — Official Python SDK for WaveSpeedAI inference platform
|
|
5
5
|
Author-email: WaveSpeedAI <support@wavespeed.ai>
|
|
6
6
|
License: MIT
|
|
@@ -62,7 +62,7 @@ import wavespeed
|
|
|
62
62
|
|
|
63
63
|
output = wavespeed.run(
|
|
64
64
|
"wavespeed-ai/z-image/turbo",
|
|
65
|
-
input={"prompt": "
|
|
65
|
+
input={"prompt": "Cat"},
|
|
66
66
|
)
|
|
67
67
|
|
|
68
68
|
print(output["outputs"][0]) # Output URL
|
|
@@ -70,7 +70,7 @@ print(output["outputs"][0]) # Output URL
|
|
|
70
70
|
|
|
71
71
|
### Authentication
|
|
72
72
|
|
|
73
|
-
Set your API key via environment variable:
|
|
73
|
+
Set your API key via environment variable (You can get your API key from [https://wavespeed.ai/accesskey](https://wavespeed.ai/accesskey)):
|
|
74
74
|
|
|
75
75
|
```bash
|
|
76
76
|
export WAVESPEED_API_KEY="your-api-key"
|
|
@@ -82,7 +82,7 @@ Or pass it directly:
|
|
|
82
82
|
from wavespeed import Client
|
|
83
83
|
|
|
84
84
|
client = Client(api_key="your-api-key")
|
|
85
|
-
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "
|
|
85
|
+
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"})
|
|
86
86
|
```
|
|
87
87
|
|
|
88
88
|
### Options
|
|
@@ -90,9 +90,39 @@ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds
|
|
|
90
90
|
```python
|
|
91
91
|
output = wavespeed.run(
|
|
92
92
|
"wavespeed-ai/z-image/turbo",
|
|
93
|
-
input={"prompt": "
|
|
94
|
-
timeout=36000.0
|
|
95
|
-
poll_interval=1.0,
|
|
93
|
+
input={"prompt": "Cat"},
|
|
94
|
+
timeout=36000.0, # Max wait time in seconds (default: 36000.0)
|
|
95
|
+
poll_interval=1.0, # Status check interval (default: 1.0)
|
|
96
|
+
enable_sync_mode=False, # Single request mode, no polling (default: False)
|
|
97
|
+
)
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
### Sync Mode
|
|
101
|
+
|
|
102
|
+
Use `enable_sync_mode=True` for a single request that waits for the result (no polling).
|
|
103
|
+
|
|
104
|
+
> **Note:** Not all models support sync mode. Check the model documentation for availability.
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
output = wavespeed.run(
|
|
108
|
+
"wavespeed-ai/z-image/turbo",
|
|
109
|
+
input={"prompt": "Cat"},
|
|
110
|
+
enable_sync_mode=True,
|
|
111
|
+
)
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
### Retry Configuration
|
|
115
|
+
|
|
116
|
+
Configure retries at the client level:
|
|
117
|
+
|
|
118
|
+
```python
|
|
119
|
+
from wavespeed import Client
|
|
120
|
+
|
|
121
|
+
client = Client(
|
|
122
|
+
api_key="your-api-key",
|
|
123
|
+
max_retries=0, # Task-level retries (default: 0)
|
|
124
|
+
max_connection_retries=5, # HTTP connection retries (default: 5)
|
|
125
|
+
retry_interval=1.0, # Base delay between retries in seconds (default: 1.0)
|
|
96
126
|
)
|
|
97
127
|
```
|
|
98
128
|
|
|
@@ -175,6 +205,25 @@ def handler(job):
|
|
|
175
205
|
return {"output": "done"}
|
|
176
206
|
```
|
|
177
207
|
|
|
208
|
+
### Concurrent Execution
|
|
209
|
+
|
|
210
|
+
Enable concurrent job processing with `concurrency_modifier`:
|
|
211
|
+
|
|
212
|
+
```python
|
|
213
|
+
import wavespeed.serverless as serverless
|
|
214
|
+
|
|
215
|
+
def handler(job):
|
|
216
|
+
return {"output": job["input"]["data"]}
|
|
217
|
+
|
|
218
|
+
def concurrency_modifier(current_concurrency):
|
|
219
|
+
return 2 # Process 2 jobs concurrently
|
|
220
|
+
|
|
221
|
+
serverless.start({
|
|
222
|
+
"handler": handler,
|
|
223
|
+
"concurrency_modifier": concurrency_modifier
|
|
224
|
+
})
|
|
225
|
+
```
|
|
226
|
+
|
|
178
227
|
## Local Development
|
|
179
228
|
|
|
180
229
|
### Test with JSON Input
|
|
@@ -33,7 +33,7 @@ import wavespeed
|
|
|
33
33
|
|
|
34
34
|
output = wavespeed.run(
|
|
35
35
|
"wavespeed-ai/z-image/turbo",
|
|
36
|
-
input={"prompt": "
|
|
36
|
+
input={"prompt": "Cat"},
|
|
37
37
|
)
|
|
38
38
|
|
|
39
39
|
print(output["outputs"][0]) # Output URL
|
|
@@ -41,7 +41,7 @@ print(output["outputs"][0]) # Output URL
|
|
|
41
41
|
|
|
42
42
|
### Authentication
|
|
43
43
|
|
|
44
|
-
Set your API key via environment variable:
|
|
44
|
+
Set your API key via environment variable (You can get your API key from [https://wavespeed.ai/accesskey](https://wavespeed.ai/accesskey)):
|
|
45
45
|
|
|
46
46
|
```bash
|
|
47
47
|
export WAVESPEED_API_KEY="your-api-key"
|
|
@@ -53,7 +53,7 @@ Or pass it directly:
|
|
|
53
53
|
from wavespeed import Client
|
|
54
54
|
|
|
55
55
|
client = Client(api_key="your-api-key")
|
|
56
|
-
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "
|
|
56
|
+
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"})
|
|
57
57
|
```
|
|
58
58
|
|
|
59
59
|
### Options
|
|
@@ -61,9 +61,39 @@ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds
|
|
|
61
61
|
```python
|
|
62
62
|
output = wavespeed.run(
|
|
63
63
|
"wavespeed-ai/z-image/turbo",
|
|
64
|
-
input={"prompt": "
|
|
65
|
-
timeout=36000.0
|
|
66
|
-
poll_interval=1.0,
|
|
64
|
+
input={"prompt": "Cat"},
|
|
65
|
+
timeout=36000.0, # Max wait time in seconds (default: 36000.0)
|
|
66
|
+
poll_interval=1.0, # Status check interval (default: 1.0)
|
|
67
|
+
enable_sync_mode=False, # Single request mode, no polling (default: False)
|
|
68
|
+
)
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
### Sync Mode
|
|
72
|
+
|
|
73
|
+
Use `enable_sync_mode=True` for a single request that waits for the result (no polling).
|
|
74
|
+
|
|
75
|
+
> **Note:** Not all models support sync mode. Check the model documentation for availability.
|
|
76
|
+
|
|
77
|
+
```python
|
|
78
|
+
output = wavespeed.run(
|
|
79
|
+
"wavespeed-ai/z-image/turbo",
|
|
80
|
+
input={"prompt": "Cat"},
|
|
81
|
+
enable_sync_mode=True,
|
|
82
|
+
)
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
### Retry Configuration
|
|
86
|
+
|
|
87
|
+
Configure retries at the client level:
|
|
88
|
+
|
|
89
|
+
```python
|
|
90
|
+
from wavespeed import Client
|
|
91
|
+
|
|
92
|
+
client = Client(
|
|
93
|
+
api_key="your-api-key",
|
|
94
|
+
max_retries=0, # Task-level retries (default: 0)
|
|
95
|
+
max_connection_retries=5, # HTTP connection retries (default: 5)
|
|
96
|
+
retry_interval=1.0, # Base delay between retries in seconds (default: 1.0)
|
|
67
97
|
)
|
|
68
98
|
```
|
|
69
99
|
|
|
@@ -146,6 +176,25 @@ def handler(job):
|
|
|
146
176
|
return {"output": "done"}
|
|
147
177
|
```
|
|
148
178
|
|
|
179
|
+
### Concurrent Execution
|
|
180
|
+
|
|
181
|
+
Enable concurrent job processing with `concurrency_modifier`:
|
|
182
|
+
|
|
183
|
+
```python
|
|
184
|
+
import wavespeed.serverless as serverless
|
|
185
|
+
|
|
186
|
+
def handler(job):
|
|
187
|
+
return {"output": job["input"]["data"]}
|
|
188
|
+
|
|
189
|
+
def concurrency_modifier(current_concurrency):
|
|
190
|
+
return 2 # Process 2 jobs concurrently
|
|
191
|
+
|
|
192
|
+
serverless.start({
|
|
193
|
+
"handler": handler,
|
|
194
|
+
"concurrency_modifier": concurrency_modifier
|
|
195
|
+
})
|
|
196
|
+
```
|
|
197
|
+
|
|
149
198
|
## Local Development
|
|
150
199
|
|
|
151
200
|
### Test with JSON Input
|
|
@@ -1,2 +1,2 @@
|
|
|
1
1
|
# This file is automatically generated by setuptools_scm
|
|
2
|
-
__version__ = "0.1.
|
|
2
|
+
__version__ = "0.1.2"
|
|
@@ -41,6 +41,8 @@ def run(
|
|
|
41
41
|
*,
|
|
42
42
|
timeout: float | None = None,
|
|
43
43
|
poll_interval: float = 1.0,
|
|
44
|
+
enable_sync_mode: bool = False,
|
|
45
|
+
max_retries: int | None = None,
|
|
44
46
|
) -> dict:
|
|
45
47
|
"""Run a model and wait for the output.
|
|
46
48
|
|
|
@@ -49,6 +51,8 @@ def run(
|
|
|
49
51
|
input: Input parameters for the model.
|
|
50
52
|
timeout: Maximum time to wait for completion (None = no timeout).
|
|
51
53
|
poll_interval: Interval between status checks in seconds.
|
|
54
|
+
enable_sync_mode: If True, use synchronous mode (single request).
|
|
55
|
+
max_retries: Maximum retries for this request (overrides default setting).
|
|
52
56
|
|
|
53
57
|
Returns:
|
|
54
58
|
Dict containing "outputs" array with model outputs.
|
|
@@ -64,12 +68,28 @@ def run(
|
|
|
64
68
|
input={"prompt": "A cat sitting on a windowsill"}
|
|
65
69
|
)
|
|
66
70
|
print(output["outputs"][0]) # First output URL
|
|
71
|
+
|
|
72
|
+
# With sync mode
|
|
73
|
+
output = wavespeed.run(
|
|
74
|
+
"wavespeed-ai/z-image/turbo",
|
|
75
|
+
input={"prompt": "A cat"},
|
|
76
|
+
enable_sync_mode=True
|
|
77
|
+
)
|
|
78
|
+
|
|
79
|
+
# With retry
|
|
80
|
+
output = wavespeed.run(
|
|
81
|
+
"wavespeed-ai/z-image/turbo",
|
|
82
|
+
input={"prompt": "A cat"},
|
|
83
|
+
max_retries=3
|
|
84
|
+
)
|
|
67
85
|
"""
|
|
68
86
|
return _get_default_client().run(
|
|
69
87
|
model,
|
|
70
88
|
input=input,
|
|
71
89
|
timeout=timeout,
|
|
72
90
|
poll_interval=poll_interval,
|
|
91
|
+
enable_sync_mode=enable_sync_mode,
|
|
92
|
+
max_retries=max_retries,
|
|
73
93
|
)
|
|
74
94
|
|
|
75
95
|
|
|
@@ -0,0 +1,393 @@
|
|
|
1
|
+
"""WaveSpeed API client implementation."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import time
|
|
5
|
+
import traceback
|
|
6
|
+
from typing import Any, BinaryIO
|
|
7
|
+
|
|
8
|
+
import requests
|
|
9
|
+
|
|
10
|
+
from wavespeed.config import api as api_config
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class Client:
|
|
14
|
+
"""WaveSpeed API client.
|
|
15
|
+
|
|
16
|
+
Args:
|
|
17
|
+
api_key: WaveSpeed API key. If not provided, uses wavespeed.config.api.api_key.
|
|
18
|
+
base_url: Base URL for the API. If not provided, uses wavespeed.config.api.base_url.
|
|
19
|
+
connection_timeout: Timeout for HTTP requests in seconds.
|
|
20
|
+
max_retries: Maximum number of retries for the entire operation.
|
|
21
|
+
max_connection_retries: Maximum retries for individual HTTP requests.
|
|
22
|
+
retry_interval: Base interval between retries in seconds.
|
|
23
|
+
|
|
24
|
+
Example:
|
|
25
|
+
client = Client(api_key="your-api-key")
|
|
26
|
+
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"})
|
|
27
|
+
|
|
28
|
+
# With sync mode (single request, waits for result)
|
|
29
|
+
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"}, enable_sync_mode=True)
|
|
30
|
+
|
|
31
|
+
# With retry
|
|
32
|
+
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"}, max_retries=3)
|
|
33
|
+
"""
|
|
34
|
+
|
|
35
|
+
def __init__(
|
|
36
|
+
self,
|
|
37
|
+
api_key: str | None = None,
|
|
38
|
+
base_url: str | None = None,
|
|
39
|
+
connection_timeout: float | None = None,
|
|
40
|
+
max_retries: int | None = None,
|
|
41
|
+
max_connection_retries: int | None = None,
|
|
42
|
+
retry_interval: float | None = None,
|
|
43
|
+
) -> None:
|
|
44
|
+
"""Initialize the client."""
|
|
45
|
+
self.api_key = api_key or api_config.api_key
|
|
46
|
+
self.base_url = (base_url or api_config.base_url).rstrip("/")
|
|
47
|
+
self.connection_timeout = connection_timeout or api_config.connection_timeout
|
|
48
|
+
self.max_retries = (
|
|
49
|
+
max_retries if max_retries is not None else api_config.max_retries
|
|
50
|
+
)
|
|
51
|
+
self.max_connection_retries = (
|
|
52
|
+
max_connection_retries
|
|
53
|
+
if max_connection_retries is not None
|
|
54
|
+
else api_config.max_connection_retries
|
|
55
|
+
)
|
|
56
|
+
self.retry_interval = (
|
|
57
|
+
retry_interval if retry_interval is not None else api_config.retry_interval
|
|
58
|
+
)
|
|
59
|
+
|
|
60
|
+
def _get_headers(self) -> dict[str, str]:
|
|
61
|
+
"""Get request headers with authentication."""
|
|
62
|
+
if not self.api_key:
|
|
63
|
+
raise ValueError(
|
|
64
|
+
"API key is required. Set WAVESPEED_API_KEY environment variable "
|
|
65
|
+
"or pass api_key to Client()."
|
|
66
|
+
)
|
|
67
|
+
return {
|
|
68
|
+
"Content-Type": "application/json",
|
|
69
|
+
"Authorization": f"Bearer {self.api_key}",
|
|
70
|
+
}
|
|
71
|
+
|
|
72
|
+
def _submit(
|
|
73
|
+
self,
|
|
74
|
+
model: str,
|
|
75
|
+
input: dict[str, Any] | None,
|
|
76
|
+
enable_sync_mode: bool = False,
|
|
77
|
+
timeout: float | None = None,
|
|
78
|
+
) -> tuple[str | None, dict[str, Any] | None]:
|
|
79
|
+
"""Submit a prediction request.
|
|
80
|
+
|
|
81
|
+
Args:
|
|
82
|
+
model: Model identifier.
|
|
83
|
+
input: Input parameters.
|
|
84
|
+
enable_sync_mode: If True, wait for result in single request.
|
|
85
|
+
timeout: Request timeout in seconds.
|
|
86
|
+
|
|
87
|
+
Returns:
|
|
88
|
+
Tuple of (request_id, result). In async mode, result is None.
|
|
89
|
+
In sync mode, request_id is None and result contains the response.
|
|
90
|
+
|
|
91
|
+
Raises:
|
|
92
|
+
RuntimeError: If submission fails after retries.
|
|
93
|
+
"""
|
|
94
|
+
url = f"{self.base_url}/api/v3/{model}"
|
|
95
|
+
body = dict(input) if input else {}
|
|
96
|
+
|
|
97
|
+
if enable_sync_mode:
|
|
98
|
+
body["enable_sync_mode"] = True
|
|
99
|
+
|
|
100
|
+
request_timeout = timeout if timeout is not None else api_config.timeout
|
|
101
|
+
# Use connection timeout for connect, request_timeout for read
|
|
102
|
+
connect_timeout = (
|
|
103
|
+
min(self.connection_timeout, request_timeout)
|
|
104
|
+
if request_timeout
|
|
105
|
+
else self.connection_timeout
|
|
106
|
+
)
|
|
107
|
+
timeouts = (connect_timeout, request_timeout)
|
|
108
|
+
|
|
109
|
+
for retry in range(self.max_connection_retries + 1):
|
|
110
|
+
try:
|
|
111
|
+
response = requests.post(
|
|
112
|
+
url, json=body, headers=self._get_headers(), timeout=timeouts
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
if response.status_code != 200:
|
|
116
|
+
raise RuntimeError(
|
|
117
|
+
f"Failed to submit prediction: HTTP {response.status_code}: "
|
|
118
|
+
f"{response.text}"
|
|
119
|
+
)
|
|
120
|
+
|
|
121
|
+
result = response.json()
|
|
122
|
+
|
|
123
|
+
if enable_sync_mode:
|
|
124
|
+
return None, result
|
|
125
|
+
|
|
126
|
+
request_id = result.get("data", {}).get("id")
|
|
127
|
+
if not request_id:
|
|
128
|
+
raise RuntimeError(f"No request ID in response: {result}")
|
|
129
|
+
|
|
130
|
+
return request_id, None
|
|
131
|
+
|
|
132
|
+
except (
|
|
133
|
+
requests.exceptions.ConnectionError,
|
|
134
|
+
requests.exceptions.Timeout,
|
|
135
|
+
) as e:
|
|
136
|
+
print(
|
|
137
|
+
f"Connection error on attempt {retry + 1}/{self.max_connection_retries + 1}:"
|
|
138
|
+
)
|
|
139
|
+
traceback.print_exc()
|
|
140
|
+
|
|
141
|
+
if retry < self.max_connection_retries:
|
|
142
|
+
delay = self.retry_interval * (retry + 1)
|
|
143
|
+
print(f"Retrying in {delay} seconds...")
|
|
144
|
+
time.sleep(delay)
|
|
145
|
+
else:
|
|
146
|
+
raise RuntimeError(
|
|
147
|
+
f"Failed to submit prediction after {self.max_connection_retries + 1} attempts"
|
|
148
|
+
) from e
|
|
149
|
+
|
|
150
|
+
def _get_result(
|
|
151
|
+
self, request_id: str, timeout: float | None = None
|
|
152
|
+
) -> dict[str, Any]:
|
|
153
|
+
"""Get prediction result.
|
|
154
|
+
|
|
155
|
+
Args:
|
|
156
|
+
request_id: The prediction request ID.
|
|
157
|
+
timeout: Request timeout in seconds.
|
|
158
|
+
|
|
159
|
+
Returns:
|
|
160
|
+
Full API response.
|
|
161
|
+
|
|
162
|
+
Raises:
|
|
163
|
+
RuntimeError: If fetching result fails after retries.
|
|
164
|
+
"""
|
|
165
|
+
url = f"{self.base_url}/api/v3/predictions/{request_id}/result"
|
|
166
|
+
request_timeout = timeout if timeout is not None else api_config.timeout
|
|
167
|
+
connect_timeout = (
|
|
168
|
+
min(self.connection_timeout, request_timeout)
|
|
169
|
+
if request_timeout
|
|
170
|
+
else self.connection_timeout
|
|
171
|
+
)
|
|
172
|
+
timeouts = (connect_timeout, request_timeout)
|
|
173
|
+
|
|
174
|
+
for retry in range(self.max_connection_retries + 1):
|
|
175
|
+
try:
|
|
176
|
+
response = requests.get(
|
|
177
|
+
url, headers=self._get_headers(), timeout=timeouts
|
|
178
|
+
)
|
|
179
|
+
|
|
180
|
+
if response.status_code != 200:
|
|
181
|
+
raise RuntimeError(
|
|
182
|
+
f"Failed to get result: HTTP {response.status_code}: {response.text}"
|
|
183
|
+
)
|
|
184
|
+
|
|
185
|
+
return response.json()
|
|
186
|
+
|
|
187
|
+
except (
|
|
188
|
+
requests.exceptions.ConnectionError,
|
|
189
|
+
requests.exceptions.Timeout,
|
|
190
|
+
) as e:
|
|
191
|
+
print(
|
|
192
|
+
f"Connection error getting result on attempt {retry + 1}/{self.max_connection_retries + 1}:"
|
|
193
|
+
)
|
|
194
|
+
traceback.print_exc()
|
|
195
|
+
|
|
196
|
+
if retry < self.max_connection_retries:
|
|
197
|
+
delay = self.retry_interval * (retry + 1)
|
|
198
|
+
print(f"Retrying in {delay} seconds...")
|
|
199
|
+
time.sleep(delay)
|
|
200
|
+
else:
|
|
201
|
+
raise RuntimeError(
|
|
202
|
+
f"Failed to get result after {self.max_connection_retries + 1} attempts"
|
|
203
|
+
) from e
|
|
204
|
+
|
|
205
|
+
def _wait(
|
|
206
|
+
self,
|
|
207
|
+
request_id: str,
|
|
208
|
+
timeout: float | None,
|
|
209
|
+
poll_interval: float,
|
|
210
|
+
) -> dict[str, Any]:
|
|
211
|
+
"""Wait for prediction to complete.
|
|
212
|
+
|
|
213
|
+
Args:
|
|
214
|
+
request_id: The prediction request ID.
|
|
215
|
+
timeout: Maximum wait time in seconds (None = no timeout).
|
|
216
|
+
poll_interval: Time between polls in seconds.
|
|
217
|
+
|
|
218
|
+
Returns:
|
|
219
|
+
Dict with "outputs" array.
|
|
220
|
+
|
|
221
|
+
Raises:
|
|
222
|
+
RuntimeError: If prediction fails.
|
|
223
|
+
TimeoutError: If prediction times out.
|
|
224
|
+
"""
|
|
225
|
+
start_time = time.time()
|
|
226
|
+
|
|
227
|
+
while True:
|
|
228
|
+
# Check timeout
|
|
229
|
+
if timeout is not None:
|
|
230
|
+
elapsed = time.time() - start_time
|
|
231
|
+
if elapsed >= timeout:
|
|
232
|
+
raise TimeoutError(f"Prediction timed out after {timeout} seconds")
|
|
233
|
+
|
|
234
|
+
result = self._get_result(request_id, timeout=timeout)
|
|
235
|
+
data = result.get("data", {})
|
|
236
|
+
status = data.get("status")
|
|
237
|
+
|
|
238
|
+
if status == "completed":
|
|
239
|
+
return {"outputs": data.get("outputs", [])}
|
|
240
|
+
|
|
241
|
+
if status == "failed":
|
|
242
|
+
error = data.get("error") or "Unknown error"
|
|
243
|
+
raise RuntimeError(f"Prediction failed: {error}")
|
|
244
|
+
|
|
245
|
+
time.sleep(poll_interval)
|
|
246
|
+
|
|
247
|
+
def _is_retryable_error(self, error: Exception) -> bool:
|
|
248
|
+
"""Determine if an error is worth retrying at the task level.
|
|
249
|
+
|
|
250
|
+
Args:
|
|
251
|
+
error: The exception to check.
|
|
252
|
+
|
|
253
|
+
Returns:
|
|
254
|
+
True if the error is retryable.
|
|
255
|
+
"""
|
|
256
|
+
# Always retry timeout and connection errors
|
|
257
|
+
if isinstance(
|
|
258
|
+
error, (requests.exceptions.Timeout, requests.exceptions.ConnectionError)
|
|
259
|
+
):
|
|
260
|
+
return True
|
|
261
|
+
|
|
262
|
+
# Retry server errors (5xx) and rate limiting (429)
|
|
263
|
+
if isinstance(error, RuntimeError):
|
|
264
|
+
error_str = str(error)
|
|
265
|
+
if "HTTP 5" in error_str or "HTTP 429" in error_str:
|
|
266
|
+
return True
|
|
267
|
+
|
|
268
|
+
return False
|
|
269
|
+
|
|
270
|
+
def run(
|
|
271
|
+
self,
|
|
272
|
+
model: str,
|
|
273
|
+
input: dict[str, Any] | None = None,
|
|
274
|
+
*,
|
|
275
|
+
timeout: float | None = None,
|
|
276
|
+
poll_interval: float = 1.0,
|
|
277
|
+
enable_sync_mode: bool = False,
|
|
278
|
+
max_retries: int | None = None,
|
|
279
|
+
) -> dict[str, Any]:
|
|
280
|
+
"""Run a model and wait for the output.
|
|
281
|
+
|
|
282
|
+
Args:
|
|
283
|
+
model: Model identifier (e.g., "wavespeed-ai/flux-dev").
|
|
284
|
+
input: Input parameters for the model.
|
|
285
|
+
timeout: Maximum time to wait for completion (None = no timeout).
|
|
286
|
+
poll_interval: Interval between status checks in seconds.
|
|
287
|
+
enable_sync_mode: If True, use synchronous mode (single request).
|
|
288
|
+
max_retries: Maximum task-level retries (overrides client setting).
|
|
289
|
+
|
|
290
|
+
Returns:
|
|
291
|
+
Dict containing "outputs" array with model outputs.
|
|
292
|
+
|
|
293
|
+
Raises:
|
|
294
|
+
ValueError: If API key is not configured.
|
|
295
|
+
RuntimeError: If the prediction fails.
|
|
296
|
+
TimeoutError: If the prediction times out.
|
|
297
|
+
"""
|
|
298
|
+
task_retries = max_retries if max_retries is not None else self.max_retries
|
|
299
|
+
last_error = None
|
|
300
|
+
|
|
301
|
+
for attempt in range(task_retries + 1):
|
|
302
|
+
try:
|
|
303
|
+
request_id, sync_result = self._submit(
|
|
304
|
+
model, input, enable_sync_mode=enable_sync_mode, timeout=timeout
|
|
305
|
+
)
|
|
306
|
+
|
|
307
|
+
if enable_sync_mode:
|
|
308
|
+
# In sync mode, extract outputs from the result
|
|
309
|
+
data = sync_result.get("data", {})
|
|
310
|
+
return {"outputs": data.get("outputs", [])}
|
|
311
|
+
|
|
312
|
+
return self._wait(request_id, timeout, poll_interval)
|
|
313
|
+
|
|
314
|
+
except Exception as e:
|
|
315
|
+
last_error = e
|
|
316
|
+
is_retryable = self._is_retryable_error(e)
|
|
317
|
+
|
|
318
|
+
if not is_retryable or attempt >= task_retries:
|
|
319
|
+
raise
|
|
320
|
+
|
|
321
|
+
print(f"Task attempt {attempt + 1}/{task_retries + 1} failed: {e}")
|
|
322
|
+
delay = self.retry_interval * (attempt + 1)
|
|
323
|
+
print(f"Retrying in {delay} seconds...")
|
|
324
|
+
time.sleep(delay)
|
|
325
|
+
|
|
326
|
+
# Should not reach here, but just in case
|
|
327
|
+
if last_error:
|
|
328
|
+
raise last_error
|
|
329
|
+
raise RuntimeError(f"All {task_retries + 1} attempts failed")
|
|
330
|
+
|
|
331
|
+
def upload(self, file: str | BinaryIO, *, timeout: float | None = None) -> str:
|
|
332
|
+
"""Upload a file to WaveSpeed.
|
|
333
|
+
|
|
334
|
+
Args:
|
|
335
|
+
file: File path string or file-like object to upload.
|
|
336
|
+
timeout: Total API call timeout in seconds.
|
|
337
|
+
|
|
338
|
+
Returns:
|
|
339
|
+
URL of the uploaded file.
|
|
340
|
+
|
|
341
|
+
Raises:
|
|
342
|
+
ValueError: If API key is not configured.
|
|
343
|
+
FileNotFoundError: If file path does not exist.
|
|
344
|
+
RuntimeError: If upload fails.
|
|
345
|
+
|
|
346
|
+
Example:
|
|
347
|
+
url = client.upload("/path/to/image.png")
|
|
348
|
+
print(url)
|
|
349
|
+
"""
|
|
350
|
+
if not self.api_key:
|
|
351
|
+
raise ValueError(
|
|
352
|
+
"API key is required. Set WAVESPEED_API_KEY environment variable "
|
|
353
|
+
"or pass api_key to Client()."
|
|
354
|
+
)
|
|
355
|
+
|
|
356
|
+
url = f"{self.base_url}/api/v3/media/upload/binary"
|
|
357
|
+
headers = {"Authorization": f"Bearer {self.api_key}"}
|
|
358
|
+
timeout = timeout or api_config.timeout
|
|
359
|
+
request_timeout = (min(self.connection_timeout, timeout), timeout)
|
|
360
|
+
|
|
361
|
+
if isinstance(file, str):
|
|
362
|
+
if not os.path.exists(file):
|
|
363
|
+
raise FileNotFoundError(f"File not found: {file}")
|
|
364
|
+
with open(file, "rb") as f:
|
|
365
|
+
files = {"file": (os.path.basename(file), f)}
|
|
366
|
+
response = requests.post(
|
|
367
|
+
url, headers=headers, files=files, timeout=request_timeout
|
|
368
|
+
)
|
|
369
|
+
else:
|
|
370
|
+
filename = getattr(file, "name", "upload")
|
|
371
|
+
if isinstance(filename, str) and os.path.sep in filename:
|
|
372
|
+
filename = os.path.basename(filename)
|
|
373
|
+
files = {"file": (filename, file)}
|
|
374
|
+
response = requests.post(
|
|
375
|
+
url, headers=headers, files=files, timeout=request_timeout
|
|
376
|
+
)
|
|
377
|
+
|
|
378
|
+
if response.status_code != 200:
|
|
379
|
+
raise RuntimeError(
|
|
380
|
+
f"Failed to upload file: HTTP {response.status_code}: {response.text}"
|
|
381
|
+
)
|
|
382
|
+
|
|
383
|
+
result = response.json()
|
|
384
|
+
if result.get("code") != 200:
|
|
385
|
+
raise RuntimeError(
|
|
386
|
+
f"Upload failed: {result.get('message', 'Unknown error')}"
|
|
387
|
+
)
|
|
388
|
+
|
|
389
|
+
download_url = result.get("data", {}).get("download_url")
|
|
390
|
+
if not download_url:
|
|
391
|
+
raise RuntimeError("Upload failed: no download_url in response")
|
|
392
|
+
|
|
393
|
+
return download_url
|
|
@@ -27,6 +27,15 @@ class api:
|
|
|
27
27
|
# Total API call timeout in seconds
|
|
28
28
|
timeout: float = 36000.0
|
|
29
29
|
|
|
30
|
+
# Maximum number of retries for the entire operation (task-level retries)
|
|
31
|
+
max_retries: int = 0
|
|
32
|
+
|
|
33
|
+
# Maximum number of retries for individual HTTP requests (connection errors, timeouts)
|
|
34
|
+
max_connection_retries: int = 5
|
|
35
|
+
|
|
36
|
+
# Base interval between retries in seconds (actual delay = retry_interval * attempt)
|
|
37
|
+
retry_interval: float = 1.0
|
|
38
|
+
|
|
30
39
|
|
|
31
40
|
class serverless:
|
|
32
41
|
"""Serverless configuration options.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: wavespeed
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.2
|
|
4
4
|
Summary: WaveSpeedAI Python Client — Official Python SDK for WaveSpeedAI inference platform
|
|
5
5
|
Author-email: WaveSpeedAI <support@wavespeed.ai>
|
|
6
6
|
License: MIT
|
|
@@ -62,7 +62,7 @@ import wavespeed
|
|
|
62
62
|
|
|
63
63
|
output = wavespeed.run(
|
|
64
64
|
"wavespeed-ai/z-image/turbo",
|
|
65
|
-
input={"prompt": "
|
|
65
|
+
input={"prompt": "Cat"},
|
|
66
66
|
)
|
|
67
67
|
|
|
68
68
|
print(output["outputs"][0]) # Output URL
|
|
@@ -70,7 +70,7 @@ print(output["outputs"][0]) # Output URL
|
|
|
70
70
|
|
|
71
71
|
### Authentication
|
|
72
72
|
|
|
73
|
-
Set your API key via environment variable:
|
|
73
|
+
Set your API key via environment variable (You can get your API key from [https://wavespeed.ai/accesskey](https://wavespeed.ai/accesskey)):
|
|
74
74
|
|
|
75
75
|
```bash
|
|
76
76
|
export WAVESPEED_API_KEY="your-api-key"
|
|
@@ -82,7 +82,7 @@ Or pass it directly:
|
|
|
82
82
|
from wavespeed import Client
|
|
83
83
|
|
|
84
84
|
client = Client(api_key="your-api-key")
|
|
85
|
-
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "
|
|
85
|
+
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Cat"})
|
|
86
86
|
```
|
|
87
87
|
|
|
88
88
|
### Options
|
|
@@ -90,9 +90,39 @@ output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "A cat holds
|
|
|
90
90
|
```python
|
|
91
91
|
output = wavespeed.run(
|
|
92
92
|
"wavespeed-ai/z-image/turbo",
|
|
93
|
-
input={"prompt": "
|
|
94
|
-
timeout=36000.0
|
|
95
|
-
poll_interval=1.0,
|
|
93
|
+
input={"prompt": "Cat"},
|
|
94
|
+
timeout=36000.0, # Max wait time in seconds (default: 36000.0)
|
|
95
|
+
poll_interval=1.0, # Status check interval (default: 1.0)
|
|
96
|
+
enable_sync_mode=False, # Single request mode, no polling (default: False)
|
|
97
|
+
)
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
### Sync Mode
|
|
101
|
+
|
|
102
|
+
Use `enable_sync_mode=True` for a single request that waits for the result (no polling).
|
|
103
|
+
|
|
104
|
+
> **Note:** Not all models support sync mode. Check the model documentation for availability.
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
output = wavespeed.run(
|
|
108
|
+
"wavespeed-ai/z-image/turbo",
|
|
109
|
+
input={"prompt": "Cat"},
|
|
110
|
+
enable_sync_mode=True,
|
|
111
|
+
)
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
### Retry Configuration
|
|
115
|
+
|
|
116
|
+
Configure retries at the client level:
|
|
117
|
+
|
|
118
|
+
```python
|
|
119
|
+
from wavespeed import Client
|
|
120
|
+
|
|
121
|
+
client = Client(
|
|
122
|
+
api_key="your-api-key",
|
|
123
|
+
max_retries=0, # Task-level retries (default: 0)
|
|
124
|
+
max_connection_retries=5, # HTTP connection retries (default: 5)
|
|
125
|
+
retry_interval=1.0, # Base delay between retries in seconds (default: 1.0)
|
|
96
126
|
)
|
|
97
127
|
```
|
|
98
128
|
|
|
@@ -175,6 +205,25 @@ def handler(job):
|
|
|
175
205
|
return {"output": "done"}
|
|
176
206
|
```
|
|
177
207
|
|
|
208
|
+
### Concurrent Execution
|
|
209
|
+
|
|
210
|
+
Enable concurrent job processing with `concurrency_modifier`:
|
|
211
|
+
|
|
212
|
+
```python
|
|
213
|
+
import wavespeed.serverless as serverless
|
|
214
|
+
|
|
215
|
+
def handler(job):
|
|
216
|
+
return {"output": job["input"]["data"]}
|
|
217
|
+
|
|
218
|
+
def concurrency_modifier(current_concurrency):
|
|
219
|
+
return 2 # Process 2 jobs concurrently
|
|
220
|
+
|
|
221
|
+
serverless.start({
|
|
222
|
+
"handler": handler,
|
|
223
|
+
"concurrency_modifier": concurrency_modifier
|
|
224
|
+
})
|
|
225
|
+
```
|
|
226
|
+
|
|
178
227
|
## Local Development
|
|
179
228
|
|
|
180
229
|
### Test with JSON Input
|
|
@@ -29,6 +29,10 @@ class TestClient(unittest.TestCase):
|
|
|
29
29
|
"""Test client initialization from config."""
|
|
30
30
|
mock_config.api_key = "config-key"
|
|
31
31
|
mock_config.base_url = "https://api.wavespeed.ai"
|
|
32
|
+
mock_config.connection_timeout = 10.0
|
|
33
|
+
mock_config.max_retries = 0
|
|
34
|
+
mock_config.max_connection_retries = 5
|
|
35
|
+
mock_config.retry_interval = 1.0
|
|
32
36
|
client = Client()
|
|
33
37
|
self.assertEqual(client.api_key, "config-key")
|
|
34
38
|
|
|
@@ -56,9 +60,12 @@ class TestClient(unittest.TestCase):
|
|
|
56
60
|
mock_post.return_value = mock_response
|
|
57
61
|
|
|
58
62
|
client = Client(api_key="test-key")
|
|
59
|
-
request_id = client._submit(
|
|
63
|
+
request_id, result = client._submit(
|
|
64
|
+
"wavespeed-ai/z-image/turbo", {"prompt": "test"}
|
|
65
|
+
)
|
|
60
66
|
|
|
61
67
|
self.assertEqual(request_id, "req-123")
|
|
68
|
+
self.assertIsNone(result)
|
|
62
69
|
mock_post.assert_called_once()
|
|
63
70
|
call_args = mock_post.call_args
|
|
64
71
|
self.assertIn("wavespeed-ai/z-image/turbo", call_args[0][0])
|
|
@@ -178,6 +185,11 @@ class TestModuleLevelRun(unittest.TestCase):
|
|
|
178
185
|
# Mock config
|
|
179
186
|
mock_config.api_key = "config-key"
|
|
180
187
|
mock_config.base_url = "https://api.wavespeed.ai"
|
|
188
|
+
mock_config.connection_timeout = 10.0
|
|
189
|
+
mock_config.timeout = 36000.0
|
|
190
|
+
mock_config.max_retries = 0
|
|
191
|
+
mock_config.max_connection_retries = 5
|
|
192
|
+
mock_config.retry_interval = 1.0
|
|
181
193
|
|
|
182
194
|
# Reset default client
|
|
183
195
|
wavespeed.api._default_client = None
|
|
@@ -1,238 +0,0 @@
|
|
|
1
|
-
"""WaveSpeed API client implementation."""
|
|
2
|
-
|
|
3
|
-
import os
|
|
4
|
-
import time
|
|
5
|
-
from typing import Any, BinaryIO
|
|
6
|
-
|
|
7
|
-
import requests
|
|
8
|
-
|
|
9
|
-
from wavespeed.config import api as api_config
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
class Client:
|
|
13
|
-
"""WaveSpeed API client.
|
|
14
|
-
|
|
15
|
-
Args:
|
|
16
|
-
api_key: WaveSpeed API key. If not provided, uses wavespeed.config.api.api_key.
|
|
17
|
-
base_url: Base URL for the API. If not provided, uses wavespeed.config.api.base_url.
|
|
18
|
-
connection_timeout: Timeout for HTTP requests in seconds.
|
|
19
|
-
|
|
20
|
-
Example:
|
|
21
|
-
client = Client(api_key="your-api-key")
|
|
22
|
-
output = client.run("wavespeed-ai/z-image/turbo", input={"prompt": "Hello"})
|
|
23
|
-
"""
|
|
24
|
-
|
|
25
|
-
def __init__(
|
|
26
|
-
self,
|
|
27
|
-
api_key: str | None = None,
|
|
28
|
-
base_url: str | None = None,
|
|
29
|
-
connection_timeout: float | None = None,
|
|
30
|
-
) -> None:
|
|
31
|
-
"""Initialize the client."""
|
|
32
|
-
self.api_key = api_key or api_config.api_key
|
|
33
|
-
self.base_url = (base_url or api_config.base_url).rstrip("/")
|
|
34
|
-
self.connection_timeout = connection_timeout or api_config.connection_timeout
|
|
35
|
-
|
|
36
|
-
def _get_headers(self) -> dict[str, str]:
|
|
37
|
-
"""Get request headers with authentication."""
|
|
38
|
-
if not self.api_key:
|
|
39
|
-
raise ValueError(
|
|
40
|
-
"API key is required. Set WAVESPEED_API_KEY environment variable "
|
|
41
|
-
"or pass api_key to Client()."
|
|
42
|
-
)
|
|
43
|
-
return {
|
|
44
|
-
"Content-Type": "application/json",
|
|
45
|
-
"Authorization": f"Bearer {self.api_key}",
|
|
46
|
-
}
|
|
47
|
-
|
|
48
|
-
def _submit(self, model: str, input: dict[str, Any] | None) -> str:
|
|
49
|
-
"""Submit a prediction request.
|
|
50
|
-
|
|
51
|
-
Args:
|
|
52
|
-
model: Model identifier.
|
|
53
|
-
input: Input parameters.
|
|
54
|
-
|
|
55
|
-
Returns:
|
|
56
|
-
Request ID for polling.
|
|
57
|
-
|
|
58
|
-
Raises:
|
|
59
|
-
RuntimeError: If submission fails.
|
|
60
|
-
"""
|
|
61
|
-
url = f"{self.base_url}/api/v3/{model}"
|
|
62
|
-
body = input or {}
|
|
63
|
-
|
|
64
|
-
response = requests.post(
|
|
65
|
-
url, json=body, headers=self._get_headers(), timeout=self.connection_timeout
|
|
66
|
-
)
|
|
67
|
-
|
|
68
|
-
if response.status_code != 200:
|
|
69
|
-
raise RuntimeError(
|
|
70
|
-
f"Failed to submit prediction: HTTP {response.status_code}: "
|
|
71
|
-
f"{response.text}"
|
|
72
|
-
)
|
|
73
|
-
|
|
74
|
-
result = response.json()
|
|
75
|
-
request_id = result.get("data", {}).get("id")
|
|
76
|
-
|
|
77
|
-
if not request_id:
|
|
78
|
-
raise RuntimeError(f"No request ID in response: {result}")
|
|
79
|
-
|
|
80
|
-
return request_id
|
|
81
|
-
|
|
82
|
-
def _get_result(self, request_id: str) -> dict[str, Any]:
|
|
83
|
-
"""Get prediction result.
|
|
84
|
-
|
|
85
|
-
Args:
|
|
86
|
-
request_id: The prediction request ID.
|
|
87
|
-
|
|
88
|
-
Returns:
|
|
89
|
-
Full API response.
|
|
90
|
-
|
|
91
|
-
Raises:
|
|
92
|
-
RuntimeError: If fetching result fails.
|
|
93
|
-
"""
|
|
94
|
-
url = f"{self.base_url}/api/v3/predictions/{request_id}/result"
|
|
95
|
-
|
|
96
|
-
response = requests.get(
|
|
97
|
-
url, headers=self._get_headers(), timeout=self.connection_timeout
|
|
98
|
-
)
|
|
99
|
-
|
|
100
|
-
if response.status_code != 200:
|
|
101
|
-
raise RuntimeError(
|
|
102
|
-
f"Failed to get result: HTTP {response.status_code}: {response.text}"
|
|
103
|
-
)
|
|
104
|
-
|
|
105
|
-
return response.json()
|
|
106
|
-
|
|
107
|
-
def _wait(
|
|
108
|
-
self,
|
|
109
|
-
request_id: str,
|
|
110
|
-
timeout: float | None,
|
|
111
|
-
poll_interval: float,
|
|
112
|
-
) -> dict[str, Any]:
|
|
113
|
-
"""Wait for prediction to complete.
|
|
114
|
-
|
|
115
|
-
Args:
|
|
116
|
-
request_id: The prediction request ID.
|
|
117
|
-
timeout: Maximum wait time in seconds (None = no timeout).
|
|
118
|
-
poll_interval: Time between polls in seconds.
|
|
119
|
-
|
|
120
|
-
Returns:
|
|
121
|
-
Dict with "outputs" array.
|
|
122
|
-
|
|
123
|
-
Raises:
|
|
124
|
-
RuntimeError: If prediction fails.
|
|
125
|
-
TimeoutError: If prediction times out.
|
|
126
|
-
"""
|
|
127
|
-
start_time = time.time()
|
|
128
|
-
|
|
129
|
-
while True:
|
|
130
|
-
# Check timeout
|
|
131
|
-
if timeout is not None:
|
|
132
|
-
elapsed = time.time() - start_time
|
|
133
|
-
if elapsed >= timeout:
|
|
134
|
-
raise TimeoutError(f"Prediction timed out after {timeout} seconds")
|
|
135
|
-
|
|
136
|
-
result = self._get_result(request_id)
|
|
137
|
-
data = result.get("data", {})
|
|
138
|
-
status = data.get("status")
|
|
139
|
-
|
|
140
|
-
if status == "completed":
|
|
141
|
-
return {"outputs": data.get("outputs", [])}
|
|
142
|
-
|
|
143
|
-
if status == "failed":
|
|
144
|
-
error = data.get("error") or "Unknown error"
|
|
145
|
-
raise RuntimeError(f"Prediction failed: {error}")
|
|
146
|
-
|
|
147
|
-
time.sleep(poll_interval)
|
|
148
|
-
|
|
149
|
-
def run(
|
|
150
|
-
self,
|
|
151
|
-
model: str,
|
|
152
|
-
input: dict[str, Any] | None = None,
|
|
153
|
-
*,
|
|
154
|
-
timeout: float | None = None,
|
|
155
|
-
poll_interval: float = 1.0,
|
|
156
|
-
) -> dict[str, Any]:
|
|
157
|
-
"""Run a model and wait for the output.
|
|
158
|
-
|
|
159
|
-
Args:
|
|
160
|
-
model: Model identifier (e.g., "wavespeed-ai/flux-dev").
|
|
161
|
-
input: Input parameters for the model.
|
|
162
|
-
timeout: Maximum time to wait for completion (None = no timeout).
|
|
163
|
-
poll_interval: Interval between status checks in seconds.
|
|
164
|
-
|
|
165
|
-
Returns:
|
|
166
|
-
Dict containing "outputs" array with model outputs.
|
|
167
|
-
|
|
168
|
-
Raises:
|
|
169
|
-
ValueError: If API key is not configured.
|
|
170
|
-
RuntimeError: If the prediction fails.
|
|
171
|
-
TimeoutError: If the prediction times out.
|
|
172
|
-
"""
|
|
173
|
-
request_id = self._submit(model, input)
|
|
174
|
-
return self._wait(request_id, timeout, poll_interval)
|
|
175
|
-
|
|
176
|
-
def upload(self, file: str | BinaryIO, *, timeout: float | None = None) -> str:
|
|
177
|
-
"""Upload a file to WaveSpeed.
|
|
178
|
-
|
|
179
|
-
Args:
|
|
180
|
-
file: File path string or file-like object to upload.
|
|
181
|
-
timeout: Total API call timeout in seconds.
|
|
182
|
-
|
|
183
|
-
Returns:
|
|
184
|
-
URL of the uploaded file.
|
|
185
|
-
|
|
186
|
-
Raises:
|
|
187
|
-
ValueError: If API key is not configured.
|
|
188
|
-
FileNotFoundError: If file path does not exist.
|
|
189
|
-
RuntimeError: If upload fails.
|
|
190
|
-
|
|
191
|
-
Example:
|
|
192
|
-
url = client.upload("/path/to/image.png")
|
|
193
|
-
print(url)
|
|
194
|
-
"""
|
|
195
|
-
if not self.api_key:
|
|
196
|
-
raise ValueError(
|
|
197
|
-
"API key is required. Set WAVESPEED_API_KEY environment variable "
|
|
198
|
-
"or pass api_key to Client()."
|
|
199
|
-
)
|
|
200
|
-
|
|
201
|
-
url = f"{self.base_url}/api/v3/media/upload/binary"
|
|
202
|
-
headers = {"Authorization": f"Bearer {self.api_key}"}
|
|
203
|
-
timeout = timeout or api_config.timeout
|
|
204
|
-
request_timeout = (min(self.connection_timeout, timeout), timeout)
|
|
205
|
-
|
|
206
|
-
if isinstance(file, str):
|
|
207
|
-
if not os.path.exists(file):
|
|
208
|
-
raise FileNotFoundError(f"File not found: {file}")
|
|
209
|
-
with open(file, "rb") as f:
|
|
210
|
-
files = {"file": (os.path.basename(file), f)}
|
|
211
|
-
response = requests.post(
|
|
212
|
-
url, headers=headers, files=files, timeout=request_timeout
|
|
213
|
-
)
|
|
214
|
-
else:
|
|
215
|
-
filename = getattr(file, "name", "upload")
|
|
216
|
-
if isinstance(filename, str) and os.path.sep in filename:
|
|
217
|
-
filename = os.path.basename(filename)
|
|
218
|
-
files = {"file": (filename, file)}
|
|
219
|
-
response = requests.post(
|
|
220
|
-
url, headers=headers, files=files, timeout=request_timeout
|
|
221
|
-
)
|
|
222
|
-
|
|
223
|
-
if response.status_code != 200:
|
|
224
|
-
raise RuntimeError(
|
|
225
|
-
f"Failed to upload file: HTTP {response.status_code}: {response.text}"
|
|
226
|
-
)
|
|
227
|
-
|
|
228
|
-
result = response.json()
|
|
229
|
-
if result.get("code") != 200:
|
|
230
|
-
raise RuntimeError(
|
|
231
|
-
f"Upload failed: {result.get('message', 'Unknown error')}"
|
|
232
|
-
)
|
|
233
|
-
|
|
234
|
-
download_url = result.get("data", {}).get("download_url")
|
|
235
|
-
if not download_url:
|
|
236
|
-
raise RuntimeError("Upload failed: no download_url in response")
|
|
237
|
-
|
|
238
|
-
return download_url
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|