Support UnionTypes in Optional check

[8W9aG]

• Support type hints like str | None in inputs
• Main issue is the detection of whether something is Optional, solved by extending the Union detection to UnionType

[aron]

This is close, we still have an issue with the cog predict command and the validation.

Given this predictor:

from cog import BasePredictor, Input


class Predictor(BasePredictor):
    def setup(self):
        self.prefix = "hello"

    def predict(
        self,
        text: str | None = Input(
            description="Text to prefix with 'hello '", default=None
        ),
    ) -> str:
        return self.prefix + " " + text

Running cog predict gives:

Starting Docker image cog-hello-world-base and running setup()...
Error while loading predictor:

Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/cog/server/http.py", line 159, in create_app
    InputType, OutputType, is_async = cog_config.get_predictor_types(
  File "/usr/local/lib/python3.8/site-packages/cog/config.py", line 171, in get_predictor_types
    predictor = self._load_predictor_for_types(
  File "/usr/local/lib/python3.8/site-packages/cog/config.py", line 147, in _load_predictor_for_types
    module = load_full_predictor_from_file(module_path, module_name)
  File "/usr/local/lib/python3.8/site-packages/cog/predictor.py", line 151, in load_full_predictor_from_file
    spec.loader.exec_module(module)
  File "<frozen importlib._bootstrap_external>", line 843, in exec_module
  File "<frozen importlib._bootstrap>", line 219, in _call_with_frames_removed
  File "predict.py", line 4, in <module>
    class Predictor(BasePredictor):
  File "predict.py", line 10, in Predictor
    text: str | None = Input(
TypeError: unsupported operand type(s) for |: 'type' and 'NoneType'

ⅹ Model setup failed

[8W9aG]

This is close, we still have an issue with the cog predict command and the validation.

Given this predictor:

from cog import BasePredictor, Input


class Predictor(BasePredictor):
    def setup(self):
        self.prefix = "hello"

    def predict(
        self,
        text: str | None = Input(
            description="Text to prefix with 'hello '", default=None
        ),
    ) -> str:
        return self.prefix + " " + text

Running cog predict gives:

Starting Docker image cog-hello-world-base and running setup()...
Error while loading predictor:

Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/cog/server/http.py", line 159, in create_app
    InputType, OutputType, is_async = cog_config.get_predictor_types(
  File "/usr/local/lib/python3.8/site-packages/cog/config.py", line 171, in get_predictor_types
    predictor = self._load_predictor_for_types(
  File "/usr/local/lib/python3.8/site-packages/cog/config.py", line 147, in _load_predictor_for_types
    module = load_full_predictor_from_file(module_path, module_name)
  File "/usr/local/lib/python3.8/site-packages/cog/predictor.py", line 151, in load_full_predictor_from_file
    spec.loader.exec_module(module)
  File "<frozen importlib._bootstrap_external>", line 843, in exec_module
  File "<frozen importlib._bootstrap>", line 219, in _call_with_frames_removed
  File "predict.py", line 4, in <module>
    class Predictor(BasePredictor):
  File "predict.py", line 10, in Predictor
    text: str | None = Input(
TypeError: unsupported operand type(s) for |: 'type' and 'NoneType'

ⅹ Model setup failed

Can you provide the cog.yaml associated with this? Then I'll make an integration test and solve against that.

[aron]

I just changed the cog-examples hello world version.

build:
  python_version: "3.10"
predict: "predict.py:Predictor"

The above issue was python 3.8 which makes sense because unions wern't supported.

With 3.10 I get:

Error while loading predictor:

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/site-packages/cog/server/http.py", line 159, in create_app
    InputType, OutputType, is_async = cog_config.get_predictor_types(
  File "/usr/local/lib/python3.10/site-packages/cog/config.py", line 180, in get_predictor_types
    get_input_type(predictor),
  File "/usr/local/lib/python3.10/site-packages/cog/predictor.py", line 309, in get_input_type
    **get_input_create_model_kwargs(signature),
  File "/usr/local/lib/python3.10/site-packages/cog/predictor.py", line 229, in get_input_create_model_kwargs
    validate_input_type(InputType, name)
  File "/usr/local/lib/python3.10/site-packages/cog/predictor.py", line 209, in validate_input_type
    validate_input_type(t, name)
  File "/usr/local/lib/python3.10/site-packages/cog/predictor.py", line 216, in validate_input_type
    raise TypeError(
TypeError: Unsupported input type NoneType for parameter `text`. Supported input types are: str, int, float, bool, cog.File, cog.Path, cog.Secret, or a Union or List of those types.

ⅹ Model setup failed

[aron]

Just checked this out again locally, it works! 👍