Python: Add Ollama text to image service

python/semantic_kernel/connectors/ai/README.md

@@ -48,5 +48,6 @@ All base clients inherit from the [`AIServiceClientBase`](../../services/ai_serv
 | Ollama | [`OllamaChatCompletion`](./ollama/services/ollama_chat_completion.py) |
 |        | [`OllamaTextCompletion`](./ollama/services/ollama_text_completion.py) |
 |        | [`OllamaTextEmbedding`](./ollama/services/ollama_text_embedding.py) |
+|        | [`OllamaTextToImage`](./ollama/services/ollama_text_to_image.py) |
 | Onnx | [`OnnxGenAIChatCompletion`](./onnx/services/onnx_gen_ai_chat_completion.py) |
 |      | [`OnnxGenAITextCompletion`](./onnx/services/onnx_gen_ai_text_completion.py) |

python/semantic_kernel/connectors/ai/ollama/__init__.py

@@ -5,10 +5,12 @@
     OllamaEmbeddingPromptExecutionSettings,
     OllamaPromptExecutionSettings,
     OllamaTextPromptExecutionSettings,
+    OllamaTextToImagePromptExecutionSettings,
 )
 from semantic_kernel.connectors.ai.ollama.services.ollama_chat_completion import OllamaChatCompletion
 from semantic_kernel.connectors.ai.ollama.services.ollama_text_completion import OllamaTextCompletion
 from semantic_kernel.connectors.ai.ollama.services.ollama_text_embedding import OllamaTextEmbedding
+from semantic_kernel.connectors.ai.ollama.services.ollama_text_to_image import OllamaTextToImage
 
 __all__ = [
     "OllamaChatCompletion",
@@ -18,4 +20,6 @@
     "OllamaTextCompletion",
     "OllamaTextEmbedding",
     "OllamaTextPromptExecutionSettings",
+    "OllamaTextToImage",
+    "OllamaTextToImagePromptExecutionSettings",
 ]

python/semantic_kernel/connectors/ai/ollama/ollama_prompt_execution_settings.py

@@ -37,3 +37,11 @@ class OllamaChatPromptExecutionSettings(OllamaPromptExecutionSettings):
 
 class OllamaEmbeddingPromptExecutionSettings(OllamaPromptExecutionSettings):
     """Settings for Ollama embedding prompt execution."""
+
+
+class OllamaTextToImagePromptExecutionSettings(OllamaPromptExecutionSettings):
+    """Settings for Ollama text to image execution."""
+
+    width: int | None = None
+    height: int | None = None
+    steps: int | None = None

python/semantic_kernel/connectors/ai/ollama/ollama_settings.py

@@ -20,6 +20,7 @@ class OllamaSettings(KernelBaseSettings):
     - chat_model_id: str - The chat model ID. (Env var OLLAMA_CHAT_MODEL_ID)
     - text_model_id: str - The text model ID. (Env var OLLAMA_TEXT_MODEL_ID)
     - embedding_model_id: str - The embedding model ID. (Env var OLLAMA_EMBEDDING_MODEL_ID)
+    - image_model_id: str - The image generation model ID. (Env var OLLAMA_IMAGE_MODEL_ID)
 
     Optional settings for prefix 'OLLAMA' are:
     - host: HttpsUrl - The endpoint of the Ollama service. (Env var OLLAMA_HOST)
@@ -30,4 +31,5 @@ class OllamaSettings(KernelBaseSettings):
     chat_model_id: str | None = None
     text_model_id: str | None = None
     embedding_model_id: str | None = None
+    image_model_id: str | None = None
     host: str | None = None

python/semantic_kernel/connectors/ai/ollama/services/ollama_text_to_image.py

@@ -0,0 +1,147 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import base64
+import logging
+import sys
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any
+from warnings import warn
+
+from ollama import AsyncClient
+from pydantic import ValidationError
+
+from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import (
+    OllamaTextToImagePromptExecutionSettings,
+)
+from semantic_kernel.connectors.ai.ollama.ollama_settings import OllamaSettings
+from semantic_kernel.connectors.ai.ollama.services.ollama_base import OllamaBase
+from semantic_kernel.connectors.ai.text_to_image_client_base import TextToImageClientBase
+from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError
+from semantic_kernel.utils.feature_stage_decorator import experimental
+
+if TYPE_CHECKING:
+    from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
+
+if sys.version_info >= (3, 12):
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@experimental
+class OllamaTextToImage(OllamaBase, TextToImageClientBase):
+    """Ollama text to image client.
+
+    Make sure to have the ollama service running either locally or remotely, with an
+    image generation model pulled, for example `x/z-image-turbo`.
+    """
+
+    def __init__(
+        self,
+        service_id: str | None = None,
+        ai_model_id: str | None = None,
+        host: str | None = None,
+        client: AsyncClient | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+    ) -> None:
+        """Initialize an OllamaTextToImage service.
+
+        Args:
+            service_id (Optional[str]): Service ID tied to the execution settings. (Optional)
+            ai_model_id (Optional[str]): The model name. (Optional)
+            host (Optional[str]): URL of the Ollama server, defaults to None and
+                will use the default Ollama service address: http://127.0.0.1:11434. (Optional)
+            client (Optional[AsyncClient]): A custom Ollama client to use for the service. (Optional)
+            env_file_path (str | None): Use the environment settings file as a fallback to using env vars.
+            env_file_encoding (str | None): The encoding of the environment settings file, defaults to 'utf-8'.
+        """
+        try:
+            ollama_settings = OllamaSettings(
+                image_model_id=ai_model_id,
+                host=host,
+                env_file_path=env_file_path,
+                env_file_encoding=env_file_encoding,
+            )
+        except ValidationError as ex:
+            raise ServiceInitializationError("Failed to create Ollama settings.", ex) from ex
+
+        if not ollama_settings.image_model_id:
+            raise ServiceInitializationError("Ollama image model ID is not set.")
+
+        super().__init__(
+            service_id=service_id or ollama_settings.image_model_id,
+            ai_model_id=ollama_settings.image_model_id,
+            client=client or AsyncClient(host=ollama_settings.host),
+        )
+
+    @override
+    async def generate_image(
+        self,
+        description: str,
+        width: int | None = None,
+        height: int | None = None,
+        settings: "PromptExecutionSettings | None" = None,
+        **kwargs: Any,
+    ) -> bytes:
+        """Generate an image from a text description.
+
+        Args:
+            description: Description of the image.
+            width: Deprecated, use settings.width instead.
+            height: Deprecated, use settings.height instead.
+            settings: Execution settings for the prompt.
+            kwargs: Additional arguments passed to the Ollama generate endpoint.
+
+        Returns:
+            bytes: The raw image bytes.
+        """
+        image_settings = (
+            OllamaTextToImagePromptExecutionSettings()
+            if settings is None
+            else OllamaTextToImagePromptExecutionSettings.from_prompt_execution_settings(settings)
+        )
+
+        if width is not None:
+            warn(
+                "The 'width' argument is deprecated. Use 'settings.width' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if image_settings.width is None:
+                image_settings.width = width
+        if height is not None:
+            warn(
+                "The 'height' argument is deprecated. Use 'settings.height' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if image_settings.height is None:
+                image_settings.height = height
+
+        options = image_settings.prepare_settings_dict()
+        options.update(kwargs)
+
+        response_object = await self.client.generate(
+            model=self.ai_model_id,
+            prompt=description,
+            stream=False,
+            **options,
+        )
+
+        image = getattr(response_object, "image", None)
+        if image is None and isinstance(response_object, Mapping):
+            image = response_object.get("image")
+        if not image:
+            raise ServiceInvalidResponseError(
+                "The Ollama response did not contain image data. Make sure the configured model "
+                f"('{self.ai_model_id}') is an image generation model."
+            )
+
+        return base64.b64decode(image)
+
+    @override
+    def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]:
+        return OllamaTextToImagePromptExecutionSettings

python/tests/unit/connectors/ai/ollama/conftest.py

@@ -56,6 +56,7 @@ def ollama_unit_test_env(monkeypatch, host, exclude_list):
         "OLLAMA_CHAT_MODEL_ID": "test_chat_model_id",
         "OLLAMA_TEXT_MODEL_ID": "test_text_model_id",
         "OLLAMA_EMBEDDING_MODEL_ID": "test_embedding_model_id",
+        "OLLAMA_IMAGE_MODEL_ID": "test_image_model_id",
         "OLLAMA_HOST": host,
     }
 

python/tests/unit/connectors/ai/ollama/services/test_ollama_text_to_image.py

@@ -0,0 +1,138 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import base64
+from unittest.mock import patch
+
+import pytest
+
+from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import (
+    OllamaTextToImagePromptExecutionSettings,
+)
+from semantic_kernel.connectors.ai.ollama.services.ollama_text_to_image import OllamaTextToImage
+from semantic_kernel.contents.image_content import ImageContent
+from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError
+
+ENCODED_IMAGE = base64.b64encode(b"test_image_bytes").decode()
+
+
+def test_init_empty_service_id(model_id):
+    """Test that the service initializes correctly with an empty service id."""
+    ollama = OllamaTextToImage(ai_model_id=model_id)
+    assert ollama.service_id == model_id
+
+
+def test_custom_client(model_id, custom_client):
+    """Test that the service initializes correctly with a custom client."""
+    ollama = OllamaTextToImage(ai_model_id=model_id, client=custom_client)
+    assert ollama.client == custom_client
+
+
+def test_invalid_ollama_settings():
+    """Test that the service initializes incorrectly with invalid settings."""
+    with pytest.raises(ServiceInitializationError):
+        _ = OllamaTextToImage(ai_model_id=123)
+
+
+@pytest.mark.parametrize("exclude_list", [["OLLAMA_IMAGE_MODEL_ID"]], indirect=True)
+def test_init_empty_model_id(ollama_unit_test_env):
+    """Test that the service initializes incorrectly with an empty model id."""
+    with pytest.raises(ServiceInitializationError):
+        _ = OllamaTextToImage(env_file_path="fake_env_file_path.env")
+
+
+@patch("ollama.AsyncClient.__init__", return_value=None)  # mock_client
+@patch("ollama.AsyncClient.generate")  # mock_generate
+async def test_custom_host(mock_generate, mock_client, model_id, host, prompt):
+    """Test that the service generates an image correctly with a custom host."""
+    mock_generate.return_value = {"image": ENCODED_IMAGE}
+
+    ollama = OllamaTextToImage(ai_model_id=model_id, host=host)
+    _ = await ollama.generate_image(prompt)
+
+    mock_client.assert_called_once_with(host=host)
+
+
+@patch("ollama.AsyncClient.generate")
+async def test_generate_image(mock_generate, model_id, prompt):
+    """Test that the service decodes the base64 image returned by Ollama."""
+    mock_generate.return_value = {"image": ENCODED_IMAGE}
+    settings = OllamaTextToImagePromptExecutionSettings()
+    settings.options = {"test_key": "test_value"}
+
+    ollama = OllamaTextToImage(ai_model_id=model_id)
+    image = await ollama.generate_image(prompt, settings=settings)
+
+    assert image == b"test_image_bytes"
+    mock_generate.assert_called_once_with(
+        model=model_id,
+        prompt=prompt,
+        stream=False,
+        options={"test_key": "test_value"},
+    )
+
+
+@patch("ollama.AsyncClient.generate")
+async def test_get_image_content(mock_generate, model_id, prompt):
+    """Test that the inherited get_image_content returns ImageContent with the image data."""
+    mock_generate.return_value = {"image": ENCODED_IMAGE}
+
+    ollama = OllamaTextToImage(ai_model_id=model_id)
+    content = await ollama.get_image_content(prompt, OllamaTextToImagePromptExecutionSettings())
+
+    assert isinstance(content, ImageContent)
+    assert content.data == b"test_image_bytes"
+
+
+@patch("ollama.AsyncClient.generate")
+async def test_generate_image_with_size_settings(mock_generate, model_id, prompt):
+    """Test that width, height and steps from the settings are forwarded to Ollama."""
+    mock_generate.return_value = {"image": ENCODED_IMAGE}
+    settings = OllamaTextToImagePromptExecutionSettings(width=512, height=256, steps=4)
+
+    ollama = OllamaTextToImage(ai_model_id=model_id)
+    _ = await ollama.generate_image(prompt, settings=settings)
+
+    call_kwargs = mock_generate.call_args.kwargs
+    assert call_kwargs["width"] == 512
+    assert call_kwargs["height"] == 256
+    assert call_kwargs["steps"] == 4
+
+
+@patch("ollama.AsyncClient.generate")
+async def test_generate_image_deprecated_width_and_height_arguments(mock_generate, model_id, prompt):
+    """Test that the deprecated width and height arguments still reach Ollama, with a warning."""
+    mock_generate.return_value = {"image": ENCODED_IMAGE}
+
+    ollama = OllamaTextToImage(ai_model_id=model_id)
+    with pytest.warns(DeprecationWarning):
+        image = await ollama.generate_image(prompt, width=512, height=256)
+
+    assert image == b"test_image_bytes"
+    call_kwargs = mock_generate.call_args.kwargs
+    assert call_kwargs["width"] == 512
+    assert call_kwargs["height"] == 256
+
+
+@patch("ollama.AsyncClient.generate")
+async def test_generate_image_settings_take_precedence_over_arguments(mock_generate, model_id, prompt):
+    """Test that explicit settings win over the deprecated width and height arguments."""
+    mock_generate.return_value = {"image": ENCODED_IMAGE}
+    settings = OllamaTextToImagePromptExecutionSettings(width=1024, height=1024)
+
+    ollama = OllamaTextToImage(ai_model_id=model_id)
+    with pytest.warns(DeprecationWarning):
+        _ = await ollama.generate_image(prompt, width=512, height=256, settings=settings)
+
+    call_kwargs = mock_generate.call_args.kwargs
+    assert call_kwargs["width"] == 1024
+    assert call_kwargs["height"] == 1024
+
+
+@patch("ollama.AsyncClient.generate")
+async def test_generate_image_without_image_in_response(mock_generate, model_id, prompt):
+    """Test that a response without image data raises instead of returning empty bytes."""
+    mock_generate.return_value = {"response": "this model returns text"}
+
+    ollama = OllamaTextToImage(ai_model_id=model_id)
+    with pytest.raises(ServiceInvalidResponseError):
+        await ollama.generate_image(prompt)