chore: fix manifest imports (#56)

CHANGELOG.rst

@@ -1,8 +1,8 @@
-0.0.3 - Unreleased
+0.1.1 - Unreleased
 ---------------------
 
 
-0.0.2 - 2022-01-31
+0.1.0 - 2022-01-31
 ---------------------
 Added
 ^^^^^

README.md

@@ -17,13 +17,18 @@ Install:
 pip install manifest-ml
 ```
 
-Install with ChatGPT Support:
+Install with diffusion support:
+```bash
+pip install manifest-ml[diffusers]
+```
+
+Install with ChatGPT support:
 ```bash
 pip install manifest-ml[chatgpt]
 ```
 This installs [pyChatGPT](https://github.com/terry3041/pyChatGPT) and uses the ChatGPT session key to start a session. This key must be set as the `CHATGPT_SESSION_KEY` environment variable or passed in with `client_connection`.
 
-Install with HuggingFace API Support:
+Install with HuggingFace local model support:
 ```bash
 pip install manifest-ml[api]
 ```

manifest/manifest.py

@@ -8,14 +8,11 @@ from manifest.caches.noop import NoopCache
 from manifest.caches.redis import RedisCache
 from manifest.caches.sqlite import SQLiteCache
 from manifest.clients.ai21 import AI21Client
-from manifest.clients.chatgpt import ChatGPTClient
 from manifest.clients.cohere import CohereClient
-from manifest.clients.diffuser import DiffuserClient
 from manifest.clients.dummy import DummyClient
 from manifest.clients.huggingface import HuggingFaceClient
 from manifest.clients.openai import OpenAIClient
 from manifest.clients.toma import TOMAClient
-from manifest.clients.toma_diffuser import TOMADiffuserClient
 from manifest.request import Request
 from manifest.response import Response
 from manifest.session import Session
@@ -25,16 +22,35 @@ logger = logging.getLogger(__name__)
 
 CLIENT_CONSTRUCTORS = {
     "openai": OpenAIClient,
-    "chatgpt": ChatGPTClient,
     "cohere": CohereClient,
     "ai21": AI21Client,
     "huggingface": HuggingFaceClient,
-    "diffuser": DiffuserClient,
     "dummy": DummyClient,
     "toma": TOMAClient,
-    "tomadiffuser": TOMADiffuserClient,
 }
 
+# ChatGPT
+try:
+    from manifest.clients.chatgpt import ChatGPTClient
+
+    CLIENT_CONSTRUCTORS["chatgpt"] = ChatGPTClient
+except Exception:
+    logger.info("ChatGPT not installed. Skipping import.")
+    pass
+
+# Diffusion
+DIFFUSION_CLIENTS = ["diffuser", "tomadiffuser"]
+try:
+    from manifest.clients.diffuser import DiffuserClient
+    from manifest.clients.toma_diffuser import TOMADiffuserClient
+
+    CLIENT_CONSTRUCTORS["diffuser"] = DiffuserClient
+    CLIENT_CONSTRUCTORS["tomadiffuser"] = TOMADiffuserClient
+except Exception:
+    logger.info("Diffusion not supported. Skipping import.")
+    pass
+
+
 CACHE_CONSTRUCTORS = {
     "redis": RedisCache,
     "sqlite": SQLiteCache,
@@ -72,10 +88,17 @@ class Manifest:
         Remaining kwargs sent to client and cache.
         """
         if client_name not in CLIENT_CONSTRUCTORS:
-            raise ValueError(
-                f"Unknown client name: {client_name}. "
-                f"Choices are {list(CLIENT_CONSTRUCTORS.keys())}"
-            )
+            if client_name in DIFFUSION_CLIENTS:
+                raise ImportError(
+                    f"Diffusion client {client_name} requires the proper install. "
+                    "Make sure to run `pip install manifest-ml[diffusers]` "
+                    "or install Pillow."
+                )
+            else:
+                raise ValueError(
+                    f"Unknown client name: {client_name}. "
+                    f"Choices are {list(CLIENT_CONSTRUCTORS.keys())}"
+                )
         if cache_name not in CACHE_CONSTRUCTORS:
             raise ValueError(
                 f"Unknown cache name: {cache_name}. "

manifest/version.py

@@ -1 +1 @@
-__version__ = "0.0.2"
+__version__ = "0.1.0"

setup.py

@@ -27,17 +27,24 @@ REQUIRES_PYTHON = ">=3.8.0"
 VERSION = main_ns["__version__"]
 
 # What packages are required for this module to be executed?
-REQUIRED = ["redis>=4.3.1", "requests>=2.27.1", "sqlitedict>=2.0.0", "xxhash>=3.0.0"]
+REQUIRED = [
+    "numpy>=1.20.0",
+    "pydantic>=1.9.0",
+    "redis>=4.3.1",
+    "requests>=2.27.1",
+    "sqlitedict>=2.0.0",
+    "xxhash>=3.0.0",
+]
 
 # What packages are optional?
 EXTRAS = {
     "api": [
+        "deepspeed>=0.7.0",
         "diffusers>=0.6.0",
         "Flask>=2.1.2",
         "accelerate>=0.10.0",
         "transformers>=4.20.0,<4.26.0",
         "torch>=1.8.0",
-        "numpy>=1.20.0",
     ],
     "app": [
         "fastapi>=0.70.0",
@@ -46,6 +53,9 @@ EXTRAS = {
     "chatgpt": [
         "pyChatGPT>=0.4.3",
     ],
+    "diffusers": [
+        "pillow>=9.0.0",
+    ],
     "dev": [
         "autopep8>=1.6.0",
         "black>=22.3.0",