From 40ec6491d4b3c0667e427622fb185878d3e98127 Mon Sep 17 00:00:00 2001
From: Nicola Dall'Asen <fodark@pm.me>
Date: Wed, 13 Mar 2024 12:32:14 +0100
Subject: [PATCH] load model on detected device and use correct dtype

---
 cli_chat.py             | 6 +++---
 deepseek_vl/utils/io.py | 4 +++-
 inference.py            | 8 +++++---
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/cli_chat.py b/cli_chat.py
index bbea14a..e34a88e 100644
--- a/cli_chat.py
+++ b/cli_chat.py
@@ -8,7 +8,7 @@ from threading import Thread
 import torch
 from transformers import TextIteratorStreamer
 
-from deepseek_vl.utils.io import load_pretrained_model
+from deepseek_vl.utils.io import load_pretrained_model, get_device_and_dtype
 
 
 def load_image(image_file):
@@ -34,13 +34,13 @@ def get_help_message(image_token):
 
 @torch.inference_mode()
 def response(args, conv, pil_images, tokenizer, vl_chat_processor, vl_gpt, generation_config):
-
+    _, dtype = get_device_and_dtype()
     prompt = conv.get_prompt()
     prepare_inputs = vl_chat_processor.__call__(
         prompt=prompt,
         images=pil_images,
         force_batchify=True
-    ).to(vl_gpt.device)
+    ).to(vl_gpt.device, dtype=dtype)
 
     # run image encoder to get the image embeddings
     inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
diff --git a/deepseek_vl/utils/io.py b/deepseek_vl/utils/io.py
index a160f00..06c1a47 100644
--- a/deepseek_vl/utils/io.py
+++ b/deepseek_vl/utils/io.py
@@ -52,10 +52,12 @@ def load_pretrained_model(model_path: str):
     vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
     tokenizer = vl_chat_processor.tokenizer
 
+    device, dtype = get_device_and_dtype()
+
     vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
         model_path, trust_remote_code=True
     )
-    vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
+    vl_gpt = vl_gpt.to(device, dtype=dtype).eval()
 
     return tokenizer, vl_chat_processor, vl_gpt
 
diff --git a/inference.py b/inference.py
index b6fe38c..b22a632 100644
--- a/inference.py
+++ b/inference.py
@@ -2,7 +2,7 @@ import torch
 from transformers import AutoModelForCausalLM
 
 from deepseek_vl.models import VLChatProcessor, MultiModalityCausalLM
-from deepseek_vl.utils.io import load_pil_images
+from deepseek_vl.utils.io import load_pil_images, get_device_and_dtype
 
 
 # specify the path to the model
@@ -11,7 +11,9 @@ vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
 tokenizer = vl_chat_processor.tokenizer
 
 vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
-vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
+
+device, dtype = get_device_and_dtype()
+vl_gpt = vl_gpt.to(dtype).to(device).eval()
 
 conversation = [
     {
@@ -32,7 +34,7 @@ prepare_inputs = vl_chat_processor(
     conversations=conversation,
     images=pil_images,
     force_batchify=True
-).to(vl_gpt.device)
+).to(vl_gpt.device, dtype=dtype)
 
 # run image encoder to get the image embeddings
 inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)