add support for use_file_output

Author: dtmeadows

src/replicate/lib/_files.py

@@ -8,6 +8,8 @@
 
 import httpx
 
+from replicate.types.prediction_output import PredictionOutput
+
 from .._utils import is_mapping, is_sequence
 from .._client import ReplicateClient, AsyncReplicateClient
 
@@ -124,7 +126,7 @@ def __repr__(self) -> str:
         return f'{self.__class__.__name__}("{self.url}")'
 
 
-def transform_output(value: Any, client: ReplicateClient | AsyncReplicateClient) -> Any:
+def transform_output(value: PredictionOutput, client: ReplicateClient | AsyncReplicateClient) -> Any:
     """
     Transform the output of a prediction to a `FileOutput` object if it's a URL.
     """

src/replicate/lib/_predictions.py

@@ -21,7 +21,7 @@ def run(
     ref: Union[Model, Version, ModelVersionIdentifier, str],
     *,
     wait: Union[int, bool, NotGiven] = NOT_GIVEN,
-    _use_file_output: Optional[bool] = True,
+    use_file_output: Optional[bool] = True,
     **params: Unpack[PredictionCreateParamsWithoutVersion],
 ) -> PredictionOutput | FileOutput | Iterable[FileOutput] | Dict[str, FileOutput]:
     """
@@ -109,15 +109,18 @@ def run(
 
     # TODO: Return an iterator for completed output if the model has an output iterator array type.
 
-    return transform_output(prediction.output, client)  # type: ignore[no-any-return]
+    if use_file_output:
+        return transform_output(prediction.output, client) # type: ignore[no-any-return]
+
+    return prediction.output
 
 
 async def async_run(
     client: "AsyncReplicateClient",
     ref: Union[Model, Version, ModelVersionIdentifier, str],
     *,
     wait: Union[int, bool, NotGiven] = NOT_GIVEN,
-    _use_file_output: Optional[bool] = True,
+    use_file_output: Optional[bool] = True,
     **params: Unpack[PredictionCreateParamsWithoutVersion],
 ) -> PredictionOutput | FileOutput | Iterable[FileOutput] | Dict[str, FileOutput]:
     """
@@ -205,4 +208,7 @@ async def async_run(
 
     # TODO: Return an iterator for completed output if the model has an output iterator array type.
 
-    return transform_output(prediction.output, client)  # type: ignore[no-any-return]
+    if use_file_output:
+        return transform_output(prediction.output, client) # type: ignore[no-any-return]
+
+    return prediction.output