Use Pydantic for type annotations

[bfirsh]

This is an implementation of #193 and #205 to use Pydantic as the type annotations for model predictors and FastAPI for serving the model via HTTP.

Models now look a bit like this:

from cog import BasePredictor, Input, Path
import torch

class ColorizationPredictor(BasePredictor):
    def setup(self):
        self.model = torch.load("./weights.pth")

    def predict(self,
          image: Path = Input(title="Grayscale input image")
    ) -> Path:
        processed_input = preprocess(image)
        output = self.model(processed_input)
        return postprocess(output)

This is a follow-up of #378, where the initial working version was made.

At some point, we want to lock down and define the HTTP API, but we are deferring that to a future piece of work. The Redis queue worker remains unchanged and supports the old API for Replicate to run these models.

The branch has now been moved to this repository so we can collaborate together by opening pull requests against it. We aren't putting this work in main because the documentation for the current version of Cog is there.

Upgrading from Cog 0.0 to 0.1

[Note: This documentation should be copied to the release notes when we merge this.]

In Cog versions up to 0.0.20 you described inputs using annotations on your predict method. For example:

@cog.input("image", type=Path, help="Image to enlarge")
@cog.input("scale", type=float, default=1.5, help="Factor to scale image by")
def predict(self, image, scale):
    ...

From Cog 0.1.0 onwards, we've started using Pydantic to define input and output types. Rather than describing inputs using annotations, you now describe them with type hinting. Here's how you'd define the same inputs now:

def predict(self,
    image: Path = Input(description="Image to enlarge"),
    scale: float = Input(description="Factor to scale image by", default=1.5)
) -> Path:
    ...

The parameters that Input() takes are pretty similar to those @cog.input() used to take. Here are the differences:

• It no longer takes a type parameter; use a type hint instead.
• The help parameter has been renamed to description.
• The options parameter has been renamed to choice.
• The min option has been replaced by two options:
  • ge: greater than or equal to (direct replacement)
  • gt: greater than (a new alternative)
• The max option has been replaced by two options:
  • le: less than or equal to (direct replacement)
  • lt: less than (a new alternative)

The other major difference is that you now need to define the output type of your model. That's the -> Path bit in the example above. That might be a simple type like str, float or bool. If you need to handle multiple outputs, check out the new documentation for complex output objects. If you only have a single output, but it can be of different types depending on the run, you can use typing.Any:

from typing import Any

def predict(self) -> Any:
    ...

Todo

• Disallow users from passing extra options to cog.Input(). Remove non-essential arguments to cog.Input() #408
• Make input validation errors work for cog predict in a simple way until we lock down HTTP API. Better error message when providing invalid inputs #420
• Force users to pick an output type and throw an error if it isn't set. This is because the validation error when it isn't set is really weird, but this also seems sensible anyway. (Couldn't just be Any, but they need to be explicit!) Display better error when output type is undefined #421
• Document how to migrate @input() to Pydantic type annotations.
• Make input errors work for Redis API & Replicate Add a test for error messages from invalid input #434
  • If there is a validation error, push_error() in the redis worker should be called with an appropriate human-readable error.
  • This call all be done in Cog-land by writing an integration test for the redis worker.
• Clean up temporary files -- both input and output: Clean up temporary input files #449
• Strictly limit the allowed model input types #429
• document cog.Path and cog.File #445
• The name of the output type on a model Ben has locally is wrong. Needs investigation. Only set custom Output object if it's not already called Output #454
• Merge examples PR Update examples for new Cog types cog-examples#8
• Write release notes: https://gist.github.com/zeke/0ff121335e8c331e4343683696277ec5

Future

• Decide on and define nouns #171
• Pluralization of "input" and "output" in new Cog #428
• Only allow particular output types #456
• What happens if you use pathlib.Path as an input type? Does it work? Should we throw an error? (This might happen if you import wrong package / collide imports) [It breaks! Strictly limit the allowed model input types #429 prevents it being used as an input type]. See also Update pydantic Cog to serialize pathlib Paths #444
• What happens if you return pathlib.Path from a predict function? Does this work? Should it throw an error? Only allow particular output types #456
• Do we want to make cog predict with an existing image backwards compatible? ('cos there are old models on Replicate) With a shim, this will just work.
• Tell users to upgrade if they use @cog.input() with new Cog #427
• Remove redundant title in OpenAPI schema #457
• Set correct cog version in OpenAPI schema #446
• Fix any TODOs/FIXMEs #459

(These should be migrated to issues.) Done!

Refs

Closes #58
Closes #113
Closes #193
Closes #205
Closes #212
Closes #246
Closes #304
Closes #306
Closes #328

[zeke]

👏🏼 PY 👏🏼 DAN 👏🏼 TIC!

[bfirsh]

All rebased and merge-conflict-resolved.