Update README.md

README.md

@@ -81,7 +81,11 @@ pip install torch transformers==4.44.0
 from transformers import LlavaNextForConditionalGeneration,LlavaNextProcessor
 import torch
 
-model = LlavaNextForConditionalGeneration.from_pretrained('Bllossom/llama-3.1-Korean-Bllossom-Vision-8B', torch_dtype=torch.bfloat16)
+model = LlavaNextForConditionalGeneration.from_pretrained(
+  'Bllossom/llama-3.1-Korean-Bllossom-Vision-8B',
+  torch_dtype=torch.bfloat16,
+  device_map='auto'
+)
 processor = LlavaNextProcessor.from_pretrained('Bllossom/llama-3.1-Korean-Bllossom-Vision-8B')
 
 with torch.no_grad():
@@ -106,7 +110,7 @@ with torch.no_grad():
     )
     
     bos_token = processor.tokenizer.bos_token_id
-    chat_messages = torch.cat([torch.tensor([[bos_token]]),chat_messages],dim=-1)
+    chat_messages = torch.cat([torch.tensor([[bos_token]]),chat_messages],dim=-1).to(model.device)
     
 
     output = model.generate(
@@ -127,7 +131,11 @@ from PIL import Image
 from transformers import LlavaNextForConditionalGeneration,LlavaNextProcessor
 import torch
 
-model = LlavaNextForConditionalGeneration.from_pretrained('Bllossom/llama-3.1-Korean-Bllossom-Vision-8B', torch_dtype=torch.bfloat16)
+model = LlavaNextForConditionalGeneration.from_pretrained(
+  'Bllossom/llama-3.1-Korean-Bllossom-Vision-8B',
+  torch_dtype=torch.bfloat16,
+  device_map='auto'
+)
 processor = LlavaNextProcessor.from_pretrained('Bllossom/llama-3.1-Korean-Bllossom-Vision-8B')
 
 image = Image.open('[IMAGE_PATH]').convert('RGB')
@@ -144,7 +152,11 @@ messages = [
     {'role': 'user', 'content': f"<image>\n{instruction}"}
 ]
 
-chat_messages = processor.tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True)
+chat_messages = processor.tokenizer.apply_chat_template(
+  messages,
+  tokenize=False,
+  add_generation_prompt=True
+).to(model.device)
 
 inputs = processor(
     chat_messages,