diff --git a/demo/__init__.py b/demo/__init__.py
new file mode 100644
index 00000000..b7744fda
--- /dev/null
+++ b/demo/__init__.py
@@ -0,0 +1,3 @@
+from .predict import detect
+
+__all__ = ['detect']
\ No newline at end of file
diff --git a/requirements/cpu_requirements.txt b/requirements/cpu_requirements.txt
index 5dca6009..16030203 100644
--- a/requirements/cpu_requirements.txt
+++ b/requirements/cpu_requirements.txt
@@ -4,4 +4,5 @@ pillow == 9.5.0
 mindspore
 pylint
 pytest
-opencv-python
\ No newline at end of file
+opencv-python
+download
\ No newline at end of file
diff --git a/tests/modules/test_deploy_predict.py b/tests/modules/test_demo_predict.py
similarity index 72%
rename from tests/modules/test_deploy_predict.py
rename to tests/modules/test_demo_predict.py
index e2da271f..451630dd 100644
--- a/tests/modules/test_deploy_predict.py
+++ b/tests/modules/test_demo_predict.py
@@ -2,27 +2,35 @@
 
 sys.path.append(".")
 import os
-import numpy as np
 from download import download
 import cv2
 
-from mindspore import nn
-
-from mindyolo.utils.metrics import non_max_suppression, scale_coords, xyxy2xywh
+from mindyolo.utils.config import load_config, Config
 from mindyolo.data import COCO80_TO_COCO91_CLASS
-from deploy.infer_engine.lite import LiteModel
 from mindyolo.utils.utils import draw_result
-from deploy.predict import detect
+from demo.predict import detect
+from mindyolo.models.model_factory import create_model
 
-def test_deploy_predict():
+def test_deploy_predict(yaml_name):
     image_url = 'https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/image_cat.zip'
     path = download(image_url, './', kind="zip", replace=True)
     image_path = ('./image_cat/jpg/000000039769.jpg')
 
-    model_url = 'https://download.mindspore.cn/toolkits/mindyolo/yolov5/yolov5n_300e_mAP273-9b16bd7b-bd03027b.mindir'
-    model_path = download(model_url, './yolov5n.mindir', kind="file", replace=True)
-
-    network = LiteModel(model_path)
+    model_url = 'https://download.mindspore.cn/toolkits/mindyolo/yolov5/yolov5n_300e_mAP273-9b16bd7b.ckpt'
+    weight = download(model_url, './yolov5n.ckpt', kind="file", replace=True)
+    parent_dir = yaml_name[:6]
+    yaml_path = os.path.join('./configs', parent_dir, yaml_name)
+    cfg, _, _ = load_config(yaml_path)
+    cfg = Config(cfg)
+
+    network = create_model(
+        model_name=cfg.network.model_name,
+        model_cfg=cfg.network,
+        num_classes=cfg.data.nc,
+        sync_bn=False,
+        checkpoint_path=weight,
+    )
+    network.set_train(False)
     img = cv2.imread(image_path)
     result_dict = detect(
         network=network,
@@ -49,4 +57,4 @@ def test_deploy_predict():
 
 
 if __name__ == '__main__':
-    test_deploy_predict()
+    test_deploy_predict('yolov5n.yaml')