Implement `replicate.use()` for Python client

[aron]

The replicate.use() method is the successor to replicate.run() (we want to ensure both exist for now). It is intended as a high level interface over the Replicate HTTP API that provides a function like interface for a model.

For background on this change see: https://github.com/replicate/client-libraries-use-prototype

There exists an implementation in the existing replicate/replicate-python library on the beta branch: https://github.com/replicate/replicate-python/tree/beta

For users currently running pip install replicate --pre they get a pre-release version of replicate that doesn't include the use function. I'd like to fix this, and ideally make the stainless version the canonical implementation.

The README of the beta branch contains the API interface plus examples of usage.

https://github.com/replicate/replicate-python/tree/beta?tab=readme-ov-file#api-reference

Tasks

• Ship a working implementation in a prerelease (See feat: add replicate.use() #37)
• Cut a prerelease
• Take replicate.use() for a spin in the prerelease npm package
• Iron out type checking issues for Python 3.8 vs 3.12 compatibility
• Add tests
• Add docs

[zeke]

@dgellow I'm tentatively assigning this one to you. Let us know if this is something you can work on!

[zeke]

@aron what do we want to do about these?

https://github.com/replicate/replicate-python/blob/974dd7fec28033fe11a611ad32228aebe3f70dbc/replicate/use.py#L1-L3

[dgellow]

@aron I've been going through the (really well written) documents you linked. I noticed the use of the term "cog type". What does that mean exactly? I never heard of such a concept and don't find relevant results online (seems to be something in the golf world though).

Edit: I think I found what Cog is! https://github.com/replicate/cog Though, it's not clear to me what a cog "type" would be.

Edit 2: I think I understand fully now. Cog isn't just a simplified way to define/build docker images, it is also a specific interface for an http server that runs within the container. One endpoint seems to be predict, that endpoint does take an argument image: Path. I believe that Path type from the cog library is what "cog type" refers to

[dgellow]

Looking at: https://github.com/replicate/client-libraries-use-prototype/blob/main/README.md

prediction = flux_dev(prompt="make me a sandwich")

# The input kwargs passed to `flux_dev()`
input = prediction.input

# The prediction status.
status = prediction.status

# An iterator that yields log lines.
for log in prediction.logs():
    print(log)

# The same output as returned from `flux_dev()`
output = prediction.output()

Is a .create() missing in this example? it is mentioned in the prose

Edit: I think I understand. The prose preceding the example isn't related. That example is demonstrating access to prediction properties. Basically, my_model() returns a prediction object that is also an iterable. If you iterate on it directly, you get the output chunks. If you call str() on it, you get the concatenated result (or approximation, depending on what format the model is returning). If you instead try to access its properties, it should behave like a prediction object. If you try to call print() on it it should print a Path object (?).

Though, I'm not sure if that example is expected to work.

Ideally, we would also have a way to retrieve a prediction for a standard function call (maybe a weak-map on output to prediction) but I've not come up with anything good here yet.

Do I understand correctly that specific note is communicating the code example isn't supported yet by the current design?

[zeke]

From our meeting today:

@dgellow is planning to have a working implementation and demo of replicate.use in the Python SDK for Thursday next week. 🎉

[aron]

@dgellow thanks for taking a a run at this.

Regarding implementation and interfaces, I think if you take the use.py as a reference implementation:

https://github.com/replicate/replicate-python/blob/beta/replicate/use.py

And it's accompanying test suite:

https://github.com/replicate/replicate-python/blob/beta/tests/test_use.py

If the stainless implementation can make that test suite pass, I think we'll be most of the way there.

I implemented use.py, without exposing any of the underlying existing Replicate SDK hoping that it would make it easier to implement the same interface on top of stainless. I think the most complex bit will be the output polling/iterators which currently defers to the Prediction class in replicate-python which won't exist in stainless.

I think I understand fully now. Cog isn't just a simplified way to define/build docker images, it is also a specific interface for an http server that runs within the container. One endpoint seems to be predict, that endpoint does take an argument image: Path. I believe that Path type from the cog library is what "cog type" refers to

You got it! Cog types are a set of Python types supported by cog models.

https://cog.run/python/#input-and-output-types

[aron]

@dgellow

Is a .create() missing in this example? it is mentioned in the prose

Not in that example. There are two ways to run a prediction using the use() function.

1. You instantiate a model (Function) and call it directly to get the completed output:

   flux_dev = replicate.use("black-forest-labs/flux-dev")
   outputs = flux_dev(prompt="a unicorn jumping over the moon")

   Here the output should be whatever type the model returns, in this case it's a list of files (or the exact cog interface is list[cog.Path]).
2. You instantiate a model (Function) and use .create() to create an instance of a Run (also called a prediction). This is a pending prediction, that has some additional properties and methods. It's intended to be used if you need to create multiple predictions in a non blocking fashion or need to access prediction attributes like id or logs().

   flux_dev = replicate.use("black-forest-labs/flux-dev")
   run = flux_dev.create(prompt="a unicorn jumping over the moon")
   run.id # "xyz123"
   run.status # "processing"
   run.logs() # logs
   outputs = run.output() # blocks until output is complete.
   
   run.status = "sucessful"

The comment you mention:

Ideally, we would also have a way to retrieve a prediction for a standard function call (maybe a weak-map on output to prediction) but I've not come up with anything good here yet.

This points to a weak point in the design of the function interface (1). We want people to use (1) 99% of the time. However people often want to be able to get the id of the prediction for book keeping. If the function only returns output you can't do this.

flux_dev = replicate.use("black-forest-labs/flux-dev")
outputs = flux_dev(prompt="a unicorn jumping over the moon")

// write to database (but need an identifier…)
db.store(id, outputs)

The workaround is simple you use .create() instead:

flux_dev = replicate.use("black-forest-labs/flux-dev")
run = flux_dev.create(prompt="a unicorn jumping over the moon")
outputs = run.output() # blocks until output is complete.
    
db.store(run.id, outputs)

However it would be nice if we could keep the interface and have a way to still get the Run object later.

flux_dev = replicate.use("black-forest-labs/flux-dev")
outputs = flux_dev(prompt="a unicorn jumping over the moon")

// Lookup the run based on the output instance.
db.store(get_run(outputs).id, outputs)

I don't think this requires further thought. It's just something I've been thinking about.

[zeke]

We want people to use (1) 99% of the time.

Not sure I agree with this.

In my experience, just getting the output from a prediction is for toy apps, while people building "real" apps very often want to also get their hands on metadata about the prediction, like the id, the metrics, the inputs, etc.

We've seen this with replicate.run(), and it's kind of a footgun for new users. replicate.run is the default way of running models that we advertise everywhere on our website and in our docs, so people use it to start out. But they often end up having to rewrite their app to use replicate.predictions.create() to get their hands on metadata about the predictions.

Note that the legacy JavasScript client has an onProgress function that can be used with replicate.run() to get your hands on the prediction data, but it's unintuitive, feels like a hack, and most people don't know about it.

I bring this up to suggest that we should make it simple for users to get full metadata about their predictions, because it's not a niche case in practice.

flux_dev.create() looks like a perfectly adequate proposal for that to me 👍🏼 -- I just want to make sure we don't bury it as a secondary thing, because most serious users will eventually need to reach for it, and ultimately default to it as they use the platform for more serious projects.

[dgellow]

Hello there! Sorry for the delay, I've been out for >1 week following a back injury. I'm back part-time, I'm aiming to have a first prototype by tomorrow early next week (and need to catch up with all your responses).

[dgellow]

Just an idea, but have you considered returning a tuple?

flux_dev = replicate.use("black-forest-labs/flux-dev")
outputs, run = flux_dev(prompt="a unicorn jumping over the moon")
    
db.store(run.id, outputs)

Though if you only care about the outputs, outputs, _ = flux_dev() is maybe a bit awkward. It's easy to forget the second arg and you end up with a tuple instead of the actual outputs :(
But that would address Zeke's concerns. There is one way to do things, and moving from the "easy toy version" to the "real app version" doesn't require to change how you call the method.

[zeke]

The tuple is interesting idea, but like you said it might be a bit awkward if you only want outputs. I think flux() and flux.create() are good options. My point above was just to make sure that we document both approaches prominently.

[zeke]

Updated the body of this issue with a checklist, per #37 (comment)

[dgellow]

Cut a prerelease

@zeke #42 is ready to be merged, once published you should be able to test replicate.use() :)