Update models and scripts with toolchain version 2.3.2

onnx/convert.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # coding: utf-8
 
+import numpy as np
 from rknn.api import RKNN
 from math import exp
 from sys import exit
@@ -67,7 +68,64 @@ def convert_decoder():
                             [batch_size, decoder_seq_len, 768]] for encoder_seq_len in encoder_seq_len_list]
     # pre-process config
     print('--> Config model')
-    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True,
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3,
+                dynamic_input=input_shapes)
+    print('done')
+
+    # Load ONNX model
+    print('--> Loading model')
+    ret = rknn.load_onnx(model=ONNX_MODEL,
+                         )
+    if ret != 0:
+        print('Load model failed!')
+        exit(ret)
+    print('done')
+
+    # Build model
+    print('--> Building model')
+    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
+    if ret != 0:
+        print('Build model failed!')
+        exit(ret)
+    print('done')
+
+    #export
+    print('--> Export RKNN model')
+    ret = rknn.export_rknn(RKNN_MODEL)
+    if ret != 0:
+        print('Export RKNN model failed!')
+        exit(ret)
+    print('done')
+
+def convert_decoder_2():
+    import onnx_graphsurgeon as gs
+    ONNX_MODEL="decoder_model_merged.onnx"
+
+    graph = gs.import_onnx(onnx.load(ONNX_MODEL))
+    inp = graph.inputs[27]  # use_cache_branch
+    inp.to_constant(np.array([True], dtype=np.bool_))
+    ONNX_MODEL
+    onnx.save(gs.export_onnx(graph), "new_model.onnx")
+
+    np_true = np.array([True], dtype=np.bool_)
+    np.save("np_true.npy", np_true)
+
+
+    rknn = RKNN(verbose=True)
+
+    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
+    DATASET="dataset.txt"
+    QUANTIZE=False
+
+    # [[batch_size, encoder_seq_len], 
+    # [batch_size, encoder_seq_len, 768], 
+    # [batch_size, decoder_seq_len, 768]]
+    input_shapes =[[[batch_size, encoder_seq_len],
+                            [batch_size, encoder_seq_len, 768],
+                            [batch_size, decoder_seq_len, 768]] for encoder_seq_len in encoder_seq_len_list]
+    # pre-process config
+    print('--> Config model')
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3,
                 dynamic_input=input_shapes)
     print('done')
 
@@ -108,7 +166,7 @@ def convert_encoder():
     input_shapes = [[[batch_size, encoder_seq_len], [batch_size, encoder_seq_len, 768]] for encoder_seq_len in encoder_seq_len_list]
     # pre-process config
     print('--> Config model')
-    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True, dynamic_input=input_shapes)
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, dynamic_input=input_shapes)
     print('done')
 
     # Load ONNX model
@@ -137,49 +195,43 @@ def convert_encoder():
     print('done')
 
 def convert_vision():
+    ONNX_MODEL="vision_encoder.onnx"
+    DATASET="dataset.txt"
+    QUANTIZE=False
+    global batch_size
+    
+    ##### Build stage 1
+    from rknn.api import RKNN
     rknn = RKNN(verbose=True)
-
     ONNX_MODEL="vision_encoder.onnx"
+    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
     DATASET="dataset.txt"
     QUANTIZE=False
+    # pre-process config
+    print('--> Config model')
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3)
+    print('done')
+
+    # Load ONNX model
+    print('--> Loading model')
+    ret = rknn.load_onnx(model=ONNX_MODEL,
+                            inputs=["pixel_values"],
+                            input_size_list=[[batch_size, 3, 768, 768]],
+                            )
+    if ret != 0:
+        print('Load model failed!')
+        exit(ret)
+    print('done')
+
+    print('--> Building model stage 1')
+    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
+    if ret != 0:
+        print('Build model failed!')
+        exit(ret)
+    print('done')
 
-    # split the first Transformers block into a separate model because it's too large to fit in the rknn
-    onnx.utils.extract_model(ONNX_MODEL, "vision_encoder_part1.onnx", ['pixel_values'], ['/blocks.0/blocks.0.0/channel_block/channel_attn/Add_output_0'])
-    
-    ##### Build stage 1, this will crash the python process, so we need to run it in a separate process
-    code = f"""
-from rknn.api import RKNN
-rknn = RKNN(verbose=True)
-ONNX_MODEL="vision_encoder.onnx"
-RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
-DATASET="dataset.txt"
-QUANTIZE=False
-batch_size = {batch_size}
-# pre-process config
-print('--> Config model')
-rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True)
-print('done')
-
-# Load ONNX model
-print('--> Loading model')
-ret = rknn.load_onnx(model=ONNX_MODEL,
-                        inputs=["pixel_values"],
-                        input_size_list=[[batch_size, 3, 768, 768]],
-                        )
-if ret != 0:
-    print('Load model failed!')
-    exit(ret)
-print('done')
-
-print('--> Building model stage 1')
-ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
-if ret != 0:
-    print('Build model failed!')
-    exit(ret)
-print('done')
-    """
-    run_python_code(code)
     print("Build stage 1 done")
+    del rknn
 
     intermidiate_model = onnx.load("check3_fuse_ops.onnx")
 
@@ -210,9 +262,9 @@ print('done')
         intermidiate_model,
         pattern_rewrite_rules=rewrite_rule_set
     )
-    onnx.save(fused_model, "vision_encoder_part2.onnx")
-    ONNX_MODEL = "vision_encoder_part2.onnx"
-    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
+    onnx.save(fused_model, "vision_encoder_optimized.onnx")
+    ONNX_MODEL = "vision_encoder_optimized.onnx"
+    # RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
     del intermidiate_model
     del fused_model
 
@@ -221,14 +273,12 @@ print('done')
 
     # pre-process config
     print('--> Config model')
-    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True)
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3)
     print('done')
 
     # Load ONNX model
     print('--> Loading model')
-    ret = rknn.load_onnx(model="check3_fuse_ops.onnx",
-                         inputs=["/blocks.0/blocks.0.0/channel_block/channel_attn/Add_output_0-rs"],
-                         input_size_list=[[batch_size, 128, 1, 36864]],)
+    ret = rknn.load_onnx(model=ONNX_MODEL)
     if ret != 0:
         print('Load model failed!')
         exit(ret)
@@ -249,10 +299,7 @@ print('done')
         print('Export RKNN model failed!')
         exit(ret)
     print('done')
-    
-    
-    
-    
+    os.remove("vision_encoder_optimized.onnx")
 
 
 
@@ -266,7 +313,7 @@ def check_vision_model():
 
     # pre-process config
     print('--> Config model')
-    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True )
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3)
     print('done')
 
     # Load ONNX model
@@ -311,9 +358,6 @@ def check_vision_model():
         print('Precision check failed!')
         exit(ret)
     print('done')
-    
-    
-    
 
 
 import argparse

onnx/decoder_model.rknn

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:331a6a05a524c72ac7287a494d6cadd425266888be5ff9375649c8760417f611
-size 194821309
+oid sha256:9ccb57a522ab8b0fa73123d654807748fbaf841c6852c775eb293e054b520341
+size 207755060

onnx/encoder_model.rknn

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a36af46c308219399dfe5f1df53c2093c3247c9dc248dc3c3167ab88975cf62c
-size 87231735
+oid sha256:3121d4ff0f5fc79420e6eda1d657eb8ff36355a414fcab3f236c72b2f4e9ddd1
+size 106957934

onnx/rknnrun.py

@@ -20,7 +20,7 @@ rknn_encoder = RKNNLite(verbose=False)
 rknn_decoder_prefill = RKNNLite(verbose=False)
 
 # Load RKNN models
-ret = rknn_vision_encoder.load_rknn('./vision_encoder_part2.rknn')
+ret = rknn_vision_encoder.load_rknn('./vision_encoder.rknn')
 ret = rknn_encoder.load_rknn('./encoder_model.rknn')
 ret = rknn_decoder_prefill.load_rknn('./decoder_model.rknn')
 
@@ -31,18 +31,18 @@ ret = rknn_decoder_prefill.init_runtime()
 
 text_embed = ort.InferenceSession("embed_tokens_fp16.onnx", providers=['CPUExecutionProvider'])
 decoder_decode = ort.InferenceSession("decoder_model_merged_q4.onnx", providers=['CPUExecutionProvider'])
-vision_encoder = ort.InferenceSession("vision_encoder_part1.onnx", providers=['CPUExecutionProvider'])
+
 prompt_tokens_list = [15, 17, 21, 25]
 
 # 1. prepare inputs
-processor = AutoProcessor.from_pretrained("/home/firefly/mnt/zt-rk3588-nn/expr/Florence-2-base-ft", trust_remote_code=True)
+processor = AutoProcessor.from_pretrained("..", trust_remote_code=True)
 
 # 2. prepare image
 image = Image.open("./test.jpg")
 original_image = image.copy()
 original_size = image.size
 # resize image to 768x768
-image = image.resize((768, 768))
+# image = image.resize((768, 768))
 # 3. prepare text
 prompt = "<MORE_DETAILED_CAPTION>"
 
@@ -56,16 +56,17 @@ for i in prompt_tokens_list:
         pad_to = i
         break
 print("pad_to: ", pad_to)
-inputs = processor(text=prompt, images=image, return_tensors="np", do_resize=False, padding="max_length", max_length=pad_to + 577, truncation=True)
+
+inputs = processor(text=prompt, images=image, return_tensors="np", do_resize=True, padding="max_length", max_length=pad_to + 577, truncation=True)
 for k, v in inputs.items():
     print(k, v.shape)
 
 # 4. run vision encoder using RKNN
 start_time = time.time()
-image_features0 = vision_encoder.run(None, {
-    "pixel_values": inputs["pixel_values"]
-})[0]
-image_features = rknn_vision_encoder.inference(inputs=[image_features0.reshape(1, 128, 1, 36864)])[0]
+# image_features0 = vision_encoder.run(None, {
+#     "pixel_values": inputs["pixel_values"]
+# })[0]
+image_features = rknn_vision_encoder.inference(inputs=[inputs["pixel_values"]], data_format="nchw")[0]
 
 end_time = time.time()
 vision_encoder_time = (end_time - start_time) * 1000
@@ -90,6 +91,7 @@ batch_size, image_token_length = image_features.shape[:-1]
 image_attention_mask = np.ones((batch_size, image_token_length))
 task_prefix_embeds = inputs_embeds
 task_prefix_attention_mask = np.ones((batch_size, task_prefix_embeds.shape[1]))
+# task_prefix_attention_mask = inputs["attention_mask"]
 if len(task_prefix_attention_mask.shape) == 3:
     task_prefix_attention_mask = task_prefix_attention_mask[:, 0]
 inputs_embeds = np.concatenate([image_features, task_prefix_embeds], axis=1)
@@ -135,7 +137,7 @@ while generated_tokens.__len__() < max_new_tokens:
     
     # 使用argmax选择下一个token (贪心算法)
     next_token = np.argmax(next_token_logits, axis=-1)[0]
-    print("next_token: ", next_token)
+    print("next_token: ", processor.decode([next_token]))
     # 将新生成的token添加到结果中
     generated_tokens.append(next_token)
 
@@ -220,7 +222,7 @@ def plot_bbox(image, data):
         font = ImageFont.load_default().font_variant(size=20)  # 如果Arial不可用，使用默认字体并放大
     
     # Plot each bounding box
-    for bbox, label in zip(data['bboxes'], data['labels']):
+    for bbox, label in zip(data['bboxes'], data.get('labels', data.get('bboxes_labels'))):
         # Unpack the bounding box coordinates
         x1, y1, x2, y2 = bbox
         # Draw the rectangle with thicker outline
@@ -312,14 +314,15 @@ def draw_ocr_bboxes(image, prediction, scale=1):
     # display(image)
     image.save("result_image.jpg")
 
-
-# draw_polygons(original_image, parsed_answer['<REFERRING_EXPRESSION_SEGMENTATION>'], fill_mask=True)
-# plot_bbox(original_image, parsed_answer[prompt.split(">")[0].strip() + ">"])
-# draw_ocr_bboxes(original_image, parsed_answer["<OCR_WITH_REGION>"], scale=1)
-
+if parsed_answer.get('<REFERRING_EXPRESSION_SEGMENTATION>'):
+    draw_polygons(original_image, parsed_answer['<REFERRING_EXPRESSION_SEGMENTATION>'], fill_mask=True)
+elif parsed_answer.get("<OCR_WITH_REGION>"):
+    draw_ocr_bboxes(original_image, parsed_answer["<OCR_WITH_REGION>"], scale=1)
+else:
+    plot_bbox(original_image, parsed_answer[prompt.split(">")[0].strip() + ">"])
 
 
 # Release RKNNLite instances
 rknn_vision_encoder.release()
 rknn_encoder.release()
-rknn_decoder_prefill.release()
+rknn_decoder_prefill.release()

onnx/vision_encoder.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:463a02cf1643c26a3414096f543a5f267ea49f384c1bcff7210cee2168912a4b
+size 261704579