Update README.md

README.md

@@ -16,19 +16,30 @@ See https://huggingface.co/spaces/pdufour/Qwen2-VL-2B-Instruct-ONNX-Q4-F16 for a
 **Python**
 
 ```
+import os
+import sys
 import time
 import torch
 import numpy as np
-import onnxruntime
-from PIL import Image
-import os
-import sys
 import requests
+import onnxruntime as ort
+from PIL import Image
 from io import BytesIO
+from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer
 
+# Constants
+DEBUG = True
+PRINT = print
 
+# Try importing config, set defaults if not found
 try:
-    from export_config import INPUT_IMAGE_SIZE, IMAGE_RESIZE, MAX_SEQ_LENGTH, HEIGHT_FACTOR, WIDTH_FACTOR
+    from export_config import (
+        INPUT_IMAGE_SIZE,
+        IMAGE_RESIZE,
+        MAX_SEQ_LENGTH,
+        HEIGHT_FACTOR,
+        WIDTH_FACTOR
+    )
 except:
     INPUT_IMAGE_SIZE = [960, 960]
     HEIGHT_FACTOR = 10
@@ -36,220 +47,186 @@ except:
     IMAGE_RESIZE = [HEIGHT_FACTOR * 28, WIDTH_FACTOR * 28]
     MAX_SEQ_LENGTH = 1024
 
-path = sys.argv[1]
-script_dir = sys.argv[2]
+# Command line arguments
+model_path = sys.argv[1]
+onnx_path = sys.argv[2]
 
-onnx_model_A = os.path.join(script_dir, 'QwenVL_A.onnx')
-onnx_model_B = os.path.join(script_dir, 'QwenVL_B_q4f16.onnx')
-onnx_model_C = os.path.join(script_dir, 'QwenVL_C_q4f16.onnx')
-onnx_model_D = os.path.join(script_dir, 'QwenVL_D_q4f16.onnx')
-onnx_model_E = os.path.join(script_dir, 'QwenVL_E_q4f16.onnx')
+# ONNX model paths
+model_paths = {
+    'A': os.path.join(onnx_path, 'QwenVL_A_q4f16.onnx'),
+    'B': os.path.join(onnx_path, 'QwenVL_B_q4f16.onnx'),
+    'C': os.path.join(onnx_path, 'QwenVL_C_q4f16.onnx'),
+    'D': os.path.join(onnx_path, 'QwenVL_D_q4f16.onnx'),
+    'E': os.path.join(onnx_path, 'QwenVL_E_q4f16.onnx')
+}
 
-print("\n[PATHS] ONNX model paths:")
-print(f"  Model A: {onnx_model_A}")
-print(f"  Model B: {onnx_model_B}")
-print(f"  Model C: {onnx_model_C}")
-print(f"  Model D: {onnx_model_D}")
-print(f"  Model E: {onnx_model_E}")
+PRINT('\n[PATHS] ONNX model paths:')
+for key, path in model_paths.items():
+    PRINT(f"  Model {key}: {path}")
 
-image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg"
-query = "Describe this image."
-
-from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer
+# Test image and prompt
+TEST_IMAGE_URL = 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg'
+TEST_PROMPT = 'Describe this image.'
 
+# Initialize model and tokenizer
 with torch.inference_mode():
-    model = Qwen2VLForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float32, device_map="mps", low_cpu_mem_usage=True)
-    max_seq_len = MAX_SEQ_LENGTH
-    num_heads = model.config.num_attention_heads
+    model = Qwen2VLForConditionalGeneration.from_pretrained(
+        model_path,
+        torch_dtype=torch.float32,
+        device_map='mps',
+        low_cpu_mem_usage=DEBUG
+    )
+
+    max_length = MAX_SEQ_LENGTH
+    num_attention_heads = model.config.num_attention_heads
     num_key_value_heads = model.config.num_key_value_heads
-    head_dim = model.config.hidden_size // num_heads
+    head_dim = model.config.hidden_size // num_attention_heads
     num_layers = model.config.num_hidden_layers
     hidden_size = model.config.hidden_size
 
-
-max_single_chat_length = 12
-
-tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
-
-session_opts = onnxruntime.SessionOptions()
-session_opts.log_severity_level = 3
-session_opts.inter_op_num_threads = 0
-session_opts.intra_op_num_threads = 0
-session_opts.enable_cpu_mem_arena = True
-session_opts.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
-session_opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
-session_opts.add_session_config_entry("session.intra_op.allow_spinning", "1")
-session_opts.add_session_config_entry("session.inter_op.allow_spinning", "1")
-
-ort_session_A = onnxruntime.InferenceSession(onnx_model_A, sess_options=session_opts)
-ort_session_B = onnxruntime.InferenceSession(onnx_model_B, sess_options=session_opts)
-ort_session_C = onnxruntime.InferenceSession(onnx_model_C, sess_options=session_opts)
-ort_session_D = onnxruntime.InferenceSession(onnx_model_D, sess_options=session_opts)
-ort_session_E = onnxruntime.InferenceSession(onnx_model_E, sess_options=session_opts)
-
-in_name_A = ort_session_A.get_inputs()
-out_name_A = ort_session_A.get_outputs()
-in_name_A0 = in_name_A[0].name
-out_name_A0 = out_name_A[0].name
-
-in_name_B = ort_session_B.get_inputs()
-out_name_B = ort_session_B.get_outputs()
-in_name_B0 = in_name_B[0].name
-in_name_B1 = in_name_B[1].name
-out_name_B0 = out_name_B[0].name
-
-in_name_C = ort_session_C.get_inputs()
-out_name_C = ort_session_C.get_outputs()
-in_name_C0 = in_name_C[0].name
-out_name_C0 = out_name_C[0].name
-
-in_name_D = ort_session_D.get_inputs()
-out_name_D = ort_session_D.get_outputs()
-in_name_D0 = in_name_D[0].name
-in_name_D1 = in_name_D[1].name
-in_name_D2 = in_name_D[2].name
-in_name_D3 = in_name_D[3].name
-in_name_D4 = in_name_D[4].name
-out_name_D0 = out_name_D[0].name
-out_name_D1 = out_name_D[1].name
-
-in_name_E = ort_session_E.get_inputs()
-out_name_E = ort_session_E.get_outputs()
-in_name_E0 = in_name_E[0].name
-in_name_E1 = in_name_E[1].name
-in_name_E2 = in_name_E[2].name
-in_name_E3 = in_name_E[3].name
-in_name_E4 = in_name_E[4].name
-in_name_E5 = in_name_E[5].name
-in_name_E6 = in_name_E[6].name
-in_name_E7 = in_name_E[7].name
-out_name_E0 = out_name_E[0].name
-out_name_E1 = out_name_E[1].name
-out_name_E2 = out_name_E[2].name
-
-response = requests.get(image_url)
+MAX_ITERATIONS = 12
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=DEBUG)
+
+# ONNX session options
+session_options = ort.SessionOptions()
+session_options.log_severity_level = 3
+session_options.inter_op_num_threads = 0
+session_options.intra_op_num_threads = 0
+session_options.enable_cpu_mem_arena = DEBUG
+session_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+session_options.add_session_config_entry('session.intra_op.allow_spinning', '1')
+session_options.add_session_config_entry('session.inter_op.allow_spinning', '1')
+
+# Initialize ONNX sessions
+sessions = {
+    'A': ort.InferenceSession(model_paths['A'], sess_options=session_options),
+    'B': ort.InferenceSession(model_paths['B'], sess_options=session_options),
+    'C': ort.InferenceSession(model_paths['C'], sess_options=session_options),
+    'D': ort.InferenceSession(model_paths['D'], sess_options=session_options),
+    'E': ort.InferenceSession(model_paths['E'], sess_options=session_options)
+}
+
+# Get input/output names for each session
+inputs = {
+    'A': {
+        'input': sessions['A'].get_inputs()[0].name,
+        'output': sessions['A'].get_outputs()[0].name
+    },
+    'B': {
+        'input_ids': sessions['B'].get_inputs()[0].name,
+        'input_lengths': sessions['B'].get_inputs()[1].name,
+        'output': sessions['B'].get_outputs()[0].name
+    },
+    'C': {
+        'input': sessions['C'].get_inputs()[0].name,
+        'output': sessions['C'].get_outputs()[0].name
+    },
+    'D': {
+        'names': [inp.name for inp in sessions['D'].get_inputs()],
+        'outputs': [out.name for out in sessions['D'].get_outputs()]
+    },
+    'E': {
+        'names': [inp.name for inp in sessions['E'].get_inputs()],
+        'outputs': [out.name for out in sessions['E'].get_outputs()]
+    }
+}
+
+# Process image
+response = requests.get(TEST_IMAGE_URL)
 image = Image.open(BytesIO(response.content))
-
+image = image.resize((INPUT_IMAGE_SIZE[1], INPUT_IMAGE_SIZE[0]))
 if image.mode != 'RGB':
     image = image.convert('RGB')
 
-pixel_values = np.transpose(np.array(image).astype(np.float32), (2, 0, 1))
-pixel_values = np.expand_dims(pixel_values, axis=0) / 255.0
-use_vision = True
-
-prompt = f"\n<|im_start|>user\n<|vision_start|><|vision_end|>{query}<|im_end|>\n<|im_start|>assistant\n"
-prompt_head_len = np.array([5], dtype=np.int64)
-
-image_embed_size = WIDTH_FACTOR * HEIGHT_FACTOR
-
-token = tokenizer(prompt, return_tensors='pt')['input_ids']
-
-ids_len = np.array([token.shape[1]], dtype=np.int64)
-
-input_ids = np.zeros(max_seq_len, dtype=np.int32)
-input_ids[:ids_len[0]] = token[0, :]
-
-history_len = np.zeros(1, dtype=np.int64)
-
-past_key_states = np.zeros((num_layers, num_key_value_heads, max_seq_len, head_dim), dtype=np.float16)
-
-past_values_states = past_key_states
-
-attention_mask = np.array([-65504.0], dtype=np.float16)
-
-pos_factor = np.array([0.0], dtype=np.float16)
-
-pos_factor_v = 1 - image_embed_size + WIDTH_FACTOR
-
-dummy = np.array(0, dtype=np.int32)
-
-hidden_states = ort_session_B.run(
-    [out_name_B0],
-    {
-        in_name_B0: input_ids,
-        in_name_B1: ids_len
-    })[0]
-
-position_ids, = ort_session_C.run(
-    [out_name_C0],
-    {
-        in_name_C0: dummy
-    })
-
-if use_vision:
-
-    image_embed = ort_session_A.run(
-        [out_name_A0],
-        {in_name_A0: pixel_values})[0]
-
-    ids_len += image_embed_size
-
-    split_factor = np.array(max_seq_len - ids_len[0] - image_embed_size, dtype=np.int32)
-
-    ids_len_minus = np.array(ids_len[0] - prompt_head_len[0], dtype=np.int32)
-
-
-    hidden_states, position_ids = ort_session_D.run(
-        [out_name_D0, out_name_D1],
+image_array = np.transpose(np.array(image).astype(np.float32), (2, 0, 1))
+image_array = np.expand_dims(image_array, axis=0) / 255.
+
+use_images = DEBUG
+prompt = f"\n<|im_start|>user\n<|vision_start|><|vision_end|>{TEST_PROMPT}<|im_end|>\n<|im_start|>assistant\n"
+eos_token_id = np.array([5], dtype=np.int64)
+total_ids = WIDTH_FACTOR * HEIGHT_FACTOR
+
+# Initialize tensors
+input_ids = tokenizer(prompt, return_tensors='pt')['input_ids']
+input_lengths = np.array([input_ids.shape[1]], dtype=np.int64)
+tokens = np.zeros(max_length, dtype=np.int32)
+tokens[:input_lengths[0]] = input_ids[0, :]
+position = np.zeros(1, dtype=np.int64)
+
+# Initialize cache tensors
+key_cache = np.zeros((num_layers, num_key_value_heads, max_length, head_dim), dtype=np.float16)
+value_cache = key_cache.copy()
+logits_mask = np.array([-65504.], dtype=np.float16)
+position_mask = np.array([.0], dtype=np.float16)
+max_total_tokens = 1 - total_ids + WIDTH_FACTOR
+batch_size = np.array(0, dtype=np.int32)
+
+# Process initial inputs
+hidden_states = sessions['B'].run([inputs['B']['output']],
+    {inputs['B']['input_ids']: tokens, inputs['B']['input_lengths']: input_lengths})[0]
+batch_size, = sessions['C'].run([inputs['C']['output']], {inputs['C']['input']: batch_size})
+
+if use_images:
+    image_features = sessions['A'].run([inputs['A']['output']], {inputs['A']['input']: image_array})[0]
+    input_lengths += total_ids
+    remaining_tokens = np.array(max_length - input_lengths[0] - total_ids, dtype=np.int32)
+    tokens_to_stop = np.array(input_lengths[0] - eos_token_id[0], dtype=np.int32)
+    hidden_states, batch_size = sessions['D'].run(
+        [inputs['D']['outputs'][0], inputs['D']['outputs'][1]],
         {
-            in_name_D0: hidden_states,
-            in_name_D1: image_embed,
-            in_name_D2: ids_len,
-            in_name_D3: ids_len_minus,
-            in_name_D4: split_factor
-        })
-
-    end_time = time.time()
-
-end_time = time.time()
-num_decode = 0
-
-while (num_decode < max_single_chat_length) & (history_len < max_seq_len):
-    token_id, past_key_states, past_values_states = ort_session_E.run(
-        [out_name_E0, out_name_E1, out_name_E2],
+            inputs['D']['names'][0]: hidden_states,
+            inputs['D']['names'][1]: image_features,
+            inputs['D']['names'][2]: input_lengths,
+            inputs['D']['names'][3]: tokens_to_stop,
+            inputs['D']['names'][4]: remaining_tokens
+        }
+    )
+
+start_time = time.time()
+iterations = 0
+
+while (iterations < MAX_ITERATIONS) & (position < max_length):
+    token, key_cache, value_cache = sessions['E'].run(
+        [inputs['E']['outputs'][0], inputs['E']['outputs'][1], inputs['E']['outputs'][2]],
         {
-            in_name_E0: hidden_states,
-            in_name_E1: attention_mask,
-            in_name_E2: past_key_states,
-            in_name_E3: past_values_states,
-            in_name_E4: history_len,
-            in_name_E5: ids_len,
-            in_name_E6: position_ids,
-            in_name_E7: pos_factor
-        })
-
-    if (token_id == 151643) | (token_id == 151645):
+            inputs['E']['names'][0]: hidden_states,
+            inputs['E']['names'][1]: logits_mask,
+            inputs['E']['names'][2]: key_cache,
+            inputs['E']['names'][3]: value_cache,
+            inputs['E']['names'][4]: position,
+            inputs['E']['names'][5]: input_lengths,
+            inputs['E']['names'][6]: batch_size,
+            inputs['E']['names'][7]: position_mask
+        }
+    )
+
+    if (token == 151643) | (token == 151645):
         break
     else:
-        num_decode += 1
-        if num_decode < 2:
-            history_len += ids_len[0]
-
-            ids_len[0] = 1
-
-            attention_mask = np.array([0.0], dtype=np.float16)
-
-            if use_vision:
-                pos_factor = np.array(pos_factor_v + ids_len[0], dtype=np.float16)
+        iterations += 1
+        if iterations < 2:
+            position += input_lengths[0]
+            input_lengths[0] = 1
+            logits_mask = np.array([.0], dtype=np.float16)
+            if use_images:
+                position_mask = np.array(max_total_tokens + input_lengths[0], dtype=np.float16)
             else:
-                pos_factor = np.array(history_len[0] + 1, dtype=np.float16)
+                position_mask = np.array(position[0] + 1, dtype=np.float16)
         else:
-            history_len += 1
-            pos_factor += 1
-
-        input_ids[0] = token_id
-        hidden_states = ort_session_B.run(
-            [out_name_B0],
-            {
-                in_name_B0: input_ids,
-                in_name_B1: ids_len
-            })[0]
-
-        decoded_token = tokenizer.decode(token_id)
-        print(f"Decoded token: {decoded_token}")
-        print(decoded_token, end="", flush=True)
-
-generation_time = time.time() - end_time
+            position += 1
+            position_mask += 1
+
+        tokens[0] = token
+        hidden_states = sessions['B'].run(
+            [inputs['B']['output']],
+            {inputs['B']['input_ids']: tokens, inputs['B']['input_lengths']: input_lengths}
+        )[0]
+        decoded_token = tokenizer.decode(token)
+        PRINT(f"Decoded token: {decoded_token}")
+        PRINT(decoded_token, end='', flush=DEBUG)
+
+total_time = time.time() - start_time
 ```
 
 # Technical Information: