From 8d4d3c54bbdd0e5a4334a6e888ae1c67797f0b62 Mon Sep 17 00:00:00 2001
From: Kentaro Wada <www.kentaro.wada@gmail.com>
Date: Tue, 30 Jul 2024 20:27:48 +0900
Subject: [PATCH] Make `osam serve` optional

---
 README.md        |  8 ++++++++
 osam/__main__.py | 13 ++++++++++---
 pyproject.toml   |  6 ++++--
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 4fce9a5..fed8ea5 100644
--- a/README.md
+++ b/README.md
@@ -30,6 +30,12 @@
 pip install osam
 ```
 
+**For `osam serve`**:
+
+```bash
+pip install osam[serve]
+```
+
 ## Quickstart
 
 To run with EfficientSAM:
@@ -102,6 +108,8 @@ PIL.Image.fromarray(response.mask).save("mask.png")
 ### HTTP
 
 ```bash
+# pip install osam[serve]  # required for `osam serve`
+
 # Get up the server
 osam serve
 
diff --git a/osam/__main__.py b/osam/__main__.py
index d0d3694..28cc56f 100644
--- a/osam/__main__.py
+++ b/osam/__main__.py
@@ -8,7 +8,6 @@
 import imgviz
 import numpy as np
 import PIL.Image
-import uvicorn
 from loguru import logger
 
 from . import __version__
@@ -97,8 +96,16 @@ def rm(model_name):
 @cli.command(help="Start server")
 @click.option("--reload", is_flag=True, help="reload server on file changes")
 def serve(reload):
-    logger.info("Starting server...")
-    uvicorn.run("osam._server:app", host="127.0.0.1", port=11368, reload=reload)
+    try:
+        import uvicorn
+
+        import osam._server  # noqa: F401
+
+        logger.info("Starting server...")
+        uvicorn.run("osam._server:app", host="127.0.0.1", port=11368, reload=reload)
+    except ImportError:
+        logger.error("Run `pip install osam[serve]` to use `osam serve`")
+        sys.exit(1)
 
 
 @cli.command(help="Run a model")
diff --git a/pyproject.toml b/pyproject.toml
index 724e7aa..121ac36 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,17 +32,19 @@ classifiers = [
 dependencies = [
   "click",
   "gdown",
-  "fastapi",
   "imgviz",
   "loguru",
   "onnxruntime",
   "Pillow",
   "pydantic",
-  "uvicorn",
 ]
 dynamic = ["readme", "version"]
 
 [project.optional-dependencies]
+serve = [
+  "fastapi",
+  "uvicorn",
+]
 test = [
   "build",
   "mypy",