Update README.md

README.md

@@ -49,21 +49,33 @@ See https://llava-vl.github.io/ for more details.
 usage is as follows 
 
 ```python
-from transformers import LlavaProcessor, LlavaLlamaForCausalLM
+from transformers import LlavaProcessor, LlavaForCausalLM
+from PIL import Image
+import requests
+import torch
 
-PATH_TO_CONVERTED_WEIGHTS = "shauray/Llava-Llama-2-13B-hf"
+PATH_TO_CONVERTED_WEIGHTS = "shauray/Llava-Llama-2-7B-hf"
 
-model = LlavaLlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
-processor = LlavaProcessor.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+model = LlavaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS,
+device_map="cuda",torch_dtype=torch.float16).to("cuda")
+processor = LlavaProcessor.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
 
 url = "https://llava-vl.github.io/static/images/view.jpg"
 image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
-prompt = "How would you best describe the image given?"
-inputs = processor(text=prompt, images=image, return_tensors="pt")
+prompt = "How can you best describe this image?"
 
+inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda",
+torch.float16)
 # Generate
-generate_ids = model.generate(**inputs, max_length=30)
-tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0]
+generate_ids = model.generate(**inputs, 
+    do_sample=True,
+    max_length=1024,
+    temperature=0.1,
+    top_p=0.9,
+)
+out = processor.decode(generate_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
+
+print(out)
 
 """The photograph shows a wooden dock floating on the water, with mountains in the background. It is an idyllic scene that captures both
 nature and human-made structures at their finest moments of beauty or tranquility depending upon one's perspective as they gaze into it"""