replicate-python/replicate/run.py at main · sfrias/replicate-python · GitHub

History

154 lines (119 loc) · 4.5 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

from typing import (

TYPE_CHECKING,

Any,

AsyncIterator,

Dict,

Iterator,

List,

Optional,

Union,

)

from typing_extensions import Unpack

from replicate import identifier

from replicate.exceptions import ModelError

from replicate.helpers import transform_output

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,

use_file_output: Optional[bool] = True,

**params: Unpack["Predictions.CreatePredictionParams"],

) -> Union[Any, Iterator[Any]]: # noqa: ANN401

"""

Run a model and wait for its output.

"""

if "wait" not in params:

params["wait"] = True

is_blocking = params["wait"] != False # noqa: E712

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

if not (is_blocking and prediction.status != "starting"):

prediction.wait()

if prediction.status == "failed":

raise ModelError(prediction)

if use_file_output:

return transform_output(prediction.output, client)

return prediction.output

async def async_run(

client: "Client",

ref: Union["Model", "Version", "ModelVersionIdentifier", str],

input: Optional[Dict[str, Any]] = None,

use_file_output: Optional[bool] = True,

**params: Unpack["Predictions.CreatePredictionParams"],

) -> Union[Any, AsyncIterator[Any]]: # noqa: ANN401

"""

Run a model and wait for its output asynchronously.

"""

if "wait" not in params:

params["wait"] = True

is_blocking = params["wait"] != False # noqa: E712

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 = await Versions(client, model=(owner, name)).async_get(version_id)

if version and (iterator := _make_async_output_iterator(version, prediction)):

return iterator

if not (is_blocking and prediction.status != "starting"):

await prediction.async_wait()

if prediction.status == "failed":

raise ModelError(prediction)

if use_file_output:

return transform_output(prediction.output, client)

return prediction.output

def _has_output_iterator_array_type(version: Version) -> bool:

schema = make_schema_backwards_compatible(

version.openapi_schema, version.cog_version

)

output = schema.get("components", {}).get("schemas", {}).get("Output", {})

return (

output.get("type") == "array" and output.get("x-cog-array-type") == "iterator"

)

def _make_output_iterator(

version: Version, prediction: Prediction

) -> Optional[Iterator[Any]]:

if _has_output_iterator_array_type(version):

return prediction.output_iterator()

return None

def _make_async_output_iterator(

version: Version, prediction: Prediction

) -> Optional[AsyncIterator[Any]]:

if _has_output_iterator_array_type(version):

return prediction.async_output_iterator()

return None

__all__: List = []

replicate-python/replicate/run.py at main · sfrias/replicate-python · GitHub

Expand file tree

Breadcrumbs

run.py

Latest commit

History

Breadcrumbs

run.py

File metadata and controls