Update README.md

README.md

@@ -38,17 +38,12 @@ import requests
 from PIL import Image
 
 import torch
-from transformers import AutoProcessor, LlavaForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForPreTraining
 
 
 model_id = "OpenFace-CQUPT/Human_LLaVA"
 cuda = 0
-model = LlavaForConditionalGeneration.from_pretrained(
-    model_id, 
-    torch_dtype=torch.float16, 
-    low_cpu_mem_usage=True, 
-    trust_remote_code=True
-).to(cuda)
+model = AutoModelForPreTraining.from_pretrained("OpenFace-CQUPT/Human_LLaVA", torch_dtype=torch.float16).to(cuda)
 
 processor = AutoProcessor.from_pretrained(model_id,trust_remote_code=True)