Views
No views yet
1from PIL import Image
2from transformers import LlavaForConditionalGeneration, AutoProcessor
3from transformers import BitsAndBytesConfig
4import torch
5import matplotlib.pyplot as plt
6
7
8
9# example quantization config, add it to model load parameters to use 4bit quantization
10quantization_config = BitsAndBytesConfig(
11 # load_in_8bit=True,
12 load_in_4bit=True,
13 bnb_4bit_compute_dtype=torch.bfloat16,
14 bnb_4bit_quant_type="nf4"
15 )
16
17
18
19model_id = "Ertugrul/Pixtral-12B-Captioner-Relaxed"
20model = LlavaForConditionalGeneration.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16)
21processor = AutoProcessor.from_pretrained(model_id)
22
23# for quantization just use this instead of previous load
24# model = LlavaForConditionalGeneration.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
25
26conversation = [
27 {
28 "role": "user",
29 "content": [
30
31 {"type": "text", "text": "Describe the image.\n"},
32 {
33 "type": "image",
34 }
35 ],
36 }
37]
38
39PROMPT = processor.apply_chat_template(conversation, add_generation_prompt=True)
40
41image = Image.open(r"PATH_TO_YOUR_IMAGE")
42
43def resize_image(image, target_size=768):
44 """Resize the image to have the target size on the shortest side."""
45 width, height = image.size
46 if width < height:
47 new_width = target_size
48 new_height = int(height * (new_width / width))
49 else:
50 new_height = target_size
51 new_width = int(width * (new_height / height))
52 return image.resize((new_width, new_height), Image.LANCZOS)
53
54
55# you can try different resolutions or disable it completely
56image = resize_image(image, 768)
57
58
59inputs = processor(text=PROMPT, images=image, return_tensors="pt").to("cuda")
60
61
62with torch.no_grad():
63 with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
64 generate_ids = model.generate(**inputs, max_new_tokens=384, do_sample=True, temperature=0.3, use_cache=True, top_k=20)
65output_text = processor.batch_decode(generate_ids[:, inputs.input_ids.shape[1]:], skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
66
67print(output_text)