Introduce experimental FileOutput interface for models that output File and Path types (#348)

This PR is a proposal to add a new `FileOutput` type to the client SDK
to abstract away file outputs from Replicate models.

It can be enabled by passing the `use_file_output` flag to the `run()`
method (this will be moved to the constructor).

```python
replicate.run("black-forest-labs/flux-schnell", input={}, use_file_input =True);
```

When enabled any URLs (and soon data-uris) will be converted into a
`FileOutput type`. This is essentially an `Iterable[bytes] |
AsyncIterable[bytes]` that has two additional fields, the attribute
`url` referencing underlying URL and `read()` which will return `bytes`
with the file data loaded into memory.
  
The intention here is to make it easier to work with file outputs and
allows us to optimize the delivery of file assets to the client in
future iterations.

Usage is as follows:

```python
output = replicate.run(
  "black-forest-labs/flux-schnell",
  input={"prompt": "astronaut riding a rocket like a horse"},
  use_file_input=True,
);
```

For most basic cases you'll want to utilize either the `url` or `read()`
fields depending on whether you want to directly consume the file or
pass it on.

To access the file URL:

```python
print(output.url) #=> "https://delivery.replicate.com/..."
```

To consume the file directly:

```python
with open('output.bin', 'wb') as file:
    file.write(output.read())
```

Or for very large files they can be streamed:

```python
with open(file_path, 'wb') as file:
    for chunk in output:
        file.write(chunk)
```

Each of these methods has an equivalent `asyncio` API.

```python
async with aiofiles.open(filename, 'w') as file:
    await file.write(await output.aread())

async with aiofiles.open(filename, 'w') as file:
    await for chunk in output:
        await file.write(chunk)
```

For streaming responses from common frameworks, all support taking
`Iterator` types:

**Django**

```python
@condition(etag_func=None)
def stream_response(request):
    output = replicate.run("black-forest-labs/flux-schnell", input={...}, use_file_input =True)
    return HttpResponse(output, content_type='image/webp')
```
  
**FastAPI**

```python
@app.get("/")
async def main():
    output = replicate.run("black-forest-labs/flux-schnell", input={...}, use_file_input =True)
    return StreamingResponse(output)
```

**Flask**

```python
@app.route('/stream')
def streamed_response():
    output = replicate.run("black-forest-labs/flux-schnell", input={...}, use_file_input =True)
    return app.response_class(stream_with_context(output))
```

---------

Signed-off-by: Mattt Zmuda <[email protected]>
Co-authored-by: Mattt Zmuda <[email protected]>

Author: aron

pyproject.toml

@@ -75,6 +75,7 @@ ignore = [
     "ANN003", # Missing type annotation for `**kwargs`
     "ANN101", # Missing type annotation for self in method
     "ANN102", # Missing type annotation for cls in classmethod
+    "ANN401", # Dynamically typed expressions (typing.Any) are disallowed in {name}
     "W191",   # Indentation contains tabs
     "UP037",  # Remove quotes from type annotation
 ]
@@ -86,3 +87,7 @@ ignore = [
     "ANN201", # Missing return type annotation for public function
     "ANN202", # Missing return type annotation for private function
 ]
+
+[tool.pyright]
+venvPath = "."
+venv = ".venv"

replicate/client.py

@@ -164,25 +164,27 @@ def run(
         self,
         ref: str,
         input: Optional[Dict[str, Any]] = None,
+        use_file_output: Optional[bool] = None,
         **params: Unpack["Predictions.CreatePredictionParams"],
     ) -> Union[Any, Iterator[Any]]:  # noqa: ANN401
         """
         Run a model and wait for its output.
         """
 
-        return run(self, ref, input, **params)
+        return run(self, ref, input, use_file_output, **params)
 
     async def async_run(
         self,
         ref: str,
         input: Optional[Dict[str, Any]] = None,
+        use_file_output: Optional[bool] = None,
         **params: Unpack["Predictions.CreatePredictionParams"],
     ) -> Union[Any, AsyncIterator[Any]]:  # noqa: ANN401
         """
         Run a model and wait for its output asynchronously.
         """
 
-        return await async_run(self, ref, input, **params)
+        return await async_run(self, ref, input, use_file_output, **params)
 
     def stream(
         self,

replicate/deployment.py

@@ -3,7 +3,7 @@
 from typing_extensions import Unpack, deprecated
 
 from replicate.account import Account
-from replicate.json import async_encode_json, encode_json
+from replicate.helpers import async_encode_json, encode_json
 from replicate.pagination import Page
 from replicate.prediction import (
     Prediction,

replicate/json.py renamed to replicate/helpers.py

@@ -1,9 +1,12 @@
 import base64
 import io
 import mimetypes
+from collections.abc import Mapping, Sequence
 from pathlib import Path
 from types import GeneratorType
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, Optional
+
+import httpx
 
 if TYPE_CHECKING:
     from replicate.client import Client
@@ -108,3 +111,80 @@ def base64_encode_file(file: io.IOBase) -> str:
         mimetypes.guess_type(getattr(file, "name", ""))[0] or "application/octet-stream"
     )
     return f"data:{mime_type};base64,{encoded_body}"
+
+
+class FileOutput(httpx.SyncByteStream, httpx.AsyncByteStream):
+    """
+    An object that can be used to read the contents of an output file
+    created by running a Replicate model.
+    """
+
+    url: str
+    """
+    The file URL.
+    """
+
+    _client: "Client"
+
+    def __init__(self, url: str, client: "Client") -> None:
+        self.url = url
+        self._client = client
+
+    def read(self) -> bytes:
+        if self.url.startswith("data:"):
+            _, encoded = self.url.split(",", 1)
+            return base64.b64decode(encoded)
+
+        with self._client._client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            return response.read()
+
+    def __iter__(self) -> Iterator[bytes]:
+        if self.url.startswith("data:"):
+            yield self.read()
+            return
+
+        with self._client._client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            yield from response.iter_bytes()
+
+    async def aread(self) -> bytes:
+        if self.url.startswith("data:"):
+            _, encoded = self.url.split(",", 1)
+            return base64.b64decode(encoded)
+
+        async with self._client._async_client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            return await response.aread()
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        if self.url.startswith("data:"):
+            yield await self.aread()
+            return
+
+        async with self._client._async_client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            async for chunk in response.aiter_bytes():
+                yield chunk
+
+    def __str__(self) -> str:
+        return self.url
+
+
+def transform_output(value: Any, client: "Client") -> Any:
+    """
+    Transform the output of a prediction to a `FileOutput` object if it's a URL.
+    """
+
+    def transform(obj: Any) -> Any:
+        if isinstance(obj, Mapping):
+            return {k: transform(v) for k, v in obj.items()}
+        elif isinstance(obj, Sequence) and not isinstance(obj, str):
+            return [transform(item) for item in obj]
+        elif isinstance(obj, str) and (
+            obj.startswith("https:") or obj.startswith("data:")
+        ):
+            return FileOutput(obj, client)
+        return obj
+
+    return transform(value)

replicate/model.py

@@ -3,8 +3,8 @@
 from typing_extensions import NotRequired, TypedDict, Unpack, deprecated
 
 from replicate.exceptions import ReplicateException
+from replicate.helpers import async_encode_json, encode_json
 from replicate.identifier import ModelVersionIdentifier
-from replicate.json import async_encode_json, encode_json
 from replicate.pagination import Page
 from replicate.prediction import (
     Prediction,

replicate/prediction.py

@@ -20,7 +20,7 @@
 
 from replicate.exceptions import ModelError, ReplicateError
 from replicate.file import FileEncodingStrategy
-from replicate.json import async_encode_json, encode_json
+from replicate.helpers import async_encode_json, encode_json
 from replicate.pagination import Page
 from replicate.resource import Namespace, Resource
 from replicate.stream import EventSource

replicate/run.py

@@ -13,6 +13,7 @@
 
 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
@@ -28,6 +29,7 @@ def run(
     client: "Client",
     ref: Union["Model", "Version", "ModelVersionIdentifier", str],
     input: Optional[Dict[str, Any]] = None,
+    use_file_output: Optional[bool] = None,
     **params: Unpack["Predictions.CreatePredictionParams"],
 ) -> Union[Any, Iterator[Any]]:  # noqa: ANN401
     """
@@ -60,13 +62,17 @@ def run(
     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] = None,
     **params: Unpack["Predictions.CreatePredictionParams"],
 ) -> Union[Any, AsyncIterator[Any]]:  # noqa: ANN401
     """
@@ -99,6 +105,9 @@ async def async_run(
     if prediction.status == "failed":
         raise ModelError(prediction)
 
+    if use_file_output:
+        return transform_output(prediction.output, client)
+
     return prediction.output
 
 

replicate/stream.py

@@ -10,6 +10,7 @@
     Union,
 )
 
+import httpx
 from typing_extensions import Unpack
 
 from replicate import identifier
@@ -22,8 +23,6 @@
 
 
 if TYPE_CHECKING:
-    import httpx
-
     from replicate.client import Client
     from replicate.identifier import ModelVersionIdentifier
     from replicate.model import Model

replicate/training.py

@@ -13,8 +13,8 @@
 
 from typing_extensions import NotRequired, Unpack
 
+from replicate.helpers import async_encode_json, encode_json
 from replicate.identifier import ModelVersionIdentifier
-from replicate.json import async_encode_json, encode_json
 from replicate.model import Model
 from replicate.pagination import Page
 from replicate.resource import Namespace, Resource

requirements-dev.lock

@@ -6,6 +6,8 @@
 #   features: []
 #   all-features: false
 #   with-sources: false
+#   generate-hashes: false
+#   universal: false
 
 -e file:.
 annotated-types==0.6.0