Update hawk/models/video_llama.py

hawk/models/video_llama.py

@@ -1,5 +1,7 @@
 import logging
 import random
+import requests
+import io
 
 import torch
 from torch.cuda.amp import autocast as autocast
@@ -164,7 +166,7 @@ class VideoLLAMA(Blip2Base):
         logging.info('Loading Q-Former Motion Done')
 
         logging.info('Loading LLAMA Tokenizer')
-        self.llama_tokenizer = LlamaTokenizer.from_pretrained(llama_model, use_fast=False)
+        self.llama_tokenizer = LlamaTokenizer.from_pretrained(llama_model, use_fast=False, subfolder="llama-2-7b-chat-hf")
         if self.llama_tokenizer.pad_token is None:
             self.llama_tokenizer.pad_token = self.llama_tokenizer.unk_token 
         DEFAULT_IMAGE_PATCH_TOKEN = '<ImageHere>'
@@ -179,6 +181,7 @@ class VideoLLAMA(Blip2Base):
         if self.low_resource:
             self.llama_model = LlamaForCausalLM.from_pretrained(
                 llama_model,
+                subfolder="llama-2-7b-chat-hf",
                 torch_dtype=torch.bfloat16,
                 load_in_8bit=True,
                 device_map={'': device_8bit}
@@ -186,6 +189,7 @@ class VideoLLAMA(Blip2Base):
         else:
             self.llama_model = LlamaForCausalLM.from_pretrained(
                 llama_model,
+                subfolder="llama-2-7b-chat-hf",
                 torch_dtype=torch.bfloat16,
             )
 
@@ -844,7 +848,8 @@ class VideoLLAMA(Blip2Base):
         ckpt_path = cfg.get("ckpt", "")  # load weights of MiniGPT-4
         if ckpt_path:
             print("Load first Checkpoint: {}".format(ckpt_path))
-            ckpt = torch.load(ckpt_path, map_location="cpu")
+            response = requests.get(ckpt_path)
+            ckpt = torch.load(io.BytesIO(response.content), map_location="cpu")
             msg = model.load_state_dict(ckpt['model'], strict=False)
         ckpt_path_2 = cfg.get("ckpt_2", "")  
         if ckpt_path_2: