save memory usage

app.py

@@ -32,7 +32,7 @@ def violates_moderation(text):
     """
     if "OPENAI_API_KEY" not in os.environ:
         print("OPENAI_API_KEY not found, skip content moderation check...")
-        return True
+        return False
     
     url = "https://api.openai.com/v1/moderations"
     headers = {
@@ -206,7 +206,7 @@ def generate(state: Conversation, temperature, top_p, max_new_tokens, add_knwl,
     # generate output
     prompt = state.get_prompt().replace("USER: <image>\n", "")
     prompt = prompt.split("USER:")[-1].replace("ASSISTANT:", "")
-    image_pt = image_trans(image).to(device).unsqueeze(0)
+    image_pt = gptk_trans(image).to(device).unsqueeze(0)
     samples = {"image": image_pt, "knowledge": knwl_embd, "prompt": prompt}
     if bool(do_beam_search):
         new_text = gptk_model.generate(
@@ -358,21 +358,12 @@ def build_demo():
     return demo
 
 
-def build_model():
-    if torch.cuda.is_available():
-        device = torch.device("cuda")
-    else:
-        device = torch.device("cpu")
-
-    query_enc, _, query_trans = open_clip.create_model_and_transforms(
-        "ViT-g-14", pretrained="laion2b_s34b_b88k", precision='fp16'
-    )
-    query_enc = query_enc.to(device).eval()
-
+def build_knowledge():
     def get_knwl(knowledge_db):
         knwl_db = TextDB(Path(knowledge_db)/"knowledge_db.hdf5")
+        knwl_db.feature = knwl_db.feature
         knwl_idx = faiss.read_index(str(Path(knowledge_db)/"faiss.index"))
-        knwl_idx.add(knwl_db.feature)
+        knwl_idx.add(knwl_db.feature.astype(np.float32))
 
         return knwl_db, knwl_idx
 
@@ -381,19 +372,38 @@ def build_model():
         "act": get_knwl('knowledge/(dataset-action)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)'),
         "attr": get_knwl('knowledge/(dataset-attribute)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)'),
     }
-    d_knwl = knwl_db["obj"][0].feature.shape[-1]
 
-    _, image_trans = get_gptk_image_transform()
+    return knwl_db
+
+
+def build_query_model():
+    query_enc, _, query_trans = open_clip.create_model_and_transforms(
+        "ViT-g-14", pretrained="laion2b_s34b_b88k", precision='fp16', device=device
+    )
+    query_enc = query_enc.eval()
+
+    return query_enc, query_trans
+
+
+def build_gptk_model():
+    _, gptk_trans = get_gptk_image_transform()
     topk = {"whole": 60, "five": 24, "nine": 16}
-    gptk_model = get_gptk_model(d_knwl=d_knwl, topk=topk)
+    gptk_model = get_gptk_model(d_knwl=1024, topk=topk)
     gptk_ckpt = "model/ckpt/gptk-vicuna7b.pt"
     gptk_ckpt = torch.load(gptk_ckpt, map_location="cpu")
     gptk_model.load_state_dict(gptk_ckpt, strict=False)
     gptk_model = gptk_model.to(device).eval()
 
-    return knwl_db, query_enc, query_trans, gptk_model, image_trans, topk, device
+    return gptk_model, gptk_trans, topk
+
 
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
 
-knwl_db, query_enc, query_trans, gptk_model, image_trans, topk, device = build_model()
+gptk_model, gptk_trans, topk = build_gptk_model()
+query_enc, query_trans = build_query_model()
+knwl_db = build_knowledge()
 demo = build_demo()
 demo.queue().launch()

knowledge/text_db.py

@@ -18,7 +18,7 @@ class TextDB:
             _, d = f[f"0/feature"].shape
 
         with h5py.File(text_db, 'r') as f:
-            feature = np.zeros((db_size, d), dtype=np.float32)
+            feature = np.zeros((db_size, d), dtype=np.float16)
             text = []
             N = 0
             for i in tqdm(range(len(f)), desc="Load text DB", dynamic_ncols=True, mininterval=1.0):

model/gptk.py

@@ -7,8 +7,7 @@ import torch
 from torch import nn
 from torchvision import transforms as T
 from torchvision.transforms.functional import InterpolationMode
-from transformers import LlamaTokenizer
-from transformers import BertTokenizer
+from transformers import LlamaTokenizer, BertTokenizer, BitsAndBytesConfig
 
 import sys
 sys.path.append("./")
@@ -49,7 +48,17 @@ class GPTK(nn.Module):
         llm_config = LlamaConfig.from_pretrained(llm_model)
         llm_config.gradient_checkpointing = True
         llm_config.use_cache = True
-        self.llm_model = LlamaForCausalLM.from_pretrained(llm_model, config=llm_config, torch_dtype=torch.float16)
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4'
+        )
+        self.llm_model = LlamaForCausalLM.from_pretrained(
+            llm_model, config=llm_config, torch_dtype=torch.float16, quantization_config=quantization_config
+        )
         self.llm_tokenizer.add_special_tokens({'pad_token': '[PAD]'})
         self.llm_tokenizer.add_special_tokens({'bos_token': '</s>'})
         self.llm_tokenizer.add_special_tokens({'eos_token': '</s>'})

requirements.txt

@@ -1,7 +1,7 @@
---extra-index-url https://download.pytorch.org/whl/cu113
-torch==1.11.0+cu113
-torchvision==0.12.0+cu113
-torchaudio==0.11.0
+--extra-index-url https://download.pytorch.org/whl/cu117
+torch==1.13.1+cu117
+torchvision==0.14.1+cu117
+torchaudio==0.13.1 
 
 transformers==4.30.2
 faiss-gpu==1.7.2
@@ -11,3 +11,5 @@ open_clip_torch
 omegaconf
 h5py>=3.8.0
 spacy>=3.5.0
+bitsandbytes
+accelerate