Skip to content

Commit 6d9cad1

Browse files
aronaron
authored
Add experimental replicate.use() function (#438)
This PR introduces a new experimental `use()` function that is intended to make running a model closer to calling a function rather than an API request. Some key differences to `replicate.run()`: 1. You "import" the model using the `use()` syntax, after that you call the model like a function. 2. The output type matches the model definition. i.e. if the model uses an iterator output will be an iterator. 3. Files will be downloaded output as `Path` objects*. * We've replaced the `FileOutput` implementation with `Path` objects. However to avoid unnecessary downloading of files until they are needed we've implemented a `PathProxy` class that will defer the download until the first time the object is used. If you need the underlying URL of the `Path` object you can use the `get_path_url(path: Path) -> str` helper. To use a model: ```py from replicate import use flux_dev = use("black-forest-labs/flux-dev") outputs = flux_dev(prompt="a cat wearing an amusing hat") for output in outputs: print(output) # Path(/tmp/output.webp) ``` Models that implement iterators will return `list | str` types depending on whether they are concatenate iterator instances. Any model can be converted into an iterator by passing `streaming=True`. ```py claude = use("anthropic/claude-4-sonnet", streaming=True) output = claude(prompt="Give me a recipe for tasty smashed avocado on sourdough toast that could feed all of California.") for token in output: print(token) # "Here's a recipe" ``` You can still call `str()` on a language model to get the complete output as a string rather than iterating over tokens: ```py str(output) # "Here's a recipe to feed all of California (about 39 million people)! ..." ``` You can pass the results of one model directly into another, we'll do our best to make this work efficiently: ```py from replicate import use flux_dev = use("black-forest-labs/flux-dev") claude = use("anthropic/claude-4-sonnet") images = flux_dev(prompt="a cat wearing an amusing hat") result = claude(prompt="describe this image for me", image=images[0]) print(str(result)) # "This shows an image of a cat wearing a hat ..." ``` To create a prediction, rather than just getting output, use the `create()` method: ``` from replicate import use claude = use("anthropic/claude-4-sonnet") prediction = claude.create(prompt="Give me a recipe for tasty smashedavocado on sourdough toast that could feed all of California.") prediction.logs() # get current logs (WIP) prediction.output() # get the output ``` You can access the underlying URL for a Path object returned from a model call by using the `get_path_url()` helper. ```py from replicate import use, get_url_path flux_dev = use("black-forest-labs/flux-dev") outputs = flux_dev(prompt="a cat wearing an amusing hat") for output in outputs: print(get_url_path(output)) # "https://replicate.delivery/xyz" ```
1 parent 0110c2c commit 6d9cad1

File tree

9 files changed

+2615
-4
lines changed

9 files changed

+2615
-4
lines changed

‎README.md

Lines changed: 363 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -503,6 +503,369 @@ replicate = Client(
503503
> Never hardcode authentication credentials like API tokens into your code.
504504
> Instead, pass them as environment variables when running your program.
505505
506+
## Experimental `use()` interface
507+
508+
The latest versions of `replicate >= 1.0.8` include a new experimental `use()` function that is intended to make running a model closer to calling a function rather than an API request.
509+
510+
Some key differences to `replicate.run()`.
511+
512+
1. You "import" the model using the `use()` syntax, after that you call the model like a function.
513+
2. The output type matches the model definition.
514+
3. Baked in support for streaming for all models.
515+
4. File outputs will be represented as `PathLike` objects and downloaded to disk when used*.
516+
517+
> [!NOTE]
518+
> \* We've replaced the `FileOutput` implementation with `Path` objects. However to avoid unnecessary downloading of files until they are needed we've implemented a `PathProxy` class that will defer the download until the first time the object is used. If you need the underlying URL of the `Path` object you can use the `get_path_url(path: Path) -> str` helper.
519+
520+
### Examples
521+
522+
To use a model:
523+
524+
> [!IMPORTANT]
525+
> For now `use()` MUST be called in the top level module scope. We may relax this in future.
526+
527+
```py
528+
import replicate
529+
530+
flux_dev = replicate.use("black-forest-labs/flux-dev")
531+
outputs = flux_dev(prompt="a cat wearing an amusing hat")
532+
533+
for output in outputs:
534+
print(output) # Path(/tmp/output.webp)
535+
```
536+
537+
Models that implement iterators will return the output of the completed run as a list unless explicitly streaming (see Streaming section below). Language models that define `x-cog-iterator-display: concatenate` will return strings:
538+
539+
```py
540+
claude = replicate.use("anthropic/claude-4-sonnet")
541+
542+
output = claude(prompt="Give me a recipe for tasty smashed avocado on sourdough toast that could feed all of California.")
543+
544+
print(output) # "Here's a recipe to feed all of California (about 39 million people)! ..."
545+
```
546+
547+
You can pass the results of one model directly into another:
548+
549+
```py
550+
import replicate
551+
552+
flux_dev = replicate.use("black-forest-labs/flux-dev")
553+
claude = replicate.use("anthropic/claude-4-sonnet")
554+
555+
images = flux_dev(prompt="a cat wearing an amusing hat")
556+
557+
result = claude(prompt="describe this image for me", image=images[0])
558+
559+
print(str(result)) # "This shows an image of a cat wearing a hat ..."
560+
```
561+
562+
To create an individual prediction that has not yet resolved, use the `create()` method:
563+
564+
```
565+
claude = replicate.use("anthropic/claude-4-sonnet")
566+
567+
prediction = claude.create(prompt="Give me a recipe for tasty smashed avocado on sourdough toast that could feed all of California.")
568+
569+
prediction.logs() # get current logs (WIP)
570+
571+
prediction.output() # get the output
572+
```
573+
574+
### Streaming
575+
576+
Many models, particularly large language models (LLMs), will yield partial results as the model is running. To consume outputs from these models as they run you can pass the `streaming` argument to `use()`:
577+
578+
```py
579+
claude = replicate.use("anthropic/claude-4-sonnet", streaming=True)
580+
581+
output = claude(prompt="Give me a recipe for tasty smashed avocado on sourdough toast that could feed all of California.")
582+
583+
for chunk in output:
584+
print(chunk) # "Here's a recipe ", "to feed all", " of California"
585+
```
586+
587+
### Downloading file outputs
588+
589+
Output files are provided as Python [os.PathLike](https://docs.python.org/3.12/library/os.html#os.PathLike) objects. These are supported by most of the Python standard library like `open()` and `Path`, as well as third-party libraries like `pillow` and `ffmpeg-python`.
590+
591+
The first time the file is accessed it will be downloaded to a temporary directory on disk ready for use.
592+
593+
Here's an example of how to use the `pillow` package to convert file outputs:
594+
595+
```py
596+
import replicate
597+
from PIL import Image
598+
599+
flux_dev = replicate.use("black-forest-labs/flux-dev")
600+
601+
images = flux_dev(prompt="a cat wearing an amusing hat")
602+
for i, path in enumerate(images):
603+
with Image.open(path) as img:
604+
img.save(f"./output_{i}.png", format="PNG")
605+
```
606+
607+
For libraries that do not support `Path` or `PathLike` instances you can use `open()` as you would with any other file. For example to use `requests` to upload the file to a different location:
608+
609+
```py
610+
import replicate
611+
import requests
612+
613+
flux_dev = replicate.use("black-forest-labs/flux-dev")
614+
615+
images = flux_dev(prompt="a cat wearing an amusing hat")
616+
for path in images:
617+
with open(path, "rb") as f:
618+
r = requests.post("https://api.example.com/upload", files={"file": f})
619+
```
620+
621+
### Accessing outputs as HTTPS URLs
622+
623+
If you do not need to download the output to disk. You can access the underlying URL for a Path object returned from a model call by using the `get_path_url()` helper.
624+
625+
```py
626+
import replicate
627+
from replicate import get_url_path
628+
629+
flux_dev = replicate.use("black-forest-labs/flux-dev")
630+
outputs = flux_dev(prompt="a cat wearing an amusing hat")
631+
632+
for output in outputs:
633+
print(get_url_path(output)) # "https://replicate.delivery/xyz"
634+
```
635+
636+
### Async Mode
637+
638+
By default `use()` will return a function instance with a sync interface. You can pass `use_async=True` to have it return an `AsyncFunction` that provides an async interface.
639+
640+
```py
641+
import asyncio
642+
import replicate
643+
644+
async def main():
645+
flux_dev = replicate.use("black-forest-labs/flux-dev", use_async=True)
646+
outputs = await flux_dev(prompt="a cat wearing an amusing hat")
647+
648+
for output in outputs:
649+
print(Path(output))
650+
651+
asyncio.run(main())
652+
```
653+
654+
When used in streaming mode then an `AsyncIterator` will be returned.
655+
656+
```py
657+
import asyncio
658+
import replicate
659+
660+
async def main():
661+
claude = replicate.use("anthropic/claude-3.5-haiku", streaming=True, use_async=True)
662+
output = await claude(prompt="say hello")
663+
664+
# Stream the response as it comes in.
665+
async for token in output:
666+
print(token)
667+
668+
# Wait until model has completed. This will return either a `list` or a `str` depending
669+
# on whether the model uses AsyncIterator or ConcatenateAsyncIterator. You can check this
670+
# on the model schema by looking for `x-cog-display: concatenate`.
671+
print(await output)
672+
673+
asyncio.run(main())
674+
```
675+
676+
### Typing
677+
678+
By default `use()` knows nothing about the interface of the model. To provide a better developer experience we provide two methods to add type annotations to the function returned by the `use()` helper.
679+
680+
**1. Provide a function signature**
681+
682+
The use method accepts a function signature as an additional `hint` keyword argument. When provided it will use this signature for the `model()` and `model.create()` functions.
683+
684+
```py
685+
# Flux takes a required prompt string and optional image and seed.
686+
def hint(*, prompt: str, image: Path | None = None, seed: int | None = None) -> str: ...
687+
688+
flux_dev = use("black-forest-labs/flux-dev", hint=hint)
689+
output1 = flux_dev() # will warn that `prompt` is missing
690+
output2 = flux_dev(prompt="str") # output2 will be typed as `str`
691+
```
692+
693+
**2. Provide a class**
694+
695+
The second method requires creating a callable class with a `name` field. The name will be used as the function reference when passed to `use()`.
696+
697+
```py
698+
class FluxDev:
699+
name = "black-forest-labs/flux-dev"
700+
701+
def __call__( self, *, prompt: str, image: Path | None = None, seed: int | None = None ) -> str: ...
702+
703+
flux_dev = use(FluxDev)
704+
output1 = flux_dev() # will warn that `prompt` is missing
705+
output2 = flux_dev(prompt="str") # output2 will be typed as `str`
706+
```
707+
708+
> [!WARNING]
709+
> Currently the typing system doesn't correctly support the `streaming` flag for models that return lists or use iterators. We're working on improvements here.
710+
711+
In future we hope to provide tooling to generate and provide these models as packages to make working with them easier. For now you may wish to create your own.
712+
713+
### API Reference
714+
715+
The Replicate Python Library provides several key classes and functions for working with models in pipelines:
716+
717+
#### `use()` Function
718+
719+
Creates a callable function wrapper for a Replicate model.
720+
721+
```py
722+
def use(
723+
ref: FunctionRef,
724+
*,
725+
streaming: bool = False,
726+
use_async: bool = False
727+
) -> Function | AsyncFunction
728+
729+
def use(
730+
ref: str,
731+
*,
732+
hint: Callable | None = None,
733+
streaming: bool = False,
734+
use_async: bool = False
735+
) -> Function | AsyncFunction
736+
```
737+
738+
**Parameters:**
739+
740+
| Parameter | Type | Default | Description |
741+
|-----------|------|---------|-------------|
742+
| `ref` | `str \| FunctionRef` | Required | Model reference (e.g., "owner/model" or "owner/model:version") |
743+
| `hint` | `Callable \| None` | `None` | Function signature for type hints |
744+
| `streaming` | `bool` | `False` | Return OutputIterator for streaming results |
745+
| `use_async` | `bool` | `False` | Return AsyncFunction instead of Function |
746+
747+
**Returns:**
748+
- `Function` - Synchronous model wrapper (default)
749+
- `AsyncFunction` - Asynchronous model wrapper (when `use_async=True`)
750+
751+
#### `Function` Class
752+
753+
A synchronous wrapper for calling Replicate models.
754+
755+
**Methods:**
756+
757+
| Method | Signature | Description |
758+
|--------|-----------|-------------|
759+
| `__call__()` | `(*args, **inputs) -> Output` | Execute the model and return final output |
760+
| `create()` | `(*args, **inputs) -> Run` | Start a prediction and return Run object |
761+
762+
**Properties:**
763+
764+
| Property | Type | Description |
765+
|----------|------|-------------|
766+
| `openapi_schema` | `dict` | Model's OpenAPI schema for inputs/outputs |
767+
| `default_example` | `dict \| None` | Default example inputs (not yet implemented) |
768+
769+
#### `AsyncFunction` Class
770+
771+
An asynchronous wrapper for calling Replicate models.
772+
773+
**Methods:**
774+
775+
| Method | Signature | Description |
776+
|--------|-----------|-------------|
777+
| `__call__()` | `async (*args, **inputs) -> Output` | Execute the model and return final output |
778+
| `create()` | `async (*args, **inputs) -> AsyncRun` | Start a prediction and return AsyncRun object |
779+
780+
**Properties:**
781+
782+
| Property | Type | Description |
783+
|----------|------|-------------|
784+
| `openapi_schema()` | `async () -> dict` | Model's OpenAPI schema for inputs/outputs |
785+
| `default_example` | `dict \| None` | Default example inputs (not yet implemented) |
786+
787+
#### `Run` Class
788+
789+
Represents a running prediction with access to output and logs.
790+
791+
**Methods:**
792+
793+
| Method | Signature | Description |
794+
|--------|-----------|-------------|
795+
| `output()` | `() -> Output` | Get prediction output (blocks until complete) |
796+
| `logs()` | `() -> str \| None` | Get current prediction logs |
797+
798+
**Behavior:**
799+
- When `streaming=True`: Returns `OutputIterator` immediately
800+
- When `streaming=False`: Waits for completion and returns final result
801+
802+
#### `AsyncRun` Class
803+
804+
Asynchronous version of Run for async model calls.
805+
806+
**Methods:**
807+
808+
| Method | Signature | Description |
809+
|--------|-----------|-------------|
810+
| `output()` | `async () -> Output` | Get prediction output (awaits completion) |
811+
| `logs()` | `async () -> str \| None` | Get current prediction logs |
812+
813+
#### `OutputIterator` Class
814+
815+
Iterator wrapper for streaming model outputs.
816+
817+
**Methods:**
818+
819+
| Method | Signature | Description |
820+
|--------|-----------|-------------|
821+
| `__iter__()` | `() -> Iterator[T]` | Synchronous iteration over output chunks |
822+
| `__aiter__()` | `() -> AsyncIterator[T]` | Asynchronous iteration over output chunks |
823+
| `__str__()` | `() -> str` | Convert to string (concatenated or list representation) |
824+
| `__await__()` | `() -> List[T] \| str` | Await all results (string for concatenate, list otherwise) |
825+
826+
#### `URLPath` Class
827+
828+
A path-like object that downloads files on first access.
829+
830+
**Methods:**
831+
832+
| Method | Signature | Description |
833+
|--------|-----------|-------------|
834+
| `__fspath__()` | `() -> str` | Get local file path (downloads if needed) |
835+
| `__str__()` | `() -> str` | String representation of local path |
836+
837+
**Usage:**
838+
- Compatible with `open()`, `pathlib.Path()`, and most file operations
839+
- Downloads file automatically on first filesystem access
840+
- Cached locally in temporary directory
841+
842+
#### `get_path_url()` Function
843+
844+
Helper function to extract original URLs from `URLPath` objects.
845+
846+
```py
847+
def get_path_url(path: Any) -> str | None
848+
```
849+
850+
**Parameters:**
851+
852+
| Parameter | Type | Description |
853+
|-----------|------|-------------|
854+
| `path` | `Any` | Path object (typically `URLPath`) |
855+
856+
**Returns:**
857+
- `str` - Original URL if path is a `URLPath`
858+
- `None` - If path is not a `URLPath` or has no URL
859+
860+
### TODO
861+
862+
There are several key things still outstanding:
863+
864+
1. Support for streaming text when available (rather than polling)
865+
2. Support for streaming files when available (rather than polling)
866+
3. Support for cleaning up downloaded files.
867+
4. Support for streaming logs using `OutputIterator`.
868+
506869
## Development
507870

508871
See [CONTRIBUTING.md](CONTRIBUTING.md)

‎pyproject.toml

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ dev-dependencies = [
3434

3535
[tool.pytest.ini_options]
3636
asyncio_mode = "auto"
37+
asyncio_default_fixture_loop_scope = "function"
3738
testpaths = "tests/"
3839

3940
[tool.setuptools]
@@ -73,8 +74,6 @@ ignore = [
7374
"ANN001", # Missing type annotation for function argument
7475
"ANN002", # Missing type annotation for `*args`
7576
"ANN003", # Missing type annotation for `**kwargs`
76-
"ANN101", # Missing type annotation for self in method
77-
"ANN102", # Missing type annotation for cls in classmethod
7877
"ANN401", # Dynamically typed expressions (typing.Any) are disallowed in {name}
7978
"W191", # Indentation contains tabs
8079
"UP037", # Remove quotes from type annotation

0 commit comments

Comments
 (0)