Views
No views yet
1from transformers import BitsAndBytesConfig, AutoModelForVision2Seq, AutoProcessor
2from transformers.image_utils import load_image
3
4processor = AutoProcessor.from_pretrained("smishr-18/Idefics2-OCR", do_image_splitting=False)
5
6bnb_config = BitsAndBytesConfig(
7 load_in_4bit=True,
8 bnb_4bit_quant_type="nf4",
9 bnb_4bit_compute_dtype=torch.float16
10)
11
12device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
13model = AutoModelForVision2Seq.from_pretrained(
14 "smishr-18/Idefics2-OCR",
15 quantization_config=bnb_config,
16 device_map=device,
17 low_cpu_mem_usage=True
18 )
19
20image = load_image("https://images.pokemontcg.io/pl1/1_hires.png")
21
22messages = [
23 {
24 "role": "user",
25 "content": [
26 {"type": "text", "text": "Explain."},
27 {"type": "image"},
28 {"type": "text", "text": "What is the reflex energy in the image?"}
29 ]
30 }
31]
32
33text = processor.apply_chat_template(messages, add_generation_prompt=True)
34inputs = processor(text=[text.strip()], images=[image4], return_tensors="pt", padding=True)
35inputs = {k: v.to(device) for k, v in inputs.items()}
36
37# Generate texts
38generated_ids = model.generate(**inputs, max_new_tokens=500)
39generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)
40print(generated_texts)
41# The reflex energy in the image is 70.torch.cuda.get_device_capability()[0] >= 8 or Ampere GPUs.