Add SD3 LoRA Adapter Support (openvinotoolkit#2187)

**Details:** 
- Add SD3 LoRA Adapter Support for Text2ImagePipeline
- Verified model: 
- Base model:
https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers
   - LoRA: https://civitai.com/models/515022/sd3-first-anime-lora-test
- Test step:
- Model conversion: `optimum-cli export openvino --model
stabilityai/stable-diffusion-3-medium-diffusers --weight-format fp16
stable-diffusion-3-medium-diffusers-ov/FP16 --task text-to-image
--trust-remote-code`
- Model inference: `build/samples/cpp/image_generation/lora_text2image
stable-diffusion-3-medium-diffusers-ov/FP16/ "A woman with black hair,
green eyes, leaking flowing light, in front of a black background, she
wears a su it, necklace, gold earrings, only eyes have color"
SD3-First_Anime_Lora_test/Highstep/bl_128_1600.safetensors 0.7`

**Tickets:** 
CVS-156874

---------

Co-authored-by: Ilya Lavrenov <ilya.lavrenov@intel.com>

Author: Xiake Sun

SUPPORTED_MODELS.md

@@ -263,7 +263,7 @@ The pipeline can work with other similar topologies produced by `optimum-intel`
       <td>Supported</td>
       <td>Supported</td>
       <td>Supported</td>
-      <td>Not supported</td>
+      <td>Partially supported</td>
       <td>
         <ul>
           <li><a href="https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers"><code>stabilityai/stable-diffusion-3-medium-diffusers</code></a></li>

src/cpp/include/openvino/genai/image_generation/sd3_transformer_2d_model.hpp

@@ -11,8 +11,8 @@
 #include "openvino/runtime/infer_request.hpp"
 #include "openvino/runtime/properties.hpp"
 #include "openvino/runtime/tensor.hpp"
-
 #include "openvino/genai/visibility.hpp"
+#include "openvino/genai/lora_adapter.hpp"
 
 namespace ov {
 namespace genai {
@@ -75,6 +75,8 @@ class OPENVINO_GENAI_EXPORTS SD3Transformer2DModel {
         return compile(device, ov::AnyMap{std::forward<Properties>(properties)...});
     }
 
+    void set_adapters(const std::optional<AdapterConfig>& adapters);
+
     void set_hidden_states(const std::string& tensor_name, ov::Tensor encoder_hidden_states);
 
     ov::Tensor infer(const ov::Tensor latent, const ov::Tensor timestep);
@@ -87,6 +89,7 @@ class OPENVINO_GENAI_EXPORTS SD3Transformer2DModel {
     ov::InferRequest m_request;
     std::shared_ptr<ov::Model> m_model;
     size_t m_vae_scale_factor;
+    AdapterController m_adapter_controller;
 
     class InferenceDynamic;
     class InferenceStaticBS1;

src/cpp/src/image_generation/flux_pipeline.hpp

@@ -636,7 +636,7 @@ class FluxPipeline : public DiffusionPipeline {
         }
     }
 
-    // Returns non-empty updated adapters iff they are required to be updated
+    // Returns non-empty updated adapters if they are required to be updated
     static std::optional<AdapterConfig> derived_adapters(const AdapterConfig& adapters) {
         return ov::genai::derived_adapters(adapters, flux_adapter_normalization);
     }

src/cpp/src/image_generation/models/sd3_transformer_2d_model.cpp

@@ -9,6 +9,7 @@
 
 #include "json_utils.hpp"
 #include "utils.hpp"
+#include "lora/helper.hpp"
 
 namespace ov {
 namespace genai {
@@ -86,6 +87,12 @@ SD3Transformer2DModel& SD3Transformer2DModel::reshape(int batch_size,
 
 SD3Transformer2DModel& SD3Transformer2DModel::compile(const std::string& device, const ov::AnyMap& properties) {
     OPENVINO_ASSERT(m_model, "Model has been already compiled. Cannot re-compile already compiled model");
+    std::optional<AdapterConfig> adapters;
+    auto filtered_properties = extract_adapters_from_properties(properties, &adapters);
+    if (adapters) {
+        adapters->set_tensor_name_prefix(adapters->get_tensor_name_prefix().value_or("transformer"));
+        m_adapter_controller = AdapterController(m_model, *adapters, device);
+    }
 
     if (device.find("NPU") != std::string::npos) {
         m_impl = std::make_shared<SD3Transformer2DModel::InferenceStaticBS1>();
@@ -94,14 +101,21 @@ SD3Transformer2DModel& SD3Transformer2DModel::compile(const std::string& device,
         m_impl = std::make_shared<SD3Transformer2DModel::InferenceDynamic>();
     }
 
-    m_impl->compile(m_model, device, properties);
+    m_impl->compile(m_model, device, *filtered_properties);
 
     // release the original model
     m_model.reset();
 
     return *this;
 }
 
+void SD3Transformer2DModel::set_adapters(const std::optional<AdapterConfig>& adapters) {
+    OPENVINO_ASSERT(m_impl, "Transformer model must be compiled first");
+    if(adapters) {
+        m_impl->set_adapters(m_adapter_controller, *adapters);
+    }
+}
+
 void SD3Transformer2DModel::set_hidden_states(const std::string& tensor_name, ov::Tensor encoder_hidden_states) {
     OPENVINO_ASSERT(m_impl, "Transformer model must be compiled first");
     m_impl->set_hidden_states(tensor_name, encoder_hidden_states);

src/cpp/src/image_generation/models/sd3transformer_2d_inference.hpp

@@ -14,6 +14,7 @@ class SD3Transformer2DModel::Inference {
 public:
     virtual void compile(std::shared_ptr<ov::Model> model, const std::string& device, const ov::AnyMap& properties) = 0;
     virtual void set_hidden_states(const std::string& tensor_name, ov::Tensor encoder_hidden_states) = 0;
+    virtual void set_adapters(AdapterController& m_adapter_controller, const AdapterConfig& adapters) = 0;
     virtual ov::Tensor infer(ov::Tensor latent_model_input, ov::Tensor timestep) = 0;
 
     // utility function to resize model given optional dimensions.

src/cpp/src/image_generation/models/sd3transformer_2d_inference_dynamic.hpp

@@ -19,6 +19,11 @@ class SD3Transformer2DModel::InferenceDynamic : public SD3Transformer2DModel::In
         m_request = compiled_model.create_infer_request();
     }
 
+    virtual void set_adapters(AdapterController& m_adapter_controller, const AdapterConfig& adapters) override {
+        OPENVINO_ASSERT(m_request, "Transformer model must be compiled first");
+        m_adapter_controller.apply(m_request, adapters);
+    }
+
     virtual void set_hidden_states(const std::string& tensor_name, ov::Tensor encoder_hidden_states) override {
         OPENVINO_ASSERT(m_request, "Transformer model must be compiled first");
         m_request.set_tensor(tensor_name, encoder_hidden_states);

src/cpp/src/image_generation/models/sd3transformer_2d_inference_static_bs1.hpp

@@ -45,6 +45,13 @@ class SD3Transformer2DModel::InferenceStaticBS1 : public SD3Transformer2DModel::
         }
     }
 
+    virtual void set_adapters(AdapterController& m_adapter_controller, const AdapterConfig& adapters) override {
+        for (auto& m_request : m_requests) {
+            OPENVINO_ASSERT(m_request, "Transformer model must be compiled first");
+            m_adapter_controller.apply(m_request, adapters);
+        }
+    }
+
     virtual void set_hidden_states(const std::string& tensor_name, ov::Tensor encoder_hidden_states) override {
         OPENVINO_ASSERT(m_native_batch_size && m_native_batch_size == m_requests.size(),
                         "Transformer model must be compiled first");

src/cpp/src/image_generation/stable_diffusion_3_pipeline.hpp

@@ -13,6 +13,7 @@
 #include "openvino/genai/image_generation/sd3_transformer_2d_model.hpp"
 
 #include "utils.hpp"
+#include "lora/helper.hpp"
 
 namespace {
 
@@ -136,22 +137,19 @@ class StableDiffusion3Pipeline : public DiffusionPipeline {
         using utils::read_json_param;
 
         set_scheduler(Scheduler::from_config(root_dir / "scheduler/scheduler_config.json"));
-
         const std::string text_encoder = data["text_encoder"][1].get<std::string>();
         if (text_encoder == "CLIPTextModelWithProjection") {
             m_clip_text_encoder_1 =
                 std::make_shared<CLIPTextModelWithProjection>(root_dir / "text_encoder", device, properties);
         } else {
             OPENVINO_THROW("Unsupported '", text_encoder, "' text encoder type");
         }
-
         const std::string text_encoder_2 = data["text_encoder_2"][1].get<std::string>();
         if (text_encoder_2 == "CLIPTextModelWithProjection") {
             m_clip_text_encoder_2 = std::make_shared<CLIPTextModelWithProjection>(root_dir / "text_encoder_2", device, properties);
         } else {
             OPENVINO_THROW("Unsupported '", text_encoder_2, "' text encoder type");
         }
-
         const auto text_encoder_3_json = data["text_encoder_3"][1];
         if (!text_encoder_3_json.is_null()) {
             const std::string text_encoder_3 = text_encoder_3_json.get<std::string>();
@@ -161,7 +159,6 @@ class StableDiffusion3Pipeline : public DiffusionPipeline {
                 OPENVINO_THROW("Unsupported '", text_encoder_3, "' text encoder type");
             }
         }
-
         const std::string transformer = data["transformer"][1].get<std::string>();
         if (transformer == "SD3Transformer2DModel") {
             m_transformer = std::make_shared<SD3Transformer2DModel>(root_dir / "transformer", device, properties);
@@ -184,7 +181,6 @@ class StableDiffusion3Pipeline : public DiffusionPipeline {
 
         // initialize generation config
         initialize_generation_config(data["_class_name"].get<std::string>());
-
         update_adapters_from_properties(properties, m_generation_config.adapters);
     }
 
@@ -456,7 +452,13 @@ class StableDiffusion3Pipeline : public DiffusionPipeline {
     }
 
     void set_lora_adapters(std::optional<AdapterConfig> adapters) override {
-        OPENVINO_THROW("LORA adapters are not implemented for Stable Diffusion 3 yet");
+        if(adapters) {
+            if(auto updated_adapters = derived_adapters(*adapters)) {
+                adapters = updated_adapters;
+            }
+            // TODO: Add LoRA Adapter support for text encoders
+            m_transformer->set_adapters(adapters);
+        }
     }
 
     ov::Tensor generate(const std::string& positive_prompt,
@@ -486,6 +488,8 @@ class StableDiffusion3Pipeline : public DiffusionPipeline {
 
         check_inputs(generation_config, initial_image);
 
+        set_lora_adapters(generation_config.adapters);
+
         // 3. Prepare timesteps
         m_scheduler->set_timesteps(generation_config.num_inference_steps, generation_config.strength);
 
@@ -586,6 +590,12 @@ class StableDiffusion3Pipeline : public DiffusionPipeline {
         return m_perf_metrics;
     }
 
+protected:
+    // Returns non-empty updated adapters if they are required to be updated
+    static std::optional<AdapterConfig> derived_adapters(const AdapterConfig& adapters) {
+        return ov::genai::derived_adapters(adapters, flux_adapter_normalization);
+    }
+
 private:
     size_t get_config_in_channels() const override {
         assert(m_transformer != nullptr);

src/cpp/src/lora/adapter.cpp

@@ -1069,7 +1069,6 @@ struct AdapterControllerImpl {
         }
 
         auto state = infer_request.query_state();
-
         // TODO: Forced to use variable_id instead of index to address the state tensors, require the same order for state as for variables from plugins
 
         // Convert LoRAVarIDs to LoRAIndices to speedup search for state with a given name