多模态大模型Reyes增加batch推理方式，提升推理速度 原创

def chat_batch(
            self,
            tokenizer,
            pixel_values_list,
            questions,
            generation_config,
            histories=None,
            return_histories=False,
            num_patches_lists=None,
            IMG_START_TOKEN='<|vision_start|>',
            IMG_END_TOKEN='<|vision_end|>',
            IMG_CONTEXT_TOKEN='<|vision_pad|>',
            verbose=False,
            visual_features_list=None
    ):

        if histories isNone:
            histories = [[] for _ in questions]

        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
        self.img_context_token_id = img_context_token_id
        # Get eos_token_id from the template
        template = get_conv_template(self.template)
        template.system_message = self.system_message
        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep)
        generation_config['eos_token_id'] = eos_token_id

        queries = []
        input_ids_list = []
        attention_mask_list = []

        for idx in range(len(questions)):
            question = questions[idx]
            history = histories[idx]
            pixel_values = pixel_values_list[idx] if pixel_values_list[idx] isnotNoneelseNone
            num_patches_list = [pixel_values.shape[0]] if pixel_values isnotNoneelse []

            ifnot history and pixel_values isnotNoneand'<image>'notin question:
                question = '<image>\n' + question

            template_i = get_conv_template(self.template)
            template_i.system_message = self.system_message
            for (old_question, old_answer) in history:
                template_i.append_message(template_i.roles[0], old_question)
                template_i.append_message(template_i.roles[1], old_answer)
            template_i.append_message(template_i.roles[0], question)
            template_i.append_message(template_i.roles[1], None)
            query = template_i.get_prompt()
            # Handle image tokens
            if pixel_values isnotNone:
                for num_patches in num_patches_list:
                    tile_pos_identifiers = [f"<tile_{i}>"for i in range(1, num_patches)] + ["<tile_global_thumbnail>"]
                    image_tokens = ''
                    for tile_pos_identifier in tile_pos_identifiers:
                        image_tokens += tile_pos_identifier + IMG_CONTEXT_TOKEN * self.num_image_token
                    image_tokens = IMG_START_TOKEN + image_tokens + IMG_END_TOKEN
                    query = query.replace('<image>', image_tokens, 1)

            model_inputs = tokenizer(
                query,
                return_tensors='pt',
                padding=True,
                truncation=True
            )
            input_ids = model_inputs['input_ids'].cuda()
            attention_mask = model_inputs['attention_mask'].cuda()
            input_ids_list.append(input_ids)
            attention_mask_list.append(attention_mask)

        # Call the generate function
        generation_output = self.generate_batch(
            pixel_values_list=pixel_values_list,
            input_ids_list=input_ids_list,
            attention_mask_list=attention_mask_list,
            **generation_config
        )
        responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)

        outputs = []
        for idx, response in enumerate(responses):
            response = response.split(template.sep)[0].strip()
            histories[idx].append((questions[idx], response))
            outputs.append(response)

        if return_histories:
            return outputs, histories
        else:
            if verbose:
                for idx, query in enumerate(queries):
                    query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')
                    query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')
                    print(query_to_print, outputs[idx])
            return outputs

    @torch.no_grad()
    def generate_batch(
            self,
            pixel_values_list: Optional[List[torch.FloatTensor]] = None,
            input_ids_list: Optional[List[torch.FloatTensor]] = None,
            attention_mask_list: Optional[List[torch.LongTensor]] = None,
            visual_features: Optional[torch.FloatTensor] = None,
            generation_config: Optional[GenerationConfig] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
            **generate_kwargs,
    ) -> torch.LongTensor:
        input_embeds_list = []
        attention_mask_padded_list = []

        max_seq_length = max(input_ids.shape[1] for input_ids in input_ids_list)

        for pixel_values, input_ids, attention_mask in zip(pixel_values_list, input_ids_list, attention_mask_list):
            if pixel_values isnotNone:
                if visual_features isnotNone:
                    vit_embeds = visual_features.cuda()
                    vit_embeds = self.mlp1(vit_embeds)
                else:
                    vit_embeds = self.extract_feature(pixel_values)

                input_embeds = self.language_model.get_input_embeddings()(input_ids)
                B, N, C = input_embeds.shape
                input_embeds = input_embeds.reshape(B * N, C)

                input_ids = input_ids.reshape(B * N)
                selected = (input_ids == self.img_context_token_id)
                assert selected.sum() != 0, "No valid image context token IDs found."
                input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)

                input_embeds = input_embeds.reshape(B, N, C)
            else:
                input_embeds = self.language_model.get_input_embeddings()(input_ids)

            seq_length = input_embeds.shape[1]
            if seq_length < max_seq_length:
                pad_size = max_seq_length - seq_length
                input_embeds = F.pad(input_embeds, (0, 0, 0, pad_size))
                attention_mask = F.pad(attention_mask, (0, pad_size))

            input_embeds_list.append(input_embeds)
            attention_mask_padded_list.append(attention_mask)

        input_embeds = torch.cat(input_embeds_list, dim=0)
        attention_mask = torch.cat(attention_mask_padded_list, dim=0)

        outputs = self.language_model.generate(
            inputs_embeds=input_embeds,
            attention_mask=attention_mask,
            generation_config=generation_config,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            use_cache=True,
            **generate_kwargs,
        )

        return outputs
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.
72.
73.
74.
75.
76.
77.
78.
79.
80.
81.
82.
83.
84.
85.
86.
87.
88.
89.
90.
91.
92.
93.
94.
95.
96.
97.
98.
99.
100.
101.
102.
103.
104.
105.
106.
107.
108.
109.
110.
111.
112.
113.
114.
115.
116.
117.
118.
119.
120.
121.
122.
123.
124.
125.
126.
127.
128.
129.
130.
131.
132.
133.
134.
135.
136.
137.
138.
139.
140.
141.
142.
143.
144.
145.
146.
147.
148.
149.
150.
151.
152.
153.
154.
155.

import torch
from modelscope import AutoTokenizer, AutoModel
from PIL import Image
import torchvision.transforms as T
from torchvision.transforms.functional import InterpolationMode

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def build_transform(input_size):
    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    transform = T.Compose([
        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB'else img),
        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(mean=MEAN, std=STD)
    ])
    return transform


def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    best_ratio_diff = float('inf')
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        if ratio_diff < best_ratio_diff:
            best_ratio_diff = ratio_diff
            best_ratio = ratio
        elif ratio_diff == best_ratio_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best_ratio = ratio
    return best_ratio


def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height

    # calculate the existing image aspect ratio
    target_ratios = set(
        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
        i * j <= max_num and i * j >= min_num)
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

    # find the closest aspect ratio to the target
    target_aspect_ratio = find_closest_aspect_ratio(
        aspect_ratio, target_ratios, orig_width, orig_height, image_size)

    # calculate the target width and height
    target_width = image_size * target_aspect_ratio[0]
    target_height = image_size * target_aspect_ratio[1]
    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

    # resize the image
    resized_img = image.resize((target_width, target_height))
    processed_images = []
    for i in range(blocks):
        box = (
            (i % (target_width // image_size)) * image_size,
            (i // (target_width // image_size)) * image_size,
            ((i % (target_width // image_size)) + 1) * image_size,
            ((i // (target_width // image_size)) + 1) * image_size
        )
        # split the image
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images


def load_image(image_file, input_size=448, max_num=12):
    image = Image.open(image_file).convert('RGB')
    transform = build_transform(input_size=input_size)
    images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
    pixel_values = [transform(image) for image in images]
    pixel_values = torch.stack(pixel_values)
    return pixel_values


def preprocess_image(file_path, dynamic=True, max_num=6, image_size=448):
    try:
        if dynamic:
            return load_image(file_path, max_num=max_num).to(torch.bfloat16).cuda()
        else:
            img = Image.open(file_path).convert('RGB')
            transform = build_transform(image_size)
            pixel_values = transform(img)
            return torch.stack([pixel_values]).to(torch.bfloat16).cuda()
    except Exception as e:
        raise RuntimeError(f"Error processing image: {e}")


path = "Reyes-8B"

model = AutoModel.from_pretrained(
    path,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
).eval().cuda()

# print(model)

tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
generation_config = dict(max_new_tokens=2048, do_sample=False)
questions = [
    "<image>\nDescribe this image.",
    "<image>\nDescribe this image.",
    "<image>\nDescribe this image.",
]

images_path = ["t6.png","t6.png","t6.png"]


def conversation(model, tokenizer, questions, images_path,generation_config,histories):
    pixel_values_list=[]

    for i in range(len(questions)):
        if images_path[i] isnotNone:
            pixel_values = preprocess_image(file_path, dynamic=True)
            pixel_values_list.append(pixel_values)


    return model.chat_batch(tokenizer, pixel_values_list, questions, generation_config, histories, return_histories=False)

responses= conversation(model, tokenizer, questions, images_path,generation_config,histories=None)
for question, response in zip(questions, responses):
    print(f"User: {question}\nAssistant: {response}\n")
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.
72.
73.
74.
75.
76.
77.
78.
79.
80.
81.
82.
83.
84.
85.
86.
87.
88.
89.
90.
91.
92.
93.
94.
95.
96.
97.
98.
99.
100.
101.
102.
103.
104.
105.
106.
107.
108.
109.
110.
111.
112.
113.
114.
115.
116.
117.
118.
119.
120.
121.
122.
123.
124.
125.
126.
127.
128.
129.
130.
131.
132.
133.