1from transformers import AutoConfig, AutoProcessor
2from transformers.image_utils import load_image
3import onnxruntime
4import numpy as np
5from huggingface_hub import hf_hub_download
6
7# 1. Load config, processor, and model
8model_id = "onnx-community/LFM2-VL-450M-ONNX"
9config = AutoConfig.from_pretrained(model_id)
10processor = AutoProcessor.from_pretrained(model_id)
11
12local_dir = 'LFM2-VL-450M-ONNX'
13vision_model_path = hf_hub_download(model_id, "vision_encoder.onnx", subfolder="onnx", local_dir=local_dir) # Download vision graph
14hf_hub_download(model_id, "vision_encoder.onnx_data", subfolder="onnx", local_dir=local_dir) # Download vision weights
15embed_model_path = hf_hub_download(model_id, "embed_tokens.onnx", subfolder="onnx", local_dir=local_dir) # Download embed_tokens graph
16hf_hub_download(model_id, "embed_tokens.onnx_data", subfolder="onnx", local_dir=local_dir) # Download embed_tokens weights
17decoder_model_path = hf_hub_download(model_id, "decoder_model_merged.onnx", subfolder="onnx", local_dir=local_dir) # Download decoder graph
18hf_hub_download(model_id, "decoder_model_merged.onnx_data", subfolder="onnx", local_dir=local_dir) # Download decoder weights
19
20## Load sessions
21providers = ['CPUExecutionProvider']
22vision_session = onnxruntime.InferenceSession(vision_model_path, providers=providers)
23embed_session = onnxruntime.InferenceSession(embed_model_path, providers=providers)
24decoder_session = onnxruntime.InferenceSession(decoder_model_path, providers=providers)
25
26## Set config values
27text_config = config.text_config
28num_key_value_heads = text_config.num_key_value_heads
29head_dim = text_config.hidden_size // text_config.num_attention_heads
30num_hidden_layers = text_config.num_hidden_layers
31eos_token_id = text_config.eos_token_id
32hidden_size = text_config.hidden_size
33conv_L_cache = text_config.conv_L_cache
34layer_types = text_config.layer_types
35image_token_index = config.image_token_index
36
37# 2. Prepare inputs
38image_url = "https://www.ilankelman.org/stopsigns/australia.jpg"
39image = load_image(image_url)
40messages = [
41 {
42 "role": "user",
43 "content": [
44 {"type": "image", "image": image},
45 {"type": "text", "text": "What is in this image?"},
46 ],
47 },
48]
49inputs = processor.apply_chat_template(
50 messages,
51 add_generation_prompt=True,
52 return_tensors="pt",
53 return_dict=True,
54 tokenize=True,
55)
56
57input_ids = inputs['input_ids'].numpy()
58attention_mask = inputs['attention_mask'].numpy()
59has_vision_inputs = 'pixel_values' in inputs
60pixel_values = inputs['pixel_values'].numpy() if has_vision_inputs else None
61pixel_attention_mask = inputs['pixel_attention_mask'].numpy().astype(np.int64) if has_vision_inputs else None
62spatial_shapes = inputs['spatial_shapes'].numpy() if has_vision_inputs else None
63
64batch_size = input_ids.shape[0]
65past_cache_values = {}
66for i in range(num_hidden_layers):
67 if layer_types[i] == 'full_attention':
68 for kv in ('key', 'value'):
69 past_cache_values[f'past_key_values.{i}.{kv}'] = np.zeros([batch_size, num_key_value_heads, 0, head_dim], dtype=np.float32)
70 elif layer_types[i] == 'conv':
71 past_cache_values[f'past_conv.{i}'] = np.zeros([batch_size, hidden_size, conv_L_cache], dtype=np.float32)
72 else:
73 raise ValueError(f"Unsupported layer type: {layer_types[i]}")
74
75# 3. Generation loop
76max_new_tokens = 1024
77generated_tokens = np.array([[]], dtype=np.int64)
78image_features = None
79for i in range(max_new_tokens):
80 inputs_embeds = embed_session.run(None, {'input_ids': input_ids})[0]
81
82 if has_vision_inputs and image_features is None:
83 ## Only compute vision features if not already computed
84 image_features = vision_session.run(None, dict(
85 pixel_values=pixel_values,
86 pixel_attention_mask=pixel_attention_mask,
87 spatial_shapes=spatial_shapes,
88 ))[0]
89
90 ## Merge text and vision embeddings
91 inputs_embeds[input_ids == image_token_index] = image_features.reshape(-1, image_features.shape[-1])
92
93 logits, *present_cache_values = decoder_session.run(None, dict(
94 inputs_embeds=inputs_embeds,
95 attention_mask=attention_mask,
96 **past_cache_values,
97 ))
98
99 ## Update values for next generation loop
100 input_ids = logits[:, -1].argmax(-1, keepdims=True)
101 attention_mask = np.concatenate([attention_mask, np.ones((batch_size, 1), dtype=attention_mask.dtype)], axis=-1)
102 for j, key in enumerate(past_cache_values):
103 past_cache_values[key] = present_cache_values[j]
104
105 generated_tokens = np.concatenate([generated_tokens, input_ids], axis=-1)
106 if np.isin(input_ids, eos_token_id).any():
107 break
108
109 ## (Optional) Streaming
110 print(processor.decode(input_ids[0], skip_special_tokens=False), end='', flush=True)
111print()
112
113# 4. Output result
114print(processor.batch_decode(generated_tokens, skip_special_tokens=False)[0])