Update models/videochat.py

models/videochat.py

@@ -51,7 +51,13 @@ class VideoChat(Blip2Base):
 
         self.tokenizer = self.init_tokenizer()
         self.low_resource = low_resource
-
+        self.llama_model = LlamaForCausalLM.from_pretrained(
+                    llama_model_path,
+                    torch_dtype=torch.float16,
+                    use_auth_token=os.environ["HF_TOKEN"],
+                    load_in_8bit=True,
+                    device_map="auto"
+                )
         self.vit_precision = vit_precision
         print(f'Loading VIT. Use fp16: {vit_precision}')
         self.visual_encoder, self.ln_vision = self.init_vision_encoder(
@@ -133,6 +139,7 @@ class VideoChat(Blip2Base):
                     use_auth_token=os.environ["HF_TOKEN"],
                 )
         else:
+                '''    
                 self.llama_model = LlamaForCausalLM.from_pretrained(
                     llama_model_path,
                     torch_dtype=torch.float16,
@@ -140,6 +147,7 @@ class VideoChat(Blip2Base):
                     load_in_8bit=True,
                     device_map="auto"
                 )
+                '''
 
         print("freeze LLAMA")
         for name, param in self.llama_model.named_parameters():