Update script.py

script.py

@@ -25,6 +25,16 @@ from io import BytesIO
 import re
 
 
+# Line 138 uncomment the cuda() to use GPUs
+
+device = "cpu"
+# device = "auto"
+
+prompt = "What are the things I should be cautious about when I visit here?"
+image_file = "Great-Room-4.jpg"
+
+model_path = "liuhaotian/llava-v1.6-mistral-7b"
+
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
@@ -53,17 +63,8 @@ def load_images(image_files):
         out.append(image)
     return out
 
-model_path = "liuhaotian/llava-v1.6-mistral-7b"
-model_name = get_model_name_from_path(model_path)
-# tokenizer = AutoTokenizer.from_pretrained(model_path)
-# model = LlavaMistralForCausalLM.from_pretrained(
-#                     model_path,
-#                     low_cpu_mem_usage=True,
-#                     # offload_folder="/content/sample_data"
-#                 )
 
-prompt = "What are the things I should be cautious about when I visit here?"
-image_file = "Great-Room-4.jpg"
+model_name = get_model_name_from_path(model_path)
 
 args = type('Args', (), {
     "model_path": model_path,
@@ -79,9 +80,8 @@ args = type('Args', (), {
     "max_new_tokens": 512
 })()
 
-
 tokenizer, model, image_processor, context_len = load_pretrained_model(
-        model_path, None, model_name, device_map="cpu"
+        model_path, None, model_name, device_map=device
     )
 
 qs = args.query