Update README.md

README.md

@@ -30,23 +30,34 @@ from PIL import Image
 
 import torch
 from transformers import AutoProcessor
-from heron.models.git_llm.git_llama import GitLlamaForCausalLM
+from heron.models.git_llm.git_llama import GitLlamaConfig, GitLlamaForCausalLM
 
 device_id = 0
 
 # prepare a pretrained model
-model = GitLlamaForCausalLM.from_pretrained('turing-motors/heron-chat-git-Llama-2-7b-v0')
+MODEL_NAME = 'turing-motors/heron-chat-git-Llama-2-7b-v0'
+
+git_config = GitLlamaConfig.from_pretrained(MODEL_NAME)
+git_config.set_vision_configs(
+    num_image_with_embedding=1, vision_model_name=git_config.vision_model_name
+)
+model = GitLlamaForCausalLM.from_pretrained(
+    MODEL_NAME, config=git_config, torch_dtype=torch.float16
+)
+
+
+model = GitLlamaForCausalLM.from_pretrained(MODEL_NAME)
 model.eval()
 model.to(f"cuda:{device_id}")
 
 # prepare a processor
-processor = AutoProcessor.from_pretrained('turing-motors/heron-chat-git-Llama-2-7b-v0')
+processor = AutoProcessor.from_pretrained(MODEL_NAME)
 
 # prepare inputs
 url = "https://www.barnorama.com/wp-content/uploads/2016/12/03-Confusing-Pictures.jpg"
 image = Image.open(requests.get(url, stream=True).raw)
 
-text = f"##human: What is this picture？\n##gpt: "
+text = f"##human: What is this picture?\n##gpt: "
 
 # do preprocessing
 inputs = processor(