FilesExpand file tree

 main

/

client.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

135 lines (121 loc) · 5.01 KB

 main

/

client.py

Copy path

Top

File metadata and controls

• Code
• Blame

135 lines (121 loc) · 5.01 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

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

import os

import re

from json import JSONDecodeError

from typing import Any, Iterator, Union

import requests

from requests.adapters import HTTPAdapter, Retry

from replicate.__about__ import __version__

from replicate.exceptions import ModelError, ReplicateError

from replicate.model import ModelCollection

from replicate.prediction import PredictionCollection

class Client:

def __init__(self, api_token=None) -> None:

super().__init__()

# Client is instantiated at import time, so do as little as possible.

# This includes resolving environment variables -- they might be set programmatically.

self.api_token = api_token

self.base_url = os.environ.get(

"REPLICATE_API_BASE_URL", "https://api.replicate.com"

)

self.poll_interval = float(os.environ.get("REPLICATE_POLL_INTERVAL", "0.5"))

# TODO: make thread safe

self.read_session = requests.Session()

read_retries = Retry(

total=5,

backoff_factor=2,

# Only retry 500s on GET so we don't unintionally mutute data

method_whitelist=["GET"],

# https://support.cloudflare.com/hc/en-us/articles/115003011431-Troubleshooting-Cloudflare-5XX-errors

status_forcelist=[

429,

500,

502,

503,

504,

520,

521,

522,

523,

524,

526,

527,

],

)

self.read_session.mount("http://", HTTPAdapter(max_retries=read_retries))

self.read_session.mount("https://", HTTPAdapter(max_retries=read_retries))

self.write_session = requests.Session()

write_retries = Retry(

total=5,

backoff_factor=2,

method_whitelist=["POST", "PUT"],

# Only retry POST/PUT requests on rate limits, so we don't unintionally mutute data

status_forcelist=[429],

)

self.write_session.mount("http://", HTTPAdapter(max_retries=write_retries))

self.write_session.mount("https://", HTTPAdapter(max_retries=write_retries))

def _request(self, method: str, path: str, **kwargs):

# from requests.Session

if method in ["GET", "OPTIONS"]:

kwargs.setdefault("allow_redirects", True)

if method in ["HEAD"]:

kwargs.setdefault("allow_redirects", False)

kwargs.setdefault("headers", {})

kwargs["headers"].update(self._headers())

session = self.read_session

if method in ["POST", "PUT", "DELETE", "PATCH"]:

session = self.write_session

resp = session.request(method, self.base_url + path, **kwargs)

if 400 <= resp.status_code < 600:

try:

raise ReplicateError(resp.json()["detail"])

except (JSONDecodeError, KeyError):

pass

raise ReplicateError(f"HTTP error: {resp.status_code, resp.reason}")

return resp

def _headers(self):

return {

"Authorization": f"Token {self._api_token()}",

"User-Agent": f"replicate-python@{__version__}",

}

def _api_token(self):

token = self.api_token

# Evaluate lazily in case environment variable is set with dotenv, or something

if token is None:

token = os.environ.get("REPLICATE_API_TOKEN")

if not token:

raise ReplicateError(

"""No API token provided. You need to set the REPLICATE_API_TOKEN environment variable or create a client with `replicate.Client(api_token=...)`.

You can find your API key on https://replicate.com"""

)

return token

@property

def models(self) -> ModelCollection:

return ModelCollection(client=self)

@property

def predictions(self) -> PredictionCollection:

return PredictionCollection(client=self)

def run(self, model_version, **kwargs) -> Union[Any, Iterator[Any]]:

"""

Run a model in the format owner/name:version.

"""

# Split model_version into owner, name, version in format owner/name:version

m = re.match(r"^(?P<model>[^/]+/[^:]+):(?P<version>.+)$", model_version)

if not m:

raise ReplicateError(

f"Invalid model_version: {model_version}. Expected format: owner/name:version"

)

model = self.models.get(m.group("model"))

version = model.versions.get(m.group("version"))

prediction = self.predictions.create(version=version, **kwargs)

# Return an iterator of the output

schema = version.get_transformed_schema()

output = schema["components"]["schemas"]["Output"]

if (

output.get("type") == "array"

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

):

return prediction.output_iterator()

prediction.wait()

if prediction.status == "failed":

raise ModelError(prediction.error)

return prediction.output