1from pathlib import Path
2
3from qwen_vl_utils import process_vision_info
4from transformers import AutoModelForImageTextToText, AutoProcessor
5
6model_id = "MCG-NJU/TimeLens2-8B"
7video_path = "/path/to/video.mp4"
8query = "A man opens the refrigerator."
9
10model = AutoModelForImageTextToText.from_pretrained(
11 model_id,
12 torch_dtype="auto",
13 device_map="auto",
14 attn_implementation="flash_attention_2",
15)
16processor = AutoProcessor.from_pretrained(model_id)
17
18prompt = (
19 f'Given the query: "{query}", return ALL time spans (in seconds) where the query is relevant.\n'
20 "Output format MUST be a JSON array of [start, end] pairs.\n"
21)
22messages = [
23 {
24 "role": "user",
25 "content": [
26 {
27 "type": "video",
28 "video": Path(video_path).resolve().as_uri(),
29 "fps": 2.0,
30 "min_pixels": 32 * 32,
31 "max_pixels": 480 * 480,
32 "total_pixels": 128000 * 32 * 32,
33 },
34 {"type": "text", "text": prompt},
35 ],
36 }
37]
38
39text = processor.apply_chat_template(
40 messages,
41 tokenize=False,
42 add_generation_prompt=True,
43)
44images, videos, video_kwargs = process_vision_info(
45 messages,
46 image_patch_size=16,
47 return_video_kwargs=True,
48 return_video_metadata=True,
49)
50
51if videos is not None:
52 videos, video_metadatas = zip(*videos)
53 videos, video_metadatas = list(videos), list(video_metadatas)
54else:
55 video_metadatas = None
56
57inputs = processor(
58 text=text,
59 images=images,
60 videos=videos,
61 video_metadata=video_metadatas,
62 do_resize=False,
63 return_tensors="pt",
64 **video_kwargs,
65).to(model.device)
66
67output_ids = model.generate(
68 **inputs,
69 max_new_tokens=4096,
70 temperature=0.01,
71 top_p=0.001,
72 top_k=1,
73 repetition_penalty=1.0,
74)
75output_ids = [
76 output[len(input_ids) :]
77 for input_ids, output in zip(inputs.input_ids, output_ids)
78]
79response = processor.batch_decode(
80 output_ids,
81 skip_special_tokens=True,
82 clean_up_tokenization_spaces=False,
83)
84print(response[0])