From 198bdeb025f491d7b87cec4599a5d6f7fd9227ac Mon Sep 17 00:00:00 2001 From: nyx-c-language Date: Thu, 24 Apr 2025 08:41:21 +0800 Subject: [PATCH] fix pir No.23 --- slm/model_zoo/ernie-vil2.0/deploy/python/infer.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/slm/model_zoo/ernie-vil2.0/deploy/python/infer.py b/slm/model_zoo/ernie-vil2.0/deploy/python/infer.py index 26b91ef6075a..a638f74b6ade 100644 --- a/slm/model_zoo/ernie-vil2.0/deploy/python/infer.py +++ b/slm/model_zoo/ernie-vil2.0/deploy/python/infer.py @@ -20,6 +20,10 @@ from PIL import Image from paddlenlp.transformers import ErnieViLProcessor +from paddlenlp.utils.env import ( + PADDLE_INFERENCE_MODEL_SUFFIX, + PADDLE_INFERENCE_WEIGHTS_SUFFIX, +) def parse_arguments(): @@ -76,11 +80,11 @@ def __init__(self, args): def create_fd_runtime(self, args): option = fd.RuntimeOption() if args.encode_type == "text": - model_path = os.path.join(args.model_dir, "get_text_features.pdmodel") - params_path = os.path.join(args.model_dir, "get_text_features.pdiparams") + model_path = os.path.join(args.model_dir, f"get_text_features{PADDLE_INFERENCE_MODEL_SUFFIX}") + params_path = os.path.join(args.model_dir, f"get_text_features{PADDLE_INFERENCE_WEIGHTS_SUFFIX}") else: - model_path = os.path.join(args.model_dir, "get_image_features.pdmodel") - params_path = os.path.join(args.model_dir, "get_image_features.pdiparams") + model_path = os.path.join(args.model_dir, f"get_image_features{PADDLE_INFERENCE_MODEL_SUFFIX}") + params_path = os.path.join(args.model_dir, f"get_image_features{PADDLE_INFERENCE_WEIGHTS_SUFFIX}") option.set_model_path(model_path, params_path) if args.device == "kunlunxin": option.use_kunlunxin()