Dockerize runtime

Dockerfile

@@ -0,0 +1,51 @@
+FROM registry.access.redhat.com/ubi9/ubi-minimal:latest as base
+
+RUN microdnf update -y && \
+    microdnf install -y \
+        python3-devel python-pip && \
+    pip install --upgrade --no-cache-dir pip wheel && \
+    microdnf clean all
+
+FROM base as builder
+WORKDIR /build
+
+RUN pip install --no-cache tox
+COPY README.md .
+COPY pyproject.toml .
+COPY tox.ini .
+COPY caikit_computer_vision caikit_computer_vision
+# .git is required for setuptools-scm get the version
+RUN --mount=source=.git,target=.git,type=bind \
+    --mount=type=cache,target=/root/.cache/pip \
+     tox -e build
+
+
+FROM base as deploy
+
+RUN python -m venv --upgrade-deps /opt/caikit/
+
+ENV VIRTUAL_ENV=/opt/caikit
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+COPY --from=builder /build/dist/caikit_computer_vision*.whl /tmp/
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install /tmp/caikit_computer_vision*.whl && \
+    rm /tmp/caikit_computer_vision*.whl
+
+COPY LICENSE /opt/caikit/
+COPY README.md /opt/caikit/
+
+RUN groupadd --system caikit --gid 1001 && \
+    adduser --system --uid 1001 --gid 0 --groups caikit \
+    --home-dir /caikit --shell /sbin/nologin \
+    --comment "Caikit User" caikit
+
+USER caikit
+
+ENV RUNTIME_LIBRARY=caikit_computer_vision
+# Optional: use `CONFIG_FILES` and the /caikit/ volume to explicitly provide a configuration file and models
+# ENV CONFIG_FILES=/caikit/caikit.yml
+VOLUME ["/caikit/"]
+WORKDIR /caikit
+
+CMD ["python"]

caikit_computer_vision/data_model/__init__.py

@@ -18,3 +18,4 @@
 from .image_classification import *
 from .image_segmentation import *
 from .object_detection import *
+from .text_to_image import *

caikit_computer_vision/data_model/tasks.py

@@ -26,6 +26,7 @@
 from .image_classification import ImageClassificationResult
 from .image_segmentation import ImageSegmentationResult
 from .object_detection import ObjectDetectionResult
+from .text_to_image import TextToImageResult
 
 
 # TODO - add support for image DM primitives
@@ -61,3 +62,14 @@ class ImageSegmentationTask(TaskBase):
     Note that at the moment, this task encapsulates all segmentation types,
     I.e., instance, object, semantic, etc...
     """
+
+
+@task(
+    required_parameters={"inputs": str},
+    output_type=TextToImageResult,
+)
+class TextToImageTask(TaskBase):
+    """The text to image task is responsible for taking an input text prompt, along with
+    other optional image generation parameters, e.g., image height and width,
+    and generating an image.
+    """

caikit_computer_vision/data_model/text_to_image.py

@@ -0,0 +1,36 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Data structures for text to image."""
+
+
+# Standard
+from typing import List
+
+# Third Party
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
+
+# First Party
+from caikit.core import DataObjectBase, dataobject
+from caikit.interfaces.common.data_model import ProducerId
+from caikit.interfaces.vision import data_model as caikit_dm
+import alog
+
+log = alog.use_channel("DATAM")
+
+
+@dataobject(package="caikit_data_model.caikit_computer_vision")
+class TextToImageResult(DataObjectBase):
+    # TODO: Align on the output format
+    output: Annotated[caikit_dm.Image, FieldNumber(1)]
+    producer_id: Annotated[ProducerId, FieldNumber(2)]

caikit_computer_vision/modules/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Local
-from . import object_detection, segmentation
+from . import object_detection, segmentation, text_to_image

caikit_computer_vision/modules/text_to_image/__init__.py

@@ -0,0 +1,16 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Local
+from .tti_stub import TTIStub

caikit_computer_vision/modules/text_to_image/tti_stub.py

@@ -0,0 +1,76 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Stub module for text to image for testing runtime interfaces.
+"""
+# Standard
+from typing import Union, get_args
+import os
+
+# Third Party
+import numpy as np
+
+# First Party
+from caikit.core.modules import ModuleBase, ModuleConfig, ModuleSaver, module
+from caikit.interfaces.vision import data_model as caikit_dm
+import alog
+
+# Local
+from ...data_model import TextToImageResult
+from ...data_model.tasks import TextToImageTask
+
+log = alog.use_channel("TTI_STUB")
+
+
+@module(
+    id="28aa938b-1a33-11a0-11a3-bb9c3b1cbb11",
+    name="Stub module for Text to Image",
+    version="0.1.0",
+    task=TextToImageTask,
+)
+class TTIStub(ModuleBase):
+    def __init__(
+        self,
+        model_name,
+    ) -> "TTIStub":
+        log.debug("STUB - initializing text to image instance")
+        super().__init__()
+        self.model_name = model_name
+
+    @classmethod
+    def load(cls, model_path: Union[str, "ModuleConfig"]) -> "TTIStub":
+        config = ModuleConfig.load(model_path)
+        return cls.bootstrap(config.model_name)
+
+    @classmethod
+    def bootstrap(cls, model_name: str) -> "TTIStub":
+        return cls(model_name)
+
+    def save(self, model_path: str):
+        saver = ModuleSaver(
+            self,
+            model_path=model_path,
+        )
+        with saver:
+            saver.update_config({"model_name": self.model_name})
+
+    def run(self, inputs: str, height: int, width: int) -> TextToImageResult:
+        """Generates an image matching the provided height and width."""
+        log.debug("STUB - running text to image inference")
+        r_channel = np.full((height, width), 0, dtype=np.uint8)
+        g_channel = np.full((height, width), 100, dtype=np.uint8)
+        b_channel = np.full((height, width), 200, dtype=np.uint8)
+        img = np.stack((r_channel, g_channel, b_channel), axis=2)
+        return TextToImageResult(
+            output=caikit_dm.Image(img),
+        )

tests/modules/text_to_image/test_tti_stub.py

@@ -0,0 +1,41 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Standard
+from tempfile import TemporaryDirectory
+import os
+
+# Local
+from caikit_computer_vision.modules.text_to_image import TTIStub
+import caikit_computer_vision
+
+
+def test_tti_stub():
+    """Ensure that the stubs for load / save / run work as expected."""
+    # Make sure we can bootstrap a model
+    model = TTIStub.bootstrap("foo")
+    assert isinstance(model, TTIStub)
+
+    # Make sure we can run a fake inference on it
+    pred = model.run("This is a prompt", height=500, width=550)
+    pil_img = pred.output.as_pil()
+    assert pil_img.width == 550
+    assert pil_img.height == 500
+
+    # Make sure we can save the model
+    model_dirname = "my_model"
+    with TemporaryDirectory() as tmpdirname:
+        model_path = os.path.join(tmpdirname, model_dirname)
+        model.save(model_path)
+        reloaded_model = model.load(model_path)

tox.ini

@@ -42,4 +42,11 @@ passenv =
 setenv =
     FLIT_USERNAME = __token__
 commands = flit publish
-skip_install = True
+skip_install = True
+
+[testenv:build]
+description = build wheel
+deps =
+    build
+commands = python -m build
+skip_install = True