deepseek-ai

/

DeepSeek-OCR

like 2.72k

Follow

DeepSeek 106k

Hello,
From this notebook: https://colab.research.google.com/drive/1zT5-1waOC7PUrn9OJhDMiJcqG6vWaB9u?usp=sharing
Replace the third cell in Google Colab by this code (you can remove the second cell as flash-attention2 isn't supported on T4) and it should work properly:

from transformers import AutoModel, AutoTokenizer
import torch
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
model_name = 'deepseek-ai/DeepSeek-OCR'

if torch.cuda.is_available():
    print(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")
    device = "cuda"
else:
    print("CUDA is not available. Using CPU.")
    device = "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, _attn_implementation='eager', trust_remote_code=True, use_safetensors=True, torch_dtype=torch.bfloat16, device_map=device)
model = model.eval()

Thank you so much! It worked perfectly and solved my issue.