Skip to content

Commit

Permalink
fix: change onnx and trt default model name to ViT-B-32::openai (#793)
Browse files Browse the repository at this point in the history
* fix: change onnx default model name to ViT-B-32::openai

* fix: trt default model name

* fix: update trt model list

* fix: trt file path
  • Loading branch information
ZiniuYu authored Aug 3, 2022
1 parent 3402b1d commit e48a7a3
Show file tree
Hide file tree
Showing 5 changed files with 23 additions and 5 deletions.
2 changes: 1 addition & 1 deletion server/clip_server/executors/clip_onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
class CLIPEncoder(Executor):
def __init__(
self,
name: str = 'ViT-B/32',
name: str = 'ViT-B-32::openai',
device: Optional[str] = None,
num_worker_preprocess: int = 4,
minibatch_size: int = 32,
Expand Down
2 changes: 1 addition & 1 deletion server/clip_server/executors/clip_tensorrt.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
class CLIPEncoder(Executor):
def __init__(
self,
name: str = 'ViT-B/32',
name: str = 'ViT-B-32::openai',
device: str = 'cuda',
num_worker_preprocess: int = 4,
minibatch_size: int = 32,
Expand Down
2 changes: 1 addition & 1 deletion server/clip_server/model/clip_onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ def __init__(self, name: str, model_path: str = None):
if name in _MODELS:
if not model_path:
cache_dir = os.path.expanduser(
f'~/.cache/clip/{name.replace("/", "-")}'
f'~/.cache/clip/{name.replace("/", "-").replace("::", "-")}'
)
textual_model_name, textual_model_md5 = _MODELS[name][0]
self._textual_path = download_model(
Expand Down
18 changes: 17 additions & 1 deletion server/clip_server/model/clip_trt.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,20 @@
from clip_server.model.clip_onnx import _MODELS as ONNX_MODELS

_MODELS = [
'RN50::openai',
'RN50::yfcc15m',
'RN50::cc12m',
'RN101::openai',
'RN101::yfcc15m',
'RN50x4::openai',
'ViT-B-32::openai',
'ViT-B-32::laion2b_e16',
'ViT-B-32::laion400m_e31',
'ViT-B-32::laion400m_e32',
'ViT-B-16::openai',
'ViT-B-16::laion400m_e31',
'ViT-B-16::laion400m_e32',
# older version name format
'RN50',
'RN101',
'RN50x4',
Expand All @@ -41,7 +55,9 @@ def __init__(
super().__init__(name)

if name in _MODELS:
cache_dir = os.path.expanduser(f'~/.cache/clip/{name.replace("/", "-")}')
cache_dir = os.path.expanduser(
f'~/.cache/clip/{name.replace("/", "-").replace("::", "-")}'
)

self._textual_path = os.path.join(
cache_dir,
Expand Down
4 changes: 3 additions & 1 deletion tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,9 @@ def make_flow(port_generator, request):
f = Flow(port=port_generator()).add(
name=request.param,
uses=CLIPEncoder,
uses_with={'model_path': os.path.expanduser('~/.cache/clip/ViT-B-32')},
uses_with={
'model_path': os.path.expanduser('~/.cache/clip/ViT-B-32-openai')
},
)
with f:
yield f
Expand Down

0 comments on commit e48a7a3

Please sign in to comment.