initial working stable-diffusion pipeline

openvinotoolkit · Jan 3, 2025 · 9a1863b · 9a1863b
1 parent 62ec715
commit 9a1863b
Show file tree

Hide file tree

Showing 7 changed files with 43 additions and 20 deletions.
diff --git a/src/cpp/include/openvino/genai/image_generation/unet2d_condition_model.hpp b/src/cpp/include/openvino/genai/image_generation/unet2d_condition_model.hpp
@@ -28,7 +28,7 @@ class OPENVINO_GENAI_EXPORTS UNet2DConditionModel {
         int time_cond_proj_dim = -1;
 
         explicit Config(const std::filesystem::path& config_path);
-        //Config() = default;
+        Config() = default;
     };
 
     explicit UNet2DConditionModel(const std::filesystem::path& root_dir);
@@ -96,6 +96,7 @@ class OPENVINO_GENAI_EXPORTS UNet2DConditionModel {
         return guidance_scale > 1.0f && m_config.time_cond_proj_dim < 0;
     }
 
+
 private:
     class UNetInference;
     std::shared_ptr<UNetInference> m_impl;

diff --git a/src/cpp/src/image_generation/models/autoencoder_kl.cpp b/src/cpp/src/image_generation/models/autoencoder_kl.cpp
@@ -179,6 +179,7 @@ AutoencoderKL::AutoencoderKL(const std::string& vae_encoder_model,
     }
 }
 
+//AutoencoderKL::AutoencoderKL(const AutoencoderKL&) = default;
 AutoencoderKL::AutoencoderKL(const AutoencoderKL& original_model){
     std::cout << "bla1" << std::endl;
     encoder_compiled_model = original_model.encoder_compiled_model;
@@ -192,6 +193,7 @@ AutoencoderKL::AutoencoderKL(const AutoencoderKL& original_model){
     }
     m_encoder_model = original_model.m_encoder_model;
     m_decoder_model = original_model.m_decoder_model;
+    m_config = original_model.m_config;
 }
 
 AutoencoderKL& AutoencoderKL::reshape(int batch_size, int height, int width) {

diff --git a/src/cpp/src/image_generation/models/unet2d_condition_model.cpp b/src/cpp/src/image_generation/models/unet2d_condition_model.cpp
@@ -32,6 +32,7 @@ UNet2DConditionModel::UNet2DConditionModel(const std::filesystem::path& root_dir
     m_config(root_dir / "config.json") {
     ov::Core core = utils::singleton_core();
     m_model = core.read_model((root_dir / "openvino_model.xml").string());
+    //ov::save_model(m_model, "static_unet_model.xml");
     m_vae_scale_factor = get_vae_scale_factor(root_dir.parent_path() / "vae_decoder" / "config.json");
 }
 
@@ -61,7 +62,18 @@ UNet2DConditionModel::UNet2DConditionModel(const std::string& model,
     compile(device, properties);
 }
 
-UNet2DConditionModel::UNet2DConditionModel(const UNet2DConditionModel&) = default;
+UNet2DConditionModel::UNet2DConditionModel(const UNet2DConditionModel& original_model) {
+    m_config = original_model.m_config;
+    m_adapter_controller = original_model.m_adapter_controller;
+    m_model = original_model.m_model;
+    m_vae_scale_factor = original_model.m_vae_scale_factor;
+    if (typeid(m_impl) == typeid(UNet2DConditionModel::UNetInferenceStaticBS1)) {
+        m_impl = std::make_shared<UNet2DConditionModel::UNetInferenceStaticBS1>(original_model.m_impl->get_compiled_model());
+    } else {
+        m_impl = std::make_shared<UNet2DConditionModel::UNetInferenceDynamic>(original_model.m_impl->get_compiled_model());
+    }
+
+}
 
 const UNet2DConditionModel::Config& UNet2DConditionModel::get_config() const {
     return m_config;
@@ -84,7 +96,6 @@ UNet2DConditionModel& UNet2DConditionModel::compile(const std::string& device, c
     if (device == "NPU") {
         m_impl = std::make_shared<UNet2DConditionModel::UNetInferenceStaticBS1>();
     } else {
-        std::cout << "impl UNetInferenceDynamic" << std::endl;
         m_impl = std::make_shared<UNet2DConditionModel::UNetInferenceDynamic>();
     }
 
@@ -94,7 +105,6 @@ UNet2DConditionModel& UNet2DConditionModel::compile(const std::string& device, c
         m_adapter_controller = AdapterController(m_model, *adapters, device);
         m_impl->compile(m_model, device, *filtered_properties);
     } else {
-        std::cout << "impl compile" << std::endl;
         m_impl->compile(m_model, device, properties);
     }
 
@@ -106,7 +116,6 @@ UNet2DConditionModel& UNet2DConditionModel::compile(const std::string& device, c
 
 void UNet2DConditionModel::set_hidden_states(const std::string& tensor_name, ov::Tensor encoder_hidden_states) {
     OPENVINO_ASSERT(m_impl, "UNet model must be compiled first");
-    std::cout << "UNet2DConditionModel::set_hidden_states " << typeid(m_impl).name() << std::endl;
     m_impl->set_hidden_states(tensor_name, encoder_hidden_states);
 }
 

diff --git a/src/cpp/src/image_generation/models/unet_inference.hpp b/src/cpp/src/image_generation/models/unet_inference.hpp
@@ -15,6 +15,7 @@ class UNet2DConditionModel::UNetInference {
     virtual void set_hidden_states(const std::string& tensor_name, ov::Tensor encoder_hidden_states) = 0;
     virtual void set_adapters(AdapterController& adapter_controller, const AdapterConfig& adapters) = 0;
     virtual ov::Tensor infer(ov::Tensor sample, ov::Tensor timestep) = 0;
+    virtual std::shared_ptr<ov::CompiledModel> get_compiled_model() = 0;
 
     // utility function to resize model given optional dimensions.
     static void reshape(std::shared_ptr<ov::Model> model,
@@ -62,6 +63,8 @@ class UNet2DConditionModel::UNetInference {
 
         model->reshape(name_to_shape);
     }
+    UNetInference(const UNetInference & );
+    UNetInference() = default;
 };
 
 }  // namespace genai

diff --git a/src/cpp/src/image_generation/models/unet_inference_dynamic.hpp b/src/cpp/src/image_generation/models/unet_inference_dynamic.hpp
@@ -21,13 +21,11 @@ class UNet2DConditionModel::UNetInferenceDynamic : public UNet2DConditionModel::
 
         compiled_model = std::make_shared<ov::CompiledModel>(core.compile_model(model, device, properties));
         ov::genai::utils::print_compiled_model_properties(*compiled_model, "UNet 2D Condition dynamic model");
-        std::cout << "unet dynamic compile" << std::endl;
         m_request = compiled_model->create_infer_request();
     }
 
     virtual void set_hidden_states(const std::string& tensor_name, ov::Tensor encoder_hidden_states) override
     {
-        std::cout << "unet dynamic set_hidden_states" << std::endl;
         OPENVINO_ASSERT(m_request, "UNet model must be compiled first");
         m_request.set_tensor(tensor_name, encoder_hidden_states);
     }
@@ -40,27 +38,26 @@ class UNet2DConditionModel::UNetInferenceDynamic : public UNet2DConditionModel::
 
     virtual ov::Tensor infer(ov::Tensor sample, ov::Tensor timestep) override
     {
-        std::cout << "unet dynamic infer" << std::endl;
         OPENVINO_ASSERT(m_request, "UNet model must be compiled first. Cannot infer non-compiled model");
-        std::cout << sample.get_shape().to_string() << std::endl; 
         m_request.set_tensor("sample", sample);
-        std::cout << timestep.get_shape().to_string() << std::endl; 
         m_request.set_tensor("timestep", timestep);
         ov::CompiledModel test =  m_request.get_compiled_model();
-        std::cout << "compiled model ok" << test.inputs().at(0).get_any_name() << std::endl; 
+        ov::genai::utils::print_compiled_model_properties(test, "UNet 2D Condition TEST");
         m_request.infer();
-        std::cout << "unet dynamic inference complated" << std::endl; 
-
         return m_request.get_output_tensor();
     }
 
-    UNetInferenceDynamic(const UNetInferenceDynamic & origin_model){
-        OPENVINO_ASSERT(origin_model.compiled_model, "Source model must be compiled first");
-        compiled_model = origin_model.compiled_model;
+    UNetInferenceDynamic(std::shared_ptr<ov::CompiledModel> origin_compiled_model){
+        compiled_model = origin_compiled_model;
         m_request = compiled_model->create_infer_request();
     }
+
     UNetInferenceDynamic() = default;
 
+    std::shared_ptr<ov::CompiledModel> get_compiled_model(){
+        return compiled_model;
+    }
+
 private:
 
     ov::InferRequest m_request;

diff --git a/src/cpp/src/image_generation/models/unet_inference_static_bs1.hpp b/src/cpp/src/image_generation/models/unet_inference_static_bs1.hpp
@@ -14,14 +14,16 @@ namespace genai {
 class UNet2DConditionModel::UNetInferenceStaticBS1 : public UNet2DConditionModel::UNetInference {
 public:
     UNetInferenceStaticBS1() = default;
-    UNetInferenceStaticBS1(const UNetInferenceStaticBS1 & origin_model){
-        OPENVINO_ASSERT(origin_model.compiled_model, "Source model must be compiled first");
-        compiled_model = origin_model.compiled_model;
-        size_t m_native_batch_size = origin_model.m_native_batch_size;
+
+    UNetInferenceStaticBS1(const std::shared_ptr<ov::CompiledModel> & origin_compiled_model){
+        OPENVINO_ASSERT(origin_compiled_model, "Source model must be compiled first");
+        compiled_model = origin_compiled_model;
+        m_native_batch_size = compiled_model->input("sample").get_shape()[0];
         for (int i = 0; i < m_native_batch_size; i++) {
             m_requests[i] = compiled_model->create_infer_request();
         }
     }
+
     virtual void compile(std::shared_ptr<ov::Model> model,
                          const std::string& device,
                          const ov::AnyMap& properties) override {
@@ -145,6 +147,9 @@ class UNet2DConditionModel::UNetInferenceStaticBS1 : public UNet2DConditionModel
 
         return out_sample;
     }
+    std::shared_ptr<ov::CompiledModel> get_compiled_model(){
+        return compiled_model;
+    }
 
 private:
     std::shared_ptr<ov::CompiledModel> compiled_model;

diff --git a/src/cpp/src/utils.cpp b/src/cpp/src/utils.cpp
@@ -407,6 +407,12 @@ void print_compiled_model_properties(ov::CompiledModel& compiled_Model, const ch
                 std::cout << "  " << cfg << ": " << prop.as<std::string>() << std::endl;
             }
         }
+        for (const auto& input : compiled_Model.inputs()) {
+            std::cout << "Input name: " << input.get_any_name() << ", shape: " << input.get_partial_shape().to_string() << std::endl;
+        }
+        for (const auto& out : compiled_Model.outputs()) {
+            std::cout << "Output name: " << out.get_any_name() << ", shape: " << out.get_partial_shape().to_string() << std::endl;
+        }
 
         ov::Core core;
         std::vector<std::string> exeTargets;