feat: add mps support

README.md

@@ -112,7 +112,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
@@ -121,7 +121,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 = [
  {

cli_chat.py

@@ -8,7 +8,7 @@
 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)

deepseek_vl/__init__.py

@@ -16,3 +16,14 @@
 # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+
+# check if python version is above 3.10
+import sys
+if sys.version_info >=(3, 10):
+ print("Python version is above 3.10, patching the collections module.")
+ # Monkey patch collections
+ import collections
+ import collections.abc
+ for type_name in collections.abc.__all__:
+ setattr(collections, type_name, getattr(collections.abc, type_name))

deepseek_vl/utils/io.py

@@ -27,14 +27,37 @@
 from deepseek_vl.models import MultiModalityCausalLM, VLChatProcessor
 
 
+def get_device_and_dtype():
+ """
+ Get the device and dtype for the model.
+ """
+
+ if torch.cuda.is_available():
+ print("Using CUDA and BFloat16")
+ device = torch.device("cuda")
+ dtype = torch.bfloat16
+ elif torch.backends.mps.is_available():
+ print("Using MPS and FP16")
+ device = torch.device("mps")
+ dtype = torch.float16
+ else:
+ print("Using CPU and FP32")
+ device = torch.device("cpu")
+ dtype = torch.float32
+
+ return device, dtype
+
+
 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
 

inference.py

@@ -2,7 +2,7 @@
 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 @@
 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 @@
  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)

pyproject.toml

@@ -13,7 +13,7 @@ readme = "README.md"
 requires-python = ">=3.8"
 dependencies = [
  "torch>=2.0.1",
- "transformers>=4.38.2",
+ "transformers==4.38.1",
  "timm>=0.9.16",
  "gradio>=4.13.0",
  "accelerate",

requirements.txt

@@ -1,5 +1,5 @@
 torch>=2.0.1
-transformers>=4.38.2
+transformers==4.38.1
 timm>=0.9.16
 gradio>=4.13.0
 accelerate