infer_qwen2_vl.py

from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info

def read_json(file_path): 
    with open(file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    return data

def write_json(file_path, data):
    with open(file_path, 'w', encoding='utf-8') as file:
        json.dump(data, file, ensure_ascii=False, indent=4)

# default: Load the model on the available device(s)
model_path = '/inspire/hdd/ws-ba572160-47f8-4ca1-984e-d6bcdeb95dbb/a100-maybe/albus/ICCV_2025/qvq/models/QVQ-72B-Preview'
model = Qwen2VLForConditionalGeneration.from_pretrained(
    model_path, torch_dtype="auto", device_map="auto"
)

# default processer
processor = AutoProcessor.from_pretrained(model_path)

# The default range for the number of visual tokens per image in the model is 4-16384. You can set min_pixels and max_pixels according to your needs, such as a token count range of 256-1280, to balance speed and memory usage.
# min_pixels = 256*28*28
# max_pixels = 1280*28*28
#processor = AutoProcessor.from_pretrained("Qwen/QVQ-72B-Preview", min_pixels=min_pixels, max_pixels=max_pixels)

import glob
from PIL import Image
import argparse
import os

# parser = argparse.ArgumentParser(description="Process a dataset with specific index range.")
# parser.add_argument("--batch_size", type=int, default = 1,help="batch size")
# #parser.add_argument("--index", type=int, default = 0,help="index")
# args = parser.parse_args()


folder = "/inspire/hdd/ws-ba572160-47f8-4ca1-984e-d6bcdeb95dbb/a100-maybe/albus/ICCV_2025/qvq/dataset"

file_names = os.listdir(folder)

num_image = len(file_names)

begin, end, batch_size= 0, 10, 2
print(f"beigin : {begin}, end : {end}, batch_size : {batch_size}")
messages = [
    {
        "role": "system",
        "content": [
            {"type": "text", "text": "You are a helpful and harmless assistant. You are Qwen developed by Alibaba. You should think step-by-step."}
        ],
    },
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/QVQ/demo.png",
            },
            {"type": "text", "text": "Please describe in detail the content of the picture."},
        ],
    }
]

from tqdm import tqdm
# Preparation for inference
ans = []
counter = 0
for batch_idx in tqdm(range(begin, end, batch_size)):
    batch = file_names[batch_idx: min(batch_idx + batch_size, end)]
    print(f"data index range : {batch_idx} ~ {min(batch_idx + batch_size, end)}")
    image_inputs_batch, video_inputs_batch = [], []
    for idx,i in enumerate(batch):
        #img = batch[i]
        print('gain image successfully !')
        messages[1]["content"][0]["image"] = i
        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        image_inputs, video_inputs = process_vision_info(messages)
        image_inputs_batch.append(image_inputs)
        video_inputs_batch.append(video_inputs)
    inputs = processor(
        text=[text],
        images=image_inputs_batch,
        videos=video_inputs_batch,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to("cuda")

    # Inference: Generation of the output
    generated_ids = model.generate(**inputs, max_new_tokens=8192)
    generated_ids_trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    #ans.append(output_text)
    save_path = "output.json"
    counter = counter + 1
    if counter % 1 == 0:
        print(f"Saving data at iteration {idx + 1}")
        write_json(save_path, data)