forked from replicate/replicate-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathrun.py
More file actions
109 lines (83 loc) · 3.39 KB
/
Copy pathrun.py
File metadata and controls
109 lines (83 loc) · 3.39 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
from typing_extensions import Unpack
from replicate import identifier
from replicate.exceptions import ModelError
from replicate.model import Model
from replicate.prediction import Prediction
from replicate.schema import make_schema_backwards_compatible
from replicate.version import Version, Versions
if TYPE_CHECKING:
from replicate.client import Client
from replicate.identifier import ModelVersionIdentifier
from replicate.prediction import Predictions
def run(
client: "Client",
ref: Union["Model", "Version", "ModelVersionIdentifier", str],
input: Optional[Dict[str, Any]] = None,
**params: Unpack["Predictions.CreatePredictionParams"],
) -> Union[Any, Iterator[Any]]: # noqa: ANN401
"""
Run a model and wait for its output.
"""
version, owner, name, version_id = identifier._resolve(ref)
if version_id is not None:
prediction = client.predictions.create(
version=version_id, input=input or {}, **params
)
elif owner and name:
prediction = client.models.predictions.create(
model=(owner, name), input=input or {}, **params
)
else:
raise ValueError(
f"Invalid argument: {ref}. Expected model, version, or reference in the format owner/name or owner/name:version"
)
if not version and (owner and name and version_id):
version = Versions(client, model=(owner, name)).get(version_id)
if version and (iterator := _make_output_iterator(version, prediction)):
return iterator
prediction.wait()
if prediction.status == "failed":
raise ModelError(prediction.error)
return prediction.output
async def async_run(
client: "Client",
ref: Union["Model", "Version", "ModelVersionIdentifier", str],
input: Optional[Dict[str, Any]] = None,
**params: Unpack["Predictions.CreatePredictionParams"],
) -> Union[Any, Iterator[Any]]: # noqa: ANN401
"""
Run a model and wait for its output asynchronously.
"""
version, owner, name, version_id = identifier._resolve(ref)
if version or version_id:
prediction = await client.predictions.async_create(
version=(version or version_id), input=input or {}, **params
)
elif owner and name:
prediction = await client.models.predictions.async_create(
model=(owner, name), input=input or {}, **params
)
else:
raise ValueError(
f"Invalid argument: {ref}. Expected model, version, or reference in the format owner/name or owner/name:version"
)
if not version and (owner and name and version_id):
version = Versions(client, model=(owner, name)).get(version_id)
if version and (iterator := _make_output_iterator(version, prediction)):
return iterator
prediction.wait()
if prediction.status == "failed":
raise ModelError(prediction.error)
return prediction.output
def _make_output_iterator(
version: Version, prediction: Prediction
) -> Optional[Iterator[Any]]:
schema = make_schema_backwards_compatible(
version.openapi_schema, version.cog_version
)
output = schema["components"]["schemas"]["Output"]
if output.get("type") == "array" and output.get("x-cog-array-type") == "iterator":
return prediction.output_iterator()
return None
__all__: List = []