瀏覽代碼

add img embedding.py

luojunhui 4 周之前
父節點
當前提交
f2aba3bc61
共有 1 個文件被更改,包括 2 次插入2 次删除
  1. 2 2
      applications/clip_embedding/clip_model.py

+ 2 - 2
applications/clip_embedding/clip_model.py

@@ -1,11 +1,11 @@
 import os
 import os
 import torch
 import torch
 
 
-from transformers import AutoModel, AutoConfig, CLIPImageProcessor, AutoTokenizer
+from transformers import AutoModel, AutoConfig, AutoTokenizer
 
 
 MODEL_NAME = "BAAI/EVA-CLIP-8B"
 MODEL_NAME = "BAAI/EVA-CLIP-8B"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
+DTYPE = torch.int8 if DEVICE == "cuda" else torch.float32
 MAX_BATCH = int(os.getenv("MAX_BATCH", "32"))
 MAX_BATCH = int(os.getenv("MAX_BATCH", "32"))
 
 
 print(f"[model_config] Loading {MODEL_NAME} on {DEVICE} dtype={DTYPE} ...")
 print(f"[model_config] Loading {MODEL_NAME} on {DEVICE} dtype={DTYPE} ...")