WIP: New provider for huggingface_image

docs/source/users/index.md

@@ -42,11 +42,14 @@ Jupyter AI supports the following model providers:
 | AI21                | `ai21`               | `AI21_API_KEY`             | `ai21`                          |
 | Anthropic           | `anthropic`          | `ANTHROPIC_API_KEY`        | `anthropic`                     |
 | Cohere              | `cohere`             | `COHERE_API_KEY`           | `cohere`                        |
-| HuggingFace Hub     | `huggingface_hub`    | `HUGGINGFACEHUB_API_TOKEN` | `huggingface_hub`, `ipywidgets` |
+| HuggingFace Hub     | `huggingface_hub`, `huggingface_image` | `HUGGINGFACEHUB_API_TOKEN` | `huggingface_hub`, `ipywidgets` |
 | OpenAI              | `openai`             | `OPENAI_API_KEY`           | `openai`                        |
 | OpenAI (chat)       | `openai-chat`        | `OPENAI_API_KEY`           | `openai`                        |
 | SageMaker Endpoints | `sagemaker-endpoint` | N/A                        | `boto3`                         |
 
+The `huggingface_hub` provider is for models that output text. The `huggingface_image`
+provider is for models such as Stable Diffusion that output images.
+
 To use SageMaker's models, you will need to authenticate via
 [boto3](https://github.com/boto/boto3).
 

examples/images.ipynb

@@ -0,0 +1,78 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "24f3f446-2b1d-4802-a47c-d298c06fc86e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%load_ext jupyter_ai"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2e178db2-55f7-45c4-a63c-9892ddf61495",
+   "metadata": {},
+   "source": [
+    "## Rapidly experimenting with different HF Hub models\n",
+    "We can call models on HuggingFace Hub that return images:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "0837ad88-7759-42e8-894c-82747c4bb3ce",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'Cannot determine model provider from model ID huggingface_image:stabilityai/stable-diffusion-2-1.'"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "%%ai huggingface_image:stabilityai/stable-diffusion-2-1\n",
+    "It's an astronaut with a boombox"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "73239301-af47-4c98-9911-2687008d0bab",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

packages/jupyter-ai-magics/jupyter_ai_magics/huggingface_image.py

@@ -0,0 +1,108 @@
+"""Wrapper around HuggingFace APIs that produce images."""
+from typing import Any, Dict, List, Mapping, Optional
+
+from pydantic import Extra, root_validator
+
+from langchain.llms.base import LLM
+from langchain.llms.utils import enforce_stop_tokens
+from langchain.utils import get_from_dict_or_env
+
+DEFAULT_REPO_ID = "stabilityai/stable-diffusion-2-1"
+VALID_TASKS = ("text-to-image")
+
+
+class HuggingFaceImage(LLM):
+    """Wrapper around HuggingFaceHub models that generate images.
+    Based on HuggingFaceHub, for text models.
+
+    To use, you should have the ``huggingface_hub`` python package installed, and the
+    environment variable ``HUGGINGFACEHUB_API_TOKEN`` set with your API token, or pass
+    it as a named parameter to the constructor.
+
+    Only supports `text-to-image` for now.
+    """
+
+    client: Any  #: :meta private:
+    repo_id: str = DEFAULT_REPO_ID
+    """Model name to use."""
+    task: Optional[str] = None
+    """Task to call the model with. Should be a task that returns `text-to-image`."""
+    model_kwargs: Optional[dict] = None
+    """Key word arguments to pass to the model."""
+
+    huggingfacehub_api_token: Optional[str] = None
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        huggingfacehub_api_token = get_from_dict_or_env(
+            values, "huggingfacehub_api_token", "HUGGINGFACEHUB_API_TOKEN"
+        )
+        try:
+            from huggingface_hub.inference_api import InferenceApi
+
+            repo_id = values["repo_id"]
+            client = InferenceApi(
+                repo_id=repo_id,
+                token=huggingfacehub_api_token,
+                task=values.get("task"),
+            )
+            if client.task not in VALID_TASKS:
+                raise ValueError(
+                    f"Got invalid task {client.task}, "
+                    f"currently only {VALID_TASKS} are supported"
+                )
+            values["client"] = client
+        except ImportError:
+            raise ValueError(
+                "Could not import huggingface_hub python package. "
+                "Please install it with `pip install huggingface_hub`."
+            )
+        return values
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Get the identifying parameters."""
+        _model_kwargs = self.model_kwargs or {}
+        return {
+            **{"repo_id": self.repo_id, "task": self.task},
+            **{"model_kwargs": _model_kwargs},
+        }
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "huggingface_image"
+
+    def _call(self, prompt: str) -> str:
+        """Call out to HuggingFace Hub's inference endpoint.
+
+        Args:
+            prompt: The prompt to pass into the model.
+
+        Returns:
+            The image generated by the model.
+
+        Example:
+            .. code-block:: python
+
+                response = hf("Tell me a joke.")
+        """
+        _model_kwargs = self.model_kwargs or {}
+        response = self.client(inputs=prompt, params=_model_kwargs)
+        if "error" in response:
+            raise ValueError(f"Error raised by inference API: {response['error']}")
+        if self.client.task == "text-to-image":
+            image = response.images[0]
+        else:
+            raise ValueError(
+                f"Got invalid task {self.client.task}, "
+                f"currently only {VALID_TASKS} are supported"
+            )
+
+        return "".join(list(image.getdata()))

packages/jupyter-ai-magics/jupyter_ai_magics/providers.py

@@ -11,6 +11,8 @@
     SagemakerEndpoint
 )
 
+from .huggingface_image import HuggingFaceImage
+
 from pydantic import BaseModel, Extra
 from langchain.chat_models import ChatOpenAI
 
@@ -137,6 +139,17 @@ class HfHubProvider(BaseProvider, HuggingFaceHub):
     pypi_package_deps = ["huggingface_hub", "ipywidgets"]
     auth_strategy = EnvAuthStrategy(name="HUGGINGFACEHUB_API_TOKEN")
 
+class HfImageProvider(BaseProvider, HuggingFaceImage):
+    id = "huggingface_image"
+    name = "HuggingFace Image"
+    models = ["*"]
+    model_id_key = "repo_id"
+    # ipywidgets needed to suppress tqdm warning
+    # https://stackoverflow.com/questions/67998191
+    # tqdm is a dependency of huggingface_hub
+    pypi_package_deps = ["huggingface_hub", "ipywidgets"]
+    auth_strategy = EnvAuthStrategy(name="HUGGINGFACEHUB_API_TOKEN")
+
 class OpenAIProvider(BaseProvider, OpenAI):
     id = "openai"
     name = "OpenAI"

packages/jupyter-ai-magics/pyproject.toml

@@ -51,6 +51,7 @@ ai21 = "jupyter_ai_magics:AI21Provider"
 anthropic = "jupyter_ai_magics:AnthropicProvider"
 cohere = "jupyter_ai_magics:CohereProvider"
 huggingface_hub = "jupyter_ai_magics:HfHubProvider"
+huggingface_image = "jupyter_ai_magics:HfImageProvider"
 openai = "jupyter_ai_magics:OpenAIProvider"
 openai-chat = "jupyter_ai_magics:ChatOpenAIProvider"
 sagemaker-endpoint = "jupyter_ai_magics:SmEndpointProvider"