[Bugfix] Fix some build errors for SD (#434)

* Update utils.cmake * Update vae.h * Update test_lite_sd_pipeline.cpp
DefTruth · Sep 4, 2024 · e3f653e · e3f653e
1 parent 6256f13
commit e3f653e
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 8 deletions.
diff --git a/cmake/utils.cmake b/cmake/utils.cmake
@@ -46,12 +46,15 @@ function(add_lite_ai_toolkit_shared_library version soversion)
         include(cmake/onnxruntime.cmake)
         set(LITE_SRCS ${LITE_SRCS} ${ORT_SRCS})
         set(LITE_DEPENDENCIES ${LITE_DEPENDENCIES} onnxruntime ddim_scheduler_cpp)
+        link_directories(${CMAKE_SOURCE_DIR}/lite/bin)
     endif ()
 
     if (ENABLE_TENSORRT)
         include(cmake/tensorrt.cmake)
         set(LITE_SRCS ${LITE_SRCS} ${TRT_SRCS})
-        set(LITE_DEPENDENCIES ${LITE_DEPENDENCIES}  cudart nvinfer nvonnxparser nvinfer_plugin ddim_scheduler_cpp)
+        set(LITE_DEPENDENCIES ${LITE_DEPENDENCIES} cudart nvinfer nvonnxparser 
+                                                   nvinfer_plugin ddim_scheduler_cpp)
+        link_directories(${CMAKE_SOURCE_DIR}/lite/bin)
     endif ()
 
     if (ENABLE_MNN)

diff --git a/examples/lite/sd/test_lite_sd_pipeline.cpp b/examples/lite/sd/test_lite_sd_pipeline.cpp
@@ -5,6 +5,7 @@
 
 static void test_default()
 {
+#ifdef ENABLE_ONNXRUNTIME
     std::string clip_onnx_path = "../../../examples/hub/onnx/sd/clip_model.onnx";
     std::string unet_onnx_path = "../../../examples/hub/onnx/sd/unet_model.onnx";
     std::string vae_onnx_path = "../../../examples/examples/hub/onnx/sd/vae_model.onnx";
@@ -21,12 +22,13 @@ static void test_default()
     pipeline->inference(prompt,negative_prompt,save_path,scheduler_config_path);
 
     delete pipeline;
-
+#endif
 }
 
 
 static void test_trt_pipeline()
 {
+#ifdef ENABLE_TENSORRT
     // 记录时间
     std::chrono::steady_clock::time_point start_time = std::chrono::steady_clock::now();
 
@@ -52,18 +54,18 @@ static void test_trt_pipeline()
     std::cout << "Elapsed time: " << elapsed_seconds.count() << " seconds" << std::endl;
 
     delete pipeline;
+#endif
 
 }
 
 static void test_lite()
 {
+    test_default();
     test_trt_pipeline();
-
-//    test_default();
 }
 
 int main()
 {
     test_lite();
     return 0;
-}
+}
diff --git a/lite/ort/sd/vae.h b/lite/ort/sd/vae.h
@@ -4,9 +4,6 @@
 
 #ifndef LITE_AI_TOOLKIT_VAE_H
 #define LITE_AI_TOOLKIT_VAE_H
-#include "onnxruntime_cxx_api.h"
-#include "opencv4/opencv2/opencv.hpp"
-#include "iostream"
 #include "lite/ort/core/ort_core.h"
 
 namespace ortsd