Implement experimental FileOutput interface

replicate/client.py

@@ -164,25 +164,27 @@ def run(
         self,
         ref: str,
         input: Optional[Dict[str, Any]] = None,
+        use_file_output: bool = False,
         **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: bool = False,
         **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/run.py

@@ -17,6 +17,7 @@
 from replicate.prediction import Prediction
 from replicate.schema import make_schema_backwards_compatible
 from replicate.version import Version, Versions
+from replicate.stream import FileOutput
 
 if TYPE_CHECKING:
     from replicate.client import Client
@@ -28,6 +29,7 @@ def run(
     client: "Client",
     ref: Union["Model", "Version", "ModelVersionIdentifier", str],
     input: Optional[Dict[str, Any]] = None,
+    use_file_output: bool = False,
     **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: bool = False,
     **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
 
 
@@ -130,4 +139,22 @@ def _make_async_output_iterator(
     return None
 
 
+def transform(obj, func):
+    if isinstance(obj, dict):
+        return {k: transform(v, func) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [transform(item, func) for item in obj]
+    else:
+        return func(obj)
+
+
+def transform_output(value: Any, client: "Client"):
+    def wrapper(x):
+        if isinstance(x, str) and (x.startswith("https:") or x.startswith("data:")):
+            return FileOutput(x, client)
+        return x
+
+    return transform(value, wrapper)
+
+
 __all__: List = []

replicate/stream.py

@@ -33,7 +33,7 @@
     from replicate.version import Version
 
 
-class FileOutputProvider:
+class FileOutput(httpx.ByteStream, httpx.AsyncByteStream):
     url: str
     client: "Client"
 
@@ -42,52 +42,32 @@ def __init__(self, url: str, client: "Client"):
         self.client = client
 
     def read(self) -> bytes:
-        with self.stream() as file:
-            return file.read()
-
-    @contextmanager
-    def stream(self) -> Iterator["FileOutput"]:
         with self.client._client.stream("GET", self.url) as response:
             response.raise_for_status()
-            yield FileOutput(response)
+            return response.read()
 
-    @asynccontextmanager
-    async def astream(self) -> AsyncIterator["FileOutput"]:
-        async with self.client._async_client.stream("GET", self.url) as response:
+    def __iter__(self) -> Iterator[bytes]:
+        with self.client._client.stream("GET", self.url) as response:
             response.raise_for_status()
-            yield FileOutput(response)
+            for chunk in response.iter_bytes():
+                yield chunk
 
     async def aread(self) -> bytes:
-        async with self.astream() as file:
-            return await file.aread()
-
-    def __repr__(self) -> str:
-        return self.url
-
-
-class FileOutput(httpx.ByteStream, httpx.AsyncByteStream):
-    def __init__(self, response: httpx.Response):
-        self.response = response
-
-    def __iter__(self) -> Iterator[bytes]:
-        for bytes in self.response.iter_bytes():
-            yield bytes
-
-    def close(self):
-        return self.response.close()
-
-    def read(self):
-        return self.response.read()
+        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]:
-        async for bytes in self.response.aiter_bytes():
-            yield bytes
+        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
 
-    async def aclose(self):
-        return await self.response.aclose()
+    def __str__(self) -> str:
+        return self.url
 
-    async def aread(self):
-        return await self.response.aread()
+    def __repr__(self) -> str:
+        return self.url
 
 
 class ServerSentEvent(pydantic.BaseModel):  # type: ignore

tests/test_run.py

@@ -5,9 +5,11 @@
 import pytest
 import respx
 
+from typing import cast
 import replicate
 from replicate.client import Client
 from replicate.exceptions import ModelError, ReplicateError
+from replicate.stream import FileOutput
 
 
 @pytest.mark.vcr("run.yaml")
@@ -73,7 +75,7 @@ async def test_run_concurrently(mock_replicate_api_token, record_mode):
     results = await asyncio.gather(*tasks)
     assert len(results) == len(prompts)
     assert all(isinstance(result, list) for result in results)
-    assert all(len(result) > 0 for result in results)
+    assert all(len(results) > 0 for result in results)
 
 
 @pytest.mark.vcr("run.yaml")
@@ -253,3 +255,174 @@ def prediction_with_status(status: str) -> dict:
     assert str(excinfo.value) == "OOM"
     assert excinfo.value.prediction.error == "OOM"
     assert excinfo.value.prediction.status == "failed"
+
+
+@pytest.mark.asyncio
+async def test_run_with_file_output(mock_replicate_api_token):
+    def prediction_with_status(
+        status: str, output: str | list[str] | None = None
+    ) -> dict:
+        return {
+            "id": "p1",
+            "model": "test/example",
+            "version": "v1",
+            "urls": {
+                "get": "https://api.replicate.com/v1/predictions/p1",
+                "cancel": "https://api.replicate.com/v1/predictions/p1/cancel",
+            },
+            "created_at": "2023-10-05T12:00:00.000000Z",
+            "source": "api",
+            "status": status,
+            "input": {"text": "world"},
+            "output": output,
+            "error": "OOM" if status == "failed" else None,
+            "logs": "",
+        }
+
+    router = respx.Router(base_url="https://api.replicate.com/v1")
+    router.route(method="POST", path="/predictions").mock(
+        return_value=httpx.Response(
+            201,
+            json=prediction_with_status("processing"),
+        )
+    )
+    router.route(method="GET", path="/predictions/p1").mock(
+        return_value=httpx.Response(
+            200,
+            json=prediction_with_status(
+                "succeeded", "https://api.replicate.com/v1/assets/output.txt"
+            ),
+        )
+    )
+    router.route(
+        method="GET",
+        path="/models/test/example/versions/v1",
+    ).mock(
+        return_value=httpx.Response(
+            201,
+            json={
+                "id": "f2d6b24e6002f25f77ae89c2b0a5987daa6d0bf751b858b94b8416e8542434d1",
+                "created_at": "2024-07-18T00:35:56.210272Z",
+                "cog_version": "0.9.10",
+                "openapi_schema": {
+                    "openapi": "3.0.2",
+                },
+            },
+        )
+    )
+    router.route(method="GET", path="/assets/output.txt").mock(
+        return_value=httpx.Response(200, content=b"Hello, world!")
+    )
+
+    client = Client(
+        api_token="test-token", transport=httpx.MockTransport(router.handler)
+    )
+    client.poll_interval = 0.001
+
+    output = cast(
+        FileOutput,
+        client.run(
+            "test/example:v1",
+            input={
+                "text": "Hello, world!",
+            },
+            use_file_output=True,
+        ),
+    )
+
+    assert output.url == "https://api.replicate.com/v1/assets/output.txt"
+
+    assert output.read() == b"Hello, world!"
+    for chunk in output:
+        assert chunk == b"Hello, world!"
+
+    assert await output.aread() == b"Hello, world!"
+    async for chunk in output:
+        assert chunk == b"Hello, world!"
+
+
+@pytest.mark.asyncio
+async def test_run_with_file_output_array(mock_replicate_api_token):
+    def prediction_with_status(
+        status: str, output: str | list[str] | None = None
+    ) -> dict:
+        return {
+            "id": "p1",
+            "model": "test/example",
+            "version": "v1",
+            "urls": {
+                "get": "https://api.replicate.com/v1/predictions/p1",
+                "cancel": "https://api.replicate.com/v1/predictions/p1/cancel",
+            },
+            "created_at": "2023-10-05T12:00:00.000000Z",
+            "source": "api",
+            "status": status,
+            "input": {"text": "world"},
+            "output": output,
+            "error": "OOM" if status == "failed" else None,
+            "logs": "",
+        }
+
+    router = respx.Router(base_url="https://api.replicate.com/v1")
+    router.route(method="POST", path="/predictions").mock(
+        return_value=httpx.Response(
+            201,
+            json=prediction_with_status("processing"),
+        )
+    )
+    router.route(method="GET", path="/predictions/p1").mock(
+        return_value=httpx.Response(
+            200,
+            json=prediction_with_status(
+                "succeeded",
+                [
+                    "https://api.replicate.com/v1/assets/hello.txt",
+                    "https://api.replicate.com/v1/assets/world.txt",
+                ],
+            ),
+        )
+    )
+    router.route(
+        method="GET",
+        path="/models/test/example/versions/v1",
+    ).mock(
+        return_value=httpx.Response(
+            201,
+            json={
+                "id": "f2d6b24e6002f25f77ae89c2b0a5987daa6d0bf751b858b94b8416e8542434d1",
+                "created_at": "2024-07-18T00:35:56.210272Z",
+                "cog_version": "0.9.10",
+                "openapi_schema": {
+                    "openapi": "3.0.2",
+                },
+            },
+        )
+    )
+    router.route(method="GET", path="/assets/hello.txt").mock(
+        return_value=httpx.Response(200, content=b"Hello,")
+    )
+    router.route(method="GET", path="/assets/world.txt").mock(
+        return_value=httpx.Response(200, content=b" world!")
+    )
+
+    client = Client(
+        api_token="test-token", transport=httpx.MockTransport(router.handler)
+    )
+    client.poll_interval = 0.001
+
+    [output1, output2] = cast(
+        list[FileOutput],
+        client.run(
+            "test/example:v1",
+            input={
+                "text": "Hello, world!",
+            },
+            use_file_output=True,
+        ),
+    )
+
+    assert output1.url == "https://api.replicate.com/v1/assets/hello.txt"
+    assert output2.url == "https://api.replicate.com/v1/assets/world.txt"
+
+    assert output1.read() == b"Hello,"
+    assert output2.read() == b" world!"